Skip to content

Commit

Permalink
Add EDSC support for loop.for operations
Browse files Browse the repository at this point in the history 
This CL adds support for loop.for operations in EDSC and adds a test.
This will be used in a followup commit to implement lowering of vector_transfer ops so that it works more generally and is not subject to affine constraints.

PiperOrigin-RevId: 275349796
  • Loading branch information
@tensorflower-gardener
Nicolas Vasilache authored and tensorflower-gardener committed Oct 17, 2019
1 parent dae0ae6 commit b65c8bb
Show file tree
Hide file tree
Showing 6 changed files with 135 additions and 39 deletions.
4 changes: 2 additions & 2 deletions mlir/bindings/python/pybind.cpp
View file
Expand Up @@ -306,7 +306,7 @@ struct PythonLoopContext {

PythonValueHandle enter() {
ValueHandle iv(lb.value.getType());
builder = new LoopBuilder(&iv, lb.value, ub.value, step);
builder = new AffineLoopNestBuilder(&iv, lb.value, ub.value, step);
return iv;
}

Expand All @@ -318,7 +318,7 @@ struct PythonLoopContext {

PythonValueHandle lb, ub;
int64_t step;
LoopBuilder *builder = nullptr;
AffineLoopNestBuilder *builder = nullptr;
};

struct PythonLoopNestContext {
Expand Down
19 changes: 11 additions & 8 deletions mlir/g3doc/EDSC.md
View file
Expand Up @@ -40,8 +40,8 @@ using store = InstructionBuilder<StoreOp>;

## LoopBuilder and AffineLoopNestBuilder

`mlir::edsc::AffineLoopNestBuilder` provides an interface to allow writing concise and
structured loop nests.
`mlir::edsc::AffineLoopNestBuilder` provides an interface to allow writing
concise and structured loop nests.

```c++
ScopedContext scope(f.get());
Expand All @@ -53,10 +53,10 @@ structured loop nests.
f13(constant_float(llvm::APFloat(13.0f), f32Type)),
i7(constant_int(7, 32)),
i13(constant_int(13, 32));
LoopBuilder(&i, lb, ub, 3)([&]{
AffineLoopNestBuilder(&i, lb, ub, 3)([&]{
lb * index_t(3) + ub;
lb + index_t(3);
LoopBuilder(&j, lb, ub, 2)([&]{
AffineLoopNestBuilder(&j, lb, ub, 2)([&]{
ceilDiv(index_t(31) * floorDiv(i + j * index_t(3), index_t(32)),
index_t(32));
((f7 + f13) / f7) % f13 - f7 * f13;
Expand Down Expand Up @@ -99,10 +99,10 @@ generated IR resembles the following, for a 4-D memref of `vector<4xi8>`:

``` {.mlir}
// CHECK-LABEL: func @t1(%lhs: memref<3x4x5x6xvector<4xi8>>, %rhs: memref<3x4x5x6xvector<4xi8>>, %result: memref<3x4x5x6xvector<4xi8>>) -> () {
// CHECK: for {{.*}} = 0 to 3 {
// CHECK: for {{.*}} = 0 to 4 {
// CHECK: for {{.*}} = 0 to 5 {
// CHECK: for {{.*}}= 0 to 6 {
// CHECK: affine.for {{.*}} = 0 to 3 {
// CHECK: affine.for {{.*}} = 0 to 4 {
// CHECK: affine.for {{.*}} = 0 to 5 {
// CHECK: affine.for {{.*}}= 0 to 6 {
// CHECK: {{.*}} = load %arg1[{{.*}}] : memref<3x4x5x6xvector<4xi8>>
// CHECK: {{.*}} = load %arg0[{{.*}}] : memref<3x4x5x6xvector<4xi8>>
// CHECK: {{.*}} = addi {{.*}} : vector<4xi8>
Expand All @@ -119,6 +119,9 @@ or the following, for a 0-D `memref<f32>`:
// CHECK: store {{.*}}, %arg2[] : memref<f32>
```

Similar APIs are provided to emit the lower-level `loop.for` op with
`LoopNestBuilder`. See the `builder-api-test.cpp` test for more usage examples.

Since the implementation of declarative builders is in C++, it is also available
to program the IR with an embedded-DSL flavor directly integrated in MLIR. We
make use of these properties in the tutorial.
Expand Down
37 changes: 33 additions & 4 deletions mlir/include/mlir/EDSC/Builders.h
View file
Expand Up @@ -24,6 +24,7 @@
#define MLIR_EDSC_BUILDERS_H_

#include "mlir/Dialect/AffineOps/AffineOps.h"
#include "mlir/Dialect/LoopOps/LoopOps.h"
#include "mlir/Dialect/StandardOps/Ops.h"
#include "mlir/Dialect/VectorOps/VectorOps.h"
#include "mlir/IR/Builders.h"
Expand Down Expand Up @@ -161,8 +162,14 @@ class LoopBuilder : public NestedBuilder {
/// Constructs a new AffineForOp and captures the associated induction
/// variable. A ValueHandle pointer is passed as the first argument and is the
/// *only* way to capture the loop induction variable.
LoopBuilder(ValueHandle *iv, ArrayRef<ValueHandle> lbHandles,
ArrayRef<ValueHandle> ubHandles, int64_t step);
static LoopBuilder makeAffine(ValueHandle *iv,
ArrayRef<ValueHandle> lbHandles,
ArrayRef<ValueHandle> ubHandles, int64_t step);
/// Constructs a new loop::ForOp and captures the associated induction
/// variable. A ValueHandle pointer is passed as the first argument and is the
/// *only* way to capture the loop induction variable.
static LoopBuilder makeLoop(ValueHandle *iv, ValueHandle lbHandle,
ValueHandle ubHandle, ValueHandle stepHandle);
LoopBuilder(const LoopBuilder &) = delete;
LoopBuilder(LoopBuilder &&) = default;

Expand All @@ -172,7 +179,10 @@ class LoopBuilder : public NestedBuilder {
/// 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 LoopBuilder.
ValueHandle operator()(llvm::function_ref<void(void)> fun = nullptr);
void operator()(llvm::function_ref<void(void)> fun = nullptr);

private:
LoopBuilder() = default;
};

/// Explicit nested LoopBuilder. Offers a compressed multi-loop builder to avoid
Expand Down Expand Up @@ -200,15 +210,34 @@ class LoopBuilder : public NestedBuilder {
/// ```
class AffineLoopNestBuilder {
public:
// This entry point accomodates the fact that AffineForOp implicitly uses
// multiple `lbs` and `ubs` with one single `iv` and `step` to encode `max`
// and and `min` constraints respectively.
AffineLoopNestBuilder(ValueHandle *iv, ArrayRef<ValueHandle> lbs,
ArrayRef<ValueHandle> ubs, int64_t step);
AffineLoopNestBuilder(ArrayRef<ValueHandle *> ivs, ArrayRef<ValueHandle> lbs,
ArrayRef<ValueHandle> ubs, ArrayRef<int64_t> steps);

ValueHandle operator()(llvm::function_ref<void(void)> fun = nullptr);
void operator()(llvm::function_ref<void(void)> fun = nullptr);

private:
SmallVector<LoopBuilder, 4> loops;
};

/// Helper class to sugar building loop.for loop nests from ranges.
/// This is similar to edsc::AffineLoopNestBuilder except it operates on
/// loop.for.
class LoopNestBuilder {
public:
LoopNestBuilder(llvm::ArrayRef<edsc::ValueHandle *> ivs,
ArrayRef<ValueHandle> lbs, ArrayRef<ValueHandle> ubs,
ArrayRef<ValueHandle> steps);
void operator()(std::function<void(void)> fun = nullptr);

private:
llvm::SmallVector<LoopBuilder, 4> loops;
};

// This class exists solely to handle the C++ vexing parse case when
// trying to enter a Block that has already been constructed.
class Append {};
Expand Down
73 changes: 56 additions & 17 deletions mlir/lib/EDSC/Builders.cpp
View file
Expand Up @@ -162,12 +162,12 @@ static llvm::Optional<ValueHandle> emitStaticFor(ArrayRef<ValueHandle> lbs,
ubConst.getValue(), step);
}

mlir::edsc::LoopBuilder::LoopBuilder(ValueHandle *iv,
ArrayRef<ValueHandle> lbHandles,
ArrayRef<ValueHandle> ubHandles,
int64_t step) {
if (auto res = emitStaticFor(lbHandles, ubHandles, step)) {
*iv = res.getValue();
mlir::edsc::LoopBuilder mlir::edsc::LoopBuilder::makeAffine(
ValueHandle *iv, ArrayRef<ValueHandle> lbHandles,
ArrayRef<ValueHandle> ubHandles, int64_t step) {
mlir::edsc::LoopBuilder result;
if (auto staticFor = emitStaticFor(lbHandles, ubHandles, step)) {
*iv = staticFor.getValue();
} else {
SmallVector<Value *, 4> lbs(lbHandles.begin(), lbHandles.end());
SmallVector<Value *, 4> ubs(ubHandles.begin(), ubHandles.end());
Expand All @@ -177,11 +177,24 @@ mlir::edsc::LoopBuilder::LoopBuilder(ValueHandle *iv,
step);
}
auto *body = getForInductionVarOwner(iv->getValue()).getBody();
enter(body, /*prev=*/1);
result.enter(body, /*prev=*/1);
return result;
}

ValueHandle
mlir::edsc::LoopBuilder::operator()(llvm::function_ref<void(void)> fun) {
mlir::edsc::LoopBuilder
mlir::edsc::LoopBuilder::makeLoop(ValueHandle *iv, ValueHandle lbHandle,
ValueHandle ubHandle,
ValueHandle stepHandle) {
mlir::edsc::LoopBuilder result;
auto forOp =
OperationHandle::createOp<loop::ForOp>(lbHandle, ubHandle, stepHandle);
*iv = ValueHandle(forOp.getInductionVar());
auto *body = loop::getForInductionVarOwner(iv->getValue()).getBody();
result.enter(body, /*prev=*/1);
return result;
}

void mlir::edsc::LoopBuilder::operator()(llvm::function_ref<void(void)> fun) {
// Call to `exit` must be explicit and asymmetric (cannot happen in the
// destructor) because of ordering wrt comma operator.
/// The particular use case concerns nested blocks:
Expand All @@ -203,7 +216,12 @@ mlir::edsc::LoopBuilder::operator()(llvm::function_ref<void(void)> fun) {
if (fun)
fun();
exit();
return ValueHandle::null();
}

mlir::edsc::AffineLoopNestBuilder::AffineLoopNestBuilder(
ValueHandle *iv, ArrayRef<ValueHandle> lbs, ArrayRef<ValueHandle> ubs,
int64_t step) {
loops.emplace_back(LoopBuilder::makeAffine(iv, lbs, ubs, step));
}

mlir::edsc::AffineLoopNestBuilder::AffineLoopNestBuilder(
Expand All @@ -212,13 +230,12 @@ mlir::edsc::AffineLoopNestBuilder::AffineLoopNestBuilder(
assert(ivs.size() == lbs.size() && "Mismatch in number of arguments");
assert(ivs.size() == ubs.size() && "Mismatch in number of arguments");
assert(ivs.size() == steps.size() && "Mismatch in number of arguments");
for (auto it : llvm::zip(ivs, lbs, ubs, steps)) {
loops.emplace_back(std::get<0>(it), std::get<1>(it), std::get<2>(it),
std::get<3>(it));
}
for (auto it : llvm::zip(ivs, lbs, ubs, steps))
loops.emplace_back(LoopBuilder::makeAffine(
std::get<0>(it), std::get<1>(it), std::get<2>(it), std::get<3>(it)));
}

ValueHandle mlir::edsc::AffineLoopNestBuilder::operator()(
void mlir::edsc::AffineLoopNestBuilder::operator()(
llvm::function_ref<void(void)> fun) {
if (fun)
fun();
Expand All @@ -227,10 +244,32 @@ ValueHandle mlir::edsc::AffineLoopNestBuilder::operator()(
// to be asymmetric (i.e. enter() occurs on LoopBuilder construction, exit()
// occurs on calling operator()). The asymmetry is required for properly
// nesting imperfectly nested regions (see LoopBuilder::operator()).
for (auto lit = loops.rbegin(), eit = loops.rend(); lit != eit; ++lit) {
for (auto lit = loops.rbegin(), eit = loops.rend(); lit != eit; ++lit)
(*lit)();
}

mlir::edsc::LoopNestBuilder::LoopNestBuilder(ArrayRef<ValueHandle *> ivs,
ArrayRef<ValueHandle> lbs,
ArrayRef<ValueHandle> ubs,
ArrayRef<ValueHandle> steps) {
assert(ivs.size() == lbs.size() && "expected size of ivs and lbs to match");
assert(ivs.size() == ubs.size() && "expected size of ivs and ubs to match");
assert(ivs.size() == steps.size() &&
"expected size of ivs and steps to match");
loops.reserve(ivs.size());
for (auto it : llvm::zip(ivs, lbs, ubs, steps)) {
loops.emplace_back(LoopBuilder::makeLoop(std::get<0>(it), std::get<1>(it),
std::get<2>(it), std::get<3>(it)));
}
return ValueHandle::null();
assert(loops.size() == ivs.size() && "Mismatch loops vs ivs size");
}

void LoopNestBuilder::LoopNestBuilder::operator()(
std::function<void(void)> fun) {
if (fun)
fun();
for (auto &lit : reverse(loops))
lit({});
}

mlir::edsc::BlockBuilder::BlockBuilder(BlockHandle bh, Append) {
Expand Down
1 change: 1 addition & 0 deletions mlir/lib/EDSC/CMakeLists.txt
View file
Expand Up @@ -17,6 +17,7 @@ add_dependencies(MLIREDSC MLIRReferenceImplementationTestGen)
target_link_libraries(MLIREDSC
PUBLIC
MLIRAffineOps
MLIRLoopOps
MLIRStandardOps
MLIRTransformUtils
MLIRVectorOps
Expand Down
40 changes: 32 additions & 8 deletions mlir/test/EDSC/builder-api-test.cpp
View file
Expand Up @@ -70,10 +70,10 @@ TEST_FUNC(builder_dynamic_for_func_args) {
ValueHandle f13(constant_float(llvm::APFloat(13.0f), f32Type));
ValueHandle i7(constant_int(7, 32));
ValueHandle i13(constant_int(13, 32));
LoopBuilder(&i, lb, ub, 3)([&] {
AffineLoopNestBuilder(&i, lb, ub, 3)([&] {
lb *index_t(3) + ub;
lb + index_t(3);
LoopBuilder(&j, lb, ub, 2)([&] {
AffineLoopNestBuilder(&j, lb, ub, 2)([&] {
ceilDiv(index_t(31) * floorDiv(i + j * index_t(3), index_t(32)),
index_t(32));
((f7 + f13) / f7) % f13 - f7 *f13;
Expand Down Expand Up @@ -118,7 +118,7 @@ TEST_FUNC(builder_dynamic_for) {
ScopedContext scope(builder, f.getLoc());
ValueHandle i(indexType), a(f.getArgument(0)), b(f.getArgument(1)),
c(f.getArgument(2)), d(f.getArgument(3));
LoopBuilder(&i, a - b, c + d, 2)();
AffineLoopNestBuilder(&i, a - b, c + d, 2)();

// clang-format off
// CHECK-LABEL: func @builder_dynamic_for(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
Expand All @@ -130,6 +130,30 @@ TEST_FUNC(builder_dynamic_for) {
f.erase();
}

TEST_FUNC(builder_loop_for) {
using namespace edsc;
using namespace edsc::op;
using namespace edsc::intrinsics;
auto indexType = IndexType::get(&globalContext());
auto f = makeFunction("builder_loop_for", {},
{indexType, indexType, indexType, indexType});

OpBuilder builder(f.getBody());
ScopedContext scope(builder, f.getLoc());
ValueHandle i(indexType), a(f.getArgument(0)), b(f.getArgument(1)),
c(f.getArgument(2)), d(f.getArgument(3));
LoopNestBuilder(&i, a - b, c + d, a)();

// clang-format off
// CHECK-LABEL: func @builder_loop_for(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
// CHECK-DAG: [[r0:%[0-9]+]] = affine.apply ()[s0, s1] -> (s0 - s1)()[%{{.*}}, %{{.*}}]
// CHECK-DAG: [[r1:%[0-9]+]] = affine.apply ()[s0, s1] -> (s0 + s1)()[%{{.*}}, %{{.*}}]
// CHECK-NEXT: loop.for %{{.*}} = [[r0]] to [[r1]] step {{.*}} {
// clang-format on
f.print(llvm::outs());
f.erase();
}

TEST_FUNC(builder_max_min_for) {
using namespace edsc;
using namespace edsc::op;
Expand All @@ -142,7 +166,7 @@ TEST_FUNC(builder_max_min_for) {
ScopedContext scope(builder, f.getLoc());
ValueHandle i(indexType), lb1(f.getArgument(0)), lb2(f.getArgument(1)),
ub1(f.getArgument(2)), ub2(f.getArgument(3));
LoopBuilder(&i, {lb1, lb2}, {ub1, ub2}, 1)();
AffineLoopNestBuilder(&i, {lb1, lb2}, {ub1, ub2}, 1)();
ret();

// clang-format off
Expand Down Expand Up @@ -344,10 +368,10 @@ TEST_FUNC(builder_helpers) {
ub2 = vA.ub(2);
step2 = vA.step(2);
AffineLoopNestBuilder({&i, &j}, {lb0, lb1}, {ub0, ub1}, {step0, step1})([&]{
LoopBuilder(&k1, lb2, ub2, step2)([&]{
AffineLoopNestBuilder(&k1, lb2, ub2, step2)([&]{
C(i, j, k1) = f7 + A(i, j, k1) + B(i, j, k1);
});
LoopBuilder(&k2, lb2, ub2, step2)([&]{
AffineLoopNestBuilder(&k2, lb2, ub2, step2)([&]{
C(i, j, k2) += A(i, j, k2) + B(i, j, k2);
});
});
Expand Down Expand Up @@ -699,7 +723,7 @@ TEST_FUNC(indirect_access) {
IndexHandle i, N(vC.ub(0));

// clang-format off
LoopBuilder(&i, zero, N, 1)([&]{
AffineLoopNestBuilder(&i, zero, N, 1)([&]{
C((ValueHandle)D(i)) = A((ValueHandle)B(i));
});
// clang-format on
Expand Down Expand Up @@ -733,7 +757,7 @@ TEST_FUNC(empty_map_load_store) {
IndexHandle iv;

// clang-format off
LoopBuilder(&iv, zero, one, 1)([&]{
AffineLoopNestBuilder(&iv, zero, one, 1)([&]{
res() = input();
});
// clang-format on
Expand Down

0 comments on commit b65c8bb

Please sign in to comment.