[Mlir-commits] [mlir] [mlir][NVVM] Add constant memory space identifier (PR #111141)

[Matthias Springer]

https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/111141

Also use these enums in `BasicPtxBuilderInferface.cpp`.

>From 3841c0c21e73b49723400b27c151862f9501396f Mon Sep 17 00:00:00 2001
From: Matthias Springer <mspringer at nvidia.com>
Date: Fri, 4 Oct 2024 14:27:25 +0200
Subject: [PATCH] [mlir][NVVM] Add constant memory space identifier

Also use these enums in `BasicPtxBuilderInferface.cpp`.
---
 mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h          | 4 +++-
 mlir/lib/Dialect/LLVMIR/IR/BasicPtxBuilderInterface.cpp | 5 ++---
 2 files changed, 5 insertions(+), 4 deletions(-)

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';