[mlir][vector] Improve vector distribute integration test and fix block distribution

[ThomasRaoux]

Summary

Fix semantic in the distribute integration test based on offline feedback. This exposed a bug in block distribution, we need to make sure the id is multiplied by the stride of the vector. Fix the transformation and unit test.

[aartbik]

mlir/integration_test/Dialect/Vector/CPU/test-vector-distribute.mlir

22–26

Not in this CL, but probably during a later cleanup, I would rename the "distribution" part to something better. Loop distribution is typically reserved for

for { for

s1            ->        s1
s2                   for

} s2

what is done here is more stripmining, blocking or tiling or chunking in 1-D, or something named like that.

37–43

c32 because 2x32=64, right?

I am not super convinced I find this intermediate step much easier to understand than generating the chunked loop right away, but I hope you convince me in the discourse discussion

[mehdi_amini]

mlir/test/Dialect/Vector/vector-distribution.mlir

3

Why removing the "-LABEL" here?

[ThomasRaoux]

Changes to preserve semantic in the integration test based on discourse:
https://llvm.discourse.group/t/vector-vector-distribution-large-vector-to-small-vector/1983/12

[ThomasRaoux]

I changed the integration test to match what I described here: https://llvm.discourse.group/t/vector-vector-distribution-large-vector-to-small-vector/1983/12

If it makes sense to all of you I can do the rest of the changes.

mlir/integration_test/Dialect/Vector/CPU/test-vector-distribute.mlir

22–26

Makes sense. I can rename it in a future patch once we get more agreement on the design,

37–43

Correct, right now the extract_map expects contiguous IDs. (%arg5 : 32 goes from 0 to 31).
About iterative vs all at once transformation let's keep talking on Discourse :)

mlir/test/Dialect/Vector/vector-distribution.mlir

3

That was done by mistake. I do need to change the CHECK-LABEL in the 3rd test otherwise the [[MAP]] variable gets reset after the CHECK-LABEL.

[nicolasvasilache]

Let's also update the ops' documentation to remove this section:

For instance, the following code:

   %a = vector.transfer_read %A[%c0]: memref<32xf32>, vector<32xf32>
   %b = vector.transfer_read %B[%c0]: memref<32xf32>, vector<32xf32>
   %c = addf %a, %b: vector<32xf32>
   vector.transfer_write %c, %C[%c0]: memref<32xf32>, vector<32xf32>

   can be rewritten to:
   %a = vector.transfer_read %A[%c0]: memref<32xf32>, vector<32xf32>
   %b = vector.transfer_read %B[%c0]: memref<32xf32>, vector<32xf32>
   %ea = vector.extract_map %a[%id : 32] : vector<32xf32> to vector<1xf32>
   %eb = vector.extract_map %b[%id : 32] : vector<32xf32> to vector<1xf32>
   %ec = addf %ea, %eb : vector<1xf32>
   %c = vector.insert_map %ec, %id, 32 : vector<1xf32> to vector<32xf32>
   vector.transfer_write %c, %C[%c0]: memref<32xf32>, vector<32xf32>

   Where %id can be an induction variable or an SPMD id going from 0 to 31.

   And then be rewritten to:
   %a = vector.transfer_read %A[%id]: memref<32xf32>, vector<1xf32>
   %b = vector.transfer_read %B[%id]: memref<32xf32>, vector<1xf32>
   %c = addf %a, %b: vector<1xf32>
   vector.transfer_write %c, %C[%id]: memref<32xf32>, vector<1xf32>

and instead focus on the actual op semantics, with an example e.g.

Examples:

%idx0 = ... : index // dynamic computation producing the value 0 of index type
%idx1 = ... : index // dynamic computation producing the value 1 of index type
%0 = constant dense<0, 1, 2, 3>: vector<4xi32>
%1 = vector.extract_map %0[%idx0 : 2] : vector<4xi32> to vector<2xi32> // extracts values [0, 1]
%2 = vector.extract_map %0[%idx1 : 2] : vector<4xi32> to vector<2xi32> // extracts values [1, 2]
... (same for insert variant)

As the op semantics evolve, the examples will too.

Thanks @ThomasRaoux !

mlir/integration_test/Dialect/Vector/CPU/test-vector-distribute.mlir

37–43

As discussed on discourse, this is transient state internal to the test pass and should not be exposed.

