search

nod-ai / iree-amd-aie

IREE plugin repository for the AMD AIE accelerator
Apache License 2.0
69 stars 30 forks source link

Optimize double buffering and loop pipelining #828

Open yzhang93 opened 1 month ago

yzhang93 commented 1 month ago

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