This is an archive of the discontinued LLVM Phabricator instance.

Differential D120441

[mlir][Linalg] Add support for tileFuseAndDistribute on tensors.
ClosedPublic

Authored by hanchung on Feb 23 2022, 4:06 PM.

Details

Reviewers
nicolasvasilache
mravishankar
gysit
Commits
rG748bf4bb289b: [mlir][Linalg] Add support for tileFuseAndDistribute on tensors.
Summary

This extends TileAndFuse to handle distribution on tensors.

Diff Detail

Event Timeline

hanchung created this revision.Feb 23 2022, 4:06 PM
Herald added subscribers: sdasgup3, wenzhicui, wrengr and 19 others. · View Herald TranscriptFeb 23 2022, 4:06 PM
hanchung requested review of this revision.Feb 23 2022, 4:06 PM
Herald added a project: Restricted Project. · View Herald TranscriptFeb 23 2022, 4:06 PM
Herald added subscribers: limo1996, stephenneuendorffer. · View Herald Transcript
hanchung added inline comments.Feb 23 2022, 4:08 PM
mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
345

We need LinalgLoopDistributionOptions, so we moved the whole section right after distribution section. Only this line is new to the change.

Harbormaster completed remote builds in B151159: Diff 410966.Feb 23 2022, 4:29 PM
gysit accepted this revision.Feb 24 2022, 12:13 AM

LGTM. I would probably consistently use the term tileDistribution instead of just distribution. However, I am also fine with just distribution if you prefer the short name.

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

nit: /// Tile the root operation using the given tileSizes, tileInterchange, and tileDistribution.

345

