• <xmp id="om0om">
  • <table id="om0om"><noscript id="om0om"></noscript></table>
  • 數據中心/云端

    使用 Numbast 實現 CUDA C++ 生態系統與 Python 開發者之間的無縫連接

    通過支持使用 Python 編寫 CUDA 內核函數,類似于在 C++中實現內核函數的方式,Numba 彌合了 Python 生態系統與 CUDA 性能之間的差距。

    但是,CUDA C++開發者可以訪問許多目前未在 Python 中公開的庫,包括 CUDA 核心計算庫(CCCL)、cuRAND 以及頭文件實現的數字類型,例如 bfloat16 等。

    雖然每個 CUDA C++ 庫都可以用自己的方式介紹給 Python,但是手動為每個庫進行綁定是一項費力、重復的工作,并且容易出現不一致。例如,float16 和 bfloat16 數據類型定義了 60 多個類似的獨立函數,這兩種類型都需要多次類似的綁定。

    此外,當底層 CUDA C++庫引入新功能時,手動創建的綁定通常會不同步。

    解決方案:Numbast?

    Numba 建立自動化工作流,將 CUDA C/C++API 轉換為 Numba 綁定。

    高級別的頂層聲明從 CUDA C++ 頭文件中讀取、序列化并傳遞至 Python API。然后,Numba 綁定生成器會迭代這些聲明,并為每個 API 生成 Numba 擴展程序。

    演示:C++聲明簡單結構?

    為展示 Numbast 的實際應用,以下示例展示了如何為 demo myfloat16 類型創建 Numba 綁定。這些 C++ 聲明以 CUDA float16 頭文件中的聲明為靈感,提供了一個簡化版本來演示在實踐中生成的綁定。

    C++聲明?

    此演示使用 C++ 語法顯示了以下元素:

    // 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);
    • A struct declaration, which has
      • 設備構造函數
      • 一些設備方法,包括轉換和算術運算符
    • 兩種函數聲明:算術運算符重載和平方根函數。

    有關 Numbast 中支持的語言功能的更多信息,請參閱 支持的 CUDA C++ 聲明

    使用 Numbast 設置腳本?

    Numbast 的使用通常涉及兩個步驟:

    1. 使用 AST_Canopy 解析頭文件。
    2. 從解析后的報文頭生成 Numba 綁定。

    以下代碼示例通過實施以下兩個步驟來設置 Numba 綁定:

    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])

    數據模型 是 Numba 表示底層數據的不同方式。用于 myfloat16PrimitiveModel 模型非常適合標量。StructModel 模型(此處未使用)適用于類和結構。其他數據模型的使用較少。

    以最自然的方式使用?

    在 CUDA C++中,您可以構建一個 myfloat16 對象并像下面這樣使用它:

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

    在 Numba 內核中,您可以按原樣使用它們:

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

    得益于 Numba 中的類型推斷,代碼甚至比原始 C++ 更簡潔。

    第一個支持的綁定:bfloat16 數據類型?

    Numbast 支持的第一個 Numba 綁定是一種新的 bfloat16 數據類型。bfloat16 數據類型可以與 PyTorch 的 torch.bfloat16 數據類型進行互操作,因此您可以使用這種新數據類型高效地開發自定義計算內核。

    以下代碼示例展示了如何使用新的 bfloat16 數據類型開發對 Torch 張量執行計算的 Numba 核函數。它將 PyTorch 數組傳遞到 Numba 計算核函數中,并使用通過 Numba 綁定的 CUDA 內部函數執行數學運算。

    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-forge 下載 Numbast 和 bfloat16 Numba 綁定:

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

    架構?

    Numbast 由兩個組件組成:

    • AST_Canopy:解析和序列化 C++頭文件的底層
    • Numbast:面向用戶的層,可使用解析后的結果并動態構建 Numba 綁定。

    AST_Canopy:聲明解析器?

    在森林生態學中, 樹冠 是指森林生態區的上層。AST_Canopy 是一個軟件包,用于檢查抽象語法樹森林中的頂層聲明,從中提取信息并將其傳遞給 Python 層。在這里,頂層是指 CUDA C++ 庫向用戶公開的面向用戶的公共 API。

    Diagram shows that the C++ core layer is libastcanopy, which depends on ClangTooling, Its functionalities are bridged to AST_Canopy through the binding layer Pylibastcanopy.
    圖 1.的分層架構 AST_Canopy

    圖 1 顯示了 AST_Canopy 的架構:

    • clangTooling:用于支持編寫獨立工具(如 Numbast)的 Clang 庫。
    • libastcanopy:使用 clangTooling 實現聲明解析邏輯。
    • pylibastcanopy:在 Python 中直接公開 libastcanopy API 的綁定。
    • AST_Canopypylibastcanopy 之上的層,可提供愉悅的 Python 用戶體驗。

    除了報文頭解析和序列化之外,AST_Canopy 還提供以下特性:

    • 運行時環境檢測: 自動檢測通過 Conda 包安裝的 libstdcxx 和 CUDA 頭文件,并相應地設置 clang 編譯器。
    • 計算能力解析的靈活性:支持根據不同的計算能力配置 AST 解析。一些頭文件會根據計算能力有條件地公開代碼,此功能支持頭文件序列化和運行時環境不同的情況。

    Numbast:綁定生成器?

    Numbast 位于 AST_Canopy 的下游,AST_Canopy 會使用聲明信息并自動生成 Numba 綁定。Numbast 的存在目的是在 C++ 和 Python 語法之間提供翻譯層。正如演示所示,大多數簡單的 C++ 語法都在 Python 中找到對應的自然語言(表 1)。

    運營 CUDA C++ Numba
    對象構建 auto hpi = myfloat16(3.14) hpi = myfloat16(3.14)
    屬性訪問 auto data = hpi.data data = hpi.data
    函數調用 auto r = hsqrt(hpi) r = hsqrt(hpi)
    類型轉換 auto fpi = float(hpi); fpi = types.float32(hpi)
    算術運算 auto pi2 = hpi + hpi pi2 = hpi + hpi
    表 1. CUDA C++ 語法與 Python 語法的映射

    Numba 的類型系統與 C 和 C++ 類語言有許多共同之處。當然,還有一些 Python 中不存在的功能,例如指針語義和基于模板的元編程。

    Numbast 是封裝 C++ 和 Python 異同的中間層。

    降低:全局?

    使用 Numbast 生成的綁定可通過 Numba 中名為 外部函數調用 (FFI)的功能降低。可在原生 CUDA 函數調用上生成與 Numba ABI 兼容的 shim 函數,然后使用 NVRTC 進行編譯。預計性能與 CUDA C++ 開發者相同,但需減去 FFI 的性能。

    未來版本的 Numba-cuda 將引入 鏈路時間優化(LTO) 支持,進一步消除加速 Numba 內核與原生 CUDA C++ 之間的性能差距。

    注意事項?

    AST_Canopy 和 Numbast 都有值得注意的注意事項。AST_Canopy 依賴于 clangToolingclangTooling 尚不支持的新 CUDA 語言功能和依賴新語言功能的庫可能無法正確解析。然而,大多數庫和頭文件都使用 clangTooling 支持的功能。

    結束語?

    在本文中,我們介紹了新的 Numba 綁定生成工具 Numbast。我們展示了通過使用 Numbast,您可以快速從不斷增長的 CUDA C++ 功能集中受益。

    Numbast v0.1.0 為 Numba 提供了新的數據類型 bfloat16。您可以期待看到 Numbast 生成的更多綁定,包括新的數據類型、NVSHMEM 綁定和 CCCL 綁定。

    ?

    0

    標簽

    人人超碰97caoporen国产