996 lines · cpp
1//===- GPUTransformOps.cpp - Implementation of GPU transform ops ----------===//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 9#include "mlir/Dialect/GPU/TransformOps/GPUTransformOps.h"10 11#include "mlir/Conversion/GPUCommon/GPUCommonPass.h"12#include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"13#include "mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h"14#include "mlir/Conversion/LLVMCommon/TypeConverter.h"15#include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"16#include "mlir/Dialect/AMDGPU/Utils/Chipset.h"17#include "mlir/Dialect/Arith/IR/Arith.h"18#include "mlir/Dialect/GPU/IR/GPUDialect.h"19#include "mlir/Dialect/GPU/TransformOps/Utils.h"20#include "mlir/Dialect/GPU/Transforms/Passes.h"21#include "mlir/Dialect/LLVMIR/NVVMDialect.h"22#include "mlir/Dialect/LLVMIR/ROCDLDialect.h"23#include "mlir/Dialect/MemRef/IR/MemRef.h"24#include "mlir/Dialect/SCF/IR/DeviceMappingInterface.h"25#include "mlir/Dialect/SCF/IR/SCF.h"26#include "mlir/Dialect/Transform/IR/TransformDialect.h"27#include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"28#include "mlir/Dialect/Utils/IndexingUtils.h"29#include "mlir/Dialect/Vector/IR/VectorOps.h"30#include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"31#include "mlir/IR/AffineExpr.h"32#include "mlir/IR/Builders.h"33#include "mlir/IR/BuiltinAttributes.h"34#include "mlir/IR/IRMapping.h"35#include "mlir/IR/MLIRContext.h"36#include "mlir/IR/OpDefinition.h"37#include "mlir/IR/Visitors.h"38#include "mlir/Support/LLVM.h"39#include "mlir/Transforms/DialectConversion.h"40#include "llvm/ADT/STLExtras.h"41#include "llvm/ADT/SmallVector.h"42#include "llvm/ADT/TypeSwitch.h"43#include "llvm/Support/DebugLog.h"44#include "llvm/Support/ErrorHandling.h"45#include "llvm/Support/InterleavedRange.h"46#include "llvm/Support/LogicalResult.h"47#include <optional>48#include <type_traits>49 50using namespace mlir;51using namespace mlir::gpu;52using namespace mlir::transform;53using namespace mlir::transform::gpu;54 55#define DEBUG_TYPE "gpu-transforms"56 57//===----------------------------------------------------------------------===//58// Apply...ConversionPatternsOp59//===----------------------------------------------------------------------===//60 61void transform::ApplyGPUToNVVMConversionPatternsOp::populatePatterns(62 TypeConverter &typeConverter, RewritePatternSet &patterns) {63 auto &llvmTypeConverter = static_cast<LLVMTypeConverter &>(typeConverter);64 // NVVM uses alloca in the default address space to represent private65 // memory allocations, so drop private annotations. NVVM uses address66 // space 3 for shared memory. NVVM uses the default address space to67 // represent global memory.68 // Used in populateGpuToNVVMConversionPatternsso attaching here for now.69 // TODO: We should have a single to_nvvm_type_converter.70 populateGpuMemorySpaceAttributeConversions(71 llvmTypeConverter, [](AddressSpace space) -> unsigned {72 switch (space) {73 case AddressSpace::Global:74 return static_cast<unsigned>(NVVM::NVVMMemorySpace::Global);75 case AddressSpace::Workgroup:76 return static_cast<unsigned>(NVVM::NVVMMemorySpace::Shared);77 case AddressSpace::Private:78 return 0;79 }80 llvm_unreachable("unknown address space enum value");81 return static_cast<unsigned>(NVVM::NVVMMemorySpace::Generic);82 });83 // Used in GPUToNVVM/WmmaOpsToNvvm.cpp so attaching here for now.84 // TODO: We should have a single to_nvvm_type_converter.85 llvmTypeConverter.addConversion(86 [&](MMAMatrixType type) -> Type { return convertMMAToLLVMType(type); });87 // Set higher benefit, so patterns will run before generic LLVM lowering.88 populateGpuToNVVMConversionPatterns(llvmTypeConverter, patterns,89 getBenefit());90}91 92LogicalResult93transform::ApplyGPUToNVVMConversionPatternsOp::verifyTypeConverter(94 transform::TypeConverterBuilderOpInterface builder) {95 if (builder.getTypeConverterType() != "LLVMTypeConverter")96 return emitOpError("expected LLVMTypeConverter");97 return success();98}99 100void transform::ApplyGPUWwmaToNVVMConversionPatternsOp::populatePatterns(101 TypeConverter &typeConverter, RewritePatternSet &patterns) {102 auto &llvmTypeConverter = static_cast<LLVMTypeConverter &>(typeConverter);103 populateGpuWMMAToNVVMConversionPatterns(llvmTypeConverter, patterns);104}105 106LogicalResult107transform::ApplyGPUWwmaToNVVMConversionPatternsOp::verifyTypeConverter(108 transform::TypeConverterBuilderOpInterface builder) {109 if (builder.getTypeConverterType() != "LLVMTypeConverter")110 return emitOpError("expected LLVMTypeConverter");111 return success();112}113 114void transform::ApplyGPUSubgroupReduceToNVVMConversionPatternsOp::115 populatePatterns(TypeConverter &typeConverter,116 RewritePatternSet &patterns) {117 auto &llvmTypeConverter = static_cast<LLVMTypeConverter &>(typeConverter);118 populateGpuSubgroupReduceOpLoweringPattern(llvmTypeConverter, patterns);119}120 121LogicalResult transform::ApplyGPUSubgroupReduceToNVVMConversionPatternsOp::122 verifyTypeConverter(transform::TypeConverterBuilderOpInterface builder) {123 if (builder.getTypeConverterType() != "LLVMTypeConverter")124 return emitOpError("expected LLVMTypeConverter");125 return success();126}127 128void transform::ApplyGPUToROCDLConversionPatternsOp::populatePatterns(129 TypeConverter &typeConverter, RewritePatternSet &patterns) {130 auto &llvmTypeConverter = static_cast<LLVMTypeConverter &>(typeConverter);131 populateGpuMemorySpaceAttributeConversions(132 llvmTypeConverter, [](AddressSpace space) {133 switch (space) {134 case AddressSpace::Global:135 return ROCDL::ROCDLDialect::kGlobalMemoryAddressSpace;136 case AddressSpace::Workgroup:137 return ROCDL::ROCDLDialect::kSharedMemoryAddressSpace;138 case AddressSpace::Private:139 return ROCDL::ROCDLDialect::kPrivateMemoryAddressSpace;140 }141 llvm_unreachable("unknown address space enum value");142 });143 FailureOr<amdgpu::Chipset> maybeChipset =144 amdgpu::Chipset::parse(getChipset());145 assert(llvm::succeeded(maybeChipset) && "expected valid chipset");146 populateGpuToROCDLConversionPatterns(147 llvmTypeConverter, patterns, mlir::gpu::amd::Runtime::HIP, *maybeChipset);148}149 150LogicalResult151transform::ApplyGPUToROCDLConversionPatternsOp::verifyTypeConverter(152 transform::TypeConverterBuilderOpInterface builder) {153 FailureOr<amdgpu::Chipset> maybeChipset =154 amdgpu::Chipset::parse(getChipset());155 if (failed(maybeChipset)) {156 return emitOpError("Invalid chipset name: " + getChipset());157 }158 if (builder.getTypeConverterType() != "LLVMTypeConverter")159 return emitOpError("expected LLVMTypeConverter");160 return success();161}162 163//===----------------------------------------------------------------------===//164// Apply...PatternsOp165//===----------------------------------------------------------------------===//s166 167void ApplyGPURewritePatternsOp::populatePatterns(RewritePatternSet &patterns) {168 populateGpuRewritePatterns(patterns);169}170 171void transform::ApplyGPUPromoteShuffleToAMDGPUPatternsOp::populatePatterns(172 RewritePatternSet &patterns) {173 std::optional<StringRef> chipsetName = getChipset();174 std::optional<amdgpu::Chipset> maybeChipset;175 if (chipsetName) {176 FailureOr<amdgpu::Chipset> parsedChipset =177 amdgpu::Chipset::parse(*chipsetName);178 assert(llvm::succeeded(parsedChipset) && "expected valid chipset");179 maybeChipset = parsedChipset;180 }181 182 populateGpuPromoteShuffleToAMDGPUPatterns(patterns, maybeChipset);183}184 185//===----------------------------------------------------------------------===//186// ApplyUnrollVectorsSubgroupMmaOp187//===----------------------------------------------------------------------===//188 189/// Pick an unrolling order that will allow tensorcore operation to reuse LHS190/// register.191static std::optional<SmallVector<int64_t>>192gpuMmaUnrollOrder(vector::ContractionOp contract) {193 SmallVector<int64_t> order;194 // First make reduction the outer dimensions.195 for (auto [index, iter] : llvm::enumerate(contract.getIteratorTypes())) {196 if (vector::isReductionIterator(iter)) {197 order.push_back(index);198 }199 }200 201 llvm::SmallDenseSet<int64_t> dims;202 for (AffineExpr expr : contract.getIndexingMapsArray()[0].getResults()) {203 dims.insert(cast<AffineDimExpr>(expr).getPosition());204 }205 // Then parallel dimensions that are part of Lhs as we want to re-use Lhs.206 for (auto [index, iter] : llvm::enumerate(contract.getIteratorTypes())) {207 if (vector::isParallelIterator(iter) && dims.count(index)) {208 order.push_back(index);209 }210 }211 // Then the remaining parallel loops.212 for (auto [index, iter] : llvm::enumerate(contract.getIteratorTypes())) {213 if (vector::isParallelIterator(iter) && !dims.count(index)) {214 order.push_back(index);215 }216 }217 return order;218}219 220/// Returns the target vector size for the target operation based on the native221/// vector size specified with `m`, `n`, and `k`.222static std::optional<SmallVector<int64_t>>223getSubgroupMmaNativeVectorSize(Operation *op, int64_t m, int64_t n, int64_t k) {224 if (auto contract = dyn_cast<vector::ContractionOp>(op)) {225 int64_t contractRank = contract.getIteratorTypes().size();226 if (contractRank < 3)227 return std::nullopt;228 SmallVector<int64_t> nativeSize(contractRank - 3, 1);229 nativeSize.append({m, n, k});230 return nativeSize;231 }232 if (auto writeOp = dyn_cast<vector::TransferWriteOp>(op)) {233 int64_t writeRank = writeOp.getVectorType().getRank();234 if (writeRank < 2)235 return std::nullopt;236 SmallVector<int64_t> nativeSize(writeRank - 2, 1);237 nativeSize.append({m, n});238 return nativeSize;239 }240 if (auto readOp = dyn_cast<vector::TransferReadOp>(op)) {241 // Transfer read ops may need different shapes based on how they are being242 // used. For simplicity just match the shape used by the extract strided op.243 VectorType sliceType;244 for (Operation *users : op->getUsers()) {245 auto extract = dyn_cast<vector::ExtractStridedSliceOp>(users);246 if (!extract)247 return std::nullopt;248 auto vecType = cast<VectorType>(extract.getResult().getType());249 if (sliceType && sliceType != vecType)250 return std::nullopt;251 sliceType = vecType;252 }253 return llvm::to_vector(sliceType.getShape());254 }255 if ((OpTrait::hasElementwiseMappableTraits(op) && op->getNumResults() == 1)) {256 if (auto vecType = dyn_cast<VectorType>(op->getResultTypes()[0])) {257 // TODO: The condition for unrolling elementwise should be restricted258 // only to operations that need unrolling (connected to the contract).259 if (vecType.getRank() < 2)260 return std::nullopt;261 262 // First check whether there is a slice to infer the shape from. This is263 // required for cases where the accumulator type differs from the input264 // types, in which case we will see an `arith.ext_` between the contract265 // and transfer_read which needs to be unrolled.266 VectorType sliceType;267 for (Operation *users : op->getUsers()) {268 auto extract = dyn_cast<vector::ExtractStridedSliceOp>(users);269 if (!extract)270 return std::nullopt;271 auto vecType = cast<VectorType>(extract.getResult().getType());272 if (sliceType && sliceType != vecType)273 return std::nullopt;274 sliceType = vecType;275 }276 if (sliceType)277 return llvm::to_vector(sliceType.getShape());278 279 // Else unroll for trailing elementwise.280 SmallVector<int64_t> nativeSize(vecType.getRank() - 2, 1);281 // Map elementwise ops to the output shape.282 nativeSize.append({m, n});283 return nativeSize;284 }285 }286 return std::nullopt;287}288 289void transform::ApplyUnrollVectorsSubgroupMmaOp::populatePatterns(290 RewritePatternSet &patterns) {291 auto unrollOrder = [](Operation *op) -> std::optional<SmallVector<int64_t>> {292 auto contract = dyn_cast<vector::ContractionOp>(op);293 if (!contract)294 return std::nullopt;295 return gpuMmaUnrollOrder(contract);296 };297 298 int64_t m = getM();299 int64_t n = getN();300 int64_t k = getK();301 auto nativeShapeFn =302 [m, n, k](Operation *op) -> std::optional<SmallVector<int64_t>> {303 return getSubgroupMmaNativeVectorSize(op, m, n, k);304 };305 vector::populateVectorUnrollPatterns(306 patterns, vector::UnrollVectorOptions()307 .setNativeShapeFn(nativeShapeFn)308 .setUnrollTraversalOrderFn(unrollOrder));309}310 311//===----------------------------------------------------------------------===//312// EliminateBarriersOp313//===----------------------------------------------------------------------===//314 315void EliminateBarriersOp::populatePatterns(RewritePatternSet &patterns) {316 populateGpuEliminateBarriersPatterns(patterns);317}318 319//===----------------------------------------------------------------------===//320// Block and thread mapping utilities.321//===----------------------------------------------------------------------===//322 323namespace {324/// Local types used for mapping verification.325struct MappingKind {};326struct BlockMappingKind : MappingKind {};327struct ThreadMappingKind : MappingKind {};328} // namespace329 330static DiagnosedSilenceableFailure331definiteFailureHelper(std::optional<TransformOpInterface> transformOp,332 Operation *target, const Twine &message) {333 if (transformOp.has_value())334 return transformOp->emitDefiniteFailure() << message;335 return emitDefiniteFailure(target, message);336}337 338/// Check if given mapping attributes are one of the desired attributes339template <typename MappingKindType>340static DiagnosedSilenceableFailure341checkMappingAttributeTypes(std::optional<TransformOpInterface> transformOp,342 scf::ForallOp forallOp) {343 if (!forallOp.getMapping().has_value()) {344 return definiteFailureHelper(transformOp, forallOp,345 "scf.forall op requires a mapping attribute");346 }347 348 bool hasBlockMapping = llvm::any_of(forallOp.getMapping().value(),349 llvm::IsaPred<GPUBlockMappingAttr>);350 bool hasWarpgroupMapping = llvm::any_of(351 forallOp.getMapping().value(), llvm::IsaPred<GPUWarpgroupMappingAttr>);352 bool hasWarpMapping = llvm::any_of(forallOp.getMapping().value(),353 llvm::IsaPred<GPUWarpMappingAttr>);354 bool hasThreadMapping = llvm::any_of(forallOp.getMapping().value(),355 llvm::IsaPred<GPUThreadMappingAttr>);356 bool hasLaneMapping = llvm::any_of(forallOp.getMapping().value(),357 llvm::IsaPred<GPULaneMappingAttr>);358 int64_t countMappingTypes = 0;359 countMappingTypes += hasBlockMapping ? 1 : 0;360 countMappingTypes += hasWarpgroupMapping ? 1 : 0;361 countMappingTypes += hasWarpMapping ? 1 : 0;362 countMappingTypes += hasThreadMapping ? 1 : 0;363 countMappingTypes += hasLaneMapping ? 1 : 0;364 if (countMappingTypes > 1) {365 return definiteFailureHelper(366 transformOp, forallOp,367 "cannot mix different mapping types, use nesting");368 }369 if (std::is_same<MappingKindType, BlockMappingKind>::value &&370 !hasBlockMapping) {371 return definiteFailureHelper(372 transformOp, forallOp,373 "scf.forall op requires a mapping attribute of kind 'block'");374 }375 if (std::is_same<MappingKindType, ThreadMappingKind>::value &&376 !hasLaneMapping && !hasThreadMapping && !hasWarpMapping &&377 !hasWarpgroupMapping) {378 return definiteFailureHelper(transformOp, forallOp,379 "scf.forall op requires a mapping attribute "380 "of kind 'thread' or 'warp'");381 }382 383 DenseSet<Attribute> seen;384 for (Attribute map : forallOp.getMapping()->getValue()) {385 if (seen.contains(map)) {386 return definiteFailureHelper(387 transformOp, forallOp,388 "duplicate attribute, cannot map different loops "389 "to the same mapping id");390 }391 seen.insert(map);392 }393 394 auto isLinear = [](DeviceMappingAttrInterface attr) {395 return attr.isLinearMapping();396 };397 if (llvm::any_of(forallOp.getDeviceMappingAttrs(), isLinear) &&398 !llvm::all_of(forallOp.getDeviceMappingAttrs(), isLinear)) {399 return definiteFailureHelper(400 transformOp, forallOp,401 "cannot mix linear and non-linear mapping modes");402 }403 404 FailureOr<DeviceMaskingAttrInterface> maybeMaskingAttr =405 forallOp.getDeviceMaskingAttr();406 if (succeeded(maybeMaskingAttr) && *maybeMaskingAttr &&407 !forallOp.usesLinearMapping()) {408 return definiteFailureHelper(409 transformOp, forallOp,410 "device masking is only available in linear mapping mode");411 }412 413 return DiagnosedSilenceableFailure::success();414}415 416template <typename MappingKindType>417static DiagnosedSilenceableFailure418verifyGpuMapping(std::optional<TransformOpInterface> transformOp,419 scf::ForallOp forallOp) {420 // Check the types of the mapping attributes match.421 DiagnosedSilenceableFailure typeRes =422 checkMappingAttributeTypes<MappingKindType>(transformOp, forallOp);423 if (!typeRes.succeeded())424 return typeRes;425 426 // Perform other non-types verifications.427 if (!forallOp.isNormalized())428 return definiteFailureHelper(transformOp, forallOp,429 "unsupported non-normalized loops");430 if (forallOp.getNumResults() > 0)431 return definiteFailureHelper(transformOp, forallOp,432 "only bufferized scf.forall can be mapped");433 bool useLinearMapping = forallOp.usesLinearMapping();434 // TODO: This would be more natural with support for Optional<EnumParameter>435 // in GPUDeviceMappingAttr.436 int64_t maxNumMappingsSupported =437 useLinearMapping ? (getMaxEnumValForMappingId() -438 static_cast<uint64_t>(MappingId::DimZ))439 : 3;440 if (forallOp.getRank() > maxNumMappingsSupported) {441 return definiteFailureHelper(transformOp, forallOp,442 "scf.forall with rank > ")443 << maxNumMappingsSupported444 << " does not lower for the specified mapping attribute type";445 }446 auto numParallelIterations =447 getConstantIntValues(forallOp.getMixedUpperBound());448 if (!forallOp.isNormalized() || !numParallelIterations.has_value()) {449 return definiteFailureHelper(450 transformOp, forallOp,451 "requires statically sized, normalized forall op");452 }453 return DiagnosedSilenceableFailure::success();454}455 456/// Struct to return the result of the rewrite of a forall operation.457struct ForallRewriteResult {458 SmallVector<int64_t> mappingSizes;459 SmallVector<Value> mappingIds;460};461 462/// Helper to replace ids of dimensions known to be 1 by 0 to simplify the IR.463template <typename OpTy, typename OperationOrBlock>464static void465replaceUnitMappingIdsHelper(RewriterBase &rewriter, Location loc,466 OperationOrBlock *parent, Value replacement,467 ArrayRef<int64_t> availableMappingSizes) {468 parent->walk([&](OpTy idOp) {469 if (availableMappingSizes[static_cast<int64_t>(idOp.getDimension())] == 1)470 rewriter.replaceAllUsesWith(idOp.getResult(), replacement);471 });472}473 474static DiagnosedSilenceableFailure rewriteOneForallCommonImpl(475 RewriterBase &rewriter, std::optional<TransformOpInterface> transformOp,476 scf::ForallOp forallOp, ArrayRef<int64_t> availableMappingSizes,477 ForallRewriteResult &result, const GpuIdBuilder &gpuIdBuilder) {478 LDBG() << "--start rewriteOneForallCommonImpl";479 480 // Step 1. Complete the mapping to a full mapping (with 1s) if necessary.481 auto numParallelIterations =482 getConstantIntValues(forallOp.getMixedUpperBound());483 assert(forallOp.isNormalized() && numParallelIterations.has_value() &&484 "requires statically sized, normalized forall op");485 SmallVector<int64_t> tmpMappingSizes = numParallelIterations.value();486 SmallVector<DeviceMappingAttrInterface> forallMappingAttrsVec =487 forallOp.getDeviceMappingAttrs();488 SetVector<Attribute> forallMappingAttrs;489 forallMappingAttrs.insert_range(forallMappingAttrsVec);490 auto comparator = [](Attribute a, Attribute b) -> bool {491 return cast<DeviceMappingAttrInterface>(a).getMappingId() <492 cast<DeviceMappingAttrInterface>(b).getMappingId();493 };494 495 // Step 1.b. In the linear case, compute the max mapping to avoid needlessly496 // mapping all dimensions. In the 3-D mapping case we need to map all497 // dimensions.498 DeviceMappingAttrInterface maxMapping = cast<DeviceMappingAttrInterface>(499 *llvm::max_element(forallMappingAttrs, comparator));500 DeviceMappingAttrInterface maxLinearMapping;501 if (maxMapping.isLinearMapping())502 maxLinearMapping = maxMapping;503 for (auto attr : gpuIdBuilder.mappingAttributes) {504 // If attr overflows, just skip.505 if (maxLinearMapping && comparator(maxLinearMapping, attr))506 continue;507 // Try to insert. If element was already present, just continue.508 if (!forallMappingAttrs.insert(attr))509 continue;510 // Otherwise, we have a new insertion without a size -> use size 1.511 tmpMappingSizes.push_back(1);512 }513 LDBG() << "----tmpMappingSizes extracted from scf.forall op: "514 << llvm::interleaved(tmpMappingSizes);515 516 // Step 2. sort the values by the corresponding DeviceMappingAttrInterface.517 SmallVector<int64_t> forallMappingSizes = getValuesSortedByKey(518 forallMappingAttrs.getArrayRef(), tmpMappingSizes, comparator);519 LDBG() << "----forallMappingSizes: " << llvm::interleaved(forallMappingSizes);520 LDBG() << "----forallMappingAttrs: " << llvm::interleaved(forallMappingAttrs);521 522 // Step 3. Generate the mappingIdOps using the provided generator.523 Location loc = forallOp.getLoc();524 OpBuilder::InsertionGuard guard(rewriter);525 rewriter.setInsertionPoint(forallOp);526 SmallVector<int64_t> originalBasis(availableMappingSizes);527 bool originalBasisWasProvided = !originalBasis.empty();528 if (!originalBasisWasProvided) {529 LDBG() << "----originalBasis was not provided, deriving it and there will "530 "be no "531 "predication";532 originalBasis = forallMappingSizes;533 while (originalBasis.size() < 3)534 originalBasis.push_back(1);535 } else {536 LDBG() << "----originalBasis was provided, using it, there will be "537 "predication";538 }539 LDBG() << "------originalBasis: " << llvm::interleaved(originalBasis);540 541 IdBuilderResult builderResult =542 gpuIdBuilder.idBuilder(rewriter, loc, forallMappingSizes, originalBasis);543 if (!builderResult.errorMsg.empty())544 return definiteFailureHelper(transformOp, forallOp, builderResult.errorMsg);545 546 LDBG() << builderResult;547 548 // Step 4. Map the induction variables to the mappingIdOps, this may involve549 // a permutation.550 SmallVector<Value> mappingIdOps = builderResult.mappingIdOps;551 IRMapping bvm;552 for (auto [iv, dim] : llvm::zip_equal(553 forallOp.getInductionVars(),554 forallMappingAttrs.getArrayRef().take_front(forallOp.getRank()))) {555 auto mappingAttr = cast<DeviceMappingAttrInterface>(dim);556 Value peIdOp = mappingIdOps[mappingAttr.getRelativeIndex()];557 LDBG() << "----map: " << iv << " to " << peIdOp;558 bvm.map(iv, peIdOp);559 }560 561 // Step 5. If the originalBasis is already known, create conditionals to562 // predicate the region. Otherwise, the current forall determines the563 // originalBasis and no predication occurs.564 Value predicate;565 if (originalBasisWasProvided) {566 for (Value tmpPredicate : builderResult.predicateOps) {567 predicate = predicate ? arith::AndIOp::create(rewriter, loc, predicate,568 tmpPredicate)569 : tmpPredicate;570 }571 }572 573 // Step 6. Move the body of forallOp.574 // Erase the terminator first, it will not be used.575 rewriter.eraseOp(forallOp.getTerminator());576 Block *targetBlock;577 Block::iterator insertionPoint;578 if (predicate) {579 // Step 6.a. If predicated, move at the beginning.580 auto ifOp = scf::IfOp::create(rewriter, loc, predicate,581 /*withElseRegion=*/false);582 targetBlock = ifOp.thenBlock();583 insertionPoint = ifOp.thenBlock()->begin();584 } else {585 // Step 6.b. Otherwise, move inline just at the rewriter insertion586 // point.587 targetBlock = forallOp->getBlock();588 insertionPoint = rewriter.getInsertionPoint();589 }590 Block &sourceBlock = forallOp.getRegion().front();591 targetBlock->getOperations().splice(insertionPoint,592 sourceBlock.getOperations());593 594 // Step 7. RAUW indices.595 for (Value loopIndex : forallOp.getInductionVars()) {596 Value threadIdx = bvm.lookup(loopIndex);597 rewriter.replaceAllUsesWith(loopIndex, threadIdx);598 }599 600 // Step 8. Erase old op.601 rewriter.eraseOp(forallOp);602 603 LDBG() << "----result forallMappingSizes: "604 << llvm::interleaved(forallMappingSizes);605 LDBG() << "----result mappingIdOps: " << llvm::interleaved(mappingIdOps);606 607 result = ForallRewriteResult{forallMappingSizes, mappingIdOps};608 return DiagnosedSilenceableFailure::success();609}610 611//===----------------------------------------------------------------------===//612// MapForallToBlocks613//===----------------------------------------------------------------------===//614 615DiagnosedSilenceableFailure mlir::transform::gpu::mapForallToBlocksImpl(616 RewriterBase &rewriter, TransformOpInterface transformOp,617 scf::ForallOp forallOp, SmallVectorImpl<int64_t> &gridDims,618 const GpuIdBuilder &gpuIdBuilder) {619 LDBG() << "Start mapForallToBlocksImpl";620 621 {622 // GPU-specific verifications. There is no better place to anchor623 // those right now: the ForallOp is target-independent and the transform624 // op does not apply to individual ForallOp.625 DiagnosedSilenceableFailure diag =626 verifyGpuMapping<BlockMappingKind>(transformOp, forallOp);627 if (!diag.succeeded())628 return diag;629 }630 631 Location loc = forallOp.getLoc();632 Block *parentBlock = forallOp->getBlock();633 Value zero;634 {635 // Create an early zero index value for replacements and immediately reset636 // the insertion point.637 OpBuilder::InsertionGuard guard(rewriter);638 rewriter.setInsertionPointToStart(parentBlock);639 zero = arith::ConstantIndexOp::create(rewriter, loc, 0);640 }641 642 ForallRewriteResult rewriteResult;643 DiagnosedSilenceableFailure diag = rewriteOneForallCommonImpl(644 rewriter, transformOp, forallOp,645 /*availableMappingSizes=*/gridDims, rewriteResult, gpuIdBuilder);646 647 // Return if anything goes wrong, use silenceable failure as a match648 // failure.649 if (!diag.succeeded())650 return diag;651 652 // If gridDims was not provided already, set it from the return.653 if (gridDims.empty()) {654 gridDims = rewriteResult.mappingSizes;655 while (gridDims.size() < 3)656 gridDims.push_back(1);657 }658 assert(gridDims.size() == 3 && "Need 3-D gridDims");659 660 // Replace ids of dimensions known to be 1 by 0 to simplify the IR.661 // Here, the result of mapping determines the available mapping sizes.662 replaceUnitMappingIdsHelper<BlockDimOp>(rewriter, loc, parentBlock, zero,663 rewriteResult.mappingSizes);664 665 return DiagnosedSilenceableFailure::success();666}667 668DiagnosedSilenceableFailure669mlir::transform::gpu::findTopLevelForallOp(Operation *target,670 scf::ForallOp &topLevelForallOp,671 TransformOpInterface transformOp) {672 auto walkResult = target->walk([&](scf::ForallOp forallOp) {673 if (forallOp->getParentOfType<scf::ForallOp>())674 return WalkResult::advance();675 if (topLevelForallOp)676 // TODO: Handle multiple forall if they are independent.677 return WalkResult::interrupt();678 topLevelForallOp = forallOp;679 return WalkResult::advance();680 });681 682 if (walkResult.wasInterrupted() || !topLevelForallOp)683 return transformOp.emitSilenceableError()684 << "could not find a unique topLevel scf.forall";685 return DiagnosedSilenceableFailure::success();686}687 688DiagnosedSilenceableFailure transform::MapForallToBlocks::applyToOne(689 transform::TransformRewriter &rewriter, Operation *target,690 ApplyToEachResultList &results, transform::TransformState &state) {691 LaunchOp gpuLaunch = dyn_cast<LaunchOp>(target);692 auto transformOp = cast<TransformOpInterface>(getOperation());693 694 if (!getGenerateGpuLaunch() && !gpuLaunch) {695 DiagnosedSilenceableFailure diag =696 emitSilenceableError()697 << "Given target is not gpu.launch, set `generate_gpu_launch` "698 "attribute";699 diag.attachNote(target->getLoc()) << "when applied to this payload op";700 return diag;701 }702 703 scf::ForallOp topLevelForallOp;704 DiagnosedSilenceableFailure diag = mlir::transform::gpu::findTopLevelForallOp(705 target, topLevelForallOp, transformOp);706 if (!diag.succeeded()) {707 diag.attachNote(target->getLoc()) << "when applied to this payload op";708 return diag;709 }710 assert(topLevelForallOp && "expect an scf.forall");711 712 SmallVector<int64_t> gridDims{getGridDims()};713 if (!getGenerateGpuLaunch() && gridDims.size() != 3)714 return transformOp.emitDefiniteFailure("transform require size-3 mapping");715 716 OpBuilder::InsertionGuard guard(rewriter);717 rewriter.setInsertionPoint(topLevelForallOp);718 719 // Generate gpu launch here and move the forall inside720 if (getGenerateGpuLaunch()) {721 DiagnosedSilenceableFailure diag =722 createGpuLaunch(rewriter, target->getLoc(), transformOp, gpuLaunch);723 if (!diag.succeeded())724 return diag;725 726 rewriter.setInsertionPointToStart(&gpuLaunch.getBody().front());727 Operation *newForallOp = rewriter.clone(*topLevelForallOp);728 rewriter.eraseOp(topLevelForallOp);729 topLevelForallOp = cast<scf::ForallOp>(newForallOp);730 }731 732 // The BlockIdBuilder adapts to whatever is thrown at it.733 bool useLinearMapping = false;734 if (topLevelForallOp.getMapping())735 useLinearMapping = topLevelForallOp.usesLinearMapping();736 737 FailureOr<DeviceMaskingAttrInterface> maybeMaskingAttr =738 topLevelForallOp.getDeviceMaskingAttr();739 assert(succeeded(maybeMaskingAttr) && "unexpected failed maybeMaskingAttr");740 assert((!*maybeMaskingAttr || useLinearMapping) &&741 "masking requires linear mapping");742 743 GpuBlockIdBuilder gpuBlockIdBuilder(getContext(), useLinearMapping,744 *maybeMaskingAttr);745 746 diag = mlir::transform::gpu::mapForallToBlocksImpl(747 rewriter, transformOp, topLevelForallOp, gridDims, gpuBlockIdBuilder);748 if (!diag.succeeded())749 return diag;750 751 // Set the GPU launch configuration for the grid dims late, this is752 // subject to IR inspection.753 diag = alterGpuLaunch(rewriter, gpuLaunch,754 cast<TransformOpInterface>(getOperation()), gridDims[0],755 gridDims[1], gridDims[2]);756 757 results.push_back(gpuLaunch);758 return diag;759}760 761LogicalResult transform::MapForallToBlocks::verify() {762 if (!getGridDims().empty() && getGridDims().size() != 3) {763 return emitOpError() << "transform requires empty or size-3 grid_dims";764 }765 return success();766}767 768//===----------------------------------------------------------------------===//769// MapNestedForallToThreads770//===----------------------------------------------------------------------===//771 772static DiagnosedSilenceableFailure checkMappingSpec(773 std::optional<TransformOpInterface> transformOp, scf::ForallOp forallOp,774 ArrayRef<int64_t> numParallelIterations, ArrayRef<int64_t> blockOrGridSizes,775 int factor, bool useLinearMapping = false) {776 if (!useLinearMapping && blockOrGridSizes.front() % factor != 0) {777 auto diag = definiteFailureHelper(778 transformOp, forallOp,779 Twine("3-D mapping: size of threadIdx.x must be a multiple of ") +780 Twine(factor));781 return diag;782 }783 if (computeProduct(numParallelIterations) * factor >784 computeProduct(blockOrGridSizes)) {785 auto diag = definiteFailureHelper(786 transformOp, forallOp,787 Twine("the number of required parallel resources (blocks or "788 "threads) ") +789 Twine(computeProduct(numParallelIterations) * factor) +790 " overflows the number of available resources " +791 Twine(computeProduct(blockOrGridSizes)));792 return diag;793 }794 return DiagnosedSilenceableFailure::success();795}796 797static DiagnosedSilenceableFailure798getThreadIdBuilder(std::optional<TransformOpInterface> transformOp,799 scf::ForallOp forallOp, ArrayRef<int64_t> blockSizes,800 int64_t warpSize, GpuIdBuilder &gpuIdBuilder) {801 DeviceMappingAttrInterface mappingAttr =802 forallOp.getDeviceMappingAttrs().front();803 bool useLinearMapping = mappingAttr.isLinearMapping();804 805 // Sanity checks that may result in runtime verification errors.806 auto numParallelIterations =807 getConstantIntValues((forallOp.getMixedUpperBound()));808 if (!forallOp.isNormalized() || !numParallelIterations.has_value()) {809 return definiteFailureHelper(810 transformOp, forallOp,811 "requires statically sized, normalized forall op");812 }813 int64_t factor = 1;814 if (isa<GPUWarpgroupMappingAttr>(mappingAttr)) {815 factor = GpuWarpgroupIdBuilder::kNumWarpsPerGroup * warpSize;816 } else if (isa<GPUWarpMappingAttr>(mappingAttr)) {817 factor = warpSize;818 }819 DiagnosedSilenceableFailure diag =820 checkMappingSpec(transformOp, forallOp, numParallelIterations.value(),821 blockSizes, factor, useLinearMapping);822 if (!diag.succeeded())823 return diag;824 825 FailureOr<DeviceMaskingAttrInterface> maybeMaskingAttr =826 forallOp.getDeviceMaskingAttr();827 assert(succeeded(maybeMaskingAttr) && "unexpected failed maybeMaskingAttr");828 assert((!*maybeMaskingAttr || useLinearMapping) &&829 "masking requires linear mapping");830 831 // Start mapping.832 MLIRContext *ctx = forallOp.getContext();833 gpuIdBuilder =834 TypeSwitch<DeviceMappingAttrInterface, GpuIdBuilder>(mappingAttr)835 .Case([&](GPUWarpgroupMappingAttr) {836 return GpuWarpgroupIdBuilder(ctx, warpSize, useLinearMapping,837 *maybeMaskingAttr);838 })839 .Case([&](GPUWarpMappingAttr) {840 return GpuWarpIdBuilder(ctx, warpSize, useLinearMapping,841 *maybeMaskingAttr);842 })843 .Case([&](GPUThreadMappingAttr) {844 return GpuThreadIdBuilder(ctx, useLinearMapping, *maybeMaskingAttr);845 })846 .Case([&](GPULaneMappingAttr) {847 return GpuLaneIdBuilder(ctx, warpSize, useLinearMapping,848 *maybeMaskingAttr);849 })850 .DefaultUnreachable("unknown mapping attribute");851 return DiagnosedSilenceableFailure::success();852}853 854DiagnosedSilenceableFailure mlir::transform::gpu::mapOneForallToThreadsImpl(855 RewriterBase &rewriter, std::optional<TransformOpInterface> transformOp,856 scf::ForallOp forallOp, ArrayRef<int64_t> blockSizes, int64_t warpSize,857 bool syncAfterDistribute) {858 859 {860 // GPU-specific verifications. There is no better place to anchor861 // those right now: the ForallOp is target-independent and the transform862 // op does not apply to individual ForallOp.863 DiagnosedSilenceableFailure diag =864 verifyGpuMapping<ThreadMappingKind>(transformOp, forallOp);865 if (!diag.succeeded())866 return diag;867 }868 869 GpuIdBuilder gpuIdBuilder;870 {871 // Try to construct the id builder, if it fails, return.872 DiagnosedSilenceableFailure diag = getThreadIdBuilder(873 transformOp, forallOp, blockSizes, warpSize, gpuIdBuilder);874 if (!diag.succeeded())875 return diag;876 }877 878 Location loc = forallOp.getLoc();879 OpBuilder::InsertionGuard g(rewriter);880 // Insert after to allow for syncthreads after `forall` is erased.881 rewriter.setInsertionPointAfter(forallOp);882 ForallRewriteResult rewriteResult;883 DiagnosedSilenceableFailure diag = rewriteOneForallCommonImpl(884 rewriter, transformOp, forallOp, blockSizes, rewriteResult, gpuIdBuilder);885 if (!diag.succeeded())886 return diag;887 // Add a syncthreads if needed. TODO: warpsync888 if (syncAfterDistribute)889 BarrierOp::create(rewriter, loc);890 891 return DiagnosedSilenceableFailure::success();892}893 894DiagnosedSilenceableFailure mlir::transform::gpu::mapNestedForallToThreadsImpl(895 RewriterBase &rewriter, std::optional<TransformOpInterface> transformOp,896 Operation *target, ArrayRef<int64_t> blockDims, int64_t warpSize,897 bool syncAfterDistribute) {898 LDBG() << "Start mapNestedForallToThreadsImpl";899 if (blockDims.size() != 3) {900 return definiteFailureHelper(transformOp, target,901 "requires size-3 thread mapping");902 }903 904 // Create an early zero index value for replacements.905 Location loc = target->getLoc();906 Value zero = arith::ConstantIndexOp::create(rewriter, loc, 0);907 DiagnosedSilenceableFailure diag = DiagnosedSilenceableFailure::success();908 WalkResult walkResult = target->walk([&](scf::ForallOp forallOp) {909 diag = mlir::transform::gpu::mapOneForallToThreadsImpl(910 rewriter, transformOp, forallOp, blockDims, warpSize,911 syncAfterDistribute);912 if (diag.isDefiniteFailure())913 return WalkResult::interrupt();914 if (diag.succeeded())915 return WalkResult::skip();916 return WalkResult::advance();917 });918 if (walkResult.wasInterrupted())919 return diag;920 921 // Replace ids of dimensions known to be 1 by 0 to simplify the IR.922 // Here, the result of mapping determines the available mapping sizes.923 replaceUnitMappingIdsHelper<ThreadIdOp>(rewriter, loc, target, zero,924 blockDims);925 926 return DiagnosedSilenceableFailure::success();927}928 929DiagnosedSilenceableFailure transform::MapNestedForallToThreads::applyToOne(930 transform::TransformRewriter &rewriter, Operation *target,931 ApplyToEachResultList &results, TransformState &state) {932 LaunchOp gpuLaunch = dyn_cast<LaunchOp>(target);933 auto transformOp = cast<TransformOpInterface>(getOperation());934 935 // Basic high-level verifications.936 if (!gpuLaunch)937 return emitSilenceableError() << "Given target is not a gpu.launch";938 939 // Mapping to block ids.940 SmallVector<int64_t> blockDims{getBlockDims()};941 DiagnosedSilenceableFailure diag =942 checkGpuLimits(transformOp, std::nullopt, std::nullopt, std::nullopt,943 blockDims[0], blockDims[1], blockDims[2]);944 if (diag.isSilenceableFailure()) {945 diag.attachNote(getLoc()) << getBlockDimsAttrName() << " is too large";946 return diag;947 }948 949 // Set the GPU launch configuration for the block dims early, this is not950 // subject to IR inspection.951 diag = alterGpuLaunch(rewriter, gpuLaunch, transformOp, std::nullopt,952 std::nullopt, std::nullopt, blockDims[0], blockDims[1],953 blockDims[2]);954 955 rewriter.setInsertionPointToStart(&gpuLaunch.getBody().front());956 diag =957 mapNestedForallToThreadsImpl(rewriter, transformOp, gpuLaunch, blockDims,958 getWarpSize(), getSyncAfterDistribute());959 960 results.push_back(gpuLaunch.getOperation());961 return diag;962}963 964//===----------------------------------------------------------------------===//965// Transform op registration966//===----------------------------------------------------------------------===//967 968namespace {969/// Registers new ops and declares PDL as dependent dialect since the970/// additional ops are using PDL types for operands and results.971class GPUTransformDialectExtension972 : public transform::TransformDialectExtension<973 GPUTransformDialectExtension> {974public:975 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(GPUTransformDialectExtension)976 977 GPUTransformDialectExtension() {978 declareGeneratedDialect<GPUDialect>();979 declareGeneratedDialect<amdgpu::AMDGPUDialect>();980 declareGeneratedDialect<arith::ArithDialect>();981 declareGeneratedDialect<scf::SCFDialect>();982 registerTransformOps<983#define GET_OP_LIST984#include "mlir/Dialect/GPU/TransformOps/GPUTransformOps.cpp.inc"985 >();986 }987};988} // namespace989 990#define GET_OP_CLASSES991#include "mlir/Dialect/GPU/TransformOps/GPUTransformOps.cpp.inc"992 993void mlir::gpu::registerTransformDialectExtension(DialectRegistry ®istry) {994 registry.addExtensions<GPUTransformDialectExtension>();995}996