Automated rollback of commit cc202b8a7099f3523b079e204f2cdf3d51c404cc
PiperOrigin-RevId: 275461067
diff --git a/bindings/python/pybind.cpp b/bindings/python/pybind.cpp
index cc8fc6e..258455d 100644
--- a/bindings/python/pybind.cpp
+++ b/bindings/python/pybind.cpp
@@ -306,7 +306,7 @@
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;
}
@@ -318,7 +318,7 @@
PythonValueHandle lb, ub;
int64_t step;
- LoopBuilder *builder = nullptr;
+ AffineLoopNestBuilder *builder = nullptr;
};
struct PythonLoopNestContext {
diff --git a/g3doc/EDSC.md b/g3doc/EDSC.md
index dfbbce2..264dcaf 100644
--- a/g3doc/EDSC.md
+++ b/g3doc/EDSC.md
@@ -40,8 +40,8 @@
## 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());
@@ -53,10 +53,10 @@
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;
@@ -99,10 +99,10 @@
``` {.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>
@@ -119,6 +119,9 @@
// 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.
diff --git a/include/mlir/EDSC/Builders.h b/include/mlir/EDSC/Builders.h
index 9fbca89..5a80571 100644
--- a/include/mlir/EDSC/Builders.h
+++ b/include/mlir/EDSC/Builders.h
@@ -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"
@@ -161,8 +162,14 @@
/// 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;
@@ -172,7 +179,10 @@
/// 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
@@ -200,15 +210,34 @@
/// ```
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 {};
diff --git a/lib/EDSC/Builders.cpp b/lib/EDSC/Builders.cpp
index ba14a14..f1dceec 100644
--- a/lib/EDSC/Builders.cpp
+++ b/lib/EDSC/Builders.cpp
@@ -162,12 +162,12 @@
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());
@@ -177,11 +177,24 @@
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:
@@ -203,7 +216,12 @@
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(
@@ -212,13 +230,12 @@
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();
@@ -227,10 +244,32 @@
// 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) {
diff --git a/lib/EDSC/CMakeLists.txt b/lib/EDSC/CMakeLists.txt
index 5f3dd9f..967d6ad 100644
--- a/lib/EDSC/CMakeLists.txt
+++ b/lib/EDSC/CMakeLists.txt
@@ -17,6 +17,7 @@
target_link_libraries(MLIREDSC
PUBLIC
MLIRAffineOps
+ MLIRLoopOps
MLIRStandardOps
MLIRTransformUtils
MLIRVectorOps
diff --git a/test/EDSC/CMakeLists.txt b/test/EDSC/CMakeLists.txt
index 2820b90..4c739f8 100644
--- a/test/EDSC/CMakeLists.txt
+++ b/test/EDSC/CMakeLists.txt
@@ -7,8 +7,9 @@
target_link_libraries(mlir-edsc-builder-api-test
PRIVATE
MLIRAffineOps
- MLIRIR
MLIREDSC
+ MLIRIR
+ MLIRLoopOps
MLIRStandardOps
MLIRTransforms
LLVMCore
@@ -19,6 +20,7 @@
whole_archive_link(mlir-edsc-builder-api-test
MLIRAffineOps
+ MLIRLoopOps
MLIRStandardOps
MLIRTransforms
)
diff --git a/test/EDSC/builder-api-test.cpp b/test/EDSC/builder-api-test.cpp
index 73a7366..f08e083 100644
--- a/test/EDSC/builder-api-test.cpp
+++ b/test/EDSC/builder-api-test.cpp
@@ -70,10 +70,10 @@
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;
@@ -118,7 +118,7 @@
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) {
@@ -130,6 +130,30 @@
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;
@@ -142,7 +166,7 @@
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
@@ -344,10 +368,10 @@
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);
});
});
@@ -699,7 +723,7 @@
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
@@ -733,7 +757,7 @@
IndexHandle iv;
// clang-format off
- LoopBuilder(&iv, zero, one, 1)([&]{
+ AffineLoopNestBuilder(&iv, zero, one, 1)([&]{
res() = input();
});
// clang-format on
