MLIR 22.0.0git
VectorToGPU.cpp File Reference

Go to the source code of this file.

Classes

class  mlir::impl::ConvertVectorToGPUBase< DerivedT >

Namespaces

namespace  mlir
 Include the generated interface declarations.
namespace  mlir::impl
 Attribute collections provide a dictionary-like interface.

Macros

#define DEBUG_TYPE   "vector-to-gpu"
#define GEN_PASS_DEF_CONVERTVECTORTOGPU

Functions

template<typename TransferOpType>
static void getXferIndices (RewriterBase &rewriter, TransferOpType xferOp, AffineMap offsetMap, ArrayRef< Value > dimValues, SmallVector< Value, 4 > &indices)
 For a vector TransferOpType xferOp, an empty indices vector, and an AffineMap representing offsets to apply to indices, the function fills indices with the original indices plus the offsets.
static bool contractSupportsMMAMatrixType (vector::ContractionOp contract, bool useNvGpu)
static bool isTransposeMatrixLoadMap (AffineMap permutationMap)
static std::optional< int64_tgetStaticallyKnownRowStride (ShapedType type)
static bool transferReadSupportsMMAMatrixType (vector::TransferReadOp readOp)
static bool transferWriteSupportsMMAMatrixType (vector::TransferWriteOp writeOp)
static bool constantSupportsMMAMatrixType (arith::ConstantOp constantOp)
 Return true if the constant is a splat to a 2D vector so that it can be converted to a MMA constant matrix op.
static bool broadcastSupportsMMAMatrixType (vector::BroadcastOp broadcastOp)
 Return true if this is a broadcast from scalar to a 2D vector.
template<typename ExtOpTy>
static bool integerExtendSupportsMMAMatrixType (ExtOpTy extOp)
 Return true if this integer extend op can be folded into a contract op.
static bool fpExtendSupportsMMAMatrixType (arith::ExtFOp extOp)
static std::optional< gpu::MMAElementwiseOp > convertElementwiseOpToMMA (Operation *op)
 Return the MMA elementwise enum associated with op if it is supported.
static bool elementwiseSupportsMMAMatrixType (Operation *op)
 Return true if the op is supported as elementwise op on MMAMatrix type.
static bool extractStridedSliceSupportsMMAMatrixType (vector::ExtractStridedSliceOp op)
 Returns true if the extract strided slice op is supported with mma.sync path.
static bool supportsMMaMatrixType (Operation *op, bool useNvGpu)
static SetVector< Operation * > getSliceContract (Operation *op, const BackwardSliceOptions &backwardSliceOptions, const ForwardSliceOptions &forwardSliceOptions)
 Return an unsorted slice handling scf.for region differently than getSlice.
static SetVector< Operation * > getOpToConvert (mlir::Operation *op, bool useNvGpu)
static const char * inferFragType (Operation *op)
static LogicalResult convertTransferReadOp (RewriterBase &rewriter, vector::TransferReadOp op, llvm::DenseMap< Value, Value > &valueMapping)
static LogicalResult convertTransferWriteOp (RewriterBase &rewriter, vector::TransferWriteOp op, llvm::DenseMap< Value, Value > &valueMapping)
static VectorType getMmaSyncVectorOperandType (const nvgpu::FragmentElementInfo &regInfo)
 Returns the vector type which represents a matrix fragment.
static LogicalResult convertConstantOpMmaSync (RewriterBase &rewriter, arith::ConstantOp op, llvm::DenseMap< Value, Value > &valueMapping)
 Convert a 2D splat ConstantOp to a SubgroupMmaConstantMatrix op.
static FailureOr< boolisTransposed (vector::TransferReadOp op)
 Check if the loaded matrix operand requires transposed.
static LogicalResult creatLdMatrixCompatibleLoads (RewriterBase &rewriter, vector::TransferReadOp op, llvm::DenseMap< Value, Value > &valueMapping)
static LogicalResult createNonLdMatrixLoads (RewriterBase &rewriter, vector::TransferReadOp op, llvm::DenseMap< Value, Value > &valueMapping)
static bool isSharedMemory (MemRefType type)
 Return true if this is a shared memory memref type.
static LogicalResult convertTransferReadToLoads (RewriterBase &rewriter, vector::TransferReadOp op, llvm::DenseMap< Value, Value > &valueMapping)
 Converts a vector.transfer_read operation directly to either a vector.load or a nvgpu.ldmatrix operation.
