256 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3from mlir.ir import *4import mlir.ir as ir5from mlir.dialects import gpu, func, arith, math6from mlir.extras import types as T7import mlir.dialects.gpu.passes8from mlir.passmanager import *9 10 11def run(f):12 print("\nTEST:", f.__name__)13 with Context(), Location.unknown():14 f()15 return f16 17 18# CHECK-LABEL: testGPUPass19# CHECK: SUCCESS20@run21def testGPUPass():22 PassManager.parse("any(gpu-kernel-outlining)")23 print("SUCCESS")24 25 26# CHECK-LABEL: testMMAElementWiseAttr27@run28def testMMAElementWiseAttr():29 module = Module.create()30 with InsertionPoint(module.body):31 gpu.BlockDimOp(gpu.Dimension.y)32 # CHECK: %block_dim_y = gpu.block_dim y33 print(module)34 pass35 36 37# CHECK-LABEL: testObjectAttr38@run39def testObjectAttr():40 target = Attribute.parse("#nvvm.target")41 format = gpu.CompilationTarget.Fatbin42 object = b"BC\xc0\xde5\x14\x00\x00\x05\x00\x00\x00b\x0c0$MY\xbef"43 properties = DictAttr.get({"O": IntegerAttr.get(IntegerType.get_signless(32), 2)})44 o = gpu.ObjectAttr.get(target, format, object, properties)45 # CHECK: #gpu.object<#nvvm.target, properties = {O = 2 : i32}, "BC\C0\DE5\14\00\00\05\00\00\00b\0C0$MY\BEf">46 print(o)47 assert o.object == object48 49 o = gpu.ObjectAttr.get(target, format, object)50 # CHECK: #gpu.object<#nvvm.target, "BC\C0\DE5\14\00\00\05\00\00\00b\0C0$MY\BEf">51 print(o)52 53 object = (54 b"//\n// Generated by LLVM NVPTX Back-End\n//\n\n.version 6.0\n.target sm_50"55 )56 o = gpu.ObjectAttr.get(target, format, object)57 # CHECK: #gpu.object<#nvvm.target, "//\0A// Generated by LLVM NVPTX Back-End\0A//\0A\0A.version 6.0\0A.target sm_50">58 print(o)59 assert o.object == object60 61 object = b"BC\xc0\xde5\x14\x00\x00\x05\x00\x00\x00b\x0c0$MY\xbef"62 kernelTable = Attribute.parse(63 '#gpu.kernel_table<[#gpu.kernel_metadata<"kernel", () -> ()>]>'64 )65 o = gpu.ObjectAttr.get(target, format, object, kernels=kernelTable)66 # CHECK: #gpu.object<#nvvm.target, kernels = <[#gpu.kernel_metadata<"kernel", () -> ()>]>, "BC\C0\DE5\14\00\00\05\00\00\00b\0C0$MY\BEf">67 print(o)68 assert o.kernels == kernelTable69 70 71# CHECK-LABEL: testGPUFuncOp72@run73def testGPUFuncOp():74 assert gpu.GPUFuncOp.__doc__ is not None75 module = Module.create()76 with InsertionPoint(module.body):77 gpu_module_name = StringAttr.get("gpu_module")78 gpumodule = gpu.GPUModuleOp(gpu_module_name)79 block = gpumodule.bodyRegion.blocks.append()80 81 def builder(func: gpu.GPUFuncOp) -> None:82 gpu.GlobalIdOp(gpu.Dimension.x)83 gpu.ReturnOp([])84 85 with InsertionPoint(block):86 name = StringAttr.get("kernel0")87 func_type = ir.FunctionType.get(inputs=[], results=[])88 type_attr = TypeAttr.get(func_type)89 func = gpu.GPUFuncOp(type_attr, name)90 func.attributes["sym_name"] = name91 func.attributes["gpu.kernel"] = UnitAttr.get()92 93 try:94 func.entry_block95 assert False, "Expected RuntimeError"96 except RuntimeError as e:97 assert (98 str(e)99 == "Entry block does not exist for kernel0. Do you need to call the add_entry_block() method on this GPUFuncOp?"100 )101 102 block = func.add_entry_block()103 with InsertionPoint(block):104 builder(func)105 106 try:107 func.add_entry_block()108 assert False, "Expected RuntimeError"109 except RuntimeError as e:110 assert str(e) == "Entry block already exists for kernel0"111 112 func = gpu.GPUFuncOp(113 func_type,114 sym_name="kernel1",115 kernel=True,116 body_builder=builder,117 known_block_size=[1, 2, 3],118 known_grid_size=DenseI32ArrayAttr.get([4, 5, 6]),119 )120 121 assert func.name.value == "kernel1"122 assert func.function_type.value == func_type123 assert func.arg_attrs == None124 assert func.res_attrs == None125 assert func.arguments == []126 assert func.entry_block == func.body.blocks[0]127 assert func.is_kernel128 assert func.known_block_size == DenseI32ArrayAttr.get(129 [1, 2, 3]130 ), func.known_block_size131 assert func.known_grid_size == DenseI32ArrayAttr.get(132 [4, 5, 6]133 ), func.known_grid_size134 135 func = gpu.GPUFuncOp(136 ir.FunctionType.get(inputs=[T.index()], results=[]),137 sym_name="non_kernel_func",138 body_builder=builder,139 arg_attrs=[{"gpu.some_attribute": ir.StringAttr.get("foo")}],140 )141 assert not func.is_kernel142 assert func.known_block_size is None143 assert func.known_grid_size is None144 145 print(module)146 147 # CHECK: gpu.module @gpu_module148 # CHECK: gpu.func @kernel0() kernel {149 # CHECK: %[[VAL_0:.*]] = gpu.global_id x150 # CHECK: gpu.return151 # CHECK: }152 # CHECK: gpu.func @kernel1() kernel attributes153 # CHECK-SAME: known_block_size = array<i32: 1, 2, 3>154 # CHECK-SAME: known_grid_size = array<i32: 4, 5, 6>155 # CHECK: %[[VAL_0:.*]] = gpu.global_id x156 # CHECK: gpu.return157 # CHECK: }158 # CHECK: gpu.func @non_kernel_func(159 # CHECK-SAME: %[[ARG0:.*]]: index {gpu.some_attribute = "foo"}) {160 # CHECK: %[[GLOBAL_ID_0:.*]] = gpu.global_id x161 # CHECK: gpu.return162 # CHECK: }163 164 165# CHECK-LABEL: testGPULaunchFuncOp166@run167def testGPULaunchFuncOp():168 module = Module.create()169 170 module.operation.attributes["gpu.container_module"] = UnitAttr.get()171 with InsertionPoint(module.body):172 gpu_module = gpu.GPUModuleOp("gpu_module")173 block = gpu_module.bodyRegion.blocks.append()174 175 with InsertionPoint(block):176 gpu_func = gpu.GPUFuncOp(177 FunctionType.get([], []),178 "kernel",179 body_builder=lambda func: gpu.return_([]),180 kernel=True,181 )182 183 with InsertionPoint(module.body):184 host = func.FuncOp(type=FunctionType.get([], []), name="host")185 186 with InsertionPoint(host.add_entry_block()):187 c1 = arith.constant(T.index(), 1)188 grid_sizes = (1, 1, 1)189 block_sizes = (1, 1, 1)190 token = gpu.wait()191 token = gpu.launch_func(192 async_dependencies=[token],193 kernel=[gpu_module.sym_name.value, gpu_func.name.value],194 grid_size=grid_sizes,195 block_size=block_sizes,196 kernel_operands=[],197 )198 gpu.wait(async_dependencies=[token])199 func.ReturnOp([])200 201 print(module)202 203 # CHECK-LABEL: gpu.module @gpu_module {204 # CHECK: gpu.func @kernel() kernel {205 # CHECK: gpu.return206 # CHECK: }207 # CHECK: }208 209 # CHECK-LABEL: func.func @host() {210 # CHECK: %[[CONSTANT_0:.*]] = arith.constant 1 : index211 # CHECK: %[[WAIT_0:.*]] = gpu.wait async212 # CHECK: %[[CONSTANT_1:.*]] = arith.constant 1 : index213 # CHECK: %[[CONSTANT_2:.*]] = arith.constant 1 : index214 # CHECK: %[[CONSTANT_3:.*]] = arith.constant 1 : index215 # CHECK: %[[CONSTANT_4:.*]] = arith.constant 1 : index216 # CHECK: %[[CONSTANT_5:.*]] = arith.constant 1 : index217 # CHECK: %[[CONSTANT_6:.*]] = arith.constant 1 : index218 # CHECK: %[[LAUNCH_FUNC_0:.*]] = gpu.launch_func async {{\[}}%[[WAIT_0]]] @gpu_module::@kernel blocks in (%[[CONSTANT_1]], %[[CONSTANT_2]], %[[CONSTANT_3]]) threads in (%[[CONSTANT_4]], %[[CONSTANT_5]], %[[CONSTANT_6]])219 # CHECK: %[[WAIT_1:.*]] = gpu.wait async {{\[}}%[[LAUNCH_FUNC_0]]]220 # CHECK: return221 # CHECK: }222 223 224# CHECK-LABEL: testGPULaunchOp225@run226def testGPULaunchOp():227 module = Module.create()228 229 with InsertionPoint(module.body):230 host = func.FuncOp(type=FunctionType.get([T.f32()], []), name="gpu_printf")231 232 entry_block = host.add_entry_block()233 with InsertionPoint(entry_block):234 c1 = arith.constant(T.index(), 1)235 grid_sizes = (c1, c1, c1)236 block_sizes = (c1, c1, c1)237 238 launch = gpu.launch(grid_sizes, block_sizes)239 240 op = launch(lambda *args: gpu.printf("%f", args[0]))241 242 with InsertionPoint(entry_block):243 func.ReturnOp([])244 245 print(module)246 247 # CHECK-LABEL: func.func @gpu_printf(248 # CHECK-SAME: %[[ARG0:.*]]: f32) {249 # CHECK: %[[CONSTANT_0:.*]] = arith.constant 1 : index250 # CHECK: gpu.launch blocks(%[[VAL_0:.*]], %[[VAL_1:.*]], %[[VAL_2:.*]]) in (%[[VAL_3:.*]] = %[[CONSTANT_0]], %[[VAL_4:.*]] = %[[CONSTANT_0]], %[[VAL_5:.*]] = %[[CONSTANT_0]]) threads(%[[VAL_6:.*]], %[[VAL_7:.*]], %[[VAL_8:.*]]) in (%[[VAL_9:.*]] = %[[CONSTANT_0]], %[[VAL_10:.*]] = %[[CONSTANT_0]], %[[VAL_11:.*]] = %[[CONSTANT_0]]) {251 # CHECK: gpu.printf "%[[VAL_12:.*]]", %[[VAL_0]] : index252 # CHECK: gpu.terminator253 # CHECK: }254 # CHECK: return255 # CHECK: }256