[Mlir-commits] [mlir] [mlir] Support ROCDL::ReadlaneOp (PR #116593)

[Kyle Wang]

https://github.com/knwng updated https://github.com/llvm/llvm-project/pull/116593

>From 4910e2a31d5e480842b834d879a852d59a6a0fb3 Mon Sep 17 00:00:00 2001
From: Kyle Wang <ec1wng at gmail.com>
Date: Mon, 18 Nov 2024 01:18:00 -0800
Subject: [PATCH 1/2] Support ROCDL::ReadlaneOp

---
 mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td | 20 ++++++++++++++++++++
 mlir/test/Target/LLVMIR/rocdl.mlir           | 19 +++++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index 3695708439d91f..21ddabb614c818 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
@@ -197,6 +197,26 @@ def ROCDL_BallotOp :
   let assemblyFormat = "$pred attr-dict `:` type($res)";
 }
 
+def ROCDL_ReadlaneOp : ROCDL_Op<"readlane">,
+  Results<(outs LLVM_Type:$res)>,
+  Arguments<(ins LLVM_Type:$src0,
+                 I32:$src1)> {
+  let summary = "Get the value in the specific lane";
+
+  let description = [{
+    Get the value in lane `src1` from input `src0`.
+  }];
+
+  string llvmBuilder = [{
+    $res = createIntrinsicCall(builder,
+           llvm::Intrinsic::amdgcn_readlane, {$src0, $src1}, {$_resultType});
+  }];
+
+  let assemblyFormat = [{
+    $src0 `,` $src1  attr-dict `:` `(` type($src0) `,` type($src1) `)` `->` type($res)
+   }];
+}
+
 //===----------------------------------------------------------------------===//
 // Thread index and Block index
 
diff --git a/mlir/test/Target/LLVMIR/rocdl.mlir b/mlir/test/Target/LLVMIR/rocdl.mlir
index 2f34070147be47..0620c23b5fdad7 100644
--- a/mlir/test/Target/LLVMIR/rocdl.mlir
+++ b/mlir/test/Target/LLVMIR/rocdl.mlir
@@ -118,6 +118,25 @@ llvm.func @rocdl.ballot64(%pred : i1) -> i64 {
   llvm.return %0 : i64
 }
 
+llvm.func @rocdl.readlane(%src0 : f32, %src1: f64, %src2: i32, %src3: vector<2 x f32>) -> f32 {
+  %idx = llvm.mlir.constant(0 : i32) : i32
+
+  // CHECK-LABEL: rocdl.readlane
+  // CHECK: call float @llvm.amdgcn.readlane.f32(float %{{.*}}, i32 0)
+  %0 = rocdl.readlane %src0, %idx : (f32, i32) -> f32
+
+  // CHECK: call double @llvm.amdgcn.readlane.f64(double %{{.*}}, i32 0)
+  %1 = rocdl.readlane %src1, %idx : (f64, i32) -> f64
+
+  // CHECK: call i32 @llvm.amdgcn.readlane.i32(i32 %{{.*}}, i32 0)
+  %2 = rocdl.readlane %src2, %idx : (i32, i32) -> i32
+
+  // CHECK: call <2 x float> @llvm.amdgcn.readlane.v2f32(<2 x float> %{{.*}}, i32 0)
+  %3 = rocdl.readlane %src3, %idx : (vector<2 x f32>, i32) -> vector<2 x f32>
+
+  llvm.return %0 : f32
+}
+
 llvm.func @rocdl.waitcnt() {
   // CHECK-LABEL: rocdl.waitcnt
   // CHECK-NEXT: call void @llvm.amdgcn.s.waitcnt(i32 0)

>From 54f2f03422a936179a141b80d6ee0cb015fbd2e9 Mon Sep 17 00:00:00 2001
From: Kyle Wang <ec1wng at gmail.com>
Date: Tue, 19 Nov 2024 00:51:49 -0800
Subject: [PATCH 2/2] Use ROCDL_IntrOp to define ROCDL_ReadlaneOp; add smoke
 testing

---
 mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td | 11 +++--------
 mlir/test/Dialect/LLVMIR/rocdl.mlir          | 11 +++++++++++
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index 21ddabb614c818..71dac3ad39b7b1 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
@@ -197,21 +197,16 @@ def ROCDL_BallotOp :
   let assemblyFormat = "$pred attr-dict `:` type($res)";
 }
 
-def ROCDL_ReadlaneOp : ROCDL_Op<"readlane">,
-  Results<(outs LLVM_Type:$res)>,
+def ROCDL_ReadlaneOp : ROCDL_IntrOp<"readlane", [], [0], [AllTypesMatch<["res", "src0"]>], 1>,
   Arguments<(ins LLVM_Type:$src0,
                  I32:$src1)> {
-  let summary = "Get the value in the specific lane";
+  let results = (outs LLVM_Type:$res);
+  let summary = "Get the value in the specific lane.";
 
   let description = [{
     Get the value in lane `src1` from input `src0`.
   }];
 
-  string llvmBuilder = [{
-    $res = createIntrinsicCall(builder,
-           llvm::Intrinsic::amdgcn_readlane, {$src0, $src1}, {$_resultType});
-  }];
-
   let assemblyFormat = [{
     $src0 `,` $src1  attr-dict `:` `(` type($src0) `,` type($src1) `)` `->` type($res)
    }];
diff --git a/mlir/test/Dialect/LLVMIR/rocdl.mlir b/mlir/test/Dialect/LLVMIR/rocdl.mlir
index 4afa839aa3ea1b..92789246edb4f3 100644
--- a/mlir/test/Dialect/LLVMIR/rocdl.mlir
+++ b/mlir/test/Dialect/LLVMIR/rocdl.mlir
@@ -388,6 +388,17 @@ llvm.func @rocdl.s.wait.dscnt() {
 
 // -----
 
+llvm.func @rocdl.readlane(%src : f32) -> f32 {
+  %cst0 = llvm.mlir.constant(0 : i32) : i32
+
+  // CHECK-LABEL: rocdl.readlane
+  // CHECK: rocdl.readlane %{{.*}} %{{.*}}
+  %ret = rocdl.readlane %src, %cst0 : (f32, i32) -> f32
+  llvm.return %ret : f32
+}
+
+// -----
+
 // expected-error at below {{attribute attached to unexpected op}}
 func.func private @expected_llvm_func() attributes { rocdl.kernel }