[Mlir-commits] [mlir] [MLIR][ROCDL] Add permlane16.swap and permanlane32.swap (PR #153804)

[Tim Gymnich]

https://github.com/tgymnich updated https://github.com/llvm/llvm-project/pull/153804

>From 2802ec440be392a57f9451dafa8e7eeab0d8a1d4 Mon Sep 17 00:00:00 2001
From: Tim Gymnich <tim at gymni.ch>
Date: Fri, 15 Aug 2025 13:19:30 +0000
Subject: [PATCH] add rocdl.permlane16.swap and rocdl.permanlane32.swap

---
 mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td | 46 ++++++++++++++++++++
 mlir/test/Dialect/LLVMIR/rocdl.mlir          | 16 +++++++
 2 files changed, 62 insertions(+)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index d6761f4da21ff..89f21b82d793b 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
@@ -780,6 +780,52 @@ def ROCDL_PermlaneX16Op : ROCDL_IntrOp<"permlanex16", [], [0],
   }];
 }
 
+class ROCDL_ConcretePair<Type elem0, Type elem1> :
+  Type<And<[
+    CPred<"::llvm::isa<::mlir::LLVM::LLVMStructType>($_self)">,
+    SubstLeaves<
+         "$_self",
+         "::llvm::cast<::mlir::LLVM::LLVMStructType>($_self).getTypeAtIndex(0)",
+         elem0.predicate>,
+    SubstLeaves<
+         "$_self",
+         "::llvm::cast<::mlir::LLVM::LLVMStructType>($_self).getTypeAtIndex(0)",
+         elem1.predicate>,
+    // CPred<"$_self.getTypeAtIndex(0) == " # elem0 # "">,
+    // CPred<"$_self.getTypeAtIndex(1) == " # elem1 # "">
+  ]> , "LLVM pair type">,
+  BuildableType<"::mlir::LLVM::LLVMStructType::getLiteral($_builder.getContext(), {" # elem0.builderCall # ", " # elem1.builderCall # "})">;
+
+// Permlane16 swap intrinsic operation
+def ROCDL_Permlane16SwapOp : ROCDL_IntrOp<"permlane16.swap", [], [],
+    [], 1, 0, 0, 0,
+    [2, 3], ["fi", "boundControl"]>,
+  Arguments<(ins I32:$old, I32:$src, I1Attr:$fi, I1Attr:$boundControl)> {
+  let results = (outs ROCDL_ConcretePair<I32, I32>:$res);
+  let assemblyFormat = [{
+    attr-dict $old `,` $src `,` $fi `,` $boundControl `:` `(` type($old) `,` type($src) `)` `->` type($res)
+  }];
+  let description = [{
+    Performs a `permlane16.swap` operation with the given operands, applying the
+    permutation specified by $fi to the provided inputs.
+  }];
+}
+
+// Permlane32 swap intrinsic operation
+def ROCDL_Permlane32SwapOp : ROCDL_IntrOp<"permlane32.swap", [], [],
+    [], 1, 0, 0, 0,
+    [2, 3], ["fi", "boundControl"]>,
+  Arguments<(ins I32:$old, I32:$src, I1Attr:$fi, I1Attr:$boundControl)> {
+  let results = (outs ROCDL_ConcretePair<I32, I32>:$res);
+  let assemblyFormat = [{
+    attr-dict $old `,` $src `,` $fi `,` $boundControl `:` `(` type($old) `,` type($src) `)` `->` type($res)
+  }];
+  let description = [{
+    Performs a `permlane32.swap` operation with the given operands, applying the
+    permutation specified by $fi to the provided inputs.
+  }];
+}
+
 class ROCDL_ConcreteVector<Type elem, int length> :
   FixedVectorOfLengthAndType<[length], [elem]>,
   BuildableType<
diff --git a/mlir/test/Dialect/LLVMIR/rocdl.mlir b/mlir/test/Dialect/LLVMIR/rocdl.mlir
index db5271c57f573..782ef4e154440 100644
--- a/mlir/test/Dialect/LLVMIR/rocdl.mlir
+++ b/mlir/test/Dialect/LLVMIR/rocdl.mlir
@@ -1009,6 +1009,22 @@ llvm.func @rocdl.permlanex16(%src : f32) -> f32 {
 
 // -----
 
+llvm.func @rocdl.permlane16.swap(%src : i32) -> !llvm.struct<(i32, i32)> {
+  // CHECK-LABEL: rocdl.permlane16.swap
+  // CHECK: rocdl.permlane16.swap %{{.*}} %{{.*}}
+  %res = rocdl.permlane16.swap %src, %src, 0, -1  : (i32, i32) -> !llvm.struct<(i32, i32)>
+  llvm.return %res : !llvm.struct<(i32, i32)>
+}
+
+llvm.func @rocdl.permlane32.swap(%src : i32) -> !llvm.struct<(i32, i32)> {
+  // CHECK-LABEL: rocdl.permlane32.swap
+  // CHECK: rocdl.permlane32.swap %{{.*}} %{{.*}}
+  %res = rocdl.permlane32.swap %src, %src, 0, -1  : (i32, i32) -> !llvm.struct<(i32, i32)>
+  llvm.return %res : !llvm.struct<(i32, i32)>
+}
+
+// -----
+
 // expected-error at below {{attribute attached to unexpected op}}
 func.func private @expected_llvm_func() attributes { rocdl.kernel }