719 lines · cpp
1//===- TestTilingInterfaceTransformOps.cpp - Test `TilingInterface` ------===//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// This file defines transform dialect operations used for testing10// TilingInterface11//12//===----------------------------------------------------------------------===//13 14#include "mlir/Dialect/Affine/IR/AffineOps.h"15#include "mlir/Dialect/Index/IR/IndexDialect.h"16#include "mlir/Dialect/SCF/IR/SCF.h"17#include "mlir/Dialect/SCF/Transforms/TileUsingInterface.h"18#include "mlir/Dialect/Transform/IR/TransformAttrs.h"19#include "mlir/Dialect/Transform/IR/TransformDialect.h"20#include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"21#include "mlir/Dialect/Utils/StaticValueUtils.h"22#include "mlir/Dialect/Utils/StructuredOpsUtils.h"23#include "mlir/IR/Dominance.h"24#include "mlir/IR/OpImplementation.h"25#include "mlir/Interfaces/TilingInterface.h"26#include "llvm/Support/Debug.h"27 28#define DEBUG_TYPE "test-tiling-interface"29 30#define GET_OP_CLASSES31#include "TestTilingInterfaceTransformOps.h.inc"32 33using namespace mlir;34using namespace mlir::transform;35 36//===----------------------------------------------------------------------===//37// TestFuseAndYieldOp38//===----------------------------------------------------------------------===//39 40static llvm::SmallDenseSet<Operation *> collectTiledAndFusedOps(Operation *op) {41 SmallVector<Operation *> worklist;42 llvm::SmallDenseSet<Operation *> producers;43 worklist.push_back(op);44 producers.insert(op);45 while (!worklist.empty()) {46 Operation *current = worklist.pop_back_val();47 for (OpOperand &operand : current->getOpOperands()) {48 Operation *producer = operand.get().getDefiningOp();49 if (!producer || !isa<TilingInterface>(producer) ||50 producers.contains(producer))51 continue;52 worklist.push_back(producer);53 producers.insert(producer);54 }55 }56 return producers;57}58 59/// Apply a tile and fuse transformation to all payload ops and store both the60/// tiled operation as well as the created tile loops.61template <typename Range>62static LogicalResult63applyTileAndFuseToAll(RewriterBase &rewriter, Operation *transformOp,64 Range &&payloadOps, unsigned numLoops,65 scf::SCFTilingOptions tilingOptions,66 TransformResults &transformResults) {67 SmallVector<Operation *> tiledOps;68 SmallVector<SmallVector<Operation *>> loopOps(numLoops);69 70 for (Operation *target : payloadOps) {71 auto tilingInterfaceOp = dyn_cast<TilingInterface>(target);72 if (!tilingInterfaceOp)73 return transformOp->emitError("only TilingInterface ops are supported");74 DominanceInfo dominanceInfo(tilingInterfaceOp);75 76 llvm::SmallDenseSet<Operation *> tiledAndFusedOps =77 collectTiledAndFusedOps(tilingInterfaceOp);78 llvm::DenseSet<Operation *> yieldReplacementsFor;79 for (auto op : tiledAndFusedOps) {80 if (llvm::any_of(op->getUsers(), [&](Operation *user) {81 return dominanceInfo.properlyDominates(tilingInterfaceOp, user);82 })) {83 yieldReplacementsFor.insert(op);84 }85 }86 87 scf::SCFTileAndFuseOptions tileAndFuseOptions;88 tileAndFuseOptions.setTilingOptions(tilingOptions);89 90 scf::SCFTileAndFuseOptions::ControlFnTy controlFn =91 [&](tensor::ExtractSliceOp candidateSliceOp, OpResult originalProducer,92 bool isDestinationOperand)93 -> std::optional<scf::SCFTileAndFuseOptions::ControlFnResult> {94 Operation *owner = originalProducer.getOwner();95 bool yieldProducerReplacement = yieldReplacementsFor.contains(owner);96 return scf::SCFTileAndFuseOptions::ControlFnResult{97 yieldProducerReplacement};98 };99 tileAndFuseOptions.setFusionControlFn(controlFn);100 101 rewriter.setInsertionPoint(target);102 FailureOr<scf::SCFTileAndFuseResult> tiledResults =103 scf::tileConsumerAndFuseProducersUsingSCF(rewriter, tilingInterfaceOp,104 tileAndFuseOptions);105 if (failed(tiledResults))106 return failure();107 108 // Perform the replacement of tiled and fused values.109 SmallVector<Operation *> opsToReplace{target};110 llvm::append_range(opsToReplace, tiledResults->fusedProducers);111 for (Operation *toReplace : opsToReplace) {112 for (OpResult res : toReplace->getResults())113 if (auto replacement = tiledResults->replacements.lookup(res)) {114 Operation *replacementOp = replacement.getDefiningOp();115 rewriter.replaceUsesWithIf(res, replacement, [&](OpOperand &use) {116 Operation *user = use.getOwner();117 return dominanceInfo.properlyDominates(replacementOp, user) &&118 user->getParentOp() == replacementOp->getParentOp();119 });120 }121 122 if (toReplace->use_empty()) {123 rewriter.eraseOp(toReplace);124 }125 }126 127 // Report back the relevant handles to the transform op.128 tiledOps.push_back(tiledResults->tiledAndFusedOps.front());129 assert(tiledResults->loops.size() == numLoops &&130 "Mismatched number of loops, tile and fuse transform should have "131 "failed");132 for (unsigned int i = 0; i < numLoops; ++i)133 loopOps[i].push_back(tiledResults->loops[i]);134 }135 136 transformResults.set(transformOp->getOpResult(0), tiledOps);137 for (unsigned int i = 0; i < numLoops; ++i)138 transformResults.set(transformOp->getOpResult(i + 1), loopOps[i]);139 140 return success();141}142 143DiagnosedSilenceableFailure144transform::TestFuseAndYieldOp::apply(TransformRewriter &rewriter,145 TransformResults &transformResults,146 TransformState &state) {147 SmallVector<int64_t> tileSizes =148 extractFromIntegerArrayAttr<int64_t>(getTileSizes());149 SmallVector<int64_t> tileInterchange =150 extractFromIntegerArrayAttr<int64_t>(getTileInterchange());151 152 SmallVector<OpFoldResult> tileSizesOfr =153 getAsIndexOpFoldResult(rewriter.getContext(), tileSizes);154 155 scf::SCFTilingOptions tilingOptions;156 tilingOptions.setTileSizes(tileSizesOfr).setInterchange(tileInterchange);157 if (getUseForall()) {158 tilingOptions.setLoopType(scf::SCFTilingOptions::LoopType::ForallOp);159 }160 161 LogicalResult result = applyTileAndFuseToAll(162 rewriter, getOperation(), state.getPayloadOps(getTarget()),163 tileSizes.size() - llvm::count(tileSizes, 0), tilingOptions,164 transformResults);165 return failed(result) ? DiagnosedSilenceableFailure::definiteFailure()166 : DiagnosedSilenceableFailure::success();167}168 169//===----------------------------------------------------------------------===//170// TestFuseConsumerOp171//===----------------------------------------------------------------------===//172 173/// Fuse the consumer and store both the original consumer operation as well as174/// the fused consumer operation.175static LogicalResult176applyFuseConsumer(RewriterBase &rewriter, Operation *transformOp,177 Operation *consumer,178 MutableArrayRef<LoopLikeOpInterface> loops,179 TransformResults &transformResults) {180 SmallVector<Operation *> fusedConsumerOps;181 rewriter.setInsertionPoint(consumer);182 183 FailureOr<scf::SCFFuseConsumerOfSliceResult> fuseConsumerResults =184 scf::tileAndFuseConsumer(rewriter, consumer, loops);185 if (failed(fuseConsumerResults))186 return consumer->emitOpError("failed to fuse consumer of slice");187 188 // Report back the relevant handles to the transform op.189 for (OpOperand *tiledAndFusedConsumerOperand :190 fuseConsumerResults->tiledAndFusedConsumerOperands) {191 fusedConsumerOps.push_back(tiledAndFusedConsumerOperand->getOwner());192 }193 transformResults.set(transformOp->getOpResult(0), fusedConsumerOps);194 for (auto [index, loop] : llvm::enumerate(loops)) {195 transformResults.set(transformOp->getOpResult(index + 1), {loop});196 }197 return success();198}199 200DiagnosedSilenceableFailure201transform::TestFuseConsumerOp::apply(TransformRewriter &rewriter,202 TransformResults &transformResults,203 TransformState &state) {204 Operation *consumer = *state.getPayloadOps(getConsumer()).begin();205 206 SmallVector<LoopLikeOpInterface> loops;207 // Since the matcher works inside-out, we need to iterate the loops in208 // reverse.209 for (auto loop : llvm::reverse(getLoops())) {210 auto loopLikeOp =211 dyn_cast<LoopLikeOpInterface>(*state.getPayloadOps(loop).begin());212 if (!loopLikeOp) {213 return DiagnosedSilenceableFailure::definiteFailure();214 }215 loops.push_back(loopLikeOp);216 }217 LogicalResult result = applyFuseConsumer(rewriter, getOperation(), consumer,218 loops, transformResults);219 return failed(result) ? DiagnosedSilenceableFailure::definiteFailure()220 : DiagnosedSilenceableFailure::success();221}222 223void transform::TestFuseConsumerOp::getEffects(224 SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {225 consumesHandle(getConsumerMutable(), effects);226 consumesHandle(getLoopsMutable(), effects);227 producesHandle(getOperation()->getOpResults(), effects);228 modifiesPayload(effects);229}230 231//===----------------------------------------------------------------------===//232// TestFuseConsumerUsingSliceOp233//===----------------------------------------------------------------------===//234 235/// Apply fusing of consumer transformation to all payload ops and store both236/// the original consumer operation as well as the fused consumer operation.237static LogicalResult applyFuseConsumerUsingSlices(238 RewriterBase &rewriter, Operation *transformOp,239 ArrayRef<Operation *> slices, MutableArrayRef<LoopLikeOpInterface> loops,240 uint32_t numConsumerToFuse, TransformResults &transformResults) {241 SmallVector<Operation *> originalConsumerOps;242 SmallVector<Operation *> fusedConsumerOps;243 244 rewriter.setInsertionPoint(slices.front());245 246 while (numConsumerToFuse--) {247 FailureOr<scf::SCFFuseConsumerOfSliceResult> fuseConsumerResults =248 scf::tileAndFuseConsumerOfSlices(rewriter, slices, loops);249 250 if (failed(fuseConsumerResults))251 return slices.front()->emitOpError("failed to fuse consumer of slice");252 253 // Report back the relevant handles to the transform op.254 for (OpOperand *origConsumerOperand :255 fuseConsumerResults->origConsumerOperands) {256 originalConsumerOps.push_back(origConsumerOperand->getOwner());257 }258 for (OpOperand *tiledAndFusedConsumerOperand :259 fuseConsumerResults->tiledAndFusedConsumerOperands) {260 fusedConsumerOps.push_back(tiledAndFusedConsumerOperand->getOwner());261 }262 }263 264 transformResults.set(transformOp->getOpResult(0), originalConsumerOps);265 transformResults.set(transformOp->getOpResult(1), fusedConsumerOps);266 return success();267}268 269DiagnosedSilenceableFailure transform::TestFuseConsumerUsingSliceOp::apply(270 TransformRewriter &rewriter, TransformResults &transformResults,271 TransformState &state) {272 SmallVector<Operation *> slices;273 for (auto op : getTargets()) {274 auto sliceOp = *state.getPayloadOps(op).begin();275 slices.push_back(sliceOp);276 }277 278 SmallVector<LoopLikeOpInterface> loops;279 for (auto op : llvm::reverse(getLoops())) {280 auto loopLikeOp =281 dyn_cast<LoopLikeOpInterface>(*state.getPayloadOps(op).begin());282 if (!loopLikeOp) {283 return DiagnosedSilenceableFailure::definiteFailure();284 }285 loops.push_back(loopLikeOp);286 }287 LogicalResult result =288 applyFuseConsumerUsingSlices(rewriter, getOperation(), slices, loops,289 getNumConsumerToFuse(), transformResults);290 return failed(result) ? DiagnosedSilenceableFailure::definiteFailure()291 : DiagnosedSilenceableFailure::success();292}293 294void transform::TestFuseConsumerUsingSliceOp::getEffects(295 SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {296 consumesHandle(getTargetsMutable(), effects);297 consumesHandle(getLoopsMutable(), effects);298 producesHandle(getOperation()->getOpResults(), effects);299 modifiesPayload(effects);300}301 302//===----------------------------------------------------------------------===//303// TestTileUsingForallOp304//===----------------------------------------------------------------------===//305 306/// Apply a tiling transformation to all payload ops and store both the307/// tiled operation as well as the created tile loops.308template <typename Range>309static LogicalResult310applyTileToAll(RewriterBase &rewriter, Operation *transformOp,311 Range &&payloadOps, ArrayRef<OpFoldResult> tileSizes,312 ArrayRef<int64_t> interchange, std::optional<ArrayAttr> mapping,313 TransformResults &transformResults) {314 SmallVector<Operation *> tiledOps;315 SmallVector<Operation *> loopOps;316 317 for (Operation *target : payloadOps) {318 auto tilingInterfaceOp = dyn_cast<TilingInterface>(target);319 if (!tilingInterfaceOp)320 return transformOp->emitError("only TilingInterface ops are supported");321 scf::SCFTilingOptions tilingOptions;322 tilingOptions.setTileSizes(tileSizes).setInterchange(interchange);323 if (mapping) {324 tilingOptions.setMapping(mapping.value().getValue());325 }326 tilingOptions.setLoopType(scf::SCFTilingOptions::LoopType::ForallOp);327 328 rewriter.setInsertionPoint(target);329 FailureOr<scf::SCFTilingResult> tiledResults =330 scf::tileUsingSCF(rewriter, tilingInterfaceOp, tilingOptions);331 if (failed(tiledResults))332 return failure();333 334 // Perform the replacement of tiled and fused values.335 rewriter.replaceOp(tilingInterfaceOp, tiledResults->replacements);336 337 // Report back the relevant handles to the transform op.338 tiledOps.push_back(tiledResults->tiledOps.front());339 for (Operation *loop : tiledResults->loops)340 loopOps.push_back(loop);341 }342 343 transformResults.set(transformOp->getOpResult(0), tiledOps);344 for (auto [index, loop] : llvm::enumerate(loopOps))345 transformResults.set(transformOp->getOpResult(index + 1), {loop});346 347 return success();348}349 350DiagnosedSilenceableFailure351transform::TestTileUsingForallOp::apply(TransformRewriter &rewriter,352 TransformResults &transformResults,353 TransformState &state) {354 SmallVector<int64_t> tileSizes =355 extractFromIntegerArrayAttr<int64_t>(getTileSizes());356 SmallVector<int64_t> interchange =357 extractFromIntegerArrayAttr<int64_t>(getInterchange());358 SmallVector<OpFoldResult> tileSizesOfr =359 getAsIndexOpFoldResult(rewriter.getContext(), tileSizes);360 361 LogicalResult result =362 applyTileToAll(rewriter, getOperation(), state.getPayloadOps(getTarget()),363 tileSizesOfr, interchange, getMapping(), transformResults);364 return failed(result) ? DiagnosedSilenceableFailure::definiteFailure()365 : DiagnosedSilenceableFailure::success();366}367 368void transform::TestTileUsingForallOp::getEffects(369 SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {370 consumesHandle(getTargetMutable(), effects);371 producesHandle(getOperation()->getOpResults(), effects);372 modifiesPayload(effects);373}374 375//===----------------------------------------------------------------------===//376// TestFuseUsingForallOp377//===----------------------------------------------------------------------===//378 379/// Apply a tiling transformation to all payload ops and store both the380/// tiled operation as well as the created tile loops.381template <typename Range>382static LogicalResult applyTilingToAll(383 RewriterBase &rewriter, Operation *transformOp, Range &&payloadOps,384 unsigned numLoops, TransformResults &transformResults,385 function_ref<FailureOr<scf::SCFTileAndFuseResult>(TilingInterface)>386 applyFn) {387 SmallVector<Operation *> tiledLinalgOps;388 SmallVector<SmallVector<Operation *>> loopOps(1);389 390 for (Operation *target : payloadOps) {391 auto tilingInterfaceOp = dyn_cast<TilingInterface>(target);392 if (!tilingInterfaceOp)393 return transformOp->emitError("only TilingInterface ops are supported");394 395 rewriter.setInsertionPoint(target);396 FailureOr<scf::SCFTileAndFuseResult> tiledResults =397 applyFn(tilingInterfaceOp);398 if (failed(tiledResults))399 return failure();400 401 // Perform the replacement of tiled and fused values.402 SmallVector<Operation *> opsToReplace{target};403 llvm::append_range(opsToReplace, tiledResults->fusedProducers);404 for (Operation *toReplace : opsToReplace) {405 for (OpResult res : toReplace->getResults())406 if (auto replacement = tiledResults->replacements.lookup(res))407 rewriter.replaceAllUsesWith(res, replacement);408 if (toReplace->use_empty())409 rewriter.eraseOp(toReplace);410 }411 412 // Report back the relevant handles to the transform op.413 tiledLinalgOps.push_back(tiledResults->tiledAndFusedOps.front());414 assert(tiledResults->loops.size() == 1 &&415 cast<scf::ForallOp>(tiledResults->loops[0]).getRank() == numLoops &&416 "Mismatched number of loops, tile and fuse transform should have "417 "failed");418 loopOps[0] = {tiledResults->loops[0]};419 }420 421 transformResults.set(transformOp->getOpResult(0), tiledLinalgOps);422 if (!loopOps.empty())423 transformResults.set(transformOp->getOpResult(1), loopOps[0]);424 425 return success();426}427 428DiagnosedSilenceableFailure429transform::TestFuseUsingForallOp::apply(TransformRewriter &rewriter,430 TransformResults &transformResults,431 TransformState &state) {432 SmallVector<int64_t> tileSizes =433 extractFromIntegerArrayAttr<int64_t>(getTileSizes());434 SmallVector<int64_t> tileInterchange =435 extractFromIntegerArrayAttr<int64_t>(getInterchange());436 437 scf::SCFTilingOptions tilingOptions;438 tilingOptions.interchangeVector = tileInterchange;439 SmallVector<OpFoldResult> tileSizesOfr =440 getAsIndexOpFoldResult(rewriter.getContext(), tileSizes);441 tilingOptions = tilingOptions.setTileSizes(tileSizesOfr);442 tilingOptions.setLoopType(scf::SCFTilingOptions::LoopType::ForallOp);443 scf::SCFTileAndFuseOptions tileAndFuseOptions;444 tileAndFuseOptions.tilingOptions = tilingOptions;445 LogicalResult result = applyTilingToAll(446 rewriter, getOperation(), state.getPayloadOps(getRootOp()),447 tileSizes.size() - llvm::count(tileSizes, 0), transformResults,448 [&](TilingInterface tilingInterfaceOp)449 -> FailureOr<scf::SCFTileAndFuseResult> {450 return tileConsumerAndFuseProducersUsingSCF(rewriter, tilingInterfaceOp,451 tileAndFuseOptions);452 });453 return failed(result) ? DiagnosedSilenceableFailure::definiteFailure()454 : DiagnosedSilenceableFailure::success();455}456 457void transform::TestFuseUsingForallOp::getEffects(458 SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {459 consumesHandle(getRootOpMutable(), effects);460 producesHandle(getOperation()->getOpResults(), effects);461 modifiesPayload(effects);462}463 464//===----------------------------------------------------------------------===//465// TestTileAndFuseOuterParallelPartialReduction466//===----------------------------------------------------------------------===//467 468DiagnosedSilenceableFailure469transform::TestTileAndFuseOuterParallelPartialReductionOp::apply(470 TransformRewriter &rewriter, TransformResults &transformResults,471 TransformState &state) {472 auto target =473 dyn_cast<TilingInterface>(*state.getPayloadOps(getRootOp()).begin());474 if (!target) {475 emitOpError("expected root operation to implement `TilingInterface`");476 return DiagnosedSilenceableFailure::definiteFailure();477 }478 479 SmallVector<unsigned> reductionDims =480 extractFromIntegerArrayAttr<unsigned>(getReductionDims());481 if (reductionDims.empty()) {482 for (auto [index, iterator] :483 llvm::enumerate(target.getLoopIteratorTypes()))484 if (iterator == utils::IteratorType::reduction)485 reductionDims.push_back(index);486 }487 488 if (reductionDims.empty()) {489 emitOpError(490 "no reduction dimension specified or found in the target operation");491 return DiagnosedSilenceableFailure::definiteFailure();492 }493 494 SmallVector<int64_t> reductionTileSizes =495 extractFromIntegerArrayAttr<int64_t>(getTileSizes());496 if (reductionTileSizes.size() != reductionDims.size()) {497 emitOpError(498 "missing tile sizes for reduction dimensions that are to be tiled");499 return DiagnosedSilenceableFailure::definiteFailure();500 }501 502 // Adjust tile sizes so that it corresponds to the reduction iterator types.503 SmallVector<OpFoldResult> tileSizes;504 int reductionTileSizeNum = 0;505 OpFoldResult zero = rewriter.getIndexAttr(0);506 for (auto iterator : target.getLoopIteratorTypes()) {507 if (iterator == utils::IteratorType::parallel) {508 tileSizes.push_back(zero);509 continue;510 }511 tileSizes.push_back(512 rewriter.getIndexAttr(reductionTileSizes[reductionTileSizeNum++]));513 }514 515 scf::SCFTilingOptions tilingOptions;516 tilingOptions.setTileSizes(tileSizes)517 .setLoopType(scf::SCFTilingOptions::LoopType::ForallOp)518 .setReductionTilingStrategy(519 ReductionTilingStrategy::PartialReductionOuterParallel)520 .setReductionDims(reductionDims);521 if (auto mapping = getMapping()) {522 tilingOptions.setMapping(getMapping().value());523 }524 525 LogicalResult result = applyTileAndFuseToAll(526 rewriter, getOperation(), state.getPayloadOps(getRootOp()),527 /*numLoops =*/1, tilingOptions, transformResults);528 529 return failed(result) ? DiagnosedSilenceableFailure::definiteFailure()530 : DiagnosedSilenceableFailure::success();531}532 533//===----------------------------------------------------------------------===//534// TestTileAndFuseOuterParallelPartialReduction535//===----------------------------------------------------------------------===//536 537DiagnosedSilenceableFailure transform::TestTileUsingCustomLoopOp::apply(538 TransformRewriter &transformRewriter, TransformResults &transformResults,539 TransformState &state) {540 auto target =541 dyn_cast<TilingInterface>(*state.getPayloadOps(getRootOp()).begin());542 if (!target) {543 emitOpError("expected root operation to implement `TilingInterface`");544 return DiagnosedSilenceableFailure::definiteFailure();545 }546 547 OpFoldResult oneOfr = transformRewriter.getIndexAttr(1);548 549 scf::SCFTilingOptions::GenerateLoopHeaderFn loopHeaderFn =550 [&](RewriterBase &rewriter, Location loc, ArrayRef<Range> loopRanges,551 ArrayRef<OpFoldResult> givenTileSizes,552 ValueRange outerDestinationTensors)553 -> FailureOr<scf::SCFTilingOptions::CustomLoopHeaderInfo> {554 // Check that the strides are all 1 (to make it easier in the test).555 if (llvm::any_of(loopRanges, [](Range r) {556 return !isConstantIntValue(r.stride, 1);557 })) {558 return emitOpError("unable to handle loop ranges with strides != 1");559 }560 // Check number of tile sizes is equal to loop dimensions.561 if (loopRanges.size() != givenTileSizes.size()) {562 return emitOpError("expected number of tile sizes to be same as the "563 "number of loops in the operation");564 }565 // For testing disallow any of the tile sizes being 0.566 if (llvm::any_of(givenTileSizes, isZeroInteger)) {567 return emitOpError("unhandled case of zero tile size");568 }569 // For testing, only handle tensor tiling.570 if (outerDestinationTensors.empty()) {571 return emitOpError("expected destination tensors");572 }573 574 // Compute the number of iterations for each of the loops.575 AffineExpr s0, s1, s2;576 bindSymbols(rewriter.getContext(), s0, s1, s2);577 AffineExpr numItersExpr = (s1 - s0).ceilDiv(s2); // (ub - lb) / tileSize578 579 SmallVector<OpFoldResult> allNumIters;580 allNumIters.reserve(loopRanges.size());581 for (auto [loopRange, tileSize] :582 llvm::zip_equal(loopRanges, givenTileSizes)) {583 OpFoldResult numIters = affine::makeComposedFoldedAffineApply(584 rewriter, loc, numItersExpr,585 {loopRange.offset, loopRange.size, tileSize});586 allNumIters.push_back(numIters);587 }588 if (allNumIters.empty()) {589 return emitOpError("invalid empty tile sizes and loop ranges");590 }591 592 AffineExpr mulExpr = s0 * s1;593 OpFoldResult cumulative = oneOfr;594 for (auto numIters : allNumIters) {595 cumulative = affine::makeComposedFoldedAffineApply(596 rewriter, loc, mulExpr, {cumulative, numIters});597 }598 599 Value zeroVal = arith::ConstantIndexOp::create(rewriter, loc, 0);600 Value oneVal = arith::ConstantIndexOp::create(rewriter, loc, 1);601 Value ub = getValueOrCreateConstantIndexOp(rewriter, loc, cumulative);602 603 SmallVector<OpFoldResult> offsets;604 SmallVector<OpFoldResult> sizes;605 SmallVector<Value> innerDestinationTensors;606 offsets.reserve(loopRanges.size());607 sizes.reserve(loopRanges.size());608 609 AffineExpr d0;610 bindDims(rewriter.getContext(), d0);611 AffineExpr offsetExpr = s0 + d0 * s1; // lb + iv * tileSize612 AffineMap minMap =613 AffineMap::get(1, 2, {s0 - d0, s1},614 rewriter.getContext()); // min(ub - offset, tileSize)615 auto forOp = scf::ForOp::create(616 rewriter, loc, zeroVal, ub, oneVal, outerDestinationTensors,617 [&](OpBuilder &b, Location bodyLoc, Value linearizedIv,618 ValueRange destinations) {619 auto delinearizeOp = affine::AffineDelinearizeIndexOp::create(620 b, bodyLoc, linearizedIv, allNumIters);621 for (auto [normalizedIv, range, tileSize] : llvm::zip_equal(622 delinearizeOp.getResults(), loopRanges, givenTileSizes)) {623 624 OpFoldResult normalizedIvOfr = getAsOpFoldResult(normalizedIv);625 OpFoldResult offset = affine::makeComposedFoldedAffineApply(626 b, bodyLoc, offsetExpr,627 {normalizedIvOfr, range.offset, tileSize});628 offsets.push_back(offset);629 630 OpFoldResult size = affine::makeComposedFoldedAffineMin(631 b, bodyLoc, minMap, {offset, range.size, tileSize});632 sizes.push_back(size);633 }634 innerDestinationTensors = llvm::to_vector(destinations);635 });636 rewriter.setInsertionPointToEnd(forOp.getBody());637 return scf::SCFTilingOptions::CustomLoopHeaderInfo{638 {cast<LoopLikeOpInterface>(forOp.getOperation())},639 offsets,640 sizes,641 innerDestinationTensors};642 };643 644 scf::SCFTilingOptions::GenerateLoopTerminatorFn terminatorFn =645 [&](RewriterBase &rewriter, Location loc,646 ArrayRef<LoopLikeOpInterface> loops, ValueRange tiledResults,647 ArrayRef<SmallVector<OpFoldResult>> resultOffsets,648 ArrayRef<SmallVector<OpFoldResult>> resultSizes,649 ValueRange destinationTensors) -> LogicalResult {650 SmallVector<Value> yieldValues;651 yieldValues.reserve(destinationTensors.size());652 for (auto [tiledResult, offsets, sizes, destination] : llvm::zip_equal(653 tiledResults, resultOffsets, resultSizes, destinationTensors)) {654 SmallVector<OpFoldResult> strides(offsets.size(), oneOfr);655 Value insertedVal = tensor::InsertSliceOp::create(656 rewriter, loc, tiledResult, destination, offsets, sizes, strides);657 yieldValues.push_back(insertedVal);658 }659 scf::YieldOp::create(rewriter, loc, yieldValues);660 return success();661 };662 663 scf::SCFTilingOptions tilingOptions;664 SmallVector<int64_t> staticTileSizes =665 extractFromIntegerArrayAttr<int64_t>(getTileSizes());666 SmallVector<OpFoldResult> tileSizes =667 getAsIndexOpFoldResult(transformRewriter.getContext(), staticTileSizes);668 tilingOptions.setTileSizes(tileSizes)669 .setLoopType(scf::SCFTilingOptions::LoopType::CustomOp)670 .setCustomLoopGenerationFns(loopHeaderFn, terminatorFn);671 672 OpBuilder::InsertionGuard g(transformRewriter);673 transformRewriter.setInsertionPoint(target);674 FailureOr<scf::SCFTilingResult> tiledResults =675 scf::tileUsingSCF(transformRewriter, target, tilingOptions);676 if (failed(tiledResults)) {677 return DiagnosedSilenceableFailure::definiteFailure();678 }679 transformRewriter.replaceOp(target, tiledResults->replacements);680 transformResults.set(getOperation()->getResult(0), tiledResults->tiledOps);681 transformResults.set(getOperation()->getResult(1), tiledResults->loops);682 683 return DiagnosedSilenceableFailure::success();684}685 686#define GET_OP_CLASSES687#include "TestTilingInterfaceTransformOps.cpp.inc"688 689namespace {690class TestTilingInterfaceDialectExtension691 : public transform::TransformDialectExtension<692 TestTilingInterfaceDialectExtension> {693public:694 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(695 TestTilingInterfaceDialectExtension)696 697 using Base::Base;698 699 void init() {700 declareDependentDialect<affine::AffineDialect>();701 declareDependentDialect<index::IndexDialect>();702 declareDependentDialect<scf::SCFDialect>();703 declareDependentDialect<tensor::TensorDialect>();704 705 registerTransformOps<706#define GET_OP_LIST707#include "TestTilingInterfaceTransformOps.cpp.inc"708 >();709 }710};711} // namespace712 713namespace test {714void registerTestTilingInterfaceTransformDialectExtension(715 DialectRegistry ®istry) {716 registry.addExtensions<TestTilingInterfaceDialectExtension>();717}718} // namespace test719