使用 Numbast 实现 CUDA C++ 生态系统与 Python 开发者之间的无缝连接

// demo.cuh
struct __attribute__((aligned(2))) myfloat16
{public:
half data;
__host__ __device__ myfloat16();

__host__ __device__ myfloat16(double val);

__host__ __device__ operator float() const;
};
__host__ __device__ myfloat16 operator+(const myfloat16 &lh, const myfloat16 &rh);

__host__ __device__ myfloat16 hsqrt(const myfloat16 a);

import os
from ast_canopy import parse_declarations_from_source
from numbast import bind_cxx_struct, bind_cxx_function, MemoryShimWriter

from numba import types, cuda
from numba.core.datamodel.models import PrimitiveModel

import numpy as np
# Step 1:
# Use `AST_Canopy` to parse demo.cuh as AST, read all declarations from it.
source = os.path.join(os.path.dirname(__file__), "demo.cuh")
# Assume that you want to generate bindings for a machine with "sm_80"
# capability.
structs, functions, *_ = parse_declarations_from_source(source, , "sm_80")

shim_writer = MemoryShimWriter(f'#include "{source}"')
# Step 2:
# Make Numba bindings from the declarations.
# New type "myfloat16" is a Number type, data model is `PrimitiveModel`.
myfloat16 = bind_cxx_struct(shim_writer, structs[0], types.Number, PrimitiveModel)# bind_cxx_function returns the generated bindings to the C++ declaration.# The first function binds to an operator, and it’s bound to `operator.add`. You can directly use `myfloat16 + myfloat16` in kernels.
bind_cxx_function(shim_writer, functions[0])# The second function is `hsqrt`, with which Numbast creates a new Python handle and returns it in the return value.
hsqrt = bind_cxx_function(shim_writer, functions[1])

__global__ void kernel()
{
  auto one = myfloat16(1.0);
  auto two = myfloat16(2.0);
  auto three = one + two;
  auto sqrt3 = hsqrt(three);
}

@cuda.jit(link=shim_writer.links())
def kernel():
  one = myfloat16(1.0)
  two = myfloat16(2.0)
  three = one + two
  sqrt3 = hsqrt(three)

from numba import float32
import numba.cuda as cuda
import torch
from numbast_extensions.bf16 import get_shims, hsin, nv_bfloat16


@cuda.jit(link=get_shims())
def torch_add_sin(a, b, out):
  i, j = cuda.grid(2)
  if i < out.shape[0] and j < out.shape[1]:
      # Arithmetic of bfloat16 type
      sum = a[i, j] + b[i, j]
      # Bfloat16 native intrinsics
      sin_of_sum = hsin(sum)
      # bf16 to f32 upcast
      f32 = float32(sin_of_sum)
      # f32 to bf16 downcast
      bf16 = nv_bfloat16(f32)
      # Assignment to external array
      out[i, j] = bf16


a = torch.ones([2, 2], device=torch.device("cuda:0"), dtype=torch.bfloat16)
b = torch.ones([2, 2], device=torch.device("cuda:0"), dtype=torch.bfloat16)
expected = torch.sin(a + b)

out = torch.zeros([2, 2], device=torch.device("cuda:0"), dtype=torch.bfloat16)

threadsperblock = (16, 16)
blockspergrid = (1, 1)
torch_add_sin[blockspergrid, threadsperblock](a, b, out)
assert torch.equal(expected, out)

conda install -c nvidia -c rapidsai -c conda-forge ml_dtypes numbast-extensions