gemv: coalesce batched DMA into a single iterated BD per column#127
gemv: coalesce batched DMA into a single iterated BD per column#127atassis wants to merge 4 commits into
Conversation
andrej
left a comment
There was a problem hiding this comment.
Thanks, looking good! Reducing the number of runtime sequence instructions seems useful and the performance gain seems significant.
I only have some questions about the tensor access patterns. I have a feeling / hope that this could be radically simplified if indeed they are just linear transfers.
| # Backpressure replaces the per-batch drain wait, so the A/C ObjectFifos must | ||
| # be deep enough (>=2) for the producer not to overrun the consumer. |
There was a problem hiding this comment.
What do you mean by the producer overrunning the consumer? Proper lock synchronization and backpressure should i.m.o. lead to at most a deadlock, and "overrunning" (buffers?) sounds concerning.
There was a problem hiding this comment.
Yeah, "overrun" is a bad word for it, sorry. The ObjectFifo lock backpressures, so a fast producer just blocks (worst case it stalls), it can't corrupt anything. depth>=2 is only there to overlap the fill with compute, so it's a perf thing, not correctness. Reworded the comment.
| # dims have a size cap (1023) while one dim is size-uncapped; TAP sizes are | ||
| # outermost-first and the verifier reverses them, so [1, num_batches, run_hi, | ||
| # run_lo] places num_batches in the uncapped dim and the contiguous run in the two | ||
| # wrap dims. Access-equivalent to the unroll (covered by test_gemv_batched). |
There was a problem hiding this comment.
| # wrap dims. Access-equivalent to the unroll (covered by test_gemv_batched). | |
| # wrap dims. |
There was a problem hiding this comment.
Done. Also dropped the duplicated sentence.
| def split_run(run, lim=MAX_WRAP, gran=GRAN_ELEMS): | ||
| """Factor a contiguous run into (hi, lo), both <= lim and lo a multiple of gran | ||
| (the address-granularity-aligned inner size), lo maximal. None if no such | ||
| split exists (caller then falls back to the per-batch path).""" | ||
| lo_start = (lim // gran) * gran | ||
| for lo in range(lo_start, 0, -gran): | ||
| if run % lo == 0 and (run // lo) <= lim: | ||
| return (run // lo, lo) | ||
| return None |
There was a problem hiding this comment.
@atassis Does this split up a purely linear, contiguous transfer? There is a special case to not use data layout transformation dimensions at all. If left off, you should be able to specify only transfer lengths, and no dimensions.
@hunhoffe My understanding is this splits up a contiguous transfer into multiple dimensions. Is this logic still required after MLIR-AIE PRs #2924 and #3036?
There was a problem hiding this comment.
So within a batch the run is contiguous, but the batch stride is the whole matrix (M*K) and each column only takes its (M//cols)*K slice, so for cols>1 the batches have gaps between them - it's linear only when cols==1. On #2924/#3036: yeah those would let me drop the inner split, but we pin mlir_aie e4f35d6 (Mar 31) and both landed after that (Apr 7 / May 11), so on this pin the split is still needed. I left a note that it can become [num_batches, A_run] once the pin moves past #3036.
There was a problem hiding this comment.
@atassis Just updated the pinned MLIR-AIE. Can that simplify this?
| MAX_WRAP = 1023 | ||
| MAX_STRIDE = (1 << 20) - 1 # conservative element-stride bound for the wrap dims | ||
| GRAN_ELEMS = 2 # 4-byte shim granularity / 2-byte bf16 element |
There was a problem hiding this comment.
| MAX_WRAP = 1023 | |
| MAX_STRIDE = (1 << 20) - 1 # conservative element-stride bound for the wrap dims | |
| GRAN_ELEMS = 2 # 4-byte shim granularity / 2-byte bf16 element | |
| MAX_WRAP = 1023 | |
| MAX_STRIDE = (1 << 20) - 1 # conservative element-stride bound for the wrap dims | |
| GRAN_ELEMS = 2 # 4-byte shim granularity / 2-byte bf16 element | |
| # FIXME: Pull this information in to the target model in MLIR-AIE; also used in https://github.com/Xilinx/mlir-aie/blob/58ad59e8b3a601f540a653d1222d1dfc6303d3b0/lib/Dialect/AIEX/IR/AIEXDialect.cpp#L191-L199 |
(No need to do in this PR but let's leave the comment as a future to do)
There was a problem hiding this comment.
Added the FIXME pointing at verifyStridesWraps in AIEXDialect.cpp. Kept the constants, still needed on this pin.
| return TensorAccessPattern( | ||
| tensor_dims=L3_ty.__args__[0], | ||
| offset=col_off, | ||
| sizes=[1, num_batches, run_hi, run_lo], | ||
| strides=[0, bstride, run_lo, 1], | ||
| ) |
There was a problem hiding this comment.
I didn't spend too much time thinking through this but at a glance this looks like it might just be a linear transfer (if bstride==run_hi*run_lo). If that's the case, a lot of this code can be deleted and you should be able to just not supply any TAP to the rt.fill at all, which will equal a linear transfer.
If the transfer is more complicated, please explain in a comment! Thanks
There was a problem hiding this comment.
Yeah I checked that - bstride is M*K and the run is (M//cols)*K, so they're equal only when cols==1. For cols>1 it's a real strided gather, that's why the batch dim has to stay. The run_hi/run_lo split is just to get under the 1023 wrap cap (more in the other thread). So I explained it in a comment instead of deleting.
| # Fallback (also the num_batches==1 path): stock per-batch unroll. | ||
| for batch in range(num_batches): | ||
| tg_ac = rt.task_group() | ||
| for col in range(cols): | ||
| rt.fill( | ||
| A_L3L1_fifos[col].prod(), | ||
| A, | ||
| A_taps[col][batch], | ||
| task_group=tg_ac, | ||
| ) | ||
| for col in range(cols): | ||
| rt.drain( | ||
| C_L1L3_fifos[col].cons(), | ||
| C, | ||
| C_taps[col][batch], | ||
| task_group=tg_ac, | ||
| wait=True, | ||
| ) | ||
| rt.finish_task_group(tg_ac) |
There was a problem hiding this comment.
Since this is essentially identical to the if-branch above except for number of wait/finish_task_group calls, could we reduce code duplication by instead looping over a new variable num_waits, which in this else-case would be num_waits==num_batches and in the if-case would be num_waits==1, and then also swap out the TAPs as appropriate.
There was a problem hiding this comment.
Done - one loop over num_waits now (1 for coalesced, num_batches otherwise), and swap the tap.
There was a problem hiding this comment.
Awesome! thanks, I prefer that
CI Test Results3ab7ee2 (2026_06_29_17_34_01) IRON - CI SummaryExamplesiron/applications/llama_3.2_1b
Smalliron/operators/axpy
iron/operators/dequant
iron/operators/elementwise_add
iron/operators/elementwise_mul
iron/operators/gelu
iron/operators/gemm
iron/operators/gemv
iron/operators/layer_norm
iron/operators/mem_copy
iron/operators/mha
iron/operators/relu
iron/operators/rms_norm
iron/operators/rope
iron/operators/sigmoid
iron/operators/silu
iron/operators/softmax
iron/operators/swiglu_decode
iron/operators/swiglu_prefill
iron/operators/tanh
iron/operators/transpose
Krackan - SmallIRONTested on iron/operators/axpy
iron/operators/dequant
iron/operators/elementwise_add
iron/operators/elementwise_mul
iron/operators/gelu
iron/operators/gemm
iron/operators/gemv
iron/operators/layer_norm
iron/operators/mem_copy
iron/operators/mha
iron/operators/relu
iron/operators/rms_norm
iron/operators/rope
iron/operators/sigmoid
iron/operators/silu
iron/operators/softmax
iron/operators/swiglu_decode
iron/operators/swiglu_prefill
iron/operators/tanh
iron/operators/transpose
Trends: IRON Trendsiron/operators/axpytest_axpy[input_length_2048-num_aie_columns_1-tile_size_2048-scalar_factor_3.0]
test_axpy[input_length_2048-num_aie_columns_2-tile_size_1024-scalar_factor_3.0]
test_axpy[input_length_2048-num_aie_columns_4-tile_size_512-scalar_factor_3.0]
test_axpy[input_length_2048-num_aie_columns_8-tile_size_256-scalar_factor_3.0]
iron/operators/dequanttest_dequant[input_length_2048-num_aie_columns_1-num_channels_1-tile_size_2048-group_size_32]
test_dequant[input_length_2048-num_aie_columns_1-num_channels_2-tile_size_1024-group_size_32]
test_dequant[input_length_2048-num_aie_columns_2-num_channels_1-tile_size_1024-group_size_32]
test_dequant[input_length_2048-num_aie_columns_2-num_channels_2-tile_size_512-group_size_32]
test_dequant[input_length_2048-num_aie_columns_4-num_channels_1-tile_size_512-group_size_32]
test_dequant[input_length_2048-num_aie_columns_4-num_channels_2-tile_size_256-group_size_32]
test_dequant[input_length_2048-num_aie_columns_8-num_channels_1-tile_size_256-group_size_32]
test_dequant[input_length_2048-num_aie_columns_8-num_channels_2-tile_size_128-group_size_32]
iron/operators/elementwise_addtest_elementwise_add[input_length_2048-num_aie_columns_1-tile_size_2048]
test_elementwise_add[input_length_2048-num_aie_columns_2-tile_size_1024]
test_elementwise_add[input_length_2048-num_aie_columns_4-tile_size_512]
test_elementwise_add[input_length_2048-num_aie_columns_8-tile_size_256]
iron/operators/elementwise_multest_elementwise_mul[input_length_2048-num_aie_columns_1-tile_size_2048]
test_elementwise_mul[input_length_2048-num_aie_columns_2-tile_size_1024]
test_elementwise_mul[input_length_2048-num_aie_columns_4-tile_size_512]
test_elementwise_mul[input_length_2048-num_aie_columns_8-tile_size_256]
iron/operators/gelutest_gelu[input_length_2048-num_aie_columns_1-num_channels_1-tile_size_2048]
test_gelu[input_length_2048-num_aie_columns_1-num_channels_2-tile_size_1024]
test_gelu[input_length_2048-num_aie_columns_2-num_channels_1-tile_size_1024]
test_gelu[input_length_2048-num_aie_columns_2-num_channels_2-tile_size_512]
test_gelu[input_length_2048-num_aie_columns_4-num_channels_1-tile_size_512]
test_gelu[input_length_2048-num_aie_columns_4-num_channels_2-tile_size_256]
test_gelu[input_length_2048-num_aie_columns_8-num_channels_1-tile_size_256]
test_gelu[input_length_2048-num_aie_columns_8-num_channels_2-tile_size_128]
iron/operators/gemmtest_gemm[M_1792-K_896-N_1152-num_aie_columns_8-b_col_maj_False-c_col_maj_True-m_64-k_32-n_48-trace_size_0-partition_N_1]
test_gemm[M_192-K_384-N_64-num_aie_columns_4-b_col_maj_False-c_col_maj_False-m_48-k_96-n_16-trace_size_0-partition_N_1]
test_gemm[M_192-K_384-N_64-num_aie_columns_4-b_col_maj_True-c_col_maj_True-m_48-k_96-n_16-trace_size_0-partition_N_1]
test_gemm[M_2048-K_2048-N_2048-num_aie_columns_1-b_col_maj_False-c_col_maj_False-m_64-k_64-n_64-trace_size_0-partition_N_1]
test_gemm[M_2048-K_2048-N_2048-num_aie_columns_2-b_col_maj_True-c_col_maj_False-m_64-k_64-n_64-trace_size_0-partition_N_1]
test_gemm[M_2048-K_2048-N_2048-num_aie_columns_8-b_col_maj_True-c_col_maj_True-m_64-k_64-n_64-trace_size_0-partition_N_1]
test_gemm[M_384-K_1536-N_1792-num_aie_columns_4-b_col_maj_True-c_col_maj_False-m_32-k_48-n_64-trace_size_0-partition_N_1]
test_gemm[M_64-K_512-N_256-num_aie_columns_4-b_col_maj_True-c_col_maj_False-m_16-k_64-n_64-trace_size_0-partition_N_4]
test_gemm[M_896-K_1792-N_640-num_aie_columns_8-b_col_maj_False-c_col_maj_True-m_32-k_64-n_80-trace_size_0-partition_N_1]
iron/operators/gemvtest_gemv[M_128-K_128-num_aie_columns_1-tile_size_input_32-tile_size_output_128]
test_gemv[M_2048-K_8192-num_aie_columns_1-tile_size_input_1-tile_size_output_2048]
test_gemv[M_2048-K_8192-num_aie_columns_2-tile_size_input_1-tile_size_output_1024]
test_gemv[M_2048-K_8192-num_aie_columns_4-tile_size_input_1-tile_size_output_512]
test_gemv[M_2048-K_8192-num_aie_columns_8-tile_size_input_1-tile_size_output_256]
test_gemv[M_8192-K_2048-num_aie_columns_1-tile_size_input_4-tile_size_output_1024]
test_gemv[M_8192-K_2048-num_aie_columns_2-tile_size_input_4-tile_size_output_1024]
test_gemv[M_8192-K_2048-num_aie_columns_4-tile_size_input_4-tile_size_output_1024]
test_gemv[M_8192-K_2048-num_aie_columns_8-tile_size_input_4-tile_size_output_1024]
test_gemv_batched[M_1024-K_1024-num_aie_columns_1-tile_size_input_1-tile_size_output_64-num_batches_2]No metrics available. test_gemv_batched[M_1026-K_64-num_aie_columns_1-tile_size_input_1-tile_size_output_2-num_batches_2]No metrics available. test_gemv_batched[M_256-K_128-num_aie_columns_1-tile_size_input_1-tile_size_output_256-num_batches_4]No metrics available. test_gemv_batched[M_256-K_128-num_aie_columns_8-tile_size_input_1-tile_size_output_32-num_batches_100]No metrics available. test_gemv_batched[M_448-K_64-num_aie_columns_8-tile_size_input_1-tile_size_output_56-num_batches_192]No metrics available. test_gemv_batched[M_512-K_64-num_aie_columns_8-tile_size_input_4-tile_size_output_64-num_batches_32]No metrics available. test_gemv_batched[M_64-K_1536-num_aie_columns_1-tile_size_input_1-tile_size_output_64-num_batches_8]No metrics available. iron/operators/layer_normtest_layer_norm[input_length_2048-num_aie_columns_1-num_channels_1-tile_size_2048]
test_layer_norm[input_length_2048-num_aie_columns_1-num_channels_2-tile_size_1024]
test_layer_norm[input_length_2048-num_aie_columns_2-num_channels_1-tile_size_1024]
test_layer_norm[input_length_2048-num_aie_columns_2-num_channels_2-tile_size_512]
test_layer_norm[input_length_2048-num_aie_columns_4-num_channels_1-tile_size_512]
test_layer_norm[input_length_2048-num_aie_columns_4-num_channels_2-tile_size_256]
test_layer_norm[input_length_2048-num_aie_columns_8-num_channels_1-tile_size_256]
test_layer_norm[input_length_2048-num_aie_columns_8-num_channels_2-tile_size_128]
iron/operators/mem_copytest_mem_copy[input_length_2048-num_cores_1-num_channels_1-bypass_False-tile_size_2048]
test_mem_copy[input_length_2048-num_cores_16-num_channels_2-bypass_False-tile_size_128]
test_mem_copy[input_length_2048-num_cores_2-num_channels_1-bypass_False-tile_size_1024]
test_mem_copy[input_length_2048-num_cores_2-num_channels_2-bypass_False-tile_size_1024]
test_mem_copy[input_length_2048-num_cores_4-num_channels_1-bypass_False-tile_size_512]
test_mem_copy[input_length_2048-num_cores_4-num_channels_2-bypass_False-tile_size_512]
test_mem_copy[input_length_2048-num_cores_8-num_channels_1-bypass_False-tile_size_256]
test_mem_copy[input_length_2048-num_cores_8-num_channels_2-bypass_False-tile_size_256]
iron/operators/mhatest_mha[seq_len_16384-dim_64-num_heads_1-num_pipelines_8-num_kv_heads_0]
iron/operators/rms_normtest_rms_norm[input_length_2048-num_aie_columns_1-num_channels_1-tile_size_2048-weighted_False]
test_rms_norm[input_length_2048-num_aie_columns_1-num_channels_1-tile_size_2048-weighted_True]
test_rms_norm[input_length_2048-num_aie_columns_1-num_channels_2-tile_size_1024-weighted_False]
test_rms_norm[input_length_2048-num_aie_columns_1-num_channels_2-tile_size_1024-weighted_True]
test_rms_norm[input_length_2048-num_aie_columns_2-num_channels_1-tile_size_1024-weighted_False]
test_rms_norm[input_length_2048-num_aie_columns_2-num_channels_1-tile_size_1024-weighted_True]
test_rms_norm[input_length_2048-num_aie_columns_2-num_channels_2-tile_size_512-weighted_False]
test_rms_norm[input_length_2048-num_aie_columns_2-num_channels_2-tile_size_512-weighted_True]
test_rms_norm[input_length_2048-num_aie_columns_4-num_channels_1-tile_size_512-weighted_False]
test_rms_norm[input_length_2048-num_aie_columns_4-num_channels_1-tile_size_512-weighted_True]
test_rms_norm[input_length_2048-num_aie_columns_4-num_channels_2-tile_size_256-weighted_False]
test_rms_norm[input_length_2048-num_aie_columns_4-num_channels_2-tile_size_256-weighted_True]
test_rms_norm[input_length_2048-num_aie_columns_8-num_channels_1-tile_size_256-weighted_False]
test_rms_norm[input_length_2048-num_aie_columns_8-num_channels_1-tile_size_256-weighted_True]
test_rms_norm[input_length_2048-num_aie_columns_8-num_channels_2-tile_size_128-weighted_False]
iron/operators/ropetest_rope[rows_32-cols_512-angle_rows_32-aie_columns_1-method_type_0]
test_rope[rows_32-cols_512-angle_rows_32-aie_columns_2-method_type_0]
test_rope[rows_32-cols_512-angle_rows_32-aie_columns_4-method_type_0]
test_rope[rows_32-cols_512-angle_rows_32-aie_columns_8-method_type_0]
test_rope[rows_32-cols_512-angle_rows_8-aie_columns_1-method_type_0]
test_rope[rows_32-cols_512-angle_rows_8-aie_columns_2-method_type_0]
test_rope[rows_32-cols_512-angle_rows_8-aie_columns_4-method_type_0]
test_rope[rows_32-cols_512-angle_rows_8-aie_columns_8-method_type_0]
iron/operators/softmaxtest_softmax[input_length_32768-num_aie_columns_2-num_channels_2-tile_size_1024]
test_softmax[input_length_32768-num_aie_columns_2-num_channels_2-tile_size_2048]
test_softmax[input_length_32768-num_aie_columns_2-num_channels_2-tile_size_512]
iron/operators/swiglu_decodetest_swiglu_decode[embedding_dim_1024-hidden_dim_3584]
test_swiglu_decode[embedding_dim_2048-hidden_dim_2048]
iron/operators/swiglu_prefilltest_swiglu_prefill[seq_len_256-embedding_dim_2048-hidden_dim_2048-prio_accuracy_False]
iron/operators/transposetest_transpose[M_2048-N_64-aie_columns_1-channels_1-m_64-n_64-s_8-num_batches_1]
test_transpose[M_2048-N_64-aie_columns_1-channels_1-m_64-n_64-s_8-num_batches_2]
test_transpose[M_2048-N_64-aie_columns_1-channels_1-m_64-n_64-s_8]
test_transpose[M_2048-N_64-aie_columns_1-channels_2-m_64-n_64-s_8-num_batches_1]
test_transpose[M_2048-N_64-aie_columns_1-channels_2-m_64-n_64-s_8]
Krackan - ExamplesIRONTested on iron/applications/llama_3.2_1b
Trends: IRON Trendsiron/applications/llama_3.2_1btest_llama_3_2_1b[llama_3.2_1b_prompt_1024_tokens_1]
test_llama_3_2_1b[llama_3.2_1b_prompt_1024_tokens_40]
test_llama_3_2_1b[llama_3.2_1b_prompt_13_tokens_1]
test_llama_3_2_1b[llama_3.2_1b_prompt_13_tokens_40]
Phoenix - SmallIRONTested on iron/operators/axpy
iron/operators/dequant
iron/operators/elementwise_add
iron/operators/elementwise_mul
iron/operators/gelu
iron/operators/gemm
iron/operators/gemv
iron/operators/layer_norm
iron/operators/mem_copy
iron/operators/relu
iron/operators/rms_norm
iron/operators/rope
iron/operators/sigmoid
iron/operators/silu
iron/operators/softmax
iron/operators/swiglu_decode
iron/operators/swiglu_prefill
iron/operators/tanh
iron/operators/transpose
Trends: IRON Trendsiron/operators/axpytest_axpy[input_length_2048-num_aie_columns_1-tile_size_2048-scalar_factor_3.0]
test_axpy[input_length_2048-num_aie_columns_2-tile_size_1024-scalar_factor_3.0]
test_axpy[input_length_2048-num_aie_columns_4-tile_size_512-scalar_factor_3.0]
iron/operators/dequanttest_dequant[input_length_2048-num_aie_columns_1-num_channels_1-tile_size_2048-group_size_32]
test_dequant[input_length_2048-num_aie_columns_1-num_channels_2-tile_size_1024-group_size_32]
test_dequant[input_length_2048-num_aie_columns_2-num_channels_1-tile_size_1024-group_size_32]
test_dequant[input_length_2048-num_aie_columns_2-num_channels_2-tile_size_512-group_size_32]
test_dequant[input_length_2048-num_aie_columns_4-num_channels_1-tile_size_512-group_size_32]
test_dequant[input_length_2048-num_aie_columns_4-num_channels_2-tile_size_256-group_size_32]
iron/operators/elementwise_addtest_elementwise_add[input_length_2048-num_aie_columns_1-tile_size_2048]
test_elementwise_add[input_length_2048-num_aie_columns_2-tile_size_1024]
test_elementwise_add[input_length_2048-num_aie_columns_4-tile_size_512]
iron/operators/elementwise_multest_elementwise_mul[input_length_2048-num_aie_columns_1-tile_size_2048]
test_elementwise_mul[input_length_2048-num_aie_columns_2-tile_size_1024]
test_elementwise_mul[input_length_2048-num_aie_columns_4-tile_size_512]
iron/operators/gelutest_gelu[input_length_2048-num_aie_columns_1-num_channels_1-tile_size_2048]
test_gelu[input_length_2048-num_aie_columns_1-num_channels_2-tile_size_1024]
test_gelu[input_length_2048-num_aie_columns_2-num_channels_1-tile_size_1024]
test_gelu[input_length_2048-num_aie_columns_2-num_channels_2-tile_size_512]
test_gelu[input_length_2048-num_aie_columns_4-num_channels_1-tile_size_512]
test_gelu[input_length_2048-num_aie_columns_4-num_channels_2-tile_size_256]
iron/operators/gemmtest_gemm[M_192-K_384-N_64-num_aie_columns_4-b_col_maj_False-c_col_maj_False-m_48-k_96-n_16-trace_size_0-partition_N_1]
test_gemm[M_192-K_384-N_64-num_aie_columns_4-b_col_maj_True-c_col_maj_True-m_48-k_96-n_16-trace_size_0-partition_N_1]
test_gemm[M_2048-K_2048-N_2048-num_aie_columns_1-b_col_maj_False-c_col_maj_False-m_64-k_64-n_64-trace_size_0-partition_N_1]
test_gemm[M_2048-K_2048-N_2048-num_aie_columns_2-b_col_maj_True-c_col_maj_False-m_64-k_64-n_64-trace_size_0-partition_N_1]
test_gemm[M_384-K_1536-N_1792-num_aie_columns_4-b_col_maj_True-c_col_maj_False-m_32-k_48-n_64-trace_size_0-partition_N_1]
test_gemm[M_64-K_512-N_256-num_aie_columns_4-b_col_maj_True-c_col_maj_False-m_16-k_64-n_64-trace_size_0-partition_N_4]
iron/operators/gemvtest_gemv[M_128-K_128-num_aie_columns_1-tile_size_input_32-tile_size_output_128]
test_gemv[M_2048-K_8192-num_aie_columns_1-tile_size_input_1-tile_size_output_2048]
test_gemv[M_2048-K_8192-num_aie_columns_2-tile_size_input_1-tile_size_output_1024]
test_gemv[M_2048-K_8192-num_aie_columns_4-tile_size_input_1-tile_size_output_512]
test_gemv[M_8192-K_2048-num_aie_columns_1-tile_size_input_4-tile_size_output_1024]
test_gemv[M_8192-K_2048-num_aie_columns_2-tile_size_input_4-tile_size_output_1024]
test_gemv[M_8192-K_2048-num_aie_columns_4-tile_size_input_4-tile_size_output_1024]
test_gemv_batched[M_1024-K_1024-num_aie_columns_1-tile_size_input_1-tile_size_output_64-num_batches_2]No metrics available. test_gemv_batched[M_1026-K_64-num_aie_columns_1-tile_size_input_1-tile_size_output_2-num_batches_2]No metrics available. test_gemv_batched[M_256-K_128-num_aie_columns_1-tile_size_input_1-tile_size_output_256-num_batches_4]No metrics available. test_gemv_batched[M_64-K_1536-num_aie_columns_1-tile_size_input_1-tile_size_output_64-num_batches_8]No metrics available. iron/operators/layer_normtest_layer_norm[input_length_2048-num_aie_columns_1-num_channels_1-tile_size_2048]
test_layer_norm[input_length_2048-num_aie_columns_1-num_channels_2-tile_size_1024]
test_layer_norm[input_length_2048-num_aie_columns_2-num_channels_1-tile_size_1024]
test_layer_norm[input_length_2048-num_aie_columns_2-num_channels_2-tile_size_512]
test_layer_norm[input_length_2048-num_aie_columns_4-num_channels_1-tile_size_512]
test_layer_norm[input_length_2048-num_aie_columns_4-num_channels_2-tile_size_256]
iron/operators/mem_copytest_mem_copy[input_length_2048-num_cores_1-num_channels_1-bypass_False-tile_size_2048]
test_mem_copy[input_length_2048-num_cores_2-num_channels_1-bypass_False-tile_size_1024]
test_mem_copy[input_length_2048-num_cores_2-num_channels_2-bypass_False-tile_size_1024]
test_mem_copy[input_length_2048-num_cores_4-num_channels_1-bypass_False-tile_size_512]
test_mem_copy[input_length_2048-num_cores_4-num_channels_2-bypass_False-tile_size_512]
test_mem_copy[input_length_2048-num_cores_8-num_channels_2-bypass_False-tile_size_256]
iron/operators/rms_normtest_rms_norm[input_length_2048-num_aie_columns_1-num_channels_1-tile_size_2048-weighted_False]
test_rms_norm[input_length_2048-num_aie_columns_1-num_channels_1-tile_size_2048-weighted_True]
test_rms_norm[input_length_2048-num_aie_columns_1-num_channels_2-tile_size_1024-weighted_False]
test_rms_norm[input_length_2048-num_aie_columns_1-num_channels_2-tile_size_1024-weighted_True]
test_rms_norm[input_length_2048-num_aie_columns_2-num_channels_1-tile_size_1024-weighted_False]
test_rms_norm[input_length_2048-num_aie_columns_2-num_channels_1-tile_size_1024-weighted_True]
test_rms_norm[input_length_2048-num_aie_columns_2-num_channels_2-tile_size_512-weighted_False]
test_rms_norm[input_length_2048-num_aie_columns_2-num_channels_2-tile_size_512-weighted_True]
test_rms_norm[input_length_2048-num_aie_columns_4-num_channels_1-tile_size_512-weighted_False]
test_rms_norm[input_length_2048-num_aie_columns_4-num_channels_1-tile_size_512-weighted_True]
test_rms_norm[input_length_2048-num_aie_columns_4-num_channels_2-tile_size_256-weighted_False]
iron/operators/ropetest_rope[rows_32-cols_512-angle_rows_32-aie_columns_1-method_type_0]
test_rope[rows_32-cols_512-angle_rows_32-aie_columns_2-method_type_0]
test_rope[rows_32-cols_512-angle_rows_32-aie_columns_4-method_type_0]
test_rope[rows_32-cols_512-angle_rows_8-aie_columns_1-method_type_0]
test_rope[rows_32-cols_512-angle_rows_8-aie_columns_2-method_type_0]
test_rope[rows_32-cols_512-angle_rows_8-aie_columns_4-method_type_0]
iron/operators/softmaxtest_softmax[input_length_32768-num_aie_columns_2-num_channels_2-tile_size_1024]
test_softmax[input_length_32768-num_aie_columns_2-num_channels_2-tile_size_2048]
test_softmax[input_length_32768-num_aie_columns_2-num_channels_2-tile_size_512]
iron/operators/swiglu_decodetest_swiglu_decode[embedding_dim_1024-hidden_dim_3584]
test_swiglu_decode[embedding_dim_2048-hidden_dim_2048]
iron/operators/swiglu_prefilltest_swiglu_prefill[seq_len_256-embedding_dim_2048-hidden_dim_2048-prio_accuracy_False]
iron/operators/transposetest_transpose[M_2048-N_64-aie_columns_1-channels_1-m_64-n_64-s_8-num_batches_1]
test_transpose[M_2048-N_64-aie_columns_1-channels_1-m_64-n_64-s_8-num_batches_2]
test_transpose[M_2048-N_64-aie_columns_1-channels_1-m_64-n_64-s_8]
test_transpose[M_2048-N_64-aie_columns_1-channels_2-m_64-n_64-s_8-num_batches_1]
test_transpose[M_2048-N_64-aie_columns_1-channels_2-m_64-n_64-s_8]
Phoenix - ExamplesIRONTested on Trends: IRON Trends |
…t, dedup branches) Responds to andrej's review on amd#127: - Explain the coalesced TAP is NOT a single linear transfer: contiguous within a batch but strided by the full matrix across batches (M*K != run for cols>1), so the batch iteration dim and the TAP are required. The inner run-split is only to fit the 1023 wrap cap; note it is removable once IRON's mlir_aie pin moves past mlir-aie #3036. - Add FIXME to source the shim BD bounds from the MLIR-AIE target model (AIEXDialect.cpp). - Reword the depth>=2 rationale: ObjectFifo backpressure blocks (never overruns/corrupts); depth>=2 is a performance/overlap guard, not correctness. - Dedup the coalesced/fallback branches into one loop over num_waits (1 vs num_batches).
|
Pushed all of the above. I also added |
|
Thanks for looking into those suggestions. I did just update the pinned MLIR-AIE. If you can look into whether that helps you simplify any of the linearizing code, that would be great. After that, or if it doesn't end up helping you, I'll run the workflows and merge :) |
The batched GEMV unrolled num_batches host DMA descriptors per column (one fill + one drain + a task-group wait per batch). Express the batch as a single iterated BD instead, the same B-unroll -> BD-iteration idiom GEMM already uses for tiling: place num_batches in the size-uncapped descriptor dim and split the contiguous run across the two wrap dims (each <= 1023), per the AIE shim BD limits (AIEXDialect.cpp verifyStridesWraps). The run split keeps the inner size and strides 4-byte-aligned (granularity-aware), and the per-batch drain wait is dropped in favour of ObjectFifo backpressure (fifo depth >= 2 asserted; the A-fill and C-drain run on separate shim channels, paced by the fifo locks). num_batches == 1 and any config that cannot be coalesced (run with no aligned split under the wrap cap, or a batch stride that is too large or unaligned) fall back to the existing per-batch path and are byte-identical. The coalesced descriptor accesses the exact same DRAM elements in the same order as the unroll (access-equivalent), so this is a descriptor-count / build-time and correctness change, not a runtime change. (cherry picked from commit d4bde97)
The test suite had no num_batches > 1 GEMV coverage. Add a batched golden reference (num_batches independent matrix-vector products, stacked contiguously) and a parametrized device test covering: the coalesced path with large num_batches (the size-uncapped dim) and a multi-dimension run split; a run that requires an aligned (even) inner split; and the per-batch fallback (batch stride over the limit). (cherry picked from commit 6d9a845)
…t, dedup branches) Responds to andrej's review on amd#127: - Explain the coalesced TAP is NOT a single linear transfer: contiguous within a batch but strided by the full matrix across batches (M*K != run for cols>1), so the batch iteration dim and the TAP are required. The inner run-split is only to fit the 1023 wrap cap; note it is removable once IRON's mlir_aie pin moves past mlir-aie #3036. - Add FIXME to source the shim BD bounds from the MLIR-AIE target model (AIEXDialect.cpp). - Reword the depth>=2 rationale: ObjectFifo backpressure blocks (never overruns/corrupts); depth>=2 is a performance/overlap guard, not correctness. - Dedup the coalesced/fallback branches into one loop over num_waits (1 vs num_batches).
Verified on the v1.3.4 pin (mlir-aie #2924 + #3036): the contiguous-run linearization only bypasses the 1023 wrap cap when the whole transfer is contiguous (cols==1, bstride==run). For a real cols>1 gather the inner run stays a wrap dim -- [1, num_batches, 4096] fails aiecc 'Size 4096 exceeds [0:1023]'. Correcting the earlier speculative FUTURE note; the split_run stays.
|
@andrej Hi there |


Replaces the GEMV operator's per-batch host DMA descriptor unroll with a single iterated descriptor per column.
GEMM already folds many transfers into one multi-dimensional iterated descriptor (via
TensorTiler2D); GEMV did not, so a batched GEMV emittednum_batchesseparate fill/drain descriptors (plus a per-batch task-groupwait) per AIE column, linear in the batch count. The test suite also had nonum_batches > 1GEMV coverage, so the batched path shipped untested.Added
num_batches > 1golden reference (generate_golden_reference_batched) and a parametrized device test (test_gemv_batched) covering: the coalesced path with largenum_batches(the size-uncapped descriptor dimension) and a multi-dimension run split; a run that requires an alignment-aware (even) inner split; the per-batch fallback (a batch stride over the limit); and an attention-style shape withtile_size_input > 1andnum_batches= a head count.Changed
num_batchesis placed in the size-uncapped descriptor dimension and the contiguous per-batch run is split across the two wrap dimensions (each <= 1023), per the AIE shim BD limits (AIEXDialect.cppverifyStridesWraps). The run split keeps the innermost size and the strides 4-byte-aligned (the shim's address granularity, which is enforced even for linear transfers).waitis replaced by ObjectFifo backpressure. Depth 2 is sufficient (rather than depthnum_batches): the A-fill (MM2S) and C-drain (S2MM) run on separate shim channels and the pipeline is paced by the fifo locks, so the producer cannot overrun the consumer; this is a streaming pipeline, not a store-all buffer. The depth invariant is asserted.num_batches == 1, and any configuration that cannot be coalesced (a run with no aligned split under the wrap cap, or a batch stride that is too large or not aligned), fall back to the existing per-batch path and are byte-identical. The coalesced path therefore introduces no new failure modes relative to the unroll.Removed
Motivation and measurements
This is a descriptor-count / build-time and correctness change, not a runtime speedup: the transfer is access-equivalent (same bytes, same order), so it does not change data-movement time.
The descriptor count for a batched GEMV drops from O(num_batches) to O(1) per column, which shrinks the descriptor-bound lowering phase of the build. Measured cold compile of a single GEMV (shape M=448, K=64, columns=8, build only, no device run):
The coalesced build is roughly constant in
num_batcheswhile the unrolled build grows linearly, so the saving scales with the batch count.Validation
num_batches == 1generates byte-identical lowered output to the previous code.num_batchesof 4, 8, 32, 100, and 192, plus the aligned-split shape (M=1026), the fallback shape (batch stride > limit), and an attention-style shape (tile_size_input=4); the existingnum_batches == 1tests are unchanged.llama_3.2_1bapplication, which calls GEMV withnum_batches = n_heads, builds and generates coherent output on NPU2 with the coalesced path.swiglu_decodeusesnum_batches=1(byte-identical, unaffected); thellama_3.2_1bapplication calls GEMV withnum_batches = n_headsfor the attention scores/context GEMVs - those now coalesce, and the coalesced descriptor is legal and access-equivalent at those shapes, so behavior is unchanged.PR Merge Checklist
develcommit and pointing todevel.