[mlir][linalg] fix crash when promoting rank-reducing memref.subviews

[christopherbate]

Summary

This change adds support for promoting linalg operation operands that
are produced by rank-reducing memref.subview ops. This can occur when
linalg-fold-unit-dims transforms are applied before promotion.

[mravishankar]

Cool! Thanks!

Just a nit on the message with the patch. The use of rank-reducing subviews could from anywhere. The fact that this might/might-not have come from fold-unit-dims pass is orthogonal.

[christopherbate]

In D127086#3560717, @mravishankar wrote:

Cool! Thanks!

Just a nit on the message with the patch. The use of rank-reducing subviews could from anywhere. The fact that this might/might-not have come from fold-unit-dims pass is orthogonal.

Sounds good, I'll update the message before pushing

[christopherbate]

Update commit message

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

 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/SmallBitVector.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 using namespace mlir::scf;
 
FailureOr<PromotionInfo> mlir::linalg::promoteSubviewAsNewBuffer(
     OpBuilder &b, Location loc, memref::SubViewOp subView,
     const AllocBufferCallbackFn &allocationFn, DataLayout &layout) {
   auto viewType = subView.getType();
   auto rank = viewType.getRank();
   SmallVector<Value, 4> fullSizes;
   SmallVector<OpFoldResult> partialSizes;
   fullSizes.reserve(rank);
   partialSizes.reserve(rank);
+  llvm::SmallBitVector droppedDims = subView.getDroppedDims();
+  int64_t resultDimIdx = 0;
   for (const auto &en : llvm::enumerate(subView.getOrCreateRanges(b, loc))) {
+    if (droppedDims[en.index()])
+      continue;
     auto rangeValue = en.value();
     // Try to extract a tight constant.
     LLVM_DEBUG(llvm::dbgs() << "Extract tightest: " << rangeValue.size << "\n");
     FailureOr<int64_t> upperBound =
         getConstantUpperBoundForIndex(rangeValue.size);
     Value size =
         failed(upperBound)
             ? rangeValue.size
             : b.create<arith::ConstantIndexOp>(loc, upperBound.getValue());
     LLVM_DEBUG(llvm::dbgs() << "Extracted tightest: " << size << "\n");
     fullSizes.push_back(size);
     partialSizes.push_back(
-        b.createOrFold<memref::DimOp>(loc, subView, en.index()));
+        b.createOrFold<memref::DimOp>(loc, subView, resultDimIdx++));
   }
   SmallVector<int64_t, 4> dynSizes(fullSizes.size(), -1);
   // If a callback is not specified, then use the default implementation for
   // allocating the promoted buffer.
   Optional<Value> fullLocalView = allocationFn(b, subView, fullSizes, layout);
   if (!fullLocalView)
     return failure();
   SmallVector<OpFoldResult, 4> zeros(fullSizes.size(), b.getIndexAttr(0));

mlir/test/Dialect/Linalg/promote.mlir

