diff options
| author | Kyle Wang <ec1wng@gmail.com> | 2024-11-19 15:36:30 -0800 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2024-11-19 17:36:30 -0600 |
| commit | d8bd7f11c8c781646406e76731dd8d76ed5425dd (patch) | |
| tree | 5428ccb6dc9861a575d748c1bd616cf4b6d7238b | |
| parent | 3a63407686313f46f9abc664fd10b01f4359ee27 (diff) | |
| download | llvm-d8bd7f11c8c781646406e76731dd8d76ed5425dd.zip llvm-d8bd7f11c8c781646406e76731dd8d76ed5425dd.tar.gz llvm-d8bd7f11c8c781646406e76731dd8d76ed5425dd.tar.bz2 | |
[mlir] Support ROCDL::ReadlaneOp (#116593)
Support ROCDL::ReadlaneOp to solve
https://github.com/ROCm/triton-internal/issues/411.
| -rw-r--r-- | mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td | 15 | ||||
| -rw-r--r-- | mlir/test/Dialect/LLVMIR/rocdl.mlir | 11 | ||||
| -rw-r--r-- | mlir/test/Target/LLVMIR/rocdl.mlir | 19 |
3 files changed, 45 insertions, 0 deletions
diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td index 3695708..71dac3a 100644 --- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td +++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td @@ -197,6 +197,21 @@ def ROCDL_BallotOp : let assemblyFormat = "$pred attr-dict `:` type($res)"; } +def ROCDL_ReadlaneOp : ROCDL_IntrOp<"readlane", [], [0], [AllTypesMatch<["res", "src0"]>], 1>, + Arguments<(ins LLVM_Type:$src0, + I32:$src1)> { + 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`. + }]; + + let assemblyFormat = [{ + $src0 `,` $src1 attr-dict `:` `(` type($src0) `,` type($src1) `)` `->` type($res) + }]; +} + //===----------------------------------------------------------------------===// // Thread index and Block index diff --git a/mlir/test/Dialect/LLVMIR/rocdl.mlir b/mlir/test/Dialect/LLVMIR/rocdl.mlir index 4afa839..9278924 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@below {{attribute attached to unexpected op}} func.func private @expected_llvm_func() attributes { rocdl.kernel } diff --git a/mlir/test/Target/LLVMIR/rocdl.mlir b/mlir/test/Target/LLVMIR/rocdl.mlir index 2f34070..0620c23b 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) |