[mlir] [linalg] Only promote selected buffers.

[poechsel]

Summary

The promotion transformation is promoting all input and output buffers of the transformed op. The user might want to only promote some of these buffers.

[poechsel]

Clang formatting suggestions.

[ftynse]

mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransforms.h

124

Please document this function

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

342

Please reserve the space before populating a vector in a loop

343

int64_t i = 0, e = linOp.getNumInputsAndOutputBuffers(); i < e; ++i) is a common LLVM pattern that avoids calling the function in a loop

376

Nit: !linalgMarker.empty() may be more efficient

mlir/test/lib/DeclarativeTransforms/TestLinalgTransformPatterns.td

153

I don't follow this. The test name suggests you promote the _first_ view, but the list contains operands with indices 0 and 1... The .mlir file also seems to check only for one copy. Consider using only [0] here and CHECK-NOT for extra copies before matmul.

[poechsel]

Address comments

[poechsel]

mlir/test/lib/DeclarativeTransforms/TestLinalgTransformPatterns.td

153

Nice catch!

I believe there is a problem with test/Dialect/Linalg/transform-patterns.mlir and that the test is passing even when it should fail:

• in my case I should have had a fallback on line 439 where I'm checking the ssa address of arg1 with the one it should have if no promotion happens.
• since D77676 promotion shouldn't generate linalg.slice operations but this test is still checking for their presence (ex: line 393).

[poechsel]

Fix test

[ftynse]

mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransforms.h

127

Nit: typo "successful". And grammo: put a comma after an if-clause.

mlir/test/lib/DeclarativeTransforms/TestLinalgTransformPatterns.td

153

Indeed, after further investigation, the test above uses wrong syntax for FileCheck. There must be no spaces between CHECK and : @nicolasvasilache.

[rriddle]

mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransformPatterns.td

119

Please wrap this at 80 characters, it seems to be edging over here.

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

345

nit: You could use std::iota here.

[thopre]

mlir/test/Dialect/Linalg/transform-patterns.mlir

439–446

These checks do not work because they refer to undefined variables because each CHECK-NOT is checked independently. Since these are patterns that should not occur the variables defined in them have nothing to be set to.

Currently these will always succeed because the use of an undefined variable is interpreted by FileCheck as being a failure to match, and thus those CHECK-NOT are satisfied. I'd suggest replacing them with CHECK-NEXT as generated by update_test_checks.py. Alternatively you could join all CHECK-NOT in a CHECK-NOT block with {{(.*[[:space:]])+}} but that would become quite unreadable.

mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransformPatterns.td

 
 //===----------------------------------------------------------------------===//
 // Linalg promote subview operands precondition and DRR.
 //===----------------------------------------------------------------------===//
 def PreconditionPromoteSubviewsLinalgOp : CPred<
   "succeeded(promoteSubviewsLinalgOpPrecondition(op))">;
 def PromoteSubviewsLinalgOp : NativeCodeCall<
   "promoteSubviewsLinalgOp($_builder, op)">;
+
+class PromoteSelectedSubviewsLinalgOp<list<int> operands, string marker=""> :
+  NativeCodeCall<"promoteSelectedSubviewsLinalgOpAndSetMarker($_builder, op, {" #
+    StrJoinInt<operands>.result # "}, \"" # marker # "\")">;
+
 #endif // LINALG_TRANSFORMS

mlir/include/mlir/Dialect/Linalg/Transforms/LinalgTransforms.h

SmallVector<Value, 0> permuteGenericLinalgOp(PatternRewriter &rewriter,
                                              ArrayRef<unsigned> permutation,
                                              StringRef linalgMarker);
 
 /// Promote std.subviews feeding linalg operations.
 LogicalResult promoteSubviewsLinalgOpPrecondition(Operation *op);
 SmallVector<Value, 0> promoteSubviewsLinalgOp(PatternRewriter &rewriter,
                                               Operation *op);
 
