Automated rollback of commit 575405f4d6762830c1c4520569de4e4ed3c8eed5

PiperOrigin-RevId: 275461067

2 files changed, 57 insertions, 17 deletions

diff --git a/mlir/lib/EDSC/Builders.cpp b/mlir/lib/EDSC/Builders.cpp
index ba14a147cf8..f1dceec8f11 100644
--- a/mlir/lib/EDSC/Builders.cpp
+++ b/mlir/lib/EDSC/Builders.cpp
@@ -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());
@@ -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:
@@ -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(
@@ -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();
@@ -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) {
diff --git a/mlir/lib/EDSC/CMakeLists.txt b/mlir/lib/EDSC/CMakeLists.txt
index 5f3dd9f1ee1..967d6add293 100644
--- a/mlir/lib/EDSC/CMakeLists.txt
+++ b/mlir/lib/EDSC/CMakeLists.txt
@@ -17,6 +17,7 @@ add_dependencies(MLIREDSC MLIRReferenceImplementationTestGen)
 target_link_libraries(MLIREDSC
   PUBLIC
     MLIRAffineOps
+    MLIRLoopOps
     MLIRStandardOps
     MLIRTransformUtils
     MLIRVectorOps