[mlir][linalg] Lower linalg to affine loops

[nicolasvasilache]

Looks great and is as simple as it should be, thanks!
Accepting conditioned on the nit above.

cc @ftynse @AlexEichenberger @mravishankar

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

114

please fold these 3 defs into their uses

[rriddle]

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

12

Why is this header necessary?

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

25

Is this really necessary here?

112

const auto -> Value

[asaadaldien]

Resolve comments...

[merge_guards_bot]

Unit tests: pass. 61146 tests passed, 0 failed and 728 were skipped.

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

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

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

[asaadaldien]

lint fixes...

[merge_guards_bot]

Unit tests: pass. 61146 tests passed, 0 failed and 728 were skipped.

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

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

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

[ftynse]

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

68

MLIR uses /// for class documentation

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

113

Does something guarantee that there is a defining op?

[mravishankar]

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

113

You can probably change emitLoopRanges that generates the ranges argument at the called of this function to return Operation * instead of Value.

114

Nit: I think you can use lbs.emplace_back(rangeOp.min()). Same below

[asaadaldien]

Resolve comments

[asaadaldien]

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

113

Added assertions for the underlying type

113

We can keep it as ranges to have now to keep interface same as underlying RangeBuilders. We can revisit using ranges here in a different diff.

[asaadaldien]

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

113

We can keep it as ranges for now to keep interface same as underlying RangeBuilders.
We can revisit using ranges here in a different diff.

[merge_guards_bot]

Unit tests: pass. 61146 tests passed, 0 failed and 728 were skipped.

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

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

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

[asaadaldien]

applying clang-format

[merge_guards_bot]

Unit tests: pass. 61146 tests passed, 0 failed and 728 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

[jpienaar]

This seems to break on GCC 5.5:

