[Mlir-commits] [mlir] [MLIR][NVVM] Add support for mapa MLIR Ops (PR #124514)

Durgadoss R llvmlistbot at llvm.org
Tue Jan 28 00:18:15 PST 2025


================
@@ -2512,6 +2512,31 @@ def NVVM_WgmmaMmaAsyncOp : NVVM_Op<"wgmma.mma_async",
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// NVVM Mapa Op
+//===----------------------------------------------------------------------===//
+
+def NVVM_MapaOp: NVVM_Op<"mapa",
+    [TypesMatchWith<"`res` and `a` should have the same type",
+                    "a", "res", "$_self">]> {
+  let results = (outs AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerShared]>:$res);
+  let arguments = (ins AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerShared]>:$a, I32:$b);
+
+  string llvmBuilder = [{
+    int addrSpace = llvm::cast<LLVMPointerType>(op.getA().getType()).getAddressSpace();
+    
+    bool isSharedMemory = addrSpace == NVVM::NVVMMemorySpace::kSharedMemorySpace;
+
+    if(isSharedMemory)
+      $res = createIntrinsicCall(builder,
+        llvm::Intrinsic::nvvm_mapa_shared_cluster, {$a, $b});
+    else
+      $res = createIntrinsicCall(builder, llvm::Intrinsic::nvvm_mapa, {$a, $b});
----------------
durga4github wrote:

nit:
This can be simplified a bit, since intrinsic id is the only difference between the two calls.

auto id = isSharedMemory ? llvm::Intrinsic::nvvm_mapa_shared_cluster : nvvm_mapa;
$res = createIntrinsicCall(builder, id, {$a, $b});


https://github.com/llvm/llvm-project/pull/124514


More information about the Mlir-commits mailing list