MLIR: lib/Conversion/VectorToGPU/VectorToGPU.cpp File Reference (original) (raw)

Go to the source code of this file.

Namespaces
Macros
#define DEBUG_TYPE "vector-to-gpu"
#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
#define DBGSNL() (llvm::dbgs() << "\n")
#define GEN_PASS_DEF_CONVERTVECTORTOGPU
Functions
template
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. More...
static bool contractSupportsMMAMatrixType (vector::ContractionOp contract, bool useNvGpu)
static bool isTransposeMatrixLoadMap (AffineMap permutationMap)
static std::optional< int64_t > getStaticallyKnownRowStride (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. More...
static bool broadcastSupportsMMAMatrixType (vector::BroadcastOp broadcastOp)
Return true if this is a broadcast from scalar to a 2D vector. More...
template
static bool integerExtendSupportsMMAMatrixType (ExtOpTy extOp)
Return true if this integer extend op can be folded into a contract op. More...
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. More...
static bool elementwiseSupportsMMAMatrixType (Operation *op)
Return true if the op is supported as elementwise op on MMAMatrix type. More...
static bool extractStridedSliceSupportsMMAMatrixType (vector::ExtractStridedSliceOp op)
Returns true if the extract strided slice op is supported with mma.sync path. More...
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. More...
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. More...
static LogicalResult convertConstantOpMmaSync (RewriterBase &rewriter, arith::ConstantOp op, llvm::DenseMap< Value, Value > &valueMapping)
Convert a 2D splat ConstantOp to a SubgroupMmaConstantMatrix op. More...
static FailureOr< bool > isTransposed (vector::TransferReadOp op)
Check if the loaded matrix operand requires transposed. More...
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. More...
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. More...
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. More...
static LogicalResult convertBroadcastOp (RewriterBase &rewriter, vector::BroadcastOp op, llvm::DenseMap< Value, Value > &valueMapping)
Convert a vector.broadcast from scalar to a SubgroupMmaConstantMatrix op. More...
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. More...

◆ DBGS

#define DBGS ( ) (llvm::dbgs() << "[" DEBUG_TYPE "]: ")

◆ DBGSNL

#define DBGSNL ( ) (llvm::dbgs() << "\n")

◆ DEBUG_TYPE

#define DEBUG_TYPE "vector-to-gpu"

◆ GEN_PASS_DEF_CONVERTVECTORTOGPU

#define GEN_PASS_DEF_CONVERTVECTORTOGPU

◆ broadcastSupportsMMAMatrixType()

static bool broadcastSupportsMMAMatrixType ( vector::BroadcastOp broadcastOp) static

◆ constantSupportsMMAMatrixType()

static bool constantSupportsMMAMatrixType ( arith::ConstantOp constantOp) static

◆ contractSupportsMMAMatrixType()

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

◆ convertBroadcastOp()

◆ convertConstantOp()

◆ convertConstantOpMmaSync()

◆ convertContractOp()

◆ convertContractOpToMmaSync()

◆ convertElementwiseOp()

◆ convertElementwiseOpToMMA()

static std::optionalgpu::MMAElementwiseOp convertElementwiseOpToMMA ( Operation * op) static

◆ convertExtractStridedSlice()

◆ convertForOp()

◆ convertTransferReadOp()

Definition at line 545 of file VectorToGPU.cpp.

References mlir::OpBuilder::create(), DBGS, mlir::gpu::MMAMatrixType::get(), mlir::get(), mlir::Builder::getIndexAttr(), mlir::AffineMap::getPermutationMap(), mlir::AffineMap::getResult(), getStaticallyKnownRowStride(), mlir::Builder::getUnitAttr(), inferFragType(), isTransposeMatrixLoadMap(), mlir::RewriterBase::notifyMatchFailure(), mlir::OpBuilder::setInsertionPoint(), mlir::quant::QuantizationFlags::Signed, and transferReadSupportsMMAMatrixType().

Referenced by mlir::convertVectorToMMAOps().

◆ convertTransferReadToLoads()

◆ convertTransferWriteOp()

◆ convertTransferWriteToStores()

Definition at line 902 of file VectorToGPU.cpp.

References mlir::OpBuilder::create(), DBGS, mlir::RewriterBase::eraseOp(), mlir::Builder::getIndexAttr(), mlir::Builder::getIndexType(), mlir::nvgpu::getLaneIdAndValueIdToOperandCoord(), mlir::nvgpu::getMmaSyncRegisterType(), getMmaSyncVectorOperandType(), mlir::nvgpu::getWarpMatrixInfo(), mlir::RewriterBase::notifyMatchFailure(), and mlir::OpBuilder::setInsertionPoint().

Referenced by mlir::convertVectorToNVVMCompatibleMMASync().

◆ convertYieldOp()

◆ createNonLdMatrixLoads()

Definition at line 771 of file VectorToGPU.cpp.

References mlir::OpBuilder::create(), mlir::get(), mlir::Builder::getIndexAttr(), mlir::Builder::getIndexType(), mlir::nvgpu::getLaneIdAndValueIdToOperandCoord(), mlir::nvgpu::getMmaSyncRegisterType(), getMmaSyncVectorOperandType(), mlir::nvgpu::getWarpMatrixInfo(), mlir::Builder::getZeroAttr(), mlir::RewriterBase::notifyMatchFailure(), and mlir::OpBuilder::setInsertionPoint().

Referenced by convertTransferReadToLoads().

◆ creatLdMatrixCompatibleLoads()

Definition at line 709 of file VectorToGPU.cpp.

References mlir::OpBuilder::create(), DBGS, mlir::nvgpu::getLaneIdToLdMatrixMatrixCoord(), mlir::nvgpu::getLdMatrixParams(), mlir::nvgpu::getMmaSyncRegisterType(), getMmaSyncVectorOperandType(), mlir::nvgpu::getWarpMatrixInfo(), isTransposed(), mlir::RewriterBase::notifyMatchFailure(), mlir::OpBuilder::setInsertionPoint(), and mlir::xegpu::transpose().

Referenced by convertTransferReadToLoads().

◆ elementwiseSupportsMMAMatrixType()

static bool elementwiseSupportsMMAMatrixType ( Operation * op) static

◆ extractStridedSliceSupportsMMAMatrixType()

static bool extractStridedSliceSupportsMMAMatrixType ( vector::ExtractStridedSliceOp op) static

◆ fpExtendSupportsMMAMatrixType()

static bool fpExtendSupportsMMAMatrixType ( arith::ExtFOp extOp) static

◆ getMmaSyncVectorOperandType()

◆ getOpToConvert()

◆ getSliceContract()

◆ getStaticallyKnownRowStride()

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

◆ getXferIndices()

template

◆ inferFragType()

static const char* inferFragType ( Operation * op) static

◆ integerExtendSupportsMMAMatrixType()

template

static bool integerExtendSupportsMMAMatrixType ( ExtOpTy extOp) static

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

Definition at line 202 of file VectorToGPU.cpp.

◆ isSharedMemory()

static bool isSharedMemory ( MemRefType type) static

◆ isTransposed()

static FailureOr 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 682 of file VectorToGPU.cpp.

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

Referenced by creatLdMatrixCompatibleLoads().

◆ isTransposeMatrixLoadMap()

static bool isTransposeMatrixLoadMap ( AffineMap permutationMap) static

◆ populateFromInt64AttrArray()

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

◆ replaceForOpWithNewSignature()

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

◆ supportsMMaMatrixType()

static bool supportsMMaMatrixType ( Operation * op, bool useNvGpu ) static

Definition at line 277 of file VectorToGPU.cpp.

References broadcast(), broadcastSupportsMMAMatrixType(), mlir::nvgpu::canLowerToWarpMatrixOperation(), constantSupportsMMAMatrixType(), contract(), contractSupportsMMAMatrixType(), elementwiseSupportsMMAMatrixType(), extractStridedSliceSupportsMMAMatrixType(), fpExtendSupportsMMAMatrixType(), transferReadSupportsMMAMatrixType(), and transferWriteSupportsMMAMatrixType().

Referenced by getOpToConvert().

◆ transferReadSupportsMMAMatrixType()

static bool transferReadSupportsMMAMatrixType ( vector::TransferReadOp readOp) static

◆ transferWriteSupportsMMAMatrixType()

static bool transferWriteSupportsMMAMatrixType ( vector::TransferWriteOp writeOp) static