[llvm-branch-commits] [mlir] 11ea2e2 - [mlir][Linalg] NFC: Expose some utility functions used for promotion.
via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Mon Nov 23 10:40:29 PST 2020
Author: MaheshRavishankar
Date: 2020-11-23T10:35:42-08:00
New Revision: 11ea2e2448a5f071a04463f94b7bccfe0a32d264
URL: https://github.com/llvm/llvm-project/commit/11ea2e2448a5f071a04463f94b7bccfe0a32d264
DIFF: https://github.com/llvm/llvm-project/commit/11ea2e2448a5f071a04463f94b7bccfe0a32d264.diff
LOG: [mlir][Linalg] NFC: Expose some utility functions used for promotion.
Exposing some utility functions from Linalg to allow for promotion of
fused views outside of the core tile+fuse logic.
This is an alternative to patch D91322 which adds the promotion logic
to the tileAndFuse method. Downside with that approach is that it is
not easily customizable based on needs.
Differential Revision: https://reviews.llvm.org/D91503
Added:
Modified:
mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index d2d6cef1954d..d67e81ceab87 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -234,6 +234,20 @@ struct LinalgPromotionOptions {
}
};
+/// Creates a new buffer using the `allocationFn` provided. The size of this
+/// buffer is the smallest constant bounding size along each dimension that can
+/// be computed for the size of the result of `subView`. Returns the allocated
+/// buffer as `fullLocalView` and the view that matches the size of the result
+/// of subview operation as `partialLocalView`.
+struct PromotionInfo {
+ Value fullLocalView;
+ Value partialLocalView;
+};
+Optional<PromotionInfo>
+promoteSubviewAsNewBuffer(OpBuilder &b, Location loc, SubViewOp subView,
+ AllocBufferCallbackFn allocationFn,
+ OperationFolder *folder = nullptr);
+
/// Promotes the `subViews` into a new buffer allocated at the insertion point
/// `b`. Promotion occurs in 3 steps:
/// 1. Create a new buffer for a full tile (i.e. not clipped at the boundary).
diff --git a/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h b/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
index a6b8afdce9d3..fb916d3962e3 100644
--- a/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
@@ -17,6 +17,7 @@
#include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/SetVector.h"
using mlir::edsc::intrinsics::AffineIndexedValue;
@@ -82,6 +83,13 @@ bool isProducerLastWriteOfView(const LinalgDependenceGraph &graph,
bool isFusableInto(const LinalgDependenceGraph &graph, LinalgOp consumer,
Value consumedView, LinalgOp producer);
+using FusableOpDependencesTy = llvm::MapVector<
+ Operation *,
+ SmallVector<LinalgDependenceGraph::LinalgDependenceGraphElem, 1>>;
+FusableOpDependencesTy
+findAllFusableDependences(ArrayRef<LinalgOp> ops,
+ const LinalgDependenceGraph &dependenceGraph);
+
/// Fuses producer into consumer if the producer is structurally feasible and
/// the fusion would not violate dependencies.
/// Implements the fusion part of the "tileAndFuse on buffers"
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
index 8ec71eec75b8..3160e8f8be0b 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
@@ -542,10 +542,6 @@ static AffineMap pruneReductionDimsFromMap(ArrayRef<Attribute> iteratorTypes,
return getProjectedMap(map, projectedDims);
}
-using FusableOpDependencesTy = llvm::MapVector<
- Operation *,
- SmallVector<LinalgDependenceGraph::LinalgDependenceGraphElem, 1>>;
-
/// Returns the mapping from iterations in the consumer that write to the same
/// location as the iterations in the producer. To do so use
/// - indexing map of the fused view in the consumer : consumerIndexMap
@@ -729,10 +725,9 @@ collectFusableLoops(ArrayRef<LinalgOp> ops,
return fusableLoops;
}
-/// Find all dependences that are to be fusable.
-static FusableOpDependencesTy
-findAllFusableDependences(ArrayRef<LinalgOp> ops,
- const LinalgDependenceGraph &dependenceGraph) {
+/// Find all dependences that are fusable.
+FusableOpDependencesTy mlir::linalg::findAllFusableDependences(
+ ArrayRef<LinalgOp> ops, const LinalgDependenceGraph &dependenceGraph) {
FusableOpDependencesTy fusableDependences;
// TODO: Currently fusion would not be legal if the fusable dependence is to
// the same producer but
diff erent indexing map in the consumer. Fix this, but
@@ -836,6 +831,7 @@ fuseOperations(OpBuilder &builder, LinalgOp tiledOp,
fusedLoopsAndRanges[loop] = getRangeFromOperandShape(
builder, tiledOp.getLoc(), shapeDim.shape, shapeDim.dimension);
}
+
SmallVector<LinalgOp, 1> fusedOps(fusionCandidates.size());
for (auto candidate : enumerate(llvm::reverse(fusionCandidates))) {
LinalgOp fusedOp = fuse(builder, candidate.value(), fusedLoopsAndRanges);
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp b/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
index e002336ed1c6..a824f6eb620f 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
@@ -166,11 +166,6 @@ struct LinalgOpInstancePromotionOptions {
/// Alignment of promoted buffer.
Optional<unsigned> alignment;
};
-
-struct PromotionInfo {
- Value fullLocalView;
- Value partialLocalView;
-};
} // namespace
LinalgOpInstancePromotionOptions::LinalgOpInstancePromotionOptions(
@@ -233,10 +228,10 @@ LinalgOpInstancePromotionOptions::LinalgOpInstancePromotionOptions(
// To account for general boundary effects, padding must be performed on the
// boundary tiles. For now this is done with an unconditional `fill` op followed
// by a partial `copy` op.
-static Optional<PromotionInfo>
-promoteSubviewAsNewBuffer(OpBuilder &b, Location loc, SubViewOp subView,
- LinalgOpInstancePromotionOptions const &options,
- OperationFolder *folder) {
+Optional<PromotionInfo> mlir::linalg::promoteSubviewAsNewBuffer(
+ OpBuilder &b, Location loc, SubViewOp subView,
+ AllocBufferCallbackFn allocationFn, OperationFolder *folder) {
+ ScopedContext scopedContext(b, loc);
auto viewType = subView.getType();
auto rank = viewType.getRank();
SmallVector<Value, 4> fullSizes, partialSizes;
@@ -254,8 +249,7 @@ promoteSubviewAsNewBuffer(OpBuilder &b, Location loc, SubViewOp subView,
SmallVector<int64_t, 4> dynSizes(fullSizes.size(), -1);
// If a callback is not specified, then use the default implementation for
// allocating the promoted buffer.
- Optional<Value> fullLocalView =
- options.allocationFn(b, subView, fullSizes, folder);
+ Optional<Value> fullLocalView = allocationFn(b, subView, fullSizes, folder);
if (!fullLocalView)
return {};
auto zero = folded_std_constant_index(folder, 0);
@@ -279,8 +273,8 @@ promoteSubViews(OpBuilder &b, Location loc,
for (auto v : options.subViews) {
SubViewOp subView = cast<SubViewOp>(v.second.getDefiningOp());
- Optional<PromotionInfo> promotionInfo =
- promoteSubviewAsNewBuffer(b, loc, subView, options, folder);
+ Optional<PromotionInfo> promotionInfo = promoteSubviewAsNewBuffer(
+ b, loc, subView, options.allocationFn, folder);
if (!promotionInfo)
return {};
promotionInfoMap[v.first] = *promotionInfo;
More information about the llvm-branch-commits
mailing list