[mlir][VectorToSCF] 128 byte alignment of alloc ops

[limo1996]

Summary

Added 128 byte alignment to alloc ops created in VectorToSCF pass.
128b alignment was already introduced to this pass but not to all alloc
ops. This commit changes that by adding 128b alignment to the remaining ops.

[aartbik]

can you say a bit more on the underlying motivation in the description (e.g. performance improvement)?

Also, do we want to allocate all data this way, what about an 1x1 or 2x2 array allocation? Should size play a role in the heuristics?

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

581

let's move all occurrences of 128 (including the one at function entry) to a symbolic constant somewhere

[limo1996]

Aart's comments addressed

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

 #include "mlir/IR/Location.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OperationSupport.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/Types.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/Passes.h"
 
+#define ALIGNMENT_SIZE 128
+
 using namespace mlir;
 using namespace mlir::edsc;
 using namespace mlir::edsc::intrinsics;
 using vector::TransferReadOp;
 using vector::TransferWriteOp;
 
 namespace {
 /// Helper class captures the common information needed to lower N>1-D vector

 static Value setAllocAtFunctionEntry(MemRefType memRefMinorVectorType,
                                      Operation *op) {
   auto &b = ScopedContext::getBuilderRef();
   OpBuilder::InsertionGuard guard(b);
   Operation *scope =
       op->getParentWithTrait<OpTrait::AutomaticAllocationScope>();
   assert(scope && "Expected op to be inside automatic allocation scope");
   b.setInsertionPointToStart(&scope->getRegion(0).front());
-  Value res =
-      std_alloca(memRefMinorVectorType, ValueRange{}, b.getI64IntegerAttr(128));
+  Value res = std_alloca(memRefMinorVectorType, ValueRange{},
+                         b.getI64IntegerAttr(ALIGNMENT_SIZE));
   return res;
 }
 
 template <>
 LogicalResult NDTransferOpHelper<TransferReadOp>::doReplace() {
   Value alloc, result;
   if (options.unroll)
     result = std_splat(vectorType, xferOp.padding());
LogicalResult VectorTransferRewriter<TransferReadOp>::matchAndRewrite(
   auto lbs = vectorBoundsCapture.getLbs();
   auto ubs = vectorBoundsCapture.getUbs();
   SmallVector<Value, 8> steps;
   steps.reserve(vectorBoundsCapture.getSteps().size());
   for (auto step : vectorBoundsCapture.getSteps())
     steps.push_back(std_constant_index(step));
 
   // 2. Emit alloc-copy-load-dealloc.
-  Value tmp = std_alloc(tmpMemRefType(transfer));
+  Value tmp = std_alloc(tmpMemRefType(transfer), ValueRange{},
+                        rewriter.getI64IntegerAttr(ALIGNMENT_SIZE));
   StdIndexedValue local(tmp);
   Value vec = vector_type_cast(tmp);
   loopNestBuilder(lbs, ubs, steps, [&](ValueRange loopIvs) {
     auto ivs = llvm::to_vector<8>(loopIvs);
     // Swap the ivs which will reorder memory accesses.
     if (coalescedIdx >= 0)
       std::swap(ivs.back(), ivs[coalescedIdx]);
     // Computes clippedScalarAccessExprs in the loop nest scope (ivs exist).
LogicalResult VectorTransferRewriter<TransferWriteOp>::matchAndRewrite(
   auto lbs = vectorBoundsCapture.getLbs();
   auto ubs = vectorBoundsCapture.getUbs();
   SmallVector<Value, 8> steps;
   steps.reserve(vectorBoundsCapture.getSteps().size());
   for (auto step : vectorBoundsCapture.getSteps())
     steps.push_back(std_constant_index(step));
 
   // 2. Emit alloc-store-copy-dealloc.
-  Value tmp = std_alloc(tmpMemRefType(transfer));
+  Value tmp = std_alloc(tmpMemRefType(transfer), ValueRange{},
+                        rewriter.getI64IntegerAttr(ALIGNMENT_SIZE));
   StdIndexedValue local(tmp);
   Value vec = vector_type_cast(tmp);
   std_store(vectorValue, vec);
   loopNestBuilder(lbs, ubs, steps, [&](ValueRange loopIvs) {
     auto ivs = llvm::to_vector<8>(loopIvs);
     // Swap the ivs which will reorder memory accesses.
     if (coalescedIdx >= 0)
       std::swap(ivs.back(), ivs[coalescedIdx]);

mlir/test/Conversion/VectorToSCF/vector-to-loops.mlir

func @materialize_read(%M: index, %N: index, %O: index, %P: index) {
   // CHECK-DAG:  %[[C3:.*]] = constant 3 : index
   // CHECK-DAG:  %[[C4:.*]] = constant 4 : index
   // CHECK-DAG:  %[[C5:.*]] = constant 5 : index
   //     CHECK:  %{{.*}} = alloc(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) : memref<?x?x?x?xf32>
   // CHECK-NEXT:  affine.for %[[I0:.*]] = 0 to %{{.*}} step 3 {
   // CHECK-NEXT:    affine.for %[[I1:.*]] = 0 to %{{.*}} {
   // CHECK-NEXT:      affine.for %[[I2:.*]] = 0 to %{{.*}} {
   // CHECK-NEXT:        affine.for %[[I3:.*]] = 0 to %{{.*}} step 5 {
-  //      CHECK:          %[[ALLOC:.*]] = alloc() : memref<5x4x3xf32>
+  //      CHECK:          %[[ALLOC:.*]] = alloc() {alignment = 128 : i64} : memref<5x4x3xf32>
   // CHECK-NEXT:          scf.for %[[I4:.*]] = %[[C0]] to %[[C3]] step %[[C1]] {
   // CHECK-NEXT:            scf.for %[[I5:.*]] = %[[C0]] to %[[C4]] step %[[C1]] {
   // CHECK-NEXT:              scf.for %[[I6:.*]] = %[[C0]] to %[[C5]] step %[[C1]] {
   // CHECK-NEXT:                {{.*}} = affine.apply #[[$ADD]](%[[I0]], %[[I4]])
   // CHECK-NEXT:                {{.*}} = affine.apply #[[$SUB]]()[%{{.*}}]
   // CHECK-NEXT:                {{.*}} = cmpi "slt", {{.*}} : index
   // CHECK-NEXT:                {{.*}} = select
   // CHECK-NEXT:                {{.*}} = cmpi "slt", {{.*}}, %[[C0]] : index
func @materialize_write(%M: index, %N: index, %O: index, %P: index) {
   // CHECK-DAG:  %[[C3:.*]] = constant 3 : index
   // CHECK-DAG:  %[[C4:.*]] = constant 4 : index
   // CHECK-DAG:  %[[C5:.*]] = constant 5 : index
   //     CHECK:  %{{.*}} = alloc(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) : memref<?x?x?x?xf32>
   // CHECK-NEXT:  affine.for %[[I0:.*]] = 0 to %{{.*}} step 3 {
   // CHECK-NEXT:    affine.for %[[I1:.*]] = 0 to %{{.*}} step 4 {
   // CHECK-NEXT:      affine.for %[[I2:.*]] = 0 to %{{.*}} {
   // CHECK-NEXT:        affine.for %[[I3:.*]] = 0 to %{{.*}} step 5 {
-  // CHECK:               %[[ALLOC:.*]] = alloc() : memref<5x4x3xf32>
+  // CHECK:               %[[ALLOC:.*]] = alloc() {alignment = 128 : i64} : memref<5x4x3xf32>
   // CHECK-NEXT:          %[[VECTOR_VIEW:.*]] = vector.type_cast {{.*}} : memref<5x4x3xf32>
   //      CHECK:          store %{{.*}}, {{.*}} : memref<vector<5x4x3xf32>>
   // CHECK-NEXT:          scf.for %[[I4:.*]] = %[[C0]] to %[[C3]] step %[[C1]] {
   // CHECK-NEXT:            scf.for %[[I5:.*]] = %[[C0]] to %[[C4]] step %[[C1]] {
   // CHECK-NEXT:              scf.for %[[I6:.*]] = %[[C0]] to %[[C5]] step %[[C1]] {
   // CHECK-NEXT:                {{.*}} = affine.apply #[[$ADD]](%[[I0]], %[[I4]])
   // CHECK-NEXT:                {{.*}} = affine.apply #[[$SUB]]()[%{{.*}}]
   // CHECK-NEXT:                {{.*}} = cmpi "slt", {{.*}}, {{.*}} : index