[mlir][MemRef] Fix the simplification of extract_strided_metadata(subview)

[qcolombet]

Summary

Prior to this patch we were wrongly applying the sub-strides to the computation of the final offset of the subview.

Put differently, we were computing the offset as:

offset = baseOffset + sum(subOffset#i * baseStrides#i * subSizes#i)

Whereas we should be doing:

offset = baseOffset + sum(subOffset#i * baseStrides#i)

I.e., drop the subSizes#i term from the sum.

[ftynse]

Thanks! The affine now also looks similar to what we have elsewhere for strided formats.

mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp

 ///              extract_strided_metadata(subview(memref, subOffset,
 ///                                               subSizes, subStrides))`
 /// With
 ///
 /// \verbatim
 /// baseBuffer, baseOffset, baseSizes, baseStrides =
 ///     extract_strided_metadata(memref)
 /// strides#i = baseStrides#i * subSizes#i
-/// offset = baseOffset + sum(subOffset#i * strides#i)
+/// offset = baseOffset + sum(subOffset#i * baseStrides#i)
 /// sizes = subSizes
 /// \endverbatim
 ///
 /// In other words, get rid of the subview in that expression and canonicalize
 /// on its effects on the offset, the sizes, and the strides using affine.apply.
 struct ExtractStridedMetadataOpSubviewFolder
     : public OpRewritePattern<memref::ExtractStridedMetadataOp> {
 public:
LogicalResult matchAndRewrite(memref::ExtractStridedMetadataOp op,
     // Compute the new strides and offset from the base strides and offset:
     // newStride#i = baseStride#i * subStride#i
     // offset = baseOffset + sum(subOffsets#i * newStrides#i)
     SmallVector<OpFoldResult> strides;
     SmallVector<OpFoldResult> subStrides = subview.getMixedStrides();
     auto origStrides = newExtractStridedMetadata.getStrides();
 
     // Hold the affine symbols and values for the computation of the offset.
-    SmallVector<OpFoldResult> values(3 * sourceRank + 1);
-    SmallVector<AffineExpr> symbols(3 * sourceRank + 1);
+    SmallVector<OpFoldResult> values(2 * sourceRank + 1);
+    SmallVector<AffineExpr> symbols(2 * sourceRank + 1);
 
     detail::bindSymbolsList(rewriter.getContext(), symbols);
     AffineExpr expr = symbols.front();
     values[0] = ShapedType::isDynamicStrideOrOffset(sourceOffset)
                     ? getAsOpFoldResult(newExtractStridedMetadata.getOffset())
                     : rewriter.getIndexAttr(sourceOffset);
     SmallVector<OpFoldResult> subOffsets = subview.getMixedOffsets();
 
     AffineExpr s0 = rewriter.getAffineSymbolExpr(0);
     AffineExpr s1 = rewriter.getAffineSymbolExpr(1);
     for (unsigned i = 0; i < sourceRank; ++i) {
       // Compute the stride.
       OpFoldResult origStride =
           ShapedType::isDynamicStrideOrOffset(sourceStrides[i])
               ? origStrides[i]
               : OpFoldResult(rewriter.getIndexAttr(sourceStrides[i]));
       strides.push_back(makeComposedFoldedAffineApply(
           rewriter, origLoc, s0 * s1, {subStrides[i], origStride}));
 
       // Build up the computation of the offset.
-      unsigned baseIdxForDim = 1 + 3 * i;
+      unsigned baseIdxForDim = 1 + 2 * i;
       unsigned subOffsetForDim = baseIdxForDim;
-      unsigned subStrideForDim = baseIdxForDim + 1;
-      unsigned origStrideForDim = baseIdxForDim + 2;
-      expr = expr + symbols[subOffsetForDim] * symbols[subStrideForDim] *
-                        symbols[origStrideForDim];
+      unsigned origStrideForDim = baseIdxForDim + 1;
+      expr = expr + symbols[subOffsetForDim] * symbols[origStrideForDim];
       values[subOffsetForDim] = subOffsets[i];
-      values[subStrideForDim] = subStrides[i];
       values[origStrideForDim] = origStride;
     }
 
     // Compute the offset.
     OpFoldResult finalOffset =
         makeComposedFoldedAffineApply(rewriter, origLoc, expr, values);
 
     SmallVector<Value> results;

mlir/test/Dialect/MemRef/simplify-extract-strided-metadata.mlir

return %base_buffer, %offset, %sizes#0, %sizes#1, %strides#0, %strides#1 :
     memref<f32>, index, index, index, index, index
 }
 
 // -----
 
 // Check that we simplify extract_strided_metadata of subview to
 // base_buf, base_offset, base_sizes, base_strides = extract_strided_metadata
 // strides = base_stride_i * subview_stride_i
-// offset = base_offset + sum(subview_offsets_i * strides_i).
+// offset = base_offset + sum(subview_offsets_i * base_strides_i).
 //
 // This test also checks that we don't create useless arith operations
 // when subview_offsets_i is 0.
 //
 // CHECK-LABEL: func @extract_strided_metadata_of_subview
 //  CHECK-SAME: (%[[ARG:.*]]: memref<5x4xf32>)
 //
 // Materialize the offset for dimension 1.
 //   CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
 //   CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
 //   CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index
 //
 // Plain extract_strided_metadata.
 //   CHECK-DAG: %[[BASE:.*]], %[[OFFSET:.*]], %[[SIZES:.*]]:2, %[[STRIDES:.*]]:2 = memref.extract_strided_metadata %[[ARG]]
 //
 // Final offset is:
 //   origOffset + (== 0)
-//   base_stride0 * subview_stride0 * subview_offset0 + (== 4 * 1 * 0 == 0)
-//   base_stride1 * subview_stride1 * subview_offset1 (== 1 * 1 * 2)
+//   base_stride0 * subview_offset0 + (== 4 * 0 == 0)
+//   base_stride1 * subview_offset1 (== 1 * 2)
 //  == 2
 //
 // Return the new tuple.
 //       CHECK: return %[[BASE]], %[[C2]], %[[C2]], %[[C2]], %[[C4]], %[[C1]]
 func.func @extract_strided_metadata_of_subview(%base: memref<5x4xf32>)
     -> (memref<f32>, index, index, index, index, index) {
 
   %subview = memref.subview %base[0, 2][2, 2][1, 1] :

 //
 // Orig strides: [64, 4, 1]
 // Sub strides: [1, %stride, 1]
 // => New strides: [64, 4 * %stride, 1]
 // Final strides == filterOutReducedDim(new strides, 0) == [4 * %stride , 1]
 //
 // Orig offset: 0
 // Sub offsets: [3, 4, 2]
-// => Final offset: 3 * 64 + 4 * 4 * %stride + 2 * 1 + 0 == 16 * %stride + 194
+// => Final offset: 3 * 64 + 4 * 4 + 2 * 1 + 0 == 210
 //
 //   CHECK-DAG: #[[$STRIDE1_MAP:.*]] = affine_map<()[s0] -> (s0 * 4)>
-//   CHECK-DAG: #[[$OFFSET_MAP:.*]] = affine_map<()[s0] -> (s0 * 16 + 194)>
 // CHECK-LABEL: func @extract_strided_metadata_of_rank_reduced_subview_w_variable_strides
 //  CHECK-SAME: (%[[ARG:.*]]: memref<8x16x4xf32>,
 //  CHECK-SAME: %[[DYN_STRIDE:.*]]: index)
 //
+//   CHECK-DAG: %[[C210:.*]] = arith.constant 210 : index
 //   CHECK-DAG: %[[C6:.*]] = arith.constant 6 : index
 //   CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
 //   CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
 //
 //   CHECK-DAG: %[[BASE:.*]], %[[OFFSET:.*]], %[[SIZES:.*]]:3, %[[STRIDES:.*]]:3 = memref.extract_strided_metadata %[[ARG]]
 //
 //   CHECK-DAG: %[[DIM1_STRIDE:.*]] = affine.apply #[[$STRIDE1_MAP]]()[%[[DYN_STRIDE]]]
-//   CHECK-DAG: %[[FINAL_OFFSET:.*]] = affine.apply #[[$OFFSET_MAP]]()[%[[DYN_STRIDE]]]
 //
-//       CHECK: return %[[BASE]], %[[FINAL_OFFSET]], %[[C6]], %[[C3]], %[[DIM1_STRIDE]], %[[C1]]
+//       CHECK: return %[[BASE]], %[[C210]], %[[C6]], %[[C3]], %[[DIM1_STRIDE]], %[[C1]]
 func.func @extract_strided_metadata_of_rank_reduced_subview_w_variable_strides(
     %base: memref<8x16x4xf32>, %stride: index)
     -> (memref<f32>, index, index, index, index, index) {
 
   %subview = memref.subview %base[3, 4, 2][1, 6, 3][1, %stride, 1] :
     memref<8x16x4xf32> to memref<6x3xf32, strided<[4, 1], offset: 210>>
 
   %base_buffer, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %subview :

 //
 // Orig strides: [s0, s1, s2]
 // Sub strides: [subS0, subS1, subS2]
 // => New strides: [s0 * subS0, s1 * subS1, s2 * subS2]
 // ==> 1 affine map (used for each stride) with two values.
 //
 // Orig offset: origOff
 // Sub offsets: [subO0, subO1, subO2]
-// => Final offset: s0 * subS0 * subO0 + ... + s2 * subS2 * subO2 + origOff
-// ==> 1 affine map with (rank * 3 + 1) symbols
+// => Final offset: s0 * * subO0 + ... + s2 * subO2 + origOff
+// ==> 1 affine map with (rank * 2 + 1) symbols
 //
 // CHECK-DAG: #[[$STRIDE_MAP:.*]] = affine_map<()[s0, s1] -> (s0 * s1)>
-// CHECK-DAG: #[[$OFFSET_MAP:.*]] = affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s0 + (s1 * s2) * s3 + (s4 * s5) * s6 + (s7 * s8) * s9)>
+// CHECK-DAG: #[[$OFFSET_MAP:.*]] = affine_map<()[s0, s1, s2, s3, s4, s5, s6] -> (s0 + s1 * s2 + s3 * s4 + s5 * s6)>
 // CHECK-LABEL: func @extract_strided_metadata_of_subview_all_dynamic
 //  CHECK-SAME: (%[[ARG:.*]]: memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>, %[[DYN_OFFSET0:.*]]: index, %[[DYN_OFFSET1:.*]]: index, %[[DYN_OFFSET2:.*]]: index, %[[DYN_SIZE0:.*]]: index, %[[DYN_SIZE1:.*]]: index, %[[DYN_SIZE2:.*]]: index, %[[DYN_STRIDE0:.*]]: index, %[[DYN_STRIDE1:.*]]: index, %[[DYN_STRIDE2:.*]]: index)
 //
 //   CHECK-DAG: %[[BASE:.*]], %[[OFFSET:.*]], %[[SIZES:.*]]:3, %[[STRIDES:.*]]:3 = memref.extract_strided_metadata %[[ARG]]
 //
 //  CHECK-DAG: %[[FINAL_STRIDE0:.*]] = affine.apply #[[$STRIDE_MAP]]()[%[[DYN_STRIDE0]], %[[STRIDES]]#0]
 //  CHECK-DAG: %[[FINAL_STRIDE1:.*]] = affine.apply #[[$STRIDE_MAP]]()[%[[DYN_STRIDE1]], %[[STRIDES]]#1]
 //  CHECK-DAG: %[[FINAL_STRIDE2:.*]] = affine.apply #[[$STRIDE_MAP]]()[%[[DYN_STRIDE2]], %[[STRIDES]]#2]
 //
-//  CHECK-DAG: %[[FINAL_OFFSET:.*]] = affine.apply #[[$OFFSET_MAP]]()[%[[OFFSET]], %[[DYN_OFFSET0]], %[[DYN_STRIDE0]], %[[STRIDES]]#0, %[[DYN_OFFSET1]], %[[DYN_STRIDE1]], %[[STRIDES]]#1, %[[DYN_OFFSET2]], %[[DYN_STRIDE2]], %[[STRIDES]]#2]
+//  CHECK-DAG: %[[FINAL_OFFSET:.*]] = affine.apply #[[$OFFSET_MAP]]()[%[[OFFSET]], %[[DYN_OFFSET0]], %[[STRIDES]]#0, %[[DYN_OFFSET1]], %[[STRIDES]]#1, %[[DYN_OFFSET2]], %[[STRIDES]]#2]
 //
 //       CHECK: return %[[BASE]], %[[FINAL_OFFSET]], %[[DYN_SIZE0]], %[[DYN_SIZE1]], %[[DYN_SIZE2]], %[[FINAL_STRIDE0]], %[[FINAL_STRIDE1]], %[[FINAL_STRIDE2]]
 func.func @extract_strided_metadata_of_subview_all_dynamic(
     %base: memref<?x?x?xf32, strided<[?,?,?], offset:?>>,
     %offset0: index, %offset1: index, %offset2: index,
     %size0: index, %size1: index, %size2: index,
     %stride0: index, %stride1: index, %stride2: index)
     -> (memref<f32>, index, index, index, index, index, index, index) {