[Mlir-commits] [mlir] 19e5b2b - [mlir][Linalg] NFC - Simplify GenericNestLoop builder
Nicolas Vasilache
llvmlistbot at llvm.org
Wed May 20 06:48:37 PDT 2020
Author: Nicolas Vasilache
Date: 2020-05-20T09:44:15-04:00
New Revision: 19e5b2bccb4c93f85ca098528f9fa42344e8312f
URL: https://github.com/llvm/llvm-project/commit/19e5b2bccb4c93f85ca098528f9fa42344e8312f
DIFF: https://github.com/llvm/llvm-project/commit/19e5b2bccb4c93f85ca098528f9fa42344e8312f.diff
LOG: [mlir][Linalg] NFC - Simplify GenericNestLoop builder
Summary: This revision trims unnecessary complexity.
Differential Revision: https://reviews.llvm.org/D80290
Added:
Modified:
mlir/include/mlir/Dialect/Linalg/EDSC/Builders.h
mlir/include/mlir/Dialect/SCF/EDSC/Builders.h
mlir/lib/Dialect/Linalg/EDSC/Builders.cpp
mlir/lib/Dialect/Linalg/Transforms/Loops.cpp
mlir/lib/Dialect/SCF/EDSC/Builders.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Linalg/EDSC/Builders.h b/mlir/include/mlir/Dialect/Linalg/EDSC/Builders.h
index bb9933798338..8f0ebd398d75 100644
--- a/mlir/include/mlir/Dialect/Linalg/EDSC/Builders.h
+++ b/mlir/include/mlir/Dialect/Linalg/EDSC/Builders.h
@@ -32,52 +32,13 @@ class ParallelOp;
namespace edsc {
class AffineLoopNestBuilder;
+class LoopNestBuilder;
class ParallelLoopNestBuilder;
-/// A LoopRangeBuilder is a generic NestedBuilder for scf.for operations.
-/// More specifically it is meant to be used as a temporary object for
-/// representing any nested MLIR construct that is "related to" an mlir::Value
-/// (for now an induction variable).
-class LoopRangeBuilder : public NestedBuilder {
-public:
- /// Constructs a new scf.for and captures the associated induction
- /// variable. A Value pointer is passed as the first argument and is the
- /// *only* way to capture the loop induction variable.
- LoopRangeBuilder(Value *iv, Value range);
- LoopRangeBuilder(Value *iv, SubViewOp::Range range);
-
- LoopRangeBuilder(const LoopRangeBuilder &) = delete;
- LoopRangeBuilder(LoopRangeBuilder &&) = default;
-
- LoopRangeBuilder &operator=(const LoopRangeBuilder &) = delete;
- LoopRangeBuilder &operator=(LoopRangeBuilder &&) = default;
-
- /// The only purpose of this operator is to serve as a sequence point so that
- /// the evaluation of `fun` (which build IR snippets in a scoped fashion) is
- /// scoped within a LoopRangeBuilder.
- Value operator()(std::function<void(void)> fun = nullptr);
-};
-
-/// Helper class to sugar building scf.for loop nests from ranges.
-/// This is similar to edsc::AffineLoopNestBuilder except it works on ranges
-/// directly. In the current implementation it produces scf.for operations.
-class LoopNestRangeBuilder {
-public:
- LoopNestRangeBuilder(MutableArrayRef<Value> ivs, ArrayRef<Value> ranges);
- LoopNestRangeBuilder(MutableArrayRef<Value> ivs,
- ArrayRef<SubViewOp::Range> ranges);
- Value operator()(std::function<void(void)> fun = nullptr);
-
-private:
- SmallVector<LoopRangeBuilder, 4> loops;
-};
-
/// Helper template class for building scf.for and affine.loop nests from
/// ranges.
template <typename LoopTy> class GenericLoopNestRangeBuilder {
public:
- GenericLoopNestRangeBuilder(MutableArrayRef<Value> ivs,
- ArrayRef<Value> ranges);
GenericLoopNestRangeBuilder(MutableArrayRef<Value> ivs,
ArrayRef<SubViewOp::Range> ranges);
void operator()(std::function<void(void)> fun = nullptr) { (*builder)(fun); }
@@ -85,7 +46,7 @@ template <typename LoopTy> class GenericLoopNestRangeBuilder {
private:
using LoopOrAffineLoopBuilder =
typename std::conditional_t<std::is_same<LoopTy, AffineForOp>::value,
- AffineLoopNestBuilder, LoopNestRangeBuilder>;
+ AffineLoopNestBuilder, LoopNestBuilder>;
using BuilderType =
typename std::conditional_t<std::is_same<LoopTy, scf::ParallelOp>::value,
ParallelLoopNestBuilder,
diff --git a/mlir/include/mlir/Dialect/SCF/EDSC/Builders.h b/mlir/include/mlir/Dialect/SCF/EDSC/Builders.h
index cefbfd95ed5f..fa72bd623b25 100644
--- a/mlir/include/mlir/Dialect/SCF/EDSC/Builders.h
+++ b/mlir/include/mlir/Dialect/SCF/EDSC/Builders.h
@@ -65,7 +65,7 @@ class LoopNestBuilder {
ValueRange iterArgsInitValues);
LoopNestBuilder(MutableArrayRef<Value> ivs, ArrayRef<Value> lbs,
ArrayRef<Value> ubs, ArrayRef<Value> steps);
- Operation::result_range operator()(std::function<void(void)> fun = nullptr);
+ ValueRange operator()(std::function<void(void)> fun = nullptr);
private:
SmallVector<LoopBuilder, 4> loops;
diff --git a/mlir/lib/Dialect/Linalg/EDSC/Builders.cpp b/mlir/lib/Dialect/Linalg/EDSC/Builders.cpp
index afc572c16bbf..0efbf8545185 100644
--- a/mlir/lib/Dialect/Linalg/EDSC/Builders.cpp
+++ b/mlir/lib/Dialect/Linalg/EDSC/Builders.cpp
@@ -21,76 +21,6 @@ using namespace mlir::edsc::intrinsics;
using namespace mlir::linalg;
using namespace mlir::scf;
-mlir::edsc::LoopRangeBuilder::LoopRangeBuilder(Value *iv, Value range) {
- assert(range.getType() && "expected !linalg.range type");
- assert(range.getDefiningOp() && "need operations to extract range parts");
- auto rangeOp = cast<RangeOp>(range.getDefiningOp());
- auto lb = rangeOp.min();
- auto ub = rangeOp.max();
- auto step = rangeOp.step();
- ForOp forOp = OperationBuilder<ForOp>(lb, ub, step);
- *iv = forOp.getInductionVar();
- auto *body = forOp.getBody();
- enter(body);
-}
-
-mlir::edsc::LoopRangeBuilder::LoopRangeBuilder(Value *iv,
- SubViewOp::Range range) {
- ForOp forOp = OperationBuilder<ForOp>(range.offset, range.size, range.stride);
- *iv = forOp.getInductionVar();
- auto *body = forOp.getBody();
- enter(body);
-}
-
-Value mlir::edsc::LoopRangeBuilder::operator()(std::function<void(void)> fun) {
- if (fun)
- fun();
- exit();
- return Value();
-}
-
-mlir::edsc::LoopNestRangeBuilder::LoopNestRangeBuilder(
- MutableArrayRef<Value> ivs, ArrayRef<SubViewOp::Range> ranges) {
- loops.reserve(ranges.size());
- for (unsigned i = 0, e = ranges.size(); i < e; ++i)
- loops.emplace_back(&ivs[i], ranges[i]);
- assert(loops.size() == ivs.size() && "Mismatch loops vs ivs size");
-}
-
-mlir::edsc::LoopNestRangeBuilder::LoopNestRangeBuilder(
- MutableArrayRef<Value> ivs, ArrayRef<Value> ranges) {
- loops.reserve(ranges.size());
- for (unsigned i = 0, e = ranges.size(); i < e; ++i)
- loops.emplace_back(&ivs[i], ranges[i]);
- assert(loops.size() == ivs.size() && "Mismatch loops vs ivs size");
-}
-
-Value LoopNestRangeBuilder::LoopNestRangeBuilder::operator()(
- std::function<void(void)> fun) {
- if (fun)
- fun();
- for (auto &lit : reverse(loops)) {
- lit({});
- }
- return Value();
-}
-
-namespace mlir {
-namespace edsc {
-
-static void unpackRanges(ArrayRef<Value> rangeOps, SmallVectorImpl<Value> &lbs,
- SmallVectorImpl<Value> &ubs,
- SmallVectorImpl<Value> &steps) {
- for (Value range : rangeOps) {
- assert(range.getType() && "expected linalg.range type");
- assert(range.getDefiningOp() && "need operations to extract range parts");
- RangeOp rangeOp = cast<RangeOp>(range.getDefiningOp());
- lbs.emplace_back(rangeOp.min());
- ubs.emplace_back(rangeOp.max());
- steps.emplace_back(rangeOp.step());
- }
-}
-
static void unpackRanges(ArrayRef<SubViewOp::Range> ranges,
SmallVectorImpl<Value> &lbs,
SmallVectorImpl<Value> &ubs,
@@ -102,45 +32,20 @@ static void unpackRanges(ArrayRef<SubViewOp::Range> ranges,
}
}
-template <>
-GenericLoopNestRangeBuilder<scf::ForOp>::GenericLoopNestRangeBuilder(
- MutableArrayRef<Value> ivs, ArrayRef<SubViewOp::Range> ranges) {
- builder = std::make_unique<LoopNestRangeBuilder>(ivs, ranges);
-}
-
-template <>
-GenericLoopNestRangeBuilder<AffineForOp>::GenericLoopNestRangeBuilder(
- MutableArrayRef<Value> ivs, ArrayRef<SubViewOp::Range> ranges) {
- SmallVector<Value, 4> lbs, ubs, steps;
- unpackRanges(ranges, lbs, ubs, steps);
- SmallVector<int64_t, 4> constantSteps;
- constantSteps.reserve(steps.size());
- for (Value v : steps) {
- auto op = v.getDefiningOp<ConstantIndexOp>();
- assert(op && "Affine loops require constant steps");
- constantSteps.push_back(op.getValue());
- }
- builder =
- std::make_unique<AffineLoopNestBuilder>(ivs, lbs, ubs, constantSteps);
-}
+namespace mlir {
+namespace edsc {
template <>
-GenericLoopNestRangeBuilder<scf::ParallelOp>::GenericLoopNestRangeBuilder(
+GenericLoopNestRangeBuilder<scf::ForOp>::GenericLoopNestRangeBuilder(
MutableArrayRef<Value> ivs, ArrayRef<SubViewOp::Range> ranges) {
SmallVector<Value, 4> lbs, ubs, steps;
unpackRanges(ranges, lbs, ubs, steps);
- builder = std::make_unique<ParallelLoopNestBuilder>(ivs, lbs, ubs, steps);
-}
-
-template <>
-GenericLoopNestRangeBuilder<scf::ForOp>::GenericLoopNestRangeBuilder(
- MutableArrayRef<Value> ivs, ArrayRef<Value> ranges) {
- builder = std::make_unique<LoopNestRangeBuilder>(ivs, ranges);
+ builder = std::make_unique<LoopNestBuilder>(ivs, lbs, ubs, steps);
}
template <>
GenericLoopNestRangeBuilder<AffineForOp>::GenericLoopNestRangeBuilder(
- MutableArrayRef<Value> ivs, ArrayRef<Value> ranges) {
+ MutableArrayRef<Value> ivs, ArrayRef<SubViewOp::Range> ranges) {
SmallVector<Value, 4> lbs, ubs, steps;
unpackRanges(ranges, lbs, ubs, steps);
SmallVector<int64_t, 4> constantSteps;
@@ -156,7 +61,7 @@ GenericLoopNestRangeBuilder<AffineForOp>::GenericLoopNestRangeBuilder(
template <>
GenericLoopNestRangeBuilder<scf::ParallelOp>::GenericLoopNestRangeBuilder(
- MutableArrayRef<Value> ivs, ArrayRef<Value> ranges) {
+ MutableArrayRef<Value> ivs, ArrayRef<SubViewOp::Range> ranges) {
SmallVector<Value, 4> lbs, ubs, steps;
unpackRanges(ranges, lbs, ubs, steps);
builder = std::make_unique<ParallelLoopNestBuilder>(ivs, lbs, ubs, steps);
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Loops.cpp b/mlir/lib/Dialect/Linalg/Transforms/Loops.cpp
index 77947ba1101b..74da63dafee3 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Loops.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Loops.cpp
@@ -61,19 +61,17 @@ static SmallVector<Value, 4> permuteIvs(ArrayRef<Value> ivs,
// Creates a number of ranges equal to the number of results in `map`.
// The returned ranges correspond to the loop ranges, in the proper order, for
// which new loops will be created.
-static SmallVector<Value, 4> emitLoopRanges(OpBuilder &b, Location loc,
- AffineMap map,
- ArrayRef<Value> allViewSizes);
-SmallVector<Value, 4> emitLoopRanges(OpBuilder &b, Location loc, AffineMap map,
- ArrayRef<Value> allViewSizes) {
+static SmallVector<SubViewOp::Range, 4>
+emitLoopRanges(OpBuilder &b, Location loc, AffineMap map,
+ ArrayRef<Value> allViewSizes) {
// Apply `map` to get view sizes in loop order.
auto sizes = applyMapToValues(b, loc, map, allViewSizes);
// Create a new range with the applied tile sizes.
ScopedContext scope(b, loc);
- SmallVector<Value, 4> res;
+ SmallVector<SubViewOp::Range, 4> res;
for (unsigned idx = 0, e = map.getNumResults(); idx < e; ++idx) {
- res.push_back(
- linalg_range(std_constant_index(0), sizes[idx], std_constant_index(1)));
+ res.push_back(SubViewOp::Range{std_constant_index(0), sizes[idx],
+ std_constant_index(1)});
}
return res;
}
@@ -498,7 +496,7 @@ class GenerateLoopNest {
using IndexedValueTy =
typename std::conditional<std::is_same<LoopTy, AffineForOp>::value,
AffineIndexedValue, StdIndexedValue>::type;
- static void doit(ConcreteOpTy linalgOp, ArrayRef<Value> loopRanges,
+ static void doit(ConcreteOpTy linalgOp, ArrayRef<SubViewOp::Range> loopRanges,
MutableArrayRef<Value> allIvs) {
GenericLoopNestRangeBuilder<LoopTy>(allIvs, loopRanges)([&] {
SmallVector<Value, 4> allIvValues(allIvs.begin(), allIvs.end());
@@ -517,7 +515,7 @@ class GenerateLoopNest<scf::ParallelOp, ConcreteOpTy> {
public:
using IndexedValueTy = StdIndexedValue;
- static void doit(ConcreteOpTy linalgOp, ArrayRef<Value> loopRanges,
+ static void doit(ConcreteOpTy linalgOp, ArrayRef<SubViewOp::Range> loopRanges,
MutableArrayRef<Value> allIvs) {
// Only generate scf.parallel for outer consecutive "parallel"
// iterator_types.
diff --git a/mlir/lib/Dialect/SCF/EDSC/Builders.cpp b/mlir/lib/Dialect/SCF/EDSC/Builders.cpp
index 311a7f79a67c..4ce701c1d7f9 100644
--- a/mlir/lib/Dialect/SCF/EDSC/Builders.cpp
+++ b/mlir/lib/Dialect/SCF/EDSC/Builders.cpp
@@ -67,8 +67,7 @@ mlir::edsc::LoopNestBuilder::LoopNestBuilder(Value *iv, Value lb, Value ub,
loops.emplace_back(makeLoopBuilder(iv, lb, ub, step, noArgs, {}));
}
-Operation::result_range
-mlir::edsc::LoopNestBuilder::LoopNestBuilder::operator()(
+ValueRange mlir::edsc::LoopNestBuilder::LoopNestBuilder::operator()(
std::function<void(void)> fun) {
if (fun)
fun();
@@ -76,7 +75,9 @@ mlir::edsc::LoopNestBuilder::LoopNestBuilder::operator()(
for (auto &lit : reverse(loops))
lit({});
- return loops[0].getOp()->getResults();
+ if (!loops.empty())
+ return loops[0].getOp()->getResults();
+ return {};
}
LoopBuilder mlir::edsc::makeParallelLoopBuilder(MutableArrayRef<Value> ivs,
More information about the Mlir-commits
mailing list