nit: I would probably use ̀tileDistribution` to make it more symmetric with the other parameters.

399

nit: I would also use tileDistribution instead of distribution and add it to the comment.

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

nit: I think this extra braces are not needed?

This revision is now accepted and ready to land.Feb 24 2022, 12:13 AM
hanchung updated this revision to Diff 411255.Feb 24 2022, 4:14 PM

Address comments

Harbormaster completed remote builds in B151371: Diff 411255.Feb 24 2022, 4:28 PM
hanchung updated this revision to Diff 411460.Feb 25 2022, 10:52 AM

rebase

Harbormaster completed remote builds in B151509: Diff 411460.Feb 25 2022, 11:18 AM
This revision was landed with ongoing or failed builds.Feb 25 2022, 11:51 AM
Closed by commit rG748bf4bb289b: [mlir][Linalg] Add support for tileFuseAndDistribute on tensors. (authored by hanchung). · Explain Why
This revision was automatically updated to reflect the committed changes.
hanchung added a commit: rG748bf4bb289b: [mlir][Linalg] Add support for tileFuseAndDistribute on tensors..

Revision Contents

PathSize
mlir/
include/
mlir/
Dialect/
Linalg/
Transforms/
12 lines
Utils/
138 lines
lib/
Dialect/
Linalg/
Transforms/
24 lines
5 lines
test/
Dialect/
Linalg/
53 lines
lib/
Dialect/
Linalg/
25 lines

Diff 411486

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

Show First 20 Lines • Show All 632 Lines • ▼ Show 20 Lines LinalgPaddingOptions &setPaddingTransposeComputationFunction(
paddingTransposeComputationFunction = std::move(fun); paddingTransposeComputationFunction = std::move(fun);
return *this; return *this;
} }
}; };
struct LinalgTilingAndFusionOptions { struct LinalgTilingAndFusionOptions {
/// Tile sizes used to tile the root operation. /// Tile sizes used to tile the root operation.
SmallVector<int64_t> tileSizes; SmallVector<int64_t> tileSizes;
LinalgTilingAndFusionOptions &setTileSizes(ArrayRef<int64_t> ts) {
tileSizes.assign(ts.begin(), ts.end());
return *this;
}
/// Tile interchange used to permute the tile loops. /// Tile interchange used to permute the tile loops.
SmallVector<int64_t> tileInterchange; SmallVector<int64_t> tileInterchange;
/// When specified, specifies distribution of generated tile loops to
/// processors.
Optional<LinalgLoopDistributionOptions> tileDistribution = None;
LinalgTilingAndFusionOptions &
setDistributionOptions(LinalgLoopDistributionOptions distributionOptions) {
tileDistribution = std::move(distributionOptions);
return *this;
}
}; };
struct LinalgTilingOptions { struct LinalgTilingOptions {
/// Computation function that returns the tile sizes for each operation. /// Computation function that returns the tile sizes for each operation.
/// Delayed construction of constant tile sizes should occur to interoperate /// Delayed construction of constant tile sizes should occur to interoperate
/// with folding. /// with folding.
TileSizeComputationFunction tileSizeComputationFunction = nullptr; TileSizeComputationFunction tileSizeComputationFunction = nullptr;
▲ Show 20 Lines • Show All 830 Lines • Show Last 20 Lines

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

Show First 20 Lines • Show All 240 Lines • ▼ Show 20 Lines
/// `extract_slice` op (generally obtained by applying the tiling /// `extract_slice` op (generally obtained by applying the tiling
/// transformation). Assumes `producerOfTensor` is a Linalg op that produces /// transformation). Assumes `producerOfTensor` is a Linalg op that produces
/// `consumerOpOperand`. /// `consumerOpOperand`.
FailureOr<FusionInfo> fuseProducerOfTensor(OpBuilder &b, FailureOr<FusionInfo> fuseProducerOfTensor(OpBuilder &b,
OpResult producerOpResult, OpResult producerOpResult,
OpOperand &consumerOpOperand); OpOperand &consumerOpOperand);
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// Fusion on tensor utilities
//===----------------------------------------------------------------------===//
/// A struct to manage the tile loop nest specific information.
class TileLoopNest {
public:
TileLoopNest(LinalgOp rootOp) : rootOp(rootOp) {}
/// Tile the root operation using the given `tileSizes` and `tileInterchange`.
LogicalResult tileRootOp(OpBuilder &b, ArrayRef<int64_t> tileSizes,
ArrayRef<int64_t> tileInterchange);
/// Fuse the producer of `consumerOpOperand` into the tile loop nest. Returns
/// the fused producer or fails if fusion is not possible.
FailureOr<LinalgOp> fuseProducer(OpBuilder &b, OpOperand *consumerOpOperand);
/// Returns the replacement results for the original untiled root operation.
ValueRange getRootOpReplacementResults();
/// Returns the tiled root operation.
LinalgOp getRootOp() { return rootOp; }
/// Returns the tiled root operation and the fused producers.
SmallVector<LinalgOp> getAllTiledAndFusedOps();
/// Returns the loop ops generated from tiling.
ArrayRef<scf::ForOp> getLoopOps() { return tileLoopOps; }
/// Returns true if the tile loop nest has no tile loops.
bool isEmpty();
private:
/// Returns true if the tile loop nest invariants are satisfied:
/// - The `rootOp` has been tiled at least once.
/// - The number of tile loop operations and dimensions match.
/// - The innermost tile loop is the parent of `tiledOp`.
/// - The tile loops are directly nested.
// TODO: relax to support additional control flow, e.g., IfOp.
bool isValid();
/// Searches the block arguments tied to a block argument `bbArg` of the
/// innermost tile loop. Returns the block argument from outermost to
/// innermost or an empty vector if none are found.
SmallVector<BlockArgument> getTiedBBArgs(BlockArgument bbArg);
/// Returns the iteration argument of the outermost tile loop mapped to a
/// block argument `bbArg` of the innermost tile loop.
OpOperand *getTiedIterArg(BlockArgument bbArg);
/// Returns true if `bbArg` has other used than `sliceOp` and its
/// dependencies. Only if there are no other uses, the producer output
/// iteration argument may reused to pass the producer result after fusion.
bool hasOtherUses(BlockArgument bbArg, tensor::ExtractSliceOp sliceOp);
LinalgOp rootOp;
SmallVector<scf::ForOp> tileLoopOps;
DenseMap<Operation *, SmallVector<int64_t>> tiledRootAndFusedOpsLoops;
};
/// Tiles `consumerOp` and fuses its dependencies if possible. Uses the
/// `tileSizes` and `tileInterchange` parameters to control the tiling.
FailureOr<TileLoopNest>
tileConsumerAndFuseProducers(OpBuilder &b, LinalgOp consumerOp,
ArrayRef<int64_t> tileSizes,
ArrayRef<int64_t> tileInterchange);
//===----------------------------------------------------------------------===//
// Distribution utilities // Distribution utilities
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
/// Scheme used to distribute loops to processors. /// Scheme used to distribute loops to processors.
enum class DistributionMethod { enum class DistributionMethod {
/// Cyclic distribution where no assumption is made about the dynamic /// Cyclic distribution where no assumption is made about the dynamic
/// relationship between number of processors and number of iterations of the /// relationship between number of processors and number of iterations of the
/// distributed loop. Distributes the following loop /// distributed loop. Distributes the following loop
▲ Show 20 Lines • Show All 68 Lines • ▼ Show 20 Lines
}; };
/// 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 &builder, Location loc, void updateBoundsForCyclicDistribution(OpBuilder &builder, Location loc,
Value procId, Value nprocs, Value &lb, Value procId, Value nprocs, Value &lb,
Value &ub, Value &step); Value &ub, Value &step);
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// Fusion on tensor utilities
//===----------------------------------------------------------------------===//
/// A struct to manage the tile loop nest specific information.
class TileLoopNest {
public:
TileLoopNest(LinalgOp rootOp) : rootOp(rootOp) {}
/// Tile the root operation using the given `tileSizes` and `tileInterchange`,
gysitUnsubmitted
Not Done
ReplyInline Actions

nit: /// Tile the root operation using the given tileSizes, tileInterchange, and tileDistribution.

gysit: nit: /// Tile the root operation using the given `tileSizes`, `tileInterchange`, and…
/// and `tileDistribution`.
LogicalResult
tileRootOp(OpBuilder &b, ArrayRef<int64_t> tileSizes,
ArrayRef<int64_t> tileInterchange,
hanchungAuthorUnsubmitted
Done
ReplyInline Actions

We need LinalgLoopDistributionOptions, so we moved the whole section right after distribution section. Only this line is new to the change.

hanchung: We need LinalgLoopDistributionOptions, so we moved the whole section right after distribution…
gysitUnsubmitted
Not Done
ReplyInline Actions

nit: I would probably use ̀tileDistribution` to make it more symmetric with the other parameters.

