feat(sunmmio): integrate pipeline planning and codegen updates

.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.

.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

.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

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_OBJECTS:tilelang_objs>)
 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)

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.

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`**
 

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`**
 

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))

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

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

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]):

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(

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(