maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

[mlir][Linalg] NFC - Simplify GenericNestLoop builder 

Summary: This revision trims unnecessary complexity.

Differential Revision: https://reviews.llvm.org/D80290
This commit is contained in:
Nicolas Vasilache 2020-05-20 09:43:58 -04:00
parent 004a3d4f56
commit 19e5b2bccb
5 changed files with 21 additions and 156 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

43
mlir/include/mlir/Dialect/Linalg/EDSC/Builders.h
View File

@ -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 @@ public:
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,

2
mlir/include/mlir/Dialect/SCF/EDSC/Builders.h
View File

@ -65,7 +65,7 @@ public:
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;

107
mlir/lib/Dialect/Linalg/EDSC/Builders.cpp
View File

@ -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,10 +32,15 @@ static void unpackRanges(ArrayRef<SubViewOp::Range> ranges,
}
}
namespace mlir {
namespace edsc {
template <>
GenericLoopNestRangeBuilder<scf::ForOp>::GenericLoopNestRangeBuilder(
MutableArrayRef<Value> ivs, ArrayRef<SubViewOp::Range> ranges) {
builder = std::make_unique<LoopNestRangeBuilder>(ivs, ranges);
SmallVector<Value, 4> lbs, ubs, steps;
unpackRanges(ranges, lbs, ubs, steps);
builder = std::make_unique<LoopNestBuilder>(ivs, lbs, ubs, steps);
}
template <>
@ -132,36 +67,6 @@ GenericLoopNestRangeBuilder<scf::ParallelOp>::GenericLoopNestRangeBuilder(
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);
}
template <>
GenericLoopNestRangeBuilder<AffineForOp>::GenericLoopNestRangeBuilder(
MutableArrayRef<Value> ivs, ArrayRef<Value> 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);
}
template <>
GenericLoopNestRangeBuilder<scf::ParallelOp>::GenericLoopNestRangeBuilder(
MutableArrayRef<Value> ivs, ArrayRef<Value> ranges) {
SmallVector<Value, 4> lbs, ubs, steps;
unpackRanges(ranges, lbs, ubs, steps);
builder = std::make_unique<ParallelLoopNestBuilder>(ivs, lbs, ubs, steps);
}
} // namespace edsc
} // namespace mlir

18
mlir/lib/Dialect/Linalg/Transforms/Loops.cpp
View File

@ -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 @@ public:
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.

7
mlir/lib/Dialect/SCF/EDSC/Builders.cpp
View File

@ -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,