gysit: nit: I would probably use ̀tileDistribution` to make it more symmetric with the other…
Optional<LinalgLoopDistributionOptions> tileDistribution);
/// Fuse the producer of `consumerOpOperand` into the tile loop nest. Returns
/// the fused producer or fails if fusion is not possible.
FailureOr<LinalgOp> fuseProducer(OpBuilder &b, OpOperand *consumerOpOperand);
/// Returns the replacement results for the original untiled root operation.
ValueRange getRootOpReplacementResults();
/// Returns the tiled root operation.
LinalgOp getRootOp() { return rootOp; }
/// Returns the tiled root operation and the fused producers.
SmallVector<LinalgOp> getAllTiledAndFusedOps();
/// Returns the loop ops generated from tiling.
ArrayRef<scf::ForOp> getLoopOps() { return tileLoopOps; }
/// Returns true if the tile loop nest has no tile loops.
bool isEmpty();
private:
/// Returns true if the tile loop nest invariants are satisfied:
/// - The `rootOp` has been tiled at least once.
/// - The number of tile loop operations and dimensions match.
/// - The innermost tile loop is the parent of `tiledOp`.
/// - The tile loops are directly nested.
// TODO: relax to support additional control flow, e.g., IfOp.
bool isValid();
/// Searches the block arguments tied to a block argument `bbArg` of the
/// innermost tile loop. Returns the block argument from outermost to
/// innermost or an empty vector if none are found.
SmallVector<BlockArgument> getTiedBBArgs(BlockArgument bbArg);
/// Returns the iteration argument of the outermost tile loop mapped to a
/// block argument `bbArg` of the innermost tile loop.
OpOperand *getTiedIterArg(BlockArgument bbArg);
/// Returns true if `bbArg` has other used than `sliceOp` and its
/// dependencies. Only if there are no other uses, the producer output
/// iteration argument may reused to pass the producer result after fusion.
bool hasOtherUses(BlockArgument bbArg, tensor::ExtractSliceOp sliceOp);
LinalgOp rootOp;
SmallVector<scf::ForOp> tileLoopOps;
DenseMap<Operation *, SmallVector<int64_t>> tiledRootAndFusedOpsLoops;
};
/// Tiles `consumerOp` and fuses its dependencies if possible. Uses the
/// `tileSizes`, `tileInterchange`, and `tileDistribution` parameters to control
/// the tiling.
FailureOr<TileLoopNest> tileConsumerAndFuseProducers(
OpBuilder &b, LinalgOp consumerOp, ArrayRef<int64_t> tileSizes,
gysitUnsubmitted
Not Done
ReplyInline Actions

nit: I would also use tileDistribution instead of distribution and add it to the comment.

gysit: nit: I would also use `tileDistribution` instead of `distribution` and add it to the comment.
ArrayRef<int64_t> tileInterchange,
Optional<LinalgLoopDistributionOptions> tileDistribution);
//===----------------------------------------------------------------------===//
// Generic op region utilities // Generic op region utilities
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
/// A struct containing common matchers over linalg op's region. /// A struct containing common matchers over linalg op's region.
struct RegionMatcher { struct RegionMatcher {
enum class BinaryOpKind { enum class BinaryOpKind {
IAdd, IAdd,
}; };
Show All 38 Lines

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

Show First 20 Lines • Show All 263 Lines • ▼ Show 20 Linesbool TileLoopNest::hasOtherUses(BlockArgument bbArg,
// Check the block arguments, except for the innermost one, have one use. // Check the block arguments, except for the innermost one, have one use.
SmallVector<BlockArgument> bbArgs = getTiedBBArgs(bbArg); SmallVector<BlockArgument> bbArgs = getTiedBBArgs(bbArg);
return !all_of(bbArgs, [&](BlockArgument bbArg) { return !all_of(bbArgs, [&](BlockArgument bbArg) {
return bbArg.hasOneUse() || bbArg == bbArgs.back(); return bbArg.hasOneUse() || bbArg == bbArgs.back();
}); });
} }
LogicalResult TileLoopNest::tileRootOp(OpBuilder &b, LogicalResult TileLoopNest::tileRootOp(
ArrayRef<int64_t> tileSizes, OpBuilder &b, ArrayRef<int64_t> tileSizes,
ArrayRef<int64_t> tileInterchange) { ArrayRef<int64_t> tileInterchange,
Optional<LinalgLoopDistributionOptions> tileDistribution) {
// Exit if all tile sizes are zero. // Exit if all tile sizes are zero.
if (tileSizes.size() == static_cast<size_t>(count(tileSizes, 0))) if (tileSizes.size() == static_cast<size_t>(count(tileSizes, 0)))
return success(); return success();
// Tile the root operation. // Tile the root operation.
LinalgTilingOptions tilingOptions; LinalgTilingOptions tilingOptions;
tilingOptions = tilingOptions tilingOptions = tilingOptions
.setInterchange(SmallVector<unsigned>( .setInterchange(SmallVector<unsigned>(
tileInterchange.begin(), tileInterchange.end())) tileInterchange.begin(), tileInterchange.end()))
.setTileSizes(tileSizes) .setTileSizes(tileSizes)
.setLoopType(LinalgTilingLoopType::Loops); .setLoopType(LinalgTilingLoopType::Loops);
if (tileDistribution)
tilingOptions =
tilingOptions.setDistributionOptions(tileDistribution.getValue());
// TODO: Propagate RewriterBase everywhere. // TODO: Propagate RewriterBase everywhere.
IRRewriter rewriter(b); IRRewriter rewriter(b);
FailureOr<TiledLinalgOp> tiledRootOp = FailureOr<TiledLinalgOp> tiledRootOp =
tileLinalgOp(rewriter, rootOp, tilingOptions); tileLinalgOp(rewriter, rootOp, tilingOptions);
// Exit if tiling the root operation fails. // Exit if tiling the root operation fails.
if (failed(tiledRootOp)) if (failed(tiledRootOp))
▲ Show 20 Lines • Show All 109 Lines • ▼ Show 20 LinesSmallVector<LinalgOp> TileLoopNest::getAllTiledAndFusedOps() {
} }
return result; return result;
} }
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// Tile and fuse entry-points. // Tile and fuse entry-points.
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
FailureOr<TileLoopNest> FailureOr<TileLoopNest> mlir::linalg::tileConsumerAndFuseProducers(
mlir::linalg::tileConsumerAndFuseProducers(OpBuilder &b, LinalgOp consumerOp, OpBuilder &b, LinalgOp consumerOp, ArrayRef<int64_t> tileSizes,
ArrayRef<int64_t> tileSizes, ArrayRef<int64_t> tileInterchange,
ArrayRef<int64_t> tileInterchange) { Optional<LinalgLoopDistributionOptions> tileDistribution) {
assert(tileSizes.size() == tileInterchange.size() && assert(tileSizes.size() == tileInterchange.size() &&
"expect the number of tile sizes and interchange dims to match"); "expect the number of tile sizes and interchange dims to match");
assert(isPermutation(tileInterchange) && assert(isPermutation(tileInterchange) &&
"expect tile interchange is a permutation"); "expect tile interchange is a permutation");
// Create an empty tile loop nest. // Create an empty tile loop nest.
TileLoopNest tileLoopNest(consumerOp); TileLoopNest tileLoopNest(consumerOp);
Show All 18 Lines while (!candidates.empty()) {
candidates.append(fusedProducer->getInputAndOutputOperands()); candidates.append(fusedProducer->getInputAndOutputOperands());
} }
}; };
// Tile the outer parallel loops and fuse the output operands. // Tile the outer parallel loops and fuse the output operands.
SmallVector<int64_t> outerTileSizes; SmallVector<int64_t> outerTileSizes;
outerTileSizes.append(tileSizes.begin(), tileSizes.begin() + split); outerTileSizes.append(tileSizes.begin(), tileSizes.begin() + split);
outerTileSizes.append(tileSizes.size() - split, 0); outerTileSizes.append(tileSizes.size() - split, 0);
if (failed(tileLoopNest.tileRootOp(b, outerTileSizes, tileInterchange))) if (failed(tileLoopNest.tileRootOp(b, outerTileSizes, tileInterchange,
tileDistribution)))
return failure(); return failure();
fuseProducersGreedily(tileLoopNest.getRootOp().getOutputOperands()); fuseProducersGreedily(tileLoopNest.getRootOp().getOutputOperands());
// Tile the remaining loops and fuse the input operands. // Tile the remaining loops and fuse the input operands.
SmallVector<int64_t> innerTileSizes; SmallVector<int64_t> innerTileSizes;
innerTileSizes.append(split, 0); innerTileSizes.append(split, 0);
innerTileSizes.append(tileSizes.begin() + split, tileSizes.end()); innerTileSizes.append(tileSizes.begin() + split, tileSizes.end());
if (failed(tileLoopNest.tileRootOp(b, innerTileSizes, tileInterchange))) if (failed(tileLoopNest.tileRootOp(b, innerTileSizes, tileInterchange,
tileDistribution)))
return failure(); return failure();
fuseProducersGreedily(tileLoopNest.getRootOp().getInputOperands()); fuseProducersGreedily(tileLoopNest.getRootOp().getInputOperands());
// Exit if the tile loop nest is empty since all tile sizes are zero. // Exit if the tile loop nest is empty since all tile sizes are zero.
if (tileLoopNest.isEmpty()) if (tileLoopNest.isEmpty())
return failure(); return failure();
return tileLoopNest; return tileLoopNest;
} }

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

Show First 20 Lines • Show All 607 Lines • ▼ Show 20 LinesLogicalResult mlir::linalg::LinalgTileAndFuseTensorOpsPattern::matchAndRewrite(
// Check `rootInterchange` is a permutation of the `rootOp` loop dimensions. // Check `rootInterchange` is a permutation of the `rootOp` loop dimensions.
// It has to be a permutation since the tiling cannot tile the same loop // It has to be a permutation since the tiling cannot tile the same loop
// dimension multiple times. // dimension multiple times.
if (!isPermutation(rootInterchange)) if (!isPermutation(rootInterchange))
return rewriter.notifyMatchFailure( return rewriter.notifyMatchFailure(
op, "expect the tile interchange permutes the root loops"); op, "expect the tile interchange permutes the root loops");
// Tile `rootOp` and fuse its producers. // Tile `rootOp` and fuse its producers.
FailureOr<TileLoopNest> tileLoopNest = tileConsumerAndFuseProducers( FailureOr<TileLoopNest> tileLoopNest =
rewriter, rootOp, rootTileSizes, rootInterchange); tileConsumerAndFuseProducers(rewriter, rootOp, rootTileSizes,
rootInterchange, options.tileDistribution);
if (failed(tileLoopNest)) if (failed(tileLoopNest))
return rewriter.notifyMatchFailure( return rewriter.notifyMatchFailure(
op, "tileConsumerAndFuseProducers failed unexpectedly"); op, "tileConsumerAndFuseProducers failed unexpectedly");
// Replace all uses of the tiled loop operation. // Replace all uses of the tiled loop operation.
rootOp->replaceAllUsesWith(tileLoopNest->getRootOpReplacementResults()); rootOp->replaceAllUsesWith(tileLoopNest->getRootOpReplacementResults());
// Apply the filter if specified. // Apply the filter if specified.
▲ Show 20 Lines • Show All 503 Lines • Show Last 20 Lines

mlir/test/Dialect/Linalg/tile-fuse-and-distribute.mlir

  • This file was added.
// RUN: mlir-opt %s -test-linalg-transform-patterns=test-tile-fuse-and-distribute-options -split-input-file | FileCheck %s
// CHECK: #[[MULMAP:.+]] = affine_map<()[s0, s1] -> (s0 * s1)>
// CHECK: #[[ADDMAP:.+]] = affine_map<()[s0, s1] -> (s0 + s1)>
// CHECK: func @fill_matmul_tensors(
// CHECK-SAME: %[[TA:[0-9a-z]+]]: tensor<?x?xf32>
// CHECK-SAME: %[[TB:[0-9a-z]+]]: tensor<?x?xf32>) -> tensor<?x?xf32> {
func @fill_matmul_tensors(
%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>)
-> tensor<?x?xf32> {
// CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[BIDY:.*]] = gpu.block_id y
// CHECK-DAG: %[[NBLOCKSY:.*]] = gpu.grid_dim y
// CHECK-DAG: %[[BIDX:.*]] = gpu.block_id x
// CHECK-DAG: %[[NBLOCKSX:.*]] = gpu.grid_dim x
// CHECK-DAG: %[[INIT:.+]] = linalg.init_tensor
// CHECK: %[[MUL:.+]] = affine.apply #[[MULMAP]]()[%[[BIDY]], %[[C8]]]
// CHECK: %[[LBY:.+]] = affine.apply #[[ADDMAP]]()[%[[MUL]], %[[C0]]]
// CHECK: %[[STEPY:.+]] = affine.apply #[[MULMAP]]()[%[[NBLOCKSY]], %[[C8]]]
// CHECK: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[INIT]]) -> (tensor<?x?xf32>) {
// CHECK: %[[MUL:.+]] = affine.apply #[[MULMAP]]()[%[[BIDX]], %[[C8]]]
// CHECK: %[[LBX:.+]] = affine.apply #[[ADDMAP]]()[%[[MUL]], %[[C0]]]
// CHECK: %[[STEPX:.+]] = affine.apply #[[MULMAP]]()[%[[NBLOCKSX]], %[[C8]]]
// CHECK: %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor<?x?xf32>) {
// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[TC1]]
// CHECK: %[[FILL:.+]] = linalg.fill(%{{.+}}, %[[SLICE]])
// CHECK: %[[sTD2:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC2:.*]] = %[[FILL]]) -> (tensor<?x?xf32>) {
// CHECK: %[[sTA:.*]] = tensor.extract_slice %[[TA]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32>
// CHECK: %[[sTB:.*]] = tensor.extract_slice %[[TB]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32>
// CHECK: %[[sTC:.*]] = tensor.extract_slice %[[TC2]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32>
// CHECK: %[[sTD:.*]] = linalg.matmul ins(%[[sTA]], %[[sTB]] : tensor<?x?xf32>, tensor<?x?xf32>)
// CHECK-SAME: outs(%[[sTC]] : tensor<?x?xf32>) -> tensor<?x?xf32>
// CHECK: %[[TD:.*]] = tensor.insert_slice %[[sTD]] into %[[TC2]][{{.*}}] : tensor<?x?xf32> into tensor<?x?xf32>
// CHECK: scf.yield %[[TD]] : tensor<?x?xf32>
// CHECK: %[[TD2:.*]] = tensor.insert_slice %[[sTD2]] into %[[TC1]][{{.*}}] : tensor<?x?xf32> into tensor<?x?xf32>
// CHECK: scf.yield %[[TD2]] : tensor<?x?xf32>
// CHECK: scf.yield %[[TD1]] : tensor<?x?xf32>
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%cst = arith.constant 0.0 : f32
%0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
%1 = tensor.dim %arg1, %c1 : tensor<?x?xf32>
%2 = linalg.init_tensor [%0, %1] : tensor<?x?xf32>
%3 = linalg.fill(%cst, %2) : f32, tensor<?x?xf32> -> tensor<?x?xf32>
%4 = linalg.matmul {__internal_linalg_transform__ = "tensors_fuse_distribute1"}
ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>)
outs(%3: tensor<?x?xf32>)
-> tensor<?x?xf32>
// CHECK: return %[[TD0]] : tensor<?x?xf32>
return %4 : tensor<?x?xf32>
}

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

Show First 20 Lines • Show All 72 Lines • ▼ Show 20 Lines Option<bool> testMatmulToVectorPatterns2dTiling{
llvm::cl::init(false)}; llvm::cl::init(false)};
Option<bool> testPromotionOptions{*this, "test-linalg-promotion-options", Option<bool> testPromotionOptions{*this, "test-linalg-promotion-options",
llvm::cl::desc("Test promotion options"), llvm::cl::desc("Test promotion options"),
llvm::cl::init(false)}; llvm::cl::init(false)};
Option<bool> testTileAndDistributionOptions{ Option<bool> testTileAndDistributionOptions{
*this, "test-tile-and-distribute-options", *this, "test-tile-and-distribute-options",
llvm::cl::desc("Test tile and distribute options"), llvm::cl::desc("Test tile and distribute options"),
llvm::cl::init(false)}; llvm::cl::init(false)};
Option<bool> testTileFuseAndDistributionOptions{
*this, "test-tile-fuse-and-distribute-options",
llvm::cl::desc("Test tile, fuse and distribute options"),
llvm::cl::init(false)};
Option<bool> testVectorTransferForwardingPatterns{ Option<bool> testVectorTransferForwardingPatterns{
*this, "test-vector-transfer-forwarding-patterns", *this, "test-vector-transfer-forwarding-patterns",
llvm::cl::desc( llvm::cl::desc(
"Test a fused pass that forwards memref.copy to vector.transfer"), "Test a fused pass that forwards memref.copy to vector.transfer"),
llvm::cl::init(false)}; llvm::cl::init(false)};
Option<bool> testGenericToVectorPattern{ Option<bool> testGenericToVectorPattern{
*this, "test-linalg-to-vector-patterns", *this, "test-linalg-to-vector-patterns",
llvm::cl::desc("Test a set of patterns that rewrite a linalg contraction " llvm::cl::desc("Test a set of patterns that rewrite a linalg contraction "
▲ Show 20 Lines • Show All 411 Lines • ▼ Show 20 Lines patterns.add<LinalgTilingPattern>(
.setLoopType(LinalgTilingLoopType::Loops) .setLoopType(LinalgTilingLoopType::Loops)
.setDistributionOptions(cyclicNprocsEqNiters), .setDistributionOptions(cyclicNprocsEqNiters),
LinalgTransformationFilter( LinalgTransformationFilter(
StringAttr::get(context, "tensors_distribute1"), StringAttr::get(context, "tensors_distribute1"),
StringAttr::get(context, "tensors_after_distribute1"))); StringAttr::get(context, "tensors_after_distribute1")));
} }
} }
static void fillTileFuseAndDistributePatterns(MLIRContext *context,
RewritePatternSet &patterns) {
LinalgLoopDistributionOptions cyclicNprocsEqNiters;
gysitUnsubmitted
Not Done
ReplyInline Actions

nit: I think this extra braces are not needed?

gysit: nit: I think this extra braces are not needed?
cyclicNprocsEqNiters.distributionMethod.resize(2, DistributionMethod::Cyclic);
cyclicNprocsEqNiters.procInfo = getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
patterns.add<LinalgTileAndFuseTensorOpsPattern>(
MatmulOp::getOperationName(), context,
LinalgTilingAndFusionOptions()
.setTileSizes({8, 8, 4})
.setDistributionOptions(cyclicNprocsEqNiters),
LinalgTransformationFilter(
StringAttr::get(context, "tensors_fuse_distribute1"),
StringAttr::get(context, "tensors_after_fuse_distribute1")));
}
static void static void
applyMatmulToVectorPatterns(FuncOp funcOp, applyMatmulToVectorPatterns(FuncOp funcOp,
bool testMatmulToVectorPatterns1dTiling, bool testMatmulToVectorPatterns1dTiling,
bool testMatmulToVectorPatterns2dTiling) { bool testMatmulToVectorPatterns2dTiling) {
MLIRContext *ctx = funcOp.getContext(); MLIRContext *ctx = funcOp.getContext();
SmallVector<RewritePatternSet, 4> stage1Patterns; SmallVector<RewritePatternSet, 4> stage1Patterns;
if (testMatmulToVectorPatterns1dTiling) { if (testMatmulToVectorPatterns1dTiling) {
fillL1TilingAndMatmulToVectorPatterns(funcOp, "START", stage1Patterns); fillL1TilingAndMatmulToVectorPatterns(funcOp, "START", stage1Patterns);
▲ Show 20 Lines • Show All 177 Lines • ▼ Show 20 Lines if (testPromotionOptions) {
return; return;
} }
if (testTileAndDistributionOptions) { if (testTileAndDistributionOptions) {
RewritePatternSet patterns(&getContext()); RewritePatternSet patterns(&getContext());
fillTileAndDistributePatterns(&getContext(), patterns); fillTileAndDistributePatterns(&getContext(), patterns);
(void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns)); (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
return; return;
} }
if (testTileFuseAndDistributionOptions) {
RewritePatternSet patterns(&getContext());
fillTileFuseAndDistributePatterns(&getContext(), patterns);
(void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
return;
}
if (testPatterns) if (testPatterns)
return applyPatterns(getOperation()); return applyPatterns(getOperation());
if (testMatmulToVectorPatterns1dTiling || testMatmulToVectorPatterns2dTiling) if (testMatmulToVectorPatterns1dTiling || testMatmulToVectorPatterns2dTiling)
return applyMatmulToVectorPatterns(getOperation(), return applyMatmulToVectorPatterns(getOperation(),
testMatmulToVectorPatterns1dTiling, testMatmulToVectorPatterns1dTiling,
testMatmulToVectorPatterns2dTiling); testMatmulToVectorPatterns2dTiling);
if (testVectorTransferForwardingPatterns) if (testVectorTransferForwardingPatterns)
return applyVectorTransferForwardingPatterns(getOperation()); return applyVectorTransferForwardingPatterns(getOperation());
Show All 26 Lines