64 lines
1.7 KiB
Python
64 lines
1.7 KiB
Python
# REQUIRES: host-supports-nvptx
|
|
# RUN: %PYTHON %s | FileCheck %s
|
|
|
|
from mlir.ir import *
|
|
import mlir.dialects.gpu as gpu
|
|
import mlir.dialects.gpu.passes
|
|
from mlir.passmanager import *
|
|
|
|
|
|
def run(f):
|
|
print("\nTEST:", f.__name__)
|
|
with Context(), Location.unknown():
|
|
f()
|
|
return f
|
|
|
|
|
|
# CHECK-LABEL: testGPUToLLVMBin
|
|
@run
|
|
def testGPUToLLVMBin():
|
|
with Context():
|
|
module = Module.parse(
|
|
r"""
|
|
module attributes {gpu.container_module} {
|
|
gpu.module @kernel_module1 [#nvvm.target<chip = "sm_70">] {
|
|
llvm.func @kernel(%arg0: i32, %arg1: !llvm.ptr,
|
|
%arg2: !llvm.ptr, %arg3: i64, %arg4: i64,
|
|
%arg5: i64) attributes {gpu.kernel} {
|
|
llvm.return
|
|
}
|
|
}
|
|
}
|
|
"""
|
|
)
|
|
pm = PassManager("any")
|
|
pm.add("gpu-module-to-binary{format=llvm}")
|
|
pm.run(module.operation)
|
|
print(module)
|
|
# CHECK-LABEL:gpu.binary @kernel_module1
|
|
# CHECK:[#gpu.object<#nvvm.target<chip = "sm_70">, offload = "{{.*}}">]
|
|
|
|
|
|
# CHECK-LABEL: testGPUToASMBin
|
|
@run
|
|
def testGPUToASMBin():
|
|
with Context():
|
|
module = Module.parse(
|
|
r"""
|
|
module attributes {gpu.container_module} {
|
|
gpu.module @kernel_module2 [#nvvm.target<flags = {fast}>, #nvvm.target] {
|
|
llvm.func @kernel(%arg0: i32, %arg1: !llvm.ptr,
|
|
%arg2: !llvm.ptr, %arg3: i64, %arg4: i64,
|
|
%arg5: i64) attributes {gpu.kernel} {
|
|
llvm.return
|
|
}
|
|
}
|
|
}
|
|
"""
|
|
)
|
|
pm = PassManager("any")
|
|
pm.add("gpu-module-to-binary{format=isa}")
|
|
pm.run(module.operation)
|
|
print(module)
|
|
# CHECK-LABEL:gpu.binary @kernel_module2
|
|
# CHECK:[#gpu.object<#nvvm.target<flags = {fast}>, properties = {O = 2 : i32}, assembly = "{{.*}}">, #gpu.object<#nvvm.target, properties = {O = 2 : i32}, assembly = "{{.*}}">]
|