Let's please have the test do 2 things.

Input IR:

%a = vector.transfer_read %in1[%c0], %cf0: memref<?xf32>, vector<256xf32>
%b = vector.transfer_read %in2[%c0], %cf0: memref<?xf32>, vector<256xf32>
%acc = addf %a, %b: vector<256xf32>
vector.transfer_write %acc, %out[%c0]: vector<256xf32>, memref<?xf32>

Output IR:

scf.for %arg5 = %c0 to %c256 step %c8 {
  %a = vector.transfer_read %in1[%arg5], %cf0: memref<?xf32>, vector<8xf32>
  %b = vector.transfer_read %in2[%arg5], %cf0: memref<?xf32>, vector<8xf32>
  %acc = addf %a, %b: vector<8xf32>
  vector.transfer_write %acc, %out[%arg5]: vector<8xf32>, memref<?xf32>
}

The test should also run with and without the application of the test pass and produce the same result.

[ThomasRaoux]

Update integration test to start from a large vector without loop.

[ThomasRaoux]

I re-based the patch and changed the test to avoid starting from an intermediate state. Please take another look.

mlir/integration_test/Dialect/Vector/CPU/test-vector-distribute.mlir

37–43

Done. Starting from just the vector add and running with and without the transformation pass.

[nicolasvasilache]

mlir/integration_test/Dialect/Vector/CPU/test-vector-distribute.mlir

1–2

Could we add an extra RUN command that just does mlir-opt %s -test-vector-to-forloop and checks the presence of the forloop+vectors ?

6

blank line to delimit commands ?

[ThomasRaoux]

Add extra check to the integration test to make sure the transformation happened.

mlir/integration_test/Dialect/Vector/CPU/test-vector-distribute.mlir

-// RUN: mlir-opt %s -test-vector-distribute-patterns=distribution-multiplicity=32 \
-// RUN:  -convert-vector-to-scf -lower-affine -convert-scf-to-std -convert-vector-to-llvm | \
+// RUN: mlir-opt %s -test-vector-distribute-patterns -convert-vector-to-scf \
+// RUN:   -lower-affine -convert-scf-to-std -convert-vector-to-llvm | \
 // RUN: mlir-cpu-runner -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 @alloc_1d_filled_inc_f32(%arg0: index, %arg1: f32) -> memref<?xf32> {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %0 = alloc(%arg0) : memref<?xf32>
   scf.for %arg2 = %c0 to %arg0 step %c1 {
     %tmp = index_cast %arg2 : index to i32
     %tmp1 = sitofp %tmp : i32 to f32
     %tmp2 = addf %tmp1, %arg1 : f32
     store %tmp2, %0[%arg2] : memref<?xf32>
   }
   return %0 : memref<?xf32>
 }
 
-func @vector_add_cycle(%id : index, %A: memref<?xf32>, %B: memref<?xf32>, %C: memref<?xf32>) {
-  %c0 = constant 0 : index
-  %cf0 = constant 0.0 : f32
-  %a = vector.transfer_read %A[%c0], %cf0: memref<?xf32>, vector<64xf32>
-  %b = vector.transfer_read %B[%c0], %cf0: memref<?xf32>, vector<64xf32>
-  %acc = addf %a, %b: vector<64xf32>
-  vector.transfer_write %acc, %C[%c0]: vector<64xf32>, memref<?xf32>
-  return
-}
+#map0 = affine_map<()[s0] -> (s0 * 2)>
 
