From 59ec7ab997379d70cca9aa67588ae4274f1c9a3b Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Thu, 7 Aug 2025 19:13:07 +0200 Subject: [PATCH 1/3] [mlir][rocdl] Add `readfirstlane` intrinsic --- mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td | 18 ++++++++++++++++-- mlir/test/Dialect/LLVMIR/rocdl.mlir | 7 +++++++ mlir/test/Target/LLVMIR/rocdl.mlir | 17 +++++++++++++++++ 3 files changed, 40 insertions(+), 2 deletions(-) diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td index a2354e22e2745..992dc2209372c 100644 --- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td +++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td @@ -189,7 +189,21 @@ def ROCDL_BallotOp : let assemblyFormat = "$pred attr-dict `:` type($res)"; } -def ROCDL_ReadlaneOp : ROCDL_IntrOp<"readlane", [], [0], [AllTypesMatch<["res", "src0"]>], 1>, +def ROCDL_ReadfirstlaneOp : ROCDL_IntrOp<"readfirstlane", [], [0], [AllTypesMatch<["res", "src"]>, Pure], 1>, + Arguments<(ins LLVM_Type:$src)> { + let results = (outs LLVM_Type:$res); + let summary = "Get the value in first active lane."; + + let description = [{ + Returns the value in the lowest active lane of the input operand. + }]; + + let assemblyFormat = [{ + $src attr-dict `:` type($res) + }]; +} + +def ROCDL_ReadlaneOp : ROCDL_IntrOp<"readlane", [], [0], [AllTypesMatch<["res", "src0"]>, Pure], 1>, Arguments<(ins LLVM_Type:$src0, I32:$src1)> { let results = (outs LLVM_Type:$res); @@ -201,7 +215,7 @@ def ROCDL_ReadlaneOp : ROCDL_IntrOp<"readlane", [], [0], [AllTypesMatch<["res", 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 a2b2f84606ba0..db5271c57f573 100644 --- a/mlir/test/Dialect/LLVMIR/rocdl.mlir +++ b/mlir/test/Dialect/LLVMIR/rocdl.mlir @@ -981,6 +981,13 @@ llvm.func @rocdl.s.wait.expcnt() { // ----- +llvm.func @rocdl.readfirstlane(%src : f32) -> f32 { + // CHECK-LABEL: rocdl.readfirstlane + // CHECK: rocdl.readfirstlane %{{.*}} : f32 + %ret = rocdl.readfirstlane %src : f32 + llvm.return %ret : f32 +} + llvm.func @rocdl.readlane(%src : f32) -> f32 { %cst0 = llvm.mlir.constant(0 : i32) : i32 diff --git a/mlir/test/Target/LLVMIR/rocdl.mlir b/mlir/test/Target/LLVMIR/rocdl.mlir index 740990a6e589b..ce4394102b5a1 100644 --- a/mlir/test/Target/LLVMIR/rocdl.mlir +++ b/mlir/test/Target/LLVMIR/rocdl.mlir @@ -125,6 +125,23 @@ llvm.func @rocdl.ballot64(%pred : i1) -> i64 { llvm.return %0 : i64 } +llvm.func @rocdl.readfirstlane(%src0 : f32, %src1: f64, %src2: i32, %src3: vector<2 x f32>) -> f32 { + // CHECK-LABEL: rocdl.readfirstlane + // CHECK: call float @llvm.amdgcn.readfirstlane.f32(float %{{.*}}) + %0 = rocdl.readfirstlane %src0 : f32 + + // CHECK: call double @llvm.amdgcn.readfirstlane.f64(double %{{.*}}) + %1 = rocdl.readfirstlane %src1 : f64 + + // CHECK: call i32 @llvm.amdgcn.readfirstlane.i32(i32 %{{.*}}) + %2 = rocdl.readfirstlane %src2 : i32 + + // CHECK: call <2 x float> @llvm.amdgcn.readfirstlane.v2f32(<2 x float> %{{.*}}) + %3 = rocdl.readfirstlane %src3 : vector<2 x f32> + + llvm.return %0 : f32 +} + llvm.func @rocdl.readlane(%src0 : f32, %src1: f64, %src2: i32, %src3: vector<2 x f32>) -> f32 { %idx = llvm.mlir.constant(0 : i32) : i32 From 53fc37b2e845bd7d0d870d82b8e9fd2fc08e7843 Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Thu, 7 Aug 2025 19:34:38 +0200 Subject: [PATCH 2/3] drop pure --- mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td index 992dc2209372c..b2b6fd3e68559 100644 --- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td +++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td @@ -189,7 +189,7 @@ def ROCDL_BallotOp : let assemblyFormat = "$pred attr-dict `:` type($res)"; } -def ROCDL_ReadfirstlaneOp : ROCDL_IntrOp<"readfirstlane", [], [0], [AllTypesMatch<["res", "src"]>, Pure], 1>, +def ROCDL_ReadfirstlaneOp : ROCDL_IntrOp<"readfirstlane", [], [0], [AllTypesMatch<["res", "src"]>], 1>, Arguments<(ins LLVM_Type:$src)> { let results = (outs LLVM_Type:$res); let summary = "Get the value in first active lane."; From c62230bd0e017241d371ec6f3e0b48467ceb3129 Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Thu, 7 Aug 2025 20:01:54 +0200 Subject: [PATCH 3/3] drop pure --- mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td index b2b6fd3e68559..90da243ee2d07 100644 --- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td +++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td @@ -203,7 +203,7 @@ def ROCDL_ReadfirstlaneOp : ROCDL_IntrOp<"readfirstlane", [], [0], [AllTypesMatc }]; } -def ROCDL_ReadlaneOp : ROCDL_IntrOp<"readlane", [], [0], [AllTypesMatch<["res", "src0"]>, Pure], 1>, +def ROCDL_ReadlaneOp : ROCDL_IntrOp<"readlane", [], [0], [AllTypesMatch<["res", "src0"]>], 1>, Arguments<(ins LLVM_Type:$src0, I32:$src1)> { let results = (outs LLVM_Type:$res);