This is an archive of the discontinued LLVM Phabricator instance.

Differential D103161

[mlir] Add `distributionTypes` to LinalgTilingOptions.
ClosedPublic

Authored by pifon2a on May 26 2021, 5:36 AM.

Details

Reviewers
nicolasvasilache
Commits
rG74a89cba8ced: [mlir] Add `distributionTypes` to LinalgTilingOptions.

Diff Detail

Event Timeline

pifon2a created this revision.May 26 2021, 5:36 AM
Herald added subscribers: dcaballe, cota, mravishankar and 17 others. · View Herald TranscriptMay 26 2021, 5:36 AM
pifon2a requested review of this revision.May 26 2021, 5:36 AM
Herald added a project: Restricted Project. · View Herald TranscriptMay 26 2021, 5:36 AM
Herald added subscribers: limo1996, stephenneuendorffer. · View Herald Transcript
nicolasvasilache accepted this revision.May 26 2021, 5:43 AM
This revision is now accepted and ready to land.May 26 2021, 5:43 AM
Harbormaster completed remote builds in B106277: Diff 347938.May 26 2021, 6:09 AM
This revision was landed with ongoing or failed builds.May 26 2021, 8:51 AM
Closed by commit rG74a89cba8ced: [mlir] Add `distributionTypes` to LinalgTilingOptions. (authored by pifon2a). · Explain Why
This revision was automatically updated to reflect the committed changes.
pifon2a added a commit: rG74a89cba8ced: [mlir] Add `distributionTypes` to LinalgTilingOptions..
hanchung added a subscriber: hanchung.Jun 18 2021, 1:25 PM
hanchung added inline comments.
mlir/test/Dialect/Linalg/tile-tensors.mlir
132–153

Do we intend to delete the test?

Revision Contents

PathSize
mlir/
include/
mlir/
Dialect/
Linalg/
IR/
6 lines
3 lines
3 lines
Transforms/
8 lines
Utils/
3 lines
lib/
Dialect/
Linalg/
Transforms/
36 lines
Utils/
13 lines
test/
Dialect/
Linalg/
29 lines

Diff 347985

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

Show First 20 Lines • Show All 609 Lines • ▼ Show 20 Lines let extraClassDeclaration = [{
} }
ValueRange getRegionInputArgs() { ValueRange getRegionInputArgs() {
return getBody()->getArguments().slice(getNumLoops(), inputs().size()); return getBody()->getArguments().slice(getNumLoops(), inputs().size());
} }
ValueRange getRegionOutputArgs() { ValueRange getRegionOutputArgs() {
return getBody()->getArguments().take_back(outputs().size()); return getBody()->getArguments().take_back(outputs().size());
} }
void setDistributionTypes(Builder& b, ArrayRef<StringRef> types) {
assert(types.size() == getNumLoops() &&
"expected distribution type for every dimension");
distribution_typesAttr(b.getStrArrayAttr(types));
}
void setLowerBounds(ValueRange lowerBounds) { void setLowerBounds(ValueRange lowerBounds) {
unsigned numLoops = getNumLoops(); unsigned numLoops = getNumLoops();
assert(lowerBounds.size() == numLoops && assert(lowerBounds.size() == numLoops &&
"expected lower bounds for every loop dimension"); "expected lower bounds for every loop dimension");
for (unsigned i = 0; i < numLoops; ++i) for (unsigned i = 0; i < numLoops; ++i)
setOperand(i, lowerBounds[i]); setOperand(i, lowerBounds[i]);
} }
▲ Show 20 Lines • Show All 166 Lines • Show Last 20 Lines

mlir/include/mlir/Dialect/Linalg/Passes.h