+/// Similar to `promoteSubviewsLinalgOp` but only tries to promote
+/// the views corresponding to the operands specified in
+/// `operandIndicesToPromote`.
+/// If linalgMarker is specified and the transformation is successfull
+/// sets the attribute `kLinalgTransformMarker` to `linalgMarker`.
+SmallVector<Value, 0> promoteSelectedSubviewsLinalgOpAndSetMarker(
+    PatternRewriter &rewriter, Operation *op,
+    ArrayRef<int64_t> operandIndicesToPromote, StringRef linalgMarker = "");
 } // namespace linalg
 } // namespace mlir
 
 #endif // DIALECT_LINALG_TRANSFORMS_LINALGTRANSFORMS_H_

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

 mlir::linalg::promoteSubviewsLinalgOp(PatternRewriter &rewriter,
                                       Operation *op) {
   LLVM_DEBUG(dbgs() << "\n[" DEBUG_TYPE "]: Promote subviews for linalg op: "
                     << *op << ":\n");
 
   assert(succeeded(promoteSubviewsLinalgOpPrecondition(op)) &&
          "DRR failure case must be a precondition");
 
+  LinalgOp linOp = cast<LinalgOp>(op);
+  SmallVector<int64_t, 4> toPromote;
+  int64_t nBuffers = linOp.getNumInputsAndOutputBuffers();
+  toPromote.reserve(nBuffers);
+  for (int64_t i = 0; i < nBuffers; ++i)
+    toPromote.push_back(i);
+  return promoteSelectedSubviewsLinalgOpAndSetMarker(rewriter, op, toPromote);
+}
+
+SmallVector<Value, 0> mlir::linalg::promoteSelectedSubviewsLinalgOpAndSetMarker(
+    PatternRewriter &rewriter, Operation *op,
+    ArrayRef<int64_t> operandIndicesToPromote, StringRef linalgMarker) {
+  LLVM_DEBUG(dbgs() << "\n[" DEBUG_TYPE "]: Promote subviews for linalg op: "
+                    << *op << ":\n");
+
+  assert(succeeded(promoteSubviewsLinalgOpPrecondition(op)) &&
+         "DRR failure case must be a precondition");
+
   if (auto convOp = dyn_cast<linalg::ConvOp>(op)) {
     // TODO(ntv): add a level of indirection to linalg.generic.
     if (convOp.padding())
       llvm_unreachable("Unexpected conv with padding");
   }
 
   LinalgOp linOp = cast<LinalgOp>(op);
   assert(linOp.hasBufferSemantics() &&
          "expected linalg op with buffer semantics");
   SetVector<Value> subViews;
-  for (auto it : linOp.getInputsAndOutputBuffers())
-    if (auto sv = dyn_cast_or_null<SubViewOp>(it.getDefiningOp()))
+  for (int64_t index : operandIndicesToPromote)
+    if (auto sv =
+            dyn_cast_or_null<SubViewOp>(linOp.getBuffer(index).getDefiningOp()))
       subViews.insert(sv);
+
   if (!subViews.empty()) {
-    promoteSubViewOperands(rewriter, linOp, subViews);
+    auto newOp = promoteSubViewOperands(rewriter, linOp, subViews);
+    if (!linalgMarker.empty())
+      newOp.setAttr(LinalgTransforms::kLinalgTransformMarker,
+                    rewriter.getStringAttr(linalgMarker));
     return {};
   }
   llvm_unreachable("DRR failure case must be a precondition");
 }

mlir/test/Dialect/Linalg/transform-patterns.mlir

 // CHECK      :         %[[l1:.*]] = linalg.slice %[[v1]][{{%.*}}, {{%.*}}] : memref<?x?xf32>, !linalg.range, !linalg.range, memref<?x?xf32, #map{{.*}}>
 // CHECK      :         %[[a2:.*]] = alloc({{%.*}}) : memref<?xi8>
 // CHECK      :         %[[v2:.*]] = std.view %[[a2]][][{{%.*}}, {{%.*}}]: memref<?xi8> to memref<?x?xf32>
 // CHECK      :         %[[l2:.*]] = linalg.slice %[[v2]][{{%.*}}, {{%.*}}] : memref<?x?xf32>, !linalg.range, !linalg.range, memref<?x?xf32, #map{{.*}}>
 // CHECK      :         linalg.copy(%[[s0]], %[[l0]]) : memref<?x?xf32, #map{{.*}}>, memref<?x?xf32, #map{{.*}}>
 // CHECK      :         linalg.copy(%[[s1]], %[[l1]]) : memref<?x?xf32, #map{{.*}}>, memref<?x?xf32, #map{{.*}}>
 // CHECK      :         linalg.copy(%[[s2]], %[[l2]]) : memref<?x?xf32, #map{{.*}}>, memref<?x?xf32, #map{{.*}}>
 // CHECK      :         linalg.matmul(%[[v0]], %[[v1]], %[[v2]]) : memref<?x?xf32, #[[STRIDED_2D]]>, memref<?x?xf32, #[[STRIDED_2D]]>, memref<?x?xf32, #[[STRIDED_2D]]>
