diff --git a/flang/lib/Optimizer/CodeGen/CodeGen.cpp b/flang/lib/Optimizer/CodeGen/CodeGen.cpp index 76f3cbd421cb9..88946aae9f4bc 100644 --- a/flang/lib/Optimizer/CodeGen/CodeGen.cpp +++ b/flang/lib/Optimizer/CodeGen/CodeGen.cpp @@ -3202,7 +3202,8 @@ struct GlobalOpConversion : public fir::FIROpConversion { if (global.getDataAttr() && *global.getDataAttr() == cuf::DataAttribute::Shared) - g.setAddrSpace(mlir::NVVM::NVVMMemorySpace::kSharedMemorySpace); + g.setAddrSpace( + static_cast(mlir::NVVM::NVVMMemorySpace::Shared)); rewriter.eraseOp(global); return mlir::success(); diff --git a/flang/lib/Optimizer/Transforms/CUFGPUToLLVMConversion.cpp b/flang/lib/Optimizer/Transforms/CUFGPUToLLVMConversion.cpp index a40ed95391c3a..40f180a8c1657 100644 --- a/flang/lib/Optimizer/Transforms/CUFGPUToLLVMConversion.cpp +++ b/flang/lib/Optimizer/Transforms/CUFGPUToLLVMConversion.cpp @@ -221,7 +221,8 @@ static mlir::Value createAddressOfOp(mlir::ConversionPatternRewriter &rewriter, gpu::GPUModuleOp gpuMod, std::string &sharedGlobalName) { auto llvmPtrTy = mlir::LLVM::LLVMPointerType::get( - rewriter.getContext(), mlir::NVVM::NVVMMemorySpace::kSharedMemorySpace); + rewriter.getContext(), + static_cast(mlir::NVVM::NVVMMemorySpace::Shared)); if (auto g = gpuMod.lookupSymbol(sharedGlobalName)) return mlir::LLVM::AddressOfOp::create(rewriter, loc, llvmPtrTy, g.getSymName()); 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 ]> { 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 alignment, + const ::mlir::DataLayout *dataLayout, + function_ref 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(memSpace); +} +inline bool operator==(NVVMMemorySpace memSpace, unsigned as) { + return static_cast(memSpace) == as; +} +inline bool operator!=(unsigned as, NVVMMemorySpace memSpace) { + return as != static_cast(memSpace); +} +inline bool operator!=(NVVMMemorySpace memSpace, unsigned as) { + return static_cast(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; +/// 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, + DeclareAttrInterfaceMethods + ]> { + let assemblyFormat = "`<` $value `>`"; +} + //===----------------------------------------------------------------------===// // NVVM intrinsic operations //===----------------------------------------------------------------------===// @@ -3592,7 +3627,7 @@ def NVVM_MapaOp: NVVM_Op<"mapa", string llvmBuilder = [{ int addrSpace = llvm::cast(op.getA().getType()).getAddressSpace(); - bool isSharedMemory = addrSpace == NVVM::NVVMMemorySpace::kSharedMemorySpace; + bool isSharedMemory = addrSpace == static_cast (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( - NVVM::NVVMMemorySpace::kGlobalMemorySpace); + return static_cast(NVVM::NVVMMemorySpace::Global); case gpu::AddressSpace::Workgroup: - return static_cast( - NVVM::NVVMMemorySpace::kSharedMemorySpace); + return static_cast(NVVM::NVVMMemorySpace::Shared); case gpu::AddressSpace::Private: return 0; } llvm_unreachable("unknown address space enum value"); - return 0; + return static_cast(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(NVVM::NVVMMemorySpace::kSharedMemorySpace), + static_cast(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( - NVVM::NVVMMemorySpace::kGlobalMemorySpace); + return static_cast(NVVM::NVVMMemorySpace::Global); case gpu::AddressSpace::Workgroup: - return static_cast( - NVVM::NVVMMemorySpace::kSharedMemorySpace); + return static_cast(NVVM::NVVMMemorySpace::Shared); case gpu::AddressSpace::Private: return 0; } llvm_unreachable("unknown address space enum value"); - return 0; + return static_cast(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(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( - NVVM::NVVMMemorySpace::kGlobalMemorySpace); + return static_cast(NVVM::NVVMMemorySpace::Global); case AddressSpace::Workgroup: - return static_cast( - NVVM::NVVMMemorySpace::kSharedMemorySpace); + return static_cast(NVVM::NVVMMemorySpace::Shared); case AddressSpace::Private: return 0; } llvm_unreachable("unknown address space enum value"); - return 0; + return static_cast(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 alignment, - const ::mlir::DataLayout *dataLayout, - function_ref emitError) { +bool LLVM::detail::isValidLoadStoreImpl( + Type type, ptr::AtomicOrdering ordering, std::optional alignment, + const ::mlir::DataLayout *dataLayout, + function_ref 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 alignment, const ::mlir::DataLayout *dataLayout, function_ref 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 alignment, const ::mlir::DataLayout *dataLayout, function_ref 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(type); if (!ptrTy) return false; return !ptrTy.hasPtrMetadata() && ptrTy.getElementType() == nullptr && - isa(ptrTy.getMemorySpace()); + isa(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(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(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(op); unsigned as = llvm::cast(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(op); unsigned as = llvm::cast(curOp.getAddr().getType()) .getAddressSpace(); - bool isShared = as == NVVMMemorySpace::kSharedMemorySpace; + bool isShared = as == NVVMMemorySpace::Shared; bool hasMulticast = static_cast(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(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::IDArgPair( {llvm::Intrinsic::nvvm_prefetch_global_L1, args}) : NVVM::IDArgPair( {llvm::Intrinsic::nvvm_prefetch_global_L2, args}); - case MemSpace::kLocalMemorySpace: + case MemSpace::Local: return *cacheLevel == CacheLevel::L1 ? NVVM::IDArgPair( {llvm::Intrinsic::nvvm_prefetch_local_L1, args}) @@ -2185,6 +2184,66 @@ LogicalResult NVVMDialect::verifyRegionArgAttribute(Operation *op, return success(); } +//===----------------------------------------------------------------------===// +// NVVM Address Space Attr +//===----------------------------------------------------------------------===// + +unsigned NVVMMemorySpaceAttr::getAddressSpace() const { + return static_cast(getValue()); +} + +bool NVVMMemorySpaceAttr::isValidLoad( + Type type, ptr::AtomicOrdering ordering, std::optional alignment, + const ::mlir::DataLayout *dataLayout, + function_ref emitError) const { + return LLVM::detail::isValidLoadStoreImpl(type, ordering, alignment, + dataLayout, emitError); +} + +bool NVVMMemorySpaceAttr::isValidStore( + Type type, ptr::AtomicOrdering ordering, std::optional alignment, + const ::mlir::DataLayout *dataLayout, + function_ref emitError) const { + return LLVM::detail::isValidLoadStoreImpl(type, ordering, alignment, + dataLayout, emitError); +} + +bool NVVMMemorySpaceAttr::isValidAtomicOp( + ptr::AtomicBinOp op, Type type, ptr::AtomicOrdering ordering, + std::optional alignment, const ::mlir::DataLayout *dataLayout, + function_ref emitError) const { + // TODO: update this method once `ptr.atomic_rmw` is implemented. + assert(false && "unimplemented, see TODO in the source."); + return false; +} + +bool NVVMMemorySpaceAttr::isValidAtomicXchg( + Type type, ptr::AtomicOrdering successOrdering, + ptr::AtomicOrdering failureOrdering, std::optional alignment, + const ::mlir::DataLayout *dataLayout, + function_ref emitError) const { + // TODO: update this method once `ptr.atomic_cmpxchg` is implemented. + assert(false && "unimplemented, see TODO in the source."); + return false; +} + +bool NVVMMemorySpaceAttr::isValidAddrSpaceCast( + Type tgt, Type src, function_ref emitError) const { + // TODO: update this method once the `ptr.addrspace_cast` op is added to the + // dialect. + assert(false && "unimplemented, see TODO in the source."); + return false; +} + +bool NVVMMemorySpaceAttr::isValidPtrIntCast( + Type intLikeTy, Type ptrLikeTy, + function_ref emitError) const { + // TODO: update this method once the int-cast ops are added to the `ptr` + // dialect. + assert(false && "unimplemented, see TODO in the source."); + return false; +} + //===----------------------------------------------------------------------===// // NVVM target attribute. //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/NVGPU/TransformOps/NVGPUTransformOps.cpp b/mlir/lib/Dialect/NVGPU/TransformOps/NVGPUTransformOps.cpp index bc3e8b2b17fb1..46e82bd8fc8c8 100644 --- a/mlir/lib/Dialect/NVGPU/TransformOps/NVGPUTransformOps.cpp +++ b/mlir/lib/Dialect/NVGPU/TransformOps/NVGPUTransformOps.cpp @@ -53,16 +53,14 @@ void transform::ApplyNVGPUToNVVMConversionPatternsOp::populatePatterns( llvmTypeConverter, [](gpu::AddressSpace space) -> unsigned { switch (space) { case gpu::AddressSpace::Global: - return static_cast( - NVVM::NVVMMemorySpace::kGlobalMemorySpace); + return static_cast(NVVM::NVVMMemorySpace::Global); case gpu::AddressSpace::Workgroup: - return static_cast( - NVVM::NVVMMemorySpace::kSharedMemorySpace); + return static_cast(NVVM::NVVMMemorySpace::Shared); case gpu::AddressSpace::Private: return 0; } llvm_unreachable("unknown address space enum value"); - return 0; + return static_cast(NVVM::NVVMMemorySpace::Generic); }); llvmTypeConverter.addConversion( [&](nvgpu::DeviceAsyncTokenType type) -> Type { diff --git a/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp index 7f69af14df338..3d86b09b32538 100644 --- a/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp +++ b/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp @@ -253,8 +253,8 @@ getStMatrixIntrinsicId(NVVM::MMALayout layout, int32_t num, /// Return the intrinsic ID associated with st.bulk for the given address type. static llvm::Intrinsic::ID getStBulkIntrinsicId(LLVM::LLVMPointerType addrType) { - bool isSharedMemory = - addrType.getAddressSpace() == NVVM::NVVMMemorySpace::kSharedMemorySpace; + bool isSharedMemory = addrType.getAddressSpace() == + static_cast(NVVM::NVVMMemorySpace::Shared); return isSharedMemory ? llvm::Intrinsic::nvvm_st_bulk_shared_cta : llvm::Intrinsic::nvvm_st_bulk; } diff --git a/mlir/lib/Target/LLVMIR/TypeToLLVM.cpp b/mlir/lib/Target/LLVMIR/TypeToLLVM.cpp index ddd5946ce5d63..4d204744450a8 100644 --- a/mlir/lib/Target/LLVMIR/TypeToLLVM.cpp +++ b/mlir/lib/Target/LLVMIR/TypeToLLVM.cpp @@ -152,8 +152,9 @@ class TypeToLLVMIRTranslatorImpl { /// Translates the given ptr type. llvm::Type *translate(PtrLikeTypeInterface type) { - auto memSpace = dyn_cast(type.getMemorySpace()); - assert(memSpace && "expected pointer with the LLVM address space"); + auto memSpace = + dyn_cast(type.getMemorySpace()); + assert(memSpace && "expected pointer with an LLVM address space"); assert(!type.hasPtrMetadata() && "expected pointer without metadata"); return llvm::PointerType::get(context, memSpace.getAddressSpace()); } diff --git a/mlir/test/Dialect/LLVMIR/nvvm.mlir b/mlir/test/Dialect/LLVMIR/nvvm.mlir index 5209b3c1d7906..3277e62893527 100644 --- a/mlir/test/Dialect/LLVMIR/nvvm.mlir +++ b/mlir/test/Dialect/LLVMIR/nvvm.mlir @@ -620,6 +620,16 @@ func.func @prefetch_tensormap(%gen_ptr: !llvm.ptr, %const_ptr: !llvm.ptr<4>) { return } +// CHECK-LABEL: @nvvm_address_space +func.func private @nvvm_address_space( + !ptr.ptr<#nvvm.memory_space>, + !ptr.ptr<#nvvm.memory_space>, + !ptr.ptr<#nvvm.memory_space>, + !ptr.ptr<#nvvm.memory_space>, + !ptr.ptr<#nvvm.memory_space>, + !ptr.ptr<#nvvm.memory_space> + ) -> !ptr.ptr<#nvvm.memory_space> + // ----- // Just check these don't emit errors. diff --git a/mlir/test/Target/LLVMIR/ptr.mlir b/mlir/test/Target/LLVMIR/ptr.mlir index 4b29be813fa81..9b99dd8e3a3eb 100644 --- a/mlir/test/Target/LLVMIR/ptr.mlir +++ b/mlir/test/Target/LLVMIR/ptr.mlir @@ -233,3 +233,25 @@ llvm.func @ptr_add_vector_base_scalar_offset(%ptrs: vector<4x!ptr.ptr<#llvm.addr %res = ptr.ptr_add %ptrs, %offset : vector<4x!ptr.ptr<#llvm.address_space<0>>>, i32 llvm.return %res : vector<4x!ptr.ptr<#llvm.address_space<0>>> } + +// CHECK-LABEL: declare ptr @nvvm_ptr_address_space(ptr addrspace(1), ptr addrspace(3), ptr addrspace(4), ptr addrspace(5), ptr addrspace(6), ptr addrspace(7)) +llvm.func @nvvm_ptr_address_space( + !ptr.ptr<#nvvm.memory_space>, + !ptr.ptr<#nvvm.memory_space>, + !ptr.ptr<#nvvm.memory_space>, + !ptr.ptr<#nvvm.memory_space>, + !ptr.ptr<#nvvm.memory_space>, + !ptr.ptr<#nvvm.memory_space> + ) -> !ptr.ptr<#nvvm.memory_space> + +// CHECK-LABEL: define void @llvm_ops_with_ptr_nvvm_values +// CHECK-SAME: (ptr %[[ARG:.*]]) { +// CHECK-NEXT: %[[V0:.*]] = load ptr addrspace(1), ptr %[[ARG]], align 8 +// CHECK-NEXT: store ptr addrspace(1) %[[V0]], ptr %[[ARG]], align 8 +// CHECK-NEXT: ret void +// CHECK-NEXT: } +llvm.func @llvm_ops_with_ptr_nvvm_values(%arg0: !llvm.ptr) { + %1 = llvm.load %arg0 : !llvm.ptr -> !ptr.ptr<#nvvm.memory_space> + llvm.store %1, %arg0 : !ptr.ptr<#nvvm.memory_space>, !llvm.ptr + llvm.return +}