4f88c23111
This PR adds NVGPU dialects' TensorMapDescriptorType in the py bindings. This is a follow-up issue from [this PR](https://github.com/llvm/llvm-project/pull/87153#discussion_r1546193095)
44 lines
1.5 KiB
Python
44 lines
1.5 KiB
Python
# RUN: %PYTHON %s | FileCheck %s
|
|
# This is just a smoke test that the dialect is functional.
|
|
|
|
from mlir.ir import *
|
|
from mlir.dialects import nvgpu, arith, memref
|
|
|
|
|
|
def constructAndPrintInModule(f):
|
|
print("\nTEST:", f.__name__)
|
|
with Context(), Location.unknown():
|
|
module = Module.create()
|
|
with InsertionPoint(module.body):
|
|
f()
|
|
print(module)
|
|
return f
|
|
|
|
|
|
# CHECK-LABEL: testTypes
|
|
@constructAndPrintInModule
|
|
def testTypes():
|
|
tensorMemrefType = MemRefType.get(
|
|
(128, 64), F16Type.get(), memory_space=Attribute.parse("3")
|
|
)
|
|
# CHECK: !nvgpu.tensormap.descriptor<tensor = memref<128x64xf16, 3>, swizzle = swizzle_128b, l2promo = l2promo_256b, oob = nan, interleave = none>
|
|
tma_desc = nvgpu.TensorMapDescriptorType.get(
|
|
tensorMemrefType,
|
|
nvgpu.TensorMapSwizzleKind.SWIZZLE_128B,
|
|
nvgpu.TensorMapL2PromoKind.L2PROMO_256B,
|
|
nvgpu.TensorMapOOBKind.OOB_NAN,
|
|
nvgpu.TensorMapInterleaveKind.INTERLEAVE_NONE,
|
|
)
|
|
print(tma_desc)
|
|
|
|
|
|
# CHECK-LABEL: testSmoke
|
|
@constructAndPrintInModule
|
|
def testSmoke():
|
|
cst = arith.ConstantOp(value=42, result=IndexType.get())
|
|
mem_t = MemRefType.get((10, 10), F32Type.get(), memory_space=Attribute.parse("3"))
|
|
vec_t = VectorType.get((4, 1), F32Type.get())
|
|
mem = memref.AllocOp(mem_t, [], [])
|
|
# CHECK: %0 = nvgpu.ldmatrix %alloc[%c42, %c42] {numTiles = 4 : i32, transpose = false} : memref<10x10xf32, 3> -> vector<4x1xf32>
|
|
nvgpu.LdMatrixOp(vec_t, mem, [cst, cst], False, 4)
|