[mlir] Rank reducing subview conversion to LLVM

[limo1996]

Summary

This commit adjusts SubViewOp lowering to take rank reduction into account.

Depends On D88879

[nicolasvasilache]

Thanks @limo1996 , let's also add a conversion to LLVM test somewhere around here: https://github.com/llvm/llvm-project/blob/master/mlir/test/Conversion/StandardToLLVM/convert-to-llvmir.mlir#L1222

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

2997

Let's please extract this into a helper function.
If I am not mistaken this could also be reused in Ops.cpp for isRankReduced?

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2818

This should not be exposed to Ops.h and ideally deprecated.
The problem is at this time the result types of map_range and ArrayAttr::getAsValueRange are unfathomble.

This is a rare case where I'd say it's ok to copy the function where you need it (and keep it static) for now.

[limo1996]

Comments of Nicolas incorporated

[nicolasvasilache]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.h

307

The name of the function and the doc comment are quite misleading.
From this, I wouldn't know how to find this functionality in 1 month once this gets out of my short-term memory...

I'd suggest :

Given an `originalShape` and a `reducedShape` assumed to be a subset of `originalShape` with some `1` entries erased, return the vector of booleans that specifies which of the entries of `originalShape` are keep to obtain `reducedShape`.
The returned mask can be applied as a projection to `originalShape` to obtain the `reducedShape`.
This mask is useful to track which dimensions must be kept when e.g. compute MemRef strides under rank-reducing operations.
Return None if reducedShape cannot be obtained by dropping only `1` entries in `originalShape`.

llvm::Optional<SmallVector<bool, 4>> computeRankReductionMask(ArrayRef<int64_t> originalShape, ArrayRef<int64_t> reducedShape);

Feel free to further improve amd/or find a better name if one comes to mind.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2827

reducedIdx is misleading, it should be called numReducedDims.

Also, it seems that you only use this in the verification.
So the other use of the API is less intuitive.

If you implement this with returning an llvm::Optional<SmallVector<bool, 4>>, you can test on whether a value is actually returned.

2835

reads somewhat hacky with the hardcoded -2.

Let's please rewrite as:

keepMask[originalIdx] = (reducedIdx < reducedRank) && 
  (originalShape[originalIdx] == reducedShape[reducedIdx]);
if (keepMask[originalIdx])
  reducedIdx++;

2863

if (reducedRank != std::count_if(keepMask.begin(), keepMask.end(), [](bool b){ return b; }))
  return ...;

2866

if (!keepMask.hasValue())

Seems like, in the limit, this could be defined on empty shapes too.
Returning llvm::None is a clear error failure that reducedShape cannot be obtained from originalShape by dropping only 1's.

[limo1996]

Next wave of comments addressed

[nicolasvasilache]

Thanks @limo1996 !

mlir/include/mlir/Dialect/StandardOps/IR/Ops.h

public:
 
   static ParseResult parse(OpAsmParser &parser, OperationState &result);
   void print(OpAsmPrinter &p);
   LogicalResult fold(ArrayRef<Attribute> cstOperands,
                      SmallVectorImpl<OpFoldResult> &results);
   LogicalResult verify();
 };
 
+/// Extracts int64_t from the assumedArrayAttr of IntegerAttr.
+SmallVector<int64_t, 4> extractFromI64ArrayAttr(Attribute attr);
+
 /// Prints dimension and symbol list.
 void printDimAndSymbolList(Operation::operand_iterator begin,
                            Operation::operand_iterator end, unsigned numDims,
                            OpAsmPrinter &p);
 
 /// Parses dimension and symbol list and returns true if parsing failed.
 ParseResult parseDimAndSymbolList(OpAsmParser &parser,
                                   SmallVectorImpl<Value> &operands,
                                   unsigned &numDims);

mlir/integration_test/Dialect/Linalg/CPU/matmul-vs-matvec.mlir

