This can redefine PIPE_FIX (PIPE_FIX not exist in system as MACRO) to the wrong pipeline symbol, which causes abnormal AIC pipeline behavior and can make the kernel hang.
module attributes {pto.target_arch = "a2a3"} {
func.func @wait_gemm(%arg0: !pto.ptr<f16>, %arg1: !pto.ptr<f16>, %arg2: !pto.ptr<i32>, %arg3: !pto.ptr<f32>, %arg4: !pto.ptr<i64>) attributes {pto.kernel_kind = #pto.kernel_kind<cube>} {
%c0_i64 = arith.constant 0 : i64
%c64_index = arith.constant 64 : index
%c1_index = arith.constant 1 : index
%c2_index = arith.constant 2 : index
%c32_i64 = arith.constant 32 : i64
%c0_index = arith.constant 0 : index
%c1_i32 = arith.constant 1 : i32
%local_a__ssa_v0_view = pto.make_tensor_view %arg0, shape = [%c64_index, %c64_index], strides = [%c64_index, %c1_index] {layout = #pto.layout<nd>}: !pto.tensor_view<?x?xf16>
%local_b__ssa_v0_view = pto.make_tensor_view %arg1, shape = [%c64_index, %c64_index], strides = [%c64_index, %c1_index] {layout = #pto.layout<nd>}: !pto.tensor_view<?x?xf16>
%signal__ssa_v0_view = pto.make_tensor_view %arg2, shape = [%c2_index, %c1_index], strides = [%c1_index, %c2_index] {layout = #pto.layout<dn>}: !pto.tensor_view<?x?xi32>
%out__ssa_v0_view = pto.make_tensor_view %arg3, shape = [%c64_index, %c64_index], strides = [%c64_index, %c1_index] {layout = #pto.layout<nd>}: !pto.tensor_view<?x?xf32>
%0 = pto.load_scalar %arg4[%c2_index] : !pto.ptr<i64> -> i64
%my_rank__ssa_v0 = arith.trunci %0 : i64 to i32
%1 = pto.load_scalar %arg4[%c2_index] : !pto.ptr<i64> -> i64
%2 = arith.shrui %1, %c32_i64 : i64
%nranks__ssa_v0 = arith.trunci %2 : i64 to i32
%nranks__ssa_v0_idx = arith.index_cast %nranks__ssa_v0 : i32 to index
scf.for %peer__idx_v0 = %c0_index to %nranks__ssa_v0_idx step %c1_index {
%3 = arith.index_cast %my_rank__ssa_v0 : i32 to index
%4 = arith.cmpi ne, %peer__idx_v0, %3 : index
scf.if %4 {
%5 = func.call @CommRemoteOffset_i32(%arg4, %peer__idx_v0) : (!pto.ptr<i64>, index) -> index
%6 = pto.addptr %arg2, %5 : !pto.ptr<i32> -> !pto.ptr<i32>
%7 = arith.muli %c1_index, %c1_index : index
%8 = pto.make_tensor_view %6, shape = [%c2_index, %c1_index], strides = [%7, %c1_index] {layout = #pto.layout<nd>} : !pto.tensor_view<?x?xi32>
%my_rank__ssa_v0_idx = arith.index_cast %my_rank__ssa_v0 : i32 to index
%signal__ssa_v0_peer_pview = pto.partition_view %8, offsets = [%my_rank__ssa_v0_idx, %c0_index], sizes = [%c1_index, %c1_index] : !pto.tensor_view<?x?xi32> -> !pto.partition_tensor_view<1x1xi32>
pto.comm.tnotify(%signal__ssa_v0_peer_pview, %c1_i32 : !pto.partition_tensor_view<1x1xi32>, i32) {notifyOp = #pto<notify_op atomic_add>}
}
}
%9 = arith.index_cast %nranks__ssa_v0 : i32 to index
scf.for %src_rank__idx_v0 = %c0_index to %9 step %c1_index {
%10 = arith.index_cast %my_rank__ssa_v0 : i32 to index
%11 = arith.cmpi ne, %src_rank__idx_v0, %10 : index
scf.if %11 {
%signal__ssa_v0_local_pview = pto.partition_view %signal__ssa_v0_view, offsets = [%src_rank__idx_v0, %c0_index], sizes = [%c1_index, %c1_index] : !pto.tensor_view<?x?xi32> -> !pto.partition_tensor_view<1x1xi32>
pto.comm.twait(%signal__ssa_v0_local_pview, %c1_i32 : !pto.partition_tensor_view<1x1xi32>, i32) {cmp = #pto<wait_cmp ge>}
}
}
%a_mat__ssa_v0 = pto.alloc_tile addr = %c0_i64 valid_row = %c64_index valid_col = %c64_index : !pto.tile_buf<loc=mat, dtype=f16, rows=64, cols=64, v_row=?, v_col=?, blayout=col_major, slayout=row_major, fractal=512, pad=0>
%local_a__ssa_v0_pview = pto.partition_view %local_a__ssa_v0_view, offsets = [%c0_index, %c0_index], sizes = [%c64_index, %c64_index] : !pto.tensor_view<?x?xf16> -> !pto.partition_tensor_view<64x64xf16>
pto.tload ins(%local_a__ssa_v0_pview : !pto.partition_tensor_view<64x64xf16>) outs(%a_mat__ssa_v0 : !pto.tile_buf<loc=mat, dtype=f16, rows=64, cols=64, v_row=?, v_col=?, blayout=col_major, slayout=row_major, fractal=512, pad=0>)
%a_left__ssa_v0 = pto.alloc_tile addr = %c0_i64 valid_row = %c64_index valid_col = %c64_index : !pto.tile_buf<loc=left, dtype=f16, rows=64, cols=64, v_row=?, v_col=?, blayout=col_major, slayout=row_major, fractal=512, pad=0>
pto.tmov ins(%a_mat__ssa_v0 : !pto.tile_buf<loc=mat, dtype=f16, rows=64, cols=64, v_row=?, v_col=?, blayout=col_major, slayout=row_major, fractal=512, pad=0>) outs(%a_left__ssa_v0 : !pto.tile_buf<loc=left, dtype=f16, rows=64, cols=64, v_row=?, v_col=?, blayout=col_major, slayout=row_major, fractal=512, pad=0>)
%b_mat__ssa_v0 = pto.alloc_tile addr = %c0_i64 valid_row = %c64_index valid_col = %c64_index : !pto.tile_buf<loc=mat, dtype=f16, rows=64, cols=64, v_row=?, v_col=?, blayout=col_major, slayout=row_major, fractal=512, pad=0>
%local_b__ssa_v0_pview = pto.partition_view %local_b__ssa_v0_view, offsets = [%c0_index, %c0_index], sizes = [%c64_index, %c64_index] : !pto.tensor_view<?x?xf16> -> !pto.partition_tensor_view<64x64xf16>
pto.tload ins(%local_b__ssa_v0_pview : !pto.partition_tensor_view<64x64xf16>) outs(%b_mat__ssa_v0 : !pto.tile_buf<loc=mat, dtype=f16, rows=64, cols=64, v_row=?, v_col=?, blayout=col_major, slayout=row_major, fractal=512, pad=0>)
%b_right__ssa_v0 = pto.alloc_tile addr = %c0_i64 valid_row = %c64_index valid_col = %c64_index : !pto.tile_buf<loc=right, dtype=f16, rows=64, cols=64, v_row=?, v_col=?, blayout=row_major, slayout=col_major, fractal=512, pad=0>
pto.tmov ins(%b_mat__ssa_v0 : !pto.tile_buf<loc=mat, dtype=f16, rows=64, cols=64, v_row=?, v_col=?, blayout=col_major, slayout=row_major, fractal=512, pad=0>) outs(%b_right__ssa_v0 : !pto.tile_buf<loc=right, dtype=f16, rows=64, cols=64, v_row=?, v_col=?, blayout=row_major, slayout=col_major, fractal=512, pad=0>)
%acc__ssa_v0 = pto.alloc_tile addr = %c0_i64 valid_row = %c64_index valid_col = %c64_index : !pto.tile_buf<loc=acc, dtype=f32, rows=64, cols=64, v_row=?, v_col=?, blayout=col_major, slayout=row_major, fractal=1024, pad=0>
pto.tmatmul ins(%a_left__ssa_v0, %b_right__ssa_v0 : !pto.tile_buf<loc=left, dtype=f16, rows=64, cols=64, v_row=?, v_col=?, blayout=col_major, slayout=row_major, fractal=512, pad=0>, !pto.tile_buf<loc=right, dtype=f16, rows=64, cols=64, v_row=?, v_col=?, blayout=row_major, slayout=col_major, fractal=512, pad=0>) outs(%acc__ssa_v0 : !pto.tile_buf<loc=acc, dtype=f32, rows=64, cols=64, v_row=?, v_col=?, blayout=col_major, slayout=row_major, fractal=1024, pad=0>)
%out__ssa_v0_pview = pto.partition_view %out__ssa_v0_view, offsets = [%c0_index, %c0_index], sizes = [%c64_index, %c64_index] : !pto.tensor_view<?x?xf32> -> !pto.partition_tensor_view<64x64xf32>
pto.tstore ins(%acc__ssa_v0 : !pto.tile_buf<loc=acc, dtype=f32, rows=64, cols=64, v_row=?, v_col=?, blayout=col_major, slayout=row_major, fractal=1024, pad=0>) outs(%out__ssa_v0_pview : !pto.partition_tensor_view<64x64xf32>)
return
}
func.func private @CommRemoteOffset_i32(%ctx: !pto.ptr<i64>, %peer: index) -> index {
%c_r = arith.constant 2 : index
%c_w = arith.constant 4 : index
%rk_pair = pto.load_scalar %ctx[%c_r] : !pto.ptr<i64> -> i64
%rk_i32 = arith.trunci %rk_pair : i64 to i32
%rk_idx = arith.index_cast %rk_i32 : i32 to index
%lb_off = arith.addi %c_w, %rk_idx : index
%lbase = pto.load_scalar %ctx[%lb_off] : !pto.ptr<i64> -> i64
%pb_off = arith.addi %c_w, %peer : index
%pbase = pto.load_scalar %ctx[%pb_off] : !pto.ptr<i64> -> i64
%dbytes = arith.subi %pbase, %lbase : i64
%esize = arith.constant 4 : i64
%delems_i = arith.divsi %dbytes, %esize : i64
%delems = arith.index_cast %delems_i : i64 to index
return %delems : index
}
}
Component
EmitC / Codegen (lib/PTO/Transforms/PTOToEmitC.cpp)
Description
PTOAS incorrectly emits the following code when generating an AIC kernel:
This can redefine PIPE_FIX (PIPE_FIX not exist in system as MACRO) to the wrong pipeline symbol, which causes abnormal AIC pipeline behavior and can make the kernel hang.
log.txt
The macro emission appears to come from:
https://github.com/hw-native-sys/PTOAS/blame/v0.54/lib/PTO/Transforms/PTOToEmitC.cpp#L14671-L14675
Minimal reproducer and generated kernel:
pipe_fix_issue.py
wait_gemm.cpp
PyPTO version used for reproduction:
478ddad5a41011645a45019d5196e414c6dc3745
Reproduction (minimal)
module attributes {pto.target_arch = "a2a3"} { func.func @wait_gemm(%arg0: !pto.ptr<f16>, %arg1: !pto.ptr<f16>, %arg2: !pto.ptr<i32>, %arg3: !pto.ptr<f32>, %arg4: !pto.ptr<i64>) attributes {pto.kernel_kind = #pto.kernel_kind<cube>} { %c0_i64 = arith.constant 0 : i64 %c64_index = arith.constant 64 : index %c1_index = arith.constant 1 : index %c2_index = arith.constant 2 : index %c32_i64 = arith.constant 32 : i64 %c0_index = arith.constant 0 : index %c1_i32 = arith.constant 1 : i32 %local_a__ssa_v0_view = pto.make_tensor_view %arg0, shape = [%c64_index, %c64_index], strides = [%c64_index, %c1_index] {layout = #pto.layout<nd>}: !pto.tensor_view<?x?xf16> %local_b__ssa_v0_view = pto.make_tensor_view %arg1, shape = [%c64_index, %c64_index], strides = [%c64_index, %c1_index] {layout = #pto.layout<nd>}: !pto.tensor_view<?x?xf16> %signal__ssa_v0_view = pto.make_tensor_view %arg2, shape = [%c2_index, %c1_index], strides = [%c1_index, %c2_index] {layout = #pto.layout<dn>}: !pto.tensor_view<?x?xi32> %out__ssa_v0_view = pto.make_tensor_view %arg3, shape = [%c64_index, %c64_index], strides = [%c64_index, %c1_index] {layout = #pto.layout<nd>}: !pto.tensor_view<?x?xf32> %0 = pto.load_scalar %arg4[%c2_index] : !pto.ptr<i64> -> i64 %my_rank__ssa_v0 = arith.trunci %0 : i64 to i32 %1 = pto.load_scalar %arg4[%c2_index] : !pto.ptr<i64> -> i64 %2 = arith.shrui %1, %c32_i64 : i64 %nranks__ssa_v0 = arith.trunci %2 : i64 to i32 %nranks__ssa_v0_idx = arith.index_cast %nranks__ssa_v0 : i32 to index scf.for %peer__idx_v0 = %c0_index to %nranks__ssa_v0_idx step %c1_index { %3 = arith.index_cast %my_rank__ssa_v0 : i32 to index %4 = arith.cmpi ne, %peer__idx_v0, %3 : index scf.if %4 { %5 = func.call @CommRemoteOffset_i32(%arg4, %peer__idx_v0) : (!pto.ptr<i64>, index) -> index %6 = pto.addptr %arg2, %5 : !pto.ptr<i32> -> !pto.ptr<i32> %7 = arith.muli %c1_index, %c1_index : index %8 = pto.make_tensor_view %6, shape = [%c2_index, %c1_index], strides = [%7, %c1_index] {layout = #pto.layout<nd>} : !pto.tensor_view<?x?xi32> %my_rank__ssa_v0_idx = arith.index_cast %my_rank__ssa_v0 : i32 to index %signal__ssa_v0_peer_pview = pto.partition_view %8, offsets = [%my_rank__ssa_v0_idx, %c0_index], sizes = [%c1_index, %c1_index] : !pto.tensor_view<?x?xi32> -> !pto.partition_tensor_view<1x1xi32> pto.comm.tnotify(%signal__ssa_v0_peer_pview, %c1_i32 : !pto.partition_tensor_view<1x1xi32>, i32) {notifyOp = #pto<notify_op atomic_add>} } } %9 = arith.index_cast %nranks__ssa_v0 : i32 to index scf.for %src_rank__idx_v0 = %c0_index to %9 step %c1_index { %10 = arith.index_cast %my_rank__ssa_v0 : i32 to index %11 = arith.cmpi ne, %src_rank__idx_v0, %10 : index scf.if %11 { %signal__ssa_v0_local_pview = pto.partition_view %signal__ssa_v0_view, offsets = [%src_rank__idx_v0, %c0_index], sizes = [%c1_index, %c1_index] : !pto.tensor_view<?x?xi32> -> !pto.partition_tensor_view<1x1xi32> pto.comm.twait(%signal__ssa_v0_local_pview, %c1_i32 : !pto.partition_tensor_view<1x1xi32>, i32) {cmp = #pto<wait_cmp ge>} } } %a_mat__ssa_v0 = pto.alloc_tile addr = %c0_i64 valid_row = %c64_index valid_col = %c64_index : !pto.tile_buf<loc=mat, dtype=f16, rows=64, cols=64, v_row=?, v_col=?, blayout=col_major, slayout=row_major, fractal=512, pad=0> %local_a__ssa_v0_pview = pto.partition_view %local_a__ssa_v0_view, offsets = [%c0_index, %c0_index], sizes = [%c64_index, %c64_index] : !pto.tensor_view<?x?xf16> -> !pto.partition_tensor_view<64x64xf16> pto.tload ins(%local_a__ssa_v0_pview : !pto.partition_tensor_view<64x64xf16>) outs(%a_mat__ssa_v0 : !pto.tile_buf<loc=mat, dtype=f16, rows=64, cols=64, v_row=?, v_col=?, blayout=col_major, slayout=row_major, fractal=512, pad=0>) %a_left__ssa_v0 = pto.alloc_tile addr = %c0_i64 valid_row = %c64_index valid_col = %c64_index : !pto.tile_buf<loc=left, dtype=f16, rows=64, cols=64, v_row=?, v_col=?, blayout=col_major, slayout=row_major, fractal=512, pad=0> pto.tmov ins(%a_mat__ssa_v0 : !pto.tile_buf<loc=mat, dtype=f16, rows=64, cols=64, v_row=?, v_col=?, blayout=col_major, slayout=row_major, fractal=512, pad=0>) outs(%a_left__ssa_v0 : !pto.tile_buf<loc=left, dtype=f16, rows=64, cols=64, v_row=?, v_col=?, blayout=col_major, slayout=row_major, fractal=512, pad=0>) %b_mat__ssa_v0 = pto.alloc_tile addr = %c0_i64 valid_row = %c64_index valid_col = %c64_index : !pto.tile_buf<loc=mat, dtype=f16, rows=64, cols=64, v_row=?, v_col=?, blayout=col_major, slayout=row_major, fractal=512, pad=0> %local_b__ssa_v0_pview = pto.partition_view %local_b__ssa_v0_view, offsets = [%c0_index, %c0_index], sizes = [%c64_index, %c64_index] : !pto.tensor_view<?x?xf16> -> !pto.partition_tensor_view<64x64xf16> pto.tload ins(%local_b__ssa_v0_pview : !pto.partition_tensor_view<64x64xf16>) outs(%b_mat__ssa_v0 : !pto.tile_buf<loc=mat, dtype=f16, rows=64, cols=64, v_row=?, v_col=?, blayout=col_major, slayout=row_major, fractal=512, pad=0>) %b_right__ssa_v0 = pto.alloc_tile addr = %c0_i64 valid_row = %c64_index valid_col = %c64_index : !pto.tile_buf<loc=right, dtype=f16, rows=64, cols=64, v_row=?, v_col=?, blayout=row_major, slayout=col_major, fractal=512, pad=0> pto.tmov ins(%b_mat__ssa_v0 : !pto.tile_buf<loc=mat, dtype=f16, rows=64, cols=64, v_row=?, v_col=?, blayout=col_major, slayout=row_major, fractal=512, pad=0>) outs(%b_right__ssa_v0 : !pto.tile_buf<loc=right, dtype=f16, rows=64, cols=64, v_row=?, v_col=?, blayout=row_major, slayout=col_major, fractal=512, pad=0>) %acc__ssa_v0 = pto.alloc_tile addr = %c0_i64 valid_row = %c64_index valid_col = %c64_index : !pto.tile_buf<loc=acc, dtype=f32, rows=64, cols=64, v_row=?, v_col=?, blayout=col_major, slayout=row_major, fractal=1024, pad=0> pto.tmatmul ins(%a_left__ssa_v0, %b_right__ssa_v0 : !pto.tile_buf<loc=left, dtype=f16, rows=64, cols=64, v_row=?, v_col=?, blayout=col_major, slayout=row_major, fractal=512, pad=0>, !pto.tile_buf<loc=right, dtype=f16, rows=64, cols=64, v_row=?, v_col=?, blayout=row_major, slayout=col_major, fractal=512, pad=0>) outs(%acc__ssa_v0 : !pto.tile_buf<loc=acc, dtype=f32, rows=64, cols=64, v_row=?, v_col=?, blayout=col_major, slayout=row_major, fractal=1024, pad=0>) %out__ssa_v0_pview = pto.partition_view %out__ssa_v0_view, offsets = [%c0_index, %c0_index], sizes = [%c64_index, %c64_index] : !pto.tensor_view<?x?xf32> -> !pto.partition_tensor_view<64x64xf32> pto.tstore ins(%acc__ssa_v0 : !pto.tile_buf<loc=acc, dtype=f32, rows=64, cols=64, v_row=?, v_col=?, blayout=col_major, slayout=row_major, fractal=1024, pad=0>) outs(%out__ssa_v0_pview : !pto.partition_tensor_view<64x64xf32>) return } func.func private @CommRemoteOffset_i32(%ctx: !pto.ptr<i64>, %peer: index) -> index { %c_r = arith.constant 2 : index %c_w = arith.constant 4 : index %rk_pair = pto.load_scalar %ctx[%c_r] : !pto.ptr<i64> -> i64 %rk_i32 = arith.trunci %rk_pair : i64 to i32 %rk_idx = arith.index_cast %rk_i32 : i32 to index %lb_off = arith.addi %c_w, %rk_idx : index %lbase = pto.load_scalar %ctx[%lb_off] : !pto.ptr<i64> -> i64 %pb_off = arith.addi %c_w, %peer : index %pbase = pto.load_scalar %ctx[%pb_off] : !pto.ptr<i64> -> i64 %dbytes = arith.subi %pbase, %lbase : i64 %esize = arith.constant 4 : i64 %delems_i = arith.divsi %dbytes, %esize : i64 %delems = arith.index_cast %delems_i : i64 to index return %delems : index } }Expected behavior
the generated AIC kernel should run successfully
Current workaround: add compiler option -DPIPE_FIX=PIPE_FIX in CCECToolchain (pypto/runtime/simpler_setup/toolchain.py)
Actual behavior / error logs
Git commit
195fdb8
Host platform
None
Target Ascend arch (if relevant)
a3
PTOAS build level (if relevant)
None