注意：本教程仍在持续更新中


    在本教程中，我们将展示：

    • 如何用C ++和CUDA编写您的算子并对其进行即时编译
    • 运行您的自定义算子
      如果您想用几行代码来实现一个非常简单的算子，请使用code运算，请参阅help(jt.code).custom_op用于实现复杂的算子。 custom_op和内置运算的功能完全相同。
    1. import jittor as jt
    3. header ="""
    4. #pragma once
    5. #include "op.h"
    7. namespace jittor {
    9. struct CustomOp : Op {
    10.     Var* output;
    11.     CustomOp(NanoVector shape, NanoString dtype=ns_float);
    13.     const char* name() const override { return "custom"; }
    14.     DECLARE_jit_run;
    15. };
    17. } // jittor
    18. """
    20. src = """
    21. #include "var.h"
    22. #include "custom_op.h"
    24. namespace jittor {
    25. #ifndef JIT
    26. CustomOp::CustomOp(NanoVector shape, NanoString dtype) {
    27.     flags.set(NodeFlags::_cuda, 1);
    28.     flags.set(NodeFlags::_cpu, 1);
    29.     output = create_output(shape, dtype);
    30. }
    32. void CustomOp::jit_prepare() {
    33.     add_jit_define("T", output->dtype());
    34. }
    36. #else // JIT
    37. #ifdef JIT_cpu
    38. void CustomOp::jit_run() {
    39.     index_t num = output->num;
    40.     auto* __restrict__ x = output->ptr<T>();
    41.     for (index_t i=0; i<num; i++)
    42.         x[i] = (T)i;
    43. }
    44. #else
    45. // JIT_cuda
    46. __global__ void kernel(index_t n, T *x) {
    47.     int index = blockIdx.x * blockDim.x + threadIdx.x;
    48.     int stride = blockDim.x * gridDim.x;
    49.     for (int i = index; i < n; i += stride)
    50.         x[i] = (T)-i;
    51. }
    53. void CustomOp::jit_run() {
    54.     index_t num = output->num;
    55.     auto* __restrict__ x = output->ptr<T>();
    56.     int blockSize = 256;
    57.     int numBlocks = (num + blockSize - 1) / blockSize;
    58.     kernel<<<numBlocks, blockSize>>>(num, x);
    59. }
    60. #endif // JIT_cpu
    61. #endif // JIT
    63. } // jittor
    64. """
    66. my_op = jt.compile_custom_op(header, src, "custom", warp=False)

    让我们查看一下这个运算的结果。

    1. # run cpu version
    2. jt.flags.use_cuda = 0
    3. a = my_op([3,4,5], 'float').fetch_sync()
    4. assert (a.flatten() == range(3*4*5)).all()
    6. if jt.compiler.has_cuda:
    7.     # run cuda version
    8.     jt.flags.use_cuda = 1
    9.     a = my_op([3,4,5], 'float').fetch_sync()
    10.     assert (-a.flatten() == range(3*4*5)).all()