[Mlir-commits] [mlir] [MLIR][ROCDL] Added `SchedGroupBarrier` and `IglpOpt` ops (PR #112237)

[llvmlistbot at llvm.org]

https://github.com/ravil-mobile updated https://github.com/llvm/llvm-project/pull/112237

>From 1f3ad8fada5e4868ad8b980a8d43b62dffb1fd88 Mon Sep 17 00:00:00 2001
From: ravil-mobile <ravil.aviva.com at gmail.com>
Date: Mon, 14 Oct 2024 17:41:54 +0000
Subject: [PATCH] [MLIR][ROCDL] Added `sched.group.barrier` and
 `rocdl.iglp.opt`

Currently, the ROCDL dialect contains only `rocdl.sched.barrier` op.
This commit adds missing `sched.group.barrier` and `rocdl.iglp.opt` ops.
The ops are converted to the corresponding intrinsic calls during
the translation from MLIR to LLVM IRs. This intrinsics are hints to
the instruction scheduler of the AMDGPU backend.
---
 mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td | 18 +++++++++++++++
 mlir/test/Dialect/LLVMIR/rocdl.mlir          | 24 ++++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index b80d9ae88910c4..c40ae4b1016b49 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
@@ -297,6 +297,24 @@ def ROCDL_SchedBarrier : ROCDL_IntrOp<"sched.barrier", [], [], [], 0>,
     "createIntrinsicCall(builder, llvm::Intrinsic::amdgcn_sched_barrier,builder.getInt32(op.getMask()));";
 }
 
+def ROCDL_SchedGroupBarrier : ROCDL_IntrOp<"sched.group.barrier", [], [], [], 0>,
+  Arguments<(ins I32Attr:$mask, I32Attr:$size, I32Attr:$groupId)> {
+  let results = (outs);
+  let assemblyFormat = "$mask `,` $size `,` $groupId attr-dict";
+  string llvmBuilder = [{
+    createIntrinsicCall(builder,
+      llvm::Intrinsic::amdgcn_sched_group_barrier,
+      {builder.getInt32(op.getMask()), builder.getInt32(op.getSize()), builder.getInt32(op.getGroupId())});
+  }];
+}
+
+def ROCDL_IglpOpt : ROCDL_IntrOp<"iglp.opt", [], [], [], 0>,
+  Arguments<(ins I32Attr:$variant)> {
+  let results = (outs);
+  let assemblyFormat = "$variant attr-dict";
+  string llvmBuilder =
+    "createIntrinsicCall(builder, llvm::Intrinsic::amdgcn_iglp_opt,builder.getInt32(op.getVariant()));";
+}
 
 //===---------------------------------------------------------------------===//
 // Xdlops intrinsics
diff --git a/mlir/test/Dialect/LLVMIR/rocdl.mlir b/mlir/test/Dialect/LLVMIR/rocdl.mlir
index 397d66d92bc5d5..90c64af8715364 100644
--- a/mlir/test/Dialect/LLVMIR/rocdl.mlir
+++ b/mlir/test/Dialect/LLVMIR/rocdl.mlir
@@ -1,4 +1,5 @@
 // RUN: mlir-opt %s -split-input-file -verify-diagnostics | FileCheck %s
+// RUN: mlir-opt %s -split-input-file -verify-diagnostics -convert-to-llvm | mlir-translate -mlir-to-llvmir | FileCheck %s -check-prefix=ROCDL2LLVM
 
 func.func @rocdl_special_regs() -> i32 {
   // CHECK-LABEL: rocdl_special_regs
@@ -387,3 +388,26 @@ gpu.module @module_1 [#rocdl.target<O = 1, chip = "gfx900", abi = "500", link =
 
 gpu.module @module_2 [#rocdl.target<chip = "gfx900">, #rocdl.target<chip = "gfx90a">] {
 }
+
+// -----
+
+// ROCDL2LLVM: @rocdl_sched_barrier
+func.func @rocdl_sched_barrier() {
+  // ROCDL2LLVM-NEXT: call void @llvm.amdgcn.sched.barrier(i32 0)
+  rocdl.sched.barrier 0
+  llvm.return
+}
+
+// ROCDL2LLVM: @rocdl_sched_group_barrier
+func.func @rocdl_sched_group_barrier() {
+  // ROCDL2LLVM-NEXT: call void @llvm.amdgcn.sched.group.barrier(i32 8, i32 1, i32 0)
+  rocdl.sched.group.barrier 8, 1, 0
+  llvm.return
+}
+
+// ROCDL2LLVM: @rocdl_iglp_opt
+func.func @rocdl_iglp_opt() {
+  // ROCDL2LLVM-NEXT: call void @llvm.amdgcn.iglp.opt(i32 0)
+  rocdl.iglp.opt 0
+  llvm.return
+}