Optimize double buffering and loop pipelining

[yzhang93]

The current way of processing double buffering is to simply unroll the loops, and thus doesn't show performance improvement compared to single buffering.

scf.for %arg0 = %c0 to %c4 step %c2_85 {
  amdaie.use_lock(%lock_63, AcquireGreaterOrEqual(1))
  %reinterpret_cast = memref.reinterpret_cast %buffer_61 to offset: [0], sizes: [1, 1, 8, 8, 4, 4], strides: [1024, 1024, 128, 16, 4, 1] : memref<1024xi32, 2 : i32> to memref<1x1x8x8x4x4xi32, 2 : i32>
  linalg.fill ins(%c0_i32 : i32) outs(%reinterpret_cast : memref<1x1x8x8x4x4xi32, 2 : i32>)
  amdaie.use_lock(%lock_68, AcquireGreaterOrEqual(1))
  %reinterpret_cast_86 = memref.reinterpret_cast %buffer_65 to offset: [0], sizes: [1, 1, 4, 8, 4, 8], strides: [1024, 1024, 256, 32, 8, 1] : memref<1024xi32, 2 : i32> to memref<1x1x4x8x4x8xi32, 2 : i32>
  amdaie.use_lock(%lock_72, AcquireGreaterOrEqual(1))
  %reinterpret_cast_87 = memref.reinterpret_cast %buffer_69 to offset: [0], sizes: [1, 1, 8, 4, 8, 4], strides: [1024, 1024, 128, 32, 4, 1] : memref<1024xi32, 2 : i32> to memref<1x1x8x4x8x4xi32, 2 : i32>
  linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d2, d5, d3, d6, d8)>, affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d2, d1, d4, d5, d8, d7)>, affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d1, d4, d3, d6, d7)>], iterator_types = ["parallel", "parallel", "reduction", "parallel", "parallel", "reduction", "parallel", "parallel", "reduction"]} ins(%reinterpret_cast_86, %reinterpret_cast_87 : memref<1x1x4x8x4x8xi32, 2 : i32>, memref<1x1x8x4x8x4xi32, 2 : i32>) outs(%reinterpret_cast : memref<1x1x8x8x4x4xi32, 2 : i32>) attrs =  {lowering_config = #iree_codegen.lowering_config<tile_sizes = [[64, 64], [0, 0, 1], [1, 1, 0, 0, 0, 0]]>, packing_config = #amdaie.packing_config<packing_config = [{packedSizes = [32, 32, 32], transposePackIndices = [1], unpackEmpty = [false], innerPerm = [[1, 0]], outerPerm = [[0, 1]]}, {packedSizes = [0, 0, 0, 4, 4, 8], transposePackIndices = [0, 1, 2], unpackEmpty = [false, false, true], innerPerm = [[0, 1], [1, 0], [0, 1]], outerPerm = [[0, 1, 3, 2], [0, 1, 3, 2], [0, 1, 3, 2]]}]>} {
  ^bb0(%in: i32, %in_97: i32, %out: i32):
    %32 = arith.muli %in, %in_97 : i32
    %33 = arith.addi %out, %32 : i32
    linalg.yield %33 : i32
  }
  %c2_88 = arith.constant 2 : index
  scf.for %arg1 = %c0 to %c6 step %c2_88 {
    amdaie.use_lock(%lock_67, Release(1))
    amdaie.use_lock(%lock_68, AcquireGreaterOrEqual(1))
    %reinterpret_cast_97 = memref.reinterpret_cast %buffer_66 to offset: [0], sizes: [1, 1, 4, 8, 4, 8], strides: [1024, 1024, 256, 32, 8, 1] : memref<1024xi32, 2 : i32> to memref<1x1x4x8x4x8xi32, 2 : i32>
    amdaie.use_lock(%lock_71, Release(1))
    amdaie.use_lock(%lock_72, AcquireGreaterOrEqual(1))
    %reinterpret_cast_98 = memref.reinterpret_cast %buffer_70 to offset: [0], sizes: [1, 1, 8, 4, 8, 4], strides: [1024, 1024, 128, 32, 4, 1] : memref<1024xi32, 2 : i32> to memref<1x1x8x4x8x4xi32, 2 : i32>
    linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d2, d5, d3, d6, d8)>, affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d2, d1, d4, d5, d8, d7)>, affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d1, d4, d3, d6, d7)>], iterator_types = ["parallel", "parallel", "reduction", "parallel", "parallel", "reduction", "parallel", "parallel", "reduction"]} ins(%reinterpret_cast_97, %reinterpret_cast_98 : memref<1x1x4x8x4x8xi32, 2 : i32>, memref<1x1x8x4x8x4xi32, 2 : i32>) outs(%reinterpret_cast : memref<1x1x8x8x4x4xi32, 2 : i32>) attrs =  {lowering_config = #iree_codegen.lowering_config<tile_sizes = [[64, 64], [0, 0, 1], [1, 1, 0, 0, 0, 0]]>, packing_config = #amdaie.packing_config<packing_config = [{packedSizes = [32, 32, 32], transposePackIndices = [1], unpackEmpty = [false], innerPerm = [[1, 0]], outerPerm = [[0, 1]]}, {packedSizes = [0, 0, 0, 4, 4, 8], transposePackIndices = [0, 1, 2], unpackEmpty = [false, false, true], innerPerm = [[0, 1], [1, 0], [0, 1]], outerPerm = [[0, 1, 3, 2], [0, 1, 3, 2], [0, 1, 3, 2]]}]>} {
    ^bb0(%in: i32, %in_101: i32, %out: i32):
      %32 = arith.muli %in, %in_101 : i32
      %33 = arith.addi %out, %32 : i32
      linalg.yield %33 : i32
    }
    amdaie.use_lock(%lock_67, Release(1))
    amdaie.use_lock(%lock_68, AcquireGreaterOrEqual(1))
    %reinterpret_cast_99 = memref.reinterpret_cast %buffer_65 to offset: [0], sizes: [1, 1, 4, 8, 4, 8], strides: [1024, 1024, 256, 32, 8, 1] : memref<1024xi32, 2 : i32> to memref<1x1x4x8x4x8xi32, 2 : i32>
    amdaie.use_lock(%lock_71, Release(1))
    amdaie.use_lock(%lock_72, AcquireGreaterOrEqual(1))
    %reinterpret_cast_100 = memref.reinterpret_cast %buffer_69 to offset: [0], sizes: [1, 1, 8, 4, 8, 4], strides: [1024, 1024, 128, 32, 4, 1] : memref<1024xi32, 2 : i32> to memref<1x1x8x4x8x4xi32, 2 : i32>
    linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d2, d5, d3, d6, d8)>, affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d2, d1, d4, d5, d8, d7)>, affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d1, d4, d3, d6, d7)>], iterator_types = ["parallel", "parallel", "reduction", "parallel", "parallel", "reduction", "parallel", "parallel", "reduction"]} ins(%reinterpret_cast_99, %reinterpret_cast_100 : memref<1x1x4x8x4x8xi32, 2 : i32>, memref<1x1x8x4x8x4xi32, 2 : i32>) outs(%reinterpret_cast : memref<1x1x8x8x4x4xi32, 2 : i32>) attrs =  {lowering_config = #iree_codegen.lowering_config<tile_sizes = [[64, 64], [0, 0, 1], [1, 1, 0, 0, 0, 0]]>, packing_config = #amdaie.packing_config<packing_config = [{packedSizes = [32, 32, 32], transposePackIndices = [1], unpackEmpty = [false], innerPerm = [[1, 0]], outerPerm = [[0, 1]]}, {packedSizes = [0, 0, 0, 4, 4, 8], transposePackIndices = [0, 1, 2], unpackEmpty = [false, false, true], innerPerm = [[0, 1], [1, 0], [0, 1]], outerPerm = [[0, 1, 3, 2], [0, 1, 3, 2], [0, 1, 3, 2]]}]>} {
    ^bb0(%in: i32, %in_101: i32, %out: i32):
      %32 = arith.muli %in, %in_101 : i32
      %33 = arith.addi %out, %32 : i32
      linalg.yield %33 : i32
    }
  }
  amdaie.use_lock(%lock_67, Release(1))
  amdaie.use_lock(%lock_68, AcquireGreaterOrEqual(1))
  %reinterpret_cast_89 = memref.reinterpret_cast %buffer_66 to offset: [0], sizes: [1, 1, 4, 8, 4, 8], strides: [1024, 1024, 256, 32, 8, 1] : memref<1024xi32, 2 : i32> to memref<1x1x4x8x4x8xi32, 2 : i32>
  amdaie.use_lock(%lock_71, Release(1))
  amdaie.use_lock(%lock_72, AcquireGreaterOrEqual(1))
  %reinterpret_cast_90 = memref.reinterpret_cast %buffer_70 to offset: [0], sizes: [1, 1, 8, 4, 8, 4], strides: [1024, 1024, 128, 32, 4, 1] : memref<1024xi32, 2 : i32> to memref<1x1x8x4x8x4xi32, 2 : i32>
  linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d2, d5, d3, d6, d8)>, affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d2, d1, d4, d5, d8, d7)>, affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d1, d4, d3, d6, d7)>], iterator_types = ["parallel", "parallel", "reduction", "parallel", "parallel", "reduction", "parallel", "parallel", "reduction"]} ins(%reinterpret_cast_89, %reinterpret_cast_90 : memref<1x1x4x8x4x8xi32, 2 : i32>, memref<1x1x8x4x8x4xi32, 2 : i32>) outs(%reinterpret_cast : memref<1x1x8x8x4x4xi32, 2 : i32>) attrs =  {lowering_config = #iree_codegen.lowering_config<tile_sizes = [[64, 64], [0, 0, 1], [1, 1, 0, 0, 0, 0]]>, packing_config = #amdaie.packing_config<packing_config = [{packedSizes = [32, 32, 32], transposePackIndices = [1], unpackEmpty = [false], innerPerm = [[1, 0]], outerPerm = [[0, 1]]}, {packedSizes = [0, 0, 0, 4, 4, 8], transposePackIndices = [0, 1, 2], unpackEmpty = [false, false, true], innerPerm = [[0, 1], [1, 0], [0, 1]], outerPerm = [[0, 1, 3, 2], [0, 1, 3, 2], [0, 1, 3, 2]]}]>} {
  ^bb0(%in: i32, %in_97: i32, %out: i32):
    %32 = arith.muli %in, %in_97 : i32
    %33 = arith.addi %out, %32 : i32
    linalg.yield %33 : i32
  }
  amdaie.use_lock(%lock_67, Release(1))
  amdaie.use_lock(%lock_71, Release(1))
  amdaie.use_lock(%lock_64, Release(1))

We probably want to apply correct loop pipelining to hide some memory access latency (although I'm not sure how much improvement it can bring). A simple loop pipelining strategy is as

load (A[0], B[0])
load (A[1], B[1])
matmul(A[0], B[0])
write (C[0])
for i = 2, ... n {
    load(A[i], B[i])
    matmul(A[i-1], B[i-1])
    write(C[i-1])
}
matmul(A[n], B[n])
write (C[n])

Currently there are double buffers in both L2 and L1, and another question is whether we need double buffers for both inputs and outputs in L2. The idea behind this is we'd want to free some memory so that a larger tile size can be used.

CC: @jtuyls @MaheshRavishankar