AIRWrapFuncWithParallel: add a new pass that wraps a funcOp body with a scf.parallel loop with specified loop bounds

[erwei-xilinx]

• This pass is intended to be used during JIT compile, where the outer tiling size is only known at runtime.
• An scf.parallel is attached in JIT driver which represents the runtime dispatch iteration space.

Input IR:

func.func @func0(%arg0: memref<*xf32>, %arg1: memref<*xf32>, %arg2: memref<*xf32>, %arg3: i32, %arg4: i32, %arg5: i32, %arg6: i32, %arg7: i32, %arg8: i32) {
  %cst = arith.constant 0.000000e+00 : f32
  %c64 = arith.constant 64 : index
  %c32_i32 = arith.constant 32 : i32
  %0 = arith.muli %arg6, %c32_i32 : i32
  %1 = arith.index_cast %0 : i32 to index
  %2 = arith.muli %arg7, %c32_i32 : i32
  %3 = arith.index_cast %2 : i32 to index
  %4 = arith.muli %1, %c64 : index
  %reinterpret_cast = memref.reinterpret_cast %arg0 to offset: [%4], sizes: [32, 64], strides: [%c64, 1] : memref<*xf32> to memref<32x64xf32, strided<[?, 1], offset: ?>>
  %alloc = memref.alloc() : memref<32x64xf32>
  memref.copy %reinterpret_cast, %alloc : memref<32x64xf32, strided<[?, 1], offset: ?>> to memref<32x64xf32>
  %5 = bufferization.to_tensor %alloc restrict writable : memref<32x64xf32> to tensor<32x64xf32>
  %reinterpret_cast_0 = memref.reinterpret_cast %arg1 to offset: [%3], sizes: [64, 32], strides: [%c64, 1] : memref<*xf32> to memref<64x32xf32, strided<[?, 1], offset: ?>>
  %alloc_1 = memref.alloc() : memref<64x32xf32>
  memref.copy %reinterpret_cast_0, %alloc_1 : memref<64x32xf32, strided<[?, 1], offset: ?>> to memref<64x32xf32>
  %6 = bufferization.to_tensor %alloc_1 restrict writable : memref<64x32xf32> to tensor<64x32xf32>
  %7 = tensor.empty() : tensor<32x32xf32>
  %8 = linalg.fill ins(%cst : f32) outs(%7 : tensor<32x32xf32>) -> tensor<32x32xf32>
  %9 = linalg.matmul ins(%5, %6 : tensor<32x64xf32>, tensor<64x32xf32>) outs(%8 : tensor<32x32xf32>) -> tensor<32x32xf32>
  %10 = arith.addi %4, %3 : index
  %reinterpret_cast_2 = memref.reinterpret_cast %arg2 to offset: [%10], sizes: [32, 32], strides: [%c64, 1] : memref<*xf32> to memref<32x32xf32, strided<[?, 1], offset: ?>>
  bufferization.materialize_in_destination %9 in writable %reinterpret_cast_2 : (tensor<32x32xf32>, memref<32x32xf32, strided<[?, 1], offset: ?>>) -> ()
  return
}

Output IR:

module {
  func.func @func0(%arg0: memref<*xf32>, %arg1: memref<*xf32>, %arg2: memref<*xf32>, %arg3: i32, %arg4: i32, %arg5: i32, %arg6: i32, %arg7: i32, %arg8: i32) {
    %c32 = arith.constant 32 : index
    %c64 = arith.constant 64 : index
    %cst = arith.constant 0.000000e+00 : f32
    %c0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    %c4 = arith.constant 4 : index
    scf.parallel (%arg9, %arg10, %arg11) = (%c0, %c0, %c0) to (%c4, %c4, %c1) step (%c1, %c1, %c1) {
      %0 = arith.index_cast %arg9 : index to i32
      %1 = arith.index_cast %arg10 : index to i32
      %2 = arith.index_cast %0 : i32 to index
      %3 = arith.muli %2, %c32 : index
      %4 = arith.index_cast %3 : index to i32
      %5 = arith.index_cast %4 : i32 to index
      %6 = arith.index_cast %1 : i32 to index
      %7 = arith.muli %6, %c32 : index
      %8 = arith.index_cast %7 : index to i32
      %9 = arith.index_cast %8 : i32 to index
      %10 = arith.muli %5, %c64 : index
      %reinterpret_cast = memref.reinterpret_cast %arg0 to offset: [%10], sizes: [32, 64], strides: [%c64, 1] : memref<*xf32> to memref<32x64xf32, strided<[?, 1], offset: ?>>
      %alloc = memref.alloc() : memref<32x64xf32>
      memref.copy %reinterpret_cast, %alloc : memref<32x64xf32, strided<[?, 1], offset: ?>> to memref<32x64xf32>
      %11 = bufferization.to_tensor %alloc restrict writable : memref<32x64xf32> to tensor<32x64xf32>
      %reinterpret_cast_0 = memref.reinterpret_cast %arg1 to offset: [%9], sizes: [64, 32], strides: [%c64, 1] : memref<*xf32> to memref<64x32xf32, strided<[?, 1], offset: ?>>
      %alloc_1 = memref.alloc() : memref<64x32xf32>
      memref.copy %reinterpret_cast_0, %alloc_1 : memref<64x32xf32, strided<[?, 1], offset: ?>> to memref<64x32xf32>
      %12 = bufferization.to_tensor %alloc_1 restrict writable : memref<64x32xf32> to tensor<64x32xf32>
      %13 = tensor.empty() : tensor<32x32xf32>
      %14 = linalg.fill ins(%cst : f32) outs(%13 : tensor<32x32xf32>) -> tensor<32x32xf32>
      %15 = linalg.matmul ins(%11, %12 : tensor<32x64xf32>, tensor<64x32xf32>) outs(%14 : tensor<32x32xf32>) -> tensor<32x32xf32>
      %16 = arith.addi %10, %9 : index
      %reinterpret_cast_2 = memref.reinterpret_cast %arg2 to offset: [%16], sizes: [32, 32], strides: [%c64, 1] : memref<*xf32> to memref<32x32xf32, strided<[?, 1], offset: ?>>
      bufferization.materialize_in_destination %15 in writable %reinterpret_cast_2 : (tensor<32x32xf32>, memref<32x32xf32, strided<[?, 1], offset: ?>>) -> ()
      scf.reduce 
    }
    return
  }
}