-// RUN: mlir-opt %s -linalg-promote-subviews | FileCheck %s
-// RUN: mlir-opt %s -linalg-promote-subviews="test-promote-dynamic" | FileCheck %s --check-prefix=DYNAMIC
-// RUN: mlir-opt %s -linalg-promote-subviews="test-use-alloca" | FileCheck %s --check-prefix=ALLOCA
+// RUN: mlir-opt %s -linalg-promote-subviews -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -linalg-promote-subviews="test-promote-dynamic" -split-input-file | FileCheck %s --check-prefix=DYNAMIC
+// RUN: mlir-opt %s -linalg-promote-subviews="test-use-alloca" -split-input-file | FileCheck %s --check-prefix=ALLOCA
 
 #map1 = affine_map<(d0) -> (d0 + 2)>
 #map2 = affine_map<(d0) -> (d0 + 4)>
 #map3 = affine_map<(d0) -> (d0 + 3)>
 
 // CHECK-DAG: #[[$strided2D:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
 
 func.func @matmul_f32(%A: memref<?xi8>, %M: index, %N: index, %K: index) {

 //
 //       CHECK:         memref.copy %[[partialC_f64]], %[[vC_f64]] :
 //       CHECK:           memref<?x?xf64, #[[$strided2D]]> to
 //       CHECK:           memref<?x?xf64, #[[$strided2D]]>
 //
 //       CHECK:         memref.dealloc %[[tmpA_f64]] : memref<64xi8>
 //       CHECK:         memref.dealloc %[[tmpB_f64]] : memref<96xi8>
 //       CHECK:         memref.dealloc %[[tmpC_f64]] : memref<48xi8>
+
+// -----
+
+#map0 = affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3, s4] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3 + d3 * s4)>
+#map2 = affine_map<(d0, d1)[s0] -> (d0 * 128 + s0 + d1)>
+#map5 = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
+#map6 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map7 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map8 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// CHECK: promote_rank_reducing_subviews([[arg0:%.+]]: memref<{{.*}}>, [[arg1:%.+]]: memref<{{.*}}>, [[arg2:%.+]]: memref<{{.*}}>, [[lb1:%.+]]: index, [[lb2:%.+]]: index, [[lb3:%.+]]: index, [[lb4:%.+]]: index, [[lb5:%.+]]: index, [[lb6:%.+]]: index, [[ub1:%.+]]: index, [[ub2:%.+]]: index
+func.func @promote_rank_reducing_subviews(%arg0:  memref<?x?x?x64xf32, #map0>, %arg1: memref<128x3x3x64xf32, #map0>, %arg2: memref<?x?x?x128xf32>,
+                                          %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: index, %arg8: index, %ub1: index, %ub2: index) {
+  %13 = memref.subview %arg0[%arg3, 0, %arg4, %arg8] [1, 1, %ub1, 32] [1, 1, 1, 1] : memref<?x?x?x64xf32, #map0> to memref<?x32xf32, #map5>
+  %14 = memref.subview %arg1[0, %arg6, %arg7, %arg8] [128, 1, 1, 32] [1, 1, 1, 1] : memref<128x3x3x64xf32, #map0> to memref<128x32xf32, #map5>
+  %9 = memref.subview %arg2[%arg3, %arg4, %arg5, 0] [1, 1, %ub2, 128] [1, 1, 1, 1] : memref<?x?x?x128xf32> to memref<?x128xf32, #map2>
+
+  // CHECK: [[a_alloc:%.+]] = memref.alloc
+  // CHECK: [[a_view:%.+]] = memref.view [[a_alloc]]{{.*}}                 
+  // CHECK: [[a_pro_subview:%.+]] = memref.subview [[a_view]][0, 0] [[[ub1]], {{%.+}}] [1, 1]
+
+  // CHECK: memref.alloc
+  // CHECK: [[b_view:%.+]] = memref.view
+  // CHECK: [[b_pro_subview:%.+]] = memref.subview [[b_view]]
+
+  // CHECK: memref.alloc
+  // CHECK: [[c_view:%.+]] = memref.view
+  // CHECK: [[c_pro_subview:%.+]] = memref.subview [[c_view]]
+
+  // CHECK-COUNT-3: memref.copy
+  // CHECK: linalg.generic
+  // CHECK-SAME: ins([[a_pro_subview]], [[b_pro_subview]]
+  // CHECK-SAME: outs([[c_pro_subview]]
+
+  linalg.generic {indexing_maps = [#map6, #map7, #map8], iterator_types = ["parallel", "parallel", "reduction"]} ins(%13, %14 : memref<?x32xf32, #map5>, memref<128x32xf32, #map5>) outs(%9 : memref<?x128xf32, #map2>) {
+  ^bb0(%arg9: f32, %arg10: f32, %arg11: f32):
+    %15 = arith.mulf %arg9, %arg10 : f32
+    %16 = arith.addf %arg11, %15 : f32
+    linalg.yield %16 : f32
+  }
+
+  return
+}