diff --git a/.github/workflows/dist.yml b/.github/workflows/dist.yml index 98d213a94c..867ded1482 100644 --- a/.github/workflows/dist.yml +++ b/.github/workflows/dist.yml @@ -77,6 +77,10 @@ jobs: run: | uv run --no-project --with=build -m -- build --sdist --outdir=dist + - name: Verify SDist contains HiGHS + run: | + tar -tzf dist/*.tar.gz | grep -q '/3rdparty/highs/CMakeLists.txt$' + - name: Setup ccache uses: hendrikmuhs/ccache-action@v1 with: @@ -97,6 +101,11 @@ jobs: cd "${TEMP_DIR}" uv pip install -v dist/*.tar.gz python3 -c "import tilelang; print(tilelang.__version__)" + TILELANG_LIBRARY="$(python3 -c 'import pathlib, tilelang; print(next((pathlib.Path(tilelang.__file__).parent / "lib").glob("libtilelang.*")))')" + if otool -L "${TILELANG_LIBRARY}" | grep -q 'libhighs'; then + echo "libtilelang must not have a runtime dependency on libhighs" + exit 1 + fi - name: Upload SDist # Not PR to save artifact storage, as SDist is only needed for releases. diff --git a/.github/workflows/sunmmio-ci.yml b/.github/workflows/sunmmio-ci.yml index a60c4ec707..e9f418d776 100644 --- a/.github/workflows/sunmmio-ci.yml +++ b/.github/workflows/sunmmio-ci.yml @@ -185,6 +185,7 @@ jobs: git -c protocol.version=2 submodule update --init --force --recursive \ 3rdparty/composable_kernel \ 3rdparty/cutlass \ + 3rdparty/highs \ 3rdparty/tvm - name: Resolve pinned NPU-IR commit @@ -341,6 +342,7 @@ jobs: flock 9 cmake -S . -B cmake-build --fresh ${CLANG_TIDY_CMAKE_OPTIONS} ${TILELANG_SUNMMIO_CMAKE_ARGS} if [[ "${TILELANG_SUNMMIO_CMAKE_ARGS}" == *"-DUSE_SUNMMIO=ON"* ]]; then + cmake --build cmake-build --target tilelang_highs -j "$(nproc)" cmake --build cmake-build --target mlir-headers llvm-headers -j "$(nproc)" fi flock -u 9 diff --git a/.gitmodules b/.gitmodules index b886e412c7..b13f8d1f0f 100644 --- a/.gitmodules +++ b/.gitmodules @@ -10,3 +10,6 @@ [submodule "3rdparty/NPU-IR"] path = 3rdparty/NPU-IR url = git@github.com:SUNMMIO/NPU-IR.git +[submodule "3rdparty/highs"] + path = 3rdparty/highs + url = https://github.com/ERGO-Code/HiGHS.git diff --git a/3rdparty/highs b/3rdparty/highs new file mode 160000 index 0000000000..dcc25308d8 --- /dev/null +++ b/3rdparty/highs @@ -0,0 +1 @@ +Subproject commit dcc25308d890237531422fc4d0e2e8a4740c8608 diff --git a/CMakeLists.txt b/CMakeLists.txt index c998c8b32e..e2cb2cb111 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -8,6 +8,12 @@ set(CMAKE_CXX_STANDARD 17) set(CMAKE_POSITION_INDEPENDENT_CODE ON) set(CMAKE_EXPORT_COMPILE_COMMANDS ON) +include(ExternalProject) + +if(NOT CMAKE_BUILD_TYPE) + set(CMAKE_BUILD_TYPE Release) +endif() + if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND "$ENV{CIBUILDWHEEL}") # Warning came from tvm submodule string(APPEND CMAKE_CXX_FLAGS " -Wno-dangling-reference") @@ -297,7 +303,59 @@ file(GLOB TILE_LANG_SRCS src/target/intrin_rule*.cc ) +# The ILP passes and their solver are part of the Sunmmio backend. Keep them out +# of backend-disabled builds so a regular TileLang build does not require HiGHS. +set(TILELANG_SUNMMIO_ILP_SRCS + "${CMAKE_CURRENT_SOURCE_DIR}/src/transform/inject_sunmmio_pipeline_ilp.cc" + "${CMAKE_CURRENT_SOURCE_DIR}/src/transform/sunmmio_pipeline_planning_ilp.cc" +) +list(REMOVE_ITEM TILE_LANG_SRCS ${TILELANG_SUNMMIO_ILP_SRCS}) + if(USE_SUNMMIO) + list(APPEND TILE_LANG_SRCS ${TILELANG_SUNMMIO_ILP_SRCS}) + + set(TILELANG_HIGHS_INSTALL_DIR + "${CMAKE_CURRENT_BINARY_DIR}/3rdparty/highs-install" + CACHE PATH "Install prefix for the vendored HiGHS build") + set(TILELANG_HIGHS_SOURCE_DIR + "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/highs" + CACHE PATH "Source directory for the vendored HiGHS submodule") + set(TILELANG_HIGHS_INCLUDE_DIR + "${TILELANG_HIGHS_INSTALL_DIR}/include") + set(TILELANG_HIGHS_LIBRARY + "${TILELANG_HIGHS_INSTALL_DIR}/lib/${CMAKE_STATIC_LIBRARY_PREFIX}highs${CMAKE_STATIC_LIBRARY_SUFFIX}") + + if(NOT EXISTS "${TILELANG_HIGHS_SOURCE_DIR}/CMakeLists.txt") + message(FATAL_ERROR + "HiGHS source is missing. Please run `git submodule update --init --recursive` " + "to fetch 3rdparty/highs.") + endif() + + # HiGHS is an implementation detail of libtilelang. Static linking keeps the + # installed package self-contained and avoids exposing a HiGHS SONAME/ABI. + ExternalProject_Add(tilelang_highs + SOURCE_DIR "${TILELANG_HIGHS_SOURCE_DIR}" + CMAKE_ARGS + -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE} + -DCMAKE_INSTALL_PREFIX=${TILELANG_HIGHS_INSTALL_DIR} + -DCMAKE_INSTALL_LIBDIR=lib + -DBUILD_SHARED_LIBS=OFF + -DBUILD_CXX_EXE=OFF + -DBUILD_EXAMPLES=OFF + -DBUILD_TESTING=OFF + -DZLIB=OFF + BUILD_BYPRODUCTS + "${TILELANG_HIGHS_LIBRARY}" + UPDATE_DISCONNECTED TRUE + INSTALL_DIR "${TILELANG_HIGHS_INSTALL_DIR}" + ) + + if(NOT EXISTS "${TILELANG_HIGHS_INCLUDE_DIR}/highs/Highs.h") + message(STATUS "TileLang: HiGHS headers will be installed to ${TILELANG_HIGHS_INSTALL_DIR}") + endif() + list(APPEND TILE_LANG_INCLUDES ${TILELANG_HIGHS_INCLUDE_DIR}) + list(APPEND TILE_LANG_INCLUDES ${TILELANG_HIGHS_INCLUDE_DIR}/highs) + file(GLOB TILE_LANG_SUNMMIO_SRCS src/target/sunmmio/codegen_sunmmio.cc src/target/sunmmio/sunmmio_codegen_tiles_loop.cc @@ -497,6 +555,9 @@ endif() # vendored TVM GTest discovery here so `ctest` does not pick up the excluded # upstream `cpptest` target as a stale *_NOT_BUILT placeholder. if(TILELANG_BUILD_CPP_TESTS) + # TVM's config.cmake defines USE_GTEST as a normal variable, which shadows + # the cache entry under modern CMake policy behavior. + set(USE_GTEST OFF) set(USE_GTEST OFF CACHE STRING "Disable vendored TVM GTest integration in TileLang builds" FORCE) endif() add_subdirectory(${TVM_SOURCE} tvm EXCLUDE_FROM_ALL) @@ -524,9 +585,15 @@ endif() if(TARGET llvm-headers) add_dependencies(tilelang_objs llvm-headers) endif() +if(USE_SUNMMIO) + add_dependencies(tilelang_objs tilelang_highs) +endif() add_library(tilelang SHARED $) target_link_libraries(tilelang PUBLIC tvm) +if(USE_SUNMMIO) + target_link_libraries(tilelang PRIVATE ${TILELANG_HIGHS_LIBRARY}) +endif() set(TILELANG_NPUIR_TOOL_TARGETS) if(USE_SUNMMIO) diff --git a/docs/sunmmio/pipeline_cost_model_calibration.md b/docs/sunmmio/pipeline_cost_model_calibration.md new file mode 100644 index 0000000000..12f570bb1e --- /dev/null +++ b/docs/sunmmio/pipeline_cost_model_calibration.md @@ -0,0 +1,101 @@ +# SunMMIO Pipeline Cost-Model Calibration Notes + +## Status + +An experimental gem5-aligned cost model was evaluated for the SunMMIO +pipeline planner and later removed from the default implementation. Its timing +constants were derived from a small set of isolated traces and were not broad +enough to represent different shapes, data types, memory paths, simulator +configurations, or future hardware revisions. + +The production planner therefore uses the general heuristic cost model. The +calibration below is retained as reference data for a future target-specific, +configurable cost-model profile. It must not be treated as a universal hardware +contract. + +## Experimental Model + +All delays below are integral planner cycles. + +### TensorCore + +For an MMA with dimensions `(M, N, K)`, the experimental estimate was: + +```text +work_blocks = ceil(M / 32) * ceil(N / 32) * ceil(K / 32) +delay = 37 + work_blocks +``` + +The constants were fitted to isolated `tcStart -> Calc done` observations: 38 +cycles for `32x32x32` and 41 cycles for `32x32x128`. The configured warm-up +before `tcStart` was excluded. + +### ODMA + +A transfer used the following formula: + +```text +delay = first_access_latency + ceil(bytes / request_bytes) + + completion_latency +``` + +| Source and destination path | First access | Request bytes | Completion | +| --- | ---: | ---: | ---: | +| DRAM to RSRAM | 67 | 1024 | 2 | +| DRAM to another memory | 67 | 1024 | 2 | +| RSRAM to WSRAM or ASRAM | 10 | 1024 | 2 | +| Any path involving TCM | 12 | 16 | 2 | +| Other paths | 7 | 1024 | 2 | + +The DRAM-to-RSRAM choice used `dmaPreProc -> dmaEventDone` observations of 72 +and 74 cycles for 2 KiB and 4 KiB transfers. Broadcast used a separate estimate: + +```text +broadcast_delay = 52 + ceil(bytes / 512) +``` + +### VectorCore + +The experimental analyzer counted expression-tree operations and charged the +following latency per 4096-bit vector chunk: + +| Operation | Cycles | +| --- | ---: | +| Buffer load | 14 | +| Buffer store | 3 | +| Add, subtract, multiply | 4 | +| Min, max, comparison | 3 | +| Cast | 3 | +| `exp2` | 11 | +| Bitwise AND | 3 | +| In-tile sum reduction | 14 | +| Other supported in-tile reduction | 2 | + +The load value came from isolated `vle16` samples `[13, 14, 14, 14, 13]`. +The store value rounded isolated samples `[3, 3, 3, 4, 4]` down to 3. + +## Why It Was Not Kept as the Default + +- The constants describe one simulator configuration rather than a target + capability or a stable architectural rule. +- The TensorCore fit covers only two closely related BF16 MMA observations. +- Memory-path classification by buffer scope does not capture topology, + contention, alignment, burst structure, or concurrent traffic. +- The VectorCore analyzer adds expression latencies serially and cannot model + instruction overlap, issue width, reuse, or unsupported expressions. +- Constant loop extents and regions were required, rejecting valid dynamic TIR. +- Independent rounding and planner time scaling can amplify calibration error + and change the selected initiation interval. + +## Future Profile Requirements + +A calibrated model can be reintroduced as an explicit SunMMIO target profile +when it provides: + +1. Versioned parameters tied to a named hardware or simulator configuration. +2. Coverage across supported shapes, data types, alignments, and memory paths. +3. A fallback to the general heuristic model for unsupported or dynamic TIR. +4. Validation tests comparing predicted ordering and latency with measured + traces over a representative kernel suite. +5. Configuration outside planner source code so calibration updates do not + change scheduling logic. diff --git a/docs/sunmmio/sunmmio_tilelang_user_guide.md b/docs/sunmmio/sunmmio_tilelang_user_guide.md index 6a57eb38e4..e3a0e137be 100644 --- a/docs/sunmmio/sunmmio_tilelang_user_guide.md +++ b/docs/sunmmio/sunmmio_tilelang_user_guide.md @@ -532,7 +532,7 @@ T.MeshTensor(shape, placement, (nrows, ncols), dtype="float32", layout=None) - `dtype`: element type, such as `"float16"`, `"bfloat16"`, or `"float32"`. - `layout`: global data layout in DRAM. When omitted, a rank-1 tensor with a regular dtype defaults to a 1024-byte-aligned row-major layout, while a rank >= 2 tensor defaults to ZZ. A rank >= 2 tensor with an MX dtype defaults to MXZZ; rank-1 MX tensors are unsupported. Users usually do not need to pass this parameter manually. -After entering the kernel, a `MeshTensor` parameter corresponds to the local shard visible to the current core. Use `A.global_shape` for the complete logical shape, `A.local_shape` for the uniformly allocated local slot shape, and `A.get_local_extent(cid)` for a core's valid extent. +After entering the kernel, a `MeshTensor` parameter corresponds to the local shard visible to the current core. Use `A.global_shape` for the complete logical shape, `A.local_shape` for the uniformly allocated local slot shape, and `A.get_local_extent()` for the current core's valid extent. **`make_row_major`** diff --git a/docs/sunmmio/sunmmio_tilelang_user_guide_zh_cn.md b/docs/sunmmio/sunmmio_tilelang_user_guide_zh_cn.md index 781d59d3fe..9454096b7d 100644 --- a/docs/sunmmio/sunmmio_tilelang_user_guide_zh_cn.md +++ b/docs/sunmmio/sunmmio_tilelang_user_guide_zh_cn.md @@ -532,7 +532,7 @@ T.MeshTensor(shape, placement, (nrows, ncols), dtype="float32", layout=None) - `dtype`：元素类型，例如 `"float16"`、`"bfloat16"`、`"float32"`。 - `layout`：DRAM 中的全局数据布局。省略时，普通 dtype 的 rank 1 默认使用 1024-byte 对齐的 row-major，rank >= 2 默认使用 ZZ；MX dtype 的 rank >= 2 默认使用 MXZZ，rank 1 不受支持。用户通常不需要手动传入该参数。 -进入 kernel 后，`MeshTensor` 参数对应当前 core 可见的本地 shard。使用 `A.global_shape` 查询逻辑完整 shape，使用 `A.local_shape` 查询统一分配的本地 slot shape，使用 `A.get_local_extent(cid)` 查询指定 core 的有效 extent。 +进入 kernel 后，`MeshTensor` 参数对应当前 core 可见的本地 shard。使用 `A.global_shape` 查询逻辑完整 shape，使用 `A.local_shape` 查询统一分配的本地 slot shape，使用 `A.get_local_extent()` 查询当前 core 的有效 extent。 **`make_row_major`** diff --git a/examples/flash_attention/sunmmio_example_gqa_fwd_bhsd.py b/examples/flash_attention/sunmmio_example_gqa_fwd_bhsd.py index 2f2b60514e..3878b8c1b3 100644 --- a/examples/flash_attention/sunmmio_example_gqa_fwd_bhsd.py +++ b/examples/flash_attention/sunmmio_example_gqa_fwd_bhsd.py @@ -10,7 +10,7 @@ def flashattn(batch, heads, seq_len, dim, groups=1, block_M=64, block_N=64, num_ dtype = T.bfloat16 accum_dtype = T.float32 - shard_policy = T.MeshShardingPolicy(y=0, x=2) + shard_policy = T.placement.full_shard(0, 2) Q_layout = make_zz_layout(q_shape, [1, 3], (32, 32)) K_layout = make_zz_layout(kv_shape, [1, 3], (32, 32)) diff --git a/examples/gemm/example_gemm_with_mesh_tensor.py b/examples/gemm/example_gemm_with_mesh_tensor.py index cd586a8d7d..1942894fd0 100644 --- a/examples/gemm/example_gemm_with_mesh_tensor.py +++ b/examples/gemm/example_gemm_with_mesh_tensor.py @@ -8,9 +8,9 @@ def matmul(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="flo @T.prim_func def gemm( - A: T.MeshTensor((M, K), T.MeshShardingPolicy(x=1, y=0), mesh_device_config, dtype), - B: T.MeshTensor((K, N), T.MeshShardingPolicy(x=1, y=0), mesh_device_config, dtype), - C: T.MeshTensor((M, N), T.MeshShardingPolicy(x=1, y=0), mesh_device_config, dtype), + A: T.MeshTensor((M, K), T.placement.full_shard(0, 1), mesh_device_config, dtype), + B: T.MeshTensor((K, N), T.placement.full_shard(0, 1), mesh_device_config, dtype), + C: T.MeshTensor((M, N), T.placement.full_shard(0, 1), mesh_device_config, dtype), ): sharded_M, sharded_K = A.shape _, sharded_N = B.shape diff --git a/examples/gemm/sunmmio_example_gemm.py b/examples/gemm/sunmmio_example_gemm.py index 38fe305758..90d5f9a7f6 100644 --- a/examples/gemm/sunmmio_example_gemm.py +++ b/examples/gemm/sunmmio_example_gemm.py @@ -9,16 +9,16 @@ def matmul_persistent(M, N, K, block_M, block_N, block_K, num_stages, dtype=T.bf @T.prim_func def main( - A: T.MeshTensor((M, K), T.MeshShardingPolicy(y=0, x=1), dtype, layout=A_layout), - B: T.MeshTensor((K, N), T.MeshShardingPolicy(y=0, x=1), dtype, layout=B_layout), - C: T.MeshTensor((M, N), T.MeshShardingPolicy(y=0, x=1), accum_dtype, layout=C_layout), + A: T.MeshTensor((M, K), T.placement.full_shard(0, 1), dtype, layout=A_layout), + B: T.MeshTensor((K, N), T.placement.full_shard(0, 1), dtype, layout=B_layout), + C: T.MeshTensor((M, N), T.placement.full_shard(0, 1), accum_dtype, layout=C_layout), ): with T.Kernel() as (_cid): sharded_M, sharded_K = A.local_shape _, sharded_N = B.local_shape - A_shared_dist = T.alloc_shared((block_M, block_K * T.mesh_ncols()), dtype) - B_shared_dist = T.alloc_shared((block_K * T.mesh_nrows(), block_N), dtype) + A_shared_dist = T.alloc_shared((block_M, block_K * T.ncols()), dtype) + B_shared_dist = T.alloc_shared((block_K * T.nrows(), block_N), dtype) C_shared = T.alloc_shared((block_M, block_N), accum_dtype) # Each core iterates its own sharded tile grid with plain nested diff --git a/examples/sunmmio/deepseek_mla/mla_decode.py b/examples/sunmmio/deepseek_mla/mla_decode.py index 8bf922358b..878b4d24b7 100644 --- a/examples/sunmmio/deepseek_mla/mla_decode.py +++ b/examples/sunmmio/deepseek_mla/mla_decode.py @@ -30,20 +30,16 @@ import tilelang import tilelang.language as T from tilelang import tvm as tvm -from tilelang.carver.arch import driver from tilelang.engine.phase import LowerAndLegalize from tilelang.utils.target import determine_target from tilelang.layout import make_zz_layout def mla_decode(batch, heads, kv_heads, seqlen_kv, dim, pe_dim, block_N=64, block_H=16) -> "Callable": - mesh = driver.get_sunmmio_device_mesh_config() - _, ncols = mesh + ncols = T.ncols() assert kv_heads == 1, "MLA shares a single latent KV across all query heads" - assert heads % ncols == 0, "heads must be divisible by the mesh column count" assert heads % block_H == 0, "heads must be divisible by block_H" - assert (heads // ncols) % block_H == 0, "heads/ncols must be divisible by block_H" scale = (1.0 / (dim + pe_dim)) ** 0.5 * 1.44269504 # log2(e) shape_q = [batch, heads, dim] @@ -59,8 +55,8 @@ def mla_decode(batch, heads, kv_heads, seqlen_kv, dim, pe_dim, block_N=64, block # Q, Q_pe and Output all share the same policy: batch on the rows (y=0), the # head axis sharded across the columns (x=1). KV/K_pe split the sequence # axis across the columns (x=1) -> split-K. - head_policy = T.MeshShardingPolicy(y=0, x=1) - kv_policy = T.MeshShardingPolicy(y=0, x=1) + head_policy = T.placement.full_shard(0, 1) + kv_policy = T.placement.full_shard(0, 1) @T.prim_func def main( @@ -75,7 +71,7 @@ def main( _, sharded_seqlen, _, _ = KV.local_shape # This core's column owns the global head blocks # [col*blocks_per_col, (col+1)*blocks_per_col). - col = cid % T.mesh_ncols() + col = cid % T.ncols() blocks_per_col = heads_per_col // block_H # Head gather (per batch): load this core's head slice, all-gather @@ -105,11 +101,11 @@ def main( lse = T.alloc_shared([block_H], accum_dtype) # Cross-row LSE-combine scratch. - lse_dist = T.alloc_shared([T.mesh_ncols(), block_H], accum_dtype) + lse_dist = T.alloc_shared([T.ncols(), block_H], accum_dtype) lse_max = T.alloc_shared([block_H], accum_dtype) lse_denom = T.alloc_shared([block_H], accum_dtype) o_scaled = T.alloc_shared([block_H, dim], accum_dtype) - o_dist = T.alloc_shared([T.mesh_ncols(), block_H, dim], accum_dtype) + o_dist = T.alloc_shared([T.ncols(), block_H, dim], accum_dtype) o_final = T.alloc_shared([block_H, dim], accum_dtype) o_cast = T.alloc_shared([block_H, dim], dtype) @@ -163,7 +159,7 @@ def main( # --- LSE combine across the row (the seqlen-split axis). --- T.comm.all_gather(lse, lse_dist, direction="h") T.reduce_max(lse_dist, lse_max, dim=0, clear=True) - for c, i in T.Tiles([T.mesh_ncols(), block_H]): + for c, i in T.Tiles([T.ncols(), block_H]): lse_dist[c, i] = T.exp2(lse_dist[c, i] - lse_max[i]) T.reduce_sum(lse_dist, lse_denom, dim=0, clear=True) for i, j in T.Tiles([block_H, dim]): diff --git a/examples/sunmmio/elementwise/elementwise_add_dynamic.py b/examples/sunmmio/elementwise/elementwise_add_dynamic.py index e3e4285f39..cd11dc9c7d 100644 --- a/examples/sunmmio/elementwise/elementwise_add_dynamic.py +++ b/examples/sunmmio/elementwise/elementwise_add_dynamic.py @@ -15,7 +15,7 @@ def _elementwise_add_prim_func(block_M, block_N, in_dtype, out_dtype): N = T.dynamic("n") zz_layout = make_zz_layout((M, N)) - placement = T.MeshShardingPolicy(y=0, x=1) + placement = T.placement.full_shard(0, 1) @T.prim_func def elem_add( diff --git a/examples/sunmmio/elementwise/elementwise_exp2_dynamic.py b/examples/sunmmio/elementwise/elementwise_exp2_dynamic.py index 753fb2bda6..9473a82fe3 100644 --- a/examples/sunmmio/elementwise/elementwise_exp2_dynamic.py +++ b/examples/sunmmio/elementwise/elementwise_exp2_dynamic.py @@ -17,7 +17,7 @@ def _elementwise_exp2_prim_func(block_M, block_N, in_dtype, out_dtype): N = T.dynamic("n") zz_layout = make_zz_layout((M, N)) - placement = T.MeshShardingPolicy(y=0, x=1) + placement = T.placement.full_shard(0, 1) @T.prim_func def elem_exp2( diff --git a/examples/sunmmio/flash_decoding/example_gqa_decode.py b/examples/sunmmio/flash_decoding/example_gqa_decode.py index 7aa2fccfa4..20a1655851 100644 --- a/examples/sunmmio/flash_decoding/example_gqa_decode.py +++ b/examples/sunmmio/flash_decoding/example_gqa_decode.py @@ -25,16 +25,16 @@ def gqa_flashattn(batch, heads, kv_heads, seqlen_kv, dim, block_N=128): @T.prim_func def main( - Q: T.MeshTensor(shape_q, T.MeshShardingPolicy(y=0, x=1), dtype, layout=make_zz_layout(shape_q)), - K: T.MeshTensor(shape_k, T.MeshShardingPolicy(y=0, x=2), dtype, layout=make_zz_layout(shape_k, axes=(1, 3))), - V: T.MeshTensor(shape_v, T.MeshShardingPolicy(y=0, x=2), dtype, layout=make_zz_layout(shape_k, axes=(1, 3))), + Q: T.MeshTensor(shape_q, T.placement.full_shard(0, 1), dtype, layout=make_zz_layout(shape_q)), + K: T.MeshTensor(shape_k, T.placement.full_shard(0, 2), dtype, layout=make_zz_layout(shape_k, axes=(1, 3))), + V: T.MeshTensor(shape_v, T.placement.full_shard(0, 2), dtype, layout=make_zz_layout(shape_k, axes=(1, 3))), mask: T.MeshTensor( [batch, seqlen_kv], - T.MeshShardingPolicy(y=0, replicate=T.MeshReplicationType.ROW), + T.placement.row_shard(0), "uint16", layout=make_row_major([batch, seqlen_kv]), ), - Output: T.MeshTensor(shape_o, T.MeshShardingPolicy(y=0, x=1), dtype, layout=make_zz_layout(shape_o)), + Output: T.MeshTensor(shape_o, T.placement.full_shard(0, 1), dtype, layout=make_zz_layout(shape_o)), ): with T.Kernel() as (_cid): sharded_batch, sharded_heads, _ = Q.local_shape diff --git a/examples/sunmmio/gemm/bf16_gemm.py b/examples/sunmmio/gemm/bf16_gemm.py index c6f34691b0..0aeeb11992 100644 --- a/examples/sunmmio/gemm/bf16_gemm.py +++ b/examples/sunmmio/gemm/bf16_gemm.py @@ -17,16 +17,16 @@ def matmul_persistent(M, N, K, block_M, block_N, block_K, dtype=T.bfloat16, accu @T.prim_func def main( - A: T.MeshTensor((M, K), T.MeshShardingPolicy(y=0, x=1), dtype, layout=A_layout), - B: T.MeshTensor((K, N), T.MeshShardingPolicy(y=0, x=1), dtype, layout=B_layout), - C: T.MeshTensor((M, N), T.MeshShardingPolicy(y=0, x=1), accum_dtype, layout=C_layout), + A: T.MeshTensor((M, K), T.placement.full_shard(0, 1), dtype, layout=A_layout), + B: T.MeshTensor((K, N), T.placement.full_shard(0, 1), dtype, layout=B_layout), + C: T.MeshTensor((M, N), T.placement.full_shard(0, 1), accum_dtype, layout=C_layout), ): with T.Kernel() as (_cid): sharded_M, sharded_K = A.local_shape _, sharded_N = B.local_shape - A_shared_dist = T.alloc_shared((block_M, block_K * T.mesh_ncols()), dtype) - B_shared_dist = T.alloc_shared((block_K * T.mesh_nrows(), block_N), dtype) + A_shared_dist = T.alloc_shared((block_M, block_K * T.ncols()), dtype) + B_shared_dist = T.alloc_shared((block_K * T.nrows(), block_N), dtype) C_shared = T.alloc_shared((block_M, block_N), accum_dtype) for bx in T.serial(T.ceildiv(sharded_M, block_M)): diff --git a/examples/sunmmio/gemm/bf16_gemm_zzn.py b/examples/sunmmio/gemm/bf16_gemm_zzn.py index e3513d8503..5f55c082c6 100644 --- a/examples/sunmmio/gemm/bf16_gemm_zzn.py +++ b/examples/sunmmio/gemm/bf16_gemm_zzn.py @@ -43,7 +43,7 @@ def matmul_persistent_zzn(M, N, K, block_M, block_N, block_K, dtype=T.bfloat16, cluster_shape=(block_K // 32, block_N // 32), ) C_layout = make_zz_layout((M, N)) - placement = T.MeshShardingPolicy(y=0, x=1) + placement = T.placement.full_shard(0, 1) @T.prim_func def main( @@ -55,8 +55,8 @@ def main( sharded_M, sharded_K = A.local_shape _, sharded_N = B.local_shape - A_shared_dist = T.alloc_shared((block_M, block_K * T.mesh_ncols()), dtype) - B_shared_dist = T.alloc_shared((block_K * T.mesh_nrows(), block_N), dtype) + A_shared_dist = T.alloc_shared((block_M, block_K * T.ncols()), dtype) + B_shared_dist = T.alloc_shared((block_K * T.nrows(), block_N), dtype) C_shared = T.alloc_shared((block_M, block_N), accum_dtype) for bx in T.serial(T.ceildiv(sharded_M, block_M)): diff --git a/examples/sunmmio/gemm/dynamic_b16_gemm.py b/examples/sunmmio/gemm/dynamic_b16_gemm.py index 802fb02278..e878a1268c 100644 --- a/examples/sunmmio/gemm/dynamic_b16_gemm.py +++ b/examples/sunmmio/gemm/dynamic_b16_gemm.py @@ -17,7 +17,7 @@ def matmul_persistent_dynamic(K, block_M, block_N, block_K, dtype=T.bfloat16, ac A_layout = make_zz_layout((M, K)) B_layout = make_zz_layout((K, N)) C_layout = make_zz_layout((M, N)) - placement = T.MeshShardingPolicy(y=0, x=1) + placement = T.placement.full_shard(0, 1) @T.prim_func def main( @@ -29,8 +29,8 @@ def main( sharded_M, sharded_K = A.local_shape _, sharded_N = B.local_shape - A_shared_dist = T.alloc_shared((block_M, block_K * T.mesh_ncols()), dtype) - B_shared_dist = T.alloc_shared((block_K * T.mesh_nrows(), block_N), dtype) + A_shared_dist = T.alloc_shared((block_M, block_K * T.ncols()), dtype) + B_shared_dist = T.alloc_shared((block_K * T.nrows(), block_N), dtype) C_shared = T.alloc_shared((block_M, block_N), accum_dtype) for bx in T.serial(T.ceildiv(sharded_M, block_M)): diff --git a/examples/sunmmio/reduction/1d_reduction.py b/examples/sunmmio/reduction/1d_reduction.py index 34a45edf5c..f5d7b72f56 100644 --- a/examples/sunmmio/reduction/1d_reduction.py +++ b/examples/sunmmio/reduction/1d_reduction.py @@ -11,8 +11,8 @@ def reduction(M, block_M, in_dtype, out_dtype): A_layout = make_row_major((M,)) B_layout = make_aligned_row_major((1,), align_bytes=1024, dtype=out_dtype) - A_placement = T.MeshShardingPolicy(y=0, replicate=T.MeshReplicationType.ROW) - B_placement = T.MeshShardingPolicy(replicate=T.MeshReplicationType.ALL) + A_placement = T.placement.row_shard(0) + B_placement = T.placement.replicated() @T.prim_func def main( @@ -24,7 +24,7 @@ def main( A_shared = T.alloc_shared((block_M), in_dtype) Acc_shared = T.alloc_shared((block_M), out_dtype) - Acc_dist_shared = T.alloc_shared((T.mesh_ncols() * block_M), out_dtype) + Acc_dist_shared = T.alloc_shared((T.ncols() * block_M), out_dtype) B_shared = T.alloc_shared((1,), out_dtype) T.annotate_layout({B_shared: B_layout}) diff --git a/examples/sunmmio/reduction/reduction.py b/examples/sunmmio/reduction/reduction.py index 8a02a5944a..4c8fa3d175 100644 --- a/examples/sunmmio/reduction/reduction.py +++ b/examples/sunmmio/reduction/reduction.py @@ -10,7 +10,7 @@ def reduction(M, K, N, block_K, block_N, in_dtype, out_dtype): zz_layout = make_zz_layout((M, K, N)) - placement = T.MeshShardingPolicy(y=0, x=1) + placement = T.placement.full_shard(0, 1) rm_layout = make_row_major((M, K)) @T.prim_func @@ -24,7 +24,7 @@ def main( A_shared = T.alloc_shared((block_K, block_N), in_dtype) Acc_shared = T.alloc_shared((block_K, block_N), out_dtype) - Acc_dist_shared = T.alloc_shared((block_K, T.mesh_ncols() * block_N), out_dtype) + Acc_dist_shared = T.alloc_shared((block_K, T.ncols() * block_N), out_dtype) B_shared = T.alloc_shared((block_K,), out_dtype) for bx in T.serial(sharded_M): diff --git a/examples/sunmmio/rmsnorm/rmsnorm.py b/examples/sunmmio/rmsnorm/rmsnorm.py index 49d5678068..bcdf927e23 100644 --- a/examples/sunmmio/rmsnorm/rmsnorm.py +++ b/examples/sunmmio/rmsnorm/rmsnorm.py @@ -15,7 +15,7 @@ def rmsnorm_kernel(M, N, block_M, block_N, dtype: T.dtype = T.bfloat16, eps: flo # across the mesh columns, mirroring the softmax example. The RMSNorm # reduction is over N, which lives on the column axis, so each core holds a # partial sum that is combined across the row with an all_gather. - placement = T.MeshShardingPolicy(y=0, x=1) + placement = T.placement.full_shard(0, 1) accum_dtype = T.float32 @@ -33,7 +33,7 @@ def main( tile_sumsq = T.alloc_shared((block_M,), accum_dtype) local_sumsq = T.alloc_shared((block_M,), accum_dtype) - sumsq_dist = T.alloc_shared((T.mesh_ncols(), block_M), accum_dtype) + sumsq_dist = T.alloc_shared((T.ncols(), block_M), accum_dtype) total_sumsq = T.alloc_shared((block_M,), accum_dtype) inv_rms = T.alloc_shared((block_M,), accum_dtype) diff --git a/examples/sunmmio/rmsnorm/rmsnorm_pipelined.py b/examples/sunmmio/rmsnorm/rmsnorm_pipelined.py index b913e60c92..03d475434e 100644 --- a/examples/sunmmio/rmsnorm/rmsnorm_pipelined.py +++ b/examples/sunmmio/rmsnorm/rmsnorm_pipelined.py @@ -29,7 +29,7 @@ def rmsnorm_kernel(M, N, block_M, block_N, dtype: T.dtype = T.bfloat16, eps: float = 1e-12) -> "Callable": zz_layout = make_zz_layout((M, N)) - placement = T.MeshShardingPolicy(y=0, x=1) + placement = T.placement.full_shard(0, 1) accum_dtype = T.float32 @@ -50,7 +50,7 @@ def main( x_sq = T.alloc_shared((block_M, block_N), accum_dtype, scope="shared.rsram") tile_sumsq = T.alloc_shared((block_M,), accum_dtype, scope="shared.rsram") local_sumsq = T.alloc_shared((block_M,), accum_dtype, scope="shared.rsram") - sumsq_dist = T.alloc_shared((T.mesh_ncols(), block_M), accum_dtype, scope="shared.rsram") + sumsq_dist = T.alloc_shared((T.ncols(), block_M), accum_dtype, scope="shared.rsram") total_sumsq = T.alloc_shared((block_M,), accum_dtype, scope="shared.rsram") inv_rms = T.alloc_shared((block_M,), accum_dtype, scope="shared.rsram") diff --git a/examples/sunmmio/softmax/dynamic_online_softmax.py b/examples/sunmmio/softmax/dynamic_online_softmax.py index e51aec688b..4de86f3d11 100644 --- a/examples/sunmmio/softmax/dynamic_online_softmax.py +++ b/examples/sunmmio/softmax/dynamic_online_softmax.py @@ -19,7 +19,7 @@ def dynamic_softmax_kernel(block_M, block_N, dtype: T.dtype = T.bfloat16) -> "Ca M, N = T.dynamic("m"), T.dynamic("n") zz_layout = make_zz_layout((M, N)) - placement = T.MeshShardingPolicy(y=0, x=1) + placement = T.placement.full_shard(0, 1) accum_dtype = T.float32 scale = 1.44269504 # log2(e) @@ -36,7 +36,7 @@ def main(X: T.MeshTensor((M, N), placement, dtype, layout=zz_layout), Y: T.MeshT exp_x = T.alloc_shared([block_M, block_N], accum_dtype) sum_exp_x = T.alloc_shared((block_M), accum_dtype) - lse_dist = T.alloc_shared((T.mesh_ncols(), block_M), accum_dtype) + lse_dist = T.alloc_shared((T.ncols(), block_M), accum_dtype) lse_max = T.alloc_shared((block_M,), accum_dtype) lse_global = T.alloc_shared((block_M,), accum_dtype) @@ -59,7 +59,7 @@ def main(X: T.MeshTensor((M, N), placement, dtype, layout=zz_layout), Y: T.MeshT T.reduce_max(lse_dist, lse_max, dim=0, clear=True) # Get global lse - for i, j in T.Tiles([T.mesh_ncols(), block_M]): + for i, j in T.Tiles([T.ncols(), block_M]): lse_dist[i, j] = T.exp2(lse_dist[i, j] - lse_max[j]) T.reduce_sum(lse_dist, lse_global, dim=0, clear=True) for i in T.Tiles([block_M]): diff --git a/examples/sunmmio/softmax/online_softmax.py b/examples/sunmmio/softmax/online_softmax.py index 308731d584..fac4dc533d 100644 --- a/examples/sunmmio/softmax/online_softmax.py +++ b/examples/sunmmio/softmax/online_softmax.py @@ -17,7 +17,7 @@ def ref_program(x): def softmax_kernel(M, N, block_M, block_N, dtype: T.dtype = T.bfloat16) -> "Callable": zz_layout = make_zz_layout((M, N)) - placement = T.MeshShardingPolicy(y=0, x=1) + placement = T.placement.full_shard(0, 1) accum_dtype = T.float32 scale = 1.44269504 # log2(e) @@ -34,7 +34,7 @@ def main(X: T.MeshTensor((M, N), placement, dtype, layout=zz_layout), Y: T.MeshT exp_x = T.alloc_shared([block_M, block_N], accum_dtype) sum_exp_x = T.alloc_shared((block_M), accum_dtype) - lse_dist = T.alloc_shared((T.mesh_ncols(), block_M), accum_dtype) + lse_dist = T.alloc_shared((T.ncols(), block_M), accum_dtype) lse_max = T.alloc_shared((block_M,), accum_dtype) lse_global = T.alloc_shared((block_M,), accum_dtype) @@ -57,7 +57,7 @@ def main(X: T.MeshTensor((M, N), placement, dtype, layout=zz_layout), Y: T.MeshT T.reduce_max(lse_dist, lse_max, dim=0, clear=True) # Get global lse - for i, j in T.Tiles([T.mesh_ncols(), block_M]): + for i, j in T.Tiles([T.ncols(), block_M]): lse_dist[i, j] = T.exp2(lse_dist[i, j] - lse_max[j]) T.reduce_sum(lse_dist, lse_global, dim=0, clear=True) for i in T.Tiles([block_M]): diff --git a/examples/sunmmio/softmax/online_softmax_sudeck.py b/examples/sunmmio/softmax/online_softmax_sudeck.py index 4522a5564b..fd4d5722f0 100644 --- a/examples/sunmmio/softmax/online_softmax_sudeck.py +++ b/examples/sunmmio/softmax/online_softmax_sudeck.py @@ -29,7 +29,7 @@ def main(M, N) -> None: kernel = online_softmax_sudeck(M, N, block_M=256, block_N=256, dtype=T.bfloat16) - # Match the kernel's MeshTensor(placement=MeshShardingPolicy(y=0, x=1), layout=zz): + # Match the kernel's MeshTensor(placement=T.placement.full_shard(0, 1), layout=zz): # zz block layout + full 2D shard across the mesh. with sm.spec(layout=sm.layout.zz(0, 1), placement=sm.placement.full_shard(0, 1)): x_dev = x.to("sunmmio") diff --git a/pyproject.toml b/pyproject.toml index 9f2dbfa9e2..11acd3ba49 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -84,7 +84,7 @@ requires = [ build-backend = "scikit_build_core.build" [tool.scikit-build] -wheel.py-api = "cp38" +wheel.py-api = "cp39" cmake.version = ">=3.26.1" build-dir = "build" @@ -119,6 +119,8 @@ include = [ # Composable Kernel "3rdparty/composable_kernel/include", "3rdparty/composable_kernel/library", + # HiGHS is required to build the Sunmmio ILP pipeline planner. + "3rdparty/highs", "testing/**", "examples/**", ] diff --git a/src/op/builtin.cc b/src/op/builtin.cc index ed988651af..6b9d943de2 100644 --- a/src/op/builtin.cc +++ b/src/op/builtin.cc @@ -21,6 +21,7 @@ TVM_REGISTER_PASS_CONFIG_OPTION(kDisableTMALower, Bool); TVM_REGISTER_PASS_CONFIG_OPTION(kDisableSafeMemoryLegalize, Bool); TVM_REGISTER_PASS_CONFIG_OPTION(kDisableWarpSpecialized, Bool); TVM_REGISTER_PASS_CONFIG_OPTION(kDisableThreadStorageSync, Bool); +TVM_REGISTER_PASS_CONFIG_OPTION(kDisableSunmmioPipeline, Bool); TVM_REGISTER_PASS_CONFIG_OPTION(kConfigIndexBitwidth, Integer); TVM_REGISTER_PASS_CONFIG_OPTION(kEnableAggressiveSharedMemoryMerge, Bool); TVM_REGISTER_PASS_CONFIG_OPTION(kForceLetInline, Bool); @@ -36,6 +37,11 @@ TVM_REGISTER_PASS_CONFIG_OPTION(kStorageRewriteDetectInplace, Bool); TVM_REGISTER_PASS_CONFIG_OPTION(kASTPrintEnable, Bool); TVM_REGISTER_PASS_CONFIG_OPTION(kLayoutVisualizationEnable, Bool); TVM_REGISTER_PASS_CONFIG_OPTION(kLayoutVisualizationFormats, String); +TVM_REGISTER_PASS_CONFIG_OPTION(kSunmmioPipelineMode, String); +TVM_REGISTER_PASS_CONFIG_OPTION(kSunmmioFaster, Integer); +TVM_REGISTER_PASS_CONFIG_OPTION(kSunmmioILPStageShrink, Bool); +TVM_REGISTER_PASS_CONFIG_OPTION(kSunmmioILPMultiversionLifetimePruning, Bool); +TVM_REGISTER_PASS_CONFIG_OPTION(kSunmmioILPModelVCBlockingIssue, Bool); TVM_REGISTER_PASS_CONFIG_OPTION(kDeviceCompileFlags, ffi::Array); TVM_REGISTER_PASS_CONFIG_OPTION(kDisableDataRaceCheck, Bool); TVM_REGISTER_PASS_CONFIG_OPTION(kEnableLowerLDGSTG, Bool); diff --git a/src/op/builtin.h b/src/op/builtin.h index eefa462c27..9d8e1256eb 100644 --- a/src/op/builtin.h +++ b/src/op/builtin.h @@ -116,6 +116,22 @@ static constexpr const char *kDisableDataRaceCheck = static constexpr const char *kDisableThreadStorageSync = "tl.disable_thread_storage_sync"; +static constexpr const char *kDisableSunmmioPipeline = + "tl.disable_sunmmio_pipeline"; + +static constexpr const char *kSunmmioPipelineMode = "tl.sunmmio_pipeline_mode"; + +static constexpr const char *kSunmmioFaster = "tl.sunmmio_faster"; + +static constexpr const char *kSunmmioILPStageShrink = + "tl.sunmmio_ilp_stage_shrink"; + +static constexpr const char *kSunmmioILPMultiversionLifetimePruning = + "tl.sunmmio_ilp_multiversion_lifetime_pruning"; + +static constexpr const char *kSunmmioILPModelVCBlockingIssue = + "tl.sunmmio_ilp_model_vc_blocking_issue"; + /*! * \brief Force inline Let bindings during simplification. * diff --git a/src/op/copy.cc b/src/op/copy.cc index 640587e1e5..1bdf97f4e8 100644 --- a/src/op/copy.cc +++ b/src/op/copy.cc @@ -832,43 +832,244 @@ Stmt CopyNode::LowerSunmmioDramRsramCopy(const LowerArgs &T, if (!src_layout.defined() || !dst_layout.defined()) return dma(src_region, dst_region); - auto make_covered_full_rank1_range = - [&](const Buffer &buffer, const Layout &layout, - const Array &ranges) -> Optional> { - if (buffer->shape.size() != 1 || ranges.size() != 1) - return Optional>(); + constexpr int kDramRowAlignBytes = 1024; + const auto &config = GetSunmmioTileProcessorConfig(T.target); + + struct AlignedRowDmaCarrier { + bool aligned_1024{false}; + bool valid{false}; + Array physical_ranges; + Array covered_shape; + Array canonical_logical_shape; + Array canonical_carrier_shape; + std::string reason; + }; + + auto shape_of_ranges = [](const Array &ranges) { + Array shape; + for (const Range &range : ranges) + shape.push_back(range->extent); + return shape; + }; + auto format_shape = [](const Array &shape) { + std::ostringstream os; + os << "("; + for (size_t i = 0; i < shape.size(); ++i) { + if (i) + os << ", "; + os << shape[i]; + } + os << ")"; + return os.str(); + }; + auto remove_leading_singletons = [&](const Array &shape) { + Array canonical; + bool found_non_singleton = false; + for (const PrimExpr &extent : shape) { + if (!found_non_singleton && analyzer->CanProveEqual(extent, 1)) + continue; + found_non_singleton = true; + canonical.push_back(extent); + } + return canonical; + }; + auto shape_matches = [&](const Array &lhs, + const Array &rhs) { + if (lhs.size() != rhs.size()) + return false; + for (size_t i = 0; i < lhs.size(); ++i) { + if (!analyzer->CanProveEqual(lhs[i], rhs[i])) + return false; + } + return true; + }; + auto bit_count_is_aligned = [&](PrimExpr element_count, DataType dtype, + int align_bytes) { + PrimExpr bit_count = + element_count * + make_const(element_count.dtype(), dtype.bits() * dtype.lanes()); + PrimExpr align_bits = + make_const(bit_count.dtype(), static_cast(align_bytes) * 8); + return analyzer->CanProveEqual(indexmod(bit_count, align_bits), + make_zero(bit_count.dtype())); + }; + auto is_aligned_row_major = [&](const Buffer &buffer, const Layout &layout, + int align_bytes) { + if (!layout.defined() || layout->InputDim() == 0 || + sunmmio::IsMXDType(buffer->dtype)) { + return false; + } + Layout expected = sunmmio::MakeAlignedRowMajor(layout->InputShape(), + buffer->dtype, align_bytes); + if (IsSameLayout(layout, expected, analyzer)) + return true; + + const auto *actual_cute = layout.as(); + const auto *expected_cute = expected.as(); + if (!actual_cute || !expected_cute || + actual_cute->GetDimLevels().size() != + expected_cute->GetDimLevels().size()) { + return false; + } + for (size_t i = 0; i < actual_cute->GetDimLevels().size(); ++i) { + if (actual_cute->GetDimLevels()[i].IntValue() != + expected_cute->GetDimLevels()[i].IntValue()) { + return false; + } + } + + Array actual_modes = actual_cute->GetModeShape(); + Array expected_modes = expected_cute->GetModeShape(); + Array actual_strides = actual_cute->GetModeStride(); + Array expected_strides = expected_cute->GetModeStride(); + if (actual_modes.size() != expected_modes.size() || + actual_strides.size() != expected_strides.size()) { + return false; + } + for (size_t i = 0; i < actual_modes.size(); ++i) { + if (!analyzer->CanProveEqual(actual_modes[i], expected_modes[i])) + return false; + // Layout inference may choose any stride for a singleton mode; its index + // is always zero, so that stride has no effect on physical addressing. + if (!analyzer->CanProveEqual(actual_modes[i], 1) && + !analyzer->CanProveEqual(actual_strides[i], expected_strides[i])) { + return false; + } + } + return true; + }; + auto analyze_carrier = [&](const Buffer &buffer, const Layout &layout, + const Array &ranges) { + AlignedRowDmaCarrier carrier; + carrier.aligned_1024 = + is_aligned_row_major(buffer, layout, kDramRowAlignBytes); + if (!carrier.aligned_1024) + return carrier; + if (buffer->shape.empty() || ranges.size() != buffer->shape.size() || + layout->InputShape().size() != buffer->shape.size()) { + carrier.reason = "buffer, region, and layout ranks must match"; + return carrier; + } + for (size_t i = 0; i < buffer->shape.size(); ++i) { + if (!analyzer->CanProveEqual(layout->InputShape()[i], buffer->shape[i])) { + carrier.reason = "layout input shape does not match the buffer shape"; + return carrier; + } + } const auto *cute = layout.as(); - if (!cute) - return Optional>(); - Array covered_shape = cute->GetCoveredShape(); - if (covered_shape.size() != 1) - return Optional>(); - const Range &range = ranges[0]; - if (!analyzer->CanProveEqual(range->min, make_zero(range->min.dtype())) || - !analyzer->CanProveEqual(range->extent, buffer->shape[0])) { - return Optional>(); - } - if (analyzer->CanProveEqual(covered_shape[0], range->extent)) - return Optional>(); - Array covered_ranges; - covered_ranges.push_back( - Range::FromMinExtent(range->min, covered_shape[0])); - return covered_ranges; + if (!cute) { + carrier.reason = "layout is not a CuteLayout"; + return carrier; + } + carrier.covered_shape = cute->GetCoveredShape(); + if (carrier.covered_shape.size() != buffer->shape.size()) { + carrier.reason = "covered shape rank does not match the buffer rank"; + return carrier; + } + + const int innermost = static_cast(ranges.size()) - 1; + const Range &row = ranges[innermost]; + if (!analyzer->CanProveEqual(row->min, make_zero(row->min.dtype()))) { + carrier.reason = "innermost range must start at zero"; + return carrier; + } + if (!analyzer->CanProveEqual(row->extent, buffer->shape[innermost])) { + carrier.reason = "innermost range must cover the complete logical row"; + return carrier; + } + if (analyzer->CanProve(carrier.covered_shape[innermost] < + buffer->shape[innermost])) { + carrier.reason = "covered row is smaller than the logical row"; + return carrier; + } + if (!analyzer->CanProve(carrier.covered_shape[innermost] >= + buffer->shape[innermost])) { + carrier.reason = + "required equality could not be proven for symbolic extents"; + return carrier; + } + if (!bit_count_is_aligned(carrier.covered_shape[innermost], buffer->dtype, + kDramRowAlignBytes)) { + carrier.reason = "covered row byte size is not a multiple of 1024"; + return carrier; + } + + Array logical_shape = shape_of_ranges(ranges); + Array physical_shape; + for (int i = 0; i < innermost; ++i) { + carrier.physical_ranges.push_back(ranges[i]); + physical_shape.push_back(ranges[i]->extent); + } + carrier.physical_ranges.push_back( + Range::FromMinExtent(row->min, carrier.covered_shape[innermost])); + physical_shape.push_back(carrier.covered_shape[innermost]); + carrier.canonical_logical_shape = remove_leading_singletons(logical_shape); + carrier.canonical_carrier_shape = remove_leading_singletons(physical_shape); + if (carrier.canonical_logical_shape.empty() || + carrier.canonical_carrier_shape.empty()) { + carrier.reason = "effective rank must be one or two"; + return carrier; + } + if (carrier.canonical_logical_shape.size() > 2 || + carrier.canonical_carrier_shape.size() > 2) { + carrier.reason = "effective rank exceeds two"; + return carrier; + } + carrier.valid = true; + return carrier; + }; + AlignedRowDmaCarrier src_carrier = + analyze_carrier(src, src_layout, src_range); + AlignedRowDmaCarrier dst_carrier = + analyze_carrier(dst, dst_layout, dst_range); + auto reject_carrier = [&](const std::string &reason) -> Stmt { + LOG(FATAL) << "Sunmmio aligned-row DMA carrier rejected for " + << (src.scope() == "global" ? "DRAM->RSRAM" : "RSRAM->DRAM") + << " copy " << src->name << " -> " << dst->name << ":\n" + << " source logical region: " + << format_shape(shape_of_ranges(src_range)) << "\n" + << " destination logical region: " + << format_shape(shape_of_ranges(dst_range)) << "\n" + << " source canonical logical shape: " + << format_shape(src_carrier.canonical_logical_shape) << "\n" + << " destination canonical logical shape: " + << format_shape(dst_carrier.canonical_logical_shape) << "\n" + << " source carrier: " + << format_shape(src_carrier.canonical_carrier_shape) << "\n" + << " destination carrier: " + << format_shape(dst_carrier.canonical_carrier_shape) << "\n" + << " source aligned-row-major(1024): " + << src_carrier.aligned_1024 << "\n" + << " destination aligned-row-major(1024): " + << dst_carrier.aligned_1024 << "\n" + << " reason: " << reason << "\n" + << " required: matching 1024-byte aligned row-major full-row " + "regions with one or two effective dimensions"; + return Stmt(); }; - if (IsLayoutMatch(src_layout, dst_layout, analyzer)) { - Optional> covered_src_range = - make_covered_full_rank1_range(src, src_layout, src_range); - Optional> covered_dst_range = - make_covered_full_rank1_range(dst, dst_layout, dst_range); - if (covered_src_range.defined() && covered_dst_range.defined() && - analyzer->CanProveEqual(covered_src_range.value()[0]->extent, - covered_dst_range.value()[0]->extent)) { - return dma(MakeRegionExpr(src, covered_src_range.value(), - /*access_mask=*/1), - MakeRegionExpr(dst, covered_dst_range.value(), - /*access_mask=*/2)); + if (src_carrier.aligned_1024 && dst_carrier.aligned_1024) { + if (!src_carrier.valid) + return reject_carrier("source " + src_carrier.reason); + if (!dst_carrier.valid) + return reject_carrier("destination " + dst_carrier.reason); + if (src->dtype != dst->dtype) + return reject_carrier("source and destination element dtypes must match"); + if (!shape_matches(src_carrier.canonical_logical_shape, + dst_carrier.canonical_logical_shape)) { + return reject_carrier("canonical logical shapes do not match"); } + if (!shape_matches(src_carrier.canonical_carrier_shape, + dst_carrier.canonical_carrier_shape)) { + return reject_carrier("canonical carrier shapes do not match"); + } + return dma(MakeRegionExpr(src, src_carrier.physical_ranges, + /*access_mask=*/1), + MakeRegionExpr(dst, dst_carrier.physical_ranges, + /*access_mask=*/2)); + } + + if (IsLayoutMatch(src_layout, dst_layout, analyzer)) { return dma(src_region, dst_region); } @@ -962,7 +1163,6 @@ Stmt CopyNode::LowerSunmmioDramRsramCopy(const LowerArgs &T, // both sides are row-major and either has rows that are not rsram_align_bytes // aligned, the leading dimension (product of all dims but the last) must fit // the tile register height. Limitations for Sunmmio A4E - const auto &config = GetSunmmioTileProcessorConfig(T.target); bool both_row_major = config.rsram_align_bytes > 0; bool any_unaligned = false; for (const auto &[layout, dtype] : diff --git a/src/op/reduce.cc b/src/op/reduce.cc index e51afdef4d..105f0d53e7 100644 --- a/src/op/reduce.cc +++ b/src/op/reduce.cc @@ -810,6 +810,19 @@ Stmt ReduceOpNode::MakeSunmmioTileReduce(const LowerArgs &T, } body = root; + // Preserve the reduction-temporary roles structurally. LowerOpaqueBlock + // transfers this per-buffer map to the corresponding Allocate annotations. + Map reduce_temp_roles; + reduce_temp_roles.Set(acc->data, Integer(static_cast( + ReduceRegisterTempRole::kAccumulator))); + if (dst_res.defined()) { + reduce_temp_roles.Set( + dst_res.value()->data, + Integer(static_cast(ReduceRegisterTempRole::kResult))); + } + Map block_annotations; + block_annotations.Set(attr::kSunmmioReduceRegisterTemp, reduce_temp_roles); + // Finally, wrap the body in a Block so the accumulator temporaries are // allocated with the reduction scope. Array alloc_buffers; @@ -820,7 +833,7 @@ Stmt ReduceOpNode::MakeSunmmioTileReduce(const LowerArgs &T, body = BlockRealize({}, Bool(true), Block({}, {}, {}, "reduce_tile_op", body, std::nullopt, - alloc_buffers, {}, {})); + alloc_buffers, {}, block_annotations)); return body; } diff --git a/src/target/sunmmio/codegen_sunmmio.cc b/src/target/sunmmio/codegen_sunmmio.cc index 203fb52991..be2fe02a8a 100644 --- a/src/target/sunmmio/codegen_sunmmio.cc +++ b/src/target/sunmmio/codegen_sunmmio.cc @@ -43,6 +43,7 @@ namespace { class DeclBufferCollector final : public tir::StmtVisitor { public: std::unordered_map buffer_data_to_buffer; + std::unordered_map reduce_register_temp_roles; private: void VisitStmt_(const tir::DeclBufferNode *op) final { @@ -50,6 +51,24 @@ class DeclBufferCollector final : public tir::StmtVisitor { tir::StmtVisitor::VisitStmt_(op); } + void VisitStmt_(const tir::AllocateNode *op) final { + auto role_attr = op->annotations.Get(tl::attr::kSunmmioReduceRegisterTemp); + if (role_attr) { + const auto *role = role_attr.value().as(); + ICHECK(role) << tl::attr::kSunmmioReduceRegisterTemp + << " Allocate annotation expects an integer role"; + ICHECK(role->value == + static_cast(tl::ReduceRegisterTempRole::kAccumulator) || + role->value == + static_cast(tl::ReduceRegisterTempRole::kResult)) + << tl::attr::kSunmmioReduceRegisterTemp << " has unknown role value " + << role->value; + reduce_register_temp_roles[op->buffer_var.get()] = + static_cast(role->value); + } + tir::StmtVisitor::VisitStmt_(op); + } + void Record(const tir::Buffer &buffer) { if (!buffer.defined() || !buffer->data.defined()) { return; @@ -79,19 +98,6 @@ std::string GetAllocateStorageScope(const tir::Var &buffer_var) { TVM_FFI_UNREACHABLE(); } -bool IsSunmmioReduceRegisterTempBuffer(const tir::Buffer &buffer) { - if (!buffer.defined()) { - return false; - } - const std::string scope = buffer.scope(); - if (scope != "shared.rsram" && scope != "rsram") { - return false; - } - const std::string name = buffer->name; - return name.size() >= 4 && (name.rfind("_acc") == name.size() - 4 || - name.rfind("_res") == name.size() - 4); -} - bool IsSunmmioLocalVarBuffer(const tir::Buffer &buffer) { if (!buffer.defined()) { return false; @@ -582,6 +588,7 @@ void CodeGenTileLangSunMMIO::Clear() { local_var_table_.clear(); buffer_registry_.clear(); buffer_data_to_buffer_.clear(); + reduce_register_temp_roles_.clear(); attr_stack_.clear(); scoped_vars_.clear(); scoped_local_vars_.clear(); @@ -724,6 +731,40 @@ void CodeGenTileLangSunMMIO::CollectDeclBuffers(const tir::Stmt &stmt) { DeclBufferCollector collector; collector(stmt); buffer_data_to_buffer_ = std::move(collector.buffer_data_to_buffer); + reduce_register_temp_roles_ = std::move(collector.reduce_register_temp_roles); +} + +bool CodeGenTileLangSunMMIO::IsSunmmioReduceRegisterTempBuffer( + const tir::Buffer &buffer) const { + if (!buffer.defined() || !IsSunmmioRsramScope(buffer.scope())) { + return false; + } + return reduce_register_temp_roles_.count(buffer->data.get()) != 0; +} + +bool CodeGenTileLangSunMMIO::IsSunmmioReduceLoopCarriedTempBuffer( + const tir::Buffer &buffer) const { + if (!IsSunmmioReduceRegisterTempBuffer(buffer)) { + return false; + } + auto it = buffer.defined() + ? reduce_register_temp_roles_.find(buffer->data.get()) + : reduce_register_temp_roles_.end(); + return it != reduce_register_temp_roles_.end() && + it->second == + static_cast(tl::ReduceRegisterTempRole::kAccumulator); +} + +bool CodeGenTileLangSunMMIO::IsSunmmioReduceLocalTempBuffer( + const tir::Buffer &buffer) const { + if (!IsSunmmioReduceRegisterTempBuffer(buffer)) { + return false; + } + auto it = buffer.defined() + ? reduce_register_temp_roles_.find(buffer->data.get()) + : reduce_register_temp_roles_.end(); + return it != reduce_register_temp_roles_.end() && + it->second == static_cast(tl::ReduceRegisterTempRole::kResult); } void CodeGenTileLangSunMMIO::WriteCoverageReport() const { diff --git a/src/target/sunmmio/codegen_sunmmio.h b/src/target/sunmmio/codegen_sunmmio.h index cd986ab15b..31ba540676 100644 --- a/src/target/sunmmio/codegen_sunmmio.h +++ b/src/target/sunmmio/codegen_sunmmio.h @@ -410,6 +410,9 @@ class CodeGenTileLangSunMMIO final void VisitStmtTracked(const tir::Stmt &stmt); void CollectExpectedCoverage(const tir::PrimFunc &f); void CollectDeclBuffers(const tir::Stmt &stmt); + bool IsSunmmioReduceRegisterTempBuffer(const tir::Buffer &buffer) const; + bool IsSunmmioReduceLoopCarriedTempBuffer(const tir::Buffer &buffer) const; + bool IsSunmmioReduceLocalTempBuffer(const tir::Buffer &buffer) const; void MarkVisitedNodeType(const std::string &type_key); void MarkVisitedCallOpFromExpr(const tvm::PrimExpr &expr); void MarkVisitedExprRoot(const tvm::PrimExpr &expr); @@ -489,6 +492,7 @@ class CodeGenTileLangSunMMIO final std::unordered_map local_var_table_; std::unordered_map buffer_registry_; std::unordered_map buffer_data_to_buffer_; + std::unordered_map reduce_register_temp_roles_; std::vector attr_stack_; std::vector scoped_vars_; diff --git a/src/target/sunmmio/sunmmio_codegen_tiles_loop.cc b/src/target/sunmmio/sunmmio_codegen_tiles_loop.cc index 410e15ddd2..9923573ec8 100644 --- a/src/target/sunmmio/sunmmio_codegen_tiles_loop.cc +++ b/src/target/sunmmio/sunmmio_codegen_tiles_loop.cc @@ -59,6 +59,7 @@ struct TileBlockState { const TilesScopeInfo *scope{nullptr}; SunmmioMlirContext *mlir_ctx{nullptr}; std::unordered_map let_values; + std::unordered_map let_unsqueeze_axes; std::unordered_map register_tile_values; std::unordered_map register_tile_types; std::unordered_map register_unsqueeze_axes; @@ -84,8 +85,6 @@ struct TileAccessInfo { bool requires_aligned_1d_load{false}; int64_t aligned_load_bytes{0}; int64_t aligned_load_elems{0}; - int64_t aligned_load_axis{-1}; - std::vector aligned_load_shape; }; struct Aligned1DAddressInfo { @@ -392,31 +391,6 @@ bool IsRsramScope(const std::string &scope) { return scope == "shared.rsram" || scope == "rsram"; } -bool IsReduceRegisterTempBuffer(const Buffer &buffer) { - if (!buffer.defined() || !IsRsramScope(buffer.scope())) { - return false; - } - const std::string name = buffer->name; - return name.size() >= 4 && (name.rfind("_acc") == name.size() - 4 || - name.rfind("_res") == name.size() - 4); -} - -bool IsReduceLoopCarriedTempBuffer(const Buffer &buffer) { - if (!IsReduceRegisterTempBuffer(buffer)) { - return false; - } - const std::string name = buffer->name; - return name.size() >= 4 && name.rfind("_acc") == name.size() - 4; -} - -bool IsReduceLocalTempBuffer(const Buffer &buffer) { - if (!IsReduceRegisterTempBuffer(buffer)) { - return false; - } - const std::string name = buffer->name; - return name.size() >= 4 && name.rfind("_res") == name.size() - 4; -} - bool ContainsVectorCoreInTileReduce(const Stmt &stmt) { if (!stmt.defined()) { return false; @@ -794,11 +768,6 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { access->aligned_load_elems = align_elems; access->requires_aligned_1d_load = access->tile_shape[0] < access->aligned_load_elems; - access->aligned_load_axis = access->unsqueeze_axis == 1 ? 0 : 1; - access->aligned_load_shape = - access->unsqueeze_axis == 1 - ? std::vector{access->aligned_load_elems, 1} - : std::vector{1, access->aligned_load_elems}; }; auto analyze_access = [&](const Buffer &buffer, @@ -1173,7 +1142,7 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { auto note_register_unsqueeze_axis = [&](TileBlockState *state, const Buffer &buffer, int64_t axis) { - if (!IsReduceRegisterTempBuffer(buffer)) { + if (!IsSunmmioReduceRegisterTempBuffer(buffer)) { return; } auto existing = state->register_unsqueeze_axes.find(buffer.get()); @@ -1226,7 +1195,7 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { std::function visit_expr; auto register_buffer = [&](const Buffer &buffer, const ffi::Array &indices) { - if (!IsReduceLoopCarriedTempBuffer(buffer) || + if (!IsSunmmioReduceLoopCarriedTempBuffer(buffer) || state->register_tile_types.count(buffer.get())) { return; } @@ -1265,14 +1234,14 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { note_register_unsqueeze_axis(state, dst_region->buffer, static_cast(axis_imm->value)); } - if (IsReduceLoopCarriedTempBuffer(src_region->buffer) && + if (IsSunmmioReduceLoopCarriedTempBuffer(src_region->buffer) && !state->register_tile_types.count(src_region->buffer.get())) { SunMMIOType src_type = make_tile_type_from_region(src_region); state->register_tile_types[src_region->buffer.get()] = src_type; state->register_tile_values[src_region->buffer.get()] = make_register_tile_value(src_region->buffer, src_type); } - if (IsReduceLoopCarriedTempBuffer(dst_region->buffer) && + if (IsSunmmioReduceLoopCarriedTempBuffer(dst_region->buffer) && !state->register_tile_types.count(dst_region->buffer.get())) { SunMMIOType dst_type = make_tile_type_from_region(dst_region); state->register_tile_types[dst_region->buffer.get()] = dst_type; @@ -1685,24 +1654,20 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { DataType dst_dtype = CanonicalizeSuvmDType(access.buffer->dtype).with_lanes(1); std::vector vector_shape = access.tile_shape; - std::vector slice_shape = - access.unsqueeze_axis == 1 - ? std::vector{access.tile_shape[0], 1} - : std::vector{1, access.tile_shape[0]}; if (IsTileLike(value)) { SunMMIOValue tile = value; std::vector shape = ExtractStaticShape(tile.type); - if (shape != vector_shape && shape != slice_shape) { - tile = reorient_unit_tile_to_shape(tile, slice_shape); + if (shape != vector_shape) { + tile = reorient_unit_tile_to_shape(tile, vector_shape); shape = ExtractStaticShape(tile.type); } - ICHECK(shape == vector_shape || shape == slice_shape) + ICHECK(shape == vector_shape) << "Aligned 1D tile store cannot normalize RHS shape: src " - << shape_to_string(shape) << ", vector " - << shape_to_string(vector_shape) << ", slice " - << shape_to_string(slice_shape); - SunMMIOType dst_tile_type = MakeTileType(access.buffer->dtype, shape); + << shape_to_string(shape) << ", expected " + << shape_to_string(vector_shape); + SunMMIOType dst_tile_type = + MakeTileType(access.buffer->dtype, vector_shape); if (tile.dtype == dst_dtype && StaticShapesEqual(tile.type, dst_tile_type)) { return tile; @@ -2102,30 +2067,11 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { value_dtype); } - SunMMIOType aligned_2d_type = - MakeTileType(access.buffer->dtype, access.aligned_load_shape); - SunMMIOValue aligned_2d_tile = checked_tile_unsqueeze( - aligned_tile, aligned_2d_type, access.unsqueeze_axis, - CanonicalizeSuvmDType(access.buffer->dtype).with_lanes(1), - "aligned 1D load bridge"); - - std::vector slice_offsets; - slice_offsets.reserve(2); - if (access.aligned_load_axis == 0) { - slice_offsets.push_back(aligned_address.offset_elems); - slice_offsets.push_back(make_index_const(0)); - } else { - slice_offsets.push_back(make_index_const(0)); - slice_offsets.push_back(aligned_address.offset_elems); - } - + std::vector slice_offsets{aligned_address.offset_elems}; SunMMIOType sliced_tile_type = - MakeTileType(access.buffer->dtype, - access.unsqueeze_axis == 1 - ? std::vector{access.tile_shape[0], 1} - : std::vector{1, access.tile_shape[0]}); + MakeTileType(access.buffer->dtype, access.tile_shape); SunMMIOValue sliced_tile = builder_->TileSlice( - NewValueName(), aligned_2d_tile, slice_offsets, sliced_tile_type, + NewValueName(), aligned_tile, slice_offsets, sliced_tile_type, CanonicalizeSuvmDType(access.buffer->dtype).with_lanes(1)); return sliced_tile; }; @@ -2167,37 +2113,17 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { std::nullopt, CanonicalizeSuvmDType(access.buffer->dtype).with_lanes(1)); } - SunMMIOType aligned_store_2d_type = - MakeTileType(access.buffer->dtype, access.aligned_load_shape); - SunMMIOValue aligned_2d_tile = checked_tile_unsqueeze( - aligned_tile, aligned_store_2d_type, access.unsqueeze_axis, - CanonicalizeSuvmDType(access.buffer->dtype).with_lanes(1), - "aligned 1D store bridge"); - - std::vector slice_shape = - access.unsqueeze_axis == 1 - ? std::vector{access.tile_shape[0], 1} - : std::vector{1, access.tile_shape[0]}; - SunMMIOType slice_2d_type = MakeTileType(access.buffer->dtype, slice_shape); + std::vector slice_shape = access.tile_shape; + SunMMIOType slice_type = MakeTileType(access.buffer->dtype, slice_shape); SunMMIOValue src_slice = value; - if (src_slice.type.shape.size() == 1) { - src_slice = checked_tile_unsqueeze( - src_slice, slice_2d_type, access.unsqueeze_axis, slice_2d_type.dtype, - "aligned 1D store source slice"); - } else { + if (src_slice.type.shape.size() != 1) { src_slice = reorient_unit_tile_to_shape(src_slice, slice_shape); } + ICHECK(StaticShapesEqual(src_slice.type, slice_type)) + << "Aligned 1D store source must be rank-1"; - std::vector slice_offsets; - slice_offsets.reserve(2); - if (access.aligned_load_axis == 0) { - slice_offsets.push_back(aligned_address.offset_elems); - slice_offsets.push_back(make_index_const(0)); - } else { - slice_offsets.push_back(make_index_const(0)); - slice_offsets.push_back(aligned_address.offset_elems); - } + std::vector slice_offsets{aligned_address.offset_elems}; if (store_mask.has_value()) { SunMMIOValue mask = store_mask.value(); @@ -2218,22 +2144,17 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { DataType::Bool()); } SunMMIOValue old_slice = builder_->TileSlice( - NewValueName(), aligned_2d_tile, slice_offsets, slice_2d_type, + NewValueName(), aligned_tile, slice_offsets, slice_type, CanonicalizeSuvmDType(access.buffer->dtype).with_lanes(1)); src_slice = builder_->TileSelect( - NewValueName(), mask, src_slice, old_slice, slice_2d_type, + NewValueName(), mask, src_slice, old_slice, slice_type, CanonicalizeSuvmDType(access.buffer->dtype).with_lanes(1)); } - SunMMIOValue merged_2d_tile = builder_->TileInsertSlice( - NewValueName(), aligned_2d_tile, src_slice, slice_offsets, - aligned_store_2d_type, + SunMMIOValue merged_tile = builder_->TileInsertSlice( + NewValueName(), aligned_tile, src_slice, slice_offsets, + aligned_tile_type, CanonicalizeSuvmDType(access.buffer->dtype).with_lanes(1)); - - SunMMIOValue merged_tile = checked_tile_squeeze( - merged_2d_tile, aligned_tile_type, access.unsqueeze_axis, - CanonicalizeSuvmDType(access.buffer->dtype).with_lanes(1), - "aligned 1D store bridge"); builder_->TileStore(merged_tile, aligned_view, std::nullopt); return builder_->BindValueAlias( make_current_value_name(access.buffer, cache_key), merged_tile); @@ -2763,6 +2684,8 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { std::function>(const PrimExpr &, TileBlockState *)> infer_tile_expr_shape; + std::function(const PrimExpr &, TileBlockState *)> + infer_tile_expr_unsqueeze_axis; broadcast_tile_to_shape = [&](const SunMMIOValue &value, @@ -2868,16 +2791,6 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { return merge_optional_tile_shapes(infer_tile_expr_shape(lhs, state), infer_tile_expr_shape(rhs, state)); }; - auto infer_compare_shape = - [&](const PrimExpr &lhs, - const PrimExpr &rhs) -> std::optional> { - auto shape = infer_binary_shape(lhs, rhs); - if (shape.has_value() && shape->size() == 1) { - return std::vector{shape.value()[0], 1}; - } - return shape; - }; - if (const auto *var = expr.as()) { auto let_it = state->let_values.find(var); if (let_it != state->let_values.end() && IsTileLike(let_it->second)) { @@ -2924,10 +2837,7 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { TileAccessInfo access = analyze_access(load->buffer, load->indices, state); if (access.requires_aligned_1d_load) { - if (access.unsqueeze_axis == 1) { - return std::vector{access.tile_shape[0], 1}; - } - return std::vector{1, access.tile_shape[0]}; + return access.tile_shape; } std::string cache_key = make_tile_cache_key(access); auto current_it = state->current_tile_values.find(cache_key); @@ -2968,22 +2878,22 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { return infer_binary_shape(max->a, max->b); } if (const auto *eq = expr.as()) { - return infer_compare_shape(eq->a, eq->b); + return infer_binary_shape(eq->a, eq->b); } if (const auto *ne = expr.as()) { - return infer_compare_shape(ne->a, ne->b); + return infer_binary_shape(ne->a, ne->b); } if (const auto *lt = expr.as()) { - return infer_compare_shape(lt->a, lt->b); + return infer_binary_shape(lt->a, lt->b); } if (const auto *le = expr.as()) { - return infer_compare_shape(le->a, le->b); + return infer_binary_shape(le->a, le->b); } if (const auto *gt = expr.as()) { - return infer_compare_shape(gt->a, gt->b); + return infer_binary_shape(gt->a, gt->b); } if (const auto *ge = expr.as()) { - return infer_compare_shape(ge->a, ge->b); + return infer_binary_shape(ge->a, ge->b); } if (const auto *and_op = expr.as()) { return infer_binary_shape(and_op->a, and_op->b); @@ -3015,6 +2925,72 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { return std::nullopt; }; + infer_tile_expr_unsqueeze_axis = + [&](const PrimExpr &expr, + TileBlockState *state) -> std::optional { + std::optional inferred_axis; + bool has_conflict = false; + auto merge_axis = [&](std::optional axis) { + if (!axis.has_value() || has_conflict) { + return; + } + ICHECK(axis.value() == 0 || axis.value() == 1) + << "Rank-1 tile orientation expects an unsqueeze axis of 0 or 1"; + if (inferred_axis.has_value() && inferred_axis.value() != axis.value()) { + has_conflict = true; + inferred_axis.reset(); + return; + } + inferred_axis = axis; + }; + + tir::PostOrderVisit(expr, [&](const ObjectRef &obj) { + if (has_conflict) { + return; + } + if (const auto *var = obj.as()) { + auto axis_it = state->let_unsqueeze_axes.find(var); + if (axis_it != state->let_unsqueeze_axes.end()) { + merge_axis(axis_it->second); + } + return; + } + const auto *load = obj.as(); + if (!load) { + return; + } + + auto local_it = state->local_tile_values.find(load->buffer.get()); + if (local_it != state->local_tile_values.end()) { + if (ExtractStaticShape(local_it->second.type).size() == 1) { + auto axis_it = state->local_unit_tile_axes.find(load->buffer.get()); + if (axis_it != state->local_unit_tile_axes.end()) { + merge_axis(axis_it->second); + } + } + return; + } + auto register_it = state->register_tile_values.find(load->buffer.get()); + if (register_it != state->register_tile_values.end()) { + if (ExtractStaticShape(register_it->second.type).size() == 1) { + auto axis_it = + state->register_unsqueeze_axes.find(load->buffer.get()); + if (axis_it != state->register_unsqueeze_axes.end()) { + merge_axis(axis_it->second); + } + } + return; + } + + TileAccessInfo access = + analyze_access(load->buffer, load->indices, state); + if (access.tile_rank == 1 && access.tile_shape[0] > 1) { + merge_axis(access.unsqueeze_axis); + } + }); + return has_conflict ? std::nullopt : inferred_axis; + }; + lower_expr = [&](const PrimExpr &expr, TileBlockState *state, std::optional preferred_dtype) -> SunMMIOValue { MarkVisitedExprRoot(expr); @@ -3044,6 +3020,27 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { } return rewritten_condition; }; + auto orient_expr_tile_operand = + [&](const SunMMIOValue &value, const PrimExpr &source_expr, + const std::vector &result_shape) -> SunMMIOValue { + if (!IsTileLike(value)) { + return value; + } + std::vector src_shape = ExtractStaticShape(value.type); + if (src_shape.size() == 1 && result_shape.size() == 2) { + if (auto axis = infer_tile_expr_unsqueeze_axis(source_expr, state)) { + std::vector unit_shape = + axis.value() == 1 ? std::vector{src_shape[0], 1} + : std::vector{1, src_shape[0]}; + if (CanBroadcastShapeTo(unit_shape, result_shape)) { + return checked_tile_unsqueeze( + value, MakeTileType(value.dtype, unit_shape), axis.value(), + value.dtype, "orient rank-1 expression at rank-2 consumer"); + } + } + } + return orient_tile_operand_to_shape(value, result_shape); + }; auto emit_select = [&](const PrimExpr &condition, const PrimExpr &true_value_expr, const PrimExpr &false_value_expr, DataType dtype) { @@ -3072,9 +3069,7 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { IsTileLike(false_value)) { std::vector result_shape; if (IsTileLike(true_value) && IsTileLike(false_value)) { - result_shape = - merge_broadcast_shapes(ExtractStaticShape(true_value.type), - ExtractStaticShape(false_value.type)); + result_shape = tile_result_shape(true_value, false_value); } else if (IsTileLike(true_value)) { result_shape = ExtractStaticShape(true_value.type); } else if (IsTileLike(false_value)) { @@ -3083,7 +3078,8 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { result_shape = ExtractStaticShape(cond.type); } if (IsTileLike(cond)) { - cond = reorient_unit_tile_to_shape(cond, result_shape); + cond = + orient_expr_tile_operand(cond, condition_to_lower, result_shape); std::vector cond_shape = ExtractStaticShape(cond.type); if (cond_shape != result_shape) { result_shape = merge_broadcast_shapes(cond_shape, result_shape); @@ -3093,6 +3089,8 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { SunMMIOType scalar_type{ SunMMIOType::Kind::kScalar, result_dtype, 1, {}}; if (IsTileLike(true_value)) { + true_value = orient_expr_tile_operand(true_value, true_value_expr, + result_shape); true_value = broadcast_tile_to_shape( cast_value_to_dtype(true_value, result_dtype), result_shape); } else { @@ -3102,6 +3100,8 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { MakeTileType(result_dtype, result_shape), result_dtype); } if (IsTileLike(false_value)) { + false_value = orient_expr_tile_operand(false_value, false_value_expr, + result_shape); false_value = broadcast_tile_to_shape( cast_value_to_dtype(false_value, result_dtype), result_shape); } else { @@ -3145,6 +3145,9 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { SunMMIOValue value = lower_expr(let->value, state, preferred_dtype); TileBlockState let_state = *state; let_state.let_values[let->var.get()] = value; + if (auto axis = infer_tile_expr_unsqueeze_axis(let->value, state)) { + let_state.let_unsqueeze_axes[let->var.get()] = axis.value(); + } return lower_expr(let->body, &let_state, preferred_dtype); } if (const auto *load = expr.as()) { @@ -3264,17 +3267,6 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { tile_type, CanonicalizeSuvmDType(load->buffer->dtype).with_lanes(1)); } - if (access.tile_rank == 1 && scope.tile_shape.size() == 2) { - std::vector unit_shape = - access.unsqueeze_axis == 1 - ? std::vector{access.tile_shape[0], 1} - : std::vector{1, access.tile_shape[0]}; - tile = checked_tile_unsqueeze( - tile, MakeTileType(load->buffer->dtype, unit_shape), - access.unsqueeze_axis, - CanonicalizeSuvmDType(load->buffer->dtype).with_lanes(1), - "rank-1 tile load orientation"); - } } if (!access.promoted_unit_tile_view && !access.requires_aligned_1d_load) { state->current_tile_values[cache_key] = builder_->BindValueAlias( @@ -3325,11 +3317,8 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { rhs, result_type, result_dtype); } std::vector result_shape = tile_result_shape(lhs, rhs); - if (result_shape.size() == 1 && (IsTileLike(lhs) || IsTileLike(rhs))) { - result_shape = {result_shape[0], 1}; - } - lhs = orient_tile_operand_to_shape(lhs, result_shape); - rhs = orient_tile_operand_to_shape(rhs, result_shape); + lhs = orient_expr_tile_operand(lhs, lhs_expr, result_shape); + rhs = orient_expr_tile_operand(rhs, rhs_expr, result_shape); SunMMIOType tile_type = MakeTileType(result_dtype, result_shape); auto broadcast_scalar_to_tile = [&](SunMMIOValue value) { if (IsTileLike(value)) { @@ -3381,15 +3370,11 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { std::vector result_shape = forced_shape.has_value() ? forced_shape.value() : tile_result_shape(lhs, rhs); - if (!forced_shape.has_value() && result_shape.size() == 1 && - (IsTileLike(lhs) || IsTileLike(rhs))) { - result_shape = {result_shape[0], 1}; - } SunMMIOType operand_type = MakeTileType( IsTileLike(lhs) ? lhs.dtype : rhs.dtype, result_shape); DataType operand_dtype = operand_type.dtype; - lhs = orient_tile_operand_to_shape(lhs, result_shape); - rhs = orient_tile_operand_to_shape(rhs, result_shape); + lhs = orient_expr_tile_operand(lhs, lhs_expr, result_shape); + rhs = orient_expr_tile_operand(rhs, rhs_expr, result_shape); if (forced_shape.has_value()) { auto operand_shape = [&](const SunMMIOValue &value) { if (IsTileLike(value)) { @@ -3474,7 +3459,8 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { if (!IsTileLike(value)) { return ensure_logical_scalar(value); } - SunMMIOValue tile = reorient_unit_tile_to_shape(value, target_shape); + SunMMIOValue tile = + orient_expr_tile_operand(value, source_expr, target_shape); if (ExtractStaticShape(tile.type) == target_shape) { return tile; } @@ -3572,9 +3558,6 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { lhs, rhs, result_type, result_type.dtype); } std::vector fallback_shape = tile_result_shape(lhs, rhs); - if (fallback_shape.size() == 1 && (IsTileLike(lhs) || IsTileLike(rhs))) { - fallback_shape = {fallback_shape[0], 1}; - } return emit_logical_values(op, lhs, rhs, fallback_shape, lhs_expr, rhs_expr); }; @@ -3670,18 +3653,13 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { select->false_value, select->dtype); } if (const auto *cast = expr.as()) { - SunMMIOValue value = lower_expr(cast->value, state, preferred_dtype); + // An explicit cast applies after its operand has been evaluated. + SunMMIOValue value = lower_expr(cast->value, state, std::nullopt); if (IsTileLike(value)) { DataType dst_dtype = CanonicalizeSuvmDType(cast->dtype).with_lanes(1); if (value.dtype == dst_dtype) { return value; } - if (preferred_dtype.has_value() && is_float_like_dtype(value.dtype) && - is_float_like_dtype(dst_dtype) && - value.dtype == - CanonicalizeSuvmDType(preferred_dtype.value()).with_lanes(1)) { - return value; - } SunMMIOType dst_type = MakeTileType(CanonicalizeSuvmDType(cast->dtype), ExtractStaticShape(value.type)); return builder_->Cast(NewValueName(), value, dst_type, @@ -3853,6 +3831,9 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { SunMMIOValue value = lower_expr(let->value, state, std::nullopt); TileBlockState let_state = *state; let_state.let_values[let->var.get()] = value; + if (auto axis = infer_tile_expr_unsqueeze_axis(let->value, state)) { + let_state.let_unsqueeze_axes[let->var.get()] = axis.value(); + } lower_stmt(let->body, &let_state); state->tile_view_cache = let_state.tile_view_cache; state->current_tile_values = let_state.current_tile_values; @@ -3865,7 +3846,7 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { if (const auto *alloc = stmt.as()) { auto buffer_it = buffer_data_to_buffer_.find(alloc->buffer_var.get()); if (buffer_it != buffer_data_to_buffer_.end() && - IsReduceRegisterTempBuffer(buffer_it->second)) { + IsSunmmioReduceRegisterTempBuffer(buffer_it->second)) { EnterScope(); lower_reduce_stmt(alloc->body, state); ExitScope(); @@ -3876,7 +3857,7 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { "reduce register temporaries"); } if (const auto *decl = stmt.as()) { - if (IsReduceRegisterTempBuffer(decl->buffer)) { + if (IsSunmmioReduceRegisterTempBuffer(decl->buffer)) { EnterScope(); lower_reduce_stmt(decl->body, state); ExitScope(); @@ -4189,7 +4170,7 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { result_tile_type, axis, CanonicalizeSuvmDType(src_region->buffer->dtype).with_lanes(1)); - if (IsReduceLocalTempBuffer(dst_region->buffer)) { + if (IsSunmmioReduceLocalTempBuffer(dst_region->buffer)) { SunMMIOValue local = builder_->BindValueAlias( make_local_value_name(dst_region->buffer), reduced); state->local_tile_values[dst_region->buffer.get()] = local; @@ -4198,7 +4179,7 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { } bool dst_is_register = - IsReduceRegisterTempBuffer(dst_region->buffer) && + IsSunmmioReduceRegisterTempBuffer(dst_region->buffer) && state->register_tile_types.count(dst_region->buffer.get()); if (dst_is_register) { SunMMIOType dst_tile_type = @@ -4323,6 +4304,9 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { SunMMIOValue value = lower_expr(let->value, state, std::nullopt); TileBlockState let_state = *state; let_state.let_values[let->var.get()] = value; + if (auto axis = infer_tile_expr_unsqueeze_axis(let->value, state)) { + let_state.let_unsqueeze_axes[let->var.get()] = axis.value(); + } lower_reduce_stmt(let->body, &let_state); state->tile_view_cache = let_state.tile_view_cache; state->current_tile_values = let_state.current_tile_values; @@ -4520,7 +4504,7 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { if (const auto *alloc = stmt.as()) { auto buffer_it = buffer_data_to_buffer_.find(alloc->buffer_var.get()); if (buffer_it != buffer_data_to_buffer_.end() && - IsReduceRegisterTempBuffer(buffer_it->second)) { + IsSunmmioReduceRegisterTempBuffer(buffer_it->second)) { EnterScope(); lower_reduce_stmt(alloc->body, state); ExitScope(); @@ -4531,7 +4515,7 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) { "reduce register temporaries"); } if (const auto *decl = stmt.as()) { - if (IsReduceRegisterTempBuffer(decl->buffer)) { + if (IsSunmmioReduceRegisterTempBuffer(decl->buffer)) { EnterScope(); lower_reduce_stmt(decl->body, state); ExitScope(); diff --git a/src/transform/common/attr.h b/src/transform/common/attr.h index 1fb8b41cd9..b6e73c2309 100644 --- a/src/transform/common/attr.h +++ b/src/transform/common/attr.h @@ -54,6 +54,12 @@ constexpr const char *tile_interior = "tile.interior"; // Which axis of the tile shape this interior loop corresponds to (0, 1, ...) constexpr const char *tile_interior_axis = "tile.interior_axis"; +// Marks an RSRAM buffer data var as a reduction temporary lowered to an SSA +// tile register by the SunMMIO backend. The integer value is a +// ReduceRegisterTempRole. +constexpr const char *kSunmmioReduceRegisterTemp = + "tile.sunmmio_reduce_register_temp"; + } // namespace attr enum class TileLoopStage : int { @@ -63,5 +69,10 @@ enum class TileLoopStage : int { kConsumed = 3, }; +enum class ReduceRegisterTempRole : int { + kAccumulator = 1, + kResult = 2, +}; + } // namespace tl } // namespace tvm diff --git a/src/transform/common/sunmmio_pipeline_utils.h b/src/transform/common/sunmmio_pipeline_utils.h index 61cc48da22..ee3fd3d0b4 100644 --- a/src/transform/common/sunmmio_pipeline_utils.h +++ b/src/transform/common/sunmmio_pipeline_utils.h @@ -2,6 +2,12 @@ #define SUNMMIO_PIPELINE_UTILS_H #include +#include + +#include +#include +#include +#include namespace tvm { namespace tl { @@ -13,6 +19,95 @@ inline int name2id(const std::string &name) { return std::stoi(name.substr(name.find('-') + 1)); } +inline ffi::Map +BuildPipelineIterZeroSubstitutionMap(const PrimExpr &expr, + const tir::Var &pipeline_loop_var) { + std::unordered_set vars; + tir::PostOrderVisit(expr, [&](const ObjectRef &obj) { + if (const auto *var = obj.as()) { + if (!pipeline_loop_var.defined() || var != pipeline_loop_var.get()) { + vars.insert(var); + } + } + }); + + ffi::Map vmap; + for (const tir::VarNode *node : vars) { + tir::Var var = ffi::GetRef(node); + vmap.Set(var, tir::make_zero(var.dtype())); + } + return vmap; +} + +inline int DetectPipelineIterOffsetFromExpr(const PrimExpr &expr, + const tir::Var &pipeline_loop_var, + arith::Analyzer *analyzer) { + if (!pipeline_loop_var.defined() || + !tir::UsesVar(expr, + [v = pipeline_loop_var.get()](const tir::VarNode *node) { + return node == v; + })) { + return 0; + } + + PrimExpr loop_only = expr; + ffi::Map vmap = + BuildPipelineIterZeroSubstitutionMap(expr, pipeline_loop_var); + if (!vmap.empty()) { + loop_only = tir::Substitute(loop_only, vmap); + } + loop_only = analyzer->Simplify(loop_only); + + ffi::Array coeffs = arith::DetectLinearEquation( + loop_only, ffi::Array{pipeline_loop_var}); + if (coeffs.size() != 2) { + return 0; + } + + PrimExpr coeff = analyzer->Simplify(coeffs[0]); + PrimExpr base = analyzer->Simplify(coeffs[1]); + const auto *coeff_int = coeff.as(); + if (coeff_int == nullptr || coeff_int->value == 0) { + return 0; + } + + PrimExpr offset_expr = analyzer->Simplify(floordiv(base, coeff)); + PrimExpr remainder = analyzer->Simplify(floormod(base, coeff)); + if (!analyzer->CanProveEqual(remainder, tir::make_zero(remainder.dtype()))) { + return 0; + } + + const auto *offset_int = offset_expr.as(); + return offset_int == nullptr ? 0 : static_cast(offset_int->value); +} + +inline int DetectPipelineIterOffsetFromRegion(const tir::BufferRegion ®ion, + const tir::Var &pipeline_loop_var, + arith::Analyzer *analyzer) { + int result = 0; + bool found = false; + for (const Range &range : region->region) { + bool uses_loop_var = + pipeline_loop_var.defined() && + tir::UsesVar(range->min, + [v = pipeline_loop_var.get()](const tir::VarNode *node) { + return node == v; + }); + if (!uses_loop_var) { + continue; + } + int dim_offset = DetectPipelineIterOffsetFromExpr( + range->min, pipeline_loop_var, analyzer); + if (!found) { + result = dim_offset; + found = true; + } else if (result != dim_offset) { + return 0; + } + } + return found ? result : 0; +} + } // namespace tl } // namespace tvm #endif diff --git a/src/transform/inject_sunmmio_pipeline.cc b/src/transform/inject_sunmmio_pipeline.cc index ae028e3b51..c4967cf3be 100644 --- a/src/transform/inject_sunmmio_pipeline.cc +++ b/src/transform/inject_sunmmio_pipeline.cc @@ -10,6 +10,7 @@ #include "../tileview/tileview.h" #include "common/loop_fusion_utils.h" #include "common/remap_buffer_rewriter.h" +#include "sunmmio_pipeline_planning/pipeline_diagnostic.h" #include "sunmmio_pipeline_planning/stmt_read_write_collector.h" #include "sunmmio_pipeline_planning/sunmmio_pipeline_utils.h" #include "tir/transforms/ir_utils.h" @@ -31,6 +32,11 @@ #include #include +#include +#include +#include +#include + namespace tvm { namespace tl { @@ -57,9 +63,20 @@ class SunmmioMultiVersionBufferRewriter : public StmtExprMutator { substituter.replace_flag = true; for (auto &buffer : substituter.versioned_buffers_) { - substituter.buffer_remap_.Set( - buffer, - substituter.makeMultiVersionBuffer(buffer, substituter.iterations_)); + int versions = substituter.version_counts_.at(buffer.get()); + if (substituter.IsBankedBuffer(buffer)) { + int ping_versions = (versions + 1) / 2; + int pong_versions = ping_versions; + Buffer ping = substituter.makeMultiVersionBuffer(buffer, ping_versions, + "_ping", true); + Buffer pong = substituter.makeMultiVersionBuffer(buffer, pong_versions, + "_pong", false); + substituter.buffer_remap_.Set(buffer, ping); + substituter.bank_peer_buffers_[buffer.get()] = pong; + } else { + substituter.buffer_remap_.Set( + buffer, substituter.makeMultiVersionBuffer(buffer, versions)); + } } substituter.RewriteFunctionLayoutAttrs(f); @@ -95,6 +112,17 @@ class SunmmioMultiVersionBufferRewriter : public StmtExprMutator { << "Failed to derive multiversioned layout for buffer " << buffer->name << " with shape " << new_buffer->shape; new_layout_map.Set(new_buffer, derived_layout.value()); + auto peer_it = bank_peer_buffers_.find(buffer.get()); + if (peer_it != bank_peer_buffers_.end()) { + const Buffer &peer_buffer = peer_it->second; + Optional peer_layout = DeriveLayoutLikeForDType( + layout, peer_buffer->shape, peer_buffer->dtype, + Optional>(), &analyzer); + ICHECK(peer_layout.defined()) + << "Failed to derive ping/pong layout for buffer " << buffer->name + << " with shape " << peer_buffer->shape; + new_layout_map.Set(peer_buffer, peer_layout.value()); + } } f = WithAttr(std::move(f), attr::kLayoutMap, new_layout_map); } @@ -105,14 +133,8 @@ class SunmmioMultiVersionBufferRewriter : public StmtExprMutator { } Map alloc_ping_pong; - for (const auto &kv : buffer_remap_) { - const Buffer &buffer = kv.first; - if (buffer.scope() != kSunmmioScopeASRAM && - buffer.scope() != kSunmmioScopeWSRAM) { - continue; - } - const Buffer &new_buffer = kv.second; - alloc_ping_pong.Set(new_buffer->data, String("pong")); + for (const auto &kv : bank_peer_buffers_) { + alloc_ping_pong.Set(kv.second->data, String("pong")); } if (alloc_ping_pong.empty()) { @@ -123,23 +145,37 @@ class SunmmioMultiVersionBufferRewriter : public StmtExprMutator { alloc_ping_pong); } - Buffer makeMultiVersionBuffer(const Buffer &buffer, int num_version) { + bool IsBankedBuffer(const Buffer &buffer) const { + return buffer.scope() == kSunmmioScopeASRAM || + buffer.scope() == kSunmmioScopeWSRAM; + } + + Buffer makeMultiVersionBuffer(const Buffer &buffer, int num_version, + const std::string &name_suffix = "", + bool reuse_primary_var = true) { const auto *ptr_type = TVM_TYPE_AS(buffer->data->type_annotation, PointerTypeNode); Var new_var; - if (var_remap_.count(buffer->data)) { + std::string data_name = std::string(buffer->data->name_hint) + name_suffix; + std::string buffer_name = std::string(buffer->name) + name_suffix; + if (reuse_primary_var && var_remap_.count(buffer->data)) { new_var = var_remap_[buffer->data]; } else { Type new_type = PointerType(ptr_type->element_type, ptr_type->storage_scope); - new_var = Var(buffer->data->name_hint, new_type); - var_remap_.Set(buffer->data, new_var); + new_var = Var(data_name, new_type); + if (reuse_primary_var) { + var_remap_.Set(buffer->data, new_var); + } } auto shape = buffer->shape; - shape.insert(shape.begin(), num_version); + if (num_version > 1) { + shape.insert(shape.begin(), num_version); + } + buffer_has_version_axis_[new_var.get()] = num_version > 1; return Buffer(new_var, buffer->dtype, shape, {}, buffer->elem_offset, - buffer->name, buffer->data_alignment, buffer->offset_factor, - buffer->buffer_type); + String(buffer_name), buffer->data_alignment, + buffer->offset_factor, buffer->buffer_type); } BufferRegion @@ -148,8 +184,10 @@ class SunmmioMultiVersionBufferRewriter : public StmtExprMutator { if (it != buffer_remap_.end()) { Region new_region = buffer_region->region; const Buffer &new_buffer = (*it).second; - Range accessed_version = Range::FromMinExtent(0, 1); - new_region.insert(new_region.begin(), accessed_version); + if (HasVersionAxis(new_buffer)) { + Range accessed_version = Range::FromMinExtent(0, 1); + new_region.insert(new_region.begin(), accessed_version); + } return BufferRegion(new_buffer, new_region); } return buffer_region; @@ -160,13 +198,22 @@ class SunmmioMultiVersionBufferRewriter : public StmtExprMutator { auto versioned_buffers_anno = op->annotations.Get("versioned_buffers"); auto used_buffers_anno = op->annotations.Get("used_buffers"); auto iterations_anno = op->annotations.Get("iterations"); + auto writer_phases_anno = op->annotations.Get("runtime_bank_writer_phases"); + auto reader_phases_anno = op->annotations.Get("runtime_bank_reader_phases"); if (versioned_buffers_anno && used_buffers_anno && iterations_anno) { Array versioned_buffers = Downcast>(versioned_buffers_anno.value()); int iterations = Downcast(iterations_anno.value()); if (!replace_flag) { - versioned_buffers_ = versioned_buffers; - iterations_ = iterations; + for (const Buffer &buffer : versioned_buffers) { + auto [it, inserted] = + version_counts_.try_emplace(buffer.get(), iterations); + if (inserted) { + versioned_buffers_.push_back(buffer); + } else { + it->second = std::max(it->second, iterations); + } + } } else { Array new_versioned_buffers; for (const Buffer &buffer : versioned_buffers) { @@ -178,6 +225,37 @@ class SunmmioMultiVersionBufferRewriter : public StmtExprMutator { } loop.CopyOnWrite()->annotations.Set("versioned_buffers", new_versioned_buffers); + Map bank_peer_buffers; + for (const Buffer &buffer : versioned_buffers) { + auto remap_it = buffer_remap_.find(buffer); + auto peer_it = bank_peer_buffers_.find(buffer.get()); + if (remap_it != buffer_remap_.end() && + peer_it != bank_peer_buffers_.end()) { + bank_peer_buffers.Set((*remap_it).second, peer_it->second); + } + } + if (!bank_peer_buffers.empty()) { + loop.CopyOnWrite()->annotations.Set("bank_peer_buffers", + bank_peer_buffers); + } + Array version_axis_buffers; + for (const Buffer &buffer : versioned_buffers) { + auto remap_it = buffer_remap_.find(buffer); + if (remap_it == buffer_remap_.end()) { + continue; + } + const Buffer &remapped = (*remap_it).second; + if (HasVersionAxis(remapped)) { + version_axis_buffers.push_back(remapped); + } + auto peer_it = bank_peer_buffers_.find(buffer.get()); + if (peer_it != bank_peer_buffers_.end() && + HasVersionAxis(peer_it->second)) { + version_axis_buffers.push_back(peer_it->second); + } + } + loop.CopyOnWrite()->annotations.Set("version_axis_buffers", + version_axis_buffers); Array used_buffers = Downcast>(used_buffers_anno.value()); Array new_used_buffers; @@ -189,6 +267,24 @@ class SunmmioMultiVersionBufferRewriter : public StmtExprMutator { } } loop.CopyOnWrite()->annotations.Set("used_buffers", new_used_buffers); + auto remap_phase_annotation = [&](const Optional &annotation, + const char *name) { + if (!annotation) + return; + Map> phases = + Downcast>>(annotation.value()); + Map> remapped; + for (const auto &[buffer, per_command] : phases) { + auto it = buffer_remap_.find(buffer); + remapped.Set(it == buffer_remap_.end() ? buffer : (*it).second, + per_command); + } + loop.CopyOnWrite()->annotations.Set(name, remapped); + }; + remap_phase_annotation(writer_phases_anno, + "runtime_bank_writer_phases"); + remap_phase_annotation(reader_phases_anno, + "runtime_bank_reader_phases"); } } return loop; @@ -213,6 +309,10 @@ class SunmmioMultiVersionBufferRewriter : public StmtExprMutator { for (const auto &[buffer, layout] : map) { if (buffer_remap_.count(buffer)) { new_map.Set(buffer_remap_[buffer], layout); + auto peer_it = bank_peer_buffers_.find(buffer.get()); + if (peer_it != bank_peer_buffers_.end()) { + new_map.Set(peer_it->second, layout); + } } else { new_map.Set(buffer, layout); } @@ -245,15 +345,19 @@ class SunmmioMultiVersionBufferRewriter : public StmtExprMutator { }); // do block->alloc_buffers remap - Array alloc_buffers = block->alloc_buffers; - - // remove the buffers - alloc_buffers.MutateByApply([this](Buffer buf) { - if (buffer_remap_.find(buf) != buffer_remap_.end()) { - return buffer_remap_.at(buf); + Array alloc_buffers; + for (const Buffer &buf : block->alloc_buffers) { + auto remap_it = buffer_remap_.find(buf); + if (remap_it == buffer_remap_.end()) { + alloc_buffers.push_back(buf); + continue; } - return buf; - }); + alloc_buffers.push_back((*remap_it).second); + auto peer_it = bank_peer_buffers_.find(buf.get()); + if (peer_it != bank_peer_buffers_.end()) { + alloc_buffers.push_back(peer_it->second); + } + } if (!alloc_buffers.same_as(block->alloc_buffers)) { block.CopyOnWrite()->alloc_buffers = alloc_buffers; @@ -272,7 +376,9 @@ class SunmmioMultiVersionBufferRewriter : public StmtExprMutator { if (buffer_remap_.count(buffer)) { auto new_buffer = buffer_remap_[load->buffer]; auto indices = load->indices; - indices.insert(indices.begin(), 0); + if (HasVersionAxis(new_buffer)) { + indices.insert(indices.begin(), 0); + } return BufferLoad(new_buffer, indices); } auto expr = StmtExprMutator::VisitExpr_(op); @@ -288,7 +394,9 @@ class SunmmioMultiVersionBufferRewriter : public StmtExprMutator { if (buffer_remap_.count(buffer)) { auto new_buffer = buffer_remap_[store->buffer]; auto indices = store->indices; - indices.insert(indices.begin(), 0); + if (HasVersionAxis(new_buffer)) { + indices.insert(indices.begin(), 0); + } return BufferStore(new_buffer, store->value, indices); } return store; @@ -317,7 +425,9 @@ class SunmmioMultiVersionBufferRewriter : public StmtExprMutator { Buffer new_buffer = buffer_remap_[original_buffer]; Array new_ranges = original_region->GetRanges(); - new_ranges.insert(new_ranges.begin(), Range(0, 1)); + if (HasVersionAxis(new_buffer)) { + new_ranges.insert(new_ranges.begin(), Range(0, 1)); + } Array new_args; new_args.push_back(BufferLoad(new_buffer, [new_ranges]() { @@ -351,24 +461,56 @@ class SunmmioMultiVersionBufferRewriter : public StmtExprMutator { } Array versioned_buffers_; - int iterations_ = -1; + bool HasVersionAxis(const Buffer &buffer) const { + auto it = buffer_has_version_axis_.find(buffer->data.get()); + return it != buffer_has_version_axis_.end() && it->second; + } + + std::unordered_map version_counts_; bool replace_flag = false; Map buffer_remap_; Map var_remap_; Map buffer_data_to_buffer_; + std::unordered_map bank_peer_buffers_; + std::unordered_map buffer_has_version_axis_; }; class PipelineBodyRewriter : public StmtExprMutator { public: - PipelineBodyRewriter(Array used_buffers, For pipeline_loop) { + PipelineBodyRewriter(Array used_buffers, + Map bank_peer_buffers, + Array version_axis_buffers, + Map> writer_phases, + Map> reader_phases, + For pipeline_loop) { used_buffers_ = used_buffers; + bank_peer_buffers_ = std::move(bank_peer_buffers); pipeline_loop_ = std::move(pipeline_loop); + for (const Buffer &buffer : version_axis_buffers) { + version_axis_buffers_.insert(buffer.get()); + } + auto import_phases = [](const Map> &source, + auto *destination) { + for (const auto &[buffer, phases] : source) { + auto &per_command = (*destination)[buffer.get()]; + for (const auto &[command_id, phase] : phases) { + per_command[command_id->value] = Downcast(phase)->value; + } + } + }; + import_phases(writer_phases, &writer_phases_); + import_phases(reader_phases, &reader_phases_); for (auto it : used_buffers) { buffer_data_to_buffer_.Set(it->data, it); + if (bank_peer_buffers_.count(it)) { + const Buffer &peer = bank_peer_buffers_[it]; + buffer_data_to_buffer_.Set(peer->data, peer); + } } } void set_current_version(int v) { current_version_ = v; } + void set_current_command(int id) { current_command_id_ = id; } void set_loop_var_replacement(PrimExpr p) { replaced_loop_var_ = p; } @@ -384,25 +526,32 @@ class PipelineBodyRewriter : public StmtExprMutator { }; Array new_args = call->args; for (int i : arg_indices) { - // const Buffer &buffer = - // buffer_data_to_buffer_.at(Downcast(call->args[i])); - // auto it = buffer_remap_.find(buffer); - // if (it != buffer_remap_.end()) { - // const Buffer &new_buffer = (*it).second; - // const PrimExpr &old_index = call->args[i + 1]; - // LOG(INFO) << old_index; - // PrimExpr offset; - // if (new_buffer->strides.empty()) { - // offset = product(buffer->shape); - // } else { - // offset = new_buffer->strides[0]; - // } - // PrimExpr new_index = - // old_index + - // floormod(pipeline_loop_->loop_var, new_buffer->shape[0]) * - // offset; - // LOG(INFO) << new_index; - // new_args.Set(i + 1, new_index); + Var data = Downcast(call->args[i]); + if (!buffer_data_to_buffer_.count(data)) { + continue; + } + const Buffer &buffer = buffer_data_to_buffer_[data]; + if (!IsVersionedBuffer(buffer)) { + continue; + } + Buffer target = ResolveTargetBuffer(buffer, true); + if (!HasVersionAxis(target)) { + new_args.Set(i, target->data); + continue; + } + PrimExpr offset; + if (!target->strides.empty()) { + offset = target->strides[0]; + } else { + Array inner_shape; + for (size_t axis = 1; axis < target->shape.size(); ++axis) { + inner_shape.push_back(target->shape[axis]); + } + offset = product(inner_shape); + } + new_args.Set(i, target->data); + new_args.Set(i + 1, call->args[i + 1] + + Integer(CurrentVersionSlot(buffer)) * offset); } return Call(call->dtype, call->op, new_args, call->annotations, call->span); } @@ -429,9 +578,10 @@ class PipelineBodyRewriter : public StmtExprMutator { if (!count) { return store; } - auto *n = store.CopyOnWrite(); - n->indices.Set(0, current_version_); - return store; + Buffer target = ResolveTargetBuffer(store->buffer, true); + Array indices = store->indices; + RewriteIndices(store->buffer, target, &indices); + return BufferStore(target, store->value, indices); } PrimExpr VisitExpr_(const BufferLoadNode *op) final { @@ -444,12 +594,44 @@ class PipelineBodyRewriter : public StmtExprMutator { if (!count) { return load; } - auto *n = load.CopyOnWrite(); - n->indices.Set(0, current_version_); - return load; + Buffer target = ResolveTargetBuffer(load->buffer, false); + Array indices = load->indices; + RewriteIndices(load->buffer, target, &indices); + return BufferLoad(target, indices); } PrimExpr VisitExpr_(const CallNode *op) final { + if (op->op.same_as(RegionOp::Get())) { + RegionOp original_region(op->args); + Buffer source = original_region->GetBuffer(); + if (IsVersionedBuffer(source)) { + bool is_write = HasCommandPhase(writer_phases_, source); + Buffer target = ResolveTargetBuffer(source, is_write); + Array ranges = original_region->GetRanges(); + if (HasVersionAxis(target)) { + ICHECK(HasVersionAxis(source)); + ranges.Set(0, Range::FromMinExtent(CurrentVersionSlot(source), 1)); + } else if (HasVersionAxis(source)) { + Array squeezed; + for (size_t i = 1; i < ranges.size(); ++i) { + squeezed.push_back(ranges[i]); + } + ranges = squeezed; + } + + Array args; + Array mins; + for (const Range &range : ranges) { + mins.push_back(VisitExpr(range->min)); + } + args.push_back(BufferLoad(target, mins)); + args.push_back(VisitExpr(original_region->GetAccessMask())); + for (const Range &range : ranges) { + args.push_back(VisitExpr(range->extent)); + } + return Call(DataType::Handle(), RegionOp::Get(), args); + } + } Call call = Downcast(StmtExprMutator::VisitExpr_(op)); if (call->op.same_as(builtin::tvm_access_ptr())) { return RewriteBufferAccess(call, {1}); @@ -465,10 +647,78 @@ class PipelineBodyRewriter : public StmtExprMutator { return var; } + bool IsVersionedBuffer(const Buffer &buffer) const { + for (const Buffer &candidate : used_buffers_) { + if (candidate.same_as(buffer)) { + return true; + } + } + return false; + } + + bool IsBankedBuffer(const Buffer &buffer) const { + return bank_peer_buffers_.count(buffer) != 0; + } + + bool HasCommandPhase( + const std::unordered_map> + &phases, + const Buffer &buffer) const { + auto it_buffer = phases.find(buffer.get()); + return it_buffer != phases.end() && + it_buffer->second.count(current_command_id_) != 0; + } + + int CurrentPhase(const Buffer &buffer, bool is_write) const { + const auto &phases = is_write ? writer_phases_ : reader_phases_; + auto it_buffer = phases.find(buffer.get()); + if (it_buffer == phases.end()) + return 0; + auto it_command = it_buffer->second.find(current_command_id_); + return it_command == it_buffer->second.end() ? 0 : it_command->second; + } + + Buffer ResolveTargetBuffer(const Buffer &buffer, bool is_write) const { + if (!IsBankedBuffer(buffer) || + (current_version_ + CurrentPhase(buffer, is_write)) % 2 == 0) { + return buffer; + } + return bank_peer_buffers_[buffer]; + } + + int CurrentVersionSlot(const Buffer &buffer) const { + return IsBankedBuffer(buffer) ? current_version_ / 2 : current_version_; + } + + bool HasVersionAxis(const Buffer &buffer) const { + return version_axis_buffers_.count(buffer.get()) != 0; + } + + void RewriteIndices(const Buffer &source, const Buffer &target, + Array *indices) const { + if (HasVersionAxis(target)) { + ICHECK(HasVersionAxis(source)); + indices->Set(0, CurrentVersionSlot(source)); + } else if (HasVersionAxis(source)) { + Array squeezed; + for (size_t i = 1; i < indices->size(); ++i) { + squeezed.push_back((*indices)[i]); + } + *indices = squeezed; + } + } + Array used_buffers_; Map buffer_data_to_buffer_; + Map bank_peer_buffers_; + std::unordered_set version_axis_buffers_; For pipeline_loop_; int current_version_ = 0; + int current_command_id_ = -1; + std::unordered_map> + writer_phases_; + std::unordered_map> + reader_phases_; PrimExpr replaced_loop_var_; }; @@ -498,15 +748,40 @@ class SunmmioPipelineInjector : public StmtExprMutator { auto iterations_anno = op->annotations.Get("iterations"); auto used_buffers_anno = op->annotations.Get("used_buffers"); auto versioned_buffers_anno = op->annotations.Get("versioned_buffers"); + auto bank_peer_buffers_anno = op->annotations.Get("bank_peer_buffers"); + auto version_axis_buffers_anno = + op->annotations.Get("version_axis_buffers"); auto prologue_orders_anno = op->annotations.Get("prologue_orders"); auto body_orders_anno = op->annotations.Get("body_orders"); auto epilogue_orders_anno = op->annotations.Get("epilogue_orders"); + auto dynamic_epilogue_orders_anno = + op->annotations.Get("dynamic_epilogue_orders"); + auto writer_phases_anno = op->annotations.Get("runtime_bank_writer_phases"); + auto reader_phases_anno = op->annotations.Get("runtime_bank_reader_phases"); if (!iterations_anno || !used_buffers_anno || !versioned_buffers_anno || !prologue_orders_anno || !body_orders_anno) { return for_node; } + arith::Analyzer extent_analyzer; + PrimExpr simplified_extent = extent_analyzer.Simplify(for_node->extent); + const auto *static_extent = simplified_extent.as(); + auto make_sequential_fallback = [&](const std::string &reason, + bool emit_warning = true) { + Map annotations; + for (const auto &kv : for_node->annotations) { + if (kv.first != "num_stages" && kv.first != "iterations" && + kv.first != "prologue_orders" && kv.first != "body_orders" && + kv.first != "epilogue_orders") { + annotations.Set(kv.first, kv.second); + } + } + For sequential = for_node; + sequential.CopyOnWrite()->annotations = annotations; + return MakePipelineFallback(sequential, "greedy", "inject", reason, + emit_warning); + }; // Step 2: Find the body and buffer allocations of the pipeline. The body // can be direct child of the for-loop. If the for-loop has BlockRealize as // its child, the pipeline body will be the child of the block. @@ -604,22 +879,58 @@ class SunmmioPipelineInjector : public StmtExprMutator { if (epilogue_orders_anno) { epilogue_orders = Downcast>(epilogue_orders_anno.value()); } + int max_body_iter_offset = 0; + for (const String &order : body_orders) { + max_body_iter_offset = std::max(max_body_iter_offset, name2iter(order)); + } + if (static_extent != nullptr && + static_extent->value <= max_body_iter_offset) { + return make_sequential_fallback("short_extent_unsupported"); + } Array versioned_buffers = Downcast>(versioned_buffers_anno.value()); Array used_buffers = Downcast>(used_buffers_anno.value()); + Map bank_peer_buffers; + if (bank_peer_buffers_anno) { + bank_peer_buffers = + Downcast>(bank_peer_buffers_anno.value()); + } + Array version_axis_buffers; + if (version_axis_buffers_anno) { + version_axis_buffers = + Downcast>(version_axis_buffers_anno.value()); + } + Map> writer_phases; + if (writer_phases_anno) { + writer_phases = Downcast>>( + writer_phases_anno.value()); + } + Map> reader_phases; + if (reader_phases_anno) { + reader_phases = Downcast>>( + reader_phases_anno.value()); + } for (auto it : used_buffers) { pipeline_allocs.push_back(it); } - auto rewriter = PipelineBodyRewriter(versioned_buffers, for_node); + auto rewriter = PipelineBodyRewriter(versioned_buffers, bank_peer_buffers, + version_axis_buffers, writer_phases, + reader_phases, for_node); + auto version_slot = [iterations](int iter) { + ICHECK_GT(iterations, 0); + int slot = iter % iterations; + return slot < 0 ? slot + iterations : slot; + }; Array for_body; // Step 3.1: Rewrite prologue for (const auto &order_str : prologue_orders) { int iter = name2iter(order_str); int id = name2id(order_str); Stmt stmt = pipeline_body_seq->seq[id]; - rewriter.set_current_version(iter); + rewriter.set_current_command(id); + rewriter.set_current_version(version_slot(iter)); PrimExpr replaced_loop_var = 0 + iter + for_node->min; rewriter.set_loop_var_replacement(replaced_loop_var); stmt = rewriter(stmt); @@ -632,12 +943,10 @@ class SunmmioPipelineInjector : public StmtExprMutator { int iter = name2iter(order_str); PrimExpr replaced_loop_var = iterations * for_node->loop_var + iter + for_node->min; - if (iter == iterations) { - iter = 0; - } int id = name2id(order_str); Stmt stmt = pipeline_body_seq->seq[id]; - rewriter.set_current_version(iter); + rewriter.set_current_command(id); + rewriter.set_current_version(version_slot(iter)); rewriter.set_loop_var_replacement(replaced_loop_var); stmt = rewriter(stmt); body.push_back(stmt); @@ -652,6 +961,8 @@ class SunmmioPipelineInjector : public StmtExprMutator { if (epilogue_iterations == 0) { extent = extent - 1; + } else if (epilogue_iterations == -1) { + extent = floordiv(max(0, for_node->extent - 1), iterations); } For new_for_stmt = For(for_node->loop_var, PrimExpr(0), extent, ForKind::kSerial, @@ -664,30 +975,62 @@ class SunmmioPipelineInjector : public StmtExprMutator { int iter = name2iter(order_str); int id = name2id(order_str); Stmt stmt = pipeline_body_seq->seq[id]; - rewriter.set_current_version(iter); + rewriter.set_current_command(id); + rewriter.set_current_version(version_slot(iter)); PrimExpr replaced_loop_var = extent * iterations + iter + for_node->min; rewriter.set_loop_var_replacement(replaced_loop_var); stmt = rewriter(stmt); for_body.push_back(stmt); } } else { - // Dynamic epilogue loop for non-constant iterations - Var epilogue_loop_var("epilogue_i", for_node->loop_var->dtype); - Array epilogue_body; - for (size_t id = 0; id < pipeline_body_seq->size(); ++id) { - Stmt stmt = pipeline_body_seq->seq[id]; - rewriter.set_current_version( - 0); // Versioning is not deeply supported in dynamic epilogue yet - PrimExpr replaced_loop_var = - extent * iterations + epilogue_loop_var + for_node->min; - rewriter.set_loop_var_replacement(replaced_loop_var); - stmt = rewriter(stmt); - epilogue_body.push_back(stmt); + ICHECK(dynamic_epilogue_orders_anno) + << "Dynamic pipeline requires remainder-specific epilogue orders"; + Map> dynamic_orders = + Downcast>>( + dynamic_epilogue_orders_anno.value()); + auto build_dynamic_epilogue = [&](const Array &schedule) { + Array epilogue_body; + for (const String &order_str : schedule) { + int iter = name2iter(order_str); + int id = name2id(order_str); + PrimExpr logical_iter = extent * iterations + iter; + PrimExpr replaced_loop_var = logical_iter + for_node->min; + Stmt stmt = pipeline_body_seq->seq[id]; + rewriter.set_current_command(id); + rewriter.set_current_version(version_slot(iter)); + rewriter.set_loop_var_replacement(replaced_loop_var); + stmt = rewriter(stmt); + PrimExpr valid = And(GE(logical_iter, Integer(0)), + LT(logical_iter, for_node->extent)); + epilogue_body.push_back(IfThenElse(valid, stmt)); + } + return SeqStmt::Flatten(epilogue_body); + }; + + PrimExpr runtime_remainder = + floormod(for_node->extent, Integer(iterations)); + Stmt dispatched_epilogue{nullptr}; + for (int remainder = iterations - 1; remainder >= 0; --remainder) { + Optional> schedule; + for (const auto &kv : dynamic_orders) { + if (kv.first->value == remainder) { + schedule = kv.second; + break; + } + } + ICHECK(schedule.defined()) + << "Missing dynamic epilogue schedule for remainder " << remainder; + Stmt branch = build_dynamic_epilogue(schedule.value()); + if (!dispatched_epilogue.defined()) { + dispatched_epilogue = std::move(branch); + } else { + dispatched_epilogue = + IfThenElse(EQ(runtime_remainder, Integer(remainder)), + std::move(branch), std::move(dispatched_epilogue)); + } } - For dynamic_epilogue_for = For( - epilogue_loop_var, PrimExpr(0), epilogue_iterations_expr, - ForKind::kSerial, SeqStmt::Flatten(epilogue_body), std::nullopt, {}); - for_body.push_back(dynamic_epilogue_for); + ICHECK(dispatched_epilogue.defined()); + for_body.push_back(dispatched_epilogue); } return SeqStmt::Flatten(for_body); } @@ -700,12 +1043,33 @@ class SunmmioPipelineInjector : public StmtExprMutator { tvm::transform::Pass InjectSunmmioPipeline() { using namespace tir::transform; auto pass_func = [=](PrimFunc f, const IRModule &m, PassContext ctx) { - Stmt multiversioned_body = SunmmioMultiVersionBufferRewriter::Substitute(f); - auto *fptr = f.CopyOnWrite(); - fptr->body = multiversioned_body; - fptr->body = SunmmioPipelineInjector::Inject(f); - fptr->body = ConvertSSA(std::move(fptr->body)); - return f; + const PrimFunc &original = f; + try { + PrimFunc candidate = f; + auto *fptr = candidate.CopyOnWrite(); + fptr->body = SunmmioMultiVersionBufferRewriter::Substitute(candidate); + fptr->body = SunmmioPipelineInjector::Inject(candidate); + fptr->body = ConvertSSA(std::move(fptr->body)); + Optional disallowed = + PipelineFallbackValidator::FindDisallowed(fptr->body); + if (disallowed) { + return MakePipelineFunctionFallback( + original, PipelineDiagnostic{false, "greedy", "inject_validation", + "candidate_fallback", + std::string(disallowed.value())}); + } + return candidate; + } catch (const std::exception &error) { + return MakePipelineFunctionFallback( + original, + PipelineDiagnostic{false, "greedy", "inject_exception", + "candidate_rewrite_failed", error.what()}); + } catch (...) { + return MakePipelineFunctionFallback( + original, + PipelineDiagnostic{false, "greedy", "inject_exception", + "candidate_rewrite_failed", "unknown exception"}); + } }; return CreatePrimFuncPass(pass_func, 0, "tl.InjectSunmmioPipeline", {}); } diff --git a/src/transform/inject_sunmmio_pipeline_ilp.cc b/src/transform/inject_sunmmio_pipeline_ilp.cc new file mode 100644 index 0000000000..c97ea20fae --- /dev/null +++ b/src/transform/inject_sunmmio_pipeline_ilp.cc @@ -0,0 +1,1584 @@ +#include "../layout/cute_layout.h" +#include "../layout/utils.h" +#include "../op/builtin.h" +#include "../op/copy.h" +#include "../op/parallel.h" +#include "../op/region.h" +#include "../op/utils.h" +#include "../target/utils.h" +#include "../tileview/tileview.h" +#include "common/ast_traverser.h" +#include "common/loop_fusion_utils.h" +#include "common/remap_buffer_rewriter.h" +#include "common/sunmmio_pipeline_utils.h" +#include "sunmmio_pipeline_planning/pipeline_diagnostic.h" +#include "tir/transforms/ir_utils.h" +#include "tvm/ir/attrs.h" +#include "tvm/ir/expr.h" +#include "tvm/node/cast.h" +#include "tvm/node/structural_equal.h" +#include "tvm/runtime/logging.h" +#include "tvm/tir/function.h" +#include "tvm/tir/stmt.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace tvm { +namespace tl { + +using namespace tir; + +struct LetWrapper { + Var var; + PrimExpr value; +}; + +int CeilDiv(int a, int b) { + ICHECK_GT(b, 0); + return (a + b - 1) / b; +} + +void AppendUniqueBuffer(Array *buffers, const Buffer &buffer) { + if (std::find(buffers->begin(), buffers->end(), buffer) == buffers->end()) { + buffers->push_back(buffer); + } +} + +Array +DeriveRuntimeMultiversionBuffers(const Optional &runtime_buffers_anno, + const Optional &versioned_buffers_anno, + const Array &banked_buffers, + int iterations) { + bool enable_banked_multiversion = iterations > 2; + if (runtime_buffers_anno) { + return Downcast>(runtime_buffers_anno.value()); + } + if (!versioned_buffers_anno) { + Array runtime_buffers; + if (enable_banked_multiversion) { + for (const Buffer &buffer : banked_buffers) { + AppendUniqueBuffer(&runtime_buffers, buffer); + } + } + return runtime_buffers; + } + + std::unordered_set banked; + for (const Buffer &buffer : banked_buffers) { + banked.insert(buffer.get()); + } + + Array runtime_buffers; + for (const Buffer &buffer : + Downcast>(versioned_buffers_anno.value())) { + runtime_buffers.push_back(buffer); + } + if (enable_banked_multiversion) { + for (const Buffer &buffer : banked_buffers) { + if (!banked.count(buffer.get())) { + continue; + } + AppendUniqueBuffer(&runtime_buffers, buffer); + } + } + return runtime_buffers; +} + +class SunmmioILPMultiVersionBufferRewriter : public StmtExprMutator { +public: + SunmmioILPMultiVersionBufferRewriter(const PrimFunc &f) { + for (const auto &kv : f->buffer_map) { + const Buffer &buffer = kv.second; + buffer_data_to_buffer_.Set(buffer->data, buffer); + } + } + + static Stmt Substitute(PrimFunc &f) { + SunmmioILPMultiVersionBufferRewriter substituter(f); + // collect used_buffers and iterations + substituter.VisitStmt(f->body); + substituter.replace_flag = true; + + for (auto &buffer : substituter.versioned_buffers_) { + int num_versions = substituter.GetVersionCount(buffer); + if (substituter.IsBankedBuffer(buffer)) { + Buffer ping_buffer = + substituter.makeRuntimeBuffer(buffer, num_versions, "_ping", true); + Buffer pong_buffer = + substituter.makeRuntimeBuffer(buffer, num_versions, "_pong", false); + substituter.buffer_remap_.Set(buffer, ping_buffer); + substituter.bank_peer_buffers_[buffer.get()] = pong_buffer; + } else { + substituter.buffer_remap_.Set( + buffer, substituter.makeRuntimeBuffer(buffer, num_versions)); + } + } + + substituter.RewriteFunctionLayoutAttrs(f); + substituter.RecordDefaultPingPongAttrs(f); + + f.CopyOnWrite()->body = + RemapBufferRewriter::Substitute(f->body, substituter.buffer_remap_); + + return substituter.VisitStmt(f->body); + } + +private: + void RecordDefaultPingPongAttrs(PrimFunc &f) { + if (buffer_remap_.empty()) { + return; + } + + Map alloc_ping_pong; + for (const auto &kv : bank_peer_buffers_) { + const Buffer &peer_buffer = kv.second; + alloc_ping_pong.Set(peer_buffer->data, String("pong")); + } + + if (alloc_ping_pong.empty()) { + return; + } + + f = WithAttr(std::move(f), tl::attr::kSunmmioAllocPingPong, + alloc_ping_pong); + } + + void RewriteFunctionLayoutAttrs(PrimFunc &f) { + auto layout_map_opt = f->GetAttr>(attr::kLayoutMap); + if (!layout_map_opt) { + return; + } + + arith::Analyzer analyzer; + Map new_layout_map; + for (const auto &[buffer, layout] : layout_map_opt.value()) { + auto it = buffer_remap_.find(buffer); + if (it == buffer_remap_.end()) { + new_layout_map.Set(buffer, layout); + continue; + } + + const Buffer &new_buffer = (*it).second; + Optional derived_layout = DeriveLayoutLike( + layout, new_buffer->shape, Optional>(), &analyzer); + ICHECK(derived_layout.defined()) + << "Failed to derive ILP multiversioned layout for buffer " + << buffer->name << " with shape " << new_buffer->shape; + new_layout_map.Set(new_buffer, derived_layout.value()); + auto peer_it = bank_peer_buffers_.find(buffer.get()); + if (peer_it != bank_peer_buffers_.end()) { + const Buffer &peer_buffer = peer_it->second; + Optional peer_layout = DeriveLayoutLike( + layout, peer_buffer->shape, Optional>(), &analyzer); + ICHECK(peer_layout.defined()) + << "Failed to derive ILP ping/pong layout for buffer " + << buffer->name << " with shape " << peer_buffer->shape; + new_layout_map.Set(peer_buffer, peer_layout.value()); + } + } + f = WithAttr(std::move(f), attr::kLayoutMap, new_layout_map); + } + + bool HasVersionAxis(const Buffer &buffer) const { + return version_axis_buffers_.count(buffer.get()) != 0; + } + + bool IsBankedBuffer(const Buffer &buffer) const { + return banked_buffers_.count(buffer.get()) != 0; + } + + int GetVersionCount(const Buffer &buffer) const { + auto it = buffer_versions_.find(buffer.get()); + int num_versions = it == buffer_versions_.end() ? 1 : it->second; + if (HasVersionAxis(buffer) && IsBankedBuffer(buffer)) { + return CeilDiv(num_versions, 2); + } + return num_versions; + } + + Array AddDefaultRuntimeAxes(const Buffer &buffer, + Array indices) const { + if (HasVersionAxis(buffer)) { + indices.insert(indices.begin(), Integer(0)); + } + return indices; + } + + Array AddDefaultRuntimeAxes(const Buffer &buffer, + Array ranges) const { + if (HasVersionAxis(buffer)) { + ranges.insert(ranges.begin(), Range::FromMinExtent(0, 1)); + } + return ranges; + } + + Buffer makeRuntimeBuffer(const Buffer &buffer, int num_version, + const std::string &name_suffix = "", + bool reuse_primary_var = true) { + const auto *ptr_type = + TVM_TYPE_AS(buffer->data->type_annotation, PointerTypeNode); + Var new_var; + std::string data_name = std::string(buffer->data->name_hint) + name_suffix; + std::string buffer_name = std::string(buffer->name) + name_suffix; + if (reuse_primary_var && var_remap_.count(buffer->data)) { + new_var = var_remap_[buffer->data]; + } else { + Type new_type = + PointerType(ptr_type->element_type, ptr_type->storage_scope); + new_var = Var(data_name, new_type); + if (reuse_primary_var) { + var_remap_.Set(buffer->data, new_var); + } + } + auto shape = buffer->shape; + if (HasVersionAxis(buffer)) { + shape.insert(shape.begin(), num_version); + } + return Buffer(new_var, buffer->dtype, shape, {}, buffer->elem_offset, + String(buffer_name), buffer->data_alignment, + buffer->offset_factor, buffer->buffer_type); + } + + BufferRegion + RewritePipelineBufferRegion(const BufferRegion &buffer_region) const { + auto it = buffer_remap_.find(buffer_region->buffer); + if (it != buffer_remap_.end()) { + Region new_region = buffer_region->region; + if (HasVersionAxis(buffer_region->buffer)) { + new_region.insert(new_region.begin(), Range::FromMinExtent(0, 1)); + } + const Buffer &new_buffer = (*it).second; + return BufferRegion(new_buffer, new_region); + } + return buffer_region; + } + + Stmt VisitStmt_(const ForNode *op) final { + For loop = Downcast(StmtExprMutator::VisitStmt_(op)); + auto runtime_buffers_anno = + op->annotations.Get("runtime_multiversion_buffers"); + auto versioned_buffers_anno = op->annotations.Get("versioned_buffers"); + auto banked_buffers_anno = op->annotations.Get("runtime_banked_buffers"); + auto resident_banked_buffers_anno = + op->annotations.Get("runtime_resident_banked_buffers"); + auto used_buffers_anno = op->annotations.Get("used_buffers"); + auto iterations_anno = op->annotations.Get("iterations"); + auto bank_start_phases_anno = + op->annotations.Get("runtime_bank_start_phases"); + auto bank_read_delta_parities_anno = + op->annotations.Get("runtime_bank_read_delta_parities"); + auto bank_writer_phases_anno = + op->annotations.Get("runtime_bank_writer_phases"); + auto bank_reader_phases_anno = + op->annotations.Get("runtime_bank_reader_phases"); + auto bank_flip_modes_anno = op->annotations.Get("runtime_bank_flip_modes"); + if (used_buffers_anno && iterations_anno && + (runtime_buffers_anno || versioned_buffers_anno)) { + Array banked_buffers; + if (banked_buffers_anno) { + banked_buffers = Downcast>(banked_buffers_anno.value()); + } + int iterations = Downcast(iterations_anno.value())->value; + Array runtime_buffers = DeriveRuntimeMultiversionBuffers( + runtime_buffers_anno, versioned_buffers_anno, banked_buffers, + iterations); + if (!replace_flag) { + for (const Buffer &buffer : runtime_buffers) { + AppendUniqueBuffer(&versioned_buffers_, buffer); + version_axis_buffers_.insert(buffer.get()); + int &num_versions = buffer_versions_[buffer.get()]; + num_versions = std::max(num_versions, iterations); + } + for (const Buffer &buffer : banked_buffers) { + AppendUniqueBuffer(&versioned_buffers_, buffer); + banked_buffers_.insert(buffer.get()); + } + } else { + Array new_runtime_buffers; + for (const Buffer &buffer : runtime_buffers) { + if (buffer_remap_.count(buffer)) { + new_runtime_buffers.push_back(buffer_remap_[buffer]); + } else { + new_runtime_buffers.push_back(buffer); + } + } + loop.CopyOnWrite()->annotations.Set("runtime_multiversion_buffers", + new_runtime_buffers); + if (versioned_buffers_anno) { + Array versioned_buffers = + Downcast>(versioned_buffers_anno.value()); + Array new_versioned_buffers; + for (const Buffer &buffer : versioned_buffers) { + if (buffer_remap_.count(buffer)) { + new_versioned_buffers.push_back(buffer_remap_[buffer]); + } else { + new_versioned_buffers.push_back(buffer); + } + } + loop.CopyOnWrite()->annotations.Set("versioned_buffers", + new_versioned_buffers); + } + if (banked_buffers_anno) { + Array banked_buffers = + Downcast>(banked_buffers_anno.value()); + Array new_banked_buffers; + for (const Buffer &buffer : banked_buffers) { + if (buffer_remap_.count(buffer)) { + new_banked_buffers.push_back(buffer_remap_[buffer]); + } else { + new_banked_buffers.push_back(buffer); + } + } + loop.CopyOnWrite()->annotations.Set("runtime_banked_buffers", + new_banked_buffers); + } + if (resident_banked_buffers_anno) { + Array resident_buffers = + Downcast>(resident_banked_buffers_anno.value()); + Array new_resident_buffers; + for (const Buffer &buffer : resident_buffers) { + auto it = buffer_remap_.find(buffer); + new_resident_buffers.push_back( + it == buffer_remap_.end() ? buffer : (*it).second); + } + loop.CopyOnWrite()->annotations.Set("runtime_resident_banked_buffers", + new_resident_buffers); + } + if (bank_start_phases_anno) { + Map bank_start_phases = + Downcast>(bank_start_phases_anno.value()); + Map new_bank_start_phases; + for (const auto &[buffer, phase] : bank_start_phases) { + if (buffer_remap_.count(buffer)) { + new_bank_start_phases.Set(buffer_remap_[buffer], phase); + } else { + new_bank_start_phases.Set(buffer, phase); + } + } + loop.CopyOnWrite()->annotations.Set("runtime_bank_start_phases", + new_bank_start_phases); + } + if (bank_read_delta_parities_anno) { + Map bank_read_delta_parities = + Downcast>( + bank_read_delta_parities_anno.value()); + Map new_bank_read_delta_parities; + for (const auto &[buffer, parity] : bank_read_delta_parities) { + if (buffer_remap_.count(buffer)) { + new_bank_read_delta_parities.Set(buffer_remap_[buffer], parity); + } else { + new_bank_read_delta_parities.Set(buffer, parity); + } + } + loop.CopyOnWrite()->annotations.Set( + "runtime_bank_read_delta_parities", new_bank_read_delta_parities); + } + if (bank_writer_phases_anno) { + Map> bank_writer_phases = + Downcast>>( + bank_writer_phases_anno.value()); + Map> new_bank_writer_phases; + for (const auto &[buffer, per_op] : bank_writer_phases) { + if (buffer_remap_.count(buffer)) { + new_bank_writer_phases.Set(buffer_remap_[buffer], per_op); + } else { + new_bank_writer_phases.Set(buffer, per_op); + } + } + loop.CopyOnWrite()->annotations.Set("runtime_bank_writer_phases", + new_bank_writer_phases); + } + if (bank_reader_phases_anno) { + Map> bank_reader_phases = + Downcast>>( + bank_reader_phases_anno.value()); + Map> new_bank_reader_phases; + for (const auto &[buffer, per_op] : bank_reader_phases) { + if (buffer_remap_.count(buffer)) { + new_bank_reader_phases.Set(buffer_remap_[buffer], per_op); + } else { + new_bank_reader_phases.Set(buffer, per_op); + } + } + loop.CopyOnWrite()->annotations.Set("runtime_bank_reader_phases", + new_bank_reader_phases); + } + if (bank_flip_modes_anno) { + Map flip_modes = + Downcast>(bank_flip_modes_anno.value()); + Map new_flip_modes; + for (const auto &[buffer, flip] : flip_modes) { + auto it = buffer_remap_.find(buffer); + new_flip_modes.Set( + it == buffer_remap_.end() ? buffer : (*it).second, flip); + } + loop.CopyOnWrite()->annotations.Set("runtime_bank_flip_modes", + new_flip_modes); + } + if (banked_buffers_anno) { + Array banked_buffers = + Downcast>(banked_buffers_anno.value()); + Map runtime_bank_peer_buffers; + for (const Buffer &buffer : banked_buffers) { + auto remap_it = buffer_remap_.find(buffer); + auto peer_it = bank_peer_buffers_.find(buffer.get()); + if (remap_it != buffer_remap_.end() && + peer_it != bank_peer_buffers_.end()) { + runtime_bank_peer_buffers.Set((*remap_it).second, + peer_it->second); + } + } + loop.CopyOnWrite()->annotations.Set("runtime_bank_peer_buffers", + runtime_bank_peer_buffers); + } + Array used_buffers = + Downcast>(used_buffers_anno.value()); + Array new_used_buffers; + for (const Buffer &buffer : used_buffers) { + if (buffer_remap_.count(buffer)) { + new_used_buffers.push_back(buffer_remap_[buffer]); + } else { + new_used_buffers.push_back(buffer); + } + } + loop.CopyOnWrite()->annotations.Set("used_buffers", new_used_buffers); + } + } + return loop; + } + + Stmt VisitStmt_(const BlockRealizeNode *op) final { + BlockRealize block_realize = + Downcast(StmtExprMutator::VisitStmt_(op)); + Block block = block_realize->block; + if (!replace_flag) { + for (const Buffer &alloc_buffer : block->alloc_buffers) { + buffer_data_to_buffer_.Set(alloc_buffer->data, alloc_buffer); + } + return block_realize; + } + + // do block attributes remap + if (block->annotations.count(attr::kLayoutMap)) { + auto map_anno = block->annotations.Get(attr::kLayoutMap); + Map map = Downcast>(map_anno.value()); + Map new_map; + for (const auto &[buffer, layout] : map) { + if (buffer_remap_.count(buffer)) { + new_map.Set(buffer_remap_[buffer], layout); + auto peer_it = bank_peer_buffers_.find(buffer.get()); + if (peer_it != bank_peer_buffers_.end()) { + new_map.Set(peer_it->second, layout); + } + } else { + new_map.Set(buffer, layout); + } + } + block.CopyOnWrite()->annotations.Set(attr::kLayoutMap, new_map); + } + + if (block->annotations.count(attr::kTileViewMap)) { + auto map = block->annotations.Get(attr::kTileViewMap) + ->as>() + .value(); + Map new_map; + for (const auto &[var, tileView] : map) { + if (var_remap_.count(var)) { + new_map.Set(var_remap_[var], tileView); + } else { + new_map.Set(var, tileView); + } + } + block.CopyOnWrite()->annotations.Set(attr::kTileViewMap, new_map); + } + + block.CopyOnWrite()->reads.MutateByApply( + [this](const BufferRegion &buffer_region) { + return RewritePipelineBufferRegion(buffer_region); + }); + block.CopyOnWrite()->writes.MutateByApply( + [this](const BufferRegion &buffer_region) { + return RewritePipelineBufferRegion(buffer_region); + }); + + // do block->alloc_buffers remap + Array alloc_buffers; + for (const Buffer &buf : block->alloc_buffers) { + auto remap_it = buffer_remap_.find(buf); + if (remap_it != buffer_remap_.end()) { + alloc_buffers.push_back((*remap_it).second); + auto peer_it = bank_peer_buffers_.find(buf.get()); + if (peer_it != bank_peer_buffers_.end()) { + alloc_buffers.push_back(peer_it->second); + } + } else { + alloc_buffers.push_back(buf); + } + } + + if (!alloc_buffers.same_as(block->alloc_buffers)) { + block.CopyOnWrite()->alloc_buffers = alloc_buffers; + } + block_realize.CopyOnWrite()->block = block; + + return block_realize; + } + + PrimExpr VisitExpr_(const BufferLoadNode *op) final { + auto load = Downcast(StmtExprMutator::VisitExpr_(op)); + if (!replace_flag) { + return load; + } + auto buffer = load->buffer; + if (buffer_remap_.count(buffer)) { + auto new_buffer = buffer_remap_[load->buffer]; + auto indices = AddDefaultRuntimeAxes(buffer, load->indices); + return BufferLoad(new_buffer, indices); + } + auto expr = StmtExprMutator::VisitExpr_(op); + return expr; + } + + Stmt VisitStmt_(const BufferStoreNode *op) final { + auto store = Downcast(StmtExprMutator::VisitStmt_(op)); + if (!replace_flag) { + return store; + } + auto buffer = store->buffer; + if (buffer_remap_.count(buffer)) { + auto new_buffer = buffer_remap_[store->buffer]; + auto indices = AddDefaultRuntimeAxes(buffer, store->indices); + return BufferStore(new_buffer, store->value, indices); + } + return store; + } + + PrimExpr VisitExpr_(const CallNode *op) final { + if (!replace_flag) + return StmtExprMutator::VisitExpr_(op); + if (op->op.same_as(builtin::tvm_access_ptr())) { + ICHECK_EQ(op->args.size(), 5U); + Var buffer_data = Downcast(op->args[1]); + if (!var_remap_.count(buffer_data)) { + return StmtExprMutator::VisitExpr_(op); + } + Var new_data = var_remap_[buffer_data]; + return Call( + op->dtype, op->op, + {op->args[0], new_data, op->args[2], op->args[3], op->args[4]}); + } else if (op->op.same_as(RegionOp::Get())) { + RegionOp original_region(op->args); + Buffer original_buffer = original_region->GetBuffer(); + + if (!buffer_remap_.count(original_buffer)) { + return StmtExprMutator::VisitExpr_(op); + } + + Buffer new_buffer = buffer_remap_[original_buffer]; + Array new_ranges = + AddDefaultRuntimeAxes(original_buffer, original_region->GetRanges()); + + Array new_args; + new_args.push_back(BufferLoad(new_buffer, [new_ranges]() { + Array mins; + for (auto r : new_ranges) { + mins.push_back(r->min); + } + return mins; + }())); + new_args.push_back(original_region->GetAccessMask()); + for (auto r : new_ranges) { + new_args.push_back(r->extent); + } + + return Call(DataType::Handle(), RegionOp::Get(), new_args); + } + auto expr = StmtExprMutator::VisitExpr_(op); + return expr; + } + + PrimExpr VisitExpr_(const VarNode *op) final { + Var var = tvm::ffi::GetRef(op); + if (!replace_flag) { + return std::move(var); + } + if (var_remap_.count(var)) { + auto new_var = var_remap_[var]; + return std::move(new_var); + } + return std::move(var); + } + + Array versioned_buffers_; + int iterations_ = -1; + bool replace_flag = false; + Map buffer_remap_; + Map var_remap_; + Map buffer_data_to_buffer_; + std::unordered_map buffer_versions_; + std::unordered_set version_axis_buffers_; + std::unordered_set banked_buffers_; + std::unordered_map bank_peer_buffers_; +}; + +class SunmmioILPPipelineBodyRewriter : public StmtExprMutator { +public: + SunmmioILPPipelineBodyRewriter( + Array runtime_buffers, Array version_axis_buffers, + Array banked_buffers, Map bank_peer_buffers, + Map bank_start_phases, + Map bank_read_delta_parities, + Map> bank_writer_phases, + Map> bank_reader_phases, + Map bank_flip_modes, For pipeline_loop, + int iterations) { + pipeline_loop_ = std::move(pipeline_loop); + iterations_ = iterations; + for (const Buffer &it : runtime_buffers) { + rewritten_buffers_.insert(it.get()); + buffer_data_to_buffer_.Set(it->data, it); + } + for (const Buffer &it : version_axis_buffers) { + version_axis_buffers_.insert(it.get()); + } + for (const Buffer &it : banked_buffers) { + banked_buffers_.insert(it.get()); + } + for (const auto &[buffer, peer] : bank_peer_buffers) { + bank_peer_buffers_[buffer.get()] = peer; + buffer_data_to_buffer_.Set(peer->data, peer); + } + for (const auto &[buffer, phase] : bank_start_phases) { + if (const auto *imm = phase.as()) { + buffer_bank_start_phase_[buffer.get()] = imm->value; + } + } + for (const auto &[buffer, parity] : bank_read_delta_parities) { + if (const auto *imm = parity.as()) { + buffer_read_delta_parity_[buffer.get()] = imm->value & 1; + } + } + for (const auto &[buffer, per_op] : bank_writer_phases) { + auto &dst = buffer_writer_phase_[buffer.get()]; + for (const auto &[op_id, phase] : per_op) { + if (const auto *op_imm = op_id.as()) { + if (const auto *phase_imm = phase.as()) { + dst[op_imm->value] = phase_imm->value; + } + } + } + } + for (const auto &[buffer, per_op] : bank_reader_phases) { + auto &dst = buffer_reader_phase_[buffer.get()]; + for (const auto &[op_id, phase] : per_op) { + if (const auto *op_imm = op_id.as()) { + if (const auto *phase_imm = phase.as()) { + dst[op_imm->value] = phase_imm->value; + } + } + } + } + for (const auto &[buffer, flip] : bank_flip_modes) { + if (const auto *imm = flip.as()) { + buffer_bank_flip_[buffer.get()] = imm->value != 0; + } + } + } + + void set_current_stmt_id(int stmt_id) { current_stmt_id_ = stmt_id; } + + void set_loop_var_replacement(PrimExpr p) { replaced_loop_var_ = p; } + void set_logical_iter_parity_override(int parity) { + logical_iter_parity_override_ = parity; + } + void clear_logical_iter_parity_override() { + logical_iter_parity_override_ = -1; + } + void set_pipeline_loop_parity_override(int parity) { + pipeline_loop_parity_override_ = parity; + } + void clear_pipeline_loop_parity_override() { + pipeline_loop_parity_override_ = -1; + } + void clear_parity_overrides() { + clear_logical_iter_parity_override(); + clear_pipeline_loop_parity_override(); + } + + void clear_current_stmt_id() { current_stmt_id_ = -1; } + +private: + int VersionAxis(const Buffer &buffer) const { + return version_axis_buffers_.count(buffer.get()) ? 0 : -1; + } + + PrimExpr LogicalIterExpr() const { + return replaced_loop_var_ - pipeline_loop_->min; + } + + PrimExpr EffectiveVersionExpr(const Buffer &buffer, + int access_iter_offset = 0) const { + if (!version_axis_buffers_.count(buffer.get())) { + return Integer(0); + } + PrimExpr logical_iter = LogicalIterExpr() + access_iter_offset; + if (banked_buffers_.count(buffer.get())) { + auto it_flip = buffer_bank_flip_.find(buffer.get()); + bool flip = it_flip == buffer_bank_flip_.end() || it_flip->second; + int num_versions = CeilDiv(iterations_, 2); + if (!flip) { + return floormod(logical_iter, Integer(num_versions)); + } + return floormod(floordiv(logical_iter, Integer(2)), + Integer(num_versions)); + } + return floormod(logical_iter, Integer(iterations_)); + } + + int ResolveLogicalIterParity() const { + if (logical_iter_parity_override_ >= 0) { + return logical_iter_parity_override_; + } + if (pipeline_loop_parity_override_ >= 0) { + arith::Analyzer analyzer; + PrimExpr iter_offset = + analyzer.Simplify(LogicalIterExpr() - pipeline_loop_->loop_var); + if (const auto *imm = iter_offset.as()) { + int parity = (pipeline_loop_parity_override_ + imm->value) % 2; + return parity < 0 ? parity + 2 : parity; + } + } + arith::Analyzer analyzer; + PrimExpr simplified = analyzer.Simplify(LogicalIterExpr()); + if (const auto *imm = simplified.as()) { + int parity = imm->value % 2; + return parity < 0 ? parity + 2 : parity; + } + return -1; + } + + Buffer ResolveTargetBuffer(const Buffer &buffer, bool is_read = false) const { + if (!banked_buffers_.count(buffer.get())) { + return buffer; + } + // Bank annotations store a phase offset. Flipping buffers XOR it with the + // logical iteration parity; non-flipping buffers use the offset directly. + auto it_flip = buffer_bank_flip_.find(buffer.get()); + bool flip = it_flip == buffer_bank_flip_.end() || it_flip->second; + int iter_phase = 0; + if (flip) { + int logical_iter_parity = ResolveLogicalIterParity(); + ICHECK_GE(logical_iter_parity, 0) + << "Dynamic bank parity must resolve before selecting ping/pong for " + << buffer->name; + iter_phase = logical_iter_parity; + } + int bank = -1; + if (current_stmt_id_ >= 0) { + if (is_read) { + auto it_buf = buffer_reader_phase_.find(buffer.get()); + if (it_buf != buffer_reader_phase_.end()) { + auto it_stmt = it_buf->second.find(current_stmt_id_); + if (it_stmt != it_buf->second.end()) { + bank = (iter_phase + it_stmt->second) % 2; + } + } + } else { + auto it_buf = buffer_writer_phase_.find(buffer.get()); + if (it_buf != buffer_writer_phase_.end()) { + auto it_stmt = it_buf->second.find(current_stmt_id_); + if (it_stmt != it_buf->second.end()) { + bank = (iter_phase + it_stmt->second) % 2; + } + } + } + } + if (bank < 0) { + int start_phase = 0; + auto it = buffer_bank_start_phase_.find(buffer.get()); + if (it != buffer_bank_start_phase_.end()) { + start_phase = it->second; + } + int read_delta_parity = 0; + if (is_read) { + auto it_delta = buffer_read_delta_parity_.find(buffer.get()); + if (it_delta != buffer_read_delta_parity_.end()) { + read_delta_parity = it_delta->second; + } + } + bank = (iter_phase + start_phase + read_delta_parity) % 2; + } + if (bank < 0) { + bank += 2; + } + if (bank == 0) { + return buffer; + } + auto peer_it = bank_peer_buffers_.find(buffer.get()); + ICHECK(peer_it != bank_peer_buffers_.end()) + << "Missing peer buffer for banked runtime buffer " << buffer->name; + return peer_it->second; + } + + PrimExpr RewriteBufferAccess(const Call &call, + const std::vector &arg_indices) { + auto product = [](const Array &input) { + return foldl( + [](PrimExpr a, PrimExpr b, Span span) { + return mul(std::move(a), std::move(b), std::move(span)); + }, + make_const(DataType::Int(32), 1), input); + }; + auto axis_stride = [&](const Buffer &buffer, int axis) { + if (!buffer->strides.empty()) { + return buffer->strides[axis]; + } + Array suffix; + for (size_t j = axis + 1; j < buffer->shape.size(); ++j) { + suffix.push_back(buffer->shape[j]); + } + return product(suffix); + }; + Array new_args = call->args; + for (int i : arg_indices) { + Var buffer_data = Downcast(call->args[i]); + if (!buffer_data_to_buffer_.count(buffer_data)) { + continue; + } + const Buffer &buffer = buffer_data_to_buffer_[buffer_data]; + if (!rewritten_buffers_.count(buffer.get())) { + continue; + } + ICHECK_GT(call->args.size(), static_cast(i + 3)); + int access_mask = 2; + if (const auto *imm = call->args[i + 3].as()) { + access_mask = imm->value; + } + ICHECK_NE(access_mask & 3, 0) + << "tvm_access_ptr must carry a read/write access mask"; + // Read-write accesses select the writer version. For a banked + // read-write operation its reader and writer offsets must agree. + bool is_read = (access_mask & 2) == 0; + Buffer target_buffer = ResolveTargetBuffer(buffer, is_read); + PrimExpr new_index = call->args[i + 1]; + int version_axis = VersionAxis(buffer); + if (version_axis >= 0) { + PrimExpr offset = axis_stride(target_buffer, version_axis); + new_index = new_index + EffectiveVersionExpr(buffer) * offset; + } + new_args.Set(i, target_buffer->data); + new_args.Set(i + 1, new_index); + } + return Call(call->dtype, call->op, new_args, call->annotations, call->span); + } + + PrimExpr RewriteRegionExpr(const Call &call) { + RegionOp original_region(call->args); + Buffer original_buffer = original_region->GetBuffer(); + arith::Analyzer analyzer; + int access_iter_offset = DetectPipelineIterOffsetFromRegion( + BufferRegion(original_buffer, original_region->GetRanges()), + pipeline_loop_->loop_var, &analyzer); + Buffer target_buffer = original_buffer; + if (rewritten_buffers_.count(original_buffer.get())) { + int access_mask = original_region->GetAccessMask(); + ICHECK_NE(access_mask & 3, 0) + << "tl.region must carry a read/write access mask"; + // Region masks use 1=read, 2=write, 3=read-write. A destination region + // must use the per-op writer offset; treating it as a reader can map the + // producer and consumer of one logical value to different ping/pong + // banks. + bool is_read = (access_mask & 2) == 0; + target_buffer = ResolveTargetBuffer(original_buffer, is_read); + } + Array new_ranges; + for (const Range &range : original_region->GetRanges()) { + new_ranges.push_back(Range::FromMinExtent(VisitExpr(range->min), + VisitExpr(range->extent))); + } + int version_axis = VersionAxis(original_buffer); + if (version_axis >= 0) { + new_ranges.Set( + version_axis, + Range::FromMinExtent( + EffectiveVersionExpr(original_buffer, access_iter_offset), 1)); + } + return MakeRegionExpr(target_buffer, new_ranges, + original_region->GetAccessMask()); + } + + Stmt VisitStmt_(const BlockNode *op) final { + Block block = Downcast(StmtExprMutator::VisitStmt_(op)); + BlockNode *n = block.CopyOnWrite(); + // n->reads.MutateByApply([this](const BufferRegion &buffer_region) { + // return RewritePipelineBufferRegion(buffer_region); + // }); + // n->writes.MutateByApply([this](const BufferRegion &buffer_region) { + // return RewritePipelineBufferRegion(buffer_region); + // }); + return block; + } + + Stmt VisitStmt_(const BufferStoreNode *op) final { + Array original_ranges; + for (const PrimExpr &index : op->indices) { + original_ranges.push_back(Range::FromMinExtent(index, 1)); + } + arith::Analyzer analyzer; + int access_iter_offset = DetectPipelineIterOffsetFromRegion( + BufferRegion(op->buffer, original_ranges), pipeline_loop_->loop_var, + &analyzer); + int prev_stmt_id = current_stmt_id_; + current_stmt_id_ = logical_stmt_cursor_; + logical_stmt_cursor_ += 1; + BufferStore store = Downcast(StmtExprMutator::VisitStmt_(op)); + current_stmt_id_ = prev_stmt_id; + if (!rewritten_buffers_.count(store->buffer.get())) { + return store; + } + Buffer target_buffer = + ResolveTargetBuffer(store->buffer, /*is_read=*/false); + Array indices = store->indices; + int version_axis = VersionAxis(store->buffer); + if (version_axis >= 0) { + indices.Set(version_axis, + EffectiveVersionExpr(store->buffer, access_iter_offset)); + } + return BufferStore(target_buffer, store->value, indices); + } + + PrimExpr VisitExpr_(const BufferLoadNode *op) final { + Array original_ranges; + for (const PrimExpr &index : op->indices) { + if (const auto *ramp = index.as()) { + original_ranges.push_back( + Range::FromMinExtent(ramp->base, ramp->lanes)); + } else { + original_ranges.push_back(Range::FromMinExtent(index, 1)); + } + } + arith::Analyzer analyzer; + int access_iter_offset = DetectPipelineIterOffsetFromRegion( + BufferRegion(op->buffer, original_ranges), pipeline_loop_->loop_var, + &analyzer); + int prev_stmt_id = current_stmt_id_; + current_stmt_id_ = logical_stmt_cursor_; + logical_stmt_cursor_ += 1; + BufferLoad load = Downcast(StmtExprMutator::VisitExpr_(op)); + current_stmt_id_ = prev_stmt_id; + if (!rewritten_buffers_.count(load->buffer.get())) { + return load; + } + Buffer target_buffer = ResolveTargetBuffer(load->buffer, /*is_read=*/true); + Array indices = load->indices; + int version_axis = VersionAxis(load->buffer); + if (version_axis >= 0) { + indices.Set(version_axis, + EffectiveVersionExpr(load->buffer, access_iter_offset)); + } + return BufferLoad(target_buffer, indices); + } + + PrimExpr VisitExpr_(const CallNode *op) final { + // A RegionOp encodes its buffer as a transport-only BufferLoad. Rewrite + // the region as a unit before generic recursion, otherwise that carrier + // BufferLoad selects a bank once and RewriteRegionExpr selects it again. + if (op->op.same_as(RegionOp::Get())) { + return RewriteRegionExpr(tvm::ffi::GetRef(op)); + } + Call call = Downcast(StmtExprMutator::VisitExpr_(op)); + if (call->op.same_as(builtin::tvm_access_ptr())) { + return RewriteBufferAccess(call, {1}); + } + return call; + } + + PrimExpr VisitExpr_(const VarNode *op) final { + Var var = Downcast(StmtExprMutator::VisitExpr_(op)); + if (var.same_as(pipeline_loop_->loop_var)) { + return replaced_loop_var_; + } + return var; + } + + std::unordered_set rewritten_buffers_; + std::unordered_set version_axis_buffers_; + std::unordered_set banked_buffers_; + std::unordered_map bank_peer_buffers_; + std::unordered_map buffer_bank_start_phase_; + std::unordered_map buffer_bank_flip_; + std::unordered_map buffer_read_delta_parity_; + std::unordered_map> + buffer_writer_phase_; + std::unordered_map> + buffer_reader_phase_; + Map buffer_data_to_buffer_; + For pipeline_loop_; + int iterations_ = 1; + PrimExpr replaced_loop_var_; + int logical_iter_parity_override_ = -1; + int pipeline_loop_parity_override_ = -1; + int current_stmt_id_ = -1; + int logical_stmt_cursor_ = 0; +}; + +class SunmmioILPPipelineInjector : public StmtExprMutator { +public: + static Stmt Inject(const PrimFunc &func) { + auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol); + SunmmioILPPipelineInjector injector(global_symbol, func); + for (const auto &kv : func->buffer_map) { + const Buffer &buffer = kv.second; + injector.buffer_data_to_buffer_.Set(buffer->data, buffer); + } + return injector(func->body); + } + +private: + explicit SunmmioILPPipelineInjector(Optional global_symbol, + const PrimFunc &f) + : global_symbol_(std::move(global_symbol)), traverser_(f) { + traverser_.clear(); + } + + Stmt VisitStmt_(const ForNode *op) final { + // Step 1: Recursively rewrite the children first. + For for_node = Downcast(StmtExprMutator::VisitStmt_(op)); + + auto iterations_anno = op->annotations.Get("iterations"); + auto ii_anno = op->annotations.Get("ii"); + auto makespan_anno = op->annotations.Get("makespan"); + auto steady_state_max_iter_offset_anno = + op->annotations.Get("steady_state_max_iter_offset"); + auto used_buffers_anno = op->annotations.Get("used_buffers"); + auto runtime_buffers_anno = + op->annotations.Get("runtime_multiversion_buffers"); + auto versioned_buffers_anno = op->annotations.Get("versioned_buffers"); + auto banked_buffers_anno = op->annotations.Get("runtime_banked_buffers"); + auto bank_start_phases_anno = + op->annotations.Get("runtime_bank_start_phases"); + auto bank_read_delta_parities_anno = + op->annotations.Get("runtime_bank_read_delta_parities"); + auto bank_writer_phases_anno = + op->annotations.Get("runtime_bank_writer_phases"); + auto bank_reader_phases_anno = + op->annotations.Get("runtime_bank_reader_phases"); + auto bank_flip_modes_anno = op->annotations.Get("runtime_bank_flip_modes"); + auto bank_peer_buffers_anno = + op->annotations.Get("runtime_bank_peer_buffers"); + auto prologue_orders_anno = op->annotations.Get("prologue_orders"); + auto body_orders_anno = op->annotations.Get("body_orders"); + auto epilogue_orders_anno = op->annotations.Get("epilogue_orders"); + + if (!iterations_anno || !used_buffers_anno || !versioned_buffers_anno || + !prologue_orders_anno || !body_orders_anno || !epilogue_orders_anno) { + return for_node; + } + + arith::Analyzer extent_analyzer; + PrimExpr simplified_extent = extent_analyzer.Simplify(for_node->extent); + const auto *static_extent = simplified_extent.as(); + auto make_sequential_fallback = [&](const std::string &reason, + bool emit_warning = true) { + Map annotations; + for (const auto &kv : for_node->annotations) { + if (kv.first != "num_stages" && kv.first != "iterations" && + kv.first != "prologue_orders" && kv.first != "body_orders" && + kv.first != "epilogue_orders") { + annotations.Set(kv.first, kv.second); + } + } + For sequential = for_node; + sequential.CopyOnWrite()->annotations = annotations; + return MakePipelineFallback(sequential, "ilp", "inject", reason, + emit_warning); + }; + + // Step 2: Find the body and buffer allocations of the pipeline. The body + // can be direct child of the for-loop. If the for-loop has BlockRealize as + // its child, the pipeline body will be the child of the block. + Stmt pipeline_body_root{nullptr}; + bool pipeline_body_from_block = false; + Array pipeline_allocs; + if (const auto *realize = for_node->body.as()) { + const auto &block = realize->block; + for (const auto &buffer : block->alloc_buffers) { + ICHECK(buffer->IsInstance()); + buffer_data_to_buffer_.Set(buffer->data, buffer); + } + pipeline_body_root = block->body; + pipeline_allocs = block->alloc_buffers; + pipeline_body_from_block = true; + } else { + pipeline_body_root = for_node->body; + } + + const SeqStmtNode *pipeline_body_seq = nullptr; + std::vector> rewrap_fns; + std::vector loop_var_let_wrappers; + auto append_attr_wrapper = [&rewrap_fns](const AttrStmtNode *attr) { + Any node = attr->node; + String attr_key = attr->attr_key; + PrimExpr value = attr->value; + Span span = attr->span; + rewrap_fns.emplace_back( + [node = std::move(node), attr_key = std::move(attr_key), + value = std::move(value), span](Stmt body) -> Stmt { + return AttrStmt(node, attr_key, value, body, span); + }); + }; + { + Stmt current = pipeline_body_root; + while (true) { + if (const auto *seq_stmt = current.as()) { + pipeline_body_seq = seq_stmt; + break; + } + if (const auto *if_then_else = current.as()) { + ICHECK(!if_then_else->else_case.defined()) + << "InjectSoftwarePipeline: Can't handle the body of the loop " + "because the IfThenElse node has an else branch"; + PrimExpr condition = if_then_else->condition; + Span span = if_then_else->span; + rewrap_fns.emplace_back( + [condition = std::move(condition), span](Stmt body) -> Stmt { + return IfThenElse(condition, body, Stmt(), span); + }); + current = if_then_else->then_case; + continue; + } + if (const auto *let_stmt = current.as()) { + // If this Let value uses the pipeline loop var, record it and push + // inside each rewritten block later so the loop var can be + // substituted with the correct per-iteration index. Otherwise, keep + // it as a normal wrapper. + bool uses_loop_var = UsesVar( + let_stmt->value, + [v = op->loop_var.get()](const VarNode *vn) { return vn == v; }); + if (uses_loop_var) { + loop_var_let_wrappers.push_back({let_stmt->var, let_stmt->value}); + } else { + Var var = let_stmt->var; + PrimExpr value = let_stmt->value; + Span span = let_stmt->span; + rewrap_fns.emplace_back([var = std::move(var), + value = std::move(value), + span](Stmt body) -> Stmt { + return LetStmt(var, value, body, span); + }); + } + current = let_stmt->body; + continue; + } + if (const auto *attr = current.as()) { + append_attr_wrapper(attr); + current = attr->body; + continue; + } + LOG(FATAL) << "ValueError: The body of the software pipeline should be " + << "SeqStmt, got " << current->GetTypeKey(); + } + } + ICHECK(pipeline_body_seq != nullptr); + + // Step 3: Rewrite the body of loop. + int iterations = Downcast(iterations_anno.value())->value; + int ii = ii_anno ? Downcast(ii_anno.value())->value : iterations; + int makespan = + makespan_anno ? Downcast(makespan_anno.value())->value : -1; + int steady_state_max_iter_offset = + steady_state_max_iter_offset_anno + ? Downcast(steady_state_max_iter_offset_anno.value())->value + : 0; + Array prologue_orders = + Downcast>(prologue_orders_anno.value()); + Array body_orders = + Downcast>(body_orders_anno.value()); + Array epilogue_orders = + Downcast>(epilogue_orders_anno.value()); + int max_logical_iter_offset = 0; + for (const String &order : body_orders) { + max_logical_iter_offset = + std::max(max_logical_iter_offset, name2iter(order)); + } + if (static_extent != nullptr && + static_extent->value <= max_logical_iter_offset) { + return make_sequential_fallback("short_extent_unsupported"); + } + Array versioned_buffers = + Downcast>(versioned_buffers_anno.value()); + Array used_buffers = + Downcast>(used_buffers_anno.value()); + for (auto it : used_buffers) { + pipeline_allocs.push_back(it); + } + Array banked_buffers; + if (banked_buffers_anno) { + banked_buffers = Downcast>(banked_buffers_anno.value()); + } + Array runtime_buffers = DeriveRuntimeMultiversionBuffers( + runtime_buffers_anno, versioned_buffers_anno, banked_buffers, + iterations); + Array rewritten_buffers = runtime_buffers; + for (const Buffer &buffer : banked_buffers) { + AppendUniqueBuffer(&rewritten_buffers, buffer); + } + Map bank_start_phases; + if (bank_start_phases_anno) { + bank_start_phases = + Downcast>(bank_start_phases_anno.value()); + } + Map bank_read_delta_parities; + if (bank_read_delta_parities_anno) { + bank_read_delta_parities = Downcast>( + bank_read_delta_parities_anno.value()); + } + Map> bank_writer_phases; + if (bank_writer_phases_anno) { + bank_writer_phases = Downcast>>( + bank_writer_phases_anno.value()); + } + Map> bank_reader_phases; + if (bank_reader_phases_anno) { + bank_reader_phases = Downcast>>( + bank_reader_phases_anno.value()); + } + Map bank_flip_modes; + if (bank_flip_modes_anno) { + bank_flip_modes = + Downcast>(bank_flip_modes_anno.value()); + } + Map bank_peer_buffers; + if (bank_peer_buffers_anno) { + bank_peer_buffers = + Downcast>(bank_peer_buffers_anno.value()); + } + bool requires_parity_specialization = false; + for (const Buffer &buffer : banked_buffers) { + auto it = bank_flip_modes.find(buffer); + if (it == bank_flip_modes.end()) { + requires_parity_specialization = true; + break; + } + const auto *imm = (*it).second.as(); + if (imm == nullptr || imm->value != 0) { + requires_parity_specialization = true; + break; + } + } + + auto rewriter = SunmmioILPPipelineBodyRewriter( + rewritten_buffers, runtime_buffers, banked_buffers, bank_peer_buffers, + bank_start_phases, bank_read_delta_parities, bank_writer_phases, + bank_reader_phases, bank_flip_modes, for_node, iterations); + arith::Analyzer analyzer; + auto rewrite_stmt_with_logical_iter_parity = + [&](const Stmt &stmt, int stmt_id, const PrimExpr &replaced_loop_var, + int parity) -> Stmt { + rewriter.set_current_stmt_id(stmt_id); + rewriter.set_loop_var_replacement(replaced_loop_var); + rewriter.clear_pipeline_loop_parity_override(); + if (parity >= 0) { + rewriter.set_logical_iter_parity_override(parity); + } else { + rewriter.clear_logical_iter_parity_override(); + } + Stmt rewritten = rewriter(stmt); + rewriter.clear_current_stmt_id(); + rewriter.clear_parity_overrides(); + return rewritten; + }; + auto rewrite_stmt = [&](const Stmt &stmt, int stmt_id, + const PrimExpr &replaced_loop_var) -> Stmt { + if (!requires_parity_specialization) { + return rewrite_stmt_with_logical_iter_parity(stmt, stmt_id, + replaced_loop_var, -1); + } + + PrimExpr logical_iter = + analyzer.Simplify(replaced_loop_var - for_node->min); + if (logical_iter.as()) { + return rewrite_stmt_with_logical_iter_parity(stmt, stmt_id, + replaced_loop_var, -1); + } + + Stmt even_stmt = rewrite_stmt_with_logical_iter_parity( + stmt, stmt_id, replaced_loop_var, 0); + Stmt odd_stmt = rewrite_stmt_with_logical_iter_parity( + stmt, stmt_id, replaced_loop_var, 1); + + PrimExpr is_even = EQ( + floormod(replaced_loop_var - for_node->min, Integer(2)), Integer(0)); + return IfThenElse(is_even, even_stmt, odd_stmt); + }; + // A ping/pong schedule has period two. Materialize two consecutive + // steady-state bases in one super-iteration so every bank choice is fixed + // at compile time. In particular, do not put the same async producer's + // consumers in runtime even/odd branches: NPU-IR requires each token to + // have exactly one static wait consumer. + PrimExpr steady_count = + max(0, for_node->extent - steady_state_max_iter_offset); + PrimExpr super_count = floordiv(steady_count, Integer(2)); + + auto rewrite_body_at_base = [&](const PrimExpr &base, + int base_parity) -> Array { + Array result; + for (const auto &order_str : body_orders) { + int iter_offset = name2iter(order_str); + int id = name2id(order_str); + PrimExpr replaced_loop_var = base + iter_offset + for_node->min; + Stmt stmt = pipeline_body_seq->seq[id]; + int logical_iter_parity = (base_parity + iter_offset) % 2; + if (logical_iter_parity < 0) { + logical_iter_parity += 2; + } + result.push_back(rewrite_stmt_with_logical_iter_parity( + stmt, id, replaced_loop_var, logical_iter_parity)); + } + return result; + }; + + auto build_pipeline_variant = [&](bool has_steady_tail) -> Stmt { + Array variant; + + // Prologue logical iterations are constants, so their physical banks are + // independent of the runtime extent parity. + for (const auto &order_str : prologue_orders) { + int iter = name2iter(order_str); + if (iter < 0 || + (static_extent != nullptr && iter >= static_extent->value)) { + continue; + } + int id = name2id(order_str); + PrimExpr replaced_loop_var = iter + for_node->min; + variant.push_back( + rewrite_stmt(pipeline_body_seq->seq[id], id, replaced_loop_var)); + } + + if (!requires_parity_specialization) { + Array steady_body; + for (const auto &order_str : body_orders) { + int iter_offset = name2iter(order_str); + int id = name2id(order_str); + PrimExpr replaced_loop_var = + for_node->loop_var + iter_offset + for_node->min; + steady_body.push_back( + rewrite_stmt(pipeline_body_seq->seq[id], id, replaced_loop_var)); + } + variant.push_back(For(for_node->loop_var, PrimExpr(0), steady_count, + ForKind::kSerial, SeqStmt::Flatten(steady_body), + for_node->thread_binding, {})); + + PrimExpr epilogue_base = steady_count - Integer(1); + for (const auto &order_str : epilogue_orders) { + int iter_offset = name2iter(order_str); + int id = name2id(order_str); + PrimExpr replaced_loop_var = + epilogue_base + iter_offset + for_node->min; + variant.push_back( + rewrite_stmt(pipeline_body_seq->seq[id], id, replaced_loop_var)); + } + return SeqStmt::Flatten(variant); + } + + PrimExpr even_base = Integer(2) * for_node->loop_var; + Array super_body = rewrite_body_at_base(even_base, 0); + Array odd_body = rewrite_body_at_base(even_base + Integer(1), 1); + for (const Stmt &stmt : odd_body) { + super_body.push_back(stmt); + } + variant.push_back(For(for_node->loop_var, PrimExpr(0), super_count, + ForKind::kSerial, SeqStmt::Flatten(super_body), + for_node->thread_binding, {})); + + // After all complete pairs, an odd steady_count leaves base + // 2*super_count. Its parity is always even. + if (has_steady_tail) { + Array tail_body = + rewrite_body_at_base(Integer(2) * super_count, 0); + for (const Stmt &stmt : tail_body) { + variant.push_back(stmt); + } + } + + // The epilogue completes the last steady-state base. Its base parity is + // even when steady_count is odd, and odd when steady_count is even. + PrimExpr epilogue_base = steady_count - Integer(1); + int epilogue_base_parity = has_steady_tail ? 0 : 1; + for (const auto &order_str : epilogue_orders) { + int iter_offset = name2iter(order_str); + int id = name2id(order_str); + PrimExpr replaced_loop_var = + epilogue_base + iter_offset + for_node->min; + int logical_iter_parity = (epilogue_base_parity + iter_offset) % 2; + if (logical_iter_parity < 0) { + logical_iter_parity += 2; + } + variant.push_back(rewrite_stmt_with_logical_iter_parity( + pipeline_body_seq->seq[id], id, replaced_loop_var, + logical_iter_parity)); + } + return SeqStmt::Flatten(variant); + }; + + if (static_extent != nullptr) { + int steady_count_value = + static_extent->value - steady_state_max_iter_offset; + ICHECK_GT(steady_count_value, 0); + return build_pipeline_variant((steady_count_value % 2) != 0); + } + + if (!requires_parity_specialization) { + Stmt injected = build_pipeline_variant(false); + Stmt sequential = make_sequential_fallback("runtime_short_extent", false); + return IfThenElse(GT(for_node->extent, Integer(max_logical_iter_offset)), + injected, sequential); + } + + // Specialize the complete pipeline, including its producers, by the + // runtime tail parity. This outer dispatch keeps token ownership within a + // single branch. Short extents take the untouched sequential loop. + Stmt even_variant = build_pipeline_variant(false); + Stmt odd_variant = build_pipeline_variant(true); + PrimExpr has_even_steady_count = + EQ(floormod(steady_count, Integer(2)), Integer(0)); + Stmt injected = + IfThenElse(has_even_steady_count, even_variant, odd_variant); + Stmt sequential = make_sequential_fallback("runtime_short_extent", false); + return IfThenElse(GT(for_node->extent, Integer(max_logical_iter_offset)), + injected, sequential); + } + + Map buffer_data_to_buffer_; + Optional global_symbol_; + ASTTraverser traverser_; +}; + +class ResidentPingPongInitializer : public StmtMutator { +public: + static Stmt Substitute(const Stmt &body) { + ResidentPingPongInitializer rewriter; + PostOrderVisit(body, [&](const ObjectRef &node) { + const auto *loop = node.as(); + if (loop == nullptr) + return; + auto residents = loop->annotations.Get("runtime_resident_banked_buffers"); + auto peers = loop->annotations.Get("runtime_bank_peer_buffers"); + if (!residents || !peers) + return; + Map peer_map = + Downcast>(peers.value()); + for (const Buffer &buffer : Downcast>(residents.value())) { + auto it = peer_map.find(buffer); + ICHECK(it != peer_map.end()) + << "Resident banked buffer " << buffer->name + << " must have a ping/pong peer"; + rewriter.resident_buffers_.insert(buffer.get()); + rewriter.peer_remap_.Set(buffer, (*it).second); + } + }); + if (rewriter.peer_remap_.empty()) + return body; + return rewriter(body); + } + +private: + bool WritesResidentBuffer(const Stmt &stmt) const { + bool writes = false; + PostOrderVisit(stmt, [&](const ObjectRef &node) { + if (const auto *store = node.as()) { + writes = writes || resident_buffers_.count(store->buffer.get()); + return; + } + const auto *call = node.as(); + if (call == nullptr || !call->op.same_as(RegionOp::Get())) + return; + RegionOp region(call->args); + if ((region->GetAccessMask() & 2) != 0) { + writes = writes || resident_buffers_.count(region->GetBuffer().get()); + } + }); + return writes; + } + + Stmt VisitStmt_(const ForNode *op) final { + if (op->annotations.count("runtime_resident_banked_buffers")) { + ++pipeline_depth_; + Stmt result = StmtMutator::VisitStmt_(op); + --pipeline_depth_; + return result; + } + return StmtMutator::VisitStmt_(op); + } + + Stmt VisitStmt_(const SeqStmtNode *op) final { + Array rewritten; + for (const Stmt &original : op->seq) { + Stmt stmt = VisitStmt(original); + rewritten.push_back(stmt); + if (pipeline_depth_ == 0 && !stmt.as() && + WritesResidentBuffer(stmt)) { + rewritten.push_back(RemapBufferRewriter::Substitute(stmt, peer_remap_)); + } + } + return SeqStmt::Flatten(rewritten); + } + + int pipeline_depth_{0}; + std::unordered_set resident_buffers_; + Map peer_remap_; +}; + +tvm::transform::Pass InjectSunmmioPipelineILP() { + using namespace tir::transform; + auto pass_func = [=](PrimFunc f, const IRModule &m, PassContext ctx) { + const PrimFunc &original = f; + try { + PrimFunc candidate = f; + auto *fptr = candidate.CopyOnWrite(); + fptr->body = SunmmioILPMultiVersionBufferRewriter::Substitute(candidate); + fptr->body = ResidentPingPongInitializer::Substitute(fptr->body); + fptr->body = SunmmioILPPipelineInjector::Inject(candidate); + fptr->body = ConvertSSA(std::move(fptr->body)); + Optional disallowed = + PipelineFallbackValidator::FindDisallowed(fptr->body); + if (disallowed) { + return MakePipelineFunctionFallback( + original, PipelineDiagnostic{false, "ilp", "inject_validation", + "candidate_fallback", + std::string(disallowed.value())}); + } + return candidate; + } catch (const std::exception &error) { + return MakePipelineFunctionFallback( + original, + PipelineDiagnostic{false, "ilp", "inject_exception", + "candidate_rewrite_failed", error.what()}); + } catch (...) { + return MakePipelineFunctionFallback( + original, + PipelineDiagnostic{false, "ilp", "inject_exception", + "candidate_rewrite_failed", "unknown exception"}); + } + }; + return CreatePrimFuncPass(pass_func, 0, "tl.InjectSunmmioPipelineILP", {}); +} + +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + refl::GlobalDef().def("tl.transform.InjectSunmmioPipelineILP", + InjectSunmmioPipelineILP); +} + +} // namespace tl +} // namespace tvm diff --git a/src/transform/inject_sunmmio_sync.cc b/src/transform/inject_sunmmio_sync.cc index f40482b884..98f7a8920b 100644 --- a/src/transform/inject_sunmmio_sync.cc +++ b/src/transform/inject_sunmmio_sync.cc @@ -609,8 +609,8 @@ class BufferAccessCollector : public ExprVisitor { }; // Collector for asynchronous operations within a loop body. -// Identifies DMA copies, layout transforms, MMA operations, and Broadcasts that -// happen asynchronously. +// Identifies DMA copies, layout transforms, transposes, MMA operations, and +// broadcasts that happen asynchronously. struct AccessRecord { Buffer buffer; Region region; @@ -1848,7 +1848,8 @@ class InjectSyncRewriter : public StmtMutator { return SeqStmt::Flatten(stmts); } - // Handles specific async instructions (dma_copy, mma_sunmmio, broadcast). + // Handles specific async instructions (DMA, layout transform, transpose, MMA, + // and broadcast). // Assigns tokens/barriers and registers them for dependency tracking. Stmt VisitStmt_(const EvaluateNode *op) { const CallNode *call = op->value.as(); @@ -2262,6 +2263,543 @@ class LoopMissingTokenWaitRewriter : public StmtMutator { std::set tokens_to_wait_before_; }; +enum AsyncWaitDomain : uint8_t { + kUnknownWaitDomain = 0, + kODMA0WaitDomain = 1 << 0, + kODMA1WaitDomain = 1 << 1, + kTCWaitDomain = 1 << 2, +}; + +// These masks mirror only the A4E ODMA path wiring used by NPU-IR channel +// assignment. NPU-IR may further narrow a path using operation capabilities +// (for example, an RSRAM-to-RSRAM strided copy is forced onto ODMA1) or its +// final channel assignment. Ambiguous paths intentionally remain unresolved +// here, so this pass cannot prevent a later-assigned same-channel submission +// from appearing before an older wait in every case. A complete solution needs +// wait placement after channel assignment, or an equivalent channel annotation +// available at this stage. +uint8_t GetODMAReadChannelMask(const Buffer &buffer) { + if (IsGlobalBuffer(buffer)) { + return kODMA0WaitDomain; + } + if (buffer.scope() == kSunmmioScopeRSRAM || buffer.scope() == "local") { + return kODMA0WaitDomain | kODMA1WaitDomain; + } + return kUnknownWaitDomain; +} + +uint8_t GetODMAWriteChannelMask(const Buffer &buffer) { + if (IsGlobalBuffer(buffer) || buffer.scope() == kSunmmioScopeWSRAM) { + return kODMA0WaitDomain; + } + if (buffer.scope() == kSunmmioScopeASRAM) { + return kODMA1WaitDomain; + } + if (buffer.scope() == kSunmmioScopeRSRAM || buffer.scope() == "local") { + return kODMA0WaitDomain | kODMA1WaitDomain; + } + return kUnknownWaitDomain; +} + +uint8_t GetPossibleAsyncWaitDomains(const CallNode *call) { + if (!call) { + return kUnknownWaitDomain; + } + if (call->op.same_as(sunmmio_layout_transform()) || + call->op.same_as(sunmmio_transpose())) { + // A4E layout transforms and full transposes are supported only by ODMA1. + return kODMA1WaitDomain; + } + if (call->op.same_as(mma_sunmmio())) { + return kTCWaitDomain; + } + if (call->op.same_as(broadcast_())) { + const auto *direction = call->args[kBroadcastArgDirection].as(); + // A4E HLink/VLink submissions use the ODMA1/ODMA0 queues respectively. + // Group them with those queues so an older channel-wide wait is placed + // before a newer link submission, and vice versa. + // + // This is a conservative, symmetric submission-domain alias. NPU-IR + // resolves multicast token waits to distinct HLink/VLink wait targets, not + // ODMA0/ODMA1 targets. Consequently, moving an older link wait before a + // newer ODMA submission may reduce otherwise valid overlap. If the shared + // submission queue requires only a directional ordering constraint, it + // should be modeled separately from the resolved wait target rather than + // with this symmetric domain bit. + if (!direction) { + return kODMA0WaitDomain | kODMA1WaitDomain; + } + if (direction->value == 0) { + return kODMA1WaitDomain; + } + if (direction->value == 1) { + return kODMA0WaitDomain; + } + return kODMA0WaitDomain | kODMA1WaitDomain; + } + if (call->op.same_as(dma_copy())) { + BufferRegion src = NormalizeToBufferRegion(call->args[0]); + BufferRegion dst = NormalizeToBufferRegion(call->args[1]); + return GetODMAReadChannelMask(src->buffer) & + GetODMAWriteChannelMask(dst->buffer); + } + return kUnknownWaitDomain; +} + +uint8_t GetDefiniteAsyncWaitDomain(const CallNode *call) { + uint8_t possible_domains = GetPossibleAsyncWaitDomains(call); + // Leave multi-engine paths to the later channel assignment pass. Moving a + // wait for them here could serialize work that ends up on distinct engines. + if (possible_domains != kUnknownWaitDomain && + (possible_domains & (possible_domains - 1)) == 0) { + return possible_domains; + } + return kUnknownWaitDomain; +} + +// Moves a token wait before the first later submission to the same physical +// submission domain. Token waits become engine-wide waits after token +// resolution, so leaving an older wait after a newer same-domain submission +// would unnecessarily drain the newer operation as well. +class EngineAwareWaitPlacementRewriter : public StmtMutator { +public: + Stmt operator()(Stmt body) { + TokenWaitDomainCollector collector; + collector(body); + token_wait_domains_ = std::move(collector.token_wait_domains); + return this->VisitStmt(body); + } + +private: + struct SubmitSummary { + uint8_t possible_domains{kUnknownWaitDomain}; + }; + + static std::optional TryGetTokenId(const CallNode *call) { + if (!call || call->args.empty()) { + return std::nullopt; + } + if (const auto *imm = call->args[0].as()) { + return static_cast(imm->value); + } + return std::nullopt; + } + + static std::optional TryGetGeneratedTokenId(const CallNode *call) { + if (!call) { + return std::nullopt; + } + for (const PrimExpr &arg : call->args) { + const auto *token_call = arg.as(); + if (token_call && token_call->op.same_as(sync_token_id())) { + return TryGetTokenId(token_call); + } + } + return std::nullopt; + } + + static bool IsAsyncSubmit(const CallNode *call) { + return call && + (call->op.same_as(dma_copy()) || + call->op.same_as(sunmmio_layout_transform()) || + call->op.same_as(sunmmio_transpose()) || + call->op.same_as(mma_sunmmio()) || call->op.same_as(broadcast_())); + } + + static bool MatchWaitTokenStmt(const Stmt &stmt, int *token_id) { + const auto *eval = stmt.as(); + if (!eval) { + return false; + } + const auto *call = eval->value.as(); + if (!call || !call->op.same_as(wait_token())) { + return false; + } + std::optional id = TryGetTokenId(call); + if (!id) { + return false; + } + *token_id = *id; + return true; + } + + static uint8_t GetUnconditionalSubmitDomains(const Stmt &stmt) { + // Only recurse through statements whose body executes whenever the + // wrapper executes. Unrecognized control flow remains a movement barrier. + if (const auto *eval = stmt.as()) { + return GetDefiniteAsyncWaitDomain(eval->value.as()); + } + if (const auto *seq = stmt.as()) { + uint8_t domains = kUnknownWaitDomain; + for (const Stmt &child : seq->seq) { + domains |= GetUnconditionalSubmitDomains(child); + } + return domains; + } + if (const auto *attr = stmt.as()) { + return GetUnconditionalSubmitDomains(attr->body); + } + if (const auto *let = stmt.as()) { + return GetUnconditionalSubmitDomains(let->body); + } + if (const auto *decl = stmt.as()) { + return GetUnconditionalSubmitDomains(decl->body); + } + if (const auto *allocate = stmt.as()) { + return is_one(allocate->condition) + ? GetUnconditionalSubmitDomains(allocate->body) + : kUnknownWaitDomain; + } + if (const auto *allocate_const = stmt.as()) { + return GetUnconditionalSubmitDomains(allocate_const->body); + } + if (const auto *realize = stmt.as()) { + return is_one(realize->condition) + ? GetUnconditionalSubmitDomains(realize->body) + : kUnknownWaitDomain; + } + if (const auto *realize = stmt.as()) { + if (!is_one(realize->predicate)) { + return kUnknownWaitDomain; + } + // A reduction block's init is not executed on every realization. + return GetUnconditionalSubmitDomains(realize->block->body); + } + return kUnknownWaitDomain; + } + + static void PushFlatten(Array *out, const Stmt &stmt) { + if (!stmt.defined()) { + return; + } + if (const auto *seq = stmt.as()) { + for (const Stmt &child : seq->seq) { + PushFlatten(out, child); + } + return; + } + out->push_back(stmt); + } + + class TokenWaitDomainCollector : public StmtExprVisitor { + public: + void VisitExpr_(const CallNode *op) final { + uint8_t domain = GetDefiniteAsyncWaitDomain(op); + if (domain != kUnknownWaitDomain) { + if (std::optional token_id = TryGetGeneratedTokenId(op)) { + auto [it, inserted] = token_wait_domains.emplace(*token_id, domain); + if (!inserted && it->second != domain) { + it->second = kUnknownWaitDomain; + } + } + } + StmtExprVisitor::VisitExpr_(op); + } + + std::map token_wait_domains; + }; + + class GeneratedTokenCollector : public StmtExprVisitor { + public: + void VisitExpr_(const CallNode *op) final { + if (op->op.same_as(sync_token_id())) { + if (std::optional token_id = TryGetTokenId(op)) { + tokens.insert(*token_id); + } + } + StmtExprVisitor::VisitExpr_(op); + } + + std::set tokens; + }; + + class SubmitSummaryCollector : public StmtExprVisitor { + public: + void VisitExpr_(const CallNode *op) final { + if (IsAsyncSubmit(op)) { + summary.possible_domains |= GetPossibleAsyncWaitDomains(op); + } + StmtExprVisitor::VisitExpr_(op); + } + + SubmitSummary summary; + }; + + Stmt VisitStmt_(const ForNode *op) final { + ++loop_depth_; + Stmt result = StmtMutator::VisitStmt_(op); + --loop_depth_; + return result; + } + + Stmt VisitStmt_(const WhileNode *op) final { + ++loop_depth_; + Stmt result = StmtMutator::VisitStmt_(op); + --loop_depth_; + return result; + } + + Stmt VisitStmt_(const SeqStmtNode *op) final { + Array stmts; + for (const Stmt &stmt : op->seq) { + PushFlatten(&stmts, VisitStmt(stmt)); + } + + int n = static_cast(stmts.size()); + std::vector> generated_tokens(n); + std::vector submit_summaries(n); + std::vector unconditional_submit_domains(n, kUnknownWaitDomain); + for (int i = 0; i < n; ++i) { + GeneratedTokenCollector token_collector; + token_collector(stmts[i]); + generated_tokens[i] = std::move(token_collector.tokens); + + SubmitSummaryCollector submit_collector; + submit_collector(stmts[i]); + submit_summaries[i] = submit_collector.summary; + unconditional_submit_domains[i] = GetUnconditionalSubmitDomains(stmts[i]); + } + + std::vector> waits_before(n); + std::vector remove_wait(n, false); + for (int wait_index = 0; wait_index < n; ++wait_index) { + int token_id = -1; + if (!MatchWaitTokenStmt(stmts[wait_index], &token_id)) { + continue; + } + auto domain_it = token_wait_domains_.find(token_id); + if (domain_it == token_wait_domains_.end() || + domain_it->second == kUnknownWaitDomain) { + continue; + } + + int lower_bound = 0; + int last_generator = -1; + for (int i = 0; i < wait_index; ++i) { + if (generated_tokens[i].count(token_id) != 0) { + last_generator = i; + } + } + if (last_generator >= 0) { + lower_bound = last_generator + 1; + } else { + // A wait before its static generation site is a loop-carried wait for + // the previous iteration. Outside a loop, do not infer such a lifetime. + bool generated_later = false; + for (int i = wait_index + 1; i < n; ++i) { + generated_later |= generated_tokens[i].count(token_id) != 0; + } + if (loop_depth_ == 0 || !generated_later) { + continue; + } + } + + int anchor = -1; + uint8_t wait_domain = domain_it->second; + for (int i = lower_bound; i < wait_index; ++i) { + const SubmitSummary &summary = submit_summaries[i]; + if ((summary.possible_domains & wait_domain) == 0) { + continue; + } + // Transparent wrappers may contain a valid anchor, but a submit hidden + // under conditional or repeated control flow still blocks movement. + if ((unconditional_submit_domains[i] & wait_domain) != 0) { + anchor = i; + } + break; + } + if (anchor >= 0) { + waits_before[anchor].push_back(stmts[wait_index]); + remove_wait[wait_index] = true; + } + } + + Array out; + for (int i = 0; i < n; ++i) { + for (const Stmt &wait : waits_before[i]) { + out.push_back(wait); + } + if (!remove_wait[i]) { + out.push_back(stmts[i]); + } + } + return SeqStmt::Flatten(out); + } + + int loop_depth_{0}; + std::map token_wait_domains_; +}; + +// A loop-carried wait refers to a token produced by the previous iteration. +// ResolveTokens intentionally merges the preheader null token with the token +// produced in the body, so the null token alone cannot make the first wait a +// no-op after lowering. Guard such waits explicitly on there being a previous +// iteration. Counted loops use their induction variable; while loops carry a +// local boolean state. This runs after wait placement so the placement pass can +// continue to reason about plain wait statements. +// +// Loop-exit waits are intentionally left unconditional. The current pipeline +// assumes loops carrying asynchronous tokens have a non-zero runtime extent; +// zero-trip support must also predicate waits that consume loop result tokens. +class LoopCarriedWaitConditionRewriter : public StmtMutator { +public: + Stmt operator()(Stmt body) { return this->VisitStmt(body); } + +private: + struct WaitRecord { + const EvaluateNode *stmt{nullptr}; + int token_id{-1}; + int order{-1}; + }; + + static std::optional TryGetTokenId(const CallNode *call) { + if (!call || call->args.empty()) { + return std::nullopt; + } + if (const auto *imm = call->args[0].as()) { + return static_cast(imm->value); + } + return std::nullopt; + } + + static std::optional TryGetGeneratedTokenId(const CallNode *call) { + if (!call) { + return std::nullopt; + } + for (const PrimExpr &arg : call->args) { + const auto *token_call = arg.as(); + if (token_call && token_call->op.same_as(sync_token_id())) { + return TryGetTokenId(token_call); + } + } + return std::nullopt; + } + + class LoopEventCollector : public StmtVisitor { + public: + void VisitStmt_(const EvaluateNode *op) final { + int order = next_order_++; + const auto *call = op->value.as(); + if (!call) { + return; + } + if (call->op.same_as(wait_token())) { + if (std::optional token_id = TryGetTokenId(call)) { + waits.push_back(WaitRecord{op, *token_id, order}); + } + return; + } + if (nested_loop_depth_ == 0) { + std::optional token_id = TryGetGeneratedTokenId(call); + if (!token_id) { + return; + } + auto [it, inserted] = first_generation_order.emplace(*token_id, order); + if (!inserted) { + it->second = std::min(it->second, order); + } + } + } + + // Waits may consume an outer-loop token from inside a nested loop, but a + // token generated by the nested loop belongs to that loop's condition. + void VisitStmt_(const ForNode *op) final { + ++nested_loop_depth_; + StmtVisitor::VisitStmt_(op); + --nested_loop_depth_; + } + + void VisitStmt_(const WhileNode *op) final { + ++nested_loop_depth_; + StmtVisitor::VisitStmt_(op); + --nested_loop_depth_; + } + + std::vector waits; + std::map first_generation_order; + + private: + int next_order_{0}; + int nested_loop_depth_{0}; + }; + + class WaitConditionalizer : public StmtMutator { + public: + WaitConditionalizer(PrimExpr condition, + std::unordered_set waits) + : condition_(std::move(condition)), waits_(std::move(waits)) {} + + Stmt VisitStmt_(const EvaluateNode *op) final { + if (waits_.count(op) == 0) { + return ffi::GetRef(op); + } + return IfThenElse(condition_, ffi::GetRef(op)); + } + + private: + PrimExpr condition_; + std::unordered_set waits_; + }; + + static std::unordered_set + FindLoopCarriedWaits(const Stmt &body) { + LoopEventCollector collector; + collector(body); + std::unordered_set loop_carried_waits; + for (const WaitRecord &wait : collector.waits) { + auto generation_it = collector.first_generation_order.find(wait.token_id); + if (generation_it != collector.first_generation_order.end() && + wait.order < generation_it->second) { + loop_carried_waits.insert(wait.stmt); + } + } + return loop_carried_waits; + } + + Stmt VisitStmt_(const ForNode *op) final { + Stmt body = this->VisitStmt(op->body); + std::unordered_set loop_carried_waits = + FindLoopCarriedWaits(body); + + if (!loop_carried_waits.empty()) { + PrimExpr has_previous_iteration = GT(op->loop_var, op->min); + body = WaitConditionalizer(has_previous_iteration, + std::move(loop_carried_waits))(body); + } + + return For(op->loop_var, op->min, op->extent, op->kind, body, + op->thread_binding, op->annotations, std::nullopt, op->span); + } + + Stmt VisitStmt_(const WhileNode *op) final { + Stmt body = this->VisitStmt(op->body); + std::unordered_set loop_carried_waits = + FindLoopCarriedWaits(body); + if (loop_carried_waits.empty()) { + return While(op->condition, body, op->span); + } + + Buffer has_previous_iteration = + decl_buffer({Integer(1)}, DataType::Bool(), + "sunmmio_has_previous_iteration", "local.var"); + PrimExpr zero = Integer(0); + PrimExpr condition = BufferLoad(has_previous_iteration, {zero}); + body = WaitConditionalizer(condition, std::move(loop_carried_waits))(body); + Array loop_body{ + body, BufferStore(has_previous_iteration, const_true(), {zero})}; + body = SeqStmt::Flatten(loop_body); + + Stmt while_stmt = While(op->condition, body, op->span); + while_stmt = DeclBuffer(has_previous_iteration, while_stmt); + Map annotations; + annotations.Set(tl::attr::kLocalVarInit, + make_const(DataType::Bool(), false)); + return Allocate(has_previous_iteration->data, has_previous_iteration->dtype, + has_previous_iteration->shape, const_true(), while_stmt, + annotations); + } +}; + // Rewriter to inject final synchronization waits before the device function // returns. This ensures all pending asynchronous operations are completed // before the device kernel finishes, handling both explicit returns and @@ -2846,6 +3384,14 @@ class SunmmioSyncRewriter : public IRMutatorWithAnalyzer { auto loop_missing_token_wait_rewriter = LoopMissingTokenWaitRewriter(); f.CopyOnWrite()->body = loop_missing_token_wait_rewriter(f->body); + auto engine_aware_wait_placement_rewriter = + EngineAwareWaitPlacementRewriter(); + f.CopyOnWrite()->body = engine_aware_wait_placement_rewriter(f->body); + + auto loop_carried_wait_condition_rewriter = + LoopCarriedWaitConditionRewriter(); + f.CopyOnWrite()->body = loop_carried_wait_condition_rewriter(f->body); + auto device_func_wait_rewriter = DeviceFuncWaitRewriter(); f.CopyOnWrite()->body = device_func_wait_rewriter(f->body); diff --git a/src/transform/lower_opaque_block.cc b/src/transform/lower_opaque_block.cc index bedf407381..65ceacc57d 100644 --- a/src/transform/lower_opaque_block.cc +++ b/src/transform/lower_opaque_block.cc @@ -31,6 +31,7 @@ #include #include "../op/builtin.h" +#include "common/attr.h" #include "tir/transforms/ir_utils.h" namespace tvm { @@ -108,6 +109,11 @@ class OpaqueBlockLower : public StmtExprMutator { allocate_annotations.Set(tl::attr::kSunmmioAllocPingPong, (*ping_pong_it).second); } + auto reduce_temp_it = reduce_register_temp_role_map_.find(buffer->data); + if (reduce_temp_it != reduce_register_temp_role_map_.end()) { + allocate_annotations.Set(tl::attr::kSunmmioReduceRegisterTemp, + (*reduce_temp_it).second); + } body = Allocate(buffer->data, buffer->dtype, allocation_shape, const_true(), std::move(body), allocate_annotations); } @@ -264,6 +270,17 @@ class OpaqueBlockLower : public StmtExprMutator { << "` to be a PrimExpr or Map, but got " << kv.second.GetTypeKey(); } + } else if (key == tl::attr::kSunmmioReduceRegisterTemp) { + ICHECK(is_block) + << "`" << tl::attr::kSunmmioReduceRegisterTemp + << "` is only supported as a per-buffer block annotation"; + auto roles = kv.second.try_cast>(); + ICHECK(roles.has_value()) + << "Expected `" << tl::attr::kSunmmioReduceRegisterTemp + << "` to be Map, but got " << kv.second.GetTypeKey(); + for (const auto &pair : roles.value()) { + reduce_register_temp_role_map_.Set(pair.first, pair.second); + } } else if (!is_block) { // the loop annotation is preserved preserved_annotations.Set(key, kv.second); @@ -303,6 +320,9 @@ class OpaqueBlockLower : public StmtExprMutator { /*! \brief SunMMIO alloc ping-pong attrs collected from function attrs. */ Map alloc_ping_pong_map_; + + /*! \brief SunMMIO reduction role transferred to Allocate annotations. */ + Map reduce_register_temp_role_map_; }; PrimFunc TLLowerOpaqueBlock(PrimFunc f) { diff --git a/src/transform/remove_unused_sunmmio_allocations.cc b/src/transform/remove_unused_sunmmio_allocations.cc new file mode 100644 index 0000000000..b30d87f0bc --- /dev/null +++ b/src/transform/remove_unused_sunmmio_allocations.cc @@ -0,0 +1,201 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/*! + * \file remove_unused_sunmmio_allocations.cc + * \brief Remove SunMMIO allocations that survive RemoveNoOp only because of + * metadata annotations. + */ + +#include +#include +#include + +#include +#include + +#include "../layout/layout.h" +#include "../op/builtin.h" + +namespace tvm { +namespace tl { + +using namespace tir; + +class ExecutableBufferUseCollector : public StmtExprVisitor { +public: + static std::unordered_set Collect(const Stmt &body) { + ExecutableBufferUseCollector collector; + collector(body); + return std::move(collector.used_vars_); + } + +private: + void VisitStmt_(const AllocateNode *op) final { + for (const PrimExpr &extent : op->extents) { + VisitExpr(extent); + } + VisitExpr(op->condition); + VisitStmt(op->body); + } + + void VisitStmt_(const DeclBufferNode *op) final { VisitStmt(op->body); } + + void VisitExpr_(const VarNode *op) final { used_vars_.insert(op); } + + void VisitExpr_(const BufferLoadNode *op) final { + used_vars_.insert(op->buffer->data.get()); + StmtExprVisitor::VisitExpr_(op); + } + + void VisitStmt_(const BufferStoreNode *op) final { + used_vars_.insert(op->buffer->data.get()); + StmtExprVisitor::VisitStmt_(op); + } + + std::unordered_set used_vars_; +}; + +class UnusedSunmmioAllocationRemover : public StmtExprMutator { +public: + explicit UnusedSunmmioAllocationRemover( + const std::unordered_set &used_vars) + : used_vars_(used_vars) {} + +private: + bool IsUsed(const Var &var) const { return used_vars_.count(var.get()) != 0; } + + Stmt VisitStmt_(const AllocateNode *op) final { + Stmt body = VisitStmt(op->body); + if (!IsUsed(op->buffer_var)) { + return body; + } + Array extents; + for (const PrimExpr &extent : op->extents) { + extents.push_back(VisitExpr(extent)); + } + PrimExpr condition = VisitExpr(op->condition); + return Allocate(op->buffer_var, op->dtype, std::move(extents), + std::move(condition), std::move(body), op->annotations, + op->span); + } + + Stmt VisitStmt_(const DeclBufferNode *op) final { + Stmt body = VisitStmt(op->body); + if (!IsUsed(op->buffer->data)) { + return body; + } + return DeclBuffer(op->buffer, std::move(body), op->span); + } + + Stmt VisitStmt_(const BlockNode *op) final { + Block block = Downcast(StmtExprMutator::VisitStmt_(op)); + auto layout_it = block->annotations.find(attr::kLayoutMap); + if (layout_it == block->annotations.end()) { + return block; + } + + Map annotations = block->annotations; + if (auto layout_map = (*layout_it).second.as>()) { + Map filtered; + for (const auto &[buffer, layout] : layout_map.value()) { + if (IsUsed(buffer->data)) { + filtered.Set(buffer, layout); + } + } + if (filtered.empty()) { + annotations.erase(attr::kLayoutMap); + } else { + annotations.Set(attr::kLayoutMap, filtered); + } + } else if (auto layout_map = (*layout_it).second.as>()) { + Map filtered; + for (const auto &[var, layout] : layout_map.value()) { + if (IsUsed(var)) { + filtered.Set(var, layout); + } + } + if (filtered.empty()) { + annotations.erase(attr::kLayoutMap); + } else { + annotations.Set(attr::kLayoutMap, filtered); + } + } + block.CopyOnWrite()->annotations = std::move(annotations); + return block; + } + + const std::unordered_set &used_vars_; +}; + +PrimFunc RemoveUnusedSunmmioAllocationsFromFunc(PrimFunc func) { + std::unordered_set used_vars = + ExecutableBufferUseCollector::Collect(func->body); + UnusedSunmmioAllocationRemover remover(used_vars); + func.CopyOnWrite()->body = remover(func->body); + + if (auto layout_map = func->GetAttr>(attr::kLayoutMap)) { + Map filtered; + for (const auto &[buffer, layout] : layout_map.value()) { + if (used_vars.count(buffer->data.get())) { + filtered.Set(buffer, layout); + } + } + if (filtered.empty()) { + func = WithoutAttr(std::move(func), ffi::String(attr::kLayoutMap)); + } else { + func = WithAttr(std::move(func), attr::kLayoutMap, filtered); + } + } + + if (auto ping_pong = + func->GetAttr>(tl::attr::kSunmmioAllocPingPong)) { + Map filtered; + for (const auto &[var, bank] : ping_pong.value()) { + if (used_vars.count(var.get())) { + filtered.Set(var, bank); + } + } + if (filtered.empty()) { + func = WithoutAttr(std::move(func), + ffi::String(tl::attr::kSunmmioAllocPingPong)); + } else { + func = + WithAttr(std::move(func), tl::attr::kSunmmioAllocPingPong, filtered); + } + } + return func; +} + +tvm::transform::Pass RemoveUnusedSunmmioAllocations() { + auto pass_func = [](PrimFunc func, IRModule, tvm::transform::PassContext) { + return RemoveUnusedSunmmioAllocationsFromFunc(std::move(func)); + }; + return tir::transform::CreatePrimFuncPass( + pass_func, 0, "tl.RemoveUnusedSunmmioAllocations", {}); +} + +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + refl::GlobalDef().def("tl.transform.RemoveUnusedSunmmioAllocations", + RemoveUnusedSunmmioAllocations); +} + +} // namespace tl +} // namespace tvm diff --git a/src/transform/sunmmio_pipeline_planning.cc b/src/transform/sunmmio_pipeline_planning.cc index 0522922301..4e009f3fd4 100644 --- a/src/transform/sunmmio_pipeline_planning.cc +++ b/src/transform/sunmmio_pipeline_planning.cc @@ -13,12 +13,19 @@ * call. */ +#include "../op/builtin.h" +#include "../op/comm.h" #include "../op/utils.h" #include "../target/sunmmio/cost_model.h" #include "../target/sunmmio/hardware_types.h" +#include "../target/sunmmio_utils.h" +#include "sunmmio_pipeline_planning/pipeline_diagnostic.h" +#include "sunmmio_pipeline_planning/resource_types_for_ilp.h" #include "sunmmio_pipeline_planning/stmt_read_write_collector.h" #include +#include +#include #include #include #include @@ -75,6 +82,28 @@ class AccessOverlapChecker { } }; +static bool HasRepeatedCollectiveDestination(const SeqStmtNode *body) { + std::unordered_set destinations; + for (const Stmt &stmt : body->seq) { + bool repeated = false; + PostOrderVisit(stmt, [&](const ObjectRef &obj) { + const auto *call = obj.as(); + if (!call || !call->op.same_as(Op::Get("tl.broadcast_"))) { + return; + } + const BufferNode *destination = + NormalizeToBufferRegion(call->args[1])->buffer.get(); + if (!destinations.insert(destination).second) { + repeated = true; + } + }); + if (repeated) { + return true; + } + } + return false; +} + /** * \brief Pure data container representing an instruction in the pipeline. * It separates the AST analysis from scheduling and latency calculation. @@ -86,6 +115,7 @@ class PipelineInstruction { std::string name{""}; Stmt stmt; DeviceType device_type{DeviceType::Unspecified}; + int execution_resource{-1}; // True if this instruction should be placed in the prefetch queue (Shift=1) bool is_prefetch{false}; @@ -120,6 +150,144 @@ class PipelineInstruction { } }; +struct GreedyAccessInfo { + BufferRegion region; + bool is_write{false}; + + Buffer buffer() const { return region->buffer; } +}; + +static int GetGreedyExecutionResource(const PipelineInstruction &instruction) { + std::vector accesses; + accesses.reserve(instruction.reads.size() + instruction.writes.size()); + for (const BufferRegion &read : instruction.reads) { + accesses.push_back({read, false}); + } + for (const BufferRegion &write : instruction.writes) { + accesses.push_back({write, true}); + } + std::vector resources = + BuildIlpResources(instruction.stmt, instruction.device_type, accesses); + for (int resource : resources) { + if (resource == static_cast(IlpResourceType::kTensorCore) || + resource == static_cast(IlpResourceType::kVectorCore) || + resource == static_cast(IlpResourceType::kODMA0) || + resource == static_cast(IlpResourceType::kODMA1)) { + return resource; + } + } + LOG(FATAL) << "No execution resource for greedy pipeline instruction " + << instruction.name; + return -1; +} + +static int GetGreedyIssuePriority(const PipelineInstruction &instruction) { + int resource = instruction.execution_resource; + if (resource == static_cast(IlpResourceType::kODMA1)) { + return 0; + } + if (resource == static_cast(IlpResourceType::kODMA0)) { + return 1; + } + // DMA launch is asynchronous, while tensor commands block the scalar issue + // stream. Launch same-time asynchronous work before blocking computation. + if (resource == static_cast(IlpResourceType::kTensorCore)) { + return 2; + } + if (resource == static_cast(IlpResourceType::kVectorCore)) { + return 3; + } + return 4; +} + +static bool IsAllGatherInstruction(const PipelineInstruction &instruction) { + const CallNode *broadcast = nullptr; + PostOrderVisit(instruction.stmt, [&](const ObjectRef &obj) { + const auto *call = obj.as(); + if (call && call->op.same_as(Op::Get("tl.broadcast_"))) { + ICHECK(broadcast == nullptr) + << "A pipeline statement may contain at most one broadcast leaf"; + broadcast = call; + } + }); + if (broadcast == nullptr) { + return false; + } + ICHECK(broadcast->args.size() == static_cast(kBroadcastArgCount) || + broadcast->args.size() == static_cast(kBroadcastArgCount + 1)) + << "tl.broadcast_ expects its fixed arguments and optional src_core"; + return broadcast->args.size() == static_cast(kBroadcastArgCount); +} + +enum class PhysicalSramBank : int { + ASRAMPing = 0, + ASRAMPong = 1, + WSRAMPing = 2, + WSRAMPong = 3, + Count = 4, +}; + +using PerCommandBankPhases = + std::unordered_map>; + +struct GreedyBankColoring { + PerCommandBankPhases writer_phases; + PerCommandBankPhases reader_phases; + std::vector bits; +}; + +static int LookupBankPhase(const PerCommandBankPhases &phases, + const BufferNode *buffer, int command_id) { + auto it_buffer = phases.find(buffer); + if (it_buffer == phases.end()) + return 0; + auto it_command = it_buffer->second.find(command_id); + return it_command == it_buffer->second.end() ? 0 : it_command->second; +} + +static std::vector GetOccupiedSramBanks( + const PipelineInstruction &instruction, + const std::unordered_set &versioned_buffers, + int iter_mod, const PerCommandBankPhases &writer_phases, + const PerCommandBankPhases &reader_phases) { + std::array(PhysicalSramBank::Count)> occupied{}; + int version_slot = instruction.iter; + if (iter_mod > 0) { + version_slot %= iter_mod; + if (version_slot < 0) { + version_slot += iter_mod; + } + } + auto collect = [&](const BufferRegion ®ion, bool is_write) { + const String &scope = region->buffer.scope(); + int phase = LookupBankPhase(is_write ? writer_phases : reader_phases, + region->buffer.get(), instruction.id); + bool pong = versioned_buffers.count(region->buffer.get()) != 0 && + (version_slot + phase) % 2 != 0; + if (scope == kSunmmioScopeASRAM) { + occupied[static_cast(pong ? PhysicalSramBank::ASRAMPong + : PhysicalSramBank::ASRAMPing)] = true; + } else if (scope == kSunmmioScopeWSRAM) { + occupied[static_cast(pong ? PhysicalSramBank::WSRAMPong + : PhysicalSramBank::WSRAMPing)] = true; + } + }; + for (const BufferRegion ®ion : instruction.reads) { + collect(region, false); + } + for (const BufferRegion ®ion : instruction.writes) { + collect(region, true); + } + + std::vector result; + for (int i = 0; i < static_cast(PhysicalSramBank::Count); ++i) { + if (occupied[i]) { + result.push_back(static_cast(i)); + } + } + return result; +} + /** * \brief A RAW dependence edge in the single-iteration local DDG. * @@ -133,6 +301,34 @@ struct LocalDependencyEdge { int distance{0}; }; +struct TemplateOrderEdge { + int source_instruction_id{-1}; + int target_instruction_id{-1}; + int distance{0}; +}; + +enum class SemanticDependencyKind { kRAW, kWAR, kWAW }; + +struct SemanticDependencyEdge { + int source_instruction_id{-1}; + int target_instruction_id{-1}; + const BufferNode *buffer{nullptr}; + int distance{0}; + SemanticDependencyKind kind{SemanticDependencyKind::kRAW}; +}; + +static const char *SemanticDependencyKindName(SemanticDependencyKind kind) { + switch (kind) { + case SemanticDependencyKind::kRAW: + return "RAW"; + case SemanticDependencyKind::kWAR: + return "WAR"; + case SemanticDependencyKind::kWAW: + return "WAW"; + } + return "unknown"; +} + /** * \brief Aggregated access information for one logical buffer in the local DDG. */ @@ -156,6 +352,8 @@ struct BufferAccessInfo { */ struct LocalDDG { std::vector edges; + std::vector ordering_edges; + std::vector semantic_edges; std::vector> forward_predecessors; std::vector> forward_successors; std::vector> backward_predecessors; @@ -164,6 +362,95 @@ struct LocalDDG { std::vector buffer_order; }; +static bool IsRuntimeBankedBuffer(const BufferNode *buffer) { + const String &scope = tvm::ffi::GetRef(buffer).scope(); + return scope == kSunmmioScopeASRAM || scope == kSunmmioScopeWSRAM; +} + +static std::vector BuildGreedyBankColorings( + const LocalDDG &local_ddg, + const std::unordered_set &versioned_buffers, int faster, + size_t *total_candidate_count) { + using WriterKey = std::pair; + std::map> writers_by_buffer; + for (const LocalDependencyEdge &edge : local_ddg.edges) { + if (!versioned_buffers.count(edge.buffer) || + !IsRuntimeBankedBuffer(edge.buffer)) { + continue; + } + writers_by_buffer[edge.buffer].insert(edge.producer_instruction_id); + } + + // Precolors are relative bank phases. Writers 0,2,... of one buffer must + // stay together, writers 1,3,... must stay together, and the two classes + // must use opposite banks. Search one global inversion bit per buffer + // instead of independently recoloring every writer. + struct WriterColor { + int variable{-1}; + int precolor{0}; + }; + std::map writer_colors; + int search_bits = 0; + for (const BufferNode *buffer : local_ddg.buffer_order) { + auto writers_it = writers_by_buffer.find(buffer); + if (writers_it == writers_by_buffer.end()) { + continue; + } + int precolor = 0; + for (int writer_id : writers_it->second) { + writer_colors[{buffer, writer_id}] = {search_bits, precolor}; + precolor ^= 1; + } + ++search_bits; + } + + ICHECK(faster == -1 || faster > 0) + << "tl.sunmmio_faster must be -1 or a positive coloring budget"; + ICHECK_LT(search_bits, static_cast(sizeof(size_t) * 8)) + << "Too many independent greedy coloring variables: " << search_bits; + size_t total_candidates = size_t{1} << search_bits; + *total_candidate_count = total_candidates; + size_t candidate_count = + faster == -1 ? total_candidates + : std::min(total_candidates, static_cast(faster)); + std::vector result; + result.reserve(candidate_count); + for (size_t mask = 0; mask < candidate_count; ++mask) { + GreedyBankColoring coloring; + coloring.bits.resize(search_bits, 0); + for (int i = 0; i < search_bits; ++i) { + coloring.bits[i] = static_cast((mask >> i) & 1); + } + for (const auto &[writer, color] : writer_colors) { + coloring.writer_phases[writer.first][writer.second] = + coloring.bits[color.variable] ^ color.precolor; + } + + bool valid = true; + for (const LocalDependencyEdge &edge : local_ddg.edges) { + auto it_writer = + writer_colors.find({edge.buffer, edge.producer_instruction_id}); + if (it_writer == writer_colors.end()) + continue; + const WriterColor &color = it_writer->second; + int writer_phase = coloring.bits[color.variable] ^ color.precolor; + int reader_phase = writer_phase ^ (edge.distance & 1); + auto &reader_map = coloring.reader_phases[edge.buffer]; + auto [it_reader, inserted] = + reader_map.emplace(edge.consumer_instruction_id, reader_phase); + if (!inserted && it_reader->second != reader_phase) { + valid = false; + break; + } + } + if (valid) + result.push_back(std::move(coloring)); + } + if (result.empty()) + result.push_back(GreedyBankColoring{}); + return result; +} + /** * \brief Build the single-iteration local DDG from read/write regions. */ @@ -296,10 +583,281 @@ class LocalDDGBuilder { } } + std::set< + std::tuple> + semantic_unique; + auto add_semantic = [&](int source, int target, const BufferNode *buffer, + int distance, SemanticDependencyKind kind) { + auto key = std::make_tuple(source, target, buffer, distance, kind); + if (semantic_unique.insert(key).second) { + ddg.semantic_edges.push_back({source, target, buffer, distance, kind}); + } + }; + for (const LocalDependencyEdge &edge : ddg.edges) { + add_semantic(edge.producer_instruction_id, edge.consumer_instruction_id, + edge.buffer, edge.distance, SemanticDependencyKind::kRAW); + } + for (const BufferNode *buffer : ddg.buffer_order) { + const BufferAccessInfo &info = ddg.buffer_access_infos.at(buffer); + for (int reader : info.read_instruction_indices) { + for (int writer : info.write_instruction_indices) { + if (reader < writer) { + add_semantic(reader, writer, buffer, 0, + SemanticDependencyKind::kWAR); + } + } + } + for (size_t i = 0; i < info.write_instruction_indices.size(); ++i) { + for (size_t j = i + 1; j < info.write_instruction_indices.size(); ++j) { + add_semantic(info.write_instruction_indices[i], + info.write_instruction_indices[j], buffer, 0, + SemanticDependencyKind::kWAW); + } + } + if (!info.read_instruction_indices.empty() && + !info.write_instruction_indices.empty()) { + add_semantic(info.read_instruction_indices.back(), + info.write_instruction_indices.front(), buffer, 1, + SemanticDependencyKind::kWAR); + if (info.write_instruction_indices.size() > 1) { + add_semantic(info.write_instruction_indices.back(), + info.write_instruction_indices.front(), buffer, 1, + SemanticDependencyKind::kWAW); + } + } + } + + // All cores must encounter all-gather barriers in one common epoch order. + // Data hazards alone cannot enforce this when collectives use different + // buffers or ODMA directions, so preserve template order explicitly and + // close the chain across consecutive logical iterations. + std::vector all_gather_ids; + for (const PipelineInstruction &instruction : + single_iteration_instructions) { + if (IsAllGatherInstruction(instruction)) { + all_gather_ids.push_back(instruction.id); + } + } + for (size_t i = 1; i < all_gather_ids.size(); ++i) { + int source = all_gather_ids[i - 1]; + int target = all_gather_ids[i]; + ddg.ordering_edges.push_back({source, target, 0}); + } + if (all_gather_ids.size() > 1) { + int source = all_gather_ids.back(); + int target = all_gather_ids.front(); + ddg.ordering_edges.push_back({source, target, 1}); + } + return ddg; } }; +static bool +ValidateLocalDDG(const std::vector &instructions, + const LocalDDG &ddg) { + const int instruction_count = static_cast(instructions.size()); + std::set> reads_with_producer; + for (const LocalDependencyEdge &edge : ddg.edges) { + if (edge.producer_instruction_id < 0 || + edge.producer_instruction_id >= instruction_count || + edge.consumer_instruction_id < 0 || + edge.consumer_instruction_id >= instruction_count || + edge.buffer == nullptr || edge.distance < 0) { + return false; + } + reads_with_producer.insert({edge.consumer_instruction_id, edge.buffer}); + } + for (const SemanticDependencyEdge &edge : ddg.semantic_edges) { + if (edge.source_instruction_id < 0 || + edge.source_instruction_id >= instruction_count || + edge.target_instruction_id < 0 || + edge.target_instruction_id >= instruction_count || + edge.buffer == nullptr || edge.distance < 0) { + return false; + } + } + for (const TemplateOrderEdge &edge : ddg.ordering_edges) { + if (edge.source_instruction_id < 0 || + edge.source_instruction_id >= instruction_count || + edge.target_instruction_id < 0 || + edge.target_instruction_id >= instruction_count || edge.distance < 0) { + return false; + } + } + for (int id = 0; id < instruction_count; ++id) { + for (const BufferRegion &read : instructions[id].reads) { + if (IsGlobalBuffer(read->buffer)) { + continue; + } + auto access_it = ddg.buffer_access_infos.find(read->buffer.get()); + if (access_it == ddg.buffer_access_infos.end() || + access_it->second.write_instruction_indices.empty()) { + continue; + } + if (!reads_with_producer.count({id, read->buffer.get()})) { + return false; + } + } + } + return true; +} + +static void MaybeWriteGreedyGraphJson( + const std::vector &instructions, const LocalDDG &ddg, + const std::unordered_set &versioned_buffers) { + const char *path = std::getenv("TL_SUNMMIO_PIPELINE_GRAPH_JSON"); + if (path == nullptr || path[0] == '\0') { + return; + } + std::ofstream out(path); + if (!out.is_open()) { + LOG(WARNING) << "Cannot write pipeline graph JSON to " << path; + return; + } + out << "{\n \"mode\": \"greedy\",\n \"commands\": [\n"; + for (size_t i = 0; i < instructions.size(); ++i) { + const PipelineInstruction &instruction = instructions[i]; + out << " {\"id\": " << instruction.id + << ", \"iteration_offset\": 0, \"hardware\": " + << static_cast(instruction.device_type) + << ", \"resource\": " << instruction.execution_resource + << ", \"reads\": ["; + for (size_t j = 0; j < instruction.reads.size(); ++j) { + if (j != 0) + out << ", "; + out << "\"" << instruction.reads[j]->buffer->name << "\""; + } + out << "], \"writes\": ["; + for (size_t j = 0; j < instruction.writes.size(); ++j) { + if (j != 0) + out << ", "; + out << "\"" << instruction.writes[j]->buffer->name << "\""; + } + out << "]}" << (i + 1 == instructions.size() ? "\n" : ",\n"); + } + out << " ],\n \"edges\": [\n"; + size_t edge_index = 0; + size_t edge_count = ddg.semantic_edges.size() + ddg.ordering_edges.size(); + for (const SemanticDependencyEdge &edge : ddg.semantic_edges) { + out << " {\"source\": " << edge.source_instruction_id + << ", \"target\": " << edge.target_instruction_id << ", \"buffer\": \"" + << edge.buffer->name << "\", \"distance\": " << edge.distance + << ", \"kind\": \"" << SemanticDependencyKindName(edge.kind) << "\"}" + << (++edge_index == edge_count ? "\n" : ",\n"); + } + for (const TemplateOrderEdge &edge : ddg.ordering_edges) { + out << " {\"source\": " << edge.source_instruction_id + << ", \"target\": " << edge.target_instruction_id + << ", \"buffer\": null, \"distance\": " << edge.distance + << ", \"kind\": \"collective_order\"}" + << (++edge_index == edge_count ? "\n" : ",\n"); + } + out << " ],\n \"buffers\": [\n"; + for (size_t i = 0; i < ddg.buffer_order.size(); ++i) { + const BufferNode *buffer = ddg.buffer_order[i]; + const BufferAccessInfo &access = ddg.buffer_access_infos.at(buffer); + out << " {\"name\": \"" << buffer->name + << "\", \"global\": " << (access.is_global ? "true" : "false") + << ", \"loop_carried\": " + << (access.HasLoopCarriedDependence() ? "true" : "false") + << ", \"versioned\": " + << (versioned_buffers.count(buffer) ? "true" : "false") + << ", \"banked\": " + << (IsRuntimeBankedBuffer(buffer) ? "true" : "false") + << ", \"classification\": \"" + << (access.is_global + ? "global" + : (access.HasLoopCarriedDependence() ? "loop_carried" + : "local")) + << "\"}" << (i + 1 == ddg.buffer_order.size() ? "\n" : ",\n"); + } + out << " ]\n}\n"; +} + +static bool +VerifyScheduledWindow(const std::vector &expected, + const std::vector &scheduled, + int command_count) { + if (expected.size() != scheduled.size()) + return false; + std::multiset> expected_instances; + std::multiset> scheduled_instances; + for (const PipelineInstruction &instruction : expected) { + if (instruction.id < 0 || instruction.id >= command_count || + instruction.iter < 0) + return false; + expected_instances.insert({instruction.iter, instruction.id}); + } + for (const PipelineInstruction &instruction : scheduled) { + if (instruction.id < 0 || instruction.id >= command_count || + instruction.iter < 0) + return false; + scheduled_instances.insert({instruction.iter, instruction.id}); + } + return expected_instances == scheduled_instances; +} + +static bool +VerifyGreedySchedule(const std::vector &expected_prologue, + const std::vector &expected_body, + const std::vector &expected_epilogue, + const std::vector &prologue, + const std::vector &body, + const std::vector &epilogue, + bool has_epilogue, int command_count) { + if (!VerifyScheduledWindow(expected_prologue, prologue, command_count) || + !VerifyScheduledWindow(expected_body, body, command_count)) { + return false; + } + return !has_epilogue || + VerifyScheduledWindow(expected_epilogue, epilogue, command_count); +} + +static bool VerifyDynamicLogicalCoverage( + int extent, int iterations, int command_count, + const std::vector &prologue, + const std::vector &body, + const std::map> &epilogues) { + std::map, int> counts; + auto record = [&](int base, const std::vector &window, + bool predicate_invalid) { + for (const PipelineInstruction &instruction : window) { + int logical_iter = base + instruction.iter; + if (instruction.id < 0 || instruction.id >= command_count) { + return false; + } + if (logical_iter < 0 || logical_iter >= extent) { + if (predicate_invalid) + continue; + return false; + } + counts[{logical_iter, instruction.id}] += 1; + } + return true; + }; + if (!record(0, prologue, false)) + return false; + int steady_groups = std::max(0, (extent - 1) / iterations); + for (int group = 0; group < steady_groups; ++group) { + if (!record(group * iterations, body, false)) + return false; + } + int remainder = extent % iterations; + auto epilogue_it = epilogues.find(remainder); + if (epilogue_it == epilogues.end() || + !record(steady_groups * iterations, epilogue_it->second, true)) { + return false; + } + for (int logical_iter = 0; logical_iter < extent; ++logical_iter) { + for (int command = 0; command < command_count; ++command) { + if (counts[{logical_iter, command}] != 1) + return false; + } + } + return true; +} + /** * \brief Identify the prefetch instruction set on top of the local DDG. */ @@ -465,8 +1023,7 @@ class PrefetchInstructionIdentifier { } return call->op.same_as(Op::Get("tl.dma_copy")) || call->op.same_as(Op::Get("tl.broadcast_")) || - call->op.same_as(Op::Get("tl.sunmmio_layout_transform")) || - call->op.same_as(Op::Get("tl.sunmmio_transpose")); + call->op.same_as(Op::Get("tl.sunmmio_layout_transform")); } static bool WasBufferReadBeforeInstruction(const LocalDDG &local_ddg, @@ -673,7 +1230,7 @@ class PipelineWindowAssembler { class PipelineDevice { public: - explicit PipelineDevice(DeviceType type) : type(type) {} + explicit PipelineDevice(int resource) : resource(resource) {} void AssignInstruction(PipelineInstruction *instruction, float time) { current_instruction = instruction; @@ -692,7 +1249,7 @@ class PipelineDevice { } } - DeviceType type{DeviceType::Unspecified}; + int resource{-1}; bool busy{false}; PipelineInstruction *current_instruction{nullptr}; float instruction_end_time{std::numeric_limits::max()}; @@ -705,9 +1262,14 @@ class GlobalPipelineScheduler { bool debug_{false}; GlobalPipelineScheduler() { - devices_.push_back(PipelineDevice(DeviceType::ODMA)); - devices_.push_back(PipelineDevice(DeviceType::TensorCore)); - devices_.push_back(PipelineDevice(DeviceType::VectorCore)); + devices_.push_back( + PipelineDevice(static_cast(IlpResourceType::kTensorCore))); + devices_.push_back( + PipelineDevice(static_cast(IlpResourceType::kVectorCore))); + devices_.push_back( + PipelineDevice(static_cast(IlpResourceType::kODMA0))); + devices_.push_back( + PipelineDevice(static_cast(IlpResourceType::kODMA1))); } void SetVersionedBuffers( @@ -715,6 +1277,15 @@ class GlobalPipelineScheduler { versioned_buffers_ = versioned_buffers; } + void SetBankColoring(const GreedyBankColoring &coloring) { + writer_phases_ = coloring.writer_phases; + reader_phases_ = coloring.reader_phases; + } + + void SetTemplateOrderEdges(const std::vector &edges) { + template_order_edges_ = edges; + } + void BuildDependencyGraph() { int instruction_count = static_cast(instructions.size()); predecessors_.assign(instruction_count, {}); @@ -737,6 +1308,7 @@ class GlobalPipelineScheduler { BufferRegion region; int instruction_index; AccessType type; + int instance_id; }; std::unordered_map> @@ -747,18 +1319,19 @@ class GlobalPipelineScheduler { int current_index = topological_order_[ordered_index]; const PipelineInstruction ¤t_instruction = instructions[current_index]; - int current_version = GetVersionId(current_instruction); for (const BufferRegion &read_region : current_instruction.reads) { const BufferNode *buffer = read_region->buffer.get(); + int current_instance = + GetAccessInstanceId(current_instruction, buffer, false); auto history_it = buffer_access_history.find(buffer); if (history_it == buffer_access_history.end()) { continue; } auto &history = history_it->second; for (auto it = history.rbegin(); it != history.rend(); ++it) { - if (ShouldSkipVersionedCrossIteration(buffer, current_version, - it->instruction_index)) { + if (ShouldSkipVersionedAccess(buffer, current_instance, + it->instance_id)) { continue; } if (it->type == AccessType::kWrite && @@ -771,14 +1344,16 @@ class GlobalPipelineScheduler { for (const BufferRegion &write_region : current_instruction.writes) { const BufferNode *buffer = write_region->buffer.get(); + int current_instance = + GetAccessInstanceId(current_instruction, buffer, true); auto history_it = buffer_access_history.find(buffer); if (history_it == buffer_access_history.end()) { continue; } auto &history = history_it->second; for (auto it = history.rbegin(); it != history.rend(); ++it) { - if (ShouldSkipVersionedCrossIteration(buffer, current_version, - it->instruction_index)) { + if (ShouldSkipVersionedAccess(buffer, current_instance, + it->instance_id)) { continue; } if (!AccessOverlapChecker::Overlap(write_region, it->region)) { @@ -793,11 +1368,35 @@ class GlobalPipelineScheduler { for (const BufferRegion &read_region : current_instruction.reads) { buffer_access_history[read_region->buffer.get()].push_back( - {read_region, current_index, AccessType::kRead}); + {read_region, current_index, AccessType::kRead, + GetAccessInstanceId(current_instruction, read_region->buffer.get(), + false)}); } for (const BufferRegion &write_region : current_instruction.writes) { buffer_access_history[write_region->buffer.get()].push_back( - {write_region, current_index, AccessType::kWrite}); + {write_region, current_index, AccessType::kWrite, + GetAccessInstanceId(current_instruction, + write_region->buffer.get(), true)}); + } + } + + std::map, int> instance_index; + for (int index = 0; index < instruction_count; ++index) { + instance_index[{instructions[index].iter, instructions[index].id}] = + index; + } + for (const TemplateOrderEdge &edge : template_order_edges_) { + for (int source_index = 0; source_index < instruction_count; + ++source_index) { + const PipelineInstruction &source = instructions[source_index]; + if (source.id != edge.source_instruction_id) { + continue; + } + auto target = instance_index.find( + {source.iter + edge.distance, edge.target_instruction_id}); + if (target != instance_index.end()) { + AddDependency(source_index, target->second); + } } } } @@ -835,7 +1434,7 @@ class GlobalPipelineScheduler { } log_file << instruction_index << " " << instruction.name << " " << instruction.iter << " " << instruction.id << " " - << static_cast(instruction.device_type) << " " + << instruction.execution_resource << " " << static_cast(instruction.is_prefetch) << " " << bottom_level << "\n"; } @@ -888,9 +1487,14 @@ class GlobalPipelineScheduler { if (!ArePredecessorsFinished(*instruction)) { continue; } + if (!AreBanksFree(*instruction, time)) { + continue; + } for (auto &device : devices_) { - if (device.type == instruction->device_type && !device.busy) { + if (device.resource == instruction->execution_resource && + !device.busy) { device.AssignInstruction(instruction, time); + ReserveBanks(*instruction, instruction->scheduled_end); schedule.push_back(*instruction); break; } @@ -917,10 +1521,18 @@ class GlobalPipelineScheduler { float start; float end; }; - std::unordered_map> busy_intervals; + std::unordered_map> busy_intervals; + std::array, static_cast(PhysicalSramBank::Count)> + bank_busy_intervals; for (const auto &instruction : schedule) { - busy_intervals[instruction.device_type].push_back( + busy_intervals[instruction.execution_resource].push_back( {instruction.scheduled_start, instruction.scheduled_end}); + for (PhysicalSramBank bank : + GetOccupiedSramBanks(instruction, versioned_buffers_, iter_mod_, + writer_phases_, reader_phases_)) { + bank_busy_intervals[static_cast(bank)].push_back( + {instruction.scheduled_start, instruction.scheduled_end}); + } } for (auto &kv : busy_intervals) { auto &intervals = kv.second; @@ -929,6 +1541,12 @@ class GlobalPipelineScheduler { return lhs.start < rhs.start; }); } + for (auto &intervals : bank_busy_intervals) { + std::sort(intervals.begin(), intervals.end(), + [](const Interval &lhs, const Interval &rhs) { + return lhs.start < rhs.start; + }); + } std::sort(prefetch_queue.begin(), prefetch_queue.end(), [](PipelineInstruction *lhs, PipelineInstruction *rhs) { @@ -984,25 +1602,43 @@ class GlobalPipelineScheduler { } float duration = instruction->delay; - auto &intervals = busy_intervals[instruction->device_type]; + auto &intervals = busy_intervals[instruction->execution_resource]; float start_time = ready_time; - for (size_t i = 0; i <= intervals.size(); ++i) { - float gap_end = (i < intervals.size()) - ? intervals[i].start - : std::numeric_limits::max(); - if (gap_end - start_time >= duration) { - instruction->scheduled_start = start_time; - instruction->scheduled_end = start_time + duration; - instruction->finished = true; - insert_interval(intervals, {instruction->scheduled_start, - instruction->scheduled_end}); - schedule.push_back(*instruction); - scheduled_prefetch += 1; - break; + std::vector *> required_intervals{&intervals}; + for (PhysicalSramBank bank : + GetOccupiedSramBanks(*instruction, versioned_buffers_, iter_mod_, + writer_phases_, reader_phases_)) { + required_intervals.push_back( + &bank_busy_intervals[static_cast(bank)]); + } + while (!instruction->finished) { + float next_start = start_time; + for (const std::vector *resource_intervals : + required_intervals) { + for (const Interval &interval : *resource_intervals) { + if (start_time + duration <= interval.start) { + break; + } + if (start_time < interval.end && + start_time + duration > interval.start) { + next_start = std::max(next_start, interval.end); + break; + } + } } - if (i < intervals.size()) { - start_time = std::max(start_time, intervals[i].end); + if (next_start != start_time) { + start_time = next_start; + continue; + } + instruction->scheduled_start = start_time; + instruction->scheduled_end = start_time + duration; + instruction->finished = true; + for (std::vector *resource_intervals : required_intervals) { + insert_interval(*resource_intervals, {instruction->scheduled_start, + instruction->scheduled_end}); } + schedule.push_back(*instruction); + scheduled_prefetch += 1; } ICHECK(instruction->finished) << "Failed to insert prefetch instruction " << instruction->name; @@ -1027,12 +1663,17 @@ class GlobalPipelineScheduler { if (lhs.scheduled_start != rhs.scheduled_start) { return lhs.scheduled_start < rhs.scheduled_start; } + int lhs_priority = GetGreedyIssuePriority(lhs); + int rhs_priority = GetGreedyIssuePriority(rhs); + if (lhs_priority != rhs_priority) { + return lhs_priority < rhs_priority; + } return lhs.name < rhs.name; }); if (debug_ && log_file.is_open()) { for (const auto &instruction : schedule) { log_file << (instruction.is_prefetch ? "p:" : "") << instruction.name - << " " << static_cast(instruction.device_type) << " " + << " " << instruction.execution_resource << " " << instruction.scheduled_start << " " << instruction.delay << "\n"; } @@ -1041,14 +1682,24 @@ class GlobalPipelineScheduler { } private: - bool ShouldSkipVersionedCrossIteration(const BufferNode *buffer, - int current_version, - int previous_instruction_index) const { + bool ShouldSkipVersionedAccess(const BufferNode *buffer, int current_instance, + int previous_instance) const { if (versioned_buffers_.count(buffer) == 0) { return false; } - return GetVersionId(instructions[previous_instruction_index]) != - current_version; + return previous_instance != current_instance; + } + + int GetAccessInstanceId(const PipelineInstruction &instruction, + const BufferNode *buffer, bool is_write) const { + int slot = GetVersionId(instruction); + if (!IsRuntimeBankedBuffer(buffer)) + return slot; + int phase = LookupBankPhase(is_write ? writer_phases_ : reader_phases_, + buffer, instruction.id); + int bank = (slot + phase) & 1; + int version = slot / 2; + return bank * std::max(1, iter_mod_) + version; } int GetVersionId(const PipelineInstruction &instruction) const { @@ -1079,6 +1730,26 @@ class GlobalPipelineScheduler { device.current_instruction = nullptr; device.instruction_end_time = std::numeric_limits::max(); } + bank_busy_until_.fill(-1.0f); + } + + bool AreBanksFree(const PipelineInstruction &instruction, float time) const { + for (PhysicalSramBank bank : + GetOccupiedSramBanks(instruction, versioned_buffers_, iter_mod_, + writer_phases_, reader_phases_)) { + if (bank_busy_until_[static_cast(bank)] > time) { + return false; + } + } + return true; + } + + void ReserveBanks(const PipelineInstruction &instruction, float end_time) { + for (PhysicalSramBank bank : + GetOccupiedSramBanks(instruction, versioned_buffers_, iter_mod_, + writer_phases_, reader_phases_)) { + bank_busy_until_[static_cast(bank)] = end_time; + } } bool ArePredecessorsFinished(const PipelineInstruction &instruction) const { @@ -1096,7 +1767,12 @@ class GlobalPipelineScheduler { } std::vector devices_; + std::array(PhysicalSramBank::Count)> + bank_busy_until_{}; std::unordered_set versioned_buffers_; + PerCommandBankPhases writer_phases_; + PerCommandBankPhases reader_phases_; + std::vector template_order_edges_; std::vector> predecessors_; std::vector> successors_; std::vector topological_order_; @@ -1125,6 +1801,22 @@ class SunmmioPipelinePlanner : public StmtExprMutator { if (num_stages <= 0) { return StmtExprMutator::VisitStmt_(op); } + arith::Analyzer analyzer; + PrimExpr simplified_extent = analyzer.Simplify(op->extent); + const auto *extent = simplified_extent.as(); + if (extent != nullptr && extent->value < num_stages) { + For fallback = Downcast(StmtExprMutator::VisitStmt_(op)); + return MakePipelineFallback(fallback, "greedy", "planning", + "short_extent_unsupported"); + } + if (extent == nullptr && num_stages != 2) { + For fallback = Downcast(StmtExprMutator::VisitStmt_(op)); + return MakePipelineFallback( + fallback, + PipelineDiagnostic{false, "greedy", "planning", + "dynamic_version_count_unsupported", + "dynamic Greedy currently requires num_stages=2"}); + } // 2. Peel off the outer layers to find the true body sequence auto inner_stmt = op->body; @@ -1144,15 +1836,24 @@ class SunmmioPipelinePlanner : public StmtExprMutator { const SeqStmtNode *pipeline_body_seq = inner_stmt.as(); ICHECK(pipeline_body_seq) << "Pipeline body must be a SeqStmt"; ICHECK(op->kind == ForKind::kSerial) << "Pipeline loop must be serial"; - // 3. Stage 1: Build the PipelineInstruction containers std::vector single_iteration_instructions; for (size_t i = 0; i < pipeline_body_seq->seq.size(); ++i) { - PipelineInstruction instruction(static_cast(i), 0, - pipeline_body_seq->seq[i]); + const Stmt &stmt = pipeline_body_seq->seq[i]; + if (!stmt.as() && !stmt.as() && + !stmt.as()) { + // HardwareMapper intentionally handles hardware commands only. Scalar + // bookkeeping stores and conditional command groups must first be + // normalized before they can be scheduled safely. + For fallback = Downcast(StmtExprMutator::VisitStmt_(op)); + return MakePipelineFallback(fallback, "greedy", "planning", + "unsupported_statement"); + } + PipelineInstruction instruction(static_cast(i), 0, stmt); instruction.device_type = HardwareMapper::Map(instruction.stmt); instruction.ExtractRegions(stmt_rw_collector_); + instruction.execution_resource = GetGreedyExecutionResource(instruction); instruction.delay = CostModel::EstimateDelay(instruction.device_type, instruction.stmt); single_iteration_instructions.push_back(instruction); @@ -1166,6 +1867,7 @@ class SunmmioPipelinePlanner : public StmtExprMutator { for (const auto &instruction : single_iteration_instructions) { std::cout << " - ID: " << instruction.id << ", Device: " << static_cast(instruction.device_type) + << ", Resource: " << instruction.execution_resource << ", Delay: " << instruction.delay << ", Reads: " << instruction.reads.size() << ", Writes: " << instruction.writes.size() << "\n"; @@ -1174,6 +1876,11 @@ class SunmmioPipelinePlanner : public StmtExprMutator { // 4. Stage 2.1: Build the local DDG for a single iteration. LocalDDG local_ddg = LocalDDGBuilder::Build(single_iteration_instructions); + if (!ValidateLocalDDG(single_iteration_instructions, local_ddg)) { + For fallback = Downcast(StmtExprMutator::VisitStmt_(op)); + return MakePipelineFallback(fallback, "greedy", "graph_validation", + "incomplete_access_info"); + } if (debug_) { int forward_edge_count = 0; @@ -1225,6 +1932,8 @@ class SunmmioPipelinePlanner : public StmtExprMutator { std::unordered_set versioned_buffers = MultiversioningIdentifier::Identify(single_iteration_instructions, local_ddg); + MaybeWriteGreedyGraphJson(single_iteration_instructions, local_ddg, + versioned_buffers); if (debug_) { std::cout << "[Pipeline Planner] Identified " << versioned_buffers.size() @@ -1263,11 +1972,50 @@ class SunmmioPipelinePlanner : public StmtExprMutator { } // 8. Stage 4: Build the global DDG and run the two-phase scheduler. + int faster = -1; + auto pass_ctx = tvm::transform::PassContext::Current(); + auto faster_config = pass_ctx->GetConfig(kSunmmioFaster); + if (faster_config) { + faster = faster_config.value()->value; + } else if (const char *env_faster = std::getenv("TL_SUNMMIO_FASTER")) { + faster = std::stoi(env_faster); + } + size_t total_coloring_candidates = 0; + std::vector coloring_candidates = + BuildGreedyBankColorings(local_ddg, versioned_buffers, faster, + &total_coloring_candidates); + GreedyBankColoring selected_coloring = coloring_candidates.front(); + float selected_body_makespan = std::numeric_limits::max(); + for (const GreedyBankColoring &candidate : coloring_candidates) { + GlobalPipelineScheduler candidate_scheduler; + candidate_scheduler.instructions = stage_assembly.body_instructions; + candidate_scheduler.iter_mod_ = stage_assembly.iterations; + candidate_scheduler.SetVersionedBuffers(versioned_buffers); + candidate_scheduler.SetBankColoring(candidate); + candidate_scheduler.SetTemplateOrderEdges(local_ddg.ordering_edges); + candidate_scheduler.BuildDependencyGraph(); + candidate_scheduler.CalculateBottomLevels(); + std::vector candidate_schedule = + candidate_scheduler.Schedule(""); + float makespan = 0.0f; + for (const PipelineInstruction &instruction : candidate_schedule) { + makespan = std::max(makespan, instruction.scheduled_end); + } + if (makespan < selected_body_makespan || + (makespan == selected_body_makespan && + candidate.bits < selected_coloring.bits)) { + selected_body_makespan = makespan; + selected_coloring = candidate; + } + } + GlobalPipelineScheduler prologue_scheduler; prologue_scheduler.instructions = stage_assembly.prologue_instructions; prologue_scheduler.iter_mod_ = stage_assembly.iterations; prologue_scheduler.debug_ = debug_; prologue_scheduler.SetVersionedBuffers(versioned_buffers); + prologue_scheduler.SetBankColoring(selected_coloring); + prologue_scheduler.SetTemplateOrderEdges(local_ddg.ordering_edges); prologue_scheduler.BuildDependencyGraph(); prologue_scheduler.CalculateBottomLevels(); std::vector prologue_schedule = @@ -1278,6 +2026,8 @@ class SunmmioPipelinePlanner : public StmtExprMutator { body_scheduler.iter_mod_ = stage_assembly.iterations; body_scheduler.debug_ = debug_; body_scheduler.SetVersionedBuffers(versioned_buffers); + body_scheduler.SetBankColoring(selected_coloring); + body_scheduler.SetTemplateOrderEdges(local_ddg.ordering_edges); body_scheduler.BuildDependencyGraph(); body_scheduler.CalculateBottomLevels(); body_scheduler.DumpGraph("body_graph.log"); @@ -1285,15 +2035,84 @@ class SunmmioPipelinePlanner : public StmtExprMutator { body_scheduler.Schedule("body.log"); std::vector epilogue_schedule; + std::map> dynamic_epilogue_schedules; if (stage_assembly.epilogue_iterations != -1) { GlobalPipelineScheduler epilogue_scheduler; epilogue_scheduler.instructions = stage_assembly.epilogue_instructions; epilogue_scheduler.iter_mod_ = stage_assembly.iterations; epilogue_scheduler.debug_ = debug_; epilogue_scheduler.SetVersionedBuffers(versioned_buffers); + epilogue_scheduler.SetBankColoring(selected_coloring); + epilogue_scheduler.SetTemplateOrderEdges(local_ddg.ordering_edges); epilogue_scheduler.BuildDependencyGraph(); epilogue_scheduler.CalculateBottomLevels(); epilogue_schedule = epilogue_scheduler.Schedule("epilogue.log"); + } else { + for (int remainder = 0; remainder < stage_assembly.iterations; + ++remainder) { + int effective_remainder = + remainder == 0 ? stage_assembly.iterations : remainder; + PipelineStageAssembly remainder_assembly = + PipelineWindowAssembler::Assemble(single_iteration_instructions, + num_stages, + Integer(effective_remainder)); + GlobalPipelineScheduler epilogue_scheduler; + epilogue_scheduler.instructions = + remainder_assembly.epilogue_instructions; + epilogue_scheduler.iter_mod_ = stage_assembly.iterations; + epilogue_scheduler.SetVersionedBuffers(versioned_buffers); + epilogue_scheduler.SetBankColoring(selected_coloring); + epilogue_scheduler.SetTemplateOrderEdges(local_ddg.ordering_edges); + epilogue_scheduler.BuildDependencyGraph(); + epilogue_scheduler.CalculateBottomLevels(); + dynamic_epilogue_schedules[remainder] = epilogue_scheduler.Schedule(""); + if (!VerifyScheduledWindow( + remainder_assembly.epilogue_instructions, + dynamic_epilogue_schedules[remainder], + static_cast(single_iteration_instructions.size()))) { + For fallback = Downcast(StmtExprMutator::VisitStmt_(op)); + return MakePipelineFallback( + fallback, + PipelineDiagnostic{false, "greedy", "schedule_validation", + "invalid_schedule_order", + "dynamic epilogue remainder " + + std::to_string(remainder)}); + } + } + } + + if (!VerifyGreedySchedule( + stage_assembly.prologue_instructions, + stage_assembly.body_instructions, + stage_assembly.epilogue_instructions, prologue_schedule, + body_schedule, epilogue_schedule, + stage_assembly.epilogue_iterations != -1, + static_cast(single_iteration_instructions.size()))) { + For fallback = Downcast(StmtExprMutator::VisitStmt_(op)); + return MakePipelineFallback( + fallback, + PipelineDiagnostic{false, "greedy", "schedule_validation", + "invalid_schedule_order", + "scheduled window does not match assembled " + "logical command instances"}); + } + if (stage_assembly.epilogue_iterations == -1) { + int verification_extent = std::max(4, stage_assembly.iterations * 2 + 2); + for (int extent_value = 1; extent_value <= verification_extent; + ++extent_value) { + if (!VerifyDynamicLogicalCoverage( + extent_value, stage_assembly.iterations, + static_cast(single_iteration_instructions.size()), + prologue_schedule, body_schedule, dynamic_epilogue_schedules)) { + For fallback = Downcast(StmtExprMutator::VisitStmt_(op)); + return MakePipelineFallback( + fallback, + PipelineDiagnostic{false, "greedy", "schedule_validation", + "logical_iteration_out_of_bounds", + "coverage failed for representative extent " + + std::to_string(extent_value)}); + } + } } if (debug_) { @@ -1311,6 +2130,33 @@ class SunmmioPipelinePlanner : public StmtExprMutator { } } annotations.Set("iterations", stage_assembly.iterations); + SetPipelineAppliedAnnotations(&annotations, "greedy"); + annotations.Set("coloring_total_candidates", + Integer(total_coloring_candidates)); + annotations.Set("coloring_evaluated_candidates", + Integer(coloring_candidates.size())); + + Map> runtime_bank_writer_phases; + for (const auto &[buffer_node, phases] : selected_coloring.writer_phases) { + Map per_command; + for (const auto &[command_id, phase] : phases) { + per_command.Set(Integer(command_id), Integer(phase)); + } + runtime_bank_writer_phases.Set(tvm::ffi::GetRef(buffer_node), + per_command); + } + annotations.Set("runtime_bank_writer_phases", runtime_bank_writer_phases); + + Map> runtime_bank_reader_phases; + for (const auto &[buffer_node, phases] : selected_coloring.reader_phases) { + Map per_command; + for (const auto &[command_id, phase] : phases) { + per_command.Set(Integer(command_id), Integer(phase)); + } + runtime_bank_reader_phases.Set(tvm::ffi::GetRef(buffer_node), + per_command); + } + annotations.Set("runtime_bank_reader_phases", runtime_bank_reader_phases); Array orders; for (const auto &instruction : prologue_schedule) { @@ -1330,6 +2176,16 @@ class SunmmioPipelinePlanner : public StmtExprMutator { orders.push_back(instruction.name); } annotations.Set("epilogue_orders", orders); + } else { + Map> dynamic_orders; + for (const auto &[remainder, schedule] : dynamic_epilogue_schedules) { + Array remainder_orders; + for (const PipelineInstruction &instruction : schedule) { + remainder_orders.push_back(instruction.name); + } + dynamic_orders.Set(Integer(remainder), remainder_orders); + } + annotations.Set("dynamic_epilogue_orders", dynamic_orders); } Array used_buffers; diff --git a/src/transform/sunmmio_pipeline_planning/pipeline_diagnostic.h b/src/transform/sunmmio_pipeline_planning/pipeline_diagnostic.h new file mode 100644 index 0000000000..3c881f5b7c --- /dev/null +++ b/src/transform/sunmmio_pipeline_planning/pipeline_diagnostic.h @@ -0,0 +1,185 @@ +#ifndef TILELANG_TRANSFORM_SUNMMIO_PIPELINE_DIAGNOSTIC_H_ +#define TILELANG_TRANSFORM_SUNMMIO_PIPELINE_DIAGNOSTIC_H_ + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace tvm { +namespace tl { + +using namespace tir; + +constexpr const char *kPipelineRequested = "tl.sunmmio.pipeline.requested"; +constexpr const char *kPipelineApplied = "tl.sunmmio.pipeline.applied"; +constexpr const char *kPipelineMode = "tl.sunmmio.pipeline.mode"; +constexpr const char *kPipelineFallbackStage = + "tl.sunmmio.pipeline.fallback_stage"; +constexpr const char *kPipelineFallbackReason = + "tl.sunmmio.pipeline.fallback_reason"; +constexpr const char *kPipelineFallbackDetail = + "tl.sunmmio.pipeline.fallback_detail"; + +struct PipelineDiagnostic { + bool applied{false}; + std::string mode; + std::string stage; + std::string reason; + std::string detail; +}; + +inline void SetPipelineAppliedAnnotations(Map *annotations, + const std::string &mode) { + annotations->Set(kPipelineRequested, Bool(true)); + annotations->Set(kPipelineApplied, Bool(true)); + annotations->Set(kPipelineMode, String(mode)); + annotations->erase(kPipelineFallbackStage); + annotations->erase(kPipelineFallbackReason); + annotations->erase(kPipelineFallbackDetail); +} + +inline For MakePipelineFallback(const For &loop, + const PipelineDiagnostic &diagnostic, + bool emit_warning = true) { + Map annotations = loop->annotations; + annotations.Set(kPipelineRequested, Bool(true)); + annotations.Set(kPipelineApplied, Bool(diagnostic.applied)); + annotations.Set(kPipelineMode, String(diagnostic.mode)); + annotations.Set(kPipelineFallbackStage, String(diagnostic.stage)); + annotations.Set(kPipelineFallbackReason, String(diagnostic.reason)); + if (!diagnostic.detail.empty()) { + annotations.Set(kPipelineFallbackDetail, String(diagnostic.detail)); + } + if (emit_warning) { + LOG(WARNING) << "[SunmmioPipeline][" << diagnostic.mode + << "][Fallback] stage=" << diagnostic.stage + << " reason=" << diagnostic.reason + << " detail=" << diagnostic.detail + << " extent=" << loop->extent; + } + For fallback = loop; + fallback.CopyOnWrite()->annotations = annotations; + return fallback; +} + +inline For MakePipelineFallback(const For &loop, const std::string &mode, + const std::string &stage, + const std::string &reason, + bool emit_warning = true) { + return MakePipelineFallback( + loop, PipelineDiagnostic{false, mode, stage, reason, ""}, emit_warning); +} + +inline bool IsPipelineScheduleAnnotation(const String &key) { + static const std::unordered_set keys = { + "num_stages", + "iterations", + "ii", + "makespan", + "stage_count", + "steady_state_max_iter_offset", + "prologue_orders", + "body_orders", + "epilogue_orders", + "dynamic_epilogue_orders", + "used_buffers", + "versioned_buffers", + "bank_peer_buffers", + "version_axis_buffers", + "runtime_multiversion_buffers", + "runtime_banked_buffers", + "runtime_resident_banked_buffers", + "runtime_bank_peer_buffers", + "runtime_bank_start_phases", + "runtime_bank_read_delta_parities", + "runtime_bank_writer_phases", + "runtime_bank_reader_phases", + "runtime_bank_flip_modes", + }; + return keys.count(std::string(key)) != 0; +} + +class PipelineAtomicFallbackRewriter : public StmtMutator { +public: + PipelineAtomicFallbackRewriter(std::string mode, std::string stage, + std::string reason, std::string detail) + : diagnostic_{false, std::move(mode), std::move(stage), std::move(reason), + std::move(detail)} {} + + bool changed() const { return changed_; } + +private: + Stmt VisitStmt_(const ForNode *op) final { + For loop = Downcast(StmtMutator::VisitStmt_(op)); + if (!op->annotations.count(kPipelineRequested) && + !op->annotations.count("prologue_orders") && + !op->annotations.count("body_orders")) { + return loop; + } + Map annotations; + for (const auto &kv : loop->annotations) { + if (!IsPipelineScheduleAnnotation(kv.first) && + kv.first != kPipelineApplied && kv.first != kPipelineFallbackStage && + kv.first != kPipelineFallbackReason && + kv.first != kPipelineFallbackDetail) { + annotations.Set(kv.first, kv.second); + } + } + loop.CopyOnWrite()->annotations = annotations; + changed_ = true; + return MakePipelineFallback(loop, diagnostic_, false); + } + + PipelineDiagnostic diagnostic_; + bool changed_{false}; +}; + +inline PrimFunc +MakePipelineFunctionFallback(const PrimFunc &original, + const PipelineDiagnostic &diagnostic) { + PipelineAtomicFallbackRewriter rewriter(diagnostic.mode, diagnostic.stage, + diagnostic.reason, diagnostic.detail); + PrimFunc fallback = original; + auto *fptr = fallback.CopyOnWrite(); + fptr->body = rewriter(original->body); + LOG(WARNING) << "[SunmmioPipeline][" << diagnostic.mode + << "][AtomicFallback] stage=" << diagnostic.stage + << " reason=" << diagnostic.reason + << " detail=" << diagnostic.detail; + return fallback; +} + +class PipelineFallbackValidator : public StmtVisitor { +public: + static Optional FindDisallowed(const Stmt &body) { + PipelineFallbackValidator validator; + validator(body); + return validator.reason_; + } + +private: + void VisitStmt_(const ForNode *op) final { + auto reason = op->annotations.Get(kPipelineFallbackReason); + if (reason && !reason_) { + String value = Downcast(reason.value()); + if (value != "runtime_short_extent") { + reason_ = value; + } + } + StmtVisitor::VisitStmt_(op); + } + + Optional reason_; +}; + +} // namespace tl +} // namespace tvm + +#endif // TILELANG_TRANSFORM_SUNMMIO_PIPELINE_DIAGNOSTIC_H_ diff --git a/src/transform/sunmmio_pipeline_planning/resource_types_for_ilp.h b/src/transform/sunmmio_pipeline_planning/resource_types_for_ilp.h new file mode 100644 index 0000000000..877852eaca --- /dev/null +++ b/src/transform/sunmmio_pipeline_planning/resource_types_for_ilp.h @@ -0,0 +1,137 @@ +/*! + * \file resource_types_for_ilp.h + * \brief ILP-specific resource typing and resource extraction helpers. + */ +#ifndef TVM_TL_TRANSFORM_SUNMMIO_PIPELINE_PLANNING_RESOURCE_TYPES_FOR_ILP_H_ +#define TVM_TL_TRANSFORM_SUNMMIO_PIPELINE_PLANNING_RESOURCE_TYPES_FOR_ILP_H_ + +#include "../../op/utils.h" +#include "../../target/sunmmio/hardware_types.h" + +#include +#include +#include + +namespace tvm { +namespace tl { + +using namespace tir; + +enum class IlpResourceType : int { + kTensorCore = 0, + kVectorCore = 1, + kODMA0 = 2, + kODMA1 = 3, + kWsramIn = 6, + kWsramOut = 7, + kAsramIn = 8, + kAsramOut = 9, + // kRsram = 10, +}; + +template +std::vector +BuildIlpResources(const Stmt &stmt, DeviceType type, + const std::vector &accesses) { + std::vector resources; + auto add_resource = [&](int resource) { + if (std::find(resources.begin(), resources.end(), resource) == + resources.end()) { + resources.push_back(resource); + } + }; + + auto has_scope_read = [&](const char *scope) { + for (const AccessInfoLike &access : accesses) { + if (!access.is_write && access.buffer().scope() == scope) { + return true; + } + } + return false; + }; + + if (type == DeviceType::TensorCore) { + add_resource(static_cast(IlpResourceType::kTensorCore)); + if (has_scope_read("shared.wsram")) { + add_resource(static_cast(IlpResourceType::kWsramOut)); + } + if (has_scope_read("shared.asram")) { + add_resource(static_cast(IlpResourceType::kAsramOut)); + } + return resources; + } + + if (type == DeviceType::VectorCore) { + add_resource(static_cast(IlpResourceType::kVectorCore)); + return resources; + } + + if (const auto *eval = stmt.as()) { + if (const auto *call = eval->value.as()) { + if (call->op.same_as(Op::Get("tl.dma_copy")) || + call->op.same_as(Op::Get("tl.broadcast_")) || + call->op.same_as(Op::Get("tl.sunmmio_layout_transform"))) { + BufferRegion src_region = NormalizeToBufferRegion(call->args[0]); + // broadcast_ argument layout is: + // args[0] = src region + // args[1] = dst region + // args[2] = direction + // Resource typing only needs the source/destination buffers, so always + // read the destination from args[1]. Using args[2] treats the integer + // direction enum as a BufferRegion and crashes ILP planning. + BufferRegion dst_region = NormalizeToBufferRegion(call->args[1]); + if (IsGlobalBuffer(src_region->buffer)) { + if (dst_region->buffer.scope() == "shared.asram") { + LOG(FATAL) + << "ILP graph does not model DRAM -> ASRAM dma path yet."; + } + if (dst_region->buffer.scope() == "shared.wsram") { + add_resource(static_cast(IlpResourceType::kODMA0)); + add_resource(static_cast(IlpResourceType::kWsramIn)); + return resources; + } + if (dst_region->buffer.scope() == "shared.rsram" || + dst_region->buffer.scope() == "local") { + add_resource(static_cast(IlpResourceType::kODMA0)); + // add_resource(static_cast(IlpResourceType::kRsram)); + return resources; + } + } + if ((src_region->buffer.scope() == "shared.rsram" || + src_region->buffer.scope() == "local") && + dst_region->buffer.scope() == "shared.asram") { + add_resource(static_cast(IlpResourceType::kODMA1)); + add_resource(static_cast(IlpResourceType::kAsramIn)); + // add_resource(static_cast(IlpResourceType::kRsram)); + return resources; + } + add_resource(static_cast(IlpResourceType::kODMA0)); + if (src_region->buffer.scope() == "shared.rsram" || + src_region->buffer.scope() == "local" || + dst_region->buffer.scope() == "shared.rsram" || + dst_region->buffer.scope() == "local") { + // add_resource(static_cast(IlpResourceType::kRsram)); + } + if (dst_region->buffer.scope() == "shared.wsram") { + add_resource(static_cast(IlpResourceType::kWsramIn)); + } + if (dst_region->buffer.scope() == "shared.asram") { + add_resource(static_cast(IlpResourceType::kAsramIn)); + } + if (src_region->buffer.scope() == "shared.wsram") { + add_resource(static_cast(IlpResourceType::kWsramOut)); + } + if (src_region->buffer.scope() == "shared.asram") { + add_resource(static_cast(IlpResourceType::kAsramOut)); + } + return resources; + } + } + } + return resources; +} + +} // namespace tl +} // namespace tvm + +#endif // TVM_TL_TRANSFORM_SUNMMIO_PIPELINE_PLANNING_RESOURCE_TYPES_FOR_ILP_H_ diff --git a/src/transform/sunmmio_pipeline_planning/stmt_read_write_collector.h b/src/transform/sunmmio_pipeline_planning/stmt_read_write_collector.h index 879ca25478..2e4f8bc776 100644 --- a/src/transform/sunmmio_pipeline_planning/stmt_read_write_collector.h +++ b/src/transform/sunmmio_pipeline_planning/stmt_read_write_collector.h @@ -229,22 +229,27 @@ class StmtReadWriteCollector : public StmtVisitor { void VisitStmt_(const EvaluateNode *op) { const CallNode *call = op->value.as(); - if (call->op.same_as(dma_copy())) { + if (call && call->op.same_as(dma_copy())) { read_buffer_regions_.insert(NormalizeToBufferRegion(call->args[0])); write_buffer_regions_.insert(NormalizeToBufferRegion(call->args[1])); - } else if (call->op.same_as(sunmmio_layout_transform()) || - call->op.same_as(sunmmio_transpose())) { + } else if (call && (call->op.same_as(sunmmio_layout_transform()) || + call->op.same_as(sunmmio_transpose()))) { read_buffer_regions_.insert(NormalizeToBufferRegion(call->args[0])); write_buffer_regions_.insert(NormalizeToBufferRegion(call->args[1])); - } else if (call->op.same_as(mma_sunmmio())) { + } else if (call && call->op.same_as(mma_sunmmio())) { read_buffer_regions_.insert(NormalizeToBufferRegion(call->args[0])); read_buffer_regions_.insert(NormalizeToBufferRegion(call->args[1])); read_buffer_regions_.insert(NormalizeToBufferRegion(call->args[2])); write_buffer_regions_.insert(NormalizeToBufferRegion(call->args[2])); - } else if (call->op.same_as(Op::Get("tl.broadcast_"))) { + } else if (call && call->op.same_as(Op::Get("tl.broadcast_"))) { read_buffer_regions_.insert(NormalizeToBufferRegion(call->args[0])); write_buffer_regions_.insert(NormalizeToBufferRegion(call->args[1])); + } else if (call && + call->op.same_as(Op::Get("tl.vector_core_in_tile_reduce"))) { + ICHECK_GE(call->args.size(), 3U); + write_buffer_regions_.insert(NormalizeToBufferRegion(call->args[1])); + read_buffer_regions_.insert(NormalizeToBufferRegion(call->args[2])); } else { auto [read_regions, write_regions] = buffer_region_collector(op->value); for (auto it : read_regions) { diff --git a/src/transform/sunmmio_pipeline_planning_ilp.cc b/src/transform/sunmmio_pipeline_planning_ilp.cc new file mode 100644 index 0000000000..a1e618f912 --- /dev/null +++ b/src/transform/sunmmio_pipeline_planning_ilp.cc @@ -0,0 +1,4339 @@ +// Builds an ILP model that schedules a SunMMIO pipeline while respecting +// command dependencies, hardware-resource capacities, and SRAM bank conflicts. + +#include "../op/builtin.h" +#include "../op/comm.h" +#include "../op/utils.h" +#include "../target/sunmmio/cost_model.h" +#include "../target/sunmmio/hardware_types.h" +#include "common/ast_traverser.h" +#include "common/sunmmio_pipeline_utils.h" +#include "sunmmio_pipeline_planning/pipeline_diagnostic.h" +#include "sunmmio_pipeline_planning/resource_types_for_ilp.h" +#include "tvm/arith/pattern.h" +#include "tvm/ffi/reflection/registry.h" +#include "tvm/runtime/logging.h" +#include "tvm/tir/analysis.h" +#include "tvm/tir/buffer.h" +#include "tvm/tir/expr.h" +#include "tvm/tir/function.h" +#include "tvm/tir/op.h" +#include "tvm/tir/stmt.h" +#include "tvm/tir/stmt_functor.h" +#include "tvm/tir/transform.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace tvm { +namespace tl { +namespace bank_ilp_internal { + +using namespace tir; + +enum class Role : uint8_t { kConsumer, kProducer, kBoth, kUndefined }; + +struct BankFlipMode { + bool wsram_flip{true}; + bool asram_flip{true}; + + bool FlipForMem(int mem) const { + ICHECK(mem == 0 || mem == 1); + return mem == 0 ? wsram_flip : asram_flip; + } + + int Id() const { return (wsram_flip ? 2 : 0) | (asram_flip ? 1 : 0); } +}; + +struct AccessInfo { + BufferRegion region; + bool is_write{false}; + int iter_offset{0}; + + const Buffer &buffer() const { return region->buffer; } +}; + +struct CommandSpec { + int latency{0}; + std::vector resources; + std::string name; +}; + +// A FlowSpec is one logical lifetime of data in WSRAM (mem == 0) or ASRAM +// (mem == 1). It connects the command that produces the data to the command +// that consumes it, possibly in a later logical iteration (delta). The +// write/read offsets and durations describe the exact bank-occupancy windows +// relative to those command starts. resident flows enter the loop already +// live and may have a fixed physical bank; non-resident flows are assigned a +// ping/pong phase by the ILP. precolor records a bank relation known before +// solving, while fp and initial_time carry the phase/lifetime information used +// to compare resident and loop-internal data. +struct FlowSpec { + bool resident{false}; + int prod{-1}; + int cons{-1}; + int delta{0}; + int mem{0}; + std::string buffer_name; + int fixed_bank{-1}; + int precolor{-1}; + int fp{1}; + int initial_time{0}; + int w_off{0}; + int w_dur{0}; + int r_off{0}; + int r_dur{0}; + int write_resource{-1}; + int read_resource{-1}; +}; + +using SameWriteFlowKey = std::tuple; + +SameWriteFlowKey MakeSameWriteFlowKey(const FlowSpec &flow) { + return std::make_tuple(flow.prod, flow.mem); +} + +// Complete scheduling problem for one template iteration. P contains the N +// hardware commands; dep_edges and delta encode producer-to-consumer ordering +// across logical iterations; R/cap describe execution resources and their +// parallel capacities. flows adds the physical SRAM-bank lifetimes that are +// not expressible as ordinary command dependencies. The solver folds this +// infinite periodic schedule into an initiation-interval-sized time window. +struct Problem { + int N{0}; + int Tmax{0}; + std::vector R; + std::unordered_map cap; + std::vector P; + std::vector> dep_edges; + std::unordered_map delta; + std::vector flows; + std::vector versioned_buffer_names; +}; + +struct ModelVars { + HighsInt col_T{-1}; + std::vector col_t; + std::vector col_y; + std::vector col_y_half; + std::vector col_start_parity; + std::vector col_m; + std::vector> col_x; + std::vector> col_a; + std::vector internal_flow_ids; + std::vector col_z; +}; + +struct SolveResult { + bool ok{false}; + int II{0}; + int makespan{0}; + int bank_slot_period{0}; + std::vector t; + std::vector m; + std::vector y; + std::vector internal_flow_ids; + std::vector z_bank; + BankFlipMode bank_flip_mode; + bool vc_blocking_issue_modeled{true}; + int vc_blocking_issue_constraints{0}; +}; + +struct SolutionVerifyResult { + bool ok{true}; + bool node_time_ok{true}; + bool dependency_ok{true}; + bool resource_slot_ok{true}; + bool bank_slot_ok{true}; + bool bank_port_ok{true}; + std::vector errors; + std::map> resource_slot_load; + std::array, 2>, 2> bank_slot_load; + std::array, 2>, 2> bank_port_load; +}; + +struct ExpandedOrderEntry { + int iter{0}; + int id{-1}; + int absolute_start{0}; +}; + +struct TimeWindowOrderResult { + std::vector prologue; + std::vector body; + std::vector epilogue; + int steady_state_max_iter_offset{0}; +}; + +int PositiveMod(int value, int mod); + +bool CommandUsesResource(const CommandSpec &spec, int resource); + +BufferRegion MaterializeBufferRegion(const BufferRegion ®ion, + const Var &loop_var, int iter); + +namespace { + +// File-local helper to avoid colliding with the similarly named function in +// sunmmio_pipeline_planning.cc during final shared-library link. +bool PipelineRegionIntersect(const Region ®ion1, const Region ®ion2) { + ICHECK(region1.size() == region2.size()); + for (size_t i = 0; i < region1.size(); ++i) { + const Range &dim1 = region1[i]; + const Range &dim2 = region2[i]; + auto int_set1 = arith::IntSet::FromRange(dim1); + auto int_set2 = arith::IntSet::FromRange(dim2); + if (arith::Intersect({int_set1, int_set2}).IsNothing()) { + return false; + } + } + return true; +} + +const double kInf = kHighsInf; + +long long EdgeKey(int i, int j) { + return (static_cast(i) << 32) | static_cast(j); +} + +int CeilDiv(int a, int b) { + ICHECK_GT(b, 0); + return (a + b - 1) / b; +} + +int FloorDiv(int a, int b) { + ICHECK_GT(b, 0); + int quotient = a / b; + int remainder = a % b; + return remainder < 0 ? quotient - 1 : quotient; +} + +int GcdInt(int a, int b) { + a = std::abs(a); + b = std::abs(b); + while (b != 0) { + int t = a % b; + a = b; + b = t; + } + return a; +} + +std::map>> +BuildResourceUsage(const std::vector &specs) { + std::map>> usage; + for (int idx = 0; idx < static_cast(specs.size()); ++idx) { + int latency = std::max(1, specs[idx].latency); + for (int resource : specs[idx].resources) { + usage[resource].push_back({idx, latency}); + } + } + return usage; +} + +std::map BuildResourceTotals( + const std::map>> &usage) { + std::map totals; + for (const auto &kv : usage) { + int total = 0; + for (const auto &item : kv.second) { + total += item.second; + } + totals[kv.first] = total; + } + return totals; +} + +std::pair FindTargetResource(const std::map &totals) { + int target_resource = -1; + int target_total = 0; + for (const auto &kv : totals) { + if (kv.second > target_total) { + target_resource = kv.first; + target_total = kv.second; + } + } + return {target_resource, target_total}; +} + +std::vector CandidateFasters(int target_total, int target_upper = 69) { + int lower = std::max(2, CeilDiv(target_total, std::max(1, target_upper))); + int upper = std::max(lower, CeilDiv(target_total, 10)); + std::vector candidates; + for (int faster = upper; faster >= lower; --faster) { + candidates.push_back(faster); + } + return candidates; +} + +std::vector +TryFactorWithOptionalBumps(const std::vector> &items, + int factor) { + std::vector bump_indices; + for (const auto &item : items) { + int idx = item.first; + int latency = item.second; + if (latency % factor == 0) { + continue; + } + if ((latency + 1) % factor == 0) { + bump_indices.push_back(idx); + } else { + return {}; + } + } + return bump_indices; +} + +int ScaledTotalForResource(const std::vector &specs, int resource, + int faster, + const std::unordered_set &bump_indices) { + int total = 0; + for (int idx = 0; idx < static_cast(specs.size()); ++idx) { + if (std::find(specs[idx].resources.begin(), specs[idx].resources.end(), + resource) == specs[idx].resources.end()) { + continue; + } + int latency = specs[idx].latency + (bump_indices.count(idx) ? 1 : 0); + total += CeilDiv(std::max(1, latency), faster); + } + return total; +} + +std::vector CandidateGCDs(const std::vector &latencies) { + std::set> gcds; + for (int value : latencies) { + if (value <= 0) { + continue; + } + for (int d = 1; d * d <= value; ++d) { + if (value % d == 0) { + gcds.insert(d); + gcds.insert(value / d); + } + } + } + return std::vector(gcds.begin(), gcds.end()); +} + +std::pair> +TryGCDWithOptionalBumps(const std::vector> &items, + int target_total, int target_upper = 69) { + std::vector latencies; + latencies.reserve(items.size()); + for (const auto &item : items) { + latencies.push_back(item.second); + } + int target_gcd_floor = std::max(1, CeilDiv(target_total, target_upper)); + int base_g = 0; + for (int latency : latencies) { + base_g = base_g == 0 ? std::abs(latency) : GcdInt(base_g, latency); + } + if (base_g >= target_gcd_floor) { + return {base_g, {}}; + } + + for (int g : CandidateGCDs(latencies)) { + if (g < target_gcd_floor) { + continue; + } + std::vector bump_indices; + bool ok = true; + for (const auto &item : items) { + int idx = item.first; + int latency = item.second; + if (latency % g == 0) { + continue; + } + if ((latency + 1) % g == 0) { + bump_indices.push_back(idx); + } else { + ok = false; + break; + } + } + if (ok) { + return {g, bump_indices}; + } + } + return {std::max(1, base_g), {}}; +} + +std::pair> +AutoSelectSunmmioILPFaster(const std::vector &specs, + int target_upper = 69) { + // HiGHS works on integral time slots, but raw cost-model latencies can make + // the modulo model unnecessarily large. Scale time using the most heavily + // occupied resource because its total latency gives the dominant lower bound + // on II. A latency may be increased by one only when that makes every + // command on the bottleneck exactly divisible by the scale factor; this is a + // conservative timing quantization, never an optimistic shortening. + // Formally, L_r = sum_{i uses r} d_i selects r* = argmax_r L_r. For a + // candidate scale f and bump b_i in {0, 1}, the model uses + // d'_i = ceil((d_i + b_i) / f), with b_i = 1 only if f divides d_i + 1. + auto usage = BuildResourceUsage(specs); + auto totals = BuildResourceTotals(usage); + auto [target_resource, target_total] = FindTargetResource(totals); + if (target_resource < 0 || target_total <= 0) { + return {1, {}}; + } + + for (int faster : CandidateFasters(target_total, target_upper)) { + std::vector bump_vec = + TryFactorWithOptionalBumps(usage[target_resource], faster); + if (bump_vec.empty() && !usage[target_resource].empty()) { + bool all_divisible = true; + for (const auto &item : usage[target_resource]) { + if (item.second % faster != 0) { + all_divisible = false; + break; + } + } + if (!all_divisible) { + continue; + } + } + std::unordered_set bump_indices(bump_vec.begin(), bump_vec.end()); + int target_scaled_total = + ScaledTotalForResource(specs, target_resource, faster, bump_indices); + // Rounding each latency independently can change which resource is the + // bottleneck. Reject such a factor: the search bound and model size must + // remain anchored to the resource selected from the unscaled problem. + bool violations = false; + for (const auto &kv : totals) { + int scaled_total = + ScaledTotalForResource(specs, kv.first, faster, bump_indices); + if (scaled_total > target_scaled_total) { + violations = true; + break; + } + } + if (!violations) { + return {std::max(1, faster), bump_vec}; + } + } + + // Exact factorization is not always possible. A common divisor preserves + // all relative integral durations; optional +1 bumps rescue near-divisible + // cost estimates while keeping the scaled bottleneck below target_upper. + // The fallback chooses g >= ceil(L_r* / target_upper) and replaces each + // latency with (d_i + b_i) / g, so the reduced model remains integral. + auto [gcd_value, bump_vec] = TryGCDWithOptionalBumps( + usage[target_resource], target_total, target_upper); + return {std::max(1, gcd_value), bump_vec}; +} + +int GetEnvInt(const char *name, int default_value) { + const char *raw = std::getenv(name); + if (!raw || !*raw) { + return default_value; + } + return std::atoi(raw); +} + +std::string GetEnvString(const char *name) { + const char *raw = std::getenv(name); + if (!raw || !*raw) { + return ""; + } + return raw; +} + +bool GetEnvBool(const char *name, bool default_value = false) { + const char *raw = std::getenv(name); + if (!raw || !*raw) { + return default_value; + } + std::string value(raw); + for (char &c : value) { + c = static_cast(std::tolower(static_cast(c))); + } + return value == "1" || value == "true" || value == "yes" || value == "on"; +} + +std::string JsonEscape(const std::string &value) { + std::string escaped; + escaped.reserve(value.size()); + for (char c : value) { + switch (c) { + case '\\': + escaped += "\\\\"; + break; + case '"': + escaped += "\\\""; + break; + case '\n': + escaped += "\\n"; + break; + case '\r': + escaped += "\\r"; + break; + case '\t': + escaped += "\\t"; + break; + default: + escaped.push_back(c); + break; + } + } + return escaped; +} + +void WriteJsonString(std::ostream &os, const std::string &value) { + os << "\"" << JsonEscape(value) << "\""; +} + +void WriteProblemJson(const Problem &prob, const std::string &path) { + std::ofstream out(path); + ICHECK(out.is_open()) << "Failed to open ILP problem json path: " << path; + out << std::boolalpha; + out << "{\n"; + out << " \"N\": " << prob.N << ",\n"; + out << " \"Tmax\": " << prob.Tmax << ",\n"; + + out << " \"R\": ["; + for (size_t i = 0; i < prob.R.size(); ++i) { + if (i != 0) { + out << ", "; + } + out << prob.R[i]; + } + out << "],\n"; + + std::map sorted_cap(prob.cap.begin(), prob.cap.end()); + out << " \"cap\": {"; + bool first_cap = true; + for (const auto &kv : sorted_cap) { + if (!first_cap) { + out << ", "; + } + first_cap = false; + WriteJsonString(out, std::to_string(kv.first)); + out << ": " << kv.second; + } + out << "},\n"; + + out << " \"commands\": {\n"; + for (int i = 0; i < prob.N; ++i) { + out << " "; + WriteJsonString(out, std::to_string(i)); + out << ": {\"latency\": " << prob.P[i].latency << ", \"resources\": ["; + for (size_t j = 0; j < prob.P[i].resources.size(); ++j) { + if (j != 0) { + out << ", "; + } + out << prob.P[i].resources[j]; + } + out << "], \"name\": "; + WriteJsonString(out, prob.P[i].name); + out << "}"; + out << (i + 1 == prob.N ? "\n" : ",\n"); + } + out << " },\n"; + + out << " \"P\": {\n"; + for (int i = 0; i < prob.N; ++i) { + out << " "; + WriteJsonString(out, std::to_string(i)); + out << ": {\"latency\": " << prob.P[i].latency << ", \"resources\": ["; + for (size_t j = 0; j < prob.P[i].resources.size(); ++j) { + if (j != 0) { + out << ", "; + } + out << prob.P[i].resources[j]; + } + out << "]}"; + out << (i + 1 == prob.N ? "\n" : ",\n"); + } + out << " },\n"; + + out << " \"dep_edges\": ["; + for (size_t i = 0; i < prob.dep_edges.size(); ++i) { + if (i != 0) { + out << ", "; + } + out << "[" << prob.dep_edges[i].first << ", " << prob.dep_edges[i].second + << "]"; + } + out << "],\n"; + + std::map, int> sorted_delta; + for (const auto &edge : prob.dep_edges) { + auto it = prob.delta.find(EdgeKey(edge.first, edge.second)); + if (it != prob.delta.end()) { + sorted_delta[edge] = it->second; + } + } + out << " \"delta\": {"; + bool first_delta = true; + for (const auto &kv : sorted_delta) { + if (!first_delta) { + out << ", "; + } + first_delta = false; + WriteJsonString(out, std::to_string(kv.first.first) + "," + + std::to_string(kv.first.second)); + out << ": " << kv.second; + } + out << "},\n"; + + std::vector sorted_flows = prob.flows; + std::sort( + sorted_flows.begin(), sorted_flows.end(), + [](const FlowSpec &a, const FlowSpec &b) { + return std::tie(a.resident, a.prod, a.cons, a.mem, a.buffer_name, + a.fixed_bank, a.fp, a.initial_time, a.w_off, a.w_dur, + a.r_off, a.r_dur, a.write_resource, a.read_resource) < + std::tie(b.resident, b.prod, b.cons, b.mem, b.buffer_name, + b.fixed_bank, b.fp, b.initial_time, b.w_off, b.w_dur, + b.r_off, b.r_dur, b.write_resource, b.read_resource); + }); + out << " \"flows\": ["; + for (size_t i = 0; i < sorted_flows.size(); ++i) { + const FlowSpec &flow = sorted_flows[i]; + if (i != 0) { + out << ", "; + } + out << "{"; + out << "\"kind\": "; + WriteJsonString(out, flow.resident ? "resident" : "internal"); + out << ", \"prod\": " << flow.prod; + out << ", \"cons\": " << flow.cons; + out << ", \"delta\": " << flow.delta; + out << ", \"mem\": " << flow.mem; + out << ", \"buffer_name\": "; + WriteJsonString(out, flow.buffer_name); + out << ", \"fixed_bank\": " << flow.fixed_bank; + out << ", \"precolor\": " << flow.precolor; + out << ", \"fp\": " << flow.fp; + out << ", \"initial_time\": " << flow.initial_time; + out << ", \"w_off\": " << flow.w_off; + out << ", \"w_dur\": " << flow.w_dur; + out << ", \"r_off\": " << flow.r_off; + out << ", \"r_dur\": " << flow.r_dur; + out << ", \"write_resource\": " << flow.write_resource; + out << ", \"read_resource\": " << flow.read_resource; + out << "}"; + } + out << "],\n"; + out << " \"versioned_buffers\": ["; + for (size_t i = 0; i < prob.versioned_buffer_names.size(); ++i) { + if (i != 0) { + out << ", "; + } + WriteJsonString(out, prob.versioned_buffer_names[i]); + } + out << "]\n"; + out << "}\n"; +} + +void MaybeExportProblemJson(const Problem &prob, bool debug) { + std::string export_path = GetEnvString("TL_SUNMMIO_ILP_PROBLEM_JSON"); + if (export_path.empty() && debug) { + export_path = "body_ilp_problem.json"; + } + if (!export_path.empty()) { + WriteProblemJson(prob, export_path); + } +} + +std::string AddStageSuffixToPath(const std::string &path, int stage) { + if (path.empty()) { + return path; + } + std::string suffix = "_" + std::to_string(stage); + size_t dot = path.find_last_of('.'); + size_t slash = path.find_last_of("/\\"); + if (dot == std::string::npos || (slash != std::string::npos && dot < slash)) { + return path + suffix; + } + return path.substr(0, dot) + suffix + path.substr(dot); +} + +void MaybeExportProblemJsonForStage(const Problem &prob, bool debug, + int stage) { + std::string export_path = GetEnvString("TL_SUNMMIO_ILP_PROBLEM_JSON"); + if (export_path.empty() && debug) { + export_path = "body_ilp_problem.json"; + } + if (!export_path.empty()) { + WriteProblemJson(prob, AddStageSuffixToPath(export_path, stage)); + } +} + +std::string Name(int b) { return b == 0 ? "ping" : "pong"; } + +std::string MemName(int mem) { return mem == 0 ? "wsram" : "asram"; } + +std::string BankPhaseName(int phase) { + return (phase & 1) == 0 ? "ping" : "pong"; +} + +std::string ResourceName(int r) { + switch (r) { + case static_cast(IlpResourceType::kTensorCore): + return "tensor_core"; + case static_cast(IlpResourceType::kVectorCore): + return "vector_core"; + case static_cast(IlpResourceType::kODMA0): + return "odma0"; + case static_cast(IlpResourceType::kODMA1): + return "odma1"; + case static_cast(IlpResourceType::kWsramIn): + return "wsram.in"; + case static_cast(IlpResourceType::kWsramOut): + return "wsram.out"; + case static_cast(IlpResourceType::kAsramIn): + return "asram.in"; + case static_cast(IlpResourceType::kAsramOut): + return "asram.out"; + // case static_cast(IlpResourceType::kRsram): + // return "rsram"; + default: + return "resource_" + std::to_string(r); + } +} + +enum class ConflictType : int { + kNone = 0, + kNeedDifferent = 1, + kNeedSame = 2, + kImpossible = 3, +}; + +int WrapBit(int start_slot, int slot) { return slot < start_slot ? 0 : 1; } + +std::map BuildSlotRhoMap(int start_slot, int duration, int ii) { + std::map slot_rho; + for (int step = 0; step < duration; ++step) { + int slot = (start_slot + step) % ii; + slot_rho[slot] = WrapBit(start_slot, slot); + } + return slot_rho; +} + +ConflictType AnalyzeWriteReadConflict(int write_start_slot, int write_dur, + int write_parity_flip, + int read_start_slot, int read_dur, + int read_parity_flip, int ii, bool flip) { + std::map write_rho; + for (int step = 0; step < write_dur; ++step) { + int slot = (write_start_slot + step) % ii; + write_rho[slot] = + flip ? (WrapBit(write_start_slot, slot) ^ write_parity_flip) : 0; + } + + bool saw_same = false; + bool saw_diff = false; + for (int step = 0; step < read_dur; ++step) { + int slot = (read_start_slot + step) % ii; + auto it = write_rho.find(slot); + if (it == write_rho.end()) + continue; + int read_rho = + flip ? (WrapBit(read_start_slot, slot) ^ read_parity_flip) : 0; + if (it->second == read_rho) { + saw_same = true; + } else { + saw_diff = true; + } + if (saw_same && saw_diff) + return ConflictType::kImpossible; + } + + if (!saw_same && !saw_diff) + return ConflictType::kNone; + if (saw_same) + return ConflictType::kNeedDifferent; + return ConflictType::kNeedSame; +} + +ConflictType MergeConflictRequirements(ConflictType lhs, ConflictType rhs) { + if (lhs == ConflictType::kImpossible || rhs == ConflictType::kImpossible) { + return ConflictType::kImpossible; + } + if (lhs == ConflictType::kNone) + return rhs; + if (rhs == ConflictType::kNone) + return lhs; + return lhs == rhs ? lhs : ConflictType::kImpossible; +} + +ConflictType AnalyzePrecolorConflict(const FlowSpec &lhs, const FlowSpec &rhs) { + if (lhs.precolor < 0 || rhs.precolor < 0 || lhs.mem != rhs.mem || + lhs.buffer_name != rhs.buffer_name) { + return ConflictType::kNone; + } + return lhs.precolor == rhs.precolor ? ConflictType::kNeedSame + : ConflictType::kNeedDifferent; +} + +ConflictType AnalyzeFlowConflict(const FlowSpec &write_flow, + int write_start_time, + const FlowSpec &read_flow, int read_start_time, + int ii, const BankFlipMode &mode) { + ConflictType precolor = AnalyzePrecolorConflict(write_flow, read_flow); + if (write_flow.write_resource < 0 || read_flow.read_resource < 0 || + write_flow.mem != read_flow.mem) { + return precolor; + } + ConflictType port = AnalyzeWriteReadConflict( + PositiveMod(write_start_time + write_flow.w_off, ii), write_flow.w_dur, + ((write_start_time + write_flow.w_off) / ii) & 1, + PositiveMod(read_start_time + read_flow.r_off, ii), read_flow.r_dur, + ((read_start_time + read_flow.r_off) / ii) & 1, ii, + mode.FlipForMem(write_flow.mem)); + return MergeConflictRequirements(port, precolor); +} + +int ComputeFoldedOccupancy(int start_time, int duration, int II, int slot) { + if (duration <= 0 || II <= 0) { + return 0; + } + int rel = slot - PositiveMod(start_time, II); + rel %= II; + if (rel < 0) { + rel += II; + } + if (rel >= duration) { + return 0; + } + return CeilDiv(duration - rel, II); +} + +SolutionVerifyResult VerifySolution(const Problem &prob, + const SolveResult &sol) { + SolutionVerifyResult result; + if (!sol.ok) { + result.ok = false; + result.errors.push_back("solver returned non-ok solution"); + return result; + } + if (sol.II <= 0) { + result.ok = false; + result.node_time_ok = false; + result.errors.push_back("II must be positive"); + return result; + } + ICHECK_EQ(static_cast(sol.t.size()), prob.N); + ICHECK_EQ(static_cast(sol.m.size()), prob.N); + ICHECK_EQ(static_cast(sol.y.size()), prob.N); + + std::unordered_map internal_pos; + for (int i = 0; i < static_cast(sol.internal_flow_ids.size()); ++i) { + internal_pos[sol.internal_flow_ids[i]] = i; + } + ICHECK_EQ(static_cast(sol.internal_flow_ids.size()), + static_cast(sol.z_bank.size())); + + auto fail = [&](bool *flag, std::string msg) { + if (flag != nullptr) { + *flag = false; + } + result.ok = false; + result.errors.push_back(std::move(msg)); + }; + + for (int i = 0; i < prob.N; ++i) { + if (sol.m[i] < 0 || sol.m[i] >= sol.II) { + fail(&result.node_time_ok, "node " + std::to_string(i) + + " has invalid slot " + + std::to_string(sol.m[i])); + } + if (sol.t[i] < 0) { + fail(&result.node_time_ok, "node " + std::to_string(i) + + " has negative start " + + std::to_string(sol.t[i])); + } + if (sol.y[i] < 0) { + fail(&result.node_time_ok, "node " + std::to_string(i) + + " has negative iteration " + + std::to_string(sol.y[i])); + } + if (sol.t[i] != sol.y[i] * sol.II + sol.m[i]) { + fail(&result.node_time_ok, + "node " + std::to_string(i) + " violates t=y*II+m: t=" + + std::to_string(sol.t[i]) + " y=" + std::to_string(sol.y[i]) + + " m=" + std::to_string(sol.m[i])); + } + if (PositiveMod(sol.t[i], sol.II) != sol.m[i]) { + fail(&result.node_time_ok, + "node " + std::to_string(i) + " has inconsistent folded slot"); + } + if (sol.t[i] + prob.P[i].latency > sol.makespan) { + fail(&result.node_time_ok, + "node " + std::to_string(i) + " finishes after makespan"); + } + } + + for (const auto &e : prob.dep_edges) { + int src = e.first; + int dst = e.second; + int delta = prob.delta.at(EdgeKey(src, dst)); + int lhs = sol.t[dst] - sol.t[src]; + int rhs = prob.P[src].latency - delta * sol.II; + if (lhs < rhs) { + fail(&result.dependency_ok, + "dependency violated " + std::to_string(src) + "->" + + std::to_string(dst) + ": lhs=" + std::to_string(lhs) + " rhs=" + + std::to_string(rhs) + " delta=" + std::to_string(delta)); + } + } + + for (int i = 0; i < prob.N; ++i) { + for (int r : prob.P[i].resources) { + for (int s = 0; s < sol.II; ++s) { + result.resource_slot_load[r][s] += + ComputeFoldedOccupancy(sol.t[i], prob.P[i].latency, sol.II, s); + } + } + } + + for (int r : prob.R) { + int cap = prob.cap.count(r) ? prob.cap.at(r) : 1; + for (int s = 0; s < sol.II; ++s) { + int use = result.resource_slot_load[r].count(s) + ? result.resource_slot_load[r][s] + : 0; + if (use > cap) { + fail(&result.resource_slot_ok, + "resource slot overflow " + ResourceName(r) + + " slot=" + std::to_string(s) + " use=" + std::to_string(use) + + " cap=" + std::to_string(cap)); + } + } + } + + for (int a = 0; a < static_cast(prob.flows.size()); ++a) { + const FlowSpec &write_flow = prob.flows[a]; + if (write_flow.write_resource < 0 || write_flow.prod < 0) + continue; + auto it_a = internal_pos.find(a); + if (it_a == internal_pos.end()) + continue; + int z_write = sol.z_bank[it_a->second]; + int write_start_time = sol.t[write_flow.prod]; + + for (int b = 0; b < static_cast(prob.flows.size()); ++b) { + if (a == b) + continue; + const FlowSpec &read_flow = prob.flows[b]; + auto it_b = internal_pos.find(b); + if (it_b == internal_pos.end() || read_flow.cons < 0) + continue; + int z_read = sol.z_bank[it_b->second]; + int read_start_time = sol.t[read_flow.cons] + read_flow.delta * sol.II; + ConflictType conflict = + AnalyzeFlowConflict(write_flow, write_start_time, read_flow, + read_start_time, sol.II, sol.bank_flip_mode); + if (conflict == ConflictType::kNeedDifferent && z_write == z_read) { + fail(&result.bank_port_ok, + "bank port conflict requires different banks between flow " + + std::to_string(a) + " and flow " + std::to_string(b)); + } else if (conflict == ConflictType::kNeedSame && z_write != z_read) { + fail(&result.bank_port_ok, + "bank port conflict requires same banks between flow " + + std::to_string(a) + " and flow " + std::to_string(b)); + } else if (conflict == ConflictType::kImpossible) { + fail(&result.bank_port_ok, + "bank port conflict impossible between flow " + std::to_string(a) + + " and flow " + std::to_string(b)); + } + } + } + + return result; +} + +void WriteSolutionJson( + const std::string &path, const Problem &prob, const SolveResult &sol, + const SolutionVerifyResult &verify, + const std::map &runtime_bank_start_phases, + const std::map &runtime_bank_read_delta_parities, + const std::map> + &runtime_bank_reader_phases) { + std::ofstream out(path); + ICHECK(out.is_open()) << "Failed to open ILP solution json path: " << path; + out << std::boolalpha; + + std::unordered_map internal_pos; + for (int i = 0; i < static_cast(sol.internal_flow_ids.size()); ++i) { + internal_pos[sol.internal_flow_ids[i]] = i; + } + + out << "{\n"; + out << " \"ii\": " << sol.II << ",\n"; + out << " \"makespan\": " << sol.makespan << ",\n"; + out << " \"wsram_flip\": " << sol.bank_flip_mode.wsram_flip << ",\n"; + out << " \"asram_flip\": " << sol.bank_flip_mode.asram_flip << ",\n"; + out << " \"vc_blocking_issue_modeled\": " << sol.vc_blocking_issue_modeled + << ",\n"; + out << " \"vc_blocking_issue_constraints\": " + << sol.vc_blocking_issue_constraints << ",\n"; + out << " \"bank_slot_period\": " + << (sol.bank_slot_period > 0 ? sol.bank_slot_period : (2 * sol.II)) + << ",\n"; + out << " \"nodes\": {\n"; + for (int i = 0; i < prob.N; ++i) { + out << " "; + WriteJsonString(out, std::to_string(i)); + out << ": {\"start\": " << sol.t[i] << ", \"slot\": " << sol.m[i] + << ", \"iteration\": " << sol.y[i] << ", \"phases\": ["; + for (size_t p = 0; p < prob.P[i].resources.size(); ++p) { + if (p != 0) { + out << ", "; + } + int resource = prob.P[i].resources[p]; + out << "{"; + out << "\"phase_id\": " << p; + out << ", \"resource_name\": "; + WriteJsonString(out, ResourceName(resource)); + out << ", \"start\": " << sol.t[i]; + out << ", \"end\": " << (sol.t[i] + prob.P[i].latency); + out << ", \"duration\": " << prob.P[i].latency; + out << "}"; + } + out << "]}"; + out << (i + 1 == prob.N ? "\n" : ",\n"); + } + out << " },\n"; + + out << " \"flows\": ["; + bool first_flow = true; + std::unordered_set emitted_resident_buffers; + for (int v = 0; v < static_cast(prob.flows.size()); ++v) { + const auto &flow = prob.flows[v]; + if (flow.resident) { + if (!emitted_resident_buffers.insert(flow.buffer_name).second) { + continue; + } + auto it_start_bank = runtime_bank_start_phases.find(flow.buffer_name); + if (it_start_bank == runtime_bank_start_phases.end()) { + continue; + } + int bank = it_start_bank->second; + int cons_start = sol.t[flow.cons] + flow.delta * sol.II + flow.r_off; + int cons_end = cons_start + flow.r_dur; + int release_time = cons_end; + if (!first_flow) { + out << ", "; + } + first_flow = false; + out << "{"; + out << "\"idx\": " << v; + out << ", \"kind\": "; + WriteJsonString(out, "resident"); + out << ", \"prod\": -1"; + out << ", \"prod_label\": "; + WriteJsonString(out, "resident"); + out << ", \"cons\": " << flow.cons; + out << ", \"delta\": " << flow.delta; + out << ", \"buffer_name\": "; + WriteJsonString(out, flow.buffer_name); + out << ", \"write_resource\": -1"; + out << ", \"read_resource\": " << flow.read_resource; + out << ", \"write_resource_name\": "; + WriteJsonString(out, ""); + out << ", \"read_resource_name\": "; + WriteJsonString( + out, flow.read_resource < 0 ? "" : ResourceName(flow.read_resource)); + out << ", \"memory\": " << flow.mem; + out << ", \"memory_name\": "; + WriteJsonString(out, MemName(flow.mem)); + out << ", \"bank\": "; + WriteJsonString(out, Name(bank)); + out << ", \"start_bank\": "; + WriteJsonString(out, Name(bank)); + out << ", \"write_time\": 0"; + out << ", \"write_end\": 0"; + out << ", \"read_time\": " << cons_start; + out << ", \"read_end\": " << cons_end; + out << ", \"release_time\": " << release_time; + out << ", \"write_resource\": -1"; + out << ", \"read_resource\": " << flow.read_resource; + out << "}"; + continue; + } + + int bank = sol.z_bank[internal_pos.at(v)]; + int cons_start = sol.t[flow.cons] + flow.delta * sol.II + flow.r_off; + int prod_start = sol.t[flow.prod] + flow.w_off; + int prod_end = prod_start + flow.w_dur; + int cons_end = cons_start + flow.r_dur; + int release_time = std::max(prod_end, cons_end); + if (!first_flow) { + out << ", "; + } + first_flow = false; + out << "{"; + out << "\"idx\": " << v; + out << ", \"kind\": "; + WriteJsonString(out, flow.resident ? "resident" : "internal"); + out << ", \"prod\": " << flow.prod; + out << ", \"prod_label\": "; + WriteJsonString(out, + flow.prod >= 0 ? std::to_string(flow.prod) : "resident"); + out << ", \"cons\": " << flow.cons; + out << ", \"delta\": " << flow.delta; + out << ", \"buffer_name\": "; + WriteJsonString(out, flow.buffer_name); + out << ", \"write_resource\": " << flow.write_resource; + out << ", \"read_resource\": " << flow.read_resource; + out << ", \"write_resource_name\": "; + WriteJsonString( + out, flow.write_resource < 0 ? "" : ResourceName(flow.write_resource)); + out << ", \"read_resource_name\": "; + WriteJsonString( + out, flow.read_resource < 0 ? "" : ResourceName(flow.read_resource)); + out << ", \"memory\": " << flow.mem; + out << ", \"memory_name\": "; + WriteJsonString(out, MemName(flow.mem)); + out << ", \"bank\": "; + WriteJsonString(out, Name(bank)); + out << ", \"start_bank\": "; + WriteJsonString(out, Name(bank)); + out << ", \"write_time\": " << prod_start; + out << ", \"write_end\": " << prod_end; + out << ", \"read_time\": " << cons_start; + out << ", \"read_end\": " << cons_end; + out << ", \"release_time\": " << release_time; + out << ", \"write_resource\": " << flow.write_resource; + out << ", \"read_resource\": " << flow.read_resource; + out << "}"; + } + + out << "],\n"; + + out << " \"verify\": {"; + out << "\"ok\": " << verify.ok; + out << ", \"node_time_ok\": " << verify.node_time_ok; + out << ", \"dependency_ok\": " << verify.dependency_ok; + out << ", \"resource_slot_ok\": " << verify.resource_slot_ok; + out << ", \"bank_slot_ok\": " << verify.bank_slot_ok; + out << ", \"bank_port_ok\": " << verify.bank_port_ok; + out << ", \"errors\": ["; + for (size_t i = 0; i < verify.errors.size(); ++i) { + if (i != 0) { + out << ", "; + } + WriteJsonString(out, verify.errors[i]); + } + out << "]},\n"; + out << " \"bank_load\": {}\n"; + out << "}\n"; +} + +void DedupAccesses(std::vector *accesses) { + std::vector deduped; + deduped.reserve(accesses->size()); + for (const AccessInfo &access : *accesses) { + bool exists = false; + for (const AccessInfo &old : deduped) { + if (access.is_write != old.is_write || + access.iter_offset != old.iter_offset || + !access.region->buffer.same_as(old.region->buffer)) { + continue; + } + if (StructuralEqual()(access.region, old.region)) { + exists = true; + break; + } + } + if (!exists) { + deduped.push_back(access); + } + } + *accesses = std::move(deduped); +} + +} // namespace + +class SunmmioStmtAccessAnalyzer : public StmtExprVisitor { +public: + explicit SunmmioStmtAccessAnalyzer(const PrimFunc &f) { + for (const auto &kv : f->buffer_map) { + buffer_data_to_buffer_.Set(kv.second->data, kv.second); + } + } + + std::vector Collect(const Stmt &stmt, + const Var &pipeline_loop_var = Var()) { + accesses_.clear(); + pipeline_loop_var_ = pipeline_loop_var; + VisitStmt(stmt); + DedupAccesses(&accesses_); + return accesses_; + } + +private: + void AddAccess(const BufferRegion ®ion, bool is_write) { + accesses_.push_back( + AccessInfo{region, is_write, + DetectPipelineIterOffsetFromRegion( + region, pipeline_loop_var_, &analyzer_)}); + } + + void VisitStmt_(const BufferStoreNode *op) final { + Array region; + for (const PrimExpr &index : op->indices) { + region.push_back(Range::FromMinExtent(index, 1)); + } + AddAccess(BufferRegion(op->buffer, region), true); + VisitExpr(op->value); + } + + void VisitStmt_(const EvaluateNode *op) final { + if (const auto *call = op->value.as()) { + if (call->op.same_as(dma_copy())) { + AddAccess(NormalizeToBufferRegion(call->args[0]), false); + AddAccess(NormalizeToBufferRegion(call->args[1]), true); + return; + } + if (call->op.same_as(sunmmio_layout_transform())) { + AddAccess(NormalizeToBufferRegion(call->args[0]), false); + AddAccess(NormalizeToBufferRegion(call->args[1]), true); + return; + } + if (call->op.same_as(mma_sunmmio())) { + AddAccess(NormalizeToBufferRegion(call->args[0]), false); + AddAccess(NormalizeToBufferRegion(call->args[1]), false); + BufferRegion c_region = NormalizeToBufferRegion(call->args[2]); + AddAccess(c_region, false); + AddAccess(c_region, true); + return; + } + if (call->op.same_as(Op::Get("tl.broadcast_"))) { + AddAccess(NormalizeToBufferRegion(call->args[0]), false); + AddAccess(NormalizeToBufferRegion(call->args[1]), true); + return; + } + if (call->op.same_as(Op::Get("tl.vector_core_in_tile_reduce"))) { + ICHECK_GE(call->args.size(), 3U); + AddAccess(NormalizeToBufferRegion(call->args[1]), true); + AddAccess(NormalizeToBufferRegion(call->args[2]), false); + return; + } + } + VisitExpr(op->value); + } + + void VisitExpr_(const BufferLoadNode *op) final { + Array region; + for (const PrimExpr &index : op->indices) { + if (const auto *ramp = index.as()) { + region.push_back(Range::FromMinExtent(ramp->base, ramp->lanes)); + } else { + region.push_back(Range::FromMinExtent(index, 1)); + } + } + AddAccess(BufferRegion(op->buffer, region), false); + } + + void VisitExpr_(const CallNode *op) final { + if (op->op.same_as(RegionOp::Get())) { + AddAccess(NormalizeToBufferRegion(ffi::GetRef(op)), false); + return; + } + + if (op->op.same_as(builtin::address_of())) { + if (const auto *load = op->args[0].as()) { + AddAccess(BufferRegion::FullRegion(load->buffer), false); + return; + } + if (const auto *var_node = op->args[0].as()) { + Var data_var = ffi::GetRef(var_node); + auto it = buffer_data_to_buffer_.find(data_var); + if (it != buffer_data_to_buffer_.end()) { + AddAccess(BufferRegion::FullRegion((*it).second), false); + return; + } + } + } + + if (op->op.same_as(builtin::tvm_access_ptr())) { + if (const auto *buffer_var = op->args[1].as()) { + auto it = buffer_data_to_buffer_.find(ffi::GetRef(buffer_var)); + if (it != buffer_data_to_buffer_.end()) { + AddAccess(BufferRegion::FullRegion((*it).second), false); + return; + } + } + } + + StmtExprVisitor::VisitExpr_(op); + } + + arith::Analyzer analyzer_; + ffi::Map buffer_data_to_buffer_; + Var pipeline_loop_var_; + std::vector accesses_; +}; + +class SunmmioRoleMarker : public StmtVisitor { +public: + SunmmioRoleMarker(ASTTraverser &traverser, const PrimFunc &func) + : traverser_(traverser), access_analyzer_(func) { + traverser_.clear(); + } + + Role GetRole(const StmtNode *stmt) const { + auto it = map_.find(stmt); + ICHECK(it != map_.end()) + << "Cannot find role for stmt: " << stmt->GetTypeKey(); + return it->second; + } + + Role GetRole(const Stmt &stmt) const { return GetRole(stmt.get()); } + + std::vector GetAccesses(const Stmt &stmt, + const Var &pipeline_loop_var = Var()) { + return access_analyzer_.Collect(stmt, pipeline_loop_var); + } + + void VisitStmt_(const EvaluateNode *op) final { + Role role = Role::kConsumer; + if (const auto *call = op->value.as()) { + if (call->op.same_as(Op::Get("tl.dma_copy"))) { + BufferRegion src_region = NormalizeToBufferRegion(call->args[0]); + if (IsGlobalBuffer(src_region->buffer)) { + role = Role::kProducer; + } + } + } + SetRole(op, role); + } + + void VisitStmt_(const BufferStoreNode *op) final { + Role role = Role::kProducer; + // Reuse the legacy traverser path for role classification. It is less + // detailed than the ILP access collector but has proven stable on large + // kernels such as flash-attention. + traverser_.traverse_stmt(ffi::GetRef(op)); + auto reads = traverser_.read_buffer_regions_; + for (const BufferRegion &read : reads) { + if (!IsGlobalBuffer(read->buffer)) { + role = Role::kConsumer; + break; + } + } + SetRole(op, role); + } + + void VisitStmt_(const SeqStmtNode *op) final { + StmtVisitor::VisitStmt_(op); + auto role = GetRole(op->seq[0]); + for (const Stmt &stmt : op->seq) { + if (role != GetRole(stmt)) { + role = Role::kBoth; + break; + } + } + SetRole(op, role); + } + + void VisitStmt_(const IfThenElseNode *op) final { + StmtVisitor::VisitStmt_(op); + auto role = GetRole(op->then_case); + if (op->else_case.defined() && role != GetRole(op->else_case.value())) { + role = Role::kBoth; + } + SetRole(op, role); + } + + void VisitStmt_(const BlockRealizeNode *op) final { + StmtVisitor::VisitStmt_(op); + SetRole(op, GetRole(op->block)); + } + + template void HandleBodyStmt(const NodeType *op) { + StmtVisitor::VisitStmt_(op); + SetRole(op, GetRole(op->body)); + } + + void VisitStmt_(const ForNode *op) final { HandleBodyStmt(op); } + void VisitStmt_(const LetStmtNode *op) final { HandleBodyStmt(op); } + void VisitStmt_(const AttrStmtNode *op) final { HandleBodyStmt(op); } + void VisitStmt_(const AssertStmtNode *op) final { HandleBodyStmt(op); } + void VisitStmt_(const BlockNode *op) final { HandleBodyStmt(op); } + void VisitStmt_(const AllocateNode *op) final { HandleBodyStmt(op); } + void VisitStmt_(const DeclBufferNode *op) final { HandleBodyStmt(op); } + +private: + void SetRole(const StmtNode *stmt, Role role) { map_[stmt] = role; } + + std::unordered_map map_; + ASTTraverser traverser_; + SunmmioStmtAccessAnalyzer access_analyzer_; +}; + +class SunmmioExprAnalyzer : public StmtExprVisitor { +public: + SunmmioExprAnalyzer() {} + + void Analyze(const PrimExpr &expr) { + loop_cost_ = 0; + load_times = 0; + flops_ = 0; + args_.clear(); + constants_.clear(); + vars_.clear(); + StmtExprVisitor::VisitExpr(expr); + } + +private: + void VisitExpr_(const MulNode *op) final { + auto a = op->a; + auto b = op->b; + flops_ += 1; + if (const auto *a_int = a.as()) { + if (const auto *b_int = b.as()) { + return; + } + if (a_int->value <= 32) { + loop_cost_ += 2; + StmtExprVisitor::VisitExpr(op->b); + return; + } + } + if (const auto *b_int = b.as()) { + if (b_int->value <= 32) { + loop_cost_ += 2; + StmtExprVisitor::VisitExpr(op->a); + return; + } + } + loop_cost_ += 4; + StmtExprVisitor::VisitExpr(op->a); + StmtExprVisitor::VisitExpr(op->b); + } + + void VisitExpr_(const SubNode *op) final { + loop_cost_ += 4; + flops_ += 1; + StmtExprVisitor::VisitExpr(op->a); + StmtExprVisitor::VisitExpr(op->b); + } + + void VisitExpr_(const AddNode *op) final { + loop_cost_ += 4; + flops_ += 1; + StmtExprVisitor::VisitExpr(op->a); + StmtExprVisitor::VisitExpr(op->b); + } + + void VisitExpr_(const MaxNode *op) final { + loop_cost_ += 3; + flops_ += 1; + StmtExprVisitor::VisitExpr(op->a); + StmtExprVisitor::VisitExpr(op->b); + } + + void VisitExpr_(const MinNode *op) final { + loop_cost_ += 3; + flops_ += 1; + StmtExprVisitor::VisitExpr(op->a); + StmtExprVisitor::VisitExpr(op->b); + } + + void VisitExpr_(const CastNode *op) final { + loop_cost_ += 3; + StmtExprVisitor::VisitExpr(op->value); + } + + void VisitExpr_(const IntImmNode *op) final { + bool insert = true; + for (auto it : constants_) { + if (ExprDeepEqual()(it, tvm::ffi::GetRef(op))) { + insert = false; + break; + } + } + if (insert) { + constants_.push_back(tvm::ffi::GetRef(op)); + } + } + + void VisitExpr_(const FloatImmNode *op) final { + bool insert = true; + for (auto it : constants_) { + if (ExprDeepEqual()(it, tvm::ffi::GetRef(op))) { + insert = false; + break; + } + } + if (insert) { + constants_.push_back(tvm::ffi::GetRef(op)); + } + } + + void VisitExpr_(const VarNode *op) final { + bool insert = true; + for (auto it : vars_) { + if (ExprDeepEqual()(it, tvm::ffi::GetRef(op))) { + insert = false; + break; + } + } + if (insert) { + vars_.push_back(tvm::ffi::GetRef(op)); + } + } + + void VisitExpr_(const CallNode *op) final { + if (op->op.same_as(Op::Get("tir.exp2"))) { + loop_cost_ += 10; + flops_ += 3; + StmtExprVisitor::VisitExpr(op->args[0]); + } else if (op->op.same_as(Op::Get("tl.infinity"))) { + bool insert = true; + for (auto it : constants_) { + if (ExprDeepEqual()(it, + FloatImm(DataType::Float(16), + std::numeric_limits::infinity()))) { + insert = false; + break; + } + } + if (insert) { + constants_.push_back(FloatImm(DataType::Float(16), + std::numeric_limits::infinity())); + } + } else if (op->op.same_as(Op::Get("tir.if_then_else"))) { + bool insert = true; + for (auto it : args_) { + if (ExprDeepEqual()(it, op->args[0])) { + insert = false; + break; + } + } + if (insert) { + args_.push_back(op->args[0]); + } + StmtExprVisitor::VisitExpr(op->args[0]); + StmtExprVisitor::VisitExpr(op->args[1]); + StmtExprVisitor::VisitExpr(op->args[2]); + } else if (op->op.same_as(Op::Get("tir.bitwise_and"))) { + bool insert = true; + for (auto it : args_) { + if (ExprDeepEqual()(it, op->args[0])) { + insert = false; + break; + } + } + if (insert) { + args_.push_back(op->args[0]); + } + flops_ += 1; + StmtExprVisitor::VisitExpr(op->args[0]); + StmtExprVisitor::VisitExpr(op->args[1]); + } else { + ICHECK(0) << "Op " << op->op << " not supported now."; + } + } + + void VisitExpr_(const LENode *op) final { + loop_cost_ += 3; + flops_ += 1; + StmtExprVisitor::VisitExpr(op->a); + StmtExprVisitor::VisitExpr(op->b); + } + + void VisitExpr_(const BufferLoadNode *op) final { + if (load_times == 0) { + load_times++; + loop_cost_ += 14; + flops_ += 5; + } else { + load_times++; + loop_cost_ += 1; + flops_ += 1; + } + for (auto arg : op->indices) { + bool insert = true; + for (auto it : args_) { + if (ExprDeepEqual()(it, arg)) { + insert = false; + break; + } + } + if (insert) { + args_.push_back(arg); + } + } + } + +public: + float loop_cost_ = 0; + Array args_; + Array vars_; + Array constants_; + int load_times = 0; + float flops_ = 0; +}; + +class TemplateCommand { +public: + int id{-1}; + std::string name; + Stmt stmt; + Role role{Role::kUndefined}; + DeviceType type{DeviceType::Unspecified}; + std::vector accesses; + CommandSpec spec; + + TemplateCommand(int id, const Stmt &stmt) + : id(id), name("cmd_" + std::to_string(id)), stmt(stmt) {} +}; + +const CallNode *GetSingleBroadcastCall(const Stmt &stmt) { + const CallNode *broadcast = nullptr; + PostOrderVisit(stmt, [&](const ObjectRef &obj) { + const auto *call = obj.as(); + if (call && call->op.same_as(Op::Get("tl.broadcast_"))) { + ICHECK(broadcast == nullptr) + << "A pipeline statement may contain at most one broadcast leaf"; + broadcast = call; + } + }); + return broadcast; +} + +bool IsAllGatherBroadcast(const TemplateCommand &cmd) { + const CallNode *call = GetSingleBroadcastCall(cmd.stmt); + if (!call) { + return false; + } + ICHECK(call->args.size() == static_cast(kBroadcastArgCount) || + call->args.size() == static_cast(kBroadcastArgCount + 1)) + << "tl.broadcast_ expects its fixed arguments and optional src_core"; + // Before sync-token injection, a broadcast with only the fixed arguments is + // issued by every core and therefore implements an all-gather collective. + return call->args.size() == static_cast(kBroadcastArgCount); +} + +bool IsCopyStage(const TemplateCommand &cmd) { + bool has_shared_write = false; + bool has_global_read = false; + for (const AccessInfo &access : cmd.accesses) { + if (access.is_write && IsSunmmioSharedBuffer(access.buffer())) { + has_shared_write = true; + } + if (!access.is_write && IsGlobalBuffer(access.buffer())) { + has_global_read = true; + } + } + return has_shared_write && has_global_read; +} + +bool IsProducerLike(const TemplateCommand &cmd) { + return cmd.role == Role::kProducer || + (cmd.role == Role::kBoth && IsCopyStage(cmd)); +} + +bool IsConsumerLike(const TemplateCommand &cmd) { + return cmd.role == Role::kConsumer || + (cmd.role == Role::kBoth && !IsCopyStage(cmd)); +} + +std::string SummarizeStmtForName(const Stmt &stmt) { + auto buffer_scope = [](const Buffer &buffer) { + return buffer.scope().empty() ? std::string("default") + : std::string(buffer.scope()); + }; + auto buffer_label = [&](const Buffer &buffer) { + return buffer->name + "@" + buffer_scope(buffer); + }; + auto expr_kind = [&](const PrimExpr &expr) -> std::string { + if (expr.as()) { + if (const auto *imm = expr.as()) { + return imm->value == 0 ? "const0" : "const"; + } + } + if (expr.as()) { + if (const auto *imm = expr.as()) { + return imm->value == 0.0 ? "const0" : "const"; + } + } + if (const auto *load = expr.as()) { + return "copy(" + load->buffer->name + ")"; + } + if (expr.as()) { + return "max"; + } + if (expr.as()) { + return "min"; + } + if (expr.as()) { + return "add"; + } + if (expr.as()) { + return "sub"; + } + if (expr.as()) { + return "mul"; + } + if (expr.as()) { + return "div"; + } + if (expr.as()) { + return "cast"; + } + if (const auto *call = expr.as()) { + if (call->op.same_as(Op::Get("tir.exp2"))) { + return "exp2"; + } + if (call->op.same_as(Op::Get("tir.if_then_else"))) { + return "if_then_else"; + } + if (call->op.same_as(Op::Get("tir.bitwise_and"))) { + return "bitwise_and"; + } + } + return expr->GetTypeKey(); + }; + if (const auto *eval = stmt.as()) { + if (const auto *call = eval->value.as()) { + if (call->op.same_as(Op::Get("tl.dma_copy"))) { + BufferRegion src = NormalizeToBufferRegion(call->args[0]); + BufferRegion dst = NormalizeToBufferRegion(call->args[1]); + return "dma_copy(" + buffer_label(src->buffer) + "->" + + buffer_label(dst->buffer) + ")"; + } + } + return "evaluate"; + } + if (const auto *block = stmt.as()) { + if (const auto *eval = block->block->body.as()) { + if (const auto *call = eval->value.as()) { + if (call->op.same_as(Op::Get("tl.mma_sunmmio"))) { + auto A = call->args[0].as(); + auto B = call->args[1].as(); + if (A && B && A->args.size() >= 4 && B->args.size() >= 4) { + return "mma_sunmmio(" + + std::to_string(A->args[2].as()->value) + "x" + + std::to_string(B->args[3].as()->value) + "x" + + std::to_string(A->args[3].as()->value) + ")"; + } + return "mma_sunmmio"; + } + } + } + if (block->block->name_hint == "reduce_tile_op") { + std::string reduce_kind = "reduce"; + std::string reduce_dst; + bool found_reduce = false; + PostOrderVisit(block->block->body, [&](const ObjectRef &obj) { + if (found_reduce) { + return; + } + if (const auto *eval = obj.as()) { + if (const auto *call = eval->value.as()) { + if (call->op.same_as(Op::Get("tl.vector_core_in_tile_reduce")) && + !call->args.empty()) { + if (const auto *kind = call->args[0].as()) { + reduce_kind = kind->value; + } else { + reduce_kind = "reduce"; + } + if (call->args.size() >= 2) { + BufferRegion dst = NormalizeToBufferRegion(call->args[1]); + reduce_dst = buffer_label(dst->buffer); + } + found_reduce = true; + } + } + } + }); + return reduce_dst.empty() + ? "reduce_tile_op(" + reduce_kind + ")" + : "reduce_tile_op(" + reduce_kind + " -> " + reduce_dst + ")"; + } + if (!block->block->name_hint.empty()) { + return "block:" + block->block->name_hint; + } + return "block"; + } + if (const auto *loop = stmt.as()) { + std::string summary = "for"; + PostOrderVisit(loop->body, [&](const ObjectRef &obj) { + if (summary != "for") { + return; + } + if (const auto *store = obj.as()) { + summary = "for store(" + buffer_label(store->buffer) + + " := " + expr_kind(store->value) + ")"; + return; + } + if (const auto *eval = obj.as()) { + if (const auto *call = eval->value.as()) { + if (call->op.same_as(Op::Get("tl.vector_core_in_tile_reduce")) && + !call->args.empty()) { + if (const auto *kind = call->args[0].as()) { + summary = std::string("for reduce(") + kind->value + ")"; + } else { + summary = "for reduce"; + } + return; + } + } + } + }); + return summary; + } + return stmt->GetTypeKey(); +} + +int GetPingPongMemoryKind(const Buffer &buffer) { + if (buffer.scope() == "shared.wsram") { + return 0; + } + if (buffer.scope() == "shared.asram") { + return 1; + } + return -1; +} + +int GetMemoryWriteResource(int mem) { + if (mem == 0) { + return static_cast(IlpResourceType::kWsramIn); + } + if (mem == 1) { + return static_cast(IlpResourceType::kAsramIn); + } + return -1; +} + +int GetMemoryReadResource(int mem) { + if (mem == 0) { + return static_cast(IlpResourceType::kWsramOut); + } + if (mem == 1) { + return static_cast(IlpResourceType::kAsramOut); + } + return -1; +} + +bool CommandUsesResource(const CommandSpec &spec, int resource) { + return std::find(spec.resources.begin(), spec.resources.end(), resource) != + spec.resources.end(); +} + +FlowSpec MakeInternalFlowSpec(const Problem &problem, int prod, int cons, + int delta, int mem, + const std::string &buffer_name) { + FlowSpec flow; + flow.resident = false; + flow.prod = prod; + flow.cons = cons; + flow.delta = delta; + flow.mem = mem; + flow.buffer_name = buffer_name; + flow.fixed_bank = -1; + // Keep the initial ILP input coarse-grained: each SRAM flow currently counts + // as one bank-capacity unit until a more precise footprint model is wired in. + flow.fp = 1; + flow.initial_time = 0; + flow.w_off = 0; + flow.w_dur = problem.P[prod].latency; + flow.r_off = 0; + flow.r_dur = problem.P[cons].latency; + flow.write_resource = GetMemoryWriteResource(mem); + flow.read_resource = GetMemoryReadResource(mem); + return flow; +} + +FlowSpec MakeResidentFlowSpec(const Problem &problem, int cons, int mem, + const std::string &buffer_name) { + FlowSpec flow; + flow.resident = true; + flow.prod = -1; + flow.cons = cons; + flow.delta = 0; + flow.mem = mem; + flow.buffer_name = buffer_name; + flow.fixed_bank = -1; + flow.fp = 1; + flow.initial_time = 0; + flow.w_off = 0; + flow.w_dur = 0; + flow.r_off = 0; + flow.r_dur = problem.P[cons].latency; + flow.write_resource = -1; + flow.read_resource = GetMemoryReadResource(mem); + return flow; +} + +std::string ExtractBufferNameFromCommandLabel(const std::string &name) { + size_t arrow = name.find("->"); + if (arrow == std::string::npos) { + return ""; + } + size_t at = name.find('@', arrow + 2); + if (at == std::string::npos) { + return ""; + } + return name.substr(arrow + 2, at - (arrow + 2)); +} + +int PositiveMod(int value, int mod) { + if (mod <= 0) { + return value; + } + int result = value % mod; + if (result < 0) { + result += mod; + } + return result; +} + +int RuntimeVersionCount(const Buffer &buffer, int iterations) { + if (iterations <= 0 || GetPingPongMemoryKind(buffer) >= 0) { + return 1; + } + return iterations; +} + +int RuntimeBankedVersionCount(const Buffer &buffer, int iterations) { + if (iterations <= 2 || GetPingPongMemoryKind(buffer) < 0) { + return 1; + } + return CeilDiv(iterations, 2); +} + +struct ScheduledAccessWindow { + int logical_iter{0}; + int command_iter{0}; + int cmd_id{-1}; + int start{0}; + int end{0}; + bool is_write{false}; + int physical_bank{-1}; + BufferRegion region; +}; + +struct BufferInstanceLifetime { + int logical_iter{0}; + int command_iter{0}; + int first_write_cmd_id{-1}; + int write_bank{-1}; + int start{0}; + int end{0}; + std::vector access_regions; + std::vector write_regions; +}; + +std::vector DetectRuntimeMultiversionBuffers( + const std::vector &commands, + const std::vector &versioned_buffers, + const std::vector &runtime_banked_buffers, + const Var &pipeline_loop_var, const SolveResult &sol, int iterations, + bool enable_lifetime_pruning, + const std::map &runtime_bank_start_phases, + const std::map &runtime_bank_read_delta_parities, + const std::map> &runtime_bank_writer_phases, + const std::map> &runtime_bank_reader_phases, + const std::map &runtime_bank_flip_modes) { + // Runtime versions follow the selected logical iteration count. In + // particular, stage shrinking changes this value independently of the + // schedule span ceil(makespan / II). + iterations = std::max(iterations, 0); + std::unordered_set candidates; + std::unordered_set banked_candidates; + for (const Buffer &buffer : runtime_banked_buffers) { + banked_candidates.insert(buffer.get()); + } + for (const Buffer &buffer : versioned_buffers) { + bool is_banked = banked_candidates.count(buffer.get()) != 0; + int version_count = is_banked + ? RuntimeBankedVersionCount(buffer, iterations) + : RuntimeVersionCount(buffer, iterations); + if (version_count > 1) { + candidates.insert(buffer.get()); + } + } + if (candidates.empty()) { + return {}; + } + if (!enable_lifetime_pruning) { + std::vector result; + for (const Buffer &buffer : versioned_buffers) { + if (candidates.count(buffer.get())) { + result.push_back(buffer); + } + } + return result; + } + + auto resolve_access_bank = [&](const Buffer &buffer, int cmd_id, + int command_iter, bool is_write) { + if (!banked_candidates.count(buffer.get())) { + return -1; + } + int phase = 0; + bool flip = true; + auto it_flip = runtime_bank_flip_modes.find(buffer->name); + if (it_flip != runtime_bank_flip_modes.end()) { + flip = it_flip->second != 0; + } + int iter_phase = flip ? command_iter : 0; + const auto &phase_maps = + is_write ? runtime_bank_writer_phases : runtime_bank_reader_phases; + auto it_buffer = phase_maps.find(buffer->name); + if (it_buffer != phase_maps.end()) { + auto it_phase = it_buffer->second.find(cmd_id); + if (it_phase != it_buffer->second.end()) { + phase = it_phase->second; + return PositiveMod(iter_phase + phase, 2); + } + } + auto it_start = runtime_bank_start_phases.find(buffer->name); + if (it_start != runtime_bank_start_phases.end()) { + phase = it_start->second; + } + if (!is_write) { + auto it_delta = runtime_bank_read_delta_parities.find(buffer->name); + if (it_delta != runtime_bank_read_delta_parities.end()) { + phase += it_delta->second; + } + } + return PositiveMod(iter_phase + phase, 2); + }; + + int max_iter_offset = 0; + for (const TemplateCommand &cmd : commands) { + for (const AccessInfo &access : cmd.accesses) { + max_iter_offset = std::max(max_iter_offset, access.iter_offset); + } + } + + const int expanded_iters = std::max(2, iterations + max_iter_offset + 1); + std::unordered_map> + windows_by_buffer; + windows_by_buffer.reserve(candidates.size()); + + struct ExpandedCommand { + int iter{0}; + const TemplateCommand *cmd{nullptr}; + }; + + std::vector producer_ids; + std::vector body_ids; + producer_ids.reserve(commands.size()); + body_ids.reserve(commands.size()); + for (const TemplateCommand &cmd : commands) { + if (IsProducerLike(cmd)) { + producer_ids.push_back(cmd.id); + } else { + body_ids.push_back(cmd.id); + } + } + + std::vector expanded_commands; + expanded_commands.reserve(expanded_iters * commands.size()); + for (int iter = 0; iter < expanded_iters; ++iter) { + for (int id : body_ids) { + expanded_commands.push_back(ExpandedCommand{iter, &commands[id]}); + } + for (int id : producer_ids) { + expanded_commands.push_back(ExpandedCommand{iter + 1, &commands[id]}); + } + } + std::sort(expanded_commands.begin(), expanded_commands.end(), + [](const ExpandedCommand &a, const ExpandedCommand &b) { + if (a.iter != b.iter) { + return a.iter < b.iter; + } + return a.cmd->id < b.cmd->id; + }); + + for (const ExpandedCommand &expanded : expanded_commands) { + const int start = sol.t[expanded.cmd->id] + expanded.iter * sol.II; + const int end = start + expanded.cmd->spec.latency; + for (const AccessInfo &access : expanded.cmd->accesses) { + const BufferNode *buf = access.buffer().get(); + if (!candidates.count(buf)) { + continue; + } + windows_by_buffer[buf].push_back(ScheduledAccessWindow{ + expanded.iter + access.iter_offset, expanded.iter, expanded.cmd->id, + start, end, access.is_write, + resolve_access_bank(access.buffer(), expanded.cmd->id, expanded.iter, + access.is_write), + MaterializeBufferRegion(access.region, pipeline_loop_var, + expanded.iter)}); + } + } + + auto region_sets_intersect = [](const std::vector &lhs, + const std::vector &rhs) { + for (const BufferRegion &lhs_region : lhs) { + for (const BufferRegion &rhs_region : rhs) { + if (PipelineRegionIntersect(lhs_region->region, rhs_region->region)) { + return true; + } + } + } + return false; + }; + + std::vector runtime_multiversion_buffers; + for (const Buffer &buffer : versioned_buffers) { + if (!candidates.count(buffer.get())) { + continue; + } + + auto it_windows = windows_by_buffer.find(buffer.get()); + if (it_windows == windows_by_buffer.end()) { + continue; + } + + bool is_banked = banked_candidates.count(buffer.get()) != 0; + std::map, std::vector> + windows_by_instance; + for (const ScheduledAccessWindow &window : it_windows->second) { + int bank = is_banked ? window.physical_bank : -1; + windows_by_instance[{window.logical_iter, bank}].push_back(&window); + } + + std::vector lifetimes; + lifetimes.reserve(windows_by_instance.size()); + for (const auto &kv : windows_by_instance) { + int first_write_start = std::numeric_limits::max(); + const ScheduledAccessWindow *first_write_window = nullptr; + for (const ScheduledAccessWindow *window : kv.second) { + if (window->is_write) { + if (window->start < first_write_start) { + first_write_start = window->start; + first_write_window = window; + } + } + } + if (first_write_start == std::numeric_limits::max()) { + continue; + } + + BufferInstanceLifetime lifetime; + lifetime.logical_iter = kv.first.first; + lifetime.command_iter = first_write_window->command_iter; + lifetime.first_write_cmd_id = first_write_window->cmd_id; + lifetime.write_bank = kv.first.second; + lifetime.start = first_write_start; + lifetime.end = first_write_start; + for (const ScheduledAccessWindow *window : kv.second) { + if (window->end < first_write_start) { + continue; + } + lifetime.end = std::max(lifetime.end, window->end); + lifetime.access_regions.push_back(window->region); + if (window->is_write) { + lifetime.write_regions.push_back(window->region); + } + } + if (!lifetime.write_regions.empty()) { + lifetimes.push_back(std::move(lifetime)); + } + } + + std::sort( + lifetimes.begin(), lifetimes.end(), + [](const BufferInstanceLifetime &a, const BufferInstanceLifetime &b) { + if (a.start != b.start) { + return a.start < b.start; + } + return a.logical_iter < b.logical_iter; + }); + + bool needs_runtime_multiversion = false; + for (size_t i = 0; i < lifetimes.size() && !needs_runtime_multiversion; + ++i) { + for (size_t j = i + 1; j < lifetimes.size(); ++j) { + if (lifetimes[j].start >= lifetimes[i].end) { + break; + } + if (is_banked && lifetimes[i].write_bank != lifetimes[j].write_bank) { + continue; + } + if (region_sets_intersect(lifetimes[i].write_regions, + lifetimes[j].access_regions) || + region_sets_intersect(lifetimes[j].write_regions, + lifetimes[i].access_regions)) { + needs_runtime_multiversion = true; + break; + } + } + } + + if (needs_runtime_multiversion) { + runtime_multiversion_buffers.push_back(buffer); + } + } + + return runtime_multiversion_buffers; +} + +TimeWindowOrderResult +BuildTimeWindowOrders(const std::vector &commands, + int iterations, const SolveResult &sol) { + TimeWindowOrderResult result; + const int stage_count = CeilDiv(sol.makespan, std::max(1, sol.II)); + const int prologue_end = std::max(0, stage_count - 1) * sol.II; + const int body_begin = prologue_end; + const int body_end = prologue_end + sol.II; + const int epilogue_begin = body_end; + const int epilogue_end = prologue_end + sol.makespan; + + int max_iter = stage_count; + std::vector expanded; + expanded.reserve(max_iter * commands.size()); + for (int iter = 0; iter < max_iter; ++iter) { + for (const TemplateCommand &cmd : commands) { + int id = cmd.id; + expanded.push_back( + ExpandedOrderEntry{iter, id, sol.t[id] + iter * sol.II}); + } + } + + std::sort(expanded.begin(), expanded.end(), + [](const ExpandedOrderEntry &a, const ExpandedOrderEntry &b) { + if (a.absolute_start != b.absolute_start) { + return a.absolute_start < b.absolute_start; + } + if (a.iter != b.iter) { + return a.iter < b.iter; + } + return a.id < b.id; + }); + + for (const ExpandedOrderEntry &entry : expanded) { + if (entry.absolute_start < body_begin) { + result.prologue.push_back(entry); + continue; + } + if (entry.absolute_start < body_end) { + result.body.push_back(entry); + result.steady_state_max_iter_offset = + std::max(result.steady_state_max_iter_offset, entry.iter); + continue; + } + if (entry.absolute_start < epilogue_end) { + result.epilogue.push_back(entry); + } + } + return result; +} + +BufferRegion MaterializeBufferRegion(const BufferRegion ®ion, + const Var &loop_var, int iter) { + if (!loop_var.defined()) { + return region; + } + ffi::Map vmap; + vmap.Set(loop_var, make_const(loop_var.dtype(), iter)); + Array materialized; + for (const Range &rng : region->region) { + PrimExpr min = tir::Substitute(rng->min, vmap); + PrimExpr extent = tir::Substitute(rng->extent, vmap); + materialized.push_back(Range::FromMinExtent(min, extent)); + } + return BufferRegion(region->buffer, materialized); +} + +std::vector +DetectVersionedBuffers(const std::vector &commands) { + std::set used_buffers; + std::unordered_set consumer_used; + std::unordered_set producer_used; + std::unordered_set self_dependent_buffers; + std::unordered_map first_write_index; + std::unordered_map> write_indexes; + std::unordered_map first_read_index; + std::unordered_map last_read_index; + std::vector versioned_buffers; + auto mark_versioned = [&](const Buffer &buffer) { + if (std::find(versioned_buffers.begin(), versioned_buffers.end(), buffer) == + versioned_buffers.end()) { + versioned_buffers.push_back(buffer); + } + }; + + for (int i = 0; i < static_cast(commands.size()); ++i) { + bool is_producer = IsProducerLike(commands[i]); + bool is_consumer = IsConsumerLike(commands[i]); + std::unordered_set reads_in_cmd; + std::unordered_set writes_in_cmd; + for (const AccessInfo &access : commands[i].accesses) { + if (IsGlobalBuffer(access.buffer())) { + continue; + } + used_buffers.insert(access.buffer()); + const BufferNode *buf = access.buffer().get(); + if (access.is_write) { + writes_in_cmd.insert(buf); + if (is_producer) { + producer_used.insert(buf); + } + if (!first_write_index.count(buf)) { + first_write_index[buf] = i; + } + write_indexes[buf].push_back(i); + } else { + reads_in_cmd.insert(buf); + if (is_consumer) { + consumer_used.insert(buf); + } + if (!first_read_index.count(buf)) { + first_read_index[buf] = i; + } + last_read_index[buf] = i; + } + } + for (const BufferNode *buf : writes_in_cmd) { + if (reads_in_cmd.count(buf)) { + self_dependent_buffers.insert(buf); + } + } + } + + for (const Buffer &buffer : used_buffers) { + const BufferNode *buf = buffer.get(); + if (self_dependent_buffers.count(buf)) { + continue; + } + auto it_w = first_write_index.find(buf); + auto it_r = first_read_index.find(buf); + if (it_w != first_write_index.end() && it_r != first_read_index.end() && + it_w->second < it_r->second) { + mark_versioned(buffer); + continue; + } + if (consumer_used.count(buf) && producer_used.count(buf)) { + auto r = first_read_index.find(buf); + auto w = first_write_index.find(buf); + if (r != first_read_index.end() && w != first_write_index.end() && + r->second > w->second) { + mark_versioned(buffer); + continue; + } + } + auto it_last_r = last_read_index.find(buf); + if (it_w != first_write_index.end() && it_last_r != last_read_index.end() && + it_w->second < it_last_r->second && + IsCopyStage(commands[it_w->second])) { + mark_versioned(buffer); + } + } + + bool updated = true; + while (updated) { + updated = false; + for (const Buffer &buffer : used_buffers) { + if (std::find(versioned_buffers.begin(), versioned_buffers.end(), + buffer) != versioned_buffers.end()) { + continue; + } + const BufferNode *buf = buffer.get(); + if (self_dependent_buffers.count(buf)) { + continue; + } + auto it_writes = write_indexes.find(buf); + auto it_first_w = first_write_index.find(buf); + auto it_first_r = first_read_index.find(buf); + if (it_writes == write_indexes.end() || it_writes->second.empty() || + it_first_w == first_write_index.end() || + it_first_r == first_read_index.end()) { + continue; + } + bool can_propagate = it_first_w->second < it_first_r->second; + for (int idx : it_writes->second) { + for (const AccessInfo &access : commands[idx].accesses) { + if (access.is_write || IsGlobalBuffer(access.buffer())) { + continue; + } + if (first_write_index.find(access.buffer().get()) == + first_write_index.end()) { + continue; + } + if (std::find(versioned_buffers.begin(), versioned_buffers.end(), + access.buffer()) == versioned_buffers.end()) { + can_propagate = false; + break; + } + } + if (!can_propagate) { + break; + } + } + if (can_propagate) { + mark_versioned(buffer); + updated = true; + } + } + } + return versioned_buffers; +} + +void BuildTemplateDependencyGraph(const std::vector &commands, + int iter_mod, + const std::vector &versioned_buffers, + const Var &pipeline_loop_var, + const BankFlipMode &mode, Problem *problem) { + std::unordered_set versioned; + std::unordered_set banked_versioned; + std::unordered_set bank_rotating_versioned; + for (const Buffer &buffer : versioned_buffers) { + versioned.insert(buffer.get()); + int mem = GetPingPongMemoryKind(buffer); + if (mem >= 0) { + banked_versioned.insert(buffer.get()); + } + if (mem >= 0 && mode.FlipForMem(mem)) { + bank_rotating_versioned.insert(buffer.get()); + } + } + + problem->flows.clear(); + problem->dep_edges.clear(); + problem->delta.clear(); + + struct ExpandedCommand { + int template_id{-1}; + int iter{-1}; + const TemplateCommand *cmd{nullptr}; + }; + enum class AccessType : uint8_t { kRead, kWrite }; + struct AccessRecord { + BufferRegion region; + int expanded_idx{-1}; + int access_idx{-1}; + AccessType type{AccessType::kRead}; + }; + + std::map, int> best_delta; + std::map, int> + flow_key_to_index; + using ConsumerAccessKey = std::tuple; + std::set satisfied_consumer_access; + + // Pre-color distinct producer values only when one banked buffer has more + // than one writer operation. A color is a phase offset: physical_bank = + // (logical_iteration + phase) % 2. Readers inherit the phase of the write + // that produces their value. + std::unordered_map> writer_phases; + std::map, int> access_phases; + for (const Buffer &buffer : versioned_buffers) { + const BufferNode *buf = buffer.get(); + if (!banked_versioned.count(buf)) { + continue; + } + std::vector writer_ids; + for (const TemplateCommand &cmd : commands) { + bool writes_buffer = false; + for (const AccessInfo &access : cmd.accesses) { + writes_buffer = + writes_buffer || (access.is_write && access.buffer().get() == buf); + } + if (writes_buffer) { + writer_ids.push_back(cmd.id); + } + } + if (writer_ids.size() <= 1) { + continue; + } + for (size_t i = 0; i < writer_ids.size(); ++i) { + writer_phases[buf][writer_ids[i]] = static_cast(i & 1); + } + } + + for (const TemplateCommand &cmd : commands) { + for (int access_idx = 0; access_idx < static_cast(cmd.accesses.size()); + ++access_idx) { + const AccessInfo &access = cmd.accesses[access_idx]; + const BufferNode *buf = access.buffer().get(); + auto phase_group = writer_phases.find(buf); + if (phase_group == writer_phases.end()) { + continue; + } + if (access.is_write) { + access_phases[{cmd.id, access_idx}] = phase_group->second.at(cmd.id); + continue; + } + + int producer_phase = -1; + int producer_iter_delta = 0; + for (int producer_id = cmd.id - 1; producer_id >= 0 && producer_phase < 0; + --producer_id) { + const TemplateCommand &producer = commands[producer_id]; + for (int producer_access_idx = + static_cast(producer.accesses.size()) - 1; + producer_access_idx >= 0; --producer_access_idx) { + const AccessInfo &producer_access = + producer.accesses[producer_access_idx]; + if (!producer_access.is_write || + producer_access.buffer().get() != buf) { + continue; + } + if (!PipelineRegionIntersect( + MaterializeBufferRegion(producer_access.region, + pipeline_loop_var, 0) + ->region, + MaterializeBufferRegion(access.region, pipeline_loop_var, 0) + ->region)) { + continue; + } + producer_phase = phase_group->second.at(producer_id); + break; + } + } + // A read before its producer in template order consumes the previous + // logical iteration of the last matching writer. + for (int producer_id = static_cast(commands.size()) - 1; + producer_id >= cmd.id && producer_phase < 0; --producer_id) { + const TemplateCommand &producer = commands[producer_id]; + for (int producer_access_idx = + static_cast(producer.accesses.size()) - 1; + producer_access_idx >= 0; --producer_access_idx) { + const AccessInfo &producer_access = + producer.accesses[producer_access_idx]; + if (!producer_access.is_write || + producer_access.buffer().get() != buf) { + continue; + } + if (!PipelineRegionIntersect( + MaterializeBufferRegion(producer_access.region, + pipeline_loop_var, -1) + ->region, + MaterializeBufferRegion(access.region, pipeline_loop_var, 0) + ->region)) { + continue; + } + producer_phase = phase_group->second.at(producer_id); + producer_iter_delta = 1; + break; + } + } + if (producer_phase >= 0) { + access_phases[{cmd.id, access_idx}] = + (producer_phase + + (bank_rotating_versioned.count(buf) ? producer_iter_delta : 0)) & + 1; + } + } + } + + auto version_mod = [&](const BufferNode *buf) { + if (iter_mod <= 0) { + return 0; + } + if (bank_rotating_versioned.count(buf)) { + // Banked buffers always rotate ping/pong every iteration. When + // num_stages > 2 we additionally attach a runtime multiversion axis to + // each ping/pong bank, so the full physical alias period becomes + // 2 * ceil(num_stages / 2). For num_stages <= 2 this collapses to the + // original ping/pong-only period 2. + return 2 * std::max(1, CeilDiv(iter_mod, 2)); + } + return iter_mod; + }; + + auto maybe_record_edge = [&](int src_id, int dst_id, int delta) { + if (delta < 0) { + return; + } + std::pair key{src_id, dst_id}; + auto it = best_delta.find(key); + if (it == best_delta.end() || delta < it->second) { + best_delta[key] = delta; + } + }; + + auto maybe_record_flow = [&](int src_id, int dst_id, int delta, + const BufferNode *buf, int src_access_idx, + int dst_access_idx, int mem) { + if (mem < 0 || src_id == dst_id) { + return false; + } + int write_resource = GetMemoryWriteResource(mem); + int read_resource = GetMemoryReadResource(mem); + if (!CommandUsesResource(problem->P[src_id], write_resource) || + !CommandUsesResource(problem->P[dst_id], read_resource)) { + return false; + } + auto flow_key = + std::make_tuple(src_id, dst_id, buf, src_access_idx, dst_access_idx); + auto it = flow_key_to_index.find(flow_key); + if (it == flow_key_to_index.end()) { + int flow_index = static_cast(problem->flows.size()); + flow_key_to_index[flow_key] = flow_index; + FlowSpec flow = + MakeInternalFlowSpec(*problem, src_id, dst_id, delta, mem, buf->name); + auto phase = access_phases.find({src_id, src_access_idx}); + if (phase != access_phases.end()) { + flow.precolor = phase->second; + } + problem->flows.push_back(std::move(flow)); + } else { + problem->flows[it->second].delta = + std::min(problem->flows[it->second].delta, delta); + } + return true; + }; + + std::vector producer_ids; + std::vector body_ids; + producer_ids.reserve(commands.size()); + body_ids.reserve(commands.size()); + for (const TemplateCommand &cmd : commands) { + if (IsProducerLike(cmd)) { + producer_ids.push_back(cmd.id); + } else { + body_ids.push_back(cmd.id); + } + } + + int steady_state_iters = std::max(1, iter_mod); + for (const Buffer &buffer : versioned_buffers) { + steady_state_iters = + std::max(steady_state_iters, std::max(1, version_mod(buffer.get()))); + } + int expanded_iters = std::max(2, steady_state_iters + 1); + std::vector expanded_commands; + expanded_commands.reserve(expanded_iters * commands.size()); + for (int iter = 0; iter < expanded_iters; ++iter) { + for (int id : body_ids) { + expanded_commands.push_back(ExpandedCommand{id, iter, &commands[id]}); + } + for (int id : producer_ids) { + expanded_commands.push_back(ExpandedCommand{id, iter + 1, &commands[id]}); + } + } + std::sort(expanded_commands.begin(), expanded_commands.end(), + [](const ExpandedCommand &a, const ExpandedCommand &b) { + if (a.iter != b.iter) { + return a.iter < b.iter; + } + return a.template_id < b.template_id; + }); + + auto access_version = [&](const BufferNode *buf, + const ExpandedCommand &command, + const AccessInfo &access, int access_idx) { + int mod = version_mod(buf); + if (mod <= 0) { + return command.iter + access.iter_offset; + } + int logical_iter = command.iter + access.iter_offset; + auto phase = access_phases.find({command.template_id, access_idx}); + if (!banked_versioned.count(buf)) { + return PositiveMod(logical_iter, mod); + } + int bank_phase = phase == access_phases.end() ? 0 : phase->second; + bool flip = bank_rotating_versioned.count(buf) != 0; + int versions_per_bank = CeilDiv(mod, 2); + int bank = flip ? PositiveMod(logical_iter + bank_phase, 2) + : PositiveMod(bank_phase, 2); + int version_in_bank = + flip ? PositiveMod(FloorDiv(logical_iter, 2), versions_per_bank) + : PositiveMod(logical_iter, versions_per_bank); + return version_in_bank * 2 + bank; + }; + + auto materialize_access = [&](const ExpandedCommand &command, + const AccessInfo &access) { + return MaterializeBufferRegion(access.region, pipeline_loop_var, + command.iter); + }; + + std::unordered_map> + buffer_access_history; + buffer_access_history.reserve(versioned.size() + commands.size()); + + struct ResidentCandidate { + int cmd_id{-1}; + int mem{-1}; + const BufferNode *buffer{nullptr}; + int access_idx{-1}; + std::string buffer_name; + }; + std::vector resident_candidates; + + for (int curr_idx = 0; curr_idx < static_cast(expanded_commands.size()); + ++curr_idx) { + const ExpandedCommand &curr_cmd = expanded_commands[curr_idx]; + + for (int dst_access_idx = 0; + dst_access_idx < static_cast(curr_cmd.cmd->accesses.size()); + ++dst_access_idx) { + const AccessInfo &dst_access = curr_cmd.cmd->accesses[dst_access_idx]; + if (dst_access.is_write) { + continue; + } + const BufferNode *buf = dst_access.buffer().get(); + auto hist_it = buffer_access_history.find(buf); + if (hist_it == buffer_access_history.end()) { + continue; + } + BufferRegion dst_region = materialize_access(curr_cmd, dst_access); + for (auto it = hist_it->second.rbegin(); it != hist_it->second.rend(); + ++it) { + const ExpandedCommand &src_cmd = expanded_commands[it->expanded_idx]; + const AccessInfo &src_access = src_cmd.cmd->accesses[it->access_idx]; + if (versioned.count(buf) && + access_version(buf, src_cmd, src_access, it->access_idx) != + access_version(buf, curr_cmd, dst_access, dst_access_idx)) { + continue; + } + if (it->type != AccessType::kWrite || + !PipelineRegionIntersect(dst_region->region, it->region->region)) { + continue; + } + maybe_record_edge(src_cmd.template_id, curr_cmd.template_id, + curr_cmd.iter - src_cmd.iter); + int mem = GetPingPongMemoryKind(dst_access.buffer()); + bool has_concrete_flow = + maybe_record_flow(src_cmd.template_id, curr_cmd.template_id, + curr_cmd.iter - src_cmd.iter, buf, it->access_idx, + dst_access_idx, mem); + if (has_concrete_flow) { + satisfied_consumer_access.insert( + std::make_tuple(curr_cmd.template_id, mem, buf, dst_access_idx)); + } + break; + } + } + + for (int dst_access_idx = 0; + dst_access_idx < static_cast(curr_cmd.cmd->accesses.size()); + ++dst_access_idx) { + const AccessInfo &dst_access = curr_cmd.cmd->accesses[dst_access_idx]; + if (!dst_access.is_write) { + continue; + } + const BufferNode *buf = dst_access.buffer().get(); + auto hist_it = buffer_access_history.find(buf); + if (hist_it == buffer_access_history.end()) { + continue; + } + BufferRegion dst_region = materialize_access(curr_cmd, dst_access); + for (auto it = hist_it->second.rbegin(); it != hist_it->second.rend(); + ++it) { + const ExpandedCommand &src_cmd = expanded_commands[it->expanded_idx]; + const AccessInfo &src_access = src_cmd.cmd->accesses[it->access_idx]; + if (versioned.count(buf) && + access_version(buf, src_cmd, src_access, it->access_idx) != + access_version(buf, curr_cmd, dst_access, dst_access_idx)) { + continue; + } + if (!PipelineRegionIntersect(dst_region->region, it->region->region)) { + continue; + } + maybe_record_edge(src_cmd.template_id, curr_cmd.template_id, + curr_cmd.iter - src_cmd.iter); + if (it->type == AccessType::kWrite) { + break; + } + } + } + + for (int access_idx = 0; + access_idx < static_cast(curr_cmd.cmd->accesses.size()); + ++access_idx) { + const AccessInfo &access = curr_cmd.cmd->accesses[access_idx]; + buffer_access_history[access.buffer().get()].push_back(AccessRecord{ + materialize_access(curr_cmd, access), curr_idx, access_idx, + access.is_write ? AccessType::kWrite : AccessType::kRead}); + } + } + + // Every core executes an all-gather broadcast and enters its participant + // barrier before issuing the mcast. Row and column collectives can use + // different ODMA engines and buffers, so ordinary resource and data hazards + // do not preserve a common barrier encounter order across cores. Chain the + // collectives in template order, including the loop-carried last-to-first + // edge, while keeping their commands and bank precolor constraints separate. + std::vector all_gather_ids; + for (const TemplateCommand &cmd : commands) { + if (IsAllGatherBroadcast(cmd)) { + all_gather_ids.push_back(cmd.id); + } + } + for (size_t i = 1; i < all_gather_ids.size(); ++i) { + maybe_record_edge(all_gather_ids[i - 1], all_gather_ids[i], 0); + } + if (all_gather_ids.size() > 1) { + maybe_record_edge(all_gather_ids.back(), all_gather_ids.front(), 1); + } + + for (const auto &kv : best_delta) { + problem->dep_edges.push_back(kv.first); + problem->delta[EdgeKey(kv.first.first, kv.first.second)] = kv.second; + } + + for (const TemplateCommand &cmd : commands) { + for (int access_idx = 0; access_idx < static_cast(cmd.accesses.size()); + ++access_idx) { + const AccessInfo &access = cmd.accesses[access_idx]; + if (access.is_write) + continue; + int mem = GetPingPongMemoryKind(access.buffer()); + if (mem < 0 || !CommandUsesResource(problem->P[cmd.id], + GetMemoryReadResource(mem))) { + continue; + } + resident_candidates.push_back( + ResidentCandidate{cmd.id, mem, access.buffer().get(), access_idx, + access.buffer()->name}); + } + } + + std::set emitted_resident_access; + for (const ResidentCandidate &candidate : resident_candidates) { + ConsumerAccessKey key = + std::make_tuple(candidate.cmd_id, candidate.mem, candidate.buffer, + candidate.access_idx); + if (satisfied_consumer_access.count(key) != 0) { + continue; + } + if (!emitted_resident_access.insert(key).second) { + continue; + } + problem->flows.push_back(MakeResidentFlowSpec( + *problem, candidate.cmd_id, candidate.mem, candidate.buffer_name)); + } +} + +bool ValidateProblemGraph(const std::vector &commands, + const Problem &problem) { + if (problem.N != static_cast(commands.size()) || + static_cast(problem.P.size()) != problem.N) { + return false; + } + for (int id = 0; id < problem.N; ++id) { + if (commands[id].id != id) { + return false; + } + for (const AccessInfo &access : commands[id].accesses) { + if (!access.region.defined() || !access.buffer().defined()) { + return false; + } + } + } + for (const auto &edge : problem.dep_edges) { + if (edge.first < 0 || edge.first >= problem.N || edge.second < 0 || + edge.second >= problem.N) { + return false; + } + auto delta_it = problem.delta.find(EdgeKey(edge.first, edge.second)); + if (delta_it == problem.delta.end() || delta_it->second < 0) { + return false; + } + } + for (const FlowSpec &flow : problem.flows) { + if (flow.cons < 0 || flow.cons >= problem.N || + (!flow.resident && (flow.prod < 0 || flow.prod >= problem.N)) || + flow.delta < 0) { + return false; + } + } + return true; +} + +int ResourceLowerBound(const Problem &prob) { + int lb = 1; + for (int r : prob.R) { + int cap = 1; + auto it = prob.cap.find(r); + if (it != prob.cap.end()) { + cap = it->second; + } + if (cap <= 0) { + continue; + } + long long total = 0; + for (int i = 0; i < prob.N; ++i) { + if (std::find(prob.P[i].resources.begin(), prob.P[i].resources.end(), + r) != prob.P[i].resources.end()) { + total += prob.P[i].latency; + } + } + lb = std::max(lb, std::max(1, int((total + cap - 1) / cap))); + } + return lb; +} + +HighsInt AddCol(Highs &highs, double lower, double upper, double cost, + bool is_integer) { + HighsStatus st = highs.addCol(cost, lower, upper, 0, nullptr, nullptr); + ICHECK(st == HighsStatus::kOk) << "addCol failed"; + HighsInt col = highs.getNumCol() - 1; + if (is_integer) { + highs.changeColIntegrality(col, HighsVarType::kInteger); + } + return col; +} + +HighsStatus AddRow(Highs &highs, double lower, double upper, + const std::vector &idx, + const std::vector &val) { + const HighsInt *idx_ptr = idx.empty() ? nullptr : idx.data(); + const double *val_ptr = val.empty() ? nullptr : val.data(); + return highs.addRow(lower, upper, HighsInt(idx.size()), idx_ptr, val_ptr); +} + +void MergeLinearTerms(const std::vector &idx, + const std::vector &val, + std::vector &merged_idx, + std::vector &merged_val) { + std::map acc; + for (size_t k = 0; k < idx.size(); ++k) { + acc[idx[k]] += val[k]; + } + merged_idx.clear(); + merged_val.clear(); + for (const auto &kv : acc) { + if (kv.second == 0.0) { + continue; + } + merged_idx.push_back(kv.first); + merged_val.push_back(kv.second); + } +} + +void AddLeq(Highs &highs, const std::vector &idx, + const std::vector &val, double rhs) { + std::vector merged_idx; + std::vector merged_val; + MergeLinearTerms(idx, val, merged_idx, merged_val); + if (merged_idx.empty()) { + if (0.0 > rhs) { + AddRow(highs, 1.0, 0.0, {}, {}); + } + return; + } + AddRow(highs, -kInf, rhs, merged_idx, merged_val); +} + +void AddEq(Highs &highs, const std::vector &idx, + const std::vector &val, double rhs) { + std::vector merged_idx; + std::vector merged_val; + MergeLinearTerms(idx, val, merged_idx, merged_val); + AddRow(highs, rhs, rhs, merged_idx, merged_val); +} + +void AddConditionalParity(Highs &highs, HighsInt z_write, HighsInt z_read, + HighsInt write_parity, HighsInt read_parity, + HighsInt x_prod, const std::vector &x_cons, + int required_xor) { + HighsInt quotient = AddCol(highs, 0, 2, 0, true); + std::vector idx{z_write, z_read, write_parity, + read_parity, quotient, x_prod}; + std::vector upper{1.0, 1.0, 1.0, 1.0, -2.0, 4.0}; + std::vector lower{-1.0, -1.0, -1.0, -1.0, 2.0, 4.0}; + for (HighsInt x : x_cons) { + idx.push_back(x); + upper.push_back(4.0); + lower.push_back(4.0); + } + AddLeq(highs, idx, upper, required_xor + 8.0); + AddLeq(highs, idx, lower, -required_xor + 8.0); +} + +void AddConditionalBankRelation(Highs &highs, HighsInt z_write, HighsInt z_read, + HighsInt x_prod, + const std::vector &x_cons, + int required_xor) { + ICHECK(required_xor == 0 || required_xor == 1); + std::vector idx{z_write, z_read, x_prod}; + std::vector first; + std::vector second; + double first_rhs = 0.0; + double second_rhs = 0.0; + if (required_xor == 0) { + first = {1.0, -1.0, 1.0}; + second = {-1.0, 1.0, 1.0}; + first_rhs = second_rhs = 2.0; + } else { + first = {1.0, 1.0, 1.0}; + second = {-1.0, -1.0, 1.0}; + first_rhs = 3.0; + second_rhs = 1.0; + } + for (HighsInt x : x_cons) { + idx.push_back(x); + first.push_back(1.0); + second.push_back(1.0); + } + AddLeq(highs, idx, first, first_rhs); + AddLeq(highs, idx, second, second_rhs); +} + +ModelVars BuildModel(Highs &highs, const Problem &prob, int II, bool optimize_t, + int threads, const BankFlipMode &mode) { + // Variable meanings for command i and modulo slot s: + // t[i] = absolute start time in the representative schedule window + // m[i] = t[i] mod II, selected by one-hot x[i][s] + // y[i] = floor(t[i] / II) = 2*y_half[i] + start_parity[i] + // a[i][s] = number of folded copies of command i occupying slot s + // z[v] = ping/pong bank phase assigned to data flow v + // T = makespan. Feasibility searches leave T unpriced; the final solve + // minimizes it for the already-minimal II. + // Domains: t_i,y_i,yh_i,m_i,a_is,T are nonnegative integers; + // x_is,p_i,z_v are binary, and 0 <= m_i < II. + highs.clear(); + const bool enable_solver_log = GetEnvBool("TL_SUNMMIO_ILP_HIGHS_LOG", false); + highs.setOptionValue("output_flag", enable_solver_log); + highs.setOptionValue("log_to_console", enable_solver_log); + if (enable_solver_log) { + highs.setOptionValue("mip_report_level", 2); + } + highs.setOptionValue("threads", threads); + highs.setOptionValue("parallel", "on"); + highs.changeObjectiveSense(ObjSense::kMinimize); + + int max_delta = 0; + int max_latency = 0; + for (const auto &kv : prob.delta) { + max_delta = std::max(max_delta, kv.second); + } + for (const auto &spec : prob.P) { + max_latency = std::max(max_latency, spec.latency); + } + int time_ub = prob.Tmax + max_delta * II + max_latency; + + ModelVars vars; + vars.col_t.resize(prob.N); + vars.col_y.resize(prob.N); + vars.col_y_half.resize(prob.N); + vars.col_start_parity.resize(prob.N); + vars.col_m.resize(prob.N); + vars.col_x.assign(prob.N, std::vector(II, -1)); + vars.col_a.assign(prob.N, std::vector(II, -1)); + + for (int v = 0; v < static_cast(prob.flows.size()); ++v) { + vars.internal_flow_ids.push_back(v); + } + + for (int i = 0; i < prob.N; ++i) { + vars.col_t[i] = AddCol(highs, 0, time_ub, 0, true); + vars.col_y[i] = AddCol(highs, 0, time_ub, 0, true); + vars.col_y_half[i] = AddCol(highs, 0, CeilDiv(time_ub, 2), 0, true); + vars.col_start_parity[i] = AddCol(highs, 0, 1, 0, true); + vars.col_m[i] = AddCol(highs, 0, II - 1, 0, true); + } + for (int i = 0; i < prob.N; ++i) { + for (int s = 0; s < II; ++s) { + vars.col_x[i][s] = AddCol(highs, 0, 1, 0, true); + } + } + for (int i = 0; i < prob.N; ++i) { + int ub = CeilDiv(prob.P[i].latency, II); + for (int s = 0; s < II; ++s) { + vars.col_a[i][s] = AddCol(highs, 0, ub, 0, true); + } + } + vars.col_T = AddCol(highs, 0, time_ub, optimize_t ? 1.0 : 0.0, true); + + // Choose exactly one modulo start slot and link all representations of the + // same start time. t = II*y + m linearizes modulo arithmetic; splitting y + // into 2*y_half + parity exposes whether bank rotation has crossed an odd + // number of initiation intervals without introducing nonlinear arithmetic. + // sum_s x[i,s] = 1 + // m[i] = sum_s s*x[i,s] + // t[i] = II*y[i] + m[i] + // y[i] = 2*y_half[i] + start_parity[i] + for (int i = 0; i < prob.N; ++i) { + std::vector idx1; + std::vector val1; + for (int s = 0; s < II; ++s) { + idx1.push_back(vars.col_x[i][s]); + val1.push_back(1.0); + } + AddEq(highs, idx1, val1, 1.0); + + std::vector idx2{vars.col_m[i]}; + std::vector val2{1.0}; + for (int s = 0; s < II; ++s) { + idx2.push_back(vars.col_x[i][s]); + val2.push_back(-double(s)); + } + AddEq(highs, idx2, val2, 0.0); + + AddEq(highs, {vars.col_t[i], vars.col_y[i], vars.col_m[i]}, + {1.0, -double(II), -1.0}, 0.0); + AddEq(highs, {vars.col_y[i], vars.col_y_half[i], vars.col_start_parity[i]}, + {1.0, -2.0, -1.0}, 0.0); + } + + // Fold a command's [start, start + latency) interval onto the cyclic II-slot + // calendar. a[i][s] can exceed one when latency > II, representing several + // overlapped iterations of the same command in physical slot s. + // a[i,s] = sum_st x[i,st] * max(0, ceil((d_i-rel(s,st))/II)) + // where rel(s,st) = (s-st) mod II. + for (int i = 0; i < prob.N; ++i) { + int dur = prob.P[i].latency; + for (int s = 0; s < II; ++s) { + std::vector idx{vars.col_a[i][s]}; + std::vector val{1.0}; + for (int st = 0; st < II; ++st) { + int rel = (s - st) % II; + if (rel < 0) { + rel += II; + } + int cnt = 0; + if (rel < dur) { + cnt = CeilDiv(dur - rel, II); + } + if (cnt != 0) { + idx.push_back(vars.col_x[i][st]); + val.push_back(-double(cnt)); + } + } + AddEq(highs, idx, val, 0.0); + } + } + + // For dependency i -> j at iteration distance delta, require + // t[i] + latency[i] <= t[j] + delta*II. A self-edge cannot be shifted by + // start times, so it is immediately infeasible when its latency exceeds the + // available delta initiation intervals. + // t[i] + d_i <= t[j] + delta(i,j)*II + for (const auto &e : prob.dep_edges) { + int i = e.first; + int j = e.second; + int d = prob.P[i].latency; + int delta = prob.delta.at(EdgeKey(i, j)); + if (i == j) { + if (d > delta * II) { + AddRow(highs, 1.0, 0.0, {}, {}); + } + continue; + } + AddLeq(highs, {vars.col_t[i], vars.col_t[j]}, {1.0, -1.0}, + double(delta * II - d)); + } + + // T bounds every command completion, making minimization of T equivalent to + // minimizing the representative window's makespan. + // T >= t[i] + d_i, for every command i; objective: min T. + for (int i = 0; i < prob.N; ++i) { + AddLeq(highs, {vars.col_t[i], vars.col_T}, {1.0, -1.0}, + double(-prob.P[i].latency)); + } + + // At every modulo slot, sum the folded occupancies of commands using a + // physical execution resource. This models conflicts across all overlapped + // iterations, not merely commands visible in one template iteration. + // sum_{i uses r} a[i,s] <= cap[r], for every resource r and slot s. + for (int r : prob.R) { + int cap = prob.cap.count(r) ? prob.cap.at(r) : 1; + for (int s = 0; s < II; ++s) { + std::vector idx; + std::vector val; + for (int i = 0; i < prob.N; ++i) { + if (std::find(prob.P[i].resources.begin(), prob.P[i].resources.end(), + r) != prob.P[i].resources.end()) { + idx.push_back(vars.col_a[i][s]); + val.push_back(1.0); + } + } + AddLeq(highs, idx, val, double(cap)); + } + } + + const int internal_count = static_cast(vars.internal_flow_ids.size()); + std::cerr << "[ILP] bank_var_count new=" << internal_count * 2 << "\n"; + + vars.col_z.resize(internal_count); + for (int vv = 0; vv < internal_count; ++vv) { + vars.col_z[vv] = AddCol(highs, 0, 1, 0, true); + } + + { + // Multiple FlowSpecs can describe reads fed by the same physical write. + // They must share one write-bank phase; allowing independent z values + // would assign the same produced data to two banks simultaneously. + // z[v] = z[w] for flows v,w with (prod[v], mem[v]) equal. + std::map> same_write_groups; + for (int vv = 0; vv < internal_count; ++vv) { + int fid = vars.internal_flow_ids[vv]; + const FlowSpec &flow = prob.flows[fid]; + if (flow.write_resource < 0 || flow.prod < 0) + continue; + same_write_groups[MakeSameWriteFlowKey(flow)].push_back(vv); + } + for (const auto &kv : same_write_groups) { + const std::vector &flows = kv.second; + for (size_t i = 1; i < flows.size(); ++i) { + AddEq(highs, {vars.col_z[flows[0]], vars.col_z[flows[i]]}, {1.0, -1.0}, + 0.0); + } + } + } + + // Precoloring captures bank identities already implied by resident buffers + // or frontend metadata. Convert pairwise implications into equality or XOR + // constraints before considering time-dependent write/read overlap. + // NeedSame: z[a] - z[b] = 0; NeedDifferent: z[a] + z[b] = 1. + for (int a = 0; a < internal_count; ++a) { + const FlowSpec &lhs = prob.flows[vars.internal_flow_ids[a]]; + for (int b = a + 1; b < internal_count; ++b) { + const FlowSpec &rhs = prob.flows[vars.internal_flow_ids[b]]; + ConflictType conflict = AnalyzePrecolorConflict(lhs, rhs); + if (conflict == ConflictType::kNeedSame) { + AddEq(highs, {vars.col_z[a], vars.col_z[b]}, {1.0, -1.0}, 0.0); + } else if (conflict == ConflictType::kNeedDifferent) { + AddEq(highs, {vars.col_z[a], vars.col_z[b]}, {1.0, 1.0}, 1.0); + } + } + } + + auto bank_build_begin = std::chrono::steady_clock::now(); + // Compare every SRAM write lifetime with every read lifetime in the same + // memory. For each possible producer modulo slot, classify consumer slots: + // overlapping accesses may require equal banks, different banks, or may be + // impossible regardless of bank assignment. When banks rotate across IIs, + // the effective bank is z XOR start_parity, hence AddConditionalParity; + // otherwise z alone determines the bank. The x variables gate each relation + // so it is active only for the pair of start slots selected by the solver. + // If x[prod,sp] = x[cons,sc] = 1, enforce + // z_write XOR z_read XOR parity_write XOR parity_read = required_xor + // for rotating banks, or z_write XOR z_read = required_xor otherwise. + // AddConditionalParity introduces an integer quotient to linearize the XOR; + // its big-M terms deactivate the equality for unselected slot pairs. + for (int a = 0; a < internal_count; ++a) { + const auto &write_flow = prob.flows[vars.internal_flow_ids[a]]; + if (write_flow.write_resource < 0 || write_flow.prod < 0) + continue; + for (int b = 0; b < internal_count; ++b) { + if (a == b) + continue; + const auto &read_flow = prob.flows[vars.internal_flow_ids[b]]; + if (read_flow.read_resource < 0) + continue; + if (write_flow.mem != read_flow.mem) + continue; + bool flip = mode.FlipForMem(write_flow.mem); + + for (int prod_slot = 0; prod_slot < II; ++prod_slot) { + std::vector diff_cons_slots; + std::vector same_cons_slots; + std::vector impossible_cons_slots; + for (int cons_slot = 0; cons_slot < II; ++cons_slot) { + ConflictType conflict = AnalyzeWriteReadConflict( + PositiveMod(prod_slot + write_flow.w_off, II), write_flow.w_dur, + (prod_slot + write_flow.w_off) / II & 1, + PositiveMod(cons_slot + read_flow.delta * II + read_flow.r_off, + II), + read_flow.r_dur, + (cons_slot + read_flow.delta * II + read_flow.r_off) / II & 1, II, + flip); + if (conflict == ConflictType::kNone) + continue; + if (conflict == ConflictType::kNeedDifferent) { + diff_cons_slots.push_back(cons_slot); + } else if (conflict == ConflictType::kNeedSame) { + same_cons_slots.push_back(cons_slot); + } else if (conflict == ConflictType::kImpossible) { + impossible_cons_slots.push_back(cons_slot); + } + } + + HighsInt x_prod = vars.col_x[write_flow.prod][prod_slot]; + HighsInt z_write_ping = vars.col_z[a]; + HighsInt z_read_ping = vars.col_z[b]; + HighsInt write_parity = vars.col_start_parity[write_flow.prod]; + HighsInt read_parity = vars.col_start_parity[read_flow.cons]; + + if (!diff_cons_slots.empty()) { + std::vector x_cons; + x_cons.reserve(diff_cons_slots.size()); + for (int cons_slot : diff_cons_slots) { + x_cons.push_back(vars.col_x[read_flow.cons][cons_slot]); + } + if (flip) { + AddConditionalParity(highs, z_write_ping, z_read_ping, write_parity, + read_parity, x_prod, x_cons, 1); + } else { + AddConditionalBankRelation(highs, z_write_ping, z_read_ping, x_prod, + x_cons, 1); + } + } + + if (!same_cons_slots.empty()) { + std::vector x_cons; + x_cons.reserve(same_cons_slots.size()); + for (int cons_slot : same_cons_slots) { + x_cons.push_back(vars.col_x[read_flow.cons][cons_slot]); + } + if (flip) { + AddConditionalParity(highs, z_write_ping, z_read_ping, write_parity, + read_parity, x_prod, x_cons, 0); + } else { + AddConditionalBankRelation(highs, z_write_ping, z_read_ping, x_prod, + x_cons, 0); + } + } + + if (!impossible_cons_slots.empty()) { + // Forbid selecting the producer slot together with any consumer slot + // whose physical intervals conflict even when placed on opposite + // banks. Since each command has one-hot x, this single inequality + // excludes every impossible pairing collected above. + // x[prod,sp] + sum_{sc impossible} x[cons,sc] <= 1. + std::vector idx; + std::vector val; + if (x_prod >= 0) { + idx.push_back(x_prod); + val.push_back(1.0); + } + for (int cons_slot : impossible_cons_slots) { + idx.push_back(vars.col_x[read_flow.cons][cons_slot]); + val.push_back(1.0); + } + AddLeq(highs, idx, val, x_prod >= 0 ? 1.0 : 0.0); + } + } + } + } + auto bank_build_end = std::chrono::steady_clock::now(); + double bank_build_elapsed = + std::chrono::duration(bank_build_end - bank_build_begin).count(); + LOG(INFO) << "[ILP] bank_constraint_build_elapsed=" << bank_build_elapsed + << "s"; + + return vars; +} + +SolveResult SolveFixedII(const Problem &prob, int II, bool optimize_t, + int threads, const BankFlipMode &mode) { + auto solve_begin = std::chrono::steady_clock::now(); + LOG(INFO) << "[ILP] start solve II=" << II << " optimize_t=" << optimize_t + << " N=" << prob.N << " edges=" << prob.dep_edges.size() + << " flows=" << prob.flows.size(); + Highs highs; + ModelVars vars = BuildModel(highs, prob, II, optimize_t, threads, mode); + bool model_vc_blocking_issue = + tvm::transform::PassContext::Current() + ->GetConfig(tl::kSunmmioILPModelVCBlockingIssue, Bool(true)) + .value(); + int vc_issue_constraint_count = 0; + if (model_vc_blocking_issue) { + std::vector vc_commands; + std::unordered_set vc_set; + for (int cmd = 0; cmd < prob.N; ++cmd) { + if (CommandUsesResource(prob.P[cmd], + static_cast(IlpResourceType::kVectorCore))) { + vc_commands.push_back(cmd); + vc_set.insert(cmd); + } + } + if (!vc_commands.empty()) { + for (int slot = 0; slot < II; ++slot) { + std::vector idx; + std::vector val; + for (int cmd = 0; cmd < prob.N; ++cmd) { + idx.push_back(vars.col_x[cmd][slot]); + val.push_back(vc_set.count(cmd) ? 1.0 - double(prob.N) : 1.0); + } + for (int vc : vc_commands) { + idx.push_back(vars.col_a[vc][slot]); + val.push_back(double(prob.N)); + } + AddLeq(highs, idx, val, double(prob.N)); + ++vc_issue_constraint_count; + } + } + } + LOG(INFO) << "[II=" << II + << "] VC blocking-issue constraints=" << vc_issue_constraint_count; + highs.run(); + if (highs.getModelStatus() != HighsModelStatus::kOptimal) { + auto solve_end = std::chrono::steady_clock::now(); + double elapsed = + std::chrono::duration(solve_end - solve_begin).count(); + LOG(INFO) << "[II=" << II << "] infeasible/failed, elapsed=" << elapsed + << "s, status=" << int(highs.getModelStatus()); + return {}; + } + + const HighsSolution &sol = highs.getSolution(); + SolveResult res; + res.ok = true; + res.II = II; + res.bank_slot_period = 2 * II; + res.bank_flip_mode = mode; + res.vc_blocking_issue_modeled = model_vc_blocking_issue; + res.vc_blocking_issue_constraints = vc_issue_constraint_count; + if (!optimize_t) { + auto solve_end = std::chrono::steady_clock::now(); + double elapsed = + std::chrono::duration(solve_end - solve_begin).count(); + LOG(INFO) << "[II=" << II << "] feasible_only_elapsed=" << elapsed << "s"; + return res; + } + res.t.resize(prob.N); + res.m.resize(prob.N); + res.y.resize(prob.N); + for (int i = 0; i < prob.N; ++i) { + res.t[i] = int(std::llround(sol.col_value[vars.col_t[i]])); + res.m[i] = int(std::llround(sol.col_value[vars.col_m[i]])); + res.y[i] = int(std::llround(sol.col_value[vars.col_y[i]])); + } + res.makespan = int(std::llround(sol.col_value[vars.col_T])); + res.internal_flow_ids = vars.internal_flow_ids; + res.z_bank.resize(vars.internal_flow_ids.size(), 0); + for (int vv = 0; vv < static_cast(vars.internal_flow_ids.size()); ++vv) { + res.z_bank[vv] = int(std::llround(sol.col_value[vars.col_z[vv]])); + } + auto solve_end = std::chrono::steady_clock::now(); + double elapsed = + std::chrono::duration(solve_end - solve_begin).count(); + LOG(INFO) << "[II=" << II << "] solve_elapsed=" << elapsed << "s"; + return res; +} + +SolveResult FindMinimalIIForMode(const Problem &prob, int threads, + const BankFlipMode &mode, int max_ii = -1) { + int lb = std::max(1, ResourceLowerBound(prob)); + int search_begin = std::max(1, lb); + int search_end = std::max(search_begin, std::max(1, prob.Tmax)); + if (max_ii > 0) { + search_end = std::min(search_end, max_ii); + } + if (search_end < search_begin) { + return {}; + } + constexpr int kInitialWindowSpan = 10; + int best_ii = -1; + LOG(INFO) << "[ILP] search start=" << search_begin << " end=" << search_end + << " lb=" << lb << " initial_window_span=" << kInitialWindowSpan; + + for (int window_l = search_begin; window_l <= search_end; + window_l += kInitialWindowSpan + 1) { + int window_r = std::min(search_end, window_l + kInitialWindowSpan); + int l = window_l; + int r = window_r; + int window_best = -1; + LOG(INFO) << "[ILP] search window l=" << window_l << " r=" << window_r; + while (l <= r) { + int mid = (l + r) / 2; + LOG(INFO) << "[ILP] try feasible-only II=" << mid; + SolveResult feas = SolveFixedII(prob, mid, false, threads, mode); + if (feas.ok) { + window_best = mid; + LOG(INFO) << "[ILP] feasible II=" << mid; + r = mid - 1; + } else { + LOG(INFO) << "[ILP] infeasible II=" << mid; + l = mid + 1; + } + } + if (window_best >= 0) { + best_ii = window_best; + break; + } + } + if (best_ii < 0) { + LOG(INFO) << "[ILP] no feasible II found"; + return {}; + } + LOG(INFO) << "[ILP] best feasible II=" << best_ii << " mode=" << mode.Id(); + return SolveFixedII(prob, best_ii, false, threads, mode); +} + +BankFlipMode GuessBankFlipMode(const Problem &prob) { + auto guess_for_mem = [&](int mem) { + int write_resource = GetMemoryWriteResource(mem); + int read_resource = GetMemoryReadResource(mem); + int write_count = 0; + int read_count = 0; + int write_time = 0; + int read_time = 0; + for (const CommandSpec &spec : prob.P) { + if (CommandUsesResource(spec, write_resource)) { + ++write_count; + write_time += spec.latency; + } + if (CommandUsesResource(spec, read_resource)) { + ++read_count; + read_time += spec.latency; + } + } + // Bank rotation normally follows the input-side command parity. Only let + // output-side parity decide when output occupancy clearly dominates input. + constexpr int kOutDominanceRatio = 2; + bool use_output = + read_time > kOutDominanceRatio * static_cast(write_time); + return ((use_output ? read_count : write_count) & 1) != 0; + }; + return {guess_for_mem(/*mem=*/0), guess_for_mem(/*mem=*/1)}; +} + +// Converts each annotated SunMMIO pipeline loop in a PrimFunc into a periodic +// scheduling Problem, solves it with HiGHS, and returns the same TIR loop with +// prologue/body/epilogue order plus multiversion and bank-phase annotations for +// the downstream injection pass. Unsupported or infeasible loops are returned +// with an explicit pipeline-fallback diagnostic instead of a partial schedule. +class SunmmioPipelinePlannerILP : public StmtExprMutator { +public: + static Stmt Substitute(const PrimFunc &f, bool debug) { + SunmmioPipelinePlannerILP planner(f, debug); + return planner.VisitStmt(f->body); + } + +private: + SunmmioPipelinePlannerILP(const PrimFunc &f, bool debug) + : func_(f), traverser_(f), debug_(debug) {} + + Optional FindPipelineLoop(const Stmt &stmt) { + Optional result; + PostOrderVisit(stmt, [&](const ObjectRef &obj) { + if (result.defined()) { + return; + } + if (const auto *loop = obj.as()) { + if (loop->annotations.find("num_stages") != loop->annotations.end()) { + result = ffi::GetRef(loop); + } + } + }); + return result; + } + + const SeqStmtNode *GetPipelineBodySeq(const For &loop) { + Stmt current = loop->body; + if (const auto *realize = current.as()) { + current = realize->block->body; + } + while (true) { + if (const auto *seq = current.as()) { + return seq; + } + if (const auto *if_node = current.as()) { + ICHECK(!if_node->else_case.defined()); + current = if_node->then_case; + continue; + } + if (const auto *let_node = current.as()) { + current = let_node->body; + continue; + } + return nullptr; + } + } + + struct IlpLoopAnalysis { + Problem prob; + std::vector commands; + std::set used_buffers; + std::vector versioned_buffers; + std::vector runtime_multiversion_buffers; + std::vector runtime_banked_buffers; + std::vector runtime_resident_banked_buffers; + std::map runtime_bank_start_phases; + std::map runtime_bank_read_delta_parities; + std::map> runtime_bank_writer_phases; + std::map> runtime_bank_reader_phases; + std::map runtime_bank_flip_modes; + int iterations{0}; + bool graph_valid{false}; + }; + + IlpLoopAnalysis AnalyzeLoop(const For &loop, + const SeqStmtNode *pipeline_body_seq, + int forced_iterations = -1, + const BankFlipMode &mode = BankFlipMode{}) { + IlpLoopAnalysis result; + bool export_only = GetEnvBool("TL_SUNMMIO_ILP_EXPORT_ONLY", false); + int num_stages = -1; + auto it = loop->annotations.find("num_stages"); + ICHECK(it != loop->annotations.end()); + const auto &any_ref = (*it).second; + if (const auto *imm = any_ref.as()) { + num_stages = imm->value; + } + ICHECK_GT(num_stages, 0); + if (forced_iterations > 0) { + num_stages = forced_iterations; + } + result.iterations = num_stages; + + ASTTraverser traverser(func_); + SunmmioRoleMarker role_marker(traverser, func_); + SunmmioStmtAccessAnalyzer access_analyzer(func_); + result.commands.reserve(pipeline_body_seq->seq.size()); + + std::set resource_set; + for (int i = 0; i < static_cast(pipeline_body_seq->seq.size()); ++i) { + const Stmt &stmt = pipeline_body_seq->seq[i]; + TemplateCommand cmd(i, stmt); + role_marker(stmt); + cmd.role = role_marker.GetRole(stmt); + traverser.traverse_stmt(stmt); + cmd.type = HardwareMapper::Map(stmt); + cmd.accesses = access_analyzer.Collect(stmt, loop->loop_var); + cmd.spec.latency = + static_cast(std::ceil(CostModel::EstimateDelay(cmd.type, stmt))); + cmd.spec.latency = std::max(cmd.spec.latency, 1); + cmd.spec.resources = BuildIlpResources(stmt, cmd.type, cmd.accesses); + cmd.spec.name = cmd.name + ": " + SummarizeStmtForName(stmt); + for (int resource : cmd.spec.resources) { + resource_set.insert(resource); + } + for (const BufferRegion &read : traverser.read_buffer_regions_) { + if (!IsGlobalBuffer(read->buffer)) { + result.used_buffers.insert(read->buffer); + } + } + for (const BufferRegion &write : traverser.write_buffer_regions_) { + if (!IsGlobalBuffer(write->buffer)) { + result.used_buffers.insert(write->buffer); + } + } + result.commands.push_back(cmd); + } + + result.prob.N = static_cast(result.commands.size()); + result.prob.P.resize(result.prob.N); + for (const TemplateCommand &cmd : result.commands) { + result.prob.P[cmd.id] = cmd.spec; + } + + int faster = 0; + std::vector bump_indices; + { + auto pass_ctx = tvm::transform::PassContext::Current(); + auto cfg = pass_ctx->GetConfig(tl::kSunmmioFaster); + if (cfg.defined()) { + faster = static_cast(cfg.value()->value); + } + } + if (faster <= 0) { + faster = GetEnvInt("TL_SUNMMIO_FASTER", 0); + } + if (faster <= 0) { + auto auto_selected = AutoSelectSunmmioILPFaster(result.prob.P); + faster = auto_selected.first; + bump_indices = std::move(auto_selected.second); + } + if (faster <= 0) { + faster = 1; + } + + std::unordered_set bump_index_set(bump_indices.begin(), + bump_indices.end()); + int total_latency = 0; + int latency_gcd = 0; + for (int i = 0; i < static_cast(result.commands.size()); ++i) { + if (bump_index_set.count(i)) { + result.commands[i].spec.latency += 1; + } + total_latency += result.commands[i].spec.latency; + latency_gcd = latency_gcd == 0 + ? result.commands[i].spec.latency + : GcdInt(latency_gcd, result.commands[i].spec.latency); + } + + if (latency_gcd <= 0) { + latency_gcd = 1; + } + for (TemplateCommand &cmd : result.commands) { + cmd.spec.latency /= latency_gcd; + if (faster > 1) { + cmd.spec.latency = CeilDiv(cmd.spec.latency, faster); + } + } + total_latency /= latency_gcd; + if (faster > 1) { + total_latency = CeilDiv(total_latency, faster); + } + + result.prob.Tmax = total_latency + 10; + result.prob.R.assign(resource_set.begin(), resource_set.end()); + for (int resource : result.prob.R) { + result.prob.cap[resource] = 1; + } + for (const TemplateCommand &cmd : result.commands) { + result.prob.P[cmd.id] = cmd.spec; + } + result.versioned_buffers = DetectVersionedBuffers(result.commands); + std::sort( + result.versioned_buffers.begin(), result.versioned_buffers.end(), + [](const Buffer &a, const Buffer &b) { return a->name < b->name; }); + result.runtime_multiversion_buffers.clear(); + result.runtime_banked_buffers.clear(); + result.runtime_resident_banked_buffers.clear(); + result.runtime_bank_start_phases.clear(); + result.runtime_bank_read_delta_parities.clear(); + result.runtime_bank_writer_phases.clear(); + result.runtime_bank_reader_phases.clear(); + result.runtime_bank_flip_modes.clear(); + result.prob.versioned_buffer_names.clear(); + for (const Buffer &buffer : result.versioned_buffers) { + result.prob.versioned_buffer_names.push_back(buffer->name); + } + BuildTemplateDependencyGraph(result.commands, num_stages, + result.versioned_buffers, loop->loop_var, mode, + &result.prob); + result.graph_valid = ValidateProblemGraph(result.commands, result.prob); + return result; + } + + struct StageShrinkResult { + IlpLoopAnalysis analysis; + SolveResult sol; + }; + + StageShrinkResult + FindMinimalIIAcrossFlipModes(const For &loop, + const SeqStmtNode *pipeline_body_seq, + int threads, int forced_iterations = -1) { + IlpLoopAnalysis seed = + AnalyzeLoop(loop, pipeline_body_seq, forced_iterations); + // Try every physical-bank rotation policy. Both modes allocate + // ceil(iterations / 2) versions per bank. Flip advances the version every + // two iterations while rotating banks; non-flip advances the version on + // every iteration while keeping its precolored bank fixed. + BankFlipMode guessed_mode = GuessBankFlipMode(seed.prob); + std::vector modes{guessed_mode}; + for (const BankFlipMode &mode : + {BankFlipMode{true, true}, BankFlipMode{true, false}, + BankFlipMode{false, true}, BankFlipMode{false, false}}) { + if (mode.Id() != guessed_mode.Id()) { + modes.push_back(mode); + } + } + int forced_mode = GetEnvInt("TL_SUNMMIO_ILP_FORCE_BANK_FLIP_MODE", -1); + if (forced_mode >= 0) { + ICHECK_LT(forced_mode, 4) + << "TL_SUNMMIO_ILP_FORCE_BANK_FLIP_MODE must be in [0, 3]"; + modes = {BankFlipMode{(forced_mode & 2) != 0, (forced_mode & 1) != 0}}; + } + + StageShrinkResult best; + for (size_t index = 0; index < modes.size(); ++index) { + const BankFlipMode &mode = modes[index]; + IlpLoopAnalysis candidate_analysis = + AnalyzeLoop(loop, pipeline_body_seq, forced_iterations, mode); + if (!candidate_analysis.graph_valid) { + continue; + } + SolveResult candidate; + if (!best.sol.ok || index == 0) { + candidate = + FindMinimalIIForMode(candidate_analysis.prob, threads, mode); + } else if (best.sol.II > 1) { + SolveResult probe = SolveFixedII(candidate_analysis.prob, + best.sol.II - 1, false, threads, mode); + if (probe.ok) { + candidate = FindMinimalIIForMode(candidate_analysis.prob, threads, + mode, best.sol.II - 1); + } + } + if (candidate.ok && (!best.sol.ok || candidate.II < best.sol.II)) { + best.analysis = std::move(candidate_analysis); + best.sol = std::move(candidate); + } + } + if (!best.sol.ok) { + return {std::move(seed), {}}; + } + return best; + } + + bool ShouldEnableStageShrink() const { + auto pass_ctx = tvm::transform::PassContext::Current(); + auto cfg = pass_ctx->GetConfig(tl::kSunmmioILPStageShrink); + if (cfg.defined()) { + return cfg.value()->value; + } + return false; + } + + void PopulateRuntimeBankedBuffers(IlpLoopAnalysis *analysis) const { + analysis->runtime_banked_buffers.clear(); + for (const Buffer &buffer : analysis->versioned_buffers) { + if (analysis->runtime_bank_start_phases.count(buffer->name)) { + analysis->runtime_banked_buffers.push_back(buffer); + } + } + } + + void ExportStageSolutionIfRequested(const IlpLoopAnalysis &analysis, + const SolveResult &sol, int stage) const { + std::string solution_json_path = + GetEnvString("TL_SUNMMIO_ILP_SOLUTION_JSON"); + if (solution_json_path.empty() || !sol.ok) { + return; + } + IlpLoopAnalysis export_analysis = analysis; + PopulateRuntimeBankMetadata(&export_analysis, sol); + PopulateRuntimeBankedBuffers(&export_analysis); + SolutionVerifyResult verify = VerifySolution(export_analysis.prob, sol); + WriteSolutionJson(AddStageSuffixToPath(solution_json_path, stage), + export_analysis.prob, sol, verify, + export_analysis.runtime_bank_start_phases, + export_analysis.runtime_bank_read_delta_parities, + export_analysis.runtime_bank_reader_phases); + } + + StageShrinkResult SolveWithStageShrink(const For &loop, + const SeqStmtNode *pipeline_body_seq, + int threads) { + StageShrinkResult mode_result = + FindMinimalIIAcrossFlipModes(loop, pipeline_body_seq, threads); + IlpLoopAnalysis base_analysis = std::move(mode_result.analysis); + SolveResult base_sol = std::move(mode_result.sol); + MaybeExportProblemJsonForStage(base_analysis.prob, debug_, + base_analysis.iterations); + if (!base_sol.ok) { + return {std::move(base_analysis), std::move(base_sol)}; + } + + int best_iterations = base_analysis.iterations; + BankFlipMode fixed_mode = base_sol.bank_flip_mode; + for (int candidate_iterations = base_analysis.iterations - 1; + candidate_iterations >= 1; --candidate_iterations) { + IlpLoopAnalysis candidate_analysis = AnalyzeLoop( + loop, pipeline_body_seq, candidate_iterations, fixed_mode); + MaybeExportProblemJsonForStage(candidate_analysis.prob, debug_, + candidate_iterations); + SolveResult feas = SolveFixedII(candidate_analysis.prob, base_sol.II, + false, threads, fixed_mode); + if (!feas.ok) { + continue; + } + best_iterations = candidate_iterations; + } + + IlpLoopAnalysis final_analysis = + AnalyzeLoop(loop, pipeline_body_seq, best_iterations, fixed_mode); + MaybeExportProblemJsonForStage(final_analysis.prob, debug_, + final_analysis.iterations); + SolveResult final_sol = SolveFixedII(final_analysis.prob, base_sol.II, true, + threads, fixed_mode); + ExportStageSolutionIfRequested(final_analysis, final_sol, + final_analysis.iterations); + return {std::move(final_analysis), std::move(final_sol)}; + } + + void PopulateRuntimeBankMetadata(IlpLoopAnalysis *analysis, + const SolveResult &sol) const { + std::unordered_map internal_pos; + for (int i = 0; i < static_cast(sol.internal_flow_ids.size()); ++i) { + internal_pos[sol.internal_flow_ids[i]] = i; + } + std::unordered_set has_non_resident_flow; + for (int fid = 0; fid < static_cast(analysis->prob.flows.size()); + ++fid) { + const auto &flow = analysis->prob.flows[fid]; + if (!flow.resident && !flow.buffer_name.empty()) { + has_non_resident_flow.insert(flow.buffer_name); + } + } + for (int fid = 0; fid < static_cast(analysis->prob.flows.size()); + ++fid) { + const auto &flow = analysis->prob.flows[fid]; + if (flow.buffer_name.empty()) { + continue; + } + auto flow_pos = internal_pos.find(fid); + if (flow_pos == internal_pos.end()) { + continue; + } + // Export one phase-offset convention for every runtime bank annotation: + // physical_bank = (logical_iter_parity + phase_offset) % 2 + // Physical bank 0 is ping and bank 1 is pong. The offset itself is not + // a fixed ping/pong selection because it flips with the logical + // iteration. + int phase_offset = sol.z_bank[flow_pos->second]; + ICHECK(phase_offset == 0 || phase_offset == 1) + << "ILP bank phase offset must be binary for flow " << flow.prod + << " -> " << flow.cons; + analysis->runtime_bank_flip_modes[flow.buffer_name] = + sol.bank_flip_mode.FlipForMem(flow.mem) ? 1 : 0; + bool allow_resident_to_own_bank = + flow.resident && !has_non_resident_flow.count(flow.buffer_name); + // Keep per-op bank metadata as the primary source of truth. A single + // per-buffer start phase is only well-defined when either: + // (1) the flow is resident-only for that buffer, or + // (2) every contributing flow happens to agree. + // + // Non-resident flows of the same logical buffer can legitimately land on + // different banks when they represent different runtime instances. Those + // cases are consumed later through runtime_bank_writer_phases / + // runtime_bank_reader_phases, so do not force them into a single + // runtime_bank_start_phases entry here. + if (!flow.resident || allow_resident_to_own_bank) { + auto it = analysis->runtime_bank_start_phases.find(flow.buffer_name); + if (it == analysis->runtime_bank_start_phases.end()) { + analysis->runtime_bank_start_phases[flow.buffer_name] = phase_offset; + } else if (it->second != phase_offset) { + if (flow.precolor < 0) { + LOG(WARNING) + << "Ignoring conflicting aggregate bank phase for buffer " + << flow.buffer_name << ": existing=" << it->second + << " new=" << phase_offset + << ". Per-op bank metadata will be used instead."; + } + } + int delta_parity = + sol.bank_flip_mode.FlipForMem(flow.mem) ? (flow.delta & 1) : 0; + auto it_delta = + analysis->runtime_bank_read_delta_parities.find(flow.buffer_name); + if (it_delta == analysis->runtime_bank_read_delta_parities.end()) { + analysis->runtime_bank_read_delta_parities[flow.buffer_name] = + delta_parity; + } else if (it_delta->second != delta_parity) { + LOG(WARNING) + << "Ignoring conflicting aggregate read-delta parity for buffer " + << flow.buffer_name << ": existing=" << it_delta->second + << " new=" << delta_parity + << ". Per-op bank metadata will be used instead."; + } + } else { + // For non-resident runtime instances, keep only per-op metadata and do + // not force a single aggregate start phase / delta parity. + } + if (flow.write_resource >= 0 && flow.prod >= 0) { + auto &writer_map = + analysis->runtime_bank_writer_phases[flow.buffer_name]; + auto it_writer = writer_map.find(flow.prod); + if (it_writer == writer_map.end()) { + writer_map[flow.prod] = phase_offset; + } else { + ICHECK_EQ(it_writer->second, phase_offset) + << "Conflicting writer bank phase for banked buffer " + << flow.buffer_name << " op " << flow.prod; + } + } + if (flow.cons >= 0 && flow.read_resource >= 0) { + int reader_phase_offset = + flow.resident || !sol.bank_flip_mode.FlipForMem(flow.mem) + ? phase_offset + : ((phase_offset + (flow.delta & 1)) & 1); + ICHECK(reader_phase_offset == 0 || reader_phase_offset == 1); + auto &reader_map = + analysis->runtime_bank_reader_phases[flow.buffer_name]; + auto it_reader = reader_map.find(flow.cons); + if (it_reader == reader_map.end()) { + reader_map[flow.cons] = reader_phase_offset; + } else { + ICHECK_EQ(it_reader->second, reader_phase_offset) + << "Conflicting reader bank phase for banked buffer " + << flow.buffer_name << " op " << flow.cons; + } + } + } + } + + void PruneUnnecessaryRuntimeBanking(IlpLoopAnalysis *analysis, + const SolveResult &sol) { + auto mem_needs_pingpong = [&](int mem) { + int write_resource = GetMemoryWriteResource(mem); + int read_resource = GetMemoryReadResource(mem); + if (write_resource < 0 || read_resource < 0) { + return false; + } + for (int slot = 0; slot < sol.II; ++slot) { + int write_use = 0; + int read_use = 0; + for (int i = 0; i < analysis->prob.N; ++i) { + const CommandSpec &spec = analysis->prob.P[i]; + bool uses_write = + std::find(spec.resources.begin(), spec.resources.end(), + write_resource) != spec.resources.end(); + bool uses_read = + std::find(spec.resources.begin(), spec.resources.end(), + read_resource) != spec.resources.end(); + if (!uses_write && !uses_read) { + continue; + } + int occ = + ComputeFoldedOccupancy(sol.t[i], spec.latency, sol.II, slot); + if (uses_write) { + write_use += occ; + } + if (uses_read) { + read_use += occ; + } + } + if (write_use > 0 && read_use > 0) { + return true; + } + } + return false; + }; + + bool keep_wsram = mem_needs_pingpong(/*mem=*/0); + bool keep_asram = mem_needs_pingpong(/*mem=*/1); + + std::vector pruned_banked_buffers; + for (const Buffer &buffer : analysis->runtime_banked_buffers) { + int mem = GetPingPongMemoryKind(buffer); + bool keep = (mem == 0 && keep_wsram) || (mem == 1 && keep_asram); + if (keep) { + pruned_banked_buffers.push_back(buffer); + } else { + analysis->runtime_bank_start_phases.erase(buffer->name); + analysis->runtime_bank_read_delta_parities.erase(buffer->name); + analysis->runtime_bank_writer_phases.erase(buffer->name); + analysis->runtime_bank_reader_phases.erase(buffer->name); + analysis->runtime_bank_flip_modes.erase(buffer->name); + } + } + analysis->runtime_banked_buffers = std::move(pruned_banked_buffers); + } + + Stmt VisitStmt_(const ForNode *op) final { + For loop = ffi::GetRef(op); + if (op->annotations.find("num_stages") == op->annotations.end()) { + return StmtExprMutator::VisitStmt_(op); + } + + // A single logical iteration has no cross-iteration overlap to schedule. + // Leaving it in the ILP window model can manufacture an out-of-range + // prologue/epilogue iteration when the solved makespan spans two IIs. + arith::Analyzer extent_analyzer; + PrimExpr simplified_extent = extent_analyzer.Simplify(op->extent); + const auto *extent = simplified_extent.as(); + if (extent != nullptr && extent->value <= 1) { + For sequential = Downcast(StmtExprMutator::VisitStmt_(op)); + Map annotations; + for (const auto &kv : sequential->annotations) { + if (kv.first != "num_stages") { + annotations.Set(kv.first, kv.second); + } + } + sequential.CopyOnWrite()->annotations = annotations; + return MakePipelineFallback(sequential, "ilp", "planning", + "short_extent_unsupported"); + } + + const SeqStmtNode *pipeline_body_seq = GetPipelineBodySeq(loop); + ICHECK(pipeline_body_seq != nullptr) + << "Pipeline body must normalize to SeqStmt."; + for (const Stmt &stmt : pipeline_body_seq->seq) { + if (!stmt.as() && !stmt.as() && + !stmt.as()) { + For sequential = Downcast(StmtExprMutator::VisitStmt_(op)); + Map annotations; + for (const auto &kv : sequential->annotations) { + if (kv.first != "num_stages") { + annotations.Set(kv.first, kv.second); + } + } + sequential.CopyOnWrite()->annotations = annotations; + return MakePipelineFallback(sequential, "ilp", "planning", + "unsupported_statement"); + } + } + int threads = GetEnvInt("HIGHS_THREADS", 20); + IlpLoopAnalysis analysis; + SolveResult sol; + if (ShouldEnableStageShrink()) { + StageShrinkResult shrink_result = + SolveWithStageShrink(loop, pipeline_body_seq, threads); + analysis = std::move(shrink_result.analysis); + sol = std::move(shrink_result.sol); + } else { + StageShrinkResult mode_result = + FindMinimalIIAcrossFlipModes(loop, pipeline_body_seq, threads); + analysis = std::move(mode_result.analysis); + SolveResult min_ii = std::move(mode_result.sol); + MaybeExportProblemJson(analysis.prob, debug_); + if (min_ii.ok) { + sol = SolveFixedII(analysis.prob, min_ii.II, true, threads, + min_ii.bank_flip_mode); + } + } + if (!analysis.graph_valid) { + return MakePipelineFallback(loop, "ilp", "graph_validation", + "incomplete_access_info"); + } + if (GetEnvBool("TL_SUNMMIO_ILP_EXPORT_ONLY", false)) { + return loop; + } + if (!sol.ok) { + return MakePipelineFallback(loop, "ilp", "planning", "ilp_infeasible"); + } + PopulateRuntimeBankMetadata(&analysis, sol); + std::unordered_set resident_buffer_names; + for (const FlowSpec &flow : analysis.prob.flows) { + if (flow.resident && !flow.buffer_name.empty()) { + resident_buffer_names.insert(flow.buffer_name); + } + } + for (const Buffer &buffer : analysis.versioned_buffers) { + if (analysis.runtime_bank_start_phases.count(buffer->name) || + analysis.runtime_bank_writer_phases.count(buffer->name) || + analysis.runtime_bank_reader_phases.count(buffer->name)) { + analysis.runtime_banked_buffers.push_back(buffer); + } + } + for (const Buffer &buffer : analysis.used_buffers) { + if (!resident_buffer_names.count(buffer->name)) + continue; + if (std::find(analysis.runtime_banked_buffers.begin(), + analysis.runtime_banked_buffers.end(), + buffer) == analysis.runtime_banked_buffers.end()) { + analysis.runtime_banked_buffers.push_back(buffer); + } + analysis.runtime_resident_banked_buffers.push_back(buffer); + } + // Disabled per current ILP annotation semantics: + // keep the bank-rotation decision from the solved flow model as-is, and do + // not prune ping/pong before annotation. + // + // PruneUnnecessaryRuntimeBanking(&analysis, sol); + SolutionVerifyResult verify = VerifySolution(analysis.prob, sol); + if (!verify.ok) { + return MakePipelineFallback(loop, "ilp", "planning", + "schedule_verification_failed"); + } + std::string solution_json_path = + GetEnvString("TL_SUNMMIO_ILP_SOLUTION_JSON"); + if (!solution_json_path.empty()) { + WriteSolutionJson(solution_json_path, analysis.prob, sol, verify, + analysis.runtime_bank_start_phases, + analysis.runtime_bank_read_delta_parities, + analysis.runtime_bank_reader_phases); + } + if (GetEnvBool("TL_SUNMMIO_ILP_SOLVE_ONLY", false)) { + return loop; + } + if (debug_) { + LOG(INFO) << "ILP problem N=" << analysis.prob.N + << " dep_edges=" << analysis.prob.dep_edges.size() + << " flows=" << analysis.prob.flows.size() + << " solved=" << sol.ok << " ii=" << sol.II; + } + + Map annotations; + for (const auto &kv : op->annotations) { + if (kv.first != "num_stages" && kv.first != "versioned_buffers") { + annotations.Set(kv.first, kv.second); + } + } + + int stage_count = CeilDiv(sol.makespan, std::max(1, sol.II)); + SetPipelineAppliedAnnotations(&annotations, "ilp"); + auto pass_ctx = tvm::transform::PassContext::Current(); + bool enable_lifetime_pruning = + pass_ctx + ->GetConfig(tl::kSunmmioILPMultiversionLifetimePruning, + Bool(true)) + .value(); + analysis.runtime_multiversion_buffers = DetectRuntimeMultiversionBuffers( + analysis.commands, analysis.versioned_buffers, + analysis.runtime_banked_buffers, loop->loop_var, sol, + analysis.iterations, enable_lifetime_pruning, + analysis.runtime_bank_start_phases, + analysis.runtime_bank_read_delta_parities, + analysis.runtime_bank_writer_phases, + analysis.runtime_bank_reader_phases, analysis.runtime_bank_flip_modes); + annotations.Set("iterations", Integer(analysis.iterations)); + annotations.Set("ii", Integer(sol.II)); + annotations.Set("makespan", Integer(sol.makespan)); + annotations.Set("stage_count", Integer(stage_count)); + Array prologue_orders; + Array body_orders; + Array epilogue_orders; + auto command_priority = [&](int id) { + const CommandSpec &spec = analysis.prob.P[id]; + if (CommandUsesResource(spec, + static_cast(IlpResourceType::kODMA1))) { + return 0; + } + if (CommandUsesResource(spec, + static_cast(IlpResourceType::kODMA0))) { + return 1; + } + // ODMA launch is asynchronous, while tmma.mm blocks the scalar issue + // stream until the tensor command completes. Submit same-time async + // work first so it can overlap the blocking tensor command. + if (CommandUsesResource(spec, + static_cast(IlpResourceType::kTensorCore))) { + return 2; + } + if (CommandUsesResource(spec, + static_cast(IlpResourceType::kVectorCore))) { + return 3; + } + return 4; + }; + auto starts_earlier_and_resource_priority = [&](int a, int b) { + if (sol.t[a] != sol.t[b]) { + return sol.t[a] < sol.t[b]; + } + int a_priority = command_priority(a); + int b_priority = command_priority(b); + if (a_priority != b_priority) { + return a_priority < b_priority; + } + return a < b; + }; + TimeWindowOrderResult window_orders = + BuildTimeWindowOrders(analysis.commands, analysis.iterations, sol); + annotations.Set("steady_state_max_iter_offset", + Integer(window_orders.steady_state_max_iter_offset)); + + auto time_then_non_vc_first = [&](const ExpandedOrderEntry &a, + const ExpandedOrderEntry &b) { + if (a.absolute_start != b.absolute_start) { + return a.absolute_start < b.absolute_start; + } + int a_priority = command_priority(a.id); + int b_priority = command_priority(b.id); + if (a_priority != b_priority) { + return a_priority < b_priority; + } + return a.id < b.id; + }; + + std::sort(window_orders.prologue.begin(), window_orders.prologue.end(), + time_then_non_vc_first); + std::sort(window_orders.body.begin(), window_orders.body.end(), + time_then_non_vc_first); + std::sort(window_orders.epilogue.begin(), window_orders.epilogue.end(), + time_then_non_vc_first); + + for (const ExpandedOrderEntry &entry : window_orders.prologue) { + prologue_orders.push_back( + String(std::to_string(entry.iter) + "-" + std::to_string(entry.id))); + } + for (const ExpandedOrderEntry &entry : window_orders.body) { + body_orders.push_back( + String(std::to_string(entry.iter) + "-" + std::to_string(entry.id))); + } + for (const ExpandedOrderEntry &entry : window_orders.epilogue) { + epilogue_orders.push_back( + String(std::to_string(entry.iter) + "-" + std::to_string(entry.id))); + } + + annotations.Set("prologue_orders", prologue_orders); + annotations.Set("body_orders", body_orders); + annotations.Set("epilogue_orders", epilogue_orders); + + Array used_buffers_array(analysis.used_buffers.begin(), + analysis.used_buffers.end()); + annotations.Set("used_buffers", used_buffers_array); + Array versioned_buffers_array(analysis.versioned_buffers.begin(), + analysis.versioned_buffers.end()); + annotations.Set("versioned_buffers", versioned_buffers_array); + Array runtime_multiversion_buffers_array( + analysis.runtime_multiversion_buffers.begin(), + analysis.runtime_multiversion_buffers.end()); + annotations.Set("runtime_multiversion_buffers", + runtime_multiversion_buffers_array); + Array runtime_banked_buffers_array( + analysis.runtime_banked_buffers.begin(), + analysis.runtime_banked_buffers.end()); + annotations.Set("runtime_banked_buffers", runtime_banked_buffers_array); + Array runtime_resident_banked_buffers_array( + analysis.runtime_resident_banked_buffers.begin(), + analysis.runtime_resident_banked_buffers.end()); + annotations.Set("runtime_resident_banked_buffers", + runtime_resident_banked_buffers_array); + Map runtime_bank_start_phases; + for (const Buffer &buffer : analysis.runtime_banked_buffers) { + runtime_bank_start_phases.Set( + buffer, Integer(analysis.runtime_bank_start_phases.at(buffer->name))); + } + annotations.Set("runtime_bank_start_phases", runtime_bank_start_phases); + Map runtime_bank_read_delta_parities; + for (const Buffer &buffer : analysis.runtime_banked_buffers) { + auto it = analysis.runtime_bank_read_delta_parities.find(buffer->name); + if (it != analysis.runtime_bank_read_delta_parities.end()) { + runtime_bank_read_delta_parities.Set(buffer, Integer(it->second)); + } + } + annotations.Set("runtime_bank_read_delta_parities", + runtime_bank_read_delta_parities); + Map> runtime_bank_writer_phases; + for (const Buffer &buffer : analysis.runtime_banked_buffers) { + auto it = analysis.runtime_bank_writer_phases.find(buffer->name); + if (it == analysis.runtime_bank_writer_phases.end()) { + continue; + } + Map per_op; + for (const auto &op_phase : it->second) { + per_op.Set(Integer(op_phase.first), Integer(op_phase.second)); + } + runtime_bank_writer_phases.Set(buffer, per_op); + } + annotations.Set("runtime_bank_writer_phases", runtime_bank_writer_phases); + Map> runtime_bank_reader_phases; + for (const Buffer &buffer : analysis.runtime_banked_buffers) { + auto it = analysis.runtime_bank_reader_phases.find(buffer->name); + if (it == analysis.runtime_bank_reader_phases.end()) { + continue; + } + Map per_op; + for (const auto &op_phase : it->second) { + per_op.Set(Integer(op_phase.first), Integer(op_phase.second)); + } + runtime_bank_reader_phases.Set(buffer, per_op); + } + annotations.Set("runtime_bank_reader_phases", runtime_bank_reader_phases); + Map runtime_bank_flip_modes; + for (const Buffer &buffer : analysis.runtime_banked_buffers) { + auto it = analysis.runtime_bank_flip_modes.find(buffer->name); + if (it != analysis.runtime_bank_flip_modes.end()) { + runtime_bank_flip_modes.Set(buffer, Integer(it->second)); + } + } + annotations.Set("runtime_bank_flip_modes", runtime_bank_flip_modes); + + Stmt body = this->VisitStmt(op->body); + For new_loop = loop; + ForNode *loop_ptr = new_loop.CopyOnWrite(); + loop_ptr->body = body; + loop_ptr->annotations = annotations; + return new_loop; + } + + PrimFunc func_; + ASTTraverser traverser_; + bool debug_{false}; +}; + +tvm::transform::Pass SunmmioPipelinePlanningILP(bool debug = false) { + using namespace tir::transform; + auto pass_func = [=](PrimFunc f, const IRModule &m, PassContext ctx) { + PrimFuncNode *fptr = f.CopyOnWrite(); + fptr->body = SunmmioPipelinePlannerILP::Substitute(f, debug); + return f; + }; + return CreatePrimFuncPass(pass_func, 0, "tl.SunmmioPipelinePlanningILP", {}); +} + +} // namespace bank_ilp_internal + +tvm::transform::Pass SunmmioPipelinePlanningILP(bool debug = false) { + return bank_ilp_internal::SunmmioPipelinePlanningILP(debug); +} + +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + refl::GlobalDef().def("tl.transform.SunmmioPipelinePlanningILP", + SunmmioPipelinePlanningILP); +} + +} // namespace tl +} // namespace tvm diff --git a/src/transform/sunmmio_tile_loop_fusion/discovery.cc b/src/transform/sunmmio_tile_loop_fusion/discovery.cc index 1c3df86cca..f1a2a3709d 100644 --- a/src/transform/sunmmio_tile_loop_fusion/discovery.cc +++ b/src/transform/sunmmio_tile_loop_fusion/discovery.cc @@ -158,6 +158,91 @@ class RegionLoopCollector : public StmtVisitor { } }; +class LastIterationWriteCollector : public StmtVisitor { +public: + explicit LastIterationWriteCollector(std::vector execution_loops) + : execution_loops_(std::move(execution_loops)) {} + + void Collect(const Stmt &stmt) { VisitStmt(stmt); } + + void Invalidate(const Buffer &buffer) { + all_writes_last_by_depth_[buffer->data.get()] = + std::vector(execution_loops_.size(), false); + } + + int GetSafeExecutionDepth(const Buffer &buffer, int available_depth) const { + auto it = all_writes_last_by_depth_.find(buffer->data.get()); + if (it == all_writes_last_by_depth_.end()) { + return -1; + } + + int safe_depth = available_depth; + while (safe_depth < static_cast(it->second.size()) && + it->second[safe_depth]) { + ++safe_depth; + } + return safe_depth > available_depth ? safe_depth : -1; + } + +private: + bool PathImpliesLastIteration(const For &loop) const { + if (path_conditions_.empty()) { + return false; + } + + PrimExpr path_condition = Bool(true); + for (const PrimExpr &condition : path_conditions_) { + path_condition = tir::And(path_condition, condition); + } + + arith::Analyzer analyzer; + With constraint(&analyzer, path_condition); + PrimExpr last_iteration = analyzer.Simplify(loop->min + loop->extent - 1); + return analyzer.CanProve(loop->loop_var == last_iteration); + } + + void VisitStmt_(const BufferStoreNode *op) final { + std::vector write_is_last; + write_is_last.reserve(execution_loops_.size()); + for (const For &loop : execution_loops_) { + write_is_last.push_back(PathImpliesLastIteration(loop)); + } + + auto it = all_writes_last_by_depth_.find(op->buffer->data.get()); + if (it == all_writes_last_by_depth_.end()) { + all_writes_last_by_depth_.emplace(op->buffer->data.get(), + std::move(write_is_last)); + return; + } + for (size_t depth = 0; depth < it->second.size(); ++depth) { + it->second[depth] = it->second[depth] && write_is_last[depth]; + } + } + + void VisitStmt_(const IfThenElseNode *op) final { + path_conditions_.push_back(op->condition); + VisitStmt(op->then_case); + path_conditions_.pop_back(); + + if (op->else_case.defined()) { + path_conditions_.push_back(tir::Not(op->condition)); + VisitStmt(op->else_case.value()); + path_conditions_.pop_back(); + } + } + + void VisitStmt_(const BlockRealizeNode *op) final { + path_conditions_.push_back(op->predicate); + VisitStmt(op->block); + path_conditions_.pop_back(); + } + + std::vector execution_loops_; + std::vector path_conditions_; + std::unordered_map> + all_writes_last_by_depth_; +}; + Map CollectVisibleBuffers(const PrimFunc &func) { Map buffers; for (const auto &kv : func->buffer_map) { @@ -481,6 +566,7 @@ struct ActiveDefInfo { struct OverlapFacts { int rho{0}; + int max_shared_execution_depth{0}; int64_t weight_bytes{0}; }; @@ -585,6 +671,15 @@ ComputeOverlapFacts(const NormalizedBufferAccess &src_access, OverlapFacts facts; facts.rho = ComputeRequiredSharedPrefixDepth(src_access, dst_access); + facts.max_shared_execution_depth = + std::min(src_access.home_depth, dst_access.home_depth); + // A destination write restricted to the final iteration cannot clobber a + // source read from a later iteration at any proven-safe depth. + if (kind == TileScopeDependenceKind::kWAR && + dst_access.last_write_safe_depth >= 0) { + facts.max_shared_execution_depth = std::max( + facts.max_shared_execution_depth, dst_access.last_write_safe_depth); + } facts.weight_bytes = ComputeEdgeWeightBytes( src_access, dst_access, exact_overlap.value(), kind, analyzer); return facts; @@ -594,8 +689,13 @@ TileScopeDependenceEdge MakeDependenceEdge(int src_region_index, int dst_region_index, TileScopeDependenceKind kind, int src_access_index, int dst_access_index, const OverlapFacts &facts) { - return {src_region_index, dst_region_index, kind, - src_access_index, dst_access_index, facts.rho, + return {src_region_index, + dst_region_index, + kind, + src_access_index, + dst_access_index, + facts.rho, + facts.max_shared_execution_depth, facts.weight_bytes}; } @@ -645,9 +745,21 @@ AnalyzeOneTileScopeRegion(const PlannerVisibleRegionMatch &match, std::vector available_at_execution_depths; available_at_execution_depths.reserve(def_out.size()); + std::vector last_write_safe_execution_depths; + last_write_safe_execution_depths.reserve(def_out.size()); + LastIterationWriteCollector last_write_collector( + loop_collector.execution_loops); + last_write_collector.Collect(scope_entry_for); + for (const BufferRegion ®ion : opaque_access_collector.writes) { + last_write_collector.Invalidate(region->buffer); + } for (const BufferRegion ®ion : def_out) { - available_at_execution_depths.push_back( - ComputeAvailableExecutionDepth(region, loop_collector.execution_loops)); + int available_depth = + ComputeAvailableExecutionDepth(region, loop_collector.execution_loops); + available_at_execution_depths.push_back(available_depth); + last_write_safe_execution_depths.push_back( + last_write_collector.GetSafeExecutionDepth(region->buffer, + available_depth)); } TileScopeRegion summary; @@ -664,6 +776,7 @@ AnalyzeOneTileScopeRegion(const PlannerVisibleRegionMatch &match, summary.use_in = use_in; summary.def_out = def_out; summary.available_at_execution_depths = available_at_execution_depths; + summary.last_write_safe_execution_depths = last_write_safe_execution_depths; return summary; } diff --git a/src/transform/sunmmio_tile_loop_fusion/planner.cc b/src/transform/sunmmio_tile_loop_fusion/planner.cc index 3d3fd619ac..f4d9152e7c 100644 --- a/src/transform/sunmmio_tile_loop_fusion/planner.cc +++ b/src/transform/sunmmio_tile_loop_fusion/planner.cc @@ -140,6 +140,7 @@ BuildWindowPlannerInput(const SunmmioTileLoopFusionWindowProblem &problem) { planner_edge.dst_local_index = dst_it->second; planner_edge.kind = edge.kind; planner_edge.rho = edge.rho; + planner_edge.max_shared_execution_depth = edge.max_shared_execution_depth; planner_edge.weight = edge.weight_bytes; planner_edge.instance_count = ComputeRawEdgeInstanceCount( problem.regions[planner_edge.src_local_index].execution_loop_extents, diff --git a/src/transform/sunmmio_tile_loop_fusion/planner_internal.h b/src/transform/sunmmio_tile_loop_fusion/planner_internal.h index 1bea9086e5..e3453e9a3c 100644 --- a/src/transform/sunmmio_tile_loop_fusion/planner_internal.h +++ b/src/transform/sunmmio_tile_loop_fusion/planner_internal.h @@ -83,8 +83,8 @@ struct WindowPlannerRegionInfo { /*! * \brief Planner-local dependence edge with precomputed multiplicity data. * - * Each edge carries the legality kind, the shared-prefix depth \c rho, and the - * specific access/region ids needed to check coverage during scheduling. + * Each edge carries the legality kind, reuse depth, maximum ordering-safe + * shared depth, and the access/region ids needed during scheduling. */ struct WindowPlannerEdgeInfo { int src_local_index{-1}; @@ -93,6 +93,7 @@ struct WindowPlannerEdgeInfo { int buffer_region_id{-1}; std::string buffer_name; int rho{0}; + int max_shared_execution_depth{0}; int64_t weight{0}; int64_t instance_count{1}; int covered_use_index{-1}; @@ -164,12 +165,13 @@ struct ResidentValueState { }; /*! - * \brief One open shared shell frame plus the residents attached to it. + * \brief One open shared shell frame, its residents, and incomplete edges. */ struct OpenScopeFrame { std::vector shell_axes; Array shell_extents; std::vector residents; + std::vector pending_edge_indices; }; /*! \brief Full memoized solver state for one partial schedule. */ diff --git a/src/transform/sunmmio_tile_loop_fusion/planner_solver.cc b/src/transform/sunmmio_tile_loop_fusion/planner_solver.cc index 1b5531188a..fa65ff4cc8 100644 --- a/src/transform/sunmmio_tile_loop_fusion/planner_solver.cc +++ b/src/transform/sunmmio_tile_loop_fusion/planner_solver.cc @@ -67,6 +67,54 @@ bool ResidentValueLess(const ResidentValueState &lhs, rhs.home_depth, rhs.payload_bytes, rhs.instance_count); } +bool RetainsPendingIncomingEdge(const WindowPlannerInput &input, + const PlannerState &state, int close_to_depth, + int region_local_index) { + const std::vector &incoming_edges = + input.incoming_edges_by_dst[region_local_index]; + for (int depth = 0; depth < close_to_depth; ++depth) { + const std::vector &pending_edges = + state.open_scopes[depth].pending_edge_indices; + for (int edge_index : incoming_edges) { + if (std::binary_search(pending_edges.begin(), pending_edges.end(), + edge_index)) { + return true; + } + } + } + return false; +} + +void InstallPendingEdgeIfMissing(std::vector *open_scopes, + int frame_index, int edge_index) { + ICHECK_GE(frame_index, 0); + ICHECK_LT(frame_index, static_cast(open_scopes->size())); + std::vector &pending_edges = + (*open_scopes)[frame_index].pending_edge_indices; + auto it = + std::lower_bound(pending_edges.begin(), pending_edges.end(), edge_index); + if (it == pending_edges.end() || *it != edge_index) { + pending_edges.insert(it, edge_index); + } +} + +int FindPendingFrameIndex(const WindowPlannerEdgeInfo &edge, + const TileScopeRegion &source_region, + int open_to_depth) { + ICHECK_GE(edge.max_shared_execution_depth, 0); + ICHECK_LE(open_to_depth, + static_cast(source_region.execution_loop_extents.size())); + for (int depth = edge.max_shared_execution_depth + 1; depth <= open_to_depth; + ++depth) { + const auto *extent = + source_region.execution_loop_extents[depth - 1].as(); + if (extent == nullptr || extent->value != 1) { + return depth - 1; + } + } + return -1; +} + struct MutablePlannerTreeNode { bool is_scope{false}; int region_index{-1}; @@ -181,7 +229,11 @@ std::string SerializePlannerState(const PlannerState &state) { os << SerializeDynamicBitset(state.scheduled_mask); for (const OpenScopeFrame &frame : state.open_scopes) { os << '[' << JoinAxes(frame.shell_axes) << '@' - << JoinExtents(frame.shell_extents) << '|'; + << JoinExtents(frame.shell_extents) << "|P:"; + for (int edge_index : frame.pending_edge_indices) { + os << edge_index << ','; + } + os << "|R:"; for (const ResidentValueState &resident : frame.residents) { os << SerializeResident(resident) << ','; } @@ -344,8 +396,8 @@ int64_t ComputeLiveRangeDelta(const PlannerState &state) { * latter. * 3. Materialize uncovered reads that are not already resident, charging * shared-read cost when they must be fetched. - * 4. Kill overwritten residents and install the current region's new - * definitions so later actions can reuse them. + * 4. Kill overwritten residents, install new definitions, and record outgoing + * dependences that remain incomplete under a deeper multi-trip shell. * 5. Mark the region scheduled, prune dead residents, and recompute the * live-range and reorder score terms for the next state. * @@ -386,6 +438,7 @@ TransitionResult ApplyAction(const WindowPlannerInput &input, result.next_state.open_scopes.push_back( {TakeExecutionAxisPrefix(region.logical_execution_axis_keys, depth), TakeExecutionExtentPrefix(region.execution_loop_extents, depth), + {}, {}}); } @@ -474,6 +527,23 @@ TransitionResult ApplyAction(const WindowPlannerInput &input, edge.instance_count}); } + // An edge whose source is emitted below its maximum legal shared depth stays + // pending until the first deeper multi-trip frame closes. Discovery may + // raise that depth for a WAR edge when every destination write is proven to + // occur on the final iteration, which preserves safe read-before-write + // fusion without weakening the default barrier. Static unit-trip shells + // cannot interleave executions and therefore remain legally fused. + for (int edge_index : input.outgoing_edges_by_src[region_local_index]) { + const WindowPlannerEdgeInfo &edge = input.edges[edge_index]; + int pending_frame_index = + FindPendingFrameIndex(edge, region, open_to_depth); + if (pending_frame_index < 0) { + continue; + } + InstallPendingEdgeIfMissing(&result.next_state.open_scopes, + pending_frame_index, edge_index); + } + // Finalize the new planner state and score terms after scheduling this // region. result.next_state.scheduled_mask = state.scheduled_mask; @@ -551,7 +621,8 @@ MemoResult BuildSourceOrderFallbackPlan(const WindowPlannerInput &input) { * scheduled, closes the current shell stack to \p close_to_depth, and then * reopens shells for the candidate region until \p open_to_depth. Prefix reuse * is only legal when the kept shells match the candidate region's execution - * prefix axes and extents. + * prefix axes and extents, and when no incoming dependence remains pending in + * that retained prefix. * * Because every planner score term is nonnegative, the search can prune any * branch whose immediate transition delta is already no better than the best @@ -632,6 +703,10 @@ MemoResult SolveWindowPlan(const WindowPlannerInput &input, region.execution_loop_extents)) { continue; } + if (RetainsPendingIncomingEdge(input, state, close_to_depth, + region_local_index)) { + continue; + } for (int open_to_depth = close_to_depth; open_to_depth <= static_cast(region.logical_execution_axis_keys.size()); diff --git a/src/transform/sunmmio_tile_loop_fusion/types.h b/src/transform/sunmmio_tile_loop_fusion/types.h index 62bd55d933..e04448984d 100644 --- a/src/transform/sunmmio_tile_loop_fusion/types.h +++ b/src/transform/sunmmio_tile_loop_fusion/types.h @@ -73,6 +73,10 @@ struct TileScopeRegion { // Example: a rank-2 reduction region may still produce a row value available // at execution depth 1. std::vector available_at_execution_depths; + // For each `def_out`, the deepest additional execution depth through which + // every write is provably restricted to the final iteration. A value of -1 + // means that no such extension beyond availability was proven. + std::vector last_write_safe_execution_depths; }; /*! @@ -115,13 +119,14 @@ struct NormalizedBufferAccessDim { * * The planner and dependence builder both need the same semantic facts for an * access: the normalized region shape, the execution-depth provenance of each - * dimension, the depth at which the value becomes available, and the payload - * size used by the cost model. + * dimension, the depth at which the value becomes available, any proven + * final-iteration write phase, and the payload size used by the cost model. */ struct NormalizedBufferAccess { tir::BufferRegion region; std::vector dims; int home_depth{0}; + int last_write_safe_depth{-1}; int64_t payload_bytes{0}; }; @@ -149,6 +154,9 @@ struct TileScopeDependenceEdge { // rho = 1 means the edge can stay internal under a shared outer row shell; // rho = 2 means it needs a deeper tile shell, and so on. int rho{0}; + // Deepest execution-loop prefix the source and destination may share without + // interleaving different executions of this dependence instance. + int max_shared_execution_depth{0}; // Estimated payload cost of cutting this edge for one execution instance. // The planner later scales this by execution multiplicity. int64_t weight_bytes{0}; diff --git a/src/transform/sunmmio_tile_loop_fusion/utils.cc b/src/transform/sunmmio_tile_loop_fusion/utils.cc index 45b204326b..27e8696692 100644 --- a/src/transform/sunmmio_tile_loop_fusion/utils.cc +++ b/src/transform/sunmmio_tile_loop_fusion/utils.cc @@ -78,7 +78,8 @@ NormalizedBufferAccess BuildNormalizedBufferAccess( const BufferRegion ®ion, const Map &subst, const std::unordered_map &canonical_execution_vars, const std::vector &logical_execution_axis_keys, - int home_depth_override, arith::Analyzer *analyzer) { + int home_depth_override, int last_write_safe_depth, + arith::Analyzer *analyzer) { BufferRegion normalized_region = NormalizeBufferRegionByLogicalExecutionAxes(region, subst); @@ -107,7 +108,7 @@ NormalizedBufferAccess BuildNormalizedBufferAccess( static_cast(normalized_region->region.size())); } - return {normalized_region, dims, home_depth, + return {normalized_region, dims, home_depth, last_write_safe_depth, ComputeAccessPayloadBytes(normalized_region, analyzer)}; } @@ -187,7 +188,7 @@ NormalizeRegionBoundaries(const std::vector ®ions) { normalized.use_in.push_back(BuildNormalizedBufferAccess( buffer_region, subst, canonical_execution_vars, region.logical_execution_axis_keys, /*home_depth_override=*/-1, - &analyzer)); + /*last_write_safe_depth=*/-1, &analyzer)); } for (size_t i = 0; i < region.def_out.size(); ++i) { @@ -195,9 +196,14 @@ NormalizeRegionBoundaries(const std::vector ®ions) { if (i < region.available_at_execution_depths.size()) { home_depth = region.available_at_execution_depths[i]; } + int last_write_safe_depth = -1; + if (i < region.last_write_safe_execution_depths.size()) { + last_write_safe_depth = region.last_write_safe_execution_depths[i]; + } normalized.def_out.push_back(BuildNormalizedBufferAccess( region.def_out[i], subst, canonical_execution_vars, - region.logical_execution_axis_keys, home_depth, &analyzer)); + region.logical_execution_axis_keys, home_depth, last_write_safe_depth, + &analyzer)); } normalized_regions.push_back(std::move(normalized)); diff --git a/sunmmio_kernel/softplus/test_softplus_1d_dynamic_opt_validate.py b/sunmmio_kernel/softplus/test_softplus_1d_dynamic_opt_validate.py index ad4944607a..72ed65d39f 100644 --- a/sunmmio_kernel/softplus/test_softplus_1d_dynamic_opt_validate.py +++ b/sunmmio_kernel/softplus/test_softplus_1d_dynamic_opt_validate.py @@ -29,20 +29,20 @@ def softplus_1d_dynamic(block_N=256, in_dtype=T.float32, out_dtype=T.float32): N = T.dynamic("n") row_major_layout = make_row_major((N,)) - placement = T.MeshShardingPolicy(cross_mesh_dim=0) + placement = T.placement.mesh_as_line(0) @T.prim_func def elem_softplus( A: T.MeshTensor((N,), placement, dtype=in_dtype, layout=row_major_layout), # type: ignore B: T.MeshTensor((N,), placement, dtype=out_dtype, layout=row_major_layout), # type: ignore ): - with T.Kernel() as cid: + with T.Kernel(): A_shared = T.alloc_shared((block_N,), in_dtype) B_shared = T.alloc_shared((block_N,), out_dtype) - for bx in T.serial(T.ceildiv(A.get_local_extent(cid)[0], block_N)): + for bx in T.serial(T.ceildiv(A.get_local_extent()[0], block_N)): T.copy(A[bx * block_N : (bx + 1) * block_N], A_shared) - for i in T.Tiles([T.min(block_N, A.get_local_extent(cid)[0] - bx * block_N)]): + for i in T.Tiles([T.min(block_N, A.get_local_extent()[0] - bx * block_N)]): value = A_shared[i] B_shared[i] = T.max(value, 0) + T.log(1 + T.exp(-T.abs(value))) T.copy(B_shared, B[bx * block_N : (bx + 1) * block_N]) @@ -56,20 +56,20 @@ def softplus_1d_dynamic_inline(block_N=256, in_dtype=T.float32, out_dtype=T.floa N = T.dynamic("n") row_major_layout = make_row_major((N,)) - placement = T.MeshShardingPolicy(cross_mesh_dim=0) + placement = T.placement.mesh_as_line(0) @T.prim_func def elem_softplus_inline( A: T.MeshTensor((N,), placement, dtype=in_dtype, layout=row_major_layout), # type: ignore B: T.MeshTensor((N,), placement, dtype=out_dtype, layout=row_major_layout), # type: ignore ): - with T.Kernel() as cid: + with T.Kernel(): A_shared = T.alloc_shared((block_N,), in_dtype) B_shared = T.alloc_shared((block_N,), out_dtype) - for bx in T.serial(T.ceildiv(A.get_local_extent(cid)[0], block_N)): + for bx in T.serial(T.ceildiv(A.get_local_extent()[0], block_N)): T.copy(A[bx * block_N : (bx + 1) * block_N], A_shared) - for i in T.Tiles([T.min(block_N, A.get_local_extent(cid)[0] - bx * block_N)]): + for i in T.Tiles([T.min(block_N, A.get_local_extent()[0] - bx * block_N)]): B_shared[i] = T.max(A_shared[i], 0) + T.log(1 + T.exp(-T.abs(A_shared[i]))) T.copy(B_shared, B[bx * block_N : (bx + 1) * block_N]) diff --git a/testing/cpp/transform/sunmmio_tile_loop_fusion/discovery_test.cc b/testing/cpp/transform/sunmmio_tile_loop_fusion/discovery_test.cc index 664fa45028..19e8c4b595 100644 --- a/testing/cpp/transform/sunmmio_tile_loop_fusion/discovery_test.cc +++ b/testing/cpp/transform/sunmmio_tile_loop_fusion/discovery_test.cc @@ -118,6 +118,52 @@ Stmt Make2DTileRead(const Buffer &src, const std::string &axis0_name = "i", return body; } +Stmt Make2DRowSummary(const Buffer &dst, const Buffer &src, + const std::string &axis0_name = "i", + const std::string &axis1_name = "j") { + Var axis0(axis0_name, DataType::Int(32)); + Var axis1(axis1_name, DataType::Int(32)); + Var ki("ki", DataType::Int(32)); + Var kj("kj", DataType::Int(32)); + + PrimExpr row = axis0 * I(8) + ki; + PrimExpr col = axis1 * I(32) + kj; + + Stmt body = BufferStore(dst, BufferLoad(src, {row, col}), {row}); + body = For(kj, 0, I(32), ForKind::kVectorized, body, Optional(), + MakeInteriorAnnotation(1)); + body = For(ki, 0, I(8), ForKind::kSerial, body, Optional(), + MakeInteriorAnnotation(0)); + body = For(axis1, 0, I(2), ForKind::kSerial, body, Optional(), + MakeExecutionAxisAnnotation(1)); + body = For(axis0, 0, I(4), ForKind::kSerial, body, Optional(), + MakeScopeEntryAnnotationsWithTileSize({0, 1}, 8, 32, 32, 64)); + return body; +} + +Stmt Make2DRowConsumer(const Buffer &dst, const Buffer &src, + const std::string &axis0_name = "i", + const std::string &axis1_name = "j") { + Var axis0(axis0_name, DataType::Int(32)); + Var axis1(axis1_name, DataType::Int(32)); + Var ki("ki", DataType::Int(32)); + Var kj("kj", DataType::Int(32)); + + PrimExpr row = axis0 * I(8) + ki; + PrimExpr col = axis1 * I(32) + kj; + + Stmt body = BufferStore(dst, BufferLoad(src, {row}), {row, col}); + body = For(kj, 0, I(32), ForKind::kVectorized, body, Optional(), + MakeInteriorAnnotation(1)); + body = For(ki, 0, I(8), ForKind::kSerial, body, Optional(), + MakeInteriorAnnotation(0)); + body = For(axis1, 0, I(2), ForKind::kSerial, body, Optional(), + MakeExecutionAxisAnnotation(1)); + body = For(axis0, 0, I(4), ForKind::kSerial, body, Optional(), + MakeScopeEntryAnnotationsWithTileSize({0, 1}, 8, 32, 32, 64)); + return body; +} + Stmt Make2DNarrowTileCopy(const Buffer &dst, const Buffer &src, const PrimExpr &src_col_shift) { Var axis0("i", DataType::Int(32)); @@ -267,9 +313,33 @@ TEST(SunmmioTileLoopFusionDiscoveryTest, EXPECT_EQ(edge.dst_region_index, 1); EXPECT_EQ(edge.kind, TileScopeDependenceKind::kRAW); EXPECT_EQ(edge.rho, 2); + EXPECT_EQ(edge.max_shared_execution_depth, 2); EXPECT_EQ(edge.weight_bytes, 1024); } +TEST(SunmmioTileLoopFusionDiscoveryTest, + RowDependenceLimitsSharedExecutionToOuterDepth) { + Buffer a_shared = MakeSharedBuffer("A_shared", {I(32), I(64)}); + Buffer row_summary = MakeSharedBuffer("row_summary", {I(32)}); + Buffer b_shared = MakeSharedBuffer("B_shared", {I(32), I(64)}); + + PrimFunc func = MakePrimFunc(SeqStmt( + Array{Make2DRowSummary(row_summary, a_shared), + Make2DRowConsumer(b_shared, row_summary, "ii", "jj")})); + SunmmioTileLoopFusionProgram program = + BuildSunmmioTileLoopFusionProgram(func); + std::vector problems = + BuildSunmmioTileLoopFusionWindowProblems(program); + + ASSERT_EQ(problems.size(), 1U); + const TileScopeWindowGraph &graph = problems[0].graph; + ASSERT_EQ(graph.edges.size(), 1U); + const TileScopeDependenceEdge &edge = graph.edges[0]; + EXPECT_EQ(edge.kind, TileScopeDependenceKind::kRAW); + EXPECT_EQ(edge.rho, 1); + EXPECT_EQ(edge.max_shared_execution_depth, 1); +} + TEST(SunmmioTileLoopFusionDiscoveryTest, ReadAfterReadOnSameBufferProducesNoDependenceEdges) { Buffer a_shared = MakeSharedBuffer("A_shared", {I(32), I(32)}); diff --git a/testing/cpp/transform/sunmmio_tile_loop_fusion/planner_internal_test.cc b/testing/cpp/transform/sunmmio_tile_loop_fusion/planner_internal_test.cc index 82a9136fdd..7292266d54 100644 --- a/testing/cpp/transform/sunmmio_tile_loop_fusion/planner_internal_test.cc +++ b/testing/cpp/transform/sunmmio_tile_loop_fusion/planner_internal_test.cc @@ -60,6 +60,7 @@ WindowPlannerInput MakeRawCoverageInput() { first_edge.buffer_region_id = 10; first_edge.buffer_name = "debug_buffer"; first_edge.rho = 1; + first_edge.max_shared_execution_depth = 1; first_edge.weight = 64; first_edge.instance_count = 1; first_edge.covered_use_index = 0; @@ -116,6 +117,57 @@ TEST(SunmmioTileLoopFusionPlannerInternalTest, EXPECT_EQ(instance_scopes[1].residents.size(), 2U); } +TEST(SunmmioTileLoopFusionPlannerInternalTest, + PlannerStateSerializationIncludesPendingDependences) { + PlannerState first{DynamicBitset(2), std::vector(2)}; + PlannerState second = first; + first.open_scopes[1].pending_edge_indices = {1}; + second.open_scopes[1].pending_edge_indices = {2}; + + EXPECT_NE(SerializePlannerState(first), SerializePlannerState(second)); +} + +TEST(SunmmioTileLoopFusionPlannerInternalTest, + ApplyActionTracksPendingDependences) { + for (TileScopeDependenceKind kind : + {TileScopeDependenceKind::kRAW, TileScopeDependenceKind::kWAR, + TileScopeDependenceKind::kWAW}) { + WindowPlannerInput input; + auto *problem = new SunmmioTileLoopFusionWindowProblem(); + problem->regions.resize(2); + for (int i = 0; i < 2; ++i) { + problem->regions[i].global_region_index = i; + problem->regions[i].logical_execution_axis_keys = {"i", "j"}; + problem->regions[i].execution_loop_extents = { + IntImm(DataType::Int(32), 8), IntImm(DataType::Int(32), 2)}; + } + + input.problem = problem; + input.regions.resize(2); + input.incoming_edges_by_dst.resize(2); + input.outgoing_edges_by_src.resize(2); + input.predecessor_masks.assign(2, DynamicBitset(2)); + input.earlier_source_masks.assign(2, DynamicBitset(2)); + + WindowPlannerEdgeInfo edge; + edge.src_local_index = 0; + edge.dst_local_index = 1; + edge.kind = kind; + edge.max_shared_execution_depth = 1; + input.edges.push_back(edge); + input.outgoing_edges_by_src[0].push_back(0); + input.incoming_edges_by_dst[1].push_back(0); + + PlannerState initial{DynamicBitset(2), {}}; + TransitionResult transition = ApplyAction(input, initial, 0, 0, 2); + + ASSERT_EQ(transition.next_state.open_scopes.size(), 2U); + EXPECT_EQ(transition.next_state.open_scopes[1].pending_edge_indices, + std::vector({0})); + delete problem; + } +} + TEST(SunmmioTileLoopFusionPlannerInternalTest, RawCoverageOnlySuppressesExactCoveredUses) { WindowPlannerInput input = MakeRawCoverageInput(); diff --git a/testing/cpp/transform/sunmmio_tile_loop_fusion/planner_test.cc b/testing/cpp/transform/sunmmio_tile_loop_fusion/planner_test.cc index 61c7cb485d..4b8aa00e39 100644 --- a/testing/cpp/transform/sunmmio_tile_loop_fusion/planner_test.cc +++ b/testing/cpp/transform/sunmmio_tile_loop_fusion/planner_test.cc @@ -94,7 +94,8 @@ NormalizedBufferAccess MakeAccess(const Buffer &buffer, int home_depth, TileScopeDependenceEdge MakeRawEdge(int src_region_index, int dst_region_index, int src_access_index, int dst_access_index, - int rho, int64_t weight_bytes) { + int rho, int64_t weight_bytes, + int max_shared_execution_depth = -1) { TileScopeDependenceEdge edge; edge.src_region_index = src_region_index; edge.dst_region_index = dst_region_index; @@ -102,6 +103,8 @@ TileScopeDependenceEdge MakeRawEdge(int src_region_index, int dst_region_index, edge.src_access_index = src_access_index; edge.dst_access_index = dst_access_index; edge.rho = rho; + edge.max_shared_execution_depth = + max_shared_execution_depth < 0 ? rho : max_shared_execution_depth; edge.weight_bytes = weight_bytes; return edge; } @@ -285,6 +288,83 @@ TEST(SunmmioTileLoopFusionPlannerTest, SharedTileRawEdgeBuildsNestedTileShell) { EXPECT_EQ(inner_scope.children[1].region_index, 1); } +TEST(SunmmioTileLoopFusionPlannerTest, + ShallowDependenceSplitsOnlyInnerShellDespiteProfitableTileReuse) { + Buffer row_buffer = MakeSharedBuffer("row_buffer", {I(8)}); + Buffer tile_buffer = MakeSharedBuffer("tile_buffer", {I(8), I(2)}); + + SunmmioTileLoopFusionWindowProblem problem; + problem.regions = { + MakePlannerRegion(0, {"i", "j"}, MakeExtents({8, 2})), + MakePlannerRegion(1, {"i", "j"}, MakeExtents({8, 2})), + }; + problem.normalized_regions.resize(2); + problem.normalized_regions[0].def_out = { + MakeAccess(row_buffer, 1, 32), + MakeAccess(tile_buffer, 2, 1024), + }; + problem.normalized_regions[1].use_in = { + MakeAccess(row_buffer, 1, 32), + MakeAccess(tile_buffer, 2, 1024), + }; + problem.graph.edges = { + MakeRawEdge(0, 1, 0, 0, 1, 32, 1), + MakeRawEdge(0, 1, 1, 1, 2, 1024, 2), + }; + + SunmmioTileLoopFusionWindowPlan plan = PlanSingleProblem(problem); + + EXPECT_EQ(LeafRegionOrder(plan.tree), std::vector({0, 1})); + ASSERT_EQ(plan.tree.size(), 1U); + const SunmmioTileLoopFusionPlannerTreeNode &outer_scope = plan.tree[0]; + ASSERT_TRUE(outer_scope.is_scope); + EXPECT_EQ(outer_scope.shell_axes, std::vector({"i"})); + ASSERT_EQ(outer_scope.children.size(), 2U); + + EXPECT_FALSE(outer_scope.children[0].is_scope); + EXPECT_EQ(outer_scope.children[0].region_index, 0); + EXPECT_FALSE(outer_scope.children[1].is_scope); + EXPECT_EQ(outer_scope.children[1].region_index, 1); + EXPECT_GT(plan.score.write_cut_cost, 0); +} + +TEST(SunmmioTileLoopFusionPlannerTest, + UnitTripInnerShellRemainsFusedAcrossShallowDependence) { + Buffer row_buffer = MakeSharedBuffer("row_buffer", {I(8)}); + Buffer tile_buffer = MakeSharedBuffer("tile_buffer", {I(8), I(1)}); + + SunmmioTileLoopFusionWindowProblem problem; + problem.regions = { + MakePlannerRegion(0, {"i", "j"}, MakeExtents({8, 1})), + MakePlannerRegion(1, {"i", "j"}, MakeExtents({8, 1})), + }; + problem.normalized_regions.resize(2); + problem.normalized_regions[0].def_out = { + MakeAccess(row_buffer, 1, 32), + MakeAccess(tile_buffer, 2, 1024), + }; + problem.normalized_regions[1].use_in = { + MakeAccess(row_buffer, 1, 32), + MakeAccess(tile_buffer, 2, 1024), + }; + problem.graph.edges = { + MakeRawEdge(0, 1, 0, 0, 1, 32, 1), + MakeRawEdge(0, 1, 1, 1, 2, 1024, 2), + }; + + SunmmioTileLoopFusionWindowPlan plan = PlanSingleProblem(problem); + + ASSERT_EQ(plan.tree.size(), 1U); + const SunmmioTileLoopFusionPlannerTreeNode &outer_scope = plan.tree[0]; + ASSERT_TRUE(outer_scope.is_scope); + const SunmmioTileLoopFusionPlannerTreeNode &inner_scope = + ExpectSingleScopeChild(outer_scope); + EXPECT_EQ(inner_scope.shell_axes, std::vector({"i", "j"})); + ASSERT_EQ(inner_scope.children.size(), 2U); + EXPECT_EQ(inner_scope.children[0].region_index, 0); + EXPECT_EQ(inner_scope.children[1].region_index, 1); +} + TEST(SunmmioTileLoopFusionPlannerTest, LaterTileConsumerCanReorderAheadOfEarlierRowConsumer) { Buffer row_buffer = MakeSharedBuffer("row_buffer", {I(8)}); diff --git a/testing/python/sunmmio/codegen/test_aligned_row_dma_copy.py b/testing/python/sunmmio/codegen/test_aligned_row_dma_copy.py new file mode 100644 index 0000000000..33e75fc3b2 --- /dev/null +++ b/testing/python/sunmmio/codegen/test_aligned_row_dma_copy.py @@ -0,0 +1,494 @@ +import pytest +import tilelang +import tilelang.language as T +from tilelang import tvm +from tilelang.layout import make_aligned_row_major + +from testing.python.sunmmio.common.codegen_validation import ( + lower_sunmmio_kernel_to_device_tir, + validate_sunmmio_codegen_with_npuir_opt, +) +from testing.python.sunmmio.common.compile_pipeline import target + + +tilelang.env.disable_cache() + + +@target("Sunmmio") +def aligned_row_vector_copy_kernel(direction="load", rsram_rank=1, cols=64, dtype=T.bfloat16): + dram_shape = (3, cols) + rsram_shape = (cols,) if rsram_rank == 1 else (1, cols) + dram_layout = make_aligned_row_major(dram_shape, dtype, align_bytes=1024) + rsram_layout = make_aligned_row_major(rsram_shape, dtype, align_bytes=1024) + + if direction == "load": + + @T.prim_func + def main( + src: T.MeshTensor(dram_shape, T.placement.replicated(), dtype, layout=dram_layout), # type: ignore + ): + with T.Kernel(): + dst = T.alloc_shared(rsram_shape, dtype, scope="shared.rsram") + T.annotate_layout({dst: rsram_layout}) + T.copy(src[1:2, :], dst) + + return main + + @T.prim_func + def main( + dst: T.MeshTensor(dram_shape, T.placement.replicated(), dtype, layout=dram_layout), # type: ignore + ): + with T.Kernel(): + src = T.alloc_shared(rsram_shape, dtype, scope="shared.rsram") + T.annotate_layout({src: rsram_layout}) + T.copy(src, dst[1:2, :]) + + return main + + +@target("Sunmmio") +def aligned_row_matrix_copy_kernel( + direction="load", + rows=500, + cols=500, + row_start=0, + copy_rows=None, + dtype=T.bfloat16, +): + copy_rows = rows if copy_rows is None else copy_rows + dram_shape = (5, rows, cols) + rsram_shape = (copy_rows, cols) + dram_layout = make_aligned_row_major(dram_shape, dtype, align_bytes=1024) + rsram_layout = make_aligned_row_major(rsram_shape, dtype, align_bytes=1024) + + if direction == "load": + + @T.prim_func + def main( + src: T.MeshTensor(dram_shape, T.placement.replicated(), dtype, layout=dram_layout), # type: ignore + ): + with T.Kernel(): + dst = T.alloc_shared(rsram_shape, dtype, scope="shared.rsram") + T.annotate_layout({dst: rsram_layout}) + T.copy(src[2, row_start : row_start + copy_rows, :], dst) + + return main + + @T.prim_func + def main( + dst: T.MeshTensor(dram_shape, T.placement.replicated(), dtype, layout=dram_layout), # type: ignore + ): + with T.Kernel(): + src = T.alloc_shared(rsram_shape, dtype, scope="shared.rsram") + T.annotate_layout({src: rsram_layout}) + T.copy(src, dst[2, row_start : row_start + copy_rows, :]) + + return main + + +@target("Sunmmio") +def aligned_row_non_singleton_reshape_kernel(direction="load", dtype=T.bfloat16): + dram_shape, rsram_shape = (2, 32), (64,) + dram_layout = make_aligned_row_major(dram_shape, dtype, align_bytes=1024) + rsram_layout = make_aligned_row_major(rsram_shape, dtype, align_bytes=1024) + + if direction == "load": + + @T.prim_func + def main( + src: T.MeshTensor(dram_shape, T.placement.replicated(), dtype, layout=dram_layout), # type: ignore + ): + with T.Kernel(): + dst = T.alloc_shared(rsram_shape, dtype, scope="shared.rsram") + T.annotate_layout({dst: rsram_layout}) + T.copy(src[0:2, 0:32], dst[0:64]) + + return main + + @T.prim_func + def main( + dst: T.MeshTensor(dram_shape, T.placement.replicated(), dtype, layout=dram_layout), # type: ignore + ): + with T.Kernel(): + src = T.alloc_shared(rsram_shape, dtype, scope="shared.rsram") + T.annotate_layout({src: rsram_layout}) + T.copy(src[0:64], dst[0:2, 0:32]) + + return main + + +@target("Sunmmio") +def aligned_row_byte_aligned_reshape_kernel(direction="load", dtype=T.bfloat16): + dram_shape, rsram_shape = (2, 512), (1024,) + dram_layout = make_aligned_row_major(dram_shape, dtype, align_bytes=1024) + rsram_layout = make_aligned_row_major(rsram_shape, dtype, align_bytes=1024) + + if direction == "load": + + @T.prim_func + def main( + src: T.MeshTensor(dram_shape, T.placement.replicated(), dtype, layout=dram_layout), # type: ignore + ): + with T.Kernel(): + dst = T.alloc_shared(rsram_shape, dtype, scope="shared.rsram") + T.annotate_layout({dst: rsram_layout}) + T.copy(src[0:2, 0:512], dst[0:1024]) + + return main + + @T.prim_func + def main( + dst: T.MeshTensor(dram_shape, T.placement.replicated(), dtype, layout=dram_layout), # type: ignore + ): + with T.Kernel(): + src = T.alloc_shared(rsram_shape, dtype, scope="shared.rsram") + T.annotate_layout({src: rsram_layout}) + T.copy(src[0:1024], dst[0:2, 0:512]) + + return main + + +@target("Sunmmio") +def aligned_row_incompatible_singleton_kernel(direction="load", dtype=T.bfloat16): + dram_shape, rsram_shape = (64,), (64, 1) + dram_layout = make_aligned_row_major(dram_shape, dtype, align_bytes=1024) + rsram_layout = make_aligned_row_major(rsram_shape, dtype, align_bytes=1024) + + if direction == "load": + + @T.prim_func + def main( + src: T.MeshTensor(dram_shape, T.placement.replicated(), dtype, layout=dram_layout), # type: ignore + ): + with T.Kernel(): + dst = T.alloc_shared(rsram_shape, dtype, scope="shared.rsram") + T.annotate_layout({dst: rsram_layout}) + T.copy(src[0:64], dst[0:64, 0:1]) + + return main + + @T.prim_func + def main( + dst: T.MeshTensor(dram_shape, T.placement.replicated(), dtype, layout=dram_layout), # type: ignore + ): + with T.Kernel(): + src = T.alloc_shared(rsram_shape, dtype, scope="shared.rsram") + T.annotate_layout({src: rsram_layout}) + T.copy(src[0:64, 0:1], dst[0:64]) + + return main + + +@target("Sunmmio") +def aligned_row_middle_singleton_kernel(direction="load", dtype=T.bfloat16): + dram_shape, rsram_shape = (500, 1, 500), (500, 500) + dram_layout = make_aligned_row_major(dram_shape, dtype, align_bytes=1024) + rsram_layout = make_aligned_row_major(rsram_shape, dtype, align_bytes=1024) + + if direction == "load": + + @T.prim_func + def main( + src: T.MeshTensor(dram_shape, T.placement.replicated(), dtype, layout=dram_layout), # type: ignore + ): + with T.Kernel(): + dst = T.alloc_shared(rsram_shape, dtype, scope="shared.rsram") + T.annotate_layout({dst: rsram_layout}) + T.copy(src[0:500, 0:1, 0:500], dst[0:500, 0:500]) + + return main + + @T.prim_func + def main( + dst: T.MeshTensor(dram_shape, T.placement.replicated(), dtype, layout=dram_layout), # type: ignore + ): + with T.Kernel(): + src = T.alloc_shared(rsram_shape, dtype, scope="shared.rsram") + T.annotate_layout({src: rsram_layout}) + T.copy(src[0:500, 0:500], dst[0:500, 0:1, 0:500]) + + return main + + +@target("Sunmmio") +def aligned_row_partial_row_kernel(direction="load", dtype=T.bfloat16): + dram_shape, rsram_shape = (3, 64), (32,) + dram_layout = make_aligned_row_major(dram_shape, dtype, align_bytes=1024) + rsram_layout = make_aligned_row_major(rsram_shape, dtype, align_bytes=1024) + + if direction == "load": + + @T.prim_func + def main( + src: T.MeshTensor(dram_shape, T.placement.replicated(), dtype, layout=dram_layout), # type: ignore + ): + with T.Kernel(): + dst = T.alloc_shared(rsram_shape, dtype, scope="shared.rsram") + T.annotate_layout({dst: rsram_layout}) + T.copy(src[1, 0:32], dst) + + return main + + @T.prim_func + def main( + dst: T.MeshTensor(dram_shape, T.placement.replicated(), dtype, layout=dram_layout), # type: ignore + ): + with T.Kernel(): + src = T.alloc_shared(rsram_shape, dtype, scope="shared.rsram") + T.annotate_layout({src: rsram_layout}) + T.copy(src, dst[1, 0:32]) + + return main + + +@target("Sunmmio") +def aligned_row_byte_aligned_partial_row_kernel(direction="load", dtype=T.bfloat16): + dram_shape, rsram_shape = (3, 1024), (512,) + dram_layout = make_aligned_row_major(dram_shape, dtype, align_bytes=1024) + rsram_layout = make_aligned_row_major(rsram_shape, dtype, align_bytes=1024) + + if direction == "load": + + @T.prim_func + def main( + src: T.MeshTensor(dram_shape, T.placement.replicated(), dtype, layout=dram_layout), # type: ignore + ): + with T.Kernel(): + dst = T.alloc_shared(rsram_shape, dtype, scope="shared.rsram") + T.annotate_layout({dst: rsram_layout}) + T.copy(src[1, 0:512], dst) + + return main + + @T.prim_func + def main( + dst: T.MeshTensor(dram_shape, T.placement.replicated(), dtype, layout=dram_layout), # type: ignore + ): + with T.Kernel(): + src = T.alloc_shared(rsram_shape, dtype, scope="shared.rsram") + T.annotate_layout({src: rsram_layout}) + T.copy(src, dst[1, 0:512]) + + return main + + +@target("Sunmmio") +def aligned_row_effective_rank3_kernel(direction="load", dtype=T.bfloat16): + shape = (2, 3, 64) + dram_layout = make_aligned_row_major(shape, dtype, align_bytes=1024) + rsram_layout = make_aligned_row_major(shape, dtype, align_bytes=1024) + + if direction == "load": + + @T.prim_func + def main( + src: T.MeshTensor(shape, T.placement.replicated(), dtype, layout=dram_layout), # type: ignore + ): + with T.Kernel(): + dst = T.alloc_shared(shape, dtype, scope="shared.rsram") + T.annotate_layout({dst: rsram_layout}) + T.copy(src, dst) + + return main + + @T.prim_func + def main( + dst: T.MeshTensor(shape, T.placement.replicated(), dtype, layout=dram_layout), # type: ignore + ): + with T.Kernel(): + src = T.alloc_shared(shape, dtype, scope="shared.rsram") + T.annotate_layout({src: rsram_layout}) + T.copy(src, dst) + + return main + + +@target("Sunmmio") +def aligned_row_alignment_mismatch_kernel(direction="load", dtype=T.bfloat16): + dram_shape, rsram_shape = (3, 64), (2, 64) + dram_layout = make_aligned_row_major(dram_shape, dtype, align_bytes=1024) + rsram_layout = make_aligned_row_major(rsram_shape, dtype, align_bytes=64) + + if direction == "load": + + @T.prim_func + def main( + src: T.MeshTensor(dram_shape, T.placement.replicated(), dtype, layout=dram_layout), # type: ignore + ): + with T.Kernel(): + dst = T.alloc_shared(rsram_shape, dtype, scope="shared.rsram") + T.annotate_layout({dst: rsram_layout}) + T.copy(src[0:2, :], dst) + + return main + + @T.prim_func + def main( + dst: T.MeshTensor(dram_shape, T.placement.replicated(), dtype, layout=dram_layout), # type: ignore + ): + with T.Kernel(): + src = T.alloc_shared(rsram_shape, dtype, scope="shared.rsram") + T.annotate_layout({src: rsram_layout}) + T.copy(src, dst[0:2, :]) + + return main + + +@pytest.mark.parametrize("direction", ["load", "store"]) +@pytest.mark.parametrize("rsram_rank", [1, 2]) +def test_aligned_row_vector_copy_uses_1024_byte_carrier(direction, rsram_rank, tmp_path): + src = validate_sunmmio_codegen_with_npuir_opt( + aligned_row_vector_copy_kernel(direction=direction, rsram_rank=rsram_rank), + tmp_path, + mlir_filename=f"aligned_row_vector_{direction}_rank{rsram_rank}.mlir", + expected_tokens=("suvm.copy_async", "!suvm.tile_view<512xbf16>"), + opt_args=("--verify-each", "--suvm-to-llvm-pipeline"), + ) + assert src.count("suvm.copy_async") == 1 + assert "suvm.transform_layout_async" not in src + + +@pytest.mark.parametrize( + "direction,rows,cols,row_start,copy_rows,dtype,expected_view", + [ + pytest.param( + direction, + 500, + 500, + 0, + None, + T.bfloat16, + "!suvm.tile_view<500x512xbf16>", + id=f"full-bf16-{direction}", + ) + for direction in ("load", "store") + ] + + [ + pytest.param( + direction, + 20, + 500, + 7, + 5, + T.bfloat16, + "!suvm.tile_view<5x512xbf16>", + id=f"row-subset-bf16-{direction}", + ) + for direction in ("load", "store") + ] + + [ + pytest.param( + direction, + 3, + 250, + 0, + None, + T.float32, + "!suvm.tile_view<3x256xf32>", + id=f"full-fp32-{direction}", + ) + for direction in ("load", "store") + ] + + [ + pytest.param( + direction, + 3, + 512, + 0, + None, + T.bfloat16, + "!suvm.tile_view<3x512xbf16>", + id=f"already-aligned-{direction}", + ) + for direction in ("load", "store") + ], +) +def test_aligned_row_matrix_copy_uses_rank2_dma( + direction, + rows, + cols, + row_start, + copy_rows, + dtype, + expected_view, + tmp_path, +): + src = validate_sunmmio_codegen_with_npuir_opt( + aligned_row_matrix_copy_kernel( + direction=direction, + rows=rows, + cols=cols, + row_start=row_start, + copy_rows=copy_rows, + dtype=dtype, + ), + tmp_path, + mlir_filename=f"aligned_row_matrix_{direction}_{rows}_{cols}_{row_start}_{copy_rows}.mlir", + expected_tokens=("suvm.copy_async", expected_view), + opt_args=("--verify-each", "--suvm-to-llvm-pipeline"), + ) + assert src.count("suvm.copy_async") == 1 + assert "suvm.transform_layout_async" not in src + + +@pytest.mark.parametrize("direction", ["load", "store"]) +def test_alignment_mismatch_falls_back_to_existing_layout_transform_path(direction): + mod = lower_sunmmio_kernel_to_device_tir(aligned_row_alignment_mismatch_kernel(direction=direction)) + src = mod.script() + assert src.count("T.dma_copy") == 1 + assert src.count("T.sunmmio_layout_transform") == 1 + + +@pytest.mark.parametrize( + "factory,reason", + [ + pytest.param( + aligned_row_non_singleton_reshape_kernel, + "canonical logical shapes do not match", + id="non-singleton-reshape", + ), + pytest.param( + aligned_row_byte_aligned_reshape_kernel, + "canonical logical shapes do not match", + id="byte-aligned-reshape", + ), + pytest.param( + aligned_row_incompatible_singleton_kernel, + "canonical logical shapes do not match", + id="incompatible-singleton", + ), + pytest.param( + aligned_row_middle_singleton_kernel, + "effective rank exceeds two", + id="middle-singleton", + ), + pytest.param( + aligned_row_partial_row_kernel, + "innermost range must cover the complete logical row", + id="partial-row", + ), + pytest.param( + aligned_row_byte_aligned_partial_row_kernel, + "innermost range must cover the complete logical row", + id="byte-aligned-partial-row", + ), + pytest.param( + aligned_row_effective_rank3_kernel, + "effective rank exceeds two", + id="effective-rank3", + ), + ], +) +@pytest.mark.parametrize("direction", ["load", "store"]) +def test_aligned_row_carrier_rejection_is_actionable(factory, reason, direction, tmp_path): + with pytest.raises(tvm.error.InternalError, match=reason): + validate_sunmmio_codegen_with_npuir_opt( + factory(direction=direction), + tmp_path, + mlir_filename=f"aligned_row_reject_{factory.__name__}_{direction}.mlir", + opt_args=("--verify-each", "--suvm-to-llvm-pipeline"), + ) + + +if __name__ == "__main__": + tilelang.testing.main() diff --git a/testing/python/sunmmio/codegen/test_alloc_var_opt_validate.py b/testing/python/sunmmio/codegen/test_alloc_var_opt_validate.py index c55f92270a..9ad4e217c4 100644 --- a/testing/python/sunmmio/codegen/test_alloc_var_opt_validate.py +++ b/testing/python/sunmmio/codegen/test_alloc_var_opt_validate.py @@ -177,7 +177,7 @@ def alloc_var_conditional_loop_extent_copy_kernel( block_N=32, dtype=T.bfloat16, ): - shard_policy = T.MeshShardingPolicy(y=0, x=1) + shard_policy = T.placement.full_shard(0, 1) A_layout = make_zz_layout((M, N), [0, 1], (32, 32)) B_layout = make_zz_layout((M, N), [0, 1], (32, 32)) @@ -260,7 +260,7 @@ def test_alloc_var_copy_mma_control_flow_kernel_codegen_validates_with_npuir_opt "scf.while", "scf.condition", "-> (!suvm.token, !suvm.token, f32, i32)", - "-> (!suvm.token, f32, i32)", + "-> (!suvm.token, f32, i32, i1)", "!suvm.token", "suvm.copy_async", "suvm.tc.mma", diff --git a/testing/python/sunmmio/codegen/test_allocate_copy_expr_opt_validate.py b/testing/python/sunmmio/codegen/test_allocate_copy_expr_opt_validate.py index b491b9929e..9a325f6a45 100644 --- a/testing/python/sunmmio/codegen/test_allocate_copy_expr_opt_validate.py +++ b/testing/python/sunmmio/codegen/test_allocate_copy_expr_opt_validate.py @@ -48,8 +48,8 @@ def main( with T.Kernel() as _cid: sharded_M, sharded_K = A.local_shape sharded_N = B.local_shape[1] - A_shared_dist = T.alloc_shared((block_M, block_K * T.mesh_ncols()), dtype) - B_shared_dist = T.alloc_shared((block_K * T.mesh_nrows(), block_N), dtype) + A_shared_dist = T.alloc_shared((block_M, block_K * T.ncols()), dtype) + B_shared_dist = T.alloc_shared((block_K * T.nrows(), block_N), dtype) C_shared = T.alloc_shared((block_M, block_N), accum_dtype) for by in T.serial(T.ceildiv(sharded_M, block_M)): @@ -149,8 +149,8 @@ def main( with T.Kernel() as _cid: sharded_M, sharded_K = A.local_shape sharded_N = B.local_shape[1] - A_shared_dist = T.alloc_shared((block_M, block_K * T.mesh_ncols()), dtype) - B_shared_dist = T.alloc_shared((block_K * T.mesh_nrows(), block_N), dtype) + A_shared_dist = T.alloc_shared((block_M, block_K * T.ncols()), dtype) + B_shared_dist = T.alloc_shared((block_K * T.nrows(), block_N), dtype) C_shared = T.alloc_shared((block_M, block_N), accum_dtype) for by in T.serial(T.ceildiv(sharded_M, block_M)): @@ -365,7 +365,7 @@ def test_pipelined_allocate_copy_mma_codegen_propagates_ping_pong(tmp_path): "#suvm.memory_space", "#suvm.memory_space", "suvm.mcast_tok", - "suvm.ping_pong = #suvm.ping_pong", + "suvm.ping_pong = #suvm.ping_pong", ), ) diff --git a/testing/python/sunmmio/codegen/test_base_dynamic_opt_validate.py b/testing/python/sunmmio/codegen/test_base_dynamic_opt_validate.py index e71963eb65..c39f1849af 100644 --- a/testing/python/sunmmio/codegen/test_base_dynamic_opt_validate.py +++ b/testing/python/sunmmio/codegen/test_base_dynamic_opt_validate.py @@ -4,7 +4,6 @@ import tilelang.language as T import tilelang.testing import tvm_ffi -from tilelang.carver.arch import driver from tilelang.layout import ( make_nzz_layout, make_zn_layout, @@ -77,7 +76,7 @@ def dynamic_allocate_copy_mma_kernel( dtype=T.bfloat16, accum_dtype=T.float32, ): - nrows, ncols = driver.get_sunmmio_device_mesh_config() + nrows, ncols = T.nrows(), T.ncols() M_sharded = T.dynamic("m") N_sharded = T.dynamic("n") @@ -101,8 +100,8 @@ def main( with T.Kernel() as _cid: sharded_M, sharded_K = A.local_shape sharded_N = B.local_shape[1] - A_shared_dist = T.alloc_shared((block_M, block_K * T.mesh_ncols()), dtype) - B_shared_dist = T.alloc_shared((block_K * T.mesh_nrows(), block_N), dtype) + A_shared_dist = T.alloc_shared((block_M, block_K * T.ncols()), dtype) + B_shared_dist = T.alloc_shared((block_K * T.nrows(), block_N), dtype) C_shared = T.alloc_shared((block_M, block_N), accum_dtype) for by in T.serial(T.ceildiv(sharded_M, block_M)): diff --git a/testing/python/sunmmio/codegen/test_copy_region_validate.py b/testing/python/sunmmio/codegen/test_copy_region_validate.py index b79b70c169..961e3c6cae 100644 --- a/testing/python/sunmmio/codegen/test_copy_region_validate.py +++ b/testing/python/sunmmio/codegen/test_copy_region_validate.py @@ -141,7 +141,7 @@ def main( @target("Sunmmio") def let_bound_mesh_local_shape_copy_kernel(): global_shape = (256, 256) - shard_policy = T.MeshShardingPolicy(y=0, x=1) + shard_policy = T.placement.full_shard(0, 1) tensor_layout = make_zz_layout(global_shape, axes=[0, 1], block_shape=(32, 32)) @T.prim_func diff --git a/testing/python/sunmmio/codegen/test_hybrid_tail_predicate_mask_opt_validate.py b/testing/python/sunmmio/codegen/test_hybrid_tail_predicate_mask_opt_validate.py index 4de39af6c8..e851be77c2 100644 --- a/testing/python/sunmmio/codegen/test_hybrid_tail_predicate_mask_opt_validate.py +++ b/testing/python/sunmmio/codegen/test_hybrid_tail_predicate_mask_opt_validate.py @@ -45,8 +45,8 @@ def hybrid_tail_predicate_mask_kernel( assert mix_pad % 32 == 0 assert comb_pad % 32 == 0 - token_policy = T.MeshShardingPolicy(cross_mesh_dim=0) - replicated_policy = T.MeshShardingPolicy() + token_policy = T.placement.mesh_as_line(0) + replicated_policy = T.placement.replicated() gemm_mul_shape = (num_tokens, mix_pad) sqrsum_shape = (num_tokens,) diff --git a/testing/python/sunmmio/codegen/test_mask_index_dtype_opt_validate.py b/testing/python/sunmmio/codegen/test_mask_index_dtype_opt_validate.py index 3a67e05ea6..b456e16794 100644 --- a/testing/python/sunmmio/codegen/test_mask_index_dtype_opt_validate.py +++ b/testing/python/sunmmio/codegen/test_mask_index_dtype_opt_validate.py @@ -79,7 +79,7 @@ def main( @target("Sunmmio") def bf16_dynamic_rect_tail_mask_kernel(block_m=32, block_n=32, valid_rows=17, valid_cols=19): dtype = T.bfloat16 - shard_policy = T.MeshShardingPolicy(y=0, x=1) + shard_policy = T.placement.full_shard(0, 1) tensor_shape = (block_m, block_n) tensor_layout = make_zz_layout(tensor_shape, [0, 1], (block_m, block_n)) diff --git a/testing/python/sunmmio/codegen/test_mx_gemm.py b/testing/python/sunmmio/codegen/test_mx_gemm.py index 2744d91b62..e8c87a262c 100644 --- a/testing/python/sunmmio/codegen/test_mx_gemm.py +++ b/testing/python/sunmmio/codegen/test_mx_gemm.py @@ -79,8 +79,8 @@ def main( sharded_M, sharded_K = A.local_shape _, sharded_N = B.local_shape - A_shared_dist = T.alloc_shared((block_M, block_K * T.mesh_ncols()), a_dtype) - B_shared_dist = T.alloc_shared((block_K * T.mesh_nrows(), block_N), b_dtype) + A_shared_dist = T.alloc_shared((block_M, block_K * T.ncols()), a_dtype) + B_shared_dist = T.alloc_shared((block_K * T.nrows(), block_N), b_dtype) C_shared = T.alloc_shared((block_M, block_N), accum_dtype) for bx in T.serial(T.ceildiv(sharded_M, block_M)): diff --git a/testing/python/sunmmio/codegen/test_mx_quant_dequant_opt_validate.py b/testing/python/sunmmio/codegen/test_mx_quant_dequant_opt_validate.py index 4c74f502b0..eaac8f36ac 100644 --- a/testing/python/sunmmio/codegen/test_mx_quant_dequant_opt_validate.py +++ b/testing/python/sunmmio/codegen/test_mx_quant_dequant_opt_validate.py @@ -4,8 +4,6 @@ import tilelang import tilelang.language as T import tilelang.testing -from tilelang.language.mesh_tensor import MeshReplicationType -from tilelang.tileview import make_tileview from tilelang.layout import ( get_mx_scale_shape, # make_aligned_row_major, @@ -82,7 +80,7 @@ def _e8m0_scale_to_fp32(x): @target("Sunmmio") def mx_ocp_quant_dequant_full_chain_kernel_for_debug(mx_dtype, data_dtype, data_dtype_name, data_max): shape = (32, 32) - shard_policy = T.MeshShardingPolicy(replicate=MeshReplicationType.ALL) + shard_policy = T.placement.replicated() tensor_layout = make_zz_layout(shape, axes=[0, 1], block_shape=(32, 32)) mx_layout = make_mxzz_layout(shape, dtype=mx_dtype) scale_shape = _int_shape(get_mx_scale_shape(mx_layout, mx_dtype)) @@ -94,60 +92,42 @@ def main( Y: T.MeshTensor(shape, shard_policy, T.bfloat16, layout=tensor_layout), # type: ignore ): with T.Kernel(): - a_rsram = T.alloc_shared(shape, T.bfloat16) - amax = T.alloc_shared((64,), T.float32) + a = T.alloc_shared(shape, T.bfloat16) + amax = T.alloc_shared((shape[0],), T.float32) data = T.alloc_shared(shape, data_dtype) scale = T.alloc_shared(scale_shape, T.float8_e8m0fnu) - scale_bf16_vec = T.alloc_shared((64,), T.bfloat16) - scale_fp32_vec = T.alloc_shared((64,), T.float32) mx = T.alloc_shared(shape, mx_dtype) unpacked_data = T.alloc_shared(shape, data_dtype) unpacked_scale = T.alloc_shared(scale_shape, T.float8_e8m0fnu) - unpacked_scale_fp32_vec = T.alloc_shared((64,), T.float32) - y_rsram = T.alloc_shared(shape, T.bfloat16) + y = T.alloc_shared(shape, T.bfloat16) - T.copy(A, a_rsram) + T.copy(A, a) # OCP MX quantization: each logical row is one 32-element scale # group in this 32x32 test block. - T.fill(amax, T.float32(1e-4)) - T.reduce_absmax(a_rsram, amax[0:32], dim=1, clear=True) - for j in T.Tiles([scale_shape[1]]): - safe_amax = T.max(amax[j], T.float32(1e-4)) - scale_bf16_vec[j] = T.Cast("bfloat16", safe_amax * T.float32(data_max_inv)) - for j in T.Tiles([scale_shape[1]]): - scale[0, j] = _to_e8m0_scale(scale_bf16_vec[j]) - for j in T.Tiles([scale_shape[1]]): - scale_fp32_vec[j] = _e8m0_scale_to_fp32(scale[0, j]) + T.reduce_absmax(a, amax, dim=1, clear=True) + for row in T.Tiles(amax): + safe_amax = T.max(amax[row], T.float32(1e-4)) + bf16_scale = T.Cast("bfloat16", safe_amax * T.float32(data_max_inv)) + scale[0, row] = _to_e8m0_scale(bf16_scale) - T.annotate_tileview( - { - a_rsram: make_tileview(a_rsram, (8, 32), (-2, -1)), - data: make_tileview(data, (8, 32), (-2, -1)), - } - ) for row, col in T.Tiles(data): - raw = T.Cast("float32", a_rsram[row, col]) / scale_fp32_vec[row] - clamped = T.min(T.max(raw, T.float32(-data_max)), T.float32(data_max)) - data[row, col] = T.Cast(data_dtype_name, clamped) + scale_fp32 = _e8m0_scale_to_fp32(scale[0, row]) + value = T.Cast("float32", a[row, col]) / scale_fp32 + data[row, col] = T.Cast( + data_dtype_name, + T.clamp(value, T.float32(-data_max), T.float32(data_max)), + ) T.mx_pack(data, scale, mx) T.mx_unpack(mx, unpacked_data, unpacked_scale) - for j in T.Tiles([scale_shape[1]]): - unpacked_scale_fp32_vec[j] = _e8m0_scale_to_fp32(unpacked_scale[0, j]) + for row, col in T.Tiles(y): + value = T.Cast("float32", unpacked_data[row, col]) + scale_fp32 = _e8m0_scale_to_fp32(unpacked_scale[0, row]) + y[row, col] = T.Cast("bfloat16", value * scale_fp32) - T.annotate_tileview( - { - unpacked_data: make_tileview(unpacked_data, (8, 32), (-2, -1)), - y_rsram: make_tileview(y_rsram, (8, 32), (-2, -1)), - } - ) - for row, col in T.Tiles(y_rsram): - q = T.Cast("float32", unpacked_data[row, col]) - y_rsram[row, col] = T.Cast("bfloat16", q * unpacked_scale_fp32_vec[row]) - - T.copy(y_rsram, Y) + T.copy(y, Y) return main @@ -155,7 +135,7 @@ def main( @target("Sunmmio") def mx_ocp_quant_dequant_full_chain_kernel_original(mx_dtype, data_dtype, data_dtype_name, data_max): shape = (32, 32) - shard_policy = T.MeshShardingPolicy(replicate=MeshReplicationType.ALL) + shard_policy = T.placement.replicated() tensor_layout = make_zz_layout(shape, axes=[0, 1], block_shape=(32, 32)) mx_layout = make_mxzz_layout(shape, dtype=mx_dtype) scale_shape = _int_shape(get_mx_scale_shape(mx_layout, mx_dtype)) @@ -218,7 +198,7 @@ def main( @target("Sunmmio") def mx_ocp_quant_kernel_for_debug(mx_dtype, data_dtype, data_dtype_name, data_max): shape = MX_COPY_ALIGNED_SHAPE - shard_policy = T.MeshShardingPolicy(replicate=MeshReplicationType.ALL) + shard_policy = T.placement.replicated() tensor_layout = make_zz_layout(shape, axes=[0, 1], block_shape=(32, 32)) mx_layout = make_mxzz_layout(shape, dtype=mx_dtype) scale_shape = _int_shape(get_mx_scale_shape(mx_layout, mx_dtype)) @@ -234,44 +214,25 @@ def main( amax = T.alloc_shared((shape[0],), T.float32) data = T.alloc_shared(shape, data_dtype) scale = T.alloc_shared(scale_shape, T.float8_e8m0fnu) - scale_bf16_vec = T.alloc_shared((scale_shape[0] * 64,), T.bfloat16) - scale_fp32_vec = T.alloc_shared((scale_shape[0] * 64,), T.float32) mx = T.alloc_shared(shape, mx_dtype) T.copy(A, a_rsram) - T.fill(amax, T.float32(1e-4)) T.reduce_absmax(a_rsram, amax, dim=1, clear=True) - for block in T.serial(scale_shape[0]): - amax_base = block * scale_shape[1] - scale_base = block * 64 - for j in T.Tiles([scale_shape[1]]): - safe_amax = T.max(amax[amax_base + j], T.float32(1e-4)) - scale_bf16_vec[scale_base + j] = T.Cast("bfloat16", safe_amax * T.float32(data_max_inv)) - - for block in T.serial(scale_shape[0]): - scale_base = block * 64 - for j in T.Tiles([scale_shape[1]]): - scale[block, j] = _to_e8m0_scale(scale_bf16_vec[scale_base + j]) - - for block in T.serial(scale_shape[0]): - scale_base = block * 64 - for j in T.Tiles([scale_shape[1]]): - scale_fp32_vec[scale_base + j] = _e8m0_scale_to_fp32(scale[block, j]) - - T.annotate_tileview( - { - a_rsram: make_tileview(a_rsram, (8, 32), (-2, -1)), - data: make_tileview(data, (8, 32), (-2, -1)), - } - ) for block in T.serial(scale_shape[0]): row_base = block * scale_shape[1] - scale_base = block * 64 - for row, col in T.Tiles([32, 32]): - raw = T.Cast("float32", a_rsram[row_base + row, col]) / scale_fp32_vec[scale_base + row] - clamped = T.min(T.max(raw, T.float32(-data_max)), T.float32(data_max)) - data[row_base + row, col] = T.Cast(data_dtype_name, clamped) + for row in T.Tiles([scale_shape[1]]): + safe_amax = T.max(amax[row_base + row], T.float32(1e-4)) + bf16_scale = T.Cast("bfloat16", safe_amax * T.float32(data_max_inv)) + scale[block, row] = _to_e8m0_scale(bf16_scale) + + for row, col in T.Tiles([scale_shape[1], shape[1]]): + scale_fp32 = _e8m0_scale_to_fp32(scale[block, row]) + value = T.Cast("float32", a_rsram[row_base + row, col]) / scale_fp32 + data[row_base + row, col] = T.Cast( + data_dtype_name, + T.clamp(value, T.float32(-data_max), T.float32(data_max)), + ) T.mx_pack(data, scale, mx) T.copy(mx, MX) @@ -282,7 +243,7 @@ def main( @target("Sunmmio") def mx_ocp_dequant_kernel_for_debug(mx_dtype, data_dtype, data_dtype_name, data_max): shape = MX_COPY_ALIGNED_SHAPE - shard_policy = T.MeshShardingPolicy(replicate=MeshReplicationType.ALL) + shard_policy = T.placement.replicated() tensor_layout = make_zz_layout(shape, axes=[0, 1], block_shape=(32, 32)) mx_layout = make_mxzz_layout(shape, dtype=mx_dtype) scale_shape = _int_shape(get_mx_scale_shape(mx_layout, mx_dtype)) @@ -296,29 +257,17 @@ def main( mx = T.alloc_shared(shape, mx_dtype) data = T.alloc_shared(shape, data_dtype) scale = T.alloc_shared(scale_shape, T.float8_e8m0fnu) - scale_fp32_vec = T.alloc_shared((scale_shape[0] * 64,), T.float32) y_rsram = T.alloc_shared(shape, T.bfloat16) T.copy(MX, mx) T.mx_unpack(mx, data, scale) - for block in T.serial(scale_shape[0]): - scale_base = block * 64 - for j in T.Tiles([scale_shape[1]]): - scale_fp32_vec[scale_base + j] = _e8m0_scale_to_fp32(scale[block, j]) - - T.annotate_tileview( - { - data: make_tileview(data, (8, 32), (-2, -1)), - y_rsram: make_tileview(y_rsram, (8, 32), (-2, -1)), - } - ) for block in T.serial(scale_shape[0]): row_base = block * scale_shape[1] - scale_base = block * 64 - for row, col in T.Tiles([32, 32]): - q = T.Cast("float32", data[row_base + row, col]) - y_rsram[row_base + row, col] = T.Cast("bfloat16", q * scale_fp32_vec[scale_base + row]) + for row, col in T.Tiles([scale_shape[1], shape[1]]): + value = T.Cast("float32", data[row_base + row, col]) + scale_fp32 = _e8m0_scale_to_fp32(scale[block, row]) + y_rsram[row_base + row, col] = T.Cast("bfloat16", value * scale_fp32) T.copy(y_rsram, Y) @@ -330,7 +279,7 @@ def mx_ocp_quant_generic_shape_kernel_for_debug(mx_dtype, data_dtype, data_dtype shape = MX_GENERIC_SHAPE num_m_blocks = shape[0] // 32 num_n_blocks = shape[1] // 32 - shard_policy = T.MeshShardingPolicy(replicate=MeshReplicationType.ALL) + shard_policy = T.placement.replicated() tensor_layout = make_zz_layout(shape, axes=[0, 1], block_shape=(32, 32)) mx_layout = make_mxzz_layout(shape, dtype=mx_dtype) scale_shape = _int_shape(get_mx_scale_shape(mx_layout, mx_dtype)) @@ -346,49 +295,101 @@ def main( ): with T.Kernel(): a_rsram = T.alloc_shared(shape, T.bfloat16) - tile_rsram = T.alloc_shared((32, 32), T.bfloat16) - amax = T.alloc_shared((64,), T.float32) + amax = T.alloc_shared((32,), T.float32) data = T.alloc_shared(shape, data_dtype) scale = T.alloc_shared(scale_shape, T.float8_e8m0fnu) - scale_bf16_vec = T.alloc_shared((scale_shape[0] * 64,), T.bfloat16) - scale_fp32_vec = T.alloc_shared((scale_shape[0] * 64,), T.float32) mx = T.alloc_shared(shape, mx_dtype) T.copy(A, a_rsram) - T.annotate_tileview( - { - a_rsram: make_tileview(a_rsram, (8, 32), (-2, -1)), - tile_rsram: make_tileview(tile_rsram, (8, 32), (-2, -1)), - data: make_tileview(data, (8, 32), (-2, -1)), - } - ) for block_m in T.serial(num_m_blocks): for block_n in T.serial(num_n_blocks): block = block_m * num_n_blocks + block_n row_base = block_m * 32 col_base = block_n * 32 - scale_base = block * 64 - - for row, col in T.Tiles([32, 32]): - tile_rsram[row, col] = a_rsram[row_base + row, col_base + col] - T.fill(amax, T.float32(1e-4)) - T.reduce_absmax(tile_rsram, amax[0:32], dim=1, clear=True) + T.reduce_absmax( + a_rsram[row_base : row_base + 32, col_base : col_base + 32], + amax, + dim=1, + clear=True, + ) for row in T.Tiles([scale_shape[1]]): safe_amax = T.max(amax[row], T.float32(1e-4)) - scale_bf16_vec[scale_base + row] = T.Cast("bfloat16", safe_amax * T.float32(data_max_inv)) + bf16_scale = T.Cast("bfloat16", safe_amax * T.float32(data_max_inv)) + scale[block, row] = _to_e8m0_scale(bf16_scale) - for row in T.Tiles([scale_shape[1]]): - scale[block, row] = _to_e8m0_scale(scale_bf16_vec[scale_base + row]) + for row, col in T.Tiles([32, 32]): + scale_fp32 = _e8m0_scale_to_fp32(scale[block, row]) + value = T.Cast("float32", a_rsram[row_base + row, col_base + col]) / scale_fp32 + data[row_base + row, col_base + col] = T.Cast( + data_dtype_name, + T.clamp(value, T.float32(-data_max), T.float32(data_max)), + ) - for row in T.Tiles([scale_shape[1]]): - scale_fp32_vec[scale_base + row] = _e8m0_scale_to_fp32(scale[block, row]) + T.mx_pack(data, scale, mx) + T.copy(mx, MX) + + return main + + +@target("Sunmmio") +def mx_ocp_quant_sharded_kernel_for_debug(mx_dtype, data_dtype, data_dtype_name, data_max): + global_shape = (256, 256) + shard_policy = T.placement.full_shard(0, 1) + tensor_layout = make_zz_layout(global_shape, axes=[0, 1], block_shape=(32, 32)) + mx_layout = make_mxzz_layout(global_shape, dtype=mx_dtype) + data_max_inv = 1.0 / data_max + a_tensor = T.MeshTensor(global_shape, shard_policy, T.bfloat16, layout=tensor_layout) + mx_tensor = T.MeshTensor(global_shape, shard_policy, mx_dtype, layout=mx_layout) + local_m, local_n = a_tensor.local_shape + num_m_blocks = T.ceildiv(local_m, 32) + num_n_blocks = T.ceildiv(local_n, 32) + scale_shape = (num_m_blocks * num_n_blocks, 32) + + @T.prim_func + def main( + A: a_tensor, # type: ignore + MX: mx_tensor, # type: ignore + ): + with T.Kernel(): + # local_m, local_n = A.local_shape + # num_m_blocks = T.ceildiv(local_m, 32) + # num_n_blocks = T.ceildiv(local_n, 32) + # scale_shape = (num_m_blocks * num_n_blocks, 32) + + a_rsram = T.alloc_shared((local_m, local_n), T.bfloat16) + amax = T.alloc_shared((32,), T.float32) + data = T.alloc_shared((local_m, local_n), data_dtype) + scale = T.alloc_shared(scale_shape, T.float8_e8m0fnu) + mx = T.alloc_shared((local_m, local_n), mx_dtype) + + T.copy(A, a_rsram) + + for block_m in T.serial(num_m_blocks): + for block_n in T.serial(num_n_blocks): + block = block_m * num_n_blocks + block_n + row_base = block_m * 32 + col_base = block_n * 32 + + T.reduce_absmax( + a_rsram[row_base : row_base + 32, col_base : col_base + 32], + amax, + dim=1, + clear=True, + ) + for row in T.Tiles([32]): + safe_amax = T.max(amax[row], T.float32(1e-4)) + bf16_scale = T.Cast("bfloat16", safe_amax * T.float32(data_max_inv)) + scale[block, row] = _to_e8m0_scale(bf16_scale) for row, col in T.Tiles([32, 32]): - raw = T.Cast("float32", tile_rsram[row, col]) / scale_fp32_vec[scale_base + row] - clamped = T.min(T.max(raw, T.float32(-data_max)), T.float32(data_max)) - data[row_base + row, col_base + col] = T.Cast(data_dtype_name, clamped) + scale_fp32 = _e8m0_scale_to_fp32(scale[block, row]) + value = T.Cast("float32", a_rsram[row_base + row, col_base + col]) / scale_fp32 + data[row_base + row, col_base + col] = T.Cast( + data_dtype_name, + T.clamp(value, T.float32(-data_max), T.float32(data_max)), + ) T.mx_pack(data, scale, mx) T.copy(mx, MX) @@ -401,7 +402,7 @@ def mx_ocp_dequant_generic_shape_kernel_for_debug(mx_dtype, data_dtype, data_dty shape = MX_GENERIC_SHAPE num_m_blocks = shape[0] // 32 num_n_blocks = shape[1] // 32 - shard_policy = T.MeshShardingPolicy(replicate=MeshReplicationType.ALL) + shard_policy = T.placement.replicated() tensor_layout = make_zz_layout(shape, axes=[0, 1], block_shape=(32, 32)) mx_layout = make_mxzz_layout(shape, dtype=mx_dtype) scale_shape = _int_shape(get_mx_scale_shape(mx_layout, mx_dtype)) @@ -418,31 +419,21 @@ def main( mx = T.alloc_shared(shape, mx_dtype) data = T.alloc_shared(shape, data_dtype) scale = T.alloc_shared(scale_shape, T.float8_e8m0fnu) - scale_fp32_vec = T.alloc_shared((scale_shape[0] * 64,), T.float32) y_rsram = T.alloc_shared(shape, T.bfloat16) T.copy(MX, mx) T.mx_unpack(mx, data, scale) - T.annotate_tileview( - { - data: make_tileview(data, (8, 32), (-2, -1)), - y_rsram: make_tileview(y_rsram, (8, 32), (-2, -1)), - } - ) for block_m in T.serial(num_m_blocks): for block_n in T.serial(num_n_blocks): block = block_m * num_n_blocks + block_n row_base = block_m * 32 col_base = block_n * 32 - scale_base = block * 64 - - for row in T.Tiles([scale_shape[1]]): - scale_fp32_vec[scale_base + row] = _e8m0_scale_to_fp32(scale[block, row]) for row, col in T.Tiles([32, 32]): - q = T.Cast("float32", data[row_base + row, col_base + col]) - y_rsram[row_base + row, col_base + col] = T.Cast("bfloat16", q * scale_fp32_vec[scale_base + row]) + value = T.Cast("float32", data[row_base + row, col_base + col]) + scale_fp32 = _e8m0_scale_to_fp32(scale[block, row]) + y_rsram[row_base + row, col_base + col] = T.Cast("bfloat16", value * scale_fp32) T.copy(y_rsram, Y) @@ -456,7 +447,7 @@ def mx_ocp_quantized_mma_kernel_for_debug(mx_dtype, data_dtype, data_dtype_name, c_shape = (32, 32) a_k_blocks = a_shape[1] // 32 b_k_blocks = b_shape[1] // 32 - shard_policy = T.MeshShardingPolicy(replicate=MeshReplicationType.ALL) + shard_policy = T.placement.replicated() a_tensor_layout = make_zz_layout(a_shape, axes=[0, 1], block_shape=(32, 32)) b_tensor_layout = make_zz_layout(b_shape, axes=[0, 1], block_shape=(32, 32)) c_tensor_layout = make_zz_layout(c_shape, axes=[0, 1], block_shape=(32, 32)) @@ -481,96 +472,169 @@ def main( with T.Kernel(): a_rsram = T.alloc_shared(a_shape, T.bfloat16) b_rsram = T.alloc_shared(b_shape, T.bfloat16) - a_tile = T.alloc_shared((32, 32), T.bfloat16) - b_tile = T.alloc_shared((32, 32), T.bfloat16) - amax = T.alloc_shared((64,), T.float32) + amax = T.alloc_shared((32,), T.float32) a_data = T.alloc_shared(a_shape, data_dtype) b_data = T.alloc_shared(b_shape, data_dtype) a_scale = T.alloc_shared(a_scale_shape, T.float8_e8m0fnu) b_scale = T.alloc_shared(b_scale_shape, T.float8_e8m0fnu) - a_scale_bf16_vec = T.alloc_shared((a_scale_shape[0] * 64,), T.bfloat16) - b_scale_bf16_vec = T.alloc_shared((b_scale_shape[0] * 64,), T.bfloat16) - a_scale_fp32_vec = T.alloc_shared((a_scale_shape[0] * 64,), T.float32) - b_scale_fp32_vec = T.alloc_shared((b_scale_shape[0] * 64,), T.float32) a_mx = T.alloc_shared(a_shape, mx_dtype) b_mx = T.alloc_shared(b_shape, mx_dtype) - a_mx_asram = T.alloc_shared(a_shape, mx_dtype, scope="shared.asram") - b_mx_wsram = T.alloc_shared(b_shape, mx_dtype, scope="shared.wsram") + a_mx_asram = T.alloc_shared(a_shape, mx_dtype) + b_mx_wsram = T.alloc_shared(b_shape, mx_dtype) c_rsram = T.alloc_shared(c_shape, T.bfloat16) - T.annotate_layout( - { - a_data: a_tensor_layout, - a_mx: a_mx_layout, - a_mx_asram: a_mx_layout, - b_rsram: b_tensor_layout, - b_data: b_tensor_layout, - b_mx: b_mx_layout, - b_mx_wsram: b_mx_layout, - } - ) - T.copy(A, a_rsram) T.copy(B, b_rsram) - T.annotate_tileview( - { - a_rsram: make_tileview(a_rsram, (8, 32), (-2, -1)), - b_rsram: make_tileview(b_rsram, (8, 32), (-2, -1)), - a_tile: make_tileview(a_tile, (8, 32), (-2, -1)), - b_tile: make_tileview(b_tile, (8, 32), (-2, -1)), - a_data: make_tileview(a_data, (8, 32), (-2, -1)), - b_data: make_tileview(b_data, (8, 32), (-2, -1)), - } - ) - for block_k in T.serial(a_k_blocks): col_base = block_k * 32 - scale_base = block_k * 64 - - for row, col in T.Tiles([32, 32]): - a_tile[row, col] = a_rsram[row, col_base + col] - T.fill(amax, T.float32(1e-4)) - T.reduce_absmax(a_tile, amax[0:32], dim=1, clear=True) + T.reduce_absmax(a_rsram[0:32, col_base : col_base + 32], amax, dim=1, clear=True) for row in T.Tiles([a_scale_shape[1]]): safe_amax = T.max(amax[row], T.float32(1e-4)) - a_scale_bf16_vec[scale_base + row] = T.Cast("bfloat16", safe_amax * T.float32(data_max_inv)) - - for row in T.Tiles([a_scale_shape[1]]): - a_scale[block_k, row] = _to_e8m0_scale(a_scale_bf16_vec[scale_base + row]) - - for row in T.Tiles([a_scale_shape[1]]): - a_scale_fp32_vec[scale_base + row] = _e8m0_scale_to_fp32(a_scale[block_k, row]) + bf16_scale = T.Cast("bfloat16", safe_amax * T.float32(data_max_inv)) + a_scale[block_k, row] = _to_e8m0_scale(bf16_scale) for row, col in T.Tiles([32, 32]): - raw = T.Cast("float32", a_tile[row, col]) / a_scale_fp32_vec[scale_base + row] - clamped = T.min(T.max(raw, T.float32(-data_max)), T.float32(data_max)) - a_data[row, col_base + col] = T.Cast(data_dtype_name, clamped) + scale_fp32 = _e8m0_scale_to_fp32(a_scale[block_k, row]) + value = T.Cast("float32", a_rsram[row, col_base + col]) / scale_fp32 + a_data[row, col_base + col] = T.Cast( + data_dtype_name, + T.clamp(value, T.float32(-data_max), T.float32(data_max)), + ) for block_k in T.serial(b_k_blocks): col_base = block_k * 32 - scale_base = block_k * 64 - - for row, col in T.Tiles([32, 32]): - b_tile[row, col] = b_rsram[row, col_base + col] - T.fill(amax, T.float32(1e-4)) - T.reduce_absmax(b_tile, amax[0:32], dim=1, clear=True) + T.reduce_absmax(b_rsram[0:32, col_base : col_base + 32], amax, dim=1, clear=True) for row in T.Tiles([b_scale_shape[1]]): safe_amax = T.max(amax[row], T.float32(1e-4)) - b_scale_bf16_vec[scale_base + row] = T.Cast("bfloat16", safe_amax * T.float32(data_max_inv)) + bf16_scale = T.Cast("bfloat16", safe_amax * T.float32(data_max_inv)) + b_scale[block_k, row] = _to_e8m0_scale(bf16_scale) - for row in T.Tiles([b_scale_shape[1]]): - b_scale[block_k, row] = _to_e8m0_scale(b_scale_bf16_vec[scale_base + row]) + for row, col in T.Tiles([32, 32]): + scale_fp32 = _e8m0_scale_to_fp32(b_scale[block_k, row]) + value = T.Cast("float32", b_rsram[row, col_base + col]) / scale_fp32 + b_data[row, col_base + col] = T.Cast( + data_dtype_name, + T.clamp(value, T.float32(-data_max), T.float32(data_max)), + ) - for row in T.Tiles([b_scale_shape[1]]): - b_scale_fp32_vec[scale_base + row] = _e8m0_scale_to_fp32(b_scale[block_k, row]) + T.mx_pack(a_data, a_scale, a_mx) + T.mx_pack(b_data, b_scale, b_mx) + T.copy(a_mx, a_mx_asram) + T.copy(b_mx, b_mx_wsram) + T.clear(c_rsram) + T.gemm(a_mx_asram, b_mx_wsram, c_rsram, transpose_B=True) + T.copy(c_rsram, C) - for row, col in T.Tiles([32, 32]): - raw = T.Cast("float32", b_tile[row, col]) / b_scale_fp32_vec[scale_base + row] - clamped = T.min(T.max(raw, T.float32(-data_max)), T.float32(data_max)) - b_data[row, col_base + col] = T.Cast(data_dtype_name, clamped) + return main + + +@target("Sunmmio") +def mx_ocp_quantized_mma_generic_shape_kernel_for_debug( + M, + N, + K, + mx_dtype, + data_dtype, + data_dtype_name, + data_max, +): + assert M % 32 == 0 + assert N % 32 == 0 + assert K % 32 == 0 + + a_shape = (M, K) + b_shape = (N, K) + c_shape = (M, N) + num_m_blocks = M // 32 + num_n_blocks = N // 32 + num_k_blocks = K // 32 + shard_policy = T.placement.replicated() + a_tensor_layout = make_zz_layout(a_shape, axes=[0, 1], block_shape=(32, 32)) + b_tensor_layout = make_zz_layout(b_shape, axes=[0, 1], block_shape=(32, 32)) + c_tensor_layout = make_zz_layout(c_shape, axes=[0, 1], block_shape=(32, 32)) + a_mx_layout = make_mxzz_layout(a_shape, dtype=mx_dtype) + b_mx_layout = make_mxzz_layout(b_shape, dtype=mx_dtype) + a_scale_shape = _int_shape(get_mx_scale_shape(a_mx_layout, mx_dtype)) + b_scale_shape = _int_shape(get_mx_scale_shape(b_mx_layout, mx_dtype)) + data_max_inv = 1.0 / data_max + assert a_scale_shape == (num_m_blocks * num_k_blocks, 32) + assert b_scale_shape == (num_n_blocks * num_k_blocks, 32) + + @T.prim_func + def main( + A: T.MeshTensor(a_shape, shard_policy, T.bfloat16, layout=a_tensor_layout), # type: ignore + B: T.MeshTensor(b_shape, shard_policy, T.bfloat16, layout=b_tensor_layout), # type: ignore + C: T.MeshTensor(c_shape, shard_policy, T.bfloat16, layout=c_tensor_layout), # type: ignore + ): + with T.Kernel(): + a_rsram = T.alloc_shared(a_shape, T.bfloat16) + b_rsram = T.alloc_shared(b_shape, T.bfloat16) + amax = T.alloc_shared((32,), T.float32) + a_data = T.alloc_shared(a_shape, data_dtype) + b_data = T.alloc_shared(b_shape, data_dtype) + a_scale = T.alloc_shared(a_scale_shape, T.float8_e8m0fnu) + b_scale = T.alloc_shared(b_scale_shape, T.float8_e8m0fnu) + a_mx = T.alloc_shared(a_shape, mx_dtype) + b_mx = T.alloc_shared(b_shape, mx_dtype) + a_mx_asram = T.alloc_shared(a_shape, mx_dtype) + b_mx_wsram = T.alloc_shared(b_shape, mx_dtype) + c_rsram = T.alloc_shared(c_shape, T.bfloat16) + + T.copy(A, a_rsram) + T.copy(B, b_rsram) + + for block_m in T.serial(num_m_blocks): + for block_k in T.serial(num_k_blocks): + block = block_m * num_k_blocks + block_k + row_base = block_m * 32 + col_base = block_k * 32 + + T.reduce_absmax( + a_rsram[row_base : row_base + 32, col_base : col_base + 32], + amax, + dim=1, + clear=True, + ) + for row in T.Tiles([a_scale_shape[1]]): + safe_amax = T.max(amax[row], T.float32(1e-4)) + bf16_scale = T.Cast("bfloat16", safe_amax * T.float32(data_max_inv)) + a_scale[block, row] = _to_e8m0_scale(bf16_scale) + + for row, col in T.Tiles([32, 32]): + scale_fp32 = _e8m0_scale_to_fp32(a_scale[block, row]) + value = T.Cast("float32", a_rsram[row_base + row, col_base + col]) / scale_fp32 + a_data[row_base + row, col_base + col] = T.Cast( + data_dtype_name, + T.clamp(value, T.float32(-data_max), T.float32(data_max)), + ) + + for block_n in T.serial(num_n_blocks): + for block_k in T.serial(num_k_blocks): + block = block_n * num_k_blocks + block_k + row_base = block_n * 32 + col_base = block_k * 32 + + T.reduce_absmax( + b_rsram[row_base : row_base + 32, col_base : col_base + 32], + amax, + dim=1, + clear=True, + ) + for row in T.Tiles([b_scale_shape[1]]): + safe_amax = T.max(amax[row], T.float32(1e-4)) + bf16_scale = T.Cast("bfloat16", safe_amax * T.float32(data_max_inv)) + b_scale[block, row] = _to_e8m0_scale(bf16_scale) + + for row, col in T.Tiles([32, 32]): + scale_fp32 = _e8m0_scale_to_fp32(b_scale[block, row]) + value = T.Cast("float32", b_rsram[row_base + row, col_base + col]) / scale_fp32 + b_data[row_base + row, col_base + col] = T.Cast( + data_dtype_name, + T.clamp(value, T.float32(-data_max), T.float32(data_max)), + ) T.mx_pack(a_data, a_scale, a_mx) T.mx_pack(b_data, b_scale, b_mx) @@ -590,16 +654,13 @@ def mx_ocp_quantized_mma_mxznz_weight_kernel_for_debug(mx_dtype, data_dtype, dat c_shape = (32, 32) a_k_blocks = a_shape[1] // 32 b_k_blocks = b_shape[0] // 32 - shard_policy = T.MeshShardingPolicy(replicate=MeshReplicationType.ALL) + shard_policy = T.placement.replicated() a_tensor_layout = make_zz_layout(a_shape, axes=[0, 1], block_shape=(32, 32)) b_tensor_layout = make_zz_layout(b_shape, axes=[0, 1], block_shape=(32, 32)) - b_tile_layout = make_zz_layout((32, 32), axes=[0, 1], block_shape=(32, 32)) - b_data_layout = make_zz_layout(b_shape, axes=[0, 1], block_shape=(32, 32)) c_tensor_layout = make_zz_layout(c_shape, axes=[0, 1], block_shape=(32, 32)) a_tensor = T.MeshTensor(a_shape, shard_policy, T.bfloat16, layout=a_tensor_layout) b_tensor = T.MeshTensor(b_shape, shard_policy, T.bfloat16, layout=b_tensor_layout) c_tensor = T.MeshTensor(c_shape, shard_policy, T.bfloat16, layout=c_tensor_layout) - b_sharded_layout = b_tensor.meta_data["sharded_layout"] a_mx_layout = make_mxzz_layout(a_shape, dtype=mx_dtype) b_mx_layout = make_mxznz_layout(b_shape, dtype=mx_dtype) a_scale_shape = _int_shape(get_mx_scale_shape(a_mx_layout, mx_dtype)) @@ -617,96 +678,53 @@ def main( with T.Kernel(): a_rsram = T.alloc_shared(a_shape, T.bfloat16) b_rsram = T.alloc_shared(b_shape, T.bfloat16) - a_tile = T.alloc_shared((32, 32), T.bfloat16) - b_tile = T.alloc_shared((32, 32), T.bfloat16) - amax = T.alloc_shared((64,), T.float32) + amax = T.alloc_shared((32,), T.float32) a_data = T.alloc_shared(a_shape, data_dtype) b_data = T.alloc_shared(b_shape, data_dtype) a_scale = T.alloc_shared(a_scale_shape, T.float8_e8m0fnu) b_scale = T.alloc_shared(b_scale_shape, T.float8_e8m0fnu) - a_scale_bf16_vec = T.alloc_shared((a_scale_shape[0] * 64,), T.bfloat16) - b_scale_bf16_vec = T.alloc_shared((b_scale_shape[0] * 64,), T.bfloat16) - a_scale_fp32_vec = T.alloc_shared((a_scale_shape[0] * 64,), T.float32) - b_scale_fp32_vec = T.alloc_shared((b_scale_shape[0] * 64,), T.float32) a_mx = T.alloc_shared(a_shape, mx_dtype) b_mx = T.alloc_shared(b_shape, mx_dtype) - a_mx_asram = T.alloc_shared(a_shape, mx_dtype, scope="shared.asram") - b_mx_wsram = T.alloc_shared(b_shape, mx_dtype, scope="shared.wsram") + a_mx_asram = T.alloc_shared(a_shape, mx_dtype) + b_mx_wsram = T.alloc_shared(b_shape, mx_dtype) c_rsram = T.alloc_shared(c_shape, T.bfloat16) - T.annotate_layout( - { - a_data: a_tensor_layout, - a_mx: a_mx_layout, - a_mx_asram: a_mx_layout, - b_rsram: b_sharded_layout, - b_tile: b_tile_layout, - b_data: b_data_layout, - b_mx: b_mx_layout, - } - ) - T.copy(A, a_rsram) T.copy(B, b_rsram) - T.annotate_tileview( - { - a_rsram: make_tileview(a_rsram, (8, 32), (-2, -1)), - b_rsram: make_tileview(b_rsram, (8, 32), (-2, -1)), - a_tile: make_tileview(a_tile, (8, 32), (-2, -1)), - b_tile: make_tileview(b_tile, (8, 32), (-2, -1)), - a_data: make_tileview(a_data, (8, 32), (-2, -1)), - b_data: make_tileview(b_data, (8, 32), (-2, -1)), - } - ) - for block_k in T.serial(a_k_blocks): col_base = block_k * 32 - scale_base = block_k * 64 - - for row, col in T.Tiles([32, 32]): - a_tile[row, col] = a_rsram[row, col_base + col] - T.fill(amax, T.float32(1e-4)) - T.reduce_absmax(a_tile, amax[0:32], dim=1, clear=True) + T.reduce_absmax(a_rsram[0:32, col_base : col_base + 32], amax, dim=1, clear=True) for row in T.Tiles([a_scale_shape[1]]): safe_amax = T.max(amax[row], T.float32(1e-4)) - a_scale_bf16_vec[scale_base + row] = T.Cast("bfloat16", safe_amax * T.float32(data_max_inv)) - - for row in T.Tiles([a_scale_shape[1]]): - a_scale[block_k, row] = _to_e8m0_scale(a_scale_bf16_vec[scale_base + row]) - - for row in T.Tiles([a_scale_shape[1]]): - a_scale_fp32_vec[scale_base + row] = _e8m0_scale_to_fp32(a_scale[block_k, row]) + bf16_scale = T.Cast("bfloat16", safe_amax * T.float32(data_max_inv)) + a_scale[block_k, row] = _to_e8m0_scale(bf16_scale) for row, col in T.Tiles([32, 32]): - raw = T.Cast("float32", a_tile[row, col]) / a_scale_fp32_vec[scale_base + row] - clamped = T.min(T.max(raw, T.float32(-data_max)), T.float32(data_max)) - a_data[row, col_base + col] = T.Cast(data_dtype_name, clamped) + scale_fp32 = _e8m0_scale_to_fp32(a_scale[block_k, row]) + value = T.Cast("float32", a_rsram[row, col_base + col]) / scale_fp32 + a_data[row, col_base + col] = T.Cast( + data_dtype_name, + T.clamp(value, T.float32(-data_max), T.float32(data_max)), + ) for block_k in T.serial(b_k_blocks): row_base = block_k * 32 - scale_base = block_k * 64 - - for row, col in T.Tiles([32, 32]): - b_tile[row, col] = b_rsram[row_base + row, col] - T.fill(amax, T.float32(1e-4)) - T.reduce_absmax(b_tile, amax[0:32], dim=0, clear=True) + T.reduce_absmax(b_rsram[row_base : row_base + 32, 0:32], amax, dim=0, clear=True) for col in T.Tiles([b_scale_shape[1]]): safe_amax = T.max(amax[col], T.float32(1e-4)) - b_scale_bf16_vec[scale_base + col] = T.Cast("bfloat16", safe_amax * T.float32(data_max_inv)) - - for col in T.Tiles([b_scale_shape[1]]): - b_scale[block_k, col] = _to_e8m0_scale(b_scale_bf16_vec[scale_base + col]) - - for col in T.Tiles([b_scale_shape[1]]): - b_scale_fp32_vec[scale_base + col] = _e8m0_scale_to_fp32(b_scale[block_k, col]) + bf16_scale = T.Cast("bfloat16", safe_amax * T.float32(data_max_inv)) + b_scale[block_k, col] = _to_e8m0_scale(bf16_scale) for row, col in T.Tiles([32, 32]): - raw = T.Cast("float32", b_tile[row, col]) / b_scale_fp32_vec[scale_base + col] - clamped = T.min(T.max(raw, T.float32(-data_max)), T.float32(data_max)) - b_data[row_base + row, col] = T.Cast(data_dtype_name, clamped) + scale_fp32 = _e8m0_scale_to_fp32(b_scale[block_k, col]) + value = T.Cast("float32", b_rsram[row_base + row, col]) / scale_fp32 + b_data[row_base + row, col] = T.Cast( + data_dtype_name, + T.clamp(value, T.float32(-data_max), T.float32(data_max)), + ) T.mx_pack(a_data, a_scale, a_mx) T.mx_pack(b_data, b_scale, b_mx) @@ -818,6 +836,33 @@ def test_mx_ocp_quant_generic_shape_kernel_codegen_logs_mlir( _assert_e8m0_scale_casts(src) +@pytest.mark.parametrize("mx_dtype,data_dtype,data_dtype_name,data_max,mx_token", MX_OCP_QUANT_CASES) +def test_mx_ocp_quant_sharded_kernel_codegen_logs_mlir( + tmp_path, + mx_dtype, + data_dtype, + data_dtype_name, + data_max, + mx_token, +): + src = validate_sunmmio_codegen_with_npuir_opt( + mx_ocp_quant_sharded_kernel_for_debug(mx_dtype, data_dtype, data_dtype_name, data_max), + tmp_path, + mlir_filename=f"mx_ocp_quant_sharded_{data_dtype_name}_suvm.mlir", + expected_tokens=( + mx_token, + "!suvm.memtensor<64x64xbf16", + f"!suvm.memtensor<64x64x{mx_token}", + "suvm.unpack", + "suvm.tile.load", + "suvm.tile.store", + "suvm.tile.cast", + ), + ) + _assert_no_e8m0_tile_select(src) + _assert_e8m0_scale_casts(src) + + @pytest.mark.parametrize("mx_dtype,data_dtype,data_dtype_name,data_max,mx_token", MX_OCP_QUANT_CASES) def test_mx_ocp_dequant_generic_shape_kernel_codegen_logs_mlir( tmp_path, @@ -869,6 +914,52 @@ def test_mx_ocp_quantized_mma_kernel_codegen_logs_mlir( _assert_e8m0_scale_casts(src) +@pytest.mark.parametrize( + "M,N,K", + ( + pytest.param(32, 32, 64, id="m32_n32_k64"), + pytest.param(64, 64, 64, id="m64_n64_k64"), + pytest.param(64, 32, 128, id="m64_n32_k128"), + ), +) +@pytest.mark.parametrize("mx_dtype,data_dtype,data_dtype_name,data_max,mx_token", MX_OCP_QUANT_CASES) +def test_mx_ocp_quantized_mma_generic_shape_kernel_codegen_logs_mlir( + tmp_path, + M, + N, + K, + mx_dtype, + data_dtype, + data_dtype_name, + data_max, + mx_token, +): + src = validate_sunmmio_codegen_with_npuir_opt( + mx_ocp_quantized_mma_generic_shape_kernel_for_debug( + M, + N, + K, + mx_dtype, + data_dtype, + data_dtype_name, + data_max, + ), + tmp_path, + mlir_filename=f"mx_ocp_quantized_mma_m{M}_n{N}_k{K}_{data_dtype_name}_suvm.mlir", + expected_tokens=( + mx_token, + "suvm.unpack", + "suvm.copy_async", + "suvm.tc.mma", + "suvm.tile.load", + "suvm.tile.store", + "suvm.tile.cast", + ), + ) + _assert_no_e8m0_tile_select(src) + _assert_e8m0_scale_casts(src) + + @pytest.mark.parametrize("mx_dtype,data_dtype,data_dtype_name,data_max,mx_token", MX_OCP_QUANT_CASES) def test_mx_ocp_quantized_mma_mxznz_weight_kernel_codegen_logs_mlir_strict( tmp_path, diff --git a/testing/python/sunmmio/codegen/test_reduce_opt_validate.py b/testing/python/sunmmio/codegen/test_reduce_opt_validate.py index de6e3ca340..bb6ac1f6b5 100644 --- a/testing/python/sunmmio/codegen/test_reduce_opt_validate.py +++ b/testing/python/sunmmio/codegen/test_reduce_opt_validate.py @@ -5,7 +5,7 @@ import tilelang import tilelang.language as T import tilelang.testing -from tilelang.layout import make_row_major, make_zz_layout +from tilelang.layout import make_aligned_row_major, make_row_major, make_zz_layout from testing.python.sunmmio.common.compile_pipeline import target from testing.python.sunmmio.common.codegen_validation import ( @@ -155,7 +155,8 @@ def reduce_tiled_test( shard_policy = T.placement.replicated() input_layout = make_zz_layout(input_shape, [1, 2], (32, 32)) - output_layout = _dram_reduce_output_layout(out_shape_full) + output_layout = make_aligned_row_major(out_shape_full, dtype, align_bytes=1024) + output_shared_layout = make_aligned_row_major(out_shape_block, dtype, align_bytes=1024) grid_b = T.ceildiv(b, block_b) grid_m = T.ceildiv(m, block_m) grid_n = T.ceildiv(n, block_n) @@ -168,6 +169,7 @@ def main( with T.Kernel(): A_shared = T.alloc_shared((block_b, block_m, block_n), dtype, scope="shared.rsram") Out_shared = T.alloc_shared(out_shape_block, dtype, scope="shared.rsram") + T.annotate_layout({Out_shared: output_shared_layout}) if reduce_axis == 2: for bz in T.serial(grid_b): @@ -297,8 +299,9 @@ def test_reduce_generic_in_tile_codegen_generates_expected_ops(tmp_path, shape, @pytest.mark.parametrize("reduce_axis,clear", [(1, False), (2, True)]) def test_reduce_tiled_in_tile_codegen_generates_expected_ops(tmp_path, reduce_axis, clear): + shape_overrides = {"n": 128} if reduce_axis == 1 else {"m": 256} src = validate_sunmmio_codegen_loose( - reduce_tiled_test(reduce_axis=reduce_axis, clear=clear), + reduce_tiled_test(reduce_axis=reduce_axis, clear=clear, **shape_overrides), tmp_path, mlir_filename=f"reduce_tiled_axis_{reduce_axis}_suvm.mlir", expected_tokens=("suvm.copy_async", "suvm.tile.reduce"), @@ -313,7 +316,8 @@ def test_reduce_small_1d_result_uses_aligned_store_bridge(tmp_path): mlir_filename="reduce_small_1d_result_aligned_store_suvm.mlir", expected_tokens=("suvm.tile.reduce", "suvm.tile.insert_slice", "suvm.tile.store"), ) - assert_source_contains(src, ("suvm.tile.reduce", "!suvm.tile<8x1xbf16>", "!suvm.tile<32x1xbf16>")) + assert_source_contains(src, ("suvm.tile.reduce", "!suvm.tile<8x1xbf16>", "!suvm.tile<32xbf16>")) + assert "suvm.tile.unsqueeze" not in src assert "fake_tile_insert_slice" not in src assert "suvm.tile.store" in src assert "fake_tile_store" not in src diff --git a/testing/python/sunmmio/codegen/test_sliding_window_attention_gqa_fwd_bhsd_opt_validate.py b/testing/python/sunmmio/codegen/test_sliding_window_attention_gqa_fwd_bhsd_opt_validate.py index 50abd38120..3c763b5444 100644 --- a/testing/python/sunmmio/codegen/test_sliding_window_attention_gqa_fwd_bhsd_opt_validate.py +++ b/testing/python/sunmmio/codegen/test_sliding_window_attention_gqa_fwd_bhsd_opt_validate.py @@ -48,7 +48,7 @@ def sliding_window_attention_gqa_fwd_bhsd( assert global_window <= seq_len assert global_window <= block_N - shard_policy = T.MeshShardingPolicy(y=0, x=2) + shard_policy = T.placement.full_shard(0, 2) Q_layout = make_zz_layout(q_shape, [1, 3], (32, 32)) K_layout = make_zz_layout(kv_shape, [1, 3], (32, 32)) diff --git a/testing/python/sunmmio/codegen/test_summa_opt_validate.py b/testing/python/sunmmio/codegen/test_summa_opt_validate.py index 0b48e3d500..2d9dd0f68e 100644 --- a/testing/python/sunmmio/codegen/test_summa_opt_validate.py +++ b/testing/python/sunmmio/codegen/test_summa_opt_validate.py @@ -55,8 +55,8 @@ def kernel( with T.Kernel() as _cid: sharded_M, _ = A.local_shape _, sharded_N = B.local_shape - core_row = _cid // T.mesh_ncols() - core_col = _cid % T.mesh_ncols() + core_row = _cid // T.ncols() + core_col = _cid % T.ncols() A_broadcast = T.alloc_shared((block_M, block_K), dtype, scope="shared.rsram") A_shared = T.alloc_shared((block_M, block_K), dtype) @@ -69,10 +69,10 @@ def kernel( K_steps = T.ceildiv(K, block_K) for k_tile in range(K_steps): - a_src_col = k_tile % T.mesh_ncols() - b_src_row = k_tile % T.mesh_nrows() - a_local_k = (k_tile // T.mesh_ncols()) * block_K - b_local_k = (k_tile // T.mesh_nrows()) * block_K + a_src_col = k_tile % T.ncols() + b_src_row = k_tile % T.nrows() + a_local_k = (k_tile // T.ncols()) * block_K + b_local_k = (k_tile // T.nrows()) * block_K T.comm.broadcast( A[ diff --git a/testing/python/sunmmio/codegen/test_tile_ops_opt_validate.py b/testing/python/sunmmio/codegen/test_tile_ops_opt_validate.py index 1151c1c77d..6f4ae1ba19 100644 --- a/testing/python/sunmmio/codegen/test_tile_ops_opt_validate.py +++ b/testing/python/sunmmio/codegen/test_tile_ops_opt_validate.py @@ -1,4 +1,5 @@ import os +import re import tilelang import tilelang.language as T @@ -17,6 +18,7 @@ # os.environ["SUNMMIO_TEST_LOG_IR"] = "1" LOOSE_OPT_ARGS = ("--verify-each",) +STRICT_OPT_ARGS = ("--verify-each", "--suvm-to-llvm-pipeline") def validate_sunmmio_codegen_loose(kernel, tmp_path, *, mlir_filename, expected_tokens=()): @@ -188,6 +190,38 @@ def main( return main +@target("Sunmmio") +def fp32_select_then_bf16_cast_test(m=32, n=32): + input_dtype = T.float32 + output_dtype = T.bfloat16 + shard_policy = T.placement.replicated() + tensor_shape = (m, n) + tensor_layout = make_zz_layout(tensor_shape, [0, 1], tensor_shape) + + @T.prim_func + def main( + A: T.MeshTensor(tensor_shape, shard_policy, input_dtype, layout=tensor_layout), # type: ignore + C: T.MeshTensor(tensor_shape, shard_policy, output_dtype, layout=tensor_layout), # type: ignore + ): + with T.Kernel(): + A_shared = T.alloc_shared(tensor_shape, input_dtype) + C_shared = T.alloc_shared(tensor_shape, output_dtype) + + T.copy(A, A_shared) + for i, j in T.Tiles(A_shared, parallel=True): + C_shared[i, j] = T.Cast( + output_dtype, + T.if_then_else( + A_shared[i, j] > T.float32(0), + A_shared[i, j], + T.float32(0), + ), + ) + T.copy(C_shared, C) + + return main + + def test_tile_elementwise_ops_2d_codegen_validates_with_npuir_opt(tmp_path): src = validate_sunmmio_codegen_with_npuir_opt( tile_elementwise_ops_2d_test(), @@ -244,5 +278,28 @@ def test_tile_elementwise_ops_codegen_validates_loose_with_npuir_opt(tmp_path): ) +def test_fp32_select_is_evaluated_before_bf16_cast(tmp_path): + src = validate_sunmmio_codegen_with_npuir_opt( + fp32_select_then_bf16_cast_test(), + tmp_path, + mlir_filename="fp32_select_then_bf16_cast_suvm.mlir", + expected_tokens=("suvm.tile.cmpf", "suvm.tile.select", "suvm.tile.cast"), + opt_args=STRICT_OPT_ARGS, + ) + + select = re.search( + r"(?P%[\w.]+) = suvm\.tile\.select .*" + r"!suvm\.tile<[^>]*xf32>, !suvm\.tile<[^>]*xf32>" + r" -> !suvm\.tile<[^>]*xf32>", + src, + ) + assert select, src + assert re.search( + rf"suvm\.tile\.cast {re.escape(select.group('result'))} : " + r"!suvm\.tile<[^>]*xf32> -> !suvm\.tile<[^>]*xbf16>", + src, + ), src + + if __name__ == "__main__": tilelang.testing.main() diff --git a/testing/python/sunmmio/codegen/test_tile_pick_scalar_access.py b/testing/python/sunmmio/codegen/test_tile_pick_scalar_access.py index 3c07f641af..6e407f1439 100644 --- a/testing/python/sunmmio/codegen/test_tile_pick_scalar_access.py +++ b/testing/python/sunmmio/codegen/test_tile_pick_scalar_access.py @@ -167,6 +167,22 @@ def main( return main +@target("Sunmmio") +def stage_2d_short_aligned_row_kernel(rows=3, cols=64, dtype=T.bfloat16): + table_shape = (rows, cols) + table_layout = make_aligned_row_major(table_shape, dtype, align_bytes=1024) + + @T.prim_func + def main(table: T.MeshTensor(table_shape, T.placement.replicated(), dtype, layout=table_layout)): # type: ignore + with T.Kernel(): + row_shared = T.alloc_shared((cols,), dtype, scope="shared.rsram") + T.annotate_layout({row_shared: make_aligned_row_major((cols,), dtype, align_bytes=1024)}) + for row in T.serial(table.local_shape[0]): + T.copy(table[row, :], row_shared) + + return main + + @target("Sunmmio") def pick_3d_side_data_kernel( heads=2, @@ -555,6 +571,17 @@ def test_pick_scalar_access_codegen_with_explicit_rsram_staging(factory, mlir_fi ) +def test_pick_2d_short_aligned_row_expands_dma_to_covered_extent(tmp_path): + src = validate_sunmmio_codegen_with_npuir_opt( + stage_2d_short_aligned_row_kernel(), + tmp_path, + mlir_filename="pick_2d_short_aligned_row_suvm.mlir", + expected_tokens=("suvm.copy_async", "!suvm.tile_view<512xbf16>"), + opt_args=("--verify-each", "--suvm-to-llvm-pipeline"), + ) + assert "#suvm.layout<(3, 512), (512, 1)>" in src + + def test_pick_3d_predicated_1d_store_preserves_old_lanes(tmp_path): validate_sunmmio_codegen_with_npuir_opt( pick_3d_side_data_kernel(heads=2, q_blocks=5, k_blocks=256, out_tiles=2), diff --git a/testing/python/sunmmio/codegen/test_tiles_aligned_store.py b/testing/python/sunmmio/codegen/test_tiles_aligned_store.py index 9a4e6f04db..fcb39852b9 100644 --- a/testing/python/sunmmio/codegen/test_tiles_aligned_store.py +++ b/testing/python/sunmmio/codegen/test_tiles_aligned_store.py @@ -27,10 +27,10 @@ def _build_sunmmio_source_from_func(func): def _has_nonzero_1d_insert_slice_offset(src): - insert_lines = [line for line in src.splitlines() if "suvm.tile.insert_slice" in line and "] [8, 1]" in line] - if any("[8, 0] [8, 1]" in line for line in insert_lines): + insert_lines = [line for line in src.splitlines() if "suvm.tile.insert_slice" in line and "] [8]" in line] + if any("[8] [8]" in line for line in insert_lines): return True - return any("[%" in line and ", 0] [8, 1]" in line for line in insert_lines) and "arith.remsi" in src + return any("[%" in line for line in insert_lines) and "arith.remsi" in src @target("Sunmmio") @@ -176,7 +176,6 @@ def _make_row_major_padded_2d_aligned_store_func(): def test_sunmmio_codegen_aligned_1d_store_uses_nonzero_insert_slice_offset(): src = _build_sunmmio_source_from_stmt(_make_nonzero_offset_aligned_store_stmt()) assert "suvm.tile.insert_slice" in src - assert "suvm.tile.unsqueeze" in src assert "suvm.tile.store" in src assert "fake_tile_store" not in src assert "!suvm.tile<32xbf16>" in src diff --git a/testing/python/sunmmio/codegen/test_tiles_dynamic_domain_opt_validate.py b/testing/python/sunmmio/codegen/test_tiles_dynamic_domain_opt_validate.py index 8cb30a608b..7555fe772d 100644 --- a/testing/python/sunmmio/codegen/test_tiles_dynamic_domain_opt_validate.py +++ b/testing/python/sunmmio/codegen/test_tiles_dynamic_domain_opt_validate.py @@ -19,7 +19,7 @@ @target("Sunmmio") def dynamic_rank2_domain_kernel(h=4, matrix_size=32, dtype=T.float32): out_shape = (16, matrix_size, matrix_size) - token_policy = T.MeshShardingPolicy(cross_mesh_dim=0) + token_policy = T.placement.mesh_as_line(0) out_layout = make_zz_layout(out_shape, [1, 2], (32, 32)) lengths_shape = (128,) lengths_layout = make_aligned_row_major(lengths_shape, T.int32, align_bytes=1024) @@ -28,7 +28,7 @@ def dynamic_rank2_domain_kernel(h=4, matrix_size=32, dtype=T.float32): @T.prim_func def main( out: T.MeshTensor(out_shape, token_policy, dtype, layout=out_layout), # type: ignore - lengths: T.MeshTensor(lengths_shape, T.MeshShardingPolicy(), T.int32, layout=lengths_layout), # type: ignore + lengths: T.MeshTensor(lengths_shape, T.placement.replicated(), T.int32, layout=lengths_layout), # type: ignore ): with T.Kernel(): src = T.alloc_shared((matrix_size, matrix_size), dtype) diff --git a/testing/python/sunmmio/codegen/test_tiles_fallback_opt_validate.py b/testing/python/sunmmio/codegen/test_tiles_fallback_opt_validate.py index 40a4415feb..51680db901 100644 --- a/testing/python/sunmmio/codegen/test_tiles_fallback_opt_validate.py +++ b/testing/python/sunmmio/codegen/test_tiles_fallback_opt_validate.py @@ -18,7 +18,7 @@ def _matrix_output_spec(h, w, dtype): shape = (16, h, w) - return shape, T.MeshShardingPolicy(cross_mesh_dim=0), make_zz_layout(shape, [1, 2], (32, 32)) + return shape, T.placement.mesh_as_line(0), make_zz_layout(shape, [1, 2], (32, 32)) @target("Sunmmio") diff --git a/testing/python/sunmmio/codegen/test_tiles_opt_validate.py b/testing/python/sunmmio/codegen/test_tiles_opt_validate.py index 56fe5faf35..d9b52e5c6a 100644 --- a/testing/python/sunmmio/codegen/test_tiles_opt_validate.py +++ b/testing/python/sunmmio/codegen/test_tiles_opt_validate.py @@ -4,7 +4,7 @@ import tilelang import tilelang.language as T import tilelang.testing -from tilelang.layout import make_aligned_row_major, make_row_major, make_zz_layout +from tilelang.layout import make_aligned_row_major, make_zz_layout from testing.python.sunmmio.common.compile_pipeline import target from testing.python.sunmmio.common.codegen_validation import ( @@ -163,7 +163,7 @@ def tiles_broadcast( m=512, n=1024, block_b=2, - block_m=256, + block_m=512, block_n=128, dtype="bfloat16", accum_dtype="bfloat16", @@ -172,7 +172,8 @@ def tiles_broadcast( tensor_shape = (batch, m, n) tensor_layout = make_zz_layout(tensor_shape, [1, 2], (32, 32)) vector_shape = (m,) - vector_layout = make_row_major(vector_shape) + vector_layout = make_aligned_row_major(vector_shape, dtype, align_bytes=1024) + vector_shared_layout = make_aligned_row_major((block_m,), dtype, align_bytes=1024) grid_b = T.ceildiv(batch, block_b) grid_m = T.ceildiv(m, block_m) grid_n = T.ceildiv(n, block_n) @@ -189,6 +190,7 @@ def main( B_shared = T.alloc_shared((block_b, block_m, block_n), dtype) C_shared = T.alloc_shared((block_b, block_m, block_n), accum_dtype) D_shared = T.alloc_shared((block_m,), dtype) + T.annotate_layout({D_shared: vector_shared_layout}) for bz in T.serial(grid_b): for by in T.serial(grid_m): @@ -241,7 +243,7 @@ def tiles_broadcast_copy( m=512, n=1024, block_b=2, - block_m=256, + block_m=512, block_n=128, dtype="bfloat16", accum_dtype="bfloat16", @@ -250,7 +252,8 @@ def tiles_broadcast_copy( tensor_shape = (batch, m, n) tensor_layout = make_zz_layout(tensor_shape, [1, 2], (32, 32)) vector_shape = (m,) - vector_layout = make_row_major(vector_shape) + vector_layout = make_aligned_row_major(vector_shape, dtype, align_bytes=1024) + vector_shared_layout = make_aligned_row_major((block_m,), dtype, align_bytes=1024) grid_b = T.ceildiv(batch, block_b) grid_m = T.ceildiv(m, block_m) grid_n = T.ceildiv(n, block_n) @@ -267,6 +270,7 @@ def main( B_shared = T.alloc_shared((block_b, block_m, block_n), dtype) C_shared = T.alloc_shared((block_b, block_m, block_n), accum_dtype) D_shared = T.alloc_shared((block_m,), dtype) + T.annotate_layout({D_shared: vector_shared_layout}) for bz in T.serial(grid_b): for by in T.serial(grid_m): @@ -317,10 +321,12 @@ def main( @target("Sunmmio") -def tiles_1d(m=512, block_m=256, dtype="bfloat16", accum_dtype="bfloat16"): +def tiles_1d(m=512, block_m=512, dtype="bfloat16", accum_dtype="bfloat16"): shard_policy = T.placement.replicated() tensor_shape = (m,) - tensor_layout = make_row_major(tensor_shape) + tensor_layout = make_aligned_row_major(tensor_shape, dtype, align_bytes=1024) + shared_layout = make_aligned_row_major((block_m,), dtype, align_bytes=1024) + accum_shared_layout = make_aligned_row_major((block_m,), accum_dtype, align_bytes=1024) grid_m = T.ceildiv(m, block_m) @T.prim_func @@ -333,6 +339,7 @@ def main( A_shared = T.alloc_shared((block_m,), dtype) B_shared = T.alloc_shared((block_m,), dtype) C_shared = T.alloc_shared((block_m,), accum_dtype) + T.annotate_layout({A_shared: shared_layout, B_shared: shared_layout, C_shared: accum_shared_layout}) for by in T.serial(grid_m): T.clear(C_shared) @@ -385,7 +392,7 @@ def main( def tiles_rank2_first_tile_partial(rows=4, cols=4, dtype="float32"): output_shape = (32, 32) output_layout = make_zz_layout(output_shape, [0, 1], (32, 32)) - shard_policy = T.MeshShardingPolicy() + shard_policy = T.placement.replicated() @T.prim_func def main( diff --git a/testing/python/sunmmio/codegen/test_tiles_rank1_tail_mask_opt_validate.py b/testing/python/sunmmio/codegen/test_tiles_rank1_tail_mask_opt_validate.py index 4b34bdbfb4..227a06644b 100644 --- a/testing/python/sunmmio/codegen/test_tiles_rank1_tail_mask_opt_validate.py +++ b/testing/python/sunmmio/codegen/test_tiles_rank1_tail_mask_opt_validate.py @@ -26,8 +26,8 @@ def dynamic_rank1_tail_mask_kernel(vector_size=512, dtype=T.float32): @T.prim_func def main( - out: T.MeshTensor(out_shape, T.MeshShardingPolicy(), dtype, layout=out_layout), # type: ignore - lengths: T.MeshTensor(lengths_shape, T.MeshShardingPolicy(), T.int32, layout=lengths_layout), # type: ignore + out: T.MeshTensor(out_shape, T.placement.replicated(), dtype, layout=out_layout), # type: ignore + lengths: T.MeshTensor(lengths_shape, T.placement.replicated(), T.int32, layout=lengths_layout), # type: ignore ): with T.Kernel(): out_shared = T.alloc_shared(out_shape, dtype) diff --git a/testing/python/sunmmio/codegen/test_tiles_shape_predicate_projection_opt_validate.py b/testing/python/sunmmio/codegen/test_tiles_shape_predicate_projection_opt_validate.py index b8f972668b..56b36f67cf 100644 --- a/testing/python/sunmmio/codegen/test_tiles_shape_predicate_projection_opt_validate.py +++ b/testing/python/sunmmio/codegen/test_tiles_shape_predicate_projection_opt_validate.py @@ -34,7 +34,7 @@ def tiles_shape_rank1_side_buffer_predicate_kernel( assert cols % 32 == 0 assert num_tokens % 16 == 0 - token_policy = T.MeshShardingPolicy(cross_mesh_dim=0) + token_policy = T.placement.mesh_as_line(0) out_shape = (num_tokens, rows, cols) out_layout = make_zz_layout(out_shape, [1, 2], (32, 32)) a_layout = make_zz_layout((rows, cols), [0, 1], (32, 32)) diff --git a/testing/python/sunmmio/common/formal_verify.py b/testing/python/sunmmio/common/formal_verify.py index 22990878e3..237654c6c0 100644 --- a/testing/python/sunmmio/common/formal_verify.py +++ b/testing/python/sunmmio/common/formal_verify.py @@ -242,6 +242,28 @@ def check(mod: IRModule): else: mask_pattern = rf"(?:T\.int64\({mask}\)|{mask})" if has_src_core: + if not core.lstrip("-").isdigit(): + # Lowering may inline let-bound dynamic coordinates and + # rename the kernel thread variable. Verify the lowered + # broadcast ABI here; per-test checks own its exact route. + message = f"Expected broadcast_ with a dynamic source core, direction={direction}, mask={mask} not found in IRModule" + found = False + for args in broadcast_arg_lists: + if len(args) == 6: + fixed_args = args + elif len(args) == 7 and args[-1].startswith("T.sync_token_id("): + fixed_args = args[:-1] + else: + continue + if fixed_args[2] != str(direction): + continue + if mask is not None and fixed_args[3] not in {f"T.int64({mask})", str(mask)}: + continue + found = True + break + assert found, message + continue + # Match T.broadcast_(..., direction, mask, src_offset_byte, src_core, ...) escaped_core = re.escape(core).replace(r"\ ", r"\s*") pattern = rf"T\.broadcast_\(.*?,\s*.*?,\s*{direction},\s*{mask_pattern},\s*.*?,\s*{escaped_core}(?:,|\))" diff --git a/testing/python/sunmmio/jit/test_jit.py b/testing/python/sunmmio/jit/test_jit.py index 5c8059dac5..677af3651a 100644 --- a/testing/python/sunmmio/jit/test_jit.py +++ b/testing/python/sunmmio/jit/test_jit.py @@ -1171,10 +1171,18 @@ def test_sunmmio_softmax_output_dma_wait_is_hoisted_before_tile_store_loop(): for idx, line in enumerate(lines[input_dma_idx:store_idx], start=input_dma_idx) if "for i in T.serial" in line and "tile.domain" in line ) - # Allow the output-DMA pipeline's prologue null-init (before the by-loop); forbid resets inside the loop body. - assert all(null_marker not in line for line in lines[input_dma_idx:output_dma_idx]) - assert any(wait_marker in line for line in lines[input_dma_idx:tile_store_loop_idx]) - assert all(wait_marker not in line for line in lines[tile_store_loop_idx:output_dma_idx]) + output_pipeline_loop_idx = max( + idx for idx, line in enumerate(lines[:input_dma_idx]) if re.search(r"\bfor by in (?:range|T\.serial)\(", line) + ) + prologue_null_idx = max(idx for idx, line in enumerate(lines[:output_pipeline_loop_idx]) if null_marker in line) + + # The previous output DMA and the next input DMA both use ODMA0. Its + # loop-carried wait must run after the prologue null-init but before the + # next input submission, or the channel-wide wait would drain that input. + assert all(null_marker not in line for line in lines[output_pipeline_loop_idx:output_dma_idx]) + assert any(wait_marker in line for line in lines[output_pipeline_loop_idx:input_dma_idx]) + assert all(wait_marker not in line for line in lines[input_dma_idx:output_dma_idx]) + assert prologue_null_idx < output_pipeline_loop_idx < input_dma_idx < tile_store_loop_idx < output_dma_idx def test_sunmmio_base_adapter_does_not_expose_sunsim_runtime_surface(tmp_path): @@ -1265,7 +1273,7 @@ def elementwise_add_jit(M, N, block_M, block_N, in_dtype, out_dtype): """JIT version of examples/sunmmio/elementwise/elementwise_add.py.""" zz_layout = make_zz_layout((M, N)) - placement = T.MeshShardingPolicy(y=0, x=1) + placement = T.placement.full_shard(0, 1) @T.prim_func def elem_add( diff --git a/testing/python/sunmmio/language/test_comm.py b/testing/python/sunmmio/language/test_comm.py index c6e59eb454..3e93605ae9 100644 --- a/testing/python/sunmmio/language/test_comm.py +++ b/testing/python/sunmmio/language/test_comm.py @@ -258,8 +258,8 @@ def test_comm_compact_path_accepts_unresolved_mesh_shape_without_warning(): def main(): with T.Kernel(): send = T.alloc_shared((32, 32), "float32", scope="shared.rsram") - recv = T.alloc_shared((32, 32 * T.mesh_ncols()), "float32", scope="shared.rsram") - reduce_src = T.alloc_shared((32 * T.mesh_ncols(), 32), "float32", scope="shared.rsram") + recv = T.alloc_shared((32, 32 * T.ncols()), "float32", scope="shared.rsram") + reduce_src = T.alloc_shared((32 * T.ncols(), 32), "float32", scope="shared.rsram") reduce_out = T.alloc_shared((128,), "float32", scope="shared.rsram") T.comm.all_gather(send, recv, direction="h", axis=-1) T.comm.all_reduce(reduce_src, reduce_out, "sum", "h", dim=1) @@ -279,7 +279,7 @@ def test_comm_compact_path_keeps_static_checks_with_unresolved_mesh_shape(): def main(): with T.Kernel(): send = T.alloc_shared((32, 32), "float32", scope="shared.rsram") - recv = T.alloc_shared((64, 32 * T.mesh_ncols()), "float32", scope="shared.rsram") + recv = T.alloc_shared((64, 32 * T.ncols()), "float32", scope="shared.rsram") T.comm.all_gather(send, recv, direction="h", axis=-1) finally: _target_utils.set_sunmmio_region_validation(previous) diff --git a/testing/python/sunmmio/language/test_copy_legality.py b/testing/python/sunmmio/language/test_copy_legality.py index 3864118a38..71cd06d0e4 100644 --- a/testing/python/sunmmio/language/test_copy_legality.py +++ b/testing/python/sunmmio/language/test_copy_legality.py @@ -201,7 +201,7 @@ def _build_script(copy_case): @target("Sunmmio") def _make_let_bound_mesh_copy_kernel(): global_shape = (256, 256) - shard_policy = T.MeshShardingPolicy(y=0, x=1) + shard_policy = T.placement.full_shard(0, 1) tensor_layout = make_zz_layout(global_shape, axes=[0, 1], block_shape=(32, 32)) @T.prim_func @@ -219,7 +219,7 @@ def kernel( @target("Sunmmio") def _make_mismatched_let_bound_mesh_copy_kernel(): global_shape = (256, 256) - shard_policy = T.MeshShardingPolicy(y=0, x=1) + shard_policy = T.placement.full_shard(0, 1) tensor_layout = make_zz_layout(global_shape, axes=[0, 1], block_shape=(32, 32)) @T.prim_func diff --git a/testing/python/sunmmio/language/test_mesh_tensor_shape_api.py b/testing/python/sunmmio/language/test_mesh_tensor_shape_api.py index a24028c09b..f3dcabd892 100644 --- a/testing/python/sunmmio/language/test_mesh_tensor_shape_api.py +++ b/testing/python/sunmmio/language/test_mesh_tensor_shape_api.py @@ -78,7 +78,7 @@ def test_placement_rejects_dimensions_outside_tensor_rank(placement): def test_mesh_tensor_shape_api_in_kernel(): tensor = T.MeshTensor( (513, 4097), - T.MeshShardingPolicy(y=0, x=1), + T.placement.full_shard(0, 1), (4, 4), "float16", ) @@ -95,10 +95,10 @@ def test_mesh_tensor_shape_api_in_kernel(): @T.prim_func def kernel(A: tensor): - with T.Kernel(T.mesh_ncores()) as cid: + with T.Kernel(T.mesh_ncores()): global_m, global_n = A.global_shape local_m, local_n = A.local_shape - valid_m, valid_n = A.get_local_extent(cid) + valid_m, valid_n = A.get_local_extent() core0_m, core0_n = A.get_local_extent(0) core15_m, core15_n = A.get_local_extent(15) @@ -121,7 +121,7 @@ def kernel(A: tensor): def test_mesh_tensor_same_dim_row_then_col_extent(): tensor = T.MeshTensor( (65, 9), - T.MeshShardingPolicy(y=0, x=0), + T.placement.full_shard(0, 0), (4, 4), "float16", ) @@ -137,10 +137,10 @@ def test_mesh_tensor_same_dim_row_then_col_extent(): @T.prim_func def kernel(A: tensor): - with T.Kernel(T.mesh_ncores()) as cid: + with T.Kernel(T.mesh_ncores()): global_m, global_n = A.global_shape local_m, local_n = A.local_shape - valid_m, valid_n = A.get_local_extent(cid) + valid_m, valid_n = A.get_local_extent() core0_m, core0_n = A.get_local_extent(0) core1_m, core1_n = A.get_local_extent(1) core15_m, core15_n = A.get_local_extent(15) diff --git a/testing/python/sunmmio/ops/test_datapath.py b/testing/python/sunmmio/ops/test_datapath.py index d70bb07c0a..34ada3d191 100644 --- a/testing/python/sunmmio/ops/test_datapath.py +++ b/testing/python/sunmmio/ops/test_datapath.py @@ -8,7 +8,6 @@ from tilelang import tvm as tvm from tilelang.layout import make_zz_layout from tilelang.utils.target import SUNMMIO_TARGET_DESC, determine_target -from tilelang.language.mesh_tensor import MeshShardingPolicy from tvm import tir from tvm.tir import PyStmtExprVisitor import pytest @@ -280,7 +279,7 @@ def copy(K, block_M, block_N, block_K, dtype="float32", accum_dtype="float32"): _layout = make_zz_layout((128, 128), [0, 1], (32, 32)) MyTensor = T.MeshTensor( (128, 128), - sharding_policy=MeshShardingPolicy(cross_mesh_dim=0), + placement=T.placement.mesh_as_line(0), layout=_layout, ) diff --git a/testing/python/sunmmio/ops/test_reduce.py b/testing/python/sunmmio/ops/test_reduce.py index 990041d27e..75c47d1512 100644 --- a/testing/python/sunmmio/ops/test_reduce.py +++ b/testing/python/sunmmio/ops/test_reduce.py @@ -209,18 +209,35 @@ def unaligned_reduce_kernel_builder(shape, reduce_axis, dtype="float16", clear=T out_shape = list(shape[:reduce_axis]) + list(shape[reduce_axis + 1 :]) if not out_shape: out_shape = [1] + input_boundary_layout = make_aligned_row_major(shape, dtype, 1024) if len(shape) == 1 else None + if input_boundary_layout is not None: + placement = T.placement.replicated() + input_type = T.MeshTensor(shape, placement, dtype, layout=input_boundary_layout) + output_type = T.Tensor(out_shape, dtype) + else: + input_type = T.Tensor(shape, dtype) + output_type = T.Tensor(out_shape, dtype) @T.prim_func - def main(A: T.Tensor(shape, dtype), Out: T.Tensor(out_shape, dtype)): + def main(A: input_type, Out: output_type): with T.Kernel(1, threads=128) as (bx,): A_shared = T.alloc_shared(shape, dtype, scope="shared.rsram") Out_shared = T.alloc_shared(out_shape, dtype, scope="shared.rsram") + if input_boundary_layout is not None: + T.annotate_layout({A_shared: input_boundary_layout}) + T.copy(A, A_shared) if not clear: - T.copy(Out, Out_shared) + if input_boundary_layout is not None: + # A rank-1 reduction produces an effective-rank-0 (1,) result, + # which is outside the aligned-row DMA carrier contract. + T.fill(Out_shared, 0) + else: + T.copy(Out, Out_shared) apply_reduce_op(reduce_op, A_shared, Out_shared, reduce_axis, clear=clear) - T.copy(Out_shared, Out) + if input_boundary_layout is None: + T.copy(Out_shared, Out) return tvm.IRModule({"main": main}) diff --git a/testing/python/sunmmio/ops/test_transpose.py b/testing/python/sunmmio/ops/test_transpose.py index c57caa3712..b7d4c10a7e 100644 --- a/testing/python/sunmmio/ops/test_transpose.py +++ b/testing/python/sunmmio/ops/test_transpose.py @@ -44,7 +44,7 @@ def mesh_transpose_kernel( expect_transposed=True, ): """Build a replicated transpose with matching DRAM and RSRAM layouts.""" - placement = T.MeshShardingPolicy(replicate=T.MeshReplicationType.ALL) + placement = T.placement.replicated() src_layout = make_zz_layout((m, n)) if layout_family == "zz" else make_zn_layout((m, n), [0, 1], (32, 32)) transposed_layout = make_zz_layout((n, m)) if layout_family == "zz" else make_zn_layout((n, m), [0, 1], (32, 32)) output_shape = (n, m) if expect_transposed else (m, n) @@ -86,7 +86,7 @@ def mesh_transpose_order_kernel(source_constraint_first): """Build the same transpose layout constraints in either source order.""" size = 64 dtype = "bfloat16" - placement = T.MeshShardingPolicy(replicate=T.MeshReplicationType.ALL) + placement = T.placement.replicated() zn_layout = make_zn_layout((size, size), [0, 1], (32, 32)) @T.prim_func @@ -118,7 +118,7 @@ def mesh_transpose_scope_kernel(global_operand): """Build a transpose with one operand intentionally in global memory.""" size = 64 dtype = "bfloat16" - placement = T.MeshShardingPolicy(replicate=T.MeshReplicationType.ALL) + placement = T.placement.replicated() layout = make_zz_layout((size, size)) @T.prim_func diff --git a/testing/python/sunmmio/pipeline/test_comm.py b/testing/python/sunmmio/pipeline/test_comm.py index fcf1de9f12..5fc34366ec 100644 --- a/testing/python/sunmmio/pipeline/test_comm.py +++ b/testing/python/sunmmio/pipeline/test_comm.py @@ -14,7 +14,7 @@ @target("Sunmmio") def kernel_comm(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="float32"): - shard_policy = T.MeshShardingPolicy(y=0, x=1) + shard_policy = T.placement.full_shard(0, 1) A_shape = (M, K) B_shape = (K, N) diff --git a/testing/python/sunmmio/pipeline/test_flashattn.py b/testing/python/sunmmio/pipeline/test_flashattn.py index 08d741d892..602fa72940 100644 --- a/testing/python/sunmmio/pipeline/test_flashattn.py +++ b/testing/python/sunmmio/pipeline/test_flashattn.py @@ -25,7 +25,7 @@ def kernel_flashattn( dtype = T.bfloat16 # accum_dtype = T.float32 accum_dtype = T.bfloat16 - shard_policy = T.MeshShardingPolicy(y=0, x=2) + shard_policy = T.placement.full_shard(0, 2) Q_layout = make_zz_layout(q_shape, [1, 3], (32, 32)) K_layout = make_zz_layout(kv_shape, [1, 3], (32, 32)) diff --git a/testing/python/sunmmio/pipeline/test_mma_3times.py b/testing/python/sunmmio/pipeline/test_mma_3times.py index 99866c5f6a..608ef91bf5 100644 --- a/testing/python/sunmmio/pipeline/test_mma_3times.py +++ b/testing/python/sunmmio/pipeline/test_mma_3times.py @@ -9,7 +9,7 @@ @target("Sunmmio") def kernel_mma_3times_single_thread(M=16, N=16, K=16, block_M=128, block_N=128, block_K=32, dtype="float16"): - shard_policy = T.MeshShardingPolicy(y=0, x=1) + shard_policy = T.placement.full_shard(0, 1) A_shape = (M, K) B_shape = (K, N) diff --git a/testing/python/sunmmio/pipeline/test_overall.py b/testing/python/sunmmio/pipeline/test_overall.py index 16556b3035..d39a9d02b6 100644 --- a/testing/python/sunmmio/pipeline/test_overall.py +++ b/testing/python/sunmmio/pipeline/test_overall.py @@ -9,7 +9,7 @@ @target("Sunmmio") def kernel_overall(M, N, K, block_M, block_N, block_K, dtype="bfloat16", accum_dtype="float32"): - shard_policy = T.MeshShardingPolicy(y=0, x=1) + shard_policy = T.placement.full_shard(0, 1) A_shape = (M, K) B_shape = (K, N) diff --git a/testing/python/sunmmio/pipeline/test_summa.py b/testing/python/sunmmio/pipeline/test_summa.py index 437d940df6..ab6bd2e489 100644 --- a/testing/python/sunmmio/pipeline/test_summa.py +++ b/testing/python/sunmmio/pipeline/test_summa.py @@ -13,9 +13,10 @@ def summa_matmul(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtyp SUMMA (Scalable Universal Matrix Multiplication Algorithm) for a 4x4 mesh. - Grid size: (N/block_N, M/block_M) = (4, 4) + Each core accumulates its local C tiles over all global K panels. For each + panel, A is broadcast along the current core row and B along its column. """ - shard_policy = T.MeshShardingPolicy(y=0, x=1) + shard_policy = T.placement.full_shard(0, 1) A_shape = (M, K) B_shape = (K, N) @@ -30,58 +31,50 @@ def kernel( B: T.MeshTensor(B_shape, shard_policy, dtype, layout=B_layout), C: T.MeshTensor(C_shape, shard_policy, accum_dtype, layout=C_layout), ): - # Assume the current is a 4x4 processor grid (Mesh) - # Each core is responsible for outputting a 32x32 block of matrix C with T.Kernel() as _cid: - sharded_M, sharded_K = A.local_shape + sharded_M, _ = A.local_shape _, sharded_N = B.local_shape + core_row = _cid // T.mesh_ncols() + core_col = _cid % T.mesh_ncols() - # Allocate local SRAM cache - # A_shared is placed in ASRAM (usually used for A matrix cache) - # B_shared is placed in WSRAM (usually used for B matrix cache) + # Multicast lands in RSRAM before A is copied into MMA's ASRAM input. + A_broadcast = T.alloc_shared((block_M, block_K), dtype, scope="shared.rsram") A_shared = T.alloc_shared((block_M, block_K), dtype) B_shared = T.alloc_shared((block_K, block_N), dtype) - - # Local accumulator, placed in RSRAM C_local = T.alloc_shared((block_M, block_N), accum_dtype) + for bx in T.serial(T.ceildiv(sharded_M, block_M)): for by in T.serial(T.ceildiv(sharded_N, block_N)): T.clear(C_local) + K_steps = T.ceildiv(K, block_K) - # Number of iterations in K dimension. - K_steps = T.ceildiv(sharded_K, block_K) - - # Core loop of SUMMA algorithm for k_tile in range(K_steps): - # --- Step 1: Broadcast row block of matrix A --- - # Broadcast directly from DRAM to asram of each core + a_src_col = k_tile % T.mesh_ncols() + b_src_row = k_tile % T.mesh_nrows() + a_local_k = (k_tile // T.mesh_ncols()) * block_K + b_local_k = (k_tile // T.mesh_nrows()) * block_K + T.comm.broadcast( A[ bx * block_M : bx * block_M + block_M, - k_tile * block_K : k_tile * block_K + block_K, + a_local_k : a_local_k + block_K, ], - A_shared, - (0, 0), + A_broadcast, + (core_row, a_src_col), direction="h", ) - - # --- Step 2: Broadcast column block of matrix B --- - # Broadcast directly from DRAM to wsram of each core + T.copy(A_broadcast, A_shared) T.comm.broadcast( B[ - k_tile * block_K : k_tile * block_K + block_K, + b_local_k : b_local_k + block_K, by * block_N : by * block_N + block_N, ], B_shared, - (0, 0), + (b_src_row, core_col), direction="v", ) - - # --- Step 3: Local computation --- - # Each core performs local GEMM using broadcasted A_shared and B_shared T.gemm(A_shared, B_shared, C_local) - # After the loop ends, write local computation result back to DRAM T.copy(C_local, C[bx * block_M, by * block_N]) return kernel @@ -90,35 +83,18 @@ def kernel( def test_summa(is_log=False): func = summa_matmul(128, 128, 128, 32, 32, 32) - script_device_mode = """ - with T.launch_thread("blockIdx.x", 16) as bx: - T.barrier_init(T.int64(15)) - T.barrier_init(T.int64(4369)) - with T.decl_buffer((32, 32), "float16", data=A_shared.data, scope="shared.asram") as A_shared: - B_shared = T.decl_buffer((32, 32), "float16", data=B_shared.data, scope="shared.wsram") - C_local = T.decl_buffer((32, 32), data=C_local.data, scope="shared.rsram") - A_rsram_stage = T.decl_buffer((32, 32), "float16", data=A_rsram_stage.data, scope="shared.rsram") - for i0 in T.serial(8, annotations={"tile.domain": [32, 32], "tile.execution_axis": 0, "tile.execution_domain_axes": [0, 1], "tile.scope_entry": 1, "tile.tile_size": [4, 32]}): - for i1 in T.serial(1, annotations={"tile.execution_axis": 1}): - for ki in T.serial(4, annotations={"tile.interior": 1, "tile.interior_axis": 0}): - for kj in T.vectorized(32, annotations={"tile.interior": 1, "tile.interior_axis": 1}): - C_local[i0 * 4 + ki, kj] = T.float32(0.0) - T.dma_copy(T.region(A_1[0, 0], 1, 32, 32), T.region(A_rsram_stage[0, 0], 2, 32, 32), 0, T.sync_token_id(0)) - T.wait_token(0) - T.barrier_arrive_and_wait(T.int64(15)) - T.broadcast_(T.region(A_rsram_stage[0, 0], 1, 32, 32), T.region(A_shared[0, 0], 2, 32, 32), 0, 15, 0, 0, T.sync_token_id(1)) - T.barrier_arrive_and_wait(T.int64(4369)) - T.broadcast_(T.region(B_1[0, 0], 1, 32, 32), T.region(B_shared[0, 0], 2, 32, 32), 1, 15, 0, 0, T.sync_token_id(2)) - T.wait_token(1) - T.barrier_arrive_and_wait(T.int64(15)) - T.wait_token(2) - T.barrier_arrive_and_wait(T.int64(4369)) - T.mma_sunmmio(T.region(A_shared[0, 0], 1, 32, 32), T.region(B_shared[0, 0], 1, 32, 32), T.region(C_local[0, 0], 3, 32, 32), T.bool(False), T.bool(False), T.bool(False), 0, T.sync_token_id(3)) - T.wait_token(3) - T.dma_copy(T.region(C_local[0, 0], 1, 32, 32), T.region(C_1[0, 0], 2, 32, 32), 0, T.sync_token_id(4)) - T.wait_token(4) - return 0 - """ + script_device_mode = [ + 'with T.launch_thread("blockIdx.x", 16) as bx:', + 'with T.decl_buffer((32, 32), "float16", data=A_broadcast.data, scope="shared.rsram") as A_broadcast:', + 'A_shared = T.decl_buffer((32, 32), "float16", data=A_shared.data, scope="shared.asram")', + 'B_shared = T.decl_buffer((32, 32), "float16", data=B_shared.data, scope="shared.wsram")', + "for k_tile in range(4):", + "bx // 4 * 4 + k_tile", + "k_tile * 4 + bx % 4", + "T.mma_sunmmio(", + "T.wait_token(", + "T.dma_copy(T.region(C_local[0, 0]", + ] script_lower_tile_op = [ 'A = T.match_buffer(A_handle, (32, 32), "float16", strides=(32, 1))', @@ -126,30 +102,25 @@ def test_summa(is_log=False): "C = T.match_buffer(C_handle, (32, 32), strides=(32, 1))", 'bx = T.launch_thread("blockIdx.x", 16)', "for bx_1, by in T.grid(1, 1):", - "T.dma_copy(T.region(A[0, 0], 1, 32, 32), T.region(A_rsram_stage[0, 0], 2, 32, 32), 0)", - "T.broadcast_(T.region(A_rsram_stage[0, 0], 1, 32, 32), T.region(A_shared[0, 0], 2, 32, 32), 0, T.int64(15), 0, 0)", - "T.broadcast_(T.region(B[0, 0], 1, 32, 32), T.region(B_shared[0, 0], 2, 32, 32), 1, T.int64(15), 0, 0)", - "T.dma_copy(T.region(C_local[0, 0], 1, 32, 32), T.region(C[0, 0], 2, 32, 32), 0)", + "for k_tile in range(4):", + "T.dma_copy(T.region(A[bx_1 * 32, 0], 1, 32, 32), T.region(A_rsram_stage[0, 0], 2, 32, 32), 0)", + "bx // 4 * 4 + k_tile", + "k_tile * 4 + bx % 4", + "T.mma_sunmmio(", + "T.dma_copy(T.region(C_local[0, 0], 1, 32, 32), T.region(C[bx_1 * 32, by * 32], 2, 32, 32), 0)", ] script_InjectSunmmioSync = [ 'with T.launch_thread("blockIdx.x", 16) as bx:', - "T.dma_copy(T.region(A_1[0, 0], 1, 32, 32), T.region(A_rsram_stage[0, 0], 2, 32, 32), 0, T.sync_token_id(0))", - "T.wait_token(0)", - "T.barrier_init(T.int64(15))", - "T.barrier_init(T.int64(4369))", - "T.barrier_arrive_and_wait(T.int64(15))", - "T.broadcast_(T.region(A_rsram_stage[0, 0], 1, 32, 32), T.region(A_shared[0, 0], 2, 32, 32), 0, 15, 0, 0, T.sync_token_id(1))", - "T.barrier_arrive_and_wait(T.int64(4369))", - "T.broadcast_(T.region(B_1[0, 0], 1, 32, 32), T.region(B_shared[0, 0], 2, 32, 32), 1, 15, 0, 0, T.sync_token_id(2))", - "T.wait_token(1)", - "T.barrier_arrive_and_wait(T.int64(15))", - "T.wait_token(2)", - "T.barrier_arrive_and_wait(T.int64(4369))", - "T.mma_sunmmio(T.region(A_shared[0, 0], 1, 32, 32), T.region(B_shared[0, 0], 1, 32, 32), T.region(C_local[0, 0], 3, 32, 32), T.bool(False), T.bool(False), T.bool(False), 0, T.sync_token_id(3))", - "T.wait_token(3)", - "T.dma_copy(T.region(C_local[0, 0], 1, 32, 32), T.region(C_1[0, 0], 2, 32, 32), 0, T.sync_token_id(4))", - "T.wait_token(4)", + "T.barrier_init(", + "for k_tile in range(4):", + "bx // 4 * 4 + k_tile", + "k_tile * 4 + bx % 4", + "T.broadcast_(", + "T.mma_sunmmio(", + "T.sync_token_id(", + "T.wait_token(", + "T.dma_copy(T.region(C_local[0, 0]", ] test_config = { diff --git a/testing/python/sunmmio/pipeline/test_sync.py b/testing/python/sunmmio/pipeline/test_sync.py index caf127bae4..f52d4bf288 100644 --- a/testing/python/sunmmio/pipeline/test_sync.py +++ b/testing/python/sunmmio/pipeline/test_sync.py @@ -9,7 +9,7 @@ @target("Sunmmio") def kernel_sync(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="float"): - shard_policy = T.MeshShardingPolicy(y=0, x=1) + shard_policy = T.placement.full_shard(0, 1) A_shape = (M, K) B_shape = (M, K) diff --git a/testing/python/sunmmio/target/test_target.py b/testing/python/sunmmio/target/test_target.py index 53a2e72f47..33f14567e0 100644 --- a/testing/python/sunmmio/target/test_target.py +++ b/testing/python/sunmmio/target/test_target.py @@ -32,7 +32,7 @@ def get_current_target(): def test_sunmmio_target_binding(): def example_tensor_annot(shape): - MyTensor = T.MeshTensor(shape, T.MeshShardingPolicy(y=0, x=1), dtype="float32") + MyTensor = T.MeshTensor(shape, T.placement.full_shard(0, 1), dtype="float32") @T.prim_func def kernel(A: MyTensor): diff --git a/testing/python/sunmmio/transform/test_global_layout_utils.py b/testing/python/sunmmio/transform/test_global_layout_utils.py index 8bc276126e..3b5a37b6c7 100644 --- a/testing/python/sunmmio/transform/test_global_layout_utils.py +++ b/testing/python/sunmmio/transform/test_global_layout_utils.py @@ -10,7 +10,6 @@ from tilelang.utils.target import determine_target, SUNMMIO_TARGET_DESC, target_is_sunmmio import tilelang as tl import tilelang.language as T -from tilelang.language.mesh_tensor import MeshShardingPolicy, MeshReplicationType from tilelang.layout import make_row_major from tvm import tir from tvm.tir import PyStmtExprVisitor @@ -64,7 +63,7 @@ def test_global_buffer_layout_populated_for_sunmmio(): Test that global buffer layouts from tensor_meta are populated into layout_map during SunmmioLayoutInference pass for Sunmmio target. """ - policy = MeshShardingPolicy(y=0, x=1, replicate=MeshReplicationType.NONE) + policy = T.placement.full_shard(0, 1) M, N, K = 64, 64, 64 block_M, block_N, block_K = 32, 32, 32 @@ -155,7 +154,7 @@ def test_row_major_global_layout_values(): """ Test that the layout created from tensor_meta produces correct forward index mapping. """ - policy = MeshShardingPolicy(y=0, x=1, replicate=MeshReplicationType.NONE) + policy = T.placement.full_shard(0, 1) M, N, K = 64, 64, 64 block_M, block_N, block_K = 32, 32, 32 @@ -219,7 +218,7 @@ def test_dynamic_shape_global_buffer_layout(): block_M, block_N, block_K = 32, 32, 32 - policy = MeshShardingPolicy(replicate=MeshReplicationType.ALL) + policy = T.placement.replicated() A_tensor = T.MeshTensor((M_var, K_var), policy, dtype="float16") B_tensor = T.MeshTensor((K_var, 64), policy, dtype="float16") diff --git a/testing/python/sunmmio/transform/test_infer_sram_scope.py b/testing/python/sunmmio/transform/test_infer_sram_scope.py index b882474036..2cca8a6353 100644 --- a/testing/python/sunmmio/transform/test_infer_sram_scope.py +++ b/testing/python/sunmmio/transform/test_infer_sram_scope.py @@ -9,7 +9,6 @@ from tvm.tir import BufferLoad, BufferStore, Buffer, Block, Call from typing import Set from tilelang.tileview import make_tileview -from tilelang.language.mesh_tensor import MeshShardingPolicy tilelang.env.disable_cache() @@ -793,19 +792,19 @@ def matmul(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="flo def gemm( A: T.MeshTensor( (M, K), - sharding_policy=MeshShardingPolicy(cross_mesh_dim=0), + placement=T.placement.mesh_as_line(0), dtype=dtype, layout=A_layout, ), B: T.MeshTensor( (K, N), - sharding_policy=MeshShardingPolicy(cross_mesh_dim=0), + placement=T.placement.mesh_as_line(0), dtype=dtype, layout=B_layout, ), C: T.MeshTensor( (M, N), - sharding_policy=MeshShardingPolicy(cross_mesh_dim=0), + placement=T.placement.mesh_as_line(0), dtype=accum_dtype, layout=C_layout, ), @@ -888,19 +887,19 @@ def auto_insert_copy_matmul(M, N, K, block_M, block_N, block_K, dtype="float16", def gemm( A: T.MeshTensor( (M, K), - sharding_policy=MeshShardingPolicy(cross_mesh_dim=0), + placement=T.placement.mesh_as_line(0), dtype=dtype, layout=A_layout, ), B: T.MeshTensor( (K, N), - sharding_policy=MeshShardingPolicy(cross_mesh_dim=0), + placement=T.placement.mesh_as_line(0), dtype=dtype, layout=B_layout, ), C: T.MeshTensor( (M, N), - sharding_policy=MeshShardingPolicy(cross_mesh_dim=0), + placement=T.placement.mesh_as_line(0), dtype=accum_dtype, layout=C_layout, ), @@ -964,19 +963,19 @@ def sliced_conflict_matmul(dtype="float16", accum_dtype="float"): def gemm( A: T.MeshTensor( (128, 128), - sharding_policy=MeshShardingPolicy(cross_mesh_dim=0), + placement=T.placement.mesh_as_line(0), dtype=dtype, layout=A_layout, ), B: T.MeshTensor( (128, 128), - sharding_policy=MeshShardingPolicy(cross_mesh_dim=0), + placement=T.placement.mesh_as_line(0), dtype=dtype, layout=B_layout, ), C: T.MeshTensor( (128, 128), - sharding_policy=MeshShardingPolicy(cross_mesh_dim=0), + placement=T.placement.mesh_as_line(0), dtype=accum_dtype, layout=C_layout, ), diff --git a/testing/python/sunmmio/transform/test_inject_sync.py b/testing/python/sunmmio/transform/test_inject_sync.py index 5118fc1080..227ffd823b 100644 --- a/testing/python/sunmmio/transform/test_inject_sync.py +++ b/testing/python/sunmmio/transform/test_inject_sync.py @@ -530,6 +530,499 @@ def _make_while_async_to_sync_store_mod(target): return tir.transform.BindTarget(target)(mod) +def _make_loop_carried_dma_wait_domain_mod(target): + global_a_data = _pointer_var("global_a", scope="global") + global_b_data = _pointer_var("global_b", scope="global") + stage_0_data = _pointer_var("stage_0") + stage_1_data = _pointer_var("stage_1") + a_ping_data = _pointer_var("a_ping", scope="shared.asram") + a_pong_data = _pointer_var("a_pong", scope="shared.asram") + b_ping_data = _pointer_var("b_ping", scope="shared.wsram") + b_pong_data = _pointer_var("b_pong", scope="shared.wsram") + accum_data = _pointer_var("accum") + + def make_buffer(data, name, scope): + return tir.decl_buffer( + (32, 32), + "float16", + name=name, + data=data, + scope=scope, + ) + + global_a = make_buffer(global_a_data, "global_a", "global") + global_b = make_buffer(global_b_data, "global_b", "global") + stage_0 = make_buffer(stage_0_data, "stage_0", "shared.rsram") + stage_1 = make_buffer(stage_1_data, "stage_1", "shared.rsram") + a_ping = make_buffer(a_ping_data, "a_ping", "shared.asram") + a_pong = make_buffer(a_pong_data, "a_pong", "shared.asram") + b_ping = make_buffer(b_ping_data, "b_ping", "shared.wsram") + b_pong = make_buffer(b_pong_data, "b_pong", "shared.wsram") + accum = make_buffer(accum_data, "accum", "shared.rsram") + + def dma(src, dst): + return tir.Evaluate( + tir.call_intrin( + "handle", + tir.op.Op.get("tl.dma_copy"), + _region(src, 1), + _region(dst, 2), + tir.IntImm("int32", 0), + ) + ) + + def mma(a, b): + return tir.Evaluate( + tir.call_intrin( + "handle", + tir.op.Op.get("tl.mma_sunmmio"), + _region(a, 1), + _region(b, 1), + _region(accum, 3), + tir.IntImm("bool", 0), + tir.IntImm("bool", 0), + tir.IntImm("bool", 0), + tir.IntImm("int32", 0), + ) + ) + + i = tir.Var("i", "int32") + loop = tir.For( + i, + tir.IntImm("int32", 0), + tir.IntImm("int32", 4), + tir.ForKind.SERIAL, + tir.SeqStmt( + [ + dma(stage_0, a_ping), + dma(global_a, stage_1), + mma(a_pong, b_ping), + dma(stage_1, a_pong), + dma(global_b, b_ping), + mma(a_ping, b_pong), + ] + ), + ) + + buffers = [ + global_a, + global_b, + stage_0, + stage_1, + a_ping, + a_pong, + b_ping, + b_pong, + accum, + ] + body = loop + for buffer in reversed(buffers): + body = tir.DeclBuffer(buffer, body) + params = [ + global_a_data, + global_b_data, + stage_0_data, + stage_1_data, + a_ping_data, + a_pong_data, + b_ping_data, + b_pong_data, + accum_data, + ] + func = tir.PrimFunc(params, body) + func = func.with_attr("global_symbol", "main") + func = func.with_attr("tir.is_global_func", True) + mod = tvm.IRModule({"main": func}) + return tir.transform.BindTarget(target)(mod) + + +def _make_loop_exit_wait_placement_mod(target, same_domain, conditional_submit=False, submit_wrapper=None): + global_b_data = _pointer_var("global_b", scope="global") + stage_0_data = _pointer_var("stage_0") + stage_1_data = _pointer_var("stage_1") + a_ping_data = _pointer_var("a_ping", scope="shared.asram") + a_pong_data = _pointer_var("a_pong", scope="shared.asram") + b_ping_data = _pointer_var("b_ping", scope="shared.wsram") + accum_data = _pointer_var("accum") + condition = tir.Var("condition", "bool") + + def make_buffer(data, name, scope): + return tir.decl_buffer( + (32, 32), + "float16", + name=name, + data=data, + scope=scope, + ) + + global_b = make_buffer(global_b_data, "global_b", "global") + stage_0 = make_buffer(stage_0_data, "stage_0", "shared.rsram") + stage_1 = make_buffer(stage_1_data, "stage_1", "shared.rsram") + a_ping = make_buffer(a_ping_data, "a_ping", "shared.asram") + a_pong = make_buffer(a_pong_data, "a_pong", "shared.asram") + b_ping = make_buffer(b_ping_data, "b_ping", "shared.wsram") + accum = make_buffer(accum_data, "accum", "shared.rsram") + + def dma(src, dst): + return tir.Evaluate( + tir.call_intrin( + "handle", + tir.op.Op.get("tl.dma_copy"), + _region(src, 1), + _region(dst, 2), + tir.IntImm("int32", 0), + ) + ) + + mma = tir.Evaluate( + tir.call_intrin( + "handle", + tir.op.Op.get("tl.mma_sunmmio"), + _region(a_pong, 1), + _region(b_ping, 1), + _region(accum, 3), + tir.IntImm("bool", 0), + tir.IntImm("bool", 0), + tir.IntImm("bool", 0), + tir.IntImm("int32", 0), + ) + ) + + i = tir.Var("i", "int32") + loop = tir.For( + i, + tir.IntImm("int32", 0), + tir.IntImm("int32", 4), + tir.ForKind.SERIAL, + dma(stage_0, a_pong), + ) + epilogue_dma = dma(stage_1, a_ping) if same_domain else dma(global_b, b_ping) + if conditional_submit: + epilogue_dma = tir.IfThenElse(condition, epilogue_dma, None) + if submit_wrapper == "attr": + epilogue_dma = tir.AttrStmt( + stage_1_data, + "test_transparent_wrapper", + tir.IntImm("int32", 1), + epilogue_dma, + ) + elif submit_wrapper == "let": + wrapper_value = tir.Var("wrapper_value", "int32") + epilogue_dma = tir.LetStmt(wrapper_value, tir.IntImm("int32", 0), epilogue_dma) + elif submit_wrapper == "decl_buffer": + epilogue_dma = tir.DeclBuffer(stage_1, epilogue_dma) + elif submit_wrapper == "allocate": + wrapper_data = _pointer_var("wrapper_alloc") + epilogue_dma = tir.Allocate( + wrapper_data, + "float16", + [tir.IntImm("int32", 1)], + tir.IntImm("bool", 1), + epilogue_dma, + ) + elif submit_wrapper == "buffer_realize": + bounds = [tvm.ir.Range.from_min_extent(0, extent) for extent in stage_1.shape] + epilogue_dma = tir.BufferRealize(stage_1, bounds, tir.IntImm("bool", 1), epilogue_dma) + elif submit_wrapper == "block_realize": + block = tir.Block([], [], [], "wrapped_submit", epilogue_dma) + epilogue_dma = tir.BlockRealize([], tir.IntImm("bool", 1), block) + elif submit_wrapper is not None: + raise ValueError(f"Unsupported submit wrapper: {submit_wrapper}") + + buffers = [global_b, stage_0, stage_1, a_ping, a_pong, b_ping, accum] + body = tir.SeqStmt([loop, epilogue_dma, mma]) + for buffer in reversed(buffers): + body = tir.DeclBuffer(buffer, body) + params = [ + global_b_data, + stage_0_data, + stage_1_data, + a_ping_data, + a_pong_data, + b_ping_data, + accum_data, + ] + if conditional_submit: + params.append(condition) + func = tir.PrimFunc(params, body) + func = func.with_attr("global_symbol", "main") + func = func.with_attr("tir.is_global_func", True) + mod = tvm.IRModule({"main": func}) + return tir.transform.BindTarget(target)(mod) + + +def _make_loop_exit_engine_wait_placement_mod(target, engine, epilogue_engine=None): + epilogue_engine = epilogue_engine or engine + + def make_buffer(data, name, scope): + return tir.decl_buffer( + (32, 32), + "float16", + name=name, + data=data, + scope=scope, + ) + + def make_mma(a, b, accum): + return tir.Evaluate( + tir.call_intrin( + "handle", + tir.op.Op.get("tl.mma_sunmmio"), + _region(a, 1), + _region(b, 1), + _region(accum, 3), + tir.IntImm("bool", 0), + tir.IntImm("bool", 0), + tir.IntImm("bool", 0), + tir.IntImm("int32", 0), + ) + ) + + def make_broadcast(src, dst, direction): + return tir.Evaluate( + tir.call_intrin( + "handle", + tir.op.Op.get("tl.broadcast_"), + _region(src, 1), + _region(dst, 2), + tir.IntImm("int32", direction), + tir.IntImm("int64", 15), + tir.IntImm("int32", 0), + tir.IntImm("int32", 0), + ) + ) + + def make_transpose(src, dst): + return tir.Evaluate( + tir.call_intrin( + "handle", + tir.op.Op.get("tl.sunmmio_transpose"), + _region(src, 1), + _region(dst, 2), + ) + ) + + if engine == "tc": + assert epilogue_engine == "tc" + specs = [ + ("a0", "shared.asram"), + ("b0", "shared.wsram"), + ("accum0", "shared.rsram"), + ("a1", "shared.asram"), + ("b1", "shared.wsram"), + ("accum1", "shared.rsram"), + ] + data = [_pointer_var(name, scope=scope) for name, scope in specs] + buffers = [make_buffer(value, name, scope) for value, (name, scope) in zip(data, specs)] + first_async = make_mma(buffers[0], buffers[1], buffers[2]) + second_async = make_mma(buffers[3], buffers[4], buffers[5]) + consumed = buffers[2] + elif engine == "transpose": + assert epilogue_engine == "transpose" + specs = [ + ("src0", "shared.rsram"), + ("dst0", "shared.rsram"), + ("src1", "shared.rsram"), + ("dst1", "shared.rsram"), + ] + data = [_pointer_var(name, dtype="float32", scope=scope) for name, scope in specs] + buffers = [ + tir.decl_buffer( + (32, 32), + "float32", + name=name, + data=value, + scope=scope, + ) + for value, (name, scope) in zip(data, specs) + ] + first_async = make_transpose(buffers[0], buffers[1]) + second_async = make_transpose(buffers[2], buffers[3]) + consumed = buffers[1] + else: + assert engine in ("hlink", "vlink") + assert epilogue_engine in ("hlink", "vlink") + specs = [ + ("src0", "shared.rsram"), + ("dst0", "shared.rsram"), + ("src1", "shared.rsram"), + ("dst1", "shared.rsram"), + ] + data = [_pointer_var(name, scope=scope) for name, scope in specs] + buffers = [make_buffer(value, name, scope) for value, (name, scope) in zip(data, specs)] + first_async = make_broadcast(buffers[0], buffers[1], 0 if engine == "hlink" else 1) + second_async = make_broadcast(buffers[2], buffers[3], 0 if epilogue_engine == "hlink" else 1) + consumed = buffers[1] + + zero = tir.IntImm("int32", 0) + consume = tir.BufferStore( + consumed, + tir.BufferLoad(consumed, [zero, zero]), + [zero, zero], + ) + i = tir.Var("i", "int32") + loop = tir.For( + i, + zero, + tir.IntImm("int32", 4), + tir.ForKind.SERIAL, + first_async, + ) + + body = tir.SeqStmt([loop, second_async, consume]) + for buffer in reversed(buffers): + body = tir.DeclBuffer(buffer, body) + func = tir.PrimFunc(data, body) + func = func.with_attr("global_symbol", "main") + func = func.with_attr("tir.is_global_func", True) + mod = tvm.IRModule({"main": func}) + return tir.transform.BindTarget(target)(mod) + + +def _make_loop_exit_odma_link_alias_mod(target, odma_domain, link_first): + assert odma_domain in ("odma0", "odma1") + + specs = [ + ("global_src", "global"), + ("odma_src", "shared.rsram"), + ("odma_dst", "shared.rsram"), + ("link_src", "shared.rsram"), + ("link_dst", "shared.rsram"), + ] + data = [_pointer_var(name, scope=scope) for name, scope in specs] + buffers = [ + tir.decl_buffer( + (32, 32), + "float16", + name=name, + data=value, + scope=scope, + ) + for value, (name, scope) in zip(data, specs) + ] + global_src, odma_src, odma_dst, link_src, link_dst = buffers + + if odma_domain == "odma0": + odma_async = tir.Evaluate( + tir.call_intrin( + "handle", + tir.op.Op.get("tl.dma_copy"), + _region(global_src, 1), + _region(odma_dst, 2), + tir.IntImm("int32", 0), + ) + ) + link_direction = 1 + else: + odma_async = tir.Evaluate( + tir.call_intrin( + "handle", + tir.op.Op.get("tl.sunmmio_layout_transform"), + _region(odma_src, 1), + _region(odma_dst, 2), + ) + ) + link_direction = 0 + + link_async = tir.Evaluate( + tir.call_intrin( + "handle", + tir.op.Op.get("tl.broadcast_"), + _region(link_src, 1), + _region(link_dst, 2), + tir.IntImm("int32", link_direction), + tir.IntImm("int64", 15), + tir.IntImm("int32", 0), + tir.IntImm("int32", 0), + ) + ) + + first_async = link_async if link_first else odma_async + second_async = odma_async if link_first else link_async + consumed = link_dst if link_first else odma_dst + zero = tir.IntImm("int32", 0) + consume = tir.BufferStore( + consumed, + tir.BufferLoad(consumed, [zero, zero]), + [zero, zero], + ) + i = tir.Var("i", "int32") + loop = tir.For( + i, + zero, + tir.IntImm("int32", 4), + tir.ForKind.SERIAL, + first_async, + ) + + body = tir.SeqStmt([loop, second_async, consume]) + for buffer in reversed(buffers): + body = tir.DeclBuffer(buffer, body) + func = tir.PrimFunc(data, body) + func = func.with_attr("global_symbol", "main") + func = func.with_attr("tir.is_global_func", True) + mod = tvm.IRModule({"main": func}) + return tir.transform.BindTarget(target)(mod) + + +def _make_outer_loop_token_consumed_in_inner_loop_mod(target): + stage_data = _pointer_var("stage") + dst_data = _pointer_var("dst", scope="shared.asram") + stage = tir.decl_buffer( + (32, 32), + "float16", + name="stage", + data=stage_data, + scope="shared.rsram", + ) + dst = tir.decl_buffer( + (32, 32), + "float16", + name="dst", + data=dst_data, + scope="shared.asram", + ) + + zero = tir.IntImm("int32", 0) + consume = tir.BufferStore( + dst, + tir.BufferLoad(dst, [zero, zero]), + [zero, zero], + ) + produce = tir.Evaluate( + tir.call_intrin( + "handle", + tir.op.Op.get("tl.dma_copy"), + _region(stage, 1), + _region(dst, 2), + zero, + ) + ) + + i = tir.Var("i", "int32") + j = tir.Var("j", "int32") + inner = tir.For( + j, + zero, + tir.IntImm("int32", 2), + tir.ForKind.SERIAL, + consume, + ) + outer = tir.For( + i, + zero, + tir.IntImm("int32", 4), + tir.ForKind.SERIAL, + tir.SeqStmt([inner, produce]), + ) + body = tir.DeclBuffer(stage, tir.DeclBuffer(dst, outer)) + func = tir.PrimFunc([stage_data, dst_data], body) + func = func.with_attr("global_symbol", "main") + func = func.with_attr("tir.is_global_func", True) + mod = tvm.IRModule({"main": func}) + return tir.transform.BindTarget(target)(mod) + + def _make_dma_to_loop_let_consumer_mod(target): src_data = _pointer_var("src", scope="global") stage_data = _pointer_var("stage") @@ -1263,6 +1756,7 @@ def test_inject_sunmmio_sync_loop_missing_wait_before_token_reuse(): wait_idx = next(idx for idx, line in enumerate(lines) if for_idx < idx < transform_idx and f"wait_token({token})" in line) assert null_idx < for_idx < wait_idx < transform_idx + assert lines[wait_idx - 1] == "if i > 0:" def test_inject_sunmmio_sync_nested_loop_reuses_tokens_without_mixing_levels(): @@ -1303,6 +1797,18 @@ def extract_call_id(line, marker): assert sync_ids.issubset(wait_ids) +def test_inject_sunmmio_sync_outer_carried_wait_inside_inner_loop_uses_outer_guard(): + target = get_target("Sunmmio") + mod = _make_outer_loop_token_consumed_in_inner_loop_mod(target) + + mod = tilelang.transform.InjectSunmmioSync()(mod) + lines = [line.strip() for line in mod.script(show_meta=True).splitlines()] + + wait_idx = next(idx for idx, line in enumerate(lines) if line == "T.wait_token(0)") + assert lines[wait_idx - 1] == "if i > 0:" + assert "if j > 0:" not in lines + + def test_inject_sunmmio_sync_if(): def kernel(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="float32"): @T.prim_func @@ -1497,6 +2003,181 @@ def extract_call_id(line, marker): assert not barrier_wait_after_second +def test_inject_sunmmio_sync_loop_carried_wait_precedes_same_domain_submit(): + target = get_target("Sunmmio") + mod = _make_loop_carried_dma_wait_domain_mod(target) + + mod = tilelang.transform.InjectSunmmioSync()(mod) + lines = mod.script(show_meta=True).splitlines() + + dma_entries = [ + (idx, _extract_call_id(line, "sync_token_id")) for idx, line in enumerate(lines) if "dma_copy" in line and "sync_token_id(" in line + ] + mma_entries = [ + (idx, _extract_call_id(line, "sync_token_id")) + for idx, line in enumerate(lines) + if "mma_sunmmio" in line and "sync_token_id(" in line + ] + wait_entries = [(idx, _extract_call_id(line, "wait_token")) for idx, line in enumerate(lines) if "wait_token(" in line] + + assert [token for _, token in dma_entries] == [0, 1, 3, 4] + assert [token for _, token in mma_entries] == [2, 5] + + first_odma1_idx = dma_entries[0][0] + first_odma0_idx = dma_entries[1][0] + first_mma_idx = mma_entries[0][0] + carried_odma1_wait_idx = next(idx for idx, token in wait_entries if token == 3) + carried_odma0_wait_idx = next(idx for idx, token in wait_entries if token == 4) + carried_tc_wait_idx = next(idx for idx, token in wait_entries if token == 5) + + assert carried_odma1_wait_idx < first_odma1_idx + assert carried_tc_wait_idx < first_odma1_idx + assert first_odma1_idx < carried_odma0_wait_idx < first_odma0_idx + assert first_odma0_idx < first_mma_idx + for wait_idx in ( + carried_odma1_wait_idx, + carried_odma0_wait_idx, + carried_tc_wait_idx, + ): + assert lines[wait_idx - 1].strip() == "if i > 0:" + + # These waits consume tokens generated earlier in the same iteration and + # must remain unconditional. + for token in (0, 1, 2): + wait_idx = next(idx for idx, wait_token in wait_entries if wait_token == token) + assert lines[wait_idx - 1].strip() != "if i > 0:" + + +def _check_loop_exit_wait_placement(same_domain, submit_wrapper=None): + target = get_target("Sunmmio") + mod = _make_loop_exit_wait_placement_mod(target, same_domain, submit_wrapper=submit_wrapper) + + mod = tilelang.transform.InjectSunmmioSync()(mod) + lines = mod.script(show_meta=True).splitlines() + dma_entries = [ + (idx, _extract_call_id(line, "sync_token_id")) for idx, line in enumerate(lines) if "dma_copy" in line and "sync_token_id(" in line + ] + mma_idx = next(idx for idx, line in enumerate(lines) if "mma_sunmmio" in line and "sync_token_id(" in line) + exit_wait_idx = next(idx for idx, line in enumerate(lines) if idx > dma_entries[0][0] and "wait_token(0)" in line) + + assert [token for _, token in dma_entries] == [0, 1] + epilogue_dma_idx = dma_entries[1][0] + if same_domain: + assert exit_wait_idx < epilogue_dma_idx < mma_idx + else: + assert epilogue_dma_idx < exit_wait_idx < mma_idx + + +def test_inject_sunmmio_sync_moves_loop_exit_wait_before_same_domain_dma(): + _check_loop_exit_wait_placement(True) + + +def test_inject_sunmmio_sync_keeps_loop_exit_wait_after_other_domain_dma(): + _check_loop_exit_wait_placement(False) + + +def test_inject_sunmmio_sync_finds_submit_in_transparent_wrapper(): + for wrapper in ( + "attr", + "let", + "decl_buffer", + "allocate", + "buffer_realize", + "block_realize", + ): + _check_loop_exit_wait_placement(True, submit_wrapper=wrapper) + + +def _check_loop_exit_engine_wait_placement(engine, epilogue_engine=None, expect_move=True): + target = get_target("Sunmmio") + mod = _make_loop_exit_engine_wait_placement_mod(target, engine, epilogue_engine) + + mod = tilelang.transform.InjectSunmmioSync()(mod) + lines = mod.script(show_meta=True).splitlines() + marker = { + "tc": "mma_sunmmio", + "transpose": "sunmmio_transpose", + }.get(engine, "broadcast_") + async_entries = [ + (idx, _extract_call_id(line, "sync_token_id")) for idx, line in enumerate(lines) if marker in line and "sync_token_id(" in line + ] + assert [token for _, token in async_entries] == [0, 1] + + first_async_idx = async_entries[0][0] + second_async_idx = async_entries[1][0] + exit_wait_idx = next(idx for idx, line in enumerate(lines) if idx > first_async_idx and "wait_token(0)" in line) + if expect_move: + assert exit_wait_idx < second_async_idx + else: + assert second_async_idx < exit_wait_idx + + +def test_inject_sunmmio_sync_moves_loop_exit_wait_before_same_domain_tc(): + _check_loop_exit_engine_wait_placement("tc") + + +def test_inject_sunmmio_sync_moves_loop_exit_wait_before_same_domain_transpose(): + _check_loop_exit_engine_wait_placement("transpose") + + +def test_inject_sunmmio_sync_moves_loop_exit_wait_before_same_domain_hlink(): + _check_loop_exit_engine_wait_placement("hlink") + + +def test_inject_sunmmio_sync_moves_loop_exit_wait_before_same_domain_vlink(): + _check_loop_exit_engine_wait_placement("vlink") + + +def test_inject_sunmmio_sync_keeps_hlink_wait_after_vlink_submit(): + _check_loop_exit_engine_wait_placement("hlink", epilogue_engine="vlink", expect_move=False) + + +def test_inject_sunmmio_sync_treats_link_as_its_odma_submission_domain(): + target = get_target("Sunmmio") + for odma_domain in ("odma0", "odma1"): + for link_first in (False, True): + mod = _make_loop_exit_odma_link_alias_mod(target, odma_domain, link_first) + + mod = tilelang.transform.InjectSunmmioSync()(mod) + lines = mod.script(show_meta=True).splitlines() + first_async_idx = next(idx for idx, line in enumerate(lines) if "sync_token_id(0)" in line) + second_async_idx = next(idx for idx, line in enumerate(lines) if "sync_token_id(1)" in line) + exit_wait_idx = next(idx for idx, line in enumerate(lines) if idx > first_async_idx and "wait_token(0)" in line) + + assert first_async_idx < exit_wait_idx < second_async_idx + + +def test_inject_sunmmio_sync_does_not_move_wait_across_conditional_dma(): + target = get_target("Sunmmio") + mod = _make_loop_exit_wait_placement_mod(target, same_domain=True, conditional_submit=True) + + mod = tilelang.transform.InjectSunmmioSync()(mod) + lines = mod.script(show_meta=True).splitlines() + dma_indices = [idx for idx, line in enumerate(lines) if "dma_copy" in line and "sync_token_id(" in line] + exit_wait_idx = next(idx for idx, line in enumerate(lines) if idx > dma_indices[0] and "wait_token(0)" in line) + mma_idx = next(idx for idx, line in enumerate(lines) if "mma_sunmmio" in line and "sync_token_id(" in line) + + assert dma_indices[1] < exit_wait_idx < mma_idx + + +def test_inject_sunmmio_sync_does_not_treat_wrapped_conditional_dma_as_anchor(): + target = get_target("Sunmmio") + mod = _make_loop_exit_wait_placement_mod( + target, + same_domain=True, + conditional_submit=True, + submit_wrapper="attr", + ) + + mod = tilelang.transform.InjectSunmmioSync()(mod) + lines = mod.script(show_meta=True).splitlines() + dma_indices = [idx for idx, line in enumerate(lines) if "dma_copy" in line and "sync_token_id(" in line] + exit_wait_idx = next(idx for idx, line in enumerate(lines) if idx > dma_indices[0] and "wait_token(0)" in line) + mma_idx = next(idx for idx, line in enumerate(lines) if "mma_sunmmio" in line and "sync_token_id(" in line) + + assert dma_indices[1] < exit_wait_idx < mma_idx + + def test_inject_sunmmio_sync_while_loop_carried_tokens(): target = get_target("Sunmmio") mod = _make_while_pair_broadcast_mod(target) @@ -1551,6 +2232,8 @@ def extract_call_id(line, marker): wait_first_between = [idx for idx, _, token in wait_entries if token == first_token and first_bcast_idx < idx < second_bcast_idx] assert carried_wait_before_first assert wait_first_between + assert "sunmmio_has_previous_iteration" in lines[min(carried_wait_before_first) - 1] + assert "tl.local_var_init" in script barrier_wait_before_first = [idx for idx, _, _ in barrier_wait_entries if min(carried_wait_before_first) < idx < first_bcast_idx] barrier_wait_between = [idx for idx, _, _ in barrier_wait_entries if min(wait_first_between) < idx < second_bcast_idx] @@ -1591,6 +2274,8 @@ def extract_call_id(line, marker): carried_wait_before_store = [idx for idx, _, token in wait_entries if token == broadcast_token and while_idx < idx < store_idx] assert carried_wait_before_store assert min(carried_wait_before_store) < store_idx < broadcast_idx + assert "sunmmio_has_previous_iteration" in lines[min(carried_wait_before_store) - 1] + assert "tl.local_var_init" in script def test_inject_sunmmio_sync_hoists_wait_before_loop_let_consumer(): @@ -1694,6 +2379,7 @@ def test_inject_sunmmio_sync_skips_unsafe_local_var_loop_domain(): test_inject_sunmmio_sync_dynamic_pair_mask_candidates() test_inject_sunmmio_sync_loop_missing_wait_before_token_reuse() test_inject_sunmmio_sync_nested_loop_reuses_tokens_without_mixing_levels() + test_inject_sunmmio_sync_outer_carried_wait_inside_inner_loop_uses_outer_guard() test_inject_sunmmio_sync_if() test_inject_sunmmio_sync_hoists_wait_before_loop_let_consumer() test_inject_sunmmio_sync_hoists_wait_before_loop_if_condition_consumer() @@ -1701,5 +2387,17 @@ def test_inject_sunmmio_sync_skips_unsafe_local_var_loop_domain(): test_inject_sunmmio_sync_waits_for_hidden_buffer_store_index_read() test_inject_sunmmio_sync_waits_for_hidden_async_region_index_read() test_inject_sunmmio_sync_loop() + test_inject_sunmmio_sync_loop_carried_wait_precedes_same_domain_submit() + test_inject_sunmmio_sync_moves_loop_exit_wait_before_same_domain_dma() + test_inject_sunmmio_sync_keeps_loop_exit_wait_after_other_domain_dma() + test_inject_sunmmio_sync_finds_submit_in_transparent_wrapper() + test_inject_sunmmio_sync_moves_loop_exit_wait_before_same_domain_tc() + test_inject_sunmmio_sync_moves_loop_exit_wait_before_same_domain_transpose() + test_inject_sunmmio_sync_moves_loop_exit_wait_before_same_domain_hlink() + test_inject_sunmmio_sync_moves_loop_exit_wait_before_same_domain_vlink() + test_inject_sunmmio_sync_keeps_hlink_wait_after_vlink_submit() + test_inject_sunmmio_sync_treats_link_as_its_odma_submission_domain() + test_inject_sunmmio_sync_does_not_move_wait_across_conditional_dma() + test_inject_sunmmio_sync_does_not_treat_wrapped_conditional_dma_as_anchor() test_inject_sunmmio_sync_while_loop_carried_tokens() test_inject_sunmmio_sync_while_loop_carried_async_to_sync_store() diff --git a/testing/python/sunmmio/transform/test_layout_inference.py b/testing/python/sunmmio/transform/test_layout_inference.py index 45326bfef8..dbee16cbe1 100644 --- a/testing/python/sunmmio/transform/test_layout_inference.py +++ b/testing/python/sunmmio/transform/test_layout_inference.py @@ -980,14 +980,13 @@ def test_immutable_conflict_rsram_annotate_vs_gemm(): def dram_zn_to_asram_kernel(): """DRAM buffer with ZN layout (via MeshTensor) copied to ASRAM → IsZZLike check should fail.""" - from tilelang.language.mesh_tensor import MeshShardingPolicy, MeshReplicationType M, K, N = 64, 32, 64 block_M, block_K, block_N = 64, 32, 64 dtype = "float16" accum_dtype = "float32" - policy = MeshShardingPolicy(y=0, x=1, replicate=MeshReplicationType.NONE) + policy = T.placement.full_shard(0, 1) # ZN layout on DRAM — incompatible with ASRAM (Gemm.A requires ZZ-like) zn_dram = make_zn_layout((M, K), axes=[0, 1], block_shape=[32, 32]) @@ -1034,14 +1033,13 @@ def test_dram_zn_to_asram_succeeds_via_staged_rsram(): def dram_zn_to_wsram_kernel(): """DRAM buffer with ZN layout (via MeshTensor) copied to WSRAM → IsZZLike check should fail.""" - from tilelang.language.mesh_tensor import MeshShardingPolicy, MeshReplicationType M, K, N = 64, 32, 64 block_M, block_K, block_N = 64, 32, 64 dtype = "float16" accum_dtype = "float32" - policy = MeshShardingPolicy(y=0, x=1, replicate=MeshReplicationType.NONE) + policy = T.placement.full_shard(0, 1) # ZN layout on DRAM B — incompatible with WSRAM DMA (requires ZZ-like) zn_dram = make_zn_layout((K, N), axes=[0, 1], block_shape=[32, 32]) @@ -1329,18 +1327,18 @@ def matmul_persistent(M, N, K, block_M, block_N, block_K, num_stages, dtype=T.bf @T.prim_func def main( - A: T.MeshTensor((M, K), T.MeshShardingPolicy(y=0, x=1), dtype, layout=A_layout), - B: T.MeshTensor((K, N), T.MeshShardingPolicy(y=0, x=1), dtype, layout=B_layout), - C: T.MeshTensor((M, N), T.MeshShardingPolicy(y=0, x=1), accum_dtype, layout=C_layout), + A: T.MeshTensor((M, K), T.placement.full_shard(0, 1), dtype, layout=A_layout), + B: T.MeshTensor((K, N), T.placement.full_shard(0, 1), dtype, layout=B_layout), + C: T.MeshTensor((M, N), T.placement.full_shard(0, 1), accum_dtype, layout=C_layout), ): with T.Kernel() as (_cid): sharded_M, sharded_K = A.local_shape _, sharded_N = B.local_shape A_shared = T.alloc_shared((block_M, block_K), dtype) - A_shared_dist = T.alloc_shared((block_M, block_K * T.mesh_ncols()), dtype) + A_shared_dist = T.alloc_shared((block_M, block_K * T.ncols()), dtype) B_shared = T.alloc_shared((block_K, block_N), dtype) - B_shared_dist = T.alloc_shared((block_K * T.mesh_nrows(), block_N), dtype) + B_shared_dist = T.alloc_shared((block_K * T.nrows(), block_N), dtype) C_shared = T.alloc_shared((block_M, block_N), accum_dtype) # Each core iterates its own sharded tile grid with plain nested diff --git a/testing/python/sunmmio/transform/test_layout_transform.py b/testing/python/sunmmio/transform/test_layout_transform.py index 3cc4bf6af2..bd64d13dd8 100644 --- a/testing/python/sunmmio/transform/test_layout_transform.py +++ b/testing/python/sunmmio/transform/test_layout_transform.py @@ -26,7 +26,6 @@ make_row_major, make_zz_layout, ) -from tilelang.language.mesh_tensor import MeshReplicationType from tvm import tir from tvm.tir import Block from tvm.tir.stmt_functor import post_order_visit @@ -122,12 +121,12 @@ def _dram(shape, layout): these are single-core layout tests (unlike the sharded GEMM example), so sharding must not perturb the shapes the assertions rely on. """ - policy = T.MeshShardingPolicy(replicate=MeshReplicationType.ALL) + policy = T.placement.replicated() return T.MeshTensor(shape, policy, DTYPE, layout=layout) def _dram_typed(shape, dtype, layout): - policy = T.MeshShardingPolicy(replicate=MeshReplicationType.ALL) + policy = T.placement.replicated() return T.MeshTensor(shape, policy, dtype, layout=layout) diff --git a/testing/python/sunmmio/transform/test_legalize_gemm.py b/testing/python/sunmmio/transform/test_legalize_gemm.py index 082c49f010..370f7f18f4 100644 --- a/testing/python/sunmmio/transform/test_legalize_gemm.py +++ b/testing/python/sunmmio/transform/test_legalize_gemm.py @@ -83,17 +83,17 @@ def bf16_gemm_with_allgather(M=128, N=128, K=128, block_M=32, block_N=32, block_ @T.prim_func def main( - A: T.MeshTensor((M, K), T.MeshShardingPolicy(y=0, x=1), dtype, layout=A_layout), - B: T.MeshTensor((K, N), T.MeshShardingPolicy(y=0, x=1), dtype, layout=B_layout), - C: T.MeshTensor((M, N), T.MeshShardingPolicy(y=0, x=1), accum_dtype, layout=C_layout), + A: T.MeshTensor((M, K), T.placement.full_shard(0, 1), dtype, layout=A_layout), + B: T.MeshTensor((K, N), T.placement.full_shard(0, 1), dtype, layout=B_layout), + C: T.MeshTensor((M, N), T.placement.full_shard(0, 1), accum_dtype, layout=C_layout), ): with T.Kernel() as (_cid): sharded_M, sharded_K = A.local_shape _, sharded_N = B.local_shape A_shared = T.alloc_shared((block_M, block_K), dtype) - A_shared_dist = T.alloc_shared((block_M, block_K * T.mesh_ncols()), dtype) + A_shared_dist = T.alloc_shared((block_M, block_K * T.ncols()), dtype) B_shared = T.alloc_shared((block_K, block_N), dtype) - B_shared_dist = T.alloc_shared((block_K * T.mesh_nrows(), block_N), dtype) + B_shared_dist = T.alloc_shared((block_K * T.nrows(), block_N), dtype) C_shared = T.alloc_shared((block_M, block_N), accum_dtype) T.clear(C_shared) @@ -159,17 +159,17 @@ def bf16_gemm_with_copy_to_asram(M=128, N=128, K=128, block_M=32, block_N=32, bl @T.prim_func def main( - A: T.MeshTensor((M, K), T.MeshShardingPolicy(y=0, x=1), dtype, layout=A_layout), - B: T.MeshTensor((K, N), T.MeshShardingPolicy(y=0, x=1), dtype, layout=B_layout), - C: T.MeshTensor((M, N), T.MeshShardingPolicy(y=0, x=1), accum_dtype, layout=C_layout), + A: T.MeshTensor((M, K), T.placement.full_shard(0, 1), dtype, layout=A_layout), + B: T.MeshTensor((K, N), T.placement.full_shard(0, 1), dtype, layout=B_layout), + C: T.MeshTensor((M, N), T.placement.full_shard(0, 1), accum_dtype, layout=C_layout), ): with T.Kernel() as (_cid): sharded_M, sharded_K = A.local_shape _, sharded_N = B.local_shape # A staged via direct copy (no all_gather): full DRAM->ASRAM transfer. - A_dist = T.alloc_shared((block_M, block_K * T.mesh_ncols()), dtype) + A_dist = T.alloc_shared((block_M, block_K * T.ncols()), dtype) B_shared = T.alloc_shared((block_K, block_N), dtype) - B_shared_dist = T.alloc_shared((block_K * T.mesh_nrows(), block_N), dtype) + B_shared_dist = T.alloc_shared((block_K * T.nrows(), block_N), dtype) C_shared = T.alloc_shared((block_M, block_N), accum_dtype) T.clear(C_shared) @@ -235,9 +235,9 @@ def bf16_gemm_with_hoisted_copy_writer(M=128, N=128, K=128, block_M=32, block_N= @T.prim_func def main( - A: T.MeshTensor((M, K), T.MeshShardingPolicy(y=0, x=1), dtype, layout=A_layout), - B: T.MeshTensor((K, N), T.MeshShardingPolicy(y=0, x=1), dtype, layout=B_layout), - C: T.MeshTensor((M, N), T.MeshShardingPolicy(y=0, x=1), accum_dtype, layout=C_layout), + A: T.MeshTensor((M, K), T.placement.full_shard(0, 1), dtype, layout=A_layout), + B: T.MeshTensor((K, N), T.placement.full_shard(0, 1), dtype, layout=B_layout), + C: T.MeshTensor((M, N), T.placement.full_shard(0, 1), accum_dtype, layout=C_layout), ): with T.Kernel() as (_cid): sharded_M, sharded_K = A.local_shape @@ -286,9 +286,9 @@ def bf16_gemm_with_hoisted_copy_writer_bf16_acc(M=128, N=128, K=128, block_M=32, @T.prim_func def main( - A: T.MeshTensor((M, K), T.MeshShardingPolicy(y=0, x=1), dtype, layout=A_layout), - B: T.MeshTensor((K, N), T.MeshShardingPolicy(y=0, x=1), dtype, layout=B_layout), - C: T.MeshTensor((M, N), T.MeshShardingPolicy(y=0, x=1), accum_dtype, layout=C_layout), + A: T.MeshTensor((M, K), T.placement.full_shard(0, 1), dtype, layout=A_layout), + B: T.MeshTensor((K, N), T.placement.full_shard(0, 1), dtype, layout=B_layout), + C: T.MeshTensor((M, N), T.placement.full_shard(0, 1), accum_dtype, layout=C_layout), ): with T.Kernel() as (_cid): sharded_M, sharded_K = A.local_shape @@ -336,11 +336,11 @@ def bf16_gemm_with_two_independent_hoisted_writers( @T.prim_func def main( - A: T.MeshTensor((M, K), T.MeshShardingPolicy(y=0, x=1), dtype, layout=A_layout), - A2: T.MeshTensor((M, K), T.MeshShardingPolicy(y=0, x=1), dtype, layout=A_layout), - B: T.MeshTensor((K, N), T.MeshShardingPolicy(y=0, x=1), dtype, layout=B_layout), - C: T.MeshTensor((M, N), T.MeshShardingPolicy(y=0, x=1), accum_dtype, layout=C_layout), - C2: T.MeshTensor((M, N), T.MeshShardingPolicy(y=0, x=1), accum_dtype, layout=C_layout), + A: T.MeshTensor((M, K), T.placement.full_shard(0, 1), dtype, layout=A_layout), + A2: T.MeshTensor((M, K), T.placement.full_shard(0, 1), dtype, layout=A_layout), + B: T.MeshTensor((K, N), T.placement.full_shard(0, 1), dtype, layout=B_layout), + C: T.MeshTensor((M, N), T.placement.full_shard(0, 1), accum_dtype, layout=C_layout), + C2: T.MeshTensor((M, N), T.placement.full_shard(0, 1), accum_dtype, layout=C_layout), ): with T.Kernel() as (_cid): sharded_M, sharded_K = A.local_shape @@ -398,9 +398,9 @@ def bf16_gemm_with_hoisted_writer_dirty_source( @T.prim_func def main( - A: T.MeshTensor((M, K), T.MeshShardingPolicy(y=0, x=1), dtype, layout=A_layout), - B: T.MeshTensor((K, N), T.MeshShardingPolicy(y=0, x=1), dtype, layout=B_layout), - C: T.MeshTensor((M, N), T.MeshShardingPolicy(y=0, x=1), accum_dtype, layout=C_layout), + A: T.MeshTensor((M, K), T.placement.full_shard(0, 1), dtype, layout=A_layout), + B: T.MeshTensor((K, N), T.placement.full_shard(0, 1), dtype, layout=B_layout), + C: T.MeshTensor((M, N), T.placement.full_shard(0, 1), accum_dtype, layout=C_layout), ): with T.Kernel() as (_cid): sharded_M, sharded_K = A.local_shape @@ -449,17 +449,17 @@ def bf16_gemm_with_hoisted_allgather(M=128, N=128, K=128, block_M=32, block_N=32 @T.prim_func def main( - A: T.MeshTensor((M, K), T.MeshShardingPolicy(y=0, x=1), dtype, layout=A_layout), - B: T.MeshTensor((K, N), T.MeshShardingPolicy(y=0, x=1), dtype, layout=B_layout), - C: T.MeshTensor((M, N), T.MeshShardingPolicy(y=0, x=1), accum_dtype, layout=C_layout), + A: T.MeshTensor((M, K), T.placement.full_shard(0, 1), dtype, layout=A_layout), + B: T.MeshTensor((K, N), T.placement.full_shard(0, 1), dtype, layout=B_layout), + C: T.MeshTensor((M, N), T.placement.full_shard(0, 1), accum_dtype, layout=C_layout), ): with T.Kernel() as (_cid): sharded_M, sharded_K = A.local_shape _, sharded_N = B.local_shape A_shared = T.alloc_shared((block_M, block_K), dtype) - A_shared_dist = T.alloc_shared((block_M, block_K * T.mesh_ncols()), dtype) + A_shared_dist = T.alloc_shared((block_M, block_K * T.ncols()), dtype) B_shared = T.alloc_shared((block_K, block_N), dtype) - B_shared_dist = T.alloc_shared((block_K * T.mesh_nrows(), block_N), dtype) + B_shared_dist = T.alloc_shared((block_K * T.nrows(), block_N), dtype) C_shared = T.alloc_shared((block_M, block_N), accum_dtype) # A loaded + gathered ONCE, hoisted out of the K-loop. @@ -502,17 +502,17 @@ def bf16_gemm_with_hoisted_allgather_dirty_source( @T.prim_func def main( - A: T.MeshTensor((M, K), T.MeshShardingPolicy(y=0, x=1), dtype, layout=A_layout), - B: T.MeshTensor((K, N), T.MeshShardingPolicy(y=0, x=1), dtype, layout=B_layout), - C: T.MeshTensor((M, N), T.MeshShardingPolicy(y=0, x=1), accum_dtype, layout=C_layout), + A: T.MeshTensor((M, K), T.placement.full_shard(0, 1), dtype, layout=A_layout), + B: T.MeshTensor((K, N), T.placement.full_shard(0, 1), dtype, layout=B_layout), + C: T.MeshTensor((M, N), T.placement.full_shard(0, 1), accum_dtype, layout=C_layout), ): with T.Kernel() as (_cid): sharded_M, sharded_K = A.local_shape _, sharded_N = B.local_shape A_shared = T.alloc_shared((block_M, block_K), dtype) - A_shared_dist = T.alloc_shared((block_M, block_K * T.mesh_ncols()), dtype) + A_shared_dist = T.alloc_shared((block_M, block_K * T.ncols()), dtype) B_shared = T.alloc_shared((block_K, block_N), dtype) - B_shared_dist = T.alloc_shared((block_K * T.mesh_nrows(), block_N), dtype) + B_shared_dist = T.alloc_shared((block_K * T.nrows(), block_N), dtype) C_shared = T.alloc_shared((block_M, block_N), accum_dtype) T.copy(A[0, 0], A_shared) @@ -863,10 +863,10 @@ def bf16_gemm_multi_consumer_groupable(M=128, N=128, K=128, block_M=32, block_N= @T.prim_func def main( - A: T.MeshTensor((M, K), T.MeshShardingPolicy(y=0, x=1), dtype, layout=A_layout), - B: T.MeshTensor((K, N), T.MeshShardingPolicy(y=0, x=1), dtype, layout=B_layout), - C: T.MeshTensor((M, N), T.MeshShardingPolicy(y=0, x=1), accum_dtype, layout=C_layout), - C2: T.MeshTensor((M, N), T.MeshShardingPolicy(y=0, x=1), accum_dtype, layout=C_layout), + A: T.MeshTensor((M, K), T.placement.full_shard(0, 1), dtype, layout=A_layout), + B: T.MeshTensor((K, N), T.placement.full_shard(0, 1), dtype, layout=B_layout), + C: T.MeshTensor((M, N), T.placement.full_shard(0, 1), accum_dtype, layout=C_layout), + C2: T.MeshTensor((M, N), T.placement.full_shard(0, 1), accum_dtype, layout=C_layout), ): with T.Kernel() as (_cid): sharded_M, sharded_K = A.local_shape @@ -918,11 +918,11 @@ def bf16_gemm_multi_consumer_non_groupable( @T.prim_func def main( - A: T.MeshTensor((M, K), T.MeshShardingPolicy(y=0, x=1), dtype, layout=A_layout), - B: T.MeshTensor((K, N), T.MeshShardingPolicy(y=0, x=1), dtype, layout=B_layout), - B2: T.MeshTensor((K, N), T.MeshShardingPolicy(y=0, x=1), dtype, layout=B_layout), - C: T.MeshTensor((M, N), T.MeshShardingPolicy(y=0, x=1), accum_dtype, layout=C_layout), - C2: T.MeshTensor((M, N), T.MeshShardingPolicy(y=0, x=1), accum_dtype, layout=C_layout), + A: T.MeshTensor((M, K), T.placement.full_shard(0, 1), dtype, layout=A_layout), + B: T.MeshTensor((K, N), T.placement.full_shard(0, 1), dtype, layout=B_layout), + B2: T.MeshTensor((K, N), T.placement.full_shard(0, 1), dtype, layout=B_layout), + C: T.MeshTensor((M, N), T.placement.full_shard(0, 1), accum_dtype, layout=C_layout), + C2: T.MeshTensor((M, N), T.placement.full_shard(0, 1), accum_dtype, layout=C_layout), ): with T.Kernel() as (_cid): sharded_M, sharded_K = A.local_shape @@ -991,10 +991,10 @@ def bf16_flash_attention_shaped(M=128, N=128, K=128, block_M=32, block_N=32, blo @T.prim_func def main( - Q: T.MeshTensor((M, K), T.MeshShardingPolicy(y=0, x=1), dtype, layout=Q_layout), - Kt: T.MeshTensor((K, N), T.MeshShardingPolicy(y=0, x=1), dtype, layout=Kt_layout), - V: T.MeshTensor((N, N), T.MeshShardingPolicy(y=0, x=1), dtype, layout=V_layout), - O: T.MeshTensor((M, N), T.MeshShardingPolicy(y=0, x=1), accum_dtype, layout=O_layout), + Q: T.MeshTensor((M, K), T.placement.full_shard(0, 1), dtype, layout=Q_layout), + Kt: T.MeshTensor((K, N), T.placement.full_shard(0, 1), dtype, layout=Kt_layout), + V: T.MeshTensor((N, N), T.placement.full_shard(0, 1), dtype, layout=V_layout), + O: T.MeshTensor((M, N), T.placement.full_shard(0, 1), accum_dtype, layout=O_layout), ): with T.Kernel() as (_cid): sharded_M, sharded_K = Q.local_shape @@ -1111,9 +1111,9 @@ def bf16_gemm_no_reaching_writer(M=128, N=128, K=128, block_M=32, block_N=32, bl @T.prim_func def main( - A: T.MeshTensor((M, K), T.MeshShardingPolicy(y=0, x=1), dtype, layout=A_layout), - B: T.MeshTensor((K, N), T.MeshShardingPolicy(y=0, x=1), dtype, layout=B_layout), - C: T.MeshTensor((M, N), T.MeshShardingPolicy(y=0, x=1), accum_dtype, layout=C_layout), + A: T.MeshTensor((M, K), T.placement.full_shard(0, 1), dtype, layout=A_layout), + B: T.MeshTensor((K, N), T.placement.full_shard(0, 1), dtype, layout=B_layout), + C: T.MeshTensor((M, N), T.placement.full_shard(0, 1), accum_dtype, layout=C_layout), ): with T.Kernel() as (_cid): sharded_M, sharded_K = A.local_shape @@ -1150,10 +1150,10 @@ def bf16_gemm_multiple_writers(M=128, N=128, K=128, block_M=32, block_N=32, bloc @T.prim_func def main( - A: T.MeshTensor((M, K), T.MeshShardingPolicy(y=0, x=1), dtype, layout=A_layout), - A2: T.MeshTensor((M, K), T.MeshShardingPolicy(y=0, x=1), dtype, layout=A_layout), - B: T.MeshTensor((K, N), T.MeshShardingPolicy(y=0, x=1), dtype, layout=B_layout), - C: T.MeshTensor((M, N), T.MeshShardingPolicy(y=0, x=1), accum_dtype, layout=C_layout), + A: T.MeshTensor((M, K), T.placement.full_shard(0, 1), dtype, layout=A_layout), + A2: T.MeshTensor((M, K), T.placement.full_shard(0, 1), dtype, layout=A_layout), + B: T.MeshTensor((K, N), T.placement.full_shard(0, 1), dtype, layout=B_layout), + C: T.MeshTensor((M, N), T.placement.full_shard(0, 1), accum_dtype, layout=C_layout), ): with T.Kernel(T.mesh_ncores()) as cid: sharded_M, sharded_K = A.local_shape @@ -1195,9 +1195,9 @@ def bf16_gemm_diverging_scopes(M=128, N=128, K=128, block_M=32, block_N=32, bloc @T.prim_func def main( - A: T.MeshTensor((M, K), T.MeshShardingPolicy(y=0, x=1), dtype, layout=A_layout), - B: T.MeshTensor((K, N), T.MeshShardingPolicy(y=0, x=1), dtype, layout=B_layout), - C: T.MeshTensor((M, N), T.MeshShardingPolicy(y=0, x=1), accum_dtype, layout=C_layout), + A: T.MeshTensor((M, K), T.placement.full_shard(0, 1), dtype, layout=A_layout), + B: T.MeshTensor((K, N), T.placement.full_shard(0, 1), dtype, layout=B_layout), + C: T.MeshTensor((M, N), T.placement.full_shard(0, 1), accum_dtype, layout=C_layout), ): with T.Kernel() as (_cid): sharded_M, sharded_K = A.local_shape diff --git a/testing/python/sunmmio/transform/test_pipeline.py b/testing/python/sunmmio/transform/test_pipeline.py index a1a9f77b57..bbee9d3555 100644 --- a/testing/python/sunmmio/transform/test_pipeline.py +++ b/testing/python/sunmmio/transform/test_pipeline.py @@ -4,7 +4,6 @@ import tilelang.language as T from tilelang.engine.phase import * from tilelang.utils.target import SUNMMIO_TARGET_DESC -from tilelang.language.mesh_tensor import MeshShardingPolicy from examples.gemm.sunmmio_example_gemm import matmul_persistent _get_logical_shape = tvm.ffi.get_global_func("tl.CuteLayout_logical_shape") @@ -15,17 +14,17 @@ def matmul(M, N, K, block_M, block_N, block_K, num_stages, dtype="bfloat16", acc def gemm( A: T.MeshTensor( (M, K), - sharding_policy=MeshShardingPolicy(cross_mesh_dim=0), + placement=T.placement.mesh_as_line(0), dtype=dtype, ), B: T.MeshTensor( (K, N), - sharding_policy=MeshShardingPolicy(cross_mesh_dim=0), + placement=T.placement.mesh_as_line(0), dtype=dtype, ), C: T.MeshTensor( (M, N), - sharding_policy=MeshShardingPolicy(cross_mesh_dim=0), + placement=T.placement.mesh_as_line(0), dtype=accum_dtype, ), ): @@ -552,47 +551,60 @@ def main_no_split( lambda: matmul(1024, 1024, 1024, 128, 128, 32, num_stages=3), { "A_rsram_stage": [3, 128, 32], - "A_shared": [3, 128, 32], - "B_shared": [3, 32, 128], + "A_shared_ping": [2, 128, 32], + "A_shared_pong": [2, 128, 32], + "B_shared_ping": [2, 32, 128], + "B_shared_pong": [2, 32, 128], }, ), ( "flashattn", lambda: flashattn(num_stages=3), { - "K_shared": [3, 64, 128], + "K_shared_ping": [2, 64, 128], + "K_shared_pong": [2, 64, 128], "src_buffer": [3, 64, 64], - "acc_s_cast": [3, 64, 64], - "V_shared": [3, 64, 128], + "acc_s_cast_ping": [2, 64, 64], + "acc_s_cast_pong": [2, 64, 64], + "V_shared_ping": [2, 64, 128], + "V_shared_pong": [2, 64, 128], }, ), ( "flashdecoding", lambda: flashdecoding(num_stages=3), { - "K_shared": [3, 128, 128], + "K_shared_ping": [2, 128, 128], + "K_shared_pong": [2, 128, 128], "mask_local": [3, 128], "src_buffer": [3, 64, 128], - "acc_s_cast": [3, 64, 128], - "V_shared": [3, 128, 128], + "acc_s_cast_ping": [2, 64, 128], + "acc_s_cast_pong": [2, 64, 128], + "V_shared_ping": [2, 128, 128], + "V_shared_pong": [2, 128, 128], }, ), ( "flashmladecode", lambda: flashmladecode(num_stages=3), { - "KV_shared": [3, 64, 512], - "KV_shared2": [3, 64, 512], - "K_pe_shared": [3, 64, 64], + "KV_shared_ping": [2, 64, 512], + "KV_shared_pong": [2, 64, 512], + "KV_shared2_ping": [2, 64, 512], + "KV_shared2_pong": [2, 64, 512], + "K_pe_shared_ping": [2, 64, 64], + "K_pe_shared_pong": [2, 64, 64], }, ), ( "matmul_persistent", lambda: matmul_persistent(1024, 1024, 1024, 128, 128, 32, num_stages=2), { - "A_shared_dist": [2, 128, 128], + "A_shared_dist_ping": [128, 128], + "A_shared_dist_pong": [128, 128], "A_rsram_stage": [2, 128, 32], - "B_shared_dist": [2, 128, 128], + "B_shared_dist_ping": [128, 128], + "B_shared_dist_pong": [128, 128], }, ), ] @@ -667,17 +679,17 @@ def if_matmul(M, N, K, block_M, block_N, block_K, num_stages, dtype="bfloat16", def gemm( A: T.MeshTensor( (M, K), - sharding_policy=MeshShardingPolicy(cross_mesh_dim=0), + placement=T.placement.mesh_as_line(0), dtype=dtype, ), B: T.MeshTensor( (K, N), - sharding_policy=MeshShardingPolicy(cross_mesh_dim=0), + placement=T.placement.mesh_as_line(0), dtype=dtype, ), C: T.MeshTensor( (M, N), - sharding_policy=MeshShardingPolicy(cross_mesh_dim=0), + placement=T.placement.mesh_as_line(0), dtype=accum_dtype, ), ): @@ -703,7 +715,7 @@ def tvm_access_ptr(): def test( A: T.MeshTensor( (M, K), - sharding_policy=MeshShardingPolicy(cross_mesh_dim=0), + placement=T.placement.mesh_as_line(0), dtype=dtype, ), ): @@ -724,26 +736,31 @@ def test( return gemm -ERROR_CASES = [ +FALLBACK_CASES = [ ( lambda: if_matmul(1024, 1024, 1024, 128, 128, 32, num_stages=3), - "Can not identify the hardware type for a tir.IfThenElse statement.", + ("inject_validation", "candidate_fallback", "unsupported_statement"), ), ] @pytest.mark.parametrize( - "kernel,error_msg", - ERROR_CASES, + "kernel,expected_diagnostic", + FALLBACK_CASES, ) -def test_tilelang_transform_sunmmio_pipeline_error(kernel, error_msg): - with pytest.raises(tvm.error.InternalError, match=error_msg): - name = SUNMMIO_TARGET_DESC - target = tvm.target.Target(name) - - with tvm.target.Target(target): - mod = tvm.IRModule.from_expr(kernel().with_attr("global_symbol", "main")) - mod = lower_and_legalize_sunmmio_pipeline_test(mod, target) - mod = tl.transform.IfStmtBinding()(mod) - mod = tl.transform.SunmmioPipelinePlanning(debug=False)(mod) - mod = tl.transform.InjectSunmmioPipeline()(mod) +def test_tilelang_transform_sunmmio_pipeline_fallback(kernel, expected_diagnostic): + name = SUNMMIO_TARGET_DESC + target = tvm.target.Target(name) + + with tvm.target.Target(target): + mod = tvm.IRModule.from_expr(kernel().with_attr("global_symbol", "main")) + mod = lower_and_legalize_sunmmio_pipeline_test(mod, target) + mod = tl.transform.IfStmtBinding()(mod) + mod = tl.transform.SunmmioPipelinePlanning(debug=False)(mod) + mod = tl.transform.InjectSunmmioPipeline()(mod) + + stage, reason, detail = expected_diagnostic + script = mod.script() + assert f'"tl.sunmmio.pipeline.fallback_stage": "{stage}"' in script + assert f'"tl.sunmmio.pipeline.fallback_reason": "{reason}"' in script + assert f'"tl.sunmmio.pipeline.fallback_detail": "{detail}"' in script diff --git a/testing/python/sunmmio/transform/test_remove_unused_sunmmio_allocations.py b/testing/python/sunmmio/transform/test_remove_unused_sunmmio_allocations.py new file mode 100644 index 0000000000..a0a009458e --- /dev/null +++ b/testing/python/sunmmio/transform/test_remove_unused_sunmmio_allocations.py @@ -0,0 +1,59 @@ +import tilelang as tl +from tilelang import tvm +from tilelang import language as T +from tvm import tir + + +def _make_dangling_ping_pong_allocation(): + live = tir.decl_buffer((16,), "float32", name="live_ping", scope="shared.rsram") + dead = tir.decl_buffer((16,), "float32", name="dead_pong", scope="shared.rsram") + + body = tir.BufferStore(live, tir.FloatImm("float32", 1), [0]) + body = tir.DeclBuffer(live, body) + body = tir.Allocate(live.data, "float32", [16], True, body) + body = tir.Allocate( + dead.data, + "float32", + [16], + True, + body, + annotations={"tl.sunmmio_alloc_ping_pong": "pong"}, + ) + + layout = T.Layout((16,), lambda i: i) + func = tir.PrimFunc([], body) + func = func.with_attr("layout_map", {live: layout, dead: layout}) + func = func.with_attr("tl.sunmmio_alloc_ping_pong", {live.data: "ping", dead.data: "pong"}) + return tvm.IRModule.from_expr(func) + + +def _collect_buffer_declarations(func): + allocations = set() + declarations = set() + + def visit(node): + if isinstance(node, tir.Allocate): + allocations.add(node.buffer_var.name) + elif isinstance(node, tir.DeclBuffer): + declarations.add(node.buffer.name) + + tir.stmt_functor.post_order_visit(func.body, visit) + return allocations, declarations + + +def test_remove_unused_sunmmio_allocations_cleans_dangling_metadata(): + result = tl.transform.RemoveUnusedSunmmioAllocations()(_make_dangling_ping_pong_allocation())["main"] + + allocations, declarations = _collect_buffer_declarations(result) + assert allocations == {"live_ping"} + assert declarations == {"live_ping"} + + layout_map = result.attrs["layout_map"] + assert {buffer.name for buffer in layout_map} == {"live_ping"} + + ping_pong = result.attrs["tl.sunmmio_alloc_ping_pong"] + assert {var.name for var in ping_pong} == {"live_ping"} + + +if __name__ == "__main__": + test_remove_unused_sunmmio_allocations_cleans_dangling_metadata() diff --git a/testing/python/sunmmio/transform/test_resolve_mesh_symbols.py b/testing/python/sunmmio/transform/test_resolve_mesh_symbols.py index 8a21f59cda..976f6300de 100644 --- a/testing/python/sunmmio/transform/test_resolve_mesh_symbols.py +++ b/testing/python/sunmmio/transform/test_resolve_mesh_symbols.py @@ -173,7 +173,7 @@ def _symbolic_mesh_gemm_mod(): block_M, block_N, block_K = 32, 32, 32 dtype = "bfloat16" accum_dtype = "float32" - policy = T.MeshShardingPolicy(y=0, x=1) + policy = T.placement.full_shard(0, 1) a_layout = make_zz_layout((M, K), [0, 1], (32, 32)) b_layout = make_zz_layout((K, N), [0, 1], (32, 32)) c_layout = make_zz_layout((M, N), [0, 1], (32, 32)) @@ -191,9 +191,9 @@ def main( _, sharded_N = B.local_shape A_shared = T.alloc_shared((block_M, block_K), dtype) - A_shared_dist = T.alloc_shared((block_M, block_K * T.mesh_ncols()), dtype) + A_shared_dist = T.alloc_shared((block_M, block_K * T.ncols()), dtype) B_shared = T.alloc_shared((block_K, block_N), dtype) - B_shared_dist = T.alloc_shared((block_K * T.mesh_nrows(), block_N), dtype) + B_shared_dist = T.alloc_shared((block_K * T.nrows(), block_N), dtype) C_shared = T.alloc_shared((block_M, block_N), accum_dtype) for bx in T.serial(T.ceildiv(sharded_M, block_M)): @@ -220,7 +220,7 @@ def _symbolic_mesh_all_gather_mod(): def main(A: T.Tensor((block_M, block_N), dtype)): with T.Kernel() as _cid: A_shared = T.alloc_shared((block_M, block_N), dtype) - A_shared_dist = T.alloc_shared((block_M, block_N * T.mesh_ncols()), dtype) + A_shared_dist = T.alloc_shared((block_M, block_N * T.ncols()), dtype) T.copy(A, A_shared) T.comm.all_gather(A_shared, A_shared_dist, direction="horizontal", axis=-1) @@ -236,7 +236,7 @@ def _symbolic_mesh_gqa_mod(): accum_dtype = "bfloat16" q_shape = [batch, seq_len, heads, dim] kv_shape = [batch, seq_len, head_kv, dim] - policy = T.MeshShardingPolicy(y=0, x=2) + policy = T.placement.full_shard(0, 2) q_layout = make_zz_layout(q_shape, [1, 3], (32, 32)) kv_layout = make_zz_layout(kv_shape, [1, 3], (32, 32)) @@ -287,10 +287,10 @@ def test_resolve_replaces_mesh_intrinsics_in_body_and_kernel_extent(): @T.prim_func def main(): with T.Kernel() as cid: - scratch = T.alloc_shared((T.mesh_nrows(), T.mesh_ncols(), T.mesh_ncores()), "float32") + scratch = T.alloc_shared((T.nrows(), T.ncols(), T.mesh_ncores()), "float32") scratch[0, 0, 0] = T.if_then_else( cid < T.mesh_ncores(), - T.Cast("float32", T.mesh_nrows() + T.mesh_ncols() + T.mesh_ncores()), + T.Cast("float32", T.nrows() + T.ncols() + T.mesh_ncores()), T.float32(0), ) @@ -310,7 +310,7 @@ def main(): def test_resolve_updates_default_mesh_tensor_buffer_map_and_layout_metadata(): target = _sunmmio_target() - policy = T.MeshShardingPolicy(y=0, x=1) + policy = T.placement.full_shard(0, 1) layout = make_zz_layout((128, 96), [0, 1], (32, 32)) with tvm.target.Target(target): @@ -318,7 +318,7 @@ def test_resolve_updates_default_mesh_tensor_buffer_map_and_layout_metadata(): @T.prim_func def main(A: T.MeshTensor((128, 96), policy, "float16", layout=layout)): with T.Kernel() as _cid: - valid_M, valid_N = A.get_local_extent(_cid) + valid_M, valid_N = A.get_local_extent() for i in T.serial(valid_M): for j in T.serial(valid_N): A[i, j] = A[i, j] @@ -367,7 +367,7 @@ def test_mesh_symbol_query_has_no_builder_side_effect(): with builder.current_context(): builder._sunmmio_mesh_symbols_used = False - expr = T.mesh_nrows() + 1 + expr = T.nrows() + 1 assert builder._sunmmio_mesh_symbols_used builder._sunmmio_mesh_symbols_used = False @@ -386,7 +386,7 @@ def main(): scratch = T.alloc_shared((1,), "float32") scratch[0] = T.if_then_else( cid < T.mesh_ncores(), - T.Cast("float32", T.mesh_nrows() + T.mesh_ncols() + T.mesh_ncores()), + T.Cast("float32", T.nrows() + T.ncols() + T.mesh_ncores()), T.float32(0), ) diff --git a/testing/python/sunmmio/transform/test_tile_loop_fusion_rewrite.py b/testing/python/sunmmio/transform/test_tile_loop_fusion_rewrite.py index f3c5bf74ad..7371aed641 100644 --- a/testing/python/sunmmio/transform/test_tile_loop_fusion_rewrite.py +++ b/testing/python/sunmmio/transform/test_tile_loop_fusion_rewrite.py @@ -316,6 +316,158 @@ def _lowered_2d_tile_region(dst, expr, *, block_m=32, block_n=32, tile_size=(8, ).strip() +def _lowered_row_reduction_region(src, dst, *, block_m=64, block_n=64, tile_size=(4, 32)): + tile_m, tile_n = tile_size + outer_m = block_m // tile_m + outer_n = block_n // tile_n + return textwrap.dedent( + f""" + for i in T.serial( + {outer_m}, + annotations={{ + "tile.domain": [T.int32({block_m}), T.int32({block_n})], + "tile.execution_axis": T.int32(0), + "tile.execution_domain_axes": [T.int32(0), T.int32(1)], + "tile.scope_entry": T.int32(1), + "tile.tile_size": [T.int32({tile_m}), T.int32({tile_n})], + }}, + ): + for j in T.serial({outer_n}, annotations={{"tile.execution_axis": T.int32(1)}}): + with T.block("reduce_tile_op"): + T.reads() + T.writes() + dst_buffer_acc = T.alloc_buffer(({tile_m}, {tile_n}), "float32", scope="shared.rsram") + if j == 0: + for ki in T.serial({tile_m}, annotations={{"tile.interior": T.int32(1), "tile.interior_axis": T.int32(0)}}): + for kj in T.vectorized({tile_n}, annotations={{"tile.interior": T.int32(1), "tile.interior_axis": T.int32(1)}}): + dst_buffer_acc[ki, kj] = T.float32("-inf") + for ki in T.serial({tile_m}, annotations={{"tile.interior": T.int32(1), "tile.interior_axis": T.int32(0)}}): + for kj in T.vectorized({tile_n}, annotations={{"tile.interior": T.int32(1), "tile.interior_axis": T.int32(1)}}): + dst_buffer_acc[ki, kj] = T.max( + dst_buffer_acc[ki, kj], + {src}[i * {tile_m} + ki, j * {tile_n} + kj], + ) + if j == {outer_n - 1}: + for ki in T.serial({tile_m}, annotations={{"tile.interior": T.int32(1), "tile.interior_axis": T.int32(0)}}): + {dst}[i * {tile_m} + ki] = dst_buffer_acc[ki, 0] + """ + ).strip() + + +def _lowered_row_write_region(src, dst, *, block_m=64, block_n=64, tile_size=(4, 32)): + tile_m, tile_n = tile_size + outer_m = block_m // tile_m + outer_n = block_n // tile_n + return textwrap.dedent( + f""" + for i in T.serial( + {outer_m}, + annotations={{ + "tile.domain": [T.int32({block_m}), T.int32({block_n})], + "tile.execution_axis": T.int32(0), + "tile.execution_domain_axes": [T.int32(0), T.int32(1)], + "tile.scope_entry": T.int32(1), + "tile.tile_size": [T.int32({tile_m}), T.int32({tile_n})], + }}, + ): + for j in T.serial({outer_n}, annotations={{"tile.execution_axis": T.int32(1)}}): + for ki in T.serial({tile_m}, annotations={{"tile.interior": T.int32(1), "tile.interior_axis": T.int32(0)}}): + {dst}[i * {tile_m} + ki] = {src}[i * {tile_m} + ki, j * {tile_n}] + """ + ).strip() + + +def reduction_consumer_lowered_kernel(block_n=64): + block_m = 64 + dtype = "float32" + tile_size = (4, 32) + return _make_manual_lowered_primfunc( + [ + f'a_shared = T.alloc_buffer(({block_m}, {block_n}), "{dtype}", scope="shared.rsram")', + f'values = T.alloc_buffer(({block_m}, {block_n}), "{dtype}", scope="shared.rsram")', + f'row_max = T.alloc_buffer(({block_m},), "{dtype}", scope="shared.rsram")', + f'b_shared = T.alloc_buffer(({block_m}, {block_n}), "{dtype}", scope="shared.rsram")', + ], + [ + _lowered_2d_tile_region( + "values", + "a_shared[i * 4 + ki, j * 32 + kj]", + block_m=block_m, + block_n=block_n, + tile_size=tile_size, + ), + _lowered_row_reduction_region("values", "row_max", block_m=block_m, block_n=block_n, tile_size=tile_size), + _lowered_2d_tile_region( + "b_shared", + "values[i * 4 + ki, j * 32 + kj] - row_max[i * 4 + ki]", + block_m=block_m, + block_n=block_n, + tile_size=tile_size, + ), + ], + ) + + +def war_reduction_lowered_kernel(): + block_m = 64 + block_n = 64 + dtype = "float32" + tile_size = (4, 32) + return _make_manual_lowered_primfunc( + [ + f'a_shared = T.alloc_buffer(({block_m}, {block_n}), "{dtype}", scope="shared.rsram")', + f'tile_values = T.alloc_buffer(({block_m}, {block_n}), "{dtype}", scope="shared.rsram")', + f'row_values = T.alloc_buffer(({block_m},), "{dtype}", scope="shared.rsram")', + ], + [ + _lowered_2d_tile_region( + "tile_values", + "a_shared[i * 4 + ki, j * 32 + kj] + row_values[i * 4 + ki]", + block_m=block_m, + block_n=block_n, + tile_size=tile_size, + ), + _lowered_row_reduction_region( + "tile_values", + "row_values", + block_m=block_m, + block_n=block_n, + tile_size=tile_size, + ), + ], + ) + + +def unsafe_war_lowered_kernel(): + block_m = 64 + block_n = 64 + dtype = "float32" + tile_size = (4, 32) + return _make_manual_lowered_primfunc( + [ + f'a_shared = T.alloc_buffer(({block_m}, {block_n}), "{dtype}", scope="shared.rsram")', + f'tile_values = T.alloc_buffer(({block_m}, {block_n}), "{dtype}", scope="shared.rsram")', + f'row_values = T.alloc_buffer(({block_m},), "{dtype}", scope="shared.rsram")', + ], + [ + _lowered_2d_tile_region( + "tile_values", + "a_shared[i * 4 + ki, j * 32 + kj] + row_values[i * 4 + ki]", + block_m=block_m, + block_n=block_n, + tile_size=tile_size, + ), + _lowered_row_write_region( + "tile_values", + "row_values", + block_m=block_m, + block_n=block_n, + tile_size=tile_size, + ), + ], + ) + + def _lowered_predicated_2d_tile_region(dst, src, *, block_m=4, block_n=4, tile_size=(4, 32)): tile_m, tile_n = tile_size outer_m = (block_m + tile_m - 1) // tile_m @@ -835,6 +987,98 @@ def test_sunmmio_tile_loop_fusion_can_share_only_the_outer_loop_prefix(): assert all(_for_annotations(stmt).get("tile.execution_axis") == 1 for stmt in fused_body) +def test_sunmmio_tile_loop_fusion_waits_for_conditional_raw_reduction_write(): + mod = IRModule.from_expr(reduction_consumer_lowered_kernel().with_attr("global_symbol", "main")) + mod = tl.transform.SunmmioTileLoopFusion()(mod) + + outer_scope = _expect_single_match( + _find_scope_entry_loops(_root_seq(mod)), + lambda loop: "b_shared" in _collect_buffer_accesses(loop)[1], + "outer row shell containing the reduction consumer", + ) + inner_scopes = [ + stmt + for stmt in _as_seq(outer_scope.body) + if isinstance(stmt, tvm.tir.For) and _for_annotations(stmt).get("tile.execution_axis") == 1 + ] + + assert len(inner_scopes) == 2 + assert "b_shared" not in _collect_buffer_accesses(inner_scopes[0])[1] + assert "b_shared" in _collect_buffer_accesses(inner_scopes[1])[1] + + +def test_sunmmio_tile_loop_fusion_keeps_conditional_war_reduction_write_fused(): + mod = IRModule.from_expr(war_reduction_lowered_kernel().with_attr("global_symbol", "main")) + mod = tl.transform.SunmmioTileLoopFusion()(mod) + + scope_loops = _find_scope_entry_loops(_root_seq(mod)) + assert len(scope_loops) == 1 + + inner_scopes = [ + stmt + for stmt in _as_seq(scope_loops[0].body) + if isinstance(stmt, tvm.tir.For) and _for_annotations(stmt).get("tile.execution_axis") == 1 + ] + assert len(inner_scopes) == 1 + + fused_body = _as_seq(inner_scopes[0].body) + assert len(fused_body) == 2 + assert _single_write_name(fused_body[0]) == "tile_values" + reduction_block = _expect_reduce_block(fused_body[1]) + + row_write_conditions = [] + + def collect_row_write_condition(node): + if not isinstance(node, tvm.tir.IfThenElse): + return + if "row_values" in _collect_buffer_accesses(node.then_case)[1]: + row_write_conditions.append(node.condition) + + tvm.tir.stmt_functor.post_order_visit(reduction_block.body, collect_row_write_condition) + assert len(row_write_conditions) == 1 + assert tvm.ir.structural_equal( + row_write_conditions[0], + inner_scopes[0].loop_var == inner_scopes[0].extent - 1, + ) + + reads, writes = _collect_buffer_accesses(inner_scopes[0]) + assert "row_values" in reads + assert "row_values" in writes + assert "tile_values" in reads + assert "tile_values" in writes + + +def test_sunmmio_tile_loop_fusion_splits_loop_carried_war_write(): + mod = IRModule.from_expr(unsafe_war_lowered_kernel().with_attr("global_symbol", "main")) + mod = tl.transform.SunmmioTileLoopFusion()(mod) + + scope_loops = _find_scope_entry_loops(_root_seq(mod)) + assert len(scope_loops) == 2 + assert [_single_write_name(loop) for loop in scope_loops] == [ + "tile_values", + "row_values", + ] + + +def test_sunmmio_tile_loop_fusion_keeps_unit_trip_reduction_consumer_fused(): + mod = IRModule.from_expr(reduction_consumer_lowered_kernel(block_n=32).with_attr("global_symbol", "main")) + mod = tl.transform.SunmmioTileLoopFusion()(mod) + + outer_scope = _expect_single_match( + _find_scope_entry_loops(_root_seq(mod)), + lambda loop: "b_shared" in _collect_buffer_accesses(loop)[1], + "outer row shell containing the unit-trip reduction consumer", + ) + inner_scopes = [ + stmt + for stmt in _as_seq(outer_scope.body) + if isinstance(stmt, tvm.tir.For) and _for_annotations(stmt).get("tile.execution_axis") == 1 + ] + + assert len(inner_scopes) == 1 + assert "b_shared" in _collect_buffer_accesses(inner_scopes[0])[1] + + def test_sunmmio_tile_loop_fusion_rewrites_multiple_planning_groups_independently(): mod = IRModule.from_expr(two_disjoint_groups_lowered_kernel().with_attr("global_symbol", "main")) mod = tl.transform.SunmmioTileLoopFusion()(mod) diff --git a/testing/python/transform/sunmmio_mesh_kernel_new_syntax_reference.py b/testing/python/transform/sunmmio_mesh_kernel_new_syntax_reference.py new file mode 100644 index 0000000000..af792d4eb3 --- /dev/null +++ b/testing/python/transform/sunmmio_mesh_kernel_new_syntax_reference.py @@ -0,0 +1,597 @@ +""" +Reference implementations showing how to rewrite the four core SunMMIO ILP +kernels with the newer TileLang-Mesh style. + +This file is documentation-oriented: +- It is intentionally separate from the existing tests. +- It focuses on the three decisions that matter most: + 1. Mesh sharding policy + 2. Global ZZ layout choice + 3. MeshTensor parameter declaration + +Important note +-------------- +The user guide describes a newer API shape: + + A: T.MeshTensor((M, K), shard_policy, dtype, layout=A_layout) + with T.Kernel() as cid: + sharded_M, sharded_K = A.local_shape + valid_M, valid_K = A.get_local_extent() + +At the time of writing, parts of that API are still ahead of the Python DSL +implementation in this local checkout. So the examples below do two things: + +- They follow the *design intent* of the new Mesh style: + explicit `T.placement`, explicit `make_zz_layout`, explicit + `T.MeshTensor(...)`. +- They also stay close to constructs that already exist in the repository, so + the examples remain grounded in the current codebase. + +Read this file as "how these kernels should be structured", not as an ABI +contract for the exact parser behavior of the current checkout. +""" + +import tilelang.language as T +from tilelang.layout import make_zz_layout, make_row_major +from testing.python.sunmmio.common.compile_pipeline import target + + +@target("Sunmmio") +def ilp( + M=128, + N=128, + K=1024, + block_M=32, + block_N=32, + block_K=32, + dtype=T.bfloat16, + accum_dtype=T.float32, +): + shard_policy = T.placement.full_shard(0, 1) + A_layout = make_zz_layout((M, K)) + B_layout = make_zz_layout((K, N)) + C_layout = make_zz_layout((M, N)) + + @T.prim_func + def main( + A: T.MeshTensor((M, K), shard_policy, dtype, layout=A_layout), # type: ignore + B: T.MeshTensor((K, N), shard_policy, dtype, layout=B_layout), # type: ignore + C: T.MeshTensor((M, N), shard_policy, accum_dtype, layout=C_layout), # type: ignore + ): + with T.Kernel() as _cid: + sharded_M, sharded_K = A.local_shape + sharded_N = B.local_shape[1] + A_shared_dist = T.alloc_shared((block_M, block_K * T.ncols()), dtype) + B_shared_dist = T.alloc_shared((block_K * T.nrows(), block_N), dtype) + C_shared = T.alloc_shared((block_M, block_N), accum_dtype) + + for by in T.serial(T.ceildiv(sharded_M, block_M)): + for bx in T.serial(T.ceildiv(sharded_N, block_N)): + T.clear(C_shared) + for k in T.Pipelined(T.ceildiv(sharded_K, block_K), num_stages=3): + T.comm.all_gather( + A[ + by * block_M : (by + 1) * block_M, + k * block_K : (k + 1) * block_K, + ], + A_shared_dist, + direction="horizontal", + axis=-1, + ) + T.comm.all_gather( + B[ + k * block_K : (k + 1) * block_K, + bx * block_N : (bx + 1) * block_N, + ], + B_shared_dist, + direction="vertical", + axis=0, + ) + T.gemm(A_shared_dist, B_shared_dist, C_shared) + T.copy(C_shared, C[by * block_M, bx * block_N]) + + return main + + +@target("Sunmmio") +def mesh_matmul_new( + M, + N, + K, + block_M=128, + block_N=128, + block_K=32, + num_stages=2, + dtype="bfloat16", + accum_dtype="float", +): + # 1) Sharding policy + # GEMM matrices are the easiest case: + # - row mesh dimension shards tensor dim 0 + # - col mesh dimension shards tensor dim 1 + # Use the same effective partitioning as the executable strict test: + # A is row-sharded on M and replicated across mesh rows for K traversal, + # B is col-sharded on N and replicated across mesh cols for K traversal, + # C is sharded on both output axes. + a_policy = T.placement.row_shard(0) + b_policy = T.placement.col_shard(1) + c_policy = T.placement.full_shard(0, 1) + + # 2) Layout + # For GEMM, the compute-critical dimensions are exactly the matrix axes. + # So we block both dimensions with a 32x32 ZZ layout. + A_layout = make_zz_layout((M, K), [0, 1], (32, 32)) + B_layout = make_zz_layout((K, N), [0, 1], (32, 32)) + C_layout = make_zz_layout((M, N), [0, 1], (32, 32)) + + @T.prim_func + def main( + # 3) MeshTensor declaration + # New-style intent: + # T.MeshTensor(shape, shard_policy, dtype, layout=...) + # This says: + # - logical global shape is `(M, K)` + # - sharding is controlled by `shard_policy` + # - layout is explicit + A: T.MeshTensor((M, K), a_policy, (4, 4), dtype, layout=A_layout), + B: T.MeshTensor((K, N), b_policy, (4, 4), dtype, layout=B_layout), + C: T.MeshTensor((M, N), c_policy, (4, 4), accum_dtype, layout=C_layout), + ): + with T.Kernel() as _cid: + sharded_M, sharded_K = A.local_shape + _, sharded_N = B.local_shape + + A_shared = T.alloc_shared((block_M, block_K), dtype) + B_shared = T.alloc_shared((block_K, block_N), dtype) + C_shared = T.alloc_shared((block_M, block_N), accum_dtype) + + for bx in T.serial(T.ceildiv(sharded_M, block_M)): + for by in T.serial(T.ceildiv(sharded_N, block_N)): + T.clear(C_shared) + for k in T.Pipelined(T.ceildiv(sharded_K, block_K), num_stages=num_stages): + T.copy( + A[ + bx * block_M : (bx + 1) * block_M, + k * block_K : (k + 1) * block_K, + ], + A_shared, + ) + T.copy( + B[ + k * block_K : (k + 1) * block_K, + by * block_N : (by + 1) * block_N, + ], + B_shared, + ) + T.gemm(A_shared, B_shared, C_shared) + T.copy(C_shared, C[bx * block_M, by * block_N]) + + return main + + +@target("Sunmmio") +def mesh_ffn_new( + seq=128, + hidden=512, + inner_dim=512, + block_seq=32, + block_hidden=32, + block_inner=32, + num_stages=2, + dtype="bfloat16", + accum_dtype="float", +): + """Small two-projection FFN used by the SunMMIO pipeline pass tests.""" + activation_policy = T.placement.full_shard(0, 1) + weight_policy = T.placement.full_shard(0, 1) + + x_shape = (seq, hidden) + up_weight_shape = (hidden, inner_dim) + mid_shape = (seq, inner_dim) + down_weight_shape = (inner_dim, hidden) + + @T.prim_func + def main( + X: T.MeshTensor(x_shape, activation_policy, dtype, layout=make_zz_layout(x_shape)), + WUp: T.MeshTensor( + up_weight_shape, + weight_policy, + dtype, + layout=make_zz_layout(up_weight_shape), + ), + WDown: T.MeshTensor( + down_weight_shape, + weight_policy, + dtype, + layout=make_zz_layout(down_weight_shape), + ), + Mid: T.MeshTensor(mid_shape, activation_policy, dtype, layout=make_zz_layout(mid_shape)), + Y: T.MeshTensor(x_shape, activation_policy, accum_dtype, layout=make_zz_layout(x_shape)), + ): + with T.Kernel(T.mesh_ncores()): + lhs_local = T.alloc_shared((block_seq, block_hidden), dtype, scope="shared.rsram") + up_local = T.alloc_shared((block_hidden, block_inner), dtype, scope="shared.rsram") + lhs_shared = T.alloc_shared((block_seq, block_hidden * T.ncols()), dtype) + up_shared = T.alloc_shared((block_hidden * T.nrows(), block_inner), dtype) + mid_acc = T.alloc_shared((block_seq, block_inner), accum_dtype, scope="shared.rsram") + mid_tile = T.alloc_shared((block_seq, block_inner), dtype, scope="shared.rsram") + + mid_local = T.alloc_shared((block_seq, block_inner), dtype, scope="shared.rsram") + down_local = T.alloc_shared((block_inner, block_hidden), dtype, scope="shared.rsram") + mid_shared = T.alloc_shared((block_seq, block_inner * T.ncols()), dtype) + down_shared = T.alloc_shared((block_inner * T.nrows(), block_hidden), dtype) + out_acc = T.alloc_shared((block_seq, block_hidden), accum_dtype, scope="shared.rsram") + + hidden_blocks = T.ceildiv(X.local_shape[1], block_hidden) + inner_blocks = T.ceildiv(Mid.local_shape[1], block_inner) + for bm in T.serial(T.ceildiv(X.local_shape[0], block_seq)): + for bn in T.serial(inner_blocks): + T.clear(mid_acc) + for bk in T.Pipelined(hidden_blocks, num_stages=num_stages): + T.copy( + X[bm * block_seq : (bm + 1) * block_seq, bk * block_hidden : (bk + 1) * block_hidden], + lhs_local, + ) + T.copy( + WUp[bk * block_hidden : (bk + 1) * block_hidden, bn * block_inner : (bn + 1) * block_inner], + up_local, + ) + T.comm.all_gather(lhs_local, lhs_shared, direction="horizontal", axis=-1) + T.comm.all_gather(up_local, up_shared, direction="vertical", axis=0) + T.gemm(lhs_shared, up_shared, mid_acc) + for i, j in T.Tiles(mid_tile, parallel=True): + mid_tile[i, j] = T.Cast(dtype, T.max(mid_acc[i, j], T.float32(0))) + T.copy(mid_tile, Mid[bm * block_seq, bn * block_inner]) + + for bh in T.serial(hidden_blocks): + T.clear(out_acc) + for bn in T.Pipelined(inner_blocks, num_stages=num_stages): + T.copy( + Mid[bm * block_seq : (bm + 1) * block_seq, bn * block_inner : (bn + 1) * block_inner], + mid_local, + ) + T.copy( + WDown[bn * block_inner : (bn + 1) * block_inner, bh * block_hidden : (bh + 1) * block_hidden], + down_local, + ) + T.comm.all_gather(mid_local, mid_shared, direction="horizontal", axis=-1) + T.comm.all_gather(down_local, down_shared, direction="vertical", axis=0) + T.gemm(mid_shared, down_shared, out_acc) + T.copy(out_acc, Y[bm * block_seq, bh * block_hidden]) + + return main + + +@target("Sunmmio") +def mesh_flashattn_new( + batch=1, + heads=64, + seq_len=4096, + dim=128, + groups=16, + is_causal=False, + block_M=64, + block_N=64, + num_stages=2, + dtype=T.bfloat16, + accum_dtype=T.bfloat16, +): + scale = (1.0 / dim) ** 0.5 * 1.44269504 + head_kv = heads // groups + q_shape = [batch, seq_len, heads, dim] + kv_shape = [batch, seq_len, head_kv, dim] + + # 1) Sharding policy + # Attention tensors are not matrices. The usual choice is: + # - shard batch on mesh rows + # - shard heads / kv-heads on mesh cols + q_policy = T.placement.full_shard(0, 2) + kv_policy = T.placement.full_shard(0, 2) + + # 2) Layout + # For Q/K/V/O the compute-critical axes are typically sequence and dim. + Q_layout = make_zz_layout(q_shape, [1, 3], (32, 32)) + K_layout = make_zz_layout(kv_shape, [1, 3], (32, 32)) + V_layout = make_zz_layout(kv_shape, [1, 3], (32, 32)) + O_layout = make_zz_layout(q_shape, [1, 3], (32, 32)) + + @T.prim_func + def main( + Q: T.MeshTensor(q_shape, q_policy, (4, 4), dtype, layout=Q_layout), + K: T.MeshTensor(kv_shape, kv_policy, (4, 4), dtype, layout=K_layout), + V: T.MeshTensor(kv_shape, kv_policy, (4, 4), dtype, layout=V_layout), + Output: T.MeshTensor(q_shape, q_policy, (4, 4), dtype, layout=O_layout), + ): + with T.Kernel() as _cid: + sharded_batch = Q.local_shape[0] + sharded_heads = Q.local_shape[2] + local_q_tiles = T.ceildiv(Q.local_shape[1], block_M) + + Q_shared = T.alloc_shared([block_M, dim], dtype) + K_shared = T.alloc_shared([block_N, dim], dtype) + V_shared = T.alloc_shared([block_N, dim], dtype) + O_shared = T.alloc_shared([block_M, dim], dtype) + acc_s = T.alloc_shared([block_M, block_N], accum_dtype) + acc_s_cast = T.alloc_shared([block_M, block_N], dtype) + acc_o = T.alloc_shared([block_M, dim], accum_dtype) + scores_max = T.alloc_shared([block_M], accum_dtype) + scores_max_prev = T.alloc_shared([block_M], accum_dtype) + scores_scale = T.alloc_shared([block_M], accum_dtype) + scores_sum = T.alloc_shared([block_M], accum_dtype) + logsum = T.alloc_shared([block_M], accum_dtype) + + for bz in T.serial(sharded_batch): + for by in T.serial(sharded_heads): + for bx in T.serial(local_q_tiles): + T.copy(Q[bz, bx * block_M : (bx + 1) * block_M, by, :], Q_shared) + T.fill(acc_o, 0) + T.fill(logsum, 0) + T.fill(scores_max, -T.infinity(accum_dtype)) + + loop_range = ( + T.min(T.ceildiv(K.local_shape[1], block_N), T.ceildiv((bx + 1) * block_M, block_N)) + if is_causal + else T.ceildiv(K.local_shape[1], block_N) + ) + + for k in T.Pipelined(loop_range, num_stages=num_stages): + T.copy(K[bz, k * block_N : (k + 1) * block_N, by // groups, :], K_shared) + if is_causal: + for i, j in T.Tiles(acc_s, parallel=True): + acc_s[i, j] = T.if_then_else( + bx * block_M + i >= k * block_N + j, + 0, + -T.infinity(acc_s.dtype), + ) + else: + for i, j in T.Tiles(acc_s, parallel=True): + acc_s[i, j] = T.if_then_else( + k * block_N + j >= seq_len, + -T.infinity(acc_s.dtype), + 0, + ) + + T.gemm(Q_shared, K_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow) + T.copy(scores_max, scores_max_prev) + T.fill(scores_max, -T.infinity(accum_dtype)) + T.reduce_max(acc_s, scores_max, dim=1, clear=False) + for i in T.Tiles(scores_max, parallel=True): + scores_max[i] = T.max(scores_max[i], scores_max_prev[i]) + for i in T.Tiles(scores_scale, parallel=True): + scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale) + for i, j in T.Tiles(acc_s, parallel=True): + acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale) + T.reduce_sum(acc_s, scores_sum, dim=1) + for i in T.Tiles(logsum, parallel=True): + logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i] + T.copy(acc_s, acc_s_cast) + + for i, j in T.Tiles(acc_o, parallel=True): + acc_o[i, j] *= scores_scale[i] + + T.copy(V[bz, k * block_N : (k + 1) * block_N, by // groups, :], V_shared) + T.gemm(acc_s_cast, V_shared, acc_o, policy=T.GemmWarpPolicy.FullRow) + + for i, j in T.Tiles(acc_o, parallel=True): + acc_o[i, j] /= logsum[i] + T.copy(acc_o, O_shared) + T.copy(O_shared, Output[bz, bx * block_M : (bx + 1) * block_M, by, :]) + + return main + + +@target("Sunmmio") +def mesh_flashdecoding_new( + batch=1, + heads=256, + kv_heads=8, + seqlen_kv=8192, + dim=128, + block_N=128, + block_H=64, + num_split=1, + num_stages=2, + dtype=T.bfloat16, + accum_dtype=T.bfloat16, +): + scale = (1.0 / dim) ** 0.5 * 1.44269504 + shape_q = [batch, heads, dim] + shape_k = [batch, seqlen_kv, kv_heads, dim] + shape_v = [batch, seqlen_kv, kv_heads, dim] + shape_o = [batch, heads, dim] + assert heads % kv_heads == 0, "GQA requires kv_heads to divide heads" + + @T.prim_func + def main( + Q: T.MeshTensor(shape_q, T.placement.full_shard(0, 1), dtype, layout=make_zz_layout(shape_q)), + K: T.MeshTensor(shape_k, T.placement.full_shard(0, 2), dtype, layout=make_zz_layout(shape_k, axes=(1, 3))), + V: T.MeshTensor(shape_v, T.placement.full_shard(0, 2), dtype, layout=make_zz_layout(shape_k, axes=(1, 3))), + mask: T.MeshTensor( + [batch, seqlen_kv], + T.placement.row_shard(0), + dtype, + layout=make_row_major([batch, seqlen_kv]), + ), + Output: T.MeshTensor(shape_o, T.placement.full_shard(0, 1), dtype, layout=make_zz_layout(shape_o)), + ): + with T.Kernel() as (_cid): + sharded_batch, sharded_heads, _ = Q.local_shape + + Q_shared = T.alloc_shared([block_H, dim], dtype) + K_shared = T.alloc_shared([block_N, dim], dtype) + V_shared = T.alloc_shared([block_N, dim], dtype) + acc_s = T.alloc_shared([block_H, block_N], accum_dtype) + acc_s_cast = T.alloc_shared([block_H, block_N], dtype) + mask_local = T.alloc_shared([block_N], dtype) + acc_o = T.alloc_shared([block_H, dim], accum_dtype) + acc_o_cast = T.alloc_shared([block_H, dim], dtype) + scores_max = T.alloc_shared([block_H], accum_dtype) + scores_max_prev = T.alloc_shared([block_H], accum_dtype) + scores_scale = T.alloc_shared([block_H], accum_dtype) + scores_sum = T.alloc_shared([block_H], accum_dtype) + logsum = T.alloc_shared([block_H], accum_dtype) + + for bid in T.serial(sharded_batch): + for hid in T.serial(T.ceildiv(sharded_heads, block_H)): + cur_kv_head = hid + T.copy(Q[bid, hid * block_H : (hid + 1) * block_H, :], Q_shared) + T.fill(acc_o, 0) + T.fill(logsum, 0) + T.fill(scores_max, -T.infinity(accum_dtype)) + + loop_range = T.ceildiv(seqlen_kv, block_N) + for k in T.Pipelined(loop_range, num_stages=num_stages): + T.copy(K[bid, k * block_N : (k + 1) * block_N, cur_kv_head, :], K_shared) + T.copy(mask[bid, k * block_N : (k + 1) * block_N], mask_local) + T.gemm(Q_shared, K_shared, acc_s, clear_accum=True, transpose_B=True) # Not accmulate + for i, j in T.Tiles([block_H, block_N]): + acc_s[i, j] = T.if_then_else(mask_local[j] != 0, acc_s[i, j], -T.infinity(accum_dtype)) + T.copy(scores_max, scores_max_prev) + T.fill(scores_max, -T.infinity(accum_dtype)) + T.reduce_max(acc_s, scores_max, dim=1, clear=False) + for i in T.Tiles([block_H]): + scores_max[i] = T.max(scores_max[i], scores_max_prev[i]) + for i in T.Tiles([block_H]): + scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale) + for i, j in T.Tiles([block_H, block_N]): + acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale) + T.reduce_sum(acc_s, scores_sum, dim=1) + for i in T.Tiles([block_H]): + logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i] + T.copy(acc_s, acc_s_cast) + for i, j in T.Tiles([block_H, dim]): + acc_o[i, j] *= scores_scale[i] + T.copy(V[bid, k * block_N : (k + 1) * block_N, cur_kv_head, :], V_shared) + T.gemm(acc_s_cast, V_shared, acc_o) + + for i, j in T.Tiles([block_H, dim]): + acc_o_cast[i, j] = acc_o[i, j] / logsum[i] + T.copy(acc_o_cast, Output[bid, hid * block_H : (hid + 1) * block_H, :]) + + return main + + +@target("Sunmmio") +def mesh_flashmladecode_new( + batch=1, + heads=128, + kv_head_num=1, + seqlen_kv=8192, + dim=512, + pe_dim=64, + block_N=64, + block_H=64, + num_split=1, + softmax_scale=1 / 24, + num_stages=2, + dtype=T.bfloat16, + accum_dtype=T.bfloat16, +): + scale = float(softmax_scale * 1.44269504) + kv_group_num = heads // kv_head_num + valid_block_H = min(block_H, kv_group_num) + + q_shape = [batch, heads, dim] + qpe_shape = [batch, heads, pe_dim] + kv_shape = [batch, seqlen_kv, kv_head_num, dim] + kpe_shape = [batch, seqlen_kv, kv_head_num, pe_dim] + glse_shape = [batch, heads, num_split] + part_shape = [batch, heads, num_split, dim] + out_shape = [batch, heads, dim] + + # 1) Sharding policy + # Q / Q_pe / Output are head-oriented tensors. + # KV / K_pe are kv-stream tensors. + q_policy = T.placement.full_shard(0, 1) + kv_policy = T.placement.full_shard(0, 2) + + # 2) Layout + # Q and Q_pe use head-oriented layouts. + # KV and K_pe use seq + dim style layouts. + Q_layout = make_zz_layout(q_shape, [1, 2], (32, 32)) + Qpe_layout = make_zz_layout(qpe_shape, [1, 2], (32, 32)) + KV_layout = make_zz_layout(kv_shape, [1, 3], (32, 32)) + Kpe_layout = make_zz_layout(kpe_shape, [1, 3], (32, 32)) + O_layout = make_zz_layout(out_shape, [1, 2], (32, 32)) + + @T.prim_func + def main( + Q: T.MeshTensor(q_shape, q_policy, (4, 4), dtype, layout=Q_layout), + Q_pe: T.MeshTensor(qpe_shape, q_policy, (4, 4), dtype, layout=Qpe_layout), + KV: T.MeshTensor(kv_shape, kv_policy, (4, 4), dtype, layout=KV_layout), + K_pe: T.MeshTensor(kpe_shape, kv_policy, (4, 4), dtype, layout=Kpe_layout), + glse: T.MeshTensor(glse_shape, q_policy, (4, 4), dtype), + Output_partial: T.MeshTensor(part_shape, q_policy, (4, 4), dtype), + Output: T.MeshTensor(out_shape, q_policy, (4, 4), dtype, layout=O_layout), + ): + with T.Kernel() as _cid: + # This kernel has two score-producing paths: + # Q_shared @ KV_shared^T + # Q_pe_shared @ K_pe_shared^T + # + # That is why the buffer roles matter more here than in plain flashattn. + Q_shared = T.alloc_shared([block_H, dim], dtype) + Q_pe_shared = T.alloc_shared([block_H, pe_dim], dtype) + KV_shared = T.alloc_shared([block_N, dim], dtype) + KV_shared2 = T.alloc_shared([block_N, dim], dtype) + K_pe_shared = T.alloc_shared([block_N, pe_dim], dtype) + S_shared = T.alloc_shared([block_H, block_N], dtype) + O_shared = T.alloc_shared([block_H, dim], dtype) + acc_s = T.alloc_shared([block_H, block_N], accum_dtype) + acc_o = T.alloc_shared([block_H, dim], accum_dtype) + scores_max = T.alloc_shared([block_H], accum_dtype) + scores_max_prev = T.alloc_shared([block_H], accum_dtype) + scores_scale = T.alloc_shared([block_H], accum_dtype) + scores_sum = T.alloc_shared([block_H], accum_dtype) + logsum = T.alloc_shared([block_H], accum_dtype) + + sharded_batch = Q.local_shape[0] + sharded_heads = Q.local_shape[1] + + for bid in T.serial(sharded_batch): + for hid in T.serial(T.ceildiv(sharded_heads, valid_block_H)): + for sid in T.serial(num_split): + cur_kv_head = hid // (kv_group_num // block_H) + + T.copy(Q[bid, hid * valid_block_H : (hid + 1) * valid_block_H, :], Q_shared) + T.copy(Q_pe[bid, hid * valid_block_H : (hid + 1) * valid_block_H, :], Q_pe_shared) + T.fill(acc_o, 0) + T.fill(logsum, 0) + T.fill(scores_max, -T.infinity(accum_dtype)) + + loop_range = T.ceildiv((seqlen_kv // num_split), block_N) + for k in T.Pipelined(loop_range, num_stages=num_stages): + kv_start = (seqlen_kv // num_split) * sid + k * block_N + kv_end = (seqlen_kv // num_split) * sid + (k + 1) * block_N + T.copy(KV[bid, kv_start:kv_end, cur_kv_head, :], KV_shared) + T.copy(KV[bid, kv_start:kv_end, cur_kv_head, :], KV_shared2) + T.copy(K_pe[bid, kv_start:kv_end, cur_kv_head, :], K_pe_shared) + T.clear(acc_s) + T.gemm(Q_shared, KV_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullCol) + T.gemm(Q_pe_shared, K_pe_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullCol) + T.copy(scores_max, scores_max_prev) + T.fill(scores_max, -T.infinity(accum_dtype)) + T.reduce_max(acc_s, scores_max, dim=1, clear=False) + for i in T.Tiles(scores_max, parallel=True): + scores_max[i] = T.max(scores_max[i], scores_max_prev[i]) + for i in T.Tiles(scores_scale, parallel=True): + scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale) + for i, j in T.Tiles(acc_s, parallel=True): + acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale) + T.reduce_sum(acc_s, scores_sum, dim=1) + T.copy(acc_s, S_shared) + for i in T.Tiles(logsum, parallel=True): + logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i] + for i, j in T.Tiles(acc_o, parallel=True): + acc_o[i, j] *= scores_scale[i] + T.gemm(S_shared, KV_shared2, acc_o, policy=T.GemmWarpPolicy.FullCol) + + for i, j in T.Tiles(acc_o, parallel=True): + acc_o[i, j] /= logsum[i] + for i in T.Tiles(logsum, parallel=True): + logsum[i] = T.log2(logsum[i]) + scores_max[i] * scale + T.copy(logsum, glse[bid, hid * valid_block_H : (hid + 1) * valid_block_H, sid]) + T.copy(acc_o, O_shared) + T.copy(O_shared, Output_partial[bid, hid * valid_block_H : (hid + 1) * valid_block_H, sid, :]) + + return main diff --git a/testing/python/transform/test_sunmmio_pipeline_greedy_collective.py b/testing/python/transform/test_sunmmio_pipeline_greedy_collective.py new file mode 100644 index 0000000000..86d07e7ba7 --- /dev/null +++ b/testing/python/transform/test_sunmmio_pipeline_greedy_collective.py @@ -0,0 +1,327 @@ +import json + +import tilelang as tl +import tilelang.language as T +from tilelang import tvm +from tilelang.engine.phase import should_force_let_inline +from tilelang.layout import make_zz_layout +from tilelang.utils.target import SUNMMIO_TARGET_DESC +from testing.python.sunmmio.common.compile_pipeline import target +from tvm import tir + + +@target("Sunmmio") +def mesh_ffn_new( + seq=128, + hidden=512, + inner_dim=512, + block_seq=32, + block_hidden=32, + block_inner=32, + num_stages=2, + dtype="bfloat16", + accum_dtype="float", +): + activation_policy = T.placement.full_shard(0, 1) + weight_policy = T.placement.full_shard(0, 1) + x_shape = (seq, hidden) + up_weight_shape = (hidden, inner_dim) + mid_shape = (seq, inner_dim) + down_weight_shape = (inner_dim, hidden) + + @T.prim_func + def main( + X: T.MeshTensor(x_shape, activation_policy, dtype, layout=make_zz_layout(x_shape)), + WUp: T.MeshTensor(up_weight_shape, weight_policy, dtype, layout=make_zz_layout(up_weight_shape)), + WDown: T.MeshTensor(down_weight_shape, weight_policy, dtype, layout=make_zz_layout(down_weight_shape)), + Mid: T.MeshTensor(mid_shape, activation_policy, dtype, layout=make_zz_layout(mid_shape)), + Y: T.MeshTensor(x_shape, activation_policy, accum_dtype, layout=make_zz_layout(x_shape)), + ): + with T.Kernel(T.mesh_ncores()): + lhs_local = T.alloc_shared((block_seq, block_hidden), dtype, scope="shared.rsram") + up_local = T.alloc_shared((block_hidden, block_inner), dtype, scope="shared.rsram") + lhs_shared = T.alloc_shared((block_seq, block_hidden * T.ncols()), dtype) + up_shared = T.alloc_shared((block_hidden * T.nrows(), block_inner), dtype) + mid_acc = T.alloc_shared((block_seq, block_inner), accum_dtype, scope="shared.rsram") + mid_tile = T.alloc_shared((block_seq, block_inner), dtype, scope="shared.rsram") + mid_local = T.alloc_shared((block_seq, block_inner), dtype, scope="shared.rsram") + down_local = T.alloc_shared((block_inner, block_hidden), dtype, scope="shared.rsram") + mid_shared = T.alloc_shared((block_seq, block_inner * T.ncols()), dtype) + down_shared = T.alloc_shared((block_inner * T.nrows(), block_hidden), dtype) + out_acc = T.alloc_shared((block_seq, block_hidden), accum_dtype, scope="shared.rsram") + + hidden_blocks = T.ceildiv(X.local_shape[1], block_hidden) + inner_blocks = T.ceildiv(Mid.local_shape[1], block_inner) + for bm in T.serial(T.ceildiv(X.local_shape[0], block_seq)): + for bn in T.serial(inner_blocks): + T.clear(mid_acc) + for bk in T.Pipelined(hidden_blocks, num_stages=num_stages): + T.copy( + X[bm * block_seq : (bm + 1) * block_seq, bk * block_hidden : (bk + 1) * block_hidden], + lhs_local, + ) + T.copy( + WUp[bk * block_hidden : (bk + 1) * block_hidden, bn * block_inner : (bn + 1) * block_inner], + up_local, + ) + T.comm.all_gather(lhs_local, lhs_shared, direction="horizontal", axis=-1) + T.comm.all_gather(up_local, up_shared, direction="vertical", axis=0) + T.gemm(lhs_shared, up_shared, mid_acc) + for i, j in T.Tiles(mid_tile, parallel=True): + mid_tile[i, j] = T.Cast(dtype, T.max(mid_acc[i, j], T.float32(0))) + T.copy(mid_tile, Mid[bm * block_seq, bn * block_inner]) + + for bh in T.serial(hidden_blocks): + T.clear(out_acc) + for bn in T.Pipelined(inner_blocks, num_stages=num_stages): + T.copy( + Mid[bm * block_seq : (bm + 1) * block_seq, bn * block_inner : (bn + 1) * block_inner], + mid_local, + ) + T.copy( + WDown[bn * block_inner : (bn + 1) * block_inner, bh * block_hidden : (bh + 1) * block_hidden], + down_local, + ) + T.comm.all_gather(mid_local, mid_shared, direction="horizontal", axis=-1) + T.comm.all_gather(down_local, down_shared, direction="vertical", axis=0) + T.gemm(mid_shared, down_shared, out_acc) + T.copy(out_acc, Y[bm * block_seq, bh * block_hidden]) + + return main + + +def _lower_ffn(): + target = tvm.target.Target(SUNMMIO_TARGET_DESC) + with tvm.target.Target(target): + func = mesh_ffn_new(num_stages=2).with_attr("global_symbol", "main") + mod = tir.transform.BindTarget(target)(tvm.IRModule.from_expr(func)) + mod = tl.transform.ResolveSunmmioMeshSymbols()(mod) + if should_force_let_inline(): + mod = tl.transform.LetInline()(mod) + for pipeline_pass in ( + tl.transform.LegalizeNegativeIndex(), + tl.transform.InjectAssumes(), + tl.transform.Simplify(), + tl.transform.InferSramScope(), + tl.transform.LegalizeSunmmioDataPath(), + tl.transform.SunmmioLayoutInference(), + tl.transform.LegalizeSunmmioGemm(), + tl.transform.LowerTileOp(), + tl.transform.LegalizeTilesLoop(), + tl.transform.TilesLoop(), + tl.transform.LegalizeVectorizedLoop(), + tl.transform.LegalizeSafeMemoryAccess(), + tl.transform.LowerAccessPtr(), + tl.transform.Simplify(), + tl.transform.HoistNonRestrictParams(), + tl.transform.HoistBlockAnnotationsToFuncAttrs(), + ): + mod = pipeline_pass(mod) + return tl.transform.IfStmtBinding()(mod) + + +def _pipeline_loops(stmt): + loops = [] + + def visit(node): + if node is None: + return + if isinstance(node, tir.For): + if node.annotations and "tl.sunmmio.pipeline.requested" in node.annotations: + loops.append(node) + visit(node.body) + elif isinstance(node, tir.BlockRealize): + visit(node.block.body) + elif isinstance(node, tir.Block): + visit(node.body) + elif isinstance(node, tir.SeqStmt): + for child in node.seq: + visit(child) + elif isinstance(node, tir.IfThenElse): + visit(node.then_case) + visit(node.else_case) + elif isinstance(node, (tir.LetStmt, tir.AttrStmt)): + visit(node.body) + + visit(stmt) + return loops + + +def _make_dynamic_injector_fixture(iterations, schedules): + extent = tir.Var("extent", "int32") + loop_var = tir.Var("i", "int32") + commands = tir.SeqStmt( + [tir.Evaluate(tir.call_extern("int32", "dynamic_epilogue_marker", loop_var, command_id)) for command_id in range(2)] + ) + body_orders = [f"{iteration}-{command_id}" for iteration in range(iterations) for command_id in range(2)] + dynamic_orders = {tir.IntImm("int32", remainder): schedule for remainder, schedule in schedules.items()} + loop = tir.For( + loop_var, + 0, + extent, + tir.ForKind.SERIAL, + commands, + annotations={ + "iterations": iterations, + "used_buffers": [], + "versioned_buffers": [], + "prologue_orders": [], + "body_orders": body_orders, + "dynamic_epilogue_orders": dynamic_orders, + }, + ) + func = tir.PrimFunc([extent], loop).with_attr("global_symbol", "main") + return tvm.IRModule.from_expr(func), extent + + +def _dynamic_dispatch_branches(dispatch, extent, iterations): + branches = {} + current = dispatch + for remainder in range(iterations - 1): + assert isinstance(current, tir.IfThenElse) + expected_condition = tir.floormod(extent, iterations) == remainder + assert tvm.ir.structural_equal(current.condition, expected_condition, map_free_vars=True) + branches[remainder] = current.then_case + current = current.else_case + branches[iterations - 1] = current + return branches + + +def _epilogue_marker_order(stmt, extent, iterations): + base = tir.floordiv(tir.max(0, extent - 1), iterations) * iterations + analyzer = tvm.arith.Analyzer() + order = [] + + def visit(node): + if not isinstance(node, tir.Call) or not isinstance(node.op, tvm.ir.Op): + return + if node.op.name != "tir.call_extern" or node.args[0].value != "dynamic_epilogue_marker": + return + logical_iteration = analyzer.simplify(node.args[1] - base) + assert isinstance(logical_iteration, tir.IntImm) + order.append((int(logical_iteration), int(node.args[2]))) + + tir.stmt_functor.post_order_visit(stmt, visit) + return order + + +def test_dynamic_epilogue_dispatch_preserves_each_remainder_schedule(): + schedules = { + 0: ["0-0", "1-0", "0-1", "2-0", "1-1", "2-1"], + 1: ["0-1", "0-0"], + 2: ["1-0", "0-1", "0-0", "1-1"], + } + mod, extent = _make_dynamic_injector_fixture(3, schedules) + injected = tl.transform.InjectSunmmioPipeline()(mod) + body = injected["main"].body + assert isinstance(body, tir.SeqStmt) + branches = _dynamic_dispatch_branches(body.seq[-1], extent, 3) + + for remainder, schedule in schedules.items(): + expected = [tuple(map(int, order.split("-"))) for order in schedule] + assert _epilogue_marker_order(branches[remainder], extent, 3) == expected + + +def test_dynamic_extent_planner_and_injector_select_remainder_schedule(): + mod = _lower_ffn() + func = mod["main"] + dynamic_extent = tir.Var("dynamic_extent", "int32") + + def replace_pipeline_extent(node): + if isinstance(node, tir.For) and node.annotations and "num_stages" in node.annotations: + return tir.For( + node.loop_var, + node.min, + dynamic_extent, + node.kind, + node.body, + node.thread_binding, + node.annotations, + ) + return None + + body = tir.stmt_functor.ir_transform(func.body, replace_pipeline_extent, None, ["tir.For"]) + dynamic_func = tir.PrimFunc( + [*func.params, dynamic_extent], + body, + func.ret_type, + func.buffer_map, + func.attrs, + ) + planned = tl.transform.SunmmioPipelinePlanning(debug=False)(tvm.IRModule({"main": dynamic_func})) + planned_loops = _pipeline_loops(planned["main"].body) + assert len(planned_loops) == 2 + for loop in planned_loops: + assert bool(loop.annotations["tl.sunmmio.pipeline.applied"]) + assert {int(key) for key in loop.annotations["dynamic_epilogue_orders"]} == {0, 1} + + injected = tl.transform.InjectSunmmioPipeline()(planned) + assert not _pipeline_loops(injected["main"].body) + assert injected.script().count("dynamic_extent % 2") == 2 + + +def test_greedy_ffn_models_collective_order_and_relative_bank_precolors(monkeypatch, tmp_path): + graph_path = tmp_path / "greedy_ffn_graph.json" + monkeypatch.setenv("TL_SUNMMIO_PIPELINE_GRAPH_JSON", str(graph_path)) + planned = tl.transform.SunmmioPipelinePlanning(debug=False)(_lower_ffn()) + loops = _pipeline_loops(planned["main"].body) + assert len(loops) == 2 + + collective_ids = (2, 3, 5) + collective_rank = {command_id: rank for rank, command_id in enumerate(collective_ids)} + for loop in loops: + annotations = loop.annotations + assert bool(annotations["tl.sunmmio.pipeline.applied"]) + + body_orders = [tuple(map(int, str(order).split("-"))) for order in annotations["body_orders"]] + body_positions = {order: position for position, order in enumerate(body_orders)} + # Commands 0-5 (ODMA1), 1-0 (ODMA0), and 0-4 (TensorCore) all + # start at time zero. Async launches must be emitted before blocking MMA. + assert body_positions[(0, 5)] < body_positions[(0, 4)] + assert body_positions[(1, 0)] < body_positions[(0, 4)] + + for name in ("prologue_orders", "body_orders", "epilogue_orders"): + collective_orders = [ + tuple(map(int, str(order).split("-"))) for order in annotations[name] if int(str(order).split("-")[1]) in collective_ids + ] + assert collective_orders == sorted( + collective_orders, + key=lambda order: (order[0], collective_rank[order[1]]), + ) + + phase_maps = [ + {int(command_id): int(phase) for command_id, phase in phases.items()} + for phases in annotations["runtime_bank_writer_phases"].values() + ] + striped_writers = next(phases for phases in phase_maps if 2 in phases and 5 in phases) + assert striped_writers[2] != striped_writers[5] + + graph = json.loads(graph_path.read_text(encoding="utf-8")) + resources = {command["id"]: command["resource"] for command in graph["commands"]} + assert {command_id: resources[command_id] for command_id in collective_ids} == {2: 3, 3: 2, 5: 3} + assert [(edge["source"], edge["target"], edge["distance"]) for edge in graph["edges"] if edge["kind"] == "collective_order"] == [ + (2, 3, 0), + (3, 5, 0), + (5, 2, 1), + ] + + injected = tl.transform.InjectSunmmioPipeline()(planned) + broadcasts = [] + allocated_shapes = {} + + def collect_injected(node): + if isinstance(node, tir.Call) and isinstance(node.op, tvm.ir.Op) and node.op.name == "tl.broadcast_": + broadcasts.append(node) + if isinstance(node, tir.Block): + for buffer in node.alloc_buffers: + allocated_shapes[str(buffer.name)] = tuple(int(dim) for dim in buffer.shape) + + tir.stmt_functor.post_order_visit( + injected["main"].body, + collect_injected, + ) + assert len(broadcasts) >= 4 + assert allocated_shapes["lhs_local"] == (2, 32, 32) + assert allocated_shapes["up_local"] == (2, 32, 32) + assert allocated_shapes["mid_local"] == (2, 32, 32) + assert allocated_shapes["down_local"] == (2, 32, 32) diff --git a/testing/python/transform/test_sunmmio_pipeline_ilp_async_ops.py b/testing/python/transform/test_sunmmio_pipeline_ilp_async_ops.py new file mode 100644 index 0000000000..0445bd2cff --- /dev/null +++ b/testing/python/transform/test_sunmmio_pipeline_ilp_async_ops.py @@ -0,0 +1,100 @@ +import json +import os +from contextlib import contextmanager + +import tilelang as tl +from tilelang import tvm +from tilelang.utils.target import SUNMMIO_TARGET_DESC +from tvm import tir + + +def _region(buffer, indices, access_mask, extents): + return tir.Call( + "handle", + tvm.ir.Op.get("tl.tileop.region"), + [tir.BufferLoad(buffer, indices), access_mask, *extents], + ) + + +def _async_call(op_name, *args): + return tir.Evaluate(tir.Call("handle", tvm.ir.Op.get(op_name), list(args))) + + +def _make_layout_transform_pipeline(): + A = tir.decl_buffer((8, 16), "float32", name="A") + B = tir.decl_buffer((8, 16), "float32", name="B") + src = tir.decl_buffer((16,), "float32", name="src", scope="shared.rsram") + dst = tir.decl_buffer((16,), "float32", name="dst", scope="shared.rsram") + k = tir.Var("k", "int32") + + body = tir.SeqStmt( + [ + _async_call( + "tl.dma_copy", + _region(A, [k, 0], 1, [1, 16]), + _region(src, [0], 2, [16]), + 0, + ), + _async_call( + "tl.sunmmio_layout_transform", + _region(src, [0], 1, [16]), + _region(dst, [0], 2, [16]), + ), + _async_call( + "tl.dma_copy", + _region(dst, [0], 1, [16]), + _region(B, [k, 0], 2, [1, 16]), + 0, + ), + ] + ) + loop = tir.For( + k, + 0, + 8, + tir.ForKind.SERIAL, + body, + annotations={"num_stages": tir.IntImm("int32", 2)}, + ) + root = tir.Block([], [], [], "root", loop, alloc_buffers=[src, dst]) + return tir.PrimFunc( + [A.data, B.data], + tir.BlockRealize([], True, root), + buffer_map={A.data: A, B.data: B}, + ) + + +@contextmanager +def _scoped_env(updates): + old = {key: os.environ.get(key) for key in updates} + os.environ.update({key: str(value) for key, value in updates.items()}) + try: + yield + finally: + for key, value in old.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + +def test_ilp_models_layout_transform_destination_as_write(tmp_path): + problem_path = tmp_path / "layout_transform_ilp_problem.json" + target = tvm.target.Target(SUNMMIO_TARGET_DESC) + mod = tvm.IRModule.from_expr(_make_layout_transform_pipeline().with_attr("global_symbol", "main")) + + with ( + tvm.target.Target(target), + _scoped_env( + { + "TL_SUNMMIO_FASTER": "20", + "TL_SUNMMIO_ILP_PROBLEM_JSON": problem_path, + } + ), + ): + tl.transform.SunmmioPipelinePlanningILP(debug=False)(mod) + + problem = json.loads(problem_path.read_text(encoding="utf-8")) + assert problem["N"] == 3 + assert [1, 2] in problem["dep_edges"] + assert problem["delta"]["1,2"] == 0 diff --git a/testing/python/transform/test_sunmmio_pipeline_ilp_pass.py b/testing/python/transform/test_sunmmio_pipeline_ilp_pass.py new file mode 100644 index 0000000000..5150d3f26c --- /dev/null +++ b/testing/python/transform/test_sunmmio_pipeline_ilp_pass.py @@ -0,0 +1,238 @@ +"""Focused planning and injection coverage for the SunMMIO ILP pipeline.""" + +import json +import os +from contextlib import contextmanager +from pathlib import Path + +import pytest +import tilelang as tl +from tilelang import tvm +from tilelang.engine.phase import should_force_let_inline +from tilelang.utils.target import SUNMMIO_TARGET_DESC +from tvm import tir + +from testing.python.transform.sunmmio_mesh_kernel_new_syntax_reference import ( + mesh_ffn_new, + mesh_flashattn_new, + mesh_matmul_new, +) + +CASES = { + "gemm": lambda num_stages: mesh_matmul_new(1024, 1024, 1024, 128, 128, 32, num_stages=num_stages), + "flashattn": lambda num_stages: mesh_flashattn_new(num_stages=num_stages), + "ffn": lambda num_stages: mesh_ffn_new(num_stages=num_stages), +} + + +def _lower_and_legalize(mod, target): + mod = tir.transform.BindTarget(target)(mod) + mod = tl.transform.ResolveSunmmioMeshSymbols()(mod) + if should_force_let_inline(): + mod = tl.transform.LetInline()(mod) + mod = tl.transform.LegalizeNegativeIndex()(mod) + mod = tl.transform.InjectAssumes()(mod) + mod = tl.transform.Simplify()(mod) + mod = tl.transform.InferSramScope()(mod) + mod = tl.transform.LegalizeSunmmioDataPath()(mod) + mod = tl.transform.SunmmioLayoutInference()(mod) + mod = tl.transform.LegalizeSunmmioGemm()(mod) + mod = tl.transform.LowerTileOp()(mod) + mod = tl.transform.LegalizeTilesLoop()(mod) + mod = tl.transform.TilesLoop()(mod) + mod = tl.transform.LegalizeVectorizedLoop()(mod) + mod = tl.transform.LegalizeSafeMemoryAccess()(mod) + mod = tl.transform.LowerAccessPtr()(mod) + mod = tl.transform.Simplify()(mod) + mod = tl.transform.HoistNonRestrictParams()(mod) + return tl.transform.HoistBlockAnnotationsToFuncAttrs()(mod) + + +@contextmanager +def _scoped_env(updates): + old = {key: os.environ.get(key) for key in updates} + os.environ.update({key: str(value) for key, value in updates.items()}) + try: + yield + finally: + for key, value in old.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + +def _pipeline_loops(stmt): + loops = [] + + def visit(node): + if node is None: + return + if isinstance(node, tir.For): + if node.annotations and "tl.sunmmio.pipeline.requested" in node.annotations: + loops.append(node) + visit(node.body) + elif isinstance(node, tir.BlockRealize): + visit(node.block.body) + elif isinstance(node, tir.Block): + visit(node.body) + elif isinstance(node, tir.SeqStmt): + for child in node.seq: + visit(child) + elif isinstance(node, tir.IfThenElse): + visit(node.then_case) + visit(node.else_case) + elif isinstance(node, (tir.AttrStmt, tir.LetStmt)): + visit(node.body) + + visit(stmt) + return loops + + +def _lower(case_name, num_stages): + target = tvm.target.Target(SUNMMIO_TARGET_DESC) + with tvm.target.Target(target): + func = CASES[case_name](num_stages).with_attr("global_symbol", "main") + mod = tvm.IRModule.from_expr(func) + mod = _lower_and_legalize(mod, target) + return tl.transform.IfStmtBinding()(mod) + + +def _output_dir(tmp_path, case_name, num_stages, shrink): + configured_root = os.environ.get("SUNMMIO_ILP_PASS_TEST_OUTPUT") + root = Path(configured_root) if configured_root else tmp_path + return root / case_name / f"stage{num_stages}_shrink_{'on' if shrink else 'off'}" + + +def _write_ir(path, mod): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(mod.script(show_meta=True).strip() + "\n", encoding="utf-8") + + +@pytest.mark.parametrize("case_name", CASES, ids=CASES) +@pytest.mark.parametrize("num_stages", (2, 3), ids=lambda value: f"stage{value}") +@pytest.mark.parametrize("shrink", (False, True), ids=("no_shrink", "shrink")) +def test_sunmmio_pipeline_ilp_planning_matrix(tmp_path, case_name, num_stages, shrink): + """Plan 3 kernels at stages 2/3 with stage shrinking both disabled/enabled.""" + output_dir = _output_dir(tmp_path, case_name, num_stages, shrink) + output_dir.mkdir(parents=True, exist_ok=True) + problem_path = output_dir / "ilp_problem.json" + solution_path = output_dir / "ilp_solution.json" + + mod = _lower(case_name, num_stages) + _write_ir(output_dir / "00_before_planning.py", mod) + with ( + tl.transform.PassContext(config={tl.PassConfigKey.TL_SUNMMIO_ILP_STAGE_SHRINK: shrink}), + _scoped_env( + { + "TL_SUNMMIO_FASTER": "200", + "TL_SUNMMIO_ILP_PROBLEM_JSON": problem_path, + "TL_SUNMMIO_ILP_SOLUTION_JSON": solution_path, + } + ), + ): + planned = tl.transform.SunmmioPipelinePlanningILP(debug=False)(mod) + _write_ir(output_dir / "01_after_planning.py", planned) + + loops = _pipeline_loops(planned["main"].body) + assert loops + for loop in loops: + annotations = loop.annotations + assert bool(annotations["tl.sunmmio.pipeline.requested"]) + assert bool(annotations["tl.sunmmio.pipeline.applied"]) + assert str(annotations["tl.sunmmio.pipeline.mode"]) == "ilp" + iterations = int(annotations["iterations"]) + assert 1 <= iterations <= num_stages + if not shrink: + assert iterations == num_stages + assert annotations["body_orders"] + assert "runtime_bank_flip_modes" in annotations + + problem_paths = sorted(output_dir.glob("ilp_problem*.json")) + assert problem_paths + assert solution_path.is_file() + solution = json.loads(solution_path.read_text(encoding="utf-8")) + assert int(solution["ii"]) > 0 + assert solution["nodes"] + assert solution["flows"] + + +def test_sunmmio_pipeline_ilp_inject_ffn_stage2(): + """Exercise the injector separately on FFN's two collective pipelines.""" + mod = _lower("ffn", 2) + with tl.transform.PassContext(config={tl.PassConfigKey.TL_SUNMMIO_ILP_STAGE_SHRINK: False}): + planned = tl.transform.SunmmioPipelinePlanningILP(debug=False)(mod) + injected = tl.transform.InjectSunmmioPipelineILP()(planned) + + assert len(_pipeline_loops(planned["main"].body)) == 2 + script = injected.script(show_meta=True) + assert '"tl.sunmmio.pipeline.fallback_reason"' not in script + assert "_ping" in script + assert "_pong" in script + assert script.count("T.mma_sunmmio(") >= 2 + broadcasts = [] + tir.stmt_functor.post_order_visit( + injected["main"].body, + lambda node: broadcasts.append(node) + if isinstance(node, tir.Call) and isinstance(node.op, tvm.ir.Op) and node.op.name == "tl.broadcast_" + else None, + ) + assert len(broadcasts) >= 4 + + +def _make_shifted_multiversion_pipeline(): + k = tir.Var("k", "int32") + output = tir.decl_buffer((16,), "int32", name="output") + scratch = tir.decl_buffer((16,), "int32", name="scratch", scope="shared") + + producer = tir.BufferStore(scratch, tir.IntImm("int32", 7), [k + 1]) + consumer = tir.BufferStore(output, tir.BufferLoad(scratch, [k]), [k]) + annotations = { + "iterations": 3, + "ii": 1, + "makespan": 2, + "steady_state_max_iter_offset": 0, + "used_buffers": [scratch], + "versioned_buffers": [scratch], + "runtime_multiversion_buffers": [scratch], + "runtime_banked_buffers": [], + "runtime_resident_banked_buffers": [], + "runtime_bank_start_phases": {}, + "runtime_bank_read_delta_parities": {}, + "runtime_bank_writer_phases": {}, + "runtime_bank_reader_phases": {}, + "runtime_bank_flip_modes": {}, + "runtime_bank_peer_buffers": {}, + "prologue_orders": ["0-0"], + "body_orders": ["1-1", "1-0"], + "epilogue_orders": ["8-1"], + "tl.sunmmio.pipeline.requested": True, + "tl.sunmmio.pipeline.applied": True, + "tl.sunmmio.pipeline.mode": "ilp", + } + loop = tir.For( + k, + 0, + 8, + tir.ForKind.SERIAL, + tir.SeqStmt([producer, consumer]), + annotations=annotations, + ) + root = tir.Block([], [], [], "root", loop, alloc_buffers=[scratch]) + body = tir.BlockRealize([], tir.const(True, "bool"), root) + func = tir.PrimFunc([output.data], body, buffer_map={output.data: output}).with_attr("global_symbol", "main") + return tvm.IRModule.from_expr(func) + + +def test_sunmmio_pipeline_ilp_inject_shifted_access_version(): + """A producer at k writing value k+1 must match its consumer at k+1.""" + injected = tl.transform.InjectSunmmioPipelineILP()(_make_shifted_multiversion_pipeline()) + script = injected.script() + + # Value 1 is produced in command iteration 0 and consumed in iteration 1. + # Both accesses must select slot 1 of the three-version buffer. + assert "scratch[1, T.Add(0, 1)] = 7" in script + assert "scratch[(k + 1) % 3, k + 1]" in script + + # The next producer writes value k+2 and therefore advances to slot k+2. + assert "scratch[(k + 1 + 1) % 3, k + 1 + 1] = 7" in script diff --git a/tilelang/engine/phase.py b/tilelang/engine/phase.py index 83bb96cd0d..1537ddbfd7 100644 --- a/tilelang/engine/phase.py +++ b/tilelang/engine/phase.py @@ -71,6 +71,22 @@ def should_enable_ast_print(pass_ctx: PassContext | None = None) -> bool: return bool(pass_ctx and pass_ctx.config.get(tilelang.PassConfigKey.TL_AST_PRINT_ENABLE, False)) +def should_enable_sunmmio_pipeline(pass_ctx: PassContext | None = None) -> bool: + if pass_ctx is None: + pass_ctx = tilelang.transform.get_pass_context() + return not bool(pass_ctx.config.get(tilelang.PassConfigKey.TL_DISABLE_SUNMMIO_PIPELINE, False)) + + +def get_sunmmio_pipeline_mode(pass_ctx: PassContext | None = None) -> str: + if pass_ctx is None: + pass_ctx = tilelang.transform.get_pass_context() + mode = pass_ctx.config.get(tilelang.PassConfigKey.TL_SUNMMIO_PIPELINE_MODE, "greedy") + mode = str(mode).strip().lower() + if mode not in {"greedy", "ilp"}: + raise ValueError(f"Invalid Sunmmio pipeline mode: {mode}. Expected one of: 'greedy', 'ilp'.") + return mode + + def should_enable_layout_visual(pass_ctx: PassContext | None = None) -> bool: if pass_ctx is None: pass_ctx = tilelang.transform.get_pass_context() @@ -248,9 +264,16 @@ def LowerAndLegalize(mod: IRModule, target: Target) -> IRModule: def OptimizeForSunmmio(mod: IRModule, target: Target) -> IRModule: + pass_ctx = tilelang.transform.get_pass_context() mod = tilelang.transform.IfStmtBinding()(mod) - mod = tilelang.transform.SunmmioPipelinePlanning(debug=False)(mod) - mod = tilelang.transform.InjectSunmmioPipeline()(mod) + if should_enable_sunmmio_pipeline(pass_ctx): + pipeline_mode = get_sunmmio_pipeline_mode(pass_ctx) + if pipeline_mode == "ilp": + mod = tilelang.transform.SunmmioPipelinePlanningILP(debug=False)(mod) + mod = tilelang.transform.InjectSunmmioPipelineILP()(mod) + else: + mod = tilelang.transform.SunmmioPipelinePlanning(debug=False)(mod) + mod = tilelang.transform.InjectSunmmioPipeline()(mod) mod = tilelang.transform.LowerOpaqueBlock()(mod) mod = tilelang.transform.Simplify()(mod) @@ -263,6 +286,7 @@ def OptimizeForSunmmio(mod: IRModule, target: Target) -> IRModule: mod = tir.transform.RenormalizeSplitPattern()(mod) mod = tir.transform.Simplify()(mod) mod = tir.transform.RemoveNoOp()(mod) + mod = tilelang.transform.RemoveUnusedSunmmioAllocations()(mod) mod = tir.transform.HoistIfThenElse()(mod) mod = tir.transform.VerifyMemory()(mod) diff --git a/tilelang/language/__init__.py b/tilelang/language/__init__.py index 7037f80e30..df40c53a87 100644 --- a/tilelang/language/__init__.py +++ b/tilelang/language/__init__.py @@ -21,7 +21,7 @@ MeshTensor, TensorWithMeta, ) -from .mesh_symbols import mesh_nrows, mesh_ncols, mesh_ncores # noqa: F401 +from .mesh_symbols import mesh_nrows, mesh_ncols, mesh_ncores, nrows, ncols # noqa: F401 from .loop import ( Parallel, # noqa: F401 Tiles, # noqa: F401 diff --git a/tilelang/language/mesh_symbols.py b/tilelang/language/mesh_symbols.py index 7a281ac325..1837a91c4b 100644 --- a/tilelang/language/mesh_symbols.py +++ b/tilelang/language/mesh_symbols.py @@ -43,6 +43,16 @@ def mesh_ncols() -> PrimExpr: return _mesh_ncols_symbol() +def nrows() -> PrimExpr: + """Return the symbolic number of rows in the current Sunmmio mesh.""" + return mesh_nrows() + + +def ncols() -> PrimExpr: + """Return the symbolic number of columns in the current Sunmmio mesh.""" + return mesh_ncols() + + def mesh_ncores() -> PrimExpr: """Return the symbolic number of cores in the current Sunmmio mesh.""" return mesh_nrows() * mesh_ncols() diff --git a/tilelang/language/mesh_tensor.py b/tilelang/language/mesh_tensor.py index 2481fa1e4c..f32c0900e2 100644 --- a/tilelang/language/mesh_tensor.py +++ b/tilelang/language/mesh_tensor.py @@ -68,8 +68,8 @@ def local_shape(self): """Return the uniform physical local buffer shape.""" return self.meta_data["local_shape"] - def get_local_extent(self, cid): - """Return the valid local extent on core ``cid``.""" + def get_local_extent(self, cid=None): + """Return the valid local extent on ``cid`` or the current kernel core.""" return get_local_extent(self, cid) @@ -91,8 +91,8 @@ def local_shape(self): """Return the uniform physical local buffer shape.""" return self.meta_data["local_shape"] - def get_local_extent(self, cid): - """Return the valid local extent on core ``cid``.""" + def get_local_extent(self, cid=None): + """Return the valid local extent on ``cid`` or the current kernel core.""" return get_local_extent(self, cid) def __getitem__(self, keys): @@ -176,13 +176,20 @@ def lookup_mesh_tensor_meta(mesh_tensor): raise TypeError(f"Expected a MeshTensor value with metadata, got {type(mesh_tensor)}") -def get_local_extent(mesh_tensor, cid): +def get_local_extent(mesh_tensor, cid=None): """Return the valid local extent for ``mesh_tensor`` on linear core id ``cid``. + When ``cid`` is omitted inside a kernel, use its current block binding. + Full sharding preserves the physical mesh-axis order: row sharding is applied first, then column sharding is applied to the row-local extent. ``mesh_as_line`` instead uses the row-major linear core id. """ + if cid is None: + from tilelang.language.kernel import get_block_binding + + cid = get_block_binding(0) + meta = lookup_mesh_tensor_meta(mesh_tensor) global_shape = meta["global_shape"] nrows, ncols = meta["mesh_shape"] @@ -380,7 +387,7 @@ def __new__( sharding_policy: PlacementSpec | MeshShardingPolicy | None = None, ) -> TensorWithMeta: ... - def get_local_extent(self, cid) -> tuple[Any, ...]: ... + def get_local_extent(self, cid=None) -> tuple[Any, ...]: ... else: MeshTensor = MeshTensorProxy() diff --git a/tilelang/transform/__init__.py b/tilelang/transform/__init__.py index 00e1b4e9de..18cd55885d 100644 --- a/tilelang/transform/__init__.py +++ b/tilelang/transform/__init__.py @@ -730,6 +730,39 @@ def InjectSunmmioPipeline(): return _ffi_api.InjectSunmmioPipeline() # type: ignore +def RemoveUnusedSunmmioAllocations(): + """Remove allocations unused by executable SunMMIO TIR. + + Returns + ------- + fpass : tvm.transform.Pass + The result pass + """ + return _ffi_api.RemoveUnusedSunmmioAllocations() # type: ignore + + +def SunmmioPipelinePlanningILP(debug: bool = False): + """SunmmioPipelinePlanning + + Returns + ------- + fpass : tvm.transform.Pass + The result pass + """ + return _ffi_api.SunmmioPipelinePlanningILP(debug) # type: ignore + + +def InjectSunmmioPipelineILP(): + """InjectSunmmioPipeline + + Returns + ------- + fpass : tvm.transform.Pass + The result pass + """ + return _ffi_api.InjectSunmmioPipelineILP() # type: ignore + + def MergeSharedMemoryAllocationsSunmmio( enable_aggressive_merge: bool = False, asram_align_bytes: int = 2048, diff --git a/tilelang/transform/pass_config.py b/tilelang/transform/pass_config.py index 945fe73a52..75d6f49616 100644 --- a/tilelang/transform/pass_config.py +++ b/tilelang/transform/pass_config.py @@ -134,6 +134,27 @@ class PassConfigKey(str, Enum): optimization in cases where manual synchronization is preferred or when synchronization is not needed. Default: False""" + TL_DISABLE_SUNMMIO_PIPELINE = "tl.disable_sunmmio_pipeline" + """Disable Sunmmio pipeline planning and injection. Default: False""" + + TL_SUNMMIO_PIPELINE_MODE = "tl.sunmmio_pipeline_mode" + """Select Sunmmio pipeline implementation. Accepts: "greedy", "ilp". Default: "greedy".""" + + TL_SUNMMIO_FASTER = "tl.sunmmio_faster" + """Tune Sunmmio pipeline-planner search. In greedy mode, limit the number + of bank-coloring candidates; -1 searches all candidates. In ILP mode, use + the positive value to scale command-latency estimates. A non-positive or + unset ILP value selects the scaling factor automatically. Default: -1.""" + + TL_SUNMMIO_ILP_STAGE_SHRINK = "tl.sunmmio_ilp_stage_shrink" + """Enable Sunmmio ILP stage shrink: solve at original num_stages, then try smaller stages at fixed II and reoptimize on the minimum feasible stage. Default: False.""" + + TL_SUNMMIO_ILP_MULTIVERSION_LIFETIME_PRUNING = "tl.sunmmio_ilp_multiversion_lifetime_pruning" + """Prune unnecessary Sunmmio ILP runtime buffer versions using scheduled lifetimes. Disable to conservatively retain every theoretical version. Default: True.""" + + TL_SUNMMIO_ILP_MODEL_VC_BLOCKING_ISSUE = "tl.sunmmio_ilp_model_vc_blocking_issue" + """Model VC's blocking scalar-issue interval with lazy ILP constraints. Default: True.""" + TL_FORCE_LET_INLINE = "tl.force_let_inline" """Force TileLang to inline let bindings during simplification. Default: False"""