Cleanup ValueHandleArray

We just need a way to unpack ArrayRef<ValueHandle> to ArrayRef<Value*>.
No need to expose this to the user.

This reduces the cognitive overhead for the tutorial.

PiperOrigin-RevId: 240037425

6 files changed, 124 insertions, 86 deletions

diff --git a/mlir/bindings/python/pybind.cpp b/mlir/bindings/python/pybind.cpp
index 0184f8f5225..ffe1120fefe 100644
--- a/mlir/bindings/python/pybind.cpp
+++ b/mlir/bindings/python/pybind.cpp
@@ -956,7 +956,7 @@ PYBIND11_MODULE(pybind, m) {
       "ret",
       [](const std::vector<PythonValueHandle> &args) {
         std::vector<ValueHandle> values(args.begin(), args.end());
-        (intrinsics::ret(ValueHandleArray(values))); // vexing parse
+        (intrinsics::ret(ArrayRef<ValueHandle>{values})); // vexing parse
         return PythonValueHandle(nullptr);
       },
       py::arg("args") = std::vector<PythonValueHandle>());
diff --git a/mlir/include/mlir/EDSC/Helpers.h b/mlir/include/mlir/EDSC/Helpers.h
index cb717cc6d39..9b2efb4e010 100644
--- a/mlir/include/mlir/EDSC/Helpers.h
+++ b/mlir/include/mlir/EDSC/Helpers.h
@@ -36,69 +36,6 @@ template <typename Load, typename Store> class TemplatedIndexedValue;
 // load and stores.
 using IndexedValue = TemplatedIndexedValue<intrinsics::load, intrinsics::store>;
 