+
+func @promote_first_subview_matmul(%arg0: memref<?x?xf32, offset: ?, strides: [?, 1]>,
+                             %arg1: memref<?x?xf32, offset: ?, strides: [?, 1]>,
+                             %arg2: memref<?x?xf32, offset: ?, strides: [?, 1]>) {
+  %c2000 = constant 2000 : index
+  %c3000 = constant 3000 : index
+  %c4000 = constant 4000 : index
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %0 = dim %arg0, 0 : memref<?x?xf32, offset: ?, strides: [?, 1]>
+  %1 = dim %arg0, 1 : memref<?x?xf32, offset: ?, strides: [?, 1]>
+  %2 = dim %arg1, 1 : memref<?x?xf32, offset: ?, strides: [?, 1]>
+  loop.for %arg3 = %c0 to %0 step %c2000 {
+    loop.for %arg4 = %c0 to %2 step %c3000 {
+      loop.for %arg5 = %c0 to %1 step %c4000 {
+        %3 = std.subview %arg0[%arg3, %arg5][%c2000, %c4000][%c1, %c1] :
+             memref<?x?xf32, offset: ?, strides: [?, 1]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %4 = std.subview %arg1[%arg5, %arg4][%c4000, %c3000][%c1, %c1] :
+             memref<?x?xf32, offset: ?, strides: [?, 1]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
+        %5 = std.subview %arg2[%arg3, %arg4][%c2000, %c3000][%c1, %c1] :
+             memref<?x?xf32, offset: ?, strides: [?, 1]> to memref<?x?xf32, offset: ?, strides: [?, ?]>
+        linalg.matmul(%3, %4, %5) {__internal_linalg_transform__ = "_promote_first_view_"} :
+                      memref<?x?xf32, offset: ?, strides: [?, ?]>,
+                      memref<?x?xf32, offset: ?, strides: [?, ?]>,
+                      memref<?x?xf32, offset: ?, strides: [?, ?]>
+      }
+    }
+  }
+  return
+}
+// CHECK-LABEL: func @promote_first_subview_matmul
+// CHECK      :   loop.for {{.*}} = %[[c0]] to {{.*}} step %[[c2000]] {
+// CHECK      :     loop.for {{.*}} = %[[c0]] to {{.*}} step %[[c3000]] {
+// CHECK      :       loop.for {{.*}} = %[[c0]] to {{.*}} step %[[c4000]] {
+// CHECK      :         %[[s0:.*]] = std.subview {{%.*}}[{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] : memref<?x?xf32, #map{{.*}}> to memref<?x?xf32, #map{{.*}}>
+// CHECK      :         %[[s1:.*]] = std.subview {{%.*}}[{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] : memref<?x?xf32, #map{{.*}}> to memref<?x?xf32, #map{{.*}}>
+// CHECK      :         %[[s2:.*]] = std.subview {{%.*}}[{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] : memref<?x?xf32, #map{{.*}}> to memref<?x?xf32, #map{{.*}}>
+// CHECK      :         %[[a0:.*]] = alloc({{%.*}}) : memref<?xi8>
+// CHECK      :         %[[v0:.*]] = std.view %[[a0]][][{{%.*}}, {{%.*}}]: memref<?xi8> to memref<?x?xf32>
+// CHECK      :         %[[l0:.*]] = linalg.slice %[[v0]][{{%.*}}, {{%.*}}] : memref<?x?xf32>, !linalg.range, !linalg.range, memref<?x?xf32, #map{{.*}}>
+// CHECK-NOT  :         %[[a1:.*]] = alloc({{%.*}}) : memref<?xi8>
+// CHECK-NOT  :         %[[v1:.*]] = std.view %[[a1]][][{{%.*}}, {{%.*}}]: memref<?xi8> to memref<?x?xf32>
+// CHECK-NOT  :         %[[l1:.*]] = linalg.slice %[[v1]][{{%.*}}, {{%.*}}] : memref<?x?xf32>, !linalg.range, !linalg.range, memref<?x?xf32, #map{{.*}}>
+// CHECK-NOT  :         %[[a2:.*]] = alloc({{%.*}}) : memref<?xi8>
+// CHECK-NOT  :         %[[v2:.*]] = std.view %[[a2]][][{{%.*}}, {{%.*}}]: memref<?xi8> to memref<?x?xf32>
+// CHECK-NOT  :         %[[l2:.*]] = linalg.slice %[[v2]][{{%.*}}, {{%.*}}] : memref<?x?xf32>, !linalg.range, !linalg.range, memref<?x?xf32, #map{{.*}}>
+// CHECK      :         linalg.copy(%[[s0]], %[[l0]]) : memref<?x?xf32, #map{{.*}}>, memref<?x?xf32, #map{{.*}}>
+// CHECK-NOT  :         linalg.copy(%[[s1]], %[[l1]]) : memref<?x?xf32, #map{{.*}}>, memref<?x?xf32, #map{{.*}}>
+// CHECK-NOT  :         linalg.copy(%[[s2]], %[[l2]]) : memref<?x?xf32, #map{{.*}}>, memref<?x?xf32, #map{{.*}}>^
+// CHECK      :         linalg.matmul(%[[v0]], %[[s1]], %[[s2]]) : memref<?x?xf32, #[[STRIDED_2D]]>, memref<?x?xf32, #[[STRIDED_2D]]>, memref<?x?xf32, #[[STRIDED_2D]]>

mlir/test/lib/DeclarativeTransforms/TestLinalgTransformPatterns.td

 def : Pat<(MatmulOp:$op $_, $_, $_),
           (PromoteSubviewsLinalgOp),
           [(Constraint<And<[
               PreconditionPromoteSubviewsLinalgOp,
               HasOperandsOfType<"SubViewOp">,
               HasLinalgTransformMarker<"_promote_views_">]>>
            )]>;
 
+def : Pat<(MatmulOp:$op $_, $_, $_),
+          (PromoteSelectedSubviewsLinalgOp<[0], "first_view_promotion">),
+          [(Constraint<And<[
+              PreconditionPromoteSubviewsLinalgOp,
+              HasOperandsOfType<"SubViewOp">,
+              HasLinalgTransformMarker<"_promote_first_view_">]>>
+           )]>;
+
 #endif // TEST_LINALG_TRANSFORMS_PATTERNS