al.ascend_address_space 接口文档

1. 背景

为了支持Ascend级编程的需要，需要支持用户手动创建指定地址空间上的内存(buffer)，对接hivm::AddressSpace枚举。

2. 接口说明

Python al.ascend_address_space.UB al.ascend_address_space.L1 al.ascend_address_space.L0A al.ascend_address_space.L0B al.ascend_address_space.L0C

2.1. 返回值

不涉及

2.2. 入参

不涉及

3. 约束说明

• 需要配合bl.alloc使用

4. 用例示例

Python import os import triton import triton.language as tl from triton.compiler.compiler import ASTSource from triton.compiler.code_generator import ast_to_ttir import triton.extension.buffer.language as bl import triton.language.extra.cann.extension as al from triton._C.libtriton import ir, buffer_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) buffer_ir.load_dialects(context) ascend_ir.load_dialects(context) module = ast_to_ttir(kernel, src, context, Options(), {"create_address_space": al.semantic.create_address_space}, {}) return str(module) # ============== Kernel definitions ============== @triton.jit def allocate_local_buffer(XBLOCK: tl.constexpr): # this statement has no effect, just to test the builder bl.alloc(tl.float32, [XBLOCK]) bl.alloc(tl.float32, [XBLOCK, XBLOCK], al.ascend_address_space.UB) bl.alloc(tl.float32, [XBLOCK, XBLOCK], al.ascend_address_space.L1) bl.alloc(tl.float32, [XBLOCK, XBLOCK], al.ascend_address_space.L0A) bl.alloc(tl.float32, [XBLOCK, XBLOCK], al.ascend_address_space.L0B) bl.alloc(tl.float32, [XBLOCK, XBLOCK], al.ascend_address_space.L0C) bl.alloc( tl.float32, [XBLOCK, XBLOCK], al.ascend_address_space.UB, is_mem_unique=True ) # ============== Main for manual testing ============== if __name__ == "__main__": print("=" * 60) print("Test 1: Nested Scopes") print("=" * 60) mlir = compile_kernel( allocate_local_buffer, {}, {"XBLOCK": 256} ) print(f"✅ Generated MLIR ({len(mlir)} chars):\n") print(mlir)

5. 编译输出结果

Plain Text ============================================================ Test 1: Nested Scopes ============================================================ ✅ Generated MLIR (2103 chars): module { tt.func public @allocate_local_buffer() attributes {noinline = false} { %alloc = memref.alloc() : memref<256xf32> loc(#loc1) annotation.mark %alloc {effects = ["write", "read"]} : memref<256xf32> loc(#loc1) %alloc_0 = memref.alloc() : memref<256x256xf32, #hivm.address_space<ub>> loc(#loc2) annotation.mark %alloc_0 {effects = ["write", "read"]} : memref<256x256xf32, #hivm.address_space<ub>> loc(#loc2) %alloc_1 = memref.alloc() : memref<256x256xf32, #hivm.address_space<cbuf>> loc(#loc3) annotation.mark %alloc_1 {effects = ["write", "read"]} : memref<256x256xf32, #hivm.address_space<cbuf>> loc(#loc3) %alloc_2 = memref.alloc() : memref<256x256xf32, #hivm.address_space<ca>> loc(#loc4) annotation.mark %alloc_2 {effects = ["write", "read"]} : memref<256x256xf32, #hivm.address_space<ca>> loc(#loc4) %alloc_3 = memref.alloc() : memref<256x256xf32, #hivm.address_space<cb>> loc(#loc5) annotation.mark %alloc_3 {effects = ["write", "read"]} : memref<256x256xf32, #hivm.address_space<cb>> loc(#loc5) %alloc_4 = memref.alloc() : memref<256x256xf32, #hivm.address_space<cc>> loc(#loc6) annotation.mark %alloc_4 {effects = ["write", "read"]} : memref<256x256xf32, #hivm.address_space<cc>> loc(#loc6) %alloc_5 = memref.alloc() : memref<256x256xf32, #hivm.address_space<ub>> loc(#loc7) annotation.mark %alloc_5 {mem_unique} : memref<256x256xf32, #hivm.address_space<ub>> loc(#loc7) annotation.mark %alloc_5 {effects = ["write", "read"]} : memref<256x256xf32, #hivm.address_space<ub>> loc(#loc7) tt.return loc(#loc8) } loc(#loc) } loc(#loc) #loc = loc("/home/linxin/triton-test/alloc.py":41:0) #loc1 = loc("/home/linxin/triton-test/alloc.py":43:25) #loc2 = loc("/home/linxin/triton-test/alloc.py":44:43) #loc3 = loc("/home/linxin/triton-test/alloc.py":45:43) #loc4 = loc("/home/linxin/triton-test/alloc.py":46:43) #loc5 = loc("/home/linxin/triton-test/alloc.py":47:43) #loc6 = loc("/home/linxin/triton-test/alloc.py":48:43) #loc7 = loc("/home/linxin/triton-test/alloc.py":50:38) #loc8 = loc("/home/linxin/triton-test/alloc.py":49:4)