-/// An IndexHandle is a simple wrapper around a ValueHandle.
-/// IndexHandles are ubiquitous enough to justify a new type to allow simple
-/// declarations without boilerplate such as:
-///
-/// ```c++
-///    IndexHandle i, j, k;
-/// ```
-struct IndexHandle : public ValueHandle {
-  explicit IndexHandle()
-      : ValueHandle(ScopedContext::getBuilder()->getIndexType()) {}
-  explicit IndexHandle(index_t v) : ValueHandle(v) {}
-  explicit IndexHandle(Value *v) : ValueHandle(v) {
-    assert(v->getType() == ScopedContext::getBuilder()->getIndexType() &&
-           "Expected index type");
-  }
-  explicit IndexHandle(ValueHandle v) : ValueHandle(v) {
-    assert(v.getType() == ScopedContext::getBuilder()->getIndexType() &&
-           "Expected index type");
-  }
-  IndexHandle &operator=(const ValueHandle &v) {
-    assert(v.getType() == ScopedContext::getBuilder()->getIndexType() &&
-           "Expected index type");
-    /// Creating a new IndexHandle(v) and then std::swap rightly complains the
-    /// binding has already occurred and that we should use another name.
-    this->t = v.getType();
-    this->v = v.getValue();
-    return *this;
-  }
-  static SmallVector<IndexHandle, 8> makeIndexHandles(unsigned rank) {
-    return SmallVector<IndexHandle, 8>(rank);
-  }
-  static SmallVector<ValueHandle *, 8>
-  makePIndexHandles(SmallVectorImpl<IndexHandle> &ivs) {
-    SmallVector<ValueHandle *, 8> pivs;
-    for (auto &iv : ivs) {
-      pivs.push_back(&iv);
-    }
-    return pivs;
-  }
-};
-
-/// Helper structure to be used with ValueBuilder / InstructionBuilder.
-/// It serves the purpose of removing boilerplate specialization for the sole
-/// purpose of implicitly converting ArrayRef<ValueHandle> -> ArrayRef<Value*>.
-class ValueHandleArray {
-public:
-  ValueHandleArray(ArrayRef<ValueHandle> vals) {
-    values.append(vals.begin(), vals.end());
-  }
-  ValueHandleArray(ArrayRef<IndexHandle> vals) {
-    values.append(vals.begin(), vals.end());
-  }
-  ValueHandleArray(ArrayRef<index_t> vals) {
-    llvm::SmallVector<IndexHandle, 8> tmp(vals.begin(), vals.end());
-    values.append(tmp.begin(), tmp.end());
-  }
-  operator ArrayRef<Value *>() { return values; }
-
-private:
-  ValueHandleArray() = default;
-  llvm::SmallVector<Value *, 8> values;
-};
-
 // Base class for MemRefView and VectorView.
 class View {
 public:
@@ -207,16 +144,16 @@ template <typename Load, typename Store> struct TemplatedIndexedValue {
   // NOLINTNEXTLINE: unconventional-assign-operator
   InstructionHandle operator=(const TemplatedIndexedValue &rhs) {
     ValueHandle rrhs(rhs);
-    return Store(rrhs, getBase(), ValueHandleArray(indices));
+    return Store(rrhs, getBase(), {indices.begin(), indices.end()});
   }
   // NOLINTNEXTLINE: unconventional-assign-operator
   InstructionHandle operator=(ValueHandle rhs) {
-    return Store(rhs, getBase(), ValueHandleArray(indices));
+    return Store(rhs, getBase(), {indices.begin(), indices.end()});
   }
 
   /// Emits a `load` when converting to a ValueHandle.
   operator ValueHandle() const {
-    return Load(getBase(), ValueHandleArray(indices));
+    return Load(getBase(), {indices.begin(), indices.end()});
   }
 
   ValueHandle getBase() const { return base; }
@@ -297,25 +234,25 @@ template <typename Load, typename Store>
 InstructionHandle
 TemplatedIndexedValue<Load, Store>::operator+=(ValueHandle e) {
   using op::operator+;
-  return Store(*this + e, getBase(), ValueHandleArray(indices));
+  return Store(*this + e, getBase(), {indices.begin(), indices.end()});
 }
 template <typename Load, typename Store>
 InstructionHandle
 TemplatedIndexedValue<Load, Store>::operator-=(ValueHandle e) {
   using op::operator-;
-  return Store(*this - e, getBase(), ValueHandleArray(indices));
+  return Store(*this - e, getBase(), {indices.begin(), indices.end()});
 }
 template <typename Load, typename Store>
 InstructionHandle
 TemplatedIndexedValue<Load, Store>::operator*=(ValueHandle e) {
   using op::operator*;
-  return Store(*this * e, getBase(), ValueHandleArray(indices));
+  return Store(*this * e, getBase(), {indices.begin(), indices.end()});
 }
 template <typename Load, typename Store>
 InstructionHandle
 TemplatedIndexedValue<Load, Store>::operator/=(ValueHandle e) {
   using op::operator/;
-  return Store(*this / e, getBase(), ValueHandleArray(indices));
+  return Store(*this / e, getBase(), {indices.begin(), indices.end()});
 }
 
 } // namespace edsc
diff --git a/mlir/include/mlir/EDSC/Intrinsics.h b/mlir/include/mlir/EDSC/Intrinsics.h
index 32170c47e20..c46f08a68e5 100644
--- a/mlir/include/mlir/EDSC/Intrinsics.h
+++ b/mlir/include/mlir/EDSC/Intrinsics.h
@@ -33,14 +33,83 @@ class Type;
 
 namespace edsc {
 
-class BlockHandle;
-class InstructionHandle;
-class ValueHandle;
+/// An IndexHandle is a simple wrapper around a ValueHandle.
+/// IndexHandles are ubiquitous enough to justify a new type to allow simple
+/// declarations without boilerplate such as:
+///
+/// ```c++
+///    IndexHandle i, j, k;
+/// ```
+struct IndexHandle : public ValueHandle {
+  explicit IndexHandle()
+      : ValueHandle(ScopedContext::getBuilder()->getIndexType()) {}
+  explicit IndexHandle(index_t v) : ValueHandle(v) {}
+  explicit IndexHandle(Value *v) : ValueHandle(v) {
+    assert(v->getType() == ScopedContext::getBuilder()->getIndexType() &&
+           "Expected index type");
+  }
+  explicit IndexHandle(ValueHandle v) : ValueHandle(v) {
+    assert(v.getType() == ScopedContext::getBuilder()->getIndexType() &&
+           "Expected index type");
+  }
+  IndexHandle &operator=(const ValueHandle &v) {
+    assert(v.getType() == ScopedContext::getBuilder()->getIndexType() &&
+           "Expected index type");
+    /// Creating a new IndexHandle(v) and then std::swap rightly complains the
+    /// binding has already occurred and that we should use another name.
+    this->t = v.getType();
+    this->v = v.getValue();
+    return *this;
+  }
+  static SmallVector<IndexHandle, 8> makeIndexHandles(unsigned rank) {
+    return SmallVector<IndexHandle, 8>(rank);
+  }
+  static SmallVector<ValueHandle *, 8>
+  makeIndexHandlePointers(SmallVectorImpl<IndexHandle> &ivs) {
+    SmallVector<ValueHandle *, 8> pivs;
+    pivs.reserve(ivs.size());
+    for (auto &iv : ivs) {
+      pivs.push_back(&iv);
+    }
+    return pivs;
+  }
+};
 
 /// Provides a set of first class intrinsics.
 /// In the future, most of intrinsics reated to Instruction that don't contain
 /// other instructions should be Tablegen'd.
 namespace intrinsics {
+namespace detail {
+/// Helper structure to be used with ValueBuilder / InstructionBuilder.
+/// It serves the purpose of removing boilerplate specialization for the sole
+/// purpose of implicitly converting ArrayRef<ValueHandle> -> ArrayRef<Value*>.
+class ValueHandleArray {
+public:
+  ValueHandleArray(ArrayRef<ValueHandle> vals) {
+    values.append(vals.begin(), vals.end());
+  }
+  ValueHandleArray(ArrayRef<IndexHandle> vals) {
+    values.append(vals.begin(), vals.end());
+  }
+  ValueHandleArray(ArrayRef<index_t> vals) {
+    llvm::SmallVector<IndexHandle, 8> tmp(vals.begin(), vals.end());
+    values.append(tmp.begin(), tmp.end());
+  }
+  operator ArrayRef<Value *>() { return values; }
+
+private:
+  ValueHandleArray() = default;
+  llvm::SmallVector<Value *, 8> values;
+};
+
+template <typename T> inline T unpack(T value) { return value; }
+
+inline detail::ValueHandleArray unpack(ArrayRef<ValueHandle> values) {
+  return detail::ValueHandleArray(values);
+}
+
+} // namespace detail
+
 /// Helper variadic abstraction to allow extending to any MLIR op without
 /// boilerplate or Tablegen.
 /// Arguably a builder is not a ValueHandle but in practice it is only used as
@@ -51,7 +120,22 @@ namespace intrinsics {
 template <typename Op> struct ValueBuilder : public ValueHandle {
   template <typename... Args>
   ValueBuilder(Args... args)
-      : ValueHandle(ValueHandle::create<Op>(std::forward<Args>(args)...)) {}
+      : ValueHandle(ValueHandle::create<Op>(detail::unpack(args)...)) {}
+  ValueBuilder(ArrayRef<ValueHandle> vs)
+      : ValueBuilder(ValueBuilder::create<Op>(detail::unpack(vs))) {}
+  template <typename... Args>
+  ValueBuilder(ArrayRef<ValueHandle> vs, Args... args)
+      : ValueHandle(ValueHandle::create<Op>(detail::unpack(vs),
+                                            detail::unpack(args)...)) {}
+  template <typename T, typename... Args>
+  ValueBuilder(T t, ArrayRef<ValueHandle> vs, Args... args)
+      : ValueHandle(ValueHandle::create<Op>(
+            detail::unpack(t), detail::unpack(vs), detail::unpack(args)...)) {}
+  template <typename T1, typename T2, typename... Args>
+  ValueBuilder(T1 t1, T2 t2, ArrayRef<ValueHandle> vs, Args... args)
+      : ValueHandle(ValueHandle::create<Op>(
+            detail::unpack(t1), detail::unpack(t2), detail::unpack(vs),
+            detail::unpack(args)...)) {}
   ValueBuilder() : ValueHandle(ValueHandle::create<Op>()) {}
 };
 
@@ -59,7 +143,22 @@ template <typename Op> struct InstructionBuilder : public InstructionHandle {
   template <typename... Args>
   InstructionBuilder(Args... args)
       : InstructionHandle(
-            InstructionHandle::create<Op>(std::forward<Args>(args)...)) {}
+            InstructionHandle::create<Op>(detail::unpack(args)...)) {}
+  InstructionBuilder(ArrayRef<ValueHandle> vs)
+      : InstructionHandle(InstructionHandle::create<Op>(detail::unpack(vs))) {}
+  template <typename... Args>
+  InstructionBuilder(ArrayRef<ValueHandle> vs, Args... args)
+      : InstructionHandle(InstructionHandle::create<Op>(
+            detail::unpack(vs), detail::unpack(args)...)) {}
+  template <typename T, typename... Args>
+  InstructionBuilder(T t, ArrayRef<ValueHandle> vs, Args... args)
+      : InstructionHandle(InstructionHandle::create<Op>(
+            detail::unpack(t), detail::unpack(vs), detail::unpack(args)...)) {}
+  template <typename T1, typename T2, typename... Args>
+  InstructionBuilder(T1 t1, T2 t2, ArrayRef<ValueHandle> vs, Args... args)
+      : InstructionHandle(InstructionHandle::create<Op>(
+            detail::unpack(t1), detail::unpack(t2), detail::unpack(vs),
+            detail::unpack(args)...)) {}
   InstructionBuilder() : InstructionHandle(InstructionHandle::create<Op>()) {}
 };
 
diff --git a/mlir/lib/Transforms/LowerVectorTransfers.cpp b/mlir/lib/Transforms/LowerVectorTransfers.cpp
index 321c932b541..f0c204f9840 100644
--- a/mlir/lib/Transforms/LowerVectorTransfers.cpp
+++ b/mlir/lib/Transforms/LowerVectorTransfers.cpp
@@ -274,7 +274,8 @@ template <> void VectorTransferRewriter<VectorTransferReadOp>::rewrite() {
   VectorView vectorView(transfer->getVector());
   SmallVector<IndexHandle, 8> ivs =
       IndexHandle::makeIndexHandles(vectorView.rank());
-  SmallVector<ValueHandle *, 8> pivs = IndexHandle::makePIndexHandles(ivs);
+  SmallVector<ValueHandle *, 8> pivs =
+      IndexHandle::makeIndexHandlePointers(ivs);
   coalesceCopy(transfer, &pivs, &vectorView);
 
   auto lbs = vectorView.getLbs();
@@ -289,7 +290,7 @@ template <> void VectorTransferRewriter<VectorTransferReadOp>::rewrite() {
       // Computes clippedScalarAccessExprs in the loop nest scope (ivs exist).
       local(ivs) = remote(clip(transfer, view, ivs)),
   });
-  ValueHandle vectorValue = load(vec, ValueHandleArray{index_t(0)});
+  ValueHandle vectorValue = load(vec, {constant_index(0)});
   (dealloc(tmp)); // vexing parse
 
   // 3. Propagate.
@@ -329,7 +330,8 @@ template <> void VectorTransferRewriter<VectorTransferWriteOp>::rewrite() {
   VectorView vectorView(transfer->getVector());
   SmallVector<IndexHandle, 8> ivs =
       IndexHandle::makeIndexHandles(vectorView.rank());
-  SmallVector<ValueHandle *, 8> pivs = IndexHandle::makePIndexHandles(ivs);
+  SmallVector<ValueHandle *, 8> pivs =
+      IndexHandle::makeIndexHandlePointers(ivs);
   coalesceCopy(transfer, &pivs, &vectorView);
 
   auto lbs = vectorView.getLbs();
@@ -340,7 +342,7 @@ template <> void VectorTransferRewriter<VectorTransferWriteOp>::rewrite() {
   ValueHandle tmp = alloc(tmpMemRefType());
   IndexedValue local(tmp);
   ValueHandle vec = vector_type_cast(tmp, vectorMemRefType());
-  store(vectorValue, vec, ValueHandleArray{index_t(0)});
+  store(vectorValue, vec, {constant_index(0)});
   LoopNestBuilder(pivs, lbs, ubs, steps)({
       // Computes clippedScalarAccessExprs in the loop nest scope (ivs exist).
       remote(clip(transfer, view, ivs)) = local(ivs),
diff --git a/mlir/test/EDSC/builder-api-test.cpp b/mlir/test/EDSC/builder-api-test.cpp
index 8d2b960c5ef..ec6d12a0876 100644
--- a/mlir/test/EDSC/builder-api-test.cpp
+++ b/mlir/test/EDSC/builder-api-test.cpp
@@ -136,7 +136,7 @@ TEST_FUNC(builder_max_min_for) {
   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)({});
-  ret({});
+  ret();
 
   // clang-format off
   // CHECK-LABEL: func @builder_max_min_for(%arg0: index, %arg1: index, %arg2: index, %arg3: index) {
@@ -250,10 +250,10 @@ TEST_FUNC(builder_cond_branch) {
 
   BlockHandle b1, b2, functionBlock(&f->front());
   BlockBuilder(&b1, {&arg1})({
-      ret({}),
+      ret(),
   });
   BlockBuilder(&b2, {&arg2, &arg3})({
-      ret({}),
+      ret(),
   });
   // Get back to entry block and add a conditional branch
   BlockBuilder(functionBlock, Append())({
@@ -292,10 +292,10 @@ TEST_FUNC(builder_cond_branch_eager) {
   BlockHandle b1, b2;
   cond_br(funcArg, &b1, {&arg1}, {c32}, &b2, {&arg2, &arg3}, {c64, c42});
   BlockBuilder(b1, Append())({
-      ret({}),
+      ret(),
   });
   BlockBuilder(b2, Append())({
-      ret({}),
+      ret(),
   });
 
   // CHECK-LABEL: @builder_cond_branch_eager
diff --git a/mlir/test/mlir-tblgen/reference-impl.td b/mlir/test/mlir-tblgen/reference-impl.td
index fad12020793..be77c6344e5 100644
--- a/mlir/test/mlir-tblgen/reference-impl.td
+++ b/mlir/test/mlir-tblgen/reference-impl.td
@@ -11,7 +11,7 @@ def X_AddOp : Op<"x.add">,
     // TODO: extract referenceImplementation to Op.
   code referenceImplementation = [{
     auto ivs = IndexHandle::makeIndexHandles(view_A.rank());
-    auto pivs = IndexHandle::makePIndexHandles(ivs);
+    auto pivs = IndexHandle::makeIndexHandlePointers(ivs);
     IndexedValue A(arg_A), B(arg_B), C(arg_C);
     LoopNestBuilder(pivs, view_A.getLbs(), view_A.getUbs(), view_A.getSteps())({
       C(ivs) = A(ivs) + B(ivs)