[Bug] FA lit regression test with S1_TILE=512 crashes at runtime, splitting to S1_TILE=256 works

[MirkoDeVita98]

Component

PTO Dialect / ODS (include/PTO/IR)

Description

The MLIR fa_perf.pto provided in #609

uses:

Q_ROWS=3072
S1_TOTAL=8192
HEAD=128
S1_TILE=512
QK tile: 32x512xf32
P tile sent V2C: 16x512xf16
PV tile: 16x128xf32

PTOAS accepts the MLIR and bisheng builds the generated C++ shared object, but the first kernel launch crashes with an AICore exception before correctness or benchmarking can run.

The smallest working change we found was to split the S1_TILE=512 protocol into two S1_TILE=256 tiles while keeping the same total problem size:

Q_ROWS=3072
S1_TOTAL=8192
HEAD=128
S1_TILE=256
NUM_TILES=32

This changes the tiles from 16x512 vector/QK chunks to 16x256
chunks. It also changes all derived MLIR constants: QK FIFO slot size, GM
per-block scratch size, tile buffer shapes, local scratch offsets, and the
software-pipeline loop count.

Important isolation results:

Original fa_perf.pto, S1_TILE=512, Q=3072, KV=8192: crashes
Same S1_TILE=512 shape with P pipe local_slot_num=1: still crashes
S1_TILE=512 with smaller Q=2048, KV=4096: still crashes
S1_TILE=256 with Q=3072, KV=8192, NUM_TILES=32: passes

So this does not appear to be caused only by total sequence length, block count, GM slot size, or the P V2C FIFO local slot count. The failure has to do with the 512-wide QK/P tile protocol.

Right now, the performance of the patched version is poor:

Q_ROWS=3072
S1_TOTAL=8192
HEAD=128
block_dim=24
kernel_us=438.92
kernel_tflops=29.356
fused_us=188.20
fused_tflops=68.462
speedup_vs_fused=0.43x

Reproduction (minimal)

Use the fa_perf.pto in #609

Expected behavior

The patched version with S1_TILE=256 works:

module {
  func.func @cube_kernel(%arg0: !pto.ptr<f32>, %arg1: !pto.ptr<f16>, %arg2: !pto.ptr<f16>, %arg3: !pto.ptr<f16>) attributes {pto.kernel_kind = #pto.kernel_kind<cube>} {
    %c0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    %c2 = arith.constant 2 : index
    %c32 = arith.constant 32 : index
    %c128 = arith.constant 128 : index
    %c256 = arith.constant 256 : index
    %c8192 = arith.constant 8192 : index
    %c32_0 = arith.constant 32 : index
    %c96 = arith.constant 96 : index
    %0 = pto.get_block_num
    %1 = arith.index_cast %0 : i64 to index
    %2 = pto.get_block_idx
    %3 = arith.index_cast %2 : i64 to index
    %4 = arith.divsi %c96, %1 : index
    %5 = arith.remsi %c96, %1 : index
    %6 = arith.addi %4, %c1 : index
    %7 = arith.muli %3, %6 : index
    %8 = arith.addi %4, %c1 : index
    %9 = arith.muli %5, %8 : index
    %10 = arith.subi %3, %5 : index
    %11 = arith.muli %10, %4 : index
    %12 = arith.addi %9, %11 : index
    %13 = arith.cmpi slt, %3, %5 : index
    %14 = arith.select %13, %7, %12 : index
    %15 = arith.cmpi slt, %3, %5 : index
    %16 = arith.addi %4, %c1 : index
    %17 = arith.select %15, %16, %4 : index
    %18 = arith.addi %14, %17 : index
    %c131072 = arith.constant 131072 : index
    %19 = arith.muli %3, %c131072 : index
    %20 = pto.addptr %arg0, %19 : <f32> -> <f32>
    %c0_1 = arith.constant 0 : index
    %21 = pto.addptr %20, %c0_1 : <f32> -> <f32>
    %c65536 = arith.constant 65536 : index
    %22 = pto.addptr %20, %c65536 : <f32> -> <f32>
    %c98304 = arith.constant 98304 : index
    %23 = pto.addptr %20, %c98304 : <f32> -> <f32>
    %24 = pto.import_reserved_buffer{name = "fa_qk_c2v_fifo", peer_func = @vector_kernel} -> i32
    %25 = pto.initialize_l2g2l_pipe{dir_mask = 1, slot_size = 32768, slot_num = 8, local_slot_num = 1} (%21 : !pto.ptr<f32>, %24 : i32) -> !pto.pipe
    %26 = pto.import_reserved_buffer{name = "fa_pv_c2v_fifo", peer_func = @vector_kernel} -> i32
    %27 = pto.initialize_l2g2l_pipe{dir_mask = 1, slot_size = 16384, slot_num = 8, local_slot_num = 1} (%22 : !pto.ptr<f32>, %26 : i32) -> !pto.pipe
    %28 = pto.reserve_buffer{name = "fa_p_v2c_fifo", size = 16384, location = <mat>, auto = false, base = 262144} -> i32
    %c0_i32 = arith.constant 0 : i32
    pto.aic_initialize_pipe {id = 30, dir_mask = 2, slot_size = 16384, local_slot_num = 1, nosplit = false}(gm_slot_buffer = %23 : !pto.ptr<f32>, c2v_consumer_buf = %c0_i32 : i32, v2c_consumer_buf = %28 : i32)
    %c0_i64 = arith.constant 0 : i64
    %c0_i64_2 = arith.constant 0 : i64
    %29 = pto.alloc_tile addr = %c0_i64_2 : !pto.tile_buf<mat, 32x128xf16, blayout=col_major, slayout=row_major>
    %c0_i64_3 = arith.constant 0 : i64
    %30 = pto.alloc_tile addr = %c0_i64_3 : !pto.tile_buf<left, 32x128xf16, slayout=row_major>
    %c8192_i64 = arith.constant 8192 : i64
    %31 = pto.alloc_tile addr = %c8192_i64 : !pto.tile_buf<mat, 128x256xf16, slayout=col_major>
    %32 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf<right, 128x256xf16, slayout=col_major>
    %c0_i64_4 = arith.constant 0 : i64
    %33 = pto.alloc_tile addr = %c0_i64_4 : !pto.tile_buf<acc, 32x256xf32, blayout=col_major, slayout=row_major, fractal=1024>
    %c73728_i64 = arith.constant 73728 : i64
    %34 = pto.alloc_tile addr = %c73728_i64 : !pto.tile_buf<mat, 32x256xf16, blayout=col_major, slayout=row_major>
    %c8192_i64_5 = arith.constant 8192 : i64
    %35 = pto.alloc_tile addr = %c8192_i64_5 : !pto.tile_buf<left, 32x256xf16, slayout=row_major>
    %c90112_i64 = arith.constant 90112 : i64
    %36 = pto.alloc_tile addr = %c90112_i64 : !pto.tile_buf<mat, 256x128xf16, blayout=col_major, slayout=row_major>
    %37 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf<right, 256x128xf16, slayout=col_major>
    %c32768_i64 = arith.constant 32768 : i64
    %38 = pto.alloc_tile addr = %c32768_i64 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>
    %c3072 = arith.constant 3072 : index
    %39 = pto.make_tensor_view %arg1, shape = [%c3072, %c128], strides = [%c128, %c1] : !pto.tensor_view<?x?xf16>
    %40 = pto.make_tensor_view %arg2, shape = [%c128, %c8192], strides = [%c1, %c128] : !pto.tensor_view<?x?xf16>
    %41 = pto.make_tensor_view %arg3, shape = [%c8192, %c128], strides = [%c128, %c1] : !pto.tensor_view<?x?xf16>
    scf.for %arg4 = %14 to %18 step %c1 {
      %42 = arith.muli %arg4, %c32 : index
      %43 = pto.partition_view %39, offsets = [%42, %c0], sizes = [%c32, %c128] : !pto.tensor_view<?x?xf16>
      pto.tload ins(%43 : !pto.partition_tensor_view<32x128xf16>) outs(%29 : !pto.tile_buf<mat, 32x128xf16, blayout=col_major, slayout=row_major>)
      pto.tmov ins(%29 : !pto.tile_buf<mat, 32x128xf16, blayout=col_major, slayout=row_major>) outs(%30 : !pto.tile_buf<left, 32x128xf16, slayout=row_major>)
      %c0_6 = arith.constant 0 : index
      %44 = pto.partition_view %40, offsets = [%c0, %c0_6], sizes = [%c128, %c256] : !pto.tensor_view<?x?xf16>
      pto.tload ins(%44 : !pto.partition_tensor_view<128x256xf16>) outs(%31 : !pto.tile_buf<mat, 128x256xf16, slayout=col_major>)
      pto.tmov ins(%31 : !pto.tile_buf<mat, 128x256xf16, slayout=col_major>) outs(%32 : !pto.tile_buf<right, 128x256xf16, slayout=col_major>)
      pto.tmatmul ins(%30, %32 : !pto.tile_buf<left, 32x128xf16, slayout=row_major>, !pto.tile_buf<right, 128x256xf16, slayout=col_major>) outs(%33 : !pto.tile_buf<acc, 32x256xf32, blayout=col_major, slayout=row_major, fractal=1024>)
      pto.tpush(%33, %25 : !pto.tile_buf<acc, 32x256xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
      %c256_7 = arith.constant 256 : index
      %45 = pto.partition_view %40, offsets = [%c0, %c256_7], sizes = [%c128, %c256] : !pto.tensor_view<?x?xf16>
      pto.tload ins(%45 : !pto.partition_tensor_view<128x256xf16>) outs(%31 : !pto.tile_buf<mat, 128x256xf16, slayout=col_major>)
      pto.tmov ins(%31 : !pto.tile_buf<mat, 128x256xf16, slayout=col_major>) outs(%32 : !pto.tile_buf<right, 128x256xf16, slayout=col_major>)
      pto.tmatmul ins(%30, %32 : !pto.tile_buf<left, 32x128xf16, slayout=row_major>, !pto.tile_buf<right, 128x256xf16, slayout=col_major>) outs(%33 : !pto.tile_buf<acc, 32x256xf32, blayout=col_major, slayout=row_major, fractal=1024>)
      pto.tpush(%33, %25 : !pto.tile_buf<acc, 32x256xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
      %46 = pto.partition_view %41, offsets = [%c0, %c0], sizes = [%c256, %c128] : !pto.tensor_view<?x?xf16>
      pto.tload ins(%46 : !pto.partition_tensor_view<256x128xf16>) outs(%36 : !pto.tile_buf<mat, 256x128xf16, blayout=col_major, slayout=row_major>)
      %c15 = arith.constant 15 : index
      scf.for %arg5 = %c0 to %c15 step %c1 {
        %50 = arith.muli %arg5, %c2 : index
        %c2_8 = arith.constant 2 : index
        %51 = arith.addi %50, %c2_8 : index
        %52 = arith.muli %51, %c256 : index
        %53 = pto.partition_view %40, offsets = [%c0, %52], sizes = [%c128, %c256] : !pto.tensor_view<?x?xf16>
        pto.tload ins(%53 : !pto.partition_tensor_view<128x256xf16>) outs(%31 : !pto.tile_buf<mat, 128x256xf16, slayout=col_major>)
        %54 = pto.tpop_from_aiv {id = 30, split = 1} -> !pto.tile_buf<mat, 32x256xf16, blayout=col_major, slayout=row_major>
        pto.tmov ins(%54 : !pto.tile_buf<mat, 32x256xf16, blayout=col_major, slayout=row_major>) outs(%35 : !pto.tile_buf<left, 32x256xf16, slayout=row_major>)
        pto.tfree_from_aiv {id = 30, split = 1}
        pto.tmov ins(%36 : !pto.tile_buf<mat, 256x128xf16, blayout=col_major, slayout=row_major>) outs(%37 : !pto.tile_buf<right, 256x128xf16, slayout=col_major>)
        %55 = arith.addi %50, %c1 : index
        %56 = arith.muli %55, %c256 : index
        %57 = pto.partition_view %41, offsets = [%56, %c0], sizes = [%c256, %c128] : !pto.tensor_view<?x?xf16>
        pto.tload ins(%57 : !pto.partition_tensor_view<256x128xf16>) outs(%36 : !pto.tile_buf<mat, 256x128xf16, blayout=col_major, slayout=row_major>)
        pto.tmatmul ins(%35, %37 : !pto.tile_buf<left, 32x256xf16, slayout=row_major>, !pto.tile_buf<right, 256x128xf16, slayout=col_major>) outs(%38 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>)
        pto.tpush(%38, %27 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
        pto.tmov ins(%31 : !pto.tile_buf<mat, 128x256xf16, slayout=col_major>) outs(%32 : !pto.tile_buf<right, 128x256xf16, slayout=col_major>)
        pto.tmatmul ins(%30, %32 : !pto.tile_buf<left, 32x128xf16, slayout=row_major>, !pto.tile_buf<right, 128x256xf16, slayout=col_major>) outs(%33 : !pto.tile_buf<acc, 32x256xf32, blayout=col_major, slayout=row_major, fractal=1024>)
        pto.tpush(%33, %25 : !pto.tile_buf<acc, 32x256xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
        %58 = arith.muli %arg5, %c2 : index
        %59 = arith.addi %58, %c1 : index
        %c2_9 = arith.constant 2 : index
        %60 = arith.addi %59, %c2_9 : index
        %61 = arith.muli %60, %c256 : index
        %62 = pto.partition_view %40, offsets = [%c0, %61], sizes = [%c128, %c256] : !pto.tensor_view<?x?xf16>
        pto.tload ins(%62 : !pto.partition_tensor_view<128x256xf16>) outs(%31 : !pto.tile_buf<mat, 128x256xf16, slayout=col_major>)
        %63 = pto.tpop_from_aiv {id = 30, split = 1} -> !pto.tile_buf<mat, 32x256xf16, blayout=col_major, slayout=row_major>
        pto.tmov ins(%63 : !pto.tile_buf<mat, 32x256xf16, blayout=col_major, slayout=row_major>) outs(%35 : !pto.tile_buf<left, 32x256xf16, slayout=row_major>)
        pto.tfree_from_aiv {id = 30, split = 1}
        pto.tmov ins(%36 : !pto.tile_buf<mat, 256x128xf16, blayout=col_major, slayout=row_major>) outs(%37 : !pto.tile_buf<right, 256x128xf16, slayout=col_major>)
        %64 = arith.addi %59, %c1 : index
        %65 = arith.muli %64, %c256 : index
        %66 = pto.partition_view %41, offsets = [%65, %c0], sizes = [%c256, %c128] : !pto.tensor_view<?x?xf16>
        pto.tload ins(%66 : !pto.partition_tensor_view<256x128xf16>) outs(%36 : !pto.tile_buf<mat, 256x128xf16, blayout=col_major, slayout=row_major>)
        pto.tmatmul ins(%35, %37 : !pto.tile_buf<left, 32x256xf16, slayout=row_major>, !pto.tile_buf<right, 256x128xf16, slayout=col_major>) outs(%38 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>)
        pto.tpush(%38, %27 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
        pto.tmov ins(%31 : !pto.tile_buf<mat, 128x256xf16, slayout=col_major>) outs(%32 : !pto.tile_buf<right, 128x256xf16, slayout=col_major>)
        pto.tmatmul ins(%30, %32 : !pto.tile_buf<left, 32x128xf16, slayout=row_major>, !pto.tile_buf<right, 128x256xf16, slayout=col_major>) outs(%33 : !pto.tile_buf<acc, 32x256xf32, blayout=col_major, slayout=row_major, fractal=1024>)
        pto.tpush(%33, %25 : !pto.tile_buf<acc, 32x256xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
      }
      %47 = pto.tpop_from_aiv {id = 30, split = 1} -> !pto.tile_buf<mat, 32x256xf16, blayout=col_major, slayout=row_major>
      pto.tmov ins(%47 : !pto.tile_buf<mat, 32x256xf16, blayout=col_major, slayout=row_major>) outs(%35 : !pto.tile_buf<left, 32x256xf16, slayout=row_major>)
      pto.tfree_from_aiv {id = 30, split = 1}
      pto.tmov ins(%36 : !pto.tile_buf<mat, 256x128xf16, blayout=col_major, slayout=row_major>) outs(%37 : !pto.tile_buf<right, 256x128xf16, slayout=col_major>)
      %c7936 = arith.constant 7936 : index
      %48 = pto.partition_view %41, offsets = [%c7936, %c0], sizes = [%c256, %c128] : !pto.tensor_view<?x?xf16>
      pto.tload ins(%48 : !pto.partition_tensor_view<256x128xf16>) outs(%36 : !pto.tile_buf<mat, 256x128xf16, blayout=col_major, slayout=row_major>)
      pto.tmatmul ins(%35, %37 : !pto.tile_buf<left, 32x256xf16, slayout=row_major>, !pto.tile_buf<right, 256x128xf16, slayout=col_major>) outs(%38 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>)
      pto.tpush(%38, %27 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
      %49 = pto.tpop_from_aiv {id = 30, split = 1} -> !pto.tile_buf<mat, 32x256xf16, blayout=col_major, slayout=row_major>
      pto.tmov ins(%49 : !pto.tile_buf<mat, 32x256xf16, blayout=col_major, slayout=row_major>) outs(%35 : !pto.tile_buf<left, 32x256xf16, slayout=row_major>)
      pto.tfree_from_aiv {id = 30, split = 1}
      pto.tmov ins(%36 : !pto.tile_buf<mat, 256x128xf16, blayout=col_major, slayout=row_major>) outs(%37 : !pto.tile_buf<right, 256x128xf16, slayout=col_major>)
      pto.tmatmul ins(%35, %37 : !pto.tile_buf<left, 32x256xf16, slayout=row_major>, !pto.tile_buf<right, 256x128xf16, slayout=col_major>) outs(%38 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>)
      pto.tpush(%38, %27 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
    }
    return
  }
  func.func @vector_kernel(%arg0: !pto.ptr<f32>, %arg1: !pto.ptr<f32>) attributes {pto.kernel_kind = #pto.kernel_kind<vector>} {
    %c0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    %c2 = arith.constant 2 : index
    %c32 = arith.constant 32 : index
    %c16 = arith.constant 16 : index
    %c128 = arith.constant 128 : index
    %c32_0 = arith.constant 32 : index
    %c96 = arith.constant 96 : index
    %0 = pto.get_block_num
    %1 = arith.index_cast %0 : i64 to index
    %2 = pto.get_block_idx
    %3 = arith.index_cast %2 : i64 to index
    %4 = arith.divsi %c96, %1 : index
    %5 = arith.remsi %c96, %1 : index
    %6 = arith.addi %4, %c1 : index
    %7 = arith.muli %3, %6 : index
    %8 = arith.addi %4, %c1 : index
    %9 = arith.muli %5, %8 : index
    %10 = arith.subi %3, %5 : index
    %11 = arith.muli %10, %4 : index
    %12 = arith.addi %9, %11 : index
    %13 = arith.cmpi slt, %3, %5 : index
    %14 = arith.select %13, %7, %12 : index
    %15 = arith.cmpi slt, %3, %5 : index
    %16 = arith.addi %4, %c1 : index
    %17 = arith.select %15, %16, %4 : index
    %18 = arith.addi %14, %17 : index
    %c131072 = arith.constant 131072 : index
    %19 = arith.muli %3, %c131072 : index
    %20 = pto.addptr %arg0, %19 : <f32> -> <f32>
    %c0_1 = arith.constant 0 : index
    %21 = pto.addptr %20, %c0_1 : <f32> -> <f32>
    %c65536 = arith.constant 65536 : index
    %22 = pto.addptr %20, %c65536 : <f32> -> <f32>
    %c98304 = arith.constant 98304 : index
    %23 = pto.addptr %20, %c98304 : <f32> -> <f32>
    %24 = pto.reserve_buffer{name = "fa_qk_c2v_fifo", size = 32768, location = <vec>, auto = false, base = 0} -> i32
    %25 = pto.initialize_l2g2l_pipe{dir_mask = 1, slot_size = 32768, slot_num = 8, local_slot_num = 1} (%21 : !pto.ptr<f32>, %24 : i32) -> !pto.pipe
    %26 = pto.reserve_buffer{name = "fa_pv_c2v_fifo", size = 16384, location = <vec>, auto = false, base = 32768} -> i32
    %27 = pto.initialize_l2g2l_pipe{dir_mask = 1, slot_size = 16384, slot_num = 8, local_slot_num = 1} (%22 : !pto.ptr<f32>, %26 : i32) -> !pto.pipe
    %28 = pto.import_reserved_buffer{name = "fa_p_v2c_fifo", peer_func = @cube_kernel} -> i32
    %c0_i32 = arith.constant 0 : i32
    pto.aiv_initialize_pipe {id = 30, dir_mask = 2, slot_size = 16384, local_slot_num = 1, nosplit = false}(gm_slot_buffer = %23 : !pto.ptr<f32>, c2v_consumer_buf = %c0_i32 : i32, v2c_consumer_buf = %28 : i32)
    %29 = pto.get_subblock_idx
    %30 = arith.index_cast %29 : i64 to index
    %31 = arith.muli %30, %c16 : index
    %c49152_i64 = arith.constant 49152 : i64
    %32 = pto.alloc_tile addr = %c49152_i64 : !pto.tile_buf<vec, 16x256xf32>
    %c65536_i64 = arith.constant 65536 : i64
    %33 = pto.alloc_tile addr = %c65536_i64 : !pto.tile_buf<vec, 16x256xf32>
    %c81920_i64 = arith.constant 81920 : i64
    %34 = pto.alloc_tile addr = %c81920_i64 : !pto.tile_buf<vec, 16x256xf16>
    %c90112_i64 = arith.constant 90112 : i64
    %35 = pto.alloc_tile addr = %c90112_i64 : !pto.tile_buf<vec, 16x128xf32>
    %c98304_i64 = arith.constant 98304 : i64
    %36 = pto.alloc_tile addr = %c98304_i64 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>
    %c98816_i64 = arith.constant 98816 : i64
    %37 = pto.alloc_tile addr = %c98816_i64 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>
    %c99328_i64 = arith.constant 99328 : i64
    %38 = pto.alloc_tile addr = %c99328_i64 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>
    %c99840_i64 = arith.constant 99840 : i64
    %39 = pto.alloc_tile addr = %c99840_i64 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>
    %c100352_i64 = arith.constant 100352 : i64
    %40 = pto.alloc_tile addr = %c100352_i64 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>
    %c100864_i64 = arith.constant 100864 : i64
    %41 = pto.alloc_tile addr = %c100864_i64 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>
    %cst = arith.constant 0.0883883461 : f32
    %cst_2 = arith.constant 1.000000e+00 : f32
    %c3072 = arith.constant 3072 : index
    %42 = pto.make_tensor_view %arg1, shape = [%c3072, %c128], strides = [%c128, %c1] : !pto.tensor_view<?x?xf32>
    scf.for %arg2 = %14 to %18 step %c1 {
      %43 = arith.muli %arg2, %c32 : index
      %c101376_i64 = arith.constant 101376 : i64
      %44 = pto.alloc_tile addr = %c101376_i64 : !pto.tile_buf<vec, 16x256xf32>
      pto.tpop(%44, %25 : !pto.tile_buf<vec, 16x256xf32>, !pto.pipe) {split = 1}
      pto.tmuls ins(%44, %cst : !pto.tile_buf<vec, 16x256xf32>, f32) outs(%44 : !pto.tile_buf<vec, 16x256xf32>)
      pto.trowmax ins(%44, %32 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x256xf32>) outs(%37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
      %45 = pto.treshape %37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %46 = pto.treshape %36 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %47 = pto.treshape %40 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %48 = pto.treshape %38 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %49 = pto.treshape %39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      pto.trowexpandsub ins(%44, %37 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%33 : !pto.tile_buf<vec, 16x256xf32>)
      pto.tmuls ins(%45, %cst_2 : !pto.tile_buf<vec, 1x16xf32>, f32) outs(%46 : !pto.tile_buf<vec, 1x16xf32>)
      pto.texp ins(%33 : !pto.tile_buf<vec, 16x256xf32>) outs(%33 : !pto.tile_buf<vec, 16x256xf32>)
      pto.trowsum ins(%33, %32 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x256xf32>) outs(%38 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
      pto.tcvt ins(%33 {rmode = #pto<round_mode CAST_RINT>} : !pto.tile_buf<vec, 16x256xf32>) outs(%34 : !pto.tile_buf<vec, 16x256xf16>)
      pto.tpush_to_aic(%34 : !pto.tile_buf<vec, 16x256xf16>) {id = 30, split = 1}
      pto.tfree(%25 : !pto.pipe) {split = 1}
      %c101376_i64_3 = arith.constant 101376 : i64
      %50 = pto.alloc_tile addr = %c101376_i64_3 : !pto.tile_buf<vec, 16x256xf32>
      pto.tpop(%50, %25 : !pto.tile_buf<vec, 16x256xf32>, !pto.pipe) {split = 1}
      pto.tmuls ins(%50, %cst : !pto.tile_buf<vec, 16x256xf32>, f32) outs(%50 : !pto.tile_buf<vec, 16x256xf32>)
      pto.trowmax ins(%50, %32 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x256xf32>) outs(%37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
      %51 = pto.treshape %37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %52 = pto.treshape %36 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %53 = pto.treshape %41 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %54 = pto.treshape %38 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %55 = pto.treshape %39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      pto.tmax ins(%51, %52 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%51 : !pto.tile_buf<vec, 1x16xf32>)
      pto.tsub ins(%52, %51 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%53 : !pto.tile_buf<vec, 1x16xf32>)
      pto.tmuls ins(%51, %cst_2 : !pto.tile_buf<vec, 1x16xf32>, f32) outs(%52 : !pto.tile_buf<vec, 1x16xf32>)
      pto.trowexpandsub ins(%50, %37 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%33 : !pto.tile_buf<vec, 16x256xf32>)
      pto.texp ins(%53 : !pto.tile_buf<vec, 1x16xf32>) outs(%53 : !pto.tile_buf<vec, 1x16xf32>)
      pto.texp ins(%33 : !pto.tile_buf<vec, 16x256xf32>) outs(%33 : !pto.tile_buf<vec, 16x256xf32>)
      pto.tmul ins(%54, %53 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%54 : !pto.tile_buf<vec, 1x16xf32>)
      pto.trowsum ins(%33, %32 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x256xf32>) outs(%39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
      pto.tadd ins(%54, %55 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%54 : !pto.tile_buf<vec, 1x16xf32>)
      pto.tcvt ins(%33 {rmode = #pto<round_mode CAST_RINT>} : !pto.tile_buf<vec, 16x256xf32>) outs(%34 : !pto.tile_buf<vec, 16x256xf16>)
      pto.tpush_to_aic(%34 : !pto.tile_buf<vec, 16x256xf16>) {id = 30, split = 1}
      pto.tfree(%25 : !pto.pipe) {split = 1}
      %c101376_i64_4 = arith.constant 101376 : i64
      %56 = pto.alloc_tile addr = %c101376_i64_4 : !pto.tile_buf<vec, 16x128xf32>
      pto.tpop(%56, %27 : !pto.tile_buf<vec, 16x128xf32>, !pto.pipe) {split = 1}
      pto.tmov ins(%56 : !pto.tile_buf<vec, 16x128xf32>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
      pto.tfree(%27 : !pto.pipe) {split = 1}
      %c101376_i64_5 = arith.constant 101376 : i64
      %57 = pto.alloc_tile addr = %c101376_i64_5 : !pto.tile_buf<vec, 16x256xf32>
      pto.tpop(%57, %25 : !pto.tile_buf<vec, 16x256xf32>, !pto.pipe) {split = 1}
      pto.tmuls ins(%57, %cst : !pto.tile_buf<vec, 16x256xf32>, f32) outs(%57 : !pto.tile_buf<vec, 16x256xf32>)
      pto.trowmax ins(%57, %32 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x256xf32>) outs(%37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
      %58 = pto.treshape %37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %59 = pto.treshape %36 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %60 = pto.treshape %40 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %61 = pto.treshape %38 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %62 = pto.treshape %39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      pto.tmax ins(%58, %59 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%58 : !pto.tile_buf<vec, 1x16xf32>)
      pto.tsub ins(%59, %58 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%60 : !pto.tile_buf<vec, 1x16xf32>)
      pto.tmuls ins(%58, %cst_2 : !pto.tile_buf<vec, 1x16xf32>, f32) outs(%59 : !pto.tile_buf<vec, 1x16xf32>)
      pto.trowexpandsub ins(%57, %37 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%33 : !pto.tile_buf<vec, 16x256xf32>)
      pto.texp ins(%60 : !pto.tile_buf<vec, 1x16xf32>) outs(%60 : !pto.tile_buf<vec, 1x16xf32>)
      pto.texp ins(%33 : !pto.tile_buf<vec, 16x256xf32>) outs(%33 : !pto.tile_buf<vec, 16x256xf32>)
      pto.tmul ins(%61, %60 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%61 : !pto.tile_buf<vec, 1x16xf32>)
      pto.trowsum ins(%33, %32 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x256xf32>) outs(%39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
      pto.tadd ins(%61, %62 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%61 : !pto.tile_buf<vec, 1x16xf32>)
      pto.tcvt ins(%33 {rmode = #pto<round_mode CAST_RINT>} : !pto.tile_buf<vec, 16x256xf32>) outs(%34 : !pto.tile_buf<vec, 16x256xf16>)
      pto.tpush_to_aic(%34 : !pto.tile_buf<vec, 16x256xf16>) {id = 30, split = 1}
      pto.tfree(%25 : !pto.pipe) {split = 1}
      %c101376_i64_6 = arith.constant 101376 : i64
      %63 = pto.alloc_tile addr = %c101376_i64_6 : !pto.tile_buf<vec, 16x128xf32>
      pto.tpop(%63, %27 : !pto.tile_buf<vec, 16x128xf32>, !pto.pipe) {split = 1}
      pto.trowexpandmul ins(%35, %41 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
      pto.tadd ins(%35, %63 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x128xf32>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
      pto.tfree(%27 : !pto.pipe) {split = 1}
      %c101376_i64_7 = arith.constant 101376 : i64
      %64 = pto.alloc_tile addr = %c101376_i64_7 : !pto.tile_buf<vec, 16x256xf32>
      pto.tpop(%64, %25 : !pto.tile_buf<vec, 16x256xf32>, !pto.pipe) {split = 1}
      pto.tmuls ins(%64, %cst : !pto.tile_buf<vec, 16x256xf32>, f32) outs(%64 : !pto.tile_buf<vec, 16x256xf32>)
      pto.trowmax ins(%64, %32 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x256xf32>) outs(%37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
      %65 = pto.treshape %37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %66 = pto.treshape %36 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %67 = pto.treshape %41 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %68 = pto.treshape %38 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %69 = pto.treshape %39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      pto.tmax ins(%65, %66 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%65 : !pto.tile_buf<vec, 1x16xf32>)
      pto.tsub ins(%66, %65 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%67 : !pto.tile_buf<vec, 1x16xf32>)
      pto.tmuls ins(%65, %cst_2 : !pto.tile_buf<vec, 1x16xf32>, f32) outs(%66 : !pto.tile_buf<vec, 1x16xf32>)
      pto.trowexpandsub ins(%64, %37 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%33 : !pto.tile_buf<vec, 16x256xf32>)
      pto.texp ins(%67 : !pto.tile_buf<vec, 1x16xf32>) outs(%67 : !pto.tile_buf<vec, 1x16xf32>)
      pto.texp ins(%33 : !pto.tile_buf<vec, 16x256xf32>) outs(%33 : !pto.tile_buf<vec, 16x256xf32>)
      pto.tmul ins(%68, %67 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%68 : !pto.tile_buf<vec, 1x16xf32>)
      pto.trowsum ins(%33, %32 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x256xf32>) outs(%39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
      pto.tadd ins(%68, %69 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%68 : !pto.tile_buf<vec, 1x16xf32>)
      pto.tcvt ins(%33 {rmode = #pto<round_mode CAST_RINT>} : !pto.tile_buf<vec, 16x256xf32>) outs(%34 : !pto.tile_buf<vec, 16x256xf16>)
      pto.tpush_to_aic(%34 : !pto.tile_buf<vec, 16x256xf16>) {id = 30, split = 1}
      pto.tfree(%25 : !pto.pipe) {split = 1}
      %c15 = arith.constant 15 : index
      scf.for %arg3 = %c1 to %c15 step %c1 {
        %c101376_i64_10 = arith.constant 101376 : i64
        %74 = pto.alloc_tile addr = %c101376_i64_10 : !pto.tile_buf<vec, 16x128xf32>
        pto.tpop(%74, %27 : !pto.tile_buf<vec, 16x128xf32>, !pto.pipe) {split = 1}
        pto.trowexpandmul ins(%35, %40 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
        pto.tadd ins(%35, %74 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x128xf32>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
        pto.tfree(%27 : !pto.pipe) {split = 1}
        %c101376_i64_11 = arith.constant 101376 : i64
        %75 = pto.alloc_tile addr = %c101376_i64_11 : !pto.tile_buf<vec, 16x256xf32>
        pto.tpop(%75, %25 : !pto.tile_buf<vec, 16x256xf32>, !pto.pipe) {split = 1}
        pto.tmuls ins(%75, %cst : !pto.tile_buf<vec, 16x256xf32>, f32) outs(%75 : !pto.tile_buf<vec, 16x256xf32>)
        pto.trowmax ins(%75, %32 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x256xf32>) outs(%37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
        %76 = pto.treshape %37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
        %77 = pto.treshape %36 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
        %78 = pto.treshape %40 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
        %79 = pto.treshape %38 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
        %80 = pto.treshape %39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
        pto.tmax ins(%76, %77 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%76 : !pto.tile_buf<vec, 1x16xf32>)
        pto.tsub ins(%77, %76 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%78 : !pto.tile_buf<vec, 1x16xf32>)
        pto.tmuls ins(%76, %cst_2 : !pto.tile_buf<vec, 1x16xf32>, f32) outs(%77 : !pto.tile_buf<vec, 1x16xf32>)
        pto.trowexpandsub ins(%75, %37 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%33 : !pto.tile_buf<vec, 16x256xf32>)
        pto.texp ins(%78 : !pto.tile_buf<vec, 1x16xf32>) outs(%78 : !pto.tile_buf<vec, 1x16xf32>)
        pto.texp ins(%33 : !pto.tile_buf<vec, 16x256xf32>) outs(%33 : !pto.tile_buf<vec, 16x256xf32>)
        pto.tmul ins(%79, %78 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%79 : !pto.tile_buf<vec, 1x16xf32>)
        pto.trowsum ins(%33, %32 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x256xf32>) outs(%39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
        pto.tadd ins(%79, %80 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%79 : !pto.tile_buf<vec, 1x16xf32>)
        pto.tcvt ins(%33 {rmode = #pto<round_mode CAST_RINT>} : !pto.tile_buf<vec, 16x256xf32>) outs(%34 : !pto.tile_buf<vec, 16x256xf16>)
        pto.tpush_to_aic(%34 : !pto.tile_buf<vec, 16x256xf16>) {id = 30, split = 1}
        pto.tfree(%25 : !pto.pipe) {split = 1}
        %c101376_i64_12 = arith.constant 101376 : i64
        %81 = pto.alloc_tile addr = %c101376_i64_12 : !pto.tile_buf<vec, 16x128xf32>
        pto.tpop(%81, %27 : !pto.tile_buf<vec, 16x128xf32>, !pto.pipe) {split = 1}
        pto.trowexpandmul ins(%35, %41 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
        pto.tadd ins(%35, %81 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x128xf32>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
        pto.tfree(%27 : !pto.pipe) {split = 1}
        %c101376_i64_13 = arith.constant 101376 : i64
        %82 = pto.alloc_tile addr = %c101376_i64_13 : !pto.tile_buf<vec, 16x256xf32>
        pto.tpop(%82, %25 : !pto.tile_buf<vec, 16x256xf32>, !pto.pipe) {split = 1}
        pto.tmuls ins(%82, %cst : !pto.tile_buf<vec, 16x256xf32>, f32) outs(%82 : !pto.tile_buf<vec, 16x256xf32>)
        pto.trowmax ins(%82, %32 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x256xf32>) outs(%37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
        %83 = pto.treshape %37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
        %84 = pto.treshape %36 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
        %85 = pto.treshape %41 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
        %86 = pto.treshape %38 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
        %87 = pto.treshape %39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
        pto.tmax ins(%83, %84 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%83 : !pto.tile_buf<vec, 1x16xf32>)
        pto.tsub ins(%84, %83 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%85 : !pto.tile_buf<vec, 1x16xf32>)
        pto.tmuls ins(%83, %cst_2 : !pto.tile_buf<vec, 1x16xf32>, f32) outs(%84 : !pto.tile_buf<vec, 1x16xf32>)
        pto.trowexpandsub ins(%82, %37 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%33 : !pto.tile_buf<vec, 16x256xf32>)
        pto.texp ins(%85 : !pto.tile_buf<vec, 1x16xf32>) outs(%85 : !pto.tile_buf<vec, 1x16xf32>)
        pto.texp ins(%33 : !pto.tile_buf<vec, 16x256xf32>) outs(%33 : !pto.tile_buf<vec, 16x256xf32>)
        pto.tmul ins(%86, %85 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%86 : !pto.tile_buf<vec, 1x16xf32>)
        pto.trowsum ins(%33, %32 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x256xf32>) outs(%39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
        pto.tadd ins(%86, %87 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%86 : !pto.tile_buf<vec, 1x16xf32>)
        pto.tcvt ins(%33 {rmode = #pto<round_mode CAST_RINT>} : !pto.tile_buf<vec, 16x256xf32>) outs(%34 : !pto.tile_buf<vec, 16x256xf16>)
        pto.tpush_to_aic(%34 : !pto.tile_buf<vec, 16x256xf16>) {id = 30, split = 1}
        pto.tfree(%25 : !pto.pipe) {split = 1}
      }
      %c101376_i64_8 = arith.constant 101376 : i64
      %70 = pto.alloc_tile addr = %c101376_i64_8 : !pto.tile_buf<vec, 16x128xf32>
      pto.tpop(%70, %27 : !pto.tile_buf<vec, 16x128xf32>, !pto.pipe) {split = 1}
      pto.trowexpandmul ins(%35, %40 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
      pto.tadd ins(%35, %70 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x128xf32>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
      pto.tfree(%27 : !pto.pipe) {split = 1}
      %c101376_i64_9 = arith.constant 101376 : i64
      %71 = pto.alloc_tile addr = %c101376_i64_9 : !pto.tile_buf<vec, 16x128xf32>
      pto.tpop(%71, %27 : !pto.tile_buf<vec, 16x128xf32>, !pto.pipe) {split = 1}
      pto.trowexpandmul ins(%35, %41 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
      pto.tadd ins(%35, %71 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x128xf32>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
      pto.tfree(%27 : !pto.pipe) {split = 1}
      pto.trowexpanddiv ins(%35, %38 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
      %72 = arith.addi %43, %31 : index
      %73 = pto.partition_view %42, offsets = [%72, %c0], sizes = [%c16, %c128] : !pto.tensor_view<?x?xf32>
      pto.tstore ins(%35 : !pto.tile_buf<vec, 16x128xf32>) outs(%73 : !pto.partition_tensor_view<16x128xf32>)
    }
    return
  }
  func.func @call_both(%arg0: memref<256xi64>, %arg1: !pto.ptr<f32>, %arg2: !pto.ptr<f16>, %arg3: !pto.ptr<f16>, %arg4: !pto.ptr<f16>, %arg5: !pto.ptr<f32>) attributes {pto.entry} {
    pto.set_ffts %arg0 : memref<256xi64>
    call @cube_kernel(%arg1, %arg2, %arg3, %arg4) : (!pto.ptr<f32>, !pto.ptr<f16>, !pto.ptr<f16>, !pto.ptr<f16>) -> ()
    call @vector_kernel(%arg1, %arg5) : (!pto.ptr<f32>, !pto.ptr<f32>) -> ()
    return
  }
}

Actual behavior / error logs

PTOAS and bisheng both complete successfully, but the first device synchronize after launching the kernel fails:


RuntimeError: npuSynchronizeDevice: ... AclrtSynchronizeDeviceWithTimeout, error code is 507015
[Error]: The aicore execution is abnormal.
EZ9999: ... exception of fftsplus aicore error
errorStr: CCU instruction address check error
retCode=0x26, [aicore exception]

Git commit

150abd8

Host platform

Linux (aarch64)

Target Ascend arch (if relevant)

a3

PTOAS build level (if relevant)

level3

[zhangstevenunity]

how did you modify S1_TILE，please attach the pto file.

[MirkoDeVita98]

how did you modify S1_TILE，please attach the pto file.

See the expected MLIR above in the expected behaviour section. The diff between that and the S1_TILE=512 is the following:

--- examples/aot/flash_attention/experimental/compiler_team_fa.pto	2026-05-05 15:47:55.590418546 +0200
+++ examples/aot/flash_attention/experimental/ptoas_issue_artifacts/fa_compiler_team_patched_s256_q3072_s8192_pass.pto	2026-05-05 15:56:12.725459562 +0200
@@ -5,9 +5,9 @@
     %c2 = arith.constant 2 : index
     %c32 = arith.constant 32 : index
     %c128 = arith.constant 128 : index
-    %c512 = arith.constant 512 : index
+    %c256 = arith.constant 256 : index
     %c8192 = arith.constant 8192 : index
-    %c16 = arith.constant 16 : index
+    %c32_0 = arith.constant 32 : index
     %c96 = arith.constant 96 : index
     %0 = pto.get_block_num
     %1 = arith.index_cast %0 : i64 to index
@@ -28,41 +28,41 @@
     %16 = arith.addi %4, %c1 : index
     %17 = arith.select %15, %16, %4 : index
     %18 = arith.addi %14, %17 : index
-    %c229376 = arith.constant 229376 : index
-    %19 = arith.muli %3, %c229376 : index
-    %20 = pto.addptr %arg0, %19 : <f32> -> <f32>
-    %c0_0 = arith.constant 0 : index
-    %21 = pto.addptr %20, %c0_0 : <f32> -> <f32>
     %c131072 = arith.constant 131072 : index
-    %22 = pto.addptr %20, %c131072 : <f32> -> <f32>
-    %c163840 = arith.constant 163840 : index
-    %23 = pto.addptr %20, %c163840 : <f32> -> <f32>
+    %19 = arith.muli %3, %c131072 : index
+    %20 = pto.addptr %arg0, %19 : <f32> -> <f32>
+    %c0_1 = arith.constant 0 : index
+    %21 = pto.addptr %20, %c0_1 : <f32> -> <f32>
+    %c65536 = arith.constant 65536 : index
+    %22 = pto.addptr %20, %c65536 : <f32> -> <f32>
+    %c98304 = arith.constant 98304 : index
+    %23 = pto.addptr %20, %c98304 : <f32> -> <f32>
     %24 = pto.import_reserved_buffer{name = "fa_qk_c2v_fifo", peer_func = @vector_kernel} -> i32
-    %25 = pto.initialize_l2g2l_pipe{dir_mask = 1, slot_size = 65536, slot_num = 8, local_slot_num = 1} (%21 : !pto.ptr<f32>, %24 : i32) -> !pto.pipe
+    %25 = pto.initialize_l2g2l_pipe{dir_mask = 1, slot_size = 32768, slot_num = 8, local_slot_num = 1} (%21 : !pto.ptr<f32>, %24 : i32) -> !pto.pipe
     %26 = pto.import_reserved_buffer{name = "fa_pv_c2v_fifo", peer_func = @vector_kernel} -> i32
     %27 = pto.initialize_l2g2l_pipe{dir_mask = 1, slot_size = 16384, slot_num = 8, local_slot_num = 1} (%22 : !pto.ptr<f32>, %26 : i32) -> !pto.pipe
-    %28 = pto.reserve_buffer{name = "fa_p_v2c_fifo", size = 262144, location = <mat>, auto = false, base = 327680} -> i32
+    %28 = pto.reserve_buffer{name = "fa_p_v2c_fifo", size = 16384, location = <mat>, auto = false, base = 262144} -> i32
     %c0_i32 = arith.constant 0 : i32
-    pto.aic_initialize_pipe {id = 30, dir_mask = 2, slot_size = 32768, nosplit = false}(gm_slot_buffer = %23 : !pto.ptr<f32>, c2v_consumer_buf = %c0_i32 : i32, v2c_consumer_buf = %28 : i32)
+    pto.aic_initialize_pipe {id = 30, dir_mask = 2, slot_size = 16384, local_slot_num = 1, nosplit = false}(gm_slot_buffer = %23 : !pto.ptr<f32>, c2v_consumer_buf = %c0_i32 : i32, v2c_consumer_buf = %28 : i32)
     %c0_i64 = arith.constant 0 : i64
-    %c0_i64_1 = arith.constant 0 : i64
-    %29 = pto.alloc_tile addr = %c0_i64_1 : !pto.tile_buf<mat, 32x128xf16, blayout=col_major, slayout=row_major>
     %c0_i64_2 = arith.constant 0 : i64
-    %30 = pto.alloc_tile addr = %c0_i64_2 : !pto.tile_buf<left, 32x128xf16, slayout=row_major>
-    %c8192_i64 = arith.constant 8192 : i64
-    %31 = pto.alloc_tile addr = %c8192_i64 : !pto.tile_buf<mat, 128x512xf16, slayout=col_major>
-    %32 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf<right, 128x512xf16, slayout=col_major>
+    %29 = pto.alloc_tile addr = %c0_i64_2 : !pto.tile_buf<mat, 32x128xf16, blayout=col_major, slayout=row_major>
     %c0_i64_3 = arith.constant 0 : i64
-    %33 = pto.alloc_tile addr = %c0_i64_3 : !pto.tile_buf<acc, 32x512xf32, blayout=col_major, slayout=row_major, fractal=1024>
-    %c139264_i64 = arith.constant 139264 : i64
-    %34 = pto.alloc_tile addr = %c139264_i64 : !pto.tile_buf<mat, 32x512xf16, blayout=col_major, slayout=row_major>
-    %c8192_i64_4 = arith.constant 8192 : i64
-    %35 = pto.alloc_tile addr = %c8192_i64_4 : !pto.tile_buf<left, 32x512xf16, slayout=row_major>
-    %c172032_i64 = arith.constant 172032 : i64
-    %36 = pto.alloc_tile addr = %c172032_i64 : !pto.tile_buf<mat, 512x128xf16, blayout=col_major, slayout=row_major>
-    %37 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf<right, 512x128xf16, slayout=col_major>
-    %c65536_i64 = arith.constant 65536 : i64
-    %38 = pto.alloc_tile addr = %c65536_i64 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>
+    %30 = pto.alloc_tile addr = %c0_i64_3 : !pto.tile_buf<left, 32x128xf16, slayout=row_major>
+    %c8192_i64 = arith.constant 8192 : i64
+    %31 = pto.alloc_tile addr = %c8192_i64 : !pto.tile_buf<mat, 128x256xf16, slayout=col_major>
+    %32 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf<right, 128x256xf16, slayout=col_major>
+    %c0_i64_4 = arith.constant 0 : i64
+    %33 = pto.alloc_tile addr = %c0_i64_4 : !pto.tile_buf<acc, 32x256xf32, blayout=col_major, slayout=row_major, fractal=1024>
+    %c73728_i64 = arith.constant 73728 : i64
+    %34 = pto.alloc_tile addr = %c73728_i64 : !pto.tile_buf<mat, 32x256xf16, blayout=col_major, slayout=row_major>
+    %c8192_i64_5 = arith.constant 8192 : i64
+    %35 = pto.alloc_tile addr = %c8192_i64_5 : !pto.tile_buf<left, 32x256xf16, slayout=row_major>
+    %c90112_i64 = arith.constant 90112 : i64
+    %36 = pto.alloc_tile addr = %c90112_i64 : !pto.tile_buf<mat, 256x128xf16, blayout=col_major, slayout=row_major>
+    %37 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf<right, 256x128xf16, slayout=col_major>
+    %c32768_i64 = arith.constant 32768 : i64
+    %38 = pto.alloc_tile addr = %c32768_i64 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>
     %c3072 = arith.constant 3072 : index
     %39 = pto.make_tensor_view %arg1, shape = [%c3072, %c128], strides = [%c128, %c1] : !pto.tensor_view<?x?xf16>
     %40 = pto.make_tensor_view %arg2, shape = [%c128, %c8192], strides = [%c1, %c128] : !pto.tensor_view<?x?xf16>
@@ -72,76 +72,76 @@
       %43 = pto.partition_view %39, offsets = [%42, %c0], sizes = [%c32, %c128] : !pto.tensor_view<?x?xf16>
       pto.tload ins(%43 : !pto.partition_tensor_view<32x128xf16>) outs(%29 : !pto.tile_buf<mat, 32x128xf16, blayout=col_major, slayout=row_major>)
       pto.tmov ins(%29 : !pto.tile_buf<mat, 32x128xf16, blayout=col_major, slayout=row_major>) outs(%30 : !pto.tile_buf<left, 32x128xf16, slayout=row_major>)
-      %c0_5 = arith.constant 0 : index
-      %44 = pto.partition_view %40, offsets = [%c0, %c0_5], sizes = [%c128, %c512] : !pto.tensor_view<?x?xf16>
-      pto.tload ins(%44 : !pto.partition_tensor_view<128x512xf16>) outs(%31 : !pto.tile_buf<mat, 128x512xf16, slayout=col_major>)
-      pto.tmov ins(%31 : !pto.tile_buf<mat, 128x512xf16, slayout=col_major>) outs(%32 : !pto.tile_buf<right, 128x512xf16, slayout=col_major>)
-      pto.tmatmul ins(%30, %32 : !pto.tile_buf<left, 32x128xf16, slayout=row_major>, !pto.tile_buf<right, 128x512xf16, slayout=col_major>) outs(%33 : !pto.tile_buf<acc, 32x512xf32, blayout=col_major, slayout=row_major, fractal=1024>)
-      pto.tpush(%33, %25 : !pto.tile_buf<acc, 32x512xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
-      %c512_6 = arith.constant 512 : index
-      %45 = pto.partition_view %40, offsets = [%c0, %c512_6], sizes = [%c128, %c512] : !pto.tensor_view<?x?xf16>
-      pto.tload ins(%45 : !pto.partition_tensor_view<128x512xf16>) outs(%31 : !pto.tile_buf<mat, 128x512xf16, slayout=col_major>)
-      pto.tmov ins(%31 : !pto.tile_buf<mat, 128x512xf16, slayout=col_major>) outs(%32 : !pto.tile_buf<right, 128x512xf16, slayout=col_major>)
-      pto.tmatmul ins(%30, %32 : !pto.tile_buf<left, 32x128xf16, slayout=row_major>, !pto.tile_buf<right, 128x512xf16, slayout=col_major>) outs(%33 : !pto.tile_buf<acc, 32x512xf32, blayout=col_major, slayout=row_major, fractal=1024>)
-      pto.tpush(%33, %25 : !pto.tile_buf<acc, 32x512xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
-      %46 = pto.partition_view %41, offsets = [%c0, %c0], sizes = [%c512, %c128] : !pto.tensor_view<?x?xf16>
-      pto.tload ins(%46 : !pto.partition_tensor_view<512x128xf16>) outs(%36 : !pto.tile_buf<mat, 512x128xf16, blayout=col_major, slayout=row_major>)
-      %c7 = arith.constant 7 : index
-      scf.for %arg5 = %c0 to %c7 step %c1 {
+      %c0_6 = arith.constant 0 : index
+      %44 = pto.partition_view %40, offsets = [%c0, %c0_6], sizes = [%c128, %c256] : !pto.tensor_view<?x?xf16>
+      pto.tload ins(%44 : !pto.partition_tensor_view<128x256xf16>) outs(%31 : !pto.tile_buf<mat, 128x256xf16, slayout=col_major>)
+      pto.tmov ins(%31 : !pto.tile_buf<mat, 128x256xf16, slayout=col_major>) outs(%32 : !pto.tile_buf<right, 128x256xf16, slayout=col_major>)
+      pto.tmatmul ins(%30, %32 : !pto.tile_buf<left, 32x128xf16, slayout=row_major>, !pto.tile_buf<right, 128x256xf16, slayout=col_major>) outs(%33 : !pto.tile_buf<acc, 32x256xf32, blayout=col_major, slayout=row_major, fractal=1024>)
+      pto.tpush(%33, %25 : !pto.tile_buf<acc, 32x256xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
+      %c256_7 = arith.constant 256 : index
+      %45 = pto.partition_view %40, offsets = [%c0, %c256_7], sizes = [%c128, %c256] : !pto.tensor_view<?x?xf16>
+      pto.tload ins(%45 : !pto.partition_tensor_view<128x256xf16>) outs(%31 : !pto.tile_buf<mat, 128x256xf16, slayout=col_major>)
+      pto.tmov ins(%31 : !pto.tile_buf<mat, 128x256xf16, slayout=col_major>) outs(%32 : !pto.tile_buf<right, 128x256xf16, slayout=col_major>)
+      pto.tmatmul ins(%30, %32 : !pto.tile_buf<left, 32x128xf16, slayout=row_major>, !pto.tile_buf<right, 128x256xf16, slayout=col_major>) outs(%33 : !pto.tile_buf<acc, 32x256xf32, blayout=col_major, slayout=row_major, fractal=1024>)
+      pto.tpush(%33, %25 : !pto.tile_buf<acc, 32x256xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
+      %46 = pto.partition_view %41, offsets = [%c0, %c0], sizes = [%c256, %c128] : !pto.tensor_view<?x?xf16>
+      pto.tload ins(%46 : !pto.partition_tensor_view<256x128xf16>) outs(%36 : !pto.tile_buf<mat, 256x128xf16, blayout=col_major, slayout=row_major>)
+      %c15 = arith.constant 15 : index
+      scf.for %arg5 = %c0 to %c15 step %c1 {
         %50 = arith.muli %arg5, %c2 : index
-        %c2_7 = arith.constant 2 : index
-        %51 = arith.addi %50, %c2_7 : index
-        %52 = arith.muli %51, %c512 : index
-        %53 = pto.partition_view %40, offsets = [%c0, %52], sizes = [%c128, %c512] : !pto.tensor_view<?x?xf16>
-        pto.tload ins(%53 : !pto.partition_tensor_view<128x512xf16>) outs(%31 : !pto.tile_buf<mat, 128x512xf16, slayout=col_major>)
-        %54 = pto.tpop_from_aiv {id = 30, split = 1} -> !pto.tile_buf<mat, 32x512xf16, blayout=col_major, slayout=row_major>
-        pto.tmov ins(%54 : !pto.tile_buf<mat, 32x512xf16, blayout=col_major, slayout=row_major>) outs(%35 : !pto.tile_buf<left, 32x512xf16, slayout=row_major>)
+        %c2_8 = arith.constant 2 : index
+        %51 = arith.addi %50, %c2_8 : index
+        %52 = arith.muli %51, %c256 : index
+        %53 = pto.partition_view %40, offsets = [%c0, %52], sizes = [%c128, %c256] : !pto.tensor_view<?x?xf16>
+        pto.tload ins(%53 : !pto.partition_tensor_view<128x256xf16>) outs(%31 : !pto.tile_buf<mat, 128x256xf16, slayout=col_major>)
+        %54 = pto.tpop_from_aiv {id = 30, split = 1} -> !pto.tile_buf<mat, 32x256xf16, blayout=col_major, slayout=row_major>
+        pto.tmov ins(%54 : !pto.tile_buf<mat, 32x256xf16, blayout=col_major, slayout=row_major>) outs(%35 : !pto.tile_buf<left, 32x256xf16, slayout=row_major>)
         pto.tfree_from_aiv {id = 30, split = 1}
-        pto.tmov ins(%36 : !pto.tile_buf<mat, 512x128xf16, blayout=col_major, slayout=row_major>) outs(%37 : !pto.tile_buf<right, 512x128xf16, slayout=col_major>)
+        pto.tmov ins(%36 : !pto.tile_buf<mat, 256x128xf16, blayout=col_major, slayout=row_major>) outs(%37 : !pto.tile_buf<right, 256x128xf16, slayout=col_major>)
         %55 = arith.addi %50, %c1 : index
-        %56 = arith.muli %55, %c512 : index
-        %57 = pto.partition_view %41, offsets = [%56, %c0], sizes = [%c512, %c128] : !pto.tensor_view<?x?xf16>
-        pto.tload ins(%57 : !pto.partition_tensor_view<512x128xf16>) outs(%36 : !pto.tile_buf<mat, 512x128xf16, blayout=col_major, slayout=row_major>)
-        pto.tmatmul ins(%35, %37 : !pto.tile_buf<left, 32x512xf16, slayout=row_major>, !pto.tile_buf<right, 512x128xf16, slayout=col_major>) outs(%38 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>)
+        %56 = arith.muli %55, %c256 : index
+        %57 = pto.partition_view %41, offsets = [%56, %c0], sizes = [%c256, %c128] : !pto.tensor_view<?x?xf16>
+        pto.tload ins(%57 : !pto.partition_tensor_view<256x128xf16>) outs(%36 : !pto.tile_buf<mat, 256x128xf16, blayout=col_major, slayout=row_major>)
+        pto.tmatmul ins(%35, %37 : !pto.tile_buf<left, 32x256xf16, slayout=row_major>, !pto.tile_buf<right, 256x128xf16, slayout=col_major>) outs(%38 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>)
         pto.tpush(%38, %27 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
-        pto.tmov ins(%31 : !pto.tile_buf<mat, 128x512xf16, slayout=col_major>) outs(%32 : !pto.tile_buf<right, 128x512xf16, slayout=col_major>)
-        pto.tmatmul ins(%30, %32 : !pto.tile_buf<left, 32x128xf16, slayout=row_major>, !pto.tile_buf<right, 128x512xf16, slayout=col_major>) outs(%33 : !pto.tile_buf<acc, 32x512xf32, blayout=col_major, slayout=row_major, fractal=1024>)
-        pto.tpush(%33, %25 : !pto.tile_buf<acc, 32x512xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
+        pto.tmov ins(%31 : !pto.tile_buf<mat, 128x256xf16, slayout=col_major>) outs(%32 : !pto.tile_buf<right, 128x256xf16, slayout=col_major>)
+        pto.tmatmul ins(%30, %32 : !pto.tile_buf<left, 32x128xf16, slayout=row_major>, !pto.tile_buf<right, 128x256xf16, slayout=col_major>) outs(%33 : !pto.tile_buf<acc, 32x256xf32, blayout=col_major, slayout=row_major, fractal=1024>)
+        pto.tpush(%33, %25 : !pto.tile_buf<acc, 32x256xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
         %58 = arith.muli %arg5, %c2 : index
         %59 = arith.addi %58, %c1 : index
-        %c2_8 = arith.constant 2 : index
-        %60 = arith.addi %59, %c2_8 : index
-        %61 = arith.muli %60, %c512 : index
-        %62 = pto.partition_view %40, offsets = [%c0, %61], sizes = [%c128, %c512] : !pto.tensor_view<?x?xf16>
-        pto.tload ins(%62 : !pto.partition_tensor_view<128x512xf16>) outs(%31 : !pto.tile_buf<mat, 128x512xf16, slayout=col_major>)
-        %63 = pto.tpop_from_aiv {id = 30, split = 1} -> !pto.tile_buf<mat, 32x512xf16, blayout=col_major, slayout=row_major>
-        pto.tmov ins(%63 : !pto.tile_buf<mat, 32x512xf16, blayout=col_major, slayout=row_major>) outs(%35 : !pto.tile_buf<left, 32x512xf16, slayout=row_major>)
+        %c2_9 = arith.constant 2 : index
+        %60 = arith.addi %59, %c2_9 : index
+        %61 = arith.muli %60, %c256 : index
+        %62 = pto.partition_view %40, offsets = [%c0, %61], sizes = [%c128, %c256] : !pto.tensor_view<?x?xf16>
+        pto.tload ins(%62 : !pto.partition_tensor_view<128x256xf16>) outs(%31 : !pto.tile_buf<mat, 128x256xf16, slayout=col_major>)
+        %63 = pto.tpop_from_aiv {id = 30, split = 1} -> !pto.tile_buf<mat, 32x256xf16, blayout=col_major, slayout=row_major>
+        pto.tmov ins(%63 : !pto.tile_buf<mat, 32x256xf16, blayout=col_major, slayout=row_major>) outs(%35 : !pto.tile_buf<left, 32x256xf16, slayout=row_major>)
         pto.tfree_from_aiv {id = 30, split = 1}
-        pto.tmov ins(%36 : !pto.tile_buf<mat, 512x128xf16, blayout=col_major, slayout=row_major>) outs(%37 : !pto.tile_buf<right, 512x128xf16, slayout=col_major>)
+        pto.tmov ins(%36 : !pto.tile_buf<mat, 256x128xf16, blayout=col_major, slayout=row_major>) outs(%37 : !pto.tile_buf<right, 256x128xf16, slayout=col_major>)
         %64 = arith.addi %59, %c1 : index
-        %65 = arith.muli %64, %c512 : index
-        %66 = pto.partition_view %41, offsets = [%65, %c0], sizes = [%c512, %c128] : !pto.tensor_view<?x?xf16>
-        pto.tload ins(%66 : !pto.partition_tensor_view<512x128xf16>) outs(%36 : !pto.tile_buf<mat, 512x128xf16, blayout=col_major, slayout=row_major>)
-        pto.tmatmul ins(%35, %37 : !pto.tile_buf<left, 32x512xf16, slayout=row_major>, !pto.tile_buf<right, 512x128xf16, slayout=col_major>) outs(%38 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>)
+        %65 = arith.muli %64, %c256 : index
+        %66 = pto.partition_view %41, offsets = [%65, %c0], sizes = [%c256, %c128] : !pto.tensor_view<?x?xf16>
+        pto.tload ins(%66 : !pto.partition_tensor_view<256x128xf16>) outs(%36 : !pto.tile_buf<mat, 256x128xf16, blayout=col_major, slayout=row_major>)
+        pto.tmatmul ins(%35, %37 : !pto.tile_buf<left, 32x256xf16, slayout=row_major>, !pto.tile_buf<right, 256x128xf16, slayout=col_major>) outs(%38 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>)
         pto.tpush(%38, %27 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
-        pto.tmov ins(%31 : !pto.tile_buf<mat, 128x512xf16, slayout=col_major>) outs(%32 : !pto.tile_buf<right, 128x512xf16, slayout=col_major>)
-        pto.tmatmul ins(%30, %32 : !pto.tile_buf<left, 32x128xf16, slayout=row_major>, !pto.tile_buf<right, 128x512xf16, slayout=col_major>) outs(%33 : !pto.tile_buf<acc, 32x512xf32, blayout=col_major, slayout=row_major, fractal=1024>)
-        pto.tpush(%33, %25 : !pto.tile_buf<acc, 32x512xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
+        pto.tmov ins(%31 : !pto.tile_buf<mat, 128x256xf16, slayout=col_major>) outs(%32 : !pto.tile_buf<right, 128x256xf16, slayout=col_major>)
+        pto.tmatmul ins(%30, %32 : !pto.tile_buf<left, 32x128xf16, slayout=row_major>, !pto.tile_buf<right, 128x256xf16, slayout=col_major>) outs(%33 : !pto.tile_buf<acc, 32x256xf32, blayout=col_major, slayout=row_major, fractal=1024>)
+        pto.tpush(%33, %25 : !pto.tile_buf<acc, 32x256xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
       }
-      %47 = pto.tpop_from_aiv {id = 30, split = 1} -> !pto.tile_buf<mat, 32x512xf16, blayout=col_major, slayout=row_major>
-      pto.tmov ins(%47 : !pto.tile_buf<mat, 32x512xf16, blayout=col_major, slayout=row_major>) outs(%35 : !pto.tile_buf<left, 32x512xf16, slayout=row_major>)
+      %47 = pto.tpop_from_aiv {id = 30, split = 1} -> !pto.tile_buf<mat, 32x256xf16, blayout=col_major, slayout=row_major>
+      pto.tmov ins(%47 : !pto.tile_buf<mat, 32x256xf16, blayout=col_major, slayout=row_major>) outs(%35 : !pto.tile_buf<left, 32x256xf16, slayout=row_major>)
       pto.tfree_from_aiv {id = 30, split = 1}
-      pto.tmov ins(%36 : !pto.tile_buf<mat, 512x128xf16, blayout=col_major, slayout=row_major>) outs(%37 : !pto.tile_buf<right, 512x128xf16, slayout=col_major>)
-      %c7680 = arith.constant 7680 : index
-      %48 = pto.partition_view %41, offsets = [%c7680, %c0], sizes = [%c512, %c128] : !pto.tensor_view<?x?xf16>
-      pto.tload ins(%48 : !pto.partition_tensor_view<512x128xf16>) outs(%36 : !pto.tile_buf<mat, 512x128xf16, blayout=col_major, slayout=row_major>)
-      pto.tmatmul ins(%35, %37 : !pto.tile_buf<left, 32x512xf16, slayout=row_major>, !pto.tile_buf<right, 512x128xf16, slayout=col_major>) outs(%38 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>)
+      pto.tmov ins(%36 : !pto.tile_buf<mat, 256x128xf16, blayout=col_major, slayout=row_major>) outs(%37 : !pto.tile_buf<right, 256x128xf16, slayout=col_major>)
+      %c7936 = arith.constant 7936 : index
+      %48 = pto.partition_view %41, offsets = [%c7936, %c0], sizes = [%c256, %c128] : !pto.tensor_view<?x?xf16>
+      pto.tload ins(%48 : !pto.partition_tensor_view<256x128xf16>) outs(%36 : !pto.tile_buf<mat, 256x128xf16, blayout=col_major, slayout=row_major>)
+      pto.tmatmul ins(%35, %37 : !pto.tile_buf<left, 32x256xf16, slayout=row_major>, !pto.tile_buf<right, 256x128xf16, slayout=col_major>) outs(%38 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>)
       pto.tpush(%38, %27 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
-      %49 = pto.tpop_from_aiv {id = 30, split = 1} -> !pto.tile_buf<mat, 32x512xf16, blayout=col_major, slayout=row_major>
-      pto.tmov ins(%49 : !pto.tile_buf<mat, 32x512xf16, blayout=col_major, slayout=row_major>) outs(%35 : !pto.tile_buf<left, 32x512xf16, slayout=row_major>)
+      %49 = pto.tpop_from_aiv {id = 30, split = 1} -> !pto.tile_buf<mat, 32x256xf16, blayout=col_major, slayout=row_major>
+      pto.tmov ins(%49 : !pto.tile_buf<mat, 32x256xf16, blayout=col_major, slayout=row_major>) outs(%35 : !pto.tile_buf<left, 32x256xf16, slayout=row_major>)
       pto.tfree_from_aiv {id = 30, split = 1}
-      pto.tmov ins(%36 : !pto.tile_buf<mat, 512x128xf16, blayout=col_major, slayout=row_major>) outs(%37 : !pto.tile_buf<right, 512x128xf16, slayout=col_major>)
-      pto.tmatmul ins(%35, %37 : !pto.tile_buf<left, 32x512xf16, slayout=row_major>, !pto.tile_buf<right, 512x128xf16, slayout=col_major>) outs(%38 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>)
+      pto.tmov ins(%36 : !pto.tile_buf<mat, 256x128xf16, blayout=col_major, slayout=row_major>) outs(%37 : !pto.tile_buf<right, 256x128xf16, slayout=col_major>)
+      pto.tmatmul ins(%35, %37 : !pto.tile_buf<left, 32x256xf16, slayout=row_major>, !pto.tile_buf<right, 256x128xf16, slayout=col_major>) outs(%38 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>)
       pto.tpush(%38, %27 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
     }
     return
@@ -153,7 +153,7 @@
     %c32 = arith.constant 32 : index
     %c16 = arith.constant 16 : index
     %c128 = arith.constant 128 : index
-    %c16_0 = arith.constant 16 : index
+    %c32_0 = arith.constant 32 : index
     %c96 = arith.constant 96 : index
     %0 = pto.get_block_num
     %1 = arith.index_cast %0 : i64 to index
@@ -174,73 +174,73 @@
     %16 = arith.addi %4, %c1 : index
     %17 = arith.select %15, %16, %4 : index
     %18 = arith.addi %14, %17 : index
-    %c229376 = arith.constant 229376 : index
-    %19 = arith.muli %3, %c229376 : index
+    %c131072 = arith.constant 131072 : index
+    %19 = arith.muli %3, %c131072 : index
     %20 = pto.addptr %arg0, %19 : <f32> -> <f32>
     %c0_1 = arith.constant 0 : index
     %21 = pto.addptr %20, %c0_1 : <f32> -> <f32>
-    %c131072 = arith.constant 131072 : index
-    %22 = pto.addptr %20, %c131072 : <f32> -> <f32>
-    %c163840 = arith.constant 163840 : index
-    %23 = pto.addptr %20, %c163840 : <f32> -> <f32>
-    %24 = pto.reserve_buffer{name = "fa_qk_c2v_fifo", size = 65536, location = <vec>, auto = false, base = 0} -> i32
-    %25 = pto.initialize_l2g2l_pipe{dir_mask = 1, slot_size = 65536, slot_num = 8, local_slot_num = 1} (%21 : !pto.ptr<f32>, %24 : i32) -> !pto.pipe
-    %26 = pto.reserve_buffer{name = "fa_pv_c2v_fifo", size = 16384, location = <vec>, auto = false, base = 65536} -> i32
+    %c65536 = arith.constant 65536 : index
+    %22 = pto.addptr %20, %c65536 : <f32> -> <f32>
+    %c98304 = arith.constant 98304 : index
+    %23 = pto.addptr %20, %c98304 : <f32> -> <f32>
+    %24 = pto.reserve_buffer{name = "fa_qk_c2v_fifo", size = 32768, location = <vec>, auto = false, base = 0} -> i32
+    %25 = pto.initialize_l2g2l_pipe{dir_mask = 1, slot_size = 32768, slot_num = 8, local_slot_num = 1} (%21 : !pto.ptr<f32>, %24 : i32) -> !pto.pipe
+    %26 = pto.reserve_buffer{name = "fa_pv_c2v_fifo", size = 16384, location = <vec>, auto = false, base = 32768} -> i32
     %27 = pto.initialize_l2g2l_pipe{dir_mask = 1, slot_size = 16384, slot_num = 8, local_slot_num = 1} (%22 : !pto.ptr<f32>, %26 : i32) -> !pto.pipe
     %28 = pto.import_reserved_buffer{name = "fa_p_v2c_fifo", peer_func = @cube_kernel} -> i32
     %c0_i32 = arith.constant 0 : i32
-    pto.aiv_initialize_pipe {id = 30, dir_mask = 2, slot_size = 32768, nosplit = false}(gm_slot_buffer = %23 : !pto.ptr<f32>, c2v_consumer_buf = %c0_i32 : i32, v2c_consumer_buf = %28 : i32)
+    pto.aiv_initialize_pipe {id = 30, dir_mask = 2, slot_size = 16384, local_slot_num = 1, nosplit = false}(gm_slot_buffer = %23 : !pto.ptr<f32>, c2v_consumer_buf = %c0_i32 : i32, v2c_consumer_buf = %28 : i32)
     %29 = pto.get_subblock_idx
     %30 = arith.index_cast %29 : i64 to index
     %31 = arith.muli %30, %c16 : index
-    %c114688_i64 = arith.constant 114688 : i64
-    %32 = pto.alloc_tile addr = %c114688_i64 : !pto.tile_buf<vec, 16x512xf32>
-    %c147456_i64 = arith.constant 147456 : i64
-    %33 = pto.alloc_tile addr = %c147456_i64 : !pto.tile_buf<vec, 16x512xf32>
-    %c163840_i64 = arith.constant 163840 : i64
-    %34 = pto.alloc_tile addr = %c163840_i64 : !pto.tile_buf<vec, 16x512xf16>
-    %c180224_i64 = arith.constant 180224 : i64
-    %35 = pto.alloc_tile addr = %c180224_i64 : !pto.tile_buf<vec, 16x128xf32>
-    %c188416_i64 = arith.constant 188416 : i64
-    %36 = pto.alloc_tile addr = %c188416_i64 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>
-    %c188480_i64 = arith.constant 188480 : i64
-    %37 = pto.alloc_tile addr = %c188480_i64 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>
-    %c188544_i64 = arith.constant 188544 : i64
-    %38 = pto.alloc_tile addr = %c188544_i64 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>
-    %c188608_i64 = arith.constant 188608 : i64
-    %39 = pto.alloc_tile addr = %c188608_i64 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>
-    %c188672_i64 = arith.constant 188672 : i64
-    %40 = pto.alloc_tile addr = %c188672_i64 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>
-    %c188736_i64 = arith.constant 188736 : i64
-    %41 = pto.alloc_tile addr = %c188736_i64 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>
+    %c49152_i64 = arith.constant 49152 : i64
+    %32 = pto.alloc_tile addr = %c49152_i64 : !pto.tile_buf<vec, 16x256xf32>
+    %c65536_i64 = arith.constant 65536 : i64
+    %33 = pto.alloc_tile addr = %c65536_i64 : !pto.tile_buf<vec, 16x256xf32>
+    %c81920_i64 = arith.constant 81920 : i64
+    %34 = pto.alloc_tile addr = %c81920_i64 : !pto.tile_buf<vec, 16x256xf16>
+    %c90112_i64 = arith.constant 90112 : i64
+    %35 = pto.alloc_tile addr = %c90112_i64 : !pto.tile_buf<vec, 16x128xf32>
+    %c98304_i64 = arith.constant 98304 : i64
+    %36 = pto.alloc_tile addr = %c98304_i64 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>
+    %c98816_i64 = arith.constant 98816 : i64
+    %37 = pto.alloc_tile addr = %c98816_i64 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>
+    %c99328_i64 = arith.constant 99328 : i64
+    %38 = pto.alloc_tile addr = %c99328_i64 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>
+    %c99840_i64 = arith.constant 99840 : i64
+    %39 = pto.alloc_tile addr = %c99840_i64 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>
+    %c100352_i64 = arith.constant 100352 : i64
+    %40 = pto.alloc_tile addr = %c100352_i64 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>
+    %c100864_i64 = arith.constant 100864 : i64
+    %41 = pto.alloc_tile addr = %c100864_i64 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>
     %cst = arith.constant 0.0883883461 : f32
     %cst_2 = arith.constant 1.000000e+00 : f32
     %c3072 = arith.constant 3072 : index
     %42 = pto.make_tensor_view %arg1, shape = [%c3072, %c128], strides = [%c128, %c1] : !pto.tensor_view<?x?xf32>
     scf.for %arg2 = %14 to %18 step %c1 {
       %43 = arith.muli %arg2, %c32 : index
-      %c81920_i64 = arith.constant 81920 : i64
-      %44 = pto.alloc_tile addr = %c81920_i64 : !pto.tile_buf<vec, 16x512xf32>
-      pto.tpop(%44, %25 : !pto.tile_buf<vec, 16x512xf32>, !pto.pipe) {split = 1}
-      pto.tmuls ins(%44, %cst : !pto.tile_buf<vec, 16x512xf32>, f32) outs(%44 : !pto.tile_buf<vec, 16x512xf32>)
-      pto.trowmax ins(%44, %32 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x512xf32>) outs(%37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
+      %c101376_i64 = arith.constant 101376 : i64
+      %44 = pto.alloc_tile addr = %c101376_i64 : !pto.tile_buf<vec, 16x256xf32>
+      pto.tpop(%44, %25 : !pto.tile_buf<vec, 16x256xf32>, !pto.pipe) {split = 1}
+      pto.tmuls ins(%44, %cst : !pto.tile_buf<vec, 16x256xf32>, f32) outs(%44 : !pto.tile_buf<vec, 16x256xf32>)
+      pto.trowmax ins(%44, %32 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x256xf32>) outs(%37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
       %45 = pto.treshape %37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
       %46 = pto.treshape %36 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
       %47 = pto.treshape %40 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
       %48 = pto.treshape %38 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
       %49 = pto.treshape %39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
-      pto.trowexpandsub ins(%44, %37 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%33 : !pto.tile_buf<vec, 16x512xf32>)
+      pto.trowexpandsub ins(%44, %37 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%33 : !pto.tile_buf<vec, 16x256xf32>)
       pto.tmuls ins(%45, %cst_2 : !pto.tile_buf<vec, 1x16xf32>, f32) outs(%46 : !pto.tile_buf<vec, 1x16xf32>)
-      pto.texp ins(%33 : !pto.tile_buf<vec, 16x512xf32>) outs(%33 : !pto.tile_buf<vec, 16x512xf32>)
-      pto.trowsum ins(%33, %32 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x512xf32>) outs(%38 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
-      pto.tcvt ins(%33 {rmode = #pto<round_mode CAST_RINT>} : !pto.tile_buf<vec, 16x512xf32>) outs(%34 : !pto.tile_buf<vec, 16x512xf16>)
-      pto.tpush_to_aic(%34 : !pto.tile_buf<vec, 16x512xf16>) {id = 30, split = 1}
+      pto.texp ins(%33 : !pto.tile_buf<vec, 16x256xf32>) outs(%33 : !pto.tile_buf<vec, 16x256xf32>)
+      pto.trowsum ins(%33, %32 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x256xf32>) outs(%38 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
+      pto.tcvt ins(%33 {rmode = #pto<round_mode CAST_RINT>} : !pto.tile_buf<vec, 16x256xf32>) outs(%34 : !pto.tile_buf<vec, 16x256xf16>)
+      pto.tpush_to_aic(%34 : !pto.tile_buf<vec, 16x256xf16>) {id = 30, split = 1}
       pto.tfree(%25 : !pto.pipe) {split = 1}
-      %c81920_i64_3 = arith.constant 81920 : i64
-      %50 = pto.alloc_tile addr = %c81920_i64_3 : !pto.tile_buf<vec, 16x512xf32>
-      pto.tpop(%50, %25 : !pto.tile_buf<vec, 16x512xf32>, !pto.pipe) {split = 1}
-      pto.tmuls ins(%50, %cst : !pto.tile_buf<vec, 16x512xf32>, f32) outs(%50 : !pto.tile_buf<vec, 16x512xf32>)
-      pto.trowmax ins(%50, %32 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x512xf32>) outs(%37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
+      %c101376_i64_3 = arith.constant 101376 : i64
+      %50 = pto.alloc_tile addr = %c101376_i64_3 : !pto.tile_buf<vec, 16x256xf32>
+      pto.tpop(%50, %25 : !pto.tile_buf<vec, 16x256xf32>, !pto.pipe) {split = 1}
+      pto.tmuls ins(%50, %cst : !pto.tile_buf<vec, 16x256xf32>, f32) outs(%50 : !pto.tile_buf<vec, 16x256xf32>)
+      pto.trowmax ins(%50, %32 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x256xf32>) outs(%37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
       %51 = pto.treshape %37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
       %52 = pto.treshape %36 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
       %53 = pto.treshape %41 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
@@ -249,25 +249,25 @@
       pto.tmax ins(%51, %52 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%51 : !pto.tile_buf<vec, 1x16xf32>)
       pto.tsub ins(%52, %51 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%53 : !pto.tile_buf<vec, 1x16xf32>)
       pto.tmuls ins(%51, %cst_2 : !pto.tile_buf<vec, 1x16xf32>, f32) outs(%52 : !pto.tile_buf<vec, 1x16xf32>)
-      pto.trowexpandsub ins(%50, %37 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%33 : !pto.tile_buf<vec, 16x512xf32>)
+      pto.trowexpandsub ins(%50, %37 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%33 : !pto.tile_buf<vec, 16x256xf32>)
       pto.texp ins(%53 : !pto.tile_buf<vec, 1x16xf32>) outs(%53 : !pto.tile_buf<vec, 1x16xf32>)
-      pto.texp ins(%33 : !pto.tile_buf<vec, 16x512xf32>) outs(%33 : !pto.tile_buf<vec, 16x512xf32>)
+      pto.texp ins(%33 : !pto.tile_buf<vec, 16x256xf32>) outs(%33 : !pto.tile_buf<vec, 16x256xf32>)
       pto.tmul ins(%54, %53 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%54 : !pto.tile_buf<vec, 1x16xf32>)
-      pto.trowsum ins(%33, %32 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x512xf32>) outs(%39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
+      pto.trowsum ins(%33, %32 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x256xf32>) outs(%39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
       pto.tadd ins(%54, %55 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%54 : !pto.tile_buf<vec, 1x16xf32>)
-      pto.tcvt ins(%33 {rmode = #pto<round_mode CAST_RINT>} : !pto.tile_buf<vec, 16x512xf32>) outs(%34 : !pto.tile_buf<vec, 16x512xf16>)
-      pto.tpush_to_aic(%34 : !pto.tile_buf<vec, 16x512xf16>) {id = 30, split = 1}
+      pto.tcvt ins(%33 {rmode = #pto<round_mode CAST_RINT>} : !pto.tile_buf<vec, 16x256xf32>) outs(%34 : !pto.tile_buf<vec, 16x256xf16>)
+      pto.tpush_to_aic(%34 : !pto.tile_buf<vec, 16x256xf16>) {id = 30, split = 1}
       pto.tfree(%25 : !pto.pipe) {split = 1}
-      %c114688_i64_4 = arith.constant 114688 : i64
-      %56 = pto.alloc_tile addr = %c114688_i64_4 : !pto.tile_buf<vec, 16x128xf32>
+      %c101376_i64_4 = arith.constant 101376 : i64
+      %56 = pto.alloc_tile addr = %c101376_i64_4 : !pto.tile_buf<vec, 16x128xf32>
       pto.tpop(%56, %27 : !pto.tile_buf<vec, 16x128xf32>, !pto.pipe) {split = 1}
       pto.tmov ins(%56 : !pto.tile_buf<vec, 16x128xf32>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
       pto.tfree(%27 : !pto.pipe) {split = 1}
-      %c81920_i64_5 = arith.constant 81920 : i64
-      %57 = pto.alloc_tile addr = %c81920_i64_5 : !pto.tile_buf<vec, 16x512xf32>
-      pto.tpop(%57, %25 : !pto.tile_buf<vec, 16x512xf32>, !pto.pipe) {split = 1}
-      pto.tmuls ins(%57, %cst : !pto.tile_buf<vec, 16x512xf32>, f32) outs(%57 : !pto.tile_buf<vec, 16x512xf32>)
-      pto.trowmax ins(%57, %32 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x512xf32>) outs(%37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
+      %c101376_i64_5 = arith.constant 101376 : i64
+      %57 = pto.alloc_tile addr = %c101376_i64_5 : !pto.tile_buf<vec, 16x256xf32>
+      pto.tpop(%57, %25 : !pto.tile_buf<vec, 16x256xf32>, !pto.pipe) {split = 1}
+      pto.tmuls ins(%57, %cst : !pto.tile_buf<vec, 16x256xf32>, f32) outs(%57 : !pto.tile_buf<vec, 16x256xf32>)
+      pto.trowmax ins(%57, %32 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x256xf32>) outs(%37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
       %58 = pto.treshape %37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
       %59 = pto.treshape %36 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
       %60 = pto.treshape %40 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
@@ -276,26 +276,26 @@
       pto.tmax ins(%58, %59 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%58 : !pto.tile_buf<vec, 1x16xf32>)
       pto.tsub ins(%59, %58 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%60 : !pto.tile_buf<vec, 1x16xf32>)
       pto.tmuls ins(%58, %cst_2 : !pto.tile_buf<vec, 1x16xf32>, f32) outs(%59 : !pto.tile_buf<vec, 1x16xf32>)
-      pto.trowexpandsub ins(%57, %37 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%33 : !pto.tile_buf<vec, 16x512xf32>)
+      pto.trowexpandsub ins(%57, %37 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%33 : !pto.tile_buf<vec, 16x256xf32>)
       pto.texp ins(%60 : !pto.tile_buf<vec, 1x16xf32>) outs(%60 : !pto.tile_buf<vec, 1x16xf32>)
-      pto.texp ins(%33 : !pto.tile_buf<vec, 16x512xf32>) outs(%33 : !pto.tile_buf<vec, 16x512xf32>)
+      pto.texp ins(%33 : !pto.tile_buf<vec, 16x256xf32>) outs(%33 : !pto.tile_buf<vec, 16x256xf32>)
       pto.tmul ins(%61, %60 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%61 : !pto.tile_buf<vec, 1x16xf32>)
-      pto.trowsum ins(%33, %32 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x512xf32>) outs(%39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
+      pto.trowsum ins(%33, %32 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x256xf32>) outs(%39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
       pto.tadd ins(%61, %62 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%61 : !pto.tile_buf<vec, 1x16xf32>)
-      pto.tcvt ins(%33 {rmode = #pto<round_mode CAST_RINT>} : !pto.tile_buf<vec, 16x512xf32>) outs(%34 : !pto.tile_buf<vec, 16x512xf16>)
-      pto.tpush_to_aic(%34 : !pto.tile_buf<vec, 16x512xf16>) {id = 30, split = 1}
+      pto.tcvt ins(%33 {rmode = #pto<round_mode CAST_RINT>} : !pto.tile_buf<vec, 16x256xf32>) outs(%34 : !pto.tile_buf<vec, 16x256xf16>)
+      pto.tpush_to_aic(%34 : !pto.tile_buf<vec, 16x256xf16>) {id = 30, split = 1}
       pto.tfree(%25 : !pto.pipe) {split = 1}
-      %c114688_i64_6 = arith.constant 114688 : i64
-      %63 = pto.alloc_tile addr = %c114688_i64_6 : !pto.tile_buf<vec, 16x128xf32>
+      %c101376_i64_6 = arith.constant 101376 : i64
+      %63 = pto.alloc_tile addr = %c101376_i64_6 : !pto.tile_buf<vec, 16x128xf32>
       pto.tpop(%63, %27 : !pto.tile_buf<vec, 16x128xf32>, !pto.pipe) {split = 1}
       pto.trowexpandmul ins(%35, %41 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
       pto.tadd ins(%35, %63 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x128xf32>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
       pto.tfree(%27 : !pto.pipe) {split = 1}
-      %c81920_i64_7 = arith.constant 81920 : i64
-      %64 = pto.alloc_tile addr = %c81920_i64_7 : !pto.tile_buf<vec, 16x512xf32>
-      pto.tpop(%64, %25 : !pto.tile_buf<vec, 16x512xf32>, !pto.pipe) {split = 1}
-      pto.tmuls ins(%64, %cst : !pto.tile_buf<vec, 16x512xf32>, f32) outs(%64 : !pto.tile_buf<vec, 16x512xf32>)
-      pto.trowmax ins(%64, %32 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x512xf32>) outs(%37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
+      %c101376_i64_7 = arith.constant 101376 : i64
+      %64 = pto.alloc_tile addr = %c101376_i64_7 : !pto.tile_buf<vec, 16x256xf32>
+      pto.tpop(%64, %25 : !pto.tile_buf<vec, 16x256xf32>, !pto.pipe) {split = 1}
+      pto.tmuls ins(%64, %cst : !pto.tile_buf<vec, 16x256xf32>, f32) outs(%64 : !pto.tile_buf<vec, 16x256xf32>)
+      pto.trowmax ins(%64, %32 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x256xf32>) outs(%37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
       %65 = pto.treshape %37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
       %66 = pto.treshape %36 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
       %67 = pto.treshape %41 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
@@ -304,28 +304,28 @@
       pto.tmax ins(%65, %66 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%65 : !pto.tile_buf<vec, 1x16xf32>)
       pto.tsub ins(%66, %65 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%67 : !pto.tile_buf<vec, 1x16xf32>)
       pto.tmuls ins(%65, %cst_2 : !pto.tile_buf<vec, 1x16xf32>, f32) outs(%66 : !pto.tile_buf<vec, 1x16xf32>)
-      pto.trowexpandsub ins(%64, %37 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%33 : !pto.tile_buf<vec, 16x512xf32>)
+      pto.trowexpandsub ins(%64, %37 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%33 : !pto.tile_buf<vec, 16x256xf32>)
       pto.texp ins(%67 : !pto.tile_buf<vec, 1x16xf32>) outs(%67 : !pto.tile_buf<vec, 1x16xf32>)
-      pto.texp ins(%33 : !pto.tile_buf<vec, 16x512xf32>) outs(%33 : !pto.tile_buf<vec, 16x512xf32>)
+      pto.texp ins(%33 : !pto.tile_buf<vec, 16x256xf32>) outs(%33 : !pto.tile_buf<vec, 16x256xf32>)
       pto.tmul ins(%68, %67 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%68 : !pto.tile_buf<vec, 1x16xf32>)
-      pto.trowsum ins(%33, %32 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x512xf32>) outs(%39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
+      pto.trowsum ins(%33, %32 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x256xf32>) outs(%39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
       pto.tadd ins(%68, %69 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%68 : !pto.tile_buf<vec, 1x16xf32>)
-      pto.tcvt ins(%33 {rmode = #pto<round_mode CAST_RINT>} : !pto.tile_buf<vec, 16x512xf32>) outs(%34 : !pto.tile_buf<vec, 16x512xf16>)
-      pto.tpush_to_aic(%34 : !pto.tile_buf<vec, 16x512xf16>) {id = 30, split = 1}
+      pto.tcvt ins(%33 {rmode = #pto<round_mode CAST_RINT>} : !pto.tile_buf<vec, 16x256xf32>) outs(%34 : !pto.tile_buf<vec, 16x256xf16>)
+      pto.tpush_to_aic(%34 : !pto.tile_buf<vec, 16x256xf16>) {id = 30, split = 1}
       pto.tfree(%25 : !pto.pipe) {split = 1}
-      %c7 = arith.constant 7 : index
-      scf.for %arg3 = %c1 to %c7 step %c1 {
-        %c114688_i64_10 = arith.constant 114688 : i64
-        %74 = pto.alloc_tile addr = %c114688_i64_10 : !pto.tile_buf<vec, 16x128xf32>
+      %c15 = arith.constant 15 : index
+      scf.for %arg3 = %c1 to %c15 step %c1 {
+        %c101376_i64_10 = arith.constant 101376 : i64
+        %74 = pto.alloc_tile addr = %c101376_i64_10 : !pto.tile_buf<vec, 16x128xf32>
         pto.tpop(%74, %27 : !pto.tile_buf<vec, 16x128xf32>, !pto.pipe) {split = 1}
         pto.trowexpandmul ins(%35, %40 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
         pto.tadd ins(%35, %74 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x128xf32>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
         pto.tfree(%27 : !pto.pipe) {split = 1}
-        %c81920_i64_11 = arith.constant 81920 : i64
-        %75 = pto.alloc_tile addr = %c81920_i64_11 : !pto.tile_buf<vec, 16x512xf32>
-        pto.tpop(%75, %25 : !pto.tile_buf<vec, 16x512xf32>, !pto.pipe) {split = 1}
-        pto.tmuls ins(%75, %cst : !pto.tile_buf<vec, 16x512xf32>, f32) outs(%75 : !pto.tile_buf<vec, 16x512xf32>)
-        pto.trowmax ins(%75, %32 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x512xf32>) outs(%37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
+        %c101376_i64_11 = arith.constant 101376 : i64
+        %75 = pto.alloc_tile addr = %c101376_i64_11 : !pto.tile_buf<vec, 16x256xf32>
+        pto.tpop(%75, %25 : !pto.tile_buf<vec, 16x256xf32>, !pto.pipe) {split = 1}
+        pto.tmuls ins(%75, %cst : !pto.tile_buf<vec, 16x256xf32>, f32) outs(%75 : !pto.tile_buf<vec, 16x256xf32>)
+        pto.trowmax ins(%75, %32 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x256xf32>) outs(%37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
         %76 = pto.treshape %37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
         %77 = pto.treshape %36 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
         %78 = pto.treshape %40 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
@@ -334,26 +334,26 @@
         pto.tmax ins(%76, %77 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%76 : !pto.tile_buf<vec, 1x16xf32>)
         pto.tsub ins(%77, %76 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%78 : !pto.tile_buf<vec, 1x16xf32>)
         pto.tmuls ins(%76, %cst_2 : !pto.tile_buf<vec, 1x16xf32>, f32) outs(%77 : !pto.tile_buf<vec, 1x16xf32>)
-        pto.trowexpandsub ins(%75, %37 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%33 : !pto.tile_buf<vec, 16x512xf32>)
+        pto.trowexpandsub ins(%75, %37 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%33 : !pto.tile_buf<vec, 16x256xf32>)
         pto.texp ins(%78 : !pto.tile_buf<vec, 1x16xf32>) outs(%78 : !pto.tile_buf<vec, 1x16xf32>)
-        pto.texp ins(%33 : !pto.tile_buf<vec, 16x512xf32>) outs(%33 : !pto.tile_buf<vec, 16x512xf32>)
+        pto.texp ins(%33 : !pto.tile_buf<vec, 16x256xf32>) outs(%33 : !pto.tile_buf<vec, 16x256xf32>)
         pto.tmul ins(%79, %78 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%79 : !pto.tile_buf<vec, 1x16xf32>)
-        pto.trowsum ins(%33, %32 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x512xf32>) outs(%39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
+        pto.trowsum ins(%33, %32 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x256xf32>) outs(%39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
         pto.tadd ins(%79, %80 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%79 : !pto.tile_buf<vec, 1x16xf32>)
-        pto.tcvt ins(%33 {rmode = #pto<round_mode CAST_RINT>} : !pto.tile_buf<vec, 16x512xf32>) outs(%34 : !pto.tile_buf<vec, 16x512xf16>)
-        pto.tpush_to_aic(%34 : !pto.tile_buf<vec, 16x512xf16>) {id = 30, split = 1}
+        pto.tcvt ins(%33 {rmode = #pto<round_mode CAST_RINT>} : !pto.tile_buf<vec, 16x256xf32>) outs(%34 : !pto.tile_buf<vec, 16x256xf16>)
+        pto.tpush_to_aic(%34 : !pto.tile_buf<vec, 16x256xf16>) {id = 30, split = 1}
         pto.tfree(%25 : !pto.pipe) {split = 1}
-        %c114688_i64_12 = arith.constant 114688 : i64
-        %81 = pto.alloc_tile addr = %c114688_i64_12 : !pto.tile_buf<vec, 16x128xf32>
+        %c101376_i64_12 = arith.constant 101376 : i64
+        %81 = pto.alloc_tile addr = %c101376_i64_12 : !pto.tile_buf<vec, 16x128xf32>
         pto.tpop(%81, %27 : !pto.tile_buf<vec, 16x128xf32>, !pto.pipe) {split = 1}
         pto.trowexpandmul ins(%35, %41 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
         pto.tadd ins(%35, %81 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x128xf32>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
         pto.tfree(%27 : !pto.pipe) {split = 1}
-        %c81920_i64_13 = arith.constant 81920 : i64
-        %82 = pto.alloc_tile addr = %c81920_i64_13 : !pto.tile_buf<vec, 16x512xf32>
-        pto.tpop(%82, %25 : !pto.tile_buf<vec, 16x512xf32>, !pto.pipe) {split = 1}
-        pto.tmuls ins(%82, %cst : !pto.tile_buf<vec, 16x512xf32>, f32) outs(%82 : !pto.tile_buf<vec, 16x512xf32>)
-        pto.trowmax ins(%82, %32 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x512xf32>) outs(%37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
+        %c101376_i64_13 = arith.constant 101376 : i64
+        %82 = pto.alloc_tile addr = %c101376_i64_13 : !pto.tile_buf<vec, 16x256xf32>
+        pto.tpop(%82, %25 : !pto.tile_buf<vec, 16x256xf32>, !pto.pipe) {split = 1}
+        pto.tmuls ins(%82, %cst : !pto.tile_buf<vec, 16x256xf32>, f32) outs(%82 : !pto.tile_buf<vec, 16x256xf32>)
+        pto.trowmax ins(%82, %32 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x256xf32>) outs(%37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
         %83 = pto.treshape %37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
         %84 = pto.treshape %36 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
         %85 = pto.treshape %41 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
@@ -362,24 +362,24 @@
         pto.tmax ins(%83, %84 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%83 : !pto.tile_buf<vec, 1x16xf32>)
         pto.tsub ins(%84, %83 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%85 : !pto.tile_buf<vec, 1x16xf32>)
         pto.tmuls ins(%83, %cst_2 : !pto.tile_buf<vec, 1x16xf32>, f32) outs(%84 : !pto.tile_buf<vec, 1x16xf32>)
-        pto.trowexpandsub ins(%82, %37 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%33 : !pto.tile_buf<vec, 16x512xf32>)
+        pto.trowexpandsub ins(%82, %37 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%33 : !pto.tile_buf<vec, 16x256xf32>)
         pto.texp ins(%85 : !pto.tile_buf<vec, 1x16xf32>) outs(%85 : !pto.tile_buf<vec, 1x16xf32>)
-        pto.texp ins(%33 : !pto.tile_buf<vec, 16x512xf32>) outs(%33 : !pto.tile_buf<vec, 16x512xf32>)
+        pto.texp ins(%33 : !pto.tile_buf<vec, 16x256xf32>) outs(%33 : !pto.tile_buf<vec, 16x256xf32>)
         pto.tmul ins(%86, %85 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%86 : !pto.tile_buf<vec, 1x16xf32>)
-        pto.trowsum ins(%33, %32 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x512xf32>) outs(%39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
+        pto.trowsum ins(%33, %32 : !pto.tile_buf<vec, 16x256xf32>, !pto.tile_buf<vec, 16x256xf32>) outs(%39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
         pto.tadd ins(%86, %87 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%86 : !pto.tile_buf<vec, 1x16xf32>)
-        pto.tcvt ins(%33 {rmode = #pto<round_mode CAST_RINT>} : !pto.tile_buf<vec, 16x512xf32>) outs(%34 : !pto.tile_buf<vec, 16x512xf16>)
-        pto.tpush_to_aic(%34 : !pto.tile_buf<vec, 16x512xf16>) {id = 30, split = 1}
+        pto.tcvt ins(%33 {rmode = #pto<round_mode CAST_RINT>} : !pto.tile_buf<vec, 16x256xf32>) outs(%34 : !pto.tile_buf<vec, 16x256xf16>)
+        pto.tpush_to_aic(%34 : !pto.tile_buf<vec, 16x256xf16>) {id = 30, split = 1}
         pto.tfree(%25 : !pto.pipe) {split = 1}
       }
-      %c114688_i64_8 = arith.constant 114688 : i64
-      %70 = pto.alloc_tile addr = %c114688_i64_8 : !pto.tile_buf<vec, 16x128xf32>
+      %c101376_i64_8 = arith.constant 101376 : i64
+      %70 = pto.alloc_tile addr = %c101376_i64_8 : !pto.tile_buf<vec, 16x128xf32>
       pto.tpop(%70, %27 : !pto.tile_buf<vec, 16x128xf32>, !pto.pipe) {split = 1}
       pto.trowexpandmul ins(%35, %40 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
       pto.tadd ins(%35, %70 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x128xf32>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
       pto.tfree(%27 : !pto.pipe) {split = 1}
-      %c114688_i64_9 = arith.constant 114688 : i64
-      %71 = pto.alloc_tile addr = %c114688_i64_9 : !pto.tile_buf<vec, 16x128xf32>
+      %c101376_i64_9 = arith.constant 101376 : i64
+      %71 = pto.alloc_tile addr = %c101376_i64_9 : !pto.tile_buf<vec, 16x128xf32>
       pto.tpop(%71, %27 : !pto.tile_buf<vec, 16x128xf32>, !pto.pipe) {split = 1}
       pto.trowexpandmul ins(%35, %41 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
       pto.tadd ins(%35, %71 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x128xf32>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
@@ -397,4 +397,5 @@
     call @vector_kernel(%arg1, %arg5) : (!pto.ptr<f32>, !pto.ptr<f32>) -> ()
     return
   }
-}
\ No newline at end of file
+}
+

[MirkoDeVita98]

The MLIR with S1_TILE=512 is not in #609 anymore. This is the MLIR we tried that gives the AI core exception:

module {
  func.func @cube_kernel(%arg0: !pto.ptr<f32>, %arg1: !pto.ptr<f16>, %arg2: !pto.ptr<f16>, %arg3: !pto.ptr<f16>) attributes {pto.kernel_kind = #pto.kernel_kind<cube>} {
    %c0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    %c2 = arith.constant 2 : index
    %c32 = arith.constant 32 : index
    %c128 = arith.constant 128 : index
    %c512 = arith.constant 512 : index
    %c8192 = arith.constant 8192 : index
    %c16 = arith.constant 16 : index
    %c96 = arith.constant 96 : index
    %0 = pto.get_block_num
    %1 = arith.index_cast %0 : i64 to index
    %2 = pto.get_block_idx
    %3 = arith.index_cast %2 : i64 to index
    %4 = arith.divsi %c96, %1 : index
    %5 = arith.remsi %c96, %1 : index
    %6 = arith.addi %4, %c1 : index
    %7 = arith.muli %3, %6 : index
    %8 = arith.addi %4, %c1 : index
    %9 = arith.muli %5, %8 : index
    %10 = arith.subi %3, %5 : index
    %11 = arith.muli %10, %4 : index
    %12 = arith.addi %9, %11 : index
    %13 = arith.cmpi slt, %3, %5 : index
    %14 = arith.select %13, %7, %12 : index
    %15 = arith.cmpi slt, %3, %5 : index
    %16 = arith.addi %4, %c1 : index
    %17 = arith.select %15, %16, %4 : index
    %18 = arith.addi %14, %17 : index
    %c229376 = arith.constant 229376 : index
    %19 = arith.muli %3, %c229376 : index
    %20 = pto.addptr %arg0, %19 : <f32> -> <f32>
    %c0_0 = arith.constant 0 : index
    %21 = pto.addptr %20, %c0_0 : <f32> -> <f32>
    %c131072 = arith.constant 131072 : index
    %22 = pto.addptr %20, %c131072 : <f32> -> <f32>
    %c163840 = arith.constant 163840 : index
    %23 = pto.addptr %20, %c163840 : <f32> -> <f32>
    %24 = pto.import_reserved_buffer{name = "fa_qk_c2v_fifo", peer_func = @vector_kernel} -> i32
    %25 = pto.initialize_l2g2l_pipe{dir_mask = 1, slot_size = 65536, slot_num = 8, local_slot_num = 1} (%21 : !pto.ptr<f32>, %24 : i32) -> !pto.pipe
    %26 = pto.import_reserved_buffer{name = "fa_pv_c2v_fifo", peer_func = @vector_kernel} -> i32
    %27 = pto.initialize_l2g2l_pipe{dir_mask = 1, slot_size = 16384, slot_num = 8, local_slot_num = 1} (%22 : !pto.ptr<f32>, %26 : i32) -> !pto.pipe
    %28 = pto.reserve_buffer{name = "fa_p_v2c_fifo", size = 262144, location = <mat>, auto = false, base = 327680} -> i32
    %c0_i32 = arith.constant 0 : i32
    pto.aic_initialize_pipe {id = 30, dir_mask = 2, slot_size = 32768, nosplit = false}(gm_slot_buffer = %23 : !pto.ptr<f32>, c2v_consumer_buf = %c0_i32 : i32, v2c_consumer_buf = %28 : i32)
    %c0_i64 = arith.constant 0 : i64
    %c0_i64_1 = arith.constant 0 : i64
    %29 = pto.alloc_tile addr = %c0_i64_1 : !pto.tile_buf<mat, 32x128xf16, blayout=col_major, slayout=row_major>
    %c0_i64_2 = arith.constant 0 : i64
    %30 = pto.alloc_tile addr = %c0_i64_2 : !pto.tile_buf<left, 32x128xf16, slayout=row_major>
    %c8192_i64 = arith.constant 8192 : i64
    %31 = pto.alloc_tile addr = %c8192_i64 : !pto.tile_buf<mat, 128x512xf16, slayout=col_major>
    %32 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf<right, 128x512xf16, slayout=col_major>
    %c0_i64_3 = arith.constant 0 : i64
    %33 = pto.alloc_tile addr = %c0_i64_3 : !pto.tile_buf<acc, 32x512xf32, blayout=col_major, slayout=row_major, fractal=1024>
    %c139264_i64 = arith.constant 139264 : i64
    %34 = pto.alloc_tile addr = %c139264_i64 : !pto.tile_buf<mat, 32x512xf16, blayout=col_major, slayout=row_major>
    %c8192_i64_4 = arith.constant 8192 : i64
    %35 = pto.alloc_tile addr = %c8192_i64_4 : !pto.tile_buf<left, 32x512xf16, slayout=row_major>
    %c172032_i64 = arith.constant 172032 : i64
    %36 = pto.alloc_tile addr = %c172032_i64 : !pto.tile_buf<mat, 512x128xf16, blayout=col_major, slayout=row_major>
    %37 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf<right, 512x128xf16, slayout=col_major>
    %c65536_i64 = arith.constant 65536 : i64
    %38 = pto.alloc_tile addr = %c65536_i64 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>
    %c3072 = arith.constant 3072 : index
    %39 = pto.make_tensor_view %arg1, shape = [%c3072, %c128], strides = [%c128, %c1] : !pto.tensor_view<?x?xf16>
    %40 = pto.make_tensor_view %arg2, shape = [%c128, %c8192], strides = [%c1, %c128] : !pto.tensor_view<?x?xf16>
    %41 = pto.make_tensor_view %arg3, shape = [%c8192, %c128], strides = [%c128, %c1] : !pto.tensor_view<?x?xf16>
    scf.for %arg4 = %14 to %18 step %c1 {
      %42 = arith.muli %arg4, %c32 : index
      %43 = pto.partition_view %39, offsets = [%42, %c0], sizes = [%c32, %c128] : !pto.tensor_view<?x?xf16>
      pto.tload ins(%43 : !pto.partition_tensor_view<32x128xf16>) outs(%29 : !pto.tile_buf<mat, 32x128xf16, blayout=col_major, slayout=row_major>)
      pto.tmov ins(%29 : !pto.tile_buf<mat, 32x128xf16, blayout=col_major, slayout=row_major>) outs(%30 : !pto.tile_buf<left, 32x128xf16, slayout=row_major>)
      %c0_5 = arith.constant 0 : index
      %44 = pto.partition_view %40, offsets = [%c0, %c0_5], sizes = [%c128, %c512] : !pto.tensor_view<?x?xf16>
      pto.tload ins(%44 : !pto.partition_tensor_view<128x512xf16>) outs(%31 : !pto.tile_buf<mat, 128x512xf16, slayout=col_major>)
      pto.tmov ins(%31 : !pto.tile_buf<mat, 128x512xf16, slayout=col_major>) outs(%32 : !pto.tile_buf<right, 128x512xf16, slayout=col_major>)
      pto.tmatmul ins(%30, %32 : !pto.tile_buf<left, 32x128xf16, slayout=row_major>, !pto.tile_buf<right, 128x512xf16, slayout=col_major>) outs(%33 : !pto.tile_buf<acc, 32x512xf32, blayout=col_major, slayout=row_major, fractal=1024>)
      pto.tpush(%33, %25 : !pto.tile_buf<acc, 32x512xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
      %c512_6 = arith.constant 512 : index
      %45 = pto.partition_view %40, offsets = [%c0, %c512_6], sizes = [%c128, %c512] : !pto.tensor_view<?x?xf16>
      pto.tload ins(%45 : !pto.partition_tensor_view<128x512xf16>) outs(%31 : !pto.tile_buf<mat, 128x512xf16, slayout=col_major>)
      pto.tmov ins(%31 : !pto.tile_buf<mat, 128x512xf16, slayout=col_major>) outs(%32 : !pto.tile_buf<right, 128x512xf16, slayout=col_major>)
      pto.tmatmul ins(%30, %32 : !pto.tile_buf<left, 32x128xf16, slayout=row_major>, !pto.tile_buf<right, 128x512xf16, slayout=col_major>) outs(%33 : !pto.tile_buf<acc, 32x512xf32, blayout=col_major, slayout=row_major, fractal=1024>)
      pto.tpush(%33, %25 : !pto.tile_buf<acc, 32x512xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
      %46 = pto.partition_view %41, offsets = [%c0, %c0], sizes = [%c512, %c128] : !pto.tensor_view<?x?xf16>
      pto.tload ins(%46 : !pto.partition_tensor_view<512x128xf16>) outs(%36 : !pto.tile_buf<mat, 512x128xf16, blayout=col_major, slayout=row_major>)
      %c7 = arith.constant 7 : index
      scf.for %arg5 = %c0 to %c7 step %c1 {
        %50 = arith.muli %arg5, %c2 : index
        %c2_7 = arith.constant 2 : index
        %51 = arith.addi %50, %c2_7 : index
        %52 = arith.muli %51, %c512 : index
        %53 = pto.partition_view %40, offsets = [%c0, %52], sizes = [%c128, %c512] : !pto.tensor_view<?x?xf16>
        pto.tload ins(%53 : !pto.partition_tensor_view<128x512xf16>) outs(%31 : !pto.tile_buf<mat, 128x512xf16, slayout=col_major>)
        %54 = pto.tpop_from_aiv {id = 30, split = 1} -> !pto.tile_buf<mat, 32x512xf16, blayout=col_major, slayout=row_major>
        pto.tmov ins(%54 : !pto.tile_buf<mat, 32x512xf16, blayout=col_major, slayout=row_major>) outs(%35 : !pto.tile_buf<left, 32x512xf16, slayout=row_major>)
        pto.tfree_from_aiv {id = 30, split = 1}
        pto.tmov ins(%36 : !pto.tile_buf<mat, 512x128xf16, blayout=col_major, slayout=row_major>) outs(%37 : !pto.tile_buf<right, 512x128xf16, slayout=col_major>)
        %55 = arith.addi %50, %c1 : index
        %56 = arith.muli %55, %c512 : index
        %57 = pto.partition_view %41, offsets = [%56, %c0], sizes = [%c512, %c128] : !pto.tensor_view<?x?xf16>
        pto.tload ins(%57 : !pto.partition_tensor_view<512x128xf16>) outs(%36 : !pto.tile_buf<mat, 512x128xf16, blayout=col_major, slayout=row_major>)
        pto.tmatmul ins(%35, %37 : !pto.tile_buf<left, 32x512xf16, slayout=row_major>, !pto.tile_buf<right, 512x128xf16, slayout=col_major>) outs(%38 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>)
        pto.tpush(%38, %27 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
        pto.tmov ins(%31 : !pto.tile_buf<mat, 128x512xf16, slayout=col_major>) outs(%32 : !pto.tile_buf<right, 128x512xf16, slayout=col_major>)
        pto.tmatmul ins(%30, %32 : !pto.tile_buf<left, 32x128xf16, slayout=row_major>, !pto.tile_buf<right, 128x512xf16, slayout=col_major>) outs(%33 : !pto.tile_buf<acc, 32x512xf32, blayout=col_major, slayout=row_major, fractal=1024>)
        pto.tpush(%33, %25 : !pto.tile_buf<acc, 32x512xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
        %58 = arith.muli %arg5, %c2 : index
        %59 = arith.addi %58, %c1 : index
        %c2_8 = arith.constant 2 : index
        %60 = arith.addi %59, %c2_8 : index
        %61 = arith.muli %60, %c512 : index
        %62 = pto.partition_view %40, offsets = [%c0, %61], sizes = [%c128, %c512] : !pto.tensor_view<?x?xf16>
        pto.tload ins(%62 : !pto.partition_tensor_view<128x512xf16>) outs(%31 : !pto.tile_buf<mat, 128x512xf16, slayout=col_major>)
        %63 = pto.tpop_from_aiv {id = 30, split = 1} -> !pto.tile_buf<mat, 32x512xf16, blayout=col_major, slayout=row_major>
        pto.tmov ins(%63 : !pto.tile_buf<mat, 32x512xf16, blayout=col_major, slayout=row_major>) outs(%35 : !pto.tile_buf<left, 32x512xf16, slayout=row_major>)
        pto.tfree_from_aiv {id = 30, split = 1}
        pto.tmov ins(%36 : !pto.tile_buf<mat, 512x128xf16, blayout=col_major, slayout=row_major>) outs(%37 : !pto.tile_buf<right, 512x128xf16, slayout=col_major>)
        %64 = arith.addi %59, %c1 : index
        %65 = arith.muli %64, %c512 : index
        %66 = pto.partition_view %41, offsets = [%65, %c0], sizes = [%c512, %c128] : !pto.tensor_view<?x?xf16>
        pto.tload ins(%66 : !pto.partition_tensor_view<512x128xf16>) outs(%36 : !pto.tile_buf<mat, 512x128xf16, blayout=col_major, slayout=row_major>)
        pto.tmatmul ins(%35, %37 : !pto.tile_buf<left, 32x512xf16, slayout=row_major>, !pto.tile_buf<right, 512x128xf16, slayout=col_major>) outs(%38 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>)
        pto.tpush(%38, %27 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
        pto.tmov ins(%31 : !pto.tile_buf<mat, 128x512xf16, slayout=col_major>) outs(%32 : !pto.tile_buf<right, 128x512xf16, slayout=col_major>)
        pto.tmatmul ins(%30, %32 : !pto.tile_buf<left, 32x128xf16, slayout=row_major>, !pto.tile_buf<right, 128x512xf16, slayout=col_major>) outs(%33 : !pto.tile_buf<acc, 32x512xf32, blayout=col_major, slayout=row_major, fractal=1024>)
        pto.tpush(%33, %25 : !pto.tile_buf<acc, 32x512xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
      }
      %47 = pto.tpop_from_aiv {id = 30, split = 1} -> !pto.tile_buf<mat, 32x512xf16, blayout=col_major, slayout=row_major>
      pto.tmov ins(%47 : !pto.tile_buf<mat, 32x512xf16, blayout=col_major, slayout=row_major>) outs(%35 : !pto.tile_buf<left, 32x512xf16, slayout=row_major>)
      pto.tfree_from_aiv {id = 30, split = 1}
      pto.tmov ins(%36 : !pto.tile_buf<mat, 512x128xf16, blayout=col_major, slayout=row_major>) outs(%37 : !pto.tile_buf<right, 512x128xf16, slayout=col_major>)
      %c7680 = arith.constant 7680 : index
      %48 = pto.partition_view %41, offsets = [%c7680, %c0], sizes = [%c512, %c128] : !pto.tensor_view<?x?xf16>
      pto.tload ins(%48 : !pto.partition_tensor_view<512x128xf16>) outs(%36 : !pto.tile_buf<mat, 512x128xf16, blayout=col_major, slayout=row_major>)
      pto.tmatmul ins(%35, %37 : !pto.tile_buf<left, 32x512xf16, slayout=row_major>, !pto.tile_buf<right, 512x128xf16, slayout=col_major>) outs(%38 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>)
      pto.tpush(%38, %27 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
      %49 = pto.tpop_from_aiv {id = 30, split = 1} -> !pto.tile_buf<mat, 32x512xf16, blayout=col_major, slayout=row_major>
      pto.tmov ins(%49 : !pto.tile_buf<mat, 32x512xf16, blayout=col_major, slayout=row_major>) outs(%35 : !pto.tile_buf<left, 32x512xf16, slayout=row_major>)
      pto.tfree_from_aiv {id = 30, split = 1}
      pto.tmov ins(%36 : !pto.tile_buf<mat, 512x128xf16, blayout=col_major, slayout=row_major>) outs(%37 : !pto.tile_buf<right, 512x128xf16, slayout=col_major>)
      pto.tmatmul ins(%35, %37 : !pto.tile_buf<left, 32x512xf16, slayout=row_major>, !pto.tile_buf<right, 512x128xf16, slayout=col_major>) outs(%38 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>)
      pto.tpush(%38, %27 : !pto.tile_buf<acc, 32x128xf32, blayout=col_major, slayout=row_major, fractal=1024>, !pto.pipe) {split = 1}
    }
    return
  }
  func.func @vector_kernel(%arg0: !pto.ptr<f32>, %arg1: !pto.ptr<f32>) attributes {pto.kernel_kind = #pto.kernel_kind<vector>} {
    %c0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    %c2 = arith.constant 2 : index
    %c32 = arith.constant 32 : index
    %c16 = arith.constant 16 : index
    %c128 = arith.constant 128 : index
    %c16_0 = arith.constant 16 : index
    %c96 = arith.constant 96 : index
    %0 = pto.get_block_num
    %1 = arith.index_cast %0 : i64 to index
    %2 = pto.get_block_idx
    %3 = arith.index_cast %2 : i64 to index
    %4 = arith.divsi %c96, %1 : index
    %5 = arith.remsi %c96, %1 : index
    %6 = arith.addi %4, %c1 : index
    %7 = arith.muli %3, %6 : index
    %8 = arith.addi %4, %c1 : index
    %9 = arith.muli %5, %8 : index
    %10 = arith.subi %3, %5 : index
    %11 = arith.muli %10, %4 : index
    %12 = arith.addi %9, %11 : index
    %13 = arith.cmpi slt, %3, %5 : index
    %14 = arith.select %13, %7, %12 : index
    %15 = arith.cmpi slt, %3, %5 : index
    %16 = arith.addi %4, %c1 : index
    %17 = arith.select %15, %16, %4 : index
    %18 = arith.addi %14, %17 : index
    %c229376 = arith.constant 229376 : index
    %19 = arith.muli %3, %c229376 : index
    %20 = pto.addptr %arg0, %19 : <f32> -> <f32>
    %c0_1 = arith.constant 0 : index
    %21 = pto.addptr %20, %c0_1 : <f32> -> <f32>
    %c131072 = arith.constant 131072 : index
    %22 = pto.addptr %20, %c131072 : <f32> -> <f32>
    %c163840 = arith.constant 163840 : index
    %23 = pto.addptr %20, %c163840 : <f32> -> <f32>
    %24 = pto.reserve_buffer{name = "fa_qk_c2v_fifo", size = 65536, location = <vec>, auto = false, base = 0} -> i32
    %25 = pto.initialize_l2g2l_pipe{dir_mask = 1, slot_size = 65536, slot_num = 8, local_slot_num = 1} (%21 : !pto.ptr<f32>, %24 : i32) -> !pto.pipe
    %26 = pto.reserve_buffer{name = "fa_pv_c2v_fifo", size = 16384, location = <vec>, auto = false, base = 65536} -> i32
    %27 = pto.initialize_l2g2l_pipe{dir_mask = 1, slot_size = 16384, slot_num = 8, local_slot_num = 1} (%22 : !pto.ptr<f32>, %26 : i32) -> !pto.pipe
    %28 = pto.import_reserved_buffer{name = "fa_p_v2c_fifo", peer_func = @cube_kernel} -> i32
    %c0_i32 = arith.constant 0 : i32
    pto.aiv_initialize_pipe {id = 30, dir_mask = 2, slot_size = 32768, nosplit = false}(gm_slot_buffer = %23 : !pto.ptr<f32>, c2v_consumer_buf = %c0_i32 : i32, v2c_consumer_buf = %28 : i32)
    %29 = pto.get_subblock_idx
    %30 = arith.index_cast %29 : i64 to index
    %31 = arith.muli %30, %c16 : index
    %c114688_i64 = arith.constant 114688 : i64
    %32 = pto.alloc_tile addr = %c114688_i64 : !pto.tile_buf<vec, 16x512xf32>
    %c147456_i64 = arith.constant 147456 : i64
    %33 = pto.alloc_tile addr = %c147456_i64 : !pto.tile_buf<vec, 16x512xf32>
    %c163840_i64 = arith.constant 163840 : i64
    %34 = pto.alloc_tile addr = %c163840_i64 : !pto.tile_buf<vec, 16x512xf16>
    %c180224_i64 = arith.constant 180224 : i64
    %35 = pto.alloc_tile addr = %c180224_i64 : !pto.tile_buf<vec, 16x128xf32>
    %c188416_i64 = arith.constant 188416 : i64
    %36 = pto.alloc_tile addr = %c188416_i64 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>
    %c188480_i64 = arith.constant 188480 : i64
    %37 = pto.alloc_tile addr = %c188480_i64 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>
    %c188544_i64 = arith.constant 188544 : i64
    %38 = pto.alloc_tile addr = %c188544_i64 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>
    %c188608_i64 = arith.constant 188608 : i64
    %39 = pto.alloc_tile addr = %c188608_i64 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>
    %c188672_i64 = arith.constant 188672 : i64
    %40 = pto.alloc_tile addr = %c188672_i64 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>
    %c188736_i64 = arith.constant 188736 : i64
    %41 = pto.alloc_tile addr = %c188736_i64 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>
    %cst = arith.constant 0.0883883461 : f32
    %cst_2 = arith.constant 1.000000e+00 : f32
    %c3072 = arith.constant 3072 : index
    %42 = pto.make_tensor_view %arg1, shape = [%c3072, %c128], strides = [%c128, %c1] : !pto.tensor_view<?x?xf32>
    scf.for %arg2 = %14 to %18 step %c1 {
      %43 = arith.muli %arg2, %c32 : index
      %c81920_i64 = arith.constant 81920 : i64
      %44 = pto.alloc_tile addr = %c81920_i64 : !pto.tile_buf<vec, 16x512xf32>
      pto.tpop(%44, %25 : !pto.tile_buf<vec, 16x512xf32>, !pto.pipe) {split = 1}
      pto.tmuls ins(%44, %cst : !pto.tile_buf<vec, 16x512xf32>, f32) outs(%44 : !pto.tile_buf<vec, 16x512xf32>)
      pto.trowmax ins(%44, %32 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x512xf32>) outs(%37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
      %45 = pto.treshape %37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %46 = pto.treshape %36 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %47 = pto.treshape %40 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %48 = pto.treshape %38 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %49 = pto.treshape %39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      pto.trowexpandsub ins(%44, %37 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%33 : !pto.tile_buf<vec, 16x512xf32>)
      pto.tmuls ins(%45, %cst_2 : !pto.tile_buf<vec, 1x16xf32>, f32) outs(%46 : !pto.tile_buf<vec, 1x16xf32>)
      pto.texp ins(%33 : !pto.tile_buf<vec, 16x512xf32>) outs(%33 : !pto.tile_buf<vec, 16x512xf32>)
      pto.trowsum ins(%33, %32 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x512xf32>) outs(%38 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
      pto.tcvt ins(%33 {rmode = #pto<round_mode CAST_RINT>} : !pto.tile_buf<vec, 16x512xf32>) outs(%34 : !pto.tile_buf<vec, 16x512xf16>)
      pto.tpush_to_aic(%34 : !pto.tile_buf<vec, 16x512xf16>) {id = 30, split = 1}
      pto.tfree(%25 : !pto.pipe) {split = 1}
      %c81920_i64_3 = arith.constant 81920 : i64
      %50 = pto.alloc_tile addr = %c81920_i64_3 : !pto.tile_buf<vec, 16x512xf32>
      pto.tpop(%50, %25 : !pto.tile_buf<vec, 16x512xf32>, !pto.pipe) {split = 1}
      pto.tmuls ins(%50, %cst : !pto.tile_buf<vec, 16x512xf32>, f32) outs(%50 : !pto.tile_buf<vec, 16x512xf32>)
      pto.trowmax ins(%50, %32 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x512xf32>) outs(%37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
      %51 = pto.treshape %37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %52 = pto.treshape %36 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %53 = pto.treshape %41 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %54 = pto.treshape %38 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %55 = pto.treshape %39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      pto.tmax ins(%51, %52 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%51 : !pto.tile_buf<vec, 1x16xf32>)
      pto.tsub ins(%52, %51 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%53 : !pto.tile_buf<vec, 1x16xf32>)
      pto.tmuls ins(%51, %cst_2 : !pto.tile_buf<vec, 1x16xf32>, f32) outs(%52 : !pto.tile_buf<vec, 1x16xf32>)
      pto.trowexpandsub ins(%50, %37 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%33 : !pto.tile_buf<vec, 16x512xf32>)
      pto.texp ins(%53 : !pto.tile_buf<vec, 1x16xf32>) outs(%53 : !pto.tile_buf<vec, 1x16xf32>)
      pto.texp ins(%33 : !pto.tile_buf<vec, 16x512xf32>) outs(%33 : !pto.tile_buf<vec, 16x512xf32>)
      pto.tmul ins(%54, %53 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%54 : !pto.tile_buf<vec, 1x16xf32>)
      pto.trowsum ins(%33, %32 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x512xf32>) outs(%39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
      pto.tadd ins(%54, %55 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%54 : !pto.tile_buf<vec, 1x16xf32>)
      pto.tcvt ins(%33 {rmode = #pto<round_mode CAST_RINT>} : !pto.tile_buf<vec, 16x512xf32>) outs(%34 : !pto.tile_buf<vec, 16x512xf16>)
      pto.tpush_to_aic(%34 : !pto.tile_buf<vec, 16x512xf16>) {id = 30, split = 1}
      pto.tfree(%25 : !pto.pipe) {split = 1}
      %c114688_i64_4 = arith.constant 114688 : i64
      %56 = pto.alloc_tile addr = %c114688_i64_4 : !pto.tile_buf<vec, 16x128xf32>
      pto.tpop(%56, %27 : !pto.tile_buf<vec, 16x128xf32>, !pto.pipe) {split = 1}
      pto.tmov ins(%56 : !pto.tile_buf<vec, 16x128xf32>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
      pto.tfree(%27 : !pto.pipe) {split = 1}
      %c81920_i64_5 = arith.constant 81920 : i64
      %57 = pto.alloc_tile addr = %c81920_i64_5 : !pto.tile_buf<vec, 16x512xf32>
      pto.tpop(%57, %25 : !pto.tile_buf<vec, 16x512xf32>, !pto.pipe) {split = 1}
      pto.tmuls ins(%57, %cst : !pto.tile_buf<vec, 16x512xf32>, f32) outs(%57 : !pto.tile_buf<vec, 16x512xf32>)
      pto.trowmax ins(%57, %32 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x512xf32>) outs(%37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
      %58 = pto.treshape %37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %59 = pto.treshape %36 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %60 = pto.treshape %40 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %61 = pto.treshape %38 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %62 = pto.treshape %39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      pto.tmax ins(%58, %59 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%58 : !pto.tile_buf<vec, 1x16xf32>)
      pto.tsub ins(%59, %58 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%60 : !pto.tile_buf<vec, 1x16xf32>)
      pto.tmuls ins(%58, %cst_2 : !pto.tile_buf<vec, 1x16xf32>, f32) outs(%59 : !pto.tile_buf<vec, 1x16xf32>)
      pto.trowexpandsub ins(%57, %37 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%33 : !pto.tile_buf<vec, 16x512xf32>)
      pto.texp ins(%60 : !pto.tile_buf<vec, 1x16xf32>) outs(%60 : !pto.tile_buf<vec, 1x16xf32>)
      pto.texp ins(%33 : !pto.tile_buf<vec, 16x512xf32>) outs(%33 : !pto.tile_buf<vec, 16x512xf32>)
      pto.tmul ins(%61, %60 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%61 : !pto.tile_buf<vec, 1x16xf32>)
      pto.trowsum ins(%33, %32 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x512xf32>) outs(%39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
      pto.tadd ins(%61, %62 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%61 : !pto.tile_buf<vec, 1x16xf32>)
      pto.tcvt ins(%33 {rmode = #pto<round_mode CAST_RINT>} : !pto.tile_buf<vec, 16x512xf32>) outs(%34 : !pto.tile_buf<vec, 16x512xf16>)
      pto.tpush_to_aic(%34 : !pto.tile_buf<vec, 16x512xf16>) {id = 30, split = 1}
      pto.tfree(%25 : !pto.pipe) {split = 1}
      %c114688_i64_6 = arith.constant 114688 : i64
      %63 = pto.alloc_tile addr = %c114688_i64_6 : !pto.tile_buf<vec, 16x128xf32>
      pto.tpop(%63, %27 : !pto.tile_buf<vec, 16x128xf32>, !pto.pipe) {split = 1}
      pto.trowexpandmul ins(%35, %41 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
      pto.tadd ins(%35, %63 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x128xf32>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
      pto.tfree(%27 : !pto.pipe) {split = 1}
      %c81920_i64_7 = arith.constant 81920 : i64
      %64 = pto.alloc_tile addr = %c81920_i64_7 : !pto.tile_buf<vec, 16x512xf32>
      pto.tpop(%64, %25 : !pto.tile_buf<vec, 16x512xf32>, !pto.pipe) {split = 1}
      pto.tmuls ins(%64, %cst : !pto.tile_buf<vec, 16x512xf32>, f32) outs(%64 : !pto.tile_buf<vec, 16x512xf32>)
      pto.trowmax ins(%64, %32 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x512xf32>) outs(%37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
      %65 = pto.treshape %37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %66 = pto.treshape %36 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %67 = pto.treshape %41 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %68 = pto.treshape %38 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      %69 = pto.treshape %39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
      pto.tmax ins(%65, %66 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%65 : !pto.tile_buf<vec, 1x16xf32>)
      pto.tsub ins(%66, %65 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%67 : !pto.tile_buf<vec, 1x16xf32>)
      pto.tmuls ins(%65, %cst_2 : !pto.tile_buf<vec, 1x16xf32>, f32) outs(%66 : !pto.tile_buf<vec, 1x16xf32>)
      pto.trowexpandsub ins(%64, %37 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%33 : !pto.tile_buf<vec, 16x512xf32>)
      pto.texp ins(%67 : !pto.tile_buf<vec, 1x16xf32>) outs(%67 : !pto.tile_buf<vec, 1x16xf32>)
      pto.texp ins(%33 : !pto.tile_buf<vec, 16x512xf32>) outs(%33 : !pto.tile_buf<vec, 16x512xf32>)
      pto.tmul ins(%68, %67 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%68 : !pto.tile_buf<vec, 1x16xf32>)
      pto.trowsum ins(%33, %32 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x512xf32>) outs(%39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
      pto.tadd ins(%68, %69 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%68 : !pto.tile_buf<vec, 1x16xf32>)
      pto.tcvt ins(%33 {rmode = #pto<round_mode CAST_RINT>} : !pto.tile_buf<vec, 16x512xf32>) outs(%34 : !pto.tile_buf<vec, 16x512xf16>)
      pto.tpush_to_aic(%34 : !pto.tile_buf<vec, 16x512xf16>) {id = 30, split = 1}
      pto.tfree(%25 : !pto.pipe) {split = 1}
      %c7 = arith.constant 7 : index
      scf.for %arg3 = %c1 to %c7 step %c1 {
        %c114688_i64_10 = arith.constant 114688 : i64
        %74 = pto.alloc_tile addr = %c114688_i64_10 : !pto.tile_buf<vec, 16x128xf32>
        pto.tpop(%74, %27 : !pto.tile_buf<vec, 16x128xf32>, !pto.pipe) {split = 1}
        pto.trowexpandmul ins(%35, %40 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
        pto.tadd ins(%35, %74 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x128xf32>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
        pto.tfree(%27 : !pto.pipe) {split = 1}
        %c81920_i64_11 = arith.constant 81920 : i64
        %75 = pto.alloc_tile addr = %c81920_i64_11 : !pto.tile_buf<vec, 16x512xf32>
        pto.tpop(%75, %25 : !pto.tile_buf<vec, 16x512xf32>, !pto.pipe) {split = 1}
        pto.tmuls ins(%75, %cst : !pto.tile_buf<vec, 16x512xf32>, f32) outs(%75 : !pto.tile_buf<vec, 16x512xf32>)
        pto.trowmax ins(%75, %32 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x512xf32>) outs(%37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
        %76 = pto.treshape %37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
        %77 = pto.treshape %36 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
        %78 = pto.treshape %40 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
        %79 = pto.treshape %38 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
        %80 = pto.treshape %39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
        pto.tmax ins(%76, %77 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%76 : !pto.tile_buf<vec, 1x16xf32>)
        pto.tsub ins(%77, %76 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%78 : !pto.tile_buf<vec, 1x16xf32>)
        pto.tmuls ins(%76, %cst_2 : !pto.tile_buf<vec, 1x16xf32>, f32) outs(%77 : !pto.tile_buf<vec, 1x16xf32>)
        pto.trowexpandsub ins(%75, %37 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%33 : !pto.tile_buf<vec, 16x512xf32>)
        pto.texp ins(%78 : !pto.tile_buf<vec, 1x16xf32>) outs(%78 : !pto.tile_buf<vec, 1x16xf32>)
        pto.texp ins(%33 : !pto.tile_buf<vec, 16x512xf32>) outs(%33 : !pto.tile_buf<vec, 16x512xf32>)
        pto.tmul ins(%79, %78 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%79 : !pto.tile_buf<vec, 1x16xf32>)
        pto.trowsum ins(%33, %32 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x512xf32>) outs(%39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
        pto.tadd ins(%79, %80 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%79 : !pto.tile_buf<vec, 1x16xf32>)
        pto.tcvt ins(%33 {rmode = #pto<round_mode CAST_RINT>} : !pto.tile_buf<vec, 16x512xf32>) outs(%34 : !pto.tile_buf<vec, 16x512xf16>)
        pto.tpush_to_aic(%34 : !pto.tile_buf<vec, 16x512xf16>) {id = 30, split = 1}
        pto.tfree(%25 : !pto.pipe) {split = 1}
        %c114688_i64_12 = arith.constant 114688 : i64
        %81 = pto.alloc_tile addr = %c114688_i64_12 : !pto.tile_buf<vec, 16x128xf32>
        pto.tpop(%81, %27 : !pto.tile_buf<vec, 16x128xf32>, !pto.pipe) {split = 1}
        pto.trowexpandmul ins(%35, %41 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
        pto.tadd ins(%35, %81 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x128xf32>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
        pto.tfree(%27 : !pto.pipe) {split = 1}
        %c81920_i64_13 = arith.constant 81920 : i64
        %82 = pto.alloc_tile addr = %c81920_i64_13 : !pto.tile_buf<vec, 16x512xf32>
        pto.tpop(%82, %25 : !pto.tile_buf<vec, 16x512xf32>, !pto.pipe) {split = 1}
        pto.tmuls ins(%82, %cst : !pto.tile_buf<vec, 16x512xf32>, f32) outs(%82 : !pto.tile_buf<vec, 16x512xf32>)
        pto.trowmax ins(%82, %32 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x512xf32>) outs(%37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
        %83 = pto.treshape %37 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
        %84 = pto.treshape %36 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
        %85 = pto.treshape %41 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
        %86 = pto.treshape %38 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
        %87 = pto.treshape %39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major> -> !pto.tile_buf<vec, 1x16xf32>
        pto.tmax ins(%83, %84 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%83 : !pto.tile_buf<vec, 1x16xf32>)
        pto.tsub ins(%84, %83 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%85 : !pto.tile_buf<vec, 1x16xf32>)
        pto.tmuls ins(%83, %cst_2 : !pto.tile_buf<vec, 1x16xf32>, f32) outs(%84 : !pto.tile_buf<vec, 1x16xf32>)
        pto.trowexpandsub ins(%82, %37 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%33 : !pto.tile_buf<vec, 16x512xf32>)
        pto.texp ins(%85 : !pto.tile_buf<vec, 1x16xf32>) outs(%85 : !pto.tile_buf<vec, 1x16xf32>)
        pto.texp ins(%33 : !pto.tile_buf<vec, 16x512xf32>) outs(%33 : !pto.tile_buf<vec, 16x512xf32>)
        pto.tmul ins(%86, %85 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%86 : !pto.tile_buf<vec, 1x16xf32>)
        pto.trowsum ins(%33, %32 : !pto.tile_buf<vec, 16x512xf32>, !pto.tile_buf<vec, 16x512xf32>) outs(%39 : !pto.tile_buf<vec, 16x1xf32, blayout=col_major>)
        pto.tadd ins(%86, %87 : !pto.tile_buf<vec, 1x16xf32>, !pto.tile_buf<vec, 1x16xf32>) outs(%86 : !pto.tile_buf<vec, 1x16xf32>)
        pto.tcvt ins(%33 {rmode = #pto<round_mode CAST_RINT>} : !pto.tile_buf<vec, 16x512xf32>) outs(%34 : !pto.tile_buf<vec, 16x512xf16>)
        pto.tpush_to_aic(%34 : !pto.tile_buf<vec, 16x512xf16>) {id = 30, split = 1}
        pto.tfree(%25 : !pto.pipe) {split = 1}
      }
      %c114688_i64_8 = arith.constant 114688 : i64
      %70 = pto.alloc_tile addr = %c114688_i64_8 : !pto.tile_buf<vec, 16x128xf32>
      pto.tpop(%70, %27 : !pto.tile_buf<vec, 16x128xf32>, !pto.pipe) {split = 1}
      pto.trowexpandmul ins(%35, %40 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
      pto.tadd ins(%35, %70 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x128xf32>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
      pto.tfree(%27 : !pto.pipe) {split = 1}
      %c114688_i64_9 = arith.constant 114688 : i64
      %71 = pto.alloc_tile addr = %c114688_i64_9 : !pto.tile_buf<vec, 16x128xf32>
      pto.tpop(%71, %27 : !pto.tile_buf<vec, 16x128xf32>, !pto.pipe) {split = 1}
      pto.trowexpandmul ins(%35, %41 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
      pto.tadd ins(%35, %71 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x128xf32>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
      pto.tfree(%27 : !pto.pipe) {split = 1}
      pto.trowexpanddiv ins(%35, %38 : !pto.tile_buf<vec, 16x128xf32>, !pto.tile_buf<vec, 16x1xf32, blayout=col_major>) outs(%35 : !pto.tile_buf<vec, 16x128xf32>)
      %72 = arith.addi %43, %31 : index
      %73 = pto.partition_view %42, offsets = [%72, %c0], sizes = [%c16, %c128] : !pto.tensor_view<?x?xf32>
      pto.tstore ins(%35 : !pto.tile_buf<vec, 16x128xf32>) outs(%73 : !pto.partition_tensor_view<16x128xf32>)
    }
    return
  }
  func.func @call_both(%arg0: memref<256xi64>, %arg1: !pto.ptr<f32>, %arg2: !pto.ptr<f16>, %arg3: !pto.ptr<f16>, %arg4: !pto.ptr<f16>, %arg5: !pto.ptr<f32>) attributes {pto.entry} {
    pto.set_ffts %arg0 : memref<256xi64>
    call @cube_kernel(%arg1, %arg2, %arg3, %arg4) : (!pto.ptr<f32>, !pto.ptr<f16>, !pto.ptr<f16>, !pto.ptr<f16>) -> ()
    call @vector_kernel(%arg1, %arg5) : (!pto.ptr<f32>, !pto.ptr<f32>) -> ()
    return
  }
}

[zhangstevenunity]

A2/A3 Right tile of L0B，L0B size is 64KB. But the S1_TILE = 512 version, there is Right 128x512xf16 ：128 * 512 * 2 = 131072 bytes = 128KB which is not fit in L0B size.

[MirkoDeVita98]

Working cpp tile_s1 =256

#include "pto/pto-inst.hpp"
using namespace pto;

enum class PTOAutoSyncTailMode : int {
  kBarrierAll = 0,
  kSetWaitMte3ToSEvent0 = 1,
};

static AICORE inline void ptoas_auto_sync_tail(
    PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) {
  switch (mode) {
  case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0:
    set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0);
    wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0);
    break;
  case PTOAutoSyncTailMode::kBarrierAll:
  default:
    pipe_barrier(PIPE_ALL);
    break;
  }
}

AICORE void cube_kernel(__gm__ float* v1, __gm__ half* v2, __gm__ half* v3, __gm__ half* v4) {
  unsigned v5 = 7936;
  unsigned v6 = 256;
  unsigned v7 = 0;
  const int32_t v8 = 96;
  const int32_t v9 = 256;
  const int32_t v10 = 128;
  const int32_t v11 = 32;
  const int32_t v12 = 2;
  const int32_t v13 = 1;
  const int32_t v14 = 131072;
  const int32_t v15 = 65536;
  const int32_t v16 = 98304;
  const int64_t v17 = 0;
  const int64_t v18 = 8192;
  const int64_t v19 = 90112;
  const int64_t v20 = 32768;
  const int32_t v21 = 15;
  const int32_t v22 = 0;
  const int32_t v23 = 32768;
  const int32_t v24 = 262144;
  using T = float;

  #if defined(__DAV_CUBE__)
  size_t v25 = (size_t) v13;
  int64_t v26 = get_block_num();
  int32_t v27 = (int32_t) ((int64_t) v26);
  int64_t v28 = get_block_idx();
  int32_t v29 = (int32_t) ((int64_t) v28);
  int32_t v30 = v8 / v27;
  int32_t v31 = v8 % v27;
  int32_t v32 = (int32_t) ((uint32_t) v30 + (uint32_t) v13);
  bool v33 = v29 < v31;
  int32_t v34 = v33 ? (int32_t) ((uint32_t) v29 * (uint32_t) v32) : (int32_t) ((uint32_t) ((int32_t) (uint32_t) v31 * (uint32_t) v32) + (uint32_t) ((int32_t) (uint32_t) ((int32_t) (uint32_t) v29 - (uint32_t) v31) * (uint32_t) v30));
  int32_t v35 = (int32_t) ((uint32_t) v29 * (uint32_t) v14);
  __gm__ float* v36 = v1 + v35;
  
  int32_t v37 = (int32_t) ((uint32_t) v35 + (uint32_t) v15);
  __gm__ float* v38 = v1 + v37;
  
  int32_t v39 = (int32_t) ((uint32_t) v35 + (uint32_t) v16);
  __gm__ float* v40 = v1 + v39;
  
  auto v41 = TPipe<0, Direction::DIR_C2V, 32768, 8, 1, false>(v36, v22, v22);
  auto v42 = TPipe<2, Direction::DIR_C2V, 16384, 8, 1, false>(v38, v23, v22);
  auto v43 = TPipe<4, Direction::DIR_V2C, 16384, 8, 1, false>(v40, v22, v24);
  Tile<TileType::Mat, half, 32, 128, BLayout::ColMajor, 32, 128, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> v44;
  TASSIGN(v44, v17);
  Tile<TileType::Left, half, 32, 128, BLayout::RowMajor, 32, 128, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> v45;
  TASSIGN(v45, v17);
  Tile<TileType::Mat, half, 128, 256, BLayout::RowMajor, 128, 256, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> v46;
  TASSIGN(v46, v18);
  Tile<TileType::Right, half, 128, 256, BLayout::RowMajor, 128, 256, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> v47;
  TASSIGN(v47, v17);
  Tile<TileType::Acc, float, 32, 256, BLayout::ColMajor, 32, 256, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null> v48;
  TASSIGN(v48, v17);
  Tile<TileType::Left, half, 32, 256, BLayout::RowMajor, 32, 256, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> v49;
  TASSIGN(v49, v18);
  Tile<TileType::Mat, half, 256, 128, BLayout::ColMajor, 256, 128, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> v50;
  TASSIGN(v50, v19);
  Tile<TileType::Right, half, 256, 128, BLayout::RowMajor, 256, 128, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> v51;
  TASSIGN(v51, v17);
  Tile<TileType::Acc, float, 32, 128, BLayout::ColMajor, 32, 128, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null> v52;
  TASSIGN(v52, v20);
  set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0);
  set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2);
  set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3);
  set_flag(PIPE_FIX, PIPE_M, EVENT_ID2);
  set_flag(PIPE_FIX, PIPE_M, EVENT_ID3);
  for (size_t v53 = (size_t) v34; v53 < ((size_t) ((int32_t) (uint32_t) v34 + (uint32_t) (v33 ? v32 : v30))); v53 += v25) {
    pto::Shape<1, 1, 1, 32, 128> v54 = pto::Shape<1, 1, 1, 32, 128>();
    pto::Stride<4096, 4096, 4096, 128, 1> v55 = pto::Stride<4096, 4096, 4096, 128, 1>();
    GlobalTensor<half, pto::Shape<1, 1, 1, 32, 128>, pto::Stride<4096, 4096, 4096, 128, 1>, pto::Layout::ND> v56 = GlobalTensor<half, pto::Shape<1, 1, 1, 32, 128>, pto::Stride<4096, 4096, 4096, 128, 1>, pto::Layout::ND>(v2 + (v7 + (unsigned) ((int32_t) (uint32_t) ((int32_t) v53) * (uint32_t) v11) * (unsigned) v10 + v7 * (unsigned) v13), v54, v55);
    TLOAD(v44, v56);
    set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0);
    wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0);
    TMOV(v45, v44);
    pto::Shape<1, 1, 1, 128, 256> v57 = pto::Shape<1, 1, 1, 128, 256>();
    pto::Stride<128, 128, 128, 1, 128> v58 = pto::Stride<128, 128, 128, 1, 128>();
    GlobalTensor<half, pto::Shape<1, 1, 1, 128, 256>, pto::Stride<128, 128, 128, 1, 128>, pto::Layout::DN> v59 = GlobalTensor<half, pto::Shape<1, 1, 1, 128, 256>, pto::Stride<128, 128, 128, 1, 128>, pto::Layout::DN>(v3 + (v7 + v7 * (unsigned) v13 + v7 * (unsigned) v10), v57, v58);
    TLOAD(v46, v59);
    set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1);
    wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1);
    wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0);
    TMOV(v47, v46);
    set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0);
    set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0);
    wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0);
    TMATMUL(v48, v45, v47);
    set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1);
    set_flag(PIPE_M, PIPE_FIX, EVENT_ID0);
    wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0);
    TPUSH<TPipe<0, Direction::DIR_C2V, 32768, 8, 1, false>, Tile<TileType::Acc, float, 32, 256, BLayout::ColMajor, 32, 256, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v41, v48);
    set_flag(PIPE_FIX, PIPE_M, EVENT_ID0);
    pto::Shape<1, 1, 1, 128, 256> v60 = pto::Shape<1, 1, 1, 128, 256>();
    pto::Stride<128, 128, 128, 1, 128> v61 = pto::Stride<128, 128, 128, 1, 128>();
    GlobalTensor<half, pto::Shape<1, 1, 1, 128, 256>, pto::Stride<128, 128, 128, 1, 128>, pto::Layout::DN> v62 = GlobalTensor<half, pto::Shape<1, 1, 1, 128, 256>, pto::Stride<128, 128, 128, 1, 128>, pto::Layout::DN>(v3 + (v7 + v7 * (unsigned) v13 + v6 * (unsigned) v10), v60, v61);
    wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0);
    TLOAD(v46, v62);
    set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2);
    wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2);
    wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1);
    TMOV(v47, v46);
    set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1);
    set_flag(PIPE_MTE1, PIPE_M, EVENT_ID1);
    wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID1);
    wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0);
    TMATMUL(v48, v45, v47);
    set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2);
    set_flag(PIPE_M, PIPE_FIX, EVENT_ID1);
    wait_flag(PIPE_M, PIPE_FIX, EVENT_ID1);
    TPUSH<TPipe<0, Direction::DIR_C2V, 32768, 8, 1, false>, Tile<TileType::Acc, float, 32, 256, BLayout::ColMajor, 32, 256, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v41, v48);
    set_flag(PIPE_FIX, PIPE_M, EVENT_ID1);
    pto::Shape<1, 1, 1, 256, 128> v63 = pto::Shape<1, 1, 1, 256, 128>();
    pto::Stride<32768, 32768, 32768, 128, 1> v64 = pto::Stride<32768, 32768, 32768, 128, 1>();
    GlobalTensor<half, pto::Shape<1, 1, 1, 256, 128>, pto::Stride<32768, 32768, 32768, 128, 1>, pto::Layout::ND> v65 = GlobalTensor<half, pto::Shape<1, 1, 1, 256, 128>, pto::Stride<32768, 32768, 32768, 128, 1>, pto::Layout::ND>(v4 + (v7 + v7 * (unsigned) v10 + v7 * (unsigned) v13), v63, v64);
    TLOAD(v50, v65);
    wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1);
    wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2);
    wait_flag(PIPE_FIX, PIPE_M, EVENT_ID1);
    for (size_t v66 = (size_t) v22; v66 < ((size_t) v21); v66 += v25) {
      int32_t v67 = (int32_t) ((uint32_t) ((int32_t) v66) * (uint32_t) v12);
      pto::Shape<1, 1, 1, 128, 256> v68 = pto::Shape<1, 1, 1, 128, 256>();
      pto::Stride<128, 128, 128, 1, 128> v69 = pto::Stride<128, 128, 128, 1, 128>();
      GlobalTensor<half, pto::Shape<1, 1, 1, 128, 256>, pto::Stride<128, 128, 128, 1, 128>, pto::Layout::DN> v70 = GlobalTensor<half, pto::Shape<1, 1, 1, 128, 256>, pto::Stride<128, 128, 128, 1, 128>, pto::Layout::DN>(v3 + (v7 + v7 * (unsigned) v13 + (unsigned) ((int32_t) (uint32_t) ((int32_t) (uint32_t) v67 + (uint32_t) v12) * (uint32_t) v9) * (unsigned) v10), v68, v69);
      wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2);
      TLOAD(v46, v70);
      Tile<TileType::Mat, half, 32, 256, BLayout::ColMajor, 32, 256, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> v71;
      TPOP<TPipe<4, Direction::DIR_V2C, 16384, 8, 1, false>, Tile<TileType::Mat, half, 32, 256, BLayout::ColMajor, 32, 256, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v43, v71);
      set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3);
      wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3);
      TMOV(v49, v71);
      TFREE<TPipe<4, Direction::DIR_V2C, 16384, 8, 1, false>, TileSplitAxis::TILE_UP_DOWN>(v43);
      wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3);
      TMOV(v51, v50);
      set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3);
      set_flag(PIPE_MTE1, PIPE_M, EVENT_ID2);
      int32_t v72 = (int32_t) ((uint32_t) v67 + (uint32_t) v13);
      pto::Shape<1, 1, 1, 256, 128> v73 = pto::Shape<1, 1, 1, 256, 128>();
      pto::Stride<32768, 32768, 32768, 128, 1> v74 = pto::Stride<32768, 32768, 32768, 128, 1>();
      GlobalTensor<half, pto::Shape<1, 1, 1, 256, 128>, pto::Stride<32768, 32768, 32768, 128, 1>, pto::Layout::ND> v75 = GlobalTensor<half, pto::Shape<1, 1, 1, 256, 128>, pto::Stride<32768, 32768, 32768, 128, 1>, pto::Layout::ND>(v4 + (v7 + (unsigned) ((int32_t) (uint32_t) v72 * (uint32_t) v9) * (unsigned) v10 + v7 * (unsigned) v13), v73, v74);
      wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3);
      TLOAD(v50, v75);
      wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID2);
      wait_flag(PIPE_FIX, PIPE_M, EVENT_ID2);
      TMATMUL(v52, v49, v51);
      set_flag(PIPE_M, PIPE_FIX, EVENT_ID2);
      set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4);
      wait_flag(PIPE_M, PIPE_FIX, EVENT_ID2);
      TPUSH<TPipe<2, Direction::DIR_C2V, 16384, 8, 1, false>, Tile<TileType::Acc, float, 32, 128, BLayout::ColMajor, 32, 128, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v42, v52);
      set_flag(PIPE_FIX, PIPE_M, EVENT_ID4);
      wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4);
      TMOV(v47, v46);
      set_flag(PIPE_MTE1, PIPE_M, EVENT_ID3);
      set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID4);
      wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID3);
      wait_flag(PIPE_FIX, PIPE_M, EVENT_ID3);
      TMATMUL(v48, v45, v47);
      set_flag(PIPE_M, PIPE_FIX, EVENT_ID3);
      set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5);
      wait_flag(PIPE_M, PIPE_FIX, EVENT_ID3);
      TPUSH<TPipe<0, Direction::DIR_C2V, 32768, 8, 1, false>, Tile<TileType::Acc, float, 32, 256, BLayout::ColMajor, 32, 256, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v41, v48);
      set_flag(PIPE_FIX, PIPE_M, EVENT_ID5);
      pto::Shape<1, 1, 1, 128, 256> v76 = pto::Shape<1, 1, 1, 128, 256>();
      pto::Stride<128, 128, 128, 1, 128> v77 = pto::Stride<128, 128, 128, 1, 128>();
      GlobalTensor<half, pto::Shape<1, 1, 1, 128, 256>, pto::Stride<128, 128, 128, 1, 128>, pto::Layout::DN> v78 = GlobalTensor<half, pto::Shape<1, 1, 1, 128, 256>, pto::Stride<128, 128, 128, 1, 128>, pto::Layout::DN>(v3 + (v7 + v7 * (unsigned) v13 + (unsigned) ((int32_t) (uint32_t) ((int32_t) (uint32_t) v72 + (uint32_t) v12) * (uint32_t) v9) * (unsigned) v10), v76, v77);
      wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID4);
      TLOAD(v46, v78);
      Tile<TileType::Mat, half, 32, 256, BLayout::ColMajor, 32, 256, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> v79;
      TPOP<TPipe<4, Direction::DIR_V2C, 16384, 8, 1, false>, Tile<TileType::Mat, half, 32, 256, BLayout::ColMajor, 32, 256, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v43, v79);
      set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID4);
      wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID4);
      TMOV(v49, v79);
      TFREE<TPipe<4, Direction::DIR_V2C, 16384, 8, 1, false>, TileSplitAxis::TILE_UP_DOWN>(v43);
      wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5);
      TMOV(v51, v50);
      set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID5);
      set_flag(PIPE_MTE1, PIPE_M, EVENT_ID4);
      pto::Shape<1, 1, 1, 256, 128> v80 = pto::Shape<1, 1, 1, 256, 128>();
      pto::Stride<32768, 32768, 32768, 128, 1> v81 = pto::Stride<32768, 32768, 32768, 128, 1>();
      GlobalTensor<half, pto::Shape<1, 1, 1, 256, 128>, pto::Stride<32768, 32768, 32768, 128, 1>, pto::Layout::ND> v82 = GlobalTensor<half, pto::Shape<1, 1, 1, 256, 128>, pto::Stride<32768, 32768, 32768, 128, 1>, pto::Layout::ND>(v4 + (v7 + (unsigned) ((int32_t) (uint32_t) ((int32_t) (uint32_t) v72 + (uint32_t) v13) * (uint32_t) v9) * (unsigned) v10 + v7 * (unsigned) v13), v80, v81);
      wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID5);
      TLOAD(v50, v82);
      wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID4);
      wait_flag(PIPE_FIX, PIPE_M, EVENT_ID4);
      TMATMUL(v52, v49, v51);
      set_flag(PIPE_M, PIPE_FIX, EVENT_ID4);
      set_flag(PIPE_M, PIPE_MTE1, EVENT_ID6);
      wait_flag(PIPE_M, PIPE_FIX, EVENT_ID4);
      TPUSH<TPipe<2, Direction::DIR_C2V, 16384, 8, 1, false>, Tile<TileType::Acc, float, 32, 128, BLayout::ColMajor, 32, 128, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v42, v52);
      set_flag(PIPE_FIX, PIPE_M, EVENT_ID2);
      wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID6);
      TMOV(v47, v46);
      set_flag(PIPE_MTE1, PIPE_M, EVENT_ID5);
      set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2);
      wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID5);
      wait_flag(PIPE_FIX, PIPE_M, EVENT_ID5);
      TMATMUL(v48, v45, v47);
      set_flag(PIPE_M, PIPE_FIX, EVENT_ID5);
      set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3);
      wait_flag(PIPE_M, PIPE_FIX, EVENT_ID5);
      TPUSH<TPipe<0, Direction::DIR_C2V, 32768, 8, 1, false>, Tile<TileType::Acc, float, 32, 256, BLayout::ColMajor, 32, 256, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v41, v48);
      set_flag(PIPE_FIX, PIPE_M, EVENT_ID3);
    };
    set_flag(PIPE_FIX, PIPE_M, EVENT_ID6);
    set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7);
    Tile<TileType::Mat, half, 32, 256, BLayout::ColMajor, 32, 256, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> v83;
    TPOP<TPipe<4, Direction::DIR_V2C, 16384, 8, 1, false>, Tile<TileType::Mat, half, 32, 256, BLayout::ColMajor, 32, 256, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v43, v83);
    set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID5);
    wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID5);
    TMOV(v49, v83);
    TFREE<TPipe<4, Direction::DIR_V2C, 16384, 8, 1, false>, TileSplitAxis::TILE_UP_DOWN>(v43);
    wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7);
    TMOV(v51, v50);
    set_flag(PIPE_MTE1, PIPE_M, EVENT_ID6);
    set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID6);
    pto::Shape<1, 1, 1, 256, 128> v84 = pto::Shape<1, 1, 1, 256, 128>();
    pto::Stride<32768, 32768, 32768, 128, 1> v85 = pto::Stride<32768, 32768, 32768, 128, 1>();
    GlobalTensor<half, pto::Shape<1, 1, 1, 256, 128>, pto::Stride<32768, 32768, 32768, 128, 1>, pto::Layout::ND> v86 = GlobalTensor<half, pto::Shape<1, 1, 1, 256, 128>, pto::Stride<32768, 32768, 32768, 128, 1>, pto::Layout::ND>(v4 + (v7 + v5 * (unsigned) v10 + v7 * (unsigned) v13), v84, v85);
    wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID6);
    TLOAD(v50, v86);
    wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID6);
    wait_flag(PIPE_FIX, PIPE_M, EVENT_ID6);
    TMATMUL(v52, v49, v51);
    set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1);
    set_flag(PIPE_M, PIPE_FIX, EVENT_ID6);
    wait_flag(PIPE_M, PIPE_FIX, EVENT_ID6);
    TPUSH<TPipe<2, Direction::DIR_C2V, 16384, 8, 1, false>, Tile<TileType::Acc, float, 32, 128, BLayout::ColMajor, 32, 128, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v42, v52);
    set_flag(PIPE_FIX, PIPE_M, EVENT_ID7);
    Tile<TileType::Mat, half, 32, 256, BLayout::ColMajor, 32, 256, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> v87;
    TPOP<TPipe<4, Direction::DIR_V2C, 16384, 8, 1, false>, Tile<TileType::Mat, half, 32, 256, BLayout::ColMajor, 32, 256, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v43, v87);
    set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID6);
    wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID6);
    wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1);
    TMOV(v49, v87);
    TFREE<TPipe<4, Direction::DIR_V2C, 16384, 8, 1, false>, TileSplitAxis::TILE_UP_DOWN>(v43);
    TMOV(v51, v50);
    set_flag(PIPE_MTE1, PIPE_M, EVENT_ID7);
    wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID7);
    wait_flag(PIPE_FIX, PIPE_M, EVENT_ID7);
    TMATMUL(v52, v49, v51);
    set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0);
    set_flag(PIPE_M, PIPE_FIX, EVENT_ID7);
    wait_flag(PIPE_M, PIPE_FIX, EVENT_ID7);
    TPUSH<TPipe<2, Direction::DIR_C2V, 16384, 8, 1, false>, Tile<TileType::Acc, float, 32, 128, BLayout::ColMajor, 32, 128, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v42, v52);
  }
  wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0);
  wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2);
  wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3);
  wait_flag(PIPE_FIX, PIPE_M, EVENT_ID2);
  wait_flag(PIPE_FIX, PIPE_M, EVENT_ID3);
  #endif // __DAV_CUBE__

  ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll);
  return;
}

AICORE void vector_kernel(__gm__ float* v1, __gm__ float* v2) {
  unsigned v3 = 0;
  RoundMode v4 = RoundMode::CAST_RINT;
  const int32_t v5 = 96;
  const int32_t v6 = 128;
  const int32_t v7 = 16;
  const int32_t v8 = 32;
  const int32_t v9 = 1;
  const int32_t v10 = 131072;
  const int32_t v11 = 65536;
  const int32_t v12 = 98304;
  const int64_t v13 = 49152;
  const int64_t v14 = 65536;
  const int64_t v15 = 81920;
  const int64_t v16 = 90112;
  const int64_t v17 = 98304;
  const int64_t v18 = 98816;
  const int64_t v19 = 99328;
  const int64_t v20 = 99840;
  const int64_t v21 = 100352;
  const int64_t v22 = 100864;
  const float v23 = 0.0883883461f;
  const float v24 = 1.0f;
  const int64_t v25 = 101376;
  const int32_t v26 = 15;
  const int32_t v27 = 0;
  const int32_t v28 = 32768;
  const int32_t v29 = 262144;
  using T = float;

  #if defined(__DAV_VEC__)
  set_mask_norm();
  set_vector_mask(-1, -1);
  size_t v30 = (size_t) v9;
  int64_t v31 = get_block_num();
  int32_t v32 = (int32_t) ((int64_t) v31);
  int64_t v33 = get_block_idx();
  int32_t v34 = (int32_t) ((int64_t) v33);
  int32_t v35 = v5 / v32;
  int32_t v36 = v5 % v32;
  int32_t v37 = (int32_t) ((uint32_t) v35 + (uint32_t) v9);
  bool v38 = v34 < v36;
  int32_t v39 = v38 ? (int32_t) ((uint32_t) v34 * (uint32_t) v37) : (int32_t) ((uint32_t) ((int32_t) (uint32_t) v36 * (uint32_t) v37) + (uint32_t) ((int32_t) (uint32_t) ((int32_t) (uint32_t) v34 - (uint32_t) v36) * (uint32_t) v35));
  int32_t v40 = (int32_t) ((uint32_t) v34 * (uint32_t) v10);
  __gm__ float* v41 = v1 + v40;
  
  int32_t v42 = (int32_t) ((uint32_t) v40 + (uint32_t) v11);
  __gm__ float* v43 = v1 + v42;
  
  int32_t v44 = (int32_t) ((uint32_t) v40 + (uint32_t) v12);
  __gm__ float* v45 = v1 + v44;
  
  auto v46 = TPipe<0, Direction::DIR_C2V, 32768, 8, 1, false>(v41, v27, v27);
  auto v47 = TPipe<2, Direction::DIR_C2V, 16384, 8, 1, false>(v43, v28, v27);
  auto v48 = TPipe<4, Direction::DIR_V2C, 16384, 8, 1, false>(v45, v27, v29);
  int64_t v49 = get_subblockid();
  Tile<TileType::Vec, float, 16, 256, BLayout::RowMajor, 16, 256, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v50;
  TASSIGN(v50, v13);
  Tile<TileType::Vec, float, 16, 256, BLayout::RowMajor, 16, 256, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v51;
  TASSIGN(v51, v14);
  Tile<TileType::Vec, half, 16, 256, BLayout::RowMajor, 16, 256, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v52;
  TASSIGN(v52, v15);
  Tile<TileType::Vec, float, 16, 128, BLayout::RowMajor, 16, 128, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v53;
  TASSIGN(v53, v16);
  Tile<TileType::Vec, float, 16, 1, BLayout::ColMajor, 16, 1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v54;
  TASSIGN(v54, v17);
  Tile<TileType::Vec, float, 16, 1, BLayout::ColMajor, 16, 1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v55;
  TASSIGN(v55, v18);
  Tile<TileType::Vec, float, 16, 1, BLayout::ColMajor, 16, 1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v56;
  TASSIGN(v56, v19);
  Tile<TileType::Vec, float, 16, 1, BLayout::ColMajor, 16, 1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v57;
  TASSIGN(v57, v20);
  Tile<TileType::Vec, float, 16, 1, BLayout::ColMajor, 16, 1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v58;
  TASSIGN(v58, v21);
  Tile<TileType::Vec, float, 16, 1, BLayout::ColMajor, 16, 1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v59;
  TASSIGN(v59, v22);
  set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0);
  set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0);
  set_flag(PIPE_V, PIPE_MTE2, EVENT_ID7);
  set_flag(PIPE_MTE3, PIPE_V, EVENT_ID5);
  for (size_t v60 = (size_t) v39; v60 < ((size_t) ((int32_t) (uint32_t) v39 + (uint32_t) (v38 ? v37 : v35))); v60 += v30) {
    Tile<TileType::Vec, float, 16, 256, BLayout::RowMajor, 16, 256, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v61;
    TASSIGN(v61, v25);
    wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0);
    TPOP<TPipe<0, Direction::DIR_C2V, 32768, 8, 1, false>, Tile<TileType::Vec, float, 16, 256, BLayout::RowMajor, 16, 256, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v46, v61);
    set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
    wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
    TMULS(v61, v61, v23);
    pipe_barrier(PIPE_V);
    TROWMAX(v55, v61, v50);
    Tile<TileType::Vec, float, 1, 16, BLayout::RowMajor, 1, 16, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v62;
    TRESHAPE(v62, v55);
    Tile<TileType::Vec, float, 1, 16, BLayout::RowMajor, 1, 16, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v63;
    TRESHAPE(v63, v54);
    pipe_barrier(PIPE_V);
    TROWEXPANDSUB(v51, v61, v55);
    set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1);
    TMULS(v63, v62, v24);
    pipe_barrier(PIPE_V);
    TEXP(v51, v51);
    pipe_barrier(PIPE_V);
    TROWSUM(v56, v51, v50);
    wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0);
    TCVT(v52, v51, v4);
    set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
    wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0);
    TPUSH<TPipe<4, Direction::DIR_V2C, 16384, 8, 1, false>, Tile<TileType::Vec, half, 16, 256, BLayout::RowMajor, 16, 256, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v48, v52);
    set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1);
    TFREE<TPipe<0, Direction::DIR_C2V, 32768, 8, 1, false>, TileSplitAxis::TILE_UP_DOWN>(v46);
    Tile<TileType::Vec, float, 16, 256, BLayout::RowMajor, 16, 256, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v64;
    TASSIGN(v64, v25);
    wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1);
    TPOP<TPipe<0, Direction::DIR_C2V, 32768, 8, 1, false>, Tile<TileType::Vec, float, 16, 256, BLayout::RowMajor, 16, 256, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v46, v64);
    set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1);
    wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1);
    TMULS(v64, v64, v23);
    pipe_barrier(PIPE_V);
    TROWMAX(v55, v64, v50);
    Tile<TileType::Vec, float, 1, 16, BLayout::RowMajor, 1, 16, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v65;
    TRESHAPE(v65, v59);
    Tile<TileType::Vec, float, 1, 16, BLayout::RowMajor, 1, 16, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v66;
    TRESHAPE(v66, v56);
    Tile<TileType::Vec, float, 1, 16, BLayout::RowMajor, 1, 16, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v67;
    TRESHAPE(v67, v57);
    pipe_barrier(PIPE_V);
    TMAX(v62, v62, v63);
    pipe_barrier(PIPE_V);
    TSUB(v65, v63, v62);
    pipe_barrier(PIPE_V);
    TMULS(v63, v62, v24);
    TROWEXPANDSUB(v51, v64, v55);
    set_flag(PIPE_V, PIPE_MTE2, EVENT_ID2);
    TEXP(v65, v65);
    pipe_barrier(PIPE_V);
    TEXP(v51, v51);
    TMUL(v66, v66, v65);
    pipe_barrier(PIPE_V);
    TROWSUM(v57, v51, v50);
    pipe_barrier(PIPE_V);
    TADD(v66, v66, v67);
    wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1);
    TCVT(v52, v51, v4);
    set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1);
    wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1);
    TPUSH<TPipe<4, Direction::DIR_V2C, 16384, 8, 1, false>, Tile<TileType::Vec, half, 16, 256, BLayout::RowMajor, 16, 256, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v48, v52);
    set_flag(PIPE_MTE3, PIPE_V, EVENT_ID2);
    TFREE<TPipe<0, Direction::DIR_C2V, 32768, 8, 1, false>, TileSplitAxis::TILE_UP_DOWN>(v46);
    Tile<TileType::Vec, float, 16, 128, BLayout::RowMajor, 16, 128, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v68;
    TASSIGN(v68, v25);
    wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID2);
    TPOP<TPipe<2, Direction::DIR_C2V, 16384, 8, 1, false>, Tile<TileType::Vec, float, 16, 128, BLayout::RowMajor, 16, 128, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v47, v68);
    set_flag(PIPE_MTE2, PIPE_V, EVENT_ID2);
    wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID2);
    TMOV(v53, v68);
    set_flag(PIPE_V, PIPE_MTE2, EVENT_ID3);
    TFREE<TPipe<2, Direction::DIR_C2V, 16384, 8, 1, false>, TileSplitAxis::TILE_UP_DOWN>(v47);
    Tile<TileType::Vec, float, 16, 256, BLayout::RowMajor, 16, 256, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v69;
    TASSIGN(v69, v25);
    wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID3);
    TPOP<TPipe<0, Direction::DIR_C2V, 32768, 8, 1, false>, Tile<TileType::Vec, float, 16, 256, BLayout::RowMajor, 16, 256, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v46, v69);
    set_flag(PIPE_MTE2, PIPE_V, EVENT_ID3);
    wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID3);
    TMULS(v69, v69, v23);
    pipe_barrier(PIPE_V);
    TROWMAX(v55, v69, v50);
    Tile<TileType::Vec, float, 1, 16, BLayout::RowMajor, 1, 16, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v70;
    TRESHAPE(v70, v58);
    pipe_barrier(PIPE_V);
    TMAX(v62, v62, v63);
    pipe_barrier(PIPE_V);
    TSUB(v70, v63, v62);
    pipe_barrier(PIPE_V);
    TMULS(v63, v62, v24);
    TROWEXPANDSUB(v51, v69, v55);
    set_flag(PIPE_V, PIPE_MTE2, EVENT_ID4);
    TEXP(v70, v70);
    pipe_barrier(PIPE_V);
    TEXP(v51, v51);
    TMUL(v66, v66, v70);
    pipe_barrier(PIPE_V);
    TROWSUM(v57, v51, v50);
    pipe_barrier(PIPE_V);
    TADD(v66, v66, v67);
    wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID2);
    TCVT(v52, v51, v4);
    set_flag(PIPE_V, PIPE_MTE3, EVENT_ID2);
    wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID2);
    TPUSH<TPipe<4, Direction::DIR_V2C, 16384, 8, 1, false>, Tile<TileType::Vec, half, 16, 256, BLayout::RowMajor, 16, 256, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v48, v52);
    set_flag(PIPE_MTE3, PIPE_V, EVENT_ID3);
    TFREE<TPipe<0, Direction::DIR_C2V, 32768, 8, 1, false>, TileSplitAxis::TILE_UP_DOWN>(v46);
    Tile<TileType::Vec, float, 16, 128, BLayout::RowMajor, 16, 128, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v71;
    TASSIGN(v71, v25);
    wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID4);
    TPOP<TPipe<2, Direction::DIR_C2V, 16384, 8, 1, false>, Tile<TileType::Vec, float, 16, 128, BLayout::RowMajor, 16, 128, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v47, v71);
    set_flag(PIPE_MTE2, PIPE_V, EVENT_ID4);
    TROWEXPANDMUL(v53, v53, v59);
    pipe_barrier(PIPE_V);
    wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID4);
    TADD(v53, v53, v71);
    set_flag(PIPE_V, PIPE_MTE2, EVENT_ID5);
    TFREE<TPipe<2, Direction::DIR_C2V, 16384, 8, 1, false>, TileSplitAxis::TILE_UP_DOWN>(v47);
    Tile<TileType::Vec, float, 16, 256, BLayout::RowMajor, 16, 256, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v72;
    TASSIGN(v72, v25);
    wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID5);
    TPOP<TPipe<0, Direction::DIR_C2V, 32768, 8, 1, false>, Tile<TileType::Vec, float, 16, 256, BLayout::RowMajor, 16, 256, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v46, v72);
    set_flag(PIPE_MTE2, PIPE_V, EVENT_ID5);
    wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID5);
    TMULS(v72, v72, v23);
    pipe_barrier(PIPE_V);
    TROWMAX(v55, v72, v50);
    pipe_barrier(PIPE_V);
    TMAX(v62, v62, v63);
    pipe_barrier(PIPE_V);
    TSUB(v65, v63, v62);
    pipe_barrier(PIPE_V);
    TMULS(v63, v62, v24);
    TROWEXPANDSUB(v51, v72, v55);
    set_flag(PIPE_V, PIPE_MTE2, EVENT_ID6);
    TEXP(v65, v65);
    pipe_barrier(PIPE_V);
    TEXP(v51, v51);
    TMUL(v66, v66, v65);
    pipe_barrier(PIPE_V);
    TROWSUM(v57, v51, v50);
    pipe_barrier(PIPE_V);
    TADD(v66, v66, v67);
    wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID3);
    TCVT(v52, v51, v4);
    set_flag(PIPE_V, PIPE_MTE3, EVENT_ID3);
    wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID3);
    TPUSH<TPipe<4, Direction::DIR_V2C, 16384, 8, 1, false>, Tile<TileType::Vec, half, 16, 256, BLayout::RowMajor, 16, 256, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v48, v52);
    set_flag(PIPE_MTE3, PIPE_V, EVENT_ID4);
    TFREE<TPipe<0, Direction::DIR_C2V, 32768, 8, 1, false>, TileSplitAxis::TILE_UP_DOWN>(v46);
    wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID6);
    wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID4);
    for (size_t v73 = v30; v73 < ((size_t) v26); v73 += v30) {
      Tile<TileType::Vec, float, 16, 128, BLayout::RowMajor, 16, 128, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v74;
      TASSIGN(v74, v25);
      wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID7);
      TPOP<TPipe<2, Direction::DIR_C2V, 16384, 8, 1, false>, Tile<TileType::Vec, float, 16, 128, BLayout::RowMajor, 16, 128, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v47, v74);
      set_flag(PIPE_MTE2, PIPE_V, EVENT_ID6);
      TROWEXPANDMUL(v53, v53, v58);
      pipe_barrier(PIPE_V);
      wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID6);
      TADD(v53, v53, v74);
      set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1);
      TFREE<TPipe<2, Direction::DIR_C2V, 16384, 8, 1, false>, TileSplitAxis::TILE_UP_DOWN>(v47);
      Tile<TileType::Vec, float, 16, 256, BLayout::RowMajor, 16, 256, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v75;
      TASSIGN(v75, v25);
      wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1);
      TPOP<TPipe<0, Direction::DIR_C2V, 32768, 8, 1, false>, Tile<TileType::Vec, float, 16, 256, BLayout::RowMajor, 16, 256, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v46, v75);
      set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
      wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
      TMULS(v75, v75, v23);
      pipe_barrier(PIPE_V);
      TROWMAX(v55, v75, v50);
      pipe_barrier(PIPE_V);
      TMAX(v62, v62, v63);
      pipe_barrier(PIPE_V);
      TSUB(v70, v63, v62);
      pipe_barrier(PIPE_V);
      TMULS(v63, v62, v24);
      TROWEXPANDSUB(v51, v75, v55);
      set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1);
      TEXP(v70, v70);
      pipe_barrier(PIPE_V);
      TEXP(v51, v51);
      TMUL(v66, v66, v70);
      pipe_barrier(PIPE_V);
      TROWSUM(v57, v51, v50);
      pipe_barrier(PIPE_V);
      TADD(v66, v66, v67);
      wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID5);
      TCVT(v52, v51, v4);
      set_flag(PIPE_V, PIPE_MTE3, EVENT_ID4);
      wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID4);
      TPUSH<TPipe<4, Direction::DIR_V2C, 16384, 8, 1, false>, Tile<TileType::Vec, half, 16, 256, BLayout::RowMajor, 16, 256, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v48, v52);
      set_flag(PIPE_MTE3, PIPE_V, EVENT_ID6);
      TFREE<TPipe<0, Direction::DIR_C2V, 32768, 8, 1, false>, TileSplitAxis::TILE_UP_DOWN>(v46);
      Tile<TileType::Vec, float, 16, 128, BLayout::RowMajor, 16, 128, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v76;
      TASSIGN(v76, v25);
      wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1);
      TPOP<TPipe<2, Direction::DIR_C2V, 16384, 8, 1, false>, Tile<TileType::Vec, float, 16, 128, BLayout::RowMajor, 16, 128, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v47, v76);
      set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
      TROWEXPANDMUL(v53, v53, v59);
      pipe_barrier(PIPE_V);
      wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
      TADD(v53, v53, v76);
      set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1);
      TFREE<TPipe<2, Direction::DIR_C2V, 16384, 8, 1, false>, TileSplitAxis::TILE_UP_DOWN>(v47);
      Tile<TileType::Vec, float, 16, 256, BLayout::RowMajor, 16, 256, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v77;
      TASSIGN(v77, v25);
      wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1);
      TPOP<TPipe<0, Direction::DIR_C2V, 32768, 8, 1, false>, Tile<TileType::Vec, float, 16, 256, BLayout::RowMajor, 16, 256, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v46, v77);
      set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
      wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
      TMULS(v77, v77, v23);
      pipe_barrier(PIPE_V);
      TROWMAX(v55, v77, v50);
      pipe_barrier(PIPE_V);
      TMAX(v62, v62, v63);
      pipe_barrier(PIPE_V);
      TSUB(v65, v63, v62);
      pipe_barrier(PIPE_V);
      TMULS(v63, v62, v24);
      TROWEXPANDSUB(v51, v77, v55);
      set_flag(PIPE_V, PIPE_MTE2, EVENT_ID7);
      TEXP(v65, v65);
      pipe_barrier(PIPE_V);
      TEXP(v51, v51);
      TMUL(v66, v66, v65);
      pipe_barrier(PIPE_V);
      TROWSUM(v57, v51, v50);
      pipe_barrier(PIPE_V);
      TADD(v66, v66, v67);
      wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID6);
      TCVT(v52, v51, v4);
      set_flag(PIPE_V, PIPE_MTE3, EVENT_ID5);
      wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID5);
      TPUSH<TPipe<4, Direction::DIR_V2C, 16384, 8, 1, false>, Tile<TileType::Vec, half, 16, 256, BLayout::RowMajor, 16, 256, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v48, v52);
      set_flag(PIPE_MTE3, PIPE_V, EVENT_ID5);
      TFREE<TPipe<0, Direction::DIR_C2V, 32768, 8, 1, false>, TileSplitAxis::TILE_UP_DOWN>(v46);
    };
    set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1);
    set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0);
    Tile<TileType::Vec, float, 16, 128, BLayout::RowMajor, 16, 128, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v78;
    TASSIGN(v78, v25);
    wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1);
    TPOP<TPipe<2, Direction::DIR_C2V, 16384, 8, 1, false>, Tile<TileType::Vec, float, 16, 128, BLayout::RowMajor, 16, 128, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v47, v78);
    set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
    TROWEXPANDMUL(v53, v53, v58);
    pipe_barrier(PIPE_V);
    wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
    TADD(v53, v53, v78);
    set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1);
    TFREE<TPipe<2, Direction::DIR_C2V, 16384, 8, 1, false>, TileSplitAxis::TILE_UP_DOWN>(v47);
    Tile<TileType::Vec, float, 16, 128, BLayout::RowMajor, 16, 128, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v79;
    TASSIGN(v79, v25);
    wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1);
    TPOP<TPipe<2, Direction::DIR_C2V, 16384, 8, 1, false>, Tile<TileType::Vec, float, 16, 128, BLayout::RowMajor, 16, 128, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>, TileSplitAxis::TILE_UP_DOWN>(v47, v79);
    set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
    pipe_barrier(PIPE_V);
    TROWEXPANDMUL(v53, v53, v59);
    pipe_barrier(PIPE_V);
    wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0);
    TADD(v53, v53, v79);
    set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0);
    TFREE<TPipe<2, Direction::DIR_C2V, 16384, 8, 1, false>, TileSplitAxis::TILE_UP_DOWN>(v47);
    pipe_barrier(PIPE_V);
    TROWEXPANDDIV(v53, v53, v56);
    set_flag(PIPE_V, PIPE_MTE3, EVENT_ID6);
    pto::Shape<1, 1, 1, 16, 128> v80 = pto::Shape<1, 1, 1, 16, 128>();
    pto::Stride<2048, 2048, 2048, 128, 1> v81 = pto::Stride<2048, 2048, 2048, 128, 1>();
    GlobalTensor<float, pto::Shape<1, 1, 1, 16, 128>, pto::Stride<2048, 2048, 2048, 128, 1>, pto::Layout::ND> v82 = GlobalTensor<float, pto::Shape<1, 1, 1, 16, 128>, pto::Stride<2048, 2048, 2048, 128, 1>, pto::Layout::ND>(v2 + (v3 + (unsigned) ((int32_t) (uint32_t) ((int32_t) (uint32_t) ((int32_t) v60) * (uint32_t) v8) + (uint32_t) ((int32_t) (uint32_t) ((int32_t) (int64_t) v49) * (uint32_t) v7)) * (unsigned) v6 + v3 * (unsigned) v9), v80, v81);
    wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID6);
    TSTORE(v82, v53);
  }
  wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0);
  wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0);
  wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID7);
  wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID5);
  #endif // __DAV_VEC__

  ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll);
  return;
}

__global__ AICORE void call_both(__gm__ int64_t* v1, __gm__ float* v2, __gm__ half* v3, __gm__ half* v4, __gm__ half* v5, __gm__ float* v6) {
  using T = float;
  uint64_t v7 = (uint64_t) v1;
  set_ffts_base_addr(v7);
  cube_kernel(v2, v3, v4, v5);
  vector_kernel(v2, v6);
  return;
}

[MirkoDeVita98]

PTO_LIB_PATH=/sources/pto-isa \
ptoas --pto-arch=a3 --pto-level=level3 --enable-insert-sync \
  compiler_team_fa.pto \
  > /tmp/compiler_team_fa.cpp

bisheng \
  -I/sources/pto-isa/include \
  -fPIC -shared -D_FORTIFY_SOURCE=2 -O2 -std=c++17 \
  -Wno-macro-redefined -Wno-ignored-attributes -fstack-protector-strong \
  -xcce -Xhost-start -Xhost-end \
  -mllvm -cce-aicore-stack-size=0x8000 \
  -mllvm -cce-aicore-function-stack-size=0x8000 \
  -mllvm -cce-aicore-record-overflow=true \
  -mllvm -cce-aicore-addr-transform \
  -mllvm -cce-aicore-dcci-insert-for-scalar=false \
  -cce-enable-mix \
  --npu-arch=dav-2201 -DMEMORY_BASE \
  -std=gnu++17 \
  -DKERNEL_CPP="\"/tmp/compiler_team_fa.cpp\"" \
  examples/aot/flash_attention/experimental/caller.cpp \
  -o /tmp/compiler_team_fa.so

Run with:

import ctypes
import torch
import torch_npu  # noqa
from ptodsl.utils.npu_info import get_num_cube_cores, get_test_device

Q_ROWS = 3072
HEAD = 128
S1_TOTAL = 8192
NUM_Q_BLOCKS = Q_ROWS // 32
GM_ELEMS_PER_BLOCK = 229376

lib = ctypes.CDLL("/tmp/compiler_team_fa.so")
lib.call_kernel.argtypes = [
    ctypes.c_uint32,
    ctypes.c_void_p,
    ctypes.c_void_p,
    ctypes.c_void_p,
    ctypes.c_void_p,
    ctypes.c_void_p,
    ctypes.c_void_p,
]
lib.call_kernel.restype = None

def ptr(t):
    return ctypes.c_void_p(t.data_ptr())

device = get_test_device()
block_dim = min(NUM_Q_BLOCKS, get_num_cube_cores())

torch.manual_seed(0)
q = torch.randn((Q_ROWS, HEAD), dtype=torch.float16, device=device)
k = torch.randn((S1_TOTAL, HEAD), dtype=torch.float16, device=device)
v = torch.randn((S1_TOTAL, HEAD), dtype=torch.float16, device=device)
gm = torch.zeros((GM_ELEMS_PER_BLOCK * block_dim,), dtype=torch.float32, device=device)
o = torch.zeros((Q_ROWS, HEAD), dtype=torch.float32, device=device)

lib.call_kernel(
    block_dim,
    torch.npu.current_stream()._as_parameter_,
    ptr(gm),
    ptr(q),
    ptr(k),
    ptr(v),
    ptr(o),
)
torch.npu.synchronize()