bind_buffer

1.硬件背景

将tensor绑定到buffer上

2.接口说明

Python def to_buffer( tensor: tl.tensor, space: address_space = None, bind_buffer: buffer = None, _builder=None ) -> buffer:

2.1 入参

参数名	类型	必需	说明
tensor	tl.tensor	是	要转换的tensor
address_space	bl.address_space	否	buffer所在的地址空间
bind_buffer	bl.buffer	否	需要绑定到的target buffer

2.2 返回值

如果使用bind_buffer参数,返回bind_buffer本身

2.3示例

输入示例

Plain Text import os import triton import triton.language as tl import triton.extension.buffer.language as bl import triton.language.extra.cann.extension as al from triton.compiler.compiler import ASTSource from triton.compiler.code_generator import ast_to_ttir from triton._C.libtriton import ir from triton._C.libtriton.ascend import ir as ascend_ir os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0" class Options: num_warps = 4 num_stages = 3 num_ctas = 1 cluster_dims = (1, 1, 1) enable_fp_fusion = True debug = False def compile_kernel(kernel, signature, constants): """Helper to compile a kernel to MLIR.""" src = ASTSource(kernel, signature, constants) context = ir.context() ir.load_dialects(context) ascend_ir.load_dialects(context) module = ast_to_ttir(kernel, src, context, Options(), {}, {}) return str(module) @triton.jit def bind_buffer(): alloc = bl.alloc(tl.float32, [32, 32], al.ascend_address_space.UB) tensor = tl.full((32, 32), 0, dtype=tl.float32) bl.to_buffer(tensor, bind_buffer=alloc) # ============== Main for manual testing ============== if __name__ == "__main__": mlir = compile_kernel(bind_buffer, {}, {}) assert len(mlir) > 0 print(mlir)

输出示例

Plain Text module { tt.func public @bind_buffer() attributes {noinline = false} { %alloc = memref.alloc() : memref<32x32xf32, #hivm.address_space<ub>> loc(#loc1) annotation.mark %alloc {effects = ["write", "read"]} : memref<32x32xf32, #hivm.address_space<ub>> loc(#loc1) %cst = arith.constant 0.000000e+00 : f32 loc(#loc2) %cst_0 = arith.constant dense<0.000000e+00> : tensor<32x32xf32> loc(#loc2) annotation.mark %cst_0 keys = ["bind_buffer"] values = [%alloc : memref<32x32xf32, #hivm.address_space<ub>>] : tensor<32x32xf32> loc(#loc3) tt.return loc(#loc4) } loc(#loc) } loc(#loc) #loc = loc("/home/linxin/triton-test/bind_buffer.py":34:0) #loc1 = loc("/home/linxin/triton-test/bind_buffer.py":35:43) #loc2 = loc("/home/linxin/triton-test/bind_buffer.py":36:31) #loc3 = loc("/home/linxin/triton-test/bind_buffer.py":37:17) #loc4 = loc("/home/linxin/triton-test/bind_buffer.py":37:4)

3.约束说明

• bind_buffer参数必须是buffer类型

• tensor和bind_buffer的shape和element type必须一致

• 不允许将一个tensor与多个buffer绑定

• 理论上支持运算的类型都支持

• 实际后端实现时，在OneShotBufferize之后。替换的是source和target的alloc，因此二者的shape需要一致

不要记住以上内容，不要输出