319 lines · cpp
1//===- MMAUtils.cpp - MLIR NVGPU dialect utils for MMA operations----------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8#include "mlir/Dialect/NVGPU/Utils/MMAUtils.h"9 10#include "mlir/Dialect/Arith/IR/Arith.h"11#include "mlir/Dialect/LLVMIR/NVVMDialect.h"12#include "mlir/Dialect/Vector/IR/VectorOps.h"13 14using namespace mlir;15using namespace mlir::nvgpu;16 17/// There are always 4 threads per [128|256|512] bit row.18static constexpr int64_t kThreadsPerRow = 4;19static constexpr int64_t kNumRowsPerTile = 8;20 21static bool isAccumulatorOrResult(MatMulOperandRole operandType) {22 return operandType == MatMulOperandRole::C;23}24 25/// Returns the number of registers which compose a matrix fragment held by a26/// single thread.27static int64_t inferNumRegistersPerMatrixFragment(const WarpMatrixInfo &type) {28 int64_t lineSize = inferTileWidthInBits(type);29 auto shape = type.vectorType.getShape();30 return (shape[0] / kNumRowsPerTile) *31 (shape[1] * type.vectorType.getElementType().getIntOrFloatBitWidth()) /32 lineSize;33}34 35/// Returns the number of 8 x [128|256|512] bit tiles that compose the given36/// operand shape.37static std::array<int64_t, 2> getTileShape(ArrayRef<int64_t> operandShape,38 Type elementType,39 int64_t lineSizeBits) {40 // For each 8x128bit square, a thread is responsible for one 32bit register.41 return {operandShape[0] / kNumRowsPerTile,42 (operandShape[1] * elementType.getIntOrFloatBitWidth()) /43 lineSizeBits};44}45 46/// Returns the first user of the `op` that is vector.contract. If no47/// vector.contract user exists, return failure.48FailureOr<vector::ContractionOp> nvgpu::getUserContract(Operation *op) {49 for (Operation *user : op->getUsers()) {50 if (auto contractOp = dyn_cast<vector::ContractionOp>(user))51 return contractOp;52 }53 return failure();54}55 56FailureOr<WarpMatrixInfo> nvgpu::getWarpMatrixInfo(Operation *op) {57 WarpMatrixInfo info;58 59 // Determine the vector type at warp-level.60 if (vector::TransferWriteOp writeOp = dyn_cast<vector::TransferWriteOp>(op)) {61 info.vectorType = writeOp.getVectorType();62 } else if (isa<vector::TransferReadOp, vector::ContractionOp,63 vector::ExtractStridedSliceOp, arith::ConstantOp>(op)) {64 info.vectorType = cast<VectorType>(op->getResult(0).getType());65 } else {66 return op->emitError()67 << "unhandled operation type in nvgpu.mma.sync conversion path";68 }69 70 // Determine the operand role. We assume it is an accumulator/result unless it71 // is directly consumed by a `vector.contract` op.72 info.operandRole = MatMulOperandRole::C;73 FailureOr<vector::ContractionOp> contractOp = getUserContract(op);74 if (failed(contractOp))75 return info;76 77 if ((*contractOp).getLhs() == op->getResult(0))78 info.operandRole = MatMulOperandRole::A;79 else if ((*contractOp).getRhs() == op->getResult(0))80 info.operandRole = MatMulOperandRole::B;81 82 return info;83}84 85int64_t nvgpu::inferTileWidthInBits(const WarpMatrixInfo &type) {86 bool isAcc = isAccumulatorOrResult(type.operandRole);87 Type elType = type.vectorType.getElementType();88 if (isAcc && elType.getIntOrFloatBitWidth() == 32) {89 return 256;90 }91 if (elType.getIntOrFloatBitWidth() == 64) {92 return isAcc ? 512 : 256;93 }94 return 128;95}96 97FailureOr<FragmentElementInfo>98nvgpu::getMmaSyncRegisterType(const WarpMatrixInfo &type) {99 MLIRContext *ctx = type.vectorType.getContext();100 const bool isAccum = isAccumulatorOrResult(type.operandRole);101 102 Type elType = type.vectorType.getElementType();103 if (elType.isF16()) {104 return FragmentElementInfo{VectorType::get(2, Float16Type::get(ctx)), 2, 32,105 inferNumRegistersPerMatrixFragment(type)};106 }107 108 // f64 operand109 Type f64Ty = Float64Type::get(ctx);110 if (elType.isF64()) {111 return isAccum112 ? FragmentElementInfo{VectorType::get(2, f64Ty), 2, 128,113 inferNumRegistersPerMatrixFragment(type)}114 : FragmentElementInfo{f64Ty, 1, 64,115 inferNumRegistersPerMatrixFragment(type)};116 }117 118 // int8 operand119 if (elType.isInteger(8)) {120 return FragmentElementInfo{VectorType::get(4, IntegerType::get(ctx, 8)), 4,121 32, inferNumRegistersPerMatrixFragment(type)};122 }123 124 // int4 operand125 if (elType.isInteger(4)) {126 return FragmentElementInfo{VectorType::get(8, IntegerType::get(ctx, 4)), 8,127 32, inferNumRegistersPerMatrixFragment(type)};128 }129 130 // Integer 32bit acc operands131 if (elType.isInteger(32)) {132 return FragmentElementInfo{VectorType::get(2, IntegerType::get(ctx, 32)), 2,133 64, inferNumRegistersPerMatrixFragment(type)};134 }135 136 // Floating point 32bit operands137 if (elType.isF32()) {138 Type f32Ty = Float32Type::get(ctx);139 return isAccum140 ? FragmentElementInfo{VectorType::get(2, f32Ty), 2, 64,141 inferNumRegistersPerMatrixFragment(type)}142 : FragmentElementInfo{f32Ty, 1, 32,143 inferNumRegistersPerMatrixFragment(type)};144 }145 return failure();146}147 148static AffineMap getRegisterIndexToTileOffsetMap(int64_t lineSize,149 Type elementType,150 ArrayRef<int64_t> operandShape,151 bool isAccumulator,152 int64_t elementsPerRegister,153 AffineExpr logicalValueId) {154 const int64_t elementsPerLine =155 lineSize / elementType.getIntOrFloatBitWidth();156 const std::array<int64_t, 2> num8x128bTiles =157 getTileShape(operandShape, elementType, lineSize);158 AffineExpr registerIdx = logicalValueId.floorDiv(elementsPerRegister);159 return AffineMap::get(160 2, 0,161 {(registerIdx % num8x128bTiles[0]) * 8,162 (registerIdx.floorDiv(num8x128bTiles[0])) * elementsPerLine},163 elementType.getContext());164}165 166FailureOr<AffineMap>167nvgpu::getLaneIdAndValueIdToOperandCoord(OpBuilder &builder, Location loc,168 const WarpMatrixInfo &fragmentType) {169 Type elementType = fragmentType.vectorType.getElementType();170 ArrayRef<int64_t> operandShape = fragmentType.vectorType.getShape();171 FailureOr<FragmentElementInfo> regInfo = getMmaSyncRegisterType(fragmentType);172 if (failed(regInfo))173 return failure();174 175 const int64_t elementBitWidth = elementType.getIntOrFloatBitWidth();176 const int64_t elementsPerRegister =177 regInfo->registerWidthBits / elementBitWidth;178 const int64_t lineSize = inferTileWidthInBits(fragmentType);179 180 AffineExpr laneId, logicalValueIdDim;181 bindDims(builder.getContext(), laneId, logicalValueIdDim);182 183 // Determine what register logicalValueId corresponds to. Use that as a184 // linear index into the coordinate mapping `index -> (tile row, tile col)`.185 AffineMap registerIndexToTileCoord = getRegisterIndexToTileOffsetMap(186 lineSize, elementType, operandShape,187 isAccumulatorOrResult(fragmentType.operandRole), elementsPerRegister,188 logicalValueIdDim);189 190 auto makeMap = [&](ArrayRef<AffineExpr> dimExprs) -> AffineMap {191 return AffineMap::get(2, 0, dimExprs, builder.getContext());192 };193 194 auto tileRow = registerIndexToTileCoord.getResult(0);195 auto tileCol = registerIndexToTileCoord.getResult(1);196 return makeMap({tileRow + laneId.floorDiv(kThreadsPerRow),197 tileCol + (laneId % kThreadsPerRow) * elementsPerRegister +198 (logicalValueIdDim % elementsPerRegister)});199}200 201FailureOr<LdMatrixParams> nvgpu::getLdMatrixParams(const WarpMatrixInfo &type,202 bool transpose) {203 LdMatrixParams params;204 Type elType = type.vectorType.getElementType();205 params.fragmentType = type.vectorType;206 if (type.operandRole == MatMulOperandRole::A ||207 type.operandRole == MatMulOperandRole::C) {208 params.targetLayout = NVVM::MMALayout::row;209 } else {210 params.targetLayout = NVVM::MMALayout::col;211 }212 ArrayRef<int64_t> shape = type.vectorType.getShape();213 params.contiguousDimType = transpose ? vector::IteratorType::parallel214 : vector::IteratorType::reduction;215 216 if (params.contiguousDimType == vector::IteratorType::reduction) {217 params.numTiles = (shape[0] / kNumRowsPerTile) *218 ((shape[1] * elType.getIntOrFloatBitWidth()) / 128);219 } else {220 params.numTiles = (shape[1] / kNumRowsPerTile) *221 ((shape[0] * elType.getIntOrFloatBitWidth()) / 128);222 }223 224 if (params.numTiles == 0)225 return failure();226 227 return params;228}229 230FailureOr<AffineMap>231nvgpu::getLaneIdToLdMatrixMatrixCoord(OpBuilder &builder, Location loc,232 const LdMatrixParams ¶ms) {233 // One thread per 128b row.234 const int bitsPerElement = static_cast<int>(235 params.fragmentType.getElementType().getIntOrFloatBitWidth());236 const int kElementsPer128b = (128 / bitsPerElement);237 ArrayRef<int64_t> operandShape = params.fragmentType.getShape();238 AffineExpr d0 = getAffineDimExpr(0, builder.getContext());239 240 auto makeMap = [&](ArrayRef<AffineExpr> dimExprs) -> AffineMap {241 return AffineMap::get(1, 0, dimExprs, builder.getContext());242 };243 244 // Index `idx` in vectorType `operandShape` maps to the strided dimension of245 // the `srcMemref` memory of the LdMatrixOp.246 int idx =247 (params.contiguousDimType == vector::IteratorType::reduction) ? 0 : 1;248 249 // Affine expr in strided and contiguous dimension encodes the coordinate250 // mapping for the element a thread points to for warp-wide LdMatrixOp.251 AffineExpr strided = d0 % (operandShape[idx]);252 AffineExpr contiguous = d0.floorDiv(operandShape[idx]) * (kElementsPer128b);253 254 // This case corresponds to row-major matrixA or col-major matrixB or255 // row-major matrixC. This is when the memory layout in `srcMemref`256 // match mma.sync hardware vector register operand layout.257 if (params.contiguousDimType == vector::IteratorType::reduction)258 return makeMap({strided, contiguous});259 260 // This case corresponds to col-major matrixA or row-major matrixB or261 // col-major matrixC. This is when the memory layout in `srcMemref` does not262 // match mma.sync hardware vector register operand layout.263 if (params.contiguousDimType == vector::IteratorType::parallel)264 return makeMap({contiguous, strided});265 266 return failure();267}268 269bool nvgpu::canLowerToWarpMatrixOperation(vector::TransferReadOp op) {270 if (op.getMask() || op.hasOutOfBoundsDim())271 return false;272 VectorType type = op.getType();273 // The result type should be 2D. Note that it is possible to expand support so274 // that we are robust to extra unit dimensions that failed to fold, but that275 // would significantly increase downstream code complexity in the conversion276 // step. For now, we rely on other patterns to ensure canonical 2D form is277 // used when targeting the `nvgpu.mma.sync` lowering path.278 if (!type.hasStaticShape() || type.getRank() != 2)279 return false;280 281 // Currently we can't support reads on tensor types because we need stride282 // information to ensure correctness of downstream assumptions. It is possible283 // to enable this if caller can assert that tensor will be lowered in a284 // particular manner.285 auto sourceType = dyn_cast<MemRefType>(op.getBase().getType());286 if (!sourceType)287 return false;288 289 // Check that the last dimension of the read is contiguous. Note that it is290 // possible to expand support for this by scalarizing all the loads during291 // conversion.292 auto [strides, offset] = sourceType.getStridesAndOffset();293 return strides.back() == 1;294}295 296bool nvgpu::canLowerToWarpMatrixOperation(vector::TransferWriteOp op) {297 if (op.getMask() || op.hasOutOfBoundsDim() || op.getTransferRank() == 0)298 return false;299 VectorType type = op.getVectorType();300 if (!type.hasStaticShape() || type.getRank() != 2)301 return false;302 // TODO: Currently we rely on lowering to a `vector.store` operation. We could303 // support the transposed write case by lowering to scalarized `memref.store`304 // operations.305 if (!op.getPermutationMap().isMinorIdentity())306 return false;307 // Currently we can't support reads on tensor types because we need stride308 // information to ensure correctness of downstream assumptions.309 auto sourceType = dyn_cast<MemRefType>(op.getBase().getType());310 if (!sourceType)311 return false;312 313 // Check that the last dimension of the target memref is contiguous. Note that314 // it is possible to expand support for this by scalarizing all the stores315 // during conversion.316 auto [strides, offset] = sourceType.getStridesAndOffset();317 return strides.back() == 1;318}319