static LogicalResult convertTransferWriteToStores (RewriterBase &rewriter, vector::TransferWriteOp op, llvm::DenseMap< Value, Value > &valueMapping)
static void populateFromInt64AttrArray (ArrayAttr arrayAttr, SmallVectorImpl< int64_t > &results)
static LogicalResult convertExtractStridedSlice (RewriterBase &rewriter, vector::ExtractStridedSliceOp op, llvm::DenseMap< Value, Value > &valueMapping)
static LogicalResult convertContractOp (RewriterBase &rewriter, vector::ContractionOp op, llvm::DenseMap< Value, Value > &valueMapping)
static LogicalResult convertContractOpToMmaSync (RewriterBase &rewriter, vector::ContractionOp op, llvm::DenseMap< Value, Value > &valueMapping)
static LogicalResult convertConstantOp (RewriterBase &rewriter, arith::ConstantOp op, llvm::DenseMap< Value, Value > &valueMapping)
 Convert a 2D splat ConstantOp to a SubgroupMmaConstantMatrix op.
static LogicalResult convertBroadcastOp (RewriterBase &rewriter, vector::BroadcastOp op, llvm::DenseMap< Value, Value > &valueMapping)
 Convert a vector.broadcast from scalar to a SubgroupMmaConstantMatrix op.
static scf::ForOp replaceForOpWithNewSignature (RewriterBase &rewriter, scf::ForOp loop, ValueRange newInitArgs)
static LogicalResult convertForOp (RewriterBase &rewriter, scf::ForOp op, llvm::DenseMap< Value, Value > &valueMapping)
static LogicalResult convertYieldOp (RewriterBase &rewriter, scf::YieldOp op, llvm::DenseMap< Value, Value > &valueMapping)
static LogicalResult convertElementwiseOp (RewriterBase &rewriter, Operation *op, gpu::MMAElementwiseOp opType, llvm::DenseMap< Value, Value > &valueMapping)
 Convert an elementwise op to the equivalent elementwise op on MMA matrix.

Macro Definition Documentation

◆ DEBUG_TYPE

#define DEBUG_TYPE   "vector-to-gpu"

Definition at line 36 of file VectorToGPU.cpp.

◆ GEN_PASS_DEF_CONVERTVECTORTOGPU

#define GEN_PASS_DEF_CONVERTVECTORTOGPU

Definition at line 39 of file VectorToGPU.cpp.

Function Documentation

◆ broadcastSupportsMMAMatrixType()

bool broadcastSupportsMMAMatrixType ( vector::BroadcastOp broadcastOp)
static

Return true if this is a broadcast from scalar to a 2D vector.

Definition at line 191 of file VectorToGPU.cpp.

Referenced by convertBroadcastOp(), and supportsMMaMatrixType().

◆ constantSupportsMMAMatrixType()

bool constantSupportsMMAMatrixType ( arith::ConstantOp constantOp)
static

Return true if the constant is a splat to a 2D vector so that it can be converted to a MMA constant matrix op.

Definition at line 183 of file VectorToGPU.cpp.

Referenced by convertConstantOp(), and supportsMMaMatrixType().

◆ contractSupportsMMAMatrixType()

bool contractSupportsMMAMatrixType ( vector::ContractionOp contract,
bool useNvGpu )
static

◆ convertBroadcastOp()

LogicalResult convertBroadcastOp ( RewriterBase & rewriter,
vector::BroadcastOp op,
llvm::DenseMap< Value, Value > & valueMapping )
static

Convert a vector.broadcast from scalar to a SubgroupMmaConstantMatrix op.

Definition at line 1092 of file VectorToGPU.cpp.

References broadcastSupportsMMAMatrixType(), mlir::gpu::MMAMatrixType::get(), inferFragType(), mlir::OpBuilder::setInsertionPoint(), and success().

Referenced by mlir::convertVectorToMMAOps().

◆ convertConstantOp()

LogicalResult convertConstantOp ( RewriterBase & rewriter,
arith::ConstantOp op,
llvm::DenseMap< Value, Value > & valueMapping )
static

Convert a 2D splat ConstantOp to a SubgroupMmaConstantMatrix op.

Definition at line 1069 of file VectorToGPU.cpp.

References constantSupportsMMAMatrixType(), mlir::gpu::MMAMatrixType::get(), inferFragType(), mlir::OpBuilder::setInsertionPoint(), and success().