Show All 24 Lines
std::unique_ptr<OperationPass<FuncOp>> std::unique_ptr<OperationPass<FuncOp>>
createLinalgTilingPass(ArrayRef<int64_t> tileSizes = {}); createLinalgTilingPass(ArrayRef<int64_t> tileSizes = {});
std::unique_ptr<OperationPass<FuncOp>> std::unique_ptr<OperationPass<FuncOp>>
createLinalgTilingToParallelLoopsPass(ArrayRef<int64_t> tileSizes = {}); createLinalgTilingToParallelLoopsPass(ArrayRef<int64_t> tileSizes = {});
std::unique_ptr<OperationPass<FuncOp>> std::unique_ptr<OperationPass<FuncOp>>
createLinalgTilingToTiledLoopPass(ArrayRef<int64_t> tileSizes = {}); createLinalgTilingToTiledLoopPass(ArrayRef<int64_t> tileSizes = {},
ArrayRef<StringRef> distributionTypes = {});
std::unique_ptr<OperationPass<FuncOp>> std::unique_ptr<OperationPass<FuncOp>>
createLinalgPromotionPass(bool dynamicBuffers, bool useAlloca); createLinalgPromotionPass(bool dynamicBuffers, bool useAlloca);
std::unique_ptr<OperationPass<FuncOp>> createLinalgPromotionPass(); std::unique_ptr<OperationPass<FuncOp>> createLinalgPromotionPass();
std::unique_ptr<OperationPass<FuncOp>> createLinalgInlineScalarOperandsPass(); std::unique_ptr<OperationPass<FuncOp>> createLinalgInlineScalarOperandsPass();
/// Create a pass to convert Linalg tiled loops to `scf.for` and `scf.parallel` /// Create a pass to convert Linalg tiled loops to `scf.for` and `scf.parallel`
▲ Show 20 Lines • Show All 45 Lines • Show Last 20 Lines

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

