[MLIR][NVVM] Enable inlining of func's calling nvvm intrinsics

[wsmoses]

No description provided.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-llvm

Author: William Moses (wsmoses)

Changes

---

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

4 Files Affected:

• (modified) mlir/include/mlir/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.h (+7)
• (modified) mlir/include/mlir/InitAllDialects.h (+1)
• (modified) mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp (+7)
• (modified) mlir/test/Dialect/LLVMIR/inlining.mlir (+15)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.h b/mlir/include/mlir/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.h
index e99b0476a6b107..48882f435428a8 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.h
+++ b/mlir/include/mlir/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.h
@@ -23,6 +23,13 @@ namespace LLVM {
 void registerInlinerInterface(DialectRegistry &registry);
 
 } // namespace LLVM
+
+namespace NVVM {
+/// Register the `NVVMInlinerInterface` implementation of
+/// `DialectInlinerInterface` with the NVVM dialect.
+void registerInlinerInterface(DialectRegistry &registry);
+}
+
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LLVMIR_TRANSFORMS_INLINERINTERFACEIMPL_H
diff --git a/mlir/include/mlir/InitAllDialects.h b/mlir/include/mlir/InitAllDialects.h
index c102f811cce4b1..0da82825c82878 100644
--- a/mlir/include/mlir/InitAllDialects.h
+++ b/mlir/include/mlir/InitAllDialects.h
@@ -167,6 +167,7 @@ inline void registerAllDialects(DialectRegistry &registry) {
   gpu::registerBufferDeallocationOpInterfaceExternalModels(registry);
   gpu::registerValueBoundsOpInterfaceExternalModels(registry);
   LLVM::registerInlinerInterface(registry);
+  NVVM::registerInlinerInterface(registry);
   linalg::registerAllDialectInterfaceImplementations(registry);
   linalg::registerRuntimeVerifiableOpInterfaceExternalModels(registry);
   memref::registerAllocationOpInterfaceExternalModels(registry);
diff --git a/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp b/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
index b3bed5ab5f412f..233cadebeec026 100644
--- a/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
@@ -14,6 +14,7 @@
 #include "mlir/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.h"
 #include "mlir/Analysis/SliceWalk.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
@@ -815,3 +816,9 @@ void mlir::LLVM::registerInlinerInterface(DialectRegistry &registry) {
     dialect->addInterfaces<LLVMInlinerInterface>();
   });
 }
+
+void mlir::NVVM::registerInlinerInterface(DialectRegistry &registry) {
+  registry.addExtension(+[](MLIRContext *ctx, NVVM::NVVMDialect *dialect) {
+    dialect->addInterfaces<LLVMInlinerInterface>();
+  });
+}
diff --git a/mlir/test/Dialect/LLVMIR/inlining.mlir b/mlir/test/Dialect/LLVMIR/inlining.mlir
index edaac4da0b044c..8081fc5ca9353c 100644
--- a/mlir/test/Dialect/LLVMIR/inlining.mlir
+++ b/mlir/test/Dialect/LLVMIR/inlining.mlir
@@ -641,6 +641,21 @@ llvm.func @caller(%ptr : !llvm.ptr) -> i32 {
 
 // -----
 
+llvm.func @threadidx() -> i32 {
+  %tid = call i32 @llvm.nvvm.read.ptx.sreg.tid.x()
+  llvm.return %tid : i32
+}
+
+// CHECK-LABEL: func @caller
+llvm.func @caller() -> i32 {
+  // CHECK-NOT: llvm.call @private_func
+  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.tid.x
+  %z = llvm.call @private_func() : () -> (i32)
+  llvm.return %z : i32
+}
+
+// -----
+
 llvm.func @vararg_func(...) {
   llvm.return
 }

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[ftynse]

Note that the CI failure looks relevant. Maybe we are not building with NVVM on Windows, or something is wrong with registration.

[wsmoses]

Yeah I’m honestly confused by the windows failure (especially since Linux succeeds), if you have any suggestions