Skip to content

Add tensor-based TileArray programming and GEMM lowering - #4

Merged
ShangkunLi merged 2 commits into
mainfrom
split/v2-frontend
Sep 17, 2026
Merged

ShangkunLi merged 2 commits into
mainfrom
split/v2-frontend

Conversation

@ShangkunLi

Copy link
Copy Markdown
Collaborator

Summary

This PR adds tensor memory programming and standalone lowering for TileArray programs.

It introduces:

  • typed tensor arguments and tensor slices;
  • configured and dynamic load and store operations;
  • TileArray operations with multiple results;
  • lowering from a TileArray program to a Taskflow task and Neura kernel;
  • a 3×3 weight-stationary GEMM example;
  • tests for memory access, lowering, and complete GEMM compilation.

Memory access

Loads and stores can use a configured tensor slice:

value = synl.load(A[:, k], tile=array[0, y])
synl.store(value, target=C[:, x], tile=array[x + 1, 0])

They can also use an explicit runtime address:

value = synl.load(addr=address, dtype=synl.i32, tile=array[0, 1])
synl.store(value, addr=address, tile=array[1, 0])

Example: 3×3 GEMM on a 4×4 TileArray

The west column loads activations, the inner 3×3 Tiles perform MACs, and the south row stores the result:

LD → MAC → MAC → MAC
       ↓     ↓     ↓
LD → MAC → MAC → MAC
       ↓     ↓     ↓
LD → MAC → MAC → MAC
       ↓     ↓     ↓
      ST    ST    ST

Each MAC produces a partial sum flowing south and forwards its activation east:

accumulated, flowing = synl.mac(
    flowing,
    previous_sum,
    stationary=B[k, x - 1],
    tile=array[x, y],
)

Synapse records this program as a TileArrayProgram, wraps it in a Taskflow task, and lowers it to a Neura kernel for template mapping.

Supports typed tensor accesses, configured and dynamic memory operations,
and standalone GEMM compilation through the shared TileArray representation.
Includes backend dependency and CI setup with the corresponding tests.

Split from #3 at 748f7b7.
The Amoeba pin and local-document ignore rules follow b4de0c4.
@ShangkunLi
ShangkunLi requested a review from guosran September 16, 2026 12:01
Comment thread python/synapse/language/tile_array_program.py
Comment thread python/synapse/library/gemm.py
Comment thread tests/python/compiler/test_systolic_gemm.py
@guosran guosran closed this Sep 17, 2026
@guosran guosran reopened this Sep 17, 2026
@ShangkunLi
ShangkunLi merged commit f3c064b into main Sep 17, 2026
2 checks passed
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.

2 participants