MLIR  15.0.0git
Functions
VectorToGPU.cpp File Reference
#include <type_traits>
#include "NvGpuSupport.h"
#include "mlir/Conversion/VectorToGPU/VectorToGPU.h"
#include "../PassDetail.h"
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
#include "mlir/IR/Builders.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/Passes.h"
#include "llvm/ADT/TypeSwitch.h"
+ Include dependency graph for VectorToGPU.cpp:

Go to the source code of this file.

Functions

template<typename TransferOpType >
static void getXferIndices (OpBuilder &b, 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 llvm::Optional< int64_t > getMemrefConstantHorizontalStride (ShapedType type)
 
static bool transferReadSupportsMMAMatrixType (vector::TransferReadOp readOp, bool useNvGpu)
 
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...
 
static llvm::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 supportsMMaMatrixType (Operation *op, bool useNvGpu)
 
static SetVector< Operation * > getSliceContract (Operation *op, TransitiveFilter backwardFilter, TransitiveFilter forwardFilter)
 Return an unsorted slice handling scf.for region differently than getSlice. More...
 
static SetVector< Operation * > getOpToConvert (mlir::Operation *op, bool useNvGpu)
 
template<typename OpTy >
static const char * inferFragType (OpTy op)
 
static void convertTransferReadOp (vector::TransferReadOp op, llvm::DenseMap< Value, Value > &valueMapping)
 
static void convertTransferWriteOp (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 (arith::ConstantOp op, llvm::DenseMap< Value, Value > &valueMapping)
 Convert a 2D splat ConstantOp to a SubgroupMmaConstantMatrix op. More...
 
static LogicalResult creatLdMatrixCompatibleLoads (vector::TransferReadOp op, OpBuilder &builder, llvm::DenseMap< Value, Value > &valueMapping)
 
static LogicalResult createNonLdMatrixLoads (vector::TransferReadOp op, OpBuilder &builder, llvm::DenseMap< Value, Value > &valueMapping)
 
static LogicalResult convertTransferReadToLoads (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 (vector::TransferWriteOp op, llvm::DenseMap< Value, Value > &valueMapping)
 
static void convertContractOp (vector::ContractionOp op, llvm::DenseMap< Value, Value > &valueMapping)
 
static LogicalResult convertContractOpToMmaSync (vector::ContractionOp op, llvm::DenseMap< Value, Value > &valueMapping)
 
static void convertConstantOp (arith::ConstantOp op, llvm::DenseMap< Value, Value > &valueMapping)
 Convert a 2D splat ConstantOp to a SubgroupMmaConstantMatrix op. More...
 
static void convertBroadcastOp (vector::BroadcastOp op, llvm::DenseMap< Value, Value > &valueMapping)
 Convert a vector.broadcast from scalar to a SubgroupMmaConstantMatrix op. More...
 
static scf::ForOp replaceForOpWithNewSignature (OpBuilder &b, scf::ForOp loop, ValueRange newIterOperands)
 
static void convertForOp (scf::ForOp op, llvm::DenseMap< Value, Value > &valueMapping)
 
static void convertYieldOp (scf::YieldOp op, llvm::DenseMap< Value, Value > &valueMapping)
 
static void convertElementwiseOp (Operation *op, gpu::MMAElementwiseOp opType, llvm::DenseMap< Value, Value > &valueMapping)
 Convert an elementwise op to the equivalent elementwise op on MMA matrix. More...
 

Function Documentation

◆ broadcastSupportsMMAMatrixType()

static bool broadcastSupportsMMAMatrixType ( vector::BroadcastOp  broadcastOp)
static

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

Definition at line 162 of file VectorToGPU.cpp.

Referenced by convertBroadcastOp(), and supportsMMaMatrixType().

◆ constantSupportsMMAMatrixType()

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

Referenced by convertConstantOp(), and supportsMMaMatrixType().

◆ contractSupportsMMAMatrixType()

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

◆ convertBroadcastOp()

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

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

Definition at line 714 of file VectorToGPU.cpp.

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

Referenced by mlir::convertVectorToMMAOps().

◆ convertConstantOp()

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

Convert a 2D splat ConstantOp to a SubgroupMmaConstantMatrix op.

Definition at line 696 of file VectorToGPU.cpp.

References mlir::Attribute::cast(), constantSupportsMMAMatrixType(), mlir::OpBuilder::create(), mlir::gpu::MMAMatrixType::get(), mlir::Attribute::getType(), and inferFragType().

Referenced by mlir::convertVectorToMMAOps().

◆ convertConstantOpMmaSync()

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

◆ convertContractOp()

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

Definition at line 668 of file VectorToGPU.cpp.

References mlir::OpBuilder::create(), and mlir::Value::getType().

Referenced by mlir::convertVectorToMMAOps().

◆ convertContractOpToMmaSync()

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

◆ convertElementwiseOp()

static void convertElementwiseOp ( Operation op,
gpu::MMAElementwiseOp  opType,
llvm::DenseMap< Value, Value > &  valueMapping 
)
static

Convert an elementwise op to the equivalent elementwise op on MMA matrix.

Definition at line 796 of file VectorToGPU.cpp.

References mlir::OpBuilder::create(), mlir::Operation::getLoc(), mlir::Operation::getOperands(), and mlir::Operation::getResult().

Referenced by mlir::convertVectorToMMAOps().

◆ convertElementwiseOpToMMA()

static llvm::Optional<gpu::MMAElementwiseOp> convertElementwiseOpToMMA ( Operation op)
static

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

Return llvm::None otherwise.

Definition at line 170 of file VectorToGPU.cpp.

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

◆ convertForOp()

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

◆ convertTransferReadOp()

static void convertTransferReadOp ( vector::TransferReadOp  op,
llvm::DenseMap< Value, Value > &  valueMapping 
)
static

◆ convertTransferReadToLoads()

static LogicalResult convertTransferReadToLoads ( 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 600 of file VectorToGPU.cpp.

References createNonLdMatrixLoads(), creatLdMatrixCompatibleLoads(), mlir::failed(), mlir::failure(), mlir::detail::getMemorySpaceAsInt(), mlir::nvgpu::getWarpMatrixInfo(), and mlir::nvgpu::inferTileWidthInBits().

Referenced by mlir::convertVectorToNVVMCompatibleMMASync().

◆ convertTransferWriteOp()

static void convertTransferWriteOp ( vector::TransferWriteOp  op,
llvm::DenseMap< Value, Value > &  valueMapping 
)
static

◆ convertTransferWriteToStores()

static LogicalResult convertTransferWriteToStores ( vector::TransferWriteOp  op,
llvm::DenseMap< Value, Value > &  valueMapping 
)
static

◆ convertYieldOp()

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

◆ createNonLdMatrixLoads()

static LogicalResult createNonLdMatrixLoads ( vector::TransferReadOp  op,
OpBuilder builder,
llvm::DenseMap< Value, Value > &  valueMapping 
)
static

◆ creatLdMatrixCompatibleLoads()

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

◆ elementwiseSupportsMMAMatrixType()

static bool elementwiseSupportsMMAMatrixType ( Operation op)
static

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

Definition at line 185 of file VectorToGPU.cpp.

References convertElementwiseOpToMMA().

Referenced by supportsMMaMatrixType().

◆ getMemrefConstantHorizontalStride()

static llvm::Optional<int64_t> getMemrefConstantHorizontalStride ( ShapedType  type)
static

◆ getMmaSyncVectorOperandType()

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

◆ getOpToConvert()

static SetVector<Operation *> getOpToConvert ( mlir::Operation op,
bool  useNvGpu 
)
static

◆ getSliceContract()

static SetVector<Operation *> getSliceContract ( Operation op,
TransitiveFilter  backwardFilter,
TransitiveFilter  forwardFilter 
)
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 208 of file VectorToGPU.cpp.

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

Referenced by getOpToConvert().

◆ getXferIndices()

template<typename TransferOpType >
static void getXferIndices ( OpBuilder b,
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 43 of file VectorToGPU.cpp.

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

◆ inferFragType()

template<typename OpTy >
static const char* inferFragType ( OpTy  op)
static

◆ replaceForOpWithNewSignature()

static scf::ForOp replaceForOpWithNewSignature ( OpBuilder b,
scf::ForOp  loop,
ValueRange  newIterOperands 
)
static

◆ supportsMMaMatrixType()

static bool supportsMMaMatrixType ( Operation op,
bool  useNvGpu 
)
static

◆ transferReadSupportsMMAMatrixType()

static bool transferReadSupportsMMAMatrixType ( vector::TransferReadOp  readOp,
bool  useNvGpu 
)
static

◆ transferWriteSupportsMMAMatrixType()

static bool transferWriteSupportsMMAMatrixType ( vector::TransferWriteOp  writeOp)
static