[mlir][EDSC] NFC - Move StructuredIndexed and IteratorType out of Linalg

[nicolasvasilache]

Summary

This NFC revision will allow those classes to be reused to allow
building structured vector operations.

[mehdi_amini]

Build failure:

/mnt/disks/ssd0/agent/workspace/amd64_debian_testing_clang8/mlir/lib/Dialect/Linalg/EDSC/Builders.cpp:244:21: error: no member named 'IterType' in namespace 'mlir::edsc'
  SmallVector<edsc::IterType, 4> iterTypes(O.getExprs().size(),
              ~~~~~~^
/mnt/disks/ssd0/agent/workspace/amd64_debian_testing_clang8/mlir/lib/Dialect/Linalg/EDSC/Builders.cpp:252:32: error: use of undeclared identifier 'iterTypes'
    return makeGenericLinalgOp(iterTypes, {I}, {O}, fun);
                               ^
/mnt/disks/ssd0/agent/workspace/amd64_debian_testing_clang8/mlir/lib/Dialect/Linalg/EDSC/Builders.cpp:259:30: error: use of undeclared identifier 'iterTypes'
  return makeGenericLinalgOp(iterTypes, {I}, {O}, fun);

[nicolasvasilache]

Fix bad rebase.

[nicolasvasilache]

Thanks Mehdi!

[nicolasvasilache]

Ugh .. misread the green checkmark from Harbormaster as the second approval from @mehdi_amini post trivial rebase.
As a consequence this landed in "Needs Review", sorry about that ...

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

private:
   using BuilderType =
       typename std::conditional_t<std::is_same<LoopTy, loop::ParallelOp>::value,
                                   ParallelLoopNestBuilder,
                                   LoopOrAffineLoopBuilder>;
 
   std::unique_ptr<BuilderType> builder;
 };
 
-enum class IterType { Parallel, Reduction };
-
-inline StringRef toString(IterType t) {
-  switch (t) {
-  case IterType::Parallel:
-    return getParallelIteratorTypeName();
-  case IterType::Reduction:
-    return getReductionIteratorTypeName();
-  }
-  llvm_unreachable("Unsupported IterType");
-}
-
-/// A StructuredIndexed represents an indexable quantity that is either:
-/// 1. a captured value, which is suitable for buffer and tensor operands, or;
-/// 2. a captured type, which is suitable for tensor return values.
-///
-/// A StructuredIndexed itself is indexed and passed to `makeGenericLinalgOp`.
-/// It enable an idiomatic syntax for index expressions such as:
-///
-/// ```
-///      StructuredIndexed A(buffer_or_tensor_value), B(buffer_or_tensor_value),
-///        C(buffer_value_or_tensor_type);
-///      makeGenericLinalgOp({A({m, n}), B({k, n})}, {C({m, n})}, ... );
-/// ```
-struct StructuredIndexed : public ValueHandle {
-  StructuredIndexed(Type type) : ValueHandle(type) {}
-  StructuredIndexed(Value value) : ValueHandle(value) {}
-  StructuredIndexed(ValueHandle valueHandle) : ValueHandle(valueHandle) {}
-  StructuredIndexed operator()(ArrayRef<AffineExpr> indexings) {
-    return StructuredIndexed(*this, indexings);
-  }
-
-  ArrayRef<AffineExpr> getExprs() { return exprs; }
-
-private:
-  StructuredIndexed(Type t, ArrayRef<AffineExpr> indexings)
-      : ValueHandle(t), exprs(indexings.begin(), indexings.end()) {
-    assert(t.isa<RankedTensorType>() && "RankedTensor expected");
-  }
-  StructuredIndexed(Value v, ArrayRef<AffineExpr> indexings)
-      : ValueHandle(v), exprs(indexings.begin(), indexings.end()) {
-    assert((v.getType().isa<MemRefType>() ||
-            v.getType().isa<RankedTensorType>()) &&
-           "MemRef or RankedTensor expected");
-  }
-  StructuredIndexed(ValueHandle vh, ArrayRef<AffineExpr> indexings)
-      : ValueHandle(vh), exprs(indexings.begin(), indexings.end()) {}
-
-  SmallVector<AffineExpr, 4> exprs;
-};
-
 inline void defaultRegionBuilder(ArrayRef<BlockArgument> args) {}
 
 /// Build a `linalg.generic` op with the specified `inputs`, `outputs` and
 /// `region`.
 ///
 /// `otherValues` and `otherAttributes` may be passed and will be appended as
 /// operands and attributes respectively.
 ///
 /// Prerequisites:
 /// =============
 ///
 /// 1. `inputs` may contain StructuredIndexed that capture either buffer or
 /// tensor values.
 /// 2. `outputs` may contain StructuredIndexed that capture either buffer values
 /// or tensor types. If both buffer values and tensor types are present, then
 /// all buffer values must appear before any tensor type. Without this
 /// restriction output tensor results would need to be reordered, which would
 /// result in surprising behavior when combined with region definition.
 Operation *makeGenericLinalgOp(
-    ArrayRef<IterType> iteratorTypes, ArrayRef<StructuredIndexed> inputs,
+    ArrayRef<IteratorType> iteratorTypes, ArrayRef<StructuredIndexed> inputs,
     ArrayRef<StructuredIndexed> outputs,
     function_ref<void(ArrayRef<BlockArgument>)> regionBuilder =
         defaultRegionBuilder,
     ArrayRef<Value> otherValues = {}, ArrayRef<Attribute> otherAttributes = {});
 
 namespace ops {
 using edsc::StructuredIndexed;
 using edsc::ValueHandle;

mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h

 
 inline unsigned getNumIterators(ArrayAttr iteratorTypes) {
   unsigned res = 0;
   for (auto n : getAllIteratorTypeNames())
     res += getNumIterators(n, iteratorTypes);
   return res;
 }
 
+/// Typed representation for loop type strings.
+enum class IteratorType { Parallel, Reduction };
+
+inline StringRef toString(IteratorType t) {
+  switch (t) {
+  case IteratorType::Parallel:
+    return getParallelIteratorTypeName();
+  case IteratorType::Reduction:
+    return getReductionIteratorTypeName();
+  }
+  llvm_unreachable("Unsupported IteratorType");
+}
+
 } // end namespace mlir
 
 #endif // MLIR_UTILS_STRUCTUREDOPSUTILS_H

