[Mlir-commits] [mlir] ab71b77 - [mlir][rocdl] Add rocdl inlining interface (#163058)

[llvmlistbot at llvm.org]

Author: Ivan Butygin
Date: 2025-10-12T15:20:31+03:00
New Revision: ab71b7793a0480f820c62c0bbc69e6c8e7434c2a

URL: https://github.com/llvm/llvm-project/commit/ab71b7793a0480f820c62c0bbc69e6c8e7434c2a
DIFF: https://github.com/llvm/llvm-project/commit/ab71b7793a0480f820c62c0bbc69e6c8e7434c2a.diff

LOG: [mlir][rocdl] Add rocdl inlining interface (#163058)

All rocdl ops should be safe to inline.

Added: 
    mlir/test/Dialect/LLVMIR/inlining-rocdl.mlir

Modified: 
    mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp
index 17371ecb79e7d..6d54bb6b82ebf 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp
@@ -23,6 +23,7 @@
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Operation.h"
+#include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
@@ -180,6 +181,15 @@ void RawBufferAtomicUMinOp::print(mlir::OpAsmPrinter &p) {
 // ROCDLDialect initialization, type parsing, and registration.
 //===----------------------------------------------------------------------===//
 
+namespace {
+struct ROCDLInlinerInterface final : DialectInlinerInterface {
+  using DialectInlinerInterface::DialectInlinerInterface;
+  bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {
+    return true;
+  }
+};
+} // namespace
+
 // TODO: This should be the llvm.rocdl dialect once this is supported.
 void ROCDLDialect::initialize() {
   addOperations<
@@ -194,6 +204,7 @@ void ROCDLDialect::initialize() {
 
   // Support unknown operations because not all ROCDL operations are registered.
   allowUnknownOperations();
+  addInterfaces<ROCDLInlinerInterface>();
   declarePromisedInterface<gpu::TargetAttrInterface, ROCDLTargetAttr>();
 }
 

diff  --git a/mlir/test/Dialect/LLVMIR/inlining-rocdl.mlir b/mlir/test/Dialect/LLVMIR/inlining-rocdl.mlir
new file mode 100644
index 0000000000000..7fd97ef168e5b
--- /dev/null
+++ b/mlir/test/Dialect/LLVMIR/inlining-rocdl.mlir
@@ -0,0 +1,14 @@
+// RUN: mlir-opt %s --inline | FileCheck %s
+
+llvm.func @threadidx() -> i32 {
+  %tid = rocdl.workitem.id.x : i32
+  llvm.return %tid : i32
+}
+
+// CHECK-LABEL: func @caller
+llvm.func @caller() -> i32 {
+  // CHECK-NOT: llvm.call @threadidx
+  // CHECK: rocdl.workitem.id.x
+  %z = llvm.call @threadidx() : () -> (i32)
+  llvm.return %z : i32
+}