nvgpu/nvvm

Author: makslevental

mlir_utils/dialects/ext/nvgpu.py

@@ -0,0 +1,31 @@
+from textwrap import dedent
+
+from mlir.dialects._nvgpu_enum_gen import (
+    TensorMapSwizzleKind,
+    TensorMapL2PromoKind,
+    TensorMapOOBKind,
+    TensorMapInterleaveKind,
+)
+from mlir.ir import Attribute, Type
+
+
+def tensormap_descriptor(
+    tensor,
+    swizzle: TensorMapSwizzleKind = TensorMapSwizzleKind.SWIZZLE_NONE,
+    l2promo: TensorMapL2PromoKind = TensorMapL2PromoKind.L2PROMO_NONE,
+    oob: TensorMapOOBKind = TensorMapOOBKind.OOB_NAN,
+    interleave: TensorMapInterleaveKind = TensorMapInterleaveKind.INTERLEAVE_NONE,
+    context=None,
+):
+    return Type.parse(
+        dedent(
+            f"""\
+                !nvgpu.tensormap.descriptor<tensor = {tensor}, 
+                                            swizzle = {swizzle}, 
+                                            l2promo = {l2promo}, 
+                                            oob = {oob}, 
+                                            interleave = {interleave}>
+            """
+        ),
+        context=context,
+    )

tests/test_nvgpu_nvvm.py

@@ -0,0 +1,48 @@
+from textwrap import dedent
+
+import pytest
+
+import mlir_utils.types as T
+from mlir_utils.dialects.ext.arith import constant
+from mlir_utils.dialects.ext.func import func
+from mlir_utils.dialects.ext.nvgpu import tensormap_descriptor
+from mlir_utils.dialects.memref import cast
+from mlir_utils.dialects.nvgpu import tma_create_descriptor
+
+# noinspection PyUnresolvedReferences
+from mlir_utils.testing import mlir_ctx as ctx, filecheck, MLIRContext
+
+# needed since the fix isn't defined here nor conftest.py
+pytest.mark.usefixtures("ctx")
+
+
+def test_basic(ctx: MLIRContext):
+    @func
+    def create_tensor_map(
+        device_ptr_2d: T.memref(64, 128, element_type=T.f32),
+    ):
+        crd0 = constant(64, index=True)
+        crd1 = constant(128, index=True)
+        device_ptr_2d_unranked = cast(T.memref(element_type=T.f32), device_ptr_2d)
+        tensor_map_2d = tensormap_descriptor(T.memref(32, 32, T.f32, memory_space=3))
+        tensor_map_2d = tma_create_descriptor(
+            tensor_map_2d, device_ptr_2d_unranked, [crd0, crd1]
+        )
+
+    create_tensor_map.emit()
+
+    ctx.module.operation.verify()
+    correct = dedent(
+        """\
+    module {
+      func.func @create_tensor_map(%arg0: memref<64x128xf32>) {
+        %c64 = arith.constant 64 : index
+        %c128 = arith.constant 128 : index
+        %cast = memref.cast %arg0 : memref<64x128xf32> to memref<*xf32>
+        %0 = nvgpu.tma.create.descriptor %cast box[%c64, %c128] : memref<*xf32> -> <tensor = memref<32x32xf32, 3>, swizzle = none, l2promo = none, oob = nan, interleave = none>
+        return
+      }
+    }
+    """
+    )
+    filecheck(correct, ctx.module)