-// Loop over a function containinng a large add vector and distribute it so that
-// each iteration of the loop process part of the vector operation.
+// Loop over a vector add being distributed into a loop of vec2 and make sure
+// distribution is being propagated.
 func @main() {
+  %cf0 = constant 0.0 : f32
   %cf1 = constant 1.0 : f32
   %cf2 = constant 2.0 : f32
   %c0 = constant 0 : index
   %c1 = constant 1 : index
+  %c32 = constant 32 : index
   %c64 = constant 64 : index
   %out = alloc(%c64) : memref<?xf32>
   %in1 = call @alloc_1d_filled_inc_f32(%c64, %cf1) : (index, f32) -> memref<?xf32>
   %in2 = call @alloc_1d_filled_inc_f32(%c64, %cf2) : (index, f32) -> memref<?xf32>
-  scf.for %arg5 = %c0 to %c64 step %c1 {
-    call @vector_add_cycle(%arg5, %in1, %in2, %out) : (index, memref<?xf32>, memref<?xf32>, memref<?xf32>) -> ()
+  scf.for %arg5 = %c0 to %c32 step %c1 {
+    %a = vector.transfer_read %in1[%c0], %cf0: memref<?xf32>, vector<64xf32>
+    %b = vector.transfer_read %in2[%c0], %cf0: memref<?xf32>, vector<64xf32>
+    %acc = addf %a, %b: vector<64xf32>
+    %ext = vector.extract_map %acc[%arg5 : 32] : vector<64xf32> to vector<2xf32>
+    %id = affine.apply #map0()[%arg5]
+    vector.transfer_write %ext, %out[%id]: vector<2xf32>, memref<?xf32>
   }
   %converted = memref_cast %out : memref<?xf32> to memref<*xf32>
   call @print_memref_f32(%converted): (memref<*xf32>) -> ()
   // CHECK:      Unranked{{.*}}data =
   // CHECK:      [
   // CHECK-SAME:  3,  5,  7,  9,  11,  13,  15,  17,  19,  21,  23,  25,  27,
   // CHECK-SAME:  29,  31,  33,  35,  37,  39,  41,  43,  45,  47,  49,  51,
   // CHECK-SAME:  53,  55,  57,  59,  61,  63,  65,  67,  69,  71,  73,  75,
   // CHECK-SAME:  77,  79,  81,  83,  85,  87,  89,  91,  93,  95,  97,  99,
   // CHECK-SAME:  101,  103,  105,  107,  109,  111,  113,  115,  117,  119,
   // CHECK-SAME:  121,  123,  125,  127,  129]
   dealloc %out : memref<?xf32>
   dealloc %in1 : memref<?xf32>
   dealloc %in2 : memref<?xf32>
   return
 }

mlir/lib/Dialect/Vector/VectorTransforms.cpp

LogicalResult matchAndRewrite(vector::TransferReadOp read,
     if (!read.getResult().hasOneUse())
       return failure();
     auto extract =
         dyn_cast<vector::ExtractMapOp>(*read.getResult().getUsers().begin());
     if (!extract)
       return failure();
     edsc::ScopedContext scope(rewriter, read.getLoc());
     using mlir::edsc::op::operator+;
+    using mlir::edsc::op::operator*;
     using namespace mlir::edsc::intrinsics;
     SmallVector<Value, 4> indices(read.indices().begin(), read.indices().end());
-    indices.back() = indices.back() + extract.id();
+    indices.back() =
+        indices.back() +
+        (extract.id() *
+         std_constant_index(extract.getResultType().getDimSize(0)));
     Value newRead = vector_transfer_read(extract.getType(), read.memref(),
                                          indices, read.permutation_map(),
                                          read.padding(), ArrayAttr());
     newRead = rewriter.create<vector::InsertMapOp>(
         read.getLoc(), newRead, extract.id(), extract.multiplicity());
     rewriter.replaceOp(read, newRead);
     return success();
   }
 };
 
 struct TransferWriteInsertPattern
     : public OpRewritePattern<vector::TransferWriteOp> {
   TransferWriteInsertPattern(MLIRContext *context)
       : OpRewritePattern<vector::TransferWriteOp>(context) {}
   LogicalResult matchAndRewrite(vector::TransferWriteOp write,
                                 PatternRewriter &rewriter) const override {
     auto insert = write.vector().getDefiningOp<vector::InsertMapOp>();
     if (!insert)
       return failure();
     edsc::ScopedContext scope(rewriter, write.getLoc());
     using mlir::edsc::op::operator+;
+    using mlir::edsc::op::operator*;
     using namespace mlir::edsc::intrinsics;
     SmallVector<Value, 4> indices(write.indices().begin(),
                                   write.indices().end());
-    indices.back() = indices.back() + insert.id();
+    indices.back() =
+        indices.back() +
+        (insert.id() *
+         std_constant_index(insert.getSourceVectorType().getDimSize(0)));
     vector_transfer_write(insert.vector(), write.memref(), indices,
                           write.permutation_map(), ArrayAttr());
     rewriter.eraseOp(write);
     return success();
   }
 };
 
 // TODO: Add pattern to rewrite ExtractSlices(ConstantMaskOp).

mlir/test/Dialect/Vector/vector-distribution.mlir

-// RUN: mlir-opt %s -test-vector-distribute-patterns=distribution-multiplicity=32 | FileCheck %s
+// RUN: mlir-opt %s -test-vector-distribute-patterns=distribution-multiplicity=32 -split-input-file | FileCheck %s
 
 // CHECK-LABEL: func @distribute_vector_add
 //  CHECK-SAME: (%[[ID:.*]]: index
 //  CHECK-NEXT:    %[[EXA:.*]] = vector.extract_map %{{.*}}[%[[ID]] : 32] : vector<32xf32> to vector<1xf32>
 //  CHECK-NEXT:    %[[EXB:.*]] = vector.extract_map %{{.*}}[%[[ID]] : 32] : vector<32xf32> to vector<1xf32>
 //  CHECK-NEXT:    %[[ADD:.*]] = addf %[[EXA]], %[[EXB]] : vector<1xf32>
 //  CHECK-NEXT:    %[[INS:.*]] = vector.insert_map %[[ADD]], %[[ID]], 32 : vector<1xf32> to vector<32xf32>
 //  CHECK-NEXT:    return %[[INS]] : vector<32xf32>
 func @distribute_vector_add(%id : index, %A: vector<32xf32>, %B: vector<32xf32>) -> vector<32xf32> {
   %0 = addf %A, %B : vector<32xf32>
   return %0: vector<32xf32>
 }
 
+// -----
+
 // CHECK-LABEL: func @vector_add_read_write
 //  CHECK-SAME: (%[[ID:.*]]: index
 //       CHECK:    %[[EXA:.*]] = vector.transfer_read %{{.*}}[%[[ID]]], %{{.*}} : memref<32xf32>, vector<1xf32>
 //  CHECK-NEXT:    %[[EXB:.*]] = vector.transfer_read %{{.*}}[%[[ID]]], %{{.*}} : memref<32xf32>, vector<1xf32>
 //  CHECK-NEXT:    %[[ADD1:.*]] = addf %[[EXA]], %[[EXB]] : vector<1xf32>
 //  CHECK-NEXT:    %[[EXC:.*]] = vector.transfer_read %{{.*}}[%[[ID]]], %{{.*}} : memref<32xf32>, vector<1xf32>
 //  CHECK-NEXT:    %[[ADD2:.*]] = addf %[[ADD1]], %[[EXC]] : vector<1xf32>
 //  CHECK-NEXT:    vector.transfer_write %[[ADD2]], %{{.*}}[%[[ID]]] : vector<1xf32>, memref<32xf32>
 //  CHECK-NEXT:    return
 func @vector_add_read_write(%id : index, %A: memref<32xf32>, %B: memref<32xf32>, %C: memref<32xf32>, %D: memref<32xf32>) {
   %c0 = constant 0 : index
   %cf0 = constant 0.0 : f32
   %a = vector.transfer_read %A[%c0], %cf0: memref<32xf32>, vector<32xf32>
   %b = vector.transfer_read %B[%c0], %cf0: memref<32xf32>, vector<32xf32>
   %acc = addf %a, %b: vector<32xf32>
   %c = vector.transfer_read %C[%c0], %cf0: memref<32xf32>, vector<32xf32>
   %d = addf %acc, %c: vector<32xf32>
   vector.transfer_write %d, %D[%c0]: vector<32xf32>, memref<32xf32>
   return
 }
 
-// CHECK-LABEL: func @vector_add_cycle
+// -----
+
+// CHECK-DAG: #[[MAP0:map[0-9]+]] = affine_map<()[s0] -> (s0 * 2)>
+
+//       CHECK: func @vector_add_cycle
 //  CHECK-SAME: (%[[ID:.*]]: index
