[mlir][NVVM] Add constant memory space identifier

[matthias-springer]

Also use these enums in BasicPtxBuilderInferface.cpp.

[llvmbot]

@llvm/pr-subscribers-mlir-llvm

@llvm/pr-subscribers-mlir

Author: Matthias Springer (matthias-springer)

Changes

Also use these enums in BasicPtxBuilderInferface.cpp.

---

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

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h (+3-1)
• (modified) mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp (+2-3)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h b/mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h
index 08019e77ae6af8..4fd00ff929bd70 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h
@@ -35,7 +35,9 @@ enum NVVMMemorySpace {
   /// Global memory space identifier.
   kGlobalMemorySpace = 1,
   /// Shared memory space identifier.
-  kSharedMemorySpace = 3
+  kSharedMemorySpace = 3,
+  /// Constant memory space identifier.
+  kConstantMemorySpace = 4
 };
 
 /// Return the element type and number of elements associated with a wmma matrix
diff --git a/mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp b/mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp
index b109f00c3da13a..d181700f757a5b 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp
@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/LLVMIR/BasicPtxBuilderInterface.h"
+#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 
 #define DEBUG_TYPE "ptx-builder"
 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
@@ -26,8 +27,6 @@
 using namespace mlir;
 using namespace NVVM;
 
-static constexpr int64_t kSharedMemorySpace = 3;
-
 static char getRegisterType(Type type) {
   if (type.isInteger(1))
     return 'b';
@@ -43,7 +42,7 @@ static char getRegisterType(Type type) {
     return 'd';
   if (auto ptr = dyn_cast<LLVM::LLVMPointerType>(type)) {
     // Shared address spaces is addressed with 32-bit pointers.
-    if (ptr.getAddressSpace() == kSharedMemorySpace) {
+    if (ptr.getAddressSpace() == NVVMMemorySpace::kSharedMemorySpace) {
       return 'r';
     }
     return 'l';

[bjacob]

I didn't see the revert #111169 while working on that Bazel breakage in a different way: #111172. Please decide if you want to keep that (and then, that unblocks re-landing this PR, making the circular dependency a non-issue now that the two Bazel targets are one) or revert that. @matthias-springer @dklimkin

[matthias-springer]

No strong preference from my side. This was just a random cleanup while adding the constant memory space.

[bjacob]

I see. I'd suggest leaving my PR in, then, because there's this tension between CMake allowing circular dependencies and Bazel not allowing them, so the folding of the two targets removes some of the tension between CMake and Bazel here.