15 #include <type_traits>
36 #include "llvm/ADT/STLExtras.h"
37 #include "llvm/ADT/TypeSwitch.h"
39 #define DEBUG_TYPE "vector-to-gpu"
40 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
41 #define DBGSNL() (llvm::dbgs() << "\n")
44 #define GEN_PASS_DEF_CONVERTVECTORTOGPU
45 #include "mlir/Conversion/Passes.h.inc"
56 template <
typename TransferOpType>
60 indices.append(xferOp.getIndices().begin(), xferOp.getIndices().end());
62 unsigned offsetsIdx = 0;
63 for (
auto expr : xferOp.getPermutationMap().getResults()) {
64 if (
auto dim = dyn_cast<AffineDimExpr>(expr)) {
65 Value prevIdx = indices[dim.getPosition()];
67 dims.push_back(prevIdx);
70 rewriter, loc, d0 + offsetMap.
getResult(offsetsIdx++), dims);
80 auto infer = [&](MapList m) {
85 auto iteratorTypes =
contract.getIteratorTypes().getValue();
94 contract.getIndexingMapsArray() != infer({{m, k}, {k, n}, {m, n}}))
97 contract.getIndexingMapsArray() != infer({{m, k}, {n, k}, {m, n}}))
120 return permutationMap ==
AffineMap::get(nDim, 0, {innerDim, outerDim}, ctx) ||
127 auto memrefType = dyn_cast<MemRefType>(type);
131 if (memrefType.getRank() < 2)
138 int64_t stride = strides[strides.size() - 2];
139 if (stride == ShapedType::kDynamic)
146 if (readOp.getMask() || readOp.hasOutOfBoundsDim() ||
147 readOp.getVectorType().getRank() != 2)
153 if (readOp.getVectorType().getElementType().isInteger(8))
154 if (!readOp->hasOneUse() || (!isa<arith::ExtSIOp>(*readOp->user_begin()) &&
155 !isa<arith::ExtUIOp>(*readOp->user_begin())))
162 auto broadcastInnerDim =
172 if (writeOp.getTransferRank() == 0)
175 if (writeOp.getMask() || writeOp.hasOutOfBoundsDim() ||
176 writeOp.getVectorType().getRank() != 2)
181 if (!writeOp.getPermutationMap().isMinorIdentity())
189 auto vecType = dyn_cast<VectorType>(constantOp.getType());
190 if (!vecType || vecType.getRank() != 2)
192 return isa<SplatElementsAttr>(constantOp.getValue());
197 return broadcastOp.getResultVectorType().getRank() == 2;
201 template <
typename ExtOpTy>
203 if (!isa<vector::TransferReadOp>(extOp.getOperand().getDefiningOp()))
205 return llvm::all_of(extOp->getUsers(), llvm::IsaPred<vector::ContractionOp>);
212 static std::optional<gpu::MMAElementwiseOp>
214 if (isa<arith::AddFOp>(op))
215 return gpu::MMAElementwiseOp::ADDF;
216 if (isa<arith::MulFOp>(op))
217 return gpu::MMAElementwiseOp::MULF;
218 if (isa<arith::SubFOp>(op))
219 return gpu::MMAElementwiseOp::SUBF;
220 if (isa<arith::MaximumFOp>(op))
221 return gpu::MMAElementwiseOp::MAXF;
222 if (isa<arith::MinimumFOp>(op))
223 return gpu::MMAElementwiseOp::MINF;
224 if (isa<arith::DivFOp>(op))
225 return gpu::MMAElementwiseOp::DIVF;
226 if (isa<arith::AddIOp>(op))
228 if (isa<arith::MulIOp>(op))
230 if (isa<arith::SubIOp>(op))
232 if (isa<arith::DivSIOp>(op))
233 return gpu::MMAElementwiseOp::DIVS;
234 if (isa<arith::DivUIOp>(op))
235 return gpu::MMAElementwiseOp::DIVU;
236 if (isa<arith::NegFOp>(op))
237 return gpu::MMAElementwiseOp::NEGATEF;
238 if (isa<arith::ExtFOp>(op))
239 return gpu::MMAElementwiseOp::EXTF;
253 FailureOr<nvgpu::WarpMatrixInfo> warpMatrixInfo =
255 if (failed(warpMatrixInfo))
259 if (failed(contractOp))
267 cast<VectorType>((*contractOp).getRhs().getType()));
270 cast<VectorType>((*contractOp).getAcc().getType()));
276 if (isa<scf::ForOp, scf::YieldOp>(op))
278 if (
auto transferRead = dyn_cast<vector::TransferReadOp>(op))
281 if (
auto transferWrite = dyn_cast<vector::TransferWriteOp>(op))
284 if (
auto extractStridedSlice = dyn_cast<vector::ExtractStridedSliceOp>(op))
287 if (
auto contract = dyn_cast<vector::ContractionOp>(op))
289 if (
auto constant = dyn_cast<arith::ConstantOp>(op))
291 if (
auto broadcast = dyn_cast<vector::BroadcastOp>(op))
293 if (
auto signedExtend = dyn_cast<arith::ExtSIOp>(op))
294 return integerExtendSupportsMMAMatrixType<arith::ExtSIOp>(signedExtend);
295 if (
auto unsignedExtend = dyn_cast<arith::ExtUIOp>(op))
296 return integerExtendSupportsMMAMatrixType<arith::ExtUIOp>(unsignedExtend);
297 if (
auto fpExtend = dyn_cast<arith::ExtFOp>(op))
311 unsigned currentIndex = 0;
314 while (currentIndex != slice.size()) {
315 auto *currentOp = (slice)[currentIndex];
317 backwardSlice.clear();
319 slice.insert(backwardSlice.begin(), backwardSlice.end());
322 forwardSlice.clear();
327 if (
auto forOp = dyn_cast<scf::ForOp>(currentOp)) {
328 for (
Value forOpResult : forOp.getResults())
335 slice.insert(forwardSlice.begin(), forwardSlice.end());
346 return llvm::any_of(op->
getResultTypes(), llvm::IsaPred<VectorType>);
349 backwardSliceOptions.
filter = hasVectorDest;
355 forwardSliceOptions.
filter = hasVectorSrc;
359 if (opToConvert.contains(
contract.getOperation()))
366 if (llvm::any_of(dependentOps, [useNvGpu](
Operation *op) {
368 LLVM_DEBUG(
DBGS() <<
"cannot convert op: " << *op <<
"\n");
375 opToConvert.insert(dependentOps.begin(), dependentOps.end());
384 struct PrepareContractToGPUMMA
388 LogicalResult matchAndRewrite(vector::ContractionOp op,
391 Value lhs = op.getLhs(), rhs = op.getRhs(), res = op.getAcc();
395 auto infer = [&](MapList m) {
400 static constexpr std::array<int64_t, 2> perm = {1, 0};
401 auto iteratorTypes = op.getIteratorTypes().getValue();
411 if (maps == infer({{m, k}, {k, n}, {m, n}}))
413 if (maps == infer({{m, k}, {n, k}, {m, n}})) {
414 rhs = rewriter.
create<vector::TransposeOp>(loc, rhs, perm);
415 }
else if (maps == infer({{k, m}, {k, n}, {m, n}})) {
416 lhs = rewriter.
create<vector::TransposeOp>(loc, lhs, perm);
417 }
else if (maps == infer({{k, m}, {n, k}, {m, n}})) {
418 rhs = rewriter.
create<vector::TransposeOp>(loc, rhs, perm);
419 lhs = rewriter.
create<vector::TransposeOp>(loc, lhs, perm);
420 }
else if (maps == infer({{m, k}, {k, n}, {n, m}})) {
422 rhs = rewriter.
create<vector::TransposeOp>(loc, rhs, perm);
423 lhs = rewriter.
create<vector::TransposeOp>(loc, lhs, perm);
424 }
else if (maps == infer({{m, k}, {n, k}, {n, m}})) {
426 rhs = rewriter.
create<vector::TransposeOp>(loc, rhs, perm);
427 }
else if (maps == infer({{k, m}, {k, n}, {n, m}})) {
429 lhs = rewriter.
create<vector::TransposeOp>(loc, lhs, perm);
430 }
else if (maps == infer({{k, m}, {n, k}, {n, m}})) {
439 op.getIteratorTypes());
448 struct CombineTransferReadOpTranspose final
452 LogicalResult matchAndRewrite(vector::TransposeOp op,
455 Value source = op.getVector();
456 Type resultType = op.getType();
464 cast<VectorType>(source.
getType()).getElementType());
467 auto transferReadOp = source.
getDefiningOp<vector::TransferReadOp>();
472 if (transferReadOp.getTransferRank() == 0)
475 if (transferReadOp.getMask() || transferReadOp.hasOutOfBoundsDim())
481 permutationMap.
compose(transferReadOp.getPermutationMap());
486 .
create<vector::TransferReadOp>(
487 loc, resultType, transferReadOp.getSource(),
489 transferReadOp.getPadding(), transferReadOp.getMask(),
490 transferReadOp.getInBoundsAttr())
495 if (isa<arith::ExtSIOp>(extOp))
496 result = rewriter.
create<arith::ExtSIOp>(loc, op.getType(), result)
498 else if (isa<arith::ExtUIOp>(extOp))
499 result = rewriter.
create<arith::ExtUIOp>(loc, op.getType(), result)
502 result = rewriter.
create<arith::ExtFOp>(loc, op.getType(), result)
527 auto contract = dyn_cast<vector::ContractionOp>(users);
545 assert(op.getTransferRank() > 0 &&
"unexpected 0-d transfer");
547 "expected convertible operation");
549 std::optional<int64_t> stride =
551 if (!stride.has_value()) {
552 LLVM_DEBUG(
DBGS() <<
"no stride\n");
560 if (
auto cstExpr = dyn_cast<AffineConstantExpr>(map.
getResult(isTranspose))) {
561 assert(cstExpr.getValue() == 0);
566 auto elType = op.getVectorType().getElementType();
571 if (isa<arith::ExtSIOp, arith::ExtUIOp>(user)) {
573 op.
getContext(), cast<IntegerType>(elType).getWidth(),
575 : IntegerType::Unsigned);
576 mappingResult = user->getResult(0);
581 Value load = rewriter.
create<gpu::SubgroupMmaLoadMatrixOp>(
582 op.
getLoc(), type, op.getSource(), op.getIndices(),
584 isTranspose ? rewriter.
getUnitAttr() : UnitAttr());
585 valueMapping[mappingResult] = load;
587 LLVM_DEBUG(
DBGS() <<
"transfer read to: " << load <<
"\n");
598 std::optional<int64_t> stride =
600 if (!stride.has_value()) {
601 LLVM_DEBUG(
DBGS() <<
"no stride\n");
605 auto it = valueMapping.find(op.getVector());
606 if (it == valueMapping.end()) {
607 LLVM_DEBUG(
DBGS() <<
"no mapping\n");
611 Value matrix = it->second;
612 auto store = rewriter.
create<gpu::SubgroupMmaStoreMatrixOp>(
613 op.
getLoc(), matrix, op.getSource(), op.getIndices(),
617 LLVM_DEBUG(
DBGS() <<
"transfer write to: " << store <<
"\n");
619 LLVM_DEBUG(
DBGS() <<
"erase: " << op <<
"\n");
630 if (
auto vecType = dyn_cast<VectorType>(elType))
631 elType = vecType.getElementType();
642 FailureOr<nvgpu::WarpMatrixInfo> warpMatrixInfo =
644 if (failed(warpMatrixInfo)) {
645 LLVM_DEBUG(
DBGS() <<
"no warpMatrixInfo\n");
649 FailureOr<nvgpu::FragmentElementInfo> regInfo =
651 if (failed(regInfo)) {
652 LLVM_DEBUG(
DBGS() <<
"not mma sync reg info\n");
657 auto dense = dyn_cast<SplatElementsAttr>(op.getValue());
659 LLVM_DEBUG(
DBGS() <<
"not a splat\n");
681 LLVM_DEBUG(
DBGS() <<
"Failed because the result of `vector.transfer_read` "
682 "is not a 2d operand\n");
691 auto exprM = dyn_cast<AffineDimExpr>(dM);
692 auto exprN = dyn_cast<AffineDimExpr>(dN);
694 if (!exprM || !exprN) {
695 LLVM_DEBUG(
DBGS() <<
"Failed because expressions are not affine dim "
696 "expressions, then transpose cannot be determined.\n");
700 return exprM.getPosition() > exprN.getPosition();
710 FailureOr<nvgpu::WarpMatrixInfo> warpMatrixInfo =
712 if (failed(warpMatrixInfo)) {
713 LLVM_DEBUG(
DBGS() <<
"no warpMatrixInfo\n");
717 FailureOr<nvgpu::FragmentElementInfo> regInfo =
719 if (failed(regInfo)) {
720 LLVM_DEBUG(
DBGS() <<
"not mma sync reg info\n");
726 LLVM_DEBUG(
DBGS() <<
"failed to determine the transpose\n");
728 op,
"Op should likely not be converted to a nvgpu.ldmatrix call.");
731 FailureOr<nvgpu::LdMatrixParams> params =
734 if (failed(params)) {
737 <<
"failed to convert vector.transfer_read to ldmatrix. "
738 <<
"Op should likely not be converted to a nvgpu.ldmatrix call.\n");
740 op,
"failed to convert vector.transfer_read to ldmatrix; this op "
741 "likely should not be converted to a nvgpu.ldmatrix call.");
745 auto laneId = rewriter.
create<gpu::LaneIdOp>(loc,
nullptr);
746 FailureOr<AffineMap> offsets =
748 if (failed(offsets)) {
749 LLVM_DEBUG(
DBGS() <<
"no offsets\n");
756 getXferIndices<vector::TransferReadOp>(rewriter, op, *offsets, {laneId},
759 nvgpu::LdMatrixOp newOp = rewriter.
create<nvgpu::LdMatrixOp>(
760 loc, vectorType, op.getSource(), indices, *
transpose, params->numTiles);
761 valueMapping[op] = newOp->getResult(0);
772 FailureOr<nvgpu::WarpMatrixInfo> warpMatrixInfo =
774 if (failed(warpMatrixInfo))
776 FailureOr<nvgpu::FragmentElementInfo> regInfo =
778 if (failed(regInfo)) {
780 op,
"Failed to deduce register fragment type during "
781 "conversion to distributed non-ldmatrix compatible load");
784 Value laneId = rewriter.
create<gpu::LaneIdOp>(loc,
nullptr);
788 Type loadedElType = regInfo->registerLLVMType;
792 op.
getLoc(), vectorType.getElementType(),
793 rewriter.
getZeroAttr(vectorType.getElementType()));
795 rewriter.
create<vector::SplatOp>(op.
getLoc(), fill, vectorType);
797 bool isTransposeLoad = !op.getPermutationMap().isMinorIdentity();
801 if (!isTransposeLoad) {
802 if (!isa<VectorType>(loadedElType)) {
806 for (
int i = 0; i < vectorType.getShape()[0]; i++) {
808 rewriter, op.
getLoc(), *warpMatrixInfo);
812 Value logicalValueId = rewriter.
create<arith::ConstantOp>(
814 rewriter.
getIndexAttr(i * regInfo->elementsPerRegister));
816 getXferIndices<vector::TransferReadOp>(
817 rewriter, op, *coords, {laneId, logicalValueId}, newIndices);
819 Value el = rewriter.
create<vector::LoadOp>(loc, loadedElType,
820 op.getSource(), newIndices);
821 result = rewriter.
create<vector::InsertOp>(loc, el, result, i);
824 if (
auto vecType = dyn_cast<VectorType>(loadedElType)) {
825 loadedElType = vecType.getElementType();
827 for (
int i = 0; i < vectorType.getShape()[0]; i++) {
828 for (
unsigned innerIdx = 0; innerIdx < vectorType.getShape()[1];
831 Value logicalValueId = rewriter.
create<arith::ConstantOp>(
833 rewriter.
getIndexAttr(i * regInfo->elementsPerRegister + innerIdx));
835 rewriter, op.
getLoc(), *warpMatrixInfo);
840 getXferIndices<vector::TransferReadOp>(
841 rewriter, op, *coords, {laneId, logicalValueId}, newIndices);
843 op.getSource(), newIndices);
844 result = rewriter.
create<vector::InsertOp>(
857 dyn_cast_or_null<gpu::AddressSpaceAttr>(type.getMemorySpace());
858 return addressSpace &&
859 addressSpace.getValue() == gpu::GPUDialect::getWorkgroupAddressSpace();
871 FailureOr<nvgpu::WarpMatrixInfo> warpMatrixInfo =
873 if (failed(warpMatrixInfo))
876 bool isLdMatrixCompatible =
880 VectorType vecTy = op.getVectorType();
881 int64_t bitWidth = vecTy.getElementType().getIntOrFloatBitWidth();
886 if (!op.getPermutationMap().isMinorIdentity() &&
887 (bitWidth != 16 || vecTy.getDimSize(1) < 8 ||
888 vecTy.getDimSize(0) * bitWidth < 128))
889 isLdMatrixCompatible =
false;
891 if (!isLdMatrixCompatible)
904 auto it = valueMapping.find(op.getVector());
905 if (it == valueMapping.end())
907 Value matrix = it->second;
909 FailureOr<nvgpu::WarpMatrixInfo> warpMatrixInfo =
911 if (failed(warpMatrixInfo))
913 FailureOr<nvgpu::FragmentElementInfo> regInfo =
919 Value laneId = rewriter.
create<gpu::LaneIdOp>(loc,
nullptr);
921 for (
unsigned i = 0; i < vectorType.getShape()[0]; i++) {
922 Value logicalValueId = rewriter.
create<arith::ConstantOp>(
924 rewriter.
getIndexAttr(i * regInfo->elementsPerRegister));
926 rewriter, op.
getLoc(), *warpMatrixInfo);
933 getXferIndices<vector::TransferWriteOp>(
934 rewriter, op, *coords, {laneId, logicalValueId}, newIndices);
935 rewriter.
create<vector::StoreOp>(loc, el, op.getSource(), newIndices);
938 LLVM_DEBUG(
DBGS() <<
"erase: " << op <<
"\n");
945 for (
auto attr : arrayAttr)
946 results.push_back(cast<IntegerAttr>(attr).getInt());
951 vector::ExtractStridedSliceOp op,
958 FailureOr<nvgpu::WarpMatrixInfo> warpMatrixInfo =
960 if (failed(warpMatrixInfo))
963 FailureOr<nvgpu::FragmentElementInfo> mmaSyncFragmentInfo =
965 if (failed(mmaSyncFragmentInfo))
969 auto transferReadOp = op.getVector().getDefiningOp<vector::TransferReadOp>();
974 if (failed(warpMatrixInfo))
977 FailureOr<nvgpu::FragmentElementInfo> ldFragmentInfo =
979 if (failed(ldFragmentInfo))
983 (mmaSyncFragmentInfo->elementsPerRegister ==
984 ldFragmentInfo->elementsPerRegister) &&
985 "Number of elements per register should be same for load and mma.sync");
988 std::array<int64_t, 2> strides = {1,
990 std::array<int64_t, 2> sliceShape = {
991 mmaSyncFragmentInfo->numRegistersPerFragment,
992 mmaSyncFragmentInfo->elementsPerRegister};
993 auto it = valueMapping.find(transferReadOp);
994 if (it == valueMapping.end())
996 auto sourceVector = it->second;
1009 std::array<int64_t, 2> sliceOffset = {0, 0};
1011 if (offsets[0] && offsets[1])
1012 return op->
emitError() <<
"Slicing fragments in 2D is not supported. ";
1014 sliceOffset[0] = (warpVectorShape[0] / offsets[0]);
1015 else if (offsets[1])
1016 sliceOffset[0] = (warpVectorShape[1] / offsets[1]);
1018 Value newOp = rewriter.
create<vector::ExtractStridedSliceOp>(
1019 loc, sourceVector, sliceOffset, sliceShape, strides);
1021 valueMapping[op] = newOp;
1025 static LogicalResult
1031 auto itA = valueMapping.find(op.getLhs());
1032 auto itB = valueMapping.find(op.getRhs());
1033 auto itC = valueMapping.find(op.getAcc());
1034 if (itA == valueMapping.end() || itB == valueMapping.end() ||
1035 itC == valueMapping.end())
1037 Value opA = itA->second, opB = itB->second, opC = itC->second;
1038 Value matmul = rewriter.
create<gpu::SubgroupMmaComputeOp>(
1039 op.
getLoc(), opC.getType(), opA, opB, opC, UnitAttr(),
1045 static LogicalResult
1051 auto itA = valueMapping.find(op.getLhs());
1052 auto itB = valueMapping.find(op.getRhs());
1053 auto itC = valueMapping.find(op.getAcc());
1054 if (itA == valueMapping.end() || itB == valueMapping.end() ||
1055 itC == valueMapping.end())
1057 Value opA = itA->second, opB = itB->second, opC = itC->second;
1058 int64_t m = cast<VectorType>(op.getLhs().getType()).getShape()[0];
1059 int64_t n = cast<VectorType>(op.getRhs().getType()).getShape()[0];
1060 int64_t k = cast<VectorType>(op.getLhs().getType()).getShape()[1];
1068 static LogicalResult
1077 cast<SplatElementsAttr>(op.getValue()).getSplatValue<TypedAttr>();
1078 auto scalarConstant =
1079 rewriter.
create<arith::ConstantOp>(op.
getLoc(), splat.getType(), splat);
1081 auto vecType = cast<VectorType>(op.getType());
1083 vecType.getShape(), vecType.getElementType(), llvm::StringRef(fragType));
1084 auto matrix = rewriter.
create<gpu::SubgroupMmaConstantMatrixOp>(
1085 op.
getLoc(), type, scalarConstant);
1091 static LogicalResult
1100 auto vecType = op.getResultVectorType();
1102 vecType.getShape(), vecType.getElementType(), llvm::StringRef(fragType));
1103 auto matrix = rewriter.
create<gpu::SubgroupMmaConstantMatrixOp>(
1104 op.
getLoc(), type, op.getSource());
1119 auto operands = llvm::to_vector<4>(loop.getInitArgs());
1120 llvm::append_range(operands, newInitArgs);
1121 scf::ForOp newLoop = rewriter.
create<scf::ForOp>(
1122 loop.getLoc(), loop.getLowerBound(), loop.getUpperBound(), loop.getStep(),
1126 newLoop.getRegion().getBlocks().splice(
1127 newLoop.getRegion().getBlocks().begin(), loop.getRegion().getBlocks());
1128 for (
Value operand : newInitArgs)
1129 newLoop.getBody()->addArgument(operand.getType(), operand.getLoc());
1131 for (
auto it : llvm::zip(loop.getResults(), newLoop.getResults().take_front(
1132 loop.getNumResults())))
1135 LLVM_DEBUG(
DBGS() <<
"newLoop now: " << newLoop <<
"\n");
1136 LLVM_DEBUG(
DBGS() <<
"stripped scf.for: " << loop <<
"\n");
1137 LLVM_DEBUG(
DBGS() <<
"erase: " << loop);
1151 auto it = valueMapping.find(operand.value());
1152 if (it == valueMapping.end()) {
1153 LLVM_DEBUG(
DBGS() <<
"no value mapping for: " << operand.value() <<
"\n");
1156 argMapping.push_back(std::make_pair(
1157 operand.index(), op.getInitArgs().size() + newOperands.size()));
1158 newOperands.push_back(it->second);
1162 Block &loopBody = *newForOp.getBody();
1163 for (
auto mapping : argMapping) {
1164 valueMapping[newForOp.getResult(mapping.first)] =
1165 newForOp.getResult(mapping.second);
1167 newForOp.getNumInductionVars())] =
1168 loopBody.
getArgument(mapping.second + newForOp.getNumInductionVars());
1171 LLVM_DEBUG(
DBGS() <<
"scf.for to: " << newForOp <<
"\n");
1175 static LogicalResult
1182 auto yieldOperands = llvm::to_vector<4>(op.
getOperands());
1184 auto it = valueMapping.find(operand.value());
1185 if (it == valueMapping.end())
1189 yieldOperands[operand.index()] = loop.getInitArgs()[operand.index()];
1190 yieldOperands.push_back(it->second);
1192 rewriter.
create<scf::YieldOp>(op.
getLoc(), yieldOperands);
1194 LLVM_DEBUG(
DBGS() <<
"erase: " << op <<
"\n");
1200 static LogicalResult
1202 gpu::MMAElementwiseOp opType,
1209 auto it = valueMapping.find(operand);
1210 if (it == valueMapping.end())
1212 matrixOperands.push_back(it->second);
1214 auto resultType = cast<gpu::MMAMatrixType>(matrixOperands[0].
getType());
1215 if (opType == gpu::MMAElementwiseOp::EXTF) {
1219 vectorType.getElementType(),
1220 resultType.getOperand());
1223 Value newOp = rewriter.
create<gpu::SubgroupMmaElementwiseOp>(
1224 op->
getLoc(), resultType, matrixOperands, opType);
1232 patterns.
add<PrepareContractToGPUMMA, CombineTransferReadOpTranspose>(
1237 patterns.
add<CombineTransferReadOpTranspose>(patterns.
getContext());
1245 auto globalRes = LogicalResult::success();
1247 LLVM_DEBUG(
DBGS() <<
"Process op: " << *op <<
"\n");
1249 auto res = LogicalResult::success();
1250 if (
auto transferRead = dyn_cast<vector::TransferReadOp>(op)) {
1252 }
else if (
auto transferWrite = dyn_cast<vector::TransferWriteOp>(op)) {
1254 }
else if (
auto contractOp = dyn_cast<vector::ContractionOp>(op)) {
1256 }
else if (
auto constantOp = dyn_cast<arith::ConstantOp>(op)) {
1258 }
else if (
auto broadcastOp = dyn_cast<vector::BroadcastOp>(op)) {
1260 }
else if (
auto forOp = dyn_cast<scf::ForOp>(op)) {
1262 }
else if (
auto yieldOp = dyn_cast<scf::YieldOp>(op)) {
1268 globalRes = failure();
1279 .Case([&](vector::TransferReadOp transferReadOp) {
1283 .Case([&](vector::TransferWriteOp transferWriteOp) {
1287 .Case([&](vector::ExtractStridedSliceOp extractStridedSliceOp) {
1291 .Case([&](vector::ContractionOp contractionOp) {
1295 .Case([&](scf::ForOp forOp) {
1298 .Case([&](scf::YieldOp yieldOp) {
1301 .Case([&](arith::ConstantOp constOp) {
1305 return op->
emitError() <<
"unhandled vector to mma type: " << *op;
1309 <<
"failed to convert op during vector-to-nvgpu conversion";
1317 struct ConvertVectorToGPUPass
1318 :
public impl::ConvertVectorToGPUBase<ConvertVectorToGPUPass> {
1320 explicit ConvertVectorToGPUPass(
bool useNvGpu_) {
1321 useNvGpu.setValue(useNvGpu_);
1324 void runOnOperation()
override {
1329 return signalPassFailure();
1335 return signalPassFailure();
1345 return std::make_unique<ConvertVectorToGPUPass>(useNvGpu);
static MLIRContext * getContext(OpFoldResult val)
static void contract(RootOrderingGraph &graph, ArrayRef< Value > cycle, const DenseMap< Value, unsigned > &parentDepths, DenseMap< Value, Value > &actualSource, DenseMap< Value, Value > &actualTarget)
Contracts the specified cycle in the given graph in-place.
static Value broadcast(Location loc, Value toBroadcast, unsigned numElements, const TypeConverter &typeConverter, ConversionPatternRewriter &rewriter)
Broadcasts the value to vector with numElements number of elements.
static ArrayRef< int64_t > getShape(Type type)
Returns the shape of the given type.
static LogicalResult convertTransferWriteOp(RewriterBase &rewriter, vector::TransferWriteOp op, llvm::DenseMap< Value, Value > &valueMapping)
static LogicalResult convertForOp(RewriterBase &rewriter, scf::ForOp 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 bool isSharedMemory(MemRefType type)
Return true if this is a shared memory memref type.
static VectorType getMmaSyncVectorOperandType(const nvgpu::FragmentElementInfo ®Info)
Returns the vector type which represents a matrix fragment.
static const char * inferFragType(Operation *op)
static bool fpExtendSupportsMMAMatrixType(arith::ExtFOp extOp)
static bool supportsMMaMatrixType(Operation *op, bool useNvGpu)
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 m...
static bool contractSupportsMMAMatrixType(vector::ContractionOp contract, bool useNvGpu)
static void populateFromInt64AttrArray(ArrayAttr arrayAttr, SmallVectorImpl< int64_t > &results)
static bool isTransposeMatrixLoadMap(AffineMap permutationMap)
static SetVector< Operation * > getSliceContract(Operation *op, const BackwardSliceOptions &backwardSliceOptions, const ForwardSliceOptions &forwardSliceOptions)
Return an unsorted slice handling scf.for region differently than getSlice.
static LogicalResult convertTransferReadOp(RewriterBase &rewriter, vector::TransferReadOp op, llvm::DenseMap< Value, Value > &valueMapping)
static bool integerExtendSupportsMMAMatrixType(ExtOpTy extOp)
Return true if this integer extend op can be folded into a contract op.
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....
static LogicalResult convertBroadcastOp(RewriterBase &rewriter, vector::BroadcastOp op, llvm::DenseMap< Value, Value > &valueMapping)
Convert a vector.broadcast from scalar to a SubgroupMmaConstantMatrix op.
static LogicalResult creatLdMatrixCompatibleLoads(RewriterBase &rewriter, vector::TransferReadOp op, llvm::DenseMap< Value, Value > &valueMapping)
static FailureOr< bool > isTransposed(vector::TransferReadOp op)
Check if the loaded matrix operand requires transposed.
static scf::ForOp replaceForOpWithNewSignature(RewriterBase &rewriter, scf::ForOp loop, ValueRange newInitArgs)
static bool transferWriteSupportsMMAMatrixType(vector::TransferWriteOp writeOp)
static std::optional< gpu::MMAElementwiseOp > convertElementwiseOpToMMA(Operation *op)
Return the MMA elementwise enum associated with op if it is supported.
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.
static bool transferReadSupportsMMAMatrixType(vector::TransferReadOp readOp)
static bool extractStridedSliceSupportsMMAMatrixType(vector::ExtractStridedSliceOp op)
Returns true if the extract strided slice op is supported with mma.sync path.
static LogicalResult convertConstantOpMmaSync(RewriterBase &rewriter, arith::ConstantOp op, llvm::DenseMap< Value, Value > &valueMapping)
Convert a 2D splat ConstantOp to a SubgroupMmaConstantMatrix op.
static bool elementwiseSupportsMMAMatrixType(Operation *op)
Return true if the op is supported as elementwise op on MMAMatrix type.
static SetVector< Operation * > getOpToConvert(mlir::Operation *op, bool useNvGpu)
static LogicalResult convertConstantOp(RewriterBase &rewriter, arith::ConstantOp op, llvm::DenseMap< Value, Value > &valueMapping)
Convert a 2D splat ConstantOp to a SubgroupMmaConstantMatrix op.
static LogicalResult convertYieldOp(RewriterBase &rewriter, scf::YieldOp op, llvm::DenseMap< Value, Value > &valueMapping)
static LogicalResult convertTransferWriteToStores(RewriterBase &rewriter, vector::TransferWriteOp op, llvm::DenseMap< Value, Value > &valueMapping)
static std::optional< int64_t > getStaticallyKnownRowStride(ShapedType type)
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...
static LogicalResult convertExtractStridedSlice(RewriterBase &rewriter, vector::ExtractStridedSliceOp op, llvm::DenseMap< Value, Value > &valueMapping)
static bool broadcastSupportsMMAMatrixType(vector::BroadcastOp broadcastOp)
Return true if this is a broadcast from scalar to a 2D vector.
static LogicalResult createNonLdMatrixLoads(RewriterBase &rewriter, vector::TransferReadOp op, llvm::DenseMap< Value, Value > &valueMapping)
Base type for affine expression.
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
MLIRContext * getContext() const
bool isMinorIdentity() const
Returns true if this affine map is a minor identity, i.e.
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
unsigned getNumDims() const
unsigned getNumResults() const
AffineExpr getResult(unsigned idx) const
static AffineMap getPermutationMap(ArrayRef< unsigned > permutation, MLIRContext *context)
Returns an AffineMap representing a permutation.
AffineMap compose(AffineMap map) const
Returns the AffineMap resulting from composing this with map.
static SmallVector< AffineMap, 4 > inferFromExprList(ArrayRef< ArrayRef< AffineExpr >> exprsList, MLIRContext *context)
Returns a vector of AffineMaps; each with as many results as exprs.size(), as many dims as the larges...
Attributes are known-constant values of operations.
This class represents an argument of a Block.
Block represents an ordered list of Operations.
BlockArgument getArgument(unsigned i)
IntegerAttr getIndexAttr(int64_t value)
AffineExpr getAffineConstantExpr(int64_t constant)
TypedAttr getZeroAttr(Type type)
AffineExpr getAffineDimExpr(unsigned position)
MLIRContext * getContext() const
ArrayAttr getI64ArrayAttr(ArrayRef< int64_t > values)
ArrayAttr getAffineMapArrayAttr(ArrayRef< AffineMap > values)
static DenseElementsAttr get(ShapedType type, ArrayRef< Attribute > values)
Constructs a dense elements attribute from an array of element values.
This class coordinates rewriting a piece of IR outside of a pattern rewrite, providing a way to keep ...
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
MLIRContext is the top-level object for a collection of MLIR operations.
RAII guard to reset the insertion point of the builder when destroyed.
This class helps build Operations.
void setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Operation is the basic unit of execution within MLIR.
Value getOperand(unsigned idx)
bool hasTrait()
Returns true if the operation was registered with a particular trait, e.g.
bool hasOneUse()
Returns true if this operation has exactly one use.
OpResult getResult(unsigned idx)
Get the 'idx'th result of this operation.
std::enable_if_t< llvm::function_traits< std::decay_t< FnT > >::num_args==1, RetT > walk(FnT &&callback)
Walk the operation by calling the callback for each nested operation (including this one),...
MLIRContext * getContext()
Return the context this operation is associated with.
Location getLoc()
The source location the operation was defined or derived from.
Operation * getParentOp()
Returns the closest surrounding operation that contains this operation or nullptr if this is a top-le...
InFlightDiagnostic emitError(const Twine &message={})
Emit an error about fatal conditions with this operation, reporting up to any diagnostic handlers tha...
operand_type_range getOperandTypes()
result_type_range getResultTypes()
operand_range getOperands()
Returns an iterator on the underlying Value's.
user_range getUsers()
Returns a range of all users.
user_iterator user_begin()
InFlightDiagnostic emitOpError(const Twine &message={})
Emit an error with the op name prefixed, like "'dim' op " which is convenient for verifiers.
unsigned getNumResults()
Return the number of results held by this operation.
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
MLIRContext * getContext() const
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&...args)
Add an instance of each of the pattern types 'Ts' to the pattern list with the given arguments.
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
std::enable_if_t<!std::is_convertible< CallbackT, Twine >::value, LogicalResult > notifyMatchFailure(Location loc, CallbackT &&reasonCallback)
Used to notify the listener that the IR failed to be rewritten because of a match failure,...
virtual void eraseBlock(Block *block)
This method erases all operations in a block.
virtual void replaceOp(Operation *op, ValueRange newValues)
Replace the results of the given (original) operation with the specified list of values (replacements...
void replaceAllUsesWith(Value from, Value to)
Find uses of from and replace them with to.
virtual void eraseOp(Operation *op)
This method erases an operation that is known to have no uses.
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replace the results of the given (original) op with a new op that is created without verification (re...
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
This class provides an abstraction over the different types of ranges over Values.
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Type getType() const
Return the type of this value.
Operation * getDefiningOp() const
If this value is the result of an operation, return the operation that defines it.
MMAMatrix represents a matrix held by a subgroup for matrix-matrix multiply accumulate operations.
static MMAMatrixType get(ArrayRef< int64_t > shape, Type elementType, StringRef operand)
Get MMAMatrixType and verify construction Invariants.
AffineApplyOp makeComposedAffineApply(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands)
Returns a composed AffineApplyOp by composing map and operands with other AffineApplyOps supplying th...
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
int64_t inferTileWidthInBits(const WarpMatrixInfo &type)
Returns the number of bits in a single tile row.
FailureOr< vector::ContractionOp > getUserContract(Operation *op)
Returns the first user of the op that is vector.contract.
FailureOr< AffineMap > getLaneIdAndValueIdToOperandCoord(OpBuilder &builder, Location loc, const WarpMatrixInfo &fragmentType)
Returns an AffineMap which maps a two dimensions representing (laneId, logicalValueId) and returns tw...
FailureOr< WarpMatrixInfo > getWarpMatrixInfo(Operation *op)
If op is a vector.transfer_write, return the WarpMatrixInfo for the vector operand.
FailureOr< AffineMap > getLaneIdToLdMatrixMatrixCoord(OpBuilder &builder, Location loc, const LdMatrixParams ¶ms)
Returns an AffineMap which maps a single dimension representing the laneId to two results representin...
FailureOr< LdMatrixParams > getLdMatrixParams(const WarpMatrixInfo &type, bool transpose)
Given type that contains info for a warp-matrix operand and whether or not the load is a transposed l...
FailureOr< FragmentElementInfo > getMmaSyncRegisterType(const WarpMatrixInfo &type)
Returns a FragmentElementInfo struct describing the register types for the given matrix fragment type...
bool canLowerToWarpMatrixOperation(vector::TransferReadOp op)
Returns whether the vector.transfer_read instruction can be interpreted as a warp-level cooperative m...
bool isReductionIterator(Attribute attr)
Returns true if attr has "reduction" iterator type semantics.
bool isParallelIterator(Attribute attr)
Returns true if attr has "parallel" iterator type semantics.
void populateVectorContractCanonicalizeMatmulToMMT(RewritePatternSet &patterns, std::function< LogicalResult(vector::ContractionOp)> constraint=[](vector::ContractionOp) { return success();}, PatternBenefit=1)
Canonicalization of a vector.contraction a, b, c with row-major matmul semantics to a contraction wit...
static void transpose(llvm::ArrayRef< int64_t > trans, SmallVector< int64_t > &shape)
Include the generated interface declarations.
void populatePrepareVectorToMMAPatterns(RewritePatternSet &patterns, bool useNvGpu=false)
Patterns to transform vector ops into a canonical form to convert to MMA matrix operations.
Type getType(OpFoldResult ofr)
Returns the int type of the integer in ofr.
void bindDims(MLIRContext *ctx, AffineExprTy &...exprs)
Bind a list of AffineExpr references to DimExpr at positions: [0 .
void getBackwardSlice(Operation *op, SetVector< Operation * > *backwardSlice, const BackwardSliceOptions &options={})
Fills backwardSlice with the computed backward slice (i.e.
LogicalResult getStridesAndOffset(MemRefType t, SmallVectorImpl< int64_t > &strides, int64_t &offset)
Returns the strides of the MemRef if the layout map is in strided form.
LogicalResult convertVectorToNVVMCompatibleMMASync(RewriterBase &rewriter, Operation *rootOp)
Convert vector ops ops nested under rootOp to vector and GPU operaitons compatible with the nvvm....
LogicalResult applyPatternsAndFoldGreedily(Region ®ion, const FrozenRewritePatternSet &patterns, GreedyRewriteConfig config=GreedyRewriteConfig(), bool *changed=nullptr)
Rewrite ops in the given region, which must be isolated from above, by repeatedly applying the highes...
AffineExpr getAffineConstantExpr(int64_t constant, MLIRContext *context)
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
std::unique_ptr< Pass > createConvertVectorToGPUPass(bool useNvGpu=false)
Convert from vector to GPU ops.
AffineExpr getAffineDimExpr(unsigned position, MLIRContext *context)
These free functions allow clients of the API to not use classes in detail.
LogicalResult convertVectorToMMAOps(RewriterBase &rewriter, Operation *rootOp)
Convert vector ops to MMA matrix operations nested under rootOp.
SetVector< Operation * > topologicalSort(const SetVector< Operation * > &toSort)
Sorts all operations in toSort topologically while also considering region semantics.
void getForwardSlice(Operation *op, SetVector< Operation * > *forwardSlice, const ForwardSliceOptions &options={})
Fills forwardSlice with the computed forward slice (i.e.
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...
This trait tags element-wise ops on vectors or tensors.
Specifies information about the registers which compose a matrix fragment according to the PTX docume...
int64_t elementsPerRegister
int64_t numRegistersPerFragment