[mlir][ROCDL] Add synchronization primitives (#1077)

[giuseros]

This PR is adding to MLIR two LLVM intrisics:

• llvm.amdgcn.s.setprio which sets the priority of a wave for the GPU scheduler
• llvm.amdgcn.sched.barrier which sets a software barrier so that the scheduler cannot move instructions around

[llvmbot]

@llvm/pr-subscribers-mlir-llvm

@llvm/pr-subscribers-mlir

Author: Giuseppe Rossini (giuseros)

Changes

This PR is adding to MLIR two LLVM intrisics:

• llvm.amdgcn.s.setprio which sets the priority of a wave for the GPU scheduler
• llvm.amdgcn.sched.barrier which sets a software barrier so that the scheduler cannot move instructions around

---

Full diff: https://github.com/llvm/llvm-project/pull/80888.diff

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td (+12)
• (modified) mlir/test/Target/LLVMIR/rocdl.mlir (+20)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index 638e46a2f9c75..51a5c2b9f129a 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
@@ -192,6 +192,18 @@ def ROCDL_BarrierOp : ROCDL_Op<"barrier"> {
   let assemblyFormat = "attr-dict";
 }
 
+def ROCDL_SetPrioOp : ROCDL_IntrOp<"s.setprio", [], [], [], 0>,
+  Arguments<(ins I16:$priority)> {
+  let results = (outs);
+  let assemblyFormat = "$priority attr-dict";
+}
+
+def ROCDL_SchedBarrier : ROCDL_IntrOp<"sched.barrier", [], [], [], 0>,
+  Arguments<(ins I32:$mask)> {
+  let results = (outs);
+  let assemblyFormat = "$mask attr-dict";
+}
+
 
 //===---------------------------------------------------------------------===//
 // Xdlops intrinsics
diff --git a/mlir/test/Target/LLVMIR/rocdl.mlir b/mlir/test/Target/LLVMIR/rocdl.mlir
index 26123300d7488..23a9986701bac 100644
--- a/mlir/test/Target/LLVMIR/rocdl.mlir
+++ b/mlir/test/Target/LLVMIR/rocdl.mlir
@@ -90,6 +90,26 @@ llvm.func @rocdl.barrier() {
   llvm.return
 }
 
+llvm.func @rocdl.setprio() {
+  %zero = llvm.mlir.constant(0 : i16) : i16
+  %one = llvm.mlir.constant(1 : i16) : i16
+  // CHECK: call void @llvm.amdgcn.s.setprio(i16 0)
+  rocdl.s.setprio %zero
+  // CHECK-NEXT: call void @llvm.amdgcn.s.setprio(i16 1)
+  rocdl.s.setprio %one
+  llvm.return
+}
+
+llvm.func @rocdl.schedbarrier() {
+  %zero = llvm.mlir.constant(0 : i32) : i32
+  %one = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: call void @llvm.amdgcn.sched.barrier(i32 0)
+  rocdl.sched.barrier %zero
+  // CHECK-NEXT: call void @llvm.amdgcn.sched.barrier(i32 1)
+  rocdl.sched.barrier %one
+  llvm.return
+}
+
 llvm.func @rocdl.xdlops(%arg0 : f32, %arg1 : f32,
                    %arg2 : vector<32 x f32>, %arg3: i32,
                    %arg4 : vector<16 x f32>, %arg5 : vector<4xf32>,

[krzysz00]

Overall, seems reasonable to me, though you'll want to stick the trivial "can this go through mlir-opt" test over in mlir/test/Dialect/LLVMIR/rocdl.mlir as well.

[giuseros]

Hi @krzysz00 ,thanks for the review! Do you mind elaborate? Do I need to add an additional check other than // RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s ?

[krzysz00]

Yeah, there's a set of tests over in mlir/test/Dialect/LLVMIR/rocdl.mlir (note: different filename) that just consists of rocdl ops and does the trivial thing of ensuring that they come back as themselves.

[krzysz00]

Approved!