diff --git a/.claude/skills/ktir-dialect.md b/.claude/skills/ktir-dialect.md index 05a760e6..c74d2dd4 100644 --- a/.claude/skills/ktir-dialect.md +++ b/.claude/skills/ktir-dialect.md @@ -88,7 +88,7 @@ operation ::= `ktdp.construct_memory_view` $offset `,` `sizes` `:` ```mlir #set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 31 >= 0, d1 >= 0, -d1 + 64 >= 0)> %A_view = ktdp.construct_memory_view %A_start_address, sizes: [32, 64], strides: [64, 1] { - coordinate_set = #set, memory_space = #ktdp.spyre_memory_space + coordinate_set = #set, memory_space = #ktdp.memory_space } : memref<32x64xf16> ``` @@ -254,12 +254,34 @@ access_tile<1 x 64 x index> // Fully static Generic abstraction for device-specific memory space attributes within `ktdp`. Provides a uniform mechanism to associate IR values with a target-specific memory hierarchy while remaining extensible across backends. -Concrete implementations (e.g., `SpyreMemorySpaceAttr`) describe specific memory kinds (on-chip scratchpad, HBM) and optional core affinity. This is preferred over integer-based memory space annotations in memref types, which lack readability and semantic richness for distributed scratchpad architectures. +`MemorySpaceAttr` describes *where* memory is visible rather than naming device-specific +memories: `global` is reachable by every compute tile, `ct_local` is private to one, with +optional per-tile affinity via `ct_id`. This is preferred over integer-based memory space +annotations in memref types, which lack readability and semantic richness for distributed +scratchpad architectures. ```mlir -memory_space = #ktdp.spyre_memory_space +memory_space = #ktdp.memory_space +memory_space = #ktdp.memory_space +memory_space = #ktdp.memory_space ``` +> **⚠️ CONFLICTS with RFC 0682 — the RFC is stale here and needs updating.** +> The RFC specifies an abstract `KtdpMemorySpaceAttr` interface with a Spyre-specific +> `SpyreMemorySpaceAttr` implementation spelled +> `#ktdp.spyre_memory_space` (plus an `unspecified` kind). +> [ktir-mlir-frontend#58](https://github.com/torch-spyre/ktir-mlir-frontend/pull/58) +> removed that interface and renamed the attribute to the device-agnostic form above; +> `unspecified` was dropped with no replacement. This is a semantic change, not just a +> rename — the enum went from naming *devices* (HBM/LX) to naming *visibility* +> (reachable by all compute tiles vs. private to one). +> +> The RFC spelling no longer parses in the dialect at all, so **the syntax above is +> authoritative and the RFC text should be revised to match** (tracked as row 4a in +> `docs/gap_analysis.md`). Do not flag the new spelling as a spec deviation in reviews. +> `ktir-cpu` maps `global`→`HBM` and `ct_local`→`LX` at the parse boundary, keeping the +> concrete Spyre hierarchy in the interpreter and latency model. + Future extensions will include richer metadata for compute-memory affinity. --- @@ -281,17 +303,17 @@ module { %A_view = ktdp.construct_memory_view %A_start_address, sizes: [96, 64], strides: [64, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 95 >= 0, d1 >= 0, -d1 + 63 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<96x64xf16> %B_view = ktdp.construct_memory_view %B_start_address, sizes: [96, 64], strides: [64, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 95 >= 0, d1 >= 0, -d1 + 63 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<96x64xf16> %C_view = ktdp.construct_memory_view %C_start_address, sizes: [96, 64], strides: [64, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 95 >= 0, d1 >= 0, -d1 + 63 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<96x64xf16> scf.for %i = %c0 to %tile_size step %c1 { diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 43c723d2..13480edf 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -37,7 +37,16 @@ jobs: LLVM_HASH=$(curl -fsSL "https://raw.githubusercontent.com/torch-spyre/ktir-mlir-frontend/$FRONTEND_COMMIT/cmake/llvm-hash.txt") curl -fsSL "$SETUP_MLIR" -o /tmp/setup_mlir.py MLIR_DIR=$(GIT_PAT="${{ secrets.GITHUB_TOKEN }}" uv run python /tmp/setup_mlir.py --hash "$LLVM_HASH" --repo torch-spyre/ktir-mlir-frontend) - uv pip install scikit-build-core "nanobind>=2.12.0" + # The install step below uses --no-build-isolation, which bypasses + # ktir-mlir-frontend's build-system.requires — so the `cmake>=4.0` pin + # it declares upstream (torch-spyre/ktir-mlir-frontend#59) is not + # applied and must be satisfied here instead. Keep this in sync with + # that pin. cmake>=4.0 is required: CMake 3.x writes the + # full target name into its file-API reply filename, and MLIR's + # mlir_generate_type_stubs() builds a 283-char target name by joining + # DEPENDS_TARGETS, overflowing NAME_MAX (255). CMake 4.0 truncates the + # name and relies on the trailing hash. + uv pip install scikit-build-core "nanobind>=2.12.0" "cmake>=4.0" "ninja>=1.11" echo "CMAKE_ARGS=-DMLIR_DIR=$MLIR_DIR" >> $GITHUB_ENV - name: Install project and dependencies diff --git a/CLAUDE.md b/CLAUDE.md index 66c13ef5..e28c1d83 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -13,7 +13,17 @@ The full spec is also mirrored in `.claude/skills/ktir-dialect.md`. - Specifically verify: - `ktdp` dialect ops (`get_compute_tile_id`, `construct_memory_view`, `construct_distributed_memory_view`, `construct_access_tile`, `construct_indirect_access_tile`, `load`, `store`) match the spec's syntax, operands, attributes, and semantics. - `AccessTileType` element type is always `index`. - - `SpyreMemorySpaceAttr` uses `#ktdp.spyre_memory_space<...>` syntax. + - `MemorySpaceAttr` uses `#ktdp.memory_space` syntax. + **⚠️ This CONFLICTS with RFC 0682**, which still specifies a `KtdpMemorySpaceAttr` + interface with a Spyre-specific `SpyreMemorySpaceAttr` spelled + `#ktdp.spyre_memory_space` (plus an `unspecified` kind). + ktir-mlir-frontend#58 removed that interface and renamed the attribute; the RFC + has not been updated to match. The dialect is the executable authority here — the + RFC spelling no longer parses at all — so follow the syntax above and treat the + RFC as stale on this point until it is revised. Tracked as gap row 4a in + `docs/gap_analysis.md`. The interpreter still uses `HBM`/`LX` internally; + the parsers map the dialect kinds onto those names via + `KTDP_MEMORY_SPACE_KINDS` in `ktir_cpu/ir_types.py`. - `coordinate_set` uses `IntegerSetAttr` (affine integer sets). - `base_map` and `access_tile_order` use `AffineMapAttr`. - `access_tile_order` follows lexicographic semantics (rightmost = innermost). diff --git a/docs/gap_analysis.md b/docs/gap_analysis.md index af6c22eb..3e2cc4fc 100644 --- a/docs/gap_analysis.md +++ b/docs/gap_analysis.md @@ -3,7 +3,7 @@ **Date**: 2026-05-30 **Spec**: [RFC 0682 — KTIR Spec](https://github.com/torch-spyre/RFCs/blob/main/0682-KtirSpec/0682-KtirSpecRFC.md) -**Legend**: ✅ implemented — 🟡 partial — ❌ not implemented — 🧪 experimental (tracks an unmerged upstream spec PR; semantics may change) +**Legend**: ✅ implemented — 🟡 partial — ❌ not implemented — 🧪 experimental (tracks an unmerged upstream spec PR; semantics may change) — ⚠️ spec conflict (implementation intentionally diverges from the RFC; the RFC needs updating) --- @@ -23,7 +23,8 @@ | # | Spec Item | Status | Notes | |---|-----------|--------|-------| | 3 | `AccessTileType` with dynamic dimensions (`?`) | ❌ | The spec allows `access_tile` (partially/fully dynamic shapes). The parser only extracts static integer dimensions — dynamic `?` dimensions are silently dropped. | -| 4 | `MemorySpaceAttr` (generic) | 🟡 | The parser extracts `SpyreMemorySpaceAttr` (`HBM`/`LX`), but the spec describes `MemorySpaceAttr` as a generic extensible wrapper that could encapsulate other hardware backends. The implementation hardcodes Spyre-specific memory spaces only. | +| 4 | `MemorySpaceAttr` (generic) | ✅ | The RFC's concern here — that the attribute was Spyre-specific rather than a generic extensible wrapper — was addressed upstream by [ktir-mlir-frontend#58](https://github.com/torch-spyre/ktir-mlir-frontend/pull/58), which removed the `KtdpMemorySpaceAttr` interface and made the attribute itself device-agnostic. `ktir_cpu` parses the new spelling and maps it onto the interpreter's `HBM`/`LX` via `KTDP_MEMORY_SPACE_KINDS` in `ktir_cpu/ir_types.py`, alongside the `MemRef.memory_space` validation that defines those names, so Spyre specifics stay confined to the simulator and latency model. **The RFC text itself is now stale — see row 4a.** | +| 4a | `MemorySpaceAttr` spelling conflicts with RFC | ⚠️ **spec conflict** | **The implementation intentionally diverges from RFC 0682, and the RFC needs updating.** The RFC specifies `#ktdp.spyre_memory_space` with an `unspecified` kind; [ktir-mlir-frontend#58](https://github.com/torch-spyre/ktir-mlir-frontend/pull/58) replaced this with `#ktdp.memory_space` and dropped `unspecified` with no replacement. This is a semantic change, not just a rename: the enum went from naming *devices* (HBM/LX) to naming *visibility* (reachable by all compute tiles vs. private to one). The RFC spelling no longer parses in the dialect at all, so the implementation follows the dialect. **Action: RFC 0682 §MemorySpaceAttr should be revised to match; until then treat the RFC as stale on this point.** | ## C. Affine/Polyhedral Attributes diff --git a/examples/ktir/ffn_swiglu.mlir b/examples/ktir/ffn_swiglu.mlir index e925f608..1568594b 100644 --- a/examples/ktir/ffn_swiglu.mlir +++ b/examples/ktir/ffn_swiglu.mlir @@ -26,27 +26,27 @@ module { // Create memory views %x_view = ktdp.construct_memory_view %x_ptr, sizes: [1, 64], strides: [64, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 0 >= 0, d1 >= 0, -d1 + 63 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<1x64xf16> %w_gate_view = ktdp.construct_memory_view %w_gate_ptr, sizes: [64, 128], strides: [128, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 63 >= 0, d1 >= 0, -d1 + 127 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<64x128xf16> %w_up_view = ktdp.construct_memory_view %w_up_ptr, sizes: [64, 128], strides: [128, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 63 >= 0, d1 >= 0, -d1 + 127 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<64x128xf16> %w_down_view = ktdp.construct_memory_view %w_down_ptr, sizes: [128, 64], strides: [64, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 127 >= 0, d1 >= 0, -d1 + 63 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<128x64xf16> %out_view = ktdp.construct_memory_view %out_ptr, sizes: [1, 64], strides: [64, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 0 >= 0, d1 >= 0, -d1 + 63 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<1x64xf16> // Load input x diff --git a/examples/ktir/ffn_swiglu_4core.mlir b/examples/ktir/ffn_swiglu_4core.mlir index 3a9e3717..c96feaae 100644 --- a/examples/ktir/ffn_swiglu_4core.mlir +++ b/examples/ktir/ffn_swiglu_4core.mlir @@ -123,12 +123,12 @@ module { %x_view = ktdp.construct_memory_view %x_ptr, sizes: [4, 256], strides: [256, 1] { coordinate_set = #x_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<4x256xf16> %out_view = ktdp.construct_memory_view %out_ptr, sizes: [4, 256], strides: [256, 1] { coordinate_set = #x_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<4x256xf16> // ---- W_gate [256, 1024] distributed view: 4 column shards ---- @@ -139,19 +139,19 @@ module { %gate_p0 = ktdp.construct_memory_view %w_gate_ptr, sizes: [256, 256], strides: [1024, 1] { coordinate_set = #gate_col_0, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<256x256xf16> %gate_p1 = ktdp.construct_memory_view %gate_p1_base, sizes: [256, 256], strides: [1024, 1] { coordinate_set = #gate_col_1, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<256x256xf16> %gate_p2 = ktdp.construct_memory_view %gate_p2_base, sizes: [256, 256], strides: [1024, 1] { coordinate_set = #gate_col_2, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<256x256xf16> %gate_p3 = ktdp.construct_memory_view %gate_p3_base, sizes: [256, 256], strides: [1024, 1] { coordinate_set = #gate_col_3, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<256x256xf16> %w_gate_dist = ktdp.construct_distributed_memory_view @@ -166,19 +166,19 @@ module { %up_p0 = ktdp.construct_memory_view %w_up_ptr, sizes: [256, 256], strides: [1024, 1] { coordinate_set = #gate_col_0, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<256x256xf16> %up_p1 = ktdp.construct_memory_view %up_p1_base, sizes: [256, 256], strides: [1024, 1] { coordinate_set = #gate_col_1, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<256x256xf16> %up_p2 = ktdp.construct_memory_view %up_p2_base, sizes: [256, 256], strides: [1024, 1] { coordinate_set = #gate_col_2, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<256x256xf16> %up_p3 = ktdp.construct_memory_view %up_p3_base, sizes: [256, 256], strides: [1024, 1] { coordinate_set = #gate_col_3, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<256x256xf16> %w_up_dist = ktdp.construct_distributed_memory_view @@ -193,19 +193,19 @@ module { %down_p0 = ktdp.construct_memory_view %w_down_ptr, sizes: [256, 256], strides: [256, 1] { coordinate_set = #down_row_0, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<256x256xf16> %down_p1 = ktdp.construct_memory_view %down_p1_base, sizes: [256, 256], strides: [256, 1] { coordinate_set = #down_row_1, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<256x256xf16> %down_p2 = ktdp.construct_memory_view %down_p2_base, sizes: [256, 256], strides: [256, 1] { coordinate_set = #down_row_2, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<256x256xf16> %down_p3 = ktdp.construct_memory_view %down_p3_base, sizes: [256, 256], strides: [256, 1] { coordinate_set = #down_row_3, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<256x256xf16> %w_down_dist = ktdp.construct_distributed_memory_view diff --git a/examples/ktir/nested_yield.ktir b/examples/ktir/nested_yield.ktir index 0fd1c1a1..ce20c058 100644 --- a/examples/ktir/nested_yield.ktir +++ b/examples/ktir/nested_yield.ktir @@ -5,7 +5,7 @@ module { %c2 = arith.constant 2 : index %one = arith.constant dense<1.000000e+00> : tensor<4xf32> %zero = arith.constant dense<0.000000e+00> : tensor<4xf32> - %o_desc = ktdp.construct_memory_view %out, sizes: [4], strides: [1] {memory_space = #ktdp.spyre_memory_space} : memref<4xf32> + %o_desc = ktdp.construct_memory_view %out, sizes: [4], strides: [1] {memory_space = #ktdp.memory_space} : memref<4xf32> %sum = scf.for %i = %c0 to %c2 step %c1 iter_args(%acc = %zero) -> (tensor<4xf32>) { %next = arith.addf %acc, %one : tensor<4xf32> diff --git a/examples/ktir/reduce_generic.mlir b/examples/ktir/reduce_generic.mlir index 69aa59ea..18ac96a5 100644 --- a/examples/ktir/reduce_generic.mlir +++ b/examples/ktir/reduce_generic.mlir @@ -2,7 +2,7 @@ module { func.func @reduce_explicit_region(%arg0: index) attributes {grid = [1, 1]} { %c0 = arith.constant 0 : index %cst = arith.constant 0.000000e+00 : f16 - %view = ktdp.construct_memory_view %arg0, sizes : [1, 4], strides : [4, 1] {coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + 3 >= 0)>, memory_space = #ktdp.spyre_memory_space} : memref<1x4xf16> + %view = ktdp.construct_memory_view %arg0, sizes : [1, 4], strides : [4, 1] {coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + 3 >= 0)>, memory_space = #ktdp.memory_space} : memref<1x4xf16> %acc = ktdp.construct_access_tile %view[%c0, %c0] { access_tile_set = affine_set<(d0, d1) : (d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + 3 >= 0)>, access_tile_order = affine_map<(d0, d1) -> (d0, d1)> @@ -22,7 +22,7 @@ module { // linalg.reduce generic format returns tensor<1xf16>; extract scalar before splat %scalar = tensor.extract %reduced[%c0] : tensor<1xf16> %splat = tensor.splat %scalar : tensor<1x4xf16> - %out_view = ktdp.construct_memory_view %arg0, sizes : [1, 4], strides : [4, 1] {coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + 3 >= 0)>, memory_space = #ktdp.spyre_memory_space} : memref<1x4xf16> + %out_view = ktdp.construct_memory_view %arg0, sizes : [1, 4], strides : [4, 1] {coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + 3 >= 0)>, memory_space = #ktdp.memory_space} : memref<1x4xf16> %out_acc = ktdp.construct_access_tile %out_view[%c0, %c0] { access_tile_set = affine_set<(d0, d1) : (d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + 3 >= 0)>, access_tile_order = affine_map<(d0, d1) -> (d0, d1)> diff --git a/examples/ktir/reduce_multiop.mlir b/examples/ktir/reduce_multiop.mlir index ba2578be..6640ec0f 100644 --- a/examples/ktir/reduce_multiop.mlir +++ b/examples/ktir/reduce_multiop.mlir @@ -6,7 +6,7 @@ module { func.func @reduce_multiop(%arg0: index) attributes {grid = [1, 1]} { %c0 = arith.constant 0 : index %cst = arith.constant 0xFC00 : f16 // -inf (identity for max) - %view = ktdp.construct_memory_view %arg0, sizes : [1, 8], strides : [8, 1] {coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + 7 >= 0)>, memory_space = #ktdp.spyre_memory_space} : memref<1x8xf16> + %view = ktdp.construct_memory_view %arg0, sizes : [1, 8], strides : [8, 1] {coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + 7 >= 0)>, memory_space = #ktdp.memory_space} : memref<1x8xf16> %acc = ktdp.construct_access_tile %view[%c0, %c0] { access_tile_set = affine_set<(d0, d1) : (d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + 7 >= 0)>, access_tile_order = affine_map<(d0, d1) -> (d0, d1)> @@ -25,7 +25,7 @@ module { %scalar = tensor.extract %reduced[%c0] : tensor<1xf16> %splat = tensor.splat %scalar : tensor<1x8xf16> - %out_view = ktdp.construct_memory_view %arg0, sizes : [1, 8], strides : [8, 1] {coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + 7 >= 0)>, memory_space = #ktdp.spyre_memory_space} : memref<1x8xf16> + %out_view = ktdp.construct_memory_view %arg0, sizes : [1, 8], strides : [8, 1] {coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + 7 >= 0)>, memory_space = #ktdp.memory_space} : memref<1x8xf16> %out_acc = ktdp.construct_access_tile %out_view[%c0, %c0] { access_tile_set = affine_set<(d0, d1) : (d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + 7 >= 0)>, access_tile_order = affine_map<(d0, d1) -> (d0, d1)> diff --git a/examples/ktir/ring_reduce.mlir b/examples/ktir/ring_reduce.mlir index 6baa6f35..e8309bb3 100644 --- a/examples/ktir/ring_reduce.mlir +++ b/examples/ktir/ring_reduce.mlir @@ -40,7 +40,7 @@ module { %row_view = ktdp.construct_memory_view %row_ptr, sizes: [1, 128], strides: [128, 1] { coordinate_set = #row_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<1x128xf16> // (2) Access tile and load @@ -89,7 +89,7 @@ module { %out_view = ktdp.construct_memory_view %out_ptr, sizes: [1, 128], strides: [128, 1] { coordinate_set = #row_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<1x128xf16> %out_acc = ktdp.construct_access_tile %out_view[%c0, %c0] { diff --git a/examples/ktir/ring_reduce_inner_loop.mlir b/examples/ktir/ring_reduce_inner_loop.mlir index e1f403aa..16ea8eb4 100644 --- a/examples/ktir/ring_reduce_inner_loop.mlir +++ b/examples/ktir/ring_reduce_inner_loop.mlir @@ -46,7 +46,7 @@ module { %row_view = ktdp.construct_memory_view %row_ptr, sizes: [1, 128], strides: [128, 1] { coordinate_set = #row_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<1x128xf16> // (2) Access tile and load @@ -109,7 +109,7 @@ module { %out_view = ktdp.construct_memory_view %out_ptr, sizes: [1, 128], strides: [128, 1] { coordinate_set = #row_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<1x128xf16> %out_acc = ktdp.construct_access_tile %out_view[%c0, %c0] { diff --git a/examples/ktir/scalar_broadcast.mlir b/examples/ktir/scalar_broadcast.mlir index 48ee87a2..01d3c929 100644 --- a/examples/ktir/scalar_broadcast.mlir +++ b/examples/ktir/scalar_broadcast.mlir @@ -17,7 +17,7 @@ module { // Load the 1x1 scalar source from HBM. %in_view = ktdp.construct_memory_view %in_ptr, sizes: [1, 1], strides: [1, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 >= 0, d1 >= 0, -d1 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<1x1xf16> %in_tile = ktdp.construct_access_tile %in_view[%c0, %c0] { @@ -39,7 +39,7 @@ module { // Store to HBM. %out_view = ktdp.construct_memory_view %out_ptr, sizes: [4, 64], strides: [64, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 3 >= 0, d1 >= 0, -d1 + 63 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<4x64xf16> %out_tile = ktdp.construct_access_tile %out_view[%c0, %c0] { diff --git a/examples/ktir/softmax_wide.mlir b/examples/ktir/softmax_wide.mlir index 5e095ad9..a9ce5025 100644 --- a/examples/ktir/softmax_wide.mlir +++ b/examples/ktir/softmax_wide.mlir @@ -13,12 +13,12 @@ module { %end = arith.select %cmp, %end_raw, %c_R : index // end = min(end_raw, R) %input_view_1 = ktdp.construct_memory_view %input_ptr, sizes: [2, 262144], strides: [262144, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 1 >= 0, d1 >= 0, -d1 + 262143 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<2x262144xf16> %output_view_2 = ktdp.construct_memory_view %output_ptr, sizes: [2, 262144], strides: [262144, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 1 >= 0, d1 >= 0, -d1 + 262143 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<2x262144xf16> scf.for %row = %start to %end step %c1 : index { diff --git a/examples/latency/matmul_small.mlir b/examples/latency/matmul_small.mlir index eef03ea2..25c27255 100644 --- a/examples/latency/matmul_small.mlir +++ b/examples/latency/matmul_small.mlir @@ -13,15 +13,15 @@ module { %c0_i32 = arith.constant 0 : i32 %a_view = ktdp.construct_memory_view %a_ptr, sizes: [16, 64], strides: [64, 1] { - coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 15 >= 0, d1 >= 0, -d1 + 63 >= 0)>, memory_space = #ktdp.spyre_memory_space + coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 15 >= 0, d1 >= 0, -d1 + 63 >= 0)>, memory_space = #ktdp.memory_space } : memref<16x64xf16> %b_view = ktdp.construct_memory_view %b_ptr, sizes: [64, 64], strides: [64, 1] { - coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 63 >= 0, d1 >= 0, -d1 + 63 >= 0)>, memory_space = #ktdp.spyre_memory_space + coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 63 >= 0, d1 >= 0, -d1 + 63 >= 0)>, memory_space = #ktdp.memory_space } : memref<64x64xf16> %c_view = ktdp.construct_memory_view %c_ptr, sizes: [16, 64], strides: [64, 1] { - coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 15 >= 0, d1 >= 0, -d1 + 63 >= 0)>, memory_space = #ktdp.spyre_memory_space + coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 15 >= 0, d1 >= 0, -d1 + 63 >= 0)>, memory_space = #ktdp.memory_space } : memref<16x64xf16> %offs_am = arith.muli %pid_m, %BLOCK_SIZE_M : index diff --git a/examples/latency/ring_reduce_multi_group.mlir b/examples/latency/ring_reduce_multi_group.mlir index 5f2df77f..9a2fa5ee 100644 --- a/examples/latency/ring_reduce_multi_group.mlir +++ b/examples/latency/ring_reduce_multi_group.mlir @@ -51,7 +51,7 @@ module { %row_view = ktdp.construct_memory_view %row_ptr, sizes: [1, 128], strides: [128, 1] { coordinate_set = #row_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<1x128xf16> // (2) Access tile and load @@ -112,7 +112,7 @@ module { %out_view = ktdp.construct_memory_view %group_out, sizes: [1, 128], strides: [128, 1] { coordinate_set = #row_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<1x128xf16> %out_acc = ktdp.construct_access_tile %out_view[%c0, %c0] { diff --git a/examples/latency/rope_fwd_4x2.mlir b/examples/latency/rope_fwd_4x2.mlir index 7ef65247..ddc993ad 100644 --- a/examples/latency/rope_fwd_4x2.mlir +++ b/examples/latency/rope_fwd_4x2.mlir @@ -48,22 +48,22 @@ module { // --- Construct memory views --- %x_view = ktdp.construct_memory_view %x_ptr, sizes: [163840, 128], strides: [128, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 163839 >= 0, d1 >= 0, -d1 + 127 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<163840x128xf16> %cos_view = ktdp.construct_memory_view %cos_ptr, sizes: [4096, 64], strides: [64, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 4095 >= 0, d1 >= 0, -d1 + 63 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<4096x64xf16> %sin_view = ktdp.construct_memory_view %sin_ptr, sizes: [4096, 64], strides: [64, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 4095 >= 0, d1 >= 0, -d1 + 63 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<4096x64xf16> %out_view = ktdp.construct_memory_view %out_ptr, sizes: [163840, 128], strides: [128, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 163839 >= 0, d1 >= 0, -d1 + 127 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<163840x128xf16> // --- OUTER LOOP: seq-tile loop (TILE_SEQ=256) --- diff --git a/examples/latency/softmax_small.mlir b/examples/latency/softmax_small.mlir index 4c35045a..0e581dff 100644 --- a/examples/latency/softmax_small.mlir +++ b/examples/latency/softmax_small.mlir @@ -10,11 +10,11 @@ module { %c0_i32 = arith.constant 0 : index %input_view = ktdp.construct_memory_view %input_ptr, sizes: [64, 64], strides: [64, 1] { - coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 63 >= 0, d1 >= 0, -d1 + 63 >= 0)>, memory_space = #ktdp.spyre_memory_space + coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 63 >= 0, d1 >= 0, -d1 + 63 >= 0)>, memory_space = #ktdp.memory_space } : memref<64x64xf16> %output_view = ktdp.construct_memory_view %output_ptr, sizes: [64, 64], strides: [64, 1] { - coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 63 >= 0, d1 >= 0, -d1 + 63 >= 0)>, memory_space = #ktdp.spyre_memory_space + coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 63 >= 0, d1 >= 0, -d1 + 63 >= 0)>, memory_space = #ktdp.memory_space } : memref<64x64xf16> scf.for %row = %core_id to %n_rows step %c32_i32 : index { diff --git a/examples/latency/softmax_small_explicit.mlir b/examples/latency/softmax_small_explicit.mlir index 2ab70019..d869234d 100644 --- a/examples/latency/softmax_small_explicit.mlir +++ b/examples/latency/softmax_small_explicit.mlir @@ -14,11 +14,11 @@ module { %c0_i32 = arith.constant 0 : index %input_view = ktdp.construct_memory_view %input_ptr, sizes: [64, 64], strides: [64, 1] { - coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 63 >= 0, d1 >= 0, -d1 + 63 >= 0)>, memory_space = #ktdp.spyre_memory_space + coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 63 >= 0, d1 >= 0, -d1 + 63 >= 0)>, memory_space = #ktdp.memory_space } : memref<64x64xf16> %output_view = ktdp.construct_memory_view %output_ptr, sizes: [64, 64], strides: [64, 1] { - coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 63 >= 0, d1 >= 0, -d1 + 63 >= 0)>, memory_space = #ktdp.spyre_memory_space + coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 63 >= 0, d1 >= 0, -d1 + 63 >= 0)>, memory_space = #ktdp.memory_space } : memref<64x64xf16> scf.for %row = %core_id to %n_rows step %c32_i32 : index { diff --git a/examples/rfc/add-with-control-flow.mlir b/examples/rfc/add-with-control-flow.mlir index 2be94015..2bae26c9 100644 --- a/examples/rfc/add-with-control-flow.mlir +++ b/examples/rfc/add-with-control-flow.mlir @@ -16,19 +16,19 @@ module { // Construct a memory view of A from a given address %A_view = ktdp.construct_memory_view %A_start_address, sizes: [96, 64], strides: [64, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 95 >= 0, d1 >= 0, -d1 + 63 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<96x64xf16> // Construct a memory view of B from a given address %B_view = ktdp.construct_memory_view %B_start_address, sizes: [96, 64], strides: [64, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 95 >= 0, d1 >= 0, -d1 + 63 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<96x64xf16> // Construct a memory view of C from a given address %C_view = ktdp.construct_memory_view %C_start_address, sizes: [96, 64], strides: [64, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 95 >= 0, d1 >= 0, -d1 + 63 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<96x64xf16> // Looping over tile size with each iteration working over 1x64 fp16 diff --git a/examples/rfc/distributed-view-copy-dynamic.mlir b/examples/rfc/distributed-view-copy-dynamic.mlir index 193261d6..0add01a9 100644 --- a/examples/rfc/distributed-view-copy-dynamic.mlir +++ b/examples/rfc/distributed-view-copy-dynamic.mlir @@ -47,12 +47,12 @@ module { // (1) Per-partition memory views with symbolic trailing dim. %A0_view = ktdp.construct_memory_view %a0_ptr, sizes: [64, %s0], strides: [%s0, 1] { coordinate_set = #A0_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<64x?xf16> %A1_view = ktdp.construct_memory_view %a1_ptr, sizes: [64, %s0], strides: [%s0, 1] { coordinate_set = #A1_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<64x?xf16> // (1) Compose into a single logical distributed view. Inputs are @@ -64,7 +64,7 @@ module { // (1) Output view B, dynamic like the partitions (trailing dim = 2*s0). %B_view = ktdp.construct_memory_view %b_ptr, sizes: [64, %ub], strides: [%ub, 1] { coordinate_set = #B_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<64x?xf16> // (2)+(3) Iterate access tiles across the global domain. Trip count diff --git a/examples/rfc/distributed-view-copy-rowmerge-dynamic.mlir b/examples/rfc/distributed-view-copy-rowmerge-dynamic.mlir index 817f2705..59a0978f 100644 --- a/examples/rfc/distributed-view-copy-rowmerge-dynamic.mlir +++ b/examples/rfc/distributed-view-copy-rowmerge-dynamic.mlir @@ -49,12 +49,12 @@ module { // (1) Per-partition memory views with symbolic trailing dim (shared s0). %A0_view = ktdp.construct_memory_view %a0_ptr, sizes: [64, %s0], strides: [%s0, 1] { coordinate_set = #A0_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<64x?xf16> %A1_view = ktdp.construct_memory_view %a1_ptr, sizes: [64, %s0], strides: [%s0, 1] { coordinate_set = #A1_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<64x?xf16> // (1) Compose along the concrete row axis: 64 + 64 = 128 rows, s0 cols. @@ -65,7 +65,7 @@ module { // (1) Output view B, 128 x s0. %B_view = ktdp.construct_memory_view %b_ptr, sizes: [128, %s0], strides: [%s0, 1] { coordinate_set = #B_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<128x?xf16> // (2)+(3) Iterate: cols over the symbolic extent, rows over the two diff --git a/examples/rfc/distributed-view-copy.mlir b/examples/rfc/distributed-view-copy.mlir index 91ed496c..14319bbb 100644 --- a/examples/rfc/distributed-view-copy.mlir +++ b/examples/rfc/distributed-view-copy.mlir @@ -3,17 +3,18 @@ // The program copies a distributed tensor A (partitioned across HBM and two LX scratchpads) // into a contiguous output tensor B on HBM. // NOTE: Copied from KTIR-V1.ppt slide 9 - construct distributed memory view. -// NOTE: Minor correction #ktpd.spyre_memory_space -> LX. -// NOTE: The two LX partitions are tagged with distinct `core = N` indices -// per the KTDP ODS `#ktdp.spyre_memory_space` form -// (KernelTileIR/src/Dialects/KTDP/KTDPOps.td:68-104). In real hardware -// each compute core has its own private LX SRAM, so A_LX0's and -// A_LX1's byte-address spaces are independent and both bases could -// naturally be 0 (a common layout in real kernels). -// NOTE: Simulator caveat — ktir-cpu does not yet honor the `core = N` -// index: every `` memref resolves to the currently -// executing core's single scratchpad. This test is marked xfail -// until per-core LX routing is implemented. Once it lands, all +// NOTE: Minor correction — the slide's `LX0` memory space is spelled +// `#ktdp.memory_space` here. +// NOTE: The two tile-local partitions are tagged with distinct `ct_id = N` +// indices per the KTDP ODS `#ktdp.memory_space` +// form. In real hardware each compute tile has its own private LX +// SRAM, so A_LX0's and A_LX1's byte-address spaces are independent +// and both bases could naturally be 0 (a common layout in real +// kernels). +// NOTE: Simulator caveat — ktir-cpu does not yet honor the `ct_id = N` +// index: every `` memref resolves to the +// currently executing core's single scratchpad. This test is marked +// xfail until per-core LX routing is implemented. Once it lands, all // three partition bases can collapse to 0. // A is a 192x64 logical tensor partitioned across three memory spaces: @@ -56,18 +57,18 @@ module { // Note: number of entries in sizes, strides, dims in coordinate_set, shape of memref must be identical %A_HBM_view = ktdp.construct_memory_view %A_HBM_addr, sizes: [96, 64], strides: [64, 1] { coordinate_set = #A_HBM_coord_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<96x64xf16> // Note: column-major layout expressed via strides [1, 64] %A_LX0_view = ktdp.construct_memory_view %A_LX0_addr, sizes: [32, 64], strides: [1, 64] { coordinate_set = #A_LX0_coord_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<32x64xf16> %A_LX1_view = ktdp.construct_memory_view %A_LX1_addr, sizes: [64, 64], strides: [64, 1] { coordinate_set = #A_LX1_coord_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<64x64xf16> // (1) Compose the three partition views into a single logical distributed view of shape 192x64 @@ -80,7 +81,7 @@ module { // (1) Construct memory view for output B %B_view = ktdp.construct_memory_view %B_addr, sizes: [192, 64], strides: [64, 1] { coordinate_set = #B_coord_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<192x64xf16> // (2) Construct direct access tile for A over the full 192x64 global coordinate space diff --git a/examples/rfc/indirect-access-copy.mlir b/examples/rfc/indirect-access-copy.mlir index 138ea5b5..354c15c9 100644 --- a/examples/rfc/indirect-access-copy.mlir +++ b/examples/rfc/indirect-access-copy.mlir @@ -37,25 +37,25 @@ module { // Note: number of entries in sizes, strides, dims in coordinate_set, shape of memref must be identical %X = ktdp.construct_memory_view %X_addr, sizes: [64, 64], strides: [64, 1] { coordinate_set = #X_coord_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<64x64xf16> // (1) Construct memory view for IDX1 (row indices into X) %IDX1 = ktdp.construct_memory_view %IDX1_addr, sizes: [64, 64], strides: [64, 1] { coordinate_set = #IDX_coord_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<64x64xi32> // (1) Construct memory view for IDX2 (column indices into X) %IDX2 = ktdp.construct_memory_view %IDX2_addr, sizes: [64, 64], strides: [64, 1] { coordinate_set = #IDX_coord_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<64x64xi32> // (1) Construct memory view for output Y %Y_view = ktdp.construct_memory_view %Y_addr, sizes: [64, 64], strides: [64, 1] { coordinate_set = #Y_coord_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<64x64xf16> // (2) Construct indirect access tile for X[IDX1[m, k], IDX2[m, k]] diff --git a/examples/rfc/indirect-scatter.mlir b/examples/rfc/indirect-scatter.mlir index a4163a48..094c8edc 100644 --- a/examples/rfc/indirect-scatter.mlir +++ b/examples/rfc/indirect-scatter.mlir @@ -25,25 +25,25 @@ module { // (1) Construct memory view for X (source) %X_view = ktdp.construct_memory_view %X_addr, sizes: [64, 64], strides: [64, 1] { coordinate_set = #X_coord_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<64x64xf16> // (1) Construct memory view for IDX1 (row indices into Y) %IDX1 = ktdp.construct_memory_view %IDX1_addr, sizes: [64, 64], strides: [64, 1] { coordinate_set = #IDX_coord_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<64x64xi32> // (1) Construct memory view for IDX2 (column indices into Y) %IDX2 = ktdp.construct_memory_view %IDX2_addr, sizes: [64, 64], strides: [64, 1] { coordinate_set = #IDX_coord_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<64x64xi32> // (1) Construct memory view for Y (destination) %Y_view = ktdp.construct_memory_view %Y_addr, sizes: [64, 64], strides: [64, 1] { coordinate_set = #Y_coord_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<64x64xf16> // (2) Direct access tile for X[m, k] — load source contiguously diff --git a/examples/rfc/paged-tensor-copy.mlir b/examples/rfc/paged-tensor-copy.mlir index efbd12b3..53f9c50c 100644 --- a/examples/rfc/paged-tensor-copy.mlir +++ b/examples/rfc/paged-tensor-copy.mlir @@ -80,7 +80,7 @@ module { %Idx_mem_view = ktdp.construct_memory_view %Idx_start_address, sizes: [%Nb, %Ntkv_Ptkv], strides: [%Ntkv_Ptkv, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 3 >= 0, d1 >= 0, -d1 + 31 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<4x32xi32> // (1) Construct memory view for input X @@ -89,7 +89,7 @@ module { %X_mem_view = ktdp.construct_memory_view %X_start_address, sizes: [%Npages, %Nhkv, %Ptkv, %Ndkv], strides: [65536, 8192, %Ndkv, 1] { coordinate_set = #X_coord_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<10000x8x64x128xf16> // (2) Construct memory view for output Y @@ -100,7 +100,7 @@ module { %Y_mem_view = ktdp.construct_memory_view %Y_start_address, sizes: [%Nb, %Ntkv, %Nhkv, %Ndkv], strides: [%Y_str_Nb, %Y_str_Ntkv, %Ndkv , 1] { coordinate_set = #Y_coord_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<4x2048x8x128xf16> // (3) Construct indirect access tile X [Idx[b][tkv/Ptkv]] [h] [tkv%Ptkv] [dkv] diff --git a/examples/rfc/paged-tensor-write.mlir b/examples/rfc/paged-tensor-write.mlir index 6855a3d4..66e1cb4d 100644 --- a/examples/rfc/paged-tensor-write.mlir +++ b/examples/rfc/paged-tensor-write.mlir @@ -56,7 +56,7 @@ module { %Idx_mem_view = ktdp.construct_memory_view %Idx_start_address, sizes: [%Nb, %Ntkv_Ptkv], strides: [%Ntkv_Ptkv, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 3 >= 0, d1 >= 0, -d1 + 31 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<4x32xi32> // (1) Construct memory view for input X (contiguous) @@ -66,7 +66,7 @@ module { %X_mem_view = ktdp.construct_memory_view %X_start_address, sizes: [%Nb, %Ntkv, %Nhkv, %Ndkv], strides: [%X_str_Nb, %X_str_Ntkv, %Ndkv, 1] { coordinate_set = #contiguous_coord_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<4x2048x8x128xf16> // (1) Construct memory view for output Y (paged) @@ -74,7 +74,7 @@ module { %Y_mem_view = ktdp.construct_memory_view %Y_start_address, sizes: [%Npages, %Nhkv, %Ptkv, %Ndkv], strides: [65536, 8192, %Ndkv, 1] { coordinate_set = #paged_coord_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<10000x8x64x128xf16> // (2) Direct access tile for X[b, tkv, h, dkv] — load source contiguously diff --git a/examples/sdsc/sdpa_pv_ksplit.mlir b/examples/sdsc/sdpa_pv_ksplit.mlir index 580b0ccd..cb8907f2 100644 --- a/examples/sdsc/sdpa_pv_ksplit.mlir +++ b/examples/sdsc/sdpa_pv_ksplit.mlir @@ -65,13 +65,13 @@ module { %pid_out_1, %pid_in_2 = ktdp.get_compute_tile_id : index, index %c0_3 = arith.constant 0 : index %a_view_4 = ktdp.construct_memory_view %a_ptr, sizes: [1, 8192], strides: [8192, 1] { - coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 0 >= 0, d1 >= 0, -d1 + 8191 >= 0)>, memory_space = #ktdp.spyre_memory_space + coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 0 >= 0, d1 >= 0, -d1 + 8191 >= 0)>, memory_space = #ktdp.memory_space } : memref<1x8192xf16> %b_view_5 = ktdp.construct_memory_view %b_ptr, sizes: [8192, 128], strides: [128, 1] { - coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 8191 >= 0, d1 >= 0, -d1 + 127 >= 0)>, memory_space = #ktdp.spyre_memory_space + coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 8191 >= 0, d1 >= 0, -d1 + 127 >= 0)>, memory_space = #ktdp.memory_space } : memref<8192x128xf16> %c_view_6 = ktdp.construct_memory_view %c_ptr, sizes: [1, 128], strides: [128, 1] { - coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 0 >= 0, d1 >= 0, -d1 + 127 >= 0)>, memory_space = #ktdp.spyre_memory_space + coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 0 >= 0, d1 >= 0, -d1 + 127 >= 0)>, memory_space = #ktdp.memory_space } : memref<1x128xf16> %blk_in_7 = arith.constant 512 : index %off_in_8 = arith.muli %pid_in_2, %blk_in_7 : index diff --git a/examples/triton-ktir/indexed_add.mlir b/examples/triton-ktir/indexed_add.mlir index 45d9c057..70fb57b7 100644 --- a/examples/triton-ktir/indexed_add.mlir +++ b/examples/triton-ktir/indexed_add.mlir @@ -12,14 +12,14 @@ module { // Memory view for index tensor: shape [2] (num_indices), stride [1] %index_view = ktdp.construct_memory_view %index_ptr, sizes: [2], strides: [1] { coordinate_set = affine_set<(d0) : (d0 >= 0, -d0 + 1 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<2xi64> // Memory view for x: shape [128, 64, 8, 128] // strides: [64*8*128, 8*128, 128, 1] = [65536, 1024, 128, 1] %x_view = ktdp.construct_memory_view %x_ptr, sizes: [128, 64, 8, 128], strides: [65536, 1024, 128, 1] { coordinate_set = affine_set<(d0, d1, d2, d3) : (d0 >= 0, -d0 + 127 >= 0, d1 >= 0, -d1 + 63 >= 0, d2 >= 0, -d2 + 7 >= 0, d3 >= 0, -d3 + 127 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<128x64x8x128xf16> %x_tile = ktdp.construct_indirect_access_tile @@ -33,7 +33,7 @@ module { // strides: [32*8*128, 8*128, 128, 1] = [32768, 1024, 128, 1] %y_view = ktdp.construct_memory_view %y_ptr, sizes: [2, 32, 8, 128], strides: [32768, 1024, 128, 1] { coordinate_set = affine_set<(d0, d1, d2, d3) : (d0 >= 0, -d0 + 1 >= 0, d1 >= 0, -d1 + 31 >= 0, d2 >= 0, -d2 + 7 >= 0, d3 >= 0, -d3 + 127 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<2x32x8x128xf16> // Direct access tile for y: y[grid0, dim1_start, grid1, 0] covering [1, 32, 1, 128] @@ -50,7 +50,7 @@ module { // Memory view for output (same layout as y) %output_view = ktdp.construct_memory_view %output_ptr, sizes: [2, 32, 8, 128], strides: [32768, 1024, 128, 1] { coordinate_set = affine_set<(d0, d1, d2, d3) : (d0 >= 0, -d0 + 1 >= 0, d1 >= 0, -d1 + 31 >= 0, d2 >= 0, -d2 + 7 >= 0, d3 >= 0, -d3 + 127 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<2x32x8x128xf16> %output_tile = ktdp.construct_access_tile %output_view[%grid0, %dim1_start, %grid1, %c0] { diff --git a/examples/triton-ktir/layernorm_fwd_ktir.mlir b/examples/triton-ktir/layernorm_fwd_ktir.mlir index 296defd6..f0b27de9 100644 --- a/examples/triton-ktir/layernorm_fwd_ktir.mlir +++ b/examples/triton-ktir/layernorm_fwd_ktir.mlir @@ -18,27 +18,27 @@ module { %f1_f16 = arith.constant 1.0 : f16 %X_view = ktdp.construct_memory_view %X, sizes: [1151, 8192], strides: [8192, 1] { - coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 1150 >= 0, d1 >= 0, -d1 + 8191 >= 0)>, memory_space = #ktdp.spyre_memory_space + coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 1150 >= 0, d1 >= 0, -d1 + 8191 >= 0)>, memory_space = #ktdp.memory_space } : memref<1151x8192xf16> %Y_view = ktdp.construct_memory_view %Y, sizes: [1151, 8192], strides: [8192, 1] { - coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 1150 >= 0, d1 >= 0, -d1 + 8191 >= 0)>, memory_space = #ktdp.spyre_memory_space + coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 1150 >= 0, d1 >= 0, -d1 + 8191 >= 0)>, memory_space = #ktdp.memory_space } : memref<1151x8192xf16> %W_view = ktdp.construct_memory_view %W, sizes: [1151, 8192], strides: [8192, 1] { - coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 1150 >= 0, d1 >= 0, -d1 + 8191 >= 0)>, memory_space = #ktdp.spyre_memory_space + coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 1150 >= 0, d1 >= 0, -d1 + 8191 >= 0)>, memory_space = #ktdp.memory_space } : memref<1151x8192xf16> %B_view = ktdp.construct_memory_view %B, sizes: [1151, 8192], strides: [8192, 1] { - coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 1150 >= 0, d1 >= 0, -d1 + 8191 >= 0)>, memory_space = #ktdp.spyre_memory_space + coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 1150 >= 0, d1 >= 0, -d1 + 8191 >= 0)>, memory_space = #ktdp.memory_space } : memref<1151x8192xf16> %Mean_view = ktdp.construct_memory_view %Mean, sizes: [1151], strides: [1] { - coordinate_set = affine_set<(d0) : (d0 >= 0, -d0 + 1150 >= 0)>, memory_space = #ktdp.spyre_memory_space + coordinate_set = affine_set<(d0) : (d0 >= 0, -d0 + 1150 >= 0)>, memory_space = #ktdp.memory_space } : memref<1151xf16> %Rstd_view = ktdp.construct_memory_view %Rstd, sizes: [1151], strides: [1] { - coordinate_set = affine_set<(d0) : (d0 >= 0, -d0 + 1150 >= 0)>, memory_space = #ktdp.spyre_memory_space + coordinate_set = affine_set<(d0) : (d0 >= 0, -d0 + 1150 >= 0)>, memory_space = #ktdp.memory_space } : memref<1151xf16> scf.for %row = %core_id to %c1151_i32 step %c32_i32 : index { diff --git a/examples/triton-ktir/matmul_fwd_ktir.mlir b/examples/triton-ktir/matmul_fwd_ktir.mlir index 7ebe6100..32b7252c 100644 --- a/examples/triton-ktir/matmul_fwd_ktir.mlir +++ b/examples/triton-ktir/matmul_fwd_ktir.mlir @@ -14,15 +14,15 @@ module { %c0 = arith.constant 0 : index %a_view = ktdp.construct_memory_view %a_ptr, sizes: [64, 2048], strides: [2048, 1] { - coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 63 >= 0, d1 >= 0, -d1 + 2047 >= 0)>, memory_space = #ktdp.spyre_memory_space + coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 63 >= 0, d1 >= 0, -d1 + 2047 >= 0)>, memory_space = #ktdp.memory_space } : memref<64x2048xf16> %b_view = ktdp.construct_memory_view %b_ptr, sizes: [2048, 8192], strides: [8192, 1] { - coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 2047 >= 0, d1 >= 0, -d1 + 8191 >= 0)>, memory_space = #ktdp.spyre_memory_space + coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 2047 >= 0, d1 >= 0, -d1 + 8191 >= 0)>, memory_space = #ktdp.memory_space } : memref<2048x8192xf16> %c_view = ktdp.construct_memory_view %c_ptr, sizes: [64, 8192], strides: [8192, 1] { - coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 63 >= 0, d1 >= 0, -d1 + 8191 >= 0)>, memory_space = #ktdp.spyre_memory_space + coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 63 >= 0, d1 >= 0, -d1 + 8191 >= 0)>, memory_space = #ktdp.memory_space } : memref<64x8192xf16> %offs_am = arith.muli %pid_m, %BLOCK_SIZE_M : index diff --git a/examples/triton-ktir/paged_attention.mlir b/examples/triton-ktir/paged_attention.mlir index 13451664..e3ec6301 100644 --- a/examples/triton-ktir/paged_attention.mlir +++ b/examples/triton-ktir/paged_attention.mlir @@ -100,7 +100,7 @@ module { d0 >= 0, -d0 + 7 >= 0, d1 >= 0, -d1 + 31 >= 0, d2 >= 0, -d2 + 127 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<8x32x128xf16> // key_cache: [64, 16, 8, 128] strides [16384, 1024, 128, 1] @@ -112,7 +112,7 @@ module { d1 >= 0, -d1 + 15 >= 0, d2 >= 0, -d2 + 7 >= 0, d3 >= 0, -d3 + 127 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<64x16x8x128xf16> // value_cache: [64, 16, 8, 128] strides [16384, 1024, 128, 1] @@ -123,7 +123,7 @@ module { d1 >= 0, -d1 + 15 >= 0, d2 >= 0, -d2 + 7 >= 0, d3 >= 0, -d3 + 127 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<64x16x8x128xf16> // block_tables: [1, 16] strides [16, 1] @@ -132,7 +132,7 @@ module { coordinate_set = affine_set<(d0, d1) : ( d0 >= 0, -d0 + 0 >= 0, d1 >= 0, -d1 + 15 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<1x16xi32> // output: [8, 32, 128] strides [4096, 128, 1] @@ -142,7 +142,7 @@ module { d0 >= 0, -d0 + 7 >= 0, d1 >= 0, -d1 + 31 >= 0, d2 >= 0, -d2 + 127 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<8x32x128xf16> // ----------------------------------------------------------------------- diff --git a/examples/triton-ktir/sdpa_2d.mlir b/examples/triton-ktir/sdpa_2d.mlir index 679e32af..5bd269b8 100644 --- a/examples/triton-ktir/sdpa_2d.mlir +++ b/examples/triton-ktir/sdpa_2d.mlir @@ -51,7 +51,7 @@ module { // ----------------------------------------------------------------------- %q_view = ktdp.construct_memory_view %q_ptr, sizes: [32, 64], strides: [64, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 31 >= 0, d1 >= 0, -d1 + 63 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<32x64xf16> %q_tile = ktdp.construct_access_tile %q_view[%pid_m, %c0] { @@ -71,7 +71,7 @@ module { // ----------------------------------------------------------------------- %k_view = ktdp.construct_memory_view %k_ptr, sizes: [32, 64], strides: [64, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 31 >= 0, d1 >= 0, -d1 + 63 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<32x64xf16> %k_tile = ktdp.construct_access_tile %k_view[%c0, %c0] { @@ -154,7 +154,7 @@ module { // ----------------------------------------------------------------------- %v_view = ktdp.construct_memory_view %v_ptr, sizes: [32, 64], strides: [64, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 31 >= 0, d1 >= 0, -d1 + 63 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<32x64xf16> %v_tile = ktdp.construct_access_tile %v_view[%c0, %c0] { @@ -180,7 +180,7 @@ module { // ----------------------------------------------------------------------- %output_view = ktdp.construct_memory_view %output_ptr, sizes: [32, 64], strides: [64, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 31 >= 0, d1 >= 0, -d1 + 63 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<32x64xf16> %output_tile = ktdp.construct_access_tile %output_view[%pid_m, %c0] { diff --git a/examples/triton-ktir/softmax_fwd_ktir.mlir b/examples/triton-ktir/softmax_fwd_ktir.mlir index 7adbc03b..1dc33e65 100644 --- a/examples/triton-ktir/softmax_fwd_ktir.mlir +++ b/examples/triton-ktir/softmax_fwd_ktir.mlir @@ -10,11 +10,11 @@ module { %c0_i32 = arith.constant 0 : index %input_view = ktdp.construct_memory_view %input_ptr, sizes: [4096, 1024], strides: [1024, 1] { - coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 4095 >= 0, d1 >= 0, -d1 + 1023 >= 0)>, memory_space = #ktdp.spyre_memory_space + coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 4095 >= 0, d1 >= 0, -d1 + 1023 >= 0)>, memory_space = #ktdp.memory_space } : memref<4096x1024xf16> %output_view = ktdp.construct_memory_view %output_ptr, sizes: [4096, 1024], strides: [1024, 1] { - coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 4095 >= 0, d1 >= 0, -d1 + 1023 >= 0)>, memory_space = #ktdp.spyre_memory_space + coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 4095 >= 0, d1 >= 0, -d1 + 1023 >= 0)>, memory_space = #ktdp.memory_space } : memref<4096x1024xf16> scf.for %row = %core_id to %n_rows step %c32_i32 : index { diff --git a/examples/triton-ktir/vector_add_dynamic_ktir.mlir b/examples/triton-ktir/vector_add_dynamic_ktir.mlir index d497b267..bbd44baf 100644 --- a/examples/triton-ktir/vector_add_dynamic_ktir.mlir +++ b/examples/triton-ktir/vector_add_dynamic_ktir.mlir @@ -16,7 +16,7 @@ module { %n = arith.index_cast %n_elements : i32 to index %x_view = ktdp.construct_memory_view %x_ptr, sizes: [%n], strides: [1] { - coordinate_set = #set, memory_space = #ktdp.spyre_memory_space + coordinate_set = #set, memory_space = #ktdp.memory_space } : memref %x_tile = ktdp.construct_access_tile %x_view[%c0] { access_tile_order = #map, access_tile_set = #set1 @@ -24,7 +24,7 @@ module { %x = ktdp.load %x_tile : <1024xindex> -> tensor<1024xf32> %y_view = ktdp.construct_memory_view %y_ptr, sizes: [%n], strides: [1] { - coordinate_set = #set, memory_space = #ktdp.spyre_memory_space + coordinate_set = #set, memory_space = #ktdp.memory_space } : memref %y_tile = ktdp.construct_access_tile %y_view[%c0] { access_tile_order = #map, access_tile_set = #set1 @@ -34,7 +34,7 @@ module { %output = arith.addf %x, %y : tensor<1024xf32> %output_view = ktdp.construct_memory_view %output_ptr, sizes: [%n], strides: [1] { - coordinate_set = #set, memory_space = #ktdp.spyre_memory_space + coordinate_set = #set, memory_space = #ktdp.memory_space } : memref %output_tile = ktdp.construct_access_tile %output_view[%c0] { access_tile_order = #map, access_tile_set = #set1 diff --git a/examples/triton-ktir/vector_add_ktir.mlir b/examples/triton-ktir/vector_add_ktir.mlir index 14767bc3..41feaba0 100644 --- a/examples/triton-ktir/vector_add_ktir.mlir +++ b/examples/triton-ktir/vector_add_ktir.mlir @@ -9,7 +9,7 @@ module { %offset = arith.muli %core_id, %BLOCK_SIZE : index %x_view = ktdp.construct_memory_view %x_ptr, sizes: [4096], strides: [1] { - coordinate_set = affine_set<(d0) : (d0 >= 0, -d0 + 4095 >= 0)>, memory_space = #ktdp.spyre_memory_space + coordinate_set = affine_set<(d0) : (d0 >= 0, -d0 + 4095 >= 0)>, memory_space = #ktdp.memory_space } : memref<4096xf16> %x_tile = ktdp.construct_access_tile %x_view[%offset] { @@ -18,7 +18,7 @@ module { } : memref<4096xf16> -> !ktdp.access_tile<128xindex> %y_view = ktdp.construct_memory_view %y_ptr, sizes: [4096], strides: [1] { - coordinate_set = affine_set<(d0) : (d0 >= 0, -d0 + 4095 >= 0)>, memory_space = #ktdp.spyre_memory_space + coordinate_set = affine_set<(d0) : (d0 >= 0, -d0 + 4095 >= 0)>, memory_space = #ktdp.memory_space } : memref<4096xf16> %y_tile = ktdp.construct_access_tile %y_view[%offset] { @@ -32,7 +32,7 @@ module { %output = arith.addf %x, %y : tensor<128xf16> %output_view = ktdp.construct_memory_view %output_ptr, sizes: [4096], strides: [1] { - coordinate_set = affine_set<(d0) : (d0 >= 0, -d0 + 4095 >= 0)>, memory_space = #ktdp.spyre_memory_space + coordinate_set = affine_set<(d0) : (d0 >= 0, -d0 + 4095 >= 0)>, memory_space = #ktdp.memory_space } : memref<4096xf16> %output_tile = ktdp.construct_access_tile %output_view[%offset] { diff --git a/ktir_cpu/dialects/ktdp_ops.py b/ktir_cpu/dialects/ktdp_ops.py index cbbe7bf2..0d2fdf23 100644 --- a/ktir_cpu/dialects/ktdp_ops.py +++ b/ktir_cpu/dialects/ktdp_ops.py @@ -23,6 +23,7 @@ DistributedMemRef, DistributedTileRef, IndirectAccessTile, + KTDP_MEMORY_SPACE_KINDS, MemRef, Operation, Tile, @@ -338,16 +339,22 @@ def parse_construct_memory_view(op_text, parse_ctx: ParseContext): memory_space = "HBM" lx_core_id = None - # Accept both ``/`` and the RFC's per-core LX form - # ``. On real hardware each compute core has its own - # private LX SRAM, so a partition tagged `core = N` lives in core N's + # `#ktdp.memory_space`, mapped to the + # interpreter's HBM/LX names. On real hardware each compute tile has its + # own private LX SRAM, so a view tagged `ct_id = N` lives in tile N's # scratchpad — captured into lx_core_id and used at load/store time. mem_match = re.search( - r'#ktdp\.spyre_memory_space<\s*(\w+)(?:\s*,\s*core\s*=\s*(\d+))?\s*>', + r'#ktdp\.memory_space<\s*(\w+)(?:\s*,\s*ct_id\s*=\s*(\d+))?\s*>', op_text, ) if mem_match: - memory_space = mem_match.group(1) + kind = mem_match.group(1) + if kind not in KTDP_MEMORY_SPACE_KINDS: + raise ValueError( + f"Unknown #ktdp.memory_space kind {kind!r}; expected one of " + f"{sorted(KTDP_MEMORY_SPACE_KINDS)}" + ) + memory_space = KTDP_MEMORY_SPACE_KINDS[kind] if mem_match.group(2) is not None: lx_core_id = int(mem_match.group(2)) diff --git a/ktir_cpu/ir_types.py b/ktir_cpu/ir_types.py index 274ed303..8cc1785a 100644 --- a/ktir_cpu/ir_types.py +++ b/ktir_cpu/ir_types.py @@ -42,6 +42,13 @@ CoordinateSet = Union[BoxSet, AffineSet, List[Tuple[int, ...]]] +# KTDP names memory spaces by visibility (`global` reaches every compute tile, +# `ct_local` is private to one); the interpreter names the concrete Spyre +# memories those map onto, since it models their addressing and latency. +# Parsers translate on the way in — see MemRef.memory_space. +KTDP_MEMORY_SPACE_KINDS = {"global": "HBM", "ct_local": "LX"} + + @dataclass class MemRef: """Hardware-aware memory view (result of construct_memory_view). @@ -64,12 +71,12 @@ class MemRef: # corner) — relied on by ``distributed_tile_access`` for ``p_i``. coordinate_set: Optional[AffineSet] = None # Set when memory_space="LX" and a core index was specified via - # #ktdp.spyre_memory_space. None means "the executing + # #ktdp.memory_space. None means "the executing # core's own LX scratchpad" (default routing). lx_core_id: Optional[int] = None def __post_init__(self): - valid = ("HBM", "LX") + valid = tuple(KTDP_MEMORY_SPACE_KINDS.values()) if self.memory_space not in valid: raise ValueError( f"Invalid memory_space {self.memory_space!r}. Must be one of {valid}." diff --git a/ktir_cpu/mlir_frontend/parser.py b/ktir_cpu/mlir_frontend/parser.py index ff6261eb..9c4a4bb4 100644 --- a/ktir_cpu/mlir_frontend/parser.py +++ b/ktir_cpu/mlir_frontend/parser.py @@ -42,13 +42,13 @@ ShapedType, ) from mlir_ktdp.passmanager import PassManager - from tools_ktdp.ir_utils import ktdp_context, walk_module + from mlir_ktdp.tools import ktdp_context, walk_module _HAS_MLIR = True except ImportError: _HAS_MLIR = False from ..affine import AffineMap, AffineSet -from ..ir_types import IRFunction, IRModule, Operation +from ..ir_types import KTDP_MEMORY_SPACE_KINDS, IRFunction, IRModule, Operation from ..parser_ast import parse_affine_map, parse_affine_set from ..parser import KTIRParserBase @@ -324,9 +324,10 @@ def _splice(static_vals, names): if not m: raise ValueError(f"ktdp.construct_memory_view: cannot parse dtype from {result_type!r}") attributes["dtype"] = m.group(1) - # str(memory_space attr) -> "#ktdp.spyre_memory_space" or "" + # str(memory_space attr) -> "#ktdp.memory_space" or + # "", mapped to the interpreter's HBM/LX names. ms = re.search( - r'#ktdp\.spyre_memory_space<\s*(\w+)(?:\s*,\s*core\s*=\s*(\d+))?\s*>', + r'#ktdp\.memory_space<\s*(\w+)(?:\s*,\s*ct_id\s*=\s*(\d+))?\s*>', str(mlir_op.attributes["memory_space"]), ) if not ms: @@ -334,7 +335,7 @@ def _splice(static_vals, names): f"ktdp.construct_memory_view: cannot parse memory_space from " f"{mlir_op.attributes['memory_space']!r}" ) - attributes["memory_space"] = ms.group(1) + attributes["memory_space"] = KTDP_MEMORY_SPACE_KINDS[ms.group(1)] if ms.group(2) is not None: attributes["lx_core_id"] = int(ms.group(2)) # str(coordinate_set attr) → "affine_set<(d0) : ...>" @@ -819,7 +820,7 @@ class MLIRFrontendParser(KTIRParserBase): def __init__(self, adapter: Optional[MLIRTypeAdapter] = None): if not _HAS_MLIR: raise ImportError( - "mlir_ktdp / tools_ktdp not installed; " + "mlir_ktdp not installed; " "MLIRFrontendParser is unavailable." ) self._adapter = adapter or MLIRTypeAdapter() diff --git a/ktir_cpu/ops/memory_ops.py b/ktir_cpu/ops/memory_ops.py index d561aefe..dd0fa29b 100644 --- a/ktir_cpu/ops/memory_ops.py +++ b/ktir_cpu/ops/memory_ops.py @@ -340,7 +340,7 @@ def tile_view( Builds a MemRef describing a contiguous region in HBM or LX. ``lx_core_id``, when set, identifies which core's LX scratchpad - the data lives in (parsed from #ktdp.spyre_memory_space); + the data lives in (parsed from #ktdp.memory_space); load/store use it to route via context.get_lx(). """ return MemRef( diff --git a/ktir_cpu/parser_utils.py b/ktir_cpu/parser_utils.py index ca1b7705..fbcb2f87 100644 --- a/ktir_cpu/parser_utils.py +++ b/ktir_cpu/parser_utils.py @@ -276,7 +276,7 @@ def parse_attr_block(op_text: str, aliases: Optional[Dict] = None, Values are returned as Python scalars (int/float/list/str). Handles: - ``keyword<...>`` values (e.g. ``affine_map<...>``, ``affine_set<...>``, - ``#ktdp.spyre_memory_space``): ``<``/``>`` depth is counted while + ``#ktdp.memory_space``): ``<``/``>`` depth is counted while skipping ``>=`` and ``->`` operators so constraint expressions like ``d0 >= 0`` do not prematurely close the value. - ``#alias`` references: resolved via *aliases* when provided. diff --git a/notebooks/demo_gen_mlir.py b/notebooks/demo_gen_mlir.py index f587124b..3136ae71 100644 --- a/notebooks/demo_gen_mlir.py +++ b/notebooks/demo_gen_mlir.py @@ -27,7 +27,7 @@ def _mem_view(name: str, ptr: str, shape: list[int], strides: list[int], f" %{name} = ktdp.construct_memory_view %{ptr}, " f"sizes: [{sizes_str}], strides: [{strides_str}] {{\n" f" coordinate_set = affine_set<({dim_decl}) : ({bounds})>,\n" - f" memory_space = #ktdp.spyre_memory_space\n" + f" memory_space = #ktdp.memory_space\n" f" }} : memref<{shape_str}x{dtype}>" ) diff --git a/pyproject.toml b/pyproject.toml index 527c721f..d4b6e355 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,7 +30,7 @@ dev = [ ] # No PyPI release yet; built from source via setup_mlir.py (see CONTRIBUTING.md). mlir-frontend = [ - "ktir-mlir-frontend @ git+https://github.com/torch-spyre/ktir-mlir-frontend@973d67ee7d535ea41f4adbf9c6eb4ee59450672b", + "ktir-mlir-frontend @ git+https://github.com/torch-spyre/ktir-mlir-frontend@ecfb9ed73b6d40a19df8f678c56ed3074bbe2f85", ] [tool.setuptools_scm] diff --git a/tests/test_boxset_fastpath.py b/tests/test_boxset_fastpath.py index 96eda575..40003787 100644 --- a/tests/test_boxset_fastpath.py +++ b/tests/test_boxset_fastpath.py @@ -163,13 +163,13 @@ def test_empty_box(self): %Z_addr = arith.constant 128 : index %X = ktdp.construct_memory_view %X_addr, sizes: [8], strides: [1] { - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<8xf16> %Y = ktdp.construct_memory_view %Y_addr, sizes: [8], strides: [1] { - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<8xf16> %Z = ktdp.construct_memory_view %Z_addr, sizes: [8], strides: [1] { - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<8xf16> %c0 = arith.constant 0 : index diff --git a/tests/test_dialects_parse.py b/tests/test_dialects_parse.py index 622a9007..ab93b35b 100644 --- a/tests/test_dialects_parse.py +++ b/tests/test_dialects_parse.py @@ -844,7 +844,7 @@ def test_construct_memory_view(self): op = self._parse( "%view = ktdp.construct_memory_view %ptr, sizes: [1024], strides: [1]" " { coordinate_set = affine_set<(d0) : (d0 >= 0, -d0 + 1023 >= 0)>," - " memory_space = #ktdp.spyre_memory_space } : memref<1024xf16>", + " memory_space = #ktdp.memory_space } : memref<1024xf16>", args={"%ptr": "index"}, ) self.assert_op_type(op, "ktdp.construct_memory_view") @@ -855,6 +855,34 @@ def test_construct_memory_view(self): self.assert_num_operands(op, 1) self.assert_operand_names(op, "%ptr") + def test_construct_memory_view_ct_local_with_ct_id(self): + # `ct_local` maps to the interpreter's LX, and `ct_id = N` is captured + # as lx_core_id so load/store can route to that tile's scratchpad. + op = self._parse( + "%view = ktdp.construct_memory_view %ptr, sizes: [128], strides: [1]" + " { coordinate_set = affine_set<(d0) : (d0 >= 0, -d0 + 127 >= 0)>," + " memory_space = #ktdp.memory_space }" + " : memref<128xf16>", + args={"%ptr": "index"}, + ) + self.assert_attribute(op, "memory_space", "LX") + self.assert_attribute(op, "lx_core_id", 3) + + def test_construct_memory_view_rejects_pre_rename_spelling(self): + # ktir-mlir-frontend#58 replaced `#ktdp.spyre_memory_space` + # with `#ktdp.memory_space`. The old kind names must + # not be silently accepted (nor fall back to the HBM default). The + # exception type differs by frontend — the regex parser raises + # ValueError on the unknown kind, while the MLIR frontend rejects the + # enum during module parsing — so assert only that it is refused. + with pytest.raises(Exception, match="HBM"): + self._parse( + "%view = ktdp.construct_memory_view %ptr, sizes: [1024], strides: [1]" + " { coordinate_set = affine_set<(d0) : (d0 >= 0, -d0 + 1023 >= 0)>," + " memory_space = #ktdp.memory_space } : memref<1024xf16>", + args={"%ptr": "index"}, + ) + def test_construct_access_tile(self): # construct_access_tile records tile shape and all operands op = self._parse( @@ -913,7 +941,7 @@ def test_affine_set_with_symbolic_dim(self): op = self._parse( "%view = ktdp.construct_memory_view %ptr, sizes: [%n_idx], strides: [1]" " { coordinate_set = affine_set<(d0)[s0] : (d0 >= 0, -d0 + s0 - 1 >= 0)>," - " memory_space = #ktdp.spyre_memory_space } : memref", + " memory_space = #ktdp.memory_space } : memref", args={"%ptr": "index", "%n_idx": "index"}, ) self.assert_op_type(op, "ktdp.construct_memory_view") @@ -929,7 +957,7 @@ def test_construct_memory_view_dynamic_memref_type(self): op = self._parse( "%view = ktdp.construct_memory_view %ptr, sizes: [%n_idx], strides: [1]" " { coordinate_set = affine_set<(d0)[s0] : (d0 >= 0, -d0 + s0 - 1 >= 0)>," - " memory_space = #ktdp.spyre_memory_space } : memref", + " memory_space = #ktdp.memory_space } : memref", args={"%ptr": "index", "%n_idx": "index"}, ) self.assert_op_type(op, "ktdp.construct_memory_view") @@ -944,7 +972,7 @@ def test_construct_memory_view_ssa_size_as_operand(self): op = self._parse( "%view = ktdp.construct_memory_view %ptr, sizes: [%n_idx], strides: [1]" " { coordinate_set = affine_set<(d0)[s0] : (d0 >= 0, -d0 + s0 - 1 >= 0)>," - " memory_space = #ktdp.spyre_memory_space } : memref", + " memory_space = #ktdp.memory_space } : memref", args={"%ptr": "index", "%n_idx": "index"}, ) # %ptr + %n_idx = 2 operands @@ -959,7 +987,7 @@ def test_construct_memory_view_multi_dim_mixed_static_dynamic(self): "%view = ktdp.construct_memory_view %ptr, sizes: [1024, %n], strides: [%n, 1]" " { coordinate_set = affine_set<(d0, d1)[s0] : (d0 >= 0, -d0 + 1023 >= 0," " d1 >= 0, -d1 + s0 - 1 >= 0)>," - " memory_space = #ktdp.spyre_memory_space } : memref<1024x?xf16>", + " memory_space = #ktdp.memory_space } : memref<1024x?xf16>", args={"%ptr": "index", "%n": "index"}, ) self.assert_op_type(op, "ktdp.construct_memory_view") @@ -983,7 +1011,7 @@ def test_construct_memory_view_sizes_count_mismatch_rejected(self): with pytest.raises(ValueError, match=r"sizes count.*does not match|mismatch"): self._parse( "%view = ktdp.construct_memory_view %ptr, sizes: [1024, 32], strides: [1]" - " { memory_space = #ktdp.spyre_memory_space } : memref<1024xf16>", + " { memory_space = #ktdp.memory_space } : memref<1024xf16>", args={"%ptr": "index"}, ) @@ -993,7 +1021,7 @@ def test_construct_memory_view_static_dim_mismatch_rejected(self): with pytest.raises(ValueError, match=r"sizes\[0\]=512 does not match memref dimension 1024"): self._parse( "%view = ktdp.construct_memory_view %ptr, sizes: [512], strides: [1]" - " { memory_space = #ktdp.spyre_memory_space } : memref<1024xf16>", + " { memory_space = #ktdp.memory_space } : memref<1024xf16>", args={"%ptr": "index"}, ) @@ -1016,7 +1044,7 @@ def test_construct_memory_view_dynamic_dim_with_concrete_size_rejected( self._parse( f"%view = ktdp.construct_memory_view %ptr, sizes: {sizes_str}," f" strides: {strides_str}" - f" {{ memory_space = #ktdp.spyre_memory_space }} : {memref_type}", + f" {{ memory_space = #ktdp.memory_space }} : {memref_type}", args=args, ) @@ -1026,7 +1054,7 @@ def test_construct_memory_view_dynamic_dim_no_sizes_rejected(self): with pytest.raises(ValueError, match=r"dynamic dim|'\\?' dim|no sizes"): self._parse( "%view = ktdp.construct_memory_view %ptr, strides: [1]" - " { memory_space = #ktdp.spyre_memory_space } : memref", + " { memory_space = #ktdp.memory_space } : memref", args={"%ptr": "index"}, ) diff --git a/tests/test_distributed_view.py b/tests/test_distributed_view.py index 71318152..9d607b84 100644 --- a/tests/test_distributed_view.py +++ b/tests/test_distributed_view.py @@ -34,6 +34,7 @@ from ktir_cpu import KTIRInterpreter from ktir_cpu.dtypes import stick_to_elem_idx +from ktir_cpu.ir_types import KTDP_MEMORY_SPACE_KINDS from conftest import get_test_params @@ -152,8 +153,12 @@ def _build_mlir(spec: DistCopySpec) -> str: ) idx_refs = ", ".join(f"%idx{i}" for i in range(len(idx))) - p0_ms = f"#ktdp.spyre_memory_space<{p0.memory_space}>" - p1_ms = f"#ktdp.spyre_memory_space<{p1.memory_space}>" + # PartitionSpec.memory_space holds the interpreter's name (it also selects + # the backing memory in the assertions below); the IR needs the dialect + # kind, so invert the parse-side mapping. + _kind = {v: k for k, v in KTDP_MEMORY_SPACE_KINDS.items()} + p0_ms = f"#ktdp.memory_space<{_kind[p0.memory_space]}>" + p1_ms = f"#ktdp.memory_space<{_kind[p1.memory_space]}>" return f""" #P0_set = {p0_set} @@ -181,7 +186,7 @@ def _build_mlir(spec: DistCopySpec) -> str: : memref<{G[0]}x{G[1]}xf16> %B = ktdp.construct_memory_view %B_addr, sizes: [{ac[0]}, {ac[1]}], strides: [{ac[1]}, 1] {{ - coordinate_set = #ac_set, memory_space = #ktdp.spyre_memory_space + coordinate_set = #ac_set, memory_space = #ktdp.memory_space }} : memref<{ac[0]}x{ac[1]}xf16> %A_at = ktdp.construct_access_tile %A[{idx_refs}] {{ diff --git a/tests/test_indirect_access.py b/tests/test_indirect_access.py index 8ee73837..2a48281f 100644 --- a/tests/test_indirect_access.py +++ b/tests/test_indirect_access.py @@ -79,22 +79,22 @@ def _prepare_and_seed(grid_shape): %X = ktdp.construct_memory_view %X_addr, sizes: [4, 4], strides: [4, 1] { coordinate_set = #coord_set_4x4, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<4x4xf16> %IDX1 = ktdp.construct_memory_view %IDX1_addr, sizes: [4, 4], strides: [4, 1] { coordinate_set = #coord_set_4x4, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<4x4xi32> %IDX2 = ktdp.construct_memory_view %IDX2_addr, sizes: [4, 4], strides: [4, 1] { coordinate_set = #coord_set_4x4, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<4x4xi32> %Y_view = ktdp.construct_memory_view %Y_addr, sizes: [4, 4], strides: [4, 1] { coordinate_set = #coord_set_4x4, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<4x4xf16> %X_access_tile = ktdp.construct_indirect_access_tile @@ -173,12 +173,12 @@ def _prepare_and_seed(grid_shape): %X = ktdp.construct_memory_view %X_addr, sizes: [4, 4], strides: [4, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 3 >= 0, d1 >= 0, -d1 + 3 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<4x4xf16> %IDX = ktdp.construct_memory_view %IDX_addr, sizes: [4, 4], strides: [4, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 3 >= 0, d1 >= 0, -d1 + 3 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<4x4xi32> %tile = ktdp.construct_indirect_access_tile @@ -205,12 +205,12 @@ def _prepare_and_seed(grid_shape): %X = ktdp.construct_memory_view %X_addr, sizes: [4, 4], strides: [4, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 3 >= 0, d1 >= 0, -d1 + 3 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<4x4xf16> %IDX = ktdp.construct_memory_view %IDX_addr, sizes: [4, 4], strides: [4, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 3 >= 0, d1 >= 0, -d1 + 3 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<4x4xi32> %tile = ktdp.construct_indirect_access_tile @@ -328,22 +328,22 @@ def _prepare_and_seed(grid_shape): %X_view = ktdp.construct_memory_view %X_addr, sizes: [4, 4], strides: [4, 1] { coordinate_set = #coord_set_4x4, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<4x4xf16> %IDX1 = ktdp.construct_memory_view %IDX1_addr, sizes: [4, 4], strides: [4, 1] { coordinate_set = #coord_set_4x4, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<4x4xi32> %IDX2 = ktdp.construct_memory_view %IDX2_addr, sizes: [4, 4], strides: [4, 1] { coordinate_set = #coord_set_4x4, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<4x4xi32> %Y_view = ktdp.construct_memory_view %Y_addr, sizes: [4, 4], strides: [4, 1] { coordinate_set = #coord_set_4x4, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<4x4xf16> %c0 = arith.constant 0 : index @@ -497,22 +497,22 @@ def _gather_vso_mlir_4x4(vso_str: str, func_name: str) -> str: %X = ktdp.construct_memory_view %X_addr, sizes: [4, 4], strides: [4, 1] {{ coordinate_set = #coord_set_4x4, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space }} : memref<4x4xf16> %IDX1 = ktdp.construct_memory_view %IDX1_addr, sizes: [4, 4], strides: [4, 1] {{ coordinate_set = #coord_set_4x4, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space }} : memref<4x4xi32> %IDX2 = ktdp.construct_memory_view %IDX2_addr, sizes: [4, 4], strides: [4, 1] {{ coordinate_set = #coord_set_4x4, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space }} : memref<4x4xi32> %Y_view = ktdp.construct_memory_view %Y_addr, sizes: [4, 4], strides: [4, 1] {{ coordinate_set = #coord_set_4x4, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space }} : memref<4x4xf16> %X_access_tile = ktdp.construct_indirect_access_tile @@ -553,22 +553,22 @@ def _scatter_vso_mlir_4x4(vso_str: str, func_name: str) -> str: %X_view = ktdp.construct_memory_view %X_addr, sizes: [4, 4], strides: [4, 1] {{ coordinate_set = #coord_set_4x4, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space }} : memref<4x4xf16> %IDX1 = ktdp.construct_memory_view %IDX1_addr, sizes: [4, 4], strides: [4, 1] {{ coordinate_set = #coord_set_4x4, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space }} : memref<4x4xi32> %IDX2 = ktdp.construct_memory_view %IDX2_addr, sizes: [4, 4], strides: [4, 1] {{ coordinate_set = #coord_set_4x4, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space }} : memref<4x4xi32> %Y_view = ktdp.construct_memory_view %Y_addr, sizes: [4, 4], strides: [4, 1] {{ coordinate_set = #coord_set_4x4, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space }} : memref<4x4xf16> %c0 = arith.constant 0 : index @@ -660,17 +660,17 @@ def test_swap_vso(mlir_factory, func_name): %X = ktdp.construct_memory_view %X_addr, sizes: [2, 2, 2], strides: [4, 2, 1] { coordinate_set = #coord_set_2x2x2, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<2x2x2xf16> %IDX = ktdp.construct_memory_view %IDX_addr, sizes: [2, 2, 2], strides: [4, 2, 1] { coordinate_set = #coord_set_2x2x2, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<2x2x2xi32> %Y_view = ktdp.construct_memory_view %Y_addr, sizes: [2, 2, 2], strides: [4, 2, 1] { coordinate_set = #coord_set_2x2x2, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<2x2x2xf16> %X_access_tile = ktdp.construct_indirect_access_tile @@ -750,17 +750,17 @@ def _prepare_and_seed(grid_shape): %X_view = ktdp.construct_memory_view %X_addr, sizes: [2, 2, 2], strides: [4, 2, 1] { coordinate_set = #coord_set_2x2x2, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<2x2x2xf16> %IDX = ktdp.construct_memory_view %IDX_addr, sizes: [2, 2, 2], strides: [4, 2, 1] { coordinate_set = #coord_set_2x2x2, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<2x2x2xi32> %Y_view = ktdp.construct_memory_view %Y_addr, sizes: [2, 2, 2], strides: [4, 2, 1] { coordinate_set = #coord_set_2x2x2, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<2x2x2xf16> %c0 = arith.constant 0 : index @@ -974,12 +974,12 @@ def resolve(node): %src = ktdp.construct_memory_view %src_addr, sizes: [128], strides: [1] { coordinate_set = #src_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<128xf16> %dst = ktdp.construct_memory_view %dst_addr, sizes: [2, 64], strides: [64, 1] { coordinate_set = #dst_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<2x64xf16> %src_tile = ktdp.construct_indirect_access_tile diff --git a/tests/test_latency_modeling.py b/tests/test_latency_modeling.py index c9a47d36..23efada2 100644 --- a/tests/test_latency_modeling.py +++ b/tests/test_latency_modeling.py @@ -89,7 +89,7 @@ def _bw_per_core(cfg, cores_active): %x_view = ktdp.construct_memory_view %x_ptr, sizes: [{total}], strides: [1] {{ coordinate_set = affine_set<(d0) : (d0 >= 0, -d0 + {total_m1} >= 0)>, - memory_space = #ktdp.spyre_memory_space }} + memory_space = #ktdp.memory_space }} : index -> memref<{total}xf16> %x_acc = ktdp.construct_access_tile %x_view[%offset] {{ base_map = affine_map<(i) -> (i)>, @@ -100,7 +100,7 @@ def _bw_per_core(cfg, cores_active): %out_view = ktdp.construct_memory_view %out_ptr, sizes: [{total}], strides: [1] {{ coordinate_set = affine_set<(d0) : (d0 >= 0, -d0 + {total_m1} >= 0)>, - memory_space = #ktdp.spyre_memory_space }} + memory_space = #ktdp.memory_space }} : index -> memref<{total}xf16> %out_acc = ktdp.construct_access_tile %out_view[%offset] {{ base_map = affine_map<(i) -> (i)>, @@ -113,8 +113,9 @@ def _bw_per_core(cfg, cores_active): """) # Copy kernel template: loads 128 f16 elements and stores them back. -# memory_space controls whether memory_cycles are charged (LX → 0, HBM → positive). -# Instantiated as _LX_MLIR and _HBM_MLIR for test_lx_reuse_vs_hbm_reload. +# memory_space is a #ktdp.memory_space kind ("ct_local"/"global") and controls +# whether memory_cycles are charged (ct_local/LX → 0, global/HBM → positive). +# Instantiated for both kinds by test_lx_reuse_vs_hbm_reload. def _copy_mlir(func_name: str, memory_space: str) -> str: return textwrap.dedent(f"""\ module {{ @@ -125,7 +126,7 @@ def _copy_mlir(func_name: str, memory_space: str) -> str: %x_view = ktdp.construct_memory_view %x_ptr, sizes: [128], strides: [1] {{ coordinate_set = affine_set<(d0) : (d0 >= 0, -d0 + 127 >= 0)>, - memory_space = #ktdp.spyre_memory_space<{memory_space}> }} + memory_space = #ktdp.memory_space<{memory_space}> }} : index -> memref<128xf16> %c0 = arith.constant {{0 : index}} : index %x_acc = ktdp.construct_access_tile %x_view[%c0] {{ @@ -136,7 +137,7 @@ def _copy_mlir(func_name: str, memory_space: str) -> str: %out_view = ktdp.construct_memory_view %out_ptr, sizes: [128], strides: [1] {{ coordinate_set = affine_set<(d0) : (d0 >= 0, -d0 + 127 >= 0)>, - memory_space = #ktdp.spyre_memory_space<{memory_space}> }} + memory_space = #ktdp.memory_space<{memory_space}> }} : index -> memref<128xf16> %c0b = arith.constant {{0 : index}} : index %out_acc = ktdp.construct_access_tile %out_view[%c0b] {{ @@ -654,8 +655,8 @@ def test_lx_reuse_vs_hbm_reload(self): """ cfg = HardwareConfig() - lx_report = _run_inline(_copy_mlir("lx_kernel", "LX"), "lx_kernel", cfg=cfg, seed_lx=True) - hbm_report = _run_inline(_copy_mlir("hbm_kernel", "HBM"), "hbm_kernel", cfg=cfg) + lx_report = _run_inline(_copy_mlir("lx_kernel", "ct_local"), "lx_kernel", cfg=cfg, seed_lx=True) + hbm_report = _run_inline(_copy_mlir("hbm_kernel", "global"), "hbm_kernel", cfg=cfg) lx_mem = lx_report.counters[0].memory_cycles hbm_mem = hbm_report.counters[0].memory_cycles diff --git a/tests/test_lx_scoping.py b/tests/test_lx_scoping.py index a91712b0..19a24c0e 100644 --- a/tests/test_lx_scoping.py +++ b/tests/test_lx_scoping.py @@ -48,15 +48,15 @@ %pid_m, %pid_n = ktdp.get_compute_tile_id : index, index %a_view = ktdp.construct_memory_view %a_ptr, sizes: [8, 8], strides: [8, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 7 >= 0, d1 >= 0, -d1 + 7 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<8x8xf16> %b_view = ktdp.construct_memory_view %b_ptr, sizes: [8, 8], strides: [8, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 7 >= 0, d1 >= 0, -d1 + 7 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<8x8xf16> %c_view = ktdp.construct_memory_view %c_ptr, sizes: [8, 8], strides: [8, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 7 >= 0, d1 >= 0, -d1 + 7 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<8x8xf16> %c0 = arith.constant 0 : index %accum_zero = arith.constant dense<0.0> : tensor<8x8xf16> @@ -105,17 +105,17 @@ %a_view = ktdp.construct_memory_view %a_ptr, sizes: [512, 256], strides: [256, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 511 >= 0, d1 >= 0, -d1 + 255 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<512x256xf16> %b_view = ktdp.construct_memory_view %b_ptr, sizes: [256, 256], strides: [256, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 255 >= 0, d1 >= 0, -d1 + 255 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<256x256xf16> %c_view = ktdp.construct_memory_view %c_ptr, sizes: [512, 256], strides: [256, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 511 >= 0, d1 >= 0, -d1 + 255 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<512x256xf16> %c0 = arith.constant 0 : index @@ -1225,17 +1225,17 @@ def test_addf_loop_body_lx_accounting(self): %a_view = ktdp.construct_memory_view %a_ptr, sizes: [64, 64], strides: [64, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 63 >= 0, d1 >= 0, -d1 + 63 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<64x64xf16> %b_view = ktdp.construct_memory_view %b_ptr, sizes: [64, 32], strides: [32, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 63 >= 0, d1 >= 0, -d1 + 31 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<64x32xf16> %c_view = ktdp.construct_memory_view %c_ptr, sizes: [64, 32], strides: [32, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 63 >= 0, d1 >= 0, -d1 + 31 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<64x32xf16> %c0 = arith.constant 0 : index @@ -1511,9 +1511,9 @@ def correct_handler(op, context, env): %bm_end = arith.addi %pid_0, %c1_i32 : i32 %bm_end_1 = arith.minsi %bm_end, %c32_i32 : i32 - %a_desc = ktdp.construct_memory_view %arg0, sizes: [4, 128, 32], strides: [4096, 32, 1] {coordinate_set = #set, memory_space = #ktdp.spyre_memory_space} : memref<4x128x32xf32> - %b_desc = ktdp.construct_memory_view %arg1, sizes: [4, 32, 64], strides: [2048, 64, 1] {coordinate_set = #set1, memory_space = #ktdp.spyre_memory_space} : memref<4x32x64xf32> - %c_desc = ktdp.construct_memory_view %arg2, sizes: [4, 128, 64], strides: [8192, 64, 1] {coordinate_set = #set2, memory_space = #ktdp.spyre_memory_space} : memref<4x128x64xf32> + %a_desc = ktdp.construct_memory_view %arg0, sizes: [4, 128, 32], strides: [4096, 32, 1] {coordinate_set = #set, memory_space = #ktdp.memory_space} : memref<4x128x32xf32> + %b_desc = ktdp.construct_memory_view %arg1, sizes: [4, 32, 64], strides: [2048, 64, 1] {coordinate_set = #set1, memory_space = #ktdp.memory_space} : memref<4x32x64xf32> + %c_desc = ktdp.construct_memory_view %arg2, sizes: [4, 128, 64], strides: [8192, 64, 1] {coordinate_set = #set2, memory_space = #ktdp.memory_space} : memref<4x128x64xf32> scf.for %arg3 = %pid_0 to %bm_end_1 step %c1_i32 : i32 { %b = arith.divsi %arg3, %m_blocks : i32 diff --git a/tests/test_parser_utils.py b/tests/test_parser_utils.py index 5c299d90..edc32f79 100644 --- a/tests/test_parser_utils.py +++ b/tests/test_parser_utils.py @@ -21,7 +21,11 @@ import pytest -from ktir_cpu.parser_utils import parse_multi_result_lhs, parse_tensor_or_memref_type, extract_outs_operands +from ktir_cpu.parser_utils import ( + extract_outs_operands, + parse_multi_result_lhs, + parse_tensor_or_memref_type, +) # --------------------------------------------------------------------------- @@ -124,7 +128,7 @@ def test_parse_tensor_or_memref_type_rejects_invalid(type_str): ("tensor<4xf32>, %arg0", (4,), "f32"), # memref wrapper ("memref<128x32xf16>", (128, 32), "f16"), - ("memref<4x4xi32, #ktdp.spyre_memory_space>", (4, 4), "i32"), + ("memref<4x4xi32, #ktdp.memory_space>", (4, 4), "i32"), # Bare inner content (no wrapper) ("64x32xindex", (64, 32), "index"), ("256xbf16", (256,), "bf16"), diff --git a/tests/test_spec_gaps.py b/tests/test_spec_gaps.py index 2bf25af0..0fa19969 100644 --- a/tests/test_spec_gaps.py +++ b/tests/test_spec_gaps.py @@ -187,11 +187,11 @@ def test_tensor_extract_slice(): %dst = arith.constant 256 : index %src_view = ktdp.construct_memory_view %src, sizes: [8, 8], strides: [8, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 7 >= 0, d1 >= 0, -d1 + 7 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<8x8xf16> %dst_view = ktdp.construct_memory_view %dst, sizes: [4, 4], strides: [4, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 3 >= 0, d1 >= 0, -d1 + 3 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<4x4xf16> %src_access = ktdp.construct_access_tile %src_view[%c0, %c0] { access_tile_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 7 >= 0, d1 >= 0, -d1 + 7 >= 0)>, @@ -231,7 +231,7 @@ def test_scf_parallel(): %dst = arith.constant 0 : index %dst_view = ktdp.construct_memory_view %dst, sizes: [4, 1], strides: [1, 1] { coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 3 >= 0, d1 >= 0)>, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<4x1xf16> %dst_access = ktdp.construct_access_tile %dst_view[%c0, %c0] { access_tile_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 3 >= 0, d1 >= 0)>, diff --git a/tests/test_tile.py b/tests/test_tile.py index ed906d75..5cd5cf4d 100644 --- a/tests/test_tile.py +++ b/tests/test_tile.py @@ -52,7 +52,7 @@ class TestTileAccess: %Y_view = ktdp.construct_memory_view %Y_addr, sizes: [64, 64], strides: [64, 1] { coordinate_set = #Y_coord_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<64x64xf16> The parser copies sizes → shape and strides verbatim into a TileRef. diff --git a/uv.lock b/uv.lock index 94d5ecaf..4dc7e831 100644 --- a/uv.lock +++ b/uv.lock @@ -1105,7 +1105,7 @@ dev = [ [package.metadata] requires-dist = [ - { name = "ktir-mlir-frontend", marker = "extra == 'mlir-frontend'", git = "https://github.com/torch-spyre/ktir-mlir-frontend?rev=973d67ee7d535ea41f4adbf9c6eb4ee59450672b" }, + { name = "ktir-mlir-frontend", marker = "extra == 'mlir-frontend'", git = "https://github.com/torch-spyre/ktir-mlir-frontend?rev=ecfb9ed73b6d40a19df8f678c56ed3074bbe2f85" }, { name = "numpy", specifier = ">=1.20.0" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=7.0" }, { name = "pytest-cov", marker = "extra == 'dev'", specifier = ">=4.0" }, @@ -1122,7 +1122,7 @@ dev = [ [[package]] name = "ktir-mlir-frontend" version = "0.1.0" -source = { git = "https://github.com/torch-spyre/ktir-mlir-frontend?rev=973d67ee7d535ea41f4adbf9c6eb4ee59450672b#973d67ee7d535ea41f4adbf9c6eb4ee59450672b" } +source = { git = "https://github.com/torch-spyre/ktir-mlir-frontend?rev=ecfb9ed73b6d40a19df8f678c56ed3074bbe2f85#ecfb9ed73b6d40a19df8f678c56ed3074bbe2f85" } dependencies = [ { name = "nanobind" }, ]