[mlir][vector] Add unit test for vector distribute by block

[ThomasRaoux]

Summary

When distributing a vector larger than the given multiplicity, we can distribute it by block where each id gets a chunk of consecutive element along the dimension distributed. This adds a test for this case and adds extra checks to make sure we don't distribute for cases not multiple of multiplicity.

[aartbik]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

2451

please add explanation in a comment (since revision comment is harder to obtain)

[ThomasRaoux]

Add a comment.

mlir/lib/Dialect/Vector/VectorTransforms.cpp

2451

Sure, Added some explanation.

[aartbik]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

2452

typo: support

mlir/test/lib/Transforms/TestVectorTransforms.cpp

149

TABs in file?

[ThomasRaoux]

Fix typo

[ThomasRaoux]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

2452

Fixed.

mlir/test/lib/Transforms/TestVectorTransforms.cpp

149

I doubled checked that those are not tabs. I believe this is how phabricator displays that the line change is an indentation.

[nicolasvasilache]

LG thanks @ThomasRaoux

mlir/lib/Dialect/Vector/VectorTransforms.cpp

 }
 
 Optional<mlir::vector::DistributeOps>
 mlir::vector::distributPointwiseVectorOp(OpBuilder &builder, Operation *op,
                                          Value id, int64_t multiplicity) {
   OpBuilder::InsertionGuard guard(builder);
   builder.setInsertionPointAfter(op);
   Location loc = op->getLoc();
+  if (op->getNumResults() != 1)
+    return {};
   Value result = op->getResult(0);
+  VectorType type = op->getResult(0).getType().dyn_cast<VectorType>();
+  // Currently only support distributing 1-D vectors of size multiple of the
+  // given multiplicty. To handle more sizes we would need to support masking.
+  if (!type || type.getRank() != 1 || type.getNumElements() % multiplicity != 0)
+    return {};
   DistributeOps ops;
   ops.extract =
       builder.create<vector::ExtractMapOp>(loc, result, id, multiplicity);
   ops.insert =
       builder.create<vector::InsertMapOp>(loc, ops.extract, id, multiplicity);
   return ops;
 }
 

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

-// RUN: mlir-opt %s -test-vector-distribute-patterns | FileCheck %s
+// RUN: mlir-opt %s -test-vector-distribute-patterns=distribution-multiplicity=32 | 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 %{{.*}}[%{{.*}}], %{{.*}} : memref<32xf32>, vector<1xf32>
-//  CHECK-NEXT:    %[[EXB:.*]] = vector.transfer_read %{{.*}}[%{{.*}}], %{{.*}} : memref<32xf32>, vector<1xf32>
+//       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 %{{.*}}[%{{.*}}], %{{.*}} : memref<32xf32>, 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]], %{{.*}}[%{{.*}}] : vector<1xf32>, memref<32xf32>
+//  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-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-NEXT:    %[[ADD:.*]] = addf %[[EXA]], %[[EXB]] : vector<2xf32>
+//  CHECK-NEXT:    vector.transfer_write %[[ADD]], %{{.*}}[%[[ID]]] : 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>
+//  CHECK-NEXT:    return
+func @vector_negative_test(%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<16xf32>
+  %b = vector.transfer_read %B[%c0], %cf0: memref<64xf32>, vector<16xf32>
+  %acc = addf %a, %b: vector<16xf32>
+  vector.transfer_write %acc, %C[%c0]: vector<16xf32>, memref<64xf32>
+  return
+}
+
+

mlir/test/lib/Transforms/TestVectorTransforms.cpp

void runOnFunction() override {
     populateVectorToVectorCanonicalizationPatterns(patterns, ctx);
     populateVectorToVectorTransformationPatterns(patterns, ctx);
     applyPatternsAndFoldGreedily(getFunction(), patterns);
   }
 };
 
 struct TestVectorDistributePatterns
     : public PassWrapper<TestVectorDistributePatterns, FunctionPass> {
+  TestVectorDistributePatterns() = default;
+  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)};
   void runOnFunction() override {
     MLIRContext *ctx = &getContext();
     OwningRewritePatternList patterns;
     FuncOp func = getFunction();
     func.walk([&](AddFOp op) {
       OpBuilder builder(op);
       Optional<mlir::vector::DistributeOps> ops = distributPointwiseVectorOp(
-          builder, op.getOperation(), func.getArgument(0), 32);
-      assert(ops.hasValue());
+          builder, op.getOperation(), func.getArgument(0), multiplicity);
+      if (ops.hasValue()) {
         SmallPtrSet<Operation *, 1> extractOp({ops->extract});
         op.getResult().replaceAllUsesExcept(ops->insert.getResult(), extractOp);
+      }
     });
     patterns.insert<PointwiseExtractPattern>(ctx);
     populateVectorToVectorTransformationPatterns(patterns, ctx);
     applyPatternsAndFoldGreedily(getFunction(), patterns);
   }
 };
 
 struct TestVectorTransferFullPartialSplitPatterns