[mlir][Linalg] Update ReshapeOp::build to be more idiomatic

[nicolasvasilache]

Summary

This diff makes it easier to create a linalg.reshape op
and adds an EDSC builder api test to exercise the new builders.

[merge_guards_bot]

Unit tests: pass. 61744 tests passed, 0 failed and 780 were skipped.

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[ftynse]

The part that drops symbols needs documentation, including on the ::build to say that symbols will be dropped, and tests. I'm not confident that creating an AffineMap with 0 symbols from a list of expressions that might use symbols is correct.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

469

Could this just return pos?

[nicolasvasilache]

You're right, thanks made that symbol part more generic and fixed the issue.

[nicolasvasilache]

Address review comments.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt

[ftynse]

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

986

Can we have an AffineExpr with symbols here? Just to check that symbol removal machinery works.

[nicolasvasilache]

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

986

I made it into an assertion, there is no removal but just an assertion that it is symbolless.

[nicolasvasilache]

Update.

[merge_guards_bot]

Unit tests: pass. 61783 tests passed, 0 failed and 781 were skipped.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

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

 #include "mlir/EDSC/Builders.h"
 #include "mlir/EDSC/Intrinsics.h"
 
 namespace mlir {
 namespace edsc {
 namespace intrinsics {
 
 using linalg_fill = OperationBuilder<linalg::FillOp>;
-using linalg_reshape = OperationBuilder<linalg::ReshapeOp>;
+using linalg_reshape = ValueBuilder<linalg::ReshapeOp>;
 using linalg_yield = OperationBuilder<linalg::YieldOp>;
 
 } // namespace intrinsics
 } // namespace edsc
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_EDSC_INTRINSICS_H_

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

let description = [{
 
     ```mlir
        // Dimension expansion i -> (i', j') and (k) -> (k')
        %1 = linalg.reshape %0 [(i, j, k) -> (i, j), (i, j, k) -> (k)] :
          memref<?x?xf32, stride_spec> into memref<?x?x?xf32, stride_spec_2>
     ```
   }];
 
-  let builders = [OpBuilder<
-    "Builder *b, OperationState &result, Value view, "
-    "ArrayAttr reassociation, ArrayRef<NamedAttribute> attrs = {}">];
+  let builders = [
+    // Builder for a contracting reshape whose result type is computed from
+    // `view` and `reassociation`.
+    OpBuilder<"Builder *b, OperationState &result, Value view, "
+    "ArrayRef<ArrayRef<AffineExpr>> reassociation, "
+    "ArrayRef<NamedAttribute> attrs = {}">,
+    // Builder for a reshape whose result type is passed explicitly. This may be
+    // either a contracting or expanding reshape.
+    OpBuilder<"Builder *b, OperationState &result, Type resultType, Value view,"
+    "ArrayRef<ArrayRef<AffineExpr>> reassociation, "
+    "ArrayRef<NamedAttribute> attrs = {}">];
 
   let extraClassDeclaration = [{
     static StringRef getReassociationAttrName() { return "reassociation"; }
     MemRefType getViewType() { return view().getType().cast<MemRefType>(); }
   }];
 }
 
 def Linalg_SliceOp : Linalg_Op<"slice", [NoSideEffect]>,

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

 /// corresponding vector.
 /// TODO(rridle,ntv) this should be evolved into a generic
 /// `getRangeOfType<AffineMap>(ArrayAttr attrs)` that does not copy.
 static SmallVector<AffineMap, 4> getAffineMaps(ArrayAttr attrs) {
   return functional::map(
       [](Attribute a) { return a.cast<AffineMapAttr>().getValue(); }, attrs);
 }
 
-void mlir::linalg::ReshapeOp::build(Builder *b, OperationState &result,
-                                    Value view, ArrayAttr reassociation,
-                                    ArrayRef<NamedAttribute> attrs) {
-  auto maps = getAffineMaps(reassociation);
+template <typename AffineExprTy>
+unsigned getMaxPosOfType(ArrayRef<ArrayRef<AffineExpr>> exprArrays) {
+  unsigned pos = 0;
+  for (auto exprs : exprArrays) {
+    for (auto expr : exprs) {
+      expr.walk([&pos](AffineExpr e) {
+        if (auto d = e.dyn_cast<AffineExprTy>())
+          pos = std::max(pos, d.getPosition());
+      });
+    }
+  }
+  return pos;
+}
+
+static SmallVector<AffineMap, 4>
+getSymbolLessAffineMaps(ArrayRef<ArrayRef<AffineExpr>> reassociation) {
+  unsigned maxDim = getMaxPosOfType<AffineDimExpr>(reassociation);
+  unsigned maxSym = getMaxPosOfType<AffineSymbolExpr>(reassociation);
+  assert(maxSym == 0 && "Expected symbol-less expressions");
+  SmallVector<AffineMap, 4> maps;
+  maps.reserve(reassociation.size());
+  for (auto exprs : reassociation)
+    maps.push_back(AffineMap::get(maxDim + 1, 0, exprs));
+  return maps;
+}
+
+void mlir::linalg::ReshapeOp::build(
+    Builder *b, OperationState &result, Value view,
+    ArrayRef<ArrayRef<AffineExpr>> reassociation,
+  auto maps = getSymbolLessAffineMaps(reassociation);
   auto memRefType = view.getType().cast<MemRefType>();
   auto resultType = computeReshapeCollapsedType(memRefType, maps);
   build(b, result, resultType, view, attrs);
-  result.addAttribute(ReshapeOp::getReassociationAttrName(), reassociation);
+  result.addAttribute(ReshapeOp::getReassociationAttrName(),
+                      b->getAffineMapArrayAttr(maps));
+}
+
+void mlir::linalg::ReshapeOp::build(
+    Builder *b, OperationState &result, Type resultType, Value view,
+    ArrayRef<ArrayRef<AffineExpr>> reassociation,
+    ArrayRef<NamedAttribute> attrs) {
+  auto maps = getSymbolLessAffineMaps(reassociation);
+  build(b, result, resultType, view, attrs);
+  result.addAttribute(ReshapeOp::getReassociationAttrName(),
+                      b->getAffineMapArrayAttr(maps));
 }
 
 static void print(OpAsmPrinter &p, ReshapeOp op) {
   p << op.getOperationName() << " " << op.view() << " " << op.reassociation();
   p.printOptionalAttrDict(op.getAttrs(),
                           {ReshapeOp::getReassociationAttrName()});
   p << " : " << op.getViewType() << " into " << op.getResult().getType();
 }

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

 //===- builder-api-test.cpp - Tests for Declarative Builder APIs ----------===//
 //
 // Part of the MLIR 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
 
 #include "mlir/Dialect/AffineOps/AffineOps.h"
 #include "mlir/Dialect/Linalg/EDSC/Builders.h"
+#include "mlir/Dialect/Linalg/EDSC/Intrinsics.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/StandardOps/Ops.h"
 #include "mlir/EDSC/Builders.h"
 #include "mlir/EDSC/Helpers.h"
 #include "mlir/EDSC/Intrinsics.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/IntegerSet.h"
linalg_dilated_conv_nhwc(
       makeValueHandles(llvm::to_vector<3>(f.getArguments())),
       /*depth_multiplier=*/7,
       /*strides=*/{3, 4}, /*dilations=*/{5, 6});
 
   f.print(llvm::outs());
   f.erase();
 }
 
+// clang-format off
+// CHECK-LABEL: func @linalg_metadata_ops
+//       CHECK: linalg.reshape {{.*}} [(d0, d1, d2) -> (d0, d1), (d0, d1, d2) -> (d2)] : memref<4x8x16xf32> into memref<32x16xf32>
+//       CHECK: linalg.reshape {{.*}} [(d0, d1, d2) -> (d0, d1), (d0, d1, d2) -> (d2)] : memref<32x16xf32> into memref<4x8x16xf32>
+// clang-format on
+TEST_FUNC(linalg_metadata_ops) {
+  using namespace edsc;
+  using namespace edsc::intrinsics;
+
+  auto f32Type = FloatType::getF32(&globalContext());
+  auto memrefType = MemRefType::get({4, 8, 16}, f32Type, {}, 0);
+  auto f = makeFunction("linalg_metadata_ops", {}, {memrefType});
+
+  OpBuilder builder(f.getBody());
+  ScopedContext scope(builder, f.getLoc());
+  AffineExpr i, j, k;
+  bindDims(&globalContext(), i, j, k);
+  ValueHandle v(f.getArgument(0));
+  auto reshaped = linalg_reshape(v, ArrayRef<ArrayRef<AffineExpr>>{{i, j}, k});
+  linalg_reshape(memrefType, reshaped,
+                 ArrayRef<ArrayRef<AffineExpr>>{{i, j}, k});
+
+  f.print(llvm::outs());
+  f.erase();
+}
+
 int main() {
   RUN_TESTS();
   return 0;
 }