Show First 20 Lines • Show All 176 Lines • ▼ Show 20 Lines
} }
def LinalgTilingToTiledLoops def LinalgTilingToTiledLoops
: FunctionPass<"linalg-tile-to-tiled-loop"> { : FunctionPass<"linalg-tile-to-tiled-loop"> {
let summary = "Tile operations in the linalg dialect to linalg.tiled_loop"; let summary = "Tile operations in the linalg dialect to linalg.tiled_loop";
let constructor = "mlir::createLinalgTilingToTiledLoopPass()"; let constructor = "mlir::createLinalgTilingToTiledLoopPass()";
let options = [ let options = [
ListOption<"tileSizes", "linalg-tile-sizes", "int64_t", "Tile sizes", ListOption<"tileSizes", "linalg-tile-sizes", "int64_t", "Tile sizes",
"llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated">,
ListOption<"distributionTypes", "linalg-distribution-types", "std::string",
"DistributionTypes",
"llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated"> "llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated">
]; ];
let dependentDialects = [ let dependentDialects = [
"AffineDialect", "AffineDialect",
"linalg::LinalgDialect", "linalg::LinalgDialect",
"memref::MemRefDialect", "memref::MemRefDialect",
"scf::SCFDialect" "scf::SCFDialect"
]; ];
Show All 37 Lines

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

Show First 20 Lines • Show All 487 Lines • ▼ Show 20 Linesstruct LinalgTilingOptions {
Optional<LinalgLoopDistributionOptions> distribution = None; Optional<LinalgLoopDistributionOptions> distribution = None;
LinalgTilingOptions & LinalgTilingOptions &
setDistributionOptions(LinalgLoopDistributionOptions distributionOptions) { setDistributionOptions(LinalgLoopDistributionOptions distributionOptions) {
distribution = std::move(distributionOptions); distribution = std::move(distributionOptions);
return *this; return *this;
} }
/// Specification markers of how to distribute the `linalg.tiled_loop`.
SmallVector<StringRef, 2> distributionTypes = {};
LinalgTilingOptions &setDistributionTypes(ArrayRef<StringRef> types) {
distributionTypes.assign(types.begin(), types.end());
return *this;
}
/// Computation function that returns a padding value to use when padding to /// Computation function that returns a padding value to use when padding to
/// force static sizes. When `paddingValueComputationFunction` is set, padding /// force static sizes. When `paddingValueComputationFunction` is set, padding
/// operations are introduced, that guarantee the underlying op is statically /// operations are introduced, that guarantee the underlying op is statically
/// shaped and can thus be vectorized. /// shaped and can thus be vectorized.
PaddingValueComputationFunction paddingValueComputationFunction = nullptr; PaddingValueComputationFunction paddingValueComputationFunction = nullptr;
LinalgTilingOptions & LinalgTilingOptions &
setPaddingValueComputationFunction(PaddingValueComputationFunction fun) { setPaddingValueComputationFunction(PaddingValueComputationFunction fun) {
▲ Show 20 Lines • Show All 555 Lines • Show Last 20 Lines

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

Show First 20 Lines • Show All 241 Lines • ▼ Show 20 Lines
/// of loop induction variables and a range of iterArgs. /// of loop induction variables and a range of iterArgs.
template <typename LoopTy> template <typename LoopTy>
struct GenerateLoopNest { struct GenerateLoopNest {
static void doit(OpBuilder &b, Location loc, ArrayRef<Range> loopRanges, static void doit(OpBuilder &b, Location loc, ArrayRef<Range> loopRanges,
LinalgOp linalgOp, ArrayRef<Attribute> iteratorTypes, LinalgOp linalgOp, ArrayRef<Attribute> iteratorTypes,
function_ref<scf::ValueVector(OpBuilder &, Location, function_ref<scf::ValueVector(OpBuilder &, Location,
ValueRange, ValueRange)> ValueRange, ValueRange)>
bodyBuilderFn, bodyBuilderFn,
Optional<LinalgLoopDistributionOptions> = None); Optional<LinalgLoopDistributionOptions> = None,
ArrayRef<StringRef> distributionTypes = {});
}; };
} // namespace linalg } // namespace linalg
} // namespace mlir } // namespace mlir
#endif // MLIR_DIALECT_LINALG_UTILS_H_ #endif // MLIR_DIALECT_LINALG_UTILS_H_

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

Show First 20 Lines • Show All 272 Lines • ▼ Show 20 Lines for (OpOperand *opOperand : op.getOutputTensorsOpOperands()) {
} else { } else {
tensorResults.push_back(res->getResult(resultIdx)); tensorResults.push_back(res->getResult(resultIdx));
} }
++resultIdx; ++resultIdx;
} }
return scf::ValueVector(tensorResults.begin(), tensorResults.end()); return scf::ValueVector(tensorResults.begin(), tensorResults.end());
}; };
GenerateLoopNest<LoopTy>::doit(b, op.getLoc(), loopRanges, op, iteratorTypes, GenerateLoopNest<LoopTy>::doit(b, op.getLoc(), loopRanges, op, iteratorTypes,
tiledLoopBodyBuilder, options.distribution); tiledLoopBodyBuilder, options.distribution,
options.distributionTypes);
// 3. Transform IndexOp results w.r.t. the tiling. // 3. Transform IndexOp results w.r.t. the tiling.
transformIndexOps(b, res, ivs, loopIndexToRangeIndex); transformIndexOps(b, res, ivs, loopIndexToRangeIndex);
// 4. Gather the newly created loops and return them with the new op. // 4. Gather the newly created loops and return them with the new op.
SmallVector<Operation *, 8> loops; SmallVector<Operation *, 8> loops;
loops.reserve(ivs.size()); loops.reserve(ivs.size());
for (auto iv : ivs) { for (auto iv : ivs) {
▲ Show 20 Lines • Show All 133 Lines • ▼ Show 20 Lines
static void insertTilingPatterns(RewritePatternSet &patterns, static void insertTilingPatterns(RewritePatternSet &patterns,
const LinalgTilingOptions &options) { const LinalgTilingOptions &options) {
RewritePatternList<GenericOp, RewritePatternList<GenericOp,
#define GET_OP_LIST #define GET_OP_LIST
#include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc" #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
>::insert(patterns, options); >::insert(patterns, options);
} }
static void applyTilingToLoopPatterns(LinalgTilingLoopType loopType, static void
FuncOp funcOp, applyTilingToLoopPatterns(LinalgTilingLoopType loopType, FuncOp funcOp,
ArrayRef<int64_t> tileSizes) { ArrayRef<int64_t> tileSizes,
auto options = ArrayRef<StringRef> distributionTypes = {}) {
LinalgTilingOptions().setTileSizes(tileSizes).setLoopType(loopType); auto options = LinalgTilingOptions()
.setTileSizes(tileSizes)
.setLoopType(loopType)
.setDistributionTypes(distributionTypes);
MLIRContext *ctx = funcOp.getContext(); MLIRContext *ctx = funcOp.getContext();
RewritePatternSet patterns(ctx); RewritePatternSet patterns(ctx);
insertTilingPatterns(patterns, options); insertTilingPatterns(patterns, options);
(void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns)); (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
(void)applyPatternsAndFoldGreedily( (void)applyPatternsAndFoldGreedily(
funcOp, getLinalgTilingCanonicalizationPatterns(ctx)); funcOp, getLinalgTilingCanonicalizationPatterns(ctx));
// Drop the marker. // Drop the marker.
funcOp.walk([](LinalgOp op) { funcOp.walk([](LinalgOp op) {
Show All 23 Lines void runOnFunction() override {
applyTilingToLoopPatterns(LinalgTilingLoopType::ParallelLoops, applyTilingToLoopPatterns(LinalgTilingLoopType::ParallelLoops,
getFunction(), tileSizes); getFunction(), tileSizes);
} }
}; };
struct LinalgTilingToTiledLoopsPass struct LinalgTilingToTiledLoopsPass
: public LinalgTilingToTiledLoopsBase<LinalgTilingToTiledLoopsPass> { : public LinalgTilingToTiledLoopsBase<LinalgTilingToTiledLoopsPass> {
LinalgTilingToTiledLoopsPass() = default; LinalgTilingToTiledLoopsPass() = default;
LinalgTilingToTiledLoopsPass(ArrayRef<int64_t> sizes) { tileSizes = sizes; } LinalgTilingToTiledLoopsPass(ArrayRef<int64_t> sizes,
ArrayRef<StringRef> types) {
tileSizes = sizes;
distributionTypes = llvm::to_vector<2>(
llvm::map_range(types, [](StringRef ref) { return ref.str(); }));
}
void runOnFunction() override { void runOnFunction() override {
applyTilingToLoopPatterns(LinalgTilingLoopType::TiledLoops, getFunction(), applyTilingToLoopPatterns(
tileSizes); LinalgTilingLoopType::TiledLoops, getFunction(), tileSizes,
llvm::to_vector<2>(
llvm::map_range(distributionTypes,
[](std::string &str) { return StringRef(str); })));
} }
}; };
} // namespace } // namespace
std::unique_ptr<OperationPass<FuncOp>> std::unique_ptr<OperationPass<FuncOp>>
mlir::createLinalgTilingPass(ArrayRef<int64_t> tileSizes) { mlir::createLinalgTilingPass(ArrayRef<int64_t> tileSizes) {
return std::make_unique<LinalgTilingPass>(tileSizes); return std::make_unique<LinalgTilingPass>(tileSizes);
} }
std::unique_ptr<OperationPass<FuncOp>> std::unique_ptr<OperationPass<FuncOp>>
mlir::createLinalgTilingToParallelLoopsPass(ArrayRef<int64_t> tileSizes) { mlir::createLinalgTilingToParallelLoopsPass(ArrayRef<int64_t> tileSizes) {
return std::make_unique<LinalgTilingToParallelLoopsPass>(tileSizes); return std::make_unique<LinalgTilingToParallelLoopsPass>(tileSizes);
} }
std::unique_ptr<OperationPass<FuncOp>> std::unique_ptr<OperationPass<FuncOp>>
mlir::createLinalgTilingToTiledLoopPass(ArrayRef<int64_t> tileSizes) { mlir::createLinalgTilingToTiledLoopPass(ArrayRef<int64_t> tileSizes,
return std::make_unique<LinalgTilingToTiledLoopsPass>(tileSizes); ArrayRef<StringRef> distributionTypes) {
return std::make_unique<LinalgTilingToTiledLoopsPass>(tileSizes,
distributionTypes);
} }

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

Show First 20 Lines • Show All 192 Lines • ▼ Show 20 Lines
/// Specialization to build an scf "for" nest. /// Specialization to build an scf "for" nest.
template <> template <>
void GenerateLoopNest<scf::ForOp>::doit( void GenerateLoopNest<scf::ForOp>::doit(
OpBuilder &b, Location loc, ArrayRef<Range> loopRanges, LinalgOp linalgOp, OpBuilder &b, Location loc, ArrayRef<Range> loopRanges, LinalgOp linalgOp,
ArrayRef<Attribute> iteratorTypes, ArrayRef<Attribute> iteratorTypes,
function_ref<scf::ValueVector(OpBuilder &, Location, ValueRange, function_ref<scf::ValueVector(OpBuilder &, Location, ValueRange,
ValueRange)> ValueRange)>
bodyBuilderFn, bodyBuilderFn,
Optional<LinalgLoopDistributionOptions> distributionOptions) { Optional<LinalgLoopDistributionOptions> distributionOptions,
ArrayRef<StringRef> distributionTypes) {
auto iterArgInitValues = linalgOp.getOutputTensors(); auto iterArgInitValues = linalgOp.getOutputTensors();
// Create procInfo so it dominates loops, if appropriate. // Create procInfo so it dominates loops, if appropriate.
SmallVector<ProcInfo, 4> procInfo; SmallVector<ProcInfo, 4> procInfo;
SmallVector<DistributionMethod, 0> distributionMethod; SmallVector<DistributionMethod, 0> distributionMethod;
if (distributionOptions.hasValue()) { if (distributionOptions.hasValue()) {
// Collect loop ranges for parallel dimensions. // Collect loop ranges for parallel dimensions.
SmallVector<Range, 2> parallelLoopRanges; SmallVector<Range, 2> parallelLoopRanges;
for (auto iteratorType : enumerate(iteratorTypes)) for (auto iteratorType : enumerate(iteratorTypes))
Show All 31 Lines
/// Specialization to build affine "for" nest. /// Specialization to build affine "for" nest.
template <> template <>
void GenerateLoopNest<AffineForOp>::doit( void GenerateLoopNest<AffineForOp>::doit(
OpBuilder &b, Location loc, ArrayRef<Range> loopRanges, LinalgOp linalgOp, OpBuilder &b, Location loc, ArrayRef<Range> loopRanges, LinalgOp linalgOp,
ArrayRef<Attribute> iteratorTypes, ArrayRef<Attribute> iteratorTypes,
function_ref<scf::ValueVector(OpBuilder &, Location, ValueRange, function_ref<scf::ValueVector(OpBuilder &, Location, ValueRange,
ValueRange)> ValueRange)>
bodyBuilderFn, bodyBuilderFn,
Optional<LinalgLoopDistributionOptions>) { Optional<LinalgLoopDistributionOptions>, ArrayRef<StringRef>) {
auto iterArgInitValues = linalgOp.getOutputTensors(); auto iterArgInitValues = linalgOp.getOutputTensors();
assert(iterArgInitValues.empty() && "unexpected AffineForOp init values"); assert(iterArgInitValues.empty() && "unexpected AffineForOp init values");
SmallVector<Value, 4> lbs, ubs, steps; SmallVector<Value, 4> lbs, ubs, steps;
unpackRanges(loopRanges, lbs, ubs, steps); unpackRanges(loopRanges, lbs, ubs, steps);
// Affine loops require constant steps. // Affine loops require constant steps.
SmallVector<int64_t, 4> constantSteps; SmallVector<int64_t, 4> constantSteps;
constantSteps.reserve(steps.size()); constantSteps.reserve(steps.size());
Show All 12 Lines
/// Specialization to build an linalg.tiled_loop /// Specialization to build an linalg.tiled_loop
template <> template <>
void GenerateLoopNest<TiledLoopOp>::doit( void GenerateLoopNest<TiledLoopOp>::doit(
OpBuilder &b, Location loc, ArrayRef<Range> loopRanges, LinalgOp linalgOp, OpBuilder &b, Location loc, ArrayRef<Range> loopRanges, LinalgOp linalgOp,
ArrayRef<Attribute> iteratorTypes, ArrayRef<Attribute> iteratorTypes,
function_ref<scf::ValueVector(OpBuilder &, Location, ValueRange, function_ref<scf::ValueVector(OpBuilder &, Location, ValueRange,
ValueRange)> ValueRange)>
bodyBuilderFn, bodyBuilderFn,
Optional<LinalgLoopDistributionOptions>) { Optional<LinalgLoopDistributionOptions> distributionOptions,
ArrayRef<StringRef> distributionTypes) {
SmallVector<ProcInfo, 2> procInfo; SmallVector<ProcInfo, 2> procInfo;
SmallVector<Value, 4> lbs, ubs, steps; SmallVector<Value, 4> lbs, ubs, steps;
unpackRanges(loopRanges, lbs, ubs, steps); unpackRanges(loopRanges, lbs, ubs, steps);
auto wrappedBuilderFn = [&](OpBuilder &nestedBuilder, Location nestedLoc, auto wrappedBuilderFn = [&](OpBuilder &nestedBuilder, Location nestedLoc,
ValueRange ivs, ValueRange inputs, ValueRange ivs, ValueRange inputs,
ValueRange outputs) { ValueRange outputs) {
scf::ValueVector results = bodyBuilderFn(nestedBuilder, nestedLoc, ivs, scf::ValueVector results = bodyBuilderFn(nestedBuilder, nestedLoc, ivs,
linalgOp.getOutputTensors()); linalgOp.getOutputTensors());
nestedBuilder.create<linalg::YieldOp>(nestedLoc, results); nestedBuilder.create<linalg::YieldOp>(nestedLoc, results);
}; };
auto tiledLoop = b.create<TiledLoopOp>( auto tiledLoop = b.create<TiledLoopOp>(
loc, lbs, ubs, steps, linalgOp.getInputs(), linalgOp.getOutputs(), loc, lbs, ubs, steps, linalgOp.getInputs(), linalgOp.getOutputs(),
b.getArrayAttr(iteratorTypes), wrappedBuilderFn); b.getArrayAttr(iteratorTypes), wrappedBuilderFn);
if (!distributionTypes.empty())
tiledLoop.setDistributionTypes(b, distributionTypes);
// Replace inputs/outputs with the corresponding region args. // Replace inputs/outputs with the corresponding region args.
auto isInsideTiledLoop = [&](OpOperand &operand) { auto isInsideTiledLoop = [&](OpOperand &operand) {
return operand.getOwner()->getBlock() == tiledLoop.getBody(); return operand.getOwner()->getBlock() == tiledLoop.getBody();
}; };
for (auto it : for (auto it :
llvm::zip(linalgOp.getInputs(), tiledLoop.getRegionInputArgs())) llvm::zip(linalgOp.getInputs(), tiledLoop.getRegionInputArgs()))
std::get<0>(it).replaceUsesWithIf(std::get<1>(it), isInsideTiledLoop); std::get<0>(it).replaceUsesWithIf(std::get<1>(it), isInsideTiledLoop);
▲ Show 20 Lines • Show All 139 Lines • ▼ Show 20 Lines
/// Specialization for generating a mix of parallel and sequential scf loops. /// Specialization for generating a mix of parallel and sequential scf loops.
template <> template <>
void GenerateLoopNest<scf::ParallelOp>::doit( void GenerateLoopNest<scf::ParallelOp>::doit(
OpBuilder &b, Location loc, ArrayRef<Range> loopRanges, LinalgOp linalgOp, OpBuilder &b, Location loc, ArrayRef<Range> loopRanges, LinalgOp linalgOp,
ArrayRef<Attribute> iteratorTypes, ArrayRef<Attribute> iteratorTypes,
function_ref<scf::ValueVector(OpBuilder &, Location, ValueRange, function_ref<scf::ValueVector(OpBuilder &, Location, ValueRange,
ValueRange)> ValueRange)>
bodyBuilderFn, bodyBuilderFn,
Optional<LinalgLoopDistributionOptions> distributionOptions) { Optional<LinalgLoopDistributionOptions> distributionOptions,
ArrayRef<StringRef> distributionTypes) {
auto iterArgInitValues = linalgOp.getOutputTensors(); auto iterArgInitValues = linalgOp.getOutputTensors();
assert(iterArgInitValues.empty() && "unexpected ParallelOp init values"); assert(iterArgInitValues.empty() && "unexpected ParallelOp init values");
// This function may be passed more iterator types than ranges. // This function may be passed more iterator types than ranges.
assert(iteratorTypes.size() >= loopRanges.size() && assert(iteratorTypes.size() >= loopRanges.size() &&
"expected iterator type for all ranges"); "expected iterator type for all ranges");
iteratorTypes = iteratorTypes.take_front(loopRanges.size()); iteratorTypes = iteratorTypes.take_front(loopRanges.size());
SmallVector<Value, 8> lbsStorage, ubsStorage, stepsStorage, ivs; SmallVector<Value, 8> lbsStorage, ubsStorage, stepsStorage, ivs;
unsigned numLoops = iteratorTypes.size(); unsigned numLoops = iteratorTypes.size();
▲ Show 20 Lines • Show All 163 Lines • Show Last 20 Lines

mlir/test/Dialect/Linalg/tile-tensors.mlir

// RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=2,3,4" -split-input-file | FileCheck %s // RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=2,3,4" -split-input-file | FileCheck %s
// RUN: mlir-opt %s -linalg-tile-to-tiled-loop="linalg-tile-sizes=2,3,4" -split-input-file | FileCheck %s -check-prefix=TLOOP // RUN: mlir-opt %s -linalg-tile-to-tiled-loop="linalg-tile-sizes=2,3,4 linalg-distribution-types=block_x,block_y,none" -split-input-file | FileCheck %s -check-prefix=TLOOP
// CHECK-LABEL: func @matmul_tensors( // CHECK-LABEL: func @matmul_tensors(
// CHECK-SAME: %[[TA:[0-9a-z]+]]: tensor<?x?xf32> // CHECK-SAME: %[[TA:[0-9a-z]+]]: tensor<?x?xf32>
// CHECK-SAME: %[[TB:[0-9a-z]+]]: tensor<?x?xf32> // CHECK-SAME: %[[TB:[0-9a-z]+]]: tensor<?x?xf32>
// CHECK-SAME: %[[TC:[0-9a-z]+]]: tensor<?x?xf32>) -> tensor<?x?xf32> { // CHECK-SAME: %[[TC:[0-9a-z]+]]: tensor<?x?xf32>) -> tensor<?x?xf32> {
func @matmul_tensors( func @matmul_tensors(
%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>) %arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>)
-> tensor<?x?xf32> { -> tensor<?x?xf32> {
Show All 32 Lines
// TLOOP: %[[ARG_1_Y:.*]] = memref.dim %[[ARG_1]], %[[C1]] : [[TY]] // TLOOP: %[[ARG_1_Y:.*]] = memref.dim %[[ARG_1]], %[[C1]] : [[TY]]
// TLOOP: %{{.*}} = linalg.tiled_loop (%[[I:.*]], %[[J:.*]], %[[K:.*]]) = // TLOOP: %{{.*}} = linalg.tiled_loop (%[[I:.*]], %[[J:.*]], %[[K:.*]]) =
// TLOOP-SAME: (%[[C0]], %[[C0]], %[[C0]]) // TLOOP-SAME: (%[[C0]], %[[C0]], %[[C0]])
// TLOOP-SAME: to (%[[ARG_0_X]], %[[ARG_1_Y]], %[[ARG_0_Y]]) // TLOOP-SAME: to (%[[ARG_0_X]], %[[ARG_1_Y]], %[[ARG_0_Y]])
// TLOOP-SAME: step (%[[C2]], %[[C3]], %[[C4]]) // TLOOP-SAME: step (%[[C2]], %[[C3]], %[[C4]])
// TLOOP-SAME: ins (%[[A0:.*]] = %[[ARG_0]]: [[TY]], %[[A1:.*]] = %[[ARG_1]]: [[TY]]) // TLOOP-SAME: ins (%[[A0:.*]] = %[[ARG_0]]: [[TY]], %[[A1:.*]] = %[[ARG_1]]: [[TY]])
// TLOOP-SAME: outs (%[[A2:.*]] = %[[ARG_2]]: [[TY]]) // TLOOP-SAME: outs (%[[A2:.*]] = %[[ARG_2]]: [[TY]])
// TLOOP-SAME: iterators["parallel", "parallel", "reduction"] { // TLOOP-SAME: iterators["parallel", "parallel", "reduction"]
// TLOOP-SAME: distribution["block_x", "block_y", "none"] {
// TLOOP: %[[SUB_ARG_0:.*]] = subtensor %[[A0]][%[[I]], %[[K]]] // TLOOP: %[[SUB_ARG_0:.*]] = subtensor %[[A0]][%[[I]], %[[K]]]
// TLOOP: %[[SUB_ARG_1:.*]] = subtensor %[[A1]][%[[K]], %[[J]]] // TLOOP: %[[SUB_ARG_1:.*]] = subtensor %[[A1]][%[[K]], %[[J]]]
// TLOOP: %[[SUB_ARG_2:.*]] = subtensor %[[A2]][%[[I]], %[[J]]] // TLOOP: %[[SUB_ARG_2:.*]] = subtensor %[[A2]][%[[I]], %[[J]]]
// TLOOP: %[[PROD:.*]] = linalg.matmul ins(%[[SUB_ARG_0]], %[[SUB_ARG_1]] // TLOOP: %[[PROD:.*]] = linalg.matmul ins(%[[SUB_ARG_0]], %[[SUB_ARG_1]]
// TLOOP-SE: outs(%[[SUB_ARG_2]] : [[TY]]) -> [[TY]] // TLOOP-SE: outs(%[[SUB_ARG_2]] : [[TY]]) -> [[TY]]
▲ Show 20 Lines • Show All 63 Lines • ▼ Show 20 Lines
// TLOOP: %[[ARG_0_Z:.*]] = memref.dim %[[ARG_0]], %[[C2]] : [[TY]] // TLOOP: %[[ARG_0_Z:.*]] = memref.dim %[[ARG_0]], %[[C2]] : [[TY]]
// TLOOP: %{{.*}} = linalg.tiled_loop (%{{.*}}, %{{.*}}, %{{.*}}) = // TLOOP: %{{.*}} = linalg.tiled_loop (%{{.*}}, %{{.*}}, %{{.*}}) =
// TLOOP-SAME: (%[[C0]], %[[C0]], %[[C0]]) // TLOOP-SAME: (%[[C0]], %[[C0]], %[[C0]])
// TLOOP-SAME: to (%[[ARG_0_X]], %[[ARG_0_Y]], %[[ARG_0_Z]]) // TLOOP-SAME: to (%[[ARG_0_X]], %[[ARG_0_Y]], %[[ARG_0_Z]])
// TLOOP-SAME: step (%[[C2]], %[[C3]], %[[C4]]) // TLOOP-SAME: step (%[[C2]], %[[C3]], %[[C4]])
// TLOOP-SAME: ins (%{{.*}} = %[[ARG_0]]: [[TY]], %{{.*}} = %[[ARG_1]]: [[TY]]) // TLOOP-SAME: ins (%{{.*}} = %[[ARG_0]]: [[TY]], %{{.*}} = %[[ARG_1]]: [[TY]])
// TLOOP-SAME: outs (%{{.*}} = %[[INIT]]: [[TY]]) // TLOOP-SAME: outs (%{{.*}} = %[[INIT]]: [[TY]])
// TLOOP-SAME: distribution["block_x", "block_y", "none"] {
// -----
func @fill_tensors(%arg0 : index, %arg1 : index, %arg2 : f32) -> tensor<?x?xf32> {
%0 = linalg.init_tensor [%arg0, %arg1] : tensor<?x?xf32>
%1 = linalg.fill(%0, %arg2) : tensor<?x?xf32>, f32 -> tensor<?x?xf32>
return %1 : tensor<?x?xf32>
}
// CHECK: func @fill_tensors
// CHECK: %[[INIT:.+]] = linalg.init_tensor
// CHECK: %[[RESULT:.+]] = scf.for %[[IV0:[a-zA-z0-9_]+]]
// CHECK-SAME: iter_args(%[[ARG4:.+]] = %[[INIT]]) -> (tensor<?x?xf32>) {
// CHECK: %[[YIELD_1:.+]] = scf.for %[[IV1:[a-zA-Z0-9_]+]]
// CHECK-SAME: iter_args(%[[ARG6:.+]] = %[[ARG4]]) -> (tensor<?x?xf32>) {
// CHECK: %[[FILL_TILE:.+]] = subtensor %[[ARG6]][%[[IV0]], %[[IV1]]]
// CHECK: %[[RESULT_TILE:.+]] = linalg.fill(%[[FILL_TILE]], %{{.+}})
// CHECK: %[[YIELD_2:.+]] = subtensor_insert %[[RESULT_TILE]]
// CHECK-SAME: into %[[ARG6]][%[[IV0]], %[[IV1]]]
// CHECK: scf.yield %[[YIELD_2]]
// CHECK: }
// CHECK: scf.yield %[[YIELD_1]]
// CHECK: }
// CHECK: return %[[RESULT]]
hanchungUnsubmitted
Not Done
ReplyInline Actions

Do we intend to delete the test?

hanchung: Do we intend to delete the test?