mlir/include/mlir/EDSC/Builders.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_EDSC_BUILDERS_H_
 #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/IR/AffineExpr.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/Transforms/FoldUtils.h"
 
 namespace mlir {
 
 namespace edsc {
 
 struct index_type {
public:
   operator bool() { return block != nullptr; }
   operator mlir::Block *() { return block; }
   mlir::Block *getBlock() { return block; }
 
 private:
   mlir::Block *block;
 };
 
+/// A StructuredIndexed represents an indexable quantity that is either:
+/// 1. a captured value, which is suitable for buffer and tensor operands, or;
+/// 2. a captured type, which is suitable for tensor return values.
+///
+/// A StructuredIndexed itself is indexed and passed to `makeGenericLinalgOp`.
+/// It enable an idiomatic syntax for index expressions such as:
+///
+/// ```
+///      StructuredIndexed A(buffer_or_tensor_value), B(buffer_or_tensor_value),
+///        C(buffer_value_or_tensor_type);
+///      makeGenericLinalgOp({A({m, n}), B({k, n})}, {C({m, n})}, ... );
+/// ```
+struct StructuredIndexed : public ValueHandle {
+  StructuredIndexed(Type type) : ValueHandle(type) {}
+  StructuredIndexed(Value value) : ValueHandle(value) {}
+  StructuredIndexed(ValueHandle valueHandle) : ValueHandle(valueHandle) {}
+  StructuredIndexed operator()(ArrayRef<AffineExpr> indexings) {
+    return this->hasValue() ? StructuredIndexed(this->getValue(), indexings)
+                            : StructuredIndexed(this->getType(), indexings);
+  }
+
+  StructuredIndexed(Type t, ArrayRef<AffineExpr> indexings)
+      : ValueHandle(t), exprs(indexings.begin(), indexings.end()) {
+    assert(t.isa<RankedTensorType>() && "RankedTensor expected");
+  }
+  StructuredIndexed(Value v, ArrayRef<AffineExpr> indexings)
+      : ValueHandle(v), exprs(indexings.begin(), indexings.end()) {
+    assert((v.getType().isa<MemRefType>() ||
+            v.getType().isa<RankedTensorType>()) &&
+           "MemRef or RankedTensor expected");
+  }
+  StructuredIndexed(ValueHandle vh, ArrayRef<AffineExpr> indexings)
+      : ValueHandle(vh), exprs(indexings.begin(), indexings.end()) {}
+
+  ArrayRef<AffineExpr> getExprs() { return exprs; }
+
+private:
+  SmallVector<AffineExpr, 4> exprs;
+};
+
 template <typename Op, typename... Args>
 OperationHandle OperationHandle::create(Args... args) {
   return OperationHandle(ScopedContext::getBuilder()
                              .create<Op>(ScopedContext::getLocation(), args...)
                              .getOperation());
 }
 
 template <typename Op, typename... Args>

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

 //===- Builders.cpp - MLIR Declarative Linalg Builders --------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/EDSC/Builders.h"
 #include "mlir/Dialect/Linalg/EDSC/Intrinsics.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/EDSC/Builders.h"
 #include "mlir/EDSC/Intrinsics.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/Support/Functional.h"
 
 using namespace mlir;
 using namespace mlir::edsc;
for (auto expr : sidx.getExprs()) {
         if (auto d = e.dyn_cast<AffineDimExpr>())
           pos = std::max(pos, d.getPosition());
       });
     }
   }
 }
 
 Operation *mlir::edsc::makeGenericLinalgOp(
-    ArrayRef<IterType> iteratorTypes, ArrayRef<StructuredIndexed> inputs,
+    ArrayRef<IteratorType> iteratorTypes, ArrayRef<StructuredIndexed> inputs,
     ArrayRef<StructuredIndexed> outputs,
     function_ref<void(ArrayRef<BlockArgument>)> regionBuilder,
     ArrayRef<Value> otherValues, ArrayRef<Attribute> otherAttributes) {
   for (unsigned i = 0, e = outputs.size(); i + 1 < e; ++i)
     assert(!(outputs[i].getType().isa<RankedTensorType>() &&
              outputs[i + 1].getType().isa<MemRefType>()) &&
            "output tensors must be passed after output buffers");
   auto &builder = edsc::ScopedContext::getBuilder();
void mlir::edsc::ops::macRegionBuilder(ArrayRef<BlockArgument> args) {
   assert(args.size() == 3 && "expected 3 block arguments");
   ValueHandle a(args[0]), b(args[1]), c(args[2]);
   linalg_yield((c + a * b).getValue());
 }
 
 Operation *mlir::edsc::ops::linalg_pointwise(UnaryPointwiseOpBuilder unaryOp,
                                              StructuredIndexed I,
                                              StructuredIndexed O) {
-  SmallVector<edsc::IterType, 4> iterTypes(O.getExprs().size(),
-                                           edsc::IterType::Parallel);
+  SmallVector<IteratorType, 4> iterTypes(O.getExprs().size(),
+                                         IteratorType::Parallel);
   if (O.getType().isa<RankedTensorType>()) {
     auto fun = [&unaryOp](ArrayRef<BlockArgument> args) {
       assert(args.size() == 1 && "expected 1 block arguments");
       ValueHandle a(args[0]);
       linalg_yield(unaryOp(a));
     };
     return makeGenericLinalgOp(iterTypes, {I}, {O}, fun);
   }
Operation *mlir::edsc::ops::linalg_pointwise_tanh(StructuredIndexed I,
   return linalg_pointwise(unOp, I, O);
 }
 
 /// Binary pointwise operation (with broadcast) entry point.
 Operation *mlir::edsc::ops::linalg_pointwise(BinaryPointwiseOpBuilder binaryOp,
                                              StructuredIndexed I1,
                                              StructuredIndexed I2,
                                              StructuredIndexed O) {
-  SmallVector<edsc::IterType, 4> iterTypes(O.getExprs().size(),
-                                           edsc::IterType::Parallel);
+  SmallVector<IteratorType, 4> iterTypes(O.getExprs().size(),
+                                         IteratorType::Parallel);
   if (O.getType().isa<RankedTensorType>()) {
     auto fun = [&binaryOp](ArrayRef<BlockArgument> args) {
       assert(args.size() == 2 && "expected 2 block arguments");
       ValueHandle a(args[0]), b(args[1]);
       linalg_yield(binaryOp(a, b));
     };
     return makeGenericLinalgOp(iterTypes, {I1, I2}, {O}, fun);
   }

 
 Operation *mlir::edsc::ops::linalg_matmul(ValueHandle vA, ValueHandle vB,
                                           ValueHandle vC) {
   // clang-format off
   AffineExpr m, n, k;
   bindDims(ScopedContext::getContext(), m, n, k);
   StructuredIndexed A(vA), B(vB), C(vC);
   return makeGenericLinalgOp(
-    {IterType::Parallel, IterType::Parallel, IterType::Reduction},
+    {IteratorType::Parallel, IteratorType::Parallel, IteratorType::Reduction},
     {A({m, k}), B({k, n})},
     {C({m, n})},
     macRegionBuilder);
   // clang-format on
 }
 
 Operation *mlir::edsc::ops::linalg_matmul(ValueHandle vA, ValueHandle vB,
                                           RankedTensorType tC) {
   // clang-format off
   AffineExpr m, n, k;
   bindDims(ScopedContext::getContext(), m, n, k);
   StructuredIndexed A(vA), B(vB), C(tC);
   return makeGenericLinalgOp(
-    {IterType::Parallel, IterType::Parallel, IterType::Reduction},
+    {IteratorType::Parallel, IteratorType::Parallel, IteratorType::Reduction},
     {A({m, k}), B({k, n})},
     {C({m, n})},
     mulRegionBuilder);
   // clang-format on
 }
 
 Operation *mlir::edsc::ops::linalg_matmul(ValueHandle vA, ValueHandle vB,
                                           ValueHandle vC, RankedTensorType tD) {
   // clang-format off
   AffineExpr m, n, k;
   bindDims(ScopedContext::getContext(), m, n, k);
   StructuredIndexed A(vA), B(vB), C(vC), D(tD);
   return makeGenericLinalgOp(
-    {IterType::Parallel, IterType::Parallel, IterType::Reduction},
+    {IteratorType::Parallel, IteratorType::Parallel, IteratorType::Reduction},
     {A({m, k}), B({k, n}), C({m, n})},
     {D({m, n})},
     macRegionBuilder);
   // clang-format on
 }
 
 Operation *mlir::edsc::ops::linalg_conv_nhwc(ValueHandle vI, ValueHandle vW,
                                              ValueHandle vO,
                                              ArrayRef<int> strides,
                                              ArrayRef<int> dilations) {
   MLIRContext *ctx = ScopedContext::getContext();
   // TODO(ntv) some template magic to make everything rank-polymorphic.
   assert((dilations.empty() || dilations.size() == 2) && "only 2-D conv atm");
   assert((strides.empty() || strides.size() == 2) && "only 2-D conv atm");
 
   // Some short names.
-  auto par = IterType::Parallel;
-  auto red = IterType::Reduction;
+  auto par = IteratorType::Parallel;
+  auto red = IteratorType::Reduction;
   auto s = strides;
   auto d = dilations;
 
   AffineExpr b, f, h, w, kh, kw, c;
   bindDims(ctx, b, f, h, w, kh, kw, c);
   unsigned numDims = c.cast<AffineDimExpr>().getPosition() + 1;
   StructuredIndexed I(vI), W(vW), O(vO);
   // clang-format off
Operation *mlir::edsc::ops::linalg_dilated_conv_nhwc(
     ValueHandle vI, ValueHandle vW, ValueHandle vO, int depth_multiplier,
     ArrayRef<int> strides, ArrayRef<int> dilations) {
   MLIRContext *ctx = ScopedContext::getContext();
   // TODO(ntv) some template magic to make everything rank-polymorphic.
   assert((dilations.empty() || dilations.size() == 2) && "only 2-D conv atm");
   assert((strides.empty() || strides.size() == 2) && "only 2-D conv atm");
 
   // Some short names.
-  auto par = IterType::Parallel;
-  auto red = IterType::Reduction;
+  auto par = IteratorType::Parallel;
+  auto red = IteratorType::Reduction;
   auto s = strides;
   auto d = dilations;
 
   // clang-format off
   AffineExpr b, dm, c, h, w, kh, kw;
   bindDims(ctx, b, dm, c, h, w, kh, kw);
   unsigned numDims = kw.cast<AffineDimExpr>().getPosition() + 1;
   StructuredIndexed I(vI), W(vW), O(vO);