[mlir][LLVM|ptr] Add the #llvm.address_space attribute, and allow ptr translation (#156333)

This commit introduces the `#llvm.address_space` attribute. This
attribute implements the `ptr::MemorySpaceAttrInterface`, establishing
the semantics of the LLVM address space.
This allows making `!ptr.ptr` translatable to LLVM IR as long it uses
the `#llvm.address_space` attribute.
Concretely, `!ptr.ptr<#llvm.address_space<N>>` now translates to `ptr
addrspace(N)`.

Additionally, this patch makes `PtrLikeTypes` with no metadata, no
element type, and with `#llvm.address_space` memory space, compatible
with the LLVM dialect.

**Infrastructure Updates:**
- Refactor `ptr::MemorySpaceAttrInterface` to include DataLayout
parameter for better validation
- Add new utility functions `LLVM::isLoadableType()` and
`LLVM::isTypeCompatibleWithAtomicOp()`
- Update type compatibility checks to support ptr-like types with LLVM
address spaces
- Splice the `MemorySpaceAttrInterface` to its own library, so the
LLVMDialect won't depend on the PtrDialect yet

**Translation Support:**
- New `PtrToLLVMIRTranslation` module for converting ptr dialect to LLVM
IR
- Type translation support for ptr types with LLVM address spaces
- Proper address space preservation during IR lowering

Example:
```mlir
llvm.func @llvm_ops_with_ptr_values(%arg0: !llvm.ptr) {
  %1 = llvm.load %arg0 : !llvm.ptr -> !ptr.ptr<#llvm.address_space<1>>
  llvm.store %1, %arg0 : !ptr.ptr<#llvm.address_space<1>>, !llvm.ptr
  llvm.return
}
```
Translates to:
```llvmir
; ModuleID = 'LLVMDialectModule'
source_filename = "LLVMDialectModule"

define void @llvm_ops_with_ptr_values(ptr %0) {
  %2 = load ptr addrspace(1), ptr %0, align 8
  store ptr addrspace(1) %2, ptr %0, align 8
  ret void
}

!llvm.module.flags = !{!0}

!0 = !{i32 2, !"Debug Info Version", i32 3}
```

Author: fabianmcg

mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td

@@ -11,6 +11,7 @@
 
 include "mlir/Dialect/LLVMIR/LLVMDialect.td"
 include "mlir/Dialect/LLVMIR/LLVMInterfaces.td"
+include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.td"
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/CommonAttrConstraints.td"
 include "mlir/Interfaces/DataLayoutInterfaces.td"
@@ -23,6 +24,41 @@ class LLVM_Attr<string name, string attrMnemonic,
   let mnemonic = attrMnemonic;
 }
 
+//===----------------------------------------------------------------------===//
+// AddressSpaceAttr
+//===----------------------------------------------------------------------===//
+
+def LLVM_AddressSpaceAttr :
+    LLVM_Attr<"AddressSpace", "address_space", [
+    DeclareAttrInterfaceMethods<MemorySpaceAttrInterface>
+  ]> {
+  let summary = "LLVM address space";
+  let description = [{
+    The `address_space` attribute represents an LLVM address space. It takes an
+    unsigned integer parameter that specifies the address space number.
+
+    Different address spaces in LLVM can have different properties:
+    - Address space 0 is the default/generic address space
+    - Other address spaces may have specific semantics (e.g., shared memory,
+      constant memory, etc.) depending on the target architecture
+
+    Example:
+
+    ```mlir
+    // Address space 0 (default)
+    #llvm.address_space<0>
+
+    // Address space 1 (e.g., global memory on some targets)
+    #llvm.address_space<1>
+
+    // Address space 3 (e.g., shared memory on some GPU targets)
+    #llvm.address_space<3>
+    ```
+  }];
+  let parameters = (ins "unsigned":$addressSpace);
+  let assemblyFormat = "`<` $addressSpace `>`";
+}
+
 //===----------------------------------------------------------------------===//
 // CConvAttr
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h

