bolt/deps/llvm-18.1.8/mlir/test/python/dialects/gpu/module-to-binary-nvvm.py
2025-02-14 19:21:04 +01:00

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 = "{{.*}}">]