Skip to content

[Bug] Incorrectly emits PIPE_FIX macro guard in generated AIC kernel #1086

Description

@Shenggan

Component

EmitC / Codegen (lib/PTO/Transforms/PTOToEmitC.cpp)

Description

PTOAS incorrectly emits the following code when generating an AIC kernel:

#ifndef PIPE_FIX
#define PIPE_FIX PIPE_M
#endif

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

the generated AIC kernel pipeline becomes abnormal and can hang
[log.txt](https://github.com/user-attachments/files/30581523/log.txt)

Git commit

195fdb8

Host platform

None

Target Ascend arch (if relevant)

a3

PTOAS build level (if relevant)

None

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions