fixup! [mlir][Vector] Update patterns for flattening vector.xfer Ops (2/N)

Refactor to use makeComposedFoldedAffineApply

Author: banach-space

mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp

@@ -544,7 +544,7 @@ class FlattenContiguousRowMajorTransferReadPattern
     auto loc = transferReadOp.getLoc();
     Value vector = transferReadOp.getVector();
     VectorType vectorType = cast<VectorType>(vector.getType());
-    Value source = transferReadOp.getSource();
+    auto source = transferReadOp.getSource();
     MemRefType sourceType = dyn_cast<MemRefType>(source.getType());
 
     // 0. Check pre-conditions
@@ -602,26 +602,30 @@ class FlattenContiguousRowMajorTransferReadPattern
       //
       // For this example:
       //   %2 = vector.transfer_read %arg4[%c0, %arg0, %c0] (...) :
-      //   memref<1x43x2xi32>, vector<1x2xi32>
+      //      memref<1x43x2xi32>, vector<1x2xi32>
       // which would be collapsed to:
       //   %1 = vector.transfer_read %collapse_shape[%c0, %offset] (...) :
-      //   memref<1x86xi32>, vector<2xi32>
+      //      memref<1x86xi32>, vector<2xi32>
       // one would get the following offset:
       //    %offset = %arg0 * 43
+      AffineExpr offsetE, idx;
+      bindSymbols(rewriter.getContext(), offsetE, idx);
+
       int64_t outputRank = transferReadOp.getIndices().size();
-      Value offset = rewriter.create<arith::ConstantIndexOp>(loc, 0);
+      OpFoldResult offset =
+          rewriter.create<arith::ConstantIndexOp>(loc, 0).getResult();
       for (int64_t i = firstDimToCollapse; i < outputRank; ++i) {
-        Value dimIdx = rewriter.create<arith::ConstantIndexOp>(loc, i);
-        auto sourceDimSize =
-            rewriter.create<memref::DimOp>(loc, source, dimIdx);
-
-        offset = rewriter.create<arith::AddIOp>(
-            loc,
-            rewriter.create<arith::MulIOp>(loc, transferReadOp.getIndices()[i],
-                                           sourceDimSize),
-            offset);
+        int64_t dim = dyn_cast<ShapedType>(source.getType()).getDimSize(i);
+        offset = affine::makeComposedFoldedAffineApply(
+            rewriter, loc, offsetE + dim * idx,
+            {offset, transferReadOp.getIndices()[i]});
+      }
+      if (offset.is<Value>()) {
+        collapsedIndices.push_back(offset.get<Value>());
+      } else {
+        collapsedIndices.push_back(rewriter.create<arith::ConstantIndexOp>(
+            loc, *getConstantIntValue(offset)));
       }
-      collapsedIndices.push_back(offset);
     }
 
     // 3. Create new vector.transfer_read that reads from the collapsed memref

mlir/test/Dialect/Vector/vector-transfer-flatten.mlir

@@ -55,21 +55,19 @@ func.func @transfer_read_dims_mismatch_non_zero_indices(
   return
 }
 
+// CHECK: #[[$ATTR_0:.+]] = affine_map<()[s0, s1] -> (s0 * 4 + s1 * 43)>
+
 // CHECK-LABEL:   func.func @transfer_read_dims_mismatch_non_zero_indices(
 // CHECK-SAME:      %[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index,
 // CHECK-SAME:      %[[VAL_2:.*]]: memref<1x43x4x6xi32>,
 // CHECK-SAME:      %[[VAL_3:.*]]: memref<1x2x6xi32>) {
-// CHECK:           %[[VAL_4:.*]] = arith.constant 43 : index
-// CHECK:           %[[VAL_5:.*]] = arith.constant 4 : index
-// CHECK:           %[[VAL_6:.*]] = arith.constant 0 : i32
-// CHECK:           %[[VAL_7:.*]] = arith.constant 0 : index
-// CHECK:           %[[VAL_8:.*]] = memref.collapse_shape %[[VAL_2]] {{\[\[}}0], [1, 2, 3]] : memref<1x43x4x6xi32> into memref<1x1032xi32>
-// CHECK:           %[[VAL_9:.*]] = arith.muli %[[VAL_0]], %[[VAL_4]] : index
-// CHECK:           %[[VAL_10:.*]] = arith.muli %[[VAL_1]], %[[VAL_5]] : index
-// CHECK:           %[[VAL_11:.*]] = arith.addi %[[VAL_10]], %[[VAL_9]] : index
-// CHECK:           %[[VAL_12:.*]] = vector.transfer_read %[[VAL_8]]{{\[}}%[[VAL_7]], %[[VAL_11]]], %[[VAL_6]] {in_bounds = [true]} : memref<1x1032xi32>, vector<12xi32>
-// CHECK:           %[[VAL_13:.*]] = memref.collapse_shape %[[VAL_3]] {{\[\[}}0, 1, 2]] : memref<1x2x6xi32> into memref<12xi32>
-// CHECK:           vector.transfer_write %[[VAL_12]], %[[VAL_13]]{{\[}}%[[VAL_7]]] {in_bounds = [true]} : vector<12xi32>, memref<12xi32>
+// CHECK:           %[[VAL_4:.*]] = arith.constant 0 : i32
+// CHECK:           %[[VAL_5:.*]] = arith.constant 0 : index
+// CHECK:           %[[VAL_6:.*]] = memref.collapse_shape %[[VAL_2]] {{\[\[}}0], [1, 2, 3]] : memref<1x43x4x6xi32> into memref<1x1032xi32>
+// CHECK:           %[[VAL_7:.*]] = affine.apply #[[$ATTR_0]](){{\[}}%[[VAL_1]], %[[VAL_0]]]
+// CHECK:           %[[VAL_8:.*]] = vector.transfer_read %[[VAL_6]]{{\[}}%[[VAL_5]], %[[VAL_7]]], %[[VAL_4]] {in_bounds = [true]} : memref<1x1032xi32>, vector<12xi32>
+// CHECK:           %[[VAL_9:.*]] = memref.collapse_shape %[[VAL_3]] {{\[\[}}0, 1, 2]] : memref<1x2x6xi32> into memref<12xi32>
+// CHECK:           vector.transfer_write %[[VAL_8]], %[[VAL_9]]{{\[}}%[[VAL_5]]] {in_bounds = [true]} : vector<12xi32>, memref<12xi32>
 
 // -----
 

mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp

@@ -454,6 +454,7 @@ struct TestFlattenVectorTransferPatterns
   }
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<memref::MemRefDialect>();
+    registry.insert<affine::AffineDialect>();
   }
   void runOnOperation() override {
     RewritePatternSet patterns(&getContext());