mlir/mlir-core/mlir/lib/Dialect/Linalg/Utils/Utils.cpp:102:62: error: specialization of 'template<class LoopTy> mlir::edsc::GenericLoopNestRangeBuilder<LoopTy>::GenericLoopNestRangeBuilder(llvm::ArrayRef<mlir::edsc::ValueHandle*>, llvm::ArrayRef<mlir::Value>)' in different namespace [-fpermissive]
     ArrayRef<edsc::ValueHandle *> ivs, ArrayRef<Value> ranges) {
                                                              ^
mlir/mlir-core/mlir/lib/Dialect/Linalg/Utils/Utils.cpp:13:0:
/var/lib/buildkite-agent/builds/buildkite-6dd7db7bf5-48dzm-1/mlir/mlir-core/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h:72:3: error:   from definition of 'template<class LoopTy> mlir::edsc::GenericLoopNestRangeBuilder<LoopTy>::GenericLoopNestRangeBuilder(llvm::ArrayRef<mlir::edsc::ValueHandle*>, llvm::ArrayRef<mlir::Value>)' [-fpermissive]
   GenericLoopNestRangeBuilder(ArrayRef<edsc::ValueHandle *> ivs,
   ^
mlir/mlir-core/mlir/lib/Dialect/Linalg/Utils/Utils.cpp:108:56: error: specialization of 'template<class LoopTy> mlir::edsc::GenericLoopNestRangeBuilder<LoopTy>::GenericLoopNestRangeBuilder(llvm::ArrayRef<mlir::edsc::ValueHandle*>, llvm::ArrayRef<mlir::Value>)' in different namespace [-fpermissive]
     ArrayRef<ValueHandle *> ivs, ArrayRef<Value> ranges) {
                                                        ^
mlir/mlir-core/mlir/lib/Dialect/Linalg/Utils/Utils.cpp:13:0:
/var/lib/buildkite-agent/builds/buildkite-6dd7db7bf5-48dzm-1/mlir/mlir-core/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h:72:3: error:   from definition of 'template<class LoopTy> mlir::edsc::GenericLoopNestRangeBuilder<LoopTy>::GenericLoopNestRangeBuilder(llvm::ArrayRef<mlir::edsc::ValueHandle*>, llvm::ArrayRef<mlir::Value>)' [-fpermissive]
   GenericLoopNestRangeBuilder(ArrayRef<edsc::ValueHandle *> ivs,
   ^

[asaadaldien]

@jpienaar D72173 should fix that

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

 //===- Utils.h - Utilities to support the Linalg dialect --------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_LINALG_UTILS_H_
 #define MLIR_DIALECT_LINALG_UTILS_H_
 
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/LoopOps/LoopOps.h"
 #include "mlir/Dialect/StandardOps/Ops.h"
 #include "mlir/EDSC/Helpers.h"
 
 #include "llvm/ADT/SetVector.h"
 
 namespace mlir {
 class AffineExpr;
public:
   LoopNestRangeBuilder(ArrayRef<edsc::ValueHandle *> ivs,
                        ArrayRef<SubViewOp::Range> ranges);
   edsc::ValueHandle operator()(std::function<void(void)> fun = nullptr);
 
 private:
   SmallVector<LoopRangeBuilder, 4> loops;
 };
 
+/// Helper template class for building loop.for and affine.loop nests from
+/// ranges.
+template <typename LoopTy> class GenericLoopNestRangeBuilder {
+public:
+  GenericLoopNestRangeBuilder(ArrayRef<edsc::ValueHandle *> ivs,
+                              ArrayRef<Value> ranges);
+  void operator()(std::function<void(void)> fun = nullptr) { (*builder)(fun); }
+
+private:
+  typedef typename std::conditional<std::is_same<LoopTy, AffineForOp>::value,
+                                    AffineLoopNestBuilder,
+                                    LoopNestRangeBuilder>::type BuilderType;
+  std::unique_ptr<BuilderType> builder;
+};
+
 } // namespace edsc
 
 namespace linalg {
 class LinalgDependenceGraph;
 
 struct FusionInfo {
   LinalgOp originalProducer;
   LinalgOp fusedProducer;

mlir/lib/Dialect/Linalg/Transforms/LinalgToLoops.cpp

 //===- LowerToLoops.cpp - conversion from Linalg library ops to loops------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Dialect/AffineOps/AffineOps.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Linalg/Transforms/LinalgTransforms.h"
 #include "mlir/Dialect/Linalg/Utils/Intrinsics.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/LoopOps/LoopOps.h"
 #include "mlir/Dialect/StandardOps/Ops.h"
LogicalResult LinalgOpToLoopsImpl<LoopTy, IndexedValueTy, ConcreteOpTy>::doit(
   auto nWin = linalgOp.getNumWindowLoops();
   SmallVector<IndexHandle, 4> allIvs(nPar + nRed + nWin);
   SmallVector<ValueHandle *, 4> allPIvs =
       makeHandlePointers(MutableArrayRef<IndexHandle>(allIvs));
   auto loopRanges = emitLoopRanges(scope.getBuilder(), scope.getLocation(),
                                    invertedMap, getViewSizes(linalgOp));
   assert(loopRanges.size() == allIvs.size());
 
-  LoopNestRangeBuilder(allPIvs, loopRanges)([&] {
+  GenericLoopNestRangeBuilder<LoopTy>(allPIvs, loopRanges)([&] {
     auto allIvValues = extractValues(allIvs);
     LinalgScopedEmitter<IndexedValueTy, ConcreteOpTy>::emitScalarImplementation(
         allIvValues, linalgOp);
   });
   return success();
 }
 
 template <typename LoopType, typename IndexedValueType, typename ConcreteOp>

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

 #include "mlir/Dialect/Linalg/Utils/Intrinsics.h"
 #include "mlir/Dialect/LoopOps/LoopOps.h"
 #include "mlir/Dialect/StandardOps/Ops.h"
 #include "mlir/EDSC/Helpers.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/STLExtras.h"
 #include "mlir/Transforms/FoldUtils.h"
 
 using namespace mlir;
 using namespace mlir::edsc;
 using namespace mlir::edsc::intrinsics;
 using namespace mlir::linalg;
 using namespace mlir::linalg::intrinsics;
 using namespace mlir::loop;
ValueHandle LoopNestRangeBuilder::LoopNestRangeBuilder::operator()(
   if (fun)
     fun();
   for (auto &lit : reverse(loops)) {
     lit({});
   }
   return ValueHandle::null();
 }
 
+template <>
+GenericLoopNestRangeBuilder<loop::ForOp>::GenericLoopNestRangeBuilder(
+    ArrayRef<edsc::ValueHandle *> ivs, ArrayRef<Value> ranges) {
+  builder = std::make_unique<LoopNestRangeBuilder>(ivs, ranges);
+}
+
+template <>
+GenericLoopNestRangeBuilder<AffineForOp>::GenericLoopNestRangeBuilder(
+    ArrayRef<ValueHandle *> ivs, ArrayRef<Value> ranges) {
+  SmallVector<ValueHandle, 4> lbs;
+  SmallVector<ValueHandle, 4> ubs;
+  SmallVector<int64_t, 4> steps;
+  for (const Value range : ranges) {
+    assert(range.getType() && "expected linalg.range type");
+    assert(range.getDefiningOp() && "need operations to extract range parts");
+    RangeOp rangeOp = cast<RangeOp>(range.getDefiningOp());
+    lbs.emplace_back(ValueHandle(rangeOp.min()));
+    ubs.emplace_back(ValueHandle(rangeOp.max()));
+    steps.emplace_back(ValueHandle(rangeOp.step()));
+  }
+  builder = std::make_unique<AffineLoopNestBuilder>(ivs, lbs, ubs, steps);
+}
+
 static Value emitOrFoldComposedAffineApply(OpBuilder &b, Location loc,
                                            AffineMap map,
                                            ArrayRef<Value> operandsRef,
                                            OperationFolder *folder) {
   SmallVector<Value, 4> operands(operandsRef.begin(), operandsRef.end());
   fullyComposeAffineMapAndOperands(&map, &operands);
   canonicalizeMapAndOperands(&map, &operands);
   return folder ? folder->create<AffineApplyOp>(b, loc, map, operands)

mlir/test/Dialect/Linalg/affine.mlir

+// RUN: mlir-opt %s -convert-linalg-to-affine-loops | FileCheck %s
+
+// Test that we can lower all the way to LLVM without crashing, don't check results here.
+// RUN: mlir-opt %s --convert-linalg-to-llvm -o=/dev/null 2>&1
+
+// CHECK-DAG: #[[strided2D:.*]] = (d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)
+// CHECK-DAG: #[[strided3D:.*]] = (d0, d1, d2)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2 + d2)
+
+// CHECK-DAG: #[[stride2Dilation1:.*]] = (d0, d1) -> (d0 * 2 + d1)
+
+func @matmul(%arg0: memref<?xi8>, %M: index, %N: index, %K: index) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %A = view %arg0[%c0][%M, %K] : memref<?xi8> to memref<?x?xf32, offset: ?, strides: [?, 1]>
+  %B = view %arg0[%c0][%K, %N] : memref<?xi8> to memref<?x?xf32, offset: ?, strides: [?, 1]>
+  %C = view %arg0[%c0][%M, %N] : memref<?xi8> to memref<?x?xf32, offset: ?, strides: [?, 1]>
+  linalg.matmul(%A, %B, %C) : memref<?x?xf32, offset: ?, strides: [?, 1]>, memref<?x?xf32, offset: ?, strides: [?, 1]>, memref<?x?xf32, offset: ?, strides: [?, 1]>
+  return
+}
+
+// CHECK-LABEL: func @matmul(%{{.*}}: memref<?xi8>,
+// CHECK-SAME: [[M:arg[0-9]+]]: index
+// CHECK-SAME: [[N:arg[0-9]+]]: index
+// CHECK-SAME: [[K:arg[0-9]+]]: index
+//       CHECK: %[[A:.*]] = std.view %{{.*}}[{{.*}}] : memref<?xi8> to memref<?x?xf32, #[[strided2D]]>
+//       CHECK: %[[B:.*]] = std.view %{{.*}}[{{.*}}] : memref<?xi8> to memref<?x?xf32, #[[strided2D]]>
+//       CHECK: %[[C:.*]] = std.view %{{.*}}[{{.*}}] : memref<?xi8> to memref<?x?xf32, #[[strided2D]]>
+//       CHECK: affine.for %{{.*}}  = 0 to %{{.*}} {
+//       CHECK:   affine.for %{{.*}} = 0 to %{{.*}} {
+//       CHECK:     affine.for %{{.*}} = 0 to %{{.*}} {
+//   CHECK-DAG:       %[[a:.*]] = affine.load %[[A]][%{{.*}}, %{{.*}}] : memref<?x?xf32, #[[strided2D]]>
+//   CHECK-DAG:       %[[b:.*]] = affine.load %[[B]][%{{.*}}, %{{.*}}] : memref<?x?xf32, #[[strided2D]]>
+//   CHECK-DAG:       %[[inc:.*]] = mulf %[[a]], %[[b]] : f32
+//   CHECK-DAG:       %[[c:.*]] = affine.load %[[C]][%{{.*}}, %{{.*}}] : memref<?x?xf32, #[[strided2D]]>
+//   CHECK-DAG:       %[[res:.*]] = addf %[[c]], %[[inc]] : f32
+//       CHECK:       affine.store %[[res]], %[[C]][%{{.*}}, %{{.*}}] : memref<?x?xf32, #[[strided2D]]>
+
+func @conv_view3(%arg0: memref<?x?x?xf32, offset: ?, strides: [?, ?, 1]>, %arg1: memref<?x?x?xf32, offset: ?, strides: [?, ?, 1]>, %arg2: memref<?x?x?xf32, offset: ?, strides: [?, ?, 1]>) {
+  linalg.conv(%arg0, %arg1, %arg2) {strides = [2]}: memref<?x?x?xf32, offset: ?, strides: [?, ?, 1]>, memref<?x?x?xf32, offset: ?, strides: [?, ?, 1]>, memref<?x?x?xf32, offset: ?, strides: [?, ?, 1]>
+  return
+}
+
+// CHECK-LABEL: func @conv_view3(
+//  CHECK: %{{.*}}: memref<?x?x?xf32, #[[strided3D]]>, %{{.*}}: memref<?x?x?xf32, #[[strided3D]]>, %{{.*}}: memref<?x?x?xf32, #[[strided3D]]>) {
+//       CHECK:   %[[Z0:.*]] = dim %arg0, 0 : memref<?x?x?xf32, #[[strided3D]]>
+//       CHECK:   %[[Q:.*]] = dim %arg0, 1 : memref<?x?x?xf32, #[[strided3D]]>
+//       CHECK:   %[[K:.*]] = dim %arg0, 2 : memref<?x?x?xf32, #[[strided3D]]>
+//       CHECK:   %[[B:.*]] = dim %arg1, 0 : memref<?x?x?xf32, #[[strided3D]]>
+//       CHECK:   %[[X0:.*]] = dim %arg2, 1 : memref<?x?x?xf32, #[[strided3D]]>
+//       CHECK:   affine.for %{{.*}} = 0 to %[[B]] {
+//       CHECK:     affine.for %{{.*}} = 0 to %[[X0]] {
+//       CHECK:       affine.for %{{.*}} = 0 to %[[K]] {
+//       CHECK:         affine.for %{{.*}} = 0 to %[[Q]] {
+//       CHECK:           affine.for %{{.*}} = 0 to %[[Z0]] {
+//       CHECK:            %[[SUM:.*]] = affine.apply #[[stride2Dilation1]](%{{.*}}, %{{.*}})