[mlir] Support ROCDL::ReadlaneOp #116593
[mlir] Support ROCDL::ReadlaneOp #116593
Conversation
|
Thank you for submitting a Pull Request (PR) to the LLVM Project! This PR will be automatically labeled and the relevant teams will be notified. If you wish to, you can add reviewers by using the "Reviewers" section on this page. If this is not working for you, it is probably because you do not have write permissions for the repository. In which case you can instead tag reviewers by name in a comment by using If you have received no comments on your PR for a week, you can request a review by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate is once a week. Please remember that you are asking for valuable time from other developers. If you have further questions, they may be answered by the LLVM GitHub User Guide. You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums. |
|
@llvm/pr-subscribers-mlir-llvm Author: Kyle Wang (knwng) ChangesSupport ROCDL::ReadlaneOp to solve https://github.com/ROCm/triton-internal/issues/411. Full diff: https://github.com/llvm/llvm-project/pull/116593.diff 2 Files Affected:
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)
|
|
@llvm/pr-subscribers-mlir Author: Kyle Wang (knwng) ChangesSupport ROCDL::ReadlaneOp to solve https://github.com/ROCm/triton-internal/issues/411. Full diff: https://github.com/llvm/llvm-project/pull/116593.diff 2 Files Affected:
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)
|
|
Hi @krzysz00 ! As discussed before, this PR is to add readlane intrinsic to ROCDL. I'm wondering if it's worthy supporting |
| Get the value in lane `src1` from input `src0`. | ||
| }]; | ||
|
|
||
| string llvmBuilder = [{ |
There was a problem hiding this comment.
Can we use the ROCDL_IntrOp infrastructure? It'll allow round-tripping if someone gets around to that
There was a problem hiding this comment.
Just want to double-check, https://github.com/llvm/llvm-project/blob/main/llvm/lib/Target/AMDGPU/VOP2Instructions.td#L794-L796 shows the output and 1st input have the same type, while https://github.com/llvm/llvm-project/blob/main/llvm/include/llvm/IR/IntrinsicsAMDGPU.td#L2161-L2163 doesn't have that constraint. I'm wondering if we should reflect this in rocdl.readlane by adding trait AllTypesMatch<["res", "src0"]>.
| }]; | ||
|
|
||
| let assemblyFormat = [{ | ||
| $src0 `,` $src1 attr-dict `:` `(` type($src0) `,` type($src1) `)` `->` type($res) |
There was a problem hiding this comment.
You could optionally just use type($res) or type($src0), type($src1) here, but, ... it's a low-level intrinsic wrapper, it doesn't really matter.
Same for not having the parentheses.
There was a problem hiding this comment.
Actually, yeah, no, it's fine
|
@knwng Congratulations on having your first Pull Request (PR) merged into the LLVM Project! Your changes will be combined with recent changes from other authors, then tested by our build bots. If there is a problem with a build, you may receive a report in an email or a comment on this PR. Please check whether problems have been caused by your change specifically, as the builds can include changes from many authors. It is not uncommon for your change to be included in a build that fails due to someone else's changes, or infrastructure issues. How to do this, and the rest of the post-merge process, is covered in detail here. If your change does cause a problem, it may be reverted, or you can revert it yourself. This is a normal part of LLVM development. You can fix your changes and open a new PR to merge them again. If you don't get any reports, no action is required from you. Your changes are working as expected, well done! |
Support ROCDL::ReadlaneOp to solve https://github.com/ROCm/triton-internal/issues/411.