[mlir] [EDSC] Add interface for yield-for loops.

[poechsel]

Summary

ModelBuilder was missing an api to easily generate yield-for-loops.
This diffs implements an interface allowing to write:

%2:2 = loop.for %i = %start to %end step %step iter_args(%arg0 = %init0, %arg1 = %init1) -> (f32, f32) {
  %sum = addf %arg0, %arg1 : f32
  loop.yield %arg1, %sum : f32, f32
}
%3 = addf %2#0, %2#1 : f32

as

auto results =
    LoopNestBuilder(&i, start, end, step, {&arg0, &arg1},  {init0, init1})([&] {
      auto sum = arg0 + arg1;
      loop_yield(ArrayRef<ValueHandle>{arg1, sum});
    });

// Add the two values accumulated by the yield-for-loop:
ValueHandle(results[0]) + ValueHandle(results[1]);

[poechsel]

Renames 'arguments' to 'iter_args_init_values' and 'argument_handles' to 'iter_args_handles'.
Swapped position of the arguments 'arguments' and 'argument_handles' to mimic the position in mlir.

[ftynse]

I have two concerns:

• this does not seem to work with builders that have the insertion point in the middle of a block;
• it is unclear if this is intended to work for loop _nests_.

mlir/include/mlir/Dialect/LoopOps/EDSC/Builders.h

38

Nit: ValueRange instead of ArrayRef<Value> unless you have a strong reason to do otherwise.

59

Nit: no need to prefix with edsc:: if you are already inside the edsc namespace

mlir/include/mlir/Dialect/LoopOps/EDSC/Intrinsics.h

3

Spurious leading // makes you have ////

mlir/lib/Dialect/LoopOps/EDSC/Builders.cpp

62

I don't understand this assert, and the message does not really help. Do you only support building one loop with iter args? If so, have you considered modifying LoopBuilder instead of LoopNestBuilder ?

76

While it was the last to be inserted, nothing guarantees that it's the last operation in the block. The insertion point may have been pointing in the middle of the block. I would suggest to just store the results somewhere when you create the loop operation and return them here....

77

Nit: let's add some camelBack here: forLoopOp

104

Nit: size_t i = 0, e = iter_args_handles.size(); i < e; ++i to avoid calling size on every iteration. This idiom is extremely common in LLVM.

[nicolasvasilache]

mlir/lib/Dialect/LoopOps/EDSC/Builders.cpp

62

+1, you already modify the static makeLoopBuilder method inside LoopBuilder.
It seems LoopBuilder is the natural thing to extend and we can revisit LoopNestBuilder later if/when necessary.

76

Does this just go away if we focus on just LoopBuilder?

[poechsel]

Rebase, address comments.

[poechsel]

Address the for loop comment.

[ftynse]

mlir/include/mlir/Dialect/LoopOps/EDSC/Intrinsics.h

3

Ehm, we need the inverse of the change you did. We only use two slashes //, not four ///.

[poechsel]

Simplified the implementation of LoopNestBuilder::operator().

[poechsel]

// -> in intrinsics.h

[nicolasvasilache]

Looks good, thanks Pierre!

