[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

116

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

26

Is this really necessary here?

114

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

70

MLIR uses /// for class documentation

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

115

Does something guarantee that there is a defining op?

[mravishankar]

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

115

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

116

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

[asaadaldien]

Resolve comments

[asaadaldien]

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

115

Added assertions for the underlying type

115

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

115

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/AffineOps/AffineOps.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/LoopOps/LoopOps.h"
 #include "mlir/Dialect/StandardOps/Ops.h"
+#include "mlir/EDSC/Builders.h"
 #include "mlir/EDSC/Helpers.h"
 
 #include "llvm/ADT/SetVector.h"
 
 namespace mlir {
 class AffineExpr;
 class AffineMap;
 class OperationFolder;
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

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/Utils.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/Utils/Intrinsics.h"
 #include "mlir/Dialect/LoopOps/LoopOps.h"
 #include "mlir/Dialect/StandardOps/Ops.h"
+#include "mlir/EDSC/Builders.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/LLVM.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;
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 auto range : ranges) {
+    auto rangeOp = cast<RangeOp>(range.getDefiningOp());
+    auto lb = rangeOp.min();
+    auto ub = rangeOp.max();
+    auto step = rangeOp.step();
+    lbs.push_back(ValueHandle(lb));
+    ubs.push_back(ValueHandle(ub));
+    steps.push_back(ValueHandle(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]](%{{.*}}, %{{.*}})