+// RUN: mlir-opt %s -convert-linalg-to-loops -convert-linalg-to-llvm | \
+// RUN: mlir-cpu-runner -O3 -e main -entry-point-result=void \
+// RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
+// RUN: | FileCheck %s
+
+func @print_memref_f32(memref<*xf32>)
+
+func @matmul(%A: memref<?x?xf32>, %B: memref<?x?xf32>) -> (memref<?x?xf32>) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %f0 = constant 0.0 : f32
+  %x = dim %A, %c0 : memref<?x?xf32>
+  %y = dim %B, %c1 : memref<?x?xf32>
+  %C = alloc(%x, %y) : memref<?x?xf32>
+  linalg.fill(%C, %f0) : memref<?x?xf32>, f32
+  linalg.matmul ins(%A, %B: memref<?x?xf32>, memref<?x?xf32>)
+                outs(%C: memref<?x?xf32>)
+  return %C : memref<?x?xf32>
+}
+
+func @matvec(%A: memref<?x?xf32>, %B: memref<?x?xf32>) -> (memref<?x?xf32>) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %f0 = constant 0.0 : f32
+  %m = dim %A, %c0 : memref<?x?xf32>
+  %x = dim %A, %c1 : memref<?x?xf32>
+  %n = dim %B, %c1 : memref<?x?xf32>
+  %C = alloc(%m, %n) : memref<?x?xf32>
+  linalg.fill(%C, %f0) : memref<?x?xf32>, f32
+  scf.for %i = %c0 to %n step %c1 {
+    %b = subview %B[0, %i][%x, 1][1, 1] : memref<?x?xf32> to memref<?xf32, offset: ?, strides: [?]>
+    %c = subview %C[0, %i][%m, 1][1, 1] : memref<?x?xf32> to memref<?xf32, offset: ?, strides: [?]>
+    linalg.matvec ins(%A, %b: memref<?x?xf32>, memref<?xf32, offset: ?, strides: [?]>)
+                  outs(%c: memref<?xf32, offset: ?, strides: [?]>)
+  }
+  return %C : memref<?x?xf32>
+}
+
+func @main() {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %m = constant 5 : index
+  %x = constant 3 : index
+  %n = constant 2 : index
+  %val1 = constant 13.0 : f32
+  %val2 = constant 17.0 : f32
+  %A = alloc(%m, %x) : memref<?x?xf32>
+  %B = alloc(%x, %n) : memref<?x?xf32>
+  linalg.fill(%A, %val1) : memref<?x?xf32>, f32
+  linalg.fill(%B, %val2) : memref<?x?xf32>, f32
+  store %val1, %B[%c0, %c0] : memref<?x?xf32>
+  %C1 = call @matmul(%A, %B) : (memref<?x?xf32>, memref<?x?xf32>) -> memref<?x?xf32>
+  %C2 = call @matvec(%A, %B) : (memref<?x?xf32>, memref<?x?xf32>) -> memref<?x?xf32>
+  scf.for %i = %c0 to %m step %c1 {
+    scf.for %j = %c0 to %n step %c1 {
+      %e1 = load %C1[%i, %j] : memref<?x?xf32>
+      %e2 = load %C2[%i, %j] : memref<?x?xf32>
+      %c = cmpf "oeq", %e1, %e2 : f32
+      assert %c, "Matmul does not produce same output as matvec"
+    }
+  }
+  %C2_ = memref_cast %C2 : memref<?x?xf32> to memref<*xf32>
+  call @print_memref_f32(%C2_) : (memref<*xf32>) -> ()
+  return
+}
+
+// CHECK: Unranked Memref base@ = {{.*}} rank = 2 offset = 0 sizes = [5, 2] strides = [2, 1] data =
+// CHECK-NEXT:      [
+// CHECK-SAME:  [611,   663],
+// CHECK-NEXT:  [611,   663],
+// CHECK-NEXT:  [611,   663],
+// CHECK-NEXT:  [611,   663],
+// CHECK-NEXT:  [611,   663]
+// CHECK-SAME: ]

mlir/integration_test/Dialect/Linalg/CPU/rank-reducing-subview.mlir

+// RUN: mlir-opt %s -convert-linalg-to-loops -convert-linalg-to-llvm | \
+// RUN: mlir-cpu-runner -O3 -e main -entry-point-result=void \
+// RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
+// RUN: | FileCheck %s
+
+func @print_memref_f32(memref<*xf32>)
+
+func @main() {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %f0 = constant 0.0 : f32
+  %f1 = constant 1.0 : f32
+  %f2 = constant 2.0 : f32
+  %f3 = constant 3.0 : f32
+  %A = alloc(%c2, %c2) : memref<?x?xf32>
+  store %f0, %A[%c0, %c0] : memref<?x?xf32>
+  store %f1, %A[%c0, %c1] : memref<?x?xf32>
+  store %f2, %A[%c1, %c0] : memref<?x?xf32>
+  store %f3, %A[%c1, %c1] : memref<?x?xf32>
+  %B = subview %A[%c1, 0][1, %c2][1, 1] : memref<?x?xf32> to memref<?xf32, offset: ?, strides: [1]>
+  %C = subview %A[0, %c1][%c2, 1][1, 1] : memref<?x?xf32> to memref<?xf32, offset: ?, strides: [?]>
+  %A_ = memref_cast %A : memref<?x?xf32> to memref<*xf32>
+  call @print_memref_f32(%A_) : (memref<*xf32>) -> ()
+  %B_ = memref_cast %B : memref<?xf32, offset: ?, strides: [1]> to memref<*xf32>
+  call @print_memref_f32(%B_) : (memref<*xf32>) -> ()
+  %C_ = memref_cast %C : memref<?xf32, offset: ?, strides: [?]> to memref<*xf32>
+  call @print_memref_f32(%C_) : (memref<*xf32>) -> ()
+  return
+}
+
+// CHECK: Unranked Memref base@ = {{.*}} rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data =
+// CHECK-NEXT:      [
+// CHECK-SAME:  [0,   1],
+// CHECK-NEXT:  [2,   3]
+// CHECK-SAME: ]
+// CHECK: [2,  3]
+// CHECK: [1,  3]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