mlir/include/mlir/Dialect/LoopOps/EDSC/Builders.h

 LoopBuilder makeParallelLoopBuilder(ArrayRef<ValueHandle *> ivs,
                                     ArrayRef<ValueHandle> lbHandles,
                                     ArrayRef<ValueHandle> ubHandles,
                                     ArrayRef<ValueHandle> steps);
 /// 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.
 LoopBuilder makeLoopBuilder(ValueHandle *iv, ValueHandle lbHandle,
-                            ValueHandle ubHandle, ValueHandle stepHandle);
+                            ValueHandle ubHandle, ValueHandle stepHandle,
+                            ArrayRef<ValueHandle *> iter_args_handles = {},
+                            ValueRange iter_args_init_values = {});
 
 /// Helper class to sugar building loop.parallel loop nests from lower/upper
 /// bounds and step sizes.
 class ParallelLoopNestBuilder {
 public:
   ParallelLoopNestBuilder(ArrayRef<ValueHandle *> ivs,
                           ArrayRef<ValueHandle> lbs, ArrayRef<ValueHandle> ubs,
                           ArrayRef<ValueHandle> steps);
 
   void operator()(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(ArrayRef<edsc::ValueHandle *> ivs, ArrayRef<ValueHandle> lbs,
+  LoopNestBuilder(ValueHandle *iv, ValueHandle lb, ValueHandle ub,
+                  ValueHandle step,
+                  ArrayRef<ValueHandle *> iter_args_handles = {},
+                  ValueRange iter_args_init_values = {});
+  LoopNestBuilder(ArrayRef<ValueHandle *> ivs, ArrayRef<ValueHandle> lbs,
                   ArrayRef<ValueHandle> ubs, ArrayRef<ValueHandle> steps);
-  void operator()(std::function<void(void)> fun = nullptr);
+  Operation::result_range operator()(std::function<void(void)> fun = nullptr);
 
 private:
   SmallVector<LoopBuilder, 4> loops;
 };
 
 } // namespace edsc
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LOOPOPS_EDSC_BUILDERS_H_

mlir/include/mlir/Dialect/LoopOps/EDSC/Intrinsics.h

+//===- Intrinsics.h - MLIR EDSC Intrinsics for Linalg -----------*- C++ -*-===//
+////
+//// Part of the LLVM Project, under the Apache License v2.0 with LLVM
+//// Exceptions.
+//// See https://llvm.org/LICENSE.txt for license information.
+//// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+////
+////===----------------------------------------------------------------------===//
+#ifndef MLIR_DIALECT_LOOPOPS_EDSC_INTRINSICS_H_
+#define MLIR_DIALECT_LOOPOPS_EDSC_INTRINSICS_H_
+
+#include "mlir/Dialect/LoopOps/EDSC/Builders.h"
+#include "mlir/EDSC/Intrinsics.h"
+
+namespace mlir {
+namespace edsc {
+namespace intrinsics {
+using loop_yield = OperationBuilder<loop::YieldOp>;
+
+} // namespace intrinsics
+} // namespace edsc
+} // namespace mlir
+
+#endif // MLIR_DIALECT_LOOPOPS_EDSC_INTRINSICS_H_

mlir/include/mlir/EDSC/Builders.h

public:
 
   LoopBuilder &operator=(const LoopBuilder &) = delete;
   LoopBuilder &operator=(LoopBuilder &&) = 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 LoopBuilder.
   void operator()(function_ref<void(void)> fun = nullptr);
+  void setOp(Operation *op) { this->op = op; }
+  Operation *getOp() { return op; }
 
 private:
   LoopBuilder() = default;
 
   friend LoopBuilder makeAffineLoopBuilder(ValueHandle *iv,
                                            ArrayRef<ValueHandle> lbHandles,
                                            ArrayRef<ValueHandle> ubHandles,
                                            int64_t step);
   friend LoopBuilder makeParallelLoopBuilder(ArrayRef<ValueHandle *> ivs,
                                              ArrayRef<ValueHandle> lbHandles,
                                              ArrayRef<ValueHandle> ubHandles,
                                              ArrayRef<ValueHandle> steps);
   friend LoopBuilder makeLoopBuilder(ValueHandle *iv, ValueHandle lbHandle,
                                      ValueHandle ubHandle,
-                                     ValueHandle stepHandle);
+                                     ValueHandle stepHandle,
+                                     ArrayRef<ValueHandle *> iter_args_handles,
+                                     ValueRange iter_args_init_values);
+  Operation *op;
 };
 
 // This class exists solely to handle the C++ vexing parse case when
 // trying to enter a Block that has already been constructed.
 class Append {};
 
 /// A BlockBuilder is a NestedBuilder for mlir::Block*.
 /// This exists by opposition to LoopBuilder which is not related to an

mlir/lib/Dialect/LoopOps/EDSC/Builders.cpp

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)) {
+  for (auto it : llvm::zip(ivs, lbs, ubs, steps))
     loops.emplace_back(makeLoopBuilder(std::get<0>(it), std::get<1>(it),
                                        std::get<2>(it), std::get<3>(it)));
-  }
   assert(loops.size() == ivs.size() && "Mismatch loops vs ivs size");
 }
 
