[mlir][LLVM] Add LLVMAddrSpaceAttrInterface and NVVMMemorySpaceAttr #157339
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
|
- Introduces LLVMAddrSpaceAttrInterface for LLVM address space attributes - Adds NVVMMemorySpaceAttr implementing both LLVMAddrSpaceAttrInterface and MemorySpaceAttrInterface - Converts NVVM memory space constants from enum to MLIR enums - Updates all NVVM memory space references to use new attribute system - Adds support for NVVM memory spaces in ptr dialect translation To convert the memory space enum to the new enum class use: ```bash grep -r . -e "NVVMMemorySpace::kGenericMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kGenericMemorySpace/NVVMMemorySpace::Generic/g" grep -r . -e "NVVMMemorySpace::kGlobalMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kGlobalMemorySpace/NVVMMemorySpace::Global/g" grep -r . -e "NVVMMemorySpace::kSharedMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kSharedMemorySpace/NVVMMemorySpace::Shared/g" grep -r . -e "NVVMMemorySpace::kConstantMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kConstantMemorySpace/NVVMMemorySpace::Constant/g" grep -r . -e "NVVMMemorySpace::kLocalMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kLocalMemorySpace/NVVMMemorySpace::Local/g" grep -r . -e "NVVMMemorySpace::kTensorMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kTensorMemorySpace/NVVMMemorySpace::Tensor/g" grep -r . -e "NVVMMemorySpace::kSharedClusterMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kSharedClusterMemorySpace/NVVMMemorySpace::SharedCluster/g" ``` NOTE: A future patch will add support for ROCDL.
fabianmcg
force-pushed
the
llvm-as-interface
branch
from
e0e70a4 to
b3116a7
Compare
There was a problem hiding this comment.
Pull Request Overview
This patch introduces the LLVMAddrSpaceAttrInterface to provide a unified interface for LLVM address space attributes, enabling better interoperability between different memory space attribute implementations. It also migrates NVVM memory space constants from a simple enum to MLIR enum attributes that implement the new interface.
Key changes:
- Adds
LLVMAddrSpaceAttrInterfacefor standardizing LLVM address space attribute access - Converts NVVM memory space constants to MLIR enum attributes with
NVVMMemorySpaceAttr - Updates all references to use the new enum class format (e.g.,
kGlobalMemorySpace→Global)
Reviewed Changes
Copilot reviewed 16 out of 16 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| mlir/include/mlir/Dialect/LLVMIR/LLVMInterfaces.td | Defines the new LLVMAddrSpaceAttrInterface |
| mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | Adds NVVM memory space enum and attribute definitions |
| mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h | Removes old enum, adds utility comparison operators |
| mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | Implements NVVMMemorySpaceAttr methods and updates enum references |
| mlir/lib/Target/LLVMIR/TypeToLLVM.cpp | Updates to use the new interface for pointer translation |
| Various conversion files | Updates NVVM memory space constant references to new enum format |
| Test files | Adds test coverage for the new NVVM memory space attributes |
| } | ||
| llvm_unreachable("unknown address space enum value"); | ||
| return 0; | ||
| return static_cast<unsigned>(NVVM::NVVMMemorySpace::Generic); |
There was a problem hiding this comment.
This change modifies the control flow by adding a return statement in the default case. This could affect performance profiling data by changing the execution path from an unreachable code path to a valid return.
Copilot generated this review using guidance from repository custom instructions.
| } | ||
| llvm_unreachable("unknown address space enum value"); | ||
| return 0; | ||
| return static_cast<unsigned>(NVVM::NVVMMemorySpace::Generic); |
There was a problem hiding this comment.
This change modifies the control flow by adding a return statement in the default case. This could affect performance profiling data by changing the execution path from an unreachable code path to a valid return.
Copilot generated this review using guidance from repository custom instructions.
| } | ||
| llvm_unreachable("unknown address space enum value"); | ||
| return 0; | ||
| return static_cast<unsigned>(NVVM::NVVMMemorySpace::Generic); |
There was a problem hiding this comment.
This change modifies the control flow by adding a return statement in the default case. This could affect performance profiling data by changing the execution path from an unreachable code path to a valid return.
Copilot generated this review using guidance from repository custom instructions.
| } | ||
| llvm_unreachable("unknown address space enum value"); | ||
| return 0; | ||
| return static_cast<unsigned>(NVVM::NVVMMemorySpace::Generic); |
There was a problem hiding this comment.
This change modifies the control flow by adding a return statement in the default case. This could affect performance profiling data by changing the execution path from an unreachable code path to a valid return.
Copilot generated this review using guidance from repository custom instructions.
|
@llvm/pr-subscribers-flang-codegen @llvm/pr-subscribers-mlir Author: Fabian Mora (fabianmcg) ChangesThis patch introduces the To test this interface, this patch also adds:
To convert the memory space enum to the new enum class use: grep -r . -e "NVVMMemorySpace::kGenericMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kGenericMemorySpace/NVVMMemorySpace::Generic/g"
grep -r . -e "NVVMMemorySpace::kGlobalMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kGlobalMemorySpace/NVVMMemorySpace::Global/g"
grep -r . -e "NVVMMemorySpace::kSharedMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kSharedMemorySpace/NVVMMemorySpace::Shared/g"
grep -r . -e "NVVMMemorySpace::kConstantMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kConstantMemorySpace/NVVMMemorySpace::Constant/g"
grep -r . -e "NVVMMemorySpace::kLocalMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kLocalMemorySpace/NVVMMemorySpace::Local/g"
grep -r . -e "NVVMMemorySpace::kTensorMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kTensorMemorySpace/NVVMMemorySpace::Tensor/g"
grep -r . -e "NVVMMemorySpace::kSharedClusterMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kSharedClusterMemorySpace/NVVMMemorySpace::SharedCluster/g"NOTE: A future patch will add support for ROCDL, it wasn't added here to keep the patch small. Patch is 28.29 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/157339.diff 16 Files Affected:
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
index ac99b8aba073a..9a851dcd73844 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
@@ -30,6 +30,7 @@ class LLVM_Attr<string name, string attrMnemonic,
def LLVM_AddressSpaceAttr :
LLVM_Attr<"AddressSpace", "address_space", [
+ LLVM_LLVMAddrSpaceAttrInterface,
DeclareAttrInterfaceMethods<MemorySpaceAttrInterface>
]> {
let summary = "LLVM address space";
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h
index fafccf304e1b4..ce62f0751d876 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h
@@ -93,6 +93,14 @@ class TBAANodeAttr : public Attribute {
using cconv::CConv;
using linkage::Linkage;
using tailcallkind::TailCallKind;
+
+namespace detail {
+/// Checks whether the given type is an LLVM type that can be loaded or stored.
+bool isValidLoadStoreImpl(Type type, ptr::AtomicOrdering ordering,
+ std::optional<int64_t> alignment,
+ const ::mlir::DataLayout *dataLayout,
+ function_ref<InFlightDiagnostic()> emitError);
+} // namespace detail
} // namespace LLVM
} // namespace mlir
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMInterfaces.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMInterfaces.td
index 60235bcb35561..e05fb6a9bac7d 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMInterfaces.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMInterfaces.td
@@ -533,6 +533,24 @@ def LLVM_DIRecursiveTypeAttrInterface
];
}
+def LLVM_LLVMAddrSpaceAttrInterface :
+ AttrInterface<"LLVMAddrSpaceAttrInterface"> {
+ let description = [{
+ An interface for attributes that represent LLVM address spaces.
+ Implementing attributes should provide access to the address space value
+ as an unsigned integer.
+ }];
+ let cppNamespace = "::mlir::LLVM";
+ let methods = [
+ InterfaceMethod<
+ /*description=*/"Returns the address space as an unsigned integer.",
+ /*retTy=*/"unsigned",
+ /*methodName=*/"getAddressSpace",
+ /*args=*/(ins)
+ >
+ ];
+}
+
def LLVM_TargetAttrInterface
: AttrInterface<"TargetAttrInterface", [DLTIQueryInterface]> {
let description = [{
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h b/mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h
index 6137bb087c576..6bd582d66ed25 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h
@@ -19,6 +19,7 @@
#include "mlir/Dialect/LLVMIR/BasicPtxBuilderInterface.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/LLVMIR/NVVMRequiresSMTraits.h"
+#include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.h"
#include "mlir/IR/Dialect.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/Interfaces/InferIntRangeInterface.h"
@@ -30,31 +31,23 @@
namespace mlir {
namespace NVVM {
+/// Utility functions to compare NVVMMemorySpace with unsigned values.
+inline bool operator==(unsigned as, NVVMMemorySpace memSpace) {
+ return as == static_cast<unsigned>(memSpace);
+}
+inline bool operator==(NVVMMemorySpace memSpace, unsigned as) {
+ return static_cast<unsigned>(memSpace) == as;
+}
+inline bool operator!=(unsigned as, NVVMMemorySpace memSpace) {
+ return as != static_cast<unsigned>(memSpace);
+}
+inline bool operator!=(NVVMMemorySpace memSpace, unsigned as) {
+ return static_cast<unsigned>(memSpace) != as;
+}
// Shared memory has 128-bit alignment
constexpr int kSharedMemoryAlignmentBit = 128;
-/// NVVM memory space identifiers.
-enum NVVMMemorySpace {
- /// Generic memory space identifier.
- kGenericMemorySpace = 0,
- /// Global memory space identifier.
- kGlobalMemorySpace = 1,
- /// Shared memory space identifier.
- kSharedMemorySpace = 3,
- /// Constant memory space identifier.
- kConstantMemorySpace = 4,
- /// Local memory space identifier.
- kLocalMemorySpace = 5,
- /// Tensor memory space identifier.
- /// Tensor memory is available only in arch-accelerated
- /// variants from sm100 onwards.
- kTensorMemorySpace = 6,
- /// Distributed shared memory space identifier.
- /// Distributed shared memory is available only in sm90+.
- kSharedClusterMemorySpace = 7,
-};
-
/// A pair type of LLVM's Intrinsic ID and args (which are llvm values).
/// This type is returned by the getIntrinsicIDAndArgs() methods.
using IDArgPair =
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 854b4d26b4368..70ab1df876d35 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -17,6 +17,7 @@ include "mlir/IR/EnumAttr.td"
include "mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td"
include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
include "mlir/Dialect/LLVMIR/NVVMRequiresSMTraits.td"
+include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
include "mlir/Dialect/LLVMIR/BasicPtxBuilderInterface.td"
include "mlir/Interfaces/InferIntRangeInterface.td"
@@ -192,6 +193,40 @@ def CacheEvictionPriorityAttr : EnumAttr<NVVM_Dialect, CacheEvictionPriority,
let assemblyFormat = "$value";
}
+// Memory Space enum definitions
+/// Generic memory space identifier.
+def MemSpaceGeneric : I32EnumCase<"Generic", 0, "generic">;
+/// Global memory space identifier.
+def MemSpaceGlobal : I32EnumCase<"Global", 1, "global">;
+/// Shared memory space identifier.
+def MemSpaceShared : I32EnumCase<"Shared", 3, "shared">;
+/// Constant memory space identifier.
+def MemSpaceConstant : I32EnumCase<"Constant", 4, "constant">;
+/// Local memory space identifier.
+def MemSpaceLocal : I32EnumCase<"Local", 5, "local">;
+/// Tensor memory space identifier.
+/// Tensor memory is available only in arch-accelerated
+/// variants from sm100 onwards.
+def MemSpaceTensor : I32EnumCase<"Tensor", 6, "tensor">;
+/// Distributed shared memory space identifier.
+/// Distributed shared memory is available only in sm90+.
+def MemSpaceSharedCluster : I32EnumCase<"SharedCluster", 7, "shared_cluster">;
+
+def NVVMMemorySpace : I32Enum<"NVVMMemorySpace", "NVVM Memory Space",
+ [MemSpaceGeneric, MemSpaceGlobal, MemSpaceShared,
+ MemSpaceConstant, MemSpaceLocal, MemSpaceTensor,
+ MemSpaceSharedCluster]> {
+ let cppNamespace = "::mlir::NVVM";
+}
+
+def NVVMMemorySpaceAttr :
+ EnumAttr<NVVM_Dialect, NVVMMemorySpace, "memory_space", [
+ DeclareAttrInterfaceMethods<LLVM_LLVMAddrSpaceAttrInterface>,
+ DeclareAttrInterfaceMethods<MemorySpaceAttrInterface>
+ ]> {
+ let assemblyFormat = "`<` $value `>`";
+}
+
//===----------------------------------------------------------------------===//
// NVVM intrinsic operations
//===----------------------------------------------------------------------===//
@@ -3592,7 +3627,7 @@ def NVVM_MapaOp: NVVM_Op<"mapa",
string llvmBuilder = [{
int addrSpace = llvm::cast<LLVMPointerType>(op.getA().getType()).getAddressSpace();
- bool isSharedMemory = addrSpace == NVVM::NVVMMemorySpace::kSharedMemorySpace;
+ bool isSharedMemory = addrSpace == static_cast<int> (NVVM::NVVMMemorySpace::Shared);
auto intId = isSharedMemory? llvm::Intrinsic::nvvm_mapa_shared_cluster : llvm::Intrinsic::nvvm_mapa;
$res = createIntrinsicCall(builder, intId, {$a, $b});
diff --git a/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp b/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
index 93e370d91e6b9..76a7e0f3831a2 100644
--- a/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
+++ b/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
@@ -451,16 +451,14 @@ void mlir::configureGpuToNVVMTypeConverter(LLVMTypeConverter &converter) {
converter, [](gpu::AddressSpace space) -> unsigned {
switch (space) {
case gpu::AddressSpace::Global:
- return static_cast<unsigned>(
- NVVM::NVVMMemorySpace::kGlobalMemorySpace);
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Global);
case gpu::AddressSpace::Workgroup:
- return static_cast<unsigned>(
- NVVM::NVVMMemorySpace::kSharedMemorySpace);
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Shared);
case gpu::AddressSpace::Private:
return 0;
}
llvm_unreachable("unknown address space enum value");
- return 0;
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Generic);
});
// Lowering for MMAMatrixType.
converter.addConversion([&](gpu::MMAMatrixType type) -> Type {
@@ -648,7 +646,7 @@ void mlir::populateGpuToNVVMConversionPatterns(
GPUFuncOpLoweringOptions{
/*allocaAddrSpace=*/0,
/*workgroupAddrSpace=*/
- static_cast<unsigned>(NVVM::NVVMMemorySpace::kSharedMemorySpace),
+ static_cast<unsigned>(NVVM::NVVMMemorySpace::Shared),
StringAttr::get(&converter.getContext(),
NVVM::NVVMDialect::getKernelFuncAttrName()),
StringAttr::get(&converter.getContext(),
diff --git a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
index 37d12bad298df..b7e3491117e9b 100644
--- a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
+++ b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
@@ -405,16 +405,14 @@ struct ConvertNVGPUToNVVMPass
converter, [](gpu::AddressSpace space) -> unsigned {
switch (space) {
case gpu::AddressSpace::Global:
- return static_cast<unsigned>(
- NVVM::NVVMMemorySpace::kGlobalMemorySpace);
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Global);
case gpu::AddressSpace::Workgroup:
- return static_cast<unsigned>(
- NVVM::NVVMMemorySpace::kSharedMemorySpace);
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Shared);
case gpu::AddressSpace::Private:
return 0;
}
llvm_unreachable("unknown address space enum value");
- return 0;
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Generic);
});
/// device-side async tokens cannot be materialized in nvvm. We just
/// convert them to a dummy i32 type in order to easily drop them during
@@ -677,7 +675,7 @@ struct NVGPUAsyncCopyLowering
adaptor.getSrcIndices());
// Intrinsics takes a global pointer so we need an address space cast.
auto srcPointerGlobalType = LLVM::LLVMPointerType::get(
- op->getContext(), NVVM::NVVMMemorySpace::kGlobalMemorySpace);
+ op->getContext(), static_cast<unsigned>(NVVM::NVVMMemorySpace::Global));
scrPtr = LLVM::AddrSpaceCastOp::create(b, srcPointerGlobalType, scrPtr);
int64_t dstElements = adaptor.getDstElements().getZExtValue();
int64_t sizeInBytes =
diff --git a/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp b/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp
index c766539f9d91a..2561f6606067f 100644
--- a/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp
+++ b/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp
@@ -71,16 +71,14 @@ void transform::ApplyGPUToNVVMConversionPatternsOp::populatePatterns(
llvmTypeConverter, [](AddressSpace space) -> unsigned {
switch (space) {
case AddressSpace::Global:
- return static_cast<unsigned>(
- NVVM::NVVMMemorySpace::kGlobalMemorySpace);
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Global);
case AddressSpace::Workgroup:
- return static_cast<unsigned>(
- NVVM::NVVMMemorySpace::kSharedMemorySpace);
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Shared);
case AddressSpace::Private:
return 0;
}
llvm_unreachable("unknown address space enum value");
- return 0;
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Generic);
});
// Used in GPUToNVVM/WmmaOpsToNvvm.cpp so attaching here for now.
// TODO: We should have a single to_nvvm_type_converter.
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp
index 23610fbd657fd..b8331e0068880 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp
@@ -57,10 +57,10 @@ void LLVMDialect::registerAttributes() {
//===----------------------------------------------------------------------===//
/// Checks whether the given type is an LLVM type that can be loaded or stored.
-static bool isValidLoadStoreImpl(Type type, ptr::AtomicOrdering ordering,
- std::optional<int64_t> alignment,
- const ::mlir::DataLayout *dataLayout,
- function_ref<InFlightDiagnostic()> emitError) {
+bool LLVM::detail::isValidLoadStoreImpl(
+ Type type, ptr::AtomicOrdering ordering, std::optional<int64_t> alignment,
+ const ::mlir::DataLayout *dataLayout,
+ function_ref<InFlightDiagnostic()> emitError) {
if (!isLoadableType(type)) {
if (emitError)
emitError() << "type must be LLVM type with size, but got " << type;
@@ -87,14 +87,16 @@ bool AddressSpaceAttr::isValidLoad(
Type type, ptr::AtomicOrdering ordering, std::optional<int64_t> alignment,
const ::mlir::DataLayout *dataLayout,
function_ref<InFlightDiagnostic()> emitError) const {
- return isValidLoadStoreImpl(type, ordering, alignment, dataLayout, emitError);
+ return detail::isValidLoadStoreImpl(type, ordering, alignment, dataLayout,
+ emitError);
}
bool AddressSpaceAttr::isValidStore(
Type type, ptr::AtomicOrdering ordering, std::optional<int64_t> alignment,
const ::mlir::DataLayout *dataLayout,
function_ref<InFlightDiagnostic()> emitError) const {
- return isValidLoadStoreImpl(type, ordering, alignment, dataLayout, emitError);
+ return detail::isValidLoadStoreImpl(type, ordering, alignment, dataLayout,
+ emitError);
}
bool AddressSpaceAttr::isValidAtomicOp(
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
index 2dd0132a65bc4..01a16ce2d8a7f 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
@@ -703,14 +703,14 @@ const llvm::fltSemantics &LLVMPPCFP128Type::getFloatSemantics() const {
//===----------------------------------------------------------------------===//
/// Check whether type is a compatible ptr type. These are pointer-like types
-/// with no element type, no metadata, and using the LLVM AddressSpaceAttr
-/// memory space.
+/// with no element type, no metadata, and using the LLVM
+/// LLVMAddrSpaceAttrInterface memory space.
static bool isCompatiblePtrType(Type type) {
auto ptrTy = dyn_cast<PtrLikeTypeInterface>(type);
if (!ptrTy)
return false;
return !ptrTy.hasPtrMetadata() && ptrTy.getElementType() == nullptr &&
- isa<AddressSpaceAttr>(ptrTy.getMemorySpace());
+ isa<LLVMAddrSpaceAttrInterface>(ptrTy.getMemorySpace());
}
bool mlir::LLVM::isCompatibleOuterType(Type type) {
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 376e3c3e1fcbe..03cc0a611513d 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -800,8 +800,8 @@ inferMMATypeFromMNK(NVVM::MMATypes type, NVVM::MMAFrag frag, int m, int n,
LogicalResult NVVM::WMMALoadOp::verify() {
unsigned addressSpace =
llvm::cast<LLVM::LLVMPointerType>(getPtr().getType()).getAddressSpace();
- if (addressSpace != 0 && addressSpace != NVVM::kGlobalMemorySpace &&
- addressSpace != NVVM::kSharedMemorySpace)
+ if (addressSpace != 0 && addressSpace != NVVMMemorySpace::Global &&
+ addressSpace != NVVMMemorySpace::Shared)
return emitOpError("expected source pointer in memory "
"space 0, 1, 3");
@@ -821,8 +821,8 @@ LogicalResult NVVM::WMMALoadOp::verify() {
LogicalResult NVVM::WMMAStoreOp::verify() {
unsigned addressSpace =
llvm::cast<LLVM::LLVMPointerType>(getPtr().getType()).getAddressSpace();
- if (addressSpace != 0 && addressSpace != NVVM::kGlobalMemorySpace &&
- addressSpace != NVVM::kSharedMemorySpace)
+ if (addressSpace != 0 && addressSpace != NVVMMemorySpace::Global &&
+ addressSpace != NVVMMemorySpace::Shared)
return emitOpError("expected operands to be a source pointer in memory "
"space 0, 1, 3");
@@ -1339,8 +1339,8 @@ LogicalResult NVVM::PrefetchOp::verify() {
return emitOpError("cannot specify both tensormap and cache level");
if (getTensormap()) {
- if (addressSpace != MemSpace::kGenericMemorySpace &&
- addressSpace != MemSpace::kConstantMemorySpace) {
+ if (addressSpace != MemSpace::Generic &&
+ addressSpace != MemSpace::Constant) {
return emitOpError(
"prefetch tensormap requires a generic or constant pointer");
}
@@ -1350,15 +1350,14 @@ LogicalResult NVVM::PrefetchOp::verify() {
"prefetch tensormap does not support eviction priority");
}
- if (getInParamSpace() && addressSpace != MemSpace::kGenericMemorySpace) {
+ if (getInParamSpace() && addressSpace != MemSpace::Generic) {
return emitOpError(
"in_param_space can only be specified for a generic pointer");
}
} else if (cacheLevel) {
- if (addressSpace != MemSpace::kGenericMemorySpace &&
- addressSpace != MemSpace::kGlobalMemorySpace &&
- addressSpace != MemSpace::kLocalMemorySpace) {
+ if (addressSpace != MemSpace::Generic && addressSpace != MemSpace::Global &&
+ addressSpace != MemSpace::Local) {
return emitOpError("prefetch to cache level requires a generic, global, "
"or local pointer");
}
@@ -1370,7 +1369,7 @@ LogicalResult NVVM::PrefetchOp::verify() {
"cache level is L1");
}
- if (addressSpace != MemSpace::kGenericMemorySpace) {
+ if (addressSpace != MemSpace::Generic) {
return emitOpError(
"prefetch to uniform cache requires a generic pointer");
}
@@ -1381,7 +1380,7 @@ LogicalResult NVVM::PrefetchOp::verify() {
return emitOpError(
"cache eviction priority supported only for cache level L2");
- if (addressSpace != MemSpace::kGlobalMemorySpace)
+ if (addressSpace != MemSpace::Global)
return emitOpError("cache eviction priority requires a global pointer");
if (*evictPriority != NVVM::CacheEvictionPriority::EvictNormal &&
@@ -1796,7 +1795,7 @@ Tcgen05AllocOp::getIntrinsicIDAndArgs(Operation &op,
auto curOp = cast<NVVM::Tcgen05AllocOp>(op);
unsigned as = llvm::cast<LLVM::LLVMPointerType>(curOp.getAddr().getType())
.getAddressSpace();
- bool isShared = as == NVVMMemorySpace::kSharedMemorySpace;
+ bool isShared = as == NVVMMemorySpace::Shared;
bool is2CTAMode = curOp.getGroup() == CTAGroupKind::CTA_2;
llvm::Intrinsic::ID id;
@@ -1845,7 +1844,7 @@ Tcgen05CommitOp::getIntrinsicIDAndArgs(Operation &op,
auto curOp = cast<NVVM::Tcgen05CommitOp>(op);
unsigned as = llvm::cast<LLVM::LLVMPointerType>(curOp.getAddr().getType())
.getAddressSpace();
- bool isShared = as == NVVMMemorySpace::kSharedMemorySpace;
+ bool isShared = as == NVVMMemorySpace::Shared;
bool hasMulticast = static_cast<bool>(curOp.getMulticastMask());
bool is2CTAMode = curOp.getGroup() == CTAGroupKind::CTA_2;
@@ -2051,18 +2050,18 @@ PrefetchOp::getIntrinsicIDAndArgs(NVVM::PrefetchOp &op,
}
}
- switch (addressSpace) {
- case MemSpace::kGenericMemorySpace:
+ switch (static_cast<MemSpace>(addressSpace)) {
+ case MemSpace::Generic:
return *cacheLevel == CacheLevel::L1
? NVVM::IDArgPair({llvm::Intrinsic::nvvm_prefetch_L1, args})
: NVVM::IDArgPair({llvm::Intrinsic::nvvm_prefetch_L2, args});
- case MemSpace::kGlobalMemorySpace:
+ case MemSpace::Global:
return *cacheLevel == CacheLevel::L1
? NVVM::IDArg...
[truncated]
|
|
@llvm/pr-subscribers-mlir-gpu Author: Fabian Mora (fabianmcg) ChangesThis patch introduces the To test this interface, this patch also adds:
To convert the memory space enum to the new enum class use: grep -r . -e "NVVMMemorySpace::kGenericMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kGenericMemorySpace/NVVMMemorySpace::Generic/g"
grep -r . -e "NVVMMemorySpace::kGlobalMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kGlobalMemorySpace/NVVMMemorySpace::Global/g"
grep -r . -e "NVVMMemorySpace::kSharedMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kSharedMemorySpace/NVVMMemorySpace::Shared/g"
grep -r . -e "NVVMMemorySpace::kConstantMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kConstantMemorySpace/NVVMMemorySpace::Constant/g"
grep -r . -e "NVVMMemorySpace::kLocalMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kLocalMemorySpace/NVVMMemorySpace::Local/g"
grep -r . -e "NVVMMemorySpace::kTensorMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kTensorMemorySpace/NVVMMemorySpace::Tensor/g"
grep -r . -e "NVVMMemorySpace::kSharedClusterMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kSharedClusterMemorySpace/NVVMMemorySpace::SharedCluster/g"NOTE: A future patch will add support for ROCDL, it wasn't added here to keep the patch small. Patch is 28.29 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/157339.diff 16 Files Affected:
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
index ac99b8aba073a..9a851dcd73844 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
@@ -30,6 +30,7 @@ class LLVM_Attr<string name, string attrMnemonic,
def LLVM_AddressSpaceAttr :
LLVM_Attr<"AddressSpace", "address_space", [
+ LLVM_LLVMAddrSpaceAttrInterface,
DeclareAttrInterfaceMethods<MemorySpaceAttrInterface>
]> {
let summary = "LLVM address space";
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h
index fafccf304e1b4..ce62f0751d876 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h
@@ -93,6 +93,14 @@ class TBAANodeAttr : public Attribute {
using cconv::CConv;
using linkage::Linkage;
using tailcallkind::TailCallKind;
+
+namespace detail {
+/// Checks whether the given type is an LLVM type that can be loaded or stored.
+bool isValidLoadStoreImpl(Type type, ptr::AtomicOrdering ordering,
+ std::optional<int64_t> alignment,
+ const ::mlir::DataLayout *dataLayout,
+ function_ref<InFlightDiagnostic()> emitError);
+} // namespace detail
} // namespace LLVM
} // namespace mlir
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMInterfaces.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMInterfaces.td
index 60235bcb35561..e05fb6a9bac7d 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMInterfaces.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMInterfaces.td
@@ -533,6 +533,24 @@ def LLVM_DIRecursiveTypeAttrInterface
];
}
+def LLVM_LLVMAddrSpaceAttrInterface :
+ AttrInterface<"LLVMAddrSpaceAttrInterface"> {
+ let description = [{
+ An interface for attributes that represent LLVM address spaces.
+ Implementing attributes should provide access to the address space value
+ as an unsigned integer.
+ }];
+ let cppNamespace = "::mlir::LLVM";
+ let methods = [
+ InterfaceMethod<
+ /*description=*/"Returns the address space as an unsigned integer.",
+ /*retTy=*/"unsigned",
+ /*methodName=*/"getAddressSpace",
+ /*args=*/(ins)
+ >
+ ];
+}
+
def LLVM_TargetAttrInterface
: AttrInterface<"TargetAttrInterface", [DLTIQueryInterface]> {
let description = [{
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h b/mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h
index 6137bb087c576..6bd582d66ed25 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h
@@ -19,6 +19,7 @@
#include "mlir/Dialect/LLVMIR/BasicPtxBuilderInterface.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/LLVMIR/NVVMRequiresSMTraits.h"
+#include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.h"
#include "mlir/IR/Dialect.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/Interfaces/InferIntRangeInterface.h"
@@ -30,31 +31,23 @@
namespace mlir {
namespace NVVM {
+/// Utility functions to compare NVVMMemorySpace with unsigned values.
+inline bool operator==(unsigned as, NVVMMemorySpace memSpace) {
+ return as == static_cast<unsigned>(memSpace);
+}
+inline bool operator==(NVVMMemorySpace memSpace, unsigned as) {
+ return static_cast<unsigned>(memSpace) == as;
+}
+inline bool operator!=(unsigned as, NVVMMemorySpace memSpace) {
+ return as != static_cast<unsigned>(memSpace);
+}
+inline bool operator!=(NVVMMemorySpace memSpace, unsigned as) {
+ return static_cast<unsigned>(memSpace) != as;
+}
// Shared memory has 128-bit alignment
constexpr int kSharedMemoryAlignmentBit = 128;
-/// NVVM memory space identifiers.
-enum NVVMMemorySpace {
- /// Generic memory space identifier.
- kGenericMemorySpace = 0,
- /// Global memory space identifier.
- kGlobalMemorySpace = 1,
- /// Shared memory space identifier.
- kSharedMemorySpace = 3,
- /// Constant memory space identifier.
- kConstantMemorySpace = 4,
- /// Local memory space identifier.
- kLocalMemorySpace = 5,
- /// Tensor memory space identifier.
- /// Tensor memory is available only in arch-accelerated
- /// variants from sm100 onwards.
- kTensorMemorySpace = 6,
- /// Distributed shared memory space identifier.
- /// Distributed shared memory is available only in sm90+.
- kSharedClusterMemorySpace = 7,
-};
-
/// A pair type of LLVM's Intrinsic ID and args (which are llvm values).
/// This type is returned by the getIntrinsicIDAndArgs() methods.
using IDArgPair =
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 854b4d26b4368..70ab1df876d35 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -17,6 +17,7 @@ include "mlir/IR/EnumAttr.td"
include "mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td"
include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
include "mlir/Dialect/LLVMIR/NVVMRequiresSMTraits.td"
+include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
include "mlir/Dialect/LLVMIR/BasicPtxBuilderInterface.td"
include "mlir/Interfaces/InferIntRangeInterface.td"
@@ -192,6 +193,40 @@ def CacheEvictionPriorityAttr : EnumAttr<NVVM_Dialect, CacheEvictionPriority,
let assemblyFormat = "$value";
}
+// Memory Space enum definitions
+/// Generic memory space identifier.
+def MemSpaceGeneric : I32EnumCase<"Generic", 0, "generic">;
+/// Global memory space identifier.
+def MemSpaceGlobal : I32EnumCase<"Global", 1, "global">;
+/// Shared memory space identifier.
+def MemSpaceShared : I32EnumCase<"Shared", 3, "shared">;
+/// Constant memory space identifier.
+def MemSpaceConstant : I32EnumCase<"Constant", 4, "constant">;
+/// Local memory space identifier.
+def MemSpaceLocal : I32EnumCase<"Local", 5, "local">;
+/// Tensor memory space identifier.
+/// Tensor memory is available only in arch-accelerated
+/// variants from sm100 onwards.
+def MemSpaceTensor : I32EnumCase<"Tensor", 6, "tensor">;
+/// Distributed shared memory space identifier.
+/// Distributed shared memory is available only in sm90+.
+def MemSpaceSharedCluster : I32EnumCase<"SharedCluster", 7, "shared_cluster">;
+
+def NVVMMemorySpace : I32Enum<"NVVMMemorySpace", "NVVM Memory Space",
+ [MemSpaceGeneric, MemSpaceGlobal, MemSpaceShared,
+ MemSpaceConstant, MemSpaceLocal, MemSpaceTensor,
+ MemSpaceSharedCluster]> {
+ let cppNamespace = "::mlir::NVVM";
+}
+
+def NVVMMemorySpaceAttr :
+ EnumAttr<NVVM_Dialect, NVVMMemorySpace, "memory_space", [
+ DeclareAttrInterfaceMethods<LLVM_LLVMAddrSpaceAttrInterface>,
+ DeclareAttrInterfaceMethods<MemorySpaceAttrInterface>
+ ]> {
+ let assemblyFormat = "`<` $value `>`";
+}
+
//===----------------------------------------------------------------------===//
// NVVM intrinsic operations
//===----------------------------------------------------------------------===//
@@ -3592,7 +3627,7 @@ def NVVM_MapaOp: NVVM_Op<"mapa",
string llvmBuilder = [{
int addrSpace = llvm::cast<LLVMPointerType>(op.getA().getType()).getAddressSpace();
- bool isSharedMemory = addrSpace == NVVM::NVVMMemorySpace::kSharedMemorySpace;
+ bool isSharedMemory = addrSpace == static_cast<int> (NVVM::NVVMMemorySpace::Shared);
auto intId = isSharedMemory? llvm::Intrinsic::nvvm_mapa_shared_cluster : llvm::Intrinsic::nvvm_mapa;
$res = createIntrinsicCall(builder, intId, {$a, $b});
diff --git a/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp b/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
index 93e370d91e6b9..76a7e0f3831a2 100644
--- a/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
+++ b/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
@@ -451,16 +451,14 @@ void mlir::configureGpuToNVVMTypeConverter(LLVMTypeConverter &converter) {
converter, [](gpu::AddressSpace space) -> unsigned {
switch (space) {
case gpu::AddressSpace::Global:
- return static_cast<unsigned>(
- NVVM::NVVMMemorySpace::kGlobalMemorySpace);
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Global);
case gpu::AddressSpace::Workgroup:
- return static_cast<unsigned>(
- NVVM::NVVMMemorySpace::kSharedMemorySpace);
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Shared);
case gpu::AddressSpace::Private:
return 0;
}
llvm_unreachable("unknown address space enum value");
- return 0;
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Generic);
});
// Lowering for MMAMatrixType.
converter.addConversion([&](gpu::MMAMatrixType type) -> Type {
@@ -648,7 +646,7 @@ void mlir::populateGpuToNVVMConversionPatterns(
GPUFuncOpLoweringOptions{
/*allocaAddrSpace=*/0,
/*workgroupAddrSpace=*/
- static_cast<unsigned>(NVVM::NVVMMemorySpace::kSharedMemorySpace),
+ static_cast<unsigned>(NVVM::NVVMMemorySpace::Shared),
StringAttr::get(&converter.getContext(),
NVVM::NVVMDialect::getKernelFuncAttrName()),
StringAttr::get(&converter.getContext(),
diff --git a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
index 37d12bad298df..b7e3491117e9b 100644
--- a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
+++ b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
@@ -405,16 +405,14 @@ struct ConvertNVGPUToNVVMPass
converter, [](gpu::AddressSpace space) -> unsigned {
switch (space) {
case gpu::AddressSpace::Global:
- return static_cast<unsigned>(
- NVVM::NVVMMemorySpace::kGlobalMemorySpace);
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Global);
case gpu::AddressSpace::Workgroup:
- return static_cast<unsigned>(
- NVVM::NVVMMemorySpace::kSharedMemorySpace);
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Shared);
case gpu::AddressSpace::Private:
return 0;
}
llvm_unreachable("unknown address space enum value");
- return 0;
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Generic);
});
/// device-side async tokens cannot be materialized in nvvm. We just
/// convert them to a dummy i32 type in order to easily drop them during
@@ -677,7 +675,7 @@ struct NVGPUAsyncCopyLowering
adaptor.getSrcIndices());
// Intrinsics takes a global pointer so we need an address space cast.
auto srcPointerGlobalType = LLVM::LLVMPointerType::get(
- op->getContext(), NVVM::NVVMMemorySpace::kGlobalMemorySpace);
+ op->getContext(), static_cast<unsigned>(NVVM::NVVMMemorySpace::Global));
scrPtr = LLVM::AddrSpaceCastOp::create(b, srcPointerGlobalType, scrPtr);
int64_t dstElements = adaptor.getDstElements().getZExtValue();
int64_t sizeInBytes =
diff --git a/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp b/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp
index c766539f9d91a..2561f6606067f 100644
--- a/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp
+++ b/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp
@@ -71,16 +71,14 @@ void transform::ApplyGPUToNVVMConversionPatternsOp::populatePatterns(
llvmTypeConverter, [](AddressSpace space) -> unsigned {
switch (space) {
case AddressSpace::Global:
- return static_cast<unsigned>(
- NVVM::NVVMMemorySpace::kGlobalMemorySpace);
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Global);
case AddressSpace::Workgroup:
- return static_cast<unsigned>(
- NVVM::NVVMMemorySpace::kSharedMemorySpace);
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Shared);
case AddressSpace::Private:
return 0;
}
llvm_unreachable("unknown address space enum value");
- return 0;
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Generic);
});
// Used in GPUToNVVM/WmmaOpsToNvvm.cpp so attaching here for now.
// TODO: We should have a single to_nvvm_type_converter.
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp
index 23610fbd657fd..b8331e0068880 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp
@@ -57,10 +57,10 @@ void LLVMDialect::registerAttributes() {
//===----------------------------------------------------------------------===//
/// Checks whether the given type is an LLVM type that can be loaded or stored.
-static bool isValidLoadStoreImpl(Type type, ptr::AtomicOrdering ordering,
- std::optional<int64_t> alignment,
- const ::mlir::DataLayout *dataLayout,
- function_ref<InFlightDiagnostic()> emitError) {
+bool LLVM::detail::isValidLoadStoreImpl(
+ Type type, ptr::AtomicOrdering ordering, std::optional<int64_t> alignment,
+ const ::mlir::DataLayout *dataLayout,
+ function_ref<InFlightDiagnostic()> emitError) {
if (!isLoadableType(type)) {
if (emitError)
emitError() << "type must be LLVM type with size, but got " << type;
@@ -87,14 +87,16 @@ bool AddressSpaceAttr::isValidLoad(
Type type, ptr::AtomicOrdering ordering, std::optional<int64_t> alignment,
const ::mlir::DataLayout *dataLayout,
function_ref<InFlightDiagnostic()> emitError) const {
- return isValidLoadStoreImpl(type, ordering, alignment, dataLayout, emitError);
+ return detail::isValidLoadStoreImpl(type, ordering, alignment, dataLayout,
+ emitError);
}
bool AddressSpaceAttr::isValidStore(
Type type, ptr::AtomicOrdering ordering, std::optional<int64_t> alignment,
const ::mlir::DataLayout *dataLayout,
function_ref<InFlightDiagnostic()> emitError) const {
- return isValidLoadStoreImpl(type, ordering, alignment, dataLayout, emitError);
+ return detail::isValidLoadStoreImpl(type, ordering, alignment, dataLayout,
+ emitError);
}
bool AddressSpaceAttr::isValidAtomicOp(
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
index 2dd0132a65bc4..01a16ce2d8a7f 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
@@ -703,14 +703,14 @@ const llvm::fltSemantics &LLVMPPCFP128Type::getFloatSemantics() const {
//===----------------------------------------------------------------------===//
/// Check whether type is a compatible ptr type. These are pointer-like types
-/// with no element type, no metadata, and using the LLVM AddressSpaceAttr
-/// memory space.
+/// with no element type, no metadata, and using the LLVM
+/// LLVMAddrSpaceAttrInterface memory space.
static bool isCompatiblePtrType(Type type) {
auto ptrTy = dyn_cast<PtrLikeTypeInterface>(type);
if (!ptrTy)
return false;
return !ptrTy.hasPtrMetadata() && ptrTy.getElementType() == nullptr &&
- isa<AddressSpaceAttr>(ptrTy.getMemorySpace());
+ isa<LLVMAddrSpaceAttrInterface>(ptrTy.getMemorySpace());
}
bool mlir::LLVM::isCompatibleOuterType(Type type) {
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 376e3c3e1fcbe..03cc0a611513d 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -800,8 +800,8 @@ inferMMATypeFromMNK(NVVM::MMATypes type, NVVM::MMAFrag frag, int m, int n,
LogicalResult NVVM::WMMALoadOp::verify() {
unsigned addressSpace =
llvm::cast<LLVM::LLVMPointerType>(getPtr().getType()).getAddressSpace();
- if (addressSpace != 0 && addressSpace != NVVM::kGlobalMemorySpace &&
- addressSpace != NVVM::kSharedMemorySpace)
+ if (addressSpace != 0 && addressSpace != NVVMMemorySpace::Global &&
+ addressSpace != NVVMMemorySpace::Shared)
return emitOpError("expected source pointer in memory "
"space 0, 1, 3");
@@ -821,8 +821,8 @@ LogicalResult NVVM::WMMALoadOp::verify() {
LogicalResult NVVM::WMMAStoreOp::verify() {
unsigned addressSpace =
llvm::cast<LLVM::LLVMPointerType>(getPtr().getType()).getAddressSpace();
- if (addressSpace != 0 && addressSpace != NVVM::kGlobalMemorySpace &&
- addressSpace != NVVM::kSharedMemorySpace)
+ if (addressSpace != 0 && addressSpace != NVVMMemorySpace::Global &&
+ addressSpace != NVVMMemorySpace::Shared)
return emitOpError("expected operands to be a source pointer in memory "
"space 0, 1, 3");
@@ -1339,8 +1339,8 @@ LogicalResult NVVM::PrefetchOp::verify() {
return emitOpError("cannot specify both tensormap and cache level");
if (getTensormap()) {
- if (addressSpace != MemSpace::kGenericMemorySpace &&
- addressSpace != MemSpace::kConstantMemorySpace) {
+ if (addressSpace != MemSpace::Generic &&
+ addressSpace != MemSpace::Constant) {
return emitOpError(
"prefetch tensormap requires a generic or constant pointer");
}
@@ -1350,15 +1350,14 @@ LogicalResult NVVM::PrefetchOp::verify() {
"prefetch tensormap does not support eviction priority");
}
- if (getInParamSpace() && addressSpace != MemSpace::kGenericMemorySpace) {
+ if (getInParamSpace() && addressSpace != MemSpace::Generic) {
return emitOpError(
"in_param_space can only be specified for a generic pointer");
}
} else if (cacheLevel) {
- if (addressSpace != MemSpace::kGenericMemorySpace &&
- addressSpace != MemSpace::kGlobalMemorySpace &&
- addressSpace != MemSpace::kLocalMemorySpace) {
+ if (addressSpace != MemSpace::Generic && addressSpace != MemSpace::Global &&
+ addressSpace != MemSpace::Local) {
return emitOpError("prefetch to cache level requires a generic, global, "
"or local pointer");
}
@@ -1370,7 +1369,7 @@ LogicalResult NVVM::PrefetchOp::verify() {
"cache level is L1");
}
- if (addressSpace != MemSpace::kGenericMemorySpace) {
+ if (addressSpace != MemSpace::Generic) {
return emitOpError(
"prefetch to uniform cache requires a generic pointer");
}
@@ -1381,7 +1380,7 @@ LogicalResult NVVM::PrefetchOp::verify() {
return emitOpError(
"cache eviction priority supported only for cache level L2");
- if (addressSpace != MemSpace::kGlobalMemorySpace)
+ if (addressSpace != MemSpace::Global)
return emitOpError("cache eviction priority requires a global pointer");
if (*evictPriority != NVVM::CacheEvictionPriority::EvictNormal &&
@@ -1796,7 +1795,7 @@ Tcgen05AllocOp::getIntrinsicIDAndArgs(Operation &op,
auto curOp = cast<NVVM::Tcgen05AllocOp>(op);
unsigned as = llvm::cast<LLVM::LLVMPointerType>(curOp.getAddr().getType())
.getAddressSpace();
- bool isShared = as == NVVMMemorySpace::kSharedMemorySpace;
+ bool isShared = as == NVVMMemorySpace::Shared;
bool is2CTAMode = curOp.getGroup() == CTAGroupKind::CTA_2;
llvm::Intrinsic::ID id;
@@ -1845,7 +1844,7 @@ Tcgen05CommitOp::getIntrinsicIDAndArgs(Operation &op,
auto curOp = cast<NVVM::Tcgen05CommitOp>(op);
unsigned as = llvm::cast<LLVM::LLVMPointerType>(curOp.getAddr().getType())
.getAddressSpace();
- bool isShared = as == NVVMMemorySpace::kSharedMemorySpace;
+ bool isShared = as == NVVMMemorySpace::Shared;
bool hasMulticast = static_cast<bool>(curOp.getMulticastMask());
bool is2CTAMode = curOp.getGroup() == CTAGroupKind::CTA_2;
@@ -2051,18 +2050,18 @@ PrefetchOp::getIntrinsicIDAndArgs(NVVM::PrefetchOp &op,
}
}
- switch (addressSpace) {
- case MemSpace::kGenericMemorySpace:
+ switch (static_cast<MemSpace>(addressSpace)) {
+ case MemSpace::Generic:
return *cacheLevel == CacheLevel::L1
? NVVM::IDArgPair({llvm::Intrinsic::nvvm_prefetch_L1, args})
: NVVM::IDArgPair({llvm::Intrinsic::nvvm_prefetch_L2, args});
- case MemSpace::kGlobalMemorySpace:
+ case MemSpace::Global:
return *cacheLevel == CacheLevel::L1
? NVVM::IDArg...
[truncated]
|
|
@llvm/pr-subscribers-mlir-llvm Author: Fabian Mora (fabianmcg) ChangesThis patch introduces the To test this interface, this patch also adds:
To convert the memory space enum to the new enum class use: grep -r . -e "NVVMMemorySpace::kGenericMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kGenericMemorySpace/NVVMMemorySpace::Generic/g"
grep -r . -e "NVVMMemorySpace::kGlobalMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kGlobalMemorySpace/NVVMMemorySpace::Global/g"
grep -r . -e "NVVMMemorySpace::kSharedMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kSharedMemorySpace/NVVMMemorySpace::Shared/g"
grep -r . -e "NVVMMemorySpace::kConstantMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kConstantMemorySpace/NVVMMemorySpace::Constant/g"
grep -r . -e "NVVMMemorySpace::kLocalMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kLocalMemorySpace/NVVMMemorySpace::Local/g"
grep -r . -e "NVVMMemorySpace::kTensorMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kTensorMemorySpace/NVVMMemorySpace::Tensor/g"
grep -r . -e "NVVMMemorySpace::kSharedClusterMemorySpace" -l | xargs sed -i -e "s/NVVMMemorySpace::kSharedClusterMemorySpace/NVVMMemorySpace::SharedCluster/g"NOTE: A future patch will add support for ROCDL, it wasn't added here to keep the patch small. Patch is 28.29 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/157339.diff 16 Files Affected:
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
index ac99b8aba073a..9a851dcd73844 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
@@ -30,6 +30,7 @@ class LLVM_Attr<string name, string attrMnemonic,
def LLVM_AddressSpaceAttr :
LLVM_Attr<"AddressSpace", "address_space", [
+ LLVM_LLVMAddrSpaceAttrInterface,
DeclareAttrInterfaceMethods<MemorySpaceAttrInterface>
]> {
let summary = "LLVM address space";
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h
index fafccf304e1b4..ce62f0751d876 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h
@@ -93,6 +93,14 @@ class TBAANodeAttr : public Attribute {
using cconv::CConv;
using linkage::Linkage;
using tailcallkind::TailCallKind;
+
+namespace detail {
+/// Checks whether the given type is an LLVM type that can be loaded or stored.
+bool isValidLoadStoreImpl(Type type, ptr::AtomicOrdering ordering,
+ std::optional<int64_t> alignment,
+ const ::mlir::DataLayout *dataLayout,
+ function_ref<InFlightDiagnostic()> emitError);
+} // namespace detail
} // namespace LLVM
} // namespace mlir
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMInterfaces.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMInterfaces.td
index 60235bcb35561..e05fb6a9bac7d 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMInterfaces.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMInterfaces.td
@@ -533,6 +533,24 @@ def LLVM_DIRecursiveTypeAttrInterface
];
}
+def LLVM_LLVMAddrSpaceAttrInterface :
+ AttrInterface<"LLVMAddrSpaceAttrInterface"> {
+ let description = [{
+ An interface for attributes that represent LLVM address spaces.
+ Implementing attributes should provide access to the address space value
+ as an unsigned integer.
+ }];
+ let cppNamespace = "::mlir::LLVM";
+ let methods = [
+ InterfaceMethod<
+ /*description=*/"Returns the address space as an unsigned integer.",
+ /*retTy=*/"unsigned",
+ /*methodName=*/"getAddressSpace",
+ /*args=*/(ins)
+ >
+ ];
+}
+
def LLVM_TargetAttrInterface
: AttrInterface<"TargetAttrInterface", [DLTIQueryInterface]> {
let description = [{
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h b/mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h
index 6137bb087c576..6bd582d66ed25 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h
@@ -19,6 +19,7 @@
#include "mlir/Dialect/LLVMIR/BasicPtxBuilderInterface.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/LLVMIR/NVVMRequiresSMTraits.h"
+#include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.h"
#include "mlir/IR/Dialect.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/Interfaces/InferIntRangeInterface.h"
@@ -30,31 +31,23 @@
namespace mlir {
namespace NVVM {
+/// Utility functions to compare NVVMMemorySpace with unsigned values.
+inline bool operator==(unsigned as, NVVMMemorySpace memSpace) {
+ return as == static_cast<unsigned>(memSpace);
+}
+inline bool operator==(NVVMMemorySpace memSpace, unsigned as) {
+ return static_cast<unsigned>(memSpace) == as;
+}
+inline bool operator!=(unsigned as, NVVMMemorySpace memSpace) {
+ return as != static_cast<unsigned>(memSpace);
+}
+inline bool operator!=(NVVMMemorySpace memSpace, unsigned as) {
+ return static_cast<unsigned>(memSpace) != as;
+}
// Shared memory has 128-bit alignment
constexpr int kSharedMemoryAlignmentBit = 128;
-/// NVVM memory space identifiers.
-enum NVVMMemorySpace {
- /// Generic memory space identifier.
- kGenericMemorySpace = 0,
- /// Global memory space identifier.
- kGlobalMemorySpace = 1,
- /// Shared memory space identifier.
- kSharedMemorySpace = 3,
- /// Constant memory space identifier.
- kConstantMemorySpace = 4,
- /// Local memory space identifier.
- kLocalMemorySpace = 5,
- /// Tensor memory space identifier.
- /// Tensor memory is available only in arch-accelerated
- /// variants from sm100 onwards.
- kTensorMemorySpace = 6,
- /// Distributed shared memory space identifier.
- /// Distributed shared memory is available only in sm90+.
- kSharedClusterMemorySpace = 7,
-};
-
/// A pair type of LLVM's Intrinsic ID and args (which are llvm values).
/// This type is returned by the getIntrinsicIDAndArgs() methods.
using IDArgPair =
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 854b4d26b4368..70ab1df876d35 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -17,6 +17,7 @@ include "mlir/IR/EnumAttr.td"
include "mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td"
include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
include "mlir/Dialect/LLVMIR/NVVMRequiresSMTraits.td"
+include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
include "mlir/Dialect/LLVMIR/BasicPtxBuilderInterface.td"
include "mlir/Interfaces/InferIntRangeInterface.td"
@@ -192,6 +193,40 @@ def CacheEvictionPriorityAttr : EnumAttr<NVVM_Dialect, CacheEvictionPriority,
let assemblyFormat = "$value";
}
+// Memory Space enum definitions
+/// Generic memory space identifier.
+def MemSpaceGeneric : I32EnumCase<"Generic", 0, "generic">;
+/// Global memory space identifier.
+def MemSpaceGlobal : I32EnumCase<"Global", 1, "global">;
+/// Shared memory space identifier.
+def MemSpaceShared : I32EnumCase<"Shared", 3, "shared">;
+/// Constant memory space identifier.
+def MemSpaceConstant : I32EnumCase<"Constant", 4, "constant">;
+/// Local memory space identifier.
+def MemSpaceLocal : I32EnumCase<"Local", 5, "local">;
+/// Tensor memory space identifier.
+/// Tensor memory is available only in arch-accelerated
+/// variants from sm100 onwards.
+def MemSpaceTensor : I32EnumCase<"Tensor", 6, "tensor">;
+/// Distributed shared memory space identifier.
+/// Distributed shared memory is available only in sm90+.
+def MemSpaceSharedCluster : I32EnumCase<"SharedCluster", 7, "shared_cluster">;
+
+def NVVMMemorySpace : I32Enum<"NVVMMemorySpace", "NVVM Memory Space",
+ [MemSpaceGeneric, MemSpaceGlobal, MemSpaceShared,
+ MemSpaceConstant, MemSpaceLocal, MemSpaceTensor,
+ MemSpaceSharedCluster]> {
+ let cppNamespace = "::mlir::NVVM";
+}
+
+def NVVMMemorySpaceAttr :
+ EnumAttr<NVVM_Dialect, NVVMMemorySpace, "memory_space", [
+ DeclareAttrInterfaceMethods<LLVM_LLVMAddrSpaceAttrInterface>,
+ DeclareAttrInterfaceMethods<MemorySpaceAttrInterface>
+ ]> {
+ let assemblyFormat = "`<` $value `>`";
+}
+
//===----------------------------------------------------------------------===//
// NVVM intrinsic operations
//===----------------------------------------------------------------------===//
@@ -3592,7 +3627,7 @@ def NVVM_MapaOp: NVVM_Op<"mapa",
string llvmBuilder = [{
int addrSpace = llvm::cast<LLVMPointerType>(op.getA().getType()).getAddressSpace();
- bool isSharedMemory = addrSpace == NVVM::NVVMMemorySpace::kSharedMemorySpace;
+ bool isSharedMemory = addrSpace == static_cast<int> (NVVM::NVVMMemorySpace::Shared);
auto intId = isSharedMemory? llvm::Intrinsic::nvvm_mapa_shared_cluster : llvm::Intrinsic::nvvm_mapa;
$res = createIntrinsicCall(builder, intId, {$a, $b});
diff --git a/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp b/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
index 93e370d91e6b9..76a7e0f3831a2 100644
--- a/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
+++ b/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
@@ -451,16 +451,14 @@ void mlir::configureGpuToNVVMTypeConverter(LLVMTypeConverter &converter) {
converter, [](gpu::AddressSpace space) -> unsigned {
switch (space) {
case gpu::AddressSpace::Global:
- return static_cast<unsigned>(
- NVVM::NVVMMemorySpace::kGlobalMemorySpace);
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Global);
case gpu::AddressSpace::Workgroup:
- return static_cast<unsigned>(
- NVVM::NVVMMemorySpace::kSharedMemorySpace);
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Shared);
case gpu::AddressSpace::Private:
return 0;
}
llvm_unreachable("unknown address space enum value");
- return 0;
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Generic);
});
// Lowering for MMAMatrixType.
converter.addConversion([&](gpu::MMAMatrixType type) -> Type {
@@ -648,7 +646,7 @@ void mlir::populateGpuToNVVMConversionPatterns(
GPUFuncOpLoweringOptions{
/*allocaAddrSpace=*/0,
/*workgroupAddrSpace=*/
- static_cast<unsigned>(NVVM::NVVMMemorySpace::kSharedMemorySpace),
+ static_cast<unsigned>(NVVM::NVVMMemorySpace::Shared),
StringAttr::get(&converter.getContext(),
NVVM::NVVMDialect::getKernelFuncAttrName()),
StringAttr::get(&converter.getContext(),
diff --git a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
index 37d12bad298df..b7e3491117e9b 100644
--- a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
+++ b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
@@ -405,16 +405,14 @@ struct ConvertNVGPUToNVVMPass
converter, [](gpu::AddressSpace space) -> unsigned {
switch (space) {
case gpu::AddressSpace::Global:
- return static_cast<unsigned>(
- NVVM::NVVMMemorySpace::kGlobalMemorySpace);
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Global);
case gpu::AddressSpace::Workgroup:
- return static_cast<unsigned>(
- NVVM::NVVMMemorySpace::kSharedMemorySpace);
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Shared);
case gpu::AddressSpace::Private:
return 0;
}
llvm_unreachable("unknown address space enum value");
- return 0;
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Generic);
});
/// device-side async tokens cannot be materialized in nvvm. We just
/// convert them to a dummy i32 type in order to easily drop them during
@@ -677,7 +675,7 @@ struct NVGPUAsyncCopyLowering
adaptor.getSrcIndices());
// Intrinsics takes a global pointer so we need an address space cast.
auto srcPointerGlobalType = LLVM::LLVMPointerType::get(
- op->getContext(), NVVM::NVVMMemorySpace::kGlobalMemorySpace);
+ op->getContext(), static_cast<unsigned>(NVVM::NVVMMemorySpace::Global));
scrPtr = LLVM::AddrSpaceCastOp::create(b, srcPointerGlobalType, scrPtr);
int64_t dstElements = adaptor.getDstElements().getZExtValue();
int64_t sizeInBytes =
diff --git a/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp b/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp
index c766539f9d91a..2561f6606067f 100644
--- a/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp
+++ b/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp
@@ -71,16 +71,14 @@ void transform::ApplyGPUToNVVMConversionPatternsOp::populatePatterns(
llvmTypeConverter, [](AddressSpace space) -> unsigned {
switch (space) {
case AddressSpace::Global:
- return static_cast<unsigned>(
- NVVM::NVVMMemorySpace::kGlobalMemorySpace);
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Global);
case AddressSpace::Workgroup:
- return static_cast<unsigned>(
- NVVM::NVVMMemorySpace::kSharedMemorySpace);
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Shared);
case AddressSpace::Private:
return 0;
}
llvm_unreachable("unknown address space enum value");
- return 0;
+ return static_cast<unsigned>(NVVM::NVVMMemorySpace::Generic);
});
// Used in GPUToNVVM/WmmaOpsToNvvm.cpp so attaching here for now.
// TODO: We should have a single to_nvvm_type_converter.
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp
index 23610fbd657fd..b8331e0068880 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp
@@ -57,10 +57,10 @@ void LLVMDialect::registerAttributes() {
//===----------------------------------------------------------------------===//
/// Checks whether the given type is an LLVM type that can be loaded or stored.
-static bool isValidLoadStoreImpl(Type type, ptr::AtomicOrdering ordering,
- std::optional<int64_t> alignment,
- const ::mlir::DataLayout *dataLayout,
- function_ref<InFlightDiagnostic()> emitError) {
+bool LLVM::detail::isValidLoadStoreImpl(
+ Type type, ptr::AtomicOrdering ordering, std::optional<int64_t> alignment,
+ const ::mlir::DataLayout *dataLayout,
+ function_ref<InFlightDiagnostic()> emitError) {
if (!isLoadableType(type)) {
if (emitError)
emitError() << "type must be LLVM type with size, but got " << type;
@@ -87,14 +87,16 @@ bool AddressSpaceAttr::isValidLoad(
Type type, ptr::AtomicOrdering ordering, std::optional<int64_t> alignment,
const ::mlir::DataLayout *dataLayout,
function_ref<InFlightDiagnostic()> emitError) const {
- return isValidLoadStoreImpl(type, ordering, alignment, dataLayout, emitError);
+ return detail::isValidLoadStoreImpl(type, ordering, alignment, dataLayout,
+ emitError);
}
bool AddressSpaceAttr::isValidStore(
Type type, ptr::AtomicOrdering ordering, std::optional<int64_t> alignment,
const ::mlir::DataLayout *dataLayout,
function_ref<InFlightDiagnostic()> emitError) const {
- return isValidLoadStoreImpl(type, ordering, alignment, dataLayout, emitError);
+ return detail::isValidLoadStoreImpl(type, ordering, alignment, dataLayout,
+ emitError);
}
bool AddressSpaceAttr::isValidAtomicOp(
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
index 2dd0132a65bc4..01a16ce2d8a7f 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
@@ -703,14 +703,14 @@ const llvm::fltSemantics &LLVMPPCFP128Type::getFloatSemantics() const {
//===----------------------------------------------------------------------===//
/// Check whether type is a compatible ptr type. These are pointer-like types
-/// with no element type, no metadata, and using the LLVM AddressSpaceAttr
-/// memory space.
+/// with no element type, no metadata, and using the LLVM
+/// LLVMAddrSpaceAttrInterface memory space.
static bool isCompatiblePtrType(Type type) {
auto ptrTy = dyn_cast<PtrLikeTypeInterface>(type);
if (!ptrTy)
return false;
return !ptrTy.hasPtrMetadata() && ptrTy.getElementType() == nullptr &&
- isa<AddressSpaceAttr>(ptrTy.getMemorySpace());
+ isa<LLVMAddrSpaceAttrInterface>(ptrTy.getMemorySpace());
}
bool mlir::LLVM::isCompatibleOuterType(Type type) {
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 376e3c3e1fcbe..03cc0a611513d 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -800,8 +800,8 @@ inferMMATypeFromMNK(NVVM::MMATypes type, NVVM::MMAFrag frag, int m, int n,
LogicalResult NVVM::WMMALoadOp::verify() {
unsigned addressSpace =
llvm::cast<LLVM::LLVMPointerType>(getPtr().getType()).getAddressSpace();
- if (addressSpace != 0 && addressSpace != NVVM::kGlobalMemorySpace &&
- addressSpace != NVVM::kSharedMemorySpace)
+ if (addressSpace != 0 && addressSpace != NVVMMemorySpace::Global &&
+ addressSpace != NVVMMemorySpace::Shared)
return emitOpError("expected source pointer in memory "
"space 0, 1, 3");
@@ -821,8 +821,8 @@ LogicalResult NVVM::WMMALoadOp::verify() {
LogicalResult NVVM::WMMAStoreOp::verify() {
unsigned addressSpace =
llvm::cast<LLVM::LLVMPointerType>(getPtr().getType()).getAddressSpace();
- if (addressSpace != 0 && addressSpace != NVVM::kGlobalMemorySpace &&
- addressSpace != NVVM::kSharedMemorySpace)
+ if (addressSpace != 0 && addressSpace != NVVMMemorySpace::Global &&
+ addressSpace != NVVMMemorySpace::Shared)
return emitOpError("expected operands to be a source pointer in memory "
"space 0, 1, 3");
@@ -1339,8 +1339,8 @@ LogicalResult NVVM::PrefetchOp::verify() {
return emitOpError("cannot specify both tensormap and cache level");
if (getTensormap()) {
- if (addressSpace != MemSpace::kGenericMemorySpace &&
- addressSpace != MemSpace::kConstantMemorySpace) {
+ if (addressSpace != MemSpace::Generic &&
+ addressSpace != MemSpace::Constant) {
return emitOpError(
"prefetch tensormap requires a generic or constant pointer");
}
@@ -1350,15 +1350,14 @@ LogicalResult NVVM::PrefetchOp::verify() {
"prefetch tensormap does not support eviction priority");
}
- if (getInParamSpace() && addressSpace != MemSpace::kGenericMemorySpace) {
+ if (getInParamSpace() && addressSpace != MemSpace::Generic) {
return emitOpError(
"in_param_space can only be specified for a generic pointer");
}
} else if (cacheLevel) {
- if (addressSpace != MemSpace::kGenericMemorySpace &&
- addressSpace != MemSpace::kGlobalMemorySpace &&
- addressSpace != MemSpace::kLocalMemorySpace) {
+ if (addressSpace != MemSpace::Generic && addressSpace != MemSpace::Global &&
+ addressSpace != MemSpace::Local) {
return emitOpError("prefetch to cache level requires a generic, global, "
"or local pointer");
}
@@ -1370,7 +1369,7 @@ LogicalResult NVVM::PrefetchOp::verify() {
"cache level is L1");
}
- if (addressSpace != MemSpace::kGenericMemorySpace) {
+ if (addressSpace != MemSpace::Generic) {
return emitOpError(
"prefetch to uniform cache requires a generic pointer");
}
@@ -1381,7 +1380,7 @@ LogicalResult NVVM::PrefetchOp::verify() {
return emitOpError(
"cache eviction priority supported only for cache level L2");
- if (addressSpace != MemSpace::kGlobalMemorySpace)
+ if (addressSpace != MemSpace::Global)
return emitOpError("cache eviction priority requires a global pointer");
if (*evictPriority != NVVM::CacheEvictionPriority::EvictNormal &&
@@ -1796,7 +1795,7 @@ Tcgen05AllocOp::getIntrinsicIDAndArgs(Operation &op,
auto curOp = cast<NVVM::Tcgen05AllocOp>(op);
unsigned as = llvm::cast<LLVM::LLVMPointerType>(curOp.getAddr().getType())
.getAddressSpace();
- bool isShared = as == NVVMMemorySpace::kSharedMemorySpace;
+ bool isShared = as == NVVMMemorySpace::Shared;
bool is2CTAMode = curOp.getGroup() == CTAGroupKind::CTA_2;
llvm::Intrinsic::ID id;
@@ -1845,7 +1844,7 @@ Tcgen05CommitOp::getIntrinsicIDAndArgs(Operation &op,
auto curOp = cast<NVVM::Tcgen05CommitOp>(op);
unsigned as = llvm::cast<LLVM::LLVMPointerType>(curOp.getAddr().getType())
.getAddressSpace();
- bool isShared = as == NVVMMemorySpace::kSharedMemorySpace;
+ bool isShared = as == NVVMMemorySpace::Shared;
bool hasMulticast = static_cast<bool>(curOp.getMulticastMask());
bool is2CTAMode = curOp.getGroup() == CTAGroupKind::CTA_2;
@@ -2051,18 +2050,18 @@ PrefetchOp::getIntrinsicIDAndArgs(NVVM::PrefetchOp &op,
}
}
- switch (addressSpace) {
- case MemSpace::kGenericMemorySpace:
+ switch (static_cast<MemSpace>(addressSpace)) {
+ case MemSpace::Generic:
return *cacheLevel == CacheLevel::L1
? NVVM::IDArgPair({llvm::Intrinsic::nvvm_prefetch_L1, args})
: NVVM::IDArgPair({llvm::Intrinsic::nvvm_prefetch_L2, args});
- case MemSpace::kGlobalMemorySpace:
+ case MemSpace::Global:
return *cacheLevel == CacheLevel::L1
? NVVM::IDArg...
[truncated]
|
llvmbot
added
flang
fabianmcg
force-pushed
the
llvm-as-interface
branch
from
7238493 to
f032688
Compare
fabianmcg
force-pushed
the
llvm-as-interface
branch
from
f032688 to
a175c07
Compare
xlauko
changed the title
joker-eph
approved these changes
Sep 13, 2025
| @@ -3592,7 +3627,7 @@ def NVVM_MapaOp: NVVM_Op<"mapa", | |||
| string llvmBuilder = [{ | |||
| int addrSpace = llvm::cast<LLVMPointerType>(op.getA().getType()).getAddressSpace(); | |||
|
|
|||
| bool isSharedMemory = addrSpace == NVVM::NVVMMemorySpace::kSharedMemorySpace; | |||
| bool isSharedMemory = addrSpace == static_cast<int> (NVVM::NVVMMemorySpace::Shared); | |||
There was a problem hiding this comment.
@fabianmcg is this supposed to be a static_cast<unsigned>? If not, why?
There was a problem hiding this comment.
I think you're right, that should be an unsigned cast, line 3628 also should be changed to unsigned.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This patch introduces the
LLVMAddrSpaceAttrInterfacefor defining compatible LLVM address space attributesTo test this interface, this patch also adds:
Example:
Translating the above code to LLVM produces:
To convert the memory space enum to the new enum class use:
NOTE: A future patch will add support for ROCDL, it wasn't added here to keep the patch small.