[Mlir-commits] [mlir] [mlir][ROCDL] Add synchronization primitives (#1077) (PR #80888)

[Giuseppe Rossini]

https://github.com/giuseros updated https://github.com/llvm/llvm-project/pull/80888

>From ff2b2a1e5c05af5647e8123bbdbf6e21b1135227 Mon Sep 17 00:00:00 2001
From: Giuseppe Rossini <giuseppe.rossini at amd.com>
Date: Mon, 15 May 2023 10:31:34 +0100
Subject: [PATCH 1/2] [mlir][ROCDL] Add synchronization primitives (#1077)

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
---
 mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td | 12 ++++++++++++
 mlir/test/Target/LLVMIR/rocdl.mlir           | 20 ++++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index 638e46a2f9c752..51a5c2b9f129a8 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 26123300d74888..23a9986701bac3 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>,

>From b3acaf9079406009b3e9eab2bfbd8ad9a30d7858 Mon Sep 17 00:00:00 2001
From: Giuseppe Rossini <giuseppe.rossini at amd.com>
Date: Wed, 7 Feb 2024 11:45:44 +0000
Subject: [PATCH 2/2] Address review feedbacks

---
 mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td |  8 ++++++--
 mlir/test/Dialect/LLVMIR/rocdl.mlir          | 12 ++++++++++++
 mlir/test/Target/LLVMIR/rocdl.mlir           | 12 ++++--------
 3 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index 51a5c2b9f129a8..962c159e68a2ee 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
@@ -193,15 +193,19 @@ def ROCDL_BarrierOp : ROCDL_Op<"barrier"> {
 }
 
 def ROCDL_SetPrioOp : ROCDL_IntrOp<"s.setprio", [], [], [], 0>,
-  Arguments<(ins I16:$priority)> {
+  Arguments<(ins I16Attr:$priority)> {
   let results = (outs);
   let assemblyFormat = "$priority attr-dict";
+  string llvmBuilder =
+    "createIntrinsicCall(builder, llvm::Intrinsic::amdgcn_s_setprio,builder.getInt16(op.getPriority()));";
 }
 
 def ROCDL_SchedBarrier : ROCDL_IntrOp<"sched.barrier", [], [], [], 0>,
-  Arguments<(ins I32:$mask)> {
+  Arguments<(ins I32Attr:$mask)> {
   let results = (outs);
   let assemblyFormat = "$mask attr-dict";
+  string llvmBuilder =
+    "createIntrinsicCall(builder, llvm::Intrinsic::amdgcn_sched_barrier,builder.getInt32(op.getMask()));";
 }
 
 
diff --git a/mlir/test/Dialect/LLVMIR/rocdl.mlir b/mlir/test/Dialect/LLVMIR/rocdl.mlir
index 5a14df9ef9f8dc..89e8e7836c3a0c 100644
--- a/mlir/test/Dialect/LLVMIR/rocdl.mlir
+++ b/mlir/test/Dialect/LLVMIR/rocdl.mlir
@@ -35,6 +35,18 @@ func.func @rocdl.barrier() {
   llvm.return
 }
 
+func.func @rocdl.sched_barrier() {
+  // CHECK: rocdl.sched.barrier
+  rocdl.sched.barrier 0
+  llvm.return
+}
+
+func.func @rocdl.setprio() {
+  // CHECK: rocdl.s.setprio
+  rocdl.s.setprio 0
+  llvm.return
+}
+
 func.func @rocdl.xdlops(%arg0 : f32, %arg1 : f32,
                    %arg2 : vector<32xf32>, %arg3 : i32,
                    %arg4 : vector<16xf32>, %arg5 : vector<4xf32>,
diff --git a/mlir/test/Target/LLVMIR/rocdl.mlir b/mlir/test/Target/LLVMIR/rocdl.mlir
index 23a9986701bac3..06b78650c8d010 100644
--- a/mlir/test/Target/LLVMIR/rocdl.mlir
+++ b/mlir/test/Target/LLVMIR/rocdl.mlir
@@ -91,22 +91,18 @@ llvm.func @rocdl.barrier() {
 }
 
 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
+  rocdl.s.setprio 0
   // CHECK-NEXT: call void @llvm.amdgcn.s.setprio(i16 1)
-  rocdl.s.setprio %one
+  rocdl.s.setprio 1
   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
+  rocdl.sched.barrier 0
   // CHECK-NEXT: call void @llvm.amdgcn.sched.barrier(i32 1)
-  rocdl.sched.barrier %one
+  rocdl.sched.barrier 1
   llvm.return
 }