[mlir][bufferize] lower allocation alignment from 128 to 64 bytes

[cota]

Summary

While it is unlikely to matter in practice, there is no reason for this value to be larger than it should be. 64 bytes is the size of a cache line in most machines, and we can fit a full 512-bit vector in it.

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

struct BufferizationOptions {
   /// checking the results of the analysis) and post analysis steps.
   bool testAnalysisOnly = false;
 
   /// If set to `true`, the IR is annotated with details about RaW conflicts.
   /// For debugging only. Should be used together with `testAnalysisOnly`.
   bool printConflicts = false;
 
   /// Buffer alignment for new memory allocations.
-  unsigned int bufferAlignment = 128;
+  unsigned int bufferAlignment = 64;
 
   /// Initializer functions for analysis state. These can be used to
   /// initialize dialect-specific analysis state.
   SmallVector<AnalysisStateInitFn> stateInitializers;
 };
 
 /// Specify fine-grain relationship between buffers to enable more analysis.
 enum class BufferRelation {

mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir

 func.func @main() {
   //  CHECK-DAG:   %[[C0:.*]] = arith.constant 0{{.*}} : f32
   //  CHECK-DAG:   %[[C1:.*]] = arith.constant 1{{.*}} : f32
   //  CHECK-DAG:   %[[C2:.*]] = arith.constant 2{{.*}} : f32
   %v0 = arith.constant 0.0 : f32
   %v1 = arith.constant 1.0 : f32
   %v2 = arith.constant 2.0 : f32
 
-  // CHECK-NEXT:   %[[A:.*]] = memref.alloc() {alignment = 128 : i64} : memref<64xf32>
-  // CHECK-NEXT:   %[[B:.*]] = memref.alloc() {alignment = 128 : i64} : memref<64xf32>
-  // CHECK-NEXT:   %[[C:.*]] = memref.alloc() {alignment = 128 : i64} : memref<f32>
+  // CHECK-NEXT:   %[[A:.*]] = memref.alloc() {alignment = 64 : i64} : memref<64xf32>
+  // CHECK-NEXT:   %[[B:.*]] = memref.alloc() {alignment = 64 : i64} : memref<64xf32>
+  // CHECK-NEXT:   %[[C:.*]] = memref.alloc() {alignment = 64 : i64} : memref<f32>
   //  CHECK-DAG:   %[[cA:.*]] = memref.cast %[[A]] : memref<64xf32> to memref<64xf32, strided<[?], offset: ?>>
   //  CHECK-DAG:   %[[cB:.*]] = memref.cast %[[B]] : memref<64xf32> to memref<64xf32, strided<[?], offset: ?>>
   //  CHECK-DAG:   %[[cC:.*]] = memref.cast %[[C]] : memref<f32> to memref<f32, strided<[], offset: ?>>
   %A = bufferization.alloc_tensor() : tensor<64xf32>
   %B = bufferization.alloc_tensor() : tensor<64xf32>
   %C = bufferization.alloc_tensor() : tensor<f32>
 
   //  CHECK-DAG:   linalg.fill ins(%[[C1]] : f32) outs(%[[A]] : memref<64xf32>)

mlir/test/Dialect/Linalg/one-shot-bufferize.mlir

 func.func @not_inplace(
     %A : tensor<?xf32> {bufferization.writable = false})
   -> tensor<?xf32>
 {
   //     CHECK: %[[F0:.*]] = arith.constant 0.000000e+00 : f32
   %f0 = arith.constant 0.0 : f32
 
   //     CHECK: %[[D0:.*]] = memref.dim %[[A]], {{.*}} : memref<?xf32, strided<[?], offset: ?>>
-  //     CHECK: %[[ALLOC:.*]] = memref.alloc(%[[D0]]) {alignment = 128 : i64} : memref<?xf32>
+  //     CHECK: %[[ALLOC:.*]] = memref.alloc(%[[D0]]) {alignment = 64 : i64} : memref<?xf32>
   //     CHECK: linalg.fill ins(%[[F0]] : f32) outs(%[[ALLOC]] : memref<?xf32>)
   %r = linalg.fill ins(%f0 : f32) outs(%A : tensor<?xf32>) -> tensor<?xf32>
 
   // CHECK-NOT: dealloc
   //     CHECK: return %[[ALLOC]] : memref<?xf32>
   return %r: tensor<?xf32>
 }
 
func.func @matmul(
   %c32 = arith.constant 32 : index
   %cst = arith.constant 0.000000e+00 : f32
   %c128 = arith.constant 128 : index
   %c192 = arith.constant 192 : index
   %c8 = arith.constant 8 : index
   %c16 = arith.constant 16 : index
 
   // Hoisted alloc.
-  // CHECK: %[[ALLOC:.*]] = memref.alloc() {alignment = 128 : i64} : memref<128x192xf32>
+  // CHECK: %[[ALLOC:.*]] = memref.alloc() {alignment = 64 : i64} : memref<128x192xf32>
   // CHECK: memref.copy %[[C]], %[[ALLOC]]
 
   // CHECK: scf.for %[[I:.*]] =
   %0 = scf.for %arg3 = %c0 to %c128 step %c8 iter_args(%arg4 = %C) -> (tensor<128x192xf32>) {
     %1 = tensor.extract_slice %A[%arg3, 0] [8, 256] [1, 1] :
       tensor<128x256xf32> to tensor<8x256xf32>
 
     // CHECK: scf.for %[[J:.*]] =

mlir/test/Dialect/Tensor/bufferize.mlir

 func.func @tensor.reshape(%t1: tensor<?x10xf32>) -> tensor<2x2x5xf32> {
   // CHECK: %[[m1:.*]] = bufferization.to_memref %[[t1]] : memref<?x10xf32>
 
   // CHECK: %[[two:.*]] = arith.constant 2 : i64
   %two = arith.constant 2 : i64
   // CHECK: %[[five:.*]] = arith.constant 5 : i64
   %five = arith.constant 5 : i64
 
-  // CHECK: %[[alloc:.*]] = memref.alloc() {alignment = 128 : i64} : memref<3xi64>
+  // CHECK: %[[alloc:.*]] = memref.alloc() {alignment = 64 : i64} : memref<3xi64>
   // CHECK: %[[zero_idx:.*]] = arith.constant 0 : index
   // CHECK: %[[one_idx:.*]] = arith.constant 1 : index
   // CHECK: %[[two_idx:.*]] = arith.constant 2 : index
   // CHECK: memref.store %[[two]], %[[alloc]][%[[zero_idx]]] : memref<3xi64>
   // CHECK: memref.store %[[two]], %[[alloc]][%[[one_idx]]] : memref<3xi64>
   // CHECK: memref.store %[[five]], %[[alloc]][%[[two_idx]]] : memref<3xi64>
   %shape = tensor.from_elements %two, %two, %five : tensor<3xi64>
 

mlir/test/Dialect/Tensor/one-shot-bufferize.mlir

 // CHECK-LABEL: func @insert_slice_fun_not_inplace
 //  CHECK-SAME:   %[[A:[a-zA-Z0-9]*]]: memref<?xf32, strided<[?], offset: ?>>
 //  CHECK-SAME:   %[[t:[a-zA-Z0-9]*]]: memref<4xf32, strided<[?], offset: ?>>
 func.func @insert_slice_fun_not_inplace(
     %A : tensor<?xf32> {bufferization.writable = false},
     %t : tensor<4xf32> {bufferization.writable = false})
   -> tensor<?xf32>
 {
-  //      CHECK: %[[ALLOC:.*]] = memref.alloc(%{{.*}}) {alignment = 128 : i64} : memref<?xf32>
+  //      CHECK: %[[ALLOC:.*]] = memref.alloc(%{{.*}}) {alignment = 64 : i64} : memref<?xf32>
   //      CHECK: memref.copy %[[A]], %[[ALLOC]] : memref<?xf32{{.*}} to memref<?xf32>
   //      CHECK: %[[SV:.*]] = memref.subview %[[ALLOC]][0] [4] [1] : memref<?xf32> to memref<4xf32, strided<[1]>>
   //      CHECK: memref.copy %[[t]], %[[SV]] : memref<4xf32, strided{{.*}}> to memref<4xf32, strided<[1]>>
   %r0 = tensor.insert_slice %t into %A[0][4][1] : tensor<4xf32> into tensor<?xf32>
 
   //     CHECK: return %{{.*}} : memref<?xf32>
   return %r0: tensor<?xf32>
 }