diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td index 906aaca21187b..04a0b58a85211 100644 --- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td +++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td @@ -284,10 +284,56 @@ def ROCDL_BarrierWaitOp : ROCDL_ConcreteNonMemIntrOp<"s.barrier.wait", [], 0, [0 let assemblyFormat = "$id attr-dict"; } -def ROCDL_WaitDscntOp: ROCDL_ConcreteNonMemIntrOp<"s.wait.dscnt", [], 0, [0], ["id"]>, - Arguments<(ins I16Attr:$id)> { +def ROCDL_WaitDscntOp: ROCDL_ConcreteNonMemIntrOp<"s.wait.dscnt", [], 0, [0], ["count"]>, + Arguments<(ins I16Attr:$count)> { + let summary = "Wait until DSCNT is less than or equal to `count`"; + let description = [{ + Wait for the counter specified to be less-than or equal-to the `count` + before continuing. + + Available on gfx12+. + }]; let results = (outs); - let assemblyFormat = "$id attr-dict"; + let assemblyFormat = "$count attr-dict"; +} + +def ROCDL_WaitLoadcntOp: ROCDL_ConcreteNonMemIntrOp<"s.wait.loadcnt", [], 0, [0], ["count"]>, + Arguments<(ins I16Attr:$count)> { + let summary = "Wait until LOADCNT is less than or equal to `count`"; + let description = [{ + Wait for the counter specified to be less-than or equal-to the `count` + before continuing. + + Available on gfx12+. + }]; + let results = (outs); + let assemblyFormat = "$count attr-dict"; +} + +def ROCDL_WaitStorecntOp: ROCDL_ConcreteNonMemIntrOp<"s.wait.storecnt", [], 0, [0], ["count"]>, + Arguments<(ins I16Attr:$count)> { + let summary = "Wait until STORECNT is less than or equal to `count`"; + let description = [{ + Wait for the counter specified to be less-than or equal-to the `count` + before continuing. + + Available on gfx12+. + }]; + let results = (outs); + let assemblyFormat = "$count attr-dict"; +} + +def ROCDL_WaitExpcntOp: ROCDL_ConcreteNonMemIntrOp<"s.wait.expcnt", [], 0, [0], ["count"]>, + Arguments<(ins I16Attr:$count)> { + let summary = "Wait until EXPCNT is less than or equal to `count`"; + let description = [{ + Wait for the counter specified to be less-than or equal-to the `count` + before continuing. + + Available on gfx12+. + }]; + let results = (outs); + let assemblyFormat = "$count attr-dict"; } def ROCDL_SetPrioOp : ROCDL_ConcreteNonMemIntrOp<"s.setprio", [], 0, [0], ["priority"]>, diff --git a/mlir/test/Dialect/LLVMIR/rocdl.mlir b/mlir/test/Dialect/LLVMIR/rocdl.mlir index a6a29bf858e59..a2b2f84606ba0 100644 --- a/mlir/test/Dialect/LLVMIR/rocdl.mlir +++ b/mlir/test/Dialect/LLVMIR/rocdl.mlir @@ -958,6 +958,27 @@ llvm.func @rocdl.s.wait.dscnt() { llvm.return } +llvm.func @rocdl.s.wait.loadcnt() { + // CHECK-LABEL: rocdl.s.wait.loadcnt + // CHECK: rocdl.s.wait.loadcnt 0 + rocdl.s.wait.loadcnt 0 + llvm.return +} + +llvm.func @rocdl.s.wait.storecnt() { + // CHECK-LABEL: rocdl.s.wait.storecnt + // CHECK: rocdl.s.wait.storecnt 0 + rocdl.s.wait.storecnt 0 + llvm.return +} + +llvm.func @rocdl.s.wait.expcnt() { + // CHECK-LABEL: rocdl.s.wait.expcnt + // CHECK: rocdl.s.wait.expcnt 0 + rocdl.s.wait.expcnt 0 + llvm.return +} + // ----- llvm.func @rocdl.readlane(%src : f32) -> f32 { diff --git a/mlir/test/Target/LLVMIR/rocdl.mlir b/mlir/test/Target/LLVMIR/rocdl.mlir index 0742eb3620a7c..740990a6e589b 100644 --- a/mlir/test/Target/LLVMIR/rocdl.mlir +++ b/mlir/test/Target/LLVMIR/rocdl.mlir @@ -196,6 +196,27 @@ llvm.func @rocdl.s.wait.dscnt() { llvm.return } +llvm.func @rocdl.s.wait.loadcnt() { + // CHECK-LABEL: rocdl.s.wait.loadcnt + // CHECK-NEXT: call void @llvm.amdgcn.s.wait.loadcnt(i16 0) + rocdl.s.wait.loadcnt 0 + llvm.return +} + +llvm.func @rocdl.s.wait.storecnt() { + // CHECK-LABEL: rocdl.s.wait.storecnt + // CHECK-NEXT: call void @llvm.amdgcn.s.wait.storecnt(i16 0) + rocdl.s.wait.storecnt 0 + llvm.return +} + +llvm.func @rocdl.s.wait.expcnt() { + // CHECK-LABEL: rocdl.s.wait.expcnt + // CHECK-NEXT: call void @llvm.amdgcn.s.wait.expcnt(i16 0) + rocdl.s.wait.expcnt 0 + llvm.return +} + llvm.func @rocdl.setprio() { // CHECK: call void @llvm.amdgcn.s.setprio(i16 0) rocdl.s.setprio 0