Writing OP Implementation in Python - 3.0 English

Vitis AI Library User Guide (UG1354)

Document ID
Release Date
3.0 English

An add OP is used as an example in this section. For more Python examples, refer to TensorFlow2 and PyTorch examples.

  1. Create a Python package and a module file.
    % mkdir -p /tmp/demo_add_op/
    % cd /tmp/demo_add_op/
    % mkdir vart_op_imp/
    % touch vart_op_imp/__init__.py
    % touch vart_op_imp/add.py
    • The folder name should be vart_op_imp.
    • If you have more than one custom OP, create them separately and place them in the vart_op_imp folder.
  2. Update add.py as shown in the following code snippet.
    import numpy as np
    class add:
        def __init__(self, op):
        def calculate(self, output, input):
            np_output = np.array(output, copy=False)
            L = len(input)
            if L == 0:
            np_input_0 = np.array(input[0], copy=False)
            np.copyto(np_output, np_input_0)
            for i in range(1, L):
                np_output = np.add(np_output,
                                   np.array(input[i], copy=False),

    You can also use the following simplified version.

    import numpy as np
    class add:
        def __init__(self, op):
        def calculate(self, output, input):
            np_output = np.array(output, copy=False)
            L = len(input)
            assert L == 2
            np_input_0 = np.array(input[0], copy=False)
            np_input_1 = np.array(input[1], copy=False)
            np_output = np.add(np_input_0, np_input_1, out=np_output)
  3. Install the op as shown in the following code snippet.
    % mkdir -p lib
    % ln -sf ~/.local/Ubuntu.18.04.x86_64.Debug/lib/libvart_op_imp_python-cpu-op.so lib/libvart_op_imp_add.so
    % ls -l lib
    % mkdir -p dump
    % env LD_LIBRARY_PATH=$HOME/.local/Ubuntu.18.04.x86_64.Debug/lib:$PWD/lib $HOME/.local/Ubuntu.18.04.x86_64.Debug/share/vitis_ai_library/test/cpu_task/test_op_imp --graph /tmp/add.xmodel --op "add_op"

For Edge, you can execute the following command to install and test the op.

# ls -sf /usr/lib/libvart_op_imp_python-cpu-op.so /usr/lib/libvart_op_imp_add.so
# cp -r vart_op_imp /usr/lib.python3.9/site-packages
# xdputil run_op add.xmodel add_op -r ref -d dump
Note: Before you execute the above commands, copy add.xmodel, the ref folder with the sample input files, and the vart_op_imp folder to the board.

Similar to the C++ interface, the Python module must have a class whose name is the same as the type of xir::Op. In this example, add is used. The class must have a constructor with a single argument op in addition to self. It is an XIR::Op. Refer to the XIR Python API for more information. Similarly, the class add must have a member function called calculate, in addition to the self argument. The first argument must be named output, and following argument names must comply with the XIR::OpDef associated with the XIR::Op, refer to XIR API for more detail.

Note: A symbolic link to libvart_op_imp_python-cpu-op.so named libvart_op_imp_add.so is created. libvart_op_imp_python-cpu-op.so is a bridge between Python and C++. From the C++ side, it searches for libvart_op_imp_add.so and finds the libvart_op_imp_python-cpu-op.so. In libvart_op_imp_python_cpu_op.so, the Python module name vart_op_imp.add is imported and Python searches for the module as usual.