diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td index db1b7e3af62fd..6925cec42b5b4 100644 --- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td +++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td @@ -292,18 +292,74 @@ def ROCDL_BarrierOp : ROCDL_Op<"barrier"> { let assemblyFormat = "attr-dict"; } +def ROCDLBufferLDS : LLVM_PointerInAddressSpace<3>; + +def ROCDL_BarrierInitOp : ROCDL_IntrOp<"s.barrier.init", [], [], [], 0, 0, 0, 0, [1], ["id"]>, + Arguments<(ins Arg:$ptr, I32Attr:$id)> { + let description = [{ + Available on gfx1250+. + }]; + let results = (outs); + let assemblyFormat = "$ptr `,` $id attr-dict"; +} + def ROCDL_BarrierSignalOp : ROCDL_ConcreteNonMemIntrOp<"s.barrier.signal", [], 0, [0], ["id"]>, Arguments<(ins I32Attr:$id)> { let results = (outs); let assemblyFormat = "$id attr-dict"; } +def ROCDL_BarrierSignalVarOp : ROCDL_IntrOp<"s.barrier.signal.var", [], [], [], 0, 0, 0, 0, [1], ["id"]>, + Arguments<(ins Arg:$ptr, I32Attr:$id)> { + let description = [{ + Available on gfx1250+. + }]; + let results = (outs); + let assemblyFormat = "$ptr `,` $id attr-dict"; +} + +def ROCDL_BarrierJoinOp : ROCDL_IntrOp<"s.barrier.join", [], [], [], 0>, + Arguments<(ins Arg:$ptr)> { + let description = [{ + Available on gfx1250+. + }]; + let results = (outs); + let assemblyFormat = "$ptr attr-dict"; +} + +def ROCDL_BarrierLeaveOp : ROCDL_ConcreteNonMemIntrOp<"s.barrier.leave", [], 0, [0], ["id"]>, + Arguments<(ins I16Attr:$id)> { + let description = [{ + Available on gfx1250+. + }]; + let results = (outs); + let assemblyFormat = "$id attr-dict"; +} + def ROCDL_BarrierWaitOp : ROCDL_ConcreteNonMemIntrOp<"s.barrier.wait", [], 0, [0], ["id"]>, Arguments<(ins I16Attr:$id)> { let results = (outs); let assemblyFormat = "$id attr-dict"; } +def ROCDL_BarrierSignalIsfirstOp : ROCDL_ConcreteNonMemIntrOp<"s.barrier.signal.isfirst", [], 1, [0], ["id"]>, + Arguments<(ins I32Attr:$id)> { + let description = [{ + Available on gfx1250+. + }]; + let results = (outs I1:$res); + let assemblyFormat = "$id attr-dict `:` type($res)"; +} + +def ROCDL_GetBarrierStateOp : ROCDL_ConcreteNonMemIntrOp<"s.get.barrier.state", [], 1, [0], ["id"]>, + Arguments<(ins I32Attr:$id)> { + let description = [{ + Available on gfx1250+. + }]; + let results = (outs I32:$res); + let assemblyFormat = "$id attr-dict `:` type($res)"; +} + 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`"; @@ -497,7 +553,6 @@ def ROCDL_wmma_i32_16x16x32_iu4 : ROCDL_Wmma_IntrOp<"wmma.i32.16x16x32.iu4", [1] // LDS transpose intrinsics (available in GFX950) def ROCDLGlobalBuffer : LLVM_PointerInAddressSpace<1>; -def ROCDLBufferLDS : LLVM_PointerInAddressSpace<3>; class ROCDL_LDS_Read_Tr_IntrOp : ROCDL_IntrOp { diff --git a/mlir/test/Dialect/LLVMIR/rocdl.mlir b/mlir/test/Dialect/LLVMIR/rocdl.mlir index a88b59aeb61b2..358bd332883f2 100644 --- a/mlir/test/Dialect/LLVMIR/rocdl.mlir +++ b/mlir/test/Dialect/LLVMIR/rocdl.mlir @@ -951,6 +951,13 @@ llvm.func @rocdl.s.barrier() { llvm.return } +llvm.func @rocdl.s.barrier.init(%ptr : !llvm.ptr<3>) { + // CHECK-LABEL: rocdl.s.barrier.init + // CHECK: rocdl.s.barrier.init %[[PTR:.+]], 1 + rocdl.s.barrier.init %ptr, 1 + llvm.return +} + llvm.func @rocdl.s.barrier.signal() { // CHECK-LABEL: rocdl.s.barrier.signal // CHECK: rocdl.s.barrier.signal -1 @@ -958,6 +965,27 @@ llvm.func @rocdl.s.barrier.signal() { llvm.return } +llvm.func @rocdl.s.barrier.signal.var(%ptr : !llvm.ptr<3>) { + // CHECK-LABEL: rocdl.s.barrier.signal.var + // CHECK: rocdl.s.barrier.signal.var %[[PTR:.+]], 1 + rocdl.s.barrier.signal.var %ptr, 1 + llvm.return +} + +llvm.func @rocdl.s.barrier.join(%ptr : !llvm.ptr<3>) { + // CHECK-LABEL: rocdl.s.barrier.join + // CHECK: rocdl.s.barrier.join %[[PTR:.+]] + rocdl.s.barrier.join %ptr + llvm.return +} + +llvm.func @rocdl.s.barrier.leave() { + // CHECK-LABEL: rocdl.s.barrier.leave + // CHECK: rocdl.s.barrier.leave 1 + rocdl.s.barrier.leave 1 + llvm.return +} + llvm.func @rocdl.s.barrier.wait() { // CHECK-LABEL: rocdl.s.barrier.wait // CHECK: rocdl.s.barrier.wait -1 @@ -965,6 +993,20 @@ llvm.func @rocdl.s.barrier.wait() { llvm.return } +llvm.func @rocdl.s.barrier.signal.isfirst() { + // CHECK-LABEL: rocdl.s.barrier.signal.isfirst + // CHECK: rocdl.s.barrier.signal.isfirst 1 + %0 = rocdl.s.barrier.signal.isfirst 1 : i1 + llvm.return +} + +llvm.func @rocdl.s.get.barrier.state() { + // CHECK-LABEL: rocdl.s.get.barrier.state + // CHECK: rocdl.s.get.barrier.state 1 + %0 = rocdl.s.get.barrier.state 1 : i32 + llvm.return +} + llvm.func @rocdl.s.wait.dscnt() { // CHECK-LABEL: rocdl.s.wait.dscnt // CHECK: rocdl.s.wait.dscnt 0 diff --git a/mlir/test/Target/LLVMIR/rocdl.mlir b/mlir/test/Target/LLVMIR/rocdl.mlir index 1c0c2eba002aa..fdd2c91f6a5b5 100644 --- a/mlir/test/Target/LLVMIR/rocdl.mlir +++ b/mlir/test/Target/LLVMIR/rocdl.mlir @@ -192,6 +192,13 @@ llvm.func @rocdl.barrier() { llvm.return } +llvm.func @rocdl.s.barrier.init(%ptr : !llvm.ptr<3>) { + // CHECK-LABEL: rocdl.s.barrier.init + // CHECK: call void @llvm.amdgcn.s.barrier.init(ptr addrspace(3) %[[PTR:.+]], i32 1) + rocdl.s.barrier.init %ptr, 1 + llvm.return +} + llvm.func @rocdl.s.barrier.signal() { // CHECK-LABEL: rocdl.s.barrier.signal // CHECK-NEXT: call void @llvm.amdgcn.s.barrier.signal(i32 -1) @@ -199,6 +206,27 @@ llvm.func @rocdl.s.barrier.signal() { llvm.return } +llvm.func @rocdl.s.barrier.signal.var(%ptr : !llvm.ptr<3>) { + // CHECK-LABEL: rocdl.s.barrier.signal.var + // CHECK: call void @llvm.amdgcn.s.barrier.signal.var(ptr addrspace(3) %[[PTR:.+]], i32 1) + rocdl.s.barrier.signal.var %ptr, 1 + llvm.return +} + +llvm.func @rocdl.s.barrier.join(%ptr : !llvm.ptr<3>) { + // CHECK-LABEL: rocdl.s.barrier.join + // CHECK: call void @llvm.amdgcn.s.barrier.join(ptr addrspace(3) %[[PTR:.+]]) + rocdl.s.barrier.join %ptr + llvm.return +} + +llvm.func @rocdl.s.barrier.leave() { + // CHECK-LABEL: rocdl.s.barrier.leave + // CHECK: call void @llvm.amdgcn.s.barrier.leave(i16 1) + rocdl.s.barrier.leave 1 + llvm.return +} + llvm.func @rocdl.s.barrier.wait() { // CHECK-LABEL: rocdl.s.barrier.wait // CHECK-NEXT: call void @llvm.amdgcn.s.barrier.wait(i16 -1) @@ -206,6 +234,20 @@ llvm.func @rocdl.s.barrier.wait() { llvm.return } +llvm.func @rocdl.s.barrier.signal.isfirst() { + // CHECK-LABEL: rocdl.s.barrier.signal.isfirst + // CHECK: %[[OUT:.+]] = call i1 @llvm.amdgcn.s.barrier.signal.isfirst(i32 1) + %0 = rocdl.s.barrier.signal.isfirst 1 : i1 + llvm.return +} + +llvm.func @rocdl.s.get.barrier.state() { + // CHECK-LABEL: rocdl.s.get.barrier.state + // CHECK: %[[STATE:.+]] = call i32 @llvm.amdgcn.s.get.barrier.state(i32 1) + %0 = rocdl.s.get.barrier.state 1 : i32 + llvm.return +} + llvm.func @rocdl.s.wait.dscnt() { // CHECK-LABEL: rocdl.s.wait.dscnt // CHECK-NEXT: call void @llvm.amdgcn.s.wait.dscnt(i16 0)