AI Core custom operator loading

Kernel functions can deliver tasks in <<<>>> mode, which features concise code and good readability.

The following is a code sample of key steps, which is for reference only and cannot be directly copied for compilation and running:

// Device code
Template<>
extern "C" __global__ __aicore__ void add_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z)
{
    KernelAdd op;
    op.Init(x, y, z);
    op.Process();
}

int main()
{
    int N = ...;
    size_t size = N * sizeof(uint64);

    // Initialize
    aclrtSetDevice(0);

    // Create stream
    aclrtStream stream;
    aclrtCreateStream(&stream);

    // Allocate vectors in host memory
    void *h_x, *h_y, *h_z;
    aclrtMallocHost(&h_x, size);
    aclrtMallocHost(&h_y, size);
    aclrtMallocHost(&h_z, size);

    // Initialize input vectors
    ...

    // Allocate vectors in device memory
    void *d_x, *d_y, *d_z;
    aclrtMalloc(&d_x, size, ACL_MEM_MALLOC_HUGE_FIRST);
    aclrtMalloc(&d_y, size, ACL_MEM_MALLOC_HUGE_FIRST);
    aclrtMalloc(&d_z, size, ACL_MEM_MALLOC_HUGE_FIRST);

    // Copy vectors from host memory to device memory
    aclrtMemcpy(d_x, size, h_x, size, ACL_MEMCPY_HOST_TO_DEVICE);
    aclrtMemcpy(d_y, size, h_y, size, ACL_MEMCPY_HOST_TO_DEVICE);

    // Invoke kernel
    uint32_t numBlocks = 48;
    add_custom<<<numBlocks, nullptr, stream>>>(d_x, d_y, d_z);
    ...
}

You can also use the LaunchKernel API (for example, aclrtLaunchKernelWithHostArgs) provided by Runtime to deliver the kernel function. Before using this method, you need to understand the concepts of binary and function.

Binary: a dynamically loaded code container unit that contains compiled kernel code and global variables. You can use aclrtBinaryLoadFromFile or aclrtBinaryLoadFromData to load the compiled operator binary file to the NPU and obtain the corresponding binary handle.
Function: a specific executable kernel function defined in the binary file. It is the entry point that can be called by the host code and executed on the NPU. You can use aclrtBinaryGetFunction to obtain the function handle corresponding to the kernel function.

The following is a code sample of key steps for using the LaunchKernel API, which is for reference only and cannot be directly copied for compilation and running:

// Device code
extern "C" __global__ __aicore__ void add_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z)
{
    KernelAdd op;
    op.Init(x, y, z);
    op.Process();
}

int main()
{
    int N = ...;
    size_t size = N * sizeof(uint64);

    // Initialize
    aclrtSetDevice(0);

    // Create stream
    aclrtStream stream;
    aclrtCreateStream(&stream);

    // Create binary from binary file
    aclrtBinHandle bin;
    aclrtBinaryLoadFromFile("add_custom.o", nullptr, &bin);

    // Get function handle from binary
    aclrtFuncHandle add_custom;
    aclrtBinaryGetFunction(bin, "add_custom", &add_custom);

    // Allocate vectors in host memory
    void *h_x, *h_y, *h_z;
    aclrtMallocHost(&h_x, size);
    aclrtMallocHost(&h_y, size);
    aclrtMallocHost(&h_z, size);

    // Initialize input vectors
    ...

    // Allocate vectors in device memory
    void *d_x, *d_y, *d_z;
    aclrtMalloc(&d_x, size, ACL_MEM_MALLOC_HUGE_FIRST);
    aclrtMalloc(&d_y, size, ACL_MEM_MALLOC_HUGE_FIRST);
    aclrtMalloc(&d_z, size, ACL_MEM_MALLOC_HUGE_FIRST);

    // Copy vectors from host memory to device memory
    aclrtMemcpy(d_x, size, h_x, size, ACL_MEMCPY_HOST_TO_DEVICE);
    aclrtMemcpy(d_y, size, h_y, size, ACL_MEMCPY_HOST_TO_DEVICE);

    // Invoke kernel
    uint32_t numBlocks = 48;
    void* args[] = {d_x, d_y, d_z};
    size_t argsSize = 3 * sizeof(void*);
    aclrtLaunchKernelWithHostArgs(add_custom, numBlocks, stream, nullptr, args, argsSize, nullptr, 0);
    ...
}

Parent topic: CANN custom operator call