diff --git a/mlir/include/mlir/Dialect/Linalg/EDSC/Builders.h b/mlir/include/mlir/Dialect/Linalg/EDSC/Builders.h --- a/mlir/include/mlir/Dialect/Linalg/EDSC/Builders.h +++ b/mlir/include/mlir/Dialect/Linalg/EDSC/Builders.h @@ -32,60 +32,21 @@ 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 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 ivs, ArrayRef ranges); - LoopNestRangeBuilder(MutableArrayRef ivs, - ArrayRef ranges); - Value operator()(std::function fun = nullptr); - -private: - SmallVector loops; -}; - /// Helper template class for building scf.for and affine.loop nests from /// ranges. template class GenericLoopNestRangeBuilder { public: GenericLoopNestRangeBuilder(MutableArrayRef ivs, - ArrayRef ranges); - GenericLoopNestRangeBuilder(MutableArrayRef ivs, ArrayRef ranges); void operator()(std::function fun = nullptr) { (*builder)(fun); } private: using LoopOrAffineLoopBuilder = typename std::conditional_t::value, - AffineLoopNestBuilder, LoopNestRangeBuilder>; + AffineLoopNestBuilder, LoopNestBuilder>; using BuilderType = typename std::conditional_t::value, ParallelLoopNestBuilder, diff --git a/mlir/include/mlir/Dialect/SCF/EDSC/Builders.h b/mlir/include/mlir/Dialect/SCF/EDSC/Builders.h --- a/mlir/include/mlir/Dialect/SCF/EDSC/Builders.h +++ b/mlir/include/mlir/Dialect/SCF/EDSC/Builders.h @@ -65,7 +65,7 @@ ValueRange iterArgsInitValues); LoopNestBuilder(MutableArrayRef ivs, ArrayRef lbs, ArrayRef ubs, ArrayRef steps); - Operation::result_range operator()(std::function fun = nullptr); + ValueRange operator()(std::function fun = nullptr); private: SmallVector loops; diff --git a/mlir/lib/Dialect/Linalg/EDSC/Builders.cpp b/mlir/lib/Dialect/Linalg/EDSC/Builders.cpp --- a/mlir/lib/Dialect/Linalg/EDSC/Builders.cpp +++ b/mlir/lib/Dialect/Linalg/EDSC/Builders.cpp @@ -21,76 +21,6 @@ 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(range.getDefiningOp()); - auto lb = rangeOp.min(); - auto ub = rangeOp.max(); - auto step = rangeOp.step(); - ForOp forOp = OperationBuilder(lb, ub, step); - *iv = forOp.getInductionVar(); - auto *body = forOp.getBody(); - enter(body); -} - -mlir::edsc::LoopRangeBuilder::LoopRangeBuilder(Value *iv, - SubViewOp::Range range) { - ForOp forOp = OperationBuilder(range.offset, range.size, range.stride); - *iv = forOp.getInductionVar(); - auto *body = forOp.getBody(); - enter(body); -} - -Value mlir::edsc::LoopRangeBuilder::operator()(std::function fun) { - if (fun) - fun(); - exit(); - return Value(); -} - -mlir::edsc::LoopNestRangeBuilder::LoopNestRangeBuilder( - MutableArrayRef ivs, ArrayRef 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 ivs, ArrayRef 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 fun) { - if (fun) - fun(); - for (auto &lit : reverse(loops)) { - lit({}); - } - return Value(); -} - -namespace mlir { -namespace edsc { - -static void unpackRanges(ArrayRef rangeOps, SmallVectorImpl &lbs, - SmallVectorImpl &ubs, - SmallVectorImpl &steps) { - for (Value range : rangeOps) { - assert(range.getType() && "expected linalg.range type"); - assert(range.getDefiningOp() && "need operations to extract range parts"); - RangeOp rangeOp = cast(range.getDefiningOp()); - lbs.emplace_back(rangeOp.min()); - ubs.emplace_back(rangeOp.max()); - steps.emplace_back(rangeOp.step()); - } -} - static void unpackRanges(ArrayRef ranges, SmallVectorImpl &lbs, SmallVectorImpl &ubs, @@ -102,45 +32,20 @@ } } -template <> -GenericLoopNestRangeBuilder::GenericLoopNestRangeBuilder( - MutableArrayRef ivs, ArrayRef ranges) { - builder = std::make_unique(ivs, ranges); -} - -template <> -GenericLoopNestRangeBuilder::GenericLoopNestRangeBuilder( - MutableArrayRef ivs, ArrayRef ranges) { - SmallVector lbs, ubs, steps; - unpackRanges(ranges, lbs, ubs, steps); - SmallVector constantSteps; - constantSteps.reserve(steps.size()); - for (Value v : steps) { - auto op = v.getDefiningOp(); - assert(op && "Affine loops require constant steps"); - constantSteps.push_back(op.getValue()); - } - builder = - std::make_unique(ivs, lbs, ubs, constantSteps); -} +namespace mlir { +namespace edsc { template <> -GenericLoopNestRangeBuilder::GenericLoopNestRangeBuilder( +GenericLoopNestRangeBuilder::GenericLoopNestRangeBuilder( MutableArrayRef ivs, ArrayRef ranges) { SmallVector lbs, ubs, steps; unpackRanges(ranges, lbs, ubs, steps); - builder = std::make_unique(ivs, lbs, ubs, steps); -} - -template <> -GenericLoopNestRangeBuilder::GenericLoopNestRangeBuilder( - MutableArrayRef ivs, ArrayRef ranges) { - builder = std::make_unique(ivs, ranges); + builder = std::make_unique(ivs, lbs, ubs, steps); } template <> GenericLoopNestRangeBuilder::GenericLoopNestRangeBuilder( - MutableArrayRef ivs, ArrayRef ranges) { + MutableArrayRef ivs, ArrayRef ranges) { SmallVector lbs, ubs, steps; unpackRanges(ranges, lbs, ubs, steps); SmallVector constantSteps; @@ -156,7 +61,7 @@ template <> GenericLoopNestRangeBuilder::GenericLoopNestRangeBuilder( - MutableArrayRef ivs, ArrayRef ranges) { + MutableArrayRef ivs, ArrayRef ranges) { SmallVector lbs, ubs, steps; unpackRanges(ranges, lbs, ubs, steps); builder = std::make_unique(ivs, lbs, ubs, steps); diff --git a/mlir/lib/Dialect/Linalg/Transforms/Loops.cpp b/mlir/lib/Dialect/Linalg/Transforms/Loops.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Loops.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Loops.cpp @@ -61,19 +61,17 @@ // 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 emitLoopRanges(OpBuilder &b, Location loc, - AffineMap map, - ArrayRef allViewSizes); -SmallVector emitLoopRanges(OpBuilder &b, Location loc, AffineMap map, - ArrayRef allViewSizes) { +static SmallVector +emitLoopRanges(OpBuilder &b, Location loc, AffineMap map, + ArrayRef 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 res; + SmallVector 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 @@ using IndexedValueTy = typename std::conditional::value, AffineIndexedValue, StdIndexedValue>::type; - static void doit(ConcreteOpTy linalgOp, ArrayRef loopRanges, + static void doit(ConcreteOpTy linalgOp, ArrayRef loopRanges, MutableArrayRef allIvs) { GenericLoopNestRangeBuilder(allIvs, loopRanges)([&] { SmallVector allIvValues(allIvs.begin(), allIvs.end()); @@ -517,7 +515,7 @@ public: using IndexedValueTy = StdIndexedValue; - static void doit(ConcreteOpTy linalgOp, ArrayRef loopRanges, + static void doit(ConcreteOpTy linalgOp, ArrayRef loopRanges, MutableArrayRef 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 --- a/mlir/lib/Dialect/SCF/EDSC/Builders.cpp +++ b/mlir/lib/Dialect/SCF/EDSC/Builders.cpp @@ -67,8 +67,7 @@ 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 fun) { if (fun) fun(); @@ -76,7 +75,9 @@ 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 ivs,