Skip to content

fix: correct knob annotate() mem_space values and TileOp emission#60

Open
ymwangg wants to merge 1 commit into
pr-58from
fix/nkigen-builder-knob-tile-and-memspace
Open

fix: correct knob annotate() mem_space values and TileOp emission#60
ymwangg wants to merge 1 commit into
pr-58from
fix/nkigen-builder-knob-tile-and-memspace

Conversation

@ymwangg

@ymwangg ymwangg commented Jun 3, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Fix ms_map in builder.py's annotate() to use 1-based enum values matching the C++ dialect definition in NkipyAttrs.td (Hbm=1, Psum=2, Sbuf=3, SharedHbm=4)
  • Fix TileOp emission to merge tile_size and reduction_tile into a single loop_tile_size array, since TileOp doesn't have a reduction_tile attribute

Test plan

Tested with #58

  • pytest tests/test_nkigen_numerical.py tests/test_nkigen_ops.py tests/unit/test_nkigen_backend.py -n auto — 45 passed, 8 xfailed (previously 32 failed)

Two bugs in builder.py's annotate() function:

1. ms_map used 0-based values (Hbm=0, Psum=1, Sbuf=2, SharedHbm=3) but
   the C++ dialect defines them starting at 1 (NkipyAttrs.td). This caused
   annotate-memory-space pass to fail with type mismatches on return ops.

2. TileOp was called with a non-existent `reduction_tile` kwarg. The
   dialect only supports `loop_tile_size` which should contain the full
   iteration space (e.g. [M, N, K] for matmul). Now merges tile_size and
   reduction_tile into a single loop_tile_size array before emission.
@ymwangg ymwangg requested a review from shaojiex-aws June 3, 2026 21:27
@ymwangg ymwangg closed this Jun 4, 2026
@ymwangg ymwangg reopened this Jun 4, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant