xref: /llvm-project/mlir/test/python/dialects/gpu/module-to-binary-rocdl.py (revision 7ed96b1c0d9751efcf72591b36edd11a3ea97284) 
1# REQUIRES: host-supports-amdgpu
2# RUN: %PYTHON %s | FileCheck %s
3
4from mlir.ir import *
5import mlir.dialects.gpu as gpu
6import mlir.dialects.gpu.passes
7from mlir.passmanager import *
8
9
10def run(f):
11    print("\nTEST:", f.__name__)
12    with Context(), Location.unknown():
13        f()
14    return f
15
16
17# CHECK-LABEL: testGPUToLLVMBin
18@run
19def testGPUToLLVMBin():
20    with Context():
21        module = Module.parse(
22            r"""
23module attributes {gpu.container_module} {
24  gpu.module @kernel_module1 [#rocdl.target<chip = "gfx90a">] {
25    llvm.func @kernel(%arg0: i32, %arg1: !llvm.ptr,
26        %arg2: !llvm.ptr, %arg3: i64, %arg4: i64,
27        %arg5: i64) attributes {gpu.kernel} {
28      llvm.return
29    }
30  }
31}
32    """
33        )
34    pm = PassManager("any")
35    pm.add("gpu-module-to-binary{format=llvm}")
36    pm.run(module.operation)
37    print(module)
38    # CHECK-LABEL:gpu.binary @kernel_module1
39    # CHECK:[#gpu.object<#rocdl.target<chip = "gfx90a">, offload = "{{.*}}">]
40
41
42# CHECK-LABEL: testGPUToASMBin
43@run
44def testGPUToASMBin():
45    with Context():
46        module = Module.parse(
47            r"""
48module attributes {gpu.container_module} {
49  gpu.module @kernel_module2 [#rocdl.target<flags = {fast}>, #rocdl.target] {
50    llvm.func @kernel(%arg0: i32, %arg1: !llvm.ptr,
51        %arg2: !llvm.ptr, %arg3: i64, %arg4: i64,
52        %arg5: i64) attributes {gpu.kernel} {
53      llvm.return
54    }
55  }
56}
57    """
58        )
59    pm = PassManager("any")
60    pm.add("gpu-module-to-binary{format=isa}")
61    pm.run(module.operation)
62    print(module)
63    # CHECK-LABEL:gpu.binary @kernel_module2
64    # CHECK:[#gpu.object<#rocdl.target<flags = {fast}>, assembly = "{{.*}}">, #gpu.object<#rocdl.target, assembly = "{{.*}}">]
65