@@ -15,6 +15,7 @@
 #define MLIR_DIALECT_LLVMIR_LLVMATTRS_H_
 
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
+#include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
 #include <optional>

mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h

@@ -28,6 +28,7 @@ namespace mlir {
 
 class AsmParser;
 class AsmPrinter;
+class DataLayout;
 
 namespace LLVM {
 class LLVMDialect;
@@ -111,6 +112,15 @@ bool isCompatibleFloatingPointType(Type type);
 /// dialect pointers and LLVM dialect scalable vector types.
 bool isCompatibleVectorType(Type type);
 
+/// Returns `true` if the given type is a loadable type compatible with the LLVM
+/// dialect.
+bool isLoadableType(Type type);
+
+/// Returns true if the given type is supported by atomic operations. All
+/// integer, float, and pointer types with a power-of-two bitsize and a minimal
+/// size of 8 bits are supported.
+bool isTypeCompatibleWithAtomicOp(Type type, const DataLayout &dataLayout);
+
 /// Returns the element count of any LLVM-compatible vector type.
 llvm::ElementCount getVectorNumElements(Type type);
 

mlir/include/mlir/Dialect/Ptr/IR/CMakeLists.txt

@@ -7,8 +7,6 @@ mlir_tablegen(PtrOpsAttrs.cpp.inc -gen-attrdef-defs -attrdefs-dialect=ptr)
 add_mlir_dialect_tablegen_target(MLIRPtrOpsAttributesIncGen)
 
 set(LLVM_TARGET_DEFINITIONS MemorySpaceInterfaces.td)
-mlir_tablegen(MemorySpaceInterfaces.h.inc -gen-op-interface-decls)
-mlir_tablegen(MemorySpaceInterfaces.cpp.inc -gen-op-interface-defs)
 mlir_tablegen(MemorySpaceAttrInterfaces.h.inc -gen-attr-interface-decls)
 mlir_tablegen(MemorySpaceAttrInterfaces.cpp.inc -gen-attr-interface-defs)
 add_mlir_dialect_tablegen_target(MLIRPtrMemorySpaceInterfacesIncGen)

mlir/include/mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.h

@@ -17,8 +17,12 @@
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/OpDefinition.h"
 
+#include <functional>
+#include <optional>
+
 namespace mlir {
 class Operation;
+class DataLayout;
 namespace ptr {
 enum class AtomicBinOp : uint32_t;
 enum class AtomicOrdering : uint32_t;
@@ -27,6 +31,4 @@ enum class AtomicOrdering : uint32_t;
 
 #include "mlir/Dialect/Ptr/IR/MemorySpaceAttrInterfaces.h.inc"
 
-#include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.h.inc"
-
 #endif // MLIR_DIALECT_PTR_IR_MEMORYSPACEINTERFACES_H

mlir/include/mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.td

@@ -43,6 +43,7 @@ def MemorySpaceAttrInterface : AttrInterface<"MemorySpaceAttrInterface"> {
       /*args=*/        (ins "::mlir::Type":$type,
                             "::mlir::ptr::AtomicOrdering":$ordering,
                             "std::optional<int64_t>":$alignment,
+                            "const ::mlir::DataLayout *":$dataLayout,
                             "::llvm::function_ref<::mlir::InFlightDiagnostic()>":$emitError)
     >,
     InterfaceMethod<
@@ -58,6 +59,7 @@ def MemorySpaceAttrInterface : AttrInterface<"MemorySpaceAttrInterface"> {
       /*args=*/        (ins "::mlir::Type":$type,
                             "::mlir::ptr::AtomicOrdering":$ordering,
                             "std::optional<int64_t>":$alignment,
+                            "const ::mlir::DataLayout *":$dataLayout,
                             "::llvm::function_ref<::mlir::InFlightDiagnostic()>":$emitError)
     >,
     InterfaceMethod<
@@ -74,6 +76,7 @@ def MemorySpaceAttrInterface : AttrInterface<"MemorySpaceAttrInterface"> {
                             "::mlir::Type":$type,
                             "::mlir::ptr::AtomicOrdering":$ordering,
                             "std::optional<int64_t>":$alignment,
+                            "const ::mlir::DataLayout *":$dataLayout,
                             "::llvm::function_ref<::mlir::InFlightDiagnostic()>":$emitError)
     >,
     InterfaceMethod<
@@ -91,6 +94,7 @@ def MemorySpaceAttrInterface : AttrInterface<"MemorySpaceAttrInterface"> {
                             "::mlir::ptr::AtomicOrdering":$successOrdering,
                             "::mlir::ptr::AtomicOrdering":$failureOrdering,
                             "std::optional<int64_t>":$alignment,
+                            "const ::mlir::DataLayout *":$dataLayout,
                             "::llvm::function_ref<::mlir::InFlightDiagnostic()>":$emitError)
     >,
     InterfaceMethod<

mlir/include/mlir/Dialect/Ptr/IR/PtrAttrs.h

@@ -19,10 +19,9 @@
 #include "llvm/Support/TypeSize.h"
 
 #include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.h"
+#include "mlir/Dialect/Ptr/IR/PtrEnums.h"
 
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/Ptr/IR/PtrOpsAttrs.h.inc"
 
-#include "mlir/Dialect/Ptr/IR/PtrOpsEnums.h.inc"
-
 #endif // MLIR_DIALECT_PTR_IR_PTRATTRS_H

mlir/include/mlir/Dialect/Ptr/IR/PtrEnums.h

@@ -0,0 +1,21 @@
+//===- PtrEnums.h - `ptr` dialect enums -------------------------*- C++ -*-===//
+//
+// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares the `ptr` dialect enums.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_PTR_IR_PTRENUMS_H
+#define MLIR_DIALECT_PTR_IR_PTRENUMS_H
+
+#include "mlir/IR/BuiltinAttributeInterfaces.h"
+#include "mlir/IR/OpImplementation.h"
+
+#include "mlir/Dialect/Ptr/IR/PtrOpsEnums.h.inc"
+
+#endif // MLIR_DIALECT_PTR_IR_PTRENUMS_H

mlir/include/mlir/Target/LLVMIR/Dialect/All.h

@@ -25,6 +25,7 @@
 #include "mlir/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/OpenACC/OpenACCToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.h"
+#include "mlir/Target/LLVMIR/Dialect/Ptr/PtrToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/SPIRV/SPIRVToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/VCIX/VCIXToLLVMIRTranslation.h"
@@ -45,6 +46,7 @@ static inline void registerAllToLLVMIRTranslations(DialectRegistry &registry) {
   registerNVVMDialectTranslation(registry);
   registerOpenACCDialectTranslation(registry);
   registerOpenMPDialectTranslation(registry);
+  registerPtrDialectTranslation(registry);
   registerROCDLDialectTranslation(registry);
   registerSPIRVDialectTranslation(registry);
   registerVCIXDialectTranslation(registry);

mlir/include/mlir/Target/LLVMIR/Dialect/Ptr/PtrToLLVMIRTranslation.h

@@ -0,0 +1,31 @@
+//===- PtrToLLVMIRTranslation.h - `ptr` to LLVM IR --------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This provides registration calls for `ptr` dialect to LLVM IR translation.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TARGET_LLVMIR_DIALECT_PTR_PTRTOLLVMIRTRANSLATION_H
+#define MLIR_TARGET_LLVMIR_DIALECT_PTR_PTRTOLLVMIRTRANSLATION_H
+
+namespace mlir {
+
+class DialectRegistry;
+class MLIRContext;
+
+/// Register the `ptr` dialect and the translation from it to the LLVM IR in the
+/// given registry;
+void registerPtrDialectTranslation(DialectRegistry &registry);
+
+/// Register the `ptr` dialect and the translation from it in the registry
+/// associated with the given context.
+void registerPtrDialectTranslation(MLIRContext &context);
+
+} // namespace mlir
+
+#endif // MLIR_TARGET_LLVMIR_DIALECT_PTR_PTRTOLLVMIRTRANSLATION_H