matchAndRewrite(Operation *op, ArrayRef<Value> operands,
     auto subViewOp = cast<SubViewOp>(op);
 
     auto sourceMemRefType = subViewOp.source().getType().cast<MemRefType>();
     auto sourceElementTy =
         typeConverter.convertType(sourceMemRefType.getElementType())
             .dyn_cast_or_null<LLVM::LLVMType>();
 
     auto viewMemRefType = subViewOp.getType();
+    auto inferredType = SubViewOp::inferResultType(
+                            subViewOp.getSourceType(),
+                            extractFromI64ArrayAttr(subViewOp.static_offsets()),
+                            extractFromI64ArrayAttr(subViewOp.static_sizes()),
+                            extractFromI64ArrayAttr(subViewOp.static_strides()))
+                            .cast<MemRefType>();
     auto targetElementTy =
         typeConverter.convertType(viewMemRefType.getElementType())
             .dyn_cast<LLVM::LLVMType>();
     auto targetDescTy = typeConverter.convertType(viewMemRefType)
                             .dyn_cast_or_null<LLVM::LLVMType>();
     if (!sourceElementTy || !targetDescTy)
       return failure();
 
     // Extract the offset and strides from the type.
     int64_t offset;
     SmallVector<int64_t, 4> strides;
-    auto successStrides = getStridesAndOffset(viewMemRefType, strides, offset);
+    auto successStrides = getStridesAndOffset(inferredType, strides, offset);
     if (failed(successStrides))
       return failure();
 
     // Create the descriptor.
     if (!operands.front().getType().isa<LLVM::LLVMType>())
       return failure();
     MemRefDescriptor sourceMemRef(operands.front());
     auto targetMemRef = MemRefDescriptor::undef(rewriter, loc, targetDescTy);
 
     // Copy the buffer pointer from the old descriptor to the new one.
     Value extracted = sourceMemRef.allocatedPtr(rewriter, loc);
     Value bitcastPtr = rewriter.create<LLVM::BitcastOp>(
         loc, targetElementTy.getPointerTo(viewMemRefType.getMemorySpace()),
         extracted);
     targetMemRef.setAllocatedPtr(rewriter, loc, bitcastPtr);
 
     // Copy the buffer pointer from the old descriptor to the new one.
     extracted = sourceMemRef.alignedPtr(rewriter, loc);
     bitcastPtr = rewriter.create<LLVM::BitcastOp>(
         loc, targetElementTy.getPointerTo(viewMemRefType.getMemorySpace()),
         extracted);
     targetMemRef.setAlignedPtr(rewriter, loc, bitcastPtr);
 
+    auto shape = viewMemRefType.getShape();
+    auto inferredShape = inferredType.getShape();
+    size_t inferredShapeRank = inferredShape.size();
+    size_t resultShapeRank = shape.size();
+    SmallVector<bool, 4> mask(inferredShapeRank);
+    for (unsigned i = 0, j = 0; i < inferredShapeRank && j < resultShapeRank;
+         ++i)
+      if ((mask[i] = inferredShape[i] == shape[j]))
+        j++;
+
     // Extract strides needed to compute offset.
     SmallVector<Value, 4> strideValues;
-    strideValues.reserve(viewMemRefType.getRank());
-    for (int i = 0, e = viewMemRefType.getRank(); i < e; ++i)
+    strideValues.reserve(inferredShapeRank);
+    for (unsigned i = 0; i < inferredShapeRank; ++i)
       strideValues.push_back(sourceMemRef.stride(rewriter, loc, i));
 
     // Offset.
     auto llvmIndexType = typeConverter.convertType(rewriter.getIndexType());
     if (!ShapedType::isDynamicStrideOrOffset(offset)) {
       targetMemRef.setConstantOffset(rewriter, loc, offset);
     } else {
       Value baseOffset = sourceMemRef.offset(rewriter, loc);
-      for (unsigned i = 0, e = viewMemRefType.getRank(); i < e; ++i) {
+      for (unsigned i = 0; i < inferredShapeRank; ++i) {
         Value offset =
             subViewOp.isDynamicOffset(i)
                 ? operands[subViewOp.getIndexOfDynamicOffset(i)]
                 : rewriter.create<LLVM::ConstantOp>(
                       loc, llvmIndexType,
                       rewriter.getI64IntegerAttr(subViewOp.getStaticOffset(i)));
         Value mul = rewriter.create<LLVM::MulOp>(loc, offset, strideValues[i]);
         baseOffset = rewriter.create<LLVM::AddOp>(loc, baseOffset, mul);
       }
       targetMemRef.setOffset(rewriter, loc, baseOffset);
     }
 
     // Update sizes and strides.
-    for (int i = viewMemRefType.getRank() - 1; i >= 0; --i) {
+    for (int i = inferredShapeRank - 1, j = resultShapeRank - 1;
+         i >= 0 && j >= 0; --i) {
+      if (!mask[i])
+        continue;
+
       Value size =
           subViewOp.isDynamicSize(i)
               ? operands[subViewOp.getIndexOfDynamicSize(i)]
               : rewriter.create<LLVM::ConstantOp>(
                     loc, llvmIndexType,
                     rewriter.getI64IntegerAttr(subViewOp.getStaticSize(i)));
-      targetMemRef.setSize(rewriter, loc, i, size);
+      targetMemRef.setSize(rewriter, loc, j, size);
       Value stride;
       if (!ShapedType::isDynamicStrideOrOffset(strides[i])) {
         stride = rewriter.create<LLVM::ConstantOp>(
             loc, llvmIndexType, rewriter.getI64IntegerAttr(strides[i]));
       } else {
         stride =
             subViewOp.isDynamicStride(i)
                 ? operands[subViewOp.getIndexOfDynamicStride(i)]
                 : rewriter.create<LLVM::ConstantOp>(
                       loc, llvmIndexType,
                       rewriter.getI64IntegerAttr(subViewOp.getStaticStride(i)));
         stride = rewriter.create<LLVM::MulOp>(loc, stride, strideValues[i]);
       }
-      targetMemRef.setStride(rewriter, loc, i, stride);
+      targetMemRef.setStride(rewriter, loc, j, stride);
+      j--;
     }
 
     rewriter.replaceOp(op, {targetMemRef});
     return success();
   }
 };
 
 /// Conversion pattern that transforms a transpose op into:

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

unsigned expectedNumDynamicEntries =
         return isDynamic(attr.cast<IntegerAttr>().getInt());
       });
   if (values.size() != expectedNumDynamicEntries)
     return op.emitError("expected ")
            << expectedNumDynamicEntries << " dynamic " << name << " values";
   return success();
 }
 