-void mlir::edsc::LoopNestBuilder::LoopNestBuilder::operator()(
+mlir::edsc::LoopNestBuilder::LoopNestBuilder(
+    ValueHandle *iv, ValueHandle lb, ValueHandle ub, ValueHandle step,
+    ArrayRef<ValueHandle *> iter_args_handles,
+    ValueRange iter_args_init_values) {
+  assert(iter_args_init_values.size() == iter_args_handles.size() &&
+         "expected size of arguments and argument_handles to match");
+  loops.emplace_back(makeLoopBuilder(iv, lb, ub, step, iter_args_handles,
+                                     iter_args_init_values));
+}
+
+Operation::result_range
+mlir::edsc::LoopNestBuilder::LoopNestBuilder::operator()(
     std::function<void(void)> fun) {
   if (fun)
     fun();
+
   for (auto &lit : reverse(loops))
     lit({});
+
+  return loops[0].getOp()->getResults();
 }
 
 LoopBuilder mlir::edsc::makeParallelLoopBuilder(ArrayRef<ValueHandle *> ivs,
                                                 ArrayRef<ValueHandle> lbHandles,
                                                 ArrayRef<ValueHandle> ubHandles,
                                                 ArrayRef<ValueHandle> steps) {
   LoopBuilder result;
   auto opHandle = OperationHandle::create<loop::ParallelOp>(
       SmallVector<Value, 4>(lbHandles.begin(), lbHandles.end()),
       SmallVector<Value, 4>(ubHandles.begin(), ubHandles.end()),
       SmallVector<Value, 4>(steps.begin(), steps.end()));
 
   loop::ParallelOp parallelOp =
       cast<loop::ParallelOp>(*opHandle.getOperation());
   for (size_t i = 0, e = ivs.size(); i < e; ++i)
     *ivs[i] = ValueHandle(parallelOp.getBody()->getArgument(i));
   result.enter(parallelOp.getBody(), /*prev=*/1);
   return result;
 }
 
-mlir::edsc::LoopBuilder mlir::edsc::makeLoopBuilder(ValueHandle *iv,
-                                                    ValueHandle lbHandle,
-                                                    ValueHandle ubHandle,
-                                                    ValueHandle stepHandle) {
+mlir::edsc::LoopBuilder
+mlir::edsc::makeLoopBuilder(ValueHandle *iv, ValueHandle lbHandle,
+                            ValueHandle ubHandle, ValueHandle stepHandle,
+                            ArrayRef<ValueHandle *> iter_args_handles,
+                            ValueRange iter_args_init_values) {
   mlir::edsc::LoopBuilder result;
-  auto forOp =
-      OperationHandle::createOp<loop::ForOp>(lbHandle, ubHandle, stepHandle);
+  auto forOp = OperationHandle::createOp<loop::ForOp>(
+      lbHandle, ubHandle, stepHandle, iter_args_init_values);
   *iv = ValueHandle(forOp.getInductionVar());
   auto *body = loop::getForInductionVarOwner(iv->getValue()).getBody();
+  for (size_t i = 0, e = iter_args_handles.size(); i < e; ++i) {
+    // Skipping the induction variable.
+    *(iter_args_handles[i]) = ValueHandle(body->getArgument(i + 1));
+  }
+  result.setOp(forOp);
   result.enter(body, /*prev=*/1);
   return result;
 }

mlir/test/EDSC/builder-api-test.cpp

 //===- builder-api-test.cpp - Tests for Declarative Builder APIs ----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 // RUN: mlir-edsc-builder-api-test | FileCheck %s -dump-input-on-failure
 
 #include "mlir/Dialect/Affine/EDSC/Intrinsics.h"
 #include "mlir/Dialect/Linalg/EDSC/Intrinsics.h"
-#include "mlir/Dialect/LoopOps/EDSC/Builders.h"
+#include "mlir/Dialect/LoopOps/EDSC/Intrinsics.h"
 #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
 #include "mlir/Dialect/Vector/EDSC/Intrinsics.h"
 #include "mlir/EDSC/Builders.h"
 #include "mlir/EDSC/Intrinsics.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/MLIRContext.h"
auto contractionBuilder = [](ArrayRef<BlockArgument> args) {
     (linalg_yield(vector_contraction_matmul(args[0], args[1], args[2])));
   };
   linalg_generic_matmul(A, B, C, contractionBuilder);
 
   f.print(llvm::outs());
   f.erase();
 }
 
+TEST_FUNC(builder_loop_for_yield) {
+  auto indexType = IndexType::get(&globalContext());
+  auto f32Type = FloatType::getF32(&globalContext());
+  auto f = makeFunction("builder_loop_for_yield", {},
+                        {indexType, indexType, indexType, indexType});
+
+  OpBuilder builder(f.getBody());
+  ScopedContext scope(builder, f.getLoc());
+  ValueHandle init0 = std_constant_float(llvm::APFloat(1.0f), f32Type);
+  ValueHandle init1 = std_constant_float(llvm::APFloat(2.0f), f32Type);
+  ValueHandle i(indexType), a(f.getArgument(0)), b(f.getArgument(1)),
+      c(f.getArgument(2)), d(f.getArgument(3));
+  ValueHandle arg0(f32Type);
+  ValueHandle arg1(f32Type);
+  using namespace edsc::op;
+  auto results =
+      LoopNestBuilder(&i, a - b, c + d, a, {&arg0, &arg1}, {init0, init1})([&] {
+        auto sum = arg0 + arg1;
+        loop_yield(ArrayRef<ValueHandle>{arg1, sum});
+      });
+  ValueHandle(results[0]) + ValueHandle(results[1]);
+
+  // clang-format off
+  // CHECK-LABEL: func @builder_loop_for_yield(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
+  // CHECK:     [[init0:%.*]] = constant
+  // CHECK:     [[init1:%.*]] = constant
+  // CHECK-DAG:    [[r0:%[0-9]+]] = affine.apply affine_map<()[s0, s1] -> (s0 - s1)>()[%{{.*}}, %{{.*}}]
+  // CHECK-DAG:    [[r1:%[0-9]+]] = affine.apply affine_map<()[s0, s1] -> (s0 + s1)>()[%{{.*}}, %{{.*}}]
+  // CHECK-NEXT: [[res:%[0-9]+]]:2 = loop.for %{{.*}} = [[r0]] to [[r1]] step {{.*}} iter_args([[arg0:%.*]] = [[init0]], [[arg1:%.*]] = [[init1]]) -> (f32, f32) {
+  // CHECK:     [[sum:%[0-9]+]] = addf [[arg0]], [[arg1]] : f32
+  // CHECK:     loop.yield [[arg1]], [[sum]] : f32, f32
+  // CHECK:     addf [[res]]#0, [[res]]#1 : f32
+
+  // clang-format on
+  f.print(llvm::outs());
+  f.erase();
+}
+
 int main() {
   RUN_TESTS();
   return 0;
 }