This is an archive of the discontinued LLVM Phabricator instance.

Differential D148687

[mlir][Affine][NFC] Wrap dialect in "affine" namespace
ClosedPublic

Authored by springerm on Apr 18 2023, 11:41 PM.

Details

Reviewers
mehdi_amini
aartbik
mravishankar
bondhugula
ThomasRaoux
nicolasvasilache
herhut
dcaballe
rriddle
Commits
rG914b38aa4b9a: Revert '[RISCV] Remove 'sx' when parsing arch string to match the latest ISA…
rG4c48f016effd: [mlir][Affine][NFC] Wrap dialect in "affine" namespace

Diff Detail

Event Timeline

springerm created this revision.Apr 18 2023, 11:41 PM
Herald added a reviewer: aartbik. · View Herald TranscriptApr 18 2023, 11:41 PM
Herald added a reviewer: aartbik. · View Herald Transcript
Herald added a reviewer: mravishankar. · View Herald Transcript
Herald added a reviewer: bondhugula. · View Herald Transcript
Herald added a reviewer: ThomasRaoux. · View Herald Transcript
Herald added projects: Restricted Project, Restricted Project. · View Herald Transcript
Herald added subscribers: bviyer, hanchung, jsetoain and 31 others. · View Herald Transcript
springerm requested review of this revision.Apr 18 2023, 11:41 PM
Herald added a reviewer: nicolasvasilache. · View Herald TranscriptApr 18 2023, 11:41 PM
Herald added a reviewer: herhut. · View Herald Transcript
Herald added a reviewer: dcaballe. · View Herald Transcript
Herald added a project: Restricted Project. · View Herald Transcript
Herald added subscribers: • pcwang-thead, stephenneuendorffer, nicolasvasilache, jdoerfert. · View Herald Transcript
springerm added a reviewer: rriddle.Apr 18 2023, 11:48 PM
Harbormaster completed remote builds in B226523: Diff 514833.Apr 19 2023, 12:09 AM
springerm updated this revision to Diff 514846.Apr 19 2023, 12:32 AM

update

springerm updated this revision to Diff 514861.Apr 19 2023, 1:37 AM

rebase

nicolasvasilache accepted this revision.Apr 19 2023, 1:58 AM
This revision is now accepted and ready to land.Apr 19 2023, 1:58 AM
springerm updated this revision to Diff 514867.Apr 19 2023, 2:01 AM

update

Harbormaster completed remote builds in B226548: Diff 514867.Apr 19 2023, 2:41 AM
mehdi_amini accepted this revision.Apr 19 2023, 10:27 AM

Thanks!

Can you add a description with the motivation, even if only: "This cleanup aligns the affine dialect with all the other dialects." or something like that.

dcaballe accepted this revision.Apr 19 2023, 11:18 AM
springerm updated this revision to Diff 515177.Apr 19 2023, 7:11 PM

fix flang

springerm updated this revision to Diff 515179.Apr 19 2023, 7:14 PM

fix build

This revision was landed with ongoing or failed builds.Apr 19 2023, 7:20 PM
Closed by commit rG4c48f016effd: [mlir][Affine][NFC] Wrap dialect in "affine" namespace (authored by springerm). · Explain Why
This revision was automatically updated to reflect the committed changes.
springerm added a commit: rG4c48f016effd: [mlir][Affine][NFC] Wrap dialect in "affine" namespace.
joshua-arch1 added a commit: rG914b38aa4b9a: Revert '[RISCV] Remove 'sx' when parsing arch string to match the latest ISA….Apr 19 2023, 7:27 PM
bondhugula added a comment.Apr 19 2023, 7:27 PM

Thanks for doing this. This had been long pending. This also cleans up the naming/addressing of free-standing methods that were otherwise in the mlir namespace. Having mlir::affine::... for these is a big naming and readability improvement.

Harbormaster completed remote builds in B226766: Diff 515179.Apr 19 2023, 7:28 PM
ubfx added a subscriber: ubfx.Apr 22 2023, 12:46 AM

Will the Operations from the Affine dialect also be renamed? I feel like affine::AffineForOp could be renamed to affine::ForOp now.

Revision Contents

PathSize
flang/
include/
flang/
Optimizer/
Support/
24 lines
Transforms/
6 lines
lib/
Optimizer/
Transforms/
22 lines
34 lines
2 lines
5 lines
5 lines
mlir/
docs/
3 lines
Tutorials/
Toy/
2 lines
examples/
toy/
Ch5/
mlir/
23 lines
4 lines
Ch6/
mlir/
23 lines
4 lines
Ch7/
mlir/
23 lines
4 lines
include/
mlir/
Conversion/
AffineToStandard/
10 lines
6 lines
SCFToGPU/
7 lines
Dialect/
Affine/
Analysis/
4 lines
18 lines
8 lines
4 lines
10 lines
IR/
2 lines
64 lines
11 lines
2 lines
2 lines
6 lines
7 lines
2 lines
34 lines
TransformOps/
3 lines
Transforms/
5 lines
10 lines
3 lines
Bufferization/
IR/
2 lines
Linalg/
IR/
2 lines
12 lines
Utils/
5 lines
MemRef/
Transforms/
8 lines
SCF/
Transforms/
6 lines
Utils/
8 lines
SparseTensor/
Transforms/
2 lines
Tensor/
IR/
2 lines
Transforms/
2 lines
Vector/
Utils/
7 lines
2 lines
2 lines
lib/
Conversion/
AffineToStandard/
1 line
LinalgToStandard/
5 lines
SCFToGPU/
1 line
2 lines
VectorToGPU/
2 lines
VectorToSCF/
9 lines
Dialect/
Affine/
Analysis/
19 lines
9 lines
21 lines
3 lines
76 lines
IR/
1 line
114 lines
1 line
1 line
TransformOps/
1 line
Transforms/
16 lines
10 lines
31 lines
7 lines
7 lines
8 lines
9 lines
8 lines
22 lines
12 lines
7 lines
7 lines
10 lines
9 lines
8 lines
26 lines
Utils/
45 lines
123 lines
83 lines
9 lines
GPU/
TransformOps/
10 lines
Transforms/
4 lines
Linalg/
IR/
2 lines
2 lines
TransformOps/
12 lines
Transforms/
2 lines
5 lines
4 lines
6 lines
9 lines
4 lines
2 lines
26 lines
6 lines
15 lines
8 lines
8 lines
Utils/
16 lines
MemRef/
TransformOps/
2 lines
Transforms/
4 lines
2 lines
4 lines
72 lines
2 lines
1 line
SCF/
TransformOps/
3 lines
Transforms/
4 lines
1 line
4 lines
2 lines
Utils/
3 lines
SparseTensor/
Transforms/
4 lines
Tensor/
IR/
10 lines
14 lines
24 lines
Transforms/
2 lines
3 lines
12 lines
6 lines
Utils/
3 lines
Vector/
Transforms/
27 lines
4 lines
10 lines
3 lines
Utils/
6 lines
test/
lib/
Analysis/
1 line
1 line
Dialect/
Affine/
1 line
1 line
1 line
1 line
18 lines
3 lines
1 line
5 lines
9 lines
GPU/
3 lines
Linalg/
4 lines
2 lines
4 lines
4 lines
2 lines
4 lines
MemRef/
2 lines
2 lines
Vector/
16 lines
Interfaces/
TilingInterface/
5 lines
unittests/
Analysis/
Presburger/
2 lines

Diff 515180

flang/include/flang/Optimizer/Support/InitFIR.h

Show All 21 Lines
#include "mlir/Pass/Pass.h" #include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassRegistry.h" #include "mlir/Pass/PassRegistry.h"
#include "mlir/Transforms/LocationSnapshot.h" #include "mlir/Transforms/LocationSnapshot.h"
#include "mlir/Transforms/Passes.h" #include "mlir/Transforms/Passes.h"
namespace fir::support { namespace fir::support {
#define FLANG_NONCODEGEN_DIALECT_LIST \ #define FLANG_NONCODEGEN_DIALECT_LIST \
mlir::AffineDialect, FIROpsDialect, hlfir::hlfirDialect, \ mlir::affine::AffineDialect, FIROpsDialect, hlfir::hlfirDialect, \
mlir::acc::OpenACCDialect, mlir::omp::OpenMPDialect, \ mlir::acc::OpenACCDialect, mlir::omp::OpenMPDialect, \
mlir::scf::SCFDialect, mlir::arith::ArithDialect, \ mlir::scf::SCFDialect, mlir::arith::ArithDialect, \
mlir::cf::ControlFlowDialect, mlir::func::FuncDialect, \ mlir::cf::ControlFlowDialect, mlir::func::FuncDialect, \
mlir::vector::VectorDialect, mlir::math::MathDialect, \ mlir::vector::VectorDialect, mlir::math::MathDialect, \
mlir::complex::ComplexDialect, mlir::DLTIDialect mlir::complex::ComplexDialect, mlir::DLTIDialect
// The definitive list of dialects used by flang. // The definitive list of dialects used by flang.
#define FLANG_DIALECT_LIST \ #define FLANG_DIALECT_LIST \
Show All 19 Linesinline void loadDialects(mlir::MLIRContext &context) {
context.loadDialect<FLANG_DIALECT_LIST>(); context.loadDialect<FLANG_DIALECT_LIST>();
} }
/// Register the standard passes we use. This comes from registerAllPasses(), /// Register the standard passes we use. This comes from registerAllPasses(),
/// but is a smaller set since we aren't using many of the passes found there. /// but is a smaller set since we aren't using many of the passes found there.
inline void registerMLIRPassesForFortranTools() { inline void registerMLIRPassesForFortranTools() {
mlir::registerCanonicalizerPass(); mlir::registerCanonicalizerPass();
mlir::registerCSEPass(); mlir::registerCSEPass();
mlir::registerAffineLoopFusionPass(); mlir::affine::registerAffineLoopFusionPass();
mlir::registerLoopInvariantCodeMotionPass(); mlir::registerLoopInvariantCodeMotionPass();
mlir::registerLoopCoalescingPass(); mlir::affine::registerLoopCoalescingPass();
mlir::registerStripDebugInfoPass(); mlir::registerStripDebugInfoPass();
mlir::registerPrintOpStatsPass(); mlir::registerPrintOpStatsPass();
mlir::registerInlinerPass(); mlir::registerInlinerPass();
mlir::registerSCCPPass(); mlir::registerSCCPPass();
mlir::registerAffineScalarReplacementPass(); mlir::affine::registerAffineScalarReplacementPass();
mlir::registerSymbolDCEPass(); mlir::registerSymbolDCEPass();
mlir::registerLocationSnapshotPass(); mlir::registerLocationSnapshotPass();
mlir::registerAffinePipelineDataTransferPass(); mlir::affine::registerAffinePipelineDataTransferPass();
mlir::registerAffineVectorizePass(); mlir::affine::registerAffineVectorizePass();
mlir::registerAffineLoopUnrollPass(); mlir::affine::registerAffineLoopUnrollPass();
mlir::registerAffineLoopUnrollAndJamPass(); mlir::affine::registerAffineLoopUnrollAndJamPass();
mlir::registerSimplifyAffineStructuresPass(); mlir::affine::registerSimplifyAffineStructuresPass();
mlir::registerAffineLoopInvariantCodeMotionPass(); mlir::affine::registerAffineLoopInvariantCodeMotionPass();
mlir::registerAffineLoopTilingPass(); mlir::affine::registerAffineLoopTilingPass();
mlir::registerAffineDataCopyGenerationPass(); mlir::affine::registerAffineDataCopyGenerationPass();
mlir::registerConvertAffineToStandardPass(); mlir::registerConvertAffineToStandardPass();
} }
/// Register the interfaces needed to lower to LLVM IR. /// Register the interfaces needed to lower to LLVM IR.
void registerLLVMTranslation(mlir::MLIRContext &context); void registerLLVMTranslation(mlir::MLIRContext &context);
} // namespace fir::support } // namespace fir::support
#endif // FORTRAN_OPTIMIZER_SUPPORT_INITFIR_H #endif // FORTRAN_OPTIMIZER_SUPPORT_INITFIR_H

flang/include/flang/Optimizer/Transforms/Passes.td

Show First 20 Lines • Show All 58 Lines • ▼ Show 20 Lines let description = [{
`fir.load` (`fir.store`) will be converted to `affine.load` (`affine.store`) `fir.load` (`fir.store`) will be converted to `affine.load` (`affine.store`)
where possible. This conversion includes adding a dummy `fir.convert` cast where possible. This conversion includes adding a dummy `fir.convert` cast
to adapt values of type `!fir.ref<!fir.array>` to `memref`. This is done to adapt values of type `!fir.ref<!fir.array>` to `memref`. This is done
because the affine dialect presently only understands the `memref` type. because the affine dialect presently only understands the `memref` type.
}]; }];
let constructor = "::fir::createPromoteToAffinePass()"; let constructor = "::fir::createPromoteToAffinePass()";
let dependentDialects = [ let dependentDialects = [
"fir::FIROpsDialect", "mlir::func::FuncDialect", "mlir::AffineDialect" "fir::FIROpsDialect", "mlir::func::FuncDialect",
"mlir::affine::AffineDialect"
]; ];
} }
def AffineDialectDemotion : Pass<"demote-affine", "::mlir::func::FuncOp"> { def AffineDialectDemotion : Pass<"demote-affine", "::mlir::func::FuncOp"> {
let summary = "Converts `affine.{load,store}` back to fir operations"; let summary = "Converts `affine.{load,store}` back to fir operations";
let description = [{ let description = [{
Affine dialect's default lowering for loads and stores is different from Affine dialect's default lowering for loads and stores is different from
fir as it uses the `memref` type. The `memref` type is not compatible with fir as it uses the `memref` type. The `memref` type is not compatible with
the Fortran runtime. Therefore, conversion of memory operations back to the Fortran runtime. Therefore, conversion of memory operations back to
`fir.load` and `fir.store` with `!fir.ref<?>` types is required. `fir.load` and `fir.store` with `!fir.ref<?>` types is required.
}]; }];
let constructor = "::fir::createAffineDemotionPass()"; let constructor = "::fir::createAffineDemotionPass()";
let dependentDialects = [ let dependentDialects = [
"fir::FIROpsDialect", "mlir::func::FuncDialect", "mlir::AffineDialect" "fir::FIROpsDialect", "mlir::func::FuncDialect",
"mlir::affine::AffineDialect"
]; ];
} }
def AnnotateConstantOperands : Pass<"annotate-constant"> { def AnnotateConstantOperands : Pass<"annotate-constant"> {
let summary = "Annotate constant operands to all FIR operations"; let summary = "Annotate constant operands to all FIR operations";
let description = [{ let description = [{
The MLIR canonicalizer makes a distinction between constants based on how The MLIR canonicalizer makes a distinction between constants based on how
they are packaged in the IR. A constant value is wrapped in an Attr and that they are packaged in the IR. A constant value is wrapped in an Attr and that
▲ Show 20 Lines • Show All 220 Lines • Show Last 20 Lines

flang/lib/Optimizer/Transforms/AffineDemotion.cpp

Show All 40 Lines
#define DEBUG_TYPE "flang-affine-demotion" #define DEBUG_TYPE "flang-affine-demotion"
using namespace fir; using namespace fir;
using namespace mlir; using namespace mlir;
namespace { namespace {
class AffineLoadConversion : public OpConversionPattern<mlir::AffineLoadOp> { class AffineLoadConversion
: public OpConversionPattern<mlir::affine::AffineLoadOp> {
public: public:
using OpConversionPattern<mlir::AffineLoadOp>::OpConversionPattern; using OpConversionPattern<mlir::affine::AffineLoadOp>::OpConversionPattern;
LogicalResult LogicalResult
matchAndRewrite(mlir::AffineLoadOp op, OpAdaptor adaptor, matchAndRewrite(mlir::affine::AffineLoadOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override { ConversionPatternRewriter &rewriter) const override {
SmallVector<Value> indices(adaptor.getIndices()); SmallVector<Value> indices(adaptor.getIndices());
auto maybeExpandedMap = auto maybeExpandedMap = affine::expandAffineMap(rewriter, op.getLoc(),
expandAffineMap(rewriter, op.getLoc(), op.getAffineMap(), indices); op.getAffineMap(), indices);
if (!maybeExpandedMap) if (!maybeExpandedMap)
return failure(); return failure();
auto coorOp = rewriter.create<fir::CoordinateOp>( auto coorOp = rewriter.create<fir::CoordinateOp>(
op.getLoc(), fir::ReferenceType::get(op.getResult().getType()), op.getLoc(), fir::ReferenceType::get(op.getResult().getType()),
adaptor.getMemref(), *maybeExpandedMap); adaptor.getMemref(), *maybeExpandedMap);
rewriter.replaceOpWithNewOp<fir::LoadOp>(op, coorOp.getResult()); rewriter.replaceOpWithNewOp<fir::LoadOp>(op, coorOp.getResult());
return success(); return success();
} }
}; };
class AffineStoreConversion : public OpConversionPattern<mlir::AffineStoreOp> { class AffineStoreConversion
: public OpConversionPattern<mlir::affine::AffineStoreOp> {
public: public:
using OpConversionPattern<mlir::AffineStoreOp>::OpConversionPattern; using OpConversionPattern<mlir::affine::AffineStoreOp>::OpConversionPattern;
LogicalResult LogicalResult
matchAndRewrite(mlir::AffineStoreOp op, OpAdaptor adaptor, matchAndRewrite(mlir::affine::AffineStoreOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override { ConversionPatternRewriter &rewriter) const override {
SmallVector<Value> indices(op.getIndices()); SmallVector<Value> indices(op.getIndices());
auto maybeExpandedMap = auto maybeExpandedMap = affine::expandAffineMap(rewriter, op.getLoc(),
expandAffineMap(rewriter, op.getLoc(), op.getAffineMap(), indices); op.getAffineMap(), indices);
if (!maybeExpandedMap) if (!maybeExpandedMap)
return failure(); return failure();
auto coorOp = rewriter.create<fir::CoordinateOp>( auto coorOp = rewriter.create<fir::CoordinateOp>(
op.getLoc(), fir::ReferenceType::get(op.getValueToStore().getType()), op.getLoc(), fir::ReferenceType::get(op.getValueToStore().getType()),
adaptor.getMemref(), *maybeExpandedMap); adaptor.getMemref(), *maybeExpandedMap);
rewriter.replaceOpWithNewOp<fir::StoreOp>(op, adaptor.getValue(), rewriter.replaceOpWithNewOp<fir::StoreOp>(op, adaptor.getValue(),
coorOp.getResult()); coorOp.getResult());
▲ Show 20 Lines • Show All 92 Lines • Show Last 20 Lines

flang/lib/Optimizer/Transforms/AffinePromotion.cpp

Show First 20 Lines • Show All 221 Lines • ▼ Show 20 Lines if (auto op = value.getDefiningOp<mlir::arith::RemUIOp>())
return affineBinaryOp(mlir::AffineExprKind::Mod, op.getLhs(), return affineBinaryOp(mlir::AffineExprKind::Mod, op.getLhs(),
op.getRhs()); op.getRhs());
if (auto op = value.getDefiningOp<mlir::arith::ConstantOp>()) if (auto op = value.getDefiningOp<mlir::arith::ConstantOp>())
if (auto intConstant = op.getValue().dyn_cast<IntegerAttr>()) if (auto intConstant = op.getValue().dyn_cast<IntegerAttr>())
return toAffineExpr(intConstant.getInt()); return toAffineExpr(intConstant.getInt());
if (auto blockArg = value.dyn_cast<mlir::BlockArgument>()) { if (auto blockArg = value.dyn_cast<mlir::BlockArgument>()) {
affineArgs.push_back(value); affineArgs.push_back(value);
if (isa<fir::DoLoopOp>(blockArg.getOwner()->getParentOp()) || if (isa<fir::DoLoopOp>(blockArg.getOwner()->getParentOp()) ||
isa<mlir::AffineForOp>(blockArg.getOwner()->getParentOp())) isa<mlir::affine::AffineForOp>(blockArg.getOwner()->getParentOp()))
return {mlir::getAffineDimExpr(dimCount++, value.getContext())}; return {mlir::getAffineDimExpr(dimCount++, value.getContext())};
return {mlir::getAffineSymbolExpr(symCount++, value.getContext())}; return {mlir::getAffineSymbolExpr(symCount++, value.getContext())};
} }
return {}; return {};
} }
void fromCmpIOp(mlir::arith::CmpIOp cmpOp) { void fromCmpIOp(mlir::arith::CmpIOp cmpOp) {
auto lhsAffine = toAffineExpr(cmpOp.getLhs()); auto lhsAffine = toAffineExpr(cmpOp.getLhs());
▲ Show 20 Lines • Show All 153 Lines • ▼ Show 20 Lines if (auto shape = acoOp.getShape().getDefiningOp<ShapeOp>())
return populateIndexArgs(acoOp, shape, indexArgs, rewriter); return populateIndexArgs(acoOp, shape, indexArgs, rewriter);
if (auto shapeShift = acoOp.getShape().getDefiningOp<ShapeShiftOp>()) if (auto shapeShift = acoOp.getShape().getDefiningOp<ShapeShiftOp>())
return populateIndexArgs(acoOp, shapeShift, indexArgs, rewriter); return populateIndexArgs(acoOp, shapeShift, indexArgs, rewriter);
if (auto slice = acoOp.getShape().getDefiningOp<SliceOp>()) if (auto slice = acoOp.getShape().getDefiningOp<SliceOp>())
return populateIndexArgs(acoOp, slice, indexArgs, rewriter); return populateIndexArgs(acoOp, slice, indexArgs, rewriter);
} }
/// Returns affine.apply and fir.convert from array_coor and gendims /// Returns affine.apply and fir.convert from array_coor and gendims
static std::pair<mlir::AffineApplyOp, fir::ConvertOp> static std::pair<affine::AffineApplyOp, fir::ConvertOp>
createAffineOps(mlir::Value arrayRef, mlir::PatternRewriter &rewriter) { createAffineOps(mlir::Value arrayRef, mlir::PatternRewriter &rewriter) {
auto acoOp = arrayRef.getDefiningOp<ArrayCoorOp>(); auto acoOp = arrayRef.getDefiningOp<ArrayCoorOp>();
auto affineMap = auto affineMap =
createArrayIndexAffineMap(acoOp.getIndices().size(), acoOp.getContext()); createArrayIndexAffineMap(acoOp.getIndices().size(), acoOp.getContext());
SmallVector<mlir::Value> indexArgs; SmallVector<mlir::Value> indexArgs;
indexArgs.append(acoOp.getIndices().begin(), acoOp.getIndices().end()); indexArgs.append(acoOp.getIndices().begin(), acoOp.getIndices().end());
populateIndexArgs(acoOp, indexArgs, rewriter); populateIndexArgs(acoOp, indexArgs, rewriter);
auto affineApply = rewriter.create<mlir::AffineApplyOp>(acoOp.getLoc(), auto affineApply = rewriter.create<affine::AffineApplyOp>(
affineMap, indexArgs); acoOp.getLoc(), affineMap, indexArgs);
auto arrayElementType = coordinateArrayElement(acoOp); auto arrayElementType = coordinateArrayElement(acoOp);
auto newType = auto newType =
mlir::MemRefType::get({mlir::ShapedType::kDynamic}, arrayElementType); mlir::MemRefType::get({mlir::ShapedType::kDynamic}, arrayElementType);
auto arrayConvert = rewriter.create<fir::ConvertOp>(acoOp.getLoc(), newType, auto arrayConvert = rewriter.create<fir::ConvertOp>(acoOp.getLoc(), newType,
acoOp.getMemref()); acoOp.getMemref());
return std::make_pair(affineApply, arrayConvert); return std::make_pair(affineApply, arrayConvert);
} }
static void rewriteLoad(fir::LoadOp loadOp, mlir::PatternRewriter &rewriter) { static void rewriteLoad(fir::LoadOp loadOp, mlir::PatternRewriter &rewriter) {
rewriter.setInsertionPoint(loadOp); rewriter.setInsertionPoint(loadOp);
auto affineOps = createAffineOps(loadOp.getMemref(), rewriter); auto affineOps = createAffineOps(loadOp.getMemref(), rewriter);
rewriter.replaceOpWithNewOp<mlir::AffineLoadOp>( rewriter.replaceOpWithNewOp<affine::AffineLoadOp>(
loadOp, affineOps.second.getResult(), affineOps.first.getResult()); loadOp, affineOps.second.getResult(), affineOps.first.getResult());
} }
static void rewriteStore(fir::StoreOp storeOp, static void rewriteStore(fir::StoreOp storeOp,
mlir::PatternRewriter &rewriter) { mlir::PatternRewriter &rewriter) {
rewriter.setInsertionPoint(storeOp); rewriter.setInsertionPoint(storeOp);
auto affineOps = createAffineOps(storeOp.getMemref(), rewriter); auto affineOps = createAffineOps(storeOp.getMemref(), rewriter);
rewriter.replaceOpWithNewOp<mlir::AffineStoreOp>(storeOp, storeOp.getValue(), rewriter.replaceOpWithNewOp<affine::AffineStoreOp>(
affineOps.second.getResult(), storeOp, storeOp.getValue(), affineOps.second.getResult(),
affineOps.first.getResult()); affineOps.first.getResult());
} }
static void rewriteMemoryOps(Block *block, mlir::PatternRewriter &rewriter) { static void rewriteMemoryOps(Block *block, mlir::PatternRewriter &rewriter) {
for (auto &bodyOp : block->getOperations()) { for (auto &bodyOp : block->getOperations()) {
if (isa<fir::LoadOp>(bodyOp)) if (isa<fir::LoadOp>(bodyOp))
rewriteLoad(cast<fir::LoadOp>(bodyOp), rewriter); rewriteLoad(cast<fir::LoadOp>(bodyOp), rewriter);
if (isa<fir::StoreOp>(bodyOp)) if (isa<fir::StoreOp>(bodyOp))
rewriteStore(cast<fir::StoreOp>(bodyOp), rewriter); rewriteStore(cast<fir::StoreOp>(bodyOp), rewriter);
Show All 36 Lines matchAndRewrite(fir::DoLoopOp loop,
LLVM_DEBUG(llvm::dbgs() << "AffineLoopConversion: loop rewriten to:\n"; LLVM_DEBUG(llvm::dbgs() << "AffineLoopConversion: loop rewriten to:\n";
affineFor.dump();); affineFor.dump(););
rewriter.replaceOp(loop, affineFor.getOperation()->getResults()); rewriter.replaceOp(loop, affineFor.getOperation()->getResults());
return success(); return success();
} }
private: private:
std::pair<mlir::AffineForOp, mlir::Value> std::pair<affine::AffineForOp, mlir::Value>
createAffineFor(fir::DoLoopOp op, mlir::PatternRewriter &rewriter) const { createAffineFor(fir::DoLoopOp op, mlir::PatternRewriter &rewriter) const {
if (auto constantStep = constantIntegerLike(op.getStep())) if (auto constantStep = constantIntegerLike(op.getStep()))
if (*constantStep > 0) if (*constantStep > 0)
return positiveConstantStep(op, *constantStep, rewriter); return positiveConstantStep(op, *constantStep, rewriter);
return genericBounds(op, rewriter); return genericBounds(op, rewriter);
} }
// when step for the loop is positive compile time constant // when step for the loop is positive compile time constant
std::pair<mlir::AffineForOp, mlir::Value> std::pair<affine::AffineForOp, mlir::Value>
positiveConstantStep(fir::DoLoopOp op, int64_t step, positiveConstantStep(fir::DoLoopOp op, int64_t step,
mlir::PatternRewriter &rewriter) const { mlir::PatternRewriter &rewriter) const {
auto affineFor = rewriter.create<mlir::AffineForOp>( auto affineFor = rewriter.create<affine::AffineForOp>(
op.getLoc(), ValueRange(op.getLowerBound()), op.getLoc(), ValueRange(op.getLowerBound()),
mlir::AffineMap::get(0, 1, mlir::AffineMap::get(0, 1,
mlir::getAffineSymbolExpr(0, op.getContext())), mlir::getAffineSymbolExpr(0, op.getContext())),
ValueRange(op.getUpperBound()), ValueRange(op.getUpperBound()),
mlir::AffineMap::get(0, 1, mlir::AffineMap::get(0, 1,
1 + mlir::getAffineSymbolExpr(0, op.getContext())), 1 + mlir::getAffineSymbolExpr(0, op.getContext())),
step); step);
return std::make_pair(affineFor, affineFor.getInductionVar()); return std::make_pair(affineFor, affineFor.getInductionVar());
} }
std::pair<mlir::AffineForOp, mlir::Value> std::pair<affine::AffineForOp, mlir::Value>
genericBounds(fir::DoLoopOp op, mlir::PatternRewriter &rewriter) const { genericBounds(fir::DoLoopOp op, mlir::PatternRewriter &rewriter) const {
auto lowerBound = mlir::getAffineSymbolExpr(0, op.getContext()); auto lowerBound = mlir::getAffineSymbolExpr(0, op.getContext());
auto upperBound = mlir::getAffineSymbolExpr(1, op.getContext()); auto upperBound = mlir::getAffineSymbolExpr(1, op.getContext());
auto step = mlir::getAffineSymbolExpr(2, op.getContext()); auto step = mlir::getAffineSymbolExpr(2, op.getContext());
mlir::AffineMap upperBoundMap = mlir::AffineMap::get( mlir::AffineMap upperBoundMap = mlir::AffineMap::get(
0, 3, (upperBound - lowerBound + step).floorDiv(step)); 0, 3, (upperBound - lowerBound + step).floorDiv(step));
auto genericUpperBound = rewriter.create<mlir::AffineApplyOp>( auto genericUpperBound = rewriter.create<affine::AffineApplyOp>(
op.getLoc(), upperBoundMap, op.getLoc(), upperBoundMap,
ValueRange({op.getLowerBound(), op.getUpperBound(), op.getStep()})); ValueRange({op.getLowerBound(), op.getUpperBound(), op.getStep()}));
auto actualIndexMap = mlir::AffineMap::get( auto actualIndexMap = mlir::AffineMap::get(
1, 2, 1, 2,
(lowerBound + mlir::getAffineDimExpr(0, op.getContext())) * (lowerBound + mlir::getAffineDimExpr(0, op.getContext())) *
mlir::getAffineSymbolExpr(1, op.getContext())); mlir::getAffineSymbolExpr(1, op.getContext()));
auto affineFor = rewriter.create<mlir::AffineForOp>( auto affineFor = rewriter.create<affine::AffineForOp>(
op.getLoc(), ValueRange(), op.getLoc(), ValueRange(),
AffineMap::getConstantMap(0, op.getContext()), AffineMap::getConstantMap(0, op.getContext()),
genericUpperBound.getResult(), genericUpperBound.getResult(),
mlir::AffineMap::get(0, 1, mlir::AffineMap::get(0, 1,
1 + mlir::getAffineSymbolExpr(0, op.getContext())), 1 + mlir::getAffineSymbolExpr(0, op.getContext())),
1); 1);
rewriter.setInsertionPointToStart(affineFor.getBody()); rewriter.setInsertionPointToStart(affineFor.getBody());
auto actualIndex = rewriter.create<mlir::AffineApplyOp>( auto actualIndex = rewriter.create<affine::AffineApplyOp>(
op.getLoc(), actualIndexMap, op.getLoc(), actualIndexMap,
ValueRange( ValueRange(
{affineFor.getInductionVar(), op.getLowerBound(), op.getStep()})); {affineFor.getInductionVar(), op.getLowerBound(), op.getStep()}));
return std::make_pair(affineFor, actualIndex.getResult()); return std::make_pair(affineFor, actualIndex.getResult());
} }
AffineFunctionAnalysis &functionAnalysis; AffineFunctionAnalysis &functionAnalysis;
}; };
Show All 12 Lines matchAndRewrite(fir::IfOp op,
auto &ifOps = op.getThenRegion().front().getOperations(); auto &ifOps = op.getThenRegion().front().getOperations();
auto affineCondition = AffineIfCondition(op.getCondition()); auto affineCondition = AffineIfCondition(op.getCondition());
if (!affineCondition.hasIntegerSet()) { if (!affineCondition.hasIntegerSet()) {
LLVM_DEBUG( LLVM_DEBUG(
llvm::dbgs() llvm::dbgs()
<< "AffineIfConversion: couldn't calculate affine condition\n";); << "AffineIfConversion: couldn't calculate affine condition\n";);
return failure(); return failure();
} }
auto affineIf = rewriter.create<mlir::AffineIfOp>( auto affineIf = rewriter.create<affine::AffineIfOp>(
op.getLoc(), affineCondition.getIntegerSet(), op.getLoc(), affineCondition.getIntegerSet(),
affineCondition.getAffineArgs(), !op.getElseRegion().empty()); affineCondition.getAffineArgs(), !op.getElseRegion().empty());
rewriter.startRootUpdate(affineIf); rewriter.startRootUpdate(affineIf);
affineIf.getThenBlock()->getOperations().splice( affineIf.getThenBlock()->getOperations().splice(
std::prev(affineIf.getThenBlock()->end()), ifOps, ifOps.begin(), std::prev(affineIf.getThenBlock()->end()), ifOps, ifOps.begin(),
std::prev(ifOps.end())); std::prev(ifOps.end()));
if (!op.getElseRegion().empty()) { if (!op.getElseRegion().empty()) {
auto &otherOps = op.getElseRegion().front().getOperations(); auto &otherOps = op.getElseRegion().front().getOperations();
Show All 21 Lines void runOnOperation() override {
auto *context = &getContext(); auto *context = &getContext();
auto function = getOperation(); auto function = getOperation();
markAllAnalysesPreserved(); markAllAnalysesPreserved();
auto functionAnalysis = AffineFunctionAnalysis(function); auto functionAnalysis = AffineFunctionAnalysis(function);
mlir::RewritePatternSet patterns(context); mlir::RewritePatternSet patterns(context);
patterns.insert<AffineIfConversion>(context, functionAnalysis); patterns.insert<AffineIfConversion>(context, functionAnalysis);
patterns.insert<AffineLoopConversion>(context, functionAnalysis); patterns.insert<AffineLoopConversion>(context, functionAnalysis);
mlir::ConversionTarget target = *context; mlir::ConversionTarget target = *context;
target.addLegalDialect<mlir::AffineDialect, FIROpsDialect, target.addLegalDialect<mlir::affine::AffineDialect, FIROpsDialect,
mlir::scf::SCFDialect, mlir::arith::ArithDialect, mlir::scf::SCFDialect, mlir::arith::ArithDialect,
mlir::func::FuncDialect>(); mlir::func::FuncDialect>();
target.addDynamicallyLegalOp<IfOp>([&functionAnalysis](fir::IfOp op) { target.addDynamicallyLegalOp<IfOp>([&functionAnalysis](fir::IfOp op) {
return !(functionAnalysis.getChildIfAnalysis(op).canPromoteToAffine()); return !(functionAnalysis.getChildIfAnalysis(op).canPromoteToAffine());
}); });
target.addDynamicallyLegalOp<DoLoopOp>([&functionAnalysis]( target.addDynamicallyLegalOp<DoLoopOp>([&functionAnalysis](
fir::DoLoopOp op) { fir::DoLoopOp op) {
return !(functionAnalysis.getChildLoopAnalysis(op).canPromoteToAffine()); return !(functionAnalysis.getChildLoopAnalysis(op).canPromoteToAffine());
Show All 20 Lines

flang/lib/Optimizer/Transforms/CharacterConversion.cpp

Show First 20 Lines • Show All 104 Lines • ▼ Show 20 Linespublic:
void runOnOperation() override { void runOnOperation() override {
CharacterConversionOptions clOpts{useRuntimeCalls.getValue()}; CharacterConversionOptions clOpts{useRuntimeCalls.getValue()};
if (clOpts.runtimeName.empty()) { if (clOpts.runtimeName.empty()) {
auto *context = &getContext(); auto *context = &getContext();
auto *func = getOperation(); auto *func = getOperation();
mlir::RewritePatternSet patterns(context); mlir::RewritePatternSet patterns(context);
patterns.insert<CharacterConvertConversion>(context); patterns.insert<CharacterConvertConversion>(context);
mlir::ConversionTarget target(*context); mlir::ConversionTarget target(*context);
target.addLegalDialect<mlir::AffineDialect, fir::FIROpsDialect, target.addLegalDialect<mlir::affine::AffineDialect, fir::FIROpsDialect,
mlir::arith::ArithDialect, mlir::arith::ArithDialect,
mlir::func::FuncDialect>(); mlir::func::FuncDialect>();
// apply the patterns // apply the patterns
target.addIllegalOp<fir::CharConvertOp>(); target.addIllegalOp<fir::CharConvertOp>();
if (mlir::failed(mlir::applyPartialConversion(func, target, if (mlir::failed(mlir::applyPartialConversion(func, target,
std::move(patterns)))) { std::move(patterns)))) {
mlir::emitError(mlir::UnknownLoc::get(context), mlir::emitError(mlir::UnknownLoc::get(context),
Show All 15 Lines

flang/lib/Optimizer/Transforms/ControlFlowConverter.cpp

Show First 20 Lines • Show All 310 Lines • ▼ Show 20 Lines
class CfgConversion : public fir::impl::CFGConversionBase<CfgConversion> { class CfgConversion : public fir::impl::CFGConversionBase<CfgConversion> {
public: public:
void runOnOperation() override { void runOnOperation() override {
auto *context = &getContext(); auto *context = &getContext();
mlir::RewritePatternSet patterns(context); mlir::RewritePatternSet patterns(context);
patterns.insert<CfgLoopConv, CfgIfConv, CfgIterWhileConv>( patterns.insert<CfgLoopConv, CfgIfConv, CfgIterWhileConv>(
context, forceLoopToExecuteOnce); context, forceLoopToExecuteOnce);
mlir::ConversionTarget target(*context); mlir::ConversionTarget target(*context);
target.addLegalDialect<mlir::AffineDialect, mlir::cf::ControlFlowDialect, target.addLegalDialect<mlir::affine::AffineDialect,
FIROpsDialect, mlir::func::FuncDialect>(); mlir::cf::ControlFlowDialect, FIROpsDialect,
mlir::func::FuncDialect>();
// apply the patterns // apply the patterns
target.addIllegalOp<ResultOp, DoLoopOp, IfOp, IterWhileOp>(); target.addIllegalOp<ResultOp, DoLoopOp, IfOp, IterWhileOp>();
target.markUnknownOpDynamicallyLegal([](Operation *) { return true; }); target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
if (mlir::failed(mlir::applyPartialConversion(getOperation(), target, if (mlir::failed(mlir::applyPartialConversion(getOperation(), target,
std::move(patterns)))) { std::move(patterns)))) {
mlir::emitError(mlir::UnknownLoc::get(context), mlir::emitError(mlir::UnknownLoc::get(context),
"error in converting to CFG\n"); "error in converting to CFG\n");
Show All 12 Lines

flang/lib/Optimizer/Transforms/PolymorphicOpConversion.cpp

Show First 20 Lines • Show All 232 Lines • ▼ Show 20 Lines void runOnOperation() override {
mlir::RewritePatternSet patterns(context); mlir::RewritePatternSet patterns(context);
BindingTables bindingTables; BindingTables bindingTables;
buildBindingTables(bindingTables, mod); buildBindingTables(bindingTables, mod);
patterns.insert<SelectTypeConv>(context, moduleMutex); patterns.insert<SelectTypeConv>(context, moduleMutex);
patterns.insert<DispatchOpConv>(context, bindingTables); patterns.insert<DispatchOpConv>(context, bindingTables);
mlir::ConversionTarget target(*context); mlir::ConversionTarget target(*context);
target.addLegalDialect<mlir::AffineDialect, mlir::cf::ControlFlowDialect, target.addLegalDialect<mlir::affine::AffineDialect,
FIROpsDialect, mlir::func::FuncDialect>(); mlir::cf::ControlFlowDialect, FIROpsDialect,
mlir::func::FuncDialect>();
// apply the patterns // apply the patterns
target.addIllegalOp<SelectTypeOp>(); target.addIllegalOp<SelectTypeOp>();
target.addIllegalOp<DispatchOp>(); target.addIllegalOp<DispatchOp>();
target.markUnknownOpDynamicallyLegal([](Operation *) { return true; }); target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
if (mlir::failed(mlir::applyPartialConversion(getOperation(), target, if (mlir::failed(mlir::applyPartialConversion(getOperation(), target,
std::move(patterns)))) { std::move(patterns)))) {
mlir::emitError(mlir::UnknownLoc::get(context), mlir::emitError(mlir::UnknownLoc::get(context),
▲ Show 20 Lines • Show All 227 Lines • Show Last 20 Lines

mlir/docs/DialectConversion.md

Show First 20 Lines • Show All 92 Lines • ▼ Show 20 Lines MyTarget(MLIRContext &ctx) : ConversionTarget(ctx) {
/// Mark `arith.constant` op is always legal on this target. /// Mark `arith.constant` op is always legal on this target.
addLegalOp<arith::ConstantOp>(); addLegalOp<arith::ConstantOp>();
//-------------------------------------------------------------------------- //--------------------------------------------------------------------------
// Marking an operation as dynamically legal. // Marking an operation as dynamically legal.
/// Mark all operations within Affine dialect have dynamic legality /// Mark all operations within Affine dialect have dynamic legality
/// constraints. /// constraints.
addDynamicallyLegalDialect<AffineDialect>([](Operation *op) { ... }); addDynamicallyLegalDialect<affine::AffineDialect>(
[](Operation *op) { ... });
/// Mark `func.return` as dynamically legal, but provide a specific legality /// Mark `func.return` as dynamically legal, but provide a specific legality
/// callback. /// callback.
addDynamicallyLegalOp<func::ReturnOp>([](func::ReturnOp op) { ... }); addDynamicallyLegalOp<func::ReturnOp>([](func::ReturnOp op) { ... });
/// Treat unknown operations, i.e. those without a legalization action /// Treat unknown operations, i.e. those without a legalization action
/// directly set, as dynamically legal. /// directly set, as dynamically legal.
markUnknownOpDynamicallyLegal([](Operation *op) { ... }); markUnknownOpDynamicallyLegal([](Operation *op) { ... });
▲ Show 20 Lines • Show All 343 Lines • Show Last 20 Lines

mlir/docs/Tutorials/Toy/Ch-5.md

Show First 20 Lines • Show All 58 Lines • ▼ Show 20 Lines
void ToyToAffineLoweringPass::runOnOperation() { void ToyToAffineLoweringPass::runOnOperation() {
// The first thing to define is the conversion target. This will define the // The first thing to define is the conversion target. This will define the
// final target for this lowering. // final target for this lowering.
mlir::ConversionTarget target(getContext()); mlir::ConversionTarget target(getContext());
// We define the specific operations, or dialects, that are legal targets for // We define the specific operations, or dialects, that are legal targets for
// this lowering. In our case, we are lowering to a combination of the // this lowering. In our case, we are lowering to a combination of the
// `Affine`, `Arith`, `Func`, and `MemRef` dialects. // `Affine`, `Arith`, `Func`, and `MemRef` dialects.
target.addLegalDialect<AffineDialect, arith::ArithDialect, target.addLegalDialect<affine::AffineDialect, arith::ArithDialect,
func::FuncDialect, memref::MemRefDialect>(); func::FuncDialect, memref::MemRefDialect>();
// We also define the Toy dialect as Illegal so that the conversion will fail // We also define the Toy dialect as Illegal so that the conversion will fail
// if any of these operations are *not* converted. Given that we actually want // if any of these operations are *not* converted. Given that we actually want
// a partial lowering, we explicitly mark the Toy operations that don't want // a partial lowering, we explicitly mark the Toy operations that don't want
// to lower, `toy.print`, as *legal*. `toy.print` will still need its operands // to lower, `toy.print`, as *legal*. `toy.print` will still need its operands
// to be updated though (as we convert from TensorType to MemRefType), so we // to be updated though (as we convert from TensorType to MemRefType), so we
// only treat it as `legal` if its operands are legal. // only treat it as `legal` if its operands are legal.
▲ Show 20 Lines • Show All 275 Lines • Show Last 20 Lines

mlir/examples/toy/Ch5/mlir/LowerToAffineLoops.cpp

Show First 20 Lines • Show All 69 Lines • ▼ Show 20 Linesstatic void lowerOpToLoops(Operation *op, ValueRange operands,
auto alloc = insertAllocAndDealloc(memRefType, loc, rewriter); auto alloc = insertAllocAndDealloc(memRefType, loc, rewriter);
// Create a nest of affine loops, with one loop per dimension of the shape. // Create a nest of affine loops, with one loop per dimension of the shape.
// The buildAffineLoopNest function takes a callback that is used to construct // The buildAffineLoopNest function takes a callback that is used to construct
// the body of the innermost loop given a builder, a location and a range of // the body of the innermost loop given a builder, a location and a range of
// loop induction variables. // loop induction variables.
SmallVector<int64_t, 4> lowerBounds(tensorType.getRank(), /*Value=*/0); SmallVector<int64_t, 4> lowerBounds(tensorType.getRank(), /*Value=*/0);
SmallVector<int64_t, 4> steps(tensorType.getRank(), /*Value=*/1); SmallVector<int64_t, 4> steps(tensorType.getRank(), /*Value=*/1);
buildAffineLoopNest( affine::buildAffineLoopNest(
rewriter, loc, lowerBounds, tensorType.getShape(), steps, rewriter, loc, lowerBounds, tensorType.getShape(), steps,
[&](OpBuilder &nestedBuilder, Location loc, ValueRange ivs) { [&](OpBuilder &nestedBuilder, Location loc, ValueRange ivs) {
// Call the processing function with the rewriter, the memref operands, // Call the processing function with the rewriter, the memref operands,
// and the loop induction variables. This function will return the value // and the loop induction variables. This function will return the value
// to store at the current index. // to store at the current index.
Value valueToStore = processIteration(nestedBuilder, operands, ivs); Value valueToStore = processIteration(nestedBuilder, operands, ivs);
nestedBuilder.create<AffineStoreOp>(loc, valueToStore, alloc, ivs); nestedBuilder.create<affine::AffineStoreOp>(loc, valueToStore, alloc,
ivs);
}); });
// Replace this operation with the generated alloc. // Replace this operation with the generated alloc.
rewriter.replaceOp(op, alloc); rewriter.replaceOp(op, alloc);
} }
namespace { namespace {
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
Show All 14 Lines lowerOpToLoops(op, operands, rewriter,
ValueRange loopIvs) { ValueRange loopIvs) {
// Generate an adaptor for the remapped operands of the // Generate an adaptor for the remapped operands of the
// BinaryOp. This allows for using the nice named accessors // BinaryOp. This allows for using the nice named accessors
// that are generated by the ODS. // that are generated by the ODS.
typename BinaryOp::Adaptor binaryAdaptor(memRefOperands); typename BinaryOp::Adaptor binaryAdaptor(memRefOperands);
// Generate loads for the element of 'lhs' and 'rhs' at the // Generate loads for the element of 'lhs' and 'rhs' at the
// inner loop. // inner loop.
auto loadedLhs = builder.create<AffineLoadOp>( auto loadedLhs = builder.create<affine::AffineLoadOp>(
loc, binaryAdaptor.getLhs(), loopIvs); loc, binaryAdaptor.getLhs(), loopIvs);
auto loadedRhs = builder.create<AffineLoadOp>( auto loadedRhs = builder.create<affine::AffineLoadOp>(
loc, binaryAdaptor.getRhs(), loopIvs); loc, binaryAdaptor.getRhs(), loopIvs);
// Create the binary operation performed on the loaded // Create the binary operation performed on the loaded
// values. // values.
return builder.create<LoweredBinaryOp>(loc, loadedLhs, return builder.create<LoweredBinaryOp>(loc, loadedLhs,
loadedRhs); loadedRhs);
}); });
return success(); return success();
▲ Show 20 Lines • Show All 42 Lines • ▼ Show 20 Lines LogicalResult matchAndRewrite(toy::ConstantOp op,
// functor recursively walks the dimensions of the constant shape, // functor recursively walks the dimensions of the constant shape,
// generating a store when the recursion hits the base case. // generating a store when the recursion hits the base case.
SmallVector<Value, 2> indices; SmallVector<Value, 2> indices;
auto valueIt = constantValue.value_begin<FloatAttr>(); auto valueIt = constantValue.value_begin<FloatAttr>();
std::function<void(uint64_t)> storeElements = [&](uint64_t dimension) { std::function<void(uint64_t)> storeElements = [&](uint64_t dimension) {
// The last dimension is the base case of the recursion, at this point // The last dimension is the base case of the recursion, at this point
// we store the element at the given index. // we store the element at the given index.
if (dimension == valueShape.size()) { if (dimension == valueShape.size()) {
rewriter.create<AffineStoreOp>( rewriter.create<affine::AffineStoreOp>(
loc, rewriter.create<arith::ConstantOp>(loc, *valueIt++), alloc, loc, rewriter.create<arith::ConstantOp>(loc, *valueIt++), alloc,
llvm::ArrayRef(indices)); llvm::ArrayRef(indices));
return; return;
} }
// Otherwise, iterate over the current dimension and add the indices to // Otherwise, iterate over the current dimension and add the indices to
// the list. // the list.
for (uint64_t i = 0, e = valueShape[dimension]; i != e; ++i) { for (uint64_t i = 0, e = valueShape[dimension]; i != e; ++i) {
▲ Show 20 Lines • Show All 100 Lines • ▼ Show 20 Lines lowerOpToLoops(op, operands, rewriter,
// TransposeOp. This allows for using the nice named // TransposeOp. This allows for using the nice named
// accessors that are generated by the ODS. // accessors that are generated by the ODS.
toy::TransposeOpAdaptor transposeAdaptor(memRefOperands); toy::TransposeOpAdaptor transposeAdaptor(memRefOperands);
Value input = transposeAdaptor.getInput(); Value input = transposeAdaptor.getInput();
// Transpose the elements by generating a load from the // Transpose the elements by generating a load from the
// reverse indices. // reverse indices.
SmallVector<Value, 2> reverseIvs(llvm::reverse(loopIvs)); SmallVector<Value, 2> reverseIvs(llvm::reverse(loopIvs));
return builder.create<AffineLoadOp>(loc, input, return builder.create<affine::AffineLoadOp>(loc, input,
reverseIvs); reverseIvs);
}); });
return success(); return success();
} }
}; };
} // namespace } // namespace
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// ToyToAffineLoweringPass // ToyToAffineLoweringPass
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
/// This is a partial lowering to affine loops of the toy operations that are /// This is a partial lowering to affine loops of the toy operations that are
/// computationally intensive (like matmul for example...) while keeping the /// computationally intensive (like matmul for example...) while keeping the
/// rest of the code in the Toy dialect. /// rest of the code in the Toy dialect.
namespace { namespace {
struct ToyToAffineLoweringPass struct ToyToAffineLoweringPass
: public PassWrapper<ToyToAffineLoweringPass, OperationPass<ModuleOp>> { : public PassWrapper<ToyToAffineLoweringPass, OperationPass<ModuleOp>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ToyToAffineLoweringPass) MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ToyToAffineLoweringPass)
void getDependentDialects(DialectRegistry &registry) const override { void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<AffineDialect, func::FuncDialect, memref::MemRefDialect>(); registry.insert<affine::AffineDialect, func::FuncDialect,
memref::MemRefDialect>();
} }
void runOnOperation() final; void runOnOperation() final;
}; };
} // namespace } // namespace
void ToyToAffineLoweringPass::runOnOperation() { void ToyToAffineLoweringPass::runOnOperation() {
// The first thing to define is the conversion target. This will define the // The first thing to define is the conversion target. This will define the
// final target for this lowering. // final target for this lowering.
ConversionTarget target(getContext()); ConversionTarget target(getContext());
// We define the specific operations, or dialects, that are legal targets for // We define the specific operations, or dialects, that are legal targets for
// this lowering. In our case, we are lowering to a combination of the // this lowering. In our case, we are lowering to a combination of the
// `Affine`, `Arith`, `Func`, and `MemRef` dialects. // `Affine`, `Arith`, `Func`, and `MemRef` dialects.
target.addLegalDialect<AffineDialect, BuiltinDialect, arith::ArithDialect, target.addLegalDialect<affine::AffineDialect, BuiltinDialect,
func::FuncDialect, memref::MemRefDialect>(); arith::ArithDialect, func::FuncDialect,
memref::MemRefDialect>();
// We also define the Toy dialect as Illegal so that the conversion will fail // We also define the Toy dialect as Illegal so that the conversion will fail
// if any of these operations are *not* converted. Given that we actually want // if any of these operations are *not* converted. Given that we actually want
// a partial lowering, we explicitly mark the Toy operations that don't want // a partial lowering, we explicitly mark the Toy operations that don't want
// to lower, `toy.print`, as `legal`. `toy.print` will still need its operands // to lower, `toy.print`, as `legal`. `toy.print` will still need its operands
// to be updated though (as we convert from TensorType to MemRefType), so we // to be updated though (as we convert from TensorType to MemRefType), so we
// only treat it as `legal` if its operands are legal. // only treat it as `legal` if its operands are legal.
target.addIllegalDialect<toy::ToyDialect>(); target.addIllegalDialect<toy::ToyDialect>();
Show All 25 Lines

mlir/examples/toy/Ch5/toyc.cpp

Show First 20 Lines • Show All 142 Lines • ▼ Show 20 Lines if (isLoweringToAffine) {
// Add a few cleanups post lowering. // Add a few cleanups post lowering.
mlir::OpPassManager &optPM = pm.nest<mlir::func::FuncOp>(); mlir::OpPassManager &optPM = pm.nest<mlir::func::FuncOp>();
optPM.addPass(mlir::createCanonicalizerPass()); optPM.addPass(mlir::createCanonicalizerPass());
optPM.addPass(mlir::createCSEPass()); optPM.addPass(mlir::createCSEPass());
// Add optimizations if enabled. // Add optimizations if enabled.
if (enableOpt) { if (enableOpt) {
optPM.addPass(mlir::createLoopFusionPass()); optPM.addPass(mlir::affine::createLoopFusionPass());
optPM.addPass(mlir::createAffineScalarReplacementPass()); optPM.addPass(mlir::affine::createAffineScalarReplacementPass());
} }
} }
if (mlir::failed(pm.run(*module))) if (mlir::failed(pm.run(*module)))
return 4; return 4;
module->dump(); module->dump();
return 0; return 0;
Show All 36 Lines

mlir/examples/toy/Ch6/mlir/LowerToAffineLoops.cpp

Show First 20 Lines • Show All 69 Lines • ▼ Show 20 Linesstatic void lowerOpToLoops(Operation *op, ValueRange operands,
auto alloc = insertAllocAndDealloc(memRefType, loc, rewriter); auto alloc = insertAllocAndDealloc(memRefType, loc, rewriter);
// Create a nest of affine loops, with one loop per dimension of the shape. // Create a nest of affine loops, with one loop per dimension of the shape.
// The buildAffineLoopNest function takes a callback that is used to construct // The buildAffineLoopNest function takes a callback that is used to construct
// the body of the innermost loop given a builder, a location and a range of // the body of the innermost loop given a builder, a location and a range of
// loop induction variables. // loop induction variables.
SmallVector<int64_t, 4> lowerBounds(tensorType.getRank(), /*Value=*/0); SmallVector<int64_t, 4> lowerBounds(tensorType.getRank(), /*Value=*/0);
SmallVector<int64_t, 4> steps(tensorType.getRank(), /*Value=*/1); SmallVector<int64_t, 4> steps(tensorType.getRank(), /*Value=*/1);
buildAffineLoopNest( affine::buildAffineLoopNest(
rewriter, loc, lowerBounds, tensorType.getShape(), steps, rewriter, loc, lowerBounds, tensorType.getShape(), steps,
[&](OpBuilder &nestedBuilder, Location loc, ValueRange ivs) { [&](OpBuilder &nestedBuilder, Location loc, ValueRange ivs) {
// Call the processing function with the rewriter, the memref operands, // Call the processing function with the rewriter, the memref operands,
// and the loop induction variables. This function will return the value // and the loop induction variables. This function will return the value
// to store at the current index. // to store at the current index.
Value valueToStore = processIteration(nestedBuilder, operands, ivs); Value valueToStore = processIteration(nestedBuilder, operands, ivs);
nestedBuilder.create<AffineStoreOp>(loc, valueToStore, alloc, ivs); nestedBuilder.create<affine::AffineStoreOp>(loc, valueToStore, alloc,
ivs);
}); });
// Replace this operation with the generated alloc. // Replace this operation with the generated alloc.
rewriter.replaceOp(op, alloc); rewriter.replaceOp(op, alloc);
} }
namespace { namespace {
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
Show All 14 Lines lowerOpToLoops(op, operands, rewriter,
ValueRange loopIvs) { ValueRange loopIvs) {
// Generate an adaptor for the remapped operands of the // Generate an adaptor for the remapped operands of the
// BinaryOp. This allows for using the nice named accessors // BinaryOp. This allows for using the nice named accessors
// that are generated by the ODS. // that are generated by the ODS.
typename BinaryOp::Adaptor binaryAdaptor(memRefOperands); typename BinaryOp::Adaptor binaryAdaptor(memRefOperands);
// Generate loads for the element of 'lhs' and 'rhs' at the // Generate loads for the element of 'lhs' and 'rhs' at the
// inner loop. // inner loop.
auto loadedLhs = builder.create<AffineLoadOp>( auto loadedLhs = builder.create<affine::AffineLoadOp>(
loc, binaryAdaptor.getLhs(), loopIvs); loc, binaryAdaptor.getLhs(), loopIvs);
auto loadedRhs = builder.create<AffineLoadOp>( auto loadedRhs = builder.create<affine::AffineLoadOp>(
loc, binaryAdaptor.getRhs(), loopIvs); loc, binaryAdaptor.getRhs(), loopIvs);
// Create the binary operation performed on the loaded // Create the binary operation performed on the loaded
// values. // values.
return builder.create<LoweredBinaryOp>(loc, loadedLhs, return builder.create<LoweredBinaryOp>(loc, loadedLhs,
loadedRhs); loadedRhs);
}); });
return success(); return success();
▲ Show 20 Lines • Show All 42 Lines • ▼ Show 20 Lines LogicalResult matchAndRewrite(toy::ConstantOp op,
// functor recursively walks the dimensions of the constant shape, // functor recursively walks the dimensions of the constant shape,
// generating a store when the recursion hits the base case. // generating a store when the recursion hits the base case.
SmallVector<Value, 2> indices; SmallVector<Value, 2> indices;
auto valueIt = constantValue.value_begin<FloatAttr>(); auto valueIt = constantValue.value_begin<FloatAttr>();
std::function<void(uint64_t)> storeElements = [&](uint64_t dimension) { std::function<void(uint64_t)> storeElements = [&](uint64_t dimension) {
// The last dimension is the base case of the recursion, at this point // The last dimension is the base case of the recursion, at this point
// we store the element at the given index. // we store the element at the given index.
if (dimension == valueShape.size()) { if (dimension == valueShape.size()) {
rewriter.create<AffineStoreOp>( rewriter.create<affine::AffineStoreOp>(
loc, rewriter.create<arith::ConstantOp>(loc, *valueIt++), alloc, loc, rewriter.create<arith::ConstantOp>(loc, *valueIt++), alloc,
llvm::ArrayRef(indices)); llvm::ArrayRef(indices));
return; return;
} }
// Otherwise, iterate over the current dimension and add the indices to // Otherwise, iterate over the current dimension and add the indices to
// the list. // the list.
for (uint64_t i = 0, e = valueShape[dimension]; i != e; ++i) { for (uint64_t i = 0, e = valueShape[dimension]; i != e; ++i) {
▲ Show 20 Lines • Show All 100 Lines • ▼ Show 20 Lines lowerOpToLoops(op, operands, rewriter,
// TransposeOp. This allows for using the nice named // TransposeOp. This allows for using the nice named
// accessors that are generated by the ODS. // accessors that are generated by the ODS.
toy::TransposeOpAdaptor transposeAdaptor(memRefOperands); toy::TransposeOpAdaptor transposeAdaptor(memRefOperands);
Value input = transposeAdaptor.getInput(); Value input = transposeAdaptor.getInput();
// Transpose the elements by generating a load from the // Transpose the elements by generating a load from the
// reverse indices. // reverse indices.
SmallVector<Value, 2> reverseIvs(llvm::reverse(loopIvs)); SmallVector<Value, 2> reverseIvs(llvm::reverse(loopIvs));
return builder.create<AffineLoadOp>(loc, input, return builder.create<affine::AffineLoadOp>(loc, input,
reverseIvs); reverseIvs);
}); });
return success(); return success();
} }
}; };
} // namespace } // namespace
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// ToyToAffineLoweringPass // ToyToAffineLoweringPass
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
/// This is a partial lowering to affine loops of the toy operations that are /// This is a partial lowering to affine loops of the toy operations that are
/// computationally intensive (like matmul for example...) while keeping the /// computationally intensive (like matmul for example...) while keeping the
/// rest of the code in the Toy dialect. /// rest of the code in the Toy dialect.
namespace { namespace {
struct ToyToAffineLoweringPass struct ToyToAffineLoweringPass
: public PassWrapper<ToyToAffineLoweringPass, OperationPass<ModuleOp>> { : public PassWrapper<ToyToAffineLoweringPass, OperationPass<ModuleOp>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ToyToAffineLoweringPass) MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ToyToAffineLoweringPass)
void getDependentDialects(DialectRegistry &registry) const override { void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<AffineDialect, func::FuncDialect, memref::MemRefDialect>(); registry.insert<affine::AffineDialect, func::FuncDialect,
memref::MemRefDialect>();
} }
void runOnOperation() final; void runOnOperation() final;
}; };
} // namespace } // namespace
void ToyToAffineLoweringPass::runOnOperation() { void ToyToAffineLoweringPass::runOnOperation() {
// The first thing to define is the conversion target. This will define the // The first thing to define is the conversion target. This will define the
// final target for this lowering. // final target for this lowering.
ConversionTarget target(getContext()); ConversionTarget target(getContext());
// We define the specific operations, or dialects, that are legal targets for // We define the specific operations, or dialects, that are legal targets for
// this lowering. In our case, we are lowering to a combination of the // this lowering. In our case, we are lowering to a combination of the
// `Affine`, `Arith`, `Func`, and `MemRef` dialects. // `Affine`, `Arith`, `Func`, and `MemRef` dialects.
target.addLegalDialect<AffineDialect, BuiltinDialect, arith::ArithDialect, target.addLegalDialect<affine::AffineDialect, BuiltinDialect,
func::FuncDialect, memref::MemRefDialect>(); arith::ArithDialect, func::FuncDialect,
memref::MemRefDialect>();
// We also define the Toy dialect as Illegal so that the conversion will fail // We also define the Toy dialect as Illegal so that the conversion will fail
// if any of these operations are *not* converted. Given that we actually want // if any of these operations are *not* converted. Given that we actually want
// a partial lowering, we explicitly mark the Toy operations that don't want // a partial lowering, we explicitly mark the Toy operations that don't want
// to lower, `toy.print`, as `legal`. `toy.print` will still need its operands // to lower, `toy.print`, as `legal`. `toy.print` will still need its operands
// to be updated though (as we convert from TensorType to MemRefType), so we // to be updated though (as we convert from TensorType to MemRefType), so we
// only treat it as `legal` if its operands are legal. // only treat it as `legal` if its operands are legal.
target.addIllegalDialect<toy::ToyDialect>(); target.addIllegalDialect<toy::ToyDialect>();
Show All 25 Lines

mlir/examples/toy/Ch6/toyc.cpp

Show First 20 Lines • Show All 160 Lines • ▼ Show 20 Lines if (isLoweringToAffine) {
// Add a few cleanups post lowering. // Add a few cleanups post lowering.
mlir::OpPassManager &optPM = pm.nest<mlir::func::FuncOp>(); mlir::OpPassManager &optPM = pm.nest<mlir::func::FuncOp>();
optPM.addPass(mlir::createCanonicalizerPass()); optPM.addPass(mlir::createCanonicalizerPass());
optPM.addPass(mlir::createCSEPass()); optPM.addPass(mlir::createCSEPass());
// Add optimizations if enabled. // Add optimizations if enabled.
if (enableOpt) { if (enableOpt) {
optPM.addPass(mlir::createLoopFusionPass()); optPM.addPass(mlir::affine::createLoopFusionPass());
optPM.addPass(mlir::createAffineScalarReplacementPass()); optPM.addPass(mlir::affine::createAffineScalarReplacementPass());
} }
} }
if (isLoweringToLLVM) { if (isLoweringToLLVM) {
// Finish lowering the toy IR to the LLVM dialect. // Finish lowering the toy IR to the LLVM dialect.
pm.addPass(mlir::toy::createLowerToLLVMPass()); pm.addPass(mlir::toy::createLowerToLLVMPass());
// This is necessary to have line tables emitted and basic // This is necessary to have line tables emitted and basic
// debugger working. In the future we will add proper debug information // debugger working. In the future we will add proper debug information
▲ Show 20 Lines • Show All 140 Lines • Show Last 20 Lines

mlir/examples/toy/Ch7/mlir/LowerToAffineLoops.cpp

Show First 20 Lines • Show All 69 Lines • ▼ Show 20 Linesstatic void lowerOpToLoops(Operation *op, ValueRange operands,
auto alloc = insertAllocAndDealloc(memRefType, loc, rewriter); auto alloc = insertAllocAndDealloc(memRefType, loc, rewriter);
// Create a nest of affine loops, with one loop per dimension of the shape. // Create a nest of affine loops, with one loop per dimension of the shape.
// The buildAffineLoopNest function takes a callback that is used to construct // The buildAffineLoopNest function takes a callback that is used to construct
// the body of the innermost loop given a builder, a location and a range of // the body of the innermost loop given a builder, a location and a range of
// loop induction variables. // loop induction variables.
SmallVector<int64_t, 4> lowerBounds(tensorType.getRank(), /*Value=*/0); SmallVector<int64_t, 4> lowerBounds(tensorType.getRank(), /*Value=*/0);
SmallVector<int64_t, 4> steps(tensorType.getRank(), /*Value=*/1); SmallVector<int64_t, 4> steps(tensorType.getRank(), /*Value=*/1);
buildAffineLoopNest( affine::buildAffineLoopNest(
rewriter, loc, lowerBounds, tensorType.getShape(), steps, rewriter, loc, lowerBounds, tensorType.getShape(), steps,
[&](OpBuilder &nestedBuilder, Location loc, ValueRange ivs) { [&](OpBuilder &nestedBuilder, Location loc, ValueRange ivs) {
// Call the processing function with the rewriter, the memref operands, // Call the processing function with the rewriter, the memref operands,
// and the loop induction variables. This function will return the value // and the loop induction variables. This function will return the value
// to store at the current index. // to store at the current index.
Value valueToStore = processIteration(nestedBuilder, operands, ivs); Value valueToStore = processIteration(nestedBuilder, operands, ivs);
nestedBuilder.create<AffineStoreOp>(loc, valueToStore, alloc, ivs); nestedBuilder.create<affine::AffineStoreOp>(loc, valueToStore, alloc,
ivs);
}); });
// Replace this operation with the generated alloc. // Replace this operation with the generated alloc.
rewriter.replaceOp(op, alloc); rewriter.replaceOp(op, alloc);
} }
namespace { namespace {
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
Show All 14 Lines lowerOpToLoops(op, operands, rewriter,
ValueRange loopIvs) { ValueRange loopIvs) {
// Generate an adaptor for the remapped operands of the // Generate an adaptor for the remapped operands of the
// BinaryOp. This allows for using the nice named accessors // BinaryOp. This allows for using the nice named accessors
// that are generated by the ODS. // that are generated by the ODS.
typename BinaryOp::Adaptor binaryAdaptor(memRefOperands); typename BinaryOp::Adaptor binaryAdaptor(memRefOperands);
// Generate loads for the element of 'lhs' and 'rhs' at the // Generate loads for the element of 'lhs' and 'rhs' at the
// inner loop. // inner loop.
auto loadedLhs = builder.create<AffineLoadOp>( auto loadedLhs = builder.create<affine::AffineLoadOp>(
loc, binaryAdaptor.getLhs(), loopIvs); loc, binaryAdaptor.getLhs(), loopIvs);
auto loadedRhs = builder.create<AffineLoadOp>( auto loadedRhs = builder.create<affine::AffineLoadOp>(
loc, binaryAdaptor.getRhs(), loopIvs); loc, binaryAdaptor.getRhs(), loopIvs);
// Create the binary operation performed on the loaded // Create the binary operation performed on the loaded
// values. // values.
return builder.create<LoweredBinaryOp>(loc, loadedLhs, return builder.create<LoweredBinaryOp>(loc, loadedLhs,
loadedRhs); loadedRhs);
}); });
return success(); return success();
▲ Show 20 Lines • Show All 42 Lines • ▼ Show 20 Lines LogicalResult matchAndRewrite(toy::ConstantOp op,
// functor recursively walks the dimensions of the constant shape, // functor recursively walks the dimensions of the constant shape,
// generating a store when the recursion hits the base case. // generating a store when the recursion hits the base case.
SmallVector<Value, 2> indices; SmallVector<Value, 2> indices;
auto valueIt = constantValue.value_begin<FloatAttr>(); auto valueIt = constantValue.value_begin<FloatAttr>();
std::function<void(uint64_t)> storeElements = [&](uint64_t dimension) { std::function<void(uint64_t)> storeElements = [&](uint64_t dimension) {
// The last dimension is the base case of the recursion, at this point // The last dimension is the base case of the recursion, at this point
// we store the element at the given index. // we store the element at the given index.
if (dimension == valueShape.size()) { if (dimension == valueShape.size()) {
rewriter.create<AffineStoreOp>( rewriter.create<affine::AffineStoreOp>(
loc, rewriter.create<arith::ConstantOp>(loc, *valueIt++), alloc, loc, rewriter.create<arith::ConstantOp>(loc, *valueIt++), alloc,
llvm::ArrayRef(indices)); llvm::ArrayRef(indices));
return; return;
} }
// Otherwise, iterate over the current dimension and add the indices to // Otherwise, iterate over the current dimension and add the indices to
// the list. // the list.
for (uint64_t i = 0, e = valueShape[dimension]; i != e; ++i) { for (uint64_t i = 0, e = valueShape[dimension]; i != e; ++i) {
▲ Show 20 Lines • Show All 100 Lines • ▼ Show 20 Lines lowerOpToLoops(op, operands, rewriter,
// TransposeOp. This allows for using the nice named // TransposeOp. This allows for using the nice named
// accessors that are generated by the ODS. // accessors that are generated by the ODS.
toy::TransposeOpAdaptor transposeAdaptor(memRefOperands); toy::TransposeOpAdaptor transposeAdaptor(memRefOperands);
Value input = transposeAdaptor.getInput(); Value input = transposeAdaptor.getInput();
// Transpose the elements by generating a load from the // Transpose the elements by generating a load from the
// reverse indices. // reverse indices.
SmallVector<Value, 2> reverseIvs(llvm::reverse(loopIvs)); SmallVector<Value, 2> reverseIvs(llvm::reverse(loopIvs));
return builder.create<AffineLoadOp>(loc, input, return builder.create<affine::AffineLoadOp>(loc, input,
reverseIvs); reverseIvs);
}); });
return success(); return success();
} }
}; };
} // namespace } // namespace
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// ToyToAffineLoweringPass // ToyToAffineLoweringPass
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
/// This is a partial lowering to affine loops of the toy operations that are /// This is a partial lowering to affine loops of the toy operations that are
/// computationally intensive (like matmul for example...) while keeping the /// computationally intensive (like matmul for example...) while keeping the
/// rest of the code in the Toy dialect. /// rest of the code in the Toy dialect.
namespace { namespace {
struct ToyToAffineLoweringPass struct ToyToAffineLoweringPass
: public PassWrapper<ToyToAffineLoweringPass, OperationPass<ModuleOp>> { : public PassWrapper<ToyToAffineLoweringPass, OperationPass<ModuleOp>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ToyToAffineLoweringPass) MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ToyToAffineLoweringPass)
void getDependentDialects(DialectRegistry &registry) const override { void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<AffineDialect, func::FuncDialect, memref::MemRefDialect>(); registry.insert<affine::AffineDialect, func::FuncDialect,
memref::MemRefDialect>();
} }
void runOnOperation() final; void runOnOperation() final;
}; };
} // namespace } // namespace
void ToyToAffineLoweringPass::runOnOperation() { void ToyToAffineLoweringPass::runOnOperation() {
// The first thing to define is the conversion target. This will define the // The first thing to define is the conversion target. This will define the
// final target for this lowering. // final target for this lowering.
ConversionTarget target(getContext()); ConversionTarget target(getContext());
// We define the specific operations, or dialects, that are legal targets for // We define the specific operations, or dialects, that are legal targets for
// this lowering. In our case, we are lowering to a combination of the // this lowering. In our case, we are lowering to a combination of the
// `Affine`, `Arith`, `Func`, and `MemRef` dialects. // `Affine`, `Arith`, `Func`, and `MemRef` dialects.
target.addLegalDialect<AffineDialect, BuiltinDialect, arith::ArithDialect, target.addLegalDialect<affine::AffineDialect, BuiltinDialect,
func::FuncDialect, memref::MemRefDialect>(); arith::ArithDialect, func::FuncDialect,
memref::MemRefDialect>();
// We also define the Toy dialect as Illegal so that the conversion will fail // We also define the Toy dialect as Illegal so that the conversion will fail
// if any of these operations are *not* converted. Given that we actually want // if any of these operations are *not* converted. Given that we actually want
// a partial lowering, we explicitly mark the Toy operations that don't want // a partial lowering, we explicitly mark the Toy operations that don't want
// to lower, `toy.print`, as `legal`. `toy.print` will still need its operands // to lower, `toy.print`, as `legal`. `toy.print` will still need its operands
// to be updated though (as we convert from TensorType to MemRefType), so we // to be updated though (as we convert from TensorType to MemRefType), so we
// only treat it as `legal` if its operands are legal. // only treat it as `legal` if its operands are legal.
target.addIllegalDialect<toy::ToyDialect>(); target.addIllegalDialect<toy::ToyDialect>();
Show All 25 Lines

mlir/examples/toy/Ch7/toyc.cpp

Show First 20 Lines • Show All 161 Lines • ▼ Show 20 Lines if (isLoweringToAffine) {
// Add a few cleanups post lowering. // Add a few cleanups post lowering.
mlir::OpPassManager &optPM = pm.nest<mlir::func::FuncOp>(); mlir::OpPassManager &optPM = pm.nest<mlir::func::FuncOp>();
optPM.addPass(mlir::createCanonicalizerPass()); optPM.addPass(mlir::createCanonicalizerPass());
optPM.addPass(mlir::createCSEPass()); optPM.addPass(mlir::createCSEPass());
// Add optimizations if enabled. // Add optimizations if enabled.
if (enableOpt) { if (enableOpt) {
optPM.addPass(mlir::createLoopFusionPass()); optPM.addPass(mlir::affine::createLoopFusionPass());
optPM.addPass(mlir::createAffineScalarReplacementPass()); optPM.addPass(mlir::affine::createAffineScalarReplacementPass());
} }
} }
if (isLoweringToLLVM) { if (isLoweringToLLVM) {
// Finish lowering the toy IR to the LLVM dialect. // Finish lowering the toy IR to the LLVM dialect.
pm.addPass(mlir::toy::createLowerToLLVMPass()); pm.addPass(mlir::toy::createLowerToLLVMPass());
// This is necessary to have line tables emitted and basic // This is necessary to have line tables emitted and basic
// debugger working. In the future we will add proper debug information // debugger working. In the future we will add proper debug information
▲ Show 20 Lines • Show All 140 Lines • Show Last 20 Lines

mlir/include/mlir/Conversion/AffineToStandard/AffineToStandard.h

//===- AffineToStandard.h - Convert Affine to Standard dialect --*- C++ -*-===// //===- AffineToStandard.h - Convert Affine to Standard dialect --*- C++ -*-===//
// //
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information. // See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// //
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
#ifndef MLIR_CONVERSION_AFFINETOSTANDARD_AFFINETOSTANDARD_H #ifndef MLIR_CONVERSION_AFFINETOSTANDARD_AFFINETOSTANDARD_H
#define MLIR_CONVERSION_AFFINETOSTANDARD_AFFINETOSTANDARD_H #define MLIR_CONVERSION_AFFINETOSTANDARD_AFFINETOSTANDARD_H
#include "mlir/Support/LLVM.h" #include "mlir/Support/LLVM.h"
namespace mlir { namespace mlir {
class AffineForOp;
class Location; class Location;
struct LogicalResult; struct LogicalResult;
class OpBuilder; class OpBuilder;
class Pass; class Pass;
class RewritePattern; class RewritePattern;
class RewritePatternSet;
class Value; class Value;
class ValueRange; class ValueRange;
class RewritePatternSet; namespace affine {
class AffineForOp;
} // namespace affine
#define GEN_PASS_DECL_CONVERTAFFINETOSTANDARD #define GEN_PASS_DECL_CONVERTAFFINETOSTANDARD
#include "mlir/Conversion/Passes.h.inc" #include "mlir/Conversion/Passes.h.inc"
/// Collect a set of patterns to convert from the Affine dialect to the Standard /// Collect a set of patterns to convert from the Affine dialect to the Standard
/// dialect, in particular convert structured affine control flow into CFG /// dialect, in particular convert structured affine control flow into CFG
/// branch-based control flow. /// branch-based control flow.
void populateAffineToStdConversionPatterns(RewritePatternSet &patterns); void populateAffineToStdConversionPatterns(RewritePatternSet &patterns);
/// Collect a set of patterns to convert vector-related Affine ops to the Vector /// Collect a set of patterns to convert vector-related Affine ops to the Vector
/// dialect. /// dialect.
void populateAffineToVectorConversionPatterns(RewritePatternSet &patterns); void populateAffineToVectorConversionPatterns(RewritePatternSet &patterns);
/// Emit code that computes the lower bound of the given affine loop using /// Emit code that computes the lower bound of the given affine loop using
/// standard arithmetic operations. /// standard arithmetic operations.
Value lowerAffineLowerBound(AffineForOp op, OpBuilder &builder); Value lowerAffineLowerBound(affine::AffineForOp op, OpBuilder &builder);
/// Emit code that computes the upper bound of the given affine loop using /// Emit code that computes the upper bound of the given affine loop using
/// standard arithmetic operations. /// standard arithmetic operations.
Value lowerAffineUpperBound(AffineForOp op, OpBuilder &builder); Value lowerAffineUpperBound(affine::AffineForOp op, OpBuilder &builder);
/// Lowers affine control flow operations (ForStmt, IfStmt and AffineApplyOp) /// Lowers affine control flow operations (ForStmt, IfStmt and AffineApplyOp)
/// to equivalent lower-level constructs (flow of basic blocks and arithmetic /// to equivalent lower-level constructs (flow of basic blocks and arithmetic
/// primitives). /// primitives).
std::unique_ptr<Pass> createLowerAffinePass(); std::unique_ptr<Pass> createLowerAffinePass();
} // namespace mlir } // namespace mlir
#endif // MLIR_CONVERSION_AFFINETOSTANDARD_AFFINETOSTANDARD_H #endif // MLIR_CONVERSION_AFFINETOSTANDARD_AFFINETOSTANDARD_H

mlir/include/mlir/Conversion/Passes.td

Show First 20 Lines • Show All 850 Lines • ▼ Show 20 Lines let options = [
Option<"numThreadDims", "gpu-thread-dims", "unsigned", /*default=*/"1u", Option<"numThreadDims", "gpu-thread-dims", "unsigned", /*default=*/"1u",
"Number of GPU thread dimensions for mapping"> "Number of GPU thread dimensions for mapping">
]; ];
} }
def ConvertParallelLoopToGpu : Pass<"convert-parallel-loops-to-gpu"> { def ConvertParallelLoopToGpu : Pass<"convert-parallel-loops-to-gpu"> {
let summary = "Convert mapped scf.parallel ops to gpu launch operations"; let summary = "Convert mapped scf.parallel ops to gpu launch operations";
let constructor = "mlir::createParallelLoopToGpuPass()"; let constructor = "mlir::createParallelLoopToGpuPass()";
let dependentDialects = ["AffineDialect", "gpu::GPUDialect"]; let dependentDialects = ["affine::AffineDialect", "gpu::GPUDialect"];
} }
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// ShapeToStandard // ShapeToStandard
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
def ConvertShapeToStandard : Pass<"convert-shape-to-std", "ModuleOp"> { def ConvertShapeToStandard : Pass<"convert-shape-to-std", "ModuleOp"> {
let summary = "Convert operations from the shape dialect into the standard " let summary = "Convert operations from the shape dialect into the standard "
▲ Show 20 Lines • Show All 160 Lines • ▼ Show 20 Lines
// VectorToGPU // VectorToGPU
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
def ConvertVectorToGPU : Pass<"convert-vector-to-gpu"> { def ConvertVectorToGPU : Pass<"convert-vector-to-gpu"> {
let summary = "Lower the operations from the vector dialect into the GPU " let summary = "Lower the operations from the vector dialect into the GPU "
"dialect"; "dialect";
let constructor = "mlir::createConvertVectorToGPUPass()"; let constructor = "mlir::createConvertVectorToGPUPass()";
let dependentDialects = [ let dependentDialects = [
"memref::MemRefDialect", "gpu::GPUDialect", "AffineDialect", "memref::MemRefDialect", "gpu::GPUDialect", "affine::AffineDialect",
"vector::VectorDialect", "nvgpu::NVGPUDialect" "vector::VectorDialect", "nvgpu::NVGPUDialect"
]; ];
let options = [ let options = [
Option<"useNvGpu", "use-nvgpu", "bool", /*default=*/"false", Option<"useNvGpu", "use-nvgpu", "bool", /*default=*/"false",
"convert to NvGPU ops instead of GPU dialect ops"> "convert to NvGPU ops instead of GPU dialect ops">
]; ];
} }
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// VectorToSCF // VectorToSCF
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
def ConvertVectorToSCF : Pass<"convert-vector-to-scf"> { def ConvertVectorToSCF : Pass<"convert-vector-to-scf"> {
let summary = "Lower the operations from the vector dialect into the SCF " let summary = "Lower the operations from the vector dialect into the SCF "
"dialect"; "dialect";
let constructor = "mlir::createConvertVectorToSCFPass()"; let constructor = "mlir::createConvertVectorToSCFPass()";
let dependentDialects = [ let dependentDialects = [
"AffineDialect", "affine::AffineDialect",
"memref::MemRefDialect", "memref::MemRefDialect",
"scf::SCFDialect", "scf::SCFDialect",
"tensor::TensorDialect" "tensor::TensorDialect"
]; ];
let options = [ let options = [
Option<"fullUnroll", "full-unroll", "bool", /*default=*/"false", Option<"fullUnroll", "full-unroll", "bool", /*default=*/"false",
"Perform full unrolling when converting vector transfers to SCF">, "Perform full unrolling when converting vector transfers to SCF">,
Option<"targetRank", "target-rank", "unsigned", /*default=*/"1", Option<"targetRank", "target-rank", "unsigned", /*default=*/"1",
▲ Show 20 Lines • Show All 66 Lines • Show Last 20 Lines

mlir/include/mlir/Conversion/SCFToGPU/SCFToGPU.h

//===- SCFToGPU.h - Convert loop nests to GPU kernels -----------*- C++ -*-===// //===- SCFToGPU.h - Convert loop nests to GPU kernels -----------*- C++ -*-===//
// //
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information. // See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// //
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
#ifndef MLIR_CONVERSION_SCFTOGPU_SCFTOGPU_H_ #ifndef MLIR_CONVERSION_SCFTOGPU_SCFTOGPU_H_
#define MLIR_CONVERSION_SCFTOGPU_SCFTOGPU_H_ #define MLIR_CONVERSION_SCFTOGPU_SCFTOGPU_H_
#include "mlir/Support/LLVM.h" #include "mlir/Support/LLVM.h"
namespace mlir { namespace mlir {
class AffineForOp;
class ConversionTarget; class ConversionTarget;
struct LogicalResult; struct LogicalResult;
class MLIRContext; class MLIRContext;
class Value; class Value;
class Operation; class Operation;
class RewritePatternSet; class RewritePatternSet;
namespace affine {
class AffineForOp;
} // namespace affine
namespace scf { namespace scf {
class ForOp; class ForOp;
} // namespace scf } // namespace scf
/// Convert a perfect affine loop nest with the outermost loop identified by /// Convert a perfect affine loop nest with the outermost loop identified by
/// `forOp` into a gpu::Launch operation. Map `numBlockDims` outer loops to /// `forOp` into a gpu::Launch operation. Map `numBlockDims` outer loops to
/// GPU blocks and `numThreadDims` to GPU threads. The bounds of the loops that /// GPU blocks and `numThreadDims` to GPU threads. The bounds of the loops that
/// are mapped should be independent of the induction variables of the other /// are mapped should be independent of the induction variables of the other
/// mapped loops. /// mapped loops.
/// ///
/// No check on the size of the block or grid, or on the validity of /// No check on the size of the block or grid, or on the validity of
/// parallelization is performed, it is under the responsibility of the caller /// parallelization is performed, it is under the responsibility of the caller
/// to strip-mine the loops and to perform the dependence analysis before /// to strip-mine the loops and to perform the dependence analysis before
/// calling the conversion. /// calling the conversion.
// TODO: Consider removing this in favor of affine.for -> affine.parallel // TODO: Consider removing this in favor of affine.for -> affine.parallel
// detection followed by an affine.parallel -> scf.parallel -> gpu.launch // detection followed by an affine.parallel -> scf.parallel -> gpu.launch
// conversion // conversion
LogicalResult convertAffineLoopNestToGPULaunch(AffineForOp forOp, LogicalResult convertAffineLoopNestToGPULaunch(affine::AffineForOp forOp,
unsigned numBlockDims, unsigned numBlockDims,
unsigned numThreadDims); unsigned numThreadDims);
/// Adds the conversion pattern from `scf.parallel` to `gpu.launch` to the /// Adds the conversion pattern from `scf.parallel` to `gpu.launch` to the
/// provided pattern list. /// provided pattern list.
void populateParallelLoopToGPUPatterns(RewritePatternSet &patterns); void populateParallelLoopToGPUPatterns(RewritePatternSet &patterns);
/// Configures the rewrite target such that only `scf.parallel` operations that /// Configures the rewrite target such that only `scf.parallel` operations that
Show All 9 Lines

mlir/include/mlir/Dialect/Affine/Analysis/AffineAnalysis.h

Show All 15 Lines
#define MLIR_DIALECT_AFFINE_ANALYSIS_AFFINEANALYSIS_H #define MLIR_DIALECT_AFFINE_ANALYSIS_AFFINEANALYSIS_H
#include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/IR/Value.h" #include "mlir/IR/Value.h"
#include "llvm/ADT/SmallVector.h" #include "llvm/ADT/SmallVector.h"
#include <optional> #include <optional>
namespace mlir { namespace mlir {
class Operation;
namespace affine {
class AffineApplyOp; class AffineApplyOp;
class AffineForOp; class AffineForOp;
class AffineValueMap; class AffineValueMap;
class FlatAffineRelation; class FlatAffineRelation;
class FlatAffineValueConstraints; class FlatAffineValueConstraints;
class Operation;
/// A description of a (parallelizable) reduction in an affine loop. /// A description of a (parallelizable) reduction in an affine loop.
struct LoopReduction { struct LoopReduction {
/// Reduction kind. /// Reduction kind.
arith::AtomicRMWKind kind; arith::AtomicRMWKind kind;
/// Position of the iteration argument that acts as accumulator. /// Position of the iteration argument that acts as accumulator.
unsigned iterArgPosition; unsigned iterArgPosition;
▲ Show 20 Lines • Show All 147 Lines • ▼ Show 20 Lines
/// Returns in 'depCompsVec', dependence components for dependences between all /// Returns in 'depCompsVec', dependence components for dependences between all
/// load and store ops in loop nest rooted at 'forOp', at loop depths in range /// load and store ops in loop nest rooted at 'forOp', at loop depths in range
/// [1, maxLoopDepth]. /// [1, maxLoopDepth].
void getDependenceComponents( void getDependenceComponents(
AffineForOp forOp, unsigned maxLoopDepth, AffineForOp forOp, unsigned maxLoopDepth,
std::vector<SmallVector<DependenceComponent, 2>> *depCompsVec); std::vector<SmallVector<DependenceComponent, 2>> *depCompsVec);
} // namespace affine
} // namespace mlir } // namespace mlir
#endif // MLIR_DIALECT_AFFINE_ANALYSIS_AFFINEANALYSIS_H #endif // MLIR_DIALECT_AFFINE_ANALYSIS_AFFINEANALYSIS_H

mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h

Show All 16 Lines
#include "mlir/Analysis/Presburger/IntegerRelation.h" #include "mlir/Analysis/Presburger/IntegerRelation.h"
#include "mlir/Analysis/Presburger/Matrix.h" #include "mlir/Analysis/Presburger/Matrix.h"
#include "mlir/IR/AffineExpr.h" #include "mlir/IR/AffineExpr.h"
#include "mlir/IR/OpDefinition.h" #include "mlir/IR/OpDefinition.h"
#include "mlir/Support/LogicalResult.h" #include "mlir/Support/LogicalResult.h"
#include <optional> #include <optional>
namespace mlir { namespace mlir {
class AffineCondition;
class AffineForOp;
class AffineIfOp;
class AffineParallelOp;
class AffineMap; class AffineMap;
class AffineValueMap;
class IntegerSet; class IntegerSet;
class MLIRContext;
class Value;
class MemRefType; class MemRefType;
class MLIRContext;
struct MutableAffineMap; struct MutableAffineMap;
class Value;
namespace presburger { namespace presburger {
class MultiAffineFunction; class MultiAffineFunction;
} // namespace presburger } // namespace presburger
namespace affine {
class AffineCondition;
class AffineForOp;
class AffineIfOp;
class AffineParallelOp;
class AffineValueMap;
/// FlatAffineValueConstraints is an extension of FlatLinearValueConstraints /// FlatAffineValueConstraints is an extension of FlatLinearValueConstraints
/// with helper functions for Affine dialect ops. /// with helper functions for Affine dialect ops.
class FlatAffineValueConstraints : public FlatLinearValueConstraints { class FlatAffineValueConstraints : public FlatLinearValueConstraints {
public: public:
using FlatLinearValueConstraints::FlatLinearValueConstraints; using FlatLinearValueConstraints::FlatLinearValueConstraints;
/// Return the kind of this object. /// Return the kind of this object.
Kind getKind() const override { return Kind::FlatAffineValueConstraints; } Kind getKind() const override { return Kind::FlatAffineValueConstraints; }
▲ Show 20 Lines • Show All 199 Lines • ▼ Show 20 Lines
/// ///
/// For AffineValueMap, the domain and symbols have Value set corresponding to /// For AffineValueMap, the domain and symbols have Value set corresponding to
/// the Value in `map`. Returns failure if the AffineMap could not be flattened /// the Value in `map`. Returns failure if the AffineMap could not be flattened
/// (i.e., semi-affine is not yet handled). /// (i.e., semi-affine is not yet handled).
LogicalResult getRelationFromMap(AffineMap &map, FlatAffineRelation &rel); LogicalResult getRelationFromMap(AffineMap &map, FlatAffineRelation &rel);
LogicalResult getRelationFromMap(const AffineValueMap &map, LogicalResult getRelationFromMap(const AffineValueMap &map,
FlatAffineRelation &rel); FlatAffineRelation &rel);
} // namespace affine
} // namespace mlir } // namespace mlir
#endif // MLIR_DIALECT_AFFINE_ANALYSIS_AFFINESTRUCTURES_H #endif // MLIR_DIALECT_AFFINE_ANALYSIS_AFFINESTRUCTURES_H

mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h

Show All 12 Lines
#ifndef MLIR_DIALECT_AFFINE_ANALYSIS_LOOPANALYSIS_H #ifndef MLIR_DIALECT_AFFINE_ANALYSIS_LOOPANALYSIS_H
#define MLIR_DIALECT_AFFINE_ANALYSIS_LOOPANALYSIS_H #define MLIR_DIALECT_AFFINE_ANALYSIS_LOOPANALYSIS_H
#include "mlir/Support/LLVM.h" #include "mlir/Support/LLVM.h"
#include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/ArrayRef.h"
#include <optional> #include <optional>
namespace mlir { namespace mlir {
class AffineExpr; class AffineExpr;
class AffineForOp;
class AffineMap; class AffineMap;
class BlockArgument; class BlockArgument;
class MemRefType; class MemRefType;
class NestedPattern;
class Operation; class Operation;
class Value; class Value;
namespace affine {
class AffineForOp;
class NestedPattern;
/// Returns the trip count of the loop as an affine map with its corresponding /// Returns the trip count of the loop as an affine map with its corresponding
/// operands if the latter is expressible as an affine expression, and nullptr /// operands if the latter is expressible as an affine expression, and nullptr
/// otherwise. This method always succeeds as long as the lower bound is not a /// otherwise. This method always succeeds as long as the lower bound is not a
/// multi-result map. The trip count expression is simplified before returning. /// multi-result map. The trip count expression is simplified before returning.
/// This method only utilizes map composition to construct lower and upper /// This method only utilizes map composition to construct lower and upper
/// bounds before computing the trip count expressions /// bounds before computing the trip count expressions
void getTripCountMapAndOperands(AffineForOp forOp, AffineMap *map, void getTripCountMapAndOperands(AffineForOp forOp, AffineMap *map,
SmallVectorImpl<Value> *operands); SmallVectorImpl<Value> *operands);
Show All 39 Lines
/// Checks where SSA dominance would be violated if a for op's body /// Checks where SSA dominance would be violated if a for op's body
/// operations are shifted by the specified shifts. This method checks if a /// operations are shifted by the specified shifts. This method checks if a
/// 'def' and all its uses have the same shift factor. /// 'def' and all its uses have the same shift factor.
// TODO: extend this to check for memory-based dependence violation when we have // TODO: extend this to check for memory-based dependence violation when we have
// the support. // the support.
bool isOpwiseShiftValid(AffineForOp forOp, ArrayRef<uint64_t> shifts); bool isOpwiseShiftValid(AffineForOp forOp, ArrayRef<uint64_t> shifts);
} // namespace affine
} // namespace mlir } // namespace mlir
#endif // MLIR_DIALECT_AFFINE_ANALYSIS_LOOPANALYSIS_H #endif // MLIR_DIALECT_AFFINE_ANALYSIS_LOOPANALYSIS_H

mlir/include/mlir/Dialect/Affine/Analysis/NestedMatcher.h

//===- NestedMacher.h - Nested matcher for Function -------------*- C++ -*-===// //===- NestedMacher.h - Nested matcher for Function -------------*- C++ -*-===//
// //
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information. // See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// //
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
#ifndef MLIR_DIALECT_AFFINE_ANALYSIS_NESTEDMATCHER_H #ifndef MLIR_DIALECT_AFFINE_ANALYSIS_NESTEDMATCHER_H
#define MLIR_DIALECT_AFFINE_ANALYSIS_NESTEDMATCHER_H #define MLIR_DIALECT_AFFINE_ANALYSIS_NESTEDMATCHER_H
#include "mlir/IR/BuiltinOps.h" #include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/Operation.h" #include "mlir/IR/Operation.h"
#include "llvm/Support/Allocator.h" #include "llvm/Support/Allocator.h"
namespace mlir { namespace mlir {
class Operation;
namespace affine {
class NestedPattern; class NestedPattern;
class Operation;
/// An NestedPattern captures nested patterns in the IR. /// An NestedPattern captures nested patterns in the IR.
/// It is used in conjunction with a scoped NestedPatternContext which is an /// It is used in conjunction with a scoped NestedPatternContext which is an
/// llvm::BumpPtrAllocator that handles memory allocations efficiently and /// llvm::BumpPtrAllocator that handles memory allocations efficiently and
/// avoids ownership issues. /// avoids ownership issues.
/// ///
/// In order to use NestedPatterns, first create a scoped context. /// In order to use NestedPatterns, first create a scoped context.
/// When the context goes out of scope, everything is freed. /// When the context goes out of scope, everything is freed.
▲ Show 20 Lines • Show All 158 Lines • ▼ Show 20 Lines
NestedPattern For(const FilterFunctionType &filter, NestedPattern For(const FilterFunctionType &filter,
ArrayRef<NestedPattern> nested = {}); ArrayRef<NestedPattern> nested = {});
bool isParallelLoop(Operation &op); bool isParallelLoop(Operation &op);
bool isReductionLoop(Operation &op); bool isReductionLoop(Operation &op);
bool isLoadOrStore(Operation &op); bool isLoadOrStore(Operation &op);
} // namespace matcher } // namespace matcher
} // namespace affine
} // namespace mlir } // namespace mlir
#endif // MLIR_DIALECT_AFFINE_ANALYSIS_NESTEDMATCHER_H #endif // MLIR_DIALECT_AFFINE_ANALYSIS_NESTEDMATCHER_H

mlir/include/mlir/Dialect/Affine/Analysis/Utils.h

Show All 16 Lines
#define MLIR_DIALECT_AFFINE_ANALYSIS_UTILS_H #define MLIR_DIALECT_AFFINE_ANALYSIS_UTILS_H
#include "mlir/Dialect/Affine/Analysis/AffineStructures.h" #include "mlir/Dialect/Affine/Analysis/AffineStructures.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Affine/IR/AffineOps.h"
#include <memory> #include <memory>
#include <optional> #include <optional>
namespace mlir { namespace mlir {
class AffineForOp;
class AffineValueMap;
class Block; class Block;
class Location; class Location;
struct MemRefAccess;
class Operation; class Operation;
class Value; class Value;
namespace affine {
class AffineForOp;
class AffineValueMap;
struct MemRefAccess;
// LoopNestStateCollector walks loop nests and collects load and store // LoopNestStateCollector walks loop nests and collects load and store
// operations, and whether or not a region holding op other than ForOp and IfOp // operations, and whether or not a region holding op other than ForOp and IfOp
// was encountered in the loop nest. // was encountered in the loop nest.
struct LoopNestStateCollector { struct LoopNestStateCollector {
SmallVector<AffineForOp, 4> forOps; SmallVector<AffineForOp, 4> forOps;
SmallVector<Operation *, 4> loadOpInsts; SmallVector<Operation *, 4> loadOpInsts;
SmallVector<Operation *, 4> storeOpInsts; SmallVector<Operation *, 4> storeOpInsts;
bool hasNonAffineRegionOp = false; bool hasNonAffineRegionOp = false;
▲ Show 20 Lines • Show All 529 Lines • ▼ Show 20 Lines
/// Try to simplify the given affine.min or affine.max op to an affine map with /// Try to simplify the given affine.min or affine.max op to an affine map with
/// a single result and operands, taking into account the specified constraint /// a single result and operands, taking into account the specified constraint
/// set. Return failure if no simplified version could be found. /// set. Return failure if no simplified version could be found.
FailureOr<AffineValueMap> FailureOr<AffineValueMap>
simplifyConstrainedMinMaxOp(Operation *op, simplifyConstrainedMinMaxOp(Operation *op,
FlatAffineValueConstraints constraints); FlatAffineValueConstraints constraints);
} // namespace affine
} // namespace mlir } // namespace mlir
#endif // MLIR_DIALECT_AFFINE_ANALYSIS_UTILS_H #endif // MLIR_DIALECT_AFFINE_ANALYSIS_UTILS_H

mlir/include/mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.h

Show All 11 Lines
#ifndef MLIR_DIALECT_AFFINE_IR_AFFINEMEMORYOPINTERFACES_H #ifndef MLIR_DIALECT_AFFINE_IR_AFFINEMEMORYOPINTERFACES_H
#define MLIR_DIALECT_AFFINE_IR_AFFINEMEMORYOPINTERFACES_H #define MLIR_DIALECT_AFFINE_IR_AFFINEMEMORYOPINTERFACES_H
#include "mlir/IR/AffineMap.h" #include "mlir/IR/AffineMap.h"
#include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/OpDefinition.h" #include "mlir/IR/OpDefinition.h"
namespace mlir {
#include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.h.inc" #include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.h.inc"
} // namespace mlir
#endif // MLIR_DIALECT_AFFINE_IR_AFFINEMEMORYOPINTERFACES_H #endif // MLIR_DIALECT_AFFINE_IR_AFFINEMEMORYOPINTERFACES_H

mlir/include/mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.td

Show All 14 Lines
include "mlir/IR/OpBase.td" include "mlir/IR/OpBase.td"
def AffineReadOpInterface : OpInterface<"AffineReadOpInterface"> { def AffineReadOpInterface : OpInterface<"AffineReadOpInterface"> {
let description = [{ let description = [{
Interface to query characteristics of read-like ops with affine Interface to query characteristics of read-like ops with affine
restrictions. restrictions.
}]; }];
let cppNamespace = "::mlir::affine";
let methods = [ let methods = [
InterfaceMethod< InterfaceMethod<
/*desc=*/"Returns the memref operand to read from.", /*desc=*/"Returns the memref operand to read from.",
/*retTy=*/"Value", /*retTy=*/"::mlir::Value",
/*methodName=*/"getMemRef", /*methodName=*/"getMemRef",
/*args=*/(ins), /*args=*/(ins),
/*methodBody*/[{}], /*methodBody*/[{}],
/*defaultImplementation=*/ [{ /*defaultImplementation=*/ [{
ConcreteOp op = cast<ConcreteOp>(this->getOperation()); return $_op.getOperand($_op.getMemRefOperandIndex());
return op.getOperand(op.getMemRefOperandIndex());
}] }]
>, >,
InterfaceMethod< InterfaceMethod<
/*desc=*/"Returns the type of the memref operand.", /*desc=*/"Returns the type of the memref operand.",
/*retTy=*/"MemRefType", /*retTy=*/"::mlir::MemRefType",
/*methodName=*/"getMemRefType", /*methodName=*/"getMemRefType",
/*args=*/(ins), /*args=*/(ins),
/*methodBody=*/[{}], /*methodBody=*/[{}],
/*defaultImplementation=*/[{ /*defaultImplementation=*/[{
ConcreteOp op = cast<ConcreteOp>(this->getOperation()); return $_op.getMemRef().getType().template cast<::mlir::MemRefType>();
return op.getMemRef().getType().template cast<MemRefType>();
}] }]
>, >,
InterfaceMethod< InterfaceMethod<
/*desc=*/"Returns affine map operands.", /*desc=*/"Returns affine map operands.",
/*retTy=*/"Operation::operand_range", /*retTy=*/"::mlir::Operation::operand_range",
/*methodName=*/"getMapOperands", /*methodName=*/"getMapOperands",
/*args=*/(ins), /*args=*/(ins),
/*methodBody=*/[{}], /*methodBody=*/[{}],
/*defaultImplementation=*/[{ /*defaultImplementation=*/[{
ConcreteOp op = cast<ConcreteOp>(this->getOperation()); return llvm::drop_begin($_op.getOperands(), 1);
return llvm::drop_begin(op.getOperands(), 1);
}] }]
>, >,
InterfaceMethod< InterfaceMethod<
/*desc=*/[{ /*desc=*/[{
Returns the affine map used to index the memref for this operation. Returns the affine map used to index the memref for this operation.
}], }],
/*retTy=*/"AffineMap", /*retTy=*/"::mlir::AffineMap",
/*methodName=*/"getAffineMap", /*methodName=*/"getAffineMap",
/*args=*/(ins), /*args=*/(ins),
/*methodBody=*/[{}], /*methodBody=*/[{}],
/*defaultImplementation=*/[{ /*defaultImplementation=*/[{
ConcreteOp op = cast<ConcreteOp>(this->getOperation()); return $_op.getAffineMapAttr().getValue();
return op.getAffineMapAttr().getValue();
}] }]
>, >,
InterfaceMethod< InterfaceMethod<
/*desc=*/"Returns the value read by this operation.", /*desc=*/"Returns the value read by this operation.",
/*retTy=*/"Value", /*retTy=*/"::mlir::Value",
/*methodName=*/"getValue", /*methodName=*/"getValue",
/*args=*/(ins), /*args=*/(ins),
/*methodBody=*/[{}], /*methodBody=*/[{}],
/*defaultImplementation=*/[{ /*defaultImplementation=*/[{
return cast<ConcreteOp>(this->getOperation()); return $_op;
}] }]
>, >,
]; ];
} }
def AffineWriteOpInterface : OpInterface<"AffineWriteOpInterface"> { def AffineWriteOpInterface : OpInterface<"AffineWriteOpInterface"> {
let description = [{ let description = [{
Interface to query characteristics of write-like ops with affine Interface to query characteristics of write-like ops with affine
restrictions. restrictions.
}]; }];
let cppNamespace = "::mlir::affine";
let methods = [ let methods = [
InterfaceMethod< InterfaceMethod<
/*desc=*/"Returns the memref operand to write to.", /*desc=*/"Returns the memref operand to write to.",
/*retTy=*/"Value", /*retTy=*/"::mlir::Value",
/*methodName=*/"getMemRef", /*methodName=*/"getMemRef",
/*args=*/(ins), /*args=*/(ins),
/*methodBody=*/[{}], /*methodBody=*/[{}],
/*defaultImplementation=*/[{ /*defaultImplementation=*/[{
ConcreteOp op = cast<ConcreteOp>(this->getOperation()); return $_op.getOperand($_op.getMemRefOperandIndex());
return op.getOperand(op.getMemRefOperandIndex());
}] }]
>, >,
InterfaceMethod< InterfaceMethod<
/*desc=*/"Returns the type of the memref operand.", /*desc=*/"Returns the type of the memref operand.",
/*retTy=*/"MemRefType", /*retTy=*/"::mlir::MemRefType",
/*methodName=*/"getMemRefType", /*methodName=*/"getMemRefType",
/*args=*/(ins), /*args=*/(ins),
/*methodBody=*/[{}], /*methodBody=*/[{}],
/*defaultImplementation=*/[{ /*defaultImplementation=*/[{
ConcreteOp op = cast<ConcreteOp>(this->getOperation()); return $_op.getMemRef().getType().template cast<::mlir::MemRefType>();
return op.getMemRef().getType().template cast<MemRefType>();
}] }]
>, >,
InterfaceMethod< InterfaceMethod<
/*desc=*/"Returns affine map operands.", /*desc=*/"Returns affine map operands.",
/*retTy=*/"Operation::operand_range", /*retTy=*/"::mlir::Operation::operand_range",
/*methodName=*/"getMapOperands", /*methodName=*/"getMapOperands",
/*args=*/(ins), /*args=*/(ins),
/*methodBody=*/[{}], /*methodBody=*/[{}],
/*defaultImplementation=*/[{ /*defaultImplementation=*/[{
ConcreteOp op = cast<ConcreteOp>(this->getOperation()); return llvm::drop_begin($_op.getOperands(), 2);
return llvm::drop_begin(op.getOperands(), 2);
}] }]
>, >,
InterfaceMethod< InterfaceMethod<
/*desc=*/[{ /*desc=*/[{
Returns the affine map used to index the memref for this operation. Returns the affine map used to index the memref for this operation.
}], }],
/*retTy=*/"AffineMap", /*retTy=*/"::mlir::AffineMap",
/*methodName=*/"getAffineMap", /*methodName=*/"getAffineMap",
/*args=*/(ins), /*args=*/(ins),
/*methodName=*/[{}], /*methodName=*/[{}],
/*defaultImplementation=*/[{ /*defaultImplementation=*/[{
ConcreteOp op = cast<ConcreteOp>(this->getOperation()); return $_op.getAffineMapAttr().getValue();
return op.getAffineMapAttr().getValue();
}] }]
>, >,
InterfaceMethod< InterfaceMethod<
/*desc=*/"Returns the value to store.", /*desc=*/"Returns the value to store.",
/*retTy=*/"Value", /*retTy=*/"::mlir::Value",
/*methodName=*/"getValueToStore", /*methodName=*/"getValueToStore",
/*args=*/(ins), /*args=*/(ins),
/*methodBody=*/[{}], /*methodBody=*/[{}],
/*defaultImplementation=*/[{ /*defaultImplementation=*/[{
ConcreteOp op = cast<ConcreteOp>(this->getOperation()); return $_op.getOperand($_op.getStoredValOperandIndex());
return op.getOperand(op.getStoredValOperandIndex());
}] }]
>, >,
]; ];
} }
def AffineMapAccessInterface : OpInterface<"AffineMapAccessInterface"> { def AffineMapAccessInterface : OpInterface<"AffineMapAccessInterface"> {
let description = [{ let description = [{
Interface to query the AffineMap used to dereference and access a given Interface to query the AffineMap used to dereference and access a given
memref. Implementers of this interface must operate on at least one memref. Implementers of this interface must operate on at least one
memref operand. The memref argument given to this interface much match memref operand. The memref argument given to this interface much match
one of those memref operands. one of those memref operands.
}]; }];
let cppNamespace = "::mlir::affine";
let methods = [ let methods = [
InterfaceMethod< InterfaceMethod<
/*desc=*/"Returns the AffineMapAttr associated with 'memref'.", /*desc=*/"Returns the AffineMapAttr associated with 'memref'.",
/*retTy=*/"NamedAttribute", /*retTy=*/"::mlir::NamedAttribute",
/*methodName=*/"getAffineMapAttrForMemRef", /*methodName=*/"getAffineMapAttrForMemRef",
/*args=*/(ins "Value":$memref), /*args=*/(ins "::mlir::Value":$memref),
/*methodBody=*/[{}], /*methodBody=*/[{}],
/*defaultImplementation=*/[{ /*defaultImplementation=*/[{
ConcreteOp op = cast<ConcreteOp>(this->getOperation()); assert(memref == $_op.getMemRef() &&
assert(memref == op.getMemRef() &&
"Expected memref argument to match memref operand"); "Expected memref argument to match memref operand");
return {StringAttr::get(op.getContext(), op.getMapAttrStrName()), return {::mlir::StringAttr::get(
op.getAffineMapAttr()}; $_op.getContext(), $_op.getMapAttrStrName()),
$_op.getAffineMapAttr()};
}] }]
>, >,
]; ];
} }
#endif // AFFINEMEMORYOPINTERFACES #endif // AFFINEMEMORYOPINTERFACES

mlir/include/mlir/Dialect/Affine/IR/AffineOps.h

Show All 16 Lines
#include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.h" #include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.h"
#include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/IR/AffineMap.h" #include "mlir/IR/AffineMap.h"
#include "mlir/IR/Builders.h" #include "mlir/IR/Builders.h"
#include "mlir/Interfaces/ControlFlowInterfaces.h" #include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "mlir/Interfaces/LoopLikeInterface.h" #include "mlir/Interfaces/LoopLikeInterface.h"
namespace mlir { namespace mlir {
namespace affine {
class AffineApplyOp; class AffineApplyOp;
class AffineBound; class AffineBound;
class AffineValueMap; class AffineValueMap;
/// TODO: These should be renamed if they are on the mlir namespace.
/// Ideally, they should go in a mlir::affine:: namespace.
/// A utility function to check if a value is defined at the top level of an /// A utility function to check if a value is defined at the top level of an
/// op with trait `AffineScope` or is a region argument for such an op. A value /// op with trait `AffineScope` or is a region argument for such an op. A value
/// of index type defined at the top level is always a valid symbol for all its /// of index type defined at the top level is always a valid symbol for all its
/// uses. /// uses.
bool isTopLevelValue(Value value); bool isTopLevelValue(Value value);
/// A utility function to check if a value is defined at the top level of /// A utility function to check if a value is defined at the top level of
/// `region` or is an argument of `region`. A value of index type defined at the /// `region` or is an argument of `region`. A value of index type defined at the
▲ Show 20 Lines • Show All 393 Lines • ▼ Show 20 Lines
/// into `map`, maps of AffineApplyOps whose results are the values in /// into `map`, maps of AffineApplyOps whose results are the values in
/// `operands`, iteratively until no more of `operands` are the result of an /// `operands`, iteratively until no more of `operands` are the result of an
/// AffineApplyOp. When this function returns, `map` becomes the composed affine /// AffineApplyOp. When this function returns, `map` becomes the composed affine
/// map, and each Value in `operands` is guaranteed to be either a loop IV or a /// map, and each Value in `operands` is guaranteed to be either a loop IV or a
/// terminal symbol, i.e., a symbol defined at the top level or a block/function /// terminal symbol, i.e., a symbol defined at the top level or a block/function
/// argument. /// argument.
void fullyComposeAffineMapAndOperands(AffineMap *map, void fullyComposeAffineMapAndOperands(AffineMap *map,
SmallVectorImpl<Value> *operands); SmallVectorImpl<Value> *operands);
} // namespace affine
} // namespace mlir } // namespace mlir
#include "mlir/Dialect/Affine/IR/AffineOpsDialect.h.inc" #include "mlir/Dialect/Affine/IR/AffineOpsDialect.h.inc"
#define GET_OP_CLASSES #define GET_OP_CLASSES
#include "mlir/Dialect/Affine/IR/AffineOps.h.inc" #include "mlir/Dialect/Affine/IR/AffineOps.h.inc"
namespace mlir { namespace mlir {
namespace affine {
/// Returns true if the provided value is the induction variable of an /// Returns true if the provided value is the induction variable of an
/// AffineForOp. /// AffineForOp.
bool isAffineForInductionVar(Value val); bool isAffineForInductionVar(Value val);
/// Returns true if `val` is the induction variable of an AffineParallelOp. /// Returns true if `val` is the induction variable of an AffineParallelOp.
bool isAffineParallelInductionVar(Value val); bool isAffineParallelInductionVar(Value val);
/// Returns true if the provided value is the induction variable of an /// Returns true if the provided value is the induction variable of an
▲ Show 20 Lines • Show All 76 Lines • ▼ Show 20 Linesprivate:
AffineMap map; AffineMap map;
AffineBound(AffineForOp op, unsigned opStart, unsigned opEnd, AffineMap map) AffineBound(AffineForOp op, unsigned opStart, unsigned opEnd, AffineMap map)
: op(op), opStart(opStart), opEnd(opEnd), map(map) {} : op(op), opStart(opStart), opEnd(opEnd), map(map) {}
friend class AffineForOp; friend class AffineForOp;
}; };
} // namespace affine
} // namespace mlir } // namespace mlir
#endif #endif

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

Show All 15 Lines
include "mlir/Dialect/Arith/IR/ArithBase.td" include "mlir/Dialect/Arith/IR/ArithBase.td"
include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.td" include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.td"
include "mlir/Interfaces/ControlFlowInterfaces.td" include "mlir/Interfaces/ControlFlowInterfaces.td"
include "mlir/Interfaces/LoopLikeInterface.td" include "mlir/Interfaces/LoopLikeInterface.td"
include "mlir/Interfaces/SideEffectInterfaces.td" include "mlir/Interfaces/SideEffectInterfaces.td"
def Affine_Dialect : Dialect { def Affine_Dialect : Dialect {
let name = "affine"; let name = "affine";
let cppNamespace = "mlir"; let cppNamespace = "::mlir::affine";
let hasConstantMaterializer = 1; let hasConstantMaterializer = 1;
let dependentDialects = ["arith::ArithDialect"]; let dependentDialects = ["arith::ArithDialect"];
} }
// Base class for Affine dialect ops. // Base class for Affine dialect ops.
class Affine_Op<string mnemonic, list<Trait> traits = []> : class Affine_Op<string mnemonic, list<Trait> traits = []> :
Op<Affine_Dialect, mnemonic, traits>; Op<Affine_Dialect, mnemonic, traits>;
▲ Show 20 Lines • Show All 1,100 Lines • Show Last 20 Lines

mlir/include/mlir/Dialect/Affine/IR/AffineValueMap.h

Show All 12 Lines
#ifndef MLIR_DIALECT_AFFINE_IR_AFFINEVALUEMAP_H #ifndef MLIR_DIALECT_AFFINE_IR_AFFINEVALUEMAP_H
#define MLIR_DIALECT_AFFINE_IR_AFFINEVALUEMAP_H #define MLIR_DIALECT_AFFINE_IR_AFFINEVALUEMAP_H
#include "mlir/IR/AffineMap.h" #include "mlir/IR/AffineMap.h"
#include "mlir/IR/OperationSupport.h" #include "mlir/IR/OperationSupport.h"
#include "mlir/IR/Value.h" #include "mlir/IR/Value.h"
namespace mlir { namespace mlir {
namespace affine {
/// An AffineValueMap is an affine map plus its ML value operands and /// An AffineValueMap is an affine map plus its ML value operands and
/// results for analysis purposes. The structure is still a tree form that is /// results for analysis purposes. The structure is still a tree form that is
/// same as that of an affine map or an AffineApplyOp. However, its operands, /// same as that of an affine map or an AffineApplyOp. However, its operands,
/// results, and its map can themselves change as a result of /// results, and its map can themselves change as a result of
/// substitutions, simplifications, and other analysis. /// substitutions, simplifications, and other analysis.
// An affine value map can readily be constructed from an AffineApplyOp, or an // An affine value map can readily be constructed from an AffineApplyOp, or an
// AffineBound of a AffineForOp. It can be further transformed, substituted // AffineBound of a AffineForOp. It can be further transformed, substituted
▲ Show 20 Lines • Show All 55 Lines • ▼ Show 20 Linesprivate:
// TODO: make these trailing objects? // TODO: make these trailing objects?
/// The SSA operands binding to the dim's and symbols of 'map'. /// The SSA operands binding to the dim's and symbols of 'map'.
SmallVector<Value, 4> operands; SmallVector<Value, 4> operands;
/// The SSA results binding to the results of 'map'. /// The SSA results binding to the results of 'map'.
SmallVector<Value, 4> results; SmallVector<Value, 4> results;
}; };
} // namespace affine
} // namespace mlir } // namespace mlir
#endif // MLIR_DIALECT_AFFINE_IR_AFFINEVALUEMAP_H #endif // MLIR_DIALECT_AFFINE_IR_AFFINEVALUEMAP_H

mlir/include/mlir/Dialect/Affine/LoopFusionUtils.h

Show All 15 Lines
#define MLIR_DIALECT_AFFINE_LOOPFUSIONUTILS_H #define MLIR_DIALECT_AFFINE_LOOPFUSIONUTILS_H
#include "mlir/IR/Value.h" #include "mlir/IR/Value.h"
#include "mlir/Support/LLVM.h" #include "mlir/Support/LLVM.h"
#include "llvm/ADT/DenseMap.h" #include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SmallVector.h" #include "llvm/ADT/SmallVector.h"
namespace mlir { namespace mlir {
class Operation;
namespace affine {
class AffineForOp; class AffineForOp;
struct ComputationSliceState; struct ComputationSliceState;
class Operation;
// TODO: Extend this module to include utility functions for querying fusion // TODO: Extend this module to include utility functions for querying fusion
// cost/storage reduction, and for performing the loop fusion transformation. // cost/storage reduction, and for performing the loop fusion transformation.
struct FusionResult { struct FusionResult {
enum ResultEnum { enum ResultEnum {
Success, Success,
FailPrecondition, // Failed precondition for fusion. (e.g. same block). FailPrecondition, // Failed precondition for fusion. (e.g. same block).
▲ Show 20 Lines • Show All 125 Lines • ▼ Show 20 Linesbool getFusionComputeCost(AffineForOp srcForOp, LoopNestStats &srcStats,
int64_t *computeCost); int64_t *computeCost);
/// Returns in 'producerConsumerMemrefs' the memrefs involved in a /// Returns in 'producerConsumerMemrefs' the memrefs involved in a
/// producer-consumer dependence between write ops in 'srcOps' and read ops in /// producer-consumer dependence between write ops in 'srcOps' and read ops in
/// 'dstOps'. /// 'dstOps'.
void gatherProducerConsumerMemrefs(ArrayRef<Operation *> srcOps, void gatherProducerConsumerMemrefs(ArrayRef<Operation *> srcOps,
ArrayRef<Operation *> dstOps, ArrayRef<Operation *> dstOps,
DenseSet<Value> &producerConsumerMemrefs); DenseSet<Value> &producerConsumerMemrefs);
} // namespace affine
} // namespace mlir } // namespace mlir
#endif // MLIR_DIALECT_AFFINE_LOOPFUSIONUTILS_H #endif // MLIR_DIALECT_AFFINE_LOOPFUSIONUTILS_H

mlir/include/mlir/Dialect/Affine/LoopUtils.h

Show All 16 Lines
#include "mlir/IR/Block.h" #include "mlir/IR/Block.h"
#include "mlir/Support/LLVM.h" #include "mlir/Support/LLVM.h"
#include "mlir/Support/LogicalResult.h" #include "mlir/Support/LogicalResult.h"
#include "mlir/Transforms/RegionUtils.h" #include "mlir/Transforms/RegionUtils.h"
#include <optional> #include <optional>
namespace mlir { namespace mlir {
class AffineForOp;
class AffineMap; class AffineMap;
class LoopLikeOpInterface; class LoopLikeOpInterface;
struct MemRefRegion;
class OpBuilder; class OpBuilder;
class Value; class Value;
class ValueRange; class ValueRange;
namespace func { namespace func {
class FuncOp; class FuncOp;
} // namespace func } // namespace func
namespace scf { namespace scf {
class ForOp; class ForOp;
class ParallelOp; class ParallelOp;
} // namespace scf } // namespace scf
namespace affine {
class AffineForOp;
struct MemRefRegion;
/// Unrolls this for operation completely if the trip count is known to be /// Unrolls this for operation completely if the trip count is known to be
/// constant. Returns failure otherwise. /// constant. Returns failure otherwise.
LogicalResult loopUnrollFull(AffineForOp forOp); LogicalResult loopUnrollFull(AffineForOp forOp);
/// Unrolls this for operation by the specified unroll factor. Returns failure /// Unrolls this for operation by the specified unroll factor. Returns failure
/// if the loop cannot be unrolled either due to restrictions or due to invalid /// if the loop cannot be unrolled either due to restrictions or due to invalid
/// unroll factors. Requires positive loop bounds and step. If specified, /// unroll factors. Requires positive loop bounds and step. If specified,
/// annotates the Ops in each unrolled iteration by applying `annotateFn`. /// annotates the Ops in each unrolled iteration by applying `annotateFn`.
▲ Show 20 Lines • Show All 290 Lines • ▼ Show 20 Lines for (unsigned end = loops.size(); end > 0; --end) {
// If a band was found and transformed, keep looking at the loops above // If a band was found and transformed, keep looking at the loops above
// the outermost transformed loop. // the outermost transformed loop.
if (start != end - 1) if (start != end - 1)
end = start + 1; end = start + 1;
} }
return result; return result;
} }
} // namespace affine
} // namespace mlir } // namespace mlir
#endif // MLIR_DIALECT_AFFINE_LOOPUTILS_H #endif // MLIR_DIALECT_AFFINE_LOOPUTILS_H

mlir/include/mlir/Dialect/Affine/Passes.h

Show All 17 Lines
#include <limits> #include <limits>
namespace mlir { namespace mlir {
namespace func { namespace func {
class FuncOp; class FuncOp;
} // namespace func } // namespace func
namespace affine {
class AffineForOp; class AffineForOp;
/// Fusion mode to attempt. The default mode `Greedy` does both /// Fusion mode to attempt. The default mode `Greedy` does both
/// producer-consumer and sibling fusion. /// producer-consumer and sibling fusion.
enum FusionMode { Greedy, ProducerConsumer, Sibling }; enum FusionMode { Greedy, ProducerConsumer, Sibling };
#define GEN_PASS_DECL #define GEN_PASS_DECL
#include "mlir/Dialect/Affine/Passes.h.inc" #include "mlir/Dialect/Affine/Passes.h.inc"
▲ Show 20 Lines • Show All 84 Lines • ▼ Show 20 Lines
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// Registration // Registration
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
/// Generate the code for registering passes. /// Generate the code for registering passes.
#define GEN_PASS_REGISTRATION #define GEN_PASS_REGISTRATION
#include "mlir/Dialect/Affine/Passes.h.inc" #include "mlir/Dialect/Affine/Passes.h.inc"
} // namespace affine
} // namespace mlir } // namespace mlir
#endif // MLIR_DIALECT_AFFINE_PASSES_H #endif // MLIR_DIALECT_AFFINE_PASSES_H

mlir/include/mlir/Dialect/Affine/Passes.td

Show All 11 Lines
#ifndef MLIR_DIALECT_AFFINE_PASSES #ifndef MLIR_DIALECT_AFFINE_PASSES
#define MLIR_DIALECT_AFFINE_PASSES #define MLIR_DIALECT_AFFINE_PASSES
include "mlir/Pass/PassBase.td" include "mlir/Pass/PassBase.td"
def AffineDataCopyGeneration : Pass<"affine-data-copy-generate", "func::FuncOp"> { def AffineDataCopyGeneration : Pass<"affine-data-copy-generate", "func::FuncOp"> {
let summary = "Generate explicit copying for affine memory operations"; let summary = "Generate explicit copying for affine memory operations";
let constructor = "mlir::createAffineDataCopyGenerationPass()"; let constructor = "mlir::affine::createAffineDataCopyGenerationPass()";
let dependentDialects = ["memref::MemRefDialect"]; let dependentDialects = ["memref::MemRefDialect"];
let options = [ let options = [
Option<"fastMemoryCapacity", "fast-mem-capacity", "uint64_t", Option<"fastMemoryCapacity", "fast-mem-capacity", "uint64_t",
/*default=*/"std::numeric_limits<uint64_t>::max()", /*default=*/"std::numeric_limits<uint64_t>::max()",
"Set fast memory space capacity in KiB (default: unlimited)">, "Set fast memory space capacity in KiB (default: unlimited)">,
Option<"fastMemorySpace", "fast-mem-space", "unsigned", Option<"fastMemorySpace", "fast-mem-space", "unsigned",
/*default=*/"1", /*default=*/"1",
"Fast memory space identifier for copy generation (default: 1)">, "Fast memory space identifier for copy generation (default: 1)">,
▲ Show 20 Lines • Show All 118 Lines • ▼ Show 20 Lines func.func @sibling_fusion(%arg0: memref<10x10xf32>, %arg1: memref<10x10xf32>,
%5 = arith.addf %3, %4 : f32 %5 = arith.addf %3, %4 : f32
affine.store %5, %arg4[%arg5, %arg6] : memref<10x10xf32> affine.store %5, %arg4[%arg5, %arg6] : memref<10x10xf32>
} }
} }
return return
} }
``` ```
}]; }];
let constructor = "mlir::createLoopFusionPass()"; let constructor = "mlir::affine::createLoopFusionPass()";
let options = [ let options = [
Option<"computeToleranceThreshold", "fusion-compute-tolerance", "double", Option<"computeToleranceThreshold", "fusion-compute-tolerance", "double",
/*default=*/"0.30f", "Fractional increase in additional computation " /*default=*/"0.30f", "Fractional increase in additional computation "
"tolerated while fusing">, "tolerated while fusing">,
Option<"fastMemorySpace", "fusion-fast-mem-space", "unsigned", Option<"fastMemorySpace", "fusion-fast-mem-space", "unsigned",
/*default=*/"0", /*default=*/"0",
"Faster memory space number to promote fusion buffers to">, "Faster memory space number to promote fusion buffers to">,
Option<"localBufSizeThreshold", "fusion-local-buf-threshold", "uint64_t", Option<"localBufSizeThreshold", "fusion-local-buf-threshold", "uint64_t",
/*default=*/"0", "Threshold size (KiB) for promoting local buffers " /*default=*/"0", "Threshold size (KiB) for promoting local buffers "
"to fast memory space">, "to fast memory space">,
Option<"maximalFusion", "fusion-maximal", "bool", /*default=*/"false", Option<"maximalFusion", "fusion-maximal", "bool", /*default=*/"false",
"Enables maximal loop fusion">, "Enables maximal loop fusion">,
Option<"affineFusionMode", "mode", "enum FusionMode", Option<"affineFusionMode", "mode", "enum FusionMode",
"mlir::FusionMode::Greedy", "fusion mode to attempt", "mlir::affine::FusionMode::Greedy", "fusion mode to attempt",
"llvm::cl::values(clEnumValN(mlir::FusionMode::Greedy," "llvm::cl::values(clEnumValN(mlir::affine::FusionMode::Greedy,"
" \"greedy\", \"Perform greedy (both producer-consumer and sibling) fusion\"), " " \"greedy\", \"Perform greedy (both producer-consumer and sibling) fusion\"), "
"clEnumValN( mlir::FusionMode::ProducerConsumer, " "clEnumValN( mlir::affine::FusionMode::ProducerConsumer, "
"\"producer\", \"Perform only producer-consumer fusion\"), " "\"producer\", \"Perform only producer-consumer fusion\"), "
"clEnumValN( mlir::FusionMode::Sibling, " "clEnumValN( mlir::affine::FusionMode::Sibling, "
"\"sibling\", \"Perform only sibling fusion\"))">, "\"sibling\", \"Perform only sibling fusion\"))">,
]; ];
let dependentDialects = ["memref::MemRefDialect"]; let dependentDialects = ["memref::MemRefDialect"];
} }
def AffineLoopInvariantCodeMotion def AffineLoopInvariantCodeMotion
: Pass<"affine-loop-invariant-code-motion", "func::FuncOp"> { : Pass<"affine-loop-invariant-code-motion", "func::FuncOp"> {
let summary = "Hoist loop invariant instructions outside of affine loops"; let summary = "Hoist loop invariant instructions outside of affine loops";
let constructor = "mlir::createAffineLoopInvariantCodeMotionPass()"; let constructor = "mlir::affine::createAffineLoopInvariantCodeMotionPass()";
} }
def AffineLoopTiling : Pass<"affine-loop-tile", "func::FuncOp"> { def AffineLoopTiling : Pass<"affine-loop-tile", "func::FuncOp"> {
let summary = "Tile affine loop nests"; let summary = "Tile affine loop nests";
let constructor = "mlir::createLoopTilingPass()"; let constructor = "mlir::affine::createLoopTilingPass()";
let options = [ let options = [
Option<"cacheSizeInKiB", "cache-size", "uint64_t", /*default=*/"512", Option<"cacheSizeInKiB", "cache-size", "uint64_t", /*default=*/"512",
"Set size of cache to tile for in KiB (default: 512)">, "Set size of cache to tile for in KiB (default: 512)">,
Option<"separate", "separate", "bool", /*default=*/"false", Option<"separate", "separate", "bool", /*default=*/"false",
"Separate full and partial tiles (default: false)">, "Separate full and partial tiles (default: false)">,
Option<"tileSize", "tile-size", "unsigned", /*default=*/"", Option<"tileSize", "tile-size", "unsigned", /*default=*/"",
"Use this tile size for all loops">, "Use this tile size for all loops">,
ListOption<"tileSizes", "tile-sizes", "unsigned", ListOption<"tileSizes", "tile-sizes", "unsigned",
"List of tile sizes for each perfect nest " "List of tile sizes for each perfect nest "
"(overridden by -tile-size)">, "(overridden by -tile-size)">,
]; ];
} }
def AffineLoopUnroll : Pass<"affine-loop-unroll", "func::FuncOp"> { def AffineLoopUnroll : Pass<"affine-loop-unroll", "func::FuncOp"> {
let summary = "Unroll affine loops"; let summary = "Unroll affine loops";
let constructor = "mlir::createLoopUnrollPass()"; let constructor = "mlir::affine::createLoopUnrollPass()";
let options = [ let options = [
Option<"unrollFactor", "unroll-factor", "unsigned", /*default=*/"4", Option<"unrollFactor", "unroll-factor", "unsigned", /*default=*/"4",
"Use this unroll factor for all loops being unrolled">, "Use this unroll factor for all loops being unrolled">,
Option<"unrollUpToFactor", "unroll-up-to-factor", "bool", Option<"unrollUpToFactor", "unroll-up-to-factor", "bool",
/*default=*/"false", "Allow unrolling up to the factor specified">, /*default=*/"false", "Allow unrolling up to the factor specified">,
Option<"unrollFull", "unroll-full", "bool", /*default=*/"false", Option<"unrollFull", "unroll-full", "bool", /*default=*/"false",
"Fully unroll loops">, "Fully unroll loops">,
Option<"numRepetitions", "unroll-num-reps", "unsigned", /*default=*/"1", Option<"numRepetitions", "unroll-num-reps", "unsigned", /*default=*/"1",
"Unroll innermost loops repeatedly this many times">, "Unroll innermost loops repeatedly this many times">,
Option<"unrollFullThreshold", "unroll-full-threshold", "unsigned", Option<"unrollFullThreshold", "unroll-full-threshold", "unsigned",
/*default=*/"1", /*default=*/"1",
"Unroll all loops with trip count less than or equal to this">, "Unroll all loops with trip count less than or equal to this">,
Option<"cleanUpUnroll", "cleanup-unroll", "bool", /*default=*/"false", Option<"cleanUpUnroll", "cleanup-unroll", "bool", /*default=*/"false",
"Fully unroll the cleanup loop when possible.">, "Fully unroll the cleanup loop when possible.">,
]; ];
} }
def AffineLoopUnrollAndJam : Pass<"affine-loop-unroll-jam", "func::FuncOp"> { def AffineLoopUnrollAndJam : Pass<"affine-loop-unroll-jam", "func::FuncOp"> {
let summary = "Unroll and jam affine loops"; let summary = "Unroll and jam affine loops";
let constructor = "mlir::createLoopUnrollAndJamPass()"; let constructor = "mlir::affine::createLoopUnrollAndJamPass()";
let options = [ let options = [
Option<"unrollJamFactor", "unroll-jam-factor", "unsigned", Option<"unrollJamFactor", "unroll-jam-factor", "unsigned",
/*default=*/"4", /*default=*/"4",
"Use this unroll jam factor for all loops (default 4)">, "Use this unroll jam factor for all loops (default 4)">,
]; ];
} }
def AffinePipelineDataTransfer def AffinePipelineDataTransfer
▲ Show 20 Lines • Show All 57 Lines • ▼ Show 20 Lines module {
affine.store %7, %1[%3 mod 2, %3] : memref<2x32xf32, 1> affine.store %7, %1[%3 mod 2, %3] : memref<2x32xf32, 1>
memref.dealloc %2 : memref<2x1xf32> memref.dealloc %2 : memref<2x1xf32>
memref.dealloc %1 : memref<2x32xf32, 1> memref.dealloc %1 : memref<2x32xf32, 1>
return return
} }
} }
``` ```
}]; }];
let constructor = "mlir::createPipelineDataTransferPass()"; let constructor = "mlir::affine::createPipelineDataTransferPass()";
} }
def AffineScalarReplacement : Pass<"affine-scalrep", "func::FuncOp"> { def AffineScalarReplacement : Pass<"affine-scalrep", "func::FuncOp"> {
let summary = "Replace affine memref accesses by scalars by forwarding stores " let summary = "Replace affine memref accesses by scalars by forwarding stores "
"to loads and eliminating redundant loads"; "to loads and eliminating redundant loads";
let description = [{ let description = [{
This pass performs store to load forwarding and redundant load elimination This pass performs store to load forwarding and redundant load elimination
for affine memref accesses and potentially eliminates the entire memref for affine memref accesses and potentially eliminates the entire memref
Show All 29 Lines module {
%1 = arith.addf %cst, %cst : f32 %1 = arith.addf %cst, %cst : f32
} }
} }
return %0 : memref<10x10xf32> return %0 : memref<10x10xf32>
} }
} }
``` ```
}]; }];
let constructor = "mlir::createAffineScalarReplacementPass()"; let constructor = "mlir::affine::createAffineScalarReplacementPass()";
} }
def AffineVectorize : Pass<"affine-super-vectorize", "func::FuncOp"> { def AffineVectorize : Pass<"affine-super-vectorize", "func::FuncOp"> {
let summary = "Vectorize to a target independent n-D vector abstraction"; let summary = "Vectorize to a target independent n-D vector abstraction";
let dependentDialects = ["vector::VectorDialect"]; let dependentDialects = ["vector::VectorDialect"];
let options = [ let options = [
ListOption<"vectorSizes", "virtual-vector-size", "int64_t", ListOption<"vectorSizes", "virtual-vector-size", "int64_t",
"Specify an n-D virtual vector size for vectorization">, "Specify an n-D virtual vector size for vectorization">,
Show All 11 Lines Option<"vectorizeReductions", "vectorize-reductions", "bool",
/*default=*/"false", /*default=*/"false",
"Vectorize known reductions expressed via iter_args. " "Vectorize known reductions expressed via iter_args. "
"Switched off by default."> "Switched off by default.">
]; ];
} }
def AffineParallelize : Pass<"affine-parallelize", "func::FuncOp"> { def AffineParallelize : Pass<"affine-parallelize", "func::FuncOp"> {
let summary = "Convert affine.for ops into 1-D affine.parallel"; let summary = "Convert affine.for ops into 1-D affine.parallel";
let constructor = "mlir::createAffineParallelizePass()"; let constructor = "mlir::affine::createAffineParallelizePass()";
let options = [ let options = [
Option<"maxNested", "max-nested", "unsigned", /*default=*/"-1u", Option<"maxNested", "max-nested", "unsigned", /*default=*/"-1u",
"Maximum number of nested parallel loops to produce. " "Maximum number of nested parallel loops to produce. "
"Defaults to unlimited (UINT_MAX).">, "Defaults to unlimited (UINT_MAX).">,
Option<"parallelReductions", "parallel-reductions", "bool", Option<"parallelReductions", "parallel-reductions", "bool",
/*default=*/"false", /*default=*/"false",
"Whether to parallelize reduction loops. Defaults to false."> "Whether to parallelize reduction loops. Defaults to false.">
]; ];
} }
def AffineLoopNormalize : Pass<"affine-loop-normalize", "func::FuncOp"> { def AffineLoopNormalize : Pass<"affine-loop-normalize", "func::FuncOp"> {
let summary = "Apply normalization transformations to affine loop-like ops"; let summary = "Apply normalization transformations to affine loop-like ops";
let constructor = "mlir::createAffineLoopNormalizePass()"; let constructor = "mlir::affine::createAffineLoopNormalizePass()";
let options = [ let options = [
Option<"promoteSingleIter", "promote-single-iter", "bool", Option<"promoteSingleIter", "promote-single-iter", "bool",
/*default=*/"true", "Promote single iteration loops">, /*default=*/"true", "Promote single iteration loops">,
]; ];
} }
def LoopCoalescing : Pass<"affine-loop-coalescing", "func::FuncOp"> { def LoopCoalescing : Pass<"affine-loop-coalescing", "func::FuncOp"> {
let summary = "Coalesce nested loops with independent bounds into a single " let summary = "Coalesce nested loops with independent bounds into a single "
"loop"; "loop";
let constructor = "mlir::createLoopCoalescingPass()"; let constructor = "mlir::affine::createLoopCoalescingPass()";
let dependentDialects = ["arith::ArithDialect"]; let dependentDialects = ["arith::ArithDialect"];
} }
def SimplifyAffineStructures : Pass<"affine-simplify-structures", "func::FuncOp"> { def SimplifyAffineStructures : Pass<"affine-simplify-structures", "func::FuncOp"> {
let summary = "Simplify affine expressions in maps/sets and normalize " let summary = "Simplify affine expressions in maps/sets and normalize "
"memrefs"; "memrefs";
let constructor = "mlir::createSimplifyAffineStructuresPass()"; let constructor = "mlir::affine::createSimplifyAffineStructuresPass()";
} }
def AffineExpandIndexOps : Pass<"affine-expand-index-ops"> { def AffineExpandIndexOps : Pass<"affine-expand-index-ops"> {
let summary = "Lower affine operations operating on indices into more fundamental operations"; let summary = "Lower affine operations operating on indices into more fundamental operations";
let constructor = "mlir::createAffineExpandIndexOpsPass()"; let constructor = "mlir::affine::createAffineExpandIndexOpsPass()";
} }
#endif // MLIR_DIALECT_AFFINE_PASSES #endif // MLIR_DIALECT_AFFINE_PASSES

mlir/include/mlir/Dialect/Affine/TransformOps/AffineTransformOps.h

Show All 9 Lines
#define MLIR_DIALECT_AFFINE_TRANSFORMOPS_AFFINETRANSFORMOPS_H #define MLIR_DIALECT_AFFINE_TRANSFORMOPS_AFFINETRANSFORMOPS_H
#include "mlir/Dialect/PDL/IR/PDLTypes.h" #include "mlir/Dialect/PDL/IR/PDLTypes.h"
#include "mlir/Dialect/Transform/IR/TransformInterfaces.h" #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
#include "mlir/Dialect/Transform/IR/TransformTypes.h" #include "mlir/Dialect/Transform/IR/TransformTypes.h"
#include "mlir/IR/OpImplementation.h" #include "mlir/IR/OpImplementation.h"
namespace mlir { namespace mlir {
class AffineForOp;
namespace func { namespace func {
class FuncOp; class FuncOp;
} // namespace func } // namespace func
namespace affine { namespace affine {
class ForOp; class AffineForOp;
} // namespace affine } // namespace affine
} // namespace mlir } // namespace mlir
#define GET_OP_CLASSES #define GET_OP_CLASSES
#include "mlir/Dialect/Affine/TransformOps/AffineTransformOps.h.inc" #include "mlir/Dialect/Affine/TransformOps/AffineTransformOps.h.inc"
namespace mlir { namespace mlir {
class DialectRegistry; class DialectRegistry;
namespace affine { namespace affine {
void registerTransformDialectExtension(DialectRegistry &registry); void registerTransformDialectExtension(DialectRegistry &registry);
} // namespace affine } // namespace affine
} // namespace mlir } // namespace mlir
#endif // MLIR_DIALECT_AFFINE_TRANSFORMOPS_AFFINETRANSFORMOPS_H #endif // MLIR_DIALECT_AFFINE_TRANSFORMOPS_AFFINETRANSFORMOPS_H

mlir/include/mlir/Dialect/Affine/Transforms/Transforms.h

Show All 12 Lines
#ifndef MLIR_DIALECT_AFFINE_TRANSFORMS_TRANSFORMS_H #ifndef MLIR_DIALECT_AFFINE_TRANSFORMS_TRANSFORMS_H
#define MLIR_DIALECT_AFFINE_TRANSFORMS_TRANSFORMS_H #define MLIR_DIALECT_AFFINE_TRANSFORMS_TRANSFORMS_H
#include "mlir/Interfaces/ValueBoundsOpInterface.h" #include "mlir/Interfaces/ValueBoundsOpInterface.h"
#include "mlir/Support/LogicalResult.h" #include "mlir/Support/LogicalResult.h"
namespace mlir { namespace mlir {
class AffineApplyOp;
class Location; class Location;
class OpBuilder; class OpBuilder;
class OpFoldResult; class OpFoldResult;
class RewritePatternSet; class RewritePatternSet;
class RewriterBase; class RewriterBase;
class Value; class Value;
namespace presburger { namespace presburger {
enum class BoundType; enum class BoundType;
} // namespace presburger } // namespace presburger
namespace affine {
class AffineApplyOp;
/// Populate patterns that expand affine index operations into more fundamental /// Populate patterns that expand affine index operations into more fundamental
/// operations (not necessarily restricted to Affine dialect). /// operations (not necessarily restricted to Affine dialect).
void populateAffineExpandIndexOpsPatterns(RewritePatternSet &patterns); void populateAffineExpandIndexOpsPatterns(RewritePatternSet &patterns);
/// Helper function to rewrite `op`'s affine map and reorder its operands such /// Helper function to rewrite `op`'s affine map and reorder its operands such
/// that they are in increasing order of hoistability (i.e. the least hoistable) /// that they are in increasing order of hoistability (i.e. the least hoistable)
/// operands come first in the operand list. /// operands come first in the operand list.
void reorderOperandsByHoistability(RewriterBase &rewriter, AffineApplyOp op); void reorderOperandsByHoistability(RewriterBase &rewriter, AffineApplyOp op);
Show All 37 Lines
/// By default, lower/equal bounds are closed and upper bounds are open. If /// By default, lower/equal bounds are closed and upper bounds are open. If
/// `closedUB` is set to "true", upper bounds are also closed. /// `closedUB` is set to "true", upper bounds are also closed.
FailureOr<OpFoldResult> reifyShapedValueDimBound( FailureOr<OpFoldResult> reifyShapedValueDimBound(
OpBuilder &b, Location loc, presburger::BoundType type, Value value, OpBuilder &b, Location loc, presburger::BoundType type, Value value,
int64_t dim, int64_t dim,
ValueBoundsConstraintSet::StopConditionFn stopCondition = nullptr, ValueBoundsConstraintSet::StopConditionFn stopCondition = nullptr,
bool closedUB = false); bool closedUB = false);
} // namespace affine
} // namespace mlir } // namespace mlir
#endif // MLIR_DIALECT_AFFINE_TRANSFORMS_TRANSFORMS_H #endif // MLIR_DIALECT_AFFINE_TRANSFORMS_TRANSFORMS_H

mlir/include/mlir/Dialect/Affine/Utils.h

Show All 12 Lines
#ifndef MLIR_DIALECT_AFFINE_UTILS_H #ifndef MLIR_DIALECT_AFFINE_UTILS_H
#define MLIR_DIALECT_AFFINE_UTILS_H #define MLIR_DIALECT_AFFINE_UTILS_H
#include "mlir/Dialect/Affine/Analysis/AffineAnalysis.h" #include "mlir/Dialect/Affine/Analysis/AffineAnalysis.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Affine/IR/AffineOps.h"
#include <optional> #include <optional>
namespace mlir { namespace mlir {
class AffineForOp;
class AffineIfOp;
class AffineParallelOp;
class DominanceInfo; class DominanceInfo;
class Operation; class Operation;
class PostDominanceInfo; class PostDominanceInfo;
namespace func { namespace func {
class FuncOp; class FuncOp;
} // namespace func } // namespace func
namespace memref { namespace memref {
class AllocOp; class AllocOp;
} // namespace memref } // namespace memref
struct LogicalResult; struct LogicalResult;
namespace affine {
class AffineForOp;
class AffineIfOp;
class AffineParallelOp;
using ReductionLoopMap = DenseMap<Operation *, SmallVector<LoopReduction, 2>>; using ReductionLoopMap = DenseMap<Operation *, SmallVector<LoopReduction, 2>>;
/// Replaces a parallel affine.for op with a 1-d affine.parallel op. `forOp`'s /// Replaces a parallel affine.for op with a 1-d affine.parallel op. `forOp`'s
/// body is taken by the affine.parallel op and the former is erased. /// body is taken by the affine.parallel op and the former is erased.
/// (mlir::isLoopParallel can be used to detect a parallel affine.for op.) The /// (mlir::isLoopParallel can be used to detect a parallel affine.for op.) The
/// reductions specified in `parallelReductions` are also parallelized. /// reductions specified in `parallelReductions` are also parallelized.
/// Parallelization will fail in the presence of loop iteration arguments that /// Parallelization will fail in the presence of loop iteration arguments that
/// are not listed in `parallelReductions`. /// are not listed in `parallelReductions`.
▲ Show 20 Lines • Show All 332 Lines • ▼ Show 20 Lines return makeComposedFoldedAffineMax(
vals); vals);
} }
private: private:
OpBuilder &b; OpBuilder &b;
Location loc; Location loc;
}; };
} // namespace affine
} // namespace mlir } // namespace mlir
#endif // MLIR_DIALECT_AFFINE_UTILS_H #endif // MLIR_DIALECT_AFFINE_UTILS_H

mlir/include/mlir/Dialect/Affine/ViewLikeInterfaceUtils.h

Show All 10 Lines
#include "mlir/Dialect/Utils/StaticValueUtils.h" #include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/OpDefinition.h" #include "mlir/IR/OpDefinition.h"
#include "mlir/Interfaces/ViewLikeInterface.h" #include "mlir/Interfaces/ViewLikeInterface.h"
namespace mlir { namespace mlir {
class RewriterBase; class RewriterBase;
namespace affine {
/// Fills the `combinedOffsets`, `combinedSizes` and `combinedStrides` to use /// Fills the `combinedOffsets`, `combinedSizes` and `combinedStrides` to use
/// when combining a producer slice **into** a consumer slice. /// when combining a producer slice **into** a consumer slice.
/// ///
/// This function performs the following computation: /// This function performs the following computation:
/// - Combined offsets = producer_offsets * consumer_strides + consumer_offsets /// - Combined offsets = producer_offsets * consumer_strides + consumer_offsets
/// - Combined sizes = consumer_sizes /// - Combined sizes = consumer_sizes
/// - Combined strides = producer_strides * consumer_strides /// - Combined strides = producer_strides * consumer_strides
// TODO: unify this API with resolveIndicesIntoOpWithOffsetsAndStrides or // TODO: unify this API with resolveIndicesIntoOpWithOffsetsAndStrides or
▲ Show 20 Lines • Show All 67 Lines • ▼ Show 20 Lines
/// In practice, this amounts to filtering by `rankReducedSourceDims` and taking /// In practice, this amounts to filtering by `rankReducedSourceDims` and taking
/// from `sourceSizes` if a dimension is dropped, otherwise taking from /// from `sourceSizes` if a dimension is dropped, otherwise taking from
/// `destSizes`. /// `destSizes`.
void resolveSizesIntoOpWithSizes( void resolveSizesIntoOpWithSizes(
ArrayRef<OpFoldResult> sourceSizes, ArrayRef<OpFoldResult> destSizes, ArrayRef<OpFoldResult> sourceSizes, ArrayRef<OpFoldResult> destSizes,
const llvm::SmallBitVector &rankReducedSourceDims, const llvm::SmallBitVector &rankReducedSourceDims,
SmallVectorImpl<OpFoldResult> &resolvedSizes); SmallVectorImpl<OpFoldResult> &resolvedSizes);
} // namespace affine
} // namespace mlir } // namespace mlir
#endif // MLIR_DIALECT_AFFINE_VIEWLIKEINTERFACEUTILS_H #endif // MLIR_DIALECT_AFFINE_VIEWLIKEINTERFACEUTILS_H

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationBase.td

Show All 20 Lines let description = [{
specific to the bufferization passes. specific to the bufferization passes.
Overview of the bufferization infrastructure and important conceptual Overview of the bufferization infrastructure and important conceptual
details related to using the MLIR dialect conversion infrastructure can be details related to using the MLIR dialect conversion infrastructure can be
found in [bufferization](/docs/Bufferization/) and [buffer found in [bufferization](/docs/Bufferization/) and [buffer
deallocation](/docs/BufferDeallocationInternals/). deallocation](/docs/BufferDeallocationInternals/).
}]; }];
let dependentDialects = [ let dependentDialects = [
"AffineDialect", "memref::MemRefDialect", "tensor::TensorDialect" "affine::AffineDialect", "memref::MemRefDialect", "tensor::TensorDialect"
]; ];
let extraClassDeclaration = [{ let extraClassDeclaration = [{
/// Verify an attribute from this dialect on the argument at 'argIndex' for /// Verify an attribute from this dialect on the argument at 'argIndex' for
/// the region at 'regionIndex' on the given operation. Returns failure if /// the region at 'regionIndex' on the given operation. Returns failure if
/// the verification failed, success otherwise. This hook may optionally be /// the verification failed, success otherwise. This hook may optionally be
/// invoked from any operation containing a region. /// invoked from any operation containing a region.
LogicalResult verifyRegionArgAttribute(Operation *, LogicalResult verifyRegionArgAttribute(Operation *,
Show All 37 Lines

mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td

Show All 31 Lines let description = [{
[Rationale [Rationale
Document](https://mlir.llvm.org/docs/Rationale/RationaleLinalgDialect) are Document](https://mlir.llvm.org/docs/Rationale/RationaleLinalgDialect) are
are also available and should be read first before going in the details of are also available and should be read first before going in the details of
the op semantics. the op semantics.
}]; }];
let cppNamespace = "::mlir::linalg"; let cppNamespace = "::mlir::linalg";
let dependentDialects = [ let dependentDialects = [
"arith::ArithDialect", "arith::ArithDialect",
"AffineDialect", "affine::AffineDialect",
"math::MathDialect", "math::MathDialect",
"memref::MemRefDialect", "memref::MemRefDialect",
"tensor::TensorDialect", "tensor::TensorDialect",
]; ];
let useDefaultAttributePrinterParser = 1; let useDefaultAttributePrinterParser = 1;
let hasCanonicalizer = 1; let hasCanonicalizer = 1;
let hasOperationAttrVerify = 1; let hasOperationAttrVerify = 1;
let hasConstantMaterializer = 1; let hasConstantMaterializer = 1;
Show All 33 Lines

mlir/include/mlir/Dialect/Linalg/Passes.td

Show All 31 Lines Option<"foldOneTripLoopsOnly", "fold-one-trip-loops-only", "bool",
/*default=*/"false", /*default=*/"false",
"Only folds the one-trip loops from Linalg ops on tensors " "Only folds the one-trip loops from Linalg ops on tensors "
"(for testing purposes only)">, "(for testing purposes only)">,
Option<"useRankReducingSlices", "use-rank-reducing-slices", "bool", Option<"useRankReducingSlices", "use-rank-reducing-slices", "bool",
/*default=*/"false", /*default=*/"false",
"Generate rank-reducing slices instead of reassociative reshapes"> "Generate rank-reducing slices instead of reassociative reshapes">
]; ];
let dependentDialects = [ let dependentDialects = [
"linalg::LinalgDialect", "AffineDialect", "memref::MemRefDialect" "linalg::LinalgDialect", "affine::AffineDialect", "memref::MemRefDialect"
]; ];
} }
def LinalgElementwiseOpFusion : Pass<"linalg-fuse-elementwise-ops"> { def LinalgElementwiseOpFusion : Pass<"linalg-fuse-elementwise-ops"> {
let summary = "Fuse elementwise operations on tensors"; let summary = "Fuse elementwise operations on tensors";
let constructor = "mlir::createLinalgElementwiseOpFusionPass()"; let constructor = "mlir::createLinalgElementwiseOpFusionPass()";
let dependentDialects = [ let dependentDialects = [
"AffineDialect", "linalg::LinalgDialect", "memref::MemRefDialect" "affine::AffineDialect", "linalg::LinalgDialect", "memref::MemRefDialect"
]; ];
} }
def LinalgNamedOpConversion: Pass<"linalg-named-op-conversion"> { def LinalgNamedOpConversion: Pass<"linalg-named-op-conversion"> {
let summary = "Convert from one named linalg op to another."; let summary = "Convert from one named linalg op to another.";
let constructor = "mlir::createLinalgNamedOpConversionPass()"; let constructor = "mlir::createLinalgNamedOpConversionPass()";
let dependentDialects = ["linalg::LinalgDialect", "tensor::TensorDialect"]; let dependentDialects = ["linalg::LinalgDialect", "tensor::TensorDialect"];
} }
def LinalgInlineScalarOperands : Pass<"linalg-inline-scalar-operands", "func::FuncOp"> { def LinalgInlineScalarOperands : Pass<"linalg-inline-scalar-operands", "func::FuncOp"> {
let summary = "Inline scalar operands into linalg generic ops"; let summary = "Inline scalar operands into linalg generic ops";
let constructor = "mlir::createLinalgInlineScalarOperandsPass()"; let constructor = "mlir::createLinalgInlineScalarOperandsPass()";
let dependentDialects = [ let dependentDialects = [
"linalg::LinalgDialect" "linalg::LinalgDialect"
]; ];
} }
def LinalgLowerToAffineLoops : Pass<"convert-linalg-to-affine-loops", "func::FuncOp"> { def LinalgLowerToAffineLoops : Pass<"convert-linalg-to-affine-loops", "func::FuncOp"> {
let summary = "Lower the operations from the linalg dialect into affine " let summary = "Lower the operations from the linalg dialect into affine "
"loops"; "loops";
let constructor = "mlir::createConvertLinalgToAffineLoopsPass()"; let constructor = "mlir::createConvertLinalgToAffineLoopsPass()";
let dependentDialects = [ let dependentDialects = [
"AffineDialect", "linalg::LinalgDialect", "memref::MemRefDialect"]; "affine::AffineDialect", "linalg::LinalgDialect", "memref::MemRefDialect"];
} }
def LinalgLowerToLoops : Pass<"convert-linalg-to-loops", "func::FuncOp"> { def LinalgLowerToLoops : Pass<"convert-linalg-to-loops", "func::FuncOp"> {
let summary = "Lower the operations from the linalg dialect into loops"; let summary = "Lower the operations from the linalg dialect into loops";
let constructor = "mlir::createConvertLinalgToLoopsPass()"; let constructor = "mlir::createConvertLinalgToLoopsPass()";
let dependentDialects = [ let dependentDialects = [
"linalg::LinalgDialect", "linalg::LinalgDialect",
"scf::SCFDialect", "scf::SCFDialect",
"AffineDialect" "affine::AffineDialect"
]; ];
} }
def LinalgLowerToParallelLoops def LinalgLowerToParallelLoops
: Pass<"convert-linalg-to-parallel-loops", "func::FuncOp"> { : Pass<"convert-linalg-to-parallel-loops", "func::FuncOp"> {
let summary = "Lower the operations from the linalg dialect into parallel " let summary = "Lower the operations from the linalg dialect into parallel "
"loops"; "loops";
let constructor = "mlir::createConvertLinalgToParallelLoopsPass()"; let constructor = "mlir::createConvertLinalgToParallelLoopsPass()";
let dependentDialects = [ let dependentDialects = [
"AffineDialect", "affine::AffineDialect",
"linalg::LinalgDialect", "linalg::LinalgDialect",
"memref::MemRefDialect", "memref::MemRefDialect",
"scf::SCFDialect" "scf::SCFDialect"
]; ];
} }
def LinalgBufferize : Pass<"linalg-bufferize", "func::FuncOp"> { def LinalgBufferize : Pass<"linalg-bufferize", "func::FuncOp"> {
let summary = "Bufferize the linalg dialect"; let summary = "Bufferize the linalg dialect";
let constructor = "mlir::createLinalgBufferizePass()"; let constructor = "mlir::createLinalgBufferizePass()";
let dependentDialects = [ let dependentDialects = [
"AffineDialect", "affine::AffineDialect",
"bufferization::BufferizationDialect", "bufferization::BufferizationDialect",
"linalg::LinalgDialect", "linalg::LinalgDialect",
"memref::MemRefDialect", "memref::MemRefDialect",
]; ];
} }
def LinalgGeneralization : Pass<"linalg-generalize-named-ops", "func::FuncOp"> { def LinalgGeneralization : Pass<"linalg-generalize-named-ops", "func::FuncOp"> {
let summary = "Convert named ops into generic ops"; let summary = "Convert named ops into generic ops";
▲ Show 20 Lines • Show All 43 Lines • Show Last 20 Lines

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

Show All 12 Lines
#include "mlir/Dialect/Linalg/Utils/IndexingUtils.h" #include "mlir/Dialect/Linalg/Utils/IndexingUtils.h"
#include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h" #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "llvm/ADT/StringSet.h" #include "llvm/ADT/StringSet.h"
#include <optional> #include <optional>
namespace mlir { namespace mlir {
class AffineExpr; class AffineExpr;
class AffineForOp;
class AffineMap; class AffineMap;
class PatternRewriter; class PatternRewriter;
namespace affine {
class AffineForOp;
} // namespace affine
namespace tensor { namespace tensor {
class ExtractSliceOp; class ExtractSliceOp;
} // namespace tensor } // namespace tensor
namespace linalg { namespace linalg {
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// Utilities for inferring various semantics properties of Linalg ops. // Utilities for inferring various semantics properties of Linalg ops.
▲ Show 20 Lines • Show All 444 Lines • Show Last 20 Lines

mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td

Show All 18 Lines
def FoldMemRefAliasOps : Pass<"fold-memref-alias-ops"> { def FoldMemRefAliasOps : Pass<"fold-memref-alias-ops"> {
let summary = "Fold memref alias ops into consumer load/store ops"; let summary = "Fold memref alias ops into consumer load/store ops";
let description = [{ let description = [{
The pass folds loading/storing from/to memref aliasing ops to loading/storing The pass folds loading/storing from/to memref aliasing ops to loading/storing
from/to the original memref. from/to the original memref.
}]; }];
let constructor = "mlir::memref::createFoldMemRefAliasOpsPass()"; let constructor = "mlir::memref::createFoldMemRefAliasOpsPass()";
let dependentDialects = [ let dependentDialects = [
"AffineDialect", "memref::MemRefDialect", "vector::VectorDialect" "affine::AffineDialect", "memref::MemRefDialect", "vector::VectorDialect"
]; ];
} }
def MemRefEmulateWideInt : Pass<"memref-emulate-wide-int"> { def MemRefEmulateWideInt : Pass<"memref-emulate-wide-int"> {
let summary = "Emulate 2*N-bit integer operations using N-bit operations"; let summary = "Emulate 2*N-bit integer operations using N-bit operations";
let description = [{ let description = [{
Emulate memref integer operations that use too wide integer types with Emulate memref integer operations that use too wide integer types with
equivalent operations on supported narrow integer types. This is done by equivalent operations on supported narrow integer types. This is done by
▲ Show 20 Lines • Show All 115 Lines • ▼ Show 20 Lines affine.for %arg3 = %c0 to %c8 {
} }
} }
} }
return return
} }
``` ```
}]; }];
let constructor = "mlir::memref::createNormalizeMemRefsPass()"; let constructor = "mlir::memref::createNormalizeMemRefsPass()";
let dependentDialects = ["AffineDialect"]; let dependentDialects = ["affine::AffineDialect"];
} }
def ResolveRankedShapeTypeResultDims : def ResolveRankedShapeTypeResultDims :
Pass<"resolve-ranked-shaped-type-result-dims"> { Pass<"resolve-ranked-shaped-type-result-dims"> {
let summary = "Resolve memref.dim of result values of ranked shape type"; let summary = "Resolve memref.dim of result values of ranked shape type";
let description = [{ let description = [{
The pass resolves memref.dim of result of operations that The pass resolves memref.dim of result of operations that
implement the `ReifyRankedShapedTypeOpInterface` in terms of implement the `ReifyRankedShapedTypeOpInterface` in terms of
Show All 11 Linesdef ResolveShapedTypeResultDims : Pass<"resolve-shaped-type-result-dims"> {
let description = [{ let description = [{
The pass resolves memref.dim of result of operations that The pass resolves memref.dim of result of operations that
implement the `InferShapedTypeOpInterface` or implement the `InferShapedTypeOpInterface` or
`ReifyRankedShapedTypeOpInterface` in terms of shapes of its `ReifyRankedShapedTypeOpInterface` in terms of shapes of its
operands. operands.
}]; }];
let constructor = "mlir::memref::createResolveShapedTypeResultDimsPass()"; let constructor = "mlir::memref::createResolveShapedTypeResultDimsPass()";
let dependentDialects = [ let dependentDialects = [
"AffineDialect", "memref::MemRefDialect", "tensor::TensorDialect" "affine::AffineDialect", "memref::MemRefDialect", "tensor::TensorDialect"
]; ];
} }
def ExpandStridedMetadata : Pass<"expand-strided-metadata"> { def ExpandStridedMetadata : Pass<"expand-strided-metadata"> {
let summary = "Expand memref operations into easier to analyze constructs"; let summary = "Expand memref operations into easier to analyze constructs";
let description = [{ let description = [{
The pass expands memref operations that modify the metadata of a memref The pass expands memref operations that modify the metadata of a memref
(sizes, offset, strides) into a sequence of easier to analyze constructs. (sizes, offset, strides) into a sequence of easier to analyze constructs.
In particular, this pass transforms operations into explicit sequence of In particular, this pass transforms operations into explicit sequence of
operations that model the effect of this operation on the different metadata. operations that model the effect of this operation on the different metadata.
This pass uses affine constructs to materialize these effects. This pass uses affine constructs to materialize these effects.
}]; }];
let constructor = "mlir::memref::createExpandStridedMetadataPass()"; let constructor = "mlir::memref::createExpandStridedMetadataPass()";
let dependentDialects = [ let dependentDialects = [
"AffineDialect", "memref::MemRefDialect" "affine::AffineDialect", "memref::MemRefDialect"
]; ];
} }
#endif // MLIR_DIALECT_MEMREF_TRANSFORMS_PASSES #endif // MLIR_DIALECT_MEMREF_TRANSFORMS_PASSES

mlir/include/mlir/Dialect/SCF/Transforms/Passes.td

Show All 18 Lines
} }
// Note: Making these canonicalization patterns would require a dependency // Note: Making these canonicalization patterns would require a dependency
// of the SCF dialect on the Affine/Tensor/MemRef dialects or vice versa. // of the SCF dialect on the Affine/Tensor/MemRef dialects or vice versa.
def SCFForLoopCanonicalization def SCFForLoopCanonicalization
: Pass<"scf-for-loop-canonicalization"> { : Pass<"scf-for-loop-canonicalization"> {
let summary = "Canonicalize operations within scf.for loop bodies"; let summary = "Canonicalize operations within scf.for loop bodies";
let constructor = "mlir::createSCFForLoopCanonicalizationPass()"; let constructor = "mlir::createSCFForLoopCanonicalizationPass()";
let dependentDialects = ["AffineDialect", "tensor::TensorDialect", let dependentDialects = ["affine::AffineDialect", "tensor::TensorDialect",
"memref::MemRefDialect"]; "memref::MemRefDialect"];
} }
def SCFForLoopPeeling : Pass<"scf-for-loop-peeling"> { def SCFForLoopPeeling : Pass<"scf-for-loop-peeling"> {
let summary = "Peel `for` loops at their upper bounds."; let summary = "Peel `for` loops at their upper bounds.";
let constructor = "mlir::createForLoopPeelingPass()"; let constructor = "mlir::createForLoopPeelingPass()";
let options = [ let options = [
Option<"skipPartial", "skip-partial", "bool", Option<"skipPartial", "skip-partial", "bool",
/*default=*/"true", /*default=*/"true",
"Do not peel loops inside of the last, partial iteration of another " "Do not peel loops inside of the last, partial iteration of another "
"already peeled loop."> "already peeled loop.">
]; ];
let dependentDialects = ["AffineDialect"]; let dependentDialects = ["affine::AffineDialect"];
} }
def SCFForLoopSpecialization : Pass<"scf-for-loop-specialization"> { def SCFForLoopSpecialization : Pass<"scf-for-loop-specialization"> {
let summary = "Specialize `for` loops for vectorization"; let summary = "Specialize `for` loops for vectorization";
let constructor = "mlir::createForLoopSpecializationPass()"; let constructor = "mlir::createForLoopSpecializationPass()";
} }
def SCFParallelLoopFusion : Pass<"scf-parallel-loop-fusion"> { def SCFParallelLoopFusion : Pass<"scf-parallel-loop-fusion"> {
▲ Show 20 Lines • Show All 55 Lines • ▼ Show 20 Linesdef SCFParallelLoopTiling : Pass<"scf-parallel-loop-tiling"> {
let options = [ let options = [
ListOption<"tileSizes", "parallel-loop-tile-sizes", "int64_t", ListOption<"tileSizes", "parallel-loop-tile-sizes", "int64_t",
"Factors to tile parallel loops by">, "Factors to tile parallel loops by">,
Option<"noMinMaxBounds", "no-min-max-bounds", "bool", Option<"noMinMaxBounds", "no-min-max-bounds", "bool",
/*default=*/"false", /*default=*/"false",
"Perform tiling with fixed upper bound with inbound check " "Perform tiling with fixed upper bound with inbound check "
"inside the internal loops"> "inside the internal loops">
]; ];
let dependentDialects = ["AffineDialect"]; let dependentDialects = ["affine::AffineDialect"];
} }
def SCFForLoopRangeFolding : Pass<"scf-for-loop-range-folding"> { def SCFForLoopRangeFolding : Pass<"scf-for-loop-range-folding"> {
let summary = "Fold add/mul ops into loop range"; let summary = "Fold add/mul ops into loop range";
let constructor = "mlir::createForLoopRangeFoldingPass()"; let constructor = "mlir::createForLoopRangeFoldingPass()";
} }
def SCFForToWhileLoop : Pass<"scf-for-to-while"> { def SCFForToWhileLoop : Pass<"scf-for-to-while"> {
Show All 34 Lines

mlir/include/mlir/Dialect/SCF/Utils/AffineCanonicalizationUtils.h

Show All 12 Lines
#ifndef MLIR_DIALECT_SCF_UTILS_AFFINECANONICALIZATIONUTILS_H_ #ifndef MLIR_DIALECT_SCF_UTILS_AFFINECANONICALIZATIONUTILS_H_
#define MLIR_DIALECT_SCF_UTILS_AFFINECANONICALIZATIONUTILS_H_ #define MLIR_DIALECT_SCF_UTILS_AFFINECANONICALIZATIONUTILS_H_
#include "mlir/Support/LLVM.h" #include "mlir/Support/LLVM.h"
#include "mlir/Support/LogicalResult.h" #include "mlir/Support/LogicalResult.h"
namespace mlir { namespace mlir {
class AffineApplyOp;
class AffineMap; class AffineMap;
class FlatAffineValueConstraints;
struct LogicalResult; struct LogicalResult;
class Operation; class Operation;
class OpFoldResult; class OpFoldResult;
class RewriterBase; class RewriterBase;
class Value; class Value;
class ValueRange; class ValueRange;
namespace affine {
class FlatAffineValueConstraints;
} // namespace affine
namespace scf { namespace scf {
class IfOp; class IfOp;
/// Match "for loop"-like operations: If the first parameter is an iteration /// Match "for loop"-like operations: If the first parameter is an iteration
/// variable, return lower/upper bounds via the second/third parameter and the /// variable, return lower/upper bounds via the second/third parameter and the
/// step size via the last parameter. The function should return `success` in /// step size via the last parameter. The function should return `success` in
/// that case. If the first parameter is not an iteration variable, return /// that case. If the first parameter is not an iteration variable, return
/// `failure`. /// `failure`.
using LoopMatcherFn = function_ref<LogicalResult( using LoopMatcherFn = function_ref<LogicalResult(
Value, OpFoldResult &, OpFoldResult &, OpFoldResult &)>; Value, OpFoldResult &, OpFoldResult &, OpFoldResult &)>;
/// Match "for loop"-like operations from the SCF dialect. /// Match "for loop"-like operations from the SCF dialect.
LogicalResult matchForLikeLoop(Value iv, OpFoldResult &lb, OpFoldResult &ub, LogicalResult matchForLikeLoop(Value iv, OpFoldResult &lb, OpFoldResult &ub,
OpFoldResult &step); OpFoldResult &step);
/// Populate the given constraint set with induction variable constraints of a /// Populate the given constraint set with induction variable constraints of a
/// "for" loop with the given range and step. /// "for" loop with the given range and step.
LogicalResult addLoopRangeConstraints(FlatAffineValueConstraints &cstr, LogicalResult addLoopRangeConstraints(affine::FlatAffineValueConstraints &cstr,
Value iv, OpFoldResult lb, Value iv, OpFoldResult lb,
OpFoldResult ub, OpFoldResult step); OpFoldResult ub, OpFoldResult step);
/// Try to canonicalize the given affine.min/max operation in the context of /// Try to canonicalize the given affine.min/max operation in the context of
/// for `loops` with a known range. /// for `loops` with a known range.
/// ///
/// `loopMatcher` is used to retrieve loop bounds and the step size for a given /// `loopMatcher` is used to retrieve loop bounds and the step size for a given
/// iteration variable. /// iteration variable.
Show All 27 Lines

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td

Show First 20 Lines • Show All 64 Lines • ▼ Show 20 Lines func.func @kernel_matvec(%arga: tensor<?x?xf64, #SparseMatrix>,
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<?xf64> } -> tensor<?xf64>
return %0 : tensor<?xf64> return %0 : tensor<?xf64>
} }
``` ```
}]; }];
let constructor = "mlir::createSparsificationPass()"; let constructor = "mlir::createSparsificationPass()";
let dependentDialects = [ let dependentDialects = [
"AffineDialect", "affine::AffineDialect",
"arith::ArithDialect", "arith::ArithDialect",
"bufferization::BufferizationDialect", "bufferization::BufferizationDialect",
"LLVM::LLVMDialect", "LLVM::LLVMDialect",
"linalg::LinalgDialect", "linalg::LinalgDialect",
"memref::MemRefDialect", "memref::MemRefDialect",
"scf::SCFDialect", "scf::SCFDialect",
"sparse_tensor::SparseTensorDialect", "sparse_tensor::SparseTensorDialect",
]; ];
▲ Show 20 Lines • Show All 276 Lines • Show Last 20 Lines

mlir/include/mlir/Dialect/Tensor/IR/TensorBase.td

Show First 20 Lines • Show All 42 Lines • ▼ Show 20 Lines let description = [{
The `tensor` type is actually a builtin type (it lives in the builtin The `tensor` type is actually a builtin type (it lives in the builtin
dialect), and does not live in this dialect. dialect), and does not live in this dialect.
}]; }];
let hasCanonicalizer = 1; let hasCanonicalizer = 1;
let hasConstantMaterializer = 1; let hasConstantMaterializer = 1;
let dependentDialects = [ let dependentDialects = [
"AffineDialect", "affine::AffineDialect",
"arith::ArithDialect", "arith::ArithDialect",
"complex::ComplexDialect", "complex::ComplexDialect",
]; ];
} }
#endif // TENSOR_BASE #endif // TENSOR_BASE

mlir/include/mlir/Dialect/Tensor/Transforms/Passes.td

Show All 17 Lines let description = [{
At the moment, the following foldings occur when possible: At the moment, the following foldings occur when possible:
- tensor.extract_slice into vector.transfer_read - tensor.extract_slice into vector.transfer_read
- vector.transfer_write into tensor.insert_slice - vector.transfer_write into tensor.insert_slice
}]; }];
let constructor = "mlir::tensor::createFoldTensorSubsetOpsPass()"; let constructor = "mlir::tensor::createFoldTensorSubsetOpsPass()";
let dependentDialects = [ let dependentDialects = [
"AffineDialect", "tensor::TensorDialect", "vector::VectorDialect" "affine::AffineDialect", "tensor::TensorDialect", "vector::VectorDialect"
]; ];
} }
def TensorBufferize : Pass<"tensor-bufferize", "func::FuncOp"> { def TensorBufferize : Pass<"tensor-bufferize", "func::FuncOp"> {
let summary = "Bufferize the `tensor` dialect"; let summary = "Bufferize the `tensor` dialect";
let constructor = "mlir::tensor::createTensorBufferizePass()"; let constructor = "mlir::tensor::createTensorBufferizePass()";
} }
#endif // MLIR_DIALECT_TENSOR_TRANSFORMS_PASSES #endif // MLIR_DIALECT_TENSOR_TRANSFORMS_PASSES

mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h

Show All 12 Lines
#include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/BuiltinAttributes.h"
#include "mlir/Support/LLVM.h" #include "mlir/Support/LLVM.h"
#include "llvm/ADT/DenseMap.h" #include "llvm/ADT/DenseMap.h"
namespace mlir { namespace mlir {
// Forward declarations. // Forward declarations.
class AffineApplyOp;
class AffineForOp;
class AffineMap; class AffineMap;
class Block; class Block;
class Location; class Location;
class OpBuilder; class OpBuilder;
class Operation; class Operation;
class ShapedType; class ShapedType;
class Value; class Value;
class VectorType; class VectorType;
class VectorTransferOpInterface; class VectorTransferOpInterface;
namespace affine {
class AffineApplyOp;
class AffineForOp;
} // namespace affine
namespace vector { namespace vector {
/// Helper function that creates a memref::DimOp or tensor::DimOp depending on /// Helper function that creates a memref::DimOp or tensor::DimOp depending on
/// the type of `source`. /// the type of `source`.
Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim); Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim);
} // namespace vector } // namespace vector
/// Constructs a permutation map of invariant memref indices to vector /// Constructs a permutation map of invariant memref indices to vector
/// dimension. /// dimension.
▲ Show 20 Lines • Show All 88 Lines • Show Last 20 Lines

mlir/include/mlir/InitAllDialects.h

Show First 20 Lines • Show All 80 Lines • ▼ Show 20 Lines
#include "mlir/IR/Dialect.h" #include "mlir/IR/Dialect.h"
namespace mlir { namespace mlir {
/// Add all the MLIR dialects to the provided registry. /// Add all the MLIR dialects to the provided registry.
inline void registerAllDialects(DialectRegistry &registry) { inline void registerAllDialects(DialectRegistry &registry) {
// clang-format off // clang-format off
registry.insert<acc::OpenACCDialect, registry.insert<acc::OpenACCDialect,
AffineDialect, affine::AffineDialect,
arith::ArithDialect, arith::ArithDialect,
amdgpu::AMDGPUDialect, amdgpu::AMDGPUDialect,
amx::AMXDialect, amx::AMXDialect,
arm_neon::ArmNeonDialect, arm_neon::ArmNeonDialect,
async::AsyncDialect, async::AsyncDialect,
bufferization::BufferizationDialect, bufferization::BufferizationDialect,
cf::ControlFlowDialect, cf::ControlFlowDialect,
complex::ComplexDialect, complex::ComplexDialect,
▲ Show 20 Lines • Show All 72 Lines • Show Last 20 Lines

mlir/include/mlir/InitAllPasses.h

Show First 20 Lines • Show All 49 Lines • ▼ Show 20 Lines
inline void registerAllPasses() { inline void registerAllPasses() {
// General passes // General passes
registerTransformsPasses(); registerTransformsPasses();
// Conversion passes // Conversion passes
registerConversionPasses(); registerConversionPasses();
// Dialect passes // Dialect passes
registerAffinePasses(); affine::registerAffinePasses();
registerAsyncPasses(); registerAsyncPasses();
arith::registerArithPasses(); arith::registerArithPasses();
bufferization::registerBufferizationPasses(); bufferization::registerBufferizationPasses();
func::registerFuncPasses(); func::registerFuncPasses();
registerGPUPasses(); registerGPUPasses();
registerGpuSerializeToCubinPass(); registerGpuSerializeToCubinPass();
registerGpuSerializeToHsacoPass(); registerGpuSerializeToHsacoPass();
registerLinalgPasses(); registerLinalgPasses();
Show All 19 Lines

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

Show All 24 Lines
#include "mlir/Transforms/Passes.h" #include "mlir/Transforms/Passes.h"
namespace mlir { namespace mlir {
#define GEN_PASS_DEF_CONVERTAFFINETOSTANDARD #define GEN_PASS_DEF_CONVERTAFFINETOSTANDARD
#include "mlir/Conversion/Passes.h.inc" #include "mlir/Conversion/Passes.h.inc"
} // namespace mlir } // namespace mlir
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
using namespace mlir::vector; using namespace mlir::vector;
/// Given a range of values, emit the code that reduces them with "min" or "max" /// Given a range of values, emit the code that reduces them with "min" or "max"
/// depending on the provided comparison predicate. The predicate defines which /// depending on the provided comparison predicate. The predicate defines which
/// comparison to perform, "lt" for "min", "gt" for "max" and is used for the /// comparison to perform, "lt" for "min", "gt" for "max" and is used for the
/// `cmpi` operation followed by the `select` operation: /// `cmpi` operation followed by the `select` operation:
/// ///
/// %cond = arith.cmpi "predicate" %v0, %v1 /// %cond = arith.cmpi "predicate" %v0, %v1
▲ Show 20 Lines • Show All 530 Lines • Show Last 20 Lines

mlir/lib/Conversion/LinalgToStandard/LinalgToStandard.cpp

Show First 20 Lines • Show All 135 Lines • ▼ Show 20 Linesstruct ConvertLinalgToStandardPass
: public impl::ConvertLinalgToStandardBase<ConvertLinalgToStandardPass> { : public impl::ConvertLinalgToStandardBase<ConvertLinalgToStandardPass> {
void runOnOperation() override; void runOnOperation() override;
}; };
} // namespace } // namespace
void ConvertLinalgToStandardPass::runOnOperation() { void ConvertLinalgToStandardPass::runOnOperation() {
auto module = getOperation(); auto module = getOperation();
ConversionTarget target(getContext()); ConversionTarget target(getContext());
target.addLegalDialect<AffineDialect, arith::ArithDialect, func::FuncDialect, target.addLegalDialect<affine::AffineDialect, arith::ArithDialect,
memref::MemRefDialect, scf::SCFDialect>(); func::FuncDialect, memref::MemRefDialect,
scf::SCFDialect>();
target.addLegalOp<ModuleOp, func::FuncOp, func::ReturnOp>(); target.addLegalOp<ModuleOp, func::FuncOp, func::ReturnOp>();
RewritePatternSet patterns(&getContext()); RewritePatternSet patterns(&getContext());
populateLinalgToStandardConversionPatterns(patterns); populateLinalgToStandardConversionPatterns(patterns);
if (failed(applyFullConversion(module, target, std::move(patterns)))) if (failed(applyFullConversion(module, target, std::move(patterns))))
signalPassFailure(); signalPassFailure();
} }
std::unique_ptr<OperationPass<ModuleOp>> std::unique_ptr<OperationPass<ModuleOp>>
mlir::createConvertLinalgToStandardPass() { mlir::createConvertLinalgToStandardPass() {
return std::make_unique<ConvertLinalgToStandardPass>(); return std::make_unique<ConvertLinalgToStandardPass>();
} }

mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp

Show All 30 Lines
#include "mlir/Transforms/RegionUtils.h" #include "mlir/Transforms/RegionUtils.h"
#include "llvm/ADT/Sequence.h" #include "llvm/ADT/Sequence.h"
#include "llvm/Support/Debug.h" #include "llvm/Support/Debug.h"
#include <optional> #include <optional>
#define DEBUG_TYPE "loops-to-gpu" #define DEBUG_TYPE "loops-to-gpu"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
using namespace mlir::scf; using namespace mlir::scf;
// Name of internal attribute to mark visited operations during conversion. // Name of internal attribute to mark visited operations during conversion.
// //
// NOTE: The conversion originally used the following legality criteria: // NOTE: The conversion originally used the following legality criteria:
// `!parallelOp->hasAttr(gpu::getMappingAttrName())` // `!parallelOp->hasAttr(gpu::getMappingAttrName())`
// But the provided pattern might reject some cases based on more detailed // But the provided pattern might reject some cases based on more detailed
// analysis of the `mapping` attribute. // analysis of the `mapping` attribute.
▲ Show 20 Lines • Show All 650 Lines • Show Last 20 Lines

mlir/lib/Conversion/SCFToGPU/SCFToGPUPass.cpp

Show All 36 Linesstruct ForLoopMapper : public impl::ConvertAffineForToGPUBase<ForLoopMapper> {
ForLoopMapper(unsigned numBlockDims, unsigned numThreadDims) { ForLoopMapper(unsigned numBlockDims, unsigned numThreadDims) {
this->numBlockDims = numBlockDims; this->numBlockDims = numBlockDims;
this->numThreadDims = numThreadDims; this->numThreadDims = numThreadDims;
} }
void runOnOperation() override { void runOnOperation() override {
for (Operation &op : llvm::make_early_inc_range( for (Operation &op : llvm::make_early_inc_range(
getOperation().getFunctionBody().getOps())) { getOperation().getFunctionBody().getOps())) {
if (auto forOp = dyn_cast<AffineForOp>(&op)) { if (auto forOp = dyn_cast<affine::AffineForOp>(&op)) {
if (failed(convertAffineLoopNestToGPULaunch(forOp, numBlockDims, if (failed(convertAffineLoopNestToGPULaunch(forOp, numBlockDims,
numThreadDims))) numThreadDims)))
signalPassFailure(); signalPassFailure();
} }
} }
} }
}; };
Show All 29 Lines

mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp

Show First 20 Lines • Show All 60 Lines • ▼ Show 20 Linesstatic void getXferIndices(RewriterBase &rewriter, TransferOpType xferOp,
Location loc = xferOp.getLoc(); Location loc = xferOp.getLoc();
unsigned offsetsIdx = 0; unsigned offsetsIdx = 0;
for (auto expr : xferOp.getPermutationMap().getResults()) { for (auto expr : xferOp.getPermutationMap().getResults()) {
if (auto dim = expr.template dyn_cast<AffineDimExpr>()) { if (auto dim = expr.template dyn_cast<AffineDimExpr>()) {
Value prevIdx = indices[dim.getPosition()]; Value prevIdx = indices[dim.getPosition()];
SmallVector<Value, 3> dims(dimValues.begin(), dimValues.end()); SmallVector<Value, 3> dims(dimValues.begin(), dimValues.end());
dims.push_back(prevIdx); dims.push_back(prevIdx);
AffineExpr d0 = rewriter.getAffineDimExpr(offsetMap.getNumDims()); AffineExpr d0 = rewriter.getAffineDimExpr(offsetMap.getNumDims());
indices[dim.getPosition()] = makeComposedAffineApply( indices[dim.getPosition()] = affine::makeComposedAffineApply(
rewriter, loc, d0 + offsetMap.getResult(offsetsIdx++), dims); rewriter, loc, d0 + offsetMap.getResult(offsetsIdx++), dims);
continue; continue;
} }
} }
} }
// Return true if the contract op can be convert to MMA matmul. // Return true if the contract op can be convert to MMA matmul.
static bool contractSupportsMMAMatrixType(vector::ContractionOp contract, static bool contractSupportsMMAMatrixType(vector::ContractionOp contract,
▲ Show 20 Lines • Show All 1,245 Lines • Show Last 20 Lines

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

Show First 20 Lines • Show All 96 Lines • ▼ Show 20 Linesstatic void getXferIndices(OpBuilder &b, OpTy xferOp, Value iv,
indices.append(prevIndices.begin(), prevIndices.end()); indices.append(prevIndices.begin(), prevIndices.end());
Location loc = xferOp.getLoc(); Location loc = xferOp.getLoc();
bool isBroadcast = !dim.has_value(); bool isBroadcast = !dim.has_value();
if (!isBroadcast) { if (!isBroadcast) {
AffineExpr d0, d1; AffineExpr d0, d1;
bindDims(xferOp.getContext(), d0, d1); bindDims(xferOp.getContext(), d0, d1);
Value offset = adaptor.getIndices()[*dim]; Value offset = adaptor.getIndices()[*dim];
indices[*dim] = makeComposedAffineApply(b, loc, d0 + d1, {offset, iv}); indices[*dim] =
affine::makeComposedAffineApply(b, loc, d0 + d1, {offset, iv});
} }
} }
static void maybeYieldValue(OpBuilder &b, Location loc, bool hasRetVal, static void maybeYieldValue(OpBuilder &b, Location loc, bool hasRetVal,
Value value) { Value value) {
if (hasRetVal) { if (hasRetVal) {
assert(value && "Expected non-empty value"); assert(value && "Expected non-empty value");
b.create<scf::YieldOp>(loc, value); b.create<scf::YieldOp>(loc, value);
▲ Show 20 Lines • Show All 59 Lines • ▼ Show 20 Linesstatic Value generateInBoundsCheck(
Location loc = xferOp.getLoc(); Location loc = xferOp.getLoc();
ImplicitLocOpBuilder lb(xferOp.getLoc(), b); ImplicitLocOpBuilder lb(xferOp.getLoc(), b);
if (!xferOp.isDimInBounds(0) && !isBroadcast) { if (!xferOp.isDimInBounds(0) && !isBroadcast) {
Value memrefDim = Value memrefDim =
vector::createOrFoldDimOp(b, loc, xferOp.getSource(), *dim); vector::createOrFoldDimOp(b, loc, xferOp.getSource(), *dim);
AffineExpr d0, d1; AffineExpr d0, d1;
bindDims(xferOp.getContext(), d0, d1); bindDims(xferOp.getContext(), d0, d1);
Value base = xferOp.getIndices()[*dim]; Value base = xferOp.getIndices()[*dim];
Value memrefIdx = makeComposedAffineApply(b, loc, d0 + d1, {base, iv}); Value memrefIdx =
affine::makeComposedAffineApply(b, loc, d0 + d1, {base, iv});
cond = lb.create<arith::CmpIOp>(arith::CmpIPredicate::sgt, memrefDim, cond = lb.create<arith::CmpIOp>(arith::CmpIPredicate::sgt, memrefDim,
memrefIdx); memrefIdx);
} }
// Condition check 2: Masked in? // Condition check 2: Masked in?
if (auto maskCond = generateMaskCheck(b, xferOp, iv)) { if (auto maskCond = generateMaskCheck(b, xferOp, iv)) {
if (cond) if (cond)
cond = lb.create<arith::AndIOp>(cond, maskCond); cond = lb.create<arith::AndIOp>(cond, maskCond);
▲ Show 20 Lines • Show All 916 Lines • ▼ Show 20 Linesget1dMemrefIndices(OpBuilder &b, OpTy xferOp, Value iv,
assert(map.getNumResults() == 1 && assert(map.getNumResults() == 1 &&
"Expected 1 permutation map result for 1D transfer"); "Expected 1 permutation map result for 1D transfer");
if (auto expr = map.getResult(0).template dyn_cast<AffineDimExpr>()) { if (auto expr = map.getResult(0).template dyn_cast<AffineDimExpr>()) {
Location loc = xferOp.getLoc(); Location loc = xferOp.getLoc();
auto dim = expr.getPosition(); auto dim = expr.getPosition();
AffineExpr d0, d1; AffineExpr d0, d1;
bindDims(xferOp.getContext(), d0, d1); bindDims(xferOp.getContext(), d0, d1);
Value offset = memrefIndices[dim]; Value offset = memrefIndices[dim];
memrefIndices[dim] = makeComposedAffineApply(b, loc, d0 + d1, {offset, iv}); memrefIndices[dim] =
affine::makeComposedAffineApply(b, loc, d0 + d1, {offset, iv});
return dim; return dim;
} }
assert(xferOp.isBroadcastDim(0) && assert(xferOp.isBroadcastDim(0) &&
"Expected AffineDimExpr or AffineConstantExpr"); "Expected AffineDimExpr or AffineConstantExpr");
return std::nullopt; return std::nullopt;
} }
▲ Show 20 Lines • Show All 201 Lines • Show Last 20 Lines

mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp

Show All 25 Lines
#include "llvm/ADT/TypeSwitch.h" #include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/Debug.h" #include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h" #include "llvm/Support/raw_ostream.h"
#include <optional> #include <optional>
#define DEBUG_TYPE "affine-analysis" #define DEBUG_TYPE "affine-analysis"
using namespace mlir; using namespace mlir;
using namespace affine;
using namespace presburger; using namespace presburger;
/// Get the value that is being reduced by `pos`-th reduction in the loop if /// Get the value that is being reduced by `pos`-th reduction in the loop if
/// such a reduction can be performed by affine parallel loops. This assumes /// such a reduction can be performed by affine parallel loops. This assumes
/// floating-point operations are commutative. On success, `kind` will be the /// floating-point operations are commutative. On success, `kind` will be the
/// reduction kind suitable for use in affine parallel loop builder. If the /// reduction kind suitable for use in affine parallel loop builder. If the
/// reduction is not supported, returns null. /// reduction is not supported, returns null.
static Value getSupportedReduction(AffineForOp forOp, unsigned pos, static Value getSupportedReduction(AffineForOp forOp, unsigned pos,
Show All 31 Linesstatic Value getSupportedReduction(AffineForOp forOp, unsigned pos,
if (!maybeKind) if (!maybeKind)
return nullptr; return nullptr;
kind = *maybeKind; kind = *maybeKind;
return reducedVal; return reducedVal;
} }
/// Populate `supportedReductions` with descriptors of the supported reductions. /// Populate `supportedReductions` with descriptors of the supported reductions.
void mlir::getSupportedReductions( void mlir::affine::getSupportedReductions(
AffineForOp forOp, SmallVectorImpl<LoopReduction> &supportedReductions) { AffineForOp forOp, SmallVectorImpl<LoopReduction> &supportedReductions) {
unsigned numIterArgs = forOp.getNumIterOperands(); unsigned numIterArgs = forOp.getNumIterOperands();
if (numIterArgs == 0) if (numIterArgs == 0)
return; return;
supportedReductions.reserve(numIterArgs); supportedReductions.reserve(numIterArgs);
for (unsigned i = 0; i < numIterArgs; ++i) { for (unsigned i = 0; i < numIterArgs; ++i) {
arith::AtomicRMWKind kind; arith::AtomicRMWKind kind;
if (Value value = getSupportedReduction(forOp, i, kind)) if (Value value = getSupportedReduction(forOp, i, kind))
supportedReductions.emplace_back(LoopReduction{kind, i, value}); supportedReductions.emplace_back(LoopReduction{kind, i, value});
} }
} }
/// Returns true if `forOp' is a parallel loop. If `parallelReductions` is /// Returns true if `forOp' is a parallel loop. If `parallelReductions` is
/// provided, populates it with descriptors of the parallelizable reductions and /// provided, populates it with descriptors of the parallelizable reductions and
/// treats them as not preventing parallelization. /// treats them as not preventing parallelization.
bool mlir::isLoopParallel(AffineForOp forOp, bool mlir::affine::isLoopParallel(
SmallVectorImpl<LoopReduction> *parallelReductions) { AffineForOp forOp, SmallVectorImpl<LoopReduction> *parallelReductions) {
unsigned numIterArgs = forOp.getNumIterOperands(); unsigned numIterArgs = forOp.getNumIterOperands();
// Loop is not parallel if it has SSA loop-carried dependences and reduction // Loop is not parallel if it has SSA loop-carried dependences and reduction
// detection is not requested. // detection is not requested.
if (numIterArgs > 0 && !parallelReductions) if (numIterArgs > 0 && !parallelReductions)
return false; return false;
// Find supported reductions of requested. // Find supported reductions of requested.
Show All 20 Lines if (hasSingleEffect<MemoryEffects::Allocate>(defOp, v) &&
enclosingOp->isProperAncestor(defOp)) enclosingOp->isProperAncestor(defOp))
return true; return true;
// Aliasing ops. // Aliasing ops.
auto viewOp = dyn_cast<ViewLikeOpInterface>(defOp); auto viewOp = dyn_cast<ViewLikeOpInterface>(defOp);
return viewOp && isLocallyDefined(viewOp.getViewSource(), enclosingOp); return viewOp && isLocallyDefined(viewOp.getViewSource(), enclosingOp);
} }
bool mlir::isLoopMemoryParallel(AffineForOp forOp) { bool mlir::affine::isLoopMemoryParallel(AffineForOp forOp) {
// Any memref-typed iteration arguments are treated as serializing. // Any memref-typed iteration arguments are treated as serializing.
if (llvm::any_of(forOp.getResultTypes(), if (llvm::any_of(forOp.getResultTypes(),
[](Type type) { return type.isa<BaseMemRefType>(); })) [](Type type) { return type.isa<BaseMemRefType>(); }))
return false; return false;
// Collect all load and store ops in loop nest rooted at 'forOp'. // Collect all load and store ops in loop nest rooted at 'forOp'.
SmallVector<Operation *, 8> loadAndStoreOps; SmallVector<Operation *, 8> loadAndStoreOps;
auto walkResult = forOp.walk([&](Operation *op) -> WalkResult { auto walkResult = forOp.walk([&](Operation *op) -> WalkResult {
Show All 37 Linesbool mlir::affine::isLoopMemoryParallel(AffineForOp forOp) {
return true; return true;
} }
/// Returns the sequence of AffineApplyOp Operations operation in /// Returns the sequence of AffineApplyOp Operations operation in
/// 'affineApplyOps', which are reachable via a search starting from 'operands', /// 'affineApplyOps', which are reachable via a search starting from 'operands',
/// and ending at operands which are not defined by AffineApplyOps. /// and ending at operands which are not defined by AffineApplyOps.
// TODO: Add a method to AffineApplyOp which forward substitutes the // TODO: Add a method to AffineApplyOp which forward substitutes the
// AffineApplyOp into any user AffineApplyOps. // AffineApplyOp into any user AffineApplyOps.
void mlir::getReachableAffineApplyOps( void mlir::affine::getReachableAffineApplyOps(
ArrayRef<Value> operands, SmallVectorImpl<Operation *> &affineApplyOps) { ArrayRef<Value> operands, SmallVectorImpl<Operation *> &affineApplyOps) {
struct State { struct State {
// The ssa value for this node in the DFS traversal. // The ssa value for this node in the DFS traversal.
Value value; Value value;
// The operand index of 'value' to explore next during DFS traversal. // The operand index of 'value' to explore next during DFS traversal.
unsigned operandIndex; unsigned operandIndex;
}; };
SmallVector<State, 4> worklist; SmallVector<State, 4> worklist;
Show All 33 Lines
// Builds a system of constraints with dimensional variables corresponding to // Builds a system of constraints with dimensional variables corresponding to
// the loop IVs of the forOps appearing in that order. Any symbols founds in // the loop IVs of the forOps appearing in that order. Any symbols founds in
// the bound operands are added as symbols in the system. Returns failure for // the bound operands are added as symbols in the system. Returns failure for
// the yet unimplemented cases. // the yet unimplemented cases.
// TODO: Handle non-unit steps through local variables or stride information in // TODO: Handle non-unit steps through local variables or stride information in
// FlatAffineValueConstraints. (For eg., by using iv - lb % step = 0 and/or by // FlatAffineValueConstraints. (For eg., by using iv - lb % step = 0 and/or by
// introducing a method in FlatAffineValueConstraints // introducing a method in FlatAffineValueConstraints
// setExprStride(ArrayRef<int64_t> expr, int64_t stride) // setExprStride(ArrayRef<int64_t> expr, int64_t stride)
LogicalResult mlir::getIndexSet(MutableArrayRef<Operation *> ops, LogicalResult mlir::affine::getIndexSet(MutableArrayRef<Operation *> ops,
FlatAffineValueConstraints *domain) { FlatAffineValueConstraints *domain) {
SmallVector<Value, 4> indices; SmallVector<Value, 4> indices;
SmallVector<Operation *, 8> loopOps; SmallVector<Operation *, 8> loopOps;
size_t numDims = 0; size_t numDims = 0;
for (Operation *op : ops) { for (Operation *op : ops) {
if (!isa<AffineForOp, AffineIfOp, AffineParallelOp>(op)) { if (!isa<AffineForOp, AffineIfOp, AffineParallelOp>(op)) {
LLVM_DEBUG(llvm::dbgs() << "getIndexSet only handles affine.for/if/" LLVM_DEBUG(llvm::dbgs() << "getIndexSet only handles affine.for/if/"
"parallel ops"); "parallel ops");
return failure(); return failure();
▲ Show 20 Lines • Show All 340 Lines • ▼ Show 20 Lines
// 0 -1 0 0 0 0 0 50 >= 0 // 0 -1 0 0 0 0 0 50 >= 0
// 0 0 1 0 0 0 0 0 >= 0 // 0 0 1 0 0 0 0 0 >= 0
// 0 0 -1 0 0 0 0 100 >= 0 // 0 0 -1 0 0 0 0 100 >= 0
// 0 0 0 1 0 0 0 0 >= 0 // 0 0 0 1 0 0 0 0 >= 0
// 0 0 0 -1 0 0 0 50 >= 0 // 0 0 0 -1 0 0 0 50 >= 0
// //
// //
// TODO: Support AffineExprs mod/floordiv/ceildiv. // TODO: Support AffineExprs mod/floordiv/ceildiv.
DependenceResult mlir::checkMemrefAccessDependence( DependenceResult mlir::affine::checkMemrefAccessDependence(
const MemRefAccess &srcAccess, const MemRefAccess &dstAccess, const MemRefAccess &srcAccess, const MemRefAccess &dstAccess,
unsigned loopDepth, FlatAffineValueConstraints *dependenceConstraints, unsigned loopDepth, FlatAffineValueConstraints *dependenceConstraints,
SmallVector<DependenceComponent, 2> *dependenceComponents, bool allowRAR) { SmallVector<DependenceComponent, 2> *dependenceComponents, bool allowRAR) {
LLVM_DEBUG(llvm::dbgs() << "Checking for dependence at depth: " LLVM_DEBUG(llvm::dbgs() << "Checking for dependence at depth: "
<< Twine(loopDepth) << " between:\n";); << Twine(loopDepth) << " between:\n";);
LLVM_DEBUG(srcAccess.opInst->dump()); LLVM_DEBUG(srcAccess.opInst->dump());
LLVM_DEBUG(dstAccess.opInst->dump()); LLVM_DEBUG(dstAccess.opInst->dump());
▲ Show 20 Lines • Show All 60 Lines • ▼ Show 20 LinesDependenceResult mlir::affine::checkMemrefAccessDependence(
if (dependenceConstraints) if (dependenceConstraints)
*dependenceConstraints = dstRel; *dependenceConstraints = dstRel;
return DependenceResult::HasDependence; return DependenceResult::HasDependence;
} }
/// Gathers dependence components for dependences between all ops in loop nest /// Gathers dependence components for dependences between all ops in loop nest
/// rooted at 'forOp' at loop depths in range [1, maxLoopDepth]. /// rooted at 'forOp' at loop depths in range [1, maxLoopDepth].
void mlir::getDependenceComponents( void mlir::affine::getDependenceComponents(
AffineForOp forOp, unsigned maxLoopDepth, AffineForOp forOp, unsigned maxLoopDepth,
std::vector<SmallVector<DependenceComponent, 2>> *depCompsVec) { std::vector<SmallVector<DependenceComponent, 2>> *depCompsVec) {
// Collect all load and store ops in loop nest rooted at 'forOp'. // Collect all load and store ops in loop nest rooted at 'forOp'.
SmallVector<Operation *, 8> loadAndStoreOps; SmallVector<Operation *, 8> loadAndStoreOps;
forOp->walk([&](Operation *op) { forOp->walk([&](Operation *op) {
if (isa<AffineReadOpInterface, AffineWriteOpInterface>(op)) if (isa<AffineReadOpInterface, AffineWriteOpInterface>(op))
loadAndStoreOps.push_back(op); loadAndStoreOps.push_back(op);
}); });
Show All 22 Lines

mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp

Show All 25 Lines
#include "llvm/ADT/SmallVector.h" #include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Debug.h" #include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h" #include "llvm/Support/raw_ostream.h"
#include <optional> #include <optional>
#define DEBUG_TYPE "affine-structures" #define DEBUG_TYPE "affine-structures"
using namespace mlir; using namespace mlir;
using namespace affine;
using namespace presburger; using namespace presburger;
void FlatAffineValueConstraints::addInductionVarOrTerminalSymbol(Value val) { void FlatAffineValueConstraints::addInductionVarOrTerminalSymbol(Value val) {
if (containsVar(val)) if (containsVar(val))
return; return;
// Caller is expected to fully compose map/operands if necessary. // Caller is expected to fully compose map/operands if necessary.
▲ Show 20 Lines • Show All 442 Lines • ▼ Show 20 Linesvoid FlatAffineRelation::removeVarRange(VarKind kind, unsigned varStart,
unsigned intersectRangeRHS = std::max(varStart, getNumDomainDims()); unsigned intersectRangeRHS = std::max(varStart, getNumDomainDims());
if (intersectDomainLHS > intersectDomainRHS) if (intersectDomainLHS > intersectDomainRHS)
numDomainDims -= intersectDomainLHS - intersectDomainRHS; numDomainDims -= intersectDomainLHS - intersectDomainRHS;
if (intersectRangeLHS > intersectRangeRHS) if (intersectRangeLHS > intersectRangeRHS)
numRangeDims -= intersectRangeLHS - intersectRangeRHS; numRangeDims -= intersectRangeLHS - intersectRangeRHS;
} }
LogicalResult mlir::getRelationFromMap(AffineMap &map, LogicalResult mlir::affine::getRelationFromMap(AffineMap &map,
FlatAffineRelation &rel) { FlatAffineRelation &rel) {
// Get flattened affine expressions. // Get flattened affine expressions.
std::vector<SmallVector<int64_t, 8>> flatExprs; std::vector<SmallVector<int64_t, 8>> flatExprs;
FlatAffineValueConstraints localVarCst; FlatAffineValueConstraints localVarCst;
if (failed(getFlattenedAffineExprs(map, &flatExprs, &localVarCst))) if (failed(getFlattenedAffineExprs(map, &flatExprs, &localVarCst)))
return failure(); return failure();
unsigned oldDimNum = localVarCst.getNumDimVars(); unsigned oldDimNum = localVarCst.getNumDimVars();
unsigned oldCols = localVarCst.getNumCols(); unsigned oldCols = localVarCst.getNumCols();
Show All 18 Lines for (unsigned i = 0, e = map.getNumResults(); i < e; ++i) {
localVarCst.addEquality(eq); localVarCst.addEquality(eq);
} }
// Create relation and return success. // Create relation and return success.
rel = FlatAffineRelation(numDomainVars, numRangeVars, localVarCst); rel = FlatAffineRelation(numDomainVars, numRangeVars, localVarCst);
return success(); return success();
} }
LogicalResult mlir::getRelationFromMap(const AffineValueMap &map, LogicalResult mlir::affine::getRelationFromMap(const AffineValueMap &map,
FlatAffineRelation &rel) { FlatAffineRelation &rel) {
AffineMap affineMap = map.getAffineMap(); AffineMap affineMap = map.getAffineMap();
if (failed(getRelationFromMap(affineMap, rel))) if (failed(getRelationFromMap(affineMap, rel)))
return failure(); return failure();
// Set symbol values for domain dimensions and symbols. // Set symbol values for domain dimensions and symbols.
for (unsigned i = 0, e = rel.getNumDomainDims(); i < e; ++i) for (unsigned i = 0, e = rel.getNumDomainDims(); i < e; ++i)
rel.setValue(i, map.getOperand(i)); rel.setValue(i, map.getOperand(i));
for (unsigned i = rel.getNumDimVars(), e = rel.getNumDimAndSymbolVars(); for (unsigned i = rel.getNumDimVars(), e = rel.getNumDimAndSymbolVars();
i < e; ++i) i < e; ++i)
rel.setValue(i, map.getOperand(i - rel.getNumRangeDims())); rel.setValue(i, map.getOperand(i - rel.getNumRangeDims()));
return success(); return success();
} }

mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp

Show All 22 Lines
#include "llvm/ADT/DenseSet.h" #include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/SmallPtrSet.h" #include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallString.h" #include "llvm/ADT/SmallString.h"
#include <numeric> #include <numeric>
#include <optional> #include <optional>
#include <type_traits> #include <type_traits>
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
/// Returns the trip count of the loop as an affine expression if the latter is /// Returns the trip count of the loop as an affine expression if the latter is
/// expressible as an affine expression, and nullptr otherwise. The trip count /// expressible as an affine expression, and nullptr otherwise. The trip count
/// expression is simplified before returning. This method only utilizes map /// expression is simplified before returning. This method only utilizes map
/// composition to construct lower and upper bounds before computing the trip /// composition to construct lower and upper bounds before computing the trip
/// count expressions. /// count expressions.
void mlir::getTripCountMapAndOperands( void mlir::affine::getTripCountMapAndOperands(
AffineForOp forOp, AffineMap *tripCountMap, AffineForOp forOp, AffineMap *tripCountMap,
SmallVectorImpl<Value> *tripCountOperands) { SmallVectorImpl<Value> *tripCountOperands) {
MLIRContext *context = forOp.getContext(); MLIRContext *context = forOp.getContext();
int64_t step = forOp.getStep(); int64_t step = forOp.getStep();
int64_t loopSpan; int64_t loopSpan;
if (forOp.hasConstantBounds()) { if (forOp.hasConstantBounds()) {
int64_t lb = forOp.getConstantLowerBound(); int64_t lb = forOp.getConstantLowerBound();
int64_t ub = forOp.getConstantUpperBound(); int64_t ub = forOp.getConstantUpperBound();
Show All 32 Linesvoid mlir::affine::getTripCountMapAndOperands(
tripCountOperands->assign(tripCountValueMap.getOperands().begin(), tripCountOperands->assign(tripCountValueMap.getOperands().begin(),
tripCountValueMap.getOperands().end()); tripCountValueMap.getOperands().end());
} }
/// Returns the trip count of the loop if it's a constant, std::nullopt /// Returns the trip count of the loop if it's a constant, std::nullopt
/// otherwise. This method uses affine expression analysis (in turn using /// otherwise. This method uses affine expression analysis (in turn using
/// getTripCount) and is able to determine constant trip count in non-trivial /// getTripCount) and is able to determine constant trip count in non-trivial
/// cases. /// cases.
std::optional<uint64_t> mlir::getConstantTripCount(AffineForOp forOp) { std::optional<uint64_t> mlir::affine::getConstantTripCount(AffineForOp forOp) {
SmallVector<Value, 4> operands; SmallVector<Value, 4> operands;
AffineMap map; AffineMap map;
getTripCountMapAndOperands(forOp, &map, &operands); getTripCountMapAndOperands(forOp, &map, &operands);
if (!map) if (!map)
return std::nullopt; return std::nullopt;
// Take the min if all trip counts are constant. // Take the min if all trip counts are constant.
Show All 9 Lines if (auto constExpr = resultExpr.dyn_cast<AffineConstantExpr>()) {
return std::nullopt; return std::nullopt;
} }
return tripCount; return tripCount;
} }
/// Returns the greatest known integral divisor of the trip count. Affine /// Returns the greatest known integral divisor of the trip count. Affine
/// expression analysis is used (indirectly through getTripCount), and /// expression analysis is used (indirectly through getTripCount), and
/// this method is thus able to determine non-trivial divisors. /// this method is thus able to determine non-trivial divisors.
uint64_t mlir::getLargestDivisorOfTripCount(AffineForOp forOp) { uint64_t mlir::affine::getLargestDivisorOfTripCount(AffineForOp forOp) {
SmallVector<Value, 4> operands; SmallVector<Value, 4> operands;
AffineMap map; AffineMap map;
getTripCountMapAndOperands(forOp, &map, &operands); getTripCountMapAndOperands(forOp, &map, &operands);
if (!map) if (!map)
return 1; return 1;
// The largest divisor of the trip count is the GCD of the individual largest // The largest divisor of the trip count is the GCD of the individual largest
▲ Show 20 Lines • Show All 57 Lines • ▼ Show 20 Linesstatic bool isAccessIndexInvariant(Value iv, Value index) {
} }
auto composeOp = cast<AffineApplyOp>(affineApplyOps[0]); auto composeOp = cast<AffineApplyOp>(affineApplyOps[0]);
// We need yet another level of indirection because the `dim` index of the // We need yet another level of indirection because the `dim` index of the
// access may not correspond to the `dim` index of composeOp. // access may not correspond to the `dim` index of composeOp.
return !composeOp.getAffineValueMap().isFunctionOf(0, iv); return !composeOp.getAffineValueMap().isFunctionOf(0, iv);
} }
DenseSet<Value> mlir::getInvariantAccesses(Value iv, ArrayRef<Value> indices) { DenseSet<Value> mlir::affine::getInvariantAccesses(Value iv,
ArrayRef<Value> indices) {
DenseSet<Value> res; DenseSet<Value> res;
for (auto val : indices) { for (auto val : indices) {
if (isAccessIndexInvariant(iv, val)) { if (isAccessIndexInvariant(iv, val)) {
res.insert(val); res.insert(val);
} }
} }
return res; return res;
} }
▲ Show 20 Lines • Show All 135 Lines • ▼ Show 20 Lines for (auto ls : loadAndStoresMatched) {
} }
if (isVectorizableOp && !isVectorizableOp(loop, *op)) { if (isVectorizableOp && !isVectorizableOp(loop, *op)) {
return false; return false;
} }
} }
return true; return true;
} }
bool mlir::isVectorizableLoopBody(AffineForOp loop, int *memRefDim, bool mlir::affine::isVectorizableLoopBody(
NestedPattern &vectorTransferMatcher) { AffineForOp loop, int *memRefDim, NestedPattern &vectorTransferMatcher) {
*memRefDim = -1; *memRefDim = -1;
VectorizableOpFun fun([memRefDim](AffineForOp loop, Operation &op) { VectorizableOpFun fun([memRefDim](AffineForOp loop, Operation &op) {
auto load = dyn_cast<AffineLoadOp>(op); auto load = dyn_cast<AffineLoadOp>(op);
auto store = dyn_cast<AffineStoreOp>(op); auto store = dyn_cast<AffineStoreOp>(op);
int thisOpMemRefDim = -1; int thisOpMemRefDim = -1;
bool isContiguous = load ? isContiguousAccess(loop.getInductionVar(), load, bool isContiguous = load ? isContiguousAccess(loop.getInductionVar(), load,
&thisOpMemRefDim) &thisOpMemRefDim)
: isContiguousAccess(loop.getInductionVar(), store, : isContiguousAccess(loop.getInductionVar(), store,
&thisOpMemRefDim); &thisOpMemRefDim);
if (thisOpMemRefDim != -1) { if (thisOpMemRefDim != -1) {
// If memory accesses vary across different dimensions then the loop is // If memory accesses vary across different dimensions then the loop is
// not vectorizable. // not vectorizable.
if (*memRefDim != -1 && *memRefDim != thisOpMemRefDim) if (*memRefDim != -1 && *memRefDim != thisOpMemRefDim)
return false; return false;
*memRefDim = thisOpMemRefDim; *memRefDim = thisOpMemRefDim;
} }
return isContiguous; return isContiguous;
}); });
return isVectorizableLoopBodyWithOpCond(loop, fun, vectorTransferMatcher); return isVectorizableLoopBodyWithOpCond(loop, fun, vectorTransferMatcher);
} }
bool mlir::isVectorizableLoopBody(AffineForOp loop, bool mlir::affine::isVectorizableLoopBody(
NestedPattern &vectorTransferMatcher) { AffineForOp loop, NestedPattern &vectorTransferMatcher) {
return isVectorizableLoopBodyWithOpCond(loop, nullptr, vectorTransferMatcher); return isVectorizableLoopBodyWithOpCond(loop, nullptr, vectorTransferMatcher);
} }
/// Checks whether SSA dominance would be violated if a for op's body /// Checks whether SSA dominance would be violated if a for op's body
/// operations are shifted by the specified shifts. This method checks if a /// operations are shifted by the specified shifts. This method checks if a
/// 'def' and all its uses have the same shift factor. /// 'def' and all its uses have the same shift factor.
// TODO: extend this to check for memory-based dependence violation when we have // TODO: extend this to check for memory-based dependence violation when we have
// the support. // the support.
bool mlir::isOpwiseShiftValid(AffineForOp forOp, ArrayRef<uint64_t> shifts) { bool mlir::affine::isOpwiseShiftValid(AffineForOp forOp,
ArrayRef<uint64_t> shifts) {
auto *forBody = forOp.getBody(); auto *forBody = forOp.getBody();
assert(shifts.size() == forBody->getOperations().size()); assert(shifts.size() == forBody->getOperations().size());
// Work backwards over the body of the block so that the shift of a use's // Work backwards over the body of the block so that the shift of a use's
// ancestor operation in the block gets recorded before it's looked up. // ancestor operation in the block gets recorded before it's looked up.
DenseMap<Operation *, uint64_t> forBodyShift; DenseMap<Operation *, uint64_t> forBodyShift;
for (const auto &it : for (const auto &it :
llvm::enumerate(llvm::reverse(forBody->getOperations()))) { llvm::enumerate(llvm::reverse(forBody->getOperations()))) {
Show All 26 Lines

mlir/lib/Dialect/Affine/Analysis/NestedMatcher.cpp

Show All 11 Lines
#include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/STLExtras.h" #include "llvm/ADT/STLExtras.h"
#include "llvm/Support/Allocator.h" #include "llvm/Support/Allocator.h"
#include "llvm/Support/raw_ostream.h" #include "llvm/Support/raw_ostream.h"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
llvm::BumpPtrAllocator *&NestedMatch::allocator() { llvm::BumpPtrAllocator *&NestedMatch::allocator() {
thread_local llvm::BumpPtrAllocator *allocator = nullptr; thread_local llvm::BumpPtrAllocator *allocator = nullptr;
return allocator; return allocator;
} }
NestedMatch NestedMatch::build(Operation *operation, NestedMatch NestedMatch::build(Operation *operation,
ArrayRef<NestedMatch> nestedMatches) { ArrayRef<NestedMatch> nestedMatches) {
▲ Show 20 Lines • Show All 97 Lines • ▼ Show 20 Linesvoid NestedPattern::matchOne(Operation *op,
} }
} }
static bool isAffineForOp(Operation &op) { return isa<AffineForOp>(op); } static bool isAffineForOp(Operation &op) { return isa<AffineForOp>(op); }
static bool isAffineIfOp(Operation &op) { return isa<AffineIfOp>(op); } static bool isAffineIfOp(Operation &op) { return isa<AffineIfOp>(op); }
namespace mlir { namespace mlir {
namespace affine {
namespace matcher { namespace matcher {
NestedPattern Op(FilterFunctionType filter) { NestedPattern Op(FilterFunctionType filter) {
return NestedPattern({}, std::move(filter)); return NestedPattern({}, std::move(filter));
} }
NestedPattern If(const NestedPattern &child) { NestedPattern If(const NestedPattern &child) {
return NestedPattern(child, isAffineIfOp); return NestedPattern(child, isAffineIfOp);
Show All 30 Lines return NestedPattern(
nested, [=](Operation &op) { return isAffineForOp(op) && filter(op); }); nested, [=](Operation &op) { return isAffineForOp(op) && filter(op); });
} }
bool isLoadOrStore(Operation &op) { bool isLoadOrStore(Operation &op) {
return isa<AffineLoadOp, AffineStoreOp>(op); return isa<AffineLoadOp, AffineStoreOp>(op);
} }
} // namespace matcher } // namespace matcher
} // namespace affine
} // namespace mlir } // namespace mlir

mlir/lib/Dialect/Affine/Analysis/Utils.cpp

Show All 22 Lines
#include "llvm/ADT/SmallPtrSet.h" #include "llvm/ADT/SmallPtrSet.h"
#include "llvm/Support/Debug.h" #include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h" #include "llvm/Support/raw_ostream.h"
#include <optional> #include <optional>
#define DEBUG_TYPE "analysis-utils" #define DEBUG_TYPE "analysis-utils"
using namespace mlir; using namespace mlir;
using namespace affine;
using namespace presburger; using namespace presburger;
using llvm::SmallDenseMap; using llvm::SmallDenseMap;
using Node = MemRefDependenceGraph::Node; using Node = MemRefDependenceGraph::Node;
// LoopNestStateCollector walks loop nests and collects load and store // LoopNestStateCollector walks loop nests and collects load and store
// operations, and whether or not a region holding op other than ForOp and IfOp // operations, and whether or not a region holding op other than ForOp and IfOp
▲ Show 20 Lines • Show All 457 Lines • ▼ Show 20 Lines for (const auto &idAndNode : nodes) {
it = outEdges.find(idAndNode.first); it = outEdges.find(idAndNode.first);
if (it != outEdges.end()) { if (it != outEdges.end()) {
for (const auto &e : it->second) for (const auto &e : it->second)
os << " OutEdge: " << e.id << " " << e.value << "\n"; os << " OutEdge: " << e.id << " " << e.value << "\n";
} }
} }
} }
void mlir::getAffineForIVs(Operation &op, SmallVectorImpl<AffineForOp> *loops) { void mlir::affine::getAffineForIVs(Operation &op,
SmallVectorImpl<AffineForOp> *loops) {
auto *currOp = op.getParentOp(); auto *currOp = op.getParentOp();
AffineForOp currAffineForOp; AffineForOp currAffineForOp;
// Traverse up the hierarchy collecting all 'affine.for' operation while // Traverse up the hierarchy collecting all 'affine.for' operation while
// skipping over 'affine.if' operations. // skipping over 'affine.if' operations.
while (currOp) { while (currOp) {
if (AffineForOp currAffineForOp = dyn_cast<AffineForOp>(currOp)) if (AffineForOp currAffineForOp = dyn_cast<AffineForOp>(currOp))
loops->push_back(currAffineForOp); loops->push_back(currAffineForOp);
currOp = currOp->getParentOp(); currOp = currOp->getParentOp();
} }
std::reverse(loops->begin(), loops->end()); std::reverse(loops->begin(), loops->end());
} }
void mlir::getEnclosingAffineOps(Operation &op, void mlir::affine::getEnclosingAffineOps(Operation &op,
SmallVectorImpl<Operation *> *ops) { SmallVectorImpl<Operation *> *ops) {
ops->clear(); ops->clear();
Operation *currOp = op.getParentOp(); Operation *currOp = op.getParentOp();
// Traverse up the hierarchy collecting all `affine.for`, `affine.if`, and // Traverse up the hierarchy collecting all `affine.for`, `affine.if`, and
// affine.parallel operations. // affine.parallel operations.
while (currOp) { while (currOp) {
if (isa<AffineIfOp, AffineForOp, AffineParallelOp>(currOp)) if (isa<AffineIfOp, AffineForOp, AffineParallelOp>(currOp))
ops->push_back(currOp); ops->push_back(currOp);
▲ Show 20 Lines • Show All 531 Lines • ▼ Show 20 LinesLogicalResult MemRefRegion::compute(Operation *op, unsigned loopDepth,
cst.removeTrivialRedundancy(); cst.removeTrivialRedundancy();
LLVM_DEBUG(llvm::dbgs() << "Memory region:\n"); LLVM_DEBUG(llvm::dbgs() << "Memory region:\n");
LLVM_DEBUG(cst.dump()); LLVM_DEBUG(cst.dump());
return success(); return success();
} }
std::optional<int64_t> std::optional<int64_t>
mlir::getMemRefIntOrFloatEltSizeInBytes(MemRefType memRefType) { mlir::affine::getMemRefIntOrFloatEltSizeInBytes(MemRefType memRefType) {
auto elementType = memRefType.getElementType(); auto elementType = memRefType.getElementType();
unsigned sizeInBits; unsigned sizeInBits;
if (elementType.isIntOrFloat()) { if (elementType.isIntOrFloat()) {
sizeInBits = elementType.getIntOrFloatBitWidth(); sizeInBits = elementType.getIntOrFloatBitWidth();
} else if (auto vectorType = elementType.dyn_cast<VectorType>()) { } else if (auto vectorType = elementType.dyn_cast<VectorType>()) {
if (vectorType.getElementType().isIntOrFloat()) if (vectorType.getElementType().isIntOrFloat())
sizeInBits = sizeInBits =
Show All 33 Linesstd::optional<int64_t> MemRefRegion::getRegionSize() {
return *eltSize * *numElements; return *eltSize * *numElements;
} }
/// Returns the size of memref data in bytes if it's statically shaped, /// Returns the size of memref data in bytes if it's statically shaped,
/// std::nullopt otherwise. If the element of the memref has vector type, takes /// std::nullopt otherwise. If the element of the memref has vector type, takes
/// into account size of the vector as well. /// into account size of the vector as well.
// TODO: improve/complete this when we have target data. // TODO: improve/complete this when we have target data.
std::optional<uint64_t> std::optional<uint64_t>
mlir::getIntOrFloatMemRefSizeInBytes(MemRefType memRefType) { mlir::affine::getIntOrFloatMemRefSizeInBytes(MemRefType memRefType) {
if (!memRefType.hasStaticShape()) if (!memRefType.hasStaticShape())
return std::nullopt; return std::nullopt;
auto elementType = memRefType.getElementType(); auto elementType = memRefType.getElementType();
if (!elementType.isIntOrFloat() && !elementType.isa<VectorType>()) if (!elementType.isIntOrFloat() && !elementType.isa<VectorType>())
return std::nullopt; return std::nullopt;
auto sizeInBytes = getMemRefIntOrFloatEltSizeInBytes(memRefType); auto sizeInBytes = getMemRefIntOrFloatEltSizeInBytes(memRefType);
if (!sizeInBytes) if (!sizeInBytes)
return std::nullopt; return std::nullopt;
for (unsigned i = 0, e = memRefType.getRank(); i < e; i++) { for (unsigned i = 0, e = memRefType.getRank(); i < e; i++) {
sizeInBytes = *sizeInBytes * memRefType.getDimSize(i); sizeInBytes = *sizeInBytes * memRefType.getDimSize(i);
} }
return sizeInBytes; return sizeInBytes;
} }
template <typename LoadOrStoreOp> template <typename LoadOrStoreOp>
LogicalResult mlir::boundCheckLoadOrStoreOp(LoadOrStoreOp loadOrStoreOp, LogicalResult mlir::affine::boundCheckLoadOrStoreOp(LoadOrStoreOp loadOrStoreOp,
bool emitError) { bool emitError) {
static_assert(llvm::is_one_of<LoadOrStoreOp, AffineReadOpInterface, static_assert(llvm::is_one_of<LoadOrStoreOp, AffineReadOpInterface,
AffineWriteOpInterface>::value, AffineWriteOpInterface>::value,
"argument should be either a AffineReadOpInterface or a " "argument should be either a AffineReadOpInterface or a "
"AffineWriteOpInterface"); "AffineWriteOpInterface");
Operation *op = loadOrStoreOp.getOperation(); Operation *op = loadOrStoreOp.getOperation();
MemRefRegion region(op->getLoc()); MemRefRegion region(op->getLoc());
if (failed(region.compute(op, /*loopDepth=*/0, /*sliceState=*/nullptr, if (failed(region.compute(op, /*loopDepth=*/0, /*sliceState=*/nullptr,
Show All 38 Lines if (outOfBounds && emitError) {
<< "memref out of lower bound access along dimension #" << (r + 1); << "memref out of lower bound access along dimension #" << (r + 1);
} }
} }
return failure(outOfBounds); return failure(outOfBounds);
} }
// Explicitly instantiate the template so that the compiler knows we need them! // Explicitly instantiate the template so that the compiler knows we need them!
template LogicalResult template LogicalResult
mlir::boundCheckLoadOrStoreOp(AffineReadOpInterface loadOp, bool emitError); mlir::affine::boundCheckLoadOrStoreOp(AffineReadOpInterface loadOp,
bool emitError);
template LogicalResult template LogicalResult
mlir::boundCheckLoadOrStoreOp(AffineWriteOpInterface storeOp, bool emitError); mlir::affine::boundCheckLoadOrStoreOp(AffineWriteOpInterface storeOp,
bool emitError);
// Returns in 'positions' the Block positions of 'op' in each ancestor // Returns in 'positions' the Block positions of 'op' in each ancestor
// Block from the Block containing operation, stopping at 'limitBlock'. // Block from the Block containing operation, stopping at 'limitBlock'.
static void findInstPosition(Operation *op, Block *limitBlock, static void findInstPosition(Operation *op, Block *limitBlock,
SmallVectorImpl<unsigned> *positions) { SmallVectorImpl<unsigned> *positions) {
Block *block = op->getBlock(); Block *block = op->getBlock();
while (block != limitBlock) { while (block != limitBlock) {
// FIXME: This algorithm is unnecessarily O(n) and should be improved to not // FIXME: This algorithm is unnecessarily O(n) and should be improved to not
▲ Show 20 Lines • Show All 45 Lines • ▼ Show 20 Lines if (ivs.count(value) == 0) {
return failure(); return failure();
} }
} }
return success(); return success();
} }
/// Returns the innermost common loop depth for the set of operations in 'ops'. /// Returns the innermost common loop depth for the set of operations in 'ops'.
// TODO: Move this to LoopUtils. // TODO: Move this to LoopUtils.
unsigned mlir::getInnermostCommonLoopDepth( unsigned mlir::affine::getInnermostCommonLoopDepth(
ArrayRef<Operation *> ops, SmallVectorImpl<AffineForOp> *surroundingLoops) { ArrayRef<Operation *> ops, SmallVectorImpl<AffineForOp> *surroundingLoops) {
unsigned numOps = ops.size(); unsigned numOps = ops.size();
assert(numOps > 0 && "Expected at least one operation"); assert(numOps > 0 && "Expected at least one operation");
std::vector<SmallVector<AffineForOp, 4>> loops(numOps); std::vector<SmallVector<AffineForOp, 4>> loops(numOps);
unsigned loopDepthLimit = std::numeric_limits<unsigned>::max(); unsigned loopDepthLimit = std::numeric_limits<unsigned>::max();
for (unsigned i = 0; i < numOps; ++i) { for (unsigned i = 0; i < numOps; ++i) {
getAffineForIVs(*ops[i], &loops[i]); getAffineForIVs(*ops[i], &loops[i]);
Show All 15 Linesunsigned mlir::affine::getInnermostCommonLoopDepth(
return loopDepth; return loopDepth;
} }
/// Computes in 'sliceUnion' the union of all slice bounds computed at /// Computes in 'sliceUnion' the union of all slice bounds computed at
/// 'loopDepth' between all dependent pairs of ops in 'opsA' and 'opsB', and /// 'loopDepth' between all dependent pairs of ops in 'opsA' and 'opsB', and
/// then verifies if it is valid. Returns 'SliceComputationResult::Success' if /// then verifies if it is valid. Returns 'SliceComputationResult::Success' if
/// union was computed correctly, an appropriate failure otherwise. /// union was computed correctly, an appropriate failure otherwise.
SliceComputationResult SliceComputationResult
mlir::computeSliceUnion(ArrayRef<Operation *> opsA, ArrayRef<Operation *> opsB, mlir::affine::computeSliceUnion(ArrayRef<Operation *> opsA,
unsigned loopDepth, unsigned numCommonLoops, ArrayRef<Operation *> opsB, unsigned loopDepth,
bool isBackwardSlice, unsigned numCommonLoops, bool isBackwardSlice,
ComputationSliceState *sliceUnion) { ComputationSliceState *sliceUnion) {
// Compute the union of slice bounds between all pairs in 'opsA' and // Compute the union of slice bounds between all pairs in 'opsA' and
// 'opsB' in 'sliceUnionCst'. // 'opsB' in 'sliceUnionCst'.
FlatAffineValueConstraints sliceUnionCst; FlatAffineValueConstraints sliceUnionCst;
assert(sliceUnionCst.getNumDimAndSymbolVars() == 0); assert(sliceUnionCst.getNumDimAndSymbolVars() == 0);
std::vector<std::pair<Operation *, Operation *>> dependentOpPairs; std::vector<std::pair<Operation *, Operation *>> dependentOpPairs;
for (auto *i : opsA) { for (auto *i : opsA) {
MemRefAccess srcAccess(i); MemRefAccess srcAccess(i);
for (auto *j : opsB) { for (auto *j : opsB) {
Show All 20 Lines for (auto *j : opsB) {
return SliceComputationResult::GenericFailure; return SliceComputationResult::GenericFailure;
} }
if (result.value == DependenceResult::NoDependence) if (result.value == DependenceResult::NoDependence)
continue; continue;
dependentOpPairs.emplace_back(i, j); dependentOpPairs.emplace_back(i, j);
// Compute slice bounds for 'srcAccess' and 'dstAccess'. // Compute slice bounds for 'srcAccess' and 'dstAccess'.
ComputationSliceState tmpSliceState; ComputationSliceState tmpSliceState;
mlir::getComputationSliceState(i, j, &dependenceConstraints, loopDepth, mlir::affine::getComputationSliceState(i, j, &dependenceConstraints,
isBackwardSlice, &tmpSliceState); loopDepth, isBackwardSlice,
&tmpSliceState);
if (sliceUnionCst.getNumDimAndSymbolVars() == 0) { if (sliceUnionCst.getNumDimAndSymbolVars() == 0) {
// Initialize 'sliceUnionCst' with the bounds computed in previous step. // Initialize 'sliceUnionCst' with the bounds computed in previous step.
if (failed(tmpSliceState.getAsConstraints(&sliceUnionCst))) { if (failed(tmpSliceState.getAsConstraints(&sliceUnionCst))) {
LLVM_DEBUG(llvm::dbgs() LLVM_DEBUG(llvm::dbgs()
<< "Unable to compute slice bound constraints\n"); << "Unable to compute slice bound constraints\n");
return SliceComputationResult::GenericFailure; return SliceComputationResult::GenericFailure;
} }
▲ Show 20 Lines • Show All 125 Lines • ▼ Show 20 Lines if (!cExpr)
return std::nullopt; return std::nullopt;
return cExpr.getValue(); return cExpr.getValue();
} }
// Builds a map 'tripCountMap' from AffineForOp to constant trip count for loop // Builds a map 'tripCountMap' from AffineForOp to constant trip count for loop
// nest surrounding represented by slice loop bounds in 'slice'. Returns true // nest surrounding represented by slice loop bounds in 'slice'. Returns true
// on success, false otherwise (if a non-constant trip count was encountered). // on success, false otherwise (if a non-constant trip count was encountered).
// TODO: Make this work with non-unit step loops. // TODO: Make this work with non-unit step loops.
bool mlir::buildSliceTripCountMap( bool mlir::affine::buildSliceTripCountMap(
const ComputationSliceState &slice, const ComputationSliceState &slice,
llvm::SmallDenseMap<Operation *, uint64_t, 8> *tripCountMap) { llvm::SmallDenseMap<Operation *, uint64_t, 8> *tripCountMap) {
unsigned numSrcLoopIVs = slice.ivs.size(); unsigned numSrcLoopIVs = slice.ivs.size();
// Populate map from AffineForOp -> trip count // Populate map from AffineForOp -> trip count
for (unsigned i = 0; i < numSrcLoopIVs; ++i) { for (unsigned i = 0; i < numSrcLoopIVs; ++i) {
AffineForOp forOp = getForInductionVarOwner(slice.ivs[i]); AffineForOp forOp = getForInductionVarOwner(slice.ivs[i]);
auto *op = forOp.getOperation(); auto *op = forOp.getOperation();
AffineMap lbMap = slice.lbs[i]; AffineMap lbMap = slice.lbs[i];
Show All 21 Lines for (unsigned i = 0; i < numSrcLoopIVs; ++i) {
if (!tripCount.has_value()) if (!tripCount.has_value())
return false; return false;
(*tripCountMap)[op] = *tripCount; (*tripCountMap)[op] = *tripCount;
} }
return true; return true;
} }
// Return the number of iterations in the given slice. // Return the number of iterations in the given slice.
uint64_t mlir::getSliceIterationCount( uint64_t mlir::affine::getSliceIterationCount(
const llvm::SmallDenseMap<Operation *, uint64_t, 8> &sliceTripCountMap) { const llvm::SmallDenseMap<Operation *, uint64_t, 8> &sliceTripCountMap) {
uint64_t iterCount = 1; uint64_t iterCount = 1;
for (const auto &count : sliceTripCountMap) { for (const auto &count : sliceTripCountMap) {
iterCount *= count.second; iterCount *= count.second;
} }
return iterCount; return iterCount;
} }
const char *const kSliceFusionBarrierAttrName = "slice_fusion_barrier"; const char *const kSliceFusionBarrierAttrName = "slice_fusion_barrier";
// Computes slice bounds by projecting out any loop IVs from // Computes slice bounds by projecting out any loop IVs from
// 'dependenceConstraints' at depth greater than 'loopDepth', and computes slice // 'dependenceConstraints' at depth greater than 'loopDepth', and computes slice
// bounds in 'sliceState' which represent the one loop nest's IVs in terms of // bounds in 'sliceState' which represent the one loop nest's IVs in terms of
// the other loop nest's IVs, symbols and constants (using 'isBackwardsSlice'). // the other loop nest's IVs, symbols and constants (using 'isBackwardsSlice').
void mlir::getComputationSliceState( void mlir::affine::getComputationSliceState(
Operation *depSourceOp, Operation *depSinkOp, Operation *depSourceOp, Operation *depSinkOp,
FlatAffineValueConstraints *dependenceConstraints, unsigned loopDepth, FlatAffineValueConstraints *dependenceConstraints, unsigned loopDepth,
bool isBackwardSlice, ComputationSliceState *sliceState) { bool isBackwardSlice, ComputationSliceState *sliceState) {
// Get loop nest surrounding src operation. // Get loop nest surrounding src operation.
SmallVector<AffineForOp, 4> srcLoopIVs; SmallVector<AffineForOp, 4> srcLoopIVs;
getAffineForIVs(*depSourceOp, &srcLoopIVs); getAffineForIVs(*depSourceOp, &srcLoopIVs);
unsigned numSrcLoopIVs = srcLoopIVs.size(); unsigned numSrcLoopIVs = srcLoopIVs.size();
▲ Show 20 Lines • Show All 97 Lines • ▼ Show 20 Lines
// TODO: extend the slicing utility to compute slices that // TODO: extend the slicing utility to compute slices that
// aren't necessarily a one-to-one relation b/w the source and destination. The // aren't necessarily a one-to-one relation b/w the source and destination. The
// relation between the source and destination could be many-to-many in general. // relation between the source and destination could be many-to-many in general.
// TODO: the slice computation is incorrect in the cases // TODO: the slice computation is incorrect in the cases
// where the dependence from the source to the destination does not cover the // where the dependence from the source to the destination does not cover the
// entire destination index set. Subtract out the dependent destination // entire destination index set. Subtract out the dependent destination
// iterations from destination index set and check for emptiness --- this is one // iterations from destination index set and check for emptiness --- this is one
// solution. // solution.
AffineForOp AffineForOp mlir::affine::insertBackwardComputationSlice(
mlir::insertBackwardComputationSlice(Operation *srcOpInst, Operation *dstOpInst, Operation *srcOpInst, Operation *dstOpInst, unsigned dstLoopDepth,
unsigned dstLoopDepth,
ComputationSliceState *sliceState) { ComputationSliceState *sliceState) {
// Get loop nest surrounding src operation. // Get loop nest surrounding src operation.
SmallVector<AffineForOp, 4> srcLoopIVs; SmallVector<AffineForOp, 4> srcLoopIVs;
getAffineForIVs(*srcOpInst, &srcLoopIVs); getAffineForIVs(*srcOpInst, &srcLoopIVs);
unsigned numSrcLoopIVs = srcLoopIVs.size(); unsigned numSrcLoopIVs = srcLoopIVs.size();
// Get loop nest surrounding dst operation. // Get loop nest surrounding dst operation.
SmallVector<AffineForOp, 4> dstLoopIVs; SmallVector<AffineForOp, 4> dstLoopIVs;
getAffineForIVs(*dstOpInst, &dstLoopIVs); getAffineForIVs(*dstOpInst, &dstLoopIVs);
▲ Show 20 Lines • Show All 62 Lines • ▼ Show 20 Lines
} }
bool MemRefAccess::isStore() const { bool MemRefAccess::isStore() const {
return isa<AffineWriteOpInterface>(opInst); return isa<AffineWriteOpInterface>(opInst);
} }
/// Returns the nesting depth of this statement, i.e., the number of loops /// Returns the nesting depth of this statement, i.e., the number of loops
/// surrounding this statement. /// surrounding this statement.
unsigned mlir::getNestingDepth(Operation *op) { unsigned mlir::affine::getNestingDepth(Operation *op) {
Operation *currOp = op; Operation *currOp = op;
unsigned depth = 0; unsigned depth = 0;
while ((currOp = currOp->getParentOp())) { while ((currOp = currOp->getParentOp())) {
if (isa<AffineForOp>(currOp)) if (isa<AffineForOp>(currOp))
depth++; depth++;
} }
return depth; return depth;
} }
Show All 11 Linesbool MemRefAccess::operator==(const MemRefAccess &rhs) const {
AffineValueMap diff, thisMap, rhsMap; AffineValueMap diff, thisMap, rhsMap;
getAccessMap(&thisMap); getAccessMap(&thisMap);
rhs.getAccessMap(&rhsMap); rhs.getAccessMap(&rhsMap);
AffineValueMap::difference(thisMap, rhsMap, &diff); AffineValueMap::difference(thisMap, rhsMap, &diff);
return llvm::all_of(diff.getAffineMap().getResults(), return llvm::all_of(diff.getAffineMap().getResults(),
[](AffineExpr e) { return e == 0; }); [](AffineExpr e) { return e == 0; });
} }
void mlir::getAffineIVs(Operation &op, SmallVectorImpl<Value> &ivs) { void mlir::affine::getAffineIVs(Operation &op, SmallVectorImpl<Value> &ivs) {
auto *currOp = op.getParentOp(); auto *currOp = op.getParentOp();
AffineForOp currAffineForOp; AffineForOp currAffineForOp;
// Traverse up the hierarchy collecting all 'affine.for' and affine.parallel // Traverse up the hierarchy collecting all 'affine.for' and affine.parallel
// operation while skipping over 'affine.if' operations. // operation while skipping over 'affine.if' operations.
while (currOp) { while (currOp) {
if (AffineForOp currAffineForOp = dyn_cast<AffineForOp>(currOp)) if (AffineForOp currAffineForOp = dyn_cast<AffineForOp>(currOp))
ivs.push_back(currAffineForOp.getInductionVar()); ivs.push_back(currAffineForOp.getInductionVar());
else if (auto parOp = dyn_cast<AffineParallelOp>(currOp)) else if (auto parOp = dyn_cast<AffineParallelOp>(currOp))
llvm::append_range(ivs, parOp.getIVs()); llvm::append_range(ivs, parOp.getIVs());
currOp = currOp->getParentOp(); currOp = currOp->getParentOp();
} }
std::reverse(ivs.begin(), ivs.end()); std::reverse(ivs.begin(), ivs.end());
} }
/// Returns the number of surrounding loops common to 'loopsA' and 'loopsB', /// Returns the number of surrounding loops common to 'loopsA' and 'loopsB',
/// where each lists loops from outer-most to inner-most in loop nest. /// where each lists loops from outer-most to inner-most in loop nest.
unsigned mlir::getNumCommonSurroundingLoops(Operation &a, Operation &b) { unsigned mlir::affine::getNumCommonSurroundingLoops(Operation &a,
Operation &b) {
SmallVector<Value, 4> loopsA, loopsB; SmallVector<Value, 4> loopsA, loopsB;
getAffineIVs(a, loopsA); getAffineIVs(a, loopsA);
getAffineIVs(b, loopsB); getAffineIVs(b, loopsB);
unsigned minNumLoops = std::min(loopsA.size(), loopsB.size()); unsigned minNumLoops = std::min(loopsA.size(), loopsB.size());
unsigned numCommonLoops = 0; unsigned numCommonLoops = 0;
for (unsigned i = 0; i < minNumLoops; ++i) { for (unsigned i = 0; i < minNumLoops; ++i) {
if (loopsA[i] != loopsB[i]) if (loopsA[i] != loopsB[i])
▲ Show 20 Lines • Show All 42 Lines • ▼ Show 20 Lines for (const auto &region : regions) {
std::optional<int64_t> size = region.second->getRegionSize(); std::optional<int64_t> size = region.second->getRegionSize();
if (!size.has_value()) if (!size.has_value())
return std::nullopt; return std::nullopt;
totalSizeInBytes += *size; totalSizeInBytes += *size;
} }
return totalSizeInBytes; return totalSizeInBytes;
} }
std::optional<int64_t> mlir::getMemoryFootprintBytes(AffineForOp forOp, std::optional<int64_t> mlir::affine::getMemoryFootprintBytes(AffineForOp forOp,
int memorySpace) { int memorySpace) {
auto *forInst = forOp.getOperation(); auto *forInst = forOp.getOperation();
return ::getMemoryFootprintBytes( return ::getMemoryFootprintBytes(
*forInst->getBlock(), Block::iterator(forInst), *forInst->getBlock(), Block::iterator(forInst),
std::next(Block::iterator(forInst)), memorySpace); std::next(Block::iterator(forInst)), memorySpace);
} }
/// Returns whether a loop is parallel and contains a reduction loop. /// Returns whether a loop is parallel and contains a reduction loop.
bool mlir::isLoopParallelAndContainsReduction(AffineForOp forOp) { bool mlir::affine::isLoopParallelAndContainsReduction(AffineForOp forOp) {
SmallVector<LoopReduction> reductions; SmallVector<LoopReduction> reductions;
if (!isLoopParallel(forOp, &reductions)) if (!isLoopParallel(forOp, &reductions))
return false; return false;
return !reductions.empty(); return !reductions.empty();
} }
/// Returns in 'sequentialLoops' all sequential loops in loop nest rooted /// Returns in 'sequentialLoops' all sequential loops in loop nest rooted
/// at 'forOp'. /// at 'forOp'.
void mlir::getSequentialLoops(AffineForOp forOp, void mlir::affine::getSequentialLoops(
llvm::SmallDenseSet<Value, 8> *sequentialLoops) { AffineForOp forOp, llvm::SmallDenseSet<Value, 8> *sequentialLoops) {
forOp->walk([&](Operation *op) { forOp->walk([&](Operation *op) {
if (auto innerFor = dyn_cast<AffineForOp>(op)) if (auto innerFor = dyn_cast<AffineForOp>(op))
if (!isLoopParallel(innerFor)) if (!isLoopParallel(innerFor))
sequentialLoops->insert(innerFor.getInductionVar()); sequentialLoops->insert(innerFor.getInductionVar());
}); });
} }
IntegerSet mlir::simplifyIntegerSet(IntegerSet set) { IntegerSet mlir::affine::simplifyIntegerSet(IntegerSet set) {
FlatAffineValueConstraints fac(set); FlatAffineValueConstraints fac(set);
if (fac.isEmpty()) if (fac.isEmpty())
return IntegerSet::getEmptySet(set.getNumDims(), set.getNumSymbols(), return IntegerSet::getEmptySet(set.getNumDims(), set.getNumSymbols(),
set.getContext()); set.getContext());
fac.removeTrivialRedundancy(); fac.removeTrivialRedundancy();
auto simplifiedSet = fac.getAsIntegerSet(set.getContext()); auto simplifiedSet = fac.getAsIntegerSet(set.getContext());
assert(simplifiedSet && "guaranteed to succeed while roundtripping"); assert(simplifiedSet && "guaranteed to succeed while roundtripping");
▲ Show 20 Lines • Show All 69 Lines • ▼ Show 20 Lines
// (set `opBound` to `op` upper bound in terms of `invar` and "extra syms") // (set `opBound` to `op` upper bound in terms of `invar` and "extra syms")
// ... | 0 | -1 | 0 | ... | ... | = 0 // ... | 0 | -1 | 0 | ... | ... | = 0
// ------+-------+---------+-----+---------------+-------+------------------- // ------+-------+---------+-----+---------------+-------+-------------------
// (for each `op` map result r_i: set r_i to corresponding map result, // (for each `op` map result r_i: set r_i to corresponding map result,
// prove that r_i >= minOpUb via contradiction) // prove that r_i >= minOpUb via contradiction)
// ... | 0 | 0 | -1 | ... | ... | = 0 // ... | 0 | 0 | -1 | ... | ... | = 0
// 0 | 0 | 1 | -1 | 0 | -1 | >= 0 // 0 | 0 | 1 | -1 | 0 | -1 | >= 0
// //
FailureOr<AffineValueMap> FailureOr<AffineValueMap> mlir::affine::simplifyConstrainedMinMaxOp(
mlir::simplifyConstrainedMinMaxOp(Operation *op, Operation *op, FlatAffineValueConstraints constraints) {
FlatAffineValueConstraints constraints) {
bool isMin = isa<AffineMinOp>(op); bool isMin = isa<AffineMinOp>(op);
assert((isMin || isa<AffineMaxOp>(op)) && "expect AffineMin/MaxOp"); assert((isMin || isa<AffineMaxOp>(op)) && "expect AffineMin/MaxOp");
MLIRContext *ctx = op->getContext(); MLIRContext *ctx = op->getContext();
Builder builder(ctx); Builder builder(ctx);
AffineMap map = AffineMap map =
isMin ? cast<AffineMinOp>(op).getMap() : cast<AffineMaxOp>(op).getMap(); isMin ? cast<AffineMinOp>(op).getMap() : cast<AffineMaxOp>(op).getMap();
ValueRange operands = op->getOperands(); ValueRange operands = op->getOperands();
unsigned numResults = map.getNumResults(); unsigned numResults = map.getNumResults();
▲ Show 20 Lines • Show All 83 Lines • ▼ Show 20 Lines if (auto bound = constraints.getConstantBound64(BoundType::EQ, i)) {
AffineExpr expr = AffineExpr expr =
i < newMap.getNumDims() i < newMap.getNumDims()
? builder.getAffineDimExpr(i) ? builder.getAffineDimExpr(i)
: builder.getAffineSymbolExpr(i - newMap.getNumDims()); : builder.getAffineSymbolExpr(i - newMap.getNumDims());
newMap = newMap.replace(expr, builder.getAffineConstantExpr(*bound), newMap = newMap.replace(expr, builder.getAffineConstantExpr(*bound),
newMap.getNumDims(), newMap.getNumSymbols()); newMap.getNumDims(), newMap.getNumSymbols());
} }
} }
mlir::canonicalizeMapAndOperands(&newMap, &newOperands); affine::canonicalizeMapAndOperands(&newMap, &newOperands);
return AffineValueMap(newMap, newOperands); return AffineValueMap(newMap, newOperands);
} }

mlir/lib/Dialect/Affine/IR/AffineMemoryOpInterfaces.cpp

//===- AffineMemoryOpInterfaces.cpp ---------------------------------------===// //===- AffineMemoryOpInterfaces.cpp ---------------------------------------===//
// //
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information. // See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// //
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
#include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.h" #include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.h"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// Affine Memory Op Interfaces // Affine Memory Op Interfaces
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
/// Include the definitions of the affine memory op interfaces. /// Include the definitions of the affine memory op interfaces.
#include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.cpp.inc" #include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.cpp.inc"

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

Show All 20 Lines
#include "llvm/ADT/ScopeExit.h" #include "llvm/ADT/ScopeExit.h"
#include "llvm/ADT/SmallBitVector.h" #include "llvm/ADT/SmallBitVector.h"
#include "llvm/ADT/TypeSwitch.h" #include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/Debug.h" #include "llvm/Support/Debug.h"
#include <numeric> #include <numeric>
#include <optional> #include <optional>
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
#define DEBUG_TYPE "affine-ops" #define DEBUG_TYPE "affine-ops"
#include "mlir/Dialect/Affine/IR/AffineOpsDialect.cpp.inc" #include "mlir/Dialect/Affine/IR/AffineOpsDialect.cpp.inc"
/// A utility function to check if a value is defined at the top level of /// A utility function to check if a value is defined at the top level of
/// `region` or is an argument of `region`. A value of index type defined at the /// `region` or is an argument of `region`. A value of index type defined at the
/// top level of a `AffineScope` region is always a valid symbol for all /// top level of a `AffineScope` region is always a valid symbol for all
/// uses in that region. /// uses in that region.
bool mlir::isTopLevelValue(Value value, Region *region) { bool mlir::affine::isTopLevelValue(Value value, Region *region) {
if (auto arg = value.dyn_cast<BlockArgument>()) if (auto arg = value.dyn_cast<BlockArgument>())
return arg.getParentRegion() == region; return arg.getParentRegion() == region;
return value.getDefiningOp()->getParentRegion() == region; return value.getDefiningOp()->getParentRegion() == region;
} }
/// Checks if `value` known to be a legal affine dimension or symbol in `src` /// Checks if `value` known to be a legal affine dimension or symbol in `src`
/// region remains legal if the operation that uses it is inlined into `dest` /// region remains legal if the operation that uses it is inlined into `dest`
/// with the given value mapping. `legalityCheck` is either `isValidDim` or /// with the given value mapping. `legalityCheck` is either `isValidDim` or
▲ Show 20 Lines • Show All 179 Lines • ▼ Show 20 LinesOperation *AffineDialect::materializeConstant(OpBuilder &builder,
Location loc) { Location loc) {
return builder.create<arith::ConstantOp>(loc, type, value); return builder.create<arith::ConstantOp>(loc, type, value);
} }
/// A utility function to check if a value is defined at the top level of an /// A utility function to check if a value is defined at the top level of an
/// op with trait `AffineScope`. If the value is defined in an unlinked region, /// op with trait `AffineScope`. If the value is defined in an unlinked region,
/// conservatively assume it is not top-level. A value of index type defined at /// conservatively assume it is not top-level. A value of index type defined at
/// the top level is always a valid symbol. /// the top level is always a valid symbol.
bool mlir::isTopLevelValue(Value value) { bool mlir::affine::isTopLevelValue(Value value) {
if (auto arg = value.dyn_cast<BlockArgument>()) { if (auto arg = value.dyn_cast<BlockArgument>()) {
// The block owning the argument may be unlinked, e.g. when the surrounding // The block owning the argument may be unlinked, e.g. when the surrounding
// region has not yet been attached to an Op, at which point the parent Op // region has not yet been attached to an Op, at which point the parent Op
// is null. // is null.
Operation *parentOp = arg.getOwner()->getParentOp(); Operation *parentOp = arg.getOwner()->getParentOp();
return parentOp && parentOp->hasTrait<OpTrait::AffineScope>(); return parentOp && parentOp->hasTrait<OpTrait::AffineScope>();
} }
// The defining Op may live in an unlinked block so its parent Op may be null. // The defining Op may live in an unlinked block so its parent Op may be null.
Operation *parentOp = value.getDefiningOp()->getParentOp(); Operation *parentOp = value.getDefiningOp()->getParentOp();
return parentOp && parentOp->hasTrait<OpTrait::AffineScope>(); return parentOp && parentOp->hasTrait<OpTrait::AffineScope>();
} }
/// Returns the closest region enclosing `op` that is held by an operation with /// Returns the closest region enclosing `op` that is held by an operation with
/// trait `AffineScope`; `nullptr` if there is no such region. /// trait `AffineScope`; `nullptr` if there is no such region.
Region *mlir::getAffineScope(Operation *op) { Region *mlir::affine::getAffineScope(Operation *op) {
auto *curOp = op; auto *curOp = op;
while (auto *parentOp = curOp->getParentOp()) { while (auto *parentOp = curOp->getParentOp()) {
if (parentOp->hasTrait<OpTrait::AffineScope>()) if (parentOp->hasTrait<OpTrait::AffineScope>())
return curOp->getParentRegion(); return curOp->getParentRegion();
curOp = parentOp; curOp = parentOp;
} }
return nullptr; return nullptr;
} }
// A Value can be used as a dimension id iff it meets one of the following // A Value can be used as a dimension id iff it meets one of the following
// conditions: // conditions:
// *) It is valid as a symbol. // *) It is valid as a symbol.
// *) It is an induction variable. // *) It is an induction variable.
// *) It is the result of affine apply operation with dimension id arguments. // *) It is the result of affine apply operation with dimension id arguments.
bool mlir::isValidDim(Value value) { bool mlir::affine::isValidDim(Value value) {
// The value must be an index type. // The value must be an index type.
if (!value.getType().isIndex()) if (!value.getType().isIndex())
return false; return false;
if (auto *defOp = value.getDefiningOp()) if (auto *defOp = value.getDefiningOp())
return isValidDim(value, getAffineScope(defOp)); return isValidDim(value, getAffineScope(defOp));
// This value has to be a block argument for an op that has the // This value has to be a block argument for an op that has the
// `AffineScope` trait or for an affine.for or affine.parallel. // `AffineScope` trait or for an affine.for or affine.parallel.
auto *parentOp = value.cast<BlockArgument>().getOwner()->getParentOp(); auto *parentOp = value.cast<BlockArgument>().getOwner()->getParentOp();
return parentOp && (parentOp->hasTrait<OpTrait::AffineScope>() || return parentOp && (parentOp->hasTrait<OpTrait::AffineScope>() ||
isa<AffineForOp, AffineParallelOp>(parentOp)); isa<AffineForOp, AffineParallelOp>(parentOp));
} }
// Value can be used as a dimension id iff it meets one of the following // Value can be used as a dimension id iff it meets one of the following
// conditions: // conditions:
// *) It is valid as a symbol. // *) It is valid as a symbol.
// *) It is an induction variable. // *) It is an induction variable.
// *) It is the result of an affine apply operation with dimension id operands. // *) It is the result of an affine apply operation with dimension id operands.
bool mlir::isValidDim(Value value, Region *region) { bool mlir::affine::isValidDim(Value value, Region *region) {
// The value must be an index type. // The value must be an index type.
if (!value.getType().isIndex()) if (!value.getType().isIndex())
return false; return false;
// All valid symbols are okay. // All valid symbols are okay.
if (isValidSymbol(value, region)) if (isValidSymbol(value, region))
return true; return true;
▲ Show 20 Lines • Show All 60 Lines • ▼ Show 20 Lines
// A value can be used as a symbol (at all its use sites) iff it meets one of // A value can be used as a symbol (at all its use sites) iff it meets one of
// the following conditions: // the following conditions:
// *) It is a constant. // *) It is a constant.
// *) Its defining op or block arg appearance is immediately enclosed by an op // *) Its defining op or block arg appearance is immediately enclosed by an op
// with `AffineScope` trait. // with `AffineScope` trait.
// *) It is the result of an affine.apply operation with symbol operands. // *) It is the result of an affine.apply operation with symbol operands.
// *) It is a result of the dim op on a memref whose corresponding size is a // *) It is a result of the dim op on a memref whose corresponding size is a
// valid symbol. // valid symbol.
bool mlir::isValidSymbol(Value value) { bool mlir::affine::isValidSymbol(Value value) {
if (!value) if (!value)
return false; return false;
// The value must be an index type. // The value must be an index type.
if (!value.getType().isIndex()) if (!value.getType().isIndex())
return false; return false;
// Check that the value is a top level value. // Check that the value is a top level value.
Show All 12 Lines
/// *) It is the result of an affine apply operation with symbol arguments. /// *) It is the result of an affine apply operation with symbol arguments.
/// *) It is a result of the dim op on a memref whose corresponding size is /// *) It is a result of the dim op on a memref whose corresponding size is
/// a valid symbol. /// a valid symbol.
/// *) It is defined at the top level of 'region' or is its argument. /// *) It is defined at the top level of 'region' or is its argument.
/// *) It dominates `region`'s parent op. /// *) It dominates `region`'s parent op.
/// If `region` is null, conservatively assume the symbol definition scope does /// If `region` is null, conservatively assume the symbol definition scope does
/// not exist and only accept the values that would be symbols regardless of /// not exist and only accept the values that would be symbols regardless of
/// the surrounding region structure, i.e. the first three cases above. /// the surrounding region structure, i.e. the first three cases above.
bool mlir::isValidSymbol(Value value, Region *region) { bool mlir::affine::isValidSymbol(Value value, Region *region) {
// The value must be an index type. // The value must be an index type.
if (!value.getType().isIndex()) if (!value.getType().isIndex())
return false; return false;
// A top-level value is a valid symbol. // A top-level value is a valid symbol.
if (region && ::isTopLevelValue(value, region)) if (region && ::isTopLevelValue(value, region))
return true; return true;
▲ Show 20 Lines • Show All 43 Lines • ▼ Show 20 Linesstatic void printDimAndSymbolList(Operation::operand_iterator begin,
unsigned numDims, OpAsmPrinter &printer) { unsigned numDims, OpAsmPrinter &printer) {
OperandRange operands(begin, end); OperandRange operands(begin, end);
printer << '(' << operands.take_front(numDims) << ')'; printer << '(' << operands.take_front(numDims) << ')';
if (operands.size() > numDims) if (operands.size() > numDims)
printer << '[' << operands.drop_front(numDims) << ']'; printer << '[' << operands.drop_front(numDims) << ']';
} }
/// Parses dimension and symbol list and returns true if parsing failed. /// Parses dimension and symbol list and returns true if parsing failed.
ParseResult mlir::parseDimAndSymbolList(OpAsmParser &parser, ParseResult mlir::affine::parseDimAndSymbolList(
SmallVectorImpl<Value> &operands, OpAsmParser &parser, SmallVectorImpl<Value> &operands, unsigned &numDims) {
unsigned &numDims) {
SmallVector<OpAsmParser::UnresolvedOperand, 8> opInfos; SmallVector<OpAsmParser::UnresolvedOperand, 8> opInfos;
if (parser.parseOperandList(opInfos, OpAsmParser::Delimiter::Paren)) if (parser.parseOperandList(opInfos, OpAsmParser::Delimiter::Paren))
return failure(); return failure();
// Store number of dimensions for validation by caller. // Store number of dimensions for validation by caller.
numDims = opInfos.size(); numDims = opInfos.size();
// Parse the optional symbol operands. // Parse the optional symbol operands.
auto indexTy = parser.getBuilder().getIndexType(); auto indexTy = parser.getBuilder().getIndexType();
▲ Show 20 Lines • Show All 75 Lines • ▼ Show 20 LinesLogicalResult AffineApplyOp::verify() {
return success(); return success();
} }
// The result of the affine apply operation can be used as a dimension id if all // The result of the affine apply operation can be used as a dimension id if all
// its operands are valid dimension ids. // its operands are valid dimension ids.
bool AffineApplyOp::isValidDim() { bool AffineApplyOp::isValidDim() {
return llvm::all_of(getOperands(), return llvm::all_of(getOperands(),
[](Value op) { return mlir::isValidDim(op); }); [](Value op) { return affine::isValidDim(op); });
} }
// The result of the affine apply operation can be used as a dimension id if all // The result of the affine apply operation can be used as a dimension id if all
// its operands are valid dimension ids with the parent operation of `region` // its operands are valid dimension ids with the parent operation of `region`
// defining the polyhedral scope for symbols. // defining the polyhedral scope for symbols.
bool AffineApplyOp::isValidDim(Region *region) { bool AffineApplyOp::isValidDim(Region *region) {
return llvm::all_of(getOperands(), return llvm::all_of(getOperands(),
[&](Value op) { return ::isValidDim(op, region); }); [&](Value op) { return ::isValidDim(op, region); });
} }
// The result of the affine apply operation can be used as a symbol if all its // The result of the affine apply operation can be used as a symbol if all its
// operands are symbols. // operands are symbols.
bool AffineApplyOp::isValidSymbol() { bool AffineApplyOp::isValidSymbol() {
return llvm::all_of(getOperands(), return llvm::all_of(getOperands(),
[](Value op) { return mlir::isValidSymbol(op); }); [](Value op) { return affine::isValidSymbol(op); });
} }
// The result of the affine apply operation can be used as a symbol in `region` // The result of the affine apply operation can be used as a symbol in `region`
// if all its operands are symbols in `region`. // if all its operands are symbols in `region`.
bool AffineApplyOp::isValidSymbol(Region *region) { bool AffineApplyOp::isValidSymbol(Region *region) {
return llvm::all_of(getOperands(), [&](Value operand) { return llvm::all_of(getOperands(), [&](Value operand) {
return mlir::isValidSymbol(operand, region); return affine::isValidSymbol(operand, region);
}); });
} }
OpFoldResult AffineApplyOp::fold(FoldAdaptor adaptor) { OpFoldResult AffineApplyOp::fold(FoldAdaptor adaptor) {
auto map = getAffineMap(); auto map = getAffineMap();
// Fold dims and symbols to existing values. // Fold dims and symbols to existing values.
auto expr = map.getResult(0); auto expr = map.getResult(0);
▲ Show 20 Lines • Show All 491 Lines • ▼ Show 20 Linesstatic void composeAffineMapAndOperands(AffineMap *map,
*map = map->replaceDimsAndSymbols(dimReplacements, symReplacements, nDims, *map = map->replaceDimsAndSymbols(dimReplacements, symReplacements, nDims,
nSyms); nSyms);
// Canonicalize and simplify before returning. // Canonicalize and simplify before returning.
canonicalizeMapAndOperands(map, operands); canonicalizeMapAndOperands(map, operands);
*map = simplifyAffineMap(*map); *map = simplifyAffineMap(*map);
} }
void mlir::fullyComposeAffineMapAndOperands(AffineMap *map, void mlir::affine::fullyComposeAffineMapAndOperands(
SmallVectorImpl<Value> *operands) { AffineMap *map, SmallVectorImpl<Value> *operands) {
while (llvm::any_of(*operands, [](Value v) { while (llvm::any_of(*operands, [](Value v) {
return isa_and_nonnull<AffineApplyOp>(v.getDefiningOp()); return isa_and_nonnull<AffineApplyOp>(v.getDefiningOp());
})) { })) {
composeAffineMapAndOperands(map, operands); composeAffineMapAndOperands(map, operands);
} }
} }
/// Given a list of `OpFoldResult`, build the necessary operations to populate /// Given a list of `OpFoldResult`, build the necessary operations to populate
▲ Show 20 Lines • Show All 79 Lines • ▼ Show 20 LinescreateOrFold(OpBuilder &b, Location loc, ValueRange operands,
// Notify the listener now that we definitely know that the operation will // Notify the listener now that we definitely know that the operation will
// persist. Use the original listener stored in the variable. // persist. Use the original listener stored in the variable.
if (listener) if (listener)
listener->notifyOperationInserted(op); listener->notifyOperationInserted(op);
return op->getResult(0); return op->getResult(0);
} }
AffineApplyOp mlir::makeComposedAffineApply(OpBuilder &b, Location loc, AffineApplyOp mlir::affine::makeComposedAffineApply(OpBuilder &b, Location loc,
AffineMap map, AffineMap map,
ValueRange operands) { ValueRange operands) {
AffineMap normalizedMap = map; AffineMap normalizedMap = map;
SmallVector<Value, 8> normalizedOperands(operands.begin(), operands.end()); SmallVector<Value, 8> normalizedOperands(operands.begin(), operands.end());
composeAffineMapAndOperands(&normalizedMap, &normalizedOperands); composeAffineMapAndOperands(&normalizedMap, &normalizedOperands);
assert(normalizedMap); assert(normalizedMap);
return b.create<AffineApplyOp>(loc, normalizedMap, normalizedOperands); return b.create<AffineApplyOp>(loc, normalizedMap, normalizedOperands);
} }
AffineApplyOp mlir::makeComposedAffineApply(OpBuilder &b, Location loc, AffineApplyOp mlir::affine::makeComposedAffineApply(OpBuilder &b, Location loc,
AffineExpr e, ValueRange values) { AffineExpr e,
ValueRange values) {
return makeComposedAffineApply( return makeComposedAffineApply(
b, loc, AffineMap::inferFromExprList(ArrayRef<AffineExpr>{e}).front(), b, loc, AffineMap::inferFromExprList(ArrayRef<AffineExpr>{e}).front(),
values); values);
} }
/// Composes the given affine map with the given list of operands, pulling in /// Composes the given affine map with the given list of operands, pulling in
/// the maps from any affine.apply operations that supply the operands. /// the maps from any affine.apply operations that supply the operands.
static void composeMultiResultAffineMap(AffineMap &map, static void composeMultiResultAffineMap(AffineMap &map,
Show All 20 Linesstatic void composeMultiResultAffineMap(AffineMap &map,
// Canonicalize the map created from composed expressions to deduplicate the // Canonicalize the map created from composed expressions to deduplicate the
// dimension and symbol operands. // dimension and symbol operands.
operands = llvm::to_vector(llvm::concat<Value>(dims, symbols)); operands = llvm::to_vector(llvm::concat<Value>(dims, symbols));
map = AffineMap::get(dims.size(), symbols.size(), exprs, map.getContext()); map = AffineMap::get(dims.size(), symbols.size(), exprs, map.getContext());
canonicalizeMapAndOperands(&map, &operands); canonicalizeMapAndOperands(&map, &operands);
} }
OpFoldResult OpFoldResult
mlir::makeComposedFoldedAffineApply(OpBuilder &b, Location loc, AffineMap map, mlir::affine::makeComposedFoldedAffineApply(OpBuilder &b, Location loc,
AffineMap map,
ArrayRef<OpFoldResult> operands) { ArrayRef<OpFoldResult> operands) {
assert(map.getNumResults() == 1 && "building affine.apply with !=1 result"); assert(map.getNumResults() == 1 && "building affine.apply with !=1 result");
SmallVector<Operation *> constants; SmallVector<Operation *> constants;
SmallVector<Value> actualValues; SmallVector<Value> actualValues;
materializeConstants(b, loc, operands, constants, actualValues); materializeConstants(b, loc, operands, constants, actualValues);
composeAffineMapAndOperands(&map, &actualValues); composeAffineMapAndOperands(&map, &actualValues);
OpFoldResult result = createOrFold<AffineApplyOp>(b, loc, actualValues, map); OpFoldResult result = createOrFold<AffineApplyOp>(b, loc, actualValues, map);
// Constants are always folded into affine min/max because they can be // Constants are always folded into affine min/max because they can be
// represented as constant expressions, so delete them. // represented as constant expressions, so delete them.
for (Operation *op : constants) for (Operation *op : constants)
op->erase(); op->erase();
return result; return result;
} }
OpFoldResult OpFoldResult
mlir::makeComposedFoldedAffineApply(OpBuilder &b, Location loc, AffineExpr expr, mlir::affine::makeComposedFoldedAffineApply(OpBuilder &b, Location loc,
AffineExpr expr,
ArrayRef<OpFoldResult> operands) { ArrayRef<OpFoldResult> operands) {
return makeComposedFoldedAffineApply( return makeComposedFoldedAffineApply(
b, loc, AffineMap::inferFromExprList(ArrayRef<AffineExpr>{expr}).front(), b, loc, AffineMap::inferFromExprList(ArrayRef<AffineExpr>{expr}).front(),
operands); operands);
} }
SmallVector<OpFoldResult> mlir::makeComposedFoldedMultiResultAffineApply( SmallVector<OpFoldResult>
mlir::affine::makeComposedFoldedMultiResultAffineApply(
OpBuilder &b, Location loc, AffineMap map, OpBuilder &b, Location loc, AffineMap map,
ArrayRef<OpFoldResult> operands) { ArrayRef<OpFoldResult> operands) {
return llvm::to_vector(llvm::map_range( return llvm::to_vector(llvm::map_range(
llvm::seq<unsigned>(0, map.getNumResults()), [&](unsigned i) { llvm::seq<unsigned>(0, map.getNumResults()), [&](unsigned i) {
return makeComposedFoldedAffineApply(b, loc, map.getSubMap({i}), return makeComposedFoldedAffineApply(b, loc, map.getSubMap({i}),
operands); operands);
})); }));
} }
Value mlir::makeComposedAffineMin(OpBuilder &b, Location loc, AffineMap map, Value mlir::affine::makeComposedAffineMin(OpBuilder &b, Location loc,
ValueRange operands) { AffineMap map, ValueRange operands) {
SmallVector<Value> allOperands = llvm::to_vector(operands); SmallVector<Value> allOperands = llvm::to_vector(operands);
composeMultiResultAffineMap(map, allOperands); composeMultiResultAffineMap(map, allOperands);
return b.createOrFold<AffineMinOp>(loc, b.getIndexType(), map, allOperands); return b.createOrFold<AffineMinOp>(loc, b.getIndexType(), map, allOperands);
} }
template <typename OpTy> template <typename OpTy>
static OpFoldResult makeComposedFoldedMinMax(OpBuilder &b, Location loc, static OpFoldResult makeComposedFoldedMinMax(OpBuilder &b, Location loc,
AffineMap map, AffineMap map,
ArrayRef<OpFoldResult> operands) { ArrayRef<OpFoldResult> operands) {
SmallVector<Operation *> constants; SmallVector<Operation *> constants;
SmallVector<Value> actualValues; SmallVector<Value> actualValues;
materializeConstants(b, loc, operands, constants, actualValues); materializeConstants(b, loc, operands, constants, actualValues);
composeMultiResultAffineMap(map, actualValues); composeMultiResultAffineMap(map, actualValues);
OpFoldResult result = OpFoldResult result =
createOrFold<OpTy>(b, loc, actualValues, b.getIndexType(), map); createOrFold<OpTy>(b, loc, actualValues, b.getIndexType(), map);
// Constants are always folded into affine min/max because they can be // Constants are always folded into affine min/max because they can be
// represented as constant expressions, so delete them. // represented as constant expressions, so delete them.
for (Operation *op : constants) for (Operation *op : constants)
op->erase(); op->erase();
return result; return result;
} }
OpFoldResult OpFoldResult
mlir::makeComposedFoldedAffineMin(OpBuilder &b, Location loc, AffineMap map, mlir::affine::makeComposedFoldedAffineMin(OpBuilder &b, Location loc,
AffineMap map,
ArrayRef<OpFoldResult> operands) { ArrayRef<OpFoldResult> operands) {
return makeComposedFoldedMinMax<AffineMinOp>(b, loc, map, operands); return makeComposedFoldedMinMax<AffineMinOp>(b, loc, map, operands);
} }
OpFoldResult OpFoldResult
mlir::makeComposedFoldedAffineMax(OpBuilder &b, Location loc, AffineMap map, mlir::affine::makeComposedFoldedAffineMax(OpBuilder &b, Location loc,
AffineMap map,
ArrayRef<OpFoldResult> operands) { ArrayRef<OpFoldResult> operands) {
return makeComposedFoldedMinMax<AffineMaxOp>(b, loc, map, operands); return makeComposedFoldedMinMax<AffineMaxOp>(b, loc, map, operands);
} }
/// Fully compose map with operands and canonicalize the result. /// Fully compose map with operands and canonicalize the result.
/// Return the `createOrFold`'ed AffineApply op. /// Return the `createOrFold`'ed AffineApply op.
static Value createFoldedComposedAffineApply(OpBuilder &b, Location loc, static Value createFoldedComposedAffineApply(OpBuilder &b, Location loc,
AffineMap map, AffineMap map,
ValueRange operandsRef) { ValueRange operandsRef) {
SmallVector<Value, 4> operands(operandsRef.begin(), operandsRef.end()); SmallVector<Value, 4> operands(operandsRef.begin(), operandsRef.end());
fullyComposeAffineMapAndOperands(&map, &operands); fullyComposeAffineMapAndOperands(&map, &operands);
canonicalizeMapAndOperands(&map, &operands); canonicalizeMapAndOperands(&map, &operands);
return b.createOrFold<AffineApplyOp>(loc, map, operands); return b.createOrFold<AffineApplyOp>(loc, map, operands);
} }
SmallVector<Value, 4> mlir::applyMapToValues(OpBuilder &b, Location loc, SmallVector<Value, 4> mlir::affine::applyMapToValues(OpBuilder &b, Location loc,
AffineMap map, ValueRange values) { AffineMap map,
ValueRange values) {
SmallVector<Value, 4> res; SmallVector<Value, 4> res;
res.reserve(map.getNumResults()); res.reserve(map.getNumResults());
unsigned numDims = map.getNumDims(), numSym = map.getNumSymbols(); unsigned numDims = map.getNumDims(), numSym = map.getNumSymbols();
// For each `expr` in `map`, applies the `expr` to the values extracted from // For each `expr` in `map`, applies the `expr` to the values extracted from
// ranges. If the resulting application can be folded into a Value, the // ranges. If the resulting application can be folded into a Value, the
// folding occurs eagerly. // folding occurs eagerly.
for (auto expr : map.getResults()) { for (auto expr : map.getResults()) {
AffineMap map = AffineMap::get(numDims, numSym, expr); AffineMap map = AffineMap::get(numDims, numSym, expr);
▲ Show 20 Lines • Show All 119 Lines • ▼ Show 20 Lines if (it == seenSymbols.end()) {
symRemapping[i] = it->second; symRemapping[i] = it->second;
} }
} }
*mapOrSet = mapOrSet->replaceDimsAndSymbols(dimRemapping, symRemapping, *mapOrSet = mapOrSet->replaceDimsAndSymbols(dimRemapping, symRemapping,
nextDim, nextSym); nextDim, nextSym);
*operands = resultOperands; *operands = resultOperands;
} }
void mlir::canonicalizeMapAndOperands(AffineMap *map, void mlir::affine::canonicalizeMapAndOperands(
SmallVectorImpl<Value> *operands) { AffineMap *map, SmallVectorImpl<Value> *operands) {
canonicalizeMapOrSetAndOperands<AffineMap>(map, operands); canonicalizeMapOrSetAndOperands<AffineMap>(map, operands);
} }
void mlir::canonicalizeSetAndOperands(IntegerSet *set, void mlir::affine::canonicalizeSetAndOperands(
SmallVectorImpl<Value> *operands) { IntegerSet *set, SmallVectorImpl<Value> *operands) {
canonicalizeMapOrSetAndOperands<IntegerSet>(set, operands); canonicalizeMapOrSetAndOperands<IntegerSet>(set, operands);
} }
namespace { namespace {
/// Simplify AffineApply, AffineLoad, and AffineStore operations by composing /// Simplify AffineApply, AffineLoad, and AffineStore operations by composing
/// maps that supply results into them. /// maps that supply results into them.
/// ///
template <typename AffineOpTy> template <typename AffineOpTy>
▲ Show 20 Lines • Show All 1,059 Lines • ▼ Show 20 LinesSpeculation::Speculatability AffineForOp::getSpeculatability() {
// For Step != 1, the loop may not terminate. We can add more smarts here if // For Step != 1, the loop may not terminate. We can add more smarts here if
// needed. // needed.
return getStep() == 1 ? Speculation::RecursivelySpeculatable return getStep() == 1 ? Speculation::RecursivelySpeculatable
: Speculation::NotSpeculatable; : Speculation::NotSpeculatable;
} }
/// Returns true if the provided value is the induction variable of a /// Returns true if the provided value is the induction variable of a
/// AffineForOp. /// AffineForOp.
bool mlir::isAffineForInductionVar(Value val) { bool mlir::affine::isAffineForInductionVar(Value val) {
return getForInductionVarOwner(val) != AffineForOp(); return getForInductionVarOwner(val) != AffineForOp();
} }
bool mlir::isAffineParallelInductionVar(Value val) { bool mlir::affine::isAffineParallelInductionVar(Value val) {
return getAffineParallelInductionVarOwner(val) != nullptr; return getAffineParallelInductionVarOwner(val) != nullptr;
} }
bool mlir::isAffineInductionVar(Value val) { bool mlir::affine::isAffineInductionVar(Value val) {
return isAffineForInductionVar(val) || isAffineParallelInductionVar(val); return isAffineForInductionVar(val) || isAffineParallelInductionVar(val);
} }
AffineForOp mlir::getForInductionVarOwner(Value val) { AffineForOp mlir::affine::getForInductionVarOwner(Value val) {
auto ivArg = val.dyn_cast<BlockArgument>(); auto ivArg = val.dyn_cast<BlockArgument>();
if (!ivArg || !ivArg.getOwner()) if (!ivArg || !ivArg.getOwner())
return AffineForOp(); return AffineForOp();
auto *containingInst = ivArg.getOwner()->getParent()->getParentOp(); auto *containingInst = ivArg.getOwner()->getParent()->getParentOp();
if (auto forOp = dyn_cast<AffineForOp>(containingInst)) if (auto forOp = dyn_cast<AffineForOp>(containingInst))
// Check to make sure `val` is the induction variable, not an iter_arg. // Check to make sure `val` is the induction variable, not an iter_arg.
return forOp.getInductionVar() == val ? forOp : AffineForOp(); return forOp.getInductionVar() == val ? forOp : AffineForOp();
return AffineForOp(); return AffineForOp();
} }
AffineParallelOp mlir::getAffineParallelInductionVarOwner(Value val) { AffineParallelOp mlir::affine::getAffineParallelInductionVarOwner(Value val) {
auto ivArg = val.dyn_cast<BlockArgument>(); auto ivArg = val.dyn_cast<BlockArgument>();
if (!ivArg || !ivArg.getOwner()) if (!ivArg || !ivArg.getOwner())
return nullptr; return nullptr;
Operation *containingOp = ivArg.getOwner()->getParentOp(); Operation *containingOp = ivArg.getOwner()->getParentOp();
auto parallelOp = dyn_cast<AffineParallelOp>(containingOp); auto parallelOp = dyn_cast<AffineParallelOp>(containingOp);
if (parallelOp && llvm::is_contained(parallelOp.getIVs(), val)) if (parallelOp && llvm::is_contained(parallelOp.getIVs(), val))
return parallelOp; return parallelOp;
return nullptr; return nullptr;
} }
/// Extracts the induction variables from a list of AffineForOps and returns /// Extracts the induction variables from a list of AffineForOps and returns
/// them. /// them.
void mlir::extractForInductionVars(ArrayRef<AffineForOp> forInsts, void mlir::affine::extractForInductionVars(ArrayRef<AffineForOp> forInsts,
SmallVectorImpl<Value> *ivs) { SmallVectorImpl<Value> *ivs) {
ivs->reserve(forInsts.size()); ivs->reserve(forInsts.size());
for (auto forInst : forInsts) for (auto forInst : forInsts)
ivs->push_back(forInst.getInductionVar()); ivs->push_back(forInst.getInductionVar());
} }
void mlir::extractInductionVars(ArrayRef<mlir::Operation *> affineOps, void mlir::affine::extractInductionVars(ArrayRef<mlir::Operation *> affineOps,
SmallVectorImpl<mlir::Value> &ivs) { SmallVectorImpl<mlir::Value> &ivs) {
ivs.reserve(affineOps.size()); ivs.reserve(affineOps.size());
for (Operation *op : affineOps) { for (Operation *op : affineOps) {
// Add constraints from forOp's bounds. // Add constraints from forOp's bounds.
if (auto forOp = dyn_cast<AffineForOp>(op)) if (auto forOp = dyn_cast<AffineForOp>(op))
ivs.push_back(forOp.getInductionVar()); ivs.push_back(forOp.getInductionVar());
else if (auto parallelOp = dyn_cast<AffineParallelOp>(op)) else if (auto parallelOp = dyn_cast<AffineParallelOp>(op))
for (size_t i = 0; i < parallelOp.getBody()->getNumArguments(); i++) for (size_t i = 0; i < parallelOp.getBody()->getNumArguments(); i++)
ivs.push_back(parallelOp.getBody()->getArgument(i)); ivs.push_back(parallelOp.getBody()->getArgument(i));
▲ Show 20 Lines • Show All 61 Lines • ▼ Show 20 LinesbuildAffineLoopFromValues(OpBuilder &builder, Location loc, Value lb, Value ub,
if (lbConst && ubConst) if (lbConst && ubConst)
return buildAffineLoopFromConstants(builder, loc, lbConst.value(), return buildAffineLoopFromConstants(builder, loc, lbConst.value(),
ubConst.value(), step, bodyBuilderFn); ubConst.value(), step, bodyBuilderFn);
return builder.create<AffineForOp>(loc, lb, builder.getDimIdentityMap(), ub, return builder.create<AffineForOp>(loc, lb, builder.getDimIdentityMap(), ub,
builder.getDimIdentityMap(), step, builder.getDimIdentityMap(), step,
/*iterArgs=*/std::nullopt, bodyBuilderFn); /*iterArgs=*/std::nullopt, bodyBuilderFn);
} }
void mlir::buildAffineLoopNest( void mlir::affine::buildAffineLoopNest(
OpBuilder &builder, Location loc, ArrayRef<int64_t> lbs, OpBuilder &builder, Location loc, ArrayRef<int64_t> lbs,
ArrayRef<int64_t> ubs, ArrayRef<int64_t> steps, ArrayRef<int64_t> ubs, ArrayRef<int64_t> steps,
function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn) { function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn) {
buildAffineLoopNestImpl(builder, loc, lbs, ubs, steps, bodyBuilderFn, buildAffineLoopNestImpl(builder, loc, lbs, ubs, steps, bodyBuilderFn,
buildAffineLoopFromConstants); buildAffineLoopFromConstants);
} }
void mlir::buildAffineLoopNest( void mlir::affine::buildAffineLoopNest(
OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs, OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs,
ArrayRef<int64_t> steps, ArrayRef<int64_t> steps,
function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn) { function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn) {
buildAffineLoopNestImpl(builder, loc, lbs, ubs, steps, bodyBuilderFn, buildAffineLoopNestImpl(builder, loc, lbs, ubs, steps, bodyBuilderFn,
buildAffineLoopFromValues); buildAffineLoopFromValues);
} }
AffineForOp mlir::replaceForOpWithNewYields(OpBuilder &b, AffineForOp loop, AffineForOp mlir::affine::replaceForOpWithNewYields(OpBuilder &b,
AffineForOp loop,
ValueRange newIterOperands, ValueRange newIterOperands,
ValueRange newYieldedValues, ValueRange newYieldedValues,
ValueRange newIterArgs, ValueRange newIterArgs,
bool replaceLoopResults) { bool replaceLoopResults) {
assert(newIterOperands.size() == newYieldedValues.size() && assert(newIterOperands.size() == newYieldedValues.size() &&
"newIterOperands must be of the same size as newYieldedValues"); "newIterOperands must be of the same size as newYieldedValues");
// Create a new loop before the existing one, with the extra operands. // Create a new loop before the existing one, with the extra operands.
OpBuilder::InsertionGuard g(b); OpBuilder::InsertionGuard g(b);
b.setInsertionPoint(loop); b.setInsertionPoint(loop);
auto operands = llvm::to_vector<4>(loop.getIterOperands()); auto operands = llvm::to_vector<4>(loop.getIterOperands());
operands.append(newIterOperands.begin(), newIterOperands.end()); operands.append(newIterOperands.begin(), newIterOperands.end());
SmallVector<Value, 4> lbOperands(loop.getLowerBoundOperands()); SmallVector<Value, 4> lbOperands(loop.getLowerBoundOperands());
▲ Show 20 Lines • Show All 1,825 Lines • Show Last 20 Lines

mlir/lib/Dialect/Affine/IR/AffineValueMap.cpp

//===- AffineValueMap.cpp - MLIR Affine Value Map Class -------------------===// //===- AffineValueMap.cpp - MLIR Affine Value Map Class -------------------===//
// //
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information. // See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// //
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
#include "mlir/Dialect/Affine/IR/AffineValueMap.h" #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Affine/IR/AffineOps.h"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
AffineValueMap::AffineValueMap(AffineMap map, ValueRange operands, AffineValueMap::AffineValueMap(AffineMap map, ValueRange operands,
ValueRange results) ValueRange results)
: map(map), operands(operands.begin(), operands.end()), : map(map), operands(operands.begin(), operands.end()),
results(results.begin(), results.end()) {} results(results.begin(), results.end()) {}
void AffineValueMap::reset(AffineMap map, ValueRange operands, void AffineValueMap::reset(AffineMap map, ValueRange operands,
ValueRange results) { ValueRange results) {
▲ Show 20 Lines • Show All 82 Lines • Show Last 20 Lines

mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp

//===- ValueBoundsOpInterfaceImpl.cpp - Impl. of ValueBoundsOpInterface ---===// //===- ValueBoundsOpInterfaceImpl.cpp - Impl. of ValueBoundsOpInterface ---===//
// //
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information. // See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// //
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
#include "mlir/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.h" #include "mlir/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Interfaces/ValueBoundsOpInterface.h" #include "mlir/Interfaces/ValueBoundsOpInterface.h"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
namespace mlir { namespace mlir {
namespace { namespace {
struct AffineApplyOpInterface struct AffineApplyOpInterface
: public ValueBoundsOpInterface::ExternalModel<AffineApplyOpInterface, : public ValueBoundsOpInterface::ExternalModel<AffineApplyOpInterface,
AffineApplyOp> { AffineApplyOp> {
void populateBoundsForIndexValue(Operation *op, Value value, void populateBoundsForIndexValue(Operation *op, Value value,
▲ Show 20 Lines • Show All 71 Lines • Show Last 20 Lines

mlir/lib/Dialect/Affine/TransformOps/AffineTransformOps.cpp

Show All 13 Lines
#include "mlir/Dialect/Affine/LoopUtils.h" #include "mlir/Dialect/Affine/LoopUtils.h"
#include "mlir/Dialect/PDL/IR/PDL.h" #include "mlir/Dialect/PDL/IR/PDL.h"
#include "mlir/Dialect/PDL/IR/PDLTypes.h" #include "mlir/Dialect/PDL/IR/PDLTypes.h"
#include "mlir/Dialect/Transform/IR/TransformDialect.h" #include "mlir/Dialect/Transform/IR/TransformDialect.h"
#include "mlir/Dialect/Transform/IR/TransformInterfaces.h" #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
using namespace mlir::transform; using namespace mlir::transform;
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// SimplifyBoundedAffineOpsOp // SimplifyBoundedAffineOpsOp
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
LogicalResult SimplifyBoundedAffineOpsOp::verify() { LogicalResult SimplifyBoundedAffineOpsOp::verify() {
if (getLowerBounds().size() != getBoundedValues().size()) if (getLowerBounds().size() != getBoundedValues().size())
▲ Show 20 Lines • Show All 148 Lines • Show Last 20 Lines

mlir/lib/Dialect/Affine/Transforms/AffineDataCopyGeneration.cpp

Show All 29 Lines
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/MapVector.h" #include "llvm/ADT/MapVector.h"
#include "llvm/Support/CommandLine.h" #include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h" #include "llvm/Support/Debug.h"
#include <algorithm> #include <algorithm>
#include <optional> #include <optional>
namespace mlir { namespace mlir {
namespace affine {
#define GEN_PASS_DEF_AFFINEDATACOPYGENERATION #define GEN_PASS_DEF_AFFINEDATACOPYGENERATION
#include "mlir/Dialect/Affine/Passes.h.inc" #include "mlir/Dialect/Affine/Passes.h.inc"
} // namespace affine
} // namespace mlir } // namespace mlir
#define DEBUG_TYPE "affine-data-copy-generate" #define DEBUG_TYPE "affine-data-copy-generate"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
namespace { namespace {
/// Replaces all loads and stores on memref's living in 'slowMemorySpace' by /// Replaces all loads and stores on memref's living in 'slowMemorySpace' by
/// introducing copy operations to transfer data into `fastMemorySpace` and /// introducing copy operations to transfer data into `fastMemorySpace` and
/// rewriting the original load's/store's to instead load/store from the /// rewriting the original load's/store's to instead load/store from the
/// allocated fast memory buffers. Additional options specify the identifier /// allocated fast memory buffers. Additional options specify the identifier
/// corresponding to the fast memory space and the amount of fast memory space /// corresponding to the fast memory space and the amount of fast memory space
/// available. The pass traverses through the nesting structure, recursing to /// available. The pass traverses through the nesting structure, recursing to
/// inner levels if necessary to determine at what depth copies need to be /// inner levels if necessary to determine at what depth copies need to be
/// placed so that the allocated buffers fit within the memory capacity /// placed so that the allocated buffers fit within the memory capacity
/// provided. /// provided.
// TODO: We currently can't generate copies correctly when stores // TODO: We currently can't generate copies correctly when stores
// are strided. Check for strided stores. // are strided. Check for strided stores.
struct AffineDataCopyGeneration struct AffineDataCopyGeneration
: public impl::AffineDataCopyGenerationBase<AffineDataCopyGeneration> { : public affine::impl::AffineDataCopyGenerationBase<
AffineDataCopyGeneration> {
AffineDataCopyGeneration() = default; AffineDataCopyGeneration() = default;
explicit AffineDataCopyGeneration(unsigned slowMemorySpace, explicit AffineDataCopyGeneration(unsigned slowMemorySpace,
unsigned fastMemorySpace, unsigned fastMemorySpace,
unsigned tagMemorySpace, unsigned tagMemorySpace,
int minDmaTransferSize, int minDmaTransferSize,
uint64_t fastMemCapacityBytes) { uint64_t fastMemCapacityBytes) {
this->slowMemorySpace = slowMemorySpace; this->slowMemorySpace = slowMemorySpace;
this->fastMemorySpace = fastMemorySpace; this->fastMemorySpace = fastMemorySpace;
Show All 11 Lines
} // namespace } // namespace
/// Generates copies for memref's living in 'slowMemorySpace' into newly created /// Generates copies for memref's living in 'slowMemorySpace' into newly created
/// buffers in 'fastMemorySpace', and replaces memory operations to the former /// buffers in 'fastMemorySpace', and replaces memory operations to the former
/// by the latter. Only load op's handled for now. /// by the latter. Only load op's handled for now.
/// TODO: extend this to store op's. /// TODO: extend this to store op's.
std::unique_ptr<OperationPass<func::FuncOp>> std::unique_ptr<OperationPass<func::FuncOp>>
mlir::createAffineDataCopyGenerationPass(unsigned slowMemorySpace, mlir::affine::createAffineDataCopyGenerationPass(
unsigned fastMemorySpace, unsigned slowMemorySpace, unsigned fastMemorySpace, unsigned tagMemorySpace,
unsigned tagMemorySpace, int minDmaTransferSize, uint64_t fastMemCapacityBytes) {
int minDmaTransferSize,
uint64_t fastMemCapacityBytes) {
return std::make_unique<AffineDataCopyGeneration>( return std::make_unique<AffineDataCopyGeneration>(
slowMemorySpace, fastMemorySpace, tagMemorySpace, minDmaTransferSize, slowMemorySpace, fastMemorySpace, tagMemorySpace, minDmaTransferSize,
fastMemCapacityBytes); fastMemCapacityBytes);
} }
std::unique_ptr<OperationPass<func::FuncOp>> std::unique_ptr<OperationPass<func::FuncOp>>
mlir::createAffineDataCopyGenerationPass() { mlir::affine::createAffineDataCopyGenerationPass() {
return std::make_unique<AffineDataCopyGeneration>(); return std::make_unique<AffineDataCopyGeneration>();
} }
/// Generate copies for this block. The block is partitioned into separate /// Generate copies for this block. The block is partitioned into separate
/// ranges: each range is either a sequence of one or more operations starting /// ranges: each range is either a sequence of one or more operations starting
/// and ending with an affine load or store op, or just an affine.forop (which /// and ending with an affine load or store op, or just an affine.forop (which
/// could have other affine for op's nested within). /// could have other affine for op's nested within).
void AffineDataCopyGeneration::runOnBlock(Block *block, void AffineDataCopyGeneration::runOnBlock(Block *block,
▲ Show 20 Lines • Show All 139 Lines • Show Last 20 Lines

mlir/lib/Dialect/Affine/Transforms/AffineExpandIndexOps.cpp

Show All 12 Lines
#include "mlir/Dialect/Affine/Passes.h" #include "mlir/Dialect/Affine/Passes.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/Transforms/Transforms.h" #include "mlir/Dialect/Affine/Transforms/Transforms.h"
#include "mlir/Dialect/Affine/Utils.h" #include "mlir/Dialect/Affine/Utils.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
namespace mlir { namespace mlir {
namespace affine {
#define GEN_PASS_DEF_AFFINEEXPANDINDEXOPS #define GEN_PASS_DEF_AFFINEEXPANDINDEXOPS
#include "mlir/Dialect/Affine/Passes.h.inc" #include "mlir/Dialect/Affine/Passes.h.inc"
} // namespace affine
} // namespace mlir } // namespace mlir
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
namespace { namespace {
/// Lowers `affine.delinearize_index` into a sequence of division and remainder /// Lowers `affine.delinearize_index` into a sequence of division and remainder
/// operations. /// operations.
struct LowerDelinearizeIndexOps struct LowerDelinearizeIndexOps
: public OpRewritePattern<AffineDelinearizeIndexOp> { : public OpRewritePattern<AffineDelinearizeIndexOp> {
using OpRewritePattern<AffineDelinearizeIndexOp>::OpRewritePattern; using OpRewritePattern<AffineDelinearizeIndexOp>::OpRewritePattern;
LogicalResult matchAndRewrite(AffineDelinearizeIndexOp op, LogicalResult matchAndRewrite(AffineDelinearizeIndexOp op,
PatternRewriter &rewriter) const override { PatternRewriter &rewriter) const override {
FailureOr<SmallVector<Value>> multiIndex = FailureOr<SmallVector<Value>> multiIndex =
delinearizeIndex(rewriter, op->getLoc(), op.getLinearIndex(), delinearizeIndex(rewriter, op->getLoc(), op.getLinearIndex(),
llvm::to_vector(op.getBasis())); llvm::to_vector(op.getBasis()));
if (failed(multiIndex)) if (failed(multiIndex))
return failure(); return failure();
rewriter.replaceOp(op, *multiIndex); rewriter.replaceOp(op, *multiIndex);
return success(); return success();
} }
}; };
class ExpandAffineIndexOpsPass class ExpandAffineIndexOpsPass
: public impl::AffineExpandIndexOpsBase<ExpandAffineIndexOpsPass> { : public affine::impl::AffineExpandIndexOpsBase<ExpandAffineIndexOpsPass> {
public: public:
ExpandAffineIndexOpsPass() = default; ExpandAffineIndexOpsPass() = default;
void runOnOperation() override { void runOnOperation() override {
MLIRContext *context = &getContext(); MLIRContext *context = &getContext();
RewritePatternSet patterns(context); RewritePatternSet patterns(context);
populateAffineExpandIndexOpsPatterns(patterns); populateAffineExpandIndexOpsPatterns(patterns);
if (failed( if (failed(
applyPatternsAndFoldGreedily(getOperation(), std::move(patterns)))) applyPatternsAndFoldGreedily(getOperation(), std::move(patterns))))
return signalPassFailure(); return signalPassFailure();
} }
}; };
} // namespace } // namespace
void mlir::populateAffineExpandIndexOpsPatterns(RewritePatternSet &patterns) { void mlir::affine::populateAffineExpandIndexOpsPatterns(
RewritePatternSet &patterns) {
patterns.insert<LowerDelinearizeIndexOps>(patterns.getContext()); patterns.insert<LowerDelinearizeIndexOps>(patterns.getContext());
} }
std::unique_ptr<Pass> mlir::createAffineExpandIndexOpsPass() { std::unique_ptr<Pass> mlir::affine::createAffineExpandIndexOpsPass() {
return std::make_unique<ExpandAffineIndexOpsPass>(); return std::make_unique<ExpandAffineIndexOpsPass>();
} }

mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp

Show All 28 Lines
#include "llvm/ADT/DenseMap.h" #include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/DenseSet.h" #include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/SmallPtrSet.h" #include "llvm/ADT/SmallPtrSet.h"
#include "llvm/Support/CommandLine.h" #include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h" #include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h" #include "llvm/Support/raw_ostream.h"
namespace mlir { namespace mlir {
namespace affine {
#define GEN_PASS_DEF_AFFINELOOPINVARIANTCODEMOTION #define GEN_PASS_DEF_AFFINELOOPINVARIANTCODEMOTION
#include "mlir/Dialect/Affine/Passes.h.inc" #include "mlir/Dialect/Affine/Passes.h.inc"
} // namespace affine
} // namespace mlir } // namespace mlir
#define DEBUG_TYPE "licm" #define DEBUG_TYPE "licm"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
namespace { namespace {
/// Loop invariant code motion (LICM) pass. /// Loop invariant code motion (LICM) pass.
/// TODO: The pass is missing zero-trip tests. /// TODO: The pass is missing zero-trip tests.
/// TODO: Check for the presence of side effects before hoisting. /// TODO: Check for the presence of side effects before hoisting.
/// TODO: This code should be removed once the new LICM pass can handle its /// TODO: This code should be removed once the new LICM pass can handle its
/// uses. /// uses.
struct LoopInvariantCodeMotion struct LoopInvariantCodeMotion
: public impl::AffineLoopInvariantCodeMotionBase<LoopInvariantCodeMotion> { : public affine::impl::AffineLoopInvariantCodeMotionBase<
LoopInvariantCodeMotion> {
void runOnOperation() override; void runOnOperation() override;
void runOnAffineForOp(AffineForOp forOp); void runOnAffineForOp(AffineForOp forOp);
}; };
} // namespace } // namespace
static bool static bool
checkInvarianceOfNestedIfOps(AffineIfOp ifOp, Value indVar, ValueRange iterArgs, checkInvarianceOfNestedIfOps(AffineIfOp ifOp, Value indVar, ValueRange iterArgs,
SmallPtrSetImpl<Operation *> &opsWithUsers, SmallPtrSetImpl<Operation *> &opsWithUsers,
SmallPtrSetImpl<Operation *> &opsToHoist); SmallPtrSetImpl<Operation *> &opsToHoist);
static bool isOpLoopInvariant(Operation &op, Value indVar, ValueRange iterArgs, static bool isOpLoopInvariant(Operation &op, Value indVar, ValueRange iterArgs,
SmallPtrSetImpl<Operation *> &opsWithUsers, SmallPtrSetImpl<Operation *> &opsWithUsers,
SmallPtrSetImpl<Operation *> &opsToHoist); SmallPtrSetImpl<Operation *> &opsToHoist);
static bool static bool
areAllOpsInTheBlockListInvariant(Region &blockList, Value indVar, areAllOpsInTheBlockListInvariant(Region &blockList, Value indVar,
ValueRange iterArgs, ValueRange iterArgs,
SmallPtrSetImpl<Operation *> &opsWithUsers, SmallPtrSetImpl<Operation *> &opsWithUsers,
SmallPtrSetImpl<Operation *> &opsToHoist); SmallPtrSetImpl<Operation *> &opsToHoist);
// Returns true if the individual op is loop invariant. // Returns true if the individual op is loop invariant.
bool isOpLoopInvariant(Operation &op, Value indVar, ValueRange iterArgs, static bool isOpLoopInvariant(Operation &op, Value indVar, ValueRange iterArgs,
SmallPtrSetImpl<Operation *> &opsWithUsers, SmallPtrSetImpl<Operation *> &opsWithUsers,
SmallPtrSetImpl<Operation *> &opsToHoist) { SmallPtrSetImpl<Operation *> &opsToHoist) {
LLVM_DEBUG(llvm::dbgs() << "iterating on op: " << op;); LLVM_DEBUG(llvm::dbgs() << "iterating on op: " << op;);
if (auto ifOp = dyn_cast<AffineIfOp>(op)) { if (auto ifOp = dyn_cast<AffineIfOp>(op)) {
if (!checkInvarianceOfNestedIfOps(ifOp, indVar, iterArgs, opsWithUsers, if (!checkInvarianceOfNestedIfOps(ifOp, indVar, iterArgs, opsWithUsers,
opsToHoist)) opsToHoist))
return false; return false;
} else if (auto forOp = dyn_cast<AffineForOp>(op)) { } else if (auto forOp = dyn_cast<AffineForOp>(op)) {
if (!areAllOpsInTheBlockListInvariant(forOp.getLoopBody(), indVar, iterArgs, if (!areAllOpsInTheBlockListInvariant(forOp.getLoopBody(), indVar, iterArgs,
▲ Show 20 Lines • Show All 77 Lines • ▼ Show 20 Linesstatic bool isOpLoopInvariant(Operation &op, Value indVar, ValueRange iterArgs,
} }
// If no operand was loop variant, mark this op for motion. // If no operand was loop variant, mark this op for motion.
opsToHoist.insert(&op); opsToHoist.insert(&op);
return true; return true;
} }
// Checks if all ops in a region (i.e. list of blocks) are loop invariant. // Checks if all ops in a region (i.e. list of blocks) are loop invariant.
bool areAllOpsInTheBlockListInvariant( static bool
Region &blockList, Value indVar, ValueRange iterArgs, areAllOpsInTheBlockListInvariant(Region &blockList, Value indVar,
ValueRange iterArgs,
SmallPtrSetImpl<Operation *> &opsWithUsers, SmallPtrSetImpl<Operation *> &opsWithUsers,
SmallPtrSetImpl<Operation *> &opsToHoist) { SmallPtrSetImpl<Operation *> &opsToHoist) {
for (auto &b : blockList) { for (auto &b : blockList) {
for (auto &op : b) { for (auto &op : b) {
if (!isOpLoopInvariant(op, indVar, iterArgs, opsWithUsers, opsToHoist)) if (!isOpLoopInvariant(op, indVar, iterArgs, opsWithUsers, opsToHoist))
return false; return false;
} }
} }
return true; return true;
} }
// Returns true if the affine.if op can be hoisted. // Returns true if the affine.if op can be hoisted.
bool checkInvarianceOfNestedIfOps(AffineIfOp ifOp, Value indVar, static bool
ValueRange iterArgs, checkInvarianceOfNestedIfOps(AffineIfOp ifOp, Value indVar, ValueRange iterArgs,
SmallPtrSetImpl<Operation *> &opsWithUsers, SmallPtrSetImpl<Operation *> &opsWithUsers,
SmallPtrSetImpl<Operation *> &opsToHoist) { SmallPtrSetImpl<Operation *> &opsToHoist) {
if (!areAllOpsInTheBlockListInvariant(ifOp.getThenRegion(), indVar, iterArgs, if (!areAllOpsInTheBlockListInvariant(ifOp.getThenRegion(), indVar, iterArgs,
opsWithUsers, opsToHoist)) opsWithUsers, opsToHoist))
return false; return false;
if (!areAllOpsInTheBlockListInvariant(ifOp.getElseRegion(), indVar, iterArgs, if (!areAllOpsInTheBlockListInvariant(ifOp.getElseRegion(), indVar, iterArgs,
opsWithUsers, opsToHoist)) opsWithUsers, opsToHoist))
return false; return false;
Show All 40 Linesvoid LoopInvariantCodeMotion::runOnOperation() {
// the outer loop, which in turn can be further LICM'ed. // the outer loop, which in turn can be further LICM'ed.
getOperation().walk([&](AffineForOp op) { getOperation().walk([&](AffineForOp op) {
LLVM_DEBUG(op->print(llvm::dbgs() << "\nOriginal loop\n")); LLVM_DEBUG(op->print(llvm::dbgs() << "\nOriginal loop\n"));
runOnAffineForOp(op); runOnAffineForOp(op);
}); });
} }
std::unique_ptr<OperationPass<func::FuncOp>> std::unique_ptr<OperationPass<func::FuncOp>>
mlir::createAffineLoopInvariantCodeMotionPass() { mlir::affine::createAffineLoopInvariantCodeMotionPass() {
return std::make_unique<LoopInvariantCodeMotion>(); return std::make_unique<LoopInvariantCodeMotion>();
} }

mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp

Show All 11 Lines
#include "mlir/Dialect/Affine/Passes.h" #include "mlir/Dialect/Affine/Passes.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/Utils.h" #include "mlir/Dialect/Affine/Utils.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/Func/IR/FuncOps.h"
namespace mlir { namespace mlir {
namespace affine {
#define GEN_PASS_DEF_AFFINELOOPNORMALIZE #define GEN_PASS_DEF_AFFINELOOPNORMALIZE
#include "mlir/Dialect/Affine/Passes.h.inc" #include "mlir/Dialect/Affine/Passes.h.inc"
} // namespace affine
} // namespace mlir } // namespace mlir
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
namespace { namespace {
/// Normalize affine.parallel ops so that lower bounds are 0 and steps are 1. /// Normalize affine.parallel ops so that lower bounds are 0 and steps are 1.
/// As currently implemented, this pass cannot fail, but it might skip over ops /// As currently implemented, this pass cannot fail, but it might skip over ops
/// that are already in a normalized form. /// that are already in a normalized form.
struct AffineLoopNormalizePass struct AffineLoopNormalizePass
: public impl::AffineLoopNormalizeBase<AffineLoopNormalizePass> { : public affine::impl::AffineLoopNormalizeBase<AffineLoopNormalizePass> {
explicit AffineLoopNormalizePass(bool promoteSingleIter) { explicit AffineLoopNormalizePass(bool promoteSingleIter) {
this->promoteSingleIter = promoteSingleIter; this->promoteSingleIter = promoteSingleIter;
} }
void runOnOperation() override { void runOnOperation() override {
getOperation().walk([&](Operation *op) { getOperation().walk([&](Operation *op) {
if (auto affineParallel = dyn_cast<AffineParallelOp>(op)) if (auto affineParallel = dyn_cast<AffineParallelOp>(op))
normalizeAffineParallel(affineParallel); normalizeAffineParallel(affineParallel);
else if (auto affineFor = dyn_cast<AffineForOp>(op)) else if (auto affineFor = dyn_cast<AffineForOp>(op))
(void)normalizeAffineFor(affineFor, promoteSingleIter); (void)normalizeAffineFor(affineFor, promoteSingleIter);
}); });
} }
}; };
} // namespace } // namespace
std::unique_ptr<OperationPass<func::FuncOp>> std::unique_ptr<OperationPass<func::FuncOp>>
mlir::createAffineLoopNormalizePass(bool promoteSingleIter) { mlir::affine::createAffineLoopNormalizePass(bool promoteSingleIter) {
return std::make_unique<AffineLoopNormalizePass>(promoteSingleIter); return std::make_unique<AffineLoopNormalizePass>(promoteSingleIter);
} }

mlir/lib/Dialect/Affine/Transforms/AffineParallelize.cpp

Show All 21 Lines
#include "mlir/Dialect/Affine/LoopUtils.h" #include "mlir/Dialect/Affine/LoopUtils.h"
#include "mlir/Dialect/Affine/Passes.h.inc" #include "mlir/Dialect/Affine/Passes.h.inc"
#include "mlir/Dialect/Affine/Utils.h" #include "mlir/Dialect/Affine/Utils.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/Func/IR/FuncOps.h"
#include "llvm/Support/Debug.h" #include "llvm/Support/Debug.h"
#include <deque> #include <deque>
namespace mlir { namespace mlir {
namespace affine {
#define GEN_PASS_DEF_AFFINEPARALLELIZE #define GEN_PASS_DEF_AFFINEPARALLELIZE
#include "mlir/Dialect/Affine/Passes.h.inc" #include "mlir/Dialect/Affine/Passes.h.inc"
} // namespace affine
} // namespace mlir } // namespace mlir
#define DEBUG_TYPE "affine-parallel" #define DEBUG_TYPE "affine-parallel"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
namespace { namespace {
/// Convert all parallel affine.for op into 1-D affine.parallel op. /// Convert all parallel affine.for op into 1-D affine.parallel op.
struct AffineParallelize struct AffineParallelize
: public impl::AffineParallelizeBase<AffineParallelize> { : public affine::impl::AffineParallelizeBase<AffineParallelize> {
void runOnOperation() override; void runOnOperation() override;
}; };
/// Descriptor of a potentially parallelizable loop. /// Descriptor of a potentially parallelizable loop.
struct ParallelizationCandidate { struct ParallelizationCandidate {
ParallelizationCandidate(AffineForOp l, SmallVector<LoopReduction> &&r) ParallelizationCandidate(AffineForOp l, SmallVector<LoopReduction> &&r)
: loop(l), reductions(std::move(r)) {} : loop(l), reductions(std::move(r)) {}
Show All 34 Lines for (const ParallelizationCandidate &candidate : parallelizableLoops) {
} else { } else {
LLVM_DEBUG(llvm::dbgs() << "[" DEBUG_TYPE "] too many nested loops\n" LLVM_DEBUG(llvm::dbgs() << "[" DEBUG_TYPE "] too many nested loops\n"
<< loop); << loop);
} }
} }
} }
std::unique_ptr<OperationPass<func::FuncOp>> std::unique_ptr<OperationPass<func::FuncOp>>
mlir::createAffineParallelizePass() { mlir::affine::createAffineParallelizePass() {
return std::make_unique<AffineParallelize>(); return std::make_unique<AffineParallelize>();
} }

mlir/lib/Dialect/Affine/Transforms/AffineScalarReplacement.cpp

Show All 17 Lines
#include "mlir/Dialect/Affine/Utils.h" #include "mlir/Dialect/Affine/Utils.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/Dominance.h" #include "mlir/IR/Dominance.h"
#include "mlir/Support/LogicalResult.h" #include "mlir/Support/LogicalResult.h"
#include <algorithm> #include <algorithm>
namespace mlir { namespace mlir {
namespace affine {
#define GEN_PASS_DEF_AFFINESCALARREPLACEMENT #define GEN_PASS_DEF_AFFINESCALARREPLACEMENT
#include "mlir/Dialect/Affine/Passes.h.inc" #include "mlir/Dialect/Affine/Passes.h.inc"
} // namespace affine
} // namespace mlir } // namespace mlir
#define DEBUG_TYPE "affine-scalrep" #define DEBUG_TYPE "affine-scalrep"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
namespace { namespace {
struct AffineScalarReplacement struct AffineScalarReplacement
: public impl::AffineScalarReplacementBase<AffineScalarReplacement> { : public affine::impl::AffineScalarReplacementBase<
AffineScalarReplacement> {
void runOnOperation() override; void runOnOperation() override;
}; };
} // namespace } // namespace
std::unique_ptr<OperationPass<func::FuncOp>> std::unique_ptr<OperationPass<func::FuncOp>>
mlir::createAffineScalarReplacementPass() { mlir::affine::createAffineScalarReplacementPass() {
return std::make_unique<AffineScalarReplacement>(); return std::make_unique<AffineScalarReplacement>();
} }
void AffineScalarReplacement::runOnOperation() { void AffineScalarReplacement::runOnOperation() {
affineScalarReplace(getOperation(), getAnalysis<DominanceInfo>(), affineScalarReplace(getOperation(), getAnalysis<DominanceInfo>(),
getAnalysis<PostDominanceInfo>()); getAnalysis<PostDominanceInfo>());
} }

mlir/lib/Dialect/Affine/Transforms/DecomposeAffineOps.cpp

Show All 12 Lines
#include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/Transforms/Transforms.h" #include "mlir/Dialect/Affine/Transforms/Transforms.h"
#include "mlir/IR/PatternMatch.h" #include "mlir/IR/PatternMatch.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/Support/Debug.h" #include "llvm/Support/Debug.h"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
#define DEBUG_TYPE "decompose-affine-ops" #define DEBUG_TYPE "decompose-affine-ops"
#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ") #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
#define DBGSNL() (llvm::dbgs() << "\n") #define DBGSNL() (llvm::dbgs() << "\n")
/// Count the number of loops surrounding `operand` such that operand could be /// Count the number of loops surrounding `operand` such that operand could be
/// hoisted above. /// hoisted above.
/// Stop counting at the first loop over which the operand cannot be hoisted. /// Stop counting at the first loop over which the operand cannot be hoisted.
static int64_t numEnclosingInvariantLoops(OpOperand &operand) { static int64_t numEnclosingInvariantLoops(OpOperand &operand) {
int64_t count = 0; int64_t count = 0;
Operation *currentOp = operand.getOwner(); Operation *currentOp = operand.getOwner();
while (auto loopOp = currentOp->getParentOfType<LoopLikeOpInterface>()) { while (auto loopOp = currentOp->getParentOfType<LoopLikeOpInterface>()) {
if (!loopOp.isDefinedOutsideOfLoop(operand.get())) if (!loopOp.isDefinedOutsideOfLoop(operand.get()))
break; break;
currentOp = loopOp; currentOp = loopOp;
count++; count++;
} }
return count; return count;
} }
void mlir::reorderOperandsByHoistability(RewriterBase &rewriter, void mlir::affine::reorderOperandsByHoistability(RewriterBase &rewriter,
AffineApplyOp op) { AffineApplyOp op) {
SmallVector<int64_t> numInvariant = llvm::to_vector( SmallVector<int64_t> numInvariant = llvm::to_vector(
llvm::map_range(op->getOpOperands(), [&](OpOperand &operand) { llvm::map_range(op->getOpOperands(), [&](OpOperand &operand) {
return numEnclosingInvariantLoops(operand); return numEnclosingInvariantLoops(operand);
})); }));
int64_t numOperands = op.getNumOperands(); int64_t numOperands = op.getNumOperands();
SmallVector<int64_t> operandPositions = SmallVector<int64_t> operandPositions =
llvm::to_vector(llvm::seq<int64_t>(0, numOperands)); llvm::to_vector(llvm::seq<int64_t>(0, numOperands));
Show All 36 Linesstatic AffineApplyOp createSubApply(RewriterBase &rewriter,
AffineMap m = originalOp.getAffineMap(); AffineMap m = originalOp.getAffineMap();
auto rhsMap = AffineMap::get(m.getNumDims(), m.getNumSymbols(), expr, ctx); auto rhsMap = AffineMap::get(m.getNumDims(), m.getNumSymbols(), expr, ctx);
SmallVector<Value> rhsOperands = originalOp->getOperands(); SmallVector<Value> rhsOperands = originalOp->getOperands();
canonicalizeMapAndOperands(&rhsMap, &rhsOperands); canonicalizeMapAndOperands(&rhsMap, &rhsOperands);
return rewriter.create<AffineApplyOp>(originalOp.getLoc(), rhsMap, return rewriter.create<AffineApplyOp>(originalOp.getLoc(), rhsMap,
rhsOperands); rhsOperands);
} }
FailureOr<AffineApplyOp> mlir::decompose(RewriterBase &rewriter, FailureOr<AffineApplyOp> mlir::affine::decompose(RewriterBase &rewriter,
AffineApplyOp op) { AffineApplyOp op) {
// 1. Preconditions: only handle dimensionless AffineApplyOp maps with a // 1. Preconditions: only handle dimensionless AffineApplyOp maps with a
// top-level binary expression that we can reassociate (i.e. add or mul). // top-level binary expression that we can reassociate (i.e. add or mul).
AffineMap m = op.getAffineMap(); AffineMap m = op.getAffineMap();
if (m.getNumDims() > 0) if (m.getNumDims() > 0)
return rewriter.notifyMatchFailure(op, "expected no dims"); return rewriter.notifyMatchFailure(op, "expected no dims");
AffineExpr remainingExp = m.getResult(0); AffineExpr remainingExp = m.getResult(0);
auto binExpr = remainingExp.dyn_cast<AffineBinaryOpExpr>(); auto binExpr = remainingExp.dyn_cast<AffineBinaryOpExpr>();
▲ Show 20 Lines • Show All 67 Lines • Show Last 20 Lines

mlir/lib/Dialect/Affine/Transforms/LoopCoalescing.cpp

Show All 13 Lines
#include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/Utils/Utils.h" #include "mlir/Dialect/SCF/Utils/Utils.h"
#include "mlir/Transforms/Passes.h" #include "mlir/Transforms/Passes.h"
#include "mlir/Transforms/RegionUtils.h" #include "mlir/Transforms/RegionUtils.h"
#include "llvm/Support/Debug.h" #include "llvm/Support/Debug.h"
namespace mlir { namespace mlir {
namespace affine {
#define GEN_PASS_DEF_LOOPCOALESCING #define GEN_PASS_DEF_LOOPCOALESCING
#include "mlir/Dialect/Affine/Passes.h.inc" #include "mlir/Dialect/Affine/Passes.h.inc"
} // namespace affine
} // namespace mlir } // namespace mlir
#define PASS_NAME "loop-coalescing" #define PASS_NAME "loop-coalescing"
#define DEBUG_TYPE PASS_NAME #define DEBUG_TYPE PASS_NAME
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
namespace { namespace {
struct LoopCoalescingPass struct LoopCoalescingPass
: public impl::LoopCoalescingBase<LoopCoalescingPass> { : public affine::impl::LoopCoalescingBase<LoopCoalescingPass> {
void runOnOperation() override { void runOnOperation() override {
func::FuncOp func = getOperation(); func::FuncOp func = getOperation();
func.walk<WalkOrder::PreOrder>([](Operation *op) { func.walk<WalkOrder::PreOrder>([](Operation *op) {
if (auto scfForOp = dyn_cast<scf::ForOp>(op)) if (auto scfForOp = dyn_cast<scf::ForOp>(op))
(void)coalescePerfectlyNestedLoops(scfForOp); (void)coalescePerfectlyNestedLoops(scfForOp);
else if (auto affineForOp = dyn_cast<AffineForOp>(op)) else if (auto affineForOp = dyn_cast<AffineForOp>(op))
(void)coalescePerfectlyNestedLoops(affineForOp); (void)coalescePerfectlyNestedLoops(affineForOp);
}); });
} }
}; };
} // namespace } // namespace
std::unique_ptr<OperationPass<func::FuncOp>> mlir::createLoopCoalescingPass() { std::unique_ptr<OperationPass<func::FuncOp>>
mlir::affine::createLoopCoalescingPass() {
return std::make_unique<LoopCoalescingPass>(); return std::make_unique<LoopCoalescingPass>();
} }

mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp

Show All 30 Lines
#include "llvm/Support/CommandLine.h" #include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h" #include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h" #include "llvm/Support/raw_ostream.h"
#include <iomanip> #include <iomanip>
#include <optional> #include <optional>
#include <sstream> #include <sstream>
namespace mlir { namespace mlir {
namespace affine {
#define GEN_PASS_DEF_AFFINELOOPFUSION #define GEN_PASS_DEF_AFFINELOOPFUSION
#include "mlir/Dialect/Affine/Passes.h.inc" #include "mlir/Dialect/Affine/Passes.h.inc"
} // namespace affine
} // namespace mlir } // namespace mlir
#define DEBUG_TYPE "affine-loop-fusion" #define DEBUG_TYPE "affine-loop-fusion"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
namespace { namespace {
/// Loop fusion pass. This pass currently supports a greedy fusion policy, /// Loop fusion pass. This pass currently supports a greedy fusion policy,
/// which fuses loop nests with single-writer/single-reader memref dependences /// which fuses loop nests with single-writer/single-reader memref dependences
/// with the goal of improving locality. /// with the goal of improving locality.
// TODO: Support fusion of source loop nests which write to multiple // TODO: Support fusion of source loop nests which write to multiple
// memrefs, where each memref can have multiple users (if profitable). // memrefs, where each memref can have multiple users (if profitable).
// TODO: Extend this pass to check for fusion preventing dependences, // TODO: Extend this pass to check for fusion preventing dependences,
// and add support for more general loop fusion algorithms. // and add support for more general loop fusion algorithms.
struct LoopFusion : public impl::AffineLoopFusionBase<LoopFusion> { struct LoopFusion : public affine::impl::AffineLoopFusionBase<LoopFusion> {
LoopFusion() = default; LoopFusion() = default;
LoopFusion(unsigned fastMemorySpace, uint64_t localBufSizeThresholdBytes, LoopFusion(unsigned fastMemorySpace, uint64_t localBufSizeThresholdBytes,
bool maximalFusion, enum FusionMode affineFusionMode) { bool maximalFusion, enum FusionMode affineFusionMode) {
this->fastMemorySpace = fastMemorySpace; this->fastMemorySpace = fastMemorySpace;
this->localBufSizeThreshold = localBufSizeThresholdBytes / 1024; this->localBufSizeThreshold = localBufSizeThresholdBytes / 1024;
this->maximalFusion = maximalFusion; this->maximalFusion = maximalFusion;
this->affineFusionMode = affineFusionMode; this->affineFusionMode = affineFusionMode;
} }
▲ Show 20 Lines • Show All 968 Lines • ▼ Show 20 Lines do {
// Check the feasibility of fusing src loop nest into dst loop nest // Check the feasibility of fusing src loop nest into dst loop nest
// at loop depths in range [1, dstLoopDepthTest]. // at loop depths in range [1, dstLoopDepthTest].
unsigned maxLegalFusionDepth = 0; unsigned maxLegalFusionDepth = 0;
SmallVector<ComputationSliceState, 8> depthSliceUnions; SmallVector<ComputationSliceState, 8> depthSliceUnions;
depthSliceUnions.resize(dstLoopDepthTest); depthSliceUnions.resize(dstLoopDepthTest);
FusionStrategy strategy(FusionStrategy::ProducerConsumer); FusionStrategy strategy(FusionStrategy::ProducerConsumer);
for (unsigned i = 1; i <= dstLoopDepthTest; ++i) { for (unsigned i = 1; i <= dstLoopDepthTest; ++i) {
FusionResult result = mlir::canFuseLoops( FusionResult result = affine::canFuseLoops(
srcAffineForOp, dstAffineForOp, srcAffineForOp, dstAffineForOp,
/*dstLoopDepth=*/i, &depthSliceUnions[i - 1], strategy); /*dstLoopDepth=*/i, &depthSliceUnions[i - 1], strategy);
if (result.value == FusionResult::Success) if (result.value == FusionResult::Success)
maxLegalFusionDepth = i; maxLegalFusionDepth = i;
} }
if (maxLegalFusionDepth == 0) { if (maxLegalFusionDepth == 0) {
▲ Show 20 Lines • Show All 203 Lines • ▼ Show 20 Lines while (findSiblingNodeToFuse(dstNode, &visitedSibNodeIds, &idAndMemref)) {
auto sibAffineForOp = cast<AffineForOp>(sibNode->op); auto sibAffineForOp = cast<AffineForOp>(sibNode->op);
// Compute loop depth and slice union for fusion. // Compute loop depth and slice union for fusion.
SmallVector<ComputationSliceState, 8> depthSliceUnions; SmallVector<ComputationSliceState, 8> depthSliceUnions;
depthSliceUnions.resize(dstLoopDepthTest); depthSliceUnions.resize(dstLoopDepthTest);
unsigned maxLegalFusionDepth = 0; unsigned maxLegalFusionDepth = 0;
FusionStrategy strategy(memref); FusionStrategy strategy(memref);
for (unsigned i = 1; i <= dstLoopDepthTest; ++i) { for (unsigned i = 1; i <= dstLoopDepthTest; ++i) {
FusionResult result = mlir::canFuseLoops( FusionResult result = affine::canFuseLoops(
sibAffineForOp, dstAffineForOp, sibAffineForOp, dstAffineForOp,
/*dstLoopDepth=*/i, &depthSliceUnions[i - 1], strategy); /*dstLoopDepth=*/i, &depthSliceUnions[i - 1], strategy);
if (result.value == FusionResult::Success) if (result.value == FusionResult::Success)
maxLegalFusionDepth = i; maxLegalFusionDepth = i;
} }
// Skip if fusion is not feasible at any loop depths. // Skip if fusion is not feasible at any loop depths.
Show All 15 Lines while (findSiblingNodeToFuse(dstNode, &visitedSibNodeIds, &idAndMemref)) {
assert(bestDstLoopDepth > 0 && "Unexpected loop fusion depth"); assert(bestDstLoopDepth > 0 && "Unexpected loop fusion depth");
assert(!depthSliceUnions[bestDstLoopDepth - 1].isEmpty() && assert(!depthSliceUnions[bestDstLoopDepth - 1].isEmpty() &&
"Fusion depth has no computed slice union"); "Fusion depth has no computed slice union");
// Check if source loop is being inserted in the innermost // Check if source loop is being inserted in the innermost
// destination loop. Based on this, the fused loop may be optimized // destination loop. Based on this, the fused loop may be optimized
// further inside `fuseLoops`. // further inside `fuseLoops`.
bool isInnermostInsertion = (bestDstLoopDepth == dstLoopDepthTest); bool isInnermostInsertion = (bestDstLoopDepth == dstLoopDepthTest);
// Fuse computation slice of 'sibLoopNest' into 'dstLoopNest'. // Fuse computation slice of 'sibLoopNest' into 'dstLoopNest'.
mlir::fuseLoops(sibAffineForOp, dstAffineForOp, affine::fuseLoops(sibAffineForOp, dstAffineForOp,
depthSliceUnions[bestDstLoopDepth - 1], depthSliceUnions[bestDstLoopDepth - 1],
isInnermostInsertion); isInnermostInsertion);
auto dstForInst = cast<AffineForOp>(dstNode->op); auto dstForInst = cast<AffineForOp>(dstNode->op);
// Update operation position of fused loop nest (if needed). // Update operation position of fused loop nest (if needed).
if (insertPointInst != dstForInst) { if (insertPointInst != dstForInst) {
dstForInst->moveBefore(insertPointInst); dstForInst->moveBefore(insertPointInst);
} }
// Update data dependence graph state post fusion. // Update data dependence graph state post fusion.
updateStateAfterSiblingFusion(sibNode, dstNode); updateStateAfterSiblingFusion(sibNode, dstNode);
▲ Show 20 Lines • Show All 191 Lines • ▼ Show 20 Lines
} }
void LoopFusion::runOnOperation() { void LoopFusion::runOnOperation() {
for (Region &region : getOperation()->getRegions()) for (Region &region : getOperation()->getRegions())
for (Block &block : region.getBlocks()) for (Block &block : region.getBlocks())
runOnBlock(&block); runOnBlock(&block);
} }
std::unique_ptr<Pass> std::unique_ptr<Pass> mlir::affine::createLoopFusionPass(
mlir::createLoopFusionPass(unsigned fastMemorySpace, unsigned fastMemorySpace, uint64_t localBufSizeThreshold,
uint64_t localBufSizeThreshold, bool maximalFusion, bool maximalFusion, enum FusionMode affineFusionMode) {
enum FusionMode affineFusionMode) {
return std::make_unique<LoopFusion>(fastMemorySpace, localBufSizeThreshold, return std::make_unique<LoopFusion>(fastMemorySpace, localBufSizeThreshold,
maximalFusion, affineFusionMode); maximalFusion, affineFusionMode);
} }

mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp

Show All 22 Lines
#include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/Builders.h" #include "mlir/IR/Builders.h"
#include "mlir/IR/IRMapping.h" #include "mlir/IR/IRMapping.h"
#include "llvm/Support/CommandLine.h" #include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h" #include "llvm/Support/Debug.h"
#include <optional> #include <optional>
namespace mlir { namespace mlir {
namespace affine {
#define GEN_PASS_DEF_AFFINELOOPTILING #define GEN_PASS_DEF_AFFINELOOPTILING
#include "mlir/Dialect/Affine/Passes.h.inc" #include "mlir/Dialect/Affine/Passes.h.inc"
} // namespace affine
} // namespace mlir } // namespace mlir
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
#define DEBUG_TYPE "affine-loop-tile" #define DEBUG_TYPE "affine-loop-tile"
namespace { namespace {
/// A pass to perform loop tiling on all suitable loop nests of a Function. /// A pass to perform loop tiling on all suitable loop nests of a Function.
struct LoopTiling : public impl::AffineLoopTilingBase<LoopTiling> { struct LoopTiling : public affine::impl::AffineLoopTilingBase<LoopTiling> {
LoopTiling() = default; LoopTiling() = default;
explicit LoopTiling(uint64_t cacheSizeBytes, bool avoidMaxMinBounds = true) explicit LoopTiling(uint64_t cacheSizeBytes, bool avoidMaxMinBounds = true)
: avoidMaxMinBounds(avoidMaxMinBounds) { : avoidMaxMinBounds(avoidMaxMinBounds) {
this->cacheSizeInKiB = cacheSizeBytes / 1024; this->cacheSizeInKiB = cacheSizeBytes / 1024;
} }
void runOnOperation() override; void runOnOperation() override;
void getTileSizes(ArrayRef<AffineForOp> band, void getTileSizes(ArrayRef<AffineForOp> band,
SmallVectorImpl<unsigned> *tileSizes); SmallVectorImpl<unsigned> *tileSizes);
// Default tile size if nothing is provided. // Default tile size if nothing is provided.
constexpr static unsigned kDefaultTileSize = 4; constexpr static unsigned kDefaultTileSize = 4;
// If true, tile sizes are set to avoid max/min in bounds if possible. // If true, tile sizes are set to avoid max/min in bounds if possible.
bool avoidMaxMinBounds = true; bool avoidMaxMinBounds = true;
}; };
} // namespace } // namespace
/// Creates a pass to perform loop tiling on all suitable loop nests of a /// Creates a pass to perform loop tiling on all suitable loop nests of a
/// Function. /// Function.
std::unique_ptr<OperationPass<func::FuncOp>> std::unique_ptr<OperationPass<func::FuncOp>>
mlir::createLoopTilingPass(uint64_t cacheSizeBytes) { mlir::affine::createLoopTilingPass(uint64_t cacheSizeBytes) {
return std::make_unique<LoopTiling>(cacheSizeBytes); return std::make_unique<LoopTiling>(cacheSizeBytes);
} }
std::unique_ptr<OperationPass<func::FuncOp>> mlir::createLoopTilingPass() { std::unique_ptr<OperationPass<func::FuncOp>>
mlir::affine::createLoopTilingPass() {
return std::make_unique<LoopTiling>(); return std::make_unique<LoopTiling>();
} }
/// Reduces each tile size to the largest divisor of the corresponding trip /// Reduces each tile size to the largest divisor of the corresponding trip
/// count (if the trip count is known). /// count (if the trip count is known).
static void adjustToDivisorsOfTripCounts(ArrayRef<AffineForOp> band, static void adjustToDivisorsOfTripCounts(ArrayRef<AffineForOp> band,
SmallVectorImpl<unsigned> *tileSizes) { SmallVectorImpl<unsigned> *tileSizes) {
assert(band.size() == tileSizes->size() && "invalid tile size count"); assert(band.size() == tileSizes->size() && "invalid tile size count");
Show All 15 Lines
/// Checks whether hyper-rectangular loop tiling of the nest represented by /// Checks whether hyper-rectangular loop tiling of the nest represented by
/// `origLoops` is valid. The validity condition is from Irigoin and Triolet, /// `origLoops` is valid. The validity condition is from Irigoin and Triolet,
/// which states that two tiles cannot depend on each other. We simplify such /// which states that two tiles cannot depend on each other. We simplify such
/// condition to just checking whether there is any negative dependence /// condition to just checking whether there is any negative dependence
/// direction, since we have the prior knowledge that the tiling results will be /// direction, since we have the prior knowledge that the tiling results will be
/// hyper-rectangles, which are scheduled in the lexicographically increasing /// hyper-rectangles, which are scheduled in the lexicographically increasing
/// order on the vector of loop indices. This function will return failure when /// order on the vector of loop indices. This function will return failure when
/// any dependence component is negative along any of `origLoops`. /// any dependence component is negative along any of `origLoops`.
static bool checkTilingLegality(MutableArrayRef<mlir::AffineForOp> origLoops) { static bool checkTilingLegality(MutableArrayRef<AffineForOp> origLoops) {
assert(!origLoops.empty() && "no original loops provided"); assert(!origLoops.empty() && "no original loops provided");
// We first find out all dependences we intend to check. // We first find out all dependences we intend to check.
SmallVector<Operation *, 8> loadAndStoreOps; SmallVector<Operation *, 8> loadAndStoreOps;
origLoops[0]->walk([&](Operation *op) { origLoops[0]->walk([&](Operation *op) {
if (isa<AffineReadOpInterface, AffineWriteOpInterface>(op)) if (isa<AffineReadOpInterface, AffineWriteOpInterface>(op))
loadAndStoreOps.push_back(op); loadAndStoreOps.push_back(op);
}); });
▲ Show 20 Lines • Show All 171 Lines • Show Last 20 Lines

mlir/lib/Dialect/Affine/Transforms/LoopUnroll.cpp

Show All 19 Lines
#include "mlir/IR/AffineMap.h" #include "mlir/IR/AffineMap.h"
#include "mlir/IR/Builders.h" #include "mlir/IR/Builders.h"
#include "llvm/ADT/DenseMap.h" #include "llvm/ADT/DenseMap.h"
#include "llvm/Support/CommandLine.h" #include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h" #include "llvm/Support/Debug.h"
#include <optional> #include <optional>
namespace mlir { namespace mlir {
namespace affine {
#define GEN_PASS_DEF_AFFINELOOPUNROLL #define GEN_PASS_DEF_AFFINELOOPUNROLL
#include "mlir/Dialect/Affine/Passes.h.inc" #include "mlir/Dialect/Affine/Passes.h.inc"
} // namespace affine
} // namespace mlir } // namespace mlir
#define DEBUG_TYPE "affine-loop-unroll" #define DEBUG_TYPE "affine-loop-unroll"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
namespace { namespace {
// TODO: this is really a test pass and should be moved out of dialect // TODO: this is really a test pass and should be moved out of dialect
// transforms. // transforms.
/// Loop unrolling pass. Unrolls all innermost loops unless full unrolling and a /// Loop unrolling pass. Unrolls all innermost loops unless full unrolling and a
/// full unroll threshold was specified, in which case, fully unrolls all loops /// full unroll threshold was specified, in which case, fully unrolls all loops
/// with trip count less than the specified threshold. The latter is for testing /// with trip count less than the specified threshold. The latter is for testing
/// purposes, especially for testing outer loop unrolling. /// purposes, especially for testing outer loop unrolling.
struct LoopUnroll : public impl::AffineLoopUnrollBase<LoopUnroll> { struct LoopUnroll : public affine::impl::AffineLoopUnrollBase<LoopUnroll> {
// Callback to obtain unroll factors; if this has a callable target, takes // Callback to obtain unroll factors; if this has a callable target, takes
// precedence over command-line argument or passed argument. // precedence over command-line argument or passed argument.
const std::function<unsigned(AffineForOp)> getUnrollFactor; const std::function<unsigned(AffineForOp)> getUnrollFactor;
LoopUnroll() : getUnrollFactor(nullptr) {} LoopUnroll() : getUnrollFactor(nullptr) {}
LoopUnroll(const LoopUnroll &other) LoopUnroll(const LoopUnroll &other)
= default; = default;
▲ Show 20 Lines • Show All 83 Lines • ▼ Show 20 Lines if (unrollFull)
return loopUnrollFull(forOp); return loopUnrollFull(forOp);
// Otherwise, unroll by the given unroll factor. // Otherwise, unroll by the given unroll factor.
if (unrollUpToFactor) if (unrollUpToFactor)
return loopUnrollUpToFactor(forOp, unrollFactor); return loopUnrollUpToFactor(forOp, unrollFactor);
return loopUnrollByFactor(forOp, unrollFactor, /*annotateFn=*/nullptr, return loopUnrollByFactor(forOp, unrollFactor, /*annotateFn=*/nullptr,
cleanUpUnroll); cleanUpUnroll);
} }
std::unique_ptr<OperationPass<func::FuncOp>> mlir::createLoopUnrollPass( std::unique_ptr<OperationPass<func::FuncOp>> mlir::affine::createLoopUnrollPass(
int unrollFactor, bool unrollUpToFactor, bool unrollFull, int unrollFactor, bool unrollUpToFactor, bool unrollFull,
const std::function<unsigned(AffineForOp)> &getUnrollFactor) { const std::function<unsigned(AffineForOp)> &getUnrollFactor) {
return std::make_unique<LoopUnroll>( return std::make_unique<LoopUnroll>(
unrollFactor == -1 ? std::nullopt : std::optional<unsigned>(unrollFactor), unrollFactor == -1 ? std::nullopt : std::optional<unsigned>(unrollFactor),
unrollUpToFactor, unrollFull, getUnrollFactor); unrollUpToFactor, unrollFull, getUnrollFactor);
} }

mlir/lib/Dialect/Affine/Transforms/LoopUnrollAndJam.cpp

Show First 20 Lines • Show All 43 Lines • ▼ Show 20 Lines
#include "mlir/IR/AffineMap.h" #include "mlir/IR/AffineMap.h"
#include "mlir/IR/Builders.h" #include "mlir/IR/Builders.h"
#include "mlir/IR/IRMapping.h" #include "mlir/IR/IRMapping.h"
#include "llvm/ADT/DenseMap.h" #include "llvm/ADT/DenseMap.h"
#include "llvm/Support/CommandLine.h" #include "llvm/Support/CommandLine.h"
#include <optional> #include <optional>
namespace mlir { namespace mlir {
namespace affine {
#define GEN_PASS_DEF_AFFINELOOPUNROLLANDJAM #define GEN_PASS_DEF_AFFINELOOPUNROLLANDJAM
#include "mlir/Dialect/Affine/Passes.h.inc" #include "mlir/Dialect/Affine/Passes.h.inc"
} // namespace affine
} // namespace mlir } // namespace mlir
#define DEBUG_TYPE "affine-loop-unroll-jam" #define DEBUG_TYPE "affine-loop-unroll-jam"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
namespace { namespace {
/// Loop unroll jam pass. Currently, this just unroll jams the first /// Loop unroll jam pass. Currently, this just unroll jams the first
/// outer loop in a Function. /// outer loop in a Function.
struct LoopUnrollAndJam struct LoopUnrollAndJam
: public impl::AffineLoopUnrollAndJamBase<LoopUnrollAndJam> { : public affine::impl::AffineLoopUnrollAndJamBase<LoopUnrollAndJam> {
explicit LoopUnrollAndJam( explicit LoopUnrollAndJam(
std::optional<unsigned> unrollJamFactor = std::nullopt) { std::optional<unsigned> unrollJamFactor = std::nullopt) {
if (unrollJamFactor) if (unrollJamFactor)
this->unrollJamFactor = *unrollJamFactor; this->unrollJamFactor = *unrollJamFactor;
} }
void runOnOperation() override; void runOnOperation() override;
}; };
} // namespace } // namespace
std::unique_ptr<OperationPass<func::FuncOp>> std::unique_ptr<OperationPass<func::FuncOp>>
mlir::createLoopUnrollAndJamPass(int unrollJamFactor) { mlir::affine::createLoopUnrollAndJamPass(int unrollJamFactor) {
return std::make_unique<LoopUnrollAndJam>( return std::make_unique<LoopUnrollAndJam>(
unrollJamFactor == -1 ? std::nullopt unrollJamFactor == -1 ? std::nullopt
: std::optional<unsigned>(unrollJamFactor)); : std::optional<unsigned>(unrollJamFactor));
} }
void LoopUnrollAndJam::runOnOperation() { void LoopUnrollAndJam::runOnOperation() {
if (getOperation().isExternal()) if (getOperation().isExternal())
return; return;
// Currently, just the outermost loop from the first loop nest is // Currently, just the outermost loop from the first loop nest is
// unroll-and-jammed by this pass. However, runOnAffineForOp can be called on // unroll-and-jammed by this pass. However, runOnAffineForOp can be called on
// any for operation. // any for operation.
auto &entryBlock = getOperation().front(); auto &entryBlock = getOperation().front();
if (auto forOp = dyn_cast<AffineForOp>(entryBlock.front())) if (auto forOp = dyn_cast<AffineForOp>(entryBlock.front()))
(void)loopUnrollJamByFactor(forOp, unrollJamFactor); (void)loopUnrollJamByFactor(forOp, unrollJamFactor);
} }

mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp

Show All 21 Lines
#include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/Builders.h" #include "mlir/IR/Builders.h"
#include "mlir/Transforms/Passes.h" #include "mlir/Transforms/Passes.h"
#include "llvm/ADT/DenseMap.h" #include "llvm/ADT/DenseMap.h"
#include "llvm/Support/Debug.h" #include "llvm/Support/Debug.h"
namespace mlir { namespace mlir {
namespace affine {
#define GEN_PASS_DEF_AFFINEPIPELINEDATATRANSFER #define GEN_PASS_DEF_AFFINEPIPELINEDATATRANSFER
#include "mlir/Dialect/Affine/Passes.h.inc" #include "mlir/Dialect/Affine/Passes.h.inc"
} // namespace affine
} // namespace mlir } // namespace mlir
#define DEBUG_TYPE "affine-pipeline-data-transfer" #define DEBUG_TYPE "affine-pipeline-data-transfer"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
namespace { namespace {
struct PipelineDataTransfer struct PipelineDataTransfer
: public impl::AffinePipelineDataTransferBase<PipelineDataTransfer> { : public affine::impl::AffinePipelineDataTransferBase<
PipelineDataTransfer> {
void runOnOperation() override; void runOnOperation() override;
void runOnAffineForOp(AffineForOp forOp); void runOnAffineForOp(AffineForOp forOp);
std::vector<AffineForOp> forOps; std::vector<AffineForOp> forOps;
}; };
} // namespace } // namespace
/// Creates a pass to pipeline explicit movement of data across levels of the /// Creates a pass to pipeline explicit movement of data across levels of the
/// memory hierarchy. /// memory hierarchy.
std::unique_ptr<OperationPass<func::FuncOp>> std::unique_ptr<OperationPass<func::FuncOp>>
mlir::createPipelineDataTransferPass() { mlir::affine::createPipelineDataTransferPass() {
return std::make_unique<PipelineDataTransfer>(); return std::make_unique<PipelineDataTransfer>();
} }
// Returns the position of the tag memref operand given a DMA operation. // Returns the position of the tag memref operand given a DMA operation.
// Temporary utility: will be replaced when DmaStart/DmaFinish abstract op's are // Temporary utility: will be replaced when DmaStart/DmaFinish abstract op's are
// added. TODO // added. TODO
static unsigned getTagMemRefPos(Operation &dmaOp) { static unsigned getTagMemRefPos(Operation &dmaOp) {
assert((isa<AffineDmaStartOp, AffineDmaWaitOp>(dmaOp))); assert((isa<AffineDmaStartOp, AffineDmaWaitOp>(dmaOp)));
▲ Show 20 Lines • Show All 262 Lines • ▼ Show 20 Linesvoid PipelineDataTransfer::runOnAffineForOp(AffineForOp forOp) {
// Store shift for operation for later lookup for AffineApplyOp's. // Store shift for operation for later lookup for AffineApplyOp's.
DenseMap<Operation *, unsigned> instShiftMap; DenseMap<Operation *, unsigned> instShiftMap;
for (auto &pair : startWaitPairs) { for (auto &pair : startWaitPairs) {
auto *dmaStartOp = pair.first; auto *dmaStartOp = pair.first;
assert(isa<AffineDmaStartOp>(dmaStartOp)); assert(isa<AffineDmaStartOp>(dmaStartOp));
instShiftMap[dmaStartOp] = 0; instShiftMap[dmaStartOp] = 0;
// Set shifts for DMA start op's affine operand computation slices to 0. // Set shifts for DMA start op's affine operand computation slices to 0.
SmallVector<AffineApplyOp, 4> sliceOps; SmallVector<AffineApplyOp, 4> sliceOps;
mlir::createAffineComputationSlice(dmaStartOp, &sliceOps); affine::createAffineComputationSlice(dmaStartOp, &sliceOps);
if (!sliceOps.empty()) { if (!sliceOps.empty()) {
for (auto sliceOp : sliceOps) { for (auto sliceOp : sliceOps) {
instShiftMap[sliceOp.getOperation()] = 0; instShiftMap[sliceOp.getOperation()] = 0;
} }
} else { } else {
// If a slice wasn't created, the reachable affine.apply op's from its // If a slice wasn't created, the reachable affine.apply op's from its
// operands are the ones that go with it. // operands are the ones that go with it.
SmallVector<Operation *, 4> affineApplyInsts; SmallVector<Operation *, 4> affineApplyInsts;
Show All 37 Lines

mlir/lib/Dialect/Affine/Transforms/ReifyValueBounds.cpp

//===- ReifyValueBounds.cpp --- Reify value bounds with affine ops ------*-===// //===- ReifyValueBounds.cpp --- Reify value bounds with affine ops ------*-===//
// //
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information. // See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// //
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
#include "mlir/Dialect/Affine/Transforms/Transforms.h" #include "mlir/Dialect/Affine/Transforms/Transforms.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Interfaces/ValueBoundsOpInterface.h" #include "mlir/Interfaces/ValueBoundsOpInterface.h"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
static FailureOr<OpFoldResult> static FailureOr<OpFoldResult>
reifyValueBound(OpBuilder &b, Location loc, presburger::BoundType type, reifyValueBound(OpBuilder &b, Location loc, presburger::BoundType type,
Value value, std::optional<int64_t> dim, Value value, std::optional<int64_t> dim,
function_ref<bool(Value, std::optional<int64_t>)> stopCondition, function_ref<bool(Value, std::optional<int64_t>)> stopCondition,
bool closedUB) { bool closedUB) {
// Compute bound. // Compute bound.
AffineMap boundMap; AffineMap boundMap;
Show All 24 Lines if (isa<RankedTensorType>(value.getType())) {
// A memref dimension is used: generate a memref.dim. // A memref dimension is used: generate a memref.dim.
operands.push_back(b.create<memref::DimOp>(loc, value, *dim)); operands.push_back(b.create<memref::DimOp>(loc, value, *dim));
} else { } else {
llvm_unreachable("cannot generate DimOp for unsupported shaped type"); llvm_unreachable("cannot generate DimOp for unsupported shaped type");
} }
} }
// Simplify and return bound. // Simplify and return bound.
mlir::canonicalizeMapAndOperands(&boundMap, &operands); affine::canonicalizeMapAndOperands(&boundMap, &operands);
// Check for special cases where no affine.apply op is needed. // Check for special cases where no affine.apply op is needed.
if (boundMap.isSingleConstant()) { if (boundMap.isSingleConstant()) {
// Bound is a constant: return an IntegerAttr. // Bound is a constant: return an IntegerAttr.
return static_cast<OpFoldResult>( return static_cast<OpFoldResult>(
b.getIndexAttr(boundMap.getSingleConstantResult())); b.getIndexAttr(boundMap.getSingleConstantResult()));
} }
// No affine.apply op is needed if the bound is a single SSA value. // No affine.apply op is needed if the bound is a single SSA value.
if (auto expr = boundMap.getResult(0).dyn_cast<AffineDimExpr>()) if (auto expr = boundMap.getResult(0).dyn_cast<AffineDimExpr>())
return static_cast<OpFoldResult>(operands[expr.getPosition()]); return static_cast<OpFoldResult>(operands[expr.getPosition()]);
if (auto expr = boundMap.getResult(0).dyn_cast<AffineSymbolExpr>()) if (auto expr = boundMap.getResult(0).dyn_cast<AffineSymbolExpr>())
return static_cast<OpFoldResult>( return static_cast<OpFoldResult>(
operands[expr.getPosition() + boundMap.getNumDims()]); operands[expr.getPosition() + boundMap.getNumDims()]);
// General case: build affine.apply op. // General case: build affine.apply op.
return static_cast<OpFoldResult>( return static_cast<OpFoldResult>(
b.create<AffineApplyOp>(loc, boundMap, operands).getResult()); b.create<affine::AffineApplyOp>(loc, boundMap, operands).getResult());
} }
FailureOr<OpFoldResult> mlir::reifyShapedValueDimBound( FailureOr<OpFoldResult> mlir::affine::reifyShapedValueDimBound(
OpBuilder &b, Location loc, presburger::BoundType type, Value value, OpBuilder &b, Location loc, presburger::BoundType type, Value value,
int64_t dim, ValueBoundsConstraintSet::StopConditionFn stopCondition, int64_t dim, ValueBoundsConstraintSet::StopConditionFn stopCondition,
bool closedUB) { bool closedUB) {
auto reifyToOperands = [&](Value v, std::optional<int64_t> d) { auto reifyToOperands = [&](Value v, std::optional<int64_t> d) {
// We are trying to reify a bound for `value` in terms of the owning op's // We are trying to reify a bound for `value` in terms of the owning op's
// operands. Construct a stop condition that evaluates to "true" for any SSA // operands. Construct a stop condition that evaluates to "true" for any SSA
// value except for `value`. I.e., the bound will be computed in terms of // value except for `value`. I.e., the bound will be computed in terms of
// any SSA values except for `value`. The first such values are operands of // any SSA values except for `value`. The first such values are operands of
// the owner of `value`. // the owner of `value`.
return v != value; return v != value;
}; };
return reifyValueBound(b, loc, type, value, dim, return reifyValueBound(b, loc, type, value, dim,
stopCondition ? stopCondition : reifyToOperands, stopCondition ? stopCondition : reifyToOperands,
closedUB); closedUB);
} }
FailureOr<OpFoldResult> mlir::reifyIndexValueBound( FailureOr<OpFoldResult> mlir::affine::reifyIndexValueBound(
OpBuilder &b, Location loc, presburger::BoundType type, Value value, OpBuilder &b, Location loc, presburger::BoundType type, Value value,
ValueBoundsConstraintSet::StopConditionFn stopCondition, bool closedUB) { ValueBoundsConstraintSet::StopConditionFn stopCondition, bool closedUB) {
auto reifyToOperands = [&](Value v, std::optional<int64_t> d) { auto reifyToOperands = [&](Value v, std::optional<int64_t> d) {
return v != value; return v != value;
}; };
return reifyValueBound(b, loc, type, value, /*dim=*/std::nullopt, return reifyValueBound(b, loc, type, value, /*dim=*/std::nullopt,
stopCondition ? stopCondition : reifyToOperands, stopCondition ? stopCondition : reifyToOperands,
closedUB); closedUB);
} }

mlir/lib/Dialect/Affine/Transforms/SimplifyAffineStructures.cpp

Show All 14 Lines
#include "mlir/Dialect/Affine/Analysis/Utils.h" #include "mlir/Dialect/Affine/Analysis/Utils.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/Utils.h" #include "mlir/Dialect/Affine/Utils.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/IntegerSet.h" #include "mlir/IR/IntegerSet.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
namespace mlir { namespace mlir {
namespace affine {
#define GEN_PASS_DEF_SIMPLIFYAFFINESTRUCTURES #define GEN_PASS_DEF_SIMPLIFYAFFINESTRUCTURES
#include "mlir/Dialect/Affine/Passes.h.inc" #include "mlir/Dialect/Affine/Passes.h.inc"
} // namespace affine
} // namespace mlir } // namespace mlir
#define DEBUG_TYPE "simplify-affine-structure" #define DEBUG_TYPE "simplify-affine-structure"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
namespace { namespace {
/// Simplifies affine maps and sets appearing in the operations of the Function. /// Simplifies affine maps and sets appearing in the operations of the Function.
/// This part is mainly to test the simplifyAffineExpr method. In addition, /// This part is mainly to test the simplifyAffineExpr method. In addition,
/// all memrefs with non-trivial layout maps are converted to ones with trivial /// all memrefs with non-trivial layout maps are converted to ones with trivial
/// identity layout ones. /// identity layout ones.
struct SimplifyAffineStructures struct SimplifyAffineStructures
: public impl::SimplifyAffineStructuresBase<SimplifyAffineStructures> { : public affine::impl::SimplifyAffineStructuresBase<
SimplifyAffineStructures> {
void runOnOperation() override; void runOnOperation() override;
/// Utility to simplify an affine attribute and update its entry in the parent /// Utility to simplify an affine attribute and update its entry in the parent
/// operation if necessary. /// operation if necessary.
template <typename AttributeT> template <typename AttributeT>
void simplifyAndUpdateAttribute(Operation *op, StringAttr name, void simplifyAndUpdateAttribute(Operation *op, StringAttr name,
AttributeT attr) { AttributeT attr) {
auto &simplified = simplifiedAttributes[attr]; auto &simplified = simplifiedAttributes[attr];
Show All 26 Linesstruct SimplifyAffineStructures
} }
DenseMap<Attribute, Attribute> simplifiedAttributes; DenseMap<Attribute, Attribute> simplifiedAttributes;
}; };
} // namespace } // namespace
std::unique_ptr<OperationPass<func::FuncOp>> std::unique_ptr<OperationPass<func::FuncOp>>
mlir::createSimplifyAffineStructuresPass() { mlir::affine::createSimplifyAffineStructuresPass() {
return std::make_unique<SimplifyAffineStructures>(); return std::make_unique<SimplifyAffineStructures>();
} }
void SimplifyAffineStructures::runOnOperation() { void SimplifyAffineStructures::runOnOperation() {
auto func = getOperation(); auto func = getOperation();
simplifiedAttributes.clear(); simplifiedAttributes.clear();
RewritePatternSet patterns(func.getContext()); RewritePatternSet patterns(func.getContext());
AffineApplyOp::getCanonicalizationPatterns(patterns, func.getContext()); AffineApplyOp::getCanonicalizationPatterns(patterns, func.getContext());
Show All 22 Lines

mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp

Show All 25 Lines
#include "mlir/IR/IRMapping.h" #include "mlir/IR/IRMapping.h"
#include "mlir/Pass/Pass.h" #include "mlir/Pass/Pass.h"
#include "mlir/Support/LLVM.h" #include "mlir/Support/LLVM.h"
#include "llvm/ADT/STLExtras.h" #include "llvm/ADT/STLExtras.h"
#include "llvm/Support/Debug.h" #include "llvm/Support/Debug.h"
#include <optional> #include <optional>
namespace mlir { namespace mlir {
namespace affine {
#define GEN_PASS_DEF_AFFINEVECTORIZE #define GEN_PASS_DEF_AFFINEVECTORIZE
#include "mlir/Dialect/Affine/Passes.h.inc" #include "mlir/Dialect/Affine/Passes.h.inc"
} // namespace affine
} // namespace mlir } // namespace mlir
using namespace mlir; using namespace mlir;
using namespace affine;
using namespace vector; using namespace vector;
/// ///
/// Implements a high-level vectorization strategy on a Function. /// Implements a high-level vectorization strategy on a Function.
/// The abstraction used is that of super-vectors, which provide a single, /// The abstraction used is that of super-vectors, which provide a single,
/// compact, representation in the vector types, information that is expected /// compact, representation in the vector types, information that is expected
/// to reduce the impact of the phase ordering problem /// to reduce the impact of the phase ordering problem
/// ///
▲ Show 20 Lines • Show All 533 Lines • ▼ Show 20 Lines
/// Creates a vectorization pattern from the command line arguments. /// Creates a vectorization pattern from the command line arguments.
/// Up to 3-D patterns are supported. /// Up to 3-D patterns are supported.
/// If the command line argument requests a pattern of higher order, returns an /// If the command line argument requests a pattern of higher order, returns an
/// empty pattern list which will conservatively result in no vectorization. /// empty pattern list which will conservatively result in no vectorization.
static std::optional<NestedPattern> static std::optional<NestedPattern>
makePattern(const DenseSet<Operation *> &parallelLoops, int vectorRank, makePattern(const DenseSet<Operation *> &parallelLoops, int vectorRank,
ArrayRef<int64_t> fastestVaryingPattern) { ArrayRef<int64_t> fastestVaryingPattern) {
using matcher::For; using affine::matcher::For;
int64_t d0 = fastestVaryingPattern.empty() ? -1 : fastestVaryingPattern[0]; int64_t d0 = fastestVaryingPattern.empty() ? -1 : fastestVaryingPattern[0];
int64_t d1 = fastestVaryingPattern.size() < 2 ? -1 : fastestVaryingPattern[1]; int64_t d1 = fastestVaryingPattern.size() < 2 ? -1 : fastestVaryingPattern[1];
int64_t d2 = fastestVaryingPattern.size() < 3 ? -1 : fastestVaryingPattern[2]; int64_t d2 = fastestVaryingPattern.size() < 3 ? -1 : fastestVaryingPattern[2];
switch (vectorRank) { switch (vectorRank) {
case 1: case 1:
return For(isVectorizableLoopPtrFactory(parallelLoops, d0)); return For(isVectorizableLoopPtrFactory(parallelLoops, d0));
case 2: case 2:
return For(isVectorizableLoopPtrFactory(parallelLoops, d0), return For(isVectorizableLoopPtrFactory(parallelLoops, d0),
For(isVectorizableLoopPtrFactory(parallelLoops, d1))); For(isVectorizableLoopPtrFactory(parallelLoops, d1)));
case 3: case 3:
return For(isVectorizableLoopPtrFactory(parallelLoops, d0), return For(isVectorizableLoopPtrFactory(parallelLoops, d0),
For(isVectorizableLoopPtrFactory(parallelLoops, d1), For(isVectorizableLoopPtrFactory(parallelLoops, d1),
For(isVectorizableLoopPtrFactory(parallelLoops, d2)))); For(isVectorizableLoopPtrFactory(parallelLoops, d2))));
default: { default: {
return std::nullopt; return std::nullopt;
} }
} }
} }
static NestedPattern &vectorTransferPattern() { static NestedPattern &vectorTransferPattern() {
static auto pattern = matcher::Op([](Operation &op) { static auto pattern = affine::matcher::Op([](Operation &op) {
return isa<vector::TransferReadOp, vector::TransferWriteOp>(op); return isa<vector::TransferReadOp, vector::TransferWriteOp>(op);
}); });
return pattern; return pattern;
} }
namespace { namespace {
/// Base state for the vectorize pass. /// Base state for the vectorize pass.
/// Command line arguments are preempted by non-empty pass arguments. /// Command line arguments are preempted by non-empty pass arguments.
struct Vectorize : public impl::AffineVectorizeBase<Vectorize> { struct Vectorize : public affine::impl::AffineVectorizeBase<Vectorize> {
using Base::Base; using Base::Base;
void runOnOperation() override; void runOnOperation() override;
}; };
} // namespace } // namespace
static void vectorizeLoopIfProfitable(Operation *loop, unsigned depthInPattern, static void vectorizeLoopIfProfitable(Operation *loop, unsigned depthInPattern,
▲ Show 20 Lines • Show All 1,163 Lines • ▼ Show 20 Lines for (AffineForOp loop : loops[i]) {
return failure(); return failure();
} }
} }
} }
return success(); return success();
} }
namespace mlir {
/// External utility to vectorize affine loops in 'loops' using the n-D /// External utility to vectorize affine loops in 'loops' using the n-D
/// vectorization factors in 'vectorSizes'. By default, each vectorization /// vectorization factors in 'vectorSizes'. By default, each vectorization
/// factor is applied inner-to-outer to the loops of each loop nest. /// factor is applied inner-to-outer to the loops of each loop nest.
/// 'fastestVaryingPattern' can be optionally used to provide a different loop /// 'fastestVaryingPattern' can be optionally used to provide a different loop
/// vectorization order. /// vectorization order.
/// If `reductionLoops` is not empty, the given reduction loops may be /// If `reductionLoops` is not empty, the given reduction loops may be
/// vectorized along the reduction dimension. /// vectorized along the reduction dimension.
/// TODO: Vectorizing reductions is supported only for 1-D vectorization. /// TODO: Vectorizing reductions is supported only for 1-D vectorization.
void vectorizeAffineLoops(Operation *parentOp, DenseSet<Operation *> &loops, void mlir::affine::vectorizeAffineLoops(
ArrayRef<int64_t> vectorSizes, Operation *parentOp, DenseSet<Operation *> &loops,
ArrayRef<int64_t> fastestVaryingPattern, ArrayRef<int64_t> vectorSizes, ArrayRef<int64_t> fastestVaryingPattern,
const ReductionLoopMap &reductionLoops) { const ReductionLoopMap &reductionLoops) {
// Thread-safe RAII local context, BumpPtrAllocator freed on exit. // Thread-safe RAII local context, BumpPtrAllocator freed on exit.
NestedPatternContext mlContext; NestedPatternContext mlContext;
vectorizeLoops(parentOp, loops, vectorSizes, fastestVaryingPattern, vectorizeLoops(parentOp, loops, vectorSizes, fastestVaryingPattern,
reductionLoops); reductionLoops);
} }
/// External utility to vectorize affine loops from a single loop nest using an /// External utility to vectorize affine loops from a single loop nest using an
/// n-D vectorization strategy (see doc in VectorizationStrategy definition). /// n-D vectorization strategy (see doc in VectorizationStrategy definition).
Show All 25 Lines
/// ///
/// loops = {{%i0}, {%i2, %i3}}, to vectorize the outermost and the two /// loops = {{%i0}, {%i2, %i3}}, to vectorize the outermost and the two
/// innermost loops; /// innermost loops;
/// loops = {{%i1}, {%i2, %i3}}, to vectorize the middle and the two innermost /// loops = {{%i1}, {%i2, %i3}}, to vectorize the middle and the two innermost
/// loops; /// loops;
/// loops = {{%i2}}, to vectorize only the first innermost loop; /// loops = {{%i2}}, to vectorize only the first innermost loop;
/// loops = {{%i3}}, to vectorize only the second innermost loop; /// loops = {{%i3}}, to vectorize only the second innermost loop;
/// loops = {{%i1}}, to vectorize only the middle loop. /// loops = {{%i1}}, to vectorize only the middle loop.
LogicalResult LogicalResult mlir::affine::vectorizeAffineLoopNest(
vectorizeAffineLoopNest(std::vector<SmallVector<AffineForOp, 2>> &loops, std::vector<SmallVector<AffineForOp, 2>> &loops,
const VectorizationStrategy &strategy) { const VectorizationStrategy &strategy) {
// Thread-safe RAII local context, BumpPtrAllocator freed on exit. // Thread-safe RAII local context, BumpPtrAllocator freed on exit.
NestedPatternContext mlContext; NestedPatternContext mlContext;
if (failed(verifyLoopNesting(loops))) if (failed(verifyLoopNesting(loops)))
return failure(); return failure();
return vectorizeLoopNest(loops, strategy); return vectorizeLoopNest(loops, strategy);
} }
} // namespace mlir

mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp

Show All 20 Lines
#include "mlir/IR/Operation.h" #include "mlir/IR/Operation.h"
#include "llvm/Support/Debug.h" #include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h" #include "llvm/Support/raw_ostream.h"
#include <optional> #include <optional>
#define DEBUG_TYPE "loop-fusion-utils" #define DEBUG_TYPE "loop-fusion-utils"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
// Gathers all load and store memref accesses in 'opA' into 'values', where // Gathers all load and store memref accesses in 'opA' into 'values', where
// 'values[memref] == true' for each store operation. // 'values[memref] == true' for each store operation.
static void getLoadAndStoreMemRefAccesses(Operation *opA, static void getLoadAndStoreMemRefAccesses(Operation *opA,
DenseMap<Value, bool> &values) { DenseMap<Value, bool> &values) {
opA->walk([&](Operation *op) { opA->walk([&](Operation *op) {
if (auto loadOp = dyn_cast<AffineReadOpInterface>(op)) { if (auto loadOp = dyn_cast<AffineReadOpInterface>(op)) {
if (values.count(loadOp.getMemRef()) == 0) if (values.count(loadOp.getMemRef()) == 0)
▲ Show 20 Lines • Show All 203 Lines • ▼ Show 20 Linesstatic unsigned getMaxLoopDepth(ArrayRef<Operation *> srcOps,
return loopDepth; return loopDepth;
} }
// TODO: Prevent fusion of loop nests with side-effecting operations. // TODO: Prevent fusion of loop nests with side-effecting operations.
// TODO: This pass performs some computation that is the same for all the depths // TODO: This pass performs some computation that is the same for all the depths
// (e.g., getMaxLoopDepth). Implement a version of this utility that processes // (e.g., getMaxLoopDepth). Implement a version of this utility that processes
// all the depths at once or only the legal maximal depth for maximal fusion. // all the depths at once or only the legal maximal depth for maximal fusion.
FusionResult mlir::canFuseLoops(AffineForOp srcForOp, AffineForOp dstForOp, FusionResult mlir::affine::canFuseLoops(AffineForOp srcForOp,
AffineForOp dstForOp,
unsigned dstLoopDepth, unsigned dstLoopDepth,
ComputationSliceState *srcSlice, ComputationSliceState *srcSlice,
FusionStrategy fusionStrategy) { FusionStrategy fusionStrategy) {
// Return 'failure' if 'dstLoopDepth == 0'. // Return 'failure' if 'dstLoopDepth == 0'.
if (dstLoopDepth == 0) { if (dstLoopDepth == 0) {
LLVM_DEBUG(llvm::dbgs() << "Cannot fuse loop nests at depth 0\n"); LLVM_DEBUG(llvm::dbgs() << "Cannot fuse loop nests at depth 0\n");
return FusionResult::FailPrecondition; return FusionResult::FailPrecondition;
} }
// Return 'failure' if 'srcForOp' and 'dstForOp' are not in the same block. // Return 'failure' if 'srcForOp' and 'dstForOp' are not in the same block.
auto *block = srcForOp->getBlock(); auto *block = srcForOp->getBlock();
if (block != dstForOp->getBlock()) { if (block != dstForOp->getBlock()) {
Show All 38 Lines if (fusionStrategy.getStrategy() == FusionStrategy::ProducerConsumer) {
if (getMaxLoopDepth(opsA, opsB) < dstLoopDepth) { if (getMaxLoopDepth(opsA, opsB) < dstLoopDepth) {
LLVM_DEBUG(llvm::dbgs() << "Fusion would violate loop dependences\n"); LLVM_DEBUG(llvm::dbgs() << "Fusion would violate loop dependences\n");
return FusionResult::FailFusionDependence; return FusionResult::FailFusionDependence;
} }
} }
// Calculate the number of common loops surrounding 'srcForOp' and 'dstForOp'. // Calculate the number of common loops surrounding 'srcForOp' and 'dstForOp'.
unsigned numCommonLoops = unsigned numCommonLoops =
mlir::getNumCommonSurroundingLoops(*srcForOp, *dstForOp); affine::getNumCommonSurroundingLoops(*srcForOp, *dstForOp);
// Filter out ops in 'opsA' to compute the slice union based on the // Filter out ops in 'opsA' to compute the slice union based on the
// assumptions made by the fusion strategy. // assumptions made by the fusion strategy.
SmallVector<Operation *, 4> strategyOpsA; SmallVector<Operation *, 4> strategyOpsA;
switch (fusionStrategy.getStrategy()) { switch (fusionStrategy.getStrategy()) {
case FusionStrategy::Generic: case FusionStrategy::Generic:
// Generic fusion. Take into account all the memory operations to compute // Generic fusion. Take into account all the memory operations to compute
// the slice union. // the slice union.
Show All 15 Lines for (Operation *op : opsA) {
if (load && load.getMemRef() == fusionStrategy.getSiblingFusionMemRef()) if (load && load.getMemRef() == fusionStrategy.getSiblingFusionMemRef())
strategyOpsA.push_back(op); strategyOpsA.push_back(op);
} }
break; break;
} }
// Compute union of computation slices computed between all pairs of ops // Compute union of computation slices computed between all pairs of ops
// from 'forOpA' and 'forOpB'. // from 'forOpA' and 'forOpB'.
SliceComputationResult sliceComputationResult = SliceComputationResult sliceComputationResult = affine::computeSliceUnion(
mlir::computeSliceUnion(strategyOpsA, opsB, dstLoopDepth, numCommonLoops, strategyOpsA, opsB, dstLoopDepth, numCommonLoops,
isSrcForOpBeforeDstForOp, srcSlice); isSrcForOpBeforeDstForOp, srcSlice);
if (sliceComputationResult.value == SliceComputationResult::GenericFailure) { if (sliceComputationResult.value == SliceComputationResult::GenericFailure) {
LLVM_DEBUG(llvm::dbgs() << "computeSliceUnion failed\n"); LLVM_DEBUG(llvm::dbgs() << "computeSliceUnion failed\n");
return FusionResult::FailPrecondition; return FusionResult::FailPrecondition;
} }
if (sliceComputationResult.value == if (sliceComputationResult.value ==
SliceComputationResult::IncorrectSliceFailure) { SliceComputationResult::IncorrectSliceFailure) {
LLVM_DEBUG(llvm::dbgs() << "Incorrect slice computation\n"); LLVM_DEBUG(llvm::dbgs() << "Incorrect slice computation\n");
return FusionResult::FailIncorrectSlice; return FusionResult::FailIncorrectSlice;
} }
return FusionResult::Success; return FusionResult::Success;
} }
/// Patch the loop body of a forOp that is a single iteration reduction loop /// Patch the loop body of a forOp that is a single iteration reduction loop
/// into its containing block. /// into its containing block.
LogicalResult promoteSingleIterReductionLoop(AffineForOp forOp, static LogicalResult promoteSingleIterReductionLoop(AffineForOp forOp,
bool siblingFusionUser) { bool siblingFusionUser) {
// Check if the reduction loop is a single iteration loop. // Check if the reduction loop is a single iteration loop.
std::optional<uint64_t> tripCount = getConstantTripCount(forOp); std::optional<uint64_t> tripCount = getConstantTripCount(forOp);
if (!tripCount || *tripCount != 1) if (!tripCount || *tripCount != 1)
return failure(); return failure();
auto iterOperands = forOp.getIterOperands(); auto iterOperands = forOp.getIterOperands();
auto *parentOp = forOp->getParentOp(); auto *parentOp = forOp->getParentOp();
if (!isa<AffineForOp>(parentOp)) if (!isa<AffineForOp>(parentOp))
return failure(); return failure();
▲ Show 20 Lines • Show All 45 Lines • ▼ Show 20 Lines parentBlock->getOperations().splice(Block::iterator(forOp),
forOp.getBody()->getOperations()); forOp.getBody()->getOperations());
forOp.erase(); forOp.erase();
parentForOp.erase(); parentForOp.erase();
return success(); return success();
} }
/// Fuses 'srcForOp' into 'dstForOp' with destination loop block insertion point /// Fuses 'srcForOp' into 'dstForOp' with destination loop block insertion point
/// and source slice loop bounds specified in 'srcSlice'. /// and source slice loop bounds specified in 'srcSlice'.
void mlir::fuseLoops(AffineForOp srcForOp, AffineForOp dstForOp, void mlir::affine::fuseLoops(AffineForOp srcForOp, AffineForOp dstForOp,
const ComputationSliceState &srcSlice, const ComputationSliceState &srcSlice,
bool isInnermostSiblingInsertion) { bool isInnermostSiblingInsertion) {
// Clone 'srcForOp' into 'dstForOp' at 'srcSlice->insertPoint'. // Clone 'srcForOp' into 'dstForOp' at 'srcSlice->insertPoint'.
OpBuilder b(srcSlice.insertPoint->getBlock(), srcSlice.insertPoint); OpBuilder b(srcSlice.insertPoint->getBlock(), srcSlice.insertPoint);
IRMapping mapper; IRMapping mapper;
b.clone(*srcForOp, mapper); b.clone(*srcForOp, mapper);
// Update 'sliceLoopNest' upper and lower bounds from computed 'srcSlice'. // Update 'sliceLoopNest' upper and lower bounds from computed 'srcSlice'.
SmallVector<AffineForOp, 4> sliceLoops; SmallVector<AffineForOp, 4> sliceLoops;
for (unsigned i = 0, e = srcSlice.ivs.size(); i < e; ++i) { for (unsigned i = 0, e = srcSlice.ivs.size(); i < e; ++i) {
Show All 30 Lines else
// Promote any single iteration slice loops. // Promote any single iteration slice loops.
(void)promoteIfSingleIteration(forOp); (void)promoteIfSingleIteration(forOp);
} }
} }
/// Collect loop nest statistics (eg. loop trip count and operation count) /// Collect loop nest statistics (eg. loop trip count and operation count)
/// in 'stats' for loop nest rooted at 'forOp'. Returns true on success, /// in 'stats' for loop nest rooted at 'forOp'. Returns true on success,
/// returns false otherwise. /// returns false otherwise.
bool mlir::getLoopNestStats(AffineForOp forOpRoot, LoopNestStats *stats) { bool mlir::affine::getLoopNestStats(AffineForOp forOpRoot,
LoopNestStats *stats) {
auto walkResult = forOpRoot.walk([&](AffineForOp forOp) { auto walkResult = forOpRoot.walk([&](AffineForOp forOp) {
auto *childForOp = forOp.getOperation(); auto *childForOp = forOp.getOperation();
auto *parentForOp = forOp->getParentOp(); auto *parentForOp = forOp->getParentOp();
if (forOp != forOpRoot) { if (forOp != forOpRoot) {
if (!isa<AffineForOp>(parentForOp)) { if (!isa<AffineForOp>(parentForOp)) {
LLVM_DEBUG(llvm::dbgs() << "Expected parent AffineForOp\n"); LLVM_DEBUG(llvm::dbgs() << "Expected parent AffineForOp\n");
return WalkResult::interrupt(); return WalkResult::interrupt();
} }
▲ Show 20 Lines • Show All 71 Lines • ▼ Show 20 Linesstatic int64_t getComputeCostHelper(
// Returns the total number of dynamic instances of operations in loop body. // Returns the total number of dynamic instances of operations in loop body.
return tripCount * opCount; return tripCount * opCount;
} }
/// Computes the total cost of the loop nest rooted at 'forOp' using 'stats'. /// Computes the total cost of the loop nest rooted at 'forOp' using 'stats'.
/// Currently, the total cost is computed by counting the total operation /// Currently, the total cost is computed by counting the total operation
/// instance count (i.e. total number of operations in the loop body * loop /// instance count (i.e. total number of operations in the loop body * loop
/// trip count) for the entire loop nest. /// trip count) for the entire loop nest.
int64_t mlir::getComputeCost(AffineForOp forOp, LoopNestStats &stats) { int64_t mlir::affine::getComputeCost(AffineForOp forOp, LoopNestStats &stats) {
return getComputeCostHelper(forOp, stats, return getComputeCostHelper(forOp, stats,
/*tripCountOverrideMap=*/nullptr, /*tripCountOverrideMap=*/nullptr,
/*computeCostMap=*/nullptr); /*computeCostMap=*/nullptr);
} }
/// Computes and returns in 'computeCost', the total compute cost of fusing the /// Computes and returns in 'computeCost', the total compute cost of fusing the
/// 'slice' of the loop nest rooted at 'srcForOp' into 'dstForOp'. Currently, /// 'slice' of the loop nest rooted at 'srcForOp' into 'dstForOp'. Currently,
/// the total cost is computed by counting the total operation instance count /// the total cost is computed by counting the total operation instance count
/// (i.e. total number of operations in the loop body * loop trip count) for /// (i.e. total number of operations in the loop body * loop trip count) for
/// the entire loop nest. /// the entire loop nest.
bool mlir::getFusionComputeCost(AffineForOp srcForOp, LoopNestStats &srcStats, bool mlir::affine::getFusionComputeCost(AffineForOp srcForOp,
AffineForOp dstForOp, LoopNestStats &dstStats, LoopNestStats &srcStats,
AffineForOp dstForOp,
LoopNestStats &dstStats,
const ComputationSliceState &slice, const ComputationSliceState &slice,
int64_t *computeCost) { int64_t *computeCost) {
llvm::SmallDenseMap<Operation *, uint64_t, 8> sliceTripCountMap; llvm::SmallDenseMap<Operation *, uint64_t, 8> sliceTripCountMap;
DenseMap<Operation *, int64_t> computeCostMap; DenseMap<Operation *, int64_t> computeCostMap;
// Build trip count map for computation slice. // Build trip count map for computation slice.
if (!buildSliceTripCountMap(slice, &sliceTripCountMap)) if (!buildSliceTripCountMap(slice, &sliceTripCountMap))
return false; return false;
// Checks whether a store to load forwarding will happen. // Checks whether a store to load forwarding will happen.
int64_t sliceIterationCount = getSliceIterationCount(sliceTripCountMap); int64_t sliceIterationCount = getSliceIterationCount(sliceTripCountMap);
▲ Show 20 Lines • Show All 50 Lines • ▼ Show 20 Lines *computeCost =
getComputeCostHelper(dstForOp, dstStats, getComputeCostHelper(dstForOp, dstStats,
/*tripCountOverrideMap=*/nullptr, &computeCostMap); /*tripCountOverrideMap=*/nullptr, &computeCostMap);
return true; return true;
} }
/// Returns in 'producerConsumerMemrefs' the memrefs involved in a /// Returns in 'producerConsumerMemrefs' the memrefs involved in a
/// producer-consumer dependence between write ops in 'srcOps' and read ops in /// producer-consumer dependence between write ops in 'srcOps' and read ops in
/// 'dstOps'. /// 'dstOps'.
void mlir::gatherProducerConsumerMemrefs( void mlir::affine::gatherProducerConsumerMemrefs(
ArrayRef<Operation *> srcOps, ArrayRef<Operation *> dstOps, ArrayRef<Operation *> srcOps, ArrayRef<Operation *> dstOps,
DenseSet<Value> &producerConsumerMemrefs) { DenseSet<Value> &producerConsumerMemrefs) {
// Gather memrefs from stores in 'srcOps'. // Gather memrefs from stores in 'srcOps'.
DenseSet<Value> srcStoreMemRefs; DenseSet<Value> srcStoreMemRefs;
for (Operation *op : srcOps) for (Operation *op : srcOps)
if (auto storeOp = dyn_cast<AffineWriteOpInterface>(op)) if (auto storeOp = dyn_cast<AffineWriteOpInterface>(op))
srcStoreMemRefs.insert(storeOp.getMemRef()); srcStoreMemRefs.insert(storeOp.getMemRef());
// Compute the intersection between memrefs from stores in 'srcOps' and // Compute the intersection between memrefs from stores in 'srcOps' and
// memrefs from loads in 'dstOps'. // memrefs from loads in 'dstOps'.
for (Operation *op : dstOps) for (Operation *op : dstOps)
if (auto loadOp = dyn_cast<AffineReadOpInterface>(op)) if (auto loadOp = dyn_cast<AffineReadOpInterface>(op))
if (srcStoreMemRefs.count(loadOp.getMemRef()) > 0) if (srcStoreMemRefs.count(loadOp.getMemRef()) > 0)
producerConsumerMemrefs.insert(loadOp.getMemRef()); producerConsumerMemrefs.insert(loadOp.getMemRef());
} }

mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp

Show All 28 Lines
#include "llvm/ADT/SmallPtrSet.h" #include "llvm/ADT/SmallPtrSet.h"
#include "llvm/Support/Debug.h" #include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h" #include "llvm/Support/raw_ostream.h"
#include <optional> #include <optional>
#define DEBUG_TYPE "loop-utils" #define DEBUG_TYPE "loop-utils"
using namespace mlir; using namespace mlir;
using namespace affine;
using namespace presburger; using namespace presburger;
using llvm::SmallMapVector; using llvm::SmallMapVector;
namespace { namespace {
// This structure is to pass and return sets of loop parameters without // This structure is to pass and return sets of loop parameters without
// confusing the order. // confusing the order.
struct LoopParams { struct LoopParams {
Value lowerBound; Value lowerBound;
▲ Show 20 Lines • Show All 78 Lines • ▼ Show 20 Linesstatic void replaceIterArgsAndYieldResults(AffineForOp forOp) {
auto innerResults = forOp.getBody()->getTerminator()->getOperands(); auto innerResults = forOp.getBody()->getTerminator()->getOperands();
for (auto e : llvm::zip(outerResults, innerResults)) for (auto e : llvm::zip(outerResults, innerResults))
std::get<0>(e).replaceAllUsesWith(std::get<1>(e)); std::get<0>(e).replaceAllUsesWith(std::get<1>(e));
} }
/// Promotes the loop body of a forOp to its containing block if the forOp /// Promotes the loop body of a forOp to its containing block if the forOp
/// was known to have a single iteration. /// was known to have a single iteration.
// TODO: extend this for arbitrary affine bounds. // TODO: extend this for arbitrary affine bounds.
LogicalResult mlir::promoteIfSingleIteration(AffineForOp forOp) { LogicalResult mlir::affine::promoteIfSingleIteration(AffineForOp forOp) {
std::optional<uint64_t> tripCount = getConstantTripCount(forOp); std::optional<uint64_t> tripCount = getConstantTripCount(forOp);
if (!tripCount || *tripCount != 1) if (!tripCount || *tripCount != 1)
return failure(); return failure();
if (forOp.getLowerBoundMap().getNumResults() != 1) if (forOp.getLowerBoundMap().getNumResults() != 1)
return failure(); return failure();
// Replaces all IV uses to its single iteration value. // Replaces all IV uses to its single iteration value.
▲ Show 20 Lines • Show All 87 Lines • ▼ Show 20 Lines
// for better register reuse, locality or parallelism. As such, the shifts are // for better register reuse, locality or parallelism. As such, the shifts are
// typically expected to be at most of the order of the number of operations. // typically expected to be at most of the order of the number of operations.
// This method should not be used as a substitute for loop distribution/fission. // This method should not be used as a substitute for loop distribution/fission.
// This method uses an algorithm// in time linear in the number of operations // This method uses an algorithm// in time linear in the number of operations
// in the body of the for loop - (using the 'sweep line' paradigm). This method // in the body of the for loop - (using the 'sweep line' paradigm). This method
// asserts preservation of SSA dominance. A check for that as well as that for // asserts preservation of SSA dominance. A check for that as well as that for
// memory-based dependence preservation check rests with the users of this // memory-based dependence preservation check rests with the users of this
// method. // method.
LogicalResult mlir::affineForOpBodySkew(AffineForOp forOp, LogicalResult mlir::affine::affineForOpBodySkew(AffineForOp forOp,
ArrayRef<uint64_t> shifts, ArrayRef<uint64_t> shifts,
bool unrollPrologueEpilogue) { bool unrollPrologueEpilogue) {
assert(forOp.getBody()->getOperations().size() == shifts.size() && assert(forOp.getBody()->getOperations().size() == shifts.size() &&
"too few/many shifts"); "too few/many shifts");
if (forOp.getBody()->begin() == std::prev(forOp.getBody()->end())) if (forOp.getBody()->begin() == std::prev(forOp.getBody()->end()))
return success(); return success();
// If the trip counts aren't constant, we would need versioning and // If the trip counts aren't constant, we would need versioning and
// conditional guards (or context information to prevent such versioning). The // conditional guards (or context information to prevent such versioning). The
// better way to pipeline for such loops is to first tile them and extract // better way to pipeline for such loops is to first tile them and extract
▲ Show 20 Lines • Show All 112 Lines • ▼ Show 20 Lines if (unrollPrologueEpilogue && prologue)
(void)loopUnrollFull(prologue); (void)loopUnrollFull(prologue);
if (unrollPrologueEpilogue && !epilogue && epilogue != prologue) if (unrollPrologueEpilogue && !epilogue && epilogue != prologue)
(void)loopUnrollFull(epilogue); (void)loopUnrollFull(epilogue);
return success(); return success();
} }
/// Checks whether a loop nest is hyper-rectangular or not. /// Checks whether a loop nest is hyper-rectangular or not.
LogicalResult checkIfHyperRectangular(MutableArrayRef<AffineForOp> input) { static LogicalResult
checkIfHyperRectangular(MutableArrayRef<AffineForOp> input) {
FlatAffineValueConstraints cst; FlatAffineValueConstraints cst;
SmallVector<Operation *, 8> ops(input.begin(), input.end()); SmallVector<Operation *, 8> ops(input.begin(), input.end());
// 0-d or 1-d is trivially hyper-rectangular. // 0-d or 1-d is trivially hyper-rectangular.
if (input.size() <= 1) if (input.size() <= 1)
return success(); return success();
if (failed(getIndexSet(ops, &cst))) { if (failed(getIndexSet(ops, &cst))) {
LLVM_DEBUG(llvm::dbgs() << "Index set computation failed!\n"); LLVM_DEBUG(llvm::dbgs() << "Index set computation failed!\n");
return failure(); return failure();
} }
if (!cst.isHyperRectangular(0, input.size())) { if (!cst.isHyperRectangular(0, input.size())) {
LLVM_DEBUG(llvm::dbgs() LLVM_DEBUG(llvm::dbgs()
<< "Non-hyperrectangular nests not supported for tiling!\n"); << "Non-hyperrectangular nests not supported for tiling!\n");
return failure(); return failure();
} }
return success(); return success();
} }
/// Check if the input nest is supported for tiling and whether tiling would be /// Check if the input nest is supported for tiling and whether tiling would be
/// legal or not. /// legal or not.
template <typename t> template <typename t>
LogicalResult performPreTilingChecks(MutableArrayRef<AffineForOp> input, static LogicalResult performPreTilingChecks(MutableArrayRef<AffineForOp> input,
ArrayRef<t> tileSizes) { ArrayRef<t> tileSizes) {
assert(input.size() == tileSizes.size() && "Too few/many tile sizes"); assert(input.size() == tileSizes.size() && "Too few/many tile sizes");
if (llvm::any_of(input, if (llvm::any_of(input,
[](AffineForOp op) { return op.getNumResults() > 0; })) { [](AffineForOp op) { return op.getNumResults() > 0; })) {
LLVM_DEBUG(llvm::dbgs() LLVM_DEBUG(llvm::dbgs()
<< "Cannot tile nest where a loop has yield values\n"); << "Cannot tile nest where a loop has yield values\n");
return failure(); return failure();
} }
Show All 16 Linesstatic void moveLoopBodyImpl(AffineForOp src, AffineForOp dest,
Block::iterator loc) { Block::iterator loc) {
auto &ops = src.getBody()->getOperations(); auto &ops = src.getBody()->getOperations();
dest.getBody()->getOperations().splice(loc, ops, ops.begin(), dest.getBody()->getOperations().splice(loc, ops, ops.begin(),
std::prev(ops.end())); std::prev(ops.end()));
} }
/// Move the loop body of AffineForOp 'src' from 'src' to the start of dest /// Move the loop body of AffineForOp 'src' from 'src' to the start of dest
/// body. /// body.
void moveLoopBody(AffineForOp src, AffineForOp dest) { static void moveLoopBody(AffineForOp src, AffineForOp dest) {
moveLoopBodyImpl(src, dest, dest.getBody()->begin()); moveLoopBodyImpl(src, dest, dest.getBody()->begin());
} }
/// Constructs tiled loop nest, without setting the loop bounds and move the /// Constructs tiled loop nest, without setting the loop bounds and move the
/// body of the original loop nest to the tiled loop nest. /// body of the original loop nest to the tiled loop nest.
void constructTiledLoopNest(MutableArrayRef<AffineForOp> origLoops, static void constructTiledLoopNest(MutableArrayRef<AffineForOp> origLoops,
AffineForOp rootAffineForOp, unsigned width, AffineForOp rootAffineForOp, unsigned width,
MutableArrayRef<AffineForOp> tiledLoops) { MutableArrayRef<AffineForOp> tiledLoops) {
Location loc = rootAffineForOp.getLoc(); Location loc = rootAffineForOp.getLoc();
// The outermost among the loops as we add more.. // The outermost among the loops as we add more..
Operation *topLoop = rootAffineForOp.getOperation(); Operation *topLoop = rootAffineForOp.getOperation();
AffineForOp innermostPointLoop; AffineForOp innermostPointLoop;
// Add intra-tile (or point) loops. // Add intra-tile (or point) loops.
for (unsigned i = 0; i < width; i++) { for (unsigned i = 0; i < width; i++) {
▲ Show 20 Lines • Show All 330 Lines • ▼ Show 20 Lines if (mayBeConstantCount && *mayBeConstantCount < tileSizes[i]) {
AffineMap ubMap = AffineMap ubMap =
AffineMap::get(1, 0, dim + tileSizes[i] * origLoops[i].getStep()); AffineMap::get(1, 0, dim + tileSizes[i] * origLoops[i].getStep());
newLoops[width + i].setUpperBound(newLoops[i].getInductionVar(), ubMap); newLoops[width + i].setUpperBound(newLoops[i].getInductionVar(), ubMap);
} }
} }
} }
LogicalResult LogicalResult
mlir::tilePerfectlyNested(MutableArrayRef<AffineForOp> input, mlir::affine::tilePerfectlyNested(MutableArrayRef<AffineForOp> input,
ArrayRef<unsigned> tileSizes, ArrayRef<unsigned> tileSizes,
SmallVectorImpl<AffineForOp> *tiledNest) { SmallVectorImpl<AffineForOp> *tiledNest) {
if (input.empty()) if (input.empty())
return success(); return success();
if (failed(performPreTilingChecks(input, tileSizes))) if (failed(performPreTilingChecks(input, tileSizes)))
return failure(); return failure();
MutableArrayRef<AffineForOp> origLoops = input; MutableArrayRef<AffineForOp> origLoops = input;
AffineForOp rootAffineForOp = origLoops[0]; AffineForOp rootAffineForOp = origLoops[0];
Show All 24 Linesmlir::affine::tilePerfectlyNested(MutableArrayRef<AffineForOp> input,
return success(); return success();
} }
/// Tiles the specified band of perfectly nested loops creating tile-space /// Tiles the specified band of perfectly nested loops creating tile-space
/// loops and intra-tile loops, using SSA values as tiling parameters. A band /// loops and intra-tile loops, using SSA values as tiling parameters. A band
/// is a contiguous set of loops. /// is a contiguous set of loops.
// TODO: handle non hyper-rectangular spaces. // TODO: handle non hyper-rectangular spaces.
LogicalResult LogicalResult mlir::affine::tilePerfectlyNestedParametric(
mlir::tilePerfectlyNestedParametric(MutableArrayRef<AffineForOp> input, MutableArrayRef<AffineForOp> input, ArrayRef<Value> tileSizes,
ArrayRef<Value> tileSizes,
SmallVectorImpl<AffineForOp> *tiledNest) { SmallVectorImpl<AffineForOp> *tiledNest) {
if (input.empty()) if (input.empty())
return success(); return success();
if (failed(performPreTilingChecks(input, tileSizes))) if (failed(performPreTilingChecks(input, tileSizes)))
return failure(); return failure();
MutableArrayRef<AffineForOp> origLoops = input; MutableArrayRef<AffineForOp> origLoops = input;
AffineForOp rootAffineForOp = origLoops[0]; AffineForOp rootAffineForOp = origLoops[0];
Show All 23 LinesLogicalResult mlir::affine::tilePerfectlyNestedParametric(
return success(); return success();
} }
/// Get perfectly nested sequence of loops starting at root of loop nest /// Get perfectly nested sequence of loops starting at root of loop nest
/// (the first op being another AffineFor, and the second op - a terminator). /// (the first op being another AffineFor, and the second op - a terminator).
/// A loop is perfectly nested iff: the first op in the loop's body is another /// A loop is perfectly nested iff: the first op in the loop's body is another
/// AffineForOp, and the second op is a terminator). /// AffineForOp, and the second op is a terminator).
void mlir::getPerfectlyNestedLoops(SmallVectorImpl<AffineForOp> &nestedLoops, void mlir::affine::getPerfectlyNestedLoops(
AffineForOp root) { SmallVectorImpl<AffineForOp> &nestedLoops, AffineForOp root) {
for (unsigned i = 0; i < std::numeric_limits<unsigned>::max(); ++i) { for (unsigned i = 0; i < std::numeric_limits<unsigned>::max(); ++i) {
nestedLoops.push_back(root); nestedLoops.push_back(root);
Block &body = root.getRegion().front(); Block &body = root.getRegion().front();
if (body.begin() != std::prev(body.end(), 2)) if (body.begin() != std::prev(body.end(), 2))
return; return;
root = dyn_cast<AffineForOp>(&body.front()); root = dyn_cast<AffineForOp>(&body.front());
if (!root) if (!root)
return; return;
} }
} }
/// Identify valid and profitable bands of loops to tile. This is currently just /// Identify valid and profitable bands of loops to tile. This is currently just
/// a temporary placeholder to test the mechanics of tiled code generation. /// a temporary placeholder to test the mechanics of tiled code generation.
/// Returns all maximal outermost perfect loop nests to tile. /// Returns all maximal outermost perfect loop nests to tile.
void mlir::getTileableBands(func::FuncOp f, void mlir::affine::getTileableBands(
std::vector<SmallVector<AffineForOp, 6>> *bands) { func::FuncOp f, std::vector<SmallVector<AffineForOp, 6>> *bands) {
// Get maximal perfect nest of 'affine.for' insts starting from root // Get maximal perfect nest of 'affine.for' insts starting from root
// (inclusive). // (inclusive).
for (AffineForOp forOp : f.getOps<AffineForOp>()) { for (AffineForOp forOp : f.getOps<AffineForOp>()) {
SmallVector<AffineForOp, 6> band; SmallVector<AffineForOp, 6> band;
getPerfectlyNestedLoops(band, forOp); getPerfectlyNestedLoops(band, forOp);
bands->push_back(band); bands->push_back(band);
} }
} }
/// Unrolls this loop completely. /// Unrolls this loop completely.
LogicalResult mlir::loopUnrollFull(AffineForOp forOp) { LogicalResult mlir::affine::loopUnrollFull(AffineForOp forOp) {
std::optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp); std::optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
if (mayBeConstantTripCount.has_value()) { if (mayBeConstantTripCount.has_value()) {
uint64_t tripCount = *mayBeConstantTripCount; uint64_t tripCount = *mayBeConstantTripCount;
if (tripCount == 0) if (tripCount == 0)
return success(); return success();
if (tripCount == 1) if (tripCount == 1)
return promoteIfSingleIteration(forOp); return promoteIfSingleIteration(forOp);
return loopUnrollByFactor(forOp, tripCount); return loopUnrollByFactor(forOp, tripCount);
} }
return failure(); return failure();
} }
/// Unrolls this loop by the specified factor or by the trip count (if constant) /// Unrolls this loop by the specified factor or by the trip count (if constant)
/// whichever is lower. /// whichever is lower.
LogicalResult mlir::loopUnrollUpToFactor(AffineForOp forOp, LogicalResult mlir::affine::loopUnrollUpToFactor(AffineForOp forOp,
uint64_t unrollFactor) { uint64_t unrollFactor) {
std::optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp); std::optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
if (mayBeConstantTripCount.has_value() && if (mayBeConstantTripCount.has_value() &&
*mayBeConstantTripCount < unrollFactor) *mayBeConstantTripCount < unrollFactor)
return loopUnrollByFactor(forOp, *mayBeConstantTripCount); return loopUnrollByFactor(forOp, *mayBeConstantTripCount);
return loopUnrollByFactor(forOp, unrollFactor); return loopUnrollByFactor(forOp, unrollFactor);
} }
/// Generates unrolled copies of AffineForOp 'loopBodyBlock', with associated /// Generates unrolled copies of AffineForOp 'loopBodyBlock', with associated
▲ Show 20 Lines • Show All 90 Lines • ▼ Show 20 Linesstatic LogicalResult generateCleanupLoopForUnroll(AffineForOp forOp,
// Adjust upper bound of the original loop; this is the same as the lower // Adjust upper bound of the original loop; this is the same as the lower
// bound of the cleanup loop. // bound of the cleanup loop.
forOp.setUpperBound(cleanupOperands, cleanupMap); forOp.setUpperBound(cleanupOperands, cleanupMap);
return success(); return success();
} }
/// Unrolls this loop by the specified factor. Returns success if the loop /// Unrolls this loop by the specified factor. Returns success if the loop
/// is successfully unrolled. /// is successfully unrolled.
LogicalResult mlir::loopUnrollByFactor( LogicalResult mlir::affine::loopUnrollByFactor(
AffineForOp forOp, uint64_t unrollFactor, AffineForOp forOp, uint64_t unrollFactor,
function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn, function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn,
bool cleanUpUnroll) { bool cleanUpUnroll) {
assert(unrollFactor > 0 && "unroll factor should be positive"); assert(unrollFactor > 0 && "unroll factor should be positive");
std::optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp); std::optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
if (unrollFactor == 1) { if (unrollFactor == 1) {
if (mayBeConstantTripCount && *mayBeConstantTripCount == 1 && if (mayBeConstantTripCount && *mayBeConstantTripCount == 1 &&
▲ Show 20 Lines • Show All 50 Lines • ▼ Show 20 Lines generateUnrolledLoop(
/*annotateFn=*/annotateFn, /*annotateFn=*/annotateFn,
/*iterArgs=*/iterArgs, /*yieldedValues=*/yieldedValues); /*iterArgs=*/iterArgs, /*yieldedValues=*/yieldedValues);
// Promote the loop body up if this has turned into a single iteration loop. // Promote the loop body up if this has turned into a single iteration loop.
(void)promoteIfSingleIteration(forOp); (void)promoteIfSingleIteration(forOp);
return success(); return success();
} }
LogicalResult mlir::loopUnrollJamUpToFactor(AffineForOp forOp, LogicalResult mlir::affine::loopUnrollJamUpToFactor(AffineForOp forOp,
uint64_t unrollJamFactor) { uint64_t unrollJamFactor) {
std::optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp); std::optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
if (mayBeConstantTripCount.has_value() && if (mayBeConstantTripCount.has_value() &&
*mayBeConstantTripCount < unrollJamFactor) *mayBeConstantTripCount < unrollJamFactor)
return loopUnrollJamByFactor(forOp, *mayBeConstantTripCount); return loopUnrollJamByFactor(forOp, *mayBeConstantTripCount);
return loopUnrollJamByFactor(forOp, unrollJamFactor); return loopUnrollJamByFactor(forOp, unrollJamFactor);
} }
/// Check if all control operands of all loops are defined outside of `forOp` /// Check if all control operands of all loops are defined outside of `forOp`
Show All 33 Lines for (auto it = block.begin(), e = std::prev(block.end()); it != e;) {
// Process all for ops that appear next. // Process all for ops that appear next.
while (it != e && isa<AffineForOp>(&*it)) while (it != e && isa<AffineForOp>(&*it))
walk(&*it++); walk(&*it++);
} }
} }
}; };
/// Unrolls and jams this loop by the specified factor. /// Unrolls and jams this loop by the specified factor.
LogicalResult mlir::loopUnrollJamByFactor(AffineForOp forOp, LogicalResult mlir::affine::loopUnrollJamByFactor(AffineForOp forOp,
uint64_t unrollJamFactor) { uint64_t unrollJamFactor) {
assert(unrollJamFactor > 0 && "unroll jam factor should be positive"); assert(unrollJamFactor > 0 && "unroll jam factor should be positive");
std::optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp); std::optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
if (unrollJamFactor == 1) { if (unrollJamFactor == 1) {
if (mayBeConstantTripCount && *mayBeConstantTripCount == 1 && if (mayBeConstantTripCount && *mayBeConstantTripCount == 1 &&
failed(promoteIfSingleIteration(forOp))) failed(promoteIfSingleIteration(forOp)))
return failure(); return failure();
return success(); return success();
▲ Show 20 Lines • Show All 65 Lines • ▼ Show 20 Lines for (AffineForOp oldForOp : loopsWithIterArgs) {
// fix iterOperands and yield operands after cloning of sub-blocks. // fix iterOperands and yield operands after cloning of sub-blocks.
for (unsigned i = unrollJamFactor - 1; i >= 1; --i) { for (unsigned i = unrollJamFactor - 1; i >= 1; --i) {
dupIterOperands.append(oldIterOperands.begin(), oldIterOperands.end()); dupIterOperands.append(oldIterOperands.begin(), oldIterOperands.end());
dupIterArgs.append(oldIterArgs.begin(), oldIterArgs.end()); dupIterArgs.append(oldIterArgs.begin(), oldIterArgs.end());
dupYieldOperands.append(oldYieldOperands.begin(), oldYieldOperands.end()); dupYieldOperands.append(oldYieldOperands.begin(), oldYieldOperands.end());
} }
// Create a new loop with additional iterOperands, iter_args and yield // Create a new loop with additional iterOperands, iter_args and yield
// operands. This new loop will take the loop body of the original loop. // operands. This new loop will take the loop body of the original loop.
AffineForOp newForOp = mlir::replaceForOpWithNewYields( AffineForOp newForOp = affine::replaceForOpWithNewYields(
builder, oldForOp, dupIterOperands, dupYieldOperands, dupIterArgs); builder, oldForOp, dupIterOperands, dupYieldOperands, dupIterArgs);
newLoopsWithIterArgs.push_back(newForOp); newLoopsWithIterArgs.push_back(newForOp);
// `forOp` has been replaced with a new loop. // `forOp` has been replaced with a new loop.
if (oldForOp == forOp) if (oldForOp == forOp)
forOp = newForOp; forOp = newForOp;
assert(oldForOp.use_empty() && "old for op should not have any user"); assert(oldForOp.use_empty() && "old for op should not have any user");
oldForOp.erase(); oldForOp.erase();
// Update `operandMaps` for `newForOp` iterArgs and results. // Update `operandMaps` for `newForOp` iterArgs and results.
▲ Show 20 Lines • Show All 94 Lines • ▼ Show 20 LinesLogicalResult mlir::affine::loopUnrollJamByFactor(AffineForOp forOp,
// Promote the loop body up if this has turned into a single iteration loop. // Promote the loop body up if this has turned into a single iteration loop.
(void)promoteIfSingleIteration(forOp); (void)promoteIfSingleIteration(forOp);
return success(); return success();
} }
/// Performs loop interchange on 'forOpA' and 'forOpB', where 'forOpB' is /// Performs loop interchange on 'forOpA' and 'forOpB', where 'forOpB' is
/// nested within 'forOpA' as the only non-terminator operation in its block. /// nested within 'forOpA' as the only non-terminator operation in its block.
void mlir::interchangeLoops(AffineForOp forOpA, AffineForOp forOpB) { void mlir::affine::interchangeLoops(AffineForOp forOpA, AffineForOp forOpB) {
assert(&*forOpA.getBody()->begin() == forOpB.getOperation()); assert(&*forOpA.getBody()->begin() == forOpB.getOperation());
auto &forOpABody = forOpA.getBody()->getOperations(); auto &forOpABody = forOpA.getBody()->getOperations();
auto &forOpBBody = forOpB.getBody()->getOperations(); auto &forOpBBody = forOpB.getBody()->getOperations();
// 1) Splice forOpA's non-terminator operations (which is just forOpB) just // 1) Splice forOpA's non-terminator operations (which is just forOpB) just
// before forOpA (in ForOpA's parent's block) this should leave 'forOpA's // before forOpA (in ForOpA's parent's block) this should leave 'forOpA's
// body containing only the terminator. // body containing only the terminator.
forOpA->getBlock()->getOperations().splice(Block::iterator(forOpA), forOpA->getBlock()->getOperations().splice(Block::iterator(forOpA),
Show All 40 Lines for (unsigned j = 0; j < maxLoopDepth; ++j) {
return false; return false;
} }
} }
return true; return true;
} }
/// Checks if the loop interchange permutation 'loopPermMap' of the perfectly /// Checks if the loop interchange permutation 'loopPermMap' of the perfectly
/// nested sequence of loops in 'loops' would violate dependences. /// nested sequence of loops in 'loops' would violate dependences.
bool mlir::isValidLoopInterchangePermutation(ArrayRef<AffineForOp> loops, bool mlir::affine::isValidLoopInterchangePermutation(
ArrayRef<unsigned> loopPermMap) { ArrayRef<AffineForOp> loops, ArrayRef<unsigned> loopPermMap) {
// Gather dependence components for dependences between all ops in loop nest // Gather dependence components for dependences between all ops in loop nest
// rooted at 'loops[0]', at loop depths in range [1, maxLoopDepth]. // rooted at 'loops[0]', at loop depths in range [1, maxLoopDepth].
assert(loopPermMap.size() == loops.size()); assert(loopPermMap.size() == loops.size());
unsigned maxLoopDepth = loops.size(); unsigned maxLoopDepth = loops.size();
std::vector<SmallVector<DependenceComponent, 2>> depCompsVec; std::vector<SmallVector<DependenceComponent, 2>> depCompsVec;
getDependenceComponents(loops[0], maxLoopDepth, &depCompsVec); getDependenceComponents(loops[0], maxLoopDepth, &depCompsVec);
return checkLoopInterchangeDependences(depCompsVec, loops, loopPermMap); return checkLoopInterchangeDependences(depCompsVec, loops, loopPermMap);
} }
/// Returns true if `loops` is a perfectly nested loop nest, where loops appear /// Returns true if `loops` is a perfectly nested loop nest, where loops appear
/// in it from outermost to innermost. /// in it from outermost to innermost.
bool LLVM_ATTRIBUTE_UNUSED bool LLVM_ATTRIBUTE_UNUSED
mlir::isPerfectlyNested(ArrayRef<AffineForOp> loops) { mlir::affine::isPerfectlyNested(ArrayRef<AffineForOp> loops) {
assert(!loops.empty() && "no loops provided"); assert(!loops.empty() && "no loops provided");
// We already know that the block can't be empty. // We already know that the block can't be empty.
auto hasTwoElements = [](Block *block) { auto hasTwoElements = [](Block *block) {
auto secondOpIt = std::next(block->begin()); auto secondOpIt = std::next(block->begin());
return secondOpIt != block->end() && &*secondOpIt == &block->back(); return secondOpIt != block->end() && &*secondOpIt == &block->back();
}; };
auto enclosingLoop = loops.front(); auto enclosingLoop = loops.front();
for (auto loop : loops.drop_front()) { for (auto loop : loops.drop_front()) {
auto parentForOp = dyn_cast<AffineForOp>(loop->getParentOp()); auto parentForOp = dyn_cast<AffineForOp>(loop->getParentOp());
// parentForOp's body should be just this loop and the terminator. // parentForOp's body should be just this loop and the terminator.
if (parentForOp != enclosingLoop || !hasTwoElements(parentForOp.getBody())) if (parentForOp != enclosingLoop || !hasTwoElements(parentForOp.getBody()))
return false; return false;
enclosingLoop = loop; enclosingLoop = loop;
} }
return true; return true;
} }
// input[i] should move from position i -> permMap[i]. Returns the position in // input[i] should move from position i -> permMap[i]. Returns the position in
// `input` that becomes the new outermost loop. // `input` that becomes the new outermost loop.
unsigned mlir::permuteLoops(MutableArrayRef<AffineForOp> input, unsigned mlir::affine::permuteLoops(MutableArrayRef<AffineForOp> input,
ArrayRef<unsigned> permMap) { ArrayRef<unsigned> permMap) {
assert(input.size() == permMap.size() && "invalid permutation map size"); assert(input.size() == permMap.size() && "invalid permutation map size");
// Check whether the permutation spec is valid. This is a small vector - we'll // Check whether the permutation spec is valid. This is a small vector - we'll
// just sort and check if it's iota. // just sort and check if it's iota.
SmallVector<unsigned, 4> checkPermMap(permMap.begin(), permMap.end()); SmallVector<unsigned, 4> checkPermMap(permMap.begin(), permMap.end());
llvm::sort(checkPermMap); llvm::sort(checkPermMap);
if (llvm::any_of(llvm::enumerate(checkPermMap), if (llvm::any_of(llvm::enumerate(checkPermMap),
[](const auto &en) { return en.value() != en.index(); })) [](const auto &en) { return en.value() != en.index(); }))
assert(false && "invalid permutation map"); assert(false && "invalid permutation map");
▲ Show 20 Lines • Show All 52 Lines • ▼ Show 20 Linesunsigned mlir::affine::permuteLoops(MutableArrayRef<AffineForOp> input,
} }
return invPermMap[0].second; return invPermMap[0].second;
} }
// Sinks all sequential loops to the innermost levels (while preserving // Sinks all sequential loops to the innermost levels (while preserving
// relative order among them) and moves all parallel loops to the // relative order among them) and moves all parallel loops to the
// outermost (while again preserving relative order among them). // outermost (while again preserving relative order among them).
AffineForOp mlir::sinkSequentialLoops(AffineForOp forOp) { AffineForOp mlir::affine::sinkSequentialLoops(AffineForOp forOp) {
SmallVector<AffineForOp, 4> loops; SmallVector<AffineForOp, 4> loops;
getPerfectlyNestedLoops(loops, forOp); getPerfectlyNestedLoops(loops, forOp);
if (loops.size() < 2) if (loops.size() < 2)
return forOp; return forOp;
// Gather dependence components for dependences between all ops in loop nest // Gather dependence components for dependences between all ops in loop nest
// rooted at 'loops[0]', at loop depths in range [1, maxLoopDepth]. // rooted at 'loops[0]', at loop depths in range [1, maxLoopDepth].
unsigned maxLoopDepth = loops.size(); unsigned maxLoopDepth = loops.size();
▲ Show 20 Lines • Show All 118 Lines • ▼ Show 20 Linesstatic AffineForOp stripmineSink(AffineForOp forOp, SizeType factor,
// exposes the capability. It seems overkill to construct a whole function // exposes the capability. It seems overkill to construct a whole function
// dominance tree at this point. // dominance tree at this point.
auto res = stripmineSink(forOp, factor, ArrayRef<AffineForOp>(target)); auto res = stripmineSink(forOp, factor, ArrayRef<AffineForOp>(target));
assert(res.size() == 1 && "Expected 1 inner forOp"); assert(res.size() == 1 && "Expected 1 inner forOp");
return res[0]; return res[0];
} }
SmallVector<SmallVector<AffineForOp, 8>, 8> SmallVector<SmallVector<AffineForOp, 8>, 8>
mlir::tile(ArrayRef<AffineForOp> forOps, ArrayRef<uint64_t> sizes, mlir::affine::tile(ArrayRef<AffineForOp> forOps, ArrayRef<uint64_t> sizes,
ArrayRef<AffineForOp> targets) { ArrayRef<AffineForOp> targets) {
SmallVector<SmallVector<AffineForOp, 8>, 8> res; SmallVector<SmallVector<AffineForOp, 8>, 8> res;
SmallVector<AffineForOp, 8> currentTargets(targets.begin(), targets.end()); SmallVector<AffineForOp, 8> currentTargets(targets.begin(), targets.end());
for (auto it : llvm::zip(forOps, sizes)) { for (auto it : llvm::zip(forOps, sizes)) {
auto step = stripmineSink(std::get<0>(it), std::get<1>(it), currentTargets); auto step = stripmineSink(std::get<0>(it), std::get<1>(it), currentTargets);
res.push_back(step); res.push_back(step);
currentTargets = step; currentTargets = step;
} }
return res; return res;
} }
SmallVector<AffineForOp, 8> mlir::tile(ArrayRef<AffineForOp> forOps, SmallVector<AffineForOp, 8> mlir::affine::tile(ArrayRef<AffineForOp> forOps,
ArrayRef<uint64_t> sizes, ArrayRef<uint64_t> sizes,
AffineForOp target) { AffineForOp target) {
SmallVector<AffineForOp, 8> res; SmallVector<AffineForOp, 8> res;
for (auto loops : tile(forOps, sizes, ArrayRef<AffineForOp>(target))) { for (auto loops : tile(forOps, sizes, ArrayRef<AffineForOp>(target))) {
assert(loops.size() == 1); assert(loops.size() == 1);
res.push_back(loops[0]); res.push_back(loops[0]);
} }
return res; return res;
} }
LogicalResult mlir::coalesceLoops(MutableArrayRef<AffineForOp> loops) { LogicalResult mlir::affine::coalesceLoops(MutableArrayRef<AffineForOp> loops) {
if (loops.size() < 2) if (loops.size() < 2)
return success(); return success();
AffineForOp innermost = loops.back(); AffineForOp innermost = loops.back();
AffineForOp outermost = loops.front(); AffineForOp outermost = loops.front();
AffineBound ub = outermost.getUpperBound(); AffineBound ub = outermost.getUpperBound();
AffineMap origUbMap = ub.getMap(); AffineMap origUbMap = ub.getMap();
Location loc = outermost.getLoc(); Location loc = outermost.getLoc();
▲ Show 20 Lines • Show All 97 Lines • ▼ Show 20 LinesLogicalResult mlir::affine::coalesceLoops(MutableArrayRef<AffineForOp> loops) {
innermost.getBody()->back().erase(); innermost.getBody()->back().erase();
outermost.getBody()->getOperations().splice( outermost.getBody()->getOperations().splice(
Block::iterator(secondOutermostLoop.getOperation()), Block::iterator(secondOutermostLoop.getOperation()),
innermost.getBody()->getOperations()); innermost.getBody()->getOperations());
secondOutermostLoop.erase(); secondOutermostLoop.erase();
return success(); return success();
} }
void mlir::mapLoopToProcessorIds(scf::ForOp forOp, ArrayRef<Value> processorId, void mlir::affine::mapLoopToProcessorIds(scf::ForOp forOp,
ArrayRef<Value> processorId,
ArrayRef<Value> numProcessors) { ArrayRef<Value> numProcessors) {
assert(processorId.size() == numProcessors.size()); assert(processorId.size() == numProcessors.size());
if (processorId.empty()) if (processorId.empty())
return; return;
OpBuilder b(forOp); OpBuilder b(forOp);
Location loc(forOp.getLoc()); Location loc(forOp.getLoc());
AffineExpr lhs, rhs; AffineExpr lhs, rhs;
bindSymbols(forOp.getContext(), lhs, rhs); bindSymbols(forOp.getContext(), lhs, rhs);
▲ Show 20 Lines • Show All 524 Lines • ▼ Show 20 Lines for (unsigned d = 0; d < rank; d++) {
auto dimSize = memRefType.getDimSize(d); auto dimSize = memRefType.getDimSize(d);
assert(dimSize > 0 && "filtered dynamic shapes above"); assert(dimSize > 0 && "filtered dynamic shapes above");
regionCst->addBound(BoundType::LB, d, 0); regionCst->addBound(BoundType::LB, d, 0);
regionCst->addBound(BoundType::UB, d, dimSize - 1); regionCst->addBound(BoundType::UB, d, dimSize - 1);
} }
return true; return true;
} }
LogicalResult mlir::affineDataCopyGenerate(Block::iterator begin, LogicalResult
Block::iterator end, mlir::affine::affineDataCopyGenerate(Block::iterator begin, Block::iterator end,
const AffineCopyOptions &copyOptions, const AffineCopyOptions &copyOptions,
std::optional<Value> filterMemRef, std::optional<Value> filterMemRef,
DenseSet<Operation *> &copyNests) { DenseSet<Operation *> &copyNests) {
if (begin == end) if (begin == end)
return success(); return success();
assert(begin->getBlock() == std::prev(end)->getBlock() && assert(begin->getBlock() == std::prev(end)->getBlock() &&
"Inconsistent block begin/end args"); "Inconsistent block begin/end args");
assert(end != end->getBlock()->end() && "end can't be the block terminator"); assert(end != end->getBlock()->end() && "end can't be the block terminator");
Block *block = begin->getBlock(); Block *block = begin->getBlock();
▲ Show 20 Lines • Show All 169 Lines • ▼ Show 20 Lines block->getParentOp()->emitWarning(
"capacity"); "capacity");
} }
return success(); return success();
} }
// A convenience version of affineDataCopyGenerate for all ops in the body of // A convenience version of affineDataCopyGenerate for all ops in the body of
// an AffineForOp. // an AffineForOp.
LogicalResult mlir::affineDataCopyGenerate(AffineForOp forOp, LogicalResult mlir::affine::affineDataCopyGenerate(
const AffineCopyOptions &copyOptions, AffineForOp forOp, const AffineCopyOptions &copyOptions,
std::optional<Value> filterMemRef, std::optional<Value> filterMemRef, DenseSet<Operation *> &copyNests) {
DenseSet<Operation *> &copyNests) {
return affineDataCopyGenerate(forOp.getBody()->begin(), return affineDataCopyGenerate(forOp.getBody()->begin(),
std::prev(forOp.getBody()->end()), copyOptions, std::prev(forOp.getBody()->end()), copyOptions,
filterMemRef, copyNests); filterMemRef, copyNests);
} }
LogicalResult mlir::generateCopyForMemRegion( LogicalResult mlir::affine::generateCopyForMemRegion(
const MemRefRegion &memrefRegion, Operation *analyzedOp, const MemRefRegion &memrefRegion, Operation *analyzedOp,
const AffineCopyOptions &copyOptions, CopyGenerateResult &result) { const AffineCopyOptions &copyOptions, CopyGenerateResult &result) {
Block *block = analyzedOp->getBlock(); Block *block = analyzedOp->getBlock();
auto begin = analyzedOp->getIterator(); auto begin = analyzedOp->getIterator();
auto end = std::next(begin); auto end = std::next(begin);
DenseMap<Value, Value> fastBufferMap; DenseMap<Value, Value> fastBufferMap;
DenseSet<Operation *> copyNests; DenseSet<Operation *> copyNests;
Show All 27 Lines for (auto &op : *block) {
if (auto forOp = dyn_cast<AffineForOp>(op)) { if (auto forOp = dyn_cast<AffineForOp>(op)) {
depthToLoops[currLoopDepth].push_back(forOp); depthToLoops[currLoopDepth].push_back(forOp);
gatherLoopsInBlock(forOp.getBody(), currLoopDepth + 1, depthToLoops); gatherLoopsInBlock(forOp.getBody(), currLoopDepth + 1, depthToLoops);
} }
} }
} }
/// Gathers all AffineForOps in 'func.func' grouped by loop depth. /// Gathers all AffineForOps in 'func.func' grouped by loop depth.
void mlir::gatherLoops(func::FuncOp func, void mlir::affine::gatherLoops(
std::vector<SmallVector<AffineForOp, 2>> &depthToLoops) { func::FuncOp func, std::vector<SmallVector<AffineForOp, 2>> &depthToLoops) {
for (auto &block : func) for (auto &block : func)
gatherLoopsInBlock(&block, /*currLoopDepth=*/0, depthToLoops); gatherLoopsInBlock(&block, /*currLoopDepth=*/0, depthToLoops);
// Remove last loop level from output since it's empty. // Remove last loop level from output since it's empty.
if (!depthToLoops.empty()) { if (!depthToLoops.empty()) {
assert(depthToLoops.back().empty() && "Last loop level is not empty?"); assert(depthToLoops.back().empty() && "Last loop level is not empty?");
depthToLoops.pop_back(); depthToLoops.pop_back();
} }
} }
// TODO: if necessary, this can be extended to also compose in any // TODO: if necessary, this can be extended to also compose in any
// affine.applys, fold to constant if all result dimensions of the map are // affine.applys, fold to constant if all result dimensions of the map are
// constant (canonicalizeMapAndOperands below already does this for single // constant (canonicalizeMapAndOperands below already does this for single
// result bound maps), and use simplifyMap to perform algebraic simplification. // result bound maps), and use simplifyMap to perform algebraic simplification.
AffineForOp mlir::createCanonicalizedAffineForOp( AffineForOp mlir::affine::createCanonicalizedAffineForOp(
OpBuilder b, Location loc, ValueRange lbOperands, AffineMap lbMap, OpBuilder b, Location loc, ValueRange lbOperands, AffineMap lbMap,
ValueRange ubOperands, AffineMap ubMap, int64_t step) { ValueRange ubOperands, AffineMap ubMap, int64_t step) {
SmallVector<Value, 4> lowerOperands(lbOperands); SmallVector<Value, 4> lowerOperands(lbOperands);
SmallVector<Value, 4> upperOperands(ubOperands); SmallVector<Value, 4> upperOperands(ubOperands);
fullyComposeAffineMapAndOperands(&lbMap, &lowerOperands); fullyComposeAffineMapAndOperands(&lbMap, &lowerOperands);
canonicalizeMapAndOperands(&lbMap, &lowerOperands); canonicalizeMapAndOperands(&lbMap, &lowerOperands);
lbMap = removeDuplicateExprs(lbMap); lbMap = removeDuplicateExprs(lbMap);
▲ Show 20 Lines • Show All 140 Lines • ▼ Show 20 Lines operandMap.map(loopEn.value().getInductionVar(),
fullTileLoops[loopEn.index()].getInductionVar()); fullTileLoops[loopEn.index()].getInductionVar());
b = OpBuilder::atBlockTerminator(fullTileLoops.back().getBody()); b = OpBuilder::atBlockTerminator(fullTileLoops.back().getBody());
for (auto &op : inputNest.back().getBody()->without_terminator()) for (auto &op : inputNest.back().getBody()->without_terminator())
b.clone(op, operandMap); b.clone(op, operandMap);
return success(); return success();
} }
LogicalResult LogicalResult
mlir::separateFullTiles(MutableArrayRef<AffineForOp> inputNest, mlir::affine::separateFullTiles(MutableArrayRef<AffineForOp> inputNest,
SmallVectorImpl<AffineForOp> *fullTileNest) { SmallVectorImpl<AffineForOp> *fullTileNest) {
if (inputNest.empty()) if (inputNest.empty())
return success(); return success();
auto firstLoop = inputNest[0]; auto firstLoop = inputNest[0];
// Each successive for op has to be nested in the other. // Each successive for op has to be nested in the other.
auto prevLoop = firstLoop; auto prevLoop = firstLoop;
for (auto loop : inputNest.drop_front(1)) { for (auto loop : inputNest.drop_front(1)) {
▲ Show 20 Lines • Show All 43 Lines • Show Last 20 Lines

mlir/lib/Dialect/Affine/Utils/Utils.cpp

Show All 24 Lines
#include "mlir/IR/IRMapping.h" #include "mlir/IR/IRMapping.h"
#include "mlir/IR/IntegerSet.h" #include "mlir/IR/IntegerSet.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include <optional> #include <optional>
#define DEBUG_TYPE "affine-utils" #define DEBUG_TYPE "affine-utils"
using namespace mlir; using namespace mlir;
using namespace affine;
using namespace presburger; using namespace presburger;
namespace { namespace {
/// Visit affine expressions recursively and build the sequence of operations /// Visit affine expressions recursively and build the sequence of operations
/// that correspond to it. Visitation functions return an Value of the /// that correspond to it. Visitation functions return an Value of the
/// expression subtree they visited or `nullptr` on error. /// expression subtree they visited or `nullptr` on error.
class AffineApplyExpander class AffineApplyExpander
: public AffineExprVisitor<AffineApplyExpander, Value> { : public AffineExprVisitor<AffineApplyExpander, Value> {
▲ Show 20 Lines • Show All 163 Lines • ▼ Show 20 Linesprivate:
ValueRange symbolValues; ValueRange symbolValues;
Location loc; Location loc;
}; };
} // namespace } // namespace
/// Create a sequence of operations that implement the `expr` applied to the /// Create a sequence of operations that implement the `expr` applied to the
/// given dimension and symbol values. /// given dimension and symbol values.
mlir::Value mlir::expandAffineExpr(OpBuilder &builder, Location loc, mlir::Value mlir::affine::expandAffineExpr(OpBuilder &builder, Location loc,
AffineExpr expr, ValueRange dimValues, AffineExpr expr,
ValueRange dimValues,
ValueRange symbolValues) { ValueRange symbolValues) {
return AffineApplyExpander(builder, dimValues, symbolValues, loc).visit(expr); return AffineApplyExpander(builder, dimValues, symbolValues, loc).visit(expr);
} }
/// Create a sequence of operations that implement the `affineMap` applied to /// Create a sequence of operations that implement the `affineMap` applied to
/// the given `operands` (as it it were an AffineApplyOp). /// the given `operands` (as it it were an AffineApplyOp).
std::optional<SmallVector<Value, 8>> std::optional<SmallVector<Value, 8>>
mlir::expandAffineMap(OpBuilder &builder, Location loc, AffineMap affineMap, mlir::affine::expandAffineMap(OpBuilder &builder, Location loc,
ValueRange operands) { AffineMap affineMap, ValueRange operands) {
auto numDims = affineMap.getNumDims(); auto numDims = affineMap.getNumDims();
auto expanded = llvm::to_vector<8>( auto expanded = llvm::to_vector<8>(
llvm::map_range(affineMap.getResults(), llvm::map_range(affineMap.getResults(),
[numDims, &builder, loc, operands](AffineExpr expr) { [numDims, &builder, loc, operands](AffineExpr expr) {
return expandAffineExpr(builder, loc, expr, return expandAffineExpr(builder, loc, expr,
operands.take_front(numDims), operands.take_front(numDims),
operands.drop_front(numDims)); operands.drop_front(numDims));
})); }));
▲ Show 20 Lines • Show All 105 Lines • ▼ Show 20 Linesstatic AffineIfOp hoistAffineIfOp(AffineIfOp ifOp, Operation *hoistOverOp) {
elseBlock->getOperations().splice( elseBlock->getOperations().splice(
elseBlock->begin(), hoistOverOpClone->getBlock()->getOperations(), elseBlock->begin(), hoistOverOpClone->getBlock()->getOperations(),
Block::iterator(hoistOverOpClone)); Block::iterator(hoistOverOpClone));
return hoistedIfOp; return hoistedIfOp;
} }
LogicalResult LogicalResult
mlir::affineParallelize(AffineForOp forOp, mlir::affine::affineParallelize(AffineForOp forOp,
ArrayRef<LoopReduction> parallelReductions) { ArrayRef<LoopReduction> parallelReductions) {
// Fail early if there are iter arguments that are not reductions. // Fail early if there are iter arguments that are not reductions.
unsigned numReductions = parallelReductions.size(); unsigned numReductions = parallelReductions.size();
if (numReductions != forOp.getNumIterOperands()) if (numReductions != forOp.getNumIterOperands())
return failure(); return failure();
Location loc = forOp.getLoc(); Location loc = forOp.getLoc();
OpBuilder outsideBuilder(forOp); OpBuilder outsideBuilder(forOp);
AffineMap lowerBoundMap = forOp.getLowerBoundMap(); AffineMap lowerBoundMap = forOp.getLowerBoundMap();
▲ Show 20 Lines • Show All 41 Lines • ▼ Show 20 Linesmlir::affine::affineParallelize(AffineForOp forOp,
yieldOp->setOperands(reducedValues); yieldOp->setOperands(reducedValues);
newPloop.getBody()->eraseArguments(numIVs, numReductions); newPloop.getBody()->eraseArguments(numIVs, numReductions);
forOp.erase(); forOp.erase();
return success(); return success();
} }
// Returns success if any hoisting happened. // Returns success if any hoisting happened.
LogicalResult mlir::hoistAffineIfOp(AffineIfOp ifOp, bool *folded) { LogicalResult mlir::affine::hoistAffineIfOp(AffineIfOp ifOp, bool *folded) {
// Bail out early if the ifOp returns a result. TODO: Consider how to // Bail out early if the ifOp returns a result. TODO: Consider how to
// properly support this case. // properly support this case.
if (ifOp.getNumResults() != 0) if (ifOp.getNumResults() != 0)
return failure(); return failure();
// Apply canonicalization patterns and folding - this is necessary for the // Apply canonicalization patterns and folding - this is necessary for the
// hoisting check to be correct (operands should be composed), and to be more // hoisting check to be correct (operands should be composed), and to be more
// effective (no unused operands). Since the pattern rewriter's folding is // effective (no unused operands). Since the pattern rewriter's folding is
Show All 37 LinesLogicalResult mlir::affine::hoistAffineIfOp(AffineIfOp ifOp, bool *folded) {
(void)applyPatternsAndFoldGreedily( (void)applyPatternsAndFoldGreedily(
hoistedIfOp->getParentWithTrait<OpTrait::IsIsolatedFromAbove>(), hoistedIfOp->getParentWithTrait<OpTrait::IsIsolatedFromAbove>(),
frozenPatterns); frozenPatterns);
return success(); return success();
} }
// Return the min expr after replacing the given dim. // Return the min expr after replacing the given dim.
AffineExpr mlir::substWithMin(AffineExpr e, AffineExpr dim, AffineExpr min, AffineExpr mlir::affine::substWithMin(AffineExpr e, AffineExpr dim,
AffineExpr max, bool positivePath) { AffineExpr min, AffineExpr max,
bool positivePath) {
if (e == dim) if (e == dim)
return positivePath ? min : max; return positivePath ? min : max;
if (auto bin = e.dyn_cast<AffineBinaryOpExpr>()) { if (auto bin = e.dyn_cast<AffineBinaryOpExpr>()) {
AffineExpr lhs = bin.getLHS(); AffineExpr lhs = bin.getLHS();
AffineExpr rhs = bin.getRHS(); AffineExpr rhs = bin.getRHS();
if (bin.getKind() == mlir::AffineExprKind::Add) if (bin.getKind() == mlir::AffineExprKind::Add)
return substWithMin(lhs, dim, min, max, positivePath) + return substWithMin(lhs, dim, min, max, positivePath) +
substWithMin(rhs, dim, min, max, positivePath); substWithMin(rhs, dim, min, max, positivePath);
auto c1 = bin.getLHS().dyn_cast<AffineConstantExpr>(); auto c1 = bin.getLHS().dyn_cast<AffineConstantExpr>();
auto c2 = bin.getRHS().dyn_cast<AffineConstantExpr>(); auto c2 = bin.getRHS().dyn_cast<AffineConstantExpr>();
if (c1 && c1.getValue() < 0) if (c1 && c1.getValue() < 0)
return getAffineBinaryOpExpr( return getAffineBinaryOpExpr(
bin.getKind(), c1, substWithMin(rhs, dim, min, max, !positivePath)); bin.getKind(), c1, substWithMin(rhs, dim, min, max, !positivePath));
if (c2 && c2.getValue() < 0) if (c2 && c2.getValue() < 0)
return getAffineBinaryOpExpr( return getAffineBinaryOpExpr(
bin.getKind(), substWithMin(lhs, dim, min, max, !positivePath), c2); bin.getKind(), substWithMin(lhs, dim, min, max, !positivePath), c2);
return getAffineBinaryOpExpr( return getAffineBinaryOpExpr(
bin.getKind(), substWithMin(lhs, dim, min, max, positivePath), bin.getKind(), substWithMin(lhs, dim, min, max, positivePath),
substWithMin(rhs, dim, min, max, positivePath)); substWithMin(rhs, dim, min, max, positivePath));
} }
return e; return e;
} }
void mlir::normalizeAffineParallel(AffineParallelOp op) { void mlir::affine::normalizeAffineParallel(AffineParallelOp op) {
// Loops with min/max in bounds are not normalized at the moment. // Loops with min/max in bounds are not normalized at the moment.
if (op.hasMinMaxBounds()) if (op.hasMinMaxBounds())
return; return;
AffineMap lbMap = op.getLowerBoundsMap(); AffineMap lbMap = op.getLowerBoundsMap();
SmallVector<int64_t, 8> steps = op.getSteps(); SmallVector<int64_t, 8> steps = op.getSteps();
// No need to do any work if the parallel op is already normalized. // No need to do any work if the parallel op is already normalized.
bool isAlreadyNormalized = bool isAlreadyNormalized =
▲ Show 20 Lines • Show All 47 Lines • ▼ Show 20 Linesvoid mlir::affine::normalizeAffineParallel(AffineParallelOp op) {
op.setSteps(newSteps); op.setSteps(newSteps);
auto newLowerMap = AffineMap::get( auto newLowerMap = AffineMap::get(
/*dimCount=*/0, /*symbolCount=*/0, lbExprs, op.getContext()); /*dimCount=*/0, /*symbolCount=*/0, lbExprs, op.getContext());
op.setLowerBounds({}, newLowerMap); op.setLowerBounds({}, newLowerMap);
auto newUpperMap = AffineMap::get(ranges.getNumDims(), ranges.getNumSymbols(), auto newUpperMap = AffineMap::get(ranges.getNumDims(), ranges.getNumSymbols(),
ubExprs, op.getContext()); ubExprs, op.getContext());
op.setUpperBounds(ranges.getOperands(), newUpperMap); op.setUpperBounds(ranges.getOperands(), newUpperMap);
} }
LogicalResult mlir::normalizeAffineFor(AffineForOp op, bool promoteSingleIter) { LogicalResult mlir::affine::normalizeAffineFor(AffineForOp op,
bool promoteSingleIter) {
if (promoteSingleIter && succeeded(promoteIfSingleIteration(op))) if (promoteSingleIter && succeeded(promoteIfSingleIteration(op)))
return success(); return success();
// Check if the forop is already normalized. // Check if the forop is already normalized.
if (op.hasConstantLowerBound() && (op.getConstantLowerBound() == 0) && if (op.hasConstantLowerBound() && (op.getConstantLowerBound() == 0) &&
(op.getStep() == 1)) (op.getStep() == 1))
return success(); return success();
▲ Show 20 Lines • Show All 140 Lines • ▼ Show 20 Lines if (srcAccess.memref == destAccess.memref &&
return false; return false;
} }
// TODO: Check here if the memrefs alias: there is no side effect if // TODO: Check here if the memrefs alias: there is no side effect if
// `srcAccess.memref` and `destAccess.memref` don't alias. // `srcAccess.memref` and `destAccess.memref` don't alias.
return true; return true;
} }
template <typename EffectType, typename T> template <typename EffectType, typename T>
bool mlir::hasNoInterveningEffect(Operation *start, T memOp) { bool mlir::affine::hasNoInterveningEffect(Operation *start, T memOp) {
auto isLocallyAllocated = [](Value memref) { auto isLocallyAllocated = [](Value memref) {
auto *defOp = memref.getDefiningOp(); auto *defOp = memref.getDefiningOp();
return defOp && hasSingleEffect<MemoryEffects::Allocate>(defOp, memref); return defOp && hasSingleEffect<MemoryEffects::Allocate>(defOp, memref);
}; };
// A boolean representing whether an intervening operation could have impacted // A boolean representing whether an intervening operation could have impacted
// memOp. // memOp.
bool hasSideEffect = false; bool hasSideEffect = false;
▲ Show 20 Lines • Show All 176 Lines • ▼ Show 20 Lines for (auto *user : loadOp.getMemRef().getUsers()) {
// guarantees this for accesses in the same block. The load could be in a // guarantees this for accesses in the same block. The load could be in a
// nested block that is unreachable. // nested block that is unreachable.
if (storeOp->getBlock() != loadOp->getBlock() && if (storeOp->getBlock() != loadOp->getBlock() &&
!mustReachAtInnermost(srcAccess, destAccess)) !mustReachAtInnermost(srcAccess, destAccess))
continue; continue;
// 4. Ensure there is no intermediate operation which could replace the // 4. Ensure there is no intermediate operation which could replace the
// value in memory. // value in memory.
if (!mlir::hasNoInterveningEffect<MemoryEffects::Write>(storeOp, loadOp)) if (!affine::hasNoInterveningEffect<MemoryEffects::Write>(storeOp, loadOp))
continue; continue;
// We now have a candidate for forwarding. // We now have a candidate for forwarding.
assert(lastWriteStoreOp == nullptr && assert(lastWriteStoreOp == nullptr &&
"multiple simultaneous replacement stores"); "multiple simultaneous replacement stores");
lastWriteStoreOp = storeOp; lastWriteStoreOp = storeOp;
} }
Show All 10 Linesstatic LogicalResult forwardStoreToLoad(
loadOp.getValue().replaceAllUsesWith(storeVal); loadOp.getValue().replaceAllUsesWith(storeVal);
// Record the memref for a later sweep to optimize away. // Record the memref for a later sweep to optimize away.
memrefsToErase.insert(loadOp.getMemRef()); memrefsToErase.insert(loadOp.getMemRef());
// Record this to erase later. // Record this to erase later.
loadOpsToErase.push_back(loadOp); loadOpsToErase.push_back(loadOp);
return success(); return success();
} }
template bool mlir::hasNoInterveningEffect<mlir::MemoryEffects::Read, template bool
mlir::AffineReadOpInterface>( mlir::affine::hasNoInterveningEffect<mlir::MemoryEffects::Read,
mlir::Operation *, mlir::AffineReadOpInterface); affine::AffineReadOpInterface>(
mlir::Operation *, affine::AffineReadOpInterface);
// This attempts to find stores which have no impact on the final result. // This attempts to find stores which have no impact on the final result.
// A writing op writeA will be eliminated if there exists an op writeB if // A writing op writeA will be eliminated if there exists an op writeB if
// 1) writeA and writeB have mathematically equivalent affine access functions. // 1) writeA and writeB have mathematically equivalent affine access functions.
// 2) writeB postdominates writeA. // 2) writeB postdominates writeA.
// 3) There is no potential read between writeA and writeB. // 3) There is no potential read between writeA and writeB.
static void findUnusedStore(AffineWriteOpInterface writeA, static void findUnusedStore(AffineWriteOpInterface writeA,
SmallVectorImpl<Operation *> &opsToErase, SmallVectorImpl<Operation *> &opsToErase,
Show All 21 Lines if (srcAccess != destAccess)
continue; continue;
// writeB must postdominate writeA. // writeB must postdominate writeA.
if (!postDominanceInfo.postDominates(writeB, writeA)) if (!postDominanceInfo.postDominates(writeB, writeA))
continue; continue;
// There cannot be an operation which reads from memory between // There cannot be an operation which reads from memory between
// the two writes. // the two writes.
if (!mlir::hasNoInterveningEffect<MemoryEffects::Read>(writeA, writeB)) if (!affine::hasNoInterveningEffect<MemoryEffects::Read>(writeA, writeB))
continue; continue;
opsToErase.push_back(writeA); opsToErase.push_back(writeA);
break; break;
} }
} }
// The load to load forwarding / redundant load elimination is similar to the // The load to load forwarding / redundant load elimination is similar to the
Show All 19 Lines if (srcAccess != destAccess) {
continue; continue;
} }
// 2. The store has to dominate the load op to be candidate. // 2. The store has to dominate the load op to be candidate.
if (!domInfo.dominates(loadB, loadA)) if (!domInfo.dominates(loadB, loadA))
continue; continue;
// 3. There is no write between loadA and loadB. // 3. There is no write between loadA and loadB.
if (!mlir::hasNoInterveningEffect<MemoryEffects::Write>( if (!affine::hasNoInterveningEffect<MemoryEffects::Write>(
loadB.getOperation(), loadA)) loadB.getOperation(), loadA))
continue; continue;
// Check if two values have the same shape. This is needed for affine vector // Check if two values have the same shape. This is needed for affine vector
// loads. // loads.
if (loadB.getValue().getType() != loadA.getValue().getType()) if (loadB.getValue().getType() != loadA.getValue().getType())
continue; continue;
▲ Show 20 Lines • Show All 41 Lines • ▼ Show 20 Lines
// multiple sources to forward from. // multiple sources to forward from.
// //
// TODO: more forwarding can be done when support for // TODO: more forwarding can be done when support for
// loop/conditional live-out SSA values is available. // loop/conditional live-out SSA values is available.
// TODO: do general dead store elimination for memref's. This pass // TODO: do general dead store elimination for memref's. This pass
// currently only eliminates the stores only if no other loads/uses (other // currently only eliminates the stores only if no other loads/uses (other
// than dealloc) remain. // than dealloc) remain.
// //
void mlir::affineScalarReplace(func::FuncOp f, DominanceInfo &domInfo, void mlir::affine::affineScalarReplace(func::FuncOp f, DominanceInfo &domInfo,
PostDominanceInfo &postDomInfo) { PostDominanceInfo &postDomInfo) {
// Load op's whose results were replaced by those forwarded from stores. // Load op's whose results were replaced by those forwarded from stores.
SmallVector<Operation *, 8> opsToErase; SmallVector<Operation *, 8> opsToErase;
// A list of memref's that are potentially dead / could be eliminated. // A list of memref's that are potentially dead / could be eliminated.
SmallPtrSet<Value, 4> memrefsToErase; SmallPtrSet<Value, 4> memrefsToErase;
// Walk all load's and perform store to load forwarding. // Walk all load's and perform store to load forwarding.
f.walk([&](AffineReadOpInterface loadOp) { f.walk([&](AffineReadOpInterface loadOp) {
Show All 36 Lines for (auto memref : memrefsToErase) {
// Erase all stores, the dealloc, and the alloc on the memref. // Erase all stores, the dealloc, and the alloc on the memref.
for (auto *user : llvm::make_early_inc_range(memref.getUsers())) for (auto *user : llvm::make_early_inc_range(memref.getUsers()))
user->erase(); user->erase();
defOp->erase(); defOp->erase();
} }
} }
// Perform the replacement in `op`. // Perform the replacement in `op`.
LogicalResult mlir::replaceAllMemRefUsesWith(Value oldMemRef, Value newMemRef, LogicalResult mlir::affine::replaceAllMemRefUsesWith(
Operation *op, Value oldMemRef, Value newMemRef, Operation *op,
ArrayRef<Value> extraIndices, ArrayRef<Value> extraIndices, AffineMap indexRemap,
AffineMap indexRemap, ArrayRef<Value> extraOperands, ArrayRef<Value> symbolOperands,
ArrayRef<Value> extraOperands,
ArrayRef<Value> symbolOperands,
bool allowNonDereferencingOps) { bool allowNonDereferencingOps) {
unsigned newMemRefRank = newMemRef.getType().cast<MemRefType>().getRank(); unsigned newMemRefRank = newMemRef.getType().cast<MemRefType>().getRank();
(void)newMemRefRank; // unused in opt mode (void)newMemRefRank; // unused in opt mode
unsigned oldMemRefRank = oldMemRef.getType().cast<MemRefType>().getRank(); unsigned oldMemRefRank = oldMemRef.getType().cast<MemRefType>().getRank();
(void)oldMemRefRank; // unused in opt mode (void)oldMemRefRank; // unused in opt mode
if (indexRemap) { if (indexRemap) {
assert(indexRemap.getNumSymbols() == symbolOperands.size() && assert(indexRemap.getNumSymbols() == symbolOperands.size() &&
"symbolic operand count mismatch"); "symbolic operand count mismatch");
assert(indexRemap.getNumInputs() == assert(indexRemap.getNumInputs() ==
▲ Show 20 Lines • Show All 153 Lines • ▼ Show 20 LinesLogicalResult mlir::affine::replaceAllMemRefUsesWith(
// Create the new operation. // Create the new operation.
auto *repOp = builder.create(state); auto *repOp = builder.create(state);
op->replaceAllUsesWith(repOp); op->replaceAllUsesWith(repOp);
op->erase(); op->erase();
return success(); return success();
} }
LogicalResult mlir::replaceAllMemRefUsesWith( LogicalResult mlir::affine::replaceAllMemRefUsesWith(
Value oldMemRef, Value newMemRef, ArrayRef<Value> extraIndices, Value oldMemRef, Value newMemRef, ArrayRef<Value> extraIndices,
AffineMap indexRemap, ArrayRef<Value> extraOperands, AffineMap indexRemap, ArrayRef<Value> extraOperands,
ArrayRef<Value> symbolOperands, Operation *domOpFilter, ArrayRef<Value> symbolOperands, Operation *domOpFilter,
Operation *postDomOpFilter, bool allowNonDereferencingOps, Operation *postDomOpFilter, bool allowNonDereferencingOps,
bool replaceInDeallocOp) { bool replaceInDeallocOp) {
unsigned newMemRefRank = newMemRef.getType().cast<MemRefType>().getRank(); unsigned newMemRefRank = newMemRef.getType().cast<MemRefType>().getRank();
(void)newMemRefRank; // unused in opt mode (void)newMemRefRank; // unused in opt mode
unsigned oldMemRefRank = oldMemRef.getType().cast<MemRefType>().getRank(); unsigned oldMemRefRank = oldMemRef.getType().cast<MemRefType>().getRank();
▲ Show 20 Lines • Show All 99 Lines • ▼ Show 20 Lines
/// This allows applying different transformations on send and compute (for eg. /// This allows applying different transformations on send and compute (for eg.
/// different shifts/delays). /// different shifts/delays).
/// ///
/// Returns nullptr either if none of opInst's operands were the result of an /// Returns nullptr either if none of opInst's operands were the result of an
/// affine.apply and thus there was no affine computation slice to create, or if /// affine.apply and thus there was no affine computation slice to create, or if
/// all the affine.apply op's supplying operands to this opInst did not have any /// all the affine.apply op's supplying operands to this opInst did not have any
/// uses besides this opInst; otherwise returns the list of affine.apply /// uses besides this opInst; otherwise returns the list of affine.apply
/// operations created in output argument `sliceOps`. /// operations created in output argument `sliceOps`.
void mlir::createAffineComputationSlice( void mlir::affine::createAffineComputationSlice(
Operation *opInst, SmallVectorImpl<AffineApplyOp> *sliceOps) { Operation *opInst, SmallVectorImpl<AffineApplyOp> *sliceOps) {
// Collect all operands that are results of affine apply ops. // Collect all operands that are results of affine apply ops.
SmallVector<Value, 4> subOperands; SmallVector<Value, 4> subOperands;
subOperands.reserve(opInst->getNumOperands()); subOperands.reserve(opInst->getNumOperands());
for (auto operand : opInst->getOperands()) for (auto operand : opInst->getOperands())
if (isa_and_nonnull<AffineApplyOp>(operand.getDefiningOp())) if (isa_and_nonnull<AffineApplyOp>(operand.getDefiningOp()))
subOperands.push_back(operand); subOperands.push_back(operand);
▲ Show 20 Lines • Show All 291 Lines • ▼ Show 20 Lines if (newMemRefShape[newDimIdx] < 0) {
b.create<AffineApplyOp>(allocOp->getLoc(), newMap, inAffineApply); b.create<AffineApplyOp>(allocOp->getLoc(), newMap, inAffineApply);
newDynamicSizes.emplace_back(affineApp); newDynamicSizes.emplace_back(affineApp);
} }
newDimIdx++; newDimIdx++;
} }
} }
// TODO: Currently works for static memrefs with a single layout map. // TODO: Currently works for static memrefs with a single layout map.
LogicalResult mlir::normalizeMemRef(memref::AllocOp *allocOp) { LogicalResult mlir::affine::normalizeMemRef(memref::AllocOp *allocOp) {
MemRefType memrefType = allocOp->getType(); MemRefType memrefType = allocOp->getType();
OpBuilder b(*allocOp); OpBuilder b(*allocOp);
// Fetch a new memref type after normalizing the old memref to have an // Fetch a new memref type after normalizing the old memref to have an
// identity map layout. // identity map layout.
MemRefType newMemRefType = MemRefType newMemRefType =
normalizeMemRefType(memrefType, allocOp->getSymbolOperands().size()); normalizeMemRefType(memrefType, allocOp->getSymbolOperands().size());
if (newMemRefType == memrefType) if (newMemRefType == memrefType)
▲ Show 20 Lines • Show All 41 Lines • ▼ Show 20 LinesLogicalResult mlir::affine::normalizeMemRef(memref::AllocOp *allocOp) {
assert(llvm::all_of(oldMemRef.getUsers(), [&](Operation *op) { assert(llvm::all_of(oldMemRef.getUsers(), [&](Operation *op) {
return hasSingleEffect<MemoryEffects::Free>(op, oldMemRef); return hasSingleEffect<MemoryEffects::Free>(op, oldMemRef);
})); }));
oldMemRef.replaceAllUsesWith(newAlloc); oldMemRef.replaceAllUsesWith(newAlloc);
allocOp->erase(); allocOp->erase();
return success(); return success();
} }
MemRefType mlir::normalizeMemRefType(MemRefType memrefType, MemRefType mlir::affine::normalizeMemRefType(MemRefType memrefType,
unsigned numSymbolicOperands) { unsigned numSymbolicOperands) {
unsigned rank = memrefType.getRank(); unsigned rank = memrefType.getRank();
if (rank == 0) if (rank == 0)
return memrefType; return memrefType;
if (memrefType.getLayout().isIdentity()) { if (memrefType.getLayout().isIdentity()) {
// Either no maps is associated with this memref or this memref has // Either no maps is associated with this memref or this memref has
// a trivial (identity) map. // a trivial (identity) map.
return memrefType; return memrefType;
▲ Show 20 Lines • Show All 63 Lines • ▼ Show 20 LinesMemRefType mlir::affine::normalizeMemRefType(MemRefType memrefType,
MemRefType newMemRefType = MemRefType newMemRefType =
MemRefType::Builder(memrefType) MemRefType::Builder(memrefType)
.setShape(newShape) .setShape(newShape)
.setLayout(AffineMapAttr::get( .setLayout(AffineMapAttr::get(
AffineMap::getMultiDimIdentityMap(newRank, context))); AffineMap::getMultiDimIdentityMap(newRank, context)));
return newMemRefType; return newMemRefType;
} }
DivModValue mlir::getDivMod(OpBuilder &b, Location loc, Value lhs, Value rhs) { DivModValue mlir::affine::getDivMod(OpBuilder &b, Location loc, Value lhs,
Value rhs) {
DivModValue result; DivModValue result;
AffineExpr d0, d1; AffineExpr d0, d1;
bindDims(b.getContext(), d0, d1); bindDims(b.getContext(), d0, d1);
result.quotient = result.quotient =
makeComposedAffineApply(b, loc, d0.floorDiv(d1), {lhs, rhs}); affine::makeComposedAffineApply(b, loc, d0.floorDiv(d1), {lhs, rhs});
result.remainder = makeComposedAffineApply(b, loc, d0 % d1, {lhs, rhs}); result.remainder =
affine::makeComposedAffineApply(b, loc, d0 % d1, {lhs, rhs});
return result; return result;
} }
/// Create IR that computes the product of all elements in the set. /// Create IR that computes the product of all elements in the set.
static FailureOr<OpFoldResult> getIndexProduct(OpBuilder &b, Location loc, static FailureOr<OpFoldResult> getIndexProduct(OpBuilder &b, Location loc,
ArrayRef<Value> set) { ArrayRef<Value> set) {
if (set.empty()) if (set.empty())
return failure(); return failure();
OpFoldResult result = set[0]; OpFoldResult result = set[0];
AffineExpr s0, s1; AffineExpr s0, s1;
bindSymbols(b.getContext(), s0, s1); bindSymbols(b.getContext(), s0, s1);
for (unsigned i = 1, e = set.size(); i < e; i++) for (unsigned i = 1, e = set.size(); i < e; i++)
result = makeComposedFoldedAffineApply(b, loc, s0 * s1, {result, set[i]}); result = makeComposedFoldedAffineApply(b, loc, s0 * s1, {result, set[i]});
return result; return result;
} }
FailureOr<SmallVector<Value>> mlir::delinearizeIndex(OpBuilder &b, Location loc, FailureOr<SmallVector<Value>>
Value linearIndex, mlir::affine::delinearizeIndex(OpBuilder &b, Location loc, Value linearIndex,
ArrayRef<Value> basis) { ArrayRef<Value> basis) {
unsigned numDims = basis.size(); unsigned numDims = basis.size();
SmallVector<Value> divisors; SmallVector<Value> divisors;
for (unsigned i = 1; i < numDims; i++) { for (unsigned i = 1; i < numDims; i++) {
ArrayRef<Value> slice = basis.drop_front(i); ArrayRef<Value> slice = basis.drop_front(i);
FailureOr<OpFoldResult> prod = getIndexProduct(b, loc, slice); FailureOr<OpFoldResult> prod = getIndexProduct(b, loc, slice);
if (failed(prod)) if (failed(prod))
return failure(); return failure();
Show All 14 Lines

mlir/lib/Dialect/Affine/Utils/ViewLikeInterfaceUtils.cpp

//===- ViewLikeInterfaceUtils.cpp -----------------------------------------===// //===- ViewLikeInterfaceUtils.cpp -----------------------------------------===//
// //
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information. // See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// //
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
#include "mlir/Dialect/Affine/ViewLikeInterfaceUtils.h" #include "mlir/Dialect/Affine/ViewLikeInterfaceUtils.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/Utils/Utils.h" #include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/IR/PatternMatch.h" #include "mlir/IR/PatternMatch.h"
using namespace mlir; using namespace mlir;
using namespace affine;
LogicalResult mlir::mergeOffsetsSizesAndStrides( LogicalResult mlir::affine::mergeOffsetsSizesAndStrides(
OpBuilder &builder, Location loc, ArrayRef<OpFoldResult> producerOffsets, OpBuilder &builder, Location loc, ArrayRef<OpFoldResult> producerOffsets,
ArrayRef<OpFoldResult> producerSizes, ArrayRef<OpFoldResult> producerSizes,
ArrayRef<OpFoldResult> producerStrides, ArrayRef<OpFoldResult> producerStrides,
const llvm::SmallBitVector &droppedProducerDims, const llvm::SmallBitVector &droppedProducerDims,
ArrayRef<OpFoldResult> consumerOffsets, ArrayRef<OpFoldResult> consumerOffsets,
ArrayRef<OpFoldResult> consumerSizes, ArrayRef<OpFoldResult> consumerSizes,
ArrayRef<OpFoldResult> consumerStrides, ArrayRef<OpFoldResult> consumerStrides,
SmallVector<OpFoldResult> &combinedOffsets, SmallVector<OpFoldResult> &combinedOffsets,
Show All 28 Lines combinedStrides[i] = makeComposedFoldedAffineApply(
builder, loc, s0 * s1, builder, loc, s0 * s1,
{consumerStrides[consumerPos], producerStrides[i]}); {consumerStrides[consumerPos], producerStrides[i]});
consumerPos++; consumerPos++;
} }
return success(); return success();
} }
LogicalResult mlir::mergeOffsetsSizesAndStrides( LogicalResult mlir::affine::mergeOffsetsSizesAndStrides(
OpBuilder &builder, Location loc, OffsetSizeAndStrideOpInterface producer, OpBuilder &builder, Location loc, OffsetSizeAndStrideOpInterface producer,
OffsetSizeAndStrideOpInterface consumer, OffsetSizeAndStrideOpInterface consumer,
const llvm::SmallBitVector &droppedProducerDims, const llvm::SmallBitVector &droppedProducerDims,
SmallVector<OpFoldResult> &combinedOffsets, SmallVector<OpFoldResult> &combinedOffsets,
SmallVector<OpFoldResult> &combinedSizes, SmallVector<OpFoldResult> &combinedSizes,
SmallVector<OpFoldResult> &combinedStrides) { SmallVector<OpFoldResult> &combinedStrides) {
SmallVector<OpFoldResult> consumerOffsets = consumer.getMixedOffsets(); SmallVector<OpFoldResult> consumerOffsets = consumer.getMixedOffsets();
SmallVector<OpFoldResult> consumerSizes = consumer.getMixedSizes(); SmallVector<OpFoldResult> consumerSizes = consumer.getMixedSizes();
SmallVector<OpFoldResult> consumerStrides = consumer.getMixedStrides(); SmallVector<OpFoldResult> consumerStrides = consumer.getMixedStrides();
SmallVector<OpFoldResult> producerOffsets = producer.getMixedOffsets(); SmallVector<OpFoldResult> producerOffsets = producer.getMixedOffsets();
SmallVector<OpFoldResult> producerSizes = producer.getMixedSizes(); SmallVector<OpFoldResult> producerSizes = producer.getMixedSizes();
SmallVector<OpFoldResult> producerStrides = producer.getMixedStrides(); SmallVector<OpFoldResult> producerStrides = producer.getMixedStrides();
return mergeOffsetsSizesAndStrides( return mergeOffsetsSizesAndStrides(
builder, loc, producerOffsets, producerSizes, producerStrides, builder, loc, producerOffsets, producerSizes, producerStrides,
droppedProducerDims, consumerOffsets, consumerSizes, consumerStrides, droppedProducerDims, consumerOffsets, consumerSizes, consumerStrides,
combinedOffsets, combinedSizes, combinedStrides); combinedOffsets, combinedSizes, combinedStrides);
} }
void mlir::resolveIndicesIntoOpWithOffsetsAndStrides( void mlir::affine::resolveIndicesIntoOpWithOffsetsAndStrides(
RewriterBase &rewriter, Location loc, RewriterBase &rewriter, Location loc,
ArrayRef<OpFoldResult> mixedSourceOffsets, ArrayRef<OpFoldResult> mixedSourceOffsets,
ArrayRef<OpFoldResult> mixedSourceStrides, ArrayRef<OpFoldResult> mixedSourceStrides,
const llvm::SmallBitVector &rankReducedDims, const llvm::SmallBitVector &rankReducedDims,
ArrayRef<OpFoldResult> consumerIndices, ArrayRef<OpFoldResult> consumerIndices,
SmallVectorImpl<Value> &resolvedIndices) { SmallVectorImpl<Value> &resolvedIndices) {
OpFoldResult zero = rewriter.getIndexAttr(0); OpFoldResult zero = rewriter.getIndexAttr(0);
Show All 15 Lines for (auto [offset, index, stride] :
OpFoldResult ofr = makeComposedFoldedAffineApply( OpFoldResult ofr = makeComposedFoldedAffineApply(
rewriter, loc, AffineMap::get(0, 3, off + idx * str), rewriter, loc, AffineMap::get(0, 3, off + idx * str),
{offset, index, stride}); {offset, index, stride});
resolvedIndices.push_back( resolvedIndices.push_back(
getValueOrCreateConstantIndexOp(rewriter, loc, ofr)); getValueOrCreateConstantIndexOp(rewriter, loc, ofr));
} }
} }
void mlir::resolveSizesIntoOpWithSizes( void mlir::affine::resolveSizesIntoOpWithSizes(
ArrayRef<OpFoldResult> sourceSizes, ArrayRef<OpFoldResult> destSizes, ArrayRef<OpFoldResult> sourceSizes, ArrayRef<OpFoldResult> destSizes,
const llvm::SmallBitVector &rankReducedSourceDims, const llvm::SmallBitVector &rankReducedSourceDims,
SmallVectorImpl<OpFoldResult> &resolvedSizes) { SmallVectorImpl<OpFoldResult> &resolvedSizes) {
int64_t dim = 0; int64_t dim = 0;
int64_t srcRank = sourceSizes.size(); int64_t srcRank = sourceSizes.size();
for (int64_t srcDim = 0; srcDim < srcRank; ++srcDim) { for (int64_t srcDim = 0; srcDim < srcRank; ++srcDim) {
if (rankReducedSourceDims[srcDim]) { if (rankReducedSourceDims[srcDim]) {
resolvedSizes.push_back(sourceSizes[srcDim]); resolvedSizes.push_back(sourceSizes[srcDim]);
continue; continue;
} }
resolvedSizes.push_back(destSizes[dim++]); resolvedSizes.push_back(destSizes[dim++]);
} }
} }

mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp

Show First 20 Lines • Show All 55 Lines • ▼ Show 20 Linesstatic Value buildLinearThreadId(RewriterBase &rewriter, Location loc,
bindSymbols(rewriter.getContext(), BDX, BDY); bindSymbols(rewriter.getContext(), BDX, BDY);
IndexType indexType = rewriter.getIndexType(); IndexType indexType = rewriter.getIndexType();
SmallVector<OpFoldResult> threadsAndWorkGroups{ SmallVector<OpFoldResult> threadsAndWorkGroups{
rewriter.create<ThreadIdOp>(loc, indexType, Dimension::x).getResult(), rewriter.create<ThreadIdOp>(loc, indexType, Dimension::x).getResult(),
rewriter.create<ThreadIdOp>(loc, indexType, Dimension::y).getResult(), rewriter.create<ThreadIdOp>(loc, indexType, Dimension::y).getResult(),
rewriter.create<ThreadIdOp>(loc, indexType, Dimension::z).getResult()}; rewriter.create<ThreadIdOp>(loc, indexType, Dimension::z).getResult()};
threadsAndWorkGroups.push_back(blockDimsOfr[0]); threadsAndWorkGroups.push_back(blockDimsOfr[0]);
threadsAndWorkGroups.push_back(blockDimsOfr[1]); threadsAndWorkGroups.push_back(blockDimsOfr[1]);
OpFoldResult ofr = makeComposedFoldedAffineApply( OpFoldResult ofr = affine::makeComposedFoldedAffineApply(
rewriter, loc, tx + ty * BDX + tz * BDX * BDY, threadsAndWorkGroups); rewriter, loc, tx + ty * BDX + tz * BDX * BDY, threadsAndWorkGroups);
return getValueOrCreateConstantIndexOp(rewriter, loc, ofr); return getValueOrCreateConstantIndexOp(rewriter, loc, ofr);
} }
/// Builder for gpu::BlockIdOps used in mapping scf.forall to blocks. /// Builder for gpu::BlockIdOps used in mapping scf.forall to blocks.
/// The `idBuilder` method returns 3-D values used for indexing rewrites as well /// The `idBuilder` method returns 3-D values used for indexing rewrites as well
/// as 3-D sizes for predicate generation. /// as 3-D sizes for predicate generation.
struct GpuBlockIdBuilder : public GpuIdBuilder { struct GpuBlockIdBuilder : public GpuIdBuilder {
▲ Show 20 Lines • Show All 59 Lines • ▼ Show 20 Lines mappingAttributes = {GPUWarpMappingAttr::get(ctx, Warps::DimX),
GPUWarpMappingAttr::get(ctx, Warps::DimY), GPUWarpMappingAttr::get(ctx, Warps::DimY),
GPUWarpMappingAttr::get(ctx, Warps::DimZ)}; GPUWarpMappingAttr::get(ctx, Warps::DimZ)};
idBuilder = [this](RewriterBase &rewriter, Location loc, idBuilder = [this](RewriterBase &rewriter, Location loc,
ArrayRef<int64_t> forallMappingSizes) { ArrayRef<int64_t> forallMappingSizes) {
// Build the linear warp id and decompose it in the basis of // Build the linear warp id and decompose it in the basis of
// `forallMappingSizes`. // `forallMappingSizes`.
Value linearId = buildLinearThreadId(rewriter, loc, this->blockDimsOfr); Value linearId = buildLinearThreadId(rewriter, loc, this->blockDimsOfr);
AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext()); AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext());
OpFoldResult warpIdOfr = makeComposedFoldedAffineApply( OpFoldResult warpIdOfr = affine::makeComposedFoldedAffineApply(
rewriter, loc, d0.floorDiv(kWarpSize), {linearId}); rewriter, loc, d0.floorDiv(kWarpSize), {linearId});
Value warpId = getValueOrCreateConstantIndexOp(rewriter, loc, warpIdOfr); Value warpId = getValueOrCreateConstantIndexOp(rewriter, loc, warpIdOfr);
// Sizes in [x, y, z] -> [z, y x] order to properly compute strides in // Sizes in [x, y, z] -> [z, y x] order to properly compute strides in
// "row-major" order. // "row-major" order.
SmallVector<int64_t> reverseBasisSizes( SmallVector<int64_t> reverseBasisSizes(
llvm::reverse(this->availableMappingSizes)); llvm::reverse(this->availableMappingSizes));
SmallVector<int64_t> strides = computeStrides(reverseBasisSizes); SmallVector<int64_t> strides = computeStrides(reverseBasisSizes);
SmallVector<AffineExpr> delinearizingExprs = delinearize(d0, strides); SmallVector<AffineExpr> delinearizingExprs = delinearize(d0, strides);
SmallVector<Value> ids; SmallVector<Value> ids;
// Reverse back to be in [x, y, z] order. // Reverse back to be in [x, y, z] order.
for (AffineExpr e : llvm::reverse(delinearizingExprs)) for (AffineExpr e : llvm::reverse(delinearizingExprs))
ids.push_back(makeComposedAffineApply(rewriter, loc, e, warpId)); ids.push_back(
affine::makeComposedAffineApply(rewriter, loc, e, warpId));
// clang-format off // clang-format off
LDBG("----linearId: " << linearId); LDBG("----linearId: " << linearId);
LDBG("----warpId: " << warpId); LDBG("----warpId: " << warpId);
LLVM_DEBUG(llvm::interleaveComma(reverseBasisSizes, LLVM_DEBUG(llvm::interleaveComma(reverseBasisSizes,
DBGS() << "--delinearization basis: "); DBGS() << "--delinearization basis: ");
llvm::dbgs() << "\n"; llvm::dbgs() << "\n";
llvm::interleaveComma(strides, llvm::interleaveComma(strides,
Show All 38 Lines idBuilder = [this](RewriterBase &rewriter, Location loc,
SmallVector<int64_t> reverseBasisSizes(llvm::reverse(forallMappingSizes)); SmallVector<int64_t> reverseBasisSizes(llvm::reverse(forallMappingSizes));
SmallVector<int64_t> strides = computeStrides(reverseBasisSizes); SmallVector<int64_t> strides = computeStrides(reverseBasisSizes);
AffineExpr d0; AffineExpr d0;
bindDims(rewriter.getContext(), d0); bindDims(rewriter.getContext(), d0);
SmallVector<AffineExpr> delinearizingExprs = delinearize(d0, strides); SmallVector<AffineExpr> delinearizingExprs = delinearize(d0, strides);
SmallVector<Value> ids; SmallVector<Value> ids;
// Reverse back to be in [x, y, z] order. // Reverse back to be in [x, y, z] order.
for (AffineExpr e : llvm::reverse(delinearizingExprs)) for (AffineExpr e : llvm::reverse(delinearizingExprs))
ids.push_back(makeComposedAffineApply(rewriter, loc, e, linearId)); ids.push_back(
affine::makeComposedAffineApply(rewriter, loc, e, linearId));
// clang-format off // clang-format off
LLVM_DEBUG(llvm::interleaveComma(reverseBasisSizes, LLVM_DEBUG(llvm::interleaveComma(reverseBasisSizes,
DBGS() << "--delinearization basis: "); DBGS() << "--delinearization basis: ");
llvm::dbgs() << "\n"; llvm::dbgs() << "\n";
llvm::interleaveComma(strides, llvm::interleaveComma(strides,
DBGS() << "--delinearization strides: "); DBGS() << "--delinearization strides: ");
llvm::dbgs() << "\n"; llvm::dbgs() << "\n";
▲ Show 20 Lines • Show All 707 Lines • Show Last 20 Lines

mlir/lib/Dialect/GPU/Transforms/MemoryPromotion.cpp

Show First 20 Lines • Show All 75 Lines • ▼ Show 20 Lines mlir::scf::buildLoopNest(
}); });
// Map the innermost loops to threads in reverse order. // Map the innermost loops to threads in reverse order.
for (const auto &en : for (const auto &en :
llvm::enumerate(llvm::reverse(llvm::ArrayRef(ivs).take_back( llvm::enumerate(llvm::reverse(llvm::ArrayRef(ivs).take_back(
GPUDialect::getNumWorkgroupDimensions())))) { GPUDialect::getNumWorkgroupDimensions())))) {
Value v = en.value(); Value v = en.value();
auto loop = cast<scf::ForOp>(v.getParentRegion()->getParentOp()); auto loop = cast<scf::ForOp>(v.getParentRegion()->getParentOp());
mapLoopToProcessorIds(loop, {threadIds[en.index()]}, affine::mapLoopToProcessorIds(loop, {threadIds[en.index()]},
{blockDims[en.index()]}); {blockDims[en.index()]});
} }
} }
/// Emits the loop nests performing the copy to the designated location in the /// Emits the loop nests performing the copy to the designated location in the
/// beginning of the region, and from the designated location immediately before /// beginning of the region, and from the designated location immediately before
/// the terminator of the first block of the region. The region is expected to /// the terminator of the first block of the region. The region is expected to
/// have one block. This boils down to the following structure /// have one block. This boils down to the following structure
/// ///
▲ Show 20 Lines • Show All 68 Lines • Show Last 20 Lines

mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp

Show First 20 Lines • Show All 630 Lines • ▼ Show 20 LinesLinalgOp::reifyResultShapes(OpBuilder &b,
// Check that the result dim map does not contain the positions corresponding // Check that the result dim map does not contain the positions corresponding
// to the outputs. // to the outputs.
llvm::SmallBitVector outputDims(resultShapesFromInputShapesMap.getNumDims()); llvm::SmallBitVector outputDims(resultShapesFromInputShapesMap.getNumDims());
outputDims.set(resultShapesSubMapPos.first, resultShapesSubMapPos.second); outputDims.set(resultShapesSubMapPos.first, resultShapesSubMapPos.second);
HasAffineDimExprVisitor checkDimExpr(std::move(outputDims)); HasAffineDimExprVisitor checkDimExpr(std::move(outputDims));
Location loc = getOperation()->getLoc(); Location loc = getOperation()->getLoc();
IRRewriter rewriter(b); IRRewriter rewriter(b);
SmallVector<OpFoldResult> allResultDimValues = SmallVector<OpFoldResult> allResultDimValues =
makeComposedFoldedMultiResultAffineApply( affine::makeComposedFoldedMultiResultAffineApply(
rewriter, loc, resultShapesFromInputShapesMap, rewriter, loc, resultShapesFromInputShapesMap,
createFlatListOfOperandDims(b, loc)); createFlatListOfOperandDims(b, loc));
int64_t pos = 0; int64_t pos = 0;
ArrayRef<AffineExpr> shapeExprs = resultShapesFromInputShapesMap.getResults(); ArrayRef<AffineExpr> shapeExprs = resultShapesFromInputShapesMap.getResults();
for (OpOperand *opOperand : getDpsInitOperands()) { for (OpOperand *opOperand : getDpsInitOperands()) {
SmallVector<OpFoldResult> shapes; SmallVector<OpFoldResult> shapes;
for (int64_t dim : llvm::seq<int64_t>(0, getRank(opOperand))) { for (int64_t dim : llvm::seq<int64_t>(0, getRank(opOperand))) {
auto shapedType = opOperand->get().getType().cast<ShapedType>(); auto shapedType = opOperand->get().getType().cast<ShapedType>();
▲ Show 20 Lines • Show All 171 Lines • Show Last 20 Lines

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

Show First 20 Lines • Show All 574 Lines • ▼ Show 20 Lines LogicalResult matchAndRewrite(tensor::InsertSliceOp insertOp,
AffineExpr sym0, sym1; AffineExpr sym0, sym1;
bindSymbols(context, sym0, sym1); bindSymbols(context, sym0, sym1);
auto addMap = AffineMap::get(0, 2, {sym0 + sym1}, context); auto addMap = AffineMap::get(0, 2, {sym0 + sym1}, context);
// Calculate the new offsets for the insert. It should be the old offsets // Calculate the new offsets for the insert. It should be the old offsets
// plus low padding sizes. // plus low padding sizes.
SmallVector<OpFoldResult, 4> newOffsets; SmallVector<OpFoldResult, 4> newOffsets;
for (const auto &p : llvm::zip(lowPads, oldOffsets)) { for (const auto &p : llvm::zip(lowPads, oldOffsets)) {
newOffsets.push_back(makeComposedFoldedAffineApply( newOffsets.push_back(affine::makeComposedFoldedAffineApply(
rewriter, loc, addMap, {std::get<0>(p), std::get<1>(p)})); rewriter, loc, addMap, {std::get<0>(p), std::get<1>(p)}));
} }
SmallVector<OpFoldResult, 4> newSizes; SmallVector<OpFoldResult, 4> newSizes;
for (int i = 0, e = srcPadOp.getSourceType().getRank(); i < e; ++i) { for (int i = 0, e = srcPadOp.getSourceType().getRank(); i < e; ++i) {
newSizes.push_back( newSizes.push_back(
rewriter.create<tensor::DimOp>(loc, srcPadOp.getSource(), i) rewriter.create<tensor::DimOp>(loc, srcPadOp.getSource(), i)
.getResult()); .getResult());
▲ Show 20 Lines • Show All 1,522 Lines • Show Last 20 Lines

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

Show First 20 Lines • Show All 1,163 Lines • ▼ Show 20 Lines FailureOr<MultiSizeSpecification> spec = computeMultiTileSizes(
builder, target, getDimension(), targetSize, divisor); builder, target, getDimension(), targetSize, divisor);
if (failed(spec)) { if (failed(spec)) {
return emitSilenceableError() << "could not generate tile size computation"; return emitSilenceableError() << "could not generate tile size computation";
} }
AffineExpr s0 = builder.getAffineSymbolExpr(0); AffineExpr s0 = builder.getAffineSymbolExpr(0);
AffineExpr s1 = builder.getAffineSymbolExpr(1); AffineExpr s1 = builder.getAffineSymbolExpr(1);
Operation *splitPoint = Operation *splitPoint =
makeComposedAffineApply(builder, target.getLoc(), s0 * s1, affine::makeComposedAffineApply(builder, target.getLoc(), s0 * s1,
{spec->lowTileSize, spec->lowTripCount}); {spec->lowTileSize, spec->lowTripCount});
Operation *lowTileSize = spec->lowTileSize.getDefiningOp(); Operation *lowTileSize = spec->lowTileSize.getDefiningOp();
Operation *highTileSize = spec->highTileSize.getDefiningOp(); Operation *highTileSize = spec->highTileSize.getDefiningOp();
assert(lowTileSize && highTileSize && splitPoint && assert(lowTileSize && highTileSize && splitPoint &&
"tile sizes are not produced by operations"); "tile sizes are not produced by operations");
results.reserve(results.size() + 3); results.reserve(results.size() + 3);
results.push_back(lowTileSize); results.push_back(lowTileSize);
results.push_back(highTileSize); results.push_back(highTileSize);
results.push_back(splitPoint); results.push_back(splitPoint);
▲ Show 20 Lines • Show All 233 Lines • ▼ Show 20 LinespackMatmulGreedily(RewriterBase &rewriter, LinalgOp linalgOp,
for (int64_t i = 0, e = numPackedDims; i < e; ++i) { for (int64_t i = 0, e = numPackedDims; i < e; ++i) {
if (paddedSizesNextMultipleOf[i] == 0) { if (paddedSizesNextMultipleOf[i] == 0) {
adjustedPackedSizes.push_back(packedSizes[i]); adjustedPackedSizes.push_back(packedSizes[i]);
continue; continue;
} }
AffineExpr d0, s0; AffineExpr d0, s0;
bindDims(rewriter.getContext(), d0); bindDims(rewriter.getContext(), d0);
bindSymbols(rewriter.getContext(), s0); bindSymbols(rewriter.getContext(), s0);
adjustedPackedSizes.push_back(makeComposedFoldedAffineApply( adjustedPackedSizes.push_back(affine::makeComposedFoldedAffineApply(
rewriter, genericOp->getLoc(), d0.ceilDiv(s0) * s0, rewriter, genericOp->getLoc(), d0.ceilDiv(s0) * s0,
{loopRanges[adjustedPackedSizes.size()].size, {loopRanges[adjustedPackedSizes.size()].size,
rewriter.getIndexAttr(paddedSizesNextMultipleOf[i])})); rewriter.getIndexAttr(paddedSizesNextMultipleOf[i])}));
} }
LLVM_DEBUG(llvm::interleaveComma(adjustedPackedSizes, LLVM_DEBUG(llvm::interleaveComma(adjustedPackedSizes,
DBGS() << "adjustedPackedSizes: "); DBGS() << "adjustedPackedSizes: ");
DBGSNL();); DBGSNL(););
▲ Show 20 Lines • Show All 546 Lines • ▼ Show 20 Lines tilingOptions.setTileSizeComputationFunction([&](OpBuilder &b, Operation *) {
SmallVector<OpFoldResult> allShapeSizes = SmallVector<OpFoldResult> allShapeSizes =
target.createFlatListOfOperandDims(b, loc); target.createFlatListOfOperandDims(b, loc);
AffineMap map = target.getShapesToLoopsMap(); AffineMap map = target.getShapesToLoopsMap();
if (!map) if (!map)
return tileSizes; return tileSizes;
TrackingListener listener(state, *this); TrackingListener listener(state, *this);
IRRewriter rewriter(getContext(), &listener); IRRewriter rewriter(getContext(), &listener);
SmallVector<OpFoldResult> shapeSizes = SmallVector<OpFoldResult> shapeSizes =
makeComposedFoldedMultiResultAffineApply(rewriter, loc, map, affine::makeComposedFoldedMultiResultAffineApply(rewriter, loc, map,
allShapeSizes); allShapeSizes);
// If the shape size is dynamic, tile by 1. // If the shape size is dynamic, tile by 1.
// Otherwise, do not tile (i.e. tile size 0). // Otherwise, do not tile (i.e. tile size 0).
for (OpFoldResult shapeSize : shapeSizes) { for (OpFoldResult shapeSize : shapeSizes) {
tileSizes.push_back(getConstantIntValue(shapeSize) tileSizes.push_back(getConstantIntValue(shapeSize)
? b.create<arith::ConstantIndexOp>(loc, 0) ? b.create<arith::ConstantIndexOp>(loc, 0)
: b.create<arith::ConstantIndexOp>(loc, 1)); : b.create<arith::ConstantIndexOp>(loc, 1));
} }
return tileSizes; return tileSizes;
▲ Show 20 Lines • Show All 1,350 Lines • ▼ Show 20 Linesclass LinalgTransformDialectExtension
: public transform::TransformDialectExtension< : public transform::TransformDialectExtension<
LinalgTransformDialectExtension> { LinalgTransformDialectExtension> {
public: public:
using Base::Base; using Base::Base;
void init() { void init() {
declareDependentDialect<pdl::PDLDialect>(); declareDependentDialect<pdl::PDLDialect>();
declareDependentDialect<LinalgDialect>(); declareDependentDialect<LinalgDialect>();
declareGeneratedDialect<AffineDialect>(); declareGeneratedDialect<affine::AffineDialect>();
declareGeneratedDialect<arith::ArithDialect>(); declareGeneratedDialect<arith::ArithDialect>();
declareGeneratedDialect<scf::SCFDialect>(); declareGeneratedDialect<scf::SCFDialect>();
declareGeneratedDialect<vector::VectorDialect>(); declareGeneratedDialect<vector::VectorDialect>();
declareGeneratedDialect<gpu::GPUDialect>(); declareGeneratedDialect<gpu::GPUDialect>();
registerTransformOps< registerTransformOps<
#define GET_OP_LIST #define GET_OP_LIST
#include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp.inc" #include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp.inc"
Show All 18 Lines

mlir/lib/Dialect/Linalg/Transforms/BubbleUpExtractSlice.cpp

Show First 20 Lines • Show All 90 Lines • ▼ Show 20 Lines LogicalResult matchAndRewrite(tensor::ExtractSliceOp sliceOp,
SmallVector<OpFoldResult> allShapeSizes = SmallVector<OpFoldResult> allShapeSizes =
linalgOp.createFlatListOfOperandDims(rewriter, linalgLoc); linalgOp.createFlatListOfOperandDims(rewriter, linalgLoc);
AffineMap shapeSizesToLoopsMap = linalgOp.getShapesToLoopsMap(); AffineMap shapeSizesToLoopsMap = linalgOp.getShapesToLoopsMap();
if (!shapeSizesToLoopsMap) { if (!shapeSizesToLoopsMap) {
return rewriter.notifyMatchFailure( return rewriter.notifyMatchFailure(
linalgOp, "failed to get loops map from shape sizes"); linalgOp, "failed to get loops map from shape sizes");
} }
SmallVector<OpFoldResult> sizeBounds = SmallVector<OpFoldResult> sizeBounds =
makeComposedFoldedMultiResultAffineApply( affine::makeComposedFoldedMultiResultAffineApply(
rewriter, linalgLoc, shapeSizesToLoopsMap, allShapeSizes); rewriter, linalgLoc, shapeSizesToLoopsMap, allShapeSizes);
// The offsets and sizes from the slice operation only give you the tile // The offsets and sizes from the slice operation only give you the tile
// size of the output. Use that compute the tile sizes and offsets of the // size of the output. Use that compute the tile sizes and offsets of the
// loops. For loops not used to access the output, set the tile sizes to // loops. For loops not used to access the output, set the tile sizes to
// loop bounds and set the offset to 0. // loop bounds and set the offset to 0.
SmallVector<OpFoldResult> tileOffsets(sizeBounds.size(), SmallVector<OpFoldResult> tileOffsets(sizeBounds.size(),
rewriter.getIndexAttr(0)); rewriter.getIndexAttr(0));
Show All 31 Lines

mlir/lib/Dialect/Linalg/Transforms/ConvertConv2DToImg2Col.cpp

Show First 20 Lines • Show All 49 Lines • ▼ Show 20 Lines
// Delinearizes the given composite `index` by the basis specified in `factors`. // Delinearizes the given composite `index` by the basis specified in `factors`.
static SmallVector<Value> unrollIndex(OpBuilder &b, Location loc, Value index, static SmallVector<Value> unrollIndex(OpBuilder &b, Location loc, Value index,
ArrayRef<int64_t> factors) { ArrayRef<int64_t> factors) {
assert(!factors.empty() && "empty factor list"); assert(!factors.empty() && "empty factor list");
SmallVector<Value> basis; SmallVector<Value> basis;
for (int64_t f : factors) for (int64_t f : factors)
basis.push_back(b.create<arith::ConstantOp>(loc, b.getIndexAttr(f))); basis.push_back(b.create<arith::ConstantOp>(loc, b.getIndexAttr(f)));
FailureOr<SmallVector<Value>> multiIndex = FailureOr<SmallVector<Value>> multiIndex =
delinearizeIndex(b, loc, index, basis); affine::delinearizeIndex(b, loc, index, basis);
assert(!failed(multiIndex) && "Failed to linearize img2col index"); assert(!failed(multiIndex) && "Failed to linearize img2col index");
return *multiIndex; return *multiIndex;
} }
// Given indices corresponding to iterators in the output (oIndex) and filter // Given indices corresponding to iterators in the output (oIndex) and filter
// (fIndex) for a convolution, compute the convolved index for the // (fIndex) for a convolution, compute the convolved index for the
// input as `oIndex * stride + fIndex`. // input as `oIndex * stride + fIndex`.
static Value getConvolvedIndex(OpBuilder &b, Location loc, Value oIndex, static Value getConvolvedIndex(OpBuilder &b, Location loc, Value oIndex,
Value fIndex, int64_t stride) { Value fIndex, int64_t stride) {
AffineExpr oExpr, fExpr; AffineExpr oExpr, fExpr;
bindSymbols(b.getContext(), oExpr, fExpr); bindSymbols(b.getContext(), oExpr, fExpr);
AffineMap convMap = AffineMap::get(0, 2, stride * oExpr + fExpr); AffineMap convMap = AffineMap::get(0, 2, stride * oExpr + fExpr);
return makeComposedAffineApply(b, loc, convMap, ValueRange{oIndex, fIndex}); return affine::makeComposedAffineApply(b, loc, convMap,
ValueRange{oIndex, fIndex});
} }
FailureOr<std::pair<Operation *, Operation *>> FailureOr<std::pair<Operation *, Operation *>>
rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNhwcHwcfOp convOp) { rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNhwcHwcfOp convOp) {
auto inputType = convOp.getInputs()[0].getType().cast<ShapedType>(); auto inputType = convOp.getInputs()[0].getType().cast<ShapedType>();
auto filterType = convOp.getInputs()[1].getType().cast<ShapedType>(); auto filterType = convOp.getInputs()[1].getType().cast<ShapedType>();
auto outputType = convOp.getOutputs()[0].getType().cast<ShapedType>(); auto outputType = convOp.getOutputs()[0].getType().cast<ShapedType>();
▲ Show 20 Lines • Show All 462 Lines • Show Last 20 Lines

mlir/lib/Dialect/Linalg/Transforms/DecomposeLinalgOps.cpp

Show First 20 Lines • Show All 105 Lines • ▼ Show 20 Linesstatic SmallVector<OpFoldResult> getGenericOpLoopRange(OpBuilder &b,
GenericOp op) { GenericOp op) {
OpBuilder::InsertionGuard g(b); OpBuilder::InsertionGuard g(b);
b.setInsertionPoint(op); b.setInsertionPoint(op);
Location loc = op.getLoc(); Location loc = op.getLoc();
auto allShapesSizes = auto allShapesSizes =
cast<LinalgOp>(op.getOperation()).createFlatListOfOperandDims(b, loc); cast<LinalgOp>(op.getOperation()).createFlatListOfOperandDims(b, loc);
AffineMap map = op.getShapesToLoopsMap(); AffineMap map = op.getShapesToLoopsMap();
IRRewriter rewriter(b); IRRewriter rewriter(b);
return makeComposedFoldedMultiResultAffineApply(rewriter, loc, map, return affine::makeComposedFoldedMultiResultAffineApply(rewriter, loc, map,
allShapesSizes); allShapesSizes);
} }
/// Helper method to permute the list of `values` based on the `map`. /// Helper method to permute the list of `values` based on the `map`.
SmallVector<OpFoldResult> permuteValues(ArrayRef<OpFoldResult> values, SmallVector<OpFoldResult> permuteValues(ArrayRef<OpFoldResult> values,
AffineMap map) { AffineMap map) {
assert(map.isPermutation()); assert(map.isPermutation());
SmallVector<OpFoldResult> permutedValues(values.size()); SmallVector<OpFoldResult> permutedValues(values.size());
for (const auto &position : for (const auto &position :
▲ Show 20 Lines • Show All 262 Lines • Show Last 20 Lines

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

Show First 20 Lines • Show All 173 Lines • ▼ Show 20 Lines if (producer.hasIndexSemantics()) {
SmallVector<Value> fusedIndices; SmallVector<Value> fusedIndices;
fusedIndices.reserve(numFusedOpLoops); fusedIndices.reserve(numFusedOpLoops);
llvm::transform(llvm::seq<uint64_t>(0, numFusedOpLoops), llvm::transform(llvm::seq<uint64_t>(0, numFusedOpLoops),
std::back_inserter(fusedIndices), [&](uint64_t dim) { std::back_inserter(fusedIndices), [&](uint64_t dim) {
return rewriter.create<IndexOp>(producer.getLoc(), dim); return rewriter.create<IndexOp>(producer.getLoc(), dim);
}); });
for (IndexOp indexOp : for (IndexOp indexOp :
llvm::make_early_inc_range(producerBlock.getOps<IndexOp>())) { llvm::make_early_inc_range(producerBlock.getOps<IndexOp>())) {
Value newIndex = rewriter.create<mlir::AffineApplyOp>( Value newIndex = rewriter.create<affine::AffineApplyOp>(
producer.getLoc(), producer.getLoc(),
consumerToProducerLoopsMap.getSubMap(indexOp.getDim()), fusedIndices); consumerToProducerLoopsMap.getSubMap(indexOp.getDim()), fusedIndices);
mapper.map(indexOp.getResult(), newIndex); mapper.map(indexOp.getResult(), newIndex);
} }
} }
// TODO: allow fusing the producer of an output operand. // TODO: allow fusing the producer of an output operand.
assert(consumer.isDpsInput(fusedOperand) && assert(consumer.isDpsInput(fusedOperand) &&
"expected producer of input operand"); "expected producer of input operand");
▲ Show 20 Lines • Show All 523 Lines • ▼ Show 20 Lines for (IndexOp indexOp :
llvm::transform( llvm::transform(
expandedDims.drop_front(), std::back_inserter(expandedIndices), expandedDims.drop_front(), std::back_inserter(expandedIndices),
[&](int64_t dim) { return rewriter.create<IndexOp>(loc, dim); }); [&](int64_t dim) { return rewriter.create<IndexOp>(loc, dim); });
Value newIndex = rewriter.create<IndexOp>(loc, expandedDims.front()); Value newIndex = rewriter.create<IndexOp>(loc, expandedDims.front());
for (auto it : llvm::zip(expandedDimsShape, expandedIndices)) { for (auto it : llvm::zip(expandedDimsShape, expandedIndices)) {
assert(!ShapedType::isDynamic(std::get<0>(it))); assert(!ShapedType::isDynamic(std::get<0>(it)));
AffineExpr idx, acc; AffineExpr idx, acc;
bindDims(rewriter.getContext(), idx, acc); bindDims(rewriter.getContext(), idx, acc);
newIndex = rewriter.create<AffineApplyOp>( newIndex = rewriter.create<affine::AffineApplyOp>(
indexOp.getLoc(), idx + acc * std::get<0>(it), indexOp.getLoc(), idx + acc * std::get<0>(it),
ValueRange{std::get<1>(it), newIndex}); ValueRange{std::get<1>(it), newIndex});
} }
rewriter.replaceOp(indexOp, newIndex); rewriter.replaceOp(indexOp, newIndex);
} }
} }
/// Implements the fusion of a tensor.collapse_shape or a tensor.expand_shape op /// Implements the fusion of a tensor.collapse_shape or a tensor.expand_shape op
▲ Show 20 Lines • Show All 1,135 Lines • ▼ Show 20 Lines ControlFusionFn defaultControlFn = [](OpOperand *fusedOperand) {
return producer && producer->hasOneUse(); return producer && producer->hasOneUse();
}; };
// Add elementwise op fusion patterns. // Add elementwise op fusion patterns.
populateElementwiseOpsFusionPatterns(patterns, defaultControlFn); populateElementwiseOpsFusionPatterns(patterns, defaultControlFn);
populateFoldReshapeOpsByExpansionPatterns(patterns, defaultControlFn); populateFoldReshapeOpsByExpansionPatterns(patterns, defaultControlFn);
// General canonicalization patterns. // General canonicalization patterns.
AffineApplyOp::getCanonicalizationPatterns(patterns, context); affine::AffineApplyOp::getCanonicalizationPatterns(patterns, context);
GenericOp::getCanonicalizationPatterns(patterns, context); GenericOp::getCanonicalizationPatterns(patterns, context);
tensor::ExpandShapeOp::getCanonicalizationPatterns(patterns, context); tensor::ExpandShapeOp::getCanonicalizationPatterns(patterns, context);
tensor::CollapseShapeOp::getCanonicalizationPatterns(patterns, context); tensor::CollapseShapeOp::getCanonicalizationPatterns(patterns, context);
context->getLoadedDialect<LinalgDialect>()->getCanonicalizationPatterns( context->getLoadedDialect<LinalgDialect>()->getCanonicalizationPatterns(
patterns); patterns);
// Add constant folding patterns. // Add constant folding patterns.
populateConstantFoldLinalgOperations(patterns, defaultControlFn); populateConstantFoldLinalgOperations(patterns, defaultControlFn);
Show All 13 Lines

mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp

Show First 20 Lines • Show All 456 Lines • ▼ Show 20 LinesHoistPaddingAnalysis::getHoistedPackedTensorSizes(RewriterBase &rewriter,
// Upper bound the packing loop lengths to size the packed tensor. Taking // Upper bound the packing loop lengths to size the packed tensor. Taking
// upper bounds can make the sizes of the packed tensor independent of the // upper bounds can make the sizes of the packed tensor independent of the
// enclosing loops. This independence is a prerequisite for reusing the same // enclosing loops. This independence is a prerequisite for reusing the same
// buffer for all enclosing loop iterations and hoisting its allocation out // buffer for all enclosing loop iterations and hoisting its allocation out
// of the enclosing loops. // of the enclosing loops.
for (auto forOp : packingLoops) { for (auto forOp : packingLoops) {
// Compute an upper bound `ubVal` for the upper bound of `forOp`. // Compute an upper bound `ubVal` for the upper bound of `forOp`.
FailureOr<OpFoldResult> loopUb = reifyIndexValueBound( FailureOr<OpFoldResult> loopUb = affine::reifyIndexValueBound(
rewriter, loc, presburger::BoundType::UB, forOp.getUpperBound(), rewriter, loc, presburger::BoundType::UB, forOp.getUpperBound(),
/*stopCondition=*/ /*stopCondition=*/
[&](Value v, std::optional<int64_t> d) { [&](Value v, std::optional<int64_t> d) {
if (v == forOp.getUpperBound()) if (v == forOp.getUpperBound())
return false; return false;
// Compute a bound that is independent of any affine op results. // Compute a bound that is independent of any affine op results.
Operation *op = v.getDefiningOp(); Operation *op = v.getDefiningOp();
if (!op) if (!op)
return true; return true;
return !isa<AffineMinOp, AffineMaxOp, AffineApplyOp>(op); return !isa<affine::AffineMinOp, affine::AffineMaxOp,
affine::AffineApplyOp>(op);
}, },
/*closedUB=*/true); /*closedUB=*/true);
assert(succeeded(loopUb) && "could not get upper bound"); assert(succeeded(loopUb) && "could not get upper bound");
Value ubVal = getValueOrCreateConstantIndexOp(rewriter, loc, *loopUb); Value ubVal = getValueOrCreateConstantIndexOp(rewriter, loc, *loopUb);
// Compute the maximal packing loop length as (ub - lb).ceilDiv(step) and // Compute the maximal packing loop length as (ub - lb).ceilDiv(step) and
// store the result to `dynamicTensorSizes`. // store the result to `dynamicTensorSizes`.
// TODO: instead of using the lower bound of `forOp` directly, implement a // TODO: instead of using the lower bound of `forOp` directly, implement a
// lower bound computation similar to the upper bound computation. // lower bound computation similar to the upper bound computation.
AffineExpr lb, ub, step; AffineExpr lb, ub, step;
bindDims(rewriter.getContext(), lb, ub); bindDims(rewriter.getContext(), lb, ub);
bindSymbols(rewriter.getContext(), step); bindSymbols(rewriter.getContext(), step);
Value res = rewriter.createOrFold<AffineApplyOp>( Value res = rewriter.createOrFold<affine::AffineApplyOp>(
loc, (ub - lb).ceilDiv(step), loc, (ub - lb).ceilDiv(step),
ValueRange{forOp.getLowerBound(), ubVal, ValueRange{forOp.getLowerBound(), ubVal,
cast<scf::ForOp>(forOp).getStep()}); cast<scf::ForOp>(forOp).getStep()});
dynamicTensorSizes.push_back(res); dynamicTensorSizes.push_back(res);
} }
return dynamicTensorSizes; return dynamicTensorSizes;
} }
Show All 17 Linesstatic Value buildLoopIterationCount(RewriterBase &rewriter, scf::ForOp outer,
bindDims(ctx, iv, lb); bindDims(ctx, iv, lb);
bindSymbols(ctx, step); bindSymbols(ctx, step);
if (!isDefinedOutsideOrConstant(outer, forOp.getLowerBound()) || if (!isDefinedOutsideOrConstant(outer, forOp.getLowerBound()) ||
!isDefinedOutsideOrConstant(outer, forOp.getStep())) !isDefinedOutsideOrConstant(outer, forOp.getStep()))
return Value(); return Value();
Value ivVal = forOp.getInductionVar(), lbVal = forOp.getLowerBound(), Value ivVal = forOp.getInductionVar(), lbVal = forOp.getLowerBound(),
stepVal = forOp.getStep(); stepVal = forOp.getStep();
auto loc = forOp->getLoc(); auto loc = forOp->getLoc();
return rewriter.createOrFold<AffineApplyOp>( return rewriter.createOrFold<affine::AffineApplyOp>(
loc, (iv - lb).ceilDiv(step), ValueRange{ivVal, lbVal, stepVal}); loc, (iv - lb).ceilDiv(step), ValueRange{ivVal, lbVal, stepVal});
} }
// Build a packing loop nest by iteratively traversing the backward slice and // Build a packing loop nest by iteratively traversing the backward slice and
// clone the operations, iteratively stepping into the loops that we encounter. // clone the operations, iteratively stepping into the loops that we encounter.
// The implementation proceeds in a stack-like fashion: // The implementation proceeds in a stack-like fashion:
// 1. Iteratively clone and step into the loops, pushing the // 1. Iteratively clone and step into the loops, pushing the
// `hoistedPackedTensor` // `hoistedPackedTensor`
▲ Show 20 Lines • Show All 478 Lines • Show Last 20 Lines

mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp

Show First 20 Lines • Show All 88 Lines • ▼ Show 20 Lines func.walk([&](vector::TransferReadOp transferRead) {
if (!transferRead.getShapedType().isa<MemRefType>()) if (!transferRead.getShapedType().isa<MemRefType>())
return WalkResult::advance(); return WalkResult::advance();
LLVM_DEBUG(DBGS() << "Candidate for hoisting: " LLVM_DEBUG(DBGS() << "Candidate for hoisting: "
<< *transferRead.getOperation() << "\n"); << *transferRead.getOperation() << "\n");
auto loop = dyn_cast<LoopLikeOpInterface>(transferRead->getParentOp()); auto loop = dyn_cast<LoopLikeOpInterface>(transferRead->getParentOp());
LLVM_DEBUG(DBGS() << "Parent op: " << *transferRead->getParentOp() LLVM_DEBUG(DBGS() << "Parent op: " << *transferRead->getParentOp()
<< "\n"); << "\n");
if (!isa_and_nonnull<scf::ForOp, AffineForOp>(loop)) if (!isa_and_nonnull<scf::ForOp, affine::AffineForOp>(loop))
return WalkResult::advance(); return WalkResult::advance();
LLVM_DEBUG(DBGS() << "Candidate read: " << *transferRead.getOperation() LLVM_DEBUG(DBGS() << "Candidate read: " << *transferRead.getOperation()
<< "\n"); << "\n");
SetVector<Operation *> forwardSlice; SetVector<Operation *> forwardSlice;
getForwardSlice(transferRead.getOperation(), &forwardSlice); getForwardSlice(transferRead.getOperation(), &forwardSlice);
▲ Show 20 Lines • Show All 89 Lines • ▼ Show 20 Lines func.walk([&](vector::TransferReadOp transferRead) {
transferWrite.getVectorMutable().assign( transferWrite.getVectorMutable().assign(
newForOp.getResults().back()); newForOp.getResults().back());
changed = true; changed = true;
loop.erase(); loop.erase();
// Need to interrupt and restart because erasing the loop messes up // Need to interrupt and restart because erasing the loop messes up
// the walk. // the walk.
return WalkResult::interrupt(); return WalkResult::interrupt();
}) })
.Case<AffineForOp>([&](AffineForOp affineForOp) { .Case<affine::AffineForOp>([&](affine::AffineForOp affineForOp) {
auto newForOp = replaceForOpWithNewYields( auto newForOp = replaceForOpWithNewYields(
b, affineForOp, transferRead.getVector(), b, affineForOp, transferRead.getVector(),
SmallVector<Value>{transferWrite.getVector()}, SmallVector<Value>{transferWrite.getVector()},
transferWrite.getVector()); transferWrite.getVector());
// Replace all uses of the `transferRead` with the corresponding // Replace all uses of the `transferRead` with the corresponding
// basic block argument. // basic block argument.
transferRead.getVector().replaceUsesWithIf( transferRead.getVector().replaceUsesWithIf(
newForOp.getLoopBody().getArguments().back(), newForOp.getLoopBody().getArguments().back(),
Show All 16 Lines

mlir/lib/Dialect/Linalg/Transforms/Interchange.cpp

Show First 20 Lines • Show All 90 Lines • ▼ Show 20 Lines for (IndexOp indexOp :
llvm::make_early_inc_range(genericOp.getBody()->getOps<IndexOp>())) { llvm::make_early_inc_range(genericOp.getBody()->getOps<IndexOp>())) {
rewriter.setInsertionPoint(indexOp); rewriter.setInsertionPoint(indexOp);
SmallVector<Value> allIndices; SmallVector<Value> allIndices;
allIndices.reserve(genericOp.getNumLoops()); allIndices.reserve(genericOp.getNumLoops());
llvm::transform(llvm::seq<uint64_t>(0, genericOp.getNumLoops()), llvm::transform(llvm::seq<uint64_t>(0, genericOp.getNumLoops()),
std::back_inserter(allIndices), [&](uint64_t dim) { std::back_inserter(allIndices), [&](uint64_t dim) {
return rewriter.create<IndexOp>(indexOp->getLoc(), dim); return rewriter.create<IndexOp>(indexOp->getLoc(), dim);
}); });
rewriter.replaceOpWithNewOp<AffineApplyOp>( rewriter.replaceOpWithNewOp<affine::AffineApplyOp>(
indexOp, permutationMap.getSubMap(indexOp.getDim()), allIndices); indexOp, permutationMap.getSubMap(indexOp.getDim()), allIndices);
} }
} }
return genericOp; return genericOp;
} }

mlir/lib/Dialect/Linalg/Transforms/Loops.cpp

Show First 20 Lines • Show All 43 Lines • ▼ Show 20 Linesstatic SmallVector<Value> makeCanonicalAffineApplies(OpBuilder &b, Location loc,
assert(map.getNumInputs() == vals.size()); assert(map.getNumInputs() == vals.size());
SmallVector<Value> res; SmallVector<Value> res;
res.reserve(map.getNumResults()); res.reserve(map.getNumResults());
auto dims = map.getNumDims(); auto dims = map.getNumDims();
for (auto e : map.getResults()) { for (auto e : map.getResults()) {
auto exprMap = AffineMap::get(dims, map.getNumSymbols(), e); auto exprMap = AffineMap::get(dims, map.getNumSymbols(), e);
SmallVector<Value> operands(vals.begin(), vals.end()); SmallVector<Value> operands(vals.begin(), vals.end());
canonicalizeMapAndOperands(&exprMap, &operands); affine::canonicalizeMapAndOperands(&exprMap, &operands);
res.push_back(b.create<AffineApplyOp>(loc, exprMap, operands)); res.push_back(b.create<affine::AffineApplyOp>(loc, exprMap, operands));
} }
return res; return res;
} }
template <typename LoadOpTy, typename StoreOpTy, typename OpType> template <typename LoadOpTy, typename StoreOpTy, typename OpType>
static void inlineRegionAndEmitStore(OpBuilder &b, Location loc, OpType op, static void inlineRegionAndEmitStore(OpBuilder &b, Location loc, OpType op,
ArrayRef<Value> indexedValues, ArrayRef<Value> indexedValues,
ArrayRef<SmallVector<Value>> indexing, ArrayRef<SmallVector<Value>> indexing,
▲ Show 20 Lines • Show All 122 Lines • ▼ Show 20 Lines for (Operation *loopOp : loopOps) {
llvm::TypeSwitch<Operation *>(loopOp) llvm::TypeSwitch<Operation *>(loopOp)
.Case([&](scf::ParallelOp parallelOp) { .Case([&](scf::ParallelOp parallelOp) {
allIvs.append(parallelOp.getInductionVars().begin(), allIvs.append(parallelOp.getInductionVars().begin(),
parallelOp.getInductionVars().end()); parallelOp.getInductionVars().end());
}) })
.Case([&](scf::ForOp forOp) { .Case([&](scf::ForOp forOp) {
allIvs.push_back(forOp.getInductionVar()); allIvs.push_back(forOp.getInductionVar());
}) })
.Case([&](AffineForOp affineForOp) { .Case([&](affine::AffineForOp affineForOp) {
allIvs.push_back(affineForOp.getInductionVar()); allIvs.push_back(affineForOp.getInductionVar());
}) })
.Default([&](Operation *op) { assert(false && "unexpected op"); }); .Default([&](Operation *op) { assert(false && "unexpected op"); });
} }
assert(linalgOp.getNumLoops() == allIvs.size() && assert(linalgOp.getNumLoops() == allIvs.size() &&
"expected the number of loops and induction variables to match"); "expected the number of loops and induction variables to match");
// Replace the index operations in the body of the innermost loop op. // Replace the index operations in the body of the innermost loop op.
if (!loopOps.empty()) { if (!loopOps.empty()) {
LoopLikeOpInterface loopOp = loopOps.back(); LoopLikeOpInterface loopOp = loopOps.back();
for (IndexOp indexOp : for (IndexOp indexOp :
llvm::make_early_inc_range(loopOp.getLoopBody().getOps<IndexOp>())) llvm::make_early_inc_range(loopOp.getLoopBody().getOps<IndexOp>()))
rewriter.replaceOp(indexOp, allIvs[indexOp.getDim()]); rewriter.replaceOp(indexOp, allIvs[indexOp.getDim()]);
} }
} }
template <typename LoopTy> template <typename LoopTy>
static FailureOr<LinalgLoops> linalgOpToLoopsImpl(RewriterBase &rewriter, static FailureOr<LinalgLoops> linalgOpToLoopsImpl(RewriterBase &rewriter,
LinalgOp linalgOp) { LinalgOp linalgOp) {
using LoadOpTy = std::conditional_t<std::is_same<LoopTy, AffineForOp>::value, using LoadOpTy =
AffineLoadOp, memref::LoadOp>; std::conditional_t<std::is_same<LoopTy, affine::AffineForOp>::value,
using StoreOpTy = std::conditional_t<std::is_same<LoopTy, AffineForOp>::value, affine::AffineLoadOp, memref::LoadOp>;
AffineStoreOp, memref::StoreOp>; using StoreOpTy =
std::conditional_t<std::is_same<LoopTy, affine::AffineForOp>::value,
affine::AffineStoreOp, memref::StoreOp>;
// The flattened loopToOperandRangesMaps is expected to be an invertible // The flattened loopToOperandRangesMaps is expected to be an invertible
// permutation map (which is asserted in the inverse calculation). // permutation map (which is asserted in the inverse calculation).
assert(linalgOp.hasBufferSemantics() && assert(linalgOp.hasBufferSemantics() &&
"expected linalg op with buffer semantics"); "expected linalg op with buffer semantics");
auto loopRanges = linalgOp.createLoopRanges(rewriter, linalgOp.getLoc()); auto loopRanges = linalgOp.createLoopRanges(rewriter, linalgOp.getLoc());
auto iteratorTypes = linalgOp.getIteratorTypesArray(); auto iteratorTypes = linalgOp.getIteratorTypesArray();
▲ Show 20 Lines • Show All 56 Lines • ▼ Show 20 Lines
/// 1. Zero operand + returns a single AffineConstantExpr /// 1. Zero operand + returns a single AffineConstantExpr
/// 2. One operand + returns a single AffineDimExpr /// 2. One operand + returns a single AffineDimExpr
/// 3. One operand + returns a single AffineSymbolExpr /// 3. One operand + returns a single AffineSymbolExpr
// //
/// In the first case, the AffineApplyOp is replaced by a new constant. In the /// In the first case, the AffineApplyOp is replaced by a new constant. In the
/// other cases, it is replaced by its unique operand. /// other cases, it is replaced by its unique operand.
struct FoldAffineOp : public RewritePattern { struct FoldAffineOp : public RewritePattern {
FoldAffineOp(MLIRContext *context) FoldAffineOp(MLIRContext *context)
: RewritePattern(AffineApplyOp::getOperationName(), 0, context) {} : RewritePattern(affine::AffineApplyOp::getOperationName(), 0, context) {}
LogicalResult matchAndRewrite(Operation *op, LogicalResult matchAndRewrite(Operation *op,
PatternRewriter &rewriter) const override { PatternRewriter &rewriter) const override {
AffineApplyOp affineApplyOp = cast<AffineApplyOp>(op); auto affineApplyOp = cast<affine::AffineApplyOp>(op);
auto map = affineApplyOp.getAffineMap(); auto map = affineApplyOp.getAffineMap();
if (map.getNumResults() != 1 || map.getNumInputs() > 1) if (map.getNumResults() != 1 || map.getNumInputs() > 1)
return failure(); return failure();
AffineExpr expr = map.getResult(0); AffineExpr expr = map.getResult(0);
if (map.getNumInputs() == 0) { if (map.getNumInputs() == 0) {
if (auto val = expr.dyn_cast<AffineConstantExpr>()) { if (auto val = expr.dyn_cast<AffineConstantExpr>()) {
rewriter.replaceOpWithNewOp<arith::ConstantIndexOp>(op, val.getValue()); rewriter.replaceOpWithNewOp<arith::ConstantIndexOp>(op, val.getValue());
Show All 11 Lines
template <typename LoopType> template <typename LoopType>
static void lowerLinalgToLoopsImpl(func::FuncOp funcOp) { static void lowerLinalgToLoopsImpl(func::FuncOp funcOp) {
MLIRContext *context = funcOp.getContext(); MLIRContext *context = funcOp.getContext();
RewritePatternSet patterns(context); RewritePatternSet patterns(context);
patterns.add<LinalgRewritePattern<LoopType>>(context); patterns.add<LinalgRewritePattern<LoopType>>(context);
memref::DimOp::getCanonicalizationPatterns(patterns, context); memref::DimOp::getCanonicalizationPatterns(patterns, context);
tensor::DimOp::getCanonicalizationPatterns(patterns, context); tensor::DimOp::getCanonicalizationPatterns(patterns, context);
AffineApplyOp::getCanonicalizationPatterns(patterns, context); affine::AffineApplyOp::getCanonicalizationPatterns(patterns, context);
patterns.add<FoldAffineOp>(context); patterns.add<FoldAffineOp>(context);
// Just apply the patterns greedily. // Just apply the patterns greedily.
(void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns)); (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
} }
struct LowerToAffineLoops struct LowerToAffineLoops
: public impl::LinalgLowerToAffineLoopsBase<LowerToAffineLoops> { : public impl::LinalgLowerToAffineLoopsBase<LowerToAffineLoops> {
void getDependentDialects(DialectRegistry &registry) const override { void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<memref::MemRefDialect>(); registry.insert<memref::MemRefDialect>();
} }
void runOnOperation() override { void runOnOperation() override {
lowerLinalgToLoopsImpl<AffineForOp>(getOperation()); lowerLinalgToLoopsImpl<affine::AffineForOp>(getOperation());
} }
}; };
struct LowerToLoops : public impl::LinalgLowerToLoopsBase<LowerToLoops> { struct LowerToLoops : public impl::LinalgLowerToLoopsBase<LowerToLoops> {
void getDependentDialects(DialectRegistry &registry) const override { void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<memref::MemRefDialect, scf::SCFDialect>(); registry.insert<memref::MemRefDialect, scf::SCFDialect>();
} }
void runOnOperation() override { void runOnOperation() override {
Show All 23 Lines
std::unique_ptr<OperationPass<func::FuncOp>> std::unique_ptr<OperationPass<func::FuncOp>>
mlir::createConvertLinalgToAffineLoopsPass() { mlir::createConvertLinalgToAffineLoopsPass() {
return std::make_unique<LowerToAffineLoops>(); return std::make_unique<LowerToAffineLoops>();
} }
/// Emits a loop nest of `affine.for` with the proper body for `linalgOp`. /// Emits a loop nest of `affine.for` with the proper body for `linalgOp`.
FailureOr<LinalgLoops> FailureOr<LinalgLoops>
mlir::linalg::linalgOpToAffineLoops(RewriterBase &rewriter, LinalgOp linalgOp) { mlir::linalg::linalgOpToAffineLoops(RewriterBase &rewriter, LinalgOp linalgOp) {
return linalgOpToLoopsImpl<AffineForOp>(rewriter, linalgOp); return linalgOpToLoopsImpl<affine::AffineForOp>(rewriter, linalgOp);
} }
/// Emits a loop nest of `scf.for` with the proper body for `linalgOp`. /// Emits a loop nest of `scf.for` with the proper body for `linalgOp`.
FailureOr<LinalgLoops> mlir::linalg::linalgOpToLoops(RewriterBase &rewriter, FailureOr<LinalgLoops> mlir::linalg::linalgOpToLoops(RewriterBase &rewriter,
LinalgOp linalgOp) { LinalgOp linalgOp) {
return linalgOpToLoopsImpl<scf::ForOp>(rewriter, linalgOp); return linalgOpToLoopsImpl<scf::ForOp>(rewriter, linalgOp);
} }
/// Emits a loop nest of `scf.parallel` with the proper body for `linalgOp`. /// Emits a loop nest of `scf.parallel` with the proper body for `linalgOp`.
FailureOr<LinalgLoops> FailureOr<LinalgLoops>
mlir::linalg::linalgOpToParallelLoops(RewriterBase &rewriter, mlir::linalg::linalgOpToParallelLoops(RewriterBase &rewriter,
LinalgOp linalgOp) { LinalgOp linalgOp) {
return linalgOpToLoopsImpl<scf::ParallelOp>(rewriter, linalgOp); return linalgOpToLoopsImpl<scf::ParallelOp>(rewriter, linalgOp);
} }

mlir/lib/Dialect/Linalg/Transforms/Split.cpp

Show First 20 Lines • Show All 75 Lines • ▼ Show 20 Lineslinalg::splitOp(RewriterBase &rewriter, TilingInterface op, unsigned dimension,
SmallVector<OpFoldResult> offsets = llvm::to_vector(llvm::map_range( SmallVector<OpFoldResult> offsets = llvm::to_vector(llvm::map_range(
iterationSpace, [](const Range &range) { return range.offset; })); iterationSpace, [](const Range &range) { return range.offset; }));
SmallVector<OpFoldResult> sizes = llvm::to_vector(llvm::map_range( SmallVector<OpFoldResult> sizes = llvm::to_vector(llvm::map_range(
iterationSpace, [](const Range &range) { return range.size; })); iterationSpace, [](const Range &range) { return range.size; }));
// Adjust the split point so that it doesn't overflow the size. // Adjust the split point so that it doesn't overflow the size.
AffineExpr d0, d1, d2; AffineExpr d0, d1, d2;
bindDims(rewriter.getContext(), d0, d1, d2); bindDims(rewriter.getContext(), d0, d1, d2);
OpFoldResult minSplitPoint = makeComposedFoldedAffineMin( OpFoldResult minSplitPoint = affine::makeComposedFoldedAffineMin(
rewriter, op.getLoc(), rewriter, op.getLoc(),
AffineMap::inferFromExprList(ArrayRef<AffineExpr>{d0, d1 + d2}).front(), AffineMap::inferFromExprList(ArrayRef<AffineExpr>{d0, d1 + d2}).front(),
{splitPoint, offsets[dimension], sizes[dimension]}); {splitPoint, offsets[dimension], sizes[dimension]});
// Compute the size of the second part. Return early if the second part would // Compute the size of the second part. Return early if the second part would
// have an empty iteration space. // have an empty iteration space.
OpFoldResult remainingSize = makeComposedFoldedAffineApply( OpFoldResult remainingSize = affine::makeComposedFoldedAffineApply(
rewriter, op.getLoc(), d0 + d1 - d2, rewriter, op.getLoc(), d0 + d1 - d2,
{iterationSpace[dimension].offset, iterationSpace[dimension].size, {iterationSpace[dimension].offset, iterationSpace[dimension].size,
minSplitPoint}); minSplitPoint});
if (auto attr = remainingSize.dyn_cast<Attribute>()) { if (auto attr = remainingSize.dyn_cast<Attribute>()) {
if (attr.cast<IntegerAttr>().getValue().isZero()) if (attr.cast<IntegerAttr>().getValue().isZero())
return {op, TilingInterface()}; return {op, TilingInterface()};
} }
Show All 16 Lineslinalg::splitOp(RewriterBase &rewriter, TilingInterface op, unsigned dimension,
rewriter.updateRootInPlace(op, [&]() { rewriter.updateRootInPlace(op, [&]() {
unsigned numTotalOperands = op->getNumOperands(); unsigned numTotalOperands = op->getNumOperands();
unsigned numOutputOperands = firstResults.size(); unsigned numOutputOperands = firstResults.size();
op->setOperands(numTotalOperands - numOutputOperands, numOutputOperands, op->setOperands(numTotalOperands - numOutputOperands, numOutputOperands,
firstResults); firstResults);
}); });
// Create the second part. // Create the second part.
OpFoldResult totalOffset = makeComposedFoldedAffineApply( OpFoldResult totalOffset = affine::makeComposedFoldedAffineApply(
rewriter, op.getLoc(), d0 + d1, {offsets[dimension], minSplitPoint}); rewriter, op.getLoc(), d0 + d1, {offsets[dimension], minSplitPoint});
SmallVector<Value> secondResults; SmallVector<Value> secondResults;
TilingInterface secondPart = TilingInterface secondPart =
createSplitPart(rewriter, op.getLoc(), op, offsets, sizes, firstResults, createSplitPart(rewriter, op.getLoc(), op, offsets, sizes, firstResults,
dimension, remainingSize, totalOffset, secondResults); dimension, remainingSize, totalOffset, secondResults);
// Propagate any errors in part creation. // Propagate any errors in part creation.
if (!firstPart || !secondPart) if (!firstPart || !secondPart)
return {TilingInterface(), TilingInterface()}; return {TilingInterface(), TilingInterface()};
// Replace the original op with the results of the two newly created ops. // Replace the original op with the results of the two newly created ops.
rewriter.replaceOp(op, secondResults); rewriter.replaceOp(op, secondResults);
return {firstPart, secondPart}; return {firstPart, secondPart};
} }

mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp

Show All 33 Lines
#include <utility> #include <utility>
namespace mlir { namespace mlir {
#define GEN_PASS_DEF_LINALGTILINGPASS #define GEN_PASS_DEF_LINALGTILINGPASS
#include "mlir/Dialect/Linalg/Passes.h.inc" #include "mlir/Dialect/Linalg/Passes.h.inc"
} // namespace mlir } // namespace mlir
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
using namespace mlir::linalg; using namespace mlir::linalg;
using namespace mlir::scf; using namespace mlir::scf;
#define DEBUG_TYPE "linalg-tiling" #define DEBUG_TYPE "linalg-tiling"
static bool isZero(OpFoldResult v) { static bool isZero(OpFoldResult v) {
if (!v) if (!v)
return false; return false;
▲ Show 20 Lines • Show All 123 Lines • ▼ Show 20 Linesmlir::linalg::computeMultiTileSizes(OpBuilder &builder, LinalgOp op,
Value tripCount = Value tripCount =
getValueOrCreateConstantIndexOp(b, op.getLoc(), loopRanges[dimension]); getValueOrCreateConstantIndexOp(b, op.getLoc(), loopRanges[dimension]);
// Compute the tile sizes and the respective numbers of tiles. // Compute the tile sizes and the respective numbers of tiles.
AffineExpr s0 = b.getAffineSymbolExpr(0); AffineExpr s0 = b.getAffineSymbolExpr(0);
AffineExpr s1 = b.getAffineSymbolExpr(1); AffineExpr s1 = b.getAffineSymbolExpr(1);
AffineExpr s2 = b.getAffineSymbolExpr(2); AffineExpr s2 = b.getAffineSymbolExpr(2);
auto apply = [&](AffineExpr expr, ValueRange values) -> Value { auto apply = [&](AffineExpr expr, ValueRange values) -> Value {
return makeComposedAffineApply(b, b.getLoc(), expr, values); return affine::makeComposedAffineApply(b, b.getLoc(), expr, values);
}; };
Value a = apply(s0.floorDiv(s1), {tripCount, divisorValue}); Value a = apply(s0.floorDiv(s1), {tripCount, divisorValue});
Value t = apply((s0 + s1 - 1).floorDiv(s1), {targetSizeValue, divisorValue}); Value t = apply((s0 + s1 - 1).floorDiv(s1), {targetSizeValue, divisorValue});
Value d = apply((s0 + s1 - 1).floorDiv(s1), {a, t}); Value d = apply((s0 + s1 - 1).floorDiv(s1), {a, t});
Value s = apply(s0.floorDiv(s1) * s2, {a, d, divisorValue}); Value s = apply(s0.floorDiv(s1) * s2, {a, d, divisorValue});
Value v = apply(s0 % s1, {a, d}); Value v = apply(s0 % s1, {a, d});
Value u = apply(s0 - s1, {d, v}); Value u = apply(s0 - s1, {d, v});
Show All 33 Linesstatic bool canOmitTileOffsetInBoundsCheck(OpFoldResult tileSize,
if (!tileSizeConst || !numThreadsConst || !iterSizeConst) if (!tileSizeConst || !numThreadsConst || !iterSizeConst)
return false; return false;
return *tileSizeConst * (*numThreadsConst - 1) < *iterSizeConst; return *tileSizeConst * (*numThreadsConst - 1) < *iterSizeConst;
} }
/// Build an `affine_max` of all the `vals`. /// Build an `affine_max` of all the `vals`.
static OpFoldResult buildMax(OpBuilder &b, Location loc, static OpFoldResult buildMax(OpBuilder &b, Location loc,
ArrayRef<OpFoldResult> vals) { ArrayRef<OpFoldResult> vals) {
return makeComposedFoldedAffineMax( return affine::makeComposedFoldedAffineMax(
b, loc, AffineMap::getMultiDimIdentityMap(vals.size(), loc.getContext()), b, loc, AffineMap::getMultiDimIdentityMap(vals.size(), loc.getContext()),
vals); vals);
} }
/// Build an `affine_min` of all the `vals`. /// Build an `affine_min` of all the `vals`.
static OpFoldResult buildMin(OpBuilder &b, Location loc, static OpFoldResult buildMin(OpBuilder &b, Location loc,
ArrayRef<OpFoldResult> vals) { ArrayRef<OpFoldResult> vals) {
return makeComposedFoldedAffineMin( return affine::makeComposedFoldedAffineMin(
b, loc, AffineMap::getMultiDimIdentityMap(vals.size(), loc.getContext()), b, loc, AffineMap::getMultiDimIdentityMap(vals.size(), loc.getContext()),
vals); vals);
} }
/// Fill out the `tiledOffsets` and `tiledSizes` to be used to tile to a given /// Fill out the `tiledOffsets` and `tiledSizes` to be used to tile to a given
/// number of threads. /// number of threads.
static void calculateTileOffsetsAndSizes( static void calculateTileOffsetsAndSizes(
RewriterBase &b, Location loc, scf::ForallOp forallOp, RewriterBase &b, Location loc, scf::ForallOp forallOp,
▲ Show 20 Lines • Show All 715 Lines • ▼ Show 20 Linesmlir::linalg::getLinalgTilingCanonicalizationPatterns(MLIRContext *ctx) {
RewritePatternSet patterns(ctx); RewritePatternSet patterns(ctx);
populateLinalgTilingCanonicalizationPatterns(patterns); populateLinalgTilingCanonicalizationPatterns(patterns);
return patterns; return patterns;
} }
void mlir::linalg::populateLinalgTilingCanonicalizationPatterns( void mlir::linalg::populateLinalgTilingCanonicalizationPatterns(
RewritePatternSet &patterns) { RewritePatternSet &patterns) {
auto *ctx = patterns.getContext(); auto *ctx = patterns.getContext();
AffineApplyOp::getCanonicalizationPatterns(patterns, ctx); affine::AffineApplyOp::getCanonicalizationPatterns(patterns, ctx);
AffineForOp::getCanonicalizationPatterns(patterns, ctx); affine::AffineForOp::getCanonicalizationPatterns(patterns, ctx);
AffineMinOp::getCanonicalizationPatterns(patterns, ctx); affine::AffineMinOp::getCanonicalizationPatterns(patterns, ctx);
AffineMaxOp::getCanonicalizationPatterns(patterns, ctx); affine::AffineMaxOp::getCanonicalizationPatterns(patterns, ctx);
arith::ConstantIndexOp::getCanonicalizationPatterns(patterns, ctx); arith::ConstantIndexOp::getCanonicalizationPatterns(patterns, ctx);
memref::SubViewOp::getCanonicalizationPatterns(patterns, ctx); memref::SubViewOp::getCanonicalizationPatterns(patterns, ctx);
memref::ViewOp::getCanonicalizationPatterns(patterns, ctx); memref::ViewOp::getCanonicalizationPatterns(patterns, ctx);
scf::ForOp::getCanonicalizationPatterns(patterns, ctx); scf::ForOp::getCanonicalizationPatterns(patterns, ctx);
scf::ParallelOp::getCanonicalizationPatterns(patterns, ctx); scf::ParallelOp::getCanonicalizationPatterns(patterns, ctx);
Show All 18 Lines

mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp

Show All 31 Lines
static SmallVector<Value> getIndicesForAccess(OpBuilder &b, Location loc, static SmallVector<Value> getIndicesForAccess(OpBuilder &b, Location loc,
AffineMap indexingMap, AffineMap indexingMap,
ValueRange ivs) { ValueRange ivs) {
SmallVector<Value> indices; SmallVector<Value> indices;
indices.reserve(indexingMap.getNumResults()); indices.reserve(indexingMap.getNumResults());
for (auto result : indexingMap.getResults()) { for (auto result : indexingMap.getResults()) {
AffineMap m = AffineMap::get(indexingMap.getNumDims(), AffineMap m = AffineMap::get(indexingMap.getNumDims(),
indexingMap.getNumSymbols(), result); indexingMap.getNumSymbols(), result);
Value v = b.create<AffineApplyOp>(loc, m, ivs); Value v = b.create<affine::AffineApplyOp>(loc, m, ivs);
indices.push_back(v); indices.push_back(v);
} }
return indices; return indices;
} }
/// Method to inline the payload of a `linalgOp` given the iteration space /// Method to inline the payload of a `linalgOp` given the iteration space
/// point and values for the arguments of the payload. /// point and values for the arguments of the payload.
static LogicalResult inlinePayload(OpBuilder &b, LinalgOp linalgOp, static LogicalResult inlinePayload(OpBuilder &b, LinalgOp linalgOp,
▲ Show 20 Lines • Show All 50 Lines • ▼ Show 20 Lines SmallVector<Range> getIterationDomain(Operation *op, OpBuilder &b) const {
Location loc = op->getLoc(); Location loc = op->getLoc();
LinalgOp linalgOp = cast<LinalgOp>(op); LinalgOp linalgOp = cast<LinalgOp>(op);
SmallVector<OpFoldResult> allShapesSizes = SmallVector<OpFoldResult> allShapesSizes =
linalgOp.createFlatListOfOperandDims(b, loc); linalgOp.createFlatListOfOperandDims(b, loc);
AffineMap map = linalgOp.getShapesToLoopsMap(); AffineMap map = linalgOp.getShapesToLoopsMap();
return llvm::to_vector( return llvm::to_vector(
llvm::map_range(map.getResults(), [&](AffineExpr loopExpr) { llvm::map_range(map.getResults(), [&](AffineExpr loopExpr) {
OpFoldResult ofr = OpFoldResult ofr = affine::makeComposedFoldedAffineApply(
makeComposedFoldedAffineApply(b, loc, loopExpr, allShapesSizes); b, loc, loopExpr, allShapesSizes);
return Range{b.getIndexAttr(0), ofr, b.getIndexAttr(1)}; return Range{b.getIndexAttr(0), ofr, b.getIndexAttr(1)};
})); }));
} }
// Instantiate the tiled implementation of the operation. // Instantiate the tiled implementation of the operation.
FailureOr<TilingResult> FailureOr<TilingResult>
getTiledImplementation(Operation *op, OpBuilder &b, getTiledImplementation(Operation *op, OpBuilder &b,
ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> offsets,
Show All 25 Lines getResultTilePosition(Operation *op, OpBuilder &b, unsigned resultNumber,
SmallVector<OpFoldResult> &resultSizes) const { SmallVector<OpFoldResult> &resultSizes) const {
Location loc = op->getLoc(); Location loc = op->getLoc();
LinalgOp linalgOp = cast<LinalgOp>(op); LinalgOp linalgOp = cast<LinalgOp>(op);
AffineExpr d0; AffineExpr d0;
bindDims(b.getContext(), d0); bindDims(b.getContext(), d0);
SmallVector<OpFoldResult> subShapeSizes = SmallVector<OpFoldResult> subShapeSizes =
llvm::to_vector(llvm::map_range(sizes, [&](OpFoldResult ofr) { llvm::to_vector(llvm::map_range(sizes, [&](OpFoldResult ofr) {
return makeComposedFoldedAffineApply(b, loc, d0 - 1, ofr); return affine::makeComposedFoldedAffineApply(b, loc, d0 - 1, ofr);
})); }));
OpOperand *outOperand = linalgOp.getDpsInitOperand(resultNumber); OpOperand *outOperand = linalgOp.getDpsInitOperand(resultNumber);
SliceParameters sliceParams = computeSliceParameters( SliceParameters sliceParams = computeSliceParameters(
b, loc, outOperand->get(), sizes, b, loc, outOperand->get(), sizes,
linalgOp.getMatchingIndexingMap(outOperand), offsets, linalgOp.getMatchingIndexingMap(outOperand), offsets,
/*ubs*/ {}, subShapeSizes, true); /*ubs*/ {}, subShapeSizes, true);
resultOffsets = sliceParams.offsets; resultOffsets = sliceParams.offsets;
▲ Show 20 Lines • Show All 270 Lines • Show Last 20 Lines

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

Show First 20 Lines • Show All 1,368 Lines • ▼ Show 20 Lines if (failed(reductionPreconditions(linalgOp))) {
return failure(); return failure();
} }
return success(); return success();
} }
/// Converts affine.apply Ops to arithmetic operations. /// Converts affine.apply Ops to arithmetic operations.
static void convertAffineApply(RewriterBase &rewriter, LinalgOp linalgOp) { static void convertAffineApply(RewriterBase &rewriter, LinalgOp linalgOp) {
OpBuilder::InsertionGuard g(rewriter); OpBuilder::InsertionGuard g(rewriter);
auto toReplace = linalgOp.getBlock()->getOps<AffineApplyOp>(); auto toReplace = linalgOp.getBlock()->getOps<affine::AffineApplyOp>();
for (auto op : make_early_inc_range(toReplace)) { for (auto op : make_early_inc_range(toReplace)) {
rewriter.setInsertionPoint(op); rewriter.setInsertionPoint(op);
auto expanded = expandAffineExpr( auto expanded = affine::expandAffineExpr(
rewriter, op->getLoc(), op.getAffineMap().getResult(0), rewriter, op->getLoc(), op.getAffineMap().getResult(0),
op.getOperands().take_front(op.getAffineMap().getNumDims()), op.getOperands().take_front(op.getAffineMap().getNumDims()),
op.getOperands().take_back(op.getAffineMap().getNumSymbols())); op.getOperands().take_back(op.getAffineMap().getNumSymbols()));
rewriter.replaceOp(op, expanded); rewriter.replaceOp(op, expanded);
} }
} }
FailureOr<vector::TransferWriteOp> FailureOr<vector::TransferWriteOp>
▲ Show 20 Lines • Show All 473 Lines • ▼ Show 20 Lines for (unsigned i = 0; i < t1.getRank(); ++i) {
// Other cases: Take a deeper look at defining ops of values. // Other cases: Take a deeper look at defining ops of values.
auto v1 = size1.dyn_cast<Value>(); auto v1 = size1.dyn_cast<Value>();
auto v2 = size2.dyn_cast<Value>(); auto v2 = size2.dyn_cast<Value>();
if (!v1 || !v2) if (!v1 || !v2)
return false; return false;
// Case 2: Both values are identical AffineMinOps. (Should not happen if // Case 2: Both values are identical AffineMinOps. (Should not happen if
// CSE is run.) // CSE is run.)
auto minOp1 = v1.getDefiningOp<AffineMinOp>(); auto minOp1 = v1.getDefiningOp<affine::AffineMinOp>();
auto minOp2 = v2.getDefiningOp<AffineMinOp>(); auto minOp2 = v2.getDefiningOp<affine::AffineMinOp>();
if (minOp1 && minOp2 && minOp1.getAffineMap() == minOp2.getAffineMap() && if (minOp1 && minOp2 && minOp1.getAffineMap() == minOp2.getAffineMap() &&
minOp1.getOperands() == minOp2.getOperands()) minOp1.getOperands() == minOp2.getOperands())
continue; continue;
// Add additional cases as needed. // Add additional cases as needed.
} }
// All tests passed. // All tests passed.
▲ Show 20 Lines • Show All 1,055 Lines • Show Last 20 Lines

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

Show All 36 Lines
#include "llvm/ADT/TypeSwitch.h" #include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/Debug.h" #include "llvm/Support/Debug.h"
#include <optional> #include <optional>
#define DEBUG_TYPE "linalg-utils" #define DEBUG_TYPE "linalg-utils"
using namespace mlir; using namespace mlir;
using namespace presburger; using namespace presburger;
using namespace mlir::affine;
using namespace mlir::linalg; using namespace mlir::linalg;
using namespace mlir::scf; using namespace mlir::scf;
namespace { namespace {
// Helper visitor to determine whether an AffineExpr is tiled. // Helper visitor to determine whether an AffineExpr is tiled.
// This is achieved by traversing every AffineDimExpr with position `pos` and // This is achieved by traversing every AffineDimExpr with position `pos` and
// checking whether the corresponding `tileSizes[pos]` is non-zero. // checking whether the corresponding `tileSizes[pos]` is non-zero.
▲ Show 20 Lines • Show All 398 Lines • ▼ Show 20 Linesvoid GenerateLoopNest<AffineForOp>::doit(
SmallVector<int64_t, 4> constantSteps; SmallVector<int64_t, 4> constantSteps;
constantSteps.reserve(steps.size()); constantSteps.reserve(steps.size());
for (Value v : steps) { for (Value v : steps) {
auto op = v.getDefiningOp<arith::ConstantIndexOp>(); auto op = v.getDefiningOp<arith::ConstantIndexOp>();
assert(op && "Affine loops require constant steps"); assert(op && "Affine loops require constant steps");
constantSteps.push_back(op.value()); constantSteps.push_back(op.value());
} }
mlir::buildAffineLoopNest(b, loc, lbs, ubs, constantSteps, affine::buildAffineLoopNest(b, loc, lbs, ubs, constantSteps,
[&](OpBuilder &b, Location loc, ValueRange ivs) { [&](OpBuilder &b, Location loc, ValueRange ivs) {
bodyBuilderFn(b, loc, ivs, bodyBuilderFn(b, loc, ivs,
linalgOp->getOperands()); linalgOp->getOperands());
}); });
} }
/// Update the `lb`, `ub` and `step` to get per processor `lb`, `ub` and `step`. /// Update the `lb`, `ub` and `step` to get per processor `lb`, `ub` and `step`.
void updateBoundsForCyclicDistribution(OpBuilder &b, Location loc, Value procId, void updateBoundsForCyclicDistribution(OpBuilder &b, Location loc, Value procId,
Value nprocs, Value &lb, Value &ub, Value nprocs, Value &lb, Value &ub,
Value &step) { Value &step) {
AffineExpr d0, d1; AffineExpr d0, d1;
bindDims(b.getContext(), d0, d1); bindDims(b.getContext(), d0, d1);
AffineExpr s0 = getAffineSymbolExpr(0, b.getContext()); AffineExpr s0 = getAffineSymbolExpr(0, b.getContext());
lb = makeComposedAffineApply(b, loc, d0 + d1 * s0, {lb, procId, step}); lb =
step = makeComposedAffineApply(b, loc, d0 * s0, {nprocs, step}); affine::makeComposedAffineApply(b, loc, d0 + d1 * s0, {lb, procId, step});
step = affine::makeComposedAffineApply(b, loc, d0 * s0, {nprocs, step});
} }
/// Generates a loop nest consisting of scf.parallel and scf.for, depending /// Generates a loop nest consisting of scf.parallel and scf.for, depending
/// on the `iteratorTypes.` Consecutive parallel loops create a single /// on the `iteratorTypes.` Consecutive parallel loops create a single
/// scf.parallel operation; each sequential loop creates a new scf.for /// scf.parallel operation; each sequential loop creates a new scf.for
/// operation. The body of the innermost loop is populated by /// operation. The body of the innermost loop is populated by
/// `bodyBuilderFn` that accepts a range of induction variables for all /// `bodyBuilderFn` that accepts a range of induction variables for all
/// loops. `ivStorage` is used to store the partial list of induction /// loops. `ivStorage` is used to store the partial list of induction
▲ Show 20 Lines • Show All 536 Lines • Show Last 20 Lines

mlir/lib/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp

Show First 20 Lines • Show All 104 Lines • ▼ Show 20 Lines
class MemRefTransformDialectExtension class MemRefTransformDialectExtension
: public transform::TransformDialectExtension< : public transform::TransformDialectExtension<
MemRefTransformDialectExtension> { MemRefTransformDialectExtension> {
public: public:
using Base::Base; using Base::Base;
void init() { void init() {
declareDependentDialect<pdl::PDLDialect>(); declareDependentDialect<pdl::PDLDialect>();
declareGeneratedDialect<AffineDialect>(); declareGeneratedDialect<affine::AffineDialect>();
declareGeneratedDialect<arith::ArithDialect>(); declareGeneratedDialect<arith::ArithDialect>();
declareGeneratedDialect<memref::MemRefDialect>(); declareGeneratedDialect<memref::MemRefDialect>();
declareGeneratedDialect<nvgpu::NVGPUDialect>(); declareGeneratedDialect<nvgpu::NVGPUDialect>();
declareGeneratedDialect<vector::VectorDialect>(); declareGeneratedDialect<vector::VectorDialect>();
registerTransformOps< registerTransformOps<
#define GET_OP_LIST #define GET_OP_LIST
#include "mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp.inc" #include "mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp.inc"
Show All 12 Lines

mlir/lib/Dialect/MemRef/Transforms/ComposeSubView.cpp

Show First 20 Lines • Show All 105 Lines • ▼ Show 20 Lines for (auto it : llvm::zip(op.getMixedOffsets(), sourceOp.getMixedOffsets(),
} else { } else {
expr = expr =
expr + rewriter.getAffineSymbolExpr(affineApplyOperands.size()); expr + rewriter.getAffineSymbolExpr(affineApplyOperands.size());
affineApplyOperands.push_back(valueOrAttr.get<Value>()); affineApplyOperands.push_back(valueOrAttr.get<Value>());
} }
} }
AffineMap map = AffineMap::get(0, affineApplyOperands.size(), expr); AffineMap map = AffineMap::get(0, affineApplyOperands.size(), expr);
Value result = rewriter.create<AffineApplyOp>(op.getLoc(), map, Value result = rewriter.create<affine::AffineApplyOp>(
affineApplyOperands); op.getLoc(), map, affineApplyOperands);
offsets.push_back(result); offsets.push_back(result);
} }
} }
// This replaces 'op' but leaves 'sourceOp' alone; if it no longer has any // This replaces 'op' but leaves 'sourceOp' alone; if it no longer has any
// uses it can be removed by a (separate) dead code elimination pass. // uses it can be removed by a (separate) dead code elimination pass.
rewriter.replaceOpWithNewOp<memref::SubViewOp>(op, sourceOp.getSource(), rewriter.replaceOpWithNewOp<memref::SubViewOp>(op, sourceOp.getSource(),
offsets, sizes, strides); offsets, sizes, strides);
Show All 10 Lines

mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp

Show All 26 Lines
#include <optional> #include <optional>
namespace mlir { namespace mlir {
namespace memref { namespace memref {
#define GEN_PASS_DEF_EXPANDSTRIDEDMETADATA #define GEN_PASS_DEF_EXPANDSTRIDEDMETADATA
#include "mlir/Dialect/MemRef/Transforms/Passes.h.inc" #include "mlir/Dialect/MemRef/Transforms/Passes.h.inc"
} // namespace memref } // namespace memref
} // namespace mlir } // namespace mlir
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
namespace { namespace {
struct StridedMetadata { struct StridedMetadata {
Value basePtr; Value basePtr;
OpFoldResult offset; OpFoldResult offset;
SmallVector<OpFoldResult> sizes; SmallVector<OpFoldResult> sizes;
SmallVector<OpFoldResult> strides; SmallVector<OpFoldResult> strides;
▲ Show 20 Lines • Show All 906 Lines • Show Last 20 Lines

mlir/lib/Dialect/MemRef/Transforms/ExtractAddressComputations.cpp

Show First 20 Lines • Show All 189 Lines • ▼ Show 20 LinesgetGenericOpViewSizeForEachDim(RewriterBase &rewriter,
SmallVector<OpFoldResult> indices = SmallVector<OpFoldResult> indices =
getAsOpFoldResult(loadStoreLikeOp.getIndices()); getAsOpFoldResult(loadStoreLikeOp.getIndices());
SmallVector<OpFoldResult> finalSizes; SmallVector<OpFoldResult> finalSizes;
AffineExpr s0 = rewriter.getAffineSymbolExpr(0); AffineExpr s0 = rewriter.getAffineSymbolExpr(0);
AffineExpr s1 = rewriter.getAffineSymbolExpr(1); AffineExpr s1 = rewriter.getAffineSymbolExpr(1);
for (auto [srcSize, indice] : llvm::zip(srcSizes, indices)) { for (auto [srcSize, indice] : llvm::zip(srcSizes, indices)) {
finalSizes.push_back(makeComposedFoldedAffineApply(rewriter, loc, s0 - s1, finalSizes.push_back(affine::makeComposedFoldedAffineApply(
{srcSize, indice})); rewriter, loc, s0 - s1, {srcSize, indice}));
} }
return finalSizes; return finalSizes;
} }
/// Rewrite a store/load-like op so that all its indices are zeros. /// Rewrite a store/load-like op so that all its indices are zeros.
/// E.g., %ld = memref.load %base[%off0]...[%offN] /// E.g., %ld = memref.load %base[%off0]...[%offN]
/// => /// =>
/// %new_base = subview %base[%off0,.., %offN][1,..,1][1,..,1] /// %new_base = subview %base[%off0,.., %offN][1,..,1][1,..,1]
▲ Show 20 Lines • Show All 106 Lines • Show Last 20 Lines

mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp

Show First 20 Lines • Show All 79 Lines • ▼ Show 20 Lines for (ArrayRef<int64_t> groups : expandShapeOp.getReassociationIndices()) {
/// Apply permutation and create AffineApplyOp. /// Apply permutation and create AffineApplyOp.
SmallVector<OpFoldResult> dynamicIndices(groupSize); SmallVector<OpFoldResult> dynamicIndices(groupSize);
for (int64_t i = 0; i < groupSize; i++) for (int64_t i = 0; i < groupSize; i++)
dynamicIndices[i] = indices[groups[i]]; dynamicIndices[i] = indices[groups[i]];
// Creating maximally folded and composd affine.apply composes better with // Creating maximally folded and composd affine.apply composes better with
// other transformations without interleaving canonicalization passes. // other transformations without interleaving canonicalization passes.
OpFoldResult ofr = makeComposedFoldedAffineApply( OpFoldResult ofr = affine::makeComposedFoldedAffineApply(
rewriter, loc, rewriter, loc,
AffineMap::get(/*numDims=*/groupSize, AffineMap::get(/*numDims=*/groupSize,
/*numSymbols=*/0, srcIndexExpr), /*numSymbols=*/0, srcIndexExpr),
dynamicIndices); dynamicIndices);
sourceIndices.push_back( sourceIndices.push_back(
getValueOrCreateConstantIndexOp(rewriter, loc, ofr)); getValueOrCreateConstantIndexOp(rewriter, loc, ofr));
} }
return success(); return success();
Show All 33 Lines for (ArrayRef<int64_t> groups : collapseShapeOp.getReassociationIndices()) {
// Derive the index values along all dimensions of the source corresponding // Derive the index values along all dimensions of the source corresponding
// to the index wrt to collapsed shape op output. // to the index wrt to collapsed shape op output.
auto d0 = rewriter.getAffineDimExpr(0); auto d0 = rewriter.getAffineDimExpr(0);
SmallVector<AffineExpr> delinearizingExprs = delinearize(d0, suffixProduct); SmallVector<AffineExpr> delinearizingExprs = delinearize(d0, suffixProduct);
// Construct the AffineApplyOp for each delinearizingExpr. // Construct the AffineApplyOp for each delinearizingExpr.
for (int64_t i = 0; i < groupSize; i++) { for (int64_t i = 0; i < groupSize; i++) {
OpFoldResult ofr = makeComposedFoldedAffineApply( OpFoldResult ofr = affine::makeComposedFoldedAffineApply(
rewriter, loc, rewriter, loc,
AffineMap::get(/*numDims=*/1, /*numSymbols=*/0, AffineMap::get(/*numDims=*/1, /*numSymbols=*/0,
delinearizingExprs[i]), delinearizingExprs[i]),
dynamicIndices); dynamicIndices);
sourceIndices.push_back( sourceIndices.push_back(
getValueOrCreateConstantIndexOp(rewriter, loc, ofr)); getValueOrCreateConstantIndexOp(rewriter, loc, ofr));
} }
dynamicIndices.clear(); dynamicIndices.clear();
} }
if (collapseShapeOp.getReassociationIndices().empty()) { if (collapseShapeOp.getReassociationIndices().empty()) {
auto zeroAffineMap = rewriter.getConstantAffineMap(0); auto zeroAffineMap = rewriter.getConstantAffineMap(0);
int64_t srcRank = int64_t srcRank =
collapseShapeOp.getViewSource().getType().cast<MemRefType>().getRank(); collapseShapeOp.getViewSource().getType().cast<MemRefType>().getRank();
for (int64_t i = 0; i < srcRank; i++) { for (int64_t i = 0; i < srcRank; i++) {
OpFoldResult ofr = makeComposedFoldedAffineApply( OpFoldResult ofr = affine::makeComposedFoldedAffineApply(
rewriter, loc, zeroAffineMap, dynamicIndices); rewriter, loc, zeroAffineMap, dynamicIndices);
sourceIndices.push_back( sourceIndices.push_back(
getValueOrCreateConstantIndexOp(rewriter, loc, ofr)); getValueOrCreateConstantIndexOp(rewriter, loc, ofr));
} }
} }
return success(); return success();
} }
▲ Show 20 Lines • Show All 101 Lines • ▼ Show 20 Lines LogicalResult matchAndRewrite(memref::SubViewOp subView,
} }
if (!srcSubView.hasUnitStride()) { if (!srcSubView.hasUnitStride()) {
return rewriter.notifyMatchFailure(srcSubView, "requires unit strides"); return rewriter.notifyMatchFailure(srcSubView, "requires unit strides");
} }
// Resolve sizes according to dropped dims. // Resolve sizes according to dropped dims.
SmallVector<OpFoldResult> resolvedSizes; SmallVector<OpFoldResult> resolvedSizes;
llvm::SmallBitVector srcDroppedDims = srcSubView.getDroppedDims(); llvm::SmallBitVector srcDroppedDims = srcSubView.getDroppedDims();
resolveSizesIntoOpWithSizes(srcSubView.getMixedSizes(), affine::resolveSizesIntoOpWithSizes(srcSubView.getMixedSizes(),
subView.getMixedSizes(), srcDroppedDims, subView.getMixedSizes(), srcDroppedDims,
resolvedSizes); resolvedSizes);
// Resolve offsets according to source offsets and strides. // Resolve offsets according to source offsets and strides.
SmallVector<Value> resolvedOffsets; SmallVector<Value> resolvedOffsets;
resolveIndicesIntoOpWithOffsetsAndStrides( affine::resolveIndicesIntoOpWithOffsetsAndStrides(
rewriter, subView.getLoc(), srcSubView.getMixedOffsets(), rewriter, subView.getLoc(), srcSubView.getMixedOffsets(),
srcSubView.getMixedStrides(), srcDroppedDims, subView.getMixedOffsets(), srcSubView.getMixedStrides(), srcDroppedDims, subView.getMixedOffsets(),
resolvedOffsets); resolvedOffsets);
// Replace original op. // Replace original op.
rewriter.replaceOpWithNewOp<memref::SubViewOp>( rewriter.replaceOpWithNewOp<memref::SubViewOp>(
subView, subView.getType(), srcSubView.getSource(), subView, subView.getType(), srcSubView.getSource(),
getAsOpFoldResult(resolvedOffsets), resolvedSizes, getAsOpFoldResult(resolvedOffsets), resolvedSizes,
Show All 18 Lines
static SmallVector<Value> static SmallVector<Value>
calculateExpandedAccessIndices(AffineMap affineMap, calculateExpandedAccessIndices(AffineMap affineMap,
const SmallVector<Value> &indices, Location loc, const SmallVector<Value> &indices, Location loc,
PatternRewriter &rewriter) { PatternRewriter &rewriter) {
SmallVector<OpFoldResult> indicesOfr(llvm::to_vector( SmallVector<OpFoldResult> indicesOfr(llvm::to_vector(
llvm::map_range(indices, [](Value v) -> OpFoldResult { return v; }))); llvm::map_range(indices, [](Value v) -> OpFoldResult { return v; })));
SmallVector<Value> expandedIndices; SmallVector<Value> expandedIndices;
for (unsigned i = 0, e = affineMap.getNumResults(); i < e; i++) { for (unsigned i = 0, e = affineMap.getNumResults(); i < e; i++) {
OpFoldResult ofr = makeComposedFoldedAffineApply( OpFoldResult ofr = affine::makeComposedFoldedAffineApply(
rewriter, loc, affineMap.getSubMap({i}), indicesOfr); rewriter, loc, affineMap.getSubMap({i}), indicesOfr);
expandedIndices.push_back( expandedIndices.push_back(
getValueOrCreateConstantIndexOp(rewriter, loc, ofr)); getValueOrCreateConstantIndexOp(rewriter, loc, ofr));
} }
return expandedIndices; return expandedIndices;
} }
template <typename XferOp> template <typename XferOp>
▲ Show 20 Lines • Show All 45 Lines • ▼ Show 20 Lines LogicalResult preconditionResult =
preconditionsFoldSubViewOp(rewriter, loadOp, subViewOp); preconditionsFoldSubViewOp(rewriter, loadOp, subViewOp);
if (failed(preconditionResult)) if (failed(preconditionResult))
return preconditionResult; return preconditionResult;
SmallVector<Value> indices(loadOp.getIndices().begin(), SmallVector<Value> indices(loadOp.getIndices().begin(),
loadOp.getIndices().end()); loadOp.getIndices().end());
// For affine ops, we need to apply the map to get the operands to get the // For affine ops, we need to apply the map to get the operands to get the
// "actual" indices. // "actual" indices.
if (auto affineLoadOp = dyn_cast<AffineLoadOp>(loadOp.getOperation())) { if (auto affineLoadOp =
dyn_cast<affine::AffineLoadOp>(loadOp.getOperation())) {
AffineMap affineMap = affineLoadOp.getAffineMap(); AffineMap affineMap = affineLoadOp.getAffineMap();
auto expandedIndices = calculateExpandedAccessIndices( auto expandedIndices = calculateExpandedAccessIndices(
affineMap, indices, loadOp.getLoc(), rewriter); affineMap, indices, loadOp.getLoc(), rewriter);
indices.assign(expandedIndices.begin(), expandedIndices.end()); indices.assign(expandedIndices.begin(), expandedIndices.end());
} }
SmallVector<Value> sourceIndices; SmallVector<Value> sourceIndices;
resolveIndicesIntoOpWithOffsetsAndStrides( affine::resolveIndicesIntoOpWithOffsetsAndStrides(
rewriter, loadOp.getLoc(), subViewOp.getMixedOffsets(), rewriter, loadOp.getLoc(), subViewOp.getMixedOffsets(),
subViewOp.getMixedStrides(), subViewOp.getDroppedDims(), indices, subViewOp.getMixedStrides(), subViewOp.getDroppedDims(), indices,
sourceIndices); sourceIndices);
llvm::TypeSwitch<Operation *, void>(loadOp) llvm::TypeSwitch<Operation *, void>(loadOp)
.Case([&](AffineLoadOp op) { .Case([&](affine::AffineLoadOp op) {
rewriter.replaceOpWithNewOp<AffineLoadOp>(loadOp, subViewOp.getSource(), rewriter.replaceOpWithNewOp<affine::AffineLoadOp>(
sourceIndices); loadOp, subViewOp.getSource(), sourceIndices);
}) })
.Case([&](memref::LoadOp op) { .Case([&](memref::LoadOp op) {
rewriter.replaceOpWithNewOp<memref::LoadOp>( rewriter.replaceOpWithNewOp<memref::LoadOp>(
loadOp, subViewOp.getSource(), sourceIndices, op.getNontemporal()); loadOp, subViewOp.getSource(), sourceIndices, op.getNontemporal());
}) })
.Case([&](vector::TransferReadOp op) { .Case([&](vector::TransferReadOp op) {
rewriter.replaceOpWithNewOp<vector::TransferReadOp>( rewriter.replaceOpWithNewOp<vector::TransferReadOp>(
op, op.getVectorType(), subViewOp.getSource(), sourceIndices, op, op.getVectorType(), subViewOp.getSource(), sourceIndices,
Show All 19 LinesLogicalResult LoadOpOfExpandShapeOpFolder<OpTy>::matchAndRewrite(
if (!expandShapeOp) if (!expandShapeOp)
return failure(); return failure();
SmallVector<Value> indices(loadOp.getIndices().begin(), SmallVector<Value> indices(loadOp.getIndices().begin(),
loadOp.getIndices().end()); loadOp.getIndices().end());
// For affine ops, we need to apply the map to get the operands to get the // For affine ops, we need to apply the map to get the operands to get the
// "actual" indices. // "actual" indices.
if (auto affineLoadOp = dyn_cast<AffineLoadOp>(loadOp.getOperation())) { if (auto affineLoadOp =
dyn_cast<affine::AffineLoadOp>(loadOp.getOperation())) {
AffineMap affineMap = affineLoadOp.getAffineMap(); AffineMap affineMap = affineLoadOp.getAffineMap();
auto expandedIndices = calculateExpandedAccessIndices( auto expandedIndices = calculateExpandedAccessIndices(
affineMap, indices, loadOp.getLoc(), rewriter); affineMap, indices, loadOp.getLoc(), rewriter);
indices.assign(expandedIndices.begin(), expandedIndices.end()); indices.assign(expandedIndices.begin(), expandedIndices.end());
} }
SmallVector<Value> sourceIndices; SmallVector<Value> sourceIndices;
if (failed(resolveSourceIndicesExpandShape( if (failed(resolveSourceIndicesExpandShape(
loadOp.getLoc(), rewriter, expandShapeOp, indices, sourceIndices))) loadOp.getLoc(), rewriter, expandShapeOp, indices, sourceIndices)))
return failure(); return failure();
llvm::TypeSwitch<Operation *, void>(loadOp) llvm::TypeSwitch<Operation *, void>(loadOp)
.Case<AffineLoadOp, memref::LoadOp>([&](auto op) { .Case<affine::AffineLoadOp, memref::LoadOp>([&](auto op) {
rewriter.replaceOpWithNewOp<decltype(op)>( rewriter.replaceOpWithNewOp<decltype(op)>(
loadOp, expandShapeOp.getViewSource(), sourceIndices); loadOp, expandShapeOp.getViewSource(), sourceIndices);
}) })
.Default([](Operation *) { llvm_unreachable("unexpected operation."); }); .Default([](Operation *) { llvm_unreachable("unexpected operation."); });
return success(); return success();
} }
template <typename OpTy> template <typename OpTy>
LogicalResult LoadOpOfCollapseShapeOpFolder<OpTy>::matchAndRewrite( LogicalResult LoadOpOfCollapseShapeOpFolder<OpTy>::matchAndRewrite(
OpTy loadOp, PatternRewriter &rewriter) const { OpTy loadOp, PatternRewriter &rewriter) const {
auto collapseShapeOp = getMemRefOperand(loadOp) auto collapseShapeOp = getMemRefOperand(loadOp)
.template getDefiningOp<memref::CollapseShapeOp>(); .template getDefiningOp<memref::CollapseShapeOp>();
if (!collapseShapeOp) if (!collapseShapeOp)
return failure(); return failure();
SmallVector<Value> indices(loadOp.getIndices().begin(), SmallVector<Value> indices(loadOp.getIndices().begin(),
loadOp.getIndices().end()); loadOp.getIndices().end());
// For affine ops, we need to apply the map to get the operands to get the // For affine ops, we need to apply the map to get the operands to get the
// "actual" indices. // "actual" indices.
if (auto affineLoadOp = dyn_cast<AffineLoadOp>(loadOp.getOperation())) { if (auto affineLoadOp =
dyn_cast<affine::AffineLoadOp>(loadOp.getOperation())) {
AffineMap affineMap = affineLoadOp.getAffineMap(); AffineMap affineMap = affineLoadOp.getAffineMap();
auto expandedIndices = calculateExpandedAccessIndices( auto expandedIndices = calculateExpandedAccessIndices(
affineMap, indices, loadOp.getLoc(), rewriter); affineMap, indices, loadOp.getLoc(), rewriter);
indices.assign(expandedIndices.begin(), expandedIndices.end()); indices.assign(expandedIndices.begin(), expandedIndices.end());
} }
SmallVector<Value> sourceIndices; SmallVector<Value> sourceIndices;
if (failed(resolveSourceIndicesCollapseShape( if (failed(resolveSourceIndicesCollapseShape(
loadOp.getLoc(), rewriter, collapseShapeOp, indices, sourceIndices))) loadOp.getLoc(), rewriter, collapseShapeOp, indices, sourceIndices)))
return failure(); return failure();
llvm::TypeSwitch<Operation *, void>(loadOp) llvm::TypeSwitch<Operation *, void>(loadOp)
.Case<AffineLoadOp, memref::LoadOp>([&](auto op) { .Case<affine::AffineLoadOp, memref::LoadOp>([&](auto op) {
rewriter.replaceOpWithNewOp<decltype(op)>( rewriter.replaceOpWithNewOp<decltype(op)>(
loadOp, collapseShapeOp.getViewSource(), sourceIndices); loadOp, collapseShapeOp.getViewSource(), sourceIndices);
}) })
.Default([](Operation *) { llvm_unreachable("unexpected operation."); }); .Default([](Operation *) { llvm_unreachable("unexpected operation."); });
return success(); return success();
} }
template <typename OpTy> template <typename OpTy>
Show All 9 Lines LogicalResult preconditionResult =
preconditionsFoldSubViewOp(rewriter, storeOp, subViewOp); preconditionsFoldSubViewOp(rewriter, storeOp, subViewOp);
if (failed(preconditionResult)) if (failed(preconditionResult))
return preconditionResult; return preconditionResult;
SmallVector<Value> indices(storeOp.getIndices().begin(), SmallVector<Value> indices(storeOp.getIndices().begin(),
storeOp.getIndices().end()); storeOp.getIndices().end());
// For affine ops, we need to apply the map to get the operands to get the // For affine ops, we need to apply the map to get the operands to get the
// "actual" indices. // "actual" indices.
if (auto affineStoreOp = dyn_cast<AffineStoreOp>(storeOp.getOperation())) { if (auto affineStoreOp =
dyn_cast<affine::AffineStoreOp>(storeOp.getOperation())) {
AffineMap affineMap = affineStoreOp.getAffineMap(); AffineMap affineMap = affineStoreOp.getAffineMap();
auto expandedIndices = calculateExpandedAccessIndices( auto expandedIndices = calculateExpandedAccessIndices(
affineMap, indices, storeOp.getLoc(), rewriter); affineMap, indices, storeOp.getLoc(), rewriter);
indices.assign(expandedIndices.begin(), expandedIndices.end()); indices.assign(expandedIndices.begin(), expandedIndices.end());
} }
SmallVector<Value> sourceIndices; SmallVector<Value> sourceIndices;
resolveIndicesIntoOpWithOffsetsAndStrides( affine::resolveIndicesIntoOpWithOffsetsAndStrides(
rewriter, storeOp.getLoc(), subViewOp.getMixedOffsets(), rewriter, storeOp.getLoc(), subViewOp.getMixedOffsets(),
subViewOp.getMixedStrides(), subViewOp.getDroppedDims(), indices, subViewOp.getMixedStrides(), subViewOp.getDroppedDims(), indices,
sourceIndices); sourceIndices);
llvm::TypeSwitch<Operation *, void>(storeOp) llvm::TypeSwitch<Operation *, void>(storeOp)
.Case([&](AffineStoreOp op) { .Case([&](affine::AffineStoreOp op) {
rewriter.replaceOpWithNewOp<AffineStoreOp>( rewriter.replaceOpWithNewOp<affine::AffineStoreOp>(
op, op.getValue(), subViewOp.getSource(), sourceIndices); op, op.getValue(), subViewOp.getSource(), sourceIndices);
}) })
.Case([&](memref::StoreOp op) { .Case([&](memref::StoreOp op) {
rewriter.replaceOpWithNewOp<memref::StoreOp>( rewriter.replaceOpWithNewOp<memref::StoreOp>(
op, op.getValue(), subViewOp.getSource(), sourceIndices, op, op.getValue(), subViewOp.getSource(), sourceIndices,
op.getNontemporal()); op.getNontemporal());
}) })
.Case([&](vector::TransferWriteOp op) { .Case([&](vector::TransferWriteOp op) {
Show All 21 LinesLogicalResult StoreOpOfExpandShapeOpFolder<OpTy>::matchAndRewrite(
if (!expandShapeOp) if (!expandShapeOp)
return failure(); return failure();
SmallVector<Value> indices(storeOp.getIndices().begin(), SmallVector<Value> indices(storeOp.getIndices().begin(),
storeOp.getIndices().end()); storeOp.getIndices().end());
// For affine ops, we need to apply the map to get the operands to get the // For affine ops, we need to apply the map to get the operands to get the
// "actual" indices. // "actual" indices.
if (auto affineStoreOp = dyn_cast<AffineStoreOp>(storeOp.getOperation())) { if (auto affineStoreOp =
dyn_cast<affine::AffineStoreOp>(storeOp.getOperation())) {
AffineMap affineMap = affineStoreOp.getAffineMap(); AffineMap affineMap = affineStoreOp.getAffineMap();
auto expandedIndices = calculateExpandedAccessIndices( auto expandedIndices = calculateExpandedAccessIndices(
affineMap, indices, storeOp.getLoc(), rewriter); affineMap, indices, storeOp.getLoc(), rewriter);
indices.assign(expandedIndices.begin(), expandedIndices.end()); indices.assign(expandedIndices.begin(), expandedIndices.end());
} }
SmallVector<Value> sourceIndices; SmallVector<Value> sourceIndices;
if (failed(resolveSourceIndicesExpandShape( if (failed(resolveSourceIndicesExpandShape(
storeOp.getLoc(), rewriter, expandShapeOp, indices, sourceIndices))) storeOp.getLoc(), rewriter, expandShapeOp, indices, sourceIndices)))
return failure(); return failure();
llvm::TypeSwitch<Operation *, void>(storeOp) llvm::TypeSwitch<Operation *, void>(storeOp)
.Case<AffineStoreOp, memref::StoreOp>([&](auto op) { .Case<affine::AffineStoreOp, memref::StoreOp>([&](auto op) {
rewriter.replaceOpWithNewOp<decltype(op)>(storeOp, storeOp.getValue(), rewriter.replaceOpWithNewOp<decltype(op)>(storeOp, storeOp.getValue(),
expandShapeOp.getViewSource(), expandShapeOp.getViewSource(),
sourceIndices); sourceIndices);
}) })
.Default([](Operation *) { llvm_unreachable("unexpected operation."); }); .Default([](Operation *) { llvm_unreachable("unexpected operation."); });
return success(); return success();
} }
template <typename OpTy> template <typename OpTy>
LogicalResult StoreOpOfCollapseShapeOpFolder<OpTy>::matchAndRewrite( LogicalResult StoreOpOfCollapseShapeOpFolder<OpTy>::matchAndRewrite(
OpTy storeOp, PatternRewriter &rewriter) const { OpTy storeOp, PatternRewriter &rewriter) const {
auto collapseShapeOp = getMemRefOperand(storeOp) auto collapseShapeOp = getMemRefOperand(storeOp)
.template getDefiningOp<memref::CollapseShapeOp>(); .template getDefiningOp<memref::CollapseShapeOp>();
if (!collapseShapeOp) if (!collapseShapeOp)
return failure(); return failure();
SmallVector<Value> indices(storeOp.getIndices().begin(), SmallVector<Value> indices(storeOp.getIndices().begin(),
storeOp.getIndices().end()); storeOp.getIndices().end());
// For affine ops, we need to apply the map to get the operands to get the // For affine ops, we need to apply the map to get the operands to get the
// "actual" indices. // "actual" indices.
if (auto affineStoreOp = dyn_cast<AffineStoreOp>(storeOp.getOperation())) { if (auto affineStoreOp =
dyn_cast<affine::AffineStoreOp>(storeOp.getOperation())) {
AffineMap affineMap = affineStoreOp.getAffineMap(); AffineMap affineMap = affineStoreOp.getAffineMap();
auto expandedIndices = calculateExpandedAccessIndices( auto expandedIndices = calculateExpandedAccessIndices(
affineMap, indices, storeOp.getLoc(), rewriter); affineMap, indices, storeOp.getLoc(), rewriter);
indices.assign(expandedIndices.begin(), expandedIndices.end()); indices.assign(expandedIndices.begin(), expandedIndices.end());
} }
SmallVector<Value> sourceIndices; SmallVector<Value> sourceIndices;
if (failed(resolveSourceIndicesCollapseShape( if (failed(resolveSourceIndicesCollapseShape(
storeOp.getLoc(), rewriter, collapseShapeOp, indices, sourceIndices))) storeOp.getLoc(), rewriter, collapseShapeOp, indices, sourceIndices)))
return failure(); return failure();
llvm::TypeSwitch<Operation *, void>(storeOp) llvm::TypeSwitch<Operation *, void>(storeOp)
.Case<AffineStoreOp, memref::StoreOp>([&](auto op) { .Case<affine::AffineStoreOp, memref::StoreOp>([&](auto op) {
rewriter.replaceOpWithNewOp<decltype(op)>( rewriter.replaceOpWithNewOp<decltype(op)>(
storeOp, storeOp.getValue(), collapseShapeOp.getViewSource(), storeOp, storeOp.getValue(), collapseShapeOp.getViewSource(),
sourceIndices); sourceIndices);
}) })
.Default([](Operation *) { llvm_unreachable("unexpected operation."); }); .Default([](Operation *) { llvm_unreachable("unexpected operation."); });
return success(); return success();
} }
Show All 13 LinesLogicalResult NvgpuAsyncCopyOpSubViewOpFolder::matchAndRewrite(
// If the source is a subview, we need to resolve the indices. // If the source is a subview, we need to resolve the indices.
SmallVector<Value> srcindices(copyOp.getSrcIndices().begin(), SmallVector<Value> srcindices(copyOp.getSrcIndices().begin(),
copyOp.getSrcIndices().end()); copyOp.getSrcIndices().end());
SmallVector<Value> foldedSrcIndices(srcindices); SmallVector<Value> foldedSrcIndices(srcindices);
if (srcSubViewOp) { if (srcSubViewOp) {
LLVM_DEBUG(DBGS() << "srcSubViewOp : " << srcSubViewOp << "\n"); LLVM_DEBUG(DBGS() << "srcSubViewOp : " << srcSubViewOp << "\n");
resolveIndicesIntoOpWithOffsetsAndStrides( affine::resolveIndicesIntoOpWithOffsetsAndStrides(
rewriter, copyOp.getLoc(), srcSubViewOp.getMixedOffsets(), rewriter, copyOp.getLoc(), srcSubViewOp.getMixedOffsets(),
srcSubViewOp.getMixedStrides(), srcSubViewOp.getDroppedDims(), srcSubViewOp.getMixedStrides(), srcSubViewOp.getDroppedDims(),
srcindices, foldedSrcIndices); srcindices, foldedSrcIndices);
} }
// If the destination is a subview, we need to resolve the indices. // If the destination is a subview, we need to resolve the indices.
SmallVector<Value> dstindices(copyOp.getDstIndices().begin(), SmallVector<Value> dstindices(copyOp.getDstIndices().begin(),
copyOp.getDstIndices().end()); copyOp.getDstIndices().end());
SmallVector<Value> foldedDstIndices(dstindices); SmallVector<Value> foldedDstIndices(dstindices);
if (dstSubViewOp) { if (dstSubViewOp) {
LLVM_DEBUG(DBGS() << "dstSubViewOp : " << dstSubViewOp << "\n"); LLVM_DEBUG(DBGS() << "dstSubViewOp : " << dstSubViewOp << "\n");
resolveIndicesIntoOpWithOffsetsAndStrides( affine::resolveIndicesIntoOpWithOffsetsAndStrides(
rewriter, copyOp.getLoc(), dstSubViewOp.getMixedOffsets(), rewriter, copyOp.getLoc(), dstSubViewOp.getMixedOffsets(),
dstSubViewOp.getMixedStrides(), dstSubViewOp.getDroppedDims(), dstSubViewOp.getMixedStrides(), dstSubViewOp.getDroppedDims(),
dstindices, foldedDstIndices); dstindices, foldedDstIndices);
} }
// Replace the copy op with a new copy op that uses the source and destination // Replace the copy op with a new copy op that uses the source and destination
// of the subview. // of the subview.
rewriter.replaceOpWithNewOp<nvgpu::DeviceAsyncCopyOp>( rewriter.replaceOpWithNewOp<nvgpu::DeviceAsyncCopyOp>(
copyOp, nvgpu::DeviceAsyncTokenType::get(copyOp.getContext()), copyOp, nvgpu::DeviceAsyncTokenType::get(copyOp.getContext()),
(dstSubViewOp ? dstSubViewOp.getSource() : copyOp.getDst()), (dstSubViewOp ? dstSubViewOp.getSource() : copyOp.getDst()),
foldedDstIndices, foldedDstIndices,
(srcSubViewOp ? srcSubViewOp.getSource() : copyOp.getSrc()), (srcSubViewOp ? srcSubViewOp.getSource() : copyOp.getSrc()),
foldedSrcIndices, copyOp.getDstElements(), copyOp.getSrcElements(), foldedSrcIndices, copyOp.getDstElements(), copyOp.getSrcElements(),
copyOp.getBypassL1Attr()); copyOp.getBypassL1Attr());
return success(); return success();
} }
void memref::populateFoldMemRefAliasOpPatterns(RewritePatternSet &patterns) { void memref::populateFoldMemRefAliasOpPatterns(RewritePatternSet &patterns) {
patterns.add<LoadOpOfSubViewOpFolder<AffineLoadOp>, patterns.add<LoadOpOfSubViewOpFolder<affine::AffineLoadOp>,
LoadOpOfSubViewOpFolder<memref::LoadOp>, LoadOpOfSubViewOpFolder<memref::LoadOp>,
LoadOpOfSubViewOpFolder<vector::TransferReadOp>, LoadOpOfSubViewOpFolder<vector::TransferReadOp>,
LoadOpOfSubViewOpFolder<gpu::SubgroupMmaLoadMatrixOp>, LoadOpOfSubViewOpFolder<gpu::SubgroupMmaLoadMatrixOp>,
StoreOpOfSubViewOpFolder<AffineStoreOp>, StoreOpOfSubViewOpFolder<affine::AffineStoreOp>,
StoreOpOfSubViewOpFolder<memref::StoreOp>, StoreOpOfSubViewOpFolder<memref::StoreOp>,
StoreOpOfSubViewOpFolder<vector::TransferWriteOp>, StoreOpOfSubViewOpFolder<vector::TransferWriteOp>,
StoreOpOfSubViewOpFolder<gpu::SubgroupMmaStoreMatrixOp>, StoreOpOfSubViewOpFolder<gpu::SubgroupMmaStoreMatrixOp>,
LoadOpOfExpandShapeOpFolder<AffineLoadOp>, LoadOpOfExpandShapeOpFolder<affine::AffineLoadOp>,
LoadOpOfExpandShapeOpFolder<memref::LoadOp>, LoadOpOfExpandShapeOpFolder<memref::LoadOp>,
StoreOpOfExpandShapeOpFolder<AffineStoreOp>, StoreOpOfExpandShapeOpFolder<affine::AffineStoreOp>,
StoreOpOfExpandShapeOpFolder<memref::StoreOp>, StoreOpOfExpandShapeOpFolder<memref::StoreOp>,
LoadOpOfCollapseShapeOpFolder<AffineLoadOp>, LoadOpOfCollapseShapeOpFolder<affine::AffineLoadOp>,
LoadOpOfCollapseShapeOpFolder<memref::LoadOp>, LoadOpOfCollapseShapeOpFolder<memref::LoadOp>,
StoreOpOfCollapseShapeOpFolder<AffineStoreOp>, StoreOpOfCollapseShapeOpFolder<affine::AffineStoreOp>,
StoreOpOfCollapseShapeOpFolder<memref::StoreOp>, StoreOpOfCollapseShapeOpFolder<memref::StoreOp>,
SubViewOfSubViewFolder, NvgpuAsyncCopyOpSubViewOpFolder>( SubViewOfSubViewFolder, NvgpuAsyncCopyOpSubViewOpFolder>(
patterns.getContext()); patterns.getContext());
} }
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// Pass registration // Pass registration
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
Show All 19 Lines

mlir/lib/Dialect/MemRef/Transforms/MultiBuffer.cpp

Show First 20 Lines • Show All 184 Lines • ▼ Show 20 Linesmlir::memref::multiBuffer(RewriterBase &rewriter, memref::AllocOp allocOp,
// 3. Within the loop, build the modular leading index (i.e. each loop // 3. Within the loop, build the modular leading index (i.e. each loop
// iteration %iv accesses slice ((%iv - %lb) / %step) % %mb_factor). // iteration %iv accesses slice ((%iv - %lb) / %step) % %mb_factor).
rewriter.setInsertionPointToStart(&candidateLoop.getLoopBody().front()); rewriter.setInsertionPointToStart(&candidateLoop.getLoopBody().front());
Value ivVal = *inductionVar; Value ivVal = *inductionVar;
Value lbVal = getValueOrCreateConstantIndexOp(rewriter, loc, *lowerBound); Value lbVal = getValueOrCreateConstantIndexOp(rewriter, loc, *lowerBound);
Value stepVal = getValueOrCreateConstantIndexOp(rewriter, loc, *singleStep); Value stepVal = getValueOrCreateConstantIndexOp(rewriter, loc, *singleStep);
AffineExpr iv, lb, step; AffineExpr iv, lb, step;
bindDims(rewriter.getContext(), iv, lb, step); bindDims(rewriter.getContext(), iv, lb, step);
Value bufferIndex = makeComposedAffineApply( Value bufferIndex = affine::makeComposedAffineApply(
rewriter, loc, ((iv - lb).floorDiv(step)) % multiBufferingFactor, rewriter, loc, ((iv - lb).floorDiv(step)) % multiBufferingFactor,
{ivVal, lbVal, stepVal}); {ivVal, lbVal, stepVal});
LLVM_DEBUG(DBGS() << "--multi-buffered indexing: " << bufferIndex << "\n"); LLVM_DEBUG(DBGS() << "--multi-buffered indexing: " << bufferIndex << "\n");
// 4. Build the subview accessing the particular slice, taking modular // 4. Build the subview accessing the particular slice, taking modular
// rotation into account. // rotation into account.
int64_t mbMemRefTypeRank = mbMemRefType.getRank(); int64_t mbMemRefTypeRank = mbMemRefType.getRank();
IntegerAttr zero = rewriter.getIndexAttr(0); IntegerAttr zero = rewriter.getIndexAttr(0);
▲ Show 20 Lines • Show All 49 Lines • Show Last 20 Lines

mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp

Show All 23 Lines
#define GEN_PASS_DEF_NORMALIZEMEMREFS #define GEN_PASS_DEF_NORMALIZEMEMREFS
#include "mlir/Dialect/MemRef/Transforms/Passes.h.inc" #include "mlir/Dialect/MemRef/Transforms/Passes.h.inc"
} // namespace memref } // namespace memref
} // namespace mlir } // namespace mlir
#define DEBUG_TYPE "normalize-memrefs" #define DEBUG_TYPE "normalize-memrefs"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
namespace { namespace {
/// All memrefs passed across functions with non-trivial layout maps are /// All memrefs passed across functions with non-trivial layout maps are
/// converted to ones with trivial identity layout ones. /// converted to ones with trivial identity layout ones.
/// If all the memref types/uses in a function are normalizable, we treat /// If all the memref types/uses in a function are normalizable, we treat
/// such functions as normalizable. Also, if a normalizable function is known /// such functions as normalizable. Also, if a normalizable function is known
/// to call a non-normalizable function, we treat that function as /// to call a non-normalizable function, we treat that function as
▲ Show 20 Lines • Show All 500 Lines • Show Last 20 Lines

mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp

Show All 14 Lines
#include "mlir/Dialect/SCF/Transforms/Transforms.h" #include "mlir/Dialect/SCF/Transforms/Transforms.h"
#include "mlir/Dialect/SCF/Utils/Utils.h" #include "mlir/Dialect/SCF/Utils/Utils.h"
#include "mlir/Dialect/Transform/IR/TransformDialect.h" #include "mlir/Dialect/Transform/IR/TransformDialect.h"
#include "mlir/Dialect/Transform/IR/TransformInterfaces.h" #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
#include "mlir/Dialect/Transform/IR/TransformOps.h" #include "mlir/Dialect/Transform/IR/TransformOps.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h" #include "mlir/Dialect/Vector/IR/VectorOps.h"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// GetParentForOp // GetParentForOp
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
DiagnosedSilenceableFailure DiagnosedSilenceableFailure
transform::GetParentForOp::apply(transform::TransformResults &results, transform::GetParentForOp::apply(transform::TransformResults &results,
transform::TransformState &state) { transform::TransformState &state) {
SetVector<Operation *> parents; SetVector<Operation *> parents;
▲ Show 20 Lines • Show All 262 Lines • ▼ Show 20 Lines
namespace { namespace {
class SCFTransformDialectExtension class SCFTransformDialectExtension
: public transform::TransformDialectExtension< : public transform::TransformDialectExtension<
SCFTransformDialectExtension> { SCFTransformDialectExtension> {
public: public:
using Base::Base; using Base::Base;
void init() { void init() {
declareGeneratedDialect<AffineDialect>(); declareGeneratedDialect<affine::AffineDialect>();
declareGeneratedDialect<func::FuncDialect>(); declareGeneratedDialect<func::FuncDialect>();
registerTransformOps< registerTransformOps<
#define GET_OP_LIST #define GET_OP_LIST
#include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.cpp.inc" #include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.cpp.inc"
>(); >();
} }
}; };
} // namespace } // namespace
#define GET_OP_CLASSES #define GET_OP_CLASSES
#include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.cpp.inc" #include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.cpp.inc"
void mlir::scf::registerTransformDialectExtension(DialectRegistry &registry) { void mlir::scf::registerTransformDialectExtension(DialectRegistry &registry) {
registry.addExtensions<SCFTransformDialectExtension>(); registry.addExtensions<SCFTransformDialectExtension>();
} }

mlir/lib/Dialect/SCF/Transforms/LoopCanonicalization.cpp

Show First 20 Lines • Show All 173 Lines • ▼ Show 20 Linesstruct SCFForLoopCanonicalization
} }
}; };
} // namespace } // namespace
void mlir::scf::populateSCFForLoopCanonicalizationPatterns( void mlir::scf::populateSCFForLoopCanonicalizationPatterns(
RewritePatternSet &patterns) { RewritePatternSet &patterns) {
MLIRContext *ctx = patterns.getContext(); MLIRContext *ctx = patterns.getContext();
patterns patterns
.add<AffineOpSCFCanonicalizationPattern<AffineMinOp>, .add<AffineOpSCFCanonicalizationPattern<affine::AffineMinOp>,
AffineOpSCFCanonicalizationPattern<AffineMaxOp>, AffineOpSCFCanonicalizationPattern<affine::AffineMaxOp>,
DimOfIterArgFolder<tensor::DimOp>, DimOfIterArgFolder<memref::DimOp>, DimOfIterArgFolder<tensor::DimOp>, DimOfIterArgFolder<memref::DimOp>,
DimOfLoopResultFolder<tensor::DimOp>, DimOfLoopResultFolder<tensor::DimOp>,
DimOfLoopResultFolder<memref::DimOp>>(ctx); DimOfLoopResultFolder<memref::DimOp>>(ctx);
} }
std::unique_ptr<Pass> mlir::createSCFForLoopCanonicalizationPass() { std::unique_ptr<Pass> mlir::createSCFForLoopCanonicalizationPass() {
return std::make_unique<SCFForLoopCanonicalization>(); return std::make_unique<SCFForLoopCanonicalization>();
} }

mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp

Show All 28 Lines
namespace mlir { namespace mlir {
#define GEN_PASS_DEF_SCFFORLOOPPEELING #define GEN_PASS_DEF_SCFFORLOOPPEELING
#define GEN_PASS_DEF_SCFFORLOOPSPECIALIZATION #define GEN_PASS_DEF_SCFFORLOOPSPECIALIZATION
#define GEN_PASS_DEF_SCFPARALLELLOOPSPECIALIZATION #define GEN_PASS_DEF_SCFPARALLELLOOPSPECIALIZATION
#include "mlir/Dialect/SCF/Transforms/Passes.h.inc" #include "mlir/Dialect/SCF/Transforms/Passes.h.inc"
} // namespace mlir } // namespace mlir
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
using scf::ForOp; using scf::ForOp;
using scf::ParallelOp; using scf::ParallelOp;
/// Rewrite a parallel loop with bounds defined by an affine.min with a constant /// Rewrite a parallel loop with bounds defined by an affine.min with a constant
/// into 2 loops after checking if the bounds are equal to that constant. This /// into 2 loops after checking if the bounds are equal to that constant. This
/// is beneficial if the loop will almost always have the constant bound and /// is beneficial if the loop will almost always have the constant bound and
/// that version can be fully unrolled and vectorized. /// that version can be fully unrolled and vectorized.
static void specializeParallelLoopForUnrolling(ParallelOp op) { static void specializeParallelLoopForUnrolling(ParallelOp op) {
▲ Show 20 Lines • Show All 243 Lines • Show Last 20 Lines

mlir/lib/Dialect/SCF/Transforms/ParallelLoopTiling.cpp

Show First 20 Lines • Show All 124 Lines • ▼ Show 20 Lines if (noMinMaxBounds) {
newBounds.push_back(newStep); newBounds.push_back(newStep);
needInboundCheck = true; needInboundCheck = true;
continue; continue;
} }
// Otherwise, we dynamically compute the bound for // Otherwise, we dynamically compute the bound for
// each iteration of the outer loop. // each iteration of the outer loop.
newBounds.push_back( newBounds.push_back(
b.create<AffineMinOp>(op.getLoc(), b.getIndexType(), minMap, b.create<affine::AffineMinOp>(op.getLoc(), b.getIndexType(), minMap,
ValueRange{newStep, upperBound, iv})); ValueRange{newStep, upperBound, iv}));
} }
auto innerLoop = b.create<ParallelOp>( auto innerLoop = b.create<ParallelOp>(
op.getLoc(), SmallVector<Value, 2>(newBounds.size(), zero), newBounds, op.getLoc(), SmallVector<Value, 2>(newBounds.size(), zero), newBounds,
op.getStep()); op.getStep());
if (noMinMaxBounds && needInboundCheck) { if (noMinMaxBounds && needInboundCheck) {
b.setInsertionPointToStart(innerLoop.getBody()); b.setInsertionPointToStart(innerLoop.getBody());
// Insert in-bound check // Insert in-bound check
▲ Show 20 Lines • Show All 72 Lines • Show Last 20 Lines

mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

Show First 20 Lines • Show All 96 Lines • ▼ Show 20 Linesstatic OpFoldResult getBoundedTileSize(OpBuilder &b, Location loc,
// The tile size to use (to avoid out of bounds access) is minimum of // The tile size to use (to avoid out of bounds access) is minimum of
// `tileSize` and `ub - iv`, where `iv` is the induction variable of the tiled // `tileSize` and `ub - iv`, where `iv` is the induction variable of the tiled
// loop. // loop.
AffineExpr s0, s1, d0; AffineExpr s0, s1, d0;
bindDims(b.getContext(), d0); bindDims(b.getContext(), d0);
bindSymbols(b.getContext(), s0, s1); bindSymbols(b.getContext(), s0, s1);
AffineMap minMap = AffineMap::get(1, 2, {s0, s1 - d0}, b.getContext()); AffineMap minMap = AffineMap::get(1, 2, {s0, s1 - d0}, b.getContext());
Value size = getValueOrCreateConstantIndexOp(b, loc, loopRange.size); Value size = getValueOrCreateConstantIndexOp(b, loc, loopRange.size);
return makeComposedFoldedAffineMin( return affine::makeComposedFoldedAffineMin(
b, loc, minMap, SmallVector<OpFoldResult>{iv, tileSize, size}); b, loc, minMap, SmallVector<OpFoldResult>{iv, tileSize, size});
} }
/// Generate an empty loop nest that represents the tiled loop nest shell. /// Generate an empty loop nest that represents the tiled loop nest shell.
/// - `loopRanges` specifies the lb, ub and step of the untiled iteration space. /// - `loopRanges` specifies the lb, ub and step of the untiled iteration space.
/// - `tileSizeVals` is the tile sizes to use. Zero represent untiled loops. /// - `tileSizeVals` is the tile sizes to use. Zero represent untiled loops.
/// - In `offsets` and `sizes` return the multi-dimensional offset and size of /// - In `offsets` and `sizes` return the multi-dimensional offset and size of
/// the /// the
▲ Show 20 Lines • Show All 640 Lines • Show Last 20 Lines

mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp

Show All 21 Lines
#include "mlir/IR/AffineMap.h" #include "mlir/IR/AffineMap.h"
#include "mlir/IR/Matchers.h" #include "mlir/IR/Matchers.h"
#include "mlir/IR/PatternMatch.h" #include "mlir/IR/PatternMatch.h"
#include "llvm/Support/Debug.h" #include "llvm/Support/Debug.h"
#define DEBUG_TYPE "mlir-scf-affine-utils" #define DEBUG_TYPE "mlir-scf-affine-utils"
using namespace mlir; using namespace mlir;
using namespace affine;
using namespace presburger; using namespace presburger;
LogicalResult scf::matchForLikeLoop(Value iv, OpFoldResult &lb, LogicalResult scf::matchForLikeLoop(Value iv, OpFoldResult &lb,
OpFoldResult &ub, OpFoldResult &step) { OpFoldResult &ub, OpFoldResult &step) {
if (scf::ForOp forOp = scf::getForInductionVarOwner(iv)) { if (scf::ForOp forOp = scf::getForInductionVarOwner(iv)) {
lb = forOp.getLowerBound(); lb = forOp.getLowerBound();
ub = forOp.getUpperBound(); ub = forOp.getUpperBound();
step = forOp.getStep(); step = forOp.getStep();
Show All 25 Lines
} }
static FailureOr<AffineApplyOp> static FailureOr<AffineApplyOp>
canonicalizeMinMaxOp(RewriterBase &rewriter, Operation *op, canonicalizeMinMaxOp(RewriterBase &rewriter, Operation *op,
FlatAffineValueConstraints constraints) { FlatAffineValueConstraints constraints) {
RewriterBase::InsertionGuard guard(rewriter); RewriterBase::InsertionGuard guard(rewriter);
rewriter.setInsertionPoint(op); rewriter.setInsertionPoint(op);
FailureOr<AffineValueMap> simplified = FailureOr<AffineValueMap> simplified =
mlir::simplifyConstrainedMinMaxOp(op, std::move(constraints)); affine::simplifyConstrainedMinMaxOp(op, std::move(constraints));
if (failed(simplified)) if (failed(simplified))
return failure(); return failure();
return rewriter.replaceOpWithNewOp<AffineApplyOp>( return rewriter.replaceOpWithNewOp<AffineApplyOp>(
op, simplified->getAffineMap(), simplified->getOperands()); op, simplified->getAffineMap(), simplified->getOperands());
} }
LogicalResult scf::addLoopRangeConstraints(FlatAffineValueConstraints &cstr, LogicalResult scf::addLoopRangeConstraints(FlatAffineValueConstraints &cstr,
Value iv, OpFoldResult lb, Value iv, OpFoldResult lb,
▲ Show 20 Lines • Show All 147 Lines • Show Last 20 Lines

mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp

Show First 20 Lines • Show All 85 Lines • ▼ Show 20 Linesstatic Value genVectorMask(PatternRewriter &rewriter, Location loc, VL vl,
// during vector execution. Here we rely on subsequent loop optimizations to // during vector execution. Here we rely on subsequent loop optimizations to
// avoid executing the mask in all iterations, for example, by splitting the // avoid executing the mask in all iterations, for example, by splitting the
// loop into an unconditional vector loop and a scalar cleanup loop. // loop into an unconditional vector loop and a scalar cleanup loop.
auto min = AffineMap::get( auto min = AffineMap::get(
/*dimCount=*/2, /*symbolCount=*/1, /*dimCount=*/2, /*symbolCount=*/1,
{rewriter.getAffineSymbolExpr(0), {rewriter.getAffineSymbolExpr(0),
rewriter.getAffineDimExpr(0) - rewriter.getAffineDimExpr(1)}, rewriter.getAffineDimExpr(0) - rewriter.getAffineDimExpr(1)},
rewriter.getContext()); rewriter.getContext());
Value end = Value end = rewriter.createOrFold<affine::AffineMinOp>(
rewriter.createOrFold<AffineMinOp>(loc, min, ValueRange{hi, iv, step}); loc, min, ValueRange{hi, iv, step});
return rewriter.create<vector::CreateMaskOp>(loc, mtp, end); return rewriter.create<vector::CreateMaskOp>(loc, mtp, end);
} }
/// Generates a vectorized invariant. Here we rely on subsequent loop /// Generates a vectorized invariant. Here we rely on subsequent loop
/// optimizations to hoist the invariant broadcast out of the vector loop. /// optimizations to hoist the invariant broadcast out of the vector loop.
static Value genVectorInvariantValue(PatternRewriter &rewriter, VL vl, static Value genVectorInvariantValue(PatternRewriter &rewriter, VL vl,
Value val) { Value val) {
VectorType vtp = vectorType(vl, val.getType()); VectorType vtp = vectorType(vl, val.getType());
▲ Show 20 Lines • Show All 567 Lines • Show Last 20 Lines

mlir/lib/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.cpp

Show First 20 Lines • Show All 49 Lines • ▼ Show 20 Linesstatic OpFoldResult getCollapsedOutputDimFromInputShape(
unsigned endPos = map.getResults().back().cast<AffineDimExpr>().getPosition(); unsigned endPos = map.getResults().back().cast<AffineDimExpr>().getPosition();
AffineExpr expr; AffineExpr expr;
SmallVector<Value, 2> dynamicDims; SmallVector<Value, 2> dynamicDims;
for (auto dim : llvm::seq_inclusive(startPos, endPos)) { for (auto dim : llvm::seq_inclusive(startPos, endPos)) {
dynamicDims.push_back(builder.createOrFold<tensor::DimOp>(loc, src, dim)); dynamicDims.push_back(builder.createOrFold<tensor::DimOp>(loc, src, dim));
AffineExpr currExpr = builder.getAffineSymbolExpr(dim - startPos); AffineExpr currExpr = builder.getAffineSymbolExpr(dim - startPos);
expr = (expr ? expr * currExpr : currExpr); expr = (expr ? expr * currExpr : currExpr);
} }
return applyMapToValues(builder, loc, return affine::applyMapToValues(
AffineMap::get(0, endPos - startPos + 1, expr), builder, loc, AffineMap::get(0, endPos - startPos + 1, expr),
dynamicDims)[0]; dynamicDims)[0];
} }
/// Given the `src` of a collapsing reshape op and its reassociation maps, /// Given the `src` of a collapsing reshape op and its reassociation maps,
/// compute the shape of the result of the reshape. /// compute the shape of the result of the reshape.
static SmallVector<OpFoldResult, 4> getCollapsedOutputShapeFromInputShape( static SmallVector<OpFoldResult, 4> getCollapsedOutputShapeFromInputShape(
OpBuilder &builder, Location loc, Value src, OpBuilder &builder, Location loc, Value src,
ArrayRef<int64_t> dstStaticShape, ArrayRef<AffineMap> reassociation) { ArrayRef<int64_t> dstStaticShape, ArrayRef<AffineMap> reassociation) {
return llvm::to_vector<4>(llvm::map_range( return llvm::to_vector<4>(llvm::map_range(
Show All 29 Lines for (auto d :
if (d.index() + startPos == static_cast<unsigned>(dimIndex)) if (d.index() + startPos == static_cast<unsigned>(dimIndex))
continue; continue;
assert(!ShapedType::isDynamic(d.value()) && assert(!ShapedType::isDynamic(d.value()) &&
"single dimension cannot be expanded into multiple dynamic " "single dimension cannot be expanded into multiple dynamic "
"dimensions"); "dimensions");
linearizedStaticDim *= d.value(); linearizedStaticDim *= d.value();
} }
Value sourceDim = builder.create<tensor::DimOp>(loc, src, sourceDimPos); Value sourceDim = builder.create<tensor::DimOp>(loc, src, sourceDimPos);
return applyMapToValues( return affine::applyMapToValues(
builder, loc, builder, loc,
AffineMap::get( AffineMap::get(
0, 1, builder.getAffineSymbolExpr(0).floorDiv(linearizedStaticDim)), 0, 1, builder.getAffineSymbolExpr(0).floorDiv(linearizedStaticDim)),
sourceDim)[0]; sourceDim)[0];
} }
/// Given the `src` of an expanding reshape op, the reassociation maps and the /// Given the `src` of an expanding reshape op, the reassociation maps and the
/// result type, compute the shape of the result of the reshape. /// result type, compute the shape of the result of the reshape.
▲ Show 20 Lines • Show All 70 Lines • ▼ Show 20 Lines for (auto dim : llvm::seq<int64_t>(0, padOp.getSourceType().getRank())) {
return; return;
} }
int64_t staticValue = int64_t staticValue =
valueOrAttr.get<Attribute>().cast<IntegerAttr>().getInt(); valueOrAttr.get<Attribute>().cast<IntegerAttr>().getInt();
expr = expr + staticValue; expr = expr + staticValue;
}; };
addOpFoldResult(lowPad[dim]); addOpFoldResult(lowPad[dim]);
addOpFoldResult(highPad[dim]); addOpFoldResult(highPad[dim]);
shapes.push_back(applyMapToValues( shapes.push_back(affine::applyMapToValues(
b, loc, AffineMap::get(1, numSymbols, expr), mapOperands)[0]); b, loc, AffineMap::get(1, numSymbols, expr), mapOperands)[0]);
} }
reifiedReturnShapes.emplace_back(std::move(shapes)); reifiedReturnShapes.emplace_back(std::move(shapes));
return success(); return success();
} }
}; };
} // namespace } // namespace
Show All 11 Lines

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

Show First 20 Lines • Show All 1,522 Lines • ▼ Show 20 Lines for (int64_t d : grp) {
} }
} }
// result dim size = src dim size / (product(other dims in reassoc group)) // result dim size = src dim size / (product(other dims in reassoc group))
Value srcDimSz = Value srcDimSz =
rewriter.create<DimOp>(dimOp.getLoc(), expandShapeOp.getSrc(), srcDim); rewriter.create<DimOp>(dimOp.getLoc(), expandShapeOp.getSrc(), srcDim);
AffineExpr expr; AffineExpr expr;
bindSymbols(dimOp.getContext(), expr); bindSymbols(dimOp.getContext(), expr);
rewriter.replaceOpWithNewOp<AffineApplyOp>(dimOp, expr.floorDiv(product), rewriter.replaceOpWithNewOp<affine::AffineApplyOp>(
srcDimSz); dimOp, expr.floorDiv(product), srcDimSz);
return success(); return success();
} }
}; };
struct FoldDimOfCollapseShape : public OpRewritePattern<DimOp> { struct FoldDimOfCollapseShape : public OpRewritePattern<DimOp> {
using OpRewritePattern<DimOp>::OpRewritePattern; using OpRewritePattern<DimOp>::OpRewritePattern;
LogicalResult matchAndRewrite(DimOp dimOp, LogicalResult matchAndRewrite(DimOp dimOp,
Show All 21 Lines LogicalResult matchAndRewrite(DimOp dimOp,
SmallVector<AffineExpr> syms; SmallVector<AffineExpr> syms;
AffineExpr product; AffineExpr product;
for (const auto &it : llvm::enumerate(group)) { for (const auto &it : llvm::enumerate(group)) {
srcDimSizes.push_back(rewriter.create<DimOp>( srcDimSizes.push_back(rewriter.create<DimOp>(
dimOp.getLoc(), collapseShapeOp.getSrc(), it.value())); dimOp.getLoc(), collapseShapeOp.getSrc(), it.value()));
syms.push_back(rewriter.getAffineSymbolExpr(it.index())); syms.push_back(rewriter.getAffineSymbolExpr(it.index()));
product = product ? product * syms.back() : syms.back(); product = product ? product * syms.back() : syms.back();
} }
rewriter.replaceOpWithNewOp<AffineApplyOp>(dimOp, product, srcDimSizes); rewriter.replaceOpWithNewOp<affine::AffineApplyOp>(dimOp, product,
srcDimSizes);
return success(); return success();
} }
}; };
} // namespace } // namespace
void ExpandShapeOp::getCanonicalizationPatterns(RewritePatternSet &results, void ExpandShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) { MLIRContext *context) {
results.add<ComposeReassociativeReshapeOps<ExpandShapeOp>, results.add<ComposeReassociativeReshapeOps<ExpandShapeOp>,
▲ Show 20 Lines • Show All 1,981 Lines • ▼ Show 20 LinesSmallVector<OpFoldResult> PackOp::getResultShape(
ArrayRef<OpFoldResult> innerTileSizes, ArrayRef<int64_t> innerDimsPos, ArrayRef<OpFoldResult> innerTileSizes, ArrayRef<int64_t> innerDimsPos,
ArrayRef<int64_t> outerDimsPerm) { ArrayRef<int64_t> outerDimsPerm) {
SmallVector<OpFoldResult> resultDims = llvm::to_vector(sourceDims); SmallVector<OpFoldResult> resultDims = llvm::to_vector(sourceDims);
AffineExpr s0, s1; AffineExpr s0, s1;
bindSymbols(builder.getContext(), s0, s1); bindSymbols(builder.getContext(), s0, s1);
AffineExpr ceilDivExpr = s0.ceilDiv(s1); AffineExpr ceilDivExpr = s0.ceilDiv(s1);
for (auto tiledDim : llvm::enumerate(innerDimsPos)) { for (auto tiledDim : llvm::enumerate(innerDimsPos)) {
resultDims[tiledDim.value()] = makeComposedFoldedAffineApply( resultDims[tiledDim.value()] = affine::makeComposedFoldedAffineApply(
builder, loc, ceilDivExpr, builder, loc, ceilDivExpr,
{resultDims[tiledDim.value()], innerTileSizes[tiledDim.index()]}); {resultDims[tiledDim.value()], innerTileSizes[tiledDim.index()]});
} }
if (!outerDimsPerm.empty()) if (!outerDimsPerm.empty())
applyPermutationToVector(resultDims, outerDimsPerm); applyPermutationToVector(resultDims, outerDimsPerm);
resultDims.append(innerTileSizes.begin(), innerTileSizes.end()); resultDims.append(innerTileSizes.begin(), innerTileSizes.end());
SmallVector<int64_t> resultTypeShape = SmallVector<int64_t> resultTypeShape =
Show All 28 Lines
Value PackOp::createDestinationTensor(OpBuilder &b, Location loc, Value source, Value PackOp::createDestinationTensor(OpBuilder &b, Location loc, Value source,
ArrayRef<OpFoldResult> innerTileSizes, ArrayRef<OpFoldResult> innerTileSizes,
ArrayRef<int64_t> innerDimsPos, ArrayRef<int64_t> innerDimsPos,
ArrayRef<int64_t> outerDimsPerm) { ArrayRef<int64_t> outerDimsPerm) {
AffineExpr dim0, dim1; AffineExpr dim0, dim1;
bindDims(b.getContext(), dim0, dim1); bindDims(b.getContext(), dim0, dim1);
auto ceilDiv = [&](OpFoldResult v1, OpFoldResult v2) -> OpFoldResult { auto ceilDiv = [&](OpFoldResult v1, OpFoldResult v2) -> OpFoldResult {
return makeComposedFoldedAffineApply(b, loc, dim0.ceilDiv(dim1), {v1, v2}); return affine::makeComposedFoldedAffineApply(b, loc, dim0.ceilDiv(dim1),
{v1, v2});
}; };
SmallVector<OpFoldResult> mixedSizes; SmallVector<OpFoldResult> mixedSizes;
for (auto [index, value] : for (auto [index, value] :
llvm::enumerate(source.getType().cast<RankedTensorType>().getShape())) { llvm::enumerate(source.getType().cast<RankedTensorType>().getShape())) {
if (ShapedType::isDynamic(value)) if (ShapedType::isDynamic(value))
mixedSizes.push_back(b.create<DimOp>(loc, source, index).getResult()); mixedSizes.push_back(b.create<DimOp>(loc, source, index).getResult());
else else
▲ Show 20 Lines • Show All 189 Lines • ▼ Show 20 Lines
Value UnPackOp::createDestinationTensor(OpBuilder &b, Location loc, Value UnPackOp::createDestinationTensor(OpBuilder &b, Location loc,
Value source, Value source,
ArrayRef<OpFoldResult> innerTileSizes, ArrayRef<OpFoldResult> innerTileSizes,
ArrayRef<int64_t> innerDimsPos, ArrayRef<int64_t> innerDimsPos,
ArrayRef<int64_t> outerDimsPerm) { ArrayRef<int64_t> outerDimsPerm) {
AffineExpr sym0, sym1; AffineExpr sym0, sym1;
bindSymbols(b.getContext(), sym0, sym1); bindSymbols(b.getContext(), sym0, sym1);
auto dimMul = [&](OpFoldResult v1, OpFoldResult v2) -> OpFoldResult { auto dimMul = [&](OpFoldResult v1, OpFoldResult v2) -> OpFoldResult {
return makeComposedFoldedAffineApply(b, loc, sym0 * sym1, {v1, v2}); return affine::makeComposedFoldedAffineApply(b, loc, sym0 * sym1, {v1, v2});
}; };
SmallVector<OpFoldResult> mixedSizes; SmallVector<OpFoldResult> mixedSizes;
auto srcType = source.getType().cast<RankedTensorType>(); auto srcType = source.getType().cast<RankedTensorType>();
for (auto i : for (auto i :
llvm::seq<unsigned>(0, srcType.getRank() - innerTileSizes.size())) { llvm::seq<unsigned>(0, srcType.getRank() - innerTileSizes.size())) {
if (srcType.isDynamicDim(i)) if (srcType.isDynamicDim(i))
mixedSizes.push_back(b.create<DimOp>(loc, source, i).getResult()); mixedSizes.push_back(b.create<DimOp>(loc, source, i).getResult());
▲ Show 20 Lines • Show All 133 Lines • Show Last 20 Lines

mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp

Show First 20 Lines • Show All 133 Lines • ▼ Show 20 Lines applyPermToRange(origOffsets, origSizes,
invertPermutationVector(packOp.getOuterDimsPerm())); invertPermutationVector(packOp.getOuterDimsPerm()));
DenseMap<int64_t, OpFoldResult> dimAndTileMapping = DenseMap<int64_t, OpFoldResult> dimAndTileMapping =
packOp.getDimAndTileMapping(); packOp.getDimAndTileMapping();
SmallVector<OpFoldResult> srcDimValues = SmallVector<OpFoldResult> srcDimValues =
tensor::createDimValues(b, loc, packOp.getSource()); tensor::createDimValues(b, loc, packOp.getSource());
SmallVector<OpFoldResult> inputIndices, inputSizes; SmallVector<OpFoldResult> inputIndices, inputSizes;
for (auto dim : llvm::seq<int64_t>(0, inputRank)) { for (auto dim : llvm::seq<int64_t>(0, inputRank)) {
using AV = AffineValueExpr; using AV = affine::AffineValueExpr;
AffineBuilder ab(b, loc); affine::AffineBuilder ab(b, loc);
AffineExpr dim0, dim1, sym; AffineExpr dim0, dim1, sym;
bindDims(b.getContext(), dim0, dim1); bindDims(b.getContext(), dim0, dim1);
bindSymbols(b.getContext(), sym); bindSymbols(b.getContext(), sym);
if (dimAndTileMapping.count(dim)) { if (dimAndTileMapping.count(dim)) {
// If the data dimension is tiled, the i-th index is the product of // If the data dimension is tiled, the i-th index is the product of
// offset_i and tile_i, and the i-th size is the product of sizes_i and // offset_i and tile_i, and the i-th size is the product of sizes_i and
// tile_i. // tile_i.
auto avOffset = AV(dim0).bind(origOffsets[dim]); auto avOffset = AV(dim0).bind(origOffsets[dim]);
▲ Show 20 Lines • Show All 98 Lines • ▼ Show 20 Lines if (!dimAndTileMapping.count(tileDim)) {
info.sourceOffset = tileOffset; info.sourceOffset = tileOffset;
info.sourceSize = tileSize; info.sourceSize = tileSize;
info.resultOffset = zeroAttr; info.resultOffset = zeroAttr;
info.destExpandedSize = tileSize; info.destExpandedSize = tileSize;
return info; return info;
} }
Location loc = unpackOp.getLoc(); Location loc = unpackOp.getLoc();
using AV = AffineValueExpr; using AV = affine::AffineValueExpr;
AffineBuilder ab(b, loc); affine::AffineBuilder ab(b, loc);
AffineExpr dim0, dim1, sym0; AffineExpr dim0, dim1, sym0;
bindDims(b.getContext(), dim0, dim1); bindDims(b.getContext(), dim0, dim1);
bindSymbols(b.getContext(), sym0); bindSymbols(b.getContext(), sym0);
OpFoldResult innerTileSize = dimAndTileMapping[tileDim]; OpFoldResult innerTileSize = dimAndTileMapping[tileDim];
info.isAlignedToInnerTileSize = false; info.isAlignedToInnerTileSize = false;
FailureOr<int64_t> cstSize = ValueBoundsConstraintSet::computeConstantBound( FailureOr<int64_t> cstSize = ValueBoundsConstraintSet::computeConstantBound(
Show All 30 Lines if (info.isAlignedToInnerTileSize) {
// If the tiling size is 32, there will be 3 tiles. Two of them have // If the tiling size is 32, there will be 3 tiles. Two of them have
// size=32; one of them have size=2. The size is represented using // size=32; one of them have size=2. The size is represented using
// affine_min op; we need ceilDiv. // affine_min op; we need ceilDiv.
info.sourceSize = info.sourceSize =
ab.ceil(AV(dim0).bind(tileSize), AV(dim1).bind(innerTileSize)); ab.ceil(AV(dim0).bind(tileSize), AV(dim1).bind(innerTileSize));
return info; return info;
} }
DivModValue firstCoord = affine::DivModValue firstCoord = affine::getDivMod(
getDivMod(b, loc, getValueOrCreateConstantIndexOp(b, loc, tileOffset), b, loc, getValueOrCreateConstantIndexOp(b, loc, tileOffset),
getValueOrCreateConstantIndexOp(b, loc, innerTileSize)); getValueOrCreateConstantIndexOp(b, loc, innerTileSize));
OpFoldResult tileExclusiveBound = OpFoldResult tileExclusiveBound =
ab.add(AV(dim0).bind(tileOffset), AV(dim1).bind(tileSize)); ab.add(AV(dim0).bind(tileOffset), AV(dim1).bind(tileSize));
DivModValue lastCoord = getDivMod( affine::DivModValue lastCoord = affine::getDivMod(
b, loc, b, loc,
getValueOrCreateConstantIndexOp( getValueOrCreateConstantIndexOp(
b, loc, b, loc,
ab.sub(AV(dim0).bind(tileExclusiveBound), AV(dim1).bind(oneAttr))), ab.sub(AV(dim0).bind(tileExclusiveBound), AV(dim1).bind(oneAttr))),
getValueOrCreateConstantIndexOp(b, loc, innerTileSize)); getValueOrCreateConstantIndexOp(b, loc, innerTileSize));
OpFoldResult lengthMinusOne = ab.sub(AV(dim0).bind(lastCoord.quotient), OpFoldResult lengthMinusOne = ab.sub(AV(dim0).bind(lastCoord.quotient),
AV(dim1).bind(firstCoord.quotient)); AV(dim1).bind(firstCoord.quotient));
▲ Show 20 Lines • Show All 143 Lines • ▼ Show 20 LinesFailureOr<TilingResult> tensor::bubbleUpPadSlice(OpBuilder &b,
// Helper variables and functions for various arithmetic operations. These // Helper variables and functions for various arithmetic operations. These
// are used extensively for computing new offset/length and padding values. // are used extensively for computing new offset/length and padding values.
Location loc = padOp->getLoc(); Location loc = padOp->getLoc();
AffineExpr dim0, dim1; AffineExpr dim0, dim1;
bindDims(b.getContext(), dim0, dim1); bindDims(b.getContext(), dim0, dim1);
// Add two integers. // Add two integers.
auto addMap = AffineMap::get(2, 0, {dim0 + dim1}); auto addMap = AffineMap::get(2, 0, {dim0 + dim1});
auto add = [&](OpFoldResult v1, OpFoldResult v2) { auto add = [&](OpFoldResult v1, OpFoldResult v2) {
return makeComposedFoldedAffineApply(b, loc, addMap, {v1, v2}); return affine::makeComposedFoldedAffineApply(b, loc, addMap, {v1, v2});
}; };
// Subtract two integers. // Subtract two integers.
auto subMap = AffineMap::get(2, 0, {dim0 - dim1}); auto subMap = AffineMap::get(2, 0, {dim0 - dim1});
auto sub = [&](OpFoldResult v1, OpFoldResult v2) { auto sub = [&](OpFoldResult v1, OpFoldResult v2) {
return makeComposedFoldedAffineApply(b, loc, subMap, {v1, v2}); return affine::makeComposedFoldedAffineApply(b, loc, subMap, {v1, v2});
}; };
// Take the minimum of two integers. // Take the minimum of two integers.
auto idMap = AffineMap::getMultiDimIdentityMap(2, b.getContext()); auto idMap = AffineMap::getMultiDimIdentityMap(2, b.getContext());
auto min = [&](OpFoldResult v1, OpFoldResult v2) { auto min = [&](OpFoldResult v1, OpFoldResult v2) {
return makeComposedFoldedAffineMin(b, loc, idMap, {v1, v2}); return affine::makeComposedFoldedAffineMin(b, loc, idMap, {v1, v2});
}; };
// Take the maximum of two integers. // Take the maximum of two integers.
auto max = [&](OpFoldResult v1, OpFoldResult v2) { auto max = [&](OpFoldResult v1, OpFoldResult v2) {
return makeComposedFoldedAffineMax(b, loc, idMap, {v1, v2}); return affine::makeComposedFoldedAffineMax(b, loc, idMap, {v1, v2});
}; };
// Zero index-typed integer. // Zero index-typed integer.
OpFoldResult zero = b.getIndexAttr(0); OpFoldResult zero = b.getIndexAttr(0);
// Compute new offsets, lengths, low padding, high padding. // Compute new offsets, lengths, low padding, high padding.
SmallVector<OpFoldResult> newOffsets, newLengths, newStrides; SmallVector<OpFoldResult> newOffsets, newLengths, newStrides;
SmallVector<OpFoldResult> newLows, newHighs; SmallVector<OpFoldResult> newLows, newHighs;
// Set to true if the original data source is not read at all. // Set to true if the original data source is not read at all.
▲ Show 20 Lines • Show All 184 Lines • Show Last 20 Lines

mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp

Show First 20 Lines • Show All 887 Lines • ▼ Show 20 Lines for (int64_t i = 0; i < resultType.getRank(); ++i) {
if (!resultType.isDynamicDim(i)) if (!resultType.isDynamicDim(i))
continue; continue;
Value srcDim = rewriter.create<tensor::DimOp>(loc, padOp.getSource(), i); Value srcDim = rewriter.create<tensor::DimOp>(loc, padOp.getSource(), i);
Value lowPad = toValue(mixedLowPad[i]); Value lowPad = toValue(mixedLowPad[i]);
Value highPad = toValue(mixedHighPad[i]); Value highPad = toValue(mixedHighPad[i]);
AffineExpr s0, s1, s2; AffineExpr s0, s1, s2;
bindSymbols(op->getContext(), s0, s1, s2); bindSymbols(op->getContext(), s0, s1, s2);
AffineExpr sumExpr = s0 + s1 + s2; AffineExpr sumExpr = s0 + s1 + s2;
Value sum = rewriter.create<AffineApplyOp>( Value sum = rewriter.create<affine::AffineApplyOp>(
loc, sumExpr, ValueRange{srcDim, lowPad, highPad}); loc, sumExpr, ValueRange{srcDim, lowPad, highPad});
dynamicSizes.push_back(sum); dynamicSizes.push_back(sum);
} }
// Should the buffer be deallocated? // Should the buffer be deallocated?
bool dealloc = bool dealloc =
shouldDeallocateOpResult(padOp.getResult().cast<OpResult>(), options); shouldDeallocateOpResult(padOp.getResult().cast<OpResult>(), options);
// Allocate a buffer for the padded result. // Allocate a buffer for the padded result.
▲ Show 20 Lines • Show All 212 Lines • Show Last 20 Lines

mlir/lib/Dialect/Tensor/Transforms/ExtractSliceFromReshapeUtils.cpp

Show All 17 Lines
#include "mlir/Dialect/Tensor/Transforms/Transforms.h" #include "mlir/Dialect/Tensor/Transforms/Transforms.h"
#include "mlir/Dialect/Utils/ReshapeOpsUtils.h" #include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h" #include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/OpDefinition.h" #include "mlir/IR/OpDefinition.h"
#include "llvm/ADT/STLExtras.h" #include "llvm/ADT/STLExtras.h"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
using namespace mlir::tensor; using namespace mlir::tensor;
/// Get the dimension size of a value of RankedTensor type at the /// Get the dimension size of a value of RankedTensor type at the
static OpFoldResult getShapeDimSize(OpBuilder &b, Location loc, static OpFoldResult getShapeDimSize(OpBuilder &b, Location loc,
Value rankedTensor, int64_t dimIdx) { Value rankedTensor, int64_t dimIdx) {
RankedTensorType tensorType = rankedTensor.getType().cast<RankedTensorType>(); RankedTensorType tensorType = rankedTensor.getType().cast<RankedTensorType>();
if (!tensorType.isDynamicDim(dimIdx)) { if (!tensorType.isDynamicDim(dimIdx)) {
return b.getIndexAttr(tensorType.getDimSize(dimIdx)); return b.getIndexAttr(tensorType.getDimSize(dimIdx));
Show All 22 Linesstatic DimAndIndex invertSliceIndexing(OpBuilder &b, Location loc,
const DimAndIndex &dimAndIndex) { const DimAndIndex &dimAndIndex) {
AffineExpr d0, s0, s1; AffineExpr d0, s0, s1;
bindDims(b.getContext(), d0); bindDims(b.getContext(), d0);
bindSymbols(b.getContext(), s0, s1); bindSymbols(b.getContext(), s0, s1);
auto [dim, indexValue] = dimAndIndex; auto [dim, indexValue] = dimAndIndex;
assert(dim < sliceParams.size() && "slice should be non rank-reducing"); assert(dim < sliceParams.size() && "slice should be non rank-reducing");
return std::make_pair( return std::make_pair(
dim, dim,
makeComposedAffineApply( affine::makeComposedAffineApply(
b, loc, s0 + d0 * s1, b, loc, s0 + d0 * s1,
{indexValue, {indexValue,
getValueOrCreateConstantIndexOp(b, loc, sliceParams[dim].offset), getValueOrCreateConstantIndexOp(b, loc, sliceParams[dim].offset),
getValueOrCreateConstantIndexOp(b, loc, sliceParams[dim].stride)})); getValueOrCreateConstantIndexOp(b, loc, sliceParams[dim].stride)}));
} }
/// Transform `dimAndIndex` from the result tensor index space of a /// Transform `dimAndIndex` from the result tensor index space of a
/// CollapseShapeOp to the source tensor index space. /// CollapseShapeOp to the source tensor index space.
▲ Show 20 Lines • Show All 140 Lines • Show Last 20 Lines

mlir/lib/Dialect/Tensor/Transforms/FoldTensorSubsetOps.cpp

Show First 20 Lines • Show All 94 Lines • ▼ Show 20 Lines LogicalResult preconditionResult =
preconditionsFoldExtractOrInsertWithTransferOp(rewriter, readOp, preconditionsFoldExtractOrInsertWithTransferOp(rewriter, readOp,
extractSliceOp); extractSliceOp);
if (failed(preconditionResult)) if (failed(preconditionResult))
return preconditionResult; return preconditionResult;
SmallVector<Value> indices(readOp.getIndices().begin(), SmallVector<Value> indices(readOp.getIndices().begin(),
readOp.getIndices().end()); readOp.getIndices().end());
SmallVector<Value> sourceIndices; SmallVector<Value> sourceIndices;
resolveIndicesIntoOpWithOffsetsAndStrides( affine::resolveIndicesIntoOpWithOffsetsAndStrides(
rewriter, readOp.getLoc(), extractSliceOp.getMixedOffsets(), rewriter, readOp.getLoc(), extractSliceOp.getMixedOffsets(),
extractSliceOp.getMixedStrides(), extractSliceOp.getDroppedDims(), extractSliceOp.getMixedStrides(), extractSliceOp.getDroppedDims(),
indices, sourceIndices); indices, sourceIndices);
rewriter.replaceOpWithNewOp<vector::TransferReadOp>( rewriter.replaceOpWithNewOp<vector::TransferReadOp>(
readOp, readOp.getVectorType(), extractSliceOp.getSource(), sourceIndices, readOp, readOp.getVectorType(), extractSliceOp.getSource(), sourceIndices,
AffineMapAttr::get(expandDimsToRank( AffineMapAttr::get(expandDimsToRank(
readOp.getPermutationMap(), extractSliceOp.getSourceType().getRank(), readOp.getPermutationMap(), extractSliceOp.getSourceType().getRank(),
Show All 15 Lines LogicalResult preconditionResult =
preconditionsFoldExtractOrInsertWithTransferOp(rewriter, writeOp, preconditionsFoldExtractOrInsertWithTransferOp(rewriter, writeOp,
insertSliceOp); insertSliceOp);
if (failed(preconditionResult)) if (failed(preconditionResult))
return preconditionResult; return preconditionResult;
SmallVector<Value> indices(writeOp.getIndices().begin(), SmallVector<Value> indices(writeOp.getIndices().begin(),
writeOp.getIndices().end()); writeOp.getIndices().end());
SmallVector<Value> sourceIndices; SmallVector<Value> sourceIndices;
resolveIndicesIntoOpWithOffsetsAndStrides( affine::resolveIndicesIntoOpWithOffsetsAndStrides(
rewriter, writeOp.getLoc(), insertSliceOp.getMixedOffsets(), rewriter, writeOp.getLoc(), insertSliceOp.getMixedOffsets(),
insertSliceOp.getMixedStrides(), insertSliceOp.getDroppedDims(), indices, insertSliceOp.getMixedStrides(), insertSliceOp.getDroppedDims(), indices,
sourceIndices); sourceIndices);
rewriter.replaceOpWithNewOp<vector::TransferWriteOp>( rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
insertSliceOp, writeOp.getValue(), insertSliceOp.getDest(), sourceIndices, insertSliceOp, writeOp.getValue(), insertSliceOp.getDest(), sourceIndices,
AffineMapAttr::get(expandDimsToRank(writeOp.getPermutationMap(), AffineMapAttr::get(expandDimsToRank(writeOp.getPermutationMap(),
insertSliceOp.getDestType().getRank(), insertSliceOp.getDestType().getRank(),
Show All 38 Lines for (int64_t d = 0, e = insertSliceOp.getDestType().getRank(); d < e; ++d) {
} }
} }
// Resolve sizes according to dropped dims. // Resolve sizes according to dropped dims.
SmallVector<OpFoldResult> resolvedSizes; SmallVector<OpFoldResult> resolvedSizes;
// Note: the "insertSlice" case is symmetrical to the extract/subview case: // Note: the "insertSlice" case is symmetrical to the extract/subview case:
// `insertSliceOp` is passed as the "source" and `sourceInsertSliceOp` is // `insertSliceOp` is passed as the "source" and `sourceInsertSliceOp` is
// passed as the destination to the helper function. // passed as the destination to the helper function.
resolveSizesIntoOpWithSizes(insertSliceOp.getMixedSizes(), affine::resolveSizesIntoOpWithSizes(insertSliceOp.getMixedSizes(),
sourceInsertSliceOp.getMixedSizes(), sourceInsertSliceOp.getMixedSizes(),
droppedDims, resolvedSizes); droppedDims, resolvedSizes);
// If we are inside an InParallel region, temporarily set the insertion // If we are inside an InParallel region, temporarily set the insertion
// point outside: only tensor.parallel_insert_slice ops are allowed in // point outside: only tensor.parallel_insert_slice ops are allowed in
// there. // there.
if (std::is_same_v<OpTy, tensor::ParallelInsertSliceOp>) { if (std::is_same_v<OpTy, tensor::ParallelInsertSliceOp>) {
rewriter.setInsertionPoint( rewriter.setInsertionPoint(
insertSliceOp->template getParentOfType<scf::InParallelOp>()); insertSliceOp->template getParentOfType<scf::InParallelOp>());
} }
// Resolve offsets according to source offsets and strides. // Resolve offsets according to source offsets and strides.
SmallVector<Value> resolvedOffsets; SmallVector<Value> resolvedOffsets;
// Note: the "insertSlice" case is symmetrical to the extract/subview case: // Note: the "insertSlice" case is symmetrical to the extract/subview case:
// `insertSliceOp` is passed as the "source" and `sourceInsertSliceOp` is // `insertSliceOp` is passed as the "source" and `sourceInsertSliceOp` is
// passed as the destination to the helper function. // passed as the destination to the helper function.
resolveIndicesIntoOpWithOffsetsAndStrides( affine::resolveIndicesIntoOpWithOffsetsAndStrides(
rewriter, insertSliceOp.getLoc(), insertSliceOp.getMixedOffsets(), rewriter, insertSliceOp.getLoc(), insertSliceOp.getMixedOffsets(),
insertSliceOp.getMixedStrides(), droppedDims, insertSliceOp.getMixedStrides(), droppedDims,
sourceInsertSliceOp.getMixedOffsets(), resolvedOffsets); sourceInsertSliceOp.getMixedOffsets(), resolvedOffsets);
// Reset the insertion point. // Reset the insertion point.
rewriter.setInsertionPoint(insertSliceOp); rewriter.setInsertionPoint(insertSliceOp);
// Replace original op. // Replace original op.
rewriter.replaceOpWithNewOp<OpTy>( rewriter.replaceOpWithNewOp<OpTy>(
Show All 37 Lines

mlir/lib/Dialect/Tensor/Transforms/MergeConsecutiveInsertExtractSlicePatterns.cpp

Show All 23 Linesstruct MergeConsecutiveExtractSlice : public OpRewritePattern<ExtractSliceOp> {
LogicalResult matchAndRewrite(ExtractSliceOp nextOp, LogicalResult matchAndRewrite(ExtractSliceOp nextOp,
PatternRewriter &rewriter) const override { PatternRewriter &rewriter) const override {
auto prevOp = nextOp.getSource().getDefiningOp<ExtractSliceOp>(); auto prevOp = nextOp.getSource().getDefiningOp<ExtractSliceOp>();
if (!prevOp) if (!prevOp)
return failure(); return failure();
SmallVector<OpFoldResult> newOffsets, newSizes, newStrides; SmallVector<OpFoldResult> newOffsets, newSizes, newStrides;
if (failed(mergeOffsetsSizesAndStrides(rewriter, nextOp.getLoc(), prevOp, if (failed(affine::mergeOffsetsSizesAndStrides(
nextOp, prevOp.getDroppedDims(), rewriter, nextOp.getLoc(), prevOp, nextOp, prevOp.getDroppedDims(),
newOffsets, newSizes, newStrides))) newOffsets, newSizes, newStrides)))
return failure(); return failure();
rewriter.replaceOpWithNewOp<ExtractSliceOp>(nextOp, nextOp.getType(), rewriter.replaceOpWithNewOp<ExtractSliceOp>(nextOp, nextOp.getType(),
prevOp.getSource(), newOffsets, prevOp.getSource(), newOffsets,
newSizes, newStrides); newSizes, newStrides);
return success(); return success();
} }
}; };
▲ Show 20 Lines • Show All 45 Lines • Show Last 20 Lines

mlir/lib/Dialect/Tensor/Utils/Utils.cpp

Show All 29 Lines for (const auto &en : enumerate(type.getShape())) {
// Pad only the static dimensions of the result tensor type. // Pad only the static dimensions of the result tensor type.
if (ShapedType::isDynamic(en.value())) if (ShapedType::isDynamic(en.value()))
continue; continue;
// Compute the padding width. // Compute the padding width.
AffineExpr d0; AffineExpr d0;
bindDims(b.getContext(), d0); bindDims(b.getContext(), d0);
auto dimOp = b.createOrFold<tensor::DimOp>(loc, source, en.index()); auto dimOp = b.createOrFold<tensor::DimOp>(loc, source, en.index());
high[en.index()] = high[en.index()] =
makeComposedAffineApply(b, loc, en.value() - d0, {dimOp}).getResult(); affine::makeComposedAffineApply(b, loc, en.value() - d0, {dimOp})
.getResult();
} }
return b.create<PadOp>(loc, type, source, low, high, pad, nofold); return b.create<PadOp>(loc, type, source, low, high, pad, nofold);
} }
SmallVector<Value> mlir::tensor::createDynamicDimValues(OpBuilder &b, SmallVector<Value> mlir::tensor::createDynamicDimValues(OpBuilder &b,
Location loc, Location loc,
Value rankedTensor) { Value rankedTensor) {
auto tensorTy = rankedTensor.getType().cast<RankedTensorType>(); auto tensorTy = rankedTensor.getType().cast<RankedTensorType>();
▲ Show 20 Lines • Show All 57 Lines • Show Last 20 Lines

mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp

Show First 20 Lines • Show All 66 Lines • ▼ Show 20 Lines DistributedLoadStoreHelper(Value sequentialVal, Value distributedVal,
if (sequentialVectorType && distributedVectorType) if (sequentialVectorType && distributedVectorType)
distributionMap = distributionMap =
calculateImplicitMap(sequentialVectorType, distributedVectorType); calculateImplicitMap(sequentialVectorType, distributedVectorType);
} }
Value buildDistributedOffset(RewriterBase &b, Location loc, int64_t index) { Value buildDistributedOffset(RewriterBase &b, Location loc, int64_t index) {
int64_t distributedSize = distributedVectorType.getDimSize(index); int64_t distributedSize = distributedVectorType.getDimSize(index);
AffineExpr tid = getAffineSymbolExpr(0, b.getContext()); AffineExpr tid = getAffineSymbolExpr(0, b.getContext());
return b.createOrFold<AffineApplyOp>(loc, tid * distributedSize, return b.createOrFold<affine::AffineApplyOp>(loc, tid * distributedSize,
ArrayRef<Value>{laneId}); ArrayRef<Value>{laneId});
} }
/// Create a store during the process of distributing the /// Create a store during the process of distributing the
/// `vector.warp_execute_on_thread_0` op. /// `vector.warp_execute_on_thread_0` op.
/// Vector distribution assumes the following convention regarding the /// Vector distribution assumes the following convention regarding the
/// temporary buffers that are created to transition values. This **must** /// temporary buffers that are created to transition values. This **must**
/// be properly specified in the `options.warpAllocationFn`: /// be properly specified in the `options.warpAllocationFn`:
/// 1. scalars of type T transit through a memref<1xT>. /// 1. scalars of type T transit through a memref<1xT>.
▲ Show 20 Lines • Show All 423 Lines • ▼ Show 20 Lines for (auto it : llvm::zip(indexMap.getResults(), map.getResults())) {
bindDims(newWarpOp.getContext(), d0, d1); bindDims(newWarpOp.getContext(), d0, d1);
auto indexExpr = std::get<0>(it).dyn_cast<AffineDimExpr>(); auto indexExpr = std::get<0>(it).dyn_cast<AffineDimExpr>();
if (!indexExpr) if (!indexExpr)
continue; continue;
unsigned indexPos = indexExpr.getPosition(); unsigned indexPos = indexExpr.getPosition();
unsigned vectorPos = std::get<1>(it).cast<AffineDimExpr>().getPosition(); unsigned vectorPos = std::get<1>(it).cast<AffineDimExpr>().getPosition();
auto scale = auto scale =
rewriter.getAffineConstantExpr(targetType.getDimSize(vectorPos)); rewriter.getAffineConstantExpr(targetType.getDimSize(vectorPos));
indices[indexPos] = indices[indexPos] = affine::makeComposedAffineApply(
makeComposedAffineApply(rewriter, loc, d0 + scale * d1, rewriter, loc, d0 + scale * d1,
{indices[indexPos], newWarpOp.getLaneid()}); {indices[indexPos], newWarpOp.getLaneid()});
} }
newWriteOp.getIndicesMutable().assign(indices); newWriteOp.getIndicesMutable().assign(indices);
return success(); return success();
} }
/// Extract TransferWriteOps of vector<1x> into a separate warp op. /// Extract TransferWriteOps of vector<1x> into a separate warp op.
LogicalResult tryExtractOp(RewriterBase &rewriter, LogicalResult tryExtractOp(RewriterBase &rewriter,
▲ Show 20 Lines • Show All 221 Lines • ▼ Show 20 Lines for (auto it : llvm::zip(indexMap.getResults(), map.getResults())) {
bindDims(read.getContext(), d0, d1); bindDims(read.getContext(), d0, d1);
auto indexExpr = std::get<0>(it).dyn_cast<AffineDimExpr>(); auto indexExpr = std::get<0>(it).dyn_cast<AffineDimExpr>();
if (!indexExpr) if (!indexExpr)
continue; continue;
unsigned indexPos = indexExpr.getPosition(); unsigned indexPos = indexExpr.getPosition();
unsigned vectorPos = std::get<1>(it).cast<AffineDimExpr>().getPosition(); unsigned vectorPos = std::get<1>(it).cast<AffineDimExpr>().getPosition();
int64_t scale = int64_t scale =
distributedVal.getType().cast<VectorType>().getDimSize(vectorPos); distributedVal.getType().cast<VectorType>().getDimSize(vectorPos);
indices[indexPos] = indices[indexPos] = affine::makeComposedAffineApply(
makeComposedAffineApply(rewriter, read.getLoc(), d0 + scale * d1, rewriter, read.getLoc(), d0 + scale * d1,
{indices[indexPos], warpOp.getLaneid()}); {indices[indexPos], warpOp.getLaneid()});
} }
Value newRead = rewriter.create<vector::TransferReadOp>( Value newRead = rewriter.create<vector::TransferReadOp>(
read.getLoc(), distributedVal.getType(), read.getSource(), indices, read.getLoc(), distributedVal.getType(), read.getSource(), indices,
read.getPermutationMapAttr(), read.getPadding(), read.getMask(), read.getPermutationMapAttr(), read.getPadding(), read.getMask(),
read.getInBoundsAttr()); read.getInBoundsAttr());
rewriter.replaceAllUsesWith(distributedVal, newRead); rewriter.replaceAllUsesWith(distributedVal, newRead);
return success(); return success();
} }
▲ Show 20 Lines • Show All 274 Lines • ▼ Show 20 Lines if (is0dOrVec1Extract) {
return success(); return success();
} }
// 1d extract: Distribute the source vector. One lane extracts and shuffles // 1d extract: Distribute the source vector. One lane extracts and shuffles
// the value to all other lanes. // the value to all other lanes.
int64_t elementsPerLane = distributedVecType.getShape()[0]; int64_t elementsPerLane = distributedVecType.getShape()[0];
AffineExpr sym0 = getAffineSymbolExpr(0, rewriter.getContext()); AffineExpr sym0 = getAffineSymbolExpr(0, rewriter.getContext());
// tid of extracting thread: pos / elementsPerLane // tid of extracting thread: pos / elementsPerLane
Value broadcastFromTid = rewriter.create<AffineApplyOp>( Value broadcastFromTid = rewriter.create<affine::AffineApplyOp>(
loc, sym0.ceilDiv(elementsPerLane), extractOp.getPosition()); loc, sym0.ceilDiv(elementsPerLane), extractOp.getPosition());
// Extract at position: pos % elementsPerLane // Extract at position: pos % elementsPerLane
Value pos = Value pos =
elementsPerLane == 1 elementsPerLane == 1
? rewriter.create<arith::ConstantIndexOp>(loc, 0).getResult() ? rewriter.create<arith::ConstantIndexOp>(loc, 0).getResult()
: rewriter : rewriter
.create<AffineApplyOp>(loc, sym0 % elementsPerLane, .create<affine::AffineApplyOp>(loc, sym0 % elementsPerLane,
extractOp.getPosition()) extractOp.getPosition())
.getResult(); .getResult();
Value extracted = Value extracted =
rewriter.create<vector::ExtractElementOp>(loc, distributedVec, pos); rewriter.create<vector::ExtractElementOp>(loc, distributedVec, pos);
// Shuffle the extracted value to all lanes. // Shuffle the extracted value to all lanes.
Value shuffled = warpShuffleFromIdxFn( Value shuffled = warpShuffleFromIdxFn(
loc, rewriter, extracted, broadcastFromTid, newWarpOp.getWarpSize()); loc, rewriter, extracted, broadcastFromTid, newWarpOp.getWarpSize());
rewriter.replaceAllUsesWith(newWarpOp->getResult(operandNumber), shuffled); rewriter.replaceAllUsesWith(newWarpOp->getResult(operandNumber), shuffled);
▲ Show 20 Lines • Show All 48 Lines • ▼ Show 20 Lines if (vecType == distrType) {
newInsert); newInsert);
return success(); return success();
} }
// This is a distribution. Only one lane should insert. // This is a distribution. Only one lane should insert.
int64_t elementsPerLane = distrType.getShape()[0]; int64_t elementsPerLane = distrType.getShape()[0];
AffineExpr sym0 = getAffineSymbolExpr(0, rewriter.getContext()); AffineExpr sym0 = getAffineSymbolExpr(0, rewriter.getContext());
// tid of extracting thread: pos / elementsPerLane // tid of extracting thread: pos / elementsPerLane
Value insertingLane = rewriter.create<AffineApplyOp>( Value insertingLane = rewriter.create<affine::AffineApplyOp>(
loc, sym0.ceilDiv(elementsPerLane), newPos); loc, sym0.ceilDiv(elementsPerLane), newPos);
// Insert position: pos % elementsPerLane // Insert position: pos % elementsPerLane
Value pos = Value pos =
elementsPerLane == 1 elementsPerLane == 1
? rewriter.create<arith::ConstantIndexOp>(loc, 0).getResult() ? rewriter.create<arith::ConstantIndexOp>(loc, 0).getResult()
: rewriter : rewriter
.create<AffineApplyOp>(loc, sym0 % elementsPerLane, newPos) .create<affine::AffineApplyOp>(loc, sym0 % elementsPerLane,
newPos)
.getResult(); .getResult();
Value isInsertingLane = rewriter.create<arith::CmpIOp>( Value isInsertingLane = rewriter.create<arith::CmpIOp>(
loc, arith::CmpIPredicate::eq, newWarpOp.getLaneid(), insertingLane); loc, arith::CmpIPredicate::eq, newWarpOp.getLaneid(), insertingLane);
Value newResult = Value newResult =
rewriter rewriter
.create<scf::IfOp>( .create<scf::IfOp>(
loc, isInsertingLane, loc, isInsertingLane,
/*thenBuilder=*/ /*thenBuilder=*/
▲ Show 20 Lines • Show All 438 Lines • Show Last 20 Lines

mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp

Show First 20 Lines • Show All 593 Lines • ▼ Show 20 Lines LogicalResult matchAndRewrite(vector::ExtractElementOp extractOp,
if (xferOp.hasOutOfBoundsDim() && xferOp.getType().getRank() > 0) if (xferOp.hasOutOfBoundsDim() && xferOp.getType().getRank() > 0)
return failure(); return failure();
// Construct scalar load. // Construct scalar load.
SmallVector<Value> newIndices(xferOp.getIndices().begin(), SmallVector<Value> newIndices(xferOp.getIndices().begin(),
xferOp.getIndices().end()); xferOp.getIndices().end());
if (extractOp.getPosition()) { if (extractOp.getPosition()) {
AffineExpr sym0, sym1; AffineExpr sym0, sym1;
bindSymbols(extractOp.getContext(), sym0, sym1); bindSymbols(extractOp.getContext(), sym0, sym1);
OpFoldResult ofr = makeComposedFoldedAffineApply( OpFoldResult ofr = affine::makeComposedFoldedAffineApply(
rewriter, extractOp.getLoc(), sym0 + sym1, rewriter, extractOp.getLoc(), sym0 + sym1,
{newIndices[newIndices.size() - 1], extractOp.getPosition()}); {newIndices[newIndices.size() - 1], extractOp.getPosition()});
if (ofr.is<Value>()) { if (ofr.is<Value>()) {
newIndices[newIndices.size() - 1] = ofr.get<Value>(); newIndices[newIndices.size() - 1] = ofr.get<Value>();
} else { } else {
newIndices[newIndices.size() - 1] = newIndices[newIndices.size() - 1] =
rewriter.create<arith::ConstantIndexOp>(extractOp.getLoc(), rewriter.create<arith::ConstantIndexOp>(extractOp.getLoc(),
*getConstantIntValue(ofr)); *getConstantIntValue(ofr));
▲ Show 20 Lines • Show All 47 Lines • ▼ Show 20 Lines if (xferOp.hasOutOfBoundsDim() && xferOp.getType().getRank() > 0)
return failure(); return failure();
// Construct scalar load. // Construct scalar load.
SmallVector<Value> newIndices(xferOp.getIndices().begin(), SmallVector<Value> newIndices(xferOp.getIndices().begin(),
xferOp.getIndices().end()); xferOp.getIndices().end());
for (const auto &it : llvm::enumerate(extractOp.getPosition())) { for (const auto &it : llvm::enumerate(extractOp.getPosition())) {
int64_t offset = it.value().cast<IntegerAttr>().getInt(); int64_t offset = it.value().cast<IntegerAttr>().getInt();
int64_t idx = int64_t idx =
newIndices.size() - extractOp.getPosition().size() + it.index(); newIndices.size() - extractOp.getPosition().size() + it.index();
OpFoldResult ofr = makeComposedFoldedAffineApply( OpFoldResult ofr = affine::makeComposedFoldedAffineApply(
rewriter, extractOp.getLoc(), rewriter, extractOp.getLoc(),
rewriter.getAffineSymbolExpr(0) + offset, {newIndices[idx]}); rewriter.getAffineSymbolExpr(0) + offset, {newIndices[idx]});
if (ofr.is<Value>()) { if (ofr.is<Value>()) {
newIndices[idx] = ofr.get<Value>(); newIndices[idx] = ofr.get<Value>();
} else { } else {
newIndices[idx] = rewriter.create<arith::ConstantIndexOp>( newIndices[idx] = rewriter.create<arith::ConstantIndexOp>(
extractOp.getLoc(), *getConstantIntValue(ofr)); extractOp.getLoc(), *getConstantIntValue(ofr));
} }
▲ Show 20 Lines • Show All 93 Lines • Show Last 20 Lines

mlir/lib/Dialect/Vector/Transforms/VectorTransferSplitRewritePatterns.cpp

Show All 35 Lines
#define DEBUG_TYPE "vector-transfer-split" #define DEBUG_TYPE "vector-transfer-split"
using namespace mlir; using namespace mlir;
using namespace mlir::vector; using namespace mlir::vector;
static std::optional<int64_t> extractConstantIndex(Value v) { static std::optional<int64_t> extractConstantIndex(Value v) {
if (auto cstOp = v.getDefiningOp<arith::ConstantIndexOp>()) if (auto cstOp = v.getDefiningOp<arith::ConstantIndexOp>())
return cstOp.value(); return cstOp.value();
if (auto affineApplyOp = v.getDefiningOp<AffineApplyOp>()) if (auto affineApplyOp = v.getDefiningOp<affine::AffineApplyOp>())
if (affineApplyOp.getAffineMap().isSingleConstant()) if (affineApplyOp.getAffineMap().isSingleConstant())
return affineApplyOp.getAffineMap().getSingleConstantResult(); return affineApplyOp.getAffineMap().getSingleConstantResult();
return std::nullopt; return std::nullopt;
} }
// Missing foldings of scf.if make it necessary to perform poor man's folding // Missing foldings of scf.if make it necessary to perform poor man's folding
// eagerly, especially in the case of unrolling. In the future, this should go // eagerly, especially in the case of unrolling. In the future, this should go
// away once scf.if folds properly. // away once scf.if folds properly.
Show All 18 Lines xferOp.zipResultAndIndexing([&](int64_t resultIdx, int64_t indicesIdx) {
// the construction of `inBoundsCond`. // the construction of `inBoundsCond`.
if (xferOp.isDimInBounds(resultIdx)) if (xferOp.isDimInBounds(resultIdx))
return; return;
// Fold or create the check that `index + vector_size` <= `memref_size`. // Fold or create the check that `index + vector_size` <= `memref_size`.
Location loc = xferOp.getLoc(); Location loc = xferOp.getLoc();
int64_t vectorSize = xferOp.getVectorType().getDimSize(resultIdx); int64_t vectorSize = xferOp.getVectorType().getDimSize(resultIdx);
auto d0 = getAffineDimExpr(0, xferOp.getContext()); auto d0 = getAffineDimExpr(0, xferOp.getContext());
auto vs = getAffineConstantExpr(vectorSize, xferOp.getContext()); auto vs = getAffineConstantExpr(vectorSize, xferOp.getContext());
Value sum = Value sum = affine::makeComposedAffineApply(b, loc, d0 + vs,
makeComposedAffineApply(b, loc, d0 + vs, xferOp.indices()[indicesIdx]); xferOp.indices()[indicesIdx]);
Value cond = createFoldedSLE( Value cond = createFoldedSLE(
b, sum, vector::createOrFoldDimOp(b, loc, xferOp.source(), indicesIdx)); b, sum, vector::createOrFoldDimOp(b, loc, xferOp.source(), indicesIdx));
if (!cond) if (!cond)
return; return;
// Conjunction over all dims for which we are in-bounds. // Conjunction over all dims for which we are in-bounds.
if (inBoundsCond) if (inBoundsCond)
inBoundsCond = b.create<arith::AndIOp>(loc, inBoundsCond, cond); inBoundsCond = b.create<arith::AndIOp>(loc, inBoundsCond, cond);
else else
▲ Show 20 Lines • Show All 114 Lines • ▼ Show 20 Lines Value dimMemRef = vector::createOrFoldDimOp(b, xferOp.getLoc(),
xferOp.source(), indicesIdx); xferOp.source(), indicesIdx);
Value dimAlloc = b.create<memref::DimOp>(loc, alloc, resultIdx); Value dimAlloc = b.create<memref::DimOp>(loc, alloc, resultIdx);
Value index = xferOp.indices()[indicesIdx]; Value index = xferOp.indices()[indicesIdx];
AffineExpr i, j, k; AffineExpr i, j, k;
bindDims(xferOp.getContext(), i, j, k); bindDims(xferOp.getContext(), i, j, k);
SmallVector<AffineMap, 4> maps = SmallVector<AffineMap, 4> maps =
AffineMap::inferFromExprList(MapList{{i - j, k}}); AffineMap::inferFromExprList(MapList{{i - j, k}});
// affine_min(%dimMemRef - %index, %dimAlloc) // affine_min(%dimMemRef - %index, %dimAlloc)
Value affineMin = b.create<AffineMinOp>( Value affineMin = b.create<affine::AffineMinOp>(
loc, index.getType(), maps[0], ValueRange{dimMemRef, index, dimAlloc}); loc, index.getType(), maps[0], ValueRange{dimMemRef, index, dimAlloc});
sizes.push_back(affineMin); sizes.push_back(affineMin);
}); });
SmallVector<OpFoldResult> srcIndices = llvm::to_vector<4>(llvm::map_range( SmallVector<OpFoldResult> srcIndices = llvm::to_vector<4>(llvm::map_range(
xferOp.indices(), [](Value idx) -> OpFoldResult { return idx; })); xferOp.indices(), [](Value idx) -> OpFoldResult { return idx; }));
SmallVector<OpFoldResult> destIndices(memrefRank, b.getIndexAttr(0)); SmallVector<OpFoldResult> destIndices(memrefRank, b.getIndexAttr(0));
SmallVector<OpFoldResult> strides(memrefRank, b.getIndexAttr(1)); SmallVector<OpFoldResult> strides(memrefRank, b.getIndexAttr(1));
▲ Show 20 Lines • Show All 224 Lines • ▼ Show 20 Linesstatic Operation *getAutomaticAllocationScope(Operation *op) {
// Find the closest surrounding allocation scope that is not a known looping // Find the closest surrounding allocation scope that is not a known looping
// construct (putting alloca's in loops doesn't always lower to deallocation // construct (putting alloca's in loops doesn't always lower to deallocation
// until the end of the loop). // until the end of the loop).
Operation *scope = nullptr; Operation *scope = nullptr;
for (Operation *parent = op->getParentOp(); parent != nullptr; for (Operation *parent = op->getParentOp(); parent != nullptr;
parent = parent->getParentOp()) { parent = parent->getParentOp()) {
if (parent->hasTrait<OpTrait::AutomaticAllocationScope>()) if (parent->hasTrait<OpTrait::AutomaticAllocationScope>())
scope = parent; scope = parent;
if (!isa<scf::ForOp, AffineForOp>(parent)) if (!isa<scf::ForOp, affine::AffineForOp>(parent))
break; break;
} }
assert(scope && "Expected op to be inside automatic allocation scope"); assert(scope && "Expected op to be inside automatic allocation scope");
return scope; return scope;
} }
/// Split a vector.transfer operation into an in-bounds (i.e., no out-of-bounds /// Split a vector.transfer operation into an in-bounds (i.e., no out-of-bounds
/// masking) fastpath and a slowpath. /// masking) fastpath and a slowpath.
▲ Show 20 Lines • Show All 215 Lines • Show Last 20 Lines

mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp

Show First 20 Lines • Show All 112 Lines • ▼ Show 20 Linesstatic SmallVector<Value> sliceTransferIndices(ArrayRef<int64_t> elementOffsets,
SmallVector<Value> slicedIndices(indices.begin(), indices.end()); SmallVector<Value> slicedIndices(indices.begin(), indices.end());
for (const auto &dim : llvm::enumerate(permutationMap.getResults())) { for (const auto &dim : llvm::enumerate(permutationMap.getResults())) {
if (isBroadcast(dim.value())) if (isBroadcast(dim.value()))
continue; continue;
unsigned pos = dim.value().cast<AffineDimExpr>().getPosition(); unsigned pos = dim.value().cast<AffineDimExpr>().getPosition();
auto expr = getAffineDimExpr(0, builder.getContext()) + auto expr = getAffineDimExpr(0, builder.getContext()) +
getAffineConstantExpr(elementOffsets[dim.index()], ctx); getAffineConstantExpr(elementOffsets[dim.index()], ctx);
auto map = AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, expr); auto map = AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, expr);
slicedIndices[pos] = builder.create<AffineApplyOp>(loc, map, indices[pos]); slicedIndices[pos] =
builder.create<affine::AffineApplyOp>(loc, map, indices[pos]);
} }
return slicedIndices; return slicedIndices;
} }
// Clones `op` into a new operations that takes `operands` and returns // Clones `op` into a new operations that takes `operands` and returns
// `resultTypes`. // `resultTypes`.
static Operation *cloneOpWithOperandsAndTypes(OpBuilder &builder, Location loc, static Operation *cloneOpWithOperandsAndTypes(OpBuilder &builder, Location loc,
Operation *op, Operation *op,
▲ Show 20 Lines • Show All 523 Lines • Show Last 20 Lines

mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp

Show First 20 Lines • Show All 71 Lines • ▼ Show 20 Lines if (enclosingLoopToVectorDim.empty())
return AffineMap(); return AffineMap();
MLIRContext *context = MLIRContext *context =
enclosingLoopToVectorDim.begin()->getFirst()->getContext(); enclosingLoopToVectorDim.begin()->getFirst()->getContext();
SmallVector<AffineExpr> perm(enclosingLoopToVectorDim.size(), SmallVector<AffineExpr> perm(enclosingLoopToVectorDim.size(),
getAffineConstantExpr(0, context)); getAffineConstantExpr(0, context));
for (auto kvp : enclosingLoopToVectorDim) { for (auto kvp : enclosingLoopToVectorDim) {
assert(kvp.second < perm.size()); assert(kvp.second < perm.size());
auto invariants = getInvariantAccesses( auto invariants = affine::getInvariantAccesses(
cast<AffineForOp>(kvp.first).getInductionVar(), indices); cast<affine::AffineForOp>(kvp.first).getInductionVar(), indices);
unsigned numIndices = indices.size(); unsigned numIndices = indices.size();
unsigned countInvariantIndices = 0; unsigned countInvariantIndices = 0;
for (unsigned dim = 0; dim < numIndices; ++dim) { for (unsigned dim = 0; dim < numIndices; ++dim) {
if (!invariants.count(indices[dim])) { if (!invariants.count(indices[dim])) {
assert(perm[kvp.second] == getAffineConstantExpr(0, context) && assert(perm[kvp.second] == getAffineConstantExpr(0, context) &&
"permutationMap already has an entry along dim"); "permutationMap already has an entry along dim");
perm[kvp.second] = getAffineDimExpr(dim, context); perm[kvp.second] = getAffineDimExpr(dim, context);
} else { } else {
Show All 24 Lines while (current) {
} }
current = current->getParentOp(); current = current->getParentOp();
} }
return res; return res;
} }
/// Returns the enclosing AffineForOp, from closest to farthest. /// Returns the enclosing AffineForOp, from closest to farthest.
static SetVector<Operation *> getEnclosingforOps(Block *block) { static SetVector<Operation *> getEnclosingforOps(Block *block) {
return getParentsOfType<AffineForOp>(block); return getParentsOfType<affine::AffineForOp>(block);
} }
AffineMap mlir::makePermutationMap( AffineMap mlir::makePermutationMap(
Block *insertPoint, ArrayRef<Value> indices, Block *insertPoint, ArrayRef<Value> indices,
const DenseMap<Operation *, unsigned> &loopToVectorDim) { const DenseMap<Operation *, unsigned> &loopToVectorDim) {
DenseMap<Operation *, unsigned> enclosingLoopToVectorDim; DenseMap<Operation *, unsigned> enclosingLoopToVectorDim;
auto enclosingLoops = getEnclosingforOps(insertPoint); auto enclosingLoops = getEnclosingforOps(insertPoint);
for (auto *forInst : enclosingLoops) { for (auto *forInst : enclosingLoops) {
▲ Show 20 Lines • Show All 67 Lines • Show Last 20 Lines

mlir/test/lib/Analysis/TestMemRefBoundCheck.cpp

Show All 17 Lines
#include "mlir/IR/Builders.h" #include "mlir/IR/Builders.h"
#include "mlir/Pass/Pass.h" #include "mlir/Pass/Pass.h"
#include "llvm/ADT/TypeSwitch.h" #include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/Debug.h" #include "llvm/Support/Debug.h"
#define DEBUG_TYPE "memref-bound-check" #define DEBUG_TYPE "memref-bound-check"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
namespace { namespace {
/// Checks for out of bound memref access subscripts.. /// Checks for out of bound memref access subscripts..
struct TestMemRefBoundCheck struct TestMemRefBoundCheck
: public PassWrapper<TestMemRefBoundCheck, OperationPass<>> { : public PassWrapper<TestMemRefBoundCheck, OperationPass<>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestMemRefBoundCheck) MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestMemRefBoundCheck)
Show All 24 Lines

mlir/test/lib/Analysis/TestMemRefDependenceCheck.cpp

Show All 15 Lines
#include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/IR/Builders.h" #include "mlir/IR/Builders.h"
#include "mlir/Pass/Pass.h" #include "mlir/Pass/Pass.h"
#include "llvm/Support/Debug.h" #include "llvm/Support/Debug.h"
#define DEBUG_TYPE "test-memref-dependence-check" #define DEBUG_TYPE "test-memref-dependence-check"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
namespace { namespace {
// TODO: Add common surrounding loop depth-wise dependence checks. // TODO: Add common surrounding loop depth-wise dependence checks.
/// Checks dependences between all pairs of memref accesses in a Function. /// Checks dependences between all pairs of memref accesses in a Function.
struct TestMemRefDependenceCheck struct TestMemRefDependenceCheck
: public PassWrapper<TestMemRefDependenceCheck, OperationPass<>> { : public PassWrapper<TestMemRefDependenceCheck, OperationPass<>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestMemRefDependenceCheck) MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestMemRefDependenceCheck)
▲ Show 20 Lines • Show All 93 Lines • Show Last 20 Lines

mlir/test/lib/Dialect/Affine/TestAffineDataCopy.cpp

Show All 17 Lines
#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Pass/Pass.h" #include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/Passes.h" #include "mlir/Transforms/Passes.h"
#define PASS_NAME "test-affine-data-copy" #define PASS_NAME "test-affine-data-copy"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
namespace { namespace {
struct TestAffineDataCopy struct TestAffineDataCopy
: public PassWrapper<TestAffineDataCopy, OperationPass<func::FuncOp>> { : public PassWrapper<TestAffineDataCopy, OperationPass<func::FuncOp>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestAffineDataCopy) MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestAffineDataCopy)
StringRef getArgument() const final { return PASS_NAME; } StringRef getArgument() const final { return PASS_NAME; }
▲ Show 20 Lines • Show All 112 Lines • Show Last 20 Lines

mlir/test/lib/Dialect/Affine/TestAffineLoopParametricTiling.cpp

Show All 11 Lines
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
#include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/LoopUtils.h" #include "mlir/Dialect/Affine/LoopUtils.h"
#include "mlir/Dialect/Affine/Passes.h" #include "mlir/Dialect/Affine/Passes.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/Func/IR/FuncOps.h"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
#define DEBUG_TYPE "test-affine-parametric-tile" #define DEBUG_TYPE "test-affine-parametric-tile"
namespace { namespace {
struct TestAffineLoopParametricTiling struct TestAffineLoopParametricTiling
: public PassWrapper<TestAffineLoopParametricTiling, : public PassWrapper<TestAffineLoopParametricTiling,
OperationPass<func::FuncOp>> { OperationPass<func::FuncOp>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestAffineLoopParametricTiling) MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestAffineLoopParametricTiling)
▲ Show 20 Lines • Show All 77 Lines • Show Last 20 Lines

mlir/test/lib/Dialect/Affine/TestAffineLoopUnswitching.cpp

Show All 13 Lines
#include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/Utils.h" #include "mlir/Dialect/Affine/Utils.h"
#include "mlir/Pass/Pass.h" #include "mlir/Pass/Pass.h"
#include "mlir/Transforms/Passes.h" #include "mlir/Transforms/Passes.h"
#define PASS_NAME "test-affine-loop-unswitch" #define PASS_NAME "test-affine-loop-unswitch"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
namespace { namespace {
/// This pass applies the permutation on the first maximal perfect nest. /// This pass applies the permutation on the first maximal perfect nest.
struct TestAffineLoopUnswitching struct TestAffineLoopUnswitching
: public PassWrapper<TestAffineLoopUnswitching, OperationPass<>> { : public PassWrapper<TestAffineLoopUnswitching, OperationPass<>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestAffineLoopUnswitching) MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestAffineLoopUnswitching)
Show All 36 Lines

mlir/test/lib/Dialect/Affine/TestDecomposeAffineOps.cpp

Show All 15 Lines
#include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/PatternMatch.h" #include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/Pass.h" #include "mlir/Pass/Pass.h"
#include "mlir/Transforms/Passes.h" #include "mlir/Transforms/Passes.h"
#define PASS_NAME "test-decompose-affine-ops" #define PASS_NAME "test-decompose-affine-ops"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
namespace { namespace {
struct TestDecomposeAffineOps struct TestDecomposeAffineOps
: public PassWrapper<TestDecomposeAffineOps, OperationPass<func::FuncOp>> { : public PassWrapper<TestDecomposeAffineOps, OperationPass<func::FuncOp>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestDecomposeAffineOps) MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestDecomposeAffineOps)
StringRef getArgument() const final { return PASS_NAME; } StringRef getArgument() const final { return PASS_NAME; }
Show All 26 Lines

mlir/test/lib/Dialect/Affine/TestLoopFusion.cpp

Show All 14 Lines
#include "mlir/Dialect/Affine/LoopFusionUtils.h" #include "mlir/Dialect/Affine/LoopFusionUtils.h"
#include "mlir/Dialect/Affine/LoopUtils.h" #include "mlir/Dialect/Affine/LoopUtils.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Pass/Pass.h" #include "mlir/Pass/Pass.h"
#define DEBUG_TYPE "test-loop-fusion" #define DEBUG_TYPE "test-loop-fusion"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
static llvm::cl::OptionCategory clOptionsCategory(DEBUG_TYPE " options"); static llvm::cl::OptionCategory clOptionsCategory(DEBUG_TYPE " options");
static llvm::cl::opt<bool> clTestDependenceCheck( static llvm::cl::opt<bool> clTestDependenceCheck(
"test-loop-fusion-dependence-check", "test-loop-fusion-dependence-check",
llvm::cl::desc("Enable testing of loop fusion dependence check"), llvm::cl::desc("Enable testing of loop fusion dependence check"),
llvm::cl::cat(clOptionsCategory)); llvm::cl::cat(clOptionsCategory));
Show All 24 Lines
// Run fusion dependence check on 'loops[i]' and 'loops[j]' at loop depths // Run fusion dependence check on 'loops[i]' and 'loops[j]' at loop depths
// in range ['loopDepth' + 1, 'maxLoopDepth']. // in range ['loopDepth' + 1, 'maxLoopDepth'].
// Emits a remark on 'loops[i]' if a fusion-preventing dependence exists. // Emits a remark on 'loops[i]' if a fusion-preventing dependence exists.
// Returns false as IR is not transformed. // Returns false as IR is not transformed.
static bool testDependenceCheck(AffineForOp srcForOp, AffineForOp dstForOp, static bool testDependenceCheck(AffineForOp srcForOp, AffineForOp dstForOp,
unsigned i, unsigned j, unsigned loopDepth, unsigned i, unsigned j, unsigned loopDepth,
unsigned maxLoopDepth) { unsigned maxLoopDepth) {
mlir::ComputationSliceState sliceUnion; affine::ComputationSliceState sliceUnion;
for (unsigned d = loopDepth + 1; d <= maxLoopDepth; ++d) { for (unsigned d = loopDepth + 1; d <= maxLoopDepth; ++d) {
FusionResult result = FusionResult result =
mlir::canFuseLoops(srcForOp, dstForOp, d, &sliceUnion); affine::canFuseLoops(srcForOp, dstForOp, d, &sliceUnion);
if (result.value == FusionResult::FailBlockDependence) { if (result.value == FusionResult::FailBlockDependence) {
srcForOp->emitRemark("block-level dependence preventing" srcForOp->emitRemark("block-level dependence preventing"
" fusion of loop nest ") " fusion of loop nest ")
<< i << " into loop nest " << j << " at depth " << loopDepth; << i << " into loop nest " << j << " at depth " << loopDepth;
} }
} }
return false; return false;
} }
// Returns the index of 'op' in its block. // Returns the index of 'op' in its block.
static unsigned getBlockIndex(Operation &op) { static unsigned getBlockIndex(Operation &op) {
unsigned index = 0; unsigned index = 0;
for (auto &opX : *op.getBlock()) { for (auto &opX : *op.getBlock()) {
if (&op == &opX) if (&op == &opX)
break; break;
++index; ++index;
} }
return index; return index;
} }
// Returns a string representation of 'sliceUnion'. // Returns a string representation of 'sliceUnion'.
static std::string getSliceStr(const mlir::ComputationSliceState &sliceUnion) { static std::string
getSliceStr(const affine::ComputationSliceState &sliceUnion) {
std::string result; std::string result;
llvm::raw_string_ostream os(result); llvm::raw_string_ostream os(result);
// Slice insertion point format [loop-depth, operation-block-index] // Slice insertion point format [loop-depth, operation-block-index]
unsigned ipd = getNestingDepth(&*sliceUnion.insertPoint); unsigned ipd = getNestingDepth(&*sliceUnion.insertPoint);
unsigned ipb = getBlockIndex(*sliceUnion.insertPoint); unsigned ipb = getBlockIndex(*sliceUnion.insertPoint);
os << "insert point: (" << std::to_string(ipd) << ", " << std::to_string(ipb) os << "insert point: (" << std::to_string(ipd) << ", " << std::to_string(ipb)
<< ")"; << ")";
assert(sliceUnion.lbs.size() == sliceUnion.ubs.size()); assert(sliceUnion.lbs.size() == sliceUnion.ubs.size());
Show All 12 Lines
/// in range ['loopDepth' + 1, 'maxLoopDepth']. /// in range ['loopDepth' + 1, 'maxLoopDepth'].
/// Emits a string representation of the slice union as a remark on 'loops[j]' /// Emits a string representation of the slice union as a remark on 'loops[j]'
/// and marks this as incorrect slice if the slice is invalid. Returns false as /// and marks this as incorrect slice if the slice is invalid. Returns false as
/// IR is not transformed. /// IR is not transformed.
static bool testSliceComputation(AffineForOp forOpA, AffineForOp forOpB, static bool testSliceComputation(AffineForOp forOpA, AffineForOp forOpB,
unsigned i, unsigned j, unsigned loopDepth, unsigned i, unsigned j, unsigned loopDepth,
unsigned maxLoopDepth) { unsigned maxLoopDepth) {
for (unsigned d = loopDepth + 1; d <= maxLoopDepth; ++d) { for (unsigned d = loopDepth + 1; d <= maxLoopDepth; ++d) {
mlir::ComputationSliceState sliceUnion; affine::ComputationSliceState sliceUnion;
FusionResult result = mlir::canFuseLoops(forOpA, forOpB, d, &sliceUnion); FusionResult result = affine::canFuseLoops(forOpA, forOpB, d, &sliceUnion);
if (result.value == FusionResult::Success) { if (result.value == FusionResult::Success) {
forOpB->emitRemark("slice (") forOpB->emitRemark("slice (")
<< " src loop: " << i << ", dst loop: " << j << ", depth: " << d << " src loop: " << i << ", dst loop: " << j << ", depth: " << d
<< " : " << getSliceStr(sliceUnion) << ")"; << " : " << getSliceStr(sliceUnion) << ")";
} else if (result.value == FusionResult::FailIncorrectSlice) { } else if (result.value == FusionResult::FailIncorrectSlice) {
forOpB->emitRemark("Incorrect slice (") forOpB->emitRemark("Incorrect slice (")
<< " src loop: " << i << ", dst loop: " << j << ", depth: " << d << " src loop: " << i << ", dst loop: " << j << ", depth: " << d
<< " : " << getSliceStr(sliceUnion) << ")"; << " : " << getSliceStr(sliceUnion) << ")";
} }
} }
return false; return false;
} }
// Attempts to fuse 'forOpA' into 'forOpB' at loop depths in range // Attempts to fuse 'forOpA' into 'forOpB' at loop depths in range
// ['loopDepth' + 1, 'maxLoopDepth']. // ['loopDepth' + 1, 'maxLoopDepth'].
// Returns true if loops were successfully fused, false otherwise. // Returns true if loops were successfully fused, false otherwise.
static bool testLoopFusionTransformation(AffineForOp forOpA, AffineForOp forOpB, static bool testLoopFusionTransformation(AffineForOp forOpA, AffineForOp forOpB,
unsigned i, unsigned j, unsigned i, unsigned j,
unsigned loopDepth, unsigned loopDepth,
unsigned maxLoopDepth) { unsigned maxLoopDepth) {
for (unsigned d = loopDepth + 1; d <= maxLoopDepth; ++d) { for (unsigned d = loopDepth + 1; d <= maxLoopDepth; ++d) {
mlir::ComputationSliceState sliceUnion; affine::ComputationSliceState sliceUnion;
FusionResult result = mlir::canFuseLoops(forOpA, forOpB, d, &sliceUnion); FusionResult result = affine::canFuseLoops(forOpA, forOpB, d, &sliceUnion);
if (result.value == FusionResult::Success) { if (result.value == FusionResult::Success) {
mlir::fuseLoops(forOpA, forOpB, sliceUnion); affine::fuseLoops(forOpA, forOpB, sliceUnion);
// Note: 'forOpA' is removed to simplify test output. A proper loop // Note: 'forOpA' is removed to simplify test output. A proper loop
// fusion pass should check the data dependence graph and run memref // fusion pass should check the data dependence graph and run memref
// region analysis to ensure removing 'forOpA' is safe. // region analysis to ensure removing 'forOpA' is safe.
forOpA.erase(); forOpA.erase();
return true; return true;
} }
} }
return false; return false;
▲ Show 20 Lines • Show All 58 Lines • Show Last 20 Lines

mlir/test/lib/Dialect/Affine/TestLoopMapping.cpp

Show All 12 Lines
#include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/LoopUtils.h" #include "mlir/Dialect/Affine/LoopUtils.h"
#include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/Builders.h" #include "mlir/IR/Builders.h"
#include "mlir/Pass/Pass.h" #include "mlir/Pass/Pass.h"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
namespace { namespace {
struct TestLoopMappingPass struct TestLoopMappingPass
: public PassWrapper<TestLoopMappingPass, OperationPass<>> { : public PassWrapper<TestLoopMappingPass, OperationPass<>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLoopMappingPass) MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLoopMappingPass)
StringRef getArgument() const final { StringRef getArgument() const final {
return "test-mapping-to-processing-elements"; return "test-mapping-to-processing-elements";
} }
StringRef getDescription() const final { StringRef getDescription() const final {
return "test mapping a single loop on a virtual processor grid"; return "test mapping a single loop on a virtual processor grid";
} }
explicit TestLoopMappingPass() = default; explicit TestLoopMappingPass() = default;
void getDependentDialects(DialectRegistry &registry) const override { void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<AffineDialect, scf::SCFDialect>(); registry.insert<affine::AffineDialect, scf::SCFDialect>();
} }
void runOnOperation() override { void runOnOperation() override {
// SSA values for the transformation are created out of thin air by // SSA values for the transformation are created out of thin air by
// unregistered "new_processor_id_and_range" operations. This is enough to // unregistered "new_processor_id_and_range" operations. This is enough to
// emulate mapping conditions. // emulate mapping conditions.
SmallVector<Value, 8> processorIds, numProcessors; SmallVector<Value, 8> processorIds, numProcessors;
getOperation()->walk([&processorIds, &numProcessors](Operation *op) { getOperation()->walk([&processorIds, &numProcessors](Operation *op) {
Show All 21 Lines

mlir/test/lib/Dialect/Affine/TestLoopPermutation.cpp

Show All 12 Lines
#include "mlir/Dialect/Affine/Analysis/Utils.h" #include "mlir/Dialect/Affine/Analysis/Utils.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/LoopUtils.h" #include "mlir/Dialect/Affine/LoopUtils.h"
#include "mlir/Pass/Pass.h" #include "mlir/Pass/Pass.h"
#define PASS_NAME "test-loop-permutation" #define PASS_NAME "test-loop-permutation"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
namespace { namespace {
/// This pass applies the permutation on the first maximal perfect nest. /// This pass applies the permutation on the first maximal perfect nest.
struct TestLoopPermutation struct TestLoopPermutation
: public PassWrapper<TestLoopPermutation, OperationPass<>> { : public PassWrapper<TestLoopPermutation, OperationPass<>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLoopPermutation) MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLoopPermutation)
▲ Show 20 Lines • Show All 43 Lines • Show Last 20 Lines

mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp

Show All 13 Lines
#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/PatternMatch.h" #include "mlir/IR/PatternMatch.h"
#include "mlir/Interfaces/ValueBoundsOpInterface.h" #include "mlir/Interfaces/ValueBoundsOpInterface.h"
#include "mlir/Pass/Pass.h" #include "mlir/Pass/Pass.h"
#define PASS_NAME "test-affine-reify-value-bounds" #define PASS_NAME "test-affine-reify-value-bounds"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
using mlir::presburger::BoundType; using mlir::presburger::BoundType;
namespace { namespace {
/// This pass applies the permutation on the first maximal perfect nest. /// This pass applies the permutation on the first maximal perfect nest.
struct TestReifyValueBounds struct TestReifyValueBounds
: public PassWrapper<TestReifyValueBounds, OperationPass<func::FuncOp>> { : public PassWrapper<TestReifyValueBounds, OperationPass<func::FuncOp>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestReifyValueBounds) MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestReifyValueBounds)
StringRef getArgument() const final { return PASS_NAME; } StringRef getArgument() const final { return PASS_NAME; }
StringRef getDescription() const final { StringRef getDescription() const final {
return "Tests ValueBoundsOpInterface with affine dialect reification"; return "Tests ValueBoundsOpInterface with affine dialect reification";
} }
TestReifyValueBounds() = default; TestReifyValueBounds() = default;
TestReifyValueBounds(const TestReifyValueBounds &pass) : PassWrapper(pass){}; TestReifyValueBounds(const TestReifyValueBounds &pass) : PassWrapper(pass){};
void getDependentDialects(DialectRegistry &registry) const override { void getDependentDialects(DialectRegistry &registry) const override {
registry registry.insert<affine::AffineDialect, tensor::TensorDialect,
.insert<AffineDialect, tensor::TensorDialect, memref::MemRefDialect>(); memref::MemRefDialect>();
} }
void runOnOperation() override; void runOnOperation() override;
private: private:
Option<bool> reifyToFuncArgs{ Option<bool> reifyToFuncArgs{
*this, "reify-to-func-args", *this, "reify-to-func-args",
llvm::cl::desc("Reify in terms of function args"), llvm::cl::init(false)}; llvm::cl::desc("Reify in terms of function args"), llvm::cl::init(false)};
▲ Show 20 Lines • Show All 139 Lines • Show Last 20 Lines

mlir/test/lib/Dialect/Affine/TestVectorizationUtils.cpp

Show All 27 Lines
#include "llvm/ADT/STLExtras.h" #include "llvm/ADT/STLExtras.h"
#include "llvm/Support/CommandLine.h" #include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h" #include "llvm/Support/Debug.h"
#define DEBUG_TYPE "affine-super-vectorizer-test" #define DEBUG_TYPE "affine-super-vectorizer-test"
using namespace mlir; using namespace mlir;
using namespace mlir::affine;
static llvm::cl::OptionCategory clOptionsCategory(DEBUG_TYPE " options"); static llvm::cl::OptionCategory clOptionsCategory(DEBUG_TYPE " options");
static llvm::cl::list<int> clTestVectorShapeRatio( static llvm::cl::list<int> clTestVectorShapeRatio(
"vector-shape-ratio", "vector-shape-ratio",
llvm::cl::desc("Specify the HW vector size for vectorization"), llvm::cl::desc("Specify the HW vector size for vectorization"),
llvm::cl::cat(clOptionsCategory)); llvm::cl::cat(clOptionsCategory));
static llvm::cl::opt<bool> clTestForwardSlicingAnalysis( static llvm::cl::opt<bool> clTestForwardSlicingAnalysis(
▲ Show 20 Lines • Show All 50 Lines • ▼ Show 20 Linesstruct VectorizerTestPass
/// Test for 'vectorizeAffineLoopNest' utility. /// Test for 'vectorizeAffineLoopNest' utility.
void testVecAffineLoopNest(llvm::raw_ostream &outs); void testVecAffineLoopNest(llvm::raw_ostream &outs);
}; };
} // namespace } // namespace
void VectorizerTestPass::testVectorShapeRatio(llvm::raw_ostream &outs) { void VectorizerTestPass::testVectorShapeRatio(llvm::raw_ostream &outs) {
auto f = getOperation(); auto f = getOperation();
using matcher::Op; using affine::matcher::Op;
SmallVector<int64_t, 8> shape(clTestVectorShapeRatio.begin(), SmallVector<int64_t, 8> shape(clTestVectorShapeRatio.begin(),
clTestVectorShapeRatio.end()); clTestVectorShapeRatio.end());
auto subVectorType = auto subVectorType =
VectorType::get(shape, FloatType::getF32(f.getContext())); VectorType::get(shape, FloatType::getF32(f.getContext()));
// Only filter operations that operate on a strict super-vector and have one // Only filter operations that operate on a strict super-vector and have one
// return. This makes testing easier. // return. This makes testing easier.
auto filter = [&](Operation &op) { auto filter = [&](Operation &op) {
assert(subVectorType.getElementType().isF32() && assert(subVectorType.getElementType().isF32() &&
"Only f32 supported for now"); "Only f32 supported for now");
if (!matcher::operatesOnSuperVectorsOf(op, subVectorType)) { if (!mlir::matcher::operatesOnSuperVectorsOf(op, subVectorType)) {
return false; return false;
} }
if (op.getNumResults() != 1) { if (op.getNumResults() != 1) {
return false; return false;
} }
return true; return true;
}; };
auto pat = Op(filter); auto pat = Op(filter);
Show All 13 Lines for (auto m : matches) {
} else { } else {
outs << "\nmatched: " << *opInst << " with shape ratio: "; outs << "\nmatched: " << *opInst << " with shape ratio: ";
llvm::interleaveComma(MutableArrayRef<int64_t>(*ratio), outs); llvm::interleaveComma(MutableArrayRef<int64_t>(*ratio), outs);
} }
} }
} }
static NestedPattern patternTestSlicingOps() { static NestedPattern patternTestSlicingOps() {
using matcher::Op; using affine::matcher::Op;
// Match all operations with the kTestSlicingOpName name. // Match all operations with the kTestSlicingOpName name.
auto filter = [](Operation &op) { auto filter = [](Operation &op) {
// Just use a custom op name for this test, it makes life easier. // Just use a custom op name for this test, it makes life easier.
return op.getName().getStringRef() == "slicing-test-op"; return op.getName().getStringRef() == "slicing-test-op";
}; };
return Op(filter); return Op(filter);
} }
▲ Show 20 Lines • Show All 46 Lines • ▼ Show 20 Lines
static bool customOpWithAffineMapAttribute(Operation &op) { static bool customOpWithAffineMapAttribute(Operation &op) {
return op.getName().getStringRef() == return op.getName().getStringRef() ==
VectorizerTestPass::kTestAffineMapOpName; VectorizerTestPass::kTestAffineMapOpName;
} }
void VectorizerTestPass::testComposeMaps(llvm::raw_ostream &outs) { void VectorizerTestPass::testComposeMaps(llvm::raw_ostream &outs) {
auto f = getOperation(); auto f = getOperation();
using matcher::Op; using affine::matcher::Op;
auto pattern = Op(customOpWithAffineMapAttribute); auto pattern = Op(customOpWithAffineMapAttribute);
SmallVector<NestedMatch, 8> matches; SmallVector<NestedMatch, 8> matches;
pattern.match(f, &matches); pattern.match(f, &matches);
SmallVector<AffineMap, 4> maps; SmallVector<AffineMap, 4> maps;
maps.reserve(matches.size()); maps.reserve(matches.size());
for (auto m : llvm::reverse(matches)) { for (auto m : llvm::reverse(matches)) {
auto *opInst = m.getMatchedOperation(); auto *opInst = m.getMatchedOperation();
auto map = opInst->getAttr(VectorizerTestPass::kTestAffineMapAttrName) auto map = opInst->getAttr(VectorizerTestPass::kTestAffineMapAttrName)
▲ Show 20 Lines • Show All 79 Lines • Show Last 20 Lines

mlir/test/lib/Dialect/GPU/TestGpuMemoryPromotion.cpp

Show All 27 Lines
/// does not check whether the promotion is legal (e.g., amount of memory used) /// does not check whether the promotion is legal (e.g., amount of memory used)
/// or beneficial (e.g., makes previously uncoalesced loads coalesced). /// or beneficial (e.g., makes previously uncoalesced loads coalesced).
struct TestGpuMemoryPromotionPass struct TestGpuMemoryPromotionPass
: public PassWrapper<TestGpuMemoryPromotionPass, : public PassWrapper<TestGpuMemoryPromotionPass,
OperationPass<gpu::GPUFuncOp>> { OperationPass<gpu::GPUFuncOp>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestGpuMemoryPromotionPass) MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestGpuMemoryPromotionPass)
void getDependentDialects(DialectRegistry &registry) const override { void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<AffineDialect, memref::MemRefDialect, scf::SCFDialect>(); registry.insert<affine::AffineDialect, memref::MemRefDialect,
scf::SCFDialect>();
} }
StringRef getArgument() const final { return "test-gpu-memory-promotion"; } StringRef getArgument() const final { return "test-gpu-memory-promotion"; }
StringRef getDescription() const final { StringRef getDescription() const final {
return "Promotes the annotated arguments of gpu.func to workgroup memory."; return "Promotes the annotated arguments of gpu.func to workgroup memory.";
} }
void runOnOperation() override { void runOnOperation() override {
gpu::GPUFuncOp op = getOperation(); gpu::GPUFuncOp op = getOperation();
Show All 13 Lines

mlir/test/lib/Dialect/Linalg/TestDataLayoutPropagation.cpp

Show All 12 Lines
using namespace mlir; using namespace mlir;
namespace { namespace {
struct TestDataLayoutPropagationPass struct TestDataLayoutPropagationPass
: public PassWrapper<TestDataLayoutPropagationPass, OperationPass<>> { : public PassWrapper<TestDataLayoutPropagationPass, OperationPass<>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestDataLayoutPropagationPass) MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestDataLayoutPropagationPass)
void getDependentDialects(DialectRegistry &registry) const override { void getDependentDialects(DialectRegistry &registry) const override {
registry registry.insert<affine::AffineDialect, linalg::LinalgDialect,
.insert<AffineDialect, linalg::LinalgDialect, tensor::TensorDialect>(); tensor::TensorDialect>();
} }
StringRef getArgument() const final { StringRef getArgument() const final {
return "test-linalg-data-layout-propagation"; return "test-linalg-data-layout-propagation";
} }
StringRef getDescription() const final { StringRef getDescription() const final {
return "Test data layout propagation"; return "Test data layout propagation";
} }
Show All 20 Lines

mlir/test/lib/Dialect/Linalg/TestLinalgDecomposeOps.cpp

Show All 20 Lines
struct TestLinalgDecomposeOps struct TestLinalgDecomposeOps
: public PassWrapper<TestLinalgDecomposeOps, OperationPass<>> { : public PassWrapper<TestLinalgDecomposeOps, OperationPass<>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLinalgDecomposeOps) MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLinalgDecomposeOps)
TestLinalgDecomposeOps() = default; TestLinalgDecomposeOps() = default;
TestLinalgDecomposeOps(const TestLinalgDecomposeOps &pass) TestLinalgDecomposeOps(const TestLinalgDecomposeOps &pass)
: PassWrapper(pass){}; : PassWrapper(pass){};
void getDependentDialects(DialectRegistry &registry) const override { void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<AffineDialect, linalg::LinalgDialect>(); registry.insert<affine::AffineDialect, linalg::LinalgDialect>();
} }
StringRef getArgument() const final { return "test-linalg-decompose-ops"; } StringRef getArgument() const final { return "test-linalg-decompose-ops"; }
StringRef getDescription() const final { StringRef getDescription() const final {
return "Test Linalg decomposition patterns"; return "Test Linalg decomposition patterns";
} }
Option<bool> removeDeadArgsAndResults{ Option<bool> removeDeadArgsAndResults{
*this, "remove-dead-args-and-results", *this, "remove-dead-args-and-results",
Show All 23 Lines

mlir/test/lib/Dialect/Linalg/TestLinalgElementwiseFusion.cpp

Show First 20 Lines • Show All 86 Lines • ▼ Show 20 Linesstruct TestLinalgElementwiseFusion
: public PassWrapper<TestLinalgElementwiseFusion, : public PassWrapper<TestLinalgElementwiseFusion,
OperationPass<func::FuncOp>> { OperationPass<func::FuncOp>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLinalgElementwiseFusion) MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLinalgElementwiseFusion)
TestLinalgElementwiseFusion() = default; TestLinalgElementwiseFusion() = default;
TestLinalgElementwiseFusion(const TestLinalgElementwiseFusion &pass) TestLinalgElementwiseFusion(const TestLinalgElementwiseFusion &pass)
: PassWrapper(pass) {} : PassWrapper(pass) {}
void getDependentDialects(DialectRegistry &registry) const override { void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<AffineDialect, linalg::LinalgDialect, memref::MemRefDialect, registry.insert<affine::AffineDialect, linalg::LinalgDialect,
tensor::TensorDialect>(); memref::MemRefDialect, tensor::TensorDialect>();
} }
StringRef getArgument() const final { StringRef getArgument() const final {
return "test-linalg-elementwise-fusion-patterns"; return "test-linalg-elementwise-fusion-patterns";
} }
StringRef getDescription() const final { StringRef getDescription() const final {
return "Test Linalg element wise operation fusion patterns"; return "Test Linalg element wise operation fusion patterns";
} }
▲ Show 20 Lines • Show All 180 Lines • Show Last 20 Lines

mlir/test/lib/Dialect/Linalg/TestLinalgFusionTransforms.cpp

Show First 20 Lines • Show All 67 Lines • ▼ Show 20 Lines
} }
namespace { namespace {
struct TestLinalgGreedyFusion struct TestLinalgGreedyFusion
: public PassWrapper<TestLinalgGreedyFusion, OperationPass<func::FuncOp>> { : public PassWrapper<TestLinalgGreedyFusion, OperationPass<func::FuncOp>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLinalgGreedyFusion) MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLinalgGreedyFusion)
void getDependentDialects(DialectRegistry &registry) const override { void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<AffineDialect, linalg::LinalgDialect, memref::MemRefDialect, registry.insert<affine::AffineDialect, linalg::LinalgDialect,
scf::SCFDialect>(); memref::MemRefDialect, scf::SCFDialect>();
} }
StringRef getArgument() const final { return "test-linalg-greedy-fusion"; } StringRef getArgument() const final { return "test-linalg-greedy-fusion"; }
StringRef getDescription() const final { StringRef getDescription() const final {
return "Test Linalg fusion by applying a greedy test transformation."; return "Test Linalg fusion by applying a greedy test transformation.";
} }
void runOnOperation() override { void runOnOperation() override {
MLIRContext *context = &getContext(); MLIRContext *context = &getContext();
RewritePatternSet patterns = RewritePatternSet patterns =
Show All 25 Lines

mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp

Show All 33 Linesstruct TestLinalgTransforms
: public PassWrapper<TestLinalgTransforms, OperationPass<func::FuncOp>> { : public PassWrapper<TestLinalgTransforms, OperationPass<func::FuncOp>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLinalgTransforms) MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLinalgTransforms)
TestLinalgTransforms() = default; TestLinalgTransforms() = default;
TestLinalgTransforms(const TestLinalgTransforms &pass) : PassWrapper(pass) {} TestLinalgTransforms(const TestLinalgTransforms &pass) : PassWrapper(pass) {}
void getDependentDialects(DialectRegistry &registry) const override { void getDependentDialects(DialectRegistry &registry) const override {
// clang-format off // clang-format off
registry.insert<AffineDialect, registry.insert<affine::AffineDialect,
bufferization::BufferizationDialect, bufferization::BufferizationDialect,
memref::MemRefDialect, memref::MemRefDialect,
scf::SCFDialect, scf::SCFDialect,
linalg::LinalgDialect, linalg::LinalgDialect,
vector::VectorDialect, vector::VectorDialect,
gpu::GPUDialect>(); gpu::GPUDialect>();
// clang-format on // clang-format on
} }
▲ Show 20 Lines • Show All 204 Lines • Show Last 20 Lines

mlir/test/lib/Dialect/Linalg/TestPadFusion.cpp

Show All 19 Lines
using namespace mlir; using namespace mlir;
namespace { namespace {
struct TestPadFusionPass struct TestPadFusionPass
: public PassWrapper<TestPadFusionPass, OperationPass<>> { : public PassWrapper<TestPadFusionPass, OperationPass<>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestPadFusionPass) MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestPadFusionPass)
void getDependentDialects(DialectRegistry &registry) const override { void getDependentDialects(DialectRegistry &registry) const override {
registry registry.insert<affine::AffineDialect, linalg::LinalgDialect,
.insert<AffineDialect, linalg::LinalgDialect, tensor::TensorDialect>(); tensor::TensorDialect>();
} }
StringRef getArgument() const final { return "test-linalg-pad-fusion"; } StringRef getArgument() const final { return "test-linalg-pad-fusion"; }
StringRef getDescription() const final { return "Test PadOp fusion"; } StringRef getDescription() const final { return "Test PadOp fusion"; }
void runOnOperation() override { void runOnOperation() override {
MLIRContext *context = &getContext(); MLIRContext *context = &getContext();
RewritePatternSet patterns(context); RewritePatternSet patterns(context);
Show All 13 Lines

mlir/test/lib/Dialect/MemRef/TestComposeSubView.cpp

Show All 26 Lines StringRef getDescription() const final {
return "Test combining composed subviews"; return "Test combining composed subviews";
} }
void runOnOperation() override; void runOnOperation() override;
void getDependentDialects(DialectRegistry &registry) const override; void getDependentDialects(DialectRegistry &registry) const override;
}; };
void TestComposeSubViewPass::getDependentDialects( void TestComposeSubViewPass::getDependentDialects(
DialectRegistry &registry) const { DialectRegistry &registry) const {
registry.insert<AffineDialect>(); registry.insert<affine::AffineDialect>();
} }
void TestComposeSubViewPass::runOnOperation() { void TestComposeSubViewPass::runOnOperation() {
RewritePatternSet patterns(&getContext()); RewritePatternSet patterns(&getContext());
memref::populateComposeSubViewPatterns(patterns, &getContext()); memref::populateComposeSubViewPatterns(patterns, &getContext());
(void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns)); (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
} }
} // namespace } // namespace
namespace mlir { namespace mlir {
namespace test { namespace test {
void registerTestComposeSubView() { void registerTestComposeSubView() {
PassRegistration<TestComposeSubViewPass>(); PassRegistration<TestComposeSubViewPass>();
} }
} // namespace test } // namespace test
} // namespace mlir } // namespace mlir

mlir/test/lib/Dialect/MemRef/TestMultiBuffer.cpp

Show All 18 Lines
struct TestMultiBufferingPass struct TestMultiBufferingPass
: public PassWrapper<TestMultiBufferingPass, OperationPass<>> { : public PassWrapper<TestMultiBufferingPass, OperationPass<>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestMultiBufferingPass) MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestMultiBufferingPass)
TestMultiBufferingPass() = default; TestMultiBufferingPass() = default;
TestMultiBufferingPass(const TestMultiBufferingPass &pass) TestMultiBufferingPass(const TestMultiBufferingPass &pass)
: PassWrapper(pass) {} : PassWrapper(pass) {}
void getDependentDialects(DialectRegistry &registry) const override { void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<AffineDialect>(); registry.insert<affine::AffineDialect>();
} }
StringRef getArgument() const final { return "test-multi-buffering"; } StringRef getArgument() const final { return "test-multi-buffering"; }
StringRef getDescription() const final { StringRef getDescription() const final {
return "Test multi buffering transformation"; return "Test multi buffering transformation";
} }
void runOnOperation() override; void runOnOperation() override;
Option<unsigned> multiplier{ Option<unsigned> multiplier{
*this, "multiplier", *this, "multiplier",
Show All 21 Lines

mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp

Show First 20 Lines • Show All 47 Lines • ▼ Show 20 Linesstruct TestVectorToVectorLowering
StringRef getArgument() const final { StringRef getArgument() const final {
return "test-vector-to-vector-lowering"; return "test-vector-to-vector-lowering";
} }
StringRef getDescription() const final { StringRef getDescription() const final {
return "Test lowering patterns between ops in the vector dialect"; return "Test lowering patterns between ops in the vector dialect";
} }
void getDependentDialects(DialectRegistry &registry) const override { void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<AffineDialect>(); registry.insert<affine::AffineDialect>();
} }
Option<bool> unroll{*this, "unroll", llvm::cl::desc("Include unrolling"), Option<bool> unroll{*this, "unroll", llvm::cl::desc("Include unrolling"),
llvm::cl::init(false)}; llvm::cl::init(false)};
void runOnOperation() override { void runOnOperation() override {
auto *ctx = &getContext(); auto *ctx = &getContext();
RewritePatternSet patterns(ctx); RewritePatternSet patterns(ctx);
▲ Show 20 Lines • Show All 58 Lines • ▼ Show 20 Lines StringRef getArgument() const final {
return "test-vector-contraction-prepare-for-mmt-lowering"; return "test-vector-contraction-prepare-for-mmt-lowering";
} }
StringRef getDescription() const final { StringRef getDescription() const final {
return "Test vector.contraction matmul canonicalization for MMT lowering."; return "Test vector.contraction matmul canonicalization for MMT lowering.";
} }
TestVectorContractionPrepareForMMTLowering() = default; TestVectorContractionPrepareForMMTLowering() = default;
void getDependentDialects(DialectRegistry &registry) const override { void getDependentDialects(DialectRegistry &registry) const override {
registry registry.insert<affine::AffineDialect, arith::ArithDialect,
.insert<AffineDialect, arith::ArithDialect, vector::VectorDialect>(); vector::VectorDialect>();
} }
void runOnOperation() override { void runOnOperation() override {
MLIRContext *ctx = &getContext(); MLIRContext *ctx = &getContext();
RewritePatternSet patterns(ctx); RewritePatternSet patterns(ctx);
vector::populateVectorContractCanonicalizeMatmulToMMT(patterns); vector::populateVectorContractCanonicalizeMatmulToMMT(patterns);
(void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns)); (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
} }
▲ Show 20 Lines • Show All 100 Lines • ▼ Show 20 Lines MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
TestVectorTransferUnrollingPatterns) TestVectorTransferUnrollingPatterns)
TestVectorTransferUnrollingPatterns() = default; TestVectorTransferUnrollingPatterns() = default;
TestVectorTransferUnrollingPatterns( TestVectorTransferUnrollingPatterns(
const TestVectorTransferUnrollingPatterns &pass) const TestVectorTransferUnrollingPatterns &pass)
: PassWrapper(pass) {} : PassWrapper(pass) {}
void getDependentDialects(DialectRegistry &registry) const override { void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<AffineDialect>(); registry.insert<affine::AffineDialect>();
} }
StringRef getArgument() const final { StringRef getArgument() const final {
return "test-vector-transfer-unrolling-patterns"; return "test-vector-transfer-unrolling-patterns";
} }
StringRef getDescription() const final { StringRef getDescription() const final {
return "Test lowering patterns to unroll transfer ops in the vector " return "Test lowering patterns to unroll transfer ops in the vector "
"dialect"; "dialect";
} }
▲ Show 20 Lines • Show All 42 Lines • ▼ Show 20 Lines StringRef getArgument() const final {
return "test-scalar-vector-transfer-lowering"; return "test-scalar-vector-transfer-lowering";
} }
StringRef getDescription() const final { StringRef getDescription() const final {
return "Test lowering of scalar vector transfers to memref loads/stores."; return "Test lowering of scalar vector transfers to memref loads/stores.";
} }
TestScalarVectorTransferLoweringPatterns() = default; TestScalarVectorTransferLoweringPatterns() = default;
void getDependentDialects(DialectRegistry &registry) const override { void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<AffineDialect, memref::MemRefDialect, tensor::TensorDialect, registry.insert<affine::AffineDialect, memref::MemRefDialect,
vector::VectorDialect>(); tensor::TensorDialect, vector::VectorDialect>();
} }
void runOnOperation() override { void runOnOperation() override {
MLIRContext *ctx = &getContext(); MLIRContext *ctx = &getContext();
RewritePatternSet patterns(ctx); RewritePatternSet patterns(ctx);
vector::populateScalarVectorTransferLoweringPatterns(patterns); vector::populateScalarVectorTransferLoweringPatterns(patterns);
(void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns)); (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
} }
Show All 19 Linesstruct TestVectorTransferCollapseInnerMostContiguousDims
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID( MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
TestVectorTransferCollapseInnerMostContiguousDims) TestVectorTransferCollapseInnerMostContiguousDims)
TestVectorTransferCollapseInnerMostContiguousDims() = default; TestVectorTransferCollapseInnerMostContiguousDims() = default;
TestVectorTransferCollapseInnerMostContiguousDims( TestVectorTransferCollapseInnerMostContiguousDims(
const TestVectorTransferCollapseInnerMostContiguousDims &pass) = default; const TestVectorTransferCollapseInnerMostContiguousDims &pass) = default;
void getDependentDialects(DialectRegistry &registry) const override { void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<memref::MemRefDialect, AffineDialect>(); registry.insert<memref::MemRefDialect, affine::AffineDialect>();
} }
StringRef getArgument() const final { StringRef getArgument() const final {
return "test-vector-transfer-collapse-inner-most-dims"; return "test-vector-transfer-collapse-inner-most-dims";
} }
StringRef getDescription() const final { StringRef getDescription() const final {
return "Test lowering patterns that reducedes the rank of the vector " return "Test lowering patterns that reducedes the rank of the vector "
▲ Show 20 Lines • Show All 129 Lines • ▼ Show 20 Lines
} }
struct TestVectorDistribution struct TestVectorDistribution
: public PassWrapper<TestVectorDistribution, OperationPass<func::FuncOp>> { : public PassWrapper<TestVectorDistribution, OperationPass<func::FuncOp>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorDistribution) MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorDistribution)
void getDependentDialects(DialectRegistry &registry) const override { void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<scf::SCFDialect, memref::MemRefDialect, gpu::GPUDialect, registry.insert<scf::SCFDialect, memref::MemRefDialect, gpu::GPUDialect,
AffineDialect>(); affine::AffineDialect>();
} }
StringRef getArgument() const final { return "test-vector-warp-distribute"; } StringRef getArgument() const final { return "test-vector-warp-distribute"; }
StringRef getDescription() const final { StringRef getDescription() const final {
return "Test vector warp distribute transformation and lowering patterns"; return "Test vector warp distribute transformation and lowering patterns";
} }
TestVectorDistribution() = default; TestVectorDistribution() = default;
TestVectorDistribution(const TestVectorDistribution &pass) TestVectorDistribution(const TestVectorDistribution &pass)
▲ Show 20 Lines • Show All 195 Lines • Show Last 20 Lines

mlir/test/lib/Interfaces/TilingInterface/TestTilingInterface.cpp

Show First 20 Lines • Show All 395 Lines • ▼ Show 20 Lines
struct TestTilingInterfacePass struct TestTilingInterfacePass
: public PassWrapper<TestTilingInterfacePass, OperationPass<func::FuncOp>> { : public PassWrapper<TestTilingInterfacePass, OperationPass<func::FuncOp>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestTilingInterfacePass) MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestTilingInterfacePass)
TestTilingInterfacePass() = default; TestTilingInterfacePass() = default;
TestTilingInterfacePass(const TestTilingInterfacePass &pass) TestTilingInterfacePass(const TestTilingInterfacePass &pass)
: PassWrapper(pass) {} : PassWrapper(pass) {}
void getDependentDialects(DialectRegistry &registry) const override { void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<AffineDialect, linalg::LinalgDialect, memref::MemRefDialect, registry.insert<affine::AffineDialect, linalg::LinalgDialect,
scf::SCFDialect, tensor::TensorDialect>(); memref::MemRefDialect, scf::SCFDialect,
tensor::TensorDialect>();
linalg::registerTilingInterfaceExternalModels(registry); linalg::registerTilingInterfaceExternalModels(registry);
tensor::registerTilingInterfaceExternalModels(registry); tensor::registerTilingInterfaceExternalModels(registry);
} }
StringRef getArgument() const final { return "test-tiling-interface"; } StringRef getArgument() const final { return "test-tiling-interface"; }
StringRef getDescription() const final { StringRef getDescription() const final {
return "Test tiling using TilingInterface"; return "Test tiling using TilingInterface";
} }
▲ Show 20 Lines • Show All 146 Lines • Show Last 20 Lines

mlir/unittests/Analysis/Presburger/Parser.h

Show All 24 Lines
namespace mlir { namespace mlir {
namespace presburger { namespace presburger {
/// Parses an IntegerPolyhedron from a StringRef. It is expected that the string /// Parses an IntegerPolyhedron from a StringRef. It is expected that the string
/// represents a valid IntegerSet. /// represents a valid IntegerSet.
inline IntegerPolyhedron parseIntegerPolyhedron(StringRef str) { inline IntegerPolyhedron parseIntegerPolyhedron(StringRef str) {
MLIRContext context(MLIRContext::Threading::DISABLED); MLIRContext context(MLIRContext::Threading::DISABLED);
return FlatAffineValueConstraints(parseIntegerSet(str, &context)); return affine::FlatAffineValueConstraints(parseIntegerSet(str, &context));
} }
/// Parse a list of StringRefs to IntegerRelation and combine them into a /// Parse a list of StringRefs to IntegerRelation and combine them into a
/// PresburgerSet by using the union operation. It is expected that the strings /// PresburgerSet by using the union operation. It is expected that the strings
/// are all valid IntegerSet representation and that all of them have compatible /// are all valid IntegerSet representation and that all of them have compatible
/// spaces. /// spaces.
inline PresburgerSet parsePresburgerSet(ArrayRef<StringRef> strs) { inline PresburgerSet parsePresburgerSet(ArrayRef<StringRef> strs) {
assert(!strs.empty() && "strs should not be empty"); assert(!strs.empty() && "strs should not be empty");
▲ Show 20 Lines • Show All 45 Lines • Show Last 20 Lines