-/// Helper function extracts int64_t from the assumedArrayAttr of IntegerAttr.
-static SmallVector<int64_t, 4> extractFromI64ArrayAttr(Attribute attr) {
+SmallVector<int64_t, 4> mlir::extractFromI64ArrayAttr(Attribute attr) {
   return llvm::to_vector<4>(
       llvm::map_range(attr.cast<ArrayAttr>(), [](Attribute a) -> int64_t {
         return a.cast<IntegerAttr>().getInt();
       }));
 }
 
 /// Checks if `original` Type type can be rank reduced to `reduced` type.
 /// This function is slight variant of `is subsequence` algorithm where
 /// not matching dimension must be 1.
 static bool isRankReducedType(Type originalType, Type reducedType) {
   if (originalType == reducedType)
     return true;
   if (!originalType.isa<RankedTensorType>() && !originalType.isa<MemRefType>())
     return true;
   if (originalType.isa<RankedTensorType>() &&
       !reducedType.isa<RankedTensorType>())
     return true;
   if (originalType.isa<MemRefType>() && !reducedType.isa<MemRefType>())
     return true;
 
   ShapedType originalShapedType = originalType.cast<ShapedType>();
   ShapedType reducedShapedType = reducedType.cast<ShapedType>();
 
   // Rank and size logic is valid for all ShapedTypes.
   ArrayRef<int64_t> originalShape = originalShapedType.getShape();
for (unsigned originalIdx = 0; originalIdx < originalRank; ++originalIdx) {
     // Skip matching dims greedily.
     if ((keepMask[originalIdx] = originalShape[originalIdx] == reducedVal))
       reducedIdx++;
     // 1 is the only non-matching allowed.
     else if (originalShape[originalIdx] != 1)
       return false;
   }
   // Must match the reduced rank.
   if (reducedIdx != reducedRank)
     return false;
 
   // We are done for the tensor case.
   if (originalType.isa<RankedTensorType>())
     return true;
 
   // Strided layout logic is relevant for MemRefType only.
   MemRefType original = originalType.cast<MemRefType>();
   MemRefType reduced = reducedType.cast<MemRefType>();
   MLIRContext *c = original.getContext();
   int64_t originalOffset, reducedOffset;