-//       CHECK:    %[[EXA:.*]] = vector.transfer_read %{{.*}}[%[[ID]]], %{{.*}} : memref<64xf32>, vector<2xf32>
-//  CHECK-NEXT:    %[[EXB:.*]] = vector.transfer_read %{{.*}}[%[[ID]]], %{{.*}} : memref<64xf32>, vector<2xf32>
+//       CHECK:    %[[ID1:.*]] = affine.apply #[[MAP0]]()[%[[ID]]]
+//  CHECK-NEXT:    %[[EXA:.*]] = vector.transfer_read %{{.*}}[%[[ID1]]], %{{.*}} : memref<64xf32>, vector<2xf32>
+//  CHECK-NEXT:    %[[ID2:.*]] = affine.apply #[[MAP0]]()[%[[ID]]]
+//  CHECK-NEXT:    %[[EXB:.*]] = vector.transfer_read %{{.*}}[%[[ID2]]], %{{.*}} : memref<64xf32>, vector<2xf32>
 //  CHECK-NEXT:    %[[ADD:.*]] = addf %[[EXA]], %[[EXB]] : vector<2xf32>
-//  CHECK-NEXT:    vector.transfer_write %[[ADD]], %{{.*}}[%[[ID]]] : vector<2xf32>, memref<64xf32>
+//  CHECK-NEXT:    %[[ID3:.*]] = affine.apply #[[MAP0]]()[%[[ID]]]
+//  CHECK-NEXT:    vector.transfer_write %[[ADD]], %{{.*}}[%[[ID3]]] : vector<2xf32>, memref<64xf32>
 //  CHECK-NEXT:    return
 func @vector_add_cycle(%id : index, %A: memref<64xf32>, %B: memref<64xf32>, %C: memref<64xf32>) {
   %c0 = constant 0 : index
   %cf0 = constant 0.0 : f32
   %a = vector.transfer_read %A[%c0], %cf0: memref<64xf32>, vector<64xf32>
   %b = vector.transfer_read %B[%c0], %cf0: memref<64xf32>, vector<64xf32>
   %acc = addf %a, %b: vector<64xf32>
   vector.transfer_write %acc, %C[%c0]: vector<64xf32>, memref<64xf32>
   return
 }
 
+// -----
+
 // Negative test to make sure nothing is done in case the vector size is not a
 // multiple of multiplicity.
 // CHECK-LABEL: func @vector_negative_test
 //       CHECK:    %[[C0:.*]] = constant 0 : index
 //       CHECK:    %[[EXA:.*]] = vector.transfer_read %{{.*}}[%[[C0]]], %{{.*}} : memref<64xf32>, vector<16xf32>
 //  CHECK-NEXT:    %[[EXB:.*]] = vector.transfer_read %{{.*}}[%[[C0]]], %{{.*}} : memref<64xf32>, vector<16xf32>
 //  CHECK-NEXT:    %[[ADD:.*]] = addf %[[EXA]], %[[EXB]] : vector<16xf32>
 //  CHECK-NEXT:    vector.transfer_write %[[ADD]], %{{.*}}[%[[C0]]] {{.*}} : vector<16xf32>, memref<64xf32>

mlir/test/lib/Transforms/TestVectorTransforms.cpp

struct TestVectorDistributePatterns
   TestVectorDistributePatterns(const TestVectorDistributePatterns &pass) {}
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<VectorDialect>();
     registry.insert<AffineDialect>();
   }
   Option<int32_t> multiplicity{
       *this, "distribution-multiplicity",
       llvm::cl::desc("Set the multiplicity used for distributing vector"),
-      llvm::cl::init(32)};
+      llvm::cl::init(1)};
   void runOnFunction() override {
     MLIRContext *ctx = &getContext();
     OwningRewritePatternList patterns;
     FuncOp func = getFunction();
+    if (multiplicity > 1) {
       func.walk([&](AddFOp op) {
         OpBuilder builder(op);
         Optional<mlir::vector::DistributeOps> ops = distributPointwiseVectorOp(
             builder, op.getOperation(), func.getArgument(0), multiplicity);
         if (ops.hasValue()) {
           SmallPtrSet<Operation *, 1> extractOp({ops->extract});
-        op.getResult().replaceAllUsesExcept(ops->insert.getResult(), extractOp);
+          op.getResult().replaceAllUsesExcept(ops->insert.getResult(),
+                                              extractOp);
         }
       });
+    }
     patterns.insert<PointwiseExtractPattern>(ctx);
     populateVectorToVectorTransformationPatterns(patterns, ctx);
     applyPatternsAndFoldGreedily(getFunction(), patterns);
   }
 };
 
 struct TestVectorTransferFullPartialSplitPatterns
     : public PassWrapper<TestVectorTransferFullPartialSplitPatterns,