Referenced by mlir::convertVectorToMMAOps().

◆ convertConstantOpMmaSync()

LogicalResult convertConstantOpMmaSync ( RewriterBase & rewriter,
arith::ConstantOp op,
llvm::DenseMap< Value, Value > & valueMapping )
static

◆ convertContractOp()

LogicalResult convertContractOp ( RewriterBase & rewriter,
vector::ContractionOp op,
llvm::DenseMap< Value, Value > & valueMapping )
static

◆ convertContractOpToMmaSync()

LogicalResult convertContractOpToMmaSync ( RewriterBase & rewriter,
vector::ContractionOp op,
llvm::DenseMap< Value, Value > & valueMapping )
static

◆ convertElementwiseOp()

LogicalResult convertElementwiseOp ( RewriterBase & rewriter,
Operation * op,
gpu::MMAElementwiseOp opType,
llvm::DenseMap< Value, Value > & valueMapping )
static

◆ convertElementwiseOpToMMA()

std::optional< gpu::MMAElementwiseOp > convertElementwiseOpToMMA ( Operation * op)
static

Return the MMA elementwise enum associated with op if it is supported.

Return std::nullopt otherwise.

Definition at line 210 of file VectorToGPU.cpp.

Referenced by mlir::convertVectorToMMAOps(), and elementwiseSupportsMMAMatrixType().

◆ convertExtractStridedSlice()

LogicalResult convertExtractStridedSlice ( RewriterBase & rewriter,
vector::ExtractStridedSliceOp op,
llvm::DenseMap< Value, Value > & valueMapping )
static

◆ convertForOp()

LogicalResult convertForOp ( RewriterBase & rewriter,
scf::ForOp op,
llvm::DenseMap< Value, Value > & valueMapping )
static

◆ convertTransferReadOp()

◆ convertTransferReadToLoads()

LogicalResult convertTransferReadToLoads ( RewriterBase & rewriter,
vector::TransferReadOp op,
llvm::DenseMap< Value, Value > & valueMapping )
static

Converts a vector.transfer_read operation directly to either a vector.load or a nvgpu.ldmatrix operation.

This function should only be used when converting to nvgpu.mma.sync operations.

Definition at line 865 of file VectorToGPU.cpp.

References createNonLdMatrixLoads(), creatLdMatrixCompatibleLoads(), mlir::nvgpu::getWarpMatrixInfo(), isSharedMemory(), mlir::RewriterBase::notifyMatchFailure(), and mlir::OpBuilder::setInsertionPoint().

Referenced by mlir::convertVectorToNVVMCompatibleMMASync().

◆ convertTransferWriteOp()

◆ convertTransferWriteToStores()

◆ convertYieldOp()

LogicalResult convertYieldOp ( RewriterBase & rewriter,
scf::YieldOp op,
llvm::DenseMap< Value, Value > & valueMapping )
static

◆ createNonLdMatrixLoads()

◆ creatLdMatrixCompatibleLoads()

LogicalResult creatLdMatrixCompatibleLoads ( RewriterBase & rewriter,
vector::TransferReadOp op,
llvm::DenseMap< Value, Value > & valueMapping )
static

◆ elementwiseSupportsMMAMatrixType()

bool elementwiseSupportsMMAMatrixType ( Operation * op)
static

Return true if the op is supported as elementwise op on MMAMatrix type.

Definition at line 241 of file VectorToGPU.cpp.

References convertElementwiseOpToMMA().

Referenced by getOpToConvert(), and supportsMMaMatrixType().

◆ extractStridedSliceSupportsMMAMatrixType()

bool extractStridedSliceSupportsMMAMatrixType ( vector::ExtractStridedSliceOp op)
static

Returns true if the extract strided slice op is supported with mma.sync path.

Definition at line 248 of file VectorToGPU.cpp.

References mlir::nvgpu::B, mlir::nvgpu::C, mlir::nvgpu::getUserContract(), and mlir::nvgpu::getWarpMatrixInfo().

Referenced by supportsMMaMatrixType().

◆ fpExtendSupportsMMAMatrixType()

bool fpExtendSupportsMMAMatrixType ( arith::ExtFOp extOp)
static

Definition at line 205 of file VectorToGPU.cpp.

Referenced by supportsMMaMatrixType().

◆ getMmaSyncVectorOperandType()

VectorType getMmaSyncVectorOperandType ( const nvgpu::FragmentElementInfo & regInfo)
static

Returns the vector type which represents a matrix fragment.

Definition at line 627 of file VectorToGPU.cpp.

Referenced by convertConstantOpMmaSync(), convertTransferWriteToStores(), createNonLdMatrixLoads(), and creatLdMatrixCompatibleLoads().

◆ getOpToConvert()

◆ getSliceContract()

SetVector< Operation * > getSliceContract ( Operation * op,
const BackwardSliceOptions & backwardSliceOptions,
const ForwardSliceOptions & forwardSliceOptions )
static

Return an unsorted slice handling scf.for region differently than getSlice.

In scf.for we only want to include as part of the slice elements that are part of the use/def chain.

Definition at line 303 of file VectorToGPU.cpp.

References mlir::getBackwardSlice(), mlir::getForwardSlice(), and result.

Referenced by getOpToConvert().

◆ getStaticallyKnownRowStride()

std::optional< int64_t > getStaticallyKnownRowStride ( ShapedType type)
static

◆ getXferIndices()

template<typename TransferOpType>
void getXferIndices ( RewriterBase & rewriter,
TransferOpType xferOp,
AffineMap offsetMap,
ArrayRef< Value > dimValues,
SmallVector< Value, 4 > & indices )
static

For a vector TransferOpType xferOp, an empty indices vector, and an AffineMap representing offsets to apply to indices, the function fills indices with the original indices plus the offsets.

The offsets are applied by taking into account the permutation map of the transfer op. If the offsetMap has dimension placeholders, those should be provided in dimValues.

Definition at line 52 of file VectorToGPU.cpp.

References mlir::Builder::getAffineDimExpr(), mlir::AffineMap::getNumDims(), mlir::AffineMap::getResult(), indices, and mlir::affine::makeComposedAffineApply().

Referenced by convertTransferWriteToStores(), createNonLdMatrixLoads(), and creatLdMatrixCompatibleLoads().

◆ inferFragType()

◆ integerExtendSupportsMMAMatrixType()

template<typename ExtOpTy>
bool integerExtendSupportsMMAMatrixType ( ExtOpTy extOp)
static

Return true if this integer extend op can be folded into a contract op.

Definition at line 197 of file VectorToGPU.cpp.

Referenced by supportsMMaMatrixType().

◆ isSharedMemory()

bool isSharedMemory ( MemRefType type)
static

Return true if this is a shared memory memref type.

Definition at line 854 of file VectorToGPU.cpp.

Referenced by convertTransferReadToLoads(), and getStBulkIntrinsicId().

◆ isTransposed()

FailureOr< bool > isTransposed ( vector::TransferReadOp op)
static

Check if the loaded matrix operand requires transposed.

Transposed Map Example: Example 1 : (..., d0, d1) -> (d1 * 1, d0 * 2) Example 2 : (d0, d1, d2, d3) -> (d3, d2) The code below checks if the output 2D is transposed using a generalized version : (d0, d1, dn, ..., dm, ...) -> (dm, dn) Returns : true; if m > n, false o.w.

Definition at line 678 of file VectorToGPU.cpp.

References mlir::AffineMap::getNumResults(), mlir::AffineMap::getPermutationMap(), and mlir::AffineMap::getResult().

Referenced by creatLdMatrixCompatibleLoads().

◆ isTransposeMatrixLoadMap()

bool isTransposeMatrixLoadMap ( AffineMap permutationMap)
static

◆ populateFromInt64AttrArray()

void populateFromInt64AttrArray ( ArrayAttr arrayAttr,
SmallVectorImpl< int64_t > & results )
static

Definition at line 942 of file VectorToGPU.cpp.

References ArrayAttr().

Referenced by convertExtractStridedSlice().

◆ replaceForOpWithNewSignature()

scf::ForOp replaceForOpWithNewSignature ( RewriterBase & rewriter,
scf::ForOp loop,
ValueRange newInitArgs )
static

◆ supportsMMaMatrixType()

◆ transferReadSupportsMMAMatrixType()

◆ transferWriteSupportsMMAMatrixType()

bool transferWriteSupportsMMAMatrixType ( vector::TransferWriteOp writeOp)
static

Definition at line 165 of file VectorToGPU.cpp.

References getStaticallyKnownRowStride().

Referenced by convertTransferWriteOp(), and supportsMMaMatrixType().