[mlir][Linalg] Fix Linalg EDSC builders

[nicolasvasilache]

Summary

This diff fixes the fact that the method mlir::edsc::makeGenericLinalgOp
incorrectly adds 2 blocks to Linalg ops.

Tests are updated accordingly.

[merge_guards_bot]

Unit tests: pass. 61804 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

[asaadaldien]

LGTM! an inline question trying to understand the Builder blocks scoping here.

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

97

OpBuilder bb(region);

103

Does the builder create a new block instead of mutating region.begin that we have to clean up ?

[nicolasvasilache]

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

103

Yes, currently it does.
I can improve this by extending ScopedContext slightly I think.
This is because the mechanism predates regions and has not yet been updated.

[nicolasvasilache]

Address review comments.

[merge_guards_bot]

Unit tests: pass. 61912 tests passed, 0 failed and 783 were skipped.

clang-tidy: unknown.

clang-format: pass.

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

[asaadaldien]

mlir/lib/EDSC/Builders.cpp

307 ↗

(On Diff #238493)

typo s/BockBuilder/BlockBuilder

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

Operation *mlir::edsc::makeGenericLinalgOp(
   using namespace edsc;
   SmallVector<Type, 4> blockTypes;
   blockTypes.reserve(values.size());
   for (auto it : llvm::enumerate(values))
     blockTypes.push_back((it.index() < nViews)
                              ? getElementTypeOrSelf(it.value())
                              : it.value().getType());
 
-  assert(op->getRegions().front().empty());
-  op->getRegions().front().push_front(new Block);
+  assert(op->getNumRegions() == 1);
+  assert(op->getRegion(0).empty());
+  auto &region = op->getRegion(0);
+  region.push_front(new Block);
   OpBuilder bb(op->getRegions().front());
   ScopedContext scope(bb, op->getLoc());
   BlockHandle b;
   auto handles = makeValueHandles(blockTypes);
   BlockBuilder(&b, makeHandlePointers(MutableArrayRef<ValueHandle>(handles)))(
       [&] { regionBuilder(b.getBlock()->getArguments()); });
+  region.getBlocks().erase(region.begin());
+  assert(op->getNumRegions() == 1);
+  assert(op->getRegion(0).getBlocks().size() == 1);
   return op;
 }
 
 void mlir::edsc::ops::macRegionBuilder(ArrayRef<BlockArgument> args) {
   using edsc::op::operator+;
   using edsc::op::operator*;
   assert(args.size() == 3 && "expected 3 block arguments");
   ValueHandle a(args[0]), b(args[1]), c(args[2]);

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

TEST_FUNC(linalg_pointwise_test) {
   f.erase();
 }
 
 // clang-format off
 // CHECK-LABEL: func @linalg_matmul
 //       CHECK:   linalg.generic {args_in = 2 : i64, args_out = 1 : i64,
 // CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, affine_map<(d0, d1, d2) -> (d2, d1)>, affine_map<(d0, d1, d2) -> (d0, d1)>],
 // CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]}
-///      CHECK:   ^bb1(%[[a0:.*]]: f32, %[[a1:.*]]: f32, %[[a2:.*]]: f32):
+///      CHECK:   ^bb0(%[[a0:.*]]: f32, %[[a1:.*]]: f32, %[[a2:.*]]: f32):
 //       CHECK:     %[[a3:.*]] = mulf %[[a0]], %[[a1]] : f32
 //       CHECK:     %[[a4:.*]] = addf %[[a2]], %[[a3]] : f32
 //       CHECK:     linalg.yield %[[a4]] : f32
 //       CHECK:   }: memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>
 // clang-format on
 TEST_FUNC(linalg_matmul_test) {
   using namespace edsc;
   using namespace edsc::ops;

 
 // clang-format off
 // CHECK-LABEL: func @linalg_conv_nhwc
 //       CHECK:   linalg.generic {args_in = 2 : i64, args_out = 1 : i64,
 // CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d2 * 3 + d4 * 5, d3 * 4 + d5 * 6, d6)>,
 // CHECK-SAME: affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d4, d5, d6, d1)>,
 // CHECK-SAME: affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d2, d3, d1)>],
 // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction"]}
-///      CHECK:   ^bb1(%[[a0:.*]]: f32, %[[a1:.*]]: f32, %[[a2:.*]]: f32):
+///      CHECK:   ^bb0(%[[a0:.*]]: f32, %[[a1:.*]]: f32, %[[a2:.*]]: f32):
 //       CHECK:     %[[a3:.*]] = mulf %[[a0]], %[[a1]] : f32
 //       CHECK:     %[[a4:.*]] = addf %[[a2]], %[[a3]] : f32
 //       CHECK:     linalg.yield %[[a4]] : f32
 //       CHECK:   }: memref<?x?x?x?xf32>, memref<?x?x?x?xf32>, memref<?x?x?x?xf32>
 // clang-format on
 TEST_FUNC(linalg_conv_nhwc) {
   using namespace edsc;
   using namespace edsc::ops;

 
 // clang-format off
 // CHECK-LABEL: func @linalg_dilated_conv_nhwc
 //       CHECK:   linalg.generic {args_in = 2 : i64, args_out = 1 : i64,
 // CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d3 * 3 + d5 * 5, d4 * 4 + d6 * 6, d2)>,
 // CHECK-SAME: affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d5, d6, d2, d1)>,
 // CHECK-SAME: affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d3, d4, d1 + d2 * 7)>],
 // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]}
-//       CHECK:   ^bb1(%[[a0:.*]]: f32, %[[a1:.*]]: f32, %[[a2:.*]]: f32):
+//       CHECK:   ^bb0(%[[a0:.*]]: f32, %[[a1:.*]]: f32, %[[a2:.*]]: f32):
 //       CHECK:     %[[a3:.*]] = mulf %[[a0]], %[[a1]] : f32
 //       CHECK:     %[[a4:.*]] = addf %[[a2]], %[[a3]] : f32
 //       CHECK:     linalg.yield %[[a4]] : f32
 //       CHECK:   }: memref<?x?x?x?xf32>, memref<?x?x?x?xf32>, memref<?x?x?x?xf32>
 // clang-format on
 TEST_FUNC(linalg_dilated_conv_nhwc) {
   using namespace edsc;
   using namespace edsc::ops;