diff --git a/.codex/skills/pto-gym-vpto-validation/SKILL.md b/.codex/skills/pto-gym-vpto-validation/SKILL.md index 0e1451a614..721df32b6f 100644 --- a/.codex/skills/pto-gym-vpto-validation/SKILL.md +++ b/.codex/skills/pto-gym-vpto-validation/SKILL.md @@ -1,6 +1,6 @@ --- name: pto-gym-vpto-validation -description: Run PTO-Gym validation from this PTOAS repo. Use when the user asks to run PTO-Gym SIM or board validation from the current source tree. Always force PTOAS onto the VPTO LLVM path instead of relying on the repo default backend. +description: Run bundled PTO-Gym exercise/validation cases. Use when the user explicitly asks for PTO-Gym, 3rdparty/PTO-Gym, or the PTO-Gym validation scripts. Always force PTOAS onto the VPTO path instead of relying on the repo default backend. --- # PTO-Gym VPTO Validation @@ -8,20 +8,20 @@ description: Run PTO-Gym validation from this PTOAS repo. Use when the user asks Use this skill when the task is specifically about: - running `3rdparty/PTO-Gym/examples/pto/scripts/run_host_vpto_validation.sh` - running `3rdparty/PTO-Gym/examples/pto/scripts/run_host_vpto_validation_parallel.sh` -- validating PTO-Gym cases from this PTOAS source tree +- validating bundled PTO-Gym exercise cases ## Required Rule When PTO-Gym is run from this repo, do not rely on the default PTOAS backend. Always pass PTOAS flags that force the VPTO LLVM path. -The current `ptoas` CLI spellings in this repo are `--pto-backend=vpto` and -`--vpto-emit-hivm-llvm`; do not shorten `--pto-backend` to `--backend`. +The current `ptoas` CLI spelling in this repo is `--pto-backend=vpto`; do not +shorten `--pto-backend` to `--backend`. Use: ```bash -PTOAS_FLAGS='--pto-backend=vpto --vpto-emit-hivm-llvm --pto-arch a5' +PTOAS_FLAGS='--pto-backend=vpto --pto-arch a5' ``` If the caller already provides `PTOAS_FLAGS`, make sure these options are still @@ -44,7 +44,7 @@ Typical simulator environment: source /home/mouliangyu/.local/ascend/beta.2/cann-9.0.0-beta.2/set_env.sh export ASCEND_HOME_PATH=/home/mouliangyu/.local/ascend/beta.2/cann-9.0.0-beta.2 export PTOAS_BIN=$PWD/build/tools/ptoas/ptoas -export PTOAS_FLAGS='--pto-backend=vpto --vpto-emit-hivm-llvm --pto-arch a5' +export PTOAS_FLAGS='--pto-backend=vpto --pto-arch a5' ``` ## Canonical Commands diff --git a/.codex/skills/ptoas-project-development/SKILL.md b/.codex/skills/ptoas-project-development/SKILL.md index 07c34dae84..b3cff9fab2 100644 --- a/.codex/skills/ptoas-project-development/SKILL.md +++ b/.codex/skills/ptoas-project-development/SKILL.md @@ -9,6 +9,12 @@ description: Project development guidance for PTOAS. Use when Codex modifies PTO When changing any user-visible behavior, update every relevant layer in the same change. Treat ODS, verifiers, lowering, command-line behavior, bindings, docs, examples, and tests as one public contract. +## File Orientation Before Editing + +Before editing any existing file, read the file's header or top-level comments first. Check whether they describe the file's purpose, structure, section boundaries, or editing constraints, and follow those constraints when making changes. + +When a change creates or depends on a file-level structure rule, ownership boundary, or important functional description, add or update that guidance in the file's top-level comments so future edits see it before jumping into local code. + ## Layers To Keep In Sync 1. ODS and dialect definitions: `include/PTO/IR/*.td` diff --git a/.codex/skills/rewrite-kernel-with-vmi/SKILL.md b/.codex/skills/rewrite-kernel-with-vmi/SKILL.md new file mode 100644 index 0000000000..9e355964ee --- /dev/null +++ b/.codex/skills/rewrite-kernel-with-vmi/SKILL.md @@ -0,0 +1,229 @@ +--- +name: rewrite-kernel-with-vmi +description: Rewrite a complete AscendC/CANN kernel into an equivalent PTODSL Python implementation that mixes VMI compute with MI/MTE/sync orchestration. Use when Codex is given AscendC kernel code or CCE-style device code and must preserve the kernel ABI, split compute from movement/synchronization, translate pure vector/SIMT compute regions into pto.vmi semantics, translate non-compute regions into PTODSL micro-instructions, and validate the result with kernel.compile().mlir_text() plus the PTOAS VMI path. +--- + +# Rewrite Kernel With VMI + +Rewrite the source kernel by preserving observable semantics first, then choosing +the clean PTODSL spelling. Do not mechanically replay physical register +choreography when a logical VMI expression captures the same algorithm. + +## Workflow + +1. Read the complete source kernel and all helpers/macros needed to understand + its ABI, loop bounds, offsets, memory movement, synchronization, and compute. +2. Read references only as needed: + - [vmi-dsl-spec.md](references/vmi-dsl-spec.md) for current PTODSL `pto.vmi` + API spelling, including `vload` / `vstore`, `create_mask`, `vci`, `vbrc`, + and common AscendC SIMD to VMI patterns. + - [mi-dsl-spec.md](references/mi-dsl-spec.md) for PTODSL non-VMI user-guide + navigation: kernel entry, buffers, control flow, MTE, sync, masks, SIMT. + - [vmi-mlir-spec.md](references/vmi-mlir-spec.md) when VMI IR semantics, + layout, mask, or PTOAS validation details matter. +3. Extract the PTODSL function interface: + - Convert each host-visible kernel argument one-for-one. + - Convert C++ template parameters to keyword-only `pto.const_expr` arguments. + - For template/generic variants, prefer one PTODSL function with + `pto.const_expr` parameters that select dtype, constants, and specialized + branches at compile time. Avoid generating many outer wrapper/probe + functions that only differ by dtype or template value. + - Use this default entry unless the source requires otherwise: + +```python +@pto.jit( + name="", + target="a5", + backend="vpto", + mode="explicit", + kernel_kind="vector", + insert_sync=False, +) +def (..., *, CONST: pto.const_expr = ...): + ... +``` + +Compile-time selection example: + +```python +@pto.jit(...) +def kernel(x: pto.ptr(pto.f8e5m2, "ub"), y_addr: pto.i64, *, OUT_DTYPE: pto.const_expr = pto.f32): + y = pto.castptr(y_addr, pto.ptr(OUT_DTYPE, "ub")) + if OUT_DTYPE is pto.f32: + ... + elif OUT_DTYPE is pto.bf16: + ... +``` + +4. Split the source body into regions: + - Compute regions: pure vector/SIMT arithmetic, compare/select, conversion, + reduction, rearrange, or math. `membar` may stay in or near a compute + region when it only orders vector-visible UB effects. + - Non-compute regions: GM/UB/L1 movement, tile allocation, pointer/view + setup, `set_flag`/`wait_flag`, cross/intra flag sync, buffer handoff, + pipeline barriers, and host/core indexing. +5. Before writing PTODSL for each nontrivial compute region, show the user a + VMI rewrite design and wait for confirmation when the interaction allows it. + If the user is unavailable or explicitly asked for an end-to-end conversion, + continue with clearly stated assumptions. +6. Implement the PTODSL Python file: + - Use `from ptodsl import pto` and add `scalar` only when needed. + - Keep imports environment-agnostic. Do not add `Path(__file__)` parent + walks, local repo discovery, or `sys.path.insert(...)` blocks to find the + `ptodsl` package; configure installation, runner paths, or `PYTHONPATH` + outside the generated DSL file instead. + - Keep the source ABI recognizable. + - Write VMI/MI operations inline by default. Create small helper functions + only when a nontrivial instruction sequence is reused many times or when a + named helper preserves an important source-level abstraction. + - Use native Python `for` / `if` control flow. Use `pto.const_expr` and + `pto.static_range` for intentional compile-time specialization or + unrolling. Do not use `pto.for_` / `pto.if_` in new rewrites. + - Translate non-compute regions with PTODSL MI/MTE/sync operations. + - Translate compute regions with `pto.vmi` logical vectors. +7. Validate and iterate: + - First run or provide a compile helper that calls + `kernel.compile(...).mlir_text()`. + - Then pass the emitted MLIR to PTOAS with the VMI path: + `ptoas --pto-arch=a5 --pto-backend=vpto --enable-vmi -o /dev/null`. + - When the goal is to inspect lowered VMI output locally and no CANN/toolchain + environment is needed, prefer the bundled helper script: + `scripts/compile_vmi_to_vpto.sh [output.mi.pto]`. + It wraps `ptoas --pto-backend=vpto --enable-vmi --emit-vpto` and is the + default path for generating reviewable lowered artifacts from emitted + `vmi.pto` files. + - Fix PTODSL trace errors before PTOAS errors. Fix semantic mismatches before + layout/lowering workarounds. + +## VMI Rewrite Design Gate + +For every nontrivial compute region, present this compact review artifact before +coding the region: + +```text +Compute region: +Inputs: + - , offset/stride formula, dtype, shape +Outputs: + - , offset/stride formula, dtype, shape +Semantic algorithm: + - +Physical-only source details to collapse: + - +VMI plan: + for off in ...: + mask = pto.vmi.create_mask(active_lanes, size=lanes) # or size=lanes, group=groups + x = pto.vmi.vload(..., size=lanes) + y = + pto.vmi.vstore(y, ..., mask) +Lane choice: + - lanes=, dtype=, reason= +Assumptions: + - +``` + +Ask the user to confirm or correct this design when the algorithm is ambiguous, +when the source uses heavy physical packing/interleave, or when multiple VMI +forms are plausible. If there is no reply and progress is requested, implement +the stated assumptions and call them out in the result. + +## Translation Rules + +- Preserve the boundary contract. Do not drop source arguments merely because a + first rewrite does not use them yet. +- Do not embed local import bootstrapping in generated DSL code. The file should + import PTODSL normally, e.g. `from ptodsl import pto`, without walking parent + directories or mutating `sys.path` to locate a workspace checkout. +- Convert common source types as follows: + - `__gm__ half*` / semantic f16 storage -> `pto.ptr(pto.f16, "gm")` + - `__gm__ float*` -> `pto.ptr(pto.f32, "gm")` + - `int32_t` -> `pto.i32` + - `uint32_t` -> `pto.ui32` + - `int64_t` / addresses and byte counters -> `pto.i64` unless source intent is + clearly index-like. +- If a C++ boundary uses raw storage (`uint16_t*`, `uint8_t*`) but the semantic + element type is `f16`, `bf16`, fp8, or packed f4, keep the safest ABI spelling + and use `pto.castptr` internally when needed. +- Preserve templates and generic dtype choices with `pto.const_expr` whenever + possible. Use Python compile-time `if` branches and dtype variables to choose + pointer casts, VMI lane/dtype choices, conversion targets, and store paths. Only add + separate wrapper/probe functions when the ABI truly differs or the test + harness explicitly requires separate entry symbols. +- Use the current VMI surface shape: `vload(..., size=...)`, `vstore(..., mask=...)`, + `vci(..., size=...)`, `vbrc(..., size=...)`, and `create_mask(..., size=..., group=...)`. + Do not use the retired `create_group_mask` helper or legacy `result_type` + spellings in new rewrites. +- Keep source offset and stride formulas symbolic. Do not replace expressions + such as `vlForHalfNumber * 2` with a constant unless the source is already + specialized and the user asked for specialization. +- Avoid small one-off helpers around VMI/MI instruction sequences. Inline direct + PTODSL operations unless the code is long, reused repeatedly, or the source + abstraction is important enough to preserve by name. +- Use Python-native control flow for both dynamic device-side branches/loops and + ordinary structured code. Use `pto.static_range(...)` only for trace-time + static loops driven by Python values or `pto.const_expr` parameters. Do not + write `with pto.for_(...)` or `with pto.if_(...)` in new kernel rewrites. +- Use explicit `mode="explicit"` orchestration for full kernel rewrites. Do not + rely on auto-inserted sync unless the user requests an auto-mode rewrite. +- Prefer `pto.vmi.vload` from UB, logical compute, then `pto.vmi.vstore` to UB. + GM movement belongs to MTE/tile movement outside the VMI compute region. +- Use `pto.vmi.create_mask(active, size=lanes)` for dynamic tails. Put masks on + compute and store; do not assume a masked load is legal on every backend. + For grouped tails, use `create_mask(..., size=lanes, group=...)`. +- Collapse physical-only details such as `PART_P0`, `PART_EVEN`, `PART_ODD`, + packed store modes, and interleave trees when they only describe hardware + lowering of one logical vector. +- Do not introduce UB store+reload round trips just to satisfy a lowering issue + unless the user accepts the performance tradeoff. +- For TileLang input, treat it as future scope unless the user explicitly asks. + Preserve the same VMI design-gate workflow and translate the logical parallel + loop body rather than physicalizing it early. + +## Optimization Tips + +- When the logical algorithm is "an `N x VL` vector where each `VL` chunk is + multiplied by the same `VL`-lane scale vector", prefer one widened/grouped + VMI expression over a scalarized inner chunk loop. + - Good PTODSL spelling when legal on the current backend: + +```python +wide_x = pto.vmi.vcvt(x, pto.f32) # e.g. N*VL lanes +scale_wide = pto.vmi.vload(scale_ptr, scale_off, size=N * VL, stride=0, group=N) +wide_y = pto.vmi.vmul(wide_x, scale_wide, full_mask) +``` + + - This pattern is often better than: + 1. spilling the widened `N x VL` value to UB, + 2. reloading `VL` chunks in a Python loop, + 3. multiplying each chunk by the same `VL` scale vector, + 4. storing chunk-by-chunk. + - In practice, this zero-stride grouped `vload` is a useful way to express + "repeat the same `VL` vector across `N` groups" directly in VMI, keeping + the computation at full-vector width. + - Still validate with both `kernel.compile(...).mlir_text()` and PTOAS VMI + lowering, because legality depends on the source dtype, lane count, and + backend support for the grouped load shape. + +## Validation Checklist + +Before finishing a rewrite, verify: + +- Function parameters and template/constexpr parameters match the source ABI. +- Template/generic variants are represented by `pto.const_expr` compile-time + selection unless separate entry symbols are explicitly needed. +- Pointer memory spaces and semantic dtypes are justified. +- GM/UB/L1 movement sizes, offsets, strides, and padding match the source. +- Sync ordering between MTE, Vector, Cube/SIMT, and stores is preserved. +- New VMI/MI code is not hidden behind one-off helper functions. +- Generated DSL code has no manual `ptodsl` path lookup, `Path(__file__)` + parent search, or `sys.path` mutation for local package discovery. +- Control flow uses native Python `for` / `if`, with `pto.static_range` only for + intentional compile-time loops. +- Every VMI compute region has a user-visible design note or stated assumptions. +- Tail masks and lane counts match the processed element count. +- `kernel.compile(...).mlir_text()` succeeds or the remaining trace error is + reported with the exact failing construct. +- PTOAS validation uses `--pto-backend=vpto --enable-vmi`; do not diagnose VMI + legality from a plain non-VMI invocation. + +Numeric NPU or simulator validation is optional unless the user requests it. diff --git a/.codex/skills/rewrite-kernel-with-vmi/agents/openai.yaml b/.codex/skills/rewrite-kernel-with-vmi/agents/openai.yaml new file mode 100644 index 0000000000..0a9d6e12ec --- /dev/null +++ b/.codex/skills/rewrite-kernel-with-vmi/agents/openai.yaml @@ -0,0 +1,4 @@ +interface: + display_name: "Rewrite Kernel With VMI" + short_description: "AscendC kernel to PTODSL VMI/MI rewrite" + default_prompt: "Use $rewrite-kernel-with-vmi to rewrite this AscendC kernel into PTODSL VMI/MI and validate the generated MLIR." diff --git a/.codex/skills/rewrite-kernel-with-vmi/references/mi-dsl-spec.md b/.codex/skills/rewrite-kernel-with-vmi/references/mi-dsl-spec.md new file mode 100644 index 0000000000..02c588eb02 --- /dev/null +++ b/.codex/skills/rewrite-kernel-with-vmi/references/mi-dsl-spec.md @@ -0,0 +1,166 @@ +# PTODSL MI And Non-VMI Index + +Use this file as a navigation index for non-VMI parts of an AscendC kernel +rewrite. The source of truth is `ptodsl/docs/user_guide/`. + +## Kernel Entry And ABI + +Read `03-kernel-entry-and-subkernels.md` for: + +- `@pto.jit` entry vs module roles. +- `backend="vpto"` and `mode="explicit"`. +- Pointer-first host ABI: `pto.ptr(dtype, "gm")`. +- Runtime scalar parameters before `*`. +- Compile-time constants after `*` as `pto.const_expr`. +- `.compile(...).mlir_text()` usage. + +Use normal PTODSL imports in generated DSL: + +```python +from ptodsl import pto +``` + +Add `scalar` or other PTODSL imports only when the code actually needs them. Do +not add standalone-script bootstrapping that scans `Path(__file__).parents`, +finds a local `ptodsl` checkout, or calls `sys.path.insert(...)`. If a test or +manual run cannot import PTODSL, fix the execution environment instead: install +the package, invoke the repository's runner, or set `PYTHONPATH` outside the DSL +file. + +Prefer `pto.const_expr` for C++ template parameters and dtype-generic variants. +One PTODSL function can select pointer casts, dtypes, tile shapes, VMI result +types, and store paths at compile time instead of creating many wrapper/probe +functions that only differ by template value. + +Default full-kernel rewrite entry: + +```python +@pto.jit( + name="...", + target="a5", + backend="vpto", + mode="explicit", + kernel_kind="vector", + insert_sync=False, +) +def kernel(..., *, CONST: pto.const_expr = ...): + ... +``` + +Example: + +```python +@pto.jit(...) +def kernel(y_addr: pto.i64, *, OUT_DTYPE: pto.const_expr = pto.f32): + y = pto.castptr(y_addr, pto.ptr(OUT_DTYPE, "ub")) + if OUT_DTYPE is pto.f32: + ... + elif OUT_DTYPE is pto.bf16: + ... +``` + +## Types, Buffers, Views + +Read `04-type-system-and-buffer.md` for: + +- Scalar annotations: `pto.i32`, `pto.ui32`, `pto.i64`, `pto.f16`, `pto.f32`. +- Pointer types and memory spaces: `pto.ptr(dtype, "gm")`, + `pto.ptr(dtype, "ub")`, `pto.MemorySpace.*`. +- `pto.make_tensor_view`, `pto.partition_view`, `.as_ptr()`. +- `pto.alloc_tile`, tile shape and `valid_shape`. +- `pto.castptr` when raw storage pointer spelling differs from semantic dtype. + +## Control Flow And Scalar/Pointer Math + +Read `05-control-flow.md` and `06-scalar-and-pointer-ops.md` for: + +- Python-native `for` / `if` control flow, which PTODSL rewrites into + device-side structured control flow when bounds or conditions are runtime PTO + values. +- `pto.const_expr` and `pto.static_range(...)` for intentional trace-time + specialization and unrolling. +- Scalar casts, selects, index math, pointer casts, pointer offsets. +- Keeping loop bounds and offset formulas symbolic. + +Do not use `pto.for_` / `pto.if_` in new rewrites. Those explicit context-manager +forms exist in older examples and compatibility surfaces, but this skill should +prefer native Python control flow. When hardware-loop lowering is important, +preserve source-like runtime loop bounds and avoid turning loop math into Python +constants. + +## Helper Function Policy + +Write VMI and MI instruction sequences inline by default. Avoid small helpers +that wrap only a few operations, because they obscure the rewrite and can add +unhelpful call boundaries in generated IR. Introduce a helper only when: + +- The same nontrivial sequence is reused many times. +- A long block becomes materially easier to audit by naming the source-level + abstraction. +- A real PTODSL sub-kernel boundary is required, such as Cube/SIMT ownership. + +## Data Movement: MTE And Tile Movement + +Read `07-data-movement-ops.md` for: + +- Tile-level `pto.tile.load` / `pto.tile.store` in auto-mode style code. +- Explicit mode DMA: + - `pto.mte_gm_ub(gm_src, ub_dst, l2_cache_ctl, len_burst, nburst=..., loops=..., pad=...)` + - `pto.mte_ub_gm(ub_src, gm_dst, len_burst, nburst=..., loops=...)` + - `pto.mte_ub_ub(ub_src, ub_dst, len_burst, nburst=...)` + - `pto.mte_ub_l1(...)` +- Shorthands `pto.mte_load` and `pto.mte_store` when they match the canonical + grouped-DMA shape. + +Stride units matter: + +- GM/UB `mte_gm_ub` and `mte_ub_gm`: byte strides. +- UB/UB `mte_ub_ub`: 32B units for gaps/lengths. +- VMI `vload/vstore`: element offsets and element strides. + +## Legacy Vector Compute And Masks + +Read `08-compute-operations.md` and `09-predicate-and-mask-ops.md` for: + +- Existing top-level vector helpers such as `pto.vadd`, `pto.vlds`, `pto.vsts`. +- Legacy masks such as `pto.make_mask`, `pto.mask_b16`, `pto.mask_b32`. + +For this skill, use `pto.vmi` for newly translated pure compute regions unless +the user explicitly asks to keep legacy vector helpers. + +## Synchronization + +Read `10-sync-ops.md` for: + +- `pto.set_flag(pipe_from, pipe_to, event_id=...)` +- `pto.wait_flag(pipe_from, pipe_to, event_id=...)` +- `pto.pipe_barrier(pto.Pipe.ALL)` +- `pto.mem_bar(pto.BarrierType.VV_ALL)` and other barrier types. +- `pto.get_buf` / `pto.rls_buf` for double buffering. +- `pto.set_cross_flag` / `pto.wait_cross_flag`. +- `pto.set_intra_flag` / `pto.wait_intra_flag`. + +Common explicit-mode pairs: + +- DMA load before vector compute: + `set_flag(MTE2, V)` then `wait_flag(MTE2, V)`. +- Vector compute before DMA store: + `set_flag(V, MTE3)` then `wait_flag(V, MTE3)`. +- Store-to-load ordering in UB: + `pto.mem_bar(pto.BarrierType.VST_VLD)`. + +In `mode="explicit"`, do not rely on compiler-inserted synchronization unless +the user requests an auto-sync rewrite. + +## SIMT Micro-Ops + +Read `13-simt-micro-ops.md` when the AscendC compute region uses SIMT-like +scalar lane programming rather than vector SIMD. Preserve SIMT control flow and +only translate to VMI when the algorithm is actually independent SIMD lanes. + +## Examples To Inspect + +- `ptodsl/examples/tadd_launch.py`: simple PTODSL entry/compile/launch shape. +- `ptodsl/examples/flash_attention/flash_attention_cv_split.py`: mixed + orchestration, modules, sync, and explicit kernel structure. +- `ptodsl/examples/mixed_backend_kernel_module.py`: entry/module composition. diff --git a/.codex/skills/rewrite-kernel-with-vmi/references/vmi-dsl-spec.md b/.codex/skills/rewrite-kernel-with-vmi/references/vmi-dsl-spec.md new file mode 100644 index 0000000000..fbea4274e9 --- /dev/null +++ b/.codex/skills/rewrite-kernel-with-vmi/references/vmi-dsl-spec.md @@ -0,0 +1,177 @@ +# PTODSL VMI DSL Quick Reference + +This is a compact guide derived from +`ptodsl/docs/user_guide/14-vmi-virtual-instruction-set.md`. Use the user guide +as the source of truth when an operation is missing here. + +## Core Types + +- `pto.vmi.vreg(lanes, dtype)` creates a logical vector type. + `lanes` must be a multiple of 64. Common full-register choices: + - `pto.f32` / `pto.i32`: 64 lanes per 256B physical vreg. + - `pto.f16` / `pto.bf16` / `pto.i16`: 128 lanes per physical vreg. + - `pto.i8` / `pto.ui8` / fp8: 256 lanes per physical vreg. +- `pto.vmi.mask(lanes)` creates a logical per-lane mask type. + Its lane count must match the gated vector. + +PTODSL does not expose layout selection on `pto.vmi.vreg(...)` or +`pto.vmi.mask(...)`; PTOAS infers layout during lowering. + +`dtype` may come from a `pto.const_expr` parameter. Prefer compile-time dtype +selection inside one PTODSL function over many small wrappers that only vary the +VMI element type: + +```python +vec_ty = pto.vmi.vreg(lanes, OUT_DTYPE) +vec = pto.vmi.vcvt(src, to_dtype=OUT_DTYPE) +``` + +## Load And Store + +```python +vec = pto.vmi.vload(ub_src, offset, size=128) +even, odd = pto.vmi.vload(ub_src, offset, size=128, dist_mode="dintlv") +wide = pto.vmi.vload(ub_src, offset, size=128, dist_mode="unpack", to_dtype=pto.f32) +pto.vmi.vstore(vec, ub_dst, offset, mask) +pto.vmi.vstore(vec, ub_dst, offset, group=8, stride=row_stride) +``` + +Use VMI load/store only for UB-resident compute operands/results. Use MTE or +tile operations for GM movement. + +Useful options: + +- `size` is required for every `vload`. +- `dist_mode=None` or `"continuous"` is the default contiguous load/store form. +- `dist_mode="dintlv"` returns an `(even, odd)` pair. +- `dist_mode="unpack", to_dtype=` widens by one adjacent bit-width step. +- `group=...`, `stride=...` select grouped access. +- `block_stride=...`, `repeat_stride=...` select block-strided access. +- `vload` does not take `mask`. +- `group` store does not take `mask`. +- `pmode="zero"` is the default masked-store behavior; `pmode="merge"` preserves + inactive lanes. + +Backend note: prefer putting dynamic tail masks on compute/store. Do not rely on +masked loads unless the current backend explicitly supports the form. + +## Masks + +```python +mask = pto.vmi.create_mask(active_lanes, size=lanes) +gmask = pto.vmi.create_mask(active_per_group, size=lanes, group=num_groups) +``` + +Use `create_mask` for prefix-active dynamic tails. Use grouped masks for grouped +reductions or grouped broadcast patterns. + +## Index And Broadcast + +```python +idx = pto.vmi.vci(pto.i32(0), size=64, order="ASC") +bc = pto.vmi.vbrc(pto.f16(0.0), size=128) +``` + +Use `vci` for lane-wise index ramps and gather/scatter offsets. Use `vbrc` for +scalar-to-vector or group-to-vector broadcast. + +## Elementwise And Scalar Ops + +Same-shape vector ops usually infer result type: + +```python +y = pto.vmi.vadd(a, b, mask) +y = pto.vmi.vsub(a, b, mask) +y = pto.vmi.vmul(a, b, mask) +y = pto.vmi.vdiv(a, b, mask) +y = pto.vmi.vmax(a, b, mask) +y = pto.vmi.vmin(a, b, mask) +``` + +Unary ops: + +```python +y = pto.vmi.vabs(x, mask) +y = pto.vmi.vneg(x, mask) +y = pto.vmi.vrelu(x, mask) +y = pto.vmi.vexp(x, mask) +y = pto.vmi.vln(x, mask) +y = pto.vmi.vsqrt(x, mask) +``` + +Vector-scalar ops require a mask: + +```python +y = pto.vmi.vadds(x, scalar, mask) +y = pto.vmi.vmuls(x, scalar, mask) +y = pto.vmi.vmaxs(x, scalar, mask) +y = pto.vmi.vmins(x, scalar, mask) +``` + +Integer/bitwise ops include `vand`, `vor`, `vxor`, `vnot`, `vshl`, `vshr`, +`vshls`, and `vshrs`. + +## Compare, Select, Reductions + +```python +cmp_mask = pto.vmi.vcmp(lhs, rhs, seed_mask, "lt") +cmp_mask = pto.vmi.vcmps(x, scalar, seed_mask, "ge") +out = pto.vmi.vsel(cmp_mask, true_value, false_value) +``` + +Reduction result types are inferred: + +```python +sum1 = pto.vmi.vcadd(x, mask, reassoc=True) +max1 = pto.vmi.vcmax(x, mask) +sum_g = pto.vmi.vcadd(x, gmask, group=num_groups, reassoc=True) +``` + +## Conversion And Reinterpretation + +```python +wide = pto.vmi.vcvt(x_f16, to_dtype=pto.f32) +narrow = pto.vmi.vcvt(x_f32, to_dtype=pto.f16, rounding="...", saturate=...) +bits = pto.vmi.vinterpret_cast(x, pto.i32) +``` + +Use `vcvt` for numeric conversion. Use `vinterpret_cast` only for bit-level +reinterpretation. + +## Gather, Scatter, Rearrangement + +```python +values = pto.vmi.vgather(src_ub, offsets, mask) +pto.vmi.vscatter(values, dst_ub, offsets, mask) +lo, hi = pto.vmi.vintlv(a, b, mask) +even, odd = pto.vmi.vdintlv(a, b, mask) +sel = pto.vmi.vselr(source, index) +``` + +Only use explicit VMI interleave/deinterleave when it changes the logical data +ordering. If the AscendC source uses interleave only to repair physical register +layout after widening or packing, collapse it into the logical VMI value. + +## Common AscendC SIMD To VMI Patterns + +- `vlds` from UB -> `pto.vmi.vload`. +- `vsts` to UB -> `pto.vmi.vstore`. +- `vcvt f16/bf16 -> f32` -> `pto.vmi.vcvt(..., to_dtype=pto.f32)`. +- `vcvt f32 -> f16/fp8` -> `pto.vmi.vcvt(..., to_dtype=)`. +- `vmul`, `vadd`, `vsub`, `vmax`, `vmin` -> corresponding VMI elementwise op. +- `vcmp` + `vsel` -> VMI compare mask plus `vsel`. +- `vintlv`/`vdintlv` trees, `PART_P*`, `PART_EVEN/ODD`, packed store modes: + usually physical-only lowering details; collapse unless they change the + logical result order. + +## Lane Selection Heuristic + +Choose the largest contiguous logical chunk that matches the algorithm and VMI +constraints: + +1. Keep dtype equal to the logical element type at that stage. +2. Prefer one full physical-register worth of elements when the algorithm and + UB layout are contiguous. +3. Use smaller multiples of 64 for row tails or natural row/group widths. +4. For dynamic remainders, use `create_mask(..., size=lanes[, group=...])` and + keep offsets symbolic. diff --git a/.codex/skills/rewrite-kernel-with-vmi/references/vmi-mlir-spec.md b/.codex/skills/rewrite-kernel-with-vmi/references/vmi-mlir-spec.md new file mode 100644 index 0000000000..0a6f3921fc --- /dev/null +++ b/.codex/skills/rewrite-kernel-with-vmi/references/vmi-mlir-spec.md @@ -0,0 +1,132 @@ +# VMI MLIR Semantics And Validation Notes + +This is a compact rewrite-focused summary of the larger VMI draft. Prefer +`ptodsl/docs/user_guide/14-vmi-virtual-instruction-set.md` for Python API +spelling and `docs/examples/vmi/vmi-design.md` for deeper design rationale. + +## Mental Model + +VMI is a logical SIMD IR: + +- A `!pto.vmi.vreg` is a flat logical vector of `L` lanes of element type + `T`. +- A `!pto.vmi.mask` gates those logical lanes. +- Physical register count, interleave, deinterleave, pack, part, and layout + materialization are lowering concerns handled by PTOAS. + +Do not translate AscendC physical register choreography directly if a single +logical VMI value has the same observable meaning. + +## Logical Vector Size + +Physical backing size: + +```text +K = ceil(lanes * bitwidth(dtype) / 2048) +``` + +One physical vector register is 256B / 2048 bits. Typical full-register lane +counts: + +- `f32/i32`: 64 lanes. +- `f16/bf16/i16`: 128 lanes. +- `i8/ui8/fp8`: 256 lanes. + +VMI also allows compact logical vectors such as `vreg(64, pto.f16)`. Unused +physical lanes are undefined and must be masked out when they may be observed. + +## Layout + +Common layouts: + +- Contiguous: logical lane `i` appears in logical order. This is the default. +- Deinterleaved: even/odd or grouped physical backing created by widening, + deinterleave loads, or layout-sensitive lowering. + +Author PTODSL without explicit layout by default. PTODSL does not expose +explicit layout selection on `pto.vmi.vreg(...)` or `pto.vmi.mask(...)`; use +`ensure_layout` only when debugging layout assignment or when the user asked +for layout-level IR. + +## Mask And Tail Semantics + +For dynamic tails: + +```python +mask = pto.vmi.create_mask(active_lanes, size=lanes) +``` + +Use the same `lanes` as the gated vector. Put the mask on compute ops and +stores. If the backend cannot predicate a load, load the full vector from a safe +UB range and rely on masked consumers/stores. For grouped tails, use +`create_mask(..., size=lanes, group=...)`. + +Group masks represent active elements inside each group and should be used for +group reductions or grouped broadcasts rather than plain prefix tails. + +## Semantic Equivalence Rules + +Treat these AscendC details as physical-only unless they change observable +logical order or data value: + +- `PART_P0`, `PART_P1`, `PART_P2`, `PART_P3`. +- `PART_EVEN`, `PART_ODD`. +- Pack/store modes such as `PK4_B32`. +- `vintlv` / `vdintlv` trees that merely repair physical layout. +- Separate low/high or even/odd temporary registers produced by hardware + widening/narrowing. + +Preserve these details when they are semantic: + +- Real element permutation or transposition. +- Gather/scatter offsets. +- Row/column order changes. +- Reduction grouping. +- Saturation, rounding, sign/zero extension mode. +- Tail behavior that changes which output elements are written. + +## Rewrite Pattern + +Preferred surface VMI shape: + +```python +mask = pto.vmi.create_mask(active, size=lanes) +x = pto.vmi.vload(x_ub, x_off, size=lanes) +y = pto.vmi.vcvt(x, to_dtype=compute_dtype) # if needed +z = pto.vmi.vmul(y, scale, mask) # representative compute +out = pto.vmi.vcvt(z, to_dtype=dst_dtype) # if needed +pto.vmi.vstore(out, y_ub, y_off, mask) +``` + +Keep offsets symbolic and aligned with the source. Prefer one logical vector per +semantic row/block iteration. Split only when the algorithm, UB layout, or VMI +constraints require it. + +## PTOAS Validation + +For MLIR emitted by PTODSL, validate through the VMI backend path: + +```bash +ptoas --pto-arch=a5 --pto-backend=vpto --enable-vmi -o /dev/null +``` + +Optional debugging forms: + +```bash +ptoas --pto-arch=a5 --pto-backend=vpto --enable-vmi --emit-vpto -o - +ptoas --pto-arch=a5 --pto-backend=vpto --enable-vmi --emit-pto-ir -o - +``` + +Do not diagnose VMI legality from a non-VMI invocation. `--enable-vmi` requires +`--pto-backend=vpto` unless the input module already declares the VPTO backend. + +## Debugging Priorities + +1. PTODSL trace error: fix Python DSL syntax, types, missing required kwargs, + or unsupported wrapper usage. +2. VMI verifier error: fix lane/mask mismatch, dtype mismatch, invalid + inferred shape, invalid group count, or unsupported conversion. +3. VMI lowering error: check whether the logical form needs a different legal + VMI expression. Avoid adding store/reload workarounds without user approval. +4. Semantic mismatch: revisit the VMI design gate; most errors come from + accidentally preserving or dropping a physical packing/interleave detail. diff --git a/.codex/skills/rewrite-kernel-with-vmi/scripts/compile_vmi_to_vpto.sh b/.codex/skills/rewrite-kernel-with-vmi/scripts/compile_vmi_to_vpto.sh new file mode 100755 index 0000000000..8a4630c7d5 --- /dev/null +++ b/.codex/skills/rewrite-kernel-with-vmi/scripts/compile_vmi_to_vpto.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)" +PTOAS_BIN="${PTOAS_BIN:-${REPO_ROOT}/install/bin/ptoas}" + +if [[ $# -lt 1 || $# -gt 2 ]]; then + echo "Usage: $0 [output.mi.pto]" >&2 + exit 1 +fi + +INPUT="$1" +if [[ ! -f "${INPUT}" ]]; then + echo "Input file not found: ${INPUT}" >&2 + exit 1 +fi + +if [[ ! -x "${PTOAS_BIN}" ]]; then + echo "ptoas not found or not executable: ${PTOAS_BIN}" >&2 + exit 1 +fi + +if [[ $# -eq 2 ]]; then + OUTPUT="$2" +else + if [[ "${INPUT}" == *.vmi.pto ]]; then + OUTPUT="${INPUT%.vmi.pto}.mi.pto" + else + OUTPUT="${INPUT%.pto}.mi.pto" + fi +fi + +echo "Compiling ${INPUT} -> ${OUTPUT}" +"${PTOAS_BIN}" \ + --pto-arch=a5 \ + --pto-backend=vpto \ + --enable-vmi \ + --pto-level=level3 \ + --emit-vpto \ + -o "${OUTPUT}" \ + "${INPUT}" + +echo "Done: ${OUTPUT}" diff --git a/.github/scripts/compute_ptoas_version.py b/.github/scripts/compute_ptoas_version.py index 77c578a42e..518ef19e1e 100644 --- a/.github/scripts/compute_ptoas_version.py +++ b/.github/scripts/compute_ptoas_version.py @@ -30,7 +30,7 @@ def parse_args() -> argparse.Namespace: ) parser.add_argument( "--check-tag", - help="Optional release tag to validate, e.g. v0.8 or 0.8.", + help="Optional release tag to validate, e.g. ptoas-v0.8 or v0.8.", ) return parser.parse_args() @@ -45,7 +45,12 @@ def read_base_version(cmake_file: pathlib.Path) -> str: return match.group(1) def normalize_tag(tag: str) -> str: - return tag[1:] if tag.startswith("v") else tag + normalized = tag.strip() + if normalized.startswith("ptoas-"): + normalized = normalized[len("ptoas-"):] + if normalized.startswith("v"): + normalized = normalized[1:] + return normalized def main() -> int: diff --git a/.github/scripts/compute_vmi_version.py b/.github/scripts/compute_vmi_version.py new file mode 100644 index 0000000000..3a9e23b9c5 --- /dev/null +++ b/.github/scripts/compute_vmi_version.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +import argparse +import pathlib +import re +import sys + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Compute the VMI release version from a version file." + ) + parser.add_argument( + "--version-file", + default="docs/release/VMI_VERSION", + help="Path to the VMI version file.", + ) + parser.add_argument( + "--check-tag", + help="Optional release tag to validate, e.g. vmi-v0.1.0.", + ) + return parser.parse_args() + + +def read_version(version_file: pathlib.Path) -> str: + version = version_file.read_text(encoding="utf-8").strip() + if not re.fullmatch(r"[0-9]+\.[0-9]+\.[0-9]+", version): + raise ValueError(f"invalid VMI version '{version}' in {version_file}") + return version + + +def normalize_tag(tag: str) -> str: + normalized = tag.strip() + if normalized.startswith("vmi-"): + normalized = normalized[len("vmi-"):] + if normalized.startswith("v"): + normalized = normalized[1:] + return normalized + + +def main() -> int: + args = parse_args() + version_file = pathlib.Path(args.version_file) + version = read_version(version_file) + + if args.check_tag is not None: + normalized_tag = normalize_tag(args.check_tag) + if normalized_tag != version: + print( + f"release tag '{args.check_tag}' does not match computed version '{version}'", + file=sys.stderr, + ) + return 1 + + print(version) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/.github/scripts/update_ptoas_base_version.py b/.github/scripts/update_ptoas_base_version.py index 692f9e2753..e79ddafecc 100644 --- a/.github/scripts/update_ptoas_base_version.py +++ b/.github/scripts/update_ptoas_base_version.py @@ -17,7 +17,7 @@ PROJECT_VERSION_RE = re.compile( r"(project\s*\(\s*ptoas\s+VERSION\s+)([0-9]+\.[0-9]+)(\s*\))" ) -TAG_VERSION_RE = re.compile(r"v?([0-9]+)\.([0-9]+)") +TAG_VERSION_RE = re.compile(r"^(?:ptoas-)?v?([0-9]+)\.([0-9]+)$") def parse_args() -> argparse.Namespace: @@ -54,6 +54,8 @@ def parse_args() -> argparse.Namespace: def normalize_version(version: str) -> str: normalized = version.strip() + if normalized.startswith("ptoas-"): + normalized = normalized[len("ptoas-"):] if normalized.startswith("v"): normalized = normalized[1:] if not re.fullmatch(r"[0-9]+\.[0-9]+", normalized): diff --git a/.github/workflows/build_wheel.yml b/.github/workflows/build_wheel.yml index f006bd2881..25dfb054d4 100644 --- a/.github/workflows/build_wheel.yml +++ b/.github/workflows/build_wheel.yml @@ -30,6 +30,7 @@ env: jobs: build_wheel: + if: github.event_name != 'release' || startsWith(github.ref_name, 'ptoas-v') || startsWith(github.ref_name, 'vmi-v') # Non-release runs keep a single-Python matrix for faster gates; # release/nightly runs fan out to publish the full wheel set plus binary artifacts. name: Build ptoas-bin (${{ matrix.arch }}, py${{ matrix.python }}) @@ -62,8 +63,22 @@ jobs: - name: Compute PTOAS CLI version run: | if [ "${GITHUB_EVENT_NAME}" = "release" ]; then - EXPECTED_RELEASE_VERSION="$(${PY_PATH}/bin/python .github/scripts/compute_ptoas_version.py --mode release)" - PTOAS_VERSION="${GITHUB_REF_NAME#v}" + case "${GITHUB_REF_NAME}" in + ptoas-v*) + EXPECTED_RELEASE_VERSION="$(${PY_PATH}/bin/python .github/scripts/compute_ptoas_version.py --mode release)" + PTOAS_VERSION="${GITHUB_REF_NAME#ptoas-v}" + ;; + vmi-v*) + EXPECTED_RELEASE_VERSION="$(${PY_PATH}/bin/python .github/scripts/compute_vmi_version.py \ + --version-file docs/release/VMI_VERSION \ + --check-tag "${GITHUB_REF_NAME}")" + PTOAS_VERSION="${GITHUB_REF_NAME#vmi-v}" + ;; + *) + echo "release tags for the wheel workflow must start with 'ptoas-v' or 'vmi-v'" >&2 + exit 1 + ;; + esac if [ "${PTOAS_VERSION}" != "${EXPECTED_RELEASE_VERSION}" ]; then echo "release tag '${GITHUB_REF_NAME}' does not match computed version '${EXPECTED_RELEASE_VERSION}'" >&2 exit 1 @@ -180,6 +195,40 @@ jobs: export PTOAS_PYTHON_PACKAGE_VERSION="${PTOAS_VERSION}" bash $PTO_SOURCE_DIR/docker/create_wheel.sh + - name: Validate wheel payload and launcher contract + if: github.event_name == 'release' || github.event_name == 'schedule' + run: | + export PATH="${PY_PATH}/bin:$PATH" + python - <<'PY' + import os + import zipfile + from pathlib import Path + + wheel_dist = Path(os.environ["PTO_SOURCE_DIR"]) / "build" / "wheel-dist" + wheels = sorted(wheel_dist.glob("ptoas-*.whl")) + if len(wheels) != 1: + raise SystemExit(f"expected exactly one wheel in {wheel_dist}, found {len(wheels)}") + + with zipfile.ZipFile(wheels[0]) as zf: + names = set(zf.namelist()) + required = {"ptoas/__init__.py", "ptoas/_launcher.py", "pto/ptoas.so"} + missing = [name for name in required if name not in names] + if missing: + raise SystemExit(f"wheel is missing required payload files: {missing}") + if "ptoas/_runtime/bin/ptoas" in names: + raise SystemExit("wheel unexpectedly contains ptoas/_runtime/bin/ptoas") + + entry_points_name = next( + name for name in names + if name.startswith("ptoas-") and name.endswith(".dist-info/entry_points.txt") + ) + entry_points = zf.read(entry_points_name).decode("utf-8") + if "ptoas=ptoas._launcher:main" not in entry_points: + raise SystemExit("wheel entry points do not route ptoas through ptoas._launcher:main") + + print(f"validated wheel payload and launcher contract: {wheels[0].name}") + PY + - name: Repair wheel with auditwheel if: github.event_name == 'release' || github.event_name == 'schedule' run: | @@ -230,7 +279,7 @@ jobs: upload_release_assets: name: Upload release assets - if: github.event_name == 'release' || github.event_name == 'schedule' + if: (github.event_name == 'release' && (startsWith(github.ref_name, 'ptoas-v') || startsWith(github.ref_name, 'vmi-v'))) || github.event_name == 'schedule' needs: build_wheel runs-on: ubuntu-latest @@ -290,7 +339,19 @@ jobs: mv "release-artifacts/ptoas-bin-aarch64/ptoas-bin-aarch64.tar.gz" \ "release-artifacts/ptoas-bin-aarch64.tar.gz" - - name: Upload assets to GitHub Release + - name: Upload wheel assets to GitHub Release + uses: softprops/action-gh-release@v2 + with: + tag_name: ${{ env.RELEASE_TAG }} + name: ${{ env.RELEASE_NAME }} + target_commitish: ${{ github.sha }} + prerelease: ${{ env.RELEASE_PRERELEASE }} + make_latest: ${{ env.RELEASE_MAKE_LATEST }} + body: ${{ env.RELEASE_BODY }} + overwrite_files: true + files: release-artifacts/*.whl + + - name: Upload binary assets to GitHub Release uses: softprops/action-gh-release@v2 with: tag_name: ${{ env.RELEASE_TAG }} @@ -300,13 +361,11 @@ jobs: make_latest: ${{ env.RELEASE_MAKE_LATEST }} body: ${{ env.RELEASE_BODY }} overwrite_files: true - files: | - release-artifacts/*.whl - release-artifacts/*.tar.gz + files: release-artifacts/*.tar.gz bump_base_version: name: Bump base version after release - if: github.event_name == 'release' && github.event.action == 'released' + if: github.event_name == 'release' && github.event.action == 'released' && startsWith(github.ref_name, 'ptoas-v') needs: upload_release_assets runs-on: ubuntu-latest diff --git a/.github/workflows/build_wheel_mac.yml b/.github/workflows/build_wheel_mac.yml index 48afb41e48..3c7653e49e 100644 --- a/.github/workflows/build_wheel_mac.yml +++ b/.github/workflows/build_wheel_mac.yml @@ -29,6 +29,7 @@ env: jobs: build_wheel: + if: github.event_name != 'release' || startsWith(github.ref_name, 'ptoas-v') || startsWith(github.ref_name, 'vmi-v') # Non-release runs keep a single-Python matrix for faster gates; # release/nightly runs fan out to publish the full wheel set plus binary artifacts. name: Build ptoas-bin (macOS ${{ matrix.arch }}, py${{ matrix.python }}) @@ -67,8 +68,22 @@ jobs: - name: Compute PTOAS CLI version run: | if [ "${GITHUB_EVENT_NAME}" = "release" ]; then - EXPECTED_RELEASE_VERSION="$(python .github/scripts/compute_ptoas_version.py --mode release)" - PTOAS_VERSION="${GITHUB_REF_NAME#v}" + case "${GITHUB_REF_NAME}" in + ptoas-v*) + EXPECTED_RELEASE_VERSION="$(python .github/scripts/compute_ptoas_version.py --mode release)" + PTOAS_VERSION="${GITHUB_REF_NAME#ptoas-v}" + ;; + vmi-v*) + EXPECTED_RELEASE_VERSION="$(python .github/scripts/compute_vmi_version.py \ + --version-file docs/release/VMI_VERSION \ + --check-tag "${GITHUB_REF_NAME}")" + PTOAS_VERSION="${GITHUB_REF_NAME#vmi-v}" + ;; + *) + echo "release tags for the wheel workflow must start with 'ptoas-v' or 'vmi-v'" >&2 + exit 1 + ;; + esac if [ "${PTOAS_VERSION}" != "${EXPECTED_RELEASE_VERSION}" ]; then echo "release tag '${GITHUB_REF_NAME}' does not match computed version '${EXPECTED_RELEASE_VERSION}'" >&2 exit 1 @@ -186,6 +201,39 @@ jobs: printf 'Built wheel file: %s\n' "$(basename "${built_wheels[0]}")" fi + - name: Validate wheel payload and launcher contract + if: github.event_name == 'release' || github.event_name == 'schedule' + run: | + python - <<'PY' + import os + import zipfile + from pathlib import Path + + wheel_dist = Path(os.environ["PTO_SOURCE_DIR"]) / "build" / "wheel-dist" + wheels = sorted(wheel_dist.glob("ptoas-*.whl")) + if len(wheels) != 1: + raise SystemExit(f"expected exactly one wheel in {wheel_dist}, found {len(wheels)}") + + with zipfile.ZipFile(wheels[0]) as zf: + names = set(zf.namelist()) + required = {"ptoas/__init__.py", "ptoas/_launcher.py", "pto/ptoas.so"} + missing = [name for name in required if name not in names] + if missing: + raise SystemExit(f"wheel is missing required payload files: {missing}") + if "ptoas/_runtime/bin/ptoas" in names: + raise SystemExit("wheel unexpectedly contains ptoas/_runtime/bin/ptoas") + + entry_points_name = next( + name for name in names + if name.startswith("ptoas-") and name.endswith(".dist-info/entry_points.txt") + ) + entry_points = zf.read(entry_points_name).decode("utf-8") + if "ptoas=ptoas._launcher:main" not in entry_points: + raise SystemExit("wheel entry points do not route ptoas through ptoas._launcher:main") + + print(f"validated wheel payload and launcher contract: {wheels[0].name}") + PY + - name: Repair wheel with delocate if: github.event_name == 'release' || github.event_name == 'schedule' run: | @@ -353,7 +401,7 @@ jobs: upload_release_assets: name: Upload release assets - if: github.event_name == 'release' || github.event_name == 'schedule' + if: (github.event_name == 'release' && (startsWith(github.ref_name, 'ptoas-v') || startsWith(github.ref_name, 'vmi-v'))) || github.event_name == 'schedule' needs: build_wheel runs-on: ubuntu-latest diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2ca29fd57d..e4158464c7 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -325,7 +325,7 @@ jobs: env: CI_EVENT_NAME: ${{ github.event_name }} WORKFLOW_SOC_VERSION: ${{ github.event.inputs.soc_version || 'Ascend910' }} - PTOAS_BIN: ${{ env.PTO_BUILD_DIR }}/tools/ptoas/ptoas + PTOAS_BIN: ${{ env.PTO_INSTALL_DIR }}/bin/ptoas PTOBC_BIN: ${{ env.PTO_BUILD_DIR }}/tools/ptobc/ptobc PTO_BUILD_DIR: ${{ env.PTO_BUILD_DIR }} PYTHON_BIN: ${{ env.PTOAS_VENV }}/bin/python diff --git a/.github/workflows/ci_sim.yml b/.github/workflows/ci_sim.yml index e4b083ecd5..a0c8ebbd87 100644 --- a/.github/workflows/ci_sim.yml +++ b/.github/workflows/ci_sim.yml @@ -154,7 +154,7 @@ jobs: fi echo "ASCEND_HOME_PATH=${ASCEND_HOME_PATH_DETECTED}" >> "${GITHUB_ENV}" - echo "PTOAS_BIN=${GITHUB_WORKSPACE}/build/tools/ptoas/ptoas" >> "${GITHUB_ENV}" + echo "PTOAS_BIN=${GITHUB_WORKSPACE}/install/bin/ptoas" >> "${GITHUB_ENV}" - name: Ensure runner dependencies shell: bash diff --git a/CMakeLists.txt b/CMakeLists.txt index cadb9dd221..1c1302fbe7 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -261,6 +261,9 @@ add_subdirectory(tools) # ========================================================= if(BUILD_TESTING) enable_testing() + if(NOT TARGET ptoas_runtime_deps) + add_custom_target(ptoas_runtime_deps) + endif() if(PTO_ENABLE_PYTHON_BINDING) get_filename_component(_llvm_root "${LLVM_DIR}/../../.." ABSOLUTE) set(_pto_python_test_pythonpath @@ -315,7 +318,7 @@ if(BUILD_TESTING) add_subdirectory(test/lit) set(_pto_check_ctest_depends) - foreach(_target IN ITEMS PTOPythonModules pto-opt ptoas ptobc) + foreach(_target IN ITEMS PTOPythonModules ptoas_runtime_deps ptobc) if(TARGET ${_target}) list(APPEND _pto_check_ctest_depends ${_target}) endif() diff --git a/README.md b/README.md index 49b22df672..b988a7d66f 100644 --- a/README.md +++ b/README.md @@ -202,6 +202,7 @@ from mlir.dialects import pto as mlir_pto > - `ptoas-bin-*.tar.gz` 这类 compiler-only 二进制 tarball 只提供 CLI/toolchain， > **不是** PTODSL-capable Python distribution；仅解压 tarball 不能保证 > `import ptodsl` 可用。 +> - release tag 约定：`ptoas-vX.Y` 发布主工具链，`vmi-vA.B.C` 发布 VMI 文档/规范。 --- @@ -255,6 +256,11 @@ ptoas test/lit/pto/empty_func.pto --pto-arch=a5 -o outputfile.cpp # 指定构建 Level（level3 会禁用 PlanMemory/InsertSync） ptoas test/lit/pto/empty_func.pto --pto-level=level3 -o outputfile.cpp +# VPTO backend 默认启用 VMI -> VPTO 语义 pipeline +# 可使用 --enable-vmi=false 临时关闭 +# public function signature 不能直接暴露 !pto.vmi.* 类型 +ptoas test/lit/vmi/vmi_ptoas_cli_pipeline.pto --pto-arch=a5 --pto-backend=vpto --emit-vpto -o - + # 查看当前 ptoas release 版本号 ptoas --version diff --git a/ReleaseNotes.md b/ReleaseNotes.md index 3beaf43603..79882f63e0 100644 --- a/ReleaseNotes.md +++ b/ReleaseNotes.md @@ -1,7 +1,7 @@ # PTOAS (PTO Assembler & Optimizer) ## 版本 -- 版本号：v0.1.0 +- 版本号：v0.51 - 发布日期：2026-02-14 ## 变更摘要 diff --git a/_ptoas_build_backend.py b/_ptoas_build_backend.py index d30ded39cd..9c500b5b52 100644 --- a/_ptoas_build_backend.py +++ b/_ptoas_build_backend.py @@ -28,6 +28,7 @@ import hashlib import io import os +import re import shutil import subprocess import sys @@ -48,6 +49,9 @@ _LLVM_BUILD_DIR / "tools" / "mlir" / "python_packages" / "mlir_core" ) _WHEEL_DIST_DIR = _BUILD_DIR / "wheel-dist" +_PROJECT_VERSION_RE = re.compile( + r"project\s*\(\s*ptoas\s+VERSION\s+([0-9]+\.[0-9]+)\s*\)" +) def _assert_installed_ptodsl_payload() -> None: @@ -63,6 +67,17 @@ def _assert_installed_ptodsl_payload() -> None: ) +def _assert_installed_ptoas_shared_module() -> None: + installed_shared_module = _PTO_INSTALL_DIR / "lib" / "ptoas.so" + if installed_shared_module.exists(): + return + raise RuntimeError( + "PTOAS shared launcher module is missing from the PTOAS install tree. " + f"Expected to find {installed_shared_module}. " + "Wheel assembly now packages the shared module from the install tree." + ) + + def _assert_editable_ptodsl_source() -> None: """Fail fast if the editable install cannot point at the PTODSL source.""" source_init = _PTODSL_SOURCE_ROOT / "ptodsl" / "__init__.py" @@ -75,6 +90,17 @@ def _assert_editable_ptodsl_source() -> None: ) +def _default_ptoas_version() -> str: + version = os.environ.get("PTOAS_PYTHON_PACKAGE_VERSION", "").strip() + if version: + return version + cmake_file = _REPO / "CMakeLists.txt" + match = _PROJECT_VERSION_RE.search(cmake_file.read_text(encoding="utf-8")) + if not match: + raise RuntimeError(f"could not find PTOAS version in {cmake_file}") + return match.group(1) + + def get_requires_for_build_wheel(config_settings=None): return ["setuptools>=68", "wheel", "pybind11<3"] @@ -91,7 +117,7 @@ def prepare_metadata_for_build_wheel(metadata_directory, config_settings=None): """Return wheel metadata without running the full build.""" import email.message - version = os.environ.get("PTOAS_PYTHON_PACKAGE_VERSION", "0.1.0") + version = _default_ptoas_version() dist_info = Path(metadata_directory) / f"ptoas-{version}.dist-info" dist_info.mkdir(parents=True, exist_ok=True) @@ -165,6 +191,7 @@ def _cmake_configure_and_build(): ["cmake", "--build", str(_BUILD_DIR), "--target", "install"] ) _assert_installed_ptodsl_payload() + _assert_installed_ptoas_shared_module() def build_wheel(wheel_directory, config_settings=None, metadata_directory=None): @@ -174,6 +201,7 @@ def build_wheel(wheel_directory, config_settings=None, metadata_directory=None): env.update({ "PTO_SOURCE_DIR": str(_REPO), "PTO_INSTALL_DIR": str(_PTO_INSTALL_DIR), + "PTO_BUILD_DIR": str(_BUILD_DIR), "LLVM_BUILD_DIR": str(_LLVM_BUILD_DIR), "PTO_WHEEL_DIST_DIR": str(_WHEEL_DIST_DIR), # Keep wheel packaging on the same interpreter pip used to invoke the @@ -212,7 +240,7 @@ def build_editable(wheel_directory, config_settings=None, metadata_directory=Non _cmake_configure_and_build() _assert_editable_ptodsl_source() - version = os.environ.get("PTOAS_PYTHON_PACKAGE_VERSION", "0.1.0") + version = _default_ptoas_version() # Paths that must be on sys.path for the package to be importable pth_paths = [ diff --git a/docker/collect_ptoas_dist.sh b/docker/collect_ptoas_dist.sh index 75fb436558..bb55336ec5 100755 --- a/docker/collect_ptoas_dist.sh +++ b/docker/collect_ptoas_dist.sh @@ -20,7 +20,8 @@ # Output structure: # / # ptoas - Wrapper script that sets up LD_LIBRARY_PATH -# bin/ptoas - The actual ptoas binary +# bin/ptoas - Python wrapper entrypoint +# python/ptoas/ - Launcher package used by the wrapper # lib/*.so* - Required shared library dependencies # share/ptoas/TileOps - TileLang template library # tilelang_dsl/ - TileLang DSL Python package @@ -49,14 +50,22 @@ export LD_LIBRARY_PATH="${LLVM_BUILD_DIR}/lib:${PTO_INSTALL_DIR}/lib:${LD_LIBRAR PTO_BUILD_DIR="${PTO_BUILD_DIR:-${PTO_SOURCE_DIR}/build}" PTOAS_BIN="${PTO_BUILD_DIR}/tools/ptoas/ptoas" +PTOAS_SHARED_MODULE="${PTO_INSTALL_DIR}/lib/ptoas.so" PTOAS_DEPS_DIR="${PTOAS_DIST_DIR}/lib" PTOAS_TILEOPS_SRC_DIR="${PTO_INSTALL_DIR}/share/ptoas/TileOps" PTOAS_TILEOPS_DIST_DIR="${PTOAS_DIST_DIR}/share/ptoas/TileOps" PTOAS_TILELANG_DSL_SRC_DIR="${PTO_INSTALL_DIR}/tilelang_dsl" PTOAS_TILELANG_DSL_DIST_DIR="${PTOAS_DIST_DIR}/tilelang_dsl" +PTOAS_WRAPPER_PKG_SRC_DIR="${PTO_INSTALL_DIR}/ptoas" +PTOAS_PYTHON_ROOT_DIST_DIR="${PTOAS_DIST_DIR}/python" +PTOAS_WRAPPER_PKG_DIST_DIR="${PTOAS_PYTHON_ROOT_DIST_DIR}/ptoas" if [ ! -f "$PTOAS_BIN" ]; then - echo "Error: ptoas binary not found at $PTOAS_BIN" >&2 + echo "Error: ptoas wrapper not found at $PTOAS_BIN" >&2 + exit 1 +fi +if [ ! -f "$PTOAS_SHARED_MODULE" ]; then + echo "Error: shared launcher module not found at $PTOAS_SHARED_MODULE" >&2 exit 1 fi @@ -65,6 +74,10 @@ remove_rpath() { if ! has_rpath "$path"; then return fi + if ! can_scrub_rpath; then + echo "WARN: skipping RPATH/RUNPATH scrub for ${path}; install patchelf or chrpath to harden local dist artifacts" >&2 + return + fi if command -v patchelf >/dev/null 2>&1; then patchelf --remove-rpath "$path" fi @@ -93,6 +106,10 @@ has_rpath() { readelf -d "$path" 2>/dev/null | grep -Eq '(RPATH|RUNPATH)' } +can_scrub_rpath() { + command -v patchelf >/dev/null 2>&1 || command -v chrpath >/dev/null 2>&1 +} + assert_relro() { local path="$1" if ! readelf -l "$path" 2>/dev/null | grep -q 'GNU_RELRO'; then @@ -114,6 +131,9 @@ assert_no_symtab() { assert_no_rpath() { local path="$1" + if ! can_scrub_rpath; then + return + fi if has_rpath "$path"; then echo "Error: runtime search path still present in ${path}" >&2 exit 1 @@ -133,15 +153,36 @@ harden_elf() { mkdir -p \ "${PTOAS_DIST_DIR}/bin" \ "${PTOAS_DEPS_DIR}" \ + "${PTOAS_PYTHON_ROOT_DIST_DIR}" \ "$(dirname "${PTOAS_TILEOPS_DIST_DIR}")" +rm -rf "${PTOAS_WRAPPER_PKG_DIST_DIR}" +cp -R "${PTOAS_WRAPPER_PKG_SRC_DIR}" "${PTOAS_WRAPPER_PKG_DIST_DIR}" # Copy ptoas binary -echo "Copying ptoas binary..." +echo "Copying ptoas wrapper..." cp "$PTOAS_BIN" "${PTOAS_DIST_DIR}/bin/" -harden_elf "${PTOAS_DIST_DIR}/bin/ptoas" +chmod +x "${PTOAS_DIST_DIR}/bin/ptoas" -# Collect *.so dependencies (transitive closure under /llvm-workspace) +# Collect non-system *.so dependencies needed by the packaged shared runtime. echo "Collecting shared library dependencies..." +linux_runtime_dep_paths() { + local path="$1" + ldd "$path" 2>/dev/null | awk ' + /=> \// { print $3 } + /^\// { print $1 } + ' +} + +should_bundle_linux_dep() { + local path="$1" + case "$path" in + /lib/*|/lib64/*|/usr/lib/*|/usr/lib64/*) + return 1 + ;; + esac + return 0 +} + copy_so() { local f="$1" [[ -f "$f" ]] || return 0 @@ -151,17 +192,26 @@ copy_so() { cp -L -n "$f" "${PTOAS_DEPS_DIR}/" 2>/dev/null || true harden_elf "${PTOAS_DEPS_DIR}/${name}" while read -r res; do + [[ -n "$res" ]] || continue + should_bundle_linux_dep "$res" || continue copy_so "$res" - done < <(ldd "$f" 2>/dev/null | awk '/=> \/llvm-workspace\// {print $3}') + done < <(linux_runtime_dep_paths "$f") } while read -r res; do + [[ -n "$res" ]] || continue + should_bundle_linux_dep "$res" || continue + copy_so "$res" +done < <(linux_runtime_dep_paths "$PTOAS_BIN") +while read -r res; do + [[ -n "$res" ]] || continue + should_bundle_linux_dep "$res" || continue copy_so "$res" -done < <(ldd "$PTOAS_BIN" 2>/dev/null | awk '/=> \/llvm-workspace\// {print $3}') +done < <(linux_runtime_dep_paths "$PTOAS_SHARED_MODULE") while read -r packaged; do harden_elf "$packaged" -done < <(find "${PTOAS_DIST_DIR}/bin" "${PTOAS_DEPS_DIR}" -type f | sort) +done < <(find "${PTOAS_DEPS_DIR}" -type f | sort) echo "Copying TileLang runtime resources..." if [ ! -d "${PTOAS_TILEOPS_SRC_DIR}" ]; then @@ -182,6 +232,7 @@ cat > "${PTOAS_DIST_DIR}/ptoas" << 'WRAPPER_EOF' #!/bin/bash SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" export LD_LIBRARY_PATH="${SCRIPT_DIR}/lib:${LD_LIBRARY_PATH}" +export PTOAS_PYTHON_ROOT="${SCRIPT_DIR}/python" exec "${SCRIPT_DIR}/bin/ptoas" "$@" WRAPPER_EOF chmod +x "${PTOAS_DIST_DIR}/ptoas" @@ -212,6 +263,7 @@ echo "" echo "=== ptoas distribution contents ===" ls -la "${PTOAS_DIST_DIR}/" ls -la "${PTOAS_DIST_DIR}/bin/" +ls -la "${PTOAS_DIST_DIR}/python/" ls -la "${PTOAS_DIST_DIR}/share/ptoas/" ls -la "${PTOAS_TILELANG_DSL_DIST_DIR}" SO_COUNT=$(find "${PTOAS_DEPS_DIR}" -name "*.so*" 2>/dev/null | wc -l) diff --git a/docker/copy_ptoas_deps.sh b/docker/copy_ptoas_deps.sh index 26a501e64e..e9f14e5039 100644 --- a/docker/copy_ptoas_deps.sh +++ b/docker/copy_ptoas_deps.sh @@ -7,7 +7,7 @@ # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. -# Collect only *.so actually needed by ptoas (transitive closure under /llvm-workspace). +# Collect only non-system *.so actually needed by ptoas. # Expects: LLVM_BUILD_DIR, PTO_INSTALL_DIR, PTOAS_DEPS_DIR, PTO_SOURCE_DIR # Optional: PTO_BUILD_DIR (defaults to PTO_SOURCE_DIR/build) @@ -22,6 +22,10 @@ remove_rpath() { if ! has_rpath "$path"; then return fi + if ! can_scrub_rpath; then + echo "WARN: skipping RPATH/RUNPATH scrub for ${path}; install patchelf or chrpath to harden local dist artifacts" >&2 + return + fi if command -v patchelf >/dev/null 2>&1; then patchelf --remove-rpath "$path" fi @@ -50,6 +54,10 @@ has_rpath() { readelf -d "$path" 2>/dev/null | grep -Eq '(RPATH|RUNPATH)' } +can_scrub_rpath() { + command -v patchelf >/dev/null 2>&1 || command -v chrpath >/dev/null 2>&1 +} + assert_relro() { local path="$1" if ! readelf -l "$path" 2>/dev/null | grep -q 'GNU_RELRO'; then @@ -71,6 +79,9 @@ assert_no_symtab() { assert_no_rpath() { local path="$1" + if ! can_scrub_rpath; then + return + fi if has_rpath "$path"; then echo "Error: runtime search path still present in ${path}" >&2 exit 1 @@ -95,14 +106,36 @@ copy_so() { cp -L -n "$f" "${PTOAS_DEPS_DIR}/" 2>/dev/null || true harden_elf "${PTOAS_DEPS_DIR}/${name}" while read -r res; do + [[ -n "$res" ]] || continue + should_bundle_linux_dep "$res" || continue copy_so "$res" - done < <(ldd "$f" 2>/dev/null | awk '/=> \/llvm-workspace\// {print $3}') + done < <(linux_runtime_dep_paths "$f") +} + +linux_runtime_dep_paths() { + local path="$1" + ldd "$path" 2>/dev/null | awk ' + /=> \// { print $3 } + /^\// { print $1 } + ' +} + +should_bundle_linux_dep() { + local path="$1" + case "$path" in + /lib/*|/lib64/*|/usr/lib/*|/usr/lib64/*) + return 1 + ;; + esac + return 0 } mkdir -p "$PTOAS_DEPS_DIR" while read -r res; do + [[ -n "$res" ]] || continue + should_bundle_linux_dep "$res" || continue copy_so "$res" -done < <(ldd "$PTOAS_BIN" 2>/dev/null | awk '/=> \/llvm-workspace\// {print $3}') +done < <(linux_runtime_dep_paths "$PTOAS_BIN") while read -r packaged; do harden_elf "$packaged" diff --git a/docker/create_wheel.sh b/docker/create_wheel.sh index 22b27776fc..d022331d91 100755 --- a/docker/create_wheel.sh +++ b/docker/create_wheel.sh @@ -34,6 +34,14 @@ export PTOAS_PYTHON_PACKAGE_VERSION linux_runtime_dep_paths() { local path="$1" + local library_path="${2:-}" + if [[ -n "${library_path}" ]]; then + env LD_LIBRARY_PATH="${library_path}" ldd "$path" 2>/dev/null | awk ' + /=> \// { print $3 } + /^\// { print $1 } + ' + return + fi ldd "$path" 2>/dev/null | awk ' /=> \// { print $3 } /^\// { print $1 } @@ -51,12 +59,12 @@ should_bundle_linux_dep() { } assemble_linux_wheel_runtime() { - local ptoas_bin="${PTO_BUILD_DIR}/tools/ptoas/ptoas" - if [[ ! -f "${ptoas_bin}" ]]; then - ptoas_bin="${PTO_INSTALL_DIR}/bin/ptoas" + local ptoas_wrapper="${PTO_INSTALL_DIR}/bin/ptoas" + if [[ ! -f "${ptoas_wrapper}" ]]; then + ptoas_wrapper="${PTO_BUILD_DIR}/tools/ptoas/ptoas" fi - if [[ ! -f "${ptoas_bin}" ]]; then - echo "Error: ptoas binary not found in build tree or install tree" >&2 + if [[ ! -f "${ptoas_wrapper}" ]]; then + echo "Error: ptoas wrapper not found in build tree or install tree" >&2 exit 1 fi if [[ ! -d "${PTO_INSTALL_DIR}/share/ptoas/TileOps" ]]; then @@ -64,21 +72,25 @@ assemble_linux_wheel_runtime() { exit 1 fi - mkdir -p "${RUNTIME_STAGING_DIR}/bin" "${RUNTIME_STAGING_DIR}/lib" "${RUNTIME_STAGING_DIR}/share/ptoas" - cp "${ptoas_bin}" "${RUNTIME_STAGING_DIR}/bin/ptoas" + mkdir -p "${RUNTIME_STAGING_DIR}/lib" "${RUNTIME_STAGING_DIR}/share/ptoas" "${RUNTIME_STAGING_DIR}/pto" cp -R "${PTO_INSTALL_DIR}/share/ptoas/TileOps" "${RUNTIME_STAGING_DIR}/share/ptoas/TileOps" + cp "${PTO_INSTALL_DIR}/lib/ptoas.so" "${RUNTIME_STAGING_DIR}/pto/ptoas.so" + # Resolve transitive MLIR/LLVM dependencies through the build tree so + # ldd can discover non-installed libs such as libMLIRMlirOptMain.so.*. + local dep_ld_library_path="${LLVM_BUILD_DIR}/lib:${PTO_INSTALL_DIR}/lib:${LD_LIBRARY_PATH:-}" while read -r dep_path; do [[ -n "${dep_path}" ]] || continue should_bundle_linux_dep "${dep_path}" || continue cp -L -n "${dep_path}" "${RUNTIME_STAGING_DIR}/lib/" - done < <(linux_runtime_dep_paths "${ptoas_bin}" | sort -u) + done < <(linux_runtime_dep_paths "${PTO_INSTALL_DIR}/lib/ptoas.so" "${dep_ld_library_path}" | sort -u) local version_output + local version_ld_library_path="${LLVM_BUILD_DIR}/lib:${RUNTIME_STAGING_DIR}/lib:${PTO_INSTALL_DIR}/lib:${LD_LIBRARY_PATH:-}" version_output="$( env -u PYTHONPATH -u DYLD_LIBRARY_PATH \ - LD_LIBRARY_PATH="${RUNTIME_STAGING_DIR}/lib:${LD_LIBRARY_PATH:-}" \ - "${RUNTIME_STAGING_DIR}/bin/ptoas" --version | tr -d '\r' + LD_LIBRARY_PATH="${version_ld_library_path}" \ + "${ptoas_wrapper}" --version | tr -d '\r' )" echo "${version_output}" if [[ -n "${PTOAS_VERSION:-}" ]]; then @@ -149,9 +161,9 @@ cp -R "${PTOAS_WRAPPER_PKG_DIR}" "${WHEEL_STAGING_DIR}/ptoas" echo "Embedding unified runtime payload for wheel-side ptoas launcher..." mkdir -p "${WHEEL_STAGING_DIR}/ptoas/_runtime" -cp -R "${RUNTIME_STAGING_DIR}/bin" "${WHEEL_STAGING_DIR}/ptoas/_runtime/bin" cp -R "${RUNTIME_STAGING_DIR}/share" "${WHEEL_STAGING_DIR}/ptoas/_runtime/share" cp -R "${RUNTIME_STAGING_DIR}/lib" "${WHEEL_STAGING_DIR}/ptoas/_runtime/lib" +cp -R "${RUNTIME_STAGING_DIR}/pto" "${WHEEL_STAGING_DIR}/pto" echo "Removing packaging residue..." find "${WHEEL_STAGING_DIR}" \( -name '*.egg-info' -o -name '*.dist-info' \) -prune -exec rm -rf {} + @@ -208,7 +220,7 @@ wheel = "\n".join([ record_rows = [] has_ptodsl_init = False -has_ptoas_runtime_binary = False +has_ptoas_shared_module = False def hash_bytes(data: bytes) -> str: digest = hashlib.sha256(data).digest() @@ -224,8 +236,8 @@ with zipfile.ZipFile(wheel_path, "w", compression=zipfile.ZIP_DEFLATED) as zf: record_rows.append((rel, hash_bytes(data), str(len(data)))) if rel == "ptodsl/__init__.py": has_ptodsl_init = True - if rel == "ptoas/_runtime/bin/ptoas": - has_ptoas_runtime_binary = True + if rel == "pto/ptoas.so": + has_ptoas_shared_module = True entry_points = "\n".join([ "[console_scripts]", @@ -253,14 +265,16 @@ with zipfile.ZipFile(wheel_path, "w", compression=zipfile.ZIP_DEFLATED) as zf: if not has_ptodsl_init: raise SystemExit("Wheel staging payload is missing ptodsl/__init__.py") -if not has_ptoas_runtime_binary: - raise SystemExit("Wheel staging payload is missing ptoas/_runtime/bin/ptoas") +if not has_ptoas_shared_module: + raise SystemExit("Wheel staging payload is missing pto/ptoas.so") with zipfile.ZipFile(wheel_path) as zf: if "ptodsl/__init__.py" not in zf.namelist(): raise SystemExit("Built wheel is missing ptodsl/__init__.py") - if "ptoas/_runtime/bin/ptoas" not in zf.namelist(): - raise SystemExit("Built wheel is missing ptoas/_runtime/bin/ptoas") + if "pto/ptoas.so" not in zf.namelist(): + raise SystemExit("Built wheel is missing pto/ptoas.so") + if "ptoas/_runtime/bin/ptoas" in zf.namelist(): + raise SystemExit("Built wheel unexpectedly contains ptoas/_runtime/bin/ptoas") print(f"Wheel created at {wheel_path}") PY diff --git a/docker/setup.py b/docker/setup.py index d26d85cd33..6f2278f0b0 100644 --- a/docker/setup.py +++ b/docker/setup.py @@ -7,12 +7,26 @@ # See LICENSE in the root of the software repository for the full text of the License. import os +import pathlib +import re from setuptools import setup, find_namespace_packages +_PROJECT_VERSION_RE = re.compile( + r"project\s*\(\s*ptoas\s+VERSION\s+([0-9]+\.[0-9]+)\s*\)" +) + + def read_package_version() -> str: - return os.environ.get("PTOAS_PYTHON_PACKAGE_VERSION", "0.1.1") + version = os.environ.get("PTOAS_PYTHON_PACKAGE_VERSION", "").strip() + if version: + return version + cmake_file = pathlib.Path(__file__).resolve().parents[1] / "CMakeLists.txt" + match = _PROJECT_VERSION_RE.search(cmake_file.read_text(encoding="utf-8")) + if not match: + raise RuntimeError(f"could not find PTOAS version in {cmake_file}") + return match.group(1) setup( name="ptoas", diff --git a/docker/setup_mac.py b/docker/setup_mac.py index 9af5a968d8..7dcd218b75 100644 --- a/docker/setup_mac.py +++ b/docker/setup_mac.py @@ -7,6 +7,8 @@ # See LICENSE in the root of the software repository for the full text of the License. import os +import pathlib +import re from setuptools import find_namespace_packages, setup from setuptools.dist import Distribution @@ -20,7 +22,17 @@ def has_ext_modules(self): def read_package_version() -> str: - return os.environ.get("PTOAS_PYTHON_PACKAGE_VERSION", "0.1.1") + version = os.environ.get("PTOAS_PYTHON_PACKAGE_VERSION", "").strip() + if version: + return version + cmake_file = pathlib.Path(__file__).resolve().parents[1] / "CMakeLists.txt" + project_version_re = re.compile( + r"project\s*\(\s*ptoas\s+VERSION\s+([0-9]+\.[0-9]+)\s*\)" + ) + match = project_version_re.search(cmake_file.read_text(encoding="utf-8")) + if not match: + raise RuntimeError(f"could not find PTOAS version in {cmake_file}") + return match.group(1) setup( name="ptoas", diff --git a/docker/test_ptoas_cli.sh b/docker/test_ptoas_cli.sh index 3d7dfe9fca..a6f59158c6 100755 --- a/docker/test_ptoas_cli.sh +++ b/docker/test_ptoas_cli.sh @@ -16,6 +16,7 @@ # PTO_BUILD_DIR - Path to PTO build directory # LLVM_BUILD_DIR - Path to LLVM build directory # PTO_INSTALL_DIR - Path to PTO install directory +# PTOAS_BIN - Path to the Python ptoas wrapper (optional) set -e @@ -32,11 +33,20 @@ export PATH="${PTO_BUILD_DIR}/tools/ptoas:${PATH}" export LD_LIBRARY_PATH="${LLVM_BUILD_DIR}/lib:${PTO_INSTALL_DIR}/lib:${LD_LIBRARY_PATH}" export DYLD_LIBRARY_PATH="${LLVM_BUILD_DIR}/lib:${PTO_INSTALL_DIR}/lib:${DYLD_LIBRARY_PATH}" +if [ -n "${PTOAS_BIN:-}" ]; then + if [ ! -x "${PTOAS_BIN}" ]; then + echo "Error: PTOAS_BIN is not executable: ${PTOAS_BIN}" >&2 + exit 1 + fi +else + PTOAS_BIN="$(command -v ptoas)" +fi + echo "Testing ptoas CLI..." -which ptoas +echo "${PTOAS_BIN}" echo "Checking ptoas version..." -VERSION_OUTPUT="$(ptoas --version | tr -d '\r')" +VERSION_OUTPUT="$("${PTOAS_BIN}" --version | tr -d '\r')" echo "$VERSION_OUTPUT" if [ -n "${PTOAS_VERSION:-}" ]; then EXPECTED_VERSION_OUTPUT="ptoas ${PTOAS_VERSION}" @@ -52,14 +62,14 @@ fi echo "Testing MatMul sample..." cd "${PTO_SOURCE_DIR}/test/samples/MatMul/" python ./tmatmulk.py > ./tmatmulk.pto -ptoas ./tmatmulk.pto -o ./tmatmulk.cpp +"${PTOAS_BIN}" ./tmatmulk.pto -o ./tmatmulk.cpp echo "MatMul test passed" # Test Abs sample echo "Testing Abs sample..." cd "${PTO_SOURCE_DIR}/test/samples/Abs/" python ./abs.py > ./abs.pto -ptoas --enable-insert-sync ./abs.pto -o ./abs.cpp +"${PTOAS_BIN}" --enable-insert-sync ./abs.pto -o ./abs.cpp echo "Abs test passed" echo "All ptoas CLI tests passed!" diff --git a/docker/test_wheel_imports.sh b/docker/test_wheel_imports.sh index 43b4f9defd..3292f97a4a 100755 --- a/docker/test_wheel_imports.sh +++ b/docker/test_wheel_imports.sh @@ -54,6 +54,30 @@ if [[ -n "${WHEEL_GLOB}" ]] && compgen -G "${WHEEL_GLOB}" >/dev/null 2>&1; then TEST_WHEEL="$(compgen -G "${WHEEL_GLOB}" | sort | tail -n 1)" TEST_TMPDIR="$(mktemp -d /tmp/ptoas-wheel-test.XXXXXX)" echo "Installing wheel into isolated venv: ${TEST_WHEEL}" + echo "Checking wheel payload before installation..." + TEST_WHEEL="${TEST_WHEEL}" "${PYTHON_BIN}" - <<'PY' +import os +import zipfile +from pathlib import Path + +wheel = Path(os.environ["TEST_WHEEL"]) +with zipfile.ZipFile(wheel) as zf: + names = set(zf.namelist()) + required = {"ptoas/__init__.py", "ptoas/_launcher.py", "pto/ptoas.so"} + missing = sorted(required - names) + if missing: + raise SystemExit(f"wheel is missing required payload files: {missing}") + if "ptoas/_runtime/bin/ptoas" in names: + raise SystemExit("wheel unexpectedly contains ptoas/_runtime/bin/ptoas") + entry_points_name = next( + name for name in names + if name.startswith("ptoas-") and name.endswith(".dist-info/entry_points.txt") + ) + entry_points = zf.read(entry_points_name).decode("utf-8") + if "ptoas=ptoas._launcher:main" not in entry_points: + raise SystemExit("wheel entry points do not route ptoas through ptoas._launcher:main") +print(f"Wheel payload check passed: {wheel.name}") +PY "${PYTHON_BIN}" -m venv "${TEST_TMPDIR}/venv" source "${TEST_TMPDIR}/venv/bin/activate" python -m pip install --no-deps --force-reinstall "${TEST_WHEEL}" diff --git a/docs/designs/vmi-e2b-scale-broadcast-optimization.md b/docs/designs/vmi-e2b-scale-broadcast-optimization.md new file mode 100644 index 0000000000..9777162f3b --- /dev/null +++ b/docs/designs/vmi-e2b-scale-broadcast-optimization.md @@ -0,0 +1,992 @@ +# VMI E2B Scale Broadcast Optimization Study + +本文推演 VMI 是否能把 block quant 中的 scale broadcast 自动优化成 +`E2B_B16` load。结论是： + +```text +group_slot_load + group_broadcast 足以表达逻辑语义。 + +它不足以单独触发 E2B，因为 E2B 是某个 physical chunk layout 下的 +materialization，不是 dense logical broadcast 的直接 lowering。 + +如果后续 layout 已经由 consumer requirement 或 target-specific layout +optimization 选成 E2B-compatible 形态，vmi-to-vpto 可以把对应 chunk lower +成 E2B。 + +如果想从普通 dense quant IR 自动得到 CCE 的 DINTLV/E2B 形态，需要一个 +target-specific layout optimization/cost selection 阶段整体选择这套计划。 +``` + +## 1. Logical Quant Semantics + +`ComputeY1ToFP8` 的 surface VMI 语义应保持 dense quant： + +```text +for i in 0..255: + y[i] = fp8(x[i] * scale[i / 32]) +``` + +也就是 8 个 scale，每个覆盖 32 个 dense logical lanes： + +```text +s0 x32, s1 x32, ..., s7 x32 +``` + +对应 VMI 形态是： + +```text +%x = pto.vmi.load %x_base[%x_off] + : !pto.ptr -> !pto.vmi.vreg<256xf16> + +%scale_slots = pto.vmi.group_slot_load %scale_base[%scale_off], %c1 + {num_groups = 8} + : !pto.ptr -> !pto.vmi.vreg<8xbf16> + +%scale = pto.vmi.group_broadcast %scale_slots {num_groups = 8} + : !pto.vmi.vreg<8xbf16> -> !pto.vmi.vreg<256xbf16> +``` + +This form is the canonical logical IR. The source scale should be BF16 payload, +not FP16, because the CCE implementation loads `uint16_t` values and later +reinterprets them as `vector_bf16`. + +`num_groups = 16` would express a different algorithm: + +```text +16 scale values, each covering 16 dense lanes +``` + +That is not equivalent unless the input memory redundantly stores +`s0, s0, s1, s1, ...`, which is not what the CCE kernel does. + +## 2. E2B_B16 Semantics + +`E2B_B16` is a VPTO load distribution mode. For a b16 destination register it +loads 8 source elements and expands each one to 16 consecutive destination +lanes: + +```text +dst[j] = src[floor(j / 16)] for j = 0..127 +``` + +The result is: + +```text +s0 x16, s1 x16, ..., s7 x16 +``` + +So `E2B_B16` does not directly materialize the dense VMI broadcast +`8 -> 256`. It materializes one 128-lane physical view that becomes useful only +after the x data and later f32 computation have been split into compatible +physical chunks. + +## 3. Why CCE Can Use E2B + +The CCE FP16 path uses a physical implementation shape like: + +```text +vlds(x0F16, x1F16, xHalf, stride, DINTLV_B16, POST_UPDATE) +vlds(scaleForMulFP16, scale_base, 0, E2B_B16) + +vcvt(x0_even_f32, x0F16, PART_EVEN) +vcvt(x0_odd_f32, x0F16, PART_ODD) +vcvt(x1_even_f32, x1F16, PART_EVEN) +vcvt(x1_odd_f32, x1F16, PART_ODD) + +vcvt(scale_f32, (vector_bf16 &)scaleForMulFP16, PART_EVEN) +``` + +`DINTLV_B16` splits the dense 256-element row into two 128-lane physical streams. +After each stream is converted from f16 to f32, the computation is effectively +four 64-lane f32 chunks: + +```text +x0 even part +x0 odd part +x1 even part +x1 odd part +``` + +For every one of those chunks, the needed scale pattern is: + +```text +s0 x8, s1 x8, ..., s7 x8 +``` + +`E2B_B16` produces: + +```text +s0 x16, s1 x16, ..., s7 x16 +``` + +Then `vcvt PART_EVEN` produces: + +```text +s0 x8, s1 x8, ..., s7 x8 +``` + +Because every scale value is duplicated in adjacent even/odd b16 positions, +`PART_EVEN` and `PART_ODD` would produce the same f32 scale chunk. The CCE code +computes the scale chunk once and reuses it for all four x chunks. + +## 4. What Is A Legal Automatic Optimization? + +The following rewrite is not legal as a standalone local rule: + +```text +group_slot_load + group_broadcast(8 -> 256) => E2B_B16 +``` + +It is invalid because the left side is a dense 256-lane logical value, while +`E2B_B16` produces a 128-lane physical value with a different lane repetition +count. + +A legal E2B lowering must be conditional on the assigned physical layout: + +```text +if the broadcasted scale value is required in physical chunks where each chunk +needs s0 x16 ... s7 x16 at b16 width, or s0 x8 ... s7 x8 after bf16->f32, +then that chunk may be materialized with E2B_B16. +``` + +In other words: + +```text +group_slot_load + group_broadcast + is the logical source pattern + +consumer-required or target-selected layout + determines whether any physical chunk is E2B-compatible + +vmi-to-vpto + lowers only those compatible chunks to E2B +``` + +`group_slot_load` alone cannot lower to E2B. A group-slot value has only group +slots as semantic lanes. `E2B_B16` already produces broadcasted physical lanes. +The `group_broadcast` use is required to justify reading those lanes. + +## 5. Layout Selection Boundary + +Deinterleaved layout must not be inferred only because E2B would be cheaper. +The selected layout must be explicit before `vmi-to-vpto`. That layout can come +from either side: + +```text +consumer requirement: + a later op requires a particular layout. + +producer natural layout: + the producing op has a declared, deterministic natural layout that is legal + for all of its uses. +``` + +`group_broadcast` is a materialization op, so it may define or participate in an +E2B-friendly natural layout when that layout is part of the declared layout +contract. That is still a layout-assignment decision, not a hidden +`vmi-to-vpto` peephole. Do not reuse `block_elems` as an ad-hoc broadcast split +knob; `block_elems` belongs to the dense deinterleaved layout contract and has +existing producer/consumer meanings. + +Baseline layout assignment may still choose conservative contiguous layouts even +when a target-specific fused implementation exists. + +Therefore this optimization has two valid implementation levels. + +### 5.1 Compatible-Layout Lowering Shortcut + +If some earlier layout pass has already assigned an E2B-compatible physical +layout, `vmi-to-vpto` may lower the scale chunk with `E2B_B16`. + +This is a local deterministic lowering. It does not discover the CCE plan by +itself. It only avoids a generic `vsldb + vselr` materialization when the +assigned layout has already made the required physical chunk shape explicit. + +### 5.2 Producer Natural Layout + +For simple broadcasts, the producer itself may choose an E2B-friendly natural +layout when that layout satisfies every use. + +Example for b16, using an existing DINTLV-like element-parity layout: + +```text +logical 1 -> 32: + s0 x32 + +layout: + deinterleaved = 2, block_elems = 1 + +physical part 0: + s0 x16 + +physical part 1: + s0 x16 +``` + +The two physical parts can share one E2B materialization or use two identical +E2B materializations. This is a general layout choice for the broadcast result, +not a quant-specific graph rewrite. + +For a uniform `1 -> 32` or per-group `x32` broadcast, `deinterleaved = 2, +block_elems = 1` yields 16 lanes of the same group per physical part and is +closer to an even/odd `DINTLV_B16` data layout. + +For the MX quant scale: + +```text +logical 8 -> 256: + s0 x32, s1 x32, ..., s7 x32 + +layout: + deinterleaved = 2, block_elems = 1 + +physical part 0: + s0 x16, s1 x16, ..., s7 x16 + +physical part 1: + s0 x16, s1 x16, ..., s7 x16 +``` + +Each physical part is directly `E2B_B16`-compatible. +The implementation should run the E2B compatibility query over the assigned lane +mapping. It should not infer a new meaning for `block_elems`. + +### 5.3 Target-Specific Layout Optimization + +To automatically discover the complete CCE plan from canonical dense quant IR, +add an optional target-specific layout optimization before `vmi-to-vpto`. + +That pass may select a cheaper equivalent implementation for the whole quant +subgraph: + +```text +dense x load +f16/bf16 -> f32 conversion +scale group_slot_load + group_broadcast 8 -> 256 +scale bf16 -> f32 conversion +mul +fp32 -> fp8 conversion/store +``` + +The pass must rewrite or annotate the VMI layout-assigned IR so that +`vmi-to-vpto` no longer has to infer the plan from context. + +Expected selected physical plan: + +```text +x load: + vlds DINTLV_B16 into two b16 streams + +scale load: + vlds E2B_B16 into one b16 stream + +scale conversion: + vcvt PART_EVEN into one 64-lane f32 stream + +mul: + reuse that scale f32 stream for the four x f32 chunks +``` + +This is an optimization, not a correctness requirement. If the optimizer does +not fire, the canonical dense VMI program still has a valid generic lowering. + +## 6. Candidate Match Preconditions + +A target-specific optimization may match the CCE-style scale pattern only under +strict conditions: + +```text +scale_slots: + pto.vmi.group_slot_load + num_groups = 8 + source_group_stride = 1 + source element width = 16 bits + semantic type is bf16 or a bitcastable ui16 payload later interpreted as bf16 + +scale broadcast: + pto.vmi.group_broadcast + same num_groups = 8 + dense logical result has 256 b16 lanes for this case + +scale conversion: + bf16 -> f32 + conversion has no rounding/exception behavior that distinguishes duplicated + even and odd source lanes + +x path: + dense logical row has 256 f16 or bf16 lanes + the target plan can legally compute the row as four 64-lane f32 chunks + +uses: + no user observes the intermediate dense scale layout in a way that prevents + rematerialization or chunk reuse +``` + +The optimization should reject or skip the pattern if any of these conditions are +not proven. + +## 7. Correctness Sketch + +Let the logical dense lane be `i`. + +The canonical VMI scale value is: + +```text +scale_dense[i] = s[floor(i / 32)] +``` + +The CCE physical decomposition maps each dense lane into one of four f32 chunks. +For a chunk-local f32 lane `k`: + +```text +dense lane = 4 * k + delta +delta in {0, 1, 2, 3} +``` + +Then: + +```text +floor((4 * k + delta) / 32) = floor(k / 8) +``` + +So every f32 chunk needs: + +```text +scale_chunk[k] = s[floor(k / 8)] +``` + +`E2B_B16` plus `vcvt PART_EVEN` gives: + +```text +e2b_b16[j] = s[floor(j / 16)] for j = 0..127 +scale_f32[k] = e2b_b16[2 * k] + = s[floor((2 * k) / 16)] + = s[floor(k / 8)] +``` + +That matches the required `scale_chunk[k]` for all four f32 chunks. + +## 8. Recommendation + +Prefer adding a target-agnostic VMI `group_broadcast_load` logical memory op if +we want to make this optimization robust and local. The op should mean: + +```text +load one source value per logical group, then broadcast that value to every lane +in the group. +``` + +It must not mean `E2B`. `E2B_B16` is only one possible lowering when the +assigned layout is compatible. + +The unfused logical IR remains valid: + +```text +group_slot_load + group_broadcast +``` + +but a canonicalization/layout-prep pass may fuse it to: + +```text +group_broadcast_load +``` + +when the group-slot value has no separate semantic users. + +Then implement E2B support in phases: + +```text +1. Ensure the example PTO uses the correct logical semantics: + bf16 scale, num_groups = 8, dense 8 -> 256 broadcast. + +2. Add group_broadcast_load as a logical VMI memory op, plus canonicalization + from group_slot_load + group_broadcast when legal. + +3. Add a compatible-layout lowering shortcut: + when layout assignment already exposes an E2B-compatible chunk, lower the + group_broadcast_load chunk with vlds E2B_B16. + +4. Add an optional target-specific quant layout optimization: + recognize the whole dense quant subgraph and select the DINTLV/E2B plan when + it is legal and profitable. +``` + +This keeps VMI logical semantics independent from physical layout, while still +leaving a clear path to recover the CCE optimization automatically. + +## 9. Generalized E2B Broadcast Optimization + +The scale case above is one instance of a broader rule: E2B is a physical +materialization primitive for a packet of repeated group slots. It is not tied +to MX quant, but its legality depends on the physical chunk layout and on the +load distribution's carrier element width. + +### 9.1 E2B As A Packet Primitive + +For the verified `B16` case: + +```text +E2B_B16 packet: + source slots per packet = 8 + destination lanes per packet = 128 b16 lanes + repeat per source slot = 16 b16 lanes + +dst[lane] = src[base_slot + floor(lane / 16)] +``` + +This can materialize a physical chunk that needs: + +```text +s0 x16, s1 x16, ..., s7 x16 +``` + +The optimization should reason in terms of physical chunks: + +```text +logical group_broadcast + source group slot for logical lane i = floor(i / logical_group_size) + +assigned physical layout + maps physical chunk lane l to logical lane i(l) + +E2B-compatible chunk + floor(i(l) / logical_group_size) = base_slot + floor(l / 16) +``` + +If this equality holds for a b16 physical chunk, the chunk can be loaded with +`E2B_B16` instead of materializing the broadcast with `vselr`. + +### 9.2 Direct 1 -> 16 + +A logical `1 -> 16` b16 broadcast is directly compatible with one E2B group: + +```text +s0 x16 +``` + +However, `E2B_B16` is naturally an 8-group packet: + +```text +s0 x16, s1 x16, ..., s7 x16 +``` + +So a single `1 -> 16` use may lower to E2B only under one of these conditions: + +```text +packed case: + the compiler can pack eight independent 1 -> 16 broadcasts into one E2B load. + +partial-live case: + only one 16-lane group is live, and the target semantics prove inactive E2B + groups do not require valid source memory or can be safely over-read. + +full-packet case: + the logical IR actually contains eight adjacent groups, even if the current + consumer observes only one group through a layout/mask. +``` + +If these conditions are not proven, `BRC_B16`, `vdup`, or the existing generic +broadcast lowering is safer than E2B. In particular, do not introduce an E2B +load that reads seven extra source values unless the memory safety rule is +explicit. + +### 9.3 1 -> 32 Via Deinterleaved Reuse + +A dense logical `1 -> 32` b16 broadcast does not fit one E2B group in a single +contiguous physical chunk: + +```text +logical: s0 x32 +E2B group: s0 x16 +``` + +It becomes E2B-compatible when the assigned physical layout splits those 32 +logical lanes into two 16-lane physical uses: + +```text +physical use A: s0 x16 +physical use B: s0 x16 +``` + +This split can use the existing DINTLV-like element-parity layout: + +```text +#pto.vmi.layout +``` + +For logical lanes `0..31`, this maps: + +```text +even lanes 0,2,...,30 -> physical part 0 lanes 0..15 +odd lanes 1,3,...,31 -> physical part 1 lanes 0..15 +``` + +Because all 32 logical lanes carry the same `s0`, each part still sees +`s0 x16`. The lowering rule should check the resulting group index function, +not invent a new layout spelling. + +Then the compiler has two valid strategies: + +```text +reuse: + materialize one E2B group/chunk and map both physical uses to the same value. + +duplicate: + materialize the same E2B group twice if reuse would violate scheduling, + lifetime, or destructive-update constraints. +``` + +This is the mechanism behind the MX quant scale case: + +```text +dense logical scale: 8 groups, each x32 +physical f16/bf16 streams: each group appears as x16 per stream +``` + +The optimization is legal only if the 32 logical lanes are split by layout. It +is not legal as a direct E2B chunk load for a contiguous physical chunk that +genuinely needs `s0 x32` inside one chunk; that would require a separate +duplicate/interleave/concat materialization. + +### 9.4 N -> N * 16 And N -> N * 32 + +For b16 group broadcasts with consecutive slots and unit source stride: + +```text +N -> N * 16 +``` + +can be lowered by E2B in packets of 8 groups when the physical chunk sees the +groups in E2B order: + +```text +for base_slot in 0, 8, 16, ... + load src[base_slot : base_slot + 8] with E2B_B16 +``` + +Tail packets require either a proven safe masked/partial E2B form or a generic +fallback. + +For: + +```text +N -> N * 32 +``` + +E2B is profitable when the assigned layout decomposes each 32-lane logical group +into two 16-lane physical uses. That assigned layout may be the +`group_broadcast` producer's natural layout, or it may be required by a +downstream consumer. The lowering then reuses or duplicates the corresponding +E2B materialization for those two uses. This rule extends to: + +```text +N -> N * (16 * F) +``` + +when the layout decomposes each logical group into `F` physical 16-lane uses. + +### 9.5 Type Generalization + +E2B is a carrier-width load distribution. For `E2B_B16`, the load itself is +valid for 16-bit carriers: + +```text +bf16 +f16 +ui16 / si16 payloads +other 16-bit bit patterns whose consumers preserve the intended interpretation +``` + +The optimization must keep type interpretation outside the load: + +```text +bf16 scale + extf to f32: + E2B_B16 may feed vcvt bf16 -> f32. + +f16 broadcast: + E2B_B16 may materialize repeated f16 lanes if the consumer expects f16. + +ui16 payload later bitcast to bf16: + E2B_B16 may load the ui16 carrier, but the later bitcast/interpretation must + remain explicit in VMI or in the selected lowering plan. +``` + +Do not infer a floating-point type from E2B itself. `E2B_B16` only says how UB +bytes are placed into b16 lanes. + +`E2B_B32` is the b32 member of the same distribution family. The VPTO verifier +accepts `E2B_B32`, the ISA docs list E2B for `b16` and `b32`, and CCE quant code +uses `E2B_B32` in FP32 paths. It follows the same 8-source-slot packet rule: + +```text +E2B_B16: 8 source slots * 16 lanes/slot = 128 b16 lanes +E2B_B32: 8 source slots * 8 lanes/slot = 64 b32 lanes +``` + +The implemented E2B broadcast optimization therefore supports: + +```text +b16 contiguous: logical 1 -> 16 +b16 deinterleaved=2: logical 1 -> 32 +b32 contiguous: logical 1 -> 8 +b32 deinterleaved=2: logical 1 -> 16 +``` + +There is no `E2B_B8` in the documented load distribution family, so b8 +broadcasts should use other distributions or generic materialization. + +### 9.6 Broadcast Generalization + +E2B can optimize `pto.vmi.group_broadcast` when all of these are true: + +```text +source: + group slots come from consecutive memory slots + source_group_stride = 1 + slot type matches the E2B carrier width + +broadcast: + each physical chunk needs a run-length pattern compatible with the E2B repeat + count for that carrier width + +layout: + the run-length pattern is visible in the assigned layout before vmi-to-vpto + +uses: + rematerializing or reusing the E2B packet does not change observable memory or + arithmetic semantics +``` + +E2B is generally not the right primitive for ordinary scalar `pto.vmi.broadcast` +unless the scalar value is already in memory as an E2B packet or the compiler can +pack several independent scalar broadcasts into one E2B load. For a scalar +stored once in memory and needed in every lane, `BRC_B16/B32`, `BRC_BLK`, or a +register `vdup` is usually the more direct representation. + +### 9.7 Implementation Shape + +The recommended implementation order is: + +```text +1. Keep VMI semantics canonical: + group_slot_load + group_broadcast is the desugared meaning. + +2. Optionally canonicalize to group_broadcast_load: + this keeps memory source and broadcast semantics in one local op. + +3. Add an E2B compatibility query over assigned physical chunks: + given source slots, result layout, carrier width, and live lanes, answer + whether a chunk's group-index function is E2B-shaped. + +4. Lower compatible chunks to E2B packets: + generate one E2B load per needed packet, or reuse an existing packet when + multiple physical uses require identical contents. + +5. Add a later target-specific layout optimizer: + it may choose layouts that expose E2B-compatible chunks, but only by + rewriting/annotating layout-assigned VMI before vmi-to-vpto. +``` + +The compatibility query should return a reason when it rejects a candidate: + +```text +non-unit source stride +non-consecutive group slots +unsupported carrier width +tail packet lacks safe partial E2B semantics +physical lane mapping is not E2B-shaped +extra source memory read would be unsafe +consumer observes a different dense layout +``` + +This keeps the optimization auditable and prevents E2B from becoming an implicit +layout-changing peephole. + +## 10. Recognition, Solidification, Propagation, Lowering + +This section describes how an implementation should carry the optimization from +canonical VMI to VPTO without making `vmi-to-vpto` rediscover hidden context. + +### 10.1 Recognize Information + +Run recognition after hard layout assignment, when every relevant value already +has an explicit layout. + +Recognize the source shape: + +```text +%slots = pto.vmi.group_slot_load %base[%off], %stride {num_groups = G} +%bcast = pto.vmi.group_broadcast %slots {num_groups = G} +``` + +or the already-fused form: + +```text +%bcast = pto.vmi.group_broadcast_load %base[%off], %stride {num_groups = G} +``` + +Collect candidate facts: + +```text +source memory: + base pointer + offset + source_group_stride + element carrier width + memory element type + +logical broadcast: + num_groups = G + logical lanes = N + logical group size S = N / G + +assigned result layout: + physical arity + physical lanes per chunk + logical lane mapped to each physical lane + +uses: + whether the broadcast feeds elementwise ops, extf/truncf, stores, or multiple + independent consumers +``` + +Then compute an E2B packet plan per physical chunk. For `E2B_B16`, a physical +chunk is compatible when: + +```text +group_index_for_physical_lane(l) = base_slot + floor(l / 16) +``` + +for all live lanes in that chunk. + +Reject the candidate if: + +```text +source_group_stride != 1 +source slots are not consecutive +carrier width is unsupported +the assigned layout does not produce E2B-shaped chunks +tail/partial packet would read memory that is not proven valid +the group_slot_load has other non-rematerializable users +``` + +This recognition is an analysis step. It must not silently change layouts. + +### 10.2 Solidify Information + +`vmi-to-vpto` should not have to look at an arbitrary +`group_slot_load -> group_broadcast` use-def chain and decide to suppress one +load while replacing another op with E2B. The optimization pass must solidify +the decision in the layout-assigned IR. + +The preferred solidification is a target-agnostic logical memory op: + +```text +%bcast = pto.vmi.group_broadcast_load %base[%off], %stride {num_groups = G} + : !pto.ptr -> !pto.vmi.vreg +``` + +Semantic definition: + +```text +group_size = N / G +for logical lane i: + group = floor(i / group_size) + result[i] = base[off + group * stride] +``` + +This op is not target-specific and does not promise E2B. It is exactly the +fused logical form of: + +```text +%slots = pto.vmi.group_slot_load %base[%off], %stride {num_groups = G} +%bcast = pto.vmi.group_broadcast %slots {num_groups = G} +``` + +The fused op makes lowering local because the memory source, stride, group count, +result type, and assigned layout are all available on one op. A generic lowering +can still materialize it with `vsldb + vselr`; an optimized lowering may choose +`E2B_B16` for compatible physical chunks. + +The current implementation is intentionally narrower: because +`group_broadcast_load` does not yet have a generic `vsldb + vselr` lowering, +layout assignment fuses `group_slot_load + group_broadcast` only when the fused +op is already an E2B-compatible b16 candidate. Non-E2B shapes stay in the +unfused form and continue to use the existing `group_slot_load` plus +`group_broadcast` lowering path. + +Canonicalization rules: + +```text +group_slot_load + group_broadcast -> group_broadcast_load + when the group_slot_load has exactly that broadcast use, or when cloning the + load is legal and profitable for that use. + +group_broadcast_load -> group_slot_load + group_broadcast + remains a valid conceptual expansion for verification, documentation, and + generic fallback reasoning. +``` + +Solidification must preserve semantics for multi-use values: + +```text +if all uses consume only the broadcasted value: + replace with one shared group_broadcast_load. + +if only one use can benefit from the fused form: + clone/rematerialize that use-site load as group_broadcast_load and keep the + original group_slot_load for other users. + +if the group_slot_load itself has semantic group-slot users: + do not delete it; add a separate group_broadcast_load only if the extra memory + read is legal or if load cloning is otherwise proven safe. +``` + +### 10.3 Propagate Information + +After solidification, propagation should use ordinary VMI layout rules whenever +possible: + +```text +elementwise ops: + preserve the assigned layout when operands agree. + +ensure_layout: + makes layout transitions explicit when one use needs E2B-compatible chunks and + another use needs a different layout. + +rematerialization: + may clone group_broadcast_load per use-site instead of forcing a single layout + for all consumers. +``` + +For casts, propagation may need a targeted rule. The important MX quant case is: + +```text +E2B_B16 gives: + s0 x16, s1 x16, ..., s7 x16 + +bf16 -> f32 PART_EVEN gives: + s0 x8, s1 x8, ..., s7 x8 +``` + +If multiple f32 physical chunks require that same `s0 x8 ... s7 x8` pattern, +the post-assignment plan may mark them as the same rematerialized value. The +lowerer can then generate one `vcvt PART_EVEN` and map several logical physical +chunks to the same VPTO value. + +This reuse fact must be derived from the assigned lane mapping and the E2B packet +plan. It must not rely on a later CSE pass accidentally proving the duplicate. + +### 10.4 Implement Lowering + +`vmi-to-vpto` should lower `group_broadcast_load` locally. It may choose E2B +only when the op's assigned layout and source facts produce an explicit +E2B-compatible packet plan. + +For each E2B packet: + +```text +1. compute the source pointer: + base + packet_base_slot + +2. emit: + pto.vlds {dist = "E2B_B16"} + +3. map the emitted VPTO value to the physical result chunk(s) recorded in the + group_broadcast_load packet plan. +``` + +For `1 -> 32` under `deinterleaved = 2, block_elems = 1`: + +```text +logical group: + s0 x32 + +physical part 0: + s0 x16 + +physical part 1: + s0 x16 + +lowering: + emit one E2B packet if reuse is legal, or two identical E2B packets if + scheduling/lifetime constraints require duplication. +``` + +For MX quant scale after bf16->f32: + +```text +1. emit E2B_B16 for the b16 scale packet. +2. emit vcvt PART_EVEN to produce the f32 packet. +3. map that f32 packet to every physical f32 chunk whose lane mapping requires + s0 x8, s1 x8, ..., s7 x8. +4. lower mulf normally using the assigned physical chunks. +``` + +### 10.5 Where Layout Choices Happen + +There are three levels of optimization: + +```text +level 0: no E2B + canonical group_broadcast lowers through generic vselr materialization. + +level 1: E2B for already-compatible layouts + recognition sees the assigned layout is E2B-shaped and solidifies an E2B + materialization. + +level 2: choose E2B-compatible layouts + an optional layout optimization changes/rematerializes layouts before + recognition, for example selecting deinterleaved=2/block_elems=1 for a + broadcast use when all consumers can accept that layout. +``` + +The full CCE-like optimization for `ComputeY1ToFP8` is level 2: + +```text +x path: + select DINTLV-compatible layout for the dense x load/cast path. + +scale path: + select an E2B-compatible broadcast materialization. + +compute path: + keep mul/trunc/store in the selected physical chunk layout or insert explicit + layout materialization where required. +``` + +### 10.6 Test Plan + +Add focused tests in phases: + +```text +positive: + bf16 group_slot_load stride=1 + group_broadcast 8->256 assigned to + deinterleaved=2/block_elems=1 lowers scale chunks with E2B_B16. + +positive: + f16 1->16 or packed 8*(1->16) lowers to E2B only when source memory safety is + proven by full packet or supported partial semantics. + +positive: + 1->32 assigned to deinterleaved=2/block_elems=1 maps two physical uses to one + E2B packet or to two explicit duplicate packets. + +positive: + f32 1->8 lowers to E2B_B32, and f32 1->16 under deinterleaved=2/block_elems=1 + maps two physical uses to one E2B_B32 packet. + +negative: + source_group_stride != 1 falls back or diagnoses the E2B optimization. + +negative: + non-E2B-shaped assigned layout falls back to generic group_broadcast lowering. + +negative: + partial packet without proven safe memory read does not emit E2B. + +deferred: + E2B_B32 remains disabled until simulator/spec tests confirm the exact lane + mapping. +``` diff --git a/docs/designs/vmi-group-value-cast-broadcast-generalization.md b/docs/designs/vmi-group-value-cast-broadcast-generalization.md new file mode 100644 index 0000000000..fb5389d3a6 --- /dev/null +++ b/docs/designs/vmi-group-value-cast-broadcast-generalization.md @@ -0,0 +1,750 @@ +# VMI Group-Value Cast And Broadcast Generalization + +This note describes the generalized rule behind the +`group_reduce -> truncf -> group_broadcast` case. The goal is to avoid a +case-specific optimization such as "if f32->f16 and slots=8, do X". The +implementation should instead treat group-value layouts, width-changing casts, +and group broadcast as independent layout facts. + +This document intentionally keeps the current layout spelling: + +```text +#pto.vmi.layout +#pto.vmi.layout +``` + +Renaming the attribute syntax is separate work. + +## 1. Terms + +### 1.1 Group-Value Layout + +A group-value VMI value contains one logical scalar value per logical group. +The VMI type element count is therefore the group count: + +```text +!pto.vmi.vreg> +``` + +The layout fields mean: + +```text +G: + number of logical groups. This is redundant with the VMI vreg element count + and with the relevant group op attr, but the current IR spelling stores it. + +K: + number of logical group values placed in one physical chunk. + +LS: + physical lane distance, measured in element-sized lanes of T, between + adjacent logical group values inside a physical chunk. +``` + +For logical group `g`: + +```text +chunk = g / K +slot = g % K +lane = slot * LS +``` + +Example: + +```text +8xf32 group_values, K=8, LS=1: + f32 physical lanes 0,1,2,3,4,5,6,7 + +8xf16 group_values, K=8, LS=2: + f16 physical lanes 0,2,4,6,8,10,12,14 + +8xui8 group_values, K=8, LS=4: + ui8 physical lanes 0,4,8,12,16,20,24,28 +``` + +This is distinct from dense layouts. Dense values contain ordinary element +streams. Group-value layouts contain one scalar per group and cannot be read +by ordinary dense consumers without an explicit operation such as +`group_broadcast` or `group_store`. + +### 1.2 Group Size + +Do not confuse these two quantities: + +```text +G = num_groups = number of group scalar values +S = group_size = dense lanes per group = dense_element_count / G +K = slots = group scalar values per physical chunk +``` + +`group_reduce` consumes a dense value with group size `S` and produces a +group-value result with `G` logical lanes. The result layout's `K` describes +how those `G` scalars are physically placed. + +## 2. General Layout Rules + +### 2.1 Valid Group-Value Layout + +A concrete group-value layout is valid for a physical lowering path when: + +```text +G == vreg element count +K > 0 +LS > 0 +physical arity = ceil(G / K) +for every chunk: + active_slots = min(K, G - chunk * K) + if active_slots > 0: + (active_slots - 1) * LS < physical_lanes_per_chunk(T) +``` + +The final inequality is important. `K=8, LS=2` is valid for f16 because the +last active lane is 14 and an f16 physical vector chunk has 128 lanes. A +layout whose last active lane does not fit in one chunk must be rejected or +materialized through a different explicit layout conversion. + +### 2.2 Physical Lane Helper + +All group-value consumers should use one helper instead of open-coding slot +math: + +```text +getGroupValuePhysicalLane(layout, group): + require layout is group-value + K = layout.slots + LS = layout.lane_stride + chunk = group / K + slot = group % K + lane = slot * LS + return (chunk, lane) +``` + +For a chunk-local helper: + +```text +getGroupValueSlotPhysicalLane(layout, slot): + return slot * layout.lane_stride +``` + +This helper is the common contract for: + +```text +group_broadcast VSELR index generation +group_store packed or point stores +group-value casts +debug/validation lane maps +``` + +## 3. Cast Generalization + +Width-changing casts on group-value layouts preserve the group structure. They +do not change `G` or `K`. They only change the element type and, for +phase-zero casts, the lane stride. + +This is not the default layout decision for every cast with matching element +widths. A cast enters the group-value relation only when one connected source +or result value already carries a group-value layout fact from an independent +producer or consumer requirement: + +```text +producer fact: + group_reduce/group_slot_load/loop-carried group value gives the source a + group-value layout, and the cast derives the result layout. + +consumer fact: + group_broadcast/group_store or another group-value consumer requests a + group-value source value layout through the consumer operand overload, and + the cast derives the inverse source or result relation through the layout + propagator. +``` + +If neither side has a group-value fact, the cast must use the ordinary +type-based dense cast relation, such as the existing deinterleaved-source / +contiguous-result rule. The cast must not invent a group-value layout from +element types alone. + +The old assignment pass could reliably use only source facts that were already +known from a defining producer. The layout propagator removes that ordering +assumption by treating cast layout equations as bidirectional transfer +relations. Consumer-driven inverse derivation is therefore implemented by the +propagator rather than by a special `group_reduce -> truncf -> group_broadcast` +pattern. + +### 3.1 Narrowing + +For a narrowing cast: + +```text +source_bits = R * result_bits +``` + +the phase-zero layout relation is: + +```text +source: group_values(G, K, LS) +result: group_values(G, K, LS * R) +``` + +Examples: + +```text +f32 -> f16, R=2: + group_values(G, K=8, LS=1) + -> group_values(G, K=8, LS=2) + +i32 -> ui16, R=2: + group_values(G, K=8, LS=1) + -> group_values(G, K=8, LS=2) + +i32 -> ui8, R=4: + group_values(G, K=8, LS=1) + -> group_values(G, K=8, LS=4) +``` + +The corresponding VPTO conversion part for the phase-zero result is: + +```text +R=2: EVEN +R=4: P0 +``` + +This document only covers phase-zero layouts. Supporting odd or non-zero phase +would require an explicit lane offset/phase field, not another special case. + +### 3.2 Widening + +For a widening cast: + +```text +result_bits = R * source_bits +``` + +the inverse phase-zero relation is: + +```text +source: group_values(G, K, LS * R) +result: group_values(G, K, LS) +``` + +Examples: + +```text +f16 -> f32, R=2: + group_values(G, K=8, LS=2) + -> group_values(G, K=8, LS=1) + +ui8 -> ui32, R=4: + group_values(G, K=8, LS=4) + -> group_values(G, K=8, LS=1) +``` + +If the source layout does not have `lane_stride` divisible by `R`, the +phase-zero widening relation does not hold. The compiler should either choose +another legal layout relation or insert an explicit layout materialization if a +supported one exists. + +### 3.3 Support Query Shape + +The support query should not mention a particular op use-site such as +`group_broadcast`. It should validate the group-value cast relation: + +```text +source and result are both group-value layouts +source.G == result.G +source.K == result.K +source/result element widths define R +narrow: result.LS == source.LS * R +widen: source.LS == result.LS * R +source/result physical arity are computable and compatible with the lowering +active slots fit in each physical chunk +target has the required phase-zero conversion part +``` + +For the current main case this query succeeds because: + +```text +source: 8xf32 group_values(G=8, K=8, LS=1) +result: 8xf16 group_values(G=8, K=8, LS=2) +R = 2 +result.LS == source.LS * R +``` + +## 4. Group Broadcast Generalization + +`group_broadcast` consumes a group-value source and produces a dense result. +It should not inspect whether the source was produced by `group_reduce`, +`group_slot_load`, `truncf`, `trunci`, or an elementwise op. Its only layout +obligation is to select the physical source lane for each output group. + +For result logical lane `i`: + +```text +group = i / group_size +slot = group % K +index = slot * source_lane_stride +source_chunk = group / K +``` + +The VSELR index vector therefore depends on the source group-value layout: + +```text +source K=8, LS=1: + index lanes are 0,1,2,3,4,5,6,7 repeated by group_size + +source K=8, LS=2: + index lanes are 0,2,4,6,8,10,12,14 repeated by group_size + +source K=8, LS=4: + index lanes are 0,4,8,12,16,20,24,28 repeated by group_size +``` + +This rule is the same for contiguous and supported deinterleaved dense results. +The only difference is how result physical lanes map back to result logical +lanes before computing `group`. + +## 5. Assignment Flow + +Assignment should stabilize layouts before `vmi-to-vpto`. The generalized +flow is: + +```text +group_reduce/group_slot_load: + choose a group-value layout for the result: + unit packed plan -> K=8, LS=1 + row-local plan -> K=1, LS=1 + existing constrained layout wins if legal + +group-value narrow cast: + if the source already has a group-value layout: + set result layout with same G/K and LS multiplied by R + request the source layout + else: + use the ordinary type-based dense cast layout + +group-value widen cast: + if the source already has a group-value layout: + derive the result layout through the widening relation + else: + use the ordinary type-based dense cast layout + +group_broadcast: + request a concrete group-value source layout + set a dense result layout chosen by the consumer or by the broadcast support + policy + +group_store: + request a concrete group-value source layout compatible with row stride and + store dist support +``` + +No step should commute: + +```text +group_broadcast(truncf(x)) <-> truncf(group_broadcast(x)) +``` + +as a required legality mechanism. A separate cost optimization may still +choose such a rematerialization later, but the basic assigned layout must be +legal without changing the semantic op order. + +## 6. Lowering Flow + +### 6.1 Group-Value Narrow Cast + +For a phase-zero group-value narrow cast: + +```text +source: group_values(G,K,LS) +result: group_values(G,K,LS*R) +``` + +lower each physical source chunk independently: + +```text +active_slots = min(K, G - chunk * K) +mask = prefix mask active_slots in source element granularity +part = EVEN for R=2, P0 for R=4 +result_chunk = vcvt source_chunk, mask, part +``` + +Main case: + +```text +%sum32 : !pto.vreg<64xf32> with active lanes 0..7 +%mask = pto.pge_b32 "PAT_VL8" +%sum16 = pto.vcvt %sum32, %mask {part = "EVEN", rnd = "R", sat = "SAT"} + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<128xf16> +``` + +The resulting f16 group values are in lanes: + +```text +0,2,4,6,8,10,12,14 +``` + +### 6.2 Group Broadcast + +After the cast, broadcast uses the source layout only: + +```text +source layout: group_values(G=8,K=8,LS=2) +group_size: 16 +index vector: + [0 repeated 16, + 2 repeated 16, + 4 repeated 16, + 6 repeated 16, + 8 repeated 16, + 10 repeated 16, + 12 repeated 16, + 14 repeated 16] + +%b16 = pto.vselr %sum16, %index +``` + +This is a group-broadcast lowering rule, not a truncf-specific rule. + +### 6.3 Store + +The broadcast result is dense contiguous f16 in the main case: + +```text +pto.vsts %b16, %out[%off] {dist = "NORM_B16"} +``` + +If a later case stores group-value data directly, `group_store` should use the +same `getGroupValuePhysicalLane` helper or a target-specific store dist that is +proven equivalent to that lane map. + +## 7. Scenarios + +### 7.1 Main Floating-Point Broadcast Case + +VMI: + +```text +%sum32 = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +%sum16 = pto.vmi.truncf %sum32 +%b16 = pto.vmi.group_broadcast %sum16 {num_groups = 8} +pto.vmi.store %b16, %out[%off] +``` + +Assigned: + +```text +%sum32 : 8xf32 group_values(G=8,K=8,LS=1) +%sum16 : 8xf16 group_values(G=8,K=8,LS=2) +%b16 : 128xf16 contiguous +``` + +Lowering skeleton: + +```text +vcgadd/vadd -> vcvt EVEN PAT_VL8 -> vselr index stride 2 -> vsts NORM_B16 +``` + +### 7.2 Integer Narrow Direct Store + +VMI: + +```text +%narrow = pto.vmi.trunci %wide +pto.vmi.group_store %narrow, %out[%off], %c1 {num_groups = 8} +``` + +Assigned: + +```text +%wide : 8xi32 group_values(G=8,K=8,LS=1) +%narrow : 8xui8 group_values(G=8,K=8,LS=4) +``` + +A direct `group_store` may lower this through a target-specific packed store +such as `PK4_B32` when that store is equivalent to reading lanes +`0,4,8,...,28` in ui8 lane units. This is a store lowering strategy; it must +not redefine the group-value layout relation. + +### 7.3 Group Slot Load Then Broadcast + +VMI: + +```text +%slots = pto.vmi.group_slot_load %base[%off], %c1 {num_groups = 8} +%dense = pto.vmi.group_broadcast %slots {num_groups = 8} +``` + +Assigned: + +```text +%slots : 8xT group_values(G=8,K=8,LS=1) +%dense : dense result chosen by consumer/support +``` + +Broadcast uses VSELR indices `0..7` repeated by group size, or a memory-source +optimization may replace the whole load+broadcast chain with +`group_broadcast_load` when that transformation is legal. + +### 7.4 Widening A Lane-Strided Group Value + +VMI: + +```text +%wide = pto.vmi.extf %narrow +``` + +Legal phase-zero relation: + +```text +%narrow : Gxf16 group_values(G,K,LS=2) +%wide : Gxf32 group_values(G,K,LS=1) +``` + +If `%narrow` has `LS=1`, this particular phase-zero relation cannot produce a +packed f32 group-value result without either selecting a different layout or +materializing a layout conversion first. + +### 7.5 Unsupported Overflowing Layout + +For element type T with `L` physical lanes per chunk: + +```text +group_values(G,K,LS) +``` + +is unsupported when: + +```text +(active_slots - 1) * LS >= L +``` + +Example: + +```text +K=128, LS=2 for f16 +last lane = 254 +f16 lanes per chunk = 128 +``` + +This cannot be represented in one physical chunk with the stated `K`; assignment +must choose a smaller `K`, insert an explicit materialization, or reject the +layout if no legal route exists. + +## 8. Implementation Plan + +### 8.1 Shared Helpers + +Add or consolidate helpers with these responsibilities: + +```text +isConcreteGroupValueLayout(type/layout) +getGroupValueSlots(layout) +getGroupValueLaneStride(layout) +getGroupValuePhysicalLane(layout, group) +getGroupValueSlotPhysicalLane(layout, slot) +checkGroupValueLaneSpan(type/layout) +deriveGroupValueNarrowLayout(sourceLayout, factor) +deriveGroupValueWidenLayout(sourceLayout, factor) +``` + +These helpers should be used by support checks and lowering. Open-coded +`group % slots` lane selection should disappear from broadcast/store/cast +lowering unless the code is explicitly computing a logical slot number before +calling the physical-lane helper. + +### 8.2 Support Layer + +Update support checks so the same relation is used by floating-point and +integer casts: + +```text +truncf/trunci group-value narrow: + source/result group-value layouts + same G and K + result LS = source LS * narrow factor + +extf/extsi/extui group-value widen: + source/result group-value layouts + same G and K + source LS = result LS * widen factor + +group_broadcast: + source group-value layout with concrete K/LS + result dense layout with registered VSELR lowering support +``` + +The first implementation must explicitly enumerate the element-width pairs with +known VPTO conversion parts in the support helper. A pair not listed there is +unsupported and must fail with a diagnostic that says which relation failed, not +which high-level pattern was expected. + +### 8.3 Assignment + +Assignment should call the support-layer layout derivation instead of duplicating +cast-specific cases: + +```text +if source layout is group-value and cast is narrowing: + result layout = deriveGroupValueNarrowLayout(source layout, factor) + request source layout + +if source layout is group-value and cast is widening: + result layout = deriveGroupValueWidenLayout(source layout, factor) + request source layout +``` + +When the cast source has no known layout yet, normal producer facts still apply. +For the main case, `group_reduce` naturally produces `K=8, LS=1`; `truncf` +then derives `K=8, LS=2`. + +Backward derivation from a result consumer request is handled by the generic +request propagator. For example, if a consumer requests a group-value cast +result, the cast transfer must request the matching source layout when the +relation is unique and supported. Do not extend the old single walk with +order-dependent checks. + +### 8.4 Layout Request Propagation + +The generic request propagation design is described separately in +[vmi-layout-request-propagation.md](vmi-layout-request-propagation.md). This +document only needs the case-specific consequence: if consumer-driven +group-value casts become required, do not extend the current single walk with +more order-dependent checks. Add value layout requests to the propagator, using +the operand overload when the request comes from a specific operand, and let +them propagate through the cast relation. + +The old assignment implementation could avoid that framework change because +the main case was producer-driven: + +```text +group_reduce sets source group-value layout before truncf is visited +truncf derives its result layout immediately +``` + +When the propagator is used, the cast width relation remains generic: + +```text +if the source value is requested/propagated as group-value: + derive result group-value layout + request the result value layout + +if the result value is requested/propagated as group-value: + derive inverse source value layout + request the source value layout through the source operand overload +``` + +If the derived source layout conflicts with the source value's existing +`assignment.layout`, the operand overload records a use-side conflict in the +source value's `assignment.conflicts`. Apply materializes it with +`ensure_layout` if the layouts still differ. + +### 8.5 VMI To VPTO + +Lowering work: + +```text +group-value truncf/trunci: + use active slot prefix mask + use conversion part from the width factor + preserve per-chunk arity + +group-value extf/extsi/extui: + use active slot prefix mask + use conversion part from the width factor + preserve per-chunk arity + +group_broadcast: + build VSELR index vectors through getGroupValuePhysicalLane +``` + +The physical type policy must not make `lane_stride` mean only "integer carrier +packing". Floating-point group-value lane stride must be representable as a +normal floating-point VPTO vector with sparse active lanes, because the main +case needs: + +```text +vcvt f32 -> f16 EVEN +then vselr.f16 from the even f16 lanes +``` + +Existing target-specific integer store paths may keep using packed store dists +when they are proven equivalent to the same group-value lane map. + +## 9. Regression Tests + +Add or update focused lit tests: + +```text +test/lit/vmi/vmi_layout_assignment_group_reduce_s16_truncf_broadcast_store.pto + CHECK assignment: + group_reduce result: slots=8 + truncf result: slots=8, lane_stride=2 + group_broadcast consumes the truncf result + no f32 group_broadcast + ensure_layout + truncf shape + + CHECK lowering: + vcgadd/vadd before vcvt + vcvt before vselr + vselr before vsts + no remaining pto.vmi ops + +test/lit/vmi/vmi_to_vpto_group_broadcast_lane_stride_source.pto + Direct assigned-IR test where the source is + group_values(G=8,K=8,LS=2); CHECK vselr is generated. + +test/lit/vmi/vmi_layout_gate_group_value_cast_invalid.pto + Invalid relation, for example result LS not equal source LS * factor. +``` + +Keep existing integer lane-stride tests, but adjust their checks only if the +generalized physical type policy changes their printed VPTO shape. The semantic +expectation remains the same: the direct store must be equivalent to the +group-value lane map. + +## 10. First-Phase Required Support + +The first implementation must support the current phase-zero group-value cases. +The non-goals below are only for layouts and rewrites beyond the current +phase-zero relation. + +Required support: + +```text +phase-zero group-value layout: + lane = slot * LS + no non-zero lane offset + no odd-phase layout + +group-value producer facts: + group_reduce and group_slot_load seed their current natural group-value + layouts. + +group-value cast transfer: + narrowing and widening use the shared group-value cast equations for every + element-width pair explicitly listed by the VPTO support helper. + source-driven and result-consumer-driven propagation must both use the same + facts. + +group_broadcast: + consumes a concrete group-value source layout and produces a dense result + layout accepted by the broadcast support helper. + +group_broadcast_load / e2b optimization: + remains an optimization over the same logical group-value broadcast relation; + it must not be required to make the assigned layout legal. +``` + +## 11. Non-Goals + +This design does not include: + +```text +attr syntax rename +non-zero lane offset or odd-phase group-value layouts +cost search between broadcast-before-cast and cast-before-broadcast +automatic rematerialization as a legality mechanism +arbitrary LS values without VPTO conversion/store support +generic register compaction/expansion for every group-value layout +``` + +Those can be added later as separate design items once the phase-zero +group-value relation is implemented and covered by tests. diff --git a/docs/designs/vmi-implementation-manual.md b/docs/designs/vmi-implementation-manual.md new file mode 100644 index 0000000000..82ab6f75cc --- /dev/null +++ b/docs/designs/vmi-implementation-manual.md @@ -0,0 +1,4653 @@ +# VMI 实现手册 + +本文配套 `docs/designs/vmi-introduction.md` 和当前 VMI lowering 设计，回答 +“按什么顺序改哪些文件、每一步做到什么程度才算完成”。 + +本文不替代最终 ODS / C++ verifier / lit 测试。实现时如果发现本文和 ODS 或 verifier 冲突，以 +更精确的 verifier 约束为准，并同步刷新本文。 + +## 0. 当前仓库约束 + +当前仓库只有一个 MLIR dialect： + +```text +dialect name: pto +cpp namespace: ::mlir::pto +``` + +VPTO 低层 op/type 也在同一个 `pto` dialect 里，通过 `VPTOOps.td`、`VPTOTypeDefs.td` 等文件组织。 +因此第一版 VMI 不新建独立 dialect，采用同一 dialect 下的嵌套 mnemonic： + +```text +types: + !pto.vmi.vreg<...> + !pto.vmi.mask<...> + +attrs: + #pto.vmi.layout<...> + +ops: + pto.vmi.addf + pto.vmi.subf + pto.vmi.mulf + pto.vmi.ensure_layout +``` + +落地方式是：`PTO_Dialect` 仍是唯一 dialect，VMI 只是 `pto` dialect 内的一组 type/attr/op。 +如果后续要拆成真正独立的 `pto.vmi` dialect，必须先保证所有 pass、type converter、parser 测试 +和公开文档同步迁移；第一版不要做这个拆分。 + +风险点：带点 mnemonic 例如 `vmi.vreg`、`vmi.addf` 必须在 Slice 0 先用 parser round-trip 测试 +证明。如果 TableGen 的默认 type/attr parser 不接受该 spelling，就在 VMI type/attr 上实现 +custom assembly format，而不是改公开 spelling。 + +## 1. 文件布局 + +新增文件： + +```text +include/PTO/IR/VMIAttrs.td +include/PTO/IR/VMITypeDefs.td +include/PTO/IR/VMIOps.td +lib/PTO/IR/VMI.cpp +lib/PTO/Transforms/VMILayoutAssignment.cpp +lib/PTO/Transforms/VMIToVPTO.cpp +lib/PTO/Transforms/PTOValidateVMIIR.cpp +test/lit/vmi/ +``` + +修改文件： + +```text +include/PTO/IR/PTOAttrs.td +include/PTO/IR/PTOTypeDefs.td +include/PTO/IR/PTOOps.td +include/PTO/IR/CMakeLists.txt +lib/PTO/IR/CMakeLists.txt +include/PTO/Transforms/Passes.td +lib/PTO/Transforms/CMakeLists.txt +``` + +推荐 include 关系： + +```text +PTOAttrs.td + include "PTO/IR/VMIAttrs.td" + +PTOTypeDefs.td + include "PTO/IR/VMITypeDefs.td" + +PTOOps.td + include "PTO/IR/VMIOps.td" +``` + +放置顺序： + +```text +VMIAttrs.td: + include PTODialect.td, AttrTypeBase.td, EnumAttr.td + must not include PTOAttrs.td + +VMITypeDefs.td: + include PTODialect.td and can rely on PTOAttrs.td having included VMIAttrs.td + +VMIOps.td: + include after PTO_Op is defined in PTOOps.td + do not include VPTOOps.td from VMIOps.td +``` + +这样现有 `LLVM_TARGET_DEFINITIONS PTOOps.td` 的 TableGen 生成路径可以继续覆盖 VMI type、attr +和 op。只有当 TableGen 生成目标不能正确收集新增 td 时，才单独新增 `mlir_tablegen` 目标。 + +`lib/PTO/IR/VMI.cpp` 放 VMI type/attr/op verifier、parse/print helper 和公共 lane-map helper。 +不要把 VMI verifier 塞进 `VPTO.cpp`。 + +Pass 注册要求： + +```text +include/PTO/Transforms/Passes.td: + add VMILayoutAssignment + add VMIToVPTO + add PTOValidateVMIIR + +include/PTO/Transforms/Passes.h: + add explicit create*Pass declarations if generated declarations are not enough + +lib/PTO/Transforms/CMakeLists.txt: + add the three new .cpp files to PTOTransforms + keep DEPENDS PTOPassesIncGen and PTOOpsIncGen + add missing MLIR dialect libraries only when a new source actually includes them +``` + +The VPTO backend runs the VMI semantic pipeline by default. Use +`--enable-vmi=false` only as a temporary compatibility escape hatch. The +pipeline is ordered around vecscope inference as follows: + +```text +pto-validate-vmi-ir +vmi-layout-assignment +canonicalize/cse +vmi-layout-fold +canonicalize/cse +vmi-layout-rematerialize +canonicalize/cse +vmi-layout-sink-materialization +canonicalize/cse +vmi-legalize-arith-select +pto-validate-vmi-layout-ir +vmi-to-vpto +SIMT unroll/SCCP/canonicalize/CSE +vpto pointer/wrapper normalization +pto-infer-vpto-vecscope +VMI LICM +canonicalize/CSE +``` + +The default only applies when the effective backend is VPTO. Explicitly using +`--enable-vmi` with another backend is rejected because the pipeline produces +physical VPTO values and ops. + +The default VPTO/VMI user-facing entry also rejects public functions whose +signature contains `!pto.vmi.*`. +Internal/private VMI-typed functions are materialized at explicit boundary +helpers by baseline `vmi-layout-assignment` and physicalized by `vmi-to-vpto`. +A later optimization pass may specialize private signatures. A public VMI ABI +requires an explicit materialization plan and must not be inferred from the +layout solver. + +CLI coverage: + +```text +vmi_ptoas_cli_pipeline.pto: + --pto-backend=vpto lowers the VMI pipeline by default + pto.backend = "vpto" also selects the default VMI path + explicit --pto-backend=emitc with --enable-vmi is rejected + f16->f32 store lowers through the fold-consumers path, proving the driver + uses the optimized pipeline rather than only the hard skeleton + +vmi_ptoas_backend_required_invalid.pto: + default emitc backend with --enable-vmi and no pto.backend = "vpto" is rejected + +vmi_ptoas_public_abi_invalid.pto / vmi_ptoas_public_result_abi_invalid.pto: + public VMI argument/result signatures are rejected before layout assignment +``` + +## MLIR Framework Usage + +三个 correctness stage 和若干 layout optimization pass 不应该用同一种 MLIR 机制硬套。 +这里先定义实现框架选择，避免后续把 layout 求解、优化重写、结构化控制流改写和 1:N +physicalization 混在一个 pattern pass 里。 + +当前实现框架按下面的职责切开： + +```text +pto-validate-vmi-ir: + Operation::walk verifier。只看 IR 是否满足阶段不变量，不改 IR，不使用 conversion framework。 + +vmi-layout-assignment: + module-level per-SSA-value constraint solver。先收集等价类、producer natural layout 和 consumer request， + 再把结果写回 VMI type/helper op。它可以使用 IRRewriter 改 IR，但不以 TypeConverter 为主模型。 + +vmi-layout-fold / vmi-layout-rematerialize / vmi-layout-sink-materialization: + legal-to-legal VMI optimization passes。它们只消费 layout-assigned VMI IR，并继续产出 + layout-assigned VMI IR；所有新选择必须体现在 current op、type 或 helper IR 中。 + +vmi-legalize-arith-select: + canonicalize 之后的 hygiene pass。它把 scalar-condition arith.select with VMI result + 恢复成 VMI pipeline 可控的结构化控制流形态。 + +vmi-to-vpto: + MLIR OneToNTypeConversion。每个 layout-assigned VMI value 按统一 physical ordering 展开成多个 + VPTO value，并依靠 OneToN structural patterns 重写函数、return、region result、block argument 和 + branch operand。 +``` + +这三个 pass 的边界必须通过 IR 可见状态传递：layout 写在 `!pto.vmi.*` type 上，必要 materialization +写成 `pto.vmi.ensure_*`，physicalization 后不允许残留 `pto.vmi.*`、`!pto.vmi.*` 或 +`unrealized_conversion_cast`。不能把 layout 决策藏在 pass-private side table 里让后续 pass 猜。 + +源码级实现应该进一步拆成七个独立层次： + +```text +IR layer: + include/PTO/IR/VMIAttrs.td + include/PTO/IR/VMITypeDefs.td + include/PTO/IR/VMIOps.td + lib/PTO/IR/VMI.cpp + + 只定义语义、parse/print、type/op verifier 和公共 lane-map helper。 + 这一层不能知道 layout assignment 的全局选择，也不能直接依赖 VPTO lowering pass。 + +Semantic validation layer: + lib/PTO/Transforms/PTOValidateVMIIR.cpp + + 只检查阶段输入/输出是否满足 contract。它是 hard gate，不做 repair。 + +Layout solving layer: + lib/PTO/Transforms/VMILayoutAssignment.cpp + + 负责从 producer/consumer/control-flow/call 关系解出每个 logical value 的 layout， + 然后把结果写回 type 或 ensure_* helper。 + +Layout support query layer: + include/PTO/Transforms/VMILayoutSupport.h + lib/PTO/Transforms/VMILayoutSupport.cpp + + 只放跨阶段共享的纯查询：cast layout fact、group_reduce layout fact、 + ensure_* materialization support、layout-aware store support 等。它可以被 + assignment、validation、layout optimization 和 vmi-to-vpto 调用，但不能保存 + per-value 状态，不能返回 VPTO 指令计划，不能决定 clone/rematerialize，也不能 + 通过 producer/user/control-flow context 恢复 lowering 决策。 + + 加新 query 的标准是：至少两个阶段需要同一个语义事实，并且重复实现会导致 + assignment、validation、lowering 对同一个 layout shape 得出不同结论。只有 + 一个 lowering pattern 自己使用的分支应该留在该 pattern 内。 + +Layout optimization layer: + lib/PTO/Transforms/VMILayoutFold.cpp + lib/PTO/Transforms/VMILayoutRematerialize.cpp + lib/PTO/Transforms/VMILayoutSinkMaterialization.cpp + lib/PTO/Transforms/VMILegalizeArithSelect.cpp + + 负责在 layout-assigned VMI IR 内做 legal-to-legal 改写。它可以让公共 canonicalize/cse + 协助清理和合并 IR，但不能把决策藏到 side table 里。 + +Physicalization layer: + lib/PTO/Transforms/VMIToVPTO.cpp + + 负责把 layout-assigned VMI value 通过 OneToNTypeConversion 展成 VPTO physical values， + 并把每个 pto.vmi.* semantic op 改写成 VPTO op 序列。 + +Driver/test layer: + tools/ptoas/ptoas.cpp + tools/pto-test-opt/ + test/lit/vmi/ + + ptoas 对 VPTO backend 默认运行完整 pipeline；pto-test-opt 保留单 pass 和中间 IR 的调试入口。 +``` + +每层的 MLIR 框架选择如下： + +```text +ODS/TableGen: + 定义 type/attr/op surface 和 verifier hook。 + +Operation::walk: + 用于 validation 和 layout constraint collection。 + +Union-find + DenseMap: + 用于 layout assignment 的 per-SSA-value 等价类求解。 + +IRRewriter/RewriterBase: + 用于 layout assignment 之后的 type rewrite、helper insertion；cheap producer + rematerialization 属于后续 layout optimization pass。 + +OneToNTypeConverter + OneToNOpConversionPattern: + 只用于 vmi-to-vpto，把一个 logical VMI value 展成多个 VPTO value。 + +Upstream OneToN structural helpers: + func.func / func.call / func.return / common SCF region-result conversion。 + +Project-local OneToN structural patterns: + cf.br / cf.cond_br / cf.switch / scf.execute_region / scf.index_switch。 +``` + +不要把这些层次合并成一个万能 pattern pass。特别是： + +```text +layout assignment 不能依赖 OneToNTypeConverter: + 因为 layout 不是 type-only 决策，同一个 !pto.vmi.vreg<128xf32> 的不同 SSA value + 可能因 producer/consumer/control-flow 约束得到不同 layout。 + +vmi-to-vpto 不能重新做 layout solving: + 它只消费已经写在 type/helper 上的 layout 决策。遇到未 assignment 的 VMI type 必须失败。 + +structural OneToN pattern 不能知道 VMI 语义: + 它们只负责 flatten/rebuild operands、results、successor operands 和 block arguments。 + 具体 lane 语义只属于 pto.vmi.* op lowering pattern。 + +verifier 不能偷偷修 IR: + 否则后续 pass 会依赖 verifier 的隐式 repair 行为，导致 pipeline 顺序不可推理。 +``` + +一个可以直接对照代码的 pass 边界表： + +```text +pass input output +--------------------------- ---------------------------- ---------------------------- +pto-validate-vmi-ir surface VMI IR same IR, or hard failure +vmi-layout-assignment surface/layout-partial VMI layout-assigned VMI IR +layout optimization passes layout-assigned VMI IR layout-assigned VMI IR +vmi-legalize-arith-select layout-assigned VMI IR layout-assigned VMI IR +pto-validate-vmi-layout-ir layout-assigned VMI IR same IR, or hard failure +vmi-to-vpto layout-assigned VMI IR physical VPTO IR +final residual verifier physical VPTO candidate no pto.vmi.*, no !pto.vmi.* +``` + +### 代码级落点 + +当前实现应该能按文件直接审计。每个 pass 的核心类、MLIR 机制和失败边界如下： + +```text +lib/PTO/Transforms/PTOValidateVMIIR.cpp + pass: + PTOValidateVMIIRPass + PTOValidateVMILayoutIRPass + public helpers: + validateVMIProducerBoundaryIR + validateVMILayoutAssignedIR + MLIR API: + Operation::walk + func::FuncOp function type inspection + recursive TypeAttr / TypedAttr / ArrayAttr / DictionaryAttr scan + must not: + rewrite IR + create unrealized_conversion_cast + create ConversionTarget + repair illegal helper/type leakage + +lib/PTO/Transforms/VMILayoutAssignment.cpp + pass: + VMILayoutAssignmentPass + core object: + LayoutSolver + state: + DenseMap + SmallVector + SmallVector + SmallVector + SmallVector + MLIR API: + Operation::walk for fact collection + SymbolTable for direct internal calls + concrete cf/scf handlers for control-flow equivalence + IRRewriter/OpBuilder only after solving + must not: + use TypeConverter as the layout decision model + rewrite while collecting constraints + hide chosen layout in a pass-private side table + infer external VMI ABI + +lib/PTO/Transforms/VMILayoutFold.cpp +lib/PTO/Transforms/VMILayoutRematerialize.cpp +lib/PTO/Transforms/VMILayoutSinkMaterialization.cpp +lib/PTO/Transforms/VMILegalizeArithSelect.cpp + pass: + VMILayoutFoldPass + VMILayoutRematerializePass + VMILayoutSinkMaterializationPass + VMILegalizeArithSelectPass + role: + legal-to-legal layout-assigned VMI optimization and hygiene + MLIR API: + Operation::walk for local discovery + OpBuilder/RewriterBase for explicit IR rewrites + canonicalize/cse between passes for cleanup and deduplication + must not: + introduce physical VPTO register types + require vmi-to-vpto to inspect producers, users, or CFG + preserve optimization decisions outside IR + +lib/PTO/Transforms/VMIToVPTO.cpp + pass: + VMIToVPTOPass + converter: + VMIToVPTOTypeConverter : OneToNTypeConverter + pattern families: + OneToNOpConversionPattern for pto.vmi.* semantic ops + upstream func/scf OneToN structural patterns + project-local cf/scf structural OneToN patterns + MLIR API: + populateFuncTypeConversionPatterns + scf::populateSCFStructuralOneToNTypeConversions + applyPartialOneToNConversion + final residual walk + must not: + redo layout solving + inspect defining ops to recover physical parts + allow pto.vmi.pack/unpack/ensure_* to survive final output + allow unrealized_conversion_cast to survive final output +``` + +这里最重要的分界是：`vmi-layout-assignment` 解决的是 value-level layout，`vmi-to-vpto` +解决的是 type/value 1:N physicalization。前者的结果必须已经写回 `!pto.vmi.*` type 或显式 +`pto.vmi.ensure_*`；后者只能消费这些 IR-visible facts。 + +这也回答了“有没有充分利用 MLIR 自带能力”：结构化 1:N signature/control-flow conversion 必须用 +MLIR OneToN conversion；layout assignment 则不能强行塞进 converter，因为 converter 看不到 +producer natural layout、consumer request、CFG join 和 call-return slot 这些 value-level facts。 + +### Pass 级实现细则 + +这几个 pass 对 MLIR 自带能力的使用方式应该是“各用其长”，而不是都套成 converter pattern。 +实现时按下面的判断标准拆： + +```text +只检查阶段不变量: + 用 Operation::walk。不要创建 ConversionTarget，也不要 rewrite。 + +需要根据 SSA value、CFG join、call boundary 和 consumer request 决策 layout: + 用 module-level solver。MLIR conversion framework 没有 per-value layout 决策模型。 + +需要把一个 logical value 展成多个 physical value，并同步改 function/block/control-flow signature: + 用 OneToNTypeConversion。这里是 converter framework 最应该发挥作用的地方。 +``` + +#### Pass 框架细化 + +第一版实现按下面的源码和 MLIR infra 对齐。这个表是实现时的边界，不只是文档分层： + +```text +source file pass primary MLIR facility +----------------------------------------- --------------------------- --------------------------------------------- +lib/PTO/Transforms/PTOValidateVMIIR.cpp pto-validate-vmi-ir Operation::walk + recursive type/attr scan +lib/PTO/Transforms/PTOValidateVMIIR.cpp pto-validate-vmi-layout-ir Operation::walk + recursive type/attr scan +lib/PTO/Transforms/VMILayoutAssignment.cpp vmi-layout-assignment module-level union-find solver + IRRewriter +lib/PTO/Transforms/VMILayoutFold.cpp + vmi-layout-fold Pattern-free local IR rewrite +lib/PTO/Transforms/VMILayoutRematerialize.cpp + vmi-layout-rematerialize Pattern-free local IR rewrite +lib/PTO/Transforms/VMILayoutSinkMaterialization.cpp + vmi-layout-sink-materialization + Pattern-free local IR rewrite +lib/PTO/Transforms/VMILegalizeArithSelect.cpp + vmi-legalize-arith-select Operation::walk + OpBuilder rewrite +lib/PTO/Transforms/VMIToVPTO.cpp vmi-to-vpto OneToNTypeConverter + OneToNOpConversionPattern +``` + +这意味着每个 pass 的输入输出 contract 是固定的： + +```text +pto-validate-vmi-ir: + input: + surface VMI IR + legal: + pto.vmi semantic ops + !pto.vmi.vreg + !pto.vmi.mask + func/scf/cf structural ops carrying those types + illegal: + layout-assigned !pto.vmi.* type + physical !pto.vreg / !pto.mask / !pto.align type + pto.vmi.ensure_* / pack / unpack helper + VMI or physical type hidden in non-signature attribute + output: + exactly the same IR, or failure + +vmi-layout-assignment: + input: + verifier-clean surface VMI IR + legal work: + solve per-SSA layout/granularity constraints + rewrite VMI value/function/block types with explicit layout + insert pto.vmi.ensure_* only for use-site materialization + rematerialize cheap producers instead of inserting ensure_* when semantics are replay-safe + illegal work: + physicalize to !pto.vreg / !pto.mask + introduce pto.vmi.pack / pto.vmi.unpack + keep layout only in a pass-private side table + output: + layout-assigned VMI IR, or failure + +pto-validate-vmi-layout-ir: + input: + layout-assigned VMI IR + legal: + pto.vmi semantic ops + pto.vmi.ensure_layout / ensure_mask_layout / ensure_mask_granularity + !pto.vmi.vreg + !pto.vmi.mask + illegal: + surface !pto.vmi.vreg + surface !pto.vmi.mask + physical VPTO register types before vmi-to-vpto + pto.vmi.pack / pto.vmi.unpack + VMI or physical type hidden in non-signature attribute + output: + exactly the same IR, or failure + +vmi-to-vpto: + input: + layout-assigned VMI IR + legal work: + convert each VMI value to an ordered list of physical VPTO values + rewrite function signatures, block arguments, branch operands, region results and calls + lower pto.vmi semantic/helper ops to VPTO ops + illegal work: + infer missing layouts + change a chosen layout because one pattern finds a cheaper lowering + leave pto.vmi.* / !pto.vmi.* / unrealized_conversion_cast in final IR + output: + physical VPTO IR, or failure +``` + +`vmi-layout-assignment` 和 `vmi-to-vpto` 的关键差异是：前者解决“这个 SSA value 应该是什么 layout”， +后者解决“这个已经有 layout 的 SSA value 展开成哪些 physical value”。同一个 surface type 不能用 +`TypeConverter` 得到唯一答案： + +```mlir +%a = pto.vmi.broadcast %s : f32 -> !pto.vmi.vreg<128xf32> +%b = pto.vmi.extf %x : !pto.vmi.vreg<128xf16> -> !pto.vmi.vreg<128xf32> +%c = scf.if %cond -> !pto.vmi.vreg<128xf32> { + scf.yield %a : !pto.vmi.vreg<128xf32> +} else { + scf.yield %b : !pto.vmi.vreg<128xf32> +} +``` + +这里 `%a` 可以按 consumer 需要 rematerialize 成 contiguous 或 deinterleaved；`%b` 的 natural layout 是 +`deinterleaved=2`；`%c` 的 layout 必须由两个 yield 和后续 consumer 共同约束。这个选择依赖 Value、 +def-use、control-flow join 和 use-site request，不是 `!pto.vmi.vreg<128xf32> -> ...` 的 type-only 规则。 + +因此 layout pass 的代码形态应该固定为： + +```cpp +LogicalResult LayoutSolver::run() { + if (failed(collectAllVMIValues())) + return failure(); + if (failed(collectEquivalenceConstraints())) + return failure(); + if (failed(collectProducerNaturalLayouts())) + return failure(); + if (failed(collectConsumerRequests())) + return failure(); + if (failed(rewriteDataTypes())) + return failure(); + if (failed(insertDataUseMaterializations())) + return failure(); + if (failed(inferAndRewriteMaskTypes())) + return failure(); + if (failed(insertMaskUseMaterializations())) + return failure(); + rewriteFunctionTypesFromSolvedValues(); + return validateVMILayoutAssignedIR(module); +} +``` + +其中 `collect*` 阶段只能记录事实，不能边 walk 边改 IR。原因是控制流和 call boundary 会把后面才遇到的 +operand/result 合并到前面的 value class；边收集边改 type 会让后续约束看到混合状态，错误诊断也会依赖 +walk 顺序。 + +`vmi-to-vpto` 则必须是 converter pass。第一版使用的是 `OneToNTypeConversion`，因为它要同时处理 +value type 和结构签名： + +```text +!pto.vmi.vreg<128xf32, #pto.vmi.layout> + -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +func.func @f(%arg0: !pto.vmi.vreg<128xf32, layout>) -> !pto.vmi.vreg<128xf32, layout> + -> func.func @f(%arg0_0: !pto.vreg<64xf32>, %arg0_1: !pto.vreg<64xf32>) + -> (!pto.vreg<64xf32>, !pto.vreg<64xf32>) +``` + +这里不能用普通 1:1 `TypeConverter`，也不能靠每个 VMI op pattern 自己拆 operand。否则 `func.return`、 +`cf.br`、`scf.for` iter arg 这种没有 VMI defining op 的边界会漏转换。`OneToN` adaptor 才是 semantic +pattern 获取 physical parts 的唯一来源： + +```cpp +ValueRange lhsParts = adaptor.getLhs(); +ValueRange rhsParts = adaptor.getRhs(); +TypeRange resultTypes = adaptor.getResultMapping().getConvertedTypes(0); +``` + +结构化转换的实现分工如下： + +```text +upstream helper: + populateFuncTypeConversionPatterns + covers func.func / func.return / direct func.call signature conversion + + scf::populateSCFStructuralOneToNTypeConversions + covers common SCF result/yield/block-argument structural conversions + +project-local OneToN patterns: + cf.br + cf.cond_br + cf.switch + scf.execute_region + scf.index_switch +``` + +项目内 structural pattern 只能做结构搬运： + +```text +1. read OneToNTypeMapping for each original operand/result +2. flatten successor operands or region result types +3. rebuild the same cf/scf op with converted types +4. inline/move original regions when required +``` + +它们不能做下面这些事： + +```text +infer layout from operand defining op +emit vadd/vcvt/vlds/vsts +decide contiguous vs deinterleaved +special-case pto.vmi semantic op +``` + +VMI 语义只能出现在 `OneToNOpConversionPattern` 里。这样才能保证 block argument、function +argument、loop-carried value 和 branch target argument 都按同一套 physical ordering 转换。 + +`vmi-to-vpto` 的 legality 由 preflight + conversion + final gate 三段组成，而不是单靠 +`ConversionTarget`： + +```text +preflight: + verifyVMIToVPTOInputIR + rejects layout-free VMI types + verifySupportedVMIToVPTOOps + rejects unsupported semantic/materialization cases before rewrite starts + +conversion: + applyPartialOneToNConversion + applies structural and semantic OneToN patterns + +final gate: + verifyNoResidualVMIIR + rejects pto.vmi.* + rejects !pto.vmi.* in operand/result/block/function/attribute type trees + rejects pto.vmi.pack/unpack materialization helpers + rejects unrealized_conversion_cast +``` + +这比只设置 `ConversionTarget` 更直接，因为当前 OneToN 工具链的重点是 type/value expansion 和 pattern +rewriter；最终合法性必须递归检查 attribute/type tree，防止 VMI type 被藏在 nested attr 里。 + +#### `pto-validate-vmi-ir` / `pto-validate-vmi-layout-ir` + +这两个 pass 是 hard gate，不是 legalization pass。 + +使用的 MLIR 能力： + +```text +Operation::walk: + 遍历 module 内所有 op、region、block argument、operand/result type 和 attribute。 + +TypeAttr / TypedAttr recursive scan: + 拒绝把 VMI/physical VPTO type 藏在 nested attribute 中。 + +func::FuncOp function type special case: + function_type attr 是签名本身，可以按当前阶段规则检查；其它 attr 不能携带 VMI/physical type。 +``` + +不使用 `ConversionTarget` 的原因： + +```text +ConversionTarget 适合表达“哪些 op/type legal，哪些 pattern 能改掉”。 +这里我们只想回答“当前 IR 是否已经处在某个阶段边界”，失败后必须停机，而不是尝试 repair。 +如果 verifier 顺手改 IR，pipeline 的阶段不变量会变成隐式行为，后续 pass 很难审计。 +``` + +这两个 pass 的输出只能是原 IR 或 failure： + +```cpp +void runOnOperation() override { + if (failed(verifyStageInvariant(getOperation()))) + signalPassFailure(); +} +``` + +#### `vmi-layout-assignment` + +这个 pass 使用 MLIR 的 IR 遍历和 rewrite 基础设施，但不使用 `TypeConverter` 作为主模型。 + +核心原因： + +```text +TypeConverter 的输入是 Type。 +layout assignment 的输入是 Value。 + +同一个 !pto.vmi.vreg<128xf32> 可以因为不同 producer/consumer 关系得到不同 layout： + f16->f32 widen result -> deinterleaved=2 + f8 ->f32 widen result -> deinterleaved=4 + only contiguous store value -> contiguous +``` + +实现应拆成两个阶段，不要边 walk 边 rewrite： + +```text +collect: + 1. 收集所有 VMI data/mask SSA value 和 block argument。 + 2. 用 union-find 合并必须同 layout 的 value。 + 3. 记录 producer natural layout。 + 4. 记录 consumer layout/granularity request。 + 5. 记录 function return slot、call operand/result、branch operand/block argument 关系。 + +rewrite: + 1. 为每个 equivalence class 选 layout。 + 2. 改写 value/function/block/result type。 + 3. 对 use-site mismatch 插入 ensure_* 或 rematerialize cheap producer。 + 4. 运行 pto-validate-vmi-layout-ir。 +``` + +建议的数据结构边界： + +```cpp +struct DataNode { + Value value; + VMIVRegType type; + unsigned parent; + VMILayoutAttr naturalLayout; +}; + +struct MaskNode { + Value value; + VMIMaskType type; + unsigned parent; + VMILayoutAttr requestedLayout; + std::string requestedGranularity; +}; + +struct DataUseRequest { + OpOperand *operand; + VMILayoutAttr layout; +}; + +struct MaskUseRequest { + OpOperand *operand; + VMILayoutAttr layout; + std::string granularity; +}; +``` + +这里可以充分使用 MLIR 的接口，但它们只是 constraint source： + +```text +BranchOpInterface / concrete cf.* handlers: + successor operand[i] == destination block argument[i] + +RegionBranchOpInterface / concrete scf.* handlers: + region yield operand[i] == parent result[i] + loop init/result/iter_arg/yield 同 slot 等价 + +CallOpInterface + SymbolTable: + direct internal call operand/result 和 callee argument/return slot 等价 + external/indirect VMI call 先拒绝，因为缺 ABI materialization + +IRRewriter: + 只在 solve 完成后统一改 type、插 ensure_*、clone cheap producer。 +``` + +`vmi-layout-assignment` 的 pass invariant 是：所有 layout 决策必须写回 IR。后续 `vmi-to-vpto` +只能读取 `!pto.vmi.*` type 和显式 `pto.vmi.ensure_*`，不能依赖 layout solver 的 side table。 + +#### `vmi-to-vpto` + +这个 pass 应该充分使用 MLIR converter framework，具体是 `OneToNTypeConversion`，不是普通 +`DialectConversion`。 + +普通 1:1 dialect conversion 不够的地方： + +```text +!pto.vmi.vreg<128xf32, deinterleaved=2> + -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +!pto.vmi.vreg<256xf8, deinterleaved=4> + -> !pto.vreg<256xf8>, !pto.vreg<256xf8>, !pto.vreg<256xf8>, !pto.vreg<256xf8> +``` + +函数参数、返回值、block argument、branch operand、region result 都必须做同样的 1:N 展开。 +这正是 `OneToNTypeConverter`、`OneToNOpConversionPattern` 和结构化 OneToN helper 的职责。 + +实现骨架： + +```cpp +void runOnOperation() override { + ModuleOp module = getOperation(); + + if (failed(verifyVMIToVPTOInputIR(module)) || + failed(verifySupportedVMIToVPTOOps(module))) + return signalPassFailure(); + + VMIToVPTOTypeConverter typeConverter; + RewritePatternSet patterns(&getContext()); + + populateFuncTypeConversionPatterns(typeConverter, patterns); + scf::populateSCFStructuralOneToNTypeConversions(typeConverter, patterns); + populateProjectLocalCFOneToNPatterns(typeConverter, patterns); + populateVMISemanticOneToNPatterns(typeConverter, patterns); + + if (failed(applyPartialOneToNConversion(module, typeConverter, + std::move(patterns))) || + failed(verifyNoResidualVMIIR(module))) + signalPassFailure(); +} +``` + +`VMIToVPTOTypeConverter` 只做一种事：把 layout-assigned VMI type 映射到 canonical physical value list。 +它不能重新推导 layout。 + +```text +contiguous: + chunk0, chunk1, ... in logical order + +deinterleaved=2: + part0 chunks for logical lanes 0,2,4,... + part1 chunks for logical lanes 1,3,5,... + +deinterleaved=4: + part0 chunks for lanes 0,4,8,... + part1 chunks for lanes 1,5,9,... + part2 chunks for lanes 2,6,10,... + part3 chunks for lanes 3,7,11,... + +num_groups=G: + group-slot reduce result layout + physical storage is contiguous chunk order + only canonical group_slot(g) lanes contain semantic values +``` + +每个 semantic pattern 必须从 adaptor 拿 physical parts，不允许从 defining op 反推： + +```cpp +LogicalResult matchAndRewrite(VMIAddFOp op, OpAdaptor adaptor, + OneToNPatternRewriter &rewriter) const override { + ValueRange lhs = adaptor.getLhs(); + ValueRange rhs = adaptor.getRhs(); + TypeRange resultTypes = adaptor.getResultMapping().getConvertedTypes(0); + + if (lhs.size() != rhs.size() || lhs.size() != resultTypes.size()) + return rewriter.notifyMatchFailure(op, "physical arity mismatch"); + + SmallVector results; + for (auto [i, resultType] : llvm::enumerate(resultTypes)) { + results.push_back( + rewriter.create(op.getLoc(), resultType, lhs[i], rhs[i]) + .getResult()); + } + + rewriter.replaceOp(op, results, adaptor.getResultMapping()); + return success(); +} +``` + +这个约束对控制流是关键的：`scf.for` iter arg、branch target argument、function argument 都没有可用的 +defining op；它们的 physical parts 只能来自 OneToN signature/block argument conversion。 + +`vmi-to-vpto` 应有三层失败点，诊断不要混在一起： + +```text +preflight: + layout 未 assignment、unsupported semantic op、unsupported materialization path + +conversion: + pattern 缺失、arity mismatch、结构化控制流展开失败 + +final residual verifier: + 任何 pto.vmi.*、!pto.vmi.*、pto.vmi.pack/unpack/ensure_*、unrealized_conversion_cast 残留 +``` + +### `pto-validate-vmi-ir` + +`pto-validate-vmi-ir` 是边界 verifier，不使用 DialectConversion。 + +推荐使用： + +```text +Operation::walk +TypeSwitch / isa / dyn_cast +emitOpError / InFlightDiagnostic +SymbolTable, for function/call boundary checks +CallGraph or manual call graph collection, if recursive SCC needs diagnostics +DominanceInfo, if helper placement or resource dominance is checked +``` + +这个 pass 只检查 VMI producer boundary 和阶段不变量： + +```text +before layout assignment: + VMI data/mask values use surface type + no layout-assigned VMI type leaks in unless the test explicitly starts after assignment + no physical VPTO op appears in the semantic VMI region + no VMI helper op appears before the pass that is allowed to create it + no non-signature op/module TypeAttr or TypedAttr payload contains VMI or physical VPTO types + +after layout assignment: + pass: pto-validate-vmi-layout-ir + every VMI data value has a layout + every VMI mask has layout and concrete granularity + control-flow joins have stable type/layout + no non-signature op/module TypeAttr or TypedAttr payload contains VMI or physical VPTO types + +after VMI-to-VPTO: + no VMI op/type/helper remains + no unrealized_conversion_cast remains +``` + +不要把这个 pass 写成 rewrite pass。它可以收集 context 用于诊断，但不能通过局部修补让非法 IR +继续前进；否则后续 pass 会开始依赖 verifier 的隐式 repair 行为。 + +实现上要扫描的不只是 operand/result/block argument： + +```text +func.func function type: + 作为函数签名本身检查，允许出现当前阶段合法的 VMI type。 + +non-signature attributes: + module/op attribute 中只要递归包含 VMI type 或 physical VPTO type 都拒绝。这里包括 TypeAttr、 + TypedAttr，以及 ArrayAttr/DictionaryAttr 这类容器中的 nested attribute/type payload。 +``` + +这样可以堵住 hidden-state 形式的 side table，例如把 `!pto.vmi.vreg<...>` 偷存在 module attribute +里。`func.func` 的内建 `function_type` attr 是唯一例外，因为它只是函数签名的 MLIR 表达，不是额外 +隐藏状态。 + +### `vmi-layout-assignment` + +`vmi-layout-assignment` 不以 MLIR `TypeConverter` 作为主机制。 + +原因是 layout 选择不是单纯的 `Type -> TypeRange` 映射： + +```text +same surface type: + !pto.vmi.vreg<128xf32> + +possible per-value decisions: + value produced by f16->f32 widen: deinterleaved=2 + value loaded only for contiguous store: contiguous + value feeding fp8-like->f32 consumer path: deinterleaved=4 +``` + +两个 SSA value 可以有完全相同的 surface type，但因为 producer natural layout、consumer demand、 +控制流 join 和 target capability 不同，得到不同 layout。因此主模型应该是 per-SSA-value 的约束图， +而不是类型转换表。 + +推荐内部结构： + +```text +DenseMap +DenseMap +DenseMap +SmallVector +SmallVector +``` + +推荐使用的 MLIR 基础能力： + +```text +RegionBranchOpInterface: + collect scf.if/scf.for-like region entry, yield, result relations + +BranchOpInterface: + collect cf.br/cf.cond_br predecessor operand -> block argument relations + +CallOpInterface, CallableOpInterface, FunctionOpInterface: + collect call operand/result and function argument/result relations + +SymbolTable: + resolve direct calls and reject unresolved VMI signature assumptions + +DominanceInfo: + choose legal insertion points for ensure_layout, mask conversion, and rematerialization + +IRRewriter / RewriterBase: + rewrite types, insert helper ops, clone rematerializable producers +``` + +求解结果必须 materialize 回 IR，不能留在 side table： + +```text +1. Rewrite every VMI value type to a layout-assigned type. +2. Rewrite mask type to layout + b8/b16/b32 granularity. +3. Insert pto.vmi.ensure_layout where a consumer requires a different layout. +4. Insert pto.vmi.ensure_mask_layout / ensure_mask_granularity where predicate layout or granularity differs. +5. Clone rematerializable producers such as constant, broadcast, create_mask, iota-like producers when cheaper. +6. Re-run the VMI stage verifier. +``` + +这个 pass 可以用 `RewritePatternSet` 辅助局部 canonicalization，例如删除同 layout 的 +`ensure_layout`，但不能让 greedy pattern driver 决定全局 layout。全局约束必须先收敛，再做改写。 + +更具体地说，这里不用 `TypeConverter` 的原因不是 MLIR converter 不好用，而是此阶段的问题不是 +“一个旧 type 机械变成一个新 type”： + +```text +%a : !pto.vmi.vreg<128xf32> // 只被 contiguous store 消费 +%b : !pto.vmi.vreg<128xf32> // 来自 f16->f32 widen，后续继续 vadd +%c : !pto.vmi.vreg<128xf32> // 控制流 join，两个 predecessor 必须统一 layout +``` + +这三个 value 的 surface type 完全相同，但 layout 决策分别可能是 contiguous、deinterleaved=2、 +以及由 join 两侧约束共同决定。`TypeConverter` 看不到“这个 SSA value 的 producer/consumer/CFG +关系”，所以它只能作为后续 physicalization 的工具，不能作为 layout assignment 的主算法。 + +该 pass 对 MLIR 基础能力的使用边界是： + +```text +Operation::walk: + 收集所有 VMI SSA value、block argument、函数签名和 op transfer facts。 + +Union-find / DenseMap: + 表达必须同 layout 的 equivalence class。 + +SymbolTable: + 解析 direct internal func.call；带 VMI type 的 external/indirect call 先拒绝。 + +IRRewriter: + 改写 function/block/result type，插入 ensure_*。 + +verifyLayoutAssignedVMIIR: + pass 末尾 hard gate，确认所有决策已经 materialize 到 IR。 +``` + +### `vmi-to-vpto` + +`vmi-to-vpto` 应该使用 MLIR 的 1:N conversion framework，而不是普通 `DialectConversion`。 +这个 pass 的核心问题正是一个 logical VMI value physicalize 成多个 VPTO value： + +```text +!pto.vmi.vreg -> !pto.vreg... +!pto.vmi.mask -> !pto.mask... +``` + +普通 `DialectConversion` 的 `OpConversionPattern` 对 1:N fixed operand/result 支持不够直接： +pattern adaptor 可能拿到 source materialization，也可能拿到 flat converted operands；`func.return` +这类“一个 logical operand 展开成多个 physical operands”的场景也容易出现不完整展开。因此这里采用 +MLIR `OneToNTypeConversion` 工具： + +推荐组件： + +```text +OneToNTypeConverter +OneToNOpConversionPattern +OneToNPatternRewriter +OneToNTypeMapping +populateFuncTypeConversionPatterns +scf::populateSCFStructuralOneToNTypeConversions +applyPartialOneToNConversion +final residual verifier +``` + +`OneToNTypeConverter` 负责 layout-assigned VMI type 到 ordered physical VPTO value list： + +```cpp +typeConverter.addConversion([](VMIVRegType type, SmallVectorImpl &results) { + // Use getVMIPhysicalArity(type) and the shared lane-map helper. + // Append one physical !pto.vreg per part/chunk. +}); + +typeConverter.addConversion([](VMIMaskType type, SmallVectorImpl &results) { + // Use mask granularity and physical arity helper. + // Append one physical !pto.mask per part/chunk. +}); +``` + +source/target materialization 可以用 VMI helper 承接中间状态： + +```text +VMI value -> physical values: + pto.vmi.unpack + +physical values -> VMI value: + pto.vmi.pack +``` + +但它们只是 conversion materialization，不是最终 IR 的合法残留。final gate 必须拒绝： + +```text +pto.vmi.pack +pto.vmi.unpack +pto.vmi.ensure_layout +pto.vmi.ensure_mask_layout +pto.vmi.ensure_mask_granularity +unrealized_conversion_cast +``` + +`applyPartialOneToNConversion` 本身不是 legality framework；它负责应用 1:N patterns 并替换内部 +`unrealized_conversion_cast`。因此 `vmi-to-vpto` 必须在 conversion 后运行 final residual verifier， +把下面这些全部作为 hard failure： + +```text +any pto.vmi.* op +any !pto.vmi.* type +any pto.vmi.pack/unpack materialization helper +any pto.vmi.ensure_* helper +any unrealized_conversion_cast +``` + +结构转换必须覆盖： + +```text +func arguments/results and return operands: + use populateFuncTypeConversionPatterns + +call operands/results: + convert callee signature and call sites together + +block arguments and branch operands: + convert target block arguments and predecessor operands in the same conversion + current implementation provides project-local OneToN patterns for cf.br, + cf.cond_br, and cf.switch because MLIR only provides the generic + BranchOpInterface helper for ordinary 1:1 dialect conversion, not for VMI + 1:N physicalization. + +scf.if/scf.for region yields and results: + use scf::populateSCFStructuralOneToNTypeConversions + otherwise write explicit OneToN patterns around RegionBranchOpInterface relations +``` + +如果当前 LLVM/MLIR 版本没有提供对应 OneToN helper，就补项目内 custom `OneToNConversionPattern`。 +选择标准不是“少写代码”，而是能否正确处理 1:N result、block argument、region yield 和 +recursive/function SCC。 + +当前实现的结构转换分工如下： + +```text +upstream OneToN helper: + func.func / func.return / func.call + scf.if / scf.for / scf.while and common SCF structural cases + +project-local OneToN structural patterns: + cf.br + cf.cond_br + cf.switch + scf.execute_region + scf.index_switch +``` + +项目内 structural pattern 只做一件事：按照 `OneToNTypeMapping` 展平/重建 operand、result、 +successor operand 和 block argument。它们不能内嵌 VMI layout 语义，也不能通过 defining op +重新推导物理寄存器列表。VMI 语义只出现在各个 `pto.vmi.*` 的 `OneToNOpConversionPattern` 中。 + +OneToN conversion 的执行顺序： + +```text +1. Populate structural conversion patterns. +2. Populate VMI semantic op lowering patterns. +3. Populate helper lowering/materialization patterns. +4. applyPartialOneToNConversion on the module. +5. Run final residual verifier as the hard legality gate. +``` + +如果 conversion 或 final gate 失败，诊断必须区分： + +```text +unsupported VMI semantic op +unsupported layout materialization path +unconverted function/control-flow boundary +unexpected VMI helper residual +unexpected unrealized_conversion_cast +``` + +这样 pass 边界就是清楚的： + +```text +pto-validate-vmi-ir: + verifier/walk, no conversion + +vmi-layout-assignment: + global per-value layout solver, then IR materialization + +vmi-to-vpto: + OneToNTypeConversion-based 1:N physicalization and final legality gate +``` + +### Concrete Pass Skeleton + +整个 pipeline 按下面的 hard contract 串起来： + +```text +raw VMI producer + -> pto-validate-vmi-ir + -> vmi-layout-assignment + -> canonicalize/cse + -> vmi-layout-fold + -> canonicalize/cse + -> vmi-layout-rematerialize + -> canonicalize/cse + -> vmi-layout-sink-materialization + -> canonicalize/cse + -> vmi-legalize-arith-select + -> pto-validate-vmi-layout-ir + -> vmi-to-vpto + -> SIMT unroll/SCCP/canonicalize/CSE + -> vpto pointer/wrapper normalization + -> pto-infer-vpto-vecscope + -> VMI LICM + -> canonicalize/CSE + -> final residual verifier +``` + +The `ptoas` VPTO driver uses this sequence by default. The test-opt entry +remains useful for isolated pass debugging. Optimization after physicalization +must preserve the inferred resultless vecscope boundary. Pre-emission +canonicalization remains before inference as input normalization; VMI LICM and +the final canonicalize/CSE cleanup run after inference. + +各阶段之间只通过 IR 传递状态，不通过 pass-private side table 传递语义。也就是说： + +```text +layout assignment output: + VMI value type already contains layout + VMI mask type already contains layout + concrete b8/b16/b32 granularity + required layout conversion already appears as pto.vmi.ensure_* or rematerialized producer + +vmi-to-vpto input: + may contain pto.vmi.* semantic ops and helper ops + must not contain layout-free VMI type + function signatures and op/module TypeAttr or TypedAttr payloads are part of this invariant, + not just SSA operands/results + +vmi-to-vpto output: + must not contain pto.vmi.* op/type/helper + must not contain unrealized_conversion_cast + function type attributes and any other op/module TypeAttr or TypedAttr payloads must not contain !pto.vmi.* +``` + +This prevents a fragile design where `vmi-to-vpto` has to rediscover layout decisions from defining ops. A VMI value +may be a function argument, block argument, `scf.if` result, `scf.for` carried value, or branch target argument; none +of those has a useful defining op. + +#### Layout Assignment State + +`vmi-layout-assignment` should be implemented as one module-level solver object: + +```cpp +struct DataValueState { + Value value; + VMIVRegType surfaceType; + UnionFindNode eqClass; + VMILayoutAttr naturalLayout; // producer-preferred layout + SmallVector uses; // consumer requirements +}; + +struct MaskValueState { + Value value; + VMIMaskType surfaceType; + UnionFindNode eqClass; + VMILayoutAttr requestedLayout; + StringRef requestedGranularity; // b8/b16/b32 after inference + SmallVector uses; // consumer layout/granularity requests +}; + +struct LayoutUseRequest { + Operation *consumer; + VMILayoutAttr layout; + StringRef reason; // add/select/store/widen-source/etc. +}; +``` + +The solver runs in phases: + +```text +1. collect all VMI data/mask SSA values, including block arguments +2. add equivalence constraints +3. add producer natural-layout constraints +4. add consumer layout/granularity requests +5. solve each equivalence class +6. insert ensure_* for non-class-compatible uses +7. rewrite value types and function signatures +8. run pto-validate-vmi-layout-ir +``` + +Equivalence is only for cases where two logical values must have the same physical lane order: + +```text +add/sub/mul: + lhs == rhs == result + +cmpf/cmpi: + lhs == rhs + result mask requests lhs layout + element-width granularity + +select: + true_value == false_value == result + mask operand gets a use-site request for result layout + element-width granularity + +scf.if: + result[i] == then yield[i] == else yield[i] + +scf.for: + init_arg[i] == region_iter_arg[i] == yield[i] == result[i] + +cf.br/cf.cond_br: + successor operand[i] == successor block argument[i] + +direct internal func.call: + call operand[i] == callee argument[i] + call result[i] == all callee return operand[i] +``` + +Natural layout is not equivalence. For example: + +```text +extf f16 -> f32: + result natural layout = deinterleaved=2 + +extf f8 -> f32: + result natural layout = deinterleaved=4 + +truncf f32 -> f16: + result natural layout = contiguous + +truncf f32 -> fp8-like: + result natural layout = contiguous + +store: + consumer requests contiguous externally visible order +``` + +If one equivalence class has incompatible natural layouts, the pass must diagnose `VMI-LAYOUT-CONTRACT` unless an +explicit use-site `ensure_*` can represent the requested materialization. Baseline layout assignment does not +clone/rematerialize producers. The separate `vmi-layout-rematerialize` optimization may replace an `ensure_*` +with a cloned trivially replayable producer after the materialization request is visible in IR: + +```text +constant +broadcast +constant_mask +create_mask +``` + +For non-rematerializable producers, insert `pto.vmi.ensure_layout` immediately before the consumer that requested the +different layout. This is the conservative first implementation rule. It works for ordinary SSA values, block +arguments, loop-carried values, branch arguments, and call results because the helper is dominated by the value at the +use site and does not need to be hoisted across control flow. `DominanceInfo` may be used later to hoist duplicated +helpers as an optimization, but it must not be required for correctness in the first implementation. + +That helper is a real IR marker: if `vmi-to-vpto` cannot lower its requested conversion, the program fails with an +explicit unsupported materialization diagnostic. + +#### Layout Assignment Implementation Frame + +This pass is a normal `OperationPass`. It deliberately does not use `DialectConversion`, because there is +no stable `Type -> Type` rule until the pass has solved producer preference, consumer demand, and control-flow joins. +The implementation should look like this: + +```cpp +struct LayoutSolver { + ModuleOp module; + MLIRContext *ctx; + + DenseMap dataIds; + SmallVector dataNodes; + DenseMap maskIds; + SmallVector maskNodes; + + SmallVector dataUseRequests; + SmallVector maskUseRequests; + DenseMap> firstReturnOperandsByFunc; + + LogicalResult collectConstraints(); + LogicalResult rewriteIR(); +}; +``` + +The concrete state objects should carry only facts that are materialized back into IR: + +```cpp +struct DataNode { + Value value; + VMIVRegType surfaceType; + unsigned parent; + VMILayoutAttr naturalLayout; // null means no producer preference yet +}; + +struct MaskNode { + Value value; + VMIMaskType surfaceType; + unsigned parent; + VMILayoutAttr requestedLayout; + std::string requestedGranularity; // empty until b8/b16/b32 is known +}; + +struct DataUseRequest { + OpOperand *operand; + VMILayoutAttr layout; +}; + +struct MaskUseRequest { + OpOperand *operand; + VMILayoutAttr layout; + std::string granularity; +}; +``` + +Do not store hidden layout state that `vmi-to-vpto` must rediscover. After this pass, a debugger should be able to read +the IR and know the chosen layout for every VMI value from its type alone. + +The pass body should stay simple: + +```cpp +void runOnOperation() override { + LayoutSolver solver(getOperation()); + if (failed(solver.collectConstraints()) || + failed(solver.rewriteIR()) || + failed(verifyLayoutAssignedVMIIR(getOperation()))) + signalPassFailure(); +} +``` + +The current implementation should map directly to this phase order: + +```cpp +LogicalResult LayoutSolver::run() { + if (failed(collect())) + return failure(); + if (failed(addConstraints())) + return failure(); + + rewriteDataTypes(); + if (failed(insertDataUseMaterializations())) + return failure(); + + if (failed(inferMaskRequests())) + return failure(); + rewriteMaskTypes(); + if (failed(insertMaskUseMaterializations())) + return failure(); + + rewriteFunctionType(); + return validateVMILayoutAssignedIR(module); +} +``` + +This order is intentional: + +```text +collect: + only discovers VMI values and block arguments. + +addConstraints: + only records equivalence, natural layout and consumer request facts. + It must not rewrite IR, because later CFG/call constraints may still merge + two values that were already seen. + +rewriteDataTypes: + commits solved data layouts to !pto.vmi.vreg type. + +insertDataUseMaterializations: + repairs use-site layout mismatch after the producer's committed type is known. + +inferMaskRequests: + uses already committed data layouts and element widths to infer concrete mask + layout/granularity requests. + +rewriteMaskTypes: + commits mask layout and b8/b16/b32 granularity. + +insertMaskUseMaterializations: + repairs mask layout/granularity mismatch. + +rewriteFunctionType: + updates function signatures last, after argument/result value types have been + rewritten. +``` + +Do not move `rewriteFunctionType` before use-site materialization. A function signature is the public shape of the +solved value class; changing it early makes call/return diagnostics depend on walk order and can hide an unresolved +use-site mismatch. + +Constraint collection is a module walk with explicit handlers. The important point is that each handler only records +facts; it must not rewrite while walking: + +```text +Data equivalence: + pto.vmi.addf/addi: lhs == rhs == result + pto.vmi.cmpf/cmpi: lhs == rhs + pto.vmi.select: true_value == false_value == result + pto.vmi.ensure_layout: source and result are not equivalent if layouts differ + +Data natural layout: + pto.vmi.extf f16->f32: result natural = deinterleaved=2 + pto.vmi.extf fp8-like->f32: result natural = deinterleaved=4 + pto.vmi.truncf: result natural = contiguous + pto.vmi.channel_merge with C inputs: result natural = deinterleaved=C + +Data use request: + pto.vmi.store: value requested as contiguous + pto.vmi.channel_split with C results: source requested as deinterleaved=C + op requiring a common operand/result layout: request producer class layout + +Mask request: + cmp result: same data layout as operands, granularity from element width + select mask: same data layout as selected value, granularity from element width + store mask path: same data layout as stored value, granularity from element width +``` + +Control flow should be handled as equivalence, not as local op preference: + +```text +scf.if: + result[i] == then yield[i] == else yield[i] + +scf.for: + init_arg[i] == body iter_arg[i] == yield[i] == result[i] + +scf.while: + before argument[i] == condition forwarded operand[i] == after argument[i] + after yield[i] == result[i] + +scf.execute_region: + every nested scf.yield operand[i] == execute_region result[i] + +scf.index_switch: + every case/default yield operand[i] == index_switch result[i] + +cf.br: + operand[i] == destination block argument[i] + +cf.cond_br: + true operand[i] == true destination block argument[i] + false operand[i] == false destination block argument[i] + +cf.switch: + default operand[i] == default destination block argument[i] + case k operand[i] == case k destination block argument[i] + +func.call: + only direct internal callees are supported in the first implementation + call operand[i] == callee argument[i] + call result[i] == every corresponding callee return operand[i] +``` + +Function returns need one extra bookkeeping rule. A function result slot has one public layout in the function type, so +all `func.return` operands at the same index must be equivalent: + +```text +first return operand[i] == every later return operand[i] +function result type[i] is rewritten from the solved type of return operand[i] +call result[i] == every corresponding callee return operand[i] +``` + +If two return paths naturally produce incompatible layouts, the pass should report `VMI-LAYOUT-CONTRACT` instead of +silently choosing one path: + +```mlir +^a: + %x = pto.vmi.extf %f16 : !pto.vmi.vreg<128xf16> -> !pto.vmi.vreg<128xf32> + return %x : !pto.vmi.vreg<128xf32> // natural deinterleaved=2 + +^b: + %y = pto.vmi.extf %f8 : !pto.vmi.vreg<256xf8E4M3FN> -> !pto.vmi.vreg<256xf32> + return %y : !pto.vmi.vreg<256xf32> // different result shape/layout, invalid by verifier/type first +``` + +For equal result shape but incompatible producer preferences, the same rule applies: + +```text +return slot 0 from f16->f32 path: natural deinterleaved=2 +return slot 0 from f8E4M3FN->f32 path with the same logical result shape: natural deinterleaved=4 +diagnostic: VMI-LAYOUT-CONTRACT: conflicting natural layouts ... +``` + +External declarations with VMI types are not a layout problem; they are ABI materialization. The first implementation +must reject them before rewriting: + +```text +VMI-LAYOUT-CONTRACT: VMI typed function declaration requires an explicit external ABI materialization plan +``` + +The rewrite phase has three ordered steps: + +```text +1. Rewrite all data SSA value types to !pto.vmi.vreg. +2. Rewrite all mask SSA value types to !pto.vmi.mask. +3. Repair use-site mismatches by either rematerializing a cheap producer or inserting an explicit helper. +``` + +Rematerialization is allowed only when replaying the producer cannot change memory, control flow, or execution count +semantics: + +```text +allowed: + pto.vmi.constant splat + pto.vmi.broadcast + pto.vmi.constant_mask + pto.vmi.create_mask + +not allowed in the first implementation: + load + arithmetic result + conversion result + shuffle/channel_split/channel_merge result + value crossing a call boundary or block argument +``` + +If rematerialization is not legal, insert: + +```text +pto.vmi.ensure_layout +pto.vmi.ensure_mask_layout +pto.vmi.ensure_mask_granularity +``` + +These helpers make the unresolved materialization explicit. `vmi-layout-assignment` is allowed to create them; +`vmi-to-vpto` is responsible for proving and lowering them. If lowering cannot prove the physical transform, the final +diagnostic should be an unsupported layout/materialization diagnostic, not silent incorrect code. + +Layout assignment completion checks: + +```text +1. No surface !pto.vmi.vreg remains. +2. No surface !pto.vmi.mask remains. +3. Every VMI function argument, result, block argument, branch operand, call operand, and return operand has the + layout-assigned type selected by the solved equivalence class. +4. Every consumer-specific mismatch is represented by an explicit pto.vmi.ensure_* op immediately before that + consumer. Optional optimization passes may later replace selected helpers with rematerialized cheap producers. +5. External declarations with VMI types are rejected; they are not rewritten into an implicit ABI. +``` + +#### OneToN Conversion Details + +`vmi-to-vpto` should use MLIR `OneToNTypeConversion` for all structural rewriting that involves VMI values: + +```text +OneToNTypeConverter: + !pto.vmi.vreg -> !pto.vreg... + !pto.vmi.mask -> !pto.mask... + +Patterns: + framework structural OneToN patterns for func/return/scf + explicit OneToNOpConversionPattern for each pto.vmi semantic op + explicit helper patterns for pack/unpack/ensure_* + +Final gate: + reject residual pto.vmi.*, !pto.vmi.*, function signatures containing !pto.vmi.*, and unrealized_conversion_cast +``` + +The implementation is an `OperationPass` with this shape: + +```cpp +struct VMIToVPTOTypeConverter final : OneToNTypeConverter { + VMIToVPTOTypeConverter() { + addConversion([](Type t) { return t; }); + addConversion(convertVMIVRegType); + addConversion(convertVMIMaskType); + + TypeConverter::addSourceMaterialization(materializeVPTOToVMI); + TypeConverter::addArgumentMaterialization(materializeVPTOToVMI); + OneToNTypeConverter::addTargetMaterialization(materializeVMIToVPTO); + } +}; + +void runOnOperation() override { + ModuleOp module = getOperation(); + if (failed(verifyVMIToVPTOInputIR(module)) || + failed(verifySupportedVMIToVPTOOps(module))) + return signalPassFailure(); + + VMIToVPTOTypeConverter typeConverter; + RewritePatternSet patterns(module.getContext()); + populateVMIOneToNConversionPatterns(typeConverter, patterns); + + if (failed(applyPartialOneToNConversion(module, typeConverter, + std::move(patterns))) || + failed(verifyNoResidualVMIIR(module))) + signalPassFailure(); +} +``` + +The type converter must define one canonical physical ordering and every pattern must use that ordering: + +```text +!pto.vmi.vreg + -> chunks in logical order: + chunk0 lanes [0..P-1], chunk1 lanes [P..2P-1], ... + +!pto.vmi.vreg + -> part-major chunks: + part0 chunk0 lanes [0,2,4,...] + part0 chunk1 next even lanes + part1 chunk0 lanes [1,3,5,...] + part1 chunk1 next odd lanes + +!pto.vmi.vreg + -> part-major chunks: + part0 lanes [0,4,8,...] + part1 lanes [1,5,9,...] + part2 lanes [2,6,10,...] + part3 lanes [3,7,11,...] + +!pto.vmi.vreg + -> chunks in contiguous physical storage order + only derived group_slot(g) lanes contain semantic values + this layout is valid only for group reduce/broadcast exchange values + +!pto.vmi.mask + -> same part/chunk ordering as its data layout, one !pto.mask per physical part/chunk +``` + +`materializeVPTOToVMI` and `materializeVMIToVPTO` should use only `pto.vmi.pack` and `pto.vmi.unpack`. These ops are +conversion scaffolding; they are never valid final output. This makes accidental framework materialization visible in +the IR and easy to reject. + +Pattern population should be explicit: + +```cpp +void populateVMIOneToNConversionPatterns(VMIToVPTOTypeConverter &converter, + RewritePatternSet &patterns) { + populateFuncTypeConversionPatterns(converter, patterns); + scf::populateSCFStructuralOneToNTypeConversions(converter, patterns); + + patterns.add(converter, ctx); + + patterns.add(converter, ctx); + + patterns.add(converter, ctx); +} +``` + +Use upstream OneToN helpers where they exist: + +```text +func.func / func.return / func.call: + populateFuncTypeConversionPatterns + +scf.if / scf.for / scf.while and common structural SCF: + scf::populateSCFStructuralOneToNTypeConversions +``` + +Use project-local OneToN patterns where the current MLIR version does not provide a complete 1:N structural rewrite: + +```text +cf.br +cf.cond_br +cf.switch +scf.execute_region +scf.index_switch +``` + +These project-local structural patterns should not know VMI semantics. They only flatten operands/results according to +`OneToNTypeMapping`, convert successor block argument lists, and rebuild the same control-flow op. + +#### Pattern Authoring Checklist + +Every new `pto.vmi.*` lowering pattern should answer the same questions before it is added to +`populateVMIOneToNConversionPatterns`: + +```text +1. Does the op require all data operands/results to have identical physical arity? + If yes, check every ValueRange size against the result mapping before emitting VPTO ops. + +2. Does the op consume a mask? + If yes, the mask must already have concrete granularity and the same physical ordering expected by the data + operand. The pattern must not reinterpret a pred mask by lane count alone. + +3. Does the op observe contiguous logical order outside the register file? + If yes, require contiguous layout or explicitly lower the ensure_layout/materialization before using load/store + style VPTO ops. + +4. Does the op have padding lanes? + If yes, prove padding is unobservable. For load-like ops this requires a full-read safety proof or a fallback. + For store-like ops this requires a true predicate that disables padding writes. + +5. Does the op have target-specific side effects or ordering, such as squeeze/compact/store coupling? + If yes, put that check in verifySupportedVMIToVPTOOps before conversion starts, so the pass fails before partial + rewriting. + +6. Can it create pto.vmi.pack/unpack or unrealized_conversion_cast through framework materialization? + If yes, the semantic pattern still may be correct, but final residual verification must reject any leftover helper. +``` + +This gives a concrete division of labor: + +```text +verifySupportedVMIToVPTOOps: + shape/target/path support checks that should fail before any rewrite. + +OneToNOpConversionPattern: + mechanical lowering for a preflight-approved case. + +verifyNoResidualVMIIR: + final hard gate for missed patterns, illegal materializations and hidden VMI type payloads. +``` + +Do not put target capability probing in a structural pattern. For example, a `cf.br` pattern must never ask whether +`deinterleaved=4` can be materialized. It only converts successor operands. The semantic op that created or consumes +the value is responsible for proving the VPTO lowering path. + +#### Converter Use By Pass + +The implementation should be reviewable with the following rule: + +```text +pto-validate-vmi-ir: + no TypeConverter, no ConversionTarget, no rewrite. + +vmi-layout-assignment: + no TypeConverter for choosing layouts. + It may use RewriterBase after solving, but not DialectConversion as the solving model. + +vmi-to-vpto: + must use OneToNTypeConverter for VMI types. + must use OneToNOpConversionPattern for semantic VMI ops. + should use upstream func/scf OneToN helpers when available. + may add project-local structural OneToN patterns only for missing framework coverage. +``` + +The main reason is not style. It is correctness across values without defining ops: + +```mlir +^bb0(%x: !pto.vmi.vreg<128xf32, #pto.vmi.layout>): + cf.br ^bb1(%x : !pto.vmi.vreg<128xf32, #pto.vmi.layout>) + +^bb1(%y: !pto.vmi.vreg<128xf32, #pto.vmi.layout>): + %z = pto.vmi.addf %y, %y + : !pto.vmi.vreg<128xf32, #pto.vmi.layout> + ... +``` + +`%y` has no defining VMI op. Its physical values are the converted block arguments produced by OneToN block signature +conversion. Any implementation that tries to recover physical parts from a defining op is therefore incomplete for +control flow, function arguments and loop-carried values. + +When writing semantic `OneToNOpConversionPattern`, do not infer physical parts from a defining op. Use the OneToN +adaptor's per-original-operand `ValueRange`: + +```cpp +LogicalResult matchAndRewrite(VMIAddFOp op, OpAdaptor adaptor, + OneToNPatternRewriter &rewriter) const override { + ValueRange lhsParts = adaptor.getLhs(); + ValueRange rhsParts = adaptor.getRhs(); + TypeRange resultTypes = adaptor.getResultMapping().getConvertedTypes(0); + ... + rewriter.replaceOp(op, physicalResults, adaptor.getResultMapping()); +} +``` + +Every VMI semantic lowering then follows the same shape: + +```cpp +ValueRange lhsParts = adaptor.getLhs(); +ValueRange rhsParts = adaptor.getRhs(); +TypeRange resultTypes = adaptor.getResultMapping().getConvertedTypes(0); + +for each physical part index i: + emit physical VPTO op for lhsParts[i], rhsParts[i] -> resultTypes[i] + +replace op with all physical results using adaptor.getResultMapping() +``` + +This convention is mandatory for values crossing control flow. For example an `scf.for` iter arg has no defining op; +its physical parts are the converted block arguments created by OneToN signature conversion. + +The concrete pattern shape is: + +```cpp +LogicalResult matchAndRewrite(SourceOp op, OpAdaptor adaptor, + OneToNPatternRewriter &rewriter) const override { + ValueRange in0 = adaptor.getIn0(); + ValueRange in1 = adaptor.getIn1(); + TypeRange outTypes = adaptor.getResultMapping().getConvertedTypes(0); + + if (in0.size() != in1.size() || in0.size() != outTypes.size()) + return rewriter.notifyMatchFailure(op, "physical arity mismatch"); + + SmallVector results; + for (auto [i, outType] : llvm::enumerate(outTypes)) { + results.push_back(rewriter.create(op.getLoc(), outType, + in0[i], in1[i]).getResult()); + } + + rewriter.replaceOp(op, results, adaptor.getResultMapping()); + return success(); +} +``` + +For non-VMI operands, use a helper like `getSingleValue(op, adaptor.getOffset(), "...")` and fail if the framework +unexpectedly expanded them. This catches malformed conversion rules early. + +#### Semantic Lowering Buckets + +The first implementation should split VMI op lowering into four buckets: + +```text +identity/helper: + pack, unpack, ensure_layout identity/materialization cases, ensure_mask_* identity case + +per-part elementwise: + addf, addi, subf, subi, mulf, muli, divf, minf, maxf, negf, absf, absi, sqrt, exp, ln, relu, andi, ori, xori, shli, shrui, shrsi, not, cmpf, cmpi, select + +per-part predicate: + mask_and, mask_or, mask_xor, mask_not + +layout-producing conversion: + extf, truncf, bitcast + +externally ordered memory: + load, store + +value-indexed accumulation: + dhist, chist +``` + +Per-part elementwise ops are straightforward only when all operands/results already share the same assigned layout: + +```text +logical deinterleaved=2 value: + part0 contains logical lanes 0, 2, 4, ... + part1 contains logical lanes 1, 3, 5, ... + +vmi.addf/subf/mulf on two such values: + emit the matching VPTO per-part op for part0_lhs, part0_rhs + emit the matching VPTO per-part op for part1_lhs, part1_rhs +``` + +This preserves logical lane semantics because each physical part contains the same logical lane subset for all +operands and the result. + +Memory ops are different because their observable semantics are contiguous logical order: + +```text +vmi.store of deinterleaved=2: + cannot blindly store part0 then part1 as the final memory order + must use a store plan that writes logical lane 0,1,2,3,... order + or materialize source to contiguous before physical store +``` + +Therefore `store` lowering must either: + +```text +1. consume contiguous layout directly, or +2. lower ensure_layout(deinterleaved -> contiguous), then store, or +3. use target store instructions whose dist mode proves contiguous external order +``` + +The first implementation uses option 2 for full physical chunks: + +```text +vmi.load: + emit contiguous physical vlds chunks in memory order + materialize contiguous -> assigned result layout + +vmi.masked_load: + only when the full physical read footprint is proven safe + emit contiguous physical vlds chunks in memory order + select loaded lanes against passthru with the VMI mask + if enable-stable-gather-masked-load is set, reject pto.vmi.masked_load with + a stable TODO diagnostic until the VGATHER2-based strict no-read path is + implemented + +vmi.store: + materialize assigned source layout -> contiguous + emit physical vsts chunks in memory order +``` + +Current direct memory lowering may only emit VPTO vector memory ops for +UB-backed memory. Concretely, a `!pto.ptr<..., ub>` is legal, a +`!pto.ptr<..., gm>` is not; a memref with `#pto.address_space` is legal, +and a memref without a memory-space attribute is treated as unknown/local to +this stage to preserve existing local-view tests. A memref explicitly marked +GM or another non-VEC space is rejected by `vmi-to-vpto`. + +GM-backed VMI memory is still a valid semantic source/sink before this pass, +but direct lowering does not perform GM<->UB movement. That must be represented +by an earlier/lower memory access plan, scratch materialization, or UB view +normalization before `vmi-to-vpto`; otherwise the diagnostic is +`VMI-UNSUPPORTED` and names the GM-backed source/destination. + +For `deinterleaved=2`, `vldsx2 DINTLV_B*` and `vstsx2 INTLV_B*` are valid optimization candidates because the ISA has +an explicit two-stream de/interleave memory distribution mode. This should be implemented only as a peephole inside +`vmi-to-vpto` after the generic plan is correct: + +```text +vmi.load result layout deinterleaved=2: + vldsx2 DINTLV_B* can directly produce part0/part1 chunks + +vmi.store source layout deinterleaved=2: + vstsx2 INTLV_B* can directly store part0/part1 chunks in logical memory order +``` + +Do not generalize this to `deinterleaved=4` unless the two-level dist composition is proven against the ISA. The +fallback for `deinterleaved=4` remains generic layout materialization plus ordinary memory ops. + +Direct `vmi.load` is lowered as full VPTO physical reads when the source memory kind/layout is supported and the +element type has a known physical lane width, even for non-full logical vectors. Masked/expand/gather read-style +operations still require the lowering to prove that the full physical read footprint is safe, or to use a future +true masked/non-faulting fallback. The current proof handles: + +```text +source is a statically shaped memref +offset is a constant non-negative index +offset + physical_arity(result) * lanes_per_physical_part <= static memref element count +``` + +When this proof holds, masked/expand read-style operations may still issue full `pto.vlds` chunks. The extra padding +lanes are not logical VMI lanes and must remain unobservable through later VMI materialization rules. Pointer sources, +dynamic offsets, dynamic memrefs, and insufficient static footprints remain unsupported for those stricter read-style +operations: + +```text +VMI-UNSUPPORTED: pto.vmi. requires full physical chunks without padding lanes or a statically safe full-read +footprint (...; safe-read proof failed: ...) +VMI-UNSUPPORTED: pto.vmi. ... (source is GM-backed, but current direct VMI-to-VPTO memory lowering emits +pto.vlds/pto.vsts and requires UB-backed memory) +``` + +Store-style ops are different because inactive lanes can be made write-free with true predicates. `vmi.store`, +`vmi.masked_store` therefore support the explicit contiguous/deinterleaved tail-store +materialization paths described below. + +## 2. Slice 0: Type / Attr Bootstrap + +第一步只实现 VMI type、layout attr 和纯 helper，不实现任何 conversion pass。 + +### 2.1 `#pto.vmi.layout` + +定义 `VMILayoutAttr`： + +```mlir +#pto.vmi.layout +#pto.vmi.layout +#pto.vmi.layout +``` + +建议内部参数： + +```text +kind: enum { contiguous, deinterleaved } +factor: int64_t +``` + +Verifier： + +```text +contiguous: + factor must be 1 + +deinterleaved: + factor must be 2 or 4 +``` + +禁止接受其它 spelling，例如 `stride2`、`stride4`、`parity`、`mod_split`、`blocked`。 + +### 2.2 `!pto.vmi.vreg` + +定义 `VMIVRegType`： + +```mlir +!pto.vmi.vreg<128xf32> +!pto.vmi.vreg<128xf32, #pto.vmi.layout> +!pto.vmi.vreg<128xf32, #pto.vmi.layout> +``` + +建议参数： + +```text +elementCount: int64_t +elementType: Type +layout: Attribute // null means surface type before layout assignment +``` + +Verifier： + +```text +elementCount > 0 +elementType is scalar-like integer / float / index supported by VMI +layout is null or VMILayoutAttr +deinterleaved=4 only allowed when target registry later supports it; type verifier only checks shape +``` + +不要要求 `elementCount * bitwidth(elementType)` 是 256B 整数倍。 + +### 2.3 `!pto.vmi.mask` + +定义 `VMIMaskType`： + +```mlir +!pto.vmi.mask<128xpred> +!pto.vmi.mask<128xb32, #pto.vmi.layout> +!pto.vmi.mask<128xb32, #pto.vmi.layout> +``` + +建议参数： + +```text +elementCount: int64_t +granularity: enum/string { pred, b8, b16, b32 } +layout: Attribute +``` + +Verifier： + +```text +elementCount > 0 +surface mask may use pred and no layout +layout-assigned mask must use b8/b16/b32 and must have VMILayoutAttr +pred mask must not carry layout +``` + +### 2.4 Lane Map Helper + +在 C++ 中提供纯函数 helper，供 verifier、layout assignment、VMI-to-VPTO 和测试共用： + +```text +getDataLanesPerPart(elementType) +getMaskLanesPerPart(granularity) +getVMIPhysicalArity(type) +mapLogicalLaneToPhysical(type, logicalLane) +mapPhysicalLaneToLogical(type, part, chunk, lane) +isPaddingLane(type, part, chunk, lane) +``` + +这些 helper 是 hard dependency。任何 pass 不能重新手写一套 arity 公式。 + +Slice 0 完成条件： + +```text +1. VMI type/attr 能 parse/print round-trip。 + Covered by vmi_type_attr_parse.pto. +2. 非法 layout factor、非法 mask granularity、非法 element count 有 verifier diagnostic。 + Covered by vmi_layout_factor_invalid.pto, + vmi_mask_granularity_invalid.pto, vmi_type_element_count_invalid.pto, + and vmi_mask_concrete_without_layout_invalid.pto / + vmi_mask_pred_with_layout_invalid.pto. +3. helper 单测或 lit 测试覆盖 contiguous/deinterleaved=2/deinterleaved=4 和非整 tile。 + Covered by vmi_to_vpto_type_only.pto and + vmi_to_vpto_type_arity.pto. +``` + +## 3. Slice 1: Minimal VMI Op Set + +不要一次实现 75 个 semantic op。第一批只实现能跑通 widening + elementwise + store 的闭环。 + +### 3.1 必选 semantic op + +Construction： + +```text +pto.vmi.constant +pto.vmi.broadcast +pto.vmi.iota +pto.vmi.create_mask +pto.vmi.constant_mask +``` + +`pto.vmi.from_elements` belongs to the eventual construction surface, but it is +not part of Slice 1. Do not synthesize it from ad hoc scalar lane inserts until +there is an explicit vreg immediate, scalar-insert, or scratch materialization +contract. + +Mask： + +```text +pto.vmi.mask_and +pto.vmi.mask_or +pto.vmi.mask_xor +pto.vmi.mask_not +``` + +Arithmetic / conversion： + +```text +pto.vmi.addf +pto.vmi.addi +pto.vmi.subf +pto.vmi.subi +pto.vmi.mulf +pto.vmi.muli +pto.vmi.fma +pto.vmi.divf +pto.vmi.minf +pto.vmi.maxf +pto.vmi.negf +pto.vmi.absf +pto.vmi.absi +pto.vmi.sqrt +pto.vmi.exp +pto.vmi.ln +pto.vmi.relu +pto.vmi.andi +pto.vmi.ori +pto.vmi.xori +pto.vmi.shli +pto.vmi.shrui +pto.vmi.shrsi +pto.vmi.not +pto.vmi.cmpf +pto.vmi.cmpi +pto.vmi.select +pto.vmi.extf +pto.vmi.truncf +pto.vmi.bitcast +``` + +`pto.vmi.shrui` represents logical right shift and lowers to unsigned +`pto.vshr`. `pto.vmi.shrsi` represents arithmetic right shift and lowers to +signed `pto.vshr`; the physical element type selects the VPTO/VISA sign mode. +Integer div/rem, integer casts, int-float casts, and index casts are also +intentionally outside the current VMI surface until signedness, rounding, +saturation, overflow/remainder, and target lowering contracts are explicit. + +Memory： + +```text +pto.vmi.load +pto.vmi.masked_load +pto.vmi.gather +pto.vmi.expand_load +pto.vmi.store +pto.vmi.masked_store +pto.vmi.scatter +pto.vmi.compress_store +``` + +Value-indexed accumulation： + +```text +pto.vmi.vdhist +pto.vmi.vchist +``` + +`pto.vmi.vdhist` is a first-stage semantic op when histogram support is enabled. +`pto.vmi.vchist` may share the surface verifier, but its final lowering must be +gated until the target CHISTv2 high-range cumulative semantics are verified. + +Current implementation scope note: + +```text +pto.vmi.gather / scatter +pto.vmi.active_prefix_index / compress / compress_store +future scan / contract style ops +``` + +These families are not first-stage completion blockers. The dialect surface may +define them, and the lowering may keep narrow direct paths when the target VPTO +contract is already explicit. Full semantic coverage for these families remains +out of scope until cross-chunk state, duplicate-index ordering, prefix carry, +compaction state, or contraction accumulation contracts are explicitly designed. +Unsupported shapes must fail before OneToN rewrite with `VMI-UNSUPPORTED`; they +must not fall through to residual-op diagnostics. + +Permutation： + +```text +pto.vmi.shuffle +pto.vmi.channel_split +pto.vmi.channel_merge +``` + +Internal helper： + +```text +pto.vmi.ensure_layout +pto.vmi.ensure_mask_layout +pto.vmi.ensure_mask_granularity +pto.vmi.unpack +pto.vmi.pack +``` + +### 3.2 Op Verifier Rules + +Construction op verifier: + +```text +constant value must be a dense elements attr, and its element type/count must match the result vreg +broadcast scalar type must match the result element type +constant_mask value must be a dense elements attr, must have i1 element type, and its element count must match the +result mask +create_mask may produce surface pred mask or concrete layout-assigned mask +mask_and/mask_or/mask_xor/mask_not require all mask operands/results to have the same logical lane count; if any +mask is layout-assigned, all masks must carry the same layout and granularity +``` + +Elementwise op verifier： + +```text +all data operands have same logical lane count +all data operands have same element type except documented conversion op +if any operand has layout, all layouted operands/results must agree +surface op may have no layout before vmi-layout-assignment +``` + +`select` verifier： + +```text +mask lane count == true/false/result lane count +mask layout must match data layout after layout assignment +mask granularity must match selected element width after layout assignment +``` + +`extf/truncf` verifier： + +```text +source/result lane count equal +source/result element types are float +bitwidth changes in the expected direction +truncf rounding attr, when present, must be A/H and currently only applies to + f32 -> !pto.hif8 +``` + +Memory op verifier: + +```text +load memory element type must match result VMI data element type when the source is PtrType or MemRefType +store memory element type must match stored VMI data element type when the destination is PtrType or MemRefType +``` + +Histogram op verifier: + +```text +dhist/chist acc type must be !pto.vmi.vreg<256xui16> +dhist/chist result type must match acc type +source type must be !pto.vmi.vreg +mask logical lane count must match source logical lane count +surface mask may be pred; after layout assignment it must be b8 contiguous +source/result/acc must not carry layout before vmi-layout-assignment +layout-assigned dhist/chist requires contiguous source, mask, acc, and result +``` + +`shuffle` verifier： + +```text +static mask length == result lane count +each mask index selects an existing source logical lane +result element type == source element type +no padding lane may be selected +``` + +`channel_split` verifier： + +```text +result count C >= 2 +input lane count N == C * M +each result is vreg +channel c result semantics: out[c][i] = input[i * C + c] +if any source/result carries layout, all must carry layout +for C=2/4, layout-assigned source must be contiguous or deinterleaved=C +layout-assigned results must be contiguous +``` + +`channel_merge` verifier： + +```text +operand count C >= 2 +all operands have same M and element type T +result is vreg +result semantics: result[i * C + c] = input[c][i] +if any input/result carries layout, all must carry layout +layout-assigned inputs must be contiguous +for C=2/4, layout-assigned result must be contiguous or deinterleaved=C +``` + +`ensure_layout` verifier： + +```text +source/result are both VMIVRegType +same elementCount and elementType +source/result both layout-assigned +source layout may equal result layout; that is a canonical no-op +``` + +`ensure_mask_layout` verifier is identical except it uses `VMIMaskType` and preserves granularity. + +`ensure_mask_granularity` verifier： + +```text +source/result are both VMIMaskType +same elementCount +same layout +source/result granularity are b8/b16/b32 +logical predicate value must be preserved +``` + +`pack/unpack` verifier： + +```text +VMI side must be layout-assigned +physical operand/result count == getVMIPhysicalArity(VMI type) +physical data types are !pto.vreg +physical mask types are !pto.mask +ordering is the shared Physical Arity helper order +``` + +Slice 1 完成条件： + +```text +1. Every Slice 1 op parses, prints, and has negative verifier tests. + Arithmetic/mask/helper verifier coverage includes vmi_elementwise_kind_invalid.pto, + vmi_mask_logic_invalid.pto, vmi_ensure_layout_surface_invalid.pto, + vmi_unpack_arity_invalid.pto, and vmi_pack_arity_invalid.pto. +2. Helper ops are marked internal in docs and rejected by final VMI-to-VPTO gate if residual. +3. `channel_split/channel_merge` have tests proving shuffle-equivalent lane order. +``` + +## 4. Slice 2: VMI Producer Boundary Verifier + +VMI core implementation starts from VMI IR. Producer-specific import is outside this manual's core path. + +实现 `PTOValidateVMIIR.cpp` 中的 VMI boundary verifier： + +```text +recommended pass name: pto-validate-vmi-ir +anchor: func::FuncOp or ModuleOp +source file: lib/PTO/Transforms/PTOValidateVMIIR.cpp +``` + +Boundary verifier checks: + +```text +all logical vector values use !pto.vmi.vreg / !pto.vmi.mask +all logical vector behavior is represented by pto.vmi semantic ops +surface VMI values before layout assignment do not carry layout +no physical VPTO op appears before vmi-to-vpto +no hidden side table is required to interpret VMI values +scalar/tensor/debug/transform boundary has already been resolved by producer +``` + +Slice 2 完成条件： + +```text +1. VMI-native positive tests pass boundary verification. + Covered by vmi_producer_boundary_valid.pto. +2. Physical VPTO op before VMI-to-VPTO is rejected. + Covered by vmi_producer_boundary_physical_invalid.pto, including both + physical function types and physical VPTO ops. +3. Layout-assigned type before layout assignment is rejected unless the test explicitly starts after layout assignment. + Covered by vmi_producer_boundary_layout_invalid.pto and + vmi_producer_boundary_mask_layout_invalid.pto. +4. Missing VMI type/op invariants produce `VMI-PASS-INVARIANT` or a more specific diagnostic. + Covered by vmi_producer_boundary_non_vmi_op_invalid.pto, + vmi_producer_boundary_helper_invalid.pto, and the producer-boundary + TypeAttr nested/surface/layout invalid tests. +``` + +## 5. Slice 3: `vmi-layout-assignment` + +推荐实现为 pass： + +```text +recommended pass name: vmi-layout-assignment +anchor: ModuleOp +source file: lib/PTO/Transforms/VMILayoutAssignment.cpp +``` + +`vmi-layout-assignment` 必须是 module 级 pass。函数参数、`func.return` operand、 +`func.call` operand/result 和 callee signature 需要在同一个约束图里求解；函数级 pass +只能看到局部 body，无法安全地同步 callsite 和 callee。 + +### 5.1 Internal Data Model + +Build one layout node per VMI SSA value: + +```text +Operation result +BlockArgument +Region yield operand +Function argument/result +Call operand/result +``` + +Each node records: + +```text +logical type: VMIVRegType or VMIMaskType +allowed layouts: bitset {contiguous, deinterleaved2, deinterleaved4} +required mask granularity: pred/b8/b16/b32 or unknown +natural layout preference +hard constraints +``` + +No information required by later passes may live only in this data structure. After the pass, type/attr/op +operands must fully describe the result. + +### 5.2 Transfer Functions + +Minimum Slice 3 transfer functions： + +```text +constant/broadcast/create_mask/constant_mask: + rematerializable in any legal consumer layout + +mask_and/mask_or/mask_xor/mask_not: + all mask operands/results same layout and granularity + +addf/addi/subf/subi/mulf/muli/divf/minf/maxf/negf/absf/absi/sqrt/exp/ln/relu/andi/ori/xori/shli/shrui/shrsi/not/cmpf/cmpi/select: + all data operands/results same layout + mask layout follows data layout + +extf f16 -> f32: + result natural layout = deinterleaved=2 + source requires contiguous layout for the direct vcvt part=EVEN/ODD path + partial/tail source chunks are supported when they still fit in one physical + source chunk and produce the natural two-part result; source padding lanes map + only to result padding lanes + +extf f8 -> f32: + result natural layout = deinterleaved=4 + source requires contiguous layout for the direct vcvt part=P0/P1/P2/P3 path + partial/tail source chunks are supported under the same one-source-chunk + contract; source padding lanes map only to result padding lanes + +truncf f32 -> f16: + can consume deinterleaved=2 and produce contiguous + current implementation records a deinterleaved=2 source use-site request and + inserts pto.vmi.ensure_layout when the source value solved to contiguous. + partial/tail source pairs are supported when the two deinterleaved source + parts pack into one contiguous result chunk; source padding lanes map only to + result padding lanes + +truncf f32 -> fp8-like: + can consume deinterleaved=4 and produce contiguous + current implementation records a deinterleaved=4 source use-site request and + inserts pto.vmi.ensure_layout when the source value solved to contiguous. + The lowering emits four pto.vcvt operations with part=P0/P1/P2/P3, then ORs + the mutually exclusive partial destination registers into one contiguous fp8 + result. This mirrors the hardware packed-4 contract: each source part owns + one quarter of the destination byte lanes, so the final externally visible + vector remains logical lane order 0..N-1 after the merge. + default round mode is result-type specific: f8E4M3/f8E5M2 use rnd=R, hif8 + uses rnd=A. hif8 may explicitly request hybrid lowering with + pto.vmi.truncf {rounding = "H"}, which forwards rnd=H to every packed part. + +bitcast: + source and result layouts must match + source/result total logical bits must match + current implementation supports contiguous/deinterleaved layouts with identical + physical arity when every source/result physical chunk carries the same number + of logical bits. This covers full chunks and partial/tail chunks such as + 65xf32 -> 130xi16, where the second physical chunk carries 32 logical bits on + both sides, and uneven deinterleaved tails such as 129xf32 -> 129xi32. + Partial/tail bitcast remains unsupported if source padding bits would become + result logical bits. group_slots bitcast follows the same rule: it is valid + only when the source/result group_slots layout is identical and every + physical group-slot chunk carries the same logical bit footprint. + +load: + baseline result layout is deterministic from explicit layout attrs or the + producer natural layout; consumer-specific alternatives are represented by + ensure_layout and optimized later + +store: + baseline requests contiguous source layout + current implementation records a contiguous use-site request for vmi.store and + inserts pto.vmi.ensure_layout when the stored value class solved to a + non-contiguous layout. This makes externally visible memory order explicit in + IR before vmi-to-vpto. If explicit IR reaches vmi-to-vpto with a + deinterleaved=2/4 tail value, the direct lowering may still materialize it to + contiguous physical chunks first, but only when every deinterleaved part has + the same physical chunk count and therefore forms complete intlv groups. + +shuffle/channel_split/channel_merge: + default result layout contiguous unless the current op explicitly carries a + supported layout-preserving contract + current implementation supports pto.vmi.shuffle when every result physical + chunk forwards one source physical chunk with identical lane positions for + all non-padding result lanes. Result padding lanes are ignored by the + forwarding proof and remain unobservable after physicalization. This allows + whole-chunk projection/reordering under contiguous or explicit deinterleaved + layouts, including tail-prefix projections such as `[0, 1, 2, 3] -> + !pto.vmi.vreg<4xf32>`. Arbitrary lane permutation remains unsupported unless + the vselr index-vector path below can materialize it. + current implementation supports channel_split/channel_merge for 2 or 4 + channels. channel_split consumes a natural deinterleaved=C source and produces + contiguous per-channel results; channel_merge consumes contiguous per-channel + inputs and produces a natural deinterleaved=C result. The direct path also + accepts partial/tail channel groups when the virtual deinterleaved=C channel + layout has the same physical arity as the source/result representation, so + every physical group can be materialized with complete intlv/dintlv pairs. + Arity-changing partial groups such as splitting 4xf32 into two 2xf32 channels + remain unsupported. If a producer/consumer + requires dense contiguous layout, pto.vmi.ensure_layout materializes the + pto.vdintlv/pto.vintlv tree explicitly. Non-matching layouts and other channel + counts remain unsupported. +``` + +### 5.3 Solver Order + +Implement deterministic solving: + +```text +1. Collect region/SCC constraints, including scf/cf/function/call boundaries. +2. Propagate impossible layouts and required mask granularities. +3. Pick one layout per node using deterministic priority, not a cost model: + explicit layout already present on the VMI type, then unique natural layout, + then hard non-contiguous request, then contiguous. +5. Rewrite result/block/function types to layout-assigned VMI types. +6. Insert ensure_layout / ensure_mask_layout / ensure_mask_granularity at uses that need conversion. +7. Run verifier gate. +``` + +Current implementation status: + +```text +implemented: + extf source -> contiguous use-site request for supported f16/fp8-like to f32 paths + truncf f32->f16 source -> deinterleaved=2 use-site request + truncf f32->fp8-like source -> deinterleaved=4 use-site request + single-use pto.vmi.load results can adopt a consumer-requested + layout before type rewrite; this covers direct memory producers such as + load -> truncf without inserting a redundant ensure_layout + vmi.store data operand -> contiguous use-site request + explicit VMI vreg layout is preserved as an initial solver constraint + explicit concrete VMI mask layout/granularity is preserved as an initial solver constraint + channel_split source -> deinterleaved=C use-site request + channel_split results -> contiguous natural layout + channel_merge inputs -> contiguous use-site request + channel_merge result -> deinterleaved=C natural layout + shuffle without explicit layouts -> contiguous source use-site request and contiguous result natural layout + shuffle with explicit source/result layouts -> preserve explicit layouts and let vmi-to-vpto prove chunk forwarding + pto.vmi.ensure_layout insertion for non-contiguous store operands + pto.vmi.ensure_layout insertion for truncf source materialization + pto.vmi.ensure_mask_layout / ensure_mask_granularity insertion for select mask operands + pto.vmi.create_mask / constant_mask rematerialization for select mask operands when the consumer needs a + different mask layout/granularity + splat pto.vmi.constant rematerialization for data operands when the consumer needs + a different layout + pto.vmi.broadcast rematerialization for data operands when the consumer needs + a different layout + scf.execute_region result/yield layout equivalence + scf.index_switch result/yield layout equivalence + scf.while state layout equivalence + +not yet implemented: + generic per-consumer layout request table for every VMI op + producer rematerialization for non-splat data constants and other cheap producers + cost model / target capability registry +``` + +Do not implement a local greedy pattern pass that ignores block arguments or function signatures. + +### 5.4 CFG Rules + +CFG 处理分两层。第一层是必须做的 layout equivalence：同一个控制流值在 +result、yield、region/block argument 之间必须形成同一个 layout/mask 约束组。第二层才是 +layout conflict resolution：当同一个 producer 的不同 consumers 希望不同 layout 时，插入 +`ensure_layout` 或 `ensure_mask_layout`。后续 `vmi-layout-rematerialize` 可以把部分 helper +替换成重放的纯构造 producer。 + +当前可落地的最小实现先做第一层。它不尝试在 branch 边界自动插入 conversion，因此下面这些 +关系一旦因为 natural layout 或 mask granularity 冲突无法合并，必须报 `VMI-LAYOUT-CONTRACT`， +不能默默选择某一边。 + +`scf.if` equivalence： + +```text +for each result index i: + scf.if result[i] + == then scf.yield operand[i] + == else scf.yield operand[i] +``` + +如果 value 是 `!pto.vmi.vreg`，合并 data layout 约束；如果 value 是 +`!pto.vmi.mask`，合并 mask layout 和 granularity 请求。这样 `%m = scf.if ... -> +!pto.vmi.mask` 后被 `vmi.select` 消费时，select 对 `%m` 推出的 `b8/b16/b32 + layout` +会传播回两边 yield 的 mask producer。 + +`scf.for` equivalence： + +```text +for each iter_arg index i: + init_arg[i] + == region_iter_arg[i] + == scf.yield operand[i] + == scf.for result[i] +``` + +这条规则避免 loop-carried value 每次迭代改变 layout。对于 `extf f16->f32` 作为 init、 +loop body 内部 `addf` 并 yield 的 case，`extf` 的 natural layout `deinterleaved=2` +必须稳定传递到 `%acc` region arg、`scf.yield` 和 loop result。 + +`cf.br` / `cf.cond_br` equivalence： + +```text +for each successor operand index i: + branch successor operand[i] + == successor block argument[i] +``` + +当前实现覆盖标准 `cf.br`、`cf.cond_br` 和 `cf.switch`。其中 `cf.switch` 的 default operands +与 default destination block arguments 按 index 建 layout 等价关系；每个 case operand segment +与对应 case destination block arguments 按 index 建 layout 等价关系。更泛化的 +`BranchOpInterface` op 如果携带 VMI type，后续要么补对应 mapping，要么在 layout assignment +阶段明确 diagnostic，不能让 hidden default layout 穿过去。 + +当前实现支持携带 VMI value 的 `scf.execute_region`：execute_region result 与直属 region terminator +`scf.yield` operands 按 result index 合并到同一个 layout 等价类。嵌套 region 内属于其他 op 的 +`scf.yield` 不参与 execute_region 的等价关系。 + +当前实现支持携带 VMI value 的 `scf.index_switch`：default/case region `scf.yield` operands 与 +index_switch results 按 result index 合并到同一个 layout 等价类。 + +当前实现支持携带 VMI value 的 `scf.while`：init operand、before region argument、`scf.condition` +forwarded operand、after region argument、after region `scf.yield` operand 和 while result 按状态 +index 合并到同一个 layout 等价类。`scf.condition` 的 i1 condition 本身不参与 VMI layout 约束。 + +Function boundary： + +```text +internal functions may get specialized layouted signatures +external ABI must not expose VMI layout +recursive SCC requires fixed-point signature layout +``` + +当前实现支持 direct `func.call` 到同一 module 内带 body 的 `func.func`： + +```text +call operand[i] == callee argument[i] +call result[i] == every callee return operand[i] +same-result-index return operands inside one callee are equivalent +``` + +如果携带 VMI type 的 call 无法解析到带 body 的 direct callee，layout assignment 必须报 +`VMI-LAYOUT-CONTRACT`。后续如需支持 public/external ABI，必须先定义 VMI 值如何在 ABI +边界 materialize，不能把 layouted VMI type 暴露出去。 +当前实现明确拒绝携带 VMI type 的 `func.call_indirect`，因为它没有可解析的 direct internal +callee signature/body 可参与 layout constraint solving。 + +当前实现对携带 VMI type 的 external function declaration 报 `VMI-LAYOUT-CONTRACT`，因为还没有 +定义 VMI value 的外部 ABI materialization plan。没有 VMI type 的 external declaration 必须在 +`rewriteFunctionType` 中保持原签名，不能因为没有 entry block arguments 被改写成空签名。 + +`ptoas` 的默认 VPTO/VMI pipeline 拒绝 public `func.func` 的 VMI-typed signature： + +```text +VMI-LAYOUT-CONTRACT: public VMI typed function requires an explicit external ABI materialization plan +``` + +这样 test-opt 仍可覆盖 internal/private function signature physicalization，用户入口则不会把 +layout-assigned VMI 值隐式暴露成 public ABI。 + +Slice 3 完成条件： + +```text +1. All VMI values have layout-assigned types after the pass. +2. All masks have b8/b16/b32 granularity after the pass. +3. CFG and call tests prove branch/yield/signature layout equality. +4. Multi-use rematerializable producer tests prove broadcast, constant, iota, + create_mask, and constant_mask rematerialization vs ensure_layout / + ensure_mask_* is deterministic. +5. The pass runs the layout-assigned VMI hard gate before returning, including + recursive TypeAttr/TypedAttr rejection; covered by + vmi_layout_assignment_post_gate_type_attr_invalid.pto. +``` + +## 6. Slice 4: `vmi-to-vpto` + +推荐实现为 pass： + +```text +recommended pass name: vmi-to-vpto +anchor: ModuleOp +source file: lib/PTO/Transforms/VMIToVPTO.cpp +``` + +第一步实现必须先落地 MLIR OneToN conversion 框架： + +```text +VMIToVPTOTypeConverter : OneToNTypeConverter: + !pto.vmi.vreg -> ordered !pto.vreg list + !pto.vmi.mask -> ordered !pto.mask list + +Structural patterns: + populateFuncTypeConversionPatterns + scf::populateSCFStructuralOneToNTypeConversions + project-local OneToN patterns for cf.br/cf.cond_br/cf.switch + project-local OneToN patterns for scf.execute_region/scf.index_switch + +VMI patterns: + OneToNOpConversionPattern for pack/unpack/ensure_*/semantic ops + +Final residual gate: + reject pto.vmi.*, !pto.vmi.*, unrealized_conversion_cast + scan SSA types, block argument types, function signatures, and op/module TypeAttr or TypedAttr payloads +``` + +这一步可以先支持 type-only physicalization 和 `pack/unpack` helper physicalization，但不能让未实现的 VMI semantic op 静默通过。 +如果还有 `pto.vmi.*` 或 VMI type 残留，必须报 `VMI-RESIDUAL-OP`。 + +当前 slice 支持 VMI function/input/block argument 展开成 physical arguments，并支持： + +```text +pto.vmi.unpack(layouted VMI aggregate) -> physical parts: + replace with OneToN adaptor source parts + +pto.vmi.pack(physical parts) -> layouted VMI aggregate: + replace with the physical parts through resultMapping + +pto.vmi.ensure_layout / ensure_mask_layout / ensure_mask_granularity: + ensure_layout must compare the original VMI source/result layout attrs, not only the converted physical type list. + If source/result layouts are identical, replace with source parts. This identity case supports partial/tail physical + chunks because no lane reordering or packing is performed. + If deinterleaved=2 -> contiguous, emit one pto.vintlv. + If contiguous -> deinterleaved=2, emit one pto.vdintlv. + If deinterleaved=4 -> contiguous, emit the two-level pto.vintlv tree. + If contiguous -> deinterleaved=4, emit the reverse two-level pto.vdintlv tree. + ensure_mask_layout supports the same contiguous <-> deinterleaved=2/4 layout conversions with predicate + rearrange ops: + deinterleaved=2 -> contiguous: pto.pintlv_b8/b16/b32 + contiguous -> deinterleaved=2: pto.pdintlv_b8/b16/b32 + deinterleaved=4 -> contiguous: two-level pto.pintlv_b8/b16/b32 tree + contiguous -> deinterleaved=4: two-level pto.pdintlv_b8/b16/b32 tree + ensure_mask_granularity supports concrete b8/b16/b32 logical predicate-preserving conversion: + widening b8 -> b16 -> b32: split each physical chunk with pto.punpack LOWER/HIGHER + narrowing b32 -> b16 -> b8: pack physical chunk pairs with pto.ppack LOWER/HIGHER and merge halves with pto.por + b8 <-> b32 conversions are lowered as two adjacent steps through b16. + +pto.vmi.broadcast: + current direct lowering requires the physical result element width to be 8, + 16, or 32 bits, because the vdup is predicated by pto.mask. + Other semantic element types need a dedicated materialization contract before + vmi-to-vpto may lower them. + for each physical result part: + materialize pto.pset_b8/b16/b32 "PAT_ALL" from the physical result element width + emit pto.vdup(scalar, all_true_mask) + This is layout-independent because every logical lane has the same scalar value. A deinterleaved layout simply + receives one identical vdup per partition/chunk; no vintlv/vdintlv is needed. + +pto.vmi.iota: + semantics: + ASC: result[lane] = base + lane + DESC: result[lane] = base - lane + supported element types follow pto.vci: + integer 8/16/32 and f16/f32 + contiguous full-chunk direct path: + for each physical chunk c: + chunk_base = base +/- c * lanes_per_part + emit pto.vci chunk_base {order = ASC|DESC} + deinterleaved layout requires strided index materialization because physical part p contains logical lanes: + p, p + factor, p + 2 * factor, ... + The required formula is: + ASC: base + p + factor * local_lane + DESC: base - p - factor * local_lane + The current lowering materializes this per physical chunk: + local = pto.vci 0 + scaled = pto.vmuls local, factor + ASC: result = pto.vadds scaled, base + part_offset + DESC: result = pto.vsub pto.vdup(base - part_offset), scaled + Partial/tail chunks are allowed. The physical padding lanes receive the natural continuation of the generated iota + sequence and remain padding/undef at the VMI semantic level; memory writes, masks, reductions, and other + externally-visible consumers must still obey the VMI padding rules. + +pto.vmi.constant_mask: + support dense bool constants for concrete b8/b16/b32 masks. For each physical chunk: + if the active lanes form a prefix: + emit pto.pset_b8/b16/b32 PAT_ALL, PAT_ALLF, or supported PAT_VL* + if a prefix count has no supported PAT_VL token, fall back to pto.plt_b8/b16/b32 with a constant i32 count + otherwise decompose the static bitset into active runs: + run [lo, hi) = prefix(hi) & ~prefix(lo) + combine runs with pto.por under an all-true predicate + pred-only masks remain unsupported until they have a concrete b8/b16/b32 consumer granularity. + +pto.vmi.mask_and / mask_or / mask_xor / mask_not: + for each physical predicate part: + materialize pto.pset_b8/b16/b32 "PAT_ALL" from the physical mask granularity + mask_and emits pto.pand(lhs_part, rhs_part, all_true_mask) + mask_or emits pto.por(lhs_part, rhs_part, all_true_mask) + mask_xor emits pto.pxor(lhs_part, rhs_part, all_true_mask) + mask_not emits pto.pnot(source_part, all_true_mask) + +pto.vmi.addf / addi / subf / subi / mulf / muli / divf / minf / maxf / negf / absf / absi / sqrt / exp / ln / relu / andi / ori / xori / shli / shrui / shrsi / not: + current direct lowering requires the physical element width to be 8, 16, or + 32 bits, because every emitted VPTO op is predicated by a materialized + pto.mask. VMI types such as index or f64 remain valid semantic + surface types only after a dedicated lowering contract exists; until then + vmi-to-vpto must report VMI-UNSUPPORTED before OneToN conversion. + This common predicate-maskability rule is necessary but not sufficient for + every target op. Direct lowering must also preflight the concrete VPTO/VISA + element contract before OneToN rewriting: + addf/subf/mulf -> pto.vadd/vsub/vmul support f16/bf16/f32 floating types + divf -> pto.vdiv supports f16/f32 floating types + minf/maxf -> pto.vmin/vmax support f16/bf16/f32 floating types + negf/absf/sqrt/exp/ln/relu -> pto.vneg/vabs/vsqrt/vexp/vln/vrelu support f16/f32 floating types + absi -> pto.vabs supports signless/signed i8/i16/i32 integer types + bf16/f8 remain legal VMI float-like semantic types for the ops whose VMI + semantics allow them, but vmi-to-vpto must report VMI-UNSUPPORTED until a + materialization plan or wider target contract exists. + for each physical part: + materialize pto.pset_b8/b16/b32 "PAT_ALL" from the physical element width + addf/addi emit pto.vadd(lhs_part, rhs_part, all_true_mask) + subf/subi emit pto.vsub(lhs_part, rhs_part, all_true_mask) + mulf/muli emit pto.vmul(lhs_part, rhs_part, all_true_mask) + divf emits pto.vdiv(lhs_part, rhs_part, all_true_mask) + minf emits pto.vmin(lhs_part, rhs_part, all_true_mask) + maxf emits pto.vmax(lhs_part, rhs_part, all_true_mask) + negf emits pto.vneg(source_part, all_true_mask) + absf/absi emit pto.vabs(source_part, all_true_mask) + sqrt emits pto.vsqrt(source_part, all_true_mask) + exp emits pto.vexp(source_part, all_true_mask) + ln emits pto.vln(source_part, all_true_mask) + relu emits pto.vrelu(source_part, all_true_mask) + andi emits pto.vand(lhs_part, rhs_part, all_true_mask) + ori emits pto.vor(lhs_part, rhs_part, all_true_mask) + xori emits pto.vxor(lhs_part, rhs_part, all_true_mask) + shli emits pto.vshl(lhs_part, rhs_part, all_true_mask) + shrui emits pto.vshr(lhs_part, rhs_part, all_true_mask) + shrsi emits signed pto.vshr(lhs_part, rhs_part, all_true_mask) + not emits pto.vnot(source_part, all_true_mask) + +pto.vmi.fma: + semantic: + result = fused_multiply_add(lhs, rhs, acc) + It must not be decomposed to pto.vmi.mulf + pto.vmi.addf because VPTO VMULA + may produce different floating-point results from separate multiply and add. + layout assignment: + lhs, rhs, acc, and result belong to one data layout equivalence class. + verifier contract: + source/result element type must be f16, bf16, or f32 + current direct lowering: + for each physical part: + materialize pto.pset_b16/b32 "PAT_ALL" from the physical element width + emit pto.vmula(acc_part, lhs_part, rhs_part, all_true_mask) + The VMI operand order is lhs, rhs, acc; the VPTO operand order is acc, lhs, rhs. + +pto.vmi.cmpf / cmpi: + verifier contract: + source/result element width must be 8/16/32-bit so the result predicate + can be materialized as b8/b16/b32. + cmpf: f16/bf16/f32, matching VISA VCMP floating-point element types + cmpi: signless/signed/unsigned i8/i16/i32, matching VISA VCMP integer element types + for each physical part: + materialize pto.pset_b8/b16/b32 "PAT_ALL" as the seed predicate + canonicalize predicate to VPTO cmp_mode eq/ne/lt/le/gt/ge + emit pto.vcmp(lhs_part, rhs_part, seed_mask, cmp_mode) + supported cmpf predicates: + eq/ne/lt/le/gt/ge pass through + oeq -> eq + one -> ne + olt -> lt + ole -> le + ogt -> gt + oge -> ge + supported cmpi predicates: + eq/ne pass through + ult -> lt on unsigned integer carriers + ule -> le on unsigned integer carriers + ugt -> gt on unsigned integer carriers + uge -> ge on unsigned integer carriers + slt -> lt + sle -> le + sgt -> gt + sge -> ge + if the physical vreg element signedness does not match the predicate, insert pto.vbitcast to the matching + si/ui integer carrier before pto.vcmp. + unsupported cmpi bare relational predicates lt/le/gt/ge must emit VMI-UNSUPPORTED because integer signedness + must be explicit. + unsupported floating-point predicates such as ord/uno/ult/ule/ugt/uge must emit VMI-UNSUPPORTED until NaN-aware + predicate construction is designed. + +pto.vmi.vcmp / vcmps: + unified vmi_new integer compare uses type-driven signedness: + signed/signless integer element types map lt/le/gt/ge to legacy slt/sle/sgt/sge + unsigned integer element types map lt/le/gt/ge to legacy ult/ule/ugt/uge + explicit integer predicates slt/sle/sgt/sge/ult/ule/ugt/uge are rejected at the unified op level + legacy cmpi remains predicate-driven and requires explicit signed/unsigned relational forms. + +pto.vmi.active_prefix_index: + semantic: + idx[i] = popcount(mask[0 .. i)) + result element type must be signless i8/i16/i32, and concrete mask granularity must match the result element width. + current direct lowering: + only contiguous layout + only one physical result/mask chunk + result and mask chunks must be full, with no padding logical lanes + materialize a zero vreg carrier with pto.vdup + emit pto.vusqz(carrier, mask) + unsupported cases: + partial/tail chunks because padding mask lanes could affect the observable prefix + multi-chunk contiguous values need cross-chunk prefix carry + deinterleaved layouts need logical-lane-order prefix reconstruction + both must report VMI-UNSUPPORTED before OneToN conversion + +pto.vmi.compress: + semantic: + keep source lanes whose mask lane is true and compact them in logical lane order; inactive tail lanes are zero/undef + at the VMI semantic level unless consumed by an operation that defines them. + current direct lowering: + source/result/mask must be contiguous + source/result/mask must each materialize to one physical chunk + source chunk must be full, with no padding logical lanes + emit pto.vsqz(source, mask) + unsupported cases: + partial/tail chunks because padding mask lanes could be squeezed into the observable result prefix + multi-chunk values need cross-chunk compaction and SQZN/carry planning + deinterleaved layouts need logical-lane-order compaction before physical part placement + compress_store is not implied by register compress; store-coupled VSQZ #st=1 and VSTUR require a separate + producer/consumer pairing plan + +pto.vmi.compress_store: + semantic: + store source lanes whose mask lane is true as a dense logical memory stream: + k = 0 + for lane in logical order: + if mask[lane]: + base[offset + k] = value[lane] + k += 1 + layout assignment: + value use is requested as contiguous + mask use is requested as contiguous with granularity derived from value element width + current direct lowering: + value and mask must be contiguous + value and mask must each materialize to one physical chunk + the value chunk must be full, with no padding logical lanes + destination must be a UB !pto.ptr because pto.vstur is pointer-only and UB-only + lower as: + store_base = pto.addptr destination, offset + squeezed = pto.vsqz(value, mask) + align0 = pto.init_align + align1 = pto.vstur align0, squeezed, store_base, "POST_UPDATE" + pto.vstar align1, store_base + The pto.vstur user is the required consumer that lets the VPTO LLVM emitter + set VSQZ #st=1. A plain register pto.vsqz must not be assumed to enqueue + SQZN for store. + unsupported cases: + memref or GM destination until an explicit pointer/materialization plan exists + partial/tail physical chunks, because padding mask lanes could be squeezed into memory + multi-chunk values, because they need cross-chunk active-count compaction and SQZN/VSTUR state planning + deinterleaved layouts, because compaction must be in logical lane order + +pto.vmi.reduce_addi: + semantic: + acc = init[0] + for lane in logical order: + if mask[lane]: + acc = acc + source[lane] // integer wraparound addition + result[0] = acc + layout assignment: + source use is requested as contiguous + init use is requested as contiguous + result natural layout is contiguous + mask use is requested as contiguous with granularity derived from source element width + current direct lowering: + source element width must be 32 bits; narrower vcadd widens its result and needs a separate result type plan + source must materialize to one or more full physical chunks with no padding logical lanes + init/result must be 1-lane VMI vectors and each materialize to one physical chunk + mask must materialize to the same number of physical chunks as source + lower as: + first_lane = pto.pge_b32 "PAT_VL1" + acc = init + for each source_chunk, mask_chunk in physical order: + reduced = pto.vcadd(source_chunk, mask_chunk) + acc = pto.vadd(reduced, acc, first_lane) + result = acc + unsupported cases: + i8/i16 until widening result and init conversion are designed + partial/tail source chunks because padding lanes must not participate + floating-point add reduction without pto.vmi.reduce_addf {reassoc} + +pto.vmi.reduce_addf: + semantic: + requires {reassoc}; without it the verifier rejects the op + acc = init[0] + for lane in any reassociated tree over active logical lanes: + acc = acc + source[lane] + result[0] = acc + layout assignment: + source use is requested as contiguous + init use is requested as contiguous + result natural layout is contiguous + mask use is requested as contiguous with granularity derived from source element width + current direct lowering: + source element type must be f32 + source must materialize to one or more full physical chunks with no padding logical lanes + init/result must be 1-lane VMI vectors and each materialize to one physical chunk + mask must materialize to the same number of b32 physical chunks as source + lower as: + first_lane = pto.pge_b32 "PAT_VL1" + acc = init + for each source_chunk, mask_chunk in physical order: + reduced = pto.vcadd(source_chunk, mask_chunk) + acc = pto.vadd(reduced, acc, first_lane) + result = acc + unsupported cases: + missing reassoc attr + f16 until accumulator precision and rounding contract are designed + partial/tail source chunks because padding lanes must not participate + +pto.vmi.group_load / pto.vmi.group_store: + semantic: + num_groups is the only static grouping attribute. + N = logical lane count; G = num_groups; S = N / G. + group_load reads each logical group as one contiguous row: + result[g * S + i] = source[offset + g * row_stride + i] + for 0 <= g < G and 0 <= i < S + group_store writes the inverse row mapping: + destination[offset + g * row_stride + i] = value[g * S + i] + row_stride is an index operand, measured in elements, and may be dynamic. + Tail/valid-lane information is not an attr; it must be represented by a + mask in the producing/consuming computation. The current direct + group_load/group_store path is for full physical chunks. + layout assignment: + group_load result natural layout is contiguous + group_store value use is requested as contiguous + current direct lowering: + source/value element width must be maskable by b8/b16/b32 + layout must be contiguous with full physical chunks + num_groups must evenly divide N, and the derived group size S must be a + multiple of the physical lanes + per part, so every physical chunk belongs to exactly one group + lower each physical chunk with pto.vlds/pto.vsts at: + offset + group * row_stride + chunk_in_group * lanes_per_part + unsupported cases: + derived group size splitting a physical chunk, because this needs partial-vreg + lane insertion/extraction or a gather/scatter plan + partial/tail physical chunks + GM-backed direct vector load/store paths not already accepted by the normal + VMI memory access plan + +pto.vmi.group_reduce_addf: + semantic: + requires {reassoc} + N = logical lane count; G = num_groups; S = N / G + L = physical lanes per 256B chunk for the element type. + The result carries #pto.vmi.layout, a group-slot + group-slot layout. It is not a dense vector layout: only slot lanes have + semantic values. Supported K values are: + K = 8 for VCGADD-style packed results, where group g is stored in + physical chunk floor(g / 8), lane g % 8. + K = 1 for row-local VCADD results, where group g is stored in physical + chunk g, lane 0. + for each group g: + result[group_slot(g)] = + reduce_add(source[g * S .. (g + 1) * S), mask in same range) + Non-slot lanes are not consumed by pto.vmi.group_broadcast. The current + direct lowering materializes them as zero where the hardware path does not + already define them. + The result remains a VMI vector with the same element type as the source, + but its logical lane count is G: one scalar result per group. Its layout + is an explicit group-slot layout that describes where those G scalars are + placed in physical registers. + layout assignment: + source use is requested as contiguous + result natural layout is #pto.vmi.layout + mask use is requested as contiguous with granularity derived from source + element width + current direct lowering: + source/result element type must be f16 or f32 + source and mask must have compatible full physical chunks. The result is + `GxT` group-slot data and may have different physical arity from the + source tile. + if S=8 for f32, lower each physical chunk with pto.vcgadd. This is the + hardware 32B VLane group reduction path for f32: each source chunk produces + eight 8-lane group sums in the low lanes of that physical chunk. The + lowering preserves this natural no-pack result. + Otherwise: + derived group size S must be a multiple of physical lanes per part + lower each source chunk with pto.vcadd, combine chunks in the same group + with pto.vadd under PAT_VL1, then place group g in the slot lane defined by + K. All other result chunks/lane values + are zero. + unsupported cases: + missing reassoc attr + integer element types, which use the corresponding typed integer op + derived group size S that neither divides nor is a multiple of L + +pto.vmi.group_reduce_addi / group_reduce_maxi / group_reduce_mini: + semantic: + source and result use the same i8/i16/i32 element type + the result has one group-slot value per logical group + integer addition has same-type wraparound semantics + layout assignment: + use the same registered group-block table as floating-point group reduction + packed 32B-block cases use slots=8 + aligned full-row cases use slots=1 + current direct lowering: + packed cases use pto.vcgadd/pto.vcgmax/pto.vcgmin and same-type combines + aligned full-row max/min cases use pto.vcmax/pto.vcmin + aligned full-row i8/i16 add cases use widening pto.vcadd partials and + widened pto.vadd combines, then pto.vbitcast the low bits back to the + declared VMI result type + the widening is internal and is not exposed in the VMI type contract + unsupported cases: + element types other than i8/i16/i32 + group sizes outside the registered high-performance group-block classes + +pto.vmi.group_broadcast: + semantic: + source logical lane count is G; result logical lane count is N. + S = N / G. + source must carry #pto.vmi.layout. For each + group g, the source value is read from the slot lane defined by K. The + result broadcasts it back to each logical group: + result[g * S + i] = source[group_slot(g)] + layout assignment: + source use is requested as #pto.vmi.layout + result is consumer-driven. If no consumer requests another layout, it + defaults to contiguous. + current direct lowering: + source must carry #pto.vmi.layout with one + logical lane per group + result may be contiguous with full physical chunks + result may also be deinterleaved when S is large enough that every physical + result chunk stays inside one logical group, for example N=512, G=2, S=256, + L=64, deinterleaved=4. If the source is + #pto.vmi.layout, the source physical part is + selected by group id rather than by source chunk id. + derived group size S must divide or be a multiple of L for canonical + group-slot addressing + if result is contiguous and S < L, each physical chunk contains multiple group + slots. Lower by + creating an index vector [0...0, 1...1, ...] and applying pto.vselr to the + corresponding source chunk. + if S >= L and each result physical chunk belongs to one group, lower by + duplicating the first lane of that group's source chunk with pto.vdup LOWEST. + unsupported cases: + partial/tail physical chunks + derived group size S that neither divides nor is a multiple of L + deinterleaved small-group broadcast where one physical result chunk needs + values from multiple source chunks + +pto.vmi.reduce_maxf / reduce_minf / reduce_maxi / reduce_mini: + semantic: + acc = init[0] + for each active logical lane in logical lane order: + reduce_max*: acc = max(acc, source[lane]) + reduce_min*: acc = min(acc, source[lane]) + result[0] = acc + inactive lanes inside each physical chunk follow VPTO identities: + reduce_maxf uses pto.vcmax, where inactive FP lanes behave as -INF + reduce_minf uses pto.vcmin, where inactive FP lanes behave as +INF + NaN and signed-zero behavior follows pto.vcmax/pto.vcmin for the chunk + reduction and pto.vmax/pto.vmin for serial chunk accumulation. The index + lane produced by pto.vcmax/pto.vcmin is ignored because VMI exposes only the + 1-lane value result. + layout assignment: + source use is requested as contiguous + init use is requested as contiguous + result natural layout is contiguous + mask use is requested as contiguous with granularity derived from source element width + current direct lowering: + source element type must be f16/f32 for the floating ops or i8/i16/i32 for + the integer ops + source must materialize to one or more full physical chunks with no padding logical lanes + init/result must be 1-lane VMI vectors and each materialize to one physical chunk + mask must materialize to the same number of physical chunks as source + lower reduce_maxf as: + first_lane = pto.pge_b16/b32 "PAT_VL1" + acc = init + for each source_chunk, mask_chunk in physical order: + reduced = pto.vcmax(source_chunk, mask_chunk) + acc = pto.vmax(reduced, acc, first_lane) + result = acc + lower reduce_minf as: + first_lane = pto.pge_b16/b32 "PAT_VL1" + acc = init + for each source_chunk, mask_chunk in physical order: + reduced = pto.vcmin(source_chunk, mask_chunk) + acc = pto.vmin(reduced, acc, first_lane) + result = acc + unsupported cases: + bf16/fp8/f64 until VPTO reduction and combine semantics are designed + partial/tail source chunks because padding lanes must not participate + integer widths other than i8/i16/i32 + +pto.vmi.select: + current direct lowering is a storage-width select rather than a semantic + arithmetic op: source/result physical elements must be b8/b16/b32-maskable, + but signedness and float-vs-integer interpretation are not inspected. + for each physical part: + consume the corresponding physical predicate part + emit pto.vsel(true_part, false_part, predicate_part) + +pto.vmi.extf, direct path: + support 16-bit float-like contiguous source part -> f32 deinterleaved=2 result parts + materialize pto.pset_b16 "PAT_ALL" + emit pto.vcvt(source_part, mask, part=EVEN/ODD) + partial/tail is valid when the logical lanes fit in the one physical source + part; PAT_ALL may convert padding lanes, but those lanes remain padding in + the deinterleaved result + support 8-bit contiguous source part -> f32 deinterleaved=4 result parts + materialize pto.pset_b8 "PAT_ALL" + emit pto.vcvt(source_part, mask, part=P0/P1/P2/P3) + the same padding rule applies + reject other extf width/layout shapes until their exact part plan is implemented + +pto.vmi.truncf, direct path: + support f32 deinterleaved=2 source parts -> 16-bit contiguous result part + materialize pto.pset_b32 "PAT_ALL" for the source conversion + emit pto.vcvt(even_f32_part, mask, rnd=R, sat=SAT, part=EVEN) + emit pto.vcvt(odd_f32_part, mask, rnd=R, sat=SAT, part=ODD) + materialize pto.pset_b16 "PAT_ALL" + merge mutually exclusive part results with pto.vor + partial/tail is valid when the two source parts pack into one physical + result part; converted padding lanes remain result padding + support f32 deinterleaved=4 source parts -> 8-bit contiguous result part + materialize pto.pset_b32 "PAT_ALL" for the source conversion + emit pto.vcvt(p0_f32_part, mask, rnd=, sat=SAT, part=P0) + emit pto.vcvt(p1_f32_part, mask, rnd=, sat=SAT, part=P1) + emit pto.vcvt(p2_f32_part, mask, rnd=, sat=SAT, part=P2) + emit pto.vcvt(p3_f32_part, mask, rnd=, sat=SAT, part=P3) + result round is R for f8E4M3/f8E5M2, A for default hif8, or H for + hif8 truncf with {rounding = "H"} + materialize pto.pset_b8 "PAT_ALL" + merge mutually exclusive part results with pto.vor + partial/tail is valid when the four source parts pack into one physical + result part; converted padding lanes remain result padding + reject other truncf width/layout shapes until their exact pack plan is implemented + +pto.vmi.bitcast: + for each physical part: + emit pto.vbitcast(source_part) -> result_part_type + source/result layouts must match, physical arity must match, and every + corresponding physical chunk must carry the same number of logical bits. + This includes contiguous, deinterleaved, and identical group_slots layouts. + Padding bits may map only to result padding bits; any shape where source + padding would become result logical data remains unsupported. + +pto.vmi.channel_split / pto.vmi.channel_merge: + support 2-way and 4-way channel transforms for contiguous per-channel values + and matching deinterleaved=C merged values. + + channel_split C=2: + if the source layout is already deinterleaved=2, forward physical chunks + directly to the two contiguous channel results. + if the source layout is contiguous, source logical vector must physicalize + as 2*N contiguous chunks. For each pair of dense chunks: + %ch0_i, %ch1_i = pto.vdintlv %dense_2i, %dense_2i_plus_1 + Results are returned in per-channel order: + channel0 chunks..., channel1 chunks... + + channel_split C=4: + if the source layout is already deinterleaved=4, forward physical chunks + directly to the four contiguous channel results. + if the source layout is contiguous, source logical vector must physicalize + as 4*N contiguous chunks. The lowering is the same two-level pto.vdintlv + tree used by contiguous -> deinterleaved=4 materialization, but the + partition-major output is interpreted as four separate contiguous channel + results. + + channel_merge C=2/C=4: + inputs are consumed as per-channel contiguous chunks. + If the result layout is deinterleaved=C, the physical chunks are forwarded + directly in partition-major order. + If the result layout is contiguous, the lowering uses the reverse + pto.vintlv tree and returns dense contiguous chunks for the merged result. + + Unsupported: + channel counts other than 2 or 4 + non-matching channel input/result layouts + arity-changing or uneven partial physical channel groups that cannot form + complete intlv/dintlv groups + +pto.vmi.shuffle: + first try whole physical chunk forwarding cases: + source/result layouts are assigned + every non-padding lane in a result physical chunk maps to the same source physical chunk + source lane number equals result lane number inside the physical chunk + result padding lanes are ignored and remain semantically unobservable + + If forwarding fails, try vci-materializable vselr per physical chunk: + every result physical chunk has no padding lane + every lane in a result physical chunk maps to the same source physical chunk + source lane indices inside the chunk form one ASC or DESC consecutive sequence + materialize the index vector with pto.vci(base_lane, ASC|DESC) + emit pto.vselr(source_chunk, index_vector) + + Examples: + identity 128xf32 -> 128xf32: + indices = [0, 1, ..., 127] + forward dense chunks 0 and 1 + + second physical chunk 128xf32 -> 64xf32: + indices = [64, 65, ..., 127] + forward dense chunk 1 + + tail prefix 128xf32 -> 4xf32: + indices = [0, 1, 2, 3] + forward dense chunk 0 + lanes 4..63 of the physical result are padding lanes and are not part of + the logical vmi value + + chunk swap 128xf32 -> 128xf32: + indices = [64, 65, ..., 127, 0, 1, ..., 63] + forward dense chunks in order 1, 0 + + reverse one 64xf32 chunk: + indices = [63, 62, ..., 0] + index = pto.vci 63 {order = DESC} : i32 -> !pto.vreg<64xi32> + result = pto.vselr source_chunk, index + + Unsupported: + partial physical chunk projection whose observable result lanes are not + padding-safe forwarding, e.g. [1, 2, 3, 4] -> 4xf32 when it would require + shifting lanes rather than forwarding a whole physical chunk + broadcast, duplicate lanes, arbitrary non-affine permutation + current implementation emits VMI-UNSUPPORTED for these cases before + OneToN conversion, instead of leaving a generic residual VMI op. +``` + +`func.return` 携带 VMI operand 时必须通过 OneToN func/return structural pattern 展开成 physical +return operands。不能只取第一个 physical part；这种错误会导致函数类型已经返回两个 physical value， +但 `func.return` 只返回一个 value。 + +### 6.1 Type Conversion + +Use one shared physicalization helper: + +```text +VMIVRegType -> N physical !pto.vreg +VMIMaskType -> N physical !pto.mask +``` + +Physical result ordering must be: + +```text +contiguous: + chunk0, chunk1, ... + +deinterleaved=K: + p0_chunk0, p0_chunk1, ..., p1_chunk0, ..., p(K-1)_chunkN +``` + +### 6.2 Structural Conversion + +The pass must convert: + +```text +operation results +block arguments +branch operands +cf.br / cf.cond_br successor block signatures +scf.if results and yields +scf.for iter_args and yields +func arguments/results +call operands/results +return operands +cf.br / cf.cond_br / cf.switch block arguments and successor operands +scf.execute_region results and yields: + current implementation uses a project-local OneToN structural pattern. +scf.index_switch results and yields: + current implementation uses a project-local OneToN structural pattern. +``` + +Do not rely on a defining op to recover parts. Any VMI value may come from a block argument or function +argument, so `unpack` must be valid on arbitrary layout-assigned VMI SSA values before final lowering. + +### 6.3 Op Lowering + +Internal helper lowering： + +```text +unpack: + replace with physical values in helper ordering + +pack: + materialize one logical VMI aggregate before it is immediately consumed by another VMI helper + must not remain after final gate + +ensure_layout: + preflight: + source/result must have computable physical arity + source/result physical arity must match + identity source/result layouts do not require full chunks + if source/result layouts differ, either: + every source/result physical chunk is full, with no padding lanes; or + source/result both have complete contiguous/deinterleaved=2/4 materialization groups and their materialized + physical arity still equals the original VMI physical arity + arity-changing partial/tail layout conversion remains unsupported because it would need an explicit padding + packing/drop plan + otherwise report VMI-UNSUPPORTED before OneToN conversion + + compare the original VMI source/result layout attrs: + same layout: + forward the converted source parts + deinterleaved=2 -> contiguous: + %d0, %d1 = pto.vintlv %p0, %p1 + contiguous -> deinterleaved=2: + %p0, %p1 = pto.vdintlv %d0, %d1 + deinterleaved=4 -> contiguous: + %a0, %a1 = pto.vintlv %p0, %p2 + %b0, %b1 = pto.vintlv %p1, %p3 + %d0, %d1 = pto.vintlv %a0, %b0 + %d2, %d3 = pto.vintlv %a1, %b1 + contiguous -> deinterleaved=4: + %a0, %b0 = pto.vdintlv %d0, %d1 + %a1, %b1 = pto.vdintlv %d2, %d3 + %p0, %p2 = pto.vdintlv %a0, %a1 + %p1, %p3 = pto.vdintlv %b0, %b1 + + It is a bug to treat layout conversion as identity merely because both sides convert to the same + number of physical !pto.vreg values with the same type. For example: + !pto.vmi.vreg<128xf32, deinterleaved=2> + !pto.vmi.vreg<128xf32, contiguous> + both physicalize to two !pto.vreg<64xf32> values, but their logical lane order differs. + +ensure_mask_layout: + preflight: + source/result must have computable physical arity + source/result physical arity must match + if source/result layouts differ, every source/result physical predicate chunk must be full, with no padding lanes + identity source/result layouts do not require full chunks + otherwise report VMI-UNSUPPORTED before OneToN conversion + + same-layout: + forward source parts + deinterleaved=2 -> contiguous: + use pto.pintlv_b8/b16/b32 on each partition pair + contiguous -> deinterleaved=2: + use pto.pdintlv_b8/b16/b32 on each dense pair + deinterleaved=4 -> contiguous: + use the same two-level tree as data layout conversion, replacing pto.vintlv with pto.pintlv_b8/b16/b32 + contiguous -> deinterleaved=4: + use the reverse two-level tree, replacing pto.vdintlv with pto.pdintlv_b8/b16/b32 + source/result granularity must be identical; granularity conversion belongs to ensure_mask_granularity. + +ensure_mask_granularity: + source/result layout and logical lane count must match. + source/result granularity must be concrete b8/b16/b32. + identity conversion forwards physical parts. + widening conversion: + b8 -> b16 or b16 -> b32 uses pto.punpack LOWER/HIGHER for each source physical chunk. + each source physical mask chunk can produce up to two result chunks in logical order. + narrowing conversion: + b32 -> b16 or b16 -> b8 uses pto.ppack LOWER for the low source chunk. + if a high source chunk exists, use pto.ppack HIGHER and merge the two partial masks with pto.por under PAT_ALL. + this handles odd tail groups because the missing high half is padding and remains zero. + multi-step conversion: + b8 -> b32 is b8 -> b16 -> b32. + b32 -> b8 is b32 -> b16 -> b8. +``` + +Elementwise lowering： + +```text +for each physical part: + lower add/cmp/select to corresponding VPTO op sequence + preserve source/result physical ordering + cmp predicates must be canonicalized before creating pto.vcmp: + cmpf eq/ne/lt/le/gt/ge pass through + ordered FP aliases oeq/one/olt/ole/ogt/oge map to eq/ne/lt/le/gt/ge + cmpi eq/ne pass through + unsigned integer aliases ult/ule/ugt/uge map to lt/le/gt/ge on unsigned carriers + signed integer aliases slt/sle/sgt/sge map to lt/le/gt/ge + insert pto.vbitcast to si/ui integer carriers when the physical vreg element signedness does not match + cmpi bare relational predicates lt/le/gt/ge are unsupported because integer signedness must be explicit + unordered/NaN-sensitive FP predicates are unsupported until represented explicitly +``` + +Producer lowering： + +```text +broadcast: + TypeConverter gives the ordered result physical types. + For each result physical vreg: + create all-true mask with the vreg element width + emit pto.vdup scalar -> that physical vreg + + This is valid for contiguous and deinterleaved layouts because splat has no lane-order dependence. + +constant: + Splat dense constants use the same path as broadcast: + create scalar arith.constant from the splat attribute + emit pto.vdup per physical result part + require the same 8/16/32-bit physical result element-width precondition as + broadcast + Non-splat dense constants need an explicit constant materialization strategy or must remain unsupported with a + precise diagnostic; do not synthesize an arbitrary lane sequence by scalar inserts unless that path is designed. + +create_mask / constant_mask: + constant active_lanes create_mask lowers per physical mask part: + clamp active_lanes to [0, logical lane count] + compute active prefix count for each physical mask chunk with the VMI lane-map helper + emit pto.pge_b8/b16/b32 PAT_ALL, PAT_ALLF, or supported PAT_VL* + if a chunk prefix count has no supported PAT_VL token, fall back to pto.plt_b8/b16/b32 with a constant i32 count + Dynamic active_lanes with contiguous layout lowers by chaining pto.plt_b8/b16/b32 over the physical chunks: + active_i32 = arith.index_cast active_lanes : index to i32 + active_i32 = minui(maxsi(active_i32, 0), logical_lane_count) + mask0, remaining0 = pto.plt_b* active_i32 + mask1, remaining1 = pto.plt_b* remaining0 + ... + Dynamic active_lanes with deinterleaved layout remaps one logical prefix into per-part dynamic lane counts before + chaining pto.plt_b*: + active_i32 = minui(maxsi(index_cast(active_lanes), 0), logical_lane_count) + part_count(part) = (active_i32 + factor - 1 - part) / factor + then chain pto.plt_b* independently for each partition in VMI physical order: + p0 chunks..., p1 chunks..., ... + dense constant_mask lowers per physical mask part: + first map logical lanes to physical predicate lanes using the assigned VMI layout + prefix chunks emit pto.pset_b8/b16/b32 PAT_ALL, PAT_ALLF, or supported PAT_VL* + if a prefix count has no supported PAT_VL token, emit pto.plt_b8/b16/b32 with a constant i32 count + non-prefix chunks are decomposed into static active runs: + prefix(hi) = pto.pge/plt for the run end + prefix(lo) = pto.pge/plt for the run begin + run = prefix(hi) & ~prefix(lo) using pto.pnot + pto.pand + chunk = run0 | run1 | ... using pto.por + +Unsupported diagnostics: + unexpected residual dynamic pto.vmi.create_mask after OneToN conversion: + VMI-UNSUPPORTED: dynamic pto.vmi.create_mask active_lanes could not be lowered by the current runtime predicate + generation plan + This is a final-gate diagnostic for malformed or newly unsupported dynamic shapes. The supported dynamic + contiguous/deinterleaved=2/deinterleaved=4 paths above must lower before this residual gate. + + non-splat pto.vmi.constant: + VMI-UNSUPPORTED: non-splat pto.vmi.constant requires a vreg immediate or scratch materialization plan + + unsupported partial/tail masked/expand read-style op: + VMI-UNSUPPORTED: pto.vmi. requires full physical chunks without padding lanes or a statically safe + full-read footprint (...; safe-read proof failed: ...) + GM-backed direct pto.vmi.load/masked_load/expand_load: + VMI-UNSUPPORTED: pto.vmi. ... (source is GM-backed, but current direct VMI-to-VPTO memory lowering + emits pto.vlds/pto.vsts and requires UB-backed memory) + unsupported partial/tail pto.vmi.store/masked_store: + VMI-UNSUPPORTED: pto.vmi. requires an 8/16/32-bit predicate-maskable element type and either full + physical chunks or contiguous/deinterleaved tail-store materialization, with UB-backed destination; unsupported + cases include values such as f64/index that have no b64 predicate representation, GM-backed destinations that + still need a memory movement/materialization plan, and uneven deinterleaved physical groups that cannot form + complete intlv groups + + unsupported non-identity partial/tail pto.vmi.ensure_layout: + VMI-UNSUPPORTED: pto.vmi.ensure_layout cannot materialize the requested data layout conversion; unsupported cases + include arity-changing partial/tail conversion and uneven deinterleaved groups that cannot form complete intlv + groups + If the helper has a single consumer, the main diagnostic is emitted on the + consumer op and operand, including both the actual operand VMI type and the + required VMI type. For example, pto.vmi.truncf operand #0 can report + `!pto.vmi.vreg<128xf32, contiguous>` vs. + `!pto.vmi.vreg<128xf32, deinterleaved=4>` for f32->fp8. The failed + pto.vmi.ensure_layout conversion is attached as a note. + + unsupported non-identity partial/tail pto.vmi.ensure_mask_layout: + VMI-UNSUPPORTED: pto.vmi.ensure_mask_layout cannot materialize the requested mask layout conversion; unsupported + cases include arity-changing partial/tail conversion and uneven deinterleaved groups that cannot form complete + predicate intlv groups + + unsupported pto.vmi.ensure_mask_granularity: + VMI-UNSUPPORTED: non-identity mask granularity materialization requires concrete b8/b16/b32 masks with matching + lane count and layout (...) + + unsupported pto.vmi.extf direct path shape: + VMI-UNSUPPORTED: pto.vmi.extf supports only one contiguous 16-bit float-like or fp8-like physical source chunk to f32 + deinterleaved=2/4 results; partial/tail is allowed only when source padding maps to result padding + + unsupported pto.vmi.truncf direct path shape: + VMI-UNSUPPORTED: pto.vmi.truncf supports only f32 deinterleaved=2 source parts to one contiguous f16 result chunk + or f32 deinterleaved=4 source parts to one contiguous fp8-like result chunk + + unsupported pto.vmi.bitcast shape: + VMI-UNSUPPORTED: pto.vmi.bitcast requires matching source/result layouts with identical physical + arity and matching per-chunk logical bit footprints (...) + + unsupported pto.vmi.channel_split / pto.vmi.channel_merge channel count: + VMI-UNSUPPORTED: pto.vmi.channel_split supports only 2 or 4 channels + VMI-UNSUPPORTED: pto.vmi.channel_merge supports only 2 or 4 channels + unsupported pto.vmi.channel_split / pto.vmi.channel_merge layout: + VMI-UNSUPPORTED: pto.vmi.channel_split requires source layout to be contiguous or matching deinterleaved channel + layout, and every result layout to be contiguous + VMI-UNSUPPORTED: pto.vmi.channel_merge requires every input layout to be contiguous and result layout to be + contiguous or matching deinterleaved channel layout +``` + +Width conversion lowering： + +```text +f16 -> f32: + supported direct path when source is contiguous and result is deinterleaved=2: + pto.vcvt part=EVEN produces logical lanes 0,2,4,... + pto.vcvt part=ODD produces logical lanes 1,3,5,... + source/result physical arity must be 1 -> 2 + +f8 -> f32: + supported direct path when source is contiguous and result is deinterleaved=4: + pto.vcvt part=P0/P1/P2/P3 produces the four modulo-4 lane partitions + source/result physical arity must be 1 -> 4 + +f32 -> f16: + supported direct path when source is deinterleaved=2 and result is contiguous: + pto.vcvt part=EVEN consumes even/source part 0 + pto.vcvt part=ODD consumes odd/source part 1 + pto.vor merges mutually exclusive f16 part results into one contiguous vreg + source/result physical arity must be 2 -> 1 + current default conversion attrs are rnd=R, sat=SAT + +f32 -> 8-bit fp-like: + supported direct path when source is deinterleaved=4 and result is contiguous: + pto.vcvt part=P0/P1/P2/P3 consumes the four source partitions + pto.vor merges mutually exclusive byte-lane part results into one + contiguous vreg + source/result physical arity must be 4 -> 1 + current default conversion attrs are rnd=R for f8E4M3/f8E5M2 and rnd=A for + hif8. pto.vmi.truncf {rounding = "H"} is accepted only for f32 -> hif8 + and forwards rnd=H to the emitted pto.vcvt operations. +``` + +Memory lowering： + +```text +vmi.load: + current direct memory path first reads contiguous physical chunks. The logical lane count must be an exact multiple + of the physical vreg lane count. + For each contiguous physical chunk i: + offset_i = base_offset + i * lanesPerPart + dense_i = pto.vlds base[offset_i] + + If the requested VMI result layout is contiguous, return the dense chunks directly. + If the requested VMI result layout is deinterleaved=2: + prefer pto.vldsx2 "DINTLV_B8/B16/B32" per physical chunk group: + %p0_i, %p1_i = pto.vldsx2 base[offset_i], "DINTLV_B*" + return results in VMI partition-major order: + p0_chunk0, p0_chunk1, ..., p1_chunk0, p1_chunk1, ... + If the requested VMI result layout is deinterleaved=4 with exactly four physical parts: + use dense pto.vlds chunks followed by the reverse two-level pto.vdintlv tree. + + For larger multi-chunk deinterleaved=4 loads, apply the same conversion per contiguous chunk group and return + physical parts in VMI partition-major order: + deinterleaved=4: p0_chunks..., p1_chunks..., p2_chunks..., p3_chunks... + +vmi.store: + direct lowering requires value element width to be 8, 16, or 32 bits so the + emitted pto.vsts/pto.vstsx2 predicate can be materialized as b8/b16/b32. + contiguous layout with full physical chunks: + offset_i = base_offset + i * lanesPerPart + mask_i = pto.pset_b8/b16/b32 "PAT_ALL" + pto.vsts value_i, base[offset_i], mask_i + contiguous layout with a final partial physical chunk: + full chunks still use PAT_ALL + the final chunk computes valid_lanes = logical_lane_count - chunk_i * lanesPerPart + tail_mask_i = pto.plt_b8/b16/b32(valid_lanes) + pto.vsts tail_value_i, base[offset_i], tail_mask_i + padding lanes therefore have no externally visible store effect. + +deinterleaved store: + deinterleaved=2 with full physical chunks: + prefer pto.vstsx2 "INTLV_B8/B16/B32" per physical chunk group: + pto.vstsx2 p0_i, p1_i, base[offset_i], "INTLV_B*", all_true_mask + offset_i = base_offset + i * 2 * lanesPerPart + the vstsx2 dist mode writes logical lane 0,1,2,3,... order externally. + + current safe path lowers through proven register materialization before store: + deinterleaved=4 with exactly four physical parts: + use the two-level pto.vintlv tree, then store %d0/%d1/%d2/%d3 as contiguous chunks + + Larger multi-chunk deinterleaved=4 values use the same conversion per chunk group. The final store order is dense + chunk order, so external memory observes logical lane 0,1,2,... order. + +vmi.masked_load: + semantics: + if mask[lane] is true, result[lane] = memory[base + lane] + if mask[lane] is false, result[lane] = passthru[lane] + inactive mask lanes do not by themselves permit unsafe memory reads + current direct path: + result, passthru, and mask are requested as contiguous + full physical chunks can always use pto.vlds because every loaded lane is logical + partial/tail chunks require the same statically safe full-read proof as vmi.load + for each contiguous physical chunk i: + loaded_i = pto.vlds base[offset_i] + result_i = pto.vsel loaded_i, passthru_i, mask_i + unsupported cases: + non-contiguous layouts + unsafe partial/tail read footprints + target true masked/non-faulting load and guarded/scratch fallback + +vmi.stride_load: + semantics: + result lane order is contiguous VMI logical order + source addresses are described by the VPTO block/repeat stride operands + mask false lanes are inactive for the underlying block-strided load + layout assignment: + result natural layout is contiguous + mask use is requested as contiguous with granularity derived from result element width + current direct path: + source must be !pto.ptr + result and mask must be one contiguous physical chunk + base = pto.addptr source, offset + result = pto.vsldb base, block_stride, repeat_stride, mask + unsupported cases: + multi-chunk result or mask + non-contiguous layouts + memref/gm source + +vmi.stride_store: + semantics: + value lane order is contiguous VMI logical order + destination addresses are described by the VPTO block/repeat stride operands + mask false lanes do not write memory + layout assignment: + value use is requested as contiguous + mask use is requested as contiguous with granularity derived from value element width + current direct path: + destination must be !pto.ptr + value and mask must be one contiguous physical chunk + base = pto.addptr destination, offset + updated_base = pto.vsstb value, base, block_stride, repeat_stride, mask + The updated base result is intentionally unused by VMI lowering, but the + post-update VPTO form matches CCE block-strided staging behavior. + unsupported cases: + multi-chunk value or mask + non-contiguous layouts + memref/gm destination + +vmi.gather: + semantics: + if mask[lane] is true, result[lane] = memory[base + indices[lane]] + if mask[lane] is false, result[lane] = passthru[lane] and no memory read occurs for that lane + indices are interpreted in element units, not bytes + layout assignment: + result natural layout is contiguous + indices and passthru uses are requested as contiguous + mask use is requested as contiguous with granularity derived from result element width + current direct path: + source must be !pto.ptr + supported 32-bit mode: + T must be a 32-bit element type + indices must be signless or unsigned i32 + result / indices / passthru / mask must be contiguous full physical chunks + mask granularity must be b32 + for each physical chunk i: + gathered_i = pto.vgather2_bc source, indices_i, mask_i + result_i = pto.vsel gathered_i, passthru_i, mask_i + supported ui16 mode: + T must be ui16 + indices must be unsigned i16 + result / indices / passthru / mask must be one contiguous physical chunk + mask granularity must be b16 + gathered = pto.vgather2 source, indices, mask + result = pto.vsel gathered, passthru, mask + VPTO LLVM emitter bitcasts the physical index register from <128xi16> + to the installed Bisheng intrinsic ABI <64xi32>; this is the same + 256B register payload viewed as the wrapper-level vector_u16 index + container. + reason for vsel: + VPTO gather false predicate lanes do not read memory but produce zero; VMI false lanes preserve passthru. + unsupported cases: + f16/b16/f8/i8 result element types + partial/tail chunks + non-contiguous layouts + memref/gm source + guarded/scratch fallback + +vmi.scatter: + semantics: + if mask[lane] is true, memory[base + indices[lane]] = value[lane] + if mask[lane] is false, no memory write occurs for that lane + indices are interpreted in element units, not bytes + all active lanes must have pairwise-distinct indices; duplicate active indices violate the VMI scatter contract + layout assignment: + value and indices uses are requested as contiguous + mask use is requested as contiguous with granularity derived from value element width + current direct path: + destination must be !pto.ptr + T must be a 32-bit element type + indices must be signless or unsigned i32 + value / indices / mask must be contiguous full physical chunks + mask granularity must be b32 + for each physical chunk i: + pto.vscatter value_i, destination, indices_i, mask_i + unsupported cases: + f16/b16/f8/i8 value element types + partial/tail chunks + non-contiguous layouts + memref/gm destination + ordered duplicate-index fallback + +vmi.expand_load: + semantics: + k = 0 + for lane in logical order: + if mask[lane]: + result[lane] = memory[base + k] + k += 1 + else: + result[lane] = passthru[lane] + layout assignment: + result natural layout is contiguous + passthru use is requested as contiguous + mask use is requested as contiguous with granularity derived from result element width + current direct path: + static all-active path: + pto.vmi.create_mask with constant active_lanes >= logical lane count + dense all-true pto.vmi.constant_mask + in that case expand_load degenerates to ordinary vmi.load: + for each contiguous physical chunk i: + loaded_i = pto.vlds base[offset_i] + result_i = loaded_i + partial/tail chunks still require the same statically safe full-read proof as vmi.load. + runtime-mask path: + source must be !pto.ptr + T must be a 32-bit element type + result / passthru / mask must be contiguous one full physical chunk + mask granularity must be b32 + base_i = pto.addptr source, offset + indices_i = pto.vusqz(zero_i32_carrier, mask_i) + loaded_i = pto.vgather2_bc base_i, indices_i, mask_i + result_i = pto.vsel loaded_i, passthru_i, mask_i + unsupported cases: + runtime masks across multiple physical chunks + runtime masks on non-32-bit element types + non-contiguous layouts + unsafe partial/tail read footprints + guarded load or scratch fallback + +vmi.masked_store: + semantics: + if mask[lane] is true, store value[lane] + if mask[lane] is false, no memory write occurs for that logical lane + current full-footprint path: + value and mask are requested as contiguous at the use site + mask granularity is derived from value element width + for each contiguous physical chunk i: + offset_i = base_offset + i * lanesPerPart + pto.vsts value_i, base[offset_i], mask_i + contiguous layout with a final partial physical chunk: + full chunks store with the user mask directly + the final chunk computes tail_valid_i with pto.plt_b8/b16/b32(valid_lanes) + store_mask_i = pto.pand user_mask_i, tail_valid_i, all_true_mask_i + pto.vsts tail_value_i, base[offset_i], store_mask_i + padding lanes and user-inactive lanes therefore both have no write effect. + If the incoming value/mask are deinterleaved, layout assignment inserts + ensure_layout/ensure_mask_layout or the vmi-to-vpto pattern materializes the same contiguous representation before + emitting stores. This preserves logical memory order and keeps inactive lanes write-free. + +non-full chunks: + vmi.store and vmi.masked_store support contiguous tail chunks by predicating the final pto.vsts with + a prefix valid mask. masked_store additionally ANDs the user mask with the tail-valid mask. + deinterleaved=2/4 tail store/masked_store is supported only through explicit layout materialization to + contiguous chunks first. This requires every deinterleaved part to have the same physical chunk count, so the + materializer can build complete vintlv/pintlv groups. After materialization, each contiguous chunk is predicated by + the logical tail-valid mask; chunks whose active logical lane count is zero are not emitted as stores. Uneven + deinterleaved groups, such as 129xf32 with deinterleaved=2, remain unsupported until a padding/scratch plan can + assemble only the observable contiguous chunks. + vmi.load support partial/tail chunks only when the direct full physical read is statically safe: + statically shaped memref source, constant non-negative offset, and enough elements for the + whole physical read footprint. Padding lanes must never become observable. Other partial/tail load cases still need + scratch/guarded/true-masked load planning. +``` + +Histogram lowering： + +```text +vmi.dhist semantics: + source lanes are ui8 samples + mask selects active source lanes + acc/result are complete logical 256-bin ui16 histograms + result[b] = acc[b] + count(active source lanes whose value equals b) + +layout assignment: + source layout = contiguous + mask layout = contiguous, granularity b8 + acc/result layout = contiguous !pto.vmi.vreg<256xui16> + +physicalization: + acc/result physical arity is 2 because 256xui16 is 512B + part0 represents logical bins 0..127 + part1 represents logical bins 128..255 +``` + +`vmi-to-vpto` lowering for `pto.vmi.vdhist` is local and deterministic from the +op and assigned types: + +```text +lo = converted acc part0 +hi = converted acc part1 + +for each converted source physical chunk c in logical order: + chunk_mask = converted b8 mask chunk c + + if source chunk c contains padding lanes because N is not a multiple of 256: + valid = pto.pge/plt_b8 prefix mask for the valid logical lanes in this chunk + chunk_mask = pto.pand chunk_mask, valid + + lo = pto.dhistv2 lo, src_c, chunk_mask, #bin=0 + hi = pto.dhistv2 hi, src_c, chunk_mask, #bin=1 + +return physical result parts [lo, hi] +``` + +Required preflight: + +```text +acc/result element type is ui16 and logical lane count is exactly 256 +source element type is ui8 +source and mask logical lane counts match +source/mask are contiguous +mask granularity is b8 +source physical chunks are 256-lane ui8 chunks; final partial chunk is allowed +only when the lowering can construct the valid-lane prefix mask +``` + +Diagnostics: + +```text +VMI-UNSUPPORTED: pto.vmi.vdhist requires contiguous ui8 source, b8 mask, and +contiguous 256xui16 accumulator/result + +VMI-UNSUPPORTED: pto.vmi.vdhist final partial source chunk requires valid-lane +b8 mask materialization +``` + +`pto.vmi.vchist` has the same verifier and assignment requirements as `pto.vmi.vdhist`. +A5 hardware `chistv2` high-range semantics have been confirmed as **global cumulative** +(bin=1 result automatically accumulates bin0's total count), so `pto.vmi.vchist` lowers +via the same template as `pto.vmi.vdhist` — the only difference is emitting `pto.chistv2` +in place of `pto.dhistv2`. No software compensation is needed. + +If future hardware switches to range-local cumulative semantics, the chist pattern +will need a broadcast+add compensation path. In that scenario, introduce a +`-vmi-chist-mode` attribute or runtime probe as a separate evolution step. + +Reference lit tests: `vmi_to_vpto_chist.pto` (mirrors `vmi_to_vpto_dhist.pto`). + +Do not classify histogram as `group_reduce`. Its result location is selected +by source values, not by lane/group position, and its low/high split is caused +by the physical `128xui16` VPTO result width. + +Final hard gate： + +```text +no pto.vmi op remains +no !pto.vmi.* type remains, including in function signatures +no UnrealizedConversionCastOp remains +physical arity matches helper for every lowered value +``` + +Slice 4 完成条件： + +```text +1. `f16 -> f32 -> add -> store` lowers with deinterleaved=2 and stores contiguous logical order. + Covered by vmi_to_vpto_e2e_widen_add_store.pto. +2. `f8 -> f32 -> add -> store` lowers with deinterleaved=4 and stores contiguous logical order. + Covered by vmi_to_vpto_e2e_widen_add_store.pto. +3. Non-full memory physical arity and valid lane map are tested. + Covered by vmi_to_vpto_load_nonfull.pto, vmi_to_vpto_load_nonfull_memref.pto, + vmi_to_vpto_store_deint_invalid.pto, + vmi_to_vpto_load_safe_tail_memref.pto, + vmi_to_vpto_load_safe_tail_memref_negative_offset.pto, + vmi_to_vpto_masked_load_safe_tail_memref.pto, + vmi_to_vpto_masked_load_safe_tail_memref_negative_offset_invalid.pto, + vmi_to_vpto_expand_load_all_active.pto, + vmi_to_vpto_expand_load_all_active_negative_offset_invalid.pto, and multi-chunk load/store layout tests. +4. Full-footprint load/store direct path lowers through pto.vlds/pto.vsts or deinterleaved=2 x2 dist + instructions with offset 0. + Covered by the load/store direct-path and layout-folding tests. +5. Internal func.call boundaries expand callee signatures, call operands/results, and returned VMI values together. + Covered by vmi_layout_assignment_call_boundary.pto, vmi_layout_assignment_indirect_call_invalid.pto, + and vmi_to_vpto_call_boundary.pto. +6. Structured control-flow carrying VMI values expands iter args, yields, results, masks, and returns together. + Covered by vmi_layout_assignment_cf_switch.pto, + vmi_layout_assignment_scf_execute_region.pto, + vmi_layout_assignment_scf_index_switch.pto, + vmi_layout_assignment_scf_while.pto, vmi_to_vpto_cf_branch.pto, + vmi_to_vpto_scf_for.pto, vmi_to_vpto_scf_if.pto, and the user-facing + vmi_ptoas_cli_control_flow.pto. +7. Final gate rejects residual VMI helper and unrealized casts. + Covered by vmi_to_vpto_ensure_identity.pto, + vmi_to_vpto_ensure_layout_partial_invalid.pto, + vmi_to_vpto_truncf_fp8_128_contiguous_invalid.pto, + vmi_to_vpto_ensure_mask_layout_partial_invalid.pto, + vmi_to_vpto_unsupported_op_invalid.pto, + vmi_to_vpto_unrealized_cast_residual_invalid.pto, + vmi_to_vpto_type_attr_residual_invalid.pto, and per-feature unsupported + tests. +8. Same-family indirect memory ops reject unsupported direct-lowering shapes consistently. + Covered by vmi_to_vpto_gather_scatter_shape_invalid.pto together with the existing gather/scatter positive and + per-feature negative tests. +9. Same-family reduction ops reject unsupported direct-lowering shapes consistently. + Covered by vmi_to_vpto_reduce_shape_invalid.pto together with the existing reduce add/min/max positive and + per-feature tests, including vmi_to_vpto_reduce_addi_i16_invalid.pto for narrow integer rejection and + vmi_to_vpto_reduce_addf_f16.pto for f16 floating-point reduction lowering. +10. VMI op/type verifiers reject unsupported element types before OneToN rewriting. + Covered by vmi_to_vpto_bf16_arith.pto, vmi_to_vpto_math_element_type_invalid.pto, + vmi_to_vpto_cmp_select.pto, vmi_to_vpto_cmp_element_type_invalid.pto, + vmi_to_vpto_fma.pto, vmi_to_vpto_fma_element_type_invalid.pto, and + vmi_to_vpto_unary_math.pto for negf/absf/absi/sqrt/exp/ln/relu, plus + vmi_to_vpto_relu_element_type_invalid.pto. +11. Same-family mask logic ops lower through the physical mask granularity instead of assuming b32 masks. + Covered by vmi_to_vpto_mask_logic.pto for mask_and/mask_or/mask_xor/mask_not on b32 masks produced by + cmpf and on direct b8/b16 mask operands. +12. `pto.vmi.vdhist` lowers one logical 256-bin histogram into two VPTO low/high + bin-range histogram accumulator chains, and tail source chunks are masked + with a valid-lane b8 prefix. `pto.vmi.vchist` uses the same lowering template + as `pto.vmi.vdhist` (A5 hardware confirmed global cumulative semantics). + Covered by vmi_to_vpto_dhist.pto, vmi_to_vpto_dhist_tail_mask.pto, and + vmi_to_vpto_chist.pto. +``` + +## 7. Slice 5: Memory Padding + +The Slice 4 direct path lowers `pto.vmi.load` through plain `pto.vlds` when the +memory source itself is supported and the element type has a known physical lane +width. This includes non-full logical vectors; the operation is treated as a +direct full physical read of the selected VPTO chunk(s). Masked/expand/gather +read-like operations still use the richer access plan because their masks or +lane maps carry additional semantic constraints. + +Implement an internal `VMIMemoryAccessPlan`: + +```text +base +logical lane count +logical_shape +permutation_map +lane-to-address map in element units +validMask +paddingValue +safeReadProof +writeMask +target capability decision +fallback resource decision +``` + +Current implementation status: + +```text +lib/PTO/Transforms/VMIToVPTO.cpp + VMIMemoryAccessPlan + VMIMemorySafeReadProof + VMIMemoryLogicalShape + VMIMemoryLaneAddressMap + VMIMemoryFallbackDecision + +currently routed through the plan: + contiguous identity logical_shape/permutation/lane-to-address map in element units + explicit rejection of non-identity memref layouts until subview/affine lane maps are represented + covered by vmi_to_vpto_memref_layout_invalid.pto, including a memref.subview-produced strided view + subview diagnostics name the missing normalized base/offset/stride lane-to-address plan + target true masked/non-faulting load capability query + current result is missing capability because pto.vlds has no mask operand + covered by vmi_to_vpto_masked_load_nonfull_invalid.pto + stable gather masked-load option + covered by vmi_to_vpto_stable_gather_masked_load_todo_invalid.pto + currently emits a TODO diagnostic instead of lowering through VGATHER2 + direct pto.vmi.load source/layout capability check for full physical reads + pto.vmi.masked_load partial/tail safe full-read proof + pto.vmi.expand_load static all-active safe full-read proof + VMI-to-VPTO rewrite match guard for supported direct load sources + pto.vmi.store direct write target decision with all-true writeMask kind + pto.vmi.masked_store direct write target decision with explicit writeMask kind + unsafe masked/expand partial/tail read fallback decision as RequiredUnavailable diagnostic + covered by vmi_to_vpto_masked_load_nonfull_invalid.pto and + vmi_to_vpto_expand_load_all_active_negative_offset_invalid.pto + +currently not implemented by the plan: + paddingValue materialization (intentionally unsupported in the first implementation stage) + non-all-true validMask direct masked/non-faulting load lowering + scratch/guarded fallback lowering or allocation + lowering for non-identity logical_shape/permutation_map/lane-to-address maps, including subview or affine lane maps + writeMask fallback planning beyond the existing contiguous tail-store predicate path +``` + +Important first-stage contract: + +```text +VMI physical tail lanes and transfer paddingValue are different concepts. + +Physical tail lanes: + arise because pto.vreg is fixed at 256 bytes + are outside the logical VMI lane count + may be read/computed only when the extra lanes remain unobservable + +transfer_read-style paddingValue: + is an observable logical result for invalid/OOB transfer lanes + cannot be dropped or replaced by arbitrary physical tail contents + is not materialized by the first-stage VMI implementation + +Therefore any frontend path that still needs transfer_read paddingValue +semantics must stop before direct VMI-to-VPTO lowering with VMI-UNSUPPORTED, +unless it has already canonicalized to an all-valid load/masked_load subset +whose invalid lanes are proven absent. +``` + +Read-like memory decision tree： + +```text +safeReadProof full && validMask all true: + direct load + +safeReadProof full && validMask not all true: + first-stage: VMI-UNSUPPORTED because paddingValue materialization is not implemented + future: full load + padding materialization + select + +target true masked/non-faulting load: + first-stage: VMI-UNSUPPORTED because true masked/non-faulting load and paddingValue materialization are not implemented + future: masked load + padding materialization + +otherwise: + first-stage: VMI-UNSUPPORTED with the missing fallback reason + future: split safe regions, scratch fill/copy/load, guarded fallback, or diagnostic +``` + +Write-like memory decision tree： + +```text +writeMask all true && full footprint safe-writable: + direct store + +target true masked store: + masked store + +otherwise: + split/guarded/scatter-like fallback or diagnostic +``` + +Slice 5 完成条件： + +```text +1. Unsafe partial/tail read-like ops never lower to a potentially invalid full + read unless the physical footprint is statically proven safe. +2. PaddingValue materialization is not required in the first implementation + stage. Any path that would require paddingValue, true masked/non-faulting + load, scratch fill/copy/load, or guarded fallback must report + `VMI-UNSUPPORTED` with the missing fallback reason. +3. Non-identity logical_shape/permutation_map/lane-to-address maps, including + subview or affine lane maps, are explicitly rejected before lowering. +4. Store-like partial/tail writes are supported only by the existing + full-chunk or contiguous/deinterleaved tail-store predicate paths. Other + writeMask fallback paths must report `VMI-UNSUPPORTED`. +``` + +## 8. Layout Fact And Lowering Support Helpers + +Keep layout facts separate from layout assignment policy and VPTO lowering +choices. Shared layout helpers expose legal/preferred layout facts; they do +not select a global lowering plan and are not a target capability registry. + +```text +getPreferredCastLayoutFact(sourceType, resultType) +getPreferredGroupReduceLayoutFact(sourceType, numGroups) +canMaterializeDataLayout(sourceType, resultType) +canMaterializeMaskLayout(sourceType, resultType) +supportsMemoryAccessProof(proof) +supportsPrefixPopcount(maskType) +supportsReductionScanContract(op) +getScratchResource(plan) +``` + +Support and materialization helpers must expose actionable reasons. A pass must +not silently choose scalar fallback when fallback is disabled. + +Current implementation status: + +```text +include/PTO/Transforms/VMILayoutSupport.h +lib/PTO/Transforms/VMILayoutSupport.cpp + central table-driven source for legal/preferred layout facts: + dense/group load and store layouts + masked load/store data-mask layout relations + ensure data/mask layout materialization pairs + cast, bitcast, reduce, group_reduce, group_broadcast, histogram layouts + +lib/PTO/Transforms/VMIToVPTO.cpp + local op lowering support helpers: + true masked-load and fallback diagnostics for currently unimplemented paths + pointer-only constraints for concrete VPTO gather/stride/scatter paths + padding-safety and full-physical-chunk requirements + +still legacy helper-based and should migrate into layout/support tables when +they become layout facts: + full layout materialization plans and padding-safety checks + adjacent ppack/punpack mask granularity materialization plans + prefix popcount and full reduction/scan/contract shape checks +``` + +## 9. Diagnostics + +Centralize diagnostic codes in one header or utility file: + +```text +VMI-UNSUPPORTED +VMI-LAYOUT-CONTRACT +VMI-PASS-INVARIANT +VMI-RESIDUAL-OP +``` + +Current implementation defines these codes and their `": "` prefixes in `include/PTO/IR/VMIUtils.h`. Transform and +CLI code must reference those constants instead of spelling the diagnostic code strings locally; a source grep for the +four code strings should find only the central definitions. + +Every diagnostic should include: + +```text +source op +logical VMI type +producer natural layout, if any +consumer required layout, if any +missing capability or disabled option +available materialization paths, if known +``` + +## 10. Lit Test Layout + +Use a dedicated directory: + +```text +test/lit/vmi/ +``` + +Minimum test files: + +```text +vmi_type_attr_parse.mlir +vmi_type_attr_invalid.mlir +vmi_op_verifier_basic.mlir +vmi_producer_boundary.mlir +vmi_layout_assignment_widen.mlir +vmi_layout_assignment_cfg.mlir +vmi_layout_assignment_broadcast_remat.mlir +vmi_layout_assignment_iota_remat.mlir +vmi_layout_assignment_mask_remat.mlir +vmi_to_vpto_deinterleaved2.mlir +vmi_to_vpto_deinterleaved4.mlir +vmi_to_vpto_compaction_deint_invalid.mlir +vmi_to_vpto_load_safe_tail_memref.mlir +vmi_to_vpto_masked_load_safe_tail_memref.mlir +vmi_to_vpto_store_tail.mlir +vmi_to_vpto_dhist.mlir +vmi_to_vpto_dhist_tail_mask.mlir +vmi_to_vpto_chist.mlir +vmi_pipeline_hard_gates.mlir +``` + +Each pass test must use `FileCheck` to prove both positive output and negative absence: + +```text +CHECK: pto.vmi.addf +CHECK-NOT: pto.vadd +CHECK-NOT: unrealized_conversion_cast +``` + +Final lowering tests must check: + +```text +CHECK-NOT: pto.vmi. +CHECK-NOT: unrealized_conversion_cast +``` + +## 11. Implementation Order + +Recommended merge order: + +```text +1. VMI type/attr + helper + parse/verify tests. +2. Slice 1 op shells + verifier tests. +3. VMI producer boundary verifier. +4. layout assignment for straight-line code. +5. layout assignment for scf/cf/function boundaries. +6. vmi-to-vpto type conversion + pack/unpack/unpackable block args. +7. deinterleaved=2 f16 widen end-to-end. +8. deinterleaved=4 f8 widen end-to-end. +9. load/store padding-safe lowering. +10. remaining semantic op families. +``` + +Do not merge a pass that leaves hidden side tables as a required interpretation mechanism. Temporary internal +analysis structures are fine only if the pass materializes the final state into IR before returning. + +## 12. Review Checklist Before Coding Each Slice + +Before implementation: + +```text +1. Is the op/type syntax written in ODS and tested by parser round-trip? +2. Does every verifier rule have a negative test? +3. Does every pass have a post-pass hard gate? +4. Are CFG block arguments and function signatures covered? +5. Does any lowering rely on a defining op that block arguments do not have? +6. Does memory lowering prove safe footprint separately from valid lane mask? +7. Does mask granularity follow consumer element width? +8. Does final VPTO lowering leave zero VMI op/type/helper or unrealized-cast residuals? +``` + +If any answer is no, the slice is not ready to be treated as complete. + +## 13. Adding One VMI Op End To End + +新增一个 `pto.vmi.*` op 时，不要只补 ODS 和 lowering pattern。它必须穿过固定的七个落点， +否则很容易出现 verifier 能过、layout pass 不知道怎么约束、或控制流 physicalization 后残留 VMI type。 + +```text +1. ODS surface: + include/PTO/IR/VMIOps.td + +2. semantic verifier: + lib/PTO/IR/VMI.cpp + +3. layout assignment facts: + lib/PTO/Transforms/VMILayoutAssignment.cpp + +4. shared layout support, when the fact crosses stages: + include/PTO/Transforms/VMILayoutSupport.h + lib/PTO/Transforms/VMILayoutSupport.cpp + +5. vmi-to-vpto preflight: + lib/PTO/Transforms/VMIToVPTO.cpp::verifySupportedVMIToVPTOOps + +6. OneToN lowering pattern: + lib/PTO/Transforms/VMIToVPTO.cpp::populateVMIOneToNConversionPatterns + +7. focused lit tests: + test/lit/vmi/ +``` + +这七个落点的职责不同： + +```text +ODS: + 只定义 op 形状、operand/result type 类别、assembly format、interface 和 verifier hook。 + +VMI.cpp verifier: + 检查局部语义，例如元素类型、rank、lane count、predicate 字符串、source/result bit 数关系。 + 不能依赖 def-use 图，不能决定 layout。 + +LayoutAssignment: + 只收集 value-level layout/granularity 事实： + - producer natural layout + - operands that must share layout with result + - consumer required layout + - mask consumer required granularity + 不能在 collect 阶段改 IR。 + +VMILayoutSupport: + 只放跨 assignment、validation、optimization、lowering 中至少两个阶段共享的纯查询。 + 典型内容是 cast layout fact、group_reduce layout fact、ensure_* materialization support。 + 不能返回 VPTO instruction sequence、不能决定 clone/rematerialize、不能读取 producer/user context。 + 只有一个 lowering pattern 自己使用的判断不要抽到这里。 + +VMIToVPTO preflight: + 在 rewrite 前拒绝当前 lowering 不支持但语义合法的 case。 + 典型例子是 partial physical chunk、non-prefix mask constant、dynamic create_mask、unsupported shuffle。 + +OneToN pattern: + 从 adaptor 读取 physical parts，按已经确定的 layout 发 VPTO op。 + 不能重新推断 layout，也不能通过 defining op 找 physical parts。 + +lit: + 至少覆盖 parser/verify、layout assignment、positive lowering、negative unsupported diagnostic。 +``` + +### Layout Fact Template + +新增 op 时先给它归类，再写 layout 约束。不要从 VPTO 指令形态反推 VMI layout；layout 的来源必须是 +logical vector 语义和当前物理指令的天然限制。 + +```text +elementwise same-shape op: + examples: + addf/addi/subf/mulf/andi/shli/shrui/shrsi/absf/absi/sqrt + layout rule: + all data operands and result are in one equivalence class + lowering rule: + emit one VPTO op per physical part + +compare op: + examples: + cmpf/cmpi + layout rule: + lhs/rhs data layout unified + result mask requested to the same data layout + result mask granularity comes from lhs/rhs element width + lowering rule: + emit one vcmp per data part, producing corresponding mask part + +mask logical op: + examples: + mask_and/mask_or/mask_xor/mask_not + layout rule: + all mask operands/results share layout and granularity + lowering rule: + emit one predicate op per physical mask part + +layout-changing producer: + examples: + extf f16->f32, extf f8->f32, truncf f32->f16, truncf f32->fp8-like + layout rule: + source/request side follows instruction input contract + result natural layout follows instruction output contract + lowering rule: + emit the instruction sequence that preserves logical lane order under that layout + +memory consumer/producer: + examples: + load/store/load/store + layout rule: + load result natural layout is chosen by memory dist capability + store value operand requests the layout that memory dist can consume + lowering rule: + direct path only when every physical chunk has no padding lane and footprint is safe + +structural boundary: + examples: + scf.if result/yield, scf.for iter args, cf.br successor operands, func.call + layout rule: + semantically identical incoming/outgoing values are unified + lowering rule: + handled by OneToN structural patterns, not by op semantic lowering +``` + +代码里 `LayoutSolver::addConstraints()` 应该只表达上面的事实。例如一个普通 elementwise binary op +只需要： + +```cpp +if (auto addf = dyn_cast(op)) { + if (failed(unite(addf.getLhs(), addf.getRhs(), op)) || + failed(unite(addf.getLhs(), addf.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); +} +``` + +一个 layout-changing op 不应该把 source/result 直接 `unite`，而是明确写 producer/consumer 合同： + +```cpp +if (auto extf = dyn_cast(op)) { + requestDataUse(extf.getSourceMutable(), getContiguousLayout()); + if (failed(setNaturalLayout(extf.getResult(), + VMILayoutAttr::getDeinterleaved(ctx, factor), + op))) + return WalkResult::interrupt(); + return WalkResult::advance(); +} +``` + +### OneToN Pattern Template + +`vmi-to-vpto` pattern 的输入不再是 logical VMI value，而是 adaptor 里已经 flatten 好的 physical parts。 +pattern 只做三件事： + +```text +1. 从 adaptor 取每个 logical operand 的 physical part list。 +2. 从 resultMapping 取每个 logical result 对应的 physical result type list。 +3. 按 part 顺序创建 VPTO op，并用 resultMapping replace 原 op。 +``` + +普通 elementwise binary op 的代码形态应该接近： + +```cpp +LogicalResult matchAndRewrite(VMIAddFOp op, OpAdaptor adaptor, + OneToNPatternRewriter &rewriter) const override { + ValueRange lhsParts = adaptor.getLhs(); + ValueRange rhsParts = adaptor.getRhs(); + TypeRange resultTypes = adaptor.getResultMapping().getConvertedTypes(0); + + if (lhsParts.size() != rhsParts.size() || lhsParts.size() != resultTypes.size()) + return rewriter.notifyMatchFailure(op, "physical arity mismatch"); + + SmallVector results; + for (auto [lhs, rhs, resultType] : llvm::zip_equal(lhsParts, rhsParts, resultTypes)) + results.push_back(rewriter.create(op.getLoc(), resultType, lhs, rhs)); + + rewriter.replaceOp(op, results, adaptor.getResultMapping()); + return success(); +} +``` + +这里不能调用 `op.getLhs().getDefiningOp()` 去找物理寄存器。原因是 VMI value 可以来自： + +```text +function argument +block argument +scf.for iter arg +scf.if result +cf.br successor argument +func.call result +``` + +这些 value 很多没有 VMI defining op。physical parts 的唯一合法来源是 OneToN adaptor 和 +OneToNTypeMapping。 + +### Control-Flow Checklist + +每新增一个 op，不一定要写新的控制流 pattern；但必须检查它的结果或 operand 是否可能跨边界。 +如果只是普通 VMI value，那么已有 structural OneToN pattern 应该负责边界 physicalization： + +```text +func.func / func.call / func.return: + upstream func OneToN conversion + +scf.if / scf.for / scf.while / scf.yield: + upstream SCF OneToN structural conversion plus layout solver equivalence constraints + +cf.br / cf.cond_br / cf.switch: + project-local OneToN patterns flatten successor operands and rewrite destination block signatures + +scf.execute_region / scf.index_switch: + project-local OneToN patterns flatten region results +``` + +新增 op 的测试要至少放一个跨边界用例，证明 op 的 result 不是只在 straight-line IR 中工作： + +```mlir +%r = scf.if %cond -> !pto.vmi.vreg<128xf32> { + %x = pto.vmi.addf %a, %b : ... -> !pto.vmi.vreg<128xf32> + scf.yield %x : !pto.vmi.vreg<128xf32> +} else { + scf.yield %c : !pto.vmi.vreg<128xf32> +} +pto.vmi.store %r, %ptr, %off : ... +``` + +对应 lowering test 必须检查： + +```text +CHECK-NOT: pto.vmi. +CHECK-NOT: !pto.vmi. +CHECK-NOT: unrealized_conversion_cast +``` + +如果这个测试失败，通常不是该 op 的 VPTO pattern 本身错，而是 layout assignment 没有把 yield/result/consumer +约束统一，或者 OneToN structural pattern 漏了某种 region/control-flow op。 + +### Preflight Versus Pattern Failure + +语义合法但当前还没有物理实现的 case，应该在 `verifySupportedVMIToVPTOOps()` 里给稳定 diagnostic， +不要让 pattern 随机 `notifyMatchFailure()` 后落成 generic conversion failure。 + +```text +use verifier failure: + op 本身语义非法，任何 target 都不应该接受。 + examples: + absf on integer element + shrui on signed integer element + shrsi on unsigned integer element + bitcast total bits mismatch + +use VMI-LAYOUT-CONTRACT: + 多个 producer/consumer/control-flow 约束互相冲突。 + examples: + one value simultaneously required as contiguous and deinterleaved=2 + one mask simultaneously required as b16 and b32 + +use VMI-UNSUPPORTED in preflight: + VMI semantics are valid, but current VPTO materialization is not implemented. + examples: + partial/tail memory access + pred-only constant mask without concrete b8/b16/b32 granularity + shuffle that requires vselr index-vector materialization + bitcast with mismatched layouts or per-chunk logical bit footprints + +use VMI-RESIDUAL-OP: + conversion framework finished but VMI op/type/helper/cast remains. + This is a pass bug or missing pattern, not a user semantic error. +``` + +Pattern-local `notifyMatchFailure()` is still useful for debugging competing patterns, but it must not be the only +user-visible explanation for a known unsupported VMI semantic case. diff --git a/docs/designs/vmi-introduction.md b/docs/designs/vmi-introduction.md new file mode 100644 index 0000000000..1ff03f2c8e --- /dev/null +++ b/docs/designs/vmi-introduction.md @@ -0,0 +1,1076 @@ +# VMI 介绍 + +本文介绍 VMI 的设计入口：VMI 解决什么问题，layout 有哪些，pass pipeline +如何分工，以及这些机制分别应对哪些典型场景。更完整的逐 case lowering 结果见 +`docs/designs/vmi-layout-lowering-cases.md`。 + +示例是设计级 IR，保留关键 type、layout、helper op 和 VPTO op 形状， +省略 module wrapper、完整 operand list 和不影响讨论的 SSA 细节。 + +## 1. VMI 表达什么 + +VMI 是 VPTO 之前的逻辑向量层。它让前端先表达“我要对 `NxT` 的逻辑向量做什么”， +再由 layout assignment 决定这个逻辑向量如何拆到 256B 物理 vector register 上。 +当 VPTO 指令因为物理 register 宽度只能暴露半宽接口时，VMI 也负责提供完整的 +逻辑语义。例如 `ui8` histogram 的完整结果是 `256xui16`，物理 VPTO histogram +一次只能返回 `128xui16`；VMI surface 应该表达完整 histogram，low/high bin +range 拆分属于 lowering 细节。 + +Surface VMI 类型不携带布局： + +```mlir +!pto.vmi.vreg<128xf32> +!pto.vmi.mask<128xpred> +``` + +Layout-assigned VMI 类型携带具体布局和 mask granularity： + +```mlir +!pto.vmi.vreg<128xf32, #pto.vmi.layout> +!pto.vmi.mask<128xb32, #pto.vmi.layout> +``` + +VMI 的核心约束是：`vmi-to-vpto` 只从当前 op 的 attrs、operands、types、 +layouts 和显式 helper ops 做 lowering，不读取隐藏 plan/recipe，也不通过 +defining op 或 sibling user 恢复上下文。 + +## 2. Layout 类型 + +### 2.1 `contiguous` + +```mlir +#pto.vmi.layout +``` + +含义：logical lane 按顺序落入物理 register list。 + +```text +logical lanes: 0 1 2 ... 63 | 64 65 ... 127 +physical part: p0 | p1 +``` + +典型场景： + +```text +dense load/store +普通 elementwise compute +一个 group 天然适配当前 reduce op 时的 reduction input +caller/callee 约定 dense order 时的 control-flow/function boundary +``` + +### 2.2 `deinterleaved = F, block_elems = B` + +```mlir +#pto.vmi.layout +#pto.vmi.layout +``` + +`block_elems` 缺省为 `1`。逻辑 lane 到物理 part 的映射是： + +```text +logical lane i +block q = i / B +in-block lane r = i % B +part p = q % F +part block t = q / F + +physical part p, physical lane t * B + r +``` + +`deinterleaved=2` 的直观例子： + +```text +logical lanes: 0 1 2 3 4 5 ... +physical part0: 0 2 4 ... +physical part1: 1 3 5 ... +``` + +`deinterleaved=4, block_elems=8` 的直观例子： + +```text +logical group S=32: + lanes 0.. 7 -> part0 lanes 0..7 + lanes 8..15 -> part1 lanes 0..7 + lanes 16..23 -> part2 lanes 0..7 + lanes 24..31 -> part3 lanes 0..7 +``` + +典型场景： + +```text +f16 -> f32: + vcvt 天然产生 even/odd 两个 f32 part，所以结果使用 deinterleaved=2。 + +f32 -> f16: + vcvt 需要 f32 source 先拆成 even/odd 两个 part，所以 source 使用 + deinterleaved=2。 + +S=32 group_reduce f32: + 一个 group 有 32 个 f32 element。高效 reduce path 消费四个 8-lane block， + 所以 source/mask 使用 deinterleaved=4, block_elems=8。 +``` + +`block_elems=8` 表示一种按 32B row fragment 组织的输入形态，不表示 +S=32 reduce 只能接受这一种形态。如果同一个 value 还要服务 narrow cast 等 +element-parity consumer，assignment 可以选择 `deinterleaved=4, block_elems=1` +作为共同 layout，再由 lowering 生成对应的物理指令序列。 + +`deinterleaved` 只描述最终物理 part 中有哪些 logical lane，不描述这个 layout +由哪条指令生成。不同 producer 可以用不同方式直接产生同一个 layout；如果不能 +直接产生，后续 lowering 再通过显式 materialization helper 把 source layout +转换成 consumer 需要的 layout。具体 lowering 形状见 case catalog。 + +### 2.3 `num_groups = G, slots = K` + +```mlir +#pto.vmi.layout +#pto.vmi.layout +#pto.vmi.layout +``` + +这是 group-slot result layout。它不表示全部 `N` 个 logical lane 都有语义值。 +只有 `G` 个 group 结果 slot 有语义值。 + +```text +slot_block(g) = g / K +slot_lane(g) = (g % K) * lane_stride + +physical part slot_block(g) 的 lane slot_lane(g) 保存 group g 的结果 +``` + +`lane_stride` 缺省为 1，单位是 logical element-sized physical slot。 +它描述 group result 在物理存储中的固定间距，不改变 VMI 的逻辑元素类型。 +例如 `ui8 lane_stride=4` 表示 group slot 存在 byte lane 0, 4, 8, ... +这种形态可以 lower 为 `PK4_B32` store，物理上使用 b32 carrier 的 low byte。 + +`num_groups=16, slots=8` 的例子： + +```text +part0 lane0..7 = group result 0..7 +part1 lane0..7 = group result 8..15 +other lanes = 对普通 dense consumer 来说未定义 +``` + +为什么 group 信息也要放进 layout： + +```text +group_reduce 自身有 num_groups，但它的结果可能继续跨过 truncf、 +group_broadcast、group_store、scf.if、scf.for、function call 或多个 consumer。 + +这些后续 op 不应该回看 producer attr。value layout 因此需要记录有多少个 +group result，以及这些 result 如何 packed 到 physical slot。 +``` + +典型场景： + +```text +group_reduce result +group_slot_load result +group_store input +group_broadcast input +group-slot control-flow/function boundary +部分 row-local cast 路径，通常使用 slots=1 +``` + +## 3. Pass Pipeline + +```text +pto-validate-vmi-ir + -> vmi-layout-assignment + -> canonicalize/cse + -> vmi-layout-fold + -> canonicalize/cse + -> vmi-layout-rematerialize + -> canonicalize/cse + -> vmi-layout-sink-materialization + -> canonicalize/cse + -> vmi-legalize-arith-select + -> pto-validate-vmi-layout-ir + -> vmi-to-vpto +``` + +### 3.1 `pto-validate-vmi-ir` + +检查 surface VMI 边界。 + +合法输入： + +```mlir +%x = pto.vmi.load %src[%off] + : !pto.ptr -> !pto.vmi.vreg<128xf16> +``` + +非法输入： + +```mlir +%x = pto.vmi.load %src[%off] + : !pto.ptr + -> !pto.vmi.vreg<128xf16, #pto.vmi.layout> +``` + +原因：具体 layout 由 `vmi-layout-assignment` 产生，不应该由 surface frontend +提前写入。 + +### 3.2 `vmi-layout-assignment` + +这是硬合法化 pass。它选择具体 value layout、具体 mask granularity， +并在 layout 不匹配的 use site 插入显式 helper op。 + +这个 pass 的工作顺序是固定的： + +```text +1. 做少量 VMI 内部规整，让后续 layout 规则面对稳定形态。 +2. 为 data value 建 union-find 求解器，并收集 data 约束和 data use request。 +3. 把可采纳的 consumer request 提升为 producer/result 的最终 layout。 +4. 改写所有 data value type，让 !pto.vmi.vreg 携带具体 layout。 +5. 对仍不匹配的 data use 插入 pto.vmi.ensure_layout。 +6. 基于已经确定的 data layout 推导 mask layout 和 predicate granularity。 +7. 改写所有 mask type，并对不匹配的 mask use 插入 ensure_mask_*。 +8. 同步更新 function type、call boundary 和 block argument type。 +9. 校验 layout-assigned VMI IR。 +``` + +Data 和 mask 分两轮求解。原因是 mask layout 通常依赖对应 data operand 或 result +的 layout；例如 `cmpf` 产生的 mask 跟比较输入的 data layout 对齐， +`select`/`reduce`/`masked_load` 消费的 mask 也要跟对应 data value 的 lane +layout 和元素 bitwidth 对齐。 + +Data 求解器为每个 `!pto.vmi.vreg` 建一个节点： + +```text +DataNode: + value = 对应 SSA value + original type = surface VMI type + parent = union-find parent + naturalLayout = 当前等价类选择的自然 layout，可能为空 +``` + +遍历 IR 时，每个 op 向 data 求解器贡献三类信息。 + +第一类是 layout 等价约束。它表示几个 value 必须使用同一个 physical layout， +也就是 union-find 中的同一个等价类。典型来源： + +```text +layout-transparent elementwise: + addf/addi/subf/subi/mulf/muli/fma/divf/minf/maxf/... + L(operands...) = L(result) + +unary elementwise: + negf/absf/absi/sqrt/exp/ln/relu/not + L(source) = L(result) + +select: + L(true_value) = L(false_value) = L(result) + +bitcast: + L(source) = L(result) + +structured control flow: + scf.if result = then/else yield operand + scf.for result = init operand = iter_arg = yield operand + scf.while result = init/before/condition/after/yield carried value + +cf branch: + branch operand = destination block argument + +function boundary: + call operand = callee argument + call result = callee return operand + multiple returns of the same function agree per result index +``` + +这一步只说明“这些 value 如果存在布局，就必须一致”。它不等价于把某个 +consumer 的 request 无条件推过所有 producer 或控制流。 + +等价类可以画成“同一个框里的 value 共用一个 layout 变量”。例如普通 +elementwise 链： + +```text +surface VMI: + + %x = pto.vmi.load ... + %k = pto.vmi.broadcast ... + %y = pto.vmi.mulf %x, %k + %q = pto.vmi.truncf %y + +data layout 等价类: + + class C0 + +--------------------------------------+ + | %x %k %y | + | load broadcast mulf result | + +--------------------------------------+ + ^ + | + use request from truncf source: + wants deinterleaved=4 + +若 %y 的 producer chain 可采纳该 request，assignment 可以选择: + + L(C0) = deinterleaved=4 +``` + +控制流 join 也是等价类，但 request adoption 的含义不同： + +```text +surface VMI: + + %y = scf.if %c -> !pto.vmi.vreg<128xf32> { + scf.yield %a + } else { + scf.yield %b + } + %q = pto.vmi.truncf %y + +data layout 等价类: + + class C1 + +--------------------------------------+ + | %a %b %y | + | then yield else yield if result | + +--------------------------------------+ + ^ + | + use request from truncf source: + wants deinterleaved=4 + +scf.if result 不是 consumer-driven adoption 的可采纳 producer。 +若 C1 不能直接选择 deinterleaved=4，assignment 保持 C1 的布局， +并在 use site materialize: + + %y_for_q = pto.vmi.ensure_layout %y : L(C1) -> deinterleaved=4 + %q = pto.vmi.truncf %y_for_q +``` + +多 consumer 冲突时，等价类仍然只有一个 layout： + +```text +surface VMI: + + %y = pto.vmi.mulf %x, %k + pto.vmi.store %y, %out0 + %q = pto.vmi.truncf %y + +data layout 等价类: + + class C2 + +-----------------------------+ + | %x %k %y | + +-----------------------------+ + |\ + | \ use request from truncf: deinterleaved=4 + | + +--- use request from store: contiguous + +两个 use request 不一致时，不能让 %y 同时拥有两个 layout。 +baseline assignment 保留 C2 已有的 natural layout；若没有 natural layout， +则使用默认 contiguous。与该 layout 不匹配的 edge 会插 ensure_layout。 +``` + +第二类是 result 自然布局。某些 op 的结果本身有目标相关的自然布局： + +```text +普通 reduce / compress / shuffle: + result 通常是 contiguous。 + +group_reduce: + source 需要适配 group reduce 指令形态； + result 使用 group_slots(num_groups, slots) 描述 group-slot result。 + +cast: + widening/narrowing 根据 cast support 决定 source request 和 result layout。 + +group_load / group_slot_load / group_broadcast_load: + result 根据 group size、row stride 和目标能力选择 contiguous、deinterleaved + 或 group_slots。group_broadcast_load 表达“每个 logical group load 一个值并 + 广播到组内 lanes”的逻辑语义；E2B 只是兼容 layout 下的一种 lowering。 + +stride_load: + result 是 contiguous。block/repeat stride 只描述 memory address map， + 不改变 register 内 logical lane order。 + +active_prefix_index: + result 使用 contiguous。 +``` + +若同一个等价类已经有自然布局，再设置不同自然布局会报 layout contract 冲突。 + +第三类是 operand 使用请求。consumer 不直接修改 operand 的 type，而是记录 +“这个 use site 希望 operand 是什么 layout”： + +```text +store / masked_store value: + wants contiguous + +ordinary reduce source/init: + wants contiguous + +group_reduce source: + wants preferred group-reduce source layout + +group_store value: + wants preferred group result layout + +stride_store value: + wants contiguous。block/repeat stride 只描述 memory write address map， + 不表示 source vreg 是 lane-strided 或 NZ layout。 + +truncf/trunci/extf/extsi/extui source: + wants cast support 给出的 source layout + +channel_split / channel_merge / shuffle: + wants 各自 lowering 需要的 source/input layout +``` + +收集完这些信息后，assignment 才尝试做 consumer-driven adoption。它逐个查看 +use request：如果 operand 的 producer 可以直接用 consumer 需要的 layout 产生 +同一个逻辑向量，并且多 use 时所有 use 都请求同一个 layout，那么这个 request +会被提升为该 value 所在 data 等价类的最终 layout。 + +可采纳 producer 是受限集合： + +```text +load +broadcast / constant / iota +layout-transparent elementwise +select +bitcast +``` + +这就是 request 看起来能穿过 elemwise 的原因： + +```mlir +%x = pto.vmi.load ... +%k = pto.vmi.broadcast ... +%y = pto.vmi.mulf %x, %k +%q = pto.vmi.truncf %y +``` + +`mulf` 先把 `%x`、`%k`、`%y` 合成同一个 data 等价类。`truncf` 对 `%y` +的 source use 请求 `deinterleaved=4` 时，这个 request 作用到 `%y` 所在等价类； +因为 `mulf` 是可采纳 producer，assignment 可以把整个等价类选成 +`deinterleaved=4`，从而让 load/broadcast/mulf 直接在这个 layout 下产生数据。 + +控制流边界也会形成等价类，但它不是任意 request 的自动传播通道： + +```mlir +%y = scf.if %c -> !pto.vmi.vreg<128xf32> { + scf.yield %a +} else { + scf.yield %b +} +%q = pto.vmi.truncf %y +``` + +`%y`、`%a`、`%b` 的 layout 必须一致；但 `scf.if` result 本身不是 +consumer-driven adoption 的可采纳 producer。若 `%q` 需要的 layout 无法成为 +这个等价类的最终布局，assignment 会在 `%q` 的 use site 插 +`pto.vmi.ensure_layout`，而不是隐式重写两个 branch 的内部计算。 + +Data layout 确定后，pass 会把每个 `!pto.vmi.vreg` 改写成 +`!pto.vmi.vreg`。如果某个记录过的 use request 仍然和 operand +当前 layout 不一致，pass 在该 consumer 前插显式 materialization： + +```mlir +%x_req = pto.vmi.ensure_layout %x + : !pto.vmi.vreg + -> !pto.vmi.vreg +consumer %x_req +``` + +这个规则也处理多 consumer 冲突： + +```mlir +%y = pto.vmi.mulf %x, %k +pto.vmi.store %y, %out0 // wants contiguous +%q = pto.vmi.truncf %y // wants deinterleaved=4 source +``` + +一个 SSA value 只能属于一个 data layout 等价类。若两个 use 不能共同满足， +baseline assignment 保留一个等价类 layout，并在不匹配 use 前插 +`ensure_layout`。后续 `vmi-layout-fold`、`vmi-layout-rematerialize` +和 `vmi-layout-sink-materialization` 可以在显式 helper op 上做优化，但 +`vmi-to-vpto` 不读取隐藏 plan 或 sibling user。 + +Mask 求解发生在 data type 改写之后。它同样维护 union-find 等价类，但节点记录 +两件事： + +```text +mask layout +predicate granularity: b8 / b16 / b32 +``` + +mask request 从已经带 layout 的 data value 推导： + +```text +cmpf/cmpi result: + mask layout = lhs data layout + granularity = lhs element bitwidth 对应的 predicate 粒度 + +select mask: + mask layout = result data layout + granularity = result element bitwidth 对应的 predicate 粒度 + +reduce / group_reduce / masked_load / expand_load mask: + mask layout = source/result data layout + granularity = 对应 data element bitwidth 的 predicate 粒度 +``` + +若 mask use 的 layout 或 granularity 不匹配，pass 显式插 +`pto.vmi.ensure_mask_layout` 或 `pto.vmi.ensure_mask_granularity`。 + +完成 data/mask 改写和 helper 插入后，pass 会同步更新 function type。直接 +internal call 会把 call operand/result 与 callee argument/return operand 合成 +同一布局约束；带 VMI type 的 external declaration 或 indirect call 没有可见 +body，当前需要显式 ABI materialization 设计，因此 layout assignment 会拒绝。 +这个阶段之后，IR 不再依赖隐藏 plan；后续 pass 和 `vmi-to-vpto` 都只读取 type +上的 layout 和显式 `ensure_*` helper。 + +### 3.3 `vmi-layout-fold` + +当 consumer 可以直接保持同样的外部效果时，把显式 materialization 折进 +consumer。 + +变换前： + +```mlir +%dense = pto.vmi.ensure_layout %x + : !pto.vmi.vreg<128xf32, #pto.vmi.layout> + -> !pto.vmi.vreg<128xf32, #pto.vmi.layout> +pto.vmi.store %dense, %dst[%off] +``` + +变换后： + +```mlir +pto.vmi.store %x, %dst[%off] + : !pto.vmi.vreg<128xf32, #pto.vmi.layout>, !pto.ptr +``` + +可能的 VPTO 形状： + +```text +fold 前：vintlv + vsts + vsts +fold 后：vstsx2，使用交错 store mode +``` + +### 3.4 `vmi-layout-rematerialize` + +通过 clone 低成本、layout-polymorphic 的 producer 来替换 `ensure_*`。 + +变换前： + +```mlir +%s = pto.vmi.broadcast %scale + : f32 -> !pto.vmi.vreg<128xf32, #pto.vmi.layout> +%s_split = pto.vmi.ensure_layout %s + : !pto.vmi.vreg<128xf32, #pto.vmi.layout> + -> !pto.vmi.vreg<128xf32, #pto.vmi.layout> +``` + +变换后： + +```mlir +%s_split = pto.vmi.broadcast %scale + : f32 -> !pto.vmi.vreg<128xf32, #pto.vmi.layout> +``` + +预期可 rematerialize 的 producer： + +```text +splat constant +broadcast +iota +create_mask +create_group_mask +constant_mask +``` + +这个 pass 不 rematerialize： + +```text +load / masked_load / group_load / group_slot_load / group_broadcast_load +stride_load +reduce / group_reduce +control-flow results +``` + +### 3.5 `vmi-layout-sink-materialization` + +把匹配的 layout 转换跨过 layout-transparent elementwise op。 + +变换前： + +```mlir +%a_dense = pto.vmi.ensure_layout %a : deinterleaved=2 -> contiguous +%b_dense = pto.vmi.ensure_layout %b : deinterleaved=2 -> contiguous +%y_dense = pto.vmi.addf %a_dense, %b_dense : contiguous +``` + +变换后： + +```mlir +%y_split = pto.vmi.addf %a, %b : deinterleaved=2 +%y_dense = pto.vmi.ensure_layout %y_split : deinterleaved=2 -> contiguous +``` + +效果： + +```text +两个 input materialization -> 一个 result materialization +``` + +这个 pass 不会 sink 穿过 cast、load、store、reduce、group_broadcast 或 +control-flow op。 + +### 3.6 `vmi-legalize-arith-select` + +Canonicalization 可能把简单的 `scf.if` 折成 `arith.select`。VMI 希望把 +control-flow lowering 保持在结构化控制流里，所以这个 pass 会把 VMI value 上的 +`arith.select` 改回 `scf.if`。 + +```mlir +%r = arith.select %cond, %a, %b + : !pto.vmi.vreg<128xf32, #pto.vmi.layout> +``` + +改成： + +```mlir +%r = scf.if %cond + -> !pto.vmi.vreg<128xf32, #pto.vmi.layout> { + scf.yield %a : !pto.vmi.vreg<128xf32, #pto.vmi.layout> +} else { + scf.yield %b : !pto.vmi.vreg<128xf32, #pto.vmi.layout> +} +``` + +### 3.7 `pto-validate-vmi-layout-ir` + +检查 post-assignment gate： + +```text +每个 VMI 数据值都有 concrete layout +每个 VMI mask 都有 concrete granularity 和 layout +helper op 有支持的 materialization path +semantic op/layout 组合有支持的 local lowering +vmi-to-vpto 之前没有物理 VPTO value 泄漏到 VMI IR 中 +``` + +非法例子： + +```mlir +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8, reassoc} + : ... -> !pto.vmi.vreg<8xf32, #pto.vmi.layout> + +pto.vmi.store %sum, %dst[%off] + : !pto.vmi.vreg<8xf32, #pto.vmi.layout>, + !pto.ptr +``` + +原因： + +```text +dense store 不能把 group_slots 当 dense vector 读取。 +应使用 group_store、group_broadcast 或显式支持的 group-to-dense op。 +``` + +### 3.8 `vmi-to-vpto` + +把 layout-assigned VMI value 转换成有序物理 VPTO value 列表，并对每个 +VMI op 做 local lowering。 + +例子： + +```text +!pto.vmi.vreg<128xf32, #pto.vmi.layout> + -> 两个 physical !pto.vreg<64xf32> part + +!pto.vmi.vreg<128xf32, #pto.vmi.layout> + -> 两个 physical !pto.vreg<64xf32> part + part0 携带 even lanes，part1 携带 odd lanes + +!pto.vmi.vreg<32xf32, #pto.vmi.layout> + -> 四个 physical part + part0 携带 group 0..7，part1 携带 group 8..15，... +``` + +`VMILayoutSupport` 不是 pass。它是 assignment、validation、optimization 和 +lowering 共享的查询库，用来避免重复实现 layout fact 和 supported +materialization 检查。 + +## 4. 典型场景 + +### 4.1 Dense Cast 与 Store + +```text +surface: + load f16，语义上连续 + extf 到 f32 + dense store f32 + +assignment: + load result = contiguous + extf result = deinterleaved=2 + store use = ensure_layout(deinterleaved=2 -> contiguous) + +baseline VPTO: + vlds + vcvt even / vcvt odd + vintlv + vsts + vsts + +fold-consumers 后的优化 VPTO: + vlds + vcvt even / vcvt odd + vstsx2，使用 interleaving store +``` + +这个场景说明为什么需要 `deinterleaved=2`，以及为什么 store-consumer folding +有价值。 + +### 4.2 Narrow Cast 与 Store + +```text +surface: + load f32 + truncf 到 f16 + dense store f16 + +assignment: + load result = deinterleaved=2 + truncf result = contiguous + +VPTO: + vldsx2 deinterleaving load + vcvt even / vcvt odd + vor + vsts +``` + +这个场景说明 memory op 可以直接产生 consumer 需要的 layout，但不需要保存隐藏 +plan。 + +### 4.3 一个 Producer 同时服务 Dense 和 Group Consumer + +```mlir +%x32 = pto.vmi.extf %x16 +%sum = pto.vmi.group_reduce_addf %x32, %mask {num_groups = 8, reassoc} +pto.vmi.group_store %sum, %sum_out[%off], %c1 {num_groups = 8} +pto.vmi.store %x32, %dense_out[%off] +``` + +Assignment 形状： + +```text +%x32 layout = deinterleaved=2 +group_reduce 直接消费 %x32 +dense store 获得 ensure_layout(%x32 -> contiguous) +``` + +VPTO 形状： + +```text +vcvt even/odd +vcgadd + vcgadd + vadd -> group_store result +vintlv + dense stores -> 产生 dense store 结果 +``` + +这个场景说明为什么需要 use-site materialization。producer 不需要选择一个能同时 +满足所有 consumer 的唯一 layout。 + +### 4.4 按 Group Size 区分的 Group Reduce + +对于 `N` 个 f32 lane 和 `G = num_groups`，group size 是 `S = N / G`。 + +```text +S=8: + input layout 可以是 contiguous。 + group_reduce result 通常使用 layout。 + +S=16: + 如果 input 来自 f16->f32 vcvt，layout 可以是 deinterleaved=2。 + 如果 input 从 dense 拆出，layout 可以是 deinterleaved=2, block_elems=8。 + result 通常使用 layout。 + +S=32: + input layout 使用 deinterleaved=4, block_elems=8。 + VPTO 形状是四个部分 group reduction 后接 add tree。 + result 通常使用 layout。 + +S=64: + row-local path 在可行时让每个 group 使用一条 physical row。 + result 可以使用 layout，避免 unsupported packing。 +``` + +S=32 例子： + +```text +assignment: + source/mask = deinterleaved=4, block_elems=8 + result = group_slots(num_groups=8, slots=8) + +VPTO: + vdintlv / pdintlv_b32 + vcgadd x4 + 使用 PAT_VL8 做 vadd tree + 通过一次 PAT_VL8 store 完成 group_store +``` + +这个场景说明为什么需要 `block_elems`。 + +### 4.5 Group Result 继续作为 Dense Rows 使用 + +Surface 意图： + +```mlir +%sum32 = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8, reassoc} +%rows32 = pto.vmi.group_broadcast %sum32 {num_groups = 8} +%rows16 = pto.vmi.truncf %rows32 +pto.vmi.store %rows16, %dst[%off] +``` + +支持的 assignment 形状： + +```mlir +%sum32 = pto.vmi.group_reduce_addf ... + -> !pto.vmi.vreg<8xf32, #pto.vmi.layout> + +%rows32 = pto.vmi.group_broadcast %sum32 {num_groups = 8} + -> !pto.vmi.vreg<128xf32, #pto.vmi.layout> + +%rows32_split = pto.vmi.ensure_layout %rows32 + : contiguous -> deinterleaved=2 + +%rows16 = pto.vmi.truncf %rows32_split + : deinterleaved=2 -> contiguous + +pto.vmi.store %rows16, %dst[%off] +``` + +VPTO 形状： + +```text +group_reduce: + vcgadd partials + vadd tree + +group_broadcast: + vselr 风格 selection，把 group slots 展开到 dense row lanes + +truncf: + vcvt even/odd + merge + +store: + vsts +``` + +这个场景说明为什么 group 结果 layout 必须挂在 value 上：reduce 之后， +cast 和 broadcast 必须知道 group 结果在哪里，而不能回看 producer。 + +### 4.6 通过 Mask 表达 Tail + +VMI 通过 mask 表达 tail，不通过 padding 表达 tail。 + +```mlir +%mask = pto.vmi.create_mask %active_lanes +%x = pto.vmi.masked_load %src[%off], %mask +%y = pto.vmi.mulf %x, %scale +pto.vmi.masked_store %y, %dst[%off], %mask +``` + +Grouped tail： + +```mlir +%gmask = pto.vmi.create_group_mask %active_elems_per_group + {num_groups = 8, group_size = 32} +%sum = pto.vmi.group_reduce_addf %x, %gmask {num_groups = 8, reassoc} +``` + +同一个 semantic mask 面对 f8/f16/f32 user 时，可能需要不同 concrete +granularity。Assignment 会通过 mask helper op 显式表达这些转换。 + +### 4.7 控制流和函数边界 + +Concrete layout 必须显式跨过 CFG 和内部 function boundary。 + +```mlir +%r = scf.if %cond + -> !pto.vmi.vreg<128xf32, #pto.vmi.layout> { + %a_dense = pto.vmi.ensure_layout %a : deinterleaved=2 -> contiguous + scf.yield %a_dense +} else { + %b_dense = pto.vmi.ensure_layout %b : deinterleaved=2 -> contiguous + scf.yield %b_dense +} +``` + +`vmi-to-vpto` 之后，region result 会变成多个物理 VPTO value： + +```text +scf.if -> (!pto.vreg<64xf32>, !pto.vreg<64xf32>) +``` + +这个场景说明为什么 layout 应该是 type 的一部分，而不是依赖 defining op。 + +### 4.8 完整 Histogram 语义 + +VPTO 的 histogram 指令一次读取 `256xui8` source，但结果只能写 +`128xui16` accumulator。完整 `ui8` histogram 有 256 个 bin，因此物理 VPTO +接口需要通过 `#bin = 0/1` 分两次统计低半区和高半区。 + +VMI surface 不暴露这个物理 split： + +```mlir +%hist = pto.vmi.vdhist %acc, %src, %mask + : !pto.vmi.vreg<256xui16>, + !pto.vmi.vreg, + !pto.vmi.mask + -> !pto.vmi.vreg<256xui16> +``` + +语义是完整 256-bin distribution histogram： + +```text +for b = 0..255: + hist[b] = acc[b] + count(i where mask[i] && src[i] == b) +``` + +Assignment 形状： + +```text +src/mask = contiguous, b8 mask granularity +acc/result = contiguous 256xui16 logical value +``` + +VPTO 形状： + +```text +acc/result part0 = bins 0..127 +acc/result part1 = bins 128..255 + +for each 256-lane source chunk: + part0 = dhistv2(part0, src_chunk, mask_chunk, #bin=0) + part1 = dhistv2(part1, src_chunk, mask_chunk, #bin=1) +``` + +这说明 VMI 的易用性不只来自 layout assignment。对于这种 value-indexed +accumulation，VMI 还应该隐藏 VPTO 为了物理 vreg 宽度暴露出来的 range +selector、lo/hi accumulator 和多条物理指令。 + +`pto.vmi.vchist` 可以使用相同 surface 形状，但当前必须先验证 VPTO `CHISTv2` +在 high range 上返回的是全局累计还是 range-local 累计。这个差异会影响是否需要 +额外给 high half 加上 low half 的总计数，因此不能只按 op 名字猜 lowering。 + +### 4.9 Block-Strided UB Staging + +有些 CCE kernel 并不是在 register 内做任意 byte shuffle，而是先把结果写到 +UB scratch，再用 block-strided vector load/store materialize 目标 UB layout。 +`quant_minimum` 的 MXFP8 NZ case 是典型例子： + +```text +compute: + row-major ND FP8 scratch + +row-wise staging: + for row in 0..31: + q8_row = vmi.stride_load(nd + row * 64, + block_stride=1, repeat_stride=1) + vmi.stride_store(q8_row, nz + row * 32, + block_stride=33, repeat_stride=1) + +copy-out: + 2D MTE copies two 1024B NZ planes from UB to GM +``` + +这里 `q8_row` 的 VMI value 仍然是 contiguous `64xf8` 逻辑向量： + +```mlir +%q8_row = pto.vmi.stride_load %nd[%nd_off], %c1_i16, %c1_i16, %mask + : !pto.ptr, i16, i16, !pto.vmi.mask<64xpred> + -> !pto.vmi.vreg<64xf8E4M3FN> + +pto.vmi.stride_store %q8_row, %nz[%nz_off], %c33_i16, %c1_i16, %mask + : !pto.vmi.vreg<64xf8E4M3FN>, !pto.ptr, i16, i16, + !pto.vmi.mask<64xpred> +``` + +Assignment 形状： + +```text +stride_load result = contiguous +stride_load mask = contiguous, granularity follows result element width +stride_store value = contiguous +stride_store mask = contiguous, granularity follows value element width +``` + +VPTO 形状： + +```text +base_in = pto.addptr nd, nd_off +q8_row = pto.vsldb base_in, block_stride=1, repeat_stride=1, mask + +base_out = pto.addptr nz, nz_off +updated = pto.vsstb q8_row, base_out, block_stride=33, repeat_stride=1, mask + -> updated_base +``` + +这个场景说明：memory layout transformation 不一定要变成 VMI data layout。 +只要 VMI op 的语义是“从哪些地址读/写哪些 logical lane”，register value +仍然可以保持 contiguous，`vmi-to-vpto` 也仍然是 local lowering。 + +## 5. 当前边界 + +当前设计方向： + +```text +surface VMI: + 描述不带 layout 的逻辑向量语义。 + +layout assignment: + 选择 layout、mask granularity 和显式 materialization helper。 + +optimization: + 只在结果 IR 仍然可以 local lowering 时改写显式 helper。 + +vmi-to-vpto: + 严格 lower 它看到的 assigned/optimized IR。 +``` + +暂不支持或有意收紧的范围： + +```text +group_slots value 的普通 dense store: + 非法，除非先经过 group_broadcast 或其他显式 group-to-dense op。 + +packed group_slots f32->f16 cast: + 非法，除非 assignment 能把它 commute 到 group_broadcast 之后，或者使用 + 支持的 row-local slots=1 path。 + +FP4 packed input/output: + packed FP4 不属于当前 VMI surface。PTO/VPTO 已有 !pto.f4E1M2x2 + 和 !pto.f4E2M1x2 packed 物理类型，且这些类型的 shape 语义是 + packed pair/byte 数，不是 logical FP4 lane 数。在 VMI 中直接写 + vreg 会让 N 表示物理 packed byte 还是逻辑 FP4 元素 + 产生歧义，因此 verifier 会直接拒绝 + vmi.vreg<...x!pto.f4E1M2x2/!pto.f4E2M1x2>。 + + 当前 VMI surface 不包含专用 FP4 packed-memory op。FP4 packed IO + 需要先作为独立语义重新设计，不能进入当前 dialect surface。 + +extract: + 暂不作为支持的 VMI surface。 + +padding transfer_read: + 当前 tail 设计不需要；tail 使用 mask。 + +scan / contract / compress / active_prefix_index: + dialect surface 中可以存在，但除非补充具体 case，否则不属于第一阶段聚焦的 + layout/lowering 实现集合。 + +gather / scatter: + 当前只覆盖 UB pointer、contiguous layout 和已明确支持的 element/index 宽度。 + `ui16` gather 可承接 E8M0 byte-pair reorder；它不是通用 byte shuffle。 +``` + +设计目标是优先保证语义完整：只要 VMI 接受某个 case，所需的 layout 沟通就必须 +在 IR 中显式表达，并且能被 `vmi-to-vpto` local lowering。 diff --git a/docs/designs/vmi-lane-stride-generalization-design.md b/docs/designs/vmi-lane-stride-generalization-design.md new file mode 100644 index 0000000000..e5134bede3 --- /dev/null +++ b/docs/designs/vmi-lane-stride-generalization-design.md @@ -0,0 +1,903 @@ +# VMI Lane-Stride Layout Generalization Design + +本文定义 `lane_stride` 从 group-slot 专用属性泛化为 VMI layout 的通用 +物理 lane 映射轴。目标不是只优化 `64xf16 -> 64xf32`，而是给 dense +value、group-slot value、类型转换、broadcast materialization 和 load/store +rematerialization 提供统一表达。 + +## 1. Problem + +当前文档对 `lane_stride` 的语义是： + +```text +logical lane-sized physical slot 之间有固定间距 +``` + +但实现只允许它出现在： + +```text +#pto.vmi.layout +``` + +并且现有 helper 会把 `ui8 lane_stride=4` 这类 group-slot lowering 映射为 +b32 carrier。这导致两个问题： + +1. dense value 无法表达“64 个 f16 logical lanes 放在一个 128xf16 物理向量 + 的偶数 lane 上”。 +2. `lane_stride` 的 layout 语义和 group-slot carrier lowering 被混在一起。 + +泛化后必须保持以下边界： + +```text +lane_stride: + layout lane map, does not change logical element type + +carrier packing: + one lowering strategy for selected group-slot integer stores +``` + +## 2. Semantic Model + +### 2.1 Dense Layout + +Dense layout 仍然表示每个 logical lane 都有语义值。第一阶段只增加 +`lane_stride` 一个新轴： + +```text +deinterleave factor F +block elems B +lane stride LS +``` + +建议 surface spelling： + +```text +#pto.vmi.layout +#pto.vmi.layout + +#pto.vmi.layout +#pto.vmi.layout +``` + +Defaults: + +```text +F = 1 for contiguous +B = 1 +LS = 1 +``` + +Dense lane map: + +```text +logical lane i + +block q = i / B +in-block lane r = i % B +part p = q % F +part block t = q / F + +dense lane index in part = t * B + r +physical part p, physical lane dense lane index * LS +``` + +The current stage intentionally describes only phase-zero strided dense layouts. +For `lane_stride = 2`, that means semantic lanes occupy even physical lanes. + +An optional future `lane_offset` or `lane_phase` field is useful only after the +IR has a concrete zero-copy view or producer whose logical lane `i` is +intentionally represented at physical lane `2 * i + 1` or another non-zero +phase. The current stage has no such producer. The field should +not be added just because the target has a `vcvt ODD` instruction. + +`vcvt ODD` is needed in two different situations: + +```text +1. Full conversion of a packed contiguous source. + Example: contiguous f16 -> deinterleaved=2 f32 uses EVEN and ODD. + This is not an odd-phase dense source layout; it is the normal multi-part + lowering of a packed source. + +2. Single-part conversion of a future zero-copy odd-lane view. + Example: if a logical deinterleave/extract result were represented as + f16 lane_stride=2, lane_offset=1 instead of being compacted, then converting + that view to f32 contiguous would use ODD. This requires an explicit VMI + producer or consumer contract; current-stage dense stride does not + create such values. +``` + +The current design implements case 1 with existing conversion lowering and case +2 only as a non-goal extension. The useful dense-stride optimization in this +stage uses phase-zero layout and therefore selects `EVEN` for `W=2`. + +### 2.2 Deinterleaved vs Lane Stride + +Use `deinterleaved` when multiple semantic residue classes or physical parts of +the same dense logical value are all present. + +Use dense `lane_stride` when one semantic stream is stored sparsely inside each +physical part and the skipped lanes have no semantic value for this VMI value. + +Decision rule: + +```text +all residue classes are semantic: + use deinterleaved + +only one phase-zero residue class is semantic: + use lane_stride + +multiple parts are semantic and each part is internally strided: + use deinterleaved + lane_stride +``` + +Examples: + +```text +contiguous f16 -> f32 full dense widen: + source lanes 0,1,2,3,... are all semantic + result naturally has even/odd conversion parts + use result deinterleaved=2 + +64xf16 -> 64xf32 where the f32 consumer wants contiguous: + the vcvt layout support may request source lane_stride=2 + if the source producer/rematerialization can satisfy that request, source + lanes 0,2,4,... become semantic and lanes 1,3,5,... are holes for this value + extf result can then be contiguous through one EVEN conversion + +group-reduce or dense consumer that needs two/four logical fragments: + the fragments are semantic parts of the same dense value + use deinterleaved=2/4, not lane_stride +``` + +Do not use `lane_stride` to describe a full packed value that happens to need an +ODD conversion part. Do not use `deinterleaved` to describe holes inside one +physical part. + +Important distinction: + +```text +one hardware vcvt output: + always one contiguous VPTO output register + +VMI ext result layout: + describes how one or more hardware output registers map back to logical lane + order +``` + +For `W=2`, with logical f16 lanes named by their logical indices: + +```text +source contiguous: + physical lanes: 0, 1, 2, 3, 4, 5, ... + vcvt EVEN output carries logical lanes 0, 2, 4, ... + vcvt ODD output carries logical lanes 1, 3, 5, ... + VMI result layout is deinterleaved=2 unless another materialization + interleaves the two outputs. + +source lane_stride=2: + physical lanes: 0, _, 1, _, 2, _, ... + vcvt EVEN output carries logical lanes 0, 1, 2, ... + VMI result layout is contiguous. +``` + +So "vcvt output is contiguous" does not by itself mean the VMI `extf` result is +contiguous. The result layout depends on the logical lane mapping of the source +layout and the selected conversion parts. + +### 2.3 Group-Slot Layout + +Group-slot layout remains non-dense. Only `G` group result slots have semantic +values: + +```text +#pto.vmi.layout +#pto.vmi.layout +``` + +Existing mapping is preserved: + +```text +slot_block(g) = g / K +slot_lane(g) = (g % K) * LS +``` + +This remains a group-slot placement property. It does not make non-slot lanes +semantic. Existing `ui8 lane_stride=4` to b32 carrier lowering is still legal, +but it is not the definition of `lane_stride`. + +Group-slot `lane_offset` is not needed in the current stage. It should remain +out of scope unless a real group-slot producer needs non-zero phase. + +## 3. Physical Capacity + +`lane_stride` increases the number of physical lane slots needed by a dense +part, but it does not change the VMI logical element type. + +For one dense physical part in the current stage: + +```text +logical lanes in this part = M +required physical lanes = (M - 1) * LS + 1 +``` + +The number of VPTO physical registers for each part is: + +```text +ceil(required physical lanes / lanes_per_vpto_register(T)) +``` + +Total physical arity: + +```text +deinterleave factor F * registers per part +``` + +Example: + +```text +!vmi.vreg<64xf16, contiguous, lane_stride=2> + +lanes_per_vpto_register(f16) = 128 +required physical lanes = 63 * 2 + 1 = 127 +physical arity = 1 +``` + +The 64 logical f16 lanes occupy physical f16 lanes `0, 2, 4, ... 126` of one +`!pto.vreg<128xf16>`. The other lanes are undefined unless another layout +value gives them semantics. + +Some lowerings represent the same lane map with wider carrier slots instead of +logical-element lanes. For example, a b16 value with `lane_stride=2` may be +lowered as the low b16 element of each b32 carrier slot when using +`UNPK_B16`/`PK_B32` or register pack/unpack materialization. This does not +change the VMI logical element type; it is a VPTO lowering representation choice. + +## 4. Type And Operation Generalization + +The design is element-type agnostic. Dense `lane_stride` applies to any VMI +element type whose physical VPTO lane count is known: + +```text +f8, f16, bf16, f32 +i8, ui8, i16, ui16, i32, ui32 +pred masks at an explicit predicate granularity +``` + +An op may support a strided dense layout only when its VPTO lowering can +preserve the lane map. Unsupported combinations are rejected by layout support +queries, not silently repaired in `vmi-to-vpto`. + +### 4.1 VPTO Pack/Unpack Support Boundary + +Dense `lane_stride` is not a generic VPTO load/store operand. It is supported +only when the lane map matches a concrete VPTO distribution or register +materializer. + +Direct compact memory support: + +| Dense lane_stride | compact load | compact store | +|---:|---|---| +| 2, b8 | `vlds UNPK_B8` | `vsts PK_B16` | +| 2, b16 | `vlds UNPK_B16` | `vsts PK_B32` | +| 2, b32 | `vlds UNPK_B32` | `vsts PK_B64` | +| 4, b8 | `vlds UNPK4` | `vsts PK4_B32` | +| 4, b16/b32 | no direct dist | no direct dist | + +Direct scalar broadcast load target capability: + +```text +lane_stride=2/4, b8/b16/b32: + vlds BRC_B8/B16/B32 +``` + +The current stage does not add a VMI scalar broadcast-load op. BRC is therefore +a target capability for a separate scalar broadcast-load semantic, not part of +the current `vmi.load -> ensure_layout` compact-stream fold. + +Register fallback between contiguous and dense `lane_stride` should use the +register-side counterpart of these distributions: + +```text +contiguous -> lane_stride: + vsunpack/vzunpack-style placement into wider slots + +lane_stride -> contiguous: + vpack-style extraction from wider slots +``` + +`vintlv`/`vdintlv` remain the materializers for two-stream +interleave/deinterleave layouts; they are not the primary fallback for dense +`lane_stride`. + +### 4.2 Layout-Transparent Dense Ops + +Layout-transparent dense ops include ordinary elementwise arithmetic and +select-like ops when every dense data operand/result has the same layout: + +```text +add/mul/fma/min/max/select: + operands and result require identical dense layout key + key includes F, B, and LS +``` + +No physical shuffle is implied by these ops. + +### 4.3 Widening Conversion + +Let a widening conversion increase element storage width by ratio `W`: + +```text +f16 -> f32: W = 2 +bf16 -> f32: W = 2 +i16 -> i32: W = 2 +ui16 -> ui32: W = 2 +f8 -> f32: W = 4 +i8 -> i32: W = 4 +ui8 -> ui32: W = 4 +ui8 -> ui16: W = 2 +``` + +For a phase-zero source dense layout with `lane_stride = LS`, a single hardware +conversion part is sufficient when: + +```text +LS % W == 0 +``` + +The selected hardware part in the current stage is: + +```text +part = 0 +``` + +The result layout after conversion is: + +```text +result lane_stride = LS / W +``` + +For a future phase-aware layout with `lane_offset = O`, the generic relation is: + +```text +part = O % W +result lane_stride = LS / W +result lane_offset = (O - part) / W +``` + +That future relation should be enabled only when a real odd/non-zero-phase VMI +producer or consumer exists. + +Examples: + +```text +f16 source: contiguous, lane_stride=2 +extf to f32: + use vcvt EVEN + result contiguous + +f16 source: contiguous, lane_stride=4 +extf to f32: + use vcvt EVEN + result contiguous, lane_stride=2 +``` + +If `LS < W` or `LS % W != 0`, the conversion may need multiple hardware parts +and may naturally produce a deinterleaved result. The current contiguous source +case is the common example: + +```text +f16 source: contiguous, lane_stride=1 +extf to f32: + use vcvt EVEN and vcvt ODD + result deinterleaved=2 +``` + +Assignment chooses one preferred fact for the op before lowering. Consumer +requests are handled by the existing use-site materialization path after the +op's assigned result layout is fixed. + +The preferred direction for this optimization is not "notice the input is +already strided". The conversion op can be the layout-entry point and compute a +single preferred layout fact for the current op instance. The choice must be +arity-driven, not special-cased by a spelling such as `64xf32`. + +For source/result logical lane count `N`, let: + +```text +natural result layout: + source dense factor F, lane_stride 1 + result dense factor F * W, lane_stride 1 + +compact result layout: + result keeps source dense factor F and uses lane_stride 1 + source uses lane_stride W inside the same dense factor F +``` + +The self-preferred widening rule is: + +```text +if physical_arity(compact result) < physical_arity(natural result) + and target supports the required source lane_stride relation: + choose compact result and request source lane_stride=W +else: + choose natural result deinterleaved by W +``` + +For ordinary contiguous `f16 -> f32` this gives: + +```text +64xf32: + compact arity = 1 + natural deinterleaved=2 arity = 2 + choose source lane_stride=2, result contiguous + +128xf32: + compact arity = 2 + natural deinterleaved=2 arity = 2 + choose natural result deinterleaved=2 + +256xf32: + compact arity = 4 + natural deinterleaved=2 arity = 4 + choose natural result deinterleaved=2 +``` + +If the source is already deinterleaved by `F`, the natural result factor is +`F * W`. For example, `deinterleaved=2 f16 -> f32` naturally produces +`deinterleaved=4 f32`. + +The same arity rule applies to other widening ratios and types. For example, +`ui8 -> ui32` has `W=4`; a lane-stride source is preferred only when the +contiguous result has fewer physical chunks than the natural +`deinterleaved=4` result and the target supports the `lane_stride=4` relation. + +The two layout facts are therefore: + +```text +baseline fact: + source contiguous + result deinterleaved=W + arity: physical_arity(result deinterleaved=W) + +lane-stride fact: + source lane_stride=W + result contiguous + arity: physical_arity(result contiguous) + source layout request: explicit +``` + +In the current single-preference framework, `ext` should publish one preferred +fact. The lane-stride fact is an op-local preference: assignment records the +required source/result relation in the IR and inserts `ensure_layout` at the +source use if the producer is not already in that layout. Later +rematerialization or fold passes may remove that helper when a concrete producer +rewrite exists; otherwise the helper is either lowered by a registered +contiguous/lane-stride materializer or rejected before `vmi-to-vpto`. + +This keeps the optimization in layout assignment/rematerialization, not in a +late `vmi-to-vpto` peephole, and stays within the existing single-preference +assignment model. + +### 4.4 Narrowing Conversion + +Narrowing uses the same arity-driven idea in the opposite direction. If source +element width is `W` times the result element width, a single hardware narrowing +part can produce a phase-zero strided result when: + +```text +result lane_stride = source lane_stride * W +part = 0 +``` + +This covers more than f32-to-f16. The same relation applies to: + +```text +f32 -> f16/bf16 +i32 -> i16/i8 +ui32 -> ui16/ui8 +ui16 -> ui8 +``` + +The natural narrowing relation is the inverse of natural widening: + +```text +source dense factor F * W, lane_stride 1 +result dense factor F, lane_stride 1 +``` + +The compact-store-oriented relation is: + +```text +source keeps dense factor F and lane_stride 1 +result keeps dense factor F and uses lane_stride W +``` + +Narrowing has the same candidate family as widening. The arity comparison is +made on the source side, because the compact relation keeps the source +contiguous while the natural relation may require a deinterleaved source. + +The self-preferred narrowing rule is: + +```text +if physical_arity(compact contiguous source) + < physical_arity(natural deinterleaved source) + and physical_arity(compact source) == physical_arity(strided result) + and target supports the source-contiguous/result-lane_stride relation: + choose source contiguous, result lane_stride=W +else: + choose natural deinterleaved-source to contiguous-result relation +``` + +Use-site requests may still select the strided relation when a later consumer +can directly consume it: + +```text +if a consumer requests result lane_stride=W + and target supports source-contiguous/result-lane_stride narrowing: + request source contiguous + set or rematerialize result lane_stride=W +``` + +For ordinary `f32 -> f16`: + +```text +64xf32 -> 64xf16: + natural source deinterleaved=2 arity = 2 + compact source contiguous arity = 1 + choose source contiguous, result lane_stride=2 + +128xf32 -> 128xf16: + natural source deinterleaved=2 arity = 2 + compact source contiguous arity = 2 + choose natural source deinterleaved=2, result contiguous +``` + +So trunc should not blindly create a lane-stride result for every narrowing. +It should apply the same arity/support checks as ext. A consumer may still +request a strided result when that layout is useful, such as an unmasked compact +store lowered with `PK`/`PK4`. For masked stores, the value and mask must share +the same lane map before a direct packed masked store is legal. + +The exact supported parts are target-op dependent. The layout assignment layer +should ask the op support interface whether a given source/result layout pair is +legal, rather than encoding type-specific shortcuts. + +### 4.5 Broadcast Materialization + +Broadcast remains a logical operation. `lane_stride` only describes the chosen +materialized layout. + +Scalar or group broadcast can materialize to a dense layout only when the +broadcast lowering or rematerialization support query accepts that lane map: + +```text +logical broadcast: + lane i gets value group(i) + +materialized layout: + lane i is stored at physical lane map(i) +``` + +This keeps E2B-style optimizations in the layout/rematerialization layer. A +group broadcast load may choose a dense strided layout when that layout directly +matches a consumer or a target instruction. If another consumer needs a +different layout, rematerialization may clone the broadcast or insert +`ensure_layout`. + +`group_broadcast_load` is also a VMI semantic, not an E2B semantic. It means: + +```text +for each logical group g: + load one scalar from source[offset + g * source_group_stride] + broadcast that scalar to all lanes in group g +``` + +E2B is a target lowering choice for the subset where that logical memory pattern, +the group size, the element width, and the assigned result layout match the E2B +packet semantics. Other lowering strategies may implement the same VMI +operation, so support queries should report "E2B is applicable" instead of +rewriting the VMI meaning to "this op is E2B". + +### 4.6 Masked Lane-Stride Stores + +Masks are logical predicates. A `masked_store` mask bit denotes whether a +logical element participates in the store; it is not automatically a predicate +for the physical lane slot that happens to carry that element after layout +assignment. + +For dense `lane_stride`, this distinction matters. With `lane_stride=2`, +logical lane `i` is carried in physical lane `2*i`. A packed store then +compacts those even physical lanes into a contiguous memory stream. A user mask +that is still contiguous cannot be passed directly to that packed store, because +the packed-store predicate is interpreted after the value lanes have been +compacted. + +A direct masked compact store is therefore legal only when the compiler has +assigned the value and mask the same lane map. That may happen because the mask +producer can directly produce the requested lane map, because assignment inserts +a mask `ensure_layout`, or because rematerialization rebuilds the mask producer +for that lane map. Without that compiler-derived proof, assignment should keep +a layout that the existing masked-store path can lower, even if the corresponding +unmasked store could use a dense lane-stride `PK` instruction. + +## 5. Assignment And Optimization Boundary + +The assignment pipeline should keep the existing responsibility split: + +```text +layout assignment: + collect consumer requests + ask producer/op support + assign explicit layout attrs + insert ensure_layout for use-local conflicts + +rematerialization: + clone cheap producers for incompatible use-site layouts + replace ensure_layout(producer) when producer can directly create target layout + +layout fold: + erase or fuse materialization helpers when the producer already has the + requested lane map + +vmi-to-vpto: + lower explicit assigned layouts only + no hidden layout selection policy +``` + +Dense `lane_stride` is therefore an assigned layout fact, not a lowering-side +pattern. An entry op such as `extf` may prefer it from the conversion ratio +alone; producer-specific rewrites are handled later by fold/rematerialization +passes over explicit helpers. The selected layout is fixed before +`vmi-to-vpto`, and `vmi-to-vpto` does not rediscover the preference. + +## 6. End-To-End Case Walkthroughs + +These cases are the intended test for the design. They show when dense +`lane_stride` is useful and when it should lose to the existing deinterleaved +plan. + +The logical programs in this section are pre-assignment VMI and do not carry +concrete layouts. Layouts shown under "baseline plan" or "lane-stride plan" are +possible assignment results, not layouts written in the input program. + +### 6.1 Contiguous Load, Ext, Contiguous Store + +Logical program: + +```text +%x16 = vmi.load %in : 64xf16 +%x32 = vmi.extf %x16 : 64xf16 -> 64xf32 +vmi.store %x32, %out : dense contiguous memory effect +``` + +Baseline plan: + +```text +load result: + contiguous f16 + +ext relation: + source contiguous f16 + result deinterleaved=2 f32 + lower: vcvt EVEN + vcvt ODD + +store: + needs contiguous f32 + requires result materialization deinterleaved=2 -> contiguous +``` + +Lane-stride plan: + +```text +load result: + lane_stride=2 f16 + +ext relation: + source lane_stride=2 f16 + result contiguous f32 + lower: vcvt EVEN + +store: + consumes contiguous f32 directly +``` + +Assignment chooses the lane-stride plan for this shape because the contiguous +`64xf32` result uses one physical chunk while the natural deinterleaved result +uses two physical chunks. This decision is made by the cast arity rule, not by +a pattern that names `64xf32` directly. + +The load side then has two concrete outcomes: + +```text +accepted direct load fold: + the original load has only the lane-stride use + compact load semantics match a supported UNPK dist + vmi-layout-fold changes the VMI load result layout in place + +no direct load fold: + keep the explicit source ensure_layout + lower it through register pack/unpack if that materialization is supported + otherwise validation rejects the unsupported assigned relation +``` + +This case proves that `extf` can be the layout-entry point, while `load` support +is still decided by the load/ensure fold or by the explicit materialization +helper. + +### 6.2 Broadcast, Ext, Contiguous Store + +Logical program: + +```text +%b16 = vmi.broadcast %s : 1xf16 -> 64xf16 +%b32 = vmi.extf %b16 : 64xf16 -> 64xf32 +vmi.store %b32, %out +``` + +Baseline plan: + +```text +broadcast materializes contiguous f16 +ext produces deinterleaved=2 f32 through EVEN + ODD +store materializes deinterleaved=2 -> contiguous +``` + +Lane-stride plan: + +```text +broadcast rematerializes directly as lane_stride=2 f16 +ext produces contiguous f32 through one EVEN +store consumes contiguous f32 +``` + +Here the lane-stride plan is accepted because broadcast is a rematerializable +producer: it can be rebuilt with the requested physical lane map instead of +requiring a register layout conversion. This is the kind of producer where +`vcvt` should drive a source `lane_stride=2` request. + +### 6.3 Ext Feeding A Deinterleaved Consumer + +Logical program: + +```text +%x16 = producer : 128xf16 +%x32 = vmi.extf %x16 : 128xf16 -> 128xf32 +%r = vmi.group_reduce %x32 // requests deinterleaved=2 +``` + +Baseline plan: + +```text +source contiguous f16 +result deinterleaved=2 f32 +consumer consumes result directly +``` + +Lane-stride plan: + +```text +source lane_stride=2 f16 +result contiguous f32 +consumer then needs contiguous -> deinterleaved=2 materialization +``` + +The baseline plan should win. A lane-stride fact is not useful when it creates a +layout the consumer does not want. The cast arity rule also does not prefer +lane_stride here: `128xf32` contiguous and `128xf32 deinterleaved=2` both use +two physical chunks. + +### 6.4 One Ext Result Feeding Store And Reduce + +Logical program: + +```text +%x16 = cheap_or_expensive_producer : 128xf16 +%x32 = vmi.extf %x16 : 128xf16 -> 128xf32 +vmi.store %x32, %out // requests contiguous +vmi.group_reduce %x32 // requests deinterleaved=2 +``` + +If `%x16` is not cheap to rematerialize: + +```text +assign ext result deinterleaved=2 for the reduce +insert ensure_layout at the store use +``` + +If `%x16` and `extf` are cheap to rematerialize: + +```text +shared path: + source contiguous -> ext result deinterleaved=2 -> reduce + +store-only remat path: + rematerialized source lane_stride=2 -> ext result contiguous -> store +``` + +This is a rematerialization decision, not a local `vcvt` peephole. + +### 6.5 Group Broadcast Load Feeding Ext + +Logical program: + +```text +%g16 = vmi.group_broadcast_load %scale : logical dense 64xf16 +%g32 = vmi.extf %g16 : 64xf16 -> 64xf32 +consumer requests contiguous %g32 +``` + +The lane-stride plan is accepted only if the group broadcast load lowering can +emit the requested lane map directly: + +```text +group broadcast load result lane_stride=2 f16 +ext result contiguous f32 +``` + +If the broadcast load can only produce contiguous or deinterleaved packets for +the target element width, assignment should keep those layouts and let later +materialization/rematerialization handle the consumer conflict. Dense +`lane_stride` is a requestable layout, not a guarantee that every producer can +create it. + +## 7. Compatibility Rules + +Two dense layouts are identical only if all lane-map fields match: + +```text +F, B, LS +``` + +Two dense layouts may be related by an explicit materialization only if a +registered relation can lower the map conversion. Examples: + +```text +contiguous <-> deinterleaved=2 +deinterleaved=2 <-> deinterleaved=4 when supported by existing intlv/dintlv +contiguous <-> contiguous, lane_stride=2 when pack/unpack materialization or +producer rematerialization supports it +``` + +The baseline assignment must not assume an arbitrary dense-to-dense +`ensure_layout` is free or legal. Unsupported materializations should fail in +verification or remain unselected by support queries. + +## 8. Non-Goals + +This design does not: + +1. Turn memory layout into strided memory semantics. Dense VMI `lane_stride` + describes register materialization, not GM/UB address stride. +2. Make non-slot lanes of group-slot layouts semantic. +3. Require every VPTO op to support every strided layout. +4. Encode `64xf16 -> 64xf32` as a one-off `vcvt EVEN` peephole. + +## 9. First Useful Optimization + +The motivating case becomes one instance of the generic rule: + +```text +source: + requested as !vmi.vreg<64xf16, contiguous, lane_stride=2> + +op: + extf f16 -> f32, W=2 + +result: + !vmi.vreg<64xf32, contiguous> + +lowering: + one vcvt EVEN +``` + +The same mechanism also covers: + +```text +bf16 -> f32 with phase-zero lane_stride=2 +ui8 -> ui16 with lane_stride=2 +ui8 -> ui32 with lane_stride=4 +f8 -> f32 with lane_stride=4 +narrowing conversions that intentionally produce phase-zero strided results +broadcast materialization into a consumer-required strided dense layout +``` diff --git a/docs/designs/vmi-lane-stride-generalization-implementation.md b/docs/designs/vmi-lane-stride-generalization-implementation.md new file mode 100644 index 0000000000..96521e4a3f --- /dev/null +++ b/docs/designs/vmi-lane-stride-generalization-implementation.md @@ -0,0 +1,1690 @@ +# VMI Lane-Stride Layout Generalization Implementation Plan + +本文给出 `lane_stride` 泛化的实现路径。设计目标是把 lane-strided dense +layout 作为一等 layout fact 固化、传播、rematerialize 和 lower，而不是在 +`vmi-to-vpto` 中识别单个 `64xf16 -> 64xf32` pattern。 + +## 1. Implementation Principles + +1. `lane_stride` is a lane-map field. +2. Dense `lane_stride` does not change the VMI logical element type. +3. Group-slot carrier packing is a separate lowering helper. +4. Layout assignment decides layout before `vmi-to-vpto`. +5. `vmi-to-vpto` only lowers explicit assigned layout attrs. + +Pre-existing baseline before this design: + +```text +dense contiguous/deinterleaved layouts: + did not carry lane_stride + +regular VMI load/store: + did not support dense lane_stride + support contiguous and selected deinterleaved lowering/materialization paths + +VPTO load/store: + pto.vlds/pto.vsts have a dist string and the VPTO surface supports several + distribution families, but there is no generic lane_stride operand + +group-slot lane_stride: + already existed and was used by selected group-store packed-byte lowering +``` + +Any dense lane-stride load/store support must enter explicitly by mapping a VMI +lane-stride layout to a specific supported VPTO dist family or materialization +sequence. It must not be inferred in `vmi-to-vpto` from a one-off producer or +consumer pattern. + +Current stage status: + +| Area | Status | Notes | +|---|---|---| +| Dense layout attrs | Supported | Dense contiguous/deinterleaved layouts carry `lane_stride`; group-slot carrier layout remains separate. | +| Direct compact load/store | Supported for selected phase-zero maps | LS=2 b8/b16/b32 through `UNPK_B8/B16/B32` and `PK_B16/B32/B64`; LS=4 b8 through `UNPK4` and `PK4_B32`. | +| Load/store layout folds | Supported with one-load/one-store preservation | `load -> ensure_layout(lane_stride)` rewrites the original load layout when all uses agree; `ensure_layout(lane_stride -> contiguous) -> store` lets the VMI store consume the lane-stride value. | +| Dense widening ext | Supported | `getPreferredCastLayoutFact` chooses the arity-reducing source `lane_stride=W` / result contiguous relation when it beats the natural deinterleaved result; otherwise it keeps the natural relation. | +| Dense narrowing trunc | Supported for dense natural paths | `getPreferredCastLayoutFact` uses the same arity rule in the inverse direction, so trunc keeps the natural deinterleaved-source / contiguous-result relation unless a compact relation actually reduces arity. | +| Masked compact store | Partially supported | Legal only when value and mask have the same lane map and the mask can be compacted for the selected store dist. | +| Masked trunc tail | Not optimized yet | Keep the existing legal path until mask lane-stride assignment/materialization is available. | +| Register fallback | Partially supported | Only same-physical-arity contiguous `<->` lane_stride paths with legal pack/unpack carriers. Arity-changing fallback is not in scope for this stage. | +| Group broadcast load | Supported only through specific strategies | `group_broadcast_load` remains a VMI semantic; E2B is one strategy with exact shape/layout constraints. | + +Remaining design/implementation work from this discussion is intentionally +limited to these areas: + +| Area | Work to settle | Required proof before enabling | +|---|---|---| +| Cast assignment | Keep `getPreferredCastLayoutFact` as the single op-local preferred relation helper, but make it shape-aware: compute the natural relation, compute the compact lane-stride relation, and select compact only when physical arity improves. | `64xf16 -> 64xf32` chooses source `lane_stride=2` and result contiguous; `128/256xf16 -> f32` keep natural `deinterleaved=2`; dense trunc keeps the natural relation unless compact arity wins. | +| Masked store | Let `masked_store` request the same lane map for value and mask, or keep the existing legal path when the mask cannot be assigned/rematerialized into that lane map. | No path may lower a lane-stride value with a stale contiguous user mask; lowering must compact the assigned mask into the packed-store predicate. | +| Group broadcast load | Keep `group_broadcast_load` as a VMI logical operation and make E2B only one support/lowering strategy selected by shape, element width, stride, and assigned result layout. | A failed E2B match must mean "this lowering strategy is unavailable", not "the VMI op is invalid" unless no fallback strategy is registered. | + +Known support boundaries that are not part of this discussion's remaining-work +queue: + +```text +b32 contiguous <-> lane_stride register fallback through generic vpack/vunpack +generic scalar broadcast-load VMI semantic for BRC +dense lane-stride masked_load +arity-changing register fallback +LS=4 b16/b32 direct compact load/store +LS > 4 direct compact load/store +non-zero lane_offset / lane_phase +ordinary load cloning/rematerialization without safe-read proof +global cost search across conflicting consumer layouts +partial-chunk dense lane-stride direct memory beyond the current full-chunk gate +``` + +### 1.1 VPTO Dist Capability Boundary + +VPTO already exposes fixed distribution families that can implement specific +layout-producing or layout-consuming memory operations: + +```text +vlds: + NORM + BRC_B8/B16/B32 + US_B8/B16 + DS_B8/B16 + UNPK_B8/B16/B32 + BRC_BLK + E2B_B16/B32 + UNPK4 + SPLT4CHN + SPLT2CHN_B8/B16 + +vldsx2: + BDINTLV + DINTLV_B8/B16/B32 + +vsts: + NORM_B8/B16/B32 + 1PT_B8/B16/B32 + PK_B16/B32/B64 + PK4_B32 + MRG4CHN_B8 + MRG2CHN_B8/B16 + +vstsx2: + INTLV_B8/B16/B32 +``` + +These are not equivalent to an arbitrary dense `lane_stride` operand: + +```text +DINTLV/INTLV: + two-stream deinterleave/interleave memory operation + maps naturally to VMI deinterleaved layouts, not to one sparse semantic stream + +US/DS: + fixed 2x upsample/downsample load families for b8/b16 + can serve selected lane-map producers when the semantic mapping matches exactly + +UNPK/PK/PK4: + fixed slot-pack/slot-unpack memory families + directly express selected dense lane_stride layouts such as b16 LS=2 and + b8 LS=4, but not arbitrary LS=N + +BRC/E2B/BRC_BLK: + fixed broadcast or group-expansion load families + useful when logical broadcast plus assigned layout matches the family + +MRG/SPLT: + fixed channel merge/split families + useful only for matching channel layouts +``` + +So VPTO has enough surface area to support selected dense lane-stride memory +optimizations, but VMI must model them as explicit support cases: + +```text +VMI layout fact + op semantics + element width + -> exact VPTO dist family + or materialization/rematerialization sequence + or unsupported +``` + +Concrete support matrix for dense phase-zero `lane_stride`: + +| Dense lane_stride | Compact stream load -> dense LS | Single-scalar broadcast load -> dense LS | Dense LS -> compact stream store | +|---:|---|---|---| +| 2 | direct for b8/b16/b32 through `vlds UNPK_B8/B16/B32` | target dist exists as `vlds BRC_B8/B16/B32`; needs a separate single-scalar broadcast-load VMI semantic | direct for b8 through `vsts PK_B16`, b16 through `vsts PK_B32`, and b32 through `vsts PK_B64` | +| 4 | direct for b8 through `vlds UNPK4` | target dist exists as `vlds BRC_B8/B16/B32`; needs a separate single-scalar broadcast-load VMI semantic | direct for b8 through `vsts PK4_B32` | + +| VMI memory semantic | Element width | VPTO op/dist | VMI result layout | Direct dense `lane_stride` support | +|---|---:|---|---|---| +| load one scalar and every logical lane uses it | b8 | `vlds BRC_B8` | any dense phase-zero lane map | target dist exists; needs a separate VMI scalar broadcast-load semantic | +| load one scalar and every logical lane uses it | b16 | `vlds BRC_B16` | any dense phase-zero lane map | target dist exists; needs a separate VMI scalar broadcast-load semantic | +| load one scalar and every logical lane uses it | b32 | `vlds BRC_B32` | any dense phase-zero lane map | target dist exists; needs a separate VMI scalar broadcast-load semantic | +| load compact stream `x[i]` into semantic lane `2*i` | b8 | `vlds UNPK_B8` | `contiguous, lane_stride=2` | yes | +| load compact stream `x[i]` into semantic lane `2*i` | b16 | `vlds UNPK_B16` | `contiguous, lane_stride=2` | yes | +| load compact stream `x[i]` into semantic lane `2*i` | b32 | `vlds UNPK_B32` | `contiguous, lane_stride=2` | yes | +| load compact stream `x[i]` into semantic lane `4*i` | b8 | `vlds UNPK4` | `contiguous, lane_stride=4` | yes | +| load compact stream `x[i]` into semantic lane `4*i` | b16/b32 | none | `contiguous, lane_stride=4` | no direct VPTO dist | +| load compact stream `x[i]` into semantic lane `K*i`, `K > 4` | any | none | `contiguous, lane_stride=K` | no direct VPTO dist | +| load memory `x[2*i]` into logical lane `i` | b8 | `vlds DS_B8` | contiguous | no; this is memory downsample | +| load memory `x[2*i]` into logical lane `i` | b16 | `vlds DS_B16` | contiguous | no; this is memory downsample | +| load alternating memory stream into even/odd logical streams | b8 | `vldsx2 DINTLV_B8` | two compact streams or deinterleaved=2 | no; not one sparse stream | +| load alternating memory stream into even/odd logical streams | b16 | `vldsx2 DINTLV_B16` | two compact streams or deinterleaved=2 | no; not one sparse stream | +| load alternating memory stream into even/odd logical streams | b32 | `vldsx2 DINTLV_B32` | two compact streams or deinterleaved=2 | no; not one sparse stream | +| store semantic lane `2*i` as compact memory `x[i]` | b8 | `vsts PK_B16` | source `contiguous, lane_stride=2` | yes | +| store semantic lane `2*i` as compact memory `x[i]` | b16 | `vsts PK_B32` | source `contiguous, lane_stride=2` | yes | +| store semantic lane `2*i` as compact memory `x[i]` | b32 | `vsts PK_B64` | source `contiguous, lane_stride=2` | yes | +| store semantic lane `4*i` as compact memory `x[i]` | b8 | `vsts PK4_B32` | source `contiguous, lane_stride=4` | yes | +| store semantic lane `4*i` as compact memory `x[i]` | b16/b32 | none | source `contiguous, lane_stride=4` | no direct VPTO dist | +| store semantic lane `K*i` as compact memory `x[i]`, `K > 4` | any | none | source `contiguous, lane_stride=K` | no direct VPTO dist | +| store two compact streams as alternating memory | b8 | `vstsx2 INTLV_B8` | two compact streams or deinterleaved=2 | no; not one sparse stream | +| store two compact streams as alternating memory | b16 | `vstsx2 INTLV_B16` | two compact streams or deinterleaved=2 | no; not one sparse stream | +| store two compact streams as alternating memory | b32 | `vstsx2 INTLV_B32` | two compact streams or deinterleaved=2 | no; not one sparse stream | + +Masked compact stores have an extra legality rule. The `vmi.masked_store` +predicate is a logical-lane predicate, while a VPTO packed store consumes a +predicate in the compacted store coordinate after the sparse lanes have been +packed. Therefore a lane-stride value cannot be paired with an unrelated +contiguous mask and lowered directly to `PK`/`PK4`. + +The direct masked-store path is legal only when all of these hold: + +```text +value source layout == mask source layout +value/mask physical arity matches +mask granularity matches the logical value element width before compaction +target has a predicate compaction path for the packed-store dist +``` + +For example, an f16 value with `lane_stride=2` places logical lanes in even +physical lanes. If a user mask remains contiguous, mask bit `i` still denotes +logical lane `i`, not physical lane `2*i`. Passing that mask directly to +`vsts PK_B32` would gate the wrong compact positions for tail or sparse masks. +The current legal path requires the mask to carry the same lane map as the value +and then compacts it with predicate unpack operations before emitting the +packed store. Ordinary unmasked `vmi.store` is different: lowering creates the +compact prefix predicate itself, so there is no user mask to reinterpret. + +Until masked-store assignment can request and prove the same lane map for value +and mask, assignment must keep masked-tail narrowing on an existing legal path +instead of choosing a lane-stride trunc result solely because the store could +otherwise use `PK`. + +Current-stage implementation: + +```text +lib/PTO/Transforms/VMILayoutAssignment.cpp + +VMIMaskedStoreOp keeps the existing conservative request: + requestDataUse(value, contiguous) + requestMaskUse(mask, contiguous, elementGranularity) + +trunc assignment does not inspect masked_store users and does not preserve a +special masked-store guard. It records the source/result relation returned by +getPreferredCastLayoutFact. If that conflicts with a masked_store contiguous +request, normal assignment conflict handling inserts the required +ensure_layout. +``` + +Future lane-stride `masked_store` support must be added as an explicit +consumer-owned extension, not as a trunc special case. The future dataflow must +prove that value and mask share the same lane map before a packed masked store +is legal: + +```text +%n = vmi.trunc* %wide + : source contiguous -> result contiguous, lane_stride = W + +%m_ls = vmi.ensure_mask_layout %m + : mask contiguous -> mask contiguous, lane_stride = W + +vmi.masked_store %n, %dst[%off], %m_ls +``` + +That future extension would need the same local VMI proof before lowering can do +the mechanical predicate compaction: + +```text +vmi-layout-fold: + may fold ensure_layout(value) + ensure_mask_layout(mask) into masked_store + only through canFoldContiguousMaskedStoreMaterialization + +vmi-to-vpto: + sees valueLayout == maskLayout + calls createDenseLaneStrideStorePredicate + emits LOWER punpack on the mask + emits vsts PK_B16/PK_B32/PK4_B32 as selected by value element width/layout +``` + +Future negative tests should cover the fallback: + +```text +fallback: + mask cannot be assigned/materialized to the candidate lane_stride + CHECK masked_store keeps contiguous value/mask request + CHECK no PK/PK4 masked compact store is emitted with a stale contiguous mask +``` + +The remaining VPTO dist tokens are fixed non-lane-stride operations: + +```text +UNPK_B8/B16/B32: + compact load into one element per 16/32/64-bit slot, giving lane_stride=2 for + b8/b16/b32 dense values + +UNPK4: + compact load into one b8 element per 32-bit slot, giving lane_stride=4 for b8 + +PK_B16/B32/B64 and PK4_B32: + compact store from one active low element per 16/32/64-bit slot. PK_B32 is + exactly the direct compact store for a b16 value with lane_stride=2, and + PK4_B32 is exactly the direct compact store for a b8 value with lane_stride=4 + +MRG4CHN_B8 and MRG2CHN_B8/B16: + fixed channel merge stores, not generic lane_stride stores + +SPLT4CHN and SPLT2CHN_B8/B16: + fixed channel split loads, not generic lane_stride loads + +BRC_BLK and E2B_B16/B32: + usable only after their exact block/group expansion semantic is modeled as a + VMI broadcast producer; do not count them as generic dense lane_stride load +``` + +### 1.2 Contiguous/Lane-Stride Fallback Materialization + +Direct load/store support is preferred. When a value already lives in VPTO +registers and a consumer requires the other layout, `ensure_layout` provides the +fallback conversion between contiguous and dense phase-zero `lane_stride`. + +For `contiguous -> lane_stride`, use register unpack placement when the VPTO +surface supports the required carrier type: + +```text +LS=2: + use vzunpack/vsunpack-style widening placement + b8 contiguous -> b16 slots with low b8 semantic + b16 contiguous -> b32 slots with low b16 semantic + b32 contiguous -> b64 slots with low b32 semantic + +LS=4: + for b8, apply two LS=2 unpack placements: + b8 contiguous -> b16 slots -> b32 slots with low b8 semantic +``` + +For `lane_stride -> contiguous`, use register pack when the VPTO surface +supports the required carrier type: + +```text +LS=2: + use vpack-style narrowing placement + low b8 from each b16 slot -> b8 contiguous + low b16 from each b32 slot -> b16 contiguous + low b32 from each b64 slot -> b32 contiguous + +LS=4: + for b8, apply two LS=2 pack placements: + low b8 from each b32 slot -> b16 slots -> b8 contiguous +``` + +This is the register-side counterpart of `UNPK`/`PK` memory distributions. Do +not use `vintlv`/`vdintlv` as the primary fallback for dense `lane_stride`; those +belong to two-stream interleave/deinterleave layouts. + +Current checked-in VPTO coverage: + +```text +register pack: + vpack supports integer 32 -> u16 and integer 16 -> u8 + so b16 LS=2 -> contiguous and b8 LS=2/4 -> contiguous are directly covered + when the VMI source/result physical arity is the same + b32 LS=2 -> contiguous needs 64 -> 32 pack support or another materializer + +register unpack: + vsunpack/vzunpack support integer widening by 2x + so integer b8/b16 contiguous -> LS=2 and b8 contiguous -> LS=4 are covered + when the VMI source/result physical arity is the same + +floating-point lane_stride: + b8/b16 FloatType values use bit-preserving vbitcast to unsigned integer + carriers around the same pack/unpack sequence; non-FloatType low precision + types need a VPTO vbitcast contract before enabling this fallback + +arity-changing lane_stride materialization: + contiguous -> lane_stride can be expressed as multiple unpack parts, and + lane_stride -> contiguous needs an explicit multi-part merge/pack plan. + The current stage rejects those helpers instead of guessing a cross + physical-chunk materialization. +``` + +This fallback is a materialization cost, not a layout preference. Assignment may +insert the `ensure_layout`; later folding/rematerialization should remove it when +the producer or consumer has direct support: + +```text +load -> ensure_layout(lane_stride) + fold into a VMI load whose result has the requested lane_stride; vmi-to-vpto + later lowers that load to UNPK when the element width and stride match. + BRC remains the target dist for a separate scalar broadcast-load VMI semantic. + +ensure_layout(lane_stride) -> store + fold into a VMI store that directly consumes the lane_stride value; vmi-to-vpto + later lowers that store to PK/PK4 when the element width and stride match + +ordinary producer -> ensure_layout(contiguous <-> lane_stride) + lower to register pack/unpack materialization when the element width is + supported +``` + +### 1.3 Pass Responsibilities + +Dense `lane_stride` should use the existing helper-driven layout pipeline. Do +not add a separate global candidate solver for the current stage. + +```text +pto-validate-vmi-ir: + verify surface syntax before assignment + reject malformed dense lane_stride attrs once the parser accepts them + keep lane_offset unavailable in the public attr + +vmi-layout-assignment: + assign explicit dense layouts, including lane_stride, on VMI value types + use op support queries to choose local cast relations: + widening compares natural deinterleaved result arity with compact + contiguous result arity; when compact wins, request source lane_stride=W + and set result contiguous + narrowing supports the inverse relation; when arity or a supported consumer + request chooses a strided result, request source contiguous and set result + lane_stride=W + keep unsupported or conflicting uses legal by inserting ensure_layout + serialize all decisions as type attrs or helper ops + do not clone producers, fold memory ops, or solve a global cost problem + +canonicalize/cse: + remove dead helpers and merge identical rematerialized values when normal MLIR + canonicalization can prove equivalence + no lane_stride-specific decision logic + +vmi-layout-rematerialize: + consume producer -> ensure_layout shapes + clone/rematerialize cheap producers directly in the requested lane_stride + layout when the producer support query says it can create that layout + examples: scalar broadcast, splat constants, iota, layout-transparent chains, + widening ext, and supported mask producers + do not rematerialize ordinary loads unless the load form has an explicit + safe-read proof and direct UNPK lowering support + +vmi-layout-fold: + consume helper-adjacent producer/consumer shapes + fold ensure_layout(lane_stride) feeding store into a VMI store that directly + consumes the lane_stride value when the support table has a direct compact + store lowering; this is still a VMI store, not a VPTO PK op + fold load -> ensure_layout when the load can directly produce the requested + lane map with UNPK and the rewrite preserves one load at the original + program point + fold identity lane-map conversions + leave unsupported conversions as explicit ensure_layout for validation or + vmi-to-vpto materialization + +vmi-layout-sink-materialization: + move ensure_layout across pure layout-transparent ops when all operands/results + can keep one identical dense lane map + reduce duplicated contiguous <-> lane_stride materializations + do not sink through cast, load, store, reduce, group_broadcast, or control flow + +pto-validate-vmi-layout-ir: + verify every dense value has a supported layout attr + verify ensure_layout has a supported materialization path: + identity + contiguous <-> lane_stride through register pack/unpack when supported + existing contiguous <-> deinterleaved relations + verify direct layout-aware load/store choices: + LS=2 b8/b16/b32 through UNPK/PK + LS=4 b8 through UNPK4/PK4 + BRC only after a scalar broadcast-load VMI semantic is modeled + reject unsupported direct cases such as LS=4 b16/b32 compact load/store + +vmi-to-vpto: + lower only from assigned type attrs, helper ops, and op attributes + emit direct vlds/vsts dist for UNPK/PK-supported memory cases + lower surviving contiguous <-> lane_stride ensure_layout through register + pack/unpack materialization when the VPTO verifier supports the carrier path + lower widening/narrowing casts according to the assigned source/result + lane_stride relation and concrete vcvt part + emit diagnostics instead of inventing hidden layout conversions +``` + +Implementation impact by pass/component: + +| Component or pass | Lane-stride implementation work | +|---|---| +| `VMILayoutAttr` ODS/C++ helpers | Yes. Add dense `laneStride` storage, parse/print, verifier, equality, lane-map helpers, and keep it separate from group-slot carrier packing. | +| VMI type physicalization helpers | Yes. Compute dense physical arity from `laneStride`; expose carrier-slot lowering helpers for pack/unpack paths without changing the VMI logical element type. | +| `VMILayoutSupport` / target capability helpers | Yes. Add support queries for dense `lane_stride` layouts, cast layout relations, direct UNPK/PK memory support, and contiguous `<->` lane-stride materialization support. BRC remains target capability for a separate scalar broadcast-load semantic. | +| `pto-validate-vmi-ir` | No lane-stride-specific pass algorithm. It relies on attr/op verifier updates; keep the existing surface-IR validation role. | +| `vmi-layout-assignment` | Yes. Assign dense lane-stride layouts when support queries choose them; insert `ensure_layout` for incompatible uses; serialize all decisions in types/helpers. | +| `canonicalize/cse` between VMI passes | No implementation. It remains ordinary cleanup for dead helpers and identical rematerialized producers. | +| `vmi-layout-rematerialize` | Yes. Teach producer rematerialization to create requested dense lane-stride layouts for cheap/safe producers. Do not add ordinary load remat without safe-read proof. | +| `vmi-layout-fold` | Yes. Fold `ensure_layout` into layout-aware VMI consumers, especially stores that can consume lane_stride and later lower to `PK/PK4`; fold `load -> ensure_layout` into a direct layout-aware load when it can preserve one load at the original program point; fold identity lane-map conversions. | +| `vmi-layout-sink-materialization` | Minimal generic update. It should compare dense layout keys including `laneStride` and reuse existing layout-transparent sinking; do not add cast/load/store/reduce-specific lane-stride patterns here. | +| `vmi-legalize-arith-select` | No implementation. Lane stride does not change scalar-condition select legalization. | +| `pto-validate-vmi-layout-ir` | Yes. Reject unsupported assigned layouts/helpers before lowering, including unsupported LS=4 b16/b32 compact load/store and unsupported register pack/unpack materializations. | +| `vmi-to-vpto` | Yes. Lower assigned dense lane-stride layouts, direct `UNPK/PK` memory cases, register pack/unpack `ensure_layout`, and lane-stride-aware ext/trunc lowering. | +| VPTO op verifier/emitter | Only if needed by the selected support matrix. Existing `vlds/vsts` dist tokens are already present; extending register fallback to b32 or floating-point carriers requires verifier/emitter support for the corresponding pack/unpack or bitcast form. | +| Lower VPTO/backend passes after `vmi-to-vpto` | No lane-stride-specific implementation. They see ordinary VPTO ops and existing dist tokens. | + +Any pass not listed above should not implement lane-stride-specific logic in the +current stage. New behavior must enter through the explicit layout attr, +support queries, helper ops, validation, or `vmi-to-vpto` lowering. + +Current-stage component checklist: + +This checklist records the components that participate in the current-stage +lane-stride implementation. It is not the remaining-work queue; remaining work +is limited to the masked-store and group-broadcast-load items above. + +```text +include/PTO/IR/VMIAttrs.td +lib/PTO/IR/VMI.cpp + add laneStride storage for dense contiguous/deinterleaved layouts + keep group-slot laneStride parse/print compatibility + add getContiguous(ctx, laneStride) and getDeinterleaved(..., laneStride) + split helpers into isDenseLaneStrided(), isGroupSlotLaneStrided(), + getLaneStride(), and exact dense lane-map equality helpers + update attr verifier so laneStride > 0 and lane_offset is not accepted + +lib/PTO/IR/VMI.cpp +lib/PTO/Transforms/VMIToVPTO.cpp + replace the current "hasLaneStride implies unsigned carrier widening" helper + with: + logical-element physicalization for ordinary dense VPTO values + selected carrier-slot physicalization for pack/unpack materializations + existing group-slot packed-byte carrier lowering + +include/PTO/Transforms/VMILayoutSupport.h +lib/PTO/Transforms/VMILayoutSupport.cpp + extend dense store layout facts with lane_stride=2/4 cases + extend VMILayoutMaterializationSupportKind with: + ContiguousToLaneStrideViaUnpack + LaneStrideToContiguousViaPack + LaneStrideToLaneStrideViaContiguous, only if needed + update getPreferredCastLayoutFact: + keep an internal baseline natural relation for dense widening/narrowing + compute the compact lane-stride relation from the same conversion ratio + select compact only when source/result physical arities match and the + relevant arity is strictly smaller than the baseline relation + use the returned source/result layouts for both ext and trunc assignment + update getWidenSourceLayoutForResultLayout for dense lane_stride result/source + update getStoreLayoutFact and canFoldContiguousStoreMaterialization for + LS=2 b8/b16/b32 -> PK_B16/B32/B64 + LS=4 b8 -> PK4_B32 + update canMaterializeDataLayout for contiguous <-> dense lane_stride through + register pack/unpack when the element/carrier path is supported + +lib/PTO/Transforms/VMILayoutAssignment.cpp + teach natural/preferred layout collection to accept dense lane_stride facts + from VMILayoutSupport + keep conflict handling unchanged: insert ensure_layout at mismatched uses + do not add producer cloning, memory folding, or global cost selection here + +lib/PTO/Transforms/VMILayoutRematerialize.cpp + allow cheap producers to be cloned with dense lane_stride result types when + VMILayoutSupport says the producer can directly create that lane map + keep ordinary load/group_load/masked_load cloning blocked until a safe-read + proof is added for the specific rematerialized memory operation + +lib/PTO/Transforms/VMILayoutFold.cpp + add producer-side fold for load -> ensure_layout: + replace the load result layout with the ensure target layout when the load + has no other incompatible uses and VMILayoutSupport has direct UNPK + support + erase the ensure_layout and keep a single load at the original program point + do not clone ordinary loads in this fold + add fold for ensure_layout(lane_stride -> contiguous) feeding pto.vmi.store or + pto.vmi.masked_store into a VMI store that consumes the lane_stride source + directly; this pass does not emit or model VPTO PK. VMIToVPTO later selects + the corresponding PK/PK4 store lowering from the assigned VMI store contract + masked_store direct fold additionally requires the mask to carry the same + dense lane_stride layout and a compactable element-width granularity: + LS=2 b8/b16 and LS=4 b8 are supported through LOWER punpack mask compaction + LS=2 b32 is left as explicit materialization until b32 lane-stride mask + compaction is specified and implemented + a contiguous user mask is not enough, even if the value layout can be + compact-stored; assignment/rematerialization must first derive the same + lane map for the mask + fold exact dense lane-map identity helpers + do not fold unsupported LS=4 b16/b32 cases + +lib/PTO/Transforms/VMILayoutSinkMaterialization.cpp + include laneStride in dense layout equality/support checks + reuse existing layout-transparent sinking logic + do not add lane-stride-specific sinking through casts or memory ops + +lib/PTO/Transforms/PTOValidateVMIIR.cpp + no new lane-stride algorithm + validation changes should come from attr/op verifier and VMILayoutSupport + diagnostics at the layout gate + +lib/PTO/Transforms/VMIToVPTO.cpp + update OneToN physical type conversion for dense laneStride and carrier slots + lower direct compact loads: + LS=2 b8/b16/b32 -> vlds UNPK_B8/B16/B32 + LS=4 b8 -> vlds UNPK4 + lower direct compact stores: + LS=2 b8/b16/b32 -> vsts PK_B16/B32/B64 + LS=4 b8 -> vsts PK4_B32 + lower direct compact masked_stores: + LS=2 b8/b16 -> LOWER punpack mask compaction + vsts PK_B16/B32 + LS=4 b8 -> two LOWER punpack steps + vsts PK4_B32 + LS=2 b32 -> no direct masked compact store until b32 lane-stride mask + compaction is specified and implemented + lower surviving ensure_layout contiguous <-> lane_stride through vpack and + vsunpack/vzunpack when the carrier path is legal + lower lane-stride-aware ext by selecting the concrete vcvt part from + the assigned source/result relation + lower lane-stride-aware trunc by selecting the concrete vcvt part from + the assigned source/result relation + +lib/PTO/IR/VPTO.cpp +lib/PTO/Transforms/VPTOLLVMEmitter.cpp +lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp + no change for existing vlds/vsts dist tokens + extend vpack/vsunpack/vzunpack verifier/emitter only if the first implemented + fallback needs currently unsupported b64->b32 or floating-point carrier paths + +test/lit/vmi + add parser/verifier tests for dense laneStride attrs + add assignment tests for ext lane-stride facts + add fold/remat/sink tests for helper-driven rewrites + add vmi-to-vpto checks for UNPK/PK and vpack/unpack fallback + add negative tests for unsupported LS=4 b16/b32 compact load/store +``` + +Load/`ensure_layout` fold algorithm: + +```text +input shape: + %x0 = pto.vmi.load ... : !vmi.vreg + %x1 = pto.vmi.ensure_layout %x0 + : !vmi.vreg + -> !vmi.vreg + +preconditions: + the load result has no other use that requires the old layout + the load semantics are a compact logical stream + VMILayoutSupport says the target layout has a direct load lowering: + compact stream: + LS=2 b8/b16/b32 -> UNPK_B8/B16/B32 + LS=4 b8 -> UNPK4 + masks/passthroughs, if present, already have compatible assigned layouts + +rewrite: + replace the original load op in place, or create the replacement load at the + same insertion point and erase the old load + the replacement load result type is the ensure target type + all ensure users use the replacement load result + erase the ensure_layout + +output shape: + %x = pto.vmi.load ... : !vmi.vreg + +lowering: + vmi-to-vpto emits the corresponding vlds UNPK dist +``` + +This fold changes the assigned result layout of the existing load; it does not +clone the load at the helper use-site. If the original load has both contiguous +and lane-stride consumers, the fold must leave the helper in place unless a +separate rematerialization step has a safe-read proof to clone the load. + +### 1.4 Scenario Ownership + +Each optimization scenario has exactly one owning pass. Other passes may verify +or lower the resulting explicit IR, but should not solve the same rewrite again. + +| Scenario | Example shape | Owning pass | Non-owners | +|---|---|---|---| +| Assign a layout request | `ext -> store` where store wants contiguous | `vmi-layout-assignment` inserts explicit layouts/helpers | Assignment does not clone, fold, or lower | +| Direct load produces requested lane map | `load(contiguous) -> ensure_layout(lane_stride=2)` | `vmi-layout-fold` rewrites the original load result layout when UNPK support exists | Remat must not clone this load without safe-read proof | +| Direct store consumes lane map | `ensure_layout(lane_stride -> contiguous) -> store` | `vmi-layout-fold` rewrites the VMI store to consume the lane_stride source directly when direct compact-store support exists | `vmi-to-vpto` emits the actual `vsts PK/PK4` | +| Cheap producer can produce target layout | `broadcast -> ensure_layout(lane_stride=2)` | `vmi-layout-rematerialize` rebuilds broadcast with lane-stride result | Fold does not rebuild arbitrary producers | +| Cast chooses arity-reducing relation | `64xf16 -> 64xf32` or a supported narrowing with smaller strided result | `vmi-layout-assignment` chooses the cast source/result layout relation | Remat only handles later use-site requests; fold only handles adjacent load/store helpers | +| Cast can move materialization to cheap source | `ext/trunc -> ensure_layout(requested layout)` with source broadcast/load-fold case | `vmi-layout-rematerialize` rebuilds the cast with the requested relation | Assignment may already choose the self-preferred relation; fold only handles the load/store subcase | +| Layout-transparent op has ensured operands | `ensure(a), ensure(b) -> add` | `vmi-layout-sink-materialization` sinks matching helpers to the result | Remat handles the opposite shape `add -> ensure` | +| Surviving supported helper | `ensure_layout(contiguous <-> lane_stride)` after optimizations | `vmi-to-vpto` lowers to register pack/unpack | Earlier passes are allowed to leave it explicit | +| Unsupported helper or layout | `lane_stride=4 b16 compact store` | `pto-validate-vmi-layout-ir` rejects before lowering | `vmi-to-vpto` should not invent a repair | +| Multi-consumer value with incompatible layouts | one load feeds contiguous user and lane-stride user | baseline keeps helper; optional remat only with safe-read proof | Fold must not silently duplicate memory effects | + +Examples: + +```text +load fold, owned by vmi-layout-fold: + before: + %x0 = pto.vmi.load ... : contiguous + %x1 = pto.vmi.ensure_layout %x0 : contiguous -> lane_stride=2 + after: + %x1 = pto.vmi.load ... : lane_stride=2 + vmi-to-vpto: + %x1 = pto.vlds ... {dist = "UNPK_B8/B16/B32 or UNPK4"} + +store fold, owned by vmi-layout-fold: + before: + %x_c = pto.vmi.ensure_layout %x_ls : lane_stride=2 -> contiguous + pto.vmi.store %x_c, %dst + after: + pto.vmi.store %x_ls, %dst // VMI store consumes lane_stride source + vmi-to-vpto: + pto.vsts %x_ls, %dst {dist = "PK_B16/B32/B64 or PK4_B32"} + +broadcast remat, owned by vmi-layout-rematerialize: + before: + %b0 = pto.vmi.broadcast %s : contiguous + %b1 = pto.vmi.ensure_layout %b0 : contiguous -> lane_stride=2 + after: + %b1 = pto.vmi.broadcast %s : lane_stride=2 + +elementwise sink, owned by vmi-layout-sink-materialization: + before: + %a1 = ensure_layout %a0 -> lane_stride=2 + %b1 = ensure_layout %b0 -> lane_stride=2 + %c1 = pto.vmi.addf %a1, %b1 + after: + %c0 = pto.vmi.addf %a0, %b0 + %c1 = ensure_layout %c0 -> lane_stride=2 +``` + +## 2. IR Attribute Changes + +### 2.1 Extend `VMILayoutAttr` + +Current storage reuses `blockElems` as group-slot `lane_stride`. Generalization +should first split lane stride from block elems: + +```c++ +kind +factor +blockElems +slots +laneStride +``` + +Meaning by kind: + +```text +contiguous: + factor = 1 + blockElems = 1 + slots = 0 + laneStride >= 1 + +deinterleaved: + factor = F + blockElems = B + slots = 0 + laneStride >= 1 + +num_groups: + factor = G + blockElems = 1 + slots = K + laneStride >= 1 +``` + +Do not add a public `laneOffset` field in the current stage. The targeted +optimization only needs phase-zero strided dense layouts. A future +phase field is justified only when there is a concrete VMI value whose logical +lane map is intentionally non-zero-phase, for example a zero-copy +deinterleave/extract view that keeps the odd lanes in place or a narrowing +conversion whose consumer explicitly requires an odd-lane result. + +Recommended helpers: + +```c++ +bool isDense() const; +bool hasDenseLaneStride() const; +bool hasGroupSlotLaneStride() const; +int64_t getLaneStride() const; +VMILayoutAttr withLaneStride(int64_t stride) const; +``` + +Keep old constructor defaults source-compatible where possible: + +```c++ +getContiguous(ctx) +getDeinterleaved(ctx, factor, blockElems = 1, + laneStride = 1) +getGroupSlots(ctx, numGroups, slots = 0, laneStride = 1) +``` + +### 2.2 Parser And Printer + +Accepted dense spellings: + +```text +#pto.vmi.layout +#pto.vmi.layout + +#pto.vmi.layout +#pto.vmi.layout +#pto.vmi.layout +``` + +Existing group-slot spellings remain valid: + +```text +#pto.vmi.layout +#pto.vmi.layout +``` + +Printing omits defaults: + +```text +lane_stride = 1 is omitted +``` + +### 2.3 Verifier + +Verifier rules: + +```text +all layouts: + laneStride > 0 + +contiguous: + factor == 1 + blockElems == 1 + slots == 0 + +deinterleaved: + factor in supported dense factors + blockElems > 0 + slots == 0 + +num_groups: + factor > 0 + slots >= 0 + blockElems == 1 +``` + +The verifier should not require every strided layout to fit one VPTO register. +Fit depends on the VMI type shape and element type, so it belongs in type +physicalization and op support checks. + +## 3. Physicalization Helpers + +### 3.1 Separate Element Carrier From Lane Map + +Replace the current shared helper shape: + +```c++ +getVMIPhysicalElementType(type) +``` + +with two concepts: + +```c++ +getVMILogicalStorageElementType(type) +getVMIPhysicalCarrierElementType(type, loweringKind) +``` + +Dense lane-strided values keep the VMI logical element type. The lowering may +represent the same lane map either as logical-element lanes or as wider carrier +slots when the selected VPTO instruction is a pack/unpack family. + +Logical-element lane representation: + +```text +!vmi.vreg + -> !pto.vreg<128xf16> physical register +``` + +Carrier-slot representation for pack/unpack lowering: + +```text +!vmi.vreg + -> low ui16 in each ui32 slot for vpack/PK_B32-style lowering + +!vmi.vreg + -> low ui8 in each ui32 slot for PK4_B32-style lowering +``` + +Group-slot packed stores also request a wider carrier in the specific lowering +path: + +```text +!vmi.vreg + group_store -> b32 carrier + PK4_B32 +``` + +Do not let dense `hasLaneStride()` imply unsigned-integer carrier widening +globally. Carrier widening is a property of a selected materialization or +load/store lowering, not of the VMI logical type itself. + +### 3.2 Physical Arity + +Add a dense lane-map helper: + +```c++ +struct DenseLaneMap { + int64_t deinterleaveFactor; + int64_t blockElems; + int64_t laneStride; +}; + +int64_t getPhysicalLaneForDenseLogicalLane(DenseLaneMap map, + int64_t logicalLane); +``` + +For a VMI vreg type: + +```text +lanesPerVPTO = getVPTOPhysicalLanes(elementType) +lanesInDensePart = ceil(N / F) with block-aware distribution +requiredLanes = O + (lanesInDensePart - 1) * LS + 1 +registersPerDensePart = ceil(requiredLanes / lanesPerVPTO) +physicalArity = F * registersPerDensePart +``` + +For the current stage, require full block divisibility for dense +deinterleaved strided layouts, matching existing direct lowering restrictions: + +```text +N % (F * B) == 0 +``` + +Relaxing tail handling is outside the current stage and should be enabled only +with an explicit materialization/lowering proof. + +## 4. Layout Support Interface + +Extend support queries to include dense strided layouts: + +```text +supportsResultLayout(op, resultIndex, layout) +supportsOperandLayout(op, operandIndex, layout) +supportsLayoutRelation(op, operandLayouts, resultLayouts) +``` + +The important change is relation support. Some ops are not independently +described by "operand supports layout X" and "result supports layout Y"; they +support specific pairs. + +Examples: + +```text +elementwise: + all dense operands/results must use identical dense layout key + +extf/extui/extsi: + source/result layouts must satisfy a widening relation. Assignment chooses + between the natural deinterleaved relation and the compact-result + lane-stride-source relation by comparing physical arity, not by matching a + concrete lane count such as 64. + +truncf/trunci: + source/result layouts must satisfy a narrowing relation. Assignment uses the + inverse relation conservatively: keep the natural deinterleaved-source to + contiguous-result relation unless arity or a supported consumer request + selects a strided result relation. Masked-store consumers may only use the + strided result relation when the value and mask can be assigned/materialized + to the same lane map. + +broadcast/group_broadcast: + result may use a dense layout only when the materialization lowering has an + explicit support case for that lane map + +load: + default result contiguous + producer rematerialization may create selected strided layouts if a direct + load/mask sequence can produce that lane map + +store: + memory effect is contiguous unless the op is an explicit logical interleave + store; a strided source requires store lowering support or ensure_layout +``` + +Assignment should still insert `ensure_layout` for incompatible use-local +requests. Rematerialization/fold can later remove it. + +### 4.1 Cast Relation Helper Shape + +Keep `getPreferredCastLayoutFact` as the assignment entry point for dense +widening and narrowing casts, but make the helper return the actual preferred +source/result relation for the current shape. Internally it first builds the +natural relation: + +```text +widen: + source contiguous + result deinterleaved=W + +narrow: + source deinterleaved=W + result contiguous +``` + +Then it computes the compact relation: + +```text +widen: + source contiguous, lane_stride=W + result contiguous + +narrow: + source contiguous + result contiguous, lane_stride=W +``` + +The compact relation is selected only when its source/result physical arities +match and it strictly reduces the relevant baseline arity: + +```text +widen: + physical_arity(compact result) < physical_arity(natural result) + +narrow: + physical_arity(compact source) < physical_arity(natural source) +``` + +If the compact relation does not win, the helper returns the natural relation. +`vmi-layout-assignment` calls this helper for `extf/extui/extsi` and +`truncf/trunci`, requests the returned source layout, and records the returned +result layout. + +The support query must validate the returned pair before assignment commits it: + +```text +supportsExtRelation(sourceTypeWithLayout, resultTypeWithLayout) +supportsTruncRelation(sourceTypeWithLayout, resultTypeWithLayout) +``` + +The validation step is a legality check, not a second optimizer. + +### 4.2 Current Framework Fit + +The existing assignment pass already has use-site requests. For example, +`pto.vmi.store` requests a contiguous source operand, and assignment can insert +`ensure_layout` when the stored value is assigned another layout. + +The dense-stride `ext` optimization should keep the same model: the cast op is +the layout-entry point and stores one preferred source/result relation. The old +preferred relation was: + +```text +extf: + request source contiguous + set result deinterleaved=W +``` + +The current stage keeps the existing single-preference framework and lets +`ext` choose one fact for the current op: + +```text +baseline fact: + source contiguous + result deinterleaved=W + +lane-stride fact: + source lane_stride=W + result contiguous +``` + +The `ext` support query chooses between these facts from op-local information: + +```text +conversion ratio W +target support for one selected hardware conversion part +physical arity of the natural result layout +physical arity of the compact contiguous result layout +requested result layout when a consumer materialization/remat path provides one +``` + +It does not inspect the defining source producer. If compact result arity is +strictly smaller than natural result arity and the target supports the +single-part relation, it selects the lane-stride fact. If it selects the +lane-stride fact and the source is not already in that layout, assignment +inserts an explicit source `ensure_layout`. Later passes either discharge that +helper by rematerializing/folding a concrete producer, lower it with a +registered pack/unpack materializer, or let validation reject the unsupported +relation. + +## 5. Widening Conversion Lowering + +Let: + +```text +W = result element storage bits / source element storage bits +``` + +For a dense source layout: + +```text +source lane_stride = LS +``` + +Single-part lowering is legal when: + +```text +LS % W == 0 +``` + +Then: + +```text +hardware part = 0 +result lane_stride = LS / W +``` + +The current stage only emits the zero-phase single-part conversion. +`vcvt ODD` remains necessary for full packed contiguous conversion, but that is +handled by the existing multi-part relation: + +```text +source contiguous, lane_stride=1 +result deinterleaved=W +``` + +Do not add a phase field merely to name that existing ODD instruction. Add a +phase field only when an assigned VMI layout needs to represent a concrete +zero-copy value/view already resident in odd/non-zero-phase lanes. + +The support query for the conversion should accept the pair only when the +requested result layout equals this computed result lane map, including +deinterleave/block fields. + +The support query should expose helpers for both legal facts, but assignment +chooses one immediately: + +```text +baseline fact: + source contiguous + result deinterleaved=W + natural result arity = physical_arity(result deinterleaved=W) + +lane-stride fact: + result contiguous + source same dense shape with lane_stride = W + compact result arity = physical_arity(result contiguous) +``` + +Assignment uses this deterministic rule: + +```text +if compact result arity < natural result arity + and the lane-stride fact is supported: + choose lane-stride fact +else: + choose baseline fact +``` + +For example, for `f16 -> f32`, the `extf` op chooses +`source lane_stride=2 -> result contiguous` for `64xf32`, because the compact +result has one physical chunk while the natural `deinterleaved=2` result has two +physical chunks. For `128xf32` and `256xf32`, both layouts have the same result +arity, so assignment chooses the natural `deinterleaved=2` result. The source +producer is handled by the explicit source `ensure_layout` and later +fold/rematerialization; it is not part of the cast support query. + +Current contiguous widening remains a separate legal relation: + +```text +source dense contiguous, lane_stride=1 +result deinterleaved=W, lane_stride=1 +``` + +Implementation steps: + +1. Factor conversion ratio calculation by storage bit width. +2. Add helper that computes the natural result layout and its physical arity. +3. Add helper that computes the compact result layout, required source + lane-stride layout, and compact result physical arity. +4. Teach `VMIToVPTO` conversion lowering to emit only the selected hardware + part when the relation is single-part. +5. Keep existing multi-part lowering for contiguous-to-deinterleaved cases. +6. Add diagnostics when an assigned conversion layout pair has no lowering. + +Hardware part names should be abstracted: + +```text +W=2: + part 0 -> EVEN + part 1 -> ODD + +W=4: + part 0..3 -> target-specific conversion part names or the existing sequence +``` + +Do not special-case f16/f32 in the matcher. The type only determines `W` and +the concrete VPTO conversion opcode. + +## 6. Narrowing Conversion Lowering + +Let: + +```text +W = source element storage bits / result element storage bits +``` + +For a single-part narrowing relation: + +```text +result lane_stride = source lane_stride * W +hardwarePart = 0 for the current stage +``` + +The narrowing assignment relation is the inverse of widening, but it must not +blindly choose a lane-stride result. Build two facts: + +```text +baseline fact: + source deinterleaved=W + result contiguous + natural result arity = physical_arity(result contiguous) + +lane-stride fact: + source contiguous + result contiguous, lane_stride=W + strided result arity = physical_arity(result lane_stride=W) +``` + +Then choose a strided result only when it is justified: + +```text +if strided result arity < natural result arity + and the lane-stride fact is supported: + choose lane-stride fact +else if a consumer/requested result layout is the strided result + and the lane-stride fact is supported: + choose or rematerialize lane-stride fact +else: + choose baseline fact +``` + +This keeps trunc symmetric with ext while avoiding the earlier mistake of +producing lane_stride solely because the operation is a narrowing cast. A +consumer may still request or preserve a strided result. For example, an +ordinary store with direct `PK` support can consume a supported lane-stride +result, and rematerialization/fold may keep that relation. A masked store may +do so only when the mask can be assigned/materialized to the same lane map. + +Implementation steps: + +1. Share ratio, dense-factor, lane-map, and physical-arity helpers with + widening. +2. Add helper that computes the natural source/result relation and result + arity. +3. Add helper that computes the strided-result relation and result arity. +4. Add support query for valid narrowing layout pairs. +5. Teach assignment/rematerialization to select the strided fact for explicit + result requests, direct compact-store consumers, or true arity reductions. +6. Lower single-part narrowing directly when the target has a part-selecting + narrow instruction. +7. Preserve existing deinterleaved-to-contiguous narrowing for the packed full + result case. + +This is the same family as the recently discussed `d4 -> c -> d2 -> vcvt -> c` +optimization: if a cast op has a direct source/result layout relation, +assignment/rematerialization should expose that relation before lowering. + +## 7. Ensure-Layout And Rematerialization + +### 7.1 `ensure_layout` + +`ensure_layout` remains the explicit use-site materialization op. + +Verifier/lowering policy: + +```text +same source and target dense lane map: + fold away + +known dense relation: + lower contiguous <-> lane_stride through register pack/unpack when supported + lower contiguous/deinterleaved relations through existing intlv/dintlv paths + +producer can rematerialize target layout: + rematerialization should replace ensure_layout(producer) + +unknown relation: + reject before vmi-to-vpto +``` + +Avoid adding a generic "any dense layout to any dense layout" promise unless the +target really has a lowering for it. + +### 7.2 Rematerialization + +The current checked-in `vmi-layout-rematerialize` cheap producers are: + +```text +data: + VMIExtFOp / VMIExtSIOp / VMIExtUIOp when the source layout can be + materialized for the requested result relation + VMIFmaOp + binary layout-transparent ops: + addf/addi/subf/subi/mulf/muli/divf/minf/maxf/andi/ori/xori/shli/shrui/shrsi + unary layout-transparent ops: + negf/absf/absi/sqrt/exp/ln/relu/not + VMIConstantOp only when the DenseElementsAttr is a splat + VMIBroadcastOp + VMIIotaOp + +mask: + VMICreateMaskOp + VMICreateGroupMaskOp + VMIConstantMaskOp + +special rewrite: + selected VMITruncFOp / VMITruncIOp through source/result ensure_layout when + the cast relation is a supported narrowing relation +``` + +Not included as cheap producers in the current pass: + +```text +load / masked_load / group_load / group_slot_load / group_broadcast / +group_broadcast_load / store / reduce / control-flow ops +``` + +Loads need a separate policy. `load -> ensure_layout` should be folded in +`vmi-layout-fold` when one original load can directly produce the requested +layout. A normal load should not be cloned/rematerialized unless a later safe-read +proof explicitly permits that clone. + +Relationship between cheap producers and dense `lane_stride`: + +```text +assignment: + creates the target layout request explicitly, usually as ensure_layout(... -> + lane_stride) or as a cast source/result relation. For casts, assignment may + itself choose the arity-reducing lane-stride relation; remat only reacts to + later use-site layout requests. + +rematerialize: + does not choose lane_stride as a preference + only consumes the explicit helper/request and rebuilds the producer with the + requested lane_stride result type when the producer is cheap and locally legal +``` + +Required rematerialize changes for dense `lane_stride`: + +```text +materializeDataLayout: + no special producer logic, but canMaterializeDataLayout must understand + contiguous <-> lane_stride through register pack/unpack + +splat constant / broadcast / iota: + rebuild the op with the requested lane_stride result type + lowering later materializes that layout directly or through ensure_layout + +layout-transparent unary/binary/fma: + rebuild the op with the requested lane_stride result type + materialize each operand to the same lane_stride layout before rebuilding + this relies on canMaterializeDataLayout for operand conversions + +widening ext: + update getWidenSourceLayoutForResultLayout so a requested result layout derives + the required source lane_stride: + result contiguous, W=2 -> source lane_stride=2 + result lane_stride=R, W=2 -> source lane_stride=2*R + remat then inserts/uses source ensure_layout and rebuilds ext with the + requested result layout + +narrowing trunc: + add getNarrowSourceLayoutForResultLayout or an equivalent relation helper. + For a requested result lane_stride=R and narrowing ratio W, derive the source + layout that can produce that result with a selected hardware part: + result lane_stride=W, W=2 -> source contiguous + result lane_stride=R, W=2 -> source lane_stride=R/W when divisible + remat then inserts/uses source ensure_layout and rebuilds trunc with the + requested result layout + +trunc source-ensure rewrite: + extend the existing source-ensure rewrite to recognize lane_stride narrowing + relations for VMITruncFOp and VMITruncIOp, not only deinterleaved narrowing + relations + +mask producers: + only participate after mask layout support defines the corresponding + lane-stride or predicate-granularity relation; otherwise unchanged +``` + +Example: + +```text +before remat: + %b0 = pto.vmi.broadcast %s : !vmi.vreg<64xf16, contiguous> + %b1 = pto.vmi.ensure_layout %b0 + : contiguous -> contiguous, lane_stride=2 + %y = pto.vmi.extf %b1 : f16 -> f32 + +after remat: + %b1 = pto.vmi.broadcast %s + : !vmi.vreg<64xf16, contiguous, lane_stride=2> + %y = pto.vmi.extf %b1 : f16 -> f32 +``` + +This removes a register layout materialization and lets `vmi-to-vpto` lower the +ext as the single selected conversion part. It is still driven by the explicit +layout request; remat does not inspect sibling consumers or choose lane_stride by +itself. + +Do lane-stride cast rematerialization only in these cases: + +```text +required shape: + cast result is followed by ensure_layout to a requested dense result layout + widening or narrowing ratio W > 1 + the requested source/result layout pair is accepted by the cast relation + helper + the cast with that source/result layout can lower as one selected conversion + part or the existing multi-part relation + +acceptance/safety gate: + the source-side lane_stride request must be discharged by a concrete local + rewrite, not merely moved from result side to source side + accepted cases: + the source already has the required lane_stride + the source producer is in the checked-in cheap producer list and can be + rebuilt with the required lane_stride + the source is load -> ensure_layout and vmi-layout-fold can replace it with + a single original-position layout-aware VMI load + a layout-transparent chain can be sunk/rematerialized until one of the above + concrete producer cases is reached + +do not apply: + result consumer already accepts the natural cast layout + requested cast layout relation is unsupported + source is an ordinary load with other incompatible consumers and no safe-read + proof to clone it + the rewrite only moves an expensive materialization from result side to source + side without exposing a direct lowering +``` + +Typical accepted shapes: + +```text +broadcast -> ext -> ensure_layout(contiguous) -> store + remat broadcast as lane_stride=W + ext lowers with one conversion part + no source-side ensure_layout remains + +load -> ensure_layout(lane_stride=W) -> ext -> store + fold load into a layout-aware VMI load + vmi-to-vpto later emits the matching UNPK dist + ext lowers with one conversion part + no extra load is cloned + +elementwise cheap chain -> ext -> ensure_layout(contiguous) + remat/sink the chain to lane_stride=W only when the chain reaches a concrete + cheap producer or direct load-fold case + +trunc -> ensure_layout(lane_stride=W) -> compact store + remat/rebuild trunc with the requested lane_stride result when the source + layout relation is supported + store fold may then consume the lane_stride result directly + +trunc -> ensure_layout(lane_stride=W) -> masked_store + only accepted after mask layout assignment can provide the same lane map for + the predicate; otherwise keep the conservative contiguous masked-store path +``` + +## 8. Broadcast And E2B Interaction + +Do not encode E2B in `lane_stride`, and do not define +`vmi.group_broadcast_load` in terms of E2B. The VMI operation is a logical +fused memory operation: + +```text +for each logical group g: + scalar = source[offset + g * source_group_stride] + for each lane i in group g: + result[i] = scalar +``` + +The result layout is assigned separately. It may be contiguous, +deinterleaved, or dense lane-strided if the consumer asks for that lane map and +the target support table accepts it. E2B is only one VPTO lowering strategy for +a restricted subset of this logical operation. + +The layering should be: + +```text +logical group broadcast load + -> assigned dense layout, possibly lane-strided + -> support query chooses a lowering strategy + -> selected VPTO dist, if any +``` + +For the current E2B strategy, the support query checks: + +```text +source is direct memory +source_group_stride is constant 1 +num_groups is a multiple of 8 +element storage width +logical group size derived from num_groups +assigned result layout: + contiguous for the direct packet size + or deinterleaved=2, block_elems=1 for the split packet size +``` + +Then it may choose an E2B packet: + +```text +b16 contiguous: direct 1 -> 16 packet +b16 deinterleaved=2: two logical halves / 1 -> 32 reuse +b16 dense lane_stride=2: direct phase-zero strided consumer packet +b32 contiguous or strided: target-specific packet size +``` + +If those conditions do not hold, the operation is still a valid VMI semantic if +some other lowering exists, such as `group_slot_load + group_broadcast`, scalar +loads plus broadcast, or future target-specific broadcast-load support. The +failure is only "this E2B lowering strategy is not applicable", not "the VMI +operation means E2B". + +Concrete implementation plan for lane-stride `group_broadcast_load`: + +```text +include/PTO/Transforms/VMILayoutSupport.h +lib/PTO/Transforms/VMILayoutSupport.cpp + +1. Split semantic support from E2B strategy checks: + + getGroupBroadcastLoadSupport(capabilities, op) + try getE2BGroupBroadcastLoadSupport(capabilities, op) + if success: + return {kind = E2BVlds} + return failure("no registered group_broadcast_load lowering strategy; " + "E2B rejected because ...") + + getE2BGroupBroadcastLoadSupport(capabilities, op) + contains the current E2B constraints: + source is !pto.ptr direct memory + element width is b16 or b32 + source_group_stride is constant 1 + num_groups is a multiple of 8 + group size matches direct or split E2B packet size + result layout is contiguous or deinterleaved=2/block_elems=1 + result has full physical chunks + +2. Keep VMIGroupBroadcastLoadSupportKind strategy-specific: + E2BVlds means "lower this VMI semantic using E2B" + It must not be used as the definition of the VMI op. +``` + +```text +lib/PTO/Transforms/VMILayoutAssignment.cpp + +3. Rename strategy helpers so the direction is clear: + + isE2BGroupBroadcastLoadCandidate + -> isE2BGroupBroadcastLoadStrategyApplicable + + getPreferredGroupBroadcastLoadLayout + -> getPreferredE2BGroupBroadcastLoadLayout + +4. Fusion from group_slot_load + group_broadcast to group_broadcast_load remains + guarded by E2B applicability. If E2B is not applicable, do not create a + fused group_broadcast_load merely because the VMI semantic would be valid. + That avoids producing an op with no registered lowering strategy. + +5. Layout assignment for an explicit group_broadcast_load uses the support + query: + if E2B strategy applies: + assign the E2B-preferred result layout + else: + leave the op to validation unless a fallback strategy is added +``` + +```text +lib/PTO/Transforms/VMIToVPTO.cpp + +6. Replace duplicated local E2B legality checks with: + support = getGroupBroadcastLoadSupport(capabilities, op) + switch support.kind: + E2BVlds: + emit the existing E2B packet sequence + + The E2B lowering code may still assert/recheck structural invariants needed + for indexing, but user-facing diagnostics should come from the support query. + +7. Diagnostics must name the strategy: + good: "group_broadcast_load has no registered lowering strategy; E2B + rejected because source_group_stride is not constant 1" + bad: "group_broadcast_load requires constant unit source_group_stride" + + The second form is only valid inside an E2B-specific diagnostic. +``` + +Required group-broadcast-load tests: + +```text +E2B positive: + explicit group_broadcast_load with b16/b32, stride=1, matching group size, + and assigned contiguous/deinterleaved result layout + CHECK vmi-to-vpto emits E2B_B16/E2B_B32 + +E2B strategy rejection: + source_group_stride != 1, wrong group size, or unsupported element width + CHECK validation/lowering diagnostic says no registered lowering strategy and + reports E2B as the rejected strategy + +fusion guard: + group_slot_load + group_broadcast shape that is not E2B-applicable + CHECK assignment does not fuse it into group_broadcast_load + +semantic boundary: + explicit group_broadcast_load that is not E2B-applicable + CHECK failure wording does not redefine the op as E2B and does not imply the + logical VMI semantic itself is E2B +``` + +This keeps broadcast optimization generic across type width and layout, instead +of hardcoding one `ComputeY1ToFP8` scale pattern. + +## 9. Tests + +Use the following as the coverage matrix for current-stage support plus the +masked-store and group-broadcast-load follow-up items. It is not a separate +list of all remaining implementation work. + +Parser/verifier: + +```text +parse/print contiguous lane_stride +parse/print deinterleaved + block_elems + lane_stride +``` + +Physicalization: + +```text +64xf16 contiguous lane_stride=2 has one physical 128xf16 part +ui16 contiguous lane_stride=2 may lower through low ui16 in ui32 carrier slots +when the selected materialization is vpack/PK_B32 +ui8 contiguous lane_stride=4 may lower through low ui8 in ui32 carrier slots +when the selected materialization is PK4_B32 +65xf16 contiguous lane_stride=2 is rejected by direct full-chunk-only paths, or +covered only by an arity-changing materialization test outside this discussion +group-slot ui8 lane_stride=4 keeps existing carrier lowering behavior +``` + +Conversion lowering: + +```text +f16 lane_stride=2 -> f32 contiguous emits one EVEN conversion +bf16 lane_stride=2 -> f32 contiguous follows the same relation +ui8 lane_stride=2 -> ui16 contiguous follows W=2 +ui8 lane_stride=4 -> ui32 contiguous follows W=4 when target supports it +contiguous f16 -> deinterleaved=2 f32 still emits EVEN + ODD +f32 contiguous -> f16 lane_stride=2 emits the selected narrowing part when the +assigned relation is supported +f32 deinterleaved=2 -> f16 contiguous keeps the existing packed full-result +narrowing relation +ui16 lane_stride=2 -> contiguous can materialize with vpack 32->16 carrier path +ui8 lane_stride=4 -> contiguous can materialize with two vpack stages +``` + +Assignment/rematerialization: + +```text +extf records a strided dense source relation when compact result arity is +smaller than natural result arity +extf 64xf16 -> 64xf32 chooses source lane_stride=2, result contiguous +extf 128xf16 -> 128xf32 chooses result deinterleaved=2 +extf 256xf16 -> 256xf32 chooses result deinterleaved=2 +truncf records a strided result relation only when the conservative +self-preference/support rule or a supported consumer request selects it; it +does not choose lane_stride solely because the op narrows +layout-transparent op propagates the same strided layout through operands/result +ensure_layout is folded when source and target lane maps match +rematerialization clones a cheap broadcast for two different dense layouts +``` + +End-to-end assignment cases: + +```text +contiguous load -> ext -> contiguous store: + uses lane_stride only when the source ensure_layout can be folded to the + original load, rematerialized from a cheap producer, or lowered by a supported + register materializer + +cheap broadcast -> ext -> contiguous store: + rematerializes broadcast as lane_stride=2 and lowers ext with one EVEN part + +producer -> ext -> deinterleaved reduce: + keeps source contiguous and result deinterleaved=2 + +cheap producer -> ext feeding both store and reduce: + keeps shared deinterleaved path for reduce and rematerializes a contiguous + result path for store only through the checked cheap-producer remat path + +group_broadcast_load -> ext -> contiguous consumer: + chooses lane_stride only if group_broadcast_load supports that lane map +``` + +Negative tests: + +```text +assigned ext layout pair where LS % W != 0 and no multi-part relation exists +assigned trunc layout pair where result lane_stride is not compatible with the +narrowing ratio +ordinary dense op with mismatched lane_stride operands +store consuming strided dense layout without a supported store/materialization +masked_store consuming lane_stride value with a stale contiguous user mask is +rejected or kept on the conservative contiguous path +``` + +## 10. Suggested Patch Order + +1. Add attr fields, parser/printer, verifier, and round-trip tests. +2. Split dense lane-map physicalization from group-slot carrier packing. +3. Update physical arity/unpack helpers for dense lane stride. +4. Extend support queries and assignment layout keys. +5. Implement widening arity-driven self-preference, single-part relation, and + tests. +6. Implement narrowing inverse relation support, consumer-request handling, and + tests. +7. Teach rematerialization/fold about exact dense lane-map equality. +8. Add broadcast/E2B recognition improvements that consume assigned lane maps. + +Each step should keep existing group-slot `lane_stride` tests passing. The first +functional optimization can be the `f16/bf16 lane_stride=2 -> f32 contiguous` +single-part conversion, but the IR and helper changes should already be generic +over type width and lane-map fields. diff --git a/docs/designs/vmi-layout-assignment-implementation.md b/docs/designs/vmi-layout-assignment-implementation.md new file mode 100644 index 0000000000..f0c83db4c7 --- /dev/null +++ b/docs/designs/vmi-layout-assignment-implementation.md @@ -0,0 +1,2323 @@ +# VMI Layout Assignment Implementation Plan + +本文是 `vmi-layout-assignment` 和 `vmi-to-vpto` 的实现计划。它配套 +`vmi-layout-assignment-lowering-design.md`，并以 +`vmi-layout-lowering-cases.md` 为测试和验收来源。 + +不使用早期 VMI 草稿作为设计输入。 + +## 1. Pipeline + +Recommended pass pipeline: + +```text +pto-validate-vmi-ir + -> vmi-layout-assignment // hard legalization baseline + -> canonicalize/cse + -> vmi-layout-rematerialize // optional optimization + -> canonicalize/cse + -> vmi-layout-fold // optional optimization over remat-exposed helpers + -> canonicalize/cse + -> vmi-layout-sink-materialization // optional optimization + -> canonicalize/cse + -> vmi-legalize-arith-select + -> pto-validate-vmi-layout-ir + -> vmi-to-vpto + -> canonicalize/cse + -> existing VPTO lowering/codegen +``` + +Only `vmi-layout-assignment` is required for the first legal implementation. +The optimization passes may be introduced one by one. Their contract is that +they consume legal layout-assigned VMI IR and produce legal layout-assigned VMI +IR; they never move a hidden decision into `vmi-to-vpto`. + +Pass responsibilities: + +```text +pto-validate-vmi-ir: + verify surface VMI has no physical VPTO layout dependency + reject public/external VMI ABI unless explicitly enabled + +vmi-layout-assignment: + solve hard value layout constraints + choose explicit layouts visible in IR + insert ensure_layout / ensure_mask_layout / ensure_mask_granularity helpers + make internal function boundary layouts explicit + rewrite VMI types with layout attrs + +canonicalize/cse: + remove dead helpers and merge identical cloned producers where MLIR legality + permits + +vmi-layout-fold: + fold use-site materialization into consumers that can directly consume the + source layout while preserving the same logical effect + example: ensure_layout(deinterleaved=2 -> contiguous) feeding store may become + a store of deinterleaved=2 when the store has a layout-aware vstsx2 INTLV + lowering + current implementation: pto.vmi.store and the value operand of + pto.vmi.masked_store when the existing mask arity matches, fed by + ensure_layout from deinterleaved=2/4, block_elems=1 to contiguous. factor=2 + uses the store's vstsx2 INTLV lowering; factor=4 is still store-local, but it + materializes through physical interleave before vsts. + +vmi-layout-rematerialize: + replace explicit ensure_* helpers with cloned cheap layout-polymorphic + producers when the clone directly creates the requested result type + current implementation: splat pto.vmi.constant, pto.vmi.broadcast, + pto.vmi.iota, selected layout-transparent data ops, widening + pto.vmi.ext{f,si,ui}, pto.vmi.create_mask, pto.vmi.create_group_mask, and + pto.vmi.constant_mask. Relation-aware remat rewrites result-side + ensure_layout through layout-transparent producers and widening ext + producers, leaving any newly exposed producer-side helpers for the following + vmi-layout-fold. + not included in the first implementation: load, group_load, masked_load, + group_slot_load, and group_broadcast; those require separate memory, + execution-count, or source-layout proof before they can be rematerialized + +vmi-layout-sink-materialization: + move ensure_layout across pure layout-transparent elementwise chains when the + rewritten IR reduces materialization overhead and keeps every op locally legal + current implementation: sink two identical operand ensure_layout helpers + across binary add/sub/mul/div/min/max/and/or/xor/shl/shru VMI ops, three + identical operand ensure_layout helpers across fma, or one source + ensure_layout across unary neg/abs/sqrt/exp/ln/relu/not VMI ops, producing + one result ensure_layout. It also sinks compare data helpers to one result + ensure_mask_layout, and sinks select only when both selected values and the + mask carry matching explicit helpers. Matching ensure_mask_layout or + ensure_mask_granularity helpers are sunk across mask_and/mask_or/mask_xor/ + mask_not, producing one result mask helper. It does not sink through cast, + load, store, reduce, group_broadcast, or control-flow ops. + +vmi-legalize-arith-select: + restore scalar-condition arith.select with VMI result type back to scf.if + after canonicalize; canonicalize may fold simple scf.if into arith.select, + but VMI values must not cross non-VMI semantic ops before vmi-to-vpto + +pto-validate-vmi-layout-ir: + verify every VMI data/mask value has layout + verify every VMI value has an assigned layout and every non-local lowering + choice has been serialized explicitly + verify helper ops have supported materialization paths. Current + implementation checks `ensure_layout`, `ensure_mask_layout`, and + `ensure_mask_granularity` at the layout gate, so unsupported helper + materializations fail before `vmi-to-vpto`. It also checks the first + semantic local lowering families, non-contiguous + `pto.vmi.store`, block8 + `pto.vmi.group_load`, `pto.vmi.group_slot_load`, group_slots + `pto.vmi.group_store`, group_slots `pto.vmi.group_reduce_add{f|i}`, + explicit-slots `pto.vmi.group_broadcast`, `pto.vmi.truncf`, + `pto.vmi.extf`, `pto.vmi.bitcast`, and histogram family ops at the layout gate. + +vmi-to-vpto: + use OneToN type conversion + lower only from current-op attrs/operands, operand/result layouts, and helper + ops + emit VPTO or precise unsupported diagnostic +``` + +### 1.1 Hard Constraints Versus Optimizations + +Hard legalization answers "can this program be lowered correctly?" It is +allowed to be conservative: + +```text +%w = pto.vmi.extf %a // natural layout deinterleaved=2 +%t1 = pto.vmi.mulf %w, %k1 // layout-transparent, stays deinterleaved=2 +%t1_c = pto.vmi.ensure_layout %t1 // hard store contract wants contiguous +pto.vmi.store %t1_c, %OUT1 +%w_c = pto.vmi.ensure_layout %w +pto.vmi.store %w_c, %OUT2 +``` + +This is a correct legal shape. The contiguous action is explicit at each store +use, and `vmi-to-vpto` lowers the helper with register materialization such as +`vintlv` before ordinary `vsts`. + +Optimization answers "can the same external effect be cheaper?" A fold pass +may rewrite the two store uses to consume the deinterleaved values directly: + +```text +pto.vmi.store %t1, %OUT1 // value type still says deinterleaved=2 +pto.vmi.store %w, %OUT2 +``` + +This optimized shape is legal only because `pto.vmi.store` has enough local +information to lower a `deinterleaved=2` f32 value to row-major memory, for +example with `vstsx2 INTLV_B32`. The optimization does not require +`vmi-to-vpto` to inspect `%w`'s producer or the sibling store. + +The split gives later passes room to improve layout choices: + +```text +hard pass: + guarantee legality with explicit ensure_* helpers + +optimization passes: + remove, fold, clone, or sink helpers when the optimized IR is still locally + deterministic + +vmi-to-vpto: + physicalize exactly the IR it sees, with no global planning +``` + +## 2. Files To Add Or Update + +Expected implementation files: + +```text +include/PTO/IR/VMITypes.td +include/PTO/IR/VMIOps.td +include/PTO/IR/VMIAttrs.td +lib/PTO/IR/VMI.cpp + +include/PTO/Transforms/Passes.td +lib/PTO/Transforms/PTOValidateVMIIR.cpp +lib/PTO/Transforms/VMILayoutAssignment.cpp +lib/PTO/Transforms/VMIToVPTO.cpp +small layout fact/materialization helpers under lib/PTO/Transforms + +test/lit/vmi/vmi_layout_assignment_*.pto +test/lit/vmi/vmi_to_vpto_*.pto +test/vpto/cases/vmi/*/ +``` + +Exact names may follow project conventions, but the layering should remain: + +```text +IR definitions + -> validation + -> assignment + -> OneToN lowering + -> lit and sim tests +``` + +## 3. IR Types And Attributes + +### 3.1 Layout Attribute + +Represent layout as a closed attribute family: + +```text +#pto.vmi.layout +#pto.vmi.layout +#pto.vmi.layout +#pto.vmi.layout +``` + +C++ form: + +```c++ +enum class VMILayoutKind { + Contiguous, + Deinterleaved, + GroupSlots, +}; + +struct VMILayoutKey { + VMILayoutKind kind; + int64_t deinterleaveFactor = 1; + int64_t blockElems = 1; + int64_t numGroups = 0; + int64_t slots = 0; + int64_t laneStride = 1; +}; +``` + +Verifier rules: + +```text +contiguous: + no extra parameters + +deinterleaved: + F > 1 + B > 0 + direct full-chunk lowerings require N % (F * B) == 0 + +group_slots: + G > 0 + K > 0 + G % K == 0 + K fits in one physical vreg for element type + LS > 0 +``` + +Parser compatibility during migration: + +```text +#pto.vmi.layout +``` + +is the lowering contract for group-slot values. The parser still accepts +`#pto.vmi.layout` as a legacy spelling for the pre-design +implicit group layout, but `vmi-to-vpto` support queries require explicit slots. +New `vmi-layout-assignment` output must print one of: + +```text +#pto.vmi.layout +#pto.vmi.layout +#pto.vmi.layout +``` + +so `vmi-to-vpto` can lower from the assigned type without reconstructing group +slot placement from producer or consumer context. + +`lane_stride` is counted in logical element-sized physical slots and records a +regular gap between stored group slots. It is used for carrier-style packed +stores such as `ui8` group slots lowered through b32 `PK4_B32`. + +The current implementation treats this as a group-slot property. The dense +generalization is tracked separately in +`vmi-lane-stride-generalization-implementation.md`; it requires splitting dense +lane-map stride from group-slot carrier packing before `lane_stride` can be used +on `contiguous` or `deinterleaved` layouts. + +### 3.2 VMI Types + +Surface: + +```text +!pto.vmi.vreg +!pto.vmi.mask +``` + +Layout-assigned: + +```text +!pto.vmi.vreg> +!pto.vmi.mask> +``` + +Surface VMI types are legal before assignment. Layout-assigned VMI types are +required after assignment. + +### 3.3 Explicit Lowering Carriers + +Lowering decisions are carried by the current op and its types, not by a +separate lowering-plan string. The allowed carriers are: + +```text +op attrs and operands +operand/result VMI layouts +mask granularity and mask layouts +helper ops such as ensure_layout / ensure_mask_layout +cloned or rematerialized producers +diagnostics for unsupported shapes +``` + +If assignment made a non-local choice by inspecting producers, users, sibling +users, control flow, callees, or memory context, it must rewrite the IR so that +the final choice is visible through those carriers before `vmi-to-vpto`. + +Local-decision table for the current implementation: + +```text +op local decision inputs +group_load result layout, num_groups, row_stride, source type +group_slot_load result group_slots layout and source_group_stride +group_reduce_add{f|i} source/mask/result layouts, num_groups, typed reduce semantics +group_broadcast source/result layouts and num_groups +truncf source/result layouts and element widths +dhist/chist acc/source/mask/result layouts and target capability +ensure_layout always carries source/result layouts +ensure_mask_layout always carries source/result layouts +ensure_mask_granularity always carries source/result granularities +``` + +Layout/attr-only decisions today: + +```text +load result layout plus full chunk or shaped memref proof +group_store source group_slots layout plus explicit output stride +masked_load explicit passthrough, mask layout, and memory proof +masked_store/select operand/result layouts plus mask granularity +dense extf/truncf source/result layouts and element widths +``` + +Implementation rule: + +```text +validate-assigned-vmi validates assigned layouts, mask granularity, boundaries, +and helper placement. +vmi-to-vpto emits VMI-LAYOUT-CONTRACT for missing local proof. +If a layout/attr-only op later gains a second legal lowering that cannot be +distinguished from current-op information, that lowering must be represented by a +new attr, helper op, or rematerialized op before vmi-to-vpto can emit it. +Unsupported shapes that have no explicit materialization/lowering path still +diagnose through their specific capability check rather than failing with a generic +missing-lowering +error. +``` + +Examples of forbidden recovery in `vmi-to-vpto`: + +```text +group_reduce_add{f|i} cannot walk to a load/group_load producer to choose + two-vlane parity versus block8. +group_store cannot inspect the group_reduce producer; it consumes only the + assigned source layout and explicit stride. +group_broadcast cannot inspect sibling users to decide whether to rematerialize. +masked_load cannot inspect the mask producer to prove memory safety. +func.call cannot inspect the callee body to decide physical function layout. +``` + +## 4. VMI Surface Ops Required By Cases + +Initial op set from the case catalog: + +```text +load +group_load +group_slot_load +store +masked_store + +create_mask +create_group_mask + +extf +truncf +extsi +extui +trunci +addf +addi +mulf +select +broadcast + +group_reduce_addf +group_reduce_addi +group_broadcast +group_store +dhist +chist + +ensure_layout // internal +ensure_mask_layout // internal +ensure_mask_granularity // internal +``` + +Type policy before lowering: + +```text +storage / memory boundary: + f8-like, i8, f16, i16, f32, i32 may appear as load/store element types when + the target memory instruction supports the physical width. + +cast boundary: + f8-like may appear as extf/truncf source or destination. + i8 may appear as extsi/extui/trunci source or destination. Signedness is an + op semantic, not a VMI type spelling. + Current VPTO lowering supports 32-bit integer narrowing to unsigned i8 + storage, matching the available VCVTII s32/u32 -> u8 forms; signed i8 + narrowing needs a separate target lowering. + +compute / accumulator: + floating compute baseline: f16/f32, with reassoc required for reductions + that lower through pair-wise VPTO reductions. + integer compute baseline: i32 for grouped reduction; i8/i16 storage must + first cast to i32 because integer reduction instructions widen narrow inputs. + f8/i8 are not baseline accumulator/compute types. Supporting direct 8-bit + compute requires a target capability entry and a separate lowering family. +``` + +Important semantic split: + +```text +load: + pointer sources must load full physical chunks directly. Partial logical + loads require a shaped memref proof or a future guarded/scratch fallback. + +group_load: + loads group_size data elements per group + +group_slot_load: + loads one scalar per group and produces group_slots +``` + +## 5. Layout Fact Helpers And Ensure-Based Optimization Hooks + +Do not implement a target-aware lowering-plan registry shared by assignment and +lowering. The shared contract is the IR itself: assigned VMI layouts, explicit +`ensure_layout` / `ensure_mask_layout` / `ensure_mask_granularity` helpers, +semantic op attrs/operands, and target capability diagnostics. + +Small pure helpers are still useful when they remove duplicated layout math. +They must return semantic layout facts, not VPTO instruction plans, costs, +clone decisions, or multi-user plans. + +Keep the support layer small. A query belongs in `VMILayoutSupport` only when +at least two stages need the same fact and a mismatch would create an +assignment-vs-lowering bug. Current valid shared facts are: + +```text +cast layout fact: + shared by layout assignment, layout validation, and vmi-to-vpto. + Example: f32->f8 must see deinterleaved=4 source and contiguous result in + every stage. + +group_reduce layout fact: + shared by layout assignment, layout validation, and vmi-to-vpto. + Example: S=2*VLaneElems means deinterleaved=2 source/mask and + group_slots(G, slots=8) result in every stage. + +histogram layout fact: + shared by layout assignment, layout validation, and vmi-to-vpto. + Example: dhist requires contiguous Nxui8 source, contiguous b8 mask, and + contiguous 256xui16 acc/result. chist uses the same layout fact but also + requires a target capability that classifies CHISTv2 cumulative range + semantics. + +layout materialization support: + shared by layout validation, vmi-to-vpto, and helper-based optimizations. + Example: ensure_layout from deinterleaved=2 f32 to contiguous f32 is the same + materialization whether it survives to lowering or is folded into a store. + +contiguous store support: + shared by fold-consumers and vmi-to-vpto because both must preserve the same + row-major memory effect when consuming a non-contiguous value. +``` + +Do not add a support query for a single private branch such as "this exact op +uses this exact VPTO mnemonic". Keep that branch in the lowering pattern until +another stage needs the same semantic fact. This prevents `VMILayoutSupport` +from becoming a second copy of the lowering pass. + +```c++ +struct VMICastLayoutFact { + VMICastLayoutKind kind; + VMILayoutAttr sourceLayout; + VMILayoutAttr resultLayout; + int64_t factor; +}; + +struct VMIGroupReduceLayoutFact { + VMILayoutAttr sourceLayout; + VMILayoutAttr maskLayout; + VMILayoutAttr resultLayout; + int64_t groupSize; + int64_t vlaneElems; +}; + +FailureOr +getPreferredCastLayoutFact(VMIVRegType sourceType, VMIVRegType resultType); + +FailureOr +getPreferredGroupReduceLayoutFact(VMIVRegType sourceType, int64_t numGroups); + +LogicalResult canMaterializeDataLayout(VMIVRegType sourceType, + VMIVRegType resultType, + std::string *reason); +``` + +Baseline assignment uses these helpers only to produce assigned layouts and +use-site helpers. It does not clone producers, rematerialize cheap ops, choose +memory-fused layouts by cost, or specialize private function signatures for +performance. + +Optimization passes are deliberately helper-driven: + +```text +fold-consumers: + input shape: ensure_layout feeding a layout-aware consumer. + support query: can this consumer preserve the same logical memory effect from + the source layout? + output shape: the consumer directly uses the source value. + +rematerialize: + input shape: cheap producer feeding ensure_layout / ensure_mask_layout. + support query: can the cloned producer directly create the requested type? + output shape: a cloned producer at the use. + +sink-materialization: + input shape: pure elementwise op whose operands are matching ensure_* helpers. + support query: can the result helper be materialized if it remains? + output shape: the op runs in the source layout and one helper remains on the + result. +``` + +These passes may improve multi-consumer cases without asking assignment to solve +a global cost problem. Assignment guarantees a legal baseline with helpers; +optimization removes or moves those helpers locally when the rewritten IR still +contains enough information for `vmi-to-vpto`. + +Implementation-relevant layout facts: + +```text +dense store: + requests contiguous source. If the value is assigned deinterleaved, + assignment inserts ensure_layout at the store use. A later optimization may + fold ensure_layout + store into a layout-aware VMI store. `vmi-to-vpto` + later lowers that explicit store contract. + +data/mask helper materialization: + identity conversions are always legal. + contiguous <-> deinterleaved=2/4 is legal only when source/result physical + arity and physical chunk shapes make the same logical value materializable. + unsupported conversions remain explicit diagnostics. + +group_slot_load: + assigned result layout is group_slots(G, slots=8) for packed slots or + group_slots(G, slots=1) for row-local slots. Because the result type is + `GxT`, assignment does not derive this choice from result lane count. A + constant unit `source_group_stride` selects slots=8; non-unit or dynamic + stride selects slots=1 first, then the support query rejects dynamic or + unaligned row-local lowering when the target cannot materialize it. + +block8 group_load: + assigned result layout is deinterleaved=2/4 with block_elems=8 only when the + op carries the required constant stride and memory-safety proof. + +group_store: + consumes group_slots(G,K). Explicit output stride attrs/operands decide + whether slots=8 packed or slots=1 row-local stores are legal. + +group_reduce_add{f|i}: + define E = sizeof(accumulator T), VLaneElems = 32B / E, L = 256B / E, + S = N / G. T is the accumulator/reduce element type after any required + storage cast. + S=VLaneElems uses contiguous source/mask and group_slots(G, slots=8). + S=2*VLaneElems uses deinterleaved=2 source/mask and group_slots(G, slots=8). + S=4*VLaneElems uses deinterleaved=4 source/mask and group_slots(G, slots=8). + S>=L && S%L==0 uses contiguous source/mask and group_slots(G, slots=1). + +group_broadcast: + consumes group_slots(G,K) and produces one assigned dense layout. If another + consumer wants a different dense layout, assignment inserts ensure_layout. + Optimization may clone/rematerialize group_broadcast per use. + +extf/truncf: + contiguous f16/bf16 -> deinterleaved=2 f32 + contiguous f8-like -> deinterleaved=4 f32 + deinterleaved=2 f32 -> contiguous f16 + deinterleaved=4 f32 -> contiguous f8-like + group_slots(G, slots=1) f32 -> f16 remains a slot-preserving transform. + +extsi/extui/trunci: + contiguous i8/i16 -> deinterleaved i32 according to widening factor. + deinterleaved i32 -> contiguous i8/i16 according to narrowing factor. + packed group_slots integer width-changing cast is unsupported until a + slot-wise transform is represented explicitly. + +bitcast: + per-part vbitcast is valid when source/result layouts match, physical arity + matches, and every physical chunk carries the same logical bit footprint. + This includes contiguous, deinterleaved, and identical group_slots layouts. +``` + +`vmi-layout-fold`, rematerialization, sink/hoist, and private +function specialization passes consume explicit helper IR. They may replace +helpers with cheaper equivalent IR, but they must not introduce hidden lowering +plans that `vmi-to-vpto` has to rediscover from producer/user context. + +## 6. Layout Assignment Data Model + +### 6.1 Solver State + +```c++ +struct ValueLayoutState { + Value value; + Type logicalType; + std::optional chosen; + std::optional naturalLayout; + SmallVector useRequests; +}; + +struct UseRequest { + OpOperand *operand; + VMILayoutKey requestedLayout; + Operation *requestingOp; + bool hard; +}; +``` + +### 6.2 Collection Phase + +Walk the module and collect: + +```text +1. every VMI value +2. every VMI block argument +3. every VMI function argument/result +4. every VMI op with natural producer layouts or use-site layout requests +5. every branch/yield/call/return edge carrying VMI +``` + +Build SCCs over: + +```text +dataflow uses +region yields +loop iter_args +function call graph for private/internal functions +``` + +Public/external VMI function boundaries are rejected unless +`enablePublicVMIABI` is explicitly supported. + +Block arguments are first-class layout variables. Assignment must write the +chosen layout into the block argument type or specialized function signature. +`vmi-to-vpto` must never recover a block argument layout by walking to an +incoming branch, yield, or call operand. + +### 6.3 Constraint Generation + +Examples: + +```text +truncf f32->f16: + source request deinterleaved=2, block_elems=1 + result contiguous + +group_reduce S=16: + source request deinterleaved=2, block_elems=1 + result group_slots(G, slots=8) + +group_reduce S=32: + source request deinterleaved=4, block_elems=1 + result group_slots(G, slots=8) + +group_reduce S=64: + source request contiguous + result group_slots(G, slots=1) + +group_broadcast: + source request group_slots(G,K) + result receives one assigned dense layout + incompatible dense uses are represented by ensure_layout + +ordinary dense add/mul/select: + operands/results same dense layout + +group-slot add/mul: + operands/results same group_slots(G,K) + +ordinary store: + dense source required + group_slots source is illegal + +group_store: + source request group_slots(G,K) + +dhist: + acc/result request contiguous 256xui16 + source request contiguous Nxui8 + mask request contiguous b8 + +chist: + same layout requests as dhist + diagnostic unless CHISTv2 cumulative range semantics are classified +``` + +Baseline assignment does not perform consumer-driven adoption for performance. +It records natural producer layouts and hard use-site requests. If a request +does not match the assigned layout, the pass inserts an explicit helper at that +use. + +```text +natural layout producer: + extf/truncf, group_reduce, group_slot_load, group_load, dhist/chist when the + op itself carries a layout-producing contract + +layout equality producer: + dense add/mul/select and CFG-carried values tie operands/results but do not + pick a cheaper layout by cost +``` + +Memory legality constraints: + +```text +S=32 tail fast load: + requires full_footprint_readable + otherwise require gather fallback or diagnose + +compact S=12 logical S=16: + requires compact-row gather materialization + diagnose if gather fallback is disabled/missing +``` + +### 6.3.1 Request Builders + +Implement request generation as small per-op builders. The builders produce +natural layouts, use-site requests, equality constraints, and diagnostics; they +do not choose optimization plans. + +```text +buildStoreRequests: + ordinary store -> dense contiguous request + group_store -> group_slots(G,K) request plus stride/alignment capability + checks + +buildCastRequests: + extf f16->f32 -> source contiguous, result deinterleaved=2 + extf f8->f32 -> source contiguous, result deinterleaved=4 + truncf f32->f16 -> source deinterleaved=2/block_elems=1, result contiguous + truncf f32->f8 -> source deinterleaved=4/block_elems=1, result contiguous + group_slots slots=1 f32->f16 -> explicit slot-preserving transform + group_slots slots=8 width-changing cast -> diagnostic unless a packed + transform is explicitly represented + +buildGroupReduceRequests: + derive E = sizeof(accumulator type), VLaneElems = 32B / E, + L = 256B / E, and S = logical_lanes / num_groups + S=VLaneElems -> contiguous source, group_slots(G,8) result + S=2*VLaneElems -> deinterleaved=2/block_elems=1 source, + group_slots(G,8) result + S=4*VLaneElems -> deinterleaved=4/block_elems=1 source, + group_slots(G,8) result + S>=L && S%L==0 -> contiguous source, group_slots(G,1) result + 8-bit storage must be cast to an accumulator type before this request builder + other S -> diagnostic unless an explicit fallback op/helper is enabled + +buildGroupMemoryRequests: + group_load S=16/S=32 with aligned constant stride -> natural block_elems=8 + group_load row-local full chunks -> natural contiguous + group_slot_load unit stride -> group_slots(G,8) + group_slot_load aligned row-local stride -> group_slots(G,1) + unsupported dynamic/unaligned grouped memory -> diagnostic + +buildElementwiseRequests: + dense add/mul/fma/min/max/select -> all dense operands/results share one + dense layout + group-slot add/mul/select -> all operands/results share one group_slots(G,K) + dense/group_slots mixing -> diagnostic unless an explicit group_broadcast or + group_store boundary exists + +buildMaskRequests: + mask layout follows each consuming data layout + predicate granularity follows each consuming element type + create_mask/create_group_mask produce one assigned mask layout and use + ensure_mask_layout / ensure_mask_granularity for incompatible uses + masked_store requests source layout, mask layout, and store predicate + granularity explicitly + +buildHistogramRequests: + dhist -> acc/result contiguous 256xui16, source contiguous Nxui8, + mask contiguous b8 + chist -> same layout requests, plus target capability diagnostic until + CHISTv2 high-range semantics are classified + do not create group_slots or group_reduce requests; histogram result bins are + selected by source values, not by lane/group position + +buildControlFlowRequests: + region yields, branch operands, loop iter_args, call operands, and returns + create equality requests on the carried VMI layout variable + +buildFunctionBoundaryRequests: + private/internal function argument/result layouts are materialized with + callee-entry/return-site helpers in baseline assignment; signature + specialization is an optimization pass + public/external VMI arguments/results diagnose unless enablePublicVMIABI has + a real ABI contract +``` + +Request builders must record the requesting op. Diagnostics and inserted +helpers are use-site operations, so the user can see which consumer forced a +layout. + +### 6.3.2 Optimization Producer Classes + +Baseline assignment does not use producer classes to solve conflicts. It +inserts helpers. Later optimization passes may classify producers to replace +helpers with cheaper equivalent IR. + +```text +cheap rematerializable producers: + load when address operands dominate the clone site, no intervening may-alias + write exists, and any shaped memory proof is preserved + broadcast + create_mask + create_group_mask + group_broadcast + group_slot_load when the same address/no-alias/proof conditions as load hold + and the memory access remains legal at the clone site + +layout-transparent producers: + add/sub/mul/fma/min/max/neg/abs + select + bitcast + integer bitwise and shift ops + +fixed-layout producers: + extf/truncf physical conversion layouts + group_load block-fragment layouts + group_reduce result group_slots + dhist/chist result contiguous 256xui16 and source/mask contiguous b8 contract + masked_load when the physical memory-safety proof fixes a full-read lowering +``` + +Optimization conflict policy: + +```text +cheap producer: + clone for each incompatible request when cloning does not duplicate a + side-effect, cross an aliasing write, or duplicate an illegal memory read + +layout-transparent producer: + merge into the consumer-requested equivalence class; insert materialization + only at incompatible uses + +fixed-layout producer: + use explicit helper materialization only; otherwise diagnose +``` + +These classes are not assignment constraints. They are rewrite preconditions +for passes that consume `ensure_layout` and decide whether the helper can be +folded, sunk, hoisted, or replaced by rematerialization. + +### 6.4 Solving And Rewriting + +Algorithm: + +```text +1. Collect natural layouts, use-site requests, equality constraints, and + memory-safety proofs. +2. Propagate equality constraints through SCCs. +3. Choose one deterministic assigned layout per value/equivalence class: + explicit user layout, then unique producer natural layout, then hard + non-contiguous layout, then contiguous. +4. For conflicting uses, insert ensure_layout / ensure_mask_layout / + ensure_mask_granularity at the use. +5. Emit diagnostics for unsupported semantic constraints or missing explicit + materialization/memory-safety proof. +6. Rewrite VMI result/block/function types with chosen layouts. +7. Insert helper ops with source/result layout attrs. +``` + +Rewrite invariants: + +```text +No VMI data/mask value after assignment has a null layout. +Any non-local choice is represented by op attrs, operand/result layouts, a +helper op, or an explicit diagnostic. Cloned/rematerialized producers may +appear only after later layout optimization passes. +Every ensure_* helper has an explicit supported materialization path or a +diagnostic. +Every function/call boundary carrying VMI is materialized, kept in an explicit +ABI contract, or diagnosed. +``` + +### 6.5 Rewrite Artifacts + +Assignment rewrites the IR so that later lowering has no hidden choices. + +```text +type rewrite: + every VMI data/mask result and block argument receives a layout attr + +ensure rewrite: + mismatched uses get pto.vmi.ensure_layout or ensure_mask_layout at the use + site, with source and target layouts visible in the types + +granularity rewrite: + one semantic mask used by f32 and f16 consumers gets + ensure_mask_granularity at the use site + +control-flow rewrite: + scf.if/scf.for yields and block arguments are rewritten to one agreed layout; + materialization is inserted before yield when branches differ + +function rewrite: + baseline private VMI functions get callee-entry/return-site ensure_layout; + signature specialization is an optimization pass + public/external VMI functions are diagnosed +``` + +Canonical assigned IR shape for a conflicting load: + +```text +%x = pto.vmi.load ... + : ... -> !pto.vmi.vreg<256xf32, #pto.vmi.layout> + +%x_dense = pto.vmi.ensure_layout %x + : !pto.vmi.vreg<256xf32, #pto.vmi.layout> + -> !pto.vmi.vreg<256xf32, #pto.vmi.layout> + +pto.vmi.store %x_dense, ... +``` + +Optional future optimized IR shape for a cloned load with an explicit +safe-read/execution proof: + +```text +%x_s16 = pto.vmi.load ... + : ... -> !pto.vmi.vreg<256xf32, #pto.vmi.layout> + +%x_s32 = pto.vmi.load ... + : ... -> !pto.vmi.vreg<256xf32, #pto.vmi.layout> +``` + +Canonical assigned IR shape for `group_broadcast` multi-use: + +```text +%b = pto.vmi.group_broadcast %slots + : !pto.vmi.vreg<8xf32, #pto.vmi.layout> + -> !pto.vmi.vreg<256xf32, #pto.vmi.layout> + +%b_c = pto.vmi.ensure_layout %b + : !pto.vmi.vreg<256xf32, #pto.vmi.layout> + -> !pto.vmi.vreg<256xf32, #pto.vmi.layout> +``` + +If the assigned IR does not have one of these explicit shapes, `vmi-to-vpto` +must reject it instead of attempting to recover the missing decision. + +### 6.6 Case-To-Implementation Closure Matrix + +The current case catalog is sufficient for the first implementation. No new +layout kind is justified by the supported endpoints. The implementation work +should instead close the following finite matrix. Each row names the request +builder that owns the decision, the assignment artifact that must appear in IR, +and the `vmi-to-vpto` contract. + +```text +case family builder / owner assignment artifact +3.1, 3.2, 3.3 dense casts buildCastRequests dense layout on each cast result +3.29 mask width split buildMaskRequests per-use mask granularity helper +3.31, 3.32 dense fanout conflict resolver cloned load or ensure_layout + +vmi-to-vpto contract: + consume only the assigned dense layouts. It may emit VCVT and dense + materialization, but it must not choose deinterleaved=2/4 by inspecting a + later truncf, store, or group_reduce user. +``` + +```text +case family builder / owner assignment artifact +3.4 32-bit S=8 reduce buildGroupReduceRequests one_vlane contiguous lowering +3.5 32-bit S=16 reduce buildGroupReduceRequests two_vlane parity/block8 layout +3.6 32-bit S=32 reduce buildGroupReduceRequests four_vlane dintlv4/block8 layout +3.7 32-bit S=64 reduce buildGroupReduceRequests full_chunk row_local lowering +3.11.1 S=64 active-row tail buildMaskRequests active-row store/reduce masks +3.19.1 S=16 block_elems choice buildGroupReduceRequests explicit block_elems layout +3.38 multi-tile S=32 reduce buildGroupReduceRequests multiple group_slots chunks +3.26 grouped tail buildMaskRequests split grouped masks +3.44, 3.45 grouped S=32 masks buildMaskRequests explicit deinterleaved mask values + +vmi-to-vpto contract: + lower each reduce from the current op's attrs, source/mask layout, result + group_slots layout. It must not walk to the load/group_load producer to + decide parity versus block8, row-local versus packed slots, or static versus + dynamic mask generation. +``` + +```text +case family builder / owner assignment artifact +3.56 full distribution hist buildHistogramRequests contiguous src/mask/acc/result +3.57 cumulative hist boundary buildHistogramRequests capability diagnostic or classified path + +vmi-to-vpto contract: + lower dhist from the current op and assigned layouts by carrying two physical + accumulator parts for bins 0..127 and 128..255. It must not expose the VPTO + #bin range selector on the VMI surface and must not model histogram as + group_reduce. chist remains rejected until the target records whether the + high-range cumulative result is global or range-local and, for range-local + behavior, until low-total materialization is explicit. +``` + +```text +case family builder / owner assignment artifact +3.15.1 S=16 row stride 16 buildGroupMemoryRequests block_elems=8 group_load layout +3.15.2 S=16 row stride > 16 buildGroupMemoryRequests strided block_elems=8 plan +3.16.1 group_slot_load slots=8 buildGroupMemoryRequests unit-stride packed slots plan +3.16.2 group_slot_load slots=1 buildGroupMemoryRequests row-local aligned slots plan +3.27 strided group_load buildGroupMemoryRequests positive block_elems=8 plan +3.28 slots=1 non-unit load buildGroupMemoryRequests row-local group_slot_load layout +3.37 slots=1 strided store buildStoreRequests group_store stride/alignment proof +3.39 strided load fanout conflict resolver preserving layout or materialization + +vmi-to-vpto contract: + consume only explicit memory stride/alignment attrs, current op operands, + and layouts. It must not infer safe read/write placement from neighboring + compute ops. Unsupported dynamic, unaligned, or compact-row gather shapes + stay diagnostics until a gather fallback is explicit in the current op. +``` + +```text +case family builder / owner assignment artifact +3.8 reduce->truncf->broadcast conflict resolver slot cast plus dense materialization +3.10 non-load S=32 producer buildElementwiseRequests transparent deinterleaved chain +3.17 broadcast deint consumer conflict resolver use-site group_broadcast layout +3.18 dense + reduce users conflict resolver ensure_layout; optional remat/fold +3.23 broadcast multi-user conflict resolver per-op group_broadcast layout +3.33 S=16 + S=32 users conflict resolver use-site materialization; optional cloned load +3.34 S=64 slots=1 cast buildCastRequests group_slot_cast layout +3.35 slots fanout buildElementwiseRequests same group_slots layout on users +3.36 scalar slots=8/slots=1 conflict resolver explicit slots=8/slots=1 producers +3.40 scalar dense + grouped conflict resolver ensure_layout; optional broadcast remat +3.41 incompatible fixed value conflict resolver diagnostic or ensure_layout + +vmi-to-vpto contract: + each op instance is already single-plan. The lowering pass never scans + sibling users to decide whether to clone, pack, broadcast, or materialize. +``` + +```text +case family builder / owner assignment artifact +3.21 S=32 rounded tail mask buildMaskRequests rounded vector plus mask +3.24 mask/select/store buildMaskRequests explicit mask layout/granularity +3.12 scf.if before reduce buildControlFlowRequests common yielded layout +3.20 group_slots scf.if buildControlFlowRequests common group_slots layout +3.22 scf.for carried value buildControlFlowRequests fixed-point iter_arg layout +3.25 function boundary buildFunctionBoundary specialized/internal boundary +3.42 loop accumulator buildControlFlowRequests loop-carried group_slots layout +3.43 call argument materialize buildFunctionBoundary callee-entry/return helper + +vmi-to-vpto contract: + block argument, region result, call operand, and function result layouts are + visible in types or helper ops. It must not inspect branch bodies, loop + bodies, callers, or callees to discover a layout. +``` + +```text +diagnostic family builder / owner required failure +3.7.4 slots=1 unit-stride store buildStoreRequests no aligned row-local store path +3.9 dense store of group slots buildStoreRequests use group_store/group_broadcast +3.11.2 S=32 unsafe tail buildMaskRequests missing full_footprint_readable/gather +3.13 slots=8 width cast buildCastRequests no packed slot cast transform +3.14 unsupported group size buildGroupReduceRequests no supported reduce layout/lowering +3.15.3 compact S=12 buildGroupMemoryRequests no compact gather plan +3.16.1 slots=8 non-unit load buildGroupMemoryRequests no packed strided slot load path +3.16.2 slots=1 bad stride buildGroupMemoryRequests no dynamic/unaligned row-local plan +3.19.2 invalid block_elems use conflict resolver no preserving materialization +3.25.2 public/external ABI buildFunctionBoundary no stable public VMI ABI +3.27 unaligned group_load buildGroupMemoryRequests no gather/block fallback path +3.30 masked_load unsafe tail buildMaskRequests no padding/gather fallback + +vmi-to-vpto contract: + these cases must fail before or at the layout contract boundary with the + requesting op named. They must not be accepted by falling back to a generic + dense load, dense store, or producer/user inspection. +``` + +Additional cases are needed only when the scope changes: + +```text +stable gather fallback enabled: + add compact S=12 positive lowering and masked_load unsafe-tail positive + lowering before accepting either path. + +pack-to-slots=8 or unaligned row-local stores enabled: + add positive S=64 unit-stride group_store and reduce->pack->dense store cases. + +public VMI ABI enabled: + add public call/return ABI cases before removing the public-boundary + diagnostic. + +packed group-slot width cast enabled: + add slots=8 f32->f16 cast and downstream group_store/broadcast cases. +``` + +## 7. OneToN Type Conversion + +`vmi-to-vpto` should use OneToN conversion for VMI values. + +Conversion rules: + +```text +contiguous: + ceil(N / lanesPerVReg(T)) physical vregs + +deinterleaved=F: + F * ceil((N / F) / lanesPerVReg(T)) physical vregs + ordering: part-major, then chunk + +group_slots(G,K): + ceil(G / K) physical vregs + each vreg has logical slot lanes 0..K-1 live +``` + +Mask conversion: + +```text +mask layout follows data layout +mask granularity is selected from consumer element width: + f32/i32 -> b32 + f16/i16 -> b16 + f8/i8 -> b8 +``` + +If one logical mask is used by multiple widths, assignment inserts +`ensure_mask_granularity` or rematerializes the mask producer. + +## 8. VMI-to-VPTO Pattern Rules + +Each pattern uses: + +```text +op +op attrs and operand values +operand/result layouts +adaptor physical values +``` + +Each pattern rejects: + +```text +missing current-op proof for an otherwise unsafe memory lowering +missing target capability +unexpected group_slots dense consumer +``` + +Target local lowering matrix: + +```text +load, lowering=dense_load_norm: + result layout contiguous + emits pto.vlds / pto.vsts NORM paths + covers dense store users and full-chunk row-local reduce input + +load, lowering=load_dintlv2: + result layout deinterleaved=2, block_elems=1 + emits vldsx2 DINTLV_B32 or normal load + vdintlv materialization + covers f32->f16, S=16 parity reduce, f16->f32 widened values + +load, lowering=load_dintlv4: + result layout deinterleaved=4, block_elems=1 + emits two vldsx2 DINTLV_B32 plus vdintlv + covers f32->f8, S=32 dintlv4 reduce + +group_load, lowering=s16_group_load_block8_unit_stride: + result layout deinterleaved=2, block_elems=8 + emits vldsx2/BDINTLV for 8 rows of 16xf32 + covers compact logical S=16 when source_group_stride == 16 + +group_load, lowering=s16_group_load_block8_stride: + result layout deinterleaved=2, block_elems=8 + emits two vsldb strided 32B block loads + requires source_group_stride % 8 == 0 + +group_load, lowering=s32_group_load_block8_stride: + result layout deinterleaved=4, block_elems=8 + emits four vsldb strided 32B block loads + requires source_group_stride % 8 == 0 + +group_load, lowering=group_load_contiguous_chunks: + result layout contiguous + emits one vlds per physical group chunk using row_stride address arithmetic + covers the currently implemented full-chunk row-local group_load path + +group_reduce_add{f|i}, lowering=one_vlane_reduce_contiguous: + consumes contiguous accumulator type T with group size VLaneElems(T) + produces group_slots(G, slots=8) + emits one vcgadd + +group_reduce_add{f|i}, lowering=two_vlane_reduce_deinterleaved: + consumes deinterleaved=2, block_elems=1 + produces group_slots(G, slots=8) + emits two vcgadd operations and one vadd + +group_reduce_add{f|i}, lowering=two_vlane_reduce_block8: + consumes deinterleaved=2, block_elems=8 + produces group_slots(G, slots=8) + emits two vcgadd operations and one vadd + +group_reduce_add{f|i}, lowering=four_vlane_reduce_dintlv4: + consumes deinterleaved=4, block_elems=1 + produces group_slots(G, slots=8) + emits four vcgadd operations and a vadd tree + +group_reduce_add{f|i}, lowering=four_vlane_reduce_block8_stride: + consumes deinterleaved=4, block_elems=8 + produces group_slots(G, slots=8) + emits four vcgadd operations and a vadd tree + +group_reduce_add{f|i}, lowering=full_chunk_reduce_row_local: + consumes contiguous accumulator type T with group size that is a multiple of + one physical chunk L(T) + produces group_slots(G, slots=1) + target lowering emits per-row vcgadd plus vcadd; the current prototype uses + the existing row-local VCADD/VADD/VSEL sequence while preserving the same + group_slots(G, slots=1) value contract + +dhist, lowering=full_256bin_histogram: + consumes contiguous Nxui8 source and contiguous b8 mask + consumes/produces contiguous 256xui16 accumulator/result + physical result parts are [bins 0..127, bins 128..255] + emits one low-range and one high-range histogram update for each 256-lane + source chunk + final partial source chunks require an explicit valid-lane b8 mask + +chist, lowering=capability_gated_cumulative_histogram: + uses the same layout shape as dhist + rejects until target capability classifies CHISTv2 high-range cumulative + semantics and any required low-total correction materialization is explicit + +group_slot_load, lowering=group_slot_load_slots8_unit_stride: + result group_slots(G, slots=8) + requires source_group_stride == 1 + emits one packed vsldb load + +group_slot_load, lowering=group_slot_load_slots1_row_local: + result group_slots(G, slots=1) + supports aligned non-unit source_group_stride + requires constant positive source_group_stride divisible by 256 / elementBits + emits one lane-0 vsldb per group + +group_broadcast, lowering=group_broadcast_slots8_vselr: + source group_slots(G, slots=8) + result dense layout selected per use + emits vselr using assigned result layout + +group_broadcast, lowering=group_broadcast_slots1_vselr: + source group_slots(G, slots=1) + result dense layout selected per use + emits vdup/vselr row-local materialization + +truncf, lowering=group_slot_cast_slots1_f32_to_f16: + source/result group_slots(G, slots=1) + emits one lane-0 vcvt per group slot block + rejects packed slots=8 unless slot-preserving cast support exists +``` + +The target matrix is the implementation contract. The staged status below +records how much of that contract the current prototype has already enforced. + +Current staged implementation status: + +```text +group_slot_load: + vmi-to-vpto lowers from #pto.vmi.layout + and source_group_stride. + +group_reduce_addf: + explicit slots=8 VCGADD lowering is selected from contiguous source/mask + layout, slots=8 result layout, num_groups, and reassoc. + S=16 block8 assignment emits source/mask + #pto.vmi.layout, result + #pto.vmi.layout; vmi-to-vpto lowers through two + VCGADDs plus a PAT_VL8 VADD per packed result block. + S=32 block8 assignment emits source/mask + #pto.vmi.layout, result + #pto.vmi.layout; vmi-to-vpto lowers through four + VCGADDs plus a PAT_VL8 VADD tree per packed result block. + Full-chunk row-local assignment, including S=64 and S=256 f32 cases, uses + #pto.vmi.layout and has focused + layout-assignment/vmi-to-vpto lit coverage; the explicit slots=1 generic + VCADD row-local lowering is selected locally from the current op attrs and + assigned layouts. + group_reduce_addi is implemented for i8/i16/i32 values over the registered + high-performance group-block classes. VCGADD paths preserve the logical + element type. Full-chunk row-local paths use widening VCADD intermediates + internally and bitcast the final low bits back to the declared VMI result + type; widening is not part of the VMI contract. + +group_broadcast: + explicit slots=8/1 source layouts select + packed or row-local VSELR lowerings locally. Deinterleaved block-fragment + results use the result layout block_elems as the local vselr selection group, + so + `deinterleaved = 4, block_elems = 8` broadcasts one group slot across each + 32B row fragment. VSELR index vectors are materialized per physical result + chunk. For small-group results, layout assignment has already fixed the + result layout, and vmi-to-vpto computes: + `firstGroup = first logical group covered by this result chunk`, + `sourceChunk = firstGroup / slots`, and + `baseGroupSlot = firstGroup % slots`. The generated index vector selects + `baseGroupSlot .. baseGroupSlot + groupsPerResultChunk - 1`; it must not be + reused across result chunks. + +group_load: + contiguous full-chunk path is selected from a contiguous result layout. + S=16/S=32 block-aligned strided loads are selected from + #pto.vmi.layout, and lower to one + vsldb per 32B row fragment and physical chunk. The explicit block8 support + is checked by pto-validate-vmi-layout-ir before vmi-to-vpto. + The dedicated S=16 unit-stride vldsx2/BDINTLV lowering remains a local + peephole target. + S=16/S=32 group_load with a non-constant, non-positive, or non-8-f32-aligned + row_stride is rejected by vmi-layout-assignment because the stable gather + fallback is not implemented. + +truncf group-slot cast: + layout assignment and vmi-to-vpto support group_slots(G, slots=1) + f32 -> f16 from source/result layouts and element widths. The reduce->truncf + -> group_store slots=1 flow has focused lit coverage and no longer relies on + vmi-to-vpto inspecting the truncf producer. + +group_store: + row-local group_slots(G, slots=1) lowering is implemented as one lane-0 + vsts per group for packed unit-stride output, or as one 1PT store per group + for non-unit row strides. The packed path is covered by the + reduce->truncf->group_store lit case, while the point-store path is covered + by `test/lit/vmi/vmi_to_vpto_group_store_slots1_1pt.pto`. + Packed group_slots(G, slots=8) group_store is implemented only when + num_groups is a multiple of 8 and row_stride is constant 1; it emits one + PAT_VL8 store per packed slot block. Non-unit packed group stores remain a + design target unless a strided packed-lane store lowering is made explicit. +``` + +Current implementation contract for type-generic grouped reduction: + +```text +ODS/verifiers: + pto.vmi.group_reduce_addi is the integer counterpart to group_reduce_addf. + group_reduce_addi, group_reduce_maxi, and group_reduce_mini accept + i8/i16/i32 element types when the group shape matches a registered layout + table row. + extsi/extui/trunci carry integer signedness across storage/accumulator + boundaries when an algorithm explicitly wants a wider accumulator. + +Layout assignment: + compute VLaneElems and L from the accumulator/reduce element type: + VLaneElems = 32B / sizeof(accumulator T) + L = 256B / sizeof(accumulator T) + use the same S formula for f16/f32/i8/i16/i32 once the typed reduce op and target + capability say the type is legal. + route f8 storage through extf to f32 before group_reduce_addf. + keep direct i8/i16 integer reductions in their declared logical type; + extsi/extui remains available for explicitly widened algorithms. + route integer narrowing to i8 through trunci; direct i8 compute remains + illegal unless target capability and explicit op semantics define that + lowering. + diagnose direct f8/i8 compute use with a message that points at the offending + op and suggests inserting the explicit cast when the op is meant to consume + storage data. + +Layout fact helpers: + replace f32-shaped checks with width-parametric group-reduce classifiers: + one_vlane_reduce layout fact + two_vlane_reduce_deinterleaved layout fact + four_vlane_reduce_deinterleaved layout fact + full_chunk_row_local_reduce layout fact + key legality on accumulator byte width, source/mask layout, result + group_slots layout, num_groups, and target instruction capability. + +VMI-to-VPTO: + lower group_reduce_addi through the same VCGADD/VADD skeleton used for + floating-point where the target supports the integer accumulator type. + for full-chunk i8/i16 rows, use the widening VCADD result only as an internal + partial type, combine partials at that width, then bitcast back to the + declared slots=1 VMI result type. + keep VPTO lowering local: it consumes assigned layouts and current-op + attrs/operands, but does not invent a new global layout plan. + +Tests: + cover direct i8/i16/i32 grouped reductions and explicitly widened variants. + add i32 S=8/S=16/S=32/S=64 group-reduce cases. + add f8 storage -> extf -> f32 group_reduce_addf cases. + add i8/i16 full-chunk VCADD plus bitcast cases. + retain invalid f8 and unsupported group-shape diagnostics. +``` + +Examples: + +```text +group_reduce_add{f|i}, lowering=two_vlane_reduce_deinterleaved: + consume deinterleaved=2, block_elems=1 + emit two VCGADDs and one VADD + +group_reduce_add{f|i}, lowering=two_vlane_reduce_block8: + consume deinterleaved=2, block_elems=8 + emit two VCGADDs and one VADD + +group_reduce_add{f|i}, lowering=four_vlane_reduce_dintlv4: + consume deinterleaved=4 + emit four VCGADDs and reduction tree + +group_broadcast: + consume group_slots + emit VSELR or VDUP depending slots and target dense layout + +group_slot_load slots=8: + emit one packed block load for unit stride + +group_slot_load slots=1: + emit row-local lane-0 loads for constant positive 32B-aligned strides +``` + +## 9. Validation Passes + +### 9.1 Surface Validation + +Before assignment: + +```text +VMI types may omit layout. +VPTO physical op must not consume VMI values. +Public/external VMI function ABI rejected unless enabled. +Unsupported vector-to-scalar extract rejected. +``` + +### 9.2 Layout Validation + +After assignment: + +```text +Every VMI value has layout. +Every VMI mask has layout and granularity plan. +Every lowering choice is locally deterministic or explicit in attrs/layouts. +Every ensure_* helper has a materialization path. +Every control-flow edge has matching VMI layouts. +``` + +### 9.3 `vmi-to-vpto` Context Read Audit + +`vmi-to-vpto` may still read defining ops in narrowly scoped cases that do not +select a layout or plan: + +```text +allowed: + arith.constant for the current op's scalar operands + create_mask/create_group_mask internals when lowering that mask op itself + ensure_mask_layout / ensure_mask_granularity stripping for static mask facts + memref.subview only to improve an already-failed non-identity memref + diagnostic + +not allowed: + walking from a consumer to a producer to decide a lowering + walking from a consumer to a mask producer to decide whether a lowering is legal + inspecting users to choose a result layout or materialization + recovering full_footprint_readable from surrounding MTE/caller context +``` + +Current audit result: + +```text +3.44 partial S=32 create_group_mask: + assignment writes explicit contiguous and deinterleaved mask values. When + lowering the deinterleaved create_group_mask itself, vmi-to-vpto first + materializes contiguous grouped predicate chunks and then applies predicate + pdintlv in the same tree shape as the data vdintlv. It still does not walk + from group_reduce_addf to the mask defining op to choose or reject lowering. + The dynamic active_elems_per_group form is also op-local: vmi-to-vpto lowers + contiguous chunks with vci/vshrs/vshls/vsub/vcmps, then uses the same + predicate pdintlv tree for S=32 deinterleaved masks. + +masked_load: + direct lowering is load + vsel. It does not inspect the mask producer to + choose a different load form; memory safety is provided by full physical + chunks or shaped memref proof. + +memref.subview: + mentioned only after identity lane-to-address planning fails. It is not used + to recover a hidden base/stride lowering. +``` + +## 10. Diagnostics + +Implement diagnostics with stable prefixes: + +```text +VMI-LAYOUT-CONTRACT +VMI-UNSUPPORTED-PLAN +VMI-MISSING-CAPABILITY +VMI-PUBLIC-ABI +VMI-MASK-GRANULARITY +VMI-CONTROL-FLOW-LAYOUT +``` + +Minimum diagnostic payload: + +```text +op name +logical type +actual layout +requested layout +selected/missing support path +recommended rewrite or option +``` + +Example: + +```text +VMI-LAYOUT-CONTRACT: + pto.vmi.truncf requires + #pto.vmi.layout, but the source value is + fixed to #pto.vmi.layout by the selected + strided group_load layout. Register a rematerialization or preserving + materialization path, or avoid consuming this block-loaded value with truncf. +``` + +## 11. Test And Simulator Acceptance + +Each numbered endpoint in `vmi-layout-lowering-cases.md` should become: + +```text +1. a layout-assignment lit test +2. a vmi-to-vpto lit test +3. a simulator case when the VPTO sequence is supported by the current backend +4. a diagnostic lit test when the case is explicitly unsupported +``` + +Repository locations: + +```text +test/lit/vmi/ +test/vpto/cases/vmi/ +``` + +The current repository uses descriptive flat lit names rather than +case-numbered subdirectories. New tests should follow the existing prefixes: + +```text +vmi_layout_assignment_.pto +vmi_to_vpto_.pto +/kernel.pto +``` + +The case number should still be recoverable from the coverage table in this +document and from the corresponding section in `vmi-layout-lowering-cases.md`. + +### 11.1 Layout Assignment Checks + +Each positive layout-assignment test must check: + +```text +assigned data layouts +assigned mask layouts +assigned op attrs +direct vmi-to-vpto local lowering +inserted ensure_layout/rematerialized producers +control-flow/function signature specialization +``` + +Negative tests check diagnostic text. + +### 11.2 VMI-to-VPTO Checks + +Each positive vmi-to-vpto test must check: + +```text +no pto.vmi ops remain +VPTO op sequence matches the case lowering +physical value arity and ordering are correct +mask granularity is correct +stores preserve observable logical memory order +``` + +### 11.3 Simulator Checks + +Simulator cases should compare final memory against the memory result written in +the case catalog. + +Current broad runtime sweep: + +```text +WORK_SPACE=$PWD/.tmp/vmi-runtime-batch-final CASE_PREFIX='vmi/' JOBS=4 \ + test/vpto/scripts/run_host_vpto_validation_parallel.sh + +TOTAL_CASES=47 +PASS=47 FAIL=0 +summary: .tmp/vmi-runtime-batch-final/parallel-summary.tsv +result: all summary entries are PASS +``` + +The `find: Permission denied` messages printed while discovering CANN simulator +paths are environment noise and are not treated as simulator failures. + +Required groups: + +```text +dense conversion: + 3.1, 3.2, 3.3, 3.31, 3.32 + +group reduce: + 3.4, 3.5.1, 3.5.2, 3.5.3 + 3.6.1, 3.6.2, 3.6.3 + 3.7.1, 3.7.2, 3.7.3 + 3.7.4 diagnostic + +layout/rematerialization: + 3.8, 3.10, 3.17, 3.18, 3.19.1, 3.22, 3.23, 3.31, + 3.32, 3.33, 3.34, 3.35, 3.36, 3.38, 3.40, 3.41 + +mask/tail: + 3.11.1, 3.15.1, 3.15.2, 3.21, 3.24, 3.26, 3.29, + 3.30, 3.44, 3.45 + +strided/group-slot memory: + 3.27, 3.28, 3.37, 3.39 + +function/control-flow: + 3.12, 3.20, 3.22, 3.25.1, 3.42, 3.43 + +histogram: + 3.56 positive dhist layout/lowering and simulator case when backend support + is enabled + 3.57 diagnostic chist case until CHISTv2 range semantics are classified +``` + +Aggregate catalog headings are covered through their endpoint subcases: + +```text +3.11 partial tail groups: + 3.11.1 positive S=64 active-row tail + 3.11.2 diagnostic S=32 tail without full_footprint_readable + +3.15 compact S=12 written as logical S=16: + 3.15.1 positive source row stride 16 + 3.15.2 positive source row stride greater than 16 + 3.15.3 diagnostic compact source row stride 12 + +3.16 group_slot_load layout contract: + 3.16.1 packed slots=8 positive and non-unit-stride diagnostic + 3.16.2 row-local slots=1 positive plus dynamic/unaligned diagnostics + +3.25 function boundary layout specialization: + 3.25.1 private/internal boundary lit and runtime coverage + 3.25.2 public/external boundary diagnostics +``` + +Current coverage audit result: + +```text +SIM-backed positive endpoints: + 3.1, 3.2, 3.3, 3.4, 3.5.1, 3.5.2, 3.5.3, + 3.6.1, 3.6.2, 3.6.3, 3.7.1, 3.7.2, 3.7.3, + 3.8, 3.10, 3.11.1, 3.12, 3.15.1, 3.15.2, + 3.16.1 positive, 3.16.2 positive, 3.17, 3.18, + 3.19.1, 3.20, 3.21, 3.22, 3.23, 3.24, 3.25.1, 3.26, + 3.27 positive, 3.28 positive, 3.29, 3.31, 3.32, + 3.33, 3.34, 3.35, 3.36, 3.37, 3.38, 3.39, + 3.40, 3.41, 3.42, 3.43, 3.44, 3.45 + +diagnostic endpoints: + 3.7.4, 3.9, 3.11.2, 3.13, 3.14, 3.15.3, + 3.16.1 non-unit slots=8 source stride, + 3.16.2 dynamic/unaligned slots=1 source stride, + 3.19.2, 3.25.2, 3.27 unaligned source_group_stride, + 3.30 unsafe masked_load tail + +repository evidence: + all concrete lit/runtime paths listed below exist + all 47 runtime case directories contain kernel.pto, launch.cpp, main.cpp, + golden.py, and compare.py + latest broad VMI runtime sweep passed: PASS=47 FAIL=0 + latest full VMI lit sweep passed: 350/350 + this historical sweep predates 3.56/3.57; histogram endpoints require new + lit/SIM or diagnostic tests before they can be counted as implemented +``` + +Current checked-in coverage for 3.3 dense f8->f32->compute->f8: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_f8_compute_f8.pto + +runtime SIM: + test/vpto/cases/vmi/f8-compute-f8 +``` + +Current checked-in coverage for 3.1/3.2 dense f16/f32 conversion stores: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_dense_f16_f32_store.pto + +runtime SIM: + test/vpto/cases/vmi/widen-f16-to-f32-store-reduce + test/vpto/cases/vmi/quant-f32-to-f16-tail +``` + +Current checked-in coverage for basic packed group_reduce -> group_store paths +for 3.4, 3.5.1, and 3.6.1: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_group_reduce_slots8_store.pto + test/lit/vmi/vmi_layout_assignment_group_reduce_s16_store.pto + test/lit/vmi/vmi_layout_assignment_group_reduce_s32_store.pto + +runtime SIM: + test/vpto/cases/vmi/group-reduce-basic-store +``` + +Current checked-in coverage for S=16 group broadcast continuation: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_group_slots_fanout.pto + +runtime SIM: + test/vpto/cases/vmi/group-reduce-s16-broadcast-reduce-store +``` + +Current checked-in coverage for 3.35 group_slots fanout to direct group_store +and group_broadcast: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_group_slots_fanout.pto + +runtime SIM: + test/vpto/cases/vmi/group-slots-fanout-store-broadcast +``` + +Current checked-in coverage for 3.8 `group_reduce -> group_broadcast -> +truncf -> dense store` and 3.17 `group_broadcast` feeding a +deinterleaved consumer: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_group_reduce_s16_truncf_broadcast_store.pto + +runtime SIM: + test/vpto/cases/vmi/group-reduce-s16-truncf-broadcast-store +``` + +Current checked-in coverage for 3.18 one dense value with dense and +group-reduce consumers: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_dense_group_reduce_multi_consumer.pto + +runtime SIM: + test/vpto/cases/vmi/dense-group-reduce-multi-consumer +``` + +Current checked-in coverage for 3.10 non-load producer feeding S=32 +`group_reduce`: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_non_load_s32_reduce.pto + +runtime SIM: + test/vpto/cases/vmi/group-reduce-s32-add-bias-store +``` + +Current checked-in coverage for 3.23 group_broadcast with multiple dense +consumers: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_group_broadcast_multi_consumer.pto + +runtime SIM: + test/vpto/cases/vmi/group-broadcast-multi-consumer +``` + +Current checked-in coverage for S=32 contiguous group broadcast continuation: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_group_reduce_s32_broadcast_reduce.pto + +runtime SIM: + test/vpto/cases/vmi/group-reduce-s32-broadcast-reduce-store +``` + +Current checked-in coverage for 3.21 S=32 tail with a statically safe +full-read source: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_group_reduce_s32_tail_full_tile.pto + +runtime SIM: + test/vpto/cases/vmi/group-reduce-s32-tail-full-tile-store + This case has `ptoas.flags` with `--enable-vmi`, because the partial pointer + load must run through layout assignment before VPTO/LLVM emission. +``` + +Current checked-in coverage for 3.44 masked_load grouped tail feeding S=32 +reduce: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_masked_load_group_tail_s32.pto + +runtime SIM: + test/vpto/cases/vmi/masked-load-group-tail-s32-reduce-store +``` + +Current checked-in coverage for 3.45 dynamic S=32 `create_group_mask`: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_create_group_mask_s32_dynamic.pto + +runtime SIM: + test/vpto/cases/vmi/dynamic-create-group-mask-s32-reduce-store + +runtime scalar source: + active_cols is passed as a kernel i32 scalar argument and cast to index inside + vecscope before pto.vmi.create_group_mask. This is an explicit scalar ABI, + not a value recovered by vmi-to-vpto from producer/consumer context. +``` + +Current checked-in runtime coverage for 3.12 control-flow join before S=32 +`group_reduce`: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_cf_branch.pto + test/lit/vmi/vmi_to_vpto_cf_branch.pto + +runtime SIM: + test/vpto/cases/vmi/group-reduce-s32-cf-join-store +``` + +Current checked-in runtime coverage for 3.20 `group_slots` control-flow join: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_group_slots_cf_join.pto + +runtime SIM: + test/vpto/cases/vmi/group-slots-cf-join-store +``` + +Current checked-in runtime coverage for 3.22 `scf.for` loop-carried VMI layout: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_scf_for.pto + test/lit/vmi/vmi_to_vpto_scf_for.pto + +runtime SIM: + test/vpto/cases/vmi/scf-for-loop-carried-store +``` + +Current checked-in runtime coverage for 3.42 `group_slots` `scf.for` +loop-carried accumulator: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_group_slots_scf_for.pto + +runtime SIM: + test/vpto/cases/vmi/group-slots-scf-for-store +``` + +Current checked-in coverage for 3.25.1 private function result boundary: + +```text +lit: + test/lit/vmi/vmi_ptoas_private_call_inline.pto + +runtime SIM: + test/vpto/cases/vmi/private-call-inline-store + +implementation note: + after vmi-to-vpto physicalizes the private helper, ptoas inlines private + single-block helpers whose signatures contain !pto.vreg or !pto.mask. This + happens before VPTO vecscope/backend emission, so physical vector values do + not escape through a function return. +``` + +Current checked-in coverage for 3.43 internal function argument boundary +materialization: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_call_argument_boundary.pto + test/lit/vmi/vmi_ptoas_call_boundary_vecscope.pto + +runtime SIM: + test/vpto/cases/vmi/private-call-argument-boundary-store + +implementation note: + private physical helper inlining also covers void helper calls with physical + VMI arguments, so the backend no longer sees a physical VPTO vector function + ABI for this internal boundary. +``` + +Current checked-in coverage for packed group-slot RHS elementwise continuations +for 3.5.3 and 3.6.2: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_group_slot_load_dual_layout.pto + +runtime SIM: + test/vpto/cases/vmi/group-reduce-slot-add-store +``` + +Current checked-in coverage for S=64 row-local group broadcast continuation +with aligned row_stride: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_group_reduce_s64_broadcast_reduce.pto + +runtime SIM: + test/vpto/cases/vmi/group-reduce-s64-broadcast-reduce-store +``` + +Current checked-in coverage for S=64 active-row tail with aligned row_stride: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_group_reduce_s64_tail_store.pto + +runtime SIM: + test/vpto/cases/vmi/group-reduce-s64-tail-store +``` + +The companion lit case for non-unit slots=1 point-store lowering is: + +```text +test/lit/vmi/vmi_to_vpto_group_store_slots1_1pt.pto +``` + +Current checked-in coverage for S=64 row-local group-slot RHS elementwise +continuation with aligned source_group_stride and aligned output row_stride: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_group_slot_load_dual_layout.pto + +runtime SIM: + test/vpto/cases/vmi/group-reduce-s64-slot-add-store +``` + +Current checked-in coverage for 3.34 S=64 `slots = 1` group-slot f32->f16 cast: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_group_reduce_s64_truncf.pto + +runtime SIM: + test/vpto/cases/vmi/group-reduce-s64-truncf-store +``` + +The companion negative lit cases for dynamic or unaligned `%c2` slots=1 +group_slot_load, and non-unit `slots = 8` group_slot_load, are: + +```text +test/lit/vmi/vmi_to_vpto_group_slot_load_nonunit_slots8_invalid.pto +test/lit/vmi/vmi_layout_assignment_group_slot_load_slots1_dynamic_stride_invalid.pto +test/lit/vmi/vmi_layout_assignment_group_slot_load_slots1_unaligned_stride_invalid.pto +``` + +Current checked-in coverage for the strided block-load cases: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_group_load_s16_stride_store.pto + test/lit/vmi/vmi_layout_assignment_group_load_s16_unaligned_stride_invalid.pto + test/lit/vmi/vmi_layout_assignment_group_load_s32_stride_store.pto + test/lit/vmi/vmi_layout_assignment_group_load_s32_stride_broadcast_reduce.pto + test/lit/vmi/vmi_layout_assignment_group_load_s32_unaligned_stride_invalid.pto + +runtime SIM: + test/vpto/cases/vmi/group-load-s16-stride-store + test/vpto/cases/vmi/group-load-s32-stride-store + test/vpto/cases/vmi/group-load-s32-stride-broadcast-reduce +``` + +Current checked-in coverage for grouped mask S=16 tail: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_create_group_mask_s16.pto + test/lit/vmi/vmi_create_group_mask_invalid.pto + +runtime SIM: + test/vpto/cases/vmi/group-reduce-s16-group-mask-tail-store + test/vpto/cases/vmi/group-reduce-s16-stride-group-mask-tail-store + test/vpto/cases/vmi/group-reduce-s16-group-mask-broadcast-reduce-store +``` + +Current checked-in coverage for 3.24 mask/select/masked-store semantics: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_mask_select_store.pto + +runtime SIM: + test/vpto/cases/vmi/mask-select-store +``` + +Current checked-in coverage for 3.29 one semantic mask with f32 and f16 +consumers: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_mask_granularity_f32_f16_store.pto + +runtime SIM: + test/vpto/cases/vmi/mask-granularity-f32-f16-store +``` + +Current checked-in coverage for 3.31 f16->f32 feeding dense store and S=16 +reduce: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_widen_f16_store_reduce.pto + +runtime SIM: + test/vpto/cases/vmi/widen-f16-to-f32-store-reduce +``` + +Current checked-in lit coverage for the first `vmi-layout-fold` +optimization is: + +```text +test/lit/vmi/vmi_layout_fold_store.pto +test/lit/vmi/vmi_layout_fold_masked_store.pto +test/lit/vmi/vmi_layout_fold_deint4.pto +``` + +Current checked-in lit coverage for the first `vmi-layout-rematerialize` +optimization is: + +```text +test/lit/vmi/vmi_layout_rematerialize_data.pto +test/lit/vmi/vmi_layout_rematerialize_mask.pto +``` + +Current checked-in lit coverage for the first +`vmi-layout-sink-materialization` optimization is: + +```text +test/lit/vmi/vmi_layout_sink_materialization_binary.pto // unary, binary, fma, cmp, and select data ops +test/lit/vmi/vmi_layout_sink_materialization_mask.pto +``` + +Current checked-in lit coverage for canonicalized VMI control-flow restoration is: + +```text +test/lit/vmi/vmi_legalize_arith_select.pto +test/lit/vmi/vmi_ptoas_cli_control_flow.pto +``` + +Current checked-in lit coverage for the first semantic local-lowering layout gate +is: + +```text +test/lit/vmi/vmi_layout_gate_group_slot_load_support_invalid.pto +test/lit/vmi/vmi_layout_gate_group_load_support_invalid.pto +test/lit/vmi/vmi_layout_gate_group_store_support_invalid.pto +test/lit/vmi/vmi_layout_gate_group_slots_unsupported_slots_invalid.pto +test/lit/vmi/vmi_layout_gate_store_support_invalid.pto +test/lit/vmi/vmi_layout_gate_helper_materialization_shape_invalid.pto +test/lit/vmi/vmi_layout_gate_group_reduce_support_invalid.pto +test/lit/vmi/vmi_layout_gate_group_reduce_slots1_support_invalid.pto +test/lit/vmi/vmi_layout_gate_group_broadcast_support_invalid.pto +test/lit/vmi/vmi_layout_gate_truncf_support_invalid.pto +test/lit/vmi/vmi_layout_gate_extf_support_invalid.pto +test/lit/vmi/vmi_layout_gate_bitcast_support_invalid.pto +test/lit/vmi/vmi_layout_gate_bitcast_group_slots.pto +``` + +Current checked-in direct `vmi-to-vpto` preflight coverage for bitcast local +lowering is: + +```text +test/lit/vmi/vmi_to_vpto_bitcast_footprint_invalid.pto +test/lit/vmi/vmi_to_vpto_bitcast_group_slots.pto +``` + +Current checked-in coverage for 3.32 f32 feeding f8 store and S=32 reduce: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_f32_f8_store_reduce.pto + +runtime SIM: + test/vpto/cases/vmi/f32-to-f8-store-reduce +``` + +Current checked-in coverage for multi-tile group-slot arity: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_group_reduce_s32_multitile_store.pto + +runtime SIM: + test/vpto/cases/vmi/group-reduce-s32-multitile-store +``` + +Current checked-in coverage for 3.40 scalar broadcast feeding dense and grouped +users: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_broadcast_dense_group_users.pto + +runtime SIM: + test/vpto/cases/vmi/broadcast-dense-group-users +``` + +Current checked-in coverage for 3.41 non-rematerializable `masked_load` feeding +dense and grouped users: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_masked_load_dense_group_users.pto + +runtime SIM: + test/vpto/cases/vmi/masked-load-dense-group-users +``` + +Diagnostic-only cases: + +```text +3.9 dense store of group slots +3.11.2 S=32 tail without full_footprint_readable +3.7.4 S=64 slots=1 group_store with unit output stride +3.13 packed group-slot f32 -> f16 cast +3.14 unsupported group size +3.15.3 compact source row stride 12 +3.16.1 group_slot_load slots=8 non-unit stride +3.16.2 group_slot_load slots=1 dynamic or unaligned stride +3.27 S=32 source_group_stride not divisible by 8 f32 elements +3.19.2 block_elems=8 value consumed by truncf without materialization path +3.25.2 public/external VMI boundary +3.30 unsafe masked_load tail without stable masked/gather fallback +``` + +Current checked-in diagnostic coverage for 3.9/3.13/3.14: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_dense_store_group_slots_invalid.pto + test/lit/vmi/vmi_layout_assignment_packed_group_slots_truncf_invalid.pto + test/lit/vmi/vmi_layout_assignment_group_reduce_s12_invalid.pto +``` + +Current checked-in diagnostic coverage for the remaining non-SIM diagnostic +entries: + +```text +lit: + test/lit/vmi/vmi_layout_gate_helper_support_invalid.pto + test/lit/vmi/vmi_layout_gate_helper_materialization_shape_invalid.pto + test/lit/vmi/vmi_layout_assignment_group_reduce_s32_tail_no_full_tile_invalid.pto + test/lit/vmi/vmi_layout_assignment_group_load_s16_compact_stride12_invalid.pto + test/lit/vmi/vmi_to_vpto_group_slot_load_nonunit_slots8_invalid.pto + test/lit/vmi/vmi_layout_assignment_group_slot_load_slots1_dynamic_stride_invalid.pto + test/lit/vmi/vmi_layout_assignment_group_slot_load_slots1_unaligned_stride_invalid.pto + test/lit/vmi/vmi_layout_assignment_group_load_block8_truncf_invalid.pto + test/lit/vmi/vmi_to_vpto_group_store_slots1_1pt.pto + test/lit/vmi/vmi_layout_assignment_group_load_s16_unaligned_stride_invalid.pto + test/lit/vmi/vmi_layout_assignment_group_load_s32_unaligned_stride_invalid.pto + test/lit/vmi/vmi_ptoas_public_abi_invalid.pto + test/lit/vmi/vmi_ptoas_public_result_abi_invalid.pto + test/lit/vmi/vmi_layout_assignment_external_call_invalid.pto + test/lit/vmi/vmi_layout_assignment_external_decl_invalid.pto + test/lit/vmi/vmi_to_vpto_masked_load_nonfull_invalid.pto + test/lit/vmi/vmi_to_vpto_stable_gather_masked_load_todo_invalid.pto +``` + +Capability boundaries and runtime evidence notes: + +```text +private physical function ABI: + 3.25.1 and 3.43 runtime coverage is closed for private/internal single-block + helpers by inlining private physical VMI helpers after vmi-to-vpto and before + VPTO vecscope/backend emission. Public/external VMI boundaries are still + rejected until a stable VMI ABI is defined. + +memory-proof runtime coverage: + 3.21 S=32 rounded tail-mask coverage is provided by a runtime case that loads + a full 256xf32 UB pointer vector and uses a 192-lane mask to define the active + logical rows. No surrounding MTE, caller/body context, or producer/user scan is + inspected to justify partial pointer reads. +``` + +## 12. Implementation Slices + +### Slice 1: IR Skeleton And Verifiers + +```text +layout attrs +vmi.vreg/vmi.mask types +surface op definitions +surface/layout validators +``` + +### Slice 2: Straight-Line Dense Assignment/Lowering + +```text +3.1 f16->f32->store +3.2 f32->f16->store +3.3 f8->f32->compute->f8 +``` + +### Slice 3: Group Slots And Reductions + +```text +3.4 S=8 +3.5 S=16 parity/block8 +3.6 S=32 +3.7 S=64 +group_slot_load +group_broadcast +group_store +``` + +### Slice 4: Layout Conflicts And Materialization + +```text +3.8 cast commute through group_broadcast +3.18 dense/group-reduce multi-consumer +3.19 block_elems layout selection +3.23 group_broadcast multi-consumer +3.32 f32 feeding f8 store and S=32 reduce +3.33 S=16/S=32 reduce multi-consumer rematerialization +3.34 slots=1 group-slot f32->f16 cast +3.35 group_slots fanout to group_store and group_broadcast +3.36 group_slot_load expressed as explicit slots=8/slots=1 producers +3.38 multi-tile group_slots arity +3.40 scalar broadcast materialized for dense/grouped users +3.41 non-rematerializable value with ensure_layout +``` + +### Slice 5: Masks, Tail, And Memory Legality + +```text +create_mask +create_group_mask +masked_store +safe full-read proof +compact/gather diagnostics +mask granularity per use +group_load stride greater than group size +group_slot_load slots=1 aligned non-unit stride plus dynamic/unaligned diagnostic +group_store slots=1 non-unit output stride +strided group_load feeding broadcast and a second reduce +masked_load grouped tail feeding S=32 reduce +``` + +### Slice 6: Control Flow And Functions + +```text +scf.if +scf.for +group_slots across control flow +group_slots loop-carried accumulator +internal function specialization +internal function argument boundary materialization +public ABI diagnostic +``` + +### Slice 7: Histogram + +```text +3.56 full 256-bin dhist logical op +3.57 chist semantic capability diagnostic +``` + +## 13. Completion Checklist + +Current evidence for the case-catalog objective: + +```text +1. every pre-histogram catalog endpoint is mapped in section 6.6 to an + assignment owner, assignment artifact, and vmi-to-vpto contract +2. every pre-histogram SIM-backed positive endpoint is listed in section 11.3 + and has a checked-in runtime case directory +3. every existing runtime case directory contains kernel.pto, launch.cpp, + main.cpp, golden.py, and compare.py +4. the latest historical broad VMI runtime sweep passed: PASS=47 FAIL=0 +5. the latest historical full VMI lit sweep passed: 350/350 +6. every pre-histogram unsupported endpoint listed in section 11.3 has a + diagnostic lit test +7. vmi-to-vpto decisions are represented by current-op attrs/operands, + assigned layouts, helper ops, rematerialization, or diagnostics +8. no separate lowering-plan string attr is emitted or consumed +9. release docs remain untouched; this is still a design/implementation plan + under docs/designs +10. new histogram endpoints 3.56/3.57 are mapped in section 6.6, but their + implementation evidence is intentionally pending new lit/SIM or diagnostic + tests +``` diff --git a/docs/designs/vmi-layout-assignment-lowering-design.md b/docs/designs/vmi-layout-assignment-lowering-design.md new file mode 100644 index 0000000000..005fd4a25d --- /dev/null +++ b/docs/designs/vmi-layout-assignment-lowering-design.md @@ -0,0 +1,1131 @@ +# VMI Layout Assignment And Lowering Design + +本文是新的 VMI layout assignment / lowering 设计文档。它只以 +`docs/designs/vmi-layout-lowering-cases.md` 为 source of truth，不继承早期 +VMI 草稿的 layout 设计，以避免旧上下文污染。 + +目标： + +```text +VMI surface IR + -> pto-validate-vmi-ir + -> vmi-layout-assignment // hard legalization baseline + -> canonicalize/cse + -> vmi-layout-rematerialize // optional optimization + -> canonicalize/cse + -> vmi-layout-fold // optional optimization over remat-exposed helpers + -> canonicalize/cse + -> vmi-layout-sink-materialization // optional optimization + -> canonicalize/cse + -> vmi-legalize-arith-select + -> pto-validate-vmi-layout-ir + -> layout-assigned and optimized VMI IR + -> vmi-to-vpto + -> VPTO IR +``` + +核心验收约束： + +```text +vmi-to-vpto 不允许通过上下文猜 lowering。 + +任何需要 producer/consumer/control-flow/memory/mask 上下文才能决定的事， +必须在 vmi-layout-assignment 或后续 VMI layout optimization 阶段变成显式 IR： + +1. vmi.vreg/vmi.mask 的 layout +2. current-op attrs/operands that make the local lowering deterministic +3. use-site ensure_layout / ensure_mask_layout / ensure_mask_granularity +4. rematerialized or cloned producer +5. target capability diagnostic +``` + +## 0. Hard Legalization And Optimization Boundary + +Layout assignment is a stage, not necessarily one monolithic pass. The design +separates correctness from optimization: + +```text +hard legalization: + produces legal layout-assigned VMI IR for all supported semantics + inserts conservative ensure_* helpers at incompatible uses + may choose a simple canonical layout even when a fused consumer lowering exists + must diagnose unsupported semantics before vmi-to-vpto has to guess + +layout optimization: + rewrites already legal VMI IR into cheaper but equivalent VMI IR + may fold ensure_layout into a layout-aware consumer + may clone/rematerialize cheap producers for different use-site layouts + may sink or hoist layout materialization through pure elementwise chains + may specialize private VMI function signatures +``` + +The driver currently runs MLIR's normal `canonicalize` and `cse` between these +VMI-specific passes. They are allowed to clean up trivially unused helpers, +merge identical rematerialized producers, and expose simpler use-def shapes. +They are not a source of hidden lowering information; after every optimization, +the IR must still carry enough local information for `vmi-to-vpto`. + +The baseline hard pass may emit: + +```text +%x_c = pto.vmi.ensure_layout %x : deinterleaved=2 -> contiguous +pto.vmi.store %x_c +``` + +A later optimization may replace that use with: + +```text +pto.vmi.store %x : deinterleaved=2 +``` + +only if the store op itself has a local deterministic lowering for preserving the +same row-major memory effect, such as a layout-aware `vstsx2 INTLV` lowering. +Both forms are semantically complete. The second form is an optimization, not +a hard requirement for correctness. + +## 1. Source Case Coverage + +设计必须覆盖 case catalog 中的端到端场景： + +```text +dense cast: + f16 -> f32 -> store + f32 -> f16 -> store + f8 -> f32 -> compute -> f8 + f8 -> f32 accumulator -> group_reduce_addf + i8/i16 -> signed/unsigned integer cast to i32 accumulator + -> group_reduce_addi + f8/i8 appear as cast source or cast destination at compute boundaries + integer narrowing back to i8 is an explicit cast, not implicit arithmetic + f16 -> f32 shared by dense store and S=16 reduce + f32 shared by f8 store and S=32 reduce + +group reduce: + 32-bit accumulator: S=8, S=16, S=32, S=64 + 16-bit accumulator: S=16, S=32, S=64, S=128 + 8-bit storage reduces only through an explicit accumulator cast + reduce -> group_store + reduce -> group_slot_load/elemwise -> group_store + reduce -> group_broadcast -> elemwise -> reduce -> store + one group_slots result fanning out to group_store and group_broadcast + grouped tail -> broadcast -> reduce -> store + +layout conflict: + one value with dense and group-reduce consumers + one value with S=16 and S=32 group-reduce consumers + one scalar broadcast materialized for dense and grouped users, with optional rematerialization + one non-rematerializable value materialized with use-site ensure_layout + one scalar group-slot source expressed as explicit slots=8 and slots=1 producers + S=16 block_elems=1/8 layout selection + dense consumer of group_slots diagnostic + packed group-slot width-changing cast diagnostic + S=64 slots=1 group-slot width-changing cast + +control flow: + scf.if before group_reduce + group_slots across scf.if + scf.for loop-carried layout fixed point + group_slots as scf.for loop-carried accumulator + internal function boundary specialization + internal function argument boundary materialization + public/external VMI ABI diagnostic + +mask and tail: + prefix mask + group-periodic mask + dynamic group-periodic mask + masked_load tail with explicit passthrough instead of padding + masked_load grouped tail feeding group_reduce + masked select/store + one semantic mask used by multiple predicate granularities + S=32 tail with and without full_footprint_readable + compact S=12 diagnostic + +strided memory: + group_load source stride greater than logical group size + strided group_load feeding broadcast and a second group_reduce + group_slot_load slots=1 with non-unit source stride + group_store slots=1 with non-unit output stride + +value-indexed accumulation: + full 256-bin distribution histogram over Nxui8 source lanes + VPTO low/high bin range split hidden behind one logical 256xui16 VMI result + cumulative histogram is a semantic boundary until CHISTv2 range semantics are verified +``` + +### 1.1 Case-Set Sufficiency + +The current case set is sufficient to define the first implementation of layout +assignment and lowering. It covers every decision axis that has changed the +design so far: + +```text +physical dense layout: + contiguous, deinterleaved=2/4, block_elems=1/8 + +group-slot result layout: + group_slots(G, slots=8) for packed VCG results + group_slots(G, slots=1) for row-local S=64 results + +producer-driven layout: + load, group_load, group_slot_load, broadcast, create_mask, + create_group_mask + +consumer-driven pressure: + dense store, group_reduce, group_store, group_broadcast, truncf, + elementwise/select, masked_load/masked_store + +conflict resolution: + explicit ensure_layout, explicit ensure_mask_layout, explicit diagnostics + optimization passes may later replace the helpers with rematerialization or + layout-aware consumers + +control-flow propagation: + scf.if, scf.for iter_args/results, internal/private function boundaries, + public ABI rejection + +memory legality: + full_footprint_readable proof, grouped masks, predicate granularity, aligned + strided group memory, stable gather diagnostic + +value-indexed accumulation: + histogram source/result shape, b8 source mask, and fixed low/high VPTO bin + split for a logical 256-bin result +``` + +No extra layout kind should be added unless a new case proves that the existing +layouts and explicit helper contracts cannot express the logical behavior. The remaining open +items are not missing layout semantics: + +```text +dynamic active_elems_per_group runtime source: + create_group_mask layout lowering is defined and has both lit and SIM + coverage. The supported runtime source is a kernel scalar argument cast to + index inside vecscope; vmi-to-vpto does not recover this value from GM/UB + scalar loads or surrounding context. + +private vector function runtime: + private/internal single-block helpers are runtime-covered by ptoas inlining + private physical VMI helpers after vmi-to-vpto and before VPTO vecscope/backend + emission. This is a post-physicalization backend hygiene step; vmi-to-vpto + still lowers only from assigned layouts and helper ops. + +diagnostic-only cases: + compact S=12 gather fallback, packed slots=8 width-changing cast, public VMI + ABI, unsafe masked_load tail, and unaligned/dynamic group memory remain + explicit capability boundaries. +``` + +## 2. Layout Domain + +Layout is a property of a layout-assigned VMI value, not a property inferred by +the final lowering pattern. + +Type policy: + +```text +storage boundary: + f8-like/i8/f16/i16/f32/i32 may appear in load/store values when the target + memory instruction supports the physical width. + +cast boundary: + f8-like participates through extf/truncf. + i8 participates through extsi/extui/trunci. Signedness is carried by the + cast op semantics, not by a separate layout. + On the current VPTO target, 32-bit to 8-bit integer narrowing is only a + baseline lowering for unsigned i8 results because the available VCVTII forms + are s32/u32 -> u8. + +compute boundary: + baseline floating compute uses f16/f32. + baseline integer grouped reduction compute uses i32 accumulators. i8/i16 + storage must be widened first because integer reduction instructions widen + narrow inputs. + f8/i8 are not baseline accumulator/compute element types. + +value-indexed accumulation boundary: + pto.vmi.vdhist consumes ui8 source lanes and produces a logical 256xui16 + accumulator/result. It is not a group_reduce family member because result + bins are selected by source values rather than by source lane/group position. + pto.vmi.vchist uses the same surface shape only after the target CHISTv2 + range semantics are verified. +``` + +### 2.1 Dense Layouts + +```text +#pto.vmi.layout +#pto.vmi.layout +``` + +`block_elems` defaults to `1`: + +```text +#pto.vmi.layout + == #pto.vmi.layout +``` + +Dense layouts preserve one semantic value for every logical lane. + +Lane map for `deinterleaved = F, block_elems = B`: + +```text +logical lane i +block q = i / B +in-block lane r = i % B +part p = q % F +part block t = q / F + +physical part p, physical lane t * B + r +``` + +Important consequence: + +```text +deinterleaved=2, block_elems=1 +deinterleaved=2, block_elems=8 +``` + +are different layouts. They cannot be treated as compatible because `F` is the +same. + +See `vmi-lane-stride-generalization-design.md` for the planned extension that +allows dense layouts to carry `lane_stride` as an additional lane-map field. +That extension keeps dense lane stride separate from the existing group-slot +carrier lowering use case. Non-zero lane phase is left as a future extension +and is not required for the first dense-stride optimization. + +### 2.2 Group-Slot Layouts + +```text +#pto.vmi.layout +#pto.vmi.layout +``` + +Only `G` lanes have semantic values: + +```text +slot_block(g) = g / K +slot_lane(g) = (g % K) * LS +``` + +All non-slot lanes are undefined and may only be read by group-aware operations. +Ordinary dense `add/mul/store/truncf` cannot consume `group_slots`. + +`LS` defaults to 1 and is measured in logical element-sized physical slots. It +is not a new group semantic; it records regular physical spacing for each stored +group slot. For example, `ui8 lane_stride=4` maps slot values to byte lanes +0, 4, 8, ... and lets `group_store` lower through a b32 carrier `PK4_B32` +store. + +`K` is selected by the assigned producer/result contract: + +```text +S=8/16/32 packed VCG result -> slots=8 +S=64 row-local result -> slots=1 +``` + +Histogram does not add a layout family. A full logical histogram result uses: + +```text +!pto.vmi.vreg<256xui16, #pto.vmi.layout> +``` + +and physicalizes to two ordered VPTO parts: + +```text +part0 = logical bins 0..127 +part1 = logical bins 128..255 +``` + +The VPTO `#bin` selector is therefore an op-local lowering detail, not a VMI +layout attribute and not a user-visible operand on `pto.vmi.vdhist`. + +## 3. Lowering Context Must Become Explicit IR Output + +`vmi-to-vpto` may inspect only: + +```text +1. op name and explicit op attrs +2. converted operand/result types with layout +3. helper/materialization ops written by layout assignment +4. inserted helper ops +5. target capability registry +``` + +It must not: + +```text +1. walk to defining op to infer layout +2. inspect all users to choose a lowering path +3. infer memory legality from a later mask +4. decide S=16 block_elems=1 vs block_elems=8 locally +5. decide whether group_broadcast should be materialized for one or many users +6. specialize function signatures during vmi-to-vpto +``` + +Any of those decisions belongs to the layout stage before `vmi-to-vpto`. + +## 4. Explicit Assignment Products + +After `vmi-layout-assignment`, every VMI data and mask value must be in one of +these states: + +```text +layout-assigned type: + !pto.vmi.vreg> + !pto.vmi.mask> + +or explicit helper: + pto.vmi.ensure_layout + pto.vmi.ensure_mask_layout + pto.vmi.ensure_mask_granularity +``` + +`vmi-to-vpto` is allowed to choose a deterministic lowering from local +information on the current op: + +```text +current op name +current op attrs +operand/result types and layouts +current op operand values such as stride and offset +target capability and pass options +``` + +This is not context inference. What remains forbidden is walking to producers, +users, sibling users, branch/loop bodies, callees/callers, or nearby memory/MTE +ops to recover a lowering decision or a memory-safety proof. + +If a decision cannot be made from that local information, layout assignment +must rewrite the IR until the decision is explicit in attrs, operand/result +layouts, helper ops, or diagnostics. Later optimization passes may replace +helpers with cloned/rematerialized producers, but `vmi-to-vpto` must not +consume a separate string lowering-plan attr. + +### 4.1 Local Lowering Contract + +The lowering path is derived from op + assigned operand/result layouts + +explicit attrs/operands. If two legal lowerings cannot be distinguished from +that local information, the IR is missing a semantic carrier and must be +extended before that lowering is implemented. + +The shared abstraction is a layout fact classifier, not a central lowering-plan +registry. A classifier may answer questions such as: + +```text +cast layout fact: + f16/i16 -> f32/i32 requires contiguous source and deinterleaved=2 result + f8/i8 -> f32/i32 requires contiguous source and deinterleaved=4 result + f32/i32 -> f16/i16 requires deinterleaved=2 source and contiguous result + f32/i32 -> f8/i8 requires deinterleaved=4 source and contiguous result + +group_reduce layout fact: + define E = sizeof(accumulator T), VLaneElems = 32B / E, + L = 256B / E, S = N / G. + S == VLaneElems requires contiguous source/mask and + group_slots(G, slots=8) result. + S == 2 * VLaneElems requires deinterleaved=2 source/mask and + group_slots(G, slots=8) result. + S == 4 * VLaneElems requires deinterleaved=4 source/mask and + group_slots(G, slots=8) result. + S >= L && S % L == 0 requires contiguous source/mask and + group_slots(G, slots=1) result. + +memory safety fact: + full physical chunks are legal for pointer sources. Partial logical loads + need a shaped safe-tail memref proof or an explicit fallback option. +``` + +These helpers return semantic layout requirements and capability diagnostics. +They do not return VPTO instruction names, cost decisions, clone decisions, or +multi-user plans. + +The useful shared fact is the part that would otherwise be recomputed by two or +more stages and must stay identical for correctness: + +```text +cast width ratio: + assignment uses it to request source/result layouts and insert ensure_layout. + validation uses it to reject unsupported assigned cast shapes. + lowering uses it to check the local op shape before emitting VPTO. + +group_reduce lane partition: + assignment uses N/G and accumulator element width to request source/mask and + result layouts. + validation uses the same math to reject legacy or incomplete group_slots. + lowering uses the already assigned layouts to select the local VPTO sequence. + +layout materialization shape: + assignment may insert ensure_layout without proving every physical sequence. + validation and lowering use one support query to decide whether that explicit + helper is materializable on the target. + optimization uses the same query only when it wants to fold/sink/remove an + explicit helper. +``` + +The helper is not useful when it only renames one local pattern. A single +`if (is this op with this attr)` that is not shared by assignment, validation, +lowering, or an optimization should stay local to that pass. The support layer +exists to prevent divergent layout math, not to move every branch into a table. + +Forbidden non-local lowering recovery: + +```text +No pattern may recover a lowering decision or memory proof by: + - walking from group_reduce to the load/group_load producer + - walking from store/broadcast/truncf to the group_reduce producer + - scanning sibling users of a group_slots value + - inspecting branch bodies or loop bodies from a control-flow boundary + - inspecting private callee bodies while lowering a call +``` + +If the current op lacks enough local information, `vmi-to-vpto` emits +`VMI-LAYOUT-CONTRACT` at the current op and prints the op name, logical type, +assigned layouts, and the missing decision class. + +## 5. Layout Requests, Helpers, And Optimization + +The compiler must not carry a target-aware lowering-plan registry as the shared +contract between assignment, optimization, validation, and lowering. The +shared contract is: + +```text +1. assigned layouts on VMI types +2. explicit use-site helpers: ensure_layout, ensure_mask_layout, + ensure_mask_granularity +3. explicit op attrs/operands that are part of the semantic op +4. small layout fact classifiers shared only where they remove duplicated + layout math +5. target capability diagnostics +``` + +This split makes optimization simpler only when optimization is phrased as +rewriting explicit helper IR: + +```text +baseline: + %x_d2 = pto.vmi.extf %x_f16 + %a = pto.vmi.addf %x_d2, %k_d2 + %a_c = pto.vmi.ensure_layout %a : deinterleaved=2 -> contiguous + pto.vmi.store %a_c, %out0 + %x_c = pto.vmi.ensure_layout %x_d2 : deinterleaved=2 -> contiguous + pto.vmi.store %x_c, %out1 + +fold-consumers: + checks only each local ensure_layout + store use. + If VMILayoutSupport says the store can preserve row-major memory from the + source layout, rewrite that use to store the source directly. + It does not inspect sibling users of %x_d2 and does not recompute the layout + assignment. + +rematerialize: + checks only cheap producer + ensure_layout. + If the producer can directly create the requested layout, clone/rematerialize + that producer for the use. + Memory producers such as group_slot_load are excluded until a separate proof + says cloning is semantically and economically valid. + +sink-materialization: + checks only explicit ensure_* operands of a layout-transparent op. + If every operand helper is compatible, rebuild the op in the source layout and + leave one ensure_* on the result. +``` + +If an optimization needs a global cost decision, it should produce a new +explicit IR shape and then rely on canonicalize/CSE. It must not communicate a +private decision to `vmi-to-vpto`. + +### 5.1 Baseline Dense Layout Requests + +```text +f16 -> f32: + source contiguous f16 + result deinterleaved=2, block_elems=1 + +f8 -> f32: + source contiguous f8 + result deinterleaved=4, block_elems=1 + +f32 -> f16: + source deinterleaved=2, block_elems=1 + result contiguous f16 + +f32 -> f8: + source deinterleaved=4, block_elems=1 + result contiguous f8 + +elementwise dense: + all dense operands/results share the same layout + +dense store: + requests contiguous source + if the stored value is assigned deinterleaved, baseline assignment inserts + ensure_layout at the store use + +two-way interleaved memory ops: + `pto.vmi.deinterleave_load` produces two dense logical streams and requests + contiguous layouts for both results + `pto.vmi.interleave_store` consumes two dense logical streams and requests + contiguous layouts for both inputs + the deinterleave/interleave memory pattern is op semantics, not a VMI layout +``` + +### 5.2 Baseline Group Layout Requests + +```text +group_reduce_add{f|i}: + uses the group_reduce layout fact in section 4.1. + The source and mask operands request the computed dense layout. + The result is assigned group_slots(G, slots=8) or group_slots(G, slots=1). + Floating-point `group_reduce_addf` carries `reassoc`; integer + `group_reduce_addi` does not. + +group_slot_load: + result group_slots(G, slots=8) for packed slots + result group_slots(G, slots=1) for row-local slots + +group_broadcast: + source requests group_slots(G,K) + result requests one dense layout + incompatible dense consumers are represented by ensure_layout after the + broadcast result; a later optimization may clone/rematerialize the broadcast + +group_store: + source requests group_slots(G,K) + explicit output stride attrs/operands decide store legality + +group_slot_cast f32 -> f16: + slots=1 row-local source/result is legal + slots=8 packed source is illegal unless a future explicit helper or semantic + op defines the packed slot-preserving transform +``` + +### 5.3 Tail And Memory Safety + +Mask semantics and memory legality are separate: + +```text +mask: + decides which logical lanes participate in compute/store semantics + +full_footprint_readable: + decides whether a rounded-up physical load is allowed to read inactive lanes +``` + +The full-tile-readable proof must be explicit. It may be carried by a +statically shaped memref source. Pointer-source runtime kernels should load a +rounded physical vector and use a mask to express logical active lanes. +`vmi-to-vpto` consumes only the op/type-local proof carrier; it does not inspect +surrounding MTE copies, producer bodies, callers, or later consumers to decide +whether inactive physical lanes are safe to read. + +Example: + +```text +S=32 tail num_groups=6: + without full_footprint_readable: + fast DINTLV_B32 full-tile load is illegal + + with full_footprint_readable: + full 8-row physical tile may be loaded + compute mask is PAT_VL48 per physical part + group store mask is PAT_VL6 + +S=16 grouped tail active_elems_per_group=12: + low 8-lane row half uses PAT_ALL + high 8-lane row half uses lane_mod_8 < 4 + the same split applies before and after group_broadcast + +one mask used by f32 and f16 consumers: + f32 use materializes a b32 predicate + f16 use materializes a b16 predicate + vmi-to-vpto consumes the assigned per-use mask materialization +``` + +### 5.4 Case-Driven Request Matrix + +The first implementation should build requests from the following finite table. +This table is deliberately case-derived; adding a new request kind requires a +new catalog case or a proof that it is equivalent to one listed here. + +```text +dense store: + requests dense contiguous source + if source is deinterleaved, baseline assignment inserts ensure_layout at the + store use. A later optimization may fold that helper into a layout-aware + store lowering such as vstsx2. + +truncf f32 -> f16: + requests source deinterleaved=2, block_elems=1 + requests result contiguous f16 + +truncf f32 -> f8: + requests source deinterleaved=4, block_elems=1 + requests result contiguous f8 + +group_reduce_add{f|i}: + computes E = sizeof(accumulator type), VLaneElems = 32B / E, + L = 256B / E, and S = logical_lanes / num_groups + S=VLaneElems requests source contiguous and result group_slots(G, slots=8) + S=2*VLaneElems requests source deinterleaved=2 and result + group_slots(G, slots=8) + S=4*VLaneElems requests source deinterleaved=4 and result + group_slots(G, slots=8) + S>=L && S%L==0 requests source contiguous and result + group_slots(G, slots=1) + 8-bit storage reaches this request only after an explicit cast to the + accumulator type + +group_broadcast: + requests source group_slots(num_groups, slots=K) + produces one assigned dense result layout + incompatible dense consumers are represented by ensure_layout uses; a later + optimization may clone/rematerialize the group_broadcast per consumer + +group_store: + requests source group_slots(num_groups, slots=K) + explicit output stride attrs/operands decide store legality + +dense elementwise add/mul/fma/min/max/select: + requests all dense data operands and results use one dense layout + mask operands request the same data layout and the consumer element + granularity + +group-slot elementwise add/mul/select: + requests all group-slot operands and results use the same + group_slots(num_groups, slots=K) + rejects mixing dense and group_slots without explicit group_broadcast or + group_store + +group_slot_load: + requests result group_slots(num_groups, slots=8) for packed unit-stride slots + requests result group_slots(num_groups, slots=1) for row-local aligned slots + +group_load: + requests result deinterleaved=2/4, block_elems=8 for S=16/S=32 block + fragments, or contiguous for row-local full chunks + +masked_load: + requests result layout from its consumers + requests mask layout matching the result + requires explicit passthrough; padding is not synthesized + +masked_store: + requests dense source layout required by the store op + requests mask layout matching the source layout and store element granularity + does not choose memory safety for an earlier load + +create_mask/create_group_mask: + produces one assigned mask layout and granularity + incompatible mask consumers are represented by ensure_mask_layout or + ensure_mask_granularity; optimization may clone/rematerialize the mask op + +dhist: + requests acc/result contiguous !pto.vmi.vreg<256xui16> + requests source contiguous !pto.vmi.vreg + requests mask contiguous with b8 granularity + lowers each 256-lane source chunk by carrying two accumulator parts: + bins 0..127 use VPTO histogram #bin=0, bins 128..255 use #bin=1 + final partial source chunks are represented by AND-ing the user mask with a + valid-lane prefix mask before the VPTO histogram op + +chist: + same layout requests as dhist + baseline lowering is disabled until target capability records whether the + high-range VPTO cumulative result is global or range-local + +scf.if/scf.for/call/return: + requests equality across carried VMI values, yielded values, call operands, + callee arguments, and function results + baseline private/internal functions materialize at boundaries; optimization + may specialize signatures + public/external VMI boundaries are diagnostics until an ABI is defined +``` + +Important negative requests: + +```text +ordinary dense add/mul/store/truncf cannot request group_slots +packed group_slots(slots=8) cannot request width-changing cast unless a packed +slot-preserving cast transform is explicitly represented +slots=1 group_store cannot request unit-stride row-major output until a pack or +unaligned-store transform is explicitly represented +``` + +### 5.5 Optimization Hooks + +Baseline assignment resolves incompatible use-site requests by keeping one +assigned layout on the value and inserting explicit helpers at the use sites +that need another layout. It does not clone producers, rematerialize cheap +ops, choose memory-fused layouts by cost, or specialize private function +signatures for performance. + +Those choices belong to later VMI layout optimization passes. They consume +the explicit helper IR and may rewrite it when the rewrite preserves the same +logical value and externally visible memory effect: + +```text +ensure_layout + store: + fold into a layout-aware store if the store can directly consume the source + layout and still write row-major memory + +producer + ensure_layout: + clone/rematerialize the producer for that use only when the producer is cheap + or has an explicit safe-read proof + +elementwise chain + ensure_layout: + sink or hoist materialization through pure layout-transparent ops + +group_broadcast + incompatible dense consumers: + type each group_broadcast op for its consumer layout; do not force one result + layout across independent group_broadcast users + +create_mask/create_group_mask + incompatible mask consumers: + clone/rematerialize the mask producer per layout or predicate granularity + +private function boundary: + specialize function signatures only in an optimization pass; baseline + assignment materializes at boundary uses +``` + +If no helper materialization or optimization rewrite is legal, the diagnostic +must name the value's assigned layout, the use-site requested layout, and the +op that requested it. + +## 6. Layout Assignment Algorithm + +`vmi-layout-assignment` is module-level. It must see function/call/control-flow +connections before choosing layouts. + +### 6.1 Variables + +Create a layout variable for: + +```text +1. every VMI OpResult +2. every VMI BlockArgument +3. every function argument/result that is allowed to carry VMI +4. every VMI mask value +``` + +Create a use-site request for: + +```text +1. every operand use that requires a specific layout +2. every control-flow yield/branch/call/return edge +3. every memory operation that requires an explicit memory legality proof +``` + +### 6.2 Constraints + +Hard constraints: + +```text +group_slots cannot feed ordinary dense consumers +direct group-slot width-changing cast requires an explicit slot-preserving transform +public/external VMI function boundary requires a stable ABI or diagnostic +S=32 fast tail load requires full_footprint_readable or gather fallback +``` + +`slots = 1` row-local cast may satisfy the slot-preserving transform requirement. +Packed `slots = 8` f32->f16 remains a diagnostic unless a separate packed cast +or unpack/materialization transform is represented explicitly. + +Equivalence constraints: + +```text +dense add/mul/select: + operands/results use same dense layout unless an explicit materialization is + inserted at a use site + +scf.if/scf.for: + region yield operands and block arguments must have the same assigned layout + as the region result/iter_arg +``` + +Canonical baseline constraints: + +```text +S=16 group_reduce: + request deinterleaved=2; baseline uses block_elems=1 unless the producer + result already carries block_elems=8 as an explicit layout + +one dense value feeding S=16 and S=32 group_reduce: + keep the value's assigned layout and insert ensure_layout at both use sites + that need deinterleaved=2 or deinterleaved=4 + +load/group_load: + use the op's assigned result layout and explicit memory-safety attrs only + +group_broadcast: + keep one assigned dense result layout and communicate other dense use layouts + through ensure_layout +``` + +### 6.3 Solving + +Recommended solving order: + +```text +1. Build function/control-flow SCCs. +2. Collect natural producer layouts and hard use-site layout requests. +3. Propagate equality constraints through dense elementwise ops and CFG edges. +4. Choose one deterministic assigned layout for each value or equivalence + class. +5. Insert ensure_layout / ensure_mask_layout / ensure_mask_granularity at uses + whose requested layout differs from the assigned layout. +6. Emit diagnostics for unsupported semantic constraints or missing explicit + memory-safety proofs. +7. Rewrite VMI types and insert explicit helper ops. +``` + +Tie-breaking must be deterministic and deliberately simple. Suggested priority: + +```text +1. Preserve an explicit user-provided layout attr. +2. Preserve a unique producer natural layout when present. +3. Preserve an equality-class non-contiguous layout when required by a hard op. +4. Otherwise choose contiguous. +``` + +## 7. Control Flow And Functions + +### 7.1 `scf.if` + +All branch yields for one result must agree on one assigned layout. If they do +not, assignment inserts materialization before `scf.yield` where possible. +The `scf.if` result type after assignment carries that layout, so +`vmi-to-vpto` does not need to inspect either branch body. + +### 7.2 `scf.for` + +Loop-carried VMI values are fixed-point variables: + +```text +initial iter_arg layout +body block argument layout +yield operand layout +loop result layout +``` + +must converge to one layout. If a body consumer needs another layout, it is a +use-site request inside the loop body. +The loop body block argument has no defining op. Its layout is therefore part +of the block argument type after assignment, not information reconstructed from +the initial value or previous iteration during lowering. + +### 7.3 Calls + +Internal/private VMI function boundaries must make layout choices explicit in +the assigned IR. The baseline implementation keeps function arguments in a +contiguous VMI ABI and inserts callee-entry `ensure_layout` helpers when the +callee body needs another layout. Private helpers are then physicalized by +`vmi-to-vpto` and inlined before VPTO vecscope/backend emission so physical +`!pto.vreg`/`!pto.mask` values do not become a backend function ABI. A later +private-function optimization may specialize signatures directly: + +```text +func @producer() -> !vmi.vreg<256xf32, deinterleaved=4> +``` + +then physicalized by `vmi-to-vpto` into multiple VPTO function results. + +Public/external VMI function boundaries are rejected until a stable VMI ABI is +defined. + +## 8. vmi-to-vpto Contract + +`vmi-to-vpto` receives layout-assigned VMI. It performs no global reasoning. + +For each op, the pattern: + +```text +1. reads operand/result layouts +2. reads current op attrs and operand values +3. asks TypeConverter for ordered physical values +4. emits the locally implied VPTO lowering +5. fails if target capability or required local proof is absent +``` + +The pattern must not: + +```text +1. inspect all users to decide result layout +2. inspect defining ops to decide source layout +3. choose between S=16 block_elems=1 and block_elems=8 +4. decide whether a load is full_footprint_readable +5. decide function signature specialization +``` + +Allowed local reads are deliberately narrower: + +```text +arith.constant defining op: + allowed only to materialize an operand of the current op, such as + create_mask active_lanes or a constant memory offset + +current VMI op body/attrs: + allowed for op-local semantics, such as create_group_mask + active_elems_per_group when lowering the create_group_mask op itself + +helper materialization chain: + allowed only to strip ensure_mask_layout / ensure_mask_granularity for + static predicate analysis that does not choose a different layout or lowering + +diagnostic embellishment: + allowed only to improve an already-failed capability message, such as naming + memref.subview after identity lane-to-address planning has failed +``` + +Anything else is a layout-assignment responsibility. In particular, an +unsupported producer/consumer combination must be rejected before assignment +emits layout-assigned IR. Section 3.44 is the model for supported partial S=32 +grouped masks: assignment emits explicit contiguous and deinterleaved mask +values, and `vmi-to-vpto` lowers the deinterleaved mask op itself through +contiguous grouped-mask materialization followed by predicate deinterleave. It +does not walk from `group_reduce_addf` to the mask producer to choose or reject +the lowering. Dynamic `active_elems_per_group` follows the same rule: the +`create_group_mask` op lowers its own SSA scalar with vci/vshrs/vshls/vsub/vcmps +for contiguous chunks before any predicate deinterleave. + +## 9. Physical Value Ordering + +The OneToN lowering order is fixed. + +```text +contiguous: + chunk0, chunk1, ... + +deinterleaved=F: + part0_chunk0, part0_chunk1, ..., + part1_chunk0, part1_chunk1, ..., + ... + part(F-1)_chunk0, ... + +group_slots(G,K): + slot_block0, slot_block1, ... +``` + +Two physical bundle entries may alias the same VPTO SSA value when the current +op semantics prove they have the same contents, such as group_broadcast feeding both +parts of a `deinterleaved=2` broadcast result. Arity still follows the layout; +aliasing is not a different layout. + +## 10. Diagnostics + +Diagnostics are part of the design. They must name: + +```text +1. the VMI op +2. source logical type +3. assigned source layout +4. requested layout +5. missing local proof or disabled fallback +6. suggested rewrite when available +``` + +Examples: + +```text +dense store of group_slots: + use group_store, group_broadcast, or explicit group-pack + +packed group-slot f32->f16: + group_broadcast before truncf, or keep group_store as f32 + +S=32 tail without full_footprint_readable: + mark source full_footprint_readable or enable stable gather fallback + +S=32 group_load with unaligned source_group_stride: + choose a stride divisible by 8 f32 elements or enable stable gather fallback + +public VMI function boundary: + make function internal, inline before assignment, or define ABI layout +``` + +## 11. Implementation Migration Checks + +The design is useful only if the implementation removes duplicated decision +points instead of renaming them. The migration target is: + +```text +assignment: + computes assigned layouts, records use-site requests, inserts ensure_* helpers, + and diagnoses unsupported semantics + does not clone/rematerialize producers + does not choose memory-fused layouts by cost + does not inspect sibling users to optimize a value + +layout optimization: + consumes explicit ensure_* helpers + may fold ensure_layout into layout-aware consumers + may clone/rematerialize cheap producers + may sink/hoist materialization through pure elementwise chains + may specialize private function signatures + +vmi-to-vpto: + consumes current op attrs/operands, assigned operand/result layouts, and + explicit helper ops + performs local physical shape and target-capability checks + does not recover layout plans from producers, sibling users, CFG regions, or + callees/callers +``` + +Concrete implementation debt to remove: + +```text +1. Move assignment-side data/mask rematerialization into + vmi-layout-rematerialize. Baseline assignment should insert ensure_* for + mismatched uses. +2. Keep `VMILayoutSupport` as target capability and layout-shape queries, not + as a shared plan table. Group-reduce layout math now lives in + `getPreferredGroupReduceLayoutFact`. Dense cast layout shape now lives in + `getPreferredCastLayoutFact`. Helper materialization gates use + `canMaterializeDataLayout`, `canMaterializeMaskLayout`, and + `canMaterializeMaskGranularity`. +3. Assignment, validation, and lowering may call layout fact helpers, but must + not each independently derive VLaneElems/groupSize/factor/slots rules. +4. Keep store-fold, rematerialization, and sink/hoist as local rewrites over + explicit ensure_* IR. They must not walk sibling users to rediscover why the + helper exists. +5. Update pass descriptions, diagnostics, and tests so "assignment only" output + is legal with helpers, and optimized output is a separate, equivalent IR + form. +``` + +Regression tests should prove the boundary: + +```text +assignment only: + multi-consumer values keep one assigned layout and use ensure_* at mismatched + uses + +fold-consumers: + ensure_layout + store becomes a layout-aware store only when the consumer can + preserve the same row-major memory effect + +rematerialize: + cheap producer + ensure_layout becomes a cloned/rematerialized producer; with + the pass disabled, the ensure_layout form remains legal + +vmi-to-vpto: + rejects any residual need for producer/user context with VMI-LAYOUT-CONTRACT +``` + +## 12. Design Completion Criteria + +The design is complete only when: + +```text +1. every case in vmi-layout-lowering-cases.md maps to assignment requests, + explicit helpers, or a precise diagnostic +2. every VMI-to-VPTO lowering can be emitted without looking at producer/user + context +3. every unsupported case has a precise capability diagnostic +4. every control-flow/function boundary materializes, specializes in an + optimization pass, or diagnoses +5. every mask has explicit data layout and predicate granularity +6. every positive case has end-to-end lit coverage +7. every simulator-supported positive case has simulator validation +``` diff --git a/docs/designs/vmi-layout-lowering-cases.md b/docs/designs/vmi-layout-lowering-cases.md new file mode 100644 index 0000000000..0ee45da2da --- /dev/null +++ b/docs/designs/vmi-layout-lowering-cases.md @@ -0,0 +1,6418 @@ +# VMI Layout Lowering Cases + +本文是 VMI layout/lowering 的典型 case catalog，不是完整设计总文档。它只回答一个问题： +一个 VMI logical vector 在某个场景下选择某种 layout 后，`vmi-to-vpto` 必须生成什么 +VPTO 结果。这里不写动机式描述；每个场景都给出 layout assignment 和 lowering result。 + +## 1. Layout Families + +### 1.1 Dense Layout + +Dense layout 的每个 logical lane 都有语义值。 + +```text +#pto.vmi.layout +``` + +Physical ordering: + +```text +chunk c, lane l -> logical lane c * L + l +``` + +`L` is the physical lanes per 256B VPTO vector register for the element type. + +```text +#pto.vmi.layout +``` + +`block_elems` defaults to `1`. Existing spellings are shorthands: + +```text +#pto.vmi.layout + == #pto.vmi.layout + +#pto.vmi.layout + == #pto.vmi.layout +``` + +Logical-to-physical mapping: + +```text +logical lane i +block q = i / B +in_block lane r = i % B +part p = q % F +part_block t = q / F + +physical part p, physical lane t * B + r +``` + +Required invariants: + +```text +F > 0 +B > 0 +N % (F * B) == 0 for the direct full-chunk paths in this document +``` + +### 1.2 Group-Slot Layout + +Group-slot layout is not dense. Only `G` lanes have semantic values. + +```text +#pto.vmi.layout +#pto.vmi.layout +``` + +Physical slot mapping: + +```text +N = logical lane count +S = N / G // logical lanes per source group + +slot_block(g) = g / K +slot_lane(g) = (g % K) * LS +``` + +Required invariants: + +```text +G > 0 +K > 0 +G % K == 0 +K must fit in the physical vreg element count +LS > 0 +``` + +`LS` defaults to 1 and is counted in logical element-sized physical slots. It +is used when the group result value is intentionally stored with a regular lane +gap. For example, `ui8 lane_stride=4` places group slots in byte positions 0, +4, 8, ... and can be lowered to a b32 carrier plus `PK4_B32` store. + +`K` is selected by the producer/consumer layout support rule. It is not always 8. For +`VCGADD`-packed results, `K = 8` matches the eight 32B block results written to +the low lanes of one destination vreg. For row-local reductions where each +logical group already occupies one full 256B vreg, `K = 1` keeps each group's +scalar result in lane 0 of its own physical vreg and avoids an unsupported +cross-vreg scalar pack. + +Only these lanes are semantic: + +```text +physical slot block slot_block(g), lane slot_lane(g) +``` + +All other lanes are undefined for ordinary VMI consumers. They may only be read +by group-aware ops that define how to interpret group slots. + +## 2. Layout Support Selection Rules + +VMI cast ops must not hard-code one physical `vcvt` lowering as their semantic +layout rule. Layout assignment records the required value layout; target +support queries only answer whether that layout can be materialized or lowered. + +```text +dense cast: + source/result are dense layouts. + lowering may require deinterleaved(F, block_elems=1) around VCVT. + +group-slot cast: + source/result are both group_slots(G,K). + lowering preserves slot_block(g) and slot_lane(g). Width-changing casts are + legal only when slot-preserving VPTO lowering support exists, or when the cast + can be commuted through a later group-aware consumer such as group_broadcast. +``` + +Illegal consumer mix: + +```text +group_slots value -> ordinary dense store/add/mul +``` + +This must fail unless an explicit semantic op converts the group-slot value: + +```text +group_broadcast +group_store +future explicit group-pack op +``` + +Contiguous memory loads may produce a non-contiguous physical value directly +when the requested result layout is a dense deinterleaved layout. This is a +lowering choice, not a separate layout family. + +```text +pto.vmi.load -> #pto.vmi.layout + lower as: + vlds NORM for each physical chunk + +pto.vmi.load -> #pto.vmi.layout + lower as: + vldsx2 DINTLV_B* for each pair of physical chunks + +pto.vmi.load -> #pto.vmi.layout + lower as: + two vldsx2 DINTLV_B* operations for each four-chunk group + followed by two vdintlv operations to split mod4 parts + +pto.vmi.load -> #pto.vmi.layout + lower using the producer-specific path or fall back to explicit + materialization. Do not treat DINTLV_B* as a block-fragment layout. +``` + +The `deinterleaved = 4` result order remains the normal VMI physical part +order: + +```text +results = [part0 chunks..., part1 chunks..., part2 chunks..., part3 chunks...] +``` + +For one full `256xf32` tile: + +```text +%even0, %odd0 = pto.vldsx2 %base[%off0], "DINTLV_B32" +%even1, %odd1 = pto.vldsx2 %base[%off128], "DINTLV_B32" + +%part0, %part2 = pto.vdintlv %even0, %even1 +%part1, %part3 = pto.vdintlv %odd0, %odd1 + +replace pto.vmi.load with [%part0, %part1, %part2, %part3] +``` + +This optimization is legal only for full physical chunks and supported +`DINTLV_B8/B16/B32` element widths. Tail and masked loads keep their explicit +safe lowering until a masked or guarded `vldsx2` strategy is designed. + +Two-way logical interleaved memory access is represented by dedicated VMI ops, +not by exposing assigned layouts in surface IR: + +```mlir +%x, %y = pto.vmi.deinterleave_load %src[%off] + : !pto.ptr -> !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> + +pto.vmi.interleave_store %x, %y, %dst[%off] + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32>, !pto.ptr +``` + +Each VMI value is an ordinary dense logical vector. Layout assignment requests +contiguous layouts for both streams. Lowering maps full-chunk 8/16/32-bit cases +to `vldsx2 DINTLV_B*` and `vstsx2 INTLV_B*`. + +## 3. Lowering Results + +The following examples use symbolic VPTO names. `PAT_ALL_B*` means an all-true +predicate with the element granularity required by the instruction. `PAT_VLk` +means a prefix predicate for the first `k` lanes. + +Completeness rule for this section: every numbered endpoint below must contain +VMI input, assigned layouts, VPTO lowering result, and either a memory result or +an explicit diagnostic. Non-endpoint layout notes may appear only as setup for +the immediately following complete endpoints. + +```text +3.1 f16 -> f32 -> store complete +3.2 f32 -> f16 -> store complete +3.3 f8 -> f32 -> compute -> f8 complete +3.4 group_reduce S=8 -> group_store complete +3.5.1 group_reduce S=16 -> group_store complete +3.5.2 group_reduce S=16 -> broadcast -> compute -> reduce -> store + complete +3.5.3 group_reduce S=16 -> elemwise(rhs) -> group_store complete +3.6.1 group_reduce S=32 -> group_store complete +3.6.2 group_reduce S=32 -> elemwise(rhs) -> group_store complete +3.6.3 group_reduce S=32 -> broadcast -> compute -> reduce -> store + complete +3.7.1 group_reduce S=64 -> aligned group_store complete +3.7.2 group_reduce S=64 -> elemwise(rhs) -> aligned group_store + complete +3.7.3 group_reduce S=64 -> broadcast -> compute -> reduce -> store + complete +3.7.4 group_reduce S=64 -> unit-stride group_store illegal diagnostic +3.8 group_reduce -> truncf -> broadcast -> dense store complete +3.9 dense store of group slots illegal diagnostic +3.10 non-load producer feeding S=32 group_reduce complete +3.11 partial tail groups complete/diagnostic +3.12 control-flow join before group_reduce complete +3.13 packed group-slot f32 -> f16 cast illegal diagnostic +3.14 unsupported group size illegal diagnostic +3.15 compact S=12 written as logical S=16 complete/diagnostic +3.16 group_slot_load layout contract complete +3.17 group_broadcast feeding deinterleaved consumer complete +3.18 one value with dense and group-reduce consumers complete/materialization +3.19 S=16 reduce block_elems support selection complete/diagnostic +3.20 group_slots control-flow join complete +3.21 S=32 tail with full-tile-readable source complete +3.22 scf.for loop-carried layout complete +3.23 group_broadcast with multiple dense consumers complete +3.24 mask with elementwise/select/store complete +3.25 function boundary layout specialization complete +3.26 S=16 grouped tail through broadcast/reduce/store complete +3.27 S=32 group_load with stride greater than group size complete +3.28 group_slot_load slots=1 aligned non-unit stride complete +3.29 one semantic mask with f32 and f16 consumers complete +3.30 masked_load tail without padding complete/diagnostic +3.31 f16->f32 feeding dense store and S=16 reduce complete +3.32 f32 feeding f8 store and S=32 reduce complete +3.33 one dense value feeding S=16 and S=32 reduces complete/materialization +3.34 S=64 group-slot result f32->f16 cast complete +3.35 group_slots fanout to group_store and broadcast complete +3.36 same scalar source materialized as slots=8/slots=1 complete/materialization +3.37 S=64 group_store with non-unit output stride complete +3.38 multi-tile S=32 group_reduce complete +3.39 strided S=32 group_load through broadcast/reduce complete +3.40 scalar broadcast feeding dense and grouped users complete/materialization +3.41 non-rematerializable value with incompatible users complete/materialization +3.42 group_slots scf.for loop-carried accumulator complete +3.43 internal function argument boundary materialization complete +3.44 masked_load grouped tail feeding S=32 reduce complete +3.45 dynamic S=32 create_group_mask complete +3.46 extf value and derived elemwise value both stored complete/optimization +3.47-3.55 typed group-reduce generalization complete/diagnostic +3.56 full 256-bin distribution histogram complete +3.57 full 256-bin cumulative histogram design boundary +``` + +### 3.1 `f16 -> f32 -> store` + +VMI input: + +```text +%x16 = pto.vmi.load %base[%off] + : memref<128xf16> -> !pto.vmi.vreg<128xf16> +%x32 = pto.vmi.extf %x16 + : !pto.vmi.vreg<128xf16> -> !pto.vmi.vreg<128xf32> +pto.vmi.store %x32, %out[%off] +``` + +Assigned layouts: + +```text +%x16 : !pto.vmi.vreg<128xf16, #pto.vmi.layout> +%x32 : !pto.vmi.vreg<128xf32, #pto.vmi.layout> +``` + +VPTO lowering result: + +```text +%x16_0 = pto.vlds %base[%off] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<128xf16> + +%x32_p0 = pto.vcvt %x16_0, PAT_ALL_B16 {part = "EVEN"} + : !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<64xf32> +%x32_p1 = pto.vcvt %x16_0, PAT_ALL_B16 {part = "ODD"} + : !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<64xf32> + +pto.vstsx2 %x32_p0, %x32_p1, %out[%off], "INTLV_B32", PAT_ALL_B32 + : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.ptr, index, + !pto.mask +``` + +Alternative complete VPTO lowering result if `vstsx2 INTLV_B32` is unavailable: + +```text +%x16_0 = pto.vlds %base[%off] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<128xf16> + +%x32_p0 = pto.vcvt %x16_0, PAT_ALL_B16 {part = "EVEN"} + : !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<64xf32> +%x32_p1 = pto.vcvt %x16_0, PAT_ALL_B16 {part = "ODD"} + : !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<64xf32> + +%x32_d0, %x32_d1 = pto.vintlv %x32_p0, %x32_p1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +pto.vsts %x32_d0, %out[%off], PAT_ALL_B32 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +pto.vsts %x32_d1, %out[%off_plus_64], PAT_ALL_B32 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for i = 0..127: + out[off + i] = extf(base[off + i]) +``` + +### 3.2 Dense `f32 -> f16 -> store` + +VMI input: + +```text +%x32 = pto.vmi.load %base[%off] + : memref<128xf32> -> !pto.vmi.vreg<128xf32> +%x16 = pto.vmi.truncf %x32 + : !pto.vmi.vreg<128xf32> -> !pto.vmi.vreg<128xf16> +pto.vmi.store %x16, %out[%off] +``` + +Assigned layouts: + +```text +%x32 : !pto.vmi.vreg<128xf32, #pto.vmi.layout> +%x16 : !pto.vmi.vreg<128xf16, #pto.vmi.layout> +``` + +VPTO lowering result: + +```text +%x32_p0, %x32_p1 = pto.vldsx2 %base[%off], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%part0 = pto.vcvt %x32_p0, PAT_ALL_B32 + {part = "EVEN", rnd = "R", sat = "SAT"} + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<128xf16> + +%part1 = pto.vcvt %x32_p1, PAT_ALL_B32 + {part = "ODD", rnd = "R", sat = "SAT"} + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<128xf16> + +%x16_0 = pto.vor %part0, %part1, PAT_ALL_B16 + : !pto.vreg<128xf16> + +pto.vsts %x16_0, %out[%off], PAT_ALL_B16 {dist = "NORM_B16"} + : !pto.vreg<128xf16>, !pto.ptr, !pto.mask +``` + +Alternative complete VPTO lowering result if the source has already been loaded +as two contiguous f32 chunks and must be materialized to `deinterleaved=2` before +the conversion: + +```text +%x32_d0 = pto.vlds %base[%off] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<64xf32> +%x32_d1 = pto.vlds %base[%off_plus_64] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<64xf32> +%x32_p0, %x32_p1 = pto.vdintlv %x32_d0, %x32_d1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%part0 = pto.vcvt %x32_p0, PAT_ALL_B32 + {part = "EVEN", rnd = "R", sat = "SAT"} + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<128xf16> + +%part1 = pto.vcvt %x32_p1, PAT_ALL_B32 + {part = "ODD", rnd = "R", sat = "SAT"} + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<128xf16> + +%x16_0 = pto.vor %part0, %part1, PAT_ALL_B16 + : !pto.vreg<128xf16> + +pto.vsts %x16_0, %out[%off], PAT_ALL_B16 {dist = "NORM_B16"} + : !pto.vreg<128xf16>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for i = 0..127: + out[off + i] = truncf(base[off + i]) +``` + +### 3.3 Dense `f8 -> f32 -> compute -> f8` + +VMI input: + +```text +%x8 = pto.vmi.load %base[%off] +%x32 = pto.vmi.extf %x8 +%scale = pto.vmi.broadcast %scale_s : f32 -> !pto.vmi.vreg<256xf32> +%y32 = pto.vmi.mulf %x32, %scale +%y8 = pto.vmi.truncf %y32 +pto.vmi.store %y8, %out[%off] +``` + +Assigned layouts: + +```text +%x8 : !pto.vmi.vreg<256xf8, #pto.vmi.layout> +%x32 : !pto.vmi.vreg<256xf32, #pto.vmi.layout> +%scale : !pto.vmi.vreg<256xf32, #pto.vmi.layout> +%y32 : !pto.vmi.vreg<256xf32, #pto.vmi.layout> +%y8 : !pto.vmi.vreg<256xf8, #pto.vmi.layout> +``` + +VPTO lowering result: + +```text +%x8_0 = pto.vlds %base[%off] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<256xf8> + +%x32_p0 = pto.vcvt %x8_0, PAT_ALL_B8 {part = "P0"} + : !pto.vreg<256xf8>, !pto.mask -> !pto.vreg<64xf32> +%x32_p1 = pto.vcvt %x8_0, PAT_ALL_B8 {part = "P1"} + : !pto.vreg<256xf8>, !pto.mask -> !pto.vreg<64xf32> +%x32_p2 = pto.vcvt %x8_0, PAT_ALL_B8 {part = "P2"} + : !pto.vreg<256xf8>, !pto.mask -> !pto.vreg<64xf32> +%x32_p3 = pto.vcvt %x8_0, PAT_ALL_B8 {part = "P3"} + : !pto.vreg<256xf8>, !pto.mask -> !pto.vreg<64xf32> + +%scale_p0 = pto.vdup %scale_s, PAT_ALL_B32 + : f32, !pto.mask -> !pto.vreg<64xf32> +%scale_p1 = pto.vdup %scale_s, PAT_ALL_B32 + : f32, !pto.mask -> !pto.vreg<64xf32> +%scale_p2 = pto.vdup %scale_s, PAT_ALL_B32 + : f32, !pto.mask -> !pto.vreg<64xf32> +%scale_p3 = pto.vdup %scale_s, PAT_ALL_B32 + : f32, !pto.mask -> !pto.vreg<64xf32> + +%y32_p0 = pto.vmul %x32_p0, %scale_p0, PAT_ALL_B32 +%y32_p1 = pto.vmul %x32_p1, %scale_p1, PAT_ALL_B32 +%y32_p2 = pto.vmul %x32_p2, %scale_p2, PAT_ALL_B32 +%y32_p3 = pto.vmul %x32_p3, %scale_p3, PAT_ALL_B32 + +%y8_p0 = pto.vcvt %y32_p0, PAT_ALL_B32 + {part = "P0", rnd = "R", sat = "SAT"} -> !pto.vreg<256xf8> +%y8_p1 = pto.vcvt %y32_p1, PAT_ALL_B32 + {part = "P1", rnd = "R", sat = "SAT"} -> !pto.vreg<256xf8> +%y8_p2 = pto.vcvt %y32_p2, PAT_ALL_B32 + {part = "P2", rnd = "R", sat = "SAT"} -> !pto.vreg<256xf8> +%y8_p3 = pto.vcvt %y32_p3, PAT_ALL_B32 + {part = "P3", rnd = "R", sat = "SAT"} -> !pto.vreg<256xf8> + +%y8_01 = pto.vor %y8_p0, %y8_p1, PAT_ALL_B8 +%y8_23 = pto.vor %y8_p2, %y8_p3, PAT_ALL_B8 +%y8_0 = pto.vor %y8_01, %y8_23, PAT_ALL_B8 + +pto.vsts %y8_0, %out[%off], PAT_ALL_B8 {dist = "NORM_B8"} + : !pto.vreg<256xf8>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for i = 0..255: + out[off + i] = truncf(extf(base[off + i]) * scale_s) +``` + +### 3.4 `group_reduce` S=8 f32 + +VMI input: + +```text +%x = pto.vmi.load %base[%off] + : memref<64xf32> -> !pto.vmi.vreg<64xf32> +%mask = pto.vmi.create_mask %c64 : index -> !pto.vmi.mask<64xpred> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} + : !pto.vmi.vreg<64xf32>, !pto.vmi.mask<64xpred> + -> !pto.vmi.vreg<64xf32> +pto.vmi.group_store %sum, %sum_out[%group_off], %c1 {num_groups = 8} +``` + +Assigned layouts: + +```text +%x : !pto.vmi.vreg<64xf32, #pto.vmi.layout> +%mask : !pto.vmi.mask<64xpred, #pto.vmi.layout> +%sum : !pto.vmi.vreg<64xf32, + #pto.vmi.layout> +``` + +VPTO lowering result for one full 8-row tile: + +```text +%mask_chunk = pto.pge_b32 "PAT_ALL" + +%x_chunk = pto.vlds %base[%tile_off] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<64xf32> + +%sum_block = pto.vcgadd %x_chunk, %mask_chunk + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +%store8 = pto.pge_b32 "PAT_VL8" +pto.vsts %sum_block, %sum_out[%group_tile_off], %store8 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Lowering result for one chunk, per the `visa.txt` VCGADD contract: + +```text +%sum_block lane 0 = reduce %x lanes 0..7 +%sum_block lane 1 = reduce %x lanes 8..15 +... +%sum_block lane 7 = reduce %x lanes 56..63 +all non-slot lanes are non-semantic +``` + +Layout result: + +```text +G = N / 8 +K = 8 + +slot_block(g) = g / 8 +slot_lane(g) = g % 8 +``` + +Memory result: + +```text +for r = 0..7: + sum_out[group_tile_off + r] = reduce(row_r[0..7]) +``` + +### 3.5 `group_reduce` S=16 f32, load-fused split + +The facts used by this lowering are checked against the current repo: + +```text +pto.vldsx2 supports "BDINTLV". +pto.vstsx2 supports only "INTLV_B8" / "INTLV_B16" / "INTLV_B32". +visa.txt says VCGADD writes one 32B-block result continuously to destination +LSBs; the current repository golden tests follow lanes 0..7 for f32. +``` + +There are three complete consumers for this layout today: + +```text +load -> group_reduce -> group_store(sum) +load -> group_reduce -> elementwise compute on group-slot values + -> group_store +load -> group_reduce -> group_broadcast -> elementwise compute + -> group_reduce -> group_store +``` + +#### 3.5.1 Reduce And Store Group Sums + +VMI input: + +```text +%x = pto.vmi.load %base[%off] + : memref -> !pto.vmi.vreg +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = N / 16} +pto.vmi.group_store %sum, %sum_out[%group_off], %c1 {num_groups = N / 16} +``` + +Assigned layouts: + +```text +%x : !pto.vmi.vreg> + +%sum : !pto.vmi.vreg> +``` + +For each 8-row tile: + +```text +row r = 16xf32 = row_r.lo8, row_r.hi8 +``` + +VPTO lowering result for one full 8-row tile: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%sum_mask = pto.pge_b32 "PAT_VL8" + +%lo, %hi = pto.vldsx2 %base[%tile_off], "BDINTLV" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%lo lanes 0..7 = row0.lo8 +%lo lanes 8..15 = row1.lo8 +... +%lo lanes 56..63 = row7.lo8 + +%hi lanes 0..7 = row0.hi8 +%hi lanes 8..15 = row1.hi8 +... +%hi lanes 56..63 = row7.hi8 + +%lo_sum = pto.vcgadd %lo, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%hi_sum = pto.vcgadd %hi, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +%sum_block = pto.vadd %lo_sum, %hi_sum, %sum_mask + : !pto.vreg<64xf32> + +%store8 = pto.pge_b32 "PAT_VL8" +pto.vsts %sum_block, %sum_out[%group_tile_off], %store8 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +`BDINTLV` here denotes the ISA `#bdintlv` block-based interleaving load mode: +it loads `2 * VL` bytes and sends even 32B blocks to the first destination +register and odd 32B blocks to the second destination register. For f32, +one 32B block is `8xf32`, matching `block_elems = 8`. + +Tail tiles use the same dataflow with `%all_b32` replaced by masks derived from +the VMI mask for the low and high 8-lane halves of each row. + +Layout result: + +```text +G = N / 16 +K = 8 + +slot_block(g) = g / 8 +slot_lane(g) = g % 8 + +%sum_block lane 0 = reduce row0 lanes 0..15 +%sum_block lane 1 = reduce row1 lanes 0..15 +... +%sum_block lane 7 = reduce row7 lanes 0..15 +``` + +No VMI value exposes `%lo_sum` or `%hi_sum`. They are internal VPTO values. + +Memory result: + +```text +sum_out[group_tile_off + 0] = reduce row0 lanes 0..15 +sum_out[group_tile_off + 1] = reduce row1 lanes 0..15 +... +sum_out[group_tile_off + 7] = reduce row7 lanes 0..15 +``` + +This endpoint is fully specified: the only group-slot value is `%sum`; `group_store` +stores the low 8 slot lanes with an ordinary prefix store. + +#### 3.5.2 Reduce, Broadcast, Elementwise, Reduce, Store + +VMI input: + +```text +%x = pto.vmi.load %base[%off] + : memref -> !pto.vmi.vreg +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = N / 16} +%b = pto.vmi.group_broadcast %sum {num_groups = N / 16} +%y = pto.vmi.mulf %x, %b +%ysum = pto.vmi.group_reduce_addf %y, %mask {num_groups = N / 16} +pto.vmi.group_store %ysum, %out[%group_off], %c1 {num_groups = N / 16} +``` + +Assigned layouts: + +```text +%x : !pto.vmi.vreg> +%sum : !pto.vmi.vreg> +%b : !pto.vmi.vreg> +%y : !pto.vmi.vreg> +%ysum : !pto.vmi.vreg> +``` + +VPTO lowering result for one full 8-row tile: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%sum_mask = pto.pge_b32 "PAT_VL8" + +%x_lo, %x_hi = pto.vldsx2 %base[%tile_off], "BDINTLV" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%x_lo_sum = pto.vcgadd %x_lo, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%x_hi_sum = pto.vcgadd %x_hi, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +%sum_block = pto.vadd %x_lo_sum, %x_hi_sum, %sum_mask + : !pto.vreg<64xf32> + +%lane_id = pto.vci %c0_i32 : i32 -> !pto.vreg<64xi32> +%broadcast_idx = pto.vshrs %lane_id, %c3_i16, %all_b32 + : !pto.vreg<64xi32>, i16, !pto.mask -> !pto.vreg<64xi32> + +// This is the materialization of pto.vmi.group_broadcast. The group sums are +// in %sum_block lanes 0..7; vselr expands each sum to the 8 lanes of the +// corresponding row half. The following vmul/vcgadd consume an ordinary dense +// physical vector. +%b_rows = pto.vselr %sum_block, %broadcast_idx + : !pto.vreg<64xf32>, !pto.vreg<64xi32> -> !pto.vreg<64xf32> + +%y_lo = pto.vmul %x_lo, %b_rows, %all_b32 + : !pto.vreg<64xf32> +%y_hi = pto.vmul %x_hi, %b_rows, %all_b32 + : !pto.vreg<64xf32> + +%y_lo_sum = pto.vcgadd %y_lo, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%y_hi_sum = pto.vcgadd %y_hi, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +// Final per-row reduction and store. +%ysum_block = pto.vadd %y_lo_sum, %y_hi_sum, %sum_mask + : !pto.vreg<64xf32> + +%store8 = pto.pge_b32 "PAT_VL8" +pto.vsts %ysum_block, %out[%group_tile_off], %store8 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +This trace processes 8 logical rows at once. `num_groups = N / 16` means each +logical group is one `16xf32` row, and one full f32 VPTO tile covers 8 such +groups: + +```text +64 f32 lanes per physical part = 8 rows * 8 f32 lanes per half-row +``` + +Tail tiles use the same dataflow with `%all_b32` replaced by masks derived from +the VMI mask for the low and high 8-lane halves of each row. + +Physical lane result for the tile: + +```text +%x_lo lanes 0..7 = row0[0..7] +%x_lo lanes 8..15 = row1[0..7] +... +%x_lo lanes 56..63 = row7[0..7] + +%x_hi lanes 0..7 = row0[8..15] +%x_hi lanes 8..15 = row1[8..15] +... +%x_hi lanes 56..63 = row7[8..15] + +%sum_block lanes 0..7 = + reduce(row0[0..15]), reduce(row1[0..15]), ..., reduce(row7[0..15]) + +%b_rows lanes 0..7 = reduce(row0[0..15]) +%b_rows lanes 8..15 = reduce(row1[0..15]) +... +%b_rows lanes 56..63 = reduce(row7[0..15]) + +For each row `r` in this 8-row tile: + +%y_lo lanes r*8 .. r*8+7 = + row_r[0..7] * reduce(row_r[0..15]) + +%y_hi lanes r*8 .. r*8+7 = + row_r[8..15] * reduce(row_r[0..15]) + +Concretely: +%y_lo lanes 0..7 = row0[0..7] * reduce(row0[0..15]) +%y_lo lanes 8..15 = row1[0..7] * reduce(row1[0..15]) +... +%y_lo lanes 56..63 = row7[0..7] * reduce(row7[0..15]) + +%y_hi lanes 0..7 = row0[8..15] * reduce(row0[0..15]) +%y_hi lanes 8..15 = row1[8..15] * reduce(row1[0..15]) +... +%y_hi lanes 56..63 = row7[8..15] * reduce(row7[0..15]) + +%ysum_block lanes 0..7 = + reduce(%y row0), reduce(%y row1), ..., reduce(%y row7) +``` + +Memory result: + +```text +out[group_tile_off + r] = + reduce_i((row_r[i] * reduce_j(row_r[j])) for i in 0..15) + = reduce(row_r[0..15]) * reduce(row_r[0..15]) +for r = 0..7 +``` + +If a later consumer requires row-major contiguous order, `vmi-to-vpto` must +materialize: + +```text +deinterleaved=2, block_elems=8 -> contiguous +``` + +This materialization cannot be implemented with `vstsx2 INTLV_B32`, because +that instruction interleaves individual b32 elements, not 32B row halves. Until +a concrete block-interleave register materialization or store op is selected, +row-major store of this layout must be rejected with: + +```text +VMI-LAYOUT-CONTRACT: + pto.vmi.store requires materializing + #pto.vmi.layout to contiguous, but no + VPTO block-interleave materialization/store support exists. +``` + +#### 3.5.3 Reduce Result, Elementwise, Store + +This case computes a per-row reduction, applies an elementwise operation to the +reduced values themselves, and stores one result per group. There is no +`group_broadcast` in this flow because the elementwise op is not applied to the +original `8x16xf32` matrix elements. + +VMI input: + +```text +%x = pto.vmi.load %base[%off] + : memref<128xf32> -> !pto.vmi.vreg<128xf32> +%rhs = pto.vmi.group_slot_load %rhs_base[%rhs_off], %c1 {num_groups = 8} + : !pto.ptr -> !pto.vmi.vreg<128xf32> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +%outv = pto.vmi.addf %sum, %rhs +pto.vmi.group_store %outv, %out[%group_off], %c1 {num_groups = 8} +``` + +Assigned layouts: + +```text +%x for reduce: + !pto.vmi.vreg<128xf32, + #pto.vmi.layout> + +%sum: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> + +%rhs: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> + +%outv: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> +``` + +For this endpoint, the RHS is a packed per-group vector: + +```text +rhs_base[rhs_off + r] = rhs(row r), for r = 0..7 +``` + +Layout assignment must treat `group_slot_load` as a group-slot producer: one +f32 value per group is placed in the live slot lanes. It must not use +`group_load`, which loads `group_size` data elements per group instead of one +per-group scalar. + +The elementwise op runs only on the live group-slot lanes: + +```text +%sum lanes 0..7 = + reduce(row0[0..15]), reduce(row1[0..15]), ..., reduce(row7[0..15]) + +%rhs lanes 0..7 = + rhs(row0), rhs(row1), ..., rhs(row7) + +%outv lanes 0..7 = + %sum lanes 0..7 + %rhs lanes 0..7 + +lanes 8..63 remain dead/zero and are masked off by PAT_VL8. +``` + +VPTO lowering result for one full 8-row tile: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%sum_mask = pto.pge_b32 "PAT_VL8" +%one_block = pto.pge_b32 "PAT_VL1" + +// Reduction path: use BDINTLV to feed two VCG reductions. +%x_lo, %x_hi = pto.vldsx2 %base[%tile_off], "BDINTLV" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%x_lo_sum = pto.vcgadd %x_lo, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%x_hi_sum = pto.vcgadd %x_hi, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +%sum_block = pto.vadd %x_lo_sum, %x_hi_sum, %sum_mask + : !pto.vreg<64xf32> + +// Packed RHS group-slot load. %rhs_tile_base points to rhs_base[rhs_off]. +// One 32B block contains 8 f32 RHS values and materializes lanes 0..7; all +// other lanes are dead/zero. +%rhs_block = pto.vsldb %rhs_tile_base, %c0_i16, %c0_i16, %one_block + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> + +// Elementwise compute on group-slot values. Only lanes 0..7 are live. +%outv_block = pto.vadd %sum_block, %rhs_block, %sum_mask + : !pto.vreg<64xf32> + +pto.vsts %outv_block, %out[%group_tile_off], %sum_mask {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + s = reduce(row_r[0..15]) + out[group_tile_off + r] = s + rhs[r] +``` + +### 3.6 `group_reduce` S=32 f32, 4-way split + +This case covers one `8x32xf32` tile. Each logical row is 128B, so it must be +split into four 32B partial rows before `vcgadd` can reduce it efficiently. + +The canonical layout for the input is: + +```text +%x : !pto.vmi.vreg<256xf32, #pto.vmi.layout> +``` + +With `deinterleaved = 4`, physical part `p` contains columns whose logical +column index is `p mod 4`: + +```text +%x_p0 lanes r*8 .. r*8+7 = + row_r[0], row_r[4], row_r[8], ..., row_r[28] + +%x_p1 lanes r*8 .. r*8+7 = + row_r[1], row_r[5], row_r[9], ..., row_r[29] + +%x_p2 lanes r*8 .. r*8+7 = + row_r[2], row_r[6], row_r[10], ..., row_r[30] + +%x_p3 lanes r*8 .. r*8+7 = + row_r[3], row_r[7], row_r[11], ..., row_r[31] +``` + +Each physical part now has exactly 8 f32 values per row, so one `vcgadd` per +part computes one partial sum per row. The four partial sums are then added +under `PAT_VL8`. + +The full contiguous-to-4-way materialization for one tile should fuse the first +deinterleave level into the load. `vldsx2 DINTLV_B32` loads `2 * VL` bytes and +splits even/odd f32 elements into two physical vectors. Two such loads cover +the `8x32xf32` tile, and a second register `vdintlv` level splits even columns +into `mod4 = 0/2` and odd columns into `mod4 = 1/3`. + +This setup documentation is repeated inside every complete 32-wide endpoint +below. + +```text +%x_even_0, %x_odd_0 = pto.vldsx2 %base[%tile_off_0], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%x_even_1, %x_odd_1 = pto.vldsx2 %base[%tile_off_1], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%x_p0, %x_p2 = pto.vdintlv %x_even_0, %x_even_1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%x_p1, %x_p3 = pto.vdintlv %x_odd_0, %x_odd_1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +``` + +Each endpoint below inlines this materialization before the first consumer of +`%x_p0..%x_p3`. + +#### 3.6.1 Reduce And Store Group Sums + +VMI input: + +```text +%x = pto.vmi.load %base[%off] + : memref<256xf32> -> !pto.vmi.vreg<256xf32> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +pto.vmi.group_store %sum, %sum_out[%group_off], %c1 {num_groups = 8} +``` + +Assigned layouts: + +```text +%x: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> + +%sum: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> +``` + +VPTO lowering result for one full 8-row tile: + +```text +%x_even_0, %x_odd_0 = pto.vldsx2 %base[%tile_off_0], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%x_even_1, %x_odd_1 = pto.vldsx2 %base[%tile_off_1], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%x_p0, %x_p2 = pto.vdintlv %x_even_0, %x_even_1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%x_p1, %x_p3 = pto.vdintlv %x_odd_0, %x_odd_1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%all_b32 = pto.pge_b32 "PAT_ALL" +%sum_mask = pto.pge_b32 "PAT_VL8" + +%s0 = pto.vcgadd %x_p0, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s1 = pto.vcgadd %x_p1, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s2 = pto.vcgadd %x_p2, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s3 = pto.vcgadd %x_p3, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +%s01 = pto.vadd %s0, %s1, %sum_mask : !pto.vreg<64xf32> +%s23 = pto.vadd %s2, %s3, %sum_mask : !pto.vreg<64xf32> +%sum_block = pto.vadd %s01, %s23, %sum_mask : !pto.vreg<64xf32> + +pto.vsts %sum_block, %sum_out[%group_tile_off], %sum_mask {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + sum_out[group_tile_off + r] = reduce(row_r[0..31]) +``` + +#### 3.6.2 Reduce Result, Elementwise, Store + +VMI input: + +```text +%x = pto.vmi.load %base[%off] + : memref<256xf32> -> !pto.vmi.vreg<256xf32> +%rhs = pto.vmi.group_slot_load %rhs_base[%rhs_off], %c1 {num_groups = 8} + : !pto.ptr -> !pto.vmi.vreg<256xf32> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +%outv = pto.vmi.addf %sum, %rhs +pto.vmi.group_store %outv, %out[%group_off], %c1 {num_groups = 8} +``` + +Assigned layouts: + +```text +%x: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> + +%sum, %rhs, %outv: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> +``` + +VPTO lowering result for one full 8-row tile: + +```text +%x_even_0, %x_odd_0 = pto.vldsx2 %base[%tile_off_0], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%x_even_1, %x_odd_1 = pto.vldsx2 %base[%tile_off_1], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%x_p0, %x_p2 = pto.vdintlv %x_even_0, %x_even_1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%x_p1, %x_p3 = pto.vdintlv %x_odd_0, %x_odd_1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%all_b32 = pto.pge_b32 "PAT_ALL" +%sum_mask = pto.pge_b32 "PAT_VL8" +%one_block = pto.pge_b32 "PAT_VL1" + +%s0 = pto.vcgadd %x_p0, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s1 = pto.vcgadd %x_p1, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s2 = pto.vcgadd %x_p2, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s3 = pto.vcgadd %x_p3, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +%s01 = pto.vadd %s0, %s1, %sum_mask : !pto.vreg<64xf32> +%s23 = pto.vadd %s2, %s3, %sum_mask : !pto.vreg<64xf32> +%sum_block = pto.vadd %s01, %s23, %sum_mask : !pto.vreg<64xf32> + +// Packed RHS group-slot load. %rhs_tile_base points to rhs_base[rhs_off]. +%rhs_block = pto.vsldb %rhs_tile_base, %c0_i16, %c0_i16, %one_block + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> + +%outv_block = pto.vadd %sum_block, %rhs_block, %sum_mask + : !pto.vreg<64xf32> + +pto.vsts %outv_block, %out[%group_tile_off], %sum_mask {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + out[group_tile_off + r] = reduce(row_r[0..31]) + rhs[r] +``` + +#### 3.6.3 Reduce, Broadcast, Elementwise, Reduce, Store + +VMI input: + +```text +%x = pto.vmi.load %base[%off] + : memref<256xf32> -> !pto.vmi.vreg<256xf32> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +%b = pto.vmi.group_broadcast %sum {num_groups = 8} +%y = pto.vmi.mulf %x, %b +%ysum = pto.vmi.group_reduce_addf %y, %mask {num_groups = 8} +pto.vmi.group_store %ysum, %out[%group_off], %c1 {num_groups = 8} +``` + +Assigned layouts: + +```text +%x, %b, %y: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> + +%sum, %ysum: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> +``` + +VPTO lowering result for one full 8-row tile: + +```text +%x_even_0, %x_odd_0 = pto.vldsx2 %base[%tile_off_0], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%x_even_1, %x_odd_1 = pto.vldsx2 %base[%tile_off_1], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%x_p0, %x_p2 = pto.vdintlv %x_even_0, %x_even_1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%x_p1, %x_p3 = pto.vdintlv %x_odd_0, %x_odd_1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%all_b32 = pto.pge_b32 "PAT_ALL" +%sum_mask = pto.pge_b32 "PAT_VL8" + +%s0 = pto.vcgadd %x_p0, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s1 = pto.vcgadd %x_p1, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s2 = pto.vcgadd %x_p2, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s3 = pto.vcgadd %x_p3, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +%s01 = pto.vadd %s0, %s1, %sum_mask : !pto.vreg<64xf32> +%s23 = pto.vadd %s2, %s3, %sum_mask : !pto.vreg<64xf32> +%sum_block = pto.vadd %s01, %s23, %sum_mask : !pto.vreg<64xf32> + +%lane_id = pto.vci %c0_i32 : i32 -> !pto.vreg<64xi32> +%broadcast_idx = pto.vshrs %lane_id, %c3_i16, %all_b32 + : !pto.vreg<64xi32>, i16, !pto.mask -> !pto.vreg<64xi32> + +// group_broadcast materialized for each deinterleaved=4 physical part. +%b_p0 = pto.vselr %sum_block, %broadcast_idx + : !pto.vreg<64xf32>, !pto.vreg<64xi32> -> !pto.vreg<64xf32> +%b_p1 = pto.vselr %sum_block, %broadcast_idx + : !pto.vreg<64xf32>, !pto.vreg<64xi32> -> !pto.vreg<64xf32> +%b_p2 = pto.vselr %sum_block, %broadcast_idx + : !pto.vreg<64xf32>, !pto.vreg<64xi32> -> !pto.vreg<64xf32> +%b_p3 = pto.vselr %sum_block, %broadcast_idx + : !pto.vreg<64xf32>, !pto.vreg<64xi32> -> !pto.vreg<64xf32> + +%y_p0 = pto.vmul %x_p0, %b_p0, %all_b32 : !pto.vreg<64xf32> +%y_p1 = pto.vmul %x_p1, %b_p1, %all_b32 : !pto.vreg<64xf32> +%y_p2 = pto.vmul %x_p2, %b_p2, %all_b32 : !pto.vreg<64xf32> +%y_p3 = pto.vmul %x_p3, %b_p3, %all_b32 : !pto.vreg<64xf32> + +%ys0 = pto.vcgadd %y_p0, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%ys1 = pto.vcgadd %y_p1, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%ys2 = pto.vcgadd %y_p2, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%ys3 = pto.vcgadd %y_p3, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +%ys01 = pto.vadd %ys0, %ys1, %sum_mask : !pto.vreg<64xf32> +%ys23 = pto.vadd %ys2, %ys3, %sum_mask : !pto.vreg<64xf32> +%ysum_block = pto.vadd %ys01, %ys23, %sum_mask : !pto.vreg<64xf32> + +pto.vsts %ysum_block, %out[%group_tile_off], %sum_mask {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + s = reduce(row_r[0..31]) + out[group_tile_off + r] = + reduce_i(row_r[i] * s for i = 0..31) + = s * s +``` + +### 3.7 `group_reduce` S=64 f32, row-local reduction + +This case covers one `8x64xf32` tile. Each logical row is exactly 256B, so the +input does not need a deinterleaved layout: + +```text +row r = 64xf32 = one !pto.vreg<64xf32> +``` + +The reduction is two-stage but row-local: + +```text +vcgadd(row_r) -> 8 partial sums in lanes 0..7 +vcadd(PAT_VL8) -> one row sum in lane 0 +``` + +The result layout is therefore not `slots = 8`. It is: + +```text +#pto.vmi.layout +``` + +Physical slot mapping for this tile: + +```text +slot_block(r) = r +slot_lane(r) = 0 + +%sum0 lane 0 = reduce row0 lanes 0..63 +%sum1 lane 0 = reduce row1 lanes 0..63 +... +%sum7 lane 0 = reduce row7 lanes 0..63 +``` + +Trying to canonicalize this result to `slots = 8` would require packing lane 0 +from eight different physical vregs into lanes 0..7 of one vreg. This document +does not use that packing transform. `slots = 1` is the canonical layout for +S=64 row-local group reductions. + +#### 3.7.1 Reduce And Store Group Sums + +VMI input: + +```text +%x = pto.vmi.load %base[%off] + : memref<512xf32> -> !pto.vmi.vreg<512xf32> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +%c8 = arith.constant 8 : index +pto.vmi.group_store %sum, %sum_out[%group_off], %c8 {num_groups = 8} +``` + +Assigned layouts: + +```text +%x: + !pto.vmi.vreg<512xf32, #pto.vmi.layout> + +%sum: + !pto.vmi.vreg<512xf32, #pto.vmi.layout> +``` + +VPTO lowering result for one full 8-row tile: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%block8 = pto.pge_b32 "PAT_VL8" +%one_b32 = pto.pge_b32 "PAT_VL1" + +%x0 = pto.vlds %base[%row_off_0] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<64xf32> +%x1 = pto.vlds %base[%row_off_1] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<64xf32> +%x2 = pto.vlds %base[%row_off_2] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<64xf32> +%x3 = pto.vlds %base[%row_off_3] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<64xf32> +%x4 = pto.vlds %base[%row_off_4] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<64xf32> +%x5 = pto.vlds %base[%row_off_5] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<64xf32> +%x6 = pto.vlds %base[%row_off_6] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<64xf32> +%x7 = pto.vlds %base[%row_off_7] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<64xf32> + +%p0 = pto.vcgadd %x0, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%p1 = pto.vcgadd %x1, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%p2 = pto.vcgadd %x2, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%p3 = pto.vcgadd %x3, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%p4 = pto.vcgadd %x4, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%p5 = pto.vcgadd %x5, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%p6 = pto.vcgadd %x6, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%p7 = pto.vcgadd %x7, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +%sum0 = pto.vcadd %p0, %block8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%sum1 = pto.vcadd %p1, %block8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%sum2 = pto.vcadd %p2, %block8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%sum3 = pto.vcadd %p3, %block8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%sum4 = pto.vcadd %p4, %block8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%sum5 = pto.vcadd %p5, %block8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%sum6 = pto.vcadd %p6, %block8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%sum7 = pto.vcadd %p7, %block8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +pto.vsts %sum0, %sum_out[%group_tile_off_0], %one_b32 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +pto.vsts %sum1, %sum_out[%group_tile_off_1], %one_b32 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +pto.vsts %sum2, %sum_out[%group_tile_off_2], %one_b32 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +pto.vsts %sum3, %sum_out[%group_tile_off_3], %one_b32 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +pto.vsts %sum4, %sum_out[%group_tile_off_4], %one_b32 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +pto.vsts %sum5, %sum_out[%group_tile_off_5], %one_b32 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +pto.vsts %sum6, %sum_out[%group_tile_off_6], %one_b32 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +pto.vsts %sum7, %sum_out[%group_tile_off_7], %one_b32 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + sum_out[group_tile_off + r * 8] = reduce(row_r[0..63]) +``` + +#### 3.7.2 Reduce Result, Elementwise, Store + +VMI input: + +```text +%x = pto.vmi.load %base[%off] + : memref<512xf32> -> !pto.vmi.vreg<512xf32> +%rhs = pto.vmi.group_slot_load %rhs_base[%rhs_off], %c1 {num_groups = 8} + : !pto.ptr -> !pto.vmi.vreg<512xf32> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +%outv = pto.vmi.addf %sum, %rhs +%c8 = arith.constant 8 : index +pto.vmi.group_store %outv, %out[%group_off], %c8 {num_groups = 8} +``` + +Assigned layouts: + +```text +%x: + !pto.vmi.vreg<512xf32, #pto.vmi.layout> + +%sum, %rhs, %outv: + !pto.vmi.vreg<512xf32, #pto.vmi.layout> +``` + +VPTO lowering result for one full 8-row tile: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%block8 = pto.pge_b32 "PAT_VL8" +%one_b32 = pto.pge_b32 "PAT_VL1" + +%x0 = pto.vlds %base[%row_off_0] {dist = "NORM"} : !pto.ptr -> !pto.vreg<64xf32> +%x1 = pto.vlds %base[%row_off_1] {dist = "NORM"} : !pto.ptr -> !pto.vreg<64xf32> +%x2 = pto.vlds %base[%row_off_2] {dist = "NORM"} : !pto.ptr -> !pto.vreg<64xf32> +%x3 = pto.vlds %base[%row_off_3] {dist = "NORM"} : !pto.ptr -> !pto.vreg<64xf32> +%x4 = pto.vlds %base[%row_off_4] {dist = "NORM"} : !pto.ptr -> !pto.vreg<64xf32> +%x5 = pto.vlds %base[%row_off_5] {dist = "NORM"} : !pto.ptr -> !pto.vreg<64xf32> +%x6 = pto.vlds %base[%row_off_6] {dist = "NORM"} : !pto.ptr -> !pto.vreg<64xf32> +%x7 = pto.vlds %base[%row_off_7] {dist = "NORM"} : !pto.ptr -> !pto.vreg<64xf32> + +%p0 = pto.vcgadd %x0, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%p1 = pto.vcgadd %x1, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%p2 = pto.vcgadd %x2, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%p3 = pto.vcgadd %x3, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%p4 = pto.vcgadd %x4, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%p5 = pto.vcgadd %x5, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%p6 = pto.vcgadd %x6, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%p7 = pto.vcgadd %x7, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +%sum0 = pto.vcadd %p0, %block8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%sum1 = pto.vcadd %p1, %block8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%sum2 = pto.vcadd %p2, %block8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%sum3 = pto.vcadd %p3, %block8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%sum4 = pto.vcadd %p4, %block8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%sum5 = pto.vcadd %p5, %block8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%sum6 = pto.vcadd %p6, %block8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%sum7 = pto.vcadd %p7, %block8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +%rhs0 = pto.vsldb %rhs_ptr_0, %c0_i16, %c0_i16, %one_b32 + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> +%rhs1 = pto.vsldb %rhs_ptr_1, %c0_i16, %c0_i16, %one_b32 + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> +%rhs2 = pto.vsldb %rhs_ptr_2, %c0_i16, %c0_i16, %one_b32 + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> +%rhs3 = pto.vsldb %rhs_ptr_3, %c0_i16, %c0_i16, %one_b32 + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> +%rhs4 = pto.vsldb %rhs_ptr_4, %c0_i16, %c0_i16, %one_b32 + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> +%rhs5 = pto.vsldb %rhs_ptr_5, %c0_i16, %c0_i16, %one_b32 + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> +%rhs6 = pto.vsldb %rhs_ptr_6, %c0_i16, %c0_i16, %one_b32 + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> +%rhs7 = pto.vsldb %rhs_ptr_7, %c0_i16, %c0_i16, %one_b32 + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> + +%out0 = pto.vadd %sum0, %rhs0, %one_b32 : !pto.vreg<64xf32> +%out1 = pto.vadd %sum1, %rhs1, %one_b32 : !pto.vreg<64xf32> +%out2 = pto.vadd %sum2, %rhs2, %one_b32 : !pto.vreg<64xf32> +%out3 = pto.vadd %sum3, %rhs3, %one_b32 : !pto.vreg<64xf32> +%out4 = pto.vadd %sum4, %rhs4, %one_b32 : !pto.vreg<64xf32> +%out5 = pto.vadd %sum5, %rhs5, %one_b32 : !pto.vreg<64xf32> +%out6 = pto.vadd %sum6, %rhs6, %one_b32 : !pto.vreg<64xf32> +%out7 = pto.vadd %sum7, %rhs7, %one_b32 : !pto.vreg<64xf32> + +pto.vsts %out0, %out[%group_tile_off_0], %one_b32 {dist = "NORM_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +pto.vsts %out1, %out[%group_tile_off_1], %one_b32 {dist = "NORM_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +pto.vsts %out2, %out[%group_tile_off_2], %one_b32 {dist = "NORM_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +pto.vsts %out3, %out[%group_tile_off_3], %one_b32 {dist = "NORM_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +pto.vsts %out4, %out[%group_tile_off_4], %one_b32 {dist = "NORM_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +pto.vsts %out5, %out[%group_tile_off_5], %one_b32 {dist = "NORM_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +pto.vsts %out6, %out[%group_tile_off_6], %one_b32 {dist = "NORM_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +pto.vsts %out7, %out[%group_tile_off_7], %one_b32 {dist = "NORM_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + out[group_tile_off + r * 8] = reduce(row_r[0..63]) + rhs[r] +``` + +#### 3.7.3 Reduce, Broadcast, Elementwise, Reduce, Store + +VMI input: + +```text +%x = pto.vmi.load %base[%off] + : memref<512xf32> -> !pto.vmi.vreg<512xf32> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +%b = pto.vmi.group_broadcast %sum {num_groups = 8} +%y = pto.vmi.mulf %x, %b +%ysum = pto.vmi.group_reduce_addf %y, %mask {num_groups = 8} +%c8 = arith.constant 8 : index +pto.vmi.group_store %ysum, %out[%group_off], %c8 {num_groups = 8} +``` + +Assigned layouts: + +```text +%x, %b, %y: + !pto.vmi.vreg<512xf32, #pto.vmi.layout> + +%sum, %ysum: + !pto.vmi.vreg<512xf32, #pto.vmi.layout> +``` + +VPTO lowering result for one full 8-row tile: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%block8 = pto.pge_b32 "PAT_VL8" +%one_b32 = pto.pge_b32 "PAT_VL1" + +// The compiler emits this row-local block once for each r in 0..7. +%x_r = pto.vlds %base[%row_off_r] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<64xf32> + +%p_r = pto.vcgadd %x_r, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%sum_r = pto.vcadd %p_r, %block8 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +// This vdup is the lowering of pto.vmi.group_broadcast for slots=1. +%b_r = pto.vdup %sum_r, %all_b32 {position = "LOWEST"} + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +%y_r = pto.vmul %x_r, %b_r, %all_b32 : !pto.vreg<64xf32> + +%yp_r = pto.vcgadd %y_r, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%ysum_r = pto.vcadd %yp_r, %block8 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +pto.vsts %ysum_r, %out[%group_tile_off_r], %one_b32 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +The row-local block above is not a runtime loop requirement. It is the repeated +VPTO shape for row offsets `%row_off_0` through `%row_off_7` and store offsets +`%group_tile_off_0` through `%group_tile_off_7`. + +Memory result: + +```text +for r = 0..7: + s = reduce(row_r[0..63]) + out[group_tile_off + r * 8] = + reduce_i(row_r[i] * s for i = 0..63) + = s * s +``` + +#### 3.7.4 Slots=1 Store Lowers To Packed Or Point Stores + +The row-local S=64 result uses one physical vreg per group with the semantic +value in lane 0: + +```text +%sum_r lane 0 = reduce(row_r[0..63]) +``` + +The current VPTO lowering for `slots = 1` group_store has two paths. + +For unit-stride output where all groups fit in one physical vector, the +lowering packs the lane-0 values into one dense vector and stores that vector +with a normal `vsts`. + +For non-unit row strides, each group stores its lane-0 scalar with a point +store. That emits `vsts` with `dist = "1PT_B32"` for f32 and only requires the +natural 4B alignment of the scalar element. + +VMI input: + +```text +%c1 = arith.constant 1 : index +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +pto.vmi.group_store %sum, %out[%group_off], %c1 {num_groups = 8} +``` + +Current checked-in coverage for the point-store path is: + +```text +test/lit/vmi/vmi_to_vpto_group_store_slots1_1pt.pto +``` + +### 3.8 `group_reduce -> truncf -> group_broadcast -> store` + +This case keeps the source op order by representing the `f32 -> f16` cast result +as lane-strided group slots. The source reduction is a packed f32 group-slot +value in b32 lanes 0..7. After `truncf`, the f16 values occupy even b16 lanes +0, 2, 4, ..., 14, so the result layout is `slots = 8, lane_stride = 2`. + +VMI input: + +```text +%x = pto.vmi.load %base[%off] + : memref<128xf32> -> !pto.vmi.vreg<128xf32> +%sum32 = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +%sum16 = pto.vmi.truncf %sum32 +%b16 = pto.vmi.group_broadcast %sum16 {num_groups = 8} +pto.vmi.store %b16, %out[%off] +``` + +Final assigned IR: + +```text +%mask = pto.vmi.create_mask %active + : index -> !pto.vmi.mask<128xb32, + #pto.vmi.layout> + +%x = pto.vmi.load %base[%off] + : memref<128xf32> -> !pto.vmi.vreg<128xf32, + #pto.vmi.layout> + +%mask_d2 = pto.vmi.ensure_mask_layout %mask + : !pto.vmi.mask<128xb32, + #pto.vmi.layout> + -> !pto.vmi.mask<128xb32, + #pto.vmi.layout> + +%sum32 = pto.vmi.group_reduce_addf %x, %mask_d2 {num_groups = 8} + : !pto.vmi.vreg<128xf32, + #pto.vmi.layout>, + !pto.vmi.mask<128xb32, + #pto.vmi.layout> + -> !pto.vmi.vreg<8xf32, + #pto.vmi.layout> + +%sum16 = pto.vmi.truncf %sum32 + : !pto.vmi.vreg<8xf32, + #pto.vmi.layout> + -> !pto.vmi.vreg<8xf16, + #pto.vmi.layout> + +%b16 = pto.vmi.group_broadcast %sum16 {num_groups = 8} + : !pto.vmi.vreg<8xf16, + #pto.vmi.layout> + -> !pto.vmi.vreg<128xf16, + #pto.vmi.layout> + +pto.vmi.store %b16, %out[%off] + : !pto.vmi.vreg<128xf16, + #pto.vmi.layout>, !pto.ptr +``` + +The layout without `lane_stride = 2` is illegal: it would claim that f16 group +results are packed in lanes 0..7, but `vcvt` produces them in lanes +0, 2, 4, ..., 14. No lane compaction is performed in this case. + +VPTO lowering result for one full 8-row tile: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%sum_mask = pto.pge_b32 "PAT_VL8" + +%x_lo, %x_hi = pto.vldsx2 %base[%tile_off], "BDINTLV" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%x_lo_sum = pto.vcgadd %x_lo, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%x_hi_sum = pto.vcgadd %x_hi, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +%sum32_block = pto.vadd %x_lo_sum, %x_hi_sum, %sum_mask + : !pto.vreg<64xf32> + +%sum16_block = pto.vcvt %sum32_block, %sum_mask {part = "EVEN", rnd = "R", sat = "SAT"} + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<128xf16> + +%broadcast_idx = compute index vector [0 repeated 16, 2 repeated 16, + 4 repeated 16, 6 repeated 16, + 8 repeated 16, 10 repeated 16, + 12 repeated 16, 14 repeated 16] + : !pto.vreg<128xi16> + +// The index vector uses the lane_stride=2 group-slot source layout. +%b16 = pto.vselr %sum16_block, %broadcast_idx + : !pto.vreg<128xf16>, !pto.vreg<128xi16> -> !pto.vreg<128xf16> + +%all_b16 = pto.pge_b16 "PAT_ALL" +pto.vsts %b16, %out[%off], %all_b16 {dist = "NORM_B16"} + : !pto.vreg<128xf16>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + s32 = reduce(row_r[0..15]) + s16 = truncf(s32) + out[r * 16 + 0 .. r * 16 + 15] = splat(s16) +``` + +Required assignment rule: + +```text +Packed `slots = 8` group-slot `truncf` may be assigned only when the narrowing +result layout records the sub-lane gap: + + f32 group_slots(G, slots=8, lane_stride=1) + -> f16 group_slots(G, slots=8, lane_stride=2) + +Consumers of lane-strided group slots, including `group_broadcast`, must select +source lanes with `(group % slots) * lane_stride`. They must not treat +`slots = 8` as lanes 0..7 after width-changing casts. +``` + +### 3.9 Illegal Dense Consumer Of Group Slots + +VMI input: + +```text +%sum32 = pto.vmi.group_reduce_addf %x, %mask {num_groups = G} +pto.vmi.store %sum32, %out[%off] +``` + +Assigned layouts before the illegal consumer: + +```text +%sum32 : group_slots(G,K) +``` + +Required diagnostic: + +```text +VMI-LAYOUT-CONTRACT: + pto.vmi.store cannot consume #pto.vmi.layout + as a dense vector. Use pto.vmi.group_store, pto.vmi.group_broadcast, or an + explicit group-pack op. +``` + +It must not be diagnosed as: + +```text +dense store materializes group slots implicitly +``` + +That behavior would silently reinterpret a group-slot value as a dense +vector. + +### 3.10 Non-Load Producer Feeding S=32 `group_reduce` + +This case proves that layout assignment is consumer-driven. The producer of the +S=32 input is an elementwise op, not a load. The S=32 `group_reduce` still +requires the elementwise result to be `deinterleaved = 4`, and that requirement +must propagate backward through the elementwise op to both operands. + +VMI input: + +```text +%a = pto.vmi.load %base[%off] + : memref<256xf32> -> !pto.vmi.vreg<256xf32> +%bias = pto.vmi.broadcast %bias_s + : f32 -> !pto.vmi.vreg<256xf32> +%x = pto.vmi.addf %a, %bias +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +pto.vmi.group_store %sum, %out[%group_off], %c1 {num_groups = 8} +``` + +Assigned layouts: + +```text +%a, %bias, %x: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> + +%sum: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> +``` + +VPTO lowering result for one full `8x32xf32` tile: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%sum_mask = pto.pge_b32 "PAT_VL8" + +%a_even_0, %a_odd_0 = pto.vldsx2 %base[%tile_off_0], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%a_even_1, %a_odd_1 = pto.vldsx2 %base[%tile_off_1], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%a_p0, %a_p2 = pto.vdintlv %a_even_0, %a_even_1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%a_p1, %a_p3 = pto.vdintlv %a_odd_0, %a_odd_1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%bias_p0 = pto.vdup %bias_s, %all_b32 : f32, !pto.mask -> !pto.vreg<64xf32> +%bias_p1 = pto.vdup %bias_s, %all_b32 : f32, !pto.mask -> !pto.vreg<64xf32> +%bias_p2 = pto.vdup %bias_s, %all_b32 : f32, !pto.mask -> !pto.vreg<64xf32> +%bias_p3 = pto.vdup %bias_s, %all_b32 : f32, !pto.mask -> !pto.vreg<64xf32> + +%x_p0 = pto.vadd %a_p0, %bias_p0, %all_b32 : !pto.vreg<64xf32> +%x_p1 = pto.vadd %a_p1, %bias_p1, %all_b32 : !pto.vreg<64xf32> +%x_p2 = pto.vadd %a_p2, %bias_p2, %all_b32 : !pto.vreg<64xf32> +%x_p3 = pto.vadd %a_p3, %bias_p3, %all_b32 : !pto.vreg<64xf32> + +%s0 = pto.vcgadd %x_p0, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s1 = pto.vcgadd %x_p1, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s2 = pto.vcgadd %x_p2, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s3 = pto.vcgadd %x_p3, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +%s01 = pto.vadd %s0, %s1, %sum_mask : !pto.vreg<64xf32> +%s23 = pto.vadd %s2, %s3, %sum_mask : !pto.vreg<64xf32> +%sum_block = pto.vadd %s01, %s23, %sum_mask : !pto.vreg<64xf32> + +pto.vsts %sum_block, %out[%group_tile_off], %sum_mask {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + out[group_tile_off + r] = + reduce_i(base[row_r, i] + bias_s for i = 0..31) +``` + +### 3.11 Partial Tail Groups + +Tail handling must be separated by the physical input layout. Row-local S=64 +can avoid inactive rows entirely. Load-fused S=16/S=32 cannot safely do that +with the current `vldsx2` materialization unless the source is known to be +full-tile readable. + +#### 3.11.1 S=64 Active Row Tail + +VMI input: + +```text +%x = pto.vmi.load %base[%off] + : memref<384xf32> -> !pto.vmi.vreg<384xf32> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 6} +%c8 = arith.constant 8 : index +pto.vmi.group_store %sum, %out[%group_off], %c8 {num_groups = 6} +``` + +Assigned layouts: + +```text +%x: + !pto.vmi.vreg<384xf32, #pto.vmi.layout> + +%sum: + !pto.vmi.vreg<384xf32, #pto.vmi.layout> +``` + +VPTO lowering result: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%block8 = pto.pge_b32 "PAT_VL8" +%one_b32 = pto.pge_b32 "PAT_VL1" + +// Emit this row-local block for r = 0..5 only. No load or store is emitted for +// rows 6 and 7. +%x_r = pto.vlds %base[%row_off_r] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<64xf32> +%p_r = pto.vcgadd %x_r, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%sum_r = pto.vcadd %p_r, %block8 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +pto.vsts %sum_r, %out[%group_tile_off_r], %one_b32 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..5: + out[group_tile_off + r * 8] = reduce(row_r[0..63]) +``` + +#### 3.11.2 S=32 Tail Without Full-Tile Read Contract + +VMI input: + +```text +%x = pto.vmi.load %base[%off] + : memref<192xf32> -> !pto.vmi.vreg<192xf32> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 6} +pto.vmi.group_store %sum, %out[%group_off], %c1 {num_groups = 6} +``` + +Assigned layout requested by the consumer: + +```text +%x: + !pto.vmi.vreg<192xf32, #pto.vmi.layout> + +%sum: + !pto.vmi.vreg<192xf32, #pto.vmi.layout> +``` + +Required diagnostic when the source does not carry a full-tile-readable +contract: + +```text +VMI-LAYOUT-CONTRACT: + pto.vmi.group_reduce_addf with group size 32 and num_groups tail 6 requires + materializing #pto.vmi.layout. The fast lowering support + uses vldsx2 DINTLV_B32 over a full 8-row tile. This source is not marked + full-tile-readable, and the stable gather tail fallback is not implemented. +``` + +If a future option enables the stable gather tail fallback, the same VMI input +may lower by gathering only the active lanes. Until that support exists, the +converter must not silently issue the full-tile `vldsx2` loads. + +### 3.12 Control-Flow Join Before `group_reduce` + +The layout carried by a value must survive block arguments. In MLIR converter +terms, the logical VMI value lowered through control flow becomes a tuple of +physical VPTO values with one tuple type per assigned layout. + +VMI input: + +```text +%x = scf.if %cond -> !pto.vmi.vreg<256xf32> { + %a = pto.vmi.load %a_base[%a_off] + : memref<256xf32> -> !pto.vmi.vreg<256xf32> + scf.yield %a : !pto.vmi.vreg<256xf32> +} else { + %b = pto.vmi.load %b_base[%b_off] + : memref<256xf32> -> !pto.vmi.vreg<256xf32> + scf.yield %b : !pto.vmi.vreg<256xf32> +} +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +pto.vmi.group_store %sum, %out[%group_off], %c1 {num_groups = 8} +``` + +Assigned layouts: + +```text +%a, %b, %x: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> + +%sum: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> +``` + +VPTO lowering result for the join: + +```text +%x_p0, %x_p1, %x_p2, %x_p3 = + scf.if %cond + -> (!pto.vreg<64xf32>, !pto.vreg<64xf32>, + !pto.vreg<64xf32>, !pto.vreg<64xf32>) { + %a_even_0, %a_odd_0 = pto.vldsx2 %a_base[%a_tile_off_0], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %a_even_1, %a_odd_1 = pto.vldsx2 %a_base[%a_tile_off_1], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %a_p0, %a_p2 = pto.vdintlv %a_even_0, %a_even_1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %a_p1, %a_p3 = pto.vdintlv %a_odd_0, %a_odd_1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + scf.yield %a_p0, %a_p1, %a_p2, %a_p3 + : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.vreg<64xf32> + } else { + %b_even_0, %b_odd_0 = pto.vldsx2 %b_base[%b_tile_off_0], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %b_even_1, %b_odd_1 = pto.vldsx2 %b_base[%b_tile_off_1], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %b_p0, %b_p2 = pto.vdintlv %b_even_0, %b_even_1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %b_p1, %b_p3 = pto.vdintlv %b_odd_0, %b_odd_1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + scf.yield %b_p0, %b_p1, %b_p2, %b_p3 + : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.vreg<64xf32> + } +``` + +The consumer after the join uses the same S=32 reduction lowering support as +section 3.6: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%sum_mask = pto.pge_b32 "PAT_VL8" + +%s0 = pto.vcgadd %x_p0, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s1 = pto.vcgadd %x_p1, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s2 = pto.vcgadd %x_p2, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s3 = pto.vcgadd %x_p3, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +%s01 = pto.vadd %s0, %s1, %sum_mask : !pto.vreg<64xf32> +%s23 = pto.vadd %s2, %s3, %sum_mask : !pto.vreg<64xf32> +%sum_block = pto.vadd %s01, %s23, %sum_mask : !pto.vreg<64xf32> + +pto.vsts %sum_block, %out[%group_tile_off], %sum_mask {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + selected_row = cond ? a_row_r : b_row_r + out[group_tile_off + r] = reduce(selected_row[0..31]) +``` + +If the two branches cannot be assigned the same layout and no materialization +support exists before `scf.yield`, the required diagnostic is: + +```text +VMI-LAYOUT-CONTRACT: + scf.yield joins incompatible VMI layouts for !pto.vmi.vreg<256xf32>. + Expected #pto.vmi.layout on every incoming value. +``` + +### 3.13 Packed Group-Slot `f32 -> f16` Cast + +This case is intentionally illegal for the current S=16/S=32 packed +group-slot layout. It prevents the compiler from treating a width-changing +`vcvt` as if it preserved low-lane group slots. + +VMI input: + +```text +%x = pto.vmi.load %base[%off] + : memref<128xf32> -> !pto.vmi.vreg<128xf32> +%sum32 = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +%sum16 = pto.vmi.truncf %sum32 +pto.vmi.group_store %sum16, %out[%group_off], %c1 {num_groups = 8} +``` + +Assigned layouts before the illegal cast: + +```text +%x: + !pto.vmi.vreg<128xf32, + #pto.vmi.layout> + +%sum32: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> +``` + +Required diagnostic: + +```text +VMI-LAYOUT-CONTRACT: + pto.vmi.truncf cannot lower from + #pto.vmi.layout f32 to f16 because no + slot-preserving width-changing VPTO support exists. f32->f16 vcvt writes + even/odd sub-lanes, not lanes 0..7. Use group_broadcast before truncf, or + keep the group_store element type as f32. +``` + +This does not contradict section 3.8. Section 3.8 is legal because the cast is +commuted after `group_broadcast`, where the value is dense again. + +### 3.14 Unsupported Group Size + +VMI input: + +```text +%x = pto.vmi.load %base[%off] + : memref<96xf32> -> !pto.vmi.vreg<96xf32> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +pto.vmi.group_store %sum, %out[%group_off], %c1 {num_groups = 8} +``` + +Here `S = 96 / 8 = 12` f32 elements per group. The current VCG-based lowering +support uses 32B groups, i.e. 8 f32 elements per row fragment: + +```text +S = 8 -> one VCGADD block per group +S = 16 -> two 8-lane row fragments, add partial sums +S = 32 -> four 8-lane row fragments, add partial sums +S = 64 -> one full 256B row, VCGADD then VCADD +``` + +Required diagnostic: + +```text +VMI-LAYOUT-CONTRACT: + pto.vmi.group_reduce_addf with f32 group size 12 has no supported VPTO + layout/lowering path. Supported VCG-based f32 group sizes are 8, 16, 32, and 64. + A scalar/gather fallback or a rewrite to logical group size 16 with an + explicit per-group mask is required. +``` + +### 3.15 Compact S=12 Written As Logical S=16 + +If the program wants to use the S=16 lowering for data with 12 semantic f32 +elements per group, the IR must distinguish two sizes: + +```text +logical group size used by VMI ops: 16 +active elements per group: 12 +``` + +The mask is not a prefix mask over the whole vector. It is a per-group mask: + +```text +mask lane i is active iff (i % 16) < 12 +``` + +The group load surface carries the physical source stride as an SSA operand: + +```text +%x = pto.vmi.group_load %base[%off], %source_group_stride + {num_groups = G, group_size = S} + : !pto.ptr, index -> !pto.vmi.vreg +``` + +`source_group_stride` is in elements, not bytes. It is an operand because it may +come from a dynamic leading dimension, a subview, or a runtime tile descriptor. +Static strides use a constant index operand and can be canonicalized later. +`group_size` remains an attribute in this design because it selects the logical +load layout. `active_elems_per_group` belongs to the mask producer, not to the +load. + +Grouped masks use a paired `pto.vmi.create_group_mask` op. It is intentionally +separate from ordinary prefix `pto.vmi.create_mask` so the IR makes group +semantics explicit next to `pto.vmi.group_load` / `pto.vmi.group_reduce_*`: + +```text +%mask = pto.vmi.create_group_mask %active_elems_per_group + {num_groups = G, group_size = S} + : index -> !pto.vmi.mask<(G*S)xpred> +``` + +Semantics: + +```text +lane i is active iff (i % S) < active_elems_per_group +``` + +Current lowering support covers constant `active_elems_per_group`. Dynamic +grouped masks require a runtime lane-index predicate materializer and remain a +separate implementation item. + +Ordinary `pto.vmi.create_mask %active_lanes` keeps the prefix-mask meaning: + +```text +lane i is active iff i < active_lanes +``` + +#### 3.15.1 Existing Design Works If Source Row Stride Is 16 + +If memory already has a 16-f32 row stride, the user can write a logical S=16 +tile and mask off the last four lanes of every group. + +VMI input: + +```text +%stride16 = arith.constant 16 : index +%x = pto.vmi.group_load %base[%off], %stride16 + {num_groups = 8, group_size = 16} + : !pto.ptr, index -> !pto.vmi.vreg<128xf32> +%c12 = arith.constant 12 : index +%mask = pto.vmi.create_group_mask %c12 {num_groups = 8, group_size = 16} + : index -> !pto.vmi.mask<128xpred> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +pto.vmi.group_store %sum, %out[%group_off], %c1 {num_groups = 8} +``` + +Assigned layouts: + +```text +%x: + !pto.vmi.vreg<128xf32, + #pto.vmi.layout> + +%mask: + !pto.vmi.mask<128xpred, + #pto.vmi.layout> + +%sum: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> + +%x32_for_store: + pto.vmi.ensure_layout %x32 + : #pto.vmi.layout -> #pto.vmi.layout +``` + +VPTO lowering result for one `8x16xf32` tile: + +```text +%lo_mask = pto.pge_b32 "PAT_ALL" +%sum_mask = pto.pge_b32 "PAT_VL8" + +%lane = pto.vci %c0_i32 : i32 -> !pto.vreg<64xi32> +%row = pto.vshrs %lane, %c3_i16, %lo_mask + : !pto.vreg<64xi32>, i16, !pto.mask -> !pto.vreg<64xi32> +%row8 = pto.vshls %row, %c3_i16, %lo_mask + : !pto.vreg<64xi32>, i16, !pto.mask -> !pto.vreg<64xi32> +%col = pto.vsub %lane, %row8, %lo_mask + : !pto.vreg<64xi32> +%hi4_mask = pto.vcmps %col, %c4_i32, %lo_mask, "lt" + : !pto.vreg<64xi32>, i32, !pto.mask -> !pto.mask + +%lo, %hi = pto.vldsx2 %base[%tile_off], "BDINTLV" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%lo lanes r*8 .. r*8+7 = row_r[0..7] +%hi lanes r*8 .. r*8+3 = row_r[8..11] +%hi lanes r*8+4 .. r*8+7 = row_r[12..15] // inactive by mask + +%lo_sum = pto.vcgadd %lo, %lo_mask + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%hi_sum = pto.vcgadd %hi, %hi4_mask + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +%sum_block = pto.vadd %lo_sum, %hi_sum, %sum_mask + : !pto.vreg<64xf32> + +pto.vsts %sum_block, %out[%group_tile_off], %sum_mask {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + out[group_tile_off + r] = reduce(row_r[0..11]) +``` + +Design requirement added by this case: VMI mask lowering must support +group-periodic masks by generating the predicate from lane indices. It must not +rewrite this mask to `PAT_M4`: VISA defines `M4` as multiples of 4, not the +first four lanes of each 8-lane block. + +```text +lane = vci(0) +row = lane >> 3 +col = lane - (row << 3) +mask = col < 4 +``` + +#### 3.15.2 Source Row Stride Greater Than 16 + +For now, support the non-compact case where each physical row has at least 16 +f32 slots and the row stride is greater than 16. The fast strided-block path +requires the row stride to be a multiple of one 32B block: + +```text +source_group_stride % 8 == 0 +``` + +The example below uses `source_group_stride = 24`. Each row has 12 semantic +values, 4 masked-but-readable slots, and 8 extra skipped slots: + +```text +row_r[0..11] semantic +row_r[12..15] readable but inactive for the S=16 logical group +row_r[16..23] outside the logical group +``` + +VMI input: + +```text +%stride24 = arith.constant 24 : index +%x = pto.vmi.group_load %base[%off], %stride24 + {num_groups = 8, group_size = 16} + : !pto.ptr, index -> !pto.vmi.vreg<128xf32> +%c12 = arith.constant 12 : index +%mask = pto.vmi.create_group_mask %c12 {num_groups = 8, group_size = 16} + : index -> !pto.vmi.mask<128xpred> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +pto.vmi.group_store %sum, %out[%group_off], %c1 {num_groups = 8} +``` + +Assigned layouts are the same as section 3.15.1: + +```text +%x, %mask: + #pto.vmi.layout +%sum: + #pto.vmi.layout +``` + +VPTO lowering result: + +```text +%lo_mask = pto.pge_b32 "PAT_ALL" +%sum_mask = pto.pge_b32 "PAT_VL8" + +%lane = pto.vci %c0_i32 : i32 -> !pto.vreg<64xi32> +%row = pto.vshrs %lane, %c3_i16, %lo_mask + : !pto.vreg<64xi32>, i16, !pto.mask -> !pto.vreg<64xi32> +%row8 = pto.vshls %row, %c3_i16, %lo_mask + : !pto.vreg<64xi32>, i16, !pto.mask -> !pto.vreg<64xi32> +%col = pto.vsub %lane, %row8, %lo_mask + : !pto.vreg<64xi32> +%hi4_mask = pto.vcmps %col, %c4_i32, %lo_mask, "lt" + : !pto.vreg<64xi32>, i32, !pto.mask -> !pto.mask + +// source_group_stride = 24 f32 = 3 * 32B blocks. +%stride_blocks = %c3_i16 + +%base_lo = %base + tile_off +%base_hi = %base + tile_off + 8 + +%lo = pto.vsldb %base_lo, %stride_blocks, %c0_i16, %lo_mask + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> +%hi = pto.vsldb %base_hi, %stride_blocks, %c0_i16, %lo_mask + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> + +%lo lanes r*8 .. r*8+7 = row_r[0..7] +%hi lanes r*8 .. r*8+7 = row_r[8..15] + +%lo_sum = pto.vcgadd %lo, %lo_mask + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%hi_sum = pto.vcgadd %hi, %hi4_mask + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +%sum_block = pto.vadd %lo_sum, %hi_sum, %sum_mask + : !pto.vreg<64xf32> + +pto.vsts %sum_block, %out[%group_tile_off], %sum_mask {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + out[group_tile_off + r] = + reduce(base[tile_off + r * 24 + 0 .. tile_off + r * 24 + 11]) +``` + +If `source_group_stride > 16` but is not a multiple of 8 f32 elements, this +strided-block path is not legal because `vsldb` block addresses are 32B based. +That case remains unsupported until a gather materialization is selected. + +#### 3.15.3 Compact Source Row Stride 12 + +Compact storage is explicitly out of scope for the first implementation: + +```text +row0[0..11], row1[0..11], row2[0..11], ... +``` + +Required diagnostic: + +```text +VMI-LAYOUT-CONTRACT: + logical group size 16 with active_elems_per_group 12 and + source_group_stride 12 requires compact-row gather materialization. This + plan is not part of the initial VMI layout lowering. +``` + +### 3.16 `group_slot_load` Layout Contract + +`group_slot_load` is separate from `group_load`. + +```text +group_load: + loads group_size data elements per group and produces dense grouped data. + +group_slot_load: + loads one scalar value per group and produces group slots. +``` + +Surface form: + +```text +%v = pto.vmi.group_slot_load %base[%off], %source_group_stride + {num_groups = G} + : !pto.ptr, index -> !pto.vmi.vreg +``` + +Semantics: + +```text +semantic group slot g = base[off + g * source_group_stride] +``` + +The result logical lane count `N` remains the surrounding VMI value shape. Only +the `G` group slots are semantic. Layout assignment chooses the group-slot physical +placement requested by the consumer: + +```text +#pto.vmi.layout +#pto.vmi.layout +``` + +#### 3.16.1 Packed `group_slot_load`, `slots = 8` + +VMI input: + +```text +%rhs = pto.vmi.group_slot_load %rhs_base[%rhs_off], %c1 {num_groups = 8} + : !pto.ptr, index -> !pto.vmi.vreg<128xf32> +pto.vmi.group_store %rhs, %out[%group_off], %c1 {num_groups = 8} +``` + +Assigned layout: + +```text +%rhs: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> +``` + +VPTO lowering result: + +```text +%slot_mask = pto.pge_b32 "PAT_VL8" +%one_block = pto.pge_b32 "PAT_VL1" + +// source_group_stride = 1, so one 32B block contains all 8 scalar group slots. +%rhs_block = pto.vsldb %rhs_base[%rhs_off], %c0_i16, %c0_i16, %one_block + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> + +pto.vsts %rhs_block, %out[%group_off], %slot_mask {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for g = 0..7: + out[group_off + g] = rhs_base[rhs_off + g] +``` + +If `source_group_stride != 1`, this packed `slots = 8` layout requires a +strided/gather group-slot load materializer. Until that support exists, +`group_slot_load` with `slots = 8` and non-unit stride must diagnose instead of +silently using full-group `group_load`. + +#### 3.16.2 Row-Local `group_slot_load`, `slots = 1` + +VMI input: + +```text +%c8 = arith.constant 8 : index +%rhs = pto.vmi.group_slot_load %rhs_base[%rhs_off], %c8 {num_groups = 8} + : !pto.ptr, index -> !pto.vmi.vreg<512xf32> +pto.vmi.group_store %rhs, %out[%group_off], %c8 {num_groups = 8} +``` + +Assigned layout: + +```text +%rhs: + !pto.vmi.vreg<512xf32, #pto.vmi.layout> +``` + +VPTO lowering result: + +```text +%one_b32 = pto.pge_b32 "PAT_VL1" + +// Emit this shape for r = 0..7. Each result value carries one semantic slot +// in lane 0, matching the S=64 row-local group_reduce result layout. +// For f32, source_group_stride = 8 elements = 32B, so every lane-0 vsldb is +// aligned. +%rhs_r = pto.vsldb %rhs_base[%rhs_off_plus_r], %c0_i16, %c0_i16, %one_b32 + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> + +pto.vsts %rhs_r, %out[%group_off_plus_r], %one_b32 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + out[group_off + r * 8] = rhs_base[rhs_off + r * 8] +``` + +Current lowering rule: + +```text +slots = 1 group_slot_load uses one lane-0 vsldb per semantic group slot. +For f32, source_group_stride must be a positive constant divisible by 8 +elements. For f16 it must be divisible by 16 elements, and for f8 it must be +divisible by 32 elements. +``` + +### 3.17 `group_broadcast` Feeding A Deinterleaved Consumer + +This case fixes a lowering invariant: `group_broadcast` itself does not infer a +consumer-specific deinterleaved result. It produces the layout selected by +layout assignment. If a later consumer requires another layout, assignment must +insert an explicit `ensure_layout`. + +The current endpoint is: + +```text +group_reduce -> group_broadcast(contiguous f32) + -> ensure_layout(deinterleaved = 2) + -> truncf(contiguous f16) +``` + +VMI input: + +```text +%x = pto.vmi.load %base[%off] + : memref<128xf32> -> !pto.vmi.vreg<128xf32> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +%b = pto.vmi.group_broadcast %sum {num_groups = 8} +%h = pto.vmi.truncf %b +pto.vmi.store %h, %out[%off] +``` + +Assigned layouts: + +```text +%x: + !pto.vmi.vreg<128xf32, + #pto.vmi.layout> + +%sum: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> + +%b_dense: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> + +%b_split = pto.vmi.ensure_layout %b_dense: + #pto.vmi.layout + -> #pto.vmi.layout + +%h: + !pto.vmi.vreg<128xf16, #pto.vmi.layout> +``` + +VPTO lowering result: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%sum_mask = pto.pge_b32 "PAT_VL8" + +%x_lo, %x_hi = pto.vldsx2 %base[%tile_off], "BDINTLV" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%lo_sum = pto.vcgadd %x_lo, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%hi_sum = pto.vcgadd %x_hi, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%sum_block = pto.vadd %lo_sum, %hi_sum, %sum_mask + : !pto.vreg<64xf32> + +// group_broadcast lowers to two contiguous f32 chunks. +%idx_lo = materialize indices [0 repeated 16, 1 repeated 16, + 2 repeated 16, 3 repeated 16] + : !pto.vreg<64xi32> +%idx_hi = materialize indices [4 repeated 16, 5 repeated 16, + 6 repeated 16, 7 repeated 16] + : !pto.vreg<64xi32> + +%b_lo = pto.vselr %sum_block, %idx_lo + : !pto.vreg<64xf32>, !pto.vreg<64xi32> -> !pto.vreg<64xf32> +%b_hi = pto.vselr %sum_block, %idx_hi + : !pto.vreg<64xf32>, !pto.vreg<64xi32> -> !pto.vreg<64xf32> + +// ensure_layout contiguous -> deinterleaved=2 is explicit in assigned VMI. +%b_even_input, %b_odd_input = pto.vdintlv %b_lo, %b_hi + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%h_even = pto.vcvt %b_even_input, %all_b32 {part = "EVEN", rnd = "R", sat = "SAT"} + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<128xf16> +%h_odd = pto.vcvt %b_odd_input, %all_b32 {part = "ODD", rnd = "R", sat = "SAT"} + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<128xf16> + +%all_b16 = pto.pge_b16 "PAT_ALL" +%h0 = pto.vor %h_even, %h_odd, %all_b16 + : !pto.vreg<128xf16> + +pto.vsts %h0, %out[%off], %all_b16 {dist = "NORM_B16"} + : !pto.vreg<128xf16>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + s = reduce(row_r[0..15]) + out[r * 16 + 0 .. r * 16 + 15] = truncf(s) +``` + +Required assignment rule: + +```text +`group_broadcast` layout is chosen before `vmi-to-vpto`. A width-changing +consumer such as `truncf` may require a deinterleaved f32 source, but that +requirement must be represented by `ensure_layout`; `truncf` lowering must not +look through the defining `group_broadcast` and choose a hidden broadcast shape. +``` + +### 3.18 One Value With Dense And Group-Reduce Consumers + +This case forces layout assignment to handle a solvable use-site conflict. One +consumer requires an S=32 group-reduce layout; another consumer requires dense +row-major store. This is not semantically illegal. It must be solved by +explicit use-site materialization. A later optimization pass may fold the +materialization into a store or rematerialize a cheap producer when the required +support exists. + +VMI input: + +```text +%x = pto.vmi.load %base[%off] + : memref<256xf32> -> !pto.vmi.vreg<256xf32> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +pto.vmi.group_store %sum, %sum_out[%group_off], %c1 {num_groups = 8} +pto.vmi.store %x, %copy_out[%off] +``` + +Assigned layouts: + +```text +%x for group_reduce: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> + +%sum: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> + +%x for dense store: + requires #pto.vmi.layout +``` + +Baseline layout assignment keeps `%x` in the group-reduce layout and inserts +`ensure_layout` before the dense store use. A later rematerialization pass may +clone the load for the dense store if that is profitable. A later fold-consumer +pass may also fold `ensure_layout + store` into a layout-aware store lowering. + +VPTO lowering result: + +```text +%x_even_0, %x_odd_0 = pto.vldsx2 %base[%tile_off_0], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%x_even_1, %x_odd_1 = pto.vldsx2 %base[%tile_off_1], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%x_p0, %x_p2 = pto.vdintlv %x_even_0, %x_even_1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%x_p1, %x_p3 = pto.vdintlv %x_odd_0, %x_odd_1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%all_b32 = pto.pge_b32 "PAT_ALL" +%sum_mask = pto.pge_b32 "PAT_VL8" + +%s0 = pto.vcgadd %x_p0, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s1 = pto.vcgadd %x_p1, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s2 = pto.vcgadd %x_p2, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s3 = pto.vcgadd %x_p3, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +%s01 = pto.vadd %s0, %s1, %sum_mask : !pto.vreg<64xf32> +%s23 = pto.vadd %s2, %s3, %sum_mask : !pto.vreg<64xf32> +%sum_block = pto.vadd %s01, %s23, %sum_mask : !pto.vreg<64xf32> + +pto.vsts %sum_block, %sum_out[%group_off], %sum_mask {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + +// Dense store materialization for the second consumer. +%even0, %even1 = pto.vintlv %x_p0, %x_p2 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%odd0, %odd1 = pto.vintlv %x_p1, %x_p3 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%d0, %d1 = pto.vintlv %even0, %odd0 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%d2, %d3 = pto.vintlv %even1, %odd1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +pto.vsts %d0, %copy_out[%off_0], %all_b32 {dist = "NORM_B32"} +pto.vsts %d1, %copy_out[%off_64], %all_b32 {dist = "NORM_B32"} +pto.vsts %d2, %copy_out[%off_128], %all_b32 {dist = "NORM_B32"} +pto.vsts %d3, %copy_out[%off_192], %all_b32 {dist = "NORM_B32"} +``` + +Memory result: + +```text +for r = 0..7: + sum_out[group_off + r] = reduce(row_r[0..31]) + +for i = 0..255: + copy_out[off + i] = base[off + i] +``` + +If `deinterleaved = 4 -> contiguous` materialization support does not exist, the +required diagnostic is: + +```text +VMI-LAYOUT-CONTRACT: + value %x is required as #pto.vmi.layout by + pto.vmi.group_reduce_addf and as #pto.vmi.layout by + pto.vmi.store, but no materialization support exists at the store use site. +``` + +### 3.19 S=16 Reduce `block_elems` Support Selection + +S=16 f32 group reduction has two legal dense input layouts: + +```text +#pto.vmi.layout +#pto.vmi.layout +``` + +`block_elems = 1` is the element-parity layout required by f32->f16 `truncf`. +It is also a valid S=16 reduction layout: each physical part contains eight +values per row, so `VCGADD` can reduce each part and `VADD` can combine the two +partial sums. + +`block_elems = 8` is still useful when the producer is a block load shape such +as `BDINTLV` or `vsldb` over 32B row fragments. Baseline layout assignment must +express any mismatch with an explicit `ensure_layout`; producer rematerialization +or consumer folding can choose the cheaper equivalent form later. Assignment +must not hard-code S=16 reduce to `block_elems = 8`. + +#### 3.19.1 Continuous S=16 Reduce And Truncf, `block_elems = 1` + +VMI input: + +```text +%x = pto.vmi.load %base[%off] + : memref<128xf32> -> !pto.vmi.vreg<128xf32> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +pto.vmi.group_store %sum, %sum_out[%group_off], %c1 {num_groups = 8} +%h = pto.vmi.truncf %x + : !pto.vmi.vreg<128xf32> -> !pto.vmi.vreg<128xf16> +pto.vmi.store %h, %out[%off] +``` + +Assigned layouts: + +```text +%x: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> + +%sum: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> + +%h: + !pto.vmi.vreg<128xf16, #pto.vmi.layout> +``` + +Physical lane map: + +```text +%x_p0 lanes r*8 .. r*8+7 = + row_r[0], row_r[2], row_r[4], ..., row_r[14] + +%x_p1 lanes r*8 .. r*8+7 = + row_r[1], row_r[3], row_r[5], ..., row_r[15] +``` + +VPTO lowering result: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%sum_mask = pto.pge_b32 "PAT_VL8" + +%x_p0, %x_p1 = pto.vldsx2 %base[%tile_off], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%s0 = pto.vcgadd %x_p0, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s1 = pto.vcgadd %x_p1, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%sum_block = pto.vadd %s0, %s1, %sum_mask + : !pto.vreg<64xf32> + +pto.vsts %sum_block, %sum_out[%group_off], %sum_mask {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + +%h_even = pto.vcvt %x_p0, %all_b32 {part = "EVEN", rnd = "R", sat = "SAT"} + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<128xf16> +%h_odd = pto.vcvt %x_p1, %all_b32 {part = "ODD", rnd = "R", sat = "SAT"} + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<128xf16> + +%all_b16 = pto.pge_b16 "PAT_ALL" +%h0 = pto.vor %h_even, %h_odd, %all_b16 + : !pto.vreg<128xf16> +pto.vsts %h0, %out[%off], %all_b16 {dist = "NORM_B16"} + : !pto.vreg<128xf16>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + sum_out[group_off + r] = reduce(row_r[0..15]) + +for i = 0..127: + out[off + i] = truncf(base[off + i]) +``` + +#### 3.19.2 Block-Load Producer Fixed To `block_elems = 8` + +This is the real conflict case. The value is fixed to `block_elems = 8` +because the producer uses block-load support. A later `truncf` +requires element-parity `block_elems = 1`. + +VMI input: + +```text +%stride24 = arith.constant 24 : index +%x = pto.vmi.group_load %base[%off], %stride24 + {num_groups = 8, group_size = 16} + : !pto.ptr, index -> !pto.vmi.vreg<128xf32> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +pto.vmi.group_store %sum, %sum_out[%group_off], %c1 {num_groups = 8} +%h = pto.vmi.truncf %x + : !pto.vmi.vreg<128xf32> -> !pto.vmi.vreg<128xf16> +pto.vmi.store %h, %out[%off] +``` + +Assigned layouts before the conflicting `truncf` use: + +```text +%x from strided block group_load: + !pto.vmi.vreg<128xf32, + #pto.vmi.layout> + +%sum: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> +``` + +The reduction path is legal and uses the same `vsldb` block-load shape as +section 3.15.2. The `truncf` path is legal only if one of these transforms +exists: + +```text +1. rematerialize the original memory producer as block_elems=1 +2. materialize block_elems=8 -> block_elems=1 in registers +3. use an explicitly enabled scratch/reload fallback +``` + +If no such transform exists, the required diagnostic is: + +```text +VMI-LAYOUT-CONTRACT: + pto.vmi.truncf requires + #pto.vmi.layout, but the source value is + fixed to #pto.vmi.layout by the strided + group_load. Add rematerialization or preserving materialization support, or + avoid consuming this block-loaded value with truncf. +``` + +### 3.20 `group_slots` Control-Flow Join + +`group_slots` values must be allowed to cross control flow. The join type is a +group-slot physical tuple, not a dense vector. + +VMI input: + +```text +%sum = scf.if %cond -> !pto.vmi.vreg<128xf32> { + %x = pto.vmi.load %base[%off] + : memref<128xf32> -> !pto.vmi.vreg<128xf32> + %a = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} + scf.yield %a : !pto.vmi.vreg<128xf32> +} else { + %b = pto.vmi.group_slot_load %rhs_base[%rhs_off], %c1 {num_groups = 8} + : !pto.ptr, index -> !pto.vmi.vreg<128xf32> + scf.yield %b : !pto.vmi.vreg<128xf32> +} +%bias = pto.vmi.group_slot_load %bias_base[%bias_off], %c1 {num_groups = 8} + : !pto.ptr, index -> !pto.vmi.vreg<128xf32> +%outv = pto.vmi.addf %sum, %bias +pto.vmi.group_store %outv, %out[%group_off], %c1 {num_groups = 8} +``` + +Assigned layouts: + +```text +%x: + !pto.vmi.vreg<128xf32, + #pto.vmi.layout> + +%a, %b, %sum, %bias, %outv: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> +``` + +VPTO lowering result for the join: + +```text +%sum_block = scf.if %cond -> !pto.vreg<64xf32> { + %all_b32 = pto.pge_b32 "PAT_ALL" + %sum_mask = pto.pge_b32 "PAT_VL8" + + %x_lo, %x_hi = pto.vldsx2 %base[%tile_off], "BDINTLV" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %lo_sum = pto.vcgadd %x_lo, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %hi_sum = pto.vcgadd %x_hi, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %a_block = pto.vadd %lo_sum, %hi_sum, %sum_mask + : !pto.vreg<64xf32> + scf.yield %a_block : !pto.vreg<64xf32> +} else { + %one_block = pto.pge_b32 "PAT_VL1" + %b_block = pto.vsldb %rhs_base[%rhs_off], %c0_i16, %c0_i16, %one_block + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> + scf.yield %b_block : !pto.vreg<64xf32> +} + +%one_block = pto.pge_b32 "PAT_VL1" +%slot_mask = pto.pge_b32 "PAT_VL8" +%bias_block = pto.vsldb %bias_base[%bias_off], %c0_i16, %c0_i16, %one_block + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> +%out_block = pto.vadd %sum_block, %bias_block, %slot_mask + : !pto.vreg<64xf32> + +pto.vsts %out_block, %out[%group_off], %slot_mask {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + lhs = cond ? reduce(row_r[0..15]) : rhs_base[rhs_off + r] + out[group_off + r] = lhs + bias_base[bias_off + r] +``` + +### 3.21 S=32 Tail With Full-Tile-Readable Source + +This is the positive counterpart to section 3.11.2. Tail participation is +still expressed by masks, but the source must provide a static proof that +reading the rounded-up 8-row physical tile is memory-safe. That proof is +explicit for partial logical loads: it can come from a statically shaped memref +source. Pointer-source runtime kernels should instead load the rounded physical +vector and use a mask to express active logical lanes; this is not inferred from +surrounding MTE copies or caller context. + +VMI input: + +```text +%x = pto.vmi.load %base[%off] + : memref<256xf32> -> !pto.vmi.vreg<192xf32> +%mask = pto.vmi.create_mask %c192 : index -> !pto.vmi.mask<192xpred> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 6} +pto.vmi.group_store %sum, %out[%group_off], %c1 {num_groups = 6} +``` + +Equivalent pointer-source VMI input for runtime kernels: + +```text +%x = pto.vmi.load %base[%off] + : !pto.ptr -> !pto.vmi.vreg<256xf32> +%mask = pto.vmi.create_mask %c192 : index -> !pto.vmi.mask<256xpred> +``` + +Assigned layouts: + +```text +%x: + !pto.vmi.vreg<192xf32, #pto.vmi.layout> + +%mask: + !pto.vmi.mask<192xpred, + #pto.vmi.layout> + +%sum: + !pto.vmi.vreg<192xf32, #pto.vmi.layout> +``` + +VPTO lowering result: + +```text +// A statically safe full-read proof allows the load plan to read the +// rounded-up 8-row tile. Only rows 0..5 are semantically active. +%x_c0 = pto.vlds %base[%tile_off_0] + : memref<256xf32> -> !pto.vreg<64xf32> +%x_c1 = pto.vlds %base[%tile_off_1] + : memref<256xf32> -> !pto.vreg<64xf32> +%x_c2 = pto.vlds %base[%tile_off_2] + : memref<256xf32> -> !pto.vreg<64xf32> +%x_c3 = pto.vlds %base[%tile_off_3] + : memref<256xf32> -> !pto.vreg<64xf32> + +%x_lo01, %x_hi01 = pto.vdintlv %x_c0, %x_c1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%x_lo23, %x_hi23 = pto.vdintlv %x_c2, %x_c3 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%x_p0, %x_p2 = pto.vdintlv %x_lo01, %x_lo23 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%x_p1, %x_p3 = pto.vdintlv %x_hi01, %x_hi23 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%data_mask0, %_ = pto.plt_b32 %c48_i32 + : i32 -> !pto.mask, i32 +%data_mask1, %_ = pto.plt_b32 %c48_i32 + : i32 -> !pto.mask, i32 +%data_mask2, %_ = pto.plt_b32 %c48_i32 + : i32 -> !pto.mask, i32 +%data_mask3, %_ = pto.plt_b32 %c48_i32 + : i32 -> !pto.mask, i32 +%sum_mask, %_ = pto.plt_b32 %c6_i32 + : i32 -> !pto.mask, i32 + +%s0 = pto.vcgadd %x_p0, %data_mask0 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s1 = pto.vcgadd %x_p1, %data_mask1 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s2 = pto.vcgadd %x_p2, %data_mask2 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s3 = pto.vcgadd %x_p3, %data_mask3 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +%s01 = pto.vadd %s0, %s1, %sum_mask : !pto.vreg<64xf32> +%s23 = pto.vadd %s2, %s3, %sum_mask : !pto.vreg<64xf32> +%sum_block = pto.vadd %s01, %s23, %sum_mask : !pto.vreg<64xf32> + +pto.vsts %sum_block, %out[%group_off], %sum_mask {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..5: + out[group_off + r] = reduce(row_r[0..31]) +``` + +Rows 6 and 7 may be physically loaded because of the safe full-read proof, but +their lanes are not active in `%data_mask*`, and their group slots are not +stored because `%sum_mask` is produced by `plt_b32 %c6_i32`. + +### 3.22 `scf.for` Loop-Carried Layout + +Loop-carried VMI values require a layout fixed point. The iter_arg, body block +argument, yield operand, loop result, and later consumer must all agree on one +layout, or `vmi-layout-assignment` must insert a materialization at a legal +dominating use site. + +VMI input: + +```text +%init = pto.vmi.load %base[%off] + : memref<256xf32> -> !pto.vmi.vreg<256xf32> +%acc = scf.for %i = %c0 to %steps step %c1 + iter_args(%arg = %init) -> !pto.vmi.vreg<256xf32> { + %bias = pto.vmi.broadcast %bias_s + : f32 -> !pto.vmi.vreg<256xf32> + %next = pto.vmi.addf %arg, %bias + scf.yield %next : !pto.vmi.vreg<256xf32> +} +%sum = pto.vmi.group_reduce_addf %acc, %mask {num_groups = 8} +pto.vmi.group_store %sum, %out[%group_off], %c1 {num_groups = 8} +``` + +Assigned layouts: + +```text +%init, %arg, %bias, %next, %acc: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> + +%sum: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> +``` + +VPTO lowering result: + +```text +%init_even_0, %init_odd_0 = pto.vldsx2 %base[%tile_off_0], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%init_even_1, %init_odd_1 = pto.vldsx2 %base[%tile_off_1], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%init_p0, %init_p2 = pto.vdintlv %init_even_0, %init_even_1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%init_p1, %init_p3 = pto.vdintlv %init_odd_0, %init_odd_1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%acc_p0, %acc_p1, %acc_p2, %acc_p3 = + scf.for %i = %c0 to %steps step %c1 + iter_args(%arg_p0 = %init_p0, %arg_p1 = %init_p1, + %arg_p2 = %init_p2, %arg_p3 = %init_p3) + -> (!pto.vreg<64xf32>, !pto.vreg<64xf32>, + !pto.vreg<64xf32>, !pto.vreg<64xf32>) { + %all_b32 = pto.pge_b32 "PAT_ALL" + %bias_p0 = pto.vdup %bias_s, %all_b32 + : f32, !pto.mask -> !pto.vreg<64xf32> + %bias_p1 = pto.vdup %bias_s, %all_b32 + : f32, !pto.mask -> !pto.vreg<64xf32> + %bias_p2 = pto.vdup %bias_s, %all_b32 + : f32, !pto.mask -> !pto.vreg<64xf32> + %bias_p3 = pto.vdup %bias_s, %all_b32 + : f32, !pto.mask -> !pto.vreg<64xf32> + + %next_p0 = pto.vadd %arg_p0, %bias_p0, %all_b32 : !pto.vreg<64xf32> + %next_p1 = pto.vadd %arg_p1, %bias_p1, %all_b32 : !pto.vreg<64xf32> + %next_p2 = pto.vadd %arg_p2, %bias_p2, %all_b32 : !pto.vreg<64xf32> + %next_p3 = pto.vadd %arg_p3, %bias_p3, %all_b32 : !pto.vreg<64xf32> + scf.yield %next_p0, %next_p1, %next_p2, %next_p3 + : !pto.vreg<64xf32>, !pto.vreg<64xf32>, + !pto.vreg<64xf32>, !pto.vreg<64xf32> + } + +%all_b32 = pto.pge_b32 "PAT_ALL" +%sum_mask = pto.pge_b32 "PAT_VL8" +%s0 = pto.vcgadd %acc_p0, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s1 = pto.vcgadd %acc_p1, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s2 = pto.vcgadd %acc_p2, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s3 = pto.vcgadd %acc_p3, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s01 = pto.vadd %s0, %s1, %sum_mask : !pto.vreg<64xf32> +%s23 = pto.vadd %s2, %s3, %sum_mask : !pto.vreg<64xf32> +%sum_block = pto.vadd %s01, %s23, %sum_mask : !pto.vreg<64xf32> +pto.vsts %sum_block, %out[%group_off], %sum_mask {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + for c = 0..31: + acc[row_r, c] = base[row_r, c] + steps * bias_s + out[group_off + r] = reduce(acc[row_r, 0..31]) +``` + +### 3.23 `group_broadcast` With Multiple Dense Consumers + +One `group_slots` value may feed multiple `group_broadcast` uses with different +dense result layout requirements. Each `group_broadcast` op has its own result +layout, so layout assignment should type each op at its use site instead of +forcing one result layout onto all consumers. + +VMI input: + +```text +%x = pto.vmi.load %base[%off] + : memref<128xf32> -> !pto.vmi.vreg<128xf32> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} + +%b_for_mul = pto.vmi.group_broadcast %sum {num_groups = 8} +%y = pto.vmi.mulf %x, %b_for_mul +%ysum = pto.vmi.group_reduce_addf %y, %mask {num_groups = 8} +pto.vmi.group_store %ysum, %sum_out[%group_off], %c1 {num_groups = 8} + +%b_for_cast = pto.vmi.group_broadcast %sum {num_groups = 8} +%h = pto.vmi.truncf %b_for_cast +pto.vmi.store %h, %dense_out[%off] +``` + +Assigned layouts in the current implementation: + +```text +%x: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> + +%x_for_reduce: + !pto.vmi.vreg<128xf32, + #pto.vmi.layout> + +%sum, %ysum: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> + +%b_for_mul, %y: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> + +%y_for_reduce: + !pto.vmi.vreg<128xf32, + #pto.vmi.layout> + +%b_for_cast: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> + +%b_for_cast_split: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> + +%h: + !pto.vmi.vreg<128xf16, #pto.vmi.layout> +``` + +The important invariant is not that both dense consumers choose the same dense +layout. It is that each use has an explicit layout boundary: + +```text +%x_for_reduce = pto.vmi.ensure_layout %x +%y_for_reduce = pto.vmi.ensure_layout %y +%b_for_cast_split = pto.vmi.ensure_layout %b_for_cast +``` + +If a future direct `group_broadcast -> deinterleaved` support path is added, layout +assignment may assign `%b_for_mul` or `%b_for_cast` directly to that layout, but +the choice must still be visible in the assigned IR. + +VPTO lowering result: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%sum_mask = pto.pge_b32 "PAT_VL8" + +%x_lo, %x_hi = pto.vldsx2 %base[%tile_off], "BDINTLV" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%x_lo_sum = pto.vcgadd %x_lo, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%x_hi_sum = pto.vcgadd %x_hi, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%sum_block = pto.vadd %x_lo_sum, %x_hi_sum, %sum_mask + : !pto.vreg<64xf32> + +%lane_id = pto.vci %c0_i32 : i32 -> !pto.vreg<64xi32> +%broadcast_idx = pto.vshrs %lane_id, %c3_i16, %all_b32 + : !pto.vreg<64xi32>, i16, !pto.mask -> !pto.vreg<64xi32> + +// Use 1: broadcast for the multiply path. Current lowering materializes two +// contiguous f32 chunks, multiplies them with the original contiguous chunks, +// then deinterleaves the product for the second group_reduce. +%b_rows_for_mul_0 = pto.vselr %sum_block, %broadcast_idx_0 + : !pto.vreg<64xf32>, !pto.vreg<64xi32> -> !pto.vreg<64xf32> +%b_rows_for_mul_1 = pto.vselr %sum_block, %broadcast_idx_1 + : !pto.vreg<64xf32>, !pto.vreg<64xi32> -> !pto.vreg<64xf32> +%y0 = pto.vmul %x0, %b_rows_for_mul_0, %all_b32 : !pto.vreg<64xf32> +%y1 = pto.vmul %x1, %b_rows_for_mul_1, %all_b32 : !pto.vreg<64xf32> +%y_lo, %y_hi = pto.vdintlv %y0, %y1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%y_lo_sum = pto.vcgadd %y_lo, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%y_hi_sum = pto.vcgadd %y_hi, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%ysum_block = pto.vadd %y_lo_sum, %y_hi_sum, %sum_mask + : !pto.vreg<64xf32> +pto.vsts %ysum_block, %sum_out[%group_off], %sum_mask {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + +// Use 2: rematerialize broadcast for the f32->f16 parity cast path. +%b_rows_for_cast_0 = pto.vselr %sum_block, %broadcast_idx_0 + : !pto.vreg<64xf32>, !pto.vreg<64xi32> -> !pto.vreg<64xf32> +%b_rows_for_cast_1 = pto.vselr %sum_block, %broadcast_idx_1 + : !pto.vreg<64xf32>, !pto.vreg<64xi32> -> !pto.vreg<64xf32> +%cast_lo, %cast_hi = pto.vdintlv %b_rows_for_cast_0, %b_rows_for_cast_1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%h_even = pto.vcvt %cast_lo, %all_b32 + {part = "EVEN", rnd = "R", sat = "SAT"} + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<128xf16> +%h_odd = pto.vcvt %cast_hi, %all_b32 + {part = "ODD", rnd = "R", sat = "SAT"} + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<128xf16> +%all_b16 = pto.pge_b16 "PAT_ALL" +%h0 = pto.vor %h_even, %h_odd, %all_b16 : !pto.vreg<128xf16> +pto.vsts %h0, %dense_out[%off], %all_b16 {dist = "NORM_B16"} + : !pto.vreg<128xf16>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + s = reduce(row_r[0..15]) + sum_out[group_off + r] = reduce_i(row_r[i] * s for i = 0..15) + dense_out[r * 16 + 0 .. r * 16 + 15] = truncf(s) +``` + +### 3.24 Mask With Elementwise, Select, And Store + +This case separates compute masking from memory effects. A masked elementwise +operation with passthrough semantics can be represented as ordinary compute +plus `select`; a masked store uses the mask only on the store effect. + +VMI input: + +```text +%x = pto.vmi.load %base[%off] + : memref<64xf32> -> !pto.vmi.vreg<64xf32> +%rhs = pto.vmi.load %rhs_base[%off] + : memref<64xf32> -> !pto.vmi.vreg<64xf32> +%mask = pto.vmi.create_mask %c48 + : index -> !pto.vmi.mask<64xpred> +%sum = pto.vmi.addf %x, %rhs +%passthrough = pto.vmi.select %mask, %sum, %x +pto.vmi.store %passthrough, %dense_out[%off] +pto.vmi.masked_store %sum, %masked_out[%off], %mask +``` + +Assigned layouts: + +```text +%x, %rhs, %sum, %passthrough: + !pto.vmi.vreg<64xf32, #pto.vmi.layout> + +%mask: + !pto.vmi.mask<64xpred, #pto.vmi.layout> +``` + +VPTO lowering result: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%m, %_ = pto.plt_b32 %c48_i32 : i32 -> !pto.mask, i32 + +%x0 = pto.vlds %base[%off] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<64xf32> +%rhs0 = pto.vlds %rhs_base[%off] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<64xf32> +%sum0 = pto.vadd %x0, %rhs0, %all_b32 : !pto.vreg<64xf32> + +%pass0 = pto.vsel %sum0, %x0, %m + : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +pto.vsts %pass0, %dense_out[%off], %all_b32 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + +pto.vsts %sum0, %masked_out[%off], %m {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for i = 0..63: + if i < 48: + dense_out[off + i] = base[off + i] + rhs_base[off + i] + masked_out[off + i] = base[off + i] + rhs_base[off + i] + else: + dense_out[off + i] = base[off + i] + masked_out[off + i] is unchanged +``` + +### 3.25 Function Boundary Layout Specialization + +Function boundaries cannot rely on hidden layout side tables. Either the +function is internal and layout-specialized by `vmi-layout-assignment`, or a +public/external VMI boundary must diagnose until a stable VMI ABI is defined. + +#### 3.25.1 Internal Function Specialized To Consumer Layout + +VMI input: + +```text +func.func private @producer(%base: !pto.ptr, %off: index) + -> !pto.vmi.vreg<256xf32> { + %x = pto.vmi.load %base[%off] + : memref<256xf32> -> !pto.vmi.vreg<256xf32> + return %x : !pto.vmi.vreg<256xf32> +} + +func.func @caller(%base: !pto.ptr, %off: index, %out: !pto.ptr) { + %x = call @producer(%base, %off) + : (!pto.ptr, index) -> !pto.vmi.vreg<256xf32> + %sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} + pto.vmi.group_store %sum, %out[%off], %c1 {num_groups = 8} + return +} +``` + +Assigned layouts: + +```text +@producer result: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> + +%x in @caller: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> + +%sum: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> +``` + +VPTO lowering result for the function boundary: + +```text +func.func private @producer(...) + -> (!pto.vreg<64xf32>, !pto.vreg<64xf32>, + !pto.vreg<64xf32>, !pto.vreg<64xf32>) { + %x_even_0, %x_odd_0 = pto.vldsx2 %base[%tile_off_0], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %x_even_1, %x_odd_1 = pto.vldsx2 %base[%tile_off_1], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %x_p0, %x_p2 = pto.vdintlv %x_even_0, %x_even_1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %x_p1, %x_p3 = pto.vdintlv %x_odd_0, %x_odd_1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + return %x_p0, %x_p1, %x_p2, %x_p3 + : !pto.vreg<64xf32>, !pto.vreg<64xf32>, + !pto.vreg<64xf32>, !pto.vreg<64xf32> +} + +func.func @caller(...) { + %x_p0, %x_p1, %x_p2, %x_p3 = call @producer(...) + : (...) -> (!pto.vreg<64xf32>, !pto.vreg<64xf32>, + !pto.vreg<64xf32>, !pto.vreg<64xf32>) + + %all_b32 = pto.pge_b32 "PAT_ALL" + %sum_mask = pto.pge_b32 "PAT_VL8" + %s0 = pto.vcgadd %x_p0, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %s1 = pto.vcgadd %x_p1, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %s2 = pto.vcgadd %x_p2, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %s3 = pto.vcgadd %x_p3, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %s01 = pto.vadd %s0, %s1, %sum_mask : !pto.vreg<64xf32> + %s23 = pto.vadd %s2, %s3, %sum_mask : !pto.vreg<64xf32> + %sum_block = pto.vadd %s01, %s23, %sum_mask : !pto.vreg<64xf32> + pto.vsts %sum_block, %out[%group_off], %sum_mask {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +} +``` + +Memory result: + +```text +for r = 0..7: + out[off + r] = reduce(row_r[0..31]) +``` + +Runtime closure: + +```text +lit: + test/lit/vmi/vmi_ptoas_private_call_inline.pto + +runtime SIM: + test/vpto/cases/vmi/private-call-inline-store + +ptoas pipeline: + vmi-layout-assignment makes the private result layout explicit + vmi-to-vpto physicalizes the private helper result into !pto.vreg values + ptoas then inlines private physical VMI helpers before VPTO vecscope/backend + emission, so physical vector values do not escape through a function return +``` + +#### 3.25.2 Public Or External VMI Boundary + +VMI input: + +```text +func.func @public_producer(%base: !pto.ptr, %off: index) + -> !pto.vmi.vreg<256xf32> attributes {public} { + %x = pto.vmi.load %base[%off] + : memref<256xf32> -> !pto.vmi.vreg<256xf32> + return %x : !pto.vmi.vreg<256xf32> +} +``` + +Required diagnostic for the initial design: + +```text +VMI-LAYOUT-CONTRACT: + public or external function boundary returns !pto.vmi.vreg<256xf32> without a + stable VMI layout ABI. Mark the function internal for layout specialization, + inline it before vmi-layout-assignment, or define an explicit ABI layout. +``` + +### 3.26 S=16 Grouped Tail Through Broadcast, Reduce, Store + +This case extends section 3.15.1 from `reduce -> group_store` to the full +grouped compute path. It is needed because `create_group_mask` must remain a +group-periodic mask after a `group_broadcast`; it cannot collapse to a prefix +mask or an all-true mask. + +VMI input: + +```text +%stride16 = arith.constant 16 : index +%x = pto.vmi.group_load %base[%off], %stride16 + {num_groups = 8, group_size = 16} + : !pto.ptr, index -> !pto.vmi.vreg<128xf32> +%c12 = arith.constant 12 : index +%mask = pto.vmi.create_group_mask %c12 {num_groups = 8, group_size = 16} + : index -> !pto.vmi.mask<128xpred> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +%b = pto.vmi.group_broadcast %sum {num_groups = 8} +%y = pto.vmi.mulf %x, %b +%ysum = pto.vmi.group_reduce_addf %y, %mask {num_groups = 8} +pto.vmi.group_store %ysum, %out[%group_off], %c1 {num_groups = 8} +``` + +Assigned layouts: + +```text +%x, %b, %y: + !pto.vmi.vreg<128xf32, + #pto.vmi.layout> + +%mask: + !pto.vmi.mask<128xpred, + #pto.vmi.layout> + +%sum, %ysum: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> +``` + +VPTO lowering result for one `8x16xf32` tile: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%sum_mask = pto.pge_b32 "PAT_VL8" + +%lane = pto.vci %c0_i32 : i32 -> !pto.vreg<64xi32> +%row = pto.vshrs %lane, %c3_i16, %all_b32 + : !pto.vreg<64xi32>, i16, !pto.mask -> !pto.vreg<64xi32> +%row8 = pto.vshls %row, %c3_i16, %all_b32 + : !pto.vreg<64xi32>, i16, !pto.mask -> !pto.vreg<64xi32> +%col = pto.vsub %lane, %row8, %all_b32 + : !pto.vreg<64xi32> +%hi4_mask = pto.vcmps %col, %c4_i32, %all_b32, "lt" + : !pto.vreg<64xi32>, i32, !pto.mask -> !pto.mask + +%x_lo, %x_hi = pto.vldsx2 %base[%tile_off], "BDINTLV" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%x_lo_sum = pto.vcgadd %x_lo, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%x_hi_sum = pto.vcgadd %x_hi, %hi4_mask + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%sum_block = pto.vadd %x_lo_sum, %x_hi_sum, %sum_mask + : !pto.vreg<64xf32> + +%broadcast_idx = pto.vshrs %lane, %c3_i16, %all_b32 + : !pto.vreg<64xi32>, i16, !pto.mask -> !pto.vreg<64xi32> +%b_rows = pto.vselr %sum_block, %broadcast_idx + : !pto.vreg<64xf32>, !pto.vreg<64xi32> -> !pto.vreg<64xf32> + +%y_lo = pto.vmul %x_lo, %b_rows, %all_b32 : !pto.vreg<64xf32> +%y_hi = pto.vmul %x_hi, %b_rows, %hi4_mask : !pto.vreg<64xf32> + +%y_lo_sum = pto.vcgadd %y_lo, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%y_hi_sum = pto.vcgadd %y_hi, %hi4_mask + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%ysum_block = pto.vadd %y_lo_sum, %y_hi_sum, %sum_mask + : !pto.vreg<64xf32> + +pto.vsts %ysum_block, %out[%group_tile_off], %sum_mask {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + s = reduce(row_r[0..11]) + out[group_tile_off + r] = + reduce_i(row_r[i] * s for i = 0..11) + = s * s +``` + +Required assignment rule: + +```text +%mask is a grouped mask with S=16 and active_elems_per_group=12. +For the low half, the physical predicate is PAT_ALL. +For the high half, the physical predicate is lane_mod_8 < 4. +The same split must be reused for both group_reduce operations. +``` + +### 3.27 S=32 `group_load` With Stride Greater Than Group Size + +This case is the S=32 counterpart to section 3.15.2. The logical group is +`32xf32`, but rows in memory have a larger stride. The fast plan is legal only +when the stride is a multiple of one 32B f32 block. + +VMI input: + +```text +%stride40 = arith.constant 40 : index +%x = pto.vmi.group_load %base[%off], %stride40 + {num_groups = 8, group_size = 32} + : !pto.ptr, index -> !pto.vmi.vreg<256xf32> +%mask = pto.vmi.create_group_mask %c32 {num_groups = 8, group_size = 32} + : index -> !pto.vmi.mask<256xpred> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +pto.vmi.group_store %sum, %out[%group_off], %c1 {num_groups = 8} +``` + +Assigned layouts: + +```text +%x: + !pto.vmi.vreg<256xf32, + #pto.vmi.layout> + +%mask: + !pto.vmi.mask<256xpred, + #pto.vmi.layout> + +%sum: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> +``` + +VPTO lowering result: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%sum_mask = pto.pge_b32 "PAT_VL8" + +// source_group_stride = 40 f32 = 5 * 32B blocks. +%stride_blocks = %c5_i16 + +%frag0 = pto.vsldb %base_frag0, %stride_blocks, %c0_i16, %all_b32 + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> +%frag1 = pto.vsldb %base_frag1, %stride_blocks, %c0_i16, %all_b32 + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> +%frag2 = pto.vsldb %base_frag2, %stride_blocks, %c0_i16, %all_b32 + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> +%frag3 = pto.vsldb %base_frag3, %stride_blocks, %c0_i16, %all_b32 + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> + +%frag0 lanes r*8 .. r*8+7 = row_r[0..7] +%frag1 lanes r*8 .. r*8+7 = row_r[8..15] +%frag2 lanes r*8 .. r*8+7 = row_r[16..23] +%frag3 lanes r*8 .. r*8+7 = row_r[24..31] + +%s0 = pto.vcgadd %frag0, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s1 = pto.vcgadd %frag1, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s2 = pto.vcgadd %frag2, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s3 = pto.vcgadd %frag3, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +%s01 = pto.vadd %s0, %s1, %sum_mask : !pto.vreg<64xf32> +%s23 = pto.vadd %s2, %s3, %sum_mask : !pto.vreg<64xf32> +%sum_block = pto.vadd %s01, %s23, %sum_mask : !pto.vreg<64xf32> + +pto.vsts %sum_block, %out[%group_tile_off], %sum_mask {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + out[group_tile_off + r] = + reduce(base[tile_off + r * 40 + 0 .. tile_off + r * 40 + 31]) +``` + +Required diagnostic when the stride is not block-aligned: + +```text +VMI-LAYOUT-CONTRACT: + pto.vmi.group_load group_size 32 with source_group_stride not divisible by + 8 f32 elements cannot use the vsldb strided-block lowering support. Enable a + stable gather fallback or choose a block-aligned source_group_stride. +``` + +Required assignment rule: + +```text +This producer requires the S=32 block-fragment layout: + #pto.vmi.layout + +It must not be unified with the contiguous-load S=32 plan from section 3.6: + #pto.vmi.layout + +Both layouts are legal inputs to group_reduce_addf S=32, but they require +different producer materialization/lowering support. +``` + +### 3.28 `group_slot_load` `slots = 1` With Aligned Non-Unit Stride + +Section 3.16.1 diagnoses non-unit stride for the packed `slots = 8` plan. The +row-local `slots = 1` plan supports non-unit stride only when each one-lane +load can be issued as an aligned `vsldb`. In the current lowering this means +the stride is a positive compile-time constant and is divisible by the 32B +alignment expressed in source elements. + +VMI input: + +```text +%c8 = arith.constant 8 : index +%rhs = pto.vmi.group_slot_load %rhs_base[%rhs_off], %c8 {num_groups = 8} + : !pto.ptr, index -> !pto.vmi.vreg<512xf32> +pto.vmi.group_store %rhs, %out[%group_off], %c8 {num_groups = 8} +``` + +Assigned layout: + +```text +%rhs: + !pto.vmi.vreg<512xf32, #pto.vmi.layout> +``` + +VPTO lowering result: + +```text +%one_b32 = pto.pge_b32 "PAT_VL1" + +// Emit this shape for r = 0..7. The address expression is scalar/index +// arithmetic outside the vector register layout. For f32, %c8 is 32B. +%addr_r = %rhs_base + %rhs_off + r * 8 +%rhs_r = pto.vsldb %addr_r, %c0_i16, %c0_i16, %one_b32 + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> + +pto.vsts %rhs_r, %out[%group_tile_off_r], %one_b32 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + out[group_off + r * 8] = rhs_base[rhs_off + r * 8] +``` + +Required assignment rule: + +```text +If a non-unit-stride group_slot_load has only slots=1 consumers and its stride +is a positive constant divisible by the element count of 32B, select +group_slot_load_slots1_row_local. Do not diagnose it using the slots=8 +unit-stride restriction. +``` + +Required diagnostic: + +```text +%c2 = arith.constant 2 : index +%bad = pto.vmi.group_slot_load %rhs_base[%rhs_off], %c2 {num_groups = 8} + : !pto.ptr -> !pto.vmi.vreg<512xf32> + +VMI-UNSUPPORTED: pto.vmi.group_slot_load + slots=1 group_slot_load currently lowers as one lane-0 vsldb per group and + requires constant positive source_group_stride divisible by 8 elements for + 32B load alignment; packed or unaligned scalar load lowering is not + implemented. +``` + +Dynamic stride has the same status until a stable gather or scalarized packed +load plan is designed: + +```text +%bad = pto.vmi.group_slot_load %rhs_base[%rhs_off], %runtime_stride + {num_groups = 8} + : !pto.ptr -> !pto.vmi.vreg<512xf32> + +VMI-UNSUPPORTED: pto.vmi.group_slot_load + requires constant positive source_group_stride divisible by 8 elements. +``` + +### 3.29 One Semantic Mask With f32 And f16 Consumers + +One VMI mask may feed consumers with different physical predicate +granularities. Layout assignment must keep the semantic mask value single, but +materialize per-use physical masks after element type is known. + +VMI input: + +```text +%x = pto.vmi.load %base[%off] + : memref<128xf32> -> !pto.vmi.vreg<128xf32> +%mask = pto.vmi.create_mask %c96 + : index -> !pto.vmi.mask<128xpred> +pto.vmi.masked_store %x, %out32[%off], %mask +%h = pto.vmi.truncf %x + : !pto.vmi.vreg<128xf32> -> !pto.vmi.vreg<128xf16> +pto.vmi.masked_store %h, %out16[%off], %mask +``` + +Assigned layouts: + +```text +%x: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> + +%mask: + !pto.vmi.mask<128xb32, #pto.vmi.layout> + +%x_for_cast: + pto.vmi.ensure_layout %x + : #pto.vmi.layout -> #pto.vmi.layout + +%mask_for_h_store: + pto.vmi.create_mask %c96 + : index -> !pto.vmi.mask<128xb16, #pto.vmi.layout> + +%h: + !pto.vmi.vreg<128xf16, #pto.vmi.layout> +``` + +Physical mask materialization: + +```text +use at masked_store %x: + predicate granularity b32, PAT_VL96, layout contiguous + +use at vcvt %x -> %h: + predicate granularity b32, PAT_ALL. The cast may compute inactive lanes + because the following masked_store controls the external memory effect. + +use at masked_store %h: + predicate granularity b16, PAT_VL96, layout contiguous +``` + +VPTO lowering result: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%mask32_0 = pto.pge_b32 "PAT_ALL" +%mask32_1 = pto.pge_b32 "PAT_VL32" + +%x0 = pto.vlds %base[%off] + : !pto.ptr, index -> !pto.vreg<64xf32> +%x1 = pto.vlds %base[%off_plus_64] + : !pto.ptr, index -> !pto.vreg<64xf32> + +pto.vsts %x0, %out32[%off], %mask32_0 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +pto.vsts %x1, %out32[%off_plus_64], %mask32_1 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + +%x_p0, %x_p1 = pto.vdintlv %x0, %x1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%h_even = pto.vcvt %x_p0, %all_b32 {part = "EVEN", rnd = "R", sat = "SAT"} + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<128xf16> +%h_odd = pto.vcvt %x_p1, %all_b32 {part = "ODD", rnd = "R", sat = "SAT"} + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<128xf16> + +%all_b16 = pto.pset_b16 "PAT_ALL" +%h0 = pto.vor %h_even, %h_odd, %all_b16 + : !pto.vreg<128xf16> +%mask_b16, %scalar_out = pto.plt_b16 %c96_i32 + : i32 -> !pto.mask, i32 +pto.vsts %h0, %out16[%off], %mask_b16 {dist = "NORM_B16"} + : !pto.vreg<128xf16>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for i = 0..95: + out32[off + i] = base[off + i] + out16[off + i] = truncf(base[off + i]) + +for i = 96..127: + out32[off + i] is unchanged + out16[off + i] is unchanged +``` + +Required assignment rule: + +```text +`vmi-to-vpto` must not decide mask granularity by inspecting users. It consumes +the per-use typed mask materialization inserted by vmi-layout-assignment. For +a rematerializable `create_mask`, assignment may clone it as b32/b16 masks. For +a non-rematerializable mask producer, assignment must insert +`ensure_mask_granularity` or diagnose if no materialization support exists. +``` + +### 3.30 `masked_load` Tail Without Padding + +This case is the replacement for `vector.transfer_read` padding semantics in the +initial VMI surface. Tail lanes are expressed by a mask and a passthrough value; +there is no implicit padding constant in the load. The direct lowering is legal +only when every physical chunk read by `vlds` is memory-safe. + +VMI input: + +```text +%c100 = arith.constant 100 : index +%mask = pto.vmi.create_mask %c100 : index -> !pto.vmi.mask<100xpred> +%zero = pto.vmi.broadcast %c0_f32 : f32 -> !pto.vmi.vreg<100xf32> +%x = pto.vmi.masked_load %base[%c0], %mask, %zero + : memref<128xf32>, !pto.vmi.mask<100xpred>, !pto.vmi.vreg<100xf32> + -> !pto.vmi.vreg<100xf32> +pto.vmi.store %x, %out[%c0] +``` + +Assigned layouts: + +```text +%mask: + !pto.vmi.mask<100xb32, #pto.vmi.layout> + +%zero, %x: + !pto.vmi.vreg<100xf32, #pto.vmi.layout> +``` + +VPTO lowering result: + +```text +%m0 = pto.pge_b32 "PAT_ALL" +%m1 = pto.pge_b32 "PAT_VL36" + +%zero0 = pto.vdup %c0_f32, %m0 + : f32, !pto.mask -> !pto.vreg<64xf32> +%zero1 = pto.vdup %c0_f32, %m0 + : f32, !pto.mask -> !pto.vreg<64xf32> + +%l0 = pto.vlds %base[%c0] + : memref<128xf32> -> !pto.vreg<64xf32> +%x0 = pto.vsel %l0, %zero0, %m0 + : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +%l1 = pto.vlds %base[%c64] + : memref<128xf32> -> !pto.vreg<64xf32> +%x1 = pto.vsel %l1, %zero1, %m1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +pto.vsts %x0, %out[%c0], %m0 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, memref<128xf32>, !pto.mask +pto.vsts %x1, %out[%c64], %m1 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, memref<128xf32>, !pto.mask +``` + +Memory result: + +```text +for i = 0..99: + out[i] = base[i] + +for i = 100..127: + out[i] is unchanged +``` + +Required diagnostic when the source cannot prove a safe full-read footprint: + +```text +VMI-UNSUPPORTED: + pto.vmi.masked_load direct lowering requires a supported memory source, + contiguous result/passthru/mask layouts, and either full physical chunks or a + statically safe full-read footprint. Use a memref with enough static extent, + enable the future stable masked/gather load plan, or make the logical vector a + full physical chunk. +``` + +Required assignment rule: + +```text +`masked_load` requests contiguous result, passthru, and mask layouts. Padding +is not a layout decision; it is the explicit passthrough operand selected by the +user. +``` + +### 3.31 `f16 -> f32` Feeding Dense Store And S=16 Reduce + +This case proves that the `deinterleaved = 2` layout produced by widening +`f16 -> f32` is not just a store layout. It must also be a legal S=16 grouped +reduction input. Layout assignment must not force the reduce consumer to +`block_elems = 8` and then rematerialize the widened value. + +VMI input: + +```text +%x16 = pto.vmi.load %base[%off] + : memref<128xf16> -> !pto.vmi.vreg<128xf16> +%x32 = pto.vmi.extf %x16 + : !pto.vmi.vreg<128xf16> -> !pto.vmi.vreg<128xf32> +%mask = pto.vmi.create_mask %c128 : index -> !pto.vmi.mask<128xpred> +%sum = pto.vmi.group_reduce_addf %x32, %mask {num_groups = 8} +pto.vmi.group_store %sum, %sum_out[%group_off], %c1 {num_groups = 8} +pto.vmi.store %x32, %dense_out[%off] +``` + +Assigned layouts: + +```text +%x16: + !pto.vmi.vreg<128xf16, #pto.vmi.layout> + +%x32: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> + +%mask: + !pto.vmi.mask<128xb32, #pto.vmi.layout> + +%sum: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> +``` + +VPTO lowering result: + +```text +%all_b16 = pto.pge_b16 "PAT_ALL" +%all_b32 = pto.pge_b32 "PAT_ALL" +%sum_mask = pto.pge_b32 "PAT_VL8" + +%x16_0 = pto.vlds %base[%off] + : memref<128xf16> -> !pto.vreg<128xf16> +%x32_p0 = pto.vcvt %x16_0, %all_b16 {part = "EVEN"} + : !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<64xf32> +%x32_p1 = pto.vcvt %x16_0, %all_b16 {part = "ODD"} + : !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<64xf32> + +%s0 = pto.vcgadd %x32_p0, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%s1 = pto.vcgadd %x32_p1, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%sum_block = pto.vadd %s0, %s1, %sum_mask + : !pto.vreg<64xf32> + +pto.vsts %sum_block, %sum_out[%group_off], %sum_mask {dist = "NORM_B32"} + : !pto.vreg<64xf32>, memref<8xf32>, !pto.mask + +%dense0, %dense1 = pto.vintlv %x32_p0, %x32_p1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +pto.vsts %dense0, %dense_out[%off], %all_b32 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, memref<128xf32>, !pto.mask +pto.vsts %dense1, %dense_out[%off_plus_64], %all_b32 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, memref<128xf32>, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + sum_out[group_off + r] = + reduce(extf(base[off + r * 16 + 0 .. off + r * 16 + 15])) + +for i = 0..127: + dense_out[off + i] = extf(base[off + i]) +``` + +Required assignment rule: + +```text +When S=16 group_reduce consumes an existing `deinterleaved = 2` dense value, +the reduce plan must accept `block_elems = 1`. `block_elems = 8` is only a +producer-driven fast plan for block-fragment loads, not the semantic +requirement of S=16 reduction. +``` + +### 3.32 `f32` Feeding f8 Store And S=32 Reduce + +This is the `f32 -> f8` counterpart to section 3.31. A 256-lane f32 value can +serve both `truncf -> f8` and S=32 group reduction with the same +`deinterleaved = 4, block_elems = 1` layout. The value must not be forced to a +block-fragment `block_elems = 8` layout unless its producer requires that plan. + +VMI input: + +```text +%x32 = pto.vmi.load %base[%off] + : memref<256xf32> -> !pto.vmi.vreg<256xf32> +%mask = pto.vmi.create_mask %c256 : index -> !pto.vmi.mask<256xpred> +%sum = pto.vmi.group_reduce_addf %x32, %mask {num_groups = 8} +pto.vmi.group_store %sum, %sum_out[%group_off], %c1 {num_groups = 8} +%x8 = pto.vmi.truncf %x32 + : !pto.vmi.vreg<256xf32> -> !pto.vmi.vreg<256xf8> +pto.vmi.store %x8, %out8[%off] +``` + +Assigned layouts: + +```text +%x32: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> + +%mask: + !pto.vmi.mask<256xb32, #pto.vmi.layout> + +%sum: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> + +%x8: + !pto.vmi.vreg<256xf8, #pto.vmi.layout> +``` + +VPTO lowering result: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%sum_mask = pto.pge_b32 "PAT_VL8" + +%x_even_0, %x_odd_0 = pto.vldsx2 %base[%off], "DINTLV_B32" + : memref<256xf32>, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%x_even_1, %x_odd_1 = pto.vldsx2 %base[%off_plus_128], "DINTLV_B32" + : memref<256xf32>, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%x_p0, %x_p2 = pto.vdintlv %x_even_0, %x_even_1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%x_p1, %x_p3 = pto.vdintlv %x_odd_0, %x_odd_1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%s0 = pto.vcgadd %x_p0, %all_b32 : !pto.vreg<64xf32> +%s1 = pto.vcgadd %x_p1, %all_b32 : !pto.vreg<64xf32> +%s2 = pto.vcgadd %x_p2, %all_b32 : !pto.vreg<64xf32> +%s3 = pto.vcgadd %x_p3, %all_b32 : !pto.vreg<64xf32> +%s01 = pto.vadd %s0, %s1, %sum_mask : !pto.vreg<64xf32> +%s23 = pto.vadd %s2, %s3, %sum_mask : !pto.vreg<64xf32> +%sum_block = pto.vadd %s01, %s23, %sum_mask : !pto.vreg<64xf32> + +pto.vsts %sum_block, %sum_out[%group_off], %sum_mask {dist = "NORM_B32"} + : !pto.vreg<64xf32>, memref<8xf32>, !pto.mask + +%x8_p0 = pto.vcvt %x_p0, %all_b32 {part = "P0", rnd = "R", sat = "SAT"} + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<256xf8> +%x8_p1 = pto.vcvt %x_p1, %all_b32 {part = "P1", rnd = "R", sat = "SAT"} + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<256xf8> +%x8_p2 = pto.vcvt %x_p2, %all_b32 {part = "P2", rnd = "R", sat = "SAT"} + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<256xf8> +%x8_p3 = pto.vcvt %x_p3, %all_b32 {part = "P3", rnd = "R", sat = "SAT"} + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<256xf8> + +%x8_01 = pto.vor %x8_p0, %x8_p1, PAT_ALL_B8 + : !pto.vreg<256xf8> +%x8_23 = pto.vor %x8_p2, %x8_p3, PAT_ALL_B8 + : !pto.vreg<256xf8> +%x8_0 = pto.vor %x8_01, %x8_23, PAT_ALL_B8 + : !pto.vreg<256xf8> + +pto.vsts %x8_0, %out8[%off], PAT_ALL_B8 {dist = "NORM_B8"} + : !pto.vreg<256xf8>, memref<256xf8>, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + sum_out[group_off + r] = + reduce(base[off + r * 32 + 0 .. off + r * 32 + 31]) + +for i = 0..255: + out8[off + i] = truncf(base[off + i]) +``` + +Required assignment rule: + +```text +The common layout selected for `%x32` is +`#pto.vmi.layout`. This satisfies both +`truncf f32 -> f8` and S=32 `group_reduce_addf`. A later strided block-load +producer may introduce `block_elems = 8`, but that is a different case and +requires an explicit materialization/rematerialization decision. + +When `%x32` is produced by a full contiguous `pto.vmi.load`, `vmi-to-vpto` +should not first materialize four contiguous f32 chunks and then run a full +four-op `vdintlv` tree. The load lowering should fold the first deinterleave +level into two `vldsx2 DINTLV_B32` operations and then run only the second +`vdintlv` level, as shown above. The layout remains just +`deinterleaved = 4, block_elems = 1`; it does not encode the fact that `vldsx2` +was used. +``` + +### 3.33 One Dense Value Feeding S=16 And S=32 Reduces + +This case is a pure layout-assignment conflict. The same logical +`256xf32` value is consumed by two legal reductions, but their efficient input +layouts are different: + +```text +S=16 reduce over 16 groups: + #pto.vmi.layout + +S=32 reduce over 8 groups: + #pto.vmi.layout +``` + +The program is semantically legal. Baseline layout assignment solves it by +inserting an explicit use-site `ensure_layout`. A later optimization pass may +clone or rematerialize the cheap load for one use. `vmi-to-vpto` must not +inspect both users and choose one locally. + +VMI input: + +```text +%x = pto.vmi.load %base[%off] + : memref<256xf32> -> !pto.vmi.vreg<256xf32> + +%mask16 = pto.vmi.create_group_mask %c16 {num_groups = 16, group_size = 16} + : index -> !pto.vmi.mask<256xpred> +%sum16 = pto.vmi.group_reduce_addf %x, %mask16 {num_groups = 16} +pto.vmi.group_store %sum16, %out16[%group_off16], %c1 {num_groups = 16} + +%mask32 = pto.vmi.create_group_mask %c32 {num_groups = 8, group_size = 32} + : index -> !pto.vmi.mask<256xpred> +%sum32 = pto.vmi.group_reduce_addf %x, %mask32 {num_groups = 8} +pto.vmi.group_store %sum32, %out32[%group_off32], %c1 {num_groups = 8} +``` + +Assigned layouts after rematerializing the load: + +```text +%x_s16: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> + +%mask16: + !pto.vmi.mask<256xpred, #pto.vmi.layout> + +%sum16: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> + +%x_s32: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> + +%mask32: + !pto.vmi.mask<256xpred, #pto.vmi.layout> + +%sum32: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> +``` + +VPTO lowering result: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%sum8_mask = pto.pge_b32 "PAT_VL8" + +// Rematerialized S=16 use. The first vldsx2 covers rows 0..7, the second +// covers rows 8..15. Each pair is deinterleaved by element parity. +%s16_p0, %s16_p1 = pto.vldsx2 %base[%tile_off_0], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%s16_p2, %s16_p3 = pto.vldsx2 %base[%tile_off_1], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%s16_0 = pto.vcgadd %s16_p0, %all_b32 : !pto.vreg<64xf32> +%s16_1 = pto.vcgadd %s16_p1, %all_b32 : !pto.vreg<64xf32> +%s16_2 = pto.vcgadd %s16_p2, %all_b32 : !pto.vreg<64xf32> +%s16_3 = pto.vcgadd %s16_p3, %all_b32 : !pto.vreg<64xf32> + +%sum16_lo = pto.vadd %s16_0, %s16_1, %sum8_mask + : !pto.vreg<64xf32> +%sum16_hi = pto.vadd %s16_2, %s16_3, %sum8_mask + : !pto.vreg<64xf32> + +pto.vsts %sum16_lo, %out16[%group_off16], %sum8_mask {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +pto.vsts %sum16_hi, %out16[%group_off16_plus_8], %sum8_mask {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + +// Rematerialized S=32 use. Two DINTLV loads plus one register deinterleave +// level produce mod-4 columns for rows 0..7. +%x_even_0, %x_odd_0 = pto.vldsx2 %base[%tile_off_0], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%x_even_1, %x_odd_1 = pto.vldsx2 %base[%tile_off_1], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%x_p0, %x_p2 = pto.vdintlv %x_even_0, %x_even_1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%x_p1, %x_p3 = pto.vdintlv %x_odd_0, %x_odd_1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%s32_0 = pto.vcgadd %x_p0, %all_b32 : !pto.vreg<64xf32> +%s32_1 = pto.vcgadd %x_p1, %all_b32 : !pto.vreg<64xf32> +%s32_2 = pto.vcgadd %x_p2, %all_b32 : !pto.vreg<64xf32> +%s32_3 = pto.vcgadd %x_p3, %all_b32 : !pto.vreg<64xf32> + +%s32_01 = pto.vadd %s32_0, %s32_1, %sum8_mask : !pto.vreg<64xf32> +%s32_23 = pto.vadd %s32_2, %s32_3, %sum8_mask : !pto.vreg<64xf32> +%sum32_block = pto.vadd %s32_01, %s32_23, %sum8_mask : !pto.vreg<64xf32> + +pto.vsts %sum32_block, %out32[%group_off32], %sum8_mask {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..15: + out16[group_off16 + r] = + reduce(base[off + r * 16 + 0 .. off + r * 16 + 15]) + +for r = 0..7: + out32[group_off32 + r] = + reduce(base[off + r * 32 + 0 .. off + r * 32 + 31]) +``` + +Required assignment rule: + +```text +Baseline assignment inserts `ensure_layout` at the mismatched use. A later +rematerialization pass may clone a cheap producer such as load and assign each +clone independently. If no deinterleaved=2 <-> deinterleaved=4 materialization +support exists, emit a layout-contract diagnostic naming both consumers and +both required layouts. +``` + +### 3.34 S=64 Group-Slot Result `f32 -> f16` Cast + +Section 3.13 rejects direct width-changing cast for packed `slots = 8` +group-slot values. This case is the positive counterpart for row-local +`slots = 1`: each group result is already lane 0 of its own physical vreg, so a +slot-preserving cast can lower one row-local result at a time. + +VMI input: + +```text +%x = pto.vmi.load %base[%off] + : memref<512xf32> -> !pto.vmi.vreg<512xf32> +%mask = pto.vmi.create_group_mask %c64 {num_groups = 8, group_size = 64} + : index -> !pto.vmi.mask<512xpred> +%sum32 = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +%sum16 = pto.vmi.truncf %sum32 + : !pto.vmi.vreg<512xf32> -> !pto.vmi.vreg<512xf16> +pto.vmi.group_store %sum16, %out[%group_off], %c1 {num_groups = 8} +``` + +Assigned layouts: + +```text +%x: + !pto.vmi.vreg<512xf32, #pto.vmi.layout> + +%sum32: + !pto.vmi.vreg<512xf32, #pto.vmi.layout> + +%sum16: + !pto.vmi.vreg<512xf16, #pto.vmi.layout> +``` + +VPTO lowering result for one full 8-row tile: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%block8 = pto.pge_b32 "PAT_VL8" +%one_b32 = pto.pge_b32 "PAT_VL1" +%one_b16 = pto.pge_b16 "PAT_VL1" + +// The compiler emits this row-local sequence for r = 0..7. +%x_r = pto.vlds %base[%row_off_r] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<64xf32> +%p_r = pto.vcgadd %x_r, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%sum32_r = pto.vcadd %p_r, %block8 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + +// Only lane 0 is semantic. EVEN keeps f32 lane 0 in f16 lane 0; all other +// lanes are non-semantic for group_slots(num_groups=8, slots=1). +%sum16_r = pto.vcvt %sum32_r, %one_b32 {part = "EVEN", rnd = "R", sat = "SAT"} + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<128xf16> + +pto.vsts %sum16_r, %out[%group_tile_off_r], %one_b16 {dist = "NORM_B16"} + : !pto.vreg<128xf16>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + out[group_off + r] = + truncf(reduce(base[off + r * 64 + 0 .. off + r * 64 + 63])) +``` + +Required assignment rule: + +```text +Group-slot casts are layout-specific. `slots = 1` may use a slot-preserving +row-local cast because each semantic scalar is lane 0 of its own physical vreg. +This does not legalize packed `slots = 8` casts from section 3.13. +``` + +### 3.35 `group_slots` Fanout To `group_store` And `group_broadcast` + +This case fixes the fanout rule for group-slot values. A `group_slots` value may +feed multiple group-aware consumers directly. Layout assignment must not +materialize it as dense just because one later use broadcasts it. + +VMI input: + +```text +%x = pto.vmi.load %base[%off] + : memref<128xf32> -> !pto.vmi.vreg<128xf32> +%mask = pto.vmi.create_group_mask %c16 {num_groups = 8, group_size = 16} + : index -> !pto.vmi.mask<128xpred> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +pto.vmi.group_store %sum, %sum_out[%group_off], %c1 {num_groups = 8} + +%b = pto.vmi.group_broadcast %sum {num_groups = 8} +%y = pto.vmi.mulf %x, %b +%ysum = pto.vmi.group_reduce_addf %y, %mask {num_groups = 8} +pto.vmi.group_store %ysum, %out[%group_off], %c1 {num_groups = 8} +``` + +Assigned layouts: + +```text +%x: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> + +%x_for_reduce: + !pto.vmi.vreg<128xf32, + #pto.vmi.layout> + +%mask_for_reduce: + !pto.vmi.mask<128xb32, + #pto.vmi.layout> + +%sum, %ysum: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> + +%b, %y: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> + +%y_for_reduce: + !pto.vmi.vreg<128xf32, + #pto.vmi.layout> +``` + +VPTO lowering result for one full 8-row tile: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%slot8 = pto.pge_b32 "PAT_VL8" + +%x0 = pto.vlds %base[%tile_off] + : !pto.ptr, index -> !pto.vreg<64xf32> +%x1 = pto.vlds %base[%tile_off_plus_64] + : !pto.ptr, index -> !pto.vreg<64xf32> + +// ensure_layout for the first group_reduce. +%x_lo, %x_hi = pto.vdintlv %x0, %x1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%lo_sum = pto.vcgadd %x_lo, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%hi_sum = pto.vcgadd %x_hi, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%sum_block = pto.vadd %lo_sum, %hi_sum, %slot8 : !pto.vreg<64xf32> + +// First group-slot consumer: store the group slots without changing layout. +pto.vsts %sum_block, %sum_out[%group_off], %slot8 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + +// Second group-slot consumer: materialize only this use as dense grouped data. +%broadcast_idx0 = compute index vector [0 repeated 16, 1 repeated 16, + 2 repeated 16, 3 repeated 16] + : !pto.vreg<64xi32> +%broadcast_idx1 = compute index vector [4 repeated 16, 5 repeated 16, + 6 repeated 16, 7 repeated 16] + : !pto.vreg<64xi32> +%b0 = pto.vselr %sum_block, %broadcast_idx0 + : !pto.vreg<64xf32>, !pto.vreg<64xi32> -> !pto.vreg<64xf32> +%b1 = pto.vselr %sum_block, %broadcast_idx1 + : !pto.vreg<64xf32>, !pto.vreg<64xi32> -> !pto.vreg<64xf32> + +%y0 = pto.vmul %x0, %b0, %all_b32 : !pto.vreg<64xf32> +%y1 = pto.vmul %x1, %b1, %all_b32 : !pto.vreg<64xf32> + +// ensure_layout for the second group_reduce. +%y_lo, %y_hi = pto.vdintlv %y0, %y1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%y_lo_sum = pto.vcgadd %y_lo, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%y_hi_sum = pto.vcgadd %y_hi, %all_b32 + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +%ysum_block = pto.vadd %y_lo_sum, %y_hi_sum, %slot8 : !pto.vreg<64xf32> + +pto.vsts %ysum_block, %out[%group_off], %slot8 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + s = reduce(row_r[0..15]) + sum_out[group_off + r] = s + out[group_off + r] = reduce_i(row_r[i] * s for i = 0..15) +``` + +Required assignment rule: + +```text +`%sum` keeps one assigned layout: + #pto.vmi.layout + +`group_store` consumes that group-slot layout directly. +`group_broadcast` is a use-site materialization to a dense layout. It must not +rewrite the defining `group_reduce` result or the sibling `group_store` use. +``` + +### 3.36 Same Scalar Source Materialized As `slots = 8` And `slots = 1` + +The same memory scalar stream may be used by both packed S=16 group-slot +compute and row-local S=64 group-slot compute. The two uses require different +logical vector shapes and different group-slot layouts, so the source must be +rematerialized as two VMI values. There is no single `group_slots` layout that +serves both uses. + +VMI input: + +```text +%rhs16 = pto.vmi.group_slot_load %rhs_base[%rhs_off], %c1 {num_groups = 8} + : !pto.ptr, index -> !pto.vmi.vreg<128xf32> +%x16 = pto.vmi.load %base16[%off16] + : memref<128xf32> -> !pto.vmi.vreg<128xf32> +%sum16 = pto.vmi.group_reduce_addf %x16, %mask16 {num_groups = 8} +%out16v = pto.vmi.addf %sum16, %rhs16 +pto.vmi.group_store %out16v, %out16[%group_off16], %c1 {num_groups = 8} + +%rhs64 = pto.vmi.group_slot_load %rhs_base[%rhs_off], %c1 {num_groups = 8} + : !pto.ptr, index -> !pto.vmi.vreg<512xf32> +%x64 = pto.vmi.load %base64[%off64] + : memref<512xf32> -> !pto.vmi.vreg<512xf32> +%sum64 = pto.vmi.group_reduce_addf %x64, %mask64 {num_groups = 8} +%out64v = pto.vmi.addf %sum64, %rhs64 +pto.vmi.group_store %out64v, %out64[%group_off64], %c1 {num_groups = 8} +``` + +Assigned layouts: + +```text +%rhs16, %sum16, %out16v: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> + +%x16, %mask16: + #pto.vmi.layout + +%rhs64, %sum64, %out64v: + !pto.vmi.vreg<512xf32, #pto.vmi.layout> + +%x64, %mask64: + #pto.vmi.layout +``` + +VPTO lowering result: + +```text +// Packed S=16 RHS: one 32B scalar block in lanes 0..7. +%slot8 = pto.pge_b32 "PAT_VL8" +%one_b32 = pto.pge_b32 "PAT_VL1" +%rhs16_block = pto.vsldb %rhs_base[%rhs_off], %c0_i16, %c0_i16, %one_b32 + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> + +// S=16 reduction is the section 3.5.1 shape. +%x16_lo, %x16_hi = pto.vldsx2 %base16[%tile_off16], "BDINTLV" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%s16_lo = pto.vcgadd %x16_lo, PAT_ALL_B32 : !pto.vreg<64xf32> +%s16_hi = pto.vcgadd %x16_hi, PAT_ALL_B32 : !pto.vreg<64xf32> +%sum16_block = pto.vadd %s16_lo, %s16_hi, %slot8 : !pto.vreg<64xf32> +%out16_block = pto.vadd %sum16_block, %rhs16_block, %slot8 + : !pto.vreg<64xf32> +pto.vsts %out16_block, %out16[%group_off16], %slot8 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + +// Row-local S=64 RHS: a separate group_slot_load op produces one lane-0 +// value per physical row-local result. +%rhs64_r = pto.vsldb %rhs_base[%rhs_off_plus_r], %c0_i16, %c0_i16, %one_b32 + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> + +// Emit this row-local reduction/add/store shape for r = 0..7. +%x64_r = pto.vlds %base64[%row_off64_r] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<64xf32> +%p64_r = pto.vcgadd %x64_r, PAT_ALL_B32 : !pto.vreg<64xf32> +%sum64_r = pto.vcadd %p64_r, PAT_VL8_B32 : !pto.vreg<64xf32> +%out64_r = pto.vadd %sum64_r, %rhs64_r, %one_b32 : !pto.vreg<64xf32> +pto.vsts %out64_r, %out64[%group_off64_plus_r], %one_b32 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + out16[group_off16 + r] = reduce(base16[row_r, 0..15]) + rhs_base[rhs_off + r] + out64[group_off64 + r] = reduce(base64[row_r, 0..63]) + rhs_base[rhs_off + r] +``` + +Required assignment rule: + +```text +`group_slot_load` is a memory op, so the baseline rematerialization pass must +not clone it as a generic cheap producer. If two use sites need different +`group_slots` layouts, the legal first-stage shape is to write two explicit +`group_slot_load` ops, as above, or to introduce a future load-cloning +optimization with an explicit memory-safety proof. Do not invent a common +layout or make `vmi-to-vpto` inspect both users. +``` + +### 3.37 S=64 `group_store` With Non-Unit Output Stride + +Packed `slots = 8` stores currently require unit output stride. Row-local +`slots = 1` does not have that restriction because each group scalar is stored +by a separate lane-0 store. + +VMI input: + +```text +%row_stride = arith.index_cast %ld : i64 to index +%x = pto.vmi.load %base[%off] + : memref<512xf32> -> !pto.vmi.vreg<512xf32> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +pto.vmi.group_store %sum, %out[%group_off], %row_stride {num_groups = 8} +``` + +Assigned layouts: + +```text +%x: + !pto.vmi.vreg<512xf32, #pto.vmi.layout> + +%sum: + !pto.vmi.vreg<512xf32, #pto.vmi.layout> +``` + +VPTO lowering result: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%block8 = pto.pge_b32 "PAT_VL8" +%one_b32 = pto.pge_b32 "PAT_VL1" + +// Emit this row-local sequence for r = 0..7. +%x_r = pto.vlds %base[%row_off_r] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<64xf32> +%p_r = pto.vcgadd %x_r, %all_b32 : !pto.vreg<64xf32> +%sum_r = pto.vcadd %p_r, %block8 : !pto.vreg<64xf32> + +%dst_r = %out + %group_off + r * %row_stride +pto.vsts %sum_r, %dst_r, %one_b32 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + out[group_off + r * row_stride] = reduce(row_r[0..63]) +``` + +Required assignment rule: + +```text +If `group_store` has non-unit row_stride and the source can legally use +`slots = 1`, assignment may select `slots = 1` to keep the store legal. If the +source is fixed to `slots = 8`, current target support must diagnose unless a +strided packed store materializer exists. +``` + +### 3.38 Multi-Tile S=32 `group_reduce` + +The S=32 plan is not only a one-tile special case. For more than eight groups, +layout assignment keeps the same layout and `vmi-to-vpto` emits the same +8-row tile lowering sequence for each physical tile. + +VMI input: + +```text +%x = pto.vmi.load %base[%off] + : memref<512xf32> -> !pto.vmi.vreg<512xf32> +%mask = pto.vmi.create_group_mask %c32 {num_groups = 16, group_size = 32} + : index -> !pto.vmi.mask<512xpred> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 16} +pto.vmi.group_store %sum, %out[%group_off], %c1 {num_groups = 16} +``` + +Assigned layouts: + +```text +%x, %mask: + !pto.vmi.vreg<512xf32, + #pto.vmi.layout> + +%sum: + !pto.vmi.vreg<512xf32, #pto.vmi.layout> +``` + +VPTO lowering result: + +```text +// Emit this shape for tile t = 0 and tile t = 1. +// Each tile covers eight 32-f32 rows. +%tile_base_t = %base + %off + t * 256 +%tile_out_t = %out + %group_off + t * 8 + +%x_even_0_t, %x_odd_0_t = pto.vldsx2 %tile_base_t[%c0], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%x_even_1_t, %x_odd_1_t = pto.vldsx2 %tile_base_t[%c128], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%x_p0_t, %x_p2_t = pto.vdintlv %x_even_0_t, %x_even_1_t + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%x_p1_t, %x_p3_t = pto.vdintlv %x_odd_0_t, %x_odd_1_t + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%s0_t = pto.vcgadd %x_p0_t, PAT_ALL_B32 : !pto.vreg<64xf32> +%s1_t = pto.vcgadd %x_p1_t, PAT_ALL_B32 : !pto.vreg<64xf32> +%s2_t = pto.vcgadd %x_p2_t, PAT_ALL_B32 : !pto.vreg<64xf32> +%s3_t = pto.vcgadd %x_p3_t, PAT_ALL_B32 : !pto.vreg<64xf32> +%s01_t = pto.vadd %s0_t, %s1_t, PAT_VL8_B32 : !pto.vreg<64xf32> +%s23_t = pto.vadd %s2_t, %s3_t, PAT_VL8_B32 : !pto.vreg<64xf32> +%sum_block_t = pto.vadd %s01_t, %s23_t, PAT_VL8_B32 + : !pto.vreg<64xf32> + +pto.vsts %sum_block_t, %tile_out_t, PAT_VL8_B32 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..15: + out[group_off + r] = + reduce(base[off + r * 32 + 0 .. off + r * 32 + 31]) +``` + +Required assignment rule: + +```text +For `group_slots(num_groups = 16, slots = 8)`, the physical arity is +`num_groups / slots = 2`. The type conversion must expose two packed result +blocks in group order. `group_store` stores both blocks with offsets +`group_off + 0` and `group_off + 8`. +``` + +### 3.39 Strided S=32 `group_load` Through Broadcast And Second Reduce + +Section 3.27 covers strided S=32 `group_load -> group_reduce -> group_store`. +This case adds the missing dense continuation. The important layout fact is +that a strided block load naturally produces +`deinterleaved = 4, block_elems = 8`; `group_broadcast` must materialize the +broadcast into that same block-fragment layout when the broadcast feeds +elementwise compute and another S=32 group reduction. + +VMI input: + +```text +%stride40 = arith.constant 40 : index +%x = pto.vmi.group_load %base[%off], %stride40 + {num_groups = 8, group_size = 32} + : !pto.ptr, index -> !pto.vmi.vreg<256xf32> +%mask = pto.vmi.create_group_mask %c32 {num_groups = 8, group_size = 32} + : index -> !pto.vmi.mask<256xpred> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +%b = pto.vmi.group_broadcast %sum {num_groups = 8} +%y = pto.vmi.mulf %x, %b +%ysum = pto.vmi.group_reduce_addf %y, %mask {num_groups = 8} +pto.vmi.group_store %ysum, %out[%group_off], %c1 {num_groups = 8} +``` + +Assigned layouts: + +```text +%x, %mask, %b, %y: + !pto.vmi.vreg<256xf32, + #pto.vmi.layout> + +%sum, %ysum: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> +``` + +VPTO lowering result: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%slot8 = pto.pge_b32 "PAT_VL8" +%stride_blocks = %c5_i16 // 40 f32 = 5 * 32B blocks. + +%x_p0 = pto.vsldb %base_frag0, %stride_blocks, %c0_i16, %all_b32 + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> +%x_p1 = pto.vsldb %base_frag1, %stride_blocks, %c0_i16, %all_b32 + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> +%x_p2 = pto.vsldb %base_frag2, %stride_blocks, %c0_i16, %all_b32 + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> +%x_p3 = pto.vsldb %base_frag3, %stride_blocks, %c0_i16, %all_b32 + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> + +%s0 = pto.vcgadd %x_p0, %all_b32 : !pto.vreg<64xf32> +%s1 = pto.vcgadd %x_p1, %all_b32 : !pto.vreg<64xf32> +%s2 = pto.vcgadd %x_p2, %all_b32 : !pto.vreg<64xf32> +%s3 = pto.vcgadd %x_p3, %all_b32 : !pto.vreg<64xf32> +%s01 = pto.vadd %s0, %s1, %slot8 : !pto.vreg<64xf32> +%s23 = pto.vadd %s2, %s3, %slot8 : !pto.vreg<64xf32> +%sum_block = pto.vadd %s01, %s23, %slot8 : !pto.vreg<64xf32> + +%lane_id = pto.vci %c0_i32 : i32 -> !pto.vreg<64xi32> +%broadcast_idx = pto.vshrs %lane_id, %c3_i16, %all_b32 + : !pto.vreg<64xi32>, i16, !pto.mask -> !pto.vreg<64xi32> + +// Materialize the same per-row scalar into every 32B row fragment. The four +// bundle entries have the same lane contents, but the result layout remains +// deinterleaved=4, block_elems=8 because the consumer `%y = mulf %x, %b` +// operates on the block-fragment layout. +%b_p0 = pto.vselr %sum_block, %broadcast_idx + : !pto.vreg<64xf32>, !pto.vreg<64xi32> -> !pto.vreg<64xf32> +%b_p1 = pto.vselr %sum_block, %broadcast_idx + : !pto.vreg<64xf32>, !pto.vreg<64xi32> -> !pto.vreg<64xf32> +%b_p2 = pto.vselr %sum_block, %broadcast_idx + : !pto.vreg<64xf32>, !pto.vreg<64xi32> -> !pto.vreg<64xf32> +%b_p3 = pto.vselr %sum_block, %broadcast_idx + : !pto.vreg<64xf32>, !pto.vreg<64xi32> -> !pto.vreg<64xf32> + +%y_p0 = pto.vmul %x_p0, %b_p0, %all_b32 : !pto.vreg<64xf32> +%y_p1 = pto.vmul %x_p1, %b_p1, %all_b32 : !pto.vreg<64xf32> +%y_p2 = pto.vmul %x_p2, %b_p2, %all_b32 : !pto.vreg<64xf32> +%y_p3 = pto.vmul %x_p3, %b_p3, %all_b32 : !pto.vreg<64xf32> + +%ys0 = pto.vcgadd %y_p0, %all_b32 : !pto.vreg<64xf32> +%ys1 = pto.vcgadd %y_p1, %all_b32 : !pto.vreg<64xf32> +%ys2 = pto.vcgadd %y_p2, %all_b32 : !pto.vreg<64xf32> +%ys3 = pto.vcgadd %y_p3, %all_b32 : !pto.vreg<64xf32> +%ys01 = pto.vadd %ys0, %ys1, %slot8 : !pto.vreg<64xf32> +%ys23 = pto.vadd %ys2, %ys3, %slot8 : !pto.vreg<64xf32> +%ysum_block = pto.vadd %ys01, %ys23, %slot8 : !pto.vreg<64xf32> + +pto.vsts %ysum_block, %out[%group_off], %slot8 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + s = reduce(base[off + r * 40 + 0 .. off + r * 40 + 31]) + out[group_off + r] = + reduce_i(base[off + r * 40 + i] * s for i = 0..31) +``` + +Required assignment rule: + +```text +`block_elems` is part of dense layout compatibility. A broadcast result feeding +an elementwise op with `%x : deinterleaved=4, block_elems=8` must also be +assigned `deinterleaved=4, block_elems=8`. Reusing a +`deinterleaved=4, block_elems=1` broadcast would be a layout mismatch even +though both have four physical parts. +``` + +### 3.40 Scalar Broadcast Feeding Dense And Grouped Users + +This case fixes the rule for ordinary scalar broadcasts. A scalar broadcast is +not born with a physical layout. Baseline layout assignment assigns the +transfer-equivalent producer chain to the non-contiguous layout requested by the +grouped consumer and inserts an explicit materialization at the dense store use. +The later `vmi-layout-rematerialize` pass may replace that helper with a cloned +broadcast when profitable. + +VMI input: + +```text +%scale = pto.vmi.broadcast %scale_s + : f32 -> !pto.vmi.vreg<256xf32> +%x = pto.vmi.load %base[%off] + : memref<256xf32> -> !pto.vmi.vreg<256xf32> + +%copy = pto.vmi.addf %x, %scale +pto.vmi.store %copy, %copy_out[%off] + +%mask = pto.vmi.create_group_mask %c32 {num_groups = 8, group_size = 32} + : index -> !pto.vmi.mask<256xpred> +%prod = pto.vmi.mulf %x, %scale +%sum = pto.vmi.group_reduce_addf %prod, %mask {num_groups = 8} +pto.vmi.group_store %sum, %sum_out[%group_off], %c1 {num_groups = 8} +``` + +Assigned layouts: + +```text +%x, %scale, %copy, %prod: + !pto.vmi.vreg<256xf32, + #pto.vmi.layout> + +%copy_dense = pto.vmi.ensure_layout %copy: + #pto.vmi.layout + -> #pto.vmi.layout + +%mask: + !pto.vmi.mask<256xpred, + #pto.vmi.layout> + +%sum: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> +``` + +VPTO lowering result: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%slot8 = pto.pge_b32 "PAT_VL8" + +// The shared load is assigned deinterleaved=4, block_elems=8 because the +// grouped consumer dominates the useful compute layout. +%x0 = pto.vlds %base[%off] : !pto.ptr, index -> !pto.vreg<64xf32> +%x1 = pto.vlds %base[%off_plus_64] : !pto.ptr, index -> !pto.vreg<64xf32> +%x2 = pto.vlds %base[%off_plus_128] : !pto.ptr, index -> !pto.vreg<64xf32> +%x3 = pto.vlds %base[%off_plus_192] : !pto.ptr, index -> !pto.vreg<64xf32> + +%x01_lo, %x01_hi = pto.vdintlv %x0, %x1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%x23_lo, %x23_hi = pto.vdintlv %x2, %x3 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%x_p0, %x_p2 = pto.vdintlv %x01_lo, %x23_lo + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%x_p1, %x_p3 = pto.vdintlv %x01_hi, %x23_hi + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%scale_p0 = pto.vdup %scale_s, %all_b32 : f32, !pto.mask -> !pto.vreg<64xf32> +%scale_p1 = pto.vdup %scale_s, %all_b32 : f32, !pto.mask -> !pto.vreg<64xf32> +%scale_p2 = pto.vdup %scale_s, %all_b32 : f32, !pto.mask -> !pto.vreg<64xf32> +%scale_p3 = pto.vdup %scale_s, %all_b32 : f32, !pto.mask -> !pto.vreg<64xf32> + +// Dense store use: compute in deinterleaved=4, then ensure_layout materializes +// the contiguous memory order for the external effect. +%copy_p0 = pto.vadd %x_p0, %scale_p0, %all_b32 : !pto.vreg<64xf32> +%copy_p1 = pto.vadd %x_p1, %scale_p1, %all_b32 : !pto.vreg<64xf32> +%copy_p2 = pto.vadd %x_p2, %scale_p2, %all_b32 : !pto.vreg<64xf32> +%copy_p3 = pto.vadd %x_p3, %scale_p3, %all_b32 : !pto.vreg<64xf32> + +%c01_lo, %c01_hi = pto.vintlv %copy_p0, %copy_p2 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%c23_lo, %c23_hi = pto.vintlv %copy_p1, %copy_p3 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%copy0, %copy1 = pto.vintlv %c01_lo, %c23_lo + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%copy2, %copy3 = pto.vintlv %c01_hi, %c23_hi + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +pto.vsts %copy0, %copy_out[%off], %all_b32 {dist = "NORM_B32"} +pto.vsts %copy1, %copy_out[%off_plus_64], %all_b32 {dist = "NORM_B32"} +pto.vsts %copy2, %copy_out[%off_plus_128], %all_b32 {dist = "NORM_B32"} +pto.vsts %copy3, %copy_out[%off_plus_192], %all_b32 {dist = "NORM_B32"} + +// Grouped use: reuse the same deinterleaved operands directly. +%prod_p0 = pto.vmul %x_p0, %scale_p0, %all_b32 : !pto.vreg<64xf32> +%prod_p1 = pto.vmul %x_p1, %scale_p1, %all_b32 : !pto.vreg<64xf32> +%prod_p2 = pto.vmul %x_p2, %scale_p2, %all_b32 : !pto.vreg<64xf32> +%prod_p3 = pto.vmul %x_p3, %scale_p3, %all_b32 : !pto.vreg<64xf32> + +%s0 = pto.vcgadd %prod_p0, %all_b32 : !pto.vreg<64xf32> +%s1 = pto.vcgadd %prod_p1, %all_b32 : !pto.vreg<64xf32> +%s2 = pto.vcgadd %prod_p2, %all_b32 : !pto.vreg<64xf32> +%s3 = pto.vcgadd %prod_p3, %all_b32 : !pto.vreg<64xf32> +%s01 = pto.vadd %s0, %s1, %slot8 : !pto.vreg<64xf32> +%s23 = pto.vadd %s2, %s3, %slot8 : !pto.vreg<64xf32> +%sum_block = pto.vadd %s01, %s23, %slot8 : !pto.vreg<64xf32> + +pto.vsts %sum_block, %sum_out[%group_off], %slot8 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for i = 0..255: + copy_out[off + i] = base[off + i] + scale_s + +for r = 0..7: + sum_out[group_off + r] = + reduce_i(base[off + r * 32 + i] * scale_s for i = 0..31) +``` + +Required assignment rule: + +```text +`broadcast` is layout-transparent and cheaply rematerializable by the optional +`vmi-layout-rematerialize` pass, but baseline assignment does not have to force +a separate contiguous broadcast just because a dense store exists. It may +choose a common deinterleaved compute layout for transfer-equivalent elementwise +ops and insert `ensure_layout` at the dense store. The required invariant is +that this choice is explicit in the assigned IR; `vmi-to-vpto` must not infer it +by inspecting both users. +``` + +### 3.41 Non-Rematerializable Value With Incompatible Users + +This is the non-cheap counterpart to section 3.18. A `masked_load` has explicit +mask and passthrough semantics, so layout assignment should not clone it as a +normal cheap load unless a dedicated rematerialization rule proves that clone +legal. The conflict is solved by inserting `ensure_layout` at one use site. + +VMI input: + +```text +%mask = pto.vmi.create_mask %c256 : index -> !pto.vmi.mask<256xpred> +%zero = pto.vmi.broadcast %c0_f32 : f32 -> !pto.vmi.vreg<256xf32> +%x = pto.vmi.masked_load %base[%off], %mask, %zero + : memref<256xf32>, !pto.vmi.mask<256xpred>, !pto.vmi.vreg<256xf32> + -> !pto.vmi.vreg<256xf32> + +pto.vmi.store %x, %copy_out[%off] + +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +pto.vmi.group_store %sum, %sum_out[%group_off], %c1 {num_groups = 8} +``` + +Assigned layouts: + +```text +%x, %zero for masked_load/store: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> + +%mask for masked_load/store: + !pto.vmi.mask<256xpred, #pto.vmi.layout> + +%x_for_reduce = pto.vmi.ensure_layout %x + : #pto.vmi.layout + -> #pto.vmi.layout + +%mask_for_reduce = pto.vmi.ensure_mask_layout %mask + : #pto.vmi.layout + -> #pto.vmi.layout + +%sum: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> +``` + +VPTO lowering result: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%slot8 = pto.pge_b32 "PAT_VL8" + +%zero0 = pto.vdup %c0_f32, %all_b32 : f32, !pto.mask -> !pto.vreg<64xf32> +%zero1 = pto.vdup %c0_f32, %all_b32 : f32, !pto.mask -> !pto.vreg<64xf32> +%zero2 = pto.vdup %c0_f32, %all_b32 : f32, !pto.mask -> !pto.vreg<64xf32> +%zero3 = pto.vdup %c0_f32, %all_b32 : f32, !pto.mask -> !pto.vreg<64xf32> + +%l0 = pto.vlds %base[%off] : !pto.ptr, index -> !pto.vreg<64xf32> +%l1 = pto.vlds %base[%off_plus_64] : !pto.ptr, index -> !pto.vreg<64xf32> +%l2 = pto.vlds %base[%off_plus_128] : !pto.ptr, index -> !pto.vreg<64xf32> +%l3 = pto.vlds %base[%off_plus_192] : !pto.ptr, index -> !pto.vreg<64xf32> + +%x0 = pto.vsel %l0, %zero0, %all_b32 : !pto.vreg<64xf32> +%x1 = pto.vsel %l1, %zero1, %all_b32 : !pto.vreg<64xf32> +%x2 = pto.vsel %l2, %zero2, %all_b32 : !pto.vreg<64xf32> +%x3 = pto.vsel %l3, %zero3, %all_b32 : !pto.vreg<64xf32> + +pto.vsts %x0, %copy_out[%off], %all_b32 {dist = "NORM_B32"} +pto.vsts %x1, %copy_out[%off_plus_64], %all_b32 {dist = "NORM_B32"} +pto.vsts %x2, %copy_out[%off_plus_128], %all_b32 {dist = "NORM_B32"} +pto.vsts %x3, %copy_out[%off_plus_192], %all_b32 {dist = "NORM_B32"} + +// ensure_layout contiguous -> deinterleaved=4 at the reduce use. +%x01_lo, %x01_hi = pto.vdintlv %x0, %x1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%x23_lo, %x23_hi = pto.vdintlv %x2, %x3 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%x_p0, %x_p2 = pto.vdintlv %x01_lo, %x23_lo + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +%x_p1, %x_p3 = pto.vdintlv %x01_hi, %x23_hi + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + +%s0 = pto.vcgadd %x_p0, %all_b32 : !pto.vreg<64xf32> +%s1 = pto.vcgadd %x_p1, %all_b32 : !pto.vreg<64xf32> +%s2 = pto.vcgadd %x_p2, %all_b32 : !pto.vreg<64xf32> +%s3 = pto.vcgadd %x_p3, %all_b32 : !pto.vreg<64xf32> +%s01 = pto.vadd %s0, %s1, %slot8 : !pto.vreg<64xf32> +%s23 = pto.vadd %s2, %s3, %slot8 : !pto.vreg<64xf32> +%sum_block = pto.vadd %s01, %s23, %slot8 : !pto.vreg<64xf32> + +pto.vsts %sum_block, %sum_out[%group_off], %slot8 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for i = 0..255: + copy_out[off + i] = base[off + i] + +for r = 0..7: + sum_out[group_off + r] = + reduce(base[off + r * 32 + 0 .. off + r * 32 + 31]) +``` + +Required assignment rule: + +```text +For non-rematerializable producers, assignment must insert an explicit use-site +materialization helper, such as contiguous -> deinterleaved=4. If that helper +has no supported materialization, the layout gate must diagnose before +vmi-to-vpto. `vmi-to-vpto` must not clone the masked_load or choose a +materialization after seeing both users. +``` + +### 3.42 `group_slots` `scf.for` Loop-Carried Accumulator + +Section 3.22 covers dense loop-carried values. Group-slot values need a +separate case because the loop-carried block argument has no dense lane +semantics outside the live group slots. + +VMI input: + +```text +%acc0 = pto.vmi.group_slot_load %init[%group_off], %c1 {num_groups = 8} + : !pto.ptr, index -> !pto.vmi.vreg<128xf32> + +%acc = scf.for %k = %c0 to %steps step %c1 + iter_args(%arg = %acc0) -> !pto.vmi.vreg<128xf32> { + %x = pto.vmi.group_load %base[%tile_off_k], %c16 + {num_groups = 8, group_size = 16} + : !pto.ptr, index -> !pto.vmi.vreg<128xf32> + %mask = pto.vmi.create_group_mask %c16 {num_groups = 8, group_size = 16} + : index -> !pto.vmi.mask<128xpred> + %sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} + %next = pto.vmi.addf %arg, %sum + scf.yield %next : !pto.vmi.vreg<128xf32> +} + +pto.vmi.group_store %acc, %out[%group_off], %c1 {num_groups = 8} +``` + +Assigned layouts: + +```text +%acc0, %arg, %sum, %next, %acc: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> + +%x: + !pto.vmi.vreg<128xf32, + #pto.vmi.layout> + +%mask: + !pto.vmi.mask<128xpred, + #pto.vmi.layout> +``` + +VPTO lowering result: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%slot8 = pto.pge_b32 "PAT_VL8" +%one_b32 = pto.pge_b32 "PAT_VL1" + +%acc0_block = pto.vsldb %init[%group_off], %c0_i16, %c0_i16, %one_b32 + : !pto.ptr, i16, i16, !pto.mask -> !pto.vreg<64xf32> + +%acc_block = scf.for %k = %c0 to %steps step %c1 + iter_args(%arg_block = %acc0_block) -> !pto.vreg<64xf32> { + %lo, %hi = pto.vldsx2 %base[%tile_off_k], "BDINTLV" + : !pto.ptr, index -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %lo_sum = pto.vcgadd %lo, %all_b32 : !pto.vreg<64xf32> + %hi_sum = pto.vcgadd %hi, %all_b32 : !pto.vreg<64xf32> + %sum_block = pto.vadd %lo_sum, %hi_sum, %slot8 : !pto.vreg<64xf32> + %next_block = pto.vadd %arg_block, %sum_block, %slot8 : !pto.vreg<64xf32> + scf.yield %next_block : !pto.vreg<64xf32> +} + +pto.vsts %acc_block, %out[%group_off], %slot8 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + out[group_off + r] = + init[group_off + r] + + sum_k reduce(base[tile_k, row_r, 0..15]) +``` + +Required assignment rule: + +```text +Loop-carried `group_slots` values are valid. The iter_arg, body block +argument, yield operand, loop result, and final group_store operand all carry +the same `group_slots(num_groups=8, slots=8)` layout. Ordinary dense consumers +inside the loop still require an explicit `group_broadcast` or diagnostic. +``` + +### 3.43 Internal Function Argument Boundary Materialization + +Section 3.25 covers a private function returning a VMI value. A callee argument +is the other direction of the same ABI problem: the callee body may require a +layout that is different from the layout naturally produced at a call site. + +The current implementation keeps the internal function VMI signature +contiguous and makes the callee-entry materialization explicit with +`ensure_layout` / `ensure_mask_layout`. This is less aggressive than +specializing the VMI function signature to `deinterleaved = 4`, but it preserves +the same invariant: after layout assignment, `vmi-to-vpto` lowers only from +explicit type and helper information and does not inspect the callee body while +lowering a call. + +VMI input: + +```text +func.func private @consume(%x: !pto.vmi.vreg<256xf32>, + %mask: !pto.vmi.mask<256xpred>, + %out: !pto.ptr, %group_off: index) { + %sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} + pto.vmi.group_store %sum, %out[%group_off], %c1 {num_groups = 8} + return +} + +func.func @caller(%base: !pto.ptr, %off: index, + %out: !pto.ptr, %group_off: index) { + %x = pto.vmi.load %base[%off] + : !pto.ptr, index -> !pto.vmi.vreg<256xf32> + %mask = pto.vmi.create_group_mask %c32 {num_groups = 8, group_size = 32} + : index -> !pto.vmi.mask<256xpred> + call @consume(%x, %mask, %out, %group_off) + : (!pto.vmi.vreg<256xf32>, !pto.vmi.mask<256xpred>, + !pto.ptr, index) -> () + return +} +``` + +Assigned layouts: + +```text +@consume argument %x: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> + +@consume argument %mask: + !pto.vmi.mask<256xpred, #pto.vmi.layout> + +inside @consume: + %x_split = pto.vmi.ensure_layout %x + : #pto.vmi.layout + -> #pto.vmi.layout + + %mask_split = pto.vmi.ensure_mask_layout %mask + : #pto.vmi.layout + -> #pto.vmi.layout + +@caller %x and %mask: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> + !pto.vmi.mask<256xpred, #pto.vmi.layout> + +%sum: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> +``` + +VPTO lowering result for the function boundary: + +```text +func.func private @consume(%x_p0: !pto.vreg<64xf32>, + %x_p1: !pto.vreg<64xf32>, + %x_p2: !pto.vreg<64xf32>, + %x_p3: !pto.vreg<64xf32>, + %m0: !pto.mask, + %m1: !pto.mask, + %m2: !pto.mask, + %m3: !pto.mask, + %out: !pto.ptr, + %group_off: index) { + // Callee-entry lowering of ensure_layout contiguous -> deinterleaved=4, + // block_elems=8. + %x01_lo, %x01_hi = pto.vdintlv %x_p0, %x_p1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %x23_lo, %x23_hi = pto.vdintlv %x_p2, %x_p3 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %x_d0, %x_d2 = pto.vdintlv %x01_lo, %x23_lo + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %x_d1, %x_d3 = pto.vdintlv %x01_hi, %x23_hi + : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + + %m01_lo, %m01_hi = pto.pdintlv_b32 %m0, %m1 + : !pto.mask, !pto.mask -> !pto.mask, !pto.mask + %m23_lo, %m23_hi = pto.pdintlv_b32 %m2, %m3 + : !pto.mask, !pto.mask -> !pto.mask, !pto.mask + %m_d0, %m_d2 = pto.pdintlv_b32 %m01_lo, %m23_lo + : !pto.mask, !pto.mask -> !pto.mask, !pto.mask + %m_d1, %m_d3 = pto.pdintlv_b32 %m01_hi, %m23_hi + : !pto.mask, !pto.mask -> !pto.mask, !pto.mask + + %slot8 = pto.pge_b32 "PAT_VL8" + %s0 = pto.vcgadd %x_d0, %m_d0 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %s1 = pto.vcgadd %x_d1, %m_d1 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %s2 = pto.vcgadd %x_d2, %m_d2 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %s3 = pto.vcgadd %x_d3, %m_d3 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %s01 = pto.vadd %s0, %s1, %slot8 : !pto.vreg<64xf32> + %s23 = pto.vadd %s2, %s3, %slot8 : !pto.vreg<64xf32> + %sum_block = pto.vadd %s01, %s23, %slot8 : !pto.vreg<64xf32> + pto.vsts %sum_block, %out[%group_off], %slot8 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + return +} + +func.func @caller(...) { + // Caller keeps the load and group mask in the contiguous function ABI layout. + %x0 = pto.vlds %base[%off] : !pto.ptr -> !pto.vreg<64xf32> + %x1 = pto.vlds %base[%off_plus_64] : !pto.ptr -> !pto.vreg<64xf32> + %x2 = pto.vlds %base[%off_plus_128] : !pto.ptr -> !pto.vreg<64xf32> + %x3 = pto.vlds %base[%off_plus_192] : !pto.ptr -> !pto.vreg<64xf32> + + %m0 = pto.pset_b32 "PAT_ALL" : !pto.mask + %m1 = pto.pset_b32 "PAT_ALL" : !pto.mask + %m2 = pto.pset_b32 "PAT_ALL" : !pto.mask + %m3 = pto.pset_b32 "PAT_ALL" : !pto.mask + + call @consume(%x0, %x1, %x2, %x3, %m0, %m1, %m2, %m3, %out, %group_off) + : (!pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.vreg<64xf32>, + !pto.vreg<64xf32>, !pto.mask, !pto.mask, + !pto.mask, !pto.mask, !pto.ptr, index) -> () + return +} +``` + +Memory result: + +```text +for r = 0..7: + out[group_off + r] = + reduce(base[off + r * 32 + 0 .. off + r * 32 + 31]) +``` + +Required assignment rule: + +```text +Private function boundary layout is explicit in the assigned function type and +callee-entry helpers. The current endpoint chooses a contiguous VMI function +ABI and inserts callee-entry materialization for the grouped body requirement. +`vmi-to-vpto` does not inspect the callee body while lowering the call and does +not inspect callers while lowering the callee block argument. + +Future optimization may specialize private VMI function signatures directly to +`deinterleaved = 4, block_elems = 8` when all call sites agree. That +optimization must still be expressed in the assigned VMI function type before +`vmi-to-vpto` runs. +``` + +Runtime closure: + +```text +lit: + test/lit/vmi/vmi_layout_assignment_call_argument_boundary.pto + test/lit/vmi/vmi_ptoas_call_boundary_vecscope.pto + +runtime SIM: + test/vpto/cases/vmi/private-call-argument-boundary-store + +ptoas pipeline: + vmi-layout-assignment inserts explicit callee-entry materialization + vmi-to-vpto physicalizes the call operands and callee body + ptoas then inlines the private physical helper before VPTO vecscope/backend + emission, so the backend never needs a physical VPTO vector function ABI +``` + +### 3.44 `masked_load` Grouped Tail Feeding S=32 Reduce + +This case connects the explicit `masked_load` tail model from section 3.30 with +grouped reduction. The load has no padding constant hidden in the op; inactive +lanes are provided by the passthrough value and excluded from the reduction by +the same grouped mask. + +VMI input: + +```text +%c25 = arith.constant 25 : index +%mask = pto.vmi.create_group_mask %c25 {num_groups = 8, group_size = 32} + : index -> !pto.vmi.mask<256xpred> +%zero = pto.vmi.broadcast %c0_f32 : f32 -> !pto.vmi.vreg<256xf32> +%x = pto.vmi.masked_load %base[%off], %mask, %zero + : memref<256xf32>, !pto.vmi.mask<256xpred>, !pto.vmi.vreg<256xf32> + -> !pto.vmi.vreg<256xf32> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 8} +pto.vmi.group_store %sum, %out[%group_off], %c1 {num_groups = 8} +``` + +Assigned layouts: + +```text +%mask for masked_load: + !pto.vmi.mask<256xpred, #pto.vmi.layout> + +%zero, %x for masked_load: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> + +%x_for_reduce = pto.vmi.ensure_layout %x: + #pto.vmi.layout + -> #pto.vmi.layout + +%mask_for_reduce: + pto.vmi.create_group_mask %c25 {num_groups = 8, group_size = 32} + -> !pto.vmi.mask<256xpred, + #pto.vmi.layout> + +%sum: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> +``` + +Lowering: + +```text +%all_b32 = pto.pge_b32 "PAT_ALL" +%slot8 = pto.pge_b32 "PAT_VL8" + +// masked_load direct lowering stays contiguous. +%m0, %m1, %m2, %m3 = materialize contiguous create_group_mask(c25, S=32) +%z0, %z1, %z2, %z3 = vdup zero +%l0 = pto.vlds %base[%off] +%l1 = pto.vlds %base[%off_plus_64] +%l2 = pto.vlds %base[%off_plus_128] +%l3 = pto.vlds %base[%off_plus_192] +%x0 = pto.vsel %l0, %z0, %m0 : !pto.vreg<64xf32> +%x1 = pto.vsel %l1, %z1, %m1 : !pto.vreg<64xf32> +%x2 = pto.vsel %l2, %z2, %m2 : !pto.vreg<64xf32> +%x3 = pto.vsel %l3, %z3, %m3 : !pto.vreg<64xf32> + +// ensure_layout contiguous -> deinterleaved=4, block_elems=8. +%x01_lo, %x01_hi = pto.vdintlv %x0, %x1 +%x23_lo, %x23_hi = pto.vdintlv %x2, %x3 +%x_p0, %x_p2 = pto.vdintlv %x01_lo, %x23_lo +%x_p1, %x_p3 = pto.vdintlv %x01_hi, %x23_hi + +// The reduce-side grouped mask is not built by guessing the final group-slot +// predicate image. It is first materialized as the same contiguous grouped +// mask used by masked_load, then converted to the reduce layout with predicate +// deinterleave. This keeps predicate reordering identical to the data +// reordering above. +%rm0, %rm1, %rm2, %rm3 = materialize contiguous create_group_mask(c25, S=32) +%rm01_lo, %rm01_hi = pto.pdintlv_b32 %rm0, %rm1 +%rm23_lo, %rm23_hi = pto.pdintlv_b32 %rm2, %rm3 +%mask_p0, %mask_p2 = pto.pdintlv_b32 %rm01_lo, %rm23_lo +%mask_p1, %mask_p3 = pto.pdintlv_b32 %rm01_hi, %rm23_hi + +%s0 = pto.vcgadd %x_p0, %mask_p0 : !pto.vreg<64xf32> +%s1 = pto.vcgadd %x_p1, %mask_p1 : !pto.vreg<64xf32> +%s2 = pto.vcgadd %x_p2, %mask_p2 : !pto.vreg<64xf32> +%s3 = pto.vcgadd %x_p3, %mask_p3 : !pto.vreg<64xf32> +%s01 = pto.vadd %s0, %s1, %slot8 : !pto.vreg<64xf32> +%s23 = pto.vadd %s2, %s3, %slot8 : !pto.vreg<64xf32> +%sum_block = pto.vadd %s01, %s23, %slot8 : !pto.vreg<64xf32> + +pto.vsts %sum_block, %out[%group_off], %slot8 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + out[group_off + r] = + reduce(base[off + r * 32 + 0 .. off + r * 32 + 24]) +``` + +Required assignment rule: + +`masked_load` and `group_reduce` must share the same grouped mask layout. The +passthrough value defines inactive loaded lanes, while the reduce mask defines +participation. Assignment materializes two explicit mask values when needed: +one contiguous value for `masked_load`, and one deinterleaved value for +`group_reduce_addf`. `vmi-to-vpto` lowers the deinterleaved +`create_group_mask` by materializing the contiguous grouped predicate chunks +and then applying `pdintlv_b32` in the same tree shape as the data +`vdintlv`. It does not walk from `group_reduce_addf` to the mask producer to +choose or reject the support path. + +Assignment may select a deinterleaved S=32 load layout only when the rounded +physical reads are memory-safe; otherwise it must diagnose or use a future +stable gather fallback. + +Runtime coverage: + +```text +test/vpto/cases/vmi/masked-load-group-tail-s32-reduce-store +``` + +### 3.45 Dynamic S=32 `create_group_mask` + +This is the dynamic-shape form of section 3.44. The active column count is an +SSA `index`, not a constant. The semantic mask is still grouped: + +```text +lane i active iff (i % 32) < active_cols +``` + +VMI input: + +```text +%mask = pto.vmi.create_group_mask %active_cols + {num_groups = 8, group_size = 32} + : index -> !pto.vmi.mask<256xpred> +``` + +Assigned layouts: + +```text +%mask for masked_load: + !pto.vmi.mask<256xb32, #pto.vmi.layout> + +%mask for S=32 group_reduce: + !pto.vmi.mask<256xb32, + #pto.vmi.layout> +``` + +Contiguous VPTO lowering for one b32 physical chunk: + +```text +%active_i32 = arith.index_cast %active_cols : index to i32 +%active_nonneg = arith.maxsi %active_i32, %c0_i32 : i32 +%active_clamped = arith.minui %active_nonneg, %c32_i32 : i32 + +%all = pto.pset_b32 "PAT_ALL" : !pto.mask +%lane = pto.vci %c0_i32 : i32 -> !pto.vreg<64xi32> +%row = pto.vshrs %lane, %c5_i16, %all + : !pto.vreg<64xi32>, i16, !pto.mask -> !pto.vreg<64xi32> +%row_base = pto.vshls %row, %c5_i16, %all + : !pto.vreg<64xi32>, i16, !pto.mask -> !pto.vreg<64xi32> +%col = pto.vsub %lane, %row_base, %all + : !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.mask + -> !pto.vreg<64xi32> +%m = pto.vcmps %col, %active_clamped, %all, "lt" + : !pto.vreg<64xi32>, i32, !pto.mask -> !pto.mask +``` + +For `deinterleaved = 4, block_elems = 8`, lowering first emits four contiguous +chunks with the sequence above, then applies the same predicate deinterleave +tree used by section 3.44: + +```text +%rm0, %rm1, %rm2, %rm3 = dynamic contiguous grouped masks +%rm01_lo, %rm01_hi = pto.pdintlv_b32 %rm0, %rm1 +%rm23_lo, %rm23_hi = pto.pdintlv_b32 %rm2, %rm3 +%mask_p0, %mask_p2 = pto.pdintlv_b32 %rm01_lo, %rm23_lo +%mask_p1, %mask_p3 = pto.pdintlv_b32 %rm01_hi, %rm23_hi +``` + +Current coverage validates both IR lowering and runtime behavior: + +```text +test/lit/vmi/vmi_layout_assignment_create_group_mask_s32_dynamic.pto +test/vpto/cases/vmi/dynamic-create-group-mask-s32-reduce-store +``` + +The runtime case passes `active_cols` as a kernel scalar argument and casts it +to `index` inside `pto.vecscope`. This keeps scalar materialization outside +`vmi-to-vpto`; the lowering pass only consumes the current +`create_group_mask` operand. + +### 3.46 `extf` Value And Derived Elementwise Value Both Stored + +This case fixes where contiguous materialization belongs when one widened value +is used directly by a store and also by a layout-transparent elementwise chain +that is stored. + +VMI input: + +```text +%a = pto.vmi.load %in[%off] + : memref<128xf16> -> !pto.vmi.vreg<128xf16> +%k = pto.vmi.broadcast %k1 + : f32 -> !pto.vmi.vreg<128xf32> + +%w = pto.vmi.extf %a + : !pto.vmi.vreg<128xf16> -> !pto.vmi.vreg<128xf32> +%t1 = pto.vmi.mulf %w, %k + : !pto.vmi.vreg<128xf32>, !pto.vmi.vreg<128xf32> + -> !pto.vmi.vreg<128xf32> + +pto.vmi.store %t1, %out1[%off] +pto.vmi.store %w, %out2[%off] +``` + +Hard-legalized assigned layouts: + +```text +%a: + !pto.vmi.vreg<128xf16, #pto.vmi.layout> + +%w, %k, %t1: + !pto.vmi.vreg<128xf32, #pto.vmi.layout> + +%t1_c = pto.vmi.ensure_layout %t1: + #pto.vmi.layout -> #pto.vmi.layout +pto.vmi.store %t1_c, %out1[%off] + +%w_c = pto.vmi.ensure_layout %w: + #pto.vmi.layout -> #pto.vmi.layout +pto.vmi.store %w_c, %out2[%off] +``` + +Baseline VPTO lowering result: + +```text +%a0 = pto.vlds %in[%off] {dist = "NORM"} + : !pto.ptr, index -> !pto.vreg<128xf16> + +%w_p0 = pto.vcvt %a0, PAT_ALL_B16 {part = "EVEN"} + : !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<64xf32> +%w_p1 = pto.vcvt %a0, PAT_ALL_B16 {part = "ODD"} + : !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<64xf32> + +%k_p0 = pto.vdup %k1, PAT_ALL_B32 : f32, !pto.mask -> !pto.vreg<64xf32> +%k_p1 = pto.vdup %k1, PAT_ALL_B32 : f32, !pto.mask -> !pto.vreg<64xf32> + +%t1_p0 = pto.vmul %w_p0, %k_p0, PAT_ALL_B32 : !pto.vreg<64xf32> +%t1_p1 = pto.vmul %w_p1, %k_p1, PAT_ALL_B32 : !pto.vreg<64xf32> + +// ensure_layout for the first store. +%t1_0, %t1_1 = pto.vintlv %t1_p0, %t1_p1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> + -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +pto.vsts %t1_0, %out1[%off], %all_b32 {dist = "NORM_B32"} +pto.vsts %t1_1, %out1[%off_plus_64], %all_b32 {dist = "NORM_B32"} + +// ensure_layout for the second store. +%w_0, %w_1 = pto.vintlv %w_p0, %w_p1 + : !pto.vreg<64xf32>, !pto.vreg<64xf32> + -> !pto.vreg<64xf32>, !pto.vreg<64xf32> +pto.vsts %w_0, %out2[%off], %all_b32 {dist = "NORM_B32"} +pto.vsts %w_1, %out2[%off_plus_64], %all_b32 {dist = "NORM_B32"} +``` + +Memory result: + +```text +for i = 0..127: + out1[off + i] = f32(in[off + i]) * k1 + out2[off + i] = f32(in[off + i]) +``` + +Optimization pass result: + +```text +// vmi-layout-fold may remove both ensure_layout ops if the target +// supports store lowering that consumes deinterleaved=2 and writes contiguous +// row-major memory. +pto.vmi.store %t1, %out1[%off] +pto.vmi.store %w, %out2[%off] +``` + +Optimized VPTO lowering result: + +```text +pto.vstsx2 %t1_p0, %t1_p1, %out1[%off], "INTLV_B32", PAT_ALL_B32 + : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.ptr, index, + !pto.mask + +pto.vstsx2 %w_p0, %w_p1, %out2[%off], "INTLV_B32", PAT_ALL_B32 + : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.ptr, index, + !pto.mask +``` + +Required assignment and optimization rule: + +```text +Hard legalization may always preserve `%w` and `%t1` in deinterleaved=2 and +insert use-site ensure_layout before ordinary stores. This is correct because +the layout change is explicit at the store use. + +Consumer folding is optional. It may remove the ensure_layout only when the +store itself can locally prove the same contiguous memory effect from the +source layout. vmi-to-vpto must not scan the `%w` producer or both store users +to decide this. +``` + +### 3.47 Type-Parametric Group Reduce Rule + +The group-reduce layout rule is parameterized by the element width, not by f32 +case names. + +```text +E = sizeof(T) +VLaneElems = 32B / E +L = 256B / E +S = logical_lane_count / num_groups +``` + +The canonical grouped-reduce layouts are: + +```text +Packed group-slot rule: + K is the physical slot capacity of one packed group-result chunk. + For VCG-style packed reductions, K = 8. + G does not have to be divisible by K; the final chunk may be partial. + active_groups(chunk c) = min(K, G - c * K). + +S == VLaneElems: + source/mask layout = contiguous + result layout = group_slots(num_groups=G, slots=8) + +S == 2 * VLaneElems: + source/mask layout = deinterleaved=2 + result layout = group_slots(num_groups=G, slots=8) + +S == 4 * VLaneElems: + source/mask layout = deinterleaved=4 + result layout = group_slots(num_groups=G, slots=8) + +S >= L && S % L == 0: + source/mask layout = contiguous + result layout = group_slots(num_groups=G, slots=1) +``` + +Concrete shape table: + +```text +T VLaneElems L packed cases row-local cases +f32 8 64 S=8, S=16, S=32 S=64, S=128, ... +i32 8 64 S=8, S=16, S=32 S=64, S=128, ... +f16 16 128 S=16, S=32, S=64 S=128, S=256, ... +i16 16 128 S=16, S=32, S=64 S=128, S=256, ... +f8 32 256 cast to f32 before grouped reduce +i8 32 256 S=32, S=64, S=128 S=256, S=512, ... +``` + +These non-f32 cases are part of the type-generic layout/lowering design. If a +typed reduce op admits the element type and the target capability registry +accepts it, assignment must use the same `VLaneElems/L/S` formula instead of +adding per-type shape special cases. Any f32-only behavior in the current +implementation is staged implementation status, not the intended design limit. +For the current baseline, `f8` remains a storage and cast-boundary type for +group reduction. Integer `i8/i16/i32` grouped reductions are direct VMI +operations when their group shape matches a registered table row. + +### 3.48 16-bit Typed Group Reduce, `S = VLaneElems = 16` + +This case covers both `f16` and `i16`. The element width is the same, so the +layout and VPTO instruction skeleton are identical. The VMI op name carries the +semantic difference: + +```text +f16: pto.vmi.group_reduce_addf ... {reassoc} +i16: pto.vmi.group_reduce_addi ... +``` + +VMI-shaped input: + +```text +// Floating form. +%xf = pto.vmi.load %base_f16[%off] + : memref<128xf16> -> !pto.vmi.vreg<128xf16> +%mf = pto.vmi.create_group_mask %c16 {num_groups = 8, group_size = 16} + : index -> !pto.vmi.mask<128xpred> +%sumf = pto.vmi.group_reduce_addf %xf, %mf {num_groups = 8, reassoc} +pto.vmi.group_store %sumf, %out_f16[%group_off], %c1 {num_groups = 8} + +// Integer form. +%xi = pto.vmi.load %base_i16[%off] + : memref<128xi16> -> !pto.vmi.vreg<128xi16> +%mi = pto.vmi.create_group_mask %c16 {num_groups = 8, group_size = 16} + : index -> !pto.vmi.mask<128xpred> +%sumi = pto.vmi.group_reduce_addi %xi, %mi {num_groups = 8} +pto.vmi.group_store %sumi, %out_i16[%group_off], %c1 {num_groups = 8} +``` + +Assigned layouts: + +```text +%xf, %mf, %xi, %mi: + #pto.vmi.layout + +%sumf: + !pto.vmi.vreg<128xf16, #pto.vmi.layout> + +%sumi: + !pto.vmi.vreg<128xi16, #pto.vmi.layout> +``` + +VPTO lowering shape: + +```text +%x0 = pto.vlds %base[%off] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<128xT16> + +%all_b16 = pto.pge_b16 "PAT_ALL" +%slot8_b16 = pto.pge_b16 "PAT_VL8" + +%sum0 = pto.vcgadd %x0, %all_b16 + : !pto.vreg<128xT16>, !pto.mask -> !pto.vreg<128xT16> + +pto.vsts %sum0, %out[%group_off], %slot8_b16 {dist = "NORM_B16"} + : !pto.vreg<128xT16>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + out[group_off + r] = reduce_T16(base[off + r * 16 + 0 .. 15]) +``` + +### 3.49 16-bit Typed Group Reduce, `S = 2 * VLaneElems = 32` + +This case covers both `f16` and `i16`. Each logical row is 64B and must be +split into two 32B VLane fragments before `vcgadd`. + +VMI-shaped input: + +```text +%x = pto.vmi.load %base[%off] + : memref<256xT16> -> !pto.vmi.vreg<256xT16> +%mask = pto.vmi.create_group_mask %c32 {num_groups = 8, group_size = 32} + : index -> !pto.vmi.mask<256xpred> +%sum = pto.vmi.group_reduce_add{f|i} %x, %mask {num_groups = 8} +pto.vmi.group_store %sum, %out[%group_off], %c1 {num_groups = 8} +``` + +Assigned layouts: + +```text +%x, %mask: + #pto.vmi.layout + +%sum: + !pto.vmi.vreg<256xT16, #pto.vmi.layout> +``` + +VPTO lowering shape: + +```text +%x_p0, %x_p1 = pto.vldsx2 %base[%off], "DINTLV_B16" + : !pto.ptr, index -> !pto.vreg<128xT16>, !pto.vreg<128xT16> + +%all_b16 = pto.pge_b16 "PAT_ALL" +%slot8_b16 = pto.pge_b16 "PAT_VL8" + +%s0 = pto.vcgadd %x_p0, %all_b16 + : !pto.vreg<128xT16>, !pto.mask -> !pto.vreg<128xT16> +%s1 = pto.vcgadd %x_p1, %all_b16 + : !pto.vreg<128xT16>, !pto.mask -> !pto.vreg<128xT16> +%sum0 = pto.vadd %s0, %s1, %slot8_b16 + : !pto.vreg<128xT16>, !pto.vreg<128xT16>, !pto.mask + -> !pto.vreg<128xT16> + +pto.vsts %sum0, %out[%group_off], %slot8_b16 {dist = "NORM_B16"} + : !pto.vreg<128xT16>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + out[group_off + r] = reduce_T16(base[off + r * 32 + 0 .. 31]) +``` + +### 3.50 16-bit Typed Group Reduce, `S = 4 * VLaneElems = 64` + +This is the four-fragment packed case for both `f16` and `i16`. + +Assigned layouts: + +```text +%x, %mask: + #pto.vmi.layout + +%sum: + !pto.vmi.vreg<512xT16, #pto.vmi.layout> +``` + +VPTO lowering shape: + +```text +%x_p0, %x_p1, %x_p2, %x_p3 = materialize deinterleaved=4 input + : four !pto.vreg<128xT16> + +%all_b16 = pto.pge_b16 "PAT_ALL" +%slot8_b16 = pto.pge_b16 "PAT_VL8" + +%s0 = pto.vcgadd %x_p0, %all_b16 : !pto.vreg<128xT16> +%s1 = pto.vcgadd %x_p1, %all_b16 : !pto.vreg<128xT16> +%s2 = pto.vcgadd %x_p2, %all_b16 : !pto.vreg<128xT16> +%s3 = pto.vcgadd %x_p3, %all_b16 : !pto.vreg<128xT16> + +%s01 = pto.vadd %s0, %s1, %slot8_b16 : !pto.vreg<128xT16> +%s23 = pto.vadd %s2, %s3, %slot8_b16 : !pto.vreg<128xT16> +%sum0 = pto.vadd %s01, %s23, %slot8_b16 : !pto.vreg<128xT16> + +pto.vsts %sum0, %out[%group_off], %slot8_b16 {dist = "NORM_B16"} + : !pto.vreg<128xT16>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + out[group_off + r] = reduce_T16(base[off + r * 64 + 0 .. 63]) +``` + +#### 3.50.1 Partial Packed `S = 64` Reductions + +This is the same `S = 4 * VLaneElems` lowering family as section 3.50, but it +covers `G` values that do not fill every packed group-result chunk. The key +point is that `slots = 8` is a physical capacity, not a promise that every +chunk contains eight valid group results. + +The result layout remains: + +```text +!pto.vmi.vreg<(G * 64)xf16, #pto.vmi.layout> +``` + +The lowering computes per result chunk: + +```text +K = 8 +chunk c active groups A(c) = min(K, G - c * K) + +source active lanes per deinterleaved part for chunk c: + A(c) * VLaneElems = A(c) * 16 f16 lanes + +reduce input mask: + PAT_VL(A(c) * 16) + +combine/store mask: + PAT_VL(A(c)) +``` + +For full chunks, `A(c) = 8`, so the reduce input mask is `PAT_ALL` for f16 +and the combine/store mask is `PAT_VL8`. For partial chunks, masks are +required for correctness. The semantic source mask produced by +`pto.vmi.create_group_mask` must also materialize only the valid source lanes; +the reduce lowering should not treat padding lanes as active data. + +##### `G = 4`: `256xf16, num_groups = 4` + +VMI-shaped input: + +```text +%x = pto.vmi.load %base[%off] + : memref<256xf16> -> !pto.vmi.vreg<256xf16> +%mask = pto.vmi.create_group_mask %c64 {num_groups = 4, group_size = 64} + : index -> !pto.vmi.mask<256xpred> +%sum = pto.vmi.group_reduce_addf %x, %mask {num_groups = 4, reassoc} +pto.vmi.group_store %sum, %out[%group_off], %c1 {num_groups = 4} +``` + +Assigned layouts: + +```text +%x, %mask: + #pto.vmi.layout + +%sum: + !pto.vmi.vreg<256xf16, #pto.vmi.layout> +``` + +VPTO lowering shape for the only result chunk: + +```text +%x_p0, %x_p1, %x_p2, %x_p3 = materialize deinterleaved=4, block_elems=8 input + : four !pto.vreg<128xf16> + +%lane64_b16 = pto.pge_b16 "PAT_VL64" // A * 16 = 4 * 16 +%slot4_b16 = pto.pge_b16 "PAT_VL4" + +%s0 = pto.vcgadd %x_p0, %lane64_b16 : !pto.vreg<128xf16> +%s1 = pto.vcgadd %x_p1, %lane64_b16 : !pto.vreg<128xf16> +%s2 = pto.vcgadd %x_p2, %lane64_b16 : !pto.vreg<128xf16> +%s3 = pto.vcgadd %x_p3, %lane64_b16 : !pto.vreg<128xf16> + +%s01 = pto.vadd %s0, %s1, %slot4_b16 : !pto.vreg<128xf16> +%s23 = pto.vadd %s2, %s3, %slot4_b16 : !pto.vreg<128xf16> +%sum0 = pto.vadd %s01, %s23, %slot4_b16 : !pto.vreg<128xf16> + +pto.vsts %sum0, %out[%group_off], %slot4_b16 {dist = "NORM_B16"} + : !pto.vreg<128xf16>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..3: + out[group_off + r] = reduce_f16(base[off + r * 64 + 0 .. 63]) + +sum0 lanes 4..127 are not semantic for this VMI result. +``` + +##### `G = 8`: full packed chunk + +This is section 3.50. There is one result chunk with `A = 8`: + +```text +source mask = PAT_ALL // 8 * 16 = 128 f16 lanes +combine/store = PAT_VL8 +result layout = group_slots(num_groups=8, slots=8) +``` + +##### `G = 12`: full chunk plus partial chunk + +This case needs two packed result chunks: + +```text +result layout = group_slots(num_groups=12, slots=8) +result arity = ceil(12 / 8) = 2 +``` + +Chunk 0 handles groups `0..7`: + +```text +A(0) = 8 +source mask = PAT_ALL +combine/store = PAT_VL8 +``` + +Chunk 1 handles groups `8..11`: + +```text +A(1) = 4 +source mask = PAT_VL64 +combine/store = PAT_VL4 +``` + +Implementation checklist for this family: + +```text +layout attr: + slots=8 should be legal even when num_groups is not divisible by 8. + slot_block(g) = g / 8 and slot_lane(g) = g % 8 are still well-defined. + +layout assignment: + packed VCG-style group_reduce results keep slots=8. + +mask materialization: + create_group_mask must not activate padding lanes in partial chunks. + For chunk c, source active lanes are A(c) * VLaneElems. + +vmi-to-vpto group_reduce: + use A(c) from result layout slots and num_groups. + combine masks use PAT_VL(A(c)). + input vcgadd consumes the physical mask parts, which must already encode + PAT_VL(A(c) * VLaneElems) for all-true grouped masks. + +vmi-to-vpto group_store: + use A(c) to build the store predicate. + output group offset for chunk c is c * slots. +``` + +### 3.51 16-bit Typed Group Reduce, `S = L = 128` + +This is the first row-local full-physical-chunk case for both `f16` and `i16`. +The canonical result is row-local `slots = 1`, not packed `slots = 8`. + +VMI-shaped input: + +```text +%x = pto.vmi.load %base[%off] + : memref<1024xT16> -> !pto.vmi.vreg<1024xT16> +%mask = pto.vmi.create_group_mask %c128 {num_groups = 8, group_size = 128} + : index -> !pto.vmi.mask<1024xpred> +%sum = pto.vmi.group_reduce_add{f|i} %x, %mask {num_groups = 8} +pto.vmi.group_store %sum, %out[%group_off], %c1 {num_groups = 8} +``` + +Assigned layouts: + +```text +%x, %mask: + #pto.vmi.layout + +%sum: + !pto.vmi.vreg<1024xT16, #pto.vmi.layout> +``` + +VPTO lowering shape: + +```text +%all_b16 = pto.pge_b16 "PAT_ALL" +%slot1_b16 = pto.pge_b16 "PAT_VL1" + +// Repeated for r = 0..7. +%x_r = pto.vlds %base[%row_off_r] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<128xT16> + +// Floating-point keeps the same physical result type. +%sumf_r = pto.vcadd %x_r, %all_b16 + : !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<128xf16> + +// Integer VCADD widens internally; VMI restores the declared i16 type. +%wide_r = pto.vcadd %x_r, %all_b16 + : !pto.vreg<128xi16>, !pto.mask -> !pto.vreg<64xi32> +%sumi_r = pto.vbitcast %wide_r + : !pto.vreg<64xi32> -> !pto.vreg<128xi16> + +pto.vsts %sum{f|i}_r, %out[%group_off_plus_r], %slot1_b16 {dist = "NORM_B16"} + : !pto.vreg<128xT16>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + out[group_off + r] = reduce_T16(base[off + r * 128 + 0 .. 127]) +``` + +### 3.52 32-bit Typed Group Reduce + +This case covers both `f32` and `i32`. The element width is the same, so +`VLaneElems = 8` and `L = 64` for both. Floating-point uses +`group_reduce_addf` with `reassoc`; integer uses `group_reduce_addi`. + +Example for `S = 2 * VLaneElems = 16`: + +```text +%x: + !pto.vmi.vreg<128xT32, #pto.vmi.layout> + +%sum: + !pto.vmi.vreg<128xT32, #pto.vmi.layout> +``` + +VPTO lowering shape: + +```text +%x_p0, %x_p1 = pto.vldsx2 %base[%off], "DINTLV_B32" + : !pto.ptr, index -> !pto.vreg<64xT32>, !pto.vreg<64xT32> + +%all_b32 = pto.pge_b32 "PAT_ALL" +%slot8_b32 = pto.pge_b32 "PAT_VL8" + +%s0 = pto.vcgadd %x_p0, %all_b32 + : !pto.vreg<64xT32>, !pto.mask -> !pto.vreg<64xT32> +%s1 = pto.vcgadd %x_p1, %all_b32 + : !pto.vreg<64xT32>, !pto.mask -> !pto.vreg<64xT32> +%sum0 = pto.vadd %s0, %s1, %slot8_b32 + : !pto.vreg<64xT32>, !pto.vreg<64xT32>, !pto.mask + -> !pto.vreg<64xT32> + +pto.vsts %sum0, %out[%group_off], %slot8_b32 {dist = "NORM_B32"} + : !pto.vreg<64xT32>, !pto.ptr, !pto.mask +``` + +The same formula gives: + +```text +S=8: + contiguous, slots=8, one vcgadd. + +S=32: + deinterleaved=4, slots=8, four vcgadd plus vadd tree. + +S=64: + contiguous, slots=1, row-local vcgadd plus vcadd. + +S=128: + contiguous, slots=1, row-local multi-chunk accumulation. +``` + +### 3.53 Integer Semantics And Invalid Typed Reductions + +Integer group reduction is not a variant of `group_reduce_addf`; it requires a +typed integer op: + +```text +%sum = pto.vmi.group_reduce_addi %x, %mask {num_groups = G} +``` + +Required semantics: + +```text +inactive lanes contribute integer zero +addition uses the target's normal integer add behavior +wrap/saturating variants must be represented by distinct ops if both are needed +signedness does not affect add, but does affect future max/min integer reduces +``` + +Required invalid cases: + +```text +pto.vmi.group_reduce_addf with integer element type -> verifier error +pto.vmi.group_reduce_addi with floating-point element type -> verifier error +pto.vmi.group_reduce_addi with an integer width other than i8/i16/i32 + -> verifier error +S not in {VLaneElems, 2*VLaneElems, 4*VLaneElems} and not a full-chunk multiple + -> layout-contract diagnostic +``` + +### 3.54 8-bit Floating Group Reduce + +There is no direct f8 `vcgadd` grouped reduction in the current target model, +but f8 supports cast to an accumulator type. The semantic path is: + +```text +f8 storage -> cast/extf to f32 accumulator -> group_reduce_addf on f32 +``` + +Here `f8` is only the cast source and the memory element type. The reduction +itself is a f32 accumulator operation. + +The group size remains a logical-lane property. For example, reducing eight +rows of 32 f8 elements produces the same logical result as reducing eight rows +of 32 f32 accumulator elements after extension. + +VMI-shaped input: + +```text +%x8 = pto.vmi.load %base_f8[%off] + : memref<256xf8> -> !pto.vmi.vreg<256xf8> +%x32 = pto.vmi.extf %x8 + : !pto.vmi.vreg<256xf8> -> !pto.vmi.vreg<256xf32> +%mask = pto.vmi.create_group_mask %c32 {num_groups = 8, group_size = 32} + : index -> !pto.vmi.mask<256xpred> +%sum = pto.vmi.group_reduce_addf %x32, %mask {num_groups = 8, reassoc} +pto.vmi.group_store %sum, %out_f32[%group_off], %c1 {num_groups = 8} +``` + +Assigned layouts: + +```text +%x8: + !pto.vmi.vreg<256xf8, #pto.vmi.layout> + +%x32, %mask: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> + !pto.vmi.mask<256xb32, #pto.vmi.layout> + +%sum: + !pto.vmi.vreg<256xf32, #pto.vmi.layout> +``` + +VPTO lowering shape: + +```text +%x8_packed = pto.vlds %base_f8[%off] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<256xf8> + +%all_b8 = pto.pge_b8 "PAT_ALL" +%all_b32 = pto.pge_b32 "PAT_ALL" +%slot8_b32 = pto.pge_b32 "PAT_VL8" + +%x32_p0 = pto.vcvt %x8_packed, %all_b8 {part = "P0"} + : !pto.vreg<256xf8>, !pto.mask -> !pto.vreg<64xf32> +%x32_p1 = pto.vcvt %x8_packed, %all_b8 {part = "P1"} + : !pto.vreg<256xf8>, !pto.mask -> !pto.vreg<64xf32> +%x32_p2 = pto.vcvt %x8_packed, %all_b8 {part = "P2"} + : !pto.vreg<256xf8>, !pto.mask -> !pto.vreg<64xf32> +%x32_p3 = pto.vcvt %x8_packed, %all_b8 {part = "P3"} + : !pto.vreg<256xf8>, !pto.mask -> !pto.vreg<64xf32> + +%s0 = pto.vcgadd %x32_p0, %all_b32 : !pto.vreg<64xf32> +%s1 = pto.vcgadd %x32_p1, %all_b32 : !pto.vreg<64xf32> +%s2 = pto.vcgadd %x32_p2, %all_b32 : !pto.vreg<64xf32> +%s3 = pto.vcgadd %x32_p3, %all_b32 : !pto.vreg<64xf32> +%s01 = pto.vadd %s0, %s1, %slot8_b32 : !pto.vreg<64xf32> +%s23 = pto.vadd %s2, %s3, %slot8_b32 : !pto.vreg<64xf32> +%sum0 = pto.vadd %s01, %s23, %slot8_b32 : !pto.vreg<64xf32> + +pto.vsts %sum0, %out_f32[%group_off], %slot8_b32 {dist = "NORM_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask +``` + +Memory result: + +```text +for r = 0..7: + out_f32[group_off + r] = + reduce_f32(f32(base_f8[off + r * 32 + 0 .. 31])) +``` + +Direct f8 grouped reduction is invalid: + +```text +pto.vmi.group_reduce_addf %x8, %mask + : !pto.vmi.vreg<256xf8>, !pto.vmi.mask<256xpred> + -> verifier or layout-contract diagnostic +``` + +### 3.55 8-bit Integer Group Reduce + +The target exposes same-type i8 `vcgadd` for 32B-block group classes and a +widening i8-to-i16 `vcadd` for full-row reduction. VMI keeps a same-type i8 +contract in both cases: + +```text +i8 source -> group_reduce_addi -> i8 group-slot result +``` + +Packed 32B-block example: + +```text +%x8 = pto.vmi.load %base_i8[%off] + : memref<256xi8> -> !pto.vmi.vreg<256xi8> +%mask = pto.vmi.create_group_mask %c32 {num_groups = 8, group_size = 32} + : index -> !pto.vmi.mask<256xpred> +%sum = pto.vmi.group_reduce_addi %x8, %mask {num_groups = 8} +pto.vmi.group_store %sum, %out_i8[%group_off], %c1 {num_groups = 8} +``` + +Assigned layouts: + +```text +%x8, %mask: + #pto.vmi.layout + +%sum: + !pto.vmi.vreg<8xi8, #pto.vmi.layout> +``` + +VPTO lowering shape: + +```text +%all_b8 = pto.pge_b8 "PAT_ALL" +%slot8_b8 = pto.pge_b8 "PAT_VL8" +%sum0 = pto.vcgadd %x8, %all_b8 + : !pto.vreg<256xi8>, !pto.mask -> !pto.vreg<256xi8> +pto.vsts %sum0, %out_i8[%group_off], %slot8_b8 {dist = "NORM_B8"} +``` + +For an aligned full row (`S = 256`), lowering uses widening only internally: + +```text +%wide = pto.vcadd %x8, %all_b8 + : !pto.vreg<256xi8>, !pto.mask -> !pto.vreg<128xi16> +%sum_i8 = pto.vbitcast %wide + : !pto.vreg<128xi16> -> !pto.vreg<256xi8> +``` + +The low i8 lane contains the same-type wraparound result. Explicit +`extsi`/`extui` before reduction remains available when the algorithm itself +requires a wider accumulator, but widening is not required by the direct i8 +group-reduce contract. + +### 3.56 Full 256-Bin Distribution Histogram + +Histogram is not modeled as `group_reduce`. A group reduce maps source lanes to +result slots by lane/group position. A histogram maps each active source lane +to a result bin by the source value itself. + +VMI-shaped input: + +```text +%src = pto.vmi.load %src_base[%src_off] + : memref -> !pto.vmi.vreg +%mask = pto.vmi.create_mask %active_lanes + : index -> !pto.vmi.mask +%acc = pto.vmi.load %acc_base[%acc_off] + : memref<256xui16> -> !pto.vmi.vreg<256xui16> +%hist = pto.vmi.vdhist %acc, %src, %mask + : !pto.vmi.vreg<256xui16>, !pto.vmi.vreg, + !pto.vmi.mask -> !pto.vmi.vreg<256xui16> +pto.vmi.store %hist, %out[%out_off] +``` + +Logical semantics: + +```text +for b = 0..255: + hist[b] = acc[b] + +for i = 0..N-1: + if mask[i]: + hist[src[i]] += 1 +``` + +Assigned layouts: + +```text +%src: + !pto.vmi.vreg> + +%mask: + !pto.vmi.mask> + +%acc, %hist: + !pto.vmi.vreg<256xui16, #pto.vmi.layout> +``` + +The `256xui16` accumulator/result is one logical VMI value but two physical +VPTO vector registers: + +```text +physical result part0 = logical bins 0..127 +physical result part1 = logical bins 128..255 +``` + +For `N = 256`, VPTO lowering shape: + +```text +%src0 = pto.vlds %src_base[%src_off] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<256xui8> + +%acc_lo = pto.vlds %acc_base[%acc_off + 0] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<128xui16> +%acc_hi = pto.vlds %acc_base[%acc_off + 128] {dist = "NORM"} + : !pto.ptr -> !pto.vreg<128xui16> + +%hist_lo = pto.dhistv2 %acc_lo, %src0, %mask0, %bin0 + : !pto.vreg<128xui16>, !pto.vreg<256xui8>, !pto.mask, i32 + -> !pto.vreg<128xui16> +%hist_hi = pto.dhistv2 %acc_hi, %src0, %mask0, %bin1 + : !pto.vreg<128xui16>, !pto.vreg<256xui8>, !pto.mask, i32 + -> !pto.vreg<128xui16> + +pto.vsts %hist_lo, %out[%out_off + 0], %all_b16 {dist = "NORM_B16"} +pto.vsts %hist_hi, %out[%out_off + 128], %all_b16 {dist = "NORM_B16"} +``` + +Memory result: + +```text +for b = 0..127: + out[out_off + b] = acc_base[acc_off + b] + + count(i where mask[i] && src_base[src_off + i] == b) + +for b = 128..255: + out[out_off + b] = acc_base[acc_off + b] + + count(i where mask[i] && src_base[src_off + i] == b) +``` + +For `N > 256`, the source is processed in contiguous 256-lane chunks. The two +histogram accumulator parts are carried through all chunks: + +```text +%lo = %acc_lo +%hi = %acc_hi + +for source chunk c in logical order: + %chunk_mask = mask chunk c + if c is the final partial chunk: + %chunk_mask = %chunk_mask & valid-lane-prefix-for-this-chunk + + %lo = pto.dhistv2 %lo, %src_c, %chunk_mask, %bin0 + %hi = pto.dhistv2 %hi, %src_c, %chunk_mask, %bin1 + +result physical parts = [%lo, %hi] +``` + +Tail source lanes are expressed only through the b8 mask. Padding lanes in the +last physical source chunk must be masked off before `pto.dhistv2`; they are +not padding values. + +The VMI op does not expose `#bin`. `#bin` is a VPTO range selector forced by +the physical result width: + +```text +ui8 value domain = 256 bins +complete histogram = 256 x ui16 = 512B +one VPTO vreg result = 128 x ui16 = 256B +``` + +Therefore VMI represents one logical `256xui16` result and `vmi-to-vpto` +locally emits the low-range and high-range VPTO histogram updates. + +### 3.57 Full 256-Bin Cumulative Histogram + +The desired VMI surface shape mirrors `dhist`: + +```text +%hist = pto.vmi.vchist %acc, %src, %mask + : !pto.vmi.vreg<256xui16>, !pto.vmi.vreg, + !pto.vmi.mask -> !pto.vmi.vreg<256xui16> +``` + +The intended logical semantics is a full cumulative histogram: + +```text +dist[b] = count(i where mask[i] && src[i] == b) + +hist[0] = acc[0] + dist[0] +for b = 1..255: + hist[b] = acc[b] + dist[0] + dist[1] + ... + dist[b] +``` + +The current VPTO/VISA documentation only states that `CHISTv2` computes a +`uint16 Cumulative histogram` over the selected bin range. It does not state +whether the high-range call with `#bin = 1` returns: + +```text +global cumulative: + result[j] = count(src <= 128 + j) + +or range-local cumulative: + result[j] = count(128 <= src <= 128 + j) +``` + +These two interpretations have different VMI lowerings. If the hardware result +is global cumulative, the full VMI lowering is the same low/high split as +`dhist`, replacing `pto.dhistv2` with `pto.chistv2`. If the hardware result is +range-local cumulative, the high half also needs the total low-half count added +to every high-half bin: + +```text +%lo = pto.chistv2 %acc_lo, %src0, %mask0, %bin0 +%hi_local = pto.chistv2 %acc_hi, %src0, %mask0, %bin1 + +%low_total = materialize count(src <= 127) from the low-half result +%low_total_vec = broadcast %low_total to every high-half bin +%hi = pto.vadd %hi_local, %low_total_vec, %all_b16 +``` + +That correction path also requires a designed way to materialize and broadcast +the low-half total. Since baseline VMI does not support arbitrary vector +extract, the range-local CHISTv2 interpretation remains unsupported until that +materialization path is explicit. + +The baseline design therefore treats `pto.vmi.vchist` as a semantic op whose +exact lowering is gated by a target semantic capability: + +```text +if target documents or validation proves CHISTv2 high range is global: + lower as two pto.chistv2 calls +elif target documents or validation proves CHISTv2 high range is range-local: + lower as pto.chistv2 low/high plus explicit high-half correction only after + low-total materialization support is designed +else: + VMI-UNSUPPORTED: pto.vmi.vchist requires a verified CHISTv2 range semantics contract +``` + +This boundary is deliberate. `pto.vmi.vdhist` is fully defined because +distribution bins are independent across the low/high split. `pto.vmi.vchist` +has cross-range prefix semantics, so VMI must not guess the high-half behavior +from the VPTO op name alone. diff --git a/docs/designs/vmi-layout-relation-rematerialization-design.md b/docs/designs/vmi-layout-relation-rematerialization-design.md new file mode 100644 index 0000000000..9b943f8e70 --- /dev/null +++ b/docs/designs/vmi-layout-relation-rematerialization-design.md @@ -0,0 +1,236 @@ +# VMI Layout Relation-Aware Rematerialization Design + +本文描述 VMI layout optimization 中 relation-aware rematerialization 的设计。 +目标是让 `vmi-layout-assignment` 只产生 legal baseline IR，把跨 layout +relation 的优化放到显式 `ensure_layout` 上完成。 + +## 1. Motivation + +`vmi-layout-assignment` 已经负责三件 hard legalization 工作： + +```text +1. 为每个 VMI value 选择 concrete layout +2. 在不匹配的 use-site 插入 ensure_layout / ensure_mask_layout +3. 保证 vmi-to-vpto 只需要 local lowering information +``` + +对 `ext` 这类 width-changing op，assignment 的 baseline 可以保守选择： + +```text +ext f16 -> f32: + source = contiguous + result = deinterleaved=2 +``` + +如果下游 `truncf f32 -> f8` 要求 source 为 `deinterleaved=4`，assignment 会 +显式插入： + +```text +%e = pto.vmi.extf %x + : !vreg<..., layout> + -> !vreg<..., layout> + +%e4 = pto.vmi.ensure_layout %e + : !vreg<..., layout> + -> !vreg<..., layout> +``` + +这个 IR 已经合法，但不是最优。优化 pass 可以从显式 helper 出发，把 relation +应用到 producer： + +```text +ensure_layout(ext(src), resultLayout) + => ext(ensure_layout(src, derivedSourceLayout)) +``` + +这样 assignment 不需要做 consumer-driven global propagation，也不需要在多 +consumer 冲突时引入 cost model。 + +## 2. Goals + +```text +1. assignment 保持 hard legalization baseline，不做 ext relation propagation。 +2. relation-aware optimization 从显式 ensure_layout 出发。 +3. 多 consumer 冲突由 use-site helper + rematerialization 解决。 +4. vmi-to-vpto 仍只消费当前 op 的 operand/result layout，不扫描上下文。 +5. 变换必须是局部、确定、可验证的 IR rewrite。 +``` + +非目标： + +```text +1. 不做 ComputeY1 专用 pattern。 +2. 不在 assignment 中实现全局 cost model。 +3. 不通过 vmi-to-vpto 猜 producer/consumer relation。 +4. 第一阶段不做 trunc/narrow relation remat。 +``` + +## 3. Optimization Model + +relation-aware remat 以 `ensure_layout` 为唯一触发点： + +```text +%wanted = pto.vmi.ensure_layout %source : sourceLayout -> targetLayout +``` + +如果 `%source` 的 producer 可以在 `targetLayout` 或 relation 派生出的 operand +layout 下重新创建等价结果，则用 cloned producer 替换 helper。 + +### 3.1 Layout-Transparent Producer Remat + +对 layout-transparent elementwise op： + +```text +ensure_layout(op(a, b), L) + => op(ensure_layout(a, L), ensure_layout(b, L)) +``` + +适用对象包括纯 elementwise data ops： + +```text +addf/addi/subf/subi/mulf/muli/divf/minf/maxf +andi/ori/xori/shli/shrui/shrsi +negf/absf/absi/sqrt/exp/ln/relu/not +fma +select, when data operands and mask layout requirements can be kept explicit +``` + +第一阶段可以先覆盖 ComputeY1 需要的 `mulf`，但实现形态应按 op family 泛化。 + +### 3.2 Widen Ext Relation Remat + +对 widening `ext`： + +```text +ensure_layout(ext(src), resultLayout) + => ext(ensure_layout(src, sourceLayout)) +``` + +其中： + +```text +resultFactor = sourceFactor * widenFactor +``` + +例子： + +```text +ext f16 -> f32, widenFactor = 2 +target result layout = deinterleaved=4 +derived source layout = deinterleaved=2 +``` + +`deinterleaved=1` 等价于 contiguous。 + +### 3.3 Producer Fold After Remat + +relation remat 可能暴露 producer-side helper： + +```text +ensure_layout(load(...), deinterleaved=2) +``` + +这类 helper 应由 `vmi-layout-fold` 吸收到 producer 或 consumer： + +```text +load contiguous + ensure_layout to deinterleaved=2 + => load result deinterleaved=2 +``` + +因此推荐优化 pipeline 在 remat 后再次运行 fold： + +```text +vmi-layout-assignment + -> canonicalize/cse + -> vmi-layout-rematerialize + -> canonicalize/cse + -> vmi-layout-fold + -> canonicalize/cse + -> vmi-layout-sink-materialization + -> canonicalize/cse +``` + +## 4. Multi-Consumer Conflict + +如果一个 `ext` result 有两个 consumer： + +```text +consumer A requires deinterleaved=2 +consumer B requires deinterleaved=4 +``` + +assignment 不需要判断哪个更优。它可以选择稳定 baseline，例如 `deinterleaved=2`， +并为另一个 use 插入 helper： + +```text +%e2 = pto.vmi.extf %x : contiguous -> deinterleaved=2 +consumer_a(%e2) + +%e4 = pto.vmi.ensure_layout %e2 : deinterleaved=2 -> deinterleaved=4 +consumer_b(%e4) +``` + +remat 再把第二个 use 优化成 cloned producer： + +```text +%x2 = pto.vmi.ensure_layout %x : contiguous -> deinterleaved=2 +%e4 = pto.vmi.extf %x2 : deinterleaved=2 -> deinterleaved=4 +consumer_b(%e4) +``` + +原 `%e2` 仍服务 `consumer_a`。这样不需要 assignment 做全局 cost selection。 + +## 5. ComputeY1 Shape + +baseline assignment 可能产生： + +```text +%x32 = extf %x16 // result deinterleaved=2 +%s32 = extf %scale16 // result deinterleaved=2 +%m = mulf %x32, %s32 // result deinterleaved=2 +%m4 = ensure_layout %m // deinterleaved=2 -> deinterleaved=4 +%y = truncf %m4 +``` + +remat/fold 后目标 IR： + +```text +%x16_d2 = load ... // folded deinterleaved=2 load +%x32_d4 = extf %x16_d2 // deinterleaved=2 -> deinterleaved=4 + +%scale16_d2 = group_broadcast_load ... // folded/assigned deinterleaved=2 +%scale32_d4 = extf %scale16_d2 // deinterleaved=2 -> deinterleaved=4 + +%m4 = mulf %x32_d4, %scale32_d4 +%y = truncf %m4 +``` + +关键点： + +```text +1. truncf 只通过 ensure_layout 表达自己的 source layout requirement。 +2. remat 不需要识别 quant 语义。 +3. ext relation 是 local rule。 +4. load/group_broadcast_load 的物理优化由 fold 或 producer capability 处理。 +``` + +## 6. Lowering Contract + +`vmi-to-vpto` 的 contract 不变： + +```text +1. 不扫描 ext 的 users。 +2. 不扫描 producer chain 来猜 layout。 +3. 只根据当前 op 的 operand/result layout lower。 +``` + +relation-aware remat 必须在 `vmi-to-vpto` 前把 IR 显式改写为： + +```text +%x = pto.vmi.load ... -> !vreg<..., layout> +%e = pto.vmi.extf %x + : !vreg<..., layout> + -> !vreg<..., layout> +``` + +之后 lowering 只消费这个 local shape。 diff --git a/docs/designs/vmi-layout-relation-rematerialization-implementation.md b/docs/designs/vmi-layout-relation-rematerialization-implementation.md new file mode 100644 index 0000000000..de721f67a1 --- /dev/null +++ b/docs/designs/vmi-layout-relation-rematerialization-implementation.md @@ -0,0 +1,406 @@ +# VMI Layout Relation-Aware Rematerialization Implementation Plan + +本文是 `vmi-layout-relation-rematerialization-design.md` 的实现计划。目标是 +扩展现有 `vmi-layout-rematerialize` / `vmi-layout-fold` 优化，让 assignment +保持 legal baseline，并从显式 `ensure_layout` 中恢复更好的 producer layout。 + +## 1. Current Baseline + +当前 pipeline 中相关 pass： + +```text +vmi-layout-assignment: + chooses concrete layouts + inserts ensure_layout / ensure_mask_layout / ensure_mask_granularity + +vmi-layout-fold: + folds selected ensure_layout helpers into layout-aware producers/consumers + current coverage includes store-side fold, load -> ensure_layout producer + fold, and inverse nested ensure_layout fold + +vmi-layout-rematerialize: + replaces ensure_* around cheap construction producers + current data coverage: splat constant, broadcast, iota + current mask coverage: create_mask, create_group_mask, constant_mask + +vmi-layout-sink-materialization: + sinks matching operand-side helpers through pure elementwise ops + it does not currently rewrite result-side ensure_layout(op(...), L) +``` + +ComputeY1-like IR currently remains suboptimal because assignment emits: + +```text +ensure_layout(mulf(ext(...), ext(...)), deinterleaved=4) +``` + +but remat does not yet: + +```text +1. hoist result-side ensure_layout through mulf +2. rematerialize ext under a requested result layout +3. expose foldable load/group_broadcast_load helpers +``` + +## 2. Support APIs + +Add support-layer helpers in `VMILayoutSupport`. + +### 2.1 Widen Relation Query + +```cpp +FailureOr getWidenSourceLayoutForResultLayout( + VMIVRegType sourceType, + VMIVRegType resultType, + VMILayoutAttr requestedResultLayout, + std::string *reason = nullptr) const; +``` + +Semantics: + +```text +1. source/result lane count must match. +2. result element width must be an integer multiple of source element width. +3. first implementation supports widenFactor 2 and 4. +4. requestedResultLayout must be contiguous or deinterleaved(block_elems=1). +5. requested result factor F must be divisible by widenFactor K. +6. derived source factor is F / K. +7. derived source factor 1 means contiguous. +8. derived source/result layout pair must be accepted by ext support gates. +``` + +Examples: + +```text +f16 -> f32, requested result deinterleaved=4 + => source deinterleaved=2 + +f16 -> f32, requested result deinterleaved=2 + => source contiguous + +f8 -> f32, requested result deinterleaved=4 + => source contiguous +``` + +### 2.2 Ext Support Gates + +Update `getExtFSupport`, `getExtSISupport`, and `getExtUISupport` so they accept +relation-rematerialized local shapes: + +```text +source layout: + contiguous or deinterleaved(S, block_elems=1) + +result layout: + deinterleaved(S * widenFactor, block_elems=1) +``` + +Keep group_slots integer extension behavior unchanged. + +Reject: + +```text +1. result layout that is not deinterleaved for dense ext. +2. block_elems != 1 in this first implementation. +3. source/result arity that does not satisfy resultArity = factor * sourceArity. +4. unsupported element width relation. +``` + +`vmi-to-vpto` ext lowering already works from physical source/result arity. If +support admits `source deinterleaved=2 -> result deinterleaved=4`, lowering must +be covered by tests. + +## 3. Rematerialize Pass Changes + +Extend `VMILayoutRematerialize.cpp` around `VMIEnsureLayoutOp`. + +Recommended ordering for one helper: + +```text +try relation-aware ext remat +try result-side layout-transparent producer remat +try existing cheap construction remat +``` + +The pass should use a helper worklist. When one rewrite creates new +`ensure_layout` helpers, enqueue them so the same pass can continue locally. + +### 3.1 Ext Remat Pattern + +Match: + +```text +%wanted = pto.vmi.ensure_layout %old +%old = pto.vmi.extf %src +``` + +where `%wanted` has `requestedResultType`. + +Rewrite: + +```text +derivedSourceLayout = + support.getWidenSourceLayoutForResultLayout(srcType, requestedResultType, + requestedResultLayout) + +%src2 = materialize source to derivedSourceLayout +%new = pto.vmi.extf %src2 : derivedSourceType -> requestedResultType +replace %wanted with %new +``` + +Equivalent patterns are needed for: + +```text +pto.vmi.extf +pto.vmi.extsi +pto.vmi.extui +``` + +The source materialization step should: + +```text +1. reuse %src if it already has derivedSourceLayout. +2. create pto.vmi.ensure_layout otherwise. +3. enqueue the new helper for further remat/fold opportunities. +``` + +### 3.2 Layout-Transparent Result Helper Remat + +Match: + +```text +%wanted = pto.vmi.ensure_layout %old +%old = pto.vmi.mulf %lhs, %rhs +``` + +Rewrite: + +```text +%lhs2 = ensure_layout %lhs : lhsLayout -> requestedLayout +%rhs2 = ensure_layout %rhs : rhsLayout -> requestedLayout +%new = pto.vmi.mulf %lhs2, %rhs2 : requestedLayout +replace %wanted with %new +``` + +Initial op coverage: + +```text +mulf +addf/addi/subf/subi/muli/divf/minf/maxf +andi/ori/xori/shli/shrui/shrsi +negf/absf/absi/sqrt/exp/ln/relu/not +fma +``` + +Optional later coverage: + +```text +cmpf/cmpi: + result is mask, so this belongs with ensure_mask_layout. + +select: + requires coordinated data and mask layout/granularity handling. +``` + +The pass must preserve op attributes exactly. + +### 3.3 Existing Cheap Producer Remat + +Keep current behavior for: + +```text +splat pto.vmi.constant +pto.vmi.broadcast +pto.vmi.iota +pto.vmi.create_mask +pto.vmi.create_group_mask +pto.vmi.constant_mask +``` + +These remain direct remat cases and do not require relation queries. + +## 4. Fold Pass Interaction + +Relation remat may create producer-side helpers: + +```text +ensure_layout(load(...), deinterleaved=2) +ensure_layout(group_broadcast_load(...), deinterleaved=2) +``` + +`vmi-layout-fold` should absorb these when the producer can directly materialize +the requested layout. + +Existing load fold should use producer capability, not helper materialization +capability. A load may directly produce a requested contiguous or +deinterleaved=2/4 block_elems=1 result layout even when the helper conversion +from the old load layout to the requested layout would not be a legal register +materialization. + +Add fold coverage if missing for: + +```text +group_broadcast_load result layout requested as deinterleaved=2/block_elems=1 +group_slot_broadcast_load result layout requested as deinterleaved=2/block_elems=1 +``` + +The fold pass must still be local: + +```text +load/group_broadcast_load + ensure_layout + => cloned/retyped producer with requested result layout +``` + +It must not inspect downstream `ext` or `trunc`. + +## 5. Pipeline + +Use a pipeline with fold after remat: + +```text +vmi-layout-assignment + -> canonicalize/cse + -> vmi-layout-rematerialize + -> canonicalize/cse + -> vmi-layout-fold + -> canonicalize/cse + -> vmi-layout-sink-materialization + -> canonicalize/cse + -> pto-validate-vmi-layout-ir + -> vmi-to-vpto +``` + +The first fold handles helpers already emitted by assignment. The second fold +handles helpers exposed by relation-aware remat. + +If later result-side remat and operand-side sink need to alternate for longer +chains, the driver may repeat: + +```text +vmi-layout-rematerialize +canonicalize/cse +vmi-layout-fold +canonicalize/cse +``` + +Keep the first implementation single-pass unless tests prove a fixed point is +needed. + +## 6. Tests + +Add focused lit tests. + +### 6.1 Direct Ext Remat + +Input shape: + +```text +load f16 +extf f16 -> f32 +ensure_layout ext result deinterleaved=2 -> deinterleaved=4 +truncf f32 -> f8 +``` + +Check after: + +```text +vmi-layout-rematerialize +``` + +```text +extf source is deinterleaved=2 +extf result is deinterleaved=4 +old ensure_layout is gone +``` + +Check after: + +```text +vmi-layout-rematerialize -vmi-layout-fold -vmi-to-vpto +``` + +```text +load uses deinterleaved load lowering when fold is available +extf lowers from local source/result arity +``` + +### 6.2 Elementwise Result Helper Remat + +Input shape: + +```text +extf lhs -> deinterleaved=2 +extf rhs -> deinterleaved=2 +mulf lhs, rhs -> deinterleaved=2 +ensure_layout mulf result -> deinterleaved=4 +truncf +``` + +Check: + +```text +mulf is cloned/rebuilt with deinterleaved=4 operands/results +each ext is rematerialized as source deinterleaved=2 -> result deinterleaved=4 +no ensure_layout remains between mulf and truncf +``` + +### 6.3 Multi-Consumer Conflict + +Input shape: + +```text +ext result deinterleaved=2 +consumer A uses deinterleaved=2 +consumer B has ensure_layout to deinterleaved=4 +``` + +Check: + +```text +original ext remains for consumer A +new cloned ext feeds consumer B +no global layout selection is required +``` + +### 6.4 ComputeY1 + +Run: + +```text +pto-test-opt compute_y1_to_fp8_fp16_vmi.pto \ + -vmi-layout-assignment \ + -vmi-layout-rematerialize \ + -vmi-layout-fold \ + -vmi-to-vpto +``` + +Expected: + +```text +x load can become deinterleaved=2 and lower through deinterleaved load support +scale path can keep the E2B-compatible deinterleaved layout +mulf/truncf path has no deinterleaved=2 -> deinterleaved=4 helper immediately +before truncf +``` + +## 7. Non-Goals And Follow-Ups + +Do not implement in this change: + +```text +1. assignment relation propagation. +2. global layout cost model. +3. trunc/narrow relation remat. +4. cloning memory loads in remat without going through explicit fold support. +5. context-sensitive vmi-to-vpto lowering. +``` + +Follow-ups: + +```text +1. Add narrow relation remat for selected trunc patterns after widen is stable. +2. Add select/cmp mask-aware result helper remat. +3. Consider a fixed-point layout optimization pipeline if long chains need it. +4. Move repeated op-family cloning utilities into a shared helper if the pass + grows beyond the first ext/elementwise implementation. +``` diff --git a/docs/designs/vmi-layout-request-propagation.md b/docs/designs/vmi-layout-request-propagation.md new file mode 100644 index 0000000000..266b721059 --- /dev/null +++ b/docs/designs/vmi-layout-request-propagation.md @@ -0,0 +1,1681 @@ +# VMI Layout Request Propagation + +This document describes VMI layout request propagation for layout assignment +and local layout rewrites. It is intentionally independent from any single +optimization case such as `group_reduce -> truncf -> group_broadcast`. + +The propagator answers this question: + +```text +Given one or more requested layouts for VMI values, can the surrounding IR be +rewritten so those semantic values are available in those layouts? +``` + +It must not work by clearing existing layout attributes. Existing layout +attributes are the current IR state. A pass expresses desired changes by adding +value-layout requests. + +## Legacy Propagation Model + +Before the shared propagator, VMI layout propagation was split between layout +assignment and post-assignment rematerialization. + +`vmi-layout-assignment` used an equivalence-class model. The pass collected +values whose layouts must be identical, then assigned one layout to each class. +This worked for layout-transparent relations: + +```text +source/result of elementwise ops +bitcast-like ops +control-flow forwarded values +mask/data same-layout relations +``` + +In that model, assignment propagation means "same layout flows through this +relation". It does not cross relations where the connected values naturally +have different layouts. Casts, reductions, channel transforms, broadcasts, and +other width- or shape-changing operations cannot be represented as a simple +equivalence-class union. When the requested layout on one side did not match +the layout chosen for the other side, assignment had to connect the two sides +with `ensure_layout` materialization. + +`vmi-layout-rematerialize` then performed a second, optimization-oriented +propagation after assignment. It started from explicit `ensure_layout` ops, +looked at the producer feeding the `ensure_layout`, and decided whether the +layout conversion could be pushed through or removed by recomputing the +producer. This was intentionally local and peephole-like: + +```text +ensure_layout result requested as layout B + inspect producer in layout A + ask whether this producer can be recomputed for B cheaply + clone/recompute the producer or one of its operands + replace this one use with the recomputed value + continue recursively while the local cost model allows it +``` + +This pass therefore acted like an instcombine-style rematerialization pass over +layout helper IR. It did not build a graph-level layout solution first. The +propagation happened as a sequence of one-by-one IR rewrites and replacements, +driven by the `ensure_layout` currently being optimized. + +## Refactor Motivation + +The split model caused the same layout knowledge to appear in two places. +Assignment needed rules for natural/preferred layouts and same-layout +constraints. Rematerialization also needed to know which non-equivalence +operations could be crossed, and what layout should appear on the other side +after crossing. As more cast and reduction cases were added, the duplicated +parts became the risky part: + +```text +vcvt/ext/trunc layout facts +group-value cast layout facts +reduce input/result layout facts +broadcast/group-slot layout support checks +``` + +Adding another rematerialization case would have required copying more of the +assignment-side rule set into the peephole optimizer. That makes the result +order-sensitive and hard to reason about: assignment may choose one layout, +rematerialization may rediscover a related layout later, and the two passes can +silently disagree about whether the same op relation is legal. + +The refactor centralizes the reusable part as op-local transfer relations and a +value-layout propagator: + +```text +propagate(value, layout) + record the requested layout for the semantic SSA value + run transfer relations through defining ops and users + derive uniquely implied layouts on connected values + record conflicts when a second layout is needed + materialize unresolved conflicts with ensure_layout during apply +``` + +Assignment remains responsible for policy: it chooses the initial seeds from +true layout decision points such as ext/trunc/reduce results, loads, stores, and +target-specific boundary requirements. Same-layout ops are relations, not +decision roots. Consumer operand requests are not producer facts unless a +support-checked direct producer can really adopt that layout. + +Rematerialization should then consume the same transfer/support helpers instead +of owning a second copy of cast/reduce layout logic. Its job becomes local IR +cleanup after a consistent layout table exists: folding helper IR, cloning cheap +producers when that is profitable, or removing redundant `ensure_layout` chains. +It should not be the only place where non-equivalence layout propagation is +defined. + +## Core Assignments + +The propagator is a pass-local object. It does not rewrite IR while requests +are being added. Its core assignment table is value-centric: + +```text +assignments: + Value -> VMIValueLayoutAssignment + +worklist: + Primary Value/layout facts that were just added and still need to be + propagated through defining and user op transfers. Conflict layouts do not + enter this worklist unless a later rematerialization/fork planner explicitly + creates a new producer instance for that alternate layout. +``` + +The assignment uses short field names: + +```text +VMIValueLayoutAssignment: + layout: + The layout that apply must make available for this SSA value. + + conflicts: + Extra layouts required for this SSA value. There are two forms: + def-side, meaning the SSA value itself must also be available in another + layout; and use-side, meaning one operand use must see the value through + another layout. + +VMILayoutConflict: + operand: + Empty for a def-side conflict. Otherwise, the OpOperand to rewrite during + apply if the conflict remains material. Do not create a fake operand for a + def-side conflict. + + layout: + The extra layout required for this SSA value or operand use. +``` + +The implementation can represent this directly as: + +```text +VMILayoutConflict: + OpOperand *operand // nullptr means def-side + VMILayoutAttr layout +``` + +Conflict uniqueness is checked by conflict form: + +```text +def-side: + key = layout + duplicate layout is a no-op + +use-side: + key = operand + duplicate operand with the same layout is a no-op + duplicate operand with a different layout is a hard conflict +``` + +`assignment.layout` is not a promise that the defining op can directly produce +that layout. It is the layout that apply must make available for the semantic +value. Apply implements it by rewriting the source value's VMI type when the +value is type-rewriteable inside the current rewrite scope, or by creating a +primary `ensure_layout` value at a boundary when the source value's type cannot +be rewritten. + +A conflict does not overwrite `assignment.layout`. It records an additional +fork requirement: + +```text +def-side conflict: + if assignment.layout differs from conflict.layout: + materialize ensure_layout assignment.layout -> conflict.layout at the + def/boundary + +use-side conflict: + if assignment.layout differs from conflict.layout: + insert ensure_layout assignment.layout -> conflict.layout before + conflict.operand + replace only conflict.operand + else: + the conflict is an identity fork and is dropped +``` + +This keeps the `assignment.layout` model value-bound while still representing +multiple layout requirements. Operand layout requirements are not stored in a +separate global table and are not encoded inside `VMILayoutAttr`. + +The propagator does not rank competing layouts for a value. The caller decides +the initial request order. The first layout accepted for a value becomes +`assignment.layout`; subsequent different layouts become conflicts. Only +`assignment.layout` is a producer-transfer fact. A conflict means an alternate +layout must be materialized as a fork from the primary value; it is not a claim +that the original defining op result has that alternate layout. Cost-based +layout choice is outside this first implementation and would change only the +merge policy, not transfer relations. + +Do not pre-build separate tables for current layouts, requests, +materializations, rewrites, or conflicts: + +```text +current layout: + Read from value.getType() on demand. If the type has no layout, the current + layout is unknown, not contiguous. + +request: + Immediately merge into assignments. If `assignment.layout` is newly added, + enqueue the value/layout fact. Conflict layouts are recorded in + `assignment.conflicts` but are not propagated through producers by this + utility. The worklist de-duplicates by `(Value, VMILayoutAttr)` so the same + primary fact is propagated once. The first implementation does not keep a + separate request log. + +materialization: + Derive during apply by comparing each value assignment's layout with its + conflicts and with the current IR type. + +rewrite: + Derive during apply from assignments by finding defining ops whose result + types need to change. + +hard conflict: + Detect while merging one operand's required layout. Def-side conflicting + value requests are represented as conflicts; same-operand conflicting + requests still fail. +``` + +## API Shape + +The basic API is: + +```text +request(value, layout): + request that the SSA value be available in layout. If this differs from the + existing `assignment.layout`, record a def-side conflict. + +request(operand, layout): + request that operand.get() be available in layout for this operand only. This + does not create the primary assignment for the source value. If the source + value's `assignment.layout` is absent or different, record a use-side + requirement on that operand. + +run(): + propagate until the worklist is empty + +apply(): + rewrite in-scope VMI value types to their assigned layouts + materialize boundary values and conflicts with ensure_layout forks +``` + +The merge operation is the only place that mutates the core table: + +```text +request(value, layout): + if assignments[value].layout exists and differs from layout: + add assignments[value].conflicts[{def, layout}] + else if assignments[value].layout is absent: + assignments[value].layout = layout + push (value, layout) to worklist + +request(operand, layout): + value = operand.get() + if assignments[value].layout exists and differs from layout: + if operand already has a different conflict layout: + report hard conflict + add or update assignments[value].conflicts[operand] = layout + else if assignments[value].layout exists and equals layout: + no-op + else: + create assignments[value] if needed + add assignments[value].conflicts[operand] = layout +``` + +`request(operand, layout)` must not delegate to `request(value, layout)`. A +consumer operand requirement is not a producer fact. It becomes a local +materialization request unless a separate value request or transfer relation +chooses the same primary layout for the source value. + +Transfer relations call the API that matches the derived fact. For +layout-transparent ops, an operand value fact can derive the op result primary +layout, while the other operands receive operand-local layout requirements. +A result value fact similarly derives operand-local requirements. Cast inverse +relations also call `request(operand, layout)` because a result layout +determines the layout needed by that cast operand, not necessarily the +producer's global primary layout. The first implementation may promote an +operand-local request into a producer seed in two narrow cases: + +- the source value is defined by a direct layout producer, such as a supported + `load`, splat-like `constant`, `broadcast`, `iota`, or `group_broadcast` + whose source group-slot layout and requested result layout pass the + target-support query; +- the source value is defined by a single-use layout-transparent op, every data + operand can directly produce the requested layout, and neither the value nor + any data operand already has a different primary assignment or current layout. + +The second case is a local rematerialization-style choice for layout-free IR. +It does not chase arbitrary producer chains, and it must not override an +existing primary assignment, an assignment seed, or an explicit current layout. +Assignment must add producer value seeds before consumer operand requests so a +consumer request cannot steal the primary layout from an ext/trunc/reduce value +that already has a preferred layout. If an alternate layout is cheap but +requires cloning or sinking an existing producer chain, a later +rematerialization/fold pass may remove the inserted `ensure_layout` by cloning +or folding the producer at the use site. + +If a subsequent request asks for a different extra layout, that request is +recorded as a def-side or use-side +conflict depending on which overload created it. + +Existing explicit type layouts are part of the current IR state, not the +request input. A pass should not first write an explicit layout and then ask +the propagator to rediscover it. Instead, the pass requests the desired +layout directly. If an explicit layout already exists, the propagator reads it +on demand as the current state when deciding whether materialization is needed. + +An existing explicit layout is not a lock. If a request asks for a layout that +differs from the current IR type or `assignment.layout`, the propagator +records the requested value/layout fact and continues propagation. During +apply, the propagator decides whether the defining op can be rewritten to that +layout by ordinary in-scope type rewrite because sourceValue is +type-rewriteable, or whether a primary `ensure_layout` fork is needed at a +boundary. + +## Implementation Shape + +The first implementation should be a small utility, not a replacement for every +layout pass at once: + +```text +include/PTO/Transforms/VMILayoutPropagation.h +lib/PTO/Transforms/VMILayoutPropagation.cpp +``` + +The public type should expose value layout requests, operand-local value layout +requests, fixed-point propagation, and final IR application: + +```text +class VMILayoutPropagator { + LogicalResult request(Value value, VMILayoutAttr layout); + LogicalResult request(OpOperand &operand, VMILayoutAttr layout); + LogicalResult run(); + LogicalResult apply(RewriterBase &rewriter); + + VMILayoutAttr getRequestedOrCurrentLayout(Value value) const; + VMILayoutAttr getRequestedLayout(Value value) const; +}; +``` + +`getRequestedLayout` returns only `assignment.layout`, not conflict layouts. +`getRequestedOrCurrentLayout` reads `assignment.layout` first, then the current +VMI type layout. It returns an empty layout when neither exists. It must not +default to contiguous during propagation. Passes that need to materialize or +inspect extra layouts should use the value assignment during apply instead of +overloading this singular accessor. + +The current assignment pass already has pieces that map directly onto this +utility: + +```text +LayoutSolver::setNaturalLayout + becomes request(value, layout). + +LayoutSolver::requestDataUse + becomes request(operand, layout). If the request conflicts with the source + value's `assignment.layout`, the propagator records a use-side conflict on that + value. + +LayoutSolver::getExplicitDataLayout + becomes a current-layout read from the value type or assigned equivalence. + +LayoutSolver::getDataLayout + currently defaults unknown to contiguous. The propagator must not do that + until finalization/apply. + +LayoutSolver::applyConsumerDrivenDataLayouts + is removed. Consumer operand requirements are represented as use-side + conflicts and do not become natural layouts. + +LayoutSolver::rewriteDataTypes and insertDataUseMaterializations + provide the first implementation material for VMILayoutPropagator::apply in + vmi-layout-assignment. +``` + +This means the first change can be staged: + +```text +1. Add VMILayoutPropagator with assignments, per-value def/use conflicts, + value-layout fact worklist, and strict same-operand conflict checking. + +2. Move layout relation helpers into it. Run mask granularity assignment/split + before layout propagation, and keep the layout propagator focused on layout + only. Do not move all materialization logic in the first step. + +3. Let vmi-layout-assignment call the propagator for the relations that need + order-independent propagation. + +4. Factor existing assignment finalization into VMILayoutPropagator::apply + once the assignment table is authoritative. +``` + +## Propagation + +The propagator's loop is deterministic: + +```text +while worklist is not empty: + pop changed value/layout fact + inspect the defining op transfer relation, if the value is an op result + inspect each user op transfer relation + operand/source value layout known -> derive result value layouts + result value layout known -> derive operand source value layouts when + inverse is legal +``` + +The propagator should work with generic relation records, not with hard-coded +patterns: + +```text +same-layout relation: + elementwise ops, bitcast, select-compatible values + +cast width relation: + source layout known -> derive result + result layout known -> derive source when inverse is legal + +channel split/merge relation: + source/input side layout known -> derive result side layout, and vice versa + when the inverse is unique + +control-flow relation: + CFG branch/yield/region layout consistency + call/return only when function/call boundary rewrite is enabled +``` + +Each relation should use the same helper that support checks use. For example, +a cast width relation should know only about source/result layout relations of +cast ops; it must not know about `truncf -> group_broadcast` as a pattern. + +### Transfer Relation Details + +The core asset is the static transfer relation: + +```text +known value + known layout + op-local rule + -> zero or more uniquely derived layouts for connected source/result values +``` + +This is different from assignment-specific policy such as "natural layout" or +"consumer request". Those policies decide where the first request comes from. +The transfer relation only answers whether a known layout on one connected +value uniquely determines layouts on other connected values. + +The first useful transfer relation set is: + +```text +same-layout relation: + For layout-transparent ops, if one operand/result gets a requested layout, + require the same layout on the connected values. Existing assignment + `unite` logic is a source of the supported op list. + +cast width relation: + Use VMILayoutSupport cast fact helpers. Dense and group-value casts should + both generate VMICastLayoutFact pairs. Group-value casts need a support + helper that derives: + narrow: result.LS = source.LS * width_ratio + widen: source.LS = result.LS * width_ratio + +channel split/merge relation: + Split and merge have fixed source/result layout equations once the channel + count is known. +``` + +The relation providers should call `VMILayoutSupport` for legality instead of +duplicating target checks. Missing support helpers should be added there, not +open-coded in the propagator. + +### Existing Assignment Transfer Assets By Op + +`vmi-layout-assignment` already contains several transfer relations. They +should be extracted by op or op family. Each op family may expose more than +one transfer rule, but every rule must still be local to that op. + +```text +VMIAddF/AddI/SubF/SubI/MulF/MulI/DivF/MinF/MaxF/AndI/OrI/XOrI: + Existing code: + constrainElementwiseBinary(...) + + Transfer rules: + lhs layout -> rhs layout and result layout, same layout + rhs layout -> lhs layout and result layout, same layout + result layout -> lhs layout and rhs layout, same layout + + Notes: + existing code has fallback handling for unsupported group_broadcast result + layouts; that fallback is assignment policy, not the transfer rule. + +VMIFma: + Existing code: + unite(lhs, rhs), unite(lhs, acc), unite(lhs, result) + + Transfer rules: + any of lhs/rhs/acc/result layout + -> all of lhs/rhs/acc/result use the same layout + +VMINegF/AbsF/AbsI/Sqrt/Exp/Ln/Relu/Not: + Existing code: + unite(source, result) + + Transfer rules: + source layout -> result layout, same layout + result layout -> source layout, same layout + +VMIFPToSI/VMISIToFP: + Existing code: + unite(source, result) + + Transfer rules: + source layout -> result layout, same layout + result layout -> source layout, same layout + +VMICmpF/VMICmpI: + Existing code: + unite(lhs, rhs) + + Transfer rules: + lhs layout -> rhs layout, same layout + rhs layout -> lhs layout, same layout + lhs/rhs layout -> mask result layout, same layout + mask result layout -> lhs/rhs layout, same layout when unique + +VMISelect: + Existing code: + unite(trueValue, falseValue), unite(trueValue, result) + + Transfer rules: + any of trueValue/falseValue/result layout + -> all of trueValue/falseValue/result use the same layout + selected value/result layout -> mask operand layout, same layout + mask operand layout -> selected value/result layout, same layout when + unique + +VMIBitcast: + Existing code: + unite(source, result) + + Transfer rules: + source layout -> result layout, same layout + result layout -> source layout, same layout + +Widen cast transfer (VMIExtF/ExtSI/ExtUI): + Existing code: + ExtF uses getPreferredCastLayoutFact(...) for dense widening. + ExtSI/ExtUI have a source group_slots slots=8 branch. + getPreferredCastLayoutFact(...) + + Transfer rules: + These three ops are the same widening transfer relation at the layout + level. Float/integer signedness affects element-type legality and lowering, + not the layout algebra. + dense source layout -> dense result layout when the cast fact is + unique and supported + dense result layout -> dense source layout by matching the same cast facts + group-value source layout -> group-value result layout for supported + widening + group-value result layout -> group-value source layout when the matching + cast fact is unique + + Extraction work: + ExtF does not currently have the group-value branch that ExtSI/ExtUI have. + This is a legacy assignment-framework artifact. Add one widening-cast fact + generator used by all three ops and let VMILayoutSupport decide which + element types are legal. + +VMITruncF: + Existing code: + VMILayoutSupport::getPreferredCastLayoutFact(...) + TruncF source group_slots slots=1 case + + Transfer rules: + dense source layout -> dense result layout when the cast fact is + unique and supported + dense result layout -> dense source layout when the matching cast fact is + unique and supported + group-value source layout -> group-value result layout for supported + narrowing + group-value result layout -> group-value source layout when the matching + cast fact is unique + + Extraction work: + current group-value branch supports only slots=1 and is source-driven. + Add dense trunc inverse generation to the shared cast fact helper. + +VMITruncI: + Existing code: + VMILayoutSupport::getPreferredCastLayoutFact(...) + TruncI source group_slots slots=1/8 branch + + Transfer rules: + dense source layout -> dense result layout when the cast fact is + unique and supported + dense result layout -> dense source layout when the matching cast fact is + unique and supported + group-value source layout -> group-value result layout for supported + narrowing + group-value result layout -> group-value source layout when the matching + cast fact is unique + + Extraction work: + current group-value branch is source-driven. The narrow4 slots=8 case + already computes lane_stride=4; move that equation into the shared cast + fact helper. + +VMIChannelSplit: + Existing code: + VMIChannelSplitOp case in addConstraints() + + Transfer rules: + source deinterleaved=channel_count -> each result contiguous + result contiguous on every result -> source deinterleaved=channel_count + +VMIChannelMerge: + Existing code: + VMIChannelMergeOp case in addConstraints() + + Transfer rules: + every input contiguous -> result deinterleaved=channel_count + result deinterleaved=channel_count -> every input contiguous + +control-flow ops: + Existing code: + addIfConstraints, addYieldConstraints, addExecuteRegionConstraints, + addIndexSwitchConstraints, addWhileConstraints, addForConstraints, + addBranchConstraints, addReturnConstraints, addCallConstraints + + Transfer rules: + any equivalent incoming/yield/result/call value layout + -> same layout on every value in that equivalence group + +mask ops: + Existing code: + uniteMask(...) + + Transfer rules: + same-layout mask propagation mirrors data same-layout propagation +``` + +Mask layout propagation is not a separate assignment flow. Mask values +participate in the same propagator/worklist as VMI data values. Data-producing +or data-consuming ops drive their mask operands/results through same-layout +relations: + +```text +data layout L -> mask layout L +mask layout L -> data layout L when the relation is unique +``` + +Mask granularity assignment is separate from layout propagation and runs before +layout assignment. Different granularities represent different mask values, so +the granularity pass should split or materialize mask values before the layout +propagator sees them. After that split, each mask SSA value has one fixed +granularity, and the layout propagator only assigns its layout. + +Granularity must not become a second independent request dimension in the +layout worklist, and different granularity requirements must not be represented +as layout conflicts on one mask value. + +Some assignment logic is not a transfer relation and should not be moved into +the propagator as if it were one: + +```text +producer-only layout choice: + group_reduce/group_load/group_slot_load/group_broadcast_load choosing an + initial result layout is assignment policy. It can seed assignments, but it + is not derived from another operand layout. + +store-only requirements: + store/group_store/masked_store have no result value to infer. They seed an + operand layout request for their value operand when the store form requires a + concrete input layout. They are not bidirectional transfer relations. + +group_broadcast pair support: + VMILayoutSupport can validate a source/result pair, but source layout alone + does not always uniquely choose a dense result layout. Treat it as a support + check and source requirement unless another value request makes the result + layout concrete. +``` + +### Relation Mechanics + +Each op does not own a persistent layout table. The only persistent table is +the propagator's global `assignments`. + +An op relation is implemented by a transfer object: + +```text +class VMILayoutTransfer: + propagate(op, changedValue, changedLayout, propagator) +``` + +`propagate` is an op-local fact propagation method. It does not rewrite IR. +It requests layouts for connected source and result values. It should be a +thin wrapper around the op family's pure relation evaluator: + +```text +derive(op, changedPort, changedLayout, assignmentView) + -> zero or more derived port/layout facts + +propagate(op, changedValue, changedLayout, propagator): + inspect op operands/results, attrs, element types, and VMILayoutSupport + facts = derive(op, changedPort, changedLayout, propagator.assignments) + for each result fact: + call request(resultValue, derivedLayout) + for each operand fact: + call request(operand, derivedLayout) +``` + +The evaluator is query-like in the ordinary sense: it is pure, it does not +mutate `assignments`, it does not enqueue work, and it does not rewrite IR. It +is not a public propagator `query` API because the propagator API that changes +state is still `request`. This keeps the mutation point explicit while letting +propagation and apply-time validation consume the same op relation. + +When the connected value is reached through an operand, `propagate` calls the +operand overload so a merge conflict can become a use-side conflict on the +source value. + +This is the mechanism that propagates layout information. When one connected +value receives a layout, the op relation may infer layouts for other connected +values: + +```text +same-layout op: + any operand source/result layout -> all connected source/result values get + the same layout + +cast op: + source value layout -> result value layout using width ratio + result value layout -> source value layout when the inverse relation is legal + +channel split/merge: + channel count plus one side layout -> the other side layout when the inverse + is unique +``` + +Not every relation is symmetric, and not every input layout determines every +other operand. If the relation cannot derive a unique supported layout, it +emits nothing. If a requested layout differs from the source value's +`assignment.layout`, the request is recorded as a conflict on that value. The +value overload records a def-side conflict; the operand overload records a +use-side conflict. If the same operand records two different conflict layouts, +strict propagation fails. + +Block arguments are also worklist values. They are not `OpResult`s and do not +have `getDefiningOp()`, but the propagator can still process them through a +boundary transfer: + +```text +process(value, layout): + if value is an OpResult: + process the defining op result port + if value is a BlockArgument: + process the block/function boundary port + process all uses of value +``` + +The block/function boundary transfer is separate from ordinary op-result +transfer: + +```text +ordinary op result: + defining op result <-> defining op operands/results + +block argument: + block argument <-> predecessor terminator successor operands + +function argument: + function argument <-> function signature / call operands, when the pass owns + signature or interprocedural rewrite +``` + +CFG block arguments require a same-transfer. A block argument and each +predecessor terminator successor operand represent the same semantic stream, so +layout requests must propagate in both directions: + +```text +block argument layout L -> request predecessor successor operand layout L +predecessor successor operand layout L -> request block argument layout L +``` + +If a CFG edge cannot satisfy the same layout as the block argument, the +terminator operand request becomes a use-side conflict on the predecessor +source value and apply materializes that edge operand before the terminator. + +Function signatures and call sites are a separate boundary. The first +implementation can leave function/call boundary transfer out if it only rewrites +inside one function and does not update function signatures or call sites. In +that mode, function arguments are boundary source values: they propagate to +their users, and primary boundary materialization is handled by apply. + +### Cast Layout Facts + +Width-changing casts should keep source/result layout information paired. Do +not encode individual concrete vector cases directly in `propagate`, such as +`64xf16 -> 64xf32`. Generate layout facts from the source/result element widths +and the known anchor layout: + +```text +VMICastLayoutFact: + sourceLayout + resultLayout +``` + +`propagate` uses those facts mechanically: + +```text +if changedValue is the cast source: + for each fact whose sourceLayout == changedLayout: + request result value layout = fact.resultLayout + +if changedValue is the cast result: + for each fact whose resultLayout == changedLayout: + call request(sourceOperand, fact.sourceLayout) +``` + +The support layer should provide a fact generator shaped like: + +```text +getCastLayoutFacts(sourceType, resultType, anchorSide, anchorLayout) + -> zero or more VMICastLayoutFact +``` + +The `anchorSide` is source or result. Together, `anchorSide` and +`anchorLayout` limit generation to the small set of facts that can match the +currently propagated layout. + +For width-changing dense casts, legality and preference are separate tables. +The legal table records only storage element widths and paired source/result +layouts. It is `AnyN`: vector element count is not part of dense cast +legality. `N` belongs only in the preferred table when one legal relation is +chosen over another for a concrete shape. + +Storage widths are represented by an `ElementBitsPattern`, parallel to +`LayoutPattern`. A row matches when the concrete source/result storage widths +match the row's bit patterns and the concrete source/result layouts match the +row's layout patterns. Do not introduce a separate width-class enum; write the +supported bit set directly in the row. + +Legal dense rows are written as paired relations, for example: + +```text +T16 -> T32: + contiguous -> deinterleaved=2 + lane_stride=2 -> contiguous + deinterleaved=2 -> deinterleaved=4 + +T32 -> T16: + deinterleaved=2 -> contiguous + contiguous -> lane_stride=2 + deinterleaved=4 -> deinterleaved=2 +``` + +`T8 <-> T32` follows the same idea with factor 4: + +```text +T8 -> T32: + contiguous -> deinterleaved=4 + lane_stride=2 -> deinterleaved=2 + lane_stride=4 -> contiguous + +T32 -> T8: + deinterleaved=4 -> contiguous + deinterleaved=2 -> lane_stride=2 + contiguous -> lane_stride=4 +``` + +Layout legality depends on storage width and physical layout, not on whether +the op is floating-point or integer. The op support layer still checks whether +a particular VMI op and element type are valid. + +Group-slot cast rows live in the same legal table. They are written as +parameterized layout patterns; `num_groups = G` is inherited from the anchor +layout used for the query. Packed narrowing records the selected sub-lane +stride on the result, and widening uses the inverse relation: + +```text +T8 -> T16: + group_slots(G, slots=1) -> group_slots(G, slots=1) + group_slots(G, slots=8, lane_stride=2) -> group_slots(G, slots=8) + +T16 -> T32: + group_slots(G, slots=1) -> group_slots(G, slots=1) + group_slots(G, slots=8, lane_stride=2) -> group_slots(G, slots=8) + +T8 -> T32: + group_slots(G, slots=1) -> group_slots(G, slots=1) + group_slots(G, slots=8, lane_stride=4) -> group_slots(G, slots=8) + +T16 -> T8: + group_slots(G, slots=1) -> group_slots(G, slots=1) + group_slots(G, slots=8) -> group_slots(G, slots=8, lane_stride=2) + +T32 -> T16: + group_slots(G, slots=1) -> group_slots(G, slots=1) + group_slots(G, slots=8) -> group_slots(G, slots=8, lane_stride=2) + +T32 -> T8: + group_slots(G, slots=1) -> group_slots(G, slots=1) + group_slots(G, slots=8) -> group_slots(G, slots=8, lane_stride=4) +``` + +There is no separate group-slot branch in the fact query. Dense and group-slot +casts are both produced by matching the same legal table against the source or +result anchor layout. + +The preferred table is a subset of the legal table. Exact `N` rows override +the default row for the same width pair: + +```text +T16 -> T32, N=64: + lane_stride=2 -> contiguous + +T16 -> T32, default: + contiguous -> deinterleaved=2 +``` + +When a preferred row is selected, the support helper must validate it through +the same legal fact query. A preferred row that is not legal is a bug in the +table, not a fallback opportunity. + +Examples of legal dense facts: + +```text +f32 -> f8, R=4: + source 256xf32 deinterleaved=4 + -> result 256xf8 contiguous + +f32 -> f16, R=2: + source 256xf32 deinterleaved=4 + -> result 256xf16 deinterleaved=2 +``` + +The inverse direction matches the same facts: + +```text +f8 result contiguous + -> f32 source deinterleaved=4 + +f16 result deinterleaved=2 + -> f32 source deinterleaved=4 +``` + +The propagator should accept an inverse only when fact generation returns a +single supported matching fact for the anchor side. If several facts could +satisfy the same side, the relation must not guess; it should emit nothing +unless another request makes the choice concrete. + +### Reduce Layout Facts + +Plain vector reductions currently have a fixed layout relation: + +```text +reduce_*: + source contiguous + init contiguous + mask same(source) + result contiguous +``` + +`group_reduce_*` has a richer relation and should be expressed as a table. The +table is parameterized by: + +```text +G = num_groups +group_size = source_element_count / G +VcgBlockElems = elements in one 32B VCG block +``` + +The query first classifies `group_size` against `VcgBlockElems`. The class is +an enum, not a numeric encoding: + +```text +QuarterBlock group_size == VcgBlockElems / 4 +HalfBlock group_size == VcgBlockElems / 2 +OneBlock group_size == VcgBlockElems +TwoBlock group_size == 2 * VcgBlockElems +FourBlock group_size == 4 * VcgBlockElems +FullPartMultiple group_size >= 8 * VcgBlockElems && + group_size % (8 * VcgBlockElems) == 0 +``` + +Other values are unsupported. Preferred group block rows are written with a +`gb` pattern. `gb(1, 4)` means one quarter of a 32B VCG block, and `gb(4)` +means four 32B VCG blocks. `group_reduce_*` is one consumer of this shared +block classification: + +```text +gb(1, 4): + source ls(4) + mask same(source) + result gs(8) + +gb(1, 2): + source ls(2) + mask same(source) + result gs(8) + +gb(1): + source c() + mask same(source) + result gs(8) + +gb(2): + source d(2, block_elems=1) + mask same(source) + result gs(8) + +gb(4): + source d(4, block_elems=1) + mask same(source) + result gs(8) + +gbFull(): + source c() + mask same(source) + result gs(1) +``` + +The preferred table is not the whole legal relation. A concrete fact query +also exists for post-assignment validation: + +```text +getGroupReduceLayoutFactForLayouts(source, mask, result, num_groups) +``` + +It matches the assigned source/mask/result layouts against legal rows for the +classified group block. Legal rows include additional source/mask alternatives +for the same semantic row, such as `block_elems=1` for the two-block and +four-block cases. Those alternatives are part of the layout relation, not +ad-hoc support relaxations. + +The concrete query is the single source of truth for layout-driven shape +legality. It may compute physical arity from the concrete VMI types, but only +to validate the selected relation row; arity is not an independent support +policy. For example, the two-block row implies two source/mask physical parts +per result part, while the four-block row implies four. + +Checks that are intrinsic to the VMI op contract belong in the op verifier, not +in layout support: + +```text +floating-point reassociation requirement +source/result element type equality +result element count equals num_groups +integer group reduction accumulator type +mask/data compatibility +num_groups divisibility +``` + +The concrete lowering plan is derived from the returned fact. The lowering may +still defensively check the OneToN-converted `sourceParts`, `maskParts`, and +`resultTypes`, but it must not re-encode a separate group-size or layout support +table. + +### Memory Layout Facts + +Load/store layout support follows the same split: + +```text +VMILayoutSupport: + layout relation facts only + +VMIToVPTO: + target/memory/stride/lowering preconditions + defensive checks on the actual OneToN value/result ranges +``` + +Do not encode VPTO instruction names in layout support. Names such as +`Vstsx2`, `Vsldb`, `PK4_B32`, or `Slots1PointVsts` are lowering choices, not +layout facts. + +Dense `vmi.load` / `vmi.store` facts are keyed by element bits and assigned +value layout: + +```text +load: + bits(8,16,32), contiguous + bits(8,16,32), lane_stride=2 + bits(8), lane_stride=4 + bits(8,16,32), deinterleaved=2/4 + +store: + bits(8,16,32), contiguous + bits(8,16,32), lane_stride=2 + bits(8), lane_stride=4 + bits(8,16,32), deinterleaved=2/4, block_elems=1 +``` + +The fact query records only the dense memory layout pattern and element-bit +pattern. It may reject impossible layout/width combinations such as +`lane_stride=4` for non-b8 elements. It must not decide whether a particular +memref is UB-backed, whether an offset is aligned, or whether the current +lowering will materialize a fallback path. + +Two-way memory ops are separate VMI semantics: + +```text +deinterleave_load: + low contiguous + high contiguous + +interleave_store: + low contiguous + high contiguous +``` + +They are not represented as dense load/store `deinterleaved` facts. Their +current VPTO lowering can still require `!pto.ptr`, direct UB memory, full +chunks, and `vldsx2/vstsx2` element support in `VMIToVPTO`. + +Group memory facts are also table relations: + +```text +group_load: + bits(32), gb(2) -> result d(2, block_elems=8) + bits(32), gb(4) -> result d(4, block_elems=8) + +group_slot_load: + result group_slots(num_groups=G, slots=1) + result group_slots(num_groups=G, slots=8) + result group_slots(num_groups=G, slots=8, lane_stride=2) + result group_slots(num_groups=G, slots=8, lane_stride=4) + +group_store: + value group_slots(num_groups=G, slots=1) + value group_slots(num_groups=G, slots=8) + value group_slots(num_groups=G, slots=8, lane_stride=2) + value group_slots(num_groups=G, slots=8, lane_stride=4) + +group_broadcast_load: + bits(8,16,32), memContiguous() -> source group_slots(G, slots=8) + bits(8,16,32), memBlockAligned() -> source group_slots(G, slots=1) + +masked_store compact lane-stride: + bits(8), value ls(2), mask same(value) -> packed predicate b16 + bits(16), value ls(2), mask same(value) -> packed predicate b32 + bits(8), value ls(4), mask same(value) -> packed predicate b32 +``` + +The `num_groups` equality is part of the layout relation. Current lowering +requirements such as `group_load` row stride being a constant positive multiple +of 8, `group_slot_load slots=8` using unit source-group stride, or +`group_store slots=8` using unit row stride remain in `VMIToVPTO`. +`group_broadcast_load` support is expressed as the equivalent +`group_slot_load + group_broadcast` relation. `VMIToVPTO` may still choose an +E2B VPTO lowering when the matched layout/shape is the E2B-friendly case, but +E2B is not a separate layout support fact. `masked_store` still materializes +the final predicate in lowering; the support table only records which +layout/mask-shape relations are legal. + +The lowering may check actual converted arity before indexing: + +```text +if resultTypes.size() != expected arity: + notifyMatchFailure(...) +``` + +This is a crash guard for the concrete rewrite. It must not become another +support policy in `VMILayoutSupport`. + +Group-value casts use a different static relation: + +```text +narrow by R: + result.LS = source.LS * R + +widen by R: + source.LS = result.LS * R +``` + +These fact generators belong in shared support helpers so assignment, +rematerialization, validation, and lowering agree on the same relation. + +## Apply + +Apply derives concrete IR actions from `assignments` and the current IR. +These actions do not need to be stored in separate propagator tables before +apply: + +Apply must not run a second propagation-style relation query over results or +operands. By the time apply starts, `run()` has already reached a fixed point: + +```text +assignment.layout records the layout to make available for each value +assignment.conflicts records def-side and use-side alternate layouts +every def/user relation has already propagated its required layouts +every use that cannot consume assignment.layout has already become a conflict +``` + +Therefore apply does not ask the op relation again. It writes the layouts +already recorded in `assignments` into IR: + +```text +apply must not create new layout requests +apply must not discover new layout conflicts +apply only consumes assignment.layout and assignment.conflicts +``` + +```text +sourceValue: + the original SSA value before apply + +currentLayout: + the layout carried by sourceValue's current VMI type + +assignedLayout: + assignment.layout for sourceValue + +assignedValue: + the SSA value that carries assignedLayout after apply + if sourceValue is type-rewriteable, assignedValue is sourceValue after its + VMI type is rewritten to assignedLayout + otherwise assignedValue is ensure_layout sourceValue : + currentLayout -> assignedLayout when currentLayout != assignedLayout + +def-side conflict: + materialize an extra SSA value from assignedValue to conflict.layout near + the def or boundary + +use-side conflict: + materialize an extra SSA value from assignedValue to conflict.layout before + that use +``` + +Primary type rewrite is not a producer-specific optimization. It is the normal +way assignment becomes explicit in IR. It does not choose a layout and does not +ask whether the producer relation supports the layout; propagation already did +that. It only updates the VMI type to `assignment.layout`. + +Primary type rewrite is available for op results whose defining op is inside +the rewrite scope and whose result type can be changed without rewriting an +external ABI boundary. Multi-result ops should be rewritten as one op update +using the final assignments for all assigned results. Block arguments, +function arguments, values defined outside the rewrite scope, and ABI boundary +values use `ensure_layout` materialization instead. + +In this document, `type-rewriteable` means exactly: + +```text +the value is an OpResult +the defining op is inside the current rewrite scope +changing the result VMI type does not rewrite an external ABI boundary +``` + +It does not mean the defining op was queried again for a preferred layout. + +Primary materialization is not conflict-driven. A value can have no conflicts +and still need `ensure_layout` when its current IR type cannot be rewritten to +`assignment.layout`: + +```text +function argument current layout A +only in-scope use requests layout B + +assignment.layout = B +assignment.conflicts = empty + +if the function signature is not rewritten: + arg_B = ensure_layout arg : A -> B + use(arg_B) +``` + +Conflicts only describe extra layouts besides `assignment.layout`. They do not +replace the primary action that makes `assignment.layout` available. + +Producer-specific improvements, such as folding a fallback +`load -> ensure_layout` into a load with the requested layout or rematerializing +a cast across an `ensure_layout`, are not part of apply. They should run as +ordinary layout-fold/rematerialization over explicit helper IR. + +```text +1. Make assignment.layout available. + Let sourceValue be the original SSA value, currentLayout be the layout on + sourceValue's current VMI type, and assignedLayout be assignment.layout. + If sourceValue is type-rewriteable, rewrite its VMI type to assignedLayout + and use sourceValue as assignedValue. Otherwise, if + currentLayout == assignedLayout, assignedValue is sourceValue. Otherwise + insert ensure_layout sourceValue : currentLayout -> assignedLayout and use + its result as assignedValue. If that ensure_layout is not supported by + VMILayoutSupport, apply fails. + +2. Materialize def-side conflicts. + For each def-side conflict layout, if assignedValue already has that layout, + reuse assignedValue. Otherwise insert ensure_layout near the value's + definition or boundary. If that ensure_layout is not supported by + VMILayoutSupport, apply fails. The ensure_layout result is the materialized + SSA value for that layout; no persistent container is needed to represent it. + +3. Rewrite non-conflicting uses. + Uses in the rewrite scope are redirected to assignedValue unless a use-side + conflict records a different layout for that operand. Do not implement this + as an unconditional replace-all-uses. Iterate the original uses and skip + operands recorded in use-side conflicts. Do not query the user op relation + here; non-conflicting uses were already accepted during propagation. + +4. Materialize use-side conflicts. + For each use-side conflict, insert ensure_layout from assignedValue to + conflict.layout before conflict.operand and replace only that operand. + If that ensure_layout is not supported by VMILayoutSupport, apply fails. Do + not special-case the old producer layout here. Redundant chains such as + l1 -> l2 -> l1 are folded by a separate layout-fold/rematerialization pass. +``` + +Concrete insertion points: + +```text +op result: + normally rewrite the result VMI type in place. If the result is outside the + rewrite scope or crosses an ABI boundary, insert the fallback primary + ensure_layout immediately after the defining op. + +block argument: + insert the fallback primary ensure_layout at the first legal insertion point + of the owning block. + +function argument: + insert the fallback primary ensure_layout at the first legal insertion point + of the entry block. + +def-side conflict: + insert ensure_layout after assignedValue is available, using the same + def/boundary placement as the primary materialization. + +use-side conflict: + insert ensure_layout immediately before conflict.operand.getOwner() and + replace only conflict.operand. +``` + +Apply should keep a local materialization map keyed by `(Value, Layout, +placement)` so the same required layout at the same placement is not emitted +twice. Different use-side conflicts may still materialize separately when a +single def-side value would not dominate all uses. + +The source for conflict materialization is always `assignedValue`, not the +operand's old value: + +```text +def-side conflict layout C: + c = ensure_layout assignedValue : assignment.layout -> C + +use-side conflict operand op.i requiring layout C: + c = ensure_layout assignedValue : assignment.layout -> C + op.i = c +``` + +If `C == assignment.layout`, the conflict is an identity and no +`ensure_layout` is inserted. + +A def-side conflict by itself does not replace arbitrary uses of the original +SSA value. It only materializes another layout view at the definition or +boundary because a value-level request asked for that layout. Use-side +conflicts are still materialized locally from assignedValue. + +For example, a block argument with current layout `A` and requested +`assignment.layout` `B` keeps its original type unless the rewrite scope allows +changing the boundary. Apply inserts `ensure_layout A -> B` near the boundary +and rewires in-scope uses to the materialized value, except for operands that +have explicit use-side conflicts. + +For a normal defining op inside the rewrite scope, apply rewrites the result +type in place: + +```text +before: + a = producer() : layout A + use(a) + +after assignment.layout = B: + a = producer() : layout B + use(a) +``` + +For a value that cannot be rewritten in place, apply materializes the assigned +layout with `ensure_layout`: + +```text +before: + a0 = boundary_value : layout A + use(a0) + +after assignment.layout = B: + a1 = boundary_value : layout A + a = ensure_layout a1 : A -> B + use(a) +``` + +If one use still requires layout `A`, apply emits the local materialization +from the assigned value: + +```text +after assignment.layout = B, with one use-side conflict requiring A: + a = producer() : layout B + c = ensure_layout a : B -> A + use(c) +``` + +For the non-rewrite fallback, the same conflict may produce an `A -> B -> A` +chain. That chain is not a special case in apply. A separate +layout-fold/rematerialization pass may fold it back to the original boundary +value when legal. + +For `vmi-layout-assignment`, the existing apply path is already usable: + +```text +rewriteDataTypes: + Sets VMI value types to `assignment.layout`. + +insertDataUseMaterializations: + Inserts pto.vmi.ensure_layout before operand uses whose requested layout does + not match the source value type. + +rewriteMaskTypes / insertMaskUseMaterializations: + Reused after first-phase mask granularity assignment/split and mask layout + propagation. +``` + +For post-assignment optimization passes, apply must be more conservative: + +```text +rewritable value: + The value is an op result inside the rewrite scope and changing the VMI type + does not cross an external ABI boundary. + +non-rewritable value: + Function/block arguments, external boundaries, or values outside the pass + rewrite scope keep their original type. If they have a requested + assignment.layout different from the current explicit layout, apply inserts + a def-side ensure_layout inside the rewrite scope and rewires in-scope uses + to the materialized value. +``` + +Whether a value is rewritable is derived from the IR and the caller's rewrite +scope. It is not stored in `assignments`. + +## Conflicts + +The propagator distinguishes representable conflicts from hard conflicts: + +```text +def-side conflict: + A source value's assignment.layout differs from another value-level request. + Record VMILayoutConflict{def, layout}. Apply will materialize it as an + ensure_layout fork near the value definition or boundary if it remains + different. + +use-side conflict: + A source value's assignment.layout differs from one operand's required layout. + Record VMILayoutConflict{operand, layout}. Apply will materialize it as an + ensure_layout fork if it remains different. + +hard operand conflict: + The same operand is requested as two different layouts. Do not create two + forks for one operand. The first implementation fails the current + propagation request. +``` + +The initial conflict policy should be strict inside the propagator: + +```text +same value requested as two different layouts: + keep the first layout as assignment.layout, record subsequent layouts as + def-side conflicts. Propagate only the primary assignment.layout fact. + +same operand requested as two different layouts: + fail the propagation request with a diagnostic at the requesting operation. +``` + +Value-level and operand-level layout differences are not hard conflicts when +they can be represented by an unambiguous fork. They are recorded in the +source value's `conflicts` list and are materialized by +VMILayoutPropagator::apply. + +## Assignment Shape + +`vmi-layout-assignment` can use the propagator as: + +```text +collect op layout constraints and relations +request natural layouts for producers that choose concrete layouts +request layouts required by consumers +propagate the value-layout table through op relations +apply ensure_layout / ensure_mask_layout materialization +validate assigned VMI IR +``` + +Later layout optimization passes should follow the same model: + +```text +request a new layout for one or more anchor values +propagate the value-layout table through registered op relations +materialize mismatches at values or uses +run layout-fold/rematerialization to remove redundant helper IR +``` + +Manual layout clearing is unsafe because it loses boundary contracts and can +turn an already validated assigned IR back into an ambiguous pre-assignment IR. + +## First Implementation Boundary + +The first implementation is deliberately limited: + +```text +included: + data value layout propagation + value-level requests + def-side and operand-level conflicts stored inside each value assignment + strict same-operand conflict diagnostics + same-layout data op transfer + mask granularity assignment/split before layout propagation + mask layout propagation + CFG/block-argument same-transfer for control-flow values + cast width relations needed by group-value cast/broadcast + reuse of existing assignment type rewrite and data ensure_layout insertion + +not included: + function signature and call-site interprocedural rewrite + a separate operand-request table outside value assignments + cost-based layout choice + best-effort request dropping + global replacement of fold/rematerialize/sink passes +``` + +This boundary makes the design implementable without forcing all existing VMI +layout passes to move at once. + +Acceptance requirement: + +```text +Existing VMI lit and simulator regression outcomes must not regress after the +refactor. Any test that passed before the propagator refactor must still pass +after it. If a test's expected IR shape changes because the new propagation is +more canonical, update the expectation only with an explicit before/after +reason in the change description. +``` + +## Anti-Specialization Rules + +The propagator must not grow optimization-pattern-specific boundary checks. +These forms are not acceptable: + +```text +if producer is group_reduce and op is truncf and user is group_broadcast: + choose layout X + +if value is function argument and consumer is some specific op: + insert special materialization Y +``` + +Boundary handling should be expressed through generic materialization support, +not producer-specific pattern checks inside the propagator: + +```text +canMaterializeLayout(sourceType, resultType): + delegate to VMILayoutSupport::getDataLayoutMaterializationSupport +``` + +Op-specific logic is still necessary, but it must be local to one op and one +role. Do not combine several ops into one pattern: + +```text +cast transfer: + source/result width relation only + +channel split/merge transfer: + channel-count layout equation only + +group_broadcast support/request: + source operand group-value requirement and source/result support only + +group_reduce seed: + initial group-value result layout choice only + +store request/support: + required operand layout and store support only +``` + +With this structure, `group_reduce -> truncf -> group_broadcast` works because +independent op-local rules compose through `assignments`, not because the +propagator recognizes that whole chain. + +## Example: Group-Value Cast + +For a group-value cast relation: + +```text +if the source value is requested/propagated as group-value: + derive result group-value layout + request the result value layout + +if the result value is requested/propagated as group-value: + derive inverse source value layout + request the source value layout through the source operand overload +``` + +If the derived source layout conflicts with the source value's existing +`assignment.layout`, the operand overload records a use-side conflict in the +source value's `assignment.conflicts`. Apply materializes it with +`ensure_layout` if the layouts still differ. That conflict is not propagated +back through the producer. Propagating an alternate layout through the producer +would mean rematerializing or cloning that producer for the alternate layout, +which is a separate optimization and must create a real forked value. + +## Pre-Refactor Regression Baseline + +Baseline captured on 2026-07-02 before introducing the shared VMI layout +propagator implementation. + +```text +git HEAD: + 4a2a100f + +working tree notes: + unrelated pre-existing changes were present in 3rdparty/PTO-Gym and + docs/designs/vmi-layout-lowering-cases.md. + untracked local investigation files were also present and are not part of + this baseline. +``` + +VMI lit baseline: + +```bash +export PATH="/home/mouliangyu/projects/github.com/vpto-dev/llvm-project/build-shared/bin:$PATH" +python3 /home/mouliangyu/projects/github.com/vpto-dev/llvm-project/llvm/utils/lit/lit.py \ + -v -j16 build/test/lit/vmi +``` + +Result: + +```text +Total Discovered Tests: 379 +Passed: 379 +Failed: 0 +``` + +VMI simulator baseline: + +```bash +WORK_SPACE=/tmp/ptoas-vmi-baseline-latest/sim \ +CASE_PREFIX='vmi/' \ +JOBS=16 \ +test/vpto/scripts/run_host_vpto_validation_parallel.sh +``` + +Result: + +```text +Total cases: 85 +PASS: 81 +FAIL: 4 +``` + +Existing simulator failures in this baseline: + +```text +vmi/group-reduce-s16-truncf-broadcast-store +vmi/group-reduce-s64-slot-add-store +vmi/group-reduce-s64-broadcast-reduce-store +vmi/group-reduce-s64-truncf-store +``` + +The refactor acceptance point is equality-or-better against this baseline: +all 379 VMI lit tests must keep passing, and the VMI simulator run must not add +new failing cases or turn any of the 81 passing cases into failures. diff --git a/docs/designs/vmi-mxfp8-32x32-expected-lowering.md b/docs/designs/vmi-mxfp8-32x32-expected-lowering.md new file mode 100644 index 0000000000..5130da3ec6 --- /dev/null +++ b/docs/designs/vmi-mxfp8-32x32-expected-lowering.md @@ -0,0 +1,236 @@ +# VMI MXFP8 32x32 Expected VPTO Lowering + +本文记录 `test/vpto/cases/vmi/kernels/tquant-mxfp8-32x32-nd/kernel.pto` +的预期 VPTO lower 结果。输入 VMI case 在 `vecscope` 内按 8 行一组循环， +每次处理一个 `256xf32` tile，也就是 8 行 x 32 列。 + +这里写的是设计目标，不是当前 `--emit-vpto` 的实际输出。重点是把 E8M0 +scale 的内存效果写明确：每个 8x32 chunk 产生 8 个 scale byte。lowering +按 CCE 风格先写到 32B 对齐的 padded UB slot，再通过 UB->GM copy 的 +`src_stride=32B, dst_stride=8B` 消除 UB padding，使 GM 端仍然连续。 + +## Complete Expected PTO File + +```mlir +module attributes {pto.backend = "vpto", pto.target_arch = "a5"} { + module attributes {pto.backend = "vpto", pto.kernel_kind = #pto.kernel_kind, pto.target_arch = "a5"} { + func.func @vmi_tquant_mxfp8_32x32_nd_kernel(%src_gm: !pto.ptr, + %out_fp8_gm: !pto.ptr, + %out_e8m0_gm: !pto.ptr) attributes {pto.kernel} { + %false = arith.constant false + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %c3 = arith.constant 3 : index + %c4 = arith.constant 4 : index + %c5 = arith.constant 5 : index + %c6 = arith.constant 6 : index + %c7 = arith.constant 7 : index + %c8 = arith.constant 8 : index + %c16 = arith.constant 16 : index + %c24 = arith.constant 24 : index + %c32 = arith.constant 32 : index + %c64 = arith.constant 64 : index + %c128 = arith.constant 128 : index + %c192 = arith.constant 192 : index + + %c0_i32 = arith.constant 0 : i32 + %c1_i32 = arith.constant 1 : i32 + %c2_i32 = arith.constant 2 : i32 + %c3_i32 = arith.constant 3 : i32 + %c4_i32 = arith.constant 4 : i32 + %c5_i32 = arith.constant 5 : i32 + %c6_i32 = arith.constant 6 : i32 + %c7_i32 = arith.constant 7 : i32 + %c8_i32 = arith.constant 8 : i32 + %c23_i32 = arith.constant 23 : i32 + %c24_i32 = arith.constant 24 : i32 + %c40_i32 = arith.constant 40 : i32 + %c48_i32 = arith.constant 48 : i32 + %c56_i32 = arith.constant 56 : i32 + %c254_i32 = arith.constant 254 : i32 + %c2139095040_i32 = arith.constant 2139095040 : i32 + + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c4_i64 = arith.constant 4 : i64 + %c8_i64 = arith.constant 8 : i64 + %c32_i64 = arith.constant 32 : i64 + %c256_i64 = arith.constant 256 : i64 + %c1024_i64 = arith.constant 1024 : i64 + %c4096_i64 = arith.constant 4096 : i64 + %c8192_i64 = arith.constant 8192 : i64 + %c12288_i64 = arith.constant 12288 : i64 + + %ub_src = pto.castptr %c0_i64 : i64 -> !pto.ptr + %ub_out_fp8_u8 = pto.castptr %c8192_i64 : i64 -> !pto.ptr + %ub_out_fp8_f8 = pto.castptr %c8192_i64 : i64 -> !pto.ptr + %ub_out_e8m0 = pto.castptr %c12288_i64 : i64 -> !pto.ptr + + pto.copy_gm_to_ubuf %src_gm, %ub_src, %c0_i64, %c1_i64, %c4096_i64, %c0_i64, %c0_i64, %false, %c0_i64, %c4096_i64, %c4096_i64 + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i1, i64, i64, i64 + pto.copy_gm_to_ubuf %out_fp8_gm, %ub_out_fp8_u8, %c0_i64, %c1_i64, %c1024_i64, %c0_i64, %c0_i64, %false, %c0_i64, %c1024_i64, %c1024_i64 + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i1, i64, i64, i64 + pto.set_flag[, , ] + pto.wait_flag[, , ] + + pto.vecscope { + scf.for %row = %c0 to %c32 step %c8 { + %elem_off = arith.muli %row, %c32 : index + %elem_off_64 = arith.addi %elem_off, %c64 : index + %elem_off_128 = arith.addi %elem_off, %c128 : index + %elem_off_192 = arith.addi %elem_off, %c192 : index + + %x0 = pto.vlds %ub_src[%elem_off] : !pto.ptr -> !pto.vreg<64xf32> + %x1 = pto.vlds %ub_src[%elem_off_64] : !pto.ptr -> !pto.vreg<64xf32> + %x2 = pto.vlds %ub_src[%elem_off_128] : !pto.ptr -> !pto.vreg<64xf32> + %x3 = pto.vlds %ub_src[%elem_off_192] : !pto.ptr -> !pto.vreg<64xf32> + + %d0, %d1 = pto.vdintlv %x0, %x1 : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %d2, %d3 = pto.vdintlv %x2, %x3 : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %d4, %d5 = pto.vdintlv %d0, %d2 : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %d6, %d7 = pto.vdintlv %d1, %d3 : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + + %all_b32 = pto.pset_b32 "PAT_ALL" : !pto.mask + %slot8_b32 = pto.pge_b32 "PAT_VL8" : !pto.mask + %vl8_b32 = pto.pset_b32 "PAT_VL8" : !pto.mask + %vl16_b32 = pto.pset_b32 "PAT_VL16" : !pto.mask + %vl24_b32, %unused24 = pto.plt_b32 %c24_i32 : i32 -> !pto.mask, i32 + %vl32_b32 = pto.pset_b32 "PAT_VL32" : !pto.mask + %vl40_b32, %unused40 = pto.plt_b32 %c40_i32 : i32 -> !pto.mask, i32 + %vl48_b32, %unused48 = pto.plt_b32 %c48_i32 : i32 -> !pto.mask, i32 + %vl56_b32, %unused56 = pto.plt_b32 %c56_i32 : i32 -> !pto.mask, i32 + + %abs0 = pto.vabs %d4, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %abs1 = pto.vabs %d6, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %abs2 = pto.vabs %d5, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %abs3 = pto.vabs %d7, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + + %g0 = pto.vcgmax %abs0, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %g1 = pto.vcgmax %abs1, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %g2 = pto.vcgmax %abs2, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %g3 = pto.vcgmax %abs3, %all_b32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %g01 = pto.vmax %g0, %g1, %slot8_b32 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %g23 = pto.vmax %g2, %g3, %slot8_b32 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %amax = pto.vmax %g01, %g23, %slot8_b32 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + + %amax_i32 = pto.vbitcast %amax : !pto.vreg<64xf32> -> !pto.vreg<64xi32> + %exp_mask = pto.vdup %c2139095040_i32, %all_b32 : i32, !pto.mask -> !pto.vreg<64xi32> + %shift = pto.vdup %c23_i32, %all_b32 : i32, !pto.mask -> !pto.vreg<64xi32> + %emax = pto.vdup %c8_i32, %all_b32 : i32, !pto.mask -> !pto.vreg<64xi32> + %scale_exp_bias = pto.vdup %c254_i32, %all_b32 : i32, !pto.mask -> !pto.vreg<64xi32> + %exp_bits = pto.vand %amax_i32, %exp_mask, %all_b32 : !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.mask -> !pto.vreg<64xi32> + %exp = pto.vshr %exp_bits, %shift, %all_b32 : !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.mask -> !pto.vreg<64xi32> + %e8m0_payload_i32 = pto.vsub %exp, %emax, %all_b32 : !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.mask -> !pto.vreg<64xi32> + + %idx0 = pto.vdup %c0_i32, %all_b32 : i32, !pto.mask -> !pto.vreg<64xi32> + %idx1 = pto.vdup %c1_i32, %all_b32 : i32, !pto.mask -> !pto.vreg<64xi32> + %idx2 = pto.vdup %c2_i32, %all_b32 : i32, !pto.mask -> !pto.vreg<64xi32> + %idx3 = pto.vdup %c3_i32, %all_b32 : i32, !pto.mask -> !pto.vreg<64xi32> + %idx4 = pto.vdup %c4_i32, %all_b32 : i32, !pto.mask -> !pto.vreg<64xi32> + %idx5 = pto.vdup %c5_i32, %all_b32 : i32, !pto.mask -> !pto.vreg<64xi32> + %idx6 = pto.vdup %c6_i32, %all_b32 : i32, !pto.mask -> !pto.vreg<64xi32> + %idx7 = pto.vdup %c7_i32, %all_b32 : i32, !pto.mask -> !pto.vreg<64xi32> + + %not_vl8 = pto.pnot %vl8_b32, %all_b32 : !pto.mask, !pto.mask -> !pto.mask + %range_8_15 = pto.pand %vl16_b32, %not_vl8, %all_b32 : !pto.mask, !pto.mask, !pto.mask -> !pto.mask + %broadcast_idx_1 = pto.vsel %idx1, %idx0, %range_8_15 : !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.mask -> !pto.vreg<64xi32> + %not_vl16 = pto.pnot %vl16_b32, %all_b32 : !pto.mask, !pto.mask -> !pto.mask + %range_16_23 = pto.pand %vl24_b32, %not_vl16, %all_b32 : !pto.mask, !pto.mask, !pto.mask -> !pto.mask + %broadcast_idx_2 = pto.vsel %idx2, %broadcast_idx_1, %range_16_23 : !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.mask -> !pto.vreg<64xi32> + %not_vl24 = pto.pnot %vl24_b32, %all_b32 : !pto.mask, !pto.mask -> !pto.mask + %range_24_31 = pto.pand %vl32_b32, %not_vl24, %all_b32 : !pto.mask, !pto.mask, !pto.mask -> !pto.mask + %broadcast_idx_3 = pto.vsel %idx3, %broadcast_idx_2, %range_24_31 : !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.mask -> !pto.vreg<64xi32> + %not_vl32 = pto.pnot %vl32_b32, %all_b32 : !pto.mask, !pto.mask -> !pto.mask + %range_32_39 = pto.pand %vl40_b32, %not_vl32, %all_b32 : !pto.mask, !pto.mask, !pto.mask -> !pto.mask + %broadcast_idx_4 = pto.vsel %idx4, %broadcast_idx_3, %range_32_39 : !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.mask -> !pto.vreg<64xi32> + %not_vl40 = pto.pnot %vl40_b32, %all_b32 : !pto.mask, !pto.mask -> !pto.mask + %range_40_47 = pto.pand %vl48_b32, %not_vl40, %all_b32 : !pto.mask, !pto.mask, !pto.mask -> !pto.mask + %broadcast_idx_5 = pto.vsel %idx5, %broadcast_idx_4, %range_40_47 : !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.mask -> !pto.vreg<64xi32> + %not_vl48 = pto.pnot %vl48_b32, %all_b32 : !pto.mask, !pto.mask -> !pto.mask + %range_48_55 = pto.pand %vl56_b32, %not_vl48, %all_b32 : !pto.mask, !pto.mask, !pto.mask -> !pto.mask + %broadcast_idx_6 = pto.vsel %idx6, %broadcast_idx_5, %range_48_55 : !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.mask -> !pto.vreg<64xi32> + %not_vl56 = pto.pnot %vl56_b32, %all_b32 : !pto.mask, !pto.mask -> !pto.mask + %range_56_63 = pto.pand %all_b32, %not_vl56, %all_b32 : !pto.mask, !pto.mask, !pto.mask -> !pto.mask + %broadcast_idx = pto.vsel %idx7, %broadcast_idx_6, %range_56_63 : !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.mask -> !pto.vreg<64xi32> + + %scale_u16 = pto.vpack %e8m0_payload_i32, "LOWER" : !pto.vreg<64xi32> -> !pto.vreg<128xui16> + %scale_u8 = pto.vpack %scale_u16, "LOWER" : !pto.vreg<128xui16> -> !pto.vreg<256xui8> + %scale_slot = arith.divui %row, %c8 : index + %scale_ub_off = arith.muli %scale_slot, %c32 : index + %scale8_b8 = pto.pge_b8 "PAT_VL8" : !pto.mask + pto.vsts %scale_u8, %ub_out_e8m0[%scale_ub_off], %scale8_b8 {dist = "NORM_B8"} : !pto.vreg<256xui8>, !pto.ptr, !pto.mask + + %scale_exp = pto.vsub %scale_exp_bias, %e8m0_payload_i32, %all_b32 : !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.mask -> !pto.vreg<64xi32> + %scale_bits = pto.vshl %scale_exp, %shift, %all_b32 : !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.mask -> !pto.vreg<64xi32> + %scale_f32 = pto.vbitcast %scale_bits : !pto.vreg<64xi32> -> !pto.vreg<64xf32> + %scale_vec = pto.vselr %scale_f32, %broadcast_idx : !pto.vreg<64xf32>, !pto.vreg<64xi32> -> !pto.vreg<64xf32> + + %m0 = pto.vmul %d4, %scale_vec, %all_b32 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %m1 = pto.vmul %d6, %scale_vec, %all_b32 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %m2 = pto.vmul %d5, %scale_vec, %all_b32 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %m3 = pto.vmul %d7, %scale_vec, %all_b32 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + + %i0, %i1 = pto.vintlv %m0, %m2 : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %i2, %i3 = pto.vintlv %m1, %m3 : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %i4, %i5 = pto.vintlv %i0, %i2 : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %i6, %i7 = pto.vintlv %i1, %i3 : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %r0, %r1 = pto.vdintlv %i4, %i5 : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %r2, %r3 = pto.vdintlv %i6, %i7 : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %r4, %r5 = pto.vdintlv %r0, %r2 : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %r6, %r7 = pto.vdintlv %r1, %r3 : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + + %all_b8 = pto.pset_b8 "PAT_ALL" : !pto.mask + %q0 = pto.vcvt %r4, %all_b32 {part = "P0", rnd = "R", sat = "SAT"} : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<256xf8E4M3FN> + %q1 = pto.vcvt %r6, %all_b32 {part = "P1", rnd = "R", sat = "SAT"} : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<256xf8E4M3FN> + %q2 = pto.vcvt %r5, %all_b32 {part = "P2", rnd = "R", sat = "SAT"} : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<256xf8E4M3FN> + %q3 = pto.vcvt %r7, %all_b32 {part = "P3", rnd = "R", sat = "SAT"} : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<256xf8E4M3FN> + %q01 = pto.vor %q0, %q1, %all_b8 : !pto.vreg<256xf8E4M3FN>, !pto.vreg<256xf8E4M3FN>, !pto.mask -> !pto.vreg<256xf8E4M3FN> + %q012 = pto.vor %q01, %q2, %all_b8 : !pto.vreg<256xf8E4M3FN>, !pto.vreg<256xf8E4M3FN>, !pto.mask -> !pto.vreg<256xf8E4M3FN> + %q = pto.vor %q012, %q3, %all_b8 : !pto.vreg<256xf8E4M3FN>, !pto.vreg<256xf8E4M3FN>, !pto.mask -> !pto.vreg<256xf8E4M3FN> + pto.vsts %q, %ub_out_fp8_f8[%elem_off], %all_b8 : !pto.vreg<256xf8E4M3FN>, !pto.ptr, !pto.mask + } + } + + pto.set_flag[, , ] + pto.wait_flag[, , ] + pto.copy_ubuf_to_gm %ub_out_fp8_u8, %out_fp8_gm, %c0_i64, %c1_i64, %c1024_i64, %c0_i64, %c1024_i64, %c1024_i64 + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 + pto.copy_ubuf_to_gm %ub_out_e8m0, %out_e8m0_gm, %c0_i64, %c4_i64, %c8_i64, %c0_i64, %c8_i64, %c32_i64 + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 + pto.barrier + return + } + } +} +``` + +## Scale Store Contract + +上面的 lower 对每次循环执行一条 `NORM_B8` store，写到 32B 对齐的 UB +slot： + +```text +row = 0 -> UB[0..7], UB[8..31] padding +row = 8 -> UB[32..39], UB[40..63] padding +row = 16 -> UB[64..71], UB[72..95] padding +row = 24 -> UB[96..103], UB[104..127] padding +``` + +最终 copy-out 只搬每个 slot 的前 8B： + +```text +copy len = 8B +repeat = 4 +source stride = 32B +destination stride = 8B +``` + +因此 GM 端效果仍然是连续 scale 输出： + +```text +GM[0..7] <- UB[0..7] +GM[8..15] <- UB[32..39] +GM[16..23] <- UB[64..71] +GM[24..31] <- UB[96..103] +``` diff --git a/docs/designs/vmi-vmull-pair-result-lowering-design.md b/docs/designs/vmi-vmull-pair-result-lowering-design.md new file mode 100644 index 0000000000..14fda600b2 --- /dev/null +++ b/docs/designs/vmi-vmull-pair-result-lowering-design.md @@ -0,0 +1,536 @@ +# VMI VMULL Pair-Result Lowering Design + +## Status + +Proposed design for `mouliangyu/PTOAS:feature-vmi`. + +This document defines the implementation contract for adding executable +`pto.vmi.vmull` support. The initial pull request containing this document is +design-only. ODS, verifier, lowering, PTODSL, user documentation, and tests are +follow-up implementation work. + +## 1. Motivation + +The physical VPTO operation already represents a native 32-bit widening +multiply as a pair of 32-bit vector results: + +```mlir +%low, %high = pto.vmull %lhs, %rhs, %mask + : !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.mask + -> !pto.vreg<64xi32>, !pto.vreg<64xi32> +``` + +One physical operation processes 64 `i32` or `ui32` lanes. VMI must expose the +same logical operation for 64, 128, and 256 lanes and split a larger logical +operation into the required physical operations automatically. + +The current VMI definition instead returns one `Lxi64` value: + +```mlir +%result = pto.vmi.vmull %a, %b, %mask + : !pto.vmi.vreg, !pto.vmi.vreg, !pto.vmi.mask + -> !pto.vmi.vreg +``` + +That form does not match the existing physical op. It would require a separate +virtual width-axis representation to translate one logical `i64` result into +low/high `i32` physical registers. The current VMI type converter has no such +result-axis contract, and there is no `VMIVmullOp` conversion pattern. + +The proposed surface therefore represents the widened product as two logical +`Lxi32` results. This makes every physical chunk a direct pair-result lowering +while retaining VMI's logical lane count. + +## 2. Goals and non-goals + +### 2.1 Goals + +- Support logical lane counts `L in {64, 128, 256}`. +- Support signed `i32` and unsigned `ui32` inputs. +- Return logical `%low` and `%high` values, both with the same `LxT` type as + the inputs. +- Lower a contiguous 64/128/256-lane operation to exactly 1/2/4 physical + `pto.vmull` operations. +- Support same-layout contiguous and deinterleaved factor-2/factor-4 dense + values with `lane_stride = 1`; the initial deinterleaved contract requires + `block_elems = 1`. +- Preserve one common layout relation across both inputs, the mask, and both + results. +- Preserve the low-result and high-result grouping expected by MLIR 1:N type + conversion. +- Expose the pair result through PTODSL as a Python tuple. +- Reject unsupported shapes and predicate modes before conversion instead of + leaving residual VMI IR. + +### 2.2 Non-goals + +- Supporting lane counts other than 64, 128, and 256. +- Reconstructing a logical `Lxi64` value from the low/high result pair. +- Adding a new VPTO op or changing the existing VPTO emitter. +- Supporting merge predication without explicit low/high passthrough values. +- Adding partial-register or tail-specific VMULL behavior. +- Supporting deinterleaved VMULL layouts with `block_elems != 1`. +- Implementing the code in the design-only pull request. + +## 3. Proposed VMI operation contract + +### 3.1 Syntax + +```mlir +%low, %high = pto.vmi.vmull %a, %b, %mask + : !pto.vmi.vreg, !pto.vmi.vreg, !pto.vmi.mask + -> !pto.vmi.vreg, !pto.vmi.vreg +``` + +Examples: + +```mlir +%low64, %high64 = pto.vmi.vmull %a64, %b64, %mask64 + : !pto.vmi.vreg<64xi32>, !pto.vmi.vreg<64xi32>, + !pto.vmi.mask<64xpred> + -> !pto.vmi.vreg<64xi32>, !pto.vmi.vreg<64xi32> + +%low256, %high256 = pto.vmi.vmull %a256, %b256, %mask256 + : !pto.vmi.vreg<256xui32>, !pto.vmi.vreg<256xui32>, + !pto.vmi.mask<256xpred> + -> !pto.vmi.vreg<256xui32>, !pto.vmi.vreg<256xui32> +``` + +### 3.2 Type constraints + +The verifier requires all of the following: + +```text +L is exactly 64, 128, or 256 +T is exactly i32 or ui32 +type(a) == type(b) == type(low) == type(high) +lanes(mask) == L +``` + +At the surface boundary, none of these types carries an assigned layout or a +concrete mask granularity. After assignment, the data values retain identical +element type, lane count, and layout. The mask has the same logical lane count +and layout and uses `b32` granularity. + +Mixed signedness is not legal. In particular, `i32` inputs cannot return +`ui32` results, and `ui32` inputs cannot return `i32` results. + +Here `i32` means MLIR's signless 32-bit integer type and `ui32` means its +unsigned 32-bit integer type. Explicitly signed `si32` is not an alias for +`i32` in this contract, and `si32`/`si64` must be rejected. The verifier must +test `isSignless()` or `isUnsigned()` explicitly; treating every type for +which `isUnsigned()` is false as signed would incorrectly admit `si32`. + +The old single `Lxi64` result form is removed rather than supported as a second +assembly form. Repository search shows no active VMI VMULL tests or kernel +consumers, so an atomic update of ODS, PTODSL, documentation, and tests is +preferred over maintaining two contracts. + +### 3.3 Per-lane semantics + +For active lane `i`, the operation computes one 64-bit product and returns its +two 32-bit halves: + +```text +if T == i32: + product = signed_64(a[i]) * signed_64(b[i]) +else: + product = unsigned_64(a[i]) * unsigned_64(b[i]) + +low[i] = bits(product, 31, 0) +high[i] = bits(product, 63, 32) +``` + +The result values use `T` so that signedness continues to select the existing +signed or unsigned physical VMULL form. The low 32 bits have the same bit +pattern for signed and unsigned multiplication; the high 32 bits reflect the +selected signedness. + +### 3.4 Predicate mode + +The first implementation is zeroing-only: + +```text +mask[i] == true: + low[i], high[i] are the product halves + +mask[i] == false: + low[i] = 0 + high[i] = 0 +``` + +For minimum assembly churn, `OptionalAttr:$pmode` may remain in ODS, +but the verifier accepts only an omitted attribute or `pmode = "zero"`. An +omitted attribute means zeroing. + +`pmode = "merge"` must be rejected by the verifier. The operation is pure and +has no old `%low` or `%high` passthrough operands, so there is no SSA value from +which inactive lanes could be preserved. Supporting merge later requires an +explicit API decision, such as adding two passthrough operands; it cannot be +implemented correctly by the conversion pattern alone. + +## 4. Layout contract + +VMULL is an elementwise pair producer. Logical lane `i` in `%a`, `%b`, and +`%mask` maps to logical lane `i` in both `%low` and `%high`. + +The assigned layout relation is: + +```text +layout(a) == layout(b) == layout(mask) == layout(low) == layout(high) +``` + +The initial implementation supports exactly these dense layouts: + +- contiguous with `lane_stride = 1`; +- deinterleaved factor 2 with `block_elems = 1` and `lane_stride = 1`; +- deinterleaved factor 4 with `block_elems = 1` and `lane_stride = 1`. + +Group-slot layouts, dense lane strides greater than one, and deinterleaved +layouts with any other positive `block_elems` are not part of the initial +contract. `VMILayoutAttr` may represent those layouts for other operations, +but VMULL preflight must reject them with an actionable diagnostic. + +Contiguous is the default layout and gives the required canonical expansion: + +| Logical type | Physical input parts | Physical VMULL count | Physical results | +|---|---:|---:|---:| +| `64xi32` | 1 | 1 | 1 low + 1 high | +| `128xi32` | 2 | 2 | 2 low + 2 high | +| `256xi32` | 4 | 4 | 4 low + 4 high | + +For the supported `block_elems = 1` deinterleaved layouts, the exact physical +arity is: + +| Logical lanes | factor 2 | factor 4 | +|---:|---:|---:| +| 64 | 2 | 4 | +| 128 | 2 | 4 | +| 256 | 4 | 4 | + +Lowering still obtains this arity from `getVMIPhysicalArity` rather than +hard-coding the table in the conversion pattern. All five logical values must +have the same physical arity. Corresponding parts of the four data values +`%a`, `%b`, `%low`, and `%high` must have the same `!pto.vreg<64xi32/ui32>` +type, while each mask part must be the corresponding `!pto.mask`. +Restricting `block_elems` closes the initial support set explicitly; for +example, +`256xi32` with `block_elems = 65` is rejected instead of silently producing +the otherwise computable factor-2 arity 5 or factor-4 arity 7. + +### 4.1 Assignment and propagation integration + +VMULL cannot be implemented only in the final conversion pattern. The earlier +passes must establish its mask and layout contract: + +1. `VMIMaskGranularityAssignment` requests `b32` for the VMULL mask because + the data element type is 32-bit. +2. `VMILayoutAssignment` follows the ordinary elementwise `unite()` path for + `%a`, `%b`, `%low`, and `%high`. This registers the values with the layout + solver without placing them in a hard data-layout DSU equivalence class. +3. `VMILayoutPropagation` treats `VMIVmullOp` as a same-layout operation so a + layout fact on any input, mask, or result propagates to the other ports. +4. Conflicting consumer layouts are handled through the existing explicit + `ensure_layout` or `ensure_mask_layout` materialization mechanism. + +`uniteDataEquivalent` must not be used for VMULL. It is reserved for values +that are genuinely the same SSA value across control-flow, call, and function +boundaries. Using it for an elementwise producer would turn compatible layout +requests into hard natural/preferred-layout conflicts and prevent use-site +materialization. + +No VMULL-specific layout transform is needed. The relation is identity across +all ports, unlike `vintlv`/`vdintlv`, whose input and result layouts may differ. + +## 5. VMI-to-VPTO lowering + +### 5.1 Canonical contiguous expansion + +For contiguous `Lxi32`, define: + +```text +K = L / 64 +``` + +The type converter produces: + +```text +a -> [a_0, ..., a_(K-1)] +b -> [b_0, ..., b_(K-1)] +mask -> [mask_0, ..., mask_(K-1)] +low -> [low_0, ..., low_(K-1)] +high -> [high_0, ..., high_(K-1)] +``` + +The conversion creates one physical operation for every `p` in `[0, K)`: + +```text +low_p, high_p = pto.vmull(a_p, b_p, mask_p) + +a_p, b_p, low_p, high_p : !pto.vreg<64xi32> or !pto.vreg<64xui32> +mask_p : !pto.mask +``` + +### 5.2 Conversion pattern algorithm + +`OneToNVMIVmullOpPattern` performs the following checks and rewrite: + +1. Read `aParts`, `bParts`, and `maskParts` from the 1:N adaptor. +2. Obtain the converted type lists for result index 0 (`low`) and result index + 1 (`high`). +3. Require all five arities to be equal and non-zero. +4. Require every input and result part to be `!pto.vreg<64xi32>` or + `!pto.vreg<64xui32>` with matching signedness, and every mask part to be + `!pto.mask`. +5. For physical part `p`, create + `pto.vmull(aParts[p], bParts[p], maskParts[p])`. +6. Replace the two logical results with the flattened physical result list. + +The flattened replacement order is critical: + +```text +[low_0, low_1, ..., low_(K-1), high_0, high_1, ..., high_(K-1)] +``` + +It must not be interleaved as `[low_0, high_0, low_1, high_1, ...]`. +`replaceOpWithFlatConvertedValues` partitions the flat list by logical result +index, so the first complete segment belongs to `%low` and the second complete +segment belongs to `%high`. + +### 5.3 Preflight validation + +`verifySupportedVMIToVPTOOps` gains an explicit VMULL check before dialect +conversion. The check requires: + +- assigned, equal, supported dense layouts on all ports; +- contiguous, or deinterleaved factor 2/4 with `block_elems = 1`; +- `lane_stride = 1` on every port; +- `b32` mask granularity; +- element type exactly signless `i32` or unsigned `ui32`, and a legal lane + count; +- matching, computable physical arity for inputs, mask, and both results; +- physical part types compatible with `pto.vmull`. + +This gives an actionable `VMI-UNSUPPORTED` diagnostic instead of a final +`VMI-RESIDUAL-OP` failure. + +## 6. Pipeline integration + +The complete implementation crosses the following layers: + +| Layer | File | Required change | +|---|---|---| +| ODS | `include/PTO/IR/VMIOps.td` | Change one `Lxi64` result to `(low, high)` `Lxi32` results and update syntax/description | +| Verifier | `lib/PTO/IR/VMI.cpp` | Enforce legal lane counts, exact signless/unsigned types, pair equality, mask shape, and zero-only pmode | +| Mask assignment | `lib/PTO/Transforms/VMIMaskGranularityAssignment.cpp` | Request `b32` for the mask use | +| Layout assignment | `lib/PTO/Transforms/VMILayoutAssignment.cpp` | Use ordinary elementwise `unite()` bookkeeping for both inputs and both results; do not use `uniteDataEquivalent` | +| Layout propagation | `lib/PTO/Transforms/VMILayoutPropagation.cpp` | Register VMULL as a same-layout relation | +| Unified bridge | `lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp` | Keep VMULL on the direct-to-VPTO path; update comments only if needed | +| Physicalization | `lib/PTO/Transforms/VMIToVPTO.cpp` | Add preflight validation, pair-result 1:N pattern, and pattern registration | +| PTODSL | `ptodsl/ptodsl/_vmi_namespace.py` | Return two results and infer their types from the inputs | +| User docs | VMI ISA and PTODSL guide | Replace the single widened result with the pair-result contract | +| Tests | `test/lit/vmi_new`, `ptodsl/tests` | Add verifier, assignment, lowering, and frontend coverage | + +The existing `PTO_VmullOp`, `VmullOp::verify`, LLVM emitters, and physical +VMULL tests remain unchanged. + +## 7. PTODSL API + +The Python surface follows the existing two-result `vintlv`/`vdintlv` style: + +```python +low, high = pto.vmi.vmull( + a, + b, + mask, + pmode=None, +) +``` + +Both result types are inferred from the matching input type. The Python +surface does not accept explicit result types; the IR verifier requires both +result types to equal the input type. + +The old required `result_type=` keyword is removed atomically with the ODS +change. Because there are no in-repository PTODSL VMULL call sites, a temporary +dual API is not required. The user guide must clearly document that the return +value is a tuple `(low, high)`. + +## 8. Test plan + +### 8.1 ODS and verifier tests + +Add positive parse/verify coverage for: + +- `64xi32`; +- `128xi32`; +- `256xi32`; +- `64xui32` or another unsigned shape; +- omitted pmode and explicit `pmode = "zero"`. + +Add negative coverage for: + +- lane count outside `{64, 128, 256}`; +- input element type other than `i32/ui32`, with explicit `si32` and `si64` + cases so the current width-only/`isUnsigned()` behavior cannot pass; +- `si32` low/high result types even when both inputs use `si32`; +- input type or signedness mismatch; +- low/high type, signedness, or lane-count mismatch; +- mask lane-count mismatch; +- the old single `Lxi64` result form; +- `pmode = "merge"` and unknown pmode strings. + +### 8.2 Mask and layout tests + +- Verify that `!pto.vmi.mask` becomes `b32` for VMULL. +- Verify that `%a`, `%b`, `%mask`, `%low`, and `%high` receive the same layout. +- Verify that a conflicting consumer layout inserts an explicit + materialization rather than silently changing one VMULL port. +- Cover deinterleaved factor-2 and factor-4 same-layout cases with + `block_elems = 1`. Both are mandatory first-version tests and must check the + exact arities above, corresponding `%a`/`%b`/`%mask` part alignment, + low/high result grouping, and inactive or padding-lane mask alignment. +- Add preflight rejection tests for deinterleaved `block_elems != 1`, including + `256xi32` factor-2/factor-4 with `block_elems = 65`. These cases would have + arity 5/7 and ensure the implementation does not accept arbitrary layouts + merely because `getVMIPhysicalArity` can compute them. + +### 8.3 VMI-to-VPTO tests + +For contiguous layouts, check exact physical operation counts: + +```text +64xi32 -> CHECK-COUNT-1: pto.vmull +128xi32 -> CHECK-COUNT-2: pto.vmull +256xi32 -> CHECK-COUNT-4: pto.vmull +``` + +The multi-chunk checks must also capture result grouping, not only operation +count. Uses of the logical `%low` value must receive all low parts, while uses +of `%high` must receive all high parts. + +Add both signed and unsigned lowering coverage and verify `!pto.mask` on +every emitted physical operation. + +### 8.4 PTODSL and end-to-end tests + +- Verify `pto.vmi.vmull` returns a two-item tuple. +- Verify both result types are inferred for signed and unsigned inputs. +- Verify invalid input types fail with a clear IR verifier diagnostic. +- Compile at least one VMI-authored kernel through the complete VPTO pipeline + and confirm no VMI op or type remains. + +### 8.5 Numerical semantics tests + +Compile-only coverage is not sufficient for VMULL. At least one executable +numerical backend must compare both result vectors against a scalar reference +oracle. The PTO ISA CPU simulator is preferred; if it does not support the +required VMULL form, equivalent A5 execution is mandatory for completion. + +The reference oracle computes each lane independently: + +```text +if mask[i]: + if T == i32: + product = signed_64(signed_32(a[i])) * signed_64(signed_32(b[i])) + else: + product = unsigned_64(unsigned_32(a[i])) * + unsigned_64(unsigned_32(b[i])) + expected_low[i] = unsigned_32(product) + expected_high[i] = unsigned_32(product >> 32) +else: + expected_low[i] = 0 + expected_high[i] = 0 +``` + +Required signed boundary cases include: + +```text +(-1) * 2 -> low=0xfffffffe, high=0xffffffff +INT32_MIN * (-1) -> low=0x80000000, high=0x00000000 +INT32_MIN * 2 -> low=0x00000000, high=0xffffffff +INT32_MAX * INT32_MAX -> low=0x00000001, high=0x3fffffff +``` + +Required unsigned boundary cases include: + +```text +0xffffffff * 0xffffffff -> low=0x00000001, high=0xfffffffe +0x80000000 * 2 -> low=0x00000000, high=0x00000001 +``` + +The numerical suite must also include: + +- logical sizes 64, 128, and 256; +- distinct per-chunk values so a chunk-ordering error is observable; +- sparse masks spanning chunk boundaries, including lanes 0, 63, 64, 127, + 128, and 255 when present; +- inactive input lanes containing non-zero sentinel values, with both output + halves checked to be zero; +- deinterleaved factor-2 and factor-4 `block_elems = 1` cases, including + logical lanes that map to different physical parts. + +Inactive lanes must remain observable. The test kernel uses a sparse mask for +VMULL, but writes `%low` and `%high` to their output buffers with either an +unmasked store or a separate all-true store mask. It must not reuse the sparse +VMULL mask for either store. The output buffers are initialized with non-zero +sentinels, every logical lane is written and read back, and inactive lanes are +compared with zero. Otherwise a masked store could hide incorrect non-zero +VMULL results in precisely the lanes this test is intended to validate. + +These checks detect signed high-half errors, low/high swaps, incorrect flat +result grouping, physical part reordering, and failure to zero inactive lanes. + +## 9. Compatibility and rollout + +Changing one `Lxi64` result to two `Lxi32` results is an intentional breaking +change to an incomplete, currently non-lowerable VMI operation. All public +layers must change in one implementation pull request: + +```text +ODS + verifier +mask/layout assignment +VMIToVPTO preflight + conversion +PTODSL +VMI ISA and PTODSL documentation +focused regression tests +``` + +Splitting the ODS change from the conversion pattern would temporarily create +another parseable but non-executable VMULL form and is therefore not +recommended. + +## 10. Acceptance criteria + +The implementation is complete when: + +1. The only accepted VMI VMULL surface is pair-result `Lxi32/ui32`, with + `L in {64, 128, 256}`; explicit `si32` and `si64` forms are rejected. +2. Mask granularity and all five layouts are assigned consistently. +3. Contiguous 64/128/256-lane inputs produce exactly 1/2/4 physical + `pto.vmull` operations. +4. Low and high result parts are associated with the correct logical result. +5. Signed and unsigned forms select their existing physical backend forms. +6. Merge predication is rejected before physical conversion. +7. Deinterleaved factor-2 and factor-4 with `block_elems = 1` pass mandatory + arity, part-alignment, result-grouping, and mask-alignment tests, while + non-1 `block_elems` receives a preflight diagnostic. +8. CPU simulator or A5 numerical validation confirms signed/unsigned high and + low halves and zeroing for sparse inactive lanes observed through full-lane + output stores. +9. The full pipeline contains no residual VMI op or VMI type. +10. ODS, verifier, lowering, PTODSL, user documentation, and tests land + together in the follow-up implementation pull request. + +## 11. Decisions requested from review + +This proposal asks reviewers to confirm two API decisions before code is +implemented: + +1. Use two logical `Lxi32/ui32` results instead of one logical `Lxi64/ui64` + result. +2. Define the first implementation as zeroing-only and reject merge until the + operation has explicit passthrough semantics. + +Once those decisions are accepted, the implementation path is fully defined +by the contracts above. diff --git a/docs/isa/micro-isa/10-reduction-ops.md b/docs/isa/micro-isa/10-reduction-ops.md index ebe95a7bf9..b03f4696c1 100644 --- a/docs/isa/micro-isa/10-reduction-ops.md +++ b/docs/isa/micro-isa/10-reduction-ops.md @@ -46,7 +46,7 @@ for (int i = 1; i < N; i++) ### `pto.vcmax` - **syntax:** `%result = pto.vcmax %input, %mask : !pto.vreg, !pto.mask -> !pto.vreg` -- **A5 types:** i16-i32, f16, f32 +- **A5 types:** i8-i32, f16, f32 - **semantics:** Find max element with argmax. The lowest destination element stores the maximum value, the second-lowest destination element stores the index of the first maximum, and all remaining elements are zero-filled. @@ -78,7 +78,7 @@ for (int i = 2; i < N; i++) ### `pto.vcmin` - **syntax:** `%result = pto.vcmin %input, %mask : !pto.vreg, !pto.mask -> !pto.vreg` -- **A5 types:** i16-i32, f16, f32 +- **A5 types:** i8-i32, f16, f32 - **semantics:** Find min element with argmin. The lowest destination element stores the minimum value, the second-lowest destination element stores the index of the first minimum, and all remaining elements are zero-filled. @@ -208,7 +208,7 @@ VLane 4: [32..39] VLane 5: [40..47] VLane 6: [48..55] VLane 7: [56..63] ### `pto.vcgadd` - **syntax:** `%result = pto.vcgadd %input, %mask : !pto.vreg, !pto.mask -> !pto.vreg` -- **A5 types:** i16-i32, f16, f32 +- **A5 types:** i8-i32, f16, f32 - **semantics:** Sum active elements within each 32-byte VLane. The 8 VLane sums are written to result elements `0..7`; all other result elements are zero. @@ -240,7 +240,7 @@ for (int i = groups; i < N; i++) ### `pto.vcgmax` - **syntax:** `%result = pto.vcgmax %input, %mask : !pto.vreg, !pto.mask -> !pto.vreg` -- **A5 types:** i16-i32, f16, f32 +- **A5 types:** i8-i32, f16, f32 - **semantics:** Find the maximum active element within each 32-byte VLane. The 8 VLane maxima are written to result elements `0..7`; all other result elements are zero. @@ -275,7 +275,7 @@ for (int i = groups; i < N; i++) ### `pto.vcgmin` - **syntax:** `%result = pto.vcgmin %input, %mask : !pto.vreg, !pto.mask -> !pto.vreg` -- **A5 types:** i16-i32, f16, f32 +- **A5 types:** i8-i32, f16, f32 - **semantics:** Find the minimum active element within each 32-byte VLane. The 8 VLane minima are written to result elements `0..7`; all other result elements are zero. diff --git a/docs/isa/vmi-isa/00-architecture-overview.md b/docs/isa/vmi-isa/00-architecture-overview.md new file mode 100644 index 0000000000..44d915d695 --- /dev/null +++ b/docs/isa/vmi-isa/00-architecture-overview.md @@ -0,0 +1,176 @@ +# VMI Architecture Overview + +> **Status:** draft. This document covers the architecture and foundational concepts +> of the unified `pto.vmi` instruction surface. Per-op reference docs are in the +> numbered group files that follow. + +`pto.vmi` sits between high-level programming models (TileLang, pto-dsl) and +the physical `pto.mi` ISA. It exposes **logically contiguous vectors** and +**elementwise compute intent**; the physical SIMD register layout (interleave, +parity, width, part, pack, dist tokens) is held and propagated by `pto.as` and +is invisible to the user. + +``` +TileLang T.parallel(N) { C[i] = cast(A[i]) + B[i] } + │ (direct translation, elementwise semantics preserved) + ▼ +pto.vmi %w = pto.vmi.vcvt %a; %c = pto.vmi.vadd %w, %b + │ (pto.as: layout-assignment + lowering) + ▼ +pto.mi vcvt EVEN/ODD + two-way vadd + vstsx2 INTLV_B32 +``` + +- **Upper → vmi**: `T.parallel`'s logical iteration space translates directly + to `pto.vmi` logical vector ops — elementwise → Category A op, `T.cast` → + a `vcvt` with no explicit `part`, logical length `N` → + `!pto.vmi.vreg`, "all active" → auto-generated tail predicate. +- **vmi → pto.mi**: `pto.as` performs layout inference + unification + + materialization, lowering logical vectors to concrete `pto.mi` instructions + (including `part/pack/interleave/dist`). At `K=1` this degenerates to + zero-overhead pass-through. + +--- + +## Logical vs Physical + +A `pto.vmi` value is **logical** — a flat sequence of `L` lanes of type `T`. +Its physical backing is `K` hardware vector registers (256B / 2048-bit each): + +``` +K = ⌈ L · bitwidth(T) / 2048 ⌉ +``` + +At `K=1` and full-width (no partial lanes), one `pto.vmi.vreg` maps 1:1 to +one `pto.vreg`. At `K>1`, the logical value fans out across `K` physical +registers with a layout descriptor (`#pto.vmi.layout`) tracking the mapping. + +**Physical constants (A5 vector pipe):** + +``` +vector register file : 32 architectural vregs, 256 B (2048 bit) each +predicate file : 8 architectural pregs, 256 bit each, 1 bit controls 1 byte +VLane : 32 B sub-lane; 8 VLanes per vreg +E_v = 32 / sizeof(T) : lanes per VLane (f32 → 8, f16/bf16 → 16, i8 → 32) +``` + +--- + +## Type System + +### `!pto.vmi.vreg` + +Logical vector register. `L` is the logical lane count; `T` is the element type. + +| T | bits | E_v (lanes per physical vreg) | Legal L multiples | +|---|---|---|---| +| `f32` / `i32` / `ui32` / `si32` | 32 | 64 | 64 | +| `f16` / `bf16` / `i16` / `ui16` / `si16` | 16 | 128 | 64 | +| `i8` / `ui8` / `si8` / `fp8_e4m3` / `fp8_e5m2` | 8 | 256 | 64 | + +- **Full vector**: `L · bitwidth(T) == N · 2048` (integer multiple of 256B). +- **Compact/partial vector**: `L · bitwidth(T) < 2048` — still backed by one + physical vreg (256B); only the low `L` logical slots are valid. Physical + slots outside the logical value are `pad/undef` and must be masked out. + +**Common logical ↔ physical mappings:** + +| Logical type | Byte size | K | Physical vregs | Valid slots per vreg | +|---|---:|---:|---:|---| +| `V<256×f32>` | 1024B | 4 | 4 | 64 f32 each, all valid | +| `V<256×f16>` | 512B | 2 | 2 | 128 f16 each, all valid | +| `V<256×i8>` | 256B | 1 | 1 | 256 i8, all valid | +| `V<128×f32>` | 512B | 2 | 2 | 64 f32 each, all valid | +| `V<64×f16>` | 128B | 1 | 1 | low 64 f16 valid | +| `V<64×i8>` | 64B | 1 | 1 | low 64 i8 valid | + +See the [Design Doc](../PTO-vmi-design.md) for detailed physical layout +diagrams (contiguous, parity EVEN/ODD, sub-part, stride-4 interleave) for each +logical type. + +### `!pto.vmi.mask` + +Virtual predicate mask. Each logical mask lane corresponds to one logical +vector lane (`L` must match the governed vreg's `L`). + +--- + +## Category A / B / C + +Every VMI op belongs to one of three lowering categories that determine how +`pto.as` handles its physical layout: + +| Category | Layout relationship | `pto.as` behavior | Output layout | +|---|---|---|---| +| **A — Layout-passthrough** | Does not modify register layout | Fan-out: emit the same `pto.mi` op once per physical reg (`K × op`); mask follows per-reg (with `ppack`/`punpack` as needed) | Unchanged: preserves input parity/half/sub-part layout | +| **B — Layout-rewritable** | Modifies layout predictably | Fan-out along other axes; instantiate matching modes (`PART_EVEN/ODD`, `Bin_N0/N1`, `PK`/`UNPK`, `INTLV`/`DINTLV`) | Rewritten to the op's natural output layout | +| **C — Contiguous-required** | Requires stride-1 contiguous input (no in-place mode satisfies it) | `pto.as` inserts `.contiguous()` materialization (store+reload or explicit repack) before the op | Flattened contiguous chunk (`is_contiguous`) | + +> **C-class note:** C-class ops cannot tolerate a non-contiguous physical +> layout — any parity/half/sub-part arrangement must first be materialized to +> contiguous before the op runs. `pto.as` therefore treats a C-class op as a +> **layout barrier**: upstream A/B ops may keep their compact layout right up to +> the C-class boundary, where a `.contiguous()` is forced. + +--- + +## Mask & Predication (`pmode`) + +All compute ops accept an optional governing mask operand `[pmode]`. The mask +is a `!pto.vmi.mask` with the same `L` as the data operand. + +**`pmode` values:** + +| `pmode` | Inactive lane behavior | Default? | +|---|---|---| +| `"zero"` | Inactive lanes produce 0 (hardware-native ZEROING) | ✓ (default) | +| `"merge"` | Inactive lanes preserve the destination's prior value | | + +On A5, MERGE is **emulated**: the hardware predicates only in ZEROING mode, so the +compiler synthesizes merge as a predicate complement plus a `vor`/`vsel` blend +of the zeroed result with the old destination (see [Appendix C](10-appendices.md)). +On A6, some ops support native MERGE. + +**A5 load restriction**: `vload` has **no** mask operand — A5 loads are +unpredicated. A logical tail mask associated with a load is never lowered as a +"masked load"; `pto.as` migrates it to the consuming compute op, the store, or +shortens the load length. `vstore` **is** predicated on A5. + +--- + +## The `group` Attribute + +Reduce ops (`vcadd`, `vcmax`, `vcmin`) and broadcast (`vbrc`) accept an +optional `{group=C}` attribute where `C` is the **number of groups** (not the +per-group lane count): + +- **Reduce**: Splits `L` lanes into `C` groups, each producing one scalar. + Output is `V` — a compact vector of `C` scalars. +- **Broadcast**: Takes a compact `V` and fans each scalar back across + `L/C` lanes, producing `V`. + +Legal `C` values: `1`, `2`, `4`, `8` (must divide `L`; must match the result +type's `C`). + +**`group → Category` decision table** (W = bytes per sub-group): + +| W vs BlockLane (32B) | Category | Lowering | +|---|---|---| +| `W == 32B` (sub-group = 1 VLane) | B | `vcgadd`/`vcgmax`/`vcgmin` — one op per reg, no cross-reg combine | +| `W > 32B`, aligned | B | Fold `(k-1)× vadd/vmax/vmin` then `vcg*` | +| Unaligned | C | Materialize → contiguous → reduce | + +--- + +## Group Index + +| # | Group | Ops | Category | Mask | +|---|---|---|---|---| +| 1 | **Load / Store** | `vload`, `vstore` | A (+B on dintlv/unpack) | load: none; store: `Pg` | +| 2 | **Index-gen** | `vci` | A | none | +| 3 | **Eltwise Compute** | `vadd`, `vsub`, `vmul`, `vdiv`, `vmax`, `vmin`, `vabs`, `vneg`, `vrelu`, `vexp`, `vln`, `vsqrt`, `vand`, `vor`, `vxor`, `vnot`, `vshl`, `vshr`, `vadds`, `vmuls`, `vmaxs`, `vmins`, `vshls`, `vshrs`, `vcmp`, `vcmps`, `vsel`, `vselr` | A | `Pg` (except `vselr`: none) | +| 4 | **Broadcast** | `vbrc` | A (ungrouped) / B (grouped) | none | +| 5 | **Reduce** | `vcadd`, `vcmax`, `vcmin` | B (VLane-aligned) / C (unaligned) | `Pg req` | +| 6 | **Convert** | `vcvt`, `vinterpret_cast` | B / A | `Pg` / none | +| 7 | **SFU** | `vexpdif`, `vaxpy`, `vlrelu`, `vprelu`, `vmull`, `vmula`, `vchist`, `vdhist`, `vgather`, `vgatherb`, `vscatter` | A (fused) / B (vmull, vchist, vdhist) / C (gather/scatter) | `Pg` (`vchist`/`vdhist`/SFU) / `Pg` (gather/scatter) | +| 8 | **Predicate Ops** | `create_mask`, `create_group_mask` | gen | gen | +| 9 | **Data Rearrange** | `vintlv`, `vdintlv` | A | `Pg` | diff --git a/docs/isa/vmi-isa/01-load-store.md b/docs/isa/vmi-isa/01-load-store.md new file mode 100644 index 0000000000..6a9b4314f6 --- /dev/null +++ b/docs/isa/vmi-isa/01-load-store.md @@ -0,0 +1,260 @@ +# 1. Load / Store + +> **Category:** A (+B on `dintlv`/`unpack`). **Mask:** load none (A5 loads are unpredicated), store `Pg`. +> +> `vload`/`vstore` are logical memory ops. **`[dist_mode]` explicitly declares +> the access pattern**, defaulting to `continuous` (contiguous); the optional +> modes are `unpack` (widening unpack), `dintlv` (deinterleave), and `brc` +> (broadcast). + +--- + +## `pto.vmi.vload` + +- **semantics:** Load elements of type `T` from UB into a logical vector + register starting at `%source + %offset` (element offset). The default + (`continuous`) is a contiguous stride-1 read: + + ```c + for (int i = 0; i < L; i++) + dst[i] = ub[base + offset + i]; + ``` + + The access pattern is not always contiguous: depending on the attributes + (`{dist_mode}`, `{group = C}` with a `stride` operand, or + `%block_stride`), the load may instead read in a strided/scattered + fashion (e.g. per-row stride for group mode, 32B-block stride for + block-stride mode), widen/deinterleave the source, or broadcast. The exact + pattern is determined by these mutually exclusive attributes (see + attributes and lowering below). + +- **syntax:** + ```mlir + %result = pto.vmi.vload %source[%offset] : !pto.ptr -> !pto.vmi.vreg + ``` +- **syntax (`dintlv`):** + ```mlir + // fused load + deinterleave → 2 results + %even, %odd = pto.vmi.vload %source[%offset] {dist_mode = "dintlv"} + : !pto.ptr -> !pto.vmi.vreg, !pto.vmi.vreg + ``` +- **syntax (`group`):** + ```mlir + // strided group load: C rows of L/C elements, row g at base + g*stride + %result = pto.vmi.vload %source[%offset], %stride {group = C} + : !pto.ptr, index -> !pto.vmi.vreg + ``` +- **syntax (block-stride):** + ```mlir + // block-strided load: %block_stride is a dynamic i16 operand (no mask) + %result = pto.vmi.vload %source[%offset], %block_stride + : !pto.ptr, i16 -> !pto.vmi.vreg + ``` +- **operands:** + + | Operand | Type | Description | + |---|---|---| + | `source` | `!pto.ptr` | UB base pointer | + | `offset` | `index` | Element offset from base | + | `stride` | `index` | Per-row stride (element units); required with `{group}`, invalid otherwise | + | `block_stride` | `i16` | 32B-block stride between scattered blocks (block-stride mode); mutually exclusive with `{group}` and `{dist_mode}` | + +- **results:** + + | Result | Type | Description | + |---|---|---| + | `result` | `!pto.vmi.vreg` | Loaded logical vector (1 result: `continuous`/`unpack`/`brc`) | + | `even`, `odd` | two `!pto.vmi.vreg` | Deinterleaved pair (2 results, `dintlv` only) | + +- **attributes:** + + | Attribute | Values | Default | Description | + |---|---|---|---| + | `dist_mode` | `"continuous"`, `"unpack"`, `"dintlv"`, `"brc"` | `"continuous"` | Memory access pattern | + | `group` | positive integer | *(none)* | Strided group load arity; mutually exclusive with `dist_mode`; requires `stride` | + | `pmode` | `"zero"`, `"merge"` | `"zero"` | Inactive-lane behavior (applied at consumer, not on load) | + +- **lowering to `pto.mi`:** + + | `dist_mode` | Physical lowering | + |---|---| + | `"continuous"` | `K × pto.vlds {dist="NORM"}` (element-width-independent `NORM` load) | + | `"unpack"` | `K × pto.vlds {dist="UNPK_B*"}` (widening unpack; suffix from `Ptr`) | + | `"dintlv"` | `K × pto.vldsx2 {dist="DINTLV_B*"}` (dual deinterleave load); surface: 2 results `(%even, %odd)`, one per parity half | + | `"brc"` | `1 × pto.vlds {dist="BRC_B*"}` or `BRC_BLK`; broadcast-axis (1-reg backing, replicate-read) | + + **Group mode** (`{group = C}` + `stride`) has two sub-cases, decided by the + relation between `result.L` and `C`: + - **Full-group load** (`result.L > C`): each group loads `L/C` elements, + row-strided tile load: `C·(L/C) = L` elements across `C` rows, + each row `g` at offset `base + g·stride`. + - **Slot load** (`result.L == C`): each group loads **1 scalar** into the + corresponding slot, producing a compact `V`. This is the + dual of group reduce — reduce + folds lanes into slots, slot load reads those slots back into a vreg. + `C ∈ {1, 2, 4, 8}`. Not combinable with `dist_mode`. + + **Block-stride mode** (`%block_stride` operand): 2D-tile block-strided load. + Memory is read in 32B blocks with block `blk` at + `base + blk * block_stride` (scattered access); the internal repeat stride + defaults to 0. `%block_stride` is a dynamic `i16` operand. A5 loads are + unpredicated, so an implicit all-active mask is applied. Not combinable + with `dist_mode` or `group`. + + `B*` suffix is derived from `Ptr` element width: `f32/i32 → B32`, `f16/bf16/i16 → B16`, `i8/fp8 → B8`. + +- **examples:** + + ```mlir + // Continuous load (default dist_mode): UB → vreg + %v = pto.vmi.vload %ub[%offset] : !pto.ptr -> !pto.vmi.vreg<64×f32> + // → pto.as: Ptr → B32, dist_mode=continuous → pto.mi.vlds {dist="NORM"} + + // Slot load: 1 scalar per group → compact V<8×f32> (reads back reduce output) + %s = pto.vmi.vload %ub[%off], %stride {group = 8} + : !pto.ptr, index -> !pto.vmi.vreg<8×f32> + + // Full-group load: 8 rows × 8 elements, stride 64 + %t = pto.vmi.vload %ub[%off], %stride {group = 8} + : !pto.ptr, index -> !pto.vmi.vreg<64×f32> + + // Block-strided load: block_stride = 8 (dynamic i16 operand, no mask) + %vb = pto.vmi.vload %ub[%off], %c8_i16 + : !pto.ptr, i16 -> !pto.vmi.vreg<64×f32> + + // Broadcast load: scalar/block replicate into vreg + %vb = pto.vmi.vload %ub[%offset] {dist_mode = "brc"} : !pto.ptr -> !pto.vmi.vreg<64×f32> + + // Widening unpack load: narrow source expanded to wide lanes + %u = pto.vmi.vload %ub[%offset] {dist_mode = "unpack"} : !pto.ptr -> !pto.vmi.vreg<64×f32> + + // Deinterleave load: fused load + deinterleave, 2 surface results + %even, %odd = pto.vmi.vload %ub[%offset] {dist_mode = "dintlv"} + : !pto.ptr -> !pto.vmi.vreg<64×f32>, !pto.vmi.vreg<64×f32> + ``` + +- **notes:** + - **A5 loads are unpredicated.** A tail mask associated with a `vload` is + never lowered as a masked load. It migrates to the consuming compute op or + to a `vstore`. + - `dist_mode` and layout inference are orthogonal: even with `dist_mode="continuous"`, + `pto.as` may lower to `DINTLV_B*` to serve a downstream grouped reduce. + - The `pmode` attribute on `vload` governs the result lane behavior at the + *consumer*, not on the load itself. + +- **attention:** + - **Result count must match the access mode.** `dist_mode = "dintlv"` is a + fused load + deinterleave and produces **two** results `(%even, %odd)`; + all other `dist_mode` values, `{group}`, and `%block_stride` produce + **one** result. If the written result count does not match the selected + mode (e.g. a single result with `dintlv`, or two results with + `continuous`), `pto.as` rejects the op. + - **`{group}`, `%block_stride`, and `{dist_mode}` are mutually exclusive.** + Specifying more than one at once is rejected by `pto.as`. + - **`stride` operand is bound to `{group}`.** It is required with + `{group = C}` and invalid otherwise; `block_stride` is bound to the + block-stride mode and invalid otherwise. `vload` has no mask operand in + any mode (A5 loads are unpredicated). + +--- + +## `pto.vmi.vstore` + +- **semantics:** Store elements from a vector register to UB starting at + `%dest + %offset` (element offset). The default (`continuous`) is a + contiguous stride-1 write; only lanes where `mask[i] != 0` are written + (A5 stores are predicated): + + ```c + for (int i = 0; i < L; i++) + if (mask[i]) + ub[base + offset + i] = src[i]; + ``` + + The access pattern is not always contiguous: depending on the attributes + (`{dist_mode}`, `{group = C}` with a `stride` operand, or + `%block_stride`), the store may instead write in a strided/scattered + fashion (e.g. per-row stride for group mode, 32B-block stride for + block-stride mode) or interleave the values. The exact pattern is + determined by these mutually exclusive attributes (see attributes and + lowering below). + +- **syntax:** + ```mlir + pto.vmi.vstore %value, %dest[%offset], %mask : !pto.vmi.vreg, !pto.ptr, !pto.vmi.mask + ``` +- **syntax (`dintlv`):** + ```mlir + // fused interleave + store → 2 values + pto.vmi.vstore %even, %odd, %dest[%offset], %mask {dist_mode = "dintlv"} + : !pto.vmi.vreg, !pto.vmi.vreg, !pto.ptr, !pto.vmi.mask + ``` +- **syntax (`group`):** + ```mlir + // strided group store: C rows of L/C elements, row g at base + g*stride (no mask) + pto.vmi.vstore %value, %dest[%offset], %stride {group = C} + : !pto.vmi.vreg, !pto.ptr, index + ``` +- **syntax (block-stride):** + ```mlir + // block-strided store: %block_stride is a dynamic i16 operand (mask required) + pto.vmi.vstore %value, %dest[%offset], %block_stride, %mask + : !pto.vmi.vreg, !pto.ptr, i16, !pto.vmi.mask + ``` +- **operands:** + + | Operand | Type | Description | + |---|---|---| + | `value` | `!pto.vmi.vreg` | Vector value to store (1 value, `continuous`) | + | `even`, `odd` | two `!pto.vmi.vreg` | Interleaved pair to store (`dintlv` only) | + | `dest` | `!pto.ptr` | UB destination base pointer | + | `offset` | `index` | Element offset from base | + | `stride` | `index` | Per-row stride (element units); required with `{group}`, invalid otherwise | + | `block_stride` | `i16` | 32B-block stride between scattered blocks (block-stride mode); mutually exclusive with `{group}` and `{dist_mode}` | + | `mask` | `!pto.vmi.mask` | Governing predicate (variadic: 0 or 1) | + +- **results:** *(none)* + +- **attributes:** + + | Attribute | Values | Default | Description | + |---|---|---|---| + | `dist_mode` | `"continuous"`, `"dintlv"` | `"continuous"` | Memory access pattern | + | `group` | positive integer | *(none)* | Strided group store arity; mutually exclusive with `dist_mode`; requires `stride`; forbids `mask` | + | `pmode` | `"zero"`, `"merge"` | `"zero"` | Inactive-lane behavior: `"zero"` (default) stores 0; `"merge"` skips write on inactive lanes | + +- **lowering to `pto.mi`:** + + | `dist_mode` | Physical lowering | + |---|---| + | `"continuous"` | `K × pto.vsts {dist="NORM_B*"}` | + | `"dintlv"` | `K × pto.vstsx2 {dist="INTLV_B*"}`; surface consumes 2 inputs `(%even, %odd)`, interleaved at lowering | + + **Group mode** (`{group = C}` + `stride`): row-strided tile store. Not combinable with + `dist_mode` or `mask` (group stores are unpredicated). + + **Block-stride mode** (`%block_stride` operand): 2D-tile block-strided store. + Memory is written in 32B blocks with block `blk` at + `base + blk * block_stride` (scattered access); the internal repeat stride + defaults to 0. `%block_stride` is a dynamic `i16` operand. An explicit + `mask` is applied; if absent an implicit all-active mask is used. Not + combinable with `dist_mode` or `group`. + +- **examples:** + + ```mlir + // Continuous store (default): vreg → UB, masked + pto.vmi.vstore %v, %ub_out[%offset], %mask : !pto.vmi.vreg<64×f32>, !pto.ptr, !pto.vmi.mask<64> + + // Interleave store: fused interleave + store, 2 surface inputs + pto.vmi.vstore %even, %odd, %ub_out[%offset], %mask {dist_mode = "dintlv"} + : !pto.vmi.vreg<64×f32>, !pto.vmi.vreg<64×f32>, !pto.ptr, !pto.vmi.mask<64> + + // Group (strided) store: 8 rows × 8 elements, stride 64 (no mask) + pto.vmi.vstore %tile, %ub_out[%off], %stride {group = 8} + : !pto.vmi.vreg<64×f32>, !pto.ptr, index + + // Block-strided store: block_stride = 8 (dynamic i16 operand + mask) + pto.vmi.vstore %v, %ub_out[%off], %c8_i16, %mask + : !pto.vmi.vreg<64×f32>, !pto.ptr, i16, !pto.vmi.mask<64> + ``` diff --git a/docs/isa/vmi-isa/02-index-gen.md b/docs/isa/vmi-isa/02-index-gen.md new file mode 100644 index 0000000000..ec4fdafde9 --- /dev/null +++ b/docs/isa/vmi-isa/02-index-gen.md @@ -0,0 +1,58 @@ +# 2. Index-gen + +> **Category:** A. **Mask:** none. +> +> Index materialization. Produces an index vector; the single physical reg +> backing is replicate-read until a Category B/C edge needs the expanded form. + +--- + +## `pto.vmi.vci` + +- **semantics:** Generate a per-lane index/counter vector from a single scalar base such as `[base, base±1, base±2, ...]`, lane `i` gets `base + i` (ASC) or `base - i` (DESC). It is the index source for `vgather`/`vscatter` offsets. + + ```c + for (int i = 0; i < L; i++) + dst[i] = base + (order == "ASC" ? i : -i); + ``` + +- **syntax:** + ```mlir + %result = pto.vmi.vci %base {order = "ASC"} : T -> !pto.vmi.vreg + ``` +- **operands:** + + | Operand | Type | Description | + |---|---|---| + | `base` | integer or float scalar | Starting value | + +- **results:** + + | Result | Type | Description | + |---|---|---| + | `result` | `!pto.vmi.vreg` | Index vector | + +- **attributes:** + + | Attribute | Values | Default | Description | + |---|---|---|---| + | `order` | `"ASC"`, `"DESC"` | `"ASC"` | Index generation direction | + +- **lowering to `pto.mi`:** + ``` + 1 × pto.vci {ASC/DESC} per chunk + ``` + `#mi = 1/chunk`, `dep = 1`. + +- **datatypes:** `i8`/`i16`/`i32`, `f16`, `f32`; the result element type also + fixes `L` (`i32`/`f32` -> 64, `i16`/`f16` -> 128, `i8` -> 256). + +- **example:** + ```mlir + // Ascending i32 indices for a gather base + %idx = pto.vmi.vci %c0 {order = "ASC"} : i32 -> !pto.vmi.vreg<64×i32> + // Descending f32 ramp + %ramp = pto.vmi.vci %c10 {order = "DESC"} : f32 -> !pto.vmi.vreg<64×f32> + %idx = pto.vmi.vci %base {order = "ASC"} : i32 -> !pto.vmi.vreg<64×i32> + // → pto.as: pto.vci {order="ASC"}, one op per physical chunk + ``` diff --git a/docs/isa/vmi-isa/03-eltwise-compute.md b/docs/isa/vmi-isa/03-eltwise-compute.md new file mode 100644 index 0000000000..d0dbd7ad33 --- /dev/null +++ b/docs/isa/vmi-isa/03-eltwise-compute.md @@ -0,0 +1,563 @@ +# 3. Eltwise Compute + +> **Category:** A (layout-passthrough). **Mask:** `Pg` (optional governing predicate, except `vselr` which has none). +> +> Pure per-lane ops. Layout passes through unchanged. An operand whose +> cardinality along an axis is 1 becomes a broadcast (replicate-read, never +> expanded to `K` copies). Under the `K ≤ 4` core profile these fan out as +> fully-unrolled straight-line code. + +--- + +## 3.1 Binary Arithmetic + +### `pto.vmi.vadd` / `pto.vmi.vsub` / `pto.vmi.vmul` + +- **semantics:** Unified fp/int elementwise add / subtract / multiply. + + ```c + for (int i = 0; i < L; i++) + dst[i] = mask[i] ? lhs[i] + rhs[i] : (pmode_merge ? dst_old[i] : 0); + ``` + +- **syntax:** + ```mlir + %r = pto.vmi.vadd %lhs, %rhs, %mask {pmode = "zero"} : !pto.vmi.vreg, !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg + ``` +- **operands:** + + | Operand | Type | Description | + |---|---|---| + | `lhs` | `!pto.vmi.vreg` | First operand | + | `rhs` | `!pto.vmi.vreg` | Second operand | + | `mask` | `!pto.vmi.mask` (variadic) | Governing predicate (0 or 1) | + +- **results:** + + | Result | Type | Description | + |---|---|---| + | `result` | `!pto.vmi.vreg` | Elementwise result | + +- **attributes:** + + | Attribute | Values | Default | Description | + |---|---|---|---| + | `pmode` | `"zero"`, `"merge"` | `"zero"` | Inactive-lane behavior | + +- **datatypes:** `i8`–`i32`, `f16`, `bf16`, `f32` +- **lowering to `pto.mi`:** + ``` + K × pto.vadd / pto.vsub / pto.vmul (+ mask per reg, ppack/punpack if needed) + ``` + `#mi = K`, `dep = 1`, util = 100%. + +- **example:** + ```mlir + // fp32 add with deinterleaved layout + %sum = pto.vmi.vadd %a, %b + : !pto.vmi.vreg<128×f32, #pto.vmi.layout>, + !pto.vmi.vreg<128×f32, #pto.vmi.layout> + -> !pto.vmi.vreg<128×f32, #pto.vmi.layout> + // → pto.as: 2 × pto.vadd (EVEN/ODD), each with create_mask all-active mask + + // Masked add with merge mode + %s = pto.vmi.vadd %a, %b, %mask {pmode = "merge"} + : !pto.vmi.vreg<64×f32>, !pto.vmi.vreg<64×f32>, !pto.vmi.mask<64> -> !pto.vmi.vreg<64×f32> + ``` + +### `pto.vmi.vdiv` + +- **semantics:** Elementwise floating-point divide. + + ```c + for (int i = 0; i < L; i++) + dst[i] = mask[i] ? lhs[i] / rhs[i] : (pmode_merge ? dst_old[i] : 0); + ``` + +- **syntax:** + ```mlir + %r = pto.vmi.vdiv %lhs, %rhs, %mask : !pto.vmi.vreg, !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg + ``` +- **datatypes:** `f16`, `f32` only +- **lowering to `pto.mi`:** + ``` + K × pto.vdiv + ``` + `#mi = K`, `dep = 1`. + +### `pto.vmi.vmax` / `pto.vmi.vmin` + +- **semantics:** Elementwise maximum / minimum (unified fp/int). + + ```c + for (int i = 0; i < L; i++) + dst[i] = mask[i] ? max(lhs[i], rhs[i]) : (pmode_merge ? dst_old[i] : 0); + ``` + +- **syntax:** + ```mlir + %r = pto.vmi.vmax %lhs, %rhs, %mask : !pto.vmi.vreg, !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg + ``` +- **datatypes:** `i8`–`i32`, `f16`, `bf16`, `f32` +- **lowering to `pto.mi`:** + ``` + K × pto.vmax / pto.vmin + ``` + `#mi = K`, `dep = 1`. + +--- + +## 3.2 Unary Arithmetic & Activation + +### `pto.vmi.vabs` + +- **semantics:** Elementwise absolute value (unified fp/int). + + ```c + for (int i = 0; i < L; i++) + dst[i] = mask[i] ? abs(src[i]) : (pmode_merge ? dst_old[i] : 0); + ``` + +- **syntax:** + ```mlir + %r = pto.vmi.vabs %src, %mask {pmode = "zero"} : !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg + ``` +- **datatypes:** `i8`–`i32`, `f16`, `bf16`, `f32` +- **lowering to `pto.mi`:** + ``` + K × pto.vabs + ``` + `#mi = K`, `dep = 1`. + +### `pto.vmi.vneg` + +- **semantics:** Elementwise negate: `0 - x`. + + ```c + for (int i = 0; i < L; i++) + dst[i] = mask[i] ? -src[i] : (pmode_merge ? dst_old[i] : 0); + ``` + +- **syntax:** + ```mlir + %r = pto.vmi.vneg %src, %mask : !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg + ``` +- **datatypes:** `i8`–`i32`, `f16`, `bf16`, `f32` +- **lowering to `pto.mi`:** + ``` + K × pto.vneg (fp) or K × (vsub 0, src) (int) + ``` + `#mi = K`, `dep = 1`. + +### `pto.vmi.vrelu` + +- **semantics:** Elementwise ReLU: `max(0, x)`. + + ```c + for (int i = 0; i < L; i++) + dst[i] = mask[i] ? max(0, src[i]) : (pmode_merge ? dst_old[i] : 0); + ``` + +- **syntax:** + ```mlir + %r = pto.vmi.vrelu %src, %mask : !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg + ``` +- **datatypes:** `i8`–`i32`, `f16`, `bf16`, `f32` +- **lowering to `pto.mi`:** + ``` + K × pto.vrelu + ``` + `#mi = K`, `dep = 1`. + +### `pto.vmi.vexp` / `pto.vmi.vln` / `pto.vmi.vsqrt` + +- **semantics:** Elementwise transcendental: exponential, natural logarithm, square root. + + ```c + for (int i = 0; i < L; i++) + dst[i] = mask[i] ? exp(src[i]) : (pmode_merge ? dst_old[i] : 0); // vexp + for (int i = 0; i < L; i++) + dst[i] = mask[i] ? ln(src[i]) : (pmode_merge ? dst_old[i] : 0); // vln + for (int i = 0; i < L; i++) + dst[i] = mask[i] ? sqrt(src[i]) : (pmode_merge ? dst_old[i] : 0); // vsqrt + ``` + +- **syntax:** + ```mlir + %r = pto.vmi.vexp %src, %mask : !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg + ``` +- **datatypes:** `f16`, `f32` only +- **lowering to `pto.mi`:** + ``` + K × pto.vexp / pto.vln / pto.vsqrt + ``` + `#mi = K`, `dep = 1`. + +--- + +## 3.3 Bitwise Ops + +> **Mask-operand support (planned):** `vand` / `vor` / `vxor` / `vnot` will be +> extended to accept **mask** operands in addition to vector registers. When +> the operands are masks, the op performs a per-lane **predicate boolean** +> operation (AND / OR / XOR / NOT) on the mask lanes and produces a mask +> result, rather than an elementwise data bitwise op on a vreg. This reuses the +> same op names for both vreg-bitwise and mask-boolean forms; the operand type +> selects the mode. There is no separate predicate-logic op (e.g. `pand`/ +> `por`/`pnot`); mask boolean logic is expressed through these ops. + +### `pto.vmi.vand` / `pto.vmi.vor` / `pto.vmi.vxor` + +- **semantics:** Elementwise bitwise AND / OR / XOR. Operands and result are + vregs by default; will also support mask-typed operands, performing a per-lane + predicate boolean op and yielding a mask (the data operands themselves are + masks, distinct from the governing `mask`). + + ```c + for (int i = 0; i < L; i++) + dst[i] = mask[i] ? (lhs[i] & rhs[i]) : (pmode_merge ? dst_old[i] : 0); + ``` + +- **syntax:** + ```mlir + %r = pto.vmi.vand %lhs, %rhs, %mask : !pto.vmi.vreg, !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg + ``` +- **datatypes:** `i8`–`i32` (integer bitwise) +- **lowering to `pto.mi`:** + ``` + K × pto.vand / pto.vor / pto.vxor + ``` + `#mi = K`, `dep = 1`. + +### `pto.vmi.vnot` + +- **semantics:** Elementwise bitwise NOT. Operand and result are vregs by + default; will also support a mask-typed operand, performing a per-lane predicate + complement and yielding a mask (the data operand itself is a mask, distinct + from the governing `mask`). + + ```c + for (int i = 0; i < L; i++) + dst[i] = mask[i] ? ~src[i] : (pmode_merge ? dst_old[i] : 0); + ``` + +- **syntax:** + ```mlir + %r = pto.vmi.vnot %src, %mask : !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg + ``` +- **datatypes:** `i8`–`i32` +- **lowering to `pto.mi`:** + ``` + K × pto.vnot + ``` + `#mi = K`, `dep = 1`. + +--- + +## 3.4 Shift Ops + +### `pto.vmi.vshl` / `pto.vmi.vshr` + +- **semantics:** Elementwise left shift (`vshl`) or signedness-aware right + shift (`vshr`). The shift count is per-lane from `rhs`. `vshr` performs a + logical right shift for explicit unsigned element types and an arithmetic + right shift for signed or signless element types. + + ```c + for (int i = 0; i < L; i++) + dst[i] = mask[i] ? (lhs[i] << rhs[i]) : (pmode_merge ? dst_old[i] : 0); // vshl + for (int i = 0; i < L; i++) + dst[i] = mask[i] ? (lhs[i] >> rhs[i]) : (pmode_merge ? dst_old[i] : 0); // vshr (type-directed) + ``` + +- **syntax:** + ```mlir + %r = pto.vmi.vshl %lhs, %rhs, %mask : !pto.vmi.vreg, !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg + ``` +- **datatypes:** `i8`–`i32` +- **lowering to `pto.mi`:** + ``` + K × pto.vshl / pto.vshr + ``` + `#mi = K`, `dep = 1`. + +--- + +## 3.5 Vec-Scalar Ops + +Vec-scalar ops broadcast a scalar to all lanes (R6 implicit broadcast). The +scalar type must match the vector element type. + +### `pto.vmi.vadds` / `pto.vmi.vmuls` / `pto.vmi.vmaxs` / `pto.vmi.vmins` + +- **semantics:** Elementwise vector-scalar add / multiply / max / min. + + ```c + for (int i = 0; i < L; i++) + dst[i] = mask[i] ? src[i] + scalar : (pmode_merge ? dst_old[i] : 0); + ``` + +- **syntax:** + ```mlir + %r = pto.vmi.vadds %src, %scalar, %mask {pmode = "merge"} : !pto.vmi.vreg, T, !pto.vmi.mask -> !pto.vmi.vreg + ``` +- **operands:** + + | Operand | Type | Description | + |---|---|---| + | `src` | `!pto.vmi.vreg` | Vector operand | + | `scalar` | `T` | Scalar (implicitly broadcast to all lanes) | + | `mask` | `!pto.vmi.mask` | Governing predicate | + +- **results:** + + | Result | Type | Description | + |---|---|---| + | `result` | `!pto.vmi.vreg` | Elementwise result | + +- **datatypes:** `i8`–`i32`, `f16`, `bf16`, `f32` +- **lowering to `pto.mi`:** + ``` + K × pto.vadds / pto.vmuls / pto.vmaxs / pto.vmins + ``` + `#mi = K`, `dep = 1`. No extra reg for scalar. + +- **example:** + ```mlir + %scaled = pto.vmi.vmuls %x, %scale, %mask + : !pto.vmi.vreg<64×f32>, f32, !pto.vmi.mask<64> -> !pto.vmi.vreg<64×f32> + ``` + +### `pto.vmi.vshls` / `pto.vmi.vshrs` + +- **semantics:** Elementwise vector-scalar shift. + + ```c + for (int i = 0; i < L; i++) + dst[i] = mask[i] ? (src[i] << scalar) : (pmode_merge ? dst_old[i] : 0); // vshls + for (int i = 0; i < L; i++) + dst[i] = mask[i] ? (src[i] >> scalar) : (pmode_merge ? dst_old[i] : 0); // vshrs + ``` + +- **syntax:** + ```mlir + %r = pto.vmi.vshls %src, %scalar, %mask : !pto.vmi.vreg, T, !pto.vmi.mask -> !pto.vmi.vreg + ``` +- **datatypes:** `i8`–`i32` +- **lowering to `pto.mi`:** + ``` + K × pto.vshls / pto.vshrs + ``` + `#mi = K`, `dep = 1`. + +--- + +## 3.6 Compare & Select + +### `pto.vmi.vcmp` + +- **semantics:** Elementwise compare → predicate mask. The `seed` mask is the + governing predicate `Pg`: where `seed[i] = 0` the result lane is 0 (zeroing); + where `seed[i] = 1` the comparison is evaluated. + + ```c + for (int i = 0; i < L; i++) + dst[i] = seed[i] ? cmp(lhs[i], rhs[i]) : 0; + ``` + +- **syntax:** + ```mlir + %r = pto.vmi.vcmp %lhs, %rhs, %seed {cmp = "lt"} : !pto.vmi.vreg, !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.mask + ``` +- **operands:** + + | Operand | Type | Description | + |---|---|---| + | `lhs` | `!pto.vmi.vreg` | First operand | + | `rhs` | `!pto.vmi.vreg` | Second operand | + | `seed` | `!pto.vmi.mask` | Governing predicate (required) | + +- **results:** + + | Result | Type | Description | + |---|---|---| + | `result` | `!pto.vmi.mask` | Predicate mask (same L, granularity derived from T) | + +- **attributes:** + + | Attribute | Values | Default | Description | + |---|---|---|---| + | `cmp` | `eq`, `ne`, `lt`, `le`, `gt`, `ge` (unordered fp+int) | *(required)* | Comparison mode | + | | `oeq`, `one`, `olt`, `ole`, `ogt`, `oge` (ordered fp) | | FP ordered forms | + | | `slt`, `sle`, `sgt`, `sge` (signed int) | | Signed integer forms | + | `pmode` | `"zero"`, `"merge"` | `"zero"` | Inactive-lane behavior | + +- **datatypes:** `i8`–`i32`, `f16`, `bf16`, `f32` +- **lowering to `pto.mi`:** + ``` + K × pto.vcmp {cmp_mode} + ``` + `#mi = K`, `dep = 1`. +1 preg per live mask result. + +- **example:** + ```mlir + // f32 less-than compare over deinterleaved layout + %lt = pto.vmi.vcmp %a, %b, %seed {cmp = "lt"} + : !pto.vmi.vreg<128×f32, #pto.vmi.layout>, + !pto.vmi.vreg<128×f32, #pto.vmi.layout>, + !pto.vmi.mask<128×b32, #pto.vmi.layout> + -> !pto.vmi.mask<128×b32, #pto.vmi.layout> + + // i32 signed greater-than-or-equal over deinterleaved layout + %ge = pto.vmi.vcmp %a, %b, %seed {cmp = "sge"} + : !pto.vmi.vreg<128×i32>, !pto.vmi.vreg<128×i32>, !pto.vmi.mask<128×b32> + -> !pto.vmi.mask<128×b32> + + // bf16 contiguous equality compare (K=1) + %eq = pto.vmi.vcmp %a, %b, %seed {cmp = "eq"} + : !pto.vmi.vreg<128×bf16>, !pto.vmi.vreg<128×bf16>, !pto.vmi.mask<128×b16> + -> !pto.vmi.mask<128×b16> + ``` + +### `pto.vmi.vcmps` + +- **semantics:** Elementwise vector-scalar compare → predicate mask. + + ```c + for (int i = 0; i < L; i++) + dst[i] = seed[i] ? cmp(src[i], scalar) : 0; + ``` + +- **syntax:** + ```mlir + %r = pto.vmi.vcmps %src, %scalar, %seed {cmp = "ge"} : !pto.vmi.vreg, T, !pto.vmi.mask -> !pto.vmi.mask + ``` +- **operands:** + + | Operand | Type | Description | + |---|---|---| + | `src` | `!pto.vmi.vreg` | Vector operand | + | `scalar` | `T` | Scalar to compare against | + | `seed` | `!pto.vmi.mask` | Governing predicate (required) | + +- **results:** + + | Result | Type | Description | + |---|---|---| + | `result` | `!pto.vmi.mask` | Predicate mask | + +- **attributes:** Same `cmp` / `pmode` as `vcmp`. +- **datatypes:** `i8`–`i32`, `f16`, `bf16`, `f32` +- **lowering to `pto.mi`:** + ``` + K × pto.vcmps {cmp_mode} + ``` + `#mi = K`, `dep = 1`. + +- **example:** + ```mlir + %ges = pto.vmi.vcmps %a, %c0, %seed {cmp = "ge"} + : !pto.vmi.vreg<64×f32>, f32, !pto.vmi.mask<64> -> !pto.vmi.mask<64> + ``` + +### `pto.vmi.vsel` + +- **semantics:** Per-lane selection driven by a predicate mask. + + ```c + for (int i = 0; i < L; i++) + dst[i] = mask[i] ? true_val[i] : false_val[i]; + ``` + +- **syntax:** + ```mlir + %r = pto.vmi.vsel %mask, %true_val, %false_val {pmode = "zero"} : !pto.vmi.mask, !pto.vmi.vreg, !pto.vmi.vreg -> !pto.vmi.vreg + ``` +- **operands:** + + | Operand | Type | Description | + |---|---|---| + | `mask` | `!pto.vmi.mask` | Selector predicate (required) | + | `true_val` | `!pto.vmi.vreg` | Value when mask[i] = 1 | + | `false_val` | `!pto.vmi.vreg` | Value when mask[i] = 0 | + +- **results:** + + | Result | Type | Description | + |---|---|---| + | `result` | `!pto.vmi.vreg` | Selected result | + +- **attributes:** + + | Attribute | Values | Default | Description | + |---|---|---|---| + | `pmode` | `"zero"`, `"merge"` | `"zero"` | Result handling when selector inactive: `"merge"` retains `false_value` lanes | + +- **datatypes:** `i8`–`i32`, `f16`, `bf16`, `f32` +- **lowering to `pto.mi`:** + ``` + K × pto.vsel + ``` + `#mi = K`, `dep = 1`. + +- **example:** + ```mlir + %out = pto.vmi.vsel %mask, %x, %y {pmode = "zero"} + : !pto.vmi.mask<256×b16>, !pto.vmi.vreg<256×ui16>, !pto.vmi.vreg<256×ui16> + -> !pto.vmi.vreg<256×ui16> + ``` + +### `pto.vmi.vselr` + +- **semantics:** Dynamic lane permutation: `result[i] = source[index[i]]`. + + ```c + for (int i = 0; i < L; i++) + dst[i] = src[index[i]]; + ``` + +- **syntax:** + ```mlir + %r = pto.vmi.vselr %source, %index : !pto.vmi.vreg, !pto.vmi.vreg -> !pto.vmi.vreg + ``` +- **operands:** + + | Operand | Type | Description | + |---|---|---| + | `source` | `!pto.vmi.vreg` | Source vector to permute from | + | `index` | `!pto.vmi.vreg` | Per-lane source lane index | + +- **results:** + + | Result | Type | Description | + |---|---|---| + | `result` | `!pto.vmi.vreg` | Permuted result | + +- **datatypes:** `i8`–`i32`, `f16`, `bf16`, `f32` +- **lowering to `pto.mi`:** + ``` + K × pto.vselr (+ index reg setup) + ``` + `#mi = K`, `dep = 1` (+1 for index setup). +1 index vreg. + +- **notes:** + - This is the permute/gather class — it is the register-resident realization + of a grouped broadcast. + - `vselr` takes no mask; the index vector encodes the permutation directly. + - Not A5-native `vselrv2` (that form is not available on A5). + +- **example:** + ```mlir + %r = pto.vmi.vselr %src, %idx + : !pto.vmi.vreg<64×f16>, !pto.vmi.vreg<4×i16> -> !pto.vmi.vreg<4×f16> + ``` + +--- + +## 3.7 Carry / Borrow Ops (Not Provided) + +Vector carry/borrow arithmetic (e.g. multi-word add-with-carry across +lanes) is **not provided** on the current surface. It will be added directly +as `i64` element-wise ops once the `i64` support plan is finalized and the +hardware path is confirmed. Until then, widening to `i64` scalar emulation +or fusing at the `pto.mi` layer is the workaround. diff --git a/docs/isa/vmi-isa/04-broadcast.md b/docs/isa/vmi-isa/04-broadcast.md new file mode 100644 index 0000000000..53e51f07dc --- /dev/null +++ b/docs/isa/vmi-isa/04-broadcast.md @@ -0,0 +1,93 @@ +# 4. Broadcast + +> **Category:** A (ungrouped scalar→vector), B (grouped `{group}`). +> **Mask:** none. +> +> `vbrc` is the logical scalar→vector / compact→full broadcast. The ungrouped +> form (single scalar fanned over `L` lanes) is cheap (`vdup`); the grouped form +> (per-group scalar fan-back) has no single native instruction and is a +> cost-model decision. + +--- + +## `pto.vmi.vbrc` + +- **semantics:** Broadcast a scalar or group-slot compact value across lanes. + + **Ungrouped:** One value replicated to all `L` lanes. + ```c + for (int i = 0; i < L; i++) + dst[i] = src[0]; + ``` + + **Grouped (`{group = C}`):** Each of the `C` compact scalar slots is + fanned back across `L/C` lanes. + ```c + int gs = L / C; // lanes per group + for (int g = 0; g < C; g++) + for (int i = 0; i < gs; i++) + dst[g * gs + i] = src[g]; + ``` + +- **syntax:** + ```mlir + // Ungrouped: scalar → full vector + %r = pto.vmi.vbrc %scalar : f32 -> !pto.vmi.vreg<64×f32> + + // Ungrouped: 1-lane vreg → full vector + %r = pto.vmi.vbrc %val : !pto.vmi.vreg<1×f32> -> !pto.vmi.vreg<256×f32> + + // Grouped: compact group-slot → dense vector + %r = pto.vmi.vbrc %source {group = 128} : !pto.vmi.vreg<128×f32> -> !pto.vmi.vreg<1024×f32> + ``` +- **operands:** + + | Operand | Type | Description | + |---|---|---| + | `value` | `T` (scalar) or `!pto.vmi.vreg` | Broadcast source | + +- **results:** + + | Result | Type | Description | + |---|---|---| + | `result` | `!pto.vmi.vreg` | Broadcast result | + +- **attributes:** + + | Attribute | Values | Default | Description | + |---|---|---|---| + | `group` | positive integer | *(none — ungrouped)* | Number of group slots; must equal `input.L` for group mode | + +- **datatypes:** `i8`–`i32`, `f16`, `bf16`, `f32` +- **lowering to `pto.mi`:** + + | Form | Physical lowering | `#mi` | `dep` | + |---|---|---|---| + | Ungrouped (scalar) | `1 × pto.vdup` (register-resident), or `vsts`+`vlds BRC_*` (UB roundtrip) | `1` | `1` | + | Ungrouped (1-lane vreg) | `1 × pto.vdup {position="LOWEST"}` per physical reg | `K` | `1` | + | Grouped (`{group}`) | **Cost-model decision**: UB roundtrip (`vsts` partials + `vlds BRC_BLK`) **or** `vselr` gather **or** masked recompute | varies | 2–3 | + +- **examples:** + ```mlir + // Ungrouped: scalar → full vector + %bc = pto.vmi.vbrc %maxe : f32 -> !pto.vmi.vreg<64×f32> + // → pto.as: pto.vdup %maxe (one op, register-resident) + + // Ungrouped: 1-lane vreg → full vector (rank-0 broadcast) + %bc = pto.vmi.vbrc %scalar : !pto.vmi.vreg<1×f32> -> !pto.vmi.vreg<256×f32> + // → pto.as: 4 × pto.vdup {position="LOWEST"} (K=4) + + // Grouped: 128 compact slots → 1024-lane dense vector + %bc = pto.vmi.vbrc %source {group = 128} + : !pto.vmi.vreg<128×f32> -> !pto.vmi.vreg<1024×f32> + // → pto.as: 16 × pto.vselr (vselr gather realization) + ``` + +- **notes:** + - Fused `reduce→broadcast` (`vcadd`+`vbrc`) is the recognized fusion pattern: + `pto.as` emits them back-to-back and keeps the result as a broadcast axis + rather than materializing `K` copies. + - Prefer `vdup` over a UB `BRC` reload for a single scalar. + - Grouped broadcast has **no single native `pto.mi` op** — `pto.as` picks + UB roundtrip (default, `vsts` partials + `vlds BRC_BLK`), `vselr` gather + (when group count and K are tiny), or masked recompute (very small groups). diff --git a/docs/isa/vmi-isa/05-reduce.md b/docs/isa/vmi-isa/05-reduce.md new file mode 100644 index 0000000000..fd6e78d5d7 --- /dev/null +++ b/docs/isa/vmi-isa/05-reduce.md @@ -0,0 +1,126 @@ +# 5. Reduce + +> **Category:** B (VLane-aligned), C (unaligned sub-VLane). +> **Mask:** `Pg req` (governing mask is a required operand). +> +> Reduction ops collapse lanes into compact scalars, governed by a mask. +> `{group=C}` controls the number of sub-groups. Inactive lane behavior: +> `vcadd` treats inactive as 0; `vcmax`/`vcmin` treat inactive as `-∞`/`+∞` +> (fp) or type min/max (int). + +--- + +## `pto.vmi.vcadd` + +- **semantics:** Masked add-reduction. When `{group=C}` is absent, reduces all + `L` active lanes to a single scalar (`V<1×T>`). + + ```c + // Without group: full reduction to scalar + T sum = 0; + for (int i = 0; i < L; i++) + if (mask[i]) sum += src[i]; + dst[0] = sum; + + // With {group=C}: per-group reduction + int gs = L / C; // lanes per group + for (int g = 0; g < C; g++) { + T sum = 0; + for (int i = 0; i < gs; i++) + if (mask[g*gs + i]) sum += src[g*gs + i]; + dst[g] = sum; + } + ``` + +- **syntax:** + ```mlir + %r = pto.vmi.vcadd %src, %mask {group = C, reassoc} : !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg + ``` +- **operands:** + + | Operand | Type | Description | + |---|---|---| + | `src` | `!pto.vmi.vreg` | Source vector | + | `mask` | `!pto.vmi.mask` | Governing predicate (required) | + +- **results:** + + | Result | Type | Description | + |---|---|---| + | `result` | `!pto.vmi.vreg` | Compact scalar vector (`C = 1` if no group) | + +- **attributes:** + + | Attribute | Values | Default | Description | + |---|---|---|---| + | `group` | `1`, `2`, `4`, `8` | `1` (full reduce) | Number of sub-groups | + | `reassoc` | *(unit attr)* | *(absent)* | Permit reassociation (**required** for fp sources) | + | `pmode` | `"zero"`, `"merge"` | `"zero"` | Inactive-result behavior | + +- **datatypes:** `i8`–`i32`, `f16`, `f32` +- **lowering to `pto.mi`:** + + | Group / W | Category | Physical lowering | `#mi` | `dep` | + |---|---|---|---|---| + | No group (`C=1`), `K=1` | B | `1 × pto.vcadd` | `1` | `1` | + | No group, `K>1` (fold) | B | `(K-1) × vadd` + `1 × vcadd` | `K` | `K` | + | No group, `K>1` (partial) | B | `K × vcadd` + combine | `K` | `1+⌈log₂K⌉` | + | `group=8` (W=32B, VLane-aligned) | B | `K × pto.vcgadd` | `K` | `1` | + | `group=2/4` (W=64B/128B aligned) | B | `(k-1) × vadd` fold + `vcgadd` | `K+k-1` | `k` | + +- **example:** + ```mlir + // Full sum reduction (to scalar) + %sum = pto.vmi.vcadd %x, %mask {reassoc} + : !pto.vmi.vreg<64×f32>, !pto.vmi.mask<64> -> !pto.vmi.vreg<1×f32> + + // Grouped: 256-lane → 8 groups of 32, each VLane-aligned (W=32B) + %sums = pto.vmi.vcadd %x, %mask {group = 8} + : !pto.vmi.vreg<256×f16>, !pto.vmi.mask<256> -> !pto.vmi.vreg<8×f16> + ``` + +--- + +## `pto.vmi.vcmax` / `pto.vmi.vcmin` + +- **semantics:** Masked max/min reduction. + + ```c + // vcmax: inactive lanes treated as -∞ + T best = -INF; + for (int i = 0; i < L; i++) + if (mask[i]) best = max(best, src[i]); + dst[0] = best; + + // vcmin: inactive lanes treated as +∞ + T best = +INF; + for (int i = 0; i < L; i++) + if (mask[i]) best = min(best, src[i]); + dst[0] = best; + ``` + +- **syntax:** + ```mlir + %r = pto.vmi.vcmax %src, %mask {group = C} : !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg + ``` +- **operands:** Same as `vcadd` (without `reassoc`). +- **results:** Same as `vcadd`. +- **attributes:** `group`, `pmode` (same as `vcadd`, no `reassoc`). +- **datatypes:** `i16`–`i32`, `f16`, `f32` +- **lowering to `pto.mi`:** + + | Group / W | Physical lowering | + |---|---| + | No group, fold | `(K-1) × vmax` + `1 × vcmax` | + | VLane-aligned | `K × pto.vcgmax` / `K × pto.vcgmin` | + +- **example:** + ```mlir + // Full max reduction + %mx = pto.vmi.vcmax %x, %mask + : !pto.vmi.vreg<64×f32>, !pto.vmi.mask<64> -> !pto.vmi.vreg<1×f32> + + // Grouped: 8-sub-group max (MX block-scale exponent pattern) + %maxe = pto.vmi.vcmax %exp, %mask {group = 8} + : !pto.vmi.vreg<256×ui16>, !pto.vmi.mask<256> -> !pto.vmi.vreg<8×ui16> + ``` diff --git a/docs/isa/vmi-isa/06-convert.md b/docs/isa/vmi-isa/06-convert.md new file mode 100644 index 0000000000..f7a4a2830d --- /dev/null +++ b/docs/isa/vmi-isa/06-convert.md @@ -0,0 +1,139 @@ +# 6. Convert + +> **Category:** B (`vcvt`), A (`vinterpret_cast`). +> **Mask:** `Pg` (`vcvt`), none (`vinterpret_cast`). +> +> One logical `vcvt` whose target dtype IS the layout. `pto.as` expands it into +> the dtype-specific cast chain + part/width staging + matching store +> distribution. The author never spells `EVEN`/`ODD`, `P0`–`P3`, `PK`/`UNPK`, +> or `VL/2` addresses. + +--- + +## `pto.vmi.vcvt` + +- **semantics:** Unified elementwise type conversion. The conversion direction + is derived from source and destination element types: + + | Direction | Condition | Replaces | + |---|---|---| + | fp → fp, `|dst| > |src|` | Floating-point widening | `extf` | + | fp → fp, `|dst| < |src|` | Floating-point narrowing | `truncf` | + | fp → int | Float to signed integer | `fptosi` | + | int → fp | Signed integer to float | `sitofp` | + | int -> int, `|dst| > |src|` | Integer extension (sign from source element type) | `extsi` / `extui` | + | int → int, `|dst| < |src|` | Saturating integer truncation | `trunci` | + +- **syntax:** + ```mlir + %r = pto.vmi.vcvt %src {rounding = "H"} : !pto.vmi.vreg -> !pto.vmi.vreg + ``` +- **operands:** + + | Operand | Type | Description | + |---|---|---| + | `src` | `!pto.vmi.vreg` | Source vector | + +- **results:** + + | Result | Type | Description | + |---|---|---| + | `result` | `!pto.vmi.vreg` | Converted vector (same `L`, different `T`) | + +- **attributes:** + + | Attribute | Values | Valid for | Description | + |---|---|---|---| + | `rounding` | `"A"` (away-from-zero), `"H"` (half-up) | fp narrowing | Rounding mode | + | `saturate` | `"SAT"` | any narrowing | Saturating on overflow | + | `pmode` | `"zero"`, `"merge"` | all | Inactive-lane behavior | + +- **datatypes:** Source and destination from `{f32, f16, bf16, fp8_e4m3, fp8_e5m2, i32, i16, i8, ui32, ui16, ui8}` +- **lowering to `pto.mi`:** + + | Conversion | Physical lowering | `#mi` | `dep` | + |---|---|---|---| + | 16↔32 (radix-2) | `2K × vcvt EVEN/ODD` + predicate `ppack`/`punpack` companion | `2K` | `2` | + | 8↔32 (radix-4) | widen: `UNPK_B8` + `vintlv` + `vcvt P0` + `punpack`; narrow: `PK4_B32` store (or `vselr` gather) + `ppack` | `2–3` | `2–3` | + | f32→fp8 quant | `1 cast` + `PK4_B32` | `K` | `1` | + | f32→int8 quant | 3-stage cast + `PK4_B32` | `~3K` | `3` | + | int↔int (same width) | `K × vtrc` or `K × vcvt` | `K` | `1` | + +- **example:** + ```mlir + // fp16 → fp32 widen (radix-2, produces parity EVEN/ODD) + %w = pto.vmi.vcvt %a + : !pto.vmi.vreg<128×f16, #pto.vmi.layout> + -> !pto.vmi.vreg<128×f32, #pto.vmi.layout> + // → pto.as: 2 × pto.vcvt EVEN/ODD + ppack (parity companion) + + // fp32 → fp16 narrow with half-up rounding + %n = pto.vmi.vcvt %y {rounding = "H"} + : !pto.vmi.vreg<64×f32> -> !pto.vmi.vreg<64×f16> + + // ui8 -> i16 unsigned extension + %z = pto.vmi.vcvt %a + : !pto.vmi.vreg<256×ui8> -> !pto.vmi.vreg<256×i16> + + // f32 → fp8 quantized narrow + %q = pto.vmi.vcvt %s + : !pto.vmi.vreg<64×f32> -> !pto.vmi.vreg<64×fp8_e4m3> + ``` + +- **notes:** + - `vcvt` **does not change lane count** — `src.L == dst.L` always. The + physical register count `K` changes because `bitwidth(T)` changes. + - Integer signedness is determined by the **element type**. + - The `part`/`parity`/`width` axes are lowering-only; the user never writes + `EVEN`/`ODD`/`P0..P3`. + - Radix-4 (8↔32) is **not** a stacked predicate chain and **not** a UB + roundtrip; the 1↔4 lane spread rides data load/store distribution + (`UNPK_B*`/`PK4_B32`) or a `vselr` byte-gather. + +--- + +## `pto.vmi.vinterpret_cast` + +- **semantics:** Bitwise reinterpretation of a vector register — same bits, + different element type. No data movement, no layout change. + + ```c + // Same bits, reinterpreted element-by-element + memcpy(&dst, &src, L * sizeof(T_src)); + ``` + +- **syntax:** + ```mlir + %r = pto.vmi.vinterpret_cast %src : !pto.vmi.vreg -> !pto.vmi.vreg + ``` +- **operands:** + + | Operand | Type | Description | + |---|---|---| + | `src` | `!pto.vmi.vreg` | Source vector | + +- **results:** + + | Result | Type | Description | + |---|---|---| + | `result` | `!pto.vmi.vreg` | Bit-reinterpreted vector | + +- **attributes:** *(none)* +- **datatypes:** Any `T_src`, `T_dst` with `L · bitwidth(T_src) == L · bitwidth(T_dst)` +- **lowering to `pto.mi`:** + ``` + K × pto.vbitcast (or no-op if same physical layout) + ``` + `#mi = 0` or `K`, `dep = 0` or `1`. + +- **notes:** + - **Category A** — layout-transparent, no new axis produced. + - This is **not** `vcvt` — no dtype cast chain, no `part`/`parity`/`width` + axis, no `[pmode]`. + - The user must ensure semantic legality (e.g., `f32` → `i32` bitcast is + valid; `f32` → `f16` is not — use `vcvt` for that). + +- **example:** + ```mlir + %r = pto.vmi.vinterpret_cast %a : !pto.vmi.vreg<64×f32> -> !pto.vmi.vreg<64×i32> + ``` diff --git a/docs/isa/vmi-isa/07-sfu.md b/docs/isa/vmi-isa/07-sfu.md new file mode 100644 index 0000000000..b43fe2df3b --- /dev/null +++ b/docs/isa/vmi-isa/07-sfu.md @@ -0,0 +1,439 @@ +# 7. SFU + +> **Category:** A (fused arithmetic, `vmull`), B (`vchist`, `vdhist`), C (gather/scatter). +> **Mask:** `Pg` on all except sort-like ops. +> +> Special-function / domain-accelerator ops. Mixed categories: `vchist` +> produces a `half` axis (B); `vdhist` yields a plain per-bin count (B); +> gather/scatter are Category C tile/permute ops; fused activation/arithmetic +> ops and pair-result `vmull` are Category A layout-passthrough operations. + +--- + +## 7.1 Fused Arithmetic + +### `pto.vmi.vexpdif` + +- **semantics:** Fused `exp(x − max)` for softmax numerical stability. Single + hardware instruction. + + ```c + for (int i = 0; i < L; i++) + dst[i] = mask[i] ? exp(x[i] - max[i]) : (pmode_merge ? dst_old[i] : 0); + ``` + +- **syntax:** + ```mlir + %e = pto.vmi.vexpdif %x, %max, %mask : !pto.vmi.vreg, !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg + ``` +- **operands:** + + | Operand | Type | Description | + |---|---|---| + | `x` | `!pto.vmi.vreg` | Input (`f16` or `f32`) | + | `max` | `!pto.vmi.vreg` | Subtracted max (always `f32`) | + | `mask` | `!pto.vmi.mask` | Governing predicate | + +- **results:** + + | Result | Type | Description | + |---|---|---| + | `result` | `!pto.vmi.vreg` | `exp(x − max)` (always `f32`) | + +- **attributes:** `pmode` (`"zero"` / `"merge"`) +- **datatypes:** Input `x`: `f16`, `f32`; `max` and result: always `f32` +- **lowering to `pto.mi`:** + ``` + K × pto.vexpdif + ``` + `#mi = K`, `dep = 1`. Fuses `vsub` + `vexp`. + +- **example:** + ```mlir + %e = pto.vmi.vexpdif %x, %max, %mask + : !pto.vmi.vreg<64×f32>, !pto.vmi.vreg<64×f32>, !pto.vmi.mask<64> + -> !pto.vmi.vreg<64×f32> + ``` + +### `pto.vmi.vaxpy` + +- **semantics:** Fused `α·x + y` (scale-add). Single hardware instruction. + + ```c + for (int i = 0; i < L; i++) + dst[i] = mask[i] ? (alpha * x[i] + acc[i]) : (pmode_merge ? dst_old[i] : 0); + ``` + +- **syntax:** + ```mlir + %y = pto.vmi.vaxpy %x, %acc, %alpha, %mask : !pto.vmi.vreg, !pto.vmi.vreg, T, !pto.vmi.mask -> !pto.vmi.vreg + ``` +- **operands:** + + | Operand | Type | Description | + |---|---|---| + | `x` | `!pto.vmi.vreg` | Input vector | + | `acc` | `!pto.vmi.vreg` | Accumulator (`y`) | + | `alpha` | `T` (float scalar) | Scale factor | + | `mask` | `!pto.vmi.mask` | Governing predicate | + +- **results:** + + | Result | Type | Description | + |---|---|---| + | `result` | `!pto.vmi.vreg` | `α·x + acc` | + +- **datatypes:** `f16`, `f32` +- **lowering to `pto.mi`:** + ``` + K × pto.vaxpy + ``` + `#mi = K`, `dep = 1`. Fuses `vmuls` + `vadd`. + +### `pto.vmi.vlrelu` + +- **semantics:** Leaky ReLU: `y = x > 0 ? x : slope × x`. The slope is a + scalar shared across all lanes. + + ```c + for (int i = 0; i < L; i++) + dst[i] = mask[i] ? (src[i] > 0 ? src[i] : slope * src[i]) : (pmode_merge ? dst_old[i] : 0); + ``` + +- **syntax:** + ```mlir + %y = pto.vmi.vlrelu %x, %slope, %mask : !pto.vmi.vreg, T, !pto.vmi.mask -> !pto.vmi.vreg + ``` +- **operands:** + + | Operand | Type | Description | + |---|---|---| + | `x` | `!pto.vmi.vreg` | Input | + | `slope` | `T` (float scalar) | Negative-slope multiplier | + | `mask` | `!pto.vmi.mask` | Governing predicate | + +- **datatypes:** `f16`, `f32` +- **lowering to `pto.mi`:** + ``` + K × pto.vlrelu + ``` + `#mi = K`, `dep = 1`. + +### `pto.vmi.vprelu` + +- **semantics:** Parametric ReLU: `y = max(x, 0) + alpha × min(x, 0)`. The + `alpha` is a per-lane parameter vector (not a shared scalar). + + ```c + for (int i = 0; i < L; i++) + dst[i] = mask[i] ? (max(src[i], 0) + alpha[i] * min(src[i], 0)) : (pmode_merge ? dst_old[i] : 0); + ``` + +- **syntax:** + ```mlir + %y = pto.vmi.vprelu %x, %alpha, %mask : !pto.vmi.vreg, !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg + ``` +- **operands:** + + | Operand | Type | Description | + |---|---|---| + | `x` | `!pto.vmi.vreg` | Input | + | `alpha` | `!pto.vmi.vreg` | Per-lane negative-slope parameter | + | `mask` | `!pto.vmi.mask` | Governing predicate | + +- **datatypes:** `f16`, `f32` +- **lowering to `pto.mi`:** + ``` + K × pto.vprelu + ``` + `#mi = K`, `dep = 1`. + +### `pto.vmi.vmull` + +- **semantics:** Widening 32-bit × 32-bit multiply. The operation returns the + low and high 32-bit halves as two logical vectors of the same type as the + inputs. Inactive lanes in both results are zero. + + ```c + for (int i = 0; i < L; i++) { + uint64_t product = T == i32 + ? (uint64_t)((int64_t)(int32_t)a[i] * (int64_t)(int32_t)b[i]) + : (uint64_t)(uint32_t)a[i] * (uint64_t)(uint32_t)b[i]; + low[i] = mask[i] ? (uint32_t)product : 0; + high[i] = mask[i] ? (uint32_t)(product >> 32) : 0; + } + ``` + +- **syntax:** + ```mlir + %low, %high = pto.vmi.vmull %a, %b, %mask + : !pto.vmi.vreg, !pto.vmi.vreg, !pto.vmi.mask + -> !pto.vmi.vreg, !pto.vmi.vreg + ``` +- **operands:** + + | Operand | Type | Description | + |---|---|---| + | `a` | `!pto.vmi.vreg` | First operand | + | `b` | `!pto.vmi.vreg` | Second operand | + | `mask` | `!pto.vmi.mask` | Governing predicate | + +- **results:** `%low` and `%high`, both `!pto.vmi.vreg` +- **datatypes:** `T` is exactly `i32` or `ui32`; `L` is exactly 64, 128, or + 256. Omitted `pmode` means zeroing; the only explicit legal value is + `pmode = "zero"`. +- **lowering to `pto.mi`:** + ``` + K × pto.vmull (produces hi+lo pair per reg) + ``` + For contiguous layout, `K = L / 64`; therefore 64, 128, and 256 lanes lower + to 1, 2, and 4 operations respectively. + +- **example:** + ```mlir + %low, %high = pto.vmi.vmull %a, %b, %mask + : !pto.vmi.vreg<64×i32>, !pto.vmi.vreg<64×i32>, !pto.vmi.mask<64> + -> !pto.vmi.vreg<64×i32>, !pto.vmi.vreg<64×i32> + ``` + +### `pto.vmi.vmula` + +- **semantics:** Fused multiply-add: `acc = acc + lhs × rhs`. Single hardware + instruction. The accumulator is both an input and output (writes back). + + ```c + for (int i = 0; i < L; i++) + dst[i] = mask[i] ? (acc[i] + lhs[i] * rhs[i]) : (pmode_merge ? dst_old[i] : 0); + ``` + +- **syntax:** + ```mlir + %acc1 = pto.vmi.vmula %acc, %lhs, %rhs, %mask : !pto.vmi.vreg, !pto.vmi.vreg, !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg + ``` +- **operands:** + + | Operand | Type | Description | + |---|---|---| + | `acc` | `!pto.vmi.vreg` | Accumulator (read-modify-write) | + | `lhs` | `!pto.vmi.vreg` | First multiply operand | + | `rhs` | `!pto.vmi.vreg` | Second multiply operand | + | `mask` | `!pto.vmi.mask` | Governing predicate | + +- **datatypes:** `i8`–`i32`, `f16`, `bf16`, `f32` +- **lowering to `pto.mi`:** + ``` + K × pto.vmula + ``` + `#mi = K`, `dep = 1`. Fuses `vmul` + `vadd`. + +- **example:** + ```mlir + %acc1 = pto.vmi.vmula %acc, %a, %b, %mask + : !pto.vmi.vreg<64×f32>, !pto.vmi.vreg<64×f32>, !pto.vmi.vreg<64×f32>, + !pto.vmi.mask<64> -> !pto.vmi.vreg<64×f32> + ``` + +--- + +## 7.2 Histogram + +### `pto.vmi.vchist` + +- **semantics:** **Cumulative histogram** — the existing `chistv2` + semantics. Counts per-bin occurrences over a bin-index vector and produces + a `half`-axis (`Bin_N0`/`Bin_N1`) pair accessible through the result's + width axis. + + ```c + // Hardware chistv2: two halves (Bin_N0, Bin_N1), 256 bins total + uint16_t bins[256] = {0}; + for (int i = 0; i < L; i++) + if (mask[i]) + bins[bin_idx[i]]++; + // dst carries Bin_N0 (bins 0–127) and Bin_N1 (bins 128–255) on a half axis + ``` + +- **syntax:** + ```mlir + %h = pto.vmi.vchist %bin_idx, %mask : !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg + ``` +- **operands:** + + | Operand | Type | Description | + |---|---|---| + | `bin_idx` | `!pto.vmi.vreg` | Per-lane bin index (unsigned 8-bit) | + | `mask` | `!pto.vmi.mask` | Governing predicate | + +- **results:** + + | Result | Type | Description | + |---|---|---| + | `result` | `!pto.vmi.vreg` | Bin counts (half axis: Bin_N0/N1 pair) | + +- **attributes:** + + | Attribute | Values | Default | Description | + |---|---|---|---| + | `pmode` | `"zero"`, `"merge"` | `"zero"` | Inactive-lane behavior | +- **datatypes:** Bin index: `i8`/`ui8`; result count type: typically `i16`/`i32` +- **lowering to `pto.mi`:** + ``` + chistv2 Bin_N0 + Bin_N1 (two-half fanout) + widen/accumulate + ``` + `#mi ≈ 2K`, `dep = 2–3`. INTLV merge on store. + +- **example:** + ```mlir + // Cumulative histogram, half-axis Bin_N0/Bin_N1 + %h = pto.vmi.vchist %bin_idx, %mask + : !pto.vmi.vreg<256×i8>, !pto.vmi.mask<256> -> !pto.vmi.vreg<256×i16> + // → pto.as: Bin_N0 + Bin_N1 fanout → INTLV merge on vstore + ``` + +### `pto.vmi.vdhist` + +- **semantics:** **Distribution histogram** — count per bin over a + value/index vector, yielding a plain per-bin count vector (no `half` + axis). + + ```c + // Plain per-bin distribution count + uint16_t bins[N] = {0}; + for (int i = 0; i < L; i++) + if (mask[i]) + bins[bin_idx[i]]++; + ``` + +- **syntax:** + ```mlir + %d = pto.vmi.vdhist %bin_idx, %mask : !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg + ``` +- **operands:** + + | Operand | Type | Description | + |---|---|---| + | `bin_idx` | `!pto.vmi.vreg` | Per-lane bin index (unsigned 8-bit) | + | `mask` | `!pto.vmi.mask` | Governing predicate | + +- **results:** + + | Result | Type | Description | + |---|---|---| + | `result` | `!pto.vmi.vreg` | Plain per-bin count vector | + +- **attributes:** + + | Attribute | Values | Default | Description | + |---|---|---|---| + | `pmode` | `"zero"`, `"merge"` | `"zero"` | Inactive-lane behavior | +- **datatypes:** Bin index: `i8`/`ui8`; result count type: typically `i16`/`i32` +- **lowering to `pto.mi`:** + ``` + distribution histogram accumulate (no half-axis fanout) + ``` + `#mi ≈ K`, `dep = 2`. + +- **example:** + ```mlir + // Distribution histogram, plain per-bin count + %d = pto.vmi.vdhist %bin_idx, %mask + : !pto.vmi.vreg<256×i8>, !pto.vmi.mask<256> -> !pto.vmi.vreg<256×i16> + ``` + +--- + +## 7.3 Gather / Scatter + +> **Category C** — contiguous-required. `pto.as` materializes `.contiguous()` +> before these ops if the input layout is non-contiguous. + +### `pto.vmi.vgather` + +- **semantics:** Indexed gather from UB at B32 granularity. For each active + lane `i`, load `src[offsets[i]]`. + + ```c + for (int i = 0; i < L; i++) + dst[i] = mask[i] ? ub[base + offsets[i]] : (pmode_merge ? dst_old[i] : 0); + ``` + +- **syntax:** + ```mlir + %g = pto.vmi.vgather %src, %offsets, %mask : !pto.ptr, !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg + ``` +- **operands:** + + | Operand | Type | Description | + |---|---|---| + | `src` | `!pto.ptr` | UB base pointer | + | `offsets` | `!pto.vmi.vreg` | Per-lane element offset | + | `mask` | `!pto.vmi.mask` | Governing predicate | + +- **results:** `!pto.vmi.vreg` +- **attributes:** `pmode` +- **datatypes:** `i8`–`i32`, `f16`, `bf16`, `f32` +- **lowering to `pto.mi`:** + ``` + K × pto.vgather2 + ``` + `#mi = K`, `dep = 1`, util data-dependent. + +### `pto.vmi.vgatherb` + +- **semantics:** Byte-granularity indexed gather. Mask lane count equals result + lane count (may differ from offset lane count). + + ```c + for (int i = 0; i < L; i++) + dst[i] = mask[i] ? ub_byte[base_byte + offsets[i]] : (pmode_merge ? dst_old[i] : 0); + ``` + +- **syntax:** + ```mlir + %gb = pto.vmi.vgatherb %src, %offsets, %mask : !pto.ptr, !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg + ``` +- **datatypes:** `i8`–`i32`, `f16`, `bf16`, `f32` +- **lowering to `pto.mi`:** + ``` + K × pto.vgatherb + ``` + `#mi = K`, `dep = 1`. + +### `pto.vmi.vscatter` + +- **semantics:** Indexed scatter to UB. For each active lane `i`, + write `value[i]` to `dest[offsets[i]]`. + + ```c + for (int i = 0; i < L; i++) + if (mask[i]) + ub[base + offsets[i]] = value[i]; + ``` + +- **syntax:** + ```mlir + pto.vmi.vscatter %value, %dest, %offsets, %mask : !pto.vmi.vreg, !pto.ptr, !pto.vmi.vreg, !pto.vmi.mask + ``` +- **operands:** + + | Operand | Type | Description | + |---|---|---| + | `value` | `!pto.vmi.vreg` | Values to scatter | + | `dest` | `!pto.ptr` | UB destination base pointer | + | `offsets` | `!pto.vmi.vreg` | Per-lane element offset | + | `mask` | `!pto.vmi.mask` | Governing predicate | + +- **results:** *(none)* +- **attributes:** `pmode` +- **datatypes:** `i8`–`i32`, `f16`, `bf16`, `f32` +- **lowering to `pto.mi`:** + ``` + K × pto.vscatter + ``` + `#mi = K`, `dep = 1`. + +- **example:** + ```mlir + pto.vmi.vscatter %v, %dest, %offsets, %mask + : !pto.vmi.vreg<64×f32>, !pto.ptr, !pto.vmi.vreg<64×i32>, !pto.vmi.mask<64> + ``` diff --git a/docs/isa/vmi-isa/08-predicate-ops.md b/docs/isa/vmi-isa/08-predicate-ops.md new file mode 100644 index 0000000000..a1076906c5 --- /dev/null +++ b/docs/isa/vmi-isa/08-predicate-ops.md @@ -0,0 +1,122 @@ +# 8. Predicate Ops + +> **Category:** gen (mask producers — take no input mask). +> **Mask in:** none (they generate masks). +> +> Mask generation is expressed with two ops: `create_mask` (prefix / first-N +> tail) and `create_group_mask` (grouped prefix / grouped first-N tail). Mask +> granularity (`b8`/`b16`/`b32`) is derived from the result type, not spelled in +> the op name. +> +> `create_mask` takes a single `index` operand `active_lanes`. When +> `active_lanes ≥ L` it yields an all-active mask; when `active_lanes = N < L` +> it yields a first-N tail mask. `create_group_mask` repeats the first-N pattern +> within each of `num_groups` equal groups (group size `group_size`). + +```mlir +%act = arith.minsi %rem, %cL // min(rem, L) +%aidx = arith.index_cast %act // i32 -> index +%mask = pto.vmi.create_mask %aidx : index -> !pto.vmi.mask<128×b32> +%next = arith.subi %rem, %act // rem - min(rem, L) +``` + +--- + +## `pto.vmi.create_mask` + +- **syntax:** + ```mlir + %m = pto.vmi.create_mask %active_lanes : index -> !pto.vmi.mask + ``` +- **semantics:** Create a predicate mask where the first `active_lanes` logical + lanes are active and the rest are inactive. `active_lanes ≥ L` produces an + all-active mask; `active_lanes = N` produces a first-N tail mask. + + ```c + for (int i = 0; i < L; i++) + dst[i] = (i < active_lanes) ? 1 : 0; + ``` + +- **operands:** + + | Operand | Type | Description | + |---|---|---| + | `active_lanes` | `index` | Number of leading active lanes | + +- **results:** + + | Result | Type | Description | + |---|---|---| + | `result` | `!pto.vmi.mask` | Predicate mask | + +- **example:** + ```mlir + // All-active mask (active_lanes >= L) + %all = pto.vmi.create_mask %c128 : index -> !pto.vmi.mask<128×b32> + + // First-N tail mask (N = 64) + %tail = pto.vmi.create_mask %c64 : index -> !pto.vmi.mask<128×b32> + ``` + +--- + +## `pto.vmi.create_group_mask` + +- **syntax:** + ```mlir + %m = pto.vmi.create_group_mask %active_elems_per_group {num_groups = C, group_size = S} + : index -> !pto.vmi.mask + ``` +- **semantics:** Create a grouped predicate mask. The mask is divided into + `num_groups` equal groups of `group_size` lanes each; lane `i` is active iff + `(i % group_size) < active_elems_per_group`. When + `active_elems_per_group ≥ group_size` all lanes are active within every group + (grouped all-active); otherwise the first `active_elems_per_group` lanes are + active within each group (grouped first-N tail). + + ```c + for (int i = 0; i < L; i++) + dst[i] = ((i % group_size) < active_elems_per_group) ? 1 : 0; + ``` + +- **operands:** + + | Operand | Type | Description | + |---|---|---| + | `active_elems_per_group` | `index` | Active lanes within each group | + +- **attributes:** + + | Attribute | Values | Description | + |---|---|---| + | `num_groups` | positive integer | Number of equal groups | + | `group_size` | positive integer | Lanes per group (`L / num_groups`) | + +- **results:** + + | Result | Type | Description | + |---|---|---| + | `result` | `!pto.vmi.mask` | Grouped predicate mask | + +- **example:** + ```mlir + // Grouped all-active: 8 groups, group size 32, all lanes active per group + %all = pto.vmi.create_group_mask %c32 {num_groups = 8, group_size = 32} + : index -> !pto.vmi.mask<256×b32> + + // Grouped first-N tail: first 25 lanes per group, 8 groups + %tail = pto.vmi.create_group_mask %c25 {num_groups = 8, group_size = 32} + : index -> !pto.vmi.mask<256×b32> + ``` + +--- + +> **Mask Boolean Ops (`vand` / `vor` / `vxor` / `vnot` on masks):** +> +> There is **no dedicated predicate-logic op** (e.g. `pand`/`por`/`pxor`/`pnot`). +> Mask (predicate) boolean operations are **not yet supported**, but are planned. +> The planned approach is to **reuse the elementwise bitwise ops** `pto.vmi.vand` / +> `vor` / `vxor` / `vnot` directly on mask operands — their implementations will be +> extended to accept mask types (treated as a per-lane bit-wise boolean op on the +> predicate). This also covers the `pnot`-style predicate complement needed by MERGE +> emulation (see [Appendix C](10-appendices.md)). diff --git a/docs/isa/vmi-isa/09-data-rearrange.md b/docs/isa/vmi-isa/09-data-rearrange.md new file mode 100644 index 0000000000..5779151e2c --- /dev/null +++ b/docs/isa/vmi-isa/09-data-rearrange.md @@ -0,0 +1,109 @@ +# 9. Data Rearrange + +> **Category:** A (layout-transparent). **Mask:** `Pg`. +> +> In-register data movement and permutation. No UB access. `vintlv`/`vdintlv` +> are per-lane, dtype-preserving ops that do not change vreg layout — the output +> has the same `L` and `T` as the inputs. Commonly used for real+imaginary and +> value+index interleaving within a single vector register. + +--- + +## `pto.vmi.vintlv` + +- **semantics:** Interleave two source vectors by even/odd lanes. + + ```c + // low = {lhs[0], rhs[0], lhs[1], rhs[1], ..., lhs[L/2-1], rhs[L/2-1]} + // high = {lhs[L/2], rhs[L/2], lhs[L/2+1], rhs[L/2+1], ...} + for (int i = 0; i < L/2; i++) { + lo[2*i] = lhs[i]; + lo[2*i + 1] = rhs[i]; + hi[2*i] = lhs[L/2 + i]; + hi[2*i + 1] = rhs[L/2 + i]; + } + ``` + +- **syntax:** + ```mlir + %lo, %hi = pto.vmi.vintlv %lhs, %rhs, %mask : !pto.vmi.vreg, !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg, !pto.vmi.vreg + ``` +- **operands:** + + | Operand | Type | Description | + |---|---|---| + | `lhs` | `!pto.vmi.vreg` | First source (provides low-half even slots) | + | `rhs` | `!pto.vmi.vreg` | Second source (provides low-half odd slots) | + | `mask` | `!pto.vmi.mask` | Governing predicate | + +- **results:** + + | Result | Type | Description | + |---|---|---| + | `low` | `!pto.vmi.vreg` | Even-odd interleaved low half | + | `high` | `!pto.vmi.vreg` | Even-odd interleaved high half | + +- **attributes:** `pmode` +- **datatypes:** `i8`–`i32`, `f16`, `bf16`, `f32` +- **lowering to `pto.mi`:** + ``` + K × pto.vintlv + ``` + `#mi = K`, `dep = 1`. Layout-transparent (Category A). + +- **example:** + ```mlir + %lo, %hi = pto.vmi.vintlv %a, %b, %mask + : !pto.vmi.vreg<64×f32>, !pto.vmi.vreg<64×f32>, !pto.vmi.mask<64> + -> !pto.vmi.vreg<64×f32>, !pto.vmi.vreg<64×f32> + ``` + +--- + +## `pto.vmi.vdintlv` + +- **semantics:** Deinterleave a paired-source by even/odd lanes (AoS → SoA). + + ```c + // lhs, rhs treated as pairs: (lhs[0], rhs[0]), (lhs[1], rhs[1]), ... + // even = {lhs[0], lhs[2], lhs[4], ...} (all even-indexed slots from paired stream) + // odd = {lhs[1], lhs[3], lhs[5], ...} (all odd-indexed slots from paired stream) + // More precisely: + // low = {lhs[0], lhs[1], lhs[2], lhs[3], ...} ← original even slots from each pair + // high = {rhs[0], rhs[1], rhs[2], rhs[3], ...} ← original odd slots from each pair + // After deinterleaving: + // even[i] = (i % 2 == 0) ? lhs[i/2] : rhs[i/2] — this is the vintlv inverse + for (int i = 0; i < L/2; i++) { + even[i] = lhs[2*i]; // even slots of paired input + even[L/2 + i] = lhs[2*i + 1]; + odd[i] = rhs[2*i]; // odd slots of paired input + odd[L/2 + i] = rhs[2*i + 1]; + } + ``` + +- **syntax:** + ```mlir + %even, %odd = pto.vmi.vdintlv %lhs, %rhs, %mask : !pto.vmi.vreg, !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg, !pto.vmi.vreg + ``` +- **operands:** Same shape as `vintlv`. +- **results:** Same shape as `vintlv` (two `!pto.vmi.vreg`). +- **datatypes:** `i8`–`i32`, `f16`, `bf16`, `f32` +- **lowering to `pto.mi`:** + ``` + K × pto.vdintlv + ``` + `#mi = K`, `dep = 1`. + +- **example:** + ```mlir + %even, %odd = pto.vmi.vdintlv %x, %y, %mask + : !pto.vmi.vreg<64×f32>, !pto.vmi.vreg<64×f32>, !pto.vmi.mask<64> + -> !pto.vmi.vreg<64×f32>, !pto.vmi.vreg<64×f32> + ``` + +- **notes:** + - `vintlv` and `vdintlv` are inverses: `vdintlv(vintlv(a, b))` recovers `(a, b)`. + - Both are Category A — they do **not** change vreg layout (parity/half/width + axes pass through unchanged). + - Common use cases: real+imaginary interleave, value+index pair manipulation, + complex number arithmetic. diff --git a/docs/isa/vmi-isa/10-appendices.md b/docs/isa/vmi-isa/10-appendices.md new file mode 100644 index 0000000000..e3aa44e21e --- /dev/null +++ b/docs/isa/vmi-isa/10-appendices.md @@ -0,0 +1,82 @@ +# Appendices + +--- + +## Appendix A: Unified Ops Index + +| # | Op | Group | Category | Brief | +|---|---|---|---|---| +| 1 | `pto.vmi.vload` | 1: Load/Store | A | Logical vector load from UB | +| 2 | `pto.vmi.vstore` | 1: Load/Store | A | Logical vector store to UB | +| 3 | `pto.vmi.vci` | 2: Index-gen | A | Lane-index vector generation | +| 4 | `pto.vmi.vadd` | 3: Eltwise | A | Elementwise add (fp+int unified) | +| 5 | `pto.vmi.vsub` | 3: Eltwise | A | Elementwise subtract | +| 6 | `pto.vmi.vmul` | 3: Eltwise | A | Elementwise multiply | +| 7 | `pto.vmi.vdiv` | 3: Eltwise | A | Elementwise divide (fp only) | +| 8 | `pto.vmi.vmax` | 3: Eltwise | A | Elementwise maximum | +| 9 | `pto.vmi.vmin` | 3: Eltwise | A | Elementwise minimum | +| 10 | `pto.vmi.vabs` | 3: Eltwise | A | Elementwise absolute value | +| 11 | `pto.vmi.vneg` | 3: Eltwise | A | Elementwise negate | +| 12 | `pto.vmi.vrelu` | 3: Eltwise | A | Elementwise ReLU | +| 13 | `pto.vmi.vexp` | 3: Eltwise | A | Elementwise exponential | +| 14 | `pto.vmi.vln` | 3: Eltwise | A | Elementwise natural log | +| 15 | `pto.vmi.vsqrt` | 3: Eltwise | A | Elementwise square root | +| 16 | `pto.vmi.vand` | 3: Eltwise | A | Elementwise bitwise AND | +| 17 | `pto.vmi.vor` | 3: Eltwise | A | Elementwise bitwise OR | +| 18 | `pto.vmi.vxor` | 3: Eltwise | A | Elementwise bitwise XOR | +| 19 | `pto.vmi.vnot` | 3: Eltwise | A | Elementwise bitwise NOT | +| 20 | `pto.vmi.vshl` | 3: Eltwise | A | Elementwise left shift | +| 21 | `pto.vmi.vshr` | 3: Eltwise | A | Elementwise signedness-aware right shift | +| 22 | `pto.vmi.vadds` | 3: Eltwise | A | Vector-scalar add | +| 23 | `pto.vmi.vmuls` | 3: Eltwise | A | Vector-scalar multiply | +| 24 | `pto.vmi.vmaxs` | 3: Eltwise | A | Vector-scalar maximum | +| 25 | `pto.vmi.vmins` | 3: Eltwise | A | Vector-scalar minimum | +| 26 | `pto.vmi.vshls` | 3: Eltwise | A | Vector-scalar shift left | +| 27 | `pto.vmi.vshrs` | 3: Eltwise | A | Vector-scalar shift right | +| 28 | `pto.vmi.vcmp` | 3: Eltwise | A | Elementwise compare → mask | +| 29 | `pto.vmi.vcmps` | 3: Eltwise | A | Vector-scalar compare → mask | +| 30 | `pto.vmi.vsel` | 3: Eltwise | A | Predicate select | +| 31 | `pto.vmi.vselr` | 3: Eltwise | A | Dynamic lane permute | +| 32 | `pto.vmi.vbrc` | 4: Broadcast | A/B | Broadcast scalar/group-slot | +| 33 | `pto.vmi.vcadd` | 5: Reduce | B | Add-reduction | +| 34 | `pto.vmi.vcmax` | 5: Reduce | B | Max-reduction | +| 35 | `pto.vmi.vcmin` | 5: Reduce | B | Min-reduction | +| 36 | `pto.vmi.vcvt` | 6: Convert | B | Unified type conversion | +| 37 | `pto.vmi.vinterpret_cast` | 6: Convert | A | Bitwise reinterpret | +| 38 | `pto.vmi.vexpdif` | 7: SFU | A | Fused exp(x−max) | +| 39 | `pto.vmi.vaxpy` | 7: SFU | A | Fused α·x+y | +| 40 | `pto.vmi.vlrelu` | 7: SFU | A | Leaky ReLU | +| 41 | `pto.vmi.vprelu` | 7: SFU | A | Parametric ReLU | +| 42 | `pto.vmi.vmull` | 7: SFU | A | Pair-result widening 32×32 multiply | +| 43 | `pto.vmi.vmula` | 7: SFU | A | Fused multiply-add | +| 44 | `pto.vmi.vchist` | 7: SFU | B | Cumulative histogram (half-axis) | +| 45 | `pto.vmi.vdhist` | 7: SFU | B | Distribution histogram (plain per-bin) | +| 46 | `pto.vmi.vgather` | 7: SFU | C | Indexed gather (B32) | +| 47 | `pto.vmi.vgatherb` | 7: SFU | C | Byte-granularity indexed gather | +| 48 | `pto.vmi.vscatter` | 7: SFU | C | Indexed scatter | +| 49 | `pto.vmi.create_mask` | 8: Predicate | gen | Prefix / first-N tail mask | +| 50 | `pto.vmi.create_group_mask` | 8: Predicate | gen | Grouped predicate mask | +| 51 | `pto.vmi.vintlv` | 9: Rearrange | A | Interleave two vectors | +| 52 | `pto.vmi.vdintlv` | 9: Rearrange | A | Deinterleave two vectors | + +--- + +## Appendix C: MERGE Mode Emulation (A5) + +On A5, the hardware predicates only in **ZEROING** mode (inactive lanes → 0). +MERGE mode is emulated by `pto.as`: + +```mlir +// MERGE emulation on A5: dst = Pg ? op(...) : dst_old +%npg = pto.vmi.vnot %pg // complement predicate +%new_z = pto.vmi. %a, %b, %pg // ZEROING: inactive → 0 +%old_z = pto.vmi.vand %dst_old, %npg // keep old on inactive lanes +%dst = pto.vmi.vor %new_z, %old_z // disjoint OR → merged +``` + +Alternatively, a single `vsel %pg, %new, %dst_old` can replace the `vand`+`vor` +pair. + +**MERGE cost on A5:** `+1 vnot` (once per distinct `Pg`) + `+K vsel`/`vor`. +On A6, merge-capable ops take the mode natively — the `vnot`+`vor` emulation +collapses to the single predicated op. diff --git a/docs/release/VMI_VERSION b/docs/release/VMI_VERSION new file mode 100644 index 0000000000..17e51c385e --- /dev/null +++ b/docs/release/VMI_VERSION @@ -0,0 +1 @@ +0.1.1 diff --git a/include/PTO/IR/PTOAttrs.td b/include/PTO/IR/PTOAttrs.td index b87e1486e6..dab64426b4 100644 --- a/include/PTO/IR/PTOAttrs.td +++ b/include/PTO/IR/PTOAttrs.td @@ -38,6 +38,8 @@ class PTO_Attr traits = []> let mnemonic = attrMnemonic; } +include "PTO/IR/VMIAttrs.td" + //===----------------------------------------------------------------------===// // Address Space //===----------------------------------------------------------------------===// diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index 6ed8a63774..50c4156ddc 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -76,6 +76,7 @@ class PTO_DpsOp traits = []> class PTO_Op traits = []> : Op; +include "PTO/IR/VMIOps.td" include "PTO/IR/VPTOOps.td" //===----------------------------------------------------------------------===// diff --git a/include/PTO/IR/PTOTypeDefs.td b/include/PTO/IR/PTOTypeDefs.td index 69003cf5a4..7310f2a8e2 100644 --- a/include/PTO/IR/PTOTypeDefs.td +++ b/include/PTO/IR/PTOTypeDefs.td @@ -377,4 +377,5 @@ def F4E2M1x2Type : TypeDef { + let summary = "VMI logical vector register layout"; + let parameters = (ins + StringRefParameter<"layout kind">:$kind, + "int64_t":$factor, + "int64_t":$blockElems, + "int64_t":$slots, + "int64_t":$laneStride + ); + let hasCustomAssemblyFormat = 1; + let genVerifyDecl = 1; + + let extraClassDeclaration = [{ + static VMILayoutAttr getContiguous(::mlir::MLIRContext *context, + int64_t laneStride = 1); + static VMILayoutAttr getDeinterleaved(::mlir::MLIRContext *context, + int64_t factor, + int64_t blockElems = 1, + int64_t laneStride = 1); + static VMILayoutAttr getGroupSlots(::mlir::MLIRContext *context, + int64_t numGroups, + int64_t slots = 0, + int64_t laneStride = 1); + + bool isContiguous() const { return getKind() == "contiguous"; } + bool isDeinterleaved() const { return getKind() == "deinterleaved"; } + bool isGroupSlots() const { return getKind() == "num_groups"; } + bool isDense() const { return isContiguous() || isDeinterleaved(); } + int64_t getNumGroups() const { return getFactor(); } + bool hasDenseLaneStride() const { + return isDense() && getLaneStride() != 1; + } + bool hasGroupSlotLaneStride() const { + return isGroupSlots() && getLaneStride() != 1; + } + bool hasLaneStride() const { return getLaneStride() != 1; } + }]; +} + +//===----------------------------------------------------------------------===// +// VMI Predication Mode — [pmode] +// Controls inactive/predicated-off lane behavior on vector operations. +//===----------------------------------------------------------------------===// + +def VMI_PMode_Merge : I32EnumAttrCase<"Merge", 0, "merge">; +def VMI_PMode_Zero : I32EnumAttrCase<"Zero", 1, "zero">; +def VMIPredicationMode : I32EnumAttr<"PredicationMode", + "VMI predication mode: merge (keep old value) or zero (fill with 0)", [ + VMI_PMode_Merge, VMI_PMode_Zero +]> { + let cppNamespace = "::mlir::pto"; +} + +//===----------------------------------------------------------------------===// +// VMI Distribution Mode — {dist-mode} +// Controls the data-layout interpretation of load/store operations. +//===----------------------------------------------------------------------===// + +def VMI_DistMode_Continuous : I32EnumAttrCase<"Continuous", 0, "continuous">; +def VMI_DistMode_Unpack : I32EnumAttrCase<"Unpack", 1, "unpack">; +def VMI_DistMode_Dintlv : I32EnumAttrCase<"Dintlv", 2, "dintlv">; +def VMI_DistMode_Brc : I32EnumAttrCase<"Brc", 3, "brc">; +def VMIDistMode : I32EnumAttr<"DistMode", + "VMI distribution mode for load/store", [ + VMI_DistMode_Continuous, VMI_DistMode_Unpack, + VMI_DistMode_Dintlv, VMI_DistMode_Brc +]> { + let cppNamespace = "::mlir::pto"; +} + +//===----------------------------------------------------------------------===// +// VMI Rounding Mode — {rnd=} +// Controls the rounding direction for narrowing convert operations. +//===----------------------------------------------------------------------===// + +def VMI_Round_RTZ : I32EnumAttrCase<"RTZ", 0, "rtz">; +def VMI_Round_RNE : I32EnumAttrCase<"RNE", 1, "rne">; +def VMI_Round_Odd : I32EnumAttrCase<"Odd", 2, "odd">; +def VMIConvRoundMode : I32EnumAttr<"ConvRoundMode", + "VMI rounding mode for narrowing conversions", [ + VMI_Round_RTZ, VMI_Round_RNE, VMI_Round_Odd +]> { + let cppNamespace = "::mlir::pto"; +} + +//===----------------------------------------------------------------------===// +// VMI Saturation Mode — {sat=} +// Controls whether narrowing conversions saturate on overflow. +//===----------------------------------------------------------------------===// + +def VMI_Sat_None : I32EnumAttrCase<"None", 0, "none">; +def VMI_Sat_Sat : I32EnumAttrCase<"Sat", 1, "sat">; +def VMISatMode : I32EnumAttr<"SatMode", + "VMI saturation mode for narrowing conversions", [ + VMI_Sat_None, VMI_Sat_Sat +]> { + let cppNamespace = "::mlir::pto"; +} + +#endif // MLIR_DIALECT_PTO_IR_VMIATTRS diff --git a/include/PTO/IR/VMIOps.td b/include/PTO/IR/VMIOps.td new file mode 100644 index 0000000000..0a1dd9c878 --- /dev/null +++ b/include/PTO/IR/VMIOps.td @@ -0,0 +1,1684 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +//===- VMIOps.td - PTO VMI semantic operations -------------*- tablegen -*-===// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_DIALECT_PTO_IR_VMIOPS +#define MLIR_DIALECT_PTO_IR_VMIOPS + +include "mlir/IR/OpBase.td" +include "mlir/Interfaces/SideEffectInterfaces.td" + +def VMI_VRegTypeConstraint : Type< + CPred<"::llvm::isa<::mlir::pto::VMIVRegType>($_self)">, + "VMI logical vector register type">; + +def VMI_MaskTypeConstraint : Type< + CPred<"::llvm::isa<::mlir::pto::VMIMaskType>($_self)">, + "VMI logical mask type">; + +def VMI_ValueTypeConstraint : Type< + CPred<"::llvm::isa<::mlir::pto::VMIVRegType, ::mlir::pto::VMIMaskType>($_self)">, + "VMI logical vector or mask type">; + +def PTO_PhysicalVRegTypeConstraint : Type< + CPred<"::llvm::isa<::mlir::pto::VRegType>($_self)">, + "PTO physical vector register type">; + +def PTO_PhysicalMaskTypeConstraint : Type< + CPred<"::llvm::isa<::mlir::pto::MaskType>($_self)">, + "PTO physical mask type">; + +def PTO_PhysicalVMIPartTypeConstraint : AnyTypeOf< + [PTO_PhysicalVRegTypeConstraint, PTO_PhysicalMaskTypeConstraint], + "PTO physical vector register or mask type">; + +class VMI_Op traits = []> + : PTO_Op<"vmi." # mnemonic, traits>; + +//===--- Legacy (old) VMI ops ---===// + +def VMIConstantOp : VMI_Op<"constant", [Pure]> { + let summary = "VMI logical vector constant"; + let arguments = (ins AnyAttr:$value); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; +} + +def VMIBroadcastOp : VMI_Op<"broadcast", [Pure]> { + let summary = "Broadcast one scalar or 1-lane VMI vector to a VMI logical vector"; + let arguments = (ins AnyType:$value); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$value attr-dict `:` type($value) `->` type($result)"; +} + +def VMIIotaOp : VMI_Op<"iota", [Pure]> { + let summary = "Create a VMI logical index vector from a scalar base"; + let arguments = (ins + AnyTypeOf<[AnyInteger, AnyFloat], "integer/float scalar">:$base, + OptionalAttr:$order + ); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$base attr-dict `:` type($base) `->` type($result)"; +} + +def VMICreateMaskOp : VMI_Op<"create_mask", [Pure]> { + let summary = "Create a VMI logical prefix predicate mask"; + let arguments = (ins Index:$active_lanes); + let results = (outs VMI_MaskTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$active_lanes attr-dict `:` type($active_lanes) `->` type($result)"; +} + +def VMICreateGroupMaskOp : VMI_Op<"create_group_mask", [Pure]> { + let summary = "Create a VMI logical grouped predicate mask"; + let description = [{ + Creates a mask where lane i is active iff + `(i % group_size) < active_elems_per_group`. + }]; + let arguments = (ins + Index:$active_elems_per_group, + I64Attr:$num_groups, + I64Attr:$group_size + ); + let results = (outs VMI_MaskTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$active_elems_per_group attr-dict `:` type($active_elems_per_group) `->` type($result)"; +} + +def VMIConstantMaskOp : VMI_Op<"constant_mask", [Pure]> { + let summary = "VMI logical predicate mask constant"; + let arguments = (ins AnyAttr:$value); + let results = (outs VMI_MaskTypeConstraint:$result); + let hasVerifier = 1; +} + +def VMIMaskAndOp : VMI_Op<"mask_and", [Pure]> { + let summary = "VMI logical predicate mask and"; + let arguments = (ins VMI_MaskTypeConstraint:$lhs, VMI_MaskTypeConstraint:$rhs); + let results = (outs VMI_MaskTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)"; +} + +def VMIMaskOrOp : VMI_Op<"mask_or", [Pure]> { + let summary = "VMI logical predicate mask or"; + let arguments = (ins VMI_MaskTypeConstraint:$lhs, VMI_MaskTypeConstraint:$rhs); + let results = (outs VMI_MaskTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)"; +} + +def VMIMaskXOrOp : VMI_Op<"mask_xor", [Pure]> { + let summary = "VMI logical predicate mask xor"; + let arguments = (ins VMI_MaskTypeConstraint:$lhs, VMI_MaskTypeConstraint:$rhs); + let results = (outs VMI_MaskTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)"; +} + +def VMIMaskNotOp : VMI_Op<"mask_not", [Pure]> { + let summary = "VMI logical predicate mask not"; + let arguments = (ins VMI_MaskTypeConstraint:$source); + let results = (outs VMI_MaskTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source attr-dict `:` type($source) `->` type($result)"; +} + +def VMIAddFOp : VMI_Op<"addf", [Pure]> { + let summary = "VMI floating-point elementwise add"; + let arguments = (ins VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)"; +} + +def VMIAddIOp : VMI_Op<"addi", [Pure]> { + let summary = "VMI integer elementwise add"; + let arguments = (ins VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)"; +} + +def VMISubFOp : VMI_Op<"subf", [Pure]> { + let summary = "VMI floating-point elementwise subtract"; + let arguments = (ins VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)"; +} + +def VMISubIOp : VMI_Op<"subi", [Pure]> { + let summary = "VMI integer elementwise subtract"; + let arguments = (ins VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)"; +} + +def VMIMulFOp : VMI_Op<"mulf", [Pure]> { + let summary = "VMI floating-point elementwise multiply"; + let arguments = (ins VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)"; +} + +def VMIMulIOp : VMI_Op<"muli", [Pure]> { + let summary = "VMI integer elementwise multiply"; + let arguments = (ins VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)"; +} + +def VMIFmaOp : VMI_Op<"fma", [Pure]> { + let summary = "VMI fused floating-point multiply-add"; + let arguments = (ins VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs, + VMI_VRegTypeConstraint:$acc); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs `,` $acc attr-dict `:` type($lhs) `,` type($rhs) `,` type($acc) `->` type($result)"; +} + +def VMIDivFOp : VMI_Op<"divf", [Pure]> { + let summary = "VMI floating-point elementwise divide"; + let arguments = (ins VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)"; +} + +def VMIMinFOp : VMI_Op<"minf", [Pure]> { + let summary = "VMI floating-point elementwise minimum"; + let arguments = (ins VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)"; +} + +def VMIMaxFOp : VMI_Op<"maxf", [Pure]> { + let summary = "VMI floating-point elementwise maximum"; + let arguments = (ins VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)"; +} + +def VMINegFOp : VMI_Op<"negf", [Pure]> { + let summary = "VMI floating-point elementwise negate"; + let arguments = (ins VMI_VRegTypeConstraint:$source); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source attr-dict `:` type($source) `->` type($result)"; +} + +def VMIAbsFOp : VMI_Op<"absf", [Pure]> { + let summary = "VMI floating-point elementwise absolute value"; + let arguments = (ins VMI_VRegTypeConstraint:$source); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source attr-dict `:` type($source) `->` type($result)"; +} + +def VMIAbsIOp : VMI_Op<"absi", [Pure]> { + let summary = "VMI integer elementwise absolute value"; + let arguments = (ins VMI_VRegTypeConstraint:$source); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source attr-dict `:` type($source) `->` type($result)"; +} + +def VMISqrtOp : VMI_Op<"sqrt", [Pure]> { + let summary = "VMI floating-point elementwise square root"; + let arguments = (ins VMI_VRegTypeConstraint:$source); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source attr-dict `:` type($source) `->` type($result)"; +} + +def VMIExpOp : VMI_Op<"exp", [Pure]> { + let summary = "VMI floating-point elementwise exponential"; + let arguments = (ins VMI_VRegTypeConstraint:$source); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source attr-dict `:` type($source) `->` type($result)"; +} + +def VMILnOp : VMI_Op<"ln", [Pure]> { + let summary = "VMI floating-point elementwise natural logarithm"; + let arguments = (ins VMI_VRegTypeConstraint:$source); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source attr-dict `:` type($source) `->` type($result)"; +} + +def VMIReluOp : VMI_Op<"relu", [Pure]> { + let summary = "VMI floating-point elementwise ReLU"; + let arguments = (ins VMI_VRegTypeConstraint:$source); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source attr-dict `:` type($source) `->` type($result)"; +} + +def VMIAndIOp : VMI_Op<"andi", [Pure]> { + let summary = "VMI integer elementwise bitwise and"; + let arguments = (ins VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)"; +} + +def VMIOrIOp : VMI_Op<"ori", [Pure]> { + let summary = "VMI integer elementwise bitwise or"; + let arguments = (ins VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)"; +} + +def VMIXOrIOp : VMI_Op<"xori", [Pure]> { + let summary = "VMI integer elementwise bitwise xor"; + let arguments = (ins VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)"; +} + +def VMIShLIOp : VMI_Op<"shli", [Pure]> { + let summary = "VMI integer elementwise left shift"; + let arguments = (ins VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)"; +} + +def VMIShRUIOp : VMI_Op<"shrui", [Pure]> { + let summary = "VMI unsigned integer elementwise right shift"; + let arguments = (ins VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)"; +} + +def VMIShRSIOp : VMI_Op<"shrsi", [Pure]> { + let summary = "VMI signed integer elementwise arithmetic right shift"; + let arguments = (ins VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)"; +} + +def VMINotOp : VMI_Op<"not", [Pure]> { + let summary = "VMI integer elementwise bitwise not"; + let arguments = (ins VMI_VRegTypeConstraint:$source); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source attr-dict `:` type($source) `->` type($result)"; +} + +def VMICmpFOp : VMI_Op<"cmpf", [Pure]> { + let summary = "VMI floating-point elementwise compare"; + let arguments = (ins StrAttr:$predicate, VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs); + let results = (outs VMI_MaskTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$predicate `,` $lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)"; +} + +def VMICmpIOp : VMI_Op<"cmpi", [Pure]> { + let summary = "VMI integer elementwise compare"; + let arguments = (ins StrAttr:$predicate, VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs); + let results = (outs VMI_MaskTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$predicate `,` $lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)"; +} + +def VMISelectOp : VMI_Op<"select", [Pure]> { + let summary = "VMI elementwise select"; + let arguments = (ins VMI_MaskTypeConstraint:$mask, VMI_VRegTypeConstraint:$true_value, + VMI_VRegTypeConstraint:$false_value); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$mask `,` $true_value `,` $false_value attr-dict `:` type($mask) `,` type($true_value) `,` type($false_value) `->` type($result)"; +} + +def VMIActivePrefixIndexOp : VMI_Op<"active_prefix_index"> { + let summary = "VMI per-lane active-prefix index from a predicate mask"; + let arguments = (ins VMI_MaskTypeConstraint:$mask); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$mask attr-dict `:` type($mask) `->` type($result)"; +} + +def VMICompressOp : VMI_Op<"compress"> { + let summary = "VMI compact active source lanes according to a predicate mask"; + let arguments = (ins VMI_VRegTypeConstraint:$source, VMI_MaskTypeConstraint:$mask); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `,` $mask attr-dict `:` type($source) `,` type($mask) `->` type($result)"; +} + +def VMICompressStoreOp : VMI_Op<"compress_store", [DeclareOpInterfaceMethods]> { + let summary = "VMI store active source lanes contiguously according to a predicate mask"; + let arguments = (ins VMI_VRegTypeConstraint:$value, PtrOrMemRef:$destination, + Index:$offset, VMI_MaskTypeConstraint:$mask); + let results = (outs); + let hasVerifier = 1; + let assemblyFormat = "$value `,` $destination `[` $offset `]` `,` $mask attr-dict `:` type($value) `,` type($destination) `,` type($mask)"; +} + +def VMIReduceAddIOp : VMI_Op<"reduce_addi"> { + let summary = "VMI masked integer add reduction with a 1-lane vector init"; + let arguments = (ins VMI_VRegTypeConstraint:$source, + VMI_VRegTypeConstraint:$init, + VMI_MaskTypeConstraint:$mask); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `,` $init `,` $mask attr-dict `:` type($source) `,` type($init) `,` type($mask) `->` type($result)"; +} + +def VMIReduceAddFOp : VMI_Op<"reduce_addf"> { + let summary = "VMI masked floating-point add reduction with explicit reassociation permission"; + let arguments = (ins VMI_VRegTypeConstraint:$source, + VMI_VRegTypeConstraint:$init, + VMI_MaskTypeConstraint:$mask, + OptionalAttr:$reassoc); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `,` $init `,` $mask attr-dict `:` type($source) `,` type($init) `,` type($mask) `->` type($result)"; +} + +def VMIReduceMaxFOp : VMI_Op<"reduce_maxf"> { + let summary = "VMI masked floating-point maximum reduction with a 1-lane vector init"; + let arguments = (ins VMI_VRegTypeConstraint:$source, + VMI_VRegTypeConstraint:$init, + VMI_MaskTypeConstraint:$mask); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `,` $init `,` $mask attr-dict `:` type($source) `,` type($init) `,` type($mask) `->` type($result)"; +} + +def VMIReduceMinFOp : VMI_Op<"reduce_minf"> { + let summary = "VMI masked floating-point minimum reduction with a 1-lane vector init"; + let arguments = (ins VMI_VRegTypeConstraint:$source, + VMI_VRegTypeConstraint:$init, + VMI_MaskTypeConstraint:$mask); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `,` $init `,` $mask attr-dict `:` type($source) `,` type($init) `,` type($mask) `->` type($result)"; +} + +def VMIReduceMaxIOp : VMI_Op<"reduce_maxi"> { + let summary = "VMI masked integer maximum reduction with a 1-lane vector init"; + let arguments = (ins VMI_VRegTypeConstraint:$source, + VMI_VRegTypeConstraint:$init, + VMI_MaskTypeConstraint:$mask); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `,` $init `,` $mask attr-dict `:` type($source) `,` type($init) `,` type($mask) `->` type($result)"; +} + +def VMIReduceMinIOp : VMI_Op<"reduce_mini"> { + let summary = "VMI masked integer minimum reduction with a 1-lane vector init"; + let arguments = (ins VMI_VRegTypeConstraint:$source, + VMI_VRegTypeConstraint:$init, + VMI_MaskTypeConstraint:$mask); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `,` $init `,` $mask attr-dict `:` type($source) `,` type($init) `,` type($mask) `->` type($result)"; +} + +def VMIGroupReduceAddFOp : VMI_Op<"group_reduce_addf"> { + let summary = "VMI masked floating-point add reduction within fixed logical groups"; + let arguments = (ins VMI_VRegTypeConstraint:$source, + VMI_MaskTypeConstraint:$mask, + I64Attr:$num_groups, + OptionalAttr:$reassoc); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `,` $mask attr-dict `:` type($source) `,` type($mask) `->` type($result)"; +} + +def VMIGroupReduceMaxFOp : VMI_Op<"group_reduce_maxf"> { + let summary = "VMI masked floating-point maximum reduction within fixed logical groups"; + let arguments = (ins VMI_VRegTypeConstraint:$source, + VMI_MaskTypeConstraint:$mask, + I64Attr:$num_groups); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `,` $mask attr-dict `:` type($source) `,` type($mask) `->` type($result)"; +} + +def VMIGroupReduceMinFOp : VMI_Op<"group_reduce_minf"> { + let summary = "VMI masked floating-point minimum reduction within fixed logical groups"; + let arguments = (ins VMI_VRegTypeConstraint:$source, + VMI_MaskTypeConstraint:$mask, + I64Attr:$num_groups); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `,` $mask attr-dict `:` type($source) `,` type($mask) `->` type($result)"; +} + +def VMIGroupReduceAddIOp : VMI_Op<"group_reduce_addi"> { + let summary = "VMI masked integer add reduction within fixed logical groups"; + let arguments = (ins VMI_VRegTypeConstraint:$source, + VMI_MaskTypeConstraint:$mask, + I64Attr:$num_groups); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `,` $mask attr-dict `:` type($source) `,` type($mask) `->` type($result)"; +} + +def VMIGroupReduceMaxIOp : VMI_Op<"group_reduce_maxi"> { + let summary = "VMI masked integer maximum reduction within fixed logical groups"; + let arguments = (ins VMI_VRegTypeConstraint:$source, + VMI_MaskTypeConstraint:$mask, + I64Attr:$num_groups); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `,` $mask attr-dict `:` type($source) `,` type($mask) `->` type($result)"; +} + +def VMIGroupReduceMinIOp : VMI_Op<"group_reduce_mini"> { + let summary = "VMI masked integer minimum reduction within fixed logical groups"; + let arguments = (ins VMI_VRegTypeConstraint:$source, + VMI_MaskTypeConstraint:$mask, + I64Attr:$num_groups); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `,` $mask attr-dict `:` type($source) `,` type($mask) `->` type($result)"; +} + +def VMIGroupBroadcastOp : VMI_Op<"group_broadcast"> { + let summary = "VMI broadcast group-slot values back to each logical group"; + let arguments = (ins VMI_VRegTypeConstraint:$source, + I64Attr:$num_groups); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source attr-dict `:` type($source) `->` type($result)"; +} + +class VMIHistogramOp + : VMI_Op { + let summary = summaryText; + let arguments = (ins VMI_VRegTypeConstraint:$acc, + VMI_VRegTypeConstraint:$source, + VMI_MaskTypeConstraint:$mask); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$acc `,` $source `,` $mask attr-dict `:` type($acc) `,` type($source) `,` type($mask) `->` type($result)"; +} + +def VMIVdhistOp : VMIHistogramOp<"vdhist", + "VMI distribution histogram (dhistv2) over unsigned 8-bit source lanes">; + +def VMIVchistOp : VMIHistogramOp<"vchist", + "VMI cumulative histogram (chistv2 half-axis) over unsigned 8-bit source lanes">; + +def VMIExtFOp : VMI_Op<"extf", [Pure]> { + let summary = "VMI floating-point elementwise extension"; + let arguments = (ins VMI_VRegTypeConstraint:$source); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source attr-dict `:` type($source) `->` type($result)"; +} + +def VMITruncFOp : VMI_Op<"truncf", [Pure]> { + let summary = "VMI floating-point elementwise truncation"; + let arguments = (ins VMI_VRegTypeConstraint:$source, + OptionalAttr:$rounding); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source attr-dict `:` type($source) `->` type($result)"; +} + +def VMIFPToSIOp : VMI_Op<"fptosi", [Pure]> { + let summary = "VMI floating-point to signed integer elementwise conversion"; + let arguments = (ins VMI_VRegTypeConstraint:$source); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source attr-dict `:` type($source) `->` type($result)"; +} + +def VMISIToFPOp : VMI_Op<"sitofp", [Pure]> { + let summary = "VMI signed integer to floating-point elementwise conversion"; + let arguments = (ins VMI_VRegTypeConstraint:$source); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source attr-dict `:` type($source) `->` type($result)"; +} + +def VMIExtSIOp : VMI_Op<"extsi", [Pure]> { + let summary = "VMI signed integer elementwise extension"; + let arguments = (ins VMI_VRegTypeConstraint:$source); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source attr-dict `:` type($source) `->` type($result)"; +} + +def VMIExtUIOp : VMI_Op<"extui", [Pure]> { + let summary = "VMI unsigned integer elementwise extension"; + let arguments = (ins VMI_VRegTypeConstraint:$source); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source attr-dict `:` type($source) `->` type($result)"; +} + +def VMITruncIOp : VMI_Op<"trunci", [Pure]> { + let summary = "VMI saturating integer elementwise truncation"; + let arguments = (ins VMI_VRegTypeConstraint:$source); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source attr-dict `:` type($source) `->` type($result)"; +} + +def VMIBitcastOp : VMI_Op<"bitcast", [Pure]> { + let summary = "VMI bitwise vector reinterpretation"; + let arguments = (ins VMI_VRegTypeConstraint:$source); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source attr-dict `:` type($source) `->` type($result)"; +} + +def VMILoadOp : VMI_Op<"load", [DeclareOpInterfaceMethods]> { + let summary = "VMI logical vector load"; + let arguments = (ins PtrOrMemRef:$source, Index:$offset); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `[` $offset `]` attr-dict `:` type($source) `->` type($result)"; +} + +def VMIDeinterleaveLoadOp : VMI_Op<"deinterleave_load", [DeclareOpInterfaceMethods]> { + let summary = "VMI two-way logical deinterleave load"; + let arguments = (ins PtrOrMemRef:$source, Index:$offset); + let results = (outs VMI_VRegTypeConstraint:$low, VMI_VRegTypeConstraint:$high); + let hasVerifier = 1; + let assemblyFormat = "$source `[` $offset `]` attr-dict `:` type($source) `->` type($low) `,` type($high)"; +} + +def VMIGroupLoadOp : VMI_Op<"group_load", [DeclareOpInterfaceMethods]> { + let summary = "VMI logical grouped vector load with a row stride between groups"; + let arguments = (ins PtrOrMemRef:$source, Index:$offset, Index:$row_stride, + I64Attr:$num_groups); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `[` $offset `]` `,` $row_stride attr-dict `:` type($source) `->` type($result)"; +} + +def VMIGroupSlotLoadOp : VMI_Op<"group_slot_load", [DeclareOpInterfaceMethods]> { + let summary = "VMI load one scalar value per logical group into group slots"; + let arguments = (ins PtrOrMemRef:$source, Index:$offset, Index:$source_group_stride, + I64Attr:$num_groups); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `[` $offset `]` `,` $source_group_stride attr-dict `:` type($source) `->` type($result)"; +} + +def VMIGroupBroadcastLoadOp : VMI_Op<"group_broadcast_load", [DeclareOpInterfaceMethods]> { + let summary = "VMI load one scalar value per logical group and broadcast it to group lanes"; + let arguments = (ins PtrOrMemRef:$source, Index:$offset, Index:$source_group_stride, + I64Attr:$num_groups); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `[` $offset `]` `,` $source_group_stride attr-dict `:` type($source) `->` type($result)"; +} + +def VMIStrideLoadOp : VMI_Op<"stride_load", [DeclareOpInterfaceMethods]> { + let summary = "VMI block-strided vector load"; + let arguments = (ins PtrOrMemRef:$source, Index:$offset, + I16:$block_stride, I16:$repeat_stride, + VMI_MaskTypeConstraint:$mask); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `[` $offset `]` `,` $block_stride `,` $repeat_stride `,` $mask attr-dict `:` type($source) `,` type($block_stride) `,` type($repeat_stride) `,` type($mask) `->` type($result)"; +} + +def VMIMaskedLoadOp : VMI_Op<"masked_load", [DeclareOpInterfaceMethods]> { + let summary = "VMI logical masked vector load with passthrough lanes"; + let arguments = (ins PtrOrMemRef:$source, Index:$offset, + VMI_MaskTypeConstraint:$mask, + VMI_VRegTypeConstraint:$passthru); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `[` $offset `]` `,` $mask `,` $passthru attr-dict `:` type($source) `,` type($mask) `,` type($passthru) `->` type($result)"; +} + +def VMIGatherOp : VMI_Op<"gather", [DeclareOpInterfaceMethods]> { + let summary = "VMI logical masked indexed gather with passthrough lanes"; + let arguments = (ins PtrOrMemRef:$source, + VMI_VRegTypeConstraint:$indices, + VMI_MaskTypeConstraint:$mask, + VMI_VRegTypeConstraint:$passthru); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `[` $indices `]` `,` $mask `,` $passthru attr-dict `:` type($source) `,` type($indices) `,` type($mask) `,` type($passthru) `->` type($result)"; +} + +def VMIExpandLoadOp : VMI_Op<"expand_load", [DeclareOpInterfaceMethods]> { + let summary = "VMI load a dense active-lane stream into masked logical lanes"; + let arguments = (ins PtrOrMemRef:$source, Index:$offset, + VMI_MaskTypeConstraint:$mask, + VMI_VRegTypeConstraint:$passthru); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `[` $offset `]` `,` $mask `,` $passthru attr-dict `:` type($source) `,` type($mask) `,` type($passthru) `->` type($result)"; +} + +def VMIStoreOp : VMI_Op<"store", [DeclareOpInterfaceMethods]> { + let summary = "VMI logical vector store"; + let arguments = (ins VMI_VRegTypeConstraint:$value, PtrOrMemRef:$destination, Index:$offset); + let results = (outs); + let hasVerifier = 1; + let assemblyFormat = "$value `,` $destination `[` $offset `]` attr-dict `:` type($value) `,` type($destination)"; +} + +def VMIInterleaveStoreOp : VMI_Op<"interleave_store", [DeclareOpInterfaceMethods]> { + let summary = "VMI two-way logical interleave store"; + let arguments = (ins VMI_VRegTypeConstraint:$low, VMI_VRegTypeConstraint:$high, + PtrOrMemRef:$destination, Index:$offset); + let results = (outs); + let hasVerifier = 1; + let assemblyFormat = "$low `,` $high `,` $destination `[` $offset `]` attr-dict `:` type($low) `,` type($high) `,` type($destination)"; +} + +def VMIGroupStoreOp : VMI_Op<"group_store", [DeclareOpInterfaceMethods]> { + let summary = "VMI logical grouped vector store with a row stride between groups"; + let arguments = (ins VMI_VRegTypeConstraint:$value, PtrOrMemRef:$destination, + Index:$offset, Index:$row_stride, I64Attr:$num_groups); + let results = (outs); + let hasVerifier = 1; + let assemblyFormat = "$value `,` $destination `[` $offset `]` `,` $row_stride attr-dict `:` type($value) `,` type($destination)"; +} + +def VMIMaskedStoreOp : VMI_Op<"masked_store", [DeclareOpInterfaceMethods]> { + let summary = "VMI logical masked vector store"; + let arguments = (ins VMI_VRegTypeConstraint:$value, PtrOrMemRef:$destination, + Index:$offset, VMI_MaskTypeConstraint:$mask); + let results = (outs); + let hasVerifier = 1; + let assemblyFormat = "$value `,` $destination `[` $offset `]` `,` $mask attr-dict `:` type($value) `,` type($destination) `,` type($mask)"; +} + +def VMIStrideStoreOp : VMI_Op<"stride_store", [DeclareOpInterfaceMethods]> { + let summary = "VMI block-strided vector store"; + let arguments = (ins VMI_VRegTypeConstraint:$value, PtrOrMemRef:$destination, + Index:$offset, I16:$block_stride, I16:$repeat_stride, + VMI_MaskTypeConstraint:$mask); + let results = (outs); + let hasVerifier = 1; + let assemblyFormat = "$value `,` $destination `[` $offset `]` `,` $block_stride `,` $repeat_stride `,` $mask attr-dict `:` type($value) `,` type($destination) `,` type($block_stride) `,` type($repeat_stride) `,` type($mask)"; +} + +def VMIScatterOp : VMI_Op<"scatter", [DeclareOpInterfaceMethods]> { + let summary = "VMI logical masked indexed scatter"; + let arguments = (ins VMI_VRegTypeConstraint:$value, + PtrOrMemRef:$destination, + VMI_VRegTypeConstraint:$indices, + VMI_MaskTypeConstraint:$mask); + let results = (outs); + let hasVerifier = 1; + let assemblyFormat = "$value `,` $destination `[` $indices `]` `,` $mask attr-dict `:` type($value) `,` type($destination) `,` type($indices) `,` type($mask)"; +} + +def VMIShuffleOp : VMI_Op<"shuffle", [Pure]> { + let summary = "VMI static lane shuffle"; + let arguments = (ins VMI_VRegTypeConstraint:$source, DenseI64ArrayAttr:$indices); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `[` $indices `]` attr-dict `:` type($source) `->` type($result)"; +} + +def VMIChannelSplitOp : VMI_Op<"channel_split"> { + let summary = "VMI split interleaved logical channels"; + let arguments = (ins VMI_VRegTypeConstraint:$source); + let results = (outs Variadic:$results); + let hasVerifier = 1; +} + +def VMIChannelMergeOp : VMI_Op<"channel_merge"> { + let summary = "VMI merge logical channels by interleaving"; + let arguments = (ins Variadic:$inputs); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; +} + +def VMIEnsureLayoutOp : VMI_Op<"ensure_layout", [Pure]> { + let summary = "Internal VMI data layout materialization helper"; + let arguments = (ins VMI_VRegTypeConstraint:$source); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source attr-dict `:` type($source) `->` type($result)"; +} + +def VMIEnsureMaskLayoutOp : VMI_Op<"ensure_mask_layout", [Pure]> { + let summary = "Internal VMI mask layout materialization helper"; + let arguments = (ins VMI_MaskTypeConstraint:$source); + let results = (outs VMI_MaskTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source attr-dict `:` type($source) `->` type($result)"; +} + +def VMIEnsureMaskGranularityOp : VMI_Op<"ensure_mask_granularity", [Pure]> { + let summary = "Internal VMI mask granularity materialization helper"; + let arguments = (ins VMI_MaskTypeConstraint:$source); + let results = (outs VMI_MaskTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source attr-dict `:` type($source) `->` type($result)"; +} + +def VMIUnpackOp : VMI_Op<"unpack"> { + let summary = "Internal VMI value projection to physical parts"; + let arguments = (ins VMI_ValueTypeConstraint:$source); + let results = (outs Variadic:$parts); + let hasVerifier = 1; +} + +def VMIPackOp : VMI_Op<"pack"> { + let summary = "Internal physical parts materialized as one VMI value"; + let arguments = (ins Variadic:$parts); + let results = (outs VMI_ValueTypeConstraint:$result); + let hasVerifier = 1; +} + +//===--- Unified (new) VMI ops ---===// + +def VMIVbrcOp : VMI_Op<"vbrc", [Pure]> { + let summary = "VMI broadcast: scalar/1-lane to vector, or group-slot to dense vector"; + let arguments = (ins AnyType:$value, OptionalAttr:$group); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$value attr-dict `:` type($value) `->` type($result)"; +} + +def VMIVciOp : VMI_Op<"vci", [Pure]> { + let summary = "Create a VMI logical index vector from a scalar base"; + let arguments = (ins + AnyTypeOf<[AnyInteger, AnyFloat], "integer/float scalar">:$base, + OptionalAttr:$order + ); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$base attr-dict `:` type($base) `->` type($result)"; +} + +def VMIPsetOp : VMI_Op<"pset", [Pure]> { + let summary = "Create a predicate mask with all lanes active"; + let description = [{ + Creates a mask where all lanes are active (set to true). + Replacement for the all-active mode of `create_mask`. + + When {group = C} is present, creates a grouped all-active mask + (all lanes active within each of the C groups). + + Example: + ``` + %m = pto.vmi.pset "PAT_ALL" : !pto.vmi.mask<16> + %m2 = pto.vmi.pset "PAT_ALL" {group = 8} : !pto.vmi.mask<256xpred> + ``` + }]; + let arguments = (ins StrAttr:$pattern, OptionalAttr:$group); + let results = (outs VMI_MaskTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$pattern attr-dict `:` type($result)"; +} + +def VMIPgeOp : VMI_Op<"pge", [Pure]> { + let summary = "Create a tail predicate mask with the first N logical lanes active"; + let description = [{ + Creates a mask where the first N logical lanes are active and the + remaining lanes are inactive. Replacement for the tail mode of + `create_mask`. + + When {group = C} is present, creates a grouped tail mask where + the first N lanes within each of the C groups are active. + Replacement for `create_group_mask`. + + Example: + ``` + %m = pto.vmi.pge "PAT_VL16" : !pto.vmi.mask<32> + %m2 = pto.vmi.pge "PAT_VL25" {group = 8} : !pto.vmi.mask<256xpred> + ``` + }]; + let arguments = (ins StrAttr:$pattern, OptionalAttr:$group); + let results = (outs VMI_MaskTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$pattern attr-dict `:` type($result)"; +} + +def VMIPltOp : VMI_Op<"plt", [Pure]> { + let summary = "Data-dependent tail mask from scalar remainder"; + let description = [{ + Creates a mask where the first min(rem, L) lanes are active + and returns the remainder for the next chunk. + + The mask granularity (b8/b16/b32) is inferred from the result type. + Chaining: + + ``` + %m1, %next1 = pto.vmi.plt %rem : i32 -> !pto.vmi.mask<128>, i32 + %m2, %next2 = pto.vmi.plt %next1 : i32 -> !pto.vmi.mask<128>, i32 + ``` + + Example: + ``` + // 250 elements across 128-lane chunks + // Round 1: rem=250 → mask all-active (clamped to 128), next=122 + // Round 2: rem=122 → mask VL122, next=0 + %m1, %n1 = pto.vmi.plt %rem : i32 -> !pto.vmi.mask<128>, i32 + %m2, %n2 = pto.vmi.plt %n1 : i32 -> !pto.vmi.mask<128>, i32 + ``` + }]; + let arguments = (ins I32:$scalar); + let results = (outs VMI_MaskTypeConstraint:$mask, I32:$scalar_out); + let hasVerifier = 1; + let assemblyFormat = "$scalar attr-dict `:` type($scalar) `->` type($mask) `,` type($scalar_out)"; +} + +def VMIVaddOp : VMI_Op<"vadd", [Pure]> { + let summary = "VMI elementwise add (unified fp/int)"; + let arguments = (ins VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs, + Variadic:$mask, + OptionalAttr:$pmode); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs (`,` $mask^)? attr-dict `:` type($lhs) `,` type($rhs) (`,` type($mask)^)? `->` type($result)"; +} + +def VMIVsubOp : VMI_Op<"vsub", [Pure]> { + let summary = "VMI elementwise subtract (unified fp/int)"; + let arguments = (ins VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs, + Variadic:$mask, + OptionalAttr:$pmode); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs (`,` $mask^)? attr-dict `:` type($lhs) `,` type($rhs) (`,` type($mask)^)? `->` type($result)"; +} + +def VMIVmulOp : VMI_Op<"vmul", [Pure]> { + let summary = "VMI elementwise multiply (unified fp/int)"; + let arguments = (ins VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs, + Variadic:$mask, + OptionalAttr:$pmode); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs (`,` $mask^)? attr-dict `:` type($lhs) `,` type($rhs) (`,` type($mask)^)? `->` type($result)"; +} + +def VMIVdivOp : VMI_Op<"vdiv", [Pure]> { + let summary = "VMI elementwise divide"; + let arguments = (ins VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs, + Variadic:$mask, + OptionalAttr:$pmode); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs (`,` $mask^)? attr-dict `:` type($lhs) `,` type($rhs) (`,` type($mask)^)? `->` type($result)"; +} + +def VMIVminOp : VMI_Op<"vmin", [Pure]> { + let summary = "VMI elementwise minimum"; + let arguments = (ins VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs, + Variadic:$mask, + OptionalAttr:$pmode); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs (`,` $mask^)? attr-dict `:` type($lhs) `,` type($rhs) (`,` type($mask)^)? `->` type($result)"; +} + +def VMIVmaxOp : VMI_Op<"vmax", [Pure]> { + let summary = "VMI elementwise maximum"; + let arguments = (ins VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs, + Variadic:$mask, + OptionalAttr:$pmode); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs (`,` $mask^)? attr-dict `:` type($lhs) `,` type($rhs) (`,` type($mask)^)? `->` type($result)"; +} + +def VMIVnegOp : VMI_Op<"vneg", [Pure]> { + let summary = "VMI elementwise negate"; + let arguments = (ins VMI_VRegTypeConstraint:$source, + Variadic:$mask, + OptionalAttr:$pmode); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source (`,` $mask^)? attr-dict `:` type($source) (`,` type($mask)^)? `->` type($result)"; +} + +def VMIVabsOp : VMI_Op<"vabs", [Pure]> { + let summary = "VMI elementwise absolute value (float/int unified)"; + let description = [{ + Computes the elementwise absolute value of a logical vector. + Supports both floating-point and integer element types. + + pmode governs inactive-lane behavior: + - "merge" (default): inactive lanes pass through the source value. + - "zero": inactive lanes are zeroed. + + Example: + ``` + %r = pto.vmi.vabs %v, %m {pmode = "zero"} + : !pto.vmi.vreg<128xi32>, !pto.vmi.mask<128xpred> -> !pto.vmi.vreg<128xi32> + ``` + }]; + let arguments = (ins + VMI_VRegTypeConstraint:$source, + Variadic:$mask, + OptionalAttr:$pmode + ); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source (`,` $mask^)? attr-dict `:` type($source) (`,` type($mask)^)? `->` type($result)"; +} + +def VMIVsqrtOp : VMI_Op<"vsqrt", [Pure]> { + let summary = "VMI elementwise square root"; + let arguments = (ins VMI_VRegTypeConstraint:$source, + Variadic:$mask, + OptionalAttr:$pmode); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source (`,` $mask^)? attr-dict `:` type($source) (`,` type($mask)^)? `->` type($result)"; +} + +def VMIVexpOp : VMI_Op<"vexp", [Pure]> { + let summary = "VMI elementwise exponential"; + let arguments = (ins VMI_VRegTypeConstraint:$source, + Variadic:$mask, + OptionalAttr:$pmode); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source (`,` $mask^)? attr-dict `:` type($source) (`,` type($mask)^)? `->` type($result)"; +} + +def VMIVlnOp : VMI_Op<"vln", [Pure]> { + let summary = "VMI elementwise natural logarithm"; + let arguments = (ins VMI_VRegTypeConstraint:$source, + Variadic:$mask, + OptionalAttr:$pmode); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source (`,` $mask^)? attr-dict `:` type($source) (`,` type($mask)^)? `->` type($result)"; +} + +def VMIVreluOp : VMI_Op<"vrelu", [Pure]> { + let summary = "VMI elementwise ReLU"; + let arguments = (ins VMI_VRegTypeConstraint:$source, + Variadic:$mask, + OptionalAttr:$pmode); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source (`,` $mask^)? attr-dict `:` type($source) (`,` type($mask)^)? `->` type($result)"; +} + +def VMIVandOp : VMI_Op<"vand", [Pure]> { + let summary = "VMI elementwise bitwise and"; + let arguments = (ins VMI_ValueTypeConstraint:$lhs, VMI_ValueTypeConstraint:$rhs, + Variadic:$mask, + OptionalAttr:$pmode); + let results = (outs VMI_ValueTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs (`,` $mask^)? attr-dict `:` type($lhs) `,` type($rhs) (`,` type($mask)^)? `->` type($result)"; +} + +def VMIVorOp : VMI_Op<"vor", [Pure]> { + let summary = "VMI elementwise bitwise or"; + let arguments = (ins VMI_ValueTypeConstraint:$lhs, VMI_ValueTypeConstraint:$rhs, + Variadic:$mask, + OptionalAttr:$pmode); + let results = (outs VMI_ValueTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs (`,` $mask^)? attr-dict `:` type($lhs) `,` type($rhs) (`,` type($mask)^)? `->` type($result)"; +} + +def VMIVxorOp : VMI_Op<"vxor", [Pure]> { + let summary = "VMI elementwise bitwise xor"; + let arguments = (ins VMI_ValueTypeConstraint:$lhs, VMI_ValueTypeConstraint:$rhs, + Variadic:$mask, + OptionalAttr:$pmode); + let results = (outs VMI_ValueTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs (`,` $mask^)? attr-dict `:` type($lhs) `,` type($rhs) (`,` type($mask)^)? `->` type($result)"; +} + +def VMIVshlOp : VMI_Op<"vshl", [Pure]> { + let summary = "VMI elementwise left shift"; + let arguments = (ins VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs, + Variadic:$mask, + OptionalAttr:$pmode); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs (`,` $mask^)? attr-dict `:` type($lhs) `,` type($rhs) (`,` type($mask)^)? `->` type($result)"; +} + +def VMIVshrOp : VMI_Op<"vshr", [Pure]> { + let summary = "VMI signedness-aware elementwise right shift"; + let arguments = (ins VMI_VRegTypeConstraint:$lhs, VMI_VRegTypeConstraint:$rhs, + Variadic:$mask, + OptionalAttr:$pmode); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$lhs `,` $rhs (`,` $mask^)? attr-dict `:` type($lhs) `,` type($rhs) (`,` type($mask)^)? `->` type($result)"; +} + +def VMIVnotOp : VMI_Op<"vnot", [Pure]> { + let summary = "VMI elementwise bitwise not"; + let arguments = (ins VMI_ValueTypeConstraint:$source, + Variadic:$mask, + OptionalAttr:$pmode); + let results = (outs VMI_ValueTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source (`,` $mask^)? attr-dict `:` type($source) (`,` type($mask)^)? `->` type($result)"; +} + +def VMIVcmpOp : VMI_Op<"vcmp", [Pure]> { + let summary = "VMI elementwise compare (fp/int unified) → predicate mask"; + let description = [{ + Compares two logical vector registers elementwise and produces a predicate + mask. The comparison mode is given by the `cmp` attribute. The `seed` + operand provides the governing predicate Pg: where seed[i]=0 the result + lane is zero (pmode="zeroing", default), and where seed[i]=1 the comparison + result is written. Floating-point comparisons support eq/ne/lt/le/gt/ge + and ordered forms oeq/one/olt/ole/ogt/oge. Integer comparisons use + eq/ne/lt/le/gt/ge; signedness is selected by the integer element type + (signed/signless versus unsigned). + + Example: + ``` + %r = pto.vmi.vcmp %a, %b, %seed {cmp = "lt"} + : !pto.vmi.vreg<128xf32>, !pto.vmi.vreg<128xf32>, !pto.vmi.mask<128xb32> + -> !pto.vmi.mask<128xb32> + ``` + }]; + + let arguments = (ins + VMI_VRegTypeConstraint:$lhs, + VMI_VRegTypeConstraint:$rhs, + VMI_MaskTypeConstraint:$seed, + StrAttr:$cmp, + OptionalAttr:$pmode + ); + let results = (outs VMI_MaskTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = [{ + $lhs `,` $rhs `,` $seed attr-dict `:` type($lhs) `,` type($rhs) `,` type($seed) `->` type($result) + }]; +} + +def VMIVcmpsOp : VMI_Op<"vcmps", [Pure]> { + let summary = "VMI elementwise vector-scalar compare → predicate mask"; + let description = [{ + Compares each lane of a logical vector against a scalar and produces a + predicate mask. The scalar type must match the vector element type. The + `seed` operand provides the governing predicate Pg. + + Example: + ``` + %r = pto.vmi.vcmps %v, %s, %seed {cmp = "ge"} + : !pto.vmi.vreg<64xf16>, f16, !pto.vmi.mask<64xpred> + -> !pto.vmi.mask<64xpred> + ``` + }]; + + let arguments = (ins + VMI_VRegTypeConstraint:$src, + AnyType:$scalar, + VMI_MaskTypeConstraint:$seed, + StrAttr:$cmp, + OptionalAttr:$pmode + ); + let results = (outs VMI_MaskTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = [{ + $src `,` $scalar `,` $seed attr-dict `:` type($src) `,` type($scalar) `,` type($seed) `->` type($result) + }]; +} + +class VMI_VecScalarOp + : VMI_Op { + let summary = summaryText; + let description = [{ + Performs elementwise operation between a logical vector and a scalar. + The scalar is implicitly broadcast to all lanes (R6/M2). + + pmode governs inactive-lane behavior on the result: + - "merge" (default): inactive lanes pass through the source value. + - "zero": inactive lanes are zeroed. + + Example: + ``` + %r = pto.vmi.} # mnemonic # [{ %v, %s, %m {pmode = "merge"} + : !pto.vmi.vreg<64xf16>, f16, !pto.vmi.mask<64xpred> -> !pto.vmi.vreg<64xf16> + ``` + }]; + let arguments = (ins + VMI_VRegTypeConstraint:$src, + AnyType:$scalar, + VMI_MaskTypeConstraint:$mask, + OptionalAttr:$pmode + ); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = [{ + $src `,` $scalar `,` $mask attr-dict `:` type($src) `,` type($scalar) `,` type($mask) `->` type($result) + }]; +} + +def VMIAddSOp : VMI_VecScalarOp<"vadds", "VMI vector-scalar elementwise add">; +def VMIMulSOp : VMI_VecScalarOp<"vmuls", "VMI vector-scalar elementwise multiply">; +def VMIMaxSOp : VMI_VecScalarOp<"vmaxs", "VMI vector-scalar elementwise maximum">; +def VMIMinSOp : VMI_VecScalarOp<"vmins", "VMI vector-scalar elementwise minimum">; +def VMIShlSOp : VMI_VecScalarOp<"vshls", "VMI vector-scalar elementwise shift left">; +def VMIShrSOp : VMI_VecScalarOp<"vshrs", "VMI vector-scalar elementwise shift right">; + +def VMIvSelOp : VMI_Op<"vsel", [Pure]> { + let summary = "VMI elementwise select"; + let description = [{ + Per-lane selection driven by a predicate mask. For each lane i: + + result[i] = mask[i] ? true_value[i] : false_value[i] + + pmode governs inactive-lane behavior on the result: + - "merge" (default): inactive lanes retain the false_value element. + - "zero": inactive lanes are zeroed. + + Example: + ```mlir + %r = pto.vmi.vsel %m, %a, %b {pmode = "zero"} + : !pto.vmi.mask<256xpred>, !pto.vmi.vreg<256xui16>, !pto.vmi.vreg<256xui16> -> !pto.vmi.vreg<256xui16> + ``` + }]; + let arguments = (ins VMI_MaskTypeConstraint:$mask, VMI_VRegTypeConstraint:$true_value, + VMI_VRegTypeConstraint:$false_value, + OptionalAttr:$pmode); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$mask `,` $true_value `,` $false_value attr-dict `:` type($mask) `,` type($true_value) `,` type($false_value) `->` type($result)"; +} + +def VMIvcaddOp : VMI_Op<"vcadd", [Pure]> { + let summary = "VMI add-reduction: sum lanes to a scalar or compact group vector"; + let description = [{ + Reduces source lanes by summation, governed by a mask. + + - No init operand: inactive lanes contribute 0 (semantically nil in + addition). Result reaches a single scalar when `group` is absent, + or `C` compact scalars when `{group=C}` is specified. + - Floating-point sources MUST carry the `reassoc` attribute (hardware + vcadd uses pair-wise reassociation). + - `pmode` controls inactive-result behaviour: "zero" (default) writes 0; + "merge" preserves the destination register's old value. + }]; + let arguments = (ins + VMI_VRegTypeConstraint:$source, + VMI_MaskTypeConstraint:$mask, + OptionalAttr:$group, + OptionalAttr:$pmode, + OptionalAttr:$reassoc + ); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `,` $mask attr-dict `:` type($source) `,` type($mask) `->` type($result)"; +} + +def VMIvcmaxOp : VMI_Op<"vcmax", [Pure]> { + let summary = "VMI max-reduction: maximum across lanes to a scalar or compact group vector"; + let description = [{ + Reduces source lanes by taking the element-wise maximum, governed by a mask. + + - Inactive lanes are treated as -INF (floating-point) or the type's + minimum representable value (integer). + - `pmode` controls inactive-result behaviour as for `vcadd`. + }]; + let arguments = (ins + VMI_VRegTypeConstraint:$source, + VMI_MaskTypeConstraint:$mask, + OptionalAttr:$group, + OptionalAttr:$pmode + ); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `,` $mask attr-dict `:` type($source) `,` type($mask) `->` type($result)"; +} + +def VMIvcminOp : VMI_Op<"vcmin", [Pure]> { + let summary = "VMI min-reduction: minimum across lanes to a scalar or compact group vector"; + let description = [{ + Reduces source lanes by taking the element-wise minimum, governed by a mask. + + - Inactive lanes are treated as +INF (floating-point) or the type's + maximum representable value (integer). + - `pmode` controls inactive-result behaviour as for `vcadd`. + }]; + let arguments = (ins + VMI_VRegTypeConstraint:$source, + VMI_MaskTypeConstraint:$mask, + OptionalAttr:$group, + OptionalAttr:$pmode + ); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `,` $mask attr-dict `:` type($source) `,` type($mask) `->` type($result)"; +} + +def VMICvtOp : VMI_Op<"vcvt", [Pure]> { + let summary = "VMI unified elementwise type conversion"; + let description = [{ + Converts a VMI logical vector from one element type to another. + The conversion direction and semantics are derived from the source and result + element types: + + - fp → fp, |dst| > |src|: floating-point widening (replaces extf) + - fp → fp, |dst| < |src|: floating-point narrowing (replaces truncf) + - fp → int: float to signed integer (replaces fptosi) + - int → fp: signed integer to float (replaces sitofp) + - int → int, |dst| > |src|: integer extension (replaces extsi/extui) + - int → int, |dst| < |src|: integer truncation (replaces trunci) + + Attributes: + - `rounding`: rounding mode for fp narrowing (A=away-from-zero, + H=half-up). Valid only when dst bit-width < src bit-width for fp types. + - `saturate`: saturating behavior on overflow ("SAT"). Valid for any + narrowing conversion (fp or int). + - `pmode`: predication mode ("merge" | "zero"). + + For int→int widening, the source element type must carry signedness + (e.g. si8/ui8/si16/ui16); signless integers are rejected. + + Example: + ```mlir + %w = pto.vmi.vcvt %x : !pto.vmi.vreg<64xf16> -> !pto.vmi.vreg<64xf32> + %n = pto.vmi.vcvt %y {rounding="H"} : !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xf16> + %z = pto.vmi.vcvt %a : !pto.vmi.vreg<256xui8> -> !pto.vmi.vreg<256xui16> + ``` + }]; + let arguments = (ins VMI_VRegTypeConstraint:$source, + OptionalAttr:$rounding, + OptionalAttr:$saturate, + OptionalAttr:$pmode); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source attr-dict `:` type($source) `->` type($result)"; +} + +def VMIvLoadOp : VMI_Op<"vload", + [DeclareOpInterfaceMethods, + AttrSizedOperandSegments]> { + let summary = "VMI unified logical vector load"; + let description = [{ + Unified vector load. dist-mode selects the memory-access pattern. + + dist-mode values: + - "continuous" (default): contiguous stride-1 load → 1 result + - "dintlv": deinterleaved dual load → 2 results (%lo, %hi) + - "unpack": widening unpack load → 1 result + - "brc": broadcast load → 1 result + + group mode (mutually exclusive with dist-mode): + - {group = C}: grouped vector load with row stride → 1 result + + block-stride mode (mutually exclusive with dist-mode and group): + - {block_stride = B, repeat_stride = R}: block-strided masked load → 1 result + + pmode ("zero"|"merge") governs inactive-lane behavior on the result. + vload itself has no mask operand; A5 loads are not predicable — masks + migrate to consumer ops or vstore. + }]; + let arguments = (ins + PtrOrMemRef:$source, + Index:$offset, + Optional:$stride, + Optional:$block_stride, + Optional:$repeat_stride, + OptionalAttr:$dist_mode, + OptionalAttr:$group, + OptionalAttr:$pmode + ); + let results = (outs Variadic:$results); + let hasCustomAssemblyFormat = 1; + let hasVerifier = 1; +} + +def VMIvStoreOp : VMI_Op<"vstore", + [DeclareOpInterfaceMethods, + AttrSizedOperandSegments]> { + let summary = "VMI unified logical vector store"; + let description = [{ + Unified vector store replacing the individual store, interleave_store, + masked_store, stride_store, and group_store ops. + + dist-mode controls the memory access pattern: + - "continuous" (default): contiguous stride-1 store → 1 value + - "dintlv": interleaved dual store → 2 values (%lo, %hi) + + group mode (mutually exclusive with dist-mode): + - {group = C}: grouped vector store with row stride + + block-stride mode (mutually exclusive with dist-mode and group): + - {block_stride = B, repeat_stride = R}: block-strided masked store + + pmode governs inactive lane behavior: + - "zero" (default): inactive lanes store 0 + - "merge": inactive lanes skip write (needs mask) + + mask is variadic: 0 or 1 mask operand. + }]; + let arguments = (ins + Variadic:$values, + PtrOrMemRef:$destination, + Index:$offset, + Optional:$stride, + Optional:$block_stride, + Optional:$repeat_stride, + Variadic:$mask, + OptionalAttr:$dist_mode, + OptionalAttr:$group, + OptionalAttr:$pmode + ); + let results = (outs); + let hasCustomAssemblyFormat = 1; + let hasVerifier = 1; +} + +def VMIVinterpretCastOp : VMI_Op<"vinterpret_cast", [Pure]> { + let summary = "VMI bitwise vector reinterpretation"; + let arguments = (ins VMI_VRegTypeConstraint:$source); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source attr-dict `:` type($source) `->` type($result)"; +} + +def VMIVselrOp : VMI_Op<"vselr", [Pure]> { + let summary = "VMI dynamic lane select"; + let description = [{ + Dynamic lane permutation: for each lane i, result[i] = source[index[i]]. + Replaces the static-index `shuffle` op with a dynamic index vector operand. + Lowers 1:1 to `pto.vselr` at the VPTO level. + + Example: + ``` + %r = pto.vmi.vselr %src, %idx : !pto.vmi.vreg<64xf16>, !pto.vmi.vreg<4xi16> -> !pto.vmi.vreg<4xf16> + ``` + }]; + let arguments = (ins + VMI_VRegTypeConstraint:$source, + VMI_VRegTypeConstraint:$index + ); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = [{ + $source `,` $index attr-dict `:` type($source) `,` type($index) `->` type($result) + }]; +} + +def VMIVintlvOp : VMI_Op<"vintlv"> { + let summary = "VMI interleave two vectors lane-by-lane (Category A: layout-transparent)"; + let arguments = (ins + VMI_VRegTypeConstraint:$lhs, + VMI_VRegTypeConstraint:$rhs, + VMI_MaskTypeConstraint:$mask, + OptionalAttr:$pmode + ); + let results = (outs + VMI_VRegTypeConstraint:$low, + VMI_VRegTypeConstraint:$high + ); + let hasVerifier = 1; + let assemblyFormat = [{ + $lhs `,` $rhs `,` $mask attr-dict `:` type($lhs) `,` type($rhs) `,` type($mask) `->` type($low) `,` type($high) + }]; +} + +def VMIVdintlvOp : VMI_Op<"vdintlv"> { + let summary = "VMI deinterleave two vectors by even/odd lanes (Category A: layout-transparent)"; + let arguments = (ins + VMI_VRegTypeConstraint:$lhs, + VMI_VRegTypeConstraint:$rhs, + VMI_MaskTypeConstraint:$mask, + OptionalAttr:$pmode + ); + let results = (outs + VMI_VRegTypeConstraint:$low, + VMI_VRegTypeConstraint:$high + ); + let hasVerifier = 1; + let assemblyFormat = [{ + $lhs `,` $rhs `,` $mask attr-dict `:` type($lhs) `,` type($rhs) `,` type($mask) `->` type($low) `,` type($high) + }]; +} + +def VMIVgatherOp : VMI_Op<"vgather", [DeclareOpInterfaceMethods]> { + let summary = "VMI logical masked indexed gather (B32 granularity)"; + let description = [{ + Gathers elements from UB memory using an index vector. + Inactive lanes are controlled by pmode (no explicit passthru operand). + + pmode controls inactive-lane behavior: + - "zero" (default): inactive lanes are zeroed. + - "merge": inactive lanes preserve the prior destination value. + + Example: + ``` + %g = pto.vmi.vgather %src, %offsets, %mask + : !pto.ptr, !pto.vmi.vreg<64xi32>, !pto.vmi.mask<64> -> !pto.vmi.vreg<64xf32> + ``` + }]; + let arguments = (ins + PtrOrMemRef:$source, + VMI_VRegTypeConstraint:$offsets, + VMI_MaskTypeConstraint:$mask, + OptionalAttr:$pmode + ); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `,` $offsets `,` $mask attr-dict `:` type($source) `,` type($offsets) `,` type($mask) `->` type($result)"; +} + +def VMIVgatherbOp : VMI_Op<"vgatherb", [DeclareOpInterfaceMethods]> { + let summary = "VMI logical masked byte-granularity indexed gather"; + let description = [{ + Gathers elements from UB memory using byte-offset indices. + Mask lane count equals result lane count (may differ from offset lane count). + + Example: + ``` + %gb = pto.vmi.vgatherb %src, %offsets, %mask + : !pto.ptr, !pto.vmi.vreg<64xi32>, !pto.vmi.mask<256> -> !pto.vmi.vreg<256xi32> + ``` + }]; + let arguments = (ins + PtrOrMemRef:$source, + VMI_VRegTypeConstraint:$offsets, + VMI_MaskTypeConstraint:$mask, + OptionalAttr:$pmode + ); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$source `,` $offsets `,` $mask attr-dict `:` type($source) `,` type($offsets) `,` type($mask) `->` type($result)"; +} + +def VMIVscatterOp : VMI_Op<"vscatter", [DeclareOpInterfaceMethods]> { + let summary = "VMI logical masked indexed scatter"; + let description = [{ + Scatters elements to UB memory using an index vector. + Inactive lanes are controlled by pmode. + + Example: + ``` + pto.vmi.vscatter %v, %dest, %offsets, %mask + : !pto.vmi.vreg<64xf32>, !pto.ptr, !pto.vmi.vreg<64xi32>, !pto.vmi.mask<64> + ``` + }]; + let arguments = (ins + VMI_VRegTypeConstraint:$value, + PtrOrMemRef:$destination, + VMI_VRegTypeConstraint:$offsets, + VMI_MaskTypeConstraint:$mask, + OptionalAttr:$pmode + ); + let results = (outs); + let hasVerifier = 1; + let assemblyFormat = "$value `,` $destination `,` $offsets `,` $mask attr-dict `:` type($value) `,` type($destination) `,` type($offsets) `,` type($mask)"; +} + +def VMIVexpdifOp : VMI_Op<"vexpdif", [Pure]> { + let summary = "VMI fused exp(x - max) for softmax numerical stability"; + let description = [{ + Computes exp(x - max) in a single hardware instruction. + x may be f16 or f32; max and result are always f32 (hardware constraint). + + Example: + ``` + %e = pto.vmi.vexpdif %x, %max, %mask + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32>, !pto.vmi.mask<64> -> !pto.vmi.vreg<64xf32> + ``` + }]; + let arguments = (ins + VMI_VRegTypeConstraint:$x, + VMI_VRegTypeConstraint:$max, + VMI_MaskTypeConstraint:$mask, + OptionalAttr:$pmode + ); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$x `,` $max `,` $mask attr-dict `:` type($x) `,` type($max) `,` type($mask) `->` type($result)"; +} + +def VMIVaxpyOp : VMI_Op<"vaxpy", [Pure]> { + let summary = "VMI fused alpha*x + y (scale-add)"; + let description = [{ + Computes y = alpha * x + y in a single hardware instruction. + + Example: + ``` + %y = pto.vmi.vaxpy %x, %acc, %alpha, %mask + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32>, f32, !pto.vmi.mask<64> -> !pto.vmi.vreg<64xf32> + ``` + }]; + let arguments = (ins + VMI_VRegTypeConstraint:$x, + VMI_VRegTypeConstraint:$acc, + AnyTypeOf<[AnyFloat]>:$alpha, + VMI_MaskTypeConstraint:$mask, + OptionalAttr:$pmode + ); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; +} + +def VMIVlreluOp : VMI_Op<"vlrelu", [Pure]> { + let summary = "VMI leaky ReLU activation"; + let description = [{ + Computes y = x > 0 ? x : slope * x using a single hardware instruction. + The slope is a scalar shared across all lanes. + + Example: + ``` + %lr = pto.vmi.vlrelu %x, %slope, %mask + : !pto.vmi.vreg<64xf32>, f32, !pto.vmi.mask<64> -> !pto.vmi.vreg<64xf32> + ``` + }]; + let arguments = (ins + VMI_VRegTypeConstraint:$x, + AnyTypeOf<[AnyFloat]>:$slope, + VMI_MaskTypeConstraint:$mask, + OptionalAttr:$pmode + ); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; +} + +def VMIVpreluOp : VMI_Op<"vprelu", [Pure]> { + let summary = "VMI parametric ReLU activation"; + let description = [{ + Computes y = max(x, 0) + alpha * min(x, 0) where alpha is a per-lane + parameter vector (not a shared scalar). + + Example: + ``` + %pr = pto.vmi.vprelu %x, %alpha, %mask + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32>, !pto.vmi.mask<64> -> !pto.vmi.vreg<64xf32> + ``` + }]; + let arguments = (ins + VMI_VRegTypeConstraint:$x, + VMI_VRegTypeConstraint:$alpha, + VMI_MaskTypeConstraint:$mask, + OptionalAttr:$pmode + ); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$x `,` $alpha `,` $mask attr-dict `:` type($x) `,` type($alpha) `,` type($mask) `->` type($result)"; +} + +def VMIVmullOp : VMI_Op<"vmull", [Pure]> { + let summary = "VMI widening 32x32 multiply with low/high results"; + let description = [{ + Computes a per-lane 32-bit x 32-bit product and returns its low and high + 32-bit halves. Inactive lanes in both results are zero. The initial + contract accepts exactly 64, 128, or 256 lanes of i32 or ui32 data. + + Example: + ``` + %low, %high = pto.vmi.vmull %a, %b, %mask + : !pto.vmi.vreg<64xi32>, !pto.vmi.vreg<64xi32>, + !pto.vmi.mask<64xpred> + -> !pto.vmi.vreg<64xi32>, !pto.vmi.vreg<64xi32> + ``` + }]; + let arguments = (ins + VMI_VRegTypeConstraint:$a, + VMI_VRegTypeConstraint:$b, + VMI_MaskTypeConstraint:$mask, + OptionalAttr:$pmode + ); + let results = (outs + VMI_VRegTypeConstraint:$low, + VMI_VRegTypeConstraint:$high + ); + let hasVerifier = 1; + let assemblyFormat = "$a `,` $b `,` $mask attr-dict `:` type($a) `,` type($b) `,` type($mask) `->` type($low) `,` type($high)"; +} + +def VMIVmulaOp : VMI_Op<"vmula", [Pure]> { + let summary = "VMI fused multiply-add (acc = acc + lhs * rhs)"; + let description = [{ + Computes acc = acc + lhs * rhs in a single hardware instruction. + The accumulator operand comes first to match hardware FMA semantics. + + Supports both floating-point (f16/bf16/f32) and integer (i8-i32) types. + The mask operand is optional; when present it is only a semantic annotation + and is discarded during lowering to legacy fma. + + Example: + ``` + %acc1 = pto.vmi.vmula %acc, %a, %b, %mask + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32>, !pto.vmi.mask<64> + -> !pto.vmi.vreg<64xf32> + + // Without mask: + %acc2 = pto.vmi.vmula %acc, %a, %b + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> + -> !pto.vmi.vreg<64xf32> + ``` + }]; + let arguments = (ins + VMI_VRegTypeConstraint:$acc, + VMI_VRegTypeConstraint:$lhs, + VMI_VRegTypeConstraint:$rhs, + Variadic:$mask, + OptionalAttr:$pmode + ); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$acc `,` $lhs `,` $rhs (`,` $mask^)? attr-dict `:` type($acc) `,` type($lhs) `,` type($rhs) (`,` type($mask)^)? `->` type($result)"; +} +#endif // MLIR_DIALECT_PTO_IR_VMIOPS diff --git a/include/PTO/IR/VMITypeDefs.td b/include/PTO/IR/VMITypeDefs.td new file mode 100644 index 0000000000..4ec6bb5009 --- /dev/null +++ b/include/PTO/IR/VMITypeDefs.td @@ -0,0 +1,67 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +//===- VMITypeDefs.td - PTO VMI type definitions -----------*- tablegen -*-===// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_DIALECT_PTO_IR_VMITYPEDEFS +#define MLIR_DIALECT_PTO_IR_VMITYPEDEFS + +include "PTO/IR/PTODialect.td" +include "PTO/IR/PTOAttrs.td" + +def VMIVRegType : TypeDef { + let mnemonic = "vmi.vreg"; + let summary = "A VMI logical vector register value"; + + let parameters = (ins + "int64_t":$elementCount, + "Type":$elementType, + "mlir::Attribute":$layout + ); + + let hasCustomAssemblyFormat = 1; + let genVerifyDecl = 1; + + let extraClassDeclaration = [{ + bool hasLayout() const { return static_cast(getLayout()); } + VMILayoutAttr getLayoutAttr() const { + return ::llvm::dyn_cast_or_null(getLayout()); + } + }]; +} + +def VMIMaskType : TypeDef { + let mnemonic = "vmi.mask"; + let summary = "A VMI logical predicate mask value"; + + let parameters = (ins + "int64_t":$elementCount, + StringRefParameter<"mask granularity view">:$granularity, + "mlir::Attribute":$layout + ); + + let hasCustomAssemblyFormat = 1; + let genVerifyDecl = 1; + + let extraClassDeclaration = [{ + static bool isSupportedGranularity(::llvm::StringRef granularity); + static bool isConcreteGranularity(::llvm::StringRef granularity); + + bool hasLayout() const { return static_cast(getLayout()); } + bool isPred() const { return getGranularity() == "pred"; } + bool isB8() const { return getGranularity() == "b8"; } + bool isB16() const { return getGranularity() == "b16"; } + bool isB32() const { return getGranularity() == "b32"; } + VMILayoutAttr getLayoutAttr() const { + return ::llvm::dyn_cast_or_null(getLayout()); + } + }]; +} + +#endif // MLIR_DIALECT_PTO_IR_VMITYPEDEFS diff --git a/include/PTO/IR/VMIUtils.h b/include/PTO/IR/VMIUtils.h new file mode 100644 index 0000000000..e55e558034 --- /dev/null +++ b/include/PTO/IR/VMIUtils.h @@ -0,0 +1,53 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +//===- VMIUtils.h - PTO VMI shared helpers ----------------------*- C++ -*-===// +//===----------------------------------------------------------------------===// + +#ifndef PTO_IR_VMIUTILS_H +#define PTO_IR_VMIUTILS_H + +#include "PTO/IR/PTO.h" +#include "mlir/Support/LLVM.h" +#include "mlir/Support/LogicalResult.h" + +namespace mlir::pto { + +inline constexpr StringLiteral kVMIDiagUnsupported = "VMI-UNSUPPORTED"; +inline constexpr StringLiteral kVMIDiagLayoutContract = + "VMI-LAYOUT-CONTRACT"; +inline constexpr StringLiteral kVMIDiagPassInvariant = "VMI-PASS-INVARIANT"; +inline constexpr StringLiteral kVMIDiagResidualOp = "VMI-RESIDUAL-OP"; + +inline constexpr StringLiteral kVMIDiagUnsupportedPrefix = + "VMI-UNSUPPORTED: "; +inline constexpr StringLiteral kVMIDiagLayoutContractPrefix = + "VMI-LAYOUT-CONTRACT: "; +inline constexpr StringLiteral kVMIDiagPassInvariantPrefix = + "VMI-PASS-INVARIANT: "; +inline constexpr StringLiteral kVMIDiagResidualOpPrefix = "VMI-RESIDUAL-OP: "; + +struct VMIPhysicalLane { + int64_t part = 0; + int64_t chunk = 0; + int64_t lane = 0; +}; + +FailureOr getDataLanesPerPart(Type elementType); +FailureOr getMaskLanesPerPart(StringRef granularity); +FailureOr getVMIPhysicalArity(Type type); +FailureOr mapLogicalLaneToPhysical(Type type, + int64_t logicalLane); +FailureOr mapPhysicalLaneToLogical(Type type, int64_t part, + int64_t chunk, int64_t lane); +FailureOr isPaddingLane(Type type, int64_t part, int64_t chunk, + int64_t lane); + +} // namespace mlir::pto + +#endif // PTO_IR_VMIUTILS_H diff --git a/include/PTO/Transforms/Passes.h b/include/PTO/Transforms/Passes.h index cdbd8e1c0a..10f9756cda 100644 --- a/include/PTO/Transforms/Passes.h +++ b/include/PTO/Transforms/Passes.h @@ -93,6 +93,7 @@ std::unique_ptr createPTOFusionRegionGenPass(); LogicalResult validateIntToPtrUses(func::FuncOp func); std::unique_ptr createPTOUnrollSIMTForPass(); +std::unique_ptr createPTONarrowVPTOLoopCountersPass(); std::unique_ptr createPTOInferVPTOVecScopePass(); std::unique_ptr createVPTOExpandWrapperOpsPass(); std::unique_ptr createPTOVPTOPtrBoundaryPass(); @@ -103,12 +104,29 @@ std::unique_ptr createPTOFusionLoadStoreElisionPass(); std::unique_ptr createPTOFlattenFusionRegionPass(); std::unique_ptr createVPTOPtrNormalizePass(); std::unique_ptr createVPTOPtrCastCleanupPass(); +std::unique_ptr createVPTONormalizeEquivalentVcvtPass(); LogicalResult validateVPTOAuthoringIR(ModuleOp module, llvm::raw_ostream *diagOS = nullptr); LogicalResult validateVPTOEmissionIR(ModuleOp module, llvm::raw_ostream *diagOS = nullptr); std::unique_ptr createPTOValidateVPTOIRPass(); std::unique_ptr createPTOValidateVPTOEmissionIRPass(); +LogicalResult validateVMIProducerBoundaryIR(ModuleOp module, + llvm::raw_ostream *diagOS = nullptr); +LogicalResult validateVMILayoutAssignedIR(ModuleOp module, + llvm::raw_ostream *diagOS = nullptr, + bool verifyHelperSupport = true); +std::unique_ptr createPTOValidateVMIIRPass(); +std::unique_ptr createPTOValidateVMILayoutIRPass(); +std::unique_ptr createVMIPreAssignmentCombinePass(); +std::unique_ptr createVMIMaskGranularityAssignmentPass(); +std::unique_ptr createVMILayoutAssignmentPass(); +std::unique_ptr createVMILayoutFoldPass(); +std::unique_ptr createVMILayoutRematerializePass(); +std::unique_ptr createVMILayoutSinkMaterializationPass(); +std::unique_ptr createVMILegalizeArithSelectPass(); +std::unique_ptr createVMILowerUnifiedToLegacyPass(); +std::unique_ptr createVMIToVPTOPass(); std::unique_ptr createInsertTemplateAttributesPass(); std::unique_ptr createInsertTemplateAttributesPass( const InsertTemplateAttributesOptions &options); diff --git a/include/PTO/Transforms/Passes.td b/include/PTO/Transforms/Passes.td index 61e6e4b00d..6b4bd3b222 100644 --- a/include/PTO/Transforms/Passes.td +++ b/include/PTO/Transforms/Passes.td @@ -828,17 +828,47 @@ def PTOUnrollSIMTFor : Pass<"pto-unroll-simt-for", "func::FuncOp"> { ]; } +def PTONarrowVPTOLoopCounters + : Pass<"pto-narrow-vpto-loop-counters", "func::FuncOp"> { + let summary = + "Narrow constant-bounded scf.for counters inside VPTO vecscope regions to i16"; + let description = [{ + Rewrites `scf.for` operations nested under `pto.vecscope` or + `pto.strict_vecscope` when the lower bound, upper bound, positive step, and + post-loop induction value are compile-time integer constants representable + as signed i16 values. + + The rewritten loop uses i16 bounds and induction variable. The induction + variable is cast back to its original index or wider integer type at the + start of the body, preserving the authored body and result types while + exposing the narrow counter form expected by the A5 hardware-loop backend. + + Loops with dynamic bounds, out-of-range constants, an overflowing final + increment, non-positive steps, counters already 16 bits or narrower, or + loops outside VPTO vecscope regions are left unchanged. + }]; + let constructor = + "mlir::pto::createPTONarrowVPTOLoopCountersPass()"; + let dependentDialects = [ + "mlir::func::FuncDialect", + "mlir::scf::SCFDialect", + "mlir::arith::ArithDialect", + "mlir::pto::PTODialect" + ]; +} + def PTOInferVPTOVecScope : Pass<"pto-infer-vpto-vecscope", "func::FuncOp"> { let summary = "Infer missing pto.vecscope regions for VPTO vector operation clusters"; let description = [{ - Runs near the VPTO emission boundary after inlining, canonicalization, - CSE, pointer normalization, and wrapper-op expansion have exposed the final - SSA shape. The pass greedily clusters contiguous VPTO vector operations - into `pto.vecscope` regions while preserving explicit vector-scope carriers - and treating DMA/copy/sync, unresolved calls, terminators, and forbidden - operations as boundaries. + Runs near the VPTO emission boundary after VMI physicalization and the + existing pre-emission canonicalization, pointer normalization, and + wrapper-op expansion, but before VMI LICM and final cleanup. The pass + greedily clusters contiguous VPTO vector operations into `pto.vecscope` + regions while preserving explicit vector-scope carriers and treating + DMA/copy/sync, unresolved calls, terminators, and forbidden operations as + boundaries. The inferred `pto.vecscope` form remains resultless. Values whose type is `!pto.vreg`, `!pto.mask`, or `!pto.align` must not escape the inferred @@ -868,6 +898,222 @@ def PTOValidateVPTOIR : Pass<"pto-validate-vpto-ir", "ModuleOp"> { "mlir::scf::SCFDialect"]; } +def PTOValidateVMIIR : Pass<"pto-validate-vmi-ir", "ModuleOp"> { + let summary = "Validate VMI producer-boundary semantic IR"; + let description = [{ + Checks that VMI producer-boundary IR uses only surface VMI data/mask types, + native pto.vmi semantic ops, and structural control-flow/function ops. This + pass runs before layout assignment, so layout-assigned VMI types, VMI helper + ops, and physical VPTO register types are rejected. + }]; + let constructor = "mlir::pto::createPTOValidateVMIIRPass()"; + let dependentDialects = ["mlir::cf::ControlFlowDialect", + "mlir::func::FuncDialect", + "mlir::pto::PTODialect", + "mlir::memref::MemRefDialect", + "mlir::scf::SCFDialect"]; +} + +def PTOValidateVMILayoutIR + : Pass<"pto-validate-vmi-layout-ir", "ModuleOp"> { + let summary = "Validate layout-assigned VMI IR"; + let description = [{ + Checks the post-layout-assignment VMI stage: every VMI data value must have + a concrete VMI layout, every VMI mask must have concrete b8/b16/b32 + granularity and layout, physical VPTO register values must not appear yet, + and VMI typed values must stay inside VMI semantic/helper or structural ops. + vmi-to-vpto chooses deterministic lowerings from the current op's attrs, + operand/result types, layouts, and operand values. Non-local choices must + be represented as explicit attrs, helper ops, or diagnostics before this + stage. Later VMI layout optimization passes may replace helpers with + cloned/rematerialized producers, but the layout gate must not depend on + hidden producer/user context. + }]; + let constructor = "mlir::pto::createPTOValidateVMILayoutIRPass()"; + let dependentDialects = ["mlir::cf::ControlFlowDialect", + "mlir::func::FuncDialect", + "mlir::pto::PTODialect", + "mlir::memref::MemRefDialect", + "mlir::scf::SCFDialect"]; +} + +def VMIPreAssignmentCombine + : Pass<"vmi-pre-assignment-combine", "ModuleOp"> { + let summary = "Combine VMI operations before layout assignment"; + let description = [{ + Performs VMI-level structural combines before VMI layout assignment. This + keeps layout assignment focused on choosing and materializing layouts while + still exposing direct semantic operations to the later layout and lowering + passes. + + The pass currently rewrites the semantic pattern + `group_broadcast(group_slot_load(...))` into the equivalent + `group_broadcast_load` operation. + }]; + let constructor = "mlir::pto::createVMIPreAssignmentCombinePass()"; + let dependentDialects = ["mlir::func::FuncDialect", + "mlir::pto::PTODialect"]; +} + +def VMILayoutAssignment : Pass<"vmi-layout-assignment", "ModuleOp"> { + let summary = "Assign concrete VMI layouts"; + let description = [{ + Solves VMI layout constraints and materializes the chosen layout into VMI + types. Mask granularity is assigned by vmi-mask-granularity-assignment + before this pass, so this pass only propagates and materializes mask + layouts. + }]; + let constructor = "mlir::pto::createVMILayoutAssignmentPass()"; + let dependentDialects = ["mlir::cf::ControlFlowDialect", + "mlir::func::FuncDialect", + "mlir::pto::PTODialect", + "mlir::memref::MemRefDialect", + "mlir::scf::SCFDialect"]; +} + +def VMIMaskGranularityAssignment + : Pass<"vmi-mask-granularity-assignment", "ModuleOp"> { + let summary = "Assign concrete VMI mask granularities"; + let description = [{ + Assigns b8/b16/b32 granularity to VMI mask values before layout + assignment. This pass follows VMI semantic op contracts where mask + granularity is determined by the associated data element width. It does + not choose layouts; conflicting use granularities are represented by + rematerializing cheap mask producers when possible, otherwise with + pto.vmi.ensure_mask_granularity. + }]; + let constructor = "mlir::pto::createVMIMaskGranularityAssignmentPass()"; + let dependentDialects = ["mlir::cf::ControlFlowDialect", + "mlir::func::FuncDialect", + "mlir::pto::PTODialect", + "mlir::memref::MemRefDialect", + "mlir::scf::SCFDialect"]; +} + +def VMILayoutFold : Pass<"vmi-layout-fold", "ModuleOp"> { + let summary = "Fold VMI layout materializations"; + let description = [{ + Optimizes legal layout-assigned VMI IR by folding selected ensure_layout + helpers into layout-aware producers or consumers while preserving the same + logical effect. The pass does not choose layouts by inspecting arbitrary + producer/user context for vmi-to-vpto; it only rewrites explicit helper IR + into equivalent local forms. + }]; + let constructor = "mlir::pto::createVMILayoutFoldPass()"; + let dependentDialects = ["mlir::cf::ControlFlowDialect", + "mlir::func::FuncDialect", + "mlir::pto::PTODialect", + "mlir::memref::MemRefDialect", + "mlir::scf::SCFDialect"]; +} + +def VMILayoutRematerialize : Pass<"vmi-layout-rematerialize", "ModuleOp"> { + let summary = "Rematerialize cheap VMI producers at layout helpers"; + let description = [{ + Optimizes legal layout-assigned VMI IR by replacing selected ensure_layout, + ensure_mask_layout, and ensure_mask_granularity helpers with cloned + producers that directly create the requested result type. The pass covers + pure construction ops, selected layout-transparent data ops, and dense + widening ext relation rematerialization. Memory, control-flow, and mask-tail + proofs remain explicit in the IR. + }]; + let constructor = "mlir::pto::createVMILayoutRematerializePass()"; + let dependentDialects = ["mlir::cf::ControlFlowDialect", + "mlir::func::FuncDialect", + "mlir::pto::PTODialect", + "mlir::memref::MemRefDialect", + "mlir::scf::SCFDialect"]; +} + +def VMILayoutSinkMaterialization + : Pass<"vmi-layout-sink-materialization", "ModuleOp"> { + let summary = "Sink VMI layout materialization through transfer ops"; + let description = [{ + Optimizes legal layout-assigned VMI IR by moving matching operand + ensure_layout helpers across pure layout-transparent elementwise operations. + The rewritten IR keeps the layout conversion explicit as a result + ensure_layout, so vmi-to-vpto still lowers from local op information only. + }]; + let constructor = "mlir::pto::createVMILayoutSinkMaterializationPass()"; + let dependentDialects = ["mlir::cf::ControlFlowDialect", + "mlir::func::FuncDialect", + "mlir::pto::PTODialect", + "mlir::memref::MemRefDialect", + "mlir::scf::SCFDialect"]; +} + +def VMILegalizeArithSelect : Pass<"vmi-legalize-arith-select", "ModuleOp"> { + let summary = "Legalize canonical arith.select over VMI values"; + let description = [{ + Rewrites scalar-condition arith.select operations that produce VMI values + back to scf.if. MLIR canonicalization may fold simple scf.if regions into + arith.select, but VMI values must not cross non-VMI semantic ops before + vmi-to-vpto. This pass restores an explicit structural control-flow form + that the VMI converter already handles. + }]; + let constructor = "mlir::pto::createVMILegalizeArithSelectPass()"; + let dependentDialects = ["mlir::arith::ArithDialect", + "mlir::cf::ControlFlowDialect", + "mlir::func::FuncDialect", + "mlir::pto::PTODialect", + "mlir::memref::MemRefDialect", + "mlir::scf::SCFDialect"]; +} + +def VMILowerUnifiedToLegacy : Pass<"vmi-lower-unified-to-legacy", "ModuleOp"> { + let summary = "Lower unified VMI ops to legacy equivalents before layout assignment"; + let description = [{ + Expands most unified v-prefixed VMI ops into their legacy (mask-less) + equivalents so downstream layout and VMIToVPTO passes only see legacy ops. + + Ops lowered (Category A–C6): + A: vci, vinterpret_cast, vsel, vbrc → iota, bitcast, select, broadcast/group_broadcast + B: vadd/vsub/vmul/vdiv/vmin/vmax/vand/vor/vxor/vshl/vshr (masked binary) + vneg/vabs/vsqrt/vexp/vln/vrelu/vnot (masked unary) + → binary ops discard mask/pmode; unary ops preserve zero mode with select; + vshr selects shrui for explicit unsigned elements and shrsi otherwise + C1: vcmp/vcmps → legacy cmp + select + C2: vcvt → legacy extf/truncf/fptosi/sitofp/extsi/extui/trunci + C3: vload/vstore → legacy load/store variants + C4: pset/pge → create_mask/create_group_mask + C5: vadds/vmuls/vmaxs/vmins/vshls/vshrs + → broadcast + legacy binary, discarding mask/pmode; vshrs selects + shrui for explicit unsigned elements and shrsi otherwise + C6: vcadd/vcmax/vcmin → legacy reduce variants + + Ops NOT lowered (no legacy equivalent — require direct VMIToVPTO 1:N patterns): + plt, vhist, vintlv, vdintlv, vselr, + vgather, vgatherb, vscatter, + vexpdif, vaxpy, vlrelu, vprelu, vmull, vmula + }]; + let constructor = "mlir::pto::createVMILowerUnifiedToLegacyPass()"; + let dependentDialects = ["mlir::pto::PTODialect"]; +} + +def VMIToVPTO : Pass<"vmi-to-vpto", "ModuleOp"> { + let summary = "Convert layout-assigned VMI IR to physical VPTO IR"; + let description = [{ + Converts layout-assigned VMI aggregate data/mask types to ordered physical + VPTO register and mask value lists using MLIR native 1:N dialect conversion + APIs. This + pass is responsible for VMI 1:N type conversion, structural control-flow + and function/call signature conversion, and VMI semantic op physicalization. + }]; + let constructor = "mlir::pto::createVMIToVPTOPass()"; + let options = [ + Option<"enableStableGatherMaskedLoad", + "enable-stable-gather-masked-load", "bool", + /*default=*/"false", + "Reserve the stable VGATHER-based lowering path for VMI masked " + "loads; currently emits a TODO diagnostic when used."> + ]; + let dependentDialects = ["mlir::cf::ControlFlowDialect", + "mlir::func::FuncDialect", + "mlir::pto::PTODialect", + "mlir::memref::MemRefDialect", + "mlir::scf::SCFDialect"]; +} + def PTOValidateVPTOEmissionIR : Pass<"pto-validate-vpto-emission-ir", "ModuleOp"> { let summary = @@ -1016,4 +1262,19 @@ def VPTOPtrCastCleanup "mlir::memref::MemRefDialect"]; } +def VPTONormalizeEquivalentVcvt + : Pass<"vpto-normalize-equivalent-vcvt", "ModuleOp"> { + let summary = "Normalize equivalent VPTO vcvt part selections"; + let description = [{ + Rewrites `pto.vcvt` operations whose `EVEN` and `ODD` part selections are + provably equivalent into the canonical `EVEN` form. The pass currently + recognizes all-true masked narrow-to-wide conversions from VPTO values with + pair-wise equivalent input lanes, such as scalar/vector broadcasts and + selected broadcast load distributions. A following CSE pass can then merge + duplicate conversions. + }]; + let constructor = "mlir::pto::createVPTONormalizeEquivalentVcvtPass()"; + let dependentDialects = ["mlir::pto::PTODialect"]; +} + #endif // MLIR_DIALECT_PTO_PASSES diff --git a/include/PTO/Transforms/VMILayoutPropagation.h b/include/PTO/Transforms/VMILayoutPropagation.h new file mode 100644 index 0000000000..97853989e0 --- /dev/null +++ b/include/PTO/Transforms/VMILayoutPropagation.h @@ -0,0 +1,100 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +//===- VMILayoutPropagation.h - VMI layout request propagation -*- C++ -*-===// +//===----------------------------------------------------------------------===// + +#ifndef PTO_TRANSFORMS_VMILAYOUTPROPAGATION_H +#define PTO_TRANSFORMS_VMILAYOUTPROPAGATION_H + +#include "PTO/IR/PTO.h" + +#include "mlir/IR/Operation.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/IR/Value.h" +#include "mlir/Support/LogicalResult.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/SmallVector.h" + +#include + +namespace mlir::pto { + +struct VMILayoutConflict { + OpOperand *operand = nullptr; + VMILayoutAttr layout; +}; + +struct VMIValueLayoutAssignment { + VMILayoutAttr layout; + SmallVector conflicts; +}; + +class VMILayoutPropagator { +public: + explicit VMILayoutPropagator(Operation *scope); + + LogicalResult request(Value value, VMILayoutAttr layout); + LogicalResult request(OpOperand &operand, VMILayoutAttr layout); + void addEquivalentValues(Value lhs, Value rhs); + + LogicalResult run(); + LogicalResult apply(RewriterBase &rewriter); + + bool canUseOperandLayout(OpOperand &operand, VMILayoutAttr layout) const; + VMILayoutAttr getRequestedLayout(Value value) const; + VMILayoutAttr getRequestedOrCurrentLayout(Value value) const; + const VMIValueLayoutAssignment *lookup(Value value) const; + +private: + using LayoutFact = std::pair; + using OperandLayoutFact = std::pair; + + bool isLayoutValue(Value value) const; + VMILayoutAttr getCurrentLayout(Value value) const; + Type getTypeWithLayout(Value value, VMILayoutAttr layout) const; + bool isTypeRewriteable(Value value) const; + VMILayoutAttr getOperandLayout(OpOperand &operand) const; + bool canProduceValueLayout(Value value, VMILayoutAttr layout) const; + bool canMaterializeLayout(Value value, VMILayoutAttr sourceLayout, + VMILayoutAttr resultLayout) const; + + void enqueue(Value value, VMILayoutAttr layout); + LogicalResult addUseConflict(OpOperand &operand, + VMIValueLayoutAssignment &assignment, + VMILayoutAttr layout); + LogicalResult propagateFact(Value value, VMILayoutAttr layout); + LogicalResult propagateOperandFact(OpOperand &operand, VMILayoutAttr layout); + LogicalResult propagateThrough(Operation *op, Value changedValue, + VMILayoutAttr changedLayout, + OpOperand *changedOperand = nullptr); + LogicalResult verifyMaterializationPlan() const; + + LogicalResult materializePrimary(Value value, + const VMIValueLayoutAssignment &assignment, + RewriterBase &rewriter, + DenseMap &assignedValues); + FailureOr materializeAt(Value source, VMILayoutAttr layout, + RewriterBase &rewriter, Location loc); + LogicalResult materializeUseConflict(Value assignedValue, + VMILayoutConflict conflict, + RewriterBase &rewriter); + + Operation *scope = nullptr; + MLIRContext *ctx = nullptr; + DenseMap assignments; + SmallVector orderedValues; + SmallVector worklist; + SmallVector seenFacts; + SmallVector seenOperandFacts; + DenseMap> equivalentValues; +}; + +} // namespace mlir::pto + +#endif // PTO_TRANSFORMS_VMILAYOUTPROPAGATION_H diff --git a/include/PTO/Transforms/VMILayoutSupport.h b/include/PTO/Transforms/VMILayoutSupport.h new file mode 100644 index 0000000000..735cc63bca --- /dev/null +++ b/include/PTO/Transforms/VMILayoutSupport.h @@ -0,0 +1,422 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +//===- VMILayoutSupport.h - VMI layout support queries ------*- C++ -*-===// +//===----------------------------------------------------------------------===// + +#ifndef PTO_TRANSFORMS_VMILAYOUTSUPPORT_H +#define PTO_TRANSFORMS_VMILAYOUTSUPPORT_H + +#include "PTO/IR/PTO.h" +#include "mlir/Support/LLVM.h" + +#include "llvm/ADT/SmallVector.h" + +#include + +namespace mlir::pto { + +struct VMILoadLayoutFact { + VMILayoutAttr resultLayout; +}; + +enum class VMIDeinterleaveLoadLayoutPort { + Low, + High, +}; + +struct VMIDeinterleaveLoadLayoutFact { + VMILayoutAttr lowLayout; + VMILayoutAttr highLayout; +}; + +struct VMIStoreLayoutFact { + VMILayoutAttr valueLayout; +}; + +struct VMIMaskedStoreLayoutFact { + VMILayoutAttr valueLayout; + VMILayoutAttr maskLayout; +}; + +struct VMIMaskedLoadLayoutFact { + VMILayoutAttr resultLayout; + VMILayoutAttr maskLayout; + VMILayoutAttr passthruLayout; +}; + +struct VMIEnsureLayoutFact { + VMILayoutAttr sourceLayout; + VMILayoutAttr resultLayout; +}; + +struct VMIEnsureMaskLayoutFact { + VMILayoutAttr sourceLayout; + VMILayoutAttr resultLayout; +}; + +enum class VMICastLayoutPort { + Source, + Result, +}; + +enum class VMIInterleaveLayoutPort { + Lhs, + Rhs, + Mask, + Low, + High, +}; + +struct VMICastLayoutFact { + VMILayoutAttr sourceLayout; + VMILayoutAttr resultLayout; + int64_t sourceBits = 0; + int64_t resultBits = 0; +}; + +struct VMIMaskGranularityCastLayoutFact { + VMILayoutAttr sourceLayout; + VMILayoutAttr resultLayout; + int64_t sourceGranularityBits = 0; + int64_t resultGranularityBits = 0; +}; + +struct VMIInterleaveLayoutFact { + VMILayoutAttr lhsLayout; + VMILayoutAttr rhsLayout; + VMILayoutAttr maskLayout; + VMILayoutAttr lowLayout; + VMILayoutAttr highLayout; + int64_t elementCount = 0; + int64_t lanesPerPart = 0; +}; + +struct VMIBitcastLayoutFact { + VMILayoutAttr sourceLayout; + VMILayoutAttr resultLayout; +}; + +enum class VMIGroupBlockClass { + QuarterBlock, + HalfBlock, + OneBlock, + TwoBlock, + FourBlock, + FullPartMultiple, +}; + +struct VMIGroupReduceLayoutFact { + VMIGroupBlockClass blockClass = VMIGroupBlockClass::OneBlock; + VMILayoutAttr sourceLayout; + VMILayoutAttr maskLayout; + VMILayoutAttr resultLayout; + int64_t groupSize = 0; + int64_t lanesPerPart = 0; + int64_t vcgBlockElems = 0; +}; + +struct VMIGroupBroadcastLayoutFact { + VMIGroupBlockClass blockClass = VMIGroupBlockClass::OneBlock; + VMILayoutAttr sourceLayout; + VMILayoutAttr resultLayout; + int64_t groupSize = 0; + int64_t lanesPerPart = 0; + int64_t vcgBlockElems = 0; +}; + +enum class VMIGroupBroadcastLoadDirectKind { + E2B, + BRC, +}; + +struct VMIGroupBroadcastLoadLayoutFact { + VMIGroupBlockClass blockClass = VMIGroupBlockClass::OneBlock; + VMILayoutAttr resultLayout; + int64_t groupSize = 0; + int64_t lanesPerPart = 0; + int64_t vcgBlockElems = 0; + int64_t elementBits = 0; +}; + +struct VMIGroupBroadcastLoadDirectFact { + VMIGroupBroadcastLoadDirectKind kind = VMIGroupBroadcastLoadDirectKind::E2B; + VMIGroupBroadcastLoadLayoutFact layout; +}; + +struct VMIGroupLoadLayoutFact { + VMIGroupBlockClass blockClass = VMIGroupBlockClass::TwoBlock; + VMILayoutAttr resultLayout; + int64_t groupSize = 0; +}; + +struct VMIGroupSlotLayoutFact { + VMILayoutAttr layout; + int64_t numGroups = 0; + int64_t slots = 0; +}; + +enum class VMIGroupReduceLayoutPort { + Source, + Mask, + Result, +}; + +enum class VMIGroupBroadcastLayoutPort { + Source, + Result, +}; + +struct VMIHistogramLayoutFact { + VMILayoutAttr accLayout; + VMILayoutAttr sourceLayout; + VMILayoutAttr maskLayout; + VMILayoutAttr resultLayout; +}; + +class VMILayoutSupport { +public: + FailureOr + getLoadLayoutFact(VMIVRegType resultType, + std::string *reason = nullptr) const; + + FailureOr + getPreferredDeinterleaveLoadLayoutFact( + VMIVRegType valueType, std::string *reason = nullptr) const; + + FailureOr> + getDeinterleaveLoadLayoutFactsForLayout( + VMIVRegType valueType, VMIDeinterleaveLoadLayoutPort port, + VMILayoutAttr layout, std::string *reason = nullptr) const; + + FailureOr + getDeinterleaveLoadLayoutFactForLayouts( + VMIVRegType lowType, VMIVRegType highType, + std::string *reason = nullptr) const; + + FailureOr + getStoreLayoutFact(VMIVRegType valueType, + std::string *reason = nullptr) const; + + FailureOr + getPreferredStoreLayoutFact(VMIVRegType valueType, + std::string *reason = nullptr) const; + + FailureOr + getMaskedStoreLayoutFact(VMIVRegType valueType, VMIMaskType maskType, + std::string *reason = nullptr) const; + + FailureOr + getPreferredMaskedStoreLayoutFact(VMIVRegType valueType, + VMIMaskType maskType, + std::string *reason = nullptr) const; + + FailureOr + getMaskedLoadLayoutFact(VMIVRegType resultType, VMIMaskType maskType, + VMIVRegType passthruType, + std::string *reason = nullptr) const; + + FailureOr + getEnsureLayoutFact(VMIVRegType sourceType, VMIVRegType resultType, + std::string *reason = nullptr) const; + + FailureOr + getEnsureMaskLayoutFact(VMIMaskType sourceType, VMIMaskType resultType, + std::string *reason = nullptr) const; + + FailureOr + getPreferredCastLayoutFact(VMIVRegType sourceType, VMIVRegType resultType, + std::string *reason = nullptr) const; + + FailureOr> + getCastLayoutFactsForLayout(VMIVRegType sourceType, VMIVRegType resultType, + VMICastLayoutPort port, VMILayoutAttr layout, + std::string *reason = nullptr) const; + + FailureOr getCastLayoutFactForSourceLayout( + VMIVRegType sourceType, VMIVRegType resultType, + VMILayoutAttr sourceLayout, std::string *reason = nullptr) const; + + FailureOr getCastLayoutFactForResultLayout( + VMIVRegType sourceType, VMIVRegType resultType, + VMILayoutAttr resultLayout, std::string *reason = nullptr) const; + + FailureOr getCastLayoutFactForLayouts( + VMIVRegType sourceType, VMIVRegType resultType, VMILayoutAttr sourceLayout, + VMILayoutAttr resultLayout, std::string *reason = nullptr) const; + + FailureOr> + getMaskGranularityCastLayoutFactsForLayout( + VMIMaskType sourceType, VMIMaskType resultType, VMICastLayoutPort port, + VMILayoutAttr layout, std::string *reason = nullptr) const; + + FailureOr + getMaskGranularityCastLayoutFactForLayouts( + VMIMaskType sourceType, VMIMaskType resultType, + VMILayoutAttr sourceLayout, VMILayoutAttr resultLayout, + std::string *reason = nullptr) const; + + FailureOr getWidenSourceLayoutForResultLayout( + VMIVRegType sourceType, VMIVRegType resultType, + VMILayoutAttr requestedResultLayout, std::string *reason = nullptr) const; + + FailureOr + getPreferredVintlvLayoutFact(VMIVRegType valueType, + std::string *reason = nullptr) const; + + FailureOr + getPreferredVdintlvLayoutFact(VMIVRegType valueType, + std::string *reason = nullptr) const; + + FailureOr> + getVintlvLayoutFactsForLayout(VMIVRegType valueType, + VMIInterleaveLayoutPort port, + VMILayoutAttr layout, + std::string *reason = nullptr) const; + + FailureOr> + getVdintlvLayoutFactsForLayout(VMIVRegType valueType, + VMIInterleaveLayoutPort port, + VMILayoutAttr layout, + std::string *reason = nullptr) const; + + FailureOr getVintlvLayoutFactForLayouts( + VMIVRegType lhsType, VMIVRegType rhsType, VMIMaskType maskType, + VMIVRegType lowType, VMIVRegType highType, + std::string *reason = nullptr) const; + + FailureOr getVdintlvLayoutFactForLayouts( + VMIVRegType lhsType, VMIVRegType rhsType, VMIMaskType maskType, + VMIVRegType lowType, VMIVRegType highType, + std::string *reason = nullptr) const; + + FailureOr + getGroupSlotLoadLayoutFact(VMIVRegType resultType, int64_t numGroups, + std::string *reason = nullptr) const; + + FailureOr + getGroupLoadLayoutFact(VMIGroupLoadOp op, + std::string *reason = nullptr) const; + FailureOr + getGroupLoadLayoutFact(VMIVRegType resultType, Value rowStride, + int64_t numGroups, + std::string *reason = nullptr) const; + + FailureOr + getGroupStoreLayoutFact(VMIVRegType valueType, int64_t numGroups, + std::string *reason = nullptr) const; + + FailureOr + getPreferredGroupReduceLayoutFact(VMIVRegType sourceType, int64_t numGroups, + std::string *reason = nullptr) const; + + FailureOr getGroupReduceLayoutFactForLayouts( + VMIVRegType sourceType, VMIMaskType maskType, VMIVRegType resultType, + int64_t numGroups, std::string *reason = nullptr) const; + + FailureOr> + getGroupReduceLayoutFactsForLayout(VMIVRegType sourceType, + int64_t numGroups, + VMIGroupReduceLayoutPort port, + VMILayoutAttr layout, + std::string *reason = nullptr) const; + + FailureOr + getGroupBroadcastLayoutFactForLayouts(VMIVRegType sourceType, + VMIVRegType resultType, + int64_t numGroups, + std::string *reason = nullptr) const; + + FailureOr> + getGroupBroadcastLayoutFactsForLayout(VMIVRegType sourceType, + VMIVRegType resultType, + int64_t numGroups, + VMIGroupBroadcastLayoutPort port, + VMILayoutAttr layout, + std::string *reason = nullptr) const; + + FailureOr + getGroupBroadcastLoadLayoutFact(VMIGroupBroadcastLoadOp op, + std::string *reason = nullptr) const; + FailureOr + getGroupBroadcastLoadLayoutFact(VMIVRegType resultType, + Value sourceGroupStride, int64_t numGroups, + std::string *reason = nullptr) const; + FailureOr getGroupBroadcastLoadDirectFact( + VMIGroupBroadcastLoadOp op, std::string *reason = nullptr) const; + FailureOr getGroupBroadcastLoadDirectFact( + VMIVRegType resultType, Type sourceType, Value sourceGroupStride, + int64_t numGroups, std::string *reason = nullptr) const; + + FailureOr + getVdhistLayoutFact(VMIVdhistOp op, std::string *reason = nullptr) const; + + FailureOr + getVchistLayoutFact(VMIVchistOp op, std::string *reason = nullptr) const; + + LogicalResult getGroupReduceAddFSupport(VMIGroupReduceAddFOp op, + std::string *reason = nullptr) const; + + LogicalResult getGroupReduceMaxFSupport(VMIGroupReduceMaxFOp op, + std::string *reason = nullptr) const; + + LogicalResult getGroupReduceMinFSupport(VMIGroupReduceMinFOp op, + std::string *reason = nullptr) const; + + LogicalResult getGroupReduceAddISupport(VMIGroupReduceAddIOp op, + std::string *reason = nullptr) const; + + LogicalResult getGroupReduceMaxISupport(VMIGroupReduceMaxIOp op, + std::string *reason = nullptr) const; + + LogicalResult getGroupReduceMinISupport(VMIGroupReduceMinIOp op, + std::string *reason = nullptr) const; + + LogicalResult getGroupBroadcastSupport(VMIGroupBroadcastOp op, + std::string *reason = nullptr) const; + + LogicalResult getGroupBroadcastSupport(VMIVRegType sourceType, + VMIVRegType resultType, + int64_t numGroups, + std::string *reason = nullptr) const; + + LogicalResult getGroupBroadcastLoadSupport( + VMIGroupBroadcastLoadOp op, std::string *reason = nullptr) const; + + LogicalResult getTruncFSupport(VMITruncFOp op, + std::string *reason = nullptr) const; + + LogicalResult getExtFSupport(VMIExtFOp op, + std::string *reason = nullptr) const; + + LogicalResult getExtSISupport(VMIExtSIOp op, + std::string *reason = nullptr) const; + + LogicalResult getExtUISupport(VMIExtUIOp op, + std::string *reason = nullptr) const; + + LogicalResult getTruncISupport(VMITruncIOp op, + std::string *reason = nullptr) const; + + FailureOr + getBitcastLayoutFact(VMIBitcastOp op, + std::string *reason = nullptr) const; + + LogicalResult getBitcastSupport(VMIBitcastOp op, + std::string *reason = nullptr) const; + + LogicalResult getVdhistSupport(VMIVdhistOp op, + std::string *reason = nullptr) const; + + LogicalResult getVchistSupport(VMIVchistOp op, + std::string *reason = nullptr) const; +}; + +} // namespace mlir::pto + +#endif // PTO_TRANSFORMS_VMILAYOUTSUPPORT_H diff --git a/lib/Bindings/Python/PTOModule.cpp b/lib/Bindings/Python/PTOModule.cpp index f0603158b3..432d08e262 100644 --- a/lib/Bindings/Python/PTOModule.cpp +++ b/lib/Bindings/Python/PTOModule.cpp @@ -108,6 +108,12 @@ static py::object wrapAttributeAs(const py::module_ &m, const char *className, return cls.attr("__call__")(attr); } +static MlirAttribute optionalAttributeFromPy(py::object attr) { + if (attr.is_none()) + return MlirAttribute{nullptr}; + return py::cast(attr); +} + void populatePTODialectSubmodule(pybind11::module &m); void populatePTODialectSubmodule(pybind11::module &m) { (void)m; @@ -969,6 +975,82 @@ static void bindPTOModule(pybind11::module &m) { return cast(unwrap(self)).getGranularity().str(); }); + mlir_type_subclass( + m, "VMIVRegType", + [](MlirType type) -> bool { + return isa(unwrap(type)); + }) + .def_classmethod( + "get", + [](py::object cls, int64_t elementCount, MlirType elementType, + py::object layout, MlirContext context) -> py::object { + context = inferContextFromElementType(context, elementType); + MlirAttribute layoutAttr = optionalAttributeFromPy(layout); + MlirType t = wrap(mlir::pto::VMIVRegType::get( + unwrap(context), elementCount, unwrap(elementType), + unwrap(layoutAttr))); + return cls.attr("__call__")(t); + }, + py::arg("cls"), py::arg("element_count"), py::arg("element_type"), + py::arg("layout") = py::none(), py::arg("context") = py::none()) + .def_property_readonly( + "element_count", + [](MlirType self) -> int64_t { + return cast(unwrap(self)).getElementCount(); + }) + .def_property_readonly( + "element_type", + [](MlirType self) -> MlirType { + return wrap(cast(unwrap(self)).getElementType()); + }) + .def_property_readonly( + "layout", + [](MlirType self) -> py::object { + mlir::Attribute attr = + cast(unwrap(self)).getLayout(); + if (!attr) + return py::none(); + return py::cast(wrap(attr)); + }); + + mlir_type_subclass( + m, "VMIMaskType", + [](MlirType type) -> bool { + return isa(unwrap(type)); + }) + .def_classmethod( + "get", + [](py::object cls, int64_t elementCount, std::string granularity, + py::object layout, MlirContext context) -> py::object { + MlirAttribute layoutAttr = optionalAttributeFromPy(layout); + MlirType t = wrap(mlir::pto::VMIMaskType::get( + unwrap(context), elementCount, granularity, + unwrap(layoutAttr))); + return cls.attr("__call__")(t); + }, + py::arg("cls"), py::arg("element_count"), + py::arg("granularity") = "pred", py::arg("layout") = py::none(), + py::arg("context") = py::none()) + .def_property_readonly( + "element_count", + [](MlirType self) -> int64_t { + return cast(unwrap(self)).getElementCount(); + }) + .def_property_readonly( + "granularity", + [](MlirType self) -> std::string { + return cast(unwrap(self)).getGranularity().str(); + }) + .def_property_readonly( + "layout", + [](MlirType self) -> py::object { + mlir::Attribute attr = + cast(unwrap(self)).getLayout(); + if (!attr) + return py::none(); + return py::cast(wrap(attr)); + }); + mlir_type_subclass( m, "AlignType", [](MlirType type) -> bool { return isa(unwrap(type)); }) diff --git a/lib/PTO/IR/CMakeLists.txt b/lib/PTO/IR/CMakeLists.txt index 74b9e0bd68..4f8d995796 100644 --- a/lib/PTO/IR/CMakeLists.txt +++ b/lib/PTO/IR/CMakeLists.txt @@ -15,6 +15,7 @@ add_mlir_dialect_library(PTOIR PTO.cpp VPTO.cpp + VMI.cpp PTOAttrs.cpp PTOSyncUtils.cpp PTOTypeDefs.cpp diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 1b66634049..c7e8effcc0 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -30,6 +30,7 @@ #include "mlir/IR/Types.h" #include "mlir/Interfaces/SideEffectInterfaces.h" #include "mlir/Support/LLVM.h" +#include "mlir/Transforms/InliningUtils.h" #include "mlir/Parser/Parser.h" #include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/STLExtras.h" @@ -121,6 +122,27 @@ static bool isKnownZeroOrUnitExtent(int64_t value); static bool isByteIntegerType(Type ty); static LogicalResult verifyTileBufCommon(Operation *op, Type ty, StringRef name, bool allowLowPrecision = false); + +namespace { +struct PTOInlinerInterface : public DialectInlinerInterface { + using DialectInlinerInterface::DialectInlinerInterface; + + bool isLegalToInline(Operation *call, Operation *callable, + bool wouldBeCloned) const final { + return true; + } + + bool isLegalToInline(Region *dest, Region *src, bool wouldBeCloned, + IRMapping &valueMapping) const final { + return true; + } + + bool isLegalToInline(Operation *op, Region *dest, bool wouldBeCloned, + IRMapping &valueMapping) const final { + return true; + } +}; +} // namespace static LogicalResult verifyTileBufSameElemType(Operation *op, Type lhs, Type rhs, StringRef lhsName, StringRef rhsName); @@ -2701,6 +2723,8 @@ void PTODialect::initialize() { #define GET_ATTRDEF_LIST #include "PTO/IR/PTOAttrs.cpp.inc" >(); + + addInterfaces(); } diff --git a/lib/PTO/IR/VMI.cpp b/lib/PTO/IR/VMI.cpp new file mode 100644 index 0000000000..61e9f2a352 --- /dev/null +++ b/lib/PTO/IR/VMI.cpp @@ -0,0 +1,4216 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +//===- VMI.cpp - PTO VMI type and attribute support -----------------------===// +//===----------------------------------------------------------------------===// + +#include "PTO/IR/PTO.h" +#include "PTO/IR/PTOTypeUtils.h" +#include "PTO/IR/VMIUtils.h" + +#include "mlir/IR/DialectImplementation.h" +#include "mlir/IR/OpImplementation.h" +#include "mlir/IR/Types.h" +#include "llvm/ADT/SmallString.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/Support/raw_ostream.h" +#include +#include + +using namespace mlir; +using namespace mlir::pto; + +namespace { + +static std::string formatVMIVRegType(int64_t elementCount, Type elementType, + Attribute layout) { + std::string result; + llvm::raw_string_ostream os(result); + os << "!pto.vmi.vreg<" << elementCount << "x" << elementType; + if (layout) + os << ", " << layout; + os << ">"; + return result; +} + +static std::string formatVMIMaskType(int64_t elementCount, + StringRef granularity, Attribute layout) { + std::string result; + llvm::raw_string_ostream os(result); + os << "!pto.vmi.mask<" << elementCount << "x" << granularity; + if (layout) + os << ", " << layout; + os << ">"; + return result; +} + +static bool isSupportedVMIElementType(Type type) { + return isa(type) || + pto::isPTOLowPrecisionType(type); +} + +static bool isVMIFloatLikeType(Type type) { + return isa(type) || pto::isPTOLowPrecisionType(type); +} + +static bool isVMIIntegerLikeType(Type type) { + return isa(type); +} + +static bool isVMIF16OrF32Type(Type type) { + return type.isF16() || type.isF32(); +} + +static bool isVMIF16BF16OrF32Type(Type type) { + return type.isF16() || type.isBF16() || type.isF32(); +} + +static bool isVMIPredicateMaskableElementType(Type type) { + unsigned elementBits = pto::getPTOStorageElemBitWidth(type); + return elementBits == 8 || elementBits == 16 || elementBits == 32; +} + +static bool isVMIAnyI8I16I32Type(Type type) { + auto integerType = dyn_cast(type); + if (!integerType) + return false; + return integerType.getWidth() == 8 || integerType.getWidth() == 16 || + integerType.getWidth() == 32; +} + +static bool isVMISignedOrSignlessI8I16I32Type(Type type) { + auto integerType = dyn_cast(type); + if (!integerType || integerType.isUnsigned()) + return false; + return integerType.getWidth() == 8 || integerType.getWidth() == 16 || + integerType.getWidth() == 32; +} + +static bool isVMISignedOrSignlessIntegerType(Type type) { + auto integerType = dyn_cast(type); + return integerType && !integerType.isUnsigned(); +} + +static bool isVMIUnsignedIntegerType(Type type) { + auto integerType = dyn_cast(type); + return integerType && integerType.isUnsigned(); +} + +static bool isVMIIotaElementType(Type type) { + if (auto intType = dyn_cast(type)) + return intType.getWidth() == 8 || intType.getWidth() == 16 || + intType.getWidth() == 32; + return type.isF16() || type.isF32(); +} + +static bool isCompatibleScalarForSemanticType(Type semanticType, + Type scalarType) { + if (semanticType == scalarType) + return true; + + auto semanticInt = dyn_cast(semanticType); + auto scalarInt = dyn_cast(scalarType); + if (!semanticInt || !scalarInt || + semanticInt.getWidth() != scalarInt.getWidth()) + return false; + + if (semanticInt.isSigned()) + return scalarInt.isSigned() || scalarInt.isSignless(); + if (semanticInt.isUnsigned()) + return scalarInt.isUnsigned() || scalarInt.isSignless(); + return scalarInt.isSignless(); +} + +static unsigned getVMIElementBitWidth(Type type) { + if (isa(type)) + return 64; + return pto::getPTOStorageElemBitWidth(type); +} + +static std::optional getVMIIntegerOrFloatBitWidth(Type type) { + if (auto intType = dyn_cast(type)) + return intType.getWidth(); + if (auto floatType = dyn_cast(type)) + return floatType.getWidth(); + return std::nullopt; +} + +static int64_t divideCeilNonNegative(int64_t value, int64_t divisor) { + return value == 0 ? 0 : (value + divisor - 1) / divisor; +} + +static LogicalResult parseOptionalVMILayout(AsmParser &parser, + Attribute &layout) { + if (failed(parser.parseOptionalComma())) + return success(); + + if (failed(parser.parseAttribute(layout))) + return failure(); + if (!mlir::isa(layout)) + return parser.emitError(parser.getCurrentLocation(), + "expected #pto.vmi.layout attribute"); + return success(); +} + +static FailureOr getVMIElementCount(Type type) { + if (auto vregType = dyn_cast(type)) + return vregType.getElementCount(); + if (auto maskType = dyn_cast(type)) + return maskType.getElementCount(); + return failure(); +} + +static FailureOr getAssignedVMILayout(Type type) { + Attribute layout; + if (auto vregType = dyn_cast(type)) + layout = vregType.getLayout(); + else if (auto maskType = dyn_cast(type)) + layout = maskType.getLayout(); + else + return failure(); + + auto layoutAttr = dyn_cast_or_null(layout); + if (!layoutAttr) + return failure(); + return layoutAttr; +} + +static FailureOr getLayoutFactor(Type type) { + FailureOr layout = getAssignedVMILayout(type); + if (failed(layout)) + return failure(); + return (*layout).isDeinterleaved() ? (*layout).getFactor() : 1; +} + +static FailureOr getLayoutBlockElems(Type type) { + FailureOr layout = getAssignedVMILayout(type); + if (failed(layout)) + return failure(); + return (*layout).isDeinterleaved() ? (*layout).getBlockElems() : 1; +} + +static FailureOr getVMIPhysicalElementType(VMIVRegType type) { + Type elementType = type.getElementType(); + VMILayoutAttr layout = type.getLayoutAttr(); + if (!layout || !layout.hasGroupSlotLaneStride()) + return elementType; + + auto integerType = dyn_cast(elementType); + if (!integerType && isa(elementType)) + return elementType; + if (!integerType) + return failure(); + if (!integerType.isUnsigned()) + return failure(); + unsigned elementBits = pto::getPTOStorageElemBitWidth(elementType); + int64_t laneStride = layout.getLaneStride(); + if (elementBits == 0 || laneStride <= 1) + return failure(); + int64_t physicalBits = static_cast(elementBits) * laneStride; + if (physicalBits != 16 && physicalBits != 32) + return failure(); + return IntegerType::get(type.getContext(), physicalBits); +} + +static int64_t getMaskGranularityBitWidth(StringRef granularity) { + if (granularity == "b8") + return 8; + if (granularity == "b16") + return 16; + if (granularity == "b32") + return 32; + return 0; +} + +static StringRef getMaskGranularityForBitWidth(int64_t bits) { + switch (bits) { + case 8: + return "b8"; + case 16: + return "b16"; + case 32: + return "b32"; + default: + return ""; + } +} + +static FailureOr getVMIMaskPhysicalGranularity(VMIMaskType type) { + int64_t bits = getMaskGranularityBitWidth(type.getGranularity()); + if (bits == 0) + return failure(); + + VMILayoutAttr layout = type.getLayoutAttr(); + int64_t laneStride = layout && layout.hasLaneStride() ? layout.getLaneStride() + : 1; + StringRef physicalGranularity = + getMaskGranularityForBitWidth(bits * laneStride); + if (physicalGranularity.empty()) + return failure(); + return physicalGranularity; +} + +static FailureOr getPhysicalLanesPerPart(Type type) { + if (auto vregType = dyn_cast(type)) { + FailureOr physicalElementType = getVMIPhysicalElementType(vregType); + if (failed(physicalElementType)) + return failure(); + return getDataLanesPerPart(*physicalElementType); + } + if (auto maskType = dyn_cast(type)) { + FailureOr physicalGranularity = + getVMIMaskPhysicalGranularity(maskType); + if (failed(physicalGranularity)) + return failure(); + return getMaskLanesPerPart(*physicalGranularity); + } + return failure(); +} + +static FailureOr getDenseLaneStride(Type type) { + FailureOr layout = getAssignedVMILayout(type); + if (failed(layout)) + return failure(); + if (isa(type)) + return 1; + return (*layout).isDense() ? (*layout).getLaneStride() : 1; +} + +static bool isLayoutAssigned(VMIVRegType type) { + return static_cast(type.getLayoutAttr()); +} + +static bool isLayoutAssigned(VMIMaskType type) { + return static_cast(type.getLayoutAttr()); +} + +static LogicalResult +verifyAllSameVRegShapeAndLayout(Operation *op, ArrayRef types, + bool requireSameElement) { + if (types.empty()) + return success(); + + VMIVRegType first = types.front(); + bool anyLayout = llvm::any_of( + types, [](VMIVRegType type) { return isLayoutAssigned(type); }); + + for (VMIVRegType type : types) { + if (type.getElementCount() != first.getElementCount()) + return op->emitOpError( + "requires all VMI data values to have the same logical lane count"); + if (requireSameElement && type.getElementType() != first.getElementType()) + return op->emitOpError( + "requires all VMI data values to have the same element type"); + if (anyLayout && !isLayoutAssigned(type)) + return op->emitOpError( + "requires either all or no VMI data values to carry layout"); + if (anyLayout && type.getLayout() != first.getLayout()) + return op->emitOpError("requires all layout-assigned VMI data values to " + "have the same layout"); + } + return success(); +} + +static LogicalResult verifyAllSameVRegShapeAndLayoutPresence( + Operation *op, ArrayRef types, bool requireSameElement) { + if (types.empty()) + return success(); + + VMIVRegType first = types.front(); + bool anyLayout = llvm::any_of( + types, [](VMIVRegType type) { return isLayoutAssigned(type); }); + + for (VMIVRegType type : types) { + if (type.getElementCount() != first.getElementCount()) + return op->emitOpError( + "requires all VMI data values to have the same logical lane count"); + if (requireSameElement && type.getElementType() != first.getElementType()) + return op->emitOpError( + "requires all VMI data values to have the same element type"); + if (anyLayout && !isLayoutAssigned(type)) + return op->emitOpError( + "requires either all or no VMI data values to carry layout"); + } + return success(); +} + +static LogicalResult verifyElementwiseVRegOp(Operation *op, VMIVRegType lhs, + VMIVRegType rhs, + VMIVRegType result) { + return verifyAllSameVRegShapeAndLayout(op, {lhs, rhs, result}, + /*requireSameElement=*/true); +} + +static LogicalResult verifyFloatUnaryVRegOp(Operation *op, VMIVRegType source, + VMIVRegType result) { + if (!isVMIFloatLikeType(source.getElementType())) + return op->emitOpError("requires floating-point-like VMI element type"); + return verifyAllSameVRegShapeAndLayout(op, {source, result}, + /*requireSameElement=*/true); +} + +static LogicalResult verifyFloatTernaryVRegOp(Operation *op, VMIVRegType lhs, + VMIVRegType rhs, VMIVRegType acc, + VMIVRegType result) { + if (!isVMIFloatLikeType(lhs.getElementType())) + return op->emitOpError("requires floating-point-like VMI element type"); + return verifyAllSameVRegShapeAndLayout(op, {lhs, rhs, acc, result}, + /*requireSameElement=*/true); +} + +static LogicalResult +verifyAllSameMaskShapeLayoutAndGranularity(Operation *op, + ArrayRef types) { + if (types.empty()) + return success(); + + VMIMaskType first = types.front(); + bool anyLayout = llvm::any_of( + types, [](VMIMaskType type) { return isLayoutAssigned(type); }); + + for (VMIMaskType type : types) { + if (type.getElementCount() != first.getElementCount()) + return op->emitOpError( + "requires all VMI mask values to have the same logical lane count"); + if (type.getGranularity() != first.getGranularity()) + return op->emitOpError( + "requires all VMI mask values to have the same granularity"); + if (anyLayout && !isLayoutAssigned(type)) + return op->emitOpError( + "requires either all or no VMI mask values to carry layout"); + if (anyLayout && type.getLayout() != first.getLayout()) + return op->emitOpError( + "requires all layout-assigned VMI mask values to have the same " + "layout"); + } + return success(); +} + +static LogicalResult verifyMaskMatchesData(Operation *op, VMIMaskType maskType, + VMIVRegType dataType) { + if (maskType.getElementCount() != dataType.getElementCount()) + return op->emitOpError( + "requires mask logical lane count to match data lane count"); + + if (isLayoutAssigned(maskType) || isLayoutAssigned(dataType)) { + if (!isLayoutAssigned(maskType) || !isLayoutAssigned(dataType)) + return op->emitOpError("requires either both mask and data to carry " + "layout or neither to carry layout"); + if (maskType.getLayout() != dataType.getLayout()) + return op->emitOpError("requires mask layout to match data layout"); + } + + if (maskType.isPred()) + return success(); + + unsigned elementBitWidth = getVMIElementBitWidth(dataType.getElementType()); + int64_t maskBitWidth = getMaskGranularityBitWidth(maskType.getGranularity()); + if (elementBitWidth != 0 && maskBitWidth != 0 && + elementBitWidth != static_cast(maskBitWidth)) + return op->emitOpError( + "requires mask granularity to match data element width"); + + return success(); +} + + +static Type getMemoryElementType(Type type) { + if (auto ptrType = dyn_cast(type)) + return ptrType.getElementType(); + if (auto memrefType = dyn_cast(type)) + return memrefType.getElementType(); + return {}; +} + +static bool isUBBackedMemoryType(Type type) { + if (auto ptrType = dyn_cast(type)) + return ptrType.getMemorySpace().getAddressSpace() == AddressSpace::VEC; + + auto memrefType = dyn_cast(type); + if (!memrefType) + return false; + + Attribute memorySpace = memrefType.getMemorySpace(); + if (auto addressSpace = dyn_cast_or_null(memorySpace)) + return addressSpace.getAddressSpace() == AddressSpace::VEC; + if (auto integerSpace = dyn_cast_or_null(memorySpace)) + return integerSpace.getInt() == static_cast(AddressSpace::VEC); + return false; +} + +static LogicalResult verifyUBBackedMemory(Operation *op, Type memoryType, + StringRef role) { + if (isUBBackedMemoryType(memoryType)) + return success(); + return op->emitOpError() << "requires memory " << role + << " to be UB-backed"; +} + +static LogicalResult verifyMemoryElementMatches(Operation *op, Type memoryType, + VMIVRegType dataType, + StringRef role) { + Type memoryElementType = getMemoryElementType(memoryType); + if (!memoryElementType) + return success(); + if (memoryElementType != dataType.getElementType()) + return op->emitOpError() << "requires memory " << role + << " element type to match VMI data element type"; + return success(); +} + +static LogicalResult verifyContiguousIfLayoutAssigned(Operation *op, + VMIVRegType type, + StringRef role) { + VMILayoutAttr layout = type.getLayoutAttr(); + if (layout && !layout.isContiguous()) + return op->emitOpError() + << "requires layout-assigned " << role + << " to use #pto.vmi.layout"; + return success(); +} + +static bool isPackedByteGroupStore(Type memoryType, VMIVRegType dataType) { + Type memoryElementType = getMemoryElementType(memoryType); + if (!memoryElementType) + return false; + auto memoryIntegerType = dyn_cast(memoryElementType); + auto dataIntegerType = dyn_cast(dataType.getElementType()); + return memoryIntegerType && dataIntegerType && + memoryIntegerType.getWidth() == 8 && dataIntegerType.getWidth() == 32; +} + +static LogicalResult verifyNumGroups(Operation *op, VMIVRegType type, + int64_t numGroups) { + if (numGroups <= 0) + return op->emitOpError("requires num_groups to be positive"); + if (type.getElementCount() % numGroups != 0) + return op->emitOpError() + << "requires num_groups to evenly divide VMI logical lane count " + << type.getElementCount(); + return success(); +} + +static LogicalResult verifyPhysicalParts(Operation *op, Type vmiType, + TypeRange physicalTypes) { + FailureOr expectedArity = getVMIPhysicalArity(vmiType); + if (failed(expectedArity)) + return op->emitOpError( + "requires a layout-assigned VMI type with computable physical arity"); + if (static_cast(physicalTypes.size()) != *expectedArity) + return op->emitOpError() << "requires " << *expectedArity + << " physical parts, got " << physicalTypes.size(); + + if (auto vregType = dyn_cast(vmiType)) { + FailureOr lanesPerPart = + getPhysicalLanesPerPart(vregType); + FailureOr physicalElementType = getVMIPhysicalElementType(vregType); + if (failed(lanesPerPart) || failed(physicalElementType)) + return op->emitOpError( + "requires data element type with known physical lane count"); + for (Type physicalType : physicalTypes) { + auto partType = dyn_cast(physicalType); + if (!partType) + return op->emitOpError("requires physical data parts to be !pto.vreg"); + if (partType.getElementCount() != *lanesPerPart || + partType.getElementType() != *physicalElementType) + return op->emitOpError( + "requires physical data part type to match VMI lane-map helper"); + } + return success(); + } + + auto maskType = dyn_cast(vmiType); + if (!maskType) + return op->emitOpError("requires VMI data or mask type"); + if (maskType.isPred()) + return op->emitOpError( + "requires layout-assigned mask with concrete granularity"); + FailureOr physicalGranularity = + getVMIMaskPhysicalGranularity(maskType); + if (failed(physicalGranularity)) + return op->emitOpError( + "requires mask type with supported physical carrier granularity"); + + for (Type physicalType : physicalTypes) { + auto partType = dyn_cast(physicalType); + if (!partType) + return op->emitOpError("requires physical mask parts to be !pto.mask"); + if (partType.getGranularity() != *physicalGranularity) + return op->emitOpError( + "requires physical mask part granularity to match VMI mask carrier"); + } + return success(); +} + +static std::optional +mapDenseLogicalLaneToPartIndex(int64_t elementCount, int64_t factor, + int64_t blockElems, int64_t logicalLane, + int64_t &part) { + if (logicalLane < 0 || logicalLane >= elementCount || factor <= 0 || + blockElems <= 0) + return std::nullopt; + int64_t block = logicalLane / blockElems; + int64_t inBlockLane = logicalLane % blockElems; + part = block % factor; + int64_t partBlock = block / factor; + return partBlock * blockElems + inBlockLane; +} + +static std::optional +mapDensePartIndexToLogicalLane(int64_t elementCount, int64_t factor, + int64_t blockElems, int64_t part, + int64_t indexInPart) { + if (part < 0 || part >= factor || indexInPart < 0 || factor <= 0 || + blockElems <= 0) + return std::nullopt; + int64_t partBlock = indexInPart / blockElems; + int64_t inBlockLane = indexInPart % blockElems; + int64_t logicalBlock = partBlock * factor + part; + int64_t logicalLane = logicalBlock * blockElems + inBlockLane; + if (logicalLane >= elementCount) + return std::nullopt; + return logicalLane; +} + +static int64_t getDenseLogicalLanesInPart(int64_t elementCount, int64_t factor, + int64_t blockElems, int64_t part) { + int64_t maxIndex = -1; + for (int64_t lane = 0; lane < elementCount; ++lane) { + int64_t lanePart = 0; + std::optional index = mapDenseLogicalLaneToPartIndex( + elementCount, factor, blockElems, lane, lanePart); + if (index && lanePart == part) + maxIndex = std::max(maxIndex, *index); + } + return maxIndex + 1; +} + +} // namespace + +VMILayoutAttr VMILayoutAttr::getContiguous(MLIRContext *context, + int64_t laneStride) { + return VMILayoutAttr::get(context, "contiguous", 1, 1, 0, laneStride); +} + +VMILayoutAttr VMILayoutAttr::getDeinterleaved(MLIRContext *context, + int64_t factor, + int64_t blockElems, + int64_t laneStride) { + return VMILayoutAttr::get(context, "deinterleaved", factor, blockElems, 0, + laneStride); +} + +VMILayoutAttr VMILayoutAttr::getGroupSlots(MLIRContext *context, + int64_t numGroups, int64_t slots, + int64_t laneStride) { + return VMILayoutAttr::get(context, "num_groups", numGroups, 1, slots, + laneStride); +} + +Attribute VMILayoutAttr::parse(AsmParser &parser, Type) { + SMLoc loc = parser.getCurrentLocation(); + StringRef kind; + int64_t factor = 1; + int64_t blockElems = 1; + int64_t slots = 0; + int64_t laneStride = 1; + + if (failed(parser.parseLess()) || failed(parser.parseKeyword(&kind))) + return {}; + + if (kind == "contiguous") { + factor = 1; + while (succeeded(parser.parseOptionalComma())) { + StringRef field; + if (failed(parser.parseKeyword(&field)) || failed(parser.parseEqual()) || + field != "lane_stride" || failed(parser.parseInteger(laneStride))) { + parser.emitError(parser.getCurrentLocation(), + "expected 'lane_stride = '"); + return {}; + } + } + } else if (kind == "deinterleaved") { + if (failed(parser.parseEqual()) || failed(parser.parseInteger(factor))) + return {}; + while (succeeded(parser.parseOptionalComma())) { + StringRef field; + if (failed(parser.parseKeyword(&field)) || failed(parser.parseEqual())) + return {}; + if (field == "block_elems") { + if (failed(parser.parseInteger(blockElems))) + return {}; + } else if (field == "lane_stride") { + if (failed(parser.parseInteger(laneStride))) + return {}; + } else { + parser.emitError(parser.getCurrentLocation(), + "expected 'block_elems = ' or " + "'lane_stride = '"); + return {}; + } + } + } else if (kind == "num_groups") { + if (failed(parser.parseEqual()) || failed(parser.parseInteger(factor))) + return {}; + while (succeeded(parser.parseOptionalComma())) { + StringRef field; + if (failed(parser.parseKeyword(&field)) || failed(parser.parseEqual())) + return {}; + if (field == "slots") { + if (failed(parser.parseInteger(slots))) + return {}; + } else if (field == "lane_stride") { + if (failed(parser.parseInteger(laneStride))) + return {}; + } else { + parser.emitError(parser.getCurrentLocation(), + "expected 'slots = ' or " + "'lane_stride = '"); + return {}; + } + } + } else { + parser.emitError(parser.getCurrentLocation(), + "expected VMI layout kind 'contiguous' or " + "'deinterleaved' or 'num_groups'"); + return {}; + } + + if (failed(parser.parseGreater())) + return {}; + + return parser.getChecked(loc, parser.getContext(), kind, + factor, blockElems, slots, + laneStride); +} + +void VMILayoutAttr::print(AsmPrinter &printer) const { + printer << "<" << getKind(); + if (isContiguous()) { + if (getLaneStride() != 1) + printer << ", lane_stride = " << getLaneStride(); + } else if (isDeinterleaved()) { + printer << " = " << getFactor(); + if (getBlockElems() != 1) + printer << ", block_elems = " << getBlockElems(); + if (getLaneStride() != 1) + printer << ", lane_stride = " << getLaneStride(); + } else if (isGroupSlots()) { + printer << " = " << getFactor(); + if (getSlots() != 0) + printer << ", slots = " << getSlots(); + if (getLaneStride() != 1) + printer << ", lane_stride = " << getLaneStride(); + } + printer << ">"; +} + +LogicalResult +VMILayoutAttr::verify(function_ref emitError, + StringRef kind, int64_t factor, int64_t blockElems, + int64_t slots, int64_t laneStride) { + if (laneStride <= 0) + return emitError() << "#pto.vmi.layout<" << kind + << "> requires lane_stride to be positive"; + + if (kind == "contiguous") { + if (factor != 1 || blockElems != 1 || slots != 0) + return emitError() + << "#pto.vmi.layout requires factor, block_elems, " + "and slots to be their defaults"; + return success(); + } + + if (kind == "deinterleaved") { + if (factor != 2 && factor != 4) + return emitError() << "#pto.vmi.layout expected factor to be 2 or 4"; + if (blockElems <= 0) + return emitError() << "#pto.vmi.layout requires block_elems to be positive"; + if (slots != 0) + return emitError() << "#pto.vmi.layout requires slots to be omitted"; + return success(); + } + + if (kind == "num_groups") { + if (factor <= 0) + return emitError() << "#pto.vmi.layout requires num_groups to be positive"; + if (blockElems != 1) + return emitError() << "#pto.vmi.layout requires block_elems to be omitted"; + if (slots < 0) + return emitError() << "#pto.vmi.layout requires slots to be omitted or positive"; + return success(); + } + + return emitError() << "expected VMI layout kind to be 'contiguous' or " + "'deinterleaved' or 'num_groups'"; +} + +Type VMIVRegType::parse(AsmParser &parser) { + SmallVector shape; + Type elementType; + Attribute layout; + SMLoc loc = parser.getCurrentLocation(); + + if (failed(parser.parseLess()) || + failed(parser.parseDimensionList(shape, /*allowDynamic=*/false, + /*withTrailingX=*/true)) || + shape.size() != 1 || failed(parser.parseType(elementType)) || + failed(parseOptionalVMILayout(parser, layout)) || + failed(parser.parseGreater())) + return {}; + + return parser.getChecked(loc, parser.getContext(), shape.front(), + elementType, layout); +} + +void VMIVRegType::print(AsmPrinter &printer) const { + printer << "<" << getElementCount() << "x"; + printer.printType(getElementType()); + if (getLayout()) + printer << ", " << getLayout(); + printer << ">"; +} + +LogicalResult VMIVRegType::verify(function_ref emitError, + int64_t elementCount, Type elementType, + Attribute layout) { + if (elementCount <= 0) + return emitError() << "'" + << formatVMIVRegType(elementCount, elementType, layout) + << "' expected a positive element count"; + + if (!isSupportedVMIElementType(elementType)) + return emitError() << "'" + << formatVMIVRegType(elementCount, elementType, layout) + << "' expected an integer, index, floating-point, or " + "PTO low-precision element type"; + if (!isVMIPredicateMaskableElementType(elementType)) + return emitError() << "'" + << formatVMIVRegType(elementCount, elementType, layout) + << "' expected an 8-bit, 16-bit, or 32-bit logical " + "element type"; + if (pto::isPTOFloat4PackedType(elementType)) + return emitError() + << "'" << formatVMIVRegType(elementCount, elementType, layout) + << "' uses a packed FP4 physical pair type as a VMI logical " + "element type; packed FP4 input/output is not a supported VMI " + "surface because the logical FP4 lane count and physical packed " + "byte count are ambiguous"; + + if (layout && !mlir::isa(layout)) + return emitError() << "'" + << formatVMIVRegType(elementCount, elementType, layout) + << "' expected layout to be #pto.vmi.layout"; + if (auto layoutAttr = llvm::dyn_cast_or_null(layout)) { + if (layoutAttr.isGroupSlots() && + elementCount != layoutAttr.getNumGroups()) + return emitError() << "'" + << formatVMIVRegType(elementCount, elementType, layout) + << "' expected num_groups layout to describe exactly " + "one logical result lane per group"; + } + + return success(); +} + +bool VMIMaskType::isSupportedGranularity(StringRef granularity) { + return granularity == "pred" || isConcreteGranularity(granularity); +} + +bool VMIMaskType::isConcreteGranularity(StringRef granularity) { + return granularity == "b8" || granularity == "b16" || granularity == "b32"; +} + +Type VMIMaskType::parse(AsmParser &parser) { + SmallVector shape; + StringRef granularity; + Attribute layout; + SMLoc loc = parser.getCurrentLocation(); + + if (failed(parser.parseLess()) || + failed(parser.parseDimensionList(shape, /*allowDynamic=*/false, + /*withTrailingX=*/true)) || + shape.size() != 1 || failed(parser.parseKeyword(&granularity)) || + failed(parseOptionalVMILayout(parser, layout)) || + failed(parser.parseGreater())) + return {}; + + return parser.getChecked(loc, parser.getContext(), shape.front(), + granularity, layout); +} + +void VMIMaskType::print(AsmPrinter &printer) const { + printer << "<" << getElementCount() << "x" << getGranularity(); + if (getLayout()) + printer << ", " << getLayout(); + printer << ">"; +} + +LogicalResult VMIMaskType::verify(function_ref emitError, + int64_t elementCount, StringRef granularity, + Attribute layout) { + if (elementCount <= 0) + return emitError() << "'" + << formatVMIMaskType(elementCount, granularity, layout) + << "' expected a positive element count"; + + if (!isSupportedGranularity(granularity)) + return emitError() << "'" + << formatVMIMaskType(elementCount, granularity, layout) + << "' expected granularity to be one of pred, b8, b16, " + "b32"; + + if (layout && !mlir::isa(layout)) + return emitError() << "'" + << formatVMIMaskType(elementCount, granularity, layout) + << "' expected layout to be #pto.vmi.layout"; + + if (granularity == "pred" && layout) + return emitError() << "'" + << formatVMIMaskType(elementCount, granularity, layout) + << "' pred mask must not carry layout"; + + return success(); +} + +//===----------------------------------------------------------------------===// +// Legacy (old) VMI op verifiers +//===----------------------------------------------------------------------===// + +//===--- Legacy (old) VMI op verifiers ---===// + +LogicalResult VMIConstantOp::verify() { + auto resultType = cast(getResult().getType()); + auto denseAttr = dyn_cast(getValue()); + if (!denseAttr) + return emitOpError("requires dense elements constant attribute"); + if (denseAttr.getElementType() != resultType.getElementType()) + return emitOpError( + "requires dense constant element type to match result element type"); + if (denseAttr.getNumElements() != resultType.getElementCount()) + return emitOpError("requires dense constant element count to match result " + "logical lane count"); + return success(); +} + +LogicalResult VMIBroadcastOp::verify() { + auto resultType = cast(getResult().getType()); + Type valueType = getValue().getType(); + if (valueType == resultType.getElementType()) + return success(); + if (auto vregType = dyn_cast(valueType)) { + if (vregType.getElementCount() != 1) + return emitOpError("requires VMI vector input to have one logical lane"); + if (vregType.getElementType() != resultType.getElementType()) + return emitOpError("requires VMI vector input element type to match " + "result element type"); + return success(); + } + return emitOpError("requires scalar or VMI vector input element type to " + "match result element type"); +} + +LogicalResult VMIIotaOp::verify() { + auto resultType = cast(getResult().getType()); + Type elementType = resultType.getElementType(); + if (!isVMIIotaElementType(elementType)) + return emitOpError("requires result element type to be integer 8/16/32 " + "or f16/f32"); + if (!isCompatibleScalarForSemanticType(elementType, getBase().getType())) + return emitOpError("requires base type to match result element type"); + + if (std::optional order = getOrder()) { + if (*order != "ASC" && *order != "DESC") + return emitOpError("requires order to be ASC or DESC"); + } + return success(); +} + +LogicalResult VMICreateMaskOp::verify() { + return success(); +} + +LogicalResult VMICreateGroupMaskOp::verify() { + auto resultType = cast(getResult().getType()); + int64_t numGroups = getNumGroupsAttr().getInt(); + int64_t groupSize = getGroupSizeAttr().getInt(); + if (numGroups <= 0) + return emitOpError("requires positive num_groups"); + if (groupSize <= 0) + return emitOpError("requires positive group_size"); + if (resultType.getElementCount() != numGroups * groupSize) + return emitOpError("requires result lane count to equal num_groups * " + "group_size"); + return success(); +} + +LogicalResult VMIConstantMaskOp::verify() { + auto resultType = cast(getResult().getType()); + auto denseAttr = dyn_cast(getValue()); + if (!denseAttr) + return emitOpError("requires dense elements mask constant attribute"); + if (!denseAttr.getElementType().isInteger(1)) + return emitOpError("requires dense mask constant element type to be i1"); + if (denseAttr.getNumElements() != resultType.getElementCount()) + return emitOpError("requires dense mask constant element count to match " + "result logical lane count"); + return success(); +} + +LogicalResult VMIMaskAndOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + return verifyAllSameMaskShapeLayoutAndGranularity( + getOperation(), {lhsType, rhsType, resultType}); +} + +LogicalResult VMIMaskOrOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + return verifyAllSameMaskShapeLayoutAndGranularity( + getOperation(), {lhsType, rhsType, resultType}); +} + +LogicalResult VMIMaskXOrOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + return verifyAllSameMaskShapeLayoutAndGranularity( + getOperation(), {lhsType, rhsType, resultType}); +} + +LogicalResult VMIMaskNotOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + return verifyAllSameMaskShapeLayoutAndGranularity(getOperation(), + {sourceType, resultType}); +} + +LogicalResult VMIAddFOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIFloatLikeType(lhsType.getElementType())) + return emitOpError("requires floating-point-like VMI element type"); + if (!isVMIF16BF16OrF32Type(lhsType.getElementType())) + return emitOpError("requires f16, bf16, or f32 VMI element type"); + return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); +} + +LogicalResult VMIAddIOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIIntegerLikeType(lhsType.getElementType())) + return emitOpError("requires integer-like VMI element type"); + if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) + return emitOpError("requires i8, i16, or i32 VMI element type"); + return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); +} + +LogicalResult VMISubFOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIFloatLikeType(lhsType.getElementType())) + return emitOpError("requires floating-point-like VMI element type"); + if (!isVMIF16BF16OrF32Type(lhsType.getElementType())) + return emitOpError("requires f16, bf16, or f32 VMI element type"); + return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); +} + +LogicalResult VMISubIOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIIntegerLikeType(lhsType.getElementType())) + return emitOpError("requires integer-like VMI element type"); + if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) + return emitOpError("requires i8, i16, or i32 VMI element type"); + return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); +} + +LogicalResult VMIMulFOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIFloatLikeType(lhsType.getElementType())) + return emitOpError("requires floating-point-like VMI element type"); + if (!isVMIF16BF16OrF32Type(lhsType.getElementType())) + return emitOpError("requires f16, bf16, or f32 VMI element type"); + return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); +} + +LogicalResult VMIMulIOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIIntegerLikeType(lhsType.getElementType())) + return emitOpError("requires integer-like VMI element type"); + if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) + return emitOpError("requires i8, i16, or i32 VMI element type"); + return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); +} + +LogicalResult VMIFmaOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto accType = cast(getAcc().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIF16BF16OrF32Type(lhsType.getElementType())) + return emitOpError("requires f16, bf16, or f32 VMI element type"); + return verifyFloatTernaryVRegOp(getOperation(), lhsType, rhsType, accType, + resultType); +} + +//===----------------------------------------------------------------------===// +// Legacy elementwise op verifiers (restored for backward compatibility). +//===----------------------------------------------------------------------===// + +LogicalResult VMIDivFOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIFloatLikeType(lhsType.getElementType())) + return emitOpError("requires floating-point-like VMI element type"); + if (!isVMIF16OrF32Type(lhsType.getElementType())) + return emitOpError("requires f16 or f32 VMI element type"); + return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); +} + +LogicalResult VMIMinFOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIFloatLikeType(lhsType.getElementType())) + return emitOpError("requires floating-point-like VMI element type"); + if (!isVMIF16BF16OrF32Type(lhsType.getElementType())) + return emitOpError("requires f16, bf16, or f32 VMI element type"); + return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); +} + +LogicalResult VMIMaxFOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIFloatLikeType(lhsType.getElementType())) + return emitOpError("requires floating-point-like VMI element type"); + if (!isVMIF16BF16OrF32Type(lhsType.getElementType())) + return emitOpError("requires f16, bf16, or f32 VMI element type"); + return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); +} + +LogicalResult VMINegFOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIF16OrF32Type(sourceType.getElementType())) + return emitOpError("requires f16 or f32 VMI element type"); + return verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType); +} + +LogicalResult VMIAbsFOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIF16OrF32Type(sourceType.getElementType())) + return emitOpError("requires f16 or f32 VMI element type"); + return verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType); +} + +LogicalResult VMIAbsIOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIIntegerLikeType(sourceType.getElementType())) + return emitOpError("requires integer-like VMI element type"); + if (!isVMISignedOrSignlessI8I16I32Type(sourceType.getElementType())) + return emitOpError("requires signless or signed i8, i16, or i32 VMI " + "element type"); + return verifyAllSameVRegShapeAndLayout(getOperation(), + {sourceType, resultType}, + /*requireSameElement=*/true); +} + +LogicalResult VMISqrtOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIF16OrF32Type(sourceType.getElementType())) + return emitOpError("requires f16 or f32 VMI element type"); + return verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType); +} + +LogicalResult VMIExpOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIF16OrF32Type(sourceType.getElementType())) + return emitOpError("requires f16 or f32 VMI element type"); + return verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType); +} + +LogicalResult VMILnOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIF16OrF32Type(sourceType.getElementType())) + return emitOpError("requires f16 or f32 VMI element type"); + return verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType); +} + +LogicalResult VMIReluOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIF16OrF32Type(sourceType.getElementType())) + return emitOpError("requires f16 or f32 VMI element type"); + return verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType); +} + +LogicalResult VMIAndIOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIIntegerLikeType(lhsType.getElementType())) + return emitOpError("requires integer-like VMI element type"); + if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) + return emitOpError("requires i8, i16, or i32 VMI element type"); + return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); +} + +LogicalResult VMIOrIOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIIntegerLikeType(lhsType.getElementType())) + return emitOpError("requires integer-like VMI element type"); + if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) + return emitOpError("requires i8, i16, or i32 VMI element type"); + return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); +} + +LogicalResult VMIXOrIOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIIntegerLikeType(lhsType.getElementType())) + return emitOpError("requires integer-like VMI element type"); + if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) + return emitOpError("requires i8, i16, or i32 VMI element type"); + return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); +} + +LogicalResult VMIShLIOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIIntegerLikeType(lhsType.getElementType())) + return emitOpError("requires integer-like VMI element type"); + if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) + return emitOpError("requires i8, i16, or i32 VMI element type"); + return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); +} + +LogicalResult VMIShRUIOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + auto integerType = dyn_cast(lhsType.getElementType()); + if (!integerType || integerType.isSigned()) + return emitOpError( + "requires signless or unsigned integer VMI element type"); + if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) + return emitOpError("requires i8, i16, or i32 VMI element type"); + return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); +} + +LogicalResult VMIShRSIOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMISignedOrSignlessI8I16I32Type(lhsType.getElementType())) + return emitOpError( + "requires signless or signed i8, i16, or i32 VMI element type"); + return verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType); +} + +LogicalResult VMINotOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIIntegerLikeType(sourceType.getElementType())) + return emitOpError("requires integer-like VMI element type"); + if (!isVMIAnyI8I16I32Type(sourceType.getElementType())) + return emitOpError("requires i8, i16, or i32 VMI element type"); + return verifyAllSameVRegShapeAndLayout(getOperation(), + {sourceType, resultType}, + /*requireSameElement=*/true); +} + +//===----------------------------------------------------------------------===// + +LogicalResult VMICmpFOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIFloatLikeType(lhsType.getElementType())) + return emitOpError("requires floating-point-like VMI element type"); + if (!isVMIF16BF16OrF32Type(lhsType.getElementType())) + return emitOpError("requires f16, bf16, or f32 VMI element type"); + if (failed(verifyAllSameVRegShapeAndLayout(getOperation(), {lhsType, rhsType}, + /*requireSameElement=*/true))) + return failure(); + return verifyMaskMatchesData(getOperation(), resultType, lhsType); +} + +LogicalResult VMICmpIOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIIntegerLikeType(lhsType.getElementType())) + return emitOpError("requires integer-like VMI element type"); + if (!isVMIAnyI8I16I32Type(lhsType.getElementType())) + return emitOpError("requires i8, i16, or i32 VMI element type"); + if (failed(verifyAllSameVRegShapeAndLayout(getOperation(), {lhsType, rhsType}, + /*requireSameElement=*/true))) + return failure(); + return verifyMaskMatchesData(getOperation(), resultType, lhsType); +} + +LogicalResult VMISelectOp::verify() { + auto maskType = cast(getMask().getType()); + auto trueType = cast(getTrueValue().getType()); + auto falseType = cast(getFalseValue().getType()); + auto resultType = cast(getResult().getType()); + if (failed(verifyAllSameVRegShapeAndLayout(getOperation(), + {trueType, falseType, resultType}, + /*requireSameElement=*/true))) + return failure(); + return verifyMaskMatchesData(getOperation(), maskType, resultType); +} + +LogicalResult VMIActivePrefixIndexOp::verify() { + auto maskType = cast(getMask().getType()); + auto resultType = cast(getResult().getType()); + auto resultIntType = dyn_cast(resultType.getElementType()); + if (!resultIntType || !resultIntType.isSignless()) + return emitOpError("requires signless integer result element type"); + unsigned resultWidth = resultIntType.getWidth(); + if (resultWidth != 8 && resultWidth != 16 && resultWidth != 32) + return emitOpError("requires i8, i16, or i32 result element type"); + return verifyMaskMatchesData(getOperation(), maskType, resultType); +} + +LogicalResult VMICompressOp::verify() { + auto sourceType = cast(getSource().getType()); + auto maskType = cast(getMask().getType()); + auto resultType = cast(getResult().getType()); + if (failed(verifyAllSameVRegShapeAndLayout(getOperation(), + {sourceType, resultType}, + /*requireSameElement=*/true))) + return failure(); + return verifyMaskMatchesData(getOperation(), maskType, sourceType); +} + +LogicalResult VMICompressStoreOp::verify() { + auto valueType = cast(getValue().getType()); + auto maskType = cast(getMask().getType()); + if (failed(verifyMemoryElementMatches(getOperation(), + getDestination().getType(), valueType, + "destination"))) + return failure(); + if (failed(verifyUBBackedMemory(getOperation(), getDestination().getType(), + "destination"))) + return failure(); + return verifyMaskMatchesData(getOperation(), maskType, valueType); +} + +void VMICompressStoreOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Write::get(), &getDestinationMutable()); +} + +LogicalResult VMIReduceAddIOp::verify() { + auto sourceType = cast(getSource().getType()); + auto initType = cast(getInit().getType()); + auto maskType = cast(getMask().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIIntegerLikeType(sourceType.getElementType())) + return emitOpError("requires integer-like VMI source element type"); + auto sourceIntegerType = dyn_cast(sourceType.getElementType()); + if (!sourceIntegerType || sourceIntegerType.getWidth() != 32) + return emitOpError("requires 32-bit integer source element type"); + if (sourceType.getElementType() != initType.getElementType() || + sourceType.getElementType() != resultType.getElementType()) + return emitOpError( + "requires source, init, and result element types to match"); + if (initType.getElementCount() != 1 || resultType.getElementCount() != 1) + return emitOpError("requires init and result to be 1-lane VMI vectors"); + if (failed(verifyAllSameVRegShapeAndLayout(getOperation(), + {initType, resultType}, + /*requireSameElement=*/true))) + return failure(); + return verifyMaskMatchesData(getOperation(), maskType, sourceType); +} + +LogicalResult VMIReduceAddFOp::verify() { + auto sourceType = cast(getSource().getType()); + auto initType = cast(getInit().getType()); + auto maskType = cast(getMask().getType()); + auto resultType = cast(getResult().getType()); + if (!getOperation()->hasAttr("reassoc")) + return emitOpError( + "requires reassoc attr because VPTO vcadd performs pair-wise " + "floating-point reduction"); + if (!isVMIFloatLikeType(sourceType.getElementType())) + return emitOpError("requires floating-point-like VMI source element type"); + if (!isVMIF16OrF32Type(sourceType.getElementType())) + return emitOpError("requires f16 or f32 source element type"); + if (sourceType.getElementType() != initType.getElementType() || + sourceType.getElementType() != resultType.getElementType()) + return emitOpError( + "requires source, init, and result element types to match"); + if (initType.getElementCount() != 1 || resultType.getElementCount() != 1) + return emitOpError("requires init and result to be 1-lane VMI vectors"); + if (failed(verifyAllSameVRegShapeAndLayout(getOperation(), + {initType, resultType}, + /*requireSameElement=*/true))) + return failure(); + return verifyMaskMatchesData(getOperation(), maskType, sourceType); +} + +template LogicalResult verifyReduceMinMaxFOp(OpTy op) { + auto sourceType = cast(op.getSource().getType()); + auto initType = cast(op.getInit().getType()); + auto maskType = cast(op.getMask().getType()); + auto resultType = cast(op.getResult().getType()); + if (!isVMIFloatLikeType(sourceType.getElementType())) + return op.emitOpError( + "requires floating-point-like VMI source element type"); + if (!isVMIF16OrF32Type(sourceType.getElementType())) + return op.emitOpError("requires f16 or f32 source element type"); + if (sourceType.getElementType() != initType.getElementType() || + sourceType.getElementType() != resultType.getElementType()) + return op.emitOpError( + "requires source, init, and result element types to match"); + if (initType.getElementCount() != 1 || resultType.getElementCount() != 1) + return op.emitOpError("requires init and result to be 1-lane VMI vectors"); + if (failed(verifyAllSameVRegShapeAndLayout(op.getOperation(), + {initType, resultType}, + /*requireSameElement=*/true))) + return failure(); + return verifyMaskMatchesData(op.getOperation(), maskType, sourceType); +} + +LogicalResult VMIReduceMaxFOp::verify() { return verifyReduceMinMaxFOp(*this); } + +LogicalResult VMIReduceMinFOp::verify() { return verifyReduceMinMaxFOp(*this); } + +template LogicalResult verifyReduceMinMaxIOp(OpTy op) { + auto sourceType = cast(op.getSource().getType()); + auto initType = cast(op.getInit().getType()); + auto maskType = cast(op.getMask().getType()); + auto resultType = cast(op.getResult().getType()); + auto sourceIntegerType = dyn_cast(sourceType.getElementType()); + if (!sourceIntegerType || + !isVMIAnyI8I16I32Type(sourceType.getElementType())) + return op.emitOpError( + "requires 8-bit, 16-bit, or 32-bit integer source element type"); + if (sourceType.getElementType() != initType.getElementType() || + sourceType.getElementType() != resultType.getElementType()) + return op.emitOpError( + "requires source, init, and result element types to match"); + if (initType.getElementCount() != 1 || resultType.getElementCount() != 1) + return op.emitOpError("requires init and result to be 1-lane VMI vectors"); + if (failed(verifyAllSameVRegShapeAndLayout(op.getOperation(), + {initType, resultType}, + /*requireSameElement=*/true))) + return failure(); + return verifyMaskMatchesData(op.getOperation(), maskType, sourceType); +} + +LogicalResult VMIReduceMaxIOp::verify() { return verifyReduceMinMaxIOp(*this); } + +LogicalResult VMIReduceMinIOp::verify() { return verifyReduceMinMaxIOp(*this); } + +template +static LogicalResult verifyGroupReduceFloatOp(OpTy op, bool requiresReassoc) { + auto sourceType = cast(op.getSource().getType()); + auto maskType = cast(op.getMask().getType()); + auto resultType = cast(op.getResult().getType()); + if (requiresReassoc && !op->hasAttr("reassoc")) + return op.emitOpError( + "requires reassoc attr because grouped lowering uses pair-wise " + "floating-point reductions"); + if (!isVMIFloatLikeType(sourceType.getElementType())) + return op.emitOpError( + "requires floating-point-like VMI source element type"); + if (!isVMIF16OrF32Type(sourceType.getElementType())) + return op.emitOpError("requires f16 or f32 source element type"); + if (resultType.getElementCount() != op.getNumGroupsAttr().getInt()) + return op.emitOpError( + "requires result logical lane count to match num_groups"); + if (sourceType.getElementType() != resultType.getElementType()) + return op.emitOpError("requires source and result element types to match"); + if (auto sourceLayout = sourceType.getLayoutAttr()) { + bool supportedSourceLayout = + sourceLayout.isContiguous() || + (sourceLayout.isDeinterleaved() && sourceLayout.getFactor() == 2 && + (sourceLayout.getBlockElems() == 1 || + sourceLayout.getBlockElems() == 8)) || + (sourceLayout.isDeinterleaved() && sourceLayout.getFactor() == 4 && + (sourceLayout.getBlockElems() == 1 || + sourceLayout.getBlockElems() == 8)); + if (!supportedSourceLayout) + return op.emitOpError( + "requires layout-assigned source to use contiguous layout or " + "deinterleaved=2/4 layout with block_elems=1 or block_elems=8"); + } + if (auto resultLayout = resultType.getLayoutAttr()) { + if (!resultLayout.isGroupSlots() || + resultLayout.getNumGroups() != op.getNumGroupsAttr().getInt()) + return op.emitOpError() << "requires layout-assigned result to use " + "#pto.vmi.layout"; + } + if (failed(verifyMaskMatchesData(op.getOperation(), maskType, sourceType))) + return failure(); + return verifyNumGroups(op.getOperation(), sourceType, + op.getNumGroupsAttr().getInt()); +} + +LogicalResult VMIGroupReduceAddFOp::verify() { + return verifyGroupReduceFloatOp(*this, /*requiresReassoc=*/true); +} + +LogicalResult VMIGroupReduceMaxFOp::verify() { + return verifyGroupReduceFloatOp(*this, /*requiresReassoc=*/false); +} + +LogicalResult VMIGroupReduceMinFOp::verify() { + return verifyGroupReduceFloatOp(*this, /*requiresReassoc=*/false); +} + +template +static LogicalResult verifyGroupReduceIntegerOp(OpTy op) { + auto sourceType = cast(op.getSource().getType()); + auto maskType = cast(op.getMask().getType()); + auto resultType = cast(op.getResult().getType()); + if (!isVMIIntegerLikeType(sourceType.getElementType())) + return op.emitOpError("requires integer-like VMI source element type"); + auto intType = dyn_cast(sourceType.getElementType()); + if (!intType || !isVMIAnyI8I16I32Type(sourceType.getElementType())) + return op.emitOpError( + "requires 8-bit, 16-bit, or 32-bit integer source element type"); + if (resultType.getElementCount() != op.getNumGroupsAttr().getInt()) + return op.emitOpError( + "requires result logical lane count to match num_groups"); + if (sourceType.getElementType() != resultType.getElementType()) + return op.emitOpError("requires source and result element types to match"); + if (auto sourceLayout = sourceType.getLayoutAttr()) { + bool supportedSourceLayout = + sourceLayout.isContiguous() || + (sourceLayout.isDeinterleaved() && sourceLayout.getFactor() == 2 && + (sourceLayout.getBlockElems() == 1 || + sourceLayout.getBlockElems() == 8)) || + (sourceLayout.isDeinterleaved() && sourceLayout.getFactor() == 4 && + (sourceLayout.getBlockElems() == 1 || + sourceLayout.getBlockElems() == 8)); + if (!supportedSourceLayout) + return op.emitOpError( + "requires layout-assigned source to use contiguous layout or " + "deinterleaved=2/4 layout with block_elems=1 or block_elems=8"); + } + if (auto resultLayout = resultType.getLayoutAttr()) { + if (!resultLayout.isGroupSlots() || + resultLayout.getNumGroups() != op.getNumGroupsAttr().getInt()) + return op.emitOpError() << "requires layout-assigned result to use " + "#pto.vmi.layout"; + } + if (failed(verifyMaskMatchesData(op.getOperation(), maskType, sourceType))) + return failure(); + return verifyNumGroups(op.getOperation(), sourceType, + op.getNumGroupsAttr().getInt()); +} + +LogicalResult VMIGroupReduceAddIOp::verify() { + return verifyGroupReduceIntegerOp(*this); +} + +LogicalResult VMIGroupReduceMaxIOp::verify() { + return verifyGroupReduceIntegerOp(*this); +} + +LogicalResult VMIGroupReduceMinIOp::verify() { + return verifyGroupReduceIntegerOp(*this); +} + +//===----------------------------------------------------------------------===// +// Group 5: vcadd / vcmax / vcmin verifiers +//===----------------------------------------------------------------------===// + +LogicalResult VMIGroupBroadcastOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + int64_t numGroups = getNumGroupsAttr().getInt(); + if (sourceType.getElementCount() != numGroups) + return emitOpError( + "requires source logical lane count to match num_groups"); + if (resultType.getElementCount() % numGroups != 0) + return emitOpError( + "requires num_groups to evenly divide result logical lane count"); + if (sourceType.getElementType() != resultType.getElementType()) + return emitOpError("requires source and result element types to match"); + if (auto sourceLayout = sourceType.getLayoutAttr()) { + if (!sourceLayout.isGroupSlots() || + sourceLayout.getNumGroups() != numGroups) + return emitOpError() << "requires layout-assigned source to use " + "#pto.vmi.layout"; + } + if (auto resultLayout = resultType.getLayoutAttr()) { + if (resultLayout.isGroupSlots()) + return emitOpError( + "requires layout-assigned result to use a dense VMI layout"); + } + return verifyNumGroups(getOperation(), resultType, numGroups); +} + +template static LogicalResult verifyVMIHistogramOp(OpTy op) { + auto accType = cast(op.getAcc().getType()); + auto sourceType = cast(op.getSource().getType()); + auto maskType = cast(op.getMask().getType()); + auto resultType = cast(op.getResult().getType()); + + auto accElemType = dyn_cast(accType.getElementType()); + auto sourceElemType = dyn_cast(sourceType.getElementType()); + if (!accElemType || !accElemType.isUnsigned() || + accElemType.getWidth() != 16 || accType.getElementCount() != 256) + return op.emitOpError("requires acc type to be " + "!pto.vmi.vreg<256xui16>"); + if (resultType != accType) + return op.emitOpError("requires result type to match acc type"); + if (!sourceElemType || !sourceElemType.isUnsigned() || + sourceElemType.getWidth() != 8) + return op.emitOpError("requires source type to be " + "!pto.vmi.vreg"); + if (maskType.getElementCount() != sourceType.getElementCount()) + return op.emitOpError("requires mask logical lane count to match source"); + + if (auto accLayout = accType.getLayoutAttr()) { + if (!accLayout.isContiguous()) + return op.emitOpError("requires layout-assigned acc to use contiguous " + "layout"); + } + if (auto sourceLayout = sourceType.getLayoutAttr()) { + if (!sourceLayout.isContiguous()) + return op.emitOpError("requires layout-assigned source to use contiguous " + "layout"); + } + if (auto resultLayout = resultType.getLayoutAttr()) { + if (!resultLayout.isContiguous()) + return op.emitOpError("requires layout-assigned result to use " + "contiguous layout"); + } + if (auto maskLayout = maskType.getLayoutAttr()) { + if (!maskLayout.isContiguous()) + return op.emitOpError("requires layout-assigned mask to use contiguous " + "layout"); + if (maskType.getGranularity() != "b8") + return op.emitOpError("requires layout-assigned mask granularity b8"); + } + return success(); +} + +LogicalResult VMIVdhistOp::verify() { return verifyVMIHistogramOp(*this); } + +LogicalResult VMIVchistOp::verify() { return verifyVMIHistogramOp(*this); } + +//===----------------------------------------------------------------------===// +// Group 7: SFU verifiers +//===----------------------------------------------------------------------===// + +LogicalResult VMIExtFOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + if (sourceType.getElementCount() != resultType.getElementCount()) + return emitOpError( + "requires source and result logical lane counts to match"); + if (!isVMIFloatLikeType(sourceType.getElementType()) || + !isVMIFloatLikeType(resultType.getElementType())) + return emitOpError( + "requires floating-point-like source and result element types"); + if (getVMIElementBitWidth(sourceType.getElementType()) >= + getVMIElementBitWidth(resultType.getElementType())) + return emitOpError( + "requires result element type to be wider than source element type"); + return success(); +} + +LogicalResult VMITruncFOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + if (sourceType.getElementCount() != resultType.getElementCount()) + return emitOpError( + "requires source and result logical lane counts to match"); + if (!isVMIFloatLikeType(sourceType.getElementType()) || + !isVMIFloatLikeType(resultType.getElementType())) + return emitOpError( + "requires floating-point-like source and result element types"); + if (getVMIElementBitWidth(sourceType.getElementType()) <= + getVMIElementBitWidth(resultType.getElementType())) + return emitOpError( + "requires result element type to be narrower than source element type"); + if (auto roundingAttr = (*this)->getAttrOfType("rounding")) { + StringRef rounding = roundingAttr.getValue(); + if (rounding != "R" && rounding != "A" && rounding != "H" && + rounding != "Z") + return emitOpError("rounding attr must be R, A, H, or Z"); + } + return success(); +} + +LogicalResult VMIFPToSIOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + if (sourceType.getElementCount() != resultType.getElementCount()) + return emitOpError( + "requires source and result logical lane counts to match"); + if (!isVMIFloatLikeType(sourceType.getElementType())) + return emitOpError("requires floating-point-like source element type"); + if (!isVMISignedOrSignlessIntegerType(resultType.getElementType())) + return emitOpError("requires signed or signless integer result element " + "type"); + if (getVMIElementBitWidth(resultType.getElementType()) != 32) + return emitOpError("requires 32-bit integer result element type"); + return success(); +} + +LogicalResult VMISIToFPOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + if (sourceType.getElementCount() != resultType.getElementCount()) + return emitOpError( + "requires source and result logical lane counts to match"); + if (!isVMISignedOrSignlessIntegerType(sourceType.getElementType())) + return emitOpError( + "requires signed or signless integer source element type"); + if (!isVMIFloatLikeType(resultType.getElementType())) + return emitOpError("requires floating-point-like result element type"); + if (getVMIElementBitWidth(sourceType.getElementType()) != 32) + return emitOpError("requires 32-bit integer source element type"); + if (!resultType.getElementType().isF32()) + return emitOpError("requires f32 result element type"); + return success(); +} + +LogicalResult VMIExtSIOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + if (sourceType.getElementCount() != resultType.getElementCount()) + return emitOpError( + "requires source and result logical lane counts to match"); + if (!isVMISignedOrSignlessIntegerType(sourceType.getElementType()) || + !isVMISignedOrSignlessIntegerType(resultType.getElementType())) + return emitOpError( + "requires signed or signless integer source and result element types"); + if (getVMIElementBitWidth(sourceType.getElementType()) >= + getVMIElementBitWidth(resultType.getElementType())) + return emitOpError( + "requires result element type to be wider than source element type"); + return success(); +} + +LogicalResult VMIExtUIOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + if (sourceType.getElementCount() != resultType.getElementCount()) + return emitOpError( + "requires source and result logical lane counts to match"); + if (!isVMIUnsignedIntegerType(sourceType.getElementType()) || + !isVMIUnsignedIntegerType(resultType.getElementType())) + return emitOpError( + "requires unsigned integer source and result element types"); + if (getVMIElementBitWidth(sourceType.getElementType()) >= + getVMIElementBitWidth(resultType.getElementType())) + return emitOpError( + "requires result element type to be wider than source element type"); + return success(); +} + +LogicalResult VMITruncIOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + if (sourceType.getElementCount() != resultType.getElementCount()) + return emitOpError( + "requires source and result logical lane counts to match"); + if (!isVMIIntegerLikeType(sourceType.getElementType()) || + !isVMIIntegerLikeType(resultType.getElementType())) + return emitOpError("requires integer source and result element types"); + if (getVMIElementBitWidth(sourceType.getElementType()) <= + getVMIElementBitWidth(resultType.getElementType())) + return emitOpError( + "requires result element type to be narrower than source element type"); + return success(); +} + +LogicalResult VMIBitcastOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + std::optional sourceBits = + getVMIIntegerOrFloatBitWidth(sourceType.getElementType()); + std::optional resultBits = + getVMIIntegerOrFloatBitWidth(resultType.getElementType()); + if (!sourceBits || !resultBits) + return emitOpError( + "requires integer or floating-point source and result element types"); + if (sourceType.getElementCount() * static_cast(*sourceBits) != + resultType.getElementCount() * static_cast(*resultBits)) + return emitOpError( + "requires source and result to carry the same total number of bits"); + + if (isLayoutAssigned(sourceType) || isLayoutAssigned(resultType)) { + if (!isLayoutAssigned(sourceType) || !isLayoutAssigned(resultType)) + return emitOpError( + "requires either both source and result to carry layout or neither " + "to carry layout"); + if (sourceType.getLayout() != resultType.getLayout()) + return emitOpError("requires source and result layouts to match"); + } + + return success(); +} + +LogicalResult VMILoadOp::verify() { + if (failed(verifyMemoryElementMatches( + getOperation(), getSource().getType(), + cast(getResult().getType()), "source"))) + return failure(); + return verifyUBBackedMemory(getOperation(), getSource().getType(), "source"); +} + +void VMILoadOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Read::get(), &getSourceMutable()); +} + +LogicalResult VMIDeinterleaveLoadOp::verify() { + auto lowType = cast(getLow().getType()); + auto highType = cast(getHigh().getType()); + if (failed(verifyAllSameVRegShapeAndLayout(getOperation(), + {lowType, highType}, + /*requireSameElement=*/true))) + return failure(); + if (failed(verifyMemoryElementMatches(getOperation(), getSource().getType(), + lowType, "source"))) + return failure(); + if (failed(verifyUBBackedMemory(getOperation(), getSource().getType(), + "source"))) + return failure(); + if (failed(verifyContiguousIfLayoutAssigned(getOperation(), lowType, + "low result")) || + failed(verifyContiguousIfLayoutAssigned(getOperation(), highType, + "high result"))) + return failure(); + return success(); +} + +void VMIDeinterleaveLoadOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Read::get(), &getSourceMutable()); +} + +LogicalResult VMIGroupLoadOp::verify() { + auto resultType = cast(getResult().getType()); + if (failed(verifyMemoryElementMatches(getOperation(), getSource().getType(), + resultType, "source"))) + return failure(); + if (failed(verifyUBBackedMemory(getOperation(), getSource().getType(), + "source"))) + return failure(); + return verifyNumGroups(getOperation(), resultType, + getNumGroupsAttr().getInt()); +} + +void VMIGroupLoadOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Read::get(), &getSourceMutable()); +} + +LogicalResult VMIGroupSlotLoadOp::verify() { + auto resultType = cast(getResult().getType()); + int64_t numGroups = getNumGroupsAttr().getInt(); + if (resultType.getElementCount() != numGroups) + return emitOpError( + "requires result logical lane count to match num_groups"); + if (failed(verifyMemoryElementMatches(getOperation(), getSource().getType(), + resultType, "source"))) + return failure(); + if (failed(verifyUBBackedMemory(getOperation(), getSource().getType(), + "source"))) + return failure(); + if (auto resultLayout = resultType.getLayoutAttr()) { + if (!resultLayout.isGroupSlots() || + resultLayout.getNumGroups() != numGroups) + return emitOpError() << "requires layout-assigned result to use " + "#pto.vmi.layout"; + } + return verifyNumGroups(getOperation(), resultType, numGroups); +} + +void VMIGroupSlotLoadOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Read::get(), &getSourceMutable()); +} + +LogicalResult VMIGroupBroadcastLoadOp::verify() { + auto resultType = cast(getResult().getType()); + int64_t numGroups = getNumGroupsAttr().getInt(); + if (numGroups <= 0) + return emitOpError("requires num_groups to be positive"); + if (resultType.getElementCount() % numGroups != 0) + return emitOpError( + "requires num_groups to evenly divide result logical lane count"); + if (failed(verifyMemoryElementMatches(getOperation(), getSource().getType(), + resultType, "source"))) + return failure(); + if (failed(verifyUBBackedMemory(getOperation(), getSource().getType(), + "source"))) + return failure(); + if (auto resultLayout = resultType.getLayoutAttr()) { + if (resultLayout.isGroupSlots()) + return emitOpError( + "requires layout-assigned result to use a dense VMI layout"); + } + return verifyNumGroups(getOperation(), resultType, numGroups); +} + +void VMIGroupBroadcastLoadOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Read::get(), &getSourceMutable()); +} + +LogicalResult VMIMaskedLoadOp::verify() { + auto maskType = cast(getMask().getType()); + auto passthruType = cast(getPassthru().getType()); + auto resultType = cast(getResult().getType()); + if (failed(verifyMemoryElementMatches(getOperation(), getSource().getType(), + resultType, "source"))) + return failure(); + if (failed(verifyUBBackedMemory(getOperation(), getSource().getType(), + "source"))) + return failure(); + if (failed(verifyAllSameVRegShapeAndLayout(getOperation(), + {passthruType, resultType}, + /*requireSameElement=*/true))) + return failure(); + return verifyMaskMatchesData(getOperation(), maskType, resultType); +} + +void VMIMaskedLoadOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Read::get(), &getSourceMutable()); +} + +LogicalResult VMIGatherOp::verify() { + auto indicesType = cast(getIndices().getType()); + auto maskType = cast(getMask().getType()); + auto passthruType = cast(getPassthru().getType()); + auto resultType = cast(getResult().getType()); + if (failed(verifyMemoryElementMatches(getOperation(), getSource().getType(), + resultType, "source"))) + return failure(); + if (failed(verifyUBBackedMemory(getOperation(), getSource().getType(), + "source"))) + return failure(); + + auto indexElementType = dyn_cast(indicesType.getElementType()); + if (!indexElementType || indexElementType.isSigned() || + (indexElementType.getWidth() != 16 && indexElementType.getWidth() != 32)) + return emitOpError( + "requires signless or unsigned 16-bit or 32-bit integer indices"); + + if (failed(verifyAllSameVRegShapeAndLayout( + getOperation(), {indicesType, passthruType, resultType}, + /*requireSameElement=*/false))) + return failure(); + if (failed(verifyAllSameVRegShapeAndLayout(getOperation(), + {passthruType, resultType}, + /*requireSameElement=*/true))) + return failure(); + + auto resultIntegerType = dyn_cast(resultType.getElementType()); + if (indexElementType.getWidth() == 16 && + (!resultIntegerType || !resultIntegerType.isUnsigned() || + resultIntegerType.getWidth() != 16)) + return emitOpError( + "requires ui16 result and passthru element type when using ui16 " + "indices"); + return verifyMaskMatchesData(getOperation(), maskType, resultType); +} + +void VMIGatherOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Read::get(), &getSourceMutable()); +} + +LogicalResult VMIExpandLoadOp::verify() { + auto maskType = cast(getMask().getType()); + auto passthruType = cast(getPassthru().getType()); + auto resultType = cast(getResult().getType()); + if (failed(verifyMemoryElementMatches(getOperation(), getSource().getType(), + resultType, "source"))) + return failure(); + if (failed(verifyUBBackedMemory(getOperation(), getSource().getType(), + "source"))) + return failure(); + if (failed(verifyAllSameVRegShapeAndLayout(getOperation(), + {passthruType, resultType}, + /*requireSameElement=*/true))) + return failure(); + return verifyMaskMatchesData(getOperation(), maskType, resultType); +} + +void VMIExpandLoadOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Read::get(), &getSourceMutable()); +} + +LogicalResult VMIStoreOp::verify() { + if (failed(verifyMemoryElementMatches( + getOperation(), getDestination().getType(), + cast(getValue().getType()), "destination"))) + return failure(); + return verifyUBBackedMemory(getOperation(), getDestination().getType(), + "destination"); +} + +void VMIStoreOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Write::get(), &getDestinationMutable()); +} + +LogicalResult VMIInterleaveStoreOp::verify() { + auto lowType = cast(getLow().getType()); + auto highType = cast(getHigh().getType()); + if (failed(verifyAllSameVRegShapeAndLayout(getOperation(), + {lowType, highType}, + /*requireSameElement=*/true))) + return failure(); + if (failed(verifyMemoryElementMatches(getOperation(), + getDestination().getType(), lowType, + "destination"))) + return failure(); + if (failed(verifyUBBackedMemory(getOperation(), getDestination().getType(), + "destination"))) + return failure(); + if (failed(verifyContiguousIfLayoutAssigned(getOperation(), lowType, + "low input")) || + failed(verifyContiguousIfLayoutAssigned(getOperation(), highType, + "high input"))) + return failure(); + return success(); +} + +void VMIInterleaveStoreOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Write::get(), &getDestinationMutable()); +} + +LogicalResult VMIGroupStoreOp::verify() { + auto valueType = cast(getValue().getType()); + if (!isPackedByteGroupStore(getDestination().getType(), valueType) && + failed(verifyMemoryElementMatches(getOperation(), + getDestination().getType(), valueType, + "destination"))) + return failure(); + if (failed(verifyUBBackedMemory(getOperation(), getDestination().getType(), + "destination"))) + return failure(); + return verifyNumGroups(getOperation(), valueType, + getNumGroupsAttr().getInt()); +} + +void VMIGroupStoreOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Write::get(), &getDestinationMutable()); +} + +LogicalResult VMIStrideLoadOp::verify() { + auto resultType = cast(getResult().getType()); + auto maskType = cast(getMask().getType()); + if (failed(verifyMemoryElementMatches(getOperation(), getSource().getType(), + resultType, "source"))) + return failure(); + if (failed(verifyUBBackedMemory(getOperation(), getSource().getType(), + "source"))) + return failure(); + return verifyMaskMatchesData(getOperation(), maskType, resultType); +} + +void VMIStrideLoadOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Read::get(), &getSourceMutable()); +} + +LogicalResult VMIMaskedStoreOp::verify() { + auto valueType = cast(getValue().getType()); + auto maskType = cast(getMask().getType()); + if (failed(verifyMemoryElementMatches(getOperation(), + getDestination().getType(), valueType, + "destination"))) + return failure(); + if (failed(verifyUBBackedMemory(getOperation(), getDestination().getType(), + "destination"))) + return failure(); + return verifyMaskMatchesData(getOperation(), maskType, valueType); +} + +void VMIMaskedStoreOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Write::get(), &getDestinationMutable()); +} + +LogicalResult VMIStrideStoreOp::verify() { + auto valueType = cast(getValue().getType()); + auto maskType = cast(getMask().getType()); + if (failed(verifyMemoryElementMatches(getOperation(), + getDestination().getType(), valueType, + "destination"))) + return failure(); + if (failed(verifyUBBackedMemory(getOperation(), getDestination().getType(), + "destination"))) + return failure(); + return verifyMaskMatchesData(getOperation(), maskType, valueType); +} + +void VMIStrideStoreOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Write::get(), &getDestinationMutable()); +} + +//===----------------------------------------------------------------------===// + +LogicalResult VMIScatterOp::verify() { + auto valueType = cast(getValue().getType()); + auto indicesType = cast(getIndices().getType()); + auto maskType = cast(getMask().getType()); + if (failed(verifyMemoryElementMatches(getOperation(), + getDestination().getType(), valueType, + "destination"))) + return failure(); + if (failed(verifyUBBackedMemory(getOperation(), getDestination().getType(), + "destination"))) + return failure(); + + auto indexElementType = dyn_cast(indicesType.getElementType()); + if (!indexElementType || indexElementType.getWidth() != 32 || + indexElementType.isSigned()) + return emitOpError("requires signless or unsigned 32-bit integer indices"); + + if (failed(verifyAllSameVRegShapeAndLayout(getOperation(), + {valueType, indicesType}, + /*requireSameElement=*/false))) + return failure(); + return verifyMaskMatchesData(getOperation(), maskType, valueType); +} + +void VMIScatterOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Write::get(), &getDestinationMutable()); +} + +LogicalResult VMIShuffleOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + if (sourceType.getElementType() != resultType.getElementType()) + return emitOpError( + "requires result element type to match source element type"); + if (static_cast(getIndices().size()) != resultType.getElementCount()) + return emitOpError( + "requires shuffle index count to match result logical lane count"); + for (int64_t index : getIndices()) { + if (index < 0 || index >= sourceType.getElementCount()) + return emitOpError("requires every shuffle index to select an existing " + "source logical lane"); + } + if (isLayoutAssigned(sourceType) || isLayoutAssigned(resultType)) { + if (!isLayoutAssigned(sourceType) || !isLayoutAssigned(resultType)) + return emitOpError("requires either both source and result to carry " + "layout or neither to carry layout"); + } + return success(); +} + +LogicalResult VMIChannelSplitOp::verify() { + auto sourceType = cast(getSource().getType()); + if (getResults().size() < 2) + return emitOpError("requires at least two channel results"); + auto firstResultType = cast(getResults().front().getType()); + if (sourceType.getElementCount() != + static_cast(getResults().size()) * + firstResultType.getElementCount()) + return emitOpError("requires source lane count to equal result count times " + "per-channel lane count"); + for (Value result : getResults()) { + auto resultType = cast(result.getType()); + if (resultType.getElementCount() != firstResultType.getElementCount() || + resultType.getElementType() != sourceType.getElementType()) + return emitOpError("requires every channel result to have equal lane " + "count and source element type"); + } + bool anyLayout = isLayoutAssigned(sourceType); + for (Value result : getResults()) + anyLayout |= isLayoutAssigned(cast(result.getType())); + if (anyLayout) { + if (!isLayoutAssigned(sourceType)) + return emitOpError("requires layout-assigned channel_split source when " + "any channel result has layout"); + for (Value result : getResults()) { + auto resultType = cast(result.getType()); + if (!isLayoutAssigned(resultType)) + return emitOpError("requires every channel_split result to carry " + "layout when source has layout"); + if (!cast(resultType.getLayout()).isContiguous()) + return emitOpError( + "requires layout-assigned channel_split results to be contiguous"); + } + int64_t channels = getResults().size(); + if (channels == 2 || channels == 4) { + auto sourceLayout = cast(sourceType.getLayout()); + auto expectedLayout = + VMILayoutAttr::getDeinterleaved(getContext(), channels); + if (!sourceLayout.isContiguous() && sourceLayout != expectedLayout) + return emitOpError("requires layout-assigned channel_split source to " + "be contiguous or deinterleaved by result count"); + } + } + return success(); +} + +LogicalResult VMIChannelMergeOp::verify() { + if (getInputs().size() < 2) + return emitOpError("requires at least two channel inputs"); + auto firstInputType = cast(getInputs().front().getType()); + auto resultType = cast(getResult().getType()); + for (Value input : getInputs()) { + auto inputType = cast(input.getType()); + if (inputType.getElementCount() != firstInputType.getElementCount() || + inputType.getElementType() != firstInputType.getElementType()) + return emitOpError("requires all channel inputs to have the same lane " + "count and element type"); + } + if (resultType.getElementCount() != static_cast(getInputs().size()) * + firstInputType.getElementCount() || + resultType.getElementType() != firstInputType.getElementType()) + return emitOpError( + "requires result lane count and element type to match merged channels"); + bool anyLayout = isLayoutAssigned(resultType); + for (Value input : getInputs()) + anyLayout |= isLayoutAssigned(cast(input.getType())); + if (anyLayout) { + if (!isLayoutAssigned(resultType)) + return emitOpError("requires layout-assigned channel_merge result when " + "any channel input has layout"); + for (Value input : getInputs()) { + auto inputType = cast(input.getType()); + if (!isLayoutAssigned(inputType)) + return emitOpError("requires every channel_merge input to carry layout " + "when result has layout"); + if (!cast(inputType.getLayout()).isContiguous()) + return emitOpError( + "requires layout-assigned channel_merge inputs to be contiguous"); + } + int64_t channels = getInputs().size(); + if (channels == 2 || channels == 4) { + auto resultLayout = cast(resultType.getLayout()); + auto expectedLayout = + VMILayoutAttr::getDeinterleaved(getContext(), channels); + if (!resultLayout.isContiguous() && resultLayout != expectedLayout) + return emitOpError("requires layout-assigned channel_merge result to " + "be contiguous or deinterleaved by input count"); + } + } + return success(); +} + +LogicalResult VMIEnsureLayoutOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + if (sourceType.getElementCount() != resultType.getElementCount() || + sourceType.getElementType() != resultType.getElementType()) + return emitOpError("requires source and result to preserve VMI data shape " + "and element type"); + if (!isLayoutAssigned(sourceType) || !isLayoutAssigned(resultType)) + return emitOpError("requires source and result to be layout-assigned"); + return success(); +} + +LogicalResult VMIEnsureMaskLayoutOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + if (sourceType.getElementCount() != resultType.getElementCount() || + sourceType.getGranularity() != resultType.getGranularity()) + return emitOpError("requires source and result to preserve VMI mask shape " + "and granularity"); + if (!isLayoutAssigned(sourceType) || !isLayoutAssigned(resultType)) + return emitOpError("requires source and result to be layout-assigned"); + return success(); +} + +LogicalResult VMIEnsureMaskGranularityOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + if (sourceType.getElementCount() != resultType.getElementCount()) + return emitOpError( + "requires source and result to preserve VMI mask lane count"); + if (sourceType.isPred() || resultType.isPred()) + return emitOpError( + "requires concrete source and result mask granularities"); + if (isLayoutAssigned(sourceType) || isLayoutAssigned(resultType)) { + if (!isLayoutAssigned(sourceType) || !isLayoutAssigned(resultType)) + return emitOpError("requires either both source and result to carry " + "layout or neither to carry layout"); + } + return success(); +} + +LogicalResult VMIUnpackOp::verify() { + return verifyPhysicalParts(getOperation(), getSource().getType(), + getParts().getTypes()); +} + +LogicalResult VMIPackOp::verify() { + return verifyPhysicalParts(getOperation(), getResult().getType(), + getParts().getTypes()); +} + + +enum class CvtDirection { FpWiden, FpNarrow, FpToSi, SiToFp, IntWiden, IntNarrow }; + +// Shared helper: validates mask-data alignment and pmode value. +// Applicable to all VMI elementwise ops that carry mask + pmode. +static LogicalResult verifyVMIPmodeMask(Operation *op, VMIMaskType maskType, + VMIVRegType dataType, + std::optional pmode) { + if (failed(verifyMaskMatchesData(op, maskType, dataType))) + return failure(); + if (pmode.has_value()) { + StringRef mode = pmode.value(); + if (mode != "merge" && mode != "zero") + return op->emitOpError("pmode must be \"merge\" or \"zero\", got \"") + << mode << "\""; + } + return success(); +} +// Variadic-aware variant: skips mask validation when no mask operand is +// provided (unified v-ops allow an absent mask meaning "all-true"). +static LogicalResult verifyVMIVariadicPmodeMask(Operation *op, + ValueRange maskParts, + VMIVRegType dataType, + std::optional pmode) { + if (pmode.has_value()) { + StringRef mode = pmode.value(); + if (mode != "merge" && mode != "zero") + return op->emitOpError("pmode must be \"merge\" or \"zero\", got \"") + << mode << "\""; + } + if (maskParts.empty()) + return success(); + if (maskParts.size() != 1) + return op->emitOpError("expects at most one mask operand"); + return verifyMaskMatchesData(op, cast(maskParts.front().getType()), + dataType); +} + +//===----------------------------------------------------------------------===// +// VMI vector-scalar op verifiers (vadds/vmuls/vmaxs/vmins/vshls/vshrs) +//===----------------------------------------------------------------------===// + +/// Shared verifier for VMI vector-scalar elementwise ops. +static LogicalResult +verifyVMIVectorScalarOp(Operation *op, VMIVRegType srcType, + Type scalarType, VMIVRegType resultType, + VMIMaskType maskType, + std::optional pmode) { + Type eltTy = srcType.getElementType(); + if (!isVMIFloatLikeType(eltTy) && !isVMIIntegerLikeType(eltTy)) + return op->emitOpError( + "requires floating-point-like or integer-like VMI element type"); + + if (scalarType != eltTy) + return op->emitOpError( + "requires scalar type to match vector element type, got scalar ") + << scalarType << " vs vector element " << eltTy; + + if (failed(verifyAllSameVRegShapeAndLayout( + op, {srcType, resultType}, /*requireSameElement=*/true))) + return failure(); + + if (failed(verifyMaskMatchesData(op, maskType, resultType))) + return failure(); + + if (pmode.has_value()) { + StringRef mode = pmode.value(); + if (mode != "merge" && mode != "zero") + return op->emitOpError("unsupported pmode '") + << mode << "'; expected \"merge\" or \"zero\""; + } + + return success(); +} + +/// Shared verifier for VMI vector-scalar integer-only shift ops. +static LogicalResult +verifyVMIVectorScalarShiftOp(Operation *op, VMIVRegType srcType, + Type scalarType, VMIVRegType resultType, + VMIMaskType maskType, + std::optional pmode) { + Type eltTy = srcType.getElementType(); + if (!isVMIIntegerLikeType(eltTy)) + return op->emitOpError( + "requires integer-like VMI element type for shift"); + + return verifyVMIVectorScalarOp(op, srcType, scalarType, resultType, + maskType, pmode); +} + +LogicalResult VMIAddSOp::verify() { + return verifyVMIVectorScalarOp(getOperation(), + cast(getSrc().getType()), getScalar().getType(), + cast(getResult().getType()), + cast(getMask().getType()), getPmode()); +} + +LogicalResult VMIMulSOp::verify() { + return verifyVMIVectorScalarOp(getOperation(), + cast(getSrc().getType()), getScalar().getType(), + cast(getResult().getType()), + cast(getMask().getType()), getPmode()); +} + +LogicalResult VMIMaxSOp::verify() { + return verifyVMIVectorScalarOp(getOperation(), + cast(getSrc().getType()), getScalar().getType(), + cast(getResult().getType()), + cast(getMask().getType()), getPmode()); +} + +LogicalResult VMIMinSOp::verify() { + return verifyVMIVectorScalarOp(getOperation(), + cast(getSrc().getType()), getScalar().getType(), + cast(getResult().getType()), + cast(getMask().getType()), getPmode()); +} + +LogicalResult VMIShlSOp::verify() { + return verifyVMIVectorScalarShiftOp(getOperation(), + cast(getSrc().getType()), getScalar().getType(), + cast(getResult().getType()), + cast(getMask().getType()), getPmode()); +} + +LogicalResult VMIShrSOp::verify() { + return verifyVMIVectorScalarShiftOp(getOperation(), + cast(getSrc().getType()), getScalar().getType(), + cast(getResult().getType()), + cast(getMask().getType()), getPmode()); +} + +//===----------------------------------------------------------------------===// +// Unified (new) VMI op verifiers +//===----------------------------------------------------------------------===// + +//===----------------------------------------------------------------------===// + +/// Returns true if `cmpMode` is a comparison predicate supported by VCMP. +static bool isSupportedVCmpPredicate(StringRef cmpMode) { + return cmpMode == "eq" || cmpMode == "ne" || cmpMode == "lt" || + cmpMode == "le" || cmpMode == "gt" || cmpMode == "ge" || + cmpMode == "oeq" || cmpMode == "one" || cmpMode == "olt" || + cmpMode == "ole" || cmpMode == "ogt" || cmpMode == "oge"; +} + +//===----------------------------------------------------------------------===// + +static const std::set &validDistModes() { + static const std::set modes = {"continuous", "unpack", "dintlv", + "brc"}; + return modes; +} + +static const std::set &validPModes() { + static const std::set modes = {"zero", "merge"}; + return modes; +} + +//===--- Unified (new) VMI op verifiers ---===// + +LogicalResult VMIVbrcOp::verify() { + auto resultType = cast(getResult().getType()); + Type valueType = getValue().getType(); + + if (auto groupAttr = getGroupAttr()) { + // Group broadcast mode + int64_t numGroupsVal = groupAttr.getInt(); + auto vregType = dyn_cast(valueType); + if (!vregType) + return emitOpError("requires VMI vector input when num_groups is set"); + if (vregType.getElementCount() != numGroupsVal) + return emitOpError() + << "requires source logical lane count " << vregType.getElementCount() + << " to match num_groups " << numGroupsVal; + if (vregType.getElementType() != resultType.getElementType()) + return emitOpError("requires source and result element types to match"); + if (auto sourceLayout = vregType.getLayoutAttr()) { + if (!sourceLayout.isGroupSlots() || + sourceLayout.getNumGroups() != numGroupsVal) + return emitOpError() << "requires layout-assigned source to use " + "#pto.vmi.layout"; + } + if (auto resultLayout = resultType.getLayoutAttr()) { + if (resultLayout.isGroupSlots()) + return emitOpError( + "requires layout-assigned result to use a dense VMI layout"); + } + return verifyNumGroups(getOperation(), resultType, numGroupsVal); + } + + // Scalar/1-lane broadcast mode (no num_groups) + if (valueType == resultType.getElementType()) + return success(); + if (auto vregType = dyn_cast(valueType)) { + if (vregType.getElementCount() != 1) + return emitOpError("requires VMI vector input to have one logical lane"); + if (vregType.getElementType() != resultType.getElementType()) + return emitOpError("requires VMI vector input element type to match " + "result element type"); + return success(); + } + return emitOpError("requires scalar or VMI vector input element type to " + "match result element type"); +} + +LogicalResult VMIVciOp::verify() { + auto resultType = cast(getResult().getType()); + Type elementType = resultType.getElementType(); + if (!isVMIIotaElementType(elementType)) + return emitOpError("requires result element type to be integer 8/16/32 " + "or f16/f32"); + if (!isCompatibleScalarForSemanticType(elementType, getBase().getType())) + return emitOpError("requires base type to match result element type"); + + if (std::optional order = getOrder()) { + if (*order != "ASC" && *order != "DESC") + return emitOpError("requires order to be ASC or DESC"); + } + return success(); +} + +LogicalResult VMIPsetOp::verify() { + auto resultType = cast(getResult().getType()); + StringRef pattern = getPattern(); + if (pattern != "PAT_ALL") + return emitOpError("requires pattern to be \"PAT_ALL\""); + if (!resultType.isPred() && !isLayoutAssigned(resultType)) + return emitOpError("requires concrete mask result to carry layout"); + return success(); +} + +LogicalResult VMIPgeOp::verify() { + auto resultType = cast(getResult().getType()); + StringRef pattern = getPattern(); + if (!pattern.starts_with("PAT_VL")) + return emitOpError("requires pattern to start with \"PAT_VL\""); + int64_t activeLanes; + if (pattern.drop_front(6).getAsInteger(10, activeLanes)) + return emitOpError("requires pattern \"PAT_VL\" with integer n"); + if (activeLanes <= 0) + return emitOpError("requires positive n in pattern \"PAT_VL\""); + if (activeLanes > resultType.getElementCount()) + return emitOpError("PAT_VL active lanes ") << activeLanes + << " exceeds mask element count " << resultType.getElementCount(); + if (!resultType.isPred() && !isLayoutAssigned(resultType)) + return emitOpError("requires concrete mask result to carry layout"); + return success(); +} + +LogicalResult VMIPltOp::verify() { + auto resultType = cast(getMask().getType()); + auto scalarType = dyn_cast(getScalar().getType()); + if (!scalarType || scalarType.getWidth() != 32) + return emitOpError("requires i32 scalar input"); + auto scalarOutType = dyn_cast(getScalarOut().getType()); + if (!scalarOutType || scalarOutType.getWidth() != 32) + return emitOpError("requires i32 scalar_out result"); + if (!resultType.isPred() && !isLayoutAssigned(resultType)) + return emitOpError("requires concrete mask result to carry layout"); + return success(); +} + +LogicalResult VMIVaddOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) + return failure(); + if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), + resultType, getPmode()))) + return failure(); + return success(); +} + +LogicalResult VMIVsubOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) + return failure(); + if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), + resultType, getPmode()))) + return failure(); + return success(); +} + +LogicalResult VMIVmulOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) + return failure(); + if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), + resultType, getPmode()))) + return failure(); + return success(); +} + +LogicalResult VMIVdivOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIFloatLikeType(lhsType.getElementType())) + return emitOpError("requires floating-point-like VMI element type"); + if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) + return failure(); + if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), + resultType, getPmode()))) + return failure(); + return success(); +} + +LogicalResult VMIVminOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIFloatLikeType(lhsType.getElementType())) + return emitOpError("requires floating-point-like VMI element type"); + if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) + return failure(); + if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), + resultType, getPmode()))) + return failure(); + return success(); +} + +LogicalResult VMIVmaxOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIFloatLikeType(lhsType.getElementType())) + return emitOpError("requires floating-point-like VMI element type"); + if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) + return failure(); + if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), + resultType, getPmode()))) + return failure(); + return success(); +} + +LogicalResult VMIVnegOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + if (failed(verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType))) + return failure(); + if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), + resultType, getPmode()))) + return failure(); + return success(); +} + +LogicalResult VMIVabsOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + + Type eltTy = sourceType.getElementType(); + if (!isVMIFloatLikeType(eltTy) && !isVMIIntegerLikeType(eltTy)) + return emitOpError( + "requires floating-point-like or integer-like VMI element type"); + + if (failed(verifyAllSameVRegShapeAndLayout( + getOperation(), {sourceType, resultType}, + /*requireSameElement=*/true))) + return failure(); + + if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), + resultType, getPmode()))) + return failure(); + return success(); +} + +LogicalResult VMIVsqrtOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + if (failed(verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType))) + return failure(); + if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), + resultType, getPmode()))) + return failure(); + return success(); +} + +LogicalResult VMIVexpOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + if (failed(verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType))) + return failure(); + if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), + resultType, getPmode()))) + return failure(); + return success(); +} + +LogicalResult VMIVlnOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + if (failed(verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType))) + return failure(); + if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), + resultType, getPmode()))) + return failure(); + return success(); +} + +LogicalResult VMIVreluOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + if (failed(verifyFloatUnaryVRegOp(getOperation(), sourceType, resultType))) + return failure(); + if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), + resultType, getPmode()))) + return failure(); + return success(); +} + +LogicalResult VMIVandOp::verify() { + if (isa(getLhs().getType())) { + // Mask logic path: reject predication mask and pmode. + if (!getMask().empty()) + return emitOpError("mask logic op does not support predication mask"); + if (auto pmode = getPmode()) + return emitOpError("mask logic op does not support pmode"); + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + return verifyAllSameMaskShapeLayoutAndGranularity( + getOperation(), {lhsType, rhsType, resultType}); + } + // VReg path. + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIIntegerLikeType(lhsType.getElementType())) + return emitOpError("requires integer-like VMI element type"); + if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) + return failure(); + if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), + resultType, getPmode()))) + return failure(); + return success(); +} + +LogicalResult VMIVorOp::verify() { + if (isa(getLhs().getType())) { + // Mask logic path: reject predication mask and pmode. + if (!getMask().empty()) + return emitOpError("mask logic op does not support predication mask"); + if (auto pmode = getPmode()) + return emitOpError("mask logic op does not support pmode"); + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + return verifyAllSameMaskShapeLayoutAndGranularity( + getOperation(), {lhsType, rhsType, resultType}); + } + // VReg path. + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIIntegerLikeType(lhsType.getElementType())) + return emitOpError("requires integer-like VMI element type"); + if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) + return failure(); + if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), + resultType, getPmode()))) + return failure(); + return success(); +} + +LogicalResult VMIVxorOp::verify() { + if (isa(getLhs().getType())) { + // Mask logic path: reject predication mask and pmode. + if (!getMask().empty()) + return emitOpError("mask logic op does not support predication mask"); + if (auto pmode = getPmode()) + return emitOpError("mask logic op does not support pmode"); + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + return verifyAllSameMaskShapeLayoutAndGranularity( + getOperation(), {lhsType, rhsType, resultType}); + } + // VReg path. + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIIntegerLikeType(lhsType.getElementType())) + return emitOpError("requires integer-like VMI element type"); + if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) + return failure(); + if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), + resultType, getPmode()))) + return failure(); + return success(); +} + +LogicalResult VMIVshlOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIIntegerLikeType(lhsType.getElementType())) + return emitOpError("requires integer-like VMI element type"); + if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) + return failure(); + if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), + resultType, getPmode()))) + return failure(); + return success(); +} + +LogicalResult VMIVshrOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + auto integerType = dyn_cast(lhsType.getElementType()); + if (!integerType) + return emitOpError("requires integer VMI element type"); + if (failed(verifyElementwiseVRegOp(getOperation(), lhsType, rhsType, resultType))) + return failure(); + if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), + resultType, getPmode()))) + return failure(); + return success(); +} + +LogicalResult VMIVnotOp::verify() { + if (isa(getSource().getType())) { + // Mask logic path: reject predication mask and pmode. + if (!getMask().empty()) + return emitOpError("mask logic op does not support predication mask"); + if (auto pmode = getPmode()) + return emitOpError("mask logic op does not support pmode"); + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + return verifyAllSameMaskShapeLayoutAndGranularity( + getOperation(), {sourceType, resultType}); + } + // VReg path. + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + if (!isVMIIntegerLikeType(sourceType.getElementType())) + return emitOpError("requires integer-like VMI element type"); + if (failed(verifyAllSameVRegShapeAndLayout(getOperation(), + {sourceType, resultType}, + /*requireSameElement=*/true))) + return failure(); + if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), + resultType, getPmode()))) + return failure(); + return success(); +} + +LogicalResult VMIvSelOp::verify() { + auto maskType = cast(getMask().getType()); + auto trueType = cast(getTrueValue().getType()); + auto falseType = cast(getFalseValue().getType()); + auto resultType = cast(getResult().getType()); + if (failed(verifyAllSameVRegShapeAndLayout(getOperation(), + {trueType, falseType, resultType}, + /*requireSameElement=*/true))) + return failure(); + if (failed(verifyMaskMatchesData(getOperation(), maskType, resultType))) + return failure(); + if (auto pmode = getPmode(); pmode.has_value()) { + StringRef mode = pmode.value(); + if (mode != "merge" && mode != "zero") + return emitOpError("pmode must be \"merge\" or \"zero\", got \"") + << mode << "\""; + } + return success(); +} + +LogicalResult VMIvcaddOp::verify() { + auto sourceType = cast(getSource().getType()); + auto maskType = cast(getMask().getType()); + auto resultType = cast(getResult().getType()); + auto elemTy = sourceType.getElementType(); + + // Element type must be integer-like or float-like + bool isFloat = isVMIFloatLikeType(elemTy); + bool isInt = isVMIIntegerLikeType(elemTy); + if (!isFloat && !isInt) + return emitOpError("requires integer-like or floating-point-like VMI " + "source element type"); + + // Floating-point vcadd MUST carry reassoc + if (isFloat && !getReassoc()) + return emitOpError("floating add-reduction requires reassoc attr"); + + if (failed(verifyMaskMatchesData(getOperation(), maskType, sourceType))) + return failure(); + + // Validate group vs result lane count + if (auto groupAttr = getGroupAttr()) { + int64_t C = groupAttr.getInt(); + if (sourceType.getElementCount() % C != 0) + return emitOpError("group count ") << C << " must divide source lane count " + << sourceType.getElementCount(); + if (resultType.getElementCount() != C) + return emitOpError("result lane count must equal group count ") + << C << ", got " << resultType.getElementCount(); + if (auto resultLayout = resultType.getLayoutAttr()) { + if (!resultLayout.isGroupSlots() || + resultLayout.getNumGroups() != C) + return emitOpError() + << "layout-assigned result must use " + "#pto.vmi.layout"; + } + } else { + if (resultType.getElementCount() != 1) + return emitOpError("full reduction (no group) requires 1-lane result, got ") + << resultType.getElementCount(); + } + + // Element types must match + if (sourceType.getElementType() != resultType.getElementType()) + return emitOpError("source and result element types must match"); + + // pmode must be "zero" or "merge" if set + if (auto pmode = getPmode()) { + StringRef val = *pmode; + if (val != "zero" && val != "merge") + return emitOpError("pmode must be \"zero\" or \"merge\", got \"") << val << "\""; + } + + return success(); +} + +LogicalResult VMIvcmaxOp::verify() { + auto sourceType = cast(getSource().getType()); + auto maskType = cast(getMask().getType()); + auto resultType = cast(getResult().getType()); + auto elemTy = sourceType.getElementType(); + + bool isFloat = isVMIFloatLikeType(elemTy); + bool isInt = isVMIIntegerLikeType(elemTy); + if (!isFloat && !isInt) + return emitOpError("requires integer-like or floating-point-like VMI " + "source element type"); + + if (failed(verifyMaskMatchesData(getOperation(), maskType, sourceType))) + return failure(); + + if (auto groupAttr = getGroupAttr()) { + int64_t C = groupAttr.getInt(); + if (sourceType.getElementCount() % C != 0) + return emitOpError("group count ") << C << " must divide source lane count " + << sourceType.getElementCount(); + if (resultType.getElementCount() != C) + return emitOpError("result lane count must equal group count ") + << C << ", got " << resultType.getElementCount(); + if (auto resultLayout = resultType.getLayoutAttr()) { + if (!resultLayout.isGroupSlots() || + resultLayout.getNumGroups() != C) + return emitOpError() + << "layout-assigned result must use " + "#pto.vmi.layout"; + } + } else { + if (resultType.getElementCount() != 1) + return emitOpError("full reduction (no group) requires 1-lane result, got ") + << resultType.getElementCount(); + } + + if (sourceType.getElementType() != resultType.getElementType()) + return emitOpError("source and result element types must match"); + + if (auto pmode = getPmode()) { + StringRef val = *pmode; + if (val != "zero" && val != "merge") + return emitOpError("pmode must be \"zero\" or \"merge\", got \"") << val << "\""; + } + + return success(); +} + +LogicalResult VMIvcminOp::verify() { + auto sourceType = cast(getSource().getType()); + auto maskType = cast(getMask().getType()); + auto resultType = cast(getResult().getType()); + auto elemTy = sourceType.getElementType(); + + bool isFloat = isVMIFloatLikeType(elemTy); + bool isInt = isVMIIntegerLikeType(elemTy); + if (!isFloat && !isInt) + return emitOpError("requires integer-like or floating-point-like VMI " + "source element type"); + + if (failed(verifyMaskMatchesData(getOperation(), maskType, sourceType))) + return failure(); + + if (auto groupAttr = getGroupAttr()) { + int64_t C = groupAttr.getInt(); + if (sourceType.getElementCount() % C != 0) + return emitOpError("group count ") << C << " must divide source lane count " + << sourceType.getElementCount(); + if (resultType.getElementCount() != C) + return emitOpError("result lane count must equal group count ") + << C << ", got " << resultType.getElementCount(); + if (auto resultLayout = resultType.getLayoutAttr()) { + if (!resultLayout.isGroupSlots() || + resultLayout.getNumGroups() != C) + return emitOpError() + << "layout-assigned result must use " + "#pto.vmi.layout"; + } + } else { + if (resultType.getElementCount() != 1) + return emitOpError("full reduction (no group) requires 1-lane result, got ") + << resultType.getElementCount(); + } + + if (sourceType.getElementType() != resultType.getElementType()) + return emitOpError("source and result element types must match"); + + if (auto pmode = getPmode()) { + StringRef val = *pmode; + if (val != "zero" && val != "merge") + return emitOpError("pmode must be \"zero\" or \"merge\", got \"") << val << "\""; + } + + return success(); +} + +LogicalResult VMIVgatherOp::verify() { + auto offsetsType = cast(getOffsets().getType()); + auto maskType = cast(getMask().getType()); + auto resultType = cast(getResult().getType()); + + if (failed(verifyMemoryElementMatches(getOperation(), getSource().getType(), + resultType, "source"))) + return failure(); + + auto indexElementType = + dyn_cast(offsetsType.getElementType()); + if (!indexElementType || indexElementType.isSigned() || + (indexElementType.getWidth() != 32 && indexElementType.getWidth() != 16)) + return emitOpError( + "requires signless or unsigned 16-bit or 32-bit integer offsets"); + + if (failed(verifyAllSameVRegShapeAndLayout( + getOperation(), {offsetsType, resultType}, + /*requireSameElement=*/false))) + return failure(); + if (failed(verifyMaskMatchesData(getOperation(), maskType, resultType))) + return failure(); + + // 16-bit offsets only address the ui16 gather path (pto.vgather2 / b16 mask), + // which requires a ui16 result element type. Reject other 16-bit-offset + // results here so the error surfaces at the vgather op rather than later in + // the legacy gather it lowers to. + auto resultIntegerType = dyn_cast(resultType.getElementType()); + if (indexElementType.getWidth() == 16 && + (!resultIntegerType || !resultIntegerType.isUnsigned() || + resultIntegerType.getWidth() != 16)) + return emitOpError( + "requires ui16 result element type when using ui16 offsets"); + + if (auto pmode = getPmode()) { + if (pmode.value() != "merge" && pmode.value() != "zero") + return emitOpError("pmode must be 'merge' or 'zero'"); + } + return success(); +} + +void VMIVgatherOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Read::get(), &getSourceMutable()); +} + +LogicalResult VMIVgatherbOp::verify() { + auto offsetsType = cast(getOffsets().getType()); + auto maskType = cast(getMask().getType()); + auto resultType = cast(getResult().getType()); + + if (failed(verifyMemoryElementMatches(getOperation(), getSource().getType(), + resultType, "source"))) + return failure(); + + auto indexElementType = + dyn_cast(offsetsType.getElementType()); + if (!indexElementType || indexElementType.isSigned() || + (indexElementType.getWidth() != 32 && indexElementType.getWidth() != 16)) + return emitOpError( + "requires signless or unsigned 16-bit or 32-bit integer offsets"); + + if (failed(verifyMaskMatchesData(getOperation(), maskType, resultType))) + return failure(); + + if (auto pmode = getPmode()) { + if (pmode.value() != "merge" && pmode.value() != "zero") + return emitOpError("pmode must be 'merge' or 'zero'"); + } + return success(); +} + +void VMIVgatherbOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Read::get(), &getSourceMutable()); +} + +LogicalResult VMIVscatterOp::verify() { + auto valueType = cast(getValue().getType()); + auto offsetsType = cast(getOffsets().getType()); + auto maskType = cast(getMask().getType()); + + if (failed(verifyMemoryElementMatches(getOperation(), + getDestination().getType(), valueType, + "destination"))) + return failure(); + + auto indexElementType = + dyn_cast(offsetsType.getElementType()); + if (!indexElementType || indexElementType.getWidth() != 32 || + indexElementType.isSigned()) + return emitOpError("requires signless or unsigned 32-bit integer offsets"); + + if (failed(verifyAllSameVRegShapeAndLayout(getOperation(), + {valueType, offsetsType}, + /*requireSameElement=*/false))) + return failure(); + if (failed(verifyMaskMatchesData(getOperation(), maskType, valueType))) + return failure(); + + if (auto pmode = getPmode()) { + if (pmode.value() != "merge" && pmode.value() != "zero") + return emitOpError("pmode must be 'merge' or 'zero'"); + } + return success(); +} + +void VMIVscatterOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Write::get(), &getDestinationMutable()); +} + +LogicalResult VMIVexpdifOp::verify() { + auto xType = cast(getX().getType()); + auto maxType = cast(getMax().getType()); + auto maskType = cast(getMask().getType()); + auto resultType = cast(getResult().getType()); + + if (!isVMIFloatLikeType(xType.getElementType())) + return emitOpError("requires x element type to be f16 or f32"); + + auto maxElemType = dyn_cast(maxType.getElementType()); + if (!maxElemType || maxElemType.getWidth() != 32) + return emitOpError("requires max element type to be f32"); + + auto resultElemType = dyn_cast(resultType.getElementType()); + if (!resultElemType || resultElemType.getWidth() != 32) + return emitOpError("requires result element type to be f32"); + + if (xType.getElementCount() != maxType.getElementCount() || + xType.getElementCount() != resultType.getElementCount()) + return emitOpError( + "requires x, max, and result logical lane counts to match"); + + if (failed(verifyMaskMatchesData(getOperation(), maskType, resultType))) + return failure(); + + if (auto pmode = getPmode()) { + if (pmode.value() != "merge" && pmode.value() != "zero") + return emitOpError("pmode must be 'merge' or 'zero'"); + } + return success(); +} + +LogicalResult VMIVaxpyOp::verify() { + auto xType = cast(getX().getType()); + auto accType = cast(getAcc().getType()); + auto maskType = cast(getMask().getType()); + auto resultType = cast(getResult().getType()); + + if (!isVMIFloatLikeType(xType.getElementType())) + return emitOpError("requires vector element type to be f16 or f32"); + + if (xType != accType || accType != resultType) + return emitOpError( + "requires x, acc, and result to have identical VMI vreg types"); + + auto alphaType = cast(getAlpha().getType()); + if (alphaType != xType.getElementType()) + return emitOpError("requires alpha scalar type to match vector element type"); + + if (failed(verifyMaskMatchesData(getOperation(), maskType, resultType))) + return failure(); + + if (auto pmode = getPmode()) { + if (pmode.value() != "merge" && pmode.value() != "zero") + return emitOpError("pmode must be 'merge' or 'zero'"); + } + return success(); +} + +LogicalResult VMIVlreluOp::verify() { + auto xType = cast(getX().getType()); + auto maskType = cast(getMask().getType()); + auto resultType = cast(getResult().getType()); + + if (!isVMIFloatLikeType(xType.getElementType())) + return emitOpError("requires vector element type to be f16 or f32"); + + if (xType != resultType) + return emitOpError("requires x and result to have identical VMI vreg types"); + + auto slopeType = cast(getSlope().getType()); + if (slopeType != xType.getElementType()) + return emitOpError( + "requires slope scalar type to match vector element type"); + + if (failed(verifyMaskMatchesData(getOperation(), maskType, resultType))) + return failure(); + + if (auto pmode = getPmode()) { + if (pmode.value() != "merge" && pmode.value() != "zero") + return emitOpError("pmode must be 'merge' or 'zero'"); + } + return success(); +} + +LogicalResult VMIVpreluOp::verify() { + auto xType = cast(getX().getType()); + auto alphaType = cast(getAlpha().getType()); + auto maskType = cast(getMask().getType()); + auto resultType = cast(getResult().getType()); + + if (!isVMIFloatLikeType(xType.getElementType())) + return emitOpError("requires vector element type to be f16 or f32"); + + if (xType != alphaType || alphaType != resultType) + return emitOpError( + "requires x, alpha, and result to have identical VMI vreg types"); + + if (failed(verifyMaskMatchesData(getOperation(), maskType, resultType))) + return failure(); + + if (auto pmode = getPmode()) { + if (pmode.value() != "merge" && pmode.value() != "zero") + return emitOpError("pmode must be 'merge' or 'zero'"); + } + return success(); +} + +LogicalResult VMIVmullOp::verify() { + auto aType = cast(getA().getType()); + auto bType = cast(getB().getType()); + auto maskType = cast(getMask().getType()); + auto lowType = cast(getLow().getType()); + auto highType = cast(getHigh().getType()); + + auto isLegalElementType = [](Type type) { + auto integerType = dyn_cast(type); + return integerType && integerType.getWidth() == 32 && + (integerType.isSignless() || integerType.isUnsigned()); + }; + if (!isLegalElementType(aType.getElementType()) || + !isLegalElementType(bType.getElementType()) || + !isLegalElementType(lowType.getElementType()) || + !isLegalElementType(highType.getElementType())) + return emitOpError( + "requires a, b, low, and high element types to be exactly i32 or ui32"); + + if (aType != bType || aType != lowType || aType != highType) + return emitOpError( + "requires a, b, low, and high to have identical VMI vreg types"); + + int64_t lanes = aType.getElementCount(); + if (lanes != 64 && lanes != 128 && lanes != 256) + return emitOpError("requires logical lane count to be 64, 128, or 256"); + + if (failed(verifyMaskMatchesData(getOperation(), maskType, aType))) + return failure(); + + if (auto pmode = getPmode(); pmode && pmode.value() != "zero") + return emitOpError("pmode must be 'zero' when specified"); + return success(); +} + +LogicalResult VMIVmulaOp::verify() { + auto accType = cast(getAcc().getType()); + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto resultType = cast(getResult().getType()); + + Type eltTy = accType.getElementType(); + if (!isVMIFloatLikeType(eltTy) && !isVMIIntegerLikeType(eltTy)) + return emitOpError( + "requires floating-point-like or integer-like VMI element type"); + + if (accType != lhsType || lhsType != rhsType || rhsType != resultType) + return emitOpError( + "requires acc, lhs, rhs, and result to have identical VMI vreg types"); + + if (failed(verifyVMIVariadicPmodeMask(getOperation(), getMask(), + resultType, getPmode()))) + return failure(); + return success(); +} + +LogicalResult VMICvtOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + + // 1. Lane count must match. + if (sourceType.getElementCount() != resultType.getElementCount()) + return emitOpError( + "requires source and result logical lane counts to match"); + + Type srcElem = sourceType.getElementType(); + Type dstElem = resultType.getElementType(); + unsigned srcBits = getVMIElementBitWidth(srcElem); + unsigned dstBits = getVMIElementBitWidth(dstElem); + bool srcFp = isVMIFloatLikeType(srcElem); + bool dstFp = isVMIFloatLikeType(dstElem); + bool srcInt = isVMIIntegerLikeType(srcElem); + bool dstInt = isVMIIntegerLikeType(dstElem); + bool srcSignlessInt = + srcInt && isa(srcElem) && + !cast(srcElem).isUnsigned() && + !cast(srcElem).isSigned(); + + // 2. Classify the conversion direction. + CvtDirection dir; + if (srcFp && dstFp) { + if (dstBits > srcBits) + dir = CvtDirection::FpWiden; + else if (dstBits < srcBits) + dir = CvtDirection::FpNarrow; + else + return emitOpError( + "fp-to-fp conversion must change element bit-width"); + } else if (srcFp && dstInt) { + if (!isVMISignedOrSignlessIntegerType(dstElem)) + return emitOpError( + "fp-to-int conversion requires signed or signless integer result " + "element type"); + dir = CvtDirection::FpToSi; + } else if (srcInt && dstFp) { + if (!isVMISignedOrSignlessIntegerType(srcElem)) + return emitOpError( + "int-to-fp conversion requires signed or signless integer source " + "element type"); + dir = CvtDirection::SiToFp; + } else if (srcInt && dstInt) { + if (dstBits > srcBits) + dir = CvtDirection::IntWiden; + else if (dstBits < srcBits) + dir = CvtDirection::IntNarrow; + else + return emitOpError( + "int-to-int conversion must change element bit-width"); + } else { + return emitOpError( + "unsupported element type combination for vcvt"); + } + + // 3. Validate attributes against the conversion direction. + + // --- rounding --- + if (auto roundingAttr = (*this)->getAttrOfType("rounding")) { + if (dir != CvtDirection::FpNarrow) + return emitOpError("'rounding' attribute is only valid for " + "fp-narrowing conversions"); + StringRef rnd = roundingAttr.getValue(); + if (rnd != "R" && rnd != "A" && rnd != "H" && rnd != "Z") + return emitOpError("rounding must be 'R' (nearest-even), " + "'A' (away-from-zero), 'H' (half-up), " + "or 'Z' (toward-zero)"); + } + + // --- saturate --- + if (auto satAttr = (*this)->getAttrOfType("saturate")) { + if (dir != CvtDirection::FpNarrow && dir != CvtDirection::IntNarrow) + return emitOpError("'saturate' attribute is only valid for " + "narrowing conversions (fp or int)"); + if (satAttr.getValue() != "SAT") + return emitOpError("saturate must be 'SAT'"); + } + + // --- sign --- + // Int-widening requires a signed/unsigned source element type to + // determine sign-extension vs zero-extension. Signless integers are + // rejected. + if (dir == CvtDirection::IntWiden && srcSignlessInt) + return emitOpError("int-widening conversions require a signed or " + "unsigned integer source element type " + "(e.g. si8/ui8/si16/ui16); " + "signless integer is not allowed"); + + // --- pmode --- + if (auto pmodeAttr = (*this)->getAttrOfType("pmode")) { + StringRef pmode = pmodeAttr.getValue(); + if (pmode != "merge" && pmode != "zero") + return emitOpError("pmode must be 'merge' or 'zero'"); + } + + return success(); +} + +LogicalResult VMIVinterpretCastOp::verify() { + auto sourceType = cast(getSource().getType()); + auto resultType = cast(getResult().getType()); + std::optional sourceBits = + getVMIIntegerOrFloatBitWidth(sourceType.getElementType()); + std::optional resultBits = + getVMIIntegerOrFloatBitWidth(resultType.getElementType()); + if (!sourceBits || !resultBits) + return emitOpError( + "requires integer or floating-point source and result element types"); + if (sourceType.getElementCount() * static_cast(*sourceBits) != + resultType.getElementCount() * static_cast(*resultBits)) + return emitOpError( + "requires source and result to carry the same total number of bits"); + + if (isLayoutAssigned(sourceType) || isLayoutAssigned(resultType)) { + if (!isLayoutAssigned(sourceType) || !isLayoutAssigned(resultType)) + return emitOpError( + "requires either both source and result to carry layout or neither " + "to carry layout"); + if (sourceType.getLayout() != resultType.getLayout()) + return emitOpError("requires source and result layouts to match"); + } + + return success(); +} + +ParseResult VMIvStoreOp::parse(OpAsmParser &parser, OperationState &result) { + SmallVector preBracketOperands; + OpAsmParser::UnresolvedOperand operand; + OpAsmParser::UnresolvedOperand offsetOperand; + SmallVector postBracketOps; + + if (parser.parseOperand(operand)) + return failure(); + preBracketOperands.push_back(operand); + + bool consumedLSquare = false; + while (!consumedLSquare) { + if (succeeded(parser.parseOptionalLSquare())) { + if (parser.parseOperand(offsetOperand) || parser.parseRSquare()) + return failure(); + consumedLSquare = true; + break; + } + if (parser.parseComma()) + return failure(); + if (succeeded(parser.parseOptionalLSquare())) { + if (parser.parseOperand(offsetOperand) || parser.parseRSquare()) + return failure(); + consumedLSquare = true; + break; + } + if (parser.parseOperand(operand)) + return failure(); + preBracketOperands.push_back(operand); + } + + if (preBracketOperands.empty()) + return parser.emitError(parser.getCurrentLocation(), + "expected at least one value and one destination"); + + // Optional post-bracket operands: stride, block_stride/repeat_stride, mask. + // Up to 3, disambiguated after parsing attrs. + while (succeeded(parser.parseOptionalComma())) { + OpAsmParser::UnresolvedOperand postOp; + if (parser.parseOperand(postOp)) + return failure(); + postBracketOps.push_back(postOp); + if (postBracketOps.size() >= 3) + break; + } + + if (parser.parseOptionalAttrDict(result.attributes)) + return failure(); + + SmallVector types; + if (parser.parseColon() || parser.parseTypeList(types)) + return failure(); + + bool hasGroup = result.attributes.get("group") != nullptr; + bool hasStride = false; + bool hasBlock = false; + bool hasRepeat = false; + bool hasMask = false; + int strideIdx = -1; + int blockIdx = -1; + int repeatIdx = -1; + int maskIdx = -1; + + if (hasGroup) { + // Group mode: post-bracket ops are stride[, mask] + if (postBracketOps.size() >= 1) { + hasStride = true; + strideIdx = 0; + } + if (postBracketOps.size() >= 2) { + hasMask = true; + maskIdx = 1; + } + } else if (postBracketOps.size() >= 2) { + // Block-stride mode: post-bracket ops are block_stride, repeat_stride[, mask] + hasBlock = true; + hasRepeat = true; + blockIdx = 0; + repeatIdx = 1; + if (postBracketOps.size() >= 3) { + hasMask = true; + maskIdx = 2; + } + } else if (postBracketOps.size() == 1) { + // Single post-bracket operand without group: mask + hasMask = true; + maskIdx = 0; + } + + size_t nValues = preBracketOperands.size() - 1; + size_t nTypes = types.size(); + size_t expectedTypes = nValues + 1 + (hasMask ? 1 : 0); + + if (nTypes != expectedTypes) + return parser.emitError(parser.getCurrentLocation()) + << "expected " << expectedTypes << " types (" << nValues + << " value(s), 1 destination" << (hasMask ? ", 1 mask" : "") + << "), got " << nTypes; + + for (size_t i = 0; i < nValues; ++i) { + if (parser.resolveOperand(preBracketOperands[i], types[i], result.operands)) + return failure(); + } + + Type destType = types[nValues]; + if (parser.resolveOperand(preBracketOperands[nValues], destType, + result.operands)) + return failure(); + + if (parser.resolveOperand(offsetOperand, parser.getBuilder().getIndexType(), + result.operands)) + return failure(); + + if (hasStride && + parser.resolveOperand(postBracketOps[strideIdx], + parser.getBuilder().getIndexType(), + result.operands)) + return failure(); + + if (hasBlock && + parser.resolveOperand(postBracketOps[blockIdx], + parser.getBuilder().getIntegerType(16), + result.operands)) + return failure(); + if (hasRepeat && + parser.resolveOperand(postBracketOps[repeatIdx], + parser.getBuilder().getIntegerType(16), + result.operands)) + return failure(); + + if (hasMask) { + Type maskType = types.back(); + if (parser.resolveOperand(postBracketOps[maskIdx], maskType, + result.operands)) + return failure(); + } + + result.addAttribute("operandSegmentSizes", + parser.getBuilder().getDenseI32ArrayAttr( + {static_cast(nValues), 1, 1, + hasStride ? 1 : 0, hasBlock ? 1 : 0, + hasRepeat ? 1 : 0, hasMask ? 1 : 0})); + return success(); +} + +void VMIvStoreOp::print(OpAsmPrinter &p) { + for (auto val : getValues()) + p << ' ' << val << ", "; + p << getDestination() << '['; + p.printOperand(getOffset()); + p << ']'; + if (getStride()) { + p << ", "; + p.printOperand(getStride()); + } + if (getBlockStride()) { + p << ", "; + p.printOperand(getBlockStride()); + p << ", "; + p.printOperand(getRepeatStride()); + } + if (!getMask().empty()) { + p << ", "; + p.printOperand(getMask()[0]); + } + p.printOptionalAttrDict((*this)->getAttrs(), {"operandSegmentSizes"}); + p << " : "; + for (auto val : getValues()) + p << val.getType() << ", "; + p << getDestination().getType(); + if (!getMask().empty()) + p << ", " << getMask()[0].getType(); +} + +LogicalResult VMIvStoreOp::verify() { + // group and dist_mode are mutually exclusive + if (getGroup() && getDistMode()) + return emitOpError("group and dist_mode are mutually exclusive"); + if (getGroup() && !getStride()) + return emitOpError("group requires a stride operand"); + if (!getGroup() && getStride()) + return emitOpError("stride operand is only valid with group"); + if (getGroup() && !getMask().empty()) + return emitOpError("group mode does not support mask operand"); + + if (getGroup()) { + int64_t numGroups = getGroupAttr().getInt(); + if (numGroups <= 0) + return emitOpError("group must be positive, got ") << numGroups; + if (getValues().size() != 1) + return emitOpError("group mode requires exactly 1 value"); + return success(); + } + + // block_stride / repeat_stride: paired, mutually exclusive with + // dist_mode and group + bool hasBlock = static_cast(getBlockStride()); + bool hasRepeat = static_cast(getRepeatStride()); + if (hasBlock != hasRepeat) + return emitOpError( + "block_stride and repeat_stride must both be present or absent"); + if (hasBlock) { + if (getDistMode()) + return emitOpError( + "block_stride and dist_mode are mutually exclusive"); + if (getValues().size() != 1) + return emitOpError("block-stride mode requires exactly 1 value"); + return success(); + } + + auto distMode = getDistMode(); + bool isDintlv = distMode && *distMode == "dintlv"; + size_t nValues = getValues().size(); + if (nValues < 1) + return emitOpError("requires at least 1 value"); + if (isDintlv && nValues != 2) + return emitOpError("dist-mode \"dintlv\" requires exactly 2 values"); + if (!isDintlv && nValues != 1) + return emitOpError("requires exactly 1 value for dist-mode \"") + << (distMode ? *distMode : "continuous") << "\""; + + bool hasMask = !getMask().empty(); + if (getMask().size() > 1) + return emitOpError("at most one mask allowed"); + + if (distMode && !validDistModes().count(*distMode)) + return emitOpError("invalid dist-mode: \"") << *distMode << "\""; + if (distMode && (*distMode == "unpack" || *distMode == "brc")) + return emitOpError("dist-mode \"") + << *distMode << "\" is not valid for vstore"; + + auto pmode = getPmode(); + if (pmode && !validPModes().count(*pmode)) + return emitOpError("invalid pmode: \"") << *pmode << "\""; + + auto valueType = cast(getValues()[0].getType()); + if (failed(verifyMemoryElementMatches(getOperation(), + getDestination().getType(), valueType, + "destination"))) + return failure(); + + if (nValues == 2) { + auto loType = cast(getValues()[0].getType()); + auto hiType = cast(getValues()[1].getType()); + if (failed(verifyAllSameVRegShapeAndLayout(getOperation(), + {loType, hiType}, + /*requireSameElement=*/true))) + return failure(); + if (failed(verifyContiguousIfLayoutAssigned(getOperation(), loType, + "low input")) || + failed(verifyContiguousIfLayoutAssigned(getOperation(), hiType, + "high input"))) + return failure(); + } + + if (hasMask) { + auto maskType = cast(getMask()[0].getType()); + if (failed(verifyMaskMatchesData(getOperation(), maskType, valueType))) + return failure(); + } + + return success(); +} + +void VMIvStoreOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Write::get(), &getDestinationMutable()); +} + +LogicalResult VMIVselrOp::verify() { + auto sourceType = cast(getSource().getType()); + auto indexType = cast(getIndex().getType()); + auto resultType = cast(getResult().getType()); + + if (sourceType.getElementType() != resultType.getElementType()) + return emitOpError( + "requires result element type to match source element type"); + + if (indexType.getElementCount() != resultType.getElementCount()) + return emitOpError( + "requires index lane count to match result lane count"); + + if (!isa(indexType.getElementType())) + return emitOpError("requires index element type to be integer"); + + bool sourceHasLayout = isLayoutAssigned(sourceType); + bool indexHasLayout = isLayoutAssigned(indexType); + bool resultHasLayout = isLayoutAssigned(resultType); + if (sourceHasLayout != resultHasLayout) + return emitOpError("requires source and result to both carry layout or " + "neither carry layout"); + if (indexHasLayout && !sourceHasLayout) + return emitOpError( + "requires index to carry layout only when source does"); + + return success(); +} + +LogicalResult VMIVintlvOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto lowType = cast(getLow().getType()); + auto highType = cast(getHigh().getType()); + if (failed(verifyAllSameVRegShapeAndLayoutPresence( + getOperation(), {lhsType, rhsType, lowType, highType}, + /*requireSameElement=*/true))) + return failure(); + if (failed(verifyVMIPmodeMask(getOperation(), + cast(getMask().getType()), + lhsType, getPmode()))) + return failure(); + return success(); +} + +LogicalResult VMIVdintlvOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto lowType = cast(getLow().getType()); + auto highType = cast(getHigh().getType()); + if (failed(verifyAllSameVRegShapeAndLayoutPresence( + getOperation(), {lhsType, rhsType, lowType, highType}, + /*requireSameElement=*/true))) + return failure(); + if (failed(verifyVMIPmodeMask(getOperation(), + cast(getMask().getType()), + lhsType, getPmode()))) + return failure(); + return success(); +} + +//===----------------------------------------------------------------------===// +// VMIVabsOp verifier (unified fp/int abs, replaces absf/absi) +//===----------------------------------------------------------------------===// +// VMIVcmpOp / VMIVcmpsOp verifiers +LogicalResult VMIVcmpOp::verify() { + auto lhsType = cast(getLhs().getType()); + auto rhsType = cast(getRhs().getType()); + auto seedType = cast(getSeed().getType()); + auto resultType = cast(getResult().getType()); + + // Element type must be float-like OR integer-like (unified). + Type eltTy = lhsType.getElementType(); + if (!isVMIFloatLikeType(eltTy) && !isVMIIntegerLikeType(eltTy)) + return emitOpError("requires floating-point-like or integer-like VMI " + "element type for unified compare"); + if (isVMIIntegerLikeType(eltTy) && getCmp().starts_with("o")) + return emitOpError("requires integer compare predicate eq/ne/lt/le/gt/ge; " + "signedness is selected by the integer element type"); + + if (failed(verifyAllSameVRegShapeAndLayout(getOperation(), {lhsType, rhsType}, + /*requireSameElement=*/true))) + return failure(); + + // Validate cmp predicate. + if (!isSupportedVCmpPredicate(getCmp())) + return emitOpError("unsupported compare predicate '") + << getCmp() << "'; expected eq/ne/lt/le/gt/ge, " + << "or oeq/one/olt/ole/ogt/oge"; + + // Validate pmode. + if (auto pmode = getPmode()) { + if (pmode.value() != "zeroing" && pmode.value() != "merge") + return emitOpError("unsupported pmode '") + << pmode.value() << "'; expected \"zeroing\" or \"merge\""; + } + + // Seed mask must match data shape. + if (failed(verifyMaskMatchesData(getOperation(), seedType, lhsType))) + return failure(); + + // Result mask must match seed mask. + if (seedType.getElementCount() != resultType.getElementCount()) + return emitOpError( + "requires result mask lane count to match seed mask lane count"); + + return success(); +} + +LogicalResult VMIVcmpsOp::verify() { + auto srcType = cast(getSrc().getType()); + auto seedType = cast(getSeed().getType()); + auto resultType = cast(getResult().getType()); + + // Element type must be float-like OR integer-like (unified). + Type eltTy = srcType.getElementType(); + if (!isVMIFloatLikeType(eltTy) && !isVMIIntegerLikeType(eltTy)) + return emitOpError("requires floating-point-like or integer-like VMI " + "element type for unified compare"); + if (isVMIIntegerLikeType(eltTy) && getCmp().starts_with("o")) + return emitOpError("requires integer compare predicate eq/ne/lt/le/gt/ge; " + "signedness is selected by the integer element type"); + + // Scalar type must match vector element type. + Type scalarTy = getScalar().getType(); + if (scalarTy != eltTy) + return emitOpError("requires scalar type to match vector element type, " + "got scalar ") + << scalarTy << " vs vector element " << eltTy; + + // Validate cmp predicate. + if (!isSupportedVCmpPredicate(getCmp())) + return emitOpError("unsupported compare predicate '") + << getCmp() << "'; expected eq/ne/lt/le/gt/ge, " + << "or oeq/one/olt/ole/ogt/oge"; + + // Validate pmode. + if (auto pmode = getPmode()) { + if (pmode.value() != "zeroing" && pmode.value() != "merge") + return emitOpError("unsupported pmode '") + << pmode.value() << "'; expected \"zeroing\" or \"merge\""; + } + + // Seed mask must match data shape. + if (failed(verifyMaskMatchesData(getOperation(), seedType, srcType))) + return failure(); + + // Result mask must match seed mask. + if (seedType.getElementCount() != resultType.getElementCount()) + return emitOpError( + "requires result mask lane count to match seed mask lane count"); + + return success(); +} + +//===----------------------------------------------------------------------===// +// VMICvtOp — unified elementwise type conversion +//===----------------------------------------------------------------------===// +// VMIvLoadOp +ParseResult VMIvLoadOp::parse(OpAsmParser &parser, OperationState &result) { + OpAsmParser::UnresolvedOperand sourceOperand; + OpAsmParser::UnresolvedOperand offsetOperand; + OpAsmParser::UnresolvedOperand strideOperand; + OpAsmParser::UnresolvedOperand blockStrideOperand; + OpAsmParser::UnresolvedOperand repeatStrideOperand; + + // Parse: %source[%offset] + if (parser.parseOperand(sourceOperand) || parser.parseLSquare() || + parser.parseOperand(offsetOperand) || parser.parseRSquare()) + return failure(); + + // Optional comma-separated post-bracket operands. + // 1 operand + group attr → stride (group mode) + // 2 operands → block_stride, repeat_stride (block-stride mode) + int numPostBracket = 0; + OpAsmParser::UnresolvedOperand postOp1, postOp2; + if (succeeded(parser.parseOptionalComma())) { + if (parser.parseOperand(postOp1)) + return failure(); + numPostBracket = 1; + if (succeeded(parser.parseOptionalComma())) { + if (parser.parseOperand(postOp2)) + return failure(); + numPostBracket = 2; + } + } + + if (parser.parseOptionalAttrDict(result.attributes)) + return failure(); + + Type sourceType; + if (parser.parseColonType(sourceType)) + return failure(); + + if (parser.parseArrow()) + return failure(); + + SmallVector resultTypes; + if (parser.parseTypeList(resultTypes)) + return failure(); + + // Disambiguate post-bracket operands + bool hasStride = false; + bool hasBlock = false; + bool hasRepeat = false; + + if (numPostBracket == 2) { + // block_stride + repeat_stride pair + hasBlock = true; + hasRepeat = true; + blockStrideOperand = postOp1; + repeatStrideOperand = postOp2; + } else if (numPostBracket == 1) { + // Single post-bracket operand: only valid as group stride. + // block_stride without repeat_stride is invalid; verifier catches + // stride without group attr. + hasStride = true; + strideOperand = postOp1; + } + + if (parser.resolveOperand(sourceOperand, sourceType, result.operands)) + return failure(); + if (parser.resolveOperand(offsetOperand, parser.getBuilder().getIndexType(), + result.operands)) + return failure(); + if (hasStride && + parser.resolveOperand(strideOperand, parser.getBuilder().getIndexType(), + result.operands)) + return failure(); + if (hasBlock && + parser.resolveOperand(blockStrideOperand, + parser.getBuilder().getIntegerType(16), + result.operands)) + return failure(); + if (hasRepeat && + parser.resolveOperand(repeatStrideOperand, + parser.getBuilder().getIntegerType(16), + result.operands)) + return failure(); + + result.addAttribute("operandSegmentSizes", + parser.getBuilder().getDenseI32ArrayAttr( + {1, 1, hasStride ? 1 : 0, hasBlock ? 1 : 0, + hasRepeat ? 1 : 0})); + + result.addTypes(resultTypes); + return success(); +} + +void VMIvLoadOp::print(OpAsmPrinter &p) { + p << ' ' << getSource() << '['; + p.printOperand(getOffset()); + p << ']'; + if (getStride()) { + p << ", "; + p.printOperand(getStride()); + } + if (getBlockStride()) { + p << ", "; + p.printOperand(getBlockStride()); + p << ", "; + p.printOperand(getRepeatStride()); + } + p.printOptionalAttrDict((*this)->getAttrs(), {"operandSegmentSizes"}); + p << " : " << getSource().getType() << " -> " << getResults().getTypes(); +} + +LogicalResult VMIvLoadOp::verify() { + // group and dist_mode are mutually exclusive, except brc which supports + // group broadcast (one scalar per group → broadcast within each group). + if (getGroup() && getDistMode() && getDistMode() != "brc") + return emitOpError("group and dist_mode are mutually exclusive"); + if (getGroup() && !getStride()) + return emitOpError("group requires a stride operand"); + if (!getGroup() && getStride()) + return emitOpError("stride operand is only valid with group"); + + if (getGroup()) { + int64_t numGroups = getGroupAttr().getInt(); + if (numGroups <= 0) + return emitOpError("group must be positive, got ") << numGroups; + if (getResults().size() != 1) + return emitOpError("group mode requires exactly 1 result"); + return success(); + } + + // block_stride and repeat_stride must be paired, mutually exclusive + // with dist_mode and group + bool hasBlock = static_cast(getBlockStride()); + bool hasRepeat = static_cast(getRepeatStride()); + if (hasBlock != hasRepeat) + return emitOpError( + "block_stride and repeat_stride must both be present or absent"); + if (hasBlock) { + if (getDistMode()) + return emitOpError( + "block_stride and dist_mode are mutually exclusive"); + if (getResults().size() != 1) + return emitOpError("block-stride mode requires exactly 1 result"); + return success(); + } + + // result count vs dist-mode + auto distMode = getDistMode(); + bool isDintlv = distMode && *distMode == "dintlv"; + size_t nResults = getResults().size(); + if (isDintlv && nResults != 2) + return emitOpError("dist-mode \"dintlv\" requires exactly 2 results"); + if (!isDintlv && nResults != 1) + return emitOpError("requires exactly 1 result for dist-mode \"") + << (distMode ? *distMode : "continuous") << "\""; + + if (distMode && !validDistModes().count(*distMode)) + return emitOpError("invalid dist-mode: \"") << *distMode << "\""; + auto pmode = getPmode(); + if (pmode && !validPModes().count(*pmode)) + return emitOpError("invalid pmode: \"") << *pmode << "\""; + + bool isUnpack = distMode && *distMode == "unpack"; + for (auto res : getResults()) { + auto resType = cast(res.getType()); + // unpack: source element type intentionally differs from result + if (!isUnpack && + failed(verifyMemoryElementMatches(getOperation(), + getSource().getType(), resType, + "source"))) + return failure(); + if (isDintlv && + failed(verifyContiguousIfLayoutAssigned(getOperation(), resType, + "result"))) + return failure(); + } + + return success(); +} + +void VMIvLoadOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Read::get(), &getSourceMutable()); +} + +//===----------------------------------------------------------------------===// +// VMIvStoreOp + +FailureOr mlir::pto::getDataLanesPerPart(Type elementType) { + unsigned elementBitWidth = pto::getPTOStorageElemBitWidth(elementType); + if (elementBitWidth == 0) + return failure(); + constexpr int64_t kPhysicalVRegBits = 256 * 8; + if (kPhysicalVRegBits % elementBitWidth != 0) + return failure(); + return kPhysicalVRegBits / elementBitWidth; +} + +FailureOr mlir::pto::getMaskLanesPerPart(StringRef granularity) { + if (granularity == "b8") + return 256; + if (granularity == "b16") + return 128; + if (granularity == "b32") + return 64; + return failure(); +} + +FailureOr mlir::pto::getVMIPhysicalArity(Type type) { + FailureOr elementCount = getVMIElementCount(type); + FailureOr lanesPerPart = getPhysicalLanesPerPart(type); + FailureOr layout = getAssignedVMILayout(type); + if (failed(elementCount) || failed(lanesPerPart) || failed(layout)) + return failure(); + + if ((*layout).isGroupSlots() && (*layout).getSlots() > 0) + return divideCeilNonNegative((*layout).getNumGroups(), + (*layout).getSlots()); + + int64_t factor = (*layout).isDeinterleaved() ? (*layout).getFactor() : 1; + int64_t blockElems = + (*layout).isDeinterleaved() ? (*layout).getBlockElems() : 1; + int64_t laneStride = + isa(type) ? 1 + : ((*layout).isDense() ? (*layout).getLaneStride() + : 1); + int64_t arity = 0; + for (int64_t part = 0; part < factor; ++part) { + int64_t lanesInPart = + getDenseLogicalLanesInPart(*elementCount, factor, blockElems, part); + int64_t requiredPhysicalLanes = + lanesInPart == 0 ? 0 : (lanesInPart - 1) * laneStride + 1; + arity += divideCeilNonNegative(requiredPhysicalLanes, *lanesPerPart); + } + return arity; +} + +FailureOr +mlir::pto::mapLogicalLaneToPhysical(Type type, int64_t logicalLane) { + FailureOr elementCount = getVMIElementCount(type); + FailureOr factor = getLayoutFactor(type); + FailureOr blockElems = getLayoutBlockElems(type); + FailureOr laneStride = getDenseLaneStride(type); + FailureOr lanesPerPart = getPhysicalLanesPerPart(type); + if (failed(elementCount) || failed(factor) || failed(blockElems) || + failed(laneStride) || failed(lanesPerPart)) + return failure(); + if (logicalLane < 0 || logicalLane >= *elementCount) + return failure(); + + FailureOr layout = getAssignedVMILayout(type); + if (succeeded(layout) && (*layout).isGroupSlots() && + (*layout).getSlots() > 0) { + int64_t slots = (*layout).getSlots(); + int64_t lane = logicalLane % slots; + if (lane >= *lanesPerPart) + return failure(); + return VMIPhysicalLane{/*part=*/0, logicalLane / slots, lane}; + } + + int64_t part = 0; + std::optional indexInPart = mapDenseLogicalLaneToPartIndex( + *elementCount, *factor, *blockElems, logicalLane, part); + if (!indexInPart) + return failure(); + int64_t physicalIndex = *indexInPart * *laneStride; + return VMIPhysicalLane{part, physicalIndex / *lanesPerPart, + physicalIndex % *lanesPerPart}; +} + +FailureOr mlir::pto::mapPhysicalLaneToLogical(Type type, int64_t part, + int64_t chunk, + int64_t lane) { + FailureOr elementCount = getVMIElementCount(type); + FailureOr factor = getLayoutFactor(type); + FailureOr blockElems = getLayoutBlockElems(type); + FailureOr laneStride = getDenseLaneStride(type); + FailureOr lanesPerPart = getPhysicalLanesPerPart(type); + if (failed(elementCount) || failed(factor) || failed(blockElems) || + failed(laneStride) || failed(lanesPerPart)) + return failure(); + if (part < 0 || part >= *factor || chunk < 0 || lane < 0 || + lane >= *lanesPerPart) + return failure(); + + FailureOr layout = getAssignedVMILayout(type); + if (succeeded(layout) && (*layout).isGroupSlots() && + (*layout).getSlots() > 0) { + int64_t slots = (*layout).getSlots(); + if (part != 0 || lane >= slots) + return failure(); + int64_t logicalLane = chunk * slots + lane; + if (logicalLane >= *elementCount) + return failure(); + return logicalLane; + } + + int64_t physicalIndexInPart = chunk * *lanesPerPart + lane; + if (physicalIndexInPart % *laneStride != 0) + return failure(); + int64_t indexInPart = physicalIndexInPart / *laneStride; + std::optional logicalLane = mapDensePartIndexToLogicalLane( + *elementCount, *factor, *blockElems, part, indexInPart); + if (!logicalLane) + return failure(); + return *logicalLane; +} + +FailureOr mlir::pto::isPaddingLane(Type type, int64_t part, int64_t chunk, + int64_t lane) { + FailureOr elementCount = getVMIElementCount(type); + FailureOr factor = getLayoutFactor(type); + FailureOr blockElems = getLayoutBlockElems(type); + FailureOr laneStride = getDenseLaneStride(type); + FailureOr lanesPerPart = getPhysicalLanesPerPart(type); + if (failed(elementCount) || failed(factor) || failed(blockElems) || + failed(laneStride) || failed(lanesPerPart)) + return failure(); + if (part < 0 || part >= *factor || chunk < 0 || lane < 0 || + lane >= *lanesPerPart) + return failure(); + + FailureOr layout = getAssignedVMILayout(type); + if (succeeded(layout) && (*layout).isGroupSlots() && + (*layout).getSlots() > 0) { + int64_t slots = (*layout).getSlots(); + if (part != 0) + return true; + if (lane >= slots) + return true; + return chunk * slots + lane >= *elementCount; + } + + int64_t lanesInPart = + getDenseLogicalLanesInPart(*elementCount, *factor, *blockElems, part); + int64_t physicalIndexInPart = chunk * *lanesPerPart + lane; + if (physicalIndexInPart % *laneStride != 0) + return true; + int64_t indexInPart = physicalIndexInPart / *laneStride; + return indexInPart >= lanesInPart; +} diff --git a/lib/PTO/IR/VPTO.cpp b/lib/PTO/IR/VPTO.cpp index 29ad4e648d..e0c36d9456 100644 --- a/lib/PTO/IR/VPTO.cpp +++ b/lib/PTO/IR/VPTO.cpp @@ -1515,7 +1515,7 @@ static std::optional lookupVcvtContract(VcvtElemKind src, case VcvtElemKind::F8E5M2: return VcvtContract{/*requiresRnd=*/true, /*requiresSat=*/true, /*requiresPart=*/true, VcvtPartFamily::Packed4, - "R"}; + "RAHZ"}; case VcvtElemKind::HiF8: return VcvtContract{/*requiresRnd=*/true, /*requiresSat=*/true, /*requiresPart=*/true, VcvtPartFamily::Packed4, @@ -1745,6 +1745,10 @@ getVstsMaskGranularityOverride(StringRef dist, Type elementType) { return StringRef("b16"); if (dist == "PK_B32" || dist == "PK_B64" || dist == "PK4_B32") return StringRef("b32"); + if (dist == "PK_B64") + return StringRef("b32"); + if (dist == "PK4_B32") + return StringRef("b32"); return std::nullopt; } @@ -5473,9 +5477,10 @@ static LogicalResult verifyGroupReductionVecOp(ReductionOp op) { auto inputType = cast(op.getInput().getType()); Type elemType = inputType.getElementType(); if (auto intType = dyn_cast(elemType)) { - if (intType.getWidth() < 16 || intType.getWidth() > 32) + if (intType.getWidth() != 8 && intType.getWidth() != 16 && + intType.getWidth() != 32) return op.emitOpError( - "requires 16-bit or 32-bit integer vector element type"); + "requires 8-bit, 16-bit, or 32-bit integer vector element type"); return success(); } if (!elemType.isF16() && !elemType.isF32()) diff --git a/lib/PTO/Transforms/CMakeLists.txt b/lib/PTO/Transforms/CMakeLists.txt index d3986843e2..067183b1b9 100644 --- a/lib/PTO/Transforms/CMakeLists.txt +++ b/lib/PTO/Transforms/CMakeLists.txt @@ -30,11 +30,25 @@ add_mlir_dialect_library(PTOTransforms VPTOLLVMEmitterHelper.cpp VPTOPtrNormalize.cpp VPTOPtrCastCleanup.cpp + VPTONormalizeEquivalentVcvt.cpp VPTOExpandWrapperOps.cpp PTOVPTOPtrBoundary.cpp VPTOBufferMaterialization.cpp PTOValidateVPTOIR.cpp + PTONarrowVPTOLoopCounters.cpp PTOUnrollSIMTForPass.cpp + PTOValidateVMIIR.cpp + VMIPreAssignmentCombine.cpp + VMILegalizeArithSelect.cpp + VMIMaskGranularityAssignment.cpp + VMILowerUnifiedToLegacy.cpp + VMILayoutAssignment.cpp + VMILayoutFold.cpp + VMILayoutPropagation.cpp + VMILayoutSupport.cpp + VMILayoutRematerialize.cpp + VMILayoutSinkMaterialization.cpp + VMIToVPTO.cpp PTOInferVPTOVecScope.cpp InsertSync/PTOInsertSync.cpp diff --git a/lib/PTO/Transforms/PTONarrowVPTOLoopCounters.cpp b/lib/PTO/Transforms/PTONarrowVPTOLoopCounters.cpp new file mode 100644 index 0000000000..a35d53309f --- /dev/null +++ b/lib/PTO/Transforms/PTONarrowVPTOLoopCounters.cpp @@ -0,0 +1,159 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under +// the terms and conditions of CANN Open Software License Agreement Version 2.0 +// (the "License"). Please refer to the License for details. You may not use +// this file except in compliance with the License. THIS SOFTWARE IS PROVIDED ON +// AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS +// FOR A PARTICULAR PURPOSE. See LICENSE in the root of the software repository +// for the full text of the License. + +//===- PTONarrowVPTOLoopCounters.cpp -------------------------------------===// +// +// Narrow constant-bounded scf.for counters under VPTO vecscope regions to +// i16. The loop body continues to observe the original induction-variable +// type through a cast inserted at the beginning of the rewritten body. +// +//===----------------------------------------------------------------------===// + +#include "PTO/Transforms/Passes.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" + +#include "llvm/ADT/SmallVector.h" + +#include +#include +#include + +namespace mlir { +namespace pto { +#define GEN_PASS_DEF_PTONARROWVPTOLOOPCOUNTERS +#include "PTO/Transforms/Passes.h.inc" +} // namespace pto +} // namespace mlir + +using namespace mlir; + +namespace { + +static bool isNestedInVecScope(Operation *op) { + return op->getParentOfType() || + op->getParentOfType(); +} + +static bool isNarrowableCounterType(Type type) { + if (isa(type)) + return true; + auto integerType = dyn_cast(type); + return integerType && integerType.getWidth() > 16; +} + +static bool fitsSignedI16(int64_t value) { + return value >= std::numeric_limits::min() && + value <= std::numeric_limits::max(); +} + +static std::optional getSignedI16Constant(Value value) { + std::optional constant = getConstantIntValue(value); + if (!constant || !fitsSignedI16(*constant)) + return std::nullopt; + return constant; +} + +static bool exitValueFitsSignedI16(int64_t lower, int64_t upper, int64_t step) { + if (lower >= upper) + return true; + + int64_t distance = upper - lower; + int64_t iterationCount = (distance + step - 1) / step; + return fitsSignedI16(lower + iterationCount * step); +} + +static Value createI16Constant(PatternRewriter &rewriter, Location loc, + int64_t value) { + Type i16Type = rewriter.getI16Type(); + return rewriter.create( + loc, i16Type, rewriter.getIntegerAttr(i16Type, value)); +} + +static Value restoreInductionVariableType(PatternRewriter &rewriter, + Location loc, Value inductionVar, + Type originalType) { + if (isa(originalType)) + return rewriter.create(loc, originalType, inductionVar); + return rewriter.create(loc, originalType, inductionVar); +} + +struct NarrowVecScopeLoopCounterPattern : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(scf::ForOp forOp, + PatternRewriter &rewriter) const override { + if (!isNestedInVecScope(forOp)) + return failure(); + + Type originalCounterType = forOp.getInductionVar().getType(); + if (!isNarrowableCounterType(originalCounterType)) + return failure(); + + std::optional lower = getSignedI16Constant(forOp.getLowerBound()); + std::optional upper = getSignedI16Constant(forOp.getUpperBound()); + std::optional step = getSignedI16Constant(forOp.getStep()); + if (!lower || !upper || !step || *step <= 0 || + !exitValueFitsSignedI16(*lower, *upper, *step)) + return failure(); + + Location loc = forOp.getLoc(); + Value newLower = createI16Constant(rewriter, loc, *lower); + Value newUpper = createI16Constant(rewriter, loc, *upper); + Value newStep = createI16Constant(rewriter, loc, *step); + + auto newFor = rewriter.create(loc, newLower, newUpper, newStep, + forOp.getInitArgs()); + newFor->setAttrs(forOp->getAttrs()); + + Block *oldBody = forOp.getBody(); + Block *newBody = newFor.getBody(); + if (!newBody->empty()) + rewriter.eraseOp(newBody->getTerminator()); + + rewriter.setInsertionPointToStart(newBody); + Value restoredInductionVar = restoreInductionVariableType( + rewriter, loc, newFor.getInductionVar(), originalCounterType); + + SmallVector bodyArgumentReplacements; + bodyArgumentReplacements.push_back(restoredInductionVar); + bodyArgumentReplacements.append(newFor.getRegionIterArgs().begin(), + newFor.getRegionIterArgs().end()); + rewriter.mergeBlocks(oldBody, newBody, bodyArgumentReplacements); + + rewriter.replaceOp(forOp, newFor.getResults()); + return success(); + } +}; + +struct PTONarrowVPTOLoopCounters + : public pto::impl::PTONarrowVPTOLoopCountersBase< + PTONarrowVPTOLoopCounters> { + using pto::impl::PTONarrowVPTOLoopCountersBase< + PTONarrowVPTOLoopCounters>::PTONarrowVPTOLoopCountersBase; + + void runOnOperation() override { + RewritePatternSet patterns(&getContext()); + patterns.add(&getContext()); + if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) + signalPassFailure(); + } +}; + +} // namespace + +std::unique_ptr mlir::pto::createPTONarrowVPTOLoopCountersPass() { + return std::make_unique(); +} diff --git a/lib/PTO/Transforms/PTOValidateVMIIR.cpp b/lib/PTO/Transforms/PTOValidateVMIIR.cpp new file mode 100644 index 0000000000..84aafb1399 --- /dev/null +++ b/lib/PTO/Transforms/PTOValidateVMIIR.cpp @@ -0,0 +1,762 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +//===- PTOValidateVMIIR.cpp - VMI boundary verifier ----------------------===// +//===----------------------------------------------------------------------===// + +#include "PTO/IR/PTO.h" +#include "PTO/IR/VMIUtils.h" +#include "PTO/Transforms/Passes.h" +#include "PTO/Transforms/VMILayoutSupport.h" + +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/OperationSupport.h" +#include "mlir/Pass/Pass.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/StringSet.h" +#include "llvm/Support/raw_ostream.h" + +namespace mlir { +namespace pto { +#define GEN_PASS_DEF_PTOVALIDATEVMIIR +#define GEN_PASS_DEF_PTOVALIDATEVMILAYOUTIR +#include "PTO/Transforms/Passes.h.inc" +} // namespace pto +} // namespace mlir + +using namespace mlir; +using namespace mlir::pto; + +namespace { + +bool isVMIType(Type type) { return isa(type); } + +bool containsVMIType(Type type) { + if (isVMIType(type)) + return true; + + if (auto functionType = dyn_cast(type)) { + return llvm::any_of(functionType.getInputs(), + [](Type input) { return containsVMIType(input); }) || + llvm::any_of(functionType.getResults(), [](Type result) { + return containsVMIType(result); + }); + } + + if (auto shapedType = dyn_cast(type)) + return containsVMIType(shapedType.getElementType()); + + return false; +} + +bool containsVMIType(Attribute attr) { + if (!attr) + return false; + + if (auto typeAttr = dyn_cast(attr)) + if (containsVMIType(typeAttr.getValue())) + return true; + + if (auto typedAttr = dyn_cast(attr)) + if (containsVMIType(typedAttr.getType())) + return true; + + if (auto arrayAttr = dyn_cast(attr)) + return llvm::any_of(arrayAttr, [](Attribute element) { + return containsVMIType(element); + }); + + if (auto dictAttr = dyn_cast(attr)) + return llvm::any_of(dictAttr, [](NamedAttribute namedAttr) { + return containsVMIType(namedAttr.getValue()); + }); + + return false; +} + +bool isSurfaceVMIType(Type type) { + if (auto vregType = dyn_cast(type)) + return !vregType.getLayout(); + if (auto maskType = dyn_cast(type)) + return maskType.isPred() && !maskType.getLayout(); + return false; +} + +bool isLayoutAssignedVMIType(Type type) { + if (auto vregType = dyn_cast(type)) + return static_cast(vregType.getLayoutAttr()); + if (auto maskType = dyn_cast(type)) + return maskType.getLayoutAttr() && + VMIMaskType::isConcreteGranularity(maskType.getGranularity()); + return false; +} + +bool isVMIHelperOp(Operation *op) { + StringRef name = op->getName().getStringRef(); + return name == "pto.vmi.ensure_layout" || + name == "pto.vmi.ensure_mask_layout" || + name == "pto.vmi.ensure_mask_granularity" || name == "pto.vmi.pack" || + name == "pto.vmi.unpack"; +} + +bool isVMILayoutHelperOp(Operation *op) { + StringRef name = op->getName().getStringRef(); + return name == "pto.vmi.ensure_layout" || + name == "pto.vmi.ensure_mask_layout" || + name == "pto.vmi.ensure_mask_granularity"; +} + +bool isVMISemanticOp(Operation *op) { + StringRef name = op->getName().getStringRef(); + return name.starts_with("pto.vmi.") && !isVMIHelperOp(op); +} + +bool isStructuralOp(Operation *op) { + StringRef name = op->getName().getStringRef(); + return name == "builtin.module" || name.starts_with("func.") || + name.starts_with("scf.") || name.starts_with("cf."); +} + +bool hasVMIType(Operation *op) { + if (llvm::any_of(op->getOperandTypes(), isVMIType) || + llvm::any_of(op->getResultTypes(), isVMIType)) + return true; + for (Region ®ion : op->getRegions()) { + for (Block &block : region) { + if (llvm::any_of(block.getArgumentTypes(), isVMIType)) + return true; + } + } + return false; +} + +void mirrorDiagnostic(llvm::raw_ostream *diagOS, Twine message) { + if (diagOS) + *diagOS << message << "\n"; +} + +LogicalResult emitInvariant(Operation *op, llvm::raw_ostream *diagOS, + Twine message) { + InFlightDiagnostic diag = op->emitError() + << kVMIDiagPassInvariantPrefix << message; + (void)diag; + mirrorDiagnostic(diagOS, Twine(kVMIDiagPassInvariantPrefix) + message); + return failure(); +} + +LogicalResult emitLayoutContract(Operation *op, llvm::raw_ostream *diagOS, + Twine message) { + InFlightDiagnostic diag = op->emitError() + << kVMIDiagLayoutContractPrefix << message; + (void)diag; + mirrorDiagnostic(diagOS, Twine(kVMIDiagLayoutContractPrefix) + message); + return failure(); +} + +LogicalResult emitLayoutSupportContract(Operation *op, + llvm::raw_ostream *diagOS, + Twine message, StringRef reason) { + std::string text; + llvm::raw_string_ostream os(text); + os << message << ": " << reason; + + bool printedAny = false; + auto printValueType = [&](StringRef kind, int64_t index, Type type) { + if (!isVMIType(type)) + return; + if (!printedAny) { + os << "; VMI types:"; + printedAny = true; + } + os << " " << kind << "#" << index << "=" << type; + }; + + for (auto [index, operand] : llvm::enumerate(op->getOperands())) + printValueType("operand", static_cast(index), operand.getType()); + for (auto [index, result] : llvm::enumerate(op->getResults())) + printValueType("result", static_cast(index), result.getType()); + + os.flush(); + return emitLayoutContract(op, diagOS, text); +} + +LogicalResult +emitHelperMaterializationContract(Operation *helper, Type sourceType, + Type resultType, StringRef helperName, + StringRef reason, llvm::raw_ostream *diagOS) { + auto emitFallback = [&]() { + return emitLayoutContract( + helper, diagOS, + Twine(helperName) + + " has no registered materialization support: " + reason); + }; + + if (helper->getNumResults() != 1 || !helper->getResult(0).hasOneUse()) + return emitFallback(); + + OpOperand &use = *helper->getResult(0).use_begin(); + Operation *requester = use.getOwner(); + std::string message; + llvm::raw_string_ostream os(message); + os << requester->getName() << " operand #" << use.getOperandNumber() + << " has type " << sourceType << " but requires " << resultType << "; " + << helperName << " has no registered materialization support: " << reason; + os.flush(); + + InFlightDiagnostic diag = requester->emitError() + << kVMIDiagLayoutContractPrefix << message; + diag.attachNote(helper->getLoc()) + << "failed helper conversion " << sourceType << " -> " << resultType + << " (" << reason << ")"; + mirrorDiagnostic(diagOS, Twine(kVMIDiagLayoutContractPrefix) + message); + return failure(); +} + +LogicalResult verifyBoundaryType(Operation *owner, Type type, + llvm::raw_ostream *diagOS) { + if (isVMIType(type) && !isSurfaceVMIType(type)) + return emitInvariant( + owner, diagOS, + "VMI producer boundary requires surface !pto.vmi.vreg or " + "!pto.vmi.mask type"); + + return success(); +} + +LogicalResult verifyBoundaryTypeTree(Operation *owner, Type type, + llvm::raw_ostream *diagOS) { + if (failed(verifyBoundaryType(owner, type, diagOS))) + return failure(); + + if (auto functionType = dyn_cast(type)) { + for (Type input : functionType.getInputs()) + if (failed(verifyBoundaryTypeTree(owner, input, diagOS))) + return failure(); + for (Type result : functionType.getResults()) + if (failed(verifyBoundaryTypeTree(owner, result, diagOS))) + return failure(); + } + + if (auto shapedType = dyn_cast(type)) + return verifyBoundaryTypeTree(owner, shapedType.getElementType(), diagOS); + + return success(); +} + +LogicalResult verifyLayoutAssignedType(Operation *owner, Type type, + llvm::raw_ostream *diagOS) { + if (isVMIType(type) && !isLayoutAssignedVMIType(type)) + return emitInvariant( + owner, diagOS, + "layout-assigned VMI IR requires !pto.vmi.vreg with layout and " + "!pto.vmi.mask with b8/b16/b32 granularity plus layout"); + + return success(); +} + +LogicalResult verifyLayoutAssignedTypeTree(Operation *owner, Type type, + llvm::raw_ostream *diagOS) { + if (failed(verifyLayoutAssignedType(owner, type, diagOS))) + return failure(); + + if (auto functionType = dyn_cast(type)) { + for (Type input : functionType.getInputs()) + if (failed(verifyLayoutAssignedTypeTree(owner, input, diagOS))) + return failure(); + for (Type result : functionType.getResults()) + if (failed(verifyLayoutAssignedTypeTree(owner, result, diagOS))) + return failure(); + } + + if (auto shapedType = dyn_cast(type)) + return verifyLayoutAssignedTypeTree(owner, shapedType.getElementType(), + diagOS); + + return success(); +} + +template +LogicalResult verifyAttributeTypes(Operation *owner, Attribute attr, + llvm::raw_ostream *diagOS, + TypeVerifier verifyType) { + if (!attr) + return success(); + + if (auto typeAttr = dyn_cast(attr)) + if (failed(verifyType(owner, typeAttr.getValue(), diagOS))) + return failure(); + + if (auto typedAttr = dyn_cast(attr)) + if (failed(verifyType(owner, typedAttr.getType(), diagOS))) + return failure(); + + if (auto arrayAttr = dyn_cast(attr)) { + for (Attribute element : arrayAttr) + if (failed(verifyAttributeTypes(owner, element, diagOS, verifyType))) + return failure(); + } + + if (auto dictAttr = dyn_cast(attr)) { + for (NamedAttribute namedAttr : dictAttr) + if (failed(verifyAttributeTypes(owner, namedAttr.getValue(), diagOS, + verifyType))) + return failure(); + } + + return success(); +} + +bool isFunctionTypeAttr(Operation *op, NamedAttribute attr) { + return isa(op) && attr.getName() == "function_type"; +} + +LogicalResult verifyNoHiddenVMIAttributeType(Operation *op, NamedAttribute attr, + llvm::raw_ostream *diagOS) { + if (isFunctionTypeAttr(op, attr)) + return success(); + if (containsVMIType(attr.getValue())) + return emitInvariant(op, diagOS, + "VMI type appears in a non-signature attribute"); + return success(); +} + +LogicalResult verifyOperationTypes(Operation *op, llvm::raw_ostream *diagOS) { + if (auto funcOp = dyn_cast(op)) { + FunctionType functionType = funcOp.getFunctionType(); + for (Type type : functionType.getInputs()) + if (failed(verifyBoundaryTypeTree(op, type, diagOS))) + return failure(); + for (Type type : functionType.getResults()) + if (failed(verifyBoundaryTypeTree(op, type, diagOS))) + return failure(); + } + + for (Type type : op->getOperandTypes()) + if (failed(verifyBoundaryTypeTree(op, type, diagOS))) + return failure(); + for (Type type : op->getResultTypes()) + if (failed(verifyBoundaryTypeTree(op, type, diagOS))) + return failure(); + for (Region ®ion : op->getRegions()) { + for (Block &block : region) { + for (Type type : block.getArgumentTypes()) { + if (failed(verifyBoundaryTypeTree(op, type, diagOS))) + return failure(); + } + } + } + for (NamedAttribute attr : op->getAttrs()) { + if (failed(verifyNoHiddenVMIAttributeType(op, attr, diagOS))) + return failure(); + if (failed(verifyAttributeTypes(op, attr.getValue(), diagOS, + verifyBoundaryTypeTree))) + return failure(); + } + return success(); +} + +LogicalResult verifyLayoutAssignedOperationTypes(Operation *op, + llvm::raw_ostream *diagOS) { + if (auto funcOp = dyn_cast(op)) { + FunctionType functionType = funcOp.getFunctionType(); + for (Type type : functionType.getInputs()) + if (failed(verifyLayoutAssignedTypeTree(op, type, diagOS))) + return failure(); + for (Type type : functionType.getResults()) + if (failed(verifyLayoutAssignedTypeTree(op, type, diagOS))) + return failure(); + } + + for (Type type : op->getOperandTypes()) + if (failed(verifyLayoutAssignedTypeTree(op, type, diagOS))) + return failure(); + for (Type type : op->getResultTypes()) + if (failed(verifyLayoutAssignedTypeTree(op, type, diagOS))) + return failure(); + for (Region ®ion : op->getRegions()) { + for (Block &block : region) { + for (Type type : block.getArgumentTypes()) { + if (failed(verifyLayoutAssignedTypeTree(op, type, diagOS))) + return failure(); + } + } + } + for (NamedAttribute attr : op->getAttrs()) { + if (failed(verifyNoHiddenVMIAttributeType(op, attr, diagOS))) + return failure(); + if (failed(verifyAttributeTypes(op, attr.getValue(), diagOS, + verifyLayoutAssignedTypeTree))) + return failure(); + } + return success(); +} + +LogicalResult verifyLayoutHelperSupport(Operation *op, + llvm::raw_ostream *diagOS); + +LogicalResult verifyLayoutSemanticSupport(Operation *op, + llvm::raw_ostream *diagOS); + +LogicalResult verifyOperationBoundary(Operation *op, + llvm::raw_ostream *diagOS) { + if (failed(verifyOperationTypes(op, diagOS))) + return failure(); + + if (!hasVMIType(op)) + return success(); + + if (isVMIHelperOp(op)) + return emitInvariant( + op, diagOS, + "VMI helper op appears before layout assignment or VMI-to-VPTO"); + + if (isVMISemanticOp(op) || isStructuralOp(op)) + return success(); + + return emitInvariant(op, diagOS, + "VMI typed value is used by a non-VMI semantic op"); +} + +LogicalResult verifyLayoutAssignedOperation(Operation *op, + llvm::raw_ostream *diagOS, + bool verifyHelperSupports = true) { + if (failed(verifyLayoutAssignedOperationTypes(op, diagOS))) + return failure(); + + if (!hasVMIType(op)) + return success(); + + if (isVMIHelperOp(op)) { + if (isVMILayoutHelperOp(op)) + return verifyHelperSupports ? verifyLayoutHelperSupport(op, diagOS) + : success(); + return emitInvariant( + op, diagOS, + "VMI pack/unpack helper appears before VMI-to-VPTO physicalization"); + } + + if (isVMISemanticOp(op)) + return verifyLayoutSemanticSupport(op, diagOS); + if (isStructuralOp(op)) + return success(); + + return emitInvariant(op, diagOS, + "VMI typed value is used by a non-VMI semantic op"); +} + +LogicalResult verifyLayoutHelperSupport(Operation *op, + llvm::raw_ostream *diagOS) { + VMILayoutSupport supports; + + if (auto ensure = dyn_cast(op)) { + auto sourceType = cast(ensure.getSource().getType()); + auto resultType = cast(ensure.getResult().getType()); + std::string reason; + if (failed(supports.getEnsureLayoutFact(sourceType, resultType, &reason))) + return emitHelperMaterializationContract( + op, sourceType, resultType, "pto.vmi.ensure_layout", reason, diagOS); + return success(); + } + + if (auto ensure = dyn_cast(op)) { + auto sourceType = cast(ensure.getSource().getType()); + auto resultType = cast(ensure.getResult().getType()); + std::string reason; + if (failed( + supports.getEnsureMaskLayoutFact(sourceType, resultType, &reason))) + return emitHelperMaterializationContract(op, sourceType, resultType, + "pto.vmi.ensure_mask_layout", + reason, diagOS); + return success(); + } + + return success(); +} + +LogicalResult verifyLayoutSemanticSupport(Operation *op, + llvm::raw_ostream *diagOS) { + VMILayoutSupport supports; + + if (auto store = dyn_cast(op)) { + auto valueType = cast(store.getValue().getType()); + VMILayoutAttr layout = valueType.getLayoutAttr(); + if (!layout || layout.isContiguous()) + return success(); + + std::string reason; + if (failed(supports.getStoreLayoutFact(valueType, &reason))) + return emitLayoutSupportContract( + op, diagOS, + "pto.vmi.store has no registered contiguous-memory layout support", + reason); + return success(); + } + + if (auto load = dyn_cast(op)) { + auto resultType = cast(load.getResult().getType()); + VMILayoutAttr layout = resultType.getLayoutAttr(); + if (!layout) + return success(); + + std::string reason; + if (failed(supports.getGroupLoadLayoutFact(load, &reason))) + return emitLayoutSupportContract( + op, diagOS, + "pto.vmi.group_load has no registered layout support", reason); + return success(); + } + + if (auto load = dyn_cast(op)) { + auto resultType = cast(load.getResult().getType()); + std::string reason; + if (failed(supports.getGroupSlotLoadLayoutFact( + resultType, load.getNumGroupsAttr().getInt(), &reason))) + return emitLayoutSupportContract( + op, diagOS, + "pto.vmi.group_slot_load has no registered layout support", reason); + return success(); + } + + if (auto load = dyn_cast(op)) { + std::string reason; + if (failed(supports.getGroupBroadcastLoadSupport(load, &reason))) + return emitLayoutSupportContract( + op, diagOS, + "pto.vmi.group_broadcast_load has no registered layout support", + reason); + return success(); + } + + if (auto store = dyn_cast(op)) { + auto valueType = cast(store.getValue().getType()); + VMILayoutAttr layout = valueType.getLayoutAttr(); + if (!layout || !layout.isGroupSlots()) + return success(); + + std::string reason; + if (failed(supports.getGroupStoreLayoutFact( + valueType, store.getNumGroupsAttr().getInt(), &reason))) + return emitLayoutSupportContract( + op, diagOS, + "pto.vmi.group_store has no registered group_slots layout support", + reason); + return success(); + } + + if (auto reduce = dyn_cast(op)) { + auto resultType = cast(reduce.getResult().getType()); + VMILayoutAttr layout = resultType.getLayoutAttr(); + if (!layout || !layout.isGroupSlots()) + return success(); + + std::string reason; + if (failed(supports.getGroupReduceAddFSupport(reduce, &reason))) + return emitLayoutSupportContract( + op, diagOS, + "pto.vmi.group_reduce_addf has no registered group_slots layout " + "support", + reason); + return success(); + } + + if (auto reduce = dyn_cast(op)) { + auto resultType = cast(reduce.getResult().getType()); + VMILayoutAttr layout = resultType.getLayoutAttr(); + if (!layout || !layout.isGroupSlots()) + return success(); + + std::string reason; + if (failed(supports.getGroupReduceMaxFSupport(reduce, &reason))) + return emitLayoutSupportContract( + op, diagOS, + "pto.vmi.group_reduce_maxf has no registered group_slots layout " + "support", + reason); + return success(); + } + + if (auto reduce = dyn_cast(op)) { + auto resultType = cast(reduce.getResult().getType()); + VMILayoutAttr layout = resultType.getLayoutAttr(); + if (!layout || !layout.isGroupSlots()) + return success(); + + std::string reason; + if (failed(supports.getGroupReduceMinFSupport(reduce, &reason))) + return emitLayoutSupportContract( + op, diagOS, + "pto.vmi.group_reduce_minf has no registered group_slots layout " + "support", + reason); + return success(); + } + + if (auto reduce = dyn_cast(op)) { + auto resultType = cast(reduce.getResult().getType()); + VMILayoutAttr layout = resultType.getLayoutAttr(); + if (!layout || !layout.isGroupSlots()) + return success(); + + std::string reason; + if (failed(supports.getGroupReduceAddISupport(reduce, &reason))) + return emitLayoutSupportContract( + op, diagOS, + "pto.vmi.group_reduce_addi has no registered group_slots layout " + "support", + reason); + return success(); + } + + if (auto reduce = dyn_cast(op)) { + auto resultType = cast(reduce.getResult().getType()); + VMILayoutAttr layout = resultType.getLayoutAttr(); + if (!layout || !layout.isGroupSlots()) + return success(); + + std::string reason; + if (failed(supports.getGroupReduceMaxISupport(reduce, &reason))) + return emitLayoutSupportContract( + op, diagOS, + "pto.vmi.group_reduce_maxi has no registered group_slots layout " + "support", + reason); + return success(); + } + + if (auto reduce = dyn_cast(op)) { + auto resultType = cast(reduce.getResult().getType()); + VMILayoutAttr layout = resultType.getLayoutAttr(); + if (!layout || !layout.isGroupSlots()) + return success(); + + std::string reason; + if (failed(supports.getGroupReduceMinISupport(reduce, &reason))) + return emitLayoutSupportContract( + op, diagOS, + "pto.vmi.group_reduce_mini has no registered group_slots layout " + "support", + reason); + return success(); + } + + if (auto broadcast = dyn_cast(op)) { + auto sourceType = cast(broadcast.getSource().getType()); + VMILayoutAttr layout = sourceType.getLayoutAttr(); + if (!layout || !layout.isGroupSlots() || layout.getSlots() <= 0) + return success(); + + std::string reason; + if (failed(supports.getGroupBroadcastSupport(broadcast, &reason))) + return emitLayoutSupportContract( + op, diagOS, + "pto.vmi.group_broadcast has no registered layout support", reason); + return success(); + } + + if (auto hist = dyn_cast(op)) { + std::string reason; + if (failed(supports.getVdhistSupport(hist, &reason))) + return emitLayoutSupportContract( + op, diagOS, "pto.vmi.vdhist has no registered histogram support", + reason); + return success(); + } + + if (auto hist = dyn_cast(op)) { + std::string reason; + if (failed(supports.getVchistSupport(hist, &reason))) + return emitLayoutSupportContract( + op, diagOS, "pto.vmi.vchist has no registered histogram support", + reason); + return success(); + } + + if (auto truncf = dyn_cast(op)) { + std::string reason; + if (failed(supports.getTruncFSupport(truncf, &reason))) + return emitLayoutSupportContract( + op, diagOS, "pto.vmi.truncf has no registered layout support", + reason); + return success(); + } + + if (auto extf = dyn_cast(op)) { + std::string reason; + if (failed(supports.getExtFSupport(extf, &reason))) + return emitLayoutSupportContract( + op, diagOS, "pto.vmi.extf has no registered layout support", reason); + return success(); + } + + if (auto bitcast = dyn_cast(op)) { + std::string reason; + if (failed(supports.getBitcastSupport(bitcast, &reason))) + return emitLayoutSupportContract( + op, diagOS, "pto.vmi.bitcast has no registered layout support", + reason); + return success(); + } + + return success(); +} + +struct PTOValidateVMIIRPass + : public mlir::pto::impl::PTOValidateVMIIRBase { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PTOValidateVMIIRPass) + + void runOnOperation() override { + if (failed(validateVMIProducerBoundaryIR(getOperation(), &llvm::errs()))) + signalPassFailure(); + } +}; + +struct PTOValidateVMILayoutIRPass + : public mlir::pto::impl::PTOValidateVMILayoutIRBase< + PTOValidateVMILayoutIRPass> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PTOValidateVMILayoutIRPass) + + void runOnOperation() override { + if (failed(validateVMILayoutAssignedIR(getOperation(), &llvm::errs()))) + signalPassFailure(); + } +}; + +} // namespace + +LogicalResult +mlir::pto::validateVMIProducerBoundaryIR(ModuleOp module, + llvm::raw_ostream *diagOS) { + WalkResult result = module.walk([&](Operation *op) { + if (failed(verifyOperationBoundary(op, diagOS))) + return WalkResult::interrupt(); + return WalkResult::advance(); + }); + return failure(result.wasInterrupted()); +} + +LogicalResult mlir::pto::validateVMILayoutAssignedIR( + ModuleOp module, llvm::raw_ostream *diagOS, bool verifyHelperSupports) { + WalkResult result = module.walk([&](Operation *op) { + if (failed(verifyLayoutAssignedOperation(op, diagOS, verifyHelperSupports))) + return WalkResult::interrupt(); + return WalkResult::advance(); + }); + return failure(result.wasInterrupted()); +} + +std::unique_ptr mlir::pto::createPTOValidateVMIIRPass() { + return std::make_unique(); +} + +std::unique_ptr mlir::pto::createPTOValidateVMILayoutIRPass() { + return std::make_unique(); +} diff --git a/lib/PTO/Transforms/VMILayoutAssignment.cpp b/lib/PTO/Transforms/VMILayoutAssignment.cpp new file mode 100644 index 0000000000..5e99337006 --- /dev/null +++ b/lib/PTO/Transforms/VMILayoutAssignment.cpp @@ -0,0 +1,2011 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +//===- VMILayoutAssignment.cpp - Assign VMI layouts ----------------------===// +//===----------------------------------------------------------------------===// + +#include "PTO/IR/PTO.h" +#include "PTO/IR/PTOTypeUtils.h" +#include "PTO/IR/VMIUtils.h" +#include "PTO/Transforms/Passes.h" +#include "PTO/Transforms/VMILayoutPropagation.h" +#include "PTO/Transforms/VMILayoutSupport.h" + +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/SymbolTable.h" +#include "mlir/IR/Value.h" +#include "mlir/Pass/Pass.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/STLFunctionalExtras.h" +#include "llvm/ADT/SmallVector.h" + +namespace mlir { +namespace pto { +#define GEN_PASS_DEF_VMILAYOUTASSIGNMENT +#include "PTO/Transforms/Passes.h.inc" +} // namespace pto +} // namespace mlir + +using namespace mlir; +using namespace mlir::pto; + +namespace { + +struct DataNode { + Value value; + VMIVRegType type; + unsigned parent = 0; + VMILayoutAttr naturalLayout; + VMILayoutAttr preferredLayout; +}; + +struct MaskNode { + Value value; + VMIMaskType type; + unsigned parent = 0; + VMILayoutAttr requestedLayout; +}; + +enum class DataLayoutSeedPhase { + Explicit, + SeedStart, + GroupLoad = SeedStart, + Reduce, + GroupSlotLoad, + GroupBroadcast, + GroupBroadcastLoad, + Cast, + WeakReduce, + Store, + Other, + SeedEnd, +}; + +struct DataLayoutSeed { + Value value; + VMILayoutAttr layout; + DataLayoutSeedPhase phase = DataLayoutSeedPhase::Other; +}; + +struct DataUseRequest { + OpOperand *operand; + VMILayoutAttr layout; + bool late = false; + DataLayoutSeedPhase phase = DataLayoutSeedPhase::Other; +}; + +struct MaskUseRequest { + OpOperand *operand; + VMILayoutAttr layout; + DataLayoutSeedPhase phase = DataLayoutSeedPhase::Other; +}; + +struct GroupStoreUseRequest { + VMIGroupStoreOp store; +}; + +static std::optional getConstantIndexValue(Value value) { + if (auto constant = value.getDefiningOp()) + return constant.value(); + if (auto constant = value.getDefiningOp()) + if (auto integerAttr = dyn_cast(constant.getValue())) + return integerAttr.getInt(); + return std::nullopt; +} + +static bool isLane0SplatShuffle(VMIShuffleOp op) { + auto sourceType = cast(op.getSource().getType()); + ArrayRef indices = op.getIndices(); + return sourceType.getElementCount() == 1 && !indices.empty() && + llvm::all_of(indices, [](int64_t index) { return index == 0; }); +} + +bool containsVMIType(Type type) { + if (isa(type)) + return true; + if (auto functionType = dyn_cast(type)) { + return llvm::any_of(functionType.getInputs(), + [](Type input) { return containsVMIType(input); }) || + llvm::any_of(functionType.getResults(), + [](Type result) { return containsVMIType(result); }); + } + if (auto shapedType = dyn_cast(type)) + return containsVMIType(shapedType.getElementType()); + return false; +} + +struct LayoutSolver { + explicit LayoutSolver(ModuleOp module) + : module(module), ctx(module.getContext()) {} + + unsigned addDataValue(Value value) { + auto type = dyn_cast(value.getType()); + if (!type) + return ~0u; + auto [it, inserted] = dataIds.try_emplace(value, dataNodes.size()); + if (inserted) { + dataNodes.push_back( + DataNode{value, type, it->second, type.getLayoutAttr(), {}}); + if (type.getLayoutAttr()) + dataLayoutSeeds.push_back(DataLayoutSeed{ + value, type.getLayoutAttr(), DataLayoutSeedPhase::Explicit}); + } + return it->second; + } + + unsigned addMaskValue(Value value) { + auto type = dyn_cast(value.getType()); + if (!type) + return ~0u; + auto [it, inserted] = maskIds.try_emplace(value, maskNodes.size()); + if (inserted) + maskNodes.push_back( + MaskNode{value, type, it->second, type.getLayoutAttr()}); + return it->second; + } + + unsigned find(unsigned id) { + if (dataNodes[id].parent == id) + return id; + dataNodes[id].parent = find(dataNodes[id].parent); + return dataNodes[id].parent; + } + + unsigned findMask(unsigned id) { + if (maskNodes[id].parent == id) + return id; + maskNodes[id].parent = findMask(maskNodes[id].parent); + return maskNodes[id].parent; + } + + LogicalResult unite(Value lhs, Value rhs, Operation *op) { + (void)op; + addDataValue(lhs); + addDataValue(rhs); + return success(); + } + + LogicalResult uniteDataEquivalent(Value lhs, Value rhs, Operation *op) { + unsigned lhsId = addDataValue(lhs); + unsigned rhsId = addDataValue(rhs); + if (lhsId == ~0u || rhsId == ~0u) + return success(); + unsigned lhsRoot = find(lhsId); + unsigned rhsRoot = find(rhsId); + if (lhsRoot == rhsRoot) + return success(); + + DataNode &lhsNode = dataNodes[lhsRoot]; + DataNode &rhsNode = dataNodes[rhsRoot]; + if (lhsNode.naturalLayout && rhsNode.naturalLayout && + lhsNode.naturalLayout != rhsNode.naturalLayout) + return op->emitError() + << kVMIDiagLayoutContractPrefix << "conflicting natural layouts " + << lhsNode.naturalLayout << " and " << rhsNode.naturalLayout; + if (lhsNode.preferredLayout && rhsNode.preferredLayout && + lhsNode.preferredLayout != rhsNode.preferredLayout) + return op->emitError() + << kVMIDiagLayoutContractPrefix << "conflicting preferred layouts " + << lhsNode.preferredLayout << " and " << rhsNode.preferredLayout; + + rhsNode.parent = lhsRoot; + if (!lhsNode.naturalLayout) + lhsNode.naturalLayout = rhsNode.naturalLayout; + if (!lhsNode.preferredLayout) + lhsNode.preferredLayout = rhsNode.preferredLayout; + return success(); + } + + LogicalResult uniteMask(Value lhs, Value rhs, Operation *op) { + unsigned lhsId = addMaskValue(lhs); + unsigned rhsId = addMaskValue(rhs); + if (lhsId == ~0u || rhsId == ~0u) + return success(); + unsigned lhsRoot = findMask(lhsId); + unsigned rhsRoot = findMask(rhsId); + if (lhsRoot == rhsRoot) + return success(); + + MaskNode &lhsNode = maskNodes[lhsRoot]; + MaskNode &rhsNode = maskNodes[rhsRoot]; + if (lhsNode.requestedLayout && rhsNode.requestedLayout && + lhsNode.requestedLayout != rhsNode.requestedLayout) + return op->emitError() + << kVMIDiagLayoutContractPrefix << "conflicting mask layouts " + << lhsNode.requestedLayout << " and " << rhsNode.requestedLayout; + rhsNode.parent = lhsRoot; + if (!lhsNode.requestedLayout) + lhsNode.requestedLayout = rhsNode.requestedLayout; + return success(); + } + + LogicalResult + setNaturalLayout(Value value, VMILayoutAttr layout, Operation *op, + DataLayoutSeedPhase phase = DataLayoutSeedPhase::Other) { + unsigned id = addDataValue(value); + if (id == ~0u || !layout) + return success(); + unsigned root = find(id); + VMILayoutAttr existing = dataNodes[root].naturalLayout; + if (existing && existing != layout) + return op->emitError() + << kVMIDiagLayoutContractPrefix << "conflicting natural layouts " + << existing << " and " << layout; + dataNodes[root].naturalLayout = layout; + dataLayoutSeeds.push_back(DataLayoutSeed{value, layout, phase}); + return success(); + } + + LogicalResult + setPreferredLayout(Value value, VMILayoutAttr layout, Operation *op, + DataLayoutSeedPhase phase = DataLayoutSeedPhase::Other) { + unsigned id = addDataValue(value); + if (id == ~0u || !layout) + return success(); + unsigned root = find(id); + VMILayoutAttr existing = dataNodes[root].preferredLayout; + if (existing && existing != layout) + return op->emitError() + << kVMIDiagLayoutContractPrefix << "conflicting preferred layouts " + << existing << " and " << layout; + dataNodes[root].preferredLayout = layout; + dataLayoutSeeds.push_back(DataLayoutSeed{value, layout, phase}); + return success(); + } + + VMILayoutAttr getContiguousLayout() { + return VMILayoutAttr::getContiguous(ctx); + } + + VMILayoutAttr getPreferredDenseStoreLayout(VMIVRegType type) { + VMILayoutSupport supports; + FailureOr fact = + supports.getPreferredStoreLayoutFact(type); + if (failed(fact)) + return {}; + return fact->valueLayout; + } + + bool hasDataLayoutSeed(Value value) { + unsigned id = addDataValue(value); + if (id == ~0u) + return false; + DataNode &node = dataNodes[find(id)]; + return static_cast(node.naturalLayout || node.preferredLayout); + } + + FailureOr + getPreferredDenseMaskedStoreLayout(VMIVRegType valueType, + VMIMaskType maskType) { + VMILayoutSupport supports; + return supports.getPreferredMaskedStoreLayoutFact(valueType, maskType); + } + + VMILayoutAttr getGroupSlotsLayout(int64_t numGroups) { + return VMILayoutAttr::getGroupSlots(ctx, numGroups); + } + + VMILayoutAttr getPreferredGroupSlotsLayout(VMIVRegType type, + int64_t numGroups) { + if (VMILayoutAttr existing = type.getLayoutAttr()) + if (existing.isGroupSlots() && existing.getSlots() > 0) + return existing; + VMILayoutSupport supports; + FailureOr fact = + supports.getPreferredGroupReduceLayoutFact(type, numGroups); + if (succeeded(fact)) + return fact->resultLayout; + return getGroupSlotsLayout(numGroups); + } + + VMILayoutAttr getPreferredGroupReduceSourceLayout(VMIVRegType type, + int64_t numGroups) { + if (VMILayoutAttr existing = type.getLayoutAttr()) + return existing; + VMILayoutSupport supports; + FailureOr fact = + supports.getPreferredGroupReduceLayoutFact(type, numGroups); + if (succeeded(fact)) + return fact->sourceLayout; + return getContiguousLayout(); + } + + DataLayoutSeedPhase getGroupReduceUseSeedPhase(VMIVRegType sourceType, + int64_t numGroups, + VMIGroupReduceLayoutFact fact) { + if (!fact.sourceLayout || !fact.sourceLayout.isContiguous() || + fact.sourceLayout.getLaneStride() != 1) + return DataLayoutSeedPhase::Reduce; + + VMILayoutSupport supports; + FailureOr> resultFacts = + supports.getGroupReduceLayoutFactsForLayout( + sourceType, numGroups, VMIGroupReduceLayoutPort::Result, + fact.resultLayout); + if (succeeded(resultFacts) && resultFacts->size() > 1) + return DataLayoutSeedPhase::WeakReduce; + return DataLayoutSeedPhase::Reduce; + } + + VMILayoutAttr getPreferredGroupSlotLoadLayout(VMIGroupSlotLoadOp op) { + auto type = cast(op.getResult().getType()); + int64_t numGroups = op.getNumGroupsAttr().getInt(); + if (VMILayoutAttr existing = type.getLayoutAttr()) + if (existing.isGroupSlots() && existing.getSlots() > 0) + return existing; + std::optional sourceGroupStride = + getConstantIndexValue(op.getSourceGroupStride()); + if (sourceGroupStride && *sourceGroupStride == 1) + return VMILayoutAttr::getGroupSlots(ctx, numGroups, /*slots=*/8); + return VMILayoutAttr::getGroupSlots(ctx, numGroups, /*slots=*/1); + } + + VMILayoutAttr + getPreferredGroupBroadcastLoadLayout(VMIGroupBroadcastLoadOp op) { + auto type = cast(op.getResult().getType()); + if (VMILayoutAttr existing = type.getLayoutAttr()) + return existing; + + VMILayoutSupport supports; + FailureOr fact = + supports.getGroupBroadcastLoadDirectFact( + type, op.getSource().getType(), op.getSourceGroupStride(), + op.getNumGroupsAttr().getInt()); + if (failed(fact)) + return {}; + return fact->layout.resultLayout; + } + + bool hasDirectGroupBroadcastLoadCandidate(VMIGroupBroadcastLoadOp op) { + VMILayoutSupport supports; + return succeeded(supports.getGroupBroadcastLoadDirectFact( + cast(op.getResult().getType()), op.getSource().getType(), + op.getSourceGroupStride(), op.getNumGroupsAttr().getInt())); + } + + VMILayoutAttr getPreferredGroupBroadcastSourceLayout(Value value, + int64_t numGroups) { + auto type = dyn_cast(value.getType()); + if (!type) + return getContiguousLayout(); + if (VMILayoutAttr existing = type.getLayoutAttr()) + if (existing.isGroupSlots() && existing.getSlots() > 0) + return existing; + VMILayoutAttr solved = getDataLayout(value); + if (solved && solved.isGroupSlots() && solved.getNumGroups() == numGroups && + solved.getSlots() > 0) + return solved; + if (type.getElementCount() == numGroups) + return VMILayoutAttr::getGroupSlots(ctx, numGroups, + numGroups >= 8 ? 8 : 1); + if (auto load = value.getDefiningOp()) + return getPreferredGroupSlotLoadLayout(load); + return getPreferredGroupSlotsLayout(type, numGroups); + } + + VMILayoutAttr getPreferredGroupLoadResultLayout(VMIGroupLoadOp op) { + auto type = cast(op.getResult().getType()); + if (VMILayoutAttr existing = type.getLayoutAttr()) + return existing; + + int64_t numGroups = op.getNumGroupsAttr().getInt(); + if (numGroups <= 0 || type.getElementCount() % numGroups != 0) + return getContiguousLayout(); + + if (!type.getElementType().isF32()) + return getContiguousLayout(); + + int64_t groupSize = type.getElementCount() / numGroups; + std::optional rowStride = getConstantIndexValue(op.getRowStride()); + if (rowStride && *rowStride == groupSize) + return getContiguousLayout(); + if (!rowStride || *rowStride <= 0 || *rowStride % 8 != 0) + return getContiguousLayout(); + + if (groupSize == 16) + return VMILayoutAttr::getDeinterleaved(ctx, 2, /*blockElems=*/8); + if (groupSize == 32) + return VMILayoutAttr::getDeinterleaved(ctx, 4, /*blockElems=*/8); + + return getContiguousLayout(); + } + + LogicalResult validateGroupLoadLayoutPlan(VMIGroupLoadOp op) { + auto type = cast(op.getResult().getType()); + if (type.getLayoutAttr()) + return success(); + + int64_t numGroups = op.getNumGroupsAttr().getInt(); + if (numGroups <= 0 || type.getElementCount() % numGroups != 0) + return success(); + if (!type.getElementType().isF32()) + return success(); + + int64_t groupSize = type.getElementCount() / numGroups; + if (groupSize != 16 && groupSize != 32) + return success(); + + std::optional rowStride = getConstantIndexValue(op.getRowStride()); + if (rowStride && *rowStride == groupSize) + return success(); + if (rowStride && *rowStride > 0 && *rowStride % 8 == 0) + return success(); + + return op.emitError() + << kVMIDiagLayoutContractPrefix << "pto.vmi.group_load group_size " + << groupSize + << " requires constant positive row_stride divisible by 8 f32 " + "elements for the block8 stride plan; stable gather fallback is " + "not implemented"; + } + + VMILayoutAttr getPreferredGroupStoreUseLayout( + Value value, int64_t numGroups, Value rowStride, + llvm::function_ref getKnownLayout) { + auto type = dyn_cast(value.getType()); + if (!type) + return getContiguousLayout(); + if (VMILayoutAttr existing = type.getLayoutAttr()) + if (existing.isGroupSlots() && existing.getSlots() > 0) + return existing; + VMILayoutAttr known = getKnownLayout(value); + if (known && known.isGroupSlots() && known.getNumGroups() == numGroups && + known.getSlots() > 0) + return known; + if (known && known.isDeinterleaved() && known.getBlockElems() == 1) { + if (known.getFactor() == 2) + return known; + if (known.getFactor() == 4) + return VMILayoutAttr::getDeinterleaved(ctx, /*factor=*/2, + /*blockElems=*/1); + } + if (auto castOp = value.getDefiningOp()) { + if (isa( + castOp) && + castOp->getNumOperands() == 1 && castOp->getNumResults() == 1) { + auto sourceType = + dyn_cast(castOp->getOperand(0).getType()); + auto resultType = dyn_cast(castOp->getResult(0).getType()); + if (sourceType && resultType) { + VMILayoutAttr sourceLayout = getKnownLayout(castOp->getOperand(0)); + if (!sourceLayout) + sourceLayout = getDataLayout(castOp->getOperand(0)); + VMILayoutSupport supports; + FailureOr fact = + supports.getCastLayoutFactForSourceLayout(sourceType, resultType, + sourceLayout); + if (succeeded(fact) && fact->resultLayout.isGroupSlots() && + fact->resultLayout.getNumGroups() == numGroups && + fact->resultLayout.getSlots() > 0) + return fact->resultLayout; + } + } + } + VMILayoutAttr solved = getDataLayout(value); + if (solved && solved.isGroupSlots() && solved.getNumGroups() == numGroups && + solved.getSlots() > 0) + return solved; + if (value.getDefiningOp() || + value.getDefiningOp() || + value.getDefiningOp() || + value.getDefiningOp() || + value.getDefiningOp() || + value.getDefiningOp()) + return getPreferredGroupSlotsLayout(type, numGroups); + if (type.getElementCount() == numGroups) { + std::optional stride = getConstantIndexValue(rowStride); + bool packedSlots = + stride && *stride == 1 && static_cast(numGroups) >= 8; + return VMILayoutAttr::getGroupSlots(ctx, numGroups, packedSlots ? 8 : 1); + } + if (auto load = value.getDefiningOp()) + return getPreferredGroupSlotLoadLayout(load); + return getContiguousLayout(); + } + + VMILayoutAttr getPreferredGroupStoreUseLayout(Value value, int64_t numGroups, + Value rowStride) { + return getPreferredGroupStoreUseLayout( + value, numGroups, rowStride, + [&](Value knownValue) { return getDataLayout(knownValue); }); + } + + VMILayoutAttr getDataLayout(Value value) { + unsigned id = addDataValue(value); + if (id == ~0u) + return {}; + unsigned root = find(id); + if (dataNodes[root].naturalLayout) + return dataNodes[root].naturalLayout; + if (dataNodes[root].preferredLayout) + return dataNodes[root].preferredLayout; + return getContiguousLayout(); + } + + void requestDataUse(OpOperand &operand, VMILayoutAttr layout, + bool late = false, + DataLayoutSeedPhase phase = DataLayoutSeedPhase::Other) { + if (isa(operand.get().getType())) { + addDataValue(operand.get()); + dataUseRequests.push_back(DataUseRequest{&operand, layout, late, phase}); + } + } + + LogicalResult constrainElementwiseBinary(OpOperand &lhs, OpOperand &rhs, + Value result, Operation *op) { + if (failed(unite(lhs.get(), rhs.get(), op))) + return failure(); + return unite(lhs.get(), result, op); + } + + LogicalResult + requestMaskUse(OpOperand &operand, VMILayoutAttr layout, Operation *op, + DataLayoutSeedPhase phase = DataLayoutSeedPhase::Other) { + if (!isa(operand.get().getType())) + return success(); + if (!layout) + return op->emitError() + << kVMIDiagLayoutContractPrefix + << "cannot infer concrete mask use layout"; + maskUseRequests.push_back(MaskUseRequest{&operand, layout, phase}); + return success(); + } + + LogicalResult collect() { + module.walk([&](Operation *op) { + for (Value result : op->getResults()) { + addDataValue(result); + addMaskValue(result); + } + for (Region ®ion : op->getRegions()) + for (Block &block : region) + for (BlockArgument arg : block.getArguments()) { + addDataValue(arg); + addMaskValue(arg); + } + }); + return success(); + } + + LogicalResult addConstraints() { + WalkResult result = module.walk([&](Operation *op) -> WalkResult { + if (auto maskAnd = dyn_cast(op)) { + if (failed(uniteMask(maskAnd.getLhs(), maskAnd.getRhs(), op)) || + failed(uniteMask(maskAnd.getLhs(), maskAnd.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto maskOr = dyn_cast(op)) { + if (failed(uniteMask(maskOr.getLhs(), maskOr.getRhs(), op)) || + failed(uniteMask(maskOr.getLhs(), maskOr.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto maskXor = dyn_cast(op)) { + if (failed(uniteMask(maskXor.getLhs(), maskXor.getRhs(), op)) || + failed(uniteMask(maskXor.getLhs(), maskXor.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto maskNot = dyn_cast(op)) { + if (failed(uniteMask(maskNot.getSource(), maskNot.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto ensure = dyn_cast(op)) { + if (failed(uniteMask(ensure.getSource(), ensure.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto ensure = dyn_cast(op)) { + return WalkResult::advance(); + } + if (auto addf = dyn_cast(op)) { + if (failed(constrainElementwiseBinary(addf.getLhsMutable(), + addf.getRhsMutable(), + addf.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto addi = dyn_cast(op)) { + if (failed(constrainElementwiseBinary(addi.getLhsMutable(), + addi.getRhsMutable(), + addi.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto subf = dyn_cast(op)) { + if (failed(constrainElementwiseBinary(subf.getLhsMutable(), + subf.getRhsMutable(), + subf.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto subi = dyn_cast(op)) { + if (failed(constrainElementwiseBinary(subi.getLhsMutable(), + subi.getRhsMutable(), + subi.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto mulf = dyn_cast(op)) { + if (failed(constrainElementwiseBinary(mulf.getLhsMutable(), + mulf.getRhsMutable(), + mulf.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto muli = dyn_cast(op)) { + if (failed(constrainElementwiseBinary(muli.getLhsMutable(), + muli.getRhsMutable(), + muli.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto vmull = dyn_cast(op)) { + if (failed(constrainElementwiseBinary(vmull.getAMutable(), + vmull.getBMutable(), + vmull.getLow(), op)) || + failed(unite(vmull.getA(), vmull.getHigh(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto fma = dyn_cast(op)) { + if (failed(unite(fma.getLhs(), fma.getRhs(), op)) || + failed(unite(fma.getLhs(), fma.getAcc(), op)) || + failed(unite(fma.getLhs(), fma.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto divf = dyn_cast(op)) { + if (failed(constrainElementwiseBinary(divf.getLhsMutable(), + divf.getRhsMutable(), + divf.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto minf = dyn_cast(op)) { + if (failed(constrainElementwiseBinary(minf.getLhsMutable(), + minf.getRhsMutable(), + minf.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto maxf = dyn_cast(op)) { + if (failed(constrainElementwiseBinary(maxf.getLhsMutable(), + maxf.getRhsMutable(), + maxf.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto negf = dyn_cast(op)) { + if (failed(unite(negf.getSource(), negf.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto absf = dyn_cast(op)) { + if (failed(unite(absf.getSource(), absf.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto absi = dyn_cast(op)) { + if (failed(unite(absi.getSource(), absi.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto sqrt = dyn_cast(op)) { + if (failed(unite(sqrt.getSource(), sqrt.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto exp = dyn_cast(op)) { + if (failed(unite(exp.getSource(), exp.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto ln = dyn_cast(op)) { + if (failed(unite(ln.getSource(), ln.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto relu = dyn_cast(op)) { + if (failed(unite(relu.getSource(), relu.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto fptosi = dyn_cast(op)) { + if (failed(unite(fptosi.getSource(), fptosi.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto sitofp = dyn_cast(op)) { + if (failed(unite(sitofp.getSource(), sitofp.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto andi = dyn_cast(op)) { + if (failed(constrainElementwiseBinary(andi.getLhsMutable(), + andi.getRhsMutable(), + andi.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto ori = dyn_cast(op)) { + if (failed(constrainElementwiseBinary( + ori.getLhsMutable(), ori.getRhsMutable(), ori.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto xori = dyn_cast(op)) { + if (failed(constrainElementwiseBinary(xori.getLhsMutable(), + xori.getRhsMutable(), + xori.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto shli = dyn_cast(op)) { + if (failed(constrainElementwiseBinary(shli.getLhsMutable(), + shli.getRhsMutable(), + shli.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto shrui = dyn_cast(op)) { + if (failed(constrainElementwiseBinary(shrui.getLhsMutable(), + shrui.getRhsMutable(), + shrui.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto shrsi = dyn_cast(op)) { + if (failed(constrainElementwiseBinary(shrsi.getLhsMutable(), + shrsi.getRhsMutable(), + shrsi.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto notOp = dyn_cast(op)) { + if (failed(unite(notOp.getSource(), notOp.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto cmpf = dyn_cast(op)) { + if (failed(unite(cmpf.getLhs(), cmpf.getRhs(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto cmpi = dyn_cast(op)) { + if (failed(unite(cmpi.getLhs(), cmpi.getRhs(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto select = dyn_cast(op)) { + if (failed(unite(select.getTrueValue(), select.getFalseValue(), op)) || + failed(unite(select.getTrueValue(), select.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto activePrefix = dyn_cast(op)) { + if (failed(setNaturalLayout(activePrefix.getResult(), + getContiguousLayout(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto compress = dyn_cast(op)) { + requestDataUse(compress.getSourceMutable(), getContiguousLayout()); + if (failed(setNaturalLayout(compress.getResult(), getContiguousLayout(), + op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto reduce = dyn_cast(op)) { + requestDataUse(reduce.getSourceMutable(), getContiguousLayout(), + /*late=*/false, DataLayoutSeedPhase::Reduce); + requestDataUse(reduce.getInitMutable(), getContiguousLayout(), + /*late=*/false, DataLayoutSeedPhase::Reduce); + if (failed(requestMaskUse(reduce.getMaskMutable(), + getContiguousLayout(), op))) + return WalkResult::interrupt(); + if (failed(setNaturalLayout(reduce.getResult(), getContiguousLayout(), + op, DataLayoutSeedPhase::Reduce))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto reduce = dyn_cast(op)) { + requestDataUse(reduce.getSourceMutable(), getContiguousLayout(), + /*late=*/false, DataLayoutSeedPhase::Reduce); + requestDataUse(reduce.getInitMutable(), getContiguousLayout(), + /*late=*/false, DataLayoutSeedPhase::Reduce); + if (failed(requestMaskUse(reduce.getMaskMutable(), + getContiguousLayout(), op))) + return WalkResult::interrupt(); + if (failed(setNaturalLayout(reduce.getResult(), getContiguousLayout(), + op, DataLayoutSeedPhase::Reduce))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto reduce = dyn_cast(op)) { + requestDataUse(reduce.getSourceMutable(), getContiguousLayout(), + /*late=*/false, DataLayoutSeedPhase::Reduce); + requestDataUse(reduce.getInitMutable(), getContiguousLayout(), + /*late=*/false, DataLayoutSeedPhase::Reduce); + if (failed(requestMaskUse(reduce.getMaskMutable(), + getContiguousLayout(), op))) + return WalkResult::interrupt(); + if (failed(setNaturalLayout(reduce.getResult(), getContiguousLayout(), + op, DataLayoutSeedPhase::Reduce))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto reduce = dyn_cast(op)) { + requestDataUse(reduce.getSourceMutable(), getContiguousLayout(), + /*late=*/false, DataLayoutSeedPhase::Reduce); + requestDataUse(reduce.getInitMutable(), getContiguousLayout(), + /*late=*/false, DataLayoutSeedPhase::Reduce); + if (failed(requestMaskUse(reduce.getMaskMutable(), + getContiguousLayout(), op))) + return WalkResult::interrupt(); + if (failed(setNaturalLayout(reduce.getResult(), getContiguousLayout(), + op, DataLayoutSeedPhase::Reduce))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto reduce = dyn_cast(op)) { + requestDataUse(reduce.getSourceMutable(), getContiguousLayout(), + /*late=*/false, DataLayoutSeedPhase::Reduce); + requestDataUse(reduce.getInitMutable(), getContiguousLayout(), + /*late=*/false, DataLayoutSeedPhase::Reduce); + if (failed(requestMaskUse(reduce.getMaskMutable(), + getContiguousLayout(), op))) + return WalkResult::interrupt(); + if (failed(setNaturalLayout(reduce.getResult(), getContiguousLayout(), + op, DataLayoutSeedPhase::Reduce))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto reduce = dyn_cast(op)) { + requestDataUse(reduce.getSourceMutable(), getContiguousLayout(), + /*late=*/false, DataLayoutSeedPhase::Reduce); + requestDataUse(reduce.getInitMutable(), getContiguousLayout(), + /*late=*/false, DataLayoutSeedPhase::Reduce); + if (failed(requestMaskUse(reduce.getMaskMutable(), + getContiguousLayout(), op))) + return WalkResult::interrupt(); + if (failed(setNaturalLayout(reduce.getResult(), getContiguousLayout(), + op, DataLayoutSeedPhase::Reduce))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto reduce = dyn_cast(op)) { + auto sourceType = cast(reduce.getSource().getType()); + auto resultType = cast(reduce.getResult().getType()); + int64_t numGroups = reduce.getNumGroupsAttr().getInt(); + VMILayoutSupport supports; + FailureOr fact = + supports.getPreferredGroupReduceLayoutFact(sourceType, numGroups); + VMILayoutAttr sourceLayout = + succeeded(fact) ? fact->sourceLayout : getContiguousLayout(); + DataLayoutSeedPhase usePhase = + succeeded(fact) + ? getGroupReduceUseSeedPhase(sourceType, numGroups, *fact) + : DataLayoutSeedPhase::Reduce; + requestDataUse(reduce.getSourceMutable(), sourceLayout, /*late=*/false, + usePhase); + if (failed(requestMaskUse(reduce.getMaskMutable(), sourceLayout, op, + usePhase))) + return WalkResult::interrupt(); + if (failed(setNaturalLayout( + reduce.getResult(), + succeeded(fact) + ? fact->resultLayout + : getPreferredGroupSlotsLayout(resultType, numGroups), + op, DataLayoutSeedPhase::Reduce))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto reduce = dyn_cast(op)) { + auto sourceType = cast(reduce.getSource().getType()); + auto resultType = cast(reduce.getResult().getType()); + int64_t numGroups = reduce.getNumGroupsAttr().getInt(); + VMILayoutSupport supports; + FailureOr fact = + supports.getPreferredGroupReduceLayoutFact(sourceType, numGroups); + VMILayoutAttr sourceLayout = + succeeded(fact) ? fact->sourceLayout : getContiguousLayout(); + DataLayoutSeedPhase usePhase = + succeeded(fact) + ? getGroupReduceUseSeedPhase(sourceType, numGroups, *fact) + : DataLayoutSeedPhase::Reduce; + requestDataUse(reduce.getSourceMutable(), sourceLayout, /*late=*/false, + usePhase); + if (failed(requestMaskUse(reduce.getMaskMutable(), sourceLayout, op, + usePhase))) + return WalkResult::interrupt(); + if (failed(setNaturalLayout( + reduce.getResult(), + succeeded(fact) + ? fact->resultLayout + : getPreferredGroupSlotsLayout(resultType, numGroups), + op, DataLayoutSeedPhase::Reduce))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto reduce = dyn_cast(op)) { + auto sourceType = cast(reduce.getSource().getType()); + auto resultType = cast(reduce.getResult().getType()); + int64_t numGroups = reduce.getNumGroupsAttr().getInt(); + VMILayoutSupport supports; + FailureOr fact = + supports.getPreferredGroupReduceLayoutFact(sourceType, numGroups); + VMILayoutAttr sourceLayout = + succeeded(fact) ? fact->sourceLayout : getContiguousLayout(); + DataLayoutSeedPhase usePhase = + succeeded(fact) + ? getGroupReduceUseSeedPhase(sourceType, numGroups, *fact) + : DataLayoutSeedPhase::Reduce; + requestDataUse(reduce.getSourceMutable(), sourceLayout, /*late=*/false, + usePhase); + if (failed(requestMaskUse(reduce.getMaskMutable(), sourceLayout, op, + usePhase))) + return WalkResult::interrupt(); + if (failed(setNaturalLayout( + reduce.getResult(), + succeeded(fact) + ? fact->resultLayout + : getPreferredGroupSlotsLayout(resultType, numGroups), + op, DataLayoutSeedPhase::Reduce))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto reduce = dyn_cast(op)) { + auto sourceType = cast(reduce.getSource().getType()); + auto resultType = cast(reduce.getResult().getType()); + int64_t numGroups = reduce.getNumGroupsAttr().getInt(); + VMILayoutSupport supports; + FailureOr fact = + supports.getPreferredGroupReduceLayoutFact(sourceType, numGroups); + VMILayoutAttr sourceLayout = + succeeded(fact) ? fact->sourceLayout : getContiguousLayout(); + DataLayoutSeedPhase usePhase = + succeeded(fact) + ? getGroupReduceUseSeedPhase(sourceType, numGroups, *fact) + : DataLayoutSeedPhase::Reduce; + requestDataUse(reduce.getSourceMutable(), sourceLayout, /*late=*/false, + usePhase); + if (failed(requestMaskUse(reduce.getMaskMutable(), sourceLayout, op, + usePhase))) + return WalkResult::interrupt(); + if (failed(setNaturalLayout( + reduce.getResult(), + succeeded(fact) + ? fact->resultLayout + : getPreferredGroupSlotsLayout(resultType, numGroups), + op, DataLayoutSeedPhase::Reduce))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto reduce = dyn_cast(op)) { + auto sourceType = cast(reduce.getSource().getType()); + auto resultType = cast(reduce.getResult().getType()); + int64_t numGroups = reduce.getNumGroupsAttr().getInt(); + VMILayoutSupport supports; + FailureOr fact = + supports.getPreferredGroupReduceLayoutFact(sourceType, numGroups); + VMILayoutAttr sourceLayout = + succeeded(fact) ? fact->sourceLayout : getContiguousLayout(); + DataLayoutSeedPhase usePhase = + succeeded(fact) + ? getGroupReduceUseSeedPhase(sourceType, numGroups, *fact) + : DataLayoutSeedPhase::Reduce; + requestDataUse(reduce.getSourceMutable(), sourceLayout, /*late=*/false, + usePhase); + if (failed(requestMaskUse(reduce.getMaskMutable(), sourceLayout, op, + usePhase))) + return WalkResult::interrupt(); + if (failed(setNaturalLayout( + reduce.getResult(), + succeeded(fact) + ? fact->resultLayout + : getPreferredGroupSlotsLayout(resultType, numGroups), + op, DataLayoutSeedPhase::Reduce))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto reduce = dyn_cast(op)) { + auto sourceType = cast(reduce.getSource().getType()); + auto resultType = cast(reduce.getResult().getType()); + int64_t numGroups = reduce.getNumGroupsAttr().getInt(); + VMILayoutSupport supports; + FailureOr fact = + supports.getPreferredGroupReduceLayoutFact(sourceType, numGroups); + VMILayoutAttr sourceLayout = + succeeded(fact) ? fact->sourceLayout : getContiguousLayout(); + DataLayoutSeedPhase usePhase = + succeeded(fact) + ? getGroupReduceUseSeedPhase(sourceType, numGroups, *fact) + : DataLayoutSeedPhase::Reduce; + requestDataUse(reduce.getSourceMutable(), sourceLayout, /*late=*/false, + usePhase); + if (failed(requestMaskUse(reduce.getMaskMutable(), sourceLayout, op, + usePhase))) + return WalkResult::interrupt(); + if (failed(setNaturalLayout( + reduce.getResult(), + succeeded(fact) + ? fact->resultLayout + : getPreferredGroupSlotsLayout(resultType, numGroups), + op, DataLayoutSeedPhase::Reduce))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto broadcast = dyn_cast(op)) { + requestDataUse( + broadcast.getSourceMutable(), + getPreferredGroupBroadcastSourceLayout( + broadcast.getSource(), broadcast.getNumGroupsAttr().getInt()), + /*late=*/false, DataLayoutSeedPhase::GroupBroadcast); + return WalkResult::advance(); + } + if (auto hist = dyn_cast(op)) { + requestDataUse(hist.getAccMutable(), getContiguousLayout()); + requestDataUse(hist.getSourceMutable(), getContiguousLayout()); + if (failed(requestMaskUse(hist.getMaskMutable(), getContiguousLayout(), + op))) + return WalkResult::interrupt(); + if (failed( + setNaturalLayout(hist.getResult(), getContiguousLayout(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto hist = dyn_cast(op)) { + requestDataUse(hist.getAccMutable(), getContiguousLayout()); + requestDataUse(hist.getSourceMutable(), getContiguousLayout()); + if (failed(requestMaskUse(hist.getMaskMutable(), getContiguousLayout(), + op))) + return WalkResult::interrupt(); + if (failed( + setNaturalLayout(hist.getResult(), getContiguousLayout(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto extf = dyn_cast(op)) { + auto sourceType = cast(extf.getSource().getType()); + auto resultType = cast(extf.getResult().getType()); + VMILayoutSupport supports; + FailureOr fact = + supports.getPreferredCastLayoutFact(sourceType, resultType); + if (succeeded(fact)) { + if (failed(setPreferredLayout(extf.getResult(), fact->resultLayout, + op, DataLayoutSeedPhase::Cast))) + return WalkResult::interrupt(); + } + return WalkResult::advance(); + } + if (auto extsi = dyn_cast(op)) { + auto sourceType = cast(extsi.getSource().getType()); + auto resultType = cast(extsi.getResult().getType()); + VMILayoutSupport supports; + FailureOr fact = + supports.getPreferredCastLayoutFact(sourceType, resultType); + if (succeeded(fact)) { + if (failed(setPreferredLayout(extsi.getResult(), fact->resultLayout, + op, DataLayoutSeedPhase::Cast))) + return WalkResult::interrupt(); + } + return WalkResult::advance(); + } + if (auto extui = dyn_cast(op)) { + auto sourceType = cast(extui.getSource().getType()); + auto resultType = cast(extui.getResult().getType()); + VMILayoutSupport supports; + FailureOr fact = + supports.getPreferredCastLayoutFact(sourceType, resultType); + if (succeeded(fact)) { + if (failed(setPreferredLayout(extui.getResult(), fact->resultLayout, + op, DataLayoutSeedPhase::Cast))) + return WalkResult::interrupt(); + } + return WalkResult::advance(); + } + if (auto truncf = dyn_cast(op)) { + auto sourceType = cast(truncf.getSource().getType()); + auto resultType = cast(truncf.getResult().getType()); + VMILayoutSupport supports; + FailureOr fact = + supports.getPreferredCastLayoutFact(sourceType, resultType); + VMILayoutAttr resultLayout = getContiguousLayout(); + if (succeeded(fact)) { + resultLayout = fact->resultLayout; + } + if (failed(setPreferredLayout(truncf.getResult(), resultLayout, op, + DataLayoutSeedPhase::Cast))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto trunci = dyn_cast(op)) { + auto sourceType = cast(trunci.getSource().getType()); + auto resultType = cast(trunci.getResult().getType()); + VMILayoutSupport supports; + FailureOr fact = + supports.getPreferredCastLayoutFact(sourceType, resultType); + VMILayoutAttr resultLayout = getContiguousLayout(); + if (succeeded(fact)) { + resultLayout = fact->resultLayout; + } + if (failed(setPreferredLayout(trunci.getResult(), resultLayout, op, + DataLayoutSeedPhase::Cast))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto bitcast = dyn_cast(op)) { + if (failed(unite(bitcast.getSource(), bitcast.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto vintlv = dyn_cast(op)) { + VMILayoutSupport supports; + auto valueType = cast(vintlv.getLow().getType()); + FailureOr fact = + supports.getPreferredVintlvLayoutFact(valueType); + if (failed(fact)) + return WalkResult::advance(); + requestDataUse(vintlv.getLhsMutable(), fact->lhsLayout); + requestDataUse(vintlv.getRhsMutable(), fact->rhsLayout); + if (failed(requestMaskUse(vintlv.getMaskMutable(), fact->maskLayout, + op)) || + failed(setPreferredLayout(vintlv.getLow(), fact->lowLayout, op)) || + failed(setPreferredLayout(vintlv.getHigh(), fact->highLayout, op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto vdintlv = dyn_cast(op)) { + VMILayoutSupport supports; + auto valueType = cast(vdintlv.getLow().getType()); + FailureOr fact = + supports.getPreferredVdintlvLayoutFact(valueType); + if (failed(fact)) + return WalkResult::advance(); + requestDataUse(vdintlv.getLhsMutable(), fact->lhsLayout); + requestDataUse(vdintlv.getRhsMutable(), fact->rhsLayout); + if (failed(requestMaskUse(vdintlv.getMaskMutable(), fact->maskLayout, + op)) || + failed(setPreferredLayout(vdintlv.getLow(), fact->lowLayout, op)) || + failed(setPreferredLayout(vdintlv.getHigh(), fact->highLayout, op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto load = dyn_cast(op)) { + VMILayoutSupport supports; + FailureOr fact = + supports.getPreferredDeinterleaveLoadLayoutFact( + cast(load.getLow().getType())); + if (failed(fact)) + return WalkResult::advance(); + if (failed(setNaturalLayout(load.getLow(), fact->lowLayout, op)) || + failed(setNaturalLayout(load.getHigh(), fact->highLayout, op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto load = dyn_cast(op)) { + requestDataUse(load.getPassthruMutable(), getContiguousLayout()); + if (failed( + setNaturalLayout(load.getResult(), getContiguousLayout(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto gather = dyn_cast(op)) { + requestDataUse(gather.getIndicesMutable(), getContiguousLayout()); + requestDataUse(gather.getPassthruMutable(), getContiguousLayout()); + if (failed(requestMaskUse(gather.getMaskMutable(), + getContiguousLayout(), op))) + return WalkResult::interrupt(); + if (failed(setNaturalLayout(gather.getResult(), getContiguousLayout(), + op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto load = dyn_cast(op)) { + requestDataUse(load.getPassthruMutable(), getContiguousLayout()); + if (failed( + setNaturalLayout(load.getResult(), getContiguousLayout(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto load = dyn_cast(op)) { + if (failed(validateGroupLoadLayoutPlan(load))) + return WalkResult::interrupt(); + VMILayoutAttr layout = getPreferredGroupLoadResultLayout(load); + if (layout.isContiguous() && layout.getLaneStride() == 1) { + if (failed(setPreferredLayout(load.getResult(), layout, op))) + return WalkResult::interrupt(); + } else if (failed(setNaturalLayout(load.getResult(), layout, op, + DataLayoutSeedPhase::GroupLoad))) { + return WalkResult::interrupt(); + } + return WalkResult::advance(); + } + if (auto load = dyn_cast(op)) { + if (failed(setPreferredLayout( + load.getResult(), getPreferredGroupSlotLoadLayout(load), op, + DataLayoutSeedPhase::GroupSlotLoad))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto load = dyn_cast(op)) { + DataLayoutSeedPhase phase = + hasDirectGroupBroadcastLoadCandidate(load) + ? DataLayoutSeedPhase::GroupBroadcastLoad + : DataLayoutSeedPhase::Other; + if (failed(setNaturalLayout(load.getResult(), + getPreferredGroupBroadcastLoadLayout(load), + op, phase))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto load = dyn_cast(op)) { + if (failed( + setNaturalLayout(load.getResult(), getContiguousLayout(), op))) + return WalkResult::interrupt(); + if (failed(requestMaskUse(load.getMaskMutable(), getContiguousLayout(), + op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto store = dyn_cast(op)) { + auto valueType = cast(store.getValue().getType()); + if (!hasDataLayoutSeed(store.getValue())) + if (VMILayoutAttr layout = getPreferredDenseStoreLayout(valueType)) + requestDataUse(store.getValueMutable(), layout, /*late=*/false, + DataLayoutSeedPhase::Store); + return WalkResult::advance(); + } + if (auto store = dyn_cast(op)) { + requestDataUse(store.getLowMutable(), getContiguousLayout()); + requestDataUse(store.getHighMutable(), getContiguousLayout()); + return WalkResult::advance(); + } + if (auto store = dyn_cast(op)) { + addDataValue(store.getValue()); + groupStoreUseRequests.push_back(GroupStoreUseRequest{store}); + return WalkResult::advance(); + } + if (auto store = dyn_cast(op)) { + auto valueType = cast(store.getValue().getType()); + auto maskType = cast(store.getMask().getType()); + if (!hasDataLayoutSeed(store.getValue())) { + FailureOr fact = + getPreferredDenseMaskedStoreLayout(valueType, maskType); + if (succeeded(fact)) { + requestDataUse(store.getValueMutable(), fact->valueLayout, + /*late=*/false, + DataLayoutSeedPhase::Store); + if (failed(requestMaskUse(store.getMaskMutable(), fact->maskLayout, + op, DataLayoutSeedPhase::Store))) + return WalkResult::interrupt(); + } + } + return WalkResult::advance(); + } + if (auto store = dyn_cast(op)) { + requestDataUse(store.getValueMutable(), getContiguousLayout()); + if (failed(requestMaskUse(store.getMaskMutable(), getContiguousLayout(), + op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto scatter = dyn_cast(op)) { + requestDataUse(scatter.getValueMutable(), getContiguousLayout()); + requestDataUse(scatter.getIndicesMutable(), getContiguousLayout()); + if (failed(requestMaskUse(scatter.getMaskMutable(), + getContiguousLayout(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto store = dyn_cast(op)) { + requestDataUse(store.getValueMutable(), getContiguousLayout()); + if (failed(requestMaskUse(store.getMaskMutable(), getContiguousLayout(), + op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto split = dyn_cast(op)) { + int64_t channels = split.getNumResults(); + if (channels != 2 && channels != 4) { + split.emitError() << kVMIDiagUnsupportedPrefix + << "pto.vmi.channel_split supports only 2 or 4 " + "channels"; + return WalkResult::interrupt(); + } + requestDataUse(split.getSourceMutable(), + VMILayoutAttr::getDeinterleaved(ctx, channels)); + for (Value result : split.getResults()) + if (failed(setNaturalLayout(result, getContiguousLayout(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto merge = dyn_cast(op)) { + int64_t channels = merge.getInputs().size(); + if (channels != 2 && channels != 4) { + merge.emitError() << kVMIDiagUnsupportedPrefix + << "pto.vmi.channel_merge supports only 2 or 4 " + "channels"; + return WalkResult::interrupt(); + } + for (OpOperand &input : merge.getInputsMutable()) + requestDataUse(input, getContiguousLayout()); + if (failed(setNaturalLayout( + merge.getResult(), + VMILayoutAttr::getDeinterleaved(ctx, channels), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto shuffle = dyn_cast(op)) { + auto sourceType = cast(shuffle.getSource().getType()); + auto resultType = cast(shuffle.getResult().getType()); + if (sourceType.hasLayout() || resultType.hasLayout()) + return WalkResult::advance(); + + requestDataUse(shuffle.getSourceMutable(), getContiguousLayout()); + if (isLane0SplatShuffle(shuffle)) + return WalkResult::advance(); + if (failed(setNaturalLayout(shuffle.getResult(), getContiguousLayout(), + op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto ifOp = dyn_cast(op)) { + if (failed(addIfConstraints(ifOp))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto executeRegionOp = dyn_cast(op)) { + if (failed(addExecuteRegionConstraints(executeRegionOp))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto indexSwitchOp = dyn_cast(op)) { + if (failed(addIndexSwitchConstraints(indexSwitchOp))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto whileOp = dyn_cast(op)) { + if (failed(addWhileConstraints(whileOp))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto forOp = dyn_cast(op)) { + if (failed(addForConstraints(forOp))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto branchOp = dyn_cast(op)) { + if (failed(addBranchConstraints(branchOp.getDest(), + branchOp.getDestOperands(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto condBranchOp = dyn_cast(op)) { + if (failed(addBranchConstraints(condBranchOp.getTrueDest(), + condBranchOp.getTrueDestOperands(), + op)) || + failed(addBranchConstraints(condBranchOp.getFalseDest(), + condBranchOp.getFalseDestOperands(), + op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto switchOp = dyn_cast(op)) { + if (failed(addBranchConstraints(switchOp.getDefaultDestination(), + switchOp.getDefaultOperands(), op))) + return WalkResult::interrupt(); + for (auto [dest, operands] : llvm::zip(switchOp.getCaseDestinations(), + switchOp.getCaseOperands())) { + if (failed(addBranchConstraints(dest, operands, op))) + return WalkResult::interrupt(); + } + return WalkResult::advance(); + } + if (auto returnOp = dyn_cast(op)) { + if (failed(addReturnConstraints(returnOp))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto callOp = dyn_cast(op)) { + if (failed(addCallConstraints(callOp))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (op->getName().getStringRef() == "func.call_indirect") { + if (hasVMIValueTypes(op)) { + op->emitError() + << kVMIDiagLayoutContractPrefix + << "VMI typed call requires a direct internal callee with a body"; + return WalkResult::interrupt(); + } + return WalkResult::advance(); + } + if (auto funcOp = dyn_cast(op)) { + if (funcOp.empty() && hasVMIFunctionType(funcOp)) { + funcOp.emitError() + << kVMIDiagLayoutContractPrefix + << "VMI typed function declaration requires an explicit " + "external ABI materialization plan"; + return WalkResult::interrupt(); + } + return WalkResult::advance(); + } + return WalkResult::advance(); + }); + return failure(result.wasInterrupted()); + } + + LogicalResult uniteEquivalentValues(Value lhs, Value rhs, Operation *op) { + if (failed(uniteDataEquivalent(lhs, rhs, op))) + return failure(); + return uniteMask(lhs, rhs, op); + } + + LogicalResult addIfConstraints(scf::IfOp ifOp) { + for (OpResult result : ifOp->getResults()) { + unsigned resultNo = result.getResultNumber(); + for (Region *region : {&ifOp.getThenRegion(), &ifOp.getElseRegion()}) { + if (region->empty()) + continue; + auto yieldOp = dyn_cast(region->front().getTerminator()); + if (!yieldOp || resultNo >= yieldOp.getNumOperands()) + continue; + if (failed(uniteEquivalentValues(result, yieldOp.getOperand(resultNo), + ifOp))) + return failure(); + } + } + return success(); + } + + LogicalResult addYieldConstraints(ResultRange results, scf::YieldOp yieldOp, + Operation *op) { + for (auto [index, result] : llvm::enumerate(results)) { + if (index >= yieldOp.getNumOperands()) + break; + if (failed(uniteEquivalentValues(result, yieldOp.getOperand(index), op))) + return failure(); + } + return success(); + } + + LogicalResult addExecuteRegionConstraints(scf::ExecuteRegionOp executeOp) { + WalkResult result = executeOp.getRegion().walk([&](scf::YieldOp yieldOp) { + if (yieldOp->getParentOp() != executeOp.getOperation()) + return WalkResult::advance(); + if (failed( + addYieldConstraints(executeOp->getResults(), yieldOp, executeOp))) + return WalkResult::interrupt(); + return WalkResult::advance(); + }); + return failure(result.wasInterrupted()); + } + + LogicalResult addIndexSwitchConstraints(scf::IndexSwitchOp indexSwitchOp) { + auto addBlockTerminator = [&](Block &block) -> LogicalResult { + auto yieldOp = dyn_cast(block.getTerminator()); + if (!yieldOp) + return success(); + return addYieldConstraints(indexSwitchOp->getResults(), yieldOp, + indexSwitchOp); + }; + + if (failed(addBlockTerminator(indexSwitchOp.getDefaultBlock()))) + return failure(); + for (unsigned idx = 0, e = indexSwitchOp.getNumCases(); idx < e; ++idx) + if (failed(addBlockTerminator(indexSwitchOp.getCaseBlock(idx)))) + return failure(); + return success(); + } + + LogicalResult addWhileConstraints(scf::WhileOp whileOp) { + auto inits = whileOp.getInits(); + auto beforeArgs = whileOp.getBeforeArguments(); + Block &afterBlock = whileOp.getAfter().front(); + auto conditionOp = + dyn_cast(whileOp.getBefore().front().getTerminator()); + auto yieldOp = dyn_cast(afterBlock.getTerminator()); + + for (auto [index, init] : llvm::enumerate(inits)) { + Value anchor = init; + if (index < beforeArgs.size() && + failed(uniteEquivalentValues(anchor, beforeArgs[index], whileOp))) + return failure(); + if (conditionOp && index < conditionOp.getArgs().size() && + failed(uniteEquivalentValues(anchor, conditionOp.getArgs()[index], + whileOp))) + return failure(); + if (index < afterBlock.getNumArguments() && + failed(uniteEquivalentValues(anchor, afterBlock.getArgument(index), + whileOp))) + return failure(); + if (yieldOp && index < yieldOp.getNumOperands() && + failed(uniteEquivalentValues(anchor, yieldOp.getOperand(index), + whileOp))) + return failure(); + if (index < whileOp.getNumResults() && + failed( + uniteEquivalentValues(anchor, whileOp.getResult(index), whileOp))) + return failure(); + } + return success(); + } + + LogicalResult addForConstraints(scf::ForOp forOp) { + auto initArgs = forOp.getInitArgs(); + auto regionIterArgs = forOp.getRegionIterArgs(); + auto results = forOp.getResults(); + scf::YieldOp yieldOp = nullptr; + if (Block *body = forOp.getBody()) + yieldOp = dyn_cast(body->getTerminator()); + + for (auto [index, initArg] : llvm::enumerate(initArgs)) { + Value anchor = initArg; + if (index < regionIterArgs.size() && + failed(uniteEquivalentValues(anchor, regionIterArgs[index], forOp))) + return failure(); + if (index < results.size() && + failed(uniteEquivalentValues(anchor, results[index], forOp))) + return failure(); + if (yieldOp && index < yieldOp.getNumOperands() && + failed( + uniteEquivalentValues(anchor, yieldOp.getOperand(index), forOp))) + return failure(); + } + return success(); + } + + LogicalResult addBranchConstraints(Block *dest, OperandRange operands, + Operation *op) { + if (!dest) + return success(); + for (auto [index, operand] : llvm::enumerate(operands)) { + if (index >= dest->getNumArguments()) + break; + if (failed(uniteEquivalentValues(operand, dest->getArgument(index), op))) + return failure(); + } + return success(); + } + + LogicalResult addReturnConstraints(func::ReturnOp returnOp) { + auto func = returnOp->getParentOfType(); + if (!func) + return success(); + + auto it = firstReturnOperandsByFunc.find(func); + if (it == firstReturnOperandsByFunc.end()) { + SmallVector operands(returnOp.getOperands()); + firstReturnOperandsByFunc.try_emplace(func, std::move(operands)); + return success(); + } + + ArrayRef firstOperands = it->second; + for (auto [index, operand] : llvm::enumerate(returnOp.getOperands())) { + if (index >= firstOperands.size()) + break; + if (failed( + uniteEquivalentValues(firstOperands[index], operand, returnOp))) + return failure(); + } + return success(); + } + + bool hasVMIValueTypes(Operation *op) { + return llvm::any_of(op->getOperandTypes(), containsVMIType) || + llvm::any_of(op->getResultTypes(), containsVMIType); + } + + bool hasVMIFunctionType(func::FuncOp func) { + FunctionType type = func.getFunctionType(); + return llvm::any_of(type.getInputs(), containsVMIType) || + llvm::any_of(type.getResults(), containsVMIType); + } + + LogicalResult addCallConstraints(func::CallOp callOp) { + if (!hasVMIValueTypes(callOp)) + return success(); + + auto callee = SymbolTable::lookupNearestSymbolFrom( + callOp, callOp.getCalleeAttr()); + if (!callee || callee.empty()) + return callOp.emitError() + << kVMIDiagLayoutContractPrefix + << "VMI typed call requires a direct internal callee with a body"; + + for (auto [operand, argument] : + llvm::zip(callOp.getOperands(), callee.getArguments())) { + if (failed(uniteEquivalentValues(operand, argument, callOp))) + return failure(); + } + + SmallVector returns; + callee.walk([&](func::ReturnOp returnOp) { returns.push_back(returnOp); }); + for (func::ReturnOp returnOp : returns) { + for (auto [index, result] : llvm::enumerate(callOp.getResults())) { + if (index >= returnOp.getNumOperands()) + break; + if (failed(uniteEquivalentValues(result, returnOp.getOperand(index), + callOp))) + return failure(); + } + } + return success(); + } + + void rewriteDataTypes() { + for (DataNode &node : dataNodes) { + VMILayoutAttr layout = getDataLayout(node.value); + node.value.setType(VMIVRegType::get(ctx, node.type.getElementCount(), + node.type.getElementType(), layout)); + } + } + + FailureOr materializeLayoutValue(Value value, Type targetType, + Location loc, OpBuilder &builder) { + if (value.getType() == targetType) + return value; + + if (auto sourceType = dyn_cast(value.getType())) { + auto targetVRegType = dyn_cast(targetType); + if (!targetVRegType || + sourceType.getElementCount() != targetVRegType.getElementCount() || + sourceType.getElementType() != targetVRegType.getElementType()) + return failure(); + return builder.create(loc, targetVRegType, value) + .getResult(); + } + + if (auto sourceType = dyn_cast(value.getType())) { + auto targetMaskType = dyn_cast(targetType); + if (!targetMaskType || + sourceType.getElementCount() != targetMaskType.getElementCount() || + sourceType.getGranularity() != targetMaskType.getGranularity()) + return failure(); + return builder + .create(loc, targetMaskType, value) + .getResult(); + } + + return failure(); + } + + SmallVector getCallResultTypes(func::FuncOp func) { + SmallVector resultTypes; + bool found = false; + module.walk([&](func::CallOp call) { + if (call.getCallee() != func.getSymName()) + return; + if (!found) { + resultTypes.assign(call.getResultTypes().begin(), + call.getResultTypes().end()); + found = true; + return; + } + if (resultTypes.size() != call.getNumResults()) + return; + for (auto [index, type] : llvm::enumerate(call.getResultTypes())) + if (index < resultTypes.size() && resultTypes[index] != type) + resultTypes[index] = {}; + }); + return found ? resultTypes : SmallVector{}; + } + + LogicalResult materializeCallBoundaries() { + IRRewriter rewriter(ctx); + + WalkResult callResult = module.walk([&](func::CallOp call) -> WalkResult { + auto callee = SymbolTable::lookupNearestSymbolFrom( + call, call.getCalleeAttr()); + if (!callee || callee.empty()) + return WalkResult::advance(); + + rewriter.setInsertionPoint(call); + for (auto [index, operand] : llvm::enumerate(call.getOperands())) { + if (index >= callee.getNumArguments()) + break; + Type targetType = callee.getArgument(index).getType(); + if (!isa(targetType)) + continue; + FailureOr materialized = + materializeLayoutValue(operand, targetType, call.getLoc(), + rewriter); + if (failed(materialized)) + return WalkResult::interrupt(); + call->setOperand(index, *materialized); + } + return WalkResult::advance(); + }); + if (callResult.wasInterrupted()) + return failure(); + + WalkResult returnResult = module.walk([&](func::FuncOp func) -> WalkResult { + SmallVector resultTypes = getCallResultTypes(func); + if (resultTypes.empty()) + return WalkResult::advance(); + + WalkResult nested = func.walk([&](func::ReturnOp ret) -> WalkResult { + rewriter.setInsertionPoint(ret); + for (auto [index, operand] : llvm::enumerate(ret.getOperands())) { + if (index >= resultTypes.size()) + break; + if (!resultTypes[index]) + continue; + Type targetType = resultTypes[index]; + if (!isa(targetType)) + continue; + FailureOr materialized = + materializeLayoutValue(operand, targetType, ret.getLoc(), + rewriter); + if (failed(materialized)) + return WalkResult::interrupt(); + ret->setOperand(index, *materialized); + } + return WalkResult::advance(); + }); + return nested.wasInterrupted() ? WalkResult::interrupt() + : WalkResult::advance(); + }); + return failure(returnResult.wasInterrupted()); + } + + LogicalResult insertDataUseMaterializations() { + OpBuilder builder(ctx); + for (DataUseRequest request : dataUseRequests) { + Value value = request.operand->get(); + auto sourceType = dyn_cast(value.getType()); + if (!sourceType) + continue; + VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); + if (!sourceLayout) + return request.operand->getOwner()->emitError() + << kVMIDiagLayoutContractPrefix + << "data use materialization requires layout-assigned source " + "type"; + if (sourceLayout == request.layout) + continue; + + auto resultType = + VMIVRegType::get(ctx, sourceType.getElementCount(), + sourceType.getElementType(), request.layout); + builder.setInsertionPoint(request.operand->getOwner()); + auto ensure = builder.create( + request.operand->getOwner()->getLoc(), resultType, value); + request.operand->set(ensure.getResult()); + } + return success(); + } + + bool hasRequestedLayout(VMILayoutPropagator &propagator, Value value) { + return static_cast(propagator.getRequestedLayout(value)); + } + + bool hasLayoutAssignment(VMILayoutPropagator &propagator, Value value) { + return propagator.lookup(value) != nullptr; + } + + LogicalResult requestDataLayoutSeeds(VMILayoutPropagator &propagator, + DataLayoutSeedPhase phase, + bool skipAlreadyRequested) { + SmallVector protectedValues; + if (skipAlreadyRequested) { + for (DataLayoutSeed seed : dataLayoutSeeds) { + if (seed.phase != phase) + continue; + if (hasRequestedLayout(propagator, seed.value) && + !llvm::is_contained(protectedValues, seed.value)) + protectedValues.push_back(seed.value); + } + } + + for (DataLayoutSeed seed : dataLayoutSeeds) { + if (seed.phase != phase) + continue; + if (llvm::is_contained(protectedValues, seed.value)) + continue; + if (failed(propagator.request(seed.value, seed.layout))) + return failure(); + } + return success(); + } + + LogicalResult requestDataUseSeeds(VMILayoutPropagator &propagator, + DataLayoutSeedPhase phase, bool late) { + for (DataUseRequest request : dataUseRequests) + if (request.phase == phase && request.late == late) { + if (hasLayoutAssignment(propagator, request.operand->get())) { + VMILayoutAttr assigned = + propagator.getRequestedOrCurrentLayout(request.operand->get()); + if (propagator.canUseOperandLayout(*request.operand, assigned)) + continue; + } + if (failed(propagator.request(*request.operand, request.layout))) + return failure(); + } + return success(); + } + + LogicalResult requestMaskUseSeeds(VMILayoutPropagator &propagator, + DataLayoutSeedPhase phase) { + for (MaskUseRequest request : maskUseRequests) + if (request.phase == phase) { + if (hasLayoutAssignment(propagator, request.operand->get())) { + VMILayoutAttr assigned = + propagator.getRequestedOrCurrentLayout(request.operand->get()); + if (propagator.canUseOperandLayout(*request.operand, assigned)) + continue; + } + if (failed(propagator.request(*request.operand, request.layout))) + return failure(); + } + return success(); + } + + LogicalResult runSeedPhase(VMILayoutPropagator &propagator, + DataLayoutSeedPhase phase) { + if (failed(requestDataLayoutSeeds(propagator, phase, + /*skipAlreadyRequested=*/true))) + return failure(); + if (failed(requestDataUseSeeds(propagator, phase, /*late=*/false))) + return failure(); + if (failed(requestMaskUseSeeds(propagator, phase))) + return failure(); + return propagator.run(); + } + + LogicalResult applyLayouts() { + VMILayoutPropagator propagator(module); + for (DataNode &node : dataNodes) { + DataNode &root = dataNodes[find(dataIds.lookup(node.value))]; + propagator.addEquivalentValues(root.value, node.value); + } + for (MaskNode &node : maskNodes) { + MaskNode &root = maskNodes[findMask(maskIds.lookup(node.value))]; + propagator.addEquivalentValues(root.value, node.value); + } + if (failed(requestDataLayoutSeeds(propagator, DataLayoutSeedPhase::Explicit, + /*skipAlreadyRequested=*/false))) + return failure(); + for (MaskNode &node : maskNodes) { + MaskNode &root = maskNodes[findMask(maskIds.lookup(node.value))]; + if (root.requestedLayout && + failed(propagator.request(node.value, root.requestedLayout))) + return failure(); + } + if (failed(propagator.run())) + return failure(); + + for (int64_t phase = static_cast(DataLayoutSeedPhase::SeedStart); + phase < static_cast(DataLayoutSeedPhase::SeedEnd); ++phase) + if (failed(runSeedPhase(propagator, + static_cast(phase)))) + return failure(); + + for (GroupStoreUseRequest request : groupStoreUseRequests) { + VMIGroupStoreOp store = request.store; + VMILayoutAttr layout = getPreferredGroupStoreUseLayout( + store.getValue(), store.getNumGroupsAttr().getInt(), + store.getRowStride(), [&](Value value) { + return propagator.getRequestedOrCurrentLayout(value); + }); + if (failed(propagator.request(store.getValueMutable(), layout))) + return failure(); + } + if (failed(propagator.run())) + return failure(); + + for (DataUseRequest request : dataUseRequests) + if (request.late && + failed(propagator.request(*request.operand, request.layout))) + return failure(); + if (failed(propagator.run())) + return failure(); + + for (DataNode &node : dataNodes) + if (!propagator.getRequestedLayout(node.value) && + failed(propagator.request(node.value, getContiguousLayout()))) + return failure(); + for (MaskNode &node : maskNodes) + if (!propagator.getRequestedLayout(node.value) && + failed(propagator.request(node.value, getContiguousLayout()))) + return failure(); + if (failed(propagator.run())) + return failure(); + + IRRewriter rewriter(ctx); + return propagator.apply(rewriter); + } + + void rewriteFunctionType() { + module.walk([&](func::FuncOp func) { + if (func.empty()) + return; + + SmallVector inputs; + inputs.reserve(func.getNumArguments()); + for (BlockArgument arg : func.getArguments()) + inputs.push_back(arg.getType()); + + SmallVector results; + auto it = firstReturnOperandsByFunc.find(func); + SmallVector callResultTypes = getCallResultTypes(func); + if (!callResultTypes.empty()) { + for (Type type : callResultTypes) + results.push_back(type); + } else if (it != firstReturnOperandsByFunc.end()) { + for (Value operand : it->second) + results.push_back(operand.getType()); + } else { + FunctionType functionType = func.getFunctionType(); + for (Type type : functionType.getResults()) { + if (auto vregType = dyn_cast(type)) { + results.push_back(VMIVRegType::get(ctx, vregType.getElementCount(), + vregType.getElementType(), + getContiguousLayout())); + } else if (auto maskType = dyn_cast(type)) { + results.push_back(VMIMaskType::get(ctx, maskType.getElementCount(), + "b32", getContiguousLayout())); + } else { + results.push_back(type); + } + } + } + + func.setFunctionType(FunctionType::get(ctx, inputs, results)); + }); + } + + LogicalResult run() { + if (failed(collect())) + return failure(); + if (failed(addConstraints())) + return failure(); + if (failed(applyLayouts())) + return failure(); + if (failed(materializeCallBoundaries())) + return failure(); + rewriteFunctionType(); + return validateVMILayoutAssignedIR(module, /*diagOS=*/nullptr, + /*verifyHelperSupport=*/false); + } + + ModuleOp module; + MLIRContext *ctx; + DenseMap dataIds; + DenseMap maskIds; + DenseMap> firstReturnOperandsByFunc; + SmallVector dataNodes; + SmallVector maskNodes; + SmallVector dataLayoutSeeds; + SmallVector dataUseRequests; + SmallVector groupStoreUseRequests; + SmallVector maskUseRequests; +}; + +struct VMILayoutAssignmentPass + : public mlir::pto::impl::VMILayoutAssignmentBase { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(VMILayoutAssignmentPass) + + void runOnOperation() override { + if (failed(LayoutSolver(getOperation()).run())) + signalPassFailure(); + } +}; + +} // namespace + +std::unique_ptr mlir::pto::createVMILayoutAssignmentPass() { + return std::make_unique(); +} diff --git a/lib/PTO/Transforms/VMILayoutFold.cpp b/lib/PTO/Transforms/VMILayoutFold.cpp new file mode 100644 index 0000000000..e00dcdce9b --- /dev/null +++ b/lib/PTO/Transforms/VMILayoutFold.cpp @@ -0,0 +1,230 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +//===- VMILayoutFold.cpp - Fold VMI layout materializations --------------===// +//===----------------------------------------------------------------------===// + +#include "PTO/IR/PTO.h" +#include "PTO/IR/PTOTypeUtils.h" +#include "PTO/IR/VMIUtils.h" +#include "PTO/Transforms/Passes.h" +#include "PTO/Transforms/VMILayoutSupport.h" + +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/IRMapping.h" +#include "mlir/IR/Operation.h" +#include "mlir/Pass/Pass.h" +#include "llvm/ADT/STLExtras.h" + +namespace mlir { +namespace pto { +#define GEN_PASS_DEF_VMILAYOUTFOLD +#include "PTO/Transforms/Passes.h.inc" +} // namespace pto +} // namespace mlir + +using namespace mlir; +using namespace mlir::pto; + +namespace { + +static bool hasSameDataShapeAndElementType(VMIVRegType lhs, VMIVRegType rhs) { + return lhs && rhs && lhs.getElementCount() == rhs.getElementCount() && + lhs.getElementType() == rhs.getElementType(); +} + +static bool isUnitContiguousLayout(VMILayoutAttr layout) { + return layout && layout.isContiguous() && layout.getLaneStride() == 1; +} + +static bool isLoadProducerLayout(VMIVRegType type) { + if (!type) + return false; + VMILayoutAttr layout = type.getLayoutAttr(); + if (!layout) + return false; + if (layout.isContiguous() && layout.getLaneStride() == 1) + return true; + if (layout.isContiguous() && layout.getLaneStride() == 2) { + unsigned elementBits = pto::getPTOStorageElemBitWidth(type.getElementType()); + return elementBits == 8 || elementBits == 16 || elementBits == 32; + } + if (layout.isContiguous() && layout.getLaneStride() == 4) { + unsigned elementBits = pto::getPTOStorageElemBitWidth(type.getElementType()); + return elementBits == 8; + } + if (!layout.isDeinterleaved() || layout.getBlockElems() != 1 || + layout.getLaneStride() != 1 || + (layout.getFactor() != 2 && layout.getFactor() != 4)) + return false; + unsigned elementBits = pto::getPTOStorageElemBitWidth(type.getElementType()); + return elementBits == 8 || elementBits == 16 || elementBits == 32; +} + +static bool isFoldableLoadEnsure(VMIEnsureLayoutOp ensure) { + auto load = ensure.getSource().getDefiningOp(); + if (!load) + return false; + + auto sourceType = dyn_cast(ensure.getSource().getType()); + auto resultType = dyn_cast(ensure.getResult().getType()); + if (!hasSameDataShapeAndElementType(sourceType, resultType)) + return false; + + return isLoadProducerLayout(resultType); +} + +static void tryFoldLoadEnsures( + VMILoadOp load, SmallVectorImpl &maybeDeadEnsures) { + auto sourceType = dyn_cast(load.getResult().getType()); + if (!sourceType) + return; + + VMIVRegType targetType; + SmallVector ensures; + for (OpOperand &use : load.getResult().getUses()) { + auto ensure = dyn_cast(use.getOwner()); + if (!ensure || use.getOperandNumber() != 0 || !isFoldableLoadEnsure(ensure)) + return; + + auto resultType = cast(ensure.getResult().getType()); + if (!targetType) { + targetType = resultType; + } else if (targetType != resultType) { + return; + } + ensures.push_back(ensure); + } + + if (ensures.empty() || targetType == sourceType) + return; + + load.getResult().setType(targetType); + for (VMIEnsureLayoutOp ensure : ensures) { + ensure.getResult().replaceAllUsesWith(load.getResult()); + maybeDeadEnsures.push_back(ensure); + } +} + +static void +tryFoldNestedEnsureLayout(VMIEnsureLayoutOp ensure, + SmallVectorImpl &maybeDeadEnsures) { + auto inner = ensure.getSource().getDefiningOp(); + if (!inner) + return; + + if (inner.getSource().getType() != ensure.getResult().getType()) + return; + + ensure.getResult().replaceAllUsesWith(inner.getSource()); + maybeDeadEnsures.push_back(ensure); + maybeDeadEnsures.push_back(inner); +} + +static bool isFoldableStoreEnsure(VMIEnsureLayoutOp ensure) { + auto sourceType = dyn_cast(ensure.getSource().getType()); + auto resultType = dyn_cast(ensure.getResult().getType()); + if (!sourceType || !resultType) + return false; + + VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + if (!isUnitContiguousLayout(resultLayout) || + isUnitContiguousLayout(sourceLayout)) + return false; + + VMILayoutSupport supports; + return succeeded(supports.getStoreLayoutFact(sourceType)); +} + +static void tryFoldEnsureLayoutIntoOperand( + OpOperand &operand, SmallVectorImpl &maybeDeadEnsures) { + auto ensure = operand.get().getDefiningOp(); + if (!ensure || !isFoldableStoreEnsure(ensure)) + return; + + operand.set(ensure.getSource()); + maybeDeadEnsures.push_back(ensure); +} + +static void tryFoldEnsureLayoutIntoMaskedStore( + VMIMaskedStoreOp store, + SmallVectorImpl &maybeDeadEnsures, + SmallVectorImpl &maybeDeadMaskEnsures) { + auto ensure = store.getValue().getDefiningOp(); + if (!ensure || !isFoldableStoreEnsure(ensure)) + return; + auto maskEnsure = store.getMask().getDefiningOp(); + if (!maskEnsure) + return; + + auto sourceType = dyn_cast(ensure.getSource().getType()); + auto maskSourceType = dyn_cast(maskEnsure.getSource().getType()); + auto maskResultType = dyn_cast(maskEnsure.getResult().getType()); + if (!sourceType || !maskSourceType || !maskResultType) + return; + + VMILayoutAttr maskResultLayout = maskResultType.getLayoutAttr(); + if (!isUnitContiguousLayout(maskResultLayout)) + return; + + VMILayoutSupport supports; + if (failed(supports.getMaskedStoreLayoutFact(sourceType, maskSourceType))) + return; + + store.getValueMutable().set(ensure.getSource()); + store.getMaskMutable().set(maskEnsure.getSource()); + maybeDeadEnsures.push_back(ensure); + maybeDeadMaskEnsures.push_back(maskEnsure); +} + +struct VMILayoutFoldPass + : public mlir::pto::impl::VMILayoutFoldBase { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(VMILayoutFoldPass) + + void runOnOperation() override { + ModuleOp module = getOperation(); + SmallVector maybeDeadEnsures; + SmallVector maybeDeadMaskEnsures; + + module.walk([&](VMILoadOp load) { + tryFoldLoadEnsures(load, maybeDeadEnsures); + }); + + module.walk([&](VMIEnsureLayoutOp ensure) { + tryFoldNestedEnsureLayout(ensure, maybeDeadEnsures); + }); + + module.walk([&](Operation *op) { + if (auto store = dyn_cast(op)) + tryFoldEnsureLayoutIntoOperand(store.getValueMutable(), + maybeDeadEnsures); + if (auto maskedStore = dyn_cast(op)) + tryFoldEnsureLayoutIntoMaskedStore(maskedStore, maybeDeadEnsures, + maybeDeadMaskEnsures); + }); + + for (VMIEnsureMaskLayoutOp ensure : llvm::reverse(maybeDeadMaskEnsures)) { + if (ensure->use_empty()) + ensure.erase(); + } + for (VMIEnsureLayoutOp ensure : llvm::reverse(maybeDeadEnsures)) { + if (ensure->use_empty()) + ensure.erase(); + } + } +}; + +} // namespace + +std::unique_ptr mlir::pto::createVMILayoutFoldPass() { + return std::make_unique(); +} diff --git a/lib/PTO/Transforms/VMILayoutPropagation.cpp b/lib/PTO/Transforms/VMILayoutPropagation.cpp new file mode 100644 index 0000000000..16da667587 --- /dev/null +++ b/lib/PTO/Transforms/VMILayoutPropagation.cpp @@ -0,0 +1,1238 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +//===- VMILayoutPropagation.cpp - VMI layout request propagation ----------===// +//===----------------------------------------------------------------------===// + +#include "PTO/Transforms/VMILayoutPropagation.h" + +#include "PTO/IR/VMIUtils.h" +#include "PTO/Transforms/VMILayoutSupport.h" + +using namespace mlir; +using namespace mlir::pto; + +namespace { + +struct VMILayoutFact { + Value value; + OpOperand *operand = nullptr; + VMILayoutAttr layout; +}; + +struct VMILayoutRelation { + SmallVector facts; +}; + +static VMILayoutFact valueFact(Value value, VMILayoutAttr layout) { + return VMILayoutFact{value, /*operand=*/nullptr, layout}; +} + +static VMILayoutFact operandFact(OpOperand &operand, VMILayoutAttr layout) { + return VMILayoutFact{/*value=*/{}, &operand, layout}; +} + +static VMILayoutRelation makeRelation(SmallVector facts) { + VMILayoutRelation relation; + relation.facts = std::move(facts); + return relation; +} + +static SmallVector +makeSingleRelation(SmallVector facts) { + SmallVector relations; + relations.push_back(makeRelation(std::move(facts))); + return relations; +} + +static bool hasAmbiguousTransferTargets(ArrayRef facts) { + auto sameTarget = [](const VMILayoutFact &lhs, const VMILayoutFact &rhs) { + if (lhs.operand || rhs.operand) + return lhs.operand && lhs.operand == rhs.operand; + return lhs.value && lhs.value == rhs.value; + }; + + for (auto [index, fact] : llvm::enumerate(facts)) { + for (const VMILayoutFact &other : facts.drop_front(index + 1)) { + if (sameTarget(fact, other) && fact.layout != other.layout) + return true; + } + } + return false; +} + +static bool relationContainsOperandLayout(const VMILayoutRelation &relation, + OpOperand &operand, + VMILayoutAttr layout) { + for (const VMILayoutFact &fact : relation.facts) + if (fact.operand == &operand) + return fact.layout == layout; + return false; +} + +static bool relationContainsValueLayout(const VMILayoutRelation &relation, + Value value, VMILayoutAttr layout) { + for (const VMILayoutFact &fact : relation.facts) + if (!fact.operand && fact.value == value) + return fact.layout == layout; + return false; +} + +static Type getValueTypeWithLayout(Value value, VMILayoutAttr layout) { + if (auto type = dyn_cast(value.getType())) + return VMIVRegType::get(type.getContext(), type.getElementCount(), + type.getElementType(), layout); + if (auto type = dyn_cast(value.getType())) + return VMIMaskType::get(type.getContext(), type.getElementCount(), + type.getGranularity(), layout); + return {}; +} + +class VMILayoutTransfer { +public: + virtual ~VMILayoutTransfer() = default; + + virtual FailureOr> + query(Operation *op, Value changedValue, VMILayoutAttr changedLayout, + const VMILayoutPropagator &propagator, + OpOperand *changedOperand) const = 0; + + LogicalResult propagate(Operation *op, Value changedValue, + VMILayoutAttr changedLayout, + VMILayoutPropagator &propagator, + OpOperand *changedOperand) const { + FailureOr> relations = + query(op, changedValue, changedLayout, propagator, changedOperand); + if (failed(relations) || relations->empty()) + return success(); + if (relations->size() != 1) + return success(); + const VMILayoutRelation &relation = relations->front(); + if (hasAmbiguousTransferTargets(relation.facts)) + return success(); + for (const VMILayoutFact &fact : relation.facts) { + if (fact.operand) { + if (failed(propagator.request(*fact.operand, fact.layout))) + return failure(); + continue; + } + if (failed(propagator.request(fact.value, fact.layout))) + return failure(); + } + return success(); + } +}; + +class VMILayoutMaterializationTransfer final { +public: + FailureOr> + query(Value source, VMILayoutAttr sourceLayout, + VMILayoutAttr resultLayout) const { + if (!source || !sourceLayout || !resultLayout) + return failure(); + if (sourceLayout == resultLayout) + return makeSingleRelation(SmallVector{ + valueFact(source, sourceLayout)}); + + Type sourceAssignedType = getValueTypeWithLayout(source, sourceLayout); + Type resultAssignedType = getValueTypeWithLayout(source, resultLayout); + if (!sourceAssignedType || !resultAssignedType) + return failure(); + + VMILayoutSupport supports; + if (auto sourceVRegType = dyn_cast(sourceAssignedType)) { + auto resultVRegType = dyn_cast(resultAssignedType); + if (!resultVRegType || + failed(supports.getEnsureLayoutFact(sourceVRegType, + resultVRegType))) + return failure(); + return makeSingleRelation(SmallVector{ + valueFact(source, sourceLayout)}); + } + + if (auto sourceMaskType = dyn_cast(sourceAssignedType)) { + auto resultMaskType = dyn_cast(resultAssignedType); + if (!resultMaskType || + failed(supports.getEnsureMaskLayoutFact(sourceMaskType, + resultMaskType))) + return failure(); + return makeSingleRelation(SmallVector{ + valueFact(source, sourceLayout)}); + } + return failure(); + } +}; + +static bool isSameLayoutOp(Operation *op) { + return isa(op); +} + +static bool isCastOp(Operation *op) { + return isa(op); +} + +class VMISameLayoutTransfer final : public VMILayoutTransfer { +public: + FailureOr> + query(Operation *op, Value changedValue, VMILayoutAttr changedLayout, + const VMILayoutPropagator &propagator, + OpOperand *changedOperand) const override { + SmallVector facts; + for (OpOperand &operand : op->getOpOperands()) + facts.push_back(operandFact(operand, changedLayout)); + for (Value result : op->getResults()) + facts.push_back(valueFact(result, changedLayout)); + return makeSingleRelation(std::move(facts)); + } +}; + +class VMICastTransfer final : public VMILayoutTransfer { +public: + FailureOr> + query(Operation *op, Value changedValue, VMILayoutAttr changedLayout, + const VMILayoutPropagator &propagator, + OpOperand *changedOperand) const override { + if (op->getNumOperands() != 1 || op->getNumResults() != 1) + return failure(); + + auto sourceType = dyn_cast(op->getOperand(0).getType()); + auto resultType = dyn_cast(op->getResult(0).getType()); + if (!sourceType || !resultType) + return failure(); + + VMILayoutSupport supports; + + if (changedValue == op->getOperand(0)) { + FailureOr> facts = + supports.getCastLayoutFactsForLayout( + sourceType, resultType, VMICastLayoutPort::Source, + changedLayout); + if (failed(facts) || facts->empty()) + return failure(); + SmallVector relations; + for (const VMICastLayoutFact &fact : *facts) + relations.push_back(makeRelation(SmallVector{ + operandFact(op->getOpOperand(0), fact.sourceLayout), + valueFact(op->getResult(0), fact.resultLayout)})); + return relations; + } + + if (changedValue == op->getResult(0)) { + FailureOr> facts = + supports.getCastLayoutFactsForLayout( + sourceType, resultType, VMICastLayoutPort::Result, + changedLayout); + if (failed(facts) || facts->empty()) + return failure(); + SmallVector relations; + for (const VMICastLayoutFact &fact : *facts) + relations.push_back(makeRelation(SmallVector{ + operandFact(op->getOpOperand(0), fact.sourceLayout), + valueFact(op->getResult(0), fact.resultLayout)})); + return relations; + } + return failure(); + } +}; + +class VMIMaskGranularityCastTransfer final : public VMILayoutTransfer { +public: + FailureOr> + query(Operation *op, Value changedValue, VMILayoutAttr changedLayout, + const VMILayoutPropagator &propagator, + OpOperand *changedOperand) const override { + auto ensure = dyn_cast(op); + if (!ensure) + return failure(); + + auto sourceType = dyn_cast(ensure.getSource().getType()); + auto resultType = dyn_cast(ensure.getResult().getType()); + if (!sourceType || !resultType) + return failure(); + + VMILayoutSupport supports; + if (changedValue == ensure.getSource()) { + FailureOr> facts = + supports.getMaskGranularityCastLayoutFactsForLayout( + sourceType, resultType, VMICastLayoutPort::Source, + changedLayout); + if (failed(facts) || facts->empty()) + return failure(); + SmallVector relations; + for (const VMIMaskGranularityCastLayoutFact &fact : *facts) + relations.push_back(makeRelation(SmallVector{ + operandFact(ensure.getSourceMutable(), fact.sourceLayout), + valueFact(ensure.getResult(), fact.resultLayout)})); + return relations; + } + + if (changedValue == ensure.getResult()) { + FailureOr> facts = + supports.getMaskGranularityCastLayoutFactsForLayout( + sourceType, resultType, VMICastLayoutPort::Result, + changedLayout); + if (failed(facts) || facts->empty()) + return failure(); + SmallVector relations; + for (const VMIMaskGranularityCastLayoutFact &fact : *facts) + relations.push_back(makeRelation(SmallVector{ + operandFact(ensure.getSourceMutable(), fact.sourceLayout), + valueFact(ensure.getResult(), fact.resultLayout)})); + return relations; + } + return failure(); + } +}; + +class VMIFreeResultLayoutTransfer final : public VMILayoutTransfer { +public: + FailureOr> + query(Operation *op, Value changedValue, VMILayoutAttr changedLayout, + const VMILayoutPropagator &propagator, + OpOperand *changedOperand) const override { + if (isa(changedValue) && changedValue.getDefiningOp() == op) + return makeSingleRelation(SmallVector{ + valueFact(changedValue, changedLayout)}); + return failure(); + } +}; + +class VMILoadTransfer final : public VMILayoutTransfer { +public: + FailureOr> + query(Operation *op, Value changedValue, VMILayoutAttr changedLayout, + const VMILayoutPropagator &propagator, + OpOperand *changedOperand) const override { + auto load = dyn_cast(op); + if (!load || changedValue != load.getResult()) + return failure(); + auto resultType = dyn_cast(load.getResult().getType()); + if (!resultType) + return failure(); + auto assignedType = VMIVRegType::get( + resultType.getContext(), resultType.getElementCount(), + resultType.getElementType(), changedLayout); + VMILayoutSupport supports; + if (failed(supports.getLoadLayoutFact(assignedType))) + return failure(); + return makeSingleRelation(SmallVector{ + valueFact(load.getResult(), changedLayout)}); + } +}; + +class VMIDeinterleaveLoadTransfer final : public VMILayoutTransfer { +public: + FailureOr> + query(Operation *op, Value changedValue, VMILayoutAttr changedLayout, + const VMILayoutPropagator &propagator, + OpOperand *changedOperand) const override { + auto load = dyn_cast(op); + if (!load) + return failure(); + + VMIDeinterleaveLoadLayoutPort port; + if (changedValue == load.getLow()) { + port = VMIDeinterleaveLoadLayoutPort::Low; + } else if (changedValue == load.getHigh()) { + port = VMIDeinterleaveLoadLayoutPort::High; + } else { + return failure(); + } + + auto valueType = dyn_cast(changedValue.getType()); + if (!valueType) + return failure(); + VMILayoutSupport supports; + FailureOr> facts = + supports.getDeinterleaveLoadLayoutFactsForLayout( + valueType, port, changedLayout); + if (failed(facts) || facts->empty()) + return failure(); + + SmallVector relations; + for (const VMIDeinterleaveLoadLayoutFact &fact : *facts) { + relations.push_back(makeRelation(SmallVector{ + valueFact(load.getLow(), fact.lowLayout), + valueFact(load.getHigh(), fact.highLayout)})); + } + return relations; + } +}; + +class VMIGroupLoadTransfer final : public VMILayoutTransfer { +public: + FailureOr> + query(Operation *op, Value changedValue, VMILayoutAttr changedLayout, + const VMILayoutPropagator &propagator, + OpOperand *changedOperand) const override { + auto load = dyn_cast(op); + if (!load || changedValue != load.getResult()) + return failure(); + auto resultType = dyn_cast(load.getResult().getType()); + if (!resultType) + return failure(); + auto assignedType = VMIVRegType::get( + resultType.getContext(), resultType.getElementCount(), + resultType.getElementType(), changedLayout); + VMILayoutSupport supports; + if (failed(supports.getGroupLoadLayoutFact( + assignedType, load.getRowStride(), + load.getNumGroupsAttr().getInt()))) + return failure(); + return makeSingleRelation(SmallVector{ + valueFact(load.getResult(), changedLayout)}); + } +}; + +class VMIGroupReduceTransfer final : public VMILayoutTransfer { +public: + FailureOr> + query(Operation *op, Value changedValue, VMILayoutAttr changedLayout, + const VMILayoutPropagator &propagator, + OpOperand *changedOperand) const override { + if (auto reduce = dyn_cast(op)) + return queryReduce(reduce, changedValue, changedLayout, changedOperand); + if (auto reduce = dyn_cast(op)) + return queryReduce(reduce, changedValue, changedLayout, changedOperand); + if (auto reduce = dyn_cast(op)) + return queryReduce(reduce, changedValue, changedLayout, changedOperand); + if (auto reduce = dyn_cast(op)) + return queryReduce(reduce, changedValue, changedLayout, changedOperand); + if (auto reduce = dyn_cast(op)) + return queryReduce(reduce, changedValue, changedLayout, changedOperand); + if (auto reduce = dyn_cast(op)) + return queryReduce(reduce, changedValue, changedLayout, changedOperand); + return failure(); + } + +private: + template + FailureOr> + queryReduce(OpTy reduce, Value changedValue, VMILayoutAttr changedLayout, + OpOperand *changedOperand) const { + auto sourceType = dyn_cast(reduce.getSource().getType()); + auto resultType = dyn_cast(reduce.getResult().getType()); + if (!sourceType || !resultType) + return failure(); + + VMIGroupReduceLayoutPort port; + if (changedValue == reduce.getSource()) { + port = VMIGroupReduceLayoutPort::Source; + } else if (changedValue == reduce.getMask()) { + port = VMIGroupReduceLayoutPort::Mask; + } else if (changedValue == reduce.getResult()) { + port = VMIGroupReduceLayoutPort::Result; + } else { + return failure(); + } + + VMILayoutSupport supports; + FailureOr> facts = + supports.getGroupReduceLayoutFactsForLayout( + sourceType, reduce.getNumGroupsAttr().getInt(), port, + changedLayout); + if (failed(facts) || facts->empty()) + return failure(); + SmallVector relations; + for (const VMIGroupReduceLayoutFact &fact : *facts) { + relations.push_back(makeRelation(SmallVector{ + operandFact(reduce.getSourceMutable(), fact.sourceLayout), + operandFact(reduce.getMaskMutable(), fact.maskLayout), + valueFact(reduce.getResult(), fact.resultLayout)})); + } + return relations; + } +}; + +class VMIGroupSlotLoadTransfer final : public VMILayoutTransfer { +public: + FailureOr> + query(Operation *op, Value changedValue, VMILayoutAttr changedLayout, + const VMILayoutPropagator &propagator, + OpOperand *changedOperand) const override { + auto load = dyn_cast(op); + if (!load || changedValue != load.getResult()) + return failure(); + auto resultType = dyn_cast(load.getResult().getType()); + if (!resultType) + return failure(); + auto assignedType = VMIVRegType::get( + resultType.getContext(), resultType.getElementCount(), + resultType.getElementType(), changedLayout); + VMILayoutSupport supports; + if (failed(supports.getGroupSlotLoadLayoutFact( + assignedType, load.getNumGroupsAttr().getInt()))) + return failure(); + return makeSingleRelation(SmallVector{ + valueFact(load.getResult(), changedLayout)}); + } +}; + +class VMIGroupBroadcastLoadTransfer final : public VMILayoutTransfer { +public: + FailureOr> + query(Operation *op, Value changedValue, VMILayoutAttr changedLayout, + const VMILayoutPropagator &propagator, + OpOperand *changedOperand) const override { + auto load = dyn_cast(op); + if (!load || changedValue != load.getResult()) + return failure(); + auto resultType = dyn_cast(load.getResult().getType()); + if (!resultType) + return failure(); + auto assignedType = VMIVRegType::get( + resultType.getContext(), resultType.getElementCount(), + resultType.getElementType(), changedLayout); + VMILayoutSupport supports; + if (failed(supports.getGroupBroadcastLoadLayoutFact( + assignedType, load.getSourceGroupStride(), + load.getNumGroupsAttr().getInt()))) + return failure(); + return makeSingleRelation(SmallVector{ + valueFact(load.getResult(), changedLayout)}); + } +}; + +class VMIGroupBroadcastTransfer final : public VMILayoutTransfer { +public: + FailureOr> + query(Operation *op, Value changedValue, VMILayoutAttr changedLayout, + const VMILayoutPropagator &propagator, + OpOperand *changedOperand) const override { + auto broadcast = dyn_cast(op); + if (!broadcast) + return failure(); + auto sourceType = dyn_cast(broadcast.getSource().getType()); + auto resultType = dyn_cast(broadcast.getResult().getType()); + if (!sourceType || !resultType) + return failure(); + + int64_t numGroups = broadcast.getNumGroupsAttr().getInt(); + VMIGroupBroadcastLayoutPort port; + if (changedValue == broadcast.getSource()) { + port = VMIGroupBroadcastLayoutPort::Source; + } else if (changedValue == broadcast.getResult()) { + port = VMIGroupBroadcastLayoutPort::Result; + } else { + return failure(); + } + + VMILayoutSupport supports; + FailureOr> facts = + supports.getGroupBroadcastLayoutFactsForLayout( + sourceType, resultType, numGroups, port, changedLayout); + if (failed(facts) || facts->empty()) + return failure(); + SmallVector relations; + for (const VMIGroupBroadcastLayoutFact &fact : *facts) { + relations.push_back(makeRelation(SmallVector{ + operandFact(broadcast.getSourceMutable(), fact.sourceLayout), + valueFact(broadcast.getResult(), fact.resultLayout)})); + } + return relations; + } +}; + +class VMIInterleaveTransfer final : public VMILayoutTransfer { +public: + FailureOr> + query(Operation *op, Value changedValue, VMILayoutAttr changedLayout, + const VMILayoutPropagator &propagator, + OpOperand *changedOperand) const override { + if (auto vintlv = dyn_cast(op)) + return queryInterleave(vintlv, /*vintlv=*/true, changedValue, + changedLayout, changedOperand); + if (auto vdintlv = dyn_cast(op)) + return queryInterleave(vdintlv, /*vintlv=*/false, changedValue, + changedLayout, changedOperand); + return failure(); + } + +private: + template + FailureOr> + queryInterleave(OpTy op, bool vintlv, Value changedValue, + VMILayoutAttr changedLayout, + OpOperand *changedOperand) const { + auto lowType = dyn_cast(op.getLow().getType()); + if (!lowType) + return failure(); + + VMIInterleaveLayoutPort port; + if (changedOperand == &op.getLhsMutable() || changedValue == op.getLhs()) { + port = VMIInterleaveLayoutPort::Lhs; + } else if (changedOperand == &op.getRhsMutable() || + changedValue == op.getRhs()) { + port = VMIInterleaveLayoutPort::Rhs; + } else if (changedOperand == &op.getMaskMutable() || + changedValue == op.getMask()) { + port = VMIInterleaveLayoutPort::Mask; + } else if (changedValue == op.getLow()) { + port = VMIInterleaveLayoutPort::Low; + } else if (changedValue == op.getHigh()) { + port = VMIInterleaveLayoutPort::High; + } else { + return failure(); + } + + VMILayoutSupport supports; + FailureOr> facts = + vintlv ? supports.getVintlvLayoutFactsForLayout(lowType, port, + changedLayout) + : supports.getVdintlvLayoutFactsForLayout(lowType, port, + changedLayout); + if (failed(facts) || facts->empty()) + return failure(); + + SmallVector relations; + for (const VMIInterleaveLayoutFact &fact : *facts) { + relations.push_back(makeRelation(SmallVector{ + operandFact(op.getLhsMutable(), fact.lhsLayout), + operandFact(op.getRhsMutable(), fact.rhsLayout), + operandFact(op.getMaskMutable(), fact.maskLayout), + valueFact(op.getLow(), fact.lowLayout), + valueFact(op.getHigh(), fact.highLayout)})); + } + return relations; + } +}; + +class VMIGatherTransfer final : public VMILayoutTransfer { +public: + FailureOr> + query(Operation *op, Value changedValue, VMILayoutAttr changedLayout, + const VMILayoutPropagator &propagator, + OpOperand *changedOperand) const override { + auto gather = dyn_cast(op); + if (!gather) + return failure(); + if (changedValue != gather.getIndices() && changedValue != gather.getMask() && + changedValue != gather.getPassthru() && changedValue != gather.getResult()) + return failure(); + if (!changedLayout.isContiguous() || changedLayout.getLaneStride() != 1) + return failure(); + return makeSingleRelation(SmallVector{ + operandFact(gather.getIndicesMutable(), changedLayout), + operandFact(gather.getMaskMutable(), changedLayout), + operandFact(gather.getPassthruMutable(), changedLayout), + valueFact(gather.getResult(), changedLayout)}); + } +}; + +class VMIStoreTransfer final : public VMILayoutTransfer { +public: + FailureOr> + query(Operation *op, Value changedValue, VMILayoutAttr changedLayout, + const VMILayoutPropagator &propagator, + OpOperand *changedOperand) const override { + auto store = dyn_cast(op); + if (!store || changedValue != store.getValue()) + return failure(); + + auto valueType = dyn_cast(store.getValue().getType()); + if (!valueType) + return failure(); + + auto assignedValueType = + VMIVRegType::get(valueType.getContext(), valueType.getElementCount(), + valueType.getElementType(), changedLayout); + VMILayoutSupport supports; + VMILayoutAttr useLayout = changedLayout; + if (failed(supports.getStoreLayoutFact(assignedValueType))) + useLayout = VMILayoutAttr::getContiguous(valueType.getContext()); + + return makeSingleRelation(SmallVector{ + operandFact(store.getValueMutable(), useLayout)}); + } +}; + +class VMIMaskedLoadTransfer final : public VMILayoutTransfer { +public: + FailureOr> + query(Operation *op, Value changedValue, VMILayoutAttr changedLayout, + const VMILayoutPropagator &propagator, + OpOperand *changedOperand) const override { + auto load = dyn_cast(op); + if (!load) + return failure(); + if (changedValue != load.getResult() && changedValue != load.getMask() && + changedValue != load.getPassthru()) + return failure(); + + auto resultType = dyn_cast(load.getResult().getType()); + auto maskType = dyn_cast(load.getMask().getType()); + auto passthruType = dyn_cast(load.getPassthru().getType()); + if (!resultType || !maskType || !passthruType) + return failure(); + + auto assignedResultType = + VMIVRegType::get(resultType.getContext(), resultType.getElementCount(), + resultType.getElementType(), changedLayout); + auto assignedMaskType = + VMIMaskType::get(maskType.getContext(), maskType.getElementCount(), + maskType.getGranularity(), changedLayout); + auto assignedPassthruType = VMIVRegType::get( + passthruType.getContext(), passthruType.getElementCount(), + passthruType.getElementType(), changedLayout); + VMILayoutSupport supports; + if (failed(supports.getMaskedLoadLayoutFact( + assignedResultType, assignedMaskType, assignedPassthruType))) + return failure(); + + return makeSingleRelation(SmallVector{ + valueFact(load.getResult(), changedLayout), + operandFact(load.getMaskMutable(), changedLayout), + operandFact(load.getPassthruMutable(), changedLayout)}); + } +}; + +class VMIMaskedStoreTransfer final : public VMILayoutTransfer { +public: + FailureOr> + query(Operation *op, Value changedValue, VMILayoutAttr changedLayout, + const VMILayoutPropagator &propagator, + OpOperand *changedOperand) const override { + auto store = dyn_cast(op); + if (!store) + return failure(); + if (changedValue != store.getValue() && changedValue != store.getMask()) + return failure(); + + auto valueType = dyn_cast(store.getValue().getType()); + auto maskType = dyn_cast(store.getMask().getType()); + if (!valueType || !maskType) + return failure(); + + VMILayoutSupport supports; + VMILayoutAttr useLayout = changedLayout; + auto assignedValueType = + VMIVRegType::get(valueType.getContext(), valueType.getElementCount(), + valueType.getElementType(), useLayout); + auto assignedMaskType = + VMIMaskType::get(maskType.getContext(), maskType.getElementCount(), + maskType.getGranularity(), useLayout); + if (failed(supports.getMaskedStoreLayoutFact(assignedValueType, + assignedMaskType))) + useLayout = VMILayoutAttr::getContiguous(valueType.getContext()); + + return makeSingleRelation(SmallVector{ + operandFact(store.getValueMutable(), useLayout), + operandFact(store.getMaskMutable(), useLayout)}); + } +}; + +const VMILayoutTransfer *getTransfer(Operation *op) { + static VMISameLayoutTransfer sameLayoutTransfer; + static VMICastTransfer castTransfer; + static VMIMaskGranularityCastTransfer maskGranularityCastTransfer; + static VMIFreeResultLayoutTransfer freeResultLayoutTransfer; + static VMILoadTransfer loadTransfer; + static VMIDeinterleaveLoadTransfer deinterleaveLoadTransfer; + static VMIGroupLoadTransfer groupLoadTransfer; + static VMIGroupReduceTransfer groupReduceTransfer; + static VMIGroupSlotLoadTransfer groupSlotLoadTransfer; + static VMIGroupBroadcastLoadTransfer groupBroadcastLoadTransfer; + static VMIGroupBroadcastTransfer groupBroadcastTransfer; + static VMIInterleaveTransfer interleaveTransfer; + static VMIGatherTransfer gatherTransfer; + static VMIStoreTransfer storeTransfer; + static VMIMaskedLoadTransfer maskedLoadTransfer; + static VMIMaskedStoreTransfer maskedStoreTransfer; + + if (isa(op)) + return &freeResultLayoutTransfer; + if (isa(op)) + return &loadTransfer; + if (isa(op)) + return &deinterleaveLoadTransfer; + if (isa(op)) + return &groupLoadTransfer; + if (isa(op)) + return &groupReduceTransfer; + if (isa(op)) + return &groupSlotLoadTransfer; + if (isa(op)) + return &groupBroadcastLoadTransfer; + if (isa(op)) + return &groupBroadcastTransfer; + if (isa(op)) + return &interleaveTransfer; + if (isa(op)) + return &gatherTransfer; + if (isSameLayoutOp(op)) + return &sameLayoutTransfer; + if (isa(op)) + return &maskGranularityCastTransfer; + if (isCastOp(op)) + return &castTransfer; + if (isa(op)) + return &storeTransfer; + if (isa(op)) + return &maskedLoadTransfer; + if (isa(op)) + return &maskedStoreTransfer; + return nullptr; +} + +} // namespace + +VMILayoutPropagator::VMILayoutPropagator(Operation *scope) + : scope(scope), ctx(scope ? scope->getContext() : nullptr) {} + +bool VMILayoutPropagator::isLayoutValue(Value value) const { + return isa(value.getType()); +} + +VMILayoutAttr VMILayoutPropagator::getCurrentLayout(Value value) const { + if (auto type = dyn_cast(value.getType())) + return type.getLayoutAttr(); + if (auto type = dyn_cast(value.getType())) + return type.getLayoutAttr(); + return {}; +} + +Type VMILayoutPropagator::getTypeWithLayout(Value value, + VMILayoutAttr layout) const { + if (auto type = dyn_cast(value.getType())) + return VMIVRegType::get(ctx, type.getElementCount(), + type.getElementType(), layout); + if (auto type = dyn_cast(value.getType())) + return VMIMaskType::get(ctx, type.getElementCount(), type.getGranularity(), + layout); + return {}; +} + +bool VMILayoutPropagator::canUseOperandLayout(OpOperand &operand, + VMILayoutAttr layout) const { + if (!layout) + return false; + if (!isLayoutValue(operand.get())) + return true; + const VMILayoutTransfer *transfer = getTransfer(operand.getOwner()); + if (!transfer) + return false; + FailureOr> relations = transfer->query( + operand.getOwner(), operand.get(), layout, *this, &operand); + if (failed(relations)) + return false; + for (const VMILayoutRelation &relation : *relations) + if (relationContainsOperandLayout(relation, operand, layout)) + return true; + return false; +} + +VMILayoutAttr VMILayoutPropagator::getRequestedLayout(Value value) const { + auto it = assignments.find(value); + if (it == assignments.end()) + return {}; + return it->second.layout; +} + +VMILayoutAttr +VMILayoutPropagator::getRequestedOrCurrentLayout(Value value) const { + if (VMILayoutAttr layout = getRequestedLayout(value)) + return layout; + return getCurrentLayout(value); +} + +VMILayoutAttr VMILayoutPropagator::getOperandLayout(OpOperand &operand) const { + const VMIValueLayoutAssignment *assignment = lookup(operand.get()); + if (!assignment) + return getCurrentLayout(operand.get()); + + for (const VMILayoutConflict &conflict : assignment->conflicts) + if (conflict.operand == &operand) + return conflict.layout; + return assignment->layout; +} + +const VMIValueLayoutAssignment * +VMILayoutPropagator::lookup(Value value) const { + auto it = assignments.find(value); + if (it == assignments.end()) + return nullptr; + return &it->second; +} + +void VMILayoutPropagator::addEquivalentValues(Value lhs, Value rhs) { + if (!isLayoutValue(lhs) || !isLayoutValue(rhs) || lhs == rhs) + return; + if (isa(lhs.getType()) != isa(rhs.getType())) + return; + + auto addEdge = [&](Value from, Value to) { + SmallVector &values = equivalentValues[from]; + if (!llvm::is_contained(values, to)) + values.push_back(to); + }; + addEdge(lhs, rhs); + addEdge(rhs, lhs); +} + +void VMILayoutPropagator::enqueue(Value value, VMILayoutAttr layout) { + if (!value || !layout) + return; + LayoutFact fact(value, layout); + if (llvm::is_contained(seenFacts, fact)) + return; + seenFacts.push_back(fact); + worklist.push_back(fact); +} + +LogicalResult +VMILayoutPropagator::addUseConflict(OpOperand &operand, + VMIValueLayoutAssignment &assignment, + VMILayoutAttr layout) { + for (VMILayoutConflict &conflict : assignment.conflicts) { + if (conflict.operand != &operand) + continue; + if (conflict.layout == layout) + return success(); + return success(); + } + assignment.conflicts.push_back(VMILayoutConflict{&operand, layout}); + return success(); +} + +bool VMILayoutPropagator::canProduceValueLayout(Value value, + VMILayoutAttr layout) const { + if (!layout || !isLayoutValue(value)) + return false; + if (getCurrentLayout(value) == layout) + return true; + if (isa(value)) + return isTypeRewriteable(value); + if (auto result = dyn_cast(value)) { + Operation *definingOp = result.getDefiningOp(); + const VMILayoutTransfer *transfer = getTransfer(definingOp); + if (!transfer) + return isTypeRewriteable(value); + FailureOr> relations = transfer->query( + definingOp, value, layout, *this, /*changedOperand=*/nullptr); + if (failed(relations)) + return false; + for (const VMILayoutRelation &relation : *relations) + if (relationContainsValueLayout(relation, value, layout)) + return true; + return false; + } + return false; +} + +bool VMILayoutPropagator::canMaterializeLayout( + Value value, VMILayoutAttr sourceLayout, VMILayoutAttr resultLayout) const { + static VMILayoutMaterializationTransfer materializationTransfer; + FailureOr> relations = + materializationTransfer.query(value, sourceLayout, resultLayout); + return succeeded(relations) && !relations->empty(); +} + +LogicalResult VMILayoutPropagator::request(Value value, VMILayoutAttr layout) { + if (!layout) + return failure(); + if (!isLayoutValue(value)) + return success(); + + auto it = assignments.find(value); + if (it == assignments.end()) { + auto inserted = assignments.try_emplace(value); + it = inserted.first; + orderedValues.push_back(value); + } + VMIValueLayoutAssignment &assignment = it->second; + if (!assignment.layout) { + if (!canProduceValueLayout(value, layout)) + return success(); + assignment.layout = layout; + enqueue(value, layout); + return success(); + } + if (assignment.layout == layout) + return success(); + return success(); +} + +LogicalResult VMILayoutPropagator::request(OpOperand &operand, + VMILayoutAttr layout) { + if (!layout) + return failure(); + Value value = operand.get(); + if (!isLayoutValue(value)) + return success(); + + auto it = assignments.find(value); + if (it == assignments.end()) { + auto inserted = assignments.try_emplace(value); + it = inserted.first; + orderedValues.push_back(value); + } + + VMIValueLayoutAssignment &assignment = it->second; + if (assignment.layout == layout) + return propagateOperandFact(operand, layout); + if (!assignment.layout && isa(value)) { + if (failed(request(value, layout))) + return failure(); + if (assignment.layout == layout) + return propagateOperandFact(operand, layout); + } + if (failed(addUseConflict(operand, assignment, layout))) + return failure(); + if (getOperandLayout(operand) != layout) + return success(); + return propagateOperandFact(operand, layout); +} + +LogicalResult VMILayoutPropagator::propagateFact(Value value, + VMILayoutAttr layout) { + if (!isLayoutValue(value)) + return success(); + + auto equivalentIt = equivalentValues.find(value); + if (equivalentIt != equivalentValues.end()) + for (Value equivalent : equivalentIt->second) + if (failed(request(equivalent, layout))) + return failure(); + + if (auto result = dyn_cast(value)) { + if (failed(propagateThrough(result.getDefiningOp(), value, layout))) + return failure(); + } + + SmallVector uses; + for (OpOperand &use : value.getUses()) + uses.push_back(&use); + for (OpOperand *use : uses) + if (getOperandLayout(*use) == layout && + failed(propagateThrough(use->getOwner(), value, layout))) + return failure(); + + return success(); +} + +LogicalResult VMILayoutPropagator::propagateOperandFact(OpOperand &operand, + VMILayoutAttr layout) { + OperandLayoutFact fact(&operand, layout); + if (llvm::is_contained(seenOperandFacts, fact)) + return success(); + seenOperandFacts.push_back(fact); + return propagateThrough(operand.getOwner(), operand.get(), layout, &operand); +} + +LogicalResult VMILayoutPropagator::propagateThrough( + Operation *op, Value changedValue, VMILayoutAttr changedLayout, + OpOperand *changedOperand) { + const VMILayoutTransfer *transfer = op ? getTransfer(op) : nullptr; + if (!transfer) + return success(); + return transfer->propagate(op, changedValue, changedLayout, *this, + changedOperand); +} + +LogicalResult VMILayoutPropagator::run() { + while (!worklist.empty()) { + LayoutFact fact = worklist.pop_back_val(); + if (failed(propagateFact(fact.first, fact.second))) + return failure(); + } + return success(); +} + +LogicalResult VMILayoutPropagator::verifyMaterializationPlan() const { + for (Value value : orderedValues) { + auto it = assignments.find(value); + if (it == assignments.end()) + continue; + + const VMIValueLayoutAssignment &assignment = it->second; + if (!assignment.layout) + return emitError(value.getLoc()) + << kVMIDiagLayoutContractPrefix + << "layout assignment has conflicts but no primary layout"; + + VMILayoutAttr currentLayout = getCurrentLayout(value); + if (currentLayout != assignment.layout && !isTypeRewriteable(value)) { + if (!currentLayout) + return emitError(value.getLoc()) + << kVMIDiagLayoutContractPrefix + << "cannot materialize primary VMI layout from an unassigned " + "boundary value"; + if (!canMaterializeLayout(value, currentLayout, assignment.layout)) + return emitError(value.getLoc()) + << kVMIDiagLayoutContractPrefix + << "cannot materialize primary VMI layout " + << assignment.layout << " from " << currentLayout; + } + + for (const VMILayoutConflict &conflict : assignment.conflicts) { + if (!canMaterializeLayout(value, assignment.layout, conflict.layout)) + return emitError(conflict.operand ? conflict.operand->getOwner()->getLoc() + : value.getLoc()) + << kVMIDiagLayoutContractPrefix + << "cannot materialize requested VMI layout " + << conflict.layout << " from assigned layout " + << assignment.layout; + } + } + return success(); +} + +bool VMILayoutPropagator::isTypeRewriteable(Value value) const { + auto result = dyn_cast(value); + if (result) { + Operation *definingOp = result.getDefiningOp(); + if (!definingOp || !scope) + return false; + return definingOp == scope || scope->isAncestor(definingOp); + } + + auto arg = dyn_cast(value); + if (!arg || !scope) + return false; + Operation *parentOp = arg.getOwner()->getParentOp(); + return parentOp && (parentOp == scope || scope->isAncestor(parentOp)); +} + +FailureOr VMILayoutPropagator::materializeAt(Value source, + VMILayoutAttr layout, + RewriterBase &rewriter, + Location loc) { + VMILayoutAttr sourceLayout = getCurrentLayout(source); + if (!sourceLayout) + return failure(); + if (sourceLayout == layout) + return source; + + Type resultType = getTypeWithLayout(source, layout); + if (!resultType) + return failure(); + if (isa(source.getType())) + return rewriter.create(loc, resultType, source) + .getResult(); + if (isa(source.getType())) + return rewriter.create(loc, resultType, source) + .getResult(); + return failure(); +} + +LogicalResult VMILayoutPropagator::materializePrimary( + Value value, const VMIValueLayoutAssignment &assignment, + RewriterBase &rewriter, DenseMap &assignedValues) { + auto sourceType = dyn_cast(value.getType()); + auto sourceMaskType = dyn_cast(value.getType()); + if (!sourceType && !sourceMaskType) + return success(); + + Type assignedType = getTypeWithLayout(value, assignment.layout); + if (!assignedType) + return failure(); + if (getCurrentLayout(value) == assignment.layout) { + assignedValues[value] = value; + return success(); + } + + if (isTypeRewriteable(value)) { + value.setType(assignedType); + assignedValues[value] = value; + return success(); + } + + if (!getCurrentLayout(value)) { + Operation *owner = scope; + if (auto result = dyn_cast(value)) + owner = result.getDefiningOp(); + else if (auto arg = dyn_cast(value)) + owner = arg.getOwner()->getParentOp(); + if (!owner) + return failure(); + return owner->emitError() + << kVMIDiagLayoutContractPrefix + << "cannot materialize a primary VMI layout from an unassigned " + "boundary value"; + } + + OpBuilder::InsertionGuard guard(rewriter); + if (auto result = dyn_cast(value)) { + rewriter.setInsertionPointAfter(result.getDefiningOp()); + } else if (auto arg = dyn_cast(value)) { + rewriter.setInsertionPointToStart(arg.getOwner()); + } else { + return failure(); + } + + FailureOr materialized = + materializeAt(value, assignment.layout, rewriter, value.getLoc()); + if (failed(materialized)) + return failure(); + + if (*materialized != value) + value.replaceAllUsesExcept(*materialized, + (*materialized).getDefiningOp()); + assignedValues[value] = *materialized; + return success(); +} + +LogicalResult VMILayoutPropagator::materializeUseConflict( + Value assignedValue, VMILayoutConflict conflict, RewriterBase &rewriter) { + if (!conflict.operand) + return success(); + if (getCurrentLayout(assignedValue) == conflict.layout) { + conflict.operand->set(assignedValue); + return success(); + } + + OpBuilder::InsertionGuard guard(rewriter); + Operation *owner = conflict.operand->getOwner(); + rewriter.setInsertionPoint(owner); + FailureOr materialized = + materializeAt(assignedValue, conflict.layout, rewriter, owner->getLoc()); + if (failed(materialized)) + return owner->emitError() + << kVMIDiagLayoutContractPrefix + << "cannot materialize requested VMI operand layout " + << conflict.layout; + conflict.operand->set(*materialized); + return success(); +} + +LogicalResult VMILayoutPropagator::apply(RewriterBase &rewriter) { + DenseMap assignedValues; + for (Value value : orderedValues) { + auto it = assignments.find(value); + if (it == assignments.end()) + continue; + if (failed(materializePrimary(value, it->second, rewriter, + assignedValues))) + return failure(); + } + + for (Value value : orderedValues) { + auto it = assignments.find(value); + if (it == assignments.end()) + continue; + Value assignedValue = assignedValues.lookup(value); + if (!assignedValue) + assignedValue = value; + for (VMILayoutConflict conflict : it->second.conflicts) { + if (failed(materializeUseConflict(assignedValue, conflict, rewriter))) + return failure(); + } + } + return success(); +} diff --git a/lib/PTO/Transforms/VMILayoutRematerialize.cpp b/lib/PTO/Transforms/VMILayoutRematerialize.cpp new file mode 100644 index 0000000000..572253ac4f --- /dev/null +++ b/lib/PTO/Transforms/VMILayoutRematerialize.cpp @@ -0,0 +1,392 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +//===- VMILayoutRematerialize.cpp - Rematerialize VMI producers ----------===// +//===----------------------------------------------------------------------===// + +#include "PTO/IR/PTO.h" +#include "PTO/IR/PTOTypeUtils.h" +#include "PTO/Transforms/Passes.h" +#include "PTO/Transforms/VMILayoutSupport.h" + +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" +#include "llvm/ADT/STLExtras.h" + +#include + +namespace mlir { +namespace pto { +#define GEN_PASS_DEF_VMILAYOUTREMATERIALIZE +#include "PTO/Transforms/Passes.h.inc" +} // namespace pto +} // namespace mlir + +using namespace mlir; +using namespace mlir::pto; + +namespace { + +static bool hasConcreteLayout(VMIVRegType type) { + return type && static_cast(type.getLayoutAttr()); +} + +static bool hasConcreteLayout(VMIMaskType type) { + return type && static_cast(type.getLayoutAttr()); +} + +static Value materializeDataLayout(Value value, VMIVRegType resultType, + Location loc, OpBuilder &builder) { + auto sourceType = dyn_cast(value.getType()); + if (!sourceType || sourceType == resultType) + return value; + + return builder.create(loc, resultType, value).getResult(); +} + +template +static std::optional +rematerializeWidenExt(ExtOp op, VMIVRegType resultType, Location loc, + OpBuilder &builder) { + auto sourceType = dyn_cast(op.getSource().getType()); + if (!sourceType || !hasConcreteLayout(resultType)) + return std::nullopt; + + VMILayoutSupport supports; + FailureOr sourceLayout = + supports.getWidenSourceLayoutForResultLayout(sourceType, resultType, + resultType.getLayoutAttr()); + if (failed(sourceLayout)) + return std::nullopt; + + auto rematSourceType = + VMIVRegType::get(sourceType.getContext(), sourceType.getElementCount(), + sourceType.getElementType(), *sourceLayout); + if (sourceType != rematSourceType && + failed(supports.getEnsureLayoutFact(sourceType, rematSourceType))) + return std::nullopt; + Value rematSource = + materializeDataLayout(op.getSource(), rematSourceType, loc, builder); + return builder.create(loc, resultType, rematSource).getResult(); +} + +static std::optional rematerializeBinaryDataOp(Operation *op, + VMIVRegType resultType, + Location loc, + OpBuilder &builder) { + auto rebuild = [&](auto typedOp) -> std::optional { + auto lhsType = dyn_cast(typedOp.getLhs().getType()); + auto rhsType = dyn_cast(typedOp.getRhs().getType()); + if (!lhsType || !rhsType) + return std::nullopt; + auto lhsResultType = + VMIVRegType::get(lhsType.getContext(), lhsType.getElementCount(), + lhsType.getElementType(), resultType.getLayoutAttr()); + auto rhsResultType = + VMIVRegType::get(rhsType.getContext(), rhsType.getElementCount(), + rhsType.getElementType(), resultType.getLayoutAttr()); + Value lhs = + materializeDataLayout(typedOp.getLhs(), lhsResultType, loc, builder); + Value rhs = + materializeDataLayout(typedOp.getRhs(), rhsResultType, loc, builder); + return builder + .create>(loc, resultType, lhs, rhs) + .getResult(); + }; + + if (auto addf = dyn_cast(op)) + return rebuild(addf); + if (auto addi = dyn_cast(op)) + return rebuild(addi); + if (auto subf = dyn_cast(op)) + return rebuild(subf); + if (auto subi = dyn_cast(op)) + return rebuild(subi); + if (auto mulf = dyn_cast(op)) + return rebuild(mulf); + if (auto muli = dyn_cast(op)) + return rebuild(muli); + if (auto divf = dyn_cast(op)) + return rebuild(divf); + if (auto minf = dyn_cast(op)) + return rebuild(minf); + if (auto maxf = dyn_cast(op)) + return rebuild(maxf); + if (auto andi = dyn_cast(op)) + return rebuild(andi); + if (auto ori = dyn_cast(op)) + return rebuild(ori); + if (auto xori = dyn_cast(op)) + return rebuild(xori); + if (auto shli = dyn_cast(op)) + return rebuild(shli); + if (auto shrui = dyn_cast(op)) + return rebuild(shrui); + if (auto shrsi = dyn_cast(op)) + return rebuild(shrsi); + return std::nullopt; +} + +static std::optional rematerializeUnaryDataOp(Operation *op, + VMIVRegType resultType, + Location loc, + OpBuilder &builder) { + auto rebuild = [&](auto typedOp) -> std::optional { + auto sourceType = dyn_cast(typedOp.getSource().getType()); + if (!sourceType) + return std::nullopt; + auto sourceResultType = VMIVRegType::get( + sourceType.getContext(), sourceType.getElementCount(), + sourceType.getElementType(), resultType.getLayoutAttr()); + Value source = materializeDataLayout(typedOp.getSource(), sourceResultType, + loc, builder); + return builder + .create>(loc, resultType, source) + .getResult(); + }; + + if (auto negf = dyn_cast(op)) + return rebuild(negf); + if (auto absf = dyn_cast(op)) + return rebuild(absf); + if (auto absi = dyn_cast(op)) + return rebuild(absi); + if (auto sqrt = dyn_cast(op)) + return rebuild(sqrt); + if (auto exp = dyn_cast(op)) + return rebuild(exp); + if (auto ln = dyn_cast(op)) + return rebuild(ln); + if (auto relu = dyn_cast(op)) + return rebuild(relu); + if (auto notOp = dyn_cast(op)) + return rebuild(notOp); + return std::nullopt; +} + +static std::optional rematerializeFma(VMIFmaOp fma, + VMIVRegType resultType, + Location loc, OpBuilder &builder) { + auto lhsType = dyn_cast(fma.getLhs().getType()); + auto rhsType = dyn_cast(fma.getRhs().getType()); + auto accType = dyn_cast(fma.getAcc().getType()); + if (!lhsType || !rhsType || !accType) + return std::nullopt; + auto makeType = [&](VMIVRegType type) { + return VMIVRegType::get(type.getContext(), type.getElementCount(), + type.getElementType(), resultType.getLayoutAttr()); + }; + Value lhs = + materializeDataLayout(fma.getLhs(), makeType(lhsType), loc, builder); + Value rhs = + materializeDataLayout(fma.getRhs(), makeType(rhsType), loc, builder); + Value acc = + materializeDataLayout(fma.getAcc(), makeType(accType), loc, builder); + return builder.create(loc, resultType, lhs, rhs, acc).getResult(); +} + +static std::optional rematerializeDataProducer(Value value, + VMIVRegType resultType, + Location loc, + OpBuilder &builder) { + if (!hasConcreteLayout(resultType)) + return std::nullopt; + + if (auto extf = value.getDefiningOp()) + return rematerializeWidenExt(extf, resultType, loc, builder); + if (auto extsi = value.getDefiningOp()) + return rematerializeWidenExt(extsi, resultType, loc, builder); + if (auto extui = value.getDefiningOp()) + return rematerializeWidenExt(extui, resultType, loc, builder); + + if (Operation *op = value.getDefiningOp()) { + if (auto fma = dyn_cast(op)) + return rematerializeFma(fma, resultType, loc, builder); + if (auto result = rematerializeBinaryDataOp(op, resultType, loc, builder)) + return result; + if (auto result = rematerializeUnaryDataOp(op, resultType, loc, builder)) + return result; + } + + if (auto constant = value.getDefiningOp()) { + auto denseAttr = dyn_cast(constant.getValue()); + if (denseAttr && denseAttr.isSplat()) + return builder.create(loc, resultType, constant.getValue()) + .getResult(); + } + + if (auto broadcast = value.getDefiningOp()) + return builder.create(loc, resultType, broadcast.getValue()) + .getResult(); + + if (auto iota = value.getDefiningOp()) + return builder + .create(loc, resultType, iota.getBase(), iota.getOrderAttr()) + .getResult(); + + return std::nullopt; +} + +static std::optional rematerializeMaskProducer(Value value, + VMIMaskType resultType, + Location loc, + OpBuilder &builder) { + if (!hasConcreteLayout(resultType)) + return std::nullopt; + + if (auto createMask = value.getDefiningOp()) + return builder + .create(loc, resultType, createMask.getActiveLanes()) + .getResult(); + + if (auto createGroupMask = value.getDefiningOp()) { + return builder + .create(loc, resultType, + createGroupMask.getActiveElemsPerGroup(), + createGroupMask.getNumGroupsAttr(), + createGroupMask.getGroupSizeAttr()) + .getResult(); + } + + if (auto constantMask = value.getDefiningOp()) + return builder + .create(loc, resultType, constantMask.getValueAttr()) + .getResult(); + + return std::nullopt; +} + +static bool tryReplaceDataEnsure(VMIEnsureLayoutOp ensure) { + auto resultType = dyn_cast(ensure.getResult().getType()); + if (!resultType) + return false; + + OpBuilder builder(ensure); + auto result = rematerializeDataProducer(ensure.getSource(), resultType, + ensure->getLoc(), builder); + if (!result) + return false; + + ensure.getResult().replaceAllUsesWith(*result); + ensure.erase(); + return true; +} + +static bool tryRematerializeTruncIThroughSourceEnsure(VMITruncIOp trunc) { + auto resultType = dyn_cast(trunc.getResult().getType()); + if (!resultType || !hasConcreteLayout(resultType)) + return false; + + auto ensure = trunc.getSource().getDefiningOp(); + if (!ensure) + return false; + + auto originalSourceType = dyn_cast(ensure.getSource().getType()); + if (!originalSourceType || !hasConcreteLayout(originalSourceType)) + return false; + VMILayoutAttr originalSourceLayout = originalSourceType.getLayoutAttr(); + if (!originalSourceLayout.isDeinterleaved() || + originalSourceLayout.getBlockElems() != 1) + return false; + + VMILayoutSupport supports; + FailureOr fact = supports.getCastLayoutFactForSourceLayout( + originalSourceType, resultType, originalSourceLayout); + if (failed(fact)) + return false; + + unsigned resultBits = + pto::getPTOStorageElemBitWidth(resultType.getElementType()); + if (resultBits == 8 && + !cast(resultType.getElementType()).isUnsigned()) + return false; + + VMILayoutAttr rematResultLayout = fact->resultLayout; + auto rematResultType = + VMIVRegType::get(resultType.getContext(), resultType.getElementCount(), + resultType.getElementType(), rematResultLayout); + if (rematResultType == resultType) + return false; + + OpBuilder builder(trunc); + Value remat = builder + .create(trunc->getLoc(), rematResultType, + ensure.getSource()) + .getResult(); + Value replacement = + materializeDataLayout(remat, resultType, trunc->getLoc(), builder); + trunc.getResult().replaceAllUsesWith(replacement); + trunc.erase(); + return true; +} + +template static bool tryReplaceMaskEnsure(EnsureOp ensure) { + auto resultType = dyn_cast(ensure.getResult().getType()); + if (!resultType) + return false; + + OpBuilder builder(ensure); + auto result = rematerializeMaskProducer(ensure.getSource(), resultType, + ensure->getLoc(), builder); + if (!result) + return false; + + ensure.getResult().replaceAllUsesWith(*result); + ensure.erase(); + return true; +} + +struct VMILayoutRematerializePass + : public mlir::pto::impl::VMILayoutRematerializeBase< + VMILayoutRematerializePass> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(VMILayoutRematerializePass) + + void runOnOperation() override { + ModuleOp module = getOperation(); + bool changed = true; + while (changed) { + changed = false; + SmallVector helpers; + module.walk([&](Operation *op) { + if (isa(op)) + helpers.push_back(op); + }); + + for (Operation *op : helpers) { + if (op->getBlock() == nullptr) + continue; + + if (auto ensure = dyn_cast(op)) { + changed |= tryReplaceDataEnsure(ensure); + continue; + } + + if (auto ensure = dyn_cast(op)) { + changed |= tryReplaceMaskEnsure(ensure); + continue; + } + + if (auto ensure = dyn_cast(op)) + changed |= tryReplaceMaskEnsure(ensure); + + if (auto trunc = dyn_cast(op)) + changed |= tryRematerializeTruncIThroughSourceEnsure(trunc); + } + } + } +}; + +} // namespace + +std::unique_ptr mlir::pto::createVMILayoutRematerializePass() { + return std::make_unique(); +} diff --git a/lib/PTO/Transforms/VMILayoutSinkMaterialization.cpp b/lib/PTO/Transforms/VMILayoutSinkMaterialization.cpp new file mode 100644 index 0000000000..d3719042b3 --- /dev/null +++ b/lib/PTO/Transforms/VMILayoutSinkMaterialization.cpp @@ -0,0 +1,627 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +//===- VMILayoutSinkMaterialization.cpp - Sink VMI layout helpers --------===// +//===----------------------------------------------------------------------===// + +#include "PTO/IR/PTO.h" +#include "PTO/Transforms/Passes.h" +#include "PTO/Transforms/VMILayoutSupport.h" + +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/Operation.h" +#include "mlir/Pass/Pass.h" +#include "llvm/ADT/STLExtras.h" + +namespace mlir { +namespace pto { +#define GEN_PASS_DEF_VMILAYOUTSINKMATERIALIZATION +#include "PTO/Transforms/Passes.h.inc" +} // namespace pto +} // namespace mlir + +using namespace mlir; +using namespace mlir::pto; + +namespace { + +struct BinaryVRegOperands { + OpOperand *lhs = nullptr; + OpOperand *rhs = nullptr; +}; + +struct TernaryVRegOperands { + OpOperand *lhs = nullptr; + OpOperand *rhs = nullptr; + OpOperand *acc = nullptr; +}; + +struct SelectOperands { + OpOperand *mask = nullptr; + OpOperand *trueValue = nullptr; + OpOperand *falseValue = nullptr; +}; + +struct UnaryVRegOperand { + OpOperand *source = nullptr; +}; + +struct BinaryMaskOperands { + OpOperand *lhs = nullptr; + OpOperand *rhs = nullptr; +}; + +struct UnaryMaskOperand { + OpOperand *source = nullptr; +}; + +static std::optional +getSinkableBinaryOperands(Operation *op) { + if (auto addf = dyn_cast(op)) + return BinaryVRegOperands{&addf.getLhsMutable(), &addf.getRhsMutable()}; + if (auto addi = dyn_cast(op)) + return BinaryVRegOperands{&addi.getLhsMutable(), &addi.getRhsMutable()}; + if (auto subf = dyn_cast(op)) + return BinaryVRegOperands{&subf.getLhsMutable(), &subf.getRhsMutable()}; + if (auto subi = dyn_cast(op)) + return BinaryVRegOperands{&subi.getLhsMutable(), &subi.getRhsMutable()}; + if (auto mulf = dyn_cast(op)) + return BinaryVRegOperands{&mulf.getLhsMutable(), &mulf.getRhsMutable()}; + if (auto muli = dyn_cast(op)) + return BinaryVRegOperands{&muli.getLhsMutable(), &muli.getRhsMutable()}; + if (auto divf = dyn_cast(op)) + return BinaryVRegOperands{&divf.getLhsMutable(), &divf.getRhsMutable()}; + if (auto minf = dyn_cast(op)) + return BinaryVRegOperands{&minf.getLhsMutable(), &minf.getRhsMutable()}; + if (auto maxf = dyn_cast(op)) + return BinaryVRegOperands{&maxf.getLhsMutable(), &maxf.getRhsMutable()}; + if (auto andi = dyn_cast(op)) + return BinaryVRegOperands{&andi.getLhsMutable(), &andi.getRhsMutable()}; + if (auto ori = dyn_cast(op)) + return BinaryVRegOperands{&ori.getLhsMutable(), &ori.getRhsMutable()}; + if (auto xori = dyn_cast(op)) + return BinaryVRegOperands{&xori.getLhsMutable(), &xori.getRhsMutable()}; + if (auto shli = dyn_cast(op)) + return BinaryVRegOperands{&shli.getLhsMutable(), &shli.getRhsMutable()}; + if (auto shrui = dyn_cast(op)) + return BinaryVRegOperands{&shrui.getLhsMutable(), &shrui.getRhsMutable()}; + if (auto shrsi = dyn_cast(op)) + return BinaryVRegOperands{&shrsi.getLhsMutable(), &shrsi.getRhsMutable()}; + return std::nullopt; +} + +static std::optional +getSinkableCompareOperands(Operation *op) { + if (auto cmpf = dyn_cast(op)) + return BinaryVRegOperands{&cmpf.getLhsMutable(), &cmpf.getRhsMutable()}; + if (auto cmpi = dyn_cast(op)) + return BinaryVRegOperands{&cmpi.getLhsMutable(), &cmpi.getRhsMutable()}; + return std::nullopt; +} + +static std::optional getSinkableSelectOperands(Operation *op) { + if (auto select = dyn_cast(op)) + return SelectOperands{&select.getMaskMutable(), + &select.getTrueValueMutable(), + &select.getFalseValueMutable()}; + return std::nullopt; +} + +static std::optional +getSinkableTernaryOperands(Operation *op) { + if (auto fma = dyn_cast(op)) + return TernaryVRegOperands{&fma.getLhsMutable(), &fma.getRhsMutable(), + &fma.getAccMutable()}; + return std::nullopt; +} + +static std::optional getSinkableUnaryOperand(Operation *op) { + if (auto negf = dyn_cast(op)) + return UnaryVRegOperand{&negf.getSourceMutable()}; + if (auto absf = dyn_cast(op)) + return UnaryVRegOperand{&absf.getSourceMutable()}; + if (auto absi = dyn_cast(op)) + return UnaryVRegOperand{&absi.getSourceMutable()}; + if (auto sqrt = dyn_cast(op)) + return UnaryVRegOperand{&sqrt.getSourceMutable()}; + if (auto exp = dyn_cast(op)) + return UnaryVRegOperand{&exp.getSourceMutable()}; + if (auto ln = dyn_cast(op)) + return UnaryVRegOperand{&ln.getSourceMutable()}; + if (auto relu = dyn_cast(op)) + return UnaryVRegOperand{&relu.getSourceMutable()}; + if (auto notOp = dyn_cast(op)) + return UnaryVRegOperand{¬Op.getSourceMutable()}; + return std::nullopt; +} + +static std::optional +getSinkableBinaryMaskOperands(Operation *op) { + if (auto maskAnd = dyn_cast(op)) + return BinaryMaskOperands{&maskAnd.getLhsMutable(), + &maskAnd.getRhsMutable()}; + if (auto maskOr = dyn_cast(op)) + return BinaryMaskOperands{&maskOr.getLhsMutable(), &maskOr.getRhsMutable()}; + if (auto maskXor = dyn_cast(op)) + return BinaryMaskOperands{&maskXor.getLhsMutable(), + &maskXor.getRhsMutable()}; + return std::nullopt; +} + +static std::optional +getSinkableUnaryMaskOperand(Operation *op) { + if (auto maskNot = dyn_cast(op)) + return UnaryMaskOperand{&maskNot.getSourceMutable()}; + return std::nullopt; +} + +static bool isSameMaterialization(VMIEnsureLayoutOp ensure, + VMIVRegType resultType) { + if (!ensure || !resultType) + return false; + + auto sourceType = dyn_cast(ensure.getSource().getType()); + auto ensureResultType = dyn_cast(ensure.getResult().getType()); + if (!sourceType || !ensureResultType) + return false; + + return ensureResultType == resultType && sourceType != resultType; +} + +static bool isSameMaterialization(VMIEnsureLayoutOp lhsEnsure, + VMIEnsureLayoutOp rhsEnsure, + VMIVRegType resultType) { + if (!lhsEnsure || !rhsEnsure || !resultType) + return false; + + auto lhsSourceType = dyn_cast(lhsEnsure.getSource().getType()); + auto rhsSourceType = dyn_cast(rhsEnsure.getSource().getType()); + auto lhsResultType = dyn_cast(lhsEnsure.getResult().getType()); + auto rhsResultType = dyn_cast(rhsEnsure.getResult().getType()); + if (!lhsSourceType || !rhsSourceType || !lhsResultType || !rhsResultType) + return false; + + return lhsSourceType == rhsSourceType && lhsResultType == rhsResultType && + lhsResultType == resultType && lhsSourceType != resultType; +} + +static bool isSameMaterialization(VMIEnsureLayoutOp lhsEnsure, + VMIEnsureLayoutOp rhsEnsure, + VMIEnsureLayoutOp accEnsure, + VMIVRegType resultType) { + if (!lhsEnsure || !rhsEnsure || !accEnsure || !resultType) + return false; + + auto lhsSourceType = dyn_cast(lhsEnsure.getSource().getType()); + auto rhsSourceType = dyn_cast(rhsEnsure.getSource().getType()); + auto accSourceType = dyn_cast(accEnsure.getSource().getType()); + auto lhsResultType = dyn_cast(lhsEnsure.getResult().getType()); + auto rhsResultType = dyn_cast(rhsEnsure.getResult().getType()); + auto accResultType = dyn_cast(accEnsure.getResult().getType()); + if (!lhsSourceType || !rhsSourceType || !accSourceType || !lhsResultType || + !rhsResultType || !accResultType) + return false; + + return lhsSourceType == rhsSourceType && lhsSourceType == accSourceType && + lhsResultType == rhsResultType && lhsResultType == accResultType && + lhsResultType == resultType && lhsSourceType != resultType; +} + +static bool hasEnsureLayoutSupport(VMIVRegType sourceType, + VMIVRegType resultType) { + VMILayoutSupport supports; + return succeeded(supports.getEnsureLayoutFact(sourceType, resultType)); +} + +template +static bool isSameMaskMaterialization(EnsureOp ensure, VMIMaskType resultType) { + if (!ensure || !resultType) + return false; + + auto sourceType = dyn_cast(ensure.getSource().getType()); + auto ensureResultType = dyn_cast(ensure.getResult().getType()); + if (!sourceType || !ensureResultType) + return false; + + return ensureResultType == resultType && sourceType != resultType; +} + +template +static bool isSameMaskMaterialization(EnsureOp lhsEnsure, EnsureOp rhsEnsure, + VMIMaskType resultType) { + if (!lhsEnsure || !rhsEnsure || !resultType) + return false; + + auto lhsSourceType = dyn_cast(lhsEnsure.getSource().getType()); + auto rhsSourceType = dyn_cast(rhsEnsure.getSource().getType()); + auto lhsResultType = dyn_cast(lhsEnsure.getResult().getType()); + auto rhsResultType = dyn_cast(rhsEnsure.getResult().getType()); + if (!lhsSourceType || !rhsSourceType || !lhsResultType || !rhsResultType) + return false; + + return lhsSourceType == rhsSourceType && lhsResultType == rhsResultType && + lhsResultType == resultType && lhsSourceType != resultType; +} + +static bool hasEnsureMaskSupport(VMIEnsureMaskLayoutOp, VMIMaskType sourceType, + VMIMaskType resultType) { + VMILayoutSupport supports; + return succeeded(supports.getEnsureMaskLayoutFact(sourceType, resultType)); +} + +static bool hasEnsureMaskSupport(VMIEnsureMaskGranularityOp, + VMIMaskType sourceType, + VMIMaskType resultType) { + return sourceType.getElementCount() == resultType.getElementCount() && + sourceType.getLayoutAttr() == resultType.getLayoutAttr() && + !sourceType.isPred() && !resultType.isPred(); +} + +static bool trySinkBinaryMaterialization(Operation *op) { + std::optional operands = getSinkableBinaryOperands(op); + if (!operands || op->getNumResults() != 1) + return false; + + auto resultType = dyn_cast(op->getResult(0).getType()); + if (!resultType) + return false; + + auto lhsEnsure = operands->lhs->get().getDefiningOp(); + auto rhsEnsure = operands->rhs->get().getDefiningOp(); + if (!isSameMaterialization(lhsEnsure, rhsEnsure, resultType)) + return false; + + auto sourceType = cast(lhsEnsure.getSource().getType()); + if (!hasEnsureLayoutSupport(sourceType, resultType)) + return false; + + OpBuilder builder(op); + OperationState state(op->getLoc(), op->getName()); + state.addOperands({lhsEnsure.getSource(), rhsEnsure.getSource()}); + state.addTypes(sourceType); + state.addAttributes(op->getAttrs()); + Operation *newOp = builder.create(state); + + builder.setInsertionPointAfter(newOp); + auto resultEnsure = builder.create( + op->getLoc(), resultType, newOp->getResult(0)); + op->getResult(0).replaceAllUsesWith(resultEnsure.getResult()); + op->erase(); + + if (lhsEnsure->use_empty()) + lhsEnsure.erase(); + if (rhsEnsure != lhsEnsure && rhsEnsure->use_empty()) + rhsEnsure.erase(); + return true; +} + +static bool trySinkSelectMaterialization(Operation *op) { + std::optional operands = getSinkableSelectOperands(op); + if (!operands || op->getNumResults() != 1) + return false; + + auto resultType = dyn_cast(op->getResult(0).getType()); + if (!resultType) + return false; + + auto maskEnsure = + operands->mask->get().getDefiningOp(); + auto trueEnsure = + operands->trueValue->get().getDefiningOp(); + auto falseEnsure = + operands->falseValue->get().getDefiningOp(); + if (!maskEnsure || !trueEnsure || !falseEnsure) + return false; + + auto trueSourceType = dyn_cast(trueEnsure.getSource().getType()); + auto falseSourceType = + dyn_cast(falseEnsure.getSource().getType()); + auto trueResultType = dyn_cast(trueEnsure.getResult().getType()); + auto falseResultType = + dyn_cast(falseEnsure.getResult().getType()); + auto maskSourceType = dyn_cast(maskEnsure.getSource().getType()); + auto maskResultType = dyn_cast(maskEnsure.getResult().getType()); + if (!trueSourceType || !falseSourceType || !trueResultType || + !falseResultType || !maskSourceType || !maskResultType) + return false; + + if (trueSourceType != falseSourceType || trueResultType != falseResultType || + trueResultType != resultType || trueSourceType == resultType) + return false; + if (maskResultType != operands->mask->get().getType()) + return false; + if (maskResultType.getLayoutAttr() != resultType.getLayoutAttr() || + maskSourceType.getLayoutAttr() != trueSourceType.getLayoutAttr()) + return false; + if (maskSourceType.getElementCount() != trueSourceType.getElementCount() || + maskResultType.getElementCount() != resultType.getElementCount() || + maskSourceType.getGranularity() != maskResultType.getGranularity()) + return false; + if (!hasEnsureLayoutSupport(trueSourceType, resultType) || + !hasEnsureMaskSupport(maskEnsure, maskSourceType, maskResultType)) + return false; + + OpBuilder builder(op); + OperationState state(op->getLoc(), op->getName()); + state.addOperands({maskEnsure.getSource(), trueEnsure.getSource(), + falseEnsure.getSource()}); + state.addTypes(trueSourceType); + state.addAttributes(op->getAttrs()); + Operation *newOp = builder.create(state); + + builder.setInsertionPointAfter(newOp); + auto resultEnsure = builder.create( + op->getLoc(), resultType, newOp->getResult(0)); + op->getResult(0).replaceAllUsesWith(resultEnsure.getResult()); + op->erase(); + + if (maskEnsure->use_empty()) + maskEnsure.erase(); + if (trueEnsure->use_empty()) + trueEnsure.erase(); + if (falseEnsure != trueEnsure && falseEnsure->use_empty()) + falseEnsure.erase(); + return true; +} + +static bool trySinkCompareMaterialization(Operation *op) { + std::optional operands = getSinkableCompareOperands(op); + if (!operands || op->getNumResults() != 1) + return false; + + auto resultMaskType = dyn_cast(op->getResult(0).getType()); + if (!resultMaskType) + return false; + + auto lhsEnsure = operands->lhs->get().getDefiningOp(); + auto rhsEnsure = operands->rhs->get().getDefiningOp(); + if (!lhsEnsure || !rhsEnsure) + return false; + + auto lhsSourceType = dyn_cast(lhsEnsure.getSource().getType()); + auto rhsSourceType = dyn_cast(rhsEnsure.getSource().getType()); + auto lhsResultType = dyn_cast(lhsEnsure.getResult().getType()); + auto rhsResultType = dyn_cast(rhsEnsure.getResult().getType()); + if (!lhsSourceType || !rhsSourceType || !lhsResultType || !rhsResultType) + return false; + if (lhsSourceType != rhsSourceType || lhsResultType != rhsResultType || + lhsSourceType == lhsResultType) + return false; + if (lhsResultType.getElementCount() != resultMaskType.getElementCount() || + lhsResultType.getLayoutAttr() != resultMaskType.getLayoutAttr()) + return false; + + auto sourceMaskType = VMIMaskType::get( + op->getContext(), resultMaskType.getElementCount(), + resultMaskType.getGranularity(), lhsSourceType.getLayoutAttr()); + VMILayoutSupport supports; + if (failed(supports.getEnsureMaskLayoutFact(sourceMaskType, resultMaskType))) + return false; + + OpBuilder builder(op); + OperationState state(op->getLoc(), op->getName()); + state.addOperands({lhsEnsure.getSource(), rhsEnsure.getSource()}); + state.addTypes(sourceMaskType); + state.addAttributes(op->getAttrs()); + Operation *newOp = builder.create(state); + + builder.setInsertionPointAfter(newOp); + auto resultEnsure = builder.create( + op->getLoc(), resultMaskType, newOp->getResult(0)); + op->getResult(0).replaceAllUsesWith(resultEnsure.getResult()); + op->erase(); + + if (lhsEnsure->use_empty()) + lhsEnsure.erase(); + if (rhsEnsure != lhsEnsure && rhsEnsure->use_empty()) + rhsEnsure.erase(); + return true; +} + +static bool trySinkTernaryMaterialization(Operation *op) { + std::optional operands = getSinkableTernaryOperands(op); + if (!operands || op->getNumResults() != 1) + return false; + + auto resultType = dyn_cast(op->getResult(0).getType()); + if (!resultType) + return false; + + auto lhsEnsure = operands->lhs->get().getDefiningOp(); + auto rhsEnsure = operands->rhs->get().getDefiningOp(); + auto accEnsure = operands->acc->get().getDefiningOp(); + if (!isSameMaterialization(lhsEnsure, rhsEnsure, accEnsure, resultType)) + return false; + + auto sourceType = cast(lhsEnsure.getSource().getType()); + if (!hasEnsureLayoutSupport(sourceType, resultType)) + return false; + + OpBuilder builder(op); + OperationState state(op->getLoc(), op->getName()); + state.addOperands( + {lhsEnsure.getSource(), rhsEnsure.getSource(), accEnsure.getSource()}); + state.addTypes(sourceType); + state.addAttributes(op->getAttrs()); + Operation *newOp = builder.create(state); + + builder.setInsertionPointAfter(newOp); + auto resultEnsure = builder.create( + op->getLoc(), resultType, newOp->getResult(0)); + op->getResult(0).replaceAllUsesWith(resultEnsure.getResult()); + op->erase(); + + if (lhsEnsure->use_empty()) + lhsEnsure.erase(); + if (rhsEnsure != lhsEnsure && rhsEnsure->use_empty()) + rhsEnsure.erase(); + if (accEnsure != lhsEnsure && accEnsure != rhsEnsure && + accEnsure->use_empty()) + accEnsure.erase(); + return true; +} + +template +static bool trySinkBinaryMaskMaterialization(Operation *op) { + std::optional operands = + getSinkableBinaryMaskOperands(op); + if (!operands || op->getNumResults() != 1) + return false; + + auto resultType = dyn_cast(op->getResult(0).getType()); + if (!resultType) + return false; + + auto lhsEnsure = operands->lhs->get().getDefiningOp(); + auto rhsEnsure = operands->rhs->get().getDefiningOp(); + if (!isSameMaskMaterialization(lhsEnsure, rhsEnsure, resultType)) + return false; + + auto sourceType = cast(lhsEnsure.getSource().getType()); + if (!hasEnsureMaskSupport(lhsEnsure, sourceType, resultType)) + return false; + + OpBuilder builder(op); + OperationState state(op->getLoc(), op->getName()); + state.addOperands({lhsEnsure.getSource(), rhsEnsure.getSource()}); + state.addTypes(sourceType); + state.addAttributes(op->getAttrs()); + Operation *newOp = builder.create(state); + + builder.setInsertionPointAfter(newOp); + auto resultEnsure = + builder.create(op->getLoc(), resultType, newOp->getResult(0)); + op->getResult(0).replaceAllUsesWith(resultEnsure.getResult()); + op->erase(); + + if (lhsEnsure->use_empty()) + lhsEnsure.erase(); + if (rhsEnsure != lhsEnsure && rhsEnsure->use_empty()) + rhsEnsure.erase(); + return true; +} + +static bool trySinkUnaryMaterialization(Operation *op) { + std::optional operand = getSinkableUnaryOperand(op); + if (!operand || op->getNumResults() != 1) + return false; + + auto resultType = dyn_cast(op->getResult(0).getType()); + if (!resultType) + return false; + + auto sourceEnsure = operand->source->get().getDefiningOp(); + if (!isSameMaterialization(sourceEnsure, resultType)) + return false; + + auto sourceType = cast(sourceEnsure.getSource().getType()); + if (!hasEnsureLayoutSupport(sourceType, resultType)) + return false; + + OpBuilder builder(op); + OperationState state(op->getLoc(), op->getName()); + state.addOperands(sourceEnsure.getSource()); + state.addTypes(sourceType); + state.addAttributes(op->getAttrs()); + Operation *newOp = builder.create(state); + + builder.setInsertionPointAfter(newOp); + auto resultEnsure = builder.create( + op->getLoc(), resultType, newOp->getResult(0)); + op->getResult(0).replaceAllUsesWith(resultEnsure.getResult()); + op->erase(); + + if (sourceEnsure->use_empty()) + sourceEnsure.erase(); + return true; +} + +template +static bool trySinkUnaryMaskMaterialization(Operation *op) { + std::optional operand = getSinkableUnaryMaskOperand(op); + if (!operand || op->getNumResults() != 1) + return false; + + auto resultType = dyn_cast(op->getResult(0).getType()); + if (!resultType) + return false; + + auto sourceEnsure = operand->source->get().getDefiningOp(); + if (!isSameMaskMaterialization(sourceEnsure, resultType)) + return false; + + auto sourceType = cast(sourceEnsure.getSource().getType()); + if (!hasEnsureMaskSupport(sourceEnsure, sourceType, resultType)) + return false; + + OpBuilder builder(op); + OperationState state(op->getLoc(), op->getName()); + state.addOperands(sourceEnsure.getSource()); + state.addTypes(sourceType); + state.addAttributes(op->getAttrs()); + Operation *newOp = builder.create(state); + + builder.setInsertionPointAfter(newOp); + auto resultEnsure = + builder.create(op->getLoc(), resultType, newOp->getResult(0)); + op->getResult(0).replaceAllUsesWith(resultEnsure.getResult()); + op->erase(); + + if (sourceEnsure->use_empty()) + sourceEnsure.erase(); + return true; +} + +static bool trySinkMaskMaterialization(Operation *op) { + return trySinkBinaryMaskMaterialization(op) || + trySinkBinaryMaskMaterialization(op) || + trySinkUnaryMaskMaterialization(op) || + trySinkUnaryMaskMaterialization(op); +} + +struct VMILayoutSinkMaterializationPass + : public mlir::pto::impl::VMILayoutSinkMaterializationBase< + VMILayoutSinkMaterializationPass> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(VMILayoutSinkMaterializationPass) + + void runOnOperation() override { + ModuleOp module = getOperation(); + SmallVector candidates; + module.walk([&](Operation *op) { + if (getSinkableBinaryOperands(op) || getSinkableCompareOperands(op) || + getSinkableSelectOperands(op) || getSinkableTernaryOperands(op) || + getSinkableUnaryOperand(op) || getSinkableBinaryMaskOperands(op) || + getSinkableUnaryMaskOperand(op)) + candidates.push_back(op); + }); + + for (Operation *op : candidates) { + if (op->getBlock() == nullptr) + continue; + if (!trySinkBinaryMaterialization(op)) { + if (!trySinkCompareMaterialization(op)) { + if (!trySinkSelectMaterialization(op)) { + if (!trySinkTernaryMaterialization(op)) { + if (!trySinkUnaryMaterialization(op)) + trySinkMaskMaterialization(op); + } + } + } + } + } + } +}; + +} // namespace + +std::unique_ptr mlir::pto::createVMILayoutSinkMaterializationPass() { + return std::make_unique(); +} diff --git a/lib/PTO/Transforms/VMILayoutSupport.cpp b/lib/PTO/Transforms/VMILayoutSupport.cpp new file mode 100644 index 0000000000..5bfb2e8e5e --- /dev/null +++ b/lib/PTO/Transforms/VMILayoutSupport.cpp @@ -0,0 +1,2590 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +//===- VMILayoutSupport.cpp - VMI layout support queries --------------===// +//===----------------------------------------------------------------------===// +// +// This file is the central table-driven source for VMI layout support facts. +// Keep file-level responsibilities separated as follows: +// +// 1. Layout pattern DSL: +// Define compact syntax for expressing layouts and table keys only. +// 2. Query key derivation helpers: +// Extract op operands and derive normalized keys used to query the tables. +// Do not add new layout support facts in this section. +// 3. Rule tables: +// Add legal/preferred layout relations here. New support facts should be +// visible as table rows instead of being hidden in query helper branches. +// 4. Table matching and materialization helpers: +// Convert table rows into facts and compare derived keys with row keys. +// Do not add new layout support facts in this section. +// 5. Query implementations: +// Query functions should consume the tables, derive table keys, and check +// op-level preconditions only. Shape/type limits that define layout support +// must be visible as table keys or table rows, not hidden in query branches. +// +// When adding a new family of support rules, first extend the shared pattern +// DSL if needed, then add table rows, then expose them through a small query. +// Avoid local mini-DSLs or ad-hoc support logic that duplicates table facts. + +#include "PTO/Transforms/VMILayoutSupport.h" + +#include "PTO/IR/PTOTypeUtils.h" +#include "PTO/IR/VMIUtils.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "llvm/ADT/Twine.h" + +using namespace mlir; +using namespace mlir::pto; + +namespace { + +//===----------------------------------------------------------------------===// +// Layout pattern DSL +//===----------------------------------------------------------------------===// + +enum class LayoutPatternKind { + Contiguous, + LaneStride, + Deinterleaved, + GroupSlots, +}; + +struct LayoutPattern { + LayoutPatternKind kind = LayoutPatternKind::Contiguous; + int64_t value = 1; + int64_t blockElems = 0; + int64_t laneStride = 1; +}; + +enum class GroupBlockPatternKind { + Ratio, + FullPartMultiple, +}; + +enum class GroupMemoryPatternKind { + Any, + Contiguous, + BlockAligned, +}; + +struct GroupBlockPattern { + GroupBlockPatternKind kind = GroupBlockPatternKind::Ratio; + int64_t numerator = 1; + int64_t denominator = 1; +}; + +struct GroupMemoryPattern { + GroupMemoryPatternKind kind = GroupMemoryPatternKind::Contiguous; +}; + +struct ElementBitsPattern { + uint64_t mask = 0; +}; + +struct ElementCountPattern { + int64_t values[32] = {}; + int64_t count = 0; + bool any = false; +}; + +struct MaskGranularityPattern { + uint8_t mask = 0; +}; + +struct PhysicalChunkCountPattern { + int64_t values[4] = {}; + int64_t count = 0; +}; + +template +static constexpr PhysicalChunkCountPattern chunk() { + static_assert(sizeof...(Counts) <= 4, "too many physical chunk counts"); + static_assert(((Counts == 1 || Counts == 2 || Counts == 4) && ...), + "unsupported physical chunk count"); + return {{Counts...}, static_cast(sizeof...(Counts))}; +} + +static constexpr uint64_t elementBitsMask(int64_t bits) { + return bits == 8 ? 1ull << 0 + : bits == 16 ? 1ull << 1 + : bits == 32 ? 1ull << 2 + : bits == 64 ? 1ull << 3 + : 0; +} + +template static constexpr ElementBitsPattern bits() { + return {((uint64_t{0} | elementBitsMask(Bits)) | ...)}; +} + +template static constexpr ElementCountPattern N() { + static_assert(sizeof...(Counts) <= 32, "too many element count patterns"); + return {{Counts...}, static_cast(sizeof...(Counts)), false}; +} + +template static constexpr ElementCountPattern G() { + return N(); +} + +static constexpr ElementCountPattern anyN() { return {{}, 0, true}; } +static constexpr ElementCountPattern anyG() { return anyN(); } + +static constexpr MaskGranularityPattern mb8() { return {1u << 0}; } +static constexpr MaskGranularityPattern mb16() { return {1u << 1}; } +static constexpr MaskGranularityPattern mb32() { return {1u << 2}; } + +static bool matchesElementBitsPattern(ElementBitsPattern pattern, + int64_t bits) { + uint64_t mask = elementBitsMask(bits); + return mask != 0 && (pattern.mask & mask) != 0; +} + +static bool matchesElementBitsPattern(ElementBitsPattern pattern, + Type elementType) { + unsigned elementBits = pto::getPTOStorageElemBitWidth(elementType); + return matchesElementBitsPattern(pattern, elementBits); +} + +static bool matchesElementCountPattern(ElementCountPattern pattern, + int64_t count) { + if (pattern.any) + return true; + for (int64_t i = 0; i < pattern.count; ++i) + if (pattern.values[i] == count) + return true; + return false; +} + +static bool matchesMaskGranularityPattern(MaskGranularityPattern pattern, + StringRef granularity) { + uint8_t mask = granularity == "b8" ? 1u << 0 + : granularity == "b16" ? 1u << 1 + : granularity == "b32" ? 1u << 2 + : 0; + return mask != 0 && (pattern.mask & mask) != 0; +} + +static VMILayoutAttr materializeLayoutPattern(MLIRContext *ctx, + LayoutPattern pattern, + int64_t inheritedBlockElems = 1, + int64_t numGroups = 0) { + switch (pattern.kind) { + case LayoutPatternKind::Contiguous: + return VMILayoutAttr::getContiguous(ctx); + case LayoutPatternKind::LaneStride: + return VMILayoutAttr::getContiguous(ctx, pattern.value); + case LayoutPatternKind::Deinterleaved: { + int64_t blockElems = + pattern.blockElems > 0 ? pattern.blockElems : inheritedBlockElems; + return VMILayoutAttr::getDeinterleaved(ctx, pattern.value, blockElems); + } + case LayoutPatternKind::GroupSlots: + return numGroups > 0 + ? VMILayoutAttr::getGroupSlots(ctx, numGroups, pattern.value, + pattern.laneStride) + : VMILayoutAttr(); + } + llvm_unreachable("unknown layout pattern kind"); +} + +static bool matchesLayoutPattern(MLIRContext *ctx, LayoutPattern pattern, + VMILayoutAttr layout, int64_t numGroups = 0) { + if (!layout) + return false; + int64_t inheritedBlockElems = + layout.isDeinterleaved() ? layout.getBlockElems() : 1; + return materializeLayoutPattern(ctx, pattern, inheritedBlockElems, + numGroups) == layout; +} + +static constexpr LayoutPattern c() { + return {LayoutPatternKind::Contiguous, 1, 0, 1}; +} + +static constexpr LayoutPattern ls(int64_t laneStride) { + return {LayoutPatternKind::LaneStride, laneStride, 0, 1}; +} + +static constexpr LayoutPattern d(int64_t factor, int64_t blockElems = 0) { + return {LayoutPatternKind::Deinterleaved, factor, blockElems, 1}; +} + +static constexpr LayoutPattern gs(int64_t slots, int64_t laneStride = 1) { + return {LayoutPatternKind::GroupSlots, slots, 0, laneStride}; +} + +static constexpr GroupMemoryPattern memAny() { + return {GroupMemoryPatternKind::Any}; +} + +static constexpr GroupMemoryPattern memContiguous() { + return {GroupMemoryPatternKind::Contiguous}; +} + +static constexpr GroupMemoryPattern memBlockAligned() { + return {GroupMemoryPatternKind::BlockAligned}; +} + +static constexpr GroupBlockPattern gb(int64_t numerator) { + return {GroupBlockPatternKind::Ratio, numerator, 1}; +} + +static constexpr GroupBlockPattern gb(int64_t numerator, int64_t denominator) { + return {GroupBlockPatternKind::Ratio, numerator, denominator}; +} + +static constexpr GroupBlockPattern gbFull(int64_t fullPartMultiple = 1) { + return {GroupBlockPatternKind::FullPartMultiple, fullPartMultiple, 1}; +} + +//===----------------------------------------------------------------------===// +// Query key derivation helpers. Keep new layout facts in the rule table +// section below; helpers in this section should only derive normalized keys. +//===----------------------------------------------------------------------===// + +static std::optional getConstantIndexValue(Value value) { + if (auto constant = value.getDefiningOp()) + return constant.value(); + if (auto constant = value.getDefiningOp()) { + if (constant.getType().isIndex()) + return constant.value(); + } + return std::nullopt; +} + +static FailureOr getGroupSizeFromNumGroups(VMIVRegType type, + int64_t numGroups, + std::string *reason) { + auto fail = [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + if (numGroups <= 0) + return fail("requires num_groups to be positive"); + if (type.getElementCount() % numGroups != 0) + return fail("requires num_groups to evenly divide logical lane count"); + return type.getElementCount() / numGroups; +} + +//===----------------------------------------------------------------------===// +// Rule tables +//===----------------------------------------------------------------------===// + +struct EnsureLayoutPattern { + ElementBitsPattern elementBits; + ElementCountPattern elementCounts; + LayoutPattern sourceLayout; + LayoutPattern resultLayout; +}; + +static constexpr EnsureLayoutPattern kEnsureLayoutPatterns[] = { + {bits<16>(), N<256>(), c(), d(2, 1)}, + {bits<32>(), N<128, 256>(), c(), d(2, 1)}, + {bits<64>(), N<64, 128, 256>(), c(), d(2, 1)}, + {bits<16>(), N<256>(), d(2, 1), c()}, + {bits<32>(), N<128, 256>(), d(2, 1), c()}, + {bits<64>(), N<64, 128, 256>(), d(2, 1), c()}, + + {bits<32>(), N<256>(), c(), d(4, 1)}, + {bits<64>(), N<128, 256>(), c(), d(4, 1)}, + {bits<32>(), N<256>(), d(4, 1), c()}, + {bits<64>(), N<128, 256>(), d(4, 1), c()}, + + {bits<32>(), N<256, 512>(), d(2, 1), d(4, 1)}, + {bits<64>(), N<128, 256>(), d(2, 1), d(4, 1)}, + {bits<32>(), N<256, 512>(), d(4, 1), d(2, 1)}, + {bits<64>(), N<128, 256>(), d(4, 1), d(2, 1)}, + + {bits<8>(), N<1, 2, 4, 8, 64, 128>(), c(), ls(2)}, + {bits<16>(), N<1, 2, 4, 8, 64>(), c(), ls(2)}, + {bits<8>(), N<1, 2, 4, 8, 64, 128>(), ls(2), c()}, + {bits<16>(), N<1, 2, 4, 8, 64>(), ls(2), c()}, + + {bits<8>(), N<1, 2, 4, 8, 64>(), c(), ls(4)}, + {bits<8>(), N<1, 2, 4, 8, 64>(), ls(4), c()}, + + // Group-slot lane-stride materialization keeps num_groups and slots=8. + // Each physical part carries at most eight row-local group values, so the + // conversion is independent of the total logical element count. + {bits<8, 16>(), anyN(), gs(8), gs(8, 2)}, + {bits<8, 16>(), anyN(), gs(8, 2), gs(8)}, + {bits<8>(), anyN(), gs(8), gs(8, 4)}, + {bits<8>(), anyN(), gs(8, 4), gs(8)}, + {bits<8>(), anyN(), gs(8, 2), gs(8, 4)}, + {bits<8>(), anyN(), gs(8, 4), gs(8, 2)}, +}; + +struct EnsureMaskLayoutPattern { + MaskGranularityPattern granularity; + ElementCountPattern elementCounts; + LayoutPattern sourceLayout; + LayoutPattern resultLayout; +}; + +static constexpr EnsureMaskLayoutPattern kEnsureMaskLayoutPatterns[] = { + {mb16(), N<256>(), c(), d(2, 1)}, + {mb32(), N<128, 256, 512>(), c(), d(2, 1)}, + {mb32(), N<128>(), c(), d(2, 8)}, + {mb16(), N<256>(), d(2, 1), c()}, + {mb32(), N<128, 256, 512>(), d(2, 1), c()}, + {mb32(), N<128>(), d(2, 8), c()}, + + {mb32(), N<256>(), c(), d(4, 1)}, + {mb32(), N<256>(), c(), d(4, 8)}, + {mb32(), N<256>(), d(4, 1), c()}, + {mb32(), N<256>(), d(4, 8), c()}, + + {mb8(), N<1, 2, 4, 8, 64, 128>(), c(), ls(2)}, + {mb16(), N<1, 2, 4, 8, 64>(), c(), ls(2)}, + {mb32(), N<1, 2, 4, 8, 64>(), c(), ls(2)}, + {mb8(), N<1, 2, 4, 8, 64, 128>(), ls(2), c()}, + {mb16(), N<1, 2, 4, 8, 64>(), ls(2), c()}, + {mb32(), N<1, 2, 4, 8, 64>(), ls(2), c()}, + + {mb8(), N<1, 2, 4, 8, 64>(), c(), ls(4)}, + {mb16(), N<1, 2, 4, 8>(), c(), ls(4)}, + {mb32(), N<1, 2, 4, 8>(), c(), ls(4)}, + {mb8(), N<1, 2, 4, 8, 64>(), ls(4), c()}, + {mb16(), N<1, 2, 4, 8>(), ls(4), c()}, + {mb32(), N<1, 2, 4, 8>(), ls(4), c()}, +}; + +struct GroupBlockClassPattern { + GroupBlockPattern block; + VMIGroupBlockClass blockClass; +}; + +static constexpr GroupBlockClassPattern kGroupBlockClassPatterns[] = { + {gb(1, 4), VMIGroupBlockClass::QuarterBlock}, + {gb(1, 2), VMIGroupBlockClass::HalfBlock}, + {gb(1), VMIGroupBlockClass::OneBlock}, + {gb(2), VMIGroupBlockClass::TwoBlock}, + {gb(4), VMIGroupBlockClass::FourBlock}, + {gbFull(), VMIGroupBlockClass::FullPartMultiple}, + {gbFull(2), VMIGroupBlockClass::FullPartMultiple}, + {gbFull(4), VMIGroupBlockClass::FullPartMultiple}, +}; + +struct GroupReduceLayoutPattern { + GroupBlockPattern block; + LayoutPattern sourceLayout; + LayoutPattern resultLayout; +}; + +static constexpr GroupReduceLayoutPattern kGroupReduceLayoutPatterns[] = { + {gb(1, 4), ls(4), gs(8)}, + {gb(1, 2), ls(2), gs(8)}, + {gb(1), c(), gs(8)}, + {gb(2), d(2, 1), gs(8)}, + {gb(2), d(2, 8), gs(8)}, + {gb(4), d(4, 1), gs(8)}, + {gb(4), d(4, 8), gs(8)}, + {gbFull(), c(), gs(1)}, + {gbFull(2), d(2, 1), gs(1)}, + {gbFull(4), d(4, 1), gs(1)}, +}; + +struct PreferredCastLayoutPattern { + ElementBitsPattern sourceBits; + ElementBitsPattern resultBits; + int64_t elementCount = 0; // 0 means the default row for this bit-width pair. + LayoutPattern sourceLayout; + LayoutPattern resultLayout; +}; + +struct LegalCastLayoutPattern { + ElementBitsPattern sourceBits; + ElementBitsPattern resultBits; + LayoutPattern sourceLayout; + LayoutPattern resultLayout; +}; + +struct LegalMaskGranularityCastLayoutPattern { + MaskGranularityPattern sourceGranularity; + MaskGranularityPattern resultGranularity; + LayoutPattern sourceLayout; + LayoutPattern resultLayout; +}; + +struct InterleaveLayoutPattern { + PhysicalChunkCountPattern chunks; + LayoutPattern lhsLayout; + LayoutPattern rhsLayout; + LayoutPattern maskLayout; + LayoutPattern lowLayout; + LayoutPattern highLayout; +}; + +static constexpr PreferredCastLayoutPattern kPreferredCastLayoutPatterns[] = { + // Exact rows override the default legal relation for small shapes where the + // compact lane-stride form is the natural cast layout. + {bits<16>(), bits<32>(), 64, ls(2), c()}, + {bits<8>(), bits<32>(), 64, ls(4), c()}, + {bits<32>(), bits<16>(), 64, c(), ls(2)}, + {bits<32>(), bits<8>(), 64, c(), ls(4)}, + {bits<32>(), bits<8>(), 128, d(2), ls(2)}, + + // Default rows for the storage-width cast families. + {bits<8>(), bits<16>(), 0, c(), d(2)}, + {bits<16>(), bits<32>(), 0, c(), d(2)}, + {bits<8>(), bits<32>(), 0, c(), d(4)}, + {bits<16>(), bits<8>(), 0, d(2), c()}, + {bits<32>(), bits<16>(), 0, d(2), c()}, + {bits<32>(), bits<8>(), 0, d(4), c()}, +}; + +static constexpr LegalCastLayoutPattern kLegalCastLayoutPatterns[] = { + // 2x widening. + {bits<8>(), bits<16>(), c(), d(2)}, + {bits<8>(), bits<16>(), ls(2), c()}, + {bits<8>(), bits<16>(), d(2), d(4)}, + {bits<16>(), bits<32>(), c(), d(2)}, + {bits<16>(), bits<32>(), ls(2), c()}, + {bits<16>(), bits<32>(), d(2), d(4)}, + + // 2x narrowing. + {bits<16>(), bits<8>(), d(2), c()}, + {bits<16>(), bits<8>(), c(), ls(2)}, + {bits<16>(), bits<8>(), d(4), d(2)}, + {bits<32>(), bits<16>(), d(2), c()}, + {bits<32>(), bits<16>(), c(), ls(2)}, + {bits<32>(), bits<16>(), d(4), d(2)}, + + // 4x widening/narrowing. + {bits<8>(), bits<32>(), c(), d(4)}, + {bits<8>(), bits<32>(), ls(2), d(2)}, + {bits<8>(), bits<32>(), ls(4), c()}, + {bits<32>(), bits<8>(), d(4), c()}, + {bits<32>(), bits<8>(), c(), ls(4)}, + {bits<32>(), bits<8>(), d(2), ls(2)}, + + // Group-slot casts keep the row-local group layout. num_groups is + // inherited from the anchor layout at query time. Packed narrowing records + // the selected sub-lane stride on the result; widening is the inverse. + {bits<8>(), bits<16>(), gs(1), gs(1)}, + {bits<8>(), bits<16>(), gs(8, 2), gs(8)}, + {bits<16>(), bits<32>(), gs(1), gs(1)}, + {bits<16>(), bits<32>(), gs(8, 2), gs(8)}, + {bits<8>(), bits<32>(), gs(1), gs(1)}, + {bits<8>(), bits<32>(), gs(8, 4), gs(8)}, + {bits<16>(), bits<32>(), gs(8), gs(8)}, + {bits<8>(), bits<32>(), gs(8), gs(8)}, + {bits<16>(), bits<8>(), gs(1), gs(1)}, + {bits<16>(), bits<8>(), gs(8), gs(8, 2)}, + {bits<32>(), bits<16>(), gs(1), gs(1)}, + {bits<32>(), bits<16>(), gs(8), gs(8, 2)}, + {bits<32>(), bits<8>(), gs(1), gs(1)}, + {bits<32>(), bits<8>(), gs(8), gs(8, 4)}, +}; + +static constexpr LegalMaskGranularityCastLayoutPattern + kLegalMaskGranularityCastLayoutPatterns[] = { + // 2x widening. + {mb8(), mb16(), c(), d(2)}, + {mb8(), mb16(), ls(2), c()}, + {mb8(), mb16(), d(2), d(4)}, + {mb16(), mb32(), c(), d(2)}, + {mb16(), mb32(), ls(2), c()}, + {mb16(), mb32(), d(2), d(4)}, + + // 2x narrowing. + {mb16(), mb8(), d(2), c()}, + {mb16(), mb8(), c(), ls(2)}, + {mb16(), mb8(), d(4), d(2)}, + {mb32(), mb16(), d(2), c()}, + {mb32(), mb16(), c(), ls(2)}, + {mb32(), mb16(), d(4), d(2)}, + + // 4x widening/narrowing. + {mb8(), mb32(), c(), d(4)}, + {mb8(), mb32(), ls(2), d(2)}, + {mb8(), mb32(), ls(4), c()}, + {mb32(), mb8(), d(4), c()}, + {mb32(), mb8(), c(), ls(4)}, + {mb32(), mb8(), d(2), ls(2)}, + + // Group-slot casts keep the row-local group layout. + {mb8(), mb16(), gs(1), gs(1)}, + {mb8(), mb16(), gs(8, 2), gs(8)}, + {mb16(), mb32(), gs(1), gs(1)}, + {mb16(), mb32(), gs(8, 2), gs(8)}, + {mb8(), mb32(), gs(1), gs(1)}, + {mb8(), mb32(), gs(8, 4), gs(8)}, + {mb16(), mb32(), gs(8), gs(8)}, + {mb8(), mb32(), gs(8), gs(8)}, + {mb16(), mb8(), gs(1), gs(1)}, + {mb16(), mb8(), gs(8), gs(8, 2)}, + {mb32(), mb16(), gs(1), gs(1)}, + {mb32(), mb16(), gs(8), gs(8, 2)}, + {mb32(), mb8(), gs(1), gs(1)}, + {mb32(), mb8(), gs(8), gs(8, 4)}, +}; + +static constexpr InterleaveLayoutPattern kVdintlvLayoutPatterns[] = { + {chunk<2, 4>(), d(2, 1), d(2, 1), d(2, 1), c(), c()}, + {chunk<4>(), d(4, 1), d(4, 1), d(4, 1), d(2, 1), d(2, 1)}, + {chunk<1>(), c(), c(), c(), c(), c()}, +}; + +static constexpr InterleaveLayoutPattern kVintlvLayoutPatterns[] = { + {chunk<2, 4>(), c(), c(), c(), d(2, 1), d(2, 1)}, + {chunk<4>(), d(2, 1), d(2, 1), d(2, 1), d(4, 1), d(4, 1)}, + {chunk<1>(), c(), c(), c(), c(), c()}, +}; + +struct SupplementalCastLayoutPattern { + ElementBitsPattern sourceBits; + ElementBitsPattern resultBits; + LayoutPattern sourceLayout; + LayoutPattern resultLayout; +}; + +static constexpr SupplementalCastLayoutPattern + kSupplementalIntegerExtLayoutPatterns[] = { + {bits<8>(), bits<32>(), gs(8), gs(8)}, + {bits<16>(), bits<32>(), gs(8), gs(8)}, +}; + +static constexpr SupplementalCastLayoutPattern + kSupplementalNarrowCastLayoutPatterns[] = { + {bits<32>(), bits<8>(), gs(8), gs(8)}, + {bits<32>(), bits<16>(), gs(8), gs(8)}, +}; + +struct DenseMemoryLayoutPattern { + ElementBitsPattern elementBits; + LayoutPattern layout; + ElementCountPattern elementCounts = anyN(); + bool preferred = false; +}; + +static constexpr DenseMemoryLayoutPattern kDenseLoadLayoutPatterns[] = { + {bits<8, 16, 32>(), c()}, {bits<8, 16, 32>(), ls(2)}, {bits<8>(), ls(4)}, + {bits<8, 16, 32>(), d(2)}, {bits<8, 16, 32>(), d(4)}, +}; + +struct DeinterleaveLoadLayoutPattern { + ElementBitsPattern elementBits; + LayoutPattern lowLayout; + LayoutPattern highLayout; +}; + +static constexpr DeinterleaveLoadLayoutPattern + kDeinterleaveLoadLayoutPatterns[] = { + {bits<8, 16, 32>(), c(), c()}, +}; + +static constexpr DenseMemoryLayoutPattern kDenseStoreLayoutPatterns[] = { + {bits<8, 16, 32>(), c()}, + {bits<8>(), ls(4), N<64>(), /*preferred=*/true}, + {bits<8>(), ls(2), N<128>(), /*preferred=*/true}, + {bits<16>(), ls(2), N<64>(), /*preferred=*/true}, + {bits<8, 16, 32>(), ls(2)}, + {bits<8>(), ls(4)}, + {bits<8, 16, 32>(), d(2, 1)}, + {bits<8, 16, 32>(), d(4, 1)}, +}; + +struct DenseMaskedStoreLayoutPattern { + ElementBitsPattern elementBits; + LayoutPattern valueLayout; + LayoutPattern maskLayout; + ElementCountPattern elementCounts = anyN(); + bool preferred = false; +}; + +struct DenseMaskedLoadLayoutPattern { + ElementBitsPattern elementBits; + LayoutPattern resultLayout; + LayoutPattern maskLayout; + LayoutPattern passthruLayout; +}; + +static constexpr DenseMaskedStoreLayoutPattern + kDenseMaskedStoreLayoutPatterns[] = { + {bits<8, 16, 32>(), c(), c()}, + {bits<8>(), ls(4), ls(4), N<64>(), /*preferred=*/true}, + {bits<8>(), ls(2), ls(2), N<128>(), /*preferred=*/true}, + {bits<16>(), ls(2), ls(2), N<64>(), /*preferred=*/true}, + {bits<8, 16>(), ls(2), ls(2)}, + {bits<8>(), ls(4), ls(4)}, + {bits<8, 16, 32>(), d(2, 1), d(2, 1)}, + {bits<8, 16, 32>(), d(4, 1), d(4, 1)}, +}; + +static constexpr DenseMaskedLoadLayoutPattern + kDenseMaskedLoadLayoutPatterns[] = { + {bits<8, 16, 32>(), c(), c(), c()}, +}; + +struct GroupLoadLayoutPattern { + ElementBitsPattern elementBits; + GroupBlockPattern block = gb(2); + GroupMemoryPattern memory = memAny(); + LayoutPattern resultLayout; +}; + +static constexpr GroupLoadLayoutPattern kGroupLoadLayoutPatterns[] = { + {bits<8, 16, 32>(), gb(1, 4), memContiguous(), c()}, + {bits<8, 16, 32>(), gb(1, 2), memContiguous(), c()}, + {bits<8, 16, 32>(), gb(1), memContiguous(), c()}, + {bits<8, 16, 32>(), gb(2), memContiguous(), c()}, + {bits<8, 16, 32>(), gb(4), memContiguous(), c()}, + {bits<8, 16, 32>(), gbFull(), memAny(), c()}, + {bits<32>(), gb(2), memBlockAligned(), d(2, 8)}, + {bits<32>(), gb(4), memBlockAligned(), d(4, 8)}, +}; + +struct GroupSlotMemoryLayoutPattern { + LayoutPattern layout; +}; + +static constexpr GroupSlotMemoryLayoutPattern kGroupSlotMemoryLayoutPatterns[] = + { + {gs(1)}, + {gs(8)}, + {gs(8, 2)}, + {gs(8, 4)}, +}; + +struct GroupBroadcastLoadLayoutPattern { + GroupBlockPattern block; + ElementBitsPattern elementBits; + GroupMemoryPattern memory = memContiguous(); + LayoutPattern resultLayout; +}; + +static constexpr GroupBroadcastLoadLayoutPattern + kGroupBroadcastLoadLayoutPatterns[] = { + {gb(1, 4), bits<8, 16, 32>(), memContiguous(), ls(4)}, + {gb(1, 2), bits<8, 16, 32>(), memContiguous(), ls(2)}, + {gb(1), bits<8, 16, 32>(), memContiguous(), c()}, + {gb(2), bits<8, 16, 32>(), memContiguous(), c()}, + {gb(2), bits<8, 16, 32>(), memContiguous(), d(2, 1)}, + {gb(4), bits<8, 16, 32>(), memContiguous(), c()}, + {gb(4), bits<8, 16, 32>(), memContiguous(), d(4, 1)}, + {gbFull(), bits<8, 16, 32>(), memAny(), c()}, +}; + +struct GroupBroadcastLoadDirectPattern { + VMIGroupBroadcastLoadDirectKind kind; + ElementCountPattern numGroups; + GroupBlockPattern block; + ElementBitsPattern elementBits; + GroupMemoryPattern memory = memContiguous(); + LayoutPattern resultLayout; +}; + +static constexpr GroupBroadcastLoadDirectPattern + kGroupBroadcastLoadDirectPatterns[] = { + {VMIGroupBroadcastLoadDirectKind::E2B, G<8>(), gb(1), bits<16, 32>(), + memContiguous(), c()}, + {VMIGroupBroadcastLoadDirectKind::E2B, G<8>(), gb(2), bits<16, 32>(), + memContiguous(), d(2, 1)}, + {VMIGroupBroadcastLoadDirectKind::E2B, G<8>(), gb(4), bits<16, 32>(), + memContiguous(), d(4, 1)}, + {VMIGroupBroadcastLoadDirectKind::BRC, anyG(), gbFull(), + bits<8, 16, 32>(), memAny(), c()}, +}; + +struct GroupBroadcastLayoutPattern { + GroupBlockPattern block; + LayoutPattern sourceLayout; + LayoutPattern resultLayout; +}; + +static constexpr GroupBroadcastLayoutPattern kGroupBroadcastLayoutPatterns[] = { + {gb(1, 4), gs(8), ls(4)}, + {gb(1, 2), gs(8), ls(2)}, + {gb(1), gs(8), c()}, + {gb(1), gs(8, 2), c()}, + {gb(1), gs(8, 4), c()}, + {gb(2), gs(8), c()}, + {gb(2), gs(8), d(2, 1)}, + {gb(2), gs(8), d(2, 8)}, + {gb(4), gs(8), c()}, + {gb(4), gs(8), d(4, 1)}, + {gb(4), gs(8), d(4, 8)}, + {gbFull(), gs(8), c()}, + {gbFull(), gs(1), c()}, + {gbFull(2), gs(1), d(2, 1)}, + {gbFull(4), gs(1), d(4, 1)}, +}; + +struct HistogramLayoutPattern { + LayoutPattern accLayout; + LayoutPattern sourceLayout; + LayoutPattern maskLayout; + LayoutPattern resultLayout; +}; + +static constexpr HistogramLayoutPattern kVdhistLayoutPatterns[] = { + {c(), c(), c(), c()}, +}; + +struct WidthChangingBitcastLayoutPattern { + LayoutPattern layout; +}; + +static constexpr WidthChangingBitcastLayoutPattern + kWidthChangingBitcastLayoutPatterns[] = { + {c()}, +}; + +//===----------------------------------------------------------------------===// +// Table matching and materialization helpers +//===----------------------------------------------------------------------===// + +static VMIDeinterleaveLoadLayoutFact materializeDeinterleaveLoadLayoutFact( + MLIRContext *ctx, const DeinterleaveLoadLayoutPattern &pattern) { + return VMIDeinterleaveLoadLayoutFact{ + materializeLayoutPattern(ctx, pattern.lowLayout), + materializeLayoutPattern(ctx, pattern.highLayout)}; +} + +static bool isSameGroupBlockPattern(GroupBlockPattern lhs, + GroupBlockPattern rhs) { + return lhs.kind == rhs.kind && lhs.numerator == rhs.numerator && + lhs.denominator == rhs.denominator; +} + +static VMIGroupBlockClass +getGroupBlockClassFromPattern(GroupBlockPattern pattern) { + for (const GroupBlockClassPattern &row : kGroupBlockClassPatterns) + if (isSameGroupBlockPattern(pattern, row.block)) + return row.blockClass; + llvm_unreachable("unsupported group block pattern"); +} + +static bool matchesGroupBroadcastLoadMemoryPattern( + GroupMemoryPattern pattern, std::optional stride, + int64_t elementBits) { + switch (pattern.kind) { + case GroupMemoryPatternKind::Any: + return true; + case GroupMemoryPatternKind::Contiguous: + if (!stride) + return false; + return *stride == 1; + case GroupMemoryPatternKind::BlockAligned: { + if (!stride) + return false; + if (elementBits <= 0 || 256 % elementBits != 0) + return false; + int64_t alignedStrideElems = 256 / elementBits; + return *stride > 0 && *stride % alignedStrideElems == 0; + } + } + llvm_unreachable("unknown group memory pattern kind"); +} + +static bool matchesGroupLoadMemoryPattern(GroupMemoryPattern pattern, + std::optional rowStride, + int64_t groupSize, + int64_t elementBits) { + switch (pattern.kind) { + case GroupMemoryPatternKind::Any: + return true; + case GroupMemoryPatternKind::Contiguous: + return rowStride && *rowStride == groupSize; + case GroupMemoryPatternKind::BlockAligned: { + if (!rowStride || elementBits <= 0 || 256 % elementBits != 0) + return false; + int64_t alignedStrideElems = 256 / elementBits; + return *rowStride > 0 && *rowStride % alignedStrideElems == 0; + } + } + llvm_unreachable("unknown group memory pattern kind"); +} + +static bool isSupportedGroupSlotMemoryLayout(VMILayoutAttr layout, + int64_t numGroups) { + if (!layout || !layout.isGroupSlots() || layout.getNumGroups() != numGroups || + layout.getSlots() <= 0) + return false; + for (const GroupSlotMemoryLayoutPattern &pattern : + kGroupSlotMemoryLayoutPatterns) + if (matchesLayoutPattern(layout.getContext(), pattern.layout, layout, + numGroups)) + return true; + return false; +} + +static FailureOr getGroupBlockClass(int64_t groupSize, + int64_t vcgBlockElems) { + if (vcgBlockElems <= 0) + return failure(); + + for (const GroupBlockClassPattern &row : kGroupBlockClassPatterns) { + GroupBlockPattern block = row.block; + if (block.kind == GroupBlockPatternKind::FullPartMultiple) { + int64_t fullPartElems = 8 * vcgBlockElems; + if (groupSize >= fullPartElems && groupSize % fullPartElems == 0) + return row.blockClass; + continue; + } + + int64_t numerator = vcgBlockElems * block.numerator; + if (block.denominator <= 0 || numerator % block.denominator != 0) + continue; + if (groupSize == numerator / block.denominator) + return row.blockClass; + } + return failure(); +} + +struct GroupLayoutKey { + int64_t groupSize = 0; + int64_t lanesPerPart = 0; + int64_t vcgBlockElems = 0; + VMIGroupBlockClass blockClass = VMIGroupBlockClass::OneBlock; +}; + +struct InterleaveLayoutKey { + int64_t elementCount = 0; + int64_t lanesPerPart = 0; + int64_t physicalChunkCount = 0; +}; + +static bool matchesPhysicalChunkCountPattern( + PhysicalChunkCountPattern pattern, InterleaveLayoutKey key) { + if (key.physicalChunkCount <= 0) + return false; + for (int64_t i = 0; i < pattern.count; ++i) + if (pattern.values[i] == key.physicalChunkCount) + return true; + return false; +} + +static FailureOr +buildInterleaveLayoutKey(VMIVRegType valueType, std::string *reason) { + auto fail = [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + FailureOr lanesPerPart = + getDataLanesPerPart(valueType.getElementType()); + if (failed(lanesPerPart)) + return fail("interleave layout requires element type with known physical " + "lanes per part"); + int64_t elementCount = valueType.getElementCount(); + if (elementCount <= 0) + return fail("interleave layout requires positive logical lane count"); + int64_t physicalChunkCount = + elementCount <= *lanesPerPart + ? 1 + : (elementCount % *lanesPerPart == 0 + ? elementCount / *lanesPerPart + : 0); + return InterleaveLayoutKey{elementCount, *lanesPerPart, physicalChunkCount}; +} + +static bool matchesGroupBlockPattern(GroupBlockPattern pattern, + GroupLayoutKey key) { + if (pattern.kind == GroupBlockPatternKind::FullPartMultiple) { + if (pattern.numerator <= 0) + return false; + int64_t fullPartElems = key.lanesPerPart * pattern.numerator; + return key.groupSize >= fullPartElems && + key.groupSize % fullPartElems == 0; + } + + int64_t numerator = key.vcgBlockElems * pattern.numerator; + if (pattern.denominator <= 0 || numerator % pattern.denominator != 0) + return false; + return key.groupSize == numerator / pattern.denominator; +} + +static FailureOr +buildGroupLayoutKey(VMIVRegType type, int64_t numGroups, + const Twine &unsupportedGroupSizeReason, + std::string *reason) { + auto fail = [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + FailureOr groupSize = + getGroupSizeFromNumGroups(type, numGroups, reason); + if (failed(groupSize)) + return failure(); + FailureOr lanesPerPart = getDataLanesPerPart(type.getElementType()); + if (failed(lanesPerPart) || *lanesPerPart % 8 != 0) + return fail("requires element type with known 32B VCG block width"); + + int64_t vcgBlockElems = *lanesPerPart / 8; + FailureOr blockClass = + getGroupBlockClass(*groupSize, vcgBlockElems); + if (failed(blockClass)) + return fail(unsupportedGroupSizeReason); + + return GroupLayoutKey{*groupSize, *lanesPerPart, vcgBlockElems, *blockClass}; +} + +static VMIGroupReduceLayoutFact +materializeGroupReduceLayoutFact(MLIRContext *ctx, + const GroupReduceLayoutPattern &pattern, + int64_t groupSize, int64_t lanesPerPart, + int64_t vcgBlockElems, int64_t numGroups) { + VMIGroupReduceLayoutFact fact; + fact.blockClass = getGroupBlockClassFromPattern(pattern.block); + fact.sourceLayout = materializeLayoutPattern(ctx, pattern.sourceLayout); + fact.maskLayout = fact.sourceLayout; + fact.resultLayout = + materializeLayoutPattern(ctx, pattern.resultLayout, + /*inheritedBlockElems=*/1, numGroups); + fact.groupSize = groupSize; + fact.lanesPerPart = lanesPerPart; + fact.vcgBlockElems = vcgBlockElems; + return fact; +} + +static VMIGroupBroadcastLayoutFact materializeGroupBroadcastLayoutFact( + MLIRContext *ctx, const GroupBroadcastLayoutPattern &pattern, + int64_t groupSize, int64_t lanesPerPart, int64_t vcgBlockElems, + int64_t numGroups) { + VMIGroupBroadcastLayoutFact fact; + fact.blockClass = getGroupBlockClassFromPattern(pattern.block); + fact.sourceLayout = + materializeLayoutPattern(ctx, pattern.sourceLayout, + /*inheritedBlockElems=*/1, numGroups); + fact.resultLayout = + materializeLayoutPattern(ctx, pattern.resultLayout, + /*inheritedBlockElems=*/1, numGroups); + fact.groupSize = groupSize; + fact.lanesPerPart = lanesPerPart; + fact.vcgBlockElems = vcgBlockElems; + return fact; +} + +static VMIInterleaveLayoutFact materializeInterleaveLayoutFact( + MLIRContext *ctx, const InterleaveLayoutPattern &pattern, + InterleaveLayoutKey key) { + VMIInterleaveLayoutFact fact; + fact.lhsLayout = materializeLayoutPattern(ctx, pattern.lhsLayout); + fact.rhsLayout = materializeLayoutPattern(ctx, pattern.rhsLayout); + fact.maskLayout = materializeLayoutPattern(ctx, pattern.maskLayout); + fact.lowLayout = materializeLayoutPattern(ctx, pattern.lowLayout); + fact.highLayout = materializeLayoutPattern(ctx, pattern.highLayout); + fact.elementCount = key.elementCount; + fact.lanesPerPart = key.lanesPerPart; + return fact; +} + +} // namespace + +//===----------------------------------------------------------------------===// +// Query implementations +//===----------------------------------------------------------------------===// + +FailureOr +VMILayoutSupport::getPreferredGroupReduceLayoutFact(VMIVRegType sourceType, + int64_t numGroups, + std::string *reason) const { + auto fail = [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + FailureOr key = buildGroupLayoutKey( + sourceType, numGroups, + "group_reduce layout supports group sizes of 1/4, 1/2, 1, 2, or 4 " + "32B VCG blocks, or full physical chunk multiples", + reason); + if (failed(key)) + return failure(); + + for (const GroupReduceLayoutPattern &pattern : kGroupReduceLayoutPatterns) { + if (!matchesGroupBlockPattern(pattern.block, *key)) + continue; + return materializeGroupReduceLayoutFact(sourceType.getContext(), pattern, + key->groupSize, key->lanesPerPart, + key->vcgBlockElems, numGroups); + } + + return fail("group_reduce layout supports group sizes of 1/4, 1/2, 1, 2, " + "or 4 32B VCG blocks, or full physical chunk multiples"); +} + +FailureOr +VMILayoutSupport::getGroupReduceLayoutFactForLayouts( + VMIVRegType sourceType, VMIMaskType maskType, VMIVRegType resultType, + int64_t numGroups, std::string *reason) const { + auto fail = [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); + VMILayoutAttr maskLayout = maskType.getLayoutAttr(); + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + if (!sourceLayout || !maskLayout || !resultLayout) + return fail("requires assigned source, mask, and result layouts"); + + FailureOr key = buildGroupLayoutKey( + sourceType, numGroups, + "group_reduce layout table has no row for this group size", reason); + if (failed(key)) + return failure(); + + for (const GroupReduceLayoutPattern &pattern : kGroupReduceLayoutPatterns) { + if (!matchesGroupBlockPattern(pattern.block, *key)) + continue; + VMIGroupReduceLayoutFact candidate = materializeGroupReduceLayoutFact( + sourceType.getContext(), pattern, key->groupSize, key->lanesPerPart, + key->vcgBlockElems, numGroups); + if (candidate.sourceLayout == sourceLayout && + candidate.maskLayout == maskLayout && + candidate.resultLayout == resultLayout) + return candidate; + } + + return fail("group_reduce source/mask/result layouts do not match a legal " + "layout table row for the group size"); +} + +FailureOr> +VMILayoutSupport::getGroupReduceLayoutFactsForLayout( + VMIVRegType sourceType, int64_t numGroups, VMIGroupReduceLayoutPort port, + VMILayoutAttr layout, std::string *reason) const { + auto fail = [&](const Twine &message) + -> FailureOr> { + if (reason) + *reason = message.str(); + return failure(); + }; + + if (!layout) + return fail("requires assigned group_reduce layout query port"); + + FailureOr key = buildGroupLayoutKey( + sourceType, numGroups, + "group_reduce layout table has no row for this group size", reason); + if (failed(key)) + return failure(); + + SmallVector facts; + for (const GroupReduceLayoutPattern &pattern : kGroupReduceLayoutPatterns) { + if (!matchesGroupBlockPattern(pattern.block, *key)) + continue; + VMIGroupReduceLayoutFact candidate = materializeGroupReduceLayoutFact( + sourceType.getContext(), pattern, key->groupSize, key->lanesPerPart, + key->vcgBlockElems, numGroups); + + VMILayoutAttr candidateLayout; + switch (port) { + case VMIGroupReduceLayoutPort::Source: + candidateLayout = candidate.sourceLayout; + break; + case VMIGroupReduceLayoutPort::Mask: + candidateLayout = candidate.maskLayout; + break; + case VMIGroupReduceLayoutPort::Result: + candidateLayout = candidate.resultLayout; + break; + } + if (candidateLayout == layout) + facts.push_back(candidate); + } + + if (facts.empty()) + return fail("group_reduce layout query port does not match a legal layout " + "table row for the group size"); + return facts; +} + +FailureOr +VMILayoutSupport::getGroupBroadcastLayoutFactForLayouts( + VMIVRegType sourceType, VMIVRegType resultType, int64_t numGroups, + std::string *reason) const { + auto fail = + [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + if (!sourceLayout || !resultLayout) + return fail("requires assigned source/result layouts"); + + FailureOr key = buildGroupLayoutKey( + resultType, numGroups, + "group_broadcast layout table has no row for this group size", reason); + if (failed(key)) + return failure(); + + for (const GroupBroadcastLayoutPattern &pattern : + kGroupBroadcastLayoutPatterns) { + if (!matchesGroupBlockPattern(pattern.block, *key)) + continue; + VMIGroupBroadcastLayoutFact candidate = materializeGroupBroadcastLayoutFact( + sourceType.getContext(), pattern, key->groupSize, key->lanesPerPart, + key->vcgBlockElems, numGroups); + if (candidate.sourceLayout == sourceLayout && + candidate.resultLayout == resultLayout) + return candidate; + } + + return fail("source/result layouts do not match a supported group_broadcast " + "table row"); +} + +FailureOr> +VMILayoutSupport::getGroupBroadcastLayoutFactsForLayout( + VMIVRegType sourceType, VMIVRegType resultType, int64_t numGroups, + VMIGroupBroadcastLayoutPort port, VMILayoutAttr layout, + std::string *reason) const { + auto fail = [&](const Twine &message) + -> FailureOr> { + if (reason) + *reason = message.str(); + return failure(); + }; + + if (!layout) + return fail("requires assigned group_broadcast layout query port"); + + FailureOr key = buildGroupLayoutKey( + resultType, numGroups, + "group_broadcast layout table has no row for this group size", reason); + if (failed(key)) + return failure(); + + SmallVector facts; + for (const GroupBroadcastLayoutPattern &pattern : + kGroupBroadcastLayoutPatterns) { + if (!matchesGroupBlockPattern(pattern.block, *key)) + continue; + VMIGroupBroadcastLayoutFact candidate = materializeGroupBroadcastLayoutFact( + sourceType.getContext(), pattern, key->groupSize, key->lanesPerPart, + key->vcgBlockElems, numGroups); + + VMILayoutAttr candidateLayout; + switch (port) { + case VMIGroupBroadcastLayoutPort::Source: + candidateLayout = candidate.sourceLayout; + break; + case VMIGroupBroadcastLayoutPort::Result: + candidateLayout = candidate.resultLayout; + break; + } + if (candidateLayout == layout) + facts.push_back(candidate); + } + + if (facts.empty()) + return fail("group_broadcast layout query port does not match a legal " + "layout table row for the group size"); + return facts; +} + +FailureOr +VMILayoutSupport::getGroupBroadcastLoadLayoutFact(VMIGroupBroadcastLoadOp op, + std::string *reason) const { + return getGroupBroadcastLoadLayoutFact( + cast(op.getResult().getType()), op.getSourceGroupStride(), + op.getNumGroupsAttr().getInt(), reason); +} + +FailureOr +VMILayoutSupport::getGroupBroadcastLoadLayoutFact(VMIVRegType resultType, + Value sourceGroupStride, + int64_t numGroups, + std::string *reason) const { + auto fail = + [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + if (!resultLayout) + return fail("requires assigned result layout"); + + unsigned elementBits = + pto::getPTOStorageElemBitWidth(resultType.getElementType()); + if (elementBits == 0) + return fail("group_broadcast_load requires known element bit width"); + std::optional stride = + getConstantIndexValue(sourceGroupStride); + + FailureOr key = buildGroupLayoutKey( + resultType, numGroups, + "group_broadcast_load layout table has no row for this group size", + reason); + if (failed(key)) + return failure(); + + for (const GroupBroadcastLoadLayoutPattern &pattern : + kGroupBroadcastLoadLayoutPatterns) { + if (!matchesGroupBlockPattern(pattern.block, *key)) + continue; + if (!matchesElementBitsPattern(pattern.elementBits, elementBits)) + continue; + if (!matchesGroupBroadcastLoadMemoryPattern(pattern.memory, stride, + elementBits)) + continue; + if (!matchesLayoutPattern(resultType.getContext(), pattern.resultLayout, + resultLayout, numGroups)) + continue; + return VMIGroupBroadcastLoadLayoutFact{ + getGroupBlockClassFromPattern(pattern.block), + resultLayout, + key->groupSize, + key->lanesPerPart, + key->vcgBlockElems, + static_cast(elementBits)}; + } + + int64_t alignedStrideElems = 256 / elementBits; + return fail(Twine("group_broadcast_load requires a table row for result " + "layout, group size, and either constant unit " + "source_group_stride or constant positive " + "source_group_stride divisible by ") + + Twine(alignedStrideElems) + " elements"); +} + +FailureOr +VMILayoutSupport::getGroupBroadcastLoadDirectFact(VMIGroupBroadcastLoadOp op, + std::string *reason) const { + return getGroupBroadcastLoadDirectFact( + cast(op.getResult().getType()), op.getSource().getType(), + op.getSourceGroupStride(), op.getNumGroupsAttr().getInt(), reason); +} + +FailureOr +VMILayoutSupport::getGroupBroadcastLoadDirectFact( + VMIVRegType resultType, Type sourceType, Value sourceGroupStride, + int64_t numGroups, std::string *reason) const { + auto fail = + [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + if (!isa(sourceType)) + return fail("group_broadcast_load direct lowering requires !pto.ptr source"); + + unsigned elementBits = + pto::getPTOStorageElemBitWidth(resultType.getElementType()); + if (elementBits == 0) + return fail("group_broadcast_load requires known element bit width"); + std::optional stride = getConstantIndexValue(sourceGroupStride); + + FailureOr key = buildGroupLayoutKey( + resultType, numGroups, + "group_broadcast_load preferred layout table has no row for this group " + "size", + reason); + if (failed(key)) + return failure(); + + VMILayoutAttr existing = resultType.getLayoutAttr(); + for (const GroupBroadcastLoadDirectPattern &pattern : + kGroupBroadcastLoadDirectPatterns) { + if (!matchesElementCountPattern(pattern.numGroups, numGroups)) + continue; + if (!matchesGroupBlockPattern(pattern.block, *key)) + continue; + if (!matchesElementBitsPattern(pattern.elementBits, elementBits)) + continue; + if (!matchesGroupBroadcastLoadMemoryPattern(pattern.memory, stride, + elementBits)) + continue; + VMILayoutAttr resultLayout = materializeLayoutPattern( + resultType.getContext(), pattern.resultLayout, /*blockElems=*/1, + numGroups); + if (existing && existing != resultLayout) + continue; + return VMIGroupBroadcastLoadDirectFact{ + pattern.kind, + VMIGroupBroadcastLoadLayoutFact{ + getGroupBlockClassFromPattern(pattern.block), + resultLayout, + key->groupSize, + key->lanesPerPart, + key->vcgBlockElems, + static_cast(elementBits)}}; + } + + return fail("group_broadcast_load has no preferred direct lowering layout " + "table row"); +} + +static std::pair getCastElementBits(VMIVRegType sourceType, + VMIVRegType resultType) { + unsigned sourceBits = + pto::getPTOStorageElemBitWidth(sourceType.getElementType()); + unsigned resultBits = + pto::getPTOStorageElemBitWidth(resultType.getElementType()); + return std::pair(sourceBits, resultBits); +} + +static VMICastLayoutFact makeCastLayoutFact(int64_t sourceBits, + int64_t resultBits, + VMILayoutAttr sourceLayout, + VMILayoutAttr resultLayout) { + VMICastLayoutFact fact; + fact.sourceBits = sourceBits; + fact.resultBits = resultBits; + fact.sourceLayout = sourceLayout; + fact.resultLayout = resultLayout; + return fact; +} + +static int64_t getMaskGranularityBits(StringRef granularity) { + if (granularity == "b8") + return 8; + if (granularity == "b16") + return 16; + if (granularity == "b32") + return 32; + return 0; +} + +static VMIMaskGranularityCastLayoutFact +makeMaskGranularityCastLayoutFact(int64_t sourceBits, int64_t resultBits, + VMILayoutAttr sourceLayout, + VMILayoutAttr resultLayout) { + VMIMaskGranularityCastLayoutFact fact; + fact.sourceGranularityBits = sourceBits; + fact.resultGranularityBits = resultBits; + fact.sourceLayout = sourceLayout; + fact.resultLayout = resultLayout; + return fact; +} + +FailureOr VMILayoutSupport::getPreferredCastLayoutFact( + VMIVRegType sourceType, VMIVRegType resultType, std::string *reason) const { + auto [sourceBits, resultBits] = getCastElementBits(sourceType, resultType); + + const PreferredCastLayoutPattern *selected = nullptr; + bool selectedIsExact = false; + int64_t elementCount = sourceType.getElementCount(); + for (const PreferredCastLayoutPattern &pattern : + kPreferredCastLayoutPatterns) { + if (!matchesElementBitsPattern(pattern.sourceBits, sourceBits) || + !matchesElementBitsPattern(pattern.resultBits, resultBits)) + continue; + bool isExact = pattern.elementCount != 0; + if (isExact && pattern.elementCount != elementCount) + continue; + if (!selected || (isExact && !selectedIsExact)) { + selected = &pattern; + selectedIsExact = isExact; + continue; + } + if (isExact == selectedIsExact) { + if (reason) + *reason = "preferred cast layout table has ambiguous matching rows"; + return failure(); + } + } + + if (!selected) { + if (reason) + *reason = "requires a preferred cast layout table row"; + return failure(); + } + + MLIRContext *ctx = sourceType.getContext(); + return makeCastLayoutFact(sourceBits, resultBits, + materializeLayoutPattern(ctx, + selected->sourceLayout), + materializeLayoutPattern(ctx, + selected->resultLayout)); +} + +FailureOr> +VMILayoutSupport::getCastLayoutFactsForLayout(VMIVRegType sourceType, + VMIVRegType resultType, + VMICastLayoutPort port, + VMILayoutAttr layout, + std::string *reason) const { + auto fail = [&](const Twine &message) + -> FailureOr> { + if (reason) + *reason = message.str(); + return failure(); + }; + + auto [sourceBits, resultBits] = getCastElementBits(sourceType, resultType); + MLIRContext *ctx = sourceType.getContext(); + SmallVector facts; + + int64_t blockElems = + layout && layout.isDeinterleaved() ? layout.getBlockElems() : 1; + int64_t numGroups = + layout && layout.isGroupSlots() ? layout.getNumGroups() : 0; + for (const LegalCastLayoutPattern &pattern : kLegalCastLayoutPatterns) { + if (!matchesElementBitsPattern(pattern.sourceBits, sourceBits) || + !matchesElementBitsPattern(pattern.resultBits, resultBits)) + continue; + + VMILayoutAttr sourceLayout = materializeLayoutPattern( + ctx, pattern.sourceLayout, blockElems, numGroups); + VMILayoutAttr resultLayout = materializeLayoutPattern( + ctx, pattern.resultLayout, blockElems, numGroups); + if (!sourceLayout || !resultLayout) + continue; + + if (port == VMICastLayoutPort::Source && sourceLayout != layout) + continue; + if (port == VMICastLayoutPort::Result && resultLayout != layout) + continue; + + facts.push_back( + makeCastLayoutFact(sourceBits, resultBits, sourceLayout, resultLayout)); + } + + if (facts.empty()) { + if (port == VMICastLayoutPort::Source) + return fail("requires a legal cast relation for the source layout"); + return fail("requires a legal cast relation for the result layout"); + } + return facts; +} + +static FailureOr +getUniqueCastLayoutFact(FailureOr> facts, + std::string *reason) { + auto fail = [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + if (failed(facts)) + return failure(); + if (facts->empty()) + return fail("cast layout query produced no layout facts"); + if (facts->size() != 1) + return fail("cast layout query produced ambiguous layout facts"); + return facts->front(); +} + +FailureOr VMILayoutSupport::getCastLayoutFactForSourceLayout( + VMIVRegType sourceType, VMIVRegType resultType, VMILayoutAttr sourceLayout, + std::string *reason) const { + return getUniqueCastLayoutFact( + getCastLayoutFactsForLayout(sourceType, resultType, + VMICastLayoutPort::Source, sourceLayout, + reason), + reason); +} + +FailureOr VMILayoutSupport::getCastLayoutFactForResultLayout( + VMIVRegType sourceType, VMIVRegType resultType, VMILayoutAttr resultLayout, + std::string *reason) const { + return getUniqueCastLayoutFact( + getCastLayoutFactsForLayout(sourceType, resultType, + VMICastLayoutPort::Result, resultLayout, + reason), + reason); +} + +FailureOr VMILayoutSupport::getCastLayoutFactForLayouts( + VMIVRegType sourceType, VMIVRegType resultType, VMILayoutAttr sourceLayout, + VMILayoutAttr resultLayout, std::string *reason) const { + auto fail = [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + FailureOr> facts = + getCastLayoutFactsForLayout(sourceType, resultType, + VMICastLayoutPort::Source, sourceLayout, + reason); + if (failed(facts)) + return failure(); + + std::optional selected; + for (const VMICastLayoutFact &fact : *facts) { + if (fact.resultLayout != resultLayout) + continue; + if (selected) + return fail("cast layout query produced ambiguous layout facts"); + selected = fact; + } + if (!selected) + return fail("source/result layouts do not match a legal cast table row"); + return *selected; +} + +FailureOr> +VMILayoutSupport::getMaskGranularityCastLayoutFactsForLayout( + VMIMaskType sourceType, VMIMaskType resultType, VMICastLayoutPort port, + VMILayoutAttr layout, std::string *reason) const { + auto fail = [&](const Twine &message) + -> FailureOr> { + if (reason) + *reason = message.str(); + return failure(); + }; + + if (sourceType.getElementCount() != resultType.getElementCount()) + return fail("requires source and result mask lane counts to match"); + if (!VMIMaskType::isConcreteGranularity(sourceType.getGranularity()) || + !VMIMaskType::isConcreteGranularity(resultType.getGranularity())) + return fail("requires concrete b8/b16/b32 source and result " + "granularities"); + + int64_t sourceBits = getMaskGranularityBits(sourceType.getGranularity()); + int64_t resultBits = getMaskGranularityBits(resultType.getGranularity()); + if (sourceBits == 0 || resultBits == 0) + return fail("requires supported source/result mask granularities"); + + MLIRContext *ctx = sourceType.getContext(); + int64_t blockElems = + layout && layout.isDeinterleaved() ? layout.getBlockElems() : 1; + int64_t numGroups = + layout && layout.isGroupSlots() ? layout.getNumGroups() : 0; + SmallVector facts; + for (const LegalMaskGranularityCastLayoutPattern &pattern : + kLegalMaskGranularityCastLayoutPatterns) { + if (!matchesMaskGranularityPattern(pattern.sourceGranularity, + sourceType.getGranularity()) || + !matchesMaskGranularityPattern(pattern.resultGranularity, + resultType.getGranularity())) + continue; + + VMILayoutAttr sourceLayout = materializeLayoutPattern( + ctx, pattern.sourceLayout, blockElems, numGroups); + VMILayoutAttr resultLayout = materializeLayoutPattern( + ctx, pattern.resultLayout, blockElems, numGroups); + if (!sourceLayout || !resultLayout) + continue; + + if (port == VMICastLayoutPort::Source && sourceLayout != layout) + continue; + if (port == VMICastLayoutPort::Result && resultLayout != layout) + continue; + + facts.push_back(makeMaskGranularityCastLayoutFact( + sourceBits, resultBits, sourceLayout, resultLayout)); + } + + if (facts.empty()) { + if (port == VMICastLayoutPort::Source) + return fail("requires a legal mask granularity cast relation for the " + "source layout"); + return fail("requires a legal mask granularity cast relation for the " + "result layout"); + } + return facts; +} + +FailureOr +VMILayoutSupport::getMaskGranularityCastLayoutFactForLayouts( + VMIMaskType sourceType, VMIMaskType resultType, VMILayoutAttr sourceLayout, + VMILayoutAttr resultLayout, std::string *reason) const { + auto fail = + [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + FailureOr> facts = + getMaskGranularityCastLayoutFactsForLayout( + sourceType, resultType, VMICastLayoutPort::Source, sourceLayout, + reason); + if (failed(facts)) + return failure(); + + std::optional selected; + for (const VMIMaskGranularityCastLayoutFact &fact : *facts) { + if (fact.resultLayout != resultLayout) + continue; + if (selected) + return fail("mask granularity cast layout query produced ambiguous " + "layout facts"); + selected = fact; + } + if (!selected) + return fail("source/result layouts do not match a legal mask granularity " + "cast table row"); + return *selected; +} + +FailureOr VMILayoutSupport::getWidenSourceLayoutForResultLayout( + VMIVRegType sourceType, VMIVRegType resultType, + VMILayoutAttr requestedResultLayout, std::string *reason) const { + FailureOr fact = getCastLayoutFactForResultLayout( + sourceType, resultType, requestedResultLayout, reason); + if (failed(fact)) + return failure(); + return fact->sourceLayout; +} + +static FailureOr getPreferredInterleaveLayoutFactImpl( + ArrayRef patterns, VMIVRegType valueType, + std::string *reason) { + auto fail = [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + FailureOr key = + buildInterleaveLayoutKey(valueType, reason); + if (failed(key)) + return failure(); + + for (const InterleaveLayoutPattern &pattern : patterns) { + if (!matchesPhysicalChunkCountPattern(pattern.chunks, *key)) + continue; + return materializeInterleaveLayoutFact(valueType.getContext(), pattern, + *key); + } + + return fail("requires a preferred interleave layout table row"); +} + +static FailureOr> +getInterleaveLayoutFactsForLayoutImpl( + ArrayRef patterns, VMIVRegType valueType, + VMIInterleaveLayoutPort port, VMILayoutAttr layout, + std::string *reason) { + auto fail = [&](const Twine &message) + -> FailureOr> { + if (reason) + *reason = message.str(); + return failure(); + }; + + if (!layout) + return fail("requires assigned interleave layout query port"); + + FailureOr key = + buildInterleaveLayoutKey(valueType, reason); + if (failed(key)) + return failure(); + + SmallVector facts; + for (const InterleaveLayoutPattern &pattern : patterns) { + if (!matchesPhysicalChunkCountPattern(pattern.chunks, *key)) + continue; + VMIInterleaveLayoutFact candidate = + materializeInterleaveLayoutFact(valueType.getContext(), pattern, *key); + + VMILayoutAttr candidateLayout; + switch (port) { + case VMIInterleaveLayoutPort::Lhs: + candidateLayout = candidate.lhsLayout; + break; + case VMIInterleaveLayoutPort::Rhs: + candidateLayout = candidate.rhsLayout; + break; + case VMIInterleaveLayoutPort::Mask: + candidateLayout = candidate.maskLayout; + break; + case VMIInterleaveLayoutPort::Low: + candidateLayout = candidate.lowLayout; + break; + case VMIInterleaveLayoutPort::High: + candidateLayout = candidate.highLayout; + break; + } + if (candidateLayout == layout) + facts.push_back(candidate); + } + + if (facts.empty()) + return fail("interleave layout query port does not match a legal layout " + "table row for the vector shape"); + return facts; +} + +static FailureOr getInterleaveLayoutFactForLayoutsImpl( + ArrayRef patterns, VMIVRegType lhsType, + VMIVRegType rhsType, VMIMaskType maskType, VMIVRegType lowType, + VMIVRegType highType, std::string *reason) { + auto fail = [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + if (lhsType.getElementCount() != rhsType.getElementCount() || + lhsType.getElementCount() != lowType.getElementCount() || + lhsType.getElementCount() != highType.getElementCount() || + lhsType.getElementCount() != maskType.getElementCount()) + return fail("interleave layout requires all ports to share logical lane " + "count"); + if (lhsType.getElementType() != rhsType.getElementType() || + lhsType.getElementType() != lowType.getElementType() || + lhsType.getElementType() != highType.getElementType()) + return fail("interleave layout requires all data ports to share element " + "type"); + + VMILayoutAttr lhsLayout = lhsType.getLayoutAttr(); + VMILayoutAttr rhsLayout = rhsType.getLayoutAttr(); + VMILayoutAttr maskLayout = maskType.getLayoutAttr(); + VMILayoutAttr lowLayout = lowType.getLayoutAttr(); + VMILayoutAttr highLayout = highType.getLayoutAttr(); + if (!lhsLayout || !rhsLayout || !maskLayout || !lowLayout || !highLayout) + return fail("requires assigned lhs/rhs/mask/low/high layouts"); + + FailureOr> facts = + getInterleaveLayoutFactsForLayoutImpl( + patterns, lhsType, VMIInterleaveLayoutPort::Lhs, lhsLayout, reason); + if (failed(facts)) + return failure(); + + std::optional selected; + for (const VMIInterleaveLayoutFact &fact : *facts) { + if (fact.rhsLayout != rhsLayout || fact.maskLayout != maskLayout || + fact.lowLayout != lowLayout || fact.highLayout != highLayout) + continue; + if (selected) + return fail("interleave layout query produced ambiguous layout facts"); + selected = fact; + } + if (!selected) + return fail("lhs/rhs/mask/low/high layouts do not match a legal " + "interleave layout table row"); + return *selected; +} + +FailureOr +VMILayoutSupport::getPreferredVintlvLayoutFact( + VMIVRegType valueType, std::string *reason) const { + return getPreferredInterleaveLayoutFactImpl(kVintlvLayoutPatterns, valueType, + reason); +} + +FailureOr +VMILayoutSupport::getPreferredVdintlvLayoutFact( + VMIVRegType valueType, std::string *reason) const { + return getPreferredInterleaveLayoutFactImpl(kVdintlvLayoutPatterns, valueType, + reason); +} + +FailureOr> +VMILayoutSupport::getVintlvLayoutFactsForLayout( + VMIVRegType valueType, VMIInterleaveLayoutPort port, VMILayoutAttr layout, + std::string *reason) const { + return getInterleaveLayoutFactsForLayoutImpl( + kVintlvLayoutPatterns, valueType, port, layout, reason); +} + +FailureOr> +VMILayoutSupport::getVdintlvLayoutFactsForLayout( + VMIVRegType valueType, VMIInterleaveLayoutPort port, VMILayoutAttr layout, + std::string *reason) const { + return getInterleaveLayoutFactsForLayoutImpl( + kVdintlvLayoutPatterns, valueType, port, layout, reason); +} + +FailureOr +VMILayoutSupport::getVintlvLayoutFactForLayouts( + VMIVRegType lhsType, VMIVRegType rhsType, VMIMaskType maskType, + VMIVRegType lowType, VMIVRegType highType, std::string *reason) const { + return getInterleaveLayoutFactForLayoutsImpl( + kVintlvLayoutPatterns, lhsType, rhsType, maskType, lowType, highType, + reason); +} + +FailureOr +VMILayoutSupport::getVdintlvLayoutFactForLayouts( + VMIVRegType lhsType, VMIVRegType rhsType, VMIMaskType maskType, + VMIVRegType lowType, VMIVRegType highType, std::string *reason) const { + return getInterleaveLayoutFactForLayoutsImpl( + kVdintlvLayoutPatterns, lhsType, rhsType, maskType, lowType, highType, + reason); +} + +FailureOr +VMILayoutSupport::getLoadLayoutFact(VMIVRegType resultType, + std::string *reason) const { + auto fail = [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + VMILayoutAttr layout = resultType.getLayoutAttr(); + if (!layout) + return fail("requires assigned result layout"); + for (const DenseMemoryLayoutPattern &pattern : kDenseLoadLayoutPatterns) { + if (!matchesElementBitsPattern(pattern.elementBits, + resultType.getElementType())) + continue; + if (!matchesElementCountPattern(pattern.elementCounts, + resultType.getElementCount())) + continue; + if (!matchesLayoutPattern(resultType.getContext(), pattern.layout, layout)) + continue; + return VMILoadLayoutFact{layout}; + } + + return fail("result layout does not match a supported dense load table row"); +} + +FailureOr +VMILayoutSupport::getPreferredDeinterleaveLoadLayoutFact( + VMIVRegType valueType, std::string *reason) const { + auto fail = + [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + for (const DeinterleaveLoadLayoutPattern &pattern : + kDeinterleaveLoadLayoutPatterns) { + if (!matchesElementBitsPattern(pattern.elementBits, + valueType.getElementType())) + continue; + return materializeDeinterleaveLoadLayoutFact(valueType.getContext(), + pattern); + } + return fail("requires a preferred deinterleave_load layout table row"); +} + +FailureOr> +VMILayoutSupport::getDeinterleaveLoadLayoutFactsForLayout( + VMIVRegType valueType, VMIDeinterleaveLoadLayoutPort port, + VMILayoutAttr layout, std::string *reason) const { + auto fail = [&](const Twine &message) + -> FailureOr> { + if (reason) + *reason = message.str(); + return failure(); + }; + if (!layout) + return fail("requires assigned deinterleave_load layout query port"); + + SmallVector facts; + for (const DeinterleaveLoadLayoutPattern &pattern : + kDeinterleaveLoadLayoutPatterns) { + if (!matchesElementBitsPattern(pattern.elementBits, + valueType.getElementType())) + continue; + VMIDeinterleaveLoadLayoutFact candidate = + materializeDeinterleaveLoadLayoutFact(valueType.getContext(), pattern); + VMILayoutAttr candidateLayout = + port == VMIDeinterleaveLoadLayoutPort::Low ? candidate.lowLayout + : candidate.highLayout; + if (candidateLayout == layout) + facts.push_back(candidate); + } + if (facts.empty()) + return fail("deinterleave_load layout query port does not match a legal " + "layout table row"); + return facts; +} + +FailureOr +VMILayoutSupport::getDeinterleaveLoadLayoutFactForLayouts( + VMIVRegType lowType, VMIVRegType highType, std::string *reason) const { + auto fail = + [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + if (lowType.getElementCount() != highType.getElementCount() || + lowType.getElementType() != highType.getElementType()) + return fail("deinterleave_load layout requires low/high to share shape " + "and element type"); + + VMILayoutAttr lowLayout = lowType.getLayoutAttr(); + VMILayoutAttr highLayout = highType.getLayoutAttr(); + if (!lowLayout || !highLayout) + return fail("requires assigned low/high layouts"); + + FailureOr> facts = + getDeinterleaveLoadLayoutFactsForLayout( + lowType, VMIDeinterleaveLoadLayoutPort::Low, lowLayout, reason); + if (failed(facts)) + return failure(); + for (const VMIDeinterleaveLoadLayoutFact &fact : *facts) + if (fact.highLayout == highLayout) + return fact; + return fail("low/high layouts do not match a legal deinterleave_load layout " + "table row"); +} + +FailureOr +VMILayoutSupport::getStoreLayoutFact(VMIVRegType valueType, + std::string *reason) const { + auto fail = [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + VMILayoutAttr layout = valueType.getLayoutAttr(); + if (!layout) + return fail("requires assigned value layout"); + for (const DenseMemoryLayoutPattern &pattern : kDenseStoreLayoutPatterns) { + if (!matchesElementBitsPattern(pattern.elementBits, + valueType.getElementType())) + continue; + if (!matchesElementCountPattern(pattern.elementCounts, + valueType.getElementCount())) + continue; + if (!matchesLayoutPattern(valueType.getContext(), pattern.layout, layout)) + continue; + return VMIStoreLayoutFact{layout}; + } + + return fail("value layout does not match a supported dense store table row"); +} + +FailureOr +VMILayoutSupport::getPreferredStoreLayoutFact(VMIVRegType valueType, + std::string *reason) const { + auto fail = [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + if (valueType.getLayoutAttr()) + return getStoreLayoutFact(valueType, reason); + + for (const DenseMemoryLayoutPattern &pattern : kDenseStoreLayoutPatterns) { + if (!pattern.preferred) + continue; + if (!matchesElementBitsPattern(pattern.elementBits, + valueType.getElementType())) + continue; + if (!matchesElementCountPattern(pattern.elementCounts, + valueType.getElementCount())) + continue; + VMILayoutAttr layout = + materializeLayoutPattern(valueType.getContext(), pattern.layout); + if (!layout) + continue; + return VMIStoreLayoutFact{layout}; + } + + return fail("value type does not match a preferred dense store table row"); +} + +FailureOr VMILayoutSupport::getMaskedStoreLayoutFact( + VMIVRegType valueType, VMIMaskType maskType, std::string *reason) const { + auto fail = [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + VMILayoutAttr valueLayout = valueType.getLayoutAttr(); + VMILayoutAttr maskLayout = maskType.getLayoutAttr(); + if (!valueLayout || !maskLayout) + return fail("requires assigned value/mask layouts"); + for (const DenseMaskedStoreLayoutPattern &pattern : + kDenseMaskedStoreLayoutPatterns) { + if (!matchesElementBitsPattern(pattern.elementBits, + valueType.getElementType())) + continue; + if (!matchesElementCountPattern(pattern.elementCounts, + valueType.getElementCount())) + continue; + if (!matchesLayoutPattern(valueType.getContext(), pattern.valueLayout, + valueLayout)) + continue; + if (!matchesLayoutPattern(maskType.getContext(), pattern.maskLayout, + maskLayout)) + continue; + return VMIMaskedStoreLayoutFact{valueLayout, maskLayout}; + } + + return fail("value/mask layouts do not match a supported dense masked store " + "table row"); +} + +FailureOr +VMILayoutSupport::getPreferredMaskedStoreLayoutFact( + VMIVRegType valueType, VMIMaskType maskType, std::string *reason) const { + auto fail = + [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + VMILayoutAttr existingValueLayout = valueType.getLayoutAttr(); + VMILayoutAttr existingMaskLayout = maskType.getLayoutAttr(); + if (existingValueLayout && existingMaskLayout) + return getMaskedStoreLayoutFact(valueType, maskType, reason); + + for (const DenseMaskedStoreLayoutPattern &pattern : + kDenseMaskedStoreLayoutPatterns) { + if (!pattern.preferred) + continue; + if (!matchesElementBitsPattern(pattern.elementBits, + valueType.getElementType())) + continue; + if (!matchesElementCountPattern(pattern.elementCounts, + valueType.getElementCount())) + continue; + + VMILayoutAttr valueLayout = + materializeLayoutPattern(valueType.getContext(), pattern.valueLayout); + VMILayoutAttr maskLayout = + materializeLayoutPattern(maskType.getContext(), pattern.maskLayout); + if (!valueLayout || !maskLayout) + continue; + if (existingValueLayout && existingValueLayout != valueLayout) + continue; + if (existingMaskLayout && existingMaskLayout != maskLayout) + continue; + return VMIMaskedStoreLayoutFact{valueLayout, maskLayout}; + } + + return fail("value/mask types do not match a preferred dense masked store " + "table row"); +} + +FailureOr VMILayoutSupport::getMaskedLoadLayoutFact( + VMIVRegType resultType, VMIMaskType maskType, VMIVRegType passthruType, + std::string *reason) const { + auto fail = [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + VMILayoutAttr maskLayout = maskType.getLayoutAttr(); + VMILayoutAttr passthruLayout = passthruType.getLayoutAttr(); + if (!resultLayout || !maskLayout || !passthruLayout) + return fail("requires assigned result/mask/passthru layouts"); + for (const DenseMaskedLoadLayoutPattern &pattern : + kDenseMaskedLoadLayoutPatterns) { + if (!matchesElementBitsPattern(pattern.elementBits, + resultType.getElementType())) + continue; + if (!matchesLayoutPattern(resultType.getContext(), pattern.resultLayout, + resultLayout)) + continue; + if (!matchesLayoutPattern(maskType.getContext(), pattern.maskLayout, + maskLayout)) + continue; + if (!matchesLayoutPattern(passthruType.getContext(), + pattern.passthruLayout, passthruLayout)) + continue; + return VMIMaskedLoadLayoutFact{resultLayout, maskLayout, passthruLayout}; + } + + return fail("result/mask/passthru layouts do not match a supported dense " + "masked_load table row"); +} + +static LogicalResult matchEnsureLayoutPattern(VMIVRegType sourceType, + VMIVRegType resultType, + VMILayoutAttr sourceLayout, + VMILayoutAttr resultLayout, + std::string *reason) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + if (!sourceLayout || !resultLayout) + return fail("requires assigned source/result layouts"); + if (sourceLayout == resultLayout) + return success(); + + int64_t numGroups = + sourceLayout.isGroupSlots() + ? sourceLayout.getNumGroups() + : (resultLayout.isGroupSlots() ? resultLayout.getNumGroups() : 0); + + for (const EnsureLayoutPattern &pattern : kEnsureLayoutPatterns) { + if (!matchesElementBitsPattern(pattern.elementBits, + sourceType.getElementType())) + continue; + if (!matchesElementCountPattern(pattern.elementCounts, + sourceType.getElementCount())) + continue; + if (!matchesLayoutPattern(sourceType.getContext(), pattern.sourceLayout, + sourceLayout, numGroups)) + continue; + if (!matchesLayoutPattern(resultType.getContext(), pattern.resultLayout, + resultLayout, numGroups)) + continue; + return success(); + } + + return fail("source/result layouts do not match a supported ensure_layout " + "table row"); +} + +static LogicalResult matchEnsureMaskLayoutPattern(VMIMaskType sourceType, + VMIMaskType resultType, + VMILayoutAttr sourceLayout, + VMILayoutAttr resultLayout, + std::string *reason) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + if (!sourceLayout || !resultLayout) + return fail("requires assigned source/result layouts"); + if (sourceLayout == resultLayout) + return success(); + + for (const EnsureMaskLayoutPattern &pattern : kEnsureMaskLayoutPatterns) { + if (!matchesMaskGranularityPattern(pattern.granularity, + sourceType.getGranularity())) + continue; + if (!matchesElementCountPattern(pattern.elementCounts, + sourceType.getElementCount())) + continue; + if (!matchesLayoutPattern(sourceType.getContext(), pattern.sourceLayout, + sourceLayout)) + continue; + if (!matchesLayoutPattern(resultType.getContext(), pattern.resultLayout, + resultLayout)) + continue; + return success(); + } + + return fail("source/result mask layouts do not match a supported " + "ensure_mask_layout table row"); +} + +FailureOr VMILayoutSupport::getEnsureLayoutFact( + VMIVRegType sourceType, VMIVRegType resultType, std::string *reason) const { + VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + if (failed(matchEnsureLayoutPattern(sourceType, resultType, sourceLayout, + resultLayout, reason))) + return failure(); + return VMIEnsureLayoutFact{sourceLayout, resultLayout}; +} + +FailureOr VMILayoutSupport::getEnsureMaskLayoutFact( + VMIMaskType sourceType, VMIMaskType resultType, std::string *reason) const { + VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + if (failed(matchEnsureMaskLayoutPattern(sourceType, resultType, sourceLayout, + resultLayout, reason))) + return failure(); + return VMIEnsureMaskLayoutFact{sourceLayout, resultLayout}; +} + +FailureOr VMILayoutSupport::getGroupSlotLoadLayoutFact( + VMIVRegType resultType, int64_t numGroups, std::string *reason) const { + auto fail = [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + VMILayoutAttr layout = resultType.getLayoutAttr(); + if (!layout) + return fail("requires assigned result layout"); + + if (!isSupportedGroupSlotMemoryLayout(layout, numGroups)) + return fail("result layout does not match a supported group_slot_load " + "table row"); + + return VMIGroupSlotLayoutFact{layout, numGroups, layout.getSlots()}; +} + +FailureOr +VMILayoutSupport::getGroupLoadLayoutFact(VMIGroupLoadOp op, + std::string *reason) const { + return getGroupLoadLayoutFact(cast(op.getResult().getType()), + op.getRowStride(), + op.getNumGroupsAttr().getInt(), reason); +} + +FailureOr VMILayoutSupport::getGroupLoadLayoutFact( + VMIVRegType resultType, Value rowStride, int64_t numGroups, + std::string *reason) const { + auto fail = [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + VMILayoutAttr layout = resultType.getLayoutAttr(); + if (!layout) + return fail("requires assigned result layout"); + + unsigned elementBits = + pto::getPTOStorageElemBitWidth(resultType.getElementType()); + if (elementBits == 0) + return fail("group_load requires known element bit width"); + std::optional stride = getConstantIndexValue(rowStride); + + FailureOr key = buildGroupLayoutKey( + resultType, numGroups, + "group_load layout table has no row for this group size", reason); + if (failed(key)) + return failure(); + + for (const GroupLoadLayoutPattern &pattern : kGroupLoadLayoutPatterns) { + if (!matchesElementBitsPattern(pattern.elementBits, elementBits)) + continue; + if (!matchesGroupBlockPattern(pattern.block, *key)) + continue; + if (!matchesGroupLoadMemoryPattern(pattern.memory, stride, key->groupSize, + elementBits)) + continue; + if (!matchesLayoutPattern(resultType.getContext(), pattern.resultLayout, + layout)) + continue; + return VMIGroupLoadLayoutFact{ + getGroupBlockClassFromPattern(pattern.block), layout, key->groupSize}; + } + + return fail("result layout, group size, and row_stride do not match a " + "supported group_load table row"); +} + +FailureOr VMILayoutSupport::getGroupStoreLayoutFact( + VMIVRegType valueType, int64_t numGroups, std::string *reason) const { + auto fail = [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + VMILayoutAttr layout = valueType.getLayoutAttr(); + if (!layout) + return fail("requires assigned value layout"); + if (!isSupportedGroupSlotMemoryLayout(layout, numGroups)) + return fail("value layout does not match a supported group_store table " + "row"); + return VMIGroupSlotLayoutFact{layout, numGroups, layout.getSlots()}; +} + +LogicalResult getGroupReduceAddSupportImpl(VMIVRegType sourceType, + VMIMaskType maskType, + VMIVRegType resultType, + int64_t numGroups, + std::string *reason) { + FailureOr fact = + VMILayoutSupport().getGroupReduceLayoutFactForLayouts( + sourceType, maskType, resultType, numGroups, reason); + if (failed(fact)) + return failure(); + return success(); +} + +LogicalResult +VMILayoutSupport::getGroupReduceAddFSupport(VMIGroupReduceAddFOp op, + std::string *reason) const { + return getGroupReduceAddSupportImpl( + cast(op.getSource().getType()), + cast(op.getMask().getType()), + cast(op.getResult().getType()), + op.getNumGroupsAttr().getInt(), reason); +} + +LogicalResult +VMILayoutSupport::getGroupReduceMaxFSupport(VMIGroupReduceMaxFOp op, + std::string *reason) const { + return getGroupReduceAddSupportImpl( + cast(op.getSource().getType()), + cast(op.getMask().getType()), + cast(op.getResult().getType()), + op.getNumGroupsAttr().getInt(), reason); +} + +LogicalResult +VMILayoutSupport::getGroupReduceMinFSupport(VMIGroupReduceMinFOp op, + std::string *reason) const { + return getGroupReduceAddSupportImpl( + cast(op.getSource().getType()), + cast(op.getMask().getType()), + cast(op.getResult().getType()), + op.getNumGroupsAttr().getInt(), reason); +} + +LogicalResult +VMILayoutSupport::getGroupReduceAddISupport(VMIGroupReduceAddIOp op, + std::string *reason) const { + return getGroupReduceAddSupportImpl( + cast(op.getSource().getType()), + cast(op.getMask().getType()), + cast(op.getResult().getType()), + op.getNumGroupsAttr().getInt(), reason); +} + +LogicalResult +VMILayoutSupport::getGroupReduceMaxISupport(VMIGroupReduceMaxIOp op, + std::string *reason) const { + return getGroupReduceAddSupportImpl( + cast(op.getSource().getType()), + cast(op.getMask().getType()), + cast(op.getResult().getType()), + op.getNumGroupsAttr().getInt(), reason); +} + +LogicalResult +VMILayoutSupport::getGroupReduceMinISupport(VMIGroupReduceMinIOp op, + std::string *reason) const { + return getGroupReduceAddSupportImpl( + cast(op.getSource().getType()), + cast(op.getMask().getType()), + cast(op.getResult().getType()), + op.getNumGroupsAttr().getInt(), reason); +} + +LogicalResult VMILayoutSupport::getGroupBroadcastSupport( + VMIGroupBroadcastOp op, std::string *reason) const { + return getGroupBroadcastSupport(cast(op.getSource().getType()), + cast(op.getResult().getType()), + op.getNumGroupsAttr().getInt(), reason); +} + +LogicalResult +VMILayoutSupport::getGroupBroadcastLoadSupport(VMIGroupBroadcastLoadOp op, + std::string *reason) const { + return success(succeeded(getGroupBroadcastLoadLayoutFact(op, reason))); +} + +LogicalResult VMILayoutSupport::getGroupBroadcastSupport( + VMIVRegType sourceType, VMIVRegType resultType, int64_t numGroups, + std::string *reason) const { + return success(succeeded(getGroupBroadcastLayoutFactForLayouts( + sourceType, resultType, numGroups, reason))); +} + +static LogicalResult matchSupplementalCastLayoutPattern( + VMIVRegType sourceType, VMIVRegType resultType, + ArrayRef patterns, + std::string *reason) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + if (!sourceLayout || !resultLayout) + return fail("requires assigned source/result layouts"); + + auto [sourceBits, resultBits] = getCastElementBits(sourceType, resultType); + int64_t numGroups = + sourceLayout.isGroupSlots() + ? sourceLayout.getNumGroups() + : (resultLayout.isGroupSlots() ? resultLayout.getNumGroups() : 0); + MLIRContext *ctx = sourceType.getContext(); + for (const SupplementalCastLayoutPattern &pattern : patterns) { + if (!matchesElementBitsPattern(pattern.sourceBits, sourceBits) || + !matchesElementBitsPattern(pattern.resultBits, resultBits)) + continue; + if (!matchesLayoutPattern(ctx, pattern.sourceLayout, sourceLayout, + numGroups)) + continue; + if (!matchesLayoutPattern(ctx, pattern.resultLayout, resultLayout, + numGroups)) + continue; + return success(); + } + + return fail("source/result layouts do not match a supplemental cast table row"); +} + +static LogicalResult getNarrowCastSupport(VMIVRegType sourceType, + VMIVRegType resultType, + std::string *reason) { + VMILayoutSupport support; + if (succeeded(support.getCastLayoutFactForLayouts( + sourceType, resultType, sourceType.getLayoutAttr(), + resultType.getLayoutAttr(), reason))) + return success(); + return matchSupplementalCastLayoutPattern( + sourceType, resultType, kSupplementalNarrowCastLayoutPatterns, reason); +} + +LogicalResult VMILayoutSupport::getTruncFSupport(VMITruncFOp op, + std::string *reason) const { + return getNarrowCastSupport(cast(op.getSource().getType()), + cast(op.getResult().getType()), + reason); +} + +LogicalResult VMILayoutSupport::getExtFSupport(VMIExtFOp op, + std::string *reason) const { + auto sourceType = cast(op.getSource().getType()); + auto resultType = cast(op.getResult().getType()); + return success(succeeded(getCastLayoutFactForLayouts( + sourceType, resultType, sourceType.getLayoutAttr(), + resultType.getLayoutAttr(), reason))); +} + +template +static LogicalResult getExtISupportImpl(OpT op, std::string *reason) { + auto sourceType = cast(op.getSource().getType()); + auto resultType = cast(op.getResult().getType()); + + FailureOr fact = + VMILayoutSupport().getCastLayoutFactForLayouts( + sourceType, resultType, sourceType.getLayoutAttr(), + resultType.getLayoutAttr(), reason); + if (succeeded(fact)) + return success(); + + return matchSupplementalCastLayoutPattern( + sourceType, resultType, kSupplementalIntegerExtLayoutPatterns, reason); +} + +LogicalResult VMILayoutSupport::getExtSISupport(VMIExtSIOp op, + std::string *reason) const { + return getExtISupportImpl(op, reason); +} + +LogicalResult VMILayoutSupport::getExtUISupport(VMIExtUIOp op, + std::string *reason) const { + return getExtISupportImpl(op, reason); +} + +LogicalResult +VMILayoutSupport::getTruncISupport(VMITruncIOp op, std::string *reason) const { + return getNarrowCastSupport(cast(op.getSource().getType()), + cast(op.getResult().getType()), + reason); +} + +FailureOr +VMILayoutSupport::getBitcastLayoutFact(VMIBitcastOp op, + std::string *reason) const { + auto fail = [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + auto sourceType = cast(op.getSource().getType()); + auto resultType = cast(op.getResult().getType()); + VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + if (!sourceLayout || !resultLayout) + return fail("requires assigned source and result layouts"); + if (sourceLayout != resultLayout) + return fail("requires matching source and result layouts"); + + int64_t numGroups = + sourceLayout.isGroupSlots() ? sourceLayout.getNumGroups() : 0; + unsigned sourceElementBits = + pto::getPTOStorageElemBitWidth(sourceType.getElementType()); + unsigned resultElementBits = + pto::getPTOStorageElemBitWidth(resultType.getElementType()); + if (sourceElementBits == 0 || resultElementBits == 0) + return fail("requires source and result with known storage element width"); + // Equal-width bitcast is layout-transparent for any identical layout. Only + // width-changing bitcast needs a table row because not every layout has a + // representation-preserving carrier reinterpretation across element widths. + if (sourceElementBits != resultElementBits) { + bool matchedLayout = false; + MLIRContext *ctx = op.getContext(); + for (const WidthChangingBitcastLayoutPattern &pattern : + kWidthChangingBitcastLayoutPatterns) { + if (matchesLayoutPattern(ctx, pattern.layout, sourceLayout, numGroups) && + matchesLayoutPattern(ctx, pattern.layout, resultLayout, numGroups)) { + matchedLayout = true; + break; + } + } + if (!matchedLayout) + return fail("width-changing bitcast layout does not match a bitcast " + "layout table row"); + } + + return VMIBitcastLayoutFact{sourceLayout, resultLayout}; +} + +LogicalResult VMILayoutSupport::getBitcastSupport(VMIBitcastOp op, + std::string *reason) const { + return getBitcastLayoutFact(op, reason); +} + +template +static FailureOr +getHistogramLayoutFactImpl(OpTy op, ArrayRef patterns, + StringRef opName, std::string *reason) { + auto fail = [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + if (patterns.empty()) + return fail(opName + " histogram layout table has no row"); + + auto accType = cast(op.getAcc().getType()); + auto sourceType = cast(op.getSource().getType()); + auto maskType = cast(op.getMask().getType()); + auto resultType = cast(op.getResult().getType()); + + VMILayoutAttr accLayout = accType.getLayoutAttr(); + VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); + VMILayoutAttr maskLayout = maskType.getLayoutAttr(); + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + if (!accLayout || !sourceLayout || !maskLayout || !resultLayout) + return fail("requires assigned acc/source/mask/result layouts"); + + MLIRContext *ctx = op.getContext(); + for (const HistogramLayoutPattern &pattern : patterns) { + if (!matchesLayoutPattern(ctx, pattern.accLayout, accLayout) || + !matchesLayoutPattern(ctx, pattern.sourceLayout, sourceLayout) || + !matchesLayoutPattern(ctx, pattern.maskLayout, maskLayout) || + !matchesLayoutPattern(ctx, pattern.resultLayout, resultLayout)) + continue; + + VMIHistogramLayoutFact fact; + fact.accLayout = accLayout; + fact.sourceLayout = sourceLayout; + fact.maskLayout = maskLayout; + fact.resultLayout = resultLayout; + return fact; + } + + return fail(opName + " acc/source/mask/result layouts do not match a " + "histogram layout table row"); +} + +FailureOr +VMILayoutSupport::getVdhistLayoutFact(VMIVdhistOp op, + std::string *reason) const { + return getHistogramLayoutFactImpl(op, kVdhistLayoutPatterns, "vdhist", reason); +} + +FailureOr +VMILayoutSupport::getVchistLayoutFact(VMIVchistOp op, + std::string *reason) const { + // vchist shares the same layout constraints as vdhist (same base class, same + // signature). When kVdhistLayoutPatterns is updated, review whether vchist + // should inherit the new patterns. + return getHistogramLayoutFactImpl(op, kVdhistLayoutPatterns, "vchist", reason); +} + +LogicalResult +VMILayoutSupport::getVdhistSupport(VMIVdhistOp op, std::string *reason) const { + return getVdhistLayoutFact(op, reason); +} + +LogicalResult +VMILayoutSupport::getVchistSupport(VMIVchistOp op, std::string *reason) const { + return getVchistLayoutFact(op, reason); +} diff --git a/lib/PTO/Transforms/VMILegalizeArithSelect.cpp b/lib/PTO/Transforms/VMILegalizeArithSelect.cpp new file mode 100644 index 0000000000..ade95024f7 --- /dev/null +++ b/lib/PTO/Transforms/VMILegalizeArithSelect.cpp @@ -0,0 +1,86 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +//===- VMILegalizeArithSelect.cpp - Legalize VMI arith.select ------------===// +//===----------------------------------------------------------------------===// + +#include "PTO/IR/PTO.h" +#include "PTO/Transforms/Passes.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" +#include "llvm/ADT/STLExtras.h" + +namespace mlir { +namespace pto { +#define GEN_PASS_DEF_VMILEGALIZEARITHSELECT +#include "PTO/Transforms/Passes.h.inc" +} // namespace pto +} // namespace mlir + +using namespace mlir; +using namespace mlir::pto; + +namespace { + +static bool isVMIValueType(Type type) { + return isa(type); +} + +static bool hasScalarI1Condition(arith::SelectOp select) { + return select.getCondition().getType().isSignlessInteger(1); +} + +static void rewriteSelectToIf(arith::SelectOp select) { + OpBuilder builder(select); + auto ifOp = builder.create( + select.getLoc(), TypeRange{select.getResult().getType()}, + select.getCondition(), /*withElseRegion=*/true); + + { + OpBuilder::InsertionGuard guard(builder); + builder.setInsertionPointToStart(&ifOp.getThenRegion().front()); + builder.create(select.getLoc(), select.getTrueValue()); + builder.setInsertionPointToStart(&ifOp.getElseRegion().front()); + builder.create(select.getLoc(), select.getFalseValue()); + } + + select.getResult().replaceAllUsesWith(ifOp.getResult(0)); + select.erase(); +} + +struct VMILegalizeArithSelectPass + : public mlir::pto::impl::VMILegalizeArithSelectBase< + VMILegalizeArithSelectPass> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(VMILegalizeArithSelectPass) + + void runOnOperation() override { + ModuleOp module = getOperation(); + SmallVector selects; + module.walk([&](arith::SelectOp select) { + if (isVMIValueType(select.getResult().getType()) && + hasScalarI1Condition(select)) + selects.push_back(select); + }); + + for (arith::SelectOp select : llvm::reverse(selects)) { + if (select->getBlock() != nullptr) + rewriteSelectToIf(select); + } + } +}; + +} // namespace + +std::unique_ptr mlir::pto::createVMILegalizeArithSelectPass() { + return std::make_unique(); +} diff --git a/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp b/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp new file mode 100644 index 0000000000..bbb1c5e587 --- /dev/null +++ b/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp @@ -0,0 +1,1639 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +//===- VMILowerUnifiedToLegacy.cpp - Lower unified v-ops to legacy ops ----===// +// +// Lowers unified v-prefixed VMI ops to their legacy equivalents under the +// opt-in --vmi-two-stage-lowering flag. +// +// Category A — pure syntactic renames (4 ops): +// vci → iota +// vinterpret_cast → bitcast +// vsel → select +// vbrc → broadcast (skipped when num_groups is present) +// +// Category B — elementwise arithmetic / bitwise (18 ops): +// vadd/vsub/vmul/vdiv/vmin/vmax → legacy type-specific binary op +// vneg/vabs/vsqrt/vexp/vln/vrelu → legacy unary op +// vand/vor/vxor/vshl/vshr/vnot → legacy bitwise op +// vshr selects shrui for explicit unsigned elements and shrsi for +// signless/signed elements. +// Mask/pmode synthesis is intentionally bypassed here so two-stage lowering +// does not introduce select chains before layout assignment. +// +// Category C1 — compare + seed (2 ops): +// vcmp → cmpf/cmpi + mask_and +// vcmps → broadcast scalar + cmpf/cmpi + mask_and +// +// Category C2 — unified type conversion (1 op): +// vcvt → type-dispatch to extf/truncf/fptosi/sitofp/extsi/extui/trunci +// For fp narrowing, unified saturate=SAT is normalized away because the +// legacy truncf -> VPTO lowering already materializes saturating low-level +// vcvt forms for supported narrowing result families. +// +// Category C3 — unified load/store (2 ops): +// vload → dispatch by dist_mode/group/block_stride to +// load / deinterleave_load / group_broadcast_load{num_groups=1} / ... +// vstore → dispatch to store / masked_store / interleave_store / group_store / ... +// Skipped: dist_mode "unpack" (physical widening, no legacy equivalent). +// +// Category C4 — static mask creation (3 ops): +// pset → create_mask(all lanes) +// pge → create_mask(N lanes) +// plt → create_mask(min(rem, L)) +// +// Category C4 — static mask creation (3 ops): +// pset → create_mask(all lanes) +// pge → create_mask(N lanes) +// plt → create_mask(min(rem, L)) +// +// Category C5 — vector-scalar ops, one-step to legacy (6 ops): +// vadds/vmuls/vmaxs/vmins/vshls/vshrs +// → broadcast scalar → legacy binary +// vshrs selects shrui for explicit unsigned elements and shrsi for +// signless/signed elements. +// +// Category C3 — unified load/store (2 ops, dispatch by dist_mode/group): +// vload → load / deinterleave_load / group_load +// vstore → store / masked_store / interleave_store / group_store +// +// Category C6 — unified reduce (3 ops): +// vcadd → reduce_addf/reduce_addi or group_reduce_addf/group_reduce_addi +// vcmax → reduce_maxf/reduce_maxi or group_reduce_maxf/group_reduce_maxi +// vcmin → reduce_minf/reduce_mini or group_reduce_minf/group_reduce_mini +// +// Category C7 — fused multiply-add family → legacy fma (2 ops): +// vmula → fma (float only; mask discarded; int → skipped, no legacy int fma) +// vaxpy → broadcast + fma (float only) +// +// Category C8 — indexed gather/scatter → legacy gather/scatter (2 ops): +// vgather → gather (pmode="zero": passthru = zero constant) +// vscatter → scatter +// +// Category C9 — fused activation / softmax, decomposed to legacy chains (3 ops): +// vexpdif → [extf] + subf + exp (widen f16 x to f32 when needed) +// vlrelu → maxf + minf + broadcast + mulf + addf +// vprelu → maxf + minf + mulf + addf +// Category C7/C8/C9 bypass mask/pmode synthesis here and skip pmode="merge". +// +// Category D — no legacy equivalent (explicitly skipped, 5 ops): +// vintlv vdintlv vselr vgatherb vmull +// +//===----------------------------------------------------------------------===// + +#include "PTO/IR/PTO.h" +#include "PTO/Transforms/Passes.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/Pass/Pass.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/StringExtras.h" +#include "llvm/Support/ErrorHandling.h" + +namespace mlir { +namespace pto { +#define GEN_PASS_DEF_VMILOWERUNIFIEDTOLEGACY +#include "PTO/Transforms/Passes.h.inc" +} // namespace pto +} // namespace mlir + +using namespace mlir; +using namespace mlir::pto; + +//===----------------------------------------------------------------------===// +// Helpers +//===----------------------------------------------------------------------===// + +/// Returns the string name of a predicate mode, defaulting to "zero". +static StringRef getPmodeOrDefault(Operation *op, StringRef attrName = "pmode") { + if (auto attr = op->getAttrOfType(attrName)) + return attr.getValue(); + return "zero"; +} + +/// Returns true when the pmode on `op` is "merge" — these ops must be +/// skipped because merge semantic (inactive lane preserves OLD_DEST) cannot +/// be expressed in VMI SSA IR. +static bool hasMergePmode(Operation *op) { + return getPmodeOrDefault(op) == "merge"; +} + +/// Create a zero-valued VMIConstantOp with the same type as \p vmiType. +static Value createZeroConstant(OpBuilder &builder, Location loc, + VMIVRegType vmiType) { + Type elemType = vmiType.getElementType(); + int64_t laneCount = vmiType.getElementCount(); + auto shapedType = RankedTensorType::get({laneCount}, elemType); + + DenseElementsAttr zeroAttr; + if (auto floatType = dyn_cast(elemType)) { + zeroAttr = DenseElementsAttr::get( + shapedType, APFloat::getZero(floatType.getFloatSemantics())); + } else if (auto intType = dyn_cast(elemType)) { + zeroAttr = DenseElementsAttr::get( + shapedType, APInt::getZero(intType.getWidth())); + } else { + llvm_unreachable("unsupported VMI element type for zero constant"); + } + return builder.create(loc, vmiType, zeroAttr).getResult(); +} + + +/// Create a 1-lane VMIConstantOp with the neutral element for reduction: +/// add: 0 (int and float) +/// max: -INF (float), INT_MIN (int) +/// min: +INF (float), INT_MAX (int) +static Value createReduceNeutralInit(OpBuilder &builder, Location loc, + Type elemType, bool isAdd, bool isMax, + Attribute layout = Attribute()) { + auto oneLaneType = + VMIVRegType::get(builder.getContext(), 1, elemType, layout); + auto shapedType = RankedTensorType::get({1}, elemType); + DenseElementsAttr attr; + if (auto floatTy = dyn_cast(elemType)) { + if (isAdd) + attr = DenseElementsAttr::get( + shapedType, APFloat::getZero(floatTy.getFloatSemantics())); + else if (isMax) + attr = DenseElementsAttr::get( + shapedType, + APFloat::getInf(floatTy.getFloatSemantics(), /*Negative=*/true)); + else + attr = DenseElementsAttr::get( + shapedType, + APFloat::getInf(floatTy.getFloatSemantics(), /*Negative=*/false)); + } else { + auto intTy = cast(elemType); + if (isAdd) + attr = DenseElementsAttr::get(shapedType, + APInt::getZero(intTy.getWidth())); + else if (isMax) + attr = DenseElementsAttr::get( + shapedType, intTy.isUnsigned() + ? APInt::getZero(intTy.getWidth()) + : APInt::getSignedMinValue(intTy.getWidth())); + else + attr = DenseElementsAttr::get( + shapedType, intTy.isUnsigned() + ? APInt::getMaxValue(intTy.getWidth()) + : APInt::getSignedMaxValue(intTy.getWidth())); + } + return builder.create(loc, oneLaneType, attr).getResult(); +} + +/// Map a unified vcmp `cmp` mode to the predicate string for legacy +/// cmpf/cmpi. Float operands use ordered predicates (olt, oeq, ...); +/// integer operands select signedness from the element type. +static std::string mapCmpPredicate(StringRef cmp, Type elemType, + bool isFloat) { + if (isFloat) { + // Already ordered/unordered — pass through. + if (cmp.starts_with("o") || cmp.starts_with("u")) + return cmp.str(); + return ("o" + cmp).str(); // e.g. "lt" → "olt" + } + // Integer. + if (cmp.starts_with("s") || cmp.starts_with("u")) + return cmp.str(); + // eq/ne are valid for both fp and int without prefix. + if (cmp == "eq" || cmp == "ne") + return cmp.str(); + auto intType = dyn_cast(elemType); + if (intType && intType.isUnsigned()) + return ("u" + cmp).str(); // e.g. "lt" -> "ult" + return ("s" + cmp).str(); // e.g. "lt" -> "slt" +} + +/// Return true when \p elemType is a floating-point type. +static bool isFloatType(Type elemType) { + return isa(elemType); +} + +/// Return the element type of a VMIVRegType. +static Type getVMIElementType(Value v) { + return cast(v.getType()).getElementType(); +} + +/// Inspect the source and result element types of a vcvt and classify the +/// conversion direction. Returns one of: +/// "widen_fp", "narrow_fp", "fptosi", "sitofp", +/// "widen_int", "narrow_int" +/// conversion direction. Returns one of: +/// "widen_fp", "narrow_fp", "fptosi", "sitofp", +/// "widen_int", "narrow_int" +static StringRef classifyCvtDirection(Type srcElem, Type dstElem) { + bool srcFp = isFloatType(srcElem); + bool dstFp = isFloatType(dstElem); + unsigned srcBits = srcElem.getIntOrFloatBitWidth(); + unsigned dstBits = dstElem.getIntOrFloatBitWidth(); + + if (srcFp && dstFp) + return dstBits > srcBits ? "widen_fp" : "narrow_fp"; + if (srcFp && !dstFp) + return "fptosi"; + if (!srcFp && dstFp) + return "sitofp"; + // int → int + return dstBits > srcBits ? "widen_int" : "narrow_int"; +} + +//===----------------------------------------------------------------------===// +// Category B: binary elementwise → legacy compute, mask/pmode discarded +//===----------------------------------------------------------------------===// + +/// Lower a BINARY unified op (vadd, vsub, ...) to a legacy compute op. +/// Unified mask and pmode are intentionally discarded. +/// +/// \p createLegacy is a callable `(Location, Type, Value, Value) -> Value` +/// that emits the legacy binary op. +template +static LogicalResult +lowerBinaryIgnoringMask( + UnifiedOp op, + function_ref createLegacy) { + if (hasMergePmode(op)) + return failure(); + Location loc = op.getLoc(); + Type resultType = op.getResult().getType(); + Value lhs = op.getLhs(); + Value rhs = op.getRhs(); + + Value raw = createLegacy(loc, resultType, lhs, rhs); + op.getResult().replaceAllUsesWith(raw); + op->erase(); + return success(); +} + +/// Lower a UNARY unified op (vneg, vabs, …) to its legacy counterpart. +template +static LogicalResult +lowerMaskedUnary(UnifiedOp op, OpBuilder &builder, + function_ref createLegacy) { + if (hasMergePmode(op)) + return failure(); + + Location loc = op.getLoc(); + Type resultType = op.getResult().getType(); + Value source = op.getSource(); + + Value raw = createLegacy(loc, resultType, source); + op.getResult().replaceAllUsesWith(raw); + op->erase(); + return success(); +} + +//===----------------------------------------------------------------------===// +// Category C1 helpers: vcmp / vcmps +//===----------------------------------------------------------------------===// + +/// Returns true if `seed` is provably an all-active mask (every lane active), +/// so `mask_and(x, seed)` is the identity and the AND can be skipped. Covers a +/// `pset` (all lanes active by definition) and a `create_mask` whose +/// active_lanes is a constant >= the mask lane count. +static bool isAllActiveSeed(Value seed) { + Operation *def = seed.getDefiningOp(); + if (!def) + return false; + if (isa(def)) + return true; + if (auto cm = dyn_cast(def)) { + auto maskTy = cast(cm.getResult().getType()); + if (auto cst = cm.getActiveLanes().getDefiningOp()) + if (auto ia = dyn_cast(cst.getValue())) + return ia.getInt() >= maskTy.getElementCount(); + } + return false; +} + +/// Lower vcmp to cmpf/cmpi + mask_and. +static LogicalResult lowerVCmp(VMIVcmpOp op, OpBuilder &builder) { + if (hasMergePmode(op)) + return failure(); + + Location loc = op.getLoc(); + Type elemType = getVMIElementType(op.getLhs()); + bool isFloat = isFloatType(elemType); + StringRef cmpMode = op.getCmp(); + std::string predicate = mapCmpPredicate(cmpMode, elemType, isFloat); + + // Build legacy cmpf or cmpi. + Value rawMask; + if (isFloat) { + rawMask = builder + .create(loc, op.getResult().getType(), + builder.getStringAttr(predicate), + op.getLhs(), op.getRhs()) + .getResult(); + } else { + rawMask = builder + .create(loc, op.getResult().getType(), + builder.getStringAttr(predicate), + op.getLhs(), op.getRhs()) + .getResult(); + } + + // mask_and with seed — skipped when the seed is all-active (identity AND). + Value result = rawMask; + if (!isAllActiveSeed(op.getSeed())) + result = builder + .create(loc, op.getResult().getType(), rawMask, + op.getSeed()) + .getResult(); + + op.getResult().replaceAllUsesWith(result); + op->erase(); + return success(); +} + +/// Lower vcmps to broadcast scalar + cmpf/cmpi + mask_and. +static LogicalResult lowerVCmps(VMIVcmpsOp op, OpBuilder &builder) { + if (hasMergePmode(op)) + return failure(); + + Location loc = op.getLoc(); + Type srcVmiType = op.getSrc().getType(); + Value scalar = op.getScalar(); + Type elemType = getVMIElementType(op.getSrc()); + bool isFloat = isFloatType(elemType); + StringRef cmpMode = op.getCmp(); + std::string predicate = mapCmpPredicate(cmpMode, elemType, isFloat); + + // 1. Broadcast scalar to vector. + Value brc = builder.create(loc, srcVmiType, scalar) + .getResult(); + + // 2. Legacy cmpf or cmpi. + Value rawMask; + if (isFloat) { + rawMask = builder + .create(loc, op.getResult().getType(), + builder.getStringAttr(predicate), + op.getSrc(), brc) + .getResult(); + } else { + rawMask = builder + .create(loc, op.getResult().getType(), + builder.getStringAttr(predicate), + op.getSrc(), brc) + .getResult(); + } + + // 3. mask_and with seed — skipped when the seed is all-active (identity AND). + Value result = rawMask; + if (!isAllActiveSeed(op.getSeed())) + result = builder + .create(loc, op.getResult().getType(), rawMask, + op.getSeed()) + .getResult(); + + op.getResult().replaceAllUsesWith(result); + op->erase(); + return success(); +} + +//===----------------------------------------------------------------------===// +// Category C2 helper: vcvt +//===----------------------------------------------------------------------===// + +/// Lower vcvt by dispatching on src→dst element types. +static LogicalResult lowerVCvt(VMICvtOp op, OpBuilder &builder) { + if (hasMergePmode(op)) + return failure(); + + Type srcElem = getVMIElementType(op.getSource()); + Type dstElem = getVMIElementType(op.getResult()); + StringRef direction = classifyCvtDirection(srcElem, dstElem); + Location loc = op.getLoc(); + Type resultType = op.getResult().getType(); + Value source = op.getSource(); + Value result; + + if (direction == "widen_fp") { + result = builder.create(loc, resultType, source).getResult(); + } else if (direction == "narrow_fp") { + StringAttr roundingAttr = op.getRoundingAttr(); + result = + builder.create(loc, resultType, source, roundingAttr) + .getResult(); + } else if (direction == "fptosi") { + result = + builder.create(loc, resultType, source).getResult(); + } else if (direction == "sitofp") { + result = + builder.create(loc, resultType, source).getResult(); + } else if (direction == "widen_int") { + // Use source type signedness to decide signed vs unsigned extension. + bool useSigned = true; + if (auto intTy = dyn_cast(srcElem)) { + useSigned = intTy.isSigned(); + } + if (useSigned) + result = + builder.create(loc, resultType, source).getResult(); + else + result = + builder.create(loc, resultType, source).getResult(); + } else if (direction == "narrow_int") { + // trunci already has saturating semantics. + result = + builder.create(loc, resultType, source).getResult(); + } else { + return failure(); + } + + op.getResult().replaceAllUsesWith(result); + op->erase(); + return success(); +} + +//===----------------------------------------------------------------------===// +// Category C3 helpers: vload / vstore +//===----------------------------------------------------------------------===// + +/// Lower vload by dispatching on dist_mode. +static LogicalResult lowerVLoad(VMIvLoadOp op, OpBuilder &builder) { + // Group mode: vload {group=C} → group_load or group_slot_load + // elementCount == num_groups → group_slot_load (compact, 1 scalar per group) + // elementCount > num_groups → group_load (full groups) + // Broadcast-style slot loads (dense output, stride=1) should be decomposed + // at the VMI level: vload(compact) + vbrc. + if (op.getGroupAttr()) { + auto resultType = cast(op.getResults().front().getType()); + int64_t numGroups = op.getGroupAttr().getInt(); + // {group, dist_mode="brc"}: group broadcast — one scalar per group + // broadcast within each group (e.g. E2B). + if (op.getDistMode() && op.getDistMode() == "brc") { + auto gbl = builder.create( + op->getLoc(), resultType, op.getSource(), op.getOffset(), + op.getStride(), op.getGroupAttr()); + op.getResults().front().replaceAllUsesWith(gbl.getResult()); + } else if (resultType.getElementCount() == numGroups) { + auto slotLoad = builder.create( + op->getLoc(), resultType, op.getSource(), op.getOffset(), + op.getStride(), op.getGroupAttr()); + op.getResults().front().replaceAllUsesWith(slotLoad.getResult()); + } else { + auto groupLoad = builder.create( + op->getLoc(), resultType, op.getSource(), op.getOffset(), + op.getStride(), op.getGroupAttr()); + op.getResults().front().replaceAllUsesWith(groupLoad.getResult()); + } + op->erase(); + return success(); + } + + // Block-stride mode: vload {block_stride, repeat_stride} → stride_load + if (op.getBlockStride()) { + auto resultType = op.getResults().front().getType(); + // Create default all-active mask via create_mask with proper granularity + // and layout, so the legacy stride_load passes both-or-neither layout + // checks. pset cannot be used here because pred masks cannot carry + // layout. + auto resultVMIType = cast(resultType); + auto elemType = resultVMIType.getElementType(); + unsigned bits = 32; + if (auto it = dyn_cast(elemType)) + bits = it.getWidth(); + else if (auto ft = dyn_cast(elemType)) + bits = ft.getWidth(); + auto gran = StringAttr::get(builder.getContext(), + bits <= 8 ? "b8" : bits <= 16 ? "b16" : "b32"); + auto maskType = VMIMaskType::get(builder.getContext(), + resultVMIType.getElementCount(), gran, + resultVMIType.getLayout()); + auto fullLanes = builder.create( + op->getLoc(), builder.getIndexAttr(resultVMIType.getElementCount())); + auto mask = builder.create(op->getLoc(), maskType, + fullLanes.getResult()); + Value bs = op.getBlockStride(); + Value rs = op.getRepeatStride(); + auto strideLoad = builder.create( + op->getLoc(), resultType, op.getSource(), op.getOffset(), bs, rs, + mask.getResult()); + op.getResults().front().replaceAllUsesWith(strideLoad.getResult()); + op->erase(); + return success(); + } + + // pmode="merge" cannot be expressed by legacy load + select — skip. + if (hasMergePmode(op)) + return failure(); + + StringAttr distModeAttr = op.getDistModeAttr(); + StringRef distMode = + distModeAttr ? distModeAttr.getValue() : "continuous"; + + Location loc = op.getLoc(); + Value source = op.getSource(); + Value offset = op.getOffset(); + + if (distMode == "continuous") { + auto loadOp = builder.create( + loc, op.getResults().front().getType(), source, offset); + op.getResults().front().replaceAllUsesWith(loadOp.getResult()); + } else if (distMode == "dintlv") { + auto dloadOp = builder.create( + loc, op.getResults()[0].getType(), op.getResults()[1].getType(), + source, offset); + op.getResults()[0].replaceAllUsesWith(dloadOp.getLow()); + op.getResults()[1].replaceAllUsesWith(dloadOp.getHigh()); + } else if (distMode == "brc") { + // vload {dist_mode="brc"} -> group_broadcast_load. + // - no group attr: all-lane scalar broadcast (num_groups=1, stride=0). + // - {group = C, stride}: per-group scalar broadcast — one scalar loaded + // per group and broadcast within each group (e.g. E2B). + auto resultType = op.getResults().front().getType(); + if (auto groupAttr = op.getGroupAttr()) { + int64_t numGroups = groupAttr.getInt(); + Value stride = op.getStride(); + auto gbl = builder.create( + loc, resultType, source, offset, stride, + builder.getI64IntegerAttr(numGroups)); + op.getResults().front().replaceAllUsesWith(gbl.getResult()); + } else { + Value zeroStride = builder.create( + loc, builder.getIndexType(), builder.getIndexAttr(0)); + auto gbl = builder.create( + loc, resultType, source, offset, zeroStride, + builder.getI64IntegerAttr(1)); + op.getResults().front().replaceAllUsesWith(gbl.getResult()); + } + } else { + // "unpack" has no legacy equivalent (physical widening, lane count changes). + return failure(); + } + + op->erase(); + return success(); +} + +/// Lower vstore by dispatching on dist_mode. +static LogicalResult lowerVStore(VMIvStoreOp op, OpBuilder &builder) { + // Group mode: vstore {group=C} → group_store + if (op.getGroupAttr()) { + builder.create( + op->getLoc(), op.getValues()[0], op.getDestination(), op.getOffset(), + op.getStride(), op.getGroupAttr()); + op->erase(); + return success(); + } + + // Block-stride mode: vstore {block_stride, repeat_stride} → stride_store + if (op.getBlockStride()) { + auto valueType = cast(op.getValues()[0].getType()); + // Use existing mask or create default all-active via create_mask + // (pset cannot be used — pred masks cannot carry layout). + Value mask; + if (!op.getMask().empty()) { + mask = op.getMask()[0]; + } else { + auto elemType = valueType.getElementType(); + unsigned bits = 32; + if (auto it = dyn_cast(elemType)) + bits = it.getWidth(); + else if (auto ft = dyn_cast(elemType)) + bits = ft.getWidth(); + auto gran = StringAttr::get(builder.getContext(), + bits <= 8 ? "b8" : bits <= 16 ? "b16" : "b32"); + auto maskType = VMIMaskType::get(builder.getContext(), + valueType.getElementCount(), gran, + valueType.getLayout()); + auto fullLanes = builder.create( + op->getLoc(), builder.getIndexAttr(valueType.getElementCount())); + mask = builder.create(op->getLoc(), maskType, + fullLanes.getResult()) + .getResult(); + } + Value bs = op.getBlockStride(); + Value rs = op.getRepeatStride(); + builder.create(op->getLoc(), op.getValues()[0], + op.getDestination(), op.getOffset(), bs, rs, + mask); + op->erase(); + return success(); + } + + StringAttr distModeAttr = op.getDistModeAttr(); + StringRef distMode = + distModeAttr ? distModeAttr.getValue() : "continuous"; + + Location loc = op.getLoc(); + Value dest = op.getDestination(); + Value offset = op.getOffset(); + auto values = op.getValues(); + + if (distMode == "continuous") { + if (values.empty()) + return failure(); + Value mask = op.getMask().empty() ? Value() : op.getMask().front(); + if (mask) { + // Masked store path. + builder.create(loc, values[0], dest, offset, mask); + } else { + builder.create(loc, values[0], dest, offset); + } + } else if (distMode == "dintlv") { + if (values.size() < 2) + return failure(); + builder.create(loc, values[0], values[1], dest, + offset); + } else { + return failure(); + } + + op->erase(); + return success(); +} + +//===----------------------------------------------------------------------===// +// Category C4 helpers: pset / pge +//===----------------------------------------------------------------------===// + +/// Lower pset "PAT_ALL" → create_mask(all_lanes). +static LogicalResult lowerPset(VMIPsetOp op, OpBuilder &builder) { + // If an all-active consumer (e.g. vcmp) elided its use, drop the seed + // entirely instead of materialising a dead create_mask. + if (op.use_empty()) { + op->erase(); + return success(); + } + Location loc = op.getLoc(); + auto maskType = cast(op.getResult().getType()); + int64_t laneCount = maskType.getElementCount(); + auto indexType = IndexType::get(builder.getContext()); + Value activeLanes = builder.create( + loc, indexType, builder.getIndexAttr(laneCount)); + Value result = + builder.create(loc, maskType, activeLanes).getResult(); + op.getResult().replaceAllUsesWith(result); + op->erase(); + return success(); +} + +/// Lower pge "PAT_VLN" → create_mask(N). +/// When {group = C} is present → create_group_mask(N, num_groups=C, +/// group_size = total_lanes / C). +static LogicalResult lowerPge(VMIPgeOp op, OpBuilder &builder) { + StringRef pattern = op.getPattern(); + // Parse "PAT_VL" or fall back to "PAT_VL16". + int64_t numLanes = 16; + if (pattern.starts_with("PAT_VL")) { + StringRef numStr = pattern.drop_front(6); // strlen("PAT_VL") + if (!numStr.empty()) { + int64_t parsed = 0; + for (char c : numStr) { + if (c < '0' || c > '9') + break; + parsed = parsed * 10 + (c - '0'); + } + if (parsed > 0) + numLanes = parsed; + } + } + + Location loc = op.getLoc(); + auto maskType = cast(op.getResult().getType()); + auto indexType = IndexType::get(builder.getContext()); + Value activeLanes = builder.create( + loc, indexType, builder.getIndexAttr(numLanes)); + + if (auto groupAttr = op.getGroupAttr()) { + // Grouped tail mask → create_group_mask + int64_t numGroups = groupAttr.getInt(); + int64_t totalLanes = maskType.getElementCount(); + int64_t groupSize = totalLanes / numGroups; + Value result = builder + .create( + loc, maskType, activeLanes, + builder.getI64IntegerAttr(numGroups), + builder.getI64IntegerAttr(groupSize)) + .getResult(); + op.getResult().replaceAllUsesWith(result); + } else { + Value result = + builder.create(loc, maskType, activeLanes).getResult(); + op.getResult().replaceAllUsesWith(result); + } + op->erase(); + return success(); +} + +//===----------------------------------------------------------------------===// +// Category C5 helpers: vector-scalar ops (one-step to legacy) +//===----------------------------------------------------------------------===// + +/// Lower a unified vector-scalar op (vadds, vmuls, ...) to a legacy chain: +/// %brc = vmi.broadcast %scalar +/// %raw = legacy.op %src, %brc +template +static LogicalResult +lowerVecScalar(VecScalarOp op, OpBuilder &builder, + function_ref createLegacy) { + Location loc = op.getLoc(); + Type srcVmiType = op.getSrc().getType(); + Value src = op.getSrc(); + Value scalar = op.getScalar(); + + Value brc = builder.create(loc, srcVmiType, scalar) + .getResult(); + Value raw = createLegacy(loc, srcVmiType, src, brc); + op.getResult().replaceAllUsesWith(raw); + op->erase(); + return success(); +} + +//===----------------------------------------------------------------------===// +// Category C6 helpers: vcadd / vcmax / vcmin +//===----------------------------------------------------------------------===// + +/// Lower vcadd to legacy reduce_addf/reduce_addi or +/// group_reduce_addf/group_reduce_addi. Always succeeds for valid input +/// (vcadd verifier guarantees reassoc for float, and group 整除 source lanes). +static LogicalResult lowerVCadd(VMIvcaddOp op, OpBuilder &builder) { + if (hasMergePmode(op)) + return failure(); + + auto sourceType = cast(op.getSource().getType()); + Type elemType = sourceType.getElementType(); + bool isFloat = isa(elemType); + Location loc = op.getLoc(); + Type resultType = op.getResult().getType(); + Value source = op.getSource(); + Value mask = op.getMask(); + + if (auto groupAttr = op.getGroupAttr()) { + // Group reduce path + int64_t C = groupAttr.getInt(); + Value result; + if (isFloat) + result = + builder + .create(loc, resultType, source, mask, + builder.getI64IntegerAttr(C), + op.getReassocAttr()) + .getResult(); + else + result = + builder + .create(loc, resultType, source, mask, + builder.getI64IntegerAttr(C)) + .getResult(); + op.getResult().replaceAllUsesWith(result); + } else { + // Full reduce path + Value init = createReduceNeutralInit(builder, loc, elemType, + /*isAdd=*/true, /*isMax=*/false, + sourceType.getLayout()); + Value result; + if (isFloat) + result = + builder + .create(loc, resultType, source, init, mask, + op.getReassocAttr()) + .getResult(); + else + result = + builder + .create(loc, resultType, source, init, mask) + .getResult(); + op.getResult().replaceAllUsesWith(result); + } + op->erase(); + return success(); +} + +/// Lower vcmax to legacy full or grouped float/integer maximum reduction. +static LogicalResult lowerVcmax(VMIvcmaxOp op, OpBuilder &builder) { + if (hasMergePmode(op)) + return failure(); + + auto sourceType = cast(op.getSource().getType()); + Type elemType = sourceType.getElementType(); + bool isFloat = isa(elemType); + Location loc = op.getLoc(); + Type resultType = op.getResult().getType(); + Value source = op.getSource(); + Value mask = op.getMask(); + + if (auto groupAttr = op.getGroupAttr()) { + // Group reduce path + int64_t C = groupAttr.getInt(); + Value result; + if (isFloat) + result = + builder + .create(loc, resultType, source, mask, + builder.getI64IntegerAttr(C)) + .getResult(); + else + result = + builder + .create(loc, resultType, source, mask, + builder.getI64IntegerAttr(C)) + .getResult(); + op.getResult().replaceAllUsesWith(result); + op->erase(); + return success(); + } + + Value init = createReduceNeutralInit(builder, loc, elemType, + /*isAdd=*/false, /*isMax=*/true, + sourceType.getLayout()); + Value result; + if (isFloat) + result = builder + .create(loc, resultType, source, init, mask) + .getResult(); + else + result = builder + .create(loc, resultType, source, init, mask) + .getResult(); + op.getResult().replaceAllUsesWith(result); + op->erase(); + return success(); +} + +/// Lower vcmin to legacy full or grouped float/integer minimum reduction. +static LogicalResult lowerVcmin(VMIvcminOp op, OpBuilder &builder) { + if (hasMergePmode(op)) + return failure(); + + auto sourceType = cast(op.getSource().getType()); + Type elemType = sourceType.getElementType(); + bool isFloat = isa(elemType); + Location loc = op.getLoc(); + Type resultType = op.getResult().getType(); + Value source = op.getSource(); + Value mask = op.getMask(); + + if (auto groupAttr = op.getGroupAttr()) { + int64_t numGroups = groupAttr.getInt(); + Value result; + if (isFloat) + result = builder + .create( + loc, resultType, source, mask, + builder.getI64IntegerAttr(numGroups)) + .getResult(); + else + result = builder + .create( + loc, resultType, source, mask, + builder.getI64IntegerAttr(numGroups)) + .getResult(); + op.getResult().replaceAllUsesWith(result); + op->erase(); + return success(); + } + + Value init = createReduceNeutralInit(builder, loc, elemType, + /*isAdd=*/false, /*isMax=*/false, + sourceType.getLayout()); + Value result; + if (isFloat) + result = builder + .create(loc, resultType, source, init, mask) + .getResult(); + else + result = builder + .create(loc, resultType, source, init, mask) + .getResult(); + op.getResult().replaceAllUsesWith(result); + op->erase(); + return success(); +} + +//===----------------------------------------------------------------------===// +// Category C7 helpers: vmula / vaxpy (fused multiply-add → legacy fma) +//===----------------------------------------------------------------------===// + +/// Lower vmula (acc = acc + lhs*rhs) to legacy fma (lhs*rhs + acc). +/// The mask operand (if present) is discarded — legacy fma has no mask. +/// Legacy fma is floating-point only; integer vmula has no legacy equivalent +/// and is skipped (falls through to VMIToVPTO). +static LogicalResult lowerVmula(VMIVmulaOp op, OpBuilder &builder) { + if (hasMergePmode(op)) + return failure(); + + Type resultType = op.getResult().getType(); + auto vmiType = cast(resultType); + if (!isFloatType(vmiType.getElementType())) + return failure(); + + Location loc = op.getLoc(); + // fma computes lhs*rhs + acc, matching vmula's acc + lhs*rhs. + Value result = builder + .create(loc, resultType, op.getLhs(), op.getRhs(), + op.getAcc()) + .getResult(); + op.getResult().replaceAllUsesWith(result); + op->erase(); + return success(); +} + +/// Lower vaxpy (alpha*x + y) to broadcast(alpha) + legacy fma. +/// alpha is a scalar float, broadcast to a vector before the fma. +static LogicalResult lowerVaxpy(VMIVaxpyOp op, OpBuilder &builder) { + if (hasMergePmode(op)) + return failure(); + + Type resultType = op.getResult().getType(); + auto vmiType = cast(resultType); + if (!isFloatType(vmiType.getElementType())) + return failure(); + + Location loc = op.getLoc(); + Value alphaVec = builder + .create(loc, resultType, op.getAlpha()) + .getResult(); + // fma(alpha, x, y) == alpha*x + y. + Value raw = builder + .create(loc, resultType, alphaVec, op.getX(), + op.getAcc()) + .getResult(); + op.getResult().replaceAllUsesWith(raw); + op->erase(); + return success(); +} + +/// Lower plt(rem:i32) -> create_mask(min(rem, L)) + arith remainder chain. +/// %act = arith.minsi %rem, %cL // min(rem, L) +/// %aidx = arith.index_cast %act // i32 -> index +/// %mask = vmi.create_mask %aidx +/// %next = arith.subi %rem, %act // rem - min(rem, L) = max(rem-L, 0) +static LogicalResult lowerPlt(VMIPltOp op, OpBuilder &builder) { + Location loc = op.getLoc(); + auto maskType = cast(op.getMask().getType()); + int64_t laneCount = maskType.getElementCount(); + + auto i32Type = builder.getIntegerType(32); + Value cL = builder.create( + loc, i32Type, builder.getIntegerAttr(i32Type, laneCount)); + Value act = builder.create(loc, i32Type, op.getScalar(), cL); + Value aidx = builder.create( + loc, builder.getIndexType(), act); + Value mask = builder.create(loc, maskType, aidx).getResult(); + Value next = builder.create(loc, i32Type, op.getScalar(), act); + + op.getMask().replaceAllUsesWith(mask); + op.getScalarOut().replaceAllUsesWith(next); + op->erase(); + return success(); +} + +//===----------------------------------------------------------------------===// +// Category C8 helpers: vgather / vscatter +//===----------------------------------------------------------------------===// + +/// Lower vgather to legacy gather. Legacy gather carries an explicit passthru +/// operand for inactive lanes; pmode="zero" is modelled with a zero passthru. +/// pmode="merge" (preserve OLD_DEST) has no SSA passthru and is skipped. +static LogicalResult lowerVgather(VMIVgatherOp op, OpBuilder &builder) { + if (hasMergePmode(op)) + return failure(); + + Location loc = op.getLoc(); + auto resultType = cast(op.getResult().getType()); + // pmode="zero" (default): inactive lanes are zeroed. Legacy gather models + // inactive lanes with an explicit passthru whose element type must match the + // result, so synthesise a zero constant of the result type — the offsets + // vector cannot be reused because its element type (e.g. i32) generally + // differs from the result element type (e.g. f32). + Value passthru = createZeroConstant(builder, loc, resultType); + Value result = builder + .create(loc, resultType, op.getSource(), + op.getOffsets(), op.getMask(), + passthru) + .getResult(); + op.getResult().replaceAllUsesWith(result); + op->erase(); + return success(); +} + +/// Lower vscatter to legacy scatter. Legacy scatter only writes active lanes +/// (mask-governed), matching vscatter's default/zero pmode; merge is skipped. +static LogicalResult lowerVscatter(VMIVscatterOp op, OpBuilder &builder) { + if (hasMergePmode(op)) + return failure(); + + Location loc = op.getLoc(); + builder.create(loc, op.getValue(), op.getDestination(), + op.getOffsets(), op.getMask()); + op->erase(); + return success(); +} + +//===----------------------------------------------------------------------===// +// Category C9 helpers: vexpdif / vlrelu / vprelu (fused → legacy chains) +//===----------------------------------------------------------------------===// + +/// Lower vexpdif (exp(x - max)) to [extf] + subf + exp. +/// x may be f16 while max and result are always f32 — widen x first when its +/// element type differs from the result type. +static LogicalResult lowerVexpdif(VMIVexpdifOp op, OpBuilder &builder) { + if (hasMergePmode(op)) + return failure(); + + Location loc = op.getLoc(); + Type resultType = op.getResult().getType(); + auto vmiType = cast(resultType); + Type resElem = vmiType.getElementType(); + + Value x = op.getX(); + if (getVMIElementType(x) != resElem) + x = builder.create(loc, resultType, x).getResult(); + + Value diff = + builder.create(loc, resultType, x, op.getMax()).getResult(); + Value raw = builder.create(loc, resultType, diff).getResult(); + op.getResult().replaceAllUsesWith(raw); + op->erase(); + return success(); +} + +/// Lower vlrelu (x>0 ? x : slope*x) to max(x,0) + slope*min(x,0). +/// slope is a scalar float broadcast to a vector. +static LogicalResult lowerVlrelu(VMIVlreluOp op, OpBuilder &builder) { + if (hasMergePmode(op)) + return failure(); + + Location loc = op.getLoc(); + Type resultType = op.getResult().getType(); + auto vmiType = cast(resultType); + Value x = op.getX(); + + Value zeroConst = createZeroConstant(builder, loc, vmiType); + Value pos = + builder.create(loc, resultType, x, zeroConst).getResult(); + Value neg = + builder.create(loc, resultType, x, zeroConst).getResult(); + Value slopeVec = builder + .create(loc, resultType, op.getSlope()) + .getResult(); + Value scaledNeg = + builder.create(loc, resultType, slopeVec, neg).getResult(); + Value raw = + builder.create(loc, resultType, pos, scaledNeg).getResult(); + op.getResult().replaceAllUsesWith(raw); + op->erase(); + return success(); +} + +/// Lower vprelu (max(x,0) + alpha*min(x,0)) to legacy max/min/mul/add. +/// alpha is a per-lane vector (no broadcast needed). +static LogicalResult lowerVprelu(VMIVpreluOp op, OpBuilder &builder) { + if (hasMergePmode(op)) + return failure(); + + Location loc = op.getLoc(); + Type resultType = op.getResult().getType(); + auto vmiType = cast(resultType); + Value x = op.getX(); + + Value zeroConst = createZeroConstant(builder, loc, vmiType); + Value pos = + builder.create(loc, resultType, x, zeroConst).getResult(); + Value neg = + builder.create(loc, resultType, x, zeroConst).getResult(); + Value scaledNeg = + builder.create(loc, resultType, op.getAlpha(), neg).getResult(); + Value raw = + builder.create(loc, resultType, pos, scaledNeg).getResult(); + op.getResult().replaceAllUsesWith(raw); + op->erase(); + return success(); +} + +//===----------------------------------------------------------------------===// +// Pass definition +//===----------------------------------------------------------------------===// + +namespace { + +struct VMILowerUnifiedToLegacyPass + : public mlir::pto::impl::VMILowerUnifiedToLegacyBase< + VMILowerUnifiedToLegacyPass> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(VMILowerUnifiedToLegacyPass) + + void runOnOperation() override; + + void getDependentDialects(mlir::DialectRegistry ®istry) const override { + registry.insert(); + } +}; + +} // namespace + +void VMILowerUnifiedToLegacyPass::runOnOperation() { + ModuleOp module = getOperation(); + SmallVector worklist; + + // Collect all unified VMI ops (walk encounters them in IR order). + module.walk([&](Operation *op) { + // Category A + if (isa(op) || + // Category B — binary + isa(op) || + // Category B — unary + isa(op) || + // Category C1 + isa(op) || + // Category C2 + isa(op) || + // Category C3 + isa(op) || + // Category C4 + isa(op) || + // Category C5 + isa(op) || + // Category C6 — unified reduce (partial coverage) + isa(op) || + // Category C7 — fused multiply-add family → legacy fma + isa(op) || + // Category C8 — indexed gather / scatter + isa(op) || + // Category C9 — fused activation / softmax (legacy chains) + isa(op)) + worklist.push_back(op); + + // Category D — no legacy equivalent (require direct VMIToVPTO lowering): + // plt, vintlv, vdintlv, vselr, vgatherb, vmull + // These are intentionally NOT added to the worklist — they flow through + // to VMIToVPTO which must provide direct 1:N lowering patterns. + if (isa(op)) { + op->emitRemark("VMI unified op has no legacy equivalent — " + "requires direct VMIToVPTO 1:N lowering"); + } + }); + + for (Operation *op : llvm::reverse(worklist)) { + if (!op->getBlock()) + continue; + OpBuilder builder(op); + + // ---- Category A: pure syntactic renames ---- + + if (auto vop = dyn_cast(op)) { + // vci -> iota + builder.setInsertionPoint(op); + StringAttr orderAttr; + if (auto order = vop.getOrder()) + orderAttr = builder.getStringAttr(*order); + Value result = + builder + .create(op->getLoc(), vop.getResult().getType(), + vop.getBase(), orderAttr) + .getResult(); + vop.getResult().replaceAllUsesWith(result); + op->erase(); + continue; + } + + if (auto vop = dyn_cast(op)) { + // vinterpret_cast -> bitcast + builder.setInsertionPoint(op); + Value result = + builder + .create(op->getLoc(), vop.getResult().getType(), + vop.getSource()) + .getResult(); + vop.getResult().replaceAllUsesWith(result); + op->erase(); + continue; + } + + if (auto vop = dyn_cast(op)) { + // vsel -> select + builder.setInsertionPoint(op); + Value result = + builder + .create(op->getLoc(), vop.getResult().getType(), + vop.getMask(), vop.getTrueValue(), + vop.getFalseValue()) + .getResult(); + vop.getResult().replaceAllUsesWith(result); + op->erase(); + continue; + } + + if (auto vop = dyn_cast(op)) { + // vbrc -> broadcast; vbrc{group} -> group_broadcast + builder.setInsertionPoint(op); + Value result; + if (vop.getGroupAttr()) { + result = + builder + .create(op->getLoc(), vop.getResult().getType(), + vop.getValue(), vop.getGroupAttr()) + .getResult(); + } else { + result = + builder + .create(op->getLoc(), vop.getResult().getType(), + vop.getValue()) + .getResult(); + } + vop.getResult().replaceAllUsesWith(result); + op->erase(); + continue; + } + + // ---- Category C4: static mask creation ---- + + if (auto vop = dyn_cast(op)) { + (void)lowerPset(vop, builder); + continue; + } + + if (auto vop = dyn_cast(op)) { + (void)lowerPge(vop, builder); + continue; + } + + if (auto vop = dyn_cast(op)) { + (void)lowerPlt(vop, builder); + continue; + } + + // ---- Category C1: vcmp / vcmps ---- + + if (auto vop = dyn_cast(op)) { + (void)lowerVCmp(vop, builder); + continue; + } + + if (auto vop = dyn_cast(op)) { + (void)lowerVCmps(vop, builder); + continue; + } + + // ---- Category C2: vcvt ---- + + if (auto vop = dyn_cast(op)) { + (void)lowerVCvt(vop, builder); + continue; + } + + // ---- Category C3: vload / vstore ---- + + if (auto vop = dyn_cast(op)) { + (void)lowerVLoad(vop, builder); + continue; + } + + if (auto vop = dyn_cast(op)) { + (void)lowerVStore(vop, builder); + continue; + } + + // ---- Category C5: vector-scalar ops ---- + + if (auto vop = dyn_cast(op)) { + Type elemType = getVMIElementType(vop.getSrc()); + auto createLegacy = [&](Location loc, Type ty, Value lhs, Value rhs) -> Value { + if (isFloatType(elemType)) + return builder.create(loc, ty, lhs, rhs).getResult(); + return builder.create(loc, ty, lhs, rhs).getResult(); + }; + (void)lowerVecScalar(vop, builder, createLegacy); + continue; + } + + if (auto vop = dyn_cast(op)) { + Type elemType = getVMIElementType(vop.getSrc()); + auto createLegacy = [&](Location loc, Type ty, Value lhs, Value rhs) -> Value { + if (isFloatType(elemType)) + return builder.create(loc, ty, lhs, rhs).getResult(); + return builder.create(loc, ty, lhs, rhs).getResult(); + }; + (void)lowerVecScalar(vop, builder, createLegacy); + continue; + } + + if (auto vop = dyn_cast(op)) { + auto createLegacy = [&](Location loc, Type ty, Value lhs, + Value rhs) -> Value { + return builder.create(loc, ty, lhs, rhs).getResult(); + }; + (void)lowerVecScalar(vop, builder, createLegacy); + continue; + } + + if (auto vop = dyn_cast(op)) { + auto createLegacy = [&](Location loc, Type ty, Value lhs, + Value rhs) -> Value { + return builder.create(loc, ty, lhs, rhs).getResult(); + }; + (void)lowerVecScalar(vop, builder, createLegacy); + continue; + } + + if (auto vop = dyn_cast(op)) { + auto createLegacy = [&](Location loc, Type ty, Value lhs, + Value rhs) -> Value { + return builder.create(loc, ty, lhs, rhs).getResult(); + }; + (void)lowerVecScalar(vop, builder, createLegacy); + continue; + } + + if (auto vop = dyn_cast(op)) { + Type elemType = getVMIElementType(vop.getSrc()); + auto createLegacy = [&](Location loc, Type ty, Value lhs, + Value rhs) -> Value { + auto intType = cast(elemType); + if (intType.isUnsigned()) + return builder.create(loc, ty, lhs, rhs).getResult(); + return builder.create(loc, ty, lhs, rhs).getResult(); + }; + (void)lowerVecScalar(vop, builder, createLegacy); + continue; + } + + // ---- Category C6: unified reduce ---- + + if (auto vop = dyn_cast(op)) { + (void)lowerVCadd(vop, builder); + continue; + } + + if (auto vop = dyn_cast(op)) { + (void)lowerVcmax(vop, builder); + continue; + } + + if (auto vop = dyn_cast(op)) { + (void)lowerVcmin(vop, builder); + continue; + } + + // ---- Category C7: fused multiply-add family ---- + + if (auto vop = dyn_cast(op)) { + (void)lowerVmula(vop, builder); + continue; + } + + if (auto vop = dyn_cast(op)) { + (void)lowerVaxpy(vop, builder); + continue; + } + + // ---- Category C8: indexed gather / scatter ---- + + if (auto vop = dyn_cast(op)) { + (void)lowerVgather(vop, builder); + continue; + } + + if (auto vop = dyn_cast(op)) { + (void)lowerVscatter(vop, builder); + continue; + } + + // ---- Category C9: fused activation / softmax ---- + + if (auto vop = dyn_cast(op)) { + (void)lowerVexpdif(vop, builder); + continue; + } + + if (auto vop = dyn_cast(op)) { + (void)lowerVlrelu(vop, builder); + continue; + } + + if (auto vop = dyn_cast(op)) { + (void)lowerVprelu(vop, builder); + continue; + } + + // ---- Category B: binary elementwise, mask/pmode discarded ---- + + if (auto vop = dyn_cast(op)) { + Type elemType = getVMIElementType(vop.getResult()); + auto createLegacy = [&](Location loc, Type ty, Value lhs, Value rhs) -> Value { + if (isFloatType(elemType)) + return builder.create(loc, ty, lhs, rhs).getResult(); + return builder.create(loc, ty, lhs, rhs).getResult(); + }; + (void)lowerBinaryIgnoringMask(vop, createLegacy); + continue; + } + + if (auto vop = dyn_cast(op)) { + Type elemType = getVMIElementType(vop.getResult()); + auto createLegacy = [&](Location loc, Type ty, Value lhs, Value rhs) -> Value { + if (isFloatType(elemType)) + return builder.create(loc, ty, lhs, rhs).getResult(); + return builder.create(loc, ty, lhs, rhs).getResult(); + }; + (void)lowerBinaryIgnoringMask(vop, createLegacy); + continue; + } + + if (auto vop = dyn_cast(op)) { + Type elemType = getVMIElementType(vop.getResult()); + auto createLegacy = [&](Location loc, Type ty, Value lhs, Value rhs) -> Value { + if (isFloatType(elemType)) + return builder.create(loc, ty, lhs, rhs).getResult(); + return builder.create(loc, ty, lhs, rhs).getResult(); + }; + (void)lowerBinaryIgnoringMask(vop, createLegacy); + continue; + } + + if (auto vop = dyn_cast(op)) { + // Float-only; no legacy integer divide. + auto createLegacy = [&](Location loc, Type ty, Value lhs, + Value rhs) -> Value { + return builder.create(loc, ty, lhs, rhs).getResult(); + }; + (void)lowerBinaryIgnoringMask(vop, createLegacy); + continue; + } + + if (auto vop = dyn_cast(op)) { + auto createLegacy = [&](Location loc, Type ty, Value lhs, + Value rhs) -> Value { + return builder.create(loc, ty, lhs, rhs).getResult(); + }; + (void)lowerBinaryIgnoringMask(vop, createLegacy); + continue; + } + + if (auto vop = dyn_cast(op)) { + auto createLegacy = [&](Location loc, Type ty, Value lhs, + Value rhs) -> Value { + return builder.create(loc, ty, lhs, rhs).getResult(); + }; + (void)lowerBinaryIgnoringMask(vop, createLegacy); + continue; + } + + if (auto vop = dyn_cast(op)) { + builder.setInsertionPoint(op); + if (isa(vop.getLhs().getType())) { + // Mask logic path: lower to mask_and. + Value result = builder.create( + vop.getLoc(), vop.getResult().getType(), + vop.getLhs(), vop.getRhs()).getResult(); + vop.getResult().replaceAllUsesWith(result); + op->erase(); + continue; + } + auto createLegacy = [&](Location loc, Type ty, Value lhs, + Value rhs) -> Value { + return builder.create(loc, ty, lhs, rhs).getResult(); + }; + (void)lowerBinaryIgnoringMask(vop, createLegacy); + continue; + } + + if (auto vop = dyn_cast(op)) { + builder.setInsertionPoint(op); + if (isa(vop.getLhs().getType())) { + // Mask logic path: lower to mask_or. + Value result = builder.create( + vop.getLoc(), vop.getResult().getType(), + vop.getLhs(), vop.getRhs()).getResult(); + vop.getResult().replaceAllUsesWith(result); + op->erase(); + continue; + } + auto createLegacy = [&](Location loc, Type ty, Value lhs, + Value rhs) -> Value { + return builder.create(loc, ty, lhs, rhs).getResult(); + }; + (void)lowerBinaryIgnoringMask(vop, createLegacy); + continue; + } + + if (auto vop = dyn_cast(op)) { + builder.setInsertionPoint(op); + if (isa(vop.getLhs().getType())) { + // Mask logic path: lower to mask_xor. + Value result = builder.create( + vop.getLoc(), vop.getResult().getType(), + vop.getLhs(), vop.getRhs()).getResult(); + vop.getResult().replaceAllUsesWith(result); + op->erase(); + continue; + } + auto createLegacy = [&](Location loc, Type ty, Value lhs, + Value rhs) -> Value { + return builder.create(loc, ty, lhs, rhs).getResult(); + }; + (void)lowerBinaryIgnoringMask(vop, createLegacy); + continue; + } + + if (auto vop = dyn_cast(op)) { + auto createLegacy = [&](Location loc, Type ty, Value lhs, + Value rhs) -> Value { + return builder.create(loc, ty, lhs, rhs).getResult(); + }; + (void)lowerBinaryIgnoringMask(vop, createLegacy); + continue; + } + + if (auto vop = dyn_cast(op)) { + Type elemType = getVMIElementType(vop.getLhs()); + auto createLegacy = [&](Location loc, Type ty, Value lhs, + Value rhs) -> Value { + auto intType = cast(elemType); + if (intType.isUnsigned()) + return builder.create(loc, ty, lhs, rhs).getResult(); + return builder.create(loc, ty, lhs, rhs).getResult(); + }; + (void)lowerBinaryIgnoringMask(vop, createLegacy); + continue; + } + + // ---- Category B: masked elementwise — unary ---- + + if (auto vop = dyn_cast(op)) { + auto createLegacy = [&](Location loc, Type ty, Value src) -> Value { + return builder.create(loc, ty, src).getResult(); + }; + (void)lowerMaskedUnary(vop, builder, createLegacy); + continue; + } + + if (auto vop = dyn_cast(op)) { + Type elemType = getVMIElementType(vop.getResult()); + auto createLegacy = [&](Location loc, Type ty, Value src) -> Value { + if (isFloatType(elemType)) + return builder.create(loc, ty, src).getResult(); + return builder.create(loc, ty, src).getResult(); + }; + (void)lowerMaskedUnary(vop, builder, createLegacy); + continue; + } + + if (auto vop = dyn_cast(op)) { + auto createLegacy = [&](Location loc, Type ty, Value src) -> Value { + return builder.create(loc, ty, src).getResult(); + }; + (void)lowerMaskedUnary(vop, builder, createLegacy); + continue; + } + + if (auto vop = dyn_cast(op)) { + auto createLegacy = [&](Location loc, Type ty, Value src) -> Value { + return builder.create(loc, ty, src).getResult(); + }; + (void)lowerMaskedUnary(vop, builder, createLegacy); + continue; + } + + if (auto vop = dyn_cast(op)) { + auto createLegacy = [&](Location loc, Type ty, Value src) -> Value { + return builder.create(loc, ty, src).getResult(); + }; + (void)lowerMaskedUnary(vop, builder, createLegacy); + continue; + } + + if (auto vop = dyn_cast(op)) { + auto createLegacy = [&](Location loc, Type ty, Value src) -> Value { + return builder.create(loc, ty, src).getResult(); + }; + (void)lowerMaskedUnary(vop, builder, createLegacy); + continue; + } + + if (auto vop = dyn_cast(op)) { + builder.setInsertionPoint(op); + if (isa(vop.getSource().getType())) { + // Mask logic path: lower to mask_not. + Value result = builder.create( + vop.getLoc(), vop.getResult().getType(), + vop.getSource()).getResult(); + vop.getResult().replaceAllUsesWith(result); + op->erase(); + continue; + } + auto createLegacy = [&](Location loc, Type ty, Value src) -> Value { + return builder.create(loc, ty, src).getResult(); + }; + (void)lowerMaskedUnary(vop, builder, createLegacy); + continue; + } + } +} + +std::unique_ptr mlir::pto::createVMILowerUnifiedToLegacyPass() { + return std::make_unique(); +} diff --git a/lib/PTO/Transforms/VMIMaskGranularityAssignment.cpp b/lib/PTO/Transforms/VMIMaskGranularityAssignment.cpp new file mode 100644 index 0000000000..3cb4b3d6c0 --- /dev/null +++ b/lib/PTO/Transforms/VMIMaskGranularityAssignment.cpp @@ -0,0 +1,851 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +//===- VMIMaskGranularityAssignment.cpp - Assign VMI mask granularity -----===// +//===----------------------------------------------------------------------===// +// +// This pass assigns concrete b8/b16/b32 granularity to VMI mask values before +// layout assignment. It deliberately does not choose layouts: mask layout is +// assigned later by vmi-layout-assignment. When a mask value has conflicting +// granularity uses, this pass keeps the value's primary granularity and either +// rematerializes cheap mask producers at the use site or inserts +// pto.vmi.ensure_mask_granularity. + +#include "PTO/IR/PTO.h" +#include "PTO/IR/PTOTypeUtils.h" +#include "PTO/IR/VMIUtils.h" +#include "PTO/Transforms/Passes.h" + +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/SymbolTable.h" +#include "mlir/IR/Value.h" +#include "mlir/Pass/Pass.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/SmallVector.h" + +namespace mlir { +namespace pto { +#define GEN_PASS_DEF_VMIMASKGRANULARITYASSIGNMENT +#include "PTO/Transforms/Passes.h.inc" +} // namespace pto +} // namespace mlir + +using namespace mlir; +using namespace mlir::pto; + +namespace { + +struct MaskNode { + Value value; + VMIMaskType type; + unsigned parent = 0; + std::string granularity; +}; + +struct MaskUseRequest { + OpOperand *operand; + std::string granularity; +}; + +static unsigned getElementBitWidth(Type type) { + if (isa(type)) + return 64; + return pto::getPTOStorageElemBitWidth(type); +} + +static StringRef getMaskGranularityForElement(Type elementType) { + switch (getElementBitWidth(elementType)) { + case 8: + return "b8"; + case 16: + return "b16"; + case 32: + return "b32"; + default: + return ""; + } +} + +static bool containsVMIType(Type type) { + if (isa(type)) + return true; + if (auto functionType = dyn_cast(type)) { + return llvm::any_of(functionType.getInputs(), containsVMIType) || + llvm::any_of(functionType.getResults(), containsVMIType); + } + if (auto shapedType = dyn_cast(type)) + return containsVMIType(shapedType.getElementType()); + return false; +} + +struct MaskGranularitySolver { + explicit MaskGranularitySolver(ModuleOp module) + : module(module), ctx(module.getContext()) {} + + unsigned addMaskValue(Value value) { + auto type = dyn_cast(value.getType()); + if (!type) + return ~0u; + auto [it, inserted] = maskIds.try_emplace(value, maskNodes.size()); + if (inserted) { + std::string granularity; + if (VMIMaskType::isConcreteGranularity(type.getGranularity())) + granularity = type.getGranularity().str(); + maskNodes.push_back(MaskNode{value, type, it->second, granularity}); + } + return it->second; + } + + unsigned findMask(unsigned id) { + if (maskNodes[id].parent == id) + return id; + maskNodes[id].parent = findMask(maskNodes[id].parent); + return maskNodes[id].parent; + } + + LogicalResult uniteMask(Value lhs, Value rhs, Operation *op) { + unsigned lhsId = addMaskValue(lhs); + unsigned rhsId = addMaskValue(rhs); + if (lhsId == ~0u || rhsId == ~0u) + return success(); + unsigned lhsRoot = findMask(lhsId); + unsigned rhsRoot = findMask(rhsId); + if (lhsRoot == rhsRoot) + return success(); + + MaskNode &lhsNode = maskNodes[lhsRoot]; + MaskNode &rhsNode = maskNodes[rhsRoot]; + if (!lhsNode.granularity.empty() && !rhsNode.granularity.empty() && + lhsNode.granularity != rhsNode.granularity) + return op->emitError() << kVMIDiagLayoutContractPrefix + << "conflicting mask granularities " + << lhsNode.granularity << " and " + << rhsNode.granularity; + + rhsNode.parent = lhsRoot; + if (lhsNode.granularity.empty()) + lhsNode.granularity = rhsNode.granularity; + return success(); + } + + LogicalResult requestMask(Value mask, StringRef granularity, Operation *op) { + unsigned id = addMaskValue(mask); + if (id == ~0u) + return success(); + if (granularity.empty()) + return op->emitError() << kVMIDiagLayoutContractPrefix + << "cannot infer concrete mask granularity"; + MaskNode &node = maskNodes[findMask(id)]; + if (!node.granularity.empty() && node.granularity != granularity) + return op->emitError() + << kVMIDiagLayoutContractPrefix + << "conflicting mask granularities " << node.granularity << " and " + << granularity; + node.granularity = granularity.str(); + return success(); + } + + LogicalResult requestMaskUse(OpOperand &operand, StringRef granularity, + Operation *op) { + if (!isa(operand.get().getType())) + return success(); + if (granularity.empty()) + return op->emitError() << kVMIDiagLayoutContractPrefix + << "cannot infer concrete mask use granularity"; + maskUseRequests.push_back(MaskUseRequest{&operand, granularity.str()}); + return success(); + } + + LogicalResult collect() { + module.walk([&](Operation *op) { + for (Value result : op->getResults()) + addMaskValue(result); + for (Region ®ion : op->getRegions()) + for (Block &block : region) + for (BlockArgument arg : block.getArguments()) + addMaskValue(arg); + }); + return success(); + } + + LogicalResult addConstraints() { + WalkResult result = module.walk([&](Operation *op) -> WalkResult { + if (auto maskAnd = dyn_cast(op)) { + if (failed(uniteMask(maskAnd.getLhs(), maskAnd.getRhs(), op)) || + failed(uniteMask(maskAnd.getLhs(), maskAnd.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto maskOr = dyn_cast(op)) { + if (failed(uniteMask(maskOr.getLhs(), maskOr.getRhs(), op)) || + failed(uniteMask(maskOr.getLhs(), maskOr.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto maskXor = dyn_cast(op)) { + if (failed(uniteMask(maskXor.getLhs(), maskXor.getRhs(), op)) || + failed(uniteMask(maskXor.getLhs(), maskXor.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto maskNot = dyn_cast(op)) { + if (failed(uniteMask(maskNot.getSource(), maskNot.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto ensure = dyn_cast(op)) { + if (failed(uniteMask(ensure.getSource(), ensure.getResult(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto cmpf = dyn_cast(op)) { + auto lhsType = cast(cmpf.getLhs().getType()); + if (failed(requestMask( + cmpf.getResult(), + getMaskGranularityForElement(lhsType.getElementType()), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto cmpi = dyn_cast(op)) { + auto lhsType = cast(cmpi.getLhs().getType()); + if (failed(requestMask( + cmpi.getResult(), + getMaskGranularityForElement(lhsType.getElementType()), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto select = dyn_cast(op)) { + auto resultType = cast(select.getResult().getType()); + if (failed(requestMaskUse( + select.getMaskMutable(), + getMaskGranularityForElement(resultType.getElementType()), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto vmull = dyn_cast(op)) { + if (failed(requestMaskUse(vmull.getMaskMutable(), "b32", op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto activePrefix = dyn_cast(op)) { + auto resultType = cast(activePrefix.getResult().getType()); + if (failed(requestMaskUse( + activePrefix.getMaskMutable(), + getMaskGranularityForElement(resultType.getElementType()), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto compress = dyn_cast(op)) { + auto resultType = cast(compress.getResult().getType()); + if (failed(requestMaskUse( + compress.getMaskMutable(), + getMaskGranularityForElement(resultType.getElementType()), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto vintlv = dyn_cast(op)) { + if (failed(requestMaskUseForSource(vintlv.getMaskMutable(), + vintlv.getLhs(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto vdintlv = dyn_cast(op)) { + if (failed(requestMaskUseForSource(vdintlv.getMaskMutable(), + vdintlv.getLhs(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto reduce = dyn_cast(op)) { + if (failed(requestMaskUseForSource(reduce.getMaskMutable(), + reduce.getSource(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto reduce = dyn_cast(op)) { + if (failed(requestMaskUseForSource(reduce.getMaskMutable(), + reduce.getSource(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto reduce = dyn_cast(op)) { + if (failed(requestMaskUseForSource(reduce.getMaskMutable(), + reduce.getSource(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto reduce = dyn_cast(op)) { + if (failed(requestMaskUseForSource(reduce.getMaskMutable(), + reduce.getSource(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto reduce = dyn_cast(op)) { + if (failed(requestMaskUseForSource(reduce.getMaskMutable(), + reduce.getSource(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto reduce = dyn_cast(op)) { + if (failed(requestMaskUseForSource(reduce.getMaskMutable(), + reduce.getSource(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto reduce = dyn_cast(op)) { + if (failed(requestMaskUseForSource(reduce.getMaskMutable(), + reduce.getSource(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto reduce = dyn_cast(op)) { + if (failed(requestMaskUseForSource(reduce.getMaskMutable(), + reduce.getSource(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto reduce = dyn_cast(op)) { + if (failed(requestMaskUseForSource(reduce.getMaskMutable(), + reduce.getSource(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto reduce = dyn_cast(op)) { + if (failed(requestMaskUseForSource(reduce.getMaskMutable(), + reduce.getSource(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto reduce = dyn_cast(op)) { + if (failed(requestMaskUseForSource(reduce.getMaskMutable(), + reduce.getSource(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto reduce = dyn_cast(op)) { + if (failed(requestMaskUseForSource(reduce.getMaskMutable(), + reduce.getSource(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto hist = dyn_cast(op)) { + if (failed(requestMaskUse(hist.getMaskMutable(), "b8", op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto hist = dyn_cast(op)) { + if (failed(requestMaskUse(hist.getMaskMutable(), "b8", op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto load = dyn_cast(op)) { + auto resultType = cast(load.getResult().getType()); + if (failed(requestMaskUse( + load.getMaskMutable(), + getMaskGranularityForElement(resultType.getElementType()), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto load = dyn_cast(op)) { + auto resultType = cast(load.getResult().getType()); + if (failed(requestMaskUse( + load.getMaskMutable(), + getMaskGranularityForElement(resultType.getElementType()), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto gather = dyn_cast(op)) { + auto resultType = cast(gather.getResult().getType()); + if (failed(requestMaskUse( + gather.getMaskMutable(), + getMaskGranularityForElement(resultType.getElementType()), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto load = dyn_cast(op)) { + auto resultType = cast(load.getResult().getType()); + if (failed(requestMaskUse( + load.getMaskMutable(), + getMaskGranularityForElement(resultType.getElementType()), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto store = dyn_cast(op)) { + if (failed(requestMaskUseForSource(store.getMaskMutable(), + store.getValue(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto store = dyn_cast(op)) { + if (failed(requestMaskUseForSource(store.getMaskMutable(), + store.getValue(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto scatter = dyn_cast(op)) { + if (failed(requestMaskUseForSource(scatter.getMaskMutable(), + scatter.getValue(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto store = dyn_cast(op)) { + if (failed(requestMaskUseForSource(store.getMaskMutable(), + store.getValue(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto ifOp = dyn_cast(op)) { + if (failed(addIfConstraints(ifOp))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto executeOp = dyn_cast(op)) { + if (failed(addExecuteRegionConstraints(executeOp))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto indexSwitchOp = dyn_cast(op)) { + if (failed(addIndexSwitchConstraints(indexSwitchOp))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto whileOp = dyn_cast(op)) { + if (failed(addWhileConstraints(whileOp))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto forOp = dyn_cast(op)) { + if (failed(addForConstraints(forOp))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto branchOp = dyn_cast(op)) { + if (failed(addBranchConstraints(branchOp.getDest(), + branchOp.getDestOperands(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto condBranchOp = dyn_cast(op)) { + if (failed(addBranchConstraints(condBranchOp.getTrueDest(), + condBranchOp.getTrueDestOperands(), + op)) || + failed(addBranchConstraints(condBranchOp.getFalseDest(), + condBranchOp.getFalseOperands(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto switchOp = dyn_cast(op)) { + if (failed(addBranchConstraints(switchOp.getDefaultDestination(), + switchOp.getDefaultOperands(), op))) + return WalkResult::interrupt(); + for (auto [dest, operands] : llvm::zip(switchOp.getCaseDestinations(), + switchOp.getCaseOperands())) { + if (failed(addBranchConstraints(dest, operands, op))) + return WalkResult::interrupt(); + } + return WalkResult::advance(); + } + if (auto returnOp = dyn_cast(op)) { + if (failed(addReturnConstraints(returnOp))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (auto callOp = dyn_cast(op)) { + if (failed(addCallConstraints(callOp))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } + if (op->getName().getStringRef() == "func.call_indirect" && + hasVMIValueTypes(op)) { + op->emitError() << kVMIDiagLayoutContractPrefix + << "VMI typed call requires a direct internal callee " + "with a body"; + return WalkResult::interrupt(); + } + if (auto funcOp = dyn_cast(op)) { + if (funcOp.empty() && hasVMIFunctionType(funcOp)) { + funcOp.emitError() + << kVMIDiagLayoutContractPrefix + << "VMI typed function declaration requires an explicit " + "external ABI materialization plan"; + return WalkResult::interrupt(); + } + return WalkResult::advance(); + } + return WalkResult::advance(); + }); + return failure(result.wasInterrupted()); + } + + LogicalResult requestMaskUseForSource(OpOperand &mask, Value source, + Operation *op) { + auto sourceType = dyn_cast(source.getType()); + if (!sourceType) + return success(); + return requestMaskUse(mask, + getMaskGranularityForElement( + sourceType.getElementType()), + op); + } + + LogicalResult uniteEquivalentValues(Value lhs, Value rhs, Operation *op) { + return uniteMask(lhs, rhs, op); + } + + LogicalResult addIfConstraints(scf::IfOp ifOp) { + for (OpResult result : ifOp->getResults()) { + unsigned resultNo = result.getResultNumber(); + for (Region *region : {&ifOp.getThenRegion(), &ifOp.getElseRegion()}) { + if (region->empty()) + continue; + auto yieldOp = dyn_cast(region->front().getTerminator()); + if (!yieldOp || resultNo >= yieldOp.getNumOperands()) + continue; + if (failed(uniteEquivalentValues(result, yieldOp.getOperand(resultNo), + ifOp))) + return failure(); + } + } + return success(); + } + + LogicalResult addYieldConstraints(ResultRange results, scf::YieldOp yieldOp, + Operation *op) { + for (auto [index, result] : llvm::enumerate(results)) { + if (index >= yieldOp.getNumOperands()) + break; + if (failed(uniteEquivalentValues(result, yieldOp.getOperand(index), op))) + return failure(); + } + return success(); + } + + LogicalResult addExecuteRegionConstraints(scf::ExecuteRegionOp executeOp) { + WalkResult result = executeOp.getRegion().walk([&](scf::YieldOp yieldOp) { + if (yieldOp->getParentOp() != executeOp.getOperation()) + return WalkResult::advance(); + if (failed( + addYieldConstraints(executeOp->getResults(), yieldOp, executeOp))) + return WalkResult::interrupt(); + return WalkResult::advance(); + }); + return failure(result.wasInterrupted()); + } + + LogicalResult addIndexSwitchConstraints(scf::IndexSwitchOp indexSwitchOp) { + auto addBlockTerminator = [&](Block &block) -> LogicalResult { + auto yieldOp = dyn_cast(block.getTerminator()); + if (!yieldOp) + return success(); + return addYieldConstraints(indexSwitchOp->getResults(), yieldOp, + indexSwitchOp); + }; + + if (failed(addBlockTerminator(indexSwitchOp.getDefaultBlock()))) + return failure(); + for (unsigned idx = 0, e = indexSwitchOp.getNumCases(); idx < e; ++idx) + if (failed(addBlockTerminator(indexSwitchOp.getCaseBlock(idx)))) + return failure(); + return success(); + } + + LogicalResult addWhileConstraints(scf::WhileOp whileOp) { + auto inits = whileOp.getInits(); + auto beforeArgs = whileOp.getBeforeArguments(); + Block &afterBlock = whileOp.getAfter().front(); + auto conditionOp = + dyn_cast(whileOp.getBefore().front().getTerminator()); + auto yieldOp = dyn_cast(afterBlock.getTerminator()); + + for (auto [index, init] : llvm::enumerate(inits)) { + Value anchor = init; + if (index < beforeArgs.size() && + failed(uniteEquivalentValues(anchor, beforeArgs[index], whileOp))) + return failure(); + if (conditionOp && index < conditionOp.getArgs().size() && + failed(uniteEquivalentValues(anchor, conditionOp.getArgs()[index], + whileOp))) + return failure(); + if (index < afterBlock.getNumArguments() && + failed(uniteEquivalentValues(anchor, afterBlock.getArgument(index), + whileOp))) + return failure(); + if (yieldOp && index < yieldOp.getNumOperands() && + failed(uniteEquivalentValues(anchor, yieldOp.getOperand(index), + whileOp))) + return failure(); + if (index < whileOp.getNumResults() && + failed( + uniteEquivalentValues(anchor, whileOp.getResult(index), whileOp))) + return failure(); + } + return success(); + } + + LogicalResult addForConstraints(scf::ForOp forOp) { + auto initArgs = forOp.getInitArgs(); + auto regionIterArgs = forOp.getRegionIterArgs(); + auto results = forOp.getResults(); + scf::YieldOp yieldOp = nullptr; + if (Block *body = forOp.getBody()) + yieldOp = dyn_cast(body->getTerminator()); + + for (auto [index, initArg] : llvm::enumerate(initArgs)) { + Value anchor = initArg; + if (index < regionIterArgs.size() && + failed(uniteEquivalentValues(anchor, regionIterArgs[index], forOp))) + return failure(); + if (index < results.size() && + failed(uniteEquivalentValues(anchor, results[index], forOp))) + return failure(); + if (yieldOp && index < yieldOp.getNumOperands() && + failed( + uniteEquivalentValues(anchor, yieldOp.getOperand(index), forOp))) + return failure(); + } + return success(); + } + + LogicalResult addBranchConstraints(Block *dest, OperandRange operands, + Operation *op) { + if (!dest) + return success(); + for (auto [index, operand] : llvm::enumerate(operands)) { + if (index >= dest->getNumArguments()) + break; + if (failed(uniteEquivalentValues(operand, dest->getArgument(index), op))) + return failure(); + } + return success(); + } + + LogicalResult addReturnConstraints(func::ReturnOp returnOp) { + auto func = returnOp->getParentOfType(); + if (!func) + return success(); + + auto it = firstReturnOperandsByFunc.find(func); + if (it == firstReturnOperandsByFunc.end()) { + SmallVector operands(returnOp.getOperands()); + firstReturnOperandsByFunc.try_emplace(func, std::move(operands)); + return success(); + } + + ArrayRef firstOperands = it->second; + for (auto [index, operand] : llvm::enumerate(returnOp.getOperands())) { + if (index >= firstOperands.size()) + break; + if (failed( + uniteEquivalentValues(firstOperands[index], operand, returnOp))) + return failure(); + } + return success(); + } + + bool hasVMIValueTypes(Operation *op) { + return llvm::any_of(op->getOperandTypes(), containsVMIType) || + llvm::any_of(op->getResultTypes(), containsVMIType); + } + + bool hasVMIFunctionType(func::FuncOp func) { + FunctionType type = func.getFunctionType(); + return llvm::any_of(type.getInputs(), containsVMIType) || + llvm::any_of(type.getResults(), containsVMIType); + } + + LogicalResult addCallConstraints(func::CallOp callOp) { + if (!hasVMIValueTypes(callOp)) + return success(); + + auto callee = SymbolTable::lookupNearestSymbolFrom( + callOp, callOp.getCalleeAttr()); + if (!callee || callee.empty()) + return callOp.emitError() + << kVMIDiagLayoutContractPrefix + << "VMI typed call requires a direct internal callee with a body"; + + for (auto [operand, argument] : + llvm::zip(callOp.getOperands(), callee.getArguments())) { + if (failed(uniteEquivalentValues(operand, argument, callOp))) + return failure(); + } + + SmallVector returns; + callee.walk([&](func::ReturnOp returnOp) { returns.push_back(returnOp); }); + for (func::ReturnOp returnOp : returns) { + for (auto [index, result] : llvm::enumerate(callOp.getResults())) { + if (index >= returnOp.getNumOperands()) + break; + if (failed(uniteEquivalentValues(result, returnOp.getOperand(index), + callOp))) + return failure(); + } + } + return success(); + } + + void rewriteMaskTypes() { + for (MaskNode &node : maskNodes) { + MaskNode &root = maskNodes[findMask(maskIds.lookup(node.value))]; + StringRef granularity = + root.granularity.empty() ? StringRef("b32") : StringRef(root.granularity); + node.value.setType(VMIMaskType::get(ctx, node.type.getElementCount(), + granularity, + node.type.getLayoutAttr())); + } + } + + SmallVector getCallResultTypes(func::FuncOp func) { + SmallVector resultTypes; + bool found = false; + module.walk([&](func::CallOp call) { + if (call.getCallee() != func.getSymName()) + return; + if (!found) { + resultTypes.assign(call.getResultTypes().begin(), + call.getResultTypes().end()); + found = true; + return; + } + if (resultTypes.size() != call.getNumResults()) + return; + for (auto [index, type] : llvm::enumerate(call.getResultTypes())) + if (index < resultTypes.size() && resultTypes[index] != type) + resultTypes[index] = {}; + }); + return found ? resultTypes : SmallVector{}; + } + + void rewriteFunctionType() { + module.walk([&](func::FuncOp func) { + if (func.empty()) + return; + + SmallVector inputs; + inputs.reserve(func.getNumArguments()); + for (BlockArgument arg : func.getArguments()) + inputs.push_back(arg.getType()); + + SmallVector results; + SmallVector callResultTypes = getCallResultTypes(func); + auto it = firstReturnOperandsByFunc.find(func); + if (!callResultTypes.empty()) { + for (Type type : callResultTypes) + results.push_back(type ? type : Type{}); + } else if (it != firstReturnOperandsByFunc.end()) { + for (Value operand : it->second) + results.push_back(operand.getType()); + } else { + FunctionType functionType = func.getFunctionType(); + for (Type type : functionType.getResults()) { + if (auto maskType = dyn_cast(type)) { + StringRef granularity = + VMIMaskType::isConcreteGranularity(maskType.getGranularity()) + ? maskType.getGranularity() + : StringRef("b32"); + results.push_back(VMIMaskType::get( + ctx, maskType.getElementCount(), granularity, + maskType.getLayoutAttr())); + } else { + results.push_back(type); + } + } + } + + for (auto [index, type] : llvm::enumerate(results)) + if (!type) + results[index] = func.getFunctionType().getResult(index); + + func.setFunctionType(FunctionType::get(ctx, inputs, results)); + }); + } + + LogicalResult insertMaskUseMaterializations() { + OpBuilder builder(ctx); + for (MaskUseRequest request : maskUseRequests) { + Value value = request.operand->get(); + auto sourceType = dyn_cast(value.getType()); + if (!sourceType) + continue; + if (sourceType.getGranularity() == request.granularity) + continue; + + builder.setInsertionPoint(request.operand->getOwner()); + auto resultType = VMIMaskType::get(ctx, sourceType.getElementCount(), + request.granularity, + sourceType.getLayoutAttr()); + Value current = rematerializeMaskProducer( + value, resultType, request.operand->getOwner()->getLoc(), builder); + if (!current) + current = builder.create( + request.operand->getOwner()->getLoc(), resultType, value); + request.operand->set(current); + } + return success(); + } + + Value rematerializeMaskProducer(Value value, VMIMaskType resultType, + Location loc, OpBuilder &builder) { + if (auto createMask = value.getDefiningOp()) + return builder + .create(loc, resultType, createMask.getActiveLanes()) + .getResult(); + + if (auto createGroupMask = value.getDefiningOp()) + return builder + .create( + loc, resultType, createGroupMask.getActiveElemsPerGroup(), + createGroupMask.getNumGroupsAttr(), + createGroupMask.getGroupSizeAttr()) + .getResult(); + + if (auto constantMask = value.getDefiningOp()) + return builder + .create(loc, resultType, + constantMask.getValueAttr()) + .getResult(); + + return {}; + } + + LogicalResult run() { + if (failed(collect())) + return failure(); + if (failed(addConstraints())) + return failure(); + rewriteMaskTypes(); + rewriteFunctionType(); + return insertMaskUseMaterializations(); + } + + ModuleOp module; + MLIRContext *ctx; + DenseMap maskIds; + DenseMap> firstReturnOperandsByFunc; + SmallVector maskNodes; + SmallVector maskUseRequests; +}; + +struct VMIMaskGranularityAssignmentPass + : public mlir::pto::impl::VMIMaskGranularityAssignmentBase< + VMIMaskGranularityAssignmentPass> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(VMIMaskGranularityAssignmentPass) + + void runOnOperation() override { + if (failed(MaskGranularitySolver(getOperation()).run())) + signalPassFailure(); + } +}; + +} // namespace + +std::unique_ptr mlir::pto::createVMIMaskGranularityAssignmentPass() { + return std::make_unique(); +} diff --git a/lib/PTO/Transforms/VMIPreAssignmentCombine.cpp b/lib/PTO/Transforms/VMIPreAssignmentCombine.cpp new file mode 100644 index 0000000000..9a19a872ac --- /dev/null +++ b/lib/PTO/Transforms/VMIPreAssignmentCombine.cpp @@ -0,0 +1,79 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +//===- VMIPreAssignmentCombine.cpp - Pre-assignment VMI combines ---------===// +//===----------------------------------------------------------------------===// + +#include "PTO/IR/PTO.h" +#include "PTO/Transforms/Passes.h" + +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/Operation.h" +#include "mlir/Pass/Pass.h" +#include "llvm/ADT/SmallVector.h" + +namespace mlir { +namespace pto { +#define GEN_PASS_DEF_VMIPREASSIGNMENTCOMBINE +#include "PTO/Transforms/Passes.h.inc" +} // namespace pto +} // namespace mlir + +using namespace mlir; +using namespace mlir::pto; + +namespace { + +static LogicalResult fuseGroupSlotBroadcastLoads(ModuleOp module) { + SmallVector broadcasts; + module.walk([&](VMIGroupBroadcastOp broadcast) { + auto load = broadcast.getSource().getDefiningOp(); + if (!load || !load.getResult().hasOneUse()) + return; + if (load.getNumGroupsAttr().getInt() != + broadcast.getNumGroupsAttr().getInt()) + return; + + if (!isa(broadcast.getResult().getType())) + return; + broadcasts.push_back(broadcast); + }); + + OpBuilder builder(module.getContext()); + for (VMIGroupBroadcastOp broadcast : broadcasts) { + auto load = broadcast.getSource().getDefiningOp(); + if (!load) + continue; + + builder.setInsertionPoint(broadcast); + auto fused = builder.create( + broadcast.getLoc(), broadcast.getResult().getType(), load.getSource(), + load.getOffset(), load.getSourceGroupStride(), + broadcast.getNumGroupsAttr()); + broadcast.getResult().replaceAllUsesWith(fused.getResult()); + broadcast.erase(); + if (load->use_empty()) + load.erase(); + } + return success(); +} + +struct VMIPreAssignmentCombinePass + : pto::impl::VMIPreAssignmentCombineBase { + void runOnOperation() override { + if (failed(fuseGroupSlotBroadcastLoads(getOperation()))) + signalPassFailure(); + } +}; + +} // namespace + +std::unique_ptr mlir::pto::createVMIPreAssignmentCombinePass() { + return std::make_unique(); +} diff --git a/lib/PTO/Transforms/VMIToVPTO.cpp b/lib/PTO/Transforms/VMIToVPTO.cpp new file mode 100644 index 0000000000..e0b1a006e6 --- /dev/null +++ b/lib/PTO/Transforms/VMIToVPTO.cpp @@ -0,0 +1,12761 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +//===- VMIToVPTO.cpp - Convert VMI to physical VPTO IR -------------------===// +//===----------------------------------------------------------------------===// + +// https://discourse.llvm.org/t/matchandrewrite-hiding-virtual-functions/84933/8 +#pragma GCC diagnostic ignored "-Woverloaded-virtual" + +#include "PTO/IR/PTO.h" +#include "PTO/IR/PTOTypeUtils.h" +#include "PTO/IR/VMIUtils.h" +#include "PTO/Transforms/Passes.h" +#include "PTO/Transforms/VMILayoutSupport.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/Func/Transforms/FuncConversions.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/SCF/Transforms/Patterns.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/Operation.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/DialectConversion.h" +#include "llvm/ADT/APFloat.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/Support/raw_ostream.h" +#include +#include +#include + +namespace mlir { +namespace pto { +#define GEN_PASS_DEF_VMITOVPTO +#include "PTO/Transforms/Passes.h.inc" +} // namespace pto +} // namespace mlir + +using namespace mlir; +using namespace mlir::pto; + +namespace { + +std::optional getX2MemoryDistToken(Type elementType, + StringRef prefix); +std::optional getDenseLaneStrideLoadDistToken(VMIVRegType type); +std::optional getDenseLaneStrideStoreDistToken(VMIVRegType type); +std::optional getPointStoreDistToken(Type elementType); + +bool isVMIType(Type type) { return isa(type); } + +bool containsVMIType(Type type) { + if (isVMIType(type)) + return true; + + if (auto functionType = dyn_cast(type)) + return llvm::any_of(functionType.getInputs(), + [](Type input) { return containsVMIType(input); }) || + llvm::any_of(functionType.getResults(), + [](Type result) { return containsVMIType(result); }); + + if (auto shapedType = dyn_cast(type)) + return containsVMIType(shapedType.getElementType()); + + return false; +} + +bool hasVMIType(TypeRange types) { + return llvm::any_of(types, [](Type type) { return containsVMIType(type); }); +} + +struct VMISupportResult { + bool supported = true; + std::string reason; + + static VMISupportResult success() { return {}; } + + static VMISupportResult failure(const Twine &reason) { + VMISupportResult result; + result.supported = false; + result.reason = reason.str(); + return result; + } + + bool isSupported() const { return supported; } + + LogicalResult toLogicalResult(std::string *outReason = nullptr) const { + if (supported) + return mlir::success(); + if (outReason) + *outReason = reason; + return mlir::failure(); + } +}; + +bool hasVMIType(FunctionType type) { + return hasVMIType(type.getInputs()) || hasVMIType(type.getResults()); +} + +bool hasVMIType(Attribute attr) { + if (!attr) + return false; + + if (auto typeAttr = dyn_cast(attr)) + if (containsVMIType(typeAttr.getValue())) + return true; + + if (auto typedAttr = dyn_cast(attr)) + if (containsVMIType(typedAttr.getType())) + return true; + + if (auto arrayAttr = dyn_cast(attr)) + return llvm::any_of(arrayAttr, + [](Attribute element) { return hasVMIType(element); }); + + if (auto dictAttr = dyn_cast(attr)) + return llvm::any_of(dictAttr, [](NamedAttribute namedAttr) { + return hasVMIType(namedAttr.getValue()); + }); + + return false; +} + +bool hasVMIType(Operation *op) { + if (auto func = dyn_cast(op)) + if (hasVMIType(func.getFunctionType())) + return true; + if (hasVMIType(op->getOperandTypes()) || hasVMIType(op->getResultTypes())) + return true; + for (Region ®ion : op->getRegions()) + for (Block &block : region) + if (hasVMIType(block.getArgumentTypes())) + return true; + for (NamedAttribute attr : op->getAttrs()) + if (hasVMIType(attr.getValue())) + return true; + return false; +} + +bool isVMIOp(Operation *op) { + return op->getName().getStringRef().starts_with("pto.vmi."); +} + +StringRef getTruncFRoundModeForResult(Type resultElementType) { + return pto::isPTOHiFloat8Type(resultElementType) ? "A" : "R"; +} + +StringRef getTruncFRoundMode(VMITruncFOp op, Type resultElementType) { + if (auto roundingAttr = op->getAttrOfType("rounding")) + return roundingAttr.getValue(); + return getTruncFRoundModeForResult(resultElementType); +} + +bool isLayoutAssignedVMIType(Type type) { + if (auto vregType = dyn_cast(type)) + return static_cast(vregType.getLayoutAttr()); + if (auto maskType = dyn_cast(type)) + return maskType.getLayoutAttr() && + VMIMaskType::isConcreteGranularity(maskType.getGranularity()); + return true; +} + +LogicalResult verifyLayoutAssignedVMITypeTree(Operation *op, Type type) { + if (!isLayoutAssignedVMIType(type)) + return op->emitError() << kVMIDiagPassInvariantPrefix + << "vmi-to-vpto requires layout-assigned VMI types"; + + if (auto functionType = dyn_cast(type)) { + for (Type input : functionType.getInputs()) + if (failed(verifyLayoutAssignedVMITypeTree(op, input))) + return failure(); + for (Type result : functionType.getResults()) + if (failed(verifyLayoutAssignedVMITypeTree(op, result))) + return failure(); + } + + if (auto shapedType = dyn_cast(type)) + return verifyLayoutAssignedVMITypeTree(op, shapedType.getElementType()); + + return success(); +} + +LogicalResult verifyVMIToVPTOInputAttribute(Operation *op, Attribute attr) { + if (!attr) + return success(); + + if (auto typeAttr = dyn_cast(attr)) + if (failed(verifyLayoutAssignedVMITypeTree(op, typeAttr.getValue()))) + return failure(); + + if (auto typedAttr = dyn_cast(attr)) + if (failed(verifyLayoutAssignedVMITypeTree(op, typedAttr.getType()))) + return failure(); + + if (auto arrayAttr = dyn_cast(attr)) { + for (Attribute element : arrayAttr) + if (failed(verifyVMIToVPTOInputAttribute(op, element))) + return failure(); + } + + if (auto dictAttr = dyn_cast(attr)) { + for (NamedAttribute namedAttr : dictAttr) + if (failed(verifyVMIToVPTOInputAttribute(op, namedAttr.getValue()))) + return failure(); + } + + return success(); +} + +LogicalResult verifyVMIToVPTOInputTypes(Operation *op) { + for (Type type : op->getOperandTypes()) + if (failed(verifyLayoutAssignedVMITypeTree(op, type))) + return failure(); + for (Type type : op->getResultTypes()) + if (failed(verifyLayoutAssignedVMITypeTree(op, type))) + return failure(); + if (auto func = dyn_cast(op)) { + FunctionType functionType = func.getFunctionType(); + for (Type type : functionType.getInputs()) + if (failed(verifyLayoutAssignedVMITypeTree(op, type))) + return failure(); + for (Type type : functionType.getResults()) + if (failed(verifyLayoutAssignedVMITypeTree(op, type))) + return failure(); + } + for (Region ®ion : op->getRegions()) + for (Block &block : region) + for (Type type : block.getArgumentTypes()) + if (failed(verifyLayoutAssignedVMITypeTree(op, type))) + return failure(); + for (NamedAttribute attr : op->getAttrs()) + if (failed(verifyVMIToVPTOInputAttribute(op, attr.getValue()))) + return failure(); + return success(); +} + +LogicalResult verifyVMIToVPTOInputIR(ModuleOp module) { + WalkResult result = module.walk([&](Operation *op) { + if (failed(verifyVMIToVPTOInputTypes(op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + }); + return failure(result.wasInterrupted()); +} + +static Value materializeVPTOToVMI(OpBuilder &builder, Type resultType, + ValueRange inputs, Location loc) { + if (!isVMIType(resultType)) + return {}; + return builder.create(loc, resultType, inputs).getResult(); +} + +static SmallVector materializeVMIToVPTO(OpBuilder &builder, + TypeRange resultTypes, + ValueRange inputs, + Location loc) { + if (inputs.size() != 1 || !isVMIType(inputs.front().getType())) + return {}; + auto unpackOp = builder.create(loc, resultTypes, inputs.front()); + return SmallVector(unpackOp->getResults()); +} + +static FailureOr getVMIVRegPhysicalElementType(VMIVRegType type) { + Type elementType = type.getElementType(); + VMILayoutAttr layout = type.getLayoutAttr(); + if (!layout || !layout.hasGroupSlotLaneStride()) + return elementType; + + auto integerType = dyn_cast(elementType); + if (!integerType && isa(elementType)) + return elementType; + if (!integerType) + return failure(); + if (!integerType.isUnsigned()) + return failure(); + unsigned elementBits = pto::getPTOStorageElemBitWidth(elementType); + int64_t laneStride = layout.getLaneStride(); + if (elementBits == 0 || laneStride <= 1) + return failure(); + if (elementBits == 8 && laneStride == 2) + return elementType; + int64_t physicalBits = static_cast(elementBits) * laneStride; + if (physicalBits != 16 && physicalBits != 32) + return failure(); + return IntegerType::get(type.getContext(), physicalBits); +} + +static int64_t getMaskGranularityBits(StringRef granularity) { + if (granularity == "b8") + return 8; + if (granularity == "b16") + return 16; + if (granularity == "b32") + return 32; + return 0; +} + +static StringRef getMaskGranularityForBits(int64_t bits) { + switch (bits) { + case 8: + return "b8"; + case 16: + return "b16"; + case 32: + return "b32"; + default: + return ""; + } +} + +static FailureOr getVMIMaskPhysicalGranularity(VMIMaskType type) { + int64_t bits = getMaskGranularityBits(type.getGranularity()); + if (bits == 0) + return failure(); + + VMILayoutAttr layout = type.getLayoutAttr(); + int64_t laneStride = layout && layout.hasLaneStride() ? layout.getLaneStride() + : 1; + int64_t physicalBits = bits * laneStride; + StringRef physicalGranularity = getMaskGranularityForBits(physicalBits); + if (physicalGranularity.empty()) + return failure(); + return physicalGranularity; +} + +class VMIToVPTOTypeConverter final : public TypeConverter { +public: + VMIToVPTOTypeConverter() { + addConversion([](Type type) { return type; }); + addConversion([](VMIVRegType type, + SmallVectorImpl &results) -> LogicalResult { + FailureOr arity = getVMIPhysicalArity(type); + FailureOr physicalElementType = getVMIVRegPhysicalElementType(type); + if (failed(arity) || failed(physicalElementType)) + return failure(); + FailureOr lanesPerPart = + getDataLanesPerPart(*physicalElementType); + if (failed(lanesPerPart)) + return failure(); + for (int64_t i = 0; i < *arity; ++i) + results.push_back(VRegType::get(type.getContext(), *lanesPerPart, + *physicalElementType)); + return success(); + }); + addConversion( + [](VMIMaskType type, SmallVectorImpl &results) -> LogicalResult { + FailureOr arity = getVMIPhysicalArity(type); + FailureOr physicalGranularity = + getVMIMaskPhysicalGranularity(type); + if (failed(arity) || failed(physicalGranularity)) + return failure(); + for (int64_t i = 0; i < *arity; ++i) + results.push_back( + MaskType::get(type.getContext(), *physicalGranularity)); + return success(); + }); + TypeConverter::addSourceMaterialization(materializeVPTOToVMI); + TypeConverter::addTargetMaterialization(materializeVMIToVPTO); + } +}; + +FailureOr> +getConvertedResultTypes(Operation *op, unsigned resultIndex, + const TypeConverter &typeConverter) { + if (resultIndex >= op->getNumResults()) + return failure(); + SmallVector resultTypes; + if (failed(typeConverter.convertType(op->getResult(resultIndex).getType(), + resultTypes))) + return failure(); + return resultTypes; +} + +FailureOr> +getConvertedResultTypes(Operation *op, const TypeConverter &typeConverter) { + SmallVector resultTypes; + if (failed(typeConverter.convertTypes(op->getResultTypes(), resultTypes))) + return failure(); + return resultTypes; +} + +FailureOr> +getConvertedVRegTypesWithLayout(VMIVRegType type, VMILayoutAttr layout, + const TypeConverter &typeConverter) { + auto relayoutType = VMIVRegType::get(type.getContext(), type.getElementCount(), + type.getElementType(), layout); + SmallVector convertedTypes; + if (failed(typeConverter.convertType(relayoutType, convertedTypes))) + return failure(); + return convertedTypes; +} + +FailureOr getVRegPhysicalFootprintBytes(TypeRange types) { + int64_t totalBytes = 0; + for (Type type : types) { + auto vregType = dyn_cast(type); + if (!vregType) + return failure(); + unsigned elementBits = + pto::getPTOStorageElemBitWidth(vregType.getElementType()); + if (elementBits == 0) + return failure(); + int64_t chunkBits = vregType.getElementCount() * elementBits; + if (chunkBits % 8 != 0) + return failure(); + totalBytes += chunkBits / 8; + } + return totalBytes; +} + +FailureOr hasNoWiderFootprintThanContiguous(TypeRange assignedTypes, + TypeRange contiguousTypes) { + FailureOr assignedBytes = + getVRegPhysicalFootprintBytes(assignedTypes); + FailureOr contiguousBytes = + getVRegPhysicalFootprintBytes(contiguousTypes); + if (failed(assignedBytes) || failed(contiguousBytes)) + return failure(); + return *assignedBytes <= *contiguousBytes; +} + +void replaceOpWithFlatConvertedValues( + ConversionPatternRewriter &rewriter, Operation *op, ValueRange flatValues, + const TypeConverter &typeConverter) { + SmallVector> replacements; + replacements.reserve(op->getNumResults()); + + auto valueIt = flatValues.begin(); + for (OpResult result : op->getResults()) { + SmallVector convertedTypes; + LogicalResult converted = + typeConverter.convertType(result.getType(), convertedTypes); + assert(succeeded(converted) && "expected converted result types"); + (void)converted; + assert(std::distance(valueIt, flatValues.end()) >= + static_cast(convertedTypes.size()) && + "not enough replacement values for converted results"); + replacements.emplace_back(valueIt, valueIt + convertedTypes.size()); + valueIt += convertedTypes.size(); + } + assert(valueIt == flatValues.end() && + "too many replacement values for converted results"); + + rewriter.replaceOpWithMultiple(op, std::move(replacements)); +} + +SmallVector +flattenOneToNOperands(ArrayRef operands) { + SmallVector flat; + for (ValueRange operand : operands) + llvm::append_range(flat, operand); + return flat; +} + +bool isIdentityOneToNValueMapping(ValueRange originalValues, + ArrayRef convertedValues) { + if (originalValues.size() != convertedValues.size()) + return false; + for (auto [original, converted] : + llvm::zip_equal(originalValues, convertedValues)) { + if (converted.size() != 1 || converted.front() != original) + return false; + } + return true; +} + +TypeRange getConvertedSignatureTypes( + const TypeConverter::SignatureConversion &conversion, + unsigned originalIndex) { + TypeRange convertedTypes = conversion.getConvertedTypes(); + if (auto mapping = conversion.getInputMapping(originalIndex)) + return convertedTypes.slice(mapping->inputNo, mapping->size); + return {}; +} + +bool hasNonIdentitySignatureConversion( + TypeRange originalTypes, + const TypeConverter::SignatureConversion &conversion) { + for (auto [index, originalType] : llvm::enumerate(originalTypes)) { + TypeRange convertedTypes = getConvertedSignatureTypes(conversion, index); + if (convertedTypes.size() != 1 || convertedTypes.front() != originalType) + return true; + } + return false; +} + +FailureOr createAllTrueMaskForVReg(Location loc, VRegType vregType, + PatternRewriter &rewriter) { + MLIRContext *ctx = rewriter.getContext(); + unsigned elementBits = + pto::getPTOStorageElemBitWidth(vregType.getElementType()); + if (elementBits == 8) + return rewriter + .create(loc, MaskType::get(ctx, "b8"), + rewriter.getStringAttr("PAT_ALL")) + .getResult(); + if (elementBits == 16) + return rewriter + .create(loc, MaskType::get(ctx, "b16"), + rewriter.getStringAttr("PAT_ALL")) + .getResult(); + if (elementBits == 32) + return rewriter + .create(loc, MaskType::get(ctx, "b32"), + rewriter.getStringAttr("PAT_ALL")) + .getResult(); + return failure(); +} + +FailureOr getMaskTypeForVReg(VRegType vregType, MLIRContext *ctx) { + unsigned elementBits = + pto::getPTOStorageElemBitWidth(vregType.getElementType()); + if (elementBits == 8) + return MaskType::get(ctx, "b8"); + if (elementBits == 16) + return MaskType::get(ctx, "b16"); + if (elementBits == 32) + return MaskType::get(ctx, "b32"); + return failure(); +} + +FailureOr createAllTrueMask(Location loc, MaskType maskType, + PatternRewriter &rewriter) { + StringAttr pattern = rewriter.getStringAttr("PAT_ALL"); + MLIRContext *ctx = rewriter.getContext(); + if (maskType.isB8()) + return rewriter.create(loc, MaskType::get(ctx, "b8"), pattern) + .getResult(); + if (maskType.isB16()) + return rewriter.create(loc, MaskType::get(ctx, "b16"), pattern) + .getResult(); + if (maskType.isB32()) + return rewriter.create(loc, MaskType::get(ctx, "b32"), pattern) + .getResult(); + return failure(); +} + +FailureOr createPatternMask(Location loc, MaskType maskType, + StringRef pattern, + PatternRewriter &rewriter) { + StringAttr patternAttr = rewriter.getStringAttr(pattern); + MLIRContext *ctx = rewriter.getContext(); + if (maskType.isB8()) + return rewriter.create(loc, MaskType::get(ctx, "b8"), patternAttr) + .getResult(); + if (maskType.isB16()) + return rewriter + .create(loc, MaskType::get(ctx, "b16"), patternAttr) + .getResult(); + if (maskType.isB32()) + return rewriter + .create(loc, MaskType::get(ctx, "b32"), patternAttr) + .getResult(); + return failure(); +} + +FailureOr createPrefixMask(Location loc, MaskType maskType, + StringRef pattern, + PatternRewriter &rewriter) { + StringAttr patternAttr = rewriter.getStringAttr(pattern); + MLIRContext *ctx = rewriter.getContext(); + if (maskType.isB8()) + return rewriter.create(loc, MaskType::get(ctx, "b8"), patternAttr) + .getResult(); + if (maskType.isB16()) + return rewriter + .create(loc, MaskType::get(ctx, "b16"), patternAttr) + .getResult(); + if (maskType.isB32()) + return rewriter + .create(loc, MaskType::get(ctx, "b32"), patternAttr) + .getResult(); + return failure(); +} + +FailureOr> +createRuntimePrefixMask(Location loc, MaskType maskType, Value activeLanes, + PatternRewriter &rewriter) { + MLIRContext *ctx = rewriter.getContext(); + Type scalarType = activeLanes.getType(); + if (maskType.isB8()) { + auto op = rewriter.create(loc, MaskType::get(ctx, "b8"), + scalarType, activeLanes); + return std::make_pair(Value(op.getMask()), Value(op.getScalarOut())); + } + if (maskType.isB16()) { + auto op = rewriter.create(loc, MaskType::get(ctx, "b16"), + scalarType, activeLanes); + return std::make_pair(Value(op.getMask()), Value(op.getScalarOut())); + } + if (maskType.isB32()) { + auto op = rewriter.create(loc, MaskType::get(ctx, "b32"), + scalarType, activeLanes); + return std::make_pair(Value(op.getMask()), Value(op.getScalarOut())); + } + return failure(); +} + +LogicalResult +checkSupportedMaskableVReg(VMIVRegType type, std::string *reason = nullptr) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + FailureOr lanesPerPart = getDataLanesPerPart(type.getElementType()); + FailureOr arity = getVMIPhysicalArity(type); + if (failed(lanesPerPart) || failed(arity) || *arity < 1) + return fail("requires computable non-empty physical vreg parts"); + + return success(); +} + +Value createI32Constant(Location loc, int64_t value, + PatternRewriter &rewriter) { + return rewriter.create(loc, value, 32); +} + +Value createI16Constant(Location loc, int64_t value, + PatternRewriter &rewriter) { + return rewriter.create(loc, value, 16); +} + +FailureOr createPrefixMaskForActiveLanes(Location loc, MaskType maskType, + int64_t activeLanes, + PatternRewriter &rewriter) { + if (activeLanes <= 0) + return createPrefixMask(loc, maskType, "PAT_ALLF", rewriter); + + switch (activeLanes) { + case 1: + case 2: + case 3: + case 4: + case 8: + case 16: + case 32: + case 64: + case 128: + return createPrefixMask( + loc, maskType, (Twine("PAT_VL") + Twine(activeLanes)).str(), rewriter); + default: { + FailureOr> dynamicMask = createRuntimePrefixMask( + loc, maskType, createI32Constant(loc, activeLanes, rewriter), rewriter); + if (failed(dynamicMask)) + return failure(); + return dynamicMask->first; + } + } +} + +Value clampDynamicActiveLanes(Location loc, Value activeLanes, + int64_t maxActiveLanes, + PatternRewriter &rewriter) { + Value activeI32 = rewriter.create( + loc, rewriter.getI32Type(), activeLanes); + Value zeroI32 = createI32Constant(loc, 0, rewriter); + Value nonNegative = rewriter.create(loc, activeI32, zeroI32); + Value maxI32 = createI32Constant(loc, maxActiveLanes, rewriter); + return rewriter.create(loc, nonNegative, maxI32); +} + +Value createPartitionActiveLanes(Location loc, Value activeLanesI32, + int64_t factor, int64_t part, + PatternRewriter &rewriter) { + if (factor == 1) + return activeLanesI32; + int64_t bias = factor - 1 - part; + Value biased = activeLanesI32; + if (bias != 0) + biased = rewriter.create( + loc, biased, createI32Constant(loc, bias, rewriter)); + return rewriter.create( + loc, biased, createI32Constant(loc, factor, rewriter)); +} + +std::optional getPowerOfTwoLog2(int64_t value) { + if (value <= 0 || (value & (value - 1)) != 0) + return std::nullopt; + int64_t log2 = 0; + while (value > 1) { + value >>= 1; + ++log2; + } + return log2; +} + +std::optional getPrefixPattern(int64_t activeLanes, + int64_t lanesPerPart) { + if (activeLanes <= 0) + return std::string("PAT_ALLF"); + if (activeLanes >= lanesPerPart) + return std::string("PAT_ALL"); + switch (activeLanes) { + case 1: + case 2: + case 3: + case 4: + case 8: + case 16: + case 32: + case 64: + case 128: + return std::string("PAT_VL") + std::to_string(activeLanes); + default: + return std::nullopt; + } +} + +FailureOr getSingleValue(Operation *op, ValueRange values, + StringRef description, + PatternRewriter &rewriter) { + if (values.size() != 1) { + (void)rewriter.notifyMatchFailure(op, description); + return failure(); + } + return values.front(); +} + +static int64_t ceilDivNonNegative(int64_t lhs, int64_t rhs) { + return (lhs + rhs - 1) / rhs; +} + +FailureOr getDataLayoutFactor(VMIVRegType type) { + VMILayoutAttr layout = type.getLayoutAttr(); + if (!layout) + return failure(); + return layout.isDeinterleaved() ? layout.getFactor() : 1; +} + +FailureOr getDataChunksInPart(VMIVRegType type, int64_t part) { + FailureOr factor = getDataLayoutFactor(type); + FailureOr lanesPerPart = getDataLanesPerPart(type.getElementType()); + if (failed(factor) || failed(lanesPerPart) || part < 0 || part >= *factor) + return failure(); + + int64_t logicalLanesInPart = + (type.getElementCount() + *factor - 1 - part) / *factor; + return ceilDivNonNegative(logicalLanesInPart, *lanesPerPart); +} + +FailureOr getDataFlatPartIndex(VMIVRegType type, int64_t part, + int64_t chunk) { + FailureOr factor = getDataLayoutFactor(type); + if (failed(factor) || part < 0 || part >= *factor || chunk < 0) + return failure(); + + int64_t flatIndex = 0; + for (int64_t currentPart = 0; currentPart < part; ++currentPart) { + FailureOr chunks = getDataChunksInPart(type, currentPart); + if (failed(chunks)) + return failure(); + flatIndex += *chunks; + } + + FailureOr chunks = getDataChunksInPart(type, part); + if (failed(chunks) || chunk >= *chunks) + return failure(); + return flatIndex + chunk; +} + +FailureOr checkFullDataPhysicalChunks(VMIVRegType type, + std::string *reason) { + auto fail = [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + FailureOr lanesPerPart = getDataLanesPerPart(type.getElementType()); + if (failed(lanesPerPart)) + return fail("requires known physical lanes per part"); + + FailureOr factor = getDataLayoutFactor(type); + if (failed(factor)) + return fail("requires assigned layout"); + + for (int64_t part = 0; part < *factor; ++part) { + FailureOr chunks = getDataChunksInPart(type, part); + if (failed(chunks)) + return fail("requires known physical chunks"); + for (int64_t chunk = 0; chunk < *chunks; ++chunk) { + for (int64_t lane = 0; lane < *lanesPerPart; ++lane) { + FailureOr padding = isPaddingLane(type, part, chunk, lane); + if (failed(padding)) + return fail("failed to map physical padding lane"); + if (*padding) + return fail("found padding lane in physical chunk"); + } + } + } + + return *lanesPerPart; +} + +FailureOr getVMITypeLayoutFactor(Type type) { + Attribute layout; + if (auto vregType = dyn_cast(type)) + layout = vregType.getLayout(); + else if (auto maskType = dyn_cast(type)) + layout = maskType.getLayout(); + else + return failure(); + + auto layoutAttr = dyn_cast_or_null(layout); + if (!layoutAttr) + return failure(); + return layoutAttr.isDeinterleaved() ? layoutAttr.getFactor() : 1; +} + +FailureOr getVMITypeElementCount(Type type) { + if (auto vregType = dyn_cast(type)) + return vregType.getElementCount(); + if (auto maskType = dyn_cast(type)) + return maskType.getElementCount(); + return failure(); +} + +FailureOr getVMITypeLanesPerPart(Type type) { + if (auto vregType = dyn_cast(type)) { + FailureOr physicalElementType = + getVMIVRegPhysicalElementType(vregType); + if (failed(physicalElementType)) + return failure(); + return getDataLanesPerPart(*physicalElementType); + } + if (auto maskType = dyn_cast(type)) { + FailureOr physicalGranularity = + getVMIMaskPhysicalGranularity(maskType); + if (failed(physicalGranularity)) + return failure(); + return getMaskLanesPerPart(*physicalGranularity); + } + return failure(); +} + +FailureOr getVMITypeChunksInPart(Type type, int64_t part) { + FailureOr elementCount = getVMITypeElementCount(type); + FailureOr factor = getVMITypeLayoutFactor(type); + FailureOr lanesPerPart = getVMITypeLanesPerPart(type); + if (failed(elementCount) || failed(factor) || failed(lanesPerPart) || + part < 0 || part >= *factor) + return failure(); + + VMILayoutAttr layout; + if (auto vregType = dyn_cast(type)) + layout = vregType.getLayoutAttr(); + else if (auto maskType = dyn_cast(type)) + layout = maskType.getLayoutAttr(); + if (!layout) + return failure(); + + int64_t logicalLanesInPart = (*elementCount + *factor - 1 - part) / *factor; + int64_t laneStride = 1; + if (isa(type) && layout.isDense()) + laneStride = layout.getLaneStride(); + int64_t physicalLanes = + logicalLanesInPart == 0 ? 0 : (logicalLanesInPart - 1) * laneStride + 1; + return ceilDivNonNegative(physicalLanes, *lanesPerPart); +} + +LogicalResult checkFullVMIPhysicalChunks(Type type, std::string *reason) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + FailureOr factor = getVMITypeLayoutFactor(type); + FailureOr lanesPerPart = getVMITypeLanesPerPart(type); + if (failed(factor) || failed(lanesPerPart)) + return fail("requires assigned layout with known physical lanes per part"); + + for (int64_t part = 0; part < *factor; ++part) { + FailureOr chunks = getVMITypeChunksInPart(type, part); + if (failed(chunks)) + return fail("requires known physical chunks"); + for (int64_t chunk = 0; chunk < *chunks; ++chunk) { + for (int64_t lane = 0; lane < *lanesPerPart; ++lane) { + FailureOr padding = isPaddingLane(type, part, chunk, lane); + if (failed(padding)) + return fail("failed to map physical padding lane"); + if (*padding) + return fail("found padding lane in physical chunk"); + } + } + } + + return success(); +} + +FailureOr getContiguousMaterializationPartCount(Type type, + std::string *reason); + +FailureOr getContiguousMaterializationPartCount(Type type, + std::string *reason) { + auto fail = [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + FailureOr arity = getVMIPhysicalArity(type); + FailureOr factor = getVMITypeLayoutFactor(type); + if (failed(arity) || failed(factor)) + return fail("requires computable physical arity and assigned layout"); + + Attribute layoutAttr; + if (auto vregType = dyn_cast(type)) + layoutAttr = vregType.getLayout(); + else if (auto maskType = dyn_cast(type)) + layoutAttr = maskType.getLayout(); + else + return fail("requires VMI data or mask type"); + + auto layout = dyn_cast_or_null(layoutAttr); + if (!layout) + return fail("requires assigned layout"); + if (layout.isContiguous() && layout.getLaneStride() == 1) + return *arity; + if (!layout.isDeinterleaved() || + (layout.getFactor() != 2 && layout.getFactor() != 4)) + return fail("requires contiguous or deinterleaved=2/4 layout"); + + FailureOr chunksPerGroup = getVMITypeChunksInPart(type, 0); + if (failed(chunksPerGroup)) + return fail("requires known physical chunks per part"); + if (*chunksPerGroup == 0) + return fail("requires at least one physical chunk per part"); + + for (int64_t part = 1; part < *factor; ++part) { + FailureOr chunks = getVMITypeChunksInPart(type, part); + if (failed(chunks)) + return fail("requires known physical chunks per part"); + if (layout.getFactor() == 2 && *chunks != *chunksPerGroup) + return fail("requires every deinterleaved part to have the same " + "physical chunk count"); + } + + VMILayoutAttr contiguous = VMILayoutAttr::getContiguous(type.getContext()); + Type contiguousType; + if (auto vregType = dyn_cast(type)) { + contiguousType = + VMIVRegType::get(type.getContext(), vregType.getElementCount(), + vregType.getElementType(), contiguous); + } else { + auto maskType = cast(type); + contiguousType = + VMIMaskType::get(type.getContext(), maskType.getElementCount(), + maskType.getGranularity(), contiguous); + } + return getVMIPhysicalArity(contiguousType); +} + +LogicalResult checkCanMaterializeToContiguous(Type type, std::string *reason) { + return succeeded(getContiguousMaterializationPartCount(type, reason)) + ? success() + : failure(); +} + +std::optional getConstantIndexValue(Value value) { + if (auto constant = value.getDefiningOp()) + return constant.value(); + if (auto constant = value.getDefiningOp()) { + if (auto integerAttr = dyn_cast(constant.getValue())) + return integerAttr.getInt(); + } + return std::nullopt; +} + +bool isKnownIndexMultipleOf(Value value, int64_t multiple, int depth = 0) { + if (multiple <= 1) + return true; + if (depth > 6) + return false; + if (std::optional constant = getConstantIndexValue(value)) + return *constant % multiple == 0; + + if (auto add = value.getDefiningOp()) + return isKnownIndexMultipleOf(add.getLhs(), multiple, depth + 1) && + isKnownIndexMultipleOf(add.getRhs(), multiple, depth + 1); + if (auto sub = value.getDefiningOp()) + return isKnownIndexMultipleOf(sub.getLhs(), multiple, depth + 1) && + isKnownIndexMultipleOf(sub.getRhs(), multiple, depth + 1); + if (auto mul = value.getDefiningOp()) + return isKnownIndexMultipleOf(mul.getLhs(), multiple, depth + 1) || + isKnownIndexMultipleOf(mul.getRhs(), multiple, depth + 1); + + return false; +} + +FailureOr getStaticMemRefElementCount(Type type) { + auto memrefType = dyn_cast(type); + if (!memrefType || !memrefType.hasStaticShape()) + return failure(); + + int64_t elements = 1; + for (int64_t dim : memrefType.getShape()) + elements *= dim; + return elements; +} + +static Type getMemoryElementType(Type type) { + if (auto ptrType = dyn_cast(type)) + return ptrType.getElementType(); + if (auto memrefType = dyn_cast(type)) + return memrefType.getElementType(); + return {}; +} + +static bool isPackedByteGroupStore(Type destinationType, VRegType valueType) { + Type destinationElementType = getMemoryElementType(destinationType); + auto destinationIntegerType = + dyn_cast_or_null(destinationElementType); + auto valueIntegerType = dyn_cast(valueType.getElementType()); + return destinationIntegerType && valueIntegerType && + pto::getPTOStorageElemBitWidth(destinationIntegerType) == 8 && + pto::getPTOStorageElemBitWidth(valueIntegerType) == 32; +} + +enum class VMIMemoryValidMaskKind { + AllTrue, + ExplicitMask, +}; + +enum class VMIMemoryWriteMaskKind { + AllTrue, + ExplicitMask, +}; + +enum class VMIMemoryPermutationKind { + Identity, +}; + +enum class VMIMemoryFallbackDecisionKind { + NotRequired, + RequiredUnavailable, +}; + +struct VMIMemoryLogicalShape { + int64_t elementCount = 0; +}; + +struct VMIMemoryLaneAddressMap { + VMIMemoryPermutationKind permutation = VMIMemoryPermutationKind::Identity; + int64_t baseElementOffset = 0; + int64_t elementStride = 1; + int64_t physicalLaneFootprint = 0; + + int64_t getExclusiveEndElement() const { + return baseElementOffset + physicalLaneFootprint * elementStride; + } +}; + +struct VMIMemoryFallbackDecision { + VMIMemoryFallbackDecisionKind kind = + VMIMemoryFallbackDecisionKind::NotRequired; + std::string reason = "not required"; + + static VMIMemoryFallbackDecision notRequired() { return {}; } + + static VMIMemoryFallbackDecision requiredUnavailable(const Twine &reason) { + VMIMemoryFallbackDecision decision; + decision.kind = VMIMemoryFallbackDecisionKind::RequiredUnavailable; + decision.reason = reason.str(); + return decision; + } +}; + +struct VMIMemorySafeReadProof { + bool proven = false; + std::string reason; + std::optional constantOffset; + std::optional staticElementCount; + std::optional laneAddressMap; + int64_t physicalFootprint = 0; +}; + +struct VMIMemoryAccessPlan { + Type baseType; + VMIVRegType valueType; + std::optional constantOffset; + VMIMemoryLogicalShape logicalShape; + VMIMemoryValidMaskKind validMask = VMIMemoryValidMaskKind::AllTrue; + VMIMemoryPermutationKind permutation = VMIMemoryPermutationKind::Identity; + std::optional laneAddressMap; + Attribute paddingValue; + VMIMemoryWriteMaskKind writeMask = VMIMemoryWriteMaskKind::AllTrue; + VMIMemorySafeReadProof safeReadProof; + VMISupportResult layoutSupport; + VMIMemoryFallbackDecision fallbackDecision; +}; + +FailureOr +buildContiguousIdentityLaneAddressMap(int64_t constantOffset, + VMIVRegType resultType, + std::string *reason = nullptr) { + auto fail = [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + FailureOr lanesPerPart = + getDataLanesPerPart(resultType.getElementType()); + FailureOr arity = getVMIPhysicalArity(resultType); + if (failed(lanesPerPart) || failed(arity)) + return fail("requires computable physical read footprint"); + + VMIMemoryLaneAddressMap map; + map.baseElementOffset = constantOffset; + map.physicalLaneFootprint = *arity * *lanesPerPart; + return map; +} + +VMISupportResult requireIdentityMemRefLayout(Type memoryType, StringRef role, + Value memoryValue = {}) { + auto memrefType = dyn_cast(memoryType); + if (!memrefType || memrefType.getLayout().isIdentity()) + return VMISupportResult::success(); + std::string reason = + (Twine(role) + + " memref layout is non-identity; current VMI memory access plan " + "supports only contiguous identity lane-to-address maps") + .str(); + if (memoryValue && memoryValue.getDefiningOp()) + reason += "; memref.subview requires normalized base/offset/stride " + "lane-to-address planning"; + return VMISupportResult::failure(reason); +} + +VMIMemorySafeReadProof +computeSafeFullReadProof(Type sourceType, std::optional constantOffset, + VMIVRegType resultType) { + VMIMemorySafeReadProof proof; + proof.constantOffset = constantOffset; + + auto fail = [&](const Twine &message) { + proof.proven = false; + proof.reason = message.str(); + return proof; + }; + + if (!constantOffset) + return fail("requires constant index offset"); + + FailureOr staticElements = getStaticMemRefElementCount(sourceType); + if (failed(staticElements)) + return fail("requires statically shaped memref source"); + int64_t elements = *staticElements; + proof.staticElementCount = elements; + + if (*constantOffset < 0) + return fail("requires non-negative offset"); + + std::string addressMapReason; + FailureOr addressMap = + buildContiguousIdentityLaneAddressMap(*constantOffset, resultType, + &addressMapReason); + if (failed(addressMap)) + return fail(addressMapReason); + proof.laneAddressMap = *addressMap; + + proof.physicalFootprint = addressMap->physicalLaneFootprint; + if (addressMap->getExclusiveEndElement() > elements) + return fail(Twine("full physical read footprint [") + + Twine(addressMap->baseElementOffset) + ", " + + Twine(addressMap->getExclusiveEndElement()) + + ") exceeds static memref element count " + Twine(elements)); + + proof.proven = true; + return proof; +} + +VMIMemoryAccessPlan +buildReadAccessPlan(Value source, Type sourceType, VMIVRegType resultType, + std::optional constantOffset, + VMIMemoryValidMaskKind validMask) { + VMIMemoryAccessPlan plan; + plan.baseType = sourceType; + plan.valueType = resultType; + plan.constantOffset = constantOffset; + plan.logicalShape.elementCount = resultType.getElementCount(); + plan.validMask = validMask; + plan.permutation = VMIMemoryPermutationKind::Identity; + plan.writeMask = VMIMemoryWriteMaskKind::AllTrue; + plan.safeReadProof = + computeSafeFullReadProof(sourceType, constantOffset, resultType); + plan.laneAddressMap = plan.safeReadProof.laneAddressMap; + plan.layoutSupport = + requireIdentityMemRefLayout(sourceType, "source", source); + return plan; +} + +VMIMemoryAccessPlan +buildWriteAccessPlan(Value destination, Type destinationType, + VMIVRegType valueType, VMIMemoryWriteMaskKind writeMask) { + VMIMemoryAccessPlan plan; + plan.baseType = destinationType; + plan.valueType = valueType; + plan.logicalShape.elementCount = valueType.getElementCount(); + plan.validMask = VMIMemoryValidMaskKind::AllTrue; + plan.permutation = VMIMemoryPermutationKind::Identity; + plan.writeMask = writeMask; + plan.layoutSupport = + requireIdentityMemRefLayout(destinationType, "destination", destination); + return plan; +} + +void requireUnavailableReadFallback(VMIMemoryAccessPlan &plan) { + std::string maskedLoadReason; + if (plan.validMask == VMIMemoryValidMaskKind::ExplicitMask) + maskedLoadReason = + "; target true masked/non-faulting load is unavailable because the " + "current VPTO pto.vlds surface has no mask operand"; + std::string scratchReason = + "; scratch memory fallback resource allocation is not implemented"; + std::string guardedReason = + "; guarded memory fallback control-flow lowering is not implemented"; + plan.fallbackDecision = VMIMemoryFallbackDecision::requiredUnavailable( + Twine("partial/tail read needs a scratch, guarded, or true " + "masked/non-faulting load fallback, but no such fallback resource " + "plan is implemented") + + maskedLoadReason + scratchReason + guardedReason); +} + +FailureOr verifyFullOrSafeReadVRegChunks(Operation *op, + VMIVRegType type, + Type sourceType, Value offset, + PatternRewriter &rewriter) { + std::string fullChunkReason; + FailureOr lanesPerPart = + checkFullDataPhysicalChunks(type, &fullChunkReason); + if (succeeded(lanesPerPart)) + return *lanesPerPart; + + VMIMemorySafeReadProof safeReadProof = + computeSafeFullReadProof(sourceType, getConstantIndexValue(offset), type); + if (safeReadProof.proven) { + lanesPerPart = getDataLanesPerPart(type.getElementType()); + if (succeeded(lanesPerPart)) + return *lanesPerPart; + } + + lanesPerPart = getDataLanesPerPart(type.getElementType()); + if (succeeded(lanesPerPart)) + return *lanesPerPart; + + (void)rewriter.notifyMatchFailure( + op, Twine("memory lowering ") + fullChunkReason + + "; safe full-read proof failed: " + safeReadProof.reason); + return failure(); +} + +LogicalResult +checkSupportedLoadShape(VMIVRegType type, Value source, Type sourceType, + std::optional constantOffset, + std::string *reason) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + VMIMemoryAccessPlan accessPlan = + buildReadAccessPlan(source, sourceType, type, + constantOffset, VMIMemoryValidMaskKind::AllTrue); + if (!accessPlan.layoutSupport.isSupported()) + return fail(accessPlan.layoutSupport.reason); + + VMILayoutSupport supports; + if (failed(supports.getLoadLayoutFact(type, reason))) + return failure(); + + if (getDenseLaneStrideLoadDistToken(type)) + return success(); + + if (failed(getDataLanesPerPart(type.getElementType()))) + return fail("requires element type with known physical lane width"); + return success(); +} + +LogicalResult checkSupportedDeinterleaveLoadShape( + VMIDeinterleaveLoadOp op, + std::string *reason) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + auto lowType = cast(op.getLow().getType()); + auto highType = cast(op.getHigh().getType()); + VMILayoutSupport supports; + if (failed(supports.getDeinterleaveLoadLayoutFactForLayouts( + lowType, highType, reason))) + return failure(); + if (!getX2MemoryDistToken(lowType.getElementType(), "DINTLV")) + return fail("requires 8/16/32-bit element type for vldsx2 DINTLV"); + + VMIMemoryAccessPlan accessPlan = buildReadAccessPlan(op.getSource(), op.getSource().getType(), lowType, + getConstantIndexValue(op.getOffset()), VMIMemoryValidMaskKind::AllTrue); + if (!accessPlan.layoutSupport.isSupported()) + return fail(accessPlan.layoutSupport.reason); + + std::string fullChunkReason; + if (failed(checkFullDataPhysicalChunks(lowType, &fullChunkReason))) + return fail(Twine("requires full physical chunks; ") + fullChunkReason); + return success(); +} + +LogicalResult +checkSupportedStoreShape(VMIVRegType type, Value destination, + Type destinationType, std::string *reason) { + VMIMemoryAccessPlan accessPlan = + buildWriteAccessPlan(destination, destinationType, type, + VMIMemoryWriteMaskKind::AllTrue); + if (!accessPlan.layoutSupport.isSupported()) { + if (reason) + *reason = accessPlan.layoutSupport.reason; + return failure(); + } + + if (failed(checkSupportedMaskableVReg(type, reason))) + return failure(); + + VMILayoutSupport supports; + if (failed(supports.getStoreLayoutFact(type, reason))) + return failure(); + + if (getDenseLaneStrideStoreDistToken(type)) + return success(); + + std::string fullChunkReason; + if (succeeded(checkFullDataPhysicalChunks(type, &fullChunkReason))) + return success(); + + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + VMILayoutAttr layout = type.getLayoutAttr(); + if (!layout) + return fail("requires assigned layout"); + if (failed(getDataLanesPerPart(type.getElementType()))) + return fail("requires known physical lanes per part"); + if (layout.isContiguous() && layout.getLaneStride() == 1) + return success(); + + std::string materializationReason; + if (succeeded(checkCanMaterializeToContiguous(type, &materializationReason))) + return success(); + return fail(Twine("partial/tail store requires contiguous layout or " + "deinterleaved layout that can materialize to contiguous; " + "value ") + + fullChunkReason + ", materialization " + materializationReason); +} + +LogicalResult checkSupportedInterleaveStoreShape( + VMIInterleaveStoreOp op, + std::string *reason) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + auto lowType = cast(op.getLow().getType()); + auto highType = cast(op.getHigh().getType()); + VMILayoutAttr lowLayout = lowType.getLayoutAttr(); + VMILayoutAttr highLayout = highType.getLayoutAttr(); + if (!lowLayout || !highLayout || !lowLayout.isContiguous() || + !highLayout.isContiguous()) + return fail("requires assigned contiguous low/high input layouts"); + if (lowType.getElementCount() != highType.getElementCount() || + lowType.getElementType() != highType.getElementType()) + return fail("requires matching low/high input shape and element type"); + if (!getX2MemoryDistToken(lowType.getElementType(), "INTLV")) + return fail("requires 8/16/32-bit element type for vstsx2 INTLV"); + + VMIMemoryAccessPlan accessPlan = buildWriteAccessPlan(op.getDestination(), op.getDestination().getType(), lowType, + VMIMemoryWriteMaskKind::AllTrue); + if (!accessPlan.layoutSupport.isSupported()) + return fail(accessPlan.layoutSupport.reason); + if (failed(checkSupportedMaskableVReg(lowType, reason))) + return failure(); + + std::string fullChunkReason; + if (failed(checkFullDataPhysicalChunks(lowType, &fullChunkReason))) + return fail(Twine("requires full physical chunks; ") + fullChunkReason); + return success(); +} + +FailureOr getGroupSizeFromNumGroups(VMIVRegType type, + int64_t numGroups, + std::string *reason = nullptr) { + auto fail = [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + if (numGroups <= 0) + return fail("requires num_groups to be positive"); + if (type.getElementCount() % numGroups != 0) + return fail("requires num_groups to evenly divide logical lane count"); + return type.getElementCount() / numGroups; +} + +LogicalResult checkSupportedGroupChunkShape(VMIVRegType type, int64_t groupSize, + std::string *reason) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + VMILayoutAttr layout = type.getLayoutAttr(); + if (!layout || !layout.isContiguous()) + return fail("requires assigned contiguous layout"); + std::string fullChunkReason; + if (failed(checkFullDataPhysicalChunks(type, &fullChunkReason))) + return fail(Twine("requires full physical chunks; ") + fullChunkReason); + FailureOr lanesPerPart = getDataLanesPerPart(type.getElementType()); + if (failed(lanesPerPart)) + return fail("requires known physical lanes per part"); + if (groupSize <= 0 || type.getElementCount() % groupSize != 0) + return fail("requires derived group size to evenly divide logical lane " + "count"); + if (groupSize % *lanesPerPart != 0) + return fail("currently requires group size to be a multiple of physical " + "lanes per part"); + return success(); +} + +LogicalResult checkDeinterleaved2GroupStoreChunkShape( + VMIVRegType type, int64_t groupSize, int64_t *lanesPerPart, + int64_t *groupCount, int64_t *chunksPerGroupPerPart, + std::string *reason) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + VMILayoutAttr layout = type.getLayoutAttr(); + if (!layout || !layout.isDeinterleaved() || layout.getFactor() != 2 || + layout.getBlockElems() != 1 || layout.getLaneStride() != 1) + return fail("requires deinterleaved=2 value layout"); + std::string fullChunkReason; + if (failed(checkFullDataPhysicalChunks(type, &fullChunkReason))) + return fail(Twine("requires full physical chunks; ") + fullChunkReason); + FailureOr lanes = getDataLanesPerPart(type.getElementType()); + if (failed(lanes)) + return fail("requires known physical lanes per part"); + if (!getX2MemoryDistToken(type.getElementType(), "INTLV")) + return fail("requires 8/16/32-bit element type for vstsx2 INTLV"); + if (groupSize <= 0 || type.getElementCount() % groupSize != 0) + return fail("requires derived group size to evenly divide logical lane " + "count"); + int64_t pairLanes = 2 * *lanes; + if (groupSize % pairLanes != 0) + return fail("requires group size to be a multiple of two physical chunks"); + + FailureOr part0Chunks = getDataChunksInPart(type, /*part=*/0); + FailureOr part1Chunks = getDataChunksInPart(type, /*part=*/1); + if (failed(part0Chunks) || failed(part1Chunks) || + *part0Chunks != *part1Chunks) + return fail("requires matching deinterleaved part chunk counts"); + + *lanesPerPart = *lanes; + *groupCount = type.getElementCount() / groupSize; + *chunksPerGroupPerPart = groupSize / pairLanes; + if (*part0Chunks != *groupCount * *chunksPerGroupPerPart) + return fail("requires deinterleaved chunks to align with group rows"); + return success(); +} + +LogicalResult +checkSupportedGroupLoadShape(VMIGroupLoadOp op, std::string *reason) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + auto resultType = cast(op.getResult().getType()); + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + if (!resultLayout) + return fail("requires assigned result layout"); + FailureOr groupSize = getGroupSizeFromNumGroups( + resultType, op.getNumGroupsAttr().getInt(), reason); + if (failed(groupSize)) + return failure(); + + if (resultLayout.isContiguous()) { + VMILayoutSupport supports; + if (failed(supports.getGroupLoadLayoutFact(op, reason))) + return failure(); + if (failed(checkSupportedLoadShape(resultType, op.getSource(), + op.getSource().getType(), std::nullopt, + reason))) + return failure(); + std::optional rowStride = getConstantIndexValue(op.getRowStride()); + if (rowStride && *rowStride == *groupSize) + return success(); + return checkSupportedGroupChunkShape(resultType, *groupSize, reason); + } + + if (resultLayout.isDeinterleaved() && resultLayout.getBlockElems() == 8 && + resultType.getElementType().isF32()) { + VMILayoutSupport supports; + if (failed(supports.getGroupLoadLayoutFact(op, reason))) + return failure(); + VMIMemoryAccessPlan accessPlan = buildReadAccessPlan(op.getSource(), op.getSource().getType(), resultType, + getConstantIndexValue(op.getOffset()), VMIMemoryValidMaskKind::AllTrue); + if (!accessPlan.layoutSupport.isSupported()) + return fail(accessPlan.layoutSupport.reason); + if (!isa(op.getSource().getType())) + return fail("block8 strided group_load requires !pto.ptr source"); + if (op.getNumGroupsAttr().getInt() % 8 != 0) + return fail( + "block8 strided group_load requires num_groups multiple of 8"); + std::optional rowStride = getConstantIndexValue(op.getRowStride()); + if (!rowStride || *rowStride <= 0 || *rowStride % 8 != 0) + return fail("block8 strided group_load requires constant positive " + "row_stride divisible by 8 f32 elements"); + std::string fullChunkReason; + if (failed(checkFullDataPhysicalChunks(resultType, &fullChunkReason))) + return fail(Twine("block8 strided group_load requires full physical " + "result chunks; ") + + fullChunkReason); + return success(); + } + + return fail("requires contiguous layout or deinterleaved block8 f32 layout"); +} + +LogicalResult checkSupportedGroupSlotLoadShape( + VMIGroupSlotLoadOp op, + std::string *reason) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + auto resultType = cast(op.getResult().getType()); + VMILayoutSupport supports; + FailureOr fact = supports.getGroupSlotLoadLayoutFact( + resultType, op.getNumGroupsAttr().getInt(), reason); + if (failed(fact)) + return failure(); + + VMIMemoryAccessPlan accessPlan = buildReadAccessPlan(op.getSource(), op.getSource().getType(), resultType, + getConstantIndexValue(op.getOffset()), VMIMemoryValidMaskKind::AllTrue); + if (!accessPlan.layoutSupport.isSupported()) + return fail(accessPlan.layoutSupport.reason); + if (!isa(op.getSource().getType())) + return fail("group_slot_load requires !pto.ptr source"); + + if (fact->slots == 8) { + std::optional sourceGroupStride = + getConstantIndexValue(op.getSourceGroupStride()); + if (!sourceGroupStride || *sourceGroupStride != 1) + return fail("slots=8 group_slot_load requires constant unit " + "source_group_stride"); + return success(); + } + + unsigned elementBits = + pto::getPTOStorageElemBitWidth(resultType.getElementType()); + if (elementBits == 0 || 256 % elementBits != 0) + return fail("slots=1 group_slot_load requires supported element width"); + int64_t alignedStrideElems = 256 / elementBits; + std::optional sourceGroupStride = + getConstantIndexValue(op.getSourceGroupStride()); + if (!sourceGroupStride || *sourceGroupStride <= 0 || + *sourceGroupStride % alignedStrideElems != 0) + return fail(Twine("slots=1 group_slot_load currently lowers as one " + "lane-0 vsldb per group and requires constant " + "positive source_group_stride divisible by ") + + Twine(alignedStrideElems) + + " elements for 32B load alignment; packed or unaligned " + "scalar load lowering is not implemented"); + return success(); +} + +LogicalResult checkSupportedGroupBroadcastLoadShape( + VMIGroupBroadcastLoadOp op, + std::string *reason) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + VMILayoutSupport supports; + if (failed(supports.getGroupBroadcastLoadSupport(op, reason))) + return failure(); + VMIMemoryAccessPlan accessPlan = buildReadAccessPlan(op.getSource(), op.getSource().getType(), + cast(op.getResult().getType()), + getConstantIndexValue(op.getOffset()), VMIMemoryValidMaskKind::AllTrue); + if (!accessPlan.layoutSupport.isSupported()) + return fail(accessPlan.layoutSupport.reason); + if (!isa(op.getSource().getType())) + return fail("group_broadcast_load requires !pto.ptr source"); + return success(); +} + +LogicalResult +checkSupportedGroupStoreShape(VMIGroupStoreOp op, std::string *reason) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + auto valueType = cast(op.getValue().getType()); + VMILayoutAttr layout = valueType.getLayoutAttr(); + if (layout && layout.isGroupSlots()) { + VMILayoutSupport supports; + FailureOr fact = supports.getGroupStoreLayoutFact( + valueType, op.getNumGroupsAttr().getInt(), reason); + if (failed(fact)) + return failure(); + + VMIMemoryAccessPlan accessPlan = buildWriteAccessPlan(op.getDestination(), op.getDestination().getType(), + valueType, VMIMemoryWriteMaskKind::AllTrue); + if (!accessPlan.layoutSupport.isSupported()) + return fail(accessPlan.layoutSupport.reason); + + if (fact->slots == 1) { + unsigned elementBits = + pto::getPTOStorageElemBitWidth(valueType.getElementType()); + if (elementBits == 0 || 256 % elementBits != 0) + return fail("slots=1 group_store requires supported element width"); + std::optional rowStride = + getConstantIndexValue(op.getRowStride()); + if (rowStride && *rowStride <= 0) + return fail("slots=1 group_store requires positive row_stride when " + "row_stride is constant"); + if (!getPointStoreDistToken(valueType.getElementType())) + return fail("slots=1 group_store requires 1PT_B8/B16/B32 store " + "support"); + return success(); + } + + std::optional rowStride = getConstantIndexValue(op.getRowStride()); + if (!rowStride || *rowStride != 1) + return fail("slots=8 group_store currently requires constant unit " + "row_stride"); + return success(); + } + + FailureOr groupSize = getGroupSizeFromNumGroups( + valueType, op.getNumGroupsAttr().getInt(), reason); + if (failed(groupSize)) + return failure(); + if (failed(checkSupportedStoreShape(valueType, + op.getDestination(), + op.getDestination().getType(), reason))) + return failure(); + if (succeeded(checkSupportedGroupChunkShape(valueType, *groupSize, reason))) + return success(); + + int64_t lanesPerPart = 0; + int64_t groupCount = 0; + int64_t chunksPerGroupPerPart = 0; + return checkDeinterleaved2GroupStoreChunkShape( + valueType, *groupSize, &lanesPerPart, &groupCount, + &chunksPerGroupPerPart, reason); +} + +LogicalResult +checkSupportedMaskedLoadShape(VMIMaskedLoadOp op, std::string *reason) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + auto resultType = cast(op.getResult().getType()); + auto passthruType = cast(op.getPassthru().getType()); + auto maskType = cast(op.getMask().getType()); + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + VMILayoutAttr passthruLayout = passthruType.getLayoutAttr(); + VMILayoutAttr maskLayout = maskType.getLayoutAttr(); + VMIMemoryAccessPlan accessPlan = buildReadAccessPlan(op.getSource(), op.getSource().getType(), resultType, + getConstantIndexValue(op.getOffset()), + VMIMemoryValidMaskKind::ExplicitMask); + if (!accessPlan.layoutSupport.isSupported()) + return fail(accessPlan.layoutSupport.reason); + if (!resultLayout || !passthruLayout || !maskLayout) + return fail("requires assigned result, passthru, and mask layouts"); + if (!resultLayout.isContiguous() || !passthruLayout.isContiguous() || + !maskLayout.isContiguous()) + return fail("requires contiguous result, passthru, and mask layouts"); + + std::string fullChunkReason; + if (succeeded(checkFullDataPhysicalChunks(resultType, &fullChunkReason))) + return success(); + + if (accessPlan.safeReadProof.proven) + return success(); + requireUnavailableReadFallback(accessPlan); + return fail(Twine("partial/tail masked_load requires statically safe " + "full-read footprint; value ") + + fullChunkReason + ", safe-read proof " + + accessPlan.safeReadProof.reason + + "; fallback decision: " + accessPlan.fallbackDecision.reason); +} + +LogicalResult +checkSupportedGatherShape(VMIGatherOp op, std::string *reason) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + auto resultType = cast(op.getResult().getType()); + auto indicesType = cast(op.getIndices().getType()); + auto passthruType = cast(op.getPassthru().getType()); + auto maskType = cast(op.getMask().getType()); + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + VMILayoutAttr indicesLayout = indicesType.getLayoutAttr(); + VMILayoutAttr passthruLayout = passthruType.getLayoutAttr(); + VMILayoutAttr maskLayout = maskType.getLayoutAttr(); + if (!resultLayout || !indicesLayout || !passthruLayout || !maskLayout) + return fail("requires assigned result, indices, passthru, and mask " + "layouts"); + if (!resultLayout.isContiguous() || !indicesLayout.isContiguous() || + !passthruLayout.isContiguous() || !maskLayout.isContiguous()) + return fail("requires contiguous result, indices, passthru, and mask " + "layouts"); + + if (!isa(op.getSource().getType())) + return fail("requires !pto.ptr source because pto.vgather2_bc is " + "pointer-only"); + + unsigned resultBits = + pto::getPTOStorageElemBitWidth(resultType.getElementType()); + auto indexElementType = dyn_cast(indicesType.getElementType()); + if (!indexElementType || indexElementType.isSigned()) + return fail("requires signless or unsigned integer indices"); + bool isU16Gather = resultBits == 16 && indexElementType.isUnsigned() && + indexElementType.getWidth() == 16 && + maskType.getGranularity() == "b16"; + bool isB32Gather = resultBits == 32 && indexElementType.getWidth() == 32 && + maskType.getGranularity() == "b32"; + if (!isU16Gather && !isB32Gather) + return fail("requires either 32-bit results with 32-bit indices and b32 " + "mask, or ui16 results with ui16 indices and b16 mask"); + + FailureOr resultArity = getVMIPhysicalArity(resultType); + FailureOr indicesArity = getVMIPhysicalArity(indicesType); + FailureOr passthruArity = getVMIPhysicalArity(passthruType); + FailureOr maskArity = getVMIPhysicalArity(maskType); + if (failed(resultArity) || failed(indicesArity) || failed(passthruArity) || + failed(maskArity)) + return fail("requires computable physical arity"); + if (*resultArity != *indicesArity || *resultArity != *passthruArity || + *resultArity != *maskArity) + return fail("requires result, indices, passthru, and mask to have the " + "same physical arity"); + + if (isB32Gather) { + std::string resultReason; + std::string indicesReason; + std::string passthruReason; + std::string maskReason; + if (failed(checkFullDataPhysicalChunks(resultType, &resultReason))) + return fail(Twine("result requires full physical chunks; ") + + resultReason); + if (failed(checkFullDataPhysicalChunks(indicesType, &indicesReason))) + return fail(Twine("indices require full physical chunks; ") + + indicesReason); + if (failed(checkFullDataPhysicalChunks(passthruType, &passthruReason))) + return fail(Twine("passthru requires full physical chunks; ") + + passthruReason); + if (failed(checkFullVMIPhysicalChunks(maskType, &maskReason))) + return fail(Twine("mask requires full physical chunks; ") + maskReason); + } else if (*resultArity != 1) { + return fail("ui16 gather currently supports one physical chunk"); + } + + return success(); +} + +LogicalResult +checkSupportedScatterShape(VMIScatterOp op, std::string *reason) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + auto valueType = cast(op.getValue().getType()); + auto indicesType = cast(op.getIndices().getType()); + auto maskType = cast(op.getMask().getType()); + VMILayoutAttr valueLayout = valueType.getLayoutAttr(); + VMILayoutAttr indicesLayout = indicesType.getLayoutAttr(); + VMILayoutAttr maskLayout = maskType.getLayoutAttr(); + if (!valueLayout || !indicesLayout || !maskLayout) + return fail("requires assigned value, indices, and mask layouts"); + if (!valueLayout.isContiguous() || !indicesLayout.isContiguous() || + !maskLayout.isContiguous()) + return fail("requires contiguous value, indices, and mask layouts"); + + if (!isa(op.getDestination().getType())) + return fail("requires !pto.ptr destination because pto.vscatter is " + "pointer-only"); + + if (pto::getPTOStorageElemBitWidth(valueType.getElementType()) != 32) + return fail("currently requires 32-bit value element type so physical " + "index and value lane counts match pto.vscatter"); + auto indexElementType = dyn_cast(indicesType.getElementType()); + if (!indexElementType || indexElementType.getWidth() != 32 || + indexElementType.isSigned()) + return fail("requires signless or unsigned 32-bit indices"); + if (maskType.getGranularity() != "b32") + return fail("requires b32 mask granularity"); + + FailureOr valueArity = getVMIPhysicalArity(valueType); + FailureOr indicesArity = getVMIPhysicalArity(indicesType); + FailureOr maskArity = getVMIPhysicalArity(maskType); + if (failed(valueArity) || failed(indicesArity) || failed(maskArity)) + return fail("requires computable physical arity"); + if (*valueArity != *indicesArity || *valueArity != *maskArity) + return fail("requires value, indices, and mask to have the same physical " + "arity"); + + std::string valueReason; + std::string indicesReason; + std::string maskReason; + if (failed(checkFullDataPhysicalChunks(valueType, &valueReason))) + return fail(Twine("value requires full physical chunks; ") + valueReason); + if (failed(checkFullDataPhysicalChunks(indicesType, &indicesReason))) + return fail(Twine("indices require full physical chunks; ") + + indicesReason); + if (failed(checkFullVMIPhysicalChunks(maskType, &maskReason))) + return fail(Twine("mask requires full physical chunks; ") + maskReason); + + return success(); +} + +LogicalResult +checkSupportedStrideStoreShape(VMIStrideStoreOp op, std::string *reason) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + auto valueType = cast(op.getValue().getType()); + auto maskType = cast(op.getMask().getType()); + VMILayoutAttr valueLayout = valueType.getLayoutAttr(); + VMILayoutAttr maskLayout = maskType.getLayoutAttr(); + if (!valueLayout || !maskLayout) + return fail("requires assigned value and mask layouts"); + if (!valueLayout.isContiguous() || !maskLayout.isContiguous()) + return fail("requires contiguous value and mask layouts"); + + if (!isa(op.getDestination().getType())) + return fail("requires !pto.ptr destination because pto.vsstb is " + "pointer-only"); + if (failed(checkSupportedStoreShape(valueType, + op.getDestination(), + op.getDestination().getType(), reason))) + return failure(); + + FailureOr valueArity = getVMIPhysicalArity(valueType); + FailureOr maskArity = getVMIPhysicalArity(maskType); + if (failed(valueArity) || failed(maskArity)) + return fail("requires computable physical arity"); + if (*valueArity != 1 || *maskArity != 1) + return fail("currently supports one physical value/mask chunk"); + return success(); +} + +LogicalResult +checkSupportedStrideLoadShape(VMIStrideLoadOp op, std::string *reason) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + auto resultType = cast(op.getResult().getType()); + auto maskType = cast(op.getMask().getType()); + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + VMILayoutAttr maskLayout = maskType.getLayoutAttr(); + if (!resultLayout || !maskLayout) + return fail("requires assigned result and mask layouts"); + if (!resultLayout.isContiguous() || !maskLayout.isContiguous()) + return fail("requires contiguous result and mask layouts"); + + if (!isa(op.getSource().getType())) + return fail("requires !pto.ptr source because pto.vsldb is pointer-only"); + + FailureOr resultArity = getVMIPhysicalArity(resultType); + FailureOr maskArity = getVMIPhysicalArity(maskType); + if (failed(resultArity) || failed(maskArity)) + return fail("requires computable physical arity"); + if (*resultArity != 1 || *maskArity != 1) + return fail("currently supports one physical result/mask chunk"); + return success(); +} + +Value stripMaskMaterialization(Value value) { + while (true) { + if (auto ensure = value.getDefiningOp()) { + value = ensure.getSource(); + continue; + } + if (auto ensure = value.getDefiningOp()) { + value = ensure.getSource(); + continue; + } + return value; + } +} + +bool isStaticAllActiveMask(Value mask, int64_t expectedLanes, + std::string *reason = nullptr) { + mask = stripMaskMaterialization(mask); + auto fail = [&](const Twine &message) { + if (reason) + *reason = message.str(); + return false; + }; + + if (auto createMask = mask.getDefiningOp()) { + auto activeConstant = + createMask.getActiveLanes().getDefiningOp(); + if (!activeConstant) + return fail("create_mask active_lanes is dynamic"); + auto activeAttr = dyn_cast(activeConstant.getValue()); + if (!activeAttr) + return fail("create_mask active_lanes is not an integer constant"); + return activeAttr.getInt() >= expectedLanes + ? true + : fail("create_mask active_lanes is smaller than the logical " + "lane count"); + } + + if (auto constantMask = mask.getDefiningOp()) { + auto denseAttr = dyn_cast(constantMask.getValue()); + if (!denseAttr) + return fail("constant_mask is not a dense integer mask"); + if (denseAttr.getNumElements() != expectedLanes) + return fail("constant_mask element count does not match the logical " + "lane count"); + auto values = denseAttr.getValues(); + for (bool value : values) + if (!value) + return fail("constant_mask contains an inactive lane"); + return true; + } + + return fail("mask is not a static all-active create_mask or constant_mask"); +} + +LogicalResult +checkSupportedExpandLoadShape(VMIExpandLoadOp op, std::string *reason) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + auto resultType = cast(op.getResult().getType()); + auto passthruType = cast(op.getPassthru().getType()); + auto maskType = cast(op.getMask().getType()); + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + VMILayoutAttr passthruLayout = passthruType.getLayoutAttr(); + VMILayoutAttr maskLayout = maskType.getLayoutAttr(); + VMIMemoryAccessPlan accessPlan = buildReadAccessPlan(op.getSource(), op.getSource().getType(), resultType, + getConstantIndexValue(op.getOffset()), + VMIMemoryValidMaskKind::ExplicitMask); + if (!accessPlan.layoutSupport.isSupported()) + return fail(accessPlan.layoutSupport.reason); + if (!resultLayout || !passthruLayout || !maskLayout) + return fail("requires assigned result, passthru, and mask layouts"); + if (!resultLayout.isContiguous() || !passthruLayout.isContiguous() || + !maskLayout.isContiguous()) + return fail("requires contiguous result, passthru, and mask layouts"); + + std::string maskReason; + bool staticAllActive = isStaticAllActiveMask( + op.getMask(), resultType.getElementCount(), &maskReason); + + std::string fullChunkReason; + if (staticAllActive && + succeeded(checkFullDataPhysicalChunks(resultType, &fullChunkReason))) + return success(); + + if (staticAllActive && accessPlan.safeReadProof.proven) + return success(); + + std::string allActivePathReason; + if (!staticAllActive) { + allActivePathReason = + maskReason.empty() ? "requires static all-active mask" : maskReason; + } else { + requireUnavailableReadFallback(accessPlan); + allActivePathReason = + (Twine("requires full physical chunks or statically safe full-read " + "footprint; value ") + + fullChunkReason + ", safe-read proof " + + accessPlan.safeReadProof.reason + + "; fallback decision: " + accessPlan.fallbackDecision.reason) + .str(); + } + + if (!isa(op.getSource().getType())) + return fail(Twine("runtime-mask path requires !pto.ptr source because " + "pto.vgather2_bc is pointer-only; all-active path ") + + allActivePathReason); + if (pto::getPTOStorageElemBitWidth(resultType.getElementType()) != 32) + return fail("runtime-mask path currently requires 32-bit result element " + "type so prefix indices and gather result lane counts match"); + if (maskType.getGranularity() != "b32") + return fail("runtime-mask path requires b32 mask granularity"); + + FailureOr resultArity = getVMIPhysicalArity(resultType); + FailureOr passthruArity = getVMIPhysicalArity(passthruType); + FailureOr maskArity = getVMIPhysicalArity(maskType); + if (failed(resultArity) || failed(passthruArity) || failed(maskArity)) + return fail("runtime-mask path requires computable physical arity"); + if (*resultArity != 1 || *passthruArity != 1 || *maskArity != 1) + return fail("runtime-mask path currently supports only one physical " + "chunk because prefix indices must not reset across chunks"); + + std::string passthruReason; + std::string maskFullReason; + if (failed(checkFullDataPhysicalChunks(resultType, &fullChunkReason))) + return fail(Twine("runtime-mask result requires full physical chunks; ") + + fullChunkReason); + if (failed(checkFullDataPhysicalChunks(passthruType, &passthruReason))) + return fail(Twine("runtime-mask passthru requires full physical chunks; ") + + passthruReason); + if (failed(checkFullVMIPhysicalChunks(maskType, &maskFullReason))) + return fail(Twine("runtime-mask mask requires full physical chunks; ") + + maskFullReason); + + return success(); +} + +LogicalResult +checkSupportedMaskedStoreShape(VMIVRegType valueType, VMIMaskType maskType, + Value destination, Type destinationType, + std::string *reason) { + VMIMemoryAccessPlan accessPlan = + buildWriteAccessPlan(destination, destinationType, + valueType, VMIMemoryWriteMaskKind::ExplicitMask); + if (!accessPlan.layoutSupport.isSupported()) { + if (reason) + *reason = accessPlan.layoutSupport.reason; + return failure(); + } + + std::string valueReason; + std::string maskReason; + if (succeeded(checkFullDataPhysicalChunks(valueType, &valueReason)) && + succeeded(checkFullVMIPhysicalChunks(maskType, &maskReason))) + return success(); + + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + VMILayoutAttr valueLayout = valueType.getLayoutAttr(); + VMILayoutAttr maskLayout = maskType.getLayoutAttr(); + if (!valueLayout || !maskLayout) + return fail("requires assigned value and mask layouts"); + + FailureOr valueArity = getVMIPhysicalArity(valueType); + FailureOr maskArity = getVMIPhysicalArity(maskType); + if (failed(valueArity) || failed(maskArity) || *valueArity != *maskArity) + return fail("requires matching value/mask physical arity"); + + if (valueLayout.hasDenseLaneStride()) { + VMILayoutSupport supports; + if (succeeded( + supports.getMaskedStoreLayoutFact(valueType, maskType, reason))) + return success(); + } + + std::string valueMaterializationReason; + FailureOr valueParts = getContiguousMaterializationPartCount( + valueType, &valueMaterializationReason); + if (failed(valueParts)) + return fail(Twine("value cannot materialize to contiguous; value ") + + valueReason + ", materialization " + + valueMaterializationReason); + + std::string maskMaterializationReason; + FailureOr maskParts = getContiguousMaterializationPartCount( + maskType, &maskMaterializationReason); + if (failed(maskParts)) + return fail(Twine("mask cannot materialize to contiguous; mask ") + + maskReason + ", materialization " + maskMaterializationReason); + if (*valueParts != *maskParts) + return fail( + "requires value/mask contiguous materialization arity to match"); + return success(); +} + +FailureOr getContiguousActiveDataLanes(VMIVRegType vmiType, + int64_t chunk) { + FailureOr lanesPerPart = + getDataLanesPerPart(vmiType.getElementType()); + if (failed(lanesPerPart)) + return failure(); + + int64_t remaining = vmiType.getElementCount() - chunk * *lanesPerPart; + return std::clamp(remaining, 0, *lanesPerPart); +} + +FailureOr getActiveDataLanesInPhysicalChunk(VMIVRegType vmiType, + int64_t chunk) { + FailureOr lanesPerPart = + getDataLanesPerPart(vmiType.getElementType()); + if (failed(lanesPerPart)) + return failure(); + + int64_t active = 0; + for (int64_t lane = 0; lane < *lanesPerPart; ++lane) { + FailureOr padding = isPaddingLane(vmiType, /*part=*/0, chunk, lane); + if (failed(padding)) + return failure(); + if (!*padding) + ++active; + } + return active; +} + +FailureOr createContiguousStoreMask(Location loc, VMIVRegType vmiType, + int64_t chunk, VRegType vregType, + PatternRewriter &rewriter) { + FailureOr lanesPerPart = + getDataLanesPerPart(vmiType.getElementType()); + if (failed(lanesPerPart)) + return failure(); + + FailureOr activeLanes = getContiguousActiveDataLanes(vmiType, chunk); + if (failed(activeLanes)) + return failure(); + if (*activeLanes == *lanesPerPart) + return createAllTrueMaskForVReg(loc, vregType, rewriter); + + FailureOr maskType = + getMaskTypeForVReg(vregType, rewriter.getContext()); + if (failed(maskType)) + return failure(); + FailureOr> maskAndRemaining = createRuntimePrefixMask( + loc, *maskType, createI32Constant(loc, *activeLanes, rewriter), rewriter); + if (failed(maskAndRemaining)) + return failure(); + return maskAndRemaining->first; +} + +FailureOr createMaskedStorePredicate(Location loc, VMIVRegType vmiType, + int64_t chunk, Value userMask, + VRegType vregType, + PatternRewriter &rewriter) { + FailureOr lanesPerPart = + getDataLanesPerPart(vmiType.getElementType()); + if (failed(lanesPerPart)) + return failure(); + + FailureOr activeLanes = getContiguousActiveDataLanes(vmiType, chunk); + if (failed(activeLanes)) + return failure(); + if (*activeLanes == *lanesPerPart) + return userMask; + + auto maskType = dyn_cast(userMask.getType()); + if (!maskType) + return failure(); + FailureOr tailMask = + createContiguousStoreMask(loc, vmiType, chunk, vregType, rewriter); + FailureOr allTrue = createAllTrueMask(loc, maskType, rewriter); + if (failed(tailMask) || failed(allTrue)) + return failure(); + return rewriter.create(loc, maskType, userMask, *tailMask, *allTrue) + .getResult(); +} + +FailureOr createDenseLaneStrideStorePredicate( + Location loc, VMIVRegType vmiType, int64_t chunk, Value userMask, + StringRef targetGranularity, PatternRewriter &rewriter) { + auto sourceMaskType = dyn_cast(userMask.getType()); + if (!sourceMaskType) + return failure(); + auto targetMaskType = MaskType::get(rewriter.getContext(), targetGranularity); + Value compactMask = userMask; + VMILayoutAttr layout = vmiType.getLayoutAttr(); + if (!layout) + return failure(); + + auto lower = rewriter.getStringAttr("LOWER"); + StringRef sourceGranularity = sourceMaskType.getGranularity(); + if (sourceGranularity == targetGranularity) { + compactMask = userMask; + } else if (layout.getLaneStride() == 2) { + compactMask = + rewriter.create(loc, targetMaskType, compactMask, lower) + .getResult(); + } else if (layout.getLaneStride() == 4 && sourceGranularity == "b8" && + targetGranularity == "b32") { + auto b16MaskType = MaskType::get(rewriter.getContext(), "b16"); + compactMask = + rewriter.create(loc, b16MaskType, compactMask, lower) + .getResult(); + compactMask = + rewriter.create(loc, targetMaskType, compactMask, lower) + .getResult(); + } else { + return failure(); + } + + FailureOr activeLanes = + getActiveDataLanesInPhysicalChunk(vmiType, chunk); + FailureOr maskLanes = getMaskLanesPerPart(targetGranularity); + if (failed(activeLanes) || failed(maskLanes)) + return failure(); + if (*activeLanes == *maskLanes) + return compactMask; + + FailureOr tailMask = createPrefixMaskForActiveLanes( + loc, targetMaskType, *activeLanes, rewriter); + FailureOr allTrue = createAllTrueMask(loc, targetMaskType, rewriter); + if (failed(tailMask) || failed(allTrue)) + return failure(); + return rewriter + .create(loc, targetMaskType, compactMask, *tailMask, *allTrue) + .getResult(); +} + +FailureOr> +computeShuffleForwardingSourceParts(VMIShuffleOp op, std::string *reason) { + auto fail = [&](const Twine &message) -> FailureOr> { + if (reason) + *reason = message.str(); + return failure(); + }; + + auto sourceType = cast(op.getSource().getType()); + auto resultType = cast(op.getResult().getType()); + FailureOr lanesPerPart = + getDataLanesPerPart(sourceType.getElementType()); + if (failed(lanesPerPart)) + return fail("requires known lanes per physical part"); + + ArrayRef indices = op.getIndices(); + if (indices.empty()) + return fail("requires non-empty indices"); + + FailureOr resultFactor = getDataLayoutFactor(resultType); + if (failed(resultFactor)) + return fail("requires assigned result layout"); + + SmallVector sourceFlatIndices; + for (int64_t resultPart = 0; resultPart < *resultFactor; ++resultPart) { + FailureOr resultChunks = + getDataChunksInPart(resultType, resultPart); + if (failed(resultChunks)) + return fail("requires known result physical chunks"); + + for (int64_t resultChunk = 0; resultChunk < *resultChunks; ++resultChunk) { + std::optional sourcePart; + std::optional sourceChunk; + for (int64_t lane = 0; lane < *lanesPerPart; ++lane) { + FailureOr padding = + isPaddingLane(resultType, resultPart, resultChunk, lane); + if (failed(padding)) + return fail("failed to classify result padding lanes"); + if (*padding) + continue; + + FailureOr resultLogicalLane = + mapPhysicalLaneToLogical(resultType, resultPart, resultChunk, lane); + if (failed(resultLogicalLane) || + *resultLogicalLane >= static_cast(indices.size())) + return fail("failed to map result lane"); + + FailureOr sourcePhysical = + mapLogicalLaneToPhysical(sourceType, indices[*resultLogicalLane]); + if (failed(sourcePhysical)) + return fail("failed to map source lane"); + if (sourcePhysical->lane != lane) + return fail("requires same-lane physical chunks"); + + if (!sourcePart) { + sourcePart = sourcePhysical->part; + sourceChunk = sourcePhysical->chunk; + continue; + } + if (*sourcePart != sourcePhysical->part || + *sourceChunk != sourcePhysical->chunk) + return fail("requires one source chunk per result chunk"); + } + + if (!sourcePart || !sourceChunk) + return fail("requires at least one logical lane per result chunk"); + FailureOr sourceFlatIndex = + getDataFlatPartIndex(sourceType, *sourcePart, *sourceChunk); + if (failed(sourceFlatIndex)) + return fail("source part range is out of bounds"); + sourceFlatIndices.push_back(*sourceFlatIndex); + } + } + + return sourceFlatIndices; +} + +struct ShuffleVselrPlan { + int64_t sourceFlatIndex = 0; + int64_t baseLane = 0; + bool descending = false; +}; + +FailureOr computeShuffleLane0SplatSourcePart(VMIShuffleOp op, + std::string *reason) { + auto fail = [&](const Twine &message) -> FailureOr { + if (reason) + *reason = message.str(); + return failure(); + }; + + ArrayRef indices = op.getIndices(); + if (indices.empty()) + return fail("requires non-empty indices"); + if (!llvm::all_of(indices, [](int64_t index) { return index == 0; })) + return fail("requires every result lane to select source lane 0"); + + auto sourceType = cast(op.getSource().getType()); + FailureOr sourceLane = + mapLogicalLaneToPhysical(sourceType, 0); + if (failed(sourceLane)) + return fail("failed to map source lane 0"); + FailureOr sourceFlatIndex = + getDataFlatPartIndex(sourceType, sourceLane->part, sourceLane->chunk); + if (failed(sourceFlatIndex)) + return fail("source lane 0 part range is out of bounds"); + return *sourceFlatIndex; +} + +FailureOr> +computeShuffleVselrPlans(VMIShuffleOp op, std::string *reason) { + auto fail = + [&](const Twine &message) -> FailureOr> { + if (reason) + *reason = message.str(); + return failure(); + }; + + auto sourceType = cast(op.getSource().getType()); + auto resultType = cast(op.getResult().getType()); + FailureOr lanesPerPart = + getDataLanesPerPart(sourceType.getElementType()); + if (failed(lanesPerPart)) + return fail("requires known lanes per physical part"); + + ArrayRef indices = op.getIndices(); + if (indices.empty()) + return fail("requires non-empty indices"); + + FailureOr resultFactor = getDataLayoutFactor(resultType); + if (failed(resultFactor)) + return fail("requires assigned result layout"); + + SmallVector plans; + for (int64_t resultPart = 0; resultPart < *resultFactor; ++resultPart) { + FailureOr resultChunks = + getDataChunksInPart(resultType, resultPart); + if (failed(resultChunks)) + return fail("requires known result physical chunks"); + + for (int64_t resultChunk = 0; resultChunk < *resultChunks; ++resultChunk) { + std::optional sourcePart; + std::optional sourceChunk; + std::optional baseLane; + std::optional descending; + for (int64_t lane = 0; lane < *lanesPerPart; ++lane) { + FailureOr padding = + isPaddingLane(resultType, resultPart, resultChunk, lane); + if (failed(padding) || *padding) + return fail("requires full physical result chunks"); + + FailureOr resultLogicalLane = + mapPhysicalLaneToLogical(resultType, resultPart, resultChunk, lane); + if (failed(resultLogicalLane) || + *resultLogicalLane >= static_cast(indices.size())) + return fail("failed to map result lane"); + + FailureOr sourcePhysical = + mapLogicalLaneToPhysical(sourceType, indices[*resultLogicalLane]); + if (failed(sourcePhysical)) + return fail("failed to map source lane"); + + if (!sourcePart) { + sourcePart = sourcePhysical->part; + sourceChunk = sourcePhysical->chunk; + baseLane = sourcePhysical->lane; + continue; + } + + if (*sourcePart != sourcePhysical->part || + *sourceChunk != sourcePhysical->chunk) + return fail("requires one source chunk per result chunk"); + + int64_t ascExpected = *baseLane + lane; + int64_t descExpected = *baseLane - lane; + bool asc = sourcePhysical->lane == ascExpected; + bool desc = sourcePhysical->lane == descExpected; + if (!asc && !desc) + return fail("requires ASC or DESC affine source lane indices"); + + bool laneDescending = desc && !asc; + if (!descending) { + descending = laneDescending; + continue; + } + if (*descending != laneDescending) + return fail("requires one index order per result chunk"); + } + + FailureOr sourceFlatIndex = + getDataFlatPartIndex(sourceType, *sourcePart, *sourceChunk); + if (failed(sourceFlatIndex)) + return fail("source part range is out of bounds"); + plans.push_back(ShuffleVselrPlan{*sourceFlatIndex, *baseLane, + descending.value_or(false)}); + } + } + + return plans; +} + +struct ConstantMaskChunkMaterialization { + SmallVector activeLanes; +}; + +FailureOr> +computeConstantMaskMaterialization(VMIConstantMaskOp op, std::string *reason) { + auto fail = [&](const Twine &message) + -> FailureOr> { + if (reason) + *reason = message.str(); + return failure(); + }; + + auto denseAttr = dyn_cast(op.getValue()); + if (!denseAttr) + return fail("only dense integer mask constants are supported"); + + auto resultVMIType = cast(op.getResult().getType()); + VMILayoutAttr layout = resultVMIType.getLayoutAttr(); + if (!layout || + !VMIMaskType::isConcreteGranularity(resultVMIType.getGranularity())) + return fail("requires concrete layout and granularity"); + + FailureOr physicalGranularity = + getVMIMaskPhysicalGranularity(resultVMIType); + FailureOr lanesPerPart = + failed(physicalGranularity) + ? FailureOr(failure()) + : getMaskLanesPerPart(*physicalGranularity); + if (failed(lanesPerPart)) + return fail("requires known physical mask lanes per part"); + + auto boolValues = denseAttr.getValues(); + int64_t factor = layout.isDeinterleaved() ? layout.getFactor() : 1; + SmallVector materializations; + for (int64_t part = 0; part < factor; ++part) { + for (int64_t chunk = 0;; ++chunk) { + bool anyLane = false; + ConstantMaskChunkMaterialization materialization; + materialization.activeLanes.reserve(*lanesPerPart); + for (int64_t lane = 0; lane < *lanesPerPart; ++lane) { + FailureOr padding = + isPaddingLane(resultVMIType, part, chunk, lane); + if (failed(padding)) + return fail("failed to map physical padding lane"); + if (*padding) { + materialization.activeLanes.push_back(0); + continue; + } + anyLane = true; + + FailureOr logicalLane = + mapPhysicalLaneToLogical(resultVMIType, part, chunk, lane); + if (failed(logicalLane)) + return fail("failed to map physical lane"); + materialization.activeLanes.push_back(boolValues[*logicalLane] ? 1 : 0); + } + if (!anyLane) + break; + materializations.push_back(std::move(materialization)); + } + } + + return materializations; +} + +FailureOr> +computeGroupMaskMaterializationForType(VMICreateGroupMaskOp op, + VMIMaskType resultVMIType, + std::string *reason) { + auto fail = [&](const Twine &message) + -> FailureOr> { + if (reason) + *reason = message.str(); + return failure(); + }; + + auto activeConstant = + op.getActiveElemsPerGroup().getDefiningOp(); + if (!activeConstant) + return fail("requires constant active_elems_per_group"); + auto activeAttr = dyn_cast(activeConstant.getValue()); + if (!activeAttr) + return fail("active_elems_per_group must be an integer constant"); + + VMILayoutAttr layout = resultVMIType.getLayoutAttr(); + if (!layout || + !VMIMaskType::isConcreteGranularity(resultVMIType.getGranularity())) + return fail("requires concrete layout and granularity"); + + FailureOr physicalGranularity = + getVMIMaskPhysicalGranularity(resultVMIType); + FailureOr lanesPerPart = + failed(physicalGranularity) + ? FailureOr(failure()) + : getMaskLanesPerPart(*physicalGranularity); + if (failed(lanesPerPart)) + return fail("requires known physical mask lanes per part"); + + int64_t numGroups = op.getNumGroupsAttr().getInt(); + int64_t groupSize = op.getGroupSizeAttr().getInt(); + if (numGroups <= 0 || groupSize <= 0 || + resultVMIType.getElementCount() != numGroups * groupSize) + return fail("requires result lane count to match num_groups * group_size"); + + int64_t activeElems = activeAttr.getInt(); + if (activeElems < 0) + activeElems = 0; + if (activeElems > groupSize) + activeElems = groupSize; + + int64_t factor = layout.isDeinterleaved() ? layout.getFactor() : 1; + SmallVector materializations; + for (int64_t part = 0; part < factor; ++part) { + for (int64_t chunk = 0;; ++chunk) { + bool anyLane = false; + ConstantMaskChunkMaterialization materialization; + materialization.activeLanes.reserve(*lanesPerPart); + for (int64_t lane = 0; lane < *lanesPerPart; ++lane) { + FailureOr padding = + isPaddingLane(resultVMIType, part, chunk, lane); + if (failed(padding)) + return fail("failed to map physical padding lane"); + if (*padding) { + materialization.activeLanes.push_back(0); + continue; + } + anyLane = true; + + FailureOr logicalLane = + mapPhysicalLaneToLogical(resultVMIType, part, chunk, lane); + if (failed(logicalLane)) + return fail("failed to map physical lane"); + int64_t laneInGroup = *logicalLane % groupSize; + materialization.activeLanes.push_back(laneInGroup < activeElems ? 1 + : 0); + } + if (!anyLane) + break; + materializations.push_back(std::move(materialization)); + } + } + + return materializations; +} + +FailureOr> +computeGroupMaskMaterialization(VMICreateGroupMaskOp op, std::string *reason) { + return computeGroupMaskMaterializationForType( + op, cast(op.getResult().getType()), reason); +} + +FailureOr materializeConstantMaskChunk(Location loc, MaskType maskType, + ArrayRef activeLanes, + PatternRewriter &rewriter); + +FailureOr createPowerOfTwoRemainder(Location loc, Value value, + int64_t modulus, Value allMask, + PatternRewriter &rewriter) { + if (modulus <= 0) + return failure(); + + auto vectorType = dyn_cast(value.getType()); + if (!vectorType) + return failure(); + + std::optional shift = getPowerOfTwoLog2(modulus); + if (!shift) + return failure(); + if (*shift == 0) { + Value zero = createI32Constant(loc, 0, rewriter); + return rewriter.create(loc, vectorType, zero, allMask, + /*position=*/nullptr) + .getResult(); + } + + Value shiftScalar = createI16Constant(loc, *shift, rewriter); + Value quotient = + rewriter.create(loc, vectorType, value, shiftScalar, allMask) + .getResult(); + Value base = + rewriter.create(loc, vectorType, quotient, shiftScalar, allMask) + .getResult(); + return rewriter.create(loc, vectorType, value, base, allMask) + .getResult(); +} + +FailureOr> materializeDynamicGroupMaskForType( + VMICreateGroupMaskOp op, Value activeElemsPerGroup, + VMIMaskType resultVMIType, TypeRange resultTypes, + PatternRewriter &rewriter) { + auto fail = [&](const Twine &message) -> FailureOr> { + (void)rewriter.notifyMatchFailure(op, message); + return failure(); + }; + + VMILayoutAttr layout = resultVMIType.getLayoutAttr(); + if (!layout) + return fail("dynamic create_group_mask requires assigned layout"); + if (layout.getLaneStride() != 1) + return fail("dynamic create_group_mask requires lane_stride=1 layout"); + if (resultVMIType.getGranularity() != "b32") + return fail("dynamic create_group_mask currently requires b32 " + "granularity"); + + int64_t numGroups = op.getNumGroupsAttr().getInt(); + int64_t groupSize = op.getGroupSizeAttr().getInt(); + if (numGroups <= 0 || groupSize <= 0 || + resultVMIType.getElementCount() != numGroups * groupSize) + return fail("dynamic create_group_mask requires result lane count to " + "match num_groups * group_size"); + + FailureOr physicalGranularity = + getVMIMaskPhysicalGranularity(resultVMIType); + FailureOr lanesPerPart = + failed(physicalGranularity) + ? FailureOr(failure()) + : getMaskLanesPerPart(*physicalGranularity); + FailureOr arity = getVMIPhysicalArity(resultVMIType); + if (failed(lanesPerPart) || failed(arity) || *arity < 1) + return fail("dynamic create_group_mask requires computable physical " + "mask chunks"); + if (static_cast(resultTypes.size()) != *arity) + return fail("dynamic create_group_mask physical result count mismatch"); + + std::optional groupShift = getPowerOfTwoLog2(groupSize); + if (!groupShift) + return fail("dynamic create_group_mask currently requires power-of-two " + "group_size"); + + int64_t factor = layout.isDeinterleaved() ? layout.getFactor() : 1; + int64_t blockElems = layout.isDeinterleaved() ? layout.getBlockElems() : 1; + if (factor <= 0 || blockElems <= 0 || + static_cast(resultTypes.size()) % factor != 0) + return fail("dynamic create_group_mask physical result count does not " + "match layout factor"); + if (!getPowerOfTwoLog2(blockElems)) + return fail("dynamic create_group_mask requires power-of-two block_elems"); + + Location loc = op.getLoc(); + MLIRContext *ctx = rewriter.getContext(); + Type i32 = rewriter.getI32Type(); + auto indexVectorType = VRegType::get(ctx, *lanesPerPart, i32); + Value activeI32 = + clampDynamicActiveLanes(loc, activeElemsPerGroup, groupSize, rewriter); + + SmallVector results; + results.reserve(resultTypes.size()); + int64_t chunksPerPart = resultTypes.size() / factor; + for (int64_t part = 0; part < factor; ++part) { + for (int64_t chunk = 0; chunk < chunksPerPart; ++chunk) { + Type resultType = resultTypes[part * chunksPerPart + chunk]; + auto maskType = dyn_cast(resultType); + if (!maskType || !maskType.isB32()) + return fail("dynamic create_group_mask result must be b32 mask"); + + FailureOr allMask = createAllTrueMask(loc, maskType, rewriter); + if (failed(allMask)) + return fail("failed to create dynamic create_group_mask all mask"); + + Value chunkBase = createI32Constant(loc, chunk * *lanesPerPart, rewriter); + Value indexInPart = + rewriter.create(loc, indexVectorType, chunkBase, StringAttr{}) + .getResult(); + + Value partBlock = indexInPart; + Value inBlockLane = createI32Constant(loc, 0, rewriter); + if (blockElems != 1) { + Value blockShift = createI16Constant(loc, *getPowerOfTwoLog2(blockElems), + rewriter); + partBlock = + rewriter + .create(loc, indexVectorType, indexInPart, blockShift, + *allMask) + .getResult(); + Value blockBase = + rewriter + .create(loc, indexVectorType, partBlock, blockShift, + *allMask) + .getResult(); + inBlockLane = + rewriter + .create(loc, indexVectorType, indexInPart, blockBase, + *allMask) + .getResult(); + } + + Value factorScalar = createI32Constant(loc, factor, rewriter); + Value logicalBlock = + rewriter + .create(loc, indexVectorType, partBlock, factorScalar, + *allMask) + .getResult(); + if (part != 0) { + Value partScalar = createI32Constant(loc, part, rewriter); + logicalBlock = + rewriter + .create(loc, indexVectorType, logicalBlock, + partScalar, *allMask) + .getResult(); + } + + Value logicalLane = logicalBlock; + if (blockElems != 1) { + Value blockShift = createI16Constant(loc, *getPowerOfTwoLog2(blockElems), + rewriter); + Value logicalBlockBase = + rewriter + .create(loc, indexVectorType, logicalBlock, + blockShift, *allMask) + .getResult(); + logicalLane = + rewriter + .create(loc, indexVectorType, logicalBlockBase, + inBlockLane, *allMask) + .getResult(); + } + + FailureOr laneInGroup = createPowerOfTwoRemainder( + loc, logicalLane, groupSize, *allMask, rewriter); + if (failed(laneInGroup)) + return fail("failed to compute dynamic create_group_mask lane index"); + + Value predicate = + rewriter + .create(loc, maskType, *laneInGroup, activeI32, + *allMask, rewriter.getStringAttr("lt")) + .getResult(); + + SmallVector validLanes; + validLanes.reserve(*lanesPerPart); + bool hasPadding = false; + for (int64_t lane = 0; lane < *lanesPerPart; ++lane) { + FailureOr padding = + isPaddingLane(resultVMIType, part, chunk, lane); + if (failed(padding)) + return fail("failed to classify dynamic create_group_mask padding"); + validLanes.push_back(*padding ? 0 : 1); + hasPadding |= *padding; + } + if (hasPadding) { + FailureOr validMask = + materializeConstantMaskChunk(loc, maskType, validLanes, rewriter); + if (failed(validMask)) + return fail("failed to materialize dynamic create_group_mask padding " + "mask"); + predicate = + rewriter + .create(loc, maskType, predicate, *validMask, *allMask) + .getResult(); + } + + results.push_back(predicate); + } + } + + return results; +} + +std::optional getPrefixActiveLaneCount(ArrayRef activeLanes) { + bool seenInactive = false; + int64_t activeCount = 0; + for (int8_t active : activeLanes) { + if (active) { + if (seenInactive) + return std::nullopt; + ++activeCount; + continue; + } + seenInactive = true; + } + return activeCount; +} + +FailureOr materializePrefixMask(Location loc, MaskType maskType, + int64_t activeLanes, + int64_t lanesPerPart, + PatternRewriter &rewriter) { + std::optional pattern = + getPrefixPattern(activeLanes, lanesPerPart); + if (pattern) + return createPatternMask(loc, maskType, *pattern, rewriter); + + FailureOr> maskAndRemaining = createRuntimePrefixMask( + loc, maskType, createI32Constant(loc, activeLanes, rewriter), rewriter); + if (failed(maskAndRemaining)) + return failure(); + return maskAndRemaining->first; +} + +FailureOr materializeConstantMaskChunk(Location loc, MaskType maskType, + ArrayRef activeLanes, + PatternRewriter &rewriter) { + FailureOr lanesPerPart = + getMaskLanesPerPart(maskType.getGranularity()); + if (failed(lanesPerPart) || + static_cast(activeLanes.size()) != *lanesPerPart) + return failure(); + + if (std::optional prefixCount = + getPrefixActiveLaneCount(activeLanes)) + return materializePrefixMask(loc, maskType, *prefixCount, *lanesPerPart, + rewriter); + + FailureOr allTrue = createAllTrueMask(loc, maskType, rewriter); + if (failed(allTrue)) + return failure(); + + Value result; + int64_t lane = 0; + while (lane < *lanesPerPart) { + while (lane < *lanesPerPart && !activeLanes[lane]) + ++lane; + if (lane >= *lanesPerPart) + break; + + int64_t runBegin = lane; + while (lane < *lanesPerPart && activeLanes[lane]) + ++lane; + int64_t runEnd = lane; + + FailureOr prefixEnd = + materializePrefixMask(loc, maskType, runEnd, *lanesPerPart, rewriter); + if (failed(prefixEnd)) + return failure(); + + Value runMask = *prefixEnd; + if (runBegin != 0) { + FailureOr prefixBegin = materializePrefixMask( + loc, maskType, runBegin, *lanesPerPart, rewriter); + if (failed(prefixBegin)) + return failure(); + Value notPrefixBegin = + rewriter.create(loc, maskType, *prefixBegin, *allTrue) + .getResult(); + runMask = rewriter + .create(loc, maskType, *prefixEnd, notPrefixBegin, + *allTrue) + .getResult(); + } + + if (!result) { + result = runMask; + continue; + } + result = rewriter.create(loc, maskType, result, runMask, *allTrue) + .getResult(); + } + + if (result) + return result; + return materializePrefixMask(loc, maskType, 0, *lanesPerPart, rewriter); +} + +FailureOr createScalarOffsetConstant(Location loc, Type type, + int64_t value, + PatternRewriter &rewriter); + +Value createChunkOffset(Location loc, Value baseOffset, int64_t laneOffset, + PatternRewriter &rewriter) { + if (laneOffset == 0) + return baseOffset; + Value delta = rewriter.create(loc, laneOffset); + return rewriter.create(loc, baseOffset, delta).getResult(); +} + +Value createGroupChunkOffset(Location loc, Value baseOffset, Value rowStride, + int64_t group, int64_t inGroupLaneOffset, + PatternRewriter &rewriter) { + Value offset = baseOffset; + if (group != 0) { + Value groupIndex = rewriter.create(loc, group); + Value rowOffset = + rewriter.create(loc, rowStride, groupIndex).getResult(); + offset = rewriter.create(loc, offset, rowOffset).getResult(); + } + return createChunkOffset(loc, offset, inGroupLaneOffset, rewriter); +} + +LogicalResult checkContiguousFullGroupChunks( + Operation *op, VMIVRegType type, int64_t groupSize, int64_t *lanesPerPart, + int64_t *groupCount, int64_t *chunksPerGroup, PatternRewriter &rewriter) { + auto fail = [&](const Twine &message) { + return rewriter.notifyMatchFailure(op, message); + }; + + VMILayoutAttr layout = type.getLayoutAttr(); + if (!layout || !layout.isContiguous()) + return fail("group op requires contiguous VMI layout"); + if (failed(checkFullDataPhysicalChunks(type, nullptr))) + return fail("group op requires full physical chunks"); + FailureOr lanes = getDataLanesPerPart(type.getElementType()); + if (failed(lanes)) + return fail("group op requires known physical lanes per part"); + if (groupSize <= 0 || type.getElementCount() % groupSize != 0) + return fail("group op requires derived group size to evenly divide lane " + "count"); + if (groupSize % *lanes != 0) + return fail("group op currently requires group size to be a multiple of " + "physical lanes per part"); + + *lanesPerPart = *lanes; + *groupCount = type.getElementCount() / groupSize; + *chunksPerGroup = groupSize / *lanes; + return success(); +} + +LogicalResult checkFullGroupSlotSourceShape( + Operation *op, VMIVRegType type, int64_t groupSize, int64_t numGroups, + int64_t *lanesPerPart, int64_t *groupCount, PatternRewriter &rewriter) { + auto fail = [&](const Twine &message) { + return rewriter.notifyMatchFailure(op, message); + }; + + VMILayoutAttr layout = type.getLayoutAttr(); + if (!layout || !layout.isGroupSlots() || layout.getNumGroups() != numGroups) + return fail("group slot op requires matching num_groups VMI layout"); + if (type.getElementCount() != numGroups) + return fail("group slot op requires one logical lane per group"); + FailureOr lanes = getDataLanesPerPart(type.getElementType()); + if (failed(lanes)) + return fail("group slot op requires known physical lanes per part"); + if (groupSize <= 0) + return fail("group slot op requires positive derived group size"); + if (*lanes % groupSize != 0 && groupSize % *lanes != 0) + return fail("group slot op requires group size to divide or be a " + "multiple of physical lanes per part"); + + *lanesPerPart = *lanes; + *groupCount = numGroups; + return success(); +} + +LogicalResult checkFullGroupBroadcastResultShape( + Operation *op, VMIVRegType type, int64_t groupSize, int64_t lanesPerPart, + int64_t *layoutFactor, int64_t *groupCount, PatternRewriter &rewriter) { + auto fail = [&](const Twine &message) { + return rewriter.notifyMatchFailure(op, message); + }; + + VMILayoutAttr layout = type.getLayoutAttr(); + if (!layout) + return fail("group_broadcast result requires assigned VMI layout"); + if (layout.isGroupSlots()) + return fail("group_broadcast result requires a dense VMI layout"); + bool laneStridedDense = layout.isDense() && layout.getLaneStride() > 1; + if (!laneStridedDense && failed(checkFullDataPhysicalChunks(type, nullptr))) + return fail("group_broadcast result requires full physical chunks"); + FailureOr resultLanes = getDataLanesPerPart(type.getElementType()); + if (failed(resultLanes) || *resultLanes != lanesPerPart) + return fail("group_broadcast result requires matching physical lanes"); + if (groupSize <= 0 || type.getElementCount() % groupSize != 0) + return fail("group_broadcast result requires derived group size to evenly " + "divide lane count"); + FailureOr factor = getDataLayoutFactor(type); + if (failed(factor)) + return fail("group_broadcast result requires known layout factor"); + + if (*factor == 1) { + if (lanesPerPart % groupSize != 0 && groupSize % lanesPerPart != 0) + return fail("group_broadcast contiguous result requires group size to " + "divide or be a multiple of physical lanes per part"); + } else { + bool blockFragmentSmallGroup = + layout.isDeinterleaved() && layout.getBlockElems() > 1 && + groupSize < lanesPerPart && lanesPerPart % layout.getBlockElems() == 0; + bool deinterleavedSmallGroup = + layout.isDeinterleaved() && layout.getBlockElems() == 1 && + groupSize < lanesPerPart && groupSize >= *factor && + groupSize % *factor == 0 && lanesPerPart % (groupSize / *factor) == 0; + int64_t logicalSpanPerResultChunk = lanesPerPart * *factor; + if (!blockFragmentSmallGroup && !deinterleavedSmallGroup && + (groupSize < lanesPerPart || + groupSize % logicalSpanPerResultChunk != 0)) + return fail("group_broadcast deinterleaved result requires every " + "physical result chunk to stay within one logical group"); + } + + *layoutFactor = *factor; + *groupCount = type.getElementCount() / groupSize; + return success(); +} + +FailureOr createZeroVector(Location loc, VRegType type, + PatternRewriter &rewriter) { + FailureOr zero = + createScalarOffsetConstant(loc, type.getElementType(), 0, rewriter); + FailureOr mask = createAllTrueMaskForVReg(loc, type, rewriter); + if (failed(zero) || failed(mask)) + return failure(); + return rewriter + .create(loc, type, *zero, *mask, + /*position=*/nullptr) + .getResult(); +} + +FailureOr createLaneRangeMask(Location loc, MaskType maskType, + int64_t begin, int64_t end, + PatternRewriter &rewriter) { + FailureOr lanesPerPart = + getMaskLanesPerPart(maskType.getGranularity()); + if (failed(lanesPerPart) || begin < 0 || begin > end || end > *lanesPerPart) + return failure(); + SmallVector active(*lanesPerPart, 0); + for (int64_t lane = begin; lane < end; ++lane) + active[lane] = 1; + return materializeConstantMaskChunk(loc, maskType, active, rewriter); +} + +FailureOr createGroupSlotIndexVector(Location loc, VRegType indexType, + int64_t groupSize, + int64_t baseGroupSlot, + PatternRewriter &rewriter, + int64_t slotLaneStride = 1) { + int64_t lanesPerPart = indexType.getElementCount(); + FailureOr baseScalar = createScalarOffsetConstant( + loc, indexType.getElementType(), baseGroupSlot * slotLaneStride, + rewriter); + FailureOr maskType = + getMaskTypeForVReg(indexType, rewriter.getContext()); + FailureOr allMask = createAllTrueMaskForVReg(loc, indexType, rewriter); + if (failed(baseScalar) || failed(maskType) || failed(allMask)) + return failure(); + Value result = rewriter + .create(loc, indexType, *baseScalar, *allMask, + /*position=*/nullptr) + .getResult(); + if (groupSize >= lanesPerPart) + return result; + if (lanesPerPart % groupSize != 0) + return failure(); + + int64_t groupsPerChunk = lanesPerPart / groupSize; + for (int64_t localGroup = 1; localGroup < groupsPerChunk; ++localGroup) { + FailureOr groupScalar = createScalarOffsetConstant( + loc, indexType.getElementType(), + (baseGroupSlot + localGroup) * slotLaneStride, rewriter); + FailureOr laneMask = + createLaneRangeMask(loc, *maskType, localGroup * groupSize, + (localGroup + 1) * groupSize, rewriter); + if (failed(groupScalar) || failed(laneMask)) + return failure(); + Value splat = rewriter + .create(loc, indexType, *groupScalar, *allMask, + /*position=*/nullptr) + .getResult(); + result = rewriter.create(loc, indexType, splat, result, *laneMask) + .getResult(); + } + return result; +} + +FailureOr createMappedGroupSlotIndexVector( + Location loc, VMIVRegType resultVMIType, int64_t part, int64_t chunk, + VRegType indexType, int64_t groupSize, int64_t slots, int64_t &sourceChunk, + PatternRewriter &rewriter, int64_t slotLaneStride = 1) { + if (groupSize <= 0 || slots <= 0) + return failure(); + + int64_t lanesPerPart = indexType.getElementCount(); + SmallVector slotByLane; + slotByLane.reserve(lanesPerPart); + std::optional resolvedSourceChunk; + for (int64_t lane = 0; lane < lanesPerPart; ++lane) { + FailureOr padding = + isPaddingLane(resultVMIType, part, chunk, lane); + if (failed(padding)) + return failure(); + if (*padding) { + slotByLane.push_back(0); + continue; + } + FailureOr logicalLane = + mapPhysicalLaneToLogical(resultVMIType, part, chunk, lane); + if (failed(logicalLane)) + return failure(); + int64_t group = *logicalLane / groupSize; + int64_t candidateSourceChunk = group / slots; + if (resolvedSourceChunk && *resolvedSourceChunk != candidateSourceChunk) + return failure(); + resolvedSourceChunk = candidateSourceChunk; + slotByLane.push_back((group % slots) * slotLaneStride); + } + if (!resolvedSourceChunk) + return failure(); + sourceChunk = *resolvedSourceChunk; + + FailureOr baseScalar = createScalarOffsetConstant( + loc, indexType.getElementType(), slotByLane.front(), rewriter); + FailureOr maskType = + getMaskTypeForVReg(indexType, rewriter.getContext()); + FailureOr allMask = createAllTrueMaskForVReg(loc, indexType, rewriter); + if (failed(baseScalar) || failed(maskType) || failed(allMask)) + return failure(); + + Value result = rewriter + .create(loc, indexType, *baseScalar, *allMask, + /*position=*/nullptr) + .getResult(); + int64_t rangeBegin = 0; + while (rangeBegin < lanesPerPart) { + int64_t slot = slotByLane[rangeBegin]; + int64_t rangeEnd = rangeBegin + 1; + while (rangeEnd < lanesPerPart && slotByLane[rangeEnd] == slot) + ++rangeEnd; + if (rangeBegin != 0 || slot != slotByLane.front()) { + FailureOr slotScalar = createScalarOffsetConstant( + loc, indexType.getElementType(), slot, rewriter); + FailureOr laneMask = + createLaneRangeMask(loc, *maskType, rangeBegin, rangeEnd, rewriter); + if (failed(slotScalar) || failed(laneMask)) + return failure(); + Value splat = rewriter + .create(loc, indexType, *slotScalar, *allMask, + /*position=*/nullptr) + .getResult(); + result = rewriter.create(loc, indexType, splat, result, *laneMask) + .getResult(); + } + rangeBegin = rangeEnd; + } + return result; +} + +std::optional getX2MemoryDistToken(Type elementType, + StringRef prefix) { + unsigned elementBits = pto::getPTOStorageElemBitWidth(elementType); + if (elementBits != 8 && elementBits != 16 && elementBits != 32) + return std::nullopt; + return (Twine(prefix) + "_B" + Twine(elementBits)).str(); +} + +std::optional getDenseLaneStrideLoadDistToken(VMIVRegType type) { + VMILayoutAttr layout = type.getLayoutAttr(); + if (!layout || !layout.isContiguous()) + return std::nullopt; + unsigned elementBits = pto::getPTOStorageElemBitWidth(type.getElementType()); + if (layout.getLaneStride() == 2 && + (elementBits == 8 || elementBits == 16 || elementBits == 32)) + return (Twine("UNPK_B") + Twine(elementBits)).str(); + if (layout.getLaneStride() == 4 && elementBits == 8) + return std::string("UNPK4"); + return std::nullopt; +} + +std::optional +getLaneStrideStoreDistToken(VMILayoutAttr layout, Type elementType) { + if (!layout || !layout.hasLaneStride()) + return std::nullopt; + unsigned elementBits = pto::getPTOStorageElemBitWidth(elementType); + if (layout.getLaneStride() == 2 && elementBits == 8) + return std::string("PK_B16"); + if (layout.getLaneStride() == 2 && elementBits == 16) + return std::string("PK_B32"); + if (layout.getLaneStride() == 2 && elementBits == 32) + return std::string("PK_B64"); + if (layout.getLaneStride() == 4 && elementBits == 8) + return std::string("PK4_B32"); + return std::nullopt; +} + +std::optional getDenseLaneStrideStoreDistToken(VMIVRegType type) { + VMILayoutAttr layout = type.getLayoutAttr(); + if (!layout || !layout.isContiguous()) + return std::nullopt; + return getLaneStrideStoreDistToken(layout, type.getElementType()); +} + +std::optional +getLaneStrideStoreMaskGranularity(VMILayoutAttr layout, Type elementType) { + if (!layout || !layout.hasLaneStride()) + return std::nullopt; + unsigned elementBits = pto::getPTOStorageElemBitWidth(elementType); + if (layout.getLaneStride() == 2 && elementBits == 8) + return StringRef("b16"); + if (layout.getLaneStride() == 2 && + (elementBits == 16 || elementBits == 32)) + return StringRef("b32"); + if (layout.getLaneStride() == 4 && elementBits == 8) + return StringRef("b32"); + return std::nullopt; +} + +std::optional +getDenseLaneStrideStoreMaskGranularity(VMIVRegType type) { + VMILayoutAttr layout = type.getLayoutAttr(); + if (!layout || !layout.isContiguous()) + return std::nullopt; + return getLaneStrideStoreMaskGranularity(layout, type.getElementType()); +} + +std::optional +getDenseLaneStrideMaskedStoreMaskGranularity(VMIVRegType type) { + VMILayoutAttr layout = type.getLayoutAttr(); + if (!layout || !layout.isContiguous()) + return std::nullopt; + unsigned elementBits = pto::getPTOStorageElemBitWidth(type.getElementType()); + if (layout.getLaneStride() == 2 && elementBits == 8) + return StringRef("b16"); + if (layout.getLaneStride() == 2 && elementBits == 16) + return StringRef("b32"); + if (layout.getLaneStride() == 4 && elementBits == 8) + return StringRef("b32"); + return std::nullopt; +} + +std::optional getPointStoreDistToken(Type elementType) { + unsigned elementBits = pto::getPTOStorageElemBitWidth(elementType); + if (elementBits != 8 && elementBits != 16 && elementBits != 32) + return std::nullopt; + return (Twine("1PT_B") + Twine(elementBits)).str(); +} + +struct VPTOCmpMode { + StringRef mode; + std::optional signedness; +}; + +std::optional getVPTOCmpFMode(StringRef predicate) { + if (predicate == "eq" || predicate == "ne" || predicate == "lt" || + predicate == "le" || predicate == "gt" || predicate == "ge") + return VPTOCmpMode{predicate, std::nullopt}; + if (predicate == "oeq") + return VPTOCmpMode{StringRef("eq"), std::nullopt}; + if (predicate == "one") + return VPTOCmpMode{StringRef("ne"), std::nullopt}; + if (predicate == "olt") + return VPTOCmpMode{StringRef("lt"), std::nullopt}; + if (predicate == "ole") + return VPTOCmpMode{StringRef("le"), std::nullopt}; + if (predicate == "ogt") + return VPTOCmpMode{StringRef("gt"), std::nullopt}; + if (predicate == "oge") + return VPTOCmpMode{StringRef("ge"), std::nullopt}; + return std::nullopt; +} + +std::optional getVPTOCmpIMode(StringRef predicate) { + if (predicate == "eq" || predicate == "ne") + return VPTOCmpMode{predicate, std::nullopt}; + if (predicate == "ult") + return VPTOCmpMode{ + StringRef("lt"), IntegerType::SignednessSemantics::Unsigned}; + if (predicate == "ule") + return VPTOCmpMode{ + StringRef("le"), IntegerType::SignednessSemantics::Unsigned}; + if (predicate == "ugt") + return VPTOCmpMode{ + StringRef("gt"), IntegerType::SignednessSemantics::Unsigned}; + if (predicate == "uge") + return VPTOCmpMode{ + StringRef("ge"), IntegerType::SignednessSemantics::Unsigned}; + if (predicate == "slt") + return VPTOCmpMode{ + StringRef("lt"), IntegerType::SignednessSemantics::Signed}; + if (predicate == "sle") + return VPTOCmpMode{ + StringRef("le"), IntegerType::SignednessSemantics::Signed}; + if (predicate == "sgt") + return VPTOCmpMode{ + StringRef("gt"), IntegerType::SignednessSemantics::Signed}; + if (predicate == "sge") + return VPTOCmpMode{ + StringRef("ge"), IntegerType::SignednessSemantics::Signed}; + return std::nullopt; +} + +template +std::optional getVPTOCmpMode(StringRef predicate) { + if constexpr (std::is_same_v) + return getVPTOCmpIMode(predicate); + else + return getVPTOCmpFMode(predicate); +} + +template +StringRef getSupportedComparePredicateMessage() { + if constexpr (std::is_same_v) + return "eq/ne, unsigned integer forms ult/ule/ugt/uge, and signed " + "integer forms slt/sle/sgt/sge"; + else + return "eq/ne/lt/le/gt/ge and ordered FP forms oeq/one/olt/ole/ogt/oge"; +} + +template +LogicalResult checkSupportedComparePredicate(Operation *op, + StringRef predicate) { + if (getVPTOCmpMode(predicate)) + return success(); + return op->emitError() + << kVMIDiagUnsupportedPrefix << "compare predicate " << predicate + << " cannot be lowered to pto.vcmp; supported predicates are " + << getSupportedComparePredicateMessage(); +} + +struct OneToNVMIUnpackOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIUnpackOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + ValueRange sourceParts = adaptor.getSource(); + if (sourceParts.size() != op->getNumResults()) + return rewriter.notifyMatchFailure( + op, "converted source part count must match unpack results"); + replaceOpWithFlatConvertedValues(rewriter, op, sourceParts, + *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMIPackOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIPackOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + FailureOr arity = getVMIPhysicalArity(op.getResult().getType()); + SmallVector flatOperands = flattenOneToNOperands(adaptor.getOperands()); + if (failed(arity) || static_cast(flatOperands.size()) != *arity) + return rewriter.notifyMatchFailure( + op, "pack part count must match converted VMI result arity"); + replaceOpWithFlatConvertedValues(rewriter, op, flatOperands, + *this->getTypeConverter()); + return success(); + } +}; + +LogicalResult verifyIdentityPartForwarding(Operation *op, + ValueRange sourceParts, + TypeRange resultTypes, + PatternRewriter &rewriter) { + if (sourceParts.size() != resultTypes.size()) + return rewriter.notifyMatchFailure( + op, "source and result physical arity mismatch"); + for (auto [part, resultType] : llvm::zip_equal(sourceParts, resultTypes)) { + if (part.getType() != resultType) + return rewriter.notifyMatchFailure( + op, "helper requires non-identity physical materialization"); + } + return success(); +} + +FailureOr getUnsignedCarrierVRegType(MLIRContext *ctx, + unsigned elementBits) { + if (elementBits != 8 && elementBits != 16 && elementBits != 32) + return failure(); + auto elementType = IntegerType::get( + ctx, elementBits, IntegerType::SignednessSemantics::Unsigned); + return VRegType::get(ctx, 2048 / elementBits, elementType); +} + +FailureOr +getSignednessCarrierVRegType(VRegType inputType, + IntegerType::SignednessSemantics signedness) { + auto inputElementType = dyn_cast(inputType.getElementType()); + if (!inputElementType) + return failure(); + if ((signedness == IntegerType::SignednessSemantics::Signed && + !inputElementType.isUnsigned()) || + (signedness == IntegerType::SignednessSemantics::Unsigned && + inputElementType.isUnsigned())) + return inputType; + auto carrierElementType = IntegerType::get( + inputType.getContext(), inputElementType.getWidth(), signedness); + return VRegType::get(inputType.getContext(), inputType.getElementCount(), + carrierElementType); +} + +FailureOr bitcastVReg(Location loc, Value value, Type resultType, + PatternRewriter &rewriter) { + if (value.getType() == resultType) + return value; + auto inputType = dyn_cast(value.getType()); + auto outputType = dyn_cast(resultType); + if (!inputType || !outputType) + return failure(); + return rewriter.create(loc, outputType, value).getResult(); +} + +FailureOr getVcaddResultType(VRegType inputType) { + auto inputIntegerType = dyn_cast(inputType.getElementType()); + if (!inputIntegerType || inputIntegerType.getWidth() == 32) + return inputType; + unsigned inputWidth = inputIntegerType.getWidth(); + if (inputWidth != 8 && inputWidth != 16) + return failure(); + auto resultElementType = IntegerType::get( + inputType.getContext(), inputWidth * 2, + inputIntegerType.getSignedness()); + return VRegType::get(inputType.getContext(), + inputType.getElementCount() / 2, resultElementType); +} + +FailureOr unpackToNextCarrier(Location loc, Value source, + unsigned sourceBits, int64_t partIndex, + PatternRewriter &rewriter) { + FailureOr resultType = + getUnsignedCarrierVRegType(rewriter.getContext(), sourceBits * 2); + if (failed(resultType)) + return failure(); + Value part = rewriter.create(loc, partIndex); + return rewriter.create(loc, *resultType, source, part) + .getResult(); +} + +FailureOr packToPreviousCarrier(Location loc, Value source, + unsigned resultBits, + PatternRewriter &rewriter) { + FailureOr resultType = + getUnsignedCarrierVRegType(rewriter.getContext(), resultBits); + if (failed(resultType)) + return failure(); + return rewriter + .create(loc, *resultType, source, + rewriter.getStringAttr("LOWER")) + .getResult(); +} + +FailureOr> materializeContiguousToLaneStride( + Operation *op, ValueRange sourceParts, TypeRange resultTypes, + Type elementType, int64_t laneStride, PatternRewriter &rewriter) { + if (sourceParts.size() != resultTypes.size()) { + (void)rewriter.notifyMatchFailure( + op, "dense lane_stride unpack materialization requires matching " + "source/result physical arity"); + return failure(); + } + + unsigned elementBits = pto::getPTOStorageElemBitWidth(elementType); + if ((laneStride != 2 && laneStride != 4) || + (laneStride == 4 && elementBits != 8) || + (elementBits != 8 && elementBits != 16)) { + (void)rewriter.notifyMatchFailure( + op, "unsupported dense lane_stride unpack carrier shape"); + return failure(); + } + + MLIRContext *ctx = rewriter.getContext(); + FailureOr inputCarrier = + getUnsignedCarrierVRegType(ctx, elementBits); + if (failed(inputCarrier)) + return failure(); + + SmallVector results; + results.reserve(resultTypes.size()); + for (auto [resultIndex, resultType] : llvm::enumerate(resultTypes)) { + int64_t sourceIndex = resultIndex / laneStride; + if (sourceIndex >= static_cast(sourceParts.size())) + return failure(); + Value source = sourceParts[sourceIndex]; + FailureOr current = + bitcastVReg(op->getLoc(), source, *inputCarrier, rewriter); + if (failed(current)) + return failure(); + int64_t part = resultIndex % laneStride; + FailureOr unpacked = + unpackToNextCarrier(op->getLoc(), *current, elementBits, + laneStride == 4 ? part / 2 : part, rewriter); + if (failed(unpacked)) + return failure(); + current = *unpacked; + if (laneStride == 4) { + unpacked = unpackToNextCarrier(op->getLoc(), *current, elementBits * 2, + part % 2, rewriter); + if (failed(unpacked)) + return failure(); + current = *unpacked; + } + FailureOr result = + bitcastVReg(op->getLoc(), *current, resultType, rewriter); + if (failed(result)) + return failure(); + results.push_back(*result); + } + return results; +} + +FailureOr> materializeLaneStrideToContiguous( + Operation *op, ValueRange sourceParts, TypeRange resultTypes, + Type elementType, int64_t laneStride, PatternRewriter &rewriter) { + if (sourceParts.size() != resultTypes.size()) { + (void)rewriter.notifyMatchFailure( + op, "dense lane_stride pack materialization requires matching " + "source/result physical arity"); + return failure(); + } + + unsigned elementBits = pto::getPTOStorageElemBitWidth(elementType); + if ((laneStride != 2 && laneStride != 4) || + (laneStride == 4 && elementBits != 8) || + (elementBits != 8 && elementBits != 16)) { + (void)rewriter.notifyMatchFailure( + op, "unsupported dense lane_stride pack carrier shape"); + return failure(); + } + + unsigned carrierBits = + static_cast(elementBits * static_cast(laneStride)); + FailureOr sourceCarrier = + getUnsignedCarrierVRegType(rewriter.getContext(), carrierBits); + if (failed(sourceCarrier)) + return failure(); + + SmallVector results; + results.reserve(sourceParts.size()); + for (auto [source, resultType] : llvm::zip_equal(sourceParts, resultTypes)) { + FailureOr current = + bitcastVReg(op->getLoc(), source, *sourceCarrier, rewriter); + if (failed(current)) + return failure(); + FailureOr packed = packToPreviousCarrier(op->getLoc(), *current, + carrierBits / 2, rewriter); + if (failed(packed)) + return failure(); + current = *packed; + if (laneStride == 4) { + packed = + packToPreviousCarrier(op->getLoc(), *current, elementBits, rewriter); + if (failed(packed)) + return failure(); + current = *packed; + } + FailureOr result = + bitcastVReg(op->getLoc(), *current, resultType, rewriter); + if (failed(result)) + return failure(); + results.push_back(*result); + } + return results; +} + +FailureOr> materializeGroupSlotLaneStride( + Operation *op, ValueRange sourceParts, TypeRange resultTypes, + Type elementType, int64_t sourceStride, int64_t resultStride, + PatternRewriter &rewriter) { + auto fail = [&](const Twine &message) -> FailureOr> { + (void)rewriter.notifyMatchFailure(op, message); + return failure(); + }; + + if (sourceParts.size() != resultTypes.size() || sourceParts.empty()) + return fail("group-slot lane_stride materialization requires matching " + "non-empty source/result physical arity"); + if ((sourceStride != 1 && sourceStride != 2 && sourceStride != 4) || + (resultStride != 1 && resultStride != 2 && resultStride != 4)) + return fail("unsupported group-slot lane_stride factor"); + + unsigned elementBits = pto::getPTOStorageElemBitWidth(elementType); + int64_t maxStride = std::max(sourceStride, resultStride); + if ((elementBits != 8 && elementBits != 16) || + elementBits * maxStride > 32) + return fail("unsupported group-slot lane_stride carrier shape"); + + SmallVector results; + results.reserve(resultTypes.size()); + for (auto [source, resultType] : + llvm::zip_equal(sourceParts, resultTypes)) { + unsigned carrierBits = elementBits * sourceStride; + FailureOr carrierType = + getUnsignedCarrierVRegType(rewriter.getContext(), carrierBits); + if (failed(carrierType)) + return fail("failed to derive group-slot source carrier type"); + FailureOr current = + bitcastVReg(op->getLoc(), source, *carrierType, rewriter); + if (failed(current)) + return fail("failed to bitcast group-slot source carrier"); + + int64_t currentStride = sourceStride; + while (currentStride < resultStride) { + FailureOr unpacked = unpackToNextCarrier( + op->getLoc(), *current, carrierBits, /*partIndex=*/0, rewriter); + if (failed(unpacked)) + return fail("failed to unpack group-slot lane_stride carrier"); + current = *unpacked; + currentStride *= 2; + carrierBits *= 2; + } + while (currentStride > resultStride) { + FailureOr packed = packToPreviousCarrier( + op->getLoc(), *current, carrierBits / 2, rewriter); + if (failed(packed)) + return fail("failed to pack group-slot lane_stride carrier"); + current = *packed; + currentStride /= 2; + carrierBits /= 2; + } + + FailureOr result = + bitcastVReg(op->getLoc(), *current, resultType, rewriter); + if (failed(result)) + return fail("failed to bitcast group-slot result carrier"); + results.push_back(*result); + } + return results; +} + +FailureOr> materializeDataLayoutConversion( + Operation *op, ValueRange sourceParts, TypeRange resultTypes, + VMILayoutAttr sourceLayout, VMILayoutAttr resultLayout, + PatternRewriter &rewriter) { + if (!sourceLayout || !resultLayout) { + (void)rewriter.notifyMatchFailure( + op, "layout materialization requires assigned source/result layouts"); + return failure(); + } + + if (sourceLayout == resultLayout) { + if (failed(verifyIdentityPartForwarding(op, sourceParts, resultTypes, + rewriter))) + return failure(); + return SmallVector(sourceParts.begin(), sourceParts.end()); + } + + if (sourceLayout.isGroupSlots() && resultLayout.isGroupSlots() && + sourceLayout.getNumGroups() == resultLayout.getNumGroups() && + sourceLayout.getSlots() == 8 && resultLayout.getSlots() == 8) { + auto ensure = dyn_cast(op); + if (!ensure) + return failure(); + auto sourceType = cast(ensure.getSource().getType()); + return materializeGroupSlotLaneStride( + op, sourceParts, resultTypes, sourceType.getElementType(), + sourceLayout.getLaneStride(), resultLayout.getLaneStride(), rewriter); + } + + auto isElementDeinterleaved = [](VMILayoutAttr layout, int64_t factor) { + return layout.isDeinterleaved() && layout.getFactor() == factor && + layout.getLaneStride() == 1 && layout.getBlockElems() == 1; + }; + auto isBlock8Deinterleaved = [](VMILayoutAttr layout, int64_t factor) { + return layout.isDeinterleaved() && layout.getFactor() == factor && + layout.getBlockElems() == 8; + }; + + bool contiguousToBlock8 = + sourceLayout.isContiguous() && sourceLayout.getLaneStride() == 1 && + ((resultLayout.isDeinterleaved() && + isBlock8Deinterleaved(resultLayout, 2)) || + (resultLayout.isDeinterleaved() && + isBlock8Deinterleaved(resultLayout, 4))); + bool block8ToContiguous = + resultLayout.isContiguous() && resultLayout.getLaneStride() == 1 && + ((sourceLayout.isDeinterleaved() && + isBlock8Deinterleaved(sourceLayout, 2)) || + (sourceLayout.isDeinterleaved() && + isBlock8Deinterleaved(sourceLayout, 4))); + if (contiguousToBlock8 || block8ToContiguous) { + if (sourceParts.size() == 1) { + if (auto cast = + sourceParts.front().getDefiningOp()) { + ValueRange inputs = cast.getInputs(); + if (inputs.size() == resultTypes.size()) { + bool typesMatch = true; + for (auto [input, resultType] : llvm::zip_equal(inputs, resultTypes)) + if (input.getType() != resultType) { + typesMatch = false; + break; + } + if (typesMatch) + return SmallVector(inputs.begin(), inputs.end()); + } + } + } + if (failed(verifyIdentityPartForwarding(op, sourceParts, resultTypes, + rewriter))) + return failure(); + return SmallVector(sourceParts.begin(), sourceParts.end()); + } + + bool deint2ToContiguous = sourceLayout.isDeinterleaved() && + isElementDeinterleaved(sourceLayout, 2) && + resultLayout.isContiguous() && + resultLayout.getLaneStride() == 1; + bool contiguousToDeint2 = + sourceLayout.isContiguous() && sourceLayout.getLaneStride() == 1 && + resultLayout.isDeinterleaved() && isElementDeinterleaved(resultLayout, 2); + if (deint2ToContiguous || contiguousToDeint2) { + SmallVector results; + if (deint2ToContiguous) { + if (sourceParts.empty() || sourceParts.size() % 2 != 0 || + resultTypes.empty()) { + (void)rewriter.notifyMatchFailure( + op, "deinterleaved=2 to contiguous materialization requires " + "2*N source parts and at least one result part"); + return failure(); + } + int64_t groups = sourceParts.size() / 2; + if (resultTypes.size() > static_cast(2 * groups)) { + (void)rewriter.notifyMatchFailure( + op, "deinterleaved=2 to contiguous materialization result arity " + "exceeds source footprint"); + return failure(); + } + + results.reserve(resultTypes.size()); + for (int64_t i = 0; i < groups && results.size() < resultTypes.size(); + ++i) { + Value lhs = sourceParts[i]; + Value rhs = sourceParts[groups + i]; + if (lhs.getType() != rhs.getType()) + return rewriter.notifyMatchFailure( + op, "vintlv requires matching source part types"); + Type lowType = resultTypes[results.size()]; + Type highType = results.size() + 1 < resultTypes.size() + ? resultTypes[results.size() + 1] + : lowType; + if (lhs.getType() != lowType || lhs.getType() != highType) + return rewriter.notifyMatchFailure( + op, "vintlv requires operands and results to share one type"); + auto materialize = rewriter.create( + op->getLoc(), lowType, highType, lhs, rhs); + results.push_back(materialize.getLow()); + if (results.size() < resultTypes.size()) + results.push_back(materialize.getHigh()); + } + } else { + if (sourceParts.empty() || resultTypes.empty() || + resultTypes.size() % 2 != 0) { + (void)rewriter.notifyMatchFailure( + op, "contiguous to deinterleaved=2 materialization requires " + "at least one source part and 2*N result parts"); + return failure(); + } + int64_t groups = resultTypes.size() / 2; + if (sourceParts.size() > static_cast(2 * groups)) { + (void)rewriter.notifyMatchFailure( + op, "contiguous to deinterleaved=2 materialization source " + "footprint exceeds result arity"); + return failure(); + } + + SmallVector part0; + SmallVector part1; + part0.reserve(groups); + part1.reserve(groups); + for (int64_t i = 0; i < groups; ++i) { + size_t lhsIndex = 2 * i; + if (lhsIndex >= sourceParts.size()) + return rewriter.notifyMatchFailure( + op, "contiguous to deinterleaved=2 materialization missing " + "source part"); + size_t rhsIndex = lhsIndex + 1 < sourceParts.size() ? lhsIndex + 1 + : lhsIndex; + Value lhs = sourceParts[lhsIndex]; + Value rhs = sourceParts[rhsIndex]; + if (lhs.getType() != rhs.getType() || + lhs.getType() != resultTypes[i] || + lhs.getType() != resultTypes[groups + i]) + return rewriter.notifyMatchFailure( + op, "vdintlv requires operands and results to share one type"); + auto materialize = rewriter.create( + op->getLoc(), resultTypes[i], resultTypes[groups + i], + lhs, rhs); + part0.push_back(materialize.getLow()); + part1.push_back(materialize.getHigh()); + } + results.reserve(resultTypes.size()); + results.append(part0); + results.append(part1); + } + return results; + } + + bool deint4ToContiguous = sourceLayout.isDeinterleaved() && + isElementDeinterleaved(sourceLayout, 4) && + resultLayout.isContiguous() && + resultLayout.getLaneStride() == 1; + bool contiguousToDeint4 = + sourceLayout.isContiguous() && sourceLayout.getLaneStride() == 1 && + resultLayout.isDeinterleaved() && isElementDeinterleaved(resultLayout, 4); + if (deint4ToContiguous || contiguousToDeint4) { + auto getPartCounts = [](size_t totalParts, + int64_t factor) -> SmallVector { + SmallVector counts; + counts.reserve(factor); + size_t base = totalParts / static_cast(factor); + size_t remainder = totalParts % static_cast(factor); + for (int64_t part = 0; part < factor; ++part) + counts.push_back(base + (static_cast(part) < remainder ? 1 : 0)); + return counts; + }; + auto getPartOffsets = [](ArrayRef counts) -> SmallVector { + SmallVector offsets; + offsets.reserve(counts.size()); + size_t offset = 0; + for (size_t count : counts) { + offsets.push_back(offset); + offset += count; + } + return offsets; + }; + + SmallVector results; + if (deint4ToContiguous) { + if (sourceParts.empty() || resultTypes.empty()) { + (void)rewriter.notifyMatchFailure( + op, "deinterleaved=4 to contiguous materialization requires " + "at least one source and result part"); + return failure(); + } + if (resultTypes.size() > sourceParts.size()) { + (void)rewriter.notifyMatchFailure( + op, "deinterleaved=4 to contiguous materialization result arity " + "exceeds source footprint"); + return failure(); + } + + SmallVector sourceCounts = getPartCounts(sourceParts.size(), 4); + SmallVector sourceOffsets = getPartOffsets(sourceCounts); + auto getSourcePart = [&](size_t part, size_t group) -> Value { + if (group < sourceCounts[part]) + return sourceParts[sourceOffsets[part] + group]; + return sourceParts.back(); + }; + + results.reserve(resultTypes.size()); + size_t groups = (resultTypes.size() + 3) / 4; + for (size_t i = 0; i < groups && results.size() < resultTypes.size(); + ++i) { + Value p0 = getSourcePart(0, i); + Value p1 = getSourcePart(1, i); + Value p2 = getSourcePart(2, i); + Value p3 = getSourcePart(3, i); + Type chunkType = p0.getType(); + if (p1.getType() != chunkType || p2.getType() != chunkType || + p3.getType() != chunkType) + return rewriter.notifyMatchFailure( + op, "vintlv deinterleaved=4 requires matching source part " + "types"); + for (size_t resultIndex = results.size(); + resultIndex < resultTypes.size() && resultIndex < results.size() + 4; + ++resultIndex) { + if (resultTypes[resultIndex] != chunkType) + return rewriter.notifyMatchFailure( + op, "vintlv requires operands and results to share one type"); + } + + auto even = rewriter.create(op->getLoc(), chunkType, + chunkType, p0, p2); + auto odd = rewriter.create(op->getLoc(), chunkType, + chunkType, p1, p3); + auto low = rewriter.create(op->getLoc(), chunkType, chunkType, + even.getLow(), odd.getLow()); + auto high = rewriter.create(op->getLoc(), chunkType, + chunkType, even.getHigh(), + odd.getHigh()); + Value groupResults[] = {low.getLow(), low.getHigh(), high.getLow(), + high.getHigh()}; + for (Value result : groupResults) { + if (results.size() >= resultTypes.size()) + break; + results.push_back(result); + } + } + } else { + if (sourceParts.empty() || resultTypes.empty()) { + (void)rewriter.notifyMatchFailure( + op, "contiguous to deinterleaved=4 materialization requires " + "at least one source and result part"); + return failure(); + } + if (sourceParts.size() > resultTypes.size()) { + (void)rewriter.notifyMatchFailure( + op, "contiguous to deinterleaved=4 materialization source " + "footprint exceeds result arity"); + return failure(); + } + + SmallVector resultCounts = getPartCounts(resultTypes.size(), 4); + SmallVector resultOffsets = getPartOffsets(resultCounts); + auto getContiguousSourcePart = [&](size_t index) { + return sourceParts[std::min(index, sourceParts.size() - 1)]; + }; + SmallVector part0; + SmallVector part1; + SmallVector part2; + SmallVector part3; + part0.reserve(resultCounts[0]); + part1.reserve(resultCounts[1]); + part2.reserve(resultCounts[2]); + part3.reserve(resultCounts[3]); + size_t groups = + *std::max_element(resultCounts.begin(), resultCounts.end()); + for (size_t i = 0; i < groups; ++i) { + Value s0 = getContiguousSourcePart(4 * i); + Value s1 = getContiguousSourcePart(4 * i + 1); + Value s2 = getContiguousSourcePart(4 * i + 2); + Value s3 = getContiguousSourcePart(4 * i + 3); + Type chunkType = s0.getType(); + if (s0.getType() != s1.getType() || s0.getType() != s2.getType() || + s0.getType() != s3.getType()) + return rewriter.notifyMatchFailure( + op, "vdintlv deinterleaved=4 requires matching source part " + "types"); + for (int64_t part = 0; part < 4; ++part) { + if (i < resultCounts[part] && + resultTypes[resultOffsets[part] + i] != chunkType) + return rewriter.notifyMatchFailure( + op, "vdintlv requires operands and results to share one type"); + } + + auto low = rewriter.create( + op->getLoc(), chunkType, chunkType, s0, s1); + auto high = rewriter.create(op->getLoc(), chunkType, + chunkType, s2, s3); + auto even = rewriter.create( + op->getLoc(), chunkType, chunkType, low.getLow(), high.getLow()); + auto odd = rewriter.create( + op->getLoc(), chunkType, chunkType, low.getHigh(), high.getHigh()); + if (i < resultCounts[0]) + part0.push_back(even.getLow()); + if (i < resultCounts[1]) + part1.push_back(odd.getLow()); + if (i < resultCounts[2]) + part2.push_back(even.getHigh()); + if (i < resultCounts[3]) + part3.push_back(odd.getHigh()); + } + results.reserve(resultTypes.size()); + results.append(part0); + results.append(part1); + results.append(part2); + results.append(part3); + } + return results; + } + + if (sourceLayout.isContiguous() && sourceLayout.getLaneStride() == 1 && + resultLayout.isContiguous() && resultLayout.getLaneStride() != 1) { + auto ensure = dyn_cast(op); + if (!ensure) + return failure(); + auto sourceType = cast(ensure.getSource().getType()); + return materializeContiguousToLaneStride( + op, sourceParts, resultTypes, sourceType.getElementType(), + resultLayout.getLaneStride(), rewriter); + } + + if (sourceLayout.isContiguous() && sourceLayout.getLaneStride() != 1 && + resultLayout.isContiguous() && resultLayout.getLaneStride() == 1) { + auto ensure = dyn_cast(op); + if (!ensure) + return failure(); + auto sourceType = cast(ensure.getSource().getType()); + return materializeLaneStrideToContiguous( + op, sourceParts, resultTypes, sourceType.getElementType(), + sourceLayout.getLaneStride(), rewriter); + } + + if (sourceLayout.isDeinterleaved() && resultLayout.isDeinterleaved() && + sourceLayout.getLaneStride() == 1 && resultLayout.getLaneStride() == 1 && + sourceLayout.getBlockElems() == 1 && resultLayout.getBlockElems() == 1 && + (sourceLayout.getFactor() == 2 || sourceLayout.getFactor() == 4) && + (resultLayout.getFactor() == 2 || resultLayout.getFactor() == 4)) { + VMILayoutAttr contiguous = + VMILayoutAttr::getContiguous(rewriter.getContext()); + FailureOr> dense = materializeDataLayoutConversion( + op, sourceParts, resultTypes, sourceLayout, contiguous, rewriter); + if (failed(dense)) + return failure(); + return materializeDataLayoutConversion(op, *dense, resultTypes, contiguous, + resultLayout, rewriter); + } + + (void)rewriter.notifyMatchFailure( + op, "unsupported VMI data layout materialization"); + return failure(); +} + +FailureOr> +createPredicateDintlv(Location loc, Type lowType, Type highType, Value lhs, + Value rhs, PatternRewriter &rewriter) { + auto maskType = dyn_cast(lowType); + if (!maskType || highType != lowType) + return failure(); + if (maskType.isB8()) { + auto op = rewriter.create(loc, lowType, highType, lhs, rhs); + return std::make_pair(op.getLow(), op.getHigh()); + } + if (maskType.isB16()) { + auto op = rewriter.create(loc, lowType, highType, lhs, rhs); + return std::make_pair(op.getLow(), op.getHigh()); + } + if (maskType.isB32()) { + auto op = rewriter.create(loc, lowType, highType, lhs, rhs); + return std::make_pair(op.getLow(), op.getHigh()); + } + return failure(); +} + +FailureOr> +createPredicateIntlv(Location loc, Type lowType, Type highType, Value lhs, + Value rhs, PatternRewriter &rewriter) { + auto maskType = dyn_cast(lowType); + if (!maskType || highType != lowType) + return failure(); + if (maskType.isB8()) { + auto op = rewriter.create(loc, lowType, highType, lhs, rhs); + return std::make_pair(op.getLow(), op.getHigh()); + } + if (maskType.isB16()) { + auto op = rewriter.create(loc, lowType, highType, lhs, rhs); + return std::make_pair(op.getLow(), op.getHigh()); + } + if (maskType.isB32()) { + auto op = rewriter.create(loc, lowType, highType, lhs, rhs); + return std::make_pair(op.getLow(), op.getHigh()); + } + return failure(); +} + +FailureOr> materializeMaskLayoutConversion( + Operation *op, ValueRange sourceParts, TypeRange resultTypes, + VMILayoutAttr sourceLayout, VMILayoutAttr resultLayout, + PatternRewriter &rewriter) { + if (!sourceLayout || !resultLayout) { + (void)rewriter.notifyMatchFailure( + op, "mask layout materialization requires assigned source/result " + "layouts"); + return failure(); + } + + if (sourceLayout == resultLayout) { + if (failed(verifyIdentityPartForwarding(op, sourceParts, resultTypes, + rewriter))) + return failure(); + return SmallVector(sourceParts.begin(), sourceParts.end()); + } + + auto isElementDeinterleaved = [](VMILayoutAttr layout, int64_t factor) { + return layout.isDeinterleaved() && layout.getFactor() == factor && + layout.getLaneStride() == 1 && layout.getBlockElems() == 1; + }; + + bool deint2ToContiguous = sourceLayout.isDeinterleaved() && + isElementDeinterleaved(sourceLayout, 2) && + resultLayout.isContiguous() && + resultLayout.getLaneStride() == 1; + bool contiguousToDeint2 = + sourceLayout.isContiguous() && sourceLayout.getLaneStride() == 1 && + resultLayout.isDeinterleaved() && isElementDeinterleaved(resultLayout, 2); + if (deint2ToContiguous || contiguousToDeint2) { + if (sourceParts.size() != resultTypes.size() || sourceParts.empty() || + sourceParts.size() % 2 != 0) { + (void)rewriter.notifyMatchFailure( + op, "deinterleaved=2 mask layout materialization requires 2*N " + "parts"); + return failure(); + } + if (failed(verifyIdentityPartForwarding(op, sourceParts, resultTypes, + rewriter))) + return failure(); + + int64_t groups = sourceParts.size() / 2; + SmallVector results; + results.reserve(sourceParts.size()); + if (deint2ToContiguous) { + for (int64_t i = 0; i < groups; ++i) { + FailureOr> materialize = createPredicateIntlv( + op->getLoc(), resultTypes[2 * i], resultTypes[2 * i + 1], + sourceParts[i], sourceParts[groups + i], rewriter); + if (failed(materialize)) + return rewriter.notifyMatchFailure( + op, "unsupported predicate intlv mask type"); + results.append({materialize->first, materialize->second}); + } + } else { + SmallVector part0; + SmallVector part1; + part0.reserve(groups); + part1.reserve(groups); + for (int64_t i = 0; i < groups; ++i) { + FailureOr> materialize = createPredicateDintlv( + op->getLoc(), resultTypes[i], resultTypes[groups + i], + sourceParts[2 * i], sourceParts[2 * i + 1], rewriter); + if (failed(materialize)) + return rewriter.notifyMatchFailure( + op, "unsupported predicate dintlv mask type"); + part0.push_back(materialize->first); + part1.push_back(materialize->second); + } + results.append(part0); + results.append(part1); + } + return results; + } + + bool deint4ToContiguous = sourceLayout.isDeinterleaved() && + isElementDeinterleaved(sourceLayout, 4) && + resultLayout.isContiguous() && + resultLayout.getLaneStride() == 1; + bool contiguousToDeint4 = + sourceLayout.isContiguous() && sourceLayout.getLaneStride() == 1 && + resultLayout.isDeinterleaved() && isElementDeinterleaved(resultLayout, 4); + if (deint4ToContiguous || contiguousToDeint4) { + if (sourceParts.size() != resultTypes.size() || sourceParts.empty() || + sourceParts.size() % 4 != 0) { + (void)rewriter.notifyMatchFailure( + op, "deinterleaved=4 mask layout materialization requires 4*N " + "parts"); + return failure(); + } + if (failed(verifyIdentityPartForwarding(op, sourceParts, resultTypes, + rewriter))) + return failure(); + + SmallVector results; + results.reserve(sourceParts.size()); + int64_t groups = sourceParts.size() / 4; + if (deint4ToContiguous) { + for (int64_t i = 0; i < groups; ++i) { + Value p0 = sourceParts[i]; + Value p1 = sourceParts[groups + i]; + Value p2 = sourceParts[2 * groups + i]; + Value p3 = sourceParts[3 * groups + i]; + FailureOr> even = + createPredicateIntlv(op->getLoc(), resultTypes[4 * i], + resultTypes[4 * i + 1], p0, p2, rewriter); + FailureOr> odd = + createPredicateIntlv(op->getLoc(), resultTypes[4 * i], + resultTypes[4 * i + 1], p1, p3, rewriter); + if (failed(even) || failed(odd)) + return rewriter.notifyMatchFailure( + op, "unsupported predicate intlv mask type"); + FailureOr> low = createPredicateIntlv( + op->getLoc(), resultTypes[4 * i], resultTypes[4 * i + 1], + even->first, odd->first, rewriter); + FailureOr> high = createPredicateIntlv( + op->getLoc(), resultTypes[4 * i + 2], resultTypes[4 * i + 3], + even->second, odd->second, rewriter); + if (failed(low) || failed(high)) + return rewriter.notifyMatchFailure( + op, "unsupported predicate intlv mask type"); + results.append({low->first, low->second, high->first, high->second}); + } + } else { + SmallVector part0; + SmallVector part1; + SmallVector part2; + SmallVector part3; + part0.reserve(groups); + part1.reserve(groups); + part2.reserve(groups); + part3.reserve(groups); + for (int64_t i = 0; i < groups; ++i) { + FailureOr> low = createPredicateDintlv( + op->getLoc(), resultTypes[i], resultTypes[groups + i], + sourceParts[4 * i], sourceParts[4 * i + 1], rewriter); + FailureOr> high = createPredicateDintlv( + op->getLoc(), resultTypes[2 * groups + i], + resultTypes[3 * groups + i], sourceParts[4 * i + 2], + sourceParts[4 * i + 3], rewriter); + if (failed(low) || failed(high)) + return rewriter.notifyMatchFailure( + op, "unsupported predicate dintlv mask type"); + FailureOr> even = createPredicateDintlv( + op->getLoc(), resultTypes[i], resultTypes[2 * groups + i], + low->first, high->first, rewriter); + FailureOr> odd = createPredicateDintlv( + op->getLoc(), resultTypes[groups + i], resultTypes[3 * groups + i], + low->second, high->second, rewriter); + if (failed(even) || failed(odd)) + return rewriter.notifyMatchFailure( + op, "unsupported predicate dintlv mask type"); + part0.push_back(even->first); + part1.push_back(odd->first); + part2.push_back(even->second); + part3.push_back(odd->second); + } + results.append(part0); + results.append(part1); + results.append(part2); + results.append(part3); + } + return results; + } + + if (sourceLayout.isContiguous() && sourceLayout.getLaneStride() == 1 && + resultLayout.isContiguous() && resultLayout.getLaneStride() != 1) { + int64_t laneStride = resultLayout.getLaneStride(); + if (laneStride != 2 && laneStride != 4) + return rewriter.notifyMatchFailure( + op, "unsupported dense mask lane_stride unpack factor"); + if (static_cast(resultTypes.size()) > + static_cast(sourceParts.size()) * laneStride) + return rewriter.notifyMatchFailure( + op, "dense mask lane_stride unpack materialization result arity " + "does not fit source arity"); + SmallVector results; + results.reserve(resultTypes.size()); + auto lower = rewriter.getStringAttr("LOWER"); + auto higher = rewriter.getStringAttr("HIGHER"); + for (auto [resultIndex, resultType] : llvm::enumerate(resultTypes)) { + auto maskType = dyn_cast(resultType); + if (!maskType) + return rewriter.notifyMatchFailure( + op, "dense mask lane_stride unpack requires mask result type"); + int64_t sourceIndex = resultIndex / laneStride; + int64_t part = resultIndex % laneStride; + Value source = sourceParts[sourceIndex]; + StringAttr firstPart = laneStride == 4 ? (part >= 2 ? higher : lower) + : (part == 1 ? higher : lower); + Value current = + rewriter.create(op->getLoc(), maskType, source, firstPart); + if (laneStride == 4) + current = rewriter.create(op->getLoc(), maskType, current, + part % 2 == 0 ? lower : higher); + results.push_back(current); + } + return results; + } + + if (sourceLayout.isContiguous() && sourceLayout.getLaneStride() != 1 && + resultLayout.isContiguous() && resultLayout.getLaneStride() == 1) { + if (sourceParts.empty()) + return rewriter.notifyMatchFailure( + op, "dense mask lane_stride pack materialization requires source " + "parts"); + int64_t laneStride = sourceLayout.getLaneStride(); + if (laneStride != 2 && laneStride != 4) + return rewriter.notifyMatchFailure( + op, "unsupported dense mask lane_stride pack factor"); + if (static_cast(sourceParts.size()) > + static_cast(resultTypes.size()) * laneStride) + return rewriter.notifyMatchFailure( + op, "dense mask lane_stride pack materialization source arity does " + "not fit result arity"); + SmallVector results; + results.reserve(resultTypes.size()); + auto lower = rewriter.getStringAttr("LOWER"); + auto higher = rewriter.getStringAttr("HIGHER"); + Value allTrue; + auto mergeMasks = [&](Value lhs, Value rhs) -> FailureOr { + if (!allTrue) { + FailureOr mask = createAllTrueMask( + op->getLoc(), cast(lhs.getType()), rewriter); + if (failed(mask)) + return failure(); + allTrue = *mask; + } + return rewriter.create(op->getLoc(), lhs.getType(), lhs, rhs, + allTrue) + .getResult(); + }; + auto packPair = [&](Value lowSource, std::optional highSource, + MaskType maskType) -> FailureOr { + Value packed = + rewriter.create(op->getLoc(), maskType, lowSource, lower); + if (!highSource) + return packed; + Value higherPacked = rewriter.create( + op->getLoc(), maskType, *highSource, higher); + return mergeMasks(packed, higherPacked); + }; + for (auto [resultIndex, resultType] : + llvm::enumerate(resultTypes)) { + auto maskType = dyn_cast(resultType); + if (!maskType) + return rewriter.notifyMatchFailure( + op, "dense mask lane_stride pack requires mask result type"); + size_t base = resultIndex * static_cast(laneStride); + if (base >= sourceParts.size()) + break; + + std::optional source1; + if (base + 1 < sourceParts.size()) + source1 = sourceParts[base + 1]; + FailureOr lowHalf = packPair(sourceParts[base], source1, maskType); + if (failed(lowHalf)) + return failure(); + Value current = *lowHalf; + if (laneStride == 4) { + current = + rewriter.create(op->getLoc(), maskType, current, lower); + if (base + 2 < sourceParts.size()) { + std::optional source3; + if (base + 3 < sourceParts.size()) + source3 = sourceParts[base + 3]; + FailureOr highHalf = + packPair(sourceParts[base + 2], source3, maskType); + if (failed(highHalf)) + return failure(); + Value higherPacked = rewriter.create( + op->getLoc(), maskType, *highHalf, higher); + FailureOr merged = mergeMasks(current, higherPacked); + if (failed(merged)) + return failure(); + current = *merged; + } + } + results.push_back(current); + } + if (results.size() != resultTypes.size()) + return rewriter.notifyMatchFailure( + op, "dense mask lane_stride pack materialization result arity " + "mismatch"); + return results; + } + + (void)rewriter.notifyMatchFailure( + op, "unsupported VMI mask layout materialization"); + return failure(); +} + +int getMaskGranularityRank(StringRef granularity) { + if (granularity == "b8") + return 0; + if (granularity == "b16") + return 1; + if (granularity == "b32") + return 2; + return -1; +} + +StringRef getMaskGranularityForRank(int rank) { + switch (rank) { + case 0: + return "b8"; + case 1: + return "b16"; + case 2: + return "b32"; + default: + return ""; + } +} + +LogicalResult checkSupportedMaskGranularityMaterialization( + VMIMaskType sourceType, + VMIMaskType resultType, std::string *reason) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + if (sourceType.getElementCount() != resultType.getElementCount()) + return fail("requires source and result mask lane counts to match"); + if (sourceType.getLayoutAttr() != resultType.getLayoutAttr()) + return fail("requires source and result mask layouts to match"); + + if (!VMIMaskType::isConcreteGranularity(sourceType.getGranularity()) || + !VMIMaskType::isConcreteGranularity(resultType.getGranularity())) + return fail("requires concrete b8/b16/b32 source and result " + "granularities"); + + FailureOr sourceArity = getVMIPhysicalArity(sourceType); + FailureOr resultArity = getVMIPhysicalArity(resultType); + if (failed(sourceArity) || failed(resultArity)) + return fail("requires computable source/result physical arity"); + if (*sourceArity < 1 || *resultArity < 1) + return fail("requires non-empty source/result physical arity"); + + return success(); +} + +FailureOr> materializeAdjacentMaskGranularityConversion( + Operation *op, VMIMaskType sourceType, VMIMaskType resultType, + ValueRange sourceParts, PatternRewriter &rewriter) { + auto fail = [&](const Twine &message) -> FailureOr> { + (void)rewriter.notifyMatchFailure(op, message); + return failure(); + }; + + int sourceRank = getMaskGranularityRank(sourceType.getGranularity()); + int resultRank = getMaskGranularityRank(resultType.getGranularity()); + if (std::abs(sourceRank - resultRank) != 1) + return fail("mask granularity conversion must be adjacent"); + + FailureOr sourceArity = getVMIPhysicalArity(sourceType); + FailureOr factor = getVMITypeLayoutFactor(sourceType); + if (failed(sourceArity) || failed(factor) || + static_cast(sourceParts.size()) != *sourceArity) + return fail("source mask part count does not match source VMI type"); + + MLIRContext *ctx = op->getContext(); + auto partAttr = [&](StringRef part) { return StringAttr::get(ctx, part); }; + auto resultMaskType = MaskType::get(ctx, resultType.getGranularity()); + SmallVector results; + + int64_t sourceOffset = 0; + for (int64_t part = 0; part < *factor; ++part) { + FailureOr sourceChunks = getVMITypeChunksInPart(sourceType, part); + FailureOr resultChunks = getVMITypeChunksInPart(resultType, part); + if (failed(sourceChunks) || failed(resultChunks)) + return fail("requires computable source/result chunks per layout part"); + + if (resultRank > sourceRank) { + int64_t produced = 0; + for (int64_t chunk = 0; chunk < *sourceChunks && produced < *resultChunks; + ++chunk) { + Value source = sourceParts[sourceOffset + chunk]; + results.push_back(rewriter + .create(op->getLoc(), resultMaskType, + source, partAttr("LOWER")) + .getResult()); + ++produced; + if (produced >= *resultChunks) + break; + results.push_back(rewriter + .create(op->getLoc(), resultMaskType, + source, partAttr("HIGHER")) + .getResult()); + ++produced; + } + if (produced != *resultChunks) + return fail("widening mask granularity conversion produced the wrong " + "number of result chunks"); + } else { + Value allTrue; + int64_t consumed = 0; + for (int64_t chunk = 0; chunk < *resultChunks; ++chunk) { + if (consumed >= *sourceChunks) + return fail("narrowing mask granularity conversion ran out of " + "source chunks"); + Value lowerSource = sourceParts[sourceOffset + consumed++]; + Value packed = rewriter + .create(op->getLoc(), resultMaskType, + lowerSource, partAttr("LOWER")) + .getResult(); + if (consumed < *sourceChunks) { + Value higherSource = sourceParts[sourceOffset + consumed++]; + Value higher = rewriter + .create(op->getLoc(), resultMaskType, + higherSource, partAttr("HIGHER")) + .getResult(); + if (!allTrue) { + FailureOr mask = + createAllTrueMask(op->getLoc(), resultMaskType, rewriter); + if (failed(mask)) + return fail("failed to create all-true mask for ppack merge"); + allTrue = *mask; + } + packed = rewriter + .create(op->getLoc(), resultMaskType, packed, + higher, allTrue) + .getResult(); + } + results.push_back(packed); + } + if (consumed != *sourceChunks) + return fail("narrowing mask granularity conversion left unused source " + "chunks"); + } + + sourceOffset += *sourceChunks; + } + + FailureOr resultArity = getVMIPhysicalArity(resultType); + if (failed(resultArity) || + static_cast(results.size()) != *resultArity) + return fail("mask granularity conversion result count mismatch"); + return results; +} + +FailureOr> materializeMaskGranularityConversion( + Operation *op, VMIMaskType sourceType, VMIMaskType resultType, ValueRange sourceParts, + PatternRewriter &rewriter) { + std::string reason; + if (failed(checkSupportedMaskGranularityMaterialization(sourceType, resultType, &reason))) { + (void)rewriter.notifyMatchFailure(op, reason); + return failure(); + } + + int currentRank = getMaskGranularityRank(sourceType.getGranularity()); + int resultRank = getMaskGranularityRank(resultType.getGranularity()); + if (std::abs(currentRank - resultRank) == 1) + return materializeAdjacentMaskGranularityConversion( + op, sourceType, resultType, sourceParts, rewriter); + + VMIMaskType currentType = sourceType; + SmallVector currentParts(sourceParts.begin(), sourceParts.end()); + while (currentRank != resultRank) { + currentRank += currentRank < resultRank ? 1 : -1; + StringRef nextGranularity = getMaskGranularityForRank(currentRank); + if (nextGranularity.empty()) { + (void)rewriter.notifyMatchFailure(op, + "invalid target mask granularity rank"); + return failure(); + } + VMIMaskType nextType = + VMIMaskType::get(op->getContext(), currentType.getElementCount(), + nextGranularity, currentType.getLayoutAttr()); + FailureOr> nextParts = + materializeAdjacentMaskGranularityConversion(op, currentType, nextType, + currentParts, rewriter); + if (failed(nextParts)) + return failure(); + currentType = nextType; + currentParts = std::move(*nextParts); + } + + return currentParts; +} + +FailureOr> getConvertedMaskPartTypes(VMIMaskType type) { + FailureOr arity = getVMIPhysicalArity(type); + FailureOr physicalGranularity = + getVMIMaskPhysicalGranularity(type); + if (failed(arity) || failed(physicalGranularity) || *arity < 0) + return failure(); + SmallVector types; + types.reserve(*arity); + Type partType = MaskType::get(type.getContext(), *physicalGranularity); + for (int64_t i = 0; i < *arity; ++i) + types.push_back(partType); + return types; +} + +static FailureOr +getVMIMaskPhysicalCarrierLayout(VMIMaskType type) { + VMILayoutAttr layout = type.getLayoutAttr(); + if (!layout) + return failure(); + MLIRContext *ctx = type.getContext(); + if (layout.isContiguous()) + return VMILayoutAttr::getContiguous(ctx); + if (layout.isDeinterleaved()) + return VMILayoutAttr::getDeinterleaved(ctx, layout.getFactor(), + layout.getBlockElems()); + if (layout.isGroupSlots()) + return VMILayoutAttr::getGroupSlots(ctx, layout.getNumGroups(), + layout.getSlots()); + return failure(); +} + +static FailureOr +getVMIMaskPhysicalCarrierType(VMIMaskType type) { + FailureOr physicalGranularity = + getVMIMaskPhysicalGranularity(type); + FailureOr physicalLayout = + getVMIMaskPhysicalCarrierLayout(type); + if (failed(physicalGranularity) || failed(physicalLayout)) + return failure(); + return VMIMaskType::get(type.getContext(), type.getElementCount(), + *physicalGranularity, *physicalLayout); +} + +static bool isElementDeinterleavedLayout(VMILayoutAttr layout, + int64_t factor) { + return layout && layout.isDeinterleaved() && layout.getFactor() == factor && + layout.getLaneStride() == 1 && layout.getBlockElems() == 1; +} + +FailureOr createAllFalseMaskLike(Location loc, Value value, + PatternRewriter &rewriter) { + auto maskType = dyn_cast(value.getType()); + if (!maskType) + return failure(); + return createPrefixMask(loc, maskType, "PAT_ALLF", rewriter); +} + +FailureOr> materializeStagingDeintToContiguousMaskLayout( + Operation *op, ValueRange sourceParts, TypeRange resultTypes, + int64_t factor, PatternRewriter &rewriter) { + auto fail = [&](const Twine &message) -> FailureOr> { + (void)rewriter.notifyMatchFailure(op, message); + return failure(); + }; + if ((factor != 2 && factor != 4) || sourceParts.empty() || + sourceParts.size() % factor != 0) + return fail("staging deinterleaved mask layout requires grouped source " + "parts"); + + int64_t groups = sourceParts.size() / factor; + SmallVector results; + results.reserve(resultTypes.size()); + for (int64_t i = 0; i < groups && results.size() < resultTypes.size(); ++i) { + auto nextType = [&](int64_t offset) -> Type { + size_t index = results.size() + offset; + return index < resultTypes.size() ? resultTypes[index] + : resultTypes[results.size()]; + }; + if (factor == 2) { + FailureOr> materialized = createPredicateIntlv( + op->getLoc(), nextType(0), nextType(1), sourceParts[i], + sourceParts[groups + i], rewriter); + if (failed(materialized)) + return fail("unsupported predicate intlv staging mask type"); + results.push_back(materialized->first); + if (results.size() < resultTypes.size()) + results.push_back(materialized->second); + continue; + } + + Value p0 = sourceParts[i]; + Value p1 = sourceParts[groups + i]; + Value p2 = sourceParts[2 * groups + i]; + Value p3 = sourceParts[3 * groups + i]; + FailureOr> even = + createPredicateIntlv(op->getLoc(), nextType(0), nextType(1), p0, p2, + rewriter); + FailureOr> odd = + createPredicateIntlv(op->getLoc(), nextType(0), nextType(1), p1, p3, + rewriter); + if (failed(even) || failed(odd)) + return fail("unsupported predicate intlv staging mask type"); + FailureOr> low = createPredicateIntlv( + op->getLoc(), nextType(0), nextType(1), even->first, odd->first, + rewriter); + FailureOr> high = createPredicateIntlv( + op->getLoc(), nextType(2), nextType(3), even->second, odd->second, + rewriter); + if (failed(low) || failed(high)) + return fail("unsupported predicate intlv staging mask type"); + results.push_back(low->first); + if (results.size() < resultTypes.size()) + results.push_back(low->second); + if (results.size() < resultTypes.size()) + results.push_back(high->first); + if (results.size() < resultTypes.size()) + results.push_back(high->second); + } + if (results.size() != resultTypes.size()) + return fail("staging deinterleaved mask layout result arity mismatch"); + return results; +} + +FailureOr> materializeStagingContiguousToDeintMaskLayout( + Operation *op, ValueRange sourceParts, TypeRange resultTypes, + int64_t factor, PatternRewriter &rewriter) { + auto fail = [&](const Twine &message) -> FailureOr> { + (void)rewriter.notifyMatchFailure(op, message); + return failure(); + }; + if ((factor != 2 && factor != 4) || sourceParts.empty() || + resultTypes.size() % factor != 0) + return fail("staging contiguous mask layout requires grouped result parts"); + + int64_t groups = resultTypes.size() / factor; + if (sourceParts.size() > static_cast(groups * factor)) + return fail("staging contiguous mask layout has too many source parts"); + + SmallVector, 4> parts(factor); + for (int64_t part = 0; part < factor; ++part) + parts[part].reserve(groups); + + for (int64_t i = 0; i < groups; ++i) { + size_t sourceBase = static_cast(i * factor); + if (sourceBase >= sourceParts.size()) + return fail("staging contiguous mask layout ran out of source parts"); + + SmallVector sources; + sources.reserve(factor); + for (int64_t lane = 0; lane < factor; ++lane) { + size_t index = sourceBase + lane; + if (index < sourceParts.size()) { + sources.push_back(sourceParts[index]); + continue; + } + FailureOr zero = + createAllFalseMaskLike(op->getLoc(), sourceParts[sourceBase], + rewriter); + if (failed(zero)) + return fail("failed to create all-false staging mask"); + sources.push_back(*zero); + } + + if (factor == 2) { + FailureOr> materialized = + createPredicateDintlv(op->getLoc(), resultTypes[i], + resultTypes[groups + i], sources[0], + sources[1], rewriter); + if (failed(materialized)) + return fail("unsupported predicate dintlv staging mask type"); + parts[0].push_back(materialized->first); + parts[1].push_back(materialized->second); + continue; + } + + FailureOr> low = createPredicateDintlv( + op->getLoc(), resultTypes[i], resultTypes[groups + i], sources[0], + sources[1], rewriter); + FailureOr> high = createPredicateDintlv( + op->getLoc(), resultTypes[2 * groups + i], + resultTypes[3 * groups + i], sources[2], sources[3], rewriter); + if (failed(low) || failed(high)) + return fail("unsupported predicate dintlv staging mask type"); + FailureOr> even = createPredicateDintlv( + op->getLoc(), resultTypes[i], resultTypes[2 * groups + i], low->first, + high->first, rewriter); + FailureOr> odd = createPredicateDintlv( + op->getLoc(), resultTypes[groups + i], resultTypes[3 * groups + i], + low->second, high->second, rewriter); + if (failed(even) || failed(odd)) + return fail("unsupported predicate dintlv staging mask type"); + parts[0].push_back(even->first); + parts[1].push_back(odd->first); + parts[2].push_back(even->second); + parts[3].push_back(odd->second); + } + + SmallVector results; + results.reserve(resultTypes.size()); + for (int64_t part = 0; part < factor; ++part) { + if (parts[part].size() != static_cast(groups)) + return fail("staging contiguous mask layout result arity mismatch"); + results.append(parts[part]); + } + return results; +} + +FailureOr> materializeMaskGranularityCastLayoutConversion( + Operation *op, VMIMaskType sourceType, VMIMaskType resultType, + ValueRange sourceParts, TypeRange resultTypes, PatternRewriter &rewriter); + +FailureOr> +materializeMaskGranularityCastLayoutConversionViaContiguous( + Operation *op, VMIMaskType sourceType, VMIMaskType resultType, + ValueRange sourceParts, TypeRange resultTypes, PatternRewriter &rewriter) { + VMILayoutAttr contiguous = VMILayoutAttr::getContiguous(op->getContext()); + VMIMaskType contiguousType = + VMIMaskType::get(op->getContext(), sourceType.getElementCount(), + sourceType.getGranularity(), contiguous); + FailureOr> contiguousTypes = + getConvertedMaskPartTypes(contiguousType); + if (failed(contiguousTypes)) + return failure(); + FailureOr> contiguousParts = + materializeMaskGranularityCastLayoutConversion( + op, sourceType, contiguousType, sourceParts, *contiguousTypes, + rewriter); + if (failed(contiguousParts)) + return failure(); + return materializeMaskGranularityCastLayoutConversion( + op, contiguousType, resultType, *contiguousParts, resultTypes, rewriter); +} + +FailureOr> materializeMaskGranularityCastLayoutConversion( + Operation *op, VMIMaskType sourceType, VMIMaskType resultType, + ValueRange sourceParts, TypeRange resultTypes, PatternRewriter &rewriter) { + auto fail = [&](const Twine &message) -> FailureOr> { + (void)rewriter.notifyMatchFailure(op, message); + return failure(); + }; + + VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + if (!sourceLayout || !resultLayout) + return fail("mask granularity cast layout conversion requires layouts"); + + if (sourceLayout == resultLayout) { + if (failed(verifyIdentityPartForwarding(op, sourceParts, resultTypes, + rewriter))) + return failure(); + return SmallVector(sourceParts.begin(), sourceParts.end()); + } + + FailureOr> layoutParts = materializeMaskLayoutConversion( + op, sourceParts, resultTypes, sourceLayout, resultLayout, rewriter); + if (succeeded(layoutParts)) + return layoutParts; + + bool sourceC = sourceLayout.isContiguous() && sourceLayout.getLaneStride() == 1; + bool resultC = resultLayout.isContiguous() && resultLayout.getLaneStride() == 1; + if (isElementDeinterleavedLayout(sourceLayout, 2) && resultC) + return materializeStagingDeintToContiguousMaskLayout( + op, sourceParts, resultTypes, /*factor=*/2, rewriter); + if (sourceC && isElementDeinterleavedLayout(resultLayout, 2)) + return materializeStagingContiguousToDeintMaskLayout( + op, sourceParts, resultTypes, /*factor=*/2, rewriter); + if (isElementDeinterleavedLayout(sourceLayout, 4) && resultC) + return materializeStagingDeintToContiguousMaskLayout( + op, sourceParts, resultTypes, /*factor=*/4, rewriter); + if (sourceC && isElementDeinterleavedLayout(resultLayout, 4)) + return materializeStagingContiguousToDeintMaskLayout( + op, sourceParts, resultTypes, /*factor=*/4, rewriter); + + if (sourceLayout.isDeinterleaved() || resultLayout.isDeinterleaved()) + return materializeMaskGranularityCastLayoutConversionViaContiguous( + op, sourceType, resultType, sourceParts, resultTypes, rewriter); + + return fail("unsupported mask granularity cast layout conversion"); +} + +FailureOr> materializeMaskGranularityCastConversion( + Operation *op, VMIMaskType sourceType, VMIMaskType resultType, + ValueRange sourceParts, TypeRange resultTypes, PatternRewriter &rewriter) { + auto fail = [&](const Twine &message) -> FailureOr> { + (void)rewriter.notifyMatchFailure(op, message); + return failure(); + }; + + if (sourceType.getElementCount() != resultType.getElementCount()) + return fail("requires source and result mask lane counts to match"); + + FailureOr physicalSourceType = + getVMIMaskPhysicalCarrierType(sourceType); + FailureOr physicalResultType = + getVMIMaskPhysicalCarrierType(resultType); + if (failed(physicalSourceType) || failed(physicalResultType)) + return fail("requires source/result mask physical carrier types"); + + if (*physicalSourceType == *physicalResultType) { + if (failed(verifyIdentityPartForwarding(op, sourceParts, resultTypes, + rewriter))) + return failure(); + return SmallVector(sourceParts.begin(), sourceParts.end()); + } + + if (physicalSourceType->getLayoutAttr() == physicalResultType->getLayoutAttr()) + return materializeMaskGranularityConversion(op, *physicalSourceType, + *physicalResultType, + sourceParts, rewriter); + + VMIMaskType granularityType = + VMIMaskType::get(op->getContext(), sourceType.getElementCount(), + physicalResultType->getGranularity(), + physicalSourceType->getLayoutAttr()); + FailureOr> granularityParts = + materializeMaskGranularityConversion(op, *physicalSourceType, + granularityType, sourceParts, + rewriter); + if (failed(granularityParts)) + return failure(); + return materializeMaskGranularityCastLayoutConversion( + op, granularityType, *physicalResultType, *granularityParts, resultTypes, + rewriter); +} + +struct OneToNVMIEnsureLayoutOpPattern + : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIEnsureLayoutOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto sourceType = cast(op.getSource().getType()); + auto resultType = cast(op.getResult().getType()); + VMILayoutSupport supports; + std::string supportReason; + if (failed(supports.getEnsureLayoutFact(sourceType, resultType, + &supportReason))) + return rewriter.notifyMatchFailure( + op, + Twine("ensure_layout has no registered materialization support: ") + + supportReason); + VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + if (!sourceLayout || !resultLayout) + return rewriter.notifyMatchFailure( + op, "ensure_layout requires assigned source/result layouts"); + + ValueRange sourceParts = adaptor.getSource(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + FailureOr> results = materializeDataLayoutConversion( + op, sourceParts, resultTypes, sourceLayout, resultLayout, rewriter); + if (failed(results)) + return failure(); + replaceOpWithFlatConvertedValues(rewriter, op, *results, *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMIEnsureMaskLayoutOpPattern + : OpConversionPattern { + using OpConversionPattern< + VMIEnsureMaskLayoutOp>::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIEnsureMaskLayoutOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto sourceType = cast(op.getSource().getType()); + auto resultType = cast(op.getResult().getType()); + VMILayoutSupport supports; + std::string supportReason; + if (failed(supports.getEnsureMaskLayoutFact(sourceType, resultType, + &supportReason))) + return rewriter.notifyMatchFailure( + op, Twine("ensure_mask_layout has no registered materialization " + "support: ") + + supportReason); + if (sourceType.getGranularity() != resultType.getGranularity()) + return rewriter.notifyMatchFailure( + op, "mask layout helper cannot also change granularity"); + VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + + ValueRange sourceParts = adaptor.getSource(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + FailureOr> results = materializeMaskLayoutConversion( + op, sourceParts, resultTypes, sourceLayout, resultLayout, rewriter); + if (failed(results)) + return failure(); + replaceOpWithFlatConvertedValues(rewriter, op, *results, *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMIEnsureMaskGranularityOpPattern + : OpConversionPattern { + using OpConversionPattern< + VMIEnsureMaskGranularityOp>::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIEnsureMaskGranularityOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto sourceType = cast(op.getSource().getType()); + auto resultType = cast(op.getResult().getType()); + VMILayoutSupport supports; + bool identity = sourceType.getGranularity() == resultType.getGranularity() && + sourceType.getLayoutAttr() == resultType.getLayoutAttr(); + if (!identity) { + std::string reason; + if (failed(supports.getMaskGranularityCastLayoutFactForLayouts( + sourceType, resultType, sourceType.getLayoutAttr(), + resultType.getLayoutAttr(), &reason))) + return rewriter.notifyMatchFailure( + op, "unsupported mask granularity cast layout relation: " + reason); + } + + ValueRange sourceParts = adaptor.getSource(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + + FailureOr> results = + materializeMaskGranularityCastConversion( + op, sourceType, resultType, sourceParts, resultTypes, rewriter); + if (failed(results)) + return failure(); + if (results->size() != resultTypes.size()) + return rewriter.notifyMatchFailure( + op, "mask granularity cast result arity mismatch"); + for (auto [result, type] : llvm::zip_equal(*results, resultTypes)) + if (result.getType() != type) + return rewriter.notifyMatchFailure( + op, "mask granularity cast result type mismatch"); + replaceOpWithFlatConvertedValues(rewriter, op, *results, + *this->getTypeConverter()); + return success(); + } + +private: + ; +}; + +struct OneToNVMIBroadcastOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIBroadcastOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + ValueRange inputParts = adaptor.getValue(); + if (inputParts.size() != 1) + return rewriter.notifyMatchFailure( + op, "broadcast input must convert to one value"); + bool inputIsVReg = isa(op.getValue().getType()); + + FailureOr> maybe_resultTypes = + + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + + if (failed(maybe_resultTypes)) + + return failure(); + + SmallVector resultTypes = std::move(*maybe_resultTypes); + SmallVector results; + results.reserve(resultTypes.size()); + for (Type resultType : resultTypes) { + auto vregType = dyn_cast(resultType); + if (!vregType) + return rewriter.notifyMatchFailure(op, "broadcast result must be vreg"); + FailureOr mask = + createAllTrueMaskForVReg(op.getLoc(), vregType, rewriter); + if (failed(mask)) + return rewriter.notifyMatchFailure( + op, "unsupported element type for broadcast mask"); + StringAttr position = + inputIsVReg ? rewriter.getStringAttr("LOWEST") : StringAttr{}; + results.push_back(rewriter + .create(op.getLoc(), resultType, + inputParts.front(), *mask, position) + .getResult()); + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +FailureOr createScalarOffsetConstant(Location loc, Type type, + int64_t value, + PatternRewriter &rewriter) { + if (auto intType = dyn_cast(type)) { + return rewriter + .create(loc, IntegerAttr::get(intType, value)) + .getResult(); + } + if (auto floatType = dyn_cast(type)) { + return rewriter + .create( + loc, rewriter.getFloatAttr(floatType, static_cast(value))) + .getResult(); + } + return failure(); +} + +FailureOr createIotaChunkBase(Location loc, Value base, + int64_t laneOffset, StringRef order, + PatternRewriter &rewriter) { + if (laneOffset == 0) + return base; + + FailureOr offset = + createScalarOffsetConstant(loc, base.getType(), laneOffset, rewriter); + if (failed(offset)) + return failure(); + + if (isa(base.getType())) { + if (order == "DESC") + return rewriter.create(loc, base, *offset).getResult(); + return rewriter.create(loc, base, *offset).getResult(); + } + if (isa(base.getType())) { + if (order == "DESC") + return rewriter.create(loc, base, *offset).getResult(); + return rewriter.create(loc, base, *offset).getResult(); + } + + return failure(); +} + +FailureOr createIotaContiguousChunk(Location loc, Type resultType, + Value base, int64_t laneOffset, + StringAttr orderAttr, + PatternRewriter &rewriter) { + StringRef order = orderAttr ? orderAttr.getValue() : StringRef("ASC"); + FailureOr chunkBase = + createIotaChunkBase(loc, base, laneOffset, order, rewriter); + if (failed(chunkBase)) + return failure(); + return rewriter.create(loc, resultType, *chunkBase, orderAttr) + .getResult(); +} + +FailureOr createIotaDeinterleavedChunk(Location loc, Type resultType, + Value base, int64_t factor, + int64_t part, int64_t chunk, + int64_t lanesPerPart, + StringAttr orderAttr, + PatternRewriter &rewriter) { + auto vregType = dyn_cast(resultType); + if (!vregType) + return failure(); + + FailureOr mask = createAllTrueMaskForVReg(loc, vregType, rewriter); + FailureOr zero = + createScalarOffsetConstant(loc, base.getType(), 0, rewriter); + FailureOr factorScalar = + createScalarOffsetConstant(loc, base.getType(), factor, rewriter); + if (failed(mask) || failed(zero) || failed(factorScalar)) + return failure(); + + Value local = + rewriter.create(loc, resultType, *zero, StringAttr{}).getResult(); + Value scaled = + rewriter.create(loc, resultType, local, *factorScalar, *mask) + .getResult(); + + StringRef order = orderAttr ? orderAttr.getValue() : StringRef("ASC"); + int64_t partOffset = part + factor * chunk * lanesPerPart; + FailureOr biasedBase = + createIotaChunkBase(loc, base, partOffset, order, rewriter); + if (failed(biasedBase)) + return failure(); + + if (order == "DESC") { + Value baseVector = rewriter + .create(loc, resultType, *biasedBase, *mask, + /*position=*/nullptr) + .getResult(); + return rewriter.create(loc, resultType, baseVector, scaled, *mask) + .getResult(); + } + + return rewriter.create(loc, resultType, scaled, *biasedBase, *mask) + .getResult(); +} + +struct OneToNVMIIotaOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIIotaOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto resultVMIType = cast(op.getResult().getType()); + VMILayoutAttr layout = resultVMIType.getLayoutAttr(); + if (!layout) + return rewriter.notifyMatchFailure(op, "iota requires assigned layout"); + + FailureOr lanesPerPart = + getDataLanesPerPart(resultVMIType.getElementType()); + if (failed(lanesPerPart)) + return rewriter.notifyMatchFailure( + op, "iota requires known physical lanes per part"); + + FailureOr base = getSingleValue( + op, adaptor.getBase(), "iota base must convert to one value", rewriter); + if (failed(base)) + return failure(); + + FailureOr> maybe_resultTypes = + + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + + if (failed(maybe_resultTypes)) + + return failure(); + + SmallVector resultTypes = std::move(*maybe_resultTypes); + SmallVector results; + results.reserve(resultTypes.size()); + + if (layout.isContiguous()) { + for (auto [index, resultType] : llvm::enumerate(resultTypes)) { + if (!isa(resultType)) + return rewriter.notifyMatchFailure(op, "iota result must be vreg"); + FailureOr result = createIotaContiguousChunk( + op.getLoc(), resultType, *base, + static_cast(index) * *lanesPerPart, op.getOrderAttr(), + rewriter); + if (failed(result)) + return rewriter.notifyMatchFailure( + op, "failed to materialize contiguous iota chunk"); + results.push_back(*result); + } + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } + + int64_t factor = layout.getFactor(); + if (resultTypes.size() % factor != 0) + return rewriter.notifyMatchFailure( + op, "deinterleaved iota physical result count does not match " + "layout factor"); + int64_t chunksPerPart = resultTypes.size() / factor; + for (int64_t part = 0; part < factor; ++part) { + for (int64_t chunk = 0; chunk < chunksPerPart; ++chunk) { + Type resultType = resultTypes[part * chunksPerPart + chunk]; + FailureOr result = createIotaDeinterleavedChunk( + op.getLoc(), resultType, *base, factor, part, chunk, *lanesPerPart, + op.getOrderAttr(), rewriter); + if (failed(result)) + return rewriter.notifyMatchFailure( + op, "failed to materialize deinterleaved iota chunk"); + results.push_back(*result); + } + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMIConstantOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIConstantOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto denseAttr = dyn_cast(op.getValue()); + if (!denseAttr || !denseAttr.isSplat()) + return rewriter.notifyMatchFailure( + op, "only splat dense data constants are supported"); + auto splatAttr = dyn_cast(denseAttr.getSplatValue()); + if (!splatAttr) + return rewriter.notifyMatchFailure(op, "splat constant must be typed"); + + // arith.constant only accepts signless integer types, whereas VMI vregs may + // carry signed/unsigned element types (e.g. ui16). Remap an unsigned/signed + // integer splat to its signless equivalent; the downstream pto.vdup accepts + // a signless scalar for a signed/unsigned result element. + if (auto intAttr = dyn_cast(splatAttr)) { + if (auto intTy = dyn_cast(intAttr.getType()); + intTy && !intTy.isSignless()) + splatAttr = IntegerAttr::get(rewriter.getIntegerType(intTy.getWidth()), + intAttr.getValue()); + } + + Value scalar = + rewriter.create(op.getLoc(), splatAttr).getResult(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + SmallVector results; + results.reserve(resultTypes.size()); + for (Type resultType : resultTypes) { + auto vregType = dyn_cast(resultType); + if (!vregType) + return rewriter.notifyMatchFailure(op, "constant result must be vreg"); + FailureOr mask = + createAllTrueMaskForVReg(op.getLoc(), vregType, rewriter); + if (failed(mask)) + return rewriter.notifyMatchFailure( + op, "unsupported element type for constant mask"); + results.push_back(rewriter + .create(op.getLoc(), resultType, scalar, + *mask, + /*position=*/nullptr) + .getResult()); + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMIConstantMaskOpPattern + : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIConstantMaskOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + std::string reason; + FailureOr> materializations = + computeConstantMaskMaterialization(op, &reason); + if (failed(materializations)) + return rewriter.notifyMatchFailure(op, Twine("constant_mask ") + reason); + + SmallVector results; + results.reserve(resultTypes.size()); + for (const ConstantMaskChunkMaterialization &materialization : + *materializations) { + if (results.size() >= resultTypes.size()) + return rewriter.notifyMatchFailure( + op, "constant_mask produced too many physical masks"); + auto maskType = dyn_cast(resultTypes[results.size()]); + if (!maskType) + return rewriter.notifyMatchFailure(op, + "constant_mask result must be mask"); + FailureOr mask = materializeConstantMaskChunk( + op.getLoc(), maskType, materialization.activeLanes, rewriter); + if (failed(mask)) + return rewriter.notifyMatchFailure( + op, "failed to materialize constant_mask physical chunk"); + results.push_back(*mask); + } + + if (results.size() != resultTypes.size()) + return rewriter.notifyMatchFailure( + op, "constant_mask physical result count mismatch"); + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMICreateMaskOpPattern + : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMICreateMaskOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto activeConstant = + op.getActiveLanes().getDefiningOp(); + auto resultVMIType = cast(op.getResult().getType()); + VMILayoutAttr layout = resultVMIType.getLayoutAttr(); + if (!layout || + !VMIMaskType::isConcreteGranularity(resultVMIType.getGranularity())) + return rewriter.notifyMatchFailure( + op, "create_mask requires concrete layout and granularity"); + FailureOr physicalGranularity = + getVMIMaskPhysicalGranularity(resultVMIType); + FailureOr lanesPerPart = + failed(physicalGranularity) + ? FailureOr(failure()) + : getMaskLanesPerPart(*physicalGranularity); + if (failed(lanesPerPart)) + return rewriter.notifyMatchFailure( + op, "create_mask requires known physical mask lanes per part"); + + if (!activeConstant) { + FailureOr active = getSingleValue( + op, adaptor.getActiveLanes(), + "create_mask active_lanes must convert to one value", rewriter); + if (failed(active)) + return failure(); + + FailureOr> maybe_resultTypes = + + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + + if (failed(maybe_resultTypes)) + + return failure(); + + SmallVector resultTypes = std::move(*maybe_resultTypes); + int64_t factor = layout.isDeinterleaved() ? layout.getFactor() : 1; + if (resultTypes.size() % factor != 0) + return rewriter.notifyMatchFailure( + op, "dynamic create_mask physical result count does not match " + "layout factor"); + int64_t chunksPerPart = resultTypes.size() / factor; + Value activeI32 = clampDynamicActiveLanes( + op.getLoc(), *active, resultVMIType.getElementCount(), rewriter); + + SmallVector results; + results.reserve(resultTypes.size()); + for (int64_t part = 0; part < factor; ++part) { + Value remaining = createPartitionActiveLanes(op.getLoc(), activeI32, + factor, part, rewriter); + for (int64_t chunk = 0; chunk < chunksPerPart; ++chunk) { + Type resultType = resultTypes[part * chunksPerPart + chunk]; + auto maskType = dyn_cast(resultType); + if (!maskType) + return rewriter.notifyMatchFailure( + op, "create_mask result must be mask"); + FailureOr> maskAndRemaining = + createRuntimePrefixMask(op.getLoc(), maskType, remaining, + rewriter); + if (failed(maskAndRemaining)) + return rewriter.notifyMatchFailure( + op, "unsupported mask type for dynamic create_mask"); + results.push_back(maskAndRemaining->first); + remaining = maskAndRemaining->second; + } + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } + + auto activeAttr = dyn_cast(activeConstant.getValue()); + if (!activeAttr) + return rewriter.notifyMatchFailure( + op, "create_mask active_lanes must be an integer constant"); + + int64_t activeLanes = activeAttr.getInt(); + if (activeLanes < 0) + activeLanes = 0; + if (activeLanes > resultVMIType.getElementCount()) + activeLanes = resultVMIType.getElementCount(); + + FailureOr> maybe_resultTypes = + + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + + if (failed(maybe_resultTypes)) + + return failure(); + + SmallVector resultTypes = std::move(*maybe_resultTypes); + int64_t factor = layout.isDeinterleaved() ? layout.getFactor() : 1; + SmallVector results; + results.reserve(resultTypes.size()); + + for (int64_t part = 0; part < factor; ++part) { + for (int64_t chunk = 0;; ++chunk) { + bool anyLane = false; + int64_t activeInChunk = 0; + for (int64_t lane = 0; lane < *lanesPerPart; ++lane) { + FailureOr padding = + isPaddingLane(resultVMIType, part, chunk, lane); + if (failed(padding)) + return rewriter.notifyMatchFailure( + op, "failed to map create_mask physical padding lane"); + if (*padding) + continue; + anyLane = true; + FailureOr logicalLane = + mapPhysicalLaneToLogical(resultVMIType, part, chunk, lane); + if (failed(logicalLane)) + return rewriter.notifyMatchFailure( + op, "failed to map create_mask physical lane"); + if (*logicalLane < activeLanes) + ++activeInChunk; + } + if (!anyLane) + break; + + if (results.size() >= resultTypes.size()) + return rewriter.notifyMatchFailure( + op, "create_mask produced too many physical masks"); + auto maskType = dyn_cast(resultTypes[results.size()]); + if (!maskType) + return rewriter.notifyMatchFailure(op, + "create_mask result must be mask"); + std::optional pattern = + getPrefixPattern(activeInChunk, *lanesPerPart); + if (pattern) { + FailureOr mask = + createPrefixMask(op.getLoc(), maskType, *pattern, rewriter); + if (failed(mask)) + return rewriter.notifyMatchFailure( + op, "unsupported mask type for create_mask"); + results.push_back(*mask); + continue; + } + + FailureOr> maskAndRemaining = + createRuntimePrefixMask( + op.getLoc(), maskType, + createI32Constant(op.getLoc(), activeInChunk, rewriter), + rewriter); + if (failed(maskAndRemaining)) + return rewriter.notifyMatchFailure( + op, "unsupported mask type for create_mask plt fallback"); + results.push_back(maskAndRemaining->first); + } + } + + if (results.size() != resultTypes.size()) + return rewriter.notifyMatchFailure( + op, "create_mask physical result count mismatch"); + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMICreateGroupMaskOpPattern + : OpConversionPattern { + using OpConversionPattern< + VMICreateGroupMaskOp>::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMICreateGroupMaskOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + auto resultVMIType = cast(op.getResult().getType()); + VMILayoutAttr resultLayout = resultVMIType.getLayoutAttr(); + if (resultLayout && resultLayout.isDeinterleaved() && + resultLayout.getFactor() == 4 && resultLayout.getBlockElems() == 8) { + VMILayoutAttr contiguousLayout = + VMILayoutAttr::getContiguous(op.getContext()); + auto contiguousType = + VMIMaskType::get(op.getContext(), resultVMIType.getElementCount(), + resultVMIType.getGranularity(), contiguousLayout); + SmallVector contiguousParts; + auto activeConstant = + op.getActiveElemsPerGroup().getDefiningOp(); + if (activeConstant) { + std::string contiguousReason; + FailureOr> + contiguousMaterializations = computeGroupMaskMaterializationForType( + op, contiguousType, &contiguousReason); + if (failed(contiguousMaterializations)) + return rewriter.notifyMatchFailure(op, Twine("create_group_mask ") + + contiguousReason); + + contiguousParts.reserve(contiguousMaterializations->size()); + for (const ConstantMaskChunkMaterialization &materialization : + *contiguousMaterializations) { + if (contiguousParts.size() >= resultTypes.size()) + return rewriter.notifyMatchFailure( + op, "create_group_mask produced too many contiguous masks"); + auto maskType = + dyn_cast(resultTypes[contiguousParts.size()]); + if (!maskType) + return rewriter.notifyMatchFailure( + op, "create_group_mask result must be mask"); + FailureOr mask = materializeConstantMaskChunk( + op.getLoc(), maskType, materialization.activeLanes, rewriter); + if (failed(mask)) + return rewriter.notifyMatchFailure( + op, "failed to materialize create_group_mask contiguous chunk"); + contiguousParts.push_back(*mask); + } + } else { + FailureOr active = getSingleValue( + op, adaptor.getActiveElemsPerGroup(), + "create_group_mask active_elems_per_group must convert to one " + "value", + rewriter); + if (failed(active)) + return failure(); + FailureOr> dynamicParts = + materializeDynamicGroupMaskForType(op, *active, contiguousType, + resultTypes, rewriter); + if (failed(dynamicParts)) + return failure(); + contiguousParts = std::move(*dynamicParts); + } + + if (contiguousParts.size() != resultTypes.size()) + return rewriter.notifyMatchFailure( + op, "create_group_mask contiguous physical result count mismatch"); + FailureOr> results = materializeMaskLayoutConversion( + op, contiguousParts, resultTypes, contiguousLayout, resultLayout, + rewriter); + if (failed(results)) + return failure(); + replaceOpWithFlatConvertedValues(rewriter, op, *results, *this->getTypeConverter()); + return success(); + } + + auto activeConstant = + op.getActiveElemsPerGroup().getDefiningOp(); + if (!activeConstant) { + FailureOr active = getSingleValue( + op, adaptor.getActiveElemsPerGroup(), + "create_group_mask active_elems_per_group must convert to one value", + rewriter); + if (failed(active)) + return failure(); + + VMILayoutAttr resultLayout = resultVMIType.getLayoutAttr(); + if (resultLayout && resultLayout.isDeinterleaved() && + resultLayout.getBlockElems() == 1) { + VMILayoutAttr contiguousLayout = + VMILayoutAttr::getContiguous(op.getContext()); + auto contiguousType = + VMIMaskType::get(op.getContext(), resultVMIType.getElementCount(), + resultVMIType.getGranularity(), contiguousLayout); + FailureOr> contiguousParts = + materializeDynamicGroupMaskForType(op, *active, contiguousType, + resultTypes, rewriter); + if (failed(contiguousParts)) + return failure(); + FailureOr> results = materializeMaskLayoutConversion( + op, *contiguousParts, resultTypes, contiguousLayout, resultLayout, + rewriter); + if (failed(results)) + return failure(); + replaceOpWithFlatConvertedValues(rewriter, op, *results, + *this->getTypeConverter()); + return success(); + } + + FailureOr> results = + materializeDynamicGroupMaskForType(op, *active, resultVMIType, + resultTypes, rewriter); + if (failed(results)) + return failure(); + replaceOpWithFlatConvertedValues(rewriter, op, *results, *this->getTypeConverter()); + return success(); + } + + std::string reason; + FailureOr> materializations = + computeGroupMaskMaterialization(op, &reason); + if (failed(materializations)) + return rewriter.notifyMatchFailure(op, + Twine("create_group_mask ") + reason); + + SmallVector results; + results.reserve(resultTypes.size()); + for (const ConstantMaskChunkMaterialization &materialization : + *materializations) { + if (results.size() >= resultTypes.size()) + return rewriter.notifyMatchFailure( + op, "create_group_mask produced too many physical masks"); + auto maskType = dyn_cast(resultTypes[results.size()]); + if (!maskType) + return rewriter.notifyMatchFailure( + op, "create_group_mask result must be mask"); + FailureOr mask = materializeConstantMaskChunk( + op.getLoc(), maskType, materialization.activeLanes, rewriter); + if (failed(mask)) + return rewriter.notifyMatchFailure( + op, "failed to materialize create_group_mask physical chunk"); + results.push_back(*mask); + } + + if (results.size() != resultTypes.size()) + return rewriter.notifyMatchFailure( + op, "create_group_mask physical result count mismatch"); + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMILoadOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMILoadOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto resultVMIType = cast(op.getResult().getType()); + FailureOr source = + getSingleValue(op, adaptor.getSource(), + "load source must convert to one value", rewriter); + FailureOr offset = + getSingleValue(op, adaptor.getOffset(), + "load offset must convert to one value", rewriter); + if (failed(source) || failed(offset)) + return failure(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + VMILayoutAttr resultLayout = resultVMIType.getLayoutAttr(); + if (std::optional dist = + getDenseLaneStrideLoadDistToken(resultVMIType)) { + SmallVector results; + results.reserve(resultTypes.size()); + int64_t semanticOffset = 0; + for (auto [index, resultType] : llvm::enumerate(resultTypes)) { + if (!isa(resultType)) + return rewriter.notifyMatchFailure(op, "load result must be vreg"); + Value chunkOffset = + createChunkOffset(op.getLoc(), *offset, semanticOffset, rewriter); + results.push_back(rewriter + .create(op.getLoc(), resultType, + /*updated_base=*/Type{}, *source, + chunkOffset, + rewriter.getStringAttr(*dist)) + .getResult()); + FailureOr activeLanes = + getActiveDataLanesInPhysicalChunk(resultVMIType, index); + if (failed(activeLanes)) + return rewriter.notifyMatchFailure( + op, "failed to compute lane_stride load active lanes"); + semanticOffset += *activeLanes; + } + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } + + FailureOr lanesPerPart = verifyFullOrSafeReadVRegChunks( + op, resultVMIType, op.getSource().getType(), *offset, rewriter); + if (failed(lanesPerPart)) + return failure(); + + VMILayoutAttr contiguousLayout = + VMILayoutAttr::getContiguous(rewriter.getContext()); + FailureOr> maybeContiguousTypes = + getConvertedVRegTypesWithLayout(resultVMIType, contiguousLayout, + *this->getTypeConverter()); + if (failed(maybeContiguousTypes)) + return rewriter.notifyMatchFailure( + op, "failed to compute contiguous load footprint"); + SmallVector contiguousTypes = std::move(*maybeContiguousTypes); + FailureOr noWiderThanContiguous = + hasNoWiderFootprintThanContiguous(resultTypes, contiguousTypes); + if (failed(noWiderThanContiguous)) + return rewriter.notifyMatchFailure( + op, "failed to compare load physical footprint"); + + if (resultLayout && resultLayout.isDeinterleaved() && + resultLayout.getFactor() == 2 && *noWiderThanContiguous) { + std::optional dist = + getX2MemoryDistToken(resultVMIType.getElementType(), "DINTLV"); + if (dist && !resultTypes.empty() && resultTypes.size() % 2 == 0) { + int64_t groups = resultTypes.size() / 2; + SmallVector lows; + SmallVector highs; + lows.reserve(groups); + highs.reserve(groups); + for (int64_t group = 0; group < groups; ++group) { + Type lowType = resultTypes[group]; + Type highType = resultTypes[groups + group]; + if (lowType != highType) + return rewriter.notifyMatchFailure( + op, "vldsx2 requires matching low/high result types"); + Value chunkOffset = createChunkOffset( + op.getLoc(), *offset, group * 2 * *lanesPerPart, rewriter); + auto load = rewriter.create(op.getLoc(), lowType, highType, + *source, chunkOffset, + rewriter.getStringAttr(*dist)); + lows.push_back(load.getLow()); + highs.push_back(load.getHigh()); + } + SmallVector results; + results.reserve(resultTypes.size()); + results.append(lows); + results.append(highs); + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } + } + + if (resultLayout && resultLayout.isDeinterleaved() && + resultLayout.getFactor() == 4 && resultLayout.getBlockElems() == 1 && + *noWiderThanContiguous) { + std::optional dist = + getX2MemoryDistToken(resultVMIType.getElementType(), "DINTLV"); + if (dist && !resultTypes.empty() && resultTypes.size() % 4 == 0) { + int64_t groups = resultTypes.size() / 4; + SmallVector part0; + SmallVector part1; + SmallVector part2; + SmallVector part3; + part0.reserve(groups); + part1.reserve(groups); + part2.reserve(groups); + part3.reserve(groups); + for (int64_t group = 0; group < groups; ++group) { + Type part0Type = resultTypes[group]; + Type part1Type = resultTypes[groups + group]; + Type part2Type = resultTypes[2 * groups + group]; + Type part3Type = resultTypes[3 * groups + group]; + if (part0Type != part1Type || part0Type != part2Type || + part0Type != part3Type) + return rewriter.notifyMatchFailure( + op, "vldsx2 deinterleaved=4 load requires matching part " + "types"); + + Value firstOffset = createChunkOffset( + op.getLoc(), *offset, group * 4 * *lanesPerPart, rewriter); + Value secondOffset = createChunkOffset( + op.getLoc(), *offset, (group * 4 + 2) * *lanesPerPart, rewriter); + auto first = rewriter.create( + op.getLoc(), part0Type, part1Type, *source, firstOffset, + rewriter.getStringAttr(*dist)); + auto second = rewriter.create( + op.getLoc(), part2Type, part3Type, *source, secondOffset, + rewriter.getStringAttr(*dist)); + + auto even = + rewriter.create(op.getLoc(), part0Type, part2Type, + first.getLow(), second.getLow()); + auto odd = + rewriter.create(op.getLoc(), part1Type, part3Type, + first.getHigh(), second.getHigh()); + part0.push_back(even.getLow()); + part1.push_back(odd.getLow()); + part2.push_back(even.getHigh()); + part3.push_back(odd.getHigh()); + } + + SmallVector results; + results.reserve(resultTypes.size()); + results.append(part0); + results.append(part1); + results.append(part2); + results.append(part3); + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } + } + + SmallVector contiguousParts; + contiguousParts.reserve(contiguousTypes.size()); + for (auto [index, resultType] : llvm::enumerate(contiguousTypes)) { + auto vregType = dyn_cast(resultType); + if (!vregType) + return rewriter.notifyMatchFailure(op, "load result must be vreg"); + Value chunkOffset = createChunkOffset(op.getLoc(), *offset, + index * *lanesPerPart, rewriter); + contiguousParts.push_back(rewriter + .create(op.getLoc(), resultType, + /*updated_base=*/Type{}, + *source, chunkOffset, + /*dist=*/nullptr) + .getResult()); + } + + FailureOr> results = materializeDataLayoutConversion( + op, contiguousParts, resultTypes, contiguousLayout, + resultVMIType.getLayoutAttr(), rewriter); + if (failed(results)) + return failure(); + + replaceOpWithFlatConvertedValues(rewriter, op, *results, *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMIDeinterleaveLoadOpPattern + : OpConversionPattern { + using OpConversionPattern< + VMIDeinterleaveLoadOp>::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIDeinterleaveLoadOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto lowVMIType = cast(op.getLow().getType()); + FailureOr source = getSingleValue( + op, adaptor.getSource(), + "deinterleave_load source must convert to one value", rewriter); + FailureOr offset = getSingleValue( + op, adaptor.getOffset(), + "deinterleave_load offset must convert to one value", rewriter); + if (failed(source) || failed(offset)) + return failure(); + + FailureOr lanesPerPart = + getDataLanesPerPart(lowVMIType.getElementType()); + if (failed(lanesPerPart)) + return rewriter.notifyMatchFailure( + op, "deinterleave_load requires known physical lanes per part"); + + std::optional dist = + getX2MemoryDistToken(lowVMIType.getElementType(), "DINTLV"); + if (!dist) + return rewriter.notifyMatchFailure( + op, "deinterleave_load requires vldsx2 DINTLV element support"); + + FailureOr> maybe_lowTypes = + + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + + if (failed(maybe_lowTypes)) + + return failure(); + + SmallVector lowTypes = std::move(*maybe_lowTypes); + FailureOr> maybe_highTypes = + getConvertedResultTypes(op, 1, *this->getTypeConverter()); + if (failed(maybe_highTypes)) + return failure(); + SmallVector highTypes = std::move(*maybe_highTypes); + if (lowTypes.size() != highTypes.size()) + return rewriter.notifyMatchFailure( + op, "deinterleave_load requires matching low/high physical arity"); + + SmallVector lows; + SmallVector highs; + lows.reserve(lowTypes.size()); + highs.reserve(highTypes.size()); + for (size_t index = 0, e = lowTypes.size(); index < e; ++index) { + Type lowType = lowTypes[index]; + Type highType = highTypes[index]; + if (lowType != highType) + return rewriter.notifyMatchFailure( + op, "deinterleave_load requires matching low/high physical types"); + Value chunkOffset = createChunkOffset( + op.getLoc(), *offset, static_cast(index) * 2 * *lanesPerPart, + rewriter); + auto load = + rewriter.create(op.getLoc(), lowType, highType, *source, + chunkOffset, rewriter.getStringAttr(*dist)); + lows.push_back(load.getLow()); + highs.push_back(load.getHigh()); + } + + SmallVector results; + results.reserve(lows.size() + highs.size()); + results.append(lows); + results.append(highs); + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMIGroupLoadOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIGroupLoadOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto resultVMIType = cast(op.getResult().getType()); + FailureOr source = + getSingleValue(op, adaptor.getSource(), + "group_load source must convert to one value", rewriter); + FailureOr offset = + getSingleValue(op, adaptor.getOffset(), + "group_load offset must convert to one value", rewriter); + FailureOr rowStride = getSingleValue( + op, adaptor.getRowStride(), + "group_load row_stride must convert to one value", rewriter); + if (failed(source) || failed(offset) || failed(rowStride)) + return failure(); + + VMILayoutAttr resultLayout = resultVMIType.getLayoutAttr(); + if (resultLayout && resultLayout.isDeinterleaved() && + resultLayout.getBlockElems() == 8 && + resultVMIType.getElementType().isF32()) { + FailureOr groupSize = getGroupSizeFromNumGroups( + resultVMIType, op.getNumGroupsAttr().getInt()); + if (failed(groupSize)) + return rewriter.notifyMatchFailure( + op, "group_load requires num_groups to evenly divide lane count"); + if ((*groupSize != 16 || resultLayout.getFactor() != 2) && + (*groupSize != 32 || resultLayout.getFactor() != 4)) + return rewriter.notifyMatchFailure( + op, "block8 group_load requires S=16/factor=2 or S=32/factor=4"); + if (op.getNumGroupsAttr().getInt() % 8 != 0) + return rewriter.notifyMatchFailure( + op, "block8 group_load requires num_groups multiple of 8"); + std::optional constantRowStride = + getConstantIndexValue(op.getRowStride()); + if (!constantRowStride || *constantRowStride <= 0 || + *constantRowStride % 8 != 0) + return rewriter.notifyMatchFailure( + op, "block8 group_load requires constant positive row_stride " + "divisible by 8 f32 elements"); + if (!isa((*source).getType())) + return rewriter.notifyMatchFailure( + op, "block8 group_load requires !pto.ptr source"); + + FailureOr> maybe_resultTypes = + + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + + if (failed(maybe_resultTypes)) + + return failure(); + + SmallVector resultTypes = std::move(*maybe_resultTypes); + int64_t factor = resultLayout.getFactor(); + FailureOr chunksPerPart = getDataChunksInPart(resultVMIType, 0); + if (failed(chunksPerPart) || *chunksPerPart <= 0) + return rewriter.notifyMatchFailure( + op, "block8 group_load requires known chunks per part"); + for (int64_t part = 1; part < factor; ++part) { + FailureOr currentChunks = + getDataChunksInPart(resultVMIType, part); + if (failed(currentChunks) || *currentChunks != *chunksPerPart) + return rewriter.notifyMatchFailure( + op, "block8 group_load requires uniform chunks per part"); + } + if (static_cast(resultTypes.size()) != factor * *chunksPerPart) + return rewriter.notifyMatchFailure(op, + "block8 group_load arity mismatch"); + + auto makeI16 = [&](int64_t value) -> Value { + return rewriter.create(op.getLoc(), value, 16); + }; + Value blockStride = makeI16(*constantRowStride / 8); + Value zeroI16 = makeI16(0); + auto makePtr = [&](Value elementOffset) -> Value { + return rewriter + .create(op.getLoc(), (*source).getType(), *source, + elementOffset) + .getResult(); + }; + + SmallVector results; + results.reserve(resultTypes.size()); + constexpr int64_t kGroupsPerBlock8Load = 8; + for (int64_t part = 0; part < factor; ++part) { + for (int64_t chunk = 0; chunk < *chunksPerPart; ++chunk) { + int64_t flatIndex = part * *chunksPerPart + chunk; + auto vregType = dyn_cast(resultTypes[flatIndex]); + if (!vregType) + return rewriter.notifyMatchFailure( + op, "block8 group_load result must be vreg"); + FailureOr allMask = + createAllTrueMaskForVReg(op.getLoc(), vregType, rewriter); + if (failed(allMask)) + return rewriter.notifyMatchFailure( + op, "failed to create block8 group_load mask"); + Value chunkOffset = createGroupChunkOffset( + op.getLoc(), *offset, *rowStride, chunk * kGroupsPerBlock8Load, + part * resultLayout.getBlockElems(), rewriter); + Value chunkBase = makePtr(chunkOffset); + results.push_back(rewriter + .create(op.getLoc(), vregType, + chunkBase, blockStride, + zeroI16, *allMask) + .getResult()); + } + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } + + if (resultLayout && resultLayout.isContiguous()) { + FailureOr groupSize = getGroupSizeFromNumGroups( + resultVMIType, op.getNumGroupsAttr().getInt()); + if (failed(groupSize)) + return rewriter.notifyMatchFailure( + op, "group_load requires num_groups to evenly divide lane count"); + std::optional constantRowStride = + getConstantIndexValue(op.getRowStride()); + if (constantRowStride && *constantRowStride == *groupSize) { + FailureOr lanesPerPart = + getDataLanesPerPart(resultVMIType.getElementType()); + if (failed(lanesPerPart)) + return rewriter.notifyMatchFailure( + op, "contiguous group_load requires known physical lanes"); + + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + SmallVector results; + results.reserve(resultTypes.size()); + for (auto [index, resultType] : llvm::enumerate(resultTypes)) { + auto vregType = dyn_cast(resultType); + if (!vregType) + return rewriter.notifyMatchFailure( + op, "contiguous group_load result must be vreg"); + Value chunkOffset = createChunkOffset( + op.getLoc(), *offset, static_cast(index) * *lanesPerPart, + rewriter); + results.push_back(rewriter + .create(op.getLoc(), resultType, + /*updated_base=*/Type{}, + *source, chunkOffset, + /*dist=*/nullptr) + .getResult()); + } + replaceOpWithFlatConvertedValues(rewriter, op, results, + *this->getTypeConverter()); + return success(); + } + } + + int64_t lanesPerPart = 0; + int64_t groupCount = 0; + int64_t chunksPerGroup = 0; + FailureOr groupSize = getGroupSizeFromNumGroups( + resultVMIType, op.getNumGroupsAttr().getInt()); + if (failed(groupSize)) + return rewriter.notifyMatchFailure( + op, "group_load requires num_groups to evenly divide lane count"); + if (failed(checkContiguousFullGroupChunks(op, resultVMIType, *groupSize, + &lanesPerPart, &groupCount, + &chunksPerGroup, rewriter))) + return failure(); + + FailureOr> maybe_resultTypes = + + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + + if (failed(maybe_resultTypes)) + + return failure(); + + SmallVector resultTypes = std::move(*maybe_resultTypes); + if (static_cast(resultTypes.size()) != groupCount * chunksPerGroup) + return rewriter.notifyMatchFailure(op, "group_load arity mismatch"); + + SmallVector results; + results.reserve(resultTypes.size()); + for (auto [index, resultType] : llvm::enumerate(resultTypes)) { + auto vregType = dyn_cast(resultType); + if (!vregType) + return rewriter.notifyMatchFailure(op, + "group_load result must be vreg"); + int64_t group = index / chunksPerGroup; + int64_t chunkInGroup = index % chunksPerGroup; + Value chunkOffset = + createGroupChunkOffset(op.getLoc(), *offset, *rowStride, group, + chunkInGroup * lanesPerPart, rewriter); + results.push_back(rewriter + .create(op.getLoc(), resultType, + /*updated_base=*/Type{}, *source, + chunkOffset, + /*dist=*/nullptr) + .getResult()); + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +static LogicalResult lowerGroupSlotLoadParts( + Operation *op, Value source, Value offset, Value sourceGroupStride, + VMIVRegType resultVMIType, TypeRange resultTypes, int64_t numGroups, + ConversionPatternRewriter &rewriter, SmallVectorImpl &results) { + VMILayoutAttr layout = resultVMIType.getLayoutAttr(); + if (!layout || !layout.isGroupSlots() || layout.getSlots() <= 0) + return rewriter.notifyMatchFailure( + op, "group_slot_load requires explicit group_slots layout"); + if (!isa(source.getType())) + return rewriter.notifyMatchFailure( + op, "group_slot_load requires !pto.ptr source"); + + int64_t slots = layout.getSlots(); + int64_t expectedArity = ceilDivNonNegative(numGroups, slots); + if (static_cast(resultTypes.size()) != expectedArity) + return rewriter.notifyMatchFailure(op, "group_slot_load arity mismatch"); + + auto makeI16 = [&](int64_t value) -> Value { + return rewriter.create(op->getLoc(), value, 16); + }; + Value zeroI16 = makeI16(0); + auto makePtr = [&](Value elementOffset) -> Value { + return rewriter + .create(op->getLoc(), source.getType(), source, elementOffset) + .getResult(); + }; + + results.reserve(results.size() + resultTypes.size()); + + if (slots == 8) { + std::optional stride = getConstantIndexValue(sourceGroupStride); + if (!stride || *stride != 1) + return rewriter.notifyMatchFailure( + op, "slots=8 group_slot_load requires constant unit stride"); + for (auto [chunk, resultType] : llvm::enumerate(resultTypes)) { + auto vregType = dyn_cast(resultType); + if (!vregType) + return rewriter.notifyMatchFailure( + op, "group_slot_load result must be vreg"); + FailureOr maskType = + getMaskTypeForVReg(vregType, rewriter.getContext()); + if (failed(maskType)) + return rewriter.notifyMatchFailure( + op, "unsupported element type for group_slot_load mask"); + int64_t groupBegin = static_cast(chunk) * slots; + int64_t activeGroups = std::min(slots, numGroups - groupBegin); + if (activeGroups <= 0) + return rewriter.notifyMatchFailure( + op, "slots=8 group_slot_load has no active groups for chunk"); + std::string pattern = (Twine("PAT_VL") + Twine(activeGroups)).str(); + FailureOr slotMask = + createPrefixMask(op->getLoc(), *maskType, pattern, rewriter); + if (failed(slotMask)) + return rewriter.notifyMatchFailure( + op, "failed to create slots=8 group_slot_load mask"); + Value groupOffset = + createChunkOffset(op->getLoc(), offset, groupBegin, rewriter); + Value slotBase = makePtr(groupOffset); + results.push_back(rewriter + .create(op->getLoc(), vregType, slotBase, + zeroI16, zeroI16, *slotMask) + .getResult()); + } + return success(); + } + + if (slots != 1) + return rewriter.notifyMatchFailure( + op, "group_slot_load supports only slots=8 or slots=1"); + unsigned elementBits = + pto::getPTOStorageElemBitWidth(resultVMIType.getElementType()); + if (elementBits == 0 || 256 % elementBits != 0) + return rewriter.notifyMatchFailure( + op, "slots=1 group_slot_load requires supported element width"); + int64_t alignedStrideElems = 256 / elementBits; + std::optional constantStride = + getConstantIndexValue(sourceGroupStride); + if (!constantStride || *constantStride <= 0 || + *constantStride % alignedStrideElems != 0) + return rewriter.notifyMatchFailure( + op, Twine("slots=1 group_slot_load requires constant positive " + "source_group_stride divisible by ") + + Twine(alignedStrideElems) + + " elements for 32B lane-0 vsldb alignment"); + + for (auto [group, resultType] : llvm::enumerate(resultTypes)) { + auto vregType = dyn_cast(resultType); + if (!vregType) + return rewriter.notifyMatchFailure(op, + "group_slot_load result must be vreg"); + FailureOr maskType = + getMaskTypeForVReg(vregType, rewriter.getContext()); + if (failed(maskType)) + return rewriter.notifyMatchFailure( + op, "unsupported element type for group_slot_load mask"); + FailureOr oneBlockMask = + createPrefixMask(op->getLoc(), *maskType, "PAT_VL1", rewriter); + if (failed(oneBlockMask)) + return rewriter.notifyMatchFailure( + op, "failed to create group_slot_load mask"); + Value groupOffset = offset; + if (group != 0) { + Value groupIndex = + rewriter.create(op->getLoc(), group); + Value rowOffset = rewriter + .create( + op->getLoc(), sourceGroupStride, groupIndex) + .getResult(); + groupOffset = + rewriter.create(op->getLoc(), groupOffset, rowOffset) + .getResult(); + } + Value slotBase = makePtr(groupOffset); + results.push_back(rewriter + .create(op->getLoc(), vregType, slotBase, + zeroI16, zeroI16, *oneBlockMask) + .getResult()); + } + return success(); +} + +static LogicalResult lowerGroupBroadcastParts( + Operation *op, ValueRange sourceParts, VMIVRegType sourceVMIType, + VMIVRegType resultVMIType, TypeRange resultTypes, int64_t numGroups, + ConversionPatternRewriter &rewriter, SmallVectorImpl &results) { + FailureOr groupSize = + getGroupSizeFromNumGroups(resultVMIType, numGroups); + if (failed(groupSize)) + return rewriter.notifyMatchFailure( + op, "group_broadcast requires num_groups to evenly divide lane count"); + int64_t lanesPerPart = 0; + int64_t groupCount = 0; + if (failed(checkFullGroupSlotSourceShape(op, sourceVMIType, *groupSize, + numGroups, &lanesPerPart, + &groupCount, rewriter))) + return failure(); + int64_t resultLayoutFactor = 0; + int64_t resultGroupCount = 0; + if (failed(checkFullGroupBroadcastResultShape( + op, resultVMIType, *groupSize, lanesPerPart, &resultLayoutFactor, + &resultGroupCount, rewriter))) + return failure(); + if (resultGroupCount != groupCount) + return rewriter.notifyMatchFailure( + op, "group_broadcast requires matching source/result group slots"); + + if (sourceParts.empty() || resultTypes.empty()) + return rewriter.notifyMatchFailure(op, "group_broadcast arity mismatch"); + + auto firstSourceType = dyn_cast(sourceParts.front().getType()); + if (!firstSourceType) + return rewriter.notifyMatchFailure(op, + "group_broadcast source must be vreg"); + unsigned indexBits = + pto::getPTOStorageElemBitWidth(firstSourceType.getElementType()); + if (indexBits != 8 && indexBits != 16 && indexBits != 32) + return rewriter.notifyMatchFailure( + op, "group_broadcast requires 8/16/32-bit index elements"); + auto indexElementType = IntegerType::get(rewriter.getContext(), indexBits); + auto indexType = + VRegType::get(rewriter.getContext(), firstSourceType.getElementCount(), + indexElementType); + FailureOr allMask = + createAllTrueMaskForVReg(op->getLoc(), firstSourceType, rewriter); + if (failed(allMask)) + return rewriter.notifyMatchFailure( + op, "failed to create group_broadcast all mask"); + VMILayoutAttr resultLayout = resultVMIType.getLayoutAttr(); + VMILayoutAttr sourceLayout = sourceVMIType.getLayoutAttr(); + int64_t selectionGroupSize = *groupSize; + if (resultLayoutFactor != 1 && resultLayout && + resultLayout.isDeinterleaved() && resultLayout.getBlockElems() > 1 && + *groupSize < lanesPerPart) + selectionGroupSize = resultLayout.getBlockElems(); + auto resolveLargeGroupSource = [&](int64_t group, int64_t chunksPerGroup, + int64_t &sourceChunk, + int64_t &baseGroupSlot) { + int64_t slots = sourceLayout.getSlots(); + if (slots > 0) { + sourceChunk = group / slots; + baseGroupSlot = group % slots; + return; + } + sourceChunk = group * chunksPerGroup; + baseGroupSlot = 0; + }; + + results.clear(); + results.resize(resultTypes.size()); + for (auto [flatIndex, resultType] : llvm::enumerate(resultTypes)) { + auto resultVRegType = dyn_cast(resultType); + if (!resultVRegType || resultVRegType != firstSourceType) + return rewriter.notifyMatchFailure( + op, "group_broadcast requires uniform physical vreg types"); + int64_t sourceChunk = flatIndex; + int64_t baseGroupSlot = 0; + Value mappedGroupSlotIndex; + if (resultLayoutFactor == 1) { + bool laneStridedDense = + resultLayout && resultLayout.isDense() && + resultLayout.getLaneStride() > 1; + if (laneStridedDense) { + VMILayoutAttr sourceLayout = sourceVMIType.getLayoutAttr(); + int64_t slots = sourceLayout.getSlots(); + if (slots <= 0) { + if (sourceParts.empty() || + groupCount % static_cast(sourceParts.size()) != 0) + return rewriter.notifyMatchFailure( + op, "group_broadcast lane-stride source requires explicit " + "group_slots slots or derivable legacy slot count"); + slots = groupCount / sourceParts.size(); + } + FailureOr index = createMappedGroupSlotIndexVector( + op->getLoc(), resultVMIType, /*part=*/0, flatIndex, indexType, + *groupSize, slots, sourceChunk, rewriter, + sourceLayout.getLaneStride()); + if (failed(index)) + return rewriter.notifyMatchFailure( + op, "failed to create group_broadcast lane-stride group-slot " + "index vector"); + mappedGroupSlotIndex = *index; + } else if (*groupSize >= lanesPerPart) { + int64_t chunksPerGroup = *groupSize / lanesPerPart; + int64_t group = flatIndex / chunksPerGroup; + resolveLargeGroupSource(group, chunksPerGroup, sourceChunk, + baseGroupSlot); + } else { + VMILayoutAttr sourceLayout = sourceVMIType.getLayoutAttr(); + int64_t slots = sourceLayout.getSlots(); + if (slots <= 0) { + if (sourceParts.empty() || + groupCount % static_cast(sourceParts.size()) != 0) + return rewriter.notifyMatchFailure( + op, "group_broadcast small-group source requires explicit " + "group_slots slots or derivable legacy slot count"); + slots = groupCount / sourceParts.size(); + } + int64_t groupsPerResultChunk = lanesPerPart / *groupSize; + int64_t firstGroup = flatIndex * groupsPerResultChunk; + sourceChunk = firstGroup / slots; + baseGroupSlot = firstGroup % slots; + } + } else { + bool blockFragmentSmallGroup = + resultLayout && resultLayout.isDeinterleaved() && + resultLayout.getBlockElems() > 1 && *groupSize < lanesPerPart; + bool deinterleavedSmallGroup = + resultLayout && resultLayout.isDeinterleaved() && + resultLayout.getBlockElems() == 1 && *groupSize < lanesPerPart; + if (blockFragmentSmallGroup) { + int64_t runningFlatIndex = 0; + bool found = false; + for (int64_t part = 0; part < resultLayoutFactor && !found; ++part) { + FailureOr chunks = getDataChunksInPart(resultVMIType, part); + if (failed(chunks)) + return rewriter.notifyMatchFailure( + op, "group_broadcast failed to enumerate result chunks"); + for (int64_t chunk = 0; chunk < *chunks; + ++chunk, ++runningFlatIndex) { + if (runningFlatIndex != static_cast(flatIndex)) + continue; + int64_t groupsPerResultChunk = + lanesPerPart / resultLayout.getBlockElems(); + int64_t firstGroup = chunk * groupsPerResultChunk; + int64_t slots = sourceLayout.getSlots(); + if (slots <= 0) { + if (sourceParts.empty() || + groupCount % static_cast(sourceParts.size()) != 0) + return rewriter.notifyMatchFailure( + op, + "group_broadcast block-fragment source requires explicit " + "group_slots slots or derivable legacy slot count"); + slots = groupCount / sourceParts.size(); + } + sourceChunk = firstGroup / slots; + baseGroupSlot = firstGroup % slots; + found = true; + break; + } + } + if (!found) + return rewriter.notifyMatchFailure( + op, "group_broadcast result chunk index is out of range"); + } else if (deinterleavedSmallGroup) { + int64_t runningFlatIndex = 0; + bool found = false; + for (int64_t part = 0; part < resultLayoutFactor && !found; ++part) { + FailureOr chunks = getDataChunksInPart(resultVMIType, part); + if (failed(chunks)) + return rewriter.notifyMatchFailure( + op, "group_broadcast failed to enumerate result chunks"); + for (int64_t chunk = 0; chunk < *chunks; + ++chunk, ++runningFlatIndex) { + if (runningFlatIndex != static_cast(flatIndex)) + continue; + int64_t slots = sourceLayout.getSlots(); + if (slots <= 0) { + if (sourceParts.empty() || + groupCount % static_cast(sourceParts.size()) != 0) + return rewriter.notifyMatchFailure( + op, "group_broadcast deinterleaved small-group source " + "requires explicit group_slots slots or derivable " + "legacy slot count"); + slots = groupCount / sourceParts.size(); + } + FailureOr index = createMappedGroupSlotIndexVector( + op->getLoc(), resultVMIType, part, chunk, indexType, *groupSize, + slots, sourceChunk, rewriter, sourceLayout.getLaneStride()); + if (failed(index)) + return rewriter.notifyMatchFailure( + op, + "failed to create group_broadcast mapped group-slot index " + "vector"); + mappedGroupSlotIndex = *index; + found = true; + break; + } + } + if (!found) + return rewriter.notifyMatchFailure( + op, "group_broadcast result chunk index is out of range"); + } else { + int64_t runningFlatIndex = 0; + bool found = false; + for (int64_t part = 0; part < resultLayoutFactor && !found; ++part) { + FailureOr chunks = getDataChunksInPart(resultVMIType, part); + if (failed(chunks)) + return rewriter.notifyMatchFailure( + op, "group_broadcast failed to enumerate result chunks"); + for (int64_t chunk = 0; chunk < *chunks; + ++chunk, ++runningFlatIndex) { + if (runningFlatIndex != static_cast(flatIndex)) + continue; + FailureOr firstLogical = + mapPhysicalLaneToLogical(resultVMIType, part, chunk, 0); + FailureOr lastLogical = mapPhysicalLaneToLogical( + resultVMIType, part, chunk, lanesPerPart - 1); + if (failed(firstLogical) || failed(lastLogical)) + return rewriter.notifyMatchFailure( + op, "group_broadcast failed to map result chunk lanes"); + int64_t firstGroup = *firstLogical / *groupSize; + int64_t lastGroup = *lastLogical / *groupSize; + if (firstGroup != lastGroup) + return rewriter.notifyMatchFailure( + op, "group_broadcast result chunk crosses logical groups"); + int64_t chunksPerGroup = *groupSize / lanesPerPart; + resolveLargeGroupSource(firstGroup, chunksPerGroup, sourceChunk, + baseGroupSlot); + found = true; + break; + } + } + if (!found) + return rewriter.notifyMatchFailure( + op, "group_broadcast result chunk index is out of range"); + } + } + if (*groupSize >= lanesPerPart) { + if (sourceChunk < 0 || + sourceChunk >= static_cast(sourceParts.size())) + return rewriter.notifyMatchFailure( + op, "group_broadcast source chunk is out of range"); + if (sourceLayout.getSlots() > 1) { + FailureOr groupSlotIndex = createGroupSlotIndexVector( + op->getLoc(), indexType, selectionGroupSize, baseGroupSlot, + rewriter, sourceLayout.getLaneStride()); + if (failed(groupSlotIndex)) + return rewriter.notifyMatchFailure( + op, "failed to create group_broadcast group-slot index vector"); + results[flatIndex] = + rewriter + .create(op->getLoc(), resultType, + sourceParts[sourceChunk], *groupSlotIndex) + .getResult(); + } else { + results[flatIndex] = + rewriter + .create(op->getLoc(), resultType, + sourceParts[sourceChunk], *allMask, + rewriter.getStringAttr("LOWEST")) + .getResult(); + } + } else { + bool blockFragmentSmallGroup = resultLayout && + resultLayout.isDeinterleaved() && + resultLayout.getBlockElems() > 1; + bool deinterleavedSmallGroup = resultLayout && + resultLayout.isDeinterleaved() && + resultLayout.getBlockElems() == 1; + if (resultLayoutFactor != 1 && !blockFragmentSmallGroup && + !deinterleavedSmallGroup) + return rewriter.notifyMatchFailure( + op, "group_broadcast small-group deinterleaved result is not " + "supported"); + if (sourceChunk < 0 || + sourceChunk >= static_cast(sourceParts.size())) + return rewriter.notifyMatchFailure( + op, "group_broadcast source chunk is out of range"); + FailureOr groupSlotIndex = + mappedGroupSlotIndex + ? FailureOr(mappedGroupSlotIndex) + : createGroupSlotIndexVector(op->getLoc(), indexType, + selectionGroupSize, baseGroupSlot, + rewriter, + sourceLayout.getLaneStride()); + if (failed(groupSlotIndex)) + return rewriter.notifyMatchFailure( + op, "failed to create group_broadcast group-slot index vector"); + results[flatIndex] = + rewriter + .create(op->getLoc(), resultType, + sourceParts[sourceChunk], *groupSlotIndex) + .getResult(); + } + } + return success(); +} + +struct OneToNVMIGroupSlotLoadOpPattern + : OpConversionPattern { + using OpConversionPattern< + VMIGroupSlotLoadOp>::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIGroupSlotLoadOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto resultVMIType = cast(op.getResult().getType()); + VMILayoutAttr layout = resultVMIType.getLayoutAttr(); + if (!layout || !layout.isGroupSlots() || layout.getSlots() <= 0) + return rewriter.notifyMatchFailure( + op, "group_slot_load requires explicit group_slots layout"); + + FailureOr source = getSingleValue( + op, adaptor.getSource(), + "group_slot_load source must convert to one value", rewriter); + FailureOr offset = getSingleValue( + op, adaptor.getOffset(), + "group_slot_load offset must convert to one value", rewriter); + FailureOr sourceGroupStride = getSingleValue( + op, adaptor.getSourceGroupStride(), + "group_slot_load source_group_stride must convert to one value", + rewriter); + if (failed(source) || failed(offset) || failed(sourceGroupStride)) + return failure(); + + FailureOr> maybe_resultTypes = + + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + + if (failed(maybe_resultTypes)) + + return failure(); + + SmallVector resultTypes = std::move(*maybe_resultTypes); + int64_t numGroups = op.getNumGroupsAttr().getInt(); + + SmallVector results; + if (failed(lowerGroupSlotLoadParts(op, *source, *offset, *sourceGroupStride, + resultVMIType, resultTypes, numGroups, + rewriter, results))) + return failure(); + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMIMaskedLoadOpPattern + : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIMaskedLoadOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto resultVMIType = cast(op.getResult().getType()); + FailureOr source = getSingleValue( + op, adaptor.getSource(), "masked_load source must convert to one value", + rewriter); + FailureOr offset = getSingleValue( + op, adaptor.getOffset(), "masked_load offset must convert to one value", + rewriter); + if (failed(source) || failed(offset)) + return failure(); + + FailureOr lanesPerPart = verifyFullOrSafeReadVRegChunks( + op, resultVMIType, (*source).getType(), *offset, rewriter); + if (failed(lanesPerPart)) + return failure(); + + ValueRange maskParts = adaptor.getMask(); + ValueRange passthruParts = adaptor.getPassthru(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + if (maskParts.size() != passthruParts.size() || + passthruParts.size() != resultTypes.size()) + return rewriter.notifyMatchFailure(op, + "masked_load physical arity mismatch"); + + SmallVector results; + results.reserve(resultTypes.size()); + for (auto [index, maskPassthruAndType] : llvm::enumerate( + llvm::zip_equal(maskParts, passthruParts, resultTypes))) { + auto [mask, passthru, resultType] = maskPassthruAndType; + if (!isa(mask.getType()) || passthru.getType() != resultType || + !isa(resultType)) + return rewriter.notifyMatchFailure( + op, "masked_load physical part type mismatch"); + + Value chunkOffset = createChunkOffset(op.getLoc(), *offset, + index * *lanesPerPart, rewriter); + Value loaded = + rewriter + .create(op.getLoc(), resultType, + /*updated_base=*/Type{}, *source, chunkOffset, + /*dist=*/nullptr) + .getResult(); + results.push_back( + rewriter + .create(op.getLoc(), resultType, loaded, passthru, mask) + .getResult()); + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMIGatherOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIGatherOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + FailureOr source = + getSingleValue(op, adaptor.getSource(), + "gather source must convert to one value", rewriter); + if (failed(source)) + return failure(); + + ValueRange indicesParts = adaptor.getIndices(); + ValueRange maskParts = adaptor.getMask(); + ValueRange passthruParts = adaptor.getPassthru(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + if (indicesParts.size() != maskParts.size() || + indicesParts.size() != passthruParts.size() || + indicesParts.size() != resultTypes.size()) + return rewriter.notifyMatchFailure(op, "gather physical arity mismatch"); + + SmallVector results; + results.reserve(resultTypes.size()); + for (auto [indices, mask, passthru, resultType] : + llvm::zip_equal(indicesParts, maskParts, passthruParts, resultTypes)) { + if (!isa(indices.getType()) || !isa(mask.getType()) || + passthru.getType() != resultType || !isa(resultType)) + return rewriter.notifyMatchFailure( + op, "gather physical part type mismatch"); + + unsigned resultBits = pto::getPTOStorageElemBitWidth( + cast(resultType).getElementType()); + Value gathered = resultBits == 16 + ? rewriter + .create(op.getLoc(), resultType, + *source, indices, mask) + .getResult() + : rewriter + .create(op.getLoc(), resultType, + *source, indices, mask) + .getResult(); + results.push_back( + rewriter + .create(op.getLoc(), resultType, gathered, passthru, mask) + .getResult()); + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMIExpandLoadOpPattern + : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIExpandLoadOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto resultVMIType = cast(op.getResult().getType()); + FailureOr source = getSingleValue( + op, adaptor.getSource(), "expand_load source must convert to one value", + rewriter); + FailureOr offset = getSingleValue( + op, adaptor.getOffset(), "expand_load offset must convert to one value", + rewriter); + if (failed(source) || failed(offset)) + return failure(); + + FailureOr> maybe_resultTypes = + + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + + if (failed(maybe_resultTypes)) + + return failure(); + + SmallVector resultTypes = std::move(*maybe_resultTypes); + if (isStaticAllActiveMask(op.getMask(), resultVMIType.getElementCount())) { + FailureOr lanesPerPart = verifyFullOrSafeReadVRegChunks( + op, resultVMIType, (*source).getType(), *offset, rewriter); + if (failed(lanesPerPart)) + return failure(); + + SmallVector results; + results.reserve(resultTypes.size()); + for (auto [index, resultType] : llvm::enumerate(resultTypes)) { + if (!isa(resultType)) + return rewriter.notifyMatchFailure(op, + "expand_load result must be vreg"); + Value chunkOffset = createChunkOffset(op.getLoc(), *offset, + index * *lanesPerPart, rewriter); + results.push_back(rewriter + .create(op.getLoc(), resultType, + /*updated_base=*/Type{}, *source, + chunkOffset, + /*dist=*/nullptr) + .getResult()); + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } + + ValueRange maskParts = adaptor.getMask(); + ValueRange passthruParts = adaptor.getPassthru(); + if (resultTypes.size() != 1 || maskParts.size() != 1 || + passthruParts.size() != 1) + return rewriter.notifyMatchFailure( + op, "runtime expand_load supports only one physical chunk"); + + auto resultType = dyn_cast(resultTypes.front()); + auto maskType = dyn_cast(maskParts.front().getType()); + if (!resultType || !maskType || + passthruParts.front().getType() != resultType) + return rewriter.notifyMatchFailure( + op, "runtime expand_load requires physical result/passthru/mask"); + + auto baseType = dyn_cast((*source).getType()); + if (!baseType) + return rewriter.notifyMatchFailure(op, + "runtime expand_load requires ptr"); + Value gatherBase = rewriter + .create(op.getLoc(), (*source).getType(), + *source, *offset) + .getResult(); + auto indexType = + VRegType::get(rewriter.getContext(), resultType.getElementCount(), + rewriter.getI32Type()); + FailureOr indexSeedMask = + createAllTrueMaskForVReg(op.getLoc(), indexType, rewriter); + if (failed(indexSeedMask)) + return rewriter.notifyMatchFailure( + op, "failed to create runtime expand_load index seed mask"); + Value zero = rewriter.create(op.getLoc(), 0, 32); + Value carrier = + rewriter + .create(op.getLoc(), indexType, zero, *indexSeedMask, + /*position=*/nullptr) + .getResult(); + Value indices = + rewriter + .create(op.getLoc(), indexType, carrier, maskParts.front()) + .getResult(); + Value gathered = + rewriter + .create(op.getLoc(), resultType, gatherBase, indices, + maskParts.front()) + .getResult(); + Value result = rewriter + .create(op.getLoc(), resultType, gathered, + passthruParts.front(), maskParts.front()) + .getResult(); + replaceOpWithFlatConvertedValues(rewriter, op, SmallVector{result}, + *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMIStoreOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIStoreOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto valueVMIType = cast(op.getValue().getType()); + FailureOr lanesPerPart = + getDataLanesPerPart(valueVMIType.getElementType()); + if (failed(lanesPerPart)) + return rewriter.notifyMatchFailure( + op, "store requires known physical lanes per part"); + bool fullPhysicalChunks = + succeeded(checkFullDataPhysicalChunks(valueVMIType, nullptr)); + FailureOr destination = + getSingleValue(op, adaptor.getDestination(), + "store destination must convert to one value", rewriter); + FailureOr offset = + getSingleValue(op, adaptor.getOffset(), + "store offset must convert to one value", rewriter); + if (failed(destination) || failed(offset)) + return failure(); + + ValueRange valueParts = adaptor.getValue(); + if (std::optional dist = + getDenseLaneStrideStoreDistToken(valueVMIType)) { + std::optional maskGranularity = + getDenseLaneStrideStoreMaskGranularity(valueVMIType); + if (!maskGranularity) + return rewriter.notifyMatchFailure( + op, "unsupported lane_stride store mask granularity"); + int64_t semanticOffset = 0; + for (auto [index, value] : llvm::enumerate(valueParts)) { + auto vregType = dyn_cast(value.getType()); + if (!vregType) + return rewriter.notifyMatchFailure(op, "store value must be vreg"); + FailureOr activeLanes = + getActiveDataLanesInPhysicalChunk(valueVMIType, index); + if (failed(activeLanes)) + return rewriter.notifyMatchFailure( + op, "failed to compute lane_stride store active lanes"); + if (*activeLanes == 0) + continue; + auto maskType = MaskType::get(rewriter.getContext(), *maskGranularity); + FailureOr mask = createPrefixMaskForActiveLanes( + op.getLoc(), maskType, *activeLanes, rewriter); + if (failed(mask)) + return rewriter.notifyMatchFailure( + op, "failed to create lane_stride store mask"); + Value chunkOffset = + createChunkOffset(op.getLoc(), *offset, semanticOffset, rewriter); + rewriter.create(op.getLoc(), + /*updated_base=*/Type{}, value, *destination, + chunkOffset, rewriter.getStringAttr(*dist), + *mask); + semanticOffset += *activeLanes; + } + rewriter.eraseOp(op); + return success(); + } + + VMILayoutAttr contiguousLayout = + VMILayoutAttr::getContiguous(rewriter.getContext()); + FailureOr> maybeContiguousTypes = + getConvertedVRegTypesWithLayout(valueVMIType, contiguousLayout, + *this->getTypeConverter()); + if (failed(maybeContiguousTypes)) + return rewriter.notifyMatchFailure( + op, "failed to compute contiguous store footprint"); + SmallVector contiguousTypes = std::move(*maybeContiguousTypes); + SmallVector valuePartTypes; + valuePartTypes.reserve(valueParts.size()); + for (Value value : valueParts) + valuePartTypes.push_back(value.getType()); + FailureOr noWiderThanContiguous = + hasNoWiderFootprintThanContiguous(valuePartTypes, contiguousTypes); + if (failed(noWiderThanContiguous)) + return rewriter.notifyMatchFailure( + op, "failed to compare store physical footprint"); + + VMILayoutSupport localSupports; + FailureOr storeFact = + localSupports.getStoreLayoutFact(valueVMIType); + if (succeeded(storeFact) && storeFact->valueLayout.isDeinterleaved() && + storeFact->valueLayout.getFactor() == 2 && fullPhysicalChunks && + *noWiderThanContiguous) { + std::optional dist = + getX2MemoryDistToken(valueVMIType.getElementType(), "INTLV"); + if (dist && !valueParts.empty() && valueParts.size() % 2 == 0) { + int64_t groups = valueParts.size() / 2; + for (int64_t group = 0; group < groups; ++group) { + Value low = valueParts[group]; + Value high = valueParts[groups + group]; + if (low.getType() != high.getType()) + return rewriter.notifyMatchFailure( + op, "vstsx2 requires matching low/high value types"); + auto vregType = dyn_cast(low.getType()); + if (!vregType) + return rewriter.notifyMatchFailure(op, "store value must be vreg"); + FailureOr mask = + createAllTrueMaskForVReg(op.getLoc(), vregType, rewriter); + if (failed(mask)) + return rewriter.notifyMatchFailure( + op, "unsupported element type for store mask"); + Value chunkOffset = createChunkOffset( + op.getLoc(), *offset, group * 2 * *lanesPerPart, rewriter); + rewriter.create(op.getLoc(), low, high, *destination, + chunkOffset, rewriter.getStringAttr(*dist), + *mask); + } + rewriter.eraseOp(op); + return success(); + } + } + + FailureOr> storeParts = materializeDataLayoutConversion( + op, valueParts, contiguousTypes, valueVMIType.getLayoutAttr(), + contiguousLayout, rewriter); + if (failed(storeParts)) + return failure(); + + for (auto [index, value] : llvm::enumerate(*storeParts)) { + auto vregType = dyn_cast(value.getType()); + if (!vregType) + return rewriter.notifyMatchFailure(op, "store value must be vreg"); + if (!fullPhysicalChunks) { + FailureOr activeLanes = + getContiguousActiveDataLanes(valueVMIType, index); + if (failed(activeLanes)) + return rewriter.notifyMatchFailure( + op, "failed to compute store active lanes"); + if (*activeLanes == 0) + continue; + } + FailureOr mask = + fullPhysicalChunks + ? createAllTrueMaskForVReg(op.getLoc(), vregType, rewriter) + : createContiguousStoreMask(op.getLoc(), valueVMIType, index, + vregType, rewriter); + if (failed(mask)) + return rewriter.notifyMatchFailure( + op, "unsupported element type for store mask"); + Value chunkOffset = createChunkOffset(op.getLoc(), *offset, + index * *lanesPerPart, rewriter); + rewriter.create(op.getLoc(), + /*updated_base=*/Type{}, value, *destination, + chunkOffset, /*dist=*/nullptr, *mask); + } + + rewriter.eraseOp(op); + return success(); + } +}; + +struct OneToNVMIInterleaveStoreOpPattern + : OpConversionPattern { + using OpConversionPattern< + VMIInterleaveStoreOp>::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIInterleaveStoreOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto lowVMIType = cast(op.getLow().getType()); + FailureOr lanesPerPart = + getDataLanesPerPart(lowVMIType.getElementType()); + if (failed(lanesPerPart)) + return rewriter.notifyMatchFailure( + op, "interleave_store requires known physical lanes per part"); + + std::optional dist = + getX2MemoryDistToken(lowVMIType.getElementType(), "INTLV"); + if (!dist) + return rewriter.notifyMatchFailure( + op, "interleave_store requires vstsx2 INTLV element support"); + + FailureOr destination = getSingleValue( + op, adaptor.getDestination(), + "interleave_store destination must convert to one value", rewriter); + FailureOr offset = getSingleValue( + op, adaptor.getOffset(), + "interleave_store offset must convert to one value", rewriter); + if (failed(destination) || failed(offset)) + return failure(); + + ValueRange lowParts = adaptor.getLow(); + ValueRange highParts = adaptor.getHigh(); + if (lowParts.size() != highParts.size()) + return rewriter.notifyMatchFailure( + op, "interleave_store requires matching low/high physical arity"); + + for (size_t index = 0, e = lowParts.size(); index < e; ++index) { + Value low = lowParts[index]; + Value high = highParts[index]; + if (low.getType() != high.getType()) + return rewriter.notifyMatchFailure( + op, "interleave_store requires matching low/high physical types"); + auto vregType = dyn_cast(low.getType()); + if (!vregType) + return rewriter.notifyMatchFailure( + op, "interleave_store value must be vreg"); + FailureOr mask = + createAllTrueMaskForVReg(op.getLoc(), vregType, rewriter); + if (failed(mask)) + return rewriter.notifyMatchFailure( + op, "unsupported element type for interleave_store mask"); + Value chunkOffset = createChunkOffset( + op.getLoc(), *offset, static_cast(index) * 2 * *lanesPerPart, + rewriter); + rewriter.create(op.getLoc(), low, high, *destination, + chunkOffset, rewriter.getStringAttr(*dist), + *mask); + } + + rewriter.eraseOp(op); + return success(); + } +}; + +struct OneToNVMIGroupStoreOpPattern + : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIGroupStoreOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto valueVMIType = cast(op.getValue().getType()); + VMILayoutAttr layout = valueVMIType.getLayoutAttr(); + + FailureOr destination = getSingleValue( + op, adaptor.getDestination(), + "group_store destination must convert to one value", rewriter); + FailureOr offset = getSingleValue( + op, adaptor.getOffset(), "group_store offset must convert to one value", + rewriter); + FailureOr rowStride = getSingleValue( + op, adaptor.getRowStride(), + "group_store row_stride must convert to one value", rewriter); + if (failed(destination) || failed(offset) || failed(rowStride)) + return failure(); + + if (layout && layout.isGroupSlots() && layout.getSlots() == 1 && + layout.getNumGroups() == op.getNumGroupsAttr().getInt()) { + ValueRange valueParts = adaptor.getValue(); + if (static_cast(valueParts.size()) != layout.getNumGroups()) + return rewriter.notifyMatchFailure( + op, "slots=1 group_store arity mismatch"); + unsigned elementBits = + pto::getPTOStorageElemBitWidth(valueVMIType.getElementType()); + if (elementBits == 0 || 256 % elementBits != 0) + return rewriter.notifyMatchFailure( + op, "slots=1 group_store requires supported element width"); + std::optional constantRowStride = + getConstantIndexValue(op.getRowStride()); + FailureOr lanesPerPart = + getDataLanesPerPart(valueVMIType.getElementType()); + int64_t alignedStoreElems = 256 / elementBits; + if (constantRowStride && *constantRowStride == 1 && + succeeded(lanesPerPart) && layout.getNumGroups() <= *lanesPerPart && + isKnownIndexMultipleOf(op.getOffset(), alignedStoreElems)) { + auto firstType = dyn_cast(valueParts.front().getType()); + if (!firstType) + return rewriter.notifyMatchFailure(op, + "group_store value must be vreg"); + FailureOr maskType = + getMaskTypeForVReg(firstType, rewriter.getContext()); + FailureOr allMask = + createAllTrueMaskForVReg(op.getLoc(), firstType, rewriter); + if (failed(maskType) || failed(allMask)) + return rewriter.notifyMatchFailure( + op, "unsupported element type for packed group_store mask"); + + Value packed = + rewriter + .create(op.getLoc(), firstType, valueParts.front(), + *allMask, rewriter.getStringAttr("LOWEST")) + .getResult(); + for (int64_t group = 1; group < layout.getNumGroups(); ++group) { + auto vregType = dyn_cast(valueParts[group].getType()); + if (!vregType || vregType != firstType) + return rewriter.notifyMatchFailure( + op, "packed group_store requires uniform vreg parts"); + Value splat = + rewriter + .create(op.getLoc(), firstType, valueParts[group], + *allMask, rewriter.getStringAttr("LOWEST")) + .getResult(); + FailureOr laneMask = createLaneRangeMask( + op.getLoc(), *maskType, group, group + 1, rewriter); + if (failed(laneMask)) + return rewriter.notifyMatchFailure( + op, "failed to create packed group_store lane mask"); + packed = rewriter + .create(op.getLoc(), firstType, splat, packed, + *laneMask) + .getResult(); + } + + FailureOr storeMask = createPrefixMaskForActiveLanes( + op.getLoc(), *maskType, layout.getNumGroups(), rewriter); + if (failed(storeMask)) + return rewriter.notifyMatchFailure( + op, "failed to create packed group_store store mask"); + rewriter.create(op.getLoc(), + /*updated_base=*/Type{}, packed, *destination, + *offset, /*dist=*/nullptr, *storeMask); + rewriter.eraseOp(op); + return success(); + } + if (constantRowStride && *constantRowStride <= 0) + return rewriter.notifyMatchFailure( + op, "slots=1 group_store requires positive row_stride when " + "row_stride is constant"); + std::optional pointDist = + getPointStoreDistToken(valueVMIType.getElementType()); + if (!pointDist) + return rewriter.notifyMatchFailure( + op, "slots=1 group_store requires 1PT_B8/B16/B32 store support"); + + for (auto [group, value] : llvm::enumerate(valueParts)) { + auto vregType = dyn_cast(value.getType()); + if (!vregType) + return rewriter.notifyMatchFailure(op, + "group_store value must be vreg"); + FailureOr maskType = + getMaskTypeForVReg(vregType, rewriter.getContext()); + if (failed(maskType)) + return rewriter.notifyMatchFailure( + op, "unsupported element type for group_store mask"); + FailureOr mask = + createPrefixMask(op.getLoc(), *maskType, "PAT_VL1", rewriter); + if (failed(mask)) + return rewriter.notifyMatchFailure( + op, "failed to create slots=1 group_store mask"); + Value groupOffset = + createGroupChunkOffset(op.getLoc(), *offset, *rowStride, group, + /*chunkLaneOffset=*/0, rewriter); + rewriter.create(op.getLoc(), + /*updated_base=*/Type{}, value, *destination, + groupOffset, rewriter.getStringAttr(*pointDist), + *mask); + } + + rewriter.eraseOp(op); + return success(); + } + + if (layout && layout.isGroupSlots() && layout.getSlots() == 8 && + layout.getNumGroups() == op.getNumGroupsAttr().getInt()) { + int64_t numGroups = layout.getNumGroups(); + std::optional constantRowStride = + getConstantIndexValue(op.getRowStride()); + if (!constantRowStride || *constantRowStride != 1) + return rewriter.notifyMatchFailure( + op, "slots=8 group_store requires constant unit row_stride"); + + ValueRange valueParts = adaptor.getValue(); + if (static_cast(valueParts.size()) != + ceilDivNonNegative(numGroups, 8)) + return rewriter.notifyMatchFailure( + op, "slots=8 group_store arity mismatch"); + + if (!valueParts.empty()) { + auto firstVRegType = dyn_cast(valueParts.front().getType()); + if (!firstVRegType) + return rewriter.notifyMatchFailure(op, + "group_store value must be vreg"); + bool packedByteStore = isPackedByteGroupStore( + op.getDestination().getType(), firstVRegType); + if (packedByteStore) { + bool laneStridedPackedByteStore = layout.hasLaneStride(); + for (Value value : valueParts) { + auto vregType = dyn_cast(value.getType()); + if (!vregType || vregType != firstVRegType) + return rewriter.notifyMatchFailure( + op, "packed slots=8 group_store requires uniform vreg parts"); + } + + FailureOr maskType = + getMaskTypeForVReg(firstVRegType, rewriter.getContext()); + if (failed(maskType)) + return rewriter.notifyMatchFailure( + op, "unsupported element type for packed group_store mask"); + if (!laneStridedPackedByteStore && numGroups == 8 && + valueParts.size() == 1 && isKnownIndexMultipleOf(*offset, 32)) { + MLIRContext *ctx = rewriter.getContext(); + auto ui16 = IntegerType::get( + ctx, 16, IntegerType::SignednessSemantics::Unsigned); + auto ui8 = IntegerType::get( + ctx, 8, IntegerType::SignednessSemantics::Unsigned); + auto packed16Type = VRegType::get(ctx, 128, ui16); + auto packed8Type = VRegType::get(ctx, 256, ui8); + Value packed16 = + rewriter + .create(op.getLoc(), packed16Type, + valueParts.front(), + rewriter.getStringAttr("LOWER")) + .getResult(); + Value packed8 = + rewriter + .create(op.getLoc(), packed8Type, packed16, + rewriter.getStringAttr("LOWER")) + .getResult(); + FailureOr packedMaskType = + getMaskTypeForVReg(packed8Type, ctx); + if (failed(packedMaskType)) + return rewriter.notifyMatchFailure( + op, "failed to create packed byte group_store mask type"); + FailureOr storeMask = createPrefixMaskForActiveLanes( + op.getLoc(), *packedMaskType, numGroups, rewriter); + if (failed(storeMask)) + return rewriter.notifyMatchFailure( + op, "failed to create packed byte group_store mask"); + rewriter.create( + op.getLoc(), /*updated_base=*/Type{}, packed8, *destination, + *offset, rewriter.getStringAttr("NORM_B8"), *storeMask); + rewriter.eraseOp(op); + return success(); + } + + auto indexElementType = IntegerType::get( + rewriter.getContext(), + pto::getPTOStorageElemBitWidth(firstVRegType.getElementType())); + auto indexType = + VRegType::get(rewriter.getContext(), + firstVRegType.getElementCount(), indexElementType); + FailureOr slotIndex = createGroupSlotIndexVector( + op.getLoc(), indexType, /*groupSize=*/8, /*baseGroupSlot=*/0, + rewriter); + FailureOr allMask = + createAllTrueMaskForVReg(op.getLoc(), firstVRegType, rewriter); + if (failed(slotIndex) || failed(allMask)) + return rewriter.notifyMatchFailure( + op, "failed to create packed group_store lane selector"); + + for (int64_t blockStart = 0; blockStart < numGroups; + blockStart += 32) { + FailureOr zero = + createZeroVector(op.getLoc(), firstVRegType, rewriter); + if (failed(zero)) + return rewriter.notifyMatchFailure( + op, "failed to create packed group_store accumulator"); + Value merged = *zero; + for (int64_t localPart = 0; localPart < 4; ++localPart) { + int64_t partIndex = blockStart / 8 + localPart; + if (partIndex >= static_cast(valueParts.size())) + break; + int64_t remainingGroups = numGroups - partIndex * 8; + int64_t activeGroups = std::min(8, remainingGroups); + if (activeGroups <= 0) + break; + Value selected = + rewriter + .create(op.getLoc(), firstVRegType, + valueParts[partIndex], *slotIndex) + .getResult(); + FailureOr laneMask = + createLaneRangeMask(op.getLoc(), *maskType, localPart * 8, + localPart * 8 + activeGroups, rewriter); + if (failed(laneMask)) + return rewriter.notifyMatchFailure( + op, "failed to create packed group_store lane mask"); + merged = rewriter + .create(op.getLoc(), firstVRegType, selected, + merged, *laneMask) + .getResult(); + } + + int64_t activeGroups = + std::min(32, numGroups - blockStart); + FailureOr storeMask = createPrefixMaskForActiveLanes( + op.getLoc(), *maskType, activeGroups, rewriter); + if (failed(storeMask)) + return rewriter.notifyMatchFailure( + op, "failed to create packed group_store store mask"); + Value groupOffset = createGroupChunkOffset( + op.getLoc(), *offset, *rowStride, blockStart / 4, + /*chunkLaneOffset=*/0, rewriter); + rewriter.create( + op.getLoc(), /*updated_base=*/Type{}, merged, *destination, + groupOffset, rewriter.getStringAttr("PK4_B32"), *storeMask); + } + + rewriter.eraseOp(op); + return success(); + } + } + + if (layout.hasLaneStride()) { + std::optional dist = getLaneStrideStoreDistToken( + layout, valueVMIType.getElementType()); + std::optional maskGranularity = + getLaneStrideStoreMaskGranularity( + layout, valueVMIType.getElementType()); + if (!dist || !maskGranularity) + return rewriter.notifyMatchFailure( + op, "unsupported slots=8 lane_stride group_store packing"); + + auto maskType = + MaskType::get(rewriter.getContext(), *maskGranularity); + for (auto [slotBlock, value] : llvm::enumerate(valueParts)) { + if (!isa(value.getType())) + return rewriter.notifyMatchFailure( + op, "group_store value must be vreg"); + int64_t activeGroups = + std::min(8, numGroups - slotBlock * 8); + FailureOr mask = createPrefixMaskForActiveLanes( + op.getLoc(), maskType, activeGroups, rewriter); + if (failed(mask)) + return rewriter.notifyMatchFailure( + op, "failed to create packed slots=8 group_store mask"); + Value groupOffset = createGroupChunkOffset( + op.getLoc(), *offset, *rowStride, slotBlock * 8, + /*chunkLaneOffset=*/0, rewriter); + rewriter.create( + op.getLoc(), /*updated_base=*/Type{}, value, *destination, + groupOffset, rewriter.getStringAttr(*dist), *mask); + } + + rewriter.eraseOp(op); + return success(); + } + + for (auto [slotBlock, value] : llvm::enumerate(valueParts)) { + auto vregType = dyn_cast(value.getType()); + if (!vregType) + return rewriter.notifyMatchFailure(op, + "group_store value must be vreg"); + FailureOr maskType = + getMaskTypeForVReg(vregType, rewriter.getContext()); + if (failed(maskType)) + return rewriter.notifyMatchFailure( + op, "unsupported element type for group_store mask"); + int64_t activeGroups = std::min(8, numGroups - slotBlock * 8); + FailureOr mask = createPrefixMaskForActiveLanes( + op.getLoc(), *maskType, activeGroups, rewriter); + if (failed(mask)) + return rewriter.notifyMatchFailure( + op, "failed to create slots=8 group_store mask"); + Value groupOffset = createGroupChunkOffset( + op.getLoc(), *offset, *rowStride, slotBlock * 8, + /*chunkLaneOffset=*/0, rewriter); + rewriter.create(op.getLoc(), + /*updated_base=*/Type{}, value, *destination, + groupOffset, /*dist=*/nullptr, *mask); + } + + rewriter.eraseOp(op); + return success(); + } + + int64_t lanesPerPart = 0; + int64_t groupCount = 0; + int64_t chunksPerGroup = 0; + FailureOr groupSize = + getGroupSizeFromNumGroups(valueVMIType, op.getNumGroupsAttr().getInt()); + if (failed(groupSize)) + return rewriter.notifyMatchFailure( + op, "group_store requires num_groups to evenly divide lane count"); + + int64_t d2LanesPerPart = 0; + int64_t d2GroupCount = 0; + int64_t d2ChunksPerGroupPerPart = 0; + std::string d2Reason; + if (succeeded(checkDeinterleaved2GroupStoreChunkShape( + valueVMIType, *groupSize, &d2LanesPerPart, &d2GroupCount, + &d2ChunksPerGroupPerPart, &d2Reason))) { + std::optional dist = + getX2MemoryDistToken(valueVMIType.getElementType(), "INTLV"); + if (!dist) + return rewriter.notifyMatchFailure( + op, "group_store requires vstsx2 INTLV element support"); + + ValueRange valueParts = adaptor.getValue(); + int64_t chunksPerPart = d2GroupCount * d2ChunksPerGroupPerPart; + if (static_cast(valueParts.size()) != 2 * chunksPerPart) + return rewriter.notifyMatchFailure( + op, "deinterleaved=2 group_store arity mismatch"); + + for (int64_t group = 0; group < d2GroupCount; ++group) { + for (int64_t chunk = 0; chunk < d2ChunksPerGroupPerPart; ++chunk) { + int64_t lowIndex = group * d2ChunksPerGroupPerPart + chunk; + int64_t highIndex = chunksPerPart + lowIndex; + Value low = valueParts[lowIndex]; + Value high = valueParts[highIndex]; + if (low.getType() != high.getType()) + return rewriter.notifyMatchFailure( + op, "vstsx2 group_store requires matching low/high types"); + auto vregType = dyn_cast(low.getType()); + if (!vregType) + return rewriter.notifyMatchFailure(op, + "group_store value must be vreg"); + FailureOr mask = + createAllTrueMaskForVReg(op.getLoc(), vregType, rewriter); + if (failed(mask)) + return rewriter.notifyMatchFailure( + op, "unsupported element type for group_store mask"); + Value chunkOffset = createGroupChunkOffset( + op.getLoc(), *offset, *rowStride, group, + chunk * 2 * d2LanesPerPart, rewriter); + rewriter.create(op.getLoc(), low, high, *destination, + chunkOffset, rewriter.getStringAttr(*dist), + *mask); + } + } + + rewriter.eraseOp(op); + return success(); + } + + if (failed(checkContiguousFullGroupChunks(op, valueVMIType, *groupSize, + &lanesPerPart, &groupCount, + &chunksPerGroup, rewriter))) + return failure(); + + ValueRange valueParts = adaptor.getValue(); + if (static_cast(valueParts.size()) != groupCount * chunksPerGroup) + return rewriter.notifyMatchFailure(op, "group_store arity mismatch"); + + for (auto [index, value] : llvm::enumerate(valueParts)) { + auto vregType = dyn_cast(value.getType()); + if (!vregType) + return rewriter.notifyMatchFailure(op, + "group_store value must be vreg"); + FailureOr mask = + createAllTrueMaskForVReg(op.getLoc(), vregType, rewriter); + if (failed(mask)) + return rewriter.notifyMatchFailure( + op, "unsupported element type for group_store mask"); + int64_t group = index / chunksPerGroup; + int64_t chunkInGroup = index % chunksPerGroup; + Value chunkOffset = + createGroupChunkOffset(op.getLoc(), *offset, *rowStride, group, + chunkInGroup * lanesPerPart, rewriter); + rewriter.create(op.getLoc(), + /*updated_base=*/Type{}, value, *destination, + chunkOffset, /*dist=*/nullptr, *mask); + } + + rewriter.eraseOp(op); + return success(); + } +}; + +struct OneToNVMIMaskedStoreOpPattern + : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIMaskedStoreOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto valueVMIType = cast(op.getValue().getType()); + FailureOr lanesPerPart = + getDataLanesPerPart(valueVMIType.getElementType()); + if (failed(lanesPerPart)) + return rewriter.notifyMatchFailure( + op, "masked_store requires known physical lanes per part"); + + FailureOr destination = getSingleValue( + op, adaptor.getDestination(), + "masked_store destination must convert to one value", rewriter); + FailureOr offset = getSingleValue( + op, adaptor.getOffset(), + "masked_store offset must convert to one value", rewriter); + if (failed(destination) || failed(offset)) + return failure(); + + ValueRange valueParts = adaptor.getValue(); + ValueRange maskParts = adaptor.getMask(); + if (valueParts.size() != maskParts.size()) + return rewriter.notifyMatchFailure( + op, "masked_store value/mask physical arity mismatch"); + + auto maskVMIType = cast(op.getMask().getType()); + if (std::optional dist = + getDenseLaneStrideStoreDistToken(valueVMIType)) { + std::optional maskGranularity = + getDenseLaneStrideMaskedStoreMaskGranularity(valueVMIType); + VMILayoutAttr valueLayout = valueVMIType.getLayoutAttr(); + VMILayoutAttr maskLayout = maskVMIType.getLayoutAttr(); + if (maskGranularity && valueLayout && maskLayout && + valueLayout == maskLayout) { + int64_t semanticOffset = 0; + for (auto [index, valueAndMask] : + llvm::enumerate(llvm::zip_equal(valueParts, maskParts))) { + auto [value, mask] = valueAndMask; + auto vregType = dyn_cast(value.getType()); + if (!vregType || !isa(mask.getType())) + return rewriter.notifyMatchFailure( + op, "lane_stride masked_store parts must be vreg/mask"); + FailureOr activeLanes = + getActiveDataLanesInPhysicalChunk(valueVMIType, index); + if (failed(activeLanes)) + return rewriter.notifyMatchFailure( + op, "failed to compute lane_stride masked_store active lanes"); + if (*activeLanes == 0) + continue; + FailureOr storeMask = createDenseLaneStrideStorePredicate( + op.getLoc(), valueVMIType, index, mask, *maskGranularity, + rewriter); + if (failed(storeMask)) + return rewriter.notifyMatchFailure( + op, "failed to compact lane_stride masked_store predicate"); + Value chunkOffset = + createChunkOffset(op.getLoc(), *offset, semanticOffset, rewriter); + rewriter.create(op.getLoc(), + /*updated_base=*/Type{}, value, *destination, + chunkOffset, rewriter.getStringAttr(*dist), + *storeMask); + semanticOffset += *activeLanes; + } + + rewriter.eraseOp(op); + return success(); + } + } + + SmallVector contiguousValueTypes; + contiguousValueTypes.reserve(valueParts.size()); + for (Value value : valueParts) + contiguousValueTypes.push_back(value.getType()); + FailureOr> storeParts = materializeDataLayoutConversion( + op, valueParts, contiguousValueTypes, valueVMIType.getLayoutAttr(), + VMILayoutAttr::getContiguous(rewriter.getContext()), rewriter); + if (failed(storeParts)) + return failure(); + + SmallVector contiguousMaskTypes; + contiguousMaskTypes.reserve(maskParts.size()); + for (Value mask : maskParts) + contiguousMaskTypes.push_back(mask.getType()); + FailureOr> storeMasks = materializeMaskLayoutConversion( + op, maskParts, contiguousMaskTypes, maskVMIType.getLayoutAttr(), + VMILayoutAttr::getContiguous(rewriter.getContext()), rewriter); + if (failed(storeMasks)) + return failure(); + + if (storeParts->size() != storeMasks->size()) + return rewriter.notifyMatchFailure( + op, "masked_store converted value/mask arity mismatch"); + + for (auto [index, valueAndMask] : + llvm::enumerate(llvm::zip_equal(*storeParts, *storeMasks))) { + auto [value, mask] = valueAndMask; + auto vregType = dyn_cast(value.getType()); + if (!vregType || !isa(mask.getType())) + return rewriter.notifyMatchFailure( + op, "masked_store converted parts must be vreg/mask"); + FailureOr activeLanes = + getContiguousActiveDataLanes(valueVMIType, index); + if (failed(activeLanes)) + return rewriter.notifyMatchFailure( + op, "failed to compute masked_store active lanes"); + if (*activeLanes == 0) + continue; + FailureOr storeMask = createMaskedStorePredicate( + op.getLoc(), valueVMIType, index, mask, vregType, rewriter); + if (failed(storeMask)) + return rewriter.notifyMatchFailure( + op, "failed to materialize masked_store predicate"); + Value chunkOffset = createChunkOffset(op.getLoc(), *offset, + index * *lanesPerPart, rewriter); + rewriter.create(op.getLoc(), + /*updated_base=*/Type{}, value, *destination, + chunkOffset, /*dist=*/nullptr, *storeMask); + } + + rewriter.eraseOp(op); + return success(); + } +}; + +struct OneToNVMIGroupBroadcastLoadOpPattern + : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIGroupBroadcastLoadOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto resultVMIType = cast(op.getResult().getType()); + int64_t numGroups = op.getNumGroupsAttr().getInt(); + FailureOr source = getSingleValue( + op, adaptor.getSource(), + "group_broadcast_load source must convert to one value", rewriter); + FailureOr offset = getSingleValue( + op, adaptor.getOffset(), + "group_broadcast_load offset must convert to one value", rewriter); + FailureOr sourceGroupStride = getSingleValue( + op, adaptor.getSourceGroupStride(), + "group_broadcast_load source_group_stride must convert to one value", + rewriter); + if (failed(source) || failed(offset) || failed(sourceGroupStride)) + return failure(); + + VMILayoutSupport supports; + std::string supportReason; + FailureOr loadFact = + supports.getGroupBroadcastLoadLayoutFact(op, &supportReason); + if (failed(loadFact)) + return rewriter.notifyMatchFailure( + op, Twine("group_broadcast_load has no registered support: ") + + supportReason); + + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + + SmallVector resultTypes = std::move(*maybe_resultTypes); + FailureOr directFact = + supports.getGroupBroadcastLoadDirectFact(op); + auto getBRCDist = [&]() -> std::optional { + unsigned elementBits = + pto::getPTOStorageElemBitWidth(resultVMIType.getElementType()); + if (elementBits == 8) + return StringRef("BRC_B8"); + if (elementBits == 16) + return StringRef("BRC_B16"); + if (elementBits == 32) + return StringRef("BRC_B32"); + return std::nullopt; + }; + + if (succeeded(directFact) && + directFact->kind == VMIGroupBroadcastLoadDirectKind::BRC) { + int64_t chunksPerGroup = + directFact->layout.groupSize / directFact->layout.lanesPerPart; + std::optional brcDist = getBRCDist(); + if (!brcDist) + return rewriter.notifyMatchFailure( + op, "group_broadcast_load BRC lowering requires b8/b16/b32 " + "element type"); + if (!isa((*source).getType())) + return rewriter.notifyMatchFailure( + op, "group_broadcast_load BRC lowering requires !pto.ptr source"); + if (chunksPerGroup <= 0 || + static_cast(resultTypes.size()) != + numGroups * chunksPerGroup) + return rewriter.notifyMatchFailure( + op, "group_broadcast_load BRC physical arity mismatch"); + + SmallVector results; + results.reserve(resultTypes.size()); + for (auto [index, resultType] : llvm::enumerate(resultTypes)) { + auto vregType = dyn_cast(resultType); + if (!vregType) + return rewriter.notifyMatchFailure( + op, "group_broadcast_load BRC result must be vreg"); + int64_t group = static_cast(index) / chunksPerGroup; + Value groupOffset = + createGroupChunkOffset(op.getLoc(), *offset, *sourceGroupStride, + group, /*inGroupLaneOffset=*/0, rewriter); + results.push_back(rewriter + .create(op.getLoc(), resultType, + /*updated_base=*/Type{}, *source, + groupOffset, + rewriter.getStringAttr(*brcDist)) + .getResult()); + } + replaceOpWithFlatConvertedValues(rewriter, op, results, + *this->getTypeConverter()); + return success(); + } + + if (failed(directFact) || + directFact->kind != VMIGroupBroadcastLoadDirectKind::E2B) { + std::optional stride = + getConstantIndexValue(op.getSourceGroupStride()); + int64_t slots = (stride && *stride == 1) ? 8 : 1; + auto sourceVMIType = VMIVRegType::get( + rewriter.getContext(), numGroups, resultVMIType.getElementType(), + VMILayoutAttr::getGroupSlots(rewriter.getContext(), numGroups, + slots)); + + FailureOr sourceArity = getVMIPhysicalArity(sourceVMIType); + FailureOr sourceElementType = + getVMIVRegPhysicalElementType(sourceVMIType); + if (failed(sourceArity) || failed(sourceElementType)) + return rewriter.notifyMatchFailure( + op, "group_broadcast_load fallback cannot derive physical types"); + + SmallVector sourceTypes; + sourceTypes.reserve(*sourceArity); + FailureOr sourceLanesPerPart = + getDataLanesPerPart(*sourceElementType); + if (failed(sourceLanesPerPart)) + return rewriter.notifyMatchFailure( + op, "group_broadcast_load fallback cannot derive source lanes"); + for (int64_t i = 0; i < *sourceArity; ++i) + sourceTypes.push_back(VRegType::get( + rewriter.getContext(), *sourceLanesPerPart, *sourceElementType)); + + SmallVector sourceParts; + if (failed(lowerGroupSlotLoadParts( + op, *source, *offset, *sourceGroupStride, sourceVMIType, + sourceTypes, numGroups, rewriter, sourceParts))) + return failure(); + + SmallVector results; + if (failed(lowerGroupBroadcastParts(op, sourceParts, sourceVMIType, + resultVMIType, resultTypes, numGroups, + rewriter, results))) + return failure(); + replaceOpWithFlatConvertedValues(rewriter, op, results, + *this->getTypeConverter()); + return success(); + } + + VMILayoutAttr layout = resultVMIType.getLayoutAttr(); + bool contiguousPacketLayout = layout && layout.isContiguous(); + bool splitPacketLayout = layout && layout.isDeinterleaved() && + (layout.getFactor() == 2 || + layout.getFactor() == 4) && + layout.getBlockElems() == 1; + if (!contiguousPacketLayout && !splitPacketLayout) + return rewriter.notifyMatchFailure( + op, "group_broadcast_load E2B lowering requires " + "contiguous result layout for direct group size or " + "deinterleaved=2/4, block_elems=1 result layout for split " + "group size"); + + unsigned elementBits = directFact->layout.elementBits; + if (elementBits != 16 && elementBits != 32) + return rewriter.notifyMatchFailure( + op, "group_broadcast_load E2B lowering requires b16 or b32 " + "element type"); + StringRef e2bDist = elementBits == 16 ? "E2B_B16" : "E2B_B32"; + + std::optional stride = + getConstantIndexValue(op.getSourceGroupStride()); + if (!stride || *stride != 1) + return rewriter.notifyMatchFailure( + op, "group_broadcast_load E2B lowering requires constant unit " + "source_group_stride"); + + if (!isa((*source).getType())) + return rewriter.notifyMatchFailure( + op, "group_broadcast_load E2B lowering requires !pto.ptr source"); + + if (numGroups != 8) + return rewriter.notifyMatchFailure( + op, "group_broadcast_load E2B lowering requires num_groups = 8"); + + FailureOr chunksPerPart = getDataChunksInPart(resultVMIType, 0); + if (failed(chunksPerPart) || *chunksPerPart <= 0) + return rewriter.notifyMatchFailure( + op, "group_broadcast_load requires known chunks per part"); + int64_t factor = layout.getFactor(); + for (int64_t part = 1; part < factor; ++part) { + FailureOr currentChunks = + getDataChunksInPart(resultVMIType, part); + if (failed(currentChunks) || *currentChunks != *chunksPerPart) + return rewriter.notifyMatchFailure( + op, "group_broadcast_load requires uniform chunks per part"); + } + if (static_cast(resultTypes.size()) != factor * *chunksPerPart) + return rewriter.notifyMatchFailure( + op, "group_broadcast_load physical arity mismatch"); + if (*chunksPerPart != 1) + return rewriter.notifyMatchFailure( + op, + "group_broadcast_load expected one E2B packet in each part"); + + SmallVector packets; + packets.reserve(*chunksPerPart); + for (int64_t chunk = 0; chunk < *chunksPerPart; ++chunk) { + Type packetType = resultTypes[chunk]; + auto vregType = dyn_cast(packetType); + if (!vregType) + return rewriter.notifyMatchFailure( + op, "group_broadcast_load result must be vreg"); + Value packetOffset = + createChunkOffset(op.getLoc(), *offset, chunk * 8, rewriter); + packets.push_back(rewriter + .create(op.getLoc(), packetType, + /*updated_base=*/Type{}, *source, + packetOffset, + rewriter.getStringAttr(e2bDist)) + .getResult()); + } + + SmallVector results; + results.reserve(resultTypes.size()); + for (int64_t part = 0; part < factor; ++part) { + for (int64_t chunk = 0; chunk < *chunksPerPart; ++chunk) { + int64_t flatIndex = part * *chunksPerPart + chunk; + if (resultTypes[flatIndex] != resultTypes[chunk]) + return rewriter.notifyMatchFailure( + op, "group_broadcast_load E2B reused packet type mismatch"); + results.push_back(packets[chunk]); + } + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } + +private: + ; +}; + +struct OneToNVMIStrideLoadOpPattern + : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIStrideLoadOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + FailureOr source = getSingleValue( + op, adaptor.getSource(), "stride_load source must convert to one value", + rewriter); + FailureOr offset = getSingleValue( + op, adaptor.getOffset(), "stride_load offset must convert to one value", + rewriter); + FailureOr blockStride = getSingleValue( + op, adaptor.getBlockStride(), + "stride_load block_stride must convert to one value", rewriter); + FailureOr repeatStride = getSingleValue( + op, adaptor.getRepeatStride(), + "stride_load repeat_stride must convert to one value", rewriter); + if (failed(source) || failed(offset) || failed(blockStride) || + failed(repeatStride)) + return failure(); + + ValueRange maskParts = adaptor.getMask(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + if (resultTypes.size() != 1 || maskParts.size() != 1) + return rewriter.notifyMatchFailure( + op, "stride_load supports one physical result/mask chunk"); + auto resultType = dyn_cast(resultTypes.front()); + if (!resultType || !isa(maskParts.front().getType())) + return rewriter.notifyMatchFailure( + op, "stride_load requires physical vreg/mask parts"); + + Value base = rewriter + .create(op.getLoc(), (*source).getType(), + *source, *offset) + .getResult(); + Value loaded = + rewriter + .create(op.getLoc(), resultType, base, *blockStride, + *repeatStride, maskParts.front()) + .getResult(); + replaceOpWithFlatConvertedValues(rewriter, op, SmallVector{loaded}, + *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMIStrideStoreOpPattern + : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIStrideStoreOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + FailureOr destination = getSingleValue( + op, adaptor.getDestination(), + "stride_store destination must convert to one value", rewriter); + FailureOr offset = getSingleValue( + op, adaptor.getOffset(), + "stride_store offset must convert to one value", rewriter); + FailureOr blockStride = getSingleValue( + op, adaptor.getBlockStride(), + "stride_store block_stride must convert to one value", rewriter); + FailureOr repeatStride = getSingleValue( + op, adaptor.getRepeatStride(), + "stride_store repeat_stride must convert to one value", rewriter); + if (failed(destination) || failed(offset) || failed(blockStride) || + failed(repeatStride)) + return failure(); + + ValueRange valueParts = adaptor.getValue(); + ValueRange maskParts = adaptor.getMask(); + if (valueParts.size() != 1 || maskParts.size() != 1) + return rewriter.notifyMatchFailure( + op, "stride_store supports one physical value/mask chunk"); + if (!isa(valueParts.front().getType()) || + !isa(maskParts.front().getType())) + return rewriter.notifyMatchFailure( + op, "stride_store requires physical vreg/mask parts"); + + Value base = rewriter + .create(op.getLoc(), (*destination).getType(), + *destination, *offset) + .getResult(); + rewriter.create(op.getLoc(), base.getType(), valueParts.front(), + base, *blockStride, *repeatStride, + maskParts.front()); + rewriter.eraseOp(op); + return success(); + } +}; + +struct OneToNVMIScatterOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIScatterOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + FailureOr destination = getSingleValue( + op, adaptor.getDestination(), + "scatter destination must convert to one value", rewriter); + if (failed(destination)) + return failure(); + + ValueRange valueParts = adaptor.getValue(); + ValueRange indicesParts = adaptor.getIndices(); + ValueRange maskParts = adaptor.getMask(); + if (valueParts.size() != indicesParts.size() || + valueParts.size() != maskParts.size()) + return rewriter.notifyMatchFailure(op, "scatter physical arity mismatch"); + + for (auto [value, indices, mask] : + llvm::zip_equal(valueParts, indicesParts, maskParts)) { + if (!isa(value.getType()) || + !isa(indices.getType()) || !isa(mask.getType())) + return rewriter.notifyMatchFailure( + op, "scatter physical part type mismatch"); + rewriter.create(op.getLoc(), value, *destination, indices, + mask); + } + + rewriter.eraseOp(op); + return success(); + } +}; + +template +struct OneToNVMIBinaryOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite( + SourceOp op, + typename OpConversionPattern::OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + ValueRange lhsParts = adaptor.getLhs(); + ValueRange rhsParts = adaptor.getRhs(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + if (lhsParts.size() != rhsParts.size() || + lhsParts.size() != resultTypes.size()) + return rewriter.notifyMatchFailure(op, "physical binary arity mismatch"); + + SmallVector results; + results.reserve(resultTypes.size()); + for (auto [lhs, rhs, resultType] : + llvm::zip_equal(lhsParts, rhsParts, resultTypes)) { + auto vregType = dyn_cast(resultType); + if (!vregType || lhs.getType() != resultType || + rhs.getType() != resultType) + return rewriter.notifyMatchFailure( + op, "physical binary part type mismatch"); + FailureOr mask = + createAllTrueMaskForVReg(op.getLoc(), vregType, rewriter); + if (failed(mask)) + return rewriter.notifyMatchFailure( + op, "unsupported element type for all-true binary mask"); + results.push_back( + rewriter.create(op.getLoc(), resultType, lhs, rhs, *mask) + .getResult()); + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMIVmullOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIVmullOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + ValueRange aParts = adaptor.getA(); + ValueRange bParts = adaptor.getB(); + ValueRange maskParts = adaptor.getMask(); + FailureOr> maybeLowTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + FailureOr> maybeHighTypes = + getConvertedResultTypes(op, 1, *this->getTypeConverter()); + if (failed(maybeLowTypes) || failed(maybeHighTypes)) + return failure(); + SmallVector lowTypes = std::move(*maybeLowTypes); + SmallVector highTypes = std::move(*maybeHighTypes); + + size_t arity = aParts.size(); + if (arity == 0 || bParts.size() != arity || maskParts.size() != arity || + lowTypes.size() != arity || highTypes.size() != arity) + return rewriter.notifyMatchFailure( + op, "physical vmull arity mismatch across a, b, mask, low, and high"); + + SmallVector lows; + SmallVector highs; + lows.reserve(arity); + highs.reserve(arity); + for (size_t index = 0; index < arity; ++index) { + Type lowType = lowTypes[index]; + Type highType = highTypes[index]; + auto dataType = dyn_cast(lowType); + auto maskType = dyn_cast(maskParts[index].getType()); + if (!dataType || dataType.getElementCount() != 64 || + lowType != highType || aParts[index].getType() != lowType || + bParts[index].getType() != lowType) + return rewriter.notifyMatchFailure( + op, "vmull requires matching 64-lane physical data part types"); + auto elementType = dyn_cast(dataType.getElementType()); + if (!elementType || elementType.getWidth() != 32 || + (!elementType.isSignless() && !elementType.isUnsigned())) + return rewriter.notifyMatchFailure( + op, "vmull requires physical i32 or ui32 data parts"); + if (!maskType || !maskType.isB32()) + return rewriter.notifyMatchFailure( + op, "vmull requires a corresponding b32 mask part"); + + auto vmull = rewriter.create(op.getLoc(), lowType, highType, + aParts[index], bParts[index], + maskParts[index]); + lows.push_back(vmull.getLow()); + highs.push_back(vmull.getHigh()); + } + + SmallVector results; + results.reserve(lows.size() + highs.size()); + results.append(lows); + results.append(highs); + replaceOpWithFlatConvertedValues(rewriter, op, results, + *this->getTypeConverter()); + return success(); + } +}; + +template +struct OneToNVMIInterleaveOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite( + SourceOp op, + typename OpConversionPattern::OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + ValueRange lhsParts = adaptor.getLhs(); + ValueRange rhsParts = adaptor.getRhs(); + ValueRange maskParts = adaptor.getMask(); + FailureOr> maybeLowTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + FailureOr> maybeHighTypes = + getConvertedResultTypes(op, 1, *this->getTypeConverter()); + if (failed(maybeLowTypes) || failed(maybeHighTypes)) + return failure(); + SmallVector lowTypes = std::move(*maybeLowTypes); + SmallVector highTypes = std::move(*maybeHighTypes); + if (lhsParts.size() != rhsParts.size() || + lhsParts.size() != lowTypes.size() || + lhsParts.size() != highTypes.size()) + return rewriter.notifyMatchFailure(op, + "physical interleave arity mismatch"); + if (!maskParts.empty() && maskParts.size() != lhsParts.size()) + return rewriter.notifyMatchFailure( + op, "physical interleave mask arity mismatch"); + + auto lhsType = cast(op.getLhs().getType()); + auto rhsType = cast(op.getRhs().getType()); + auto maskType = cast(op.getMask().getType()); + auto lowType = cast(op.getLow().getType()); + auto highType = cast(op.getHigh().getType()); + VMILayoutSupport supports; + FailureOr fact; + if constexpr (std::is_same_v) { + fact = supports.getVintlvLayoutFactForLayouts( + lhsType, rhsType, maskType, lowType, highType); + } else { + fact = supports.getVdintlvLayoutFactForLayouts( + lhsType, rhsType, maskType, lowType, highType); + } + if (failed(fact)) + return rewriter.notifyMatchFailure( + op, "unsupported interleave layout relation"); + + auto isContiguous = [](VMILayoutAttr layout) { + return layout && layout.isContiguous() && layout.getLaneStride() == 1; + }; + auto getElementDeintFactor = [](VMILayoutAttr layout) -> int64_t { + if (layout && layout.isContiguous() && layout.getLaneStride() == 1) + return 1; + if (layout && layout.isDeinterleaved() && + layout.getBlockElems() == 1 && layout.getLaneStride() == 1) + return layout.getFactor(); + return 0; + }; + + bool allContiguous = isContiguous(fact->lhsLayout) && + isContiguous(fact->rhsLayout) && + isContiguous(fact->maskLayout) && + isContiguous(fact->lowLayout) && + isContiguous(fact->highLayout); + if (allContiguous) { + if (lhsParts.size() != 1 || rhsParts.size() != 1 || + lowTypes.size() != 1 || highTypes.size() != 1) + return rewriter.notifyMatchFailure( + op, "single-chunk interleave expects one physical part"); + if (!maskParts.empty() && !isa(maskParts.front().getType())) + return rewriter.notifyMatchFailure( + op, "single-chunk interleave mask part type mismatch"); + if (!isa(lowTypes.front()) || + !isa(highTypes.front()) || + lhsParts.front().getType() != lowTypes.front() || + rhsParts.front().getType() != lowTypes.front() || + highTypes.front() != lowTypes.front()) + return rewriter.notifyMatchFailure( + op, "single-chunk interleave part type mismatch"); + auto interleave = rewriter.create( + op.getLoc(), lowTypes.front(), highTypes.front(), lhsParts.front(), + rhsParts.front()); + SmallVector directResults = {interleave.getLow(), + interleave.getHigh()}; + replaceOpWithFlatConvertedValues(rewriter, op, directResults, + *this->getTypeConverter()); + return success(); + } + + int64_t inputFactor = getElementDeintFactor(fact->lhsLayout); + int64_t outputFactor = getElementDeintFactor(fact->lowLayout); + bool zeroCopyVintlv = std::is_same_v && + inputFactor > 0 && + fact->rhsLayout == fact->lhsLayout && + fact->maskLayout == fact->lhsLayout && + fact->highLayout == fact->lowLayout && + outputFactor == 2 * inputFactor; + bool zeroCopyVdintlv = std::is_same_v && + inputFactor > 0 && + fact->rhsLayout == fact->lhsLayout && + fact->maskLayout == fact->lhsLayout && + fact->highLayout == fact->lowLayout && + inputFactor == 2 * outputFactor; + if (!zeroCopyVintlv && !zeroCopyVdintlv) + return rewriter.notifyMatchFailure( + op, "unsupported interleave physical layout relation"); + + SmallVector results; + results.reserve(lhsParts.size() + rhsParts.size()); + if (zeroCopyVintlv) { + if (lhsParts.empty() || lhsParts.size() % (2 * inputFactor) != 0) + return rewriter.notifyMatchFailure( + op, "zero-copy vintlv expects input groups with even chunk count"); + size_t groupChunks = lhsParts.size() / inputFactor; + size_t halfGroupChunks = groupChunks / 2; + for (int64_t group = 0; group < inputFactor; ++group) { + size_t offset = group * groupChunks; + llvm::append_range( + results, + lhsParts.slice(offset, halfGroupChunks)); + llvm::append_range( + results, + rhsParts.slice(offset, halfGroupChunks)); + } + for (int64_t group = 0; group < inputFactor; ++group) { + size_t offset = group * groupChunks + halfGroupChunks; + llvm::append_range( + results, + lhsParts.slice(offset, halfGroupChunks)); + llvm::append_range( + results, + rhsParts.slice(offset, halfGroupChunks)); + } + } else { + if (lhsParts.empty() || lhsParts.size() % inputFactor != 0) + return rewriter.notifyMatchFailure( + op, "zero-copy vdintlv expects complete input layout groups"); + size_t groupChunks = lhsParts.size() / inputFactor; + for (int64_t group = 0; group < outputFactor; ++group) { + size_t offset = 2 * group * groupChunks; + llvm::append_range(results, lhsParts.slice(offset, groupChunks)); + llvm::append_range(results, rhsParts.slice(offset, groupChunks)); + } + for (int64_t group = 0; group < outputFactor; ++group) { + size_t offset = (2 * group + 1) * groupChunks; + llvm::append_range(results, lhsParts.slice(offset, groupChunks)); + llvm::append_range(results, rhsParts.slice(offset, groupChunks)); + } + } + + SmallVector resultTypes; + resultTypes.reserve(lowTypes.size() + highTypes.size()); + llvm::append_range(resultTypes, lowTypes); + llvm::append_range(resultTypes, highTypes); + if (results.size() != resultTypes.size()) + return rewriter.notifyMatchFailure( + op, "zero-copy interleave result arity mismatch"); + for (auto [value, resultType] : llvm::zip_equal(results, resultTypes)) { + if (value.getType() != resultType) + return rewriter.notifyMatchFailure( + op, "zero-copy interleave part type mismatch"); + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, + *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMIFmaOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIFmaOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + ValueRange lhsParts = adaptor.getLhs(); + ValueRange rhsParts = adaptor.getRhs(); + ValueRange accParts = adaptor.getAcc(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + if (lhsParts.size() != rhsParts.size() || + lhsParts.size() != accParts.size() || + lhsParts.size() != resultTypes.size()) + return rewriter.notifyMatchFailure(op, "fma physical arity mismatch"); + + SmallVector results; + results.reserve(resultTypes.size()); + for (auto [lhs, rhs, acc, resultType] : + llvm::zip_equal(lhsParts, rhsParts, accParts, resultTypes)) { + auto vregType = dyn_cast(resultType); + if (!vregType || lhs.getType() != resultType || + rhs.getType() != resultType || acc.getType() != resultType) + return rewriter.notifyMatchFailure( + op, "fma requires matching physical vreg parts"); + FailureOr mask = + createAllTrueMaskForVReg(op.getLoc(), vregType, rewriter); + if (failed(mask)) + return rewriter.notifyMatchFailure(op, + "unsupported element type for fma"); + results.push_back( + rewriter + .create(op.getLoc(), resultType, acc, lhs, rhs, *mask) + .getResult()); + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +template +struct OneToNVMIUnaryOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite( + SourceOp op, + typename OpConversionPattern::OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + ValueRange sourceParts = adaptor.getSource(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + if (sourceParts.size() != resultTypes.size()) + return rewriter.notifyMatchFailure(op, "physical unary arity mismatch"); + + SmallVector results; + results.reserve(resultTypes.size()); + for (auto [source, resultType] : + llvm::zip_equal(sourceParts, resultTypes)) { + auto vregType = dyn_cast(resultType); + if (!vregType || source.getType() != resultType) + return rewriter.notifyMatchFailure(op, + "physical unary part type mismatch"); + FailureOr mask = + createAllTrueMaskForVReg(op.getLoc(), vregType, rewriter); + if (failed(mask)) + return rewriter.notifyMatchFailure( + op, "unsupported element type for all-true unary mask"); + results.push_back( + rewriter.create(op.getLoc(), resultType, source, *mask) + .getResult()); + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +template +struct OneToNVMIMaskBinaryOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite( + SourceOp op, + typename OpConversionPattern::OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + ValueRange lhsParts = adaptor.getLhs(); + ValueRange rhsParts = adaptor.getRhs(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + if (lhsParts.size() != rhsParts.size() || + lhsParts.size() != resultTypes.size()) + return rewriter.notifyMatchFailure(op, + "physical mask binary arity mismatch"); + + SmallVector results; + results.reserve(resultTypes.size()); + for (auto [lhs, rhs, resultType] : + llvm::zip_equal(lhsParts, rhsParts, resultTypes)) { + auto maskType = dyn_cast(resultType); + if (!maskType || lhs.getType() != resultType || + rhs.getType() != resultType) + return rewriter.notifyMatchFailure( + op, "physical mask binary part type mismatch"); + FailureOr seedMask = + createAllTrueMask(op.getLoc(), maskType, rewriter); + if (failed(seedMask)) + return rewriter.notifyMatchFailure( + op, "unsupported mask type for all-true mask binary seed"); + results.push_back( + rewriter + .create(op.getLoc(), resultType, lhs, rhs, *seedMask) + .getResult()); + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +template +struct OneToNVMIMaskUnaryOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite( + SourceOp op, + typename OpConversionPattern::OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + ValueRange sourceParts = adaptor.getSource(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + if (sourceParts.size() != resultTypes.size()) + return rewriter.notifyMatchFailure(op, + "physical mask unary arity mismatch"); + + SmallVector results; + results.reserve(resultTypes.size()); + for (auto [source, resultType] : + llvm::zip_equal(sourceParts, resultTypes)) { + auto maskType = dyn_cast(resultType); + if (!maskType || source.getType() != resultType) + return rewriter.notifyMatchFailure( + op, "physical mask unary part type mismatch"); + FailureOr seedMask = + createAllTrueMask(op.getLoc(), maskType, rewriter); + if (failed(seedMask)) + return rewriter.notifyMatchFailure( + op, "unsupported mask type for all-true mask unary seed"); + results.push_back( + rewriter.create(op.getLoc(), resultType, source, *seedMask) + .getResult()); + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +template +struct OneToNVMICmpOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite( + SourceOp op, + typename OpConversionPattern::OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + std::optional cmpMode = + getVPTOCmpMode(op.getPredicate()); + if (!cmpMode) + return op.emitOpError() + << kVMIDiagUnsupportedPrefix << "compare predicate " + << op.getPredicate() + << " cannot be lowered to pto.vcmp; supported predicates are " + << getSupportedComparePredicateMessage(); + + ValueRange lhsParts = adaptor.getLhs(); + ValueRange rhsParts = adaptor.getRhs(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + if (lhsParts.size() != rhsParts.size() || + lhsParts.size() != resultTypes.size()) + return rewriter.notifyMatchFailure(op, "physical cmp arity mismatch"); + + SmallVector results; + results.reserve(resultTypes.size()); + for (auto [lhs, rhs, resultType] : + llvm::zip_equal(lhsParts, rhsParts, resultTypes)) { + auto maskType = dyn_cast(resultType); + auto lhsType = dyn_cast(lhs.getType()); + if (!maskType || lhs.getType() != rhs.getType() || !lhsType) + return rewriter.notifyMatchFailure(op, + "physical cmp part type mismatch"); + FailureOr seedMask = + createAllTrueMask(op.getLoc(), maskType, rewriter); + if (failed(seedMask)) + return rewriter.notifyMatchFailure( + op, "unsupported mask type for all-true cmp seed"); + if (cmpMode->signedness) { + FailureOr carrierType = + getSignednessCarrierVRegType(lhsType, *cmpMode->signedness); + if (failed(carrierType)) + return rewriter.notifyMatchFailure( + op, "unsupported integer compare signedness carrier"); + FailureOr carrierLhs = + bitcastVReg(op.getLoc(), lhs, *carrierType, rewriter); + FailureOr carrierRhs = + bitcastVReg(op.getLoc(), rhs, *carrierType, rewriter); + if (failed(carrierLhs) || failed(carrierRhs)) + return rewriter.notifyMatchFailure( + op, "failed to materialize integer compare signedness carrier"); + lhs = *carrierLhs; + rhs = *carrierRhs; + } + results.push_back(rewriter + .create(op.getLoc(), resultType, lhs, rhs, + *seedMask, + rewriter.getStringAttr(cmpMode->mode)) + .getResult()); + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMISelectOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMISelectOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + ValueRange maskParts = adaptor.getMask(); + ValueRange trueParts = adaptor.getTrueValue(); + ValueRange falseParts = adaptor.getFalseValue(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + if (maskParts.size() != trueParts.size() || + trueParts.size() != falseParts.size() || + trueParts.size() != resultTypes.size()) + return rewriter.notifyMatchFailure(op, "physical select arity mismatch"); + + SmallVector results; + results.reserve(resultTypes.size()); + for (auto [mask, trueValue, falseValue, resultType] : + llvm::zip_equal(maskParts, trueParts, falseParts, resultTypes)) { + if (!isa(mask.getType()) || trueValue.getType() != resultType || + falseValue.getType() != resultType || !isa(resultType)) + return rewriter.notifyMatchFailure( + op, "physical select part type mismatch"); + results.push_back(rewriter + .create(op.getLoc(), resultType, trueValue, + falseValue, mask) + .getResult()); + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMIActivePrefixIndexOpPattern + : OpConversionPattern { + using OpConversionPattern< + VMIActivePrefixIndexOp>::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIActivePrefixIndexOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + ValueRange maskParts = adaptor.getMask(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + if (maskParts.size() != 1 || resultTypes.size() != 1) + return rewriter.notifyMatchFailure( + op, "active_prefix_index supports only one physical part"); + + auto resultType = dyn_cast(resultTypes.front()); + auto maskType = dyn_cast(maskParts.front().getType()); + if (!resultType || !maskType) + return rewriter.notifyMatchFailure( + op, "active_prefix_index requires physical vreg/mask parts"); + + auto intType = dyn_cast(resultType.getElementType()); + if (!intType || !intType.isSignless()) + return rewriter.notifyMatchFailure( + op, "active_prefix_index requires signless integer result part"); + + FailureOr seedMask = + createAllTrueMaskForVReg(op.getLoc(), resultType, rewriter); + if (failed(seedMask)) + return rewriter.notifyMatchFailure( + op, "unsupported element type for active_prefix_index seed mask"); + + Value zero = rewriter.create(op.getLoc(), 0, + intType.getWidth()); + Value carrier = + rewriter + .create(op.getLoc(), resultType, zero, *seedMask, + /*position=*/nullptr) + .getResult(); + Value result = rewriter + .create(op.getLoc(), resultType, carrier, + maskParts.front()) + .getResult(); + replaceOpWithFlatConvertedValues(rewriter, op, SmallVector{result}, + *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMICompressOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMICompressOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + ValueRange sourceParts = adaptor.getSource(); + ValueRange maskParts = adaptor.getMask(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + if (sourceParts.size() != 1 || maskParts.size() != 1 || + resultTypes.size() != 1) + return rewriter.notifyMatchFailure( + op, "compress supports only one physical part"); + + auto resultType = dyn_cast(resultTypes.front()); + if (!resultType || sourceParts.front().getType() != resultType || + !isa(maskParts.front().getType())) + return rewriter.notifyMatchFailure( + op, "compress requires physical source/mask/result parts"); + + Value result = rewriter + .create(op.getLoc(), resultType, + sourceParts.front(), maskParts.front()) + .getResult(); + replaceOpWithFlatConvertedValues(rewriter, op, SmallVector{result}, + *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMICompressStoreOpPattern + : OpConversionPattern { + using OpConversionPattern< + VMICompressStoreOp>::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMICompressStoreOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + FailureOr destination = getSingleValue( + op, adaptor.getDestination(), + "compress_store destination must convert to one value", rewriter); + FailureOr offset = getSingleValue( + op, adaptor.getOffset(), + "compress_store offset must convert to one value", rewriter); + if (failed(destination) || failed(offset)) + return failure(); + + ValueRange valueParts = adaptor.getValue(); + ValueRange maskParts = adaptor.getMask(); + if (valueParts.size() != 1 || maskParts.size() != 1) + return rewriter.notifyMatchFailure( + op, "compress_store supports only one physical part"); + + auto valueType = dyn_cast(valueParts.front().getType()); + if (!valueType || !isa(maskParts.front().getType()) || + !isa((*destination).getType())) + return rewriter.notifyMatchFailure( + op, "compress_store requires physical value/mask and ptr " + "destination"); + + Value storeBase = + rewriter + .create(op.getLoc(), (*destination).getType(), + *destination, *offset) + .getResult(); + Value squeezed = rewriter + .create(op.getLoc(), valueType, + valueParts.front(), maskParts.front()) + .getResult(); + auto align = rewriter.create( + op.getLoc(), AlignType::get(rewriter.getContext())); + auto store = rewriter.create( + op.getLoc(), align.getResult().getType(), align.getResult(), squeezed, + storeBase, rewriter.getStringAttr("POST_UPDATE")); + rewriter.create(op.getLoc(), store.getAlignOut(), storeBase); + rewriter.eraseOp(op); + return success(); + } +}; + +struct OneToNVMIReduceAddIOpPattern + : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIReduceAddIOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + ValueRange sourceParts = adaptor.getSource(); + ValueRange initParts = adaptor.getInit(); + ValueRange maskParts = adaptor.getMask(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + if (sourceParts.empty() || sourceParts.size() != maskParts.size() || + initParts.size() != 1 || resultTypes.size() != 1) + return rewriter.notifyMatchFailure( + op, "reduce_addi requires matching source/mask chunks and one " + "init/result chunk"); + + auto resultType = dyn_cast(resultTypes.front()); + auto maskType = dyn_cast(maskParts.front().getType()); + if (!resultType || !maskType || initParts.front().getType() != resultType) + return rewriter.notifyMatchFailure( + op, "reduce_addi requires matching physical source/init/result " + "vregs and one mask"); + + for (Value sourcePart : sourceParts) + if (sourcePart.getType() != resultType) + return rewriter.notifyMatchFailure( + op, "reduce_addi requires every source chunk to match result " + "vreg type"); + for (Value maskPart : maskParts) + if (maskPart.getType() != maskType) + return rewriter.notifyMatchFailure( + op, "reduce_addi requires every mask chunk to have the same " + "predicate type"); + + FailureOr firstLaneMask = + createPrefixMask(op.getLoc(), maskType, "PAT_VL1", rewriter); + if (failed(firstLaneMask)) + return rewriter.notifyMatchFailure( + op, "failed to create reduce_addi first-lane mask"); + + Value accumulator = initParts.front(); + for (auto [sourcePart, maskPart] : + llvm::zip_equal(sourceParts, maskParts)) { + Value reduced = + rewriter + .create(op.getLoc(), resultType, sourcePart, maskPart) + .getResult(); + accumulator = rewriter + .create(op.getLoc(), resultType, reduced, + accumulator, *firstLaneMask) + .getResult(); + } + + replaceOpWithFlatConvertedValues( + rewriter, op, SmallVector{accumulator}, + *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMIReduceAddFOpPattern + : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIReduceAddFOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + ValueRange sourceParts = adaptor.getSource(); + ValueRange initParts = adaptor.getInit(); + ValueRange maskParts = adaptor.getMask(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + if (sourceParts.empty() || sourceParts.size() != maskParts.size() || + initParts.size() != 1 || resultTypes.size() != 1) + return rewriter.notifyMatchFailure( + op, "reduce_addf requires matching source/mask chunks and one " + "init/result chunk"); + + auto resultType = dyn_cast(resultTypes.front()); + auto maskType = dyn_cast(maskParts.front().getType()); + if (!resultType || !maskType || initParts.front().getType() != resultType) + return rewriter.notifyMatchFailure( + op, "reduce_addf requires matching physical source/init/result " + "vregs and one mask"); + + for (Value sourcePart : sourceParts) + if (sourcePart.getType() != resultType) + return rewriter.notifyMatchFailure( + op, "reduce_addf requires every source chunk to match result " + "vreg type"); + for (Value maskPart : maskParts) + if (maskPart.getType() != maskType) + return rewriter.notifyMatchFailure( + op, "reduce_addf requires every mask chunk to have the same " + "predicate type"); + + FailureOr firstLaneMask = + createPrefixMask(op.getLoc(), maskType, "PAT_VL1", rewriter); + if (failed(firstLaneMask)) + return rewriter.notifyMatchFailure( + op, "failed to create reduce_addf first-lane mask"); + + Value accumulator = initParts.front(); + for (auto [sourcePart, maskPart] : + llvm::zip_equal(sourceParts, maskParts)) { + Value reduced = + rewriter + .create(op.getLoc(), resultType, sourcePart, maskPart) + .getResult(); + accumulator = rewriter + .create(op.getLoc(), resultType, reduced, + accumulator, *firstLaneMask) + .getResult(); + } + + replaceOpWithFlatConvertedValues( + rewriter, op, SmallVector{accumulator}, + *this->getTypeConverter()); + return success(); + } +}; + +enum class GroupReduceLoweringPlan { + OneBlockVcgadd, + TwoBlockDeinterleaved2VcgaddVadd, + FourBlockDeinterleaved4VcgaddTree, + FullDeinterleaved2VcaddRows, + ContiguousVcaddRows, +}; + +FailureOr +classifyGroupReduceLoweringPlan(VMIVRegType sourceType, VMIMaskType maskType, + VMIVRegType resultType, int64_t numGroups, + std::string *reason = nullptr) { + VMILayoutSupport supports; + FailureOr fact = + supports.getGroupReduceLayoutFactForLayouts( + sourceType, maskType, resultType, numGroups, reason); + if (failed(fact)) + return failure(); + + switch (fact->blockClass) { + case VMIGroupBlockClass::QuarterBlock: + case VMIGroupBlockClass::HalfBlock: + case VMIGroupBlockClass::OneBlock: + return GroupReduceLoweringPlan::OneBlockVcgadd; + case VMIGroupBlockClass::TwoBlock: + return GroupReduceLoweringPlan::TwoBlockDeinterleaved2VcgaddVadd; + case VMIGroupBlockClass::FourBlock: + return GroupReduceLoweringPlan::FourBlockDeinterleaved4VcgaddTree; + case VMIGroupBlockClass::FullPartMultiple: + if (fact->sourceLayout && fact->sourceLayout.isDeinterleaved() && + fact->sourceLayout.getFactor() == 2 && + fact->sourceLayout.getBlockElems() == 1) + return GroupReduceLoweringPlan::FullDeinterleaved2VcaddRows; + return GroupReduceLoweringPlan::ContiguousVcaddRows; + } + llvm_unreachable("unknown group block class"); +} + +template +struct OneToNVMIGroupReduceOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(OpTy op, + typename OpConversionPattern::OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto sourceVMIType = cast(op.getSource().getType()); + auto resultVMIType = cast(op.getResult().getType()); + ValueRange sourceParts = adaptor.getSource(); + ValueRange maskParts = adaptor.getMask(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + + VMILayoutSupport supports; + std::string supportReason; + if (failed(getSupport(supports, op, &supportReason))) + return rewriter.notifyMatchFailure( + op, Twine(op->getName().getStringRef()) + + " has no layout support: " + supportReason); + auto maskVMIType = cast(op.getMask().getType()); + FailureOr plan = classifyGroupReduceLoweringPlan( + sourceVMIType, maskVMIType, resultVMIType, + op.getNumGroupsAttr().getInt(), &supportReason); + if (failed(plan)) + return rewriter.notifyMatchFailure( + op, Twine(op->getName().getStringRef()) + + " has no lowering plan: " + supportReason); + + FailureOr groupSize = getGroupSizeFromNumGroups( + sourceVMIType, op.getNumGroupsAttr().getInt()); + if (failed(groupSize)) + return rewriter.notifyMatchFailure( + op, "group reduce requires num_groups to evenly divide lane count"); + + if (*plan == GroupReduceLoweringPlan::OneBlockVcgadd) { + if (sourceParts.size() != maskParts.size() || + sourceParts.size() != resultTypes.size() || sourceParts.empty()) + return rewriter.notifyMatchFailure( + op, "vcg group_reduce path requires matching physical " + "arity"); + auto resultType = dyn_cast(resultTypes.front()); + auto maskType = dyn_cast(maskParts.front().getType()); + if (!resultType || !maskType) + return rewriter.notifyMatchFailure( + op, "vcg group_reduce path requires physical vreg/mask"); + for (auto [sourcePart, maskPart, physicalResultType] : + llvm::zip_equal(sourceParts, maskParts, resultTypes)) { + if (sourcePart.getType() != resultType || + maskPart.getType() != maskType || physicalResultType != resultType) + return rewriter.notifyMatchFailure( + op, "vcg group_reduce path requires uniform physical " + "chunk types"); + } + + SmallVector results; + results.reserve(resultTypes.size()); + for (auto [sourceIndex, sourcePart] : llvm::enumerate(sourceParts)) { + results.push_back(rewriter + .create(op.getLoc(), resultType, + sourcePart, + maskParts[sourceIndex]) + .getResult()); + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } + + if (*plan == GroupReduceLoweringPlan::TwoBlockDeinterleaved2VcgaddVadd) { + int64_t resultPartCount = resultTypes.size(); + if (static_cast(sourceParts.size()) != resultPartCount * 2 || + maskParts.size() != sourceParts.size()) + return rewriter.notifyMatchFailure( + op, "s16 block8 group_reduce arity mismatch"); + + SmallVector results; + results.reserve(resultPartCount); + auto resultType = dyn_cast(resultTypes.front()); + auto maskType = dyn_cast(maskParts.front().getType()); + if (!resultType || !maskType) + return rewriter.notifyMatchFailure( + op, "s16 block8 group_reduce requires physical vreg/mask"); + int64_t numGroups = op.getNumGroupsAttr().getInt(); + + for (int64_t resultIndex = 0; resultIndex < resultPartCount; + ++resultIndex) { + int64_t activeGroups = + std::min(8, numGroups - resultIndex * 8); + FailureOr combineMask = createPrefixMaskForActiveLanes( + op.getLoc(), maskType, activeGroups, rewriter); + if (failed(combineMask)) + return rewriter.notifyMatchFailure( + op, "failed to create s16 block8 combine mask"); + Value loSource = sourceParts[resultIndex]; + Value hiSource = sourceParts[resultPartCount + resultIndex]; + Value loMask = maskParts[resultIndex]; + Value hiMask = maskParts[resultPartCount + resultIndex]; + Type physicalResultType = resultTypes[resultIndex]; + if (physicalResultType != resultType || + loSource.getType() != resultType || + hiSource.getType() != resultType || loMask.getType() != maskType || + hiMask.getType() != maskType) + return rewriter.notifyMatchFailure( + op, "s16 block8 group_reduce requires uniform physical " + "types"); + Value lo = rewriter + .create(op.getLoc(), resultType, + loSource, loMask) + .getResult(); + Value hi = rewriter + .create(op.getLoc(), resultType, + hiSource, hiMask) + .getResult(); + results.push_back(rewriter + .create(op.getLoc(), resultType, lo, + hi, *combineMask) + .getResult()); + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } + + if (*plan == GroupReduceLoweringPlan::FourBlockDeinterleaved4VcgaddTree) { + int64_t resultPartCount = resultTypes.size(); + if (static_cast(sourceParts.size()) != resultPartCount * 4 || + maskParts.size() != sourceParts.size()) + return rewriter.notifyMatchFailure( + op, "s32 block8 group_reduce arity mismatch"); + + SmallVector results; + results.reserve(resultPartCount); + auto resultType = dyn_cast(resultTypes.front()); + auto maskType = dyn_cast(maskParts.front().getType()); + if (!resultType || !maskType) + return rewriter.notifyMatchFailure( + op, "s32 block8 group_reduce requires physical vreg/mask"); + int64_t numGroups = op.getNumGroupsAttr().getInt(); + + for (int64_t resultIndex = 0; resultIndex < resultPartCount; + ++resultIndex) { + int64_t activeGroups = + std::min(8, numGroups - resultIndex * 8); + FailureOr combineMask = createPrefixMaskForActiveLanes( + op.getLoc(), maskType, activeGroups, rewriter); + if (failed(combineMask)) + return rewriter.notifyMatchFailure( + op, "failed to create s32 block8 combine mask"); + SmallVector partials; + partials.reserve(4); + for (int64_t part = 0; part < 4; ++part) { + int64_t sourceIndex = part * resultPartCount + resultIndex; + Value source = sourceParts[sourceIndex]; + Value mask = maskParts[sourceIndex]; + Type physicalResultType = resultTypes[resultIndex]; + if (physicalResultType != resultType || + source.getType() != resultType || mask.getType() != maskType) + return rewriter.notifyMatchFailure( + op, "s32 block8 group_reduce requires uniform physical " + "types"); + partials.push_back(rewriter + .create( + op.getLoc(), resultType, source, mask) + .getResult()); + } + Value sum01 = + rewriter + .create(op.getLoc(), resultType, partials[0], + partials[1], *combineMask) + .getResult(); + Value sum23 = + rewriter + .create(op.getLoc(), resultType, partials[2], + partials[3], *combineMask) + .getResult(); + results.push_back(rewriter + .create(op.getLoc(), resultType, + sum01, sum23, *combineMask) + .getResult()); + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } + + if (*plan == GroupReduceLoweringPlan::FullDeinterleaved2VcaddRows) { + VMILayoutAttr resultLayout = resultVMIType.getLayoutAttr(); + bool rowLocalSlots1Result = resultLayout && resultLayout.isGroupSlots() && + resultLayout.getSlots() == 1; + if (!rowLocalSlots1Result) + return rewriter.notifyMatchFailure( + op, "deinterleaved=2 full group_reduce requires slots=1 result"); + + FailureOr lanesPerPart = + getDataLanesPerPart(sourceVMIType.getElementType()); + if (failed(lanesPerPart)) + return rewriter.notifyMatchFailure( + op, "deinterleaved=2 group_reduce requires known physical lanes"); + if (*groupSize % (2 * *lanesPerPart) != 0) + return rewriter.notifyMatchFailure( + op, "deinterleaved=2 group_reduce requires group size to be a " + "multiple of two physical chunks"); + int64_t groupCount = sourceVMIType.getElementCount() / *groupSize; + int64_t chunksPerGroupPerPart = *groupSize / (2 * *lanesPerPart); + int64_t chunksPerPart = groupCount * chunksPerGroupPerPart; + if (sourceParts.size() != maskParts.size() || + static_cast(sourceParts.size()) != 2 * chunksPerPart || + static_cast(resultTypes.size()) != groupCount) + return rewriter.notifyMatchFailure( + op, "deinterleaved=2 group_reduce arity mismatch"); + + SmallVector results(resultTypes.size()); + for (Type resultType : resultTypes) { + if (!isa(resultType)) + return rewriter.notifyMatchFailure( + op, "deinterleaved=2 group_reduce result must be vreg"); + } + + auto resultType = dyn_cast(resultTypes.front()); + auto maskType = dyn_cast(maskParts.front().getType()); + if (!resultType || !maskType) + return rewriter.notifyMatchFailure( + op, "deinterleaved=2 group_reduce requires physical vreg/mask"); + auto sourcePartType = dyn_cast(sourceParts.front().getType()); + if (!sourcePartType) + return rewriter.notifyMatchFailure( + op, "deinterleaved=2 group_reduce source must be vreg"); + FailureOr rowResultType = + getRowResultType(sourcePartType, resultType); + if (failed(rowResultType)) + return rewriter.notifyMatchFailure( + op, "failed to derive deinterleaved=2 row-reduction type"); + FailureOr rowMaskType = + getMaskTypeForVReg(*rowResultType, rewriter.getContext()); + if (failed(rowMaskType)) + return rewriter.notifyMatchFailure( + op, "failed to derive deinterleaved=2 combine mask type"); + FailureOr firstLaneMask = createPrefixMask( + op.getLoc(), *rowMaskType, "PAT_VL1", rewriter); + if (failed(firstLaneMask)) + return rewriter.notifyMatchFailure( + op, "failed to create deinterleaved=2 group_reduce lane mask"); + + for (int64_t group = 0; group < groupCount; ++group) { + Value accumulator; + for (int64_t chunk = 0; chunk < chunksPerGroupPerPart; ++chunk) { + int64_t loIndex = group * chunksPerGroupPerPart + chunk; + int64_t hiIndex = chunksPerPart + loIndex; + if (sourceParts[loIndex].getType() != sourcePartType || + sourceParts[hiIndex].getType() != sourcePartType || + maskParts[loIndex].getType() != maskType || + maskParts[hiIndex].getType() != maskType) + return rewriter.notifyMatchFailure( + op, "deinterleaved=2 group_reduce requires uniform physical " + "chunk types"); + + Value loReduced = + rewriter + .create(op.getLoc(), *rowResultType, + sourceParts[loIndex], + maskParts[loIndex]) + .getResult(); + Value hiReduced = + rewriter + .create(op.getLoc(), *rowResultType, + sourceParts[hiIndex], + maskParts[hiIndex]) + .getResult(); + Value pairReduced = + rewriter + .create(op.getLoc(), *rowResultType, loReduced, + hiReduced, *firstLaneMask) + .getResult(); + if (!accumulator) { + accumulator = pairReduced; + continue; + } + accumulator = + rewriter + .create(op.getLoc(), *rowResultType, pairReduced, + accumulator, *firstLaneMask) + .getResult(); + } + FailureOr finalResult = + bitcastVReg(op.getLoc(), accumulator, resultType, rewriter); + if (failed(finalResult)) + return rewriter.notifyMatchFailure( + op, "failed to restore deinterleaved=2 group result type"); + results[group] = *finalResult; + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, + *this->getTypeConverter()); + return success(); + } + + if (*plan != GroupReduceLoweringPlan::ContiguousVcaddRows) + return rewriter.notifyMatchFailure(op, + "unknown group_reduce lowering plan"); + + int64_t lanesPerPart = 0; + int64_t groupCount = 0; + int64_t chunksPerGroup = 0; + if (failed(checkContiguousFullGroupChunks(op, sourceVMIType, *groupSize, + &lanesPerPart, &groupCount, + &chunksPerGroup, rewriter))) + return failure(); + VMILayoutAttr resultLayout = resultVMIType.getLayoutAttr(); + bool rowLocalSlots1Result = resultLayout && resultLayout.isGroupSlots() && + resultLayout.getNumGroups() == groupCount && + resultLayout.getSlots() == 1; + int64_t expectedResultParts = + rowLocalSlots1Result ? groupCount : groupCount * chunksPerGroup; + if (sourceParts.size() != maskParts.size() || + static_cast(sourceParts.size()) != + groupCount * chunksPerGroup || + static_cast(resultTypes.size()) != expectedResultParts) + return rewriter.notifyMatchFailure( + op, "group_reduce requires matching source/mask/result arity"); + + SmallVector results(resultTypes.size()); + for (Type resultType : resultTypes) { + if (!isa(resultType)) + return rewriter.notifyMatchFailure( + op, "group_reduce result must be vreg"); + } + + auto resultType = dyn_cast(resultTypes.front()); + auto maskType = dyn_cast(maskParts.front().getType()); + if (!resultType || !maskType) + return rewriter.notifyMatchFailure( + op, "group_reduce requires physical vreg result and mask"); + + auto sourcePartType = dyn_cast(sourceParts.front().getType()); + if (!sourcePartType) + return rewriter.notifyMatchFailure(op, + "group_reduce source must be vreg"); + FailureOr rowResultType = + getRowResultType(sourcePartType, resultType); + if (failed(rowResultType)) + return rewriter.notifyMatchFailure( + op, "failed to derive group row-reduction type"); + FailureOr rowMaskType = + getMaskTypeForVReg(*rowResultType, rewriter.getContext()); + if (failed(rowMaskType)) + return rewriter.notifyMatchFailure( + op, "failed to derive group combine mask type"); + FailureOr firstLaneMask = createPrefixMask( + op.getLoc(), *rowMaskType, "PAT_VL1", rewriter); + if (failed(firstLaneMask)) + return rewriter.notifyMatchFailure( + op, "failed to create group_reduce masks"); + + for (int64_t group = 0; group < groupCount; ++group) { + Value accumulator; + + for (int64_t chunk = 0; chunk < chunksPerGroup; ++chunk) { + int64_t index = group * chunksPerGroup + chunk; + if (sourceParts[index].getType() != sourcePartType || + maskParts[index].getType() != maskType) + return rewriter.notifyMatchFailure( + op, "group_reduce requires uniform physical chunk types"); + Value reduced = + rewriter + .create(op.getLoc(), *rowResultType, + sourceParts[index], maskParts[index]) + .getResult(); + if (!accumulator) { + accumulator = reduced; + continue; + } + accumulator = rewriter + .create(op.getLoc(), *rowResultType, + reduced, accumulator, + *firstLaneMask) + .getResult(); + } + + FailureOr finalResult = + bitcastVReg(op.getLoc(), accumulator, resultType, rewriter); + if (failed(finalResult)) + return rewriter.notifyMatchFailure( + op, "failed to restore group result type"); + + int64_t destChunk = rowLocalSlots1Result ? group : group * chunksPerGroup; + if (rowLocalSlots1Result) { + results[destChunk] = *finalResult; + } else { + for (int64_t chunk = 0; chunk < chunksPerGroup; ++chunk) + results[destChunk + chunk] = *finalResult; + } + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } + +private: + FailureOr getRowResultType(VRegType sourceType, + VRegType resultType) const { + if constexpr (std::is_same_v) + return getVcaddResultType(sourceType); + return resultType; + } + + LogicalResult getSupport(VMILayoutSupport &supports, VMIGroupReduceAddFOp op, + std::string *reason) const { + return supports.getGroupReduceAddFSupport(op, reason); + } + + LogicalResult getSupport(VMILayoutSupport &supports, VMIGroupReduceAddIOp op, + std::string *reason) const { + return supports.getGroupReduceAddISupport(op, reason); + } + + LogicalResult getSupport(VMILayoutSupport &supports, VMIGroupReduceMaxIOp op, + std::string *reason) const { + return supports.getGroupReduceMaxISupport(op, reason); + } + + LogicalResult getSupport(VMILayoutSupport &supports, VMIGroupReduceMaxFOp op, + std::string *reason) const { + return supports.getGroupReduceMaxFSupport(op, reason); + } + + LogicalResult getSupport(VMILayoutSupport &supports, VMIGroupReduceMinFOp op, + std::string *reason) const { + return supports.getGroupReduceMinFSupport(op, reason); + } + + LogicalResult getSupport(VMILayoutSupport &supports, VMIGroupReduceMinIOp op, + std::string *reason) const { + return supports.getGroupReduceMinISupport(op, reason); + } + + ; +}; + +struct OneToNVMIGroupBroadcastOpPattern + : OpConversionPattern { + using OpConversionPattern< + VMIGroupBroadcastOp>::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIGroupBroadcastOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto sourceVMIType = cast(op.getSource().getType()); + auto resultVMIType = cast(op.getResult().getType()); + FailureOr groupSize = getGroupSizeFromNumGroups( + resultVMIType, op.getNumGroupsAttr().getInt()); + if (failed(groupSize)) + return rewriter.notifyMatchFailure( + op, + "group_broadcast requires num_groups to evenly divide lane count"); + int64_t lanesPerPart = 0; + int64_t groupCount = 0; + if (failed(checkFullGroupSlotSourceShape( + op, sourceVMIType, *groupSize, op.getNumGroupsAttr().getInt(), + &lanesPerPart, &groupCount, rewriter))) + return failure(); + int64_t resultLayoutFactor = 0; + int64_t resultGroupCount = 0; + if (failed(checkFullGroupBroadcastResultShape( + op, resultVMIType, *groupSize, lanesPerPart, &resultLayoutFactor, + &resultGroupCount, rewriter))) + return failure(); + if (resultGroupCount != groupCount) + return rewriter.notifyMatchFailure( + op, "group_broadcast requires matching source/result group slots"); + + ValueRange sourceParts = adaptor.getSource(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + if (sourceParts.empty() || resultTypes.empty()) + return rewriter.notifyMatchFailure(op, "group_broadcast arity mismatch"); + + auto firstSourceType = dyn_cast(sourceParts.front().getType()); + if (!firstSourceType) + return rewriter.notifyMatchFailure(op, + "group_broadcast source must be vreg"); + unsigned indexBits = + pto::getPTOStorageElemBitWidth(firstSourceType.getElementType()); + if (indexBits != 8 && indexBits != 16 && indexBits != 32) + return rewriter.notifyMatchFailure( + op, "group_broadcast requires 8/16/32-bit index elements"); + auto indexElementType = IntegerType::get(rewriter.getContext(), indexBits); + auto indexType = + VRegType::get(rewriter.getContext(), firstSourceType.getElementCount(), + indexElementType); + FailureOr allMask = + createAllTrueMaskForVReg(op.getLoc(), firstSourceType, rewriter); + if (failed(allMask)) + return rewriter.notifyMatchFailure( + op, "failed to create group_broadcast all mask"); + VMILayoutAttr resultLayout = resultVMIType.getLayoutAttr(); + VMILayoutAttr sourceLayout = sourceVMIType.getLayoutAttr(); + int64_t selectionGroupSize = *groupSize; + if (resultLayoutFactor != 1 && resultLayout && + resultLayout.isDeinterleaved() && resultLayout.getBlockElems() > 1 && + *groupSize < lanesPerPart) + selectionGroupSize = resultLayout.getBlockElems(); + auto resolveLargeGroupSource = [&](int64_t group, int64_t chunksPerGroup, + int64_t &sourceChunk, + int64_t &baseGroupSlot) { + int64_t slots = sourceLayout.getSlots(); + if (slots > 0) { + sourceChunk = group / slots; + baseGroupSlot = group % slots; + return; + } + sourceChunk = group * chunksPerGroup; + baseGroupSlot = 0; + }; + + SmallVector results; + results.resize(resultTypes.size()); + for (auto [flatIndex, resultType] : llvm::enumerate(resultTypes)) { + auto resultVRegType = dyn_cast(resultType); + if (!resultVRegType || resultVRegType != firstSourceType) + return rewriter.notifyMatchFailure( + op, "group_broadcast requires uniform physical vreg types"); + int64_t sourceChunk = flatIndex; + int64_t baseGroupSlot = 0; + Value mappedGroupSlotIndex; + if (resultLayoutFactor == 1) { + bool laneStridedDense = + resultLayout && resultLayout.isDense() && + resultLayout.getLaneStride() > 1; + if (laneStridedDense) { + VMILayoutAttr sourceLayout = sourceVMIType.getLayoutAttr(); + int64_t slots = sourceLayout.getSlots(); + if (slots <= 0) { + if (sourceParts.empty() || + groupCount % static_cast(sourceParts.size()) != 0) + return rewriter.notifyMatchFailure( + op, "group_broadcast lane-stride source requires explicit " + "group_slots slots or derivable legacy slot count"); + slots = groupCount / sourceParts.size(); + } + FailureOr index = createMappedGroupSlotIndexVector( + op.getLoc(), resultVMIType, /*part=*/0, flatIndex, indexType, + *groupSize, slots, sourceChunk, rewriter, + sourceLayout.getLaneStride()); + if (failed(index)) + return rewriter.notifyMatchFailure( + op, "failed to create group_broadcast lane-stride group-slot " + "index vector"); + mappedGroupSlotIndex = *index; + } else if (*groupSize >= lanesPerPart) { + int64_t chunksPerGroup = *groupSize / lanesPerPart; + int64_t group = flatIndex / chunksPerGroup; + resolveLargeGroupSource(group, chunksPerGroup, sourceChunk, + baseGroupSlot); + } else { + VMILayoutAttr sourceLayout = sourceVMIType.getLayoutAttr(); + int64_t slots = sourceLayout.getSlots(); + if (slots <= 0) { + if (sourceParts.empty() || + groupCount % static_cast(sourceParts.size()) != 0) + return rewriter.notifyMatchFailure( + op, "group_broadcast small-group source requires explicit " + "group_slots slots or derivable legacy slot count"); + slots = groupCount / sourceParts.size(); + } + int64_t groupsPerResultChunk = lanesPerPart / *groupSize; + int64_t firstGroup = flatIndex * groupsPerResultChunk; + sourceChunk = firstGroup / slots; + baseGroupSlot = firstGroup % slots; + } + } else { + bool blockFragmentSmallGroup = + resultLayout && resultLayout.isDeinterleaved() && + resultLayout.getBlockElems() > 1 && *groupSize < lanesPerPart; + bool deinterleavedSmallGroup = + resultLayout && resultLayout.isDeinterleaved() && + resultLayout.getBlockElems() == 1 && *groupSize < lanesPerPart; + if (blockFragmentSmallGroup) { + int64_t runningFlatIndex = 0; + bool found = false; + for (int64_t part = 0; part < resultLayoutFactor && !found; ++part) { + FailureOr chunks = + getDataChunksInPart(resultVMIType, part); + if (failed(chunks)) + return rewriter.notifyMatchFailure( + op, "group_broadcast failed to enumerate result chunks"); + for (int64_t chunk = 0; chunk < *chunks; + ++chunk, ++runningFlatIndex) { + if (runningFlatIndex != static_cast(flatIndex)) + continue; + int64_t groupsPerResultChunk = + lanesPerPart / resultLayout.getBlockElems(); + int64_t firstGroup = chunk * groupsPerResultChunk; + int64_t slots = sourceLayout.getSlots(); + if (slots <= 0) { + if (sourceParts.empty() || + groupCount % static_cast(sourceParts.size()) != 0) + return rewriter.notifyMatchFailure( + op, + "group_broadcast block-fragment source requires explicit " + "group_slots slots or derivable legacy slot count"); + slots = groupCount / sourceParts.size(); + } + sourceChunk = firstGroup / slots; + baseGroupSlot = firstGroup % slots; + found = true; + break; + } + } + if (!found) + return rewriter.notifyMatchFailure( + op, "group_broadcast result chunk index is out of range"); + } else if (deinterleavedSmallGroup) { + int64_t runningFlatIndex = 0; + bool found = false; + for (int64_t part = 0; part < resultLayoutFactor && !found; ++part) { + FailureOr chunks = + getDataChunksInPart(resultVMIType, part); + if (failed(chunks)) + return rewriter.notifyMatchFailure( + op, "group_broadcast failed to enumerate result chunks"); + for (int64_t chunk = 0; chunk < *chunks; + ++chunk, ++runningFlatIndex) { + if (runningFlatIndex != static_cast(flatIndex)) + continue; + int64_t slots = sourceLayout.getSlots(); + if (slots <= 0) { + if (sourceParts.empty() || + groupCount % static_cast(sourceParts.size()) != 0) + return rewriter.notifyMatchFailure( + op, "group_broadcast deinterleaved small-group source " + "requires explicit group_slots slots or derivable " + "legacy slot count"); + slots = groupCount / sourceParts.size(); + } + FailureOr index = createMappedGroupSlotIndexVector( + op.getLoc(), resultVMIType, part, chunk, indexType, + *groupSize, slots, sourceChunk, rewriter, + sourceLayout.getLaneStride()); + if (failed(index)) + return rewriter.notifyMatchFailure( + op, + "failed to create group_broadcast mapped group-slot index " + "vector"); + mappedGroupSlotIndex = *index; + found = true; + break; + } + } + if (!found) + return rewriter.notifyMatchFailure( + op, "group_broadcast result chunk index is out of range"); + } else { + int64_t runningFlatIndex = 0; + bool found = false; + for (int64_t part = 0; part < resultLayoutFactor && !found; ++part) { + FailureOr chunks = + getDataChunksInPart(resultVMIType, part); + if (failed(chunks)) + return rewriter.notifyMatchFailure( + op, "group_broadcast failed to enumerate result chunks"); + for (int64_t chunk = 0; chunk < *chunks; + ++chunk, ++runningFlatIndex) { + if (runningFlatIndex != static_cast(flatIndex)) + continue; + FailureOr firstLogical = + mapPhysicalLaneToLogical(resultVMIType, part, chunk, 0); + FailureOr lastLogical = mapPhysicalLaneToLogical( + resultVMIType, part, chunk, lanesPerPart - 1); + if (failed(firstLogical) || failed(lastLogical)) + return rewriter.notifyMatchFailure( + op, "group_broadcast failed to map result chunk lanes"); + int64_t firstGroup = *firstLogical / *groupSize; + int64_t lastGroup = *lastLogical / *groupSize; + if (firstGroup != lastGroup) + return rewriter.notifyMatchFailure( + op, "group_broadcast result chunk crosses logical groups"); + int64_t chunksPerGroup = *groupSize / lanesPerPart; + resolveLargeGroupSource(firstGroup, chunksPerGroup, sourceChunk, + baseGroupSlot); + found = true; + break; + } + } + if (!found) + return rewriter.notifyMatchFailure( + op, "group_broadcast result chunk index is out of range"); + } + } + if (*groupSize >= lanesPerPart) { + if (sourceChunk < 0 || + sourceChunk >= static_cast(sourceParts.size())) + return rewriter.notifyMatchFailure( + op, "group_broadcast source chunk is out of range"); + if (sourceLayout.getSlots() > 1) { + FailureOr groupSlotIndex = createGroupSlotIndexVector( + op.getLoc(), indexType, selectionGroupSize, baseGroupSlot, + rewriter, sourceLayout.getLaneStride()); + if (failed(groupSlotIndex)) + return rewriter.notifyMatchFailure( + op, "failed to create group_broadcast group-slot index vector"); + results[flatIndex] = + rewriter + .create(op.getLoc(), resultType, + sourceParts[sourceChunk], *groupSlotIndex) + .getResult(); + } else { + results[flatIndex] = + rewriter + .create(op.getLoc(), resultType, + sourceParts[sourceChunk], *allMask, + rewriter.getStringAttr("LOWEST")) + .getResult(); + } + } else { + bool blockFragmentSmallGroup = resultLayout && + resultLayout.isDeinterleaved() && + resultLayout.getBlockElems() > 1; + bool deinterleavedSmallGroup = resultLayout && + resultLayout.isDeinterleaved() && + resultLayout.getBlockElems() == 1; + if (resultLayoutFactor != 1 && !blockFragmentSmallGroup && + !deinterleavedSmallGroup) + return rewriter.notifyMatchFailure( + op, "group_broadcast small-group deinterleaved result is not " + "supported"); + if (sourceChunk < 0 || + sourceChunk >= static_cast(sourceParts.size())) + return rewriter.notifyMatchFailure( + op, "group_broadcast source chunk is out of range"); + FailureOr groupSlotIndex = + mappedGroupSlotIndex + ? FailureOr(mappedGroupSlotIndex) + : createGroupSlotIndexVector(op.getLoc(), indexType, + selectionGroupSize, baseGroupSlot, + rewriter, + sourceLayout.getLaneStride()); + if (failed(groupSlotIndex)) + return rewriter.notifyMatchFailure( + op, "failed to create group_broadcast group-slot index vector"); + results[flatIndex] = + rewriter + .create(op.getLoc(), resultType, + sourceParts[sourceChunk], *groupSlotIndex) + .getResult(); + } + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMIVdhistOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIVdhistOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + ValueRange accParts = adaptor.getAcc(); + ValueRange sourceParts = adaptor.getSource(); + ValueRange maskParts = adaptor.getMask(); + if (accParts.size() != 2 || sourceParts.empty() || + sourceParts.size() != maskParts.size()) + return rewriter.notifyMatchFailure( + op, "expected two accumulator parts and matching source/mask chunks"); + + auto loType = dyn_cast(accParts[0].getType()); + auto hiType = dyn_cast(accParts[1].getType()); + if (!loType || loType != hiType) + return rewriter.notifyMatchFailure(op, + "expected matching ui16 acc parts"); + auto sourceType = cast(op.getSource().getType()); + FailureOr lanesPerPart = + getDataLanesPerPart(sourceType.getElementType()); + if (failed(lanesPerPart)) + return rewriter.notifyMatchFailure(op, "failed to compute source lanes"); + + Location loc = op.getLoc(); + Value bin0 = createI32Constant(loc, 0, rewriter); + Value bin1 = createI32Constant(loc, 1, rewriter); + Value lo = accParts[0]; + Value hi = accParts[1]; + + for (size_t index = 0, e = sourceParts.size(); index < e; ++index) { + Value source = sourceParts[index]; + Value userMask = maskParts[index]; + auto maskType = dyn_cast(userMask.getType()); + if (!maskType || !maskType.isB8()) + return rewriter.notifyMatchFailure(op, "expected b8 source mask"); + + Value chunkMask = userMask; + int64_t firstLane = static_cast(index) * *lanesPerPart; + int64_t activeLanes = std::min( + *lanesPerPart, sourceType.getElementCount() - firstLane); + if (activeLanes < *lanesPerPart) { + FailureOr validMask = createPrefixMaskForActiveLanes( + loc, maskType, activeLanes, rewriter); + FailureOr allMask = createAllTrueMask(loc, maskType, rewriter); + if (failed(validMask) || failed(allMask)) + return rewriter.notifyMatchFailure( + op, "failed to materialize tail-valid b8 mask"); + chunkMask = + rewriter + .create(loc, maskType, chunkMask, *validMask, *allMask) + .getResult(); + } + + lo = rewriter.create(loc, loType, lo, source, chunkMask, bin0) + .getResult(); + hi = rewriter.create(loc, hiType, hi, source, chunkMask, bin1) + .getResult(); + } + + replaceOpWithFlatConvertedValues(rewriter, op, SmallVector{lo, hi}, + *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMIVchistOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIVchistOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + ValueRange accParts = adaptor.getAcc(); + ValueRange sourceParts = adaptor.getSource(); + ValueRange maskParts = adaptor.getMask(); + if (accParts.size() != 2 || sourceParts.empty() || + sourceParts.size() != maskParts.size()) + return rewriter.notifyMatchFailure( + op, "expected two accumulator parts and matching source/mask chunks"); + + auto loType = dyn_cast(accParts[0].getType()); + auto hiType = dyn_cast(accParts[1].getType()); + if (!loType || loType != hiType) + return rewriter.notifyMatchFailure(op, + "expected matching ui16 acc parts"); + auto sourceType = cast(op.getSource().getType()); + FailureOr lanesPerPart = + getDataLanesPerPart(sourceType.getElementType()); + if (failed(lanesPerPart)) + return rewriter.notifyMatchFailure(op, "failed to compute source lanes"); + + Location loc = op.getLoc(); + Value bin0 = createI32Constant(loc, 0, rewriter); + Value bin1 = createI32Constant(loc, 1, rewriter); + Value lo = accParts[0]; + Value hi = accParts[1]; + + for (size_t index = 0, e = sourceParts.size(); index < e; ++index) { + Value source = sourceParts[index]; + Value userMask = maskParts[index]; + auto maskType = dyn_cast(userMask.getType()); + if (!maskType || !maskType.isB8()) + return rewriter.notifyMatchFailure(op, "expected b8 source mask"); + + Value chunkMask = userMask; + int64_t firstLane = static_cast(index) * *lanesPerPart; + int64_t activeLanes = std::min( + *lanesPerPart, sourceType.getElementCount() - firstLane); + if (activeLanes < *lanesPerPart) { + FailureOr validMask = createPrefixMaskForActiveLanes( + loc, maskType, activeLanes, rewriter); + FailureOr allMask = createAllTrueMask(loc, maskType, rewriter); + if (failed(validMask) || failed(allMask)) + return rewriter.notifyMatchFailure( + op, "failed to materialize tail-valid b8 mask"); + chunkMask = + rewriter + .create(loc, maskType, chunkMask, *validMask, *allMask) + .getResult(); + } + + lo = rewriter.create(loc, loType, lo, source, chunkMask, bin0) + .getResult(); + hi = rewriter.create(loc, hiType, hi, source, chunkMask, bin1) + .getResult(); + } + + replaceOpWithFlatConvertedValues(rewriter, op, SmallVector{lo, hi}, + *this->getTypeConverter()); + return success(); + } +}; + +template +struct OneToNVMIReduceMinMaxOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite( + SourceOp op, + typename OpConversionPattern::OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + ValueRange sourceParts = adaptor.getSource(); + ValueRange initParts = adaptor.getInit(); + ValueRange maskParts = adaptor.getMask(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + if (sourceParts.empty() || sourceParts.size() != maskParts.size() || + initParts.size() != 1 || resultTypes.size() != 1) + return rewriter.notifyMatchFailure( + op, "min/max reduction requires matching source/mask chunks " + "and one init/result chunk"); + + auto resultType = dyn_cast(resultTypes.front()); + auto maskType = dyn_cast(maskParts.front().getType()); + if (!resultType || !maskType || initParts.front().getType() != resultType) + return rewriter.notifyMatchFailure( + op, "min/max reduction requires matching physical source/" + "init/result vregs and one mask"); + + for (Value sourcePart : sourceParts) + if (sourcePart.getType() != resultType) + return rewriter.notifyMatchFailure( + op, "min/max reduction requires every source chunk to " + "match result vreg type"); + for (Value maskPart : maskParts) + if (maskPart.getType() != maskType) + return rewriter.notifyMatchFailure( + op, "min/max reduction requires every mask chunk to have " + "the same predicate type"); + + FailureOr firstLaneMask = + createPrefixMask(op.getLoc(), maskType, "PAT_VL1", rewriter); + if (failed(firstLaneMask)) + return rewriter.notifyMatchFailure( + op, "failed to create min/max reduction first-lane mask"); + + Value accumulator = initParts.front(); + for (auto [sourcePart, maskPart] : + llvm::zip_equal(sourceParts, maskParts)) { + Value reduced = rewriter + .create(op.getLoc(), resultType, + sourcePart, maskPart) + .getResult(); + accumulator = rewriter + .create(op.getLoc(), resultType, reduced, + accumulator, *firstLaneMask) + .getResult(); + } + + replaceOpWithFlatConvertedValues( + rewriter, op, SmallVector{accumulator}, + *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMIExtFOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIExtFOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto sourceVMIType = cast(op.getSource().getType()); + auto resultVMIType = cast(op.getResult().getType()); + ValueRange sourceParts = adaptor.getSource(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + if (sourceParts.empty()) + return rewriter.notifyMatchFailure( + op, "extf requires at least one physical source chunk"); + + auto sourceType = dyn_cast(sourceParts.front().getType()); + if (!sourceType) + return rewriter.notifyMatchFailure(op, "expected physical extf source"); + for (Value sourcePart : sourceParts) { + auto currentSourceType = dyn_cast(sourcePart.getType()); + if (!currentSourceType || currentSourceType != sourceType) + return rewriter.notifyMatchFailure( + op, "extf source physical parts must have matching type"); + } + + SmallVector resultVRegTypes; + resultVRegTypes.reserve(resultTypes.size()); + for (Type resultType : resultTypes) { + auto resultVRegType = dyn_cast(resultType); + if (!resultVRegType || + (resultVRegTypes.empty() ? !resultVRegType.getElementType().isF32() + : resultVRegType != resultVRegTypes.front())) + return rewriter.notifyMatchFailure( + op, "unsupported physical extf result type"); + resultVRegTypes.push_back(resultVRegType); + } + + unsigned sourceBits = + pto::getPTOStorageElemBitWidth(sourceType.getElementType()); + VMILayoutAttr sourceLayout = sourceVMIType.getLayoutAttr(); + VMILayoutAttr resultLayout = resultVMIType.getLayoutAttr(); + if (sourceLayout && resultLayout && sourceLayout.isContiguous() && + resultLayout.isContiguous() && resultLayout.getLaneStride() == 1 && + ((sourceBits == 16 && sourceLayout.getLaneStride() == 2) || + (sourceBits == 8 && sourceLayout.getLaneStride() == 4)) && + resultTypes.size() == sourceParts.size()) { + StringRef part = sourceBits == 16 ? StringRef("EVEN") : StringRef("P0"); + FailureOr mask = + createAllTrueMaskForVReg(op.getLoc(), sourceType, rewriter); + if (failed(mask)) + return rewriter.notifyMatchFailure(op, + "failed to build extf seed mask"); + + SmallVector results; + results.reserve(resultTypes.size()); + for (auto [sourcePart, resultType] : + llvm::zip_equal(sourceParts, resultVRegTypes)) { + results.push_back(rewriter + .create(op.getLoc(), resultType, + sourcePart, *mask, + /*rnd=*/nullptr, /*sat=*/nullptr, + rewriter.getStringAttr(part)) + .getResult()); + } + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } + + ArrayRef parts; + int64_t factor = 0; + if (sourceBits == 16 && resultTypes.size() == 2 * sourceParts.size()) { + static constexpr StringRef kEvenOddParts[] = {"EVEN", "ODD"}; + parts = kEvenOddParts; + factor = 2; + } else if (sourceBits == 8 && + resultTypes.size() == 4 * sourceParts.size()) { + static constexpr StringRef kPacked4Parts[] = {"P0", "P1", "P2", "P3"}; + parts = kPacked4Parts; + factor = 4; + } else { + return rewriter.notifyMatchFailure( + op, "unsupported physical extf source/result width relation"); + } + + FailureOr mask = + createAllTrueMaskForVReg(op.getLoc(), sourceType, rewriter); + if (failed(mask)) + return rewriter.notifyMatchFailure(op, "failed to build extf seed mask"); + + SmallVector results; + results.reserve(resultTypes.size()); + for (int64_t partIndex = 0; partIndex < factor; ++partIndex) { + for (auto [chunkIndex, sourcePart] : llvm::enumerate(sourceParts)) { + VRegType resultType = + resultVRegTypes[partIndex * sourceParts.size() + chunkIndex]; + results.push_back( + rewriter + .create(op.getLoc(), resultType, sourcePart, *mask, + /*rnd=*/nullptr, /*sat=*/nullptr, + rewriter.getStringAttr(parts[partIndex])) + .getResult()); + } + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMITruncFOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMITruncFOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto sourceVMIType = cast(op.getSource().getType()); + auto resultVMIType = cast(op.getResult().getType()); + ValueRange sourceParts = adaptor.getSource(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + + VMILayoutAttr sourceLayout = sourceVMIType.getLayoutAttr(); + VMILayoutAttr resultLayout = resultVMIType.getLayoutAttr(); + if (sourceLayout && resultLayout && sourceLayout.isGroupSlots() && + resultLayout.isGroupSlots()) { + unsigned logicalResultBits = + pto::getPTOStorageElemBitWidth(resultVMIType.getElementType()); + if (sourceLayout.getNumGroups() != resultLayout.getNumGroups() || + sourceLayout.getSlots() != resultLayout.getSlots() || + (sourceLayout.getSlots() != 1 && sourceLayout.getSlots() != 8) || + !sourceVMIType.getElementType().isF32() || + (logicalResultBits != 16 && logicalResultBits != 8) || + sourceParts.size() != resultTypes.size()) + return rewriter.notifyMatchFailure( + op, "unsupported group-slot truncf shape"); + + SmallVector results; + results.reserve(resultTypes.size()); + const char *activeSlotPattern = + sourceLayout.getSlots() == 1 ? "PAT_VL1" : "PAT_VL8"; + FailureOr activeSlotMask = createPrefixMask( + op.getLoc(), MaskType::get(rewriter.getContext(), "b32"), + activeSlotPattern, rewriter); + if (failed(activeSlotMask)) + return rewriter.notifyMatchFailure( + op, "failed to build group-slot truncf active slot mask"); + StringAttr sat = rewriter.getStringAttr("SAT"); + for (auto [sourcePart, physicalResultType] : + llvm::zip_equal(sourceParts, resultTypes)) { + auto sourceType = dyn_cast(sourcePart.getType()); + auto resultType = dyn_cast(physicalResultType); + if (!sourceType || !sourceType.getElementType().isF32() || + !resultType) + return rewriter.notifyMatchFailure( + op, "unsupported group-slot truncf physical type"); + unsigned physicalResultBits = + pto::getPTOStorageElemBitWidth(resultType.getElementType()); + StringAttr part; + if (physicalResultBits == 16) { + part = rewriter.getStringAttr("EVEN"); + } else if (physicalResultBits == 8) { + part = rewriter.getStringAttr("P0"); + } else { + return rewriter.notifyMatchFailure( + op, "unsupported group-slot truncf physical type"); + } + StringAttr rnd = rewriter.getStringAttr( + getTruncFRoundMode(op, resultType.getElementType())); + results.push_back(rewriter + .create(op.getLoc(), resultType, + sourcePart, *activeSlotMask, rnd, + sat, part) + .getResult()); + } + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } + + if (resultTypes.empty()) + return rewriter.notifyMatchFailure(op, "truncf requires result chunks"); + + auto sourceType0 = dyn_cast(sourceParts.front().getType()); + if (!sourceType0 || !sourceType0.getElementType().isF32()) + return rewriter.notifyMatchFailure( + op, "unsupported physical truncf source/result type"); + for (Value sourcePart : sourceParts) { + auto sourceType = dyn_cast(sourcePart.getType()); + if (!sourceType || sourceType != sourceType0) + return rewriter.notifyMatchFailure( + op, "truncf source physical parts must have matching f32 type"); + } + + SmallVector resultVRegTypes; + resultVRegTypes.reserve(resultTypes.size()); + for (Type physicalResultType : resultTypes) { + auto resultType = dyn_cast(physicalResultType); + if (!resultType || + (resultVRegTypes.empty() ? pto::getPTOStorageElemBitWidth( + resultType.getElementType()) == 0 + : resultType != resultVRegTypes.front())) + return rewriter.notifyMatchFailure( + op, "unsupported physical truncf result type"); + resultVRegTypes.push_back(resultType); + } + + unsigned resultBits = pto::getPTOStorageElemBitWidth( + resultVRegTypes.front().getElementType()); + if (sourceLayout && resultLayout && sourceLayout.isContiguous() && + sourceLayout.getLaneStride() == 1 && resultLayout.isContiguous() && + resultLayout.getLaneStride() != 1 && + sourceParts.size() == resultTypes.size()) { + StringRef part; + if (resultBits == 16 && resultLayout.getLaneStride() == 2) + part = "EVEN"; + else if (resultBits == 8 && resultLayout.getLaneStride() == 4) + part = "P0"; + else + return rewriter.notifyMatchFailure( + op, "unsupported dense lane_stride truncf result layout"); + + FailureOr sourceMask = + createAllTrueMaskForVReg(op.getLoc(), sourceType0, rewriter); + if (failed(sourceMask)) + return rewriter.notifyMatchFailure(op, "failed to build truncf masks"); + + StringAttr rnd = rewriter.getStringAttr( + getTruncFRoundMode(op, resultVRegTypes.front().getElementType())); + StringAttr sat = rewriter.getStringAttr("SAT"); + StringAttr partAttr = rewriter.getStringAttr(part); + SmallVector results; + results.reserve(resultTypes.size()); + for (auto [sourcePart, resultType] : + llvm::zip_equal(sourceParts, resultVRegTypes)) { + results.push_back(rewriter + .create(op.getLoc(), resultType, + sourcePart, *sourceMask, rnd, sat, + partAttr) + .getResult()); + } + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } + + ArrayRef allParts; + int64_t factor = 0; + if (resultBits == 16) { + static constexpr StringRef kEvenOddParts[] = {"EVEN", "ODD"}; + allParts = kEvenOddParts; + factor = 2; + } else if (resultBits == 8) { + static constexpr StringRef kPacked4Parts[] = {"P0", "P1", "P2", "P3"}; + allParts = kPacked4Parts; + factor = 4; + } else { + return rewriter.notifyMatchFailure( + op, "unsupported physical truncf source/result width relation"); + } + + int64_t resultLaneStride = resultLayout && resultLayout.isContiguous() + ? resultLayout.getLaneStride() + : 1; + if (resultLaneStride <= 0 || factor % resultLaneStride != 0) + return rewriter.notifyMatchFailure( + op, "unsupported physical truncf result lane stride"); + int64_t sourceFactor = factor / resultLaneStride; + if (sourceParts.size() != sourceFactor * resultTypes.size()) + return rewriter.notifyMatchFailure( + op, "unsupported physical truncf source/result arity relation"); + + FailureOr sourceMask = + createAllTrueMaskForVReg(op.getLoc(), sourceType0, rewriter); + if (failed(sourceMask)) + return rewriter.notifyMatchFailure(op, "failed to build truncf masks"); + + StringAttr rnd = rewriter.getStringAttr( + getTruncFRoundMode(op, resultVRegTypes.front().getElementType())); + StringAttr sat = rewriter.getStringAttr("SAT"); + SmallVector results; + results.reserve(resultTypes.size()); + for (auto [chunkIndex, resultType] : llvm::enumerate(resultVRegTypes)) { + FailureOr resultMask = + createAllTrueMaskForVReg(op.getLoc(), resultType, rewriter); + if (failed(resultMask)) + return rewriter.notifyMatchFailure( + op, "failed to build truncf result mask"); + + SmallVector partials; + partials.reserve(sourceFactor); + for (int64_t partIndex = 0; partIndex < sourceFactor; ++partIndex) { + Value sourcePart = + sourceParts[partIndex * resultTypes.size() + chunkIndex]; + partials.push_back( + rewriter + .create(op.getLoc(), resultType, sourcePart, + *sourceMask, rnd, sat, + rewriter.getStringAttr( + allParts[partIndex * resultLaneStride])) + .getResult()); + } + + Value merged = partials.front(); + for (Value partial : llvm::drop_begin(partials)) + merged = rewriter + .create(op.getLoc(), resultType, merged, partial, + *resultMask) + .getResult(); + results.push_back(merged); + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +template +struct OneToNVMIExtIOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(OpT op, + typename OpConversionPattern::OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto sourceVMIType = cast(op.getSource().getType()); + auto resultVMIType = cast(op.getResult().getType()); + ValueRange sourceParts = adaptor.getSource(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + if (sourceParts.empty()) + return rewriter.notifyMatchFailure( + op, "integer extension requires at least one physical source chunk"); + + auto sourceType = dyn_cast(sourceParts.front().getType()); + if (!sourceType) + return rewriter.notifyMatchFailure( + op, "expected physical integer extension source"); + for (Value sourcePart : sourceParts) { + auto currentSourceType = dyn_cast(sourcePart.getType()); + if (!currentSourceType || currentSourceType != sourceType) + return rewriter.notifyMatchFailure( + op, "integer extension source physical parts must have matching " + "type"); + } + + VMILayoutAttr sourceLayout = sourceVMIType.getLayoutAttr(); + VMILayoutAttr resultLayout = resultVMIType.getLayoutAttr(); + if (sourceLayout && resultLayout && sourceLayout.isGroupSlots() && + resultLayout.isGroupSlots()) { + if (sourceLayout.getNumGroups() != resultLayout.getNumGroups() || + sourceLayout.getSlots() != 8 || resultLayout.getSlots() != 8 || + sourceParts.size() != resultTypes.size()) + return rewriter.notifyMatchFailure( + op, "unsupported group-slot integer extension shape"); + + unsigned sourceBits = + pto::getPTOStorageElemBitWidth(sourceVMIType.getElementType()); + unsigned resultBits = + pto::getPTOStorageElemBitWidth(resultVMIType.getElementType()); + if ((sourceBits != 8 && sourceBits != 16) || resultBits != 32) + return rewriter.notifyMatchFailure( + op, "group-slot integer extension requires 8/16-bit source and " + "32-bit result element widths"); + + FailureOr maskType = + getMaskTypeForVReg(sourceType, rewriter.getContext()); + if (failed(maskType)) + return rewriter.notifyMatchFailure( + op, "failed to create group-slot integer extension mask type"); + FailureOr slotMask = createPrefixMaskForActiveLanes( + op.getLoc(), *maskType, sourceLayout.getSlots(), rewriter); + if (failed(slotMask)) + return rewriter.notifyMatchFailure( + op, "failed to build group-slot integer extension mask"); + + SmallVector partNames; + int64_t partFactor = 0; + if (sourceBits == 16) { + partNames.assign({"EVEN", "ODD"}); + partFactor = 2; + } else { + partNames.assign({"P0", "P1", "P2", "P3"}); + partFactor = 4; + } + + SmallVector results; + results.reserve(resultTypes.size()); + for (auto [chunkIndex, sourcePart, resultType] : + llvm::enumerate(sourceParts, resultTypes)) { + auto resultVRegType = dyn_cast(resultType); + if (!resultVRegType || pto::getPTOStorageElemBitWidth( + resultVRegType.getElementType()) != 32) + return rewriter.notifyMatchFailure( + op, "unsupported group-slot integer extension result type"); + + SmallVector convertedParts; + convertedParts.reserve(partNames.size()); + for (StringRef partName : partNames) { + convertedParts.push_back( + rewriter + .create(op.getLoc(), resultVRegType, sourcePart, + *slotMask, /*rnd=*/nullptr, /*sat=*/nullptr, + rewriter.getStringAttr(partName)) + .getResult()); + } + + FailureOr resultMaskType = + getMaskTypeForVReg(resultVRegType, rewriter.getContext()); + FailureOr resultAllMask = + createAllTrueMaskForVReg(op.getLoc(), resultVRegType, rewriter); + if (failed(resultMaskType) || failed(resultAllMask)) + return rewriter.notifyMatchFailure( + op, "failed to build group-slot integer extension result seed"); + + auto indexType = VRegType::get( + rewriter.getContext(), resultVRegType.getElementCount(), + IntegerType::get(rewriter.getContext(), 32)); + int64_t groupBegin = + static_cast(chunkIndex) * sourceLayout.getSlots(); + int64_t activeSlots = std::min( + sourceLayout.getSlots(), sourceLayout.getNumGroups() - groupBegin); + if (activeSlots <= 0) + return rewriter.notifyMatchFailure( + op, "group-slot integer extension has no active slots"); + Value assembled; + for (int64_t slot = 0; slot < activeSlots; ++slot) { + int64_t partIndex = slot % partFactor; + int64_t sourceLane = slot / partFactor; + FailureOr laneIndexScalar = createScalarOffsetConstant( + op.getLoc(), indexType.getElementType(), sourceLane, rewriter); + FailureOr laneMask = createLaneRangeMask( + op.getLoc(), *resultMaskType, slot, slot + 1, rewriter); + if (failed(laneIndexScalar) || failed(laneMask)) + return rewriter.notifyMatchFailure( + op, "failed to build group-slot integer extension slot mask"); + Value laneIndex = + rewriter + .create(op.getLoc(), indexType, *laneIndexScalar, + *resultAllMask, /*position=*/nullptr) + .getResult(); + Value selected = + rewriter + .create(op.getLoc(), resultVRegType, + convertedParts[partIndex], laneIndex) + .getResult(); + if (!assembled) { + assembled = selected; + continue; + } + assembled = rewriter + .create(op.getLoc(), resultVRegType, selected, + assembled, *laneMask) + .getResult(); + } + + results.push_back(assembled); + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } + + SmallVector resultVRegTypes; + resultVRegTypes.reserve(resultTypes.size()); + for (Type resultType : resultTypes) { + auto resultVRegType = dyn_cast(resultType); + if (!resultVRegType || + !isa(resultVRegType.getElementType()) || + (!resultVRegTypes.empty() && + resultVRegType != resultVRegTypes.front())) + return rewriter.notifyMatchFailure( + op, "unsupported physical integer extension result type"); + resultVRegTypes.push_back(resultVRegType); + } + + unsigned sourceBits = + pto::getPTOStorageElemBitWidth(sourceType.getElementType()); + unsigned resultBits = pto::getPTOStorageElemBitWidth( + resultVRegTypes.front().getElementType()); + if (sourceLayout && resultLayout && sourceLayout.isContiguous() && + resultLayout.isContiguous() && resultLayout.getLaneStride() == 1 && + ((resultBits == sourceBits * 2 && sourceLayout.getLaneStride() == 2) || + (resultBits == sourceBits * 4 && sourceLayout.getLaneStride() == 4)) && + resultTypes.size() == sourceParts.size()) { + StringRef part = + resultBits == sourceBits * 2 ? StringRef("EVEN") : StringRef("P0"); + FailureOr mask = + createAllTrueMaskForVReg(op.getLoc(), sourceType, rewriter); + if (failed(mask)) + return rewriter.notifyMatchFailure( + op, "failed to build integer extension seed mask"); + + SmallVector results; + results.reserve(resultTypes.size()); + for (auto [sourcePart, resultType] : + llvm::zip_equal(sourceParts, resultVRegTypes)) { + results.push_back(rewriter + .create(op.getLoc(), resultType, + sourcePart, *mask, + /*rnd=*/nullptr, /*sat=*/nullptr, + rewriter.getStringAttr(part)) + .getResult()); + } + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } + + ArrayRef parts; + int64_t factor = 0; + if (resultBits == sourceBits * 2 && + resultTypes.size() == 2 * sourceParts.size()) { + static constexpr StringRef kEvenOddParts[] = {"EVEN", "ODD"}; + parts = kEvenOddParts; + factor = 2; + } else if (resultBits == sourceBits * 4 && + resultTypes.size() == 4 * sourceParts.size()) { + static constexpr StringRef kPacked4Parts[] = {"P0", "P1", "P2", "P3"}; + parts = kPacked4Parts; + factor = 4; + } else { + return rewriter.notifyMatchFailure( + op, "unsupported physical integer extension source/result width " + "relation"); + } + + FailureOr mask = + createAllTrueMaskForVReg(op.getLoc(), sourceType, rewriter); + if (failed(mask)) + return rewriter.notifyMatchFailure( + op, "failed to build integer extension seed mask"); + + SmallVector results; + results.reserve(resultTypes.size()); + for (int64_t partIndex = 0; partIndex < factor; ++partIndex) { + for (auto [chunkIndex, sourcePart] : llvm::enumerate(sourceParts)) { + VRegType resultType = + resultVRegTypes[partIndex * sourceParts.size() + chunkIndex]; + results.push_back( + rewriter + .create(op.getLoc(), resultType, sourcePart, *mask, + /*rnd=*/nullptr, /*sat=*/nullptr, + rewriter.getStringAttr(parts[partIndex])) + .getResult()); + } + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +// TruncI lowering support matrix +// +// Keep this comment aligned with both: +// 1. verifySupportedVMIToVPTOOps() diagnostics below, and +// 2. the actual OneToN lowering implemented in this pattern. +// +// Dense logical layouts +// - deinterleaved factor 2/4 -> contiguous +// Example: 32 -> 16 or 32 -> 8. +// Lowering shape: emit vcvt parts EVEN/ODD or P0/P1/P2/P3, then merge +// physical results when multiple source chunks contribute to one result. +// - deinterleaved factor 4 -> deinterleaved factor 2 +// Example: 32 -> 16. +// Lowering shape: emit vcvt EVEN/ODD per source chunk pair. +// - contiguous lane_stride = 1 -> contiguous lane_stride = 2/4 +// Example: 16 -> 8 lane_stride=2, 32 -> 8 lane_stride=4. +// Lowering shape: emit vcvt into the widened physical carrier selected by +// the lane-stride result type. +// +// Group-slots logical layouts: group_slots(num_groups=G, slots=1 or 8) +// - 32-bit integer -> 16-bit integer, same group_slots layout +// Lowering shape: direct vcvt with part = EVEN. +// - 32-bit integer -> 8-bit integer, same group_slots layout +// Lowering shape: direct vcvt with part = P0. +// - 16-bit unsigned integer -> 8-bit unsigned integer, slots = 8, +// result lane_stride = 2 +// Lowering shape: direct vcvt with part = EVEN. +// - 32-bit integer -> 8-bit integer, result lane_stride = 4 +// Lowering shape: no vcvt; keep/bitcast the 32-bit carrier and let the +// later store consume it as PK4_B32. +struct OneToNVMITruncIOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMITruncIOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto sourceVMIType = cast(op.getSource().getType()); + auto resultVMIType = cast(op.getResult().getType()); + ValueRange sourceParts = adaptor.getSource(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + + VMILayoutAttr sourceLayout = sourceVMIType.getLayoutAttr(); + VMILayoutAttr resultLayout = resultVMIType.getLayoutAttr(); + if (sourceLayout && resultLayout && sourceLayout.isGroupSlots() && + resultLayout.isGroupSlots()) { + unsigned sourceLogicalBits = + pto::getPTOStorageElemBitWidth(sourceVMIType.getElementType()); + unsigned resultLogicalBits = + pto::getPTOStorageElemBitWidth(resultVMIType.getElementType()); + bool supportsDirectGroupSlotTrunc = + sourceLogicalBits == 32 && + (resultLogicalBits == 16 || resultLogicalBits == 8); + bool supportsPackedU16ToU8GroupSlotTrunc = + sourceLogicalBits == 16 && resultLogicalBits == 8 && + sourceLayout.getSlots() == 8 && resultLayout.getSlots() == 8 && + resultLayout.hasLaneStride() && resultLayout.getLaneStride() == 2; + if (sourceLayout.getNumGroups() != resultLayout.getNumGroups() || + sourceLayout.getSlots() != resultLayout.getSlots() || + (sourceLayout.getSlots() != 1 && sourceLayout.getSlots() != 8) || + (!supportsDirectGroupSlotTrunc && + !supportsPackedU16ToU8GroupSlotTrunc) || + sourceParts.size() != resultTypes.size()) + return rewriter.notifyMatchFailure( + op, "unsupported group-slot trunci shape"); + + SmallVector results; + results.reserve(resultTypes.size()); + StringAttr sat = rewriter.getStringAttr("SAT"); + const char *activeSlotPattern = + sourceLayout.getSlots() == 1 ? "PAT_VL1" : "PAT_VL8"; + StringRef activeSlotGranularity = sourceLogicalBits == 16 ? "b16" : "b32"; + FailureOr activeSlotMask = createPrefixMask( + op.getLoc(), MaskType::get(rewriter.getContext(), activeSlotGranularity), + activeSlotPattern, rewriter); + if (failed(activeSlotMask)) + return rewriter.notifyMatchFailure( + op, "failed to build group-slot trunci active slot mask"); + for (auto [sourcePart, physicalResultType] : + llvm::zip_equal(sourceParts, resultTypes)) { + auto sourceType = dyn_cast(sourcePart.getType()); + auto resultType = dyn_cast(physicalResultType); + if (!sourceType || + pto::getPTOStorageElemBitWidth(sourceType.getElementType()) != + sourceLogicalBits || + !resultType) + return rewriter.notifyMatchFailure( + op, "unsupported group-slot trunci physical type"); + + if (supportsPackedU16ToU8GroupSlotTrunc) { + results.push_back(rewriter + .create(op.getLoc(), resultType, + sourcePart, *activeSlotMask, + /*rnd=*/nullptr, sat, + rewriter.getStringAttr("EVEN")) + .getResult()); + continue; + } + + unsigned physicalResultBits = + pto::getPTOStorageElemBitWidth(resultType.getElementType()); + if (resultLayout.hasLaneStride() && resultLayout.getLaneStride() == 4 && + pto::getPTOStorageElemBitWidth(resultVMIType.getElementType()) == + 8 && + physicalResultBits == 32) { + if (sourcePart.getType() == resultType) { + results.push_back(sourcePart); + } else { + results.push_back( + rewriter.create(op.getLoc(), resultType, sourcePart) + .getResult()); + } + continue; + } + + if (physicalResultBits != 16 && physicalResultBits != 8) + return rewriter.notifyMatchFailure( + op, "unsupported group-slot trunci physical type"); + + StringAttr part = physicalResultBits == 16 + ? rewriter.getStringAttr("EVEN") + : rewriter.getStringAttr("P0"); + results.push_back(rewriter + .create(op.getLoc(), resultType, + sourcePart, *activeSlotMask, + /*rnd=*/nullptr, sat, part) + .getResult()); + } + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } + + if (sourceParts.empty() || resultTypes.empty()) + return rewriter.notifyMatchFailure( + op, "trunci requires non-empty physical source and result parts"); + + auto sourceType0 = dyn_cast(sourceParts.front().getType()); + auto resultType0 = dyn_cast(resultTypes.front()); + if (!sourceType0 || !isa(sourceType0.getElementType()) || + !resultType0 || !isa(resultType0.getElementType())) + return rewriter.notifyMatchFailure( + op, "unsupported physical trunci source/result type"); + for (Value sourcePart : sourceParts) { + auto sourceType = dyn_cast(sourcePart.getType()); + if (!sourceType || sourceType != sourceType0) + return rewriter.notifyMatchFailure( + op, "trunci source physical parts must have matching integer type"); + } + for (Type resultType : resultTypes) { + auto resultVRegType = dyn_cast(resultType); + if (!resultVRegType || resultVRegType != resultType0) + return rewriter.notifyMatchFailure( + op, "trunci result physical parts must have matching integer type"); + } + + unsigned sourceBits = + pto::getPTOStorageElemBitWidth(sourceType0.getElementType()); + unsigned resultBits = + pto::getPTOStorageElemBitWidth(resultType0.getElementType()); + if (sourceBits == 0 || resultBits == 0 || sourceBits % resultBits != 0) + return rewriter.notifyMatchFailure( + op, "unsupported physical trunci source/result width relation"); + int64_t factor = sourceBits / resultBits; + if (sourceLayout && resultLayout && sourceLayout.isContiguous() && + sourceLayout.getLaneStride() == 1 && resultLayout.isContiguous() && + resultLayout.getLaneStride() == factor && + sourceParts.size() == resultTypes.size()) { + if (factor != 2 && factor != 4) + return rewriter.notifyMatchFailure( + op, "unsupported dense lane_stride trunci result layout"); + StringAttr part = rewriter.getStringAttr(factor == 2 ? "EVEN" : "P0"); + FailureOr sourceMask = + createAllTrueMaskForVReg(op.getLoc(), sourceType0, rewriter); + if (failed(sourceMask)) + return rewriter.notifyMatchFailure(op, "failed to build trunci masks"); + + StringAttr sat = rewriter.getStringAttr("SAT"); + SmallVector results; + results.reserve(resultTypes.size()); + for (auto [sourcePart, resultType] : + llvm::zip_equal(sourceParts, resultTypes)) { + results.push_back(rewriter + .create(op.getLoc(), resultType, + sourcePart, *sourceMask, + /*rnd=*/nullptr, sat, part) + .getResult()); + } + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } + + if ((factor != 2 && factor != 4) || + sourceParts.size() != resultTypes.size() * factor) + return rewriter.notifyMatchFailure( + op, "unsupported physical trunci source/result arity relation"); + + ArrayRef parts; + if (factor == 2) { + static constexpr StringRef kEvenOddParts[] = {"EVEN", "ODD"}; + parts = kEvenOddParts; + } else if (factor == 4) { + static constexpr StringRef kPacked4Parts[] = {"P0", "P1", "P2", "P3"}; + parts = kPacked4Parts; + } else { + return rewriter.notifyMatchFailure( + op, "unsupported physical trunci source/result width relation"); + } + + FailureOr sourceMask = + createAllTrueMaskForVReg(op.getLoc(), sourceType0, rewriter); + FailureOr resultMask = + createAllTrueMaskForVReg(op.getLoc(), resultType0, rewriter); + if (failed(sourceMask) || failed(resultMask)) + return rewriter.notifyMatchFailure(op, "failed to build trunci masks"); + + StringAttr sat = rewriter.getStringAttr("SAT"); + SmallVector results; + results.reserve(resultTypes.size()); + for (int64_t resultIndex = 0, resultCount = resultTypes.size(); + resultIndex < resultCount; ++resultIndex) { + Type resultType = resultTypes[resultIndex]; + SmallVector partials; + partials.reserve(parts.size()); + for (int64_t partIndex = 0; partIndex < factor; ++partIndex) { + Value sourcePart = sourceParts[resultIndex * factor + partIndex]; + partials.push_back( + rewriter + .create(op.getLoc(), resultType, sourcePart, + *sourceMask, /*rnd=*/nullptr, sat, + rewriter.getStringAttr(parts[partIndex])) + .getResult()); + } + + Value merged = partials.front(); + for (Value partial : llvm::drop_begin(partials)) + merged = rewriter + .create(op.getLoc(), resultType, merged, partial, + *resultMask) + .getResult(); + results.push_back(merged); + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMIFPToSIOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIFPToSIOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + ValueRange sourceParts = adaptor.getSource(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + if (sourceParts.size() != resultTypes.size()) + return rewriter.notifyMatchFailure( + op, "fptosi physical source/result arity mismatch"); + + SmallVector results; + results.reserve(resultTypes.size()); + StringAttr rnd = rewriter.getStringAttr("R"); + StringAttr sat = rewriter.getStringAttr("SAT"); + for (auto [sourcePart, resultType] : + llvm::zip_equal(sourceParts, resultTypes)) { + auto sourceType = dyn_cast(sourcePart.getType()); + auto resultVRegType = dyn_cast(resultType); + if (!sourceType || !sourceType.getElementType().isF32() || + !resultVRegType || + !isa(resultVRegType.getElementType()) || + pto::getPTOStorageElemBitWidth(resultVRegType.getElementType()) != 32) + return rewriter.notifyMatchFailure( + op, "fptosi requires physical f32 source and 32-bit integer " + "result chunks"); + + FailureOr mask = + createAllTrueMaskForVReg(op.getLoc(), sourceType, rewriter); + if (failed(mask)) + return rewriter.notifyMatchFailure(op, "failed to build fptosi mask"); + results.push_back(rewriter + .create(op.getLoc(), resultVRegType, + sourcePart, *mask, rnd, sat, + /*part=*/nullptr) + .getResult()); + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMISIToFPOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMISIToFPOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + ValueRange sourceParts = adaptor.getSource(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + if (sourceParts.size() != resultTypes.size()) + return rewriter.notifyMatchFailure( + op, "sitofp physical source/result arity mismatch"); + + SmallVector results; + results.reserve(resultTypes.size()); + StringAttr rnd = rewriter.getStringAttr("R"); + for (auto [sourcePart, resultType] : + llvm::zip_equal(sourceParts, resultTypes)) { + auto sourceType = dyn_cast(sourcePart.getType()); + auto resultVRegType = dyn_cast(resultType); + if (!sourceType || !isa(sourceType.getElementType()) || + pto::getPTOStorageElemBitWidth(sourceType.getElementType()) != 32 || + !resultVRegType || !resultVRegType.getElementType().isF32()) + return rewriter.notifyMatchFailure( + op, "sitofp requires physical 32-bit integer source and f32 " + "result chunks"); + + FailureOr mask = + createAllTrueMaskForVReg(op.getLoc(), sourceType, rewriter); + if (failed(mask)) + return rewriter.notifyMatchFailure(op, "failed to build sitofp mask"); + results.push_back(rewriter + .create(op.getLoc(), resultVRegType, + sourcePart, *mask, rnd, + /*sat=*/nullptr, /*part=*/nullptr) + .getResult()); + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMIBitcastOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIBitcastOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + ValueRange sourceParts = adaptor.getSource(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + if (sourceParts.size() != resultTypes.size()) + return rewriter.notifyMatchFailure(op, "physical bitcast arity mismatch"); + + SmallVector results; + results.reserve(resultTypes.size()); + for (auto [sourcePart, resultType] : + llvm::zip_equal(sourceParts, resultTypes)) { + if (!isa(sourcePart.getType()) || !isa(resultType)) + return rewriter.notifyMatchFailure( + op, "physical bitcast part type mismatch"); + results.push_back( + rewriter.create(op.getLoc(), resultType, sourcePart) + .getResult()); + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMIChannelSplitOpPattern + : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIChannelSplitOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + int64_t channels = op.getNumResults(); + if (channels != 2 && channels != 4) + return rewriter.notifyMatchFailure( + op, "channel_split only supports 2 or 4 channels"); + + auto sourceType = cast(op.getSource().getType()); + VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); + auto channelLayout = + VMILayoutAttr::getDeinterleaved(rewriter.getContext(), channels); + if (!sourceLayout || + (!sourceLayout.isContiguous() && sourceLayout != channelLayout)) + return rewriter.notifyMatchFailure( + op, + "channel_split requires contiguous or matching deinterleaved source " + "layout"); + for (Value result : op.getResults()) { + auto resultType = cast(result.getType()); + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + if (!resultLayout || !resultLayout.isContiguous()) + return rewriter.notifyMatchFailure( + op, "channel_split requires contiguous result layouts"); + } + + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + FailureOr> results = + materializeDataLayoutConversion(op, adaptor.getSource(), resultTypes, + sourceLayout, channelLayout, rewriter); + if (failed(results)) + return failure(); + + replaceOpWithFlatConvertedValues(rewriter, op, *results, *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMIChannelMergeOpPattern + : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIChannelMergeOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + int64_t channels = op.getInputs().size(); + if (channels != 2 && channels != 4) + return rewriter.notifyMatchFailure( + op, "channel_merge only supports 2 or 4 channels"); + + for (Value input : op.getInputs()) { + auto inputType = cast(input.getType()); + VMILayoutAttr inputLayout = inputType.getLayoutAttr(); + if (!inputLayout || !inputLayout.isContiguous()) + return rewriter.notifyMatchFailure( + op, "channel_merge requires contiguous input layouts"); + } + auto resultType = cast(op.getResult().getType()); + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + auto channelLayout = + VMILayoutAttr::getDeinterleaved(rewriter.getContext(), channels); + if (!resultLayout || + (!resultLayout.isContiguous() && resultLayout != channelLayout)) + return rewriter.notifyMatchFailure( + op, + "channel_merge requires contiguous or matching deinterleaved result " + "layout"); + + FailureOr> maybeResultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybeResultTypes)) + return failure(); + FailureOr> results = + materializeDataLayoutConversion( + op, flattenOneToNOperands(adaptor.getOperands()), + *maybeResultTypes, channelLayout, resultLayout, rewriter); + if (failed(results)) + return failure(); + + replaceOpWithFlatConvertedValues(rewriter, op, *results, *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNVMIShuffleOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(VMIShuffleOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + ValueRange sourceParts = adaptor.getSource(); + FailureOr> maybe_resultTypes = + getConvertedResultTypes(op, 0, *this->getTypeConverter()); + if (failed(maybe_resultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybe_resultTypes); + std::string reason; + FailureOr> sourceFlatIndices = + computeShuffleForwardingSourceParts(op, &reason); + if (succeeded(sourceFlatIndices)) { + SmallVector results; + results.reserve(resultTypes.size()); + for (int64_t sourceFlatIndex : *sourceFlatIndices) { + if (sourceFlatIndex >= static_cast(sourceParts.size())) + return rewriter.notifyMatchFailure( + op, "shuffle forwarding source part range is out of bounds"); + results.push_back(sourceParts[sourceFlatIndex]); + } + + if (failed( + verifyIdentityPartForwarding(op, results, resultTypes, rewriter))) + return failure(); + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } + + std::string splatReason; + FailureOr splatSource = + computeShuffleLane0SplatSourcePart(op, &splatReason); + if (succeeded(splatSource)) { + if (*splatSource >= static_cast(sourceParts.size())) + return rewriter.notifyMatchFailure( + op, "shuffle lane0 splat source part range is out of bounds"); + + SmallVector results; + results.reserve(resultTypes.size()); + Value sourcePart = sourceParts[*splatSource]; + for (Type resultType : resultTypes) { + auto sourceVRegType = dyn_cast(sourcePart.getType()); + auto resultVRegType = dyn_cast(resultType); + if (!sourceVRegType || !resultVRegType || + sourceVRegType != resultVRegType) + return rewriter.notifyMatchFailure( + op, "shuffle lane0 splat requires matching physical vreg type"); + FailureOr mask = + createAllTrueMaskForVReg(op.getLoc(), resultVRegType, rewriter); + if (failed(mask)) + return rewriter.notifyMatchFailure( + op, "failed to create shuffle lane0 splat mask"); + results.push_back(rewriter + .create(op.getLoc(), resultType, + sourcePart, *mask, + rewriter.getStringAttr("LOWEST")) + .getResult()); + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } + + std::string vselrReason; + FailureOr> vselrPlans = + computeShuffleVselrPlans(op, &vselrReason); + if (failed(vselrPlans)) + return rewriter.notifyMatchFailure(op, + Twine("shuffle vselr ") + vselrReason); + + if (vselrPlans->size() != resultTypes.size()) + return rewriter.notifyMatchFailure(op, "shuffle vselr arity mismatch"); + + SmallVector results; + results.reserve(resultTypes.size()); + for (auto [plan, resultType] : llvm::zip_equal(*vselrPlans, resultTypes)) { + if (plan.sourceFlatIndex >= static_cast(sourceParts.size())) + return rewriter.notifyMatchFailure( + op, "shuffle vselr source part range is out of bounds"); + + auto sourceVRegType = + dyn_cast(sourceParts[plan.sourceFlatIndex].getType()); + auto resultVRegType = dyn_cast(resultType); + if (!sourceVRegType || !resultVRegType || + sourceVRegType.getElementCount() != + resultVRegType.getElementCount() || + sourceVRegType.getElementType() != resultVRegType.getElementType()) + return rewriter.notifyMatchFailure( + op, "shuffle vselr source/result type mismatch"); + + unsigned indexBits = + pto::getPTOStorageElemBitWidth(sourceVRegType.getElementType()); + if (indexBits != 8 && indexBits != 16 && indexBits != 32) + return rewriter.notifyMatchFailure( + op, "shuffle vselr requires 8/16/32-bit index elements"); + + auto indexElementType = + IntegerType::get(rewriter.getContext(), indexBits); + Type indexType = + VRegType::get(rewriter.getContext(), sourceVRegType.getElementCount(), + indexElementType); + FailureOr base = createScalarOffsetConstant( + op.getLoc(), indexElementType, plan.baseLane, rewriter); + if (failed(base)) + return rewriter.notifyMatchFailure( + op, "failed to materialize shuffle vselr index base"); + StringAttr orderAttr = + plan.descending ? rewriter.getStringAttr("DESC") : StringAttr{}; + Value indexVector = + rewriter.create(op.getLoc(), indexType, *base, orderAttr) + .getResult(); + results.push_back(rewriter + .create(op.getLoc(), resultType, + sourceParts[plan.sourceFlatIndex], + indexVector) + .getResult()); + } + + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } +}; + +Block *convertBranchDestBlock(Block *block, ConversionPatternRewriter &rewriter, + const TypeConverter &typeConverter, + llvm::DenseMap &converted) { + auto [it, inserted] = converted.try_emplace(block, nullptr); + if (!inserted) + return it->second; + + TypeConverter::SignatureConversion argMapping(block->getNumArguments()); + if (failed(typeConverter.convertSignatureArgs(block->getArgumentTypes(), + argMapping)) || + !hasNonIdentitySignatureConversion(block->getArgumentTypes(), + argMapping)) { + it->second = block; + return block; + } + + Block *newBlock = + rewriter.applySignatureConversion(block, argMapping, &typeConverter); + it->second = newBlock; + return newBlock; +} + +struct OneToNCFBranchOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(cf::BranchOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto *converter = this->getTypeConverter(); + llvm::DenseMap convertedBlocks; + Block *dest = convertBranchDestBlock(op.getDest(), rewriter, *converter, + convertedBlocks); + SmallVector destOperands = + flattenOneToNOperands(adaptor.getDestOperands()); + + if (isIdentityOneToNValueMapping(op.getDestOperands(), + adaptor.getDestOperands()) && + dest == op.getDest()) + return failure(); + + rewriter.replaceOpWithNewOp(op, dest, destOperands); + return success(); + } +}; + +struct OneToNCFCondBranchOpPattern + : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(cf::CondBranchOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto *converter = this->getTypeConverter(); + llvm::DenseMap convertedBlocks; + Block *trueDest = convertBranchDestBlock(op.getTrueDest(), rewriter, + *converter, convertedBlocks); + Block *falseDest = convertBranchDestBlock(op.getFalseDest(), rewriter, + *converter, convertedBlocks); + + if (isIdentityOneToNValueMapping(op.getTrueDestOperands(), + adaptor.getTrueDestOperands()) && + isIdentityOneToNValueMapping(op.getFalseDestOperands(), + adaptor.getFalseDestOperands()) && + trueDest == op.getTrueDest() && falseDest == op.getFalseDest()) + return failure(); + + ValueRange condition = adaptor.getCondition(); + if (condition.size() != 1) + return rewriter.notifyMatchFailure( + op, "condition converted to multiple values"); + + SmallVector trueOperands = + flattenOneToNOperands(adaptor.getTrueDestOperands()); + SmallVector falseOperands = + flattenOneToNOperands(adaptor.getFalseDestOperands()); + + rewriter.replaceOpWithNewOp(op, condition.front(), + trueDest, trueOperands, + falseDest, falseOperands); + return success(); + } +}; + +struct OneToNCFSwitchOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(cf::SwitchOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto *converter = this->getTypeConverter(); + llvm::DenseMap convertedBlocks; + Block *defaultDest = convertBranchDestBlock( + op.getDefaultDestination(), rewriter, *converter, convertedBlocks); + + SmallVector caseDests; + caseDests.reserve(op.getCaseDestinations().size()); + for (Block *dest : op.getCaseDestinations()) + caseDests.push_back( + convertBranchDestBlock(dest, rewriter, *converter, convertedBlocks)); + + bool changed = defaultDest != op.getDefaultDestination(); + for (auto [oldDest, newDest] : + llvm::zip(op.getCaseDestinations(), caseDests)) + changed |= oldDest != newDest; + changed |= !isIdentityOneToNValueMapping(op.getDefaultOperands(), + adaptor.getDefaultOperands()); + for (auto [originalOperands, convertedOperands] : + llvm::zip(op.getCaseOperands(), adaptor.getCaseOperands())) + changed |= + !isIdentityOneToNValueMapping(originalOperands, convertedOperands); + if (!changed) + return failure(); + + ValueRange flag = adaptor.getFlag(); + if (flag.size() != 1) + return rewriter.notifyMatchFailure(op, + "flag converted to multiple values"); + + SmallVector defaultOperands; + SmallVector> caseOperandStorage; + SmallVector caseOperands; + defaultOperands = flattenOneToNOperands(adaptor.getDefaultOperands()); + + caseOperandStorage.reserve(op.getCaseOperandSegments().size()); + caseOperands.reserve(op.getCaseOperandSegments().size()); + for (ArrayRef convertedOperands : adaptor.getCaseOperands()) + caseOperandStorage.push_back(flattenOneToNOperands(convertedOperands)); + for (SmallVector &operands : caseOperandStorage) + caseOperands.push_back(operands); + + rewriter.replaceOpWithNewOp( + op, flag.front(), defaultDest, defaultOperands, op.getCaseValuesAttr(), + caseDests, caseOperands); + return success(); + } +}; + +struct OneToNSCFExecuteRegionOpPattern + : OpConversionPattern { + using OpConversionPattern< + scf::ExecuteRegionOp>::OpConversionPattern; + + LogicalResult + matchAndRewrite(scf::ExecuteRegionOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + FailureOr> maybeResultTypes = + getConvertedResultTypes(op, *this->getTypeConverter()); + if (failed(maybeResultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybeResultTypes); + if (resultTypes == op->getResultTypes()) + return failure(); + + auto newOp = + rewriter.create(op.getLoc(), resultTypes); + newOp->setAttrs(op->getAttrs()); + rewriter.inlineRegionBefore(op.getRegion(), newOp.getRegion(), + newOp.getRegion().end()); + replaceOpWithFlatConvertedValues( + rewriter, op, newOp->getResults(), *this->getTypeConverter()); + return success(); + } +}; + +struct OneToNSCFIndexSwitchOpPattern + : OpConversionPattern { + using OpConversionPattern< + scf::IndexSwitchOp>::OpConversionPattern; + + LogicalResult + matchAndRewrite(scf::IndexSwitchOp op, OneToNOpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + ValueRange arg = adaptor.getArg(); + if (arg.size() != 1) + return rewriter.notifyMatchFailure( + op, "index_switch selector converted to multiple values"); + + FailureOr> maybeResultTypes = + getConvertedResultTypes(op, *this->getTypeConverter()); + if (failed(maybeResultTypes)) + return failure(); + SmallVector resultTypes = std::move(*maybeResultTypes); + if (resultTypes == op->getResultTypes()) + return failure(); + + auto newOp = rewriter.create( + op.getLoc(), resultTypes, arg.front(), op.getCases(), op.getNumCases()); + newOp->setAttrs(op->getAttrs()); + rewriter.inlineRegionBefore(op.getDefaultRegion(), newOp.getDefaultRegion(), + newOp.getDefaultRegion().end()); + for (auto [srcRegion, dstRegion] : + llvm::zip(op.getCaseRegions(), newOp.getCaseRegions())) + rewriter.inlineRegionBefore(srcRegion, dstRegion, dstRegion.end()); + replaceOpWithFlatConvertedValues( + rewriter, op, newOp->getResults(), *this->getTypeConverter()); + return success(); + } +}; + +void populateVMIConversionPatterns( + VMIToVPTOTypeConverter &typeConverter, RewritePatternSet &patterns) { + populateFunctionOpInterfaceTypeConversionPattern(patterns, typeConverter); + populateCallOpTypeConversionPattern(patterns, typeConverter); + populateReturnOpTypeConversionPattern(patterns, typeConverter); + scf::populateSCFStructuralTypeConversions(typeConverter, patterns); + patterns.add(typeConverter, patterns.getContext()); + patterns.add( + typeConverter, patterns.getContext()); + patterns.add( + typeConverter, patterns.getContext()); + patterns.add< + OneToNVMIEnsureLayoutOpPattern, OneToNVMIEnsureMaskLayoutOpPattern, + OneToNVMIBroadcastOpPattern, OneToNVMIIotaOpPattern, + OneToNVMIConstantOpPattern, OneToNVMIConstantMaskOpPattern, + OneToNVMICreateMaskOpPattern, OneToNVMICreateGroupMaskOpPattern, + OneToNVMIMaskBinaryOpPattern, + OneToNVMIMaskBinaryOpPattern, + OneToNVMIMaskBinaryOpPattern, + OneToNVMIMaskUnaryOpPattern, OneToNVMILoadOpPattern, + OneToNVMIDeinterleaveLoadOpPattern, OneToNVMIGroupLoadOpPattern, + OneToNVMIGroupSlotLoadOpPattern, OneToNVMIStrideLoadOpPattern, + OneToNVMIMaskedLoadOpPattern, OneToNVMIGatherOpPattern, + OneToNVMIExpandLoadOpPattern, OneToNVMIStoreOpPattern, + OneToNVMIInterleaveStoreOpPattern, OneToNVMIGroupStoreOpPattern, + OneToNVMIMaskedStoreOpPattern, OneToNVMIStrideStoreOpPattern, + OneToNVMIScatterOpPattern, OneToNVMIBinaryOpPattern, + OneToNVMIBinaryOpPattern, + OneToNVMIBinaryOpPattern, + OneToNVMIBinaryOpPattern, + OneToNVMIBinaryOpPattern, + OneToNVMIBinaryOpPattern, OneToNVMIVmullOpPattern, + OneToNVMIFmaOpPattern, OneToNVMIBinaryOpPattern, + OneToNVMIBinaryOpPattern, + OneToNVMIBinaryOpPattern, + OneToNVMIUnaryOpPattern, + OneToNVMIUnaryOpPattern, + OneToNVMIUnaryOpPattern, + OneToNVMIUnaryOpPattern, + OneToNVMIUnaryOpPattern, + OneToNVMIUnaryOpPattern, + OneToNVMIUnaryOpPattern, + OneToNVMIBinaryOpPattern, + OneToNVMIBinaryOpPattern, + OneToNVMIBinaryOpPattern, + OneToNVMIBinaryOpPattern, + OneToNVMIBinaryOpPattern, + OneToNVMIBinaryOpPattern, + OneToNVMIUnaryOpPattern, + OneToNVMICmpOpPattern, OneToNVMICmpOpPattern, + OneToNVMISelectOpPattern, OneToNVMIActivePrefixIndexOpPattern, + OneToNVMICompressOpPattern, OneToNVMICompressStoreOpPattern, + OneToNVMIReduceAddIOpPattern, OneToNVMIReduceAddFOpPattern, + OneToNVMIGroupBroadcastOpPattern, OneToNVMIVdhistOpPattern, + OneToNVMIVchistOpPattern, + OneToNVMIReduceMinMaxOpPattern, + OneToNVMIReduceMinMaxOpPattern, + OneToNVMIReduceMinMaxOpPattern, + OneToNVMIReduceMinMaxOpPattern, + OneToNVMIExtFOpPattern, OneToNVMITruncFOpPattern, + OneToNVMIExtIOpPattern, OneToNVMIExtIOpPattern, + OneToNVMITruncIOpPattern, OneToNVMIFPToSIOpPattern, + OneToNVMISIToFPOpPattern, OneToNVMIBitcastOpPattern, + OneToNVMIInterleaveOpPattern, + OneToNVMIInterleaveOpPattern, + OneToNVMIChannelSplitOpPattern, OneToNVMIChannelMergeOpPattern, + OneToNVMIShuffleOpPattern>(typeConverter, patterns.getContext()); + patterns.add( + typeConverter, patterns.getContext()); + patterns.add< + OneToNVMIGroupReduceOpPattern, + OneToNVMIGroupReduceOpPattern, + OneToNVMIGroupReduceOpPattern, + OneToNVMIGroupReduceOpPattern, + OneToNVMIGroupReduceOpPattern, + OneToNVMIGroupReduceOpPattern>(typeConverter, + patterns.getContext()); + patterns.add( + typeConverter, patterns.getContext()); +} + +LogicalResult verifyNoResidualVMIIR(ModuleOp module) { + WalkResult result = module.walk([&](Operation *op) { + if (auto createMask = dyn_cast(op)) { + if (!createMask.getActiveLanes().getDefiningOp()) { + createMask.emitError() + << kVMIDiagUnsupportedPrefix + << "dynamic pto.vmi.create_mask active_lanes could not be lowered " + "by the current runtime predicate generation plan"; + return WalkResult::interrupt(); + } + } + if (auto constant = dyn_cast(op)) { + auto denseAttr = dyn_cast(constant.getValue()); + if (denseAttr && !denseAttr.isSplat()) { + constant.emitError() + << kVMIDiagUnsupportedPrefix + << "non-splat pto.vmi.constant requires a vreg immediate or " + "scratch materialization plan"; + return WalkResult::interrupt(); + } + } + if (isVMIOp(op) || hasVMIType(op)) { + op->emitError() << kVMIDiagResidualOpPrefix + << "failed to convert all VMI ops/types to VPTO"; + return WalkResult::interrupt(); + } + return WalkResult::advance(); + }); + return failure(result.wasInterrupted()); +} + +LogicalResult checkSupportedExtFShape(VMIExtFOp op, + std::string *reason = nullptr) { + VMILayoutSupport supports; + if (failed(supports.getExtFSupport(op, reason))) + return failure(); + return success(); +} + +LogicalResult checkSupportedTruncFShape(VMITruncFOp op, + std::string *reason = nullptr) { + VMILayoutSupport supports; + if (failed(supports.getTruncFSupport(op, reason))) + return failure(); + return success(); +} + +LogicalResult checkSupportedExtSIShape(VMIExtSIOp op, + std::string *reason = nullptr) { + VMILayoutSupport supports; + if (failed(supports.getExtSISupport(op, reason))) + return failure(); + return success(); +} + +LogicalResult checkSupportedExtUIShape(VMIExtUIOp op, + std::string *reason = nullptr) { + VMILayoutSupport supports; + if (failed(supports.getExtUISupport(op, reason))) + return failure(); + return success(); +} + +LogicalResult checkSupportedTruncIShape(VMITruncIOp op, + std::string *reason = nullptr) { + VMILayoutSupport supports; + if (failed(supports.getTruncISupport(op, reason))) + return failure(); + return success(); +} + +LogicalResult checkSupportedFPToSIShape(VMIFPToSIOp op, + std::string *reason = nullptr) { + auto fail = [&](const Twine &message) { + if (reason) + *reason = message.str(); + return failure(); + }; + + auto sourceType = cast(op.getSource().getType()); + auto resultType = cast(op.getResult().getType()); + VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + if (!sourceLayout || !resultLayout) + return fail("requires assigned source/result layouts"); + if (sourceLayout != resultLayout) + return fail("requires source/result layouts to match"); + if (!sourceType.getElementType().isF32()) + return fail("requires f32 source element type"); + if (!isa(resultType.getElementType()) || + pto::getPTOStorageElemBitWidth(resultType.getElementType()) != 32) + return fail("requires 32-bit integer result element type"); + FailureOr sourceArity = getVMIPhysicalArity(sourceType); + FailureOr resultArity = getVMIPhysicalArity(resultType); + if (failed(sourceArity) || failed(resultArity) || + *sourceArity != *resultArity) + return fail("requires matching computable physical arity"); + return success(); +} + +LogicalResult checkSupportedSIToFPShape(VMISIToFPOp op, + std::string *reason = nullptr) { + auto fail = [&](const Twine &message) { + if (reason) + *reason = message.str(); + return failure(); + }; + + auto sourceType = cast(op.getSource().getType()); + auto resultType = cast(op.getResult().getType()); + VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + if (!sourceLayout || !resultLayout) + return fail("requires assigned source/result layouts"); + if (sourceLayout != resultLayout) + return fail("requires source/result layouts to match"); + if (!isa(sourceType.getElementType()) || + pto::getPTOStorageElemBitWidth(sourceType.getElementType()) != 32) + return fail("requires 32-bit integer source element type"); + if (!resultType.getElementType().isF32()) + return fail("requires f32 result element type"); + FailureOr sourceArity = getVMIPhysicalArity(sourceType); + FailureOr resultArity = getVMIPhysicalArity(resultType); + if (failed(sourceArity) || failed(resultArity) || + *sourceArity != *resultArity) + return fail("requires matching computable physical arity"); + return success(); +} + +LogicalResult checkSupportedBitcastShape(VMIBitcastOp op, std::string *reason) { + VMILayoutSupport supports; + if (failed(supports.getBitcastSupport(op, reason))) + return failure(); + return success(); +} + +LogicalResult +checkSupportedChannelSplitShape(VMIChannelSplitOp op, + std::string *reason = nullptr) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + int64_t channels = op.getNumResults(); + if (channels != 2 && channels != 4) + return fail("pto.vmi.channel_split supports only 2 or 4 channels"); + + auto sourceType = cast(op.getSource().getType()); + VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); + if (!sourceLayout) + return fail("requires assigned source layout"); + auto expectedLayout = + VMILayoutAttr::getDeinterleaved(op.getContext(), channels); + if (!sourceLayout.isContiguous() && sourceLayout != expectedLayout) + return fail("requires source layout to be contiguous or matching " + "deinterleaved channel layout"); + + for (Value result : op.getResults()) { + VMILayoutAttr resultLayout = + cast(result.getType()).getLayoutAttr(); + if (!resultLayout || !resultLayout.isContiguous()) + return fail("requires every result layout to be contiguous"); + } + + FailureOr sourceArity = getVMIPhysicalArity(sourceType); + int64_t resultArity = 0; + for (Value result : op.getResults()) { + FailureOr arity = + getVMIPhysicalArity(cast(result.getType())); + if (failed(arity)) + return fail("requires computable result physical arity"); + resultArity += *arity; + } + if (failed(sourceArity)) + return fail("requires computable source physical arity"); + if (*sourceArity != resultArity) + return fail("requires source and result to have the same physical arity"); + + return success(); +} + +LogicalResult +checkSupportedChannelMergeShape(VMIChannelMergeOp op, + std::string *reason = nullptr) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + int64_t channels = op.getInputs().size(); + if (channels != 2 && channels != 4) + return fail("pto.vmi.channel_merge supports only 2 or 4 channels"); + + int64_t inputArity = 0; + for (Value input : op.getInputs()) { + auto inputType = cast(input.getType()); + VMILayoutAttr inputLayout = inputType.getLayoutAttr(); + if (!inputLayout || !inputLayout.isContiguous()) + return fail("requires every input layout to be contiguous"); + FailureOr arity = getVMIPhysicalArity(inputType); + if (failed(arity)) + return fail("requires computable input physical arity"); + inputArity += *arity; + } + + auto resultType = cast(op.getResult().getType()); + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + if (!resultLayout) + return fail("requires assigned result layout"); + auto expectedLayout = + VMILayoutAttr::getDeinterleaved(op.getContext(), channels); + if (!resultLayout.isContiguous() && resultLayout != expectedLayout) + return fail("requires result layout to be contiguous or matching " + "deinterleaved channel layout"); + + FailureOr resultArity = getVMIPhysicalArity(resultType); + if (failed(resultArity)) + return fail("requires computable result physical arity"); + if (*resultArity != inputArity) + return fail("requires source and result to have the same physical arity"); + + return success(); +} + +LogicalResult +checkSupportedActivePrefixIndexShape(VMIActivePrefixIndexOp op, + std::string *reason = nullptr) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + auto maskType = cast(op.getMask().getType()); + auto resultType = cast(op.getResult().getType()); + VMILayoutAttr maskLayout = maskType.getLayoutAttr(); + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + if (!maskLayout || !resultLayout) + return fail("requires assigned mask and result layouts"); + if (!maskLayout.isContiguous() || !resultLayout.isContiguous()) + return fail("requires contiguous mask and result layouts"); + + std::string resultFullReason; + if (failed(checkFullDataPhysicalChunks(resultType, &resultFullReason))) + return fail(Twine("requires full result physical chunks so padding mask " + "lanes cannot affect the observable prefix; ") + + resultFullReason); + + std::string maskFullReason; + if (failed(checkFullVMIPhysicalChunks(maskType, &maskFullReason))) + return fail(Twine("requires full mask physical chunks so padding mask " + "lanes cannot affect the observable prefix; ") + + maskFullReason); + + FailureOr maskArity = getVMIPhysicalArity(maskType); + FailureOr resultArity = getVMIPhysicalArity(resultType); + if (failed(maskArity) || failed(resultArity)) + return fail("requires computable mask and result physical arity"); + if (*maskArity != 1 || *resultArity != 1) + return fail("requires a single physical chunk; multi-chunk prefix needs " + "cross-chunk carry"); + + return success(); +} + +LogicalResult checkSupportedCompressShape(VMICompressOp op, + std::string *reason = nullptr) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + auto sourceType = cast(op.getSource().getType()); + auto maskType = cast(op.getMask().getType()); + auto resultType = cast(op.getResult().getType()); + VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); + VMILayoutAttr maskLayout = maskType.getLayoutAttr(); + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + if (!sourceLayout || !maskLayout || !resultLayout) + return fail("requires assigned source, mask, and result layouts"); + if (!sourceLayout.isContiguous() || !maskLayout.isContiguous() || + !resultLayout.isContiguous()) + return fail("requires contiguous source, mask, and result layouts"); + + std::string fullChunkReason; + if (failed(checkFullDataPhysicalChunks(sourceType, &fullChunkReason))) + return fail(Twine("requires full source physical chunks so padding mask " + "lanes cannot be squeezed into the result; ") + + fullChunkReason); + + FailureOr sourceArity = getVMIPhysicalArity(sourceType); + FailureOr maskArity = getVMIPhysicalArity(maskType); + FailureOr resultArity = getVMIPhysicalArity(resultType); + if (failed(sourceArity) || failed(maskArity) || failed(resultArity)) + return fail("requires computable source, mask, and result physical arity"); + if (*sourceArity != 1 || *maskArity != 1 || *resultArity != 1) + return fail("requires a single physical chunk; multi-chunk compress needs " + "cross-chunk compaction"); + + return success(); +} + +LogicalResult checkSupportedCompressStoreShape( + VMICompressStoreOp op, + std::string *reason = nullptr) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + auto valueType = cast(op.getValue().getType()); + auto maskType = cast(op.getMask().getType()); + VMILayoutAttr valueLayout = valueType.getLayoutAttr(); + VMILayoutAttr maskLayout = maskType.getLayoutAttr(); + if (!valueLayout || !maskLayout) + return fail("requires assigned value and mask layouts"); + if (!valueLayout.isContiguous() || !maskLayout.isContiguous()) + return fail("requires contiguous value and mask layouts"); + + if (!isa(op.getDestination().getType())) + return fail("requires !pto.ptr destination because pto.vstur is " + "pointer-only"); + + std::string fullChunkReason; + if (failed(checkFullDataPhysicalChunks(valueType, &fullChunkReason))) + return fail(Twine("requires full physical chunks so padding mask lanes " + "cannot be squeezed into memory; ") + + fullChunkReason); + + FailureOr valueArity = getVMIPhysicalArity(valueType); + FailureOr maskArity = getVMIPhysicalArity(maskType); + if (failed(valueArity) || failed(maskArity)) + return fail("requires computable value and mask physical arity"); + if (*valueArity != 1 || *maskArity != 1) + return fail("requires a single physical chunk; multi-chunk " + "compress_store needs cross-chunk compaction and SQZN " + "state planning"); + + return success(); +} + +template +LogicalResult +checkSupportedReduceShape(OpTy op, bool requiresReassoc, + std::string *reason = nullptr) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + if (requiresReassoc && !op->hasAttr("reassoc")) + return fail("requires reassoc attr for pair-wise floating-point vcadd"); + + auto sourceType = cast(op.getSource().getType()); + auto initType = cast(op.getInit().getType()); + auto maskType = cast(op.getMask().getType()); + auto resultType = cast(op.getResult().getType()); + VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); + VMILayoutAttr initLayout = initType.getLayoutAttr(); + VMILayoutAttr maskLayout = maskType.getLayoutAttr(); + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + if (!sourceLayout || !initLayout || !maskLayout || !resultLayout) + return fail("requires assigned source, init, mask, and result layouts"); + if (!sourceLayout.isContiguous() || !initLayout.isContiguous() || + !maskLayout.isContiguous() || !resultLayout.isContiguous()) + return fail("requires contiguous source, init, mask, and result layouts"); + + std::string fullChunkReason; + if (failed(checkFullDataPhysicalChunks(sourceType, &fullChunkReason))) + return fail(Twine("requires full source physical chunks so padding lanes " + "do not participate in the reduction; ") + + fullChunkReason); + + FailureOr sourceArity = getVMIPhysicalArity(sourceType); + FailureOr initArity = getVMIPhysicalArity(initType); + FailureOr maskArity = getVMIPhysicalArity(maskType); + FailureOr resultArity = getVMIPhysicalArity(resultType); + if (failed(sourceArity) || failed(initArity) || failed(maskArity) || + failed(resultArity)) + return fail("requires computable physical arity"); + if (*sourceArity < 1 || *maskArity != *sourceArity) + return fail("requires source and mask physical arity to match and be " + "non-empty"); + if (*initArity != 1 || *resultArity != 1) + return fail("requires one init and result physical chunk"); + + return success(); +} + +template +LogicalResult +checkSupportedGroupReduceShape(OpTy op, std::string *reason = nullptr) { + VMILayoutSupport supports; + if constexpr (std::is_same_v) { + if (succeeded(supports.getGroupReduceAddFSupport(op, reason))) + return success(); + } else if constexpr (std::is_same_v) { + if (succeeded(supports.getGroupReduceMaxFSupport(op, reason))) + return success(); + } else if constexpr (std::is_same_v) { + if (succeeded(supports.getGroupReduceMaxISupport(op, reason))) + return success(); + } else if constexpr (std::is_same_v) { + if (succeeded(supports.getGroupReduceMinFSupport(op, reason))) + return success(); + } else if constexpr (std::is_same_v) { + if (succeeded(supports.getGroupReduceMinISupport(op, reason))) + return success(); + } else { + if (succeeded(supports.getGroupReduceAddISupport(op, reason))) + return success(); + } + return failure(); +} + +LogicalResult checkSupportedGroupBroadcastShape( + VMIGroupBroadcastOp op, + std::string *reason = nullptr) { + auto sourceType = cast(op.getSource().getType()); + auto resultType = cast(op.getResult().getType()); + if (sourceType.getElementType() != resultType.getElementType()) { + if (reason) + *reason = "requires source/result element type to match"; + return failure(); + } + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + VMILayoutAttr sourceLayout = sourceType.getLayoutAttr(); + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + if (!sourceLayout || !resultLayout) + return fail("requires assigned source/result layouts"); + int64_t numGroups = op.getNumGroupsAttr().getInt(); + if (numGroups <= 0) + return fail("requires positive num_groups"); + if (sourceType.getElementCount() != numGroups) + return fail("requires source lane count to match num_groups"); + if (resultType.getElementCount() % numGroups != 0) + return fail("requires num_groups to evenly divide result lane count"); + if (!sourceLayout.isGroupSlots() || sourceLayout.getNumGroups() != numGroups) + return fail("requires matching num_groups source layout"); + if (resultLayout.isGroupSlots()) + return fail("requires dense result layout"); + + if (sourceLayout.getSlots() > 0 && sourceLayout.getSlots() != 8 && + sourceLayout.getSlots() != 1) + return fail("supports only slots=8 or slots=1 group_broadcast source " + "layouts"); + if (sourceLayout.getSlots() > 1 && numGroups % sourceLayout.getSlots() != 0) + return fail("requires full source group-slot chunks"); + + VMILayoutSupport supports; + std::string supportReason; + if (failed(supports.getGroupBroadcastSupport(op, &supportReason))) + return fail(supportReason); + + FailureOr lanesPerPart = + getDataLanesPerPart(sourceType.getElementType()); + FailureOr resultLanesPerPart = + getDataLanesPerPart(resultType.getElementType()); + if (failed(lanesPerPart) || failed(resultLanesPerPart) || + *lanesPerPart != *resultLanesPerPart) + return fail("requires matching physical lanes per part"); + FailureOr groupSize = getGroupSizeFromNumGroups( + resultType, numGroups, reason); + if (failed(groupSize)) + return failure(); + if (*lanesPerPart % *groupSize != 0 && *groupSize % *lanesPerPart != 0) + return fail("requires derived group size to divide or be a multiple of " + "physical lanes per part"); + + FailureOr resultFactor = getDataLayoutFactor(resultType); + if (failed(resultFactor)) + return fail("requires known result layout factor"); + bool laneStridedDense = + resultLayout.isDense() && resultLayout.getLaneStride() > 1; + if (!laneStridedDense) { + std::string fullChunkReason; + if (failed(checkFullDataPhysicalChunks(resultType, &fullChunkReason))) + return fail(Twine("requires full result physical chunks; ") + + fullChunkReason); + } + if (*resultFactor == 1) + return success(); + bool blockFragmentSmallGroup = + resultLayout.isDeinterleaved() && resultLayout.getBlockElems() > 1 && + *groupSize < *lanesPerPart && + *lanesPerPart % resultLayout.getBlockElems() == 0; + bool deinterleavedSmallGroup = + resultLayout.isDeinterleaved() && resultLayout.getBlockElems() == 1 && + *groupSize < *lanesPerPart && *groupSize >= *resultFactor && + *groupSize % *resultFactor == 0 && + *lanesPerPart % (*groupSize / *resultFactor) == 0; + if (blockFragmentSmallGroup || deinterleavedSmallGroup) + return success(); + int64_t logicalSpanPerResultChunk = *lanesPerPart * *resultFactor; + if (*groupSize < *lanesPerPart || *groupSize % logicalSpanPerResultChunk != 0) + return fail("deinterleaved result requires every physical result chunk to " + "stay within one logical group"); + return success(); +} + +LogicalResult checkSupportedVdhistShape(VMIVdhistOp op, + std::string *reason = nullptr) { + VMILayoutSupport supports; + if (succeeded(supports.getVdhistSupport(op, reason))) + return success(); + return failure(); +} + +LogicalResult checkSupportedVchistShape(VMIVchistOp op, + std::string *reason = nullptr) { + VMILayoutSupport supports; + if (succeeded(supports.getVchistSupport(op, reason))) + return success(); + return failure(); +} + +LogicalResult checkSupportedVmullShape(VMIVmullOp op, + std::string *reason = nullptr) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + auto aType = cast(op.getA().getType()); + auto bType = cast(op.getB().getType()); + auto lowType = cast(op.getLow().getType()); + auto highType = cast(op.getHigh().getType()); + auto maskType = cast(op.getMask().getType()); + + auto elementType = dyn_cast(aType.getElementType()); + if (!elementType || elementType.getWidth() != 32 || + (!elementType.isSignless() && !elementType.isUnsigned())) + return fail("requires element type to be exactly i32 or ui32"); + if (aType != bType || aType != lowType || aType != highType) + return fail("requires identical a, b, low, and high VMI vreg types"); + + int64_t lanes = aType.getElementCount(); + if (lanes != 64 && lanes != 128 && lanes != 256) + return fail("requires logical lane count 64, 128, or 256"); + + VMILayoutAttr layout = aType.getLayoutAttr(); + if (!layout) + return fail("requires an assigned data layout"); + bool supportedLayout = + layout.getLaneStride() == 1 && + (layout.isContiguous() || + (layout.isDeinterleaved() && layout.getBlockElems() == 1 && + (layout.getFactor() == 2 || layout.getFactor() == 4))); + if (!supportedLayout) + return fail("requires contiguous layout or deinterleaved factor 2/4 with " + "block_elems=1 and lane_stride=1"); + if (maskType.getLayoutAttr() != layout) + return fail("requires the mask and all four data values to share one " + "layout"); + if (maskType.getGranularity() != "b32") + return fail("requires b32 mask granularity"); + + FailureOr aArity = getVMIPhysicalArity(aType); + FailureOr bArity = getVMIPhysicalArity(bType); + FailureOr lowArity = getVMIPhysicalArity(lowType); + FailureOr highArity = getVMIPhysicalArity(highType); + FailureOr maskArity = getVMIPhysicalArity(maskType); + if (failed(aArity) || failed(bArity) || failed(lowArity) || + failed(highArity) || failed(maskArity) || *aArity < 1) + return fail("requires computable non-empty physical arity on every port"); + if (*aArity != *bArity || *aArity != *lowArity || *aArity != *highArity || + *aArity != *maskArity) + return fail("requires matching physical arity on a, b, mask, low, and " + "high"); + + FailureOr lanesPerPart = getDataLanesPerPart(aType.getElementType()); + FailureOr physicalElementType = getVMIVRegPhysicalElementType(aType); + FailureOr physicalMaskGranularity = + getVMIMaskPhysicalGranularity(maskType); + if (failed(lanesPerPart) || *lanesPerPart != 64 || + failed(physicalElementType) || + *physicalElementType != aType.getElementType() || + failed(physicalMaskGranularity) || *physicalMaskGranularity != "b32") + return fail("requires 64xi32/ui32 data parts with corresponding b32 mask " + "parts"); + + return success(); +} + +LogicalResult +checkSupportedFmaShape(VMIFmaOp op, std::string *reason = nullptr) { + auto fail = [&](const Twine &message) -> LogicalResult { + if (reason) + *reason = message.str(); + return failure(); + }; + + auto lhsType = cast(op.getLhs().getType()); + FailureOr arity = getVMIPhysicalArity(lhsType); + if (failed(arity) || *arity < 1) + return fail("requires computable non-empty physical arity"); + + return success(); +} + +LogicalResult +checkSupportedReluShape(VMIReluOp op, std::string *reason = nullptr) { + auto resultType = cast(op.getResult().getType()); + if (failed(checkSupportedMaskableVReg(resultType, reason))) + return failure(); + + return success(); +} + +void emitEnsureLayoutMaterializationError(VMIEnsureLayoutOp ensure, + VMIVRegType sourceType, + VMIVRegType resultType, + StringRef reason) { + if (ensure.getResult().hasOneUse()) { + OpOperand &use = *ensure.getResult().use_begin(); + Operation *requester = use.getOwner(); + InFlightDiagnostic diag = + requester->emitError() + << kVMIDiagUnsupportedPrefix << requester->getName() << " operand #" + << use.getOperandNumber() << " has type " << sourceType + << " but requires " << resultType + << "; pto.vmi.ensure_layout cannot materialize this conversion"; + diag.attachNote(ensure.getLoc()) + << "failed helper conversion " << sourceType << " -> " << resultType + << " (" << reason + << "); partial/tail layout materialization requires an explicit " + "packing plan"; + return; + } + + ensure.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.ensure_layout cannot materialize the requested data " + "layout conversion (" + << reason + << "); partial/tail layout materialization requires an explicit " + "packing plan"; +} + +LogicalResult +verifySupportedVMIToVPTOOps(ModuleOp module, + bool enableStableGatherMaskedLoad) { + auto emitMemoryUnsupported = + [&](Operation *op, StringRef opName, VMIVRegType type, Value source, + std::optional constantOffset) -> WalkResult { + std::string reason; + if (succeeded(checkSupportedLoadShape(type, source, + source.getType(), constantOffset, + &reason))) + return WalkResult::advance(); + + op->emitError() << kVMIDiagUnsupportedPrefix << opName + << " direct lowering requires a supported memory source (" + << reason << ")"; + return WalkResult::interrupt(); + }; + + auto emitMaskableUnsupported = [&](Operation *op, StringRef opName, + VMIVRegType type) -> WalkResult { + std::string reason; + if (succeeded(checkSupportedMaskableVReg(type, &reason))) + return WalkResult::advance(); + + op->emitError() + << kVMIDiagUnsupportedPrefix << opName + << " direct lowering requires physical vreg parts with b8/b16/b32 " + "predicate masks (" + << reason << ")"; + return WalkResult::interrupt(); + }; + + WalkResult result = module.walk([&](Operation *op) { + if (auto constant = dyn_cast(op)) { + auto denseAttr = dyn_cast(constant.getValue()); + if (!denseAttr || !denseAttr.isSplat()) { + constant.emitError() + << kVMIDiagUnsupportedPrefix + << "non-splat pto.vmi.constant requires a vreg immediate or " + "scratch materialization plan"; + return WalkResult::interrupt(); + } + return emitMaskableUnsupported( + op, "pto.vmi.constant", + cast(constant.getResult().getType())); + } + + if (auto broadcast = dyn_cast(op)) + return emitMaskableUnsupported( + op, "pto.vmi.broadcast", + cast(broadcast.getResult().getType())); + if (auto broadcast = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedGroupBroadcastShape(broadcast, + &reason))) + return WalkResult::advance(); + broadcast.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.group_broadcast requires full source chunks with " + "#pto.vmi.layout, a dense full result " + "layout, " + "and num_groups deriving a group size that divides or is a " + "multiple of physical chunk lanes (" + << reason << ")"; + return WalkResult::interrupt(); + } + if (auto hist = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedVdhistShape(hist, &reason))) + return WalkResult::advance(); + hist.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.vdhist requires contiguous Nxui8 source, contiguous b8 " + "mask, and contiguous 256xui16 acc/result (" + << reason << ")"; + return WalkResult::interrupt(); + } + if (auto hist = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedVchistShape(hist, &reason))) + return WalkResult::advance(); + hist.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.vchist requires contiguous Nxui8 source, contiguous b8 " + "mask, and contiguous 256xui16 acc/result (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto load = dyn_cast(op)) { + return emitMemoryUnsupported( + op, "pto.vmi.load", cast(load.getResult().getType()), + load.getSource(), getConstantIndexValue(load.getOffset())); + } + if (auto load = dyn_cast(op)) { + std::string reason; + if (succeeded( + checkSupportedDeinterleaveLoadShape(load, &reason))) + return WalkResult::advance(); + load.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.deinterleave_load lowers through pto.vldsx2 only for " + "matching contiguous full low/high result chunks with a supported " + "UB source and 8/16/32-bit element type (" + << reason << ")"; + return WalkResult::interrupt(); + } + if (auto load = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedStrideLoadShape(load, &reason))) + return WalkResult::advance(); + load.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.stride_load lowers through pto.vsldb only for one " + "contiguous physical result/mask chunk and a supported UB source (" + << reason << ")"; + return WalkResult::interrupt(); + } + if (auto load = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedGroupLoadShape(load, &reason))) + return WalkResult::advance(); + load.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.group_load requires contiguous full result chunks, a " + "supported UB source, and num_groups deriving a group size " + "aligned to physical chunks (" + << reason << ")"; + return WalkResult::interrupt(); + } + if (auto load = dyn_cast(op)) { + std::string reason; + if (succeeded( + checkSupportedGroupSlotLoadShape(load, &reason))) + return WalkResult::advance(); + load.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.group_slot_load requires explicit group_slots result " + "layout matching num_groups, a supported UB pointer source, " + "and either slots=8 with constant unit source_group_stride or " + "slots=1 row-local lowering (" + << reason << ")"; + return WalkResult::interrupt(); + } + if (auto load = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedGroupBroadcastLoadShape(load, + &reason))) + return WalkResult::advance(); + load.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.group_broadcast_load requires either the BRC full-group " + "chunk form, the E2B packet form for b16/b32 direct or split " + "group size, or the generic group-slot-load then group-broadcast " + "fallback with supported UB pointer source and source_group_stride (" + << reason << ")"; + return WalkResult::interrupt(); + } + if (auto load = dyn_cast(op)) { + if (enableStableGatherMaskedLoad) { + load.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.masked_load stable VGATHER-based lowering is reserved " + "for strict masked/tail loads but is not implemented yet"; + return WalkResult::interrupt(); + } + std::string reason; + if (succeeded(checkSupportedMaskedLoadShape(load, &reason))) + return WalkResult::advance(); + load.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.masked_load direct lowering requires a supported memory " + "source, contiguous result/passthru/mask layouts, and either " + "full physical chunks or a statically safe full-read footprint (" + << reason << ")"; + return WalkResult::interrupt(); + } + if (auto gather = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedGatherShape(gather, &reason))) + return WalkResult::advance(); + gather.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.gather lowers through pto.vgather2_bc + pto.vsel only " + "for UB pointer sources, contiguous full physical chunks, " + "32-bit result elements, i32 indices, and b32 masks (" + << reason << ")"; + return WalkResult::interrupt(); + } + if (auto load = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedExpandLoadShape(load, &reason))) + return WalkResult::advance(); + load.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.expand_load direct lowering is currently supported for " + "either a static all-active mask lowered as pto.vlds, or a " + "one-full-chunk 32-bit UB runtime mask lowered through pto.vusqz " + "+ pto.vgather2_bc + pto.vsel (" + << reason << ")"; + return WalkResult::interrupt(); + } + if (auto store = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedStoreShape(cast(store.getValue().getType()), + store.getDestination(), store.getDestination().getType(), + &reason))) + return WalkResult::advance(); + store.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.store requires an 8/16/32-bit predicate-maskable " + "element type and either full physical chunks or contiguous " + "tail-store layout, with UB-backed destination (" + << reason << ")"; + return WalkResult::interrupt(); + } + if (auto store = dyn_cast(op)) { + std::string reason; + if (succeeded( + checkSupportedInterleaveStoreShape(store, &reason))) + return WalkResult::advance(); + store.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.interleave_store lowers through pto.vstsx2 only for " + "matching contiguous full low/high input chunks with a supported " + "UB destination and 8/16/32-bit element type (" + << reason << ")"; + return WalkResult::interrupt(); + } + if (auto store = dyn_cast(op)) { + std::string reason; + if (succeeded( + checkSupportedGroupStoreShape(store, &reason))) + return WalkResult::advance(); + store.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.group_store requires contiguous full value chunks, a " + "supported UB destination, and num_groups deriving a group size " + "aligned to physical chunks (" + << reason << ")"; + return WalkResult::interrupt(); + } + if (auto store = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedMaskedStoreShape(cast(store.getValue().getType()), + cast(store.getMask().getType()), + store.getDestination(), store.getDestination().getType(), + &reason))) + return WalkResult::advance(); + store.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.masked_store requires either full physical chunks or " + "contiguous tail-store value/mask layout, with UB-backed " + "destination (" + << reason << ")"; + return WalkResult::interrupt(); + } + if (auto store = dyn_cast(op)) { + std::string reason; + if (succeeded( + checkSupportedStrideStoreShape(store, &reason))) + return WalkResult::advance(); + store.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.stride_store lowers through pto.vsstb only for one " + "contiguous physical value/mask chunk and a supported UB " + "destination (" + << reason << ")"; + return WalkResult::interrupt(); + } + if (auto scatter = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedScatterShape(scatter, &reason))) + return WalkResult::advance(); + scatter.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.scatter lowers through pto.vscatter only with a UB " + "pointer destination, contiguous full physical chunks, 32-bit " + "value elements, i32 indices, and b32 masks (" + << reason << ")"; + return WalkResult::interrupt(); + } + if (auto ensure = dyn_cast(op)) { + auto sourceType = cast(ensure.getSource().getType()); + auto resultType = cast(ensure.getResult().getType()); + std::string reason; + VMILayoutSupport supports; + if (succeeded( + supports.getEnsureLayoutFact(sourceType, resultType, &reason))) + return WalkResult::advance(); + + emitEnsureLayoutMaterializationError(ensure, sourceType, resultType, + reason); + return WalkResult::interrupt(); + } + + if (auto ensure = dyn_cast(op)) { + auto sourceType = cast(ensure.getSource().getType()); + auto resultType = cast(ensure.getResult().getType()); + std::string reason; + VMILayoutSupport supports; + if (succeeded(supports.getEnsureMaskLayoutFact(sourceType, resultType, + &reason))) + return WalkResult::advance(); + + ensure.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.ensure_mask_layout cannot materialize the requested " + "mask layout conversion (" + << reason + << "); partial/tail predicate layout materialization requires an " + "explicit packing plan"; + return WalkResult::interrupt(); + } + + if (auto ensure = dyn_cast(op)) { + auto sourceType = cast(ensure.getSource().getType()); + auto resultType = cast(ensure.getResult().getType()); + bool identity = + sourceType.getGranularity() == resultType.getGranularity() && + sourceType.getLayoutAttr() == resultType.getLayoutAttr(); + if (!identity) { + VMILayoutSupport supports; + std::string reason; + if (failed(supports.getMaskGranularityCastLayoutFactForLayouts( + sourceType, resultType, sourceType.getLayoutAttr(), + resultType.getLayoutAttr(), &reason))) { + ensure.emitError() + << kVMIDiagUnsupportedPrefix + << "mask granularity cast layout relation is unsupported (" + << reason << ")"; + return WalkResult::interrupt(); + } + } + + return WalkResult::advance(); + } + + if (auto addf = dyn_cast(op)) + return emitMaskableUnsupported( + op, "pto.vmi.addf", cast(addf.getResult().getType())); + if (auto addi = dyn_cast(op)) + return emitMaskableUnsupported( + op, "pto.vmi.addi", cast(addi.getResult().getType())); + if (auto subf = dyn_cast(op)) + return emitMaskableUnsupported( + op, "pto.vmi.subf", cast(subf.getResult().getType())); + if (auto subi = dyn_cast(op)) + return emitMaskableUnsupported( + op, "pto.vmi.subi", cast(subi.getResult().getType())); + if (auto mulf = dyn_cast(op)) + return emitMaskableUnsupported( + op, "pto.vmi.mulf", cast(mulf.getResult().getType())); + if (auto muli = dyn_cast(op)) + return emitMaskableUnsupported( + op, "pto.vmi.muli", cast(muli.getResult().getType())); + if (auto vmull = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedVmullShape(vmull, &reason))) + return WalkResult::advance(); + vmull.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.vmull requires equal 64/128/256-lane i32/ui32 data " + "ports, a matching b32 mask, and contiguous or deinterleaved " + "factor-2/factor-4 block_elems=1 lane_stride=1 layout (" + << reason << ")"; + return WalkResult::interrupt(); + } + if (auto divf = dyn_cast(op)) + return emitMaskableUnsupported( + op, "pto.vmi.divf", cast(divf.getResult().getType())); + if (auto minf = dyn_cast(op)) + return emitMaskableUnsupported( + op, "pto.vmi.minf", cast(minf.getResult().getType())); + if (auto maxf = dyn_cast(op)) + return emitMaskableUnsupported( + op, "pto.vmi.maxf", cast(maxf.getResult().getType())); + if (auto negf = dyn_cast(op)) + return emitMaskableUnsupported( + op, "pto.vmi.negf", cast(negf.getResult().getType())); + if (auto absf = dyn_cast(op)) + return emitMaskableUnsupported( + op, "pto.vmi.absf", cast(absf.getResult().getType())); + if (auto absi = dyn_cast(op)) + return emitMaskableUnsupported( + op, "pto.vmi.absi", cast(absi.getResult().getType())); + if (auto sqrt = dyn_cast(op)) + return emitMaskableUnsupported( + op, "pto.vmi.sqrt", cast(sqrt.getResult().getType())); + if (auto exp = dyn_cast(op)) + return emitMaskableUnsupported( + op, "pto.vmi.exp", cast(exp.getResult().getType())); + if (auto ln = dyn_cast(op)) + return emitMaskableUnsupported( + op, "pto.vmi.ln", cast(ln.getResult().getType())); + if (auto relu = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedReluShape(relu, &reason))) + return WalkResult::advance(); + relu.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.relu direct lowering requires physical vreg parts with " + "b8/b16/b32 predicate masks and f16/f32 element type (" + << reason << ")"; + return WalkResult::interrupt(); + } + if (auto andi = dyn_cast(op)) + return emitMaskableUnsupported( + op, "pto.vmi.andi", cast(andi.getResult().getType())); + if (auto ori = dyn_cast(op)) + return emitMaskableUnsupported( + op, "pto.vmi.ori", cast(ori.getResult().getType())); + if (auto xori = dyn_cast(op)) + return emitMaskableUnsupported( + op, "pto.vmi.xori", cast(xori.getResult().getType())); + if (auto shli = dyn_cast(op)) + return emitMaskableUnsupported( + op, "pto.vmi.shli", cast(shli.getResult().getType())); + if (auto shrui = dyn_cast(op)) + return emitMaskableUnsupported( + op, "pto.vmi.shrui", cast(shrui.getResult().getType())); + if (auto shrsi = dyn_cast(op)) + return emitMaskableUnsupported( + op, "pto.vmi.shrsi", cast(shrsi.getResult().getType())); + if (auto notOp = dyn_cast(op)) + return emitMaskableUnsupported( + op, "pto.vmi.not", cast(notOp.getResult().getType())); + if (auto select = dyn_cast(op)) + return emitMaskableUnsupported( + op, "pto.vmi.select", + cast(select.getResult().getType())); + + if (auto cmpf = dyn_cast(op)) { + WalkResult physical = emitMaskableUnsupported( + op, "pto.vmi.cmpf", cast(cmpf.getLhs().getType())); + if (physical.wasInterrupted()) + return physical; + if (succeeded(checkSupportedComparePredicate( + op, cmpf.getPredicate()))) + return WalkResult::advance(); + return WalkResult::interrupt(); + } + + if (auto cmpi = dyn_cast(op)) { + WalkResult physical = emitMaskableUnsupported( + op, "pto.vmi.cmpi", cast(cmpi.getLhs().getType())); + if (physical.wasInterrupted()) + return physical; + if (succeeded(checkSupportedComparePredicate( + op, cmpi.getPredicate()))) + return WalkResult::advance(); + return WalkResult::interrupt(); + } + + if (auto activePrefix = dyn_cast(op)) { + std::string reason; + if (succeeded( + checkSupportedActivePrefixIndexShape(activePrefix, &reason))) + return WalkResult::advance(); + activePrefix.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.active_prefix_index lowers through pto.vusqz only for " + "one contiguous physical chunk (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto compress = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedCompressShape(compress, &reason))) + return WalkResult::advance(); + compress.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.compress lowers through pto.vsqz only for one " + "contiguous full physical chunk (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto compressStore = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedCompressStoreShape(compressStore, &reason))) + return WalkResult::advance(); + compressStore.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.compress_store lowers through pto.vsqz + pto.vstur " + "only for one contiguous full physical chunk with a UB pointer " + "destination (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto reduce = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedReduceShape( + reduce, /*requiresReassoc=*/false, &reason))) + return WalkResult::advance(); + reduce.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.reduce_addi lowers through pto.vcadd only for " + "contiguous full 32-bit integer source chunks with matching " + "mask chunks and one init/result chunk (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto reduce = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedReduceShape( + reduce, /*requiresReassoc=*/true, &reason))) + return WalkResult::advance(); + reduce.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.reduce_addf lowers through pto.vcadd only with " + "reassoc, f32 contiguous full source chunks, matching mask " + "chunks, and one init/result chunk (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto reduce = dyn_cast(op)) { + std::string reason; + if (succeeded( + checkSupportedGroupReduceShape(reduce, &reason))) + return WalkResult::advance(); + reduce.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.group_reduce_addf lowers through pto.vcgadd for 32B " + "blocks or through pto.vcadd for contiguous full " + "source/mask chunks, #pto.vmi.layout " + "result " + "chunks, and num_groups deriving a group size aligned to " + "physical chunks (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto reduce = dyn_cast(op)) { + std::string reason; + if (succeeded( + checkSupportedGroupReduceShape(reduce, &reason))) + return WalkResult::advance(); + reduce.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.group_reduce_addi lowers through pto.vcgadd/vadd for " + "supported 32B block classes or through an internal widening " + "pto.vcadd path for aligned full chunks (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto reduce = dyn_cast(op)) { + std::string reason; + if (succeeded( + checkSupportedGroupReduceShape(reduce, &reason))) + return WalkResult::advance(); + reduce.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.group_reduce_maxi lowers through pto.vcgmax/vmax for " + "supported 32B block classes or through pto.vcmax for aligned " + "full chunks (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto reduce = dyn_cast(op)) { + std::string reason; + if (succeeded( + checkSupportedGroupReduceShape(reduce, &reason))) + return WalkResult::advance(); + reduce.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.group_reduce_maxf lowers through pto.vcgmax/vmax for " + "32B blocks or through pto.vcmax for contiguous full chunks, " + "matching source/mask chunks, " + "#pto.vmi.layout result chunks, and " + "num_groups deriving a group size aligned to physical chunks (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto reduce = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedGroupReduceShape(reduce, &reason))) + return WalkResult::advance(); + reduce.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.group_reduce_minf lowers through pto.vcgmin/vmin for " + "supported 32B block classes or through pto.vcmin for aligned " + "full chunks (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto reduce = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedGroupReduceShape(reduce, &reason))) + return WalkResult::advance(); + reduce.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.group_reduce_mini lowers through pto.vcgmin/vmin for " + "supported 32B block classes or through pto.vcmin for aligned " + "full chunks (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto reduce = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedReduceShape( + reduce, /*requiresReassoc=*/false, &reason))) + return WalkResult::advance(); + reduce.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.reduce_maxf lowers through pto.vcmax only for f16/f32 " + "contiguous full source chunks with matching mask chunks and one " + "init/result chunk (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto reduce = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedReduceShape( + reduce, /*requiresReassoc=*/false, &reason))) + return WalkResult::advance(); + reduce.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.reduce_minf lowers through pto.vcmin only for f16/f32 " + "contiguous full source chunks with matching mask chunks and one " + "init/result chunk (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto reduce = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedReduceShape( + reduce, /*requiresReassoc=*/false, &reason))) + return WalkResult::advance(); + reduce.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.reduce_maxi lowers through pto.vcmax only for " + "contiguous full integer source chunks with matching mask " + "chunks and one init/result chunk (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto reduce = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedReduceShape( + reduce, /*requiresReassoc=*/false, &reason))) + return WalkResult::advance(); + reduce.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.reduce_mini lowers through pto.vcmin only for " + "contiguous full integer source chunks with matching mask " + "chunks and one init/result chunk (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto fma = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedFmaShape(fma, &reason))) + return WalkResult::advance(); + fma.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.fma lowers through pto.vmula only for f16/bf16/f32 " + "element types (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto extf = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedExtFShape(extf, &reason))) + return WalkResult::advance(); + + extf.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.extf supports contiguous 16-bit float-like or fp8-like " + "physical source chunks to f32 deinterleaved=2/4 results; " + "partial/tail is allowed only when source padding maps to result " + "padding (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto truncf = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedTruncFShape(truncf, &reason))) + return WalkResult::advance(); + + truncf.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.truncf supports only f32 deinterleaved=2 source parts " + "to dense f16 results, f32 source layouts whose factor times the " + "result lane_stride matches the fp8-like narrowing factor, or f32 " + "group_slots(num_groups=G, slots=1) to f16 " + "group_slots(num_groups=G, slots=1) (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto fptosi = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedFPToSIShape(fptosi, &reason))) + return WalkResult::advance(); + + fptosi.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.fptosi supports f32 source chunks to matching 32-bit " + "integer result chunks with identical assigned layouts (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto sitofp = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedSIToFPShape(sitofp, &reason))) + return WalkResult::advance(); + + sitofp.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.sitofp supports 32-bit integer source chunks to " + "matching f32 result chunks with identical assigned layouts (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto extsi = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedExtSIShape(extsi, &reason))) + return WalkResult::advance(); + + extsi.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.extsi supports contiguous signed/signless 8-bit or " + "16-bit integer physical source chunks to 2x/4x wider integer " + "deinterleaved results, or matching " + "group_slots(num_groups=G, slots=8) 8/16-bit integer source to " + "32-bit integer result (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto extui = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedExtUIShape(extui, &reason))) + return WalkResult::advance(); + + extui.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.extui supports contiguous unsigned 8-bit or 16-bit " + "integer physical source chunks to 2x/4x wider unsigned integer " + "deinterleaved results, or matching " + "group_slots(num_groups=G, slots=8) 8/16-bit integer source to " + "32-bit integer result (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto trunci = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedTruncIShape(trunci, &reason))) + return WalkResult::advance(); + + trunci.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.trunci supports integer deinterleaved source layouts " + "whose factor is the 2x/4x narrowing multiple of the contiguous " + "or deinterleaved result layout factor, or 32-bit integer " + "group_slots(num_groups=G, slots=1 or 8) to 8/16-bit integer " + "group_slots(num_groups=G, slots=1 or 8), or 16-bit unsigned " + "integer group_slots(num_groups=G, slots=8) to 8-bit unsigned " + "integer group_slots(num_groups=G, slots=8, lane_stride=2) (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto bitcast = dyn_cast(op)) { + std::string reason; + if (succeeded(checkSupportedBitcastShape(bitcast, &reason))) + return WalkResult::advance(); + + bitcast.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.bitcast requires matching source/result layouts with " + "width-changing forms restricted to supported layout table rows (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto split = dyn_cast(op)) { + int64_t channels = split.getNumResults(); + std::string reason; + if (succeeded( + checkSupportedChannelSplitShape(split, &reason))) + return WalkResult::advance(); + + if (channels != 2 && channels != 4) + split.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.channel_split supports only 2 or 4 channels"; + else + split.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.channel_split requires source layout to be contiguous " + "or matching deinterleaved channel layout, every result layout " + "to be contiguous, and complete physical channel groups (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto merge = dyn_cast(op)) { + int64_t channels = merge.getInputs().size(); + std::string reason; + if (succeeded( + checkSupportedChannelMergeShape(merge, &reason))) + return WalkResult::advance(); + + if (channels != 2 && channels != 4) + merge.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.channel_merge supports only 2 or 4 channels"; + else + merge.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.channel_merge requires every input layout to be " + "contiguous and result layout to be contiguous or matching " + "deinterleaved channel layout, with complete physical channel " + "groups (" + << reason << ")"; + return WalkResult::interrupt(); + } + + if (auto shuffle = dyn_cast(op)) { + std::string reason; + if (succeeded(computeShuffleForwardingSourceParts(shuffle, &reason))) + return WalkResult::advance(); + std::string splatReason; + if (succeeded(computeShuffleLane0SplatSourcePart(shuffle, &splatReason))) + return WalkResult::advance(); + std::string vselrReason; + if (succeeded(computeShuffleVselrPlans(shuffle, &vselrReason))) + return WalkResult::advance(); + + shuffle.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.shuffle requires physical chunk forwarding or " + "lane0 splat or vci-materializable vselr indices (forwarding: " + << reason << "; lane0 splat: " << splatReason + << "; vselr: " << vselrReason << ")"; + return WalkResult::interrupt(); + } + + if (auto constantMask = dyn_cast(op)) { + std::string reason; + if (succeeded(computeConstantMaskMaterialization(constantMask, &reason))) + return WalkResult::advance(); + + constantMask.emitError() + << kVMIDiagUnsupportedPrefix + << "pto.vmi.constant_mask requires a dense bool constant with " + "concrete layout and b8/b16/b32 granularity (" + << reason << ")"; + return WalkResult::interrupt(); + } + + return WalkResult::advance(); + }); + return failure(result.wasInterrupted()); +} + +struct VMIToVPTOPass : public mlir::pto::impl::VMIToVPTOBase { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(VMIToVPTOPass) + + void runOnOperation() override { + ModuleOp module = getOperation(); + if (failed(verifyVMIToVPTOInputIR(module))) { + signalPassFailure(); + return; + } + if (failed(verifySupportedVMIToVPTOOps(module, + enableStableGatherMaskedLoad))) { + signalPassFailure(); + return; + } + + MLIRContext *context = module.getContext(); + VMIToVPTOTypeConverter typeConverter; + RewritePatternSet patterns(context); + + populateVMIConversionPatterns(typeConverter, patterns); + ConversionTarget target(*context); + target.markUnknownOpDynamicallyLegal([](Operation *op) { + return !isVMIOp(op) && !hasVMIType(op); + }); + if (failed(applyPartialConversion(module, target, std::move(patterns)))) { + module.emitError() << kVMIDiagResidualOpPrefix + << "failed to convert all VMI ops/types to VPTO"; + signalPassFailure(); + return; + } + if (failed(verifyNoResidualVMIIR(module))) { + signalPassFailure(); + } + } +}; + +} // namespace + +std::unique_ptr mlir::pto::createVMIToVPTOPass() { + return std::make_unique(); +} diff --git a/lib/PTO/Transforms/VPTONormalizeEquivalentVcvt.cpp b/lib/PTO/Transforms/VPTONormalizeEquivalentVcvt.cpp new file mode 100644 index 0000000000..22dbbfc094 --- /dev/null +++ b/lib/PTO/Transforms/VPTONormalizeEquivalentVcvt.cpp @@ -0,0 +1,96 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#include "PTO/IR/PTO.h" +#include "PTO/IR/PTOTypeUtils.h" +#include "PTO/Transforms/Passes.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" + +namespace mlir { +namespace pto { +#define GEN_PASS_DEF_VPTONORMALIZEEQUIVALENTVCVT +#include "PTO/Transforms/Passes.h.inc" +} // namespace pto +} // namespace mlir + +using namespace mlir; +using namespace mlir::pto; + +namespace { + +static bool isOddPart(StringRef part) { + return part == "ODD" || part == "PART_ODD"; +} + +static bool isAllTrueMask(Value mask) { + if (auto op = mask.getDefiningOp()) + return op.getPattern() == "PAT_ALL"; + if (auto op = mask.getDefiningOp()) + return op.getPattern() == "PAT_ALL"; + if (auto op = mask.getDefiningOp()) + return op.getPattern() == "PAT_ALL"; + return false; +} + +static bool isPairEquivalentLoadDist(StringRef dist) { + return dist == "BRC_B8" || dist == "BRC_B16" || dist == "BRC_B32" || + dist == "US_B8" || dist == "US_B16" || dist == "E2B_B16" || + dist == "E2B_B32"; +} + +static bool hasEvenOddEquivalentLanes(Value value) { + if (value.getDefiningOp()) + return true; + + auto load = value.getDefiningOp(); + if (!load || value != load.getResult()) + return false; + + std::optional dist = load.getDist(); + return dist && isPairEquivalentLoadDist(*dist); +} + +static bool isNarrowToWideVcvt(VcvtOp op) { + auto inputType = dyn_cast(op.getInput().getType()); + auto resultType = dyn_cast(op.getResult().getType()); + if (!inputType || !resultType) + return false; + + unsigned inputBits = getPTOStorageElemBitWidth(inputType.getElementType()); + unsigned resultBits = getPTOStorageElemBitWidth(resultType.getElementType()); + return inputBits != 0 && resultBits != 0 && inputBits < resultBits; +} + +struct VPTONormalizeEquivalentVcvtPass + : public pto::impl::VPTONormalizeEquivalentVcvtBase< + VPTONormalizeEquivalentVcvtPass> { + void runOnOperation() override { + StringAttr even = StringAttr::get(&getContext(), "EVEN"); + + getOperation().walk([&](VcvtOp op) { + std::optional part = op.getPart(); + if (!part || !isOddPart(*part)) + return; + if (!isNarrowToWideVcvt(op)) + return; + if (!isAllTrueMask(op.getMask())) + return; + if (!hasEvenOddEquivalentLanes(op.getInput())) + return; + + op.setPartAttr(even); + }); + } +}; + +} // namespace + +std::unique_ptr mlir::pto::createVPTONormalizeEquivalentVcvtPass() { + return std::make_unique(); +} diff --git a/lib/TileOps/issue_518_repro.py b/lib/TileOps/issue_518_repro.py new file mode 100644 index 0000000000..e3357d6b41 --- /dev/null +++ b/lib/TileOps/issue_518_repro.py @@ -0,0 +1,85 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""PTODSL VMI implementation of the group_count dhist benchmark.""" + +from ptodsl import pto +from ptodsl._control_flow import vecscope as _vecscope +from ptodsl._runtime import native_build as _native_build + +_native_build._source_ptoas_overrides = ( # noqa: SLF001 + lambda module_spec: { + "backend": getattr(module_spec, "backend", None) or "vpto", + "pto_level": "level3", + } +) + +SRC_ELEMS = 24064 +NBINS = 256 +NUM_GROUPS = 9 +VMI_TILE_ELEMS = 256 +NUM_VMI_TILES = SRC_ELEMS // VMI_TILE_ELEMS + +_SRC_UB = 0 +_OUT_UB = SRC_ELEMS * 4 + + +@pto.jit( + name="group_count_dhist_vmi", + target="a5", + backend="vpto", + mode="explicit", + kernel_kind="vector", + insert_sync=False, + ast_rewrite=False, +) +def group_count_dhist_vmi( + src_gm: pto.ptr(pto.si64, "gm"), + out_gm: pto.ptr(pto.ui32, "gm"), +): + src_ub = pto.castptr(pto.const(_SRC_UB, dtype=pto.ui64), pto.ptr(pto.si64, "ub")) + src32_ub = pto.castptr(pto.const(_SRC_UB, dtype=pto.ui64), pto.ptr(pto.ui32, "ub")) + out_ub = pto.castptr(pto.const(_OUT_UB, dtype=pto.ui64), pto.ptr(pto.ui32, "ub")) + + pto.mte_gm_ub(src_gm, src_ub, 0, SRC_ELEMS * 8, nburst=(1, SRC_ELEMS * 8, SRC_ELEMS * 8)) + pto.set_flag(pto.Pipe.MTE2, pto.Pipe.V, event_id=0) + pto.wait_flag(pto.Pipe.MTE2, pto.Pipe.V, event_id=0) + + with _vecscope(): + hist_source_mask = pto.vmi.create_mask(VMI_TILE_ELEMS, size=VMI_TILE_ELEMS) + hist_mask = pto.vmi.create_mask(NBINS, size=NBINS) + + hist16_zero = pto.vmi.vbrc(pto.ui16(0), result_type=pto.vmi.vreg(256, pto.ui16)) + hist32_init = pto.vmi.vbrc(pto.ui32(0), result_type=pto.vmi.vreg(256, pto.ui32)) + + tile_loop = pto.for_(0, NUM_VMI_TILES, step=1).carry(hist32=hist32_init) + with tile_loop: + idx32, _ = pto.vmi.vload( + src32_ub, + tile_loop.iv * VMI_TILE_ELEMS * 2, + size=VMI_TILE_ELEMS, + dist_mode="dintlv", + ) + valid = pto.vmi.vcmps(idx32, pto.ui32(NUM_GROUPS), hist_source_mask, "lt", result_type=pto.vmi.mask(256)) + source = pto.vmi.vcvt(idx32, pto.ui8, result_type=pto.vmi.vreg(256, pto.ui8)) + hist = pto.vmi.vdhist( + hist16_zero, + source, + valid, + result_type=pto.vmi.vreg(256, pto.ui16), + ) + hist32 = pto.vmi.vcvt(hist, pto.ui32, result_type=pto.vmi.vreg(256, pto.ui32)) + hist32 = pto.vmi.vadd(tile_loop.hist32, hist32, hist_mask, result_type=pto.vmi.vreg(256, pto.ui32)) + tile_loop.update(hist32=hist32) + + pto.vmi.vstore(tile_loop.final("hist32"), out_ub, 0, hist_mask) + + pto.set_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=1) + pto.wait_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=1) + pto.mte_ub_gm(out_ub, out_gm, NBINS * 4, nburst=(1, NBINS * 4, NBINS * 4)) + pto.mem_bar(pto.BarrierType.SS_ALL) diff --git a/ptodsl/docs/developer_guide/tilelib-debugging-playbook.md b/ptodsl/docs/developer_guide/tilelib-debugging-playbook.md index 93734059b7..db519c5165 100644 --- a/ptodsl/docs/developer_guide/tilelib-debugging-playbook.md +++ b/ptodsl/docs/developer_guide/tilelib-debugging-playbook.md @@ -12,7 +12,7 @@ usually live in different layers. Build PTOAS after C++ or TableGen changes: ```bash -ninja -C build-llvm21 tools/ptoas/ptoas +cmake --build build-llvm21 --target ptoas_runtime ``` Stage PTODSL after Python package changes: diff --git a/ptodsl/docs/user_guide/01-introduction.md b/ptodsl/docs/user_guide/01-introduction.md index 7ca942a774..8182c624b2 100644 --- a/ptodsl/docs/user_guide/01-introduction.md +++ b/ptodsl/docs/user_guide/01-introduction.md @@ -318,7 +318,7 @@ Chapter 11 walks through this example in full detail. |---------------|---------------| | New to PTODSL | Chapter 2 (Quick Start), then Chapter 3 (Kernel Entries & Modules) | | Writing your first kernel | Chapter 2 → Chapter 4 (Type System) → Chapter 5 (Control Flow) | -| Looking up a specific operation | Chapters 6–10 and Chapter 13 (organized by topic) | +| Looking up a specific operation | Chapters 6–10, Chapter 13, and Chapter 14 (organized by topic) | | Understanding the flash attention reference | Chapter 11 | **Chapter overview:** @@ -338,3 +338,4 @@ Chapter 11 walks through this example in full detail. | 11 | Flash attention walkthrough | | 12 | Additional examples | | 13 | SIMT micro-ops | +| 14 | The Virtual Micro-instruction Set (VMI): logical vector instructions, masks, and type-directed authoring | diff --git a/ptodsl/docs/user_guide/05-control-flow.md b/ptodsl/docs/user_guide/05-control-flow.md index 4ea132257b..dbac778d33 100644 --- a/ptodsl/docs/user_guide/05-control-flow.md +++ b/ptodsl/docs/user_guide/05-control-flow.md @@ -5,6 +5,7 @@ PTODSL uses a **tracing** compilation model. When you call `kernel.compile(...)` This has one critical implication for how you write loops and branches: - **Python native `for`/`if`** is rewritten to device-side control flow by default in `@pto.jit` bodies and named `@pto.cube` / `@pto.simd` / `@pto.simt` sub-kernels. A `for i in range(rows)` loop records a device loop, and a runtime `if` records both branches. +- **Assign-form Python conditional expressions** such as `x = a if cond else b` are normalized through the same AST rewrite path, so runtime conditions lower to device-side branches before the assignment is merged back into `x`. - **`pto.const_expr` / `pto.static_range`** keep compile-time Python behavior when you want trace-time specialization or unrolling. - **`pto.for_` / `pto.if_`** produce device-side control flow. The loop bound or branch condition can be a runtime value, and the hardware will execute the loop or take the branch dynamically. diff --git a/ptodsl/docs/user_guide/07-data-movement-ops.md b/ptodsl/docs/user_guide/07-data-movement-ops.md index cabe44abdf..324202673e 100644 --- a/ptodsl/docs/user_guide/07-data-movement-ops.md +++ b/ptodsl/docs/user_guide/07-data-movement-ops.md @@ -341,7 +341,7 @@ The compiler automatically computes the byte offset from the tile's shape, eleme | `tile[start:]` | Tile index | 1D tile with starting element (vector-width range) | | `buf` | `PtrType` (UB) | Pointer to buffer in UB (pointer form) | | `offset` | `Index` | Element offset (pointer form) | -| `dist` | `VLoadDist` or `None` | Optional load distribution: `NORM` (default), `UNPK_B8`/`UNPK_B16`/`UNPK_B32`, `BRC_B8`/`BRC_B16`/`BRC_B32` | +| `dist` | `VLoadDist` or `None` | Optional load distribution: `NORM` (default), `UNPK_B8`/`UNPK_B16`/`UNPK_B32`, `BRC_B8`/`BRC_B16`/`BRC_B32`, `BRC_BLK`, `E2B_B16`/`E2B_B32`, `UNPK4`, `SPLT4CHN`, `US_B8`/`US_B16`, `DS_B8`/`DS_B16` | | `post_update` | `PostUpdate` | Pointer form only. `OFF` (default) — stateless load. `ON` — returns `(vec, updated_buf)` where `updated_buf` is the buffer pointer advanced past the loaded elements | **Returns**: @@ -815,7 +815,7 @@ pto.vstas(align, ub_dst_f32, pto.const(64)) | Enum | Values | Used with | |------|--------|-----------| -| `VLoadDist` | `NORM`, `UNPK_B8`, `UNPK_B16`, `UNPK_B32`, `BRC_B8`, `BRC_B16`, `BRC_B32`, `US_B8`, `US_B16`, `DS_B8`, `DS_B16` | `vlds` | +| `VLoadDist` | `NORM`, `UNPK_B8`, `UNPK_B16`, `UNPK_B32`, `BRC_B8`, `BRC_B16`, `BRC_B32`, `BRC_BLK`, `E2B_B16`, `E2B_B32`, `UNPK4`, `SPLT4CHN`, `US_B8`, `US_B16`, `DS_B8`, `DS_B16` | `vlds` | | `VStoreDist` | `NORM_B8`, `NORM_B16`, `NORM_B32`, `1PT_B8`, `1PT_B16`, `1PT_B32`, `PK_B16`, `PK_B32`, `PK_B64`, `PK4_B32`, `MRG4CHN_B8`, `MRG2CHN_B8`, `MRG2CHN_B16` | `vsts` | | `DeinterleaveDist` | `DINTLV_B8`, `DINTLV_B16`, `DINTLV_B32`, `BDINTLV` | `vldsx2` | | `InterleaveDist` | `INTLV_B8`, `INTLV_B16`, `INTLV_B32` | `vstsx2` | diff --git a/ptodsl/docs/user_guide/08-compute-operations.md b/ptodsl/docs/user_guide/08-compute-operations.md index df0952eedf..d7288834c3 100644 --- a/ptodsl/docs/user_guide/08-compute-operations.md +++ b/ptodsl/docs/user_guide/08-compute-operations.md @@ -1374,6 +1374,19 @@ s_shifted = pto.vsubs(s_row, m_next, col_mask) **Description**: Leaky ReLU — `vec[i] >= 0 ? vec[i] : alpha * vec[i]`. +#### `pto.vshls(vec: VRegType, scalar: ScalarType, mask: MaskType) -> VRegType` +#### `pto.vshrs(vec: VRegType, scalar: ScalarType, mask: MaskType) -> VRegType` + +**Description**: Uniform integer shift by a scalar amount. PTODSL coerces +`scalar` to signless `i16`, matching the VPTO `vshls`/`vshrs` requirement. + +#### `pto.vands(vec: VRegType, scalar: ScalarType, mask: MaskType) -> VRegType` +#### `pto.vors(vec: VRegType, scalar: ScalarType, mask: MaskType) -> VRegType` +#### `pto.vxors(vec: VRegType, scalar: ScalarType, mask: MaskType) -> VRegType` + +**Description**: Vector/scalar bitwise ops. PTODSL lowers these surface helpers +as `vbr(scalar)` followed by `vand(...)`, `vor(...)`, or `vxor(...)`. + --- ### 8.2.3.1 Vector duplication: `pto.vdup` @@ -1506,6 +1519,11 @@ These combine an arithmetic operation with a math function or activation in a si **Description**: Fused multiply-add: `alpha * x[i] + y[i]`. +#### `pto.vmula(acc: VRegType, lhs: VRegType, rhs: VRegType, mask: MaskType) -> VRegType` + +**Description**: Fused multiply-add with an explicit accumulator: +`acc[i] + lhs[i] * rhs[i]`. + --- #### `pto.vaddrelu(v0: VRegType, v1: VRegType, mask: MaskType) -> VRegType` @@ -1630,7 +1648,7 @@ These ops change the element type or layout of vector registers. They are distin **Constraints**: - Source and result dtype pair must be a legal hardware conversion. Illegal pairs (e.g., unsupported narrowing/widening combinations) are rejected at frontend time. -- `f32 -> f8e4m3/f8e5m2` requires `rnd=R`, `sat`, and `part=P0/P1/P2/P3`. +- `f32 -> f8e4m3/f8e5m2` requires `rnd=R/A/H/Z`, `sat`, and `part=P0/P1/P2/P3`. - `f32 -> hif8` requires `rnd=A/H`, `sat`, and `part=P0/P1/P2/P3`. - `f16/bf16 -> f8e4m3/f8e5m2` requires `rnd=R/A/F/Z/C`, `sat`, and `part=EVEN/ODD`. - `f16 -> hif8` requires `rnd=A/H`, `sat`, and `part=EVEN/ODD`. @@ -1703,20 +1721,121 @@ packed_high = pto.vpack(vec_i32, pto.VPackPart.HIGHER) # upper 64 lanes -> 128 --- -### 8.2.8 Vector compute quick reference +### 8.2.7.1 Index generation + +#### `pto.vci(base: ScalarType | int, order: OrderMode | None = None) -> VRegType` + +**Description**: Generate a lane-index vector starting from `base`. When the +base is a Python `int`, PTODSL defaults it to `i32`. To control the result +dtype, materialize a typed scalar explicitly before calling `vci`. + +**Examples**: + +```python +idx_i32 = pto.vci(0) +idx_i8 = pto.vci(pto.i8(0), pto.OrderMode.ASC) +typed_idx = pto.vci(pto.i32(16), order=pto.OrderMode.ASC) +``` + +--- + +### 8.2.8 Vector rearrangement + +These ops rearrange data between vector registers without touching UB memory. +They are useful for switching between interleaved layouts (`x0, y0, x1, y1, +...`) and split layouts (`x...`, `y...`) inside `@pto.simd`. + +#### `pto.vintlv(lhs: VRegType, rhs: VRegType) -> tuple[VRegType, VRegType]` + +**Description**: Interleave two vectors lane-by-lane and return the result as a +pair of vector registers. The first result contains the interleaved lower half +of the logical output stream; the second result contains the upper half. + +For a vector with `N` lanes: + +- `low = [lhs[0], rhs[0], lhs[1], rhs[1], ..., lhs[N/2 - 1], rhs[N/2 - 1]]` +- `high = [lhs[N/2], rhs[N/2], lhs[N/2 + 1], rhs[N/2 + 1], ..., lhs[N - 1], rhs[N - 1]]` + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `lhs` | `VRegType` | First source vector | +| `rhs` | `VRegType` | Second source vector | + +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `low` | `VRegType` | Interleaved lower half | +| `high` | `VRegType` | Interleaved upper half | + +**Constraints**: +- `lhs` and `rhs` must have exactly the same `VRegType`. +- The two returned vectors form one logical interleaved result pair; preserve + their ordering when passing them to later ops such as `vdintlv`. + +--- + +#### `pto.vdintlv(lhs: VRegType, rhs: VRegType) -> tuple[VRegType, VRegType]` + +**Description**: Deinterleave a previously interleaved vector pair. This is the +inverse of `vintlv`: it separates the even-position and odd-position lanes of +the logical input stream into two output vectors. + +For a vector with `N` lanes: + +- `low = [lhs[0], lhs[2], lhs[4], ..., rhs[0], rhs[2], rhs[4], ...]` +- `high = [lhs[1], lhs[3], lhs[5], ..., rhs[1], rhs[3], rhs[5], ...]` + +If `(packed_low, packed_high) = pto.vintlv(a, b)`, then +`pto.vdintlv(packed_low, packed_high)` reconstructs `(a, b)`. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `lhs` | `VRegType` | Lower half of the interleaved input stream | +| `rhs` | `VRegType` | Upper half of the interleaved input stream | + +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `low` | `VRegType` | Lanes from even interleaved positions | +| `high` | `VRegType` | Lanes from odd interleaved positions | + +**Constraints**: +- `lhs` and `rhs` must have exactly the same `VRegType`. +- `lhs` and `rhs` are interpreted as an ordered pair. Swapping them changes the + reconstructed lane order. + +**Example** — interleave two channels and recover them later: + + +```python +packed_low, packed_high = pto.vintlv(vec_f32, vec_f32) +even_lanes, odd_lanes = pto.vdintlv(packed_low, packed_high) +``` + +--- + +### 8.2.9 Vector compute quick reference | Category | Operations | |----------|------------| | Unary | `vexp`, `vln`, `vsqrt`, `vabs`, `vneg`, `vrec`, `vrsqrt`, `vrelu`, `vnot` | | Binary | `vadd`, `vsub`, `vmul`, `vdiv`, `vmax`, `vmin`, `vand`, `vor`, `vxor`, `vshl`, `vshr` | -| Vector-scalar | `vadds`, `vsubs`, `vmuls`, `vmaxs`, `vmins`, `vlrelu` | +| Vector-scalar | `vadds`, `vsubs`, `vmuls`, `vmaxs`, `vmins`, `vlrelu`, `vands`, `vors`, `vxors`, `vshls`, `vshrs` | | Broadcast | `vbr`, `vdup` | | Full reduction | `vcadd`, `vcmax`, `vcmin` | | Group reduction | `vcgadd`, `vcgmax`, `vcgmin` | | Scan | `vcpadd` | -| Fused | `vexpdif`, `vaxpy`, `vaddrelu`, `vsubrelu`, `vmulscvt` | +| Fused | `vexpdif`, `vaxpy`, `vmula`, `vaddrelu`, `vsubrelu`, `vmulscvt` | | Compare/select | `vcmp`, `vcmps`, `vsel` | | Conversion | `vcvt`, `vpack`, `vbitcast`, `pbitcast` | +| Index generation | `vci` | +| Rearrangement | `vintlv`, `vdintlv` | --- diff --git a/ptodsl/docs/user_guide/14-vmi-virtual-instruction-set.md b/ptodsl/docs/user_guide/14-vmi-virtual-instruction-set.md new file mode 100644 index 0000000000..6229bb953f --- /dev/null +++ b/ptodsl/docs/user_guide/14-vmi-virtual-instruction-set.md @@ -0,0 +1,1399 @@ +# 14. The Virtual Micro-instruction Set (VMI) + +The Virtual Micro-instruction Set (VMI) is a logical SIMD vector instruction +set exposed through the `pto.vmi` namespace. It provides a complete set of +virtual micro-instructions for writing vectorized kernels directly against a +hardware-abstracted instruction set — load, store, compute, compare, reduce, +convert, rearrange, and predicate control. + +Use `pto.vmi` when you want to: + +- author kernels against a stable, logical vector instruction set +- write SIMD code that mirrors the formal VMI specification directly +- carry explicit logical vector and mask types in your authored code +- bypass the top-level PTODSL vector helpers and work one level closer to the + hardware abstraction + +VMI is not a replacement for the existing top-level vector helpers +(`pto.vadd`, `pto.vlds`, etc.). The two surfaces coexist: the top-level helpers +remain the established PTODSL vector programming surface, while `pto.vmi` is +the explicit, instruction-set-oriented alternative. + +## 14.1 VMI logical types + +VMI introduces two logical type constructors. They describe a logical vector +register and a logical predicate mask at the PTODSL level — the physical +register mapping is handled by the backend. + +### `pto.vmi.vreg(lanes, dtype) -> TypeDescriptor` + +**Description**: Creates a logical VMI vector register type descriptor. +`lanes` is the logical lane count (not a physical register count). `dtype` is +a PTODSL element type token such as `pto.f32`, `pto.f16`, `pto.i32`, etc. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `lanes` | `int` | Logical lane count. Must be a multiple of 64. See Constraints below | +| `dtype` | `DType` | Element type token (`pto.f32`, `pto.f16`, `pto.i32`, etc.) | + +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `type_descriptor` | VMI vreg type | Logical VMI vector type descriptor | + +**Constraints**: + +- `lanes` must be a multiple of 64. +- `lanes · bitwidth(dtype)` determines the physical register count + `K = ⌈ lanes · bitwidth(dtype) / 2048 ⌉`. Each physical register is + 256 B (2048 bits). +- Common legal combinations: + + | dtype | bitwidth | lanes per physical reg | example `lanes` | + |-------|----------|------------------------|-----------------| + | `f32`, `i32`, `ui32`, `si32` | 32 | 64 | 64, 128, 256 | + | `f16`, `bf16`, `i16`, `ui16`, `si16` | 16 | 128 | 64, 128, 256 | + | `i8`, `ui8`, `si8`, `fp8_e4m3`, `fp8_e5m2` | 8 | 256 | 64, 128, 256 | + +- Compact/partial vectors (`K < 1` in the formula above, e.g. `vreg(64, pto.f16)` + = 128 B) still occupy one physical register; lanes outside the logical value + are undefined and must be masked out. + +**Example**: + +```python +vec_f32 = pto.vmi.vreg(64, pto.f32) # 1 physical reg (64 × 32b = 256B) +vec_i32 = pto.vmi.vreg(64, pto.i32) # 1 physical reg +vec_f16 = pto.vmi.vreg(128, pto.f16) # 1 physical reg (128 × 16b = 256B) +vec_f32_x2 = pto.vmi.vreg(128, pto.f32) # 2 physical regs (128 × 32b = 512B) +``` + +--- + +### `pto.vmi.mask(lanes) -> TypeDescriptor` + +**Description**: Creates a logical VMI mask type descriptor. The predicate +granularity is always per-lane: one mask bit governs one vector lane. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `lanes` | `int` | Logical lane count. Must match the gated vector's lanes. See Constraints below | + +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `type_descriptor` | VMI mask type | Logical VMI mask type descriptor | + +**Constraints**: + +- `lanes` must match the lane count of the vector being gated. A `mask(64)` + gates a `vreg(64, ...)`; a `mask(128)` gates a `vreg(128, ...)`. + +**Example**: + +```python +mask64 = pto.vmi.mask(64) # gates a vreg(64, ...) +mask128 = pto.vmi.mask(128) # gates a vreg(128, ...) +``` + +**About VMI layouts.** A VMI logical vector value may span `K` physical vector +registers (256 B each). PTOAS tracks the logical lane layout internally during +lowering; PTODSL does not expose layout selection on `pto.vmi.vreg(...)` or +`pto.vmi.mask(...)`. Two common internal layouts are: + +| Layout | Description | +|--------|-------------| +| `contiguous` (default) | Stride-1 mapping: lane `i` sits at position `i mod (2048/bitwidth(T))` within physical register `⌊i / (2048/bitwidth(T))⌋` | +| `deinterleaved` | Parity split: EVEN lanes occupy the first `K/2` physical registers, ODD lanes occupy the second `K/2`. This is the natural output layout of a widening `vcvt` (e.g., `f16 → f32`) or a `vload` with `dist_mode="dintlv"` | + +Layouts are a lowering concern managed by `ptoas`. The system +propagates layouts automatically through Category A ops (most elementwise +compute) and inserts materialization at Category C boundaries. In day-to-day +authoring you do not spell layouts explicitly; just declare the logical lane +count and element type, and let PTOAS infer the internal layout. + +VMI types are mainly used as type annotations. They +are not Python callables that produce values — use `pto.vmi.vload`, +`pto.vmi.vci`, etc. to produce actual VMI vector values, and +`pto.vmi.create_mask` to produce actual VMI mask +values. + +--- + +## 14.2 Load and store + +The load/store family moves data between UB memory and VMI logical vector +registers. These are the primary entry and exit points for VMI vector data. + +PTODSL groups the `vload` / `vstore` surface into three mutually exclusive +mode families: + +- `dist_mode`: the regular logical memory surface. This covers the default + contiguous case plus other access patterns selected by `dist_mode`. +- `group`: grouped row-strided load/store. +- `block_stride`: block-strided load/store using paired + `block_stride` / `repeat_stride` operands. + +Pick exactly one family per call. Do not mix `dist_mode`, `group`, and +`block_stride` parameters in the same `vload` / `vstore`. + +### `vload` + +### `pto.vmi.vload(source, offset, *, size, dist_mode=None, to_dtype=None) -> VRegType` +### `pto.vmi.vload(source, offset, *, size, dist_mode="dintlv") -> (VRegType, VRegType)` +### `pto.vmi.vload(source, offset, *, size, group, stride) -> VRegType` +### `pto.vmi.vload(source, offset, *, size, group, stride, dist_mode="brc") -> VRegType` +### `pto.vmi.vload(source, offset, *, size, block_stride, repeat_stride) -> VRegType` + +**Description**: Loads a logical VMI vector from a UB pointer. The element +type is derived from the source pointer; `size` determines the logical lane +count. Which memory access pattern is used depends on the selected mode family. + +**Common parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `source` | `PtrType` (ub) | UB source pointer | +| `offset` | `IndexLike` | Element offset into the source buffer | +| `size` | `int` | Logical result lane count | + +**About `mask` on `vload`.** + +- `pto.vmi.vload(...)` does not take an explicit `mask` parameter. The load + surface describes how data is read from UB into a logical VMI value, not + which lanes of a later computation are active. +- Tail handling and partial-lane participation are expressed on the consumer + side, typically by passing a mask to a later compute op such as + `pto.vmi.vadd(...)`, or to the final `pto.vmi.vstore(...)`. +- In practice, if you need "load only the active lanes" behavior in authored + DSL code, write a normal `vload`, then apply your mask on the first consumer + or on the eventual store. + +**Mode 1: `dist_mode`** + +Use this family for the normal logical load surface. `dist_mode=None` means the +default contiguous load. + +| Parameter | Type | Description | +|-----------|------|-------------| +| `dist_mode` | `str` or `None` | One of `None`, `"continuous"`, `"dintlv"`, `"unpack"`, or `"brc"` | +| `to_dtype` | `DType` or `None` | Required only when `dist_mode="unpack"` | + +Dist-mode behavior: + +- `None` or `"continuous"`: contiguous stride-1 load, returning one VMI vector. +- `"dintlv"`: deinterleaved load, returning an `(even, odd)` pair. +- `"unpack"`: widening unpack load, returning one widened VMI vector. +- `"brc"`: broadcast load from one source element, returning one VMI vector. + +**Mode 2: `group`** + +Use this family for grouped row-strided accesses. + +| Parameter | Type | Description | +|-----------|------|-------------| +| `group` | `int` | Number of groups | +| `stride` | `IndexLike` | Element stride between groups | + +Grouped load behavior: + +- `size > group`: full-group load. Each group contributes `size / group` + elements. +- `size == group`: slot load. Each group contributes one scalar slot. +- `dist_mode="brc"` with `group` switches grouped load into grouped broadcast: + one source scalar is loaded per group and broadcast within that group. + +**Mode 3: `block_stride`** + +Use this family for block-strided accesses. + +| Parameter | Type | Description | +|-----------|------|-------------| +| `block_stride` | `int` | 16-bit block stride operand | +| `repeat_stride` | `int` | 16-bit repeat stride operand | + +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `vec` | `VRegType` | Single vector result for continuous, unpack, brc, group, and block-stride modes | +| `(even, odd)` | `(VRegType, VRegType)` | Two-vector result for `dist_mode="dintlv"` | + +**Examples**: + +Continuous load: + +```python +lhs = pto.vmi.vload(src_ptr, offset, size=64) +rhs = pto.vmi.vload(other_ptr, offset, size=64) +``` + +Dist-mode unpack load: + +```python +wide = pto.vmi.vload( + src_ptr, + offset, + size=128, + dist_mode="unpack", + to_dtype=pto.i16, +) +``` + +Broadcast load: + +```python +bias = pto.vmi.vload( + src_ptr, + offset, + size=64, + dist_mode="brc", +) +``` + +Group-mode load: + +```python +tile = pto.vmi.vload( + src_ptr, + offset, + size=64, + group=8, + stride=row_stride, +) +``` + +Grouped broadcast load: + +```python +grouped = pto.vmi.vload( + src_ptr, + offset, + size=64, + group=8, + stride=row_stride, + dist_mode="brc", +) +``` + +Block-stride load: + +```python +blocks = pto.vmi.vload( + src_ptr, + offset, + size=64, + block_stride=pto.i16(8), + repeat_stride=pto.i16(0), +) +``` + +**Constraints**: + +- `size` is required for every `vload` form. +- `block_stride` mode is mutually exclusive with both `dist_mode` and `group`. +- `group` and `dist_mode` are mutually exclusive except for grouped broadcast, + spelled as `group=...`, `stride=...`, `dist_mode="brc"`. +- `to_dtype` is only accepted when `dist_mode="unpack"`. +- `stride` is only accepted when `group` is provided. +- `block_stride` and `repeat_stride` must be provided together. +- The unpack form widens by exactly one adjacent bit-width step. + +--- + +### `vstore` + +### `pto.vmi.vstore(values, destination, offset, mask=None, *, dist_mode=None, pmode=None) -> None` +### `pto.vmi.vstore((even, odd), destination, offset, mask=None, *, dist_mode="dintlv", pmode=None) -> None` +### `pto.vmi.vstore(values, destination, offset, *, group, stride, pmode=None) -> None` +### `pto.vmi.vstore(values, destination, offset, mask=None, *, block_stride, repeat_stride, pmode=None) -> None` + +**Description**: Writes one logical VMI vector, or a deinterleaved pair, back +to a UB pointer. As with `vload`, the PTODSL surface is organized into the same +three mutually exclusive mode families. + +**Common parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `values` | `VRegType` or `(VRegType, VRegType)` | One VMI vector for normal forms, or an `(even, odd)` pair for `dist_mode="dintlv"` | +| `destination` | `PtrType` (ub) | UB destination pointer | +| `offset` | `IndexLike` | Element offset into the destination buffer | +| `pmode` | `str` or `None` | Optional inactive-lane mode: `"zero"` stores 0 to masked-off lanes; `"merge"` skips the write for masked-off lanes | + +**About `pmode` on `vstore`.** + +- `pmode="zero"` is the default store behavior. When a `mask` is present, + inactive lanes are written as zero. +- `pmode="merge"` preserves destination contents on inactive lanes by skipping + those writes. +- `pmode` only matters on store forms that actually use a `mask`. Group-mode + store does not take a mask operand, so there are no inactive lanes to define + there. + +**Mode 1: `dist_mode`** + +| Parameter | Type | Description | +|-----------|------|-------------| +| `mask` | VMI mask or `None` | Optional predicate mask gating which lanes are written | +| `dist_mode` | `str` or `None` | One of `None`, `"continuous"`, or `"dintlv"` | + +Dist-mode behavior: + +- `None` or `"continuous"`: contiguous store of one VMI vector. +- `"dintlv"`: interleaved store of an `(even, odd)` pair. + +**Mode 2: `group`** + +| Parameter | Type | Description | +|-----------|------|-------------| +| `group` | `int` | Number of groups | +| `stride` | `IndexLike` | Element stride between groups | + +Group store writes one grouped logical stream. On the current VMI contract, +group-mode store does not take a mask operand. + +**Mode 3: `block_stride`** + +| Parameter | Type | Description | +|-----------|------|-------------| +| `mask` | VMI mask or `None` | Optional store mask; omitting it means all lanes are active | +| `block_stride` | `int` | 16-bit block stride operand | +| `repeat_stride` | `int` | 16-bit repeat stride operand | + +**Returns**: None (side-effect operation). + +**Examples**: + +Continuous store: + +```python +pto.vmi.vstore(vec, dst_ptr, offset, mask) +``` + +Group-mode store: + +```python +pto.vmi.vstore( + tile, + dst_ptr, + offset, + group=8, + stride=row_stride, +) +``` + +Block-stride store: + +```python +pto.vmi.vstore( + vec, + dst_ptr, + offset, + mask, + block_stride=pto.i16(8), + repeat_stride=pto.i16(0), +) +``` + +**Constraints**: + +- `dist_mode`, `group`, and `block_stride` mode selection are mutually + exclusive. +- `dist_mode="dintlv"` requires `values` to be an `(even, odd)` pair. +- Group mode requires `group` and `stride`, and does not accept `mask`. +- `block_stride` and `repeat_stride` must be provided together. + +--- + +## 14.3 Index generation and broadcast + +These instructions produce a new logical vector from a scalar seed — either as +a lane-wise ramp or a uniform broadcast. + +### `pto.vmi.vci(base, *, size, order=None) -> VRegType` + +**Description**: Builds a logical lane-wise index ramp starting from a scalar +base value. Use it when you need an index vector for lane addressing, +gather/scatter offsets, or dynamic lane selection. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `base` | `ScalarType` | Typed scalar starting value for the ramp | +| `size` | `int` | Logical lane count of the result vector | +| `order` | `str` or `None` | Ramp order: `"ASC"` for ascending (default if omitted), or `"DESC"` for descending | + +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `idx` | `VRegType` | Lane-wise index vector | + +**Example**: + +```python +idx = pto.vmi.vci(pto.i32(0), size=64, order="ASC") +out = pto.vmi.vselr(src, idx) +``` + +**Constraints**: +- `base` must already carry a scalar dtype. A plain Python literal like `0` + is ambiguous, so use `pto.i32(0)`, `pto.i16(0)`, `pto.f16(0.0)`, or + `pto.f32(0.0)`. +- `size` determines the logical lane count of the result vector. + +--- + +### `pto.vmi.vbrc(value, *, size) -> VRegType` +### `pto.vmi.vbrc(value, *, size, group) -> VRegType` + +**Description**: Broadcasts a scalar value (or a compact group-shaped input) +across all lanes of a logical vector. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `value` | `ScalarType` or `VRegType` | Scalar to broadcast, or a compact VMI vector for grouped broadcast | +| `size` | `int` | Logical lane count of the expanded result | +| `group` | `int` or `None` | Group count for grouped broadcast. When provided, `value` is treated as a compact group-shaped input and expanded accordingly | + +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `result` | `VRegType` | Broadcast vector | + +**Example** — scalar broadcast: + +```python +bias = pto.vmi.vbrc(pto.f32(0.0), size=64) +``` + +**Example** — grouped broadcast: + +```python +expanded = pto.vmi.vbrc(compact, size=256, group=8) +``` + +**Constraints**: +- `value` must already carry a scalar dtype or be a VMI vector. A plain Python + literal like `0.0` is ambiguous, so use `pto.f32(0.0)` or `pto.i32(23)`. +- When `group` is provided, `value` must be a VMI vector whose lane count + matches the group count. + +--- + +## 14.4 Elementwise compute + +Elementwise instructions operate lane-by-lane on one or two VMI vector +operands. They form the arithmetic core of VMI SIMD kernels. + +### 14.4.1 Binary vector-vector + +#### `pto.vmi.vadd(lhs, rhs, mask=None, *, pmode=None) -> VRegType` +#### `pto.vmi.vsub(lhs, rhs, mask=None, *, pmode=None) -> VRegType` +#### `pto.vmi.vmul(lhs, rhs, mask=None, *, pmode=None) -> VRegType` +#### `pto.vmi.vdiv(lhs, rhs, mask=None, *, pmode=None) -> VRegType` +#### `pto.vmi.vmax(lhs, rhs, mask=None, *, pmode=None) -> VRegType` +#### `pto.vmi.vmin(lhs, rhs, mask=None, *, pmode=None) -> VRegType` +#### `pto.vmi.vand(lhs, rhs, mask=None, *, pmode=None) -> VRegType` +#### `pto.vmi.vor(lhs, rhs, mask=None, *, pmode=None) -> VRegType` +#### `pto.vmi.vxor(lhs, rhs, mask=None, *, pmode=None) -> VRegType` +#### `pto.vmi.vshl(lhs, rhs, mask=None, *, pmode=None) -> VRegType` +#### `pto.vmi.vshr(lhs, rhs, mask=None, *, pmode=None) -> VRegType` + +**Description**: Element-wise binary operation: `result[i] = lhs[i] rhs[i]` +for lanes where `mask[i]` is true (or all lanes when `mask` is omitted). + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `lhs` | `VRegType` | First operand vector | +| `rhs` | `VRegType` | Second operand vector | +| `mask` | VMI mask or `None` | Optional predicate mask gating lane participation | +| `pmode` | `str` or `None` | Optional predicate mode: `"merge"` keeps predicate-inactive lanes at their prior value; `"zero"` writes 0 | + +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `result` | `VRegType` | Result vector (same shape and element type as `lhs`) | + +**Example**: + +```python +out = pto.vmi.vadd(lhs, rhs, mask) +out = pto.vmi.vmul(scale, data, full_mask) +``` + +**Constraints**: +- `lhs` and `rhs` must have compatible shapes and element types. +- The result type is inferred from `lhs`. +- For bitwise ops (`vand`, `vor`, `vxor`, `vshl`, `vshr`), integer element + types are expected. Floating-point usage is rejected. +- `vshr` performs logical right shift for explicit unsigned element types and + arithmetic right shift for signed or signless element types. + +--- + +### 14.4.2 Unary vector + +#### `pto.vmi.vabs(source, mask=None, *, pmode=None) -> VRegType` +#### `pto.vmi.vneg(source, mask=None, *, pmode=None) -> VRegType` +#### `pto.vmi.vrelu(source, mask=None, *, pmode=None) -> VRegType` +#### `pto.vmi.vexp(source, mask=None, *, pmode=None) -> VRegType` +#### `pto.vmi.vln(source, mask=None, *, pmode=None) -> VRegType` +#### `pto.vmi.vsqrt(source, mask=None, *, pmode=None) -> VRegType` +#### `pto.vmi.vnot(source, mask=None, *, pmode=None) -> VRegType` + +**Description**: Element-wise unary operation: `result[i] = op(source[i])` for +active lanes. `vrelu` = `max(0, x)`, `vnot` = bitwise NOT (integer types +only). + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `source` | `VRegType` | Input vector | +| `mask` | VMI mask or `None` | Optional predicate mask gating lane participation | +| `pmode` | `str` or `None` | Optional predicate mode: `"merge"` keeps predicate-inactive lanes at their prior value; `"zero"` writes 0 | + +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `result` | `VRegType` | Result vector (same shape and element type as `source`) | + +**Example**: + +```python +activated = pto.vmi.vrelu(src, mask) +inverted = pto.vmi.vnot(int_vec) +``` + +--- + +### 14.4.3 Vector-scalar + +Formal `pto.vmi` vector-scalar ops in VMI v0.1: + +#### `pto.vmi.vadds(source, scalar, mask, *, pmode=None) -> VRegType` +#### `pto.vmi.vmuls(source, scalar, mask, *, pmode=None) -> VRegType` +#### `pto.vmi.vmaxs(source, scalar, mask, *, pmode=None) -> VRegType` +#### `pto.vmi.vmins(source, scalar, mask, *, pmode=None) -> VRegType` +#### `pto.vmi.vshls(source, scalar, mask, *, pmode=None) -> VRegType` +#### `pto.vmi.vshrs(source, scalar, mask, *, pmode=None) -> VRegType` + +The following are **PTODSL syntax sugar** — convenience wrappers provided by the +PTODSL authoring layer. They have **no corresponding VMI instruction**; PTODSL lowers +each to an equivalent `pto.vmi.*` form (e.g., `pto.vsubs` lowers to `pto.vmi.vadds` with a +negated scalar). Users may freely use these spellings in PTODSL programs, but tooling and +the VMI v0.1 spec only recognize the formal `pto.vmi.*` ops listed above. + +#### `pto.vsubs(source, scalar, mask) -> VRegType` +#### `pto.vands(source, scalar, mask) -> VRegType` +#### `pto.vors(source, scalar, mask) -> VRegType` +#### `pto.vxors(source, scalar, mask) -> VRegType` + +**Description**: Element-wise operation with a uniform scalar second operand: +`result[i] = source[i] scalar`. The scalar is broadcast to all active +lanes. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `source` | `VRegType` | Input vector | +| `scalar` | `ScalarType` | Scalar operand (Python number or PTODSL scalar). Automatically coerced to the vector element type | +| `mask` | VMI mask | **Required.** Predicate mask gating lane participation | +| `pmode` | `str` or `None` | Optional predicate mode: `"merge"` keeps predicate-inactive lanes at their prior value; `"zero"` writes 0 | + +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `result` | `VRegType` | Result vector (same shape and element type as `source`) | + +**Example**: + +```python +scaled = pto.vmi.vmuls(data, 0.5, full_mask) +shifted = pto.vsubs(scores, row_max, col_mask) +``` + +**Constraints**: +- `mask` is always required for vector-scalar ops — unlike binary and unary + ops, there is no mask-optional form. +- `scalar` is coerced to match the element type of `source`. +- `vsubs`, `vands`, `vors`, and `vxors` are PTODSL convenience spellings, not + dedicated formal `pto.vmi` instructions in VMI v0.1. +- `vdivs` is listed for symmetry with `vdiv`, but it is not currently surfaced + as a PTODSL API and is not part of the formal VMI v0.1 instruction set. + +--- + +## 14.5 Compare and select + +Compare instructions produce logical VMI masks from vector data. Select +instructions consume masks to pick between values lane by lane. + +### `pto.vmi.vcmp(lhs, rhs, seed, cmp, *, pmode=None) -> MaskType` + +**Description**: Element-wise vector-vector comparison producing a VMI mask: +`result[i] = seed[i] ? (lhs[i] rhs[i]) : 0`. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `lhs` | `VRegType` | First operand vector | +| `rhs` | `VRegType` | Second operand vector | +| `seed` | VMI mask | Seed mask gating which lanes participate | +| `cmp` | `str` | Comparison predicate. VMI accepts bare predicates `"eq"`, `"ne"`, `"lt"`, `"le"`, `"gt"`, `"ge"`. Floating-point compares also accept ordered forms `"oeq"`, `"one"`, `"olt"`, `"ole"`, `"ogt"`, `"oge"` | +| `pmode` | `str` or `None` | Optional predicate mode: `"merge"` keeps predicate-inactive lanes at their prior value; `"zero"` writes 0 | +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `pred` | VMI mask | Result predicate mask (same granularity and lane count as `seed`) | + +--- + +### `pto.vmi.vcmps(source, scalar, seed, cmp, *, pmode=None) -> MaskType` + +**Description**: Vector-scalar comparison: `result[i] = seed[i] ? (source[i] scalar) : 0`. +The scalar is broadcast to all lanes. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `source` | `VRegType` | Input vector | +| `scalar` | `ScalarType` | Scalar operand (coerced to the vector element type) | +| `seed` | VMI mask | Seed mask gating lane participation | +| `cmp` | `str` | Comparison predicate. Same accepted spellings as `vcmp` | +| `pmode` | `str` or `None` | Optional predicate mode: `"merge"` keeps predicate-inactive lanes at their prior value; `"zero"` writes 0 | +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `pred` | VMI mask | Result predicate mask | + +**Example**: + +```python +pred = pto.vmi.vcmp(lhs, rhs, seed_mask, "gt") +pred2 = pto.vmi.vcmps(src, 0.0, seed_mask, "ge") +``` + +--- + +### `pto.vmi.vsel(mask, true_value, false_value, *, pmode=None) -> VRegType` + +**Description**: Per-lane ternary select: `result[i] = mask[i] ? true_value[i] : false_value[i]`. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `mask` | VMI mask | Selection predicate | +| `true_value` | `VRegType` | Value taken when mask is true | +| `false_value` | `VRegType` | Value taken when mask is false | +| `pmode` | `str` or `None` | Optional predicate mode: `"merge"` keeps predicate-inactive lanes at their prior value; `"zero"` writes 0 | +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `result` | `VRegType` | Selected vector | + +--- + +### `pto.vmi.vselr(source, index) -> VRegType` + +**Description**: Dynamic per-lane selection from a source vector using an +index vector: `result[i] = source[index[i]]`. This is a gather-style select +within a single vector register. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `source` | `VRegType` | Source vector to select from | +| `index` | `VRegType` | Integer index vector (per-lane source lane indices) | +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `result` | `VRegType` | Selected vector | + +**Example**: + +```python +out = pto.vmi.vselr(src, idx) +``` + +**Constraints**: +- The result type is inferred directly from `source`. +- `index` must be an integer-typed VMI vector. + +--- + +## 14.6 Reduction + +Reduction instructions collapse a logical vector along its lane dimension, +producing a smaller logical result. + +### `pto.vmi.vcadd(source, mask, *, group=None, reassoc, pmode=None) -> VRegType` +### `pto.vmi.vcmax(source, mask, *, group=None, pmode=None) -> VRegType` +### `pto.vmi.vcmin(source, mask, *, group=None, pmode=None) -> VRegType` + +**Description**: Full-vector or grouped reduction. `vcadd` computes the sum, +`vcmax` / `vcmin` compute the maximum / minimum with their lane index. When +`group` is omitted (or `None`), the reduction is across the full vector and +the result lane count is 1. When `group` is provided, the vector is +partitioned into that many equal-sized groups and a separate reduction is +performed per group. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `source` | `VRegType` | Input vector | +| `mask` | VMI mask | **Required.** Predicate mask gating lane participation | +| `group` | `int` or `None` | Number of groups for per-group reduction. `None` means full-vector reduction | +| `reassoc` | `bool` | For `vcadd` on floating-point data only: PTODSL requires this keyword to be written explicitly as `True` or `False` | +| `pmode` | `str` or `None` | Optional predicate mode: `"merge"` keeps predicate-inactive lanes at their prior value; `"zero"` writes 0 | + +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `result` | `VRegType` | Reduced vector | + +**Example** — full-vector reduction: + +```python +total = pto.vmi.vcadd(src, mask, reassoc=True) +peak = pto.vmi.vcmax(src, mask) +``` + +**Example** — grouped reduction: + +```python +group_max = pto.vmi.vcmax( + src, + mask, + group=8, +) +``` + +**Constraints**: +- `mask` is always required. +- PTODSL infers the result type automatically: full-vector reduction returns + `!pto.vmi.vreg<1xT>`, and grouped reduction returns `!pto.vmi.vreg`, + where `T` is the source element type and `G` is `group`. +- `reassoc` is only meaningful for `vcadd` on floating-point data. +- Floating-point `vcadd` must spell `reassoc` explicitly at the PTODSL surface. +- `reassoc=None` is rejected by PTODSL; use `reassoc=True` or `reassoc=False`. +- The current VMI op encoding remains presence-based, so `reassoc=False` + lowers to the same no-attribute form as legacy callers. + +--- + +## 14.7 Conversion and reinterpretation + +### `pto.vmi.vcvt(source, to_dtype, *, rounding=None, saturate=None, pmode=None) -> VRegType` + +**Description**: Numeric type conversion between VMI vector element types. +Converts the element type of `source` to the target element type. PTODSL +infers the result vector type from the source lane count/layout and `to_dtype`. + +For int→int widening, the source element type must carry signedness +(e.g. `si8`/`ui8`/`si16`/`ui16`); signless integers are rejected. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `source` | `VRegType` | Input vector (source element type) | +| `to_dtype` | `DType` | Target element type. PTODSL derives the result vector type from the source lane count/layout and this dtype | +| `rounding` | rounding mode or `None` | Optional rounding mode token | +| `saturate` | saturate mode or `None` | Optional saturation mode token | +| `pmode` | `str` or `None` | Optional predicate mode: `"merge"` keeps predicate-inactive lanes at their prior value; `"zero"` writes 0 | + +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `result` | `VRegType` | Converted vector | + +**Example**: + +```python +wide = pto.vmi.vcvt(src_f16, pto.f32) +narrow = pto.vmi.vcvt(src_f32, pto.f16) +``` + +**Constraints**: +- The masked form of `vcvt` is not currently supported on this surface. +- The source and target dtype pair must be legal for the target backend. +- For `f32 -> f8e4m3/f8e5m2`, PTODSL accepts `rounding="R"`, `"A"`, `"H"`, + and `"Z"`; other low-level rounding tokens remain rejected on the VMI + surface. + +--- + +### `pto.vmi.vinterpret_cast(source, to_dtype) -> VRegType` + +**Description**: Bitwise reinterpretation of a VMI vector under a different +element type. The logical bit pattern is unchanged; only the element type +annotation changes. This is a reinterpretation, not a numeric conversion. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `source` | `VRegType` | Input vector | +| `to_dtype` | `DType` | **Required.** Target element type. PTODSL keeps the source lane count/layout and reinterprets each lane at the new type | + +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `result` | `VRegType` | Reinterpreted vector | + +**Example**: + +```python +as_i32 = pto.vmi.vinterpret_cast(src, pto.i32) +``` + +**Constraints**: +- `to_dtype` is always required — PTODSL must not guess a reinterpretation + target element type. +- The source and target element widths must match. + +--- + +## 14.8 SFU, fused, and indexed memory instructions + +This family covers special-function-unit ops, fused multiply-accumulate forms, +and indexed memory access (gather, scatter, histogram). They go beyond simple +elementwise arithmetic. + +### `pto.vmi.vexpdif(x, max_value, mask, *, pmode=None) -> VRegType` + +**Description**: Computes `exp(x[i] - max_value[i])` for active lanes — the +stable softmax numerator. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `x` | `VRegType` | Input vector | +| `max_value` | `VRegType` | Maximum value vector to subtract before exponentiation | +| `mask` | VMI mask | **Required.** Predicate mask gating lane participation | +| `pmode` | `str` or `None` | Optional predicate mode: `"merge"` keeps predicate-inactive lanes at their prior value; `"zero"` writes 0 | +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `result` | `VRegType` | `exp(x - max_value)` | + +--- + +### `pto.vmi.vaxpy(x, acc, alpha, mask, *, pmode=None) -> VRegType` + +**Description**: Fused multiply-add: `result[i] = alpha * x[i] + acc[i]`. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `x` | `VRegType` | Input vector | +| `acc` | `VRegType` | Accumulator vector | +| `alpha` | `ScalarType` | Scalar multiplier (coerced to the element type of `x`) | +| `mask` | VMI mask | **Required.** Predicate mask | +| `pmode` | `str` or `None` | Optional predicate mode: `"merge"` keeps predicate-inactive lanes at their prior value; `"zero"` writes 0 | +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `result` | `VRegType` | `alpha * x + acc` | + +--- + +### `pto.vmi.vlrelu(x, slope, mask, *, pmode=None) -> VRegType` + +**Description**: Leaky ReLU: `result[i] = x[i] >= 0 ? x[i] : slope * x[i]`. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `x` | `VRegType` | Input vector | +| `slope` | `ScalarType` | Negative-slope multiplier (coerced to the element type of `x`) | +| `mask` | VMI mask | **Required.** Predicate mask | +| `pmode` | `str` or `None` | Optional predicate mode: `"merge"` keeps predicate-inactive lanes at their prior value; `"zero"` writes 0 | +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `result` | `VRegType` | Leaky ReLU result | + +--- + +### `pto.vmi.vprelu(x, alpha, mask, *, pmode=None) -> VRegType` + +**Description**: Parametric ReLU with a per-lane vector alpha: +`result[i] = x[i] >= 0 ? x[i] : alpha[i] * x[i]`. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `x` | `VRegType` | Input vector | +| `alpha` | `VRegType` | Per-lane slope vector | +| `mask` | VMI mask | **Required.** Predicate mask | +| `pmode` | `str` or `None` | Optional predicate mode: `"merge"` keeps predicate-inactive lanes at their prior value; `"zero"` writes 0 | +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `result` | `VRegType` | Parametric ReLU result | + +--- + +### `pto.vmi.vmull(a, b, mask, *, pmode=None) -> (VRegType, VRegType)` + +**Description**: Widening multiply for 32-bit integer vectors. PTODSL returns +a `(low, high)` pair of 32-bit VMI vectors; `low` carries the lower 32 bits +and `high` carries the upper 32 bits. Signedness follows the inputs. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `a` | `VRegType` | First `i32` or `ui32` operand vector | +| `b` | `VRegType` | Matching second operand vector | +| `mask` | VMI mask | **Required.** Predicate mask | +| `pmode` | `str` or `None` | Optional predicate mode: `"merge"` keeps predicate-inactive lanes at their prior value; `"zero"` writes 0 | + +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `low` | `VRegType` | Lower 32 bits of the widened product | +| `high` | `VRegType` | Upper 32 bits of the widened product | + +**Example**: + +```python +low, high = pto.vmi.vmull(a32, b32, mask) +``` + +**Constraints**: +- `a` and `b` must be identical 32-bit integer VMI vectors. +- `low` and `high` each have the same lane count and signedness as `a`. + +### `pto.vmi.vmula(acc, lhs, rhs, mask, *, pmode=None) -> VRegType` + +**Description**: Fused multiply-accumulate: `result = acc + (lhs * rhs)`. +The accumulator is both the input and the result value. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `acc` | `VRegType` | Accumulator vector | +| `lhs` | `VRegType` | First multiply operand | +| `rhs` | `VRegType` | Second multiply operand | +| `mask` | VMI mask | **Required.** Predicate mask | +| `pmode` | `str` or `None` | Optional predicate mode: `"merge"` keeps predicate-inactive lanes at their prior value; `"zero"` writes 0 | +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `result` | `VRegType` | Accumulator after the multiply-add | + +**Constraints**: +- `acc`, `lhs`, `rhs`, and `result` must have identical VMI vreg types. +- Supported element types are `i8`–`i32`, `f16`, `bf16`, and `f32`. + +--- + +### `pto.vmi.vdhist(acc, source, mask) -> VRegType` + +**Description**: Distribution histogram (dhistv2). Counts per-bin +occurrences of `source` values and adds them to the accumulator `acc`, +producing a 256-bin unsigned 16-bit result. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `acc` | `VRegType` | Accumulator vector (256×ui16) | +| `source` | `VRegType` | Source values (N×ui8) | +| `mask` | VMI mask | **Required.** Predicate mask (b8 granularity) | + +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `result` | `VRegType` | Distribution histogram counts per bin | + +**Constraints**: +- PTODSL infers the result vector type from `acc`; it must be the matching + 256×ui16 histogram accumulator layout. + +--- + +### `pto.vmi.vchist(acc, source, mask) -> VRegType` + +**Description**: Cumulative histogram (chistv2 half-axis). Same signature +as `vdhist` but each bin accumulates the sum of counts for all bins ≤ its index +(cumulative / prefix-sum semantics). + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `acc` | `VRegType` | Accumulator vector (256×ui16) | +| `source` | `VRegType` | Source values (N×ui8) | +| `mask` | VMI mask | **Required.** Predicate mask (b8 granularity) | + +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `result` | `VRegType` | Cumulative histogram counts per bin | + +**Constraints**: +- PTODSL infers the result vector type from `acc`; it must be the matching + 256×ui16 histogram accumulator layout. + +--- + +### `pto.vmi.vgather(source, offsets, mask, *, pmode=None) -> VRegType` + +**Description**: Indexed gather from a UB pointer using per-lane element +offsets. Only masked-on lanes participate; masked-off lanes produce an +unspecified value. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `source` | `PtrType` (ub) | UB source pointer | +| `offsets` | `VRegType` | Per-lane element offsets (integer VMI vector) | +| `mask` | VMI mask | **Required.** Predicate mask gating lane participation | +| `pmode` | `str` or `None` | Optional predicate mode: `"merge"` keeps predicate-inactive lanes at their prior value; `"zero"` writes 0 | + +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `result` | `VRegType` | Gathered vector | + +**Constraints**: +- PTODSL infers the result lane count from `offsets` and the element type + from `source`. + +--- + +### `pto.vmi.vgatherb(source, offsets, mask, *, pmode=None) -> VRegType` + +**Description**: Block gather from a UB pointer. Each participating lane +gathers one 32-byte block using byte-level offsets. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `source` | `PtrType` (ub) | UB source pointer | +| `offsets` | `VRegType` | Per-lane byte offsets (integer VMI vector) | +| `mask` | VMI mask | **Required.** Predicate mask | +| `pmode` | `str` or `None` | Optional predicate mode: `"merge"` keeps predicate-inactive lanes at their prior value; `"zero"` writes 0 | +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `result` | `VRegType` | Block-gathered vector | + +**Constraints**: +- PTODSL infers the result element type from `source`. +- PTODSL infers the result lane count, mask granularity, and layout match from + `mask`. + +--- + +### `pto.vmi.vscatter(value, destination, offsets, mask, *, pmode=None) -> None` + +**Description**: Indexed scatter to a UB pointer. Writes vector lanes to +irregular memory locations using per-lane element offsets. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `value` | `VRegType` | Source vector to scatter | +| `destination` | `PtrType` (ub) | UB destination pointer | +| `offsets` | `VRegType` | Per-lane element offsets (integer VMI vector) | +| `mask` | VMI mask | **Required.** Predicate mask gating lane participation | +| `pmode` | `str` or `None` | Optional predicate mode: `"merge"` keeps predicate-inactive lanes at their prior value; `"zero"` writes 0 | + +**Returns**: None (side-effect operation). + +**Example**: + +```python +g = pto.vmi.vgather(src_ptr, offsets, mask) +pto.vmi.vscatter(value, dst_ptr, offsets, mask) +``` + +--- + +## 14.9 Predicate construction + +VMI provides one surface API for creating predicate masks: +`create_mask(...)`. It covers both whole-vector prefix masks and grouped +prefix masks. + +### `pto.vmi.create_mask(active_lanes, *, size) -> MaskType` +### `pto.vmi.create_mask(active_lanes, *, size, group) -> MaskType` + +**Description**: Creates a prefix-style VMI mask where the first +`active_lanes` lanes are active and all remaining lanes are inactive. This is +the primary mask constructor for tail handling and partial-vector scenarios. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `active_lanes` | `IndexLike` | Number of active lanes in the prefix | +| `size` | `int` | Total logical lane count | +| `group` | `int` or `None` | When provided, creates a grouped prefix mask instead of a whole-vector prefix mask. The group size is inferred as `size / group` | + +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `mask` | VMI mask | Prefix mask with `active_lanes` active lanes, either across the whole logical vector or within each group | + +**Example**: + +```python +full_mask = pto.vmi.create_mask(64, size=64) +tail_mask = pto.vmi.create_mask(remained, size=64) +group_mask = pto.vmi.create_mask( + active_per_group, + size=128, + group=8, +) +``` + +**Constraints**: +- When `group` is present, `size` must be divisible by `group`. +- In grouped form, `group_size` is inferred as `size / group`. +- In grouped form, `active_lanes` must be ≤ inferred `group_size`. + +--- + +## 14.10 Data rearrangement + +Rearrangement instructions reorganize data between VMI vector registers +without touching memory. They are used to switch between interleaved and +deinterleaved data layouts. + +### `pto.vmi.vintlv(lhs, rhs, mask, *, pmode=None) -> (VRegType, VRegType)` + +**Description**: Interleave two logical vectors lane-by-lane and return the +result as a pair: `low` contains the interleaved lower half, `high` contains +the upper half. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `lhs` | `VRegType` | First source vector | +| `rhs` | `VRegType` | Second source vector | +| `mask` | VMI mask | **Required.** Predicate mask gating lane participation | +| `pmode` | `str` or `None` | Optional predicate mode: `"merge"` keeps predicate-inactive lanes at their prior value; `"zero"` writes 0 | +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `low` | `VRegType` | Interleaved lower half | +| `high` | `VRegType` | Interleaved upper half | + +--- + +### `pto.vmi.vdintlv(lhs, rhs, mask, *, pmode=None) -> (VRegType, VRegType)` + +**Description**: Deinterleave a previously interleaved vector pair. This is +the inverse of `vintlv`: it separates even-positioned and odd-positioned lanes +of the logical input stream into two output vectors. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `lhs` | `VRegType` | Lower half of the interleaved input | +| `rhs` | `VRegType` | Upper half of the interleaved input | +| `mask` | VMI mask | **Required.** Predicate mask gating lane participation | +| `pmode` | `str` or `None` | Optional predicate mode: `"merge"` keeps predicate-inactive lanes at their prior value; `"zero"` writes 0 | +**Returns**: + +| Return Value | Type | Description | +|--------------|------|-------------| +| `even` | `VRegType` | Lanes from even interleaved positions | +| `odd` | `VRegType` | Lanes from odd interleaved positions | + +**Example**: + +```python +lo, hi = pto.vmi.vintlv(src, src, mask) +even, odd = pto.vmi.vdintlv(lo, hi, mask) +``` + +**Constraints**: +- `lhs` and `rhs` must have the same type. +- The two returned vectors form one logical interleaved pair. Preserve their + order when passing them to subsequent ops. + +--- + +## 14.11 Result typing rules + +VMI infers result types when the output shape is unambiguous from the inputs. +Some ops still need a small hint such as `size` or `to_dtype`, but the +surface no longer asks you to spell the full result type manually. + +**Ops that infer their result type automatically**: + +- Same-shape elementwise binary and unary ops: inferred from the first vector + operand. +- Vector-scalar ops: inferred from the source vector. +- `vci` / `vbrc`: inferred from the typed scalar or vector input plus `size` + (and `group` for grouped broadcast). +- `vcmp` / `vcmps`: inferred from the `seed` mask. +- `vsel`: inferred from `true_value`. +- `vcadd`, `vcmax`, `vcmin`: inferred from the source vector and `group`. +- `vmull`: inferred from `a`, returning a `(low, high)` pair with the same 32-bit integer type. +- `vmula`: inferred from `acc`, preserving the accumulator type. +- `vdhist`, `vchist`: inferred from `acc`. +- `vgather`: inferred from the source element type and the `offsets` lanes. +- `vcvt` (when `to_dtype` is provided): inferred from the source lane count + and target element type. +- `vload` (when `size` is provided): inferred from the source pointer element + type and `size`. +- `vstore`: no result — side-effect only. +- `vscatter`: no result — side-effect only. +- `vintlv` / `vdintlv`: inferred from the input vector types. +- `vmull`: both low/high result types are inferred from the two matching input + vector types. + +**Ops that infer when given the right hint**: + +- `vload` with `dist_mode="unpack"` requires `to_dtype` to derive the widened + element type. +- `vcvt` requires `to_dtype`. +- `vinterpret_cast` requires `to_dtype`. +- `vgatherb` is inferred from the source pointer element type and the mask. + +--- + +## 14.12 Relationship to the top-level vector surface + +`pto.vmi` and the top-level PTODSL vector helpers (`pto.vadd`, `pto.vlds`, +`pto.vcvt`, etc.) are two distinct authoring surfaces that coexist in PTODSL. + +| Aspect | Top-level helpers | `pto.vmi` | +|--------|-------------------|-----------| +| Type system | `VRegType` / `MaskType` | `pto.vmi.vreg(...)` / `pto.vmi.mask(...)` | +| Naming | `pto.vadd`, `pto.vlds` | `pto.vmi.vadd`, `pto.vmi.vload` | +| Style | PTODSL vector programming model | Formal VMI instruction set | +| Predicate creation | `pto.make_mask`, `pto.pset_b32` | `pto.vmi.create_mask` | +| Return model | Varies by op | Consistent: returns new value or `(values, ...)` tuple | + +Key differences in practice: + +- `pto.vmi.vreg(...)` is distinct from `pto.vreg_type(...)`. +- `pto.vmi.mask(...)` is distinct from `pto.mask_type(...)`. +- `pto.vmi.vadd(...)` is a formal VMI call, not a synonym for `pto.vadd(...)`. +- VMI operations return values rather than writing to destination buffers. + +Choose one surface intentionally inside a given sub-kernel or helper region, +and keep the authored style consistent. Mixing both surfaces in the same +region is possible but makes the IR intent harder to follow. + +--- + +## 14.13 Full example: elementwise vector pipeline + +The following example shows a complete VMI pipeline: load, compute under mask, +and store back. + +```python +from ptodsl import pto + +@pto.jit( + name="vmi_elementwise", + target="a5", + backend="vpto", + mode="explicit", + kernel_kind="vector", + insert_sync=False, +) +def vmi_elementwise( + src_ptr: pto.ptr(pto.f32, "ub"), + dst_ptr: pto.ptr(pto.f32, "ub"), + count: pto.i32, + scale: pto.f32, +): + full_mask = pto.vmi.create_mask(64, size=64) + + lhs = pto.vmi.vload(src_ptr, 0, size=64) + rhs = pto.vmi.vload(src_ptr, 64, size=64) + + summed = pto.vmi.vadd(lhs, rhs, full_mask) + scaled = pto.vmi.vmuls(summed, scale, full_mask) + activated = pto.vmi.vrelu(scaled, full_mask) + + pto.vmi.vstore(activated, dst_ptr, 0, full_mask) +``` + +--- + +## 14.14 VMI instruction quick reference + +| Category | Instructions | +|----------|-------------| +| Types | `vreg`, `mask` | +| Load / Store | `vload`, `vstore` | +| Index / Broadcast | `vci`, `vbrc` | +| Binary vector-vector | `vadd`, `vsub`, `vmul`, `vdiv`, `vmax`, `vmin`, `vand`, `vor`, `vxor`, `vshl`, `vshr` | +| Unary vector | `vabs`, `vneg`, `vrelu`, `vexp`, `vln`, `vsqrt`, `vnot` | +| Vector-scalar | formal `pto.vmi`: `vadds`, `vmuls`, `vmaxs`, `vmins`, `vshls`, `vshrs`; DSL convenience: `vsubs`, `vands`, `vors`, `vxors` | +| Compare / Select | `vcmp`, `vcmps`, `vsel`, `vselr` | +| Reduction | `vcadd`, `vcmax`, `vcmin` | +| Conversion | `vcvt`, `vinterpret_cast` | +| SFU / Fused | `vexpdif`, `vaxpy`, `vlrelu`, `vprelu`, `vmull`, `vmula` | +| Histogram | `vchist`, `vdhist` | +| Indexed memory | `vgather`, `vgatherb`, `vscatter` | +| Predicate construction | `create_mask` | +| Data rearrangement | `vintlv`, `vdintlv` | + +All formally listed VMI instructions above are members of the `pto.vmi` +namespace. Entries explicitly labeled as DSL convenience spellings are PTODSL +authoring helpers and are not part of the formal `pto.vmi` v0.1 inventory. diff --git a/ptodsl/ptoas/_launcher.py b/ptodsl/ptoas/_launcher.py index dc7868f99e..dc724b7ac9 100644 --- a/ptodsl/ptoas/_launcher.py +++ b/ptodsl/ptoas/_launcher.py @@ -10,7 +10,10 @@ from __future__ import annotations +import ctypes import os +import subprocess +import shutil import sys from pathlib import Path from typing import NoReturn @@ -36,6 +39,36 @@ def _has_cli_option(argv: list[str], option: str) -> bool: return False +def _resolve_wrapper_path() -> Path: + argv0 = Path(sys.argv[0]) + if argv0.exists(): + return argv0.resolve() + + found = shutil.which(argv0.name or "ptoas") + if found: + return Path(found).resolve() + + raise SystemExit(f"unable to locate the installed ptoas wrapper: {sys.argv[0]}") + + +def _resolve_shared_module_path(package_root: Path, runtime_root: Path, wrapper: Path) -> Path: + candidates = [ + package_root.parent / "pto" / "ptoas.so", + wrapper.parent / "ptoas.so", + runtime_root / "lib" / "ptoas.so", + runtime_root / "pto" / "ptoas.so", + ] + + for candidate in candidates: + if candidate.is_file() and candidate.stat().st_size > 0: + return candidate + + raise SystemExit( + "wheel/runtime is missing the packaged shared module: expected pto.ptoas " + "or a local ptoas.so next to the wrapper/install tree" + ) + + def _resolve_runtime_root(package_root: Path) -> Path: runtime_root = package_root / "_runtime" if runtime_root.exists(): @@ -46,35 +79,107 @@ def _resolve_runtime_root(package_root: Path) -> Path: return package_root.parent.parent / "install" +def _iter_runtime_library_deps(shared_module: Path, runtime_lib_dir: Path) -> list[Path]: + if not runtime_lib_dir.is_dir(): + return [] + + env = os.environ.copy() + current = env.get("LD_LIBRARY_PATH", "") + rendered = str(runtime_lib_dir) + env["LD_LIBRARY_PATH"] = os.pathsep.join( + [rendered] + [part for part in current.split(os.pathsep) if part] + ) + + output = subprocess.check_output( + ["ldd", str(shared_module)], + text=True, + stderr=subprocess.DEVNULL, + env=env, + ) + deps: list[Path] = [] + for line in output.splitlines(): + line = line.strip() + if "=>" in line: + candidate = line.split("=>", 1)[1].strip().split(" ", 1)[0] + else: + candidate = line.split(" ", 1)[0] + if not candidate.startswith("/"): + continue + dep_path = Path(candidate) + try: + if runtime_lib_dir.resolve() not in dep_path.resolve().parents: + continue + except FileNotFoundError: + continue + deps.append(dep_path) + return deps + + +def _preload_runtime_libraries(shared_module: Path, runtime_lib_dir: Path) -> None: + runtime_lib_dir = runtime_lib_dir.resolve() + if not runtime_lib_dir.is_dir(): + return + + loaded: set[Path] = set() + visiting: set[Path] = set() + + def visit(path: Path) -> None: + resolved = path.resolve() + if resolved in loaded or resolved in visiting: + return + visiting.add(resolved) + try: + for dep in _iter_runtime_library_deps(resolved, runtime_lib_dir): + visit(dep) + ctypes.CDLL(str(resolved), mode=getattr(ctypes, "RTLD_GLOBAL", 0)) + loaded.add(resolved) + finally: + visiting.discard(resolved) + + for dep in _iter_runtime_library_deps(shared_module, runtime_lib_dir): + visit(dep) + + +def _load_shared_entrypoint(shared_module: Path, runtime_lib_dir: Path): + _preload_runtime_libraries(shared_module, runtime_lib_dir) + library = ctypes.CDLL(str(shared_module), mode=getattr(ctypes, "RTLD_GLOBAL", 0)) + entrypoint = library.ptoas_entrypoint + entrypoint.argtypes = [ctypes.c_int, ctypes.POINTER(ctypes.c_char_p)] + entrypoint.restype = ctypes.c_int + return entrypoint + + def main() -> NoReturn: package_root = Path(__file__).resolve().parent runtime_root = _resolve_runtime_root(package_root) - binary = runtime_root / "bin" / "ptoas" - if not binary.is_file(): - raise SystemExit( - f"wheel runtime is missing the packaged ptoas binary: {binary}" - ) + wrapper = _resolve_wrapper_path() + shared_module = _resolve_shared_module_path(package_root, runtime_root, wrapper) python_root = package_root.parent if runtime_root.name == "_runtime" else runtime_root tileops_dir = runtime_root / "share" / "ptoas" / "TileOps" env = os.environ.copy() env["PTOAS_HOME"] = str(runtime_root) - env["PTOAS_BIN"] = str(binary) + env["PTOAS_BIN"] = str(wrapper) env["PTOAS_TILEOPS_DIR"] = str(tileops_dir) - _prepend_env_path(env, "PATH", binary.parent) + _prepend_env_path(env, "PATH", wrapper.parent) _prepend_env_path(env, "PYTHONPATH", python_root) _prepend_env_path(env, "LD_LIBRARY_PATH", runtime_root / "lib") _prepend_env_path(env, "DYLD_LIBRARY_PATH", runtime_root / "lib") + os.environ.update(env) - argv = [str(binary)] + argv = [str(wrapper)] user_args = sys.argv[1:] if not _has_cli_option(user_args, "--tilelang-path"): argv.extend(["--tilelang-path", str(tileops_dir)]) if not _has_cli_option(user_args, "--tilelang-pkg-path"): argv.extend(["--tilelang-pkg-path", str(python_root)]) argv.extend(user_args) - os.execvpe(str(binary), argv, env) + + entrypoint = _load_shared_entrypoint(shared_module, runtime_root / "lib") + argv_bytes = [os.fsencode(arg) for arg in argv] + c_argv = (ctypes.c_char_p * len(argv_bytes))(*argv_bytes) + raise SystemExit(entrypoint(len(argv_bytes), c_argv)) if __name__ == "__main__": diff --git a/ptodsl/ptodsl/_ast_rewrite.py b/ptodsl/ptodsl/_ast_rewrite.py index be766b58f5..1f5a94c4ac 100644 --- a/ptodsl/ptodsl/_ast_rewrite.py +++ b/ptodsl/ptodsl/_ast_rewrite.py @@ -40,9 +40,12 @@ def rewrite_jit_function(fn): function_def.decorator_list = [] closure_vars = inspect.getclosurevars(fn) + static_env = dict(fn.__globals__) + static_env.update(closure_vars.nonlocals) _inject_closure_defaults(function_def, closure_vars.nonlocals) _sanitize_signature_for_exec(function_def) - rewriter = _ControlFlowRewriter() + function_def = _ConditionalExpressionNormalizer().visit(function_def) + rewriter = _ControlFlowRewriter(static_env) function_def.body = rewriter.rewrite_block(function_def.body, live_after=set()) tree = ast.Module(body=[function_def], type_ignores=[]) ast.fix_missing_locations(tree) @@ -146,12 +149,65 @@ def _sanitize_signature_for_exec(function_def): function_def.returns = None +def _is_normalizable_ifexp_assign_target(node) -> bool: + return isinstance(node, ast.Name) + + +class _ConditionalExpressionNormalizer(ast.NodeTransformer): + """Normalize assign-form ``IfExp`` into statement ``if`` before rewrite.""" + + def visit_Assign(self, node): + node = self.generic_visit(node) + if not isinstance(node.value, ast.IfExp): + return node + if not node.targets or not all(_is_normalizable_ifexp_assign_target(target) for target in node.targets): + return node + return self._normalize_ifexp_assignment(node, node.value) + + def visit_AnnAssign(self, node): + node = self.generic_visit(node) + if node.value is None or not isinstance(node.value, ast.IfExp): + return node + if not _is_normalizable_ifexp_assign_target(node.target): + return node + return self._normalize_ifexp_assignment(node, node.value) + + def _normalize_ifexp_assignment(self, stmt, value): + then_stmt = copy.deepcopy(stmt) + then_stmt.value = value.body + else_stmt = copy.deepcopy(stmt) + else_stmt.value = value.orelse + if_stmt = ast.If( + test=value.test, + body=[then_stmt], + orelse=[else_stmt], + ) + return ast.copy_location(self.generic_visit(if_stmt), stmt) + + @dataclass(frozen=True) class _NameInfo: loads: set[str] stores: set[str] +@dataclass(frozen=True, order=True) +class _SubscriptSlot: + base: str + index: int + + @property + def display(self) -> str: + return f"{self.base}[{self.index}]" + + +@dataclass(frozen=True) +class _SlotInfo: + loads: set[_SubscriptSlot] + stores: set[_SubscriptSlot] + invalid_stores: tuple[str, ...] = () + + class _NameInfoVisitor(ast.NodeVisitor): def __init__(self): self.loads = set() @@ -201,6 +257,18 @@ def visit_ClassDef(self, node): self.visit(keyword.value) self.loads.update(_class_body_free_vars(node)) + def visit_ListComp(self, node): + self._visit_comprehension(node.generators, (node.elt,)) + + def visit_SetComp(self, node): + self._visit_comprehension(node.generators, (node.elt,)) + + def visit_GeneratorExp(self, node): + self._visit_comprehension(node.generators, (node.elt,)) + + def visit_DictComp(self, node): + self._visit_comprehension(node.generators, (node.key, node.value)) + def _visit_arguments_defaults(self, args): for default in args.defaults: self.visit(default) @@ -208,6 +276,21 @@ def _visit_arguments_defaults(self, args): if default is not None: self.visit(default) + def _visit_comprehension(self, generators, result_nodes): + bound = set() + for generator in generators: + self._visit_comprehension_expr(generator.iter, bound) + bound |= _target_stores(generator.target) + for if_node in generator.ifs: + self._visit_comprehension_expr(if_node, bound) + for result_node in result_nodes: + self._visit_comprehension_expr(result_node, bound) + + def _visit_comprehension_expr(self, node, bound): + info = _name_info(node) + self.loads.update(info.loads - set(bound)) + self.stores.update(info.stores - set(bound)) + def _name_info(node) -> _NameInfo: visitor = _NameInfoVisitor() @@ -219,6 +302,250 @@ def _name_info(node) -> _NameInfo: return _NameInfo(visitor.loads, visitor.stores) +class _SlotInfoVisitor(ast.NodeVisitor): + def __init__(self, static_env, static_iters=None): + self._static_env = static_env + self._static_iters = dict(static_iters or {}) + self.loads = set() + self.stores = set() + self.invalid_stores = [] + + def visit_Subscript(self, node): + if isinstance(node.ctx, ast.Load): + self.loads.update(_resolve_subscript_slots(node, self._static_iters, require_static=False)) + return + if isinstance(node.ctx, (ast.Store, ast.Del)): + slots = _resolve_subscript_slots(node, self._static_iters, require_static=True) + if slots: + self.stores.update(slots) + else: + self.invalid_stores.append(_unsupported_subscript_store_message(node)) + return + self.generic_visit(node) + + def visit_AugAssign(self, node): + if isinstance(node.target, ast.Subscript): + slots = _resolve_subscript_slots(node.target, self._static_iters, require_static=True) + if slots: + self.loads.update(slots) + self.stores.update(slots) + else: + self.invalid_stores.append(_unsupported_subscript_store_message(node.target)) + else: + self.visit(node.target) + self.visit(node.value) + + def visit_For(self, node): + if _is_pto_attr_call(node.iter, "static_range") and isinstance(node.target, ast.Name): + values = _try_eval_static_range(node.iter, self._static_env) + if values is None: + for stmt in node.body: + self.visit(stmt) + for stmt in node.orelse: + self.visit(stmt) + return + old = self._static_iters.get(node.target.id) + self._static_iters[node.target.id] = values + try: + for stmt in node.body: + self.visit(stmt) + finally: + if old is None: + self._static_iters.pop(node.target.id, None) + else: + self._static_iters[node.target.id] = old + for stmt in node.orelse: + self.visit(stmt) + return + self.generic_visit(node) + + def visit_FunctionDef(self, node): + return + + def visit_AsyncFunctionDef(self, node): + return + + def visit_Lambda(self, node): + return + + def visit_ClassDef(self, node): + return + + +def _slot_info(node, static_env, static_iters=None) -> _SlotInfo: + visitor = _SlotInfoVisitor(static_env, static_iters) + if isinstance(node, list): + for item in node: + visitor.visit(item) + else: + visitor.visit(node) + return _SlotInfo(visitor.loads, visitor.stores, tuple(visitor.invalid_stores)) + + +def _slot_live_before_block(stmts, live_after, static_env, static_iters=None) -> set[_SubscriptSlot]: + live = set(live_after) + for stmt in reversed(stmts): + live = _slot_live_before_stmt(stmt, live, static_env, static_iters or {}) + return live + + +def _slot_live_before_stmt(stmt, live_after, static_env, static_iters) -> set[_SubscriptSlot]: + if isinstance(stmt, ast.If): + test_info = _slot_info(stmt.test, static_env, static_iters) + return ( + set(test_info.loads) + | _slot_live_before_block(stmt.body, live_after, static_env, static_iters) + | _slot_live_before_block(stmt.orelse, live_after, static_env, static_iters) + ) + if isinstance(stmt, ast.For): + if _is_pto_attr_call(stmt.iter, "static_range") and isinstance(stmt.target, ast.Name): + values = _try_eval_static_range(stmt.iter, static_env) + if values is not None: + next_static_iters = dict(static_iters) + next_static_iters[stmt.target.id] = values + return ( + _slot_live_before_block(stmt.body, live_after, static_env, next_static_iters) + | _slot_live_before_block(stmt.orelse, live_after, static_env, static_iters) + ) + iter_info = _slot_info(stmt.iter, static_env, static_iters) + body_info = _slot_info(stmt.body, static_env, static_iters) + orelse_info = _slot_info(stmt.orelse, static_env, static_iters) + assigned = body_info.stores | orelse_info.stores + return ( + (set(live_after) - assigned) + | set(iter_info.loads) + | _slot_live_before_block(stmt.body, set(), static_env, static_iters) + | _slot_live_before_block(stmt.orelse, set(), static_env, static_iters) + ) + info = _slot_info(stmt, static_env, static_iters) + live = _kill_slots_for_assigned_bases(live_after, stmt) + return (set(live) - info.stores) | info.loads + + +def _read_before_assignment_slots(stmts, static_env, static_iters=None) -> set[_SubscriptSlot]: + return _slot_live_before_block(stmts, set(), static_env, static_iters) + + +def _kill_slots_for_assigned_bases(slots, stmt) -> set[_SubscriptSlot]: + assigned_bases = _assigned_name_targets(stmt) + if not assigned_bases: + return set(slots) + return { + slot + for slot in slots + if slot.base not in assigned_bases + } + + +def _assigned_name_targets(stmt) -> set[str]: + if isinstance(stmt, ast.Assign): + names = set() + for target in stmt.targets: + names.update(_simple_name_targets(target)) + return names + if isinstance(stmt, ast.AnnAssign): + return _simple_name_targets(stmt.target) + if isinstance(stmt, (ast.For, ast.AsyncFor)): + return _simple_name_targets(stmt.target) + if isinstance(stmt, (ast.With, ast.AsyncWith)): + names = set() + for item in stmt.items: + if item.optional_vars is not None: + names.update(_simple_name_targets(item.optional_vars)) + return names + return set() + + +def _simple_name_targets(target) -> set[str]: + if isinstance(target, ast.Name): + return {target.id} + if isinstance(target, (ast.Tuple, ast.List)): + names = set() + for elt in target.elts: + names.update(_simple_name_targets(elt)) + return names + return set() + + +def _resolve_subscript_slots(node, static_iters, *, require_static) -> set[_SubscriptSlot]: + if not isinstance(node.value, ast.Name): + return set() + index_values = _static_index_values(node.slice, static_iters) + if index_values is None: + return set() + return { + _SubscriptSlot(node.value.id, index) + for index in index_values + } + + +def _static_index_values(node, static_iters): + if isinstance(node, ast.Constant) and isinstance(node.value, int) and not isinstance(node.value, bool): + return (node.value,) + if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub): + values = _static_index_values(node.operand, static_iters) + if values is not None and len(values) == 1: + return (-values[0],) + if isinstance(node, ast.Name) and node.id in static_iters: + return tuple(static_iters[node.id]) + return None + + +def _unsupported_subscript_store_message(node) -> str: + try: + text = ast.unparse(node) + except Exception: + text = "" + return ( + "ast_rewrite=True only supports static subscript carry stores of the form " + f"simple_name[static_int_or_static_range_iv]; got {text!r}" + ) + + +def _try_eval_static_range(call, static_env): + if not _is_pto_attr_call(call, "static_range") or call.keywords: + return None + try: + values = [_eval_static_int(arg, static_env) for arg in call.args] + except PTODSLAstRewriteError: + return None + if len(values) == 1: + return tuple(range(values[0])) + if len(values) == 2: + return tuple(range(values[0], values[1])) + if len(values) == 3: + return tuple(range(values[0], values[1], values[2])) + return None + + +def _eval_static_int(node, static_env) -> int: + if isinstance(node, ast.Constant) and isinstance(node.value, int) and not isinstance(node.value, bool): + return node.value + if isinstance(node, ast.Name): + value = static_env.get(node.id, _MISSING_GLOBAL) + if isinstance(value, int) and not isinstance(value, bool): + return value + raise PTODSLAstRewriteError("static value is not an integer") + if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.UAdd): + return +_eval_static_int(node.operand, static_env) + if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub): + return -_eval_static_int(node.operand, static_env) + if isinstance(node, ast.BinOp): + lhs = _eval_static_int(node.left, static_env) + rhs = _eval_static_int(node.right, static_env) + if isinstance(node.op, ast.Add): + return lhs + rhs + if isinstance(node.op, ast.Sub): + return lhs - rhs + if isinstance(node.op, ast.Mult): + return lhs * rhs + if isinstance(node.op, ast.FloorDiv): + return lhs // rhs + if isinstance(node.op, ast.Mod): + return lhs % rhs + raise PTODSLAstRewriteError("unsupported static integer expression") + + class _ScopeBindingVisitor(ast.NodeVisitor): def __init__(self): self.stores = set() @@ -241,6 +568,18 @@ def visit_Lambda(self, node): def visit_ClassDef(self, node): self.stores.add(node.name) + def visit_ListComp(self, node): + self._visit_comprehension(node.generators, (node.elt,)) + + def visit_SetComp(self, node): + self._visit_comprehension(node.generators, (node.elt,)) + + def visit_GeneratorExp(self, node): + self._visit_comprehension(node.generators, (node.elt,)) + + def visit_DictComp(self, node): + self._visit_comprehension(node.generators, (node.key, node.value)) + def visit_Global(self, node): self.globals.update(node.names) @@ -257,6 +596,14 @@ def visit_ImportFrom(self, node): continue self.stores.add(alias.asname or alias.name) + def _visit_comprehension(self, generators, result_nodes): + for generator in generators: + self.visit(generator.iter) + for if_node in generator.ifs: + self.visit(if_node) + for result_node in result_nodes: + self.visit(result_node) + def _argument_names(args) -> set[str]: names = {arg.arg for arg in list(args.posonlyargs) + list(args.args) + list(args.kwonlyargs)} @@ -370,8 +717,65 @@ def _name(name: str, ctx=ast.Load()): return ast.Name(id=name, ctx=ctx) +def _slot_subscript(slot: _SubscriptSlot, ctx=ast.Load()): + return ast.Subscript( + value=_name(slot.base), + slice=ast.Constant(slot.index), + ctx=ctx, + ) + + +def _map_subscript(map_name: str, index: int, ctx=ast.Load()): + return ast.Subscript( + value=_name(map_name), + slice=ast.Constant(index), + ctx=ctx, + ) + + +class _SlotCarryRewriter(ast.NodeTransformer): + def __init__(self, slot_maps, static_env, static_iters=None): + self._slot_maps = slot_maps + self._static_env = static_env + self._static_iters = dict(static_iters or {}) + + def visit_For(self, node): + if _is_pto_attr_call(node.iter, "static_range") and isinstance(node.target, ast.Name): + values = _try_eval_static_range(node.iter, self._static_env) + old = self._static_iters.get(node.target.id) + if values is not None: + self._static_iters[node.target.id] = values + try: + node.body = [self.visit(stmt) for stmt in node.body] + finally: + if values is not None: + if old is None: + self._static_iters.pop(node.target.id, None) + else: + self._static_iters[node.target.id] = old + node.orelse = [self.visit(stmt) for stmt in node.orelse] + return node + return self.generic_visit(node) + + def visit_Subscript(self, node): + slots = _resolve_subscript_slots(node, self._static_iters, require_static=False) + if slots and len({slot.base for slot in slots}) == 1: + base = next(iter(slots)).base + if base in self._slot_maps and slots <= set(self._slot_maps[base]["slots"]): + return ast.copy_location( + ast.Subscript( + value=_name(self._slot_maps[base]["map_name"]), + slice=copy.deepcopy(node.slice), + ctx=node.ctx, + ), + node, + ) + return self.generic_visit(node) + + class _ControlFlowRewriter: - def __init__(self): + def __init__(self, static_env=None): + self._static_env = dict(static_env or {}) self._counter = 0 def _fresh(self, prefix: str) -> str: @@ -379,31 +783,47 @@ def _fresh(self, prefix: str) -> str: self._counter += 1 return value - def rewrite_block(self, stmts, *, live_after, allow_loop_control=False): + def rewrite_block(self, stmts, *, live_after, live_after_slots=None, allow_loop_control=False, static_iters=None): rewritten_reversed = [] live = set(live_after) + live_slots = set(live_after_slots or ()) + static_iters = dict(static_iters or {}) for stmt in reversed(stmts): + # Compute liveness from the authored AST before rewrite_stmt mutates + # sibling statements in-place, otherwise later rewrites can pollute + # earlier live-after analysis. + live_before = _live_before_stmt(stmt, live) + live_before_slots = _slot_live_before_stmt(stmt, live_slots, self._static_env, static_iters) rewritten = self.rewrite_stmt( stmt, live_after=live, + live_after_slots=live_slots, allow_loop_control=allow_loop_control, + static_iters=static_iters, ) rewritten_reversed[:0] = rewritten - live = _live_before_stmt(stmt, live) + live = live_before + live_slots = live_before_slots return rewritten_reversed - def rewrite_stmt(self, stmt, *, live_after, allow_loop_control=False): + def rewrite_stmt(self, stmt, *, live_after, live_after_slots=None, allow_loop_control=False, static_iters=None): + live_after_slots = set(live_after_slots or ()) + static_iters = dict(static_iters or {}) if isinstance(stmt, ast.If): return self._rewrite_if( stmt, live_after=live_after, + live_after_slots=live_after_slots, allow_loop_control=allow_loop_control, + static_iters=static_iters, ) if isinstance(stmt, ast.For): return self._rewrite_for( stmt, live_after=live_after, + live_after_slots=live_after_slots, allow_loop_control=allow_loop_control, + static_iters=static_iters, ) if isinstance(stmt, (ast.Break, ast.Continue)): if allow_loop_control: @@ -413,16 +833,22 @@ def rewrite_stmt(self, stmt, *, live_after, allow_loop_control=False): self._rewrite_nested( stmt, live_after=live_after, + live_after_slots=live_after_slots, allow_loop_control=allow_loop_control, + static_iters=static_iters, ) ] - def _rewrite_nested(self, stmt, *, live_after, allow_loop_control=False): + def _rewrite_nested(self, stmt, *, live_after, live_after_slots=None, allow_loop_control=False, static_iters=None): + live_after_slots = set(live_after_slots or ()) + static_iters = dict(static_iters or {}) if isinstance(stmt, (ast.FunctionDef, ast.AsyncFunctionDef)): stmt.body = self.rewrite_block( stmt.body, live_after=set(), + live_after_slots=set(), allow_loop_control=False, + static_iters={}, ) return stmt if isinstance(stmt, (ast.Lambda, ast.ClassDef)): @@ -435,14 +861,18 @@ def _rewrite_nested(self, stmt, *, live_after, allow_loop_control=False): self.rewrite_block( value, live_after=live_after, + live_after_slots=live_after_slots, allow_loop_control=allow_loop_control, + static_iters=static_iters, ), ) elif isinstance(value, ast.AST): self._rewrite_nested( value, live_after=live_after, + live_after_slots=live_after_slots, allow_loop_control=allow_loop_control, + static_iters=static_iters, ) elif isinstance(value, list): for item in value: @@ -450,27 +880,45 @@ def _rewrite_nested(self, stmt, *, live_after, allow_loop_control=False): self._rewrite_nested( item, live_after=live_after, + live_after_slots=live_after_slots, allow_loop_control=allow_loop_control, + static_iters=static_iters, ) return stmt - def _rewrite_if(self, stmt, *, live_after, allow_loop_control=False): + def _rewrite_if(self, stmt, *, live_after, live_after_slots=None, allow_loop_control=False, static_iters=None): + live_after_slots = set(live_after_slots or ()) + static_iters = dict(static_iters or {}) if _is_pto_attr_call(stmt.test, "const_expr"): stmt.body = self.rewrite_block( stmt.body, live_after=live_after, + live_after_slots=live_after_slots, allow_loop_control=allow_loop_control, + static_iters=static_iters, ) stmt.orelse = self.rewrite_block( stmt.orelse, live_after=live_after, + live_after_slots=live_after_slots, allow_loop_control=allow_loop_control, + static_iters=static_iters, ) return [stmt] cond_name = self._fresh("cond") then_info = _name_info(stmt.body) else_info = _name_info(stmt.orelse) + assigned_slots = ( + _slot_info(stmt.body, self._static_env, static_iters).stores + | _slot_info(stmt.orelse, self._static_env, static_iters).stores + ) + if live_after_slots & assigned_slots: + slots = ", ".join(slot.display for slot in sorted(live_after_slots & assigned_slots)) + raise PTODSLAstRewriteError( + "ast_rewrite=True does not support automatic branch merges for static subscript slots yet; " + f"rewrite {slots} with explicit scalar temporaries" + ) assigned_any = then_info.stores | else_info.stores merge_names = tuple(sorted(live_after & assigned_any)) old_value_names = { @@ -483,12 +931,16 @@ def _rewrite_if(self, stmt, *, live_after, allow_loop_control=False): then_body = self.rewrite_block( stmt.body, live_after=branch_live_after, + live_after_slots=live_after_slots, allow_loop_control=False, + static_iters=static_iters, ) else_body = self.rewrite_block( stmt.orelse, live_after=branch_live_after, + live_after_slots=live_after_slots, allow_loop_control=False, + static_iters=static_iters, ) trace_time_if = ast.If( test=_name(cond_name), @@ -612,17 +1064,28 @@ def _branch_assign(self, branch_name, names, *, old_value_names, assigned_names) ) ) - def _rewrite_for(self, stmt, *, live_after, allow_loop_control=False): + def _rewrite_for(self, stmt, *, live_after, live_after_slots=None, allow_loop_control=False, static_iters=None): + live_after_slots = set(live_after_slots or ()) + static_iters = dict(static_iters or {}) if _is_pto_attr_call(stmt.iter, "static_range"): + next_static_iters = dict(static_iters) + if isinstance(stmt.target, ast.Name): + values = _try_eval_static_range(stmt.iter, self._static_env) + if values is not None: + next_static_iters[stmt.target.id] = values stmt.body = self.rewrite_block( stmt.body, live_after=live_after, + live_after_slots=live_after_slots, allow_loop_control=True, + static_iters=next_static_iters, ) stmt.orelse = self.rewrite_block( stmt.orelse, live_after=live_after, + live_after_slots=live_after_slots, allow_loop_control=allow_loop_control, + static_iters=static_iters, ) return [stmt] @@ -638,21 +1101,58 @@ def _rewrite_for(self, stmt, *, live_after, allow_loop_control=False): start, stop, step = _range_triplet(stmt.iter) body_info = _name_info(stmt.body) + body_slot_info = _slot_info(stmt.body, self._static_env, static_iters) + if body_slot_info.invalid_stores: + raise PTODSLAstRewriteError(body_slot_info.invalid_stores[0]) reads_before = _read_before_assignment_names(stmt.body) + slot_reads_before = _read_before_assignment_slots(stmt.body, self._static_env, static_iters) assigned_live_after = body_info.stores & set(live_after) + assigned_slots_live_after = body_slot_info.stores & set(live_after_slots) loop_carried = tuple(sorted(body_info.stores & reads_before)) + loop_carried_slots = tuple(sorted(body_slot_info.stores & slot_reads_before)) unsupported_last_values = sorted(assigned_live_after - set(loop_carried)) if unsupported_last_values: raise PTODSLAstRewriteError( "ast_rewrite=True runtime for-loops cannot expose last-iteration-only values yet; " f"use explicit pto.for_(...).carry(...) for {unsupported_last_values}" ) + unsupported_last_slots = sorted(assigned_slots_live_after - set(loop_carried_slots)) + if unsupported_last_slots: + slots = [slot.display for slot in unsupported_last_slots] + raise PTODSLAstRewriteError( + "ast_rewrite=True runtime for-loops cannot expose last-iteration-only static subscript values yet; " + f"use explicit scalar temporaries for {slots}" + ) loop_name = self._fresh("loop") loop_live_after = set(live_after) | set(loop_carried) - body = self.rewrite_block(stmt.body, live_after=loop_live_after) + loop_live_after_slots = set(live_after_slots) | set(loop_carried_slots) + body = self.rewrite_block( + stmt.body, + live_after=loop_live_after, + live_after_slots=loop_live_after_slots, + static_iters=static_iters, + ) - if loop_carried: + slot_carry_names = { + slot: self._fresh(f"slot_{slot.base}_{slot.index}") + for slot in loop_carried_slots + } + slot_maps = {} + for slot in loop_carried_slots: + slot_maps.setdefault(slot.base, {"map_name": self._fresh(f"slot_{slot.base}_map"), "slots": []}) + slot_maps[slot.base]["slots"].append(slot) + for data in slot_maps.values(): + data["slots"] = tuple(sorted(data["slots"])) + + if loop_carried or loop_carried_slots: + slot_initializers = [ + ast.Assign( + targets=[_name(slot_carry_names[slot], ast.Store())], + value=_slot_subscript(slot), + ) + for slot in loop_carried_slots + ] setup = ast.Assign( targets=[_name(loop_name, ast.Store())], value=ast.Call( @@ -669,6 +1169,9 @@ def _rewrite_for(self, stmt, *, live_after, allow_loop_control=False): keywords=[ ast.keyword(arg=name, value=_name(name)) for name in loop_carried + ] + [ + ast.keyword(arg=slot_carry_names[slot], value=_name(slot_carry_names[slot])) + for slot in loop_carried_slots ], ), ) @@ -685,7 +1188,37 @@ def _rewrite_for(self, stmt, *, live_after, allow_loop_control=False): ) for name in loop_carried ) + prologue.extend( + ast.Assign( + targets=[_name(slot_carry_names[slot], ast.Store())], + value=ast.Attribute(value=_name(loop_name), attr=slot_carry_names[slot], ctx=ast.Load()), + ) + for slot in loop_carried_slots + ) + prologue.extend( + ast.Assign( + targets=[_name(data["map_name"], ast.Store())], + value=ast.Dict( + keys=[ast.Constant(slot.index) for slot in data["slots"]], + values=[_name(slot_carry_names[slot]) for slot in data["slots"]], + ), + ) + for data in slot_maps.values() + ) + if loop_carried_slots: + body = [ + _SlotCarryRewriter(slot_maps, self._static_env, static_iters).visit(stmt) + for stmt in body + ] + slot_epilogue = [ + ast.Assign( + targets=[_name(slot_carry_names[slot], ast.Store())], + value=_map_subscript(slot_maps[slot.base]["map_name"], slot.index), + ) + for slot in loop_carried_slots + ] body = prologue + body + [ + *slot_epilogue, ast.Expr( value=ast.Call( func=ast.Attribute(value=_name(loop_name), attr="update", ctx=ast.Load()), @@ -693,6 +1226,9 @@ def _rewrite_for(self, stmt, *, live_after, allow_loop_control=False): keywords=[ ast.keyword(arg=name, value=_name(name)) for name in loop_carried + ] + [ + ast.keyword(arg=slot_carry_names[slot], value=_name(slot_carry_names[slot])) + for slot in loop_carried_slots ], ) ) @@ -702,7 +1238,7 @@ def _rewrite_for(self, stmt, *, live_after, allow_loop_control=False): body=body or [ast.Pass()], type_comment=None, ) - result = [ast.copy_location(setup, stmt), ast.copy_location(with_stmt, stmt)] + result = slot_initializers + [ast.copy_location(setup, stmt), ast.copy_location(with_stmt, stmt)] for name in loop_carried: result.append( ast.Assign( @@ -714,6 +1250,17 @@ def _rewrite_for(self, stmt, *, live_after, allow_loop_control=False): ), ) ) + for slot in loop_carried_slots: + result.append( + ast.Assign( + targets=[_slot_subscript(slot, ast.Store())], + value=ast.Call( + func=ast.Attribute(value=_name(loop_name), attr="final", ctx=ast.Load()), + args=[ast.Constant(slot_carry_names[slot])], + keywords=[], + ), + ) + ) return result with_stmt = ast.With( diff --git a/ptodsl/ptodsl/_control_flow.py b/ptodsl/ptodsl/_control_flow.py index 8ff6acef14..c71f1f24a0 100644 --- a/ptodsl/ptodsl/_control_flow.py +++ b/ptodsl/ptodsl/_control_flow.py @@ -23,6 +23,7 @@ from ._bootstrap import make_context # noqa: F401 from ._runtime_index_ops import coerce_runtime_index +from ._scalar_coercion import coerce_scalar_to_type from ._surface_types import const_expr from ._tracing.active import current_session from ._surface_values import unwrap_surface_value, wrap_like_surface_value, wrap_surface_value @@ -415,9 +416,10 @@ def _assign_branch_values(self, kwargs): order = tuple(kwargs.keys()) for name, value in kwargs.items(): raw_value = unwrap_surface_value(value) - if not hasattr(raw_value, "type"): + if not hasattr(raw_value, "type") and not _is_branch_assign_literal(raw_value): raise TypeError( - "br.assign(...) expects PTO runtime values or authored surface values; " + "br.assign(...) expects PTO runtime values, authored surface values, " + "or Python scalar literals that can be inferred from the opposite branch; " f"'{name}' received {type(value).__name__}" ) raw_values[name] = raw_value @@ -489,22 +491,37 @@ def _validate_merge_spec(self): raise RuntimeError("br.assign(...) names must match across branches; " + "; ".join(pieces)) order = then_assignment["order"] + resolved_then_values = {} + resolved_else_values = {} result_types = [] for name in order: then_value = then_assignment["raw_values"][name] else_value = else_assignment["raw_values"][name] + then_value, else_value = _reconcile_branch_assignment_values( + name, + then_value, + else_value, + ) if then_value.type != else_value.type: raise RuntimeError( f"br.assign(...) type mismatch for '{name}': " f"then branch yields {then_value.type}, else branch yields {else_value.type}" ) + resolved_then_values[name] = then_value + resolved_else_values[name] = else_value result_types.append(then_value.type) return { "order": order, "result_types": result_types, - "then": then_assignment, - "else": else_assignment, + "then": { + **then_assignment, + "raw_values": resolved_then_values, + }, + "else": { + **else_assignment, + "raw_values": resolved_else_values, + }, } def _finalize_side_effect_if(self): @@ -574,6 +591,35 @@ def if_(cond) -> _IfCM: return _IfCM(cond) +def _is_branch_assign_literal(value) -> bool: + return isinstance(value, (int, float)) and not isinstance(value, bool) + + +def _reconcile_branch_assignment_values(name, then_value, else_value): + then_is_typed = hasattr(then_value, "type") + else_is_typed = hasattr(else_value, "type") + + if then_is_typed and else_is_typed: + return then_value, else_value + if then_is_typed: + return then_value, coerce_scalar_to_type( + else_value, + then_value.type, + context=f"br.assign(...) else branch value for '{name}'", + ) + if else_is_typed: + return coerce_scalar_to_type( + then_value, + else_value.type, + context=f"br.assign(...) then branch value for '{name}'", + ), else_value + + raise TypeError( + "br.assign(...) cannot infer a PTO type when both branches provide only " + f"Python literals for '{name}'; materialize one side explicitly with pto.const(...)" + ) + + # ── yield_ ──────────────────────────────────────────────────────────────────── def yield_(*vals): diff --git a/ptodsl/ptodsl/_diagnostics.py b/ptodsl/ptodsl/_diagnostics.py index 92226e000d..3fc812afe5 100644 --- a/ptodsl/ptodsl/_diagnostics.py +++ b/ptodsl/ptodsl/_diagnostics.py @@ -460,6 +460,10 @@ def unsupported_public_surface_error(name: str) -> AttributeError: "constexpr": ( "Use pto.const_expr for compile-time @pto.jit parameters and trace-time control-flow guards." ), + "copy_ubuf_to_ubuf": ( + "Use pto.mte_ub_ub(src, dst, len_burst, nburst=(n_burst, src_stride, dst_stride)) " + "for authored UB-to-UB DMA instead of the removed raw copy helper." + ), "tensor_spec": ( "Host tensor ABI hints were removed from the PTODSL public surface. Use explicit " 'GM pointers such as pto.ptr(pto.f32, "gm"), pass runtime shape/stride scalars, ' diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index 0e8f878880..cdb3eece24 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -244,9 +244,13 @@ def addptr(base_ptr, index_offset): _VLOAD_DIST_TOKENS = { "NORM", "UNPK_B8", "UNPK_B16", "UNPK_B32", - "BRC_B8", "BRC_B16", "BRC_B32", + "BRC_B8", "BRC_B16", "BRC_B32", "BRC_BLK", "US_B8", "US_B16", "DS_B8", "DS_B16", + # Extra load distributions already supported by the backend lowering. + "E2B_B16", "E2B_B32", + "UNPK4", + "SPLT4CHN", } @@ -542,8 +546,8 @@ def _vcvt_contract(requires_rnd, requires_sat, requires_part, *, part_family=Non _VCVT_CONTRACTS = { - ("f32", "f8e4m3"): _vcvt_contract(True, True, True, part_family="packed4", allowed_rnd="R"), - ("f32", "f8e5m2"): _vcvt_contract(True, True, True, part_family="packed4", allowed_rnd="R"), + ("f32", "f8e4m3"): _vcvt_contract(True, True, True, part_family="packed4", allowed_rnd="RAHZ"), + ("f32", "f8e5m2"): _vcvt_contract(True, True, True, part_family="packed4", allowed_rnd="RAHZ"), ("f32", "hif8"): _vcvt_contract(True, True, True, part_family="packed4", allowed_rnd="AH"), ("f32", "f16"): _vcvt_contract(True, True, True), ("f32", "bf16"): _vcvt_contract(True, True, True), @@ -1481,6 +1485,32 @@ def pdintlv_b32(lhs, rhs): ) +def _vreg_pair_op(op_ctor, lhs, rhs, *, context: str): + lhs_type = unwrap_surface_value(lhs).type + rhs_type = unwrap_surface_value(rhs).type + if lhs_type != rhs_type: + raise TypeError(f"{context} expects matching vreg types, got {lhs_type} and {rhs_type}") + _infer_vreg_metadata(lhs) + _infer_vreg_metadata(rhs) + op = op_ctor( + lhs_type, + lhs_type, + unwrap_surface_value(lhs), + unwrap_surface_value(rhs), + ) + return wrap_surface_value(op.low), wrap_surface_value(op.high) + + +def vintlv(lhs, rhs): + """``pto.vintlv`` – interleave two vector registers.""" + return _vreg_pair_op(_pto.VintlvOp, lhs, rhs, context="vintlv(lhs, rhs)") + + +def vdintlv(lhs, rhs): + """``pto.vdintlv`` – deinterleave two vector registers.""" + return _vreg_pair_op(_pto.VdintlvOp, lhs, rhs, context="vdintlv(lhs, rhs)") + + def vcmp(src0, src1, seed_mask, cmp_mode): """``pto.vcmp`` – vector/vector comparison producing a predicate mask.""" _, elem_type = _infer_vreg_metadata(src0) @@ -1771,32 +1801,6 @@ def vshr(lhs, rhs, mask): return _emit_binary_vec_op(_pto.VshrOp, lhs, rhs, mask) -def vshls(inp, scalar, mask): - """``pto.vshls`` – vector shift-left by scalar under mask.""" - _reject_low_precision_vreg_operands(inp, context="pto.vshls(...)") - return wrap_surface_value( - _pto.VshlsOp( - unwrap_surface_value(inp).type, - unwrap_surface_value(inp), - _coerce_i16(scalar, context="vshls"), - unwrap_surface_value(mask), - ).result - ) - - -def vshrs(inp, scalar, mask): - """``pto.vshrs`` – vector shift-right by scalar under mask.""" - _reject_low_precision_vreg_operands(inp, context="pto.vshrs(...)") - return wrap_surface_value( - _pto.VshrsOp( - unwrap_surface_value(inp).type, - unwrap_surface_value(inp), - _coerce_i16(scalar, context="vshrs"), - unwrap_surface_value(mask), - ).result - ) - - def vcmax(v, mask): """``pto.vcmax`` – cross-lane maximum reduction.""" return _emit_unary_vec_op(_pto.VcmaxOp, v, mask) @@ -2100,23 +2104,6 @@ def vselr(src0, src1): ) -def vci(index, order=None): - """``pto.vci`` – vector consecutive index generator.""" - raw_index = unwrap_surface_value(index) - if not hasattr(raw_index, "type"): - raw_index = _coerce_i32(raw_index, context="vci(index)") - result_type = _resolve(vreg_type(_elements_per_vreg(raw_index.type), raw_index.type)) - kwargs = {} - if order is not None: - token = getattr(order, "value", order) - if not isinstance(token, str): - token = str(token) - if "." in token: - token = token.rsplit(".", 1)[-1] - kwargs["order"] = token.strip().upper() - return wrap_surface_value(_pto.VciOp(result_type, raw_index, **kwargs).result) - - def vaddc(lhs, rhs, mask): """``pto.vaddc`` – vector add with carry-out predicate.""" _reject_low_precision_vreg_operands(lhs, rhs, context="pto.vaddc(...)") @@ -2182,19 +2169,6 @@ def vmrgsort4(destination, source0, source1, source2, source3, count, config): ) -def copy_ubuf_to_ubuf(source, destination, sid, n_burst, len_burst, src_stride, dst_stride): - """``pto.copy_ubuf_to_ubuf`` – raw UB-to-UB DMA primitive.""" - _pto.CopyUbufToUbufOp( - unwrap_surface_value(source), - unwrap_surface_value(destination), - _coerce_i64(sid, context="copy_ubuf_to_ubuf(sid)"), - _coerce_i64(n_burst, context="copy_ubuf_to_ubuf(n_burst)"), - _coerce_i64(len_burst, context="copy_ubuf_to_ubuf(len_burst)"), - _coerce_i64(src_stride, context="copy_ubuf_to_ubuf(src_stride)"), - _coerce_i64(dst_stride, context="copy_ubuf_to_ubuf(dst_stride)"), - ) - - def load_scalar(ptr_value, offset=0, result_type=None): """``pto.load_scalar`` – load one scalar from a pointer-like value.""" if result_type is None: @@ -2256,6 +2230,71 @@ def vlrelu(inp, alpha, mask): return _emit_vec_scalar_masked_op(_pto.VlreluOp, inp, alpha, mask, context="vlrelu") +def vshrs(inp, scalar, mask): + """``pto.vshrs`` – vector shift-right by a uniform ``i16`` amount.""" + _reject_low_precision_vreg_operands(inp, context="pto.vshrs(...)") + i16 = IntegerType.get_signless(16) + raw = unwrap_surface_value(scalar) + if hasattr(raw, "type"): + scalar_value = coerce_scalar_to_type(raw, i16, context="vshrs") + else: + scalar_value = materialize_scalar_literal(int(raw), i16, context="vshrs") + return wrap_surface_value( + _pto.VshrsOp( + unwrap_surface_value(inp).type, + unwrap_surface_value(inp), + unwrap_surface_value(scalar_value), + unwrap_surface_value(mask), + ).result + ) + + +def vshls(inp, scalar, mask): + """``pto.vshls`` – vector shift-left by a uniform ``i16`` amount.""" + _reject_low_precision_vreg_operands(inp, context="pto.vshls(...)") + i16 = IntegerType.get_signless(16) + raw = unwrap_surface_value(scalar) + if hasattr(raw, "type"): + scalar_value = coerce_scalar_to_type(raw, i16, context="vshls") + else: + scalar_value = materialize_scalar_literal(int(raw), i16, context="vshls") + return wrap_surface_value( + _pto.VshlsOp( + unwrap_surface_value(inp).type, + unwrap_surface_value(inp), + unwrap_surface_value(scalar_value), + unwrap_surface_value(mask), + ).result + ) + + +def vands(inp, scalar, mask): + """``pto.vands`` – vector AND scalar (emulated via ``vand`` + ``vbr``).""" + return vand( + inp, + vbr(_coerce_scalar_like_vector_element(inp, scalar, context="vands")), + mask, + ) + + +def vors(inp, scalar, mask): + """``pto.vors`` – vector OR scalar (emulated via ``vor`` + ``vbr``).""" + return vor( + inp, + vbr(_coerce_scalar_like_vector_element(inp, scalar, context="vors")), + mask, + ) + + +def vxors(inp, scalar, mask): + """``pto.vxors`` – vector XOR scalar (emulated via ``vxor`` + ``vbr``).""" + return vxor( + inp, + vbr(_coerce_scalar_like_vector_element(inp, scalar, context="vxors")), + mask, + ) + + def vaddrelu(lhs, rhs, mask): """``pto.vaddrelu`` – add, then apply ReLU.""" return vrelu(vadd(lhs, rhs, mask), mask) @@ -2281,6 +2320,39 @@ def vaxpy(alpha, x, y, mask): ) +def vmula(acc, lhs, rhs, mask): + """``pto.vmula`` – fused ``acc + lhs * rhs`` under mask.""" + _reject_low_precision_vreg_operands(acc, lhs, rhs, context="pto.vmula(...)") + return wrap_surface_value( + _pto.VmulaOp( + unwrap_surface_value(acc).type, + unwrap_surface_value(acc), + unwrap_surface_value(lhs), + unwrap_surface_value(rhs), + unwrap_surface_value(mask), + ).result + ) + + +def vci(base, order=None): + """``pto.vci`` – generate lane indices from a scalar base.""" + raw_base = unwrap_surface_value(base) + if hasattr(raw_base, "type"): + scalar_value = raw_base + elem_type = raw_base.type + elif isinstance(raw_base, int): + elem_type = IntegerType.get_signless(32) + scalar_value = materialize_scalar_literal(raw_base, elem_type, context="vci(base)") + else: + raise TypeError("vci(base) expects a runtime scalar or Python int") + + result_type = _resolve(vreg_type(_elements_per_vreg(elem_type), elem_type)) + kwargs = {} + if order is not None: + kwargs["order"] = order + return wrap_surface_value(_pto.VciOp(result_type, scalar_value, **kwargs).result) + + def vsel(true_v, false_v, mask): """``pto.vsel`` – element-wise select under a predicate mask.""" _reject_low_precision_vreg_operands(true_v, false_v, context="pto.vsel(...)") @@ -5818,6 +5890,7 @@ def import_reserved_buffer(name, *, peer_func): "pbitcast", "vcvt", "vpack", "vmulscvt", "ppack", "punpack", "pintlv_b8", "pintlv_b16", "pintlv_b32", "pdintlv_b8", "pdintlv_b16", "pdintlv_b32", + "vintlv", "vdintlv", "vgather2", "vgather2_bc", "vgatherb", "vscatter", "vsldb", "vsstb", "vcmp", "vcmps", "plds", "psts", "pstu", "vstar", "vstas", "vstur", "vstus", @@ -5828,8 +5901,8 @@ def import_reserved_buffer(name, *, peer_func): "vcmax", "vcadd", "vcmin", "vdup", "vexpdif", "vexp", "vln", "vsqrt", "vabs", "vneg", "vrec", "vrsqrt", "vrelu", "vnot", "vcgmax", "vcgadd", "vcgmin", "vcpadd", - "vadds", "vsubs", "vmuls", "vmaxs", "vmins", "vlrelu", - "vaxpy", "vaddrelu", "vsubrelu", + "vadds", "vsubs", "vmuls", "vmaxs", "vmins", "vlrelu", "vshrs", "vshls", "vands", "vors", "vxors", + "vaxpy", "vmula", "vci", "vaddrelu", "vsubrelu", "vsel", "make_tensor_view", "partition_view", "alloc_buffer", "alloc_tile", diff --git a/ptodsl/ptodsl/_runtime/native_build.py b/ptodsl/ptodsl/_runtime/native_build.py index 410a3997ae..476c754fa8 100644 --- a/ptodsl/ptodsl/_runtime/native_build.py +++ b/ptodsl/ptodsl/_runtime/native_build.py @@ -39,6 +39,11 @@ def _run(cmd: list[str], *, cwd: Path | None = None) -> None: ) +def _mlir_requires_enable_vmi(mlir_path: Path) -> bool: + text = mlir_path.read_text(encoding="utf-8") + return "pto.vmi." in text or "!pto.vmi." in text + + def _run_ptoas( mlir_path: Path, kernel_object: Path, @@ -59,6 +64,8 @@ def _run_ptoas( cmd.append(f"--pto-level={pto_level}") if insert_sync is True: cmd.append("--enable-insert-sync") + if _mlir_requires_enable_vmi(mlir_path): + cmd.append("--enable-vmi") cmd.extend([ "--enable-tile-op-expand", str(mlir_path), diff --git a/ptodsl/ptodsl/_surface_values.py b/ptodsl/ptodsl/_surface_values.py index 94374333f3..01592cd24f 100644 --- a/ptodsl/ptodsl/_surface_values.py +++ b/ptodsl/ptodsl/_surface_values.py @@ -147,6 +147,20 @@ def _maybe_cast_tile_buf_type(type_obj): return None +def _maybe_cast_mask_type(type_obj): + try: + return _pto.MaskType(type_obj) + except Exception: + return None + + +def _maybe_cast_vmi_mask_type(type_obj): + try: + return _pto.VMIMaskType(type_obj) + except Exception: + return None + + def wrap_surface_value( value, *, @@ -169,6 +183,8 @@ def wrap_surface_value( offsets=offsets, sizes=sizes, ) + if _maybe_cast_mask_type(type_obj) is not None or _maybe_cast_vmi_mask_type(type_obj) is not None: + return MaskValue(value) if _maybe_cast_tile_buf_type(type_obj) is not None: return TileValue(value, **(tile_metadata or {})) try: @@ -334,6 +350,10 @@ def _emit_vec_binary_op(op_name: str, lhs, rhs): return VecValue(emit_runtime_binary_op(op_name, lhs_raw, rhs_raw)) +class MaskValue(_SurfaceValue): + """Concrete authored wrapper for PTO / VMI mask SSA values.""" + + class MaskResultValue(_SurfaceValue): """Mask value that also supports `(mask, remained)` unpacking.""" diff --git a/ptodsl/ptodsl/_tracing/session.py b/ptodsl/ptodsl/_tracing/session.py index e990a26a39..bd6c99f33f 100644 --- a/ptodsl/ptodsl/_tracing/session.py +++ b/ptodsl/ptodsl/_tracing/session.py @@ -207,7 +207,13 @@ def bind_entry_block(self, entry_block) -> None: def validate_surface_value_access(self, value) -> None: """Reject inline-subkernel SSA values that escaped their outlined helper body.""" - record = self._escaped_inline_values.get(value) + try: + record = self._escaped_inline_values.get(value) + except TypeError: + raw_value = getattr(value, "_value", None) + if raw_value is None: + return + record = self._escaped_inline_values.get(raw_value) if record is None: return role, type_text = record diff --git a/ptodsl/ptodsl/_types.py b/ptodsl/ptodsl/_types.py index 24a318613a..bbfd8141f8 100644 --- a/ptodsl/ptodsl/_types.py +++ b/ptodsl/ptodsl/_types.py @@ -152,6 +152,41 @@ def resolve(self) -> Type: def __repr__(self): return f"" +class _VMIVRegDescriptor(_DType): + def __init__(self, lanes: int, elem): + self._lanes = lanes + self._elem = elem + + def resolve(self) -> Type: + elem = _ensure_tensor_storage_dtype(self._elem, context="pto.vmi.vreg(...)") + vreg_type_cls = getattr(_pto, "VMIVRegType", None) + if vreg_type_cls is None: + raise TypeError( + "The current PTO Python bindings do not expose VMIVRegType. " + "Rebuild the PTO Python extension before using pto.vmi.vreg(...)." + ) + return vreg_type_cls.get(self._lanes, elem) + + def __repr__(self): + return f"" + + +class _VMIMaskDescriptor(_DType): + def __init__(self, lanes: int): + self._lanes = lanes + self._granularity = "pred" + + def resolve(self) -> Type: + mask_type_cls = getattr(_pto, "VMIMaskType", None) + if mask_type_cls is None: + raise TypeError( + "The current PTO Python bindings do not expose VMIMaskType. " + "Rebuild the PTO Python extension before using pto.vmi.mask(...)." + ) + return mask_type_cls.get(self._lanes, self._granularity) + + def __repr__(self): + return f"" class _VecDescriptor(_DType): def __init__(self, elem, size: int): @@ -455,6 +490,16 @@ def mask_type(bits: str = "b32") -> _MaskDescriptor: return _MaskDescriptor(bits) +def vmi_vreg_type(lanes: int, elem) -> _VMIVRegDescriptor: + """Return a lazy descriptor for ``!pto.vmi.vreg``.""" + return _VMIVRegDescriptor(lanes, elem) + + +def vmi_mask_type(lanes: int) -> _VMIMaskDescriptor: + """Return a lazy descriptor for ``!pto.vmi.mask``.""" + return _VMIMaskDescriptor(lanes) + + def tile_buf_type(shape, dtype, valid_shape=None, *, blayout: str = "RowMajor", address_space: str = "ub", @@ -535,6 +580,7 @@ def part_tensor_view_type_from_dims(dims, elem) -> Type: "ui8", "ui16", "ui32", "ui64", "index", "ptr", "vreg_type", "vec_type", "mask_type", + "vmi_vreg_type", "vmi_mask_type", "tile_buf_type", "tensor_view_type", "tensor_view_type_from_dims", "part_tensor_view_type", "part_tensor_view_type_from_dims", ] diff --git a/ptodsl/ptodsl/_vmi_namespace.py b/ptodsl/ptodsl/_vmi_namespace.py new file mode 100644 index 0000000000..24233d7875 --- /dev/null +++ b/ptodsl/ptodsl/_vmi_namespace.py @@ -0,0 +1,980 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +"""Public PTODSL namespace for formal VMI APIs.""" + +from __future__ import annotations + +from collections.abc import Sequence + +from mlir.dialects import pto as _pto +from mlir.ir import BF16Type, F16Type, F32Type, Float8E4M3FNType, Float8E5M2Type, IntegerType, MemRefType + +from ._scalar_coercion import coerce_scalar_to_type +from ._surface_values import _coerce_index_value, _try_get_constant_index, unwrap_surface_value, wrap_surface_value +from ._types import _ensure_tensor_storage_dtype, _resolve, vmi_mask_type, vmi_vreg_type + + +class _UnspecifiedArgument: + def __repr__(self) -> str: + return "UNSPECIFIED" + + +_UNSPECIFIED = _UnspecifiedArgument() + + +def _missing_vmi_support_error(op_name: str) -> NotImplementedError: + return NotImplementedError( + f"{op_name} is not available in the current PTO Python bindings or " + "backend support. Rebuild PTO Python bindings and update VMI support " + "for this operation before using the PTODSL pto.vmi surface." + ) + + +def _unsupported_vmi_feature_error(op_name: str, feature: str) -> NotImplementedError: + return NotImplementedError( + f"{op_name} {feature} is not available in the current generated VMI " + "binding/backend support; update the PTO bindings or use an explicit " + "supported VMI form." + ) + + +def _generated(op_name: str): + fn = getattr(_pto, f"vmi_{op_name}", None) + if fn is None: + raise _missing_vmi_support_error(f"pto.vmi.{op_name}") + return fn + + +def _raw(value): + return unwrap_surface_value(value) + + +def _raw_sequence(values): + if _is_sequence(values): + return [_raw(value) for value in values] + return [_raw(values)] + + +def _is_sequence(value) -> bool: + return isinstance(value, Sequence) and not isinstance(value, (str, bytes)) + + +def _wrap_result(result): + if hasattr(result, "type"): + return wrap_surface_value(result) + try: + count = len(result) + except TypeError: + count = None + if count is not None: + return tuple(wrap_surface_value(result[index]) for index in range(count)) + if _is_sequence(result) or hasattr(result, "__iter__"): + return tuple(wrap_surface_value(value) for value in result) + return wrap_surface_value(result) + + +def _type_of(value): + return _raw(value).type + + +def _require_result_type(result_type, *, context: str): + if result_type is None: + raise TypeError(f"{context} requires explicit result_type") + return _resolve(result_type) + + +def _as_vmi_vreg_type(type_obj, *, context: str): + vreg_type_cls = getattr(_pto, "VMIVRegType", None) + if vreg_type_cls is None: + raise _missing_vmi_support_error("!pto.vmi.vreg") + try: + return vreg_type_cls(type_obj) + except Exception as exc: + raise TypeError(f"{context} expects a !pto.vmi.vreg value, got {type_obj}") from exc + + +def _vmi_element_type(type_obj, *, context: str): + return _as_vmi_vreg_type(type_obj, context=context).element_type + + +def _as_vmi_mask_type(type_obj, *, context: str): + mask_type_cls = getattr(_pto, "VMIMaskType", None) + if mask_type_cls is None: + raise _missing_vmi_support_error("!pto.vmi.mask") + try: + return mask_type_cls(type_obj) + except Exception as exc: + raise TypeError(f"{context} expects a !pto.vmi.mask value, got {type_obj}") from exc + + +def _vmi_mask_element_count(mask_type, *, context: str): + for attr in ("element_count", "elementCount"): + value = getattr(mask_type, attr, None) + if value is not None: + return int(value) + getter = getattr(mask_type, "getElementCount", None) + if callable(getter): + return int(getter()) + raise TypeError(f"{context} could not determine VMI mask lane count from {mask_type}") + + +def _vmi_layout_attr(type_obj): + for attr in ("layout", "layout_attr"): + value = getattr(type_obj, attr, None) + if value is not None: + return value + for getter_name in ("getLayout", "getLayoutAttr"): + getter = getattr(type_obj, getter_name, None) + if callable(getter): + value = getter() + if value is not None: + return value + return None + + +def _pointer_element_type(type_obj, *, context: str): + ptr_type_cls = getattr(_pto, "PtrType", None) + if ptr_type_cls is not None: + try: + return ptr_type_cls(type_obj).element_type + except Exception: + pass + try: + return MemRefType(type_obj).element_type + except Exception as exc: + raise TypeError(f"{context} expects a pointer or memref source, got {type_obj}") from exc + + +def _type_bit_width(type_obj, *, context: str): + if IntegerType.isinstance(type_obj): + return IntegerType(type_obj).width + if Float8E4M3FNType.isinstance(type_obj) or Float8E5M2Type.isinstance(type_obj): + return 8 + if F16Type.isinstance(type_obj) or BF16Type.isinstance(type_obj): + return 16 + if F32Type.isinstance(type_obj): + return 32 + raise TypeError(f"{context} does not support element type {type_obj}") + + +def _is_vmi_float_element_type(type_obj) -> bool: + return any( + cls.isinstance(type_obj) + for cls in (BF16Type, F16Type, F32Type, Float8E4M3FNType, Float8E5M2Type) + ) + + +def _normalize_vmi_vcvt_rounding(mode, *, context: str): + token = mode + if not isinstance(token, str): + token = str(token) + if "." in token: + token = token.rsplit(".", 1)[-1] + normalized = token.strip().upper() + allowed = {"R", "A", "H", "Z"} + if normalized not in allowed: + expected = ", ".join(sorted(allowed)) + raise ValueError( + f"{context} does not support rounding {mode!r}; expected one of {expected}" + ) + return normalized + + +def _derive_vcvt_result_type(source, to_dtype, *, context: str): + if to_dtype is None: + raise TypeError(f"{context} requires to_dtype") + source_type = _as_vmi_vreg_type(_type_of(source), context=context) + elem_type = _ensure_tensor_storage_dtype(to_dtype, context=context) + return _pto.VMIVRegType.get( + source_type.element_count, + elem_type, + layout=source_type.layout, + ) + + +def _derive_vinterpret_cast_result_type(source, to_dtype, *, context: str): + if to_dtype is None: + raise TypeError(f"{context} requires to_dtype") + source_type = _as_vmi_vreg_type(_type_of(source), context=context) + source_elem_type = source_type.element_type + target_elem_type = _ensure_tensor_storage_dtype(to_dtype, context=context) + source_bits = _type_bit_width(source_elem_type, context=context) + target_bits = _type_bit_width(target_elem_type, context=context) + if source_bits != target_bits: + raise TypeError( + f"{context} requires source and target element widths to match; got " + f"{source_elem_type} -> {target_elem_type}" + ) + return _pto.VMIVRegType.get( + source_type.element_count, + target_elem_type, + layout=source_type.layout, + ) + + +def _derive_vbrc_result_type(value, size, *, context: str): + if size is None: + raise TypeError(f"{context} requires size") + raw_value = _raw(value) + if not hasattr(raw_value, "type"): + raise TypeError( + f"{context} requires a typed scalar such as pto.f32(0.0) or " + "a VMI vector input; plain Python scalars are ambiguous" + ) + value_type = raw_value.type + if _is_vmi_vreg_type(value_type): + elem_type = _vmi_element_type(value_type, context=context) + else: + elem_type = value_type + return _pto.VMIVRegType.get(size, elem_type) + + +def _derive_vci_result_type(base, size, *, context: str): + if size is None: + raise TypeError(f"{context} requires size") + raw_base = _raw(base) + if not hasattr(raw_base, "type"): + raise TypeError( + f"{context} requires a typed scalar such as pto.i32(0) or " + "pto.f32(0.0); plain Python scalars are ambiguous" + ) + return _pto.VMIVRegType.get(size, raw_base.type) + + +def _derive_vmull_result_types(a, b, *, context: str): + lhs_type = _as_vmi_vreg_type(_type_of(a), context=context) + rhs_type = _as_vmi_vreg_type(_type_of(b), context=context) + if lhs_type != rhs_type: + raise TypeError(f"{context} requires a and b to have identical VMI vreg types") + element_type = lhs_type.element_type + if not IntegerType.isinstance(element_type): + raise TypeError(f"{context} requires 32-bit integer vectors") + integer_type = IntegerType(element_type) + if integer_type.width != 32: + raise TypeError(f"{context} requires 32-bit integer vectors") + return lhs_type, rhs_type + + +def _derive_hist_result_type(acc, *, context: str): + return _as_vmi_vreg_type(_type_of(acc), context=context) + + +def _derive_vgather_result_type(source, offsets, *, context: str): + offsets_type = _as_vmi_vreg_type(_type_of(offsets), context=context) + result_type = _pointer_element_type(_type_of(source), context=context) + return _pto.VMIVRegType.get( + offsets_type.element_count, + result_type, + layout=offsets_type.layout, + ) + + +def _derive_vgatherb_result_type(source, mask, *, context: str): + mask_type = _as_vmi_mask_type(_type_of(mask), context=context) + result_element_type = _pointer_element_type(_type_of(source), context=context) + result_layout = _vmi_layout_attr(mask_type) + return _pto.VMIVRegType.get( + _vmi_mask_element_count(mask_type, context=context), + result_element_type, + layout=result_layout, + ) + + +def _derive_vmi_reduce_result_type(source, group, *, context: str): + source_type = _as_vmi_vreg_type(_type_of(source), context=context) + result_lanes = 1 + if group is not None: + try: + result_lanes = int(group) + except (TypeError, ValueError) as exc: + raise TypeError(f"{context} requires group to be an integer when provided") from exc + if result_lanes <= 0: + raise TypeError(f"{context} requires group to be positive, got {group!r}") + return _pto.VMIVRegType.get(result_lanes, source_type.element_type) + + +def _coerce_scalar_like_vmi_element(vector_value, scalar_value, *, context: str): + elem_type = _vmi_element_type(_type_of(vector_value), context=context) + return coerce_scalar_to_type(scalar_value, elem_type, context=context) + + +def _variadic_mask(mask): + if mask is None: + return [] + return _raw_sequence(mask) + + +def _required_mask(mask, *, context: str): + if mask is None: + raise TypeError(f"{context} requires a mask operand") + return _raw(mask) + + +def _required_variadic_mask(mask, *, context: str): + if mask is None: + raise TypeError(f"{context} requires a mask operand") + return _raw_sequence(mask) + + +def _i16_value(value, *, context: str): + if value is None: + return None + return coerce_scalar_to_type(value, IntegerType.get_signless(16), context=context) + + +def _resolve_vmi_mask_type(size, *, context: str): + if size is None: + raise TypeError(f"{context} requires size") + return _resolve(vmi_mask_type(size)) + + +def _vmi_vreg_element_count(type_obj, *, context: str): + vreg_type = _as_vmi_vreg_type(type_obj, context=context) + for attr in ("element_count", "elementCount"): + value = getattr(vreg_type, attr, None) + if value is not None: + return int(value) + getter = getattr(vreg_type, "getElementCount", None) + if callable(getter): + return int(getter()) + raise TypeError(f"{context} could not determine VMI vector lane count from {type_obj}") + + +def _resolve_vmi_unpack_result_type(source, size, to_dtype, *, context: str): + if to_dtype is None: + raise TypeError(f'{context} requires to_dtype when dist_mode="unpack"') + source_type = _pointer_element_type(_type_of(source), context=context) + result_type = _ensure_tensor_storage_dtype(to_dtype, context=context) + source_bits = _type_bit_width(source_type, context=context) + result_bits = _type_bit_width(result_type, context=context) + if source_bits * 2 != result_bits: + raise TypeError( + f"{context} requires unpack to widen by exactly one step; got " + f"{source_type} -> {result_type}" + ) + return _pto.VMIVRegType.get(size, result_type) + + +def _resolve_vmi_vload_result_types(source, size, *, dist_mode, to_dtype, context: str): + if to_dtype is not None and dist_mode != "unpack": + raise TypeError(f'{context} accepts to_dtype only when dist_mode="unpack"') + if size is None: + raise TypeError(f"{context} requires size") + if dist_mode == "unpack": + return [_resolve_vmi_unpack_result_type(source, size, to_dtype, context=context)] + element_type = _pointer_element_type(_type_of(source), context=context) + resolved = _pto.VMIVRegType.get(size, element_type) + if dist_mode == "dintlv": + return [resolved, resolved] + return [resolved] + + +def _validate_vmi_load_modes( + context: str, + *, + dist_mode, + group, + stride, + block_stride, + repeat_stride, + allow_group_brc: bool, + allowed_dist_modes, +): + if dist_mode is not None and dist_mode not in allowed_dist_modes: + expected = ", ".join(repr(mode) for mode in sorted(allowed_dist_modes, key=str)) + raise TypeError(f"{context} does not support dist_mode={dist_mode!r}; expected one of {expected}") + + if group is not None: + if dist_mode is not None and (not allow_group_brc or dist_mode != "brc"): + raise TypeError(f"{context} does not allow dist_mode together with group") + if block_stride is not None or repeat_stride is not None: + raise TypeError(f"{context} does not allow block_stride together with group") + if stride is None: + raise TypeError(f"{context} with group=... requires stride") + return + + if block_stride is not None or repeat_stride is not None: + if dist_mode is not None: + raise TypeError(f"{context} does not allow dist_mode together with block_stride") + if block_stride is None or repeat_stride is None: + raise TypeError(f"{context} requires block_stride and repeat_stride together") + if stride is not None: + raise TypeError(f"{context} does not allow stride together with block_stride") + return + + if stride is not None: + raise TypeError(f"{context} accepts stride only when group is provided") + + +def _call_value(op_name: str, *args, **kwargs): + return _wrap_result(_generated(op_name)(*args, **kwargs)) + + +def _emit_binary(op_name: str, lhs, rhs, mask=None, *, pmode=None, loc=None, ip=None): + return _call_value( + op_name, + _type_of(lhs), + _raw(lhs), + _raw(rhs), + _variadic_mask(mask), + pmode=pmode, + loc=loc, + ip=ip, + ) + + +def _emit_unary(op_name: str, source, mask=None, *, pmode=None, loc=None, ip=None): + return _call_value( + op_name, + _type_of(source), + _raw(source), + _variadic_mask(mask), + pmode=pmode, + loc=loc, + ip=ip, + ) + + +def _emit_vec_scalar(op_name: str, source, scalar, mask, *, pmode=None, loc=None, ip=None): + context = f"pto.vmi.{op_name}(...)" + return _call_value( + op_name, + _type_of(source), + _raw(source), + _coerce_scalar_like_vmi_element(source, scalar, context=context), + _required_mask(mask, context=context), + pmode=pmode, + loc=loc, + ip=ip, + ) + + +def _emit_reduce( + op_name: str, + source, + mask, + *, + group=None, + pmode=None, + loc=None, + ip=None, + reassoc=_UNSPECIFIED, +): + context = f"pto.vmi.{op_name}(...)" + if op_name == "vcadd": + source_elem_type = _vmi_element_type(_type_of(source), context=context) + if reassoc is _UNSPECIFIED: + if _is_vmi_float_element_type(source_elem_type): + raise TypeError( + f"{context} on floating-point vectors requires an explicit reassoc " + "argument; spell out reassoc=True or reassoc=False" + ) + elif not isinstance(reassoc, bool): + raise TypeError( + f"{context} requires reassoc to be the Python boolean True or False; " + f"received {reassoc!r}" + ) + kwargs = {"group": group, "pmode": pmode, "loc": loc, "ip": ip} + if reassoc is not _UNSPECIFIED: + kwargs["reassoc"] = reassoc + return _call_value( + op_name, + _derive_vmi_reduce_result_type(source, group, context=context), + _raw(source), + _required_mask(mask, context=context), + **kwargs, + ) + + +class _VMINamespace: + vreg = staticmethod(vmi_vreg_type) + mask = staticmethod(vmi_mask_type) + + @staticmethod + def vload( + source, + offset, + *, + size, + to_dtype=None, + stride=None, + block_stride=None, + repeat_stride=None, + dist_mode=None, + group=None, + loc=None, + ip=None, + ): + _validate_vmi_load_modes( + "pto.vmi.vload(...)", + dist_mode=dist_mode, + group=group, + stride=stride, + block_stride=block_stride, + repeat_stride=repeat_stride, + allow_group_brc=True, + allowed_dist_modes={None, "continuous", "dintlv", "unpack", "brc"}, + ) + result_types = _resolve_vmi_vload_result_types( + source, + size, + dist_mode=dist_mode, + to_dtype=to_dtype, + context="pto.vmi.vload(...)", + ) + return _call_value( + "vload", + result_types, + _raw(source), + _coerce_index_value(offset), + stride=None if stride is None else _coerce_index_value(stride), + block_stride=_i16_value(block_stride, context="pto.vmi.vload(block_stride)"), + repeat_stride=_i16_value(repeat_stride, context="pto.vmi.vload(repeat_stride)"), + dist_mode=dist_mode, + group=group, + loc=loc, + ip=ip, + ) + + @staticmethod + def vstore( + values, + destination, + offset, + mask=None, + *, + stride=None, + block_stride=None, + repeat_stride=None, + dist_mode=None, + group=None, + pmode=None, + loc=None, + ip=None, + ): + _validate_vmi_load_modes( + "pto.vmi.vstore(...)", + dist_mode=dist_mode, + group=group, + stride=stride, + block_stride=block_stride, + repeat_stride=repeat_stride, + allow_group_brc=False, + allowed_dist_modes={None, "continuous", "dintlv"}, + ) + if group is not None and mask is not None: + raise TypeError("pto.vmi.vstore(...) group mode does not take a mask operand") + if dist_mode == "dintlv": + if not _is_sequence(values) or len(values) != 2: + raise TypeError('pto.vmi.vstore(...) with dist_mode="dintlv" requires an (even, odd) pair') + elif _is_sequence(values): + raise TypeError("pto.vmi.vstore(...) expects a single VMI vector unless dist_mode=\"dintlv\"") + return _generated("vstore")( + _raw_sequence(values), + _raw(destination), + _coerce_index_value(offset), + _variadic_mask(mask), + stride=None if stride is None else _coerce_index_value(stride), + block_stride=_i16_value(block_stride, context="pto.vmi.vstore(block_stride)"), + repeat_stride=_i16_value(repeat_stride, context="pto.vmi.vstore(repeat_stride)"), + dist_mode=dist_mode, + group=group, + pmode=pmode, + loc=loc, + ip=ip, + ) + + @staticmethod + def vci(base, *, size, order=None, loc=None, ip=None): + result_type = _derive_vci_result_type(base, size, context="pto.vmi.vci(...)") + base = coerce_scalar_to_type( + base, + _vmi_element_type(result_type, context="pto.vmi.vci(...)"), + context="pto.vmi.vci(base)", + ) + return _call_value("vci", result_type, base, order=order, loc=loc, ip=ip) + + vadd = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vadd", lhs, rhs, mask, **kw)) + vsub = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vsub", lhs, rhs, mask, **kw)) + vmul = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vmul", lhs, rhs, mask, **kw)) + vdiv = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vdiv", lhs, rhs, mask, **kw)) + vmax = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vmax", lhs, rhs, mask, **kw)) + vmin = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vmin", lhs, rhs, mask, **kw)) + vand = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vand", lhs, rhs, mask, **kw)) + vor = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vor", lhs, rhs, mask, **kw)) + vxor = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vxor", lhs, rhs, mask, **kw)) + vshl = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vshl", lhs, rhs, mask, **kw)) + vshr = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vshr", lhs, rhs, mask, **kw)) + + vabs = staticmethod(lambda source, mask=None, **kw: _emit_unary("vabs", source, mask, **kw)) + vneg = staticmethod(lambda source, mask=None, **kw: _emit_unary("vneg", source, mask, **kw)) + vrelu = staticmethod(lambda source, mask=None, **kw: _emit_unary("vrelu", source, mask, **kw)) + vexp = staticmethod(lambda source, mask=None, **kw: _emit_unary("vexp", source, mask, **kw)) + vln = staticmethod(lambda source, mask=None, **kw: _emit_unary("vln", source, mask, **kw)) + vsqrt = staticmethod(lambda source, mask=None, **kw: _emit_unary("vsqrt", source, mask, **kw)) + vnot = staticmethod(lambda source, mask=None, **kw: _emit_unary("vnot", source, mask, **kw)) + + vadds = staticmethod(lambda source, scalar, mask, **kw: _emit_vec_scalar("vadds", source, scalar, mask, **kw)) + vmuls = staticmethod(lambda source, scalar, mask, **kw: _emit_vec_scalar("vmuls", source, scalar, mask, **kw)) + vmaxs = staticmethod(lambda source, scalar, mask, **kw: _emit_vec_scalar("vmaxs", source, scalar, mask, **kw)) + vmins = staticmethod(lambda source, scalar, mask, **kw: _emit_vec_scalar("vmins", source, scalar, mask, **kw)) + vshls = staticmethod(lambda source, scalar, mask, **kw: _emit_vec_scalar("vshls", source, scalar, mask, **kw)) + vshrs = staticmethod(lambda source, scalar, mask, **kw: _emit_vec_scalar("vshrs", source, scalar, mask, **kw)) + + @staticmethod + def vcmp(lhs, rhs, seed, cmp, *, pmode=None, loc=None, ip=None): + return _call_value( + "vcmp", + _type_of(seed), + _raw(lhs), + _raw(rhs), + _raw(seed), + cmp, + pmode=pmode, + loc=loc, + ip=ip, + ) + + @staticmethod + def vcmps(source, scalar, seed, cmp, *, pmode=None, loc=None, ip=None): + context = "pto.vmi.vcmps(...)" + return _call_value( + "vcmps", + _type_of(seed), + _raw(source), + _coerce_scalar_like_vmi_element(source, scalar, context=context), + _raw(seed), + cmp, + pmode=pmode, + loc=loc, + ip=ip, + ) + + @staticmethod + def vsel(mask, true_value, false_value, *, pmode=None, loc=None, ip=None): + return _call_value( + "vsel", + _type_of(true_value), + _raw(mask), + _raw(true_value), + _raw(false_value), + pmode=pmode, + loc=loc, + ip=ip, + ) + + @staticmethod + def vselr(source, index, *, loc=None, ip=None): + return _call_value( + "vselr", + _type_of(source), + _raw(source), + _raw(index), + loc=loc, + ip=ip, + ) + + @staticmethod + def vbrc(value, *, size, group=None, loc=None, ip=None): + context = "pto.vmi.vbrc(...)" + result_type = _derive_vbrc_result_type(value, size, context=context) + raw_value = _raw(value) + if group is not None and (not hasattr(raw_value, "type") or not _is_vmi_vreg_type(raw_value.type)): + raise TypeError(f"{context} with group=... requires a VMI vector input") + if group is not None: + if isinstance(group, bool) or not isinstance(group, int): + raise TypeError(f"{context} requires group to be a positive Python integer") + if group <= 0: + raise ValueError(f"{context} requires group to be positive, got {group!r}") + if not hasattr(raw_value, "type") or not _is_vmi_vreg_type(raw_value.type): + raise TypeError(f"{context} with group=... requires a VMI vector input") + value_lanes = _vmi_vreg_element_count(raw_value.type, context=context) + if value_lanes != group: + raise ValueError( + f"{context} with group=... requires the input lane count to match group; " + f"got {value_lanes} lanes for group={group}" + ) + if not hasattr(raw_value, "type") or not _is_vmi_vreg_type(raw_value.type): + raw_value = coerce_scalar_to_type( + value, + _vmi_element_type(result_type, context=context), + context="pto.vmi.vbrc(value)", + ) + return _call_value("vbrc", result_type, raw_value, group=group, loc=loc, ip=ip) + + vcadd = staticmethod(lambda source, mask, *, group=None, pmode=None, reassoc=_UNSPECIFIED, loc=None, ip=None: _emit_reduce("vcadd", source, mask, group=group, pmode=pmode, reassoc=reassoc, loc=loc, ip=ip)) + vcmax = staticmethod(lambda source, mask, *, group=None, pmode=None, loc=None, ip=None: _emit_reduce("vcmax", source, mask, group=group, pmode=pmode, loc=loc, ip=ip)) + vcmin = staticmethod(lambda source, mask, *, group=None, pmode=None, loc=None, ip=None: _emit_reduce("vcmin", source, mask, group=group, pmode=pmode, loc=loc, ip=ip)) + + @staticmethod + def vcvt( + source, + to_dtype=None, + mask=None, + *, + rounding=None, + saturate=None, + pmode=None, + loc=None, + ip=None, + ): + if mask is not None: + raise _unsupported_vmi_feature_error("pto.vmi.vcvt", "masked form") + result_type = _derive_vcvt_result_type(source, to_dtype, context="pto.vmi.vcvt(...)") + if rounding is not None: + rounding = _normalize_vmi_vcvt_rounding( + rounding, + context="pto.vmi.vcvt(..., rounding=...)", + ) + return _call_value( + "vcvt", + result_type, + _raw(source), + rounding=rounding, + saturate=saturate, + pmode=pmode, + loc=loc, + ip=ip, + ) + + @staticmethod + def vinterpret_cast(source, to_dtype=None, *, loc=None, ip=None): + return _call_value( + "vinterpret_cast", + _derive_vinterpret_cast_result_type( + source, + to_dtype, + context="pto.vmi.vinterpret_cast(...)", + ), + _raw(source), + loc=loc, + ip=ip, + ) + + @staticmethod + def vexpdif(x, max_value, mask, *, pmode=None, loc=None, ip=None): + return _call_value( + "vexpdif", + _type_of(max_value), + _raw(x), + _raw(max_value), + _required_mask(mask, context="pto.vmi.vexpdif(...)"), + pmode=pmode, + loc=loc, + ip=ip, + ) + + @staticmethod + def vaxpy(x, acc, alpha, mask, *, pmode=None, loc=None, ip=None): + context = "pto.vmi.vaxpy(...)" + return _call_value( + "vaxpy", + _type_of(acc), + _raw(x), + _raw(acc), + _coerce_scalar_like_vmi_element(x, alpha, context=context), + _required_mask(mask, context=context), + pmode=pmode, + loc=loc, + ip=ip, + ) + + @staticmethod + def vlrelu(x, slope, mask, *, pmode=None, loc=None, ip=None): + return _call_value( + "vlrelu", + _type_of(x), + _raw(x), + _coerce_scalar_like_vmi_element(x, slope, context="pto.vmi.vlrelu(...)"), + _required_mask(mask, context="pto.vmi.vlrelu(...)"), + pmode=pmode, + loc=loc, + ip=ip, + ) + + @staticmethod + def vprelu(x, alpha, mask, *, pmode=None, loc=None, ip=None): + return _call_value( + "vprelu", + _type_of(x), + _raw(x), + _raw(alpha), + _required_mask(mask, context="pto.vmi.vprelu(...)"), + pmode=pmode, + loc=loc, + ip=ip, + ) + + @staticmethod + def vmull(a, b, mask, *, pmode=None, loc=None, ip=None): + result_type = _type_of(a) + return _call_value( + "vmull", + *_derive_vmull_result_types(a, b, context="pto.vmi.vmull(...)"), + _raw(a), + _raw(b), + _required_mask(mask, context="pto.vmi.vmull(...)"), + pmode=pmode, + loc=loc, + ip=ip, + ) + + @staticmethod + def vmula(acc, lhs, rhs, mask, *, pmode=None, loc=None, ip=None): + return _call_value( + "vmula", + _type_of(acc), + _raw(acc), + _raw(lhs), + _raw(rhs), + _required_variadic_mask(mask, context="pto.vmi.vmula(...)"), + pmode=pmode, + loc=loc, + ip=ip, + ) + + @staticmethod + def vdhist(acc, source, mask, *, loc=None, ip=None): + return _call_value( + "vdhist", + _derive_hist_result_type(acc, context="pto.vmi.vdhist(...)"), + _raw(acc), + _raw(source), + _required_mask(mask, context="pto.vmi.vdhist(...)"), + loc=loc, + ip=ip, + ) + + @staticmethod + def vchist(acc, source, mask, *, loc=None, ip=None): + return _call_value( + "vchist", + _derive_hist_result_type(acc, context="pto.vmi.vchist(...)"), + _raw(acc), + _raw(source), + _required_mask(mask, context="pto.vmi.vchist(...)"), + loc=loc, + ip=ip, + ) + + @staticmethod + def vgather(source, offsets, mask, *, pmode=None, loc=None, ip=None): + return _call_value( + "vgather", + _derive_vgather_result_type(source, offsets, context="pto.vmi.vgather(...)"), + _raw(source), + _raw(offsets), + _required_mask(mask, context="pto.vmi.vgather(...)"), + pmode=pmode, + loc=loc, + ip=ip, + ) + + @staticmethod + def vgatherb(source, offsets, mask, *, pmode=None, loc=None, ip=None): + return _call_value( + "vgatherb", + _derive_vgatherb_result_type(source, mask, context="pto.vmi.vgatherb(...)"), + _raw(source), + _raw(offsets), + _required_mask(mask, context="pto.vmi.vgatherb(...)"), + pmode=pmode, + loc=loc, + ip=ip, + ) + + @staticmethod + def vscatter(value, destination, offsets, mask, *, pmode=None, loc=None, ip=None): + return _generated("vscatter")( + _raw(value), + _raw(destination), + _raw(offsets), + _required_mask(mask, context="pto.vmi.vscatter(...)"), + pmode=pmode, + loc=loc, + ip=ip, + ) + + @staticmethod + def create_mask( + active_lanes, + *, + size, + group=None, + loc=None, + ip=None, + ): + context = "pto.vmi.create_mask(...)" + result_type = _resolve_vmi_mask_type(size, context=context) + if group is None: + return _call_value("create_mask", result_type, _coerce_index_value(active_lanes), loc=loc, ip=ip) + if isinstance(group, bool) or not isinstance(group, int): + raise TypeError(f"{context} requires group to be a positive Python integer") + if group <= 0: + raise ValueError(f"{context} requires group to be positive, got {group!r}") + if size % group != 0: + raise ValueError(f"{context} requires size to be divisible by group; got size={size!r}, group={group!r}") + group_size = size // group + active_lanes_const = _try_get_constant_index(active_lanes) + if active_lanes_const is not None and active_lanes_const > group_size: + raise ValueError( + f"{context} requires active_lanes to be <= the inferred group_size; " + f"got active_lanes={active_lanes_const!r}, group_size={group_size!r}" + ) + return _call_value( + "create_group_mask", + result_type, + _coerce_index_value(active_lanes), + group, + group_size, + loc=loc, + ip=ip, + ) + + @staticmethod + def vintlv(lhs, rhs, mask, *, pmode=None, loc=None, ip=None): + return _call_value( + "vintlv", + _type_of(lhs), + _type_of(rhs), + _raw(lhs), + _raw(rhs), + _required_mask(mask, context="pto.vmi.vintlv(...)"), + pmode=pmode, + loc=loc, + ip=ip, + ) + + @staticmethod + def vdintlv(lhs, rhs, mask, *, pmode=None, loc=None, ip=None): + return _call_value( + "vdintlv", + _type_of(lhs), + _type_of(rhs), + _raw(lhs), + _raw(rhs), + _required_mask(mask, context="pto.vmi.vdintlv(...)"), + pmode=pmode, + loc=loc, + ip=ip, + ) + + +def _is_vmi_vreg_type(type_obj) -> bool: + vreg_type_cls = getattr(_pto, "VMIVRegType", None) + if vreg_type_cls is None: + return False + try: + return vreg_type_cls.isinstance(type_obj) + except Exception: + return False + + +vmi = _VMINamespace() + +__all__ = ["vmi"] diff --git a/ptodsl/ptodsl/pto.py b/ptodsl/ptodsl/pto.py index fdd0379023..2b269821c4 100644 --- a/ptodsl/ptodsl/pto.py +++ b/ptodsl/ptodsl/pto.py @@ -78,6 +78,7 @@ ) from ._tensor_factories import empty_like # noqa: F401 from ._tile_namespace import tile # noqa: F401 +from ._vmi_namespace import vmi # noqa: F401 # ── Operations ──────────────────────────────────────────────────────────────── from ._ops import ( # noqa: F401 @@ -96,6 +97,7 @@ ppack, punpack, pintlv_b8, pintlv_b16, pintlv_b32, pdintlv_b8, pdintlv_b16, pdintlv_b32, + vintlv, vdintlv, vgather2, vgather2_bc, vgatherb, vscatter, vsldb, vsstb, vcmp, vcmps, plds, psts, pstu, vstar, vstas, vstur, vstus, @@ -108,9 +110,9 @@ vcgmax, vcgadd, vcgmin, vcpadd, vtrc, vprelu, vintlv, vdintlv, vselr, vci, vaddc, vaddcs, vmull, vbitsort, vmrgsort4, - copy_ubuf_to_ubuf, load_scalar, store_scalar, - vadds, vsubs, vmuls, vmaxs, vmins, vlrelu, - vaxpy, vaddrelu, vsubrelu, + load_scalar, store_scalar, + vadds, vsubs, vmuls, vmaxs, vmins, vlrelu, vands, vors, vxors, + vaxpy, vmula, vaddrelu, vsubrelu, vsel, make_tensor_view, partition_view, alloc_buffer, alloc_tile, @@ -180,6 +182,6 @@ def gm_ptr(elem): def __getattr__(name): - if name in {"ukernel", "tile_buf_type", "vecscope", "as_ptr", "vbrc_load", "vsts_1pt", "constexpr", "tensor_spec", "TensorSpec"}: + if name in {"ukernel", "tile_buf_type", "vecscope", "as_ptr", "vbrc_load", "vsts_1pt", "constexpr", "copy_ubuf_to_ubuf", "tensor_spec", "TensorSpec"}: raise unsupported_public_surface_error(name) raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tsort32.py b/ptodsl/ptodsl/tilelib/templates/a5/tsort32.py index 0bf57a7b20..7d3ebb0f4e 100644 --- a/ptodsl/ptodsl/tilelib/templates/a5/tsort32.py +++ b/ptodsl/ptodsl/tilelib/templates/a5/tsort32.py @@ -164,14 +164,11 @@ def template_tsort32_with_tmp(src: pto.Tile, idx: pto.Tile, tmp: pto.Tile, dst: tmp_last_offset = repeat_num_per_row * BLOCK_SIZE - BLOCK_SIZE for row in range(0, valid_rows, 1): - pto.copy_ubuf_to_ubuf( + pto.mte_ub_ub( pto.addptr(src_ptr, row * src_stride), tmp_ptr, - 0, - 1, len_burst, - 0, - 0, + nburst=(1, 0, 0), ) pad_mask, _ = pto.make_mask(dtype, BLOCK_SIZE - src_tail_per_row) pad_vec = pto.vdup(pad_value, pad_mask) @@ -225,14 +222,11 @@ def template_tsort32_with_tmp(src: pto.Tile, idx: pto.Tile, tmp: pto.Tile, dst: ) len_burst = (src_tail_per_row * elem_bytes + BLOCK_SIZE - 1) // BLOCK_SIZE - pto.copy_ubuf_to_ubuf( + pto.mte_ub_ub( pto.addptr(src_ptr, row * src_stride + tail_src_offset), tmp_ptr, - 0, - 1, len_burst, - 0, - 0, + nburst=(1, 0, 0), ) tmp_last_offset = ( diff --git a/ptodsl/tests/CMakeLists.txt b/ptodsl/tests/CMakeLists.txt index f788ce4552..a61074b89e 100644 --- a/ptodsl/tests/CMakeLists.txt +++ b/ptodsl/tests/CMakeLists.txt @@ -58,7 +58,7 @@ endforeach() if(NOT TARGET check-dsl) set(_ptodsl_check_depends) - foreach(_target IN ITEMS PTOPythonModules pto-opt ptoas) + foreach(_target IN ITEMS PTOPythonModules ptoas_runtime_deps) if(TARGET ${_target}) list(APPEND _ptodsl_check_depends ${_target}) endif() diff --git a/ptodsl/tests/support/docs_fragment_fixtures.py b/ptodsl/tests/support/docs_fragment_fixtures.py index db833700c2..4fbc021a6f 100644 --- a/ptodsl/tests/support/docs_fragment_fixtures.py +++ b/ptodsl/tests/support/docs_fragment_fixtures.py @@ -189,6 +189,51 @@ def type_system_make_mask_probe(): {SNIPPET_PLACEHOLDER} """ ), + "vmi.vector_pipeline": _fixture( + f""" + @pto.jit(target="a5") + def vmi_vector_pipeline_probe(): + src_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + other_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + dst_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + src_ptr = src_tile.as_ptr() + other_ptr = other_tile.as_ptr() + dst_ptr = dst_tile.as_ptr() + offset = pto.const(0, dtype=pto.index) + active_lanes = pto.const(64, dtype=pto.index) + vec_ty = pto.vmi.vreg(64, pto.f32) + mask_ty = pto.vmi.mask(64) + {SNIPPET_PLACEHOLDER} + """ + ), + "vmi.index_select": _fixture( + f""" + @pto.jit(target="a5") + def vmi_index_select_probe(): + src_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + src_ptr = src_tile.as_ptr() + offset = pto.const(0, dtype=pto.index) + src = pto.vmi.vload(src_ptr, offset, size=64) + {SNIPPET_PLACEHOLDER} + """ + ), + "vmi.predicate_and_rearrange": _fixture( + f""" + @pto.jit(target="a5") + def vmi_predicate_and_rearrange_probe(): + src_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + src_ptr = src_tile.as_ptr() + offset = pto.const(0, dtype=pto.index) + active_lanes = pto.const(64, dtype=pto.index) + active_per_group = pto.const(8, dtype=pto.index) + src = pto.vmi.vload(src_ptr, offset, size=64) + mask = pto.vmi.create_mask( + active_lanes, + size=64, + ) + {SNIPPET_PLACEHOLDER} + """ + ), "quick_start.make_tensor_view": _fixture( f""" @pto.jit(target="a5") diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index 6a6d9fdec4..41f1de4703 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -19,6 +19,8 @@ from ptodsl import pto, scalar from ptodsl import _types as pto_types +import ptodsl._vmi_namespace as vmi_namespace +from ptodsl._ast_rewrite import PTODSLAstRewriteError from ptodsl._bootstrap import make_context from ptodsl._kernel_signature import DeviceParameterSpec, HelperMarkerParameterSpec, RuntimeScalarParameterSpec from ptodsl._tracing.runtime import SignatureTracingRuntime @@ -1256,6 +1258,107 @@ def ast_runtime_for_branch_local_temp_probe(rows: pto.i32): _ = out +@pto.jit(target="a5") +def ast_runtime_ifexp_assign_probe(rows: pto.i32): + limit = pto.const(64, dtype=pto.index) + zero = pto.const(0, dtype=pto.index) + + for row in range(rows): + cnt = row if row < limit else limit + out = cnt + zero + _ = out + + +@pto.jit(target="a5") +def ast_runtime_ifexp_python_literal_assign_probe(rows: pto.i32): + limit = pto.const(64, dtype=pto.index) + zero = pto.const(0, dtype=pto.index) + + for row in range(rows): + cnt = row if row < limit else 64 + out = cnt + zero + _ = out + + +@pto.jit(target="a5") +def ast_runtime_for_sibling_iv_reuse_probe(rows: pto.i32, cols: pto.i32): + stride = pto.const(64, dtype=pto.index) + one = pto.const(1, dtype=pto.index) + zero = pto.const(0, dtype=pto.index) + + for t in range(rows): + for c in range(cols): + col = c * stride + sink = col + one + _ = sink + + for stage in pto.static_range(2): + acc = zero + for c in range(cols): + col = c * stride + acc = acc + col + _ = acc + _ = stage + _ = t + + +@pto.jit(target="a5") +def ast_runtime_for_static_range_name_reuse_probe(cols: pto.i32): + zero = pto.const(0, dtype=pto.index) + + for phase_a_chunk in range(cols): + d_heads = [phase_a_chunk + h for h in pto.static_range(4)] + acc = zero + for h in pto.static_range(4): + acc = acc + d_heads[h] + for mi in pto.static_range(2): + inner = zero + for h in pto.static_range(4): + inner = inner + d_heads[h] + acc = acc + inner + _ = mi + _ = acc + + +@pto.jit(target="a5") +def ast_runtime_for_static_slot_carry_probe(cols: pto.i32): + zero = pto.const(0, dtype=pto.index) + accs = [zero for _ in pto.static_range(4)] + + for c in range(cols): + for h in pto.static_range(4): + accs[h] = accs[h] + c + + total = zero + for h in pto.static_range(4): + total = total + accs[h] + _ = total + + +@pto.jit(target="a5") +def ast_runtime_for_dynamic_slot_store_error_probe(cols: pto.i32): + zero = pto.const(0, dtype=pto.index) + accs = [zero for _ in pto.static_range(4)] + + for idx in range(cols): + accs[idx] = zero + + +class _StaticSlotHolder: + pass + + +@pto.jit(target="a5") +def ast_runtime_for_complex_slot_store_error_probe(cols: pto.i32): + zero = pto.const(0, dtype=pto.index) + holder = _StaticSlotHolder() + holder.accs = [zero for _ in pto.static_range(4)] + + for _ in range(cols): + for h in pto.static_range(4): + holder.accs[h] = zero + + @pto.jit(target="a5", ast_rewrite=False) def ast_rewrite_disabled_nested_helper_python_control_probe(): def helper(enabled): @@ -2254,6 +2357,207 @@ def public_surface_exports_probe( ) +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_wrapper_dispatch_probe(): + src_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + other_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + dst_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + int_src_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.i32) + int_other_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.i32) + hist_acc_tile = pto.alloc_tile(shape=[1, 256], dtype=pto.ui16) + hist_src_tile = pto.alloc_tile(shape=[1, 256], dtype=pto.ui8) + + src_ptr = src_tile.as_ptr() + other_ptr = other_tile.as_ptr() + dst_ptr = dst_tile.as_ptr() + int_src_ptr = int_src_tile.as_ptr() + int_other_ptr = int_other_tile.as_ptr() + hist_acc_ptr = hist_acc_tile.as_ptr() + hist_src_ptr = hist_src_tile.as_ptr() + + offset = pto.const(0, dtype=pto.index) + active_lanes = pto.const(64, dtype=pto.index) + active_per_group = pto.const(8, dtype=pto.index) + + mask = pto.vmi.create_mask( + active_lanes, + size=64, + ) + group_mask = pto.vmi.create_mask( + active_per_group, + size=64, + group=8, + ) + lhs = pto.vmi.vload(src_ptr, offset, size=64) + rhs = pto.vmi.vload(other_ptr, offset, size=64) + compact = pto.vmi.vload(src_ptr, offset, size=8) + idx = pto.vmi.vci(pto.i32(0), size=64, order="ASC") + bias = pto.vmi.vbrc(pto.f32(0.0), size=64) + expanded = pto.vmi.vbrc(compact, size=64, group=8) + hist_acc = pto.vmi.vload(hist_acc_ptr, offset, size=256) + hist_src = pto.vmi.vload(hist_src_ptr, offset, size=256) + hist_mask = pto.vmi.create_mask(pto.const(256, dtype=pto.index), size=256) + added = pto.vmi.vadd(lhs, rhs, mask) + relu = pto.vmi.vrelu(added, mask) + scaled = pto.vmi.vadd(pto.vmi.vmuls(relu, 2.0, mask), bias, mask) + pred = pto.vmi.vcmp(scaled, lhs, mask, "ogt") + selected = pto.vmi.vsel(pred, scaled, expanded) + shuffled = pto.vmi.vselr(selected, idx) + total = pto.vmi.vcadd(shuffled, mask, reassoc=True) + peak = pto.vmi.vcmax(shuffled, mask) + floor = pto.vmi.vcmin(shuffled, mask) + group_peak = pto.vmi.vcmax(shuffled, group_mask, group=8) + gather = pto.vmi.vgather(src_ptr, idx, mask) + gatherb = pto.vmi.vgatherb(src_ptr, idx, mask) + hist = pto.vmi.vdhist(hist_acc, hist_src, hist_mask) + cumul = pto.vmi.vchist(hist_acc, hist_src, hist_mask) + int_lhs = pto.vmi.vload(int_src_ptr, offset, size=64) + int_rhs = pto.vmi.vload(int_other_ptr, offset, size=64) + low, high = pto.vmi.vmull(int_lhs, int_rhs, mask) + widened = pto.vmi.vadd(low, high, mask) + casted = pto.vmi.vcvt(shuffled, pto.f16) + recast = pto.vmi.vinterpret_cast( + lhs, + pto.i32, + ) + lo, hi = pto.vmi.vintlv(selected, shuffled, mask) + pto.vmi.vstore(lo, dst_ptr, offset, mask) + + _ = group_mask + _ = total + _ = peak + _ = floor + _ = group_peak + _ = gather + _ = gatherb + _ = hist + _ = cumul + _ = widened + _ = casted + _ = recast + _ = hi + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_missing_binding_probe(): + src_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + src_ptr = src_tile.as_ptr() + offset = pto.const(0, dtype=pto.index) + active_lanes = pto.const(64, dtype=pto.index) + + src = pto.vmi.vload(src_ptr, offset, size=64) + mask = pto.vmi.create_mask( + active_lanes, + size=64, + ) + _ = pto.vmi.vadd(src, src, mask) + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_masked_vcvt_probe(): + src_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + src_ptr = src_tile.as_ptr() + offset = pto.const(0, dtype=pto.index) + active_lanes = pto.const(64, dtype=pto.index) + + src = pto.vmi.vload(src_ptr, offset, size=64) + mask = pto.vmi.create_mask( + active_lanes, + size=64, + ) + _ = pto.vmi.vcvt(src, pto.f16, mask) + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_invalid_rounding_vcvt_probe(): + src_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + src_ptr = src_tile.as_ptr() + offset = pto.const(0, dtype=pto.index) + + src = pto.vmi.vload(src_ptr, offset, size=64) + _ = pto.vmi.vcvt( + src, + pto.f8e4m3, + rounding=pto.VcvtRoundMode.F, + saturate=pto.VcvtSatMode.SAT, + ) + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_round_r_vcvt_probe(): + src_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + src_ptr = src_tile.as_ptr() + offset = pto.const(0, dtype=pto.index) + + src = pto.vmi.vload(src_ptr, offset, size=64) + _ = pto.vmi.vcvt( + src, + pto.f8e4m3, + rounding=pto.VcvtRoundMode.R, + saturate=pto.VcvtSatMode.SAT, + ) + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_unpack_vload_probe(): + src_tile = pto.alloc_tile(shape=[1, 128], dtype=pto.i8) + src_ptr = src_tile.as_ptr() + offset = pto.const(0, dtype=pto.index) + + _ = pto.vmi.vload( + src_ptr, + offset, + size=128, + dist_mode="unpack", + to_dtype=pto.i16, + ) + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_brc_vload_probe(): + src_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + src_ptr = src_tile.as_ptr() + offset = pto.const(0, dtype=pto.index) + + _ = pto.vmi.vload( + src_ptr, + offset, + size=64, + dist_mode="brc", + ) + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_group_brc_vload_probe(): + src_tile = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + src_ptr = src_tile.as_ptr() + offset = pto.const(0, dtype=pto.index) + row_stride = pto.const(1, dtype=pto.index) + + _ = pto.vmi.vload( + src_ptr, + offset, + size=64, + group=8, + stride=row_stride, + dist_mode="brc", + ) + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_unpack_vload_missing_dtype_probe(): + src_tile = pto.alloc_tile(shape=[1, 128], dtype=pto.i8) + src_ptr = src_tile.as_ptr() + offset = pto.const(0, dtype=pto.index) + + _ = pto.vmi.vload( + src_ptr, + offset, + size=128, + dist_mode="unpack", + ) + + @pto.jit(target="a5") def compile_time_query_probe(): f32_bw = pto.bytewidth(pto.f32) @@ -3063,9 +3367,14 @@ def main() -> None: expect(fake_empty.shape == fake_tensor.shape, "pto.empty_like(...) should preserve the logical tensor shape") expect(not hasattr(pto, "scalar"), "pto.scalar should not remain in the public pto namespace") expect(hasattr(pto, "tile"), "pto.tile should be exported from the public namespace") + expect(hasattr(pto, "vmi"), "pto.vmi should be exported from the public namespace") expect(hasattr(pto.tile, "load"), "pto.tile.load should be exported from the public tile namespace") expect(hasattr(pto.tile, "add"), "pto.tile.add should be exported from the public tile namespace") expect(hasattr(pto.tile, "cmps"), "pto.tile.cmps should be exported from the public tile namespace") + expect(hasattr(pto.vmi, "vreg"), "pto.vmi.vreg should be exported from the public VMI namespace") + expect(hasattr(pto.vmi, "mask"), "pto.vmi.mask should be exported from the public VMI namespace") + expect(hasattr(pto.vmi, "vadd"), "pto.vmi.vadd should be exported from the public VMI namespace") + expect(hasattr(pto.vmi, "create_mask"), "pto.vmi.create_mask should be exported from the public VMI namespace") expect(not hasattr(pto, "load_tile"), "pto.load_tile should not remain on the public pto namespace") expect(not hasattr(pto, "store_tile"), "pto.store_tile should not remain on the public pto namespace") expect(hasattr(pto.tile, "matmul"), "pto.tile.matmul should be exported from the public tile namespace") @@ -3085,6 +3394,9 @@ def main() -> None: expect(not hasattr(pto, "vbrc_load"), "pto.vbrc_load should not remain on the public pto namespace") expect(not hasattr(pto, "vsts_1pt"), "pto.vsts_1pt should not remain on the public pto namespace") expect(not hasattr(pto, "constexpr"), "pto.const_expr should not remain on the public pto namespace") + expect(not hasattr(pto, "copy_ubuf_to_ubuf"), "pto.copy_ubuf_to_ubuf should not remain on the public pto namespace") + expect(not hasattr(pto, "vmi_vreg_type"), "pto.vmi_vreg_type should not remain on the public pto namespace") + expect(not hasattr(pto, "vmi_mask_type"), "pto.vmi_mask_type should not remain on the public pto namespace") expect(not hasattr(scalar, "sts"), "scalar.sts should not remain in the public scalar namespace") expect(not hasattr(scalar, "cmpi"), "scalar.cmpi should not remain in the public scalar namespace") expect(not hasattr(scalar, "cmpi_sgt"), "scalar.cmpi_sgt should not remain in the public scalar namespace") @@ -3119,6 +3431,12 @@ def main() -> None: and "Use pto.const_expr" in str(removed_constexpr), "removed pto.constexpr should diagnose pto.const_expr as the replacement", ) + removed_copy_ubuf_to_ubuf = expect_raises(AttributeError, lambda: getattr(pto, "copy_ubuf_to_ubuf")) + expect( + "pto.copy_ubuf_to_ubuf is not a supported PTODSL public interface" in str(removed_copy_ubuf_to_ubuf) + and "Use pto.mte_ub_ub" in str(removed_copy_ubuf_to_ubuf), + "removed pto.copy_ubuf_to_ubuf should diagnose pto.mte_ub_ub as the replacement", + ) for name in ("max", "min", "exp", "log", "sqrt", "abs"): expect(hasattr(scalar, name), f"scalar.{name} should be exported from the public scalar namespace") @@ -3256,6 +3574,18 @@ def main() -> None: str(pto.mask_b32.resolve()) == "!pto.mask", "pto.mask_b32 should resolve to the public 32-bit mask type", ) + expect( + str(pto.vmi.vreg(128, pto.f32).resolve()) == "!pto.vmi.vreg<128xf32>", + "pto.vmi.vreg(...) should resolve to the public logical VMI vector type", + ) + expect( + str(pto.vmi.mask(128).resolve()) == "!pto.vmi.mask<128xpred>", + "pto.vmi.mask(...) should resolve to the public logical VMI mask type", + ) + expect( + str(pto.vmi.vreg(128, pto.f8e4m3).resolve()) == "!pto.vmi.vreg<128xf8E4M3FN>", + "pto.vmi.vreg(...) should preserve low-precision authored element types", + ) lp_tile_ty = pto_types.tile_buf_type([16, 16], pto.hif8, [16, 16]) lp_tv_ty = pto_types.tensor_view_type(2, pto.f8e4m3) @@ -4122,6 +4452,29 @@ def fake_run_ptoas_cmd(cmd, *, cwd=None): "--enable-insert-sync" in explicit_ptoas_cmd, "source-backed native build should still pass explicit/effective insert-sync to ptoas", ) + ptoas_cmds.clear() + vmi_mlir_text = vmi_wrapper_dispatch_probe.compile().mlir_text() + mlir_path.write_text(vmi_mlir_text, encoding="utf-8") + expect( + native_build_runtime._mlir_requires_enable_vmi(mlir_path), + "native build should detect PTODSL-generated VMI MLIR through the enable-vmi text probe", + ) + with mock.patch.object(native_build_runtime, "resolve_ptoas_binary", return_value=Path("/tmp/fake-ptoas")), mock.patch.object( + native_build_runtime, "_run", side_effect=fake_run_ptoas_cmd + ): + native_build_runtime._run_ptoas( + mlir_path, + kernel_object, + target_arch="a5", + ) + expect( + len(ptoas_cmds) == 1, + "native build should issue exactly one ptoas command for PTODSL-generated VMI MLIR", + ) + expect( + "--enable-vmi" in ptoas_cmds[0], + "native build should auto-enable the VMI semantic pipeline when PTODSL-generated MLIR contains VMI ops", + ) expect("valid=?" not in default_text, "default alloc_tile() should keep full static valid-shape when valid_shape= is omitted") auto_mode_violation = expect_raises( RuntimeError, @@ -4836,6 +5189,81 @@ def _enter_inline_simt_with_resource_attr(): "branch-local temporaries should not be inferred as loop-carried state", ) + ast_runtime_ifexp_assign_text = ast_runtime_ifexp_assign_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify( + ast_runtime_ifexp_assign_text, + "AST-rewritten runtime IfExp assignment specialization", + ) + expect( + ast_runtime_ifexp_assign_text.count("scf.for") == 1, + "assign-form Python conditional expressions inside runtime loops should preserve the runtime loop", + ) + expect( + ast_runtime_ifexp_assign_text.count("scf.if") == 1, + "assign-form Python conditional expressions should normalize through the existing AST if rewrite", + ) + + ast_runtime_ifexp_python_literal_assign_text = ( + ast_runtime_ifexp_python_literal_assign_probe.compile().mlir_text() + ) + expect_parse_roundtrip_and_verify( + ast_runtime_ifexp_python_literal_assign_text, + "AST-rewritten runtime IfExp assignment should materialize opposite-branch Python literals", + ) + expect( + ast_runtime_ifexp_python_literal_assign_text.count("scf.for") == 1, + "assign-form Python conditional expressions with opposite-branch literals should preserve the runtime loop", + ) + expect( + ast_runtime_ifexp_python_literal_assign_text.count("scf.if") == 1, + "assign-form Python conditional expressions with opposite-branch literals should normalize through the existing AST if rewrite", + ) + + ast_runtime_for_sibling_iv_reuse_text = ast_runtime_for_sibling_iv_reuse_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify( + ast_runtime_for_sibling_iv_reuse_text, + "AST-rewritten sibling runtime for IV reuse specialization", + ) + expect( + ast_runtime_for_sibling_iv_reuse_text.count("scf.for") >= 3, + "reusing a sibling runtime loop IV name should not be misdiagnosed as loop-target live-out", + ) + + ast_runtime_for_static_range_name_reuse_text = ast_runtime_for_static_range_name_reuse_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify( + ast_runtime_for_static_range_name_reuse_text, + "AST-rewritten runtime for static_range/list-comprehension name reuse specialization", + ) + expect( + ast_runtime_for_static_range_name_reuse_text.count("scf.for") == 1, + "static_range and comprehension-local names should not be inferred as outer runtime loop carry state", + ) + + ast_runtime_for_static_slot_carry_text = ast_runtime_for_static_slot_carry_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify( + ast_runtime_for_static_slot_carry_text, + "AST-rewritten runtime for static subscript slot carry specialization", + ) + expect( + ast_runtime_for_static_slot_carry_text.count("scf.for") == 1, + "static subscript slot carry should preserve the authored runtime loop", + ) + expect( + "iter_args(" in ast_runtime_for_static_slot_carry_text + and "scf.yield" in ast_runtime_for_static_slot_carry_text, + "static subscript slot carry should lower through scf.for iter_args", + ) + expect_raises( + PTODSLAstRewriteError, + lambda: ast_runtime_for_dynamic_slot_store_error_probe.compile(), + "simple_name[static_int_or_static_range_iv]", + ) + expect_raises( + PTODSLAstRewriteError, + lambda: ast_runtime_for_complex_slot_store_error_probe.compile(), + "simple_name[static_int_or_static_range_iv]", + ) + ast_rewrite_disabled_nested_helper_python_control_text = ( ast_rewrite_disabled_nested_helper_python_control_probe.compile().mlir_text() ) @@ -5359,8 +5787,136 @@ def _enter_inline_simt_with_resource_attr(): expect_parse_roundtrip_and_verify(vmulscvt_surface_text, "public vmulscvt surface specialization") vsstb_post_update_surface_text = vsstb_post_update_surface_probe.compile().mlir_text() expect_parse_roundtrip_and_verify(vsstb_post_update_surface_text, "vsstb post-update surface specialization") + vmi_wrapper_dispatch_text = vmi_wrapper_dispatch_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify(vmi_wrapper_dispatch_text, "public VMI wrapper dispatch specialization") + vmi_unpack_vload_text = vmi_unpack_vload_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify(vmi_unpack_vload_text, "public VMI unpack vload specialization") + vmi_brc_vload_text = vmi_brc_vload_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify(vmi_brc_vload_text, "public VMI brc vload specialization") + expect( + 'dist_mode = "brc"' in vmi_brc_vload_text, + "pto.vmi.vload should preserve the authored brc dist_mode without requiring a stride operand", + ) + vmi_group_brc_vload_text = vmi_group_brc_vload_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify(vmi_group_brc_vload_text, "public VMI grouped brc vload specialization") + expect( + 'dist_mode = "brc"' in vmi_group_brc_vload_text and "group = 8" in vmi_group_brc_vload_text, + "pto.vmi.vload should allow the grouped brc form exposed by the VMI IR contract", + ) fixed_width_integer_text = fixed_width_integer_specialization_probe.compile().mlir_text() expect_parse_roundtrip_and_verify(fixed_width_integer_text, "fixed-width integer specialization") + with mock.patch.object(vmi_namespace._pto, "vmi_vadd", None): + missing_vmi_binding = expect_raises( + NotImplementedError, + vmi_missing_binding_probe.compile, + "pto.vmi.vadd", + ) + expect( + "Rebuild PTO Python bindings" in str(missing_vmi_binding), + "missing generated VMI bindings should diagnose rebuild guidance", + ) + masked_vcvt_error = expect_raises( + NotImplementedError, + vmi_masked_vcvt_probe.compile, + "pto.vmi.vcvt", + ) + expect( + "masked form" in str(masked_vcvt_error), + "unsupported VMI backend/binding forms should diagnose the unavailable feature", + ) + invalid_rounding_vcvt_error = expect_raises( + ValueError, + vmi_invalid_rounding_vcvt_probe.compile, + "rounding", + ) + expect( + "expected one of A, H, R, Z" in str(invalid_rounding_vcvt_error), + "pto.vmi.vcvt should reject unsupported VMI rounding tokens before IR verification", + ) + vmi_round_r_vcvt_text = vmi_round_r_vcvt_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify(vmi_round_r_vcvt_text, "VMI R-rounding vcvt specialization") + expect( + 'rounding = "R"' in vmi_round_r_vcvt_text, + "pto.vmi.vcvt should preserve the authored R rounding token for fp32->fp8", + ) + unpack_missing_dtype_error = expect_raises( + TypeError, + vmi_unpack_vload_missing_dtype_probe.compile, + 'to_dtype when dist_mode="unpack"', + ) + expect( + 'to_dtype when dist_mode="unpack"' in str(unpack_missing_dtype_error), + "unpack vload without to_dtype should diagnose the missing widened element type", + ) + + expected_vmi_ops = [ + "pto.vmi.create_mask", + "pto.vmi.create_group_mask", + "pto.vmi.vload", + "pto.vmi.vci", + "pto.vmi.vadd", + "pto.vmi.vrelu", + "pto.vmi.vmuls", + "pto.vmi.vcmp", + "pto.vmi.vsel", + "pto.vmi.vselr", + "pto.vmi.vcadd", + "pto.vmi.vcmax", + "pto.vmi.vcmin", + "pto.vmi.vdhist", + "pto.vmi.vchist", + "pto.vmi.vmull", + "pto.vmi.vgather", + "pto.vmi.vcvt", + "pto.vmi.vinterpret_cast", + "pto.vmi.vintlv", + "pto.vmi.vstore", + ] + for op_name in expected_vmi_ops: + expect( + op_name in vmi_wrapper_dispatch_text, + f"representative {op_name} wrapper dispatch should emit the matching generated VMI op", + ) + expect( + vmi_wrapper_dispatch_text.count("pto.vmi.vload") == 7, + "vmi wrapper dispatch probe should lower seven explicit VMI loads", + ) + expect( + "pto.backend = \"vpto\"" in vmi_wrapper_dispatch_text, + "VMI public surface probe should compile through the VMI/VPTO backend partition", + ) + expect( + "!pto.vmi.vreg<64xf32>" in vmi_wrapper_dispatch_text, + "PTODSL VMI compile probes should materialize logical VMI vector result types in MLIR", + ) + expect( + "!pto.vmi.vreg<1xf32>" in vmi_wrapper_dispatch_text, + "PTODSL VMI reduction probes should materialize reduced logical VMI vector result types in MLIR", + ) + expect( + "!pto.vmi.vreg<64xf16>" in vmi_wrapper_dispatch_text, + "PTODSL VMI conversion probes should materialize converted logical VMI vector result types in MLIR", + ) + expect( + "!pto.vmi.vreg<64xi32>" in vmi_wrapper_dispatch_text, + "PTODSL VMI index/reinterpret probes should materialize integer logical VMI vector result types in MLIR", + ) + expect( + "!pto.vmi.mask<64xpred>" in vmi_wrapper_dispatch_text, + "PTODSL VMI compare and prefix mask probes should materialize logical VMI mask types in MLIR", + ) + expect( + "!pto.vmi.mask<64xpred>" in vmi_wrapper_dispatch_text, + "PTODSL grouped mask probes should materialize grouped logical VMI mask types in MLIR", + ) + expect( + "reassoc" in vmi_wrapper_dispatch_text, + "PTODSL VMI vcadd should preserve the reassoc attr in MLIR", + ) + expect( + "!pto.vmi.vreg<128xi16>" in vmi_unpack_vload_text, + "PTODSL VMI unpack vload should infer the widened logical VMI vector result type from to_dtype", + ) expect("pto.mte_gm_ub" in public_surface_text, "mte_load(...) should lower to pto.mte_gm_ub") expect("pto.mte_ub_gm" in public_surface_text, "mte_store(...) should lower to pto.mte_ub_gm") expect(public_surface_text.count("pto.mem_bar") >= 1, "mem_bar(...) should still lower explicit memory barriers") diff --git a/ptodsl/tests/test_jit_diagnostics.py b/ptodsl/tests/test_jit_diagnostics.py index a14d997815..56b56b557b 100644 --- a/ptodsl/tests/test_jit_diagnostics.py +++ b/ptodsl/tests/test_jit_diagnostics.py @@ -82,6 +82,92 @@ def same_width_float_store_probe(): scalar.store(f16_value, bf16_tile[0, 0]) +@pto.jit(target="a5") +def vmi_float_vcadd_missing_reassoc_probe(): + src = pto.vmi.vbrc(pto.f32(0.0), size=64) + mask = pto.vmi.create_mask(64, size=64) + _ = pto.vmi.vcadd(src, mask) + + +@pto.jit(target="a5") +def vmi_float_vcadd_none_reassoc_probe(): + src = pto.vmi.vbrc(pto.f32(0.0), size=64) + mask = pto.vmi.create_mask(64, size=64) + _ = pto.vmi.vcadd(src, mask, reassoc=None) + + +@pto.jit(target="a5") +def vmi_vbrc_untyped_scalar_probe(): + _ = pto.vmi.vbrc(0.0, size=64) + + +@pto.jit(target="a5") +def vmi_vci_untyped_scalar_probe(): + _ = pto.vmi.vci(0, size=64, order="ASC") + +@pto.jit(target="a5") +def vmi_vinterpret_cast_missing_dtype_probe(): + src = pto.vmi.vbrc(pto.f32(0.0), size=64) + _ = pto.vmi.vinterpret_cast(src) + + +@pto.jit(target="a5") +def vmi_vinterpret_cast_width_mismatch_probe(): + src = pto.vmi.vbrc(pto.f32(0.0), size=64) + _ = pto.vmi.vinterpret_cast(src, to_dtype=pto.f16) + + +@pto.jit(target="a5") +def vmi_create_mask_group_mismatch_probe(): + _ = pto.vmi.create_mask( + 9, + size=64, + group=8, + ) + + +@pto.jit(target="a5") +def vmi_vload_mixed_mode_probe(): + tile = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + src = tile.as_ptr() + offset = pto.const(0, dtype=pto.index) + _ = pto.vmi.vload(src, offset, size=64, dist_mode="continuous", group=8, stride=8) + + +@pto.jit(target="a5") +def vmi_vload_brc_stride_without_group_probe(): + tile = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + src = tile.as_ptr() + offset = pto.const(0, dtype=pto.index) + _ = pto.vmi.vload(src, offset, size=64, dist_mode="brc", stride=8) + + +@pto.jit(target="a5") +def vmi_vstore_group_mask_probe(): + tile = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + src = tile.as_ptr() + offset = pto.const(0, dtype=pto.index) + vec = pto.vmi.vload(src, offset, size=64) + mask = pto.vmi.create_mask(64, size=64) + pto.vmi.vstore(vec, src, offset, mask, group=8, stride=8) + + +@pto.jit(target="a5") +def vmi_vbrc_group_lane_mismatch_probe(): + tile = pto.alloc_tile(shape=[1, 16], dtype=pto.f32) + src = tile.as_ptr() + offset = pto.const(0, dtype=pto.index) + compact = pto.vmi.vload(src, offset, size=16) + _ = pto.vmi.vbrc(compact, size=64, group=8) + +@pto.jit(target="a5") +def vmi_vload_missing_size_probe(): + tile = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + src = tile.as_ptr() + offset = pto.const(0, dtype=pto.index) + _ = pto.vmi.vload(src, offset) + + @pto.jit(target="a5") def bool_loop_bound_probe(): with pto.for_(0, True, step=1): @@ -501,6 +587,17 @@ def assign_type_mismatch_probe(): br.assign(val=rhs) +@pto.jit(target="a5") +def assign_untyped_literal_without_anchor_probe(): + cond = pto.const(1, dtype=pto.i1) + with pto.if_(cond) as br: + with br.then_: + br.assign(val=1) + with br.else_: + br.assign(val=2) + _ = br.val + + @pto.jit(target="a5") def duplicate_assign_probe(): lhs = pto.const(4, dtype=pto.i32) @@ -580,6 +677,82 @@ def main() -> None: "f16", "bf16", ) + expect_raises( + vmi_float_vcadd_missing_reassoc_probe.compile, + TypeError, + "pto.vmi.vcadd(...)", + "floating-point vectors", + "reassoc", + "reassoc=True or reassoc=False", + ) + expect_raises( + vmi_float_vcadd_none_reassoc_probe.compile, + TypeError, + "pto.vmi.vcadd(...)", + "True or False", + "received None", + ) + expect_raises( + vmi_vbrc_untyped_scalar_probe.compile, + TypeError, + "pto.vmi.vbrc(...)", + "typed scalar", + "plain Python scalars are ambiguous", + ) + expect_raises( + vmi_vci_untyped_scalar_probe.compile, + TypeError, + "pto.vmi.vci(...)", + "typed scalar", + "plain Python scalars are ambiguous", + ) + expect_raises( + vmi_vinterpret_cast_missing_dtype_probe.compile, + TypeError, + "pto.vmi.vinterpret_cast(...)", + "requires to_dtype", + ) + expect_raises( + vmi_vinterpret_cast_width_mismatch_probe.compile, + TypeError, + "pto.vmi.vinterpret_cast(...)", + "element widths to match", + ) + expect_raises( + vmi_create_mask_group_mismatch_probe.compile, + ValueError, + "pto.vmi.create_mask(...)", + "active_lanes to be <= the inferred group_size", + ) + expect_raises( + vmi_vload_mixed_mode_probe.compile, + TypeError, + "pto.vmi.vload(...)", + "dist_mode together with group", + ) + expect_raises( + vmi_vload_brc_stride_without_group_probe.compile, + TypeError, + "pto.vmi.vload(...)", + "accepts stride only when group is provided", + ) + expect_raises( + vmi_vstore_group_mask_probe.compile, + TypeError, + "pto.vmi.vstore(...)", + "group mode does not take a mask operand", + ) + expect_raises( + vmi_vbrc_group_lane_mismatch_probe.compile, + ValueError, + "pto.vmi.vbrc(...)", + "input lane count to match group", + ) + expect_raises( + vmi_vload_missing_size_probe.compile, + TypeError, + "size", + ) expect_raises( bool_loop_bound_probe.compile, TypeError, @@ -1059,6 +1232,12 @@ def inline_count_mismatch(ptr: pto.ptr(pto.f32, "gm"), rows: pto.i32): RuntimeError, "br.assign(...) type mismatch for 'val'", ) + expect_raises( + assign_untyped_literal_without_anchor_probe.compile, + TypeError, + "br.assign(...) cannot infer a PTO type", + "materialize one side explicitly with pto.const(...)", + ) expect_raises( duplicate_assign_probe.compile, RuntimeError, diff --git a/ptodsl/tests/test_ptoas_wheel_launcher.py b/ptodsl/tests/test_ptoas_wheel_launcher.py index af8497034d..fe808d6911 100644 --- a/ptodsl/tests/test_ptoas_wheel_launcher.py +++ b/ptodsl/tests/test_ptoas_wheel_launcher.py @@ -17,19 +17,23 @@ class WheelLauncherTests(unittest.TestCase): - def _make_runtime_tree(self, temp_root: Path) -> Path: + def _make_runtime_tree(self, temp_root: Path) -> tuple[Path, Path, Path]: package_root = temp_root / "site-packages" / "ptoas" runtime_root = package_root / "_runtime" (runtime_root / "bin").mkdir(parents=True, exist_ok=True) (runtime_root / "lib").mkdir(parents=True, exist_ok=True) (runtime_root / "share" / "ptoas" / "TileOps").mkdir(parents=True, exist_ok=True) + wrapper = temp_root / "bin" / "ptoas" + wrapper.parent.mkdir(parents=True, exist_ok=True) + wrapper.write_text("", encoding="utf-8") (temp_root / "site-packages" / "ptodsl").mkdir(parents=True, exist_ok=True) (temp_root / "site-packages" / "tilelang_dsl").mkdir(parents=True, exist_ok=True) (temp_root / "site-packages" / "mlir").mkdir(parents=True, exist_ok=True) - (runtime_root / "bin" / "ptoas").write_text("", encoding="utf-8") - return package_root + shared_module = runtime_root / "lib" / "ptoas.so" + shared_module.write_text("fake shared module", encoding="utf-8") + return package_root, wrapper, shared_module - def _make_editable_tree(self, temp_root: Path) -> tuple[Path, Path]: + def _make_editable_tree(self, temp_root: Path) -> tuple[Path, Path, Path]: repo_root = temp_root / "repo" package_root = repo_root / "ptodsl" / "ptoas" install_root = repo_root / "install" @@ -39,42 +43,52 @@ def _make_editable_tree(self, temp_root: Path) -> tuple[Path, Path]: (install_root / "share" / "ptoas" / "TileOps").mkdir(parents=True, exist_ok=True) (install_root / "tilelang_dsl").mkdir(parents=True, exist_ok=True) (install_root / "mlir").mkdir(parents=True, exist_ok=True) - (install_root / "bin" / "ptoas").write_text("", encoding="utf-8") - return package_root, install_root + wrapper = install_root / "bin" / "ptoas" + wrapper.write_text("", encoding="utf-8") + shared_module = install_root / "lib" / "ptoas.so" + shared_module.write_text("fake shared module", encoding="utf-8") + return package_root, install_root, wrapper def test_launcher_exports_runtime_contract_and_injects_default_paths(self): with tempfile.TemporaryDirectory() as temp_dir: temp_root = Path(temp_dir) - package_root = self._make_runtime_tree(temp_root) + package_root, wrapper, shared_module = self._make_runtime_tree(temp_root) fake_launcher = package_root / "_launcher.py" fake_launcher.write_text("", encoding="utf-8") with mock.patch.dict(_launcher.os.environ, {}, clear=True), mock.patch.object( _launcher, "__file__", str(fake_launcher) ), mock.patch.object( - _launcher.sys, "argv", ["ptoas", "--version"] - ), mock.patch.object(_launcher.os, "execvpe") as execvpe: - _launcher.main() - - execvpe.assert_called_once() - binary, argv, env = execvpe.call_args.args - self.assertEqual(binary, str(package_root / "_runtime" / "bin" / "ptoas")) - self.assertEqual(argv[:5], [ - str(package_root / "_runtime" / "bin" / "ptoas"), - "--tilelang-path", - str(package_root / "_runtime" / "share" / "ptoas" / "TileOps"), - "--tilelang-pkg-path", - str(package_root.parent), - ]) - self.assertEqual(argv[-1], "--version") - self.assertEqual(env["PTOAS_HOME"], str(package_root / "_runtime")) - self.assertEqual(env["PTOAS_BIN"], str(package_root / "_runtime" / "bin" / "ptoas")) - self.assertEqual( - env["PTOAS_TILEOPS_DIR"], - str(package_root / "_runtime" / "share" / "ptoas" / "TileOps"), - ) - self.assertTrue(env["PATH"].split(os.pathsep)[0].endswith("ptoas/_runtime/bin")) - self.assertEqual(env["PYTHONPATH"].split(os.pathsep)[0], str(package_root.parent)) + _launcher.sys, "argv", [str(wrapper), "--version"] + ), mock.patch.object(_launcher, "_load_shared_entrypoint", return_value=mock.Mock(return_value=0)) as load_entrypoint: + with self.assertRaises(SystemExit) as exc: + _launcher.main() + + self.assertEqual(exc.exception.code, 0) + load_entrypoint.assert_called_once_with(shared_module, package_root / "_runtime" / "lib") + entrypoint = load_entrypoint.return_value + self.assertEqual(entrypoint.call_count, 1) + argc, c_argv = entrypoint.call_args.args + self.assertEqual(argc, 6) + self.assertEqual([ + c_argv[i].decode("utf-8") for i in range(argc) + ], [ + str(wrapper), + "--tilelang-path", + str(package_root / "_runtime" / "share" / "ptoas" / "TileOps"), + "--tilelang-pkg-path", + str(package_root.parent), + "--version", + ]) + env = _launcher.os.environ + self.assertEqual(env["PTOAS_HOME"], str(package_root / "_runtime")) + self.assertEqual(env["PTOAS_BIN"], str(wrapper)) + self.assertEqual( + env["PTOAS_TILEOPS_DIR"], + str(package_root / "_runtime" / "share" / "ptoas" / "TileOps"), + ) + self.assertEqual(env["PATH"].split(os.pathsep)[0], str(wrapper.parent)) + self.assertEqual(env["PYTHONPATH"].split(os.pathsep)[0], str(package_root.parent)) def test_resolve_runtime_root_defaults_to_repo_install_tree(self): package_root = Path("/tmp/repo/ptodsl/ptoas") @@ -82,10 +96,56 @@ def test_resolve_runtime_root_defaults_to_repo_install_tree(self): runtime_root = _launcher._resolve_runtime_root(package_root) self.assertEqual(runtime_root, Path("/tmp/repo/install")) + def test_resolve_wrapper_path_falls_back_to_shutil_lookup(self): + with mock.patch.object(_launcher.sys, "argv", ["ptoas"]), mock.patch.object( + _launcher.shutil, "which", return_value="/tmp/bin/ptoas" + ): + wrapper = _launcher._resolve_wrapper_path() + + self.assertEqual(wrapper, Path("/tmp/bin/ptoas")) + + def test_load_shared_entrypoint_configures_in_process_ctypes_call(self): + with tempfile.TemporaryDirectory() as temp_dir: + temp_root = Path(temp_dir) + runtime_lib_dir = temp_root / "runtime" / "lib" + runtime_lib_dir.mkdir(parents=True, exist_ok=True) + dep = runtime_lib_dir / "libMLIRMlirOptMain.so.21.1" + dep.write_text("fake dep", encoding="utf-8") + shared_module = temp_root / "runtime" / "pto" / "ptoas.so" + shared_module.parent.mkdir(parents=True, exist_ok=True) + shared_module.write_text("fake shared module", encoding="utf-8") + + def fake_ldd(cmd, *, text, stderr, env): + target = Path(cmd[1]).resolve() + if target == shared_module.resolve(): + return f"\tlibMLIRMlirOptMain.so.21.1 => {dep} (0x0)\n" + return "" + + dep_library = mock.Mock() + shared_library = mock.Mock() + with mock.patch.object(_launcher.subprocess, "check_output", side_effect=fake_ldd), mock.patch.object( + _launcher.ctypes, "CDLL", side_effect=[dep_library, shared_library] + ) as load_library: + entrypoint = _launcher._load_shared_entrypoint(shared_module, runtime_lib_dir) + + self.assertEqual( + [call.args[0] for call in load_library.call_args_list], + [str(dep.resolve()), str(shared_module.resolve())], + ) + self.assertTrue( + all( + call.kwargs["mode"] == getattr(_launcher.ctypes, "RTLD_GLOBAL", 0) + for call in load_library.call_args_list + ) + ) + self.assertIs(entrypoint, shared_library.ptoas_entrypoint) + self.assertEqual(entrypoint.argtypes, [_launcher.ctypes.c_int, _launcher.ctypes.POINTER(_launcher.ctypes.c_char_p)]) + self.assertIs(entrypoint.restype, _launcher.ctypes.c_int) + def test_launcher_falls_back_to_env_install_tree_for_editable_installs(self): with tempfile.TemporaryDirectory() as temp_dir: temp_root = Path(temp_dir) - package_root, install_root = self._make_editable_tree(temp_root) + package_root, install_root, wrapper = self._make_editable_tree(temp_root) fake_launcher = package_root / "_launcher.py" fake_launcher.write_text("", encoding="utf-8") @@ -94,34 +154,41 @@ def test_launcher_falls_back_to_env_install_tree_for_editable_installs(self): {"PTO_INSTALL_DIR": str(install_root)}, clear=True, ), mock.patch.object( - _launcher.sys, "argv", ["ptoas", "--version"] - ), mock.patch.object(_launcher.os, "execvpe") as execvpe: - _launcher.main() - - execvpe.assert_called_once() - binary, argv, env = execvpe.call_args.args - self.assertEqual(binary, str(install_root / "bin" / "ptoas")) - self.assertEqual(argv[:5], [ - str(install_root / "bin" / "ptoas"), - "--tilelang-path", - str(install_root / "share" / "ptoas" / "TileOps"), - "--tilelang-pkg-path", - str(install_root), - ]) - self.assertEqual(argv[-1], "--version") - self.assertEqual(env["PTOAS_HOME"], str(install_root)) - self.assertEqual(env["PTOAS_BIN"], str(install_root / "bin" / "ptoas")) - self.assertEqual( - env["PTOAS_TILEOPS_DIR"], - str(install_root / "share" / "ptoas" / "TileOps"), - ) - self.assertTrue(env["PATH"].split(os.pathsep)[0].endswith("install/bin")) - self.assertEqual(env["PYTHONPATH"].split(os.pathsep)[0], str(install_root)) + _launcher, "__file__", str(fake_launcher) + ), mock.patch.object( + _launcher.sys, "argv", [str(wrapper), "--version"] + ), mock.patch.object(_launcher, "_load_shared_entrypoint", return_value=mock.Mock(return_value=0)) as load_entrypoint: + with self.assertRaises(SystemExit) as exc: + _launcher.main() + + self.assertEqual(exc.exception.code, 0) + load_entrypoint.assert_called_once_with(install_root / "lib" / "ptoas.so", install_root / "lib") + entrypoint = load_entrypoint.return_value + argc, c_argv = entrypoint.call_args.args + self.assertEqual([ + c_argv[i].decode("utf-8") for i in range(argc) + ], [ + str(wrapper), + "--tilelang-path", + str(install_root / "share" / "ptoas" / "TileOps"), + "--tilelang-pkg-path", + str(install_root), + "--version", + ]) + env = _launcher.os.environ + self.assertEqual(env["PTOAS_HOME"], str(install_root)) + self.assertEqual(env["PTOAS_BIN"], str(wrapper)) + self.assertEqual( + env["PTOAS_TILEOPS_DIR"], + str(install_root / "share" / "ptoas" / "TileOps"), + ) + self.assertEqual(env["PATH"].split(os.pathsep)[0], str(wrapper.parent)) + self.assertEqual(env["PYTHONPATH"].split(os.pathsep)[0], str(install_root)) def test_launcher_respects_explicit_tilelang_flags(self): with tempfile.TemporaryDirectory() as temp_dir: temp_root = Path(temp_dir) - package_root = self._make_runtime_tree(temp_root) + package_root, wrapper, shared_module = self._make_runtime_tree(temp_root) fake_launcher = package_root / "_launcher.py" fake_launcher.write_text("", encoding="utf-8") @@ -129,25 +196,92 @@ def test_launcher_respects_explicit_tilelang_flags(self): _launcher.sys, "argv", [ - "ptoas", + str(wrapper), "--tilelang-path=/tmp/custom-tileops", "--tilelang-pkg-path", "/tmp/custom-python", "--help", ], - ), mock.patch.object(_launcher.os, "execvpe") as execvpe: - _launcher.main() + ), mock.patch.object(_launcher, "_load_shared_entrypoint", return_value=mock.Mock(return_value=0)) as load_entrypoint: + with self.assertRaises(SystemExit) as exc: + _launcher.main() - execvpe.assert_called_once() - _, argv, _ = execvpe.call_args.args - self.assertEqual(argv, [ - str(package_root / "_runtime" / "bin" / "ptoas"), + self.assertEqual(exc.exception.code, 0) + entrypoint = load_entrypoint.return_value + argc, c_argv = entrypoint.call_args.args + self.assertEqual([ + c_argv[i].decode("utf-8") for i in range(argc) + ], [ + str(wrapper), "--tilelang-path=/tmp/custom-tileops", "--tilelang-pkg-path", "/tmp/custom-python", "--help", ]) + def test_launcher_prefers_local_shared_module_next_to_build_wrapper(self): + with tempfile.TemporaryDirectory() as temp_dir: + temp_root = Path(temp_dir) + package_root = temp_root / "build" / "python" / "ptoas" + package_root.mkdir(parents=True, exist_ok=True) + fake_launcher = package_root / "_launcher.py" + fake_launcher.write_text("", encoding="utf-8") + wrapper = temp_root / "build" / "tools" / "ptoas" / "ptoas" + wrapper.parent.mkdir(parents=True, exist_ok=True) + wrapper.write_text("", encoding="utf-8") + shared_module = wrapper.parent / "ptoas.so" + shared_module.write_text("fake shared module", encoding="utf-8") + + with mock.patch.dict( + _launcher.os.environ, + {"PTO_INSTALL_DIR": str(temp_root / "install")}, + clear=True, + ), mock.patch.object( + _launcher, "__file__", str(fake_launcher) + ), mock.patch.object( + _launcher.sys, "argv", [str(wrapper), "--version"] + ), mock.patch.object( + _launcher, "_load_shared_entrypoint", return_value=mock.Mock(return_value=0) + ) as load_entrypoint: + with self.assertRaises(SystemExit) as exc: + _launcher.main() + + self.assertEqual(exc.exception.code, 0) + load_entrypoint.assert_called_once_with(shared_module, temp_root / "install" / "lib") + + def test_launcher_skips_empty_placeholder_shared_module(self): + with tempfile.TemporaryDirectory() as temp_dir: + temp_root = Path(temp_dir) + package_root = temp_root / "build" / "python" / "ptoas" + package_root.mkdir(parents=True, exist_ok=True) + fake_launcher = package_root / "_launcher.py" + fake_launcher.write_text("", encoding="utf-8") + wrapper = temp_root / "build" / "tools" / "ptoas" / "ptoas" + wrapper.parent.mkdir(parents=True, exist_ok=True) + wrapper.write_text("", encoding="utf-8") + (temp_root / "build" / "python" / "pto").mkdir(parents=True, exist_ok=True) + (temp_root / "build" / "python" / "pto" / "ptoas.so").write_text("", encoding="utf-8") + install_shared_module = temp_root / "install" / "lib" / "ptoas.so" + install_shared_module.parent.mkdir(parents=True, exist_ok=True) + install_shared_module.write_text("fake shared module", encoding="utf-8") + + with mock.patch.dict( + _launcher.os.environ, + {"PTO_INSTALL_DIR": str(temp_root / "install")}, + clear=True, + ), mock.patch.object( + _launcher, "__file__", str(fake_launcher) + ), mock.patch.object( + _launcher.sys, "argv", [str(wrapper), "--version"] + ), mock.patch.object( + _launcher, "_load_shared_entrypoint", return_value=mock.Mock(return_value=0) + ) as load_entrypoint: + with self.assertRaises(SystemExit) as exc: + _launcher.main() + + self.assertEqual(exc.exception.code, 0) + load_entrypoint.assert_called_once_with(install_shared_module, temp_root / "install" / "lib") + if __name__ == "__main__": unittest.main() diff --git a/ptodsl/tests/test_tilelib_catalog.py b/ptodsl/tests/test_tilelib_catalog.py index c027e989e3..ee86ad5a9b 100644 --- a/ptodsl/tests/test_tilelib_catalog.py +++ b/ptodsl/tests/test_tilelib_catalog.py @@ -981,7 +981,7 @@ def test_tsort32_unaligned_tmp_version_renders(self): selected = select("pto.tsort32", "a5", specs) self.assertEqual(selected.name, "template_tsort32_with_tmp") mlir = selected.specialize(**specs).mlir_text() - self.assertIn("pto.copy_ubuf_to_ubuf", mlir) + self.assertIn("pto.mte_ub_ub", mlir) self.assertIn("pto.vbitsort", mlir) def test_tsort32_unaligned_tmp_uses_valid_width(self): @@ -994,7 +994,7 @@ def test_tsort32_unaligned_tmp_uses_valid_width(self): selected = select("pto.tsort32", "a5", specs) self.assertEqual(selected.name, "template_tsort32_with_tmp") mlir = selected.specialize(**specs).mlir_text() - self.assertIn("pto.copy_ubuf_to_ubuf", mlir) + self.assertIn("pto.mte_ub_ub", mlir) self.assertIn("pto.vbitsort", mlir) self.assertIn("%c130", mlir) diff --git a/ptodsl/tests/test_vector_cube_ops.py b/ptodsl/tests/test_vector_cube_ops.py index cfa7c24a35..5fd27cfec0 100644 --- a/ptodsl/tests/test_vector_cube_ops.py +++ b/ptodsl/tests/test_vector_cube_ops.py @@ -87,8 +87,8 @@ def test_public_namespace_exports_new_vector_and_cube_apis(self): "vsub", "vmin", "vand", "vor", "vxor", "vshl", "vshr", "vln", "vsqrt", "vabs", "vneg", "vrec", "vrsqrt", "vrelu", "vnot", "vcmin", "vcgmin", "vcpadd", - "vadds", "vmuls", "vmaxs", "vmins", "vlrelu", - "vaxpy", "vaddrelu", "vsubrelu", "vsel", + "vadds", "vmuls", "vmaxs", "vmins", "vlrelu", "vshls", "vshrs", "vands", "vors", "vxors", + "vaxpy", "vmula", "vci", "vaddrelu", "vsubrelu", "vsel", "mte_gm_l1", "mte_l1_ub", "mte_gm_l1_frac", "mte_l1_bt", "mte_l1_fb", "mad_acc", "mad_bias", "mad_mx", "mad_mx_acc", "mad_mx_bias", "FractalMode", "AccStoreUnitFlagCtrl", "MadUnitFlagMode", "SatMode", "Tf32Mode", "SplitMode", @@ -282,6 +282,98 @@ def test_vcgmin_and_vsel_dispatch_correctly(self): self.assertIs(output, selected) self.assertEqual(vsel_op.call_args.args, ("vec_ty", vec, other, mask)) + def test_vector_scalar_helper_dispatches_cover_shift_and_bitwise_scalar_helpers(self): + vec = SimpleNamespace(type="vec_ty") + mask = SimpleNamespace(type="mask_ty") + scalar = SimpleNamespace(type="scalar_ty") + scalar_i16 = SimpleNamespace(type="i16_ty") + shifted = object() + anded = object() + ored = object() + xored = object() + broadcast = object() + + with patch.object(_ops, "unwrap_surface_value", side_effect=_identity), \ + patch.object(_ops, "wrap_surface_value", side_effect=_identity), \ + patch.object(_ops, "_reject_low_precision_vreg_operands") as reject_lp, \ + patch.object(_ops.IntegerType, "get_signless", return_value="i16_ty") as get_signless, \ + patch.object(_ops, "coerce_scalar_to_type", return_value=scalar_i16) as coerce_i16, \ + patch.object(_ops._pto, "VshlsOp", return_value=SimpleNamespace(result=shifted)) as vshls_op, \ + patch.object(_ops._pto, "VshrsOp", return_value=SimpleNamespace(result=shifted)) as vshrs_op: + self.assertIs(_ops.vshls(vec, scalar, mask), shifted) + self.assertIs(_ops.vshrs(vec, scalar, mask), shifted) + self.assertEqual(reject_lp.call_count, 2) + self.assertEqual(get_signless.call_count, 2) + self.assertEqual(coerce_i16.call_count, 2) + self.assertEqual(vshls_op.call_args.args, ("vec_ty", vec, scalar_i16, mask)) + self.assertEqual(vshrs_op.call_args.args, ("vec_ty", vec, scalar_i16, mask)) + + with patch.object(_ops, "_coerce_scalar_like_vector_element", return_value=scalar) as coerce_scalar, \ + patch.object(_ops, "vbr", return_value=broadcast) as vbr, \ + patch.object(_ops, "vand", return_value=anded) as vand: + self.assertIs(_ops.vands(vec, scalar, mask), anded) + coerce_scalar.assert_called_once_with(vec, scalar, context="vands") + vbr.assert_called_once_with(scalar) + vand.assert_called_once_with(vec, broadcast, mask) + + with patch.object(_ops, "_coerce_scalar_like_vector_element", return_value=scalar) as coerce_scalar, \ + patch.object(_ops, "vbr", return_value=broadcast) as vbr, \ + patch.object(_ops, "vor", return_value=ored) as vor: + self.assertIs(_ops.vors(vec, scalar, mask), ored) + coerce_scalar.assert_called_once_with(vec, scalar, context="vors") + vbr.assert_called_once_with(scalar) + vor.assert_called_once_with(vec, broadcast, mask) + + with patch.object(_ops, "_coerce_scalar_like_vector_element", return_value=scalar) as coerce_scalar, \ + patch.object(_ops, "vbr", return_value=broadcast) as vbr, \ + patch.object(_ops, "vxor", return_value=xored) as vxor: + self.assertIs(_ops.vxors(vec, scalar, mask), xored) + coerce_scalar.assert_called_once_with(vec, scalar, context="vxors") + vbr.assert_called_once_with(scalar) + vxor.assert_called_once_with(vec, broadcast, mask) + + + def test_vlds_accepts_extended_distribution_tokens(self): + ptr = SimpleNamespace(type="ptr_ty") + vec = object() + + with patch.object(_ops, "unwrap_surface_value", side_effect=_identity), \ + patch.object(_ops, "wrap_surface_value", side_effect=_identity), \ + patch.object(_ops, "_infer_vreg_type_from_address_source", return_value="vec_ty"), \ + patch.object(_ops, "_coerce_index", return_value="idx"), \ + patch.object(_ops, "_normalize_post_update_mode", return_value="NO_POST_UPDATE"), \ + patch.object(_ops, "_normalize_dist_token", side_effect=lambda dist, *, allowed, context: dist) as normalize_dist, \ + patch.object(_ops._pto, "VldsOp", return_value=SimpleNamespace(result=vec)) as vlds_op: + self.assertIs(_ops.vlds(ptr, 0, dist="E2B_B16"), vec) + self.assertIs(_ops.vlds(ptr, 0, dist="BRC_BLK"), vec) + self.assertEqual(normalize_dist.call_args_list[0].args[0], "E2B_B16") + self.assertEqual(normalize_dist.call_args_list[1].args[0], "BRC_BLK") + self.assertEqual(vlds_op.call_count, 2) + + def test_f32_to_fp8_vcvt_contract_accepts_rahz_rounding(self): + contract = _ops._VCVT_CONTRACTS[("f32", "f8e4m3")] + for rnd in ("R", "A", "H", "Z"): + with self.subTest(rnd=rnd): + _ops._validate_vcvt_attrs( + "f32", + "f8e4m3", + contract, + rnd=rnd, + sat="SAT", + part="P0", + context="pto.vcvt", + ) + with self.assertRaisesRegex(ValueError, r"expected one of A, H, R, Z"): + _ops._validate_vcvt_attrs( + "f32", + "f8e4m3", + contract, + rnd="F", + sat="SAT", + part="P0", + context="pto.vcvt", + ) + def test_cube_variant_wrappers_dispatch_to_generated_ops(self): lhs = object() rhs = object() diff --git a/ptodsl/tests/test_vmi_vmull.py b/ptodsl/tests/test_vmi_vmull.py new file mode 100644 index 0000000000..3fbf2972a4 --- /dev/null +++ b/ptodsl/tests/test_vmi_vmull.py @@ -0,0 +1,72 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +from pathlib import Path +import sys + + +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "ptodsl")) + +from ptodsl import pto + + +def expect(condition: bool, message: str) -> None: + if not condition: + raise AssertionError(message) + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_vmull_inferred_probe(): + lhs_tile = pto.alloc_tile(shape=[1, 256], dtype=pto.i32) + rhs_tile = pto.alloc_tile(shape=[1, 256], dtype=pto.i32) + dst_tile = pto.alloc_tile(shape=[1, 256], dtype=pto.i32) + offset = pto.const(0, dtype=pto.index) + active_lanes = pto.const(256, dtype=pto.index) + mask = pto.vmi.create_mask(active_lanes, size=256) + lhs = pto.vmi.vload(lhs_tile.as_ptr(), offset, size=256) + rhs = pto.vmi.vload(rhs_tile.as_ptr(), offset, size=256) + low, high = pto.vmi.vmull(lhs, rhs, mask) + pto.vmi.vstore(low, dst_tile.as_ptr(), offset, mask) + _ = high + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_vmull_unsigned_zero_probe(): + lhs_tile = pto.alloc_tile(shape=[1, 128], dtype=pto.ui32) + rhs_tile = pto.alloc_tile(shape=[1, 128], dtype=pto.ui32) + offset = pto.const(0, dtype=pto.index) + active_lanes = pto.const(128, dtype=pto.index) + mask = pto.vmi.create_mask(active_lanes, size=128) + lhs = pto.vmi.vload(lhs_tile.as_ptr(), offset, size=128) + rhs = pto.vmi.vload(rhs_tile.as_ptr(), offset, size=128) + low, high = pto.vmi.vmull(lhs, rhs, mask, pmode="zero") + _ = low + _ = high + + +def main() -> None: + inferred_text = vmi_vmull_inferred_probe.compile().mlir_text() + expect(inferred_text.count("pto.vmi.vmull") == 1, "inferred probe must emit one VMI VMULL") + expect( + "-> !pto.vmi.vreg<256xi32>, !pto.vmi.vreg<256xi32>" in inferred_text, + "VMULL must infer both 256xi32 results", + ) + + unsigned_text = vmi_vmull_unsigned_zero_probe.compile().mlir_text() + expect(unsigned_text.count("pto.vmi.vmull") == 1, "unsigned probe must emit one VMI VMULL") + expect( + "-> !pto.vmi.vreg<128xui32>, !pto.vmi.vreg<128xui32>" in unsigned_text, + "VMULL must infer both 128xui32 results", + ) + expect('pmode = "zero"' in unsigned_text, "explicit zero pmode must be preserved") + print("ptodsl_vmi_vmull: PASS") + + +if __name__ == "__main__": + main() diff --git a/ptodsl/tests/test_vmi_vshr_signedness.py b/ptodsl/tests/test_vmi_vshr_signedness.py new file mode 100644 index 0000000000..f4dfe09340 --- /dev/null +++ b/ptodsl/tests/test_vmi_vshr_signedness.py @@ -0,0 +1,74 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +from pathlib import Path +import sys + + +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "ptodsl")) + +from ptodsl import pto + + +def expect(condition: bool, message: str) -> None: + if not condition: + raise AssertionError(message) + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_vshr_signed_probe(): + lhs_tile = pto.alloc_tile(shape=[1, 128], dtype=pto.si32) + rhs_tile = pto.alloc_tile(shape=[1, 128], dtype=pto.si32) + offset = pto.const(0, dtype=pto.index) + active_lanes = pto.const(128, dtype=pto.index) + mask = pto.vmi.create_mask(active_lanes, size=128) + lhs = pto.vmi.vload(lhs_tile.as_ptr(), offset, size=128) + rhs = pto.vmi.vload(rhs_tile.as_ptr(), offset, size=128) + shifted = pto.vmi.vshr(lhs, rhs, mask) + shifted_scalar = pto.vmi.vshrs(lhs, pto.si32(3), mask) + _ = shifted + _ = shifted_scalar + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_vshr_unsigned_probe(): + lhs_tile = pto.alloc_tile(shape=[1, 128], dtype=pto.ui32) + rhs_tile = pto.alloc_tile(shape=[1, 128], dtype=pto.ui32) + offset = pto.const(0, dtype=pto.index) + active_lanes = pto.const(128, dtype=pto.index) + mask = pto.vmi.create_mask(active_lanes, size=128) + lhs = pto.vmi.vload(lhs_tile.as_ptr(), offset, size=128) + rhs = pto.vmi.vload(rhs_tile.as_ptr(), offset, size=128) + shifted = pto.vmi.vshr(lhs, rhs, mask) + shifted_scalar = pto.vmi.vshrs(lhs, pto.ui32(3), mask) + _ = shifted + _ = shifted_scalar + + +def main() -> None: + signed_text = vmi_vshr_signed_probe.compile().mlir_text() + expect("pto.vmi.vshr" in signed_text, "signed probe must emit pto.vmi.vshr") + expect("pto.vmi.vshrs" in signed_text, "signed probe must emit pto.vmi.vshrs") + expect( + "!pto.vmi.vreg<128xsi32>" in signed_text, + "signed probe must preserve the explicit si32 VMI element type", + ) + + unsigned_text = vmi_vshr_unsigned_probe.compile().mlir_text() + expect("pto.vmi.vshr" in unsigned_text, "unsigned probe must emit pto.vmi.vshr") + expect("pto.vmi.vshrs" in unsigned_text, "unsigned probe must emit pto.vmi.vshrs") + expect( + "!pto.vmi.vreg<128xui32>" in unsigned_text, + "unsigned probe must preserve the explicit ui32 VMI element type", + ) + print("ptodsl_vmi_vshr_signedness: PASS") + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index e38456cbe5..b34908352a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,7 @@ backend-path = ["."] [project] name = "ptoas" -version = "0.1.0" +version = "0.51" description = "PTO Assembler & Optimizer" readme = "README.md" requires-python = ">=3.9" diff --git a/python/pto/dialects/pto.py b/python/pto/dialects/pto.py index 5442f100e4..3e66f0bdb1 100644 --- a/python/pto/dialects/pto.py +++ b/python/pto/dialects/pto.py @@ -119,6 +119,8 @@ def _export_optional_cext_symbol(name): PtrType = _pto_mod.PtrType VRegType = _pto_mod.VRegType MaskType = _pto_mod.MaskType +VMIVRegType = _pto_mod.VMIVRegType +VMIMaskType = _pto_mod.VMIMaskType AlignType = _pto_mod.AlignType AsyncSessionType = _pto_mod.AsyncSessionType AsyncEventType = _pto_mod.AsyncEventType @@ -296,6 +298,8 @@ def fence_scope_attr_builder(value, context=None): "PtrType", "VRegType", "MaskType", + "VMIVRegType", + "VMIMaskType", "AlignType", "AsyncSessionType", "AsyncEventType", diff --git a/quick_install.sh b/quick_install.sh index 38509b61d5..6cf8effddf 100755 --- a/quick_install.sh +++ b/quick_install.sh @@ -37,9 +37,12 @@ PTOAS_VERSION="${PTOAS_VERSION:-$(python "${PTO_SOURCE_DIR}/.github/scripts/comp cd "$PTO_SOURCE_DIR" +# Release (-O3) is required with LinuxHardeningCache: _FORTIFY_SOURCE=2 +# errors under -Werror when optimization is off (default cmake build type). cmake -C "${PTO_SOURCE_DIR}/cmake/LinuxHardeningCache.cmake" -G Ninja \ -S . \ -B build \ + -DCMAKE_BUILD_TYPE=Release \ -DLLVM_DIR="${LLVM_BUILD_DIR}/lib/cmake/llvm" \ -DMLIR_DIR="${LLVM_BUILD_DIR}/lib/cmake/mlir" \ -DPython3_ROOT_DIR="${PY_ROOT}" \ @@ -58,9 +61,12 @@ export PTOAS_PYTHON_PACKAGE_VERSION="${PTOAS_PYTHON_PACKAGE_VERSION:-${PTOAS_VER bash "${PTO_SOURCE_DIR}/docker/create_wheel.sh" shopt -s nullglob -wheels=("${MLIR_PY_PKG}/dist/ptoas-"*.whl) +wheels=("${PTO_SOURCE_DIR}/build/wheel-dist/ptoas-"*.whl) +if ((${#wheels[@]} == 0)); then + wheels=("${MLIR_PY_PKG}/dist/ptoas-"*.whl) +fi shopt -u nullglob -((${#wheels[@]} > 0)) || { echo "error: no ptoas-*.whl under ${MLIR_PY_PKG}/dist" >&2; exit 1; } +((${#wheels[@]} > 0)) || { echo "error: no ptoas-*.whl under build/wheel-dist or ${MLIR_PY_PKG}/dist" >&2; exit 1; } pip install --force-reinstall "${wheels[0]}" export PATH="${PTO_SOURCE_DIR}/build/tools/ptoas:${PATH}" diff --git a/test/dsl/issue_518_repro.pto b/test/dsl/issue_518_repro.pto new file mode 100644 index 0000000000..17e4d20e8b --- /dev/null +++ b/test/dsl/issue_518_repro.pto @@ -0,0 +1,72 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +module attributes {pto.target_arch = "a5"} { + module attributes {pto.backend = "vpto", pto.kernel_kind = #pto.kernel_kind, pto.target_arch = "a5"} { + func.func @group_count_dhist_vmi(%arg0: !pto.ptr, %arg1: !pto.ptr) attributes {pto.entry} { + %c0_i64 = arith.constant 0 : i64 + %0 = builtin.unrealized_conversion_cast %c0_i64 : i64 to ui64 + %1 = pto.castptr %0 : ui64 -> !pto.ptr + %c0_i64_0 = arith.constant 0 : i64 + %2 = builtin.unrealized_conversion_cast %c0_i64_0 : i64 to ui64 + %3 = pto.castptr %2 : ui64 -> !pto.ptr + %c96256_i64 = arith.constant 96256 : i64 + %4 = builtin.unrealized_conversion_cast %c96256_i64 : i64 to ui64 + %5 = pto.castptr %4 : ui64 -> !pto.ptr + %c1_i64 = arith.constant 1 : i64 + %c192512_i64 = arith.constant 192512 : i64 + %c192512_i64_1 = arith.constant 192512 : i64 + %c0_i64_2 = arith.constant 0 : i64 + %c192512_i64_3 = arith.constant 192512 : i64 + pto.mte_gm_ub %arg0, %1, %c0_i64_2, %c192512_i64_3 nburst(%c1_i64, %c192512_i64, %c192512_i64_1) {operandSegmentSizes = array} : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + pto.set_flag[, , ] + pto.wait_flag[, , ] + pto.vecscope { + %c256 = arith.constant 256 : index + %6 = pto.vmi.create_mask %c256 : index -> !pto.vmi.mask<256xpred> + %c256_7 = arith.constant 256 : index + %7 = pto.vmi.create_mask %c256_7 : index -> !pto.vmi.mask<256xpred> + %c0_i16 = arith.constant 0 : i16 + %8 = builtin.unrealized_conversion_cast %c0_i16 : i16 to ui16 + %9 = pto.vmi.vbrc %8 : ui16 -> !pto.vmi.vreg<256xui16> + %c0_i32 = arith.constant 0 : i32 + %10 = builtin.unrealized_conversion_cast %c0_i32 : i32 to ui32 + %11 = pto.vmi.vbrc %10 : ui32 -> !pto.vmi.vreg<256xui32> + %c0 = arith.constant 0 : index + %c94 = arith.constant 94 : index + %c1 = arith.constant 1 : index + %12 = scf.for %arg2 = %c0 to %c94 step %c1 iter_args(%arg3 = %11) -> (!pto.vmi.vreg<256xui32>) { + %c256_9 = arith.constant 256 : index + %13 = arith.muli %arg2, %c256_9 : index + %c2 = arith.constant 2 : index + %14 = arith.muli %13, %c2 : index + %15:2 = pto.vmi.vload %3[%14] {dist_mode = "dintlv"} : !pto.ptr -> !pto.vmi.vreg<256xui32>, !pto.vmi.vreg<256xui32> + %c9_i32 = arith.constant 9 : i32 + %16 = builtin.unrealized_conversion_cast %c9_i32 : i32 to ui32 + %17 = pto.vmi.vcmps %15#0, %16, %6 {cmp = "lt"} : !pto.vmi.vreg<256xui32>, ui32, !pto.vmi.mask<256xpred> -> !pto.vmi.mask<256xpred> + %18 = pto.vmi.vcvt %15#0 : !pto.vmi.vreg<256xui32> -> !pto.vmi.vreg<256xui8> + %19 = pto.vmi.vdhist %9, %18, %17 : !pto.vmi.vreg<256xui16>, !pto.vmi.vreg<256xui8>, !pto.vmi.mask<256xpred> -> !pto.vmi.vreg<256xui16> + %20 = pto.vmi.vcvt %19 : !pto.vmi.vreg<256xui16> -> !pto.vmi.vreg<256xui32> + %21 = pto.vmi.vadd %arg3, %20, %7 : !pto.vmi.vreg<256xui32>, !pto.vmi.vreg<256xui32>, !pto.vmi.mask<256xpred> -> !pto.vmi.vreg<256xui32> + scf.yield %21 : !pto.vmi.vreg<256xui32> + } + %c0_8 = arith.constant 0 : index + pto.vmi.vstore %12, %5[%c0_8], %7 : !pto.vmi.vreg<256xui32>, !pto.ptr, !pto.vmi.mask<256xpred> + } + pto.set_flag[, , ] + pto.wait_flag[, , ] + %c1_i64_4 = arith.constant 1 : i64 + %c1024_i64 = arith.constant 1024 : i64 + %c1024_i64_5 = arith.constant 1024 : i64 + %c1024_i64_6 = arith.constant 1024 : i64 + pto.mte_ub_gm %5, %arg1, %c1024_i64_6 nburst(%c1_i64_4, %c1024_i64, %c1024_i64_5) {operandSegmentSizes = array} : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + pto.mem_bar "SS_ALL" + return + } + } +} diff --git a/test/kernel-test/.gitignore b/test/kernel-test/.gitignore new file mode 100644 index 0000000000..4c9addafd9 --- /dev/null +++ b/test/kernel-test/.gitignore @@ -0,0 +1,11 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +sim_outputs +**/__pycache__/ +kernels/**/generated/ diff --git a/test/kernel-test/README.md b/test/kernel-test/README.md new file mode 100644 index 0000000000..3b13a7277c --- /dev/null +++ b/test/kernel-test/README.md @@ -0,0 +1,138 @@ + + +# kernel-test + +Shared test framework skeleton for multi-kernel, multi-backend validation. + +Current `rope` adapter validation assumes the shell is already running inside the +`pto` environment, because the copied rope implementation depends on packages that +are not available in the base environment. + +## Layout + +- `run.py`: unified CLI entry point +- `kernel_test/`: shared framework package +- `kernels/`: default built-in kernel root +- external kernel roots: `run.py` can also discover a user-specified directory root such as `docs/` + or one concrete kernel package directory such as `docs/rope` +- one directory per kernel, each discovered by the registry + and free to keep local `cce/`, `vmi/`, and `mi/` backend directories +- `scripts/`: shell transport wrappers for cannsim and other external runners + - see `scripts/README.md` for the entrypoint/helper split + +## Usage + +List registered kernels: + +```bash +python kernel-test/run.py --list-ops +``` + +List cases for one kernel after its adapter is added: + +```bash +python kernel-test/run.py --op rope --list-cases +``` + +List and run kernels from a user-specified external directory root: + +```bash +python kernel-test/run.py --kernel-dir docs --list-ops +python kernel-test/run.py --kernel-dir docs --op rope --list-cases +``` + +You can also point directly at one kernel package directory: + +```bash +python kernel-test/run.py --kernel-dir docs/rope --op rope --list-cases +``` + +Run direct Python correctness: + +```bash +python kernel-test/run.py --op rope --workflow correctness --backend cce +``` + +Generate PTO artifacts without launching runtime: + +```bash +python kernel-test/run.py --op dequant --backend vmi --case e4m3_f32 --emit-mlir +``` + +Run one cycle case directly through Python: + +```bash +python kernel-test/run.py --op rope --workflow cycle --backend cce --case f16_half +``` + +Run one cycle case through the generic cannsim transport: + +```bash +kernel-test/scripts/run_sim.sh \ + --output kernel-test/sim_outputs/manual/rope-cce-f16-half \ + kernel-test/run.py \ + -- \ + --op rope --workflow cycle --backend cce --case f16_half +``` + +Run multiple cycle cases, optionally in parallel, and store outputs under +`kernel-test/sim_outputs////`. The default cycle engine is +`msprof`, with `cannsim` still available as a fallback: + +```bash +kernel-test/scripts/run_cycle.sh --op rope --backend cce --parallel-sim 1 --jobs 4 +``` + +`run_cycle.sh` accepts the same external kernel discovery override: + +```bash +kernel-test/scripts/run_cycle.sh --kernel-dir docs --op rope --backend cce +``` + +If you need a non-default interpreter for cannsim transport, `run_sim.sh` still +accepts `--python-cmd`. + +If you prefer environment variables, both `run.py` and `run_cycle.sh` also honor +`KERNEL_TEST_KERNEL_DIR`. + +`--emit-mlir` uses correctness case selection, asks the selected backend to +materialize PTO outputs, and writes them under the kernel-local +`test/kernel-test/kernels//generated/` directory. Backends that do not +implement PTO artifact emission are reported as `SKIP`. + +`run_cycle.sh` now runs a kernel-local cycle analysis step after successful jobs. +For `rope`, the primary VF cycle number now comes from `msprof` `RVECEX` pipe +cycles when available; `trace`, `instr_log`, and coarse SoC cycles are reported +only as supporting references. + +## Status + +This directory currently contains the framework foundation from tasks 1 through 3: + +- package structure +- registry-based kernel discovery +- unified CLI entry point +- shared runtime types +- case selection helpers +- correctness and cycle runners +- one self-contained rope kernel directory with `cce/`, `vmi/`, and `mi/` backends +- shared rope runtime that prepares launch arguments before backend execution + +Simulator transport scripts now live under `scripts/`. `run_sim.sh` owns one +single-case cannsim run, while `run_cycle.sh` expands a kernel's selected case +set into one sim job per case. + +The first concrete adapter lives under `kernels/rope/`. Future kernels should follow +the same per-kernel directory pattern so backend-specific files can stay local to the +kernel they belong to, whether they live under the built-in `kernels/` root or an +external root such as `docs/`. Rope now keeps its copied CCE source directly under +`cce/`, its copied backend PTO files directly under `vmi/` and `mi/`, and its CCE +build configuration in `cce/CMakeLists.txt`. diff --git a/test/kernel-test/kernel_test/__init__.py b/test/kernel-test/kernel_test/__init__.py new file mode 100644 index 0000000000..86b2f6588a --- /dev/null +++ b/test/kernel-test/kernel_test/__init__.py @@ -0,0 +1,26 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""Shared Python package for the kernel-test framework.""" + +from .backends import BackendAdapter +from .registry import KernelRegistry, OperatorSpec, RegistryError, load_registry +from .results import CaseResult, RunSummary + +__all__ = [ + "BackendAdapter", + "CaseResult", + "KernelRegistry", + "OperatorSpec", + "RegistryError", + "RunSummary", + "__version__", + "load_registry", +] + +__version__ = "0.1.0" diff --git a/test/kernel-test/kernel_test/backends.py b/test/kernel-test/kernel_test/backends.py new file mode 100644 index 0000000000..4d9d599ba7 --- /dev/null +++ b/test/kernel-test/kernel_test/backends.py @@ -0,0 +1,57 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""Backend interfaces for the kernel-test framework.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Literal, Protocol + +RunPurpose = Literal["correctness", "cycle"] + + +@dataclass(frozen=True) +class ArtifactOutputs: + """Normalized artifact emission result for one kernel-test case.""" + + case_dir: str + message: str + paths: Mapping[str, str] + + +@dataclass(frozen=True) +class ArtifactPlan: + """Backend-provided compile result that the framework can materialize.""" + + generated_dir: Path + case_dir: Path + vmi_text: str | None = None + mi_text: str | None = None + alias_stem: str | None = None + + +class BackendAdapter(Protocol): + """Stable interface shared by all framework backends.""" + + name: str + + def is_supported(self, case: object, *, purpose: RunPurpose) -> tuple[bool, str | None]: + """Return support status and an optional human-readable reason.""" + + def launch(self, case: object, *, purpose: RunPurpose) -> object: + """Launch one case and return backend-specific outputs.""" + + +class ArtifactBackend(Protocol): + """Optional extension for backends that can emit PTO artifacts.""" + + def build_artifact_plan(self, case_id: str, case: object) -> ArtifactPlan: + """Build the backend-specific PTO artifact plan for one case.""" diff --git a/test/kernel-test/kernel_test/cases.py b/test/kernel-test/kernel_test/cases.py new file mode 100644 index 0000000000..44eafb4f67 --- /dev/null +++ b/test/kernel-test/kernel_test/cases.py @@ -0,0 +1,58 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""Case selection helpers for framework runners.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import TypeVar + +CaseT = TypeVar("CaseT") + + +def select_cases( + cases: Mapping[str, CaseT], + *, + case_ids: Sequence[str], + case_filter: str | None, + require_single: bool = False, +) -> dict[str, CaseT]: + """Resolve exact-case and substring filters into a concrete selection.""" + + if case_ids: + selected: dict[str, CaseT] = {} + missing: list[str] = [] + for case_id in case_ids: + if case_id not in cases: + missing.append(case_id) + continue + selected[case_id] = cases[case_id] + else: + selected = {case_id: cases[case_id] for case_id in sorted(cases.keys())} + missing = [] + + if missing: + raise ValueError(f"unknown case ids: {', '.join(missing)}") + + if case_filter: + selected = { + case_id: case + for case_id, case in selected.items() + if case_filter in case_id + } + + if not selected: + raise ValueError("no cases matched the requested selection") + + if require_single and len(selected) != 1: + raise ValueError( + f"cycle workflow requires exactly one case, but resolved {len(selected)} cases" + ) + + return selected diff --git a/test/kernel-test/kernel_test/cli.py b/test/kernel-test/kernel_test/cli.py new file mode 100644 index 0000000000..d4fdbbfdc3 --- /dev/null +++ b/test/kernel-test/kernel_test/cli.py @@ -0,0 +1,139 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""Command-line interface for the kernel-test framework.""" + +from __future__ import annotations + +import argparse +import os +from typing import Sequence + +from .cases import select_cases +from .registry import RegistryError, load_registry +from .runners import run_artifact_suite, run_correctness_suite, run_cycle_probe + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Unified kernel-test CLI") + parser.add_argument( + "--kernel-dir", + help=( + "Kernel package root to discover. Accepts either a directory that contains " + "kernel subdirectories such as docs/ or one kernel package directory such as docs/rope. " + "Defaults to test/kernel-test/kernels or $KERNEL_TEST_KERNEL_DIR." + ), + ) + parser.add_argument("--list-ops", action="store_true", help="List registered kernels") + parser.add_argument("--op", help="Kernel name to run") + parser.add_argument( + "--workflow", + choices=("correctness", "cycle"), + default="correctness", + help="Workflow name", + ) + parser.add_argument("--backend", help="Backend name") + parser.add_argument("--case", action="append", default=[], help="Case id to run") + parser.add_argument("--case-filter", help="Substring filter for cases") + parser.add_argument("--list-cases", action="store_true", help="List cases for one kernel") + parser.add_argument( + "--emit-mlir", + action="store_true", + help="Generate PTO artifacts such as vmi.pto/mi.pto under kernels/**/generated", + ) + return parser + + +def _print_lines(lines: Sequence[str]) -> None: + for line in lines: + print(line) + + +def main(argv: Sequence[str] | None = None) -> int: + parser = build_parser() + args = parser.parse_args(argv) + kernel_dir = args.kernel_dir or os.environ.get("KERNEL_TEST_KERNEL_DIR") + + try: + registry = load_registry(kernel_dir=kernel_dir) + except RegistryError as exc: + raise SystemExit(str(exc)) from exc + + if args.list_ops: + names = registry.list_names() + if not names: + print("No kernels registered yet.") + return 0 + _print_lines(names) + return 0 + + if args.list_cases: + if not args.op: + parser.error("--list-cases requires --op") + spec = registry.get(args.op) + if spec is None: + parser.error(f"unknown kernel: {args.op}") + case_ids = sorted(spec.list_cases(args.workflow).keys()) + if not case_ids: + print(f"No cases registered for kernel {spec.name} workflow={args.workflow}.") + return 0 + _print_lines(case_ids) + return 0 + + if not args.op: + parser.error("one of --list-ops or --op is required") + + if args.emit_mlir and args.workflow != "correctness": + parser.error("--emit-mlir currently requires --workflow correctness") + + spec = registry.get(args.op) + if spec is None: + parser.error(f"unknown kernel: {args.op}") + + backend_name = args.backend or spec.default_backend + if spec.backend_names and backend_name not in spec.backend_names: + parser.error( + f"unsupported backend {backend_name!r} for kernel {spec.name}; " + f"choices: {', '.join(spec.backend_names)}" + ) + + try: + cases = select_cases( + spec.list_cases(args.workflow), + case_ids=args.case, + case_filter=args.case_filter, + require_single=args.workflow == "cycle", + ) + except ValueError as exc: + parser.error(str(exc)) + + backend = spec.create_backend(backend_name) + + if args.emit_mlir: + summary = run_artifact_suite(cases, backend=backend) + print( + f"SUMMARY total={summary.total} passed={summary.passed} " + f"failed={summary.failed} skipped={summary.skipped}" + ) + return 0 if summary.all_passed else 1 + + if args.workflow == "correctness": + summary = run_correctness_suite(cases, backend=backend, verify_case=spec.verify) + print( + f"SUMMARY total={summary.total} passed={summary.passed} " + f"failed={summary.failed} skipped={summary.skipped}" + ) + return 0 if summary.all_passed else 1 + + case_id, case = next(iter(cases.items())) + marker_fields = { + "op": spec.name, + "backend": backend.name, + **dict(spec.cycle_fields(case_id, case, backend)), + } + return run_cycle_probe(case_id=case_id, case=case, backend=backend, marker_fields=marker_fields) diff --git a/test/kernel-test/kernel_test/cycle_metrics.py b/test/kernel-test/kernel_test/cycle_metrics.py new file mode 100644 index 0000000000..e2dbd375cb --- /dev/null +++ b/test/kernel-test/kernel_test/cycle_metrics.py @@ -0,0 +1,571 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""Shared cannsim cycle-metric parsing helpers for kernel-test.""" + +from __future__ import annotations + +import glob +import csv +import json +import os +import re +from collections import Counter +from dataclasses import dataclass, field +from typing import Iterable + +INSTR_LOG_RE = re.compile( + r"start:\s*(\d+),\s*tick:\s*(\d+).*?blkDim:\s*(\d+)", + re.DOTALL, +) +SOC_CYCLE_RE = re.compile( + r"\[Hardware\]\s+parallel simulation finish\.\s+sim time:\s*" + r"(?:SoC sub \d+ )?([\d.]+)s,\s*cycle:\s*(\d+)", +) + + +@dataclass(frozen=True) +class LaunchRecord: + start: int + tick: int + blk_dim: int + + @property + def duration(self) -> int: + return self.tick - self.start + + +@dataclass(frozen=True) +class InstrLogMetrics: + launches: tuple[LaunchRecord, ...] + span: int + max_dur: int + max_blk_dim: int + + @classmethod + def empty(cls) -> InstrLogMetrics: + return cls(launches=(), span=0, max_dur=0, max_blk_dim=0) + + +@dataclass(frozen=True) +class TraceMetrics: + core_vf_span: int | None = None + rvec_span: int | None = None + pushq_vf_dur: int | None = None + mte2_span: int | None = None + mte3_span: int | None = None + vector_span: int | None = None + arith_sum_dur: int | None = None + rvec_op_counts: dict[str, int] = field(default_factory=dict) + rvec_event_count: int = 0 + trace_path: str | None = None + + @property + def vf_cycles(self) -> int | None: + if self.core_vf_span is not None and self.core_vf_span > 0: + return self.core_vf_span + if self.rvec_span is not None and self.rvec_span > 0: + return self.rvec_span + if self.pushq_vf_dur is not None and self.pushq_vf_dur > 0: + return self.pushq_vf_dur + return None + + +@dataclass(frozen=True) +class MsprofMetrics: + opprof_dir: str | None = None + instr_csv_path: str | None = None + core_vf_cycles: int | None = None + arith_cycles: int | None = None + pipe_cycles: dict[str, int] = field(default_factory=dict) + instr_cycles: dict[str, int] = field(default_factory=dict) + + @classmethod + def empty(cls) -> MsprofMetrics: + return cls() + + +@dataclass(frozen=True) +class SocCycleRecord: + sim_wall_s: float + soc_cycles: int + + +@dataclass(frozen=True) +class RunMetrics: + out_dir: str + cannsim_run_dir: str | None + msprof: MsprofMetrics + instr: InstrLogMetrics + trace: TraceMetrics + soc_cycles: tuple[SocCycleRecord, ...] + steady_soc_cycles: int | None + measured_kernel_cycles: int | None + + @property + def primary_vf_cycles(self) -> int | None: + if self.msprof.core_vf_cycles is not None and self.msprof.core_vf_cycles > 0: + return self.msprof.core_vf_cycles + if self.trace.vf_cycles is not None: + return self.trace.vf_cycles + if self.instr.max_dur > 0: + return self.instr.max_dur + return self.measured_kernel_cycles + + +def find_cannsim_run_dir(out_dir: str) -> str: + pattern = os.path.join(out_dir, "cannsim_*") + candidates = [path for path in glob.glob(pattern) if os.path.isdir(path)] + if not candidates: + raise FileNotFoundError(f"No cannsim_* directory under {out_dir}") + return max(candidates, key=os.path.getmtime) + + +def maybe_find_cannsim_run_dir(out_dir: str) -> str | None: + pattern = os.path.join(out_dir, "cannsim_*") + candidates = [path for path in glob.glob(pattern) if os.path.isdir(path)] + if not candidates: + return None + return max(candidates, key=os.path.getmtime) + + +def _parse_instr_log_file(path: str) -> list[LaunchRecord]: + with open(path, encoding="utf-8", errors="replace") as fh: + text = fh.read() + records: list[LaunchRecord] = [] + for match in INSTR_LOG_RE.finditer(text): + start, tick, blk_dim = (int(match.group(i)) for i in range(1, 4)) + records.append(LaunchRecord(start=start, tick=tick, blk_dim=blk_dim)) + return records + + +def _group_launches( + launches: list[LaunchRecord], + gap_threshold: int = 3000, +) -> list[list[LaunchRecord]]: + if not launches: + return [] + sorted_launches = sorted(launches, key=lambda record: record.start) + groups: list[list[LaunchRecord]] = [[sorted_launches[0]]] + for record in sorted_launches[1:]: + prev_end = max(item.tick for item in groups[-1]) + if record.start - prev_end > gap_threshold: + groups.append([record]) + else: + groups[-1].append(record) + return groups + + +def metrics_from_launches(launches: Iterable[LaunchRecord]) -> InstrLogMetrics: + items = list(launches) + if not items: + return InstrLogMetrics.empty() + starts = [record.start for record in items] + ticks = [record.tick for record in items] + durations = [record.duration for record in items] + blk_dims = [record.blk_dim for record in items] + return InstrLogMetrics( + launches=tuple(items), + span=max(ticks) - min(starts), + max_dur=max(durations), + max_blk_dim=max(blk_dims), + ) + + +def parse_instr_log_dir(log_ca_dir: str, measured_only: bool = True) -> InstrLogMetrics: + pattern = os.path.join(log_ca_dir, "core*.veccore*.instr_log.dump") + paths = sorted(glob.glob(pattern)) + if not paths: + return InstrLogMetrics.empty() + + all_launches: list[LaunchRecord] = [] + for path in paths: + all_launches.extend(_parse_instr_log_file(path)) + + if not all_launches: + return InstrLogMetrics.empty() + + if measured_only: + max_blk = max(record.blk_dim for record in all_launches) + kernel_launches = [record for record in all_launches if record.blk_dim == max_blk] + if kernel_launches: + groups = _group_launches(kernel_launches, gap_threshold=2000) + all_launches = groups[-1] if groups else kernel_launches + + return metrics_from_launches(all_launches) + + +def parse_soc_cycles(cannsim_log: str) -> list[SocCycleRecord]: + with open(cannsim_log, encoding="utf-8", errors="replace") as fh: + text = fh.read() + + records: list[SocCycleRecord] = [] + for match in SOC_CYCLE_RE.finditer(text): + sim_wall_s = float(match.group(1)) + soc_cycles = int(match.group(2)) + records.append(SocCycleRecord(sim_wall_s=sim_wall_s, soc_cycles=soc_cycles)) + return records + + +def steady_state_soc_cycles(records: Iterable[SocCycleRecord]) -> int | None: + items = list(records) + if not items: + return None + if len(items) == 1: + return items[0].soc_cycles + return items[-1].soc_cycles + + +def find_trace_json(run_dir: str) -> str | None: + patterns = [ + os.path.join(run_dir, "report", "trace_core0.json"), + os.path.join(run_dir, "**", "trace_core0.json"), + ] + for pattern in patterns: + matches = glob.glob(pattern, recursive=True) + if matches: + return max(matches, key=os.path.getmtime) + return None + + +def find_msprof_run_dir(out_dir: str) -> str | None: + patterns = [ + os.path.join(out_dir, "msprof", "OPPROF_*"), + os.path.join(out_dir, "OPPROF_*"), + ] + candidates: list[str] = [] + for pattern in patterns: + candidates.extend(path for path in glob.glob(pattern) if os.path.isdir(path)) + if not candidates: + return None + return max(candidates, key=os.path.getmtime) + + +def find_msprof_instr_csv(opprof_dir: str) -> str | None: + patterns = [ + os.path.join(opprof_dir, "simulator", "core0.veccore0", "core0.veccore0_instr_exe.csv"), + os.path.join(opprof_dir, "**", "core*.veccore*.instr_exe.csv"), + os.path.join(opprof_dir, "**", "*instr_exe*.csv"), + ] + candidates: list[str] = [] + for pattern in patterns: + candidates.extend(path for path in glob.glob(pattern, recursive=True) if os.path.isfile(path)) + if not candidates: + return None + return max(candidates, key=os.path.getmtime) + + +def parse_msprof_metrics(out_dir: str) -> MsprofMetrics: + opprof_dir = find_msprof_run_dir(out_dir) + if not opprof_dir: + return MsprofMetrics.empty() + + instr_csv_path = find_msprof_instr_csv(opprof_dir) + if not instr_csv_path: + return MsprofMetrics(opprof_dir=opprof_dir) + + pipe_cycles: dict[str, int] = {} + instr_cycles: dict[str, int] = {} + with open(instr_csv_path, encoding="utf-8", newline="") as fh: + reader = csv.DictReader(fh) + for row in reader: + normalized = { + str(key).strip().lower(): (value or "").strip() + for key, value in row.items() + if key is not None + } + instr = ( + normalized.get("instr") + or normalized.get("instruction") + or normalized.get("instr_name") + or "" + ) + pipe = normalized.get("pipe") or normalized.get("pipeline") or "" + cycle_value = ( + normalized.get("cycles") + or normalized.get("cycle") + or normalized.get("duration") + or "0" + ) + try: + cycles = int(float(cycle_value)) + except ValueError: + continue + if pipe: + pipe_cycles[pipe] = pipe_cycles.get(pipe, 0) + cycles + if instr: + instr_cycles[instr] = instr_cycles.get(instr, 0) + cycles + + arith_instrs = { + "RV_VMUL", + "RV_VADD", + "RV_VSUB", + "RV_VDIV", + "RV_VMAX", + "RV_VMIN", + "RV_VMAC", + "RV_VMADD", + "RV_VMLS", + } + arith_cycles = sum(cycles for instr, cycles in instr_cycles.items() if instr in arith_instrs) or None + core_vf_cycles = pipe_cycles.get("RVECEX") or None + + return MsprofMetrics( + opprof_dir=opprof_dir, + instr_csv_path=instr_csv_path, + core_vf_cycles=core_vf_cycles, + arith_cycles=arith_cycles, + pipe_cycles=pipe_cycles, + instr_cycles=instr_cycles, + ) + + +def _load_trace_events(trace_path: str) -> list[dict]: + with open(trace_path, encoding="utf-8") as fh: + payload = json.load(fh) + if isinstance(payload, list): + return payload + return payload.get("traceEvents", []) + + +def _pipe_span(events: list[dict], processes: dict[int, str], *needles: str) -> int | None: + pids = {pid for pid, name in processes.items() if any(needle in name.upper() for needle in needles)} + exec_ev = [event for event in events if event.get("ph") == "X" and "dur" in event and event.get("pid") in pids] + if not exec_ev: + return None + start = min(event["ts"] for event in exec_ev) + end = max(event["ts"] + event["dur"] for event in exec_ev) + span = int(end - start) + return span if span > 0 else None + + +def parse_trace_metrics(trace_path: str) -> TraceMetrics: + if not os.path.isfile(trace_path): + return TraceMetrics() + + events = _load_trace_events(trace_path) + processes = { + event["pid"]: event["args"]["name"] + for event in events + if event.get("ph") == "M" and event.get("name") == "process_name" + } + exec_ev = [event for event in events if event.get("ph") == "X" and "dur" in event] + + pushq_pids = {pid for pid, name in processes.items() if "PUSHQ" in name.upper()} + vf_dispatch = [ + event for event in exec_ev if event.get("pid") in pushq_pids and "VF" in event.get("name", "") + ] + pushq_vf_dur = max((event.get("dur", 0) for event in vf_dispatch), default=0) or None + + rvec_pids = {pid for pid, name in processes.items() if "RVEC" in name.upper()} + rvec_events = [event for event in exec_ev if event.get("pid") in rvec_pids] + rvecex_pids = {pid for pid, name in processes.items() if "RVECEX" in name.upper()} + rvecex_events = [event for event in exec_ev if event.get("pid") in rvecex_pids] + arith_events = [ + event + for event in rvecex_events + if event.get("name") in {"RV_VMUL", "RV_VADD", "RV_VSUB", "RV_VDIV", "RV_VMAX", "RV_VMIN"} + ] + core_vf_span = None + arith_sum_dur = None + if rvecex_events: + start = min(event["ts"] for event in rvecex_events) + end = max(event["ts"] + event["dur"] for event in rvecex_events) + core_vf_span = int(end - start) + if arith_events: + arith_sum_dur = int(sum(event.get("dur", 0) for event in arith_events)) + + rvec_span = None + rvec_op_counts: dict[str, int] = {} + rvec_event_count = 0 + if rvec_events: + start = min(event["ts"] for event in rvec_events) + end = max(event["ts"] + event["dur"] for event in rvec_events) + rvec_span = int(end - start) + rvec_op_counts = dict(Counter(event.get("name", "?") for event in rvec_events)) + rvec_event_count = len(rvec_events) + + return TraceMetrics( + core_vf_span=core_vf_span, + rvec_span=rvec_span, + pushq_vf_dur=pushq_vf_dur, + mte2_span=_pipe_span(events, processes, "MTE2"), + mte3_span=_pipe_span(events, processes, "MTE3"), + vector_span=_pipe_span(events, processes, "RVEC", "VECTOR", "VEC"), + arith_sum_dur=arith_sum_dur, + rvec_op_counts=rvec_op_counts, + rvec_event_count=rvec_event_count, + trace_path=trace_path, + ) + + +def parse_marker_soc_cycles(cannsim_log: str) -> tuple[list[SocCycleRecord], list[SocCycleRecord]]: + with open(cannsim_log, encoding="utf-8", errors="replace") as fh: + lines = fh.readlines() + + in_window = False + all_records: list[SocCycleRecord] = [] + window_records: list[SocCycleRecord] = [] + + for line in lines: + if "CYCLE_MARKER" in line: + in_window = True + window_records.clear() + continue + if "CYCLE_DONE" in line: + in_window = False + continue + match = SOC_CYCLE_RE.search(line) + if not match: + continue + record = SocCycleRecord(sim_wall_s=float(match.group(1)), soc_cycles=int(match.group(2))) + all_records.append(record) + if in_window: + window_records.append(record) + + return all_records, window_records + + +def measured_kernel_soc_cycles(cannsim_log: str) -> int | None: + _, window = parse_marker_soc_cycles(cannsim_log) + if not window: + return None + return window[-1].soc_cycles + + +def parse_run_metrics(out_dir: str) -> RunMetrics: + msprof = parse_msprof_metrics(out_dir) + run_dir = maybe_find_cannsim_run_dir(out_dir) + + if run_dir is None and msprof.opprof_dir is None: + raise FileNotFoundError( + f"No cycle artifacts found under {out_dir}; expected msprof/OPPROF_* or cannsim_* outputs" + ) + + log_ca = os.path.join(run_dir, "log_ca") if run_dir else "" + cannsim_log = os.path.join(run_dir, "cannsim.log") if run_dir else "" + + instr = parse_instr_log_dir(log_ca) if log_ca and os.path.isdir(log_ca) else InstrLogMetrics.empty() + soc = parse_soc_cycles(cannsim_log) if cannsim_log and os.path.isfile(cannsim_log) else [] + measured = ( + measured_kernel_soc_cycles(cannsim_log) + if cannsim_log and os.path.isfile(cannsim_log) + else None + ) + + trace_path = find_trace_json(run_dir) if run_dir else None + trace = parse_trace_metrics(trace_path) if trace_path else TraceMetrics() + + return RunMetrics( + out_dir=out_dir, + cannsim_run_dir=run_dir, + msprof=msprof, + instr=instr, + trace=trace, + soc_cycles=tuple(soc), + steady_soc_cycles=steady_state_soc_cycles(soc), + measured_kernel_cycles=measured, + ) + + +def format_run_summary(metrics: RunMetrics, label: str | None = None) -> str: + title = label or os.path.basename(metrics.out_dir.rstrip("/")) + lines = [f"=== {title} ({metrics.out_dir}) ==="] + if metrics.msprof.opprof_dir: + lines.append(f"msprof: {metrics.msprof.opprof_dir}") + if metrics.cannsim_run_dir: + lines.append(f"cannsim run dir: {metrics.cannsim_run_dir}") + + primary = metrics.primary_vf_cycles + if primary is not None: + if metrics.msprof.core_vf_cycles and metrics.msprof.core_vf_cycles > 0: + source = "msprof RVECEX cycles" + elif metrics.trace.core_vf_span and metrics.trace.core_vf_span > 0: + source = "RVECEX span" + elif metrics.trace.rvec_span and metrics.trace.rvec_span > 0: + source = "RVEC span" + elif metrics.trace.pushq_vf_dur and metrics.trace.pushq_vf_dur > 0: + source = "PUSHQ VF dur" + elif metrics.instr.max_dur > 0: + source = "instr MaxDur" + else: + source = "SoC (fallback)" + lines.append(f"primary VF cycles: {primary} ({source})") + + if metrics.msprof.core_vf_cycles: + lines.append(f"msprof core VF cycles: {metrics.msprof.core_vf_cycles}") + if metrics.msprof.arith_cycles: + lines.append(f"msprof arith cycles: {metrics.msprof.arith_cycles}") + if metrics.msprof.pipe_cycles: + ordered_pipes = ["RVECEX", "RVECLD", "RVECST", "RVECSU", "MTE2", "MTE3", "VECTOR"] + rendered = ", ".join( + f"{pipe}={metrics.msprof.pipe_cycles[pipe]}" + for pipe in ordered_pipes + if pipe in metrics.msprof.pipe_cycles + ) + if rendered: + lines.append(f"msprof pipe cycles: {rendered}") + if metrics.trace.core_vf_span: + lines.append(f"core VF span (RVECEX): {metrics.trace.core_vf_span}") + if metrics.trace.rvec_span: + lines.append(f"RVEC span: {metrics.trace.rvec_span}") + if metrics.trace.pushq_vf_dur: + lines.append(f"PUSHQ VF dur: {metrics.trace.pushq_vf_dur}") + if metrics.trace.arith_sum_dur: + lines.append(f"arith sum dur: {metrics.trace.arith_sum_dur}") + if metrics.trace.mte2_span: + lines.append(f"MTE2 span: {metrics.trace.mte2_span}") + if metrics.trace.mte3_span: + lines.append(f"MTE3 span: {metrics.trace.mte3_span}") + if metrics.trace.rvec_op_counts: + top_ops = sorted(metrics.trace.rvec_op_counts.items(), key=lambda item: -item[1])[:6] + lines.append("top RVEC ops: " + ", ".join(f"{name}={count}" for name, count in top_ops)) + + if metrics.instr.max_dur: + lines.append( + f"instr MaxDur / span: {metrics.instr.max_dur} / {metrics.instr.span} " + f"(blkDim={metrics.instr.max_blk_dim})" + ) + elif not metrics.cannsim_run_dir or not os.path.isdir(os.path.join(metrics.cannsim_run_dir, "log_ca")): + lines.append("instr MaxDur / span: (log_ca unavailable)") + + if metrics.measured_kernel_cycles is not None: + lines.append( + f"SoC cycles (measured): {metrics.measured_kernel_cycles} " + "(coarse; ~420 is normal)" + ) + elif metrics.steady_soc_cycles is not None: + lines.append(f"SoC cycles (steady): {metrics.steady_soc_cycles} (coarse)") + + if metrics.trace.trace_path: + lines.append(f"trace: {metrics.trace.trace_path}") + elif metrics.cannsim_run_dir and not find_trace_json(metrics.cannsim_run_dir): + lines.append("trace: (none — need instr.bin + cannsim report)") + + return "\n".join(lines) + + +def format_table(rows: list[tuple[str, RunMetrics]]) -> str: + header = ( + f"{'case':<22} {'primary':>8} {'RVECEX':>8} {'RVEC':>8} " + f"{'MaxDur':>8} {'Span':>8} {'SoC':>6}" + ) + lines = [header, "-" * len(header)] + for label, metrics in rows: + primary = metrics.primary_vf_cycles + rvecex = metrics.msprof.core_vf_cycles or metrics.trace.core_vf_span + lines.append( + f"{label:<22} " + f"{primary if primary is not None else '-':>8} " + f"{rvecex if rvecex else '-':>8} " + f"{metrics.trace.rvec_span if metrics.trace.rvec_span else '-':>8} " + f"{metrics.instr.max_dur if metrics.instr.max_dur else '-':>8} " + f"{metrics.instr.span if metrics.instr.span else '-':>8} " + f"{metrics.measured_kernel_cycles if metrics.measured_kernel_cycles else '-':>6}" + ) + return "\n".join(lines) diff --git a/test/kernel-test/kernel_test/cycle_reporting.py b/test/kernel-test/kernel_test/cycle_reporting.py new file mode 100644 index 0000000000..b8a64f11c4 --- /dev/null +++ b/test/kernel-test/kernel_test/cycle_reporting.py @@ -0,0 +1,67 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""Framework contract and runner for per-kernel cycle reporting.""" + +from __future__ import annotations + +import argparse +import os +import sys +from collections.abc import Callable +from dataclasses import dataclass + +from .cycle_metrics import RunMetrics, format_run_summary, format_table, parse_run_metrics + + +@dataclass(frozen=True) +class CycleReporterSpec: + """Framework-owned contract for one kernel cycle-report adapter.""" + + name: str + default_out_dirs: Callable[[str | None], list[str]] + missing_message: str + parse_metrics: Callable[[str], RunMetrics] = parse_run_metrics + format_summary: Callable[[RunMetrics, str | None], str] = format_run_summary + format_table: Callable[[list[tuple[str, RunMetrics]]], str] = format_table + label_for_dir: Callable[[str], str] = lambda path: os.path.basename(path.rstrip("/")) + + +def run_cycle_report(spec: CycleReporterSpec, argv: list[str] | None = None) -> int: + """Run one kernel-local cycle report with shared CLI behavior.""" + + parser = argparse.ArgumentParser(description=f"Parse {spec.name} cycle metrics from kernel-test outputs") + parser.add_argument( + "out_dirs", + nargs="*", + help="case output dirs such as sim_outputs///", + ) + parser.add_argument("--table", action="store_true", help="Print compact table") + args = parser.parse_args(argv) + + dirs = args.out_dirs if args.out_dirs else spec.default_out_dirs(None) + if not dirs: + print(spec.missing_message, file=sys.stderr) + return 1 + + rows: list[tuple[str, RunMetrics]] = [] + for path in dirs: + try: + metrics = spec.parse_metrics(path) + except FileNotFoundError as exc: + print(str(exc), file=sys.stderr) + continue + label = spec.label_for_dir(path) + rows.append((label, metrics)) + if not args.table: + print(spec.format_summary(metrics, label)) + print() + + if args.table and rows: + print(spec.format_table(rows)) + return 0 if rows else 1 diff --git a/test/kernel-test/kernel_test/npu_runtime.py b/test/kernel-test/kernel_test/npu_runtime.py new file mode 100644 index 0000000000..7f6f0ec0b2 --- /dev/null +++ b/test/kernel-test/kernel_test/npu_runtime.py @@ -0,0 +1,62 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""Shared torch_npu runtime helpers for kernel-test backends.""" + +from __future__ import annotations + +import os + +_DEVICE = f"npu:{os.environ.get('NPU_DEVICE', '0')}" + + +def device_str() -> str: + return _DEVICE + + +def init_torch_npu(device: str | None = None) -> None: + global _DEVICE + + import torch + import torch_npu + + _DEVICE = device or _DEVICE + torch.npu.config.allow_internal_format = False + torch_npu.npu.set_compile_mode(jit_compile=False) + torch.npu.set_device(_DEVICE) + + +def ensure_runtime(component: str = "kernel-test") -> None: + """Initialize torch_npu and normalize runtime initialization failures.""" + + try: + init_torch_npu() + except Exception as exc: # noqa: BLE001 + raise RuntimeError( + f"failed to initialize torch_npu for {component}; run direct correctness " + "on a host with NPU access or use the simulator transport script for " + "cannsim workflows" + ) from exc + + +def empty_npu(shape, dtype): + import torch + + return torch.empty(shape, dtype=dtype, device=_DEVICE) + + +def stream_ptr() -> int: + import torch + + return torch.npu.current_stream()._as_parameter_ + + +def sync() -> None: + import torch + + torch.npu.synchronize() diff --git a/test/kernel-test/kernel_test/pto_artifacts.py b/test/kernel-test/kernel_test/pto_artifacts.py new file mode 100644 index 0000000000..dcf8e18dd5 --- /dev/null +++ b/test/kernel-test/kernel_test/pto_artifacts.py @@ -0,0 +1,125 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""Helpers for writing kernel-test PTO artifacts and lowering VMI to MI.""" + +from __future__ import annotations + +import os +from pathlib import Path +import shutil +import subprocess + +from .backends import ArtifactOutputs, ArtifactPlan + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_PTOAS_BIN = Path( + os.environ.get( + "PTOAS_BIN", + str(_REPO_ROOT / "build" / "tools" / "ptoas" / "ptoas"), + ) +) + + +def resolve_ptoas_bin() -> Path: + """Return the ptoas binary used for kernel-test artifact lowering.""" + + if _PTOAS_BIN.is_file(): + return _PTOAS_BIN + + fallback = _REPO_ROOT / "install" / "bin" / "ptoas" + if fallback.is_file(): + return fallback + + raise FileNotFoundError(f"ptoas not found: {_PTOAS_BIN}") + + +def lower_vmi_to_mi(vmi_path: Path, mi_path: Path) -> str: + """Lower one VMI PTO artifact to MI PTO text with ptoas.""" + + result = subprocess.run( + [ + str(resolve_ptoas_bin()), + "--pto-arch=a5", + "--pto-backend=vpto", + "--enable-vmi", + "--pto-level=level3", + "--emit-vpto", + "-o", + str(mi_path), + str(vmi_path), + ], + cwd=_REPO_ROOT, + check=False, + text=True, + capture_output=True, + ) + if result.returncode != 0: + raise RuntimeError( + "failed to lower VMI artifact with ptoas:\n" + f"stdout:\n{result.stdout}\n" + f"stderr:\n{result.stderr}" + ) + return mi_path.read_text(encoding="utf-8") + + +def materialize_artifact_plan( + case_id: str, + plan: ArtifactPlan, + *, + root_alias: bool = False, +) -> ArtifactOutputs: + """Write one backend artifact plan into the kernel-local generated tree.""" + + plan.generated_dir.mkdir(parents=True, exist_ok=True) + plan.case_dir.mkdir(parents=True, exist_ok=True) + + paths: dict[str, str] = {"case_dir": str(plan.case_dir)} + written_names: list[str] = [] + + vmi_path = plan.case_dir / "vmi.pto" + mi_path = plan.case_dir / "mi.pto" + if plan.vmi_text is not None: + vmi_path.write_text(plan.vmi_text, encoding="utf-8") + paths["vmi_path"] = str(vmi_path) + written_names.append("vmi.pto") + + if plan.mi_text is not None: + mi_path.write_text(plan.mi_text, encoding="utf-8") + elif plan.vmi_text is not None: + lower_vmi_to_mi(vmi_path, mi_path) + + if plan.mi_text is not None or plan.vmi_text is not None: + if mi_path.is_file(): + paths["mi_path"] = str(mi_path) + if "mi.pto" not in written_names: + written_names.append("mi.pto") + + if plan.alias_stem: + for artifact_name in tuple(written_names): + legacy_case_alias = plan.case_dir / f"{plan.alias_stem}.{artifact_name}" + if legacy_case_alias.exists(): + legacy_case_alias.unlink() + + if root_alias: + for artifact_name in tuple(written_names): + src = plan.case_dir / artifact_name + dst = plan.generated_dir / artifact_name + shutil.copyfile(src, dst) + paths[f"root_{artifact_name.replace('.', '_')}"] = str(dst) + if plan.alias_stem: + legacy_root_alias = plan.generated_dir / f"{plan.alias_stem}.{artifact_name}" + if legacy_root_alias.exists(): + legacy_root_alias.unlink() + + rendered = " and ".join(written_names) if written_names else "no PTO artifacts" + return ArtifactOutputs( + case_dir=str(plan.case_dir), + message=f"wrote {rendered} under {plan.case_dir} for {case_id}", + paths=paths, + ) diff --git a/test/kernel-test/kernel_test/registry.py b/test/kernel-test/kernel_test/registry.py new file mode 100644 index 0000000000..f75e6418a2 --- /dev/null +++ b/test/kernel-test/kernel_test/registry.py @@ -0,0 +1,230 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""Kernel discovery and registry loading for kernel-test.""" + +from __future__ import annotations + +import hashlib +from collections.abc import Callable, Iterable, Mapping +from dataclasses import dataclass +import importlib +from importlib.machinery import ModuleSpec +import os +from pathlib import Path +import sys +import types + +from .backends import BackendAdapter +from .results import CaseResult + + +DEFAULT_KERNEL_ROOT = Path(__file__).resolve().parents[1] / "kernels" + + +class RegistryError(RuntimeError): + """Raised when registry discovery fails.""" + + +@dataclass(frozen=True) +class OperatorSpec: + """Framework contract for one registered kernel operator.""" + + name: str + default_backend: str + backend_names: tuple[str, ...] + create_backend: Callable[[str], BackendAdapter] + list_cases: Callable[[str], Mapping[str, object]] + verify: Callable[[str, object, object], CaseResult] + cycle_fields: Callable[[str, object, BackendAdapter], Mapping[str, object]] + summary: str = "" + + +def _empty_verify(case_id: str, case: object, output: object) -> CaseResult: + del case_id, case, output + raise NotImplementedError("operator spec does not provide a verifier") + + +def _empty_cycle_fields( + case_id: str, + case: object, + backend: BackendAdapter, +) -> Mapping[str, object]: + del case_id, case, backend + return {} + + +def make_operator_spec( + *, + name: str, + default_backend: str = "cce", + backend_names: Iterable[str] = ("cce",), + create_backend: Callable[[str], BackendAdapter], + list_cases: Callable[[str], Mapping[str, object]], + verify: Callable[[str, object, object], CaseResult] = _empty_verify, + cycle_fields: Callable[[str, object, BackendAdapter], Mapping[str, object]] = _empty_cycle_fields, + summary: str = "", +) -> OperatorSpec: + """Build a normalized operator spec with immutable backend names.""" + + return OperatorSpec( + name=name, + default_backend=default_backend, + backend_names=tuple(backend_names), + create_backend=create_backend, + list_cases=list_cases, + verify=verify, + cycle_fields=cycle_fields, + summary=summary, + ) + + +class KernelRegistry: + """In-memory mapping of kernel name to discovered spec.""" + + def __init__(self) -> None: + self._entries: dict[str, OperatorSpec] = {} + + def register(self, spec: OperatorSpec) -> None: + existing = self._entries.get(spec.name) + if existing is not None: + raise RegistryError(f"duplicate kernel registration: {spec.name}") + self._entries[spec.name] = spec + + def get(self, name: str) -> OperatorSpec | None: + return self._entries.get(name) + + def list_names(self) -> list[str]: + return sorted(self._entries.keys()) + + +def _resolve_kernel_dir(kernel_dir: str | os.PathLike[str] | None) -> Path: + candidate = DEFAULT_KERNEL_ROOT if kernel_dir is None else Path(kernel_dir).expanduser() + resolved = candidate.resolve() + if not resolved.exists(): + raise RegistryError(f"kernel directory does not exist: {resolved}") + if not resolved.is_dir(): + raise RegistryError(f"kernel directory is not a directory: {resolved}") + return resolved + + +def _looks_like_kernel_package_dir(kernel_dir: Path) -> bool: + if not (kernel_dir / "__init__.py").is_file(): + return False + return any((kernel_dir / name).exists() for name in ("backends.py", "spec.py", "cycle_metrics.py")) + + +def _namespace_package_name(kernel_dir: Path) -> str: + digest = hashlib.sha1(str(kernel_dir).encode("utf-8")).hexdigest()[:12] + return f"_kernel_test_ext_{digest}" + + +def _ensure_namespace_package(module_name: str, search_path: Path) -> None: + existing = sys.modules.get(module_name) + if existing is not None: + module_path = getattr(existing, "__path__", None) + if module_path is None: + raise RegistryError(f"module namespace conflict while loading kernels: {module_name}") + if str(search_path) not in module_path: + module_path.append(str(search_path)) + return + + module = types.ModuleType(module_name) + spec = ModuleSpec(name=module_name, loader=None, is_package=True) + spec.submodule_search_locations = [str(search_path)] + module.__file__ = str(search_path) + module.__package__ = module_name + module.__path__ = list(spec.submodule_search_locations) + module.__spec__ = spec + sys.modules[module_name] = module + + +def _module_namespace_for_kernel_dir(kernel_dir: Path) -> tuple[str, str | None]: + namespace = _namespace_package_name(kernel_dir) + if _looks_like_kernel_package_dir(kernel_dir): + _ensure_namespace_package(namespace, kernel_dir.parent) + return namespace, kernel_dir.name + + _ensure_namespace_package(namespace, kernel_dir) + return namespace, None + + +def _iter_kernel_module_names(kernel_dir: Path) -> tuple[str, ...]: + _, single_kernel_name = _module_namespace_for_kernel_dir(kernel_dir) + if single_kernel_name is not None: + return (single_kernel_name,) + + names: list[str] = [] + for child in sorted(kernel_dir.iterdir()): + if child.name.startswith("_") or not child.is_dir(): + continue + if _looks_like_kernel_package_dir(child): + names.append(child.name) + return tuple(names) + + +def import_kernel_module( + kernel_name: str, + *, + kernel_dir: str | os.PathLike[str] | None = None, + submodule: str | None = None, +): + """Import one kernel package or submodule from the requested kernel directory.""" + + resolved_dir = _resolve_kernel_dir(kernel_dir) + namespace, single_kernel_name = _module_namespace_for_kernel_dir(resolved_dir) + if single_kernel_name is not None and kernel_name != single_kernel_name: + raise ModuleNotFoundError( + f"kernel directory {resolved_dir} only exposes package {single_kernel_name!r}, " + f"not {kernel_name!r}" + ) + + qualified_name = f"{namespace}.{single_kernel_name or kernel_name}" + if submodule: + qualified_name = f"{qualified_name}.{submodule}" + return importlib.import_module(qualified_name) + + +def _load_from_module(module_name: str, registry: KernelRegistry) -> None: + module = importlib.import_module(module_name) + + if hasattr(module, "register"): + module.register(registry) + return + + if hasattr(module, "get_operator_spec"): + registry.register(module.get_operator_spec()) + return + + if hasattr(module, "OPERATOR_SPEC"): + registry.register(module.OPERATOR_SPEC) + return + + if hasattr(module, "get_kernel_spec"): + registry.register(module.get_kernel_spec()) + return + + if hasattr(module, "KERNEL_SPEC"): + registry.register(module.KERNEL_SPEC) + return + + raise RegistryError( + f"kernel module {module_name!r} must expose register(), get_kernel_spec(), or KERNEL_SPEC" + ) + +def load_registry(kernel_dir: str | os.PathLike[str] | None = None) -> KernelRegistry: + """Load all kernel operators from the default or requested kernel directory.""" + + resolved_dir = _resolve_kernel_dir(kernel_dir) + registry = KernelRegistry() + namespace, _ = _module_namespace_for_kernel_dir(resolved_dir) + + for kernel_name in _iter_kernel_module_names(resolved_dir): + _load_from_module(f"{namespace}.{kernel_name}", registry) + + return registry diff --git a/test/kernel-test/kernel_test/results.py b/test/kernel-test/kernel_test/results.py new file mode 100644 index 0000000000..7f1cde064c --- /dev/null +++ b/test/kernel-test/kernel_test/results.py @@ -0,0 +1,36 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""Result models shared by framework runners.""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class CaseResult: + """Normalized result for one case execution.""" + + ok: bool + message: str + skipped: bool = False + + +@dataclass(frozen=True) +class RunSummary: + """Aggregate summary for one correctness run.""" + + total: int + passed: int + failed: int + skipped: int + + @property + def all_passed(self) -> bool: + return self.failed == 0 diff --git a/test/kernel-test/kernel_test/runners.py b/test/kernel-test/kernel_test/runners.py new file mode 100644 index 0000000000..eb8f6f2354 --- /dev/null +++ b/test/kernel-test/kernel_test/runners.py @@ -0,0 +1,154 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""Shared correctness and cycle runners for the kernel-test framework.""" + +from __future__ import annotations + +import os +import time +from collections.abc import Callable, Mapping + +from .backends import ArtifactOutputs, BackendAdapter +from .pto_artifacts import materialize_artifact_plan +from .results import CaseResult, RunSummary + + +def run_correctness_suite( + cases: Mapping[str, object], + *, + backend: BackendAdapter, + verify_case: Callable[[str, object, object], CaseResult], + print_prefix: str = "", + flush: bool = True, +) -> RunSummary: + """Run a correctness suite with stable PASS/FAIL/SKIP output.""" + + total = 0 + passed = 0 + failed = 0 + skipped = 0 + + for case_id in sorted(cases.keys()): + total += 1 + case = cases[case_id] + supported, reason = backend.is_supported(case, purpose="correctness") + if not supported: + skipped += 1 + result = CaseResult( + ok=True, + skipped=True, + message=reason or "backend not wired for this case", + ) + else: + outputs = backend.launch(case, purpose="correctness") + result = verify_case(case_id, case, outputs) + if result.ok: + passed += 1 + else: + failed += 1 + + status = "SKIP" if result.skipped else ("PASS" if result.ok else "FAIL") + print(f"{print_prefix}[{case_id}] {status}: {result.message}", flush=flush) + + return RunSummary(total=total, passed=passed, failed=failed, skipped=skipped) + + +def _emit_backend_artifacts( + backend: BackendAdapter, + case_id: str, + case: object, + *, + root_alias: bool, +) -> ArtifactOutputs | None: + build_plan = getattr(backend, "build_artifact_plan", None) + if build_plan is None: + return None + return materialize_artifact_plan(case_id, build_plan(case_id, case), root_alias=root_alias) + + +def run_artifact_suite( + cases: Mapping[str, object], + *, + backend: BackendAdapter, + print_prefix: str = "", + flush: bool = True, +) -> RunSummary: + """Generate PTO artifacts for one or more cases with stable output.""" + + total = 0 + passed = 0 + failed = 0 + skipped = 0 + root_alias = len(cases) == 1 + + for case_id in sorted(cases.keys()): + total += 1 + case = cases[case_id] + supported, reason = backend.is_supported(case, purpose="correctness") + if not supported: + skipped += 1 + result = CaseResult( + ok=True, + skipped=True, + message=reason or "backend not wired for this case", + ) + else: + artifacts = _emit_backend_artifacts(backend, case_id, case, root_alias=root_alias) + if artifacts is None: + skipped += 1 + result = CaseResult( + ok=True, + skipped=True, + message=f"backend={backend.name} does not implement PTO artifact emission", + ) + else: + passed += 1 + result = CaseResult(ok=True, message=artifacts.message) + + status = "SKIP" if result.skipped else ("PASS" if result.ok else "FAIL") + print(f"{print_prefix}[{case_id}] {status}: {result.message}", flush=flush) + + return RunSummary(total=total, passed=passed, failed=failed, skipped=skipped) + + +def format_cycle_fields(**fields: object) -> str: + """Format key/value fields in a stable single-line marker payload.""" + + return " ".join(f"{key}={value}" for key, value in fields.items()) + + +def run_cycle_probe( + *, + case_id: str, + case: object, + backend: BackendAdapter, + marker_fields: Mapping[str, object], + flush_wait_env: str = "KERNEL_TEST_RECORD_FLUSH_WAIT", + default_flush_wait_s: float = 0.0, +) -> int: + """Run one cycle probe with consistent marker and skip output.""" + + supported, reason = backend.is_supported(case, purpose="cycle") + if not supported: + payload = dict(marker_fields) + if reason: + payload["reason"] = reason + print(f"CYCLE_SKIP {format_cycle_fields(**payload)}", flush=True) + return 0 + + marker = {"case": case_id, **dict(marker_fields)} + rendered = format_cycle_fields(**marker) + print(f"CYCLE_MARKER {rendered}", flush=True) + backend.launch(case, purpose="cycle") + print(f"CYCLE_DONE {rendered}", flush=True) + + wait_s = float(os.environ.get(flush_wait_env, str(default_flush_wait_s))) + if wait_s > 0.0: + time.sleep(wait_s) + return 0 diff --git a/test/kernel-test/kernels/dequant/__init__.py b/test/kernel-test/kernels/dequant/__init__.py new file mode 100644 index 0000000000..5c2664ea77 --- /dev/null +++ b/test/kernel-test/kernels/dequant/__init__.py @@ -0,0 +1,31 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""dequant kernel-test adapter.""" + +from __future__ import annotations + +from kernel_test.registry import OperatorSpec, make_operator_spec + +from .backends import create_backend +from .spec import cycle_fields, list_cases, verify_case + + +def get_operator_spec() -> OperatorSpec: + """Return the dequant operator registration.""" + + return make_operator_spec( + name="dequant", + default_backend="vmi", + backend_names=("vmi",), + create_backend=create_backend, + list_cases=list_cases, + verify=verify_case, + cycle_fields=cycle_fields, + summary="Runtime correctness adapter for the dequant VMI rewrite", + ) diff --git a/test/kernel-test/kernels/dequant/anti_mx_quant_tail_axis.h b/test/kernel-test/kernels/dequant/anti_mx_quant_tail_axis.h new file mode 100644 index 0000000000..2e41e296af --- /dev/null +++ b/test/kernel-test/kernels/dequant/anti_mx_quant_tail_axis.h @@ -0,0 +1,582 @@ +#ifndef ANTI_MX_QUANT_TAIL_AXIS_H_ +#define ANTI_MX_QUANT_TAIL_AXIS_H_ + +#include "anti_mx_quant_common.h" + +namespace AntiMxQuant { +using namespace AscendC; + +template +class AntiMxQuantTailAxis { +public: + __aicore__ inline AntiMxQuantTailAxis() {} + __aicore__ inline void Init(GM_ADDR x, GM_ADDR mxScale, GM_ADDR y, + const AntiMxQuantTilingData* tilingData); + __aicore__ inline void Process(); + +private: + __aicore__ inline void ParseTilingData(const AntiMxQuantTilingData* tilingData); + __aicore__ inline void GetGmParams(); + __aicore__ inline void GetUbParams(); + + __aicore__ inline void CopyIn(int64_t rowLoopIdx, int64_t colLoopIdx, + int64_t ubFactorRowNum, int64_t ubFactorColNum, int64_t ubFactorColBlockNum); + __aicore__ inline void CopyOut(int64_t rowLoopIdx, int64_t colLoopIdx, + int64_t ubFactorRowNum, int64_t ubFactorColNum, int64_t ubFactorColBlockNum); + + __aicore__ inline void Compute(int64_t rowBlockNum, int64_t colBlockNum); + + __aicore__ inline void ComputeScale(__ubuf__ uint8_t* scaleLocalAddr, __ubuf__ float* scaleBufAddr, int64_t scaleNum); + __aicore__ inline void ComputeScale(__ubuf__ uint8_t* scaleLocalAddr, __ubuf__ bfloat16_t* scaleBufAddr, int64_t scaleNum); + + __aicore__ inline void ComputeData(__ubuf__ uint8_t* xLocalAddr, __ubuf__ float* scaleBufAddr, __ubuf__ U* yLocalAddr, uint16_t loopNum2VF); + __aicore__ inline void ComputeData(__ubuf__ uint8_t* xLocalAddr, __ubuf__ bfloat16_t* scaleBufAddr, __ubuf__ U* yLocalAddr, uint16_t loopNum2VF); + +private: + TPipe pipe_; + TQue inQueueX_; + TQue inQueueScale_; + TQue outQueueY_; + + TBuf scaleBuffer_; + + GlobalTensor xGm_; + GlobalTensor mxScaleGm_; + GlobalTensor yGm_; + + int64_t totalCoreNum_{0}; + int64_t usedCoreNum_{0}; + int64_t rowTileNum_{0}; + int64_t colTileNum_{0}; + int64_t rowNum_{1}; + int64_t colNum_{1}; + int64_t colNormalBlockNum_{0}; + int64_t colTailLen_{0}; + int64_t rowNormalBlockNum_{0}; + int64_t rowTailLen_{0}; + int64_t maxUbBlockNum_{0}; + int64_t dstType_{0}; + + int64_t coreIdx_{0}; + int64_t coreColIdx_{0}; + int64_t coreRowIdx_{0}; + int64_t xGmOffset_{0}; + int64_t scaleGmOffset_{0}; + int64_t scaleColNum_{0}; + + int64_t ubFactorColBlockNum_{0}; + int64_t ubFactorColNum_{0}; + int64_t ubFactorRowNum_{0}; + int64_t ubFactorColLoopNum_{0}; + int64_t ubFactorRowLoopNum_{0}; + int64_t ubFactorColNormalBlockNum_{0}; + int64_t ubFactorColTailBlockNum_{0}; + int64_t ubFactorColNormalLen_{0}; + int64_t ubFactorColTailLen_{0}; + int64_t ubFactorRowNormalNum_{0}; + int64_t ubFactorRowTailNum_{0}; +}; + +template +__aicore__ inline void AntiMxQuantTailAxis::Init( + GM_ADDR x, GM_ADDR mxScale, GM_ADDR y, const AntiMxQuantTilingData* tilingData) +{ + ParseTilingData(tilingData); + GetGmParams(); + GetUbParams(); + + int64_t xBufSize = 0; + int64_t scaleBufSize = 0; + int64_t yBufSize = 0; + int64_t scaleComputeBufSize = 0; + + if constexpr (IsFp8Type()) { + xBufSize = maxUbBlockNum_ * BLOCK_SIZE * sizeof(uint8_t); + scaleBufSize = ((maxUbBlockNum_ + UBBlockSize_ - 1) / UBBlockSize_) * UBBlockSize_ * sizeof(uint8_t); + yBufSize = maxUbBlockNum_ * BLOCK_SIZE * sizeof(U); + int64_t maxScaleNum = maxUbBlockNum_; + int64_t scaleLoopNum = (maxScaleNum + vfLen8 - 1) / vfLen8; + scaleComputeBufSize = scaleLoopNum * vfLen8 * sizeof(float); + } else { + xBufSize = maxUbBlockNum_ * BLOCK_SIZE / DIGIT_TWO; + scaleBufSize = ((maxUbBlockNum_ + UBBlockSize_ - 1) / UBBlockSize_) * UBBlockSize_ * sizeof(uint8_t); + yBufSize = maxUbBlockNum_ * BLOCK_SIZE * sizeof(U); + int64_t maxScaleNum = maxUbBlockNum_; + int64_t scaleLoopNum = (maxScaleNum + vfLen8 - 1) / vfLen8; + scaleComputeBufSize = scaleLoopNum * vfLen8 * sizeof(bfloat16_t); + } + + pipe_.InitBuffer(inQueueX_, DB_BUFFER, xBufSize); + pipe_.InitBuffer(inQueueScale_, DB_BUFFER, scaleBufSize); + pipe_.InitBuffer(outQueueY_, DB_BUFFER, yBufSize); + pipe_.InitBuffer(scaleBuffer_, scaleComputeBufSize); + + if constexpr (IsFp8Type()) { + xGm_.SetGlobalBuffer((__gm__ uint8_t*)x + xGmOffset_); + yGm_.SetGlobalBuffer((__gm__ U*)y + xGmOffset_); + } else { + xGm_.SetGlobalBuffer((__gm__ uint8_t*)x + xGmOffset_ / DIGIT_TWO); + yGm_.SetGlobalBuffer((__gm__ U*)y + xGmOffset_); + } + mxScaleGm_.SetGlobalBuffer((__gm__ uint8_t*)mxScale + scaleGmOffset_); +} + +template +__aicore__ inline void AntiMxQuantTailAxis::ParseTilingData(const AntiMxQuantTilingData* tilingData) +{ + totalCoreNum_ = tilingData->totalCoreNum; + usedCoreNum_ = tilingData->usedCoreNum; + rowTileNum_ = tilingData->rowTileNum; + colTileNum_ = tilingData->colTileNum; + rowNum_ = tilingData->rowNum; + colNum_ = tilingData->colNum; + colNormalBlockNum_ = tilingData->colNormalBlockNum; + colTailLen_ = tilingData->colTailLen; + rowNormalBlockNum_ = tilingData->rowNormalBlockNum; + rowTailLen_ = tilingData->rowTailLen; + maxUbBlockNum_ = tilingData->maxUbBlockNum; + dstType_ = tilingData->dstType; +} + +template +__aicore__ inline void AntiMxQuantTailAxis::GetGmParams() +{ + coreIdx_ = GetBlockIdx(); + coreColIdx_ = coreIdx_ % colTileNum_; + coreRowIdx_ = coreIdx_ / colTileNum_; + xGmOffset_ = coreRowIdx_ * rowNormalBlockNum_ * colNum_ + coreColIdx_ * colNormalBlockNum_ * SPLIT_N; + scaleColNum_ = (((((colNum_ + BLOCK_SIZE - 1) / BLOCK_SIZE) + DIGIT_TWO - 1) / DIGIT_TWO) * DIGIT_TWO); + scaleGmOffset_ = coreRowIdx_ * rowNormalBlockNum_ * scaleColNum_ + coreColIdx_ * colNormalBlockNum_ * DIGIT_SIXTEEN; +} + +template +__aicore__ inline void AntiMxQuantTailAxis::GetUbParams() +{ + if (coreColIdx_ == colTileNum_ - 1) { + ubFactorColBlockNum_ = (colTailLen_ + BLOCK_SIZE - 1) / BLOCK_SIZE; + ubFactorColNum_ = colTailLen_; + } else { + ubFactorColBlockNum_ = colNormalBlockNum_ * DIGIT_SIXTEEN; + ubFactorColNum_ = ubFactorColBlockNum_ * BLOCK_SIZE; + } + + if (coreRowIdx_ == rowTileNum_ - 1) { + ubFactorRowNum_ = rowTailLen_; + } else { + ubFactorRowNum_ = rowNormalBlockNum_; + } + + ubFactorColLoopNum_ = (ubFactorColBlockNum_ + maxUbBlockNum_ - 1) / maxUbBlockNum_; + ubFactorColNormalBlockNum_ = (ubFactorColBlockNum_ + ubFactorColLoopNum_ - 1) / ubFactorColLoopNum_; + ubFactorColNormalBlockNum_ = ((ubFactorColNormalBlockNum_ + DIGIT_TWO - 1) / DIGIT_TWO) * DIGIT_TWO; + ubFactorColTailBlockNum_ = ubFactorColBlockNum_ - (ubFactorColLoopNum_ - DIGIT_ONE) * ubFactorColNormalBlockNum_; + ubFactorColNormalLen_ = ubFactorColNormalBlockNum_ * BLOCK_SIZE; + ubFactorColTailLen_ = ubFactorColNum_ - (ubFactorColLoopNum_ - DIGIT_ONE) * ubFactorColNormalLen_; + + ubFactorRowNormalNum_ = maxUbBlockNum_ / ubFactorColNormalBlockNum_; + ubFactorRowLoopNum_ = (ubFactorRowNum_ + ubFactorRowNormalNum_ - 1) / ubFactorRowNormalNum_; + ubFactorRowNormalNum_ = (ubFactorRowNum_ + ubFactorRowLoopNum_ - 1) / ubFactorRowLoopNum_; + ubFactorRowTailNum_ = ubFactorRowNum_ - (ubFactorRowLoopNum_ - DIGIT_ONE) * ubFactorRowNormalNum_; +} + +template +__aicore__ inline void AntiMxQuantTailAxis::Process() +{ + if (coreIdx_ >= usedCoreNum_) { + return; + } + + int64_t rowLoopIdx = 0; + int64_t colLoopIdx = 0; + + for (rowLoopIdx = 0; rowLoopIdx < ubFactorRowLoopNum_ - 1; rowLoopIdx++) { + for (colLoopIdx = 0; colLoopIdx < ubFactorColLoopNum_ - 1; colLoopIdx++) { + CopyIn(rowLoopIdx, colLoopIdx, ubFactorRowNormalNum_, ubFactorColNormalLen_, ubFactorColNormalBlockNum_); + Compute(ubFactorRowNormalNum_, ubFactorColNormalBlockNum_); + CopyOut(rowLoopIdx, colLoopIdx, ubFactorRowNormalNum_, ubFactorColNormalLen_, ubFactorColNormalBlockNum_); + } + CopyIn(rowLoopIdx, colLoopIdx, ubFactorRowNormalNum_, ubFactorColTailLen_, ubFactorColTailBlockNum_); + Compute(ubFactorRowNormalNum_, ubFactorColTailBlockNum_); + CopyOut(rowLoopIdx, colLoopIdx, ubFactorRowNormalNum_, ubFactorColTailLen_, ubFactorColTailBlockNum_); + } + + for (colLoopIdx = 0; colLoopIdx < ubFactorColLoopNum_ - 1; colLoopIdx++) { + CopyIn(rowLoopIdx, colLoopIdx, ubFactorRowTailNum_, ubFactorColNormalLen_, ubFactorColNormalBlockNum_); + Compute(ubFactorRowTailNum_, ubFactorColNormalBlockNum_); + CopyOut(rowLoopIdx, colLoopIdx, ubFactorRowTailNum_, ubFactorColNormalLen_, ubFactorColNormalBlockNum_); + } + CopyIn(rowLoopIdx, colLoopIdx, ubFactorRowTailNum_, ubFactorColTailLen_, ubFactorColTailBlockNum_); + Compute(ubFactorRowTailNum_, ubFactorColTailBlockNum_); + CopyOut(rowLoopIdx, colLoopIdx, ubFactorRowTailNum_, ubFactorColTailLen_, ubFactorColTailBlockNum_); +} + +template +__aicore__ inline void AntiMxQuantTailAxis::CopyIn( + int64_t rowLoopIdx, int64_t colLoopIdx, int64_t ubFactorRowNum, int64_t ubFactorColNum, int64_t ubFactorColBlockNum) +{ + LocalTensor xLocal = inQueueX_.AllocTensor(); + int64_t xOffset = rowLoopIdx * ubFactorRowNormalNum_ * colNum_ + colLoopIdx * ubFactorColNormalLen_; + if constexpr (IsFp4Type()) { + xOffset /= DIGIT_TWO; + } + + if constexpr (IsFp8Type()) { + int64_t scaleIsNotOdd = ubFactorColBlockNum % DIGIT_TWO; + int64_t xPad = ((ubFactorColNum + UBBlockSize_ - 1) / UBBlockSize_) * UBBlockSize_ - ubFactorColNum; + DataCopyExtParams copyInParamX = { + static_cast(ubFactorRowNum), + static_cast(ubFactorColNum * sizeof(uint8_t)), + static_cast((colNum_ - ubFactorColNum) * sizeof(uint8_t)), + static_cast(scaleIsNotOdd), 0}; + DataCopyPadExtParams xPadParams{true, 0, static_cast(xPad), 0}; + DataCopyPad(xLocal, xGm_[xOffset], copyInParamX, xPadParams); + } else { + int64_t xBytes = ubFactorColNum / 2; + int64_t xPadBytes = ((xBytes + UBBlockSize_ - 1) / UBBlockSize_) * UBBlockSize_ - xBytes; + DataCopyExtParams copyInParamX = { + static_cast(ubFactorRowNum), + static_cast(xBytes), + static_cast((colNum_ - ubFactorColNum) / 2), + 0, 0}; + DataCopyPadExtParams xPadParams{true, 0, static_cast(xPadBytes), 0}; + DataCopyPad(xLocal, xGm_[xOffset], copyInParamX, xPadParams); + } + inQueueX_.EnQue(xLocal); + + LocalTensor scaleLocal = inQueueScale_.AllocTensor(); + int64_t scaleNum = (ubFactorColBlockNum + DIGIT_TWO - 1) / DIGIT_TWO * DIGIT_TWO; + int64_t scaleOffset = rowLoopIdx * ubFactorRowNormalNum_ * scaleColNum_ + colLoopIdx * ubFactorColNormalBlockNum_; + + DataCopyExtParams copyInParamScale = { + static_cast(ubFactorRowNum), + static_cast(scaleNum), + static_cast((scaleColNum_ - scaleNum)), + 0, 0}; + DataCopyPadExtParams scalePadParams{false, 0, 0, 0}; + DataCopyPad(scaleLocal, mxScaleGm_[scaleOffset], copyInParamScale, scalePadParams); + inQueueScale_.EnQue(scaleLocal); +} + +template +__aicore__ inline void AntiMxQuantTailAxis::CopyOut( + int64_t rowLoopIdx, int64_t colLoopIdx, int64_t ubFactorRowNum, int64_t ubFactorColNum, int64_t ubFactorColBlockNum) +{ + int64_t scaleIsNotOdd = ubFactorColBlockNum % DIGIT_TWO; + int64_t anotherStride = + ((ubFactorColNum + BLOCK_SIZE - 1) / BLOCK_SIZE * BLOCK_SIZE - ubFactorColNum) * sizeof(U) / UBBlockSize_; + int64_t srcStride = scaleIsNotOdd * sizeof(U) + anotherStride; + + LocalTensor yLocal = outQueueY_.DeQue(); + int64_t yOffset = rowLoopIdx * ubFactorRowNormalNum_ * colNum_ + colLoopIdx * ubFactorColNormalLen_; + + DataCopyExtParams copyOutParamY = { + static_cast(ubFactorRowNum), + static_cast(ubFactorColNum * sizeof(U)), + static_cast(srcStride), + static_cast((colNum_ - ubFactorColNum) * sizeof(U)), 0}; + DataCopyPad(yGm_[yOffset], yLocal, copyOutParamY); + outQueueY_.FreeTensor(yLocal); +} + +template +__aicore__ inline void AntiMxQuantTailAxis::Compute(int64_t rowBlockNum, int64_t colBlockNum) +{ + colBlockNum += colBlockNum % DIGIT_TWO; + int64_t totalScaleNum = rowBlockNum * colBlockNum; + uint32_t totalBlockNum = static_cast(rowBlockNum * colBlockNum); + uint16_t loopNum2VF = static_cast( + (static_cast(totalBlockNum) + static_cast(DIGIT_SIXTEEN) - 1) / static_cast(DIGIT_SIXTEEN)); + + LocalTensor xLocal = inQueueX_.DeQue(); + LocalTensor scaleLocal = inQueueScale_.DeQue(); + LocalTensor yLocal = outQueueY_.AllocTensor(); + + auto xLocalAddr = reinterpret_cast<__ubuf__ uint8_t*>(xLocal.GetPhyAddr()); + auto yLocalAddr = reinterpret_cast<__ubuf__ U*>(yLocal.GetPhyAddr()); + auto scaleLocalAddr = reinterpret_cast<__ubuf__ uint8_t*>(scaleLocal.GetPhyAddr()); + + if constexpr (IsFp8Type()) { + auto scaleBufAddr = reinterpret_cast<__ubuf__ float*>(scaleBuffer_.Get().GetPhyAddr()); + ComputeScale(scaleLocalAddr, scaleBufAddr, totalScaleNum); + ComputeData(xLocalAddr, scaleBufAddr, yLocalAddr, loopNum2VF); + } else { + auto scaleBufAddr = reinterpret_cast<__ubuf__ bfloat16_t*>(scaleBuffer_.Get().GetPhyAddr()); + ComputeScale(scaleLocalAddr, scaleBufAddr, totalScaleNum); + ComputeData(xLocalAddr, scaleBufAddr, yLocalAddr, loopNum2VF); + } + + inQueueX_.FreeTensor(xLocal); + inQueueScale_.FreeTensor(scaleLocal); + outQueueY_.EnQue(yLocal); +} + +template +__aicore__ inline void AntiMxQuantTailAxis::ComputeScale( + __ubuf__ uint8_t* scaleLocalAddr, __ubuf__ float* scaleBufAddr, int64_t scaleNum) +{ + uint16_t loopNum = static_cast((scaleNum + vfLen8 - 1) / vfLen8); + __VEC_SCOPE__ + { + Reg::MaskReg scaleMask = Reg::CreateMask(); + Reg::MaskReg storeMask = Reg::CreateMask(); + Reg::RegTensor vdScale0, vdScale1, vdu8_zero; + Reg::RegTensor vdScaleBf16_0, vdScaleBf16_1; + Reg::RegTensor vdScaleFp32_0_0, vdScaleFp32_0_1; + Reg::RegTensor vdScaleFp32_1_0, vdScaleFp32_1_1; + + Reg::Duplicate(vdu8_zero, 0); + + for (uint16_t i = 0; i < loopNum; i++) { + Reg::LoadAlign(vdScale0, scaleLocalAddr, vfLen8); + Reg::Interleave(vdScale0, vdScale1, vdScale0, vdu8_zero); + + Reg::Cast(vdScaleBf16_0, (Reg::RegTensor&)vdScale0, scaleMask); + Reg::Cast(vdScaleBf16_1, (Reg::RegTensor&)vdScale1, scaleMask); + + Reg::Cast(vdScaleFp32_0_0, vdScaleBf16_0, scaleMask); + Reg::Cast(vdScaleFp32_0_1, vdScaleBf16_0, scaleMask); + Reg::Cast(vdScaleFp32_1_0, vdScaleBf16_1, scaleMask); + Reg::Cast(vdScaleFp32_1_1, vdScaleBf16_1, scaleMask); + + Reg::Interleave(vdScaleFp32_0_0, vdScaleFp32_0_1, vdScaleFp32_0_0, vdScaleFp32_0_1); + Reg::Interleave(vdScaleFp32_1_0, vdScaleFp32_1_1, vdScaleFp32_1_0, vdScaleFp32_1_1); + + Reg::StoreAlign(scaleBufAddr, vdScaleFp32_0_0, vfLen32, storeMask); + Reg::StoreAlign(scaleBufAddr, vdScaleFp32_0_1, vfLen32, storeMask); + Reg::StoreAlign(scaleBufAddr, vdScaleFp32_1_0, vfLen32, storeMask); + Reg::StoreAlign(scaleBufAddr, vdScaleFp32_1_1, vfLen32, storeMask); + } + } +} + +template +__aicore__ inline void AntiMxQuantTailAxis::ComputeScale( + __ubuf__ uint8_t* scaleLocalAddr, __ubuf__ bfloat16_t* scaleBufAddr, int64_t scaleNum) +{ + uint16_t loopNum = static_cast((scaleNum + vfLen8 - 1) / vfLen8); + __VEC_SCOPE__ + { + Reg::MaskReg scaleMask = Reg::CreateMask(); + Reg::MaskReg storeMask = Reg::CreateMask(); + Reg::RegTensor vdScale0, vdScale1, vdu8_zero; + Reg::RegTensor vdScaleBf16_0, vdScaleBf16_1; + + Reg::Duplicate(vdu8_zero, 0); + + for (uint16_t i = 0; i < loopNum; i++) { + Reg::LoadAlign(vdScale0, scaleLocalAddr, vfLen8); + Reg::Interleave(vdScale0, vdScale1, vdScale0, vdu8_zero); + + Reg::Cast(vdScaleBf16_0, (Reg::RegTensor&)vdScale0, scaleMask); + Reg::Cast(vdScaleBf16_1, (Reg::RegTensor&)vdScale1, scaleMask); + + Reg::StoreAlign(scaleBufAddr, vdScaleBf16_0, vfLen16, storeMask); + Reg::StoreAlign(scaleBufAddr, vdScaleBf16_1, vfLen16, storeMask); + } + } +} + +template +__aicore__ inline void AntiMxQuantTailAxis::ComputeData( + __ubuf__ uint8_t* xLocalAddr, __ubuf__ float* scaleBufAddr, __ubuf__ U* yLocalAddr, uint16_t loopNum2VF) +{ + __VEC_SCOPE__ + { + Reg::MaskReg maskAll = Reg::CreateMask(); + Reg::MaskReg maskFp32 = Reg::CreateMask(); + Reg::MaskReg maskFp8 = Reg::CreateMask(); + + Reg::RegTensor vdFp8_0, vdFp8_1; + Reg::RegTensor vdFp32_0_0, vdFp32_0_1, vdFp32_0_2, vdFp32_0_3; + Reg::RegTensor vdFp32_1_0, vdFp32_1_1, vdFp32_1_2, vdFp32_1_3; + Reg::RegTensor vdScale_0, vdScale_1; + + for (uint16_t i = 0; i < loopNum2VF; i++) { + Reg::LoadAlign(vdFp8_0, vdFp8_1, xLocalAddr, vfLen8Double); + + Reg::Interleave(vdFp8_0, vdFp8_1, vdFp8_0, vdFp8_1); + Reg::Cast(vdFp32_0_0, (Reg::RegTensor&)vdFp8_0, maskFp8); + Reg::Cast(vdFp32_0_1, (Reg::RegTensor&)vdFp8_0, maskFp8); + Reg::Cast(vdFp32_0_2, (Reg::RegTensor&)vdFp8_0, maskFp8); + Reg::Cast(vdFp32_0_3, (Reg::RegTensor&)vdFp8_0, maskFp8); + Reg::Cast(vdFp32_1_0, (Reg::RegTensor&)vdFp8_1, maskFp8); + Reg::Cast(vdFp32_1_1, (Reg::RegTensor&)vdFp8_1, maskFp8); + Reg::Cast(vdFp32_1_2, (Reg::RegTensor&)vdFp8_1, maskFp8); + Reg::Cast(vdFp32_1_3, (Reg::RegTensor&)vdFp8_1, maskFp8); + + Reg::LoadAlign(vdScale_0, scaleBufAddr, elementAfterReduce_); + Reg::LoadAlign(vdScale_1, scaleBufAddr, elementAfterReduce_); + + Reg::Mul(vdFp32_0_0, vdFp32_0_0, vdScale_0, maskFp32); + Reg::Mul(vdFp32_0_1, vdFp32_0_1, vdScale_0, maskFp32); + Reg::Mul(vdFp32_0_2, vdFp32_0_2, vdScale_0, maskFp32); + Reg::Mul(vdFp32_0_3, vdFp32_0_3, vdScale_0, maskFp32); + Reg::Mul(vdFp32_1_0, vdFp32_1_0, vdScale_1, maskFp32); + Reg::Mul(vdFp32_1_1, vdFp32_1_1, vdScale_1, maskFp32); + Reg::Mul(vdFp32_1_2, vdFp32_1_2, vdScale_1, maskFp32); + Reg::Mul(vdFp32_1_3, vdFp32_1_3, vdScale_1, maskFp32); + + Reg::Interleave(vdFp32_0_0, vdFp32_0_2, vdFp32_0_0, vdFp32_0_2); + Reg::Interleave(vdFp32_0_1, vdFp32_0_3, vdFp32_0_1, vdFp32_0_3); + Reg::Interleave(vdFp32_1_0, vdFp32_1_2, vdFp32_1_0, vdFp32_1_2); + Reg::Interleave(vdFp32_1_1, vdFp32_1_3, vdFp32_1_1, vdFp32_1_3); + + if constexpr (IsFp32Type()) { + Reg::Interleave(vdFp32_0_0, vdFp32_0_1, vdFp32_0_0, vdFp32_0_1); + Reg::Interleave(vdFp32_0_2, vdFp32_0_3, vdFp32_0_2, vdFp32_0_3); + Reg::Interleave(vdFp32_1_0, vdFp32_1_1, vdFp32_1_0, vdFp32_1_1); + Reg::Interleave(vdFp32_1_2, vdFp32_1_3, vdFp32_1_2, vdFp32_1_3); + + Reg::StoreAlign(yLocalAddr, vdFp32_0_0, vfLen32, maskFp32); + Reg::StoreAlign(yLocalAddr, vdFp32_0_1, vfLen32, maskFp32); + Reg::StoreAlign(yLocalAddr, vdFp32_0_2, vfLen32, maskFp32); + Reg::StoreAlign(yLocalAddr, vdFp32_0_3, vfLen32, maskFp32); + Reg::StoreAlign(yLocalAddr, vdFp32_1_0, vfLen32, maskFp32); + Reg::StoreAlign(yLocalAddr, vdFp32_1_1, vfLen32, maskFp32); + Reg::StoreAlign(yLocalAddr, vdFp32_1_2, vfLen32, maskFp32); + Reg::StoreAlign(yLocalAddr, vdFp32_1_3, vfLen32, maskFp32); + } else if constexpr (IsBf16Type()) { + Reg::RegTensor vdBf16_0_z, vdBf16_0_o; + Reg::RegTensor vdBf16_1_z, vdBf16_1_o; + Reg::RegTensor vdBf16_2_z, vdBf16_2_o; + Reg::RegTensor vdBf16_3_z, vdBf16_3_o; + + Reg::Cast(vdBf16_0_z, vdFp32_0_0, maskAll); + Reg::Cast(vdBf16_0_o, vdFp32_0_1, maskAll); + Reg::Add(vdBf16_0_z, vdBf16_0_z, vdBf16_0_o, maskAll); + Reg::StoreAlign(yLocalAddr, vdBf16_0_z, vfLen16, maskAll); + + Reg::Cast(vdBf16_1_z, vdFp32_0_2, maskAll); + Reg::Cast(vdBf16_1_o, vdFp32_0_3, maskAll); + Reg::Add(vdBf16_1_z, vdBf16_1_z, vdBf16_1_o, maskAll); + Reg::StoreAlign(yLocalAddr, vdBf16_1_z, vfLen16, maskAll); + + Reg::Cast(vdBf16_2_z, vdFp32_1_0, maskAll); + Reg::Cast(vdBf16_2_o, vdFp32_1_1, maskAll); + Reg::Add(vdBf16_2_z, vdBf16_2_z, vdBf16_2_o, maskAll); + Reg::StoreAlign(yLocalAddr, vdBf16_2_z, vfLen16, maskAll); + + Reg::Cast(vdBf16_3_z, vdFp32_1_2, maskAll); + Reg::Cast(vdBf16_3_o, vdFp32_1_3, maskAll); + Reg::Add(vdBf16_3_z, vdBf16_3_z, vdBf16_3_o, maskAll); + Reg::StoreAlign(yLocalAddr, vdBf16_3_z, vfLen16, maskAll); + } else { + Reg::RegTensor vdFp16_0_z, vdFp16_0_o; + Reg::RegTensor vdFp16_1_z, vdFp16_1_o; + Reg::RegTensor vdFp16_2_z, vdFp16_2_o; + Reg::RegTensor vdFp16_3_z, vdFp16_3_o; + + Reg::Cast(vdFp16_0_z, vdFp32_0_0, maskAll); + Reg::Cast(vdFp16_0_o, vdFp32_0_1, maskAll); + Reg::Add(vdFp16_0_z, vdFp16_0_z, vdFp16_0_o, maskAll); + Reg::StoreAlign(yLocalAddr, vdFp16_0_z, vfLen16, maskAll); + + Reg::Cast(vdFp16_1_z, vdFp32_0_2, maskAll); + Reg::Cast(vdFp16_1_o, vdFp32_0_3, maskAll); + Reg::Add(vdFp16_1_z, vdFp16_1_z, vdFp16_1_o, maskAll); + Reg::StoreAlign(yLocalAddr, vdFp16_1_z, vfLen16, maskAll); + + Reg::Cast(vdFp16_2_z, vdFp32_1_0, maskAll); + Reg::Cast(vdFp16_2_o, vdFp32_1_1, maskAll); + Reg::Add(vdFp16_2_z, vdFp16_2_z, vdFp16_2_o, maskAll); + Reg::StoreAlign(yLocalAddr, vdFp16_2_z, vfLen16, maskAll); + + Reg::Cast(vdFp16_3_z, vdFp32_1_2, maskAll); + Reg::Cast(vdFp16_3_o, vdFp32_1_3, maskAll); + Reg::Add(vdFp16_3_z, vdFp16_3_z, vdFp16_3_o, maskAll); + Reg::StoreAlign(yLocalAddr, vdFp16_3_z, vfLen16, maskAll); + } + } + } +} + +template +__aicore__ inline void AntiMxQuantTailAxis::ComputeData( + __ubuf__ uint8_t* xLocalAddr, __ubuf__ bfloat16_t* scaleBufAddr, __ubuf__ U* yLocalAddr, uint16_t loopNum2VF) +{ + __VEC_SCOPE__ + { + Reg::MaskReg maskAll16 = Reg::CreateMask(); + Reg::MaskReg maskU8 = Reg::CreateMask(); + + Reg::RegTensor vdFp4U8_0, vdFp4U8_1, vdFp4U8_2, vdFp4U8_3; + Reg::RegTensor vdMerged0, vdMerged1, vdMerged2, vdMerged3; + Reg::RegTensor vdScale0, vdScale1, vdScale2, vdScale3; + Reg::RegTensor vdScaleTmp0, vdScaleTmp1; + + for (uint16_t i = 0; i < loopNum2VF; i++) { + Reg::LoadAlign(vdFp4U8_0, xLocalAddr, vfLen32); + Reg::LoadAlign(vdFp4U8_1, xLocalAddr, vfLen32); + Reg::LoadAlign(vdFp4U8_2, xLocalAddr, vfLen32); + Reg::LoadAlign(vdFp4U8_3, xLocalAddr, vfLen32); + + Reg::Cast(vdMerged0, (Reg::RegTensor&)vdFp4U8_0, maskU8); + Reg::Cast(vdMerged1, (Reg::RegTensor&)vdFp4U8_1, maskU8); + Reg::Cast(vdMerged2, (Reg::RegTensor&)vdFp4U8_2, maskU8); + Reg::Cast(vdMerged3, (Reg::RegTensor&)vdFp4U8_3, maskU8); + + Reg::LoadAlign(vdScaleTmp0, scaleBufAddr, elementAfterReduce_); + Reg::LoadAlign(vdScaleTmp1, scaleBufAddr, elementAfterReduce_); + + Reg::Interleave(vdScale0, vdScale1, vdScaleTmp0, vdScaleTmp0); + Reg::Interleave(vdScale2, vdScale3, vdScaleTmp1, vdScaleTmp1); + + Reg::Mul(vdMerged0, vdMerged0, vdScale0, maskAll16); + Reg::Mul(vdMerged1, vdMerged1, vdScale1, maskAll16); + Reg::Mul(vdMerged2, vdMerged2, vdScale2, maskAll16); + Reg::Mul(vdMerged3, vdMerged3, vdScale3, maskAll16); + + if constexpr (IsFp32Type()) { + Reg::RegTensor vdFp32_0_0, vdFp32_0_1; + Reg::RegTensor vdFp32_1_0, vdFp32_1_1; + Reg::RegTensor vdFp32_2_0, vdFp32_2_1; + Reg::RegTensor vdFp32_3_0, vdFp32_3_1; + + Reg::Cast(vdFp32_0_0, vdMerged0, maskAll16); + Reg::Cast(vdFp32_0_1, vdMerged0, maskAll16); + Reg::Interleave(vdFp32_0_0, vdFp32_0_1, vdFp32_0_0, vdFp32_0_1); + Reg::StoreAlign(yLocalAddr, vdFp32_0_0, vfLen32, maskAll16); + Reg::StoreAlign(yLocalAddr, vdFp32_0_1, vfLen32, maskAll16); + + Reg::Cast(vdFp32_1_0, vdMerged1, maskAll16); + Reg::Cast(vdFp32_1_1, vdMerged1, maskAll16); + Reg::Interleave(vdFp32_1_0, vdFp32_1_1, vdFp32_1_0, vdFp32_1_1); + Reg::StoreAlign(yLocalAddr, vdFp32_1_0, vfLen32, maskAll16); + Reg::StoreAlign(yLocalAddr, vdFp32_1_1, vfLen32, maskAll16); + + Reg::Cast(vdFp32_2_0, vdMerged2, maskAll16); + Reg::Cast(vdFp32_2_1, vdMerged2, maskAll16); + Reg::Interleave(vdFp32_2_0, vdFp32_2_1, vdFp32_2_0, vdFp32_2_1); + Reg::StoreAlign(yLocalAddr, vdFp32_2_0, vfLen32, maskAll16); + Reg::StoreAlign(yLocalAddr, vdFp32_2_1, vfLen32, maskAll16); + + Reg::Cast(vdFp32_3_0, vdMerged3, maskAll16); + Reg::Cast(vdFp32_3_1, vdMerged3, maskAll16); + Reg::Interleave(vdFp32_3_0, vdFp32_3_1, vdFp32_3_0, vdFp32_3_1); + Reg::StoreAlign(yLocalAddr, vdFp32_3_0, vfLen32, maskAll16); + Reg::StoreAlign(yLocalAddr, vdFp32_3_1, vfLen32, maskAll16); + } else if constexpr (IsBf16Type()) { + Reg::StoreAlign(yLocalAddr, vdMerged0, vfLen16, maskAll16); + Reg::StoreAlign(yLocalAddr, vdMerged1, vfLen16, maskAll16); + Reg::StoreAlign(yLocalAddr, vdMerged2, vfLen16, maskAll16); + Reg::StoreAlign(yLocalAddr, vdMerged3, vfLen16, maskAll16); + } else { + Reg::RegTensor vdFp16_0, vdFp16_1; + Reg::RegTensor vdFp16_2, vdFp16_3; + + Reg::Cast(vdFp16_0, vdMerged0, maskAll16); + Reg::Cast(vdFp16_1, vdMerged1, maskAll16); + Reg::StoreAlign(yLocalAddr, vdFp16_0, vfLen16, maskAll16); + Reg::StoreAlign(yLocalAddr, vdFp16_1, vfLen16, maskAll16); + + Reg::Cast(vdFp16_2, vdMerged2, maskAll16); + Reg::Cast(vdFp16_3, vdMerged3, maskAll16); + Reg::StoreAlign(yLocalAddr, vdFp16_2, vfLen16, maskAll16); + Reg::StoreAlign(yLocalAddr, vdFp16_3, vfLen16, maskAll16); + } + } + } +} + +} // namespace AntiMxQuant +#endif \ No newline at end of file diff --git a/test/kernel-test/kernels/dequant/backends.py b/test/kernel-test/kernels/dequant/backends.py new file mode 100644 index 0000000000..64aecccede --- /dev/null +++ b/test/kernel-test/kernels/dequant/backends.py @@ -0,0 +1,29 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""Backend factory for dequant.""" + +from __future__ import annotations + +from kernel_test.backends import BackendAdapter + + +def create_backend(name: str) -> BackendAdapter: + """Create one dequant backend adapter.""" + + if name == "vmi": + try: + from .vmi.backend import AntiMxQuantTailAxisVmiBackend + except (ImportError, ModuleNotFoundError) as exc: + raise RuntimeError( + "dequant vmi backend requires the ptodsl/PTOAS " + "Python environment; use the `pto` conda environment or install " + "the local ptodsl dependencies first" + ) from exc + return AntiMxQuantTailAxisVmiBackend() + raise ValueError(f"unknown dequant backend: {name}") diff --git a/test/kernel-test/kernels/dequant/reference.py b/test/kernel-test/kernels/dequant/reference.py new file mode 100644 index 0000000000..731d763c7c --- /dev/null +++ b/test/kernel-test/kernels/dequant/reference.py @@ -0,0 +1,158 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""Reference case metadata for the dequant kernel.""" + +from __future__ import annotations + +import numpy as np +import torch + +DTYPES = ("f32", "bf16", "f16") +SRC_FORMATS = ("e4m3", "e5m2") +SCALE_FORMAT = "e8m0" +DEFAULT_ROW_BLOCK_NUM = 4 +DEFAULT_COL_BLOCK_NUM = 4 +TOLERANCE = { + "f32": 1e-6, + "bf16": 4e-3, + "f16": 1e-3, +} + +_SEED = 42 +_BLOCK_SIZE = 32 +_BLOCKS_PER_LOOP = 16 +_ELEMS_PER_LOOP = _BLOCK_SIZE * _BLOCKS_PER_LOOP +_RAW_SCALE_BYTES_PER_LOOP = _BLOCKS_PER_LOOP +_E8M0_BYTES = np.array([0x7E, 0x7F, 0x80, 0x81], dtype=np.uint8) + + +def _case_id(src_fmt: str, dst_fmt: str) -> str: + return f"{src_fmt}_{dst_fmt}" + + +def _e8m0_to_f32(bits: np.ndarray) -> np.ndarray: + return np.exp2(np.asarray(bits, dtype=np.uint8).astype(np.int32) - 127).astype(np.float32) + + +def _torch_src_dtype(src_fmt: str) -> torch.dtype: + if src_fmt == "e4m3": + return torch.float8_e4m3fn + if src_fmt == "e5m2": + return torch.float8_e5m2 + raise ValueError(f"unknown src_fmt: {src_fmt}") + + +def _quantize_src_to_bits(src_fmt: str, values: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + values_t = torch.from_numpy(np.asarray(values, dtype=np.float32)) + quantized_t = values_t.to(_torch_src_dtype(src_fmt)) + x_bits = quantized_t.view(torch.uint8).numpy().copy() + x_decoded = quantized_t.float().numpy().astype(np.float32, copy=True) + return x_bits, x_decoded + + +def _cast_output(dst_fmt: str, values: np.ndarray) -> np.ndarray: + values = np.asarray(values, dtype=np.float32) + if dst_fmt == "f32": + return values.astype(np.float32, copy=True) + if dst_fmt == "f16": + return values.astype(np.float16).astype(np.float32) + if dst_fmt == "bf16": + return torch.from_numpy(values).to(torch.bfloat16).float().numpy().astype(np.float32, copy=True) + raise ValueError(f"unknown dst_fmt: {dst_fmt}") + + +def _effective_col_block_num(col_block_num: int) -> int: + return col_block_num + (col_block_num % 2) + + +def _loop_num2vf(row_block_num: int, col_block_num: int) -> int: + total_block_num = row_block_num * _effective_col_block_num(col_block_num) + return (total_block_num + _BLOCKS_PER_LOOP - 1) // _BLOCKS_PER_LOOP + + +def generate_case( + src_fmt: str, + dst_fmt: str, + *, + row_block_num: int = DEFAULT_ROW_BLOCK_NUM, + col_block_num: int = DEFAULT_COL_BLOCK_NUM, +) -> dict[str, object]: + """Generate one lightweight CPU reference case.""" + + if src_fmt not in SRC_FORMATS: + raise ValueError(f"unknown src_fmt: {src_fmt}") + if dst_fmt not in DTYPES: + raise ValueError(f"unknown dst_fmt: {dst_fmt}") + + effective_col_block_num = _effective_col_block_num(col_block_num) + total_block_num = row_block_num * effective_col_block_num + loop_num2vf = _loop_num2vf(row_block_num, col_block_num) + padded_block_num = loop_num2vf * _BLOCKS_PER_LOOP + + rng = np.random.default_rng(_SEED) + x_valid = rng.normal(0.0, 0.5, size=total_block_num * _BLOCK_SIZE).astype(np.float32) + x_padded = np.zeros(padded_block_num * _BLOCK_SIZE, dtype=np.float32) + x_padded[: x_valid.size] = x_valid + + scale_bits = np.tile(_E8M0_BYTES, (total_block_num + _E8M0_BYTES.size - 1) // _E8M0_BYTES.size)[:total_block_num] + scale_bits = scale_bits.astype(np.uint8, copy=True) + scale_values = _e8m0_to_f32(scale_bits) + scale_bits_padded = np.zeros(loop_num2vf * _RAW_SCALE_BYTES_PER_LOOP, dtype=np.uint8) + for i in range(loop_num2vf): + src_off = i * _BLOCKS_PER_LOOP + dst_off = i * _RAW_SCALE_BYTES_PER_LOOP + scale_bits_padded[dst_off : dst_off + _BLOCKS_PER_LOOP] = scale_bits[src_off : src_off + _BLOCKS_PER_LOOP] + + x_bits, x_quantized = _quantize_src_to_bits(src_fmt, x_padded) + block_scale = np.zeros(padded_block_num, dtype=np.float32) + block_scale[:total_block_num] = scale_values + y_f32 = ( + x_quantized.reshape(padded_block_num, _BLOCK_SIZE) + * block_scale.reshape(padded_block_num, 1) + ).reshape(-1) + y_expected = _cast_output(dst_fmt, y_f32) + + return { + "case_id": _case_id(src_fmt, dst_fmt), + "src_fmt": src_fmt, + "scale_fmt": SCALE_FORMAT, + "dst_fmt": dst_fmt, + "row_block_num": row_block_num, + "col_block_num": col_block_num, + "effective_col_block_num": effective_col_block_num, + "total_block_num": total_block_num, + "total_scale_num": total_block_num, + "loop_num2vf": loop_num2vf, + "x_bits": x_bits, + "x_f32": x_quantized, + "scale_bits": scale_bits_padded, + "scale_f32": scale_values, + "y_expected": y_expected, + "tolerance": TOLERANCE[dst_fmt], + "default_alias": src_fmt == "e4m3" and dst_fmt == "f32", + } + + +def generate_all( + *, + row_block_num: int = DEFAULT_ROW_BLOCK_NUM, + col_block_num: int = DEFAULT_COL_BLOCK_NUM, +) -> dict[str, dict[str, object]]: + """Generate the default dequant case matrix.""" + + return { + _case_id(src_fmt, dst_fmt): generate_case( + src_fmt, + dst_fmt, + row_block_num=row_block_num, + col_block_num=col_block_num, + ) + for src_fmt in SRC_FORMATS + for dst_fmt in DTYPES + } diff --git a/test/kernel-test/kernels/dequant/runtime.py b/test/kernel-test/kernels/dequant/runtime.py new file mode 100644 index 0000000000..1271290ef6 --- /dev/null +++ b/test/kernel-test/kernels/dequant/runtime.py @@ -0,0 +1,117 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""Shared runtime and artifact helpers for the dequant kernel.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +from kernel_test.npu_runtime import device_str, empty_npu + +if TYPE_CHECKING: + import torch + +ENTRY_SYMBOL = "anti_mx_quant_tail_axis_compute_data_probe" +DEFAULT_NAMED_STEM = "dequant" + + +@dataclass(frozen=True) +class DequantCompileArgs: + """Prepared compile inputs shared by the dequant backends.""" + + src_fmt: str + scale_fmt: str + dst_fmt: str + row_block_num: int + col_block_num: int + loop_num2vf: int + entry_symbol: str + case_dir_name: str + named_stem: str + default_alias: bool + + +@dataclass(frozen=True) +class DequantLaunchArgs: + """Prepared runtime tensors and metadata for one dequant launch.""" + + src_fmt: str + dst_fmt: str + row_block_num: int + col_block_num: int + loop_num2vf: int + x: "torch.Tensor" + scale: "torch.Tensor" + y: "torch.Tensor" + + +def case_dir_name(case: dict[str, object]) -> str: + return ( + f"fp8_{case['src_fmt']}_scale_{case['scale_fmt']}_out_{case['dst_fmt']}" + f"_rb{case['row_block_num']}_cb{case['col_block_num']}" + ) + + +def prepare_compile_args(case: dict[str, object]) -> DequantCompileArgs: + """Normalize one dequant case dict into stable compile metadata.""" + + return DequantCompileArgs( + src_fmt=str(case["src_fmt"]), + scale_fmt=str(case["scale_fmt"]), + dst_fmt=str(case["dst_fmt"]), + row_block_num=int(case["row_block_num"]), + col_block_num=int(case["col_block_num"]), + loop_num2vf=int(case.get("loop_num2vf", 1)), + entry_symbol=str(case.get("entry_symbol", ENTRY_SYMBOL)), + case_dir_name=case_dir_name(case), + named_stem=str(case.get("named_stem", DEFAULT_NAMED_STEM)), + default_alias=bool(case.get("default_alias", False)), + ) + + +def artifact_case_dir(root: Path, case: dict[str, object]) -> Path: + """Return the per-case artifact directory.""" + + return root / case_dir_name(case) + + +def torch_dtype(dst_fmt: str) -> torch.dtype: + import torch + + if dst_fmt == "f32": + return torch.float32 + if dst_fmt == "bf16": + return torch.bfloat16 + if dst_fmt == "f16": + return torch.float16 + raise ValueError(f"unknown dst_fmt: {dst_fmt}") + + +def prepare_launch_args(case: dict[str, object]) -> DequantLaunchArgs: + """Convert one dequant case into device tensors and launch metadata.""" + + import torch + + dev = device_str() + dst_dtype = torch_dtype(str(case["dst_fmt"])) + x = torch.from_numpy(case["x_bits"]).to(torch.uint8).to(dev) + scale = torch.from_numpy(case["scale_bits"]).to(torch.uint8).to(dev) + y = empty_npu(case["y_expected"].shape, dst_dtype) + return DequantLaunchArgs( + src_fmt=str(case["src_fmt"]), + dst_fmt=str(case["dst_fmt"]), + row_block_num=int(case["row_block_num"]), + col_block_num=int(case["col_block_num"]), + loop_num2vf=int(case.get("loop_num2vf", 1)), + x=x, + scale=scale, + y=y, + ) diff --git a/test/kernel-test/kernels/dequant/spec.py b/test/kernel-test/kernels/dequant/spec.py new file mode 100644 index 0000000000..6df613377a --- /dev/null +++ b/test/kernel-test/kernels/dequant/spec.py @@ -0,0 +1,93 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""Case listing and verification for dequant.""" + +from __future__ import annotations + +from kernel_test.results import CaseResult + +from .runtime import ENTRY_SYMBOL + + +def list_cases(workflow: str) -> dict[str, object]: + """Return the dequant case matrix for the requested workflow.""" + + if workflow == "cycle": + return {} + if workflow != "correctness": + raise ValueError(f"unsupported dequant workflow: {workflow}") + + try: + from .reference import generate_all + except ModuleNotFoundError as exc: + if exc.name not in {"numpy", "torch"}: + raise + cases = { + f"{src_fmt}_{dst_fmt}": { + "src_fmt": src_fmt, + "scale_fmt": "e8m0", + "dst_fmt": dst_fmt, + "row_block_num": 4, + "col_block_num": 4, + "loop_num2vf": 1, + "default_alias": src_fmt == "e4m3" and dst_fmt == "f32", + } + for src_fmt in ("e4m3", "e5m2") + for dst_fmt in ("f32", "bf16", "f16") + } + else: + cases = generate_all() + + for case in cases.values(): + case["entry_symbol"] = ENTRY_SYMBOL + return cases + + +def verify_case(case_id: str, case: object, output: object) -> CaseResult: + """Verify one dequant case against the CPU golden.""" + + import numpy as np + + y_host = output["y"].cpu() + if case["dst_fmt"] == "f32": + got = y_host.numpy().astype(np.float32) + else: + got = y_host.float().numpy().astype(np.float32) + expected = np.asarray(case["y_expected"], dtype=np.float32) + max_diff = float(np.max(np.abs(got - expected))) + tol = float(case.get("tolerance", 1e-3)) + if max_diff > tol: + return CaseResult( + ok=False, + message=( + f"maxDiff={max_diff:.6g} exceeds tol={tol:.6g} " + f"for {case['src_fmt']}+{case['scale_fmt']}->{case['dst_fmt']}" + ), + ) + + return CaseResult( + ok=True, + message=( + f"{case['src_fmt']}+{case['scale_fmt']}->{case['dst_fmt']}: " + f"maxDiff={max_diff:.6g}" + ), + ) + + +def cycle_fields(case_id: str, case: object, backend: object) -> dict[str, object]: + """Build stable case fields for future cycle support.""" + + del case_id, backend + return { + "src": case["src_fmt"], + "scale": case["scale_fmt"], + "dst": case["dst_fmt"], + "rb": case["row_block_num"], + "cb": case["col_block_num"], + } diff --git a/test/kernel-test/kernels/dequant/vmi/README.md b/test/kernel-test/kernels/dequant/vmi/README.md new file mode 100644 index 0000000000..56e6c9a6ef --- /dev/null +++ b/test/kernel-test/kernels/dequant/vmi/README.md @@ -0,0 +1,116 @@ + + +# dequant VMI backend + +This directory keeps the PTODSL VMI rewrite for the `dequant` kernel test. + +## Scope + +- source kernel: `test/kernel-test/kernels/dequant/anti_mx_quant_tail_axis.h` +- current VMI rewrite: `anti_mx_quant_tail_axis_vmi.py` +- currently covered path: FP8 input + FP32 scale buffer + `f32` / `bf16` / `f16` output +- the FP4 + BF16-scale path in the source kernel is not modeled here yet + +## Artifact layout + +Running the `backend=vmi` correctness path, or `kernel-test/run.py --emit-mlir`, +emits artifacts under: + +```text +test/kernel-test/kernels/dequant/generated/ +``` + +Per-case outputs live in a case directory, and the backend also refreshes a few +root-level aliases for quick manual inspection: + +- `vmi.pto`: latest lowered source VMI IR +- `mi.pto`: latest lowered MI IR + +This matches the current manual-debug habit: first open the stable root alias, +and only jump into the per-case directory when a specific specialization needs +to be checked. + +## Current VMI modeling + +The working FP8 path is intentionally written in surface VMI terms instead of +trying to mirror every source-side physical shuffle: + +- `x` is modeled as two dense 256-lane FP8 `vload`s +- each half is widened by `pto.vmi.vcvt(..., pto.f32)` +- `ComputeScale` writes a compact converted scale buffer: one scale value per + 32 input elements +- `ComputeData` consumes that compact buffer with `vload(dist_mode="brc", + group=8)`, matching the source `DIST_E2B_B32` scale-load contract +- dequant itself is just `vmul` +- output uses `vstore`, with an extra `vcvt` only for `bf16` / `f16` + +This version is functionally correct and avoids the earlier invalid formulation +that mixed deinterleaved VMI data with a surface `vintlv`. + +## Static efficiency comparison + +The notes below compare the generated MI with the original handwritten ASC/CCE +for the FP8 path. This is a static code-shape comparison, not a profiler result. + +### Overall conclusion + +The current lowered MI is likely slower than the original ASC implementation. + +- `f32` output: probably somewhat slower +- `bf16` / `f16` output: likely more noticeably slower + +The main reason is not the math itself, but that the generated MI still needs +to materialize more explicit layout transforms around load/broadcast/store. + +### Where the original ASC is tighter + +The original source uses specialized distribution modes that already match the +intended data choreography fairly well: + +- `DIST_DINTLV_B8` for the FP8 payload load +- `DIST_E2B_B32` for scale expansion +- direct store-side layout/stype handling for `f32` +- relatively compact cast-pack-store handling for `bf16` / `f16` + +So a lot of the "how to arrange lanes" work is implicit in the source +instruction choice itself. + +### What the current lowered MI does + +For the `f32` path, the generated MI roughly expands into: + +- 2 dense `vlds` for FP8 input +- 8 `vcvt` ops (`P0`..`P3` for each 256-lane half) +- compact scale conversion stores +- 2 direct `E2B_B32` scale loads for group-broadcast scale consumption +- 8 `vmul` +- multiple `vintlv` plus `pintlv_b32` +- 8 `vsts` + +Compared with the source ASC, the extra cost is mainly on store-side re-layout +before `vsts`. + +For the `bf16` / `f16` path, the gap is usually larger because lowering still +has to do more explicit rearrangement before the final packed 16-bit store. + +## Interpretation + +The current result should be treated as: + +- semantically correct VMI surface modeling +- good for validating VMI expression and lowering correctness +- not yet the most efficient MI shape PTOAS could theoretically emit + +The biggest optimization opportunities are likely: + +- reduce redundant store-side layout materialization for `bf16` / `f16` +- improve lowering so the final MI better preserves the compactness of the + source distribution-based implementation diff --git a/test/kernel-test/kernels/dequant/vmi/anti_mx_quant_tail_axis_vmi.py b/test/kernel-test/kernels/dequant/vmi/anti_mx_quant_tail_axis_vmi.py new file mode 100644 index 0000000000..3c059f7ca5 --- /dev/null +++ b/test/kernel-test/kernels/dequant/vmi/anti_mx_quant_tail_axis_vmi.py @@ -0,0 +1,373 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""VMI rewrite of AntiMxQuantTailAxis scale/dequant helpers. + +Source regions: + test/kernel-test/kernels/dequant/anti_mx_quant_tail_axis.h:313 + test/kernel-test/kernels/dequant/anti_mx_quant_tail_axis.h:380 + +Current modeling notes: + 1. ``ComputeScale`` is expressed in terms of the logical e8m0 scale value + rather than replaying the source's temporary zero-fill + interleave + choreography. For the fp8 path we materialize `scaleBuffer` as one + contiguous 16-lane compact FP32 scale vector per loop: the 16 raw + block-scale bytes are converted to FP32 and kept compact in UB. + 2. The source does ``DIST_DINTLV_B8`` and then immediately interleaves the + two FP8 registers before FP8-to-FP32 part casts. Logically that is + equivalent to two contiguous 256-lane FP8 halves, so the VMI rewrite + models this as two dense loads followed by two `vcvt`s. + 3. ``scaleBufAddr`` therefore holds one contiguous 16-lane compact scale + vector per loop, and ``ComputeData`` reloads each 8-scale half with a + grouped ``vload`` before grouped ``vbrc`` expansion and multiply. + 4. The runtime `dequant_vmi_*` kernels now call `ComputeScale` and + `ComputeData` explicitly, matching the source `Compute(...)` structure. +""" + +from ptodsl import pto + +_ELEMS_PER_LOOP = 512 +_HALF_ELEMS_PER_LOOP = _ELEMS_PER_LOOP // 2 +_BLOCK_SIZE = 32 +_BLOCKS_PER_DATA_LOOP = _ELEMS_PER_LOOP // _BLOCK_SIZE +_SCALE_LANES_PER_LOOP = _BLOCKS_PER_DATA_LOOP +_COMPACT_SCALE_LANES_PER_LOOP = _SCALE_LANES_PER_LOOP +_SCALE_LANES_PER_HALF = _SCALE_LANES_PER_LOOP // 2 +_SCALE_REPEAT_GROUPS = _SCALE_LANES_PER_HALF +_RAW_SCALE_BYTES_PER_LOOP = _BLOCKS_PER_DATA_LOOP + +_X_BASE_ADDR = 0 +_RAW_SCALE_BASE_ADDR = 0x4000 +_SCALE_BASE_ADDR = 0x8000 +_Y_BASE_ADDR = 0x10000 + + +def _compute_scale_ub( + scale_local_addr: pto.i64, + scale_buf_addr: pto.i64, + *, + DST_FMT: pto.const_expr = "f32", + LOOP_NUM2VF: pto.const_expr = 1, +): + scale_src_ptr = pto.castptr(scale_local_addr, pto.ptr(pto.ui8, "ub")) + + if pto.const_expr(DST_FMT == "f32"): + scale_dst_ptr = pto.castptr(scale_buf_addr, pto.ptr(pto.f32, "ub")) + dst_dtype = pto.f32 + elif pto.const_expr(DST_FMT == "bf16"): + scale_dst_ptr = pto.castptr(scale_buf_addr, pto.ptr(pto.bf16, "ub")) + dst_dtype = pto.bf16 + else: + raise ValueError(f"unsupported ComputeScale DST_FMT specialization: {DST_FMT}") + + shift = pto.vmi.vbrc(pto.i32(23), size=_BLOCKS_PER_DATA_LOOP) + scale_mask = pto.vmi.create_mask(_SCALE_LANES_PER_LOOP, size=_SCALE_LANES_PER_LOOP) + + for i in range(LOOP_NUM2VF): + raw_scale_off = i * _RAW_SCALE_BYTES_PER_LOOP + dst_off = i * _COMPACT_SCALE_LANES_PER_LOOP + + raw_scale = pto.vmi.vload( + scale_src_ptr, + raw_scale_off, + size=_BLOCKS_PER_DATA_LOOP, + ) + scale_u32 = pto.vmi.vcvt(raw_scale, pto.ui32) + scale_i32 = pto.vmi.vinterpret_cast( + scale_u32, + pto.i32, + ) + scale_bits = pto.vmi.vshl(scale_i32, shift) + scale_compact = pto.vmi.vinterpret_cast( + scale_bits, + pto.f32, + ) + if pto.const_expr(DST_FMT == "f32"): + pto.vmi.vstore(scale_compact, scale_dst_ptr, dst_off, scale_mask) + else: + pto.vmi.vstore(pto.vmi.vcvt(scale_compact, dst_dtype), scale_dst_ptr, dst_off, scale_mask) + + +def _compute_data_ub( + x_local_addr: pto.i64, + scale_buf_addr: pto.i64, + y_local_addr: pto.i64, + *, + SRC_FMT: pto.const_expr = "e4m3", + DST_FMT: pto.const_expr = "f32", + LOOP_NUM2VF: pto.const_expr = 1, +): + if pto.const_expr(SRC_FMT == "e4m3"): + x_ptr = pto.castptr(x_local_addr, pto.ptr(pto.f8e4m3, "ub")) + elif pto.const_expr(SRC_FMT == "e5m2"): + x_ptr = pto.castptr(x_local_addr, pto.ptr(pto.f8e5m2, "ub")) + else: + raise ValueError(f"unsupported SRC_FMT specialization: {SRC_FMT}") + + scale_ptr = pto.castptr(scale_buf_addr, pto.ptr(pto.f32, "ub")) + + if pto.const_expr(DST_FMT == "f32"): + y_ptr = pto.castptr(y_local_addr, pto.ptr(pto.f32, "ub")) + dst_dtype = pto.f32 + elif pto.const_expr(DST_FMT == "bf16"): + y_ptr = pto.castptr(y_local_addr, pto.ptr(pto.bf16, "ub")) + dst_dtype = pto.bf16 + elif pto.const_expr(DST_FMT == "f16"): + y_ptr = pto.castptr(y_local_addr, pto.ptr(pto.f16, "ub")) + dst_dtype = pto.f16 + else: + raise ValueError(f"unsupported DST_FMT specialization: {DST_FMT}") + + mask256 = pto.vmi.create_mask(_HALF_ELEMS_PER_LOOP, size=_HALF_ELEMS_PER_LOOP) + + for i in range(LOOP_NUM2VF): + x_off = i * _ELEMS_PER_LOOP + scale_off = i * _COMPACT_SCALE_LANES_PER_LOOP + y_off = i * _ELEMS_PER_LOOP + + x_lo_f8 = pto.vmi.vload(x_ptr, x_off, size=_HALF_ELEMS_PER_LOOP) + x_hi_f8 = pto.vmi.vload(x_ptr, x_off + _HALF_ELEMS_PER_LOOP, size=_HALF_ELEMS_PER_LOOP) + x_lo_f32 = pto.vmi.vcvt(x_lo_f8, pto.f32) + x_hi_f32 = pto.vmi.vcvt(x_hi_f8, pto.f32) + + scale_lo_slots = pto.vmi.vload( + scale_ptr, + scale_off, + size=_SCALE_LANES_PER_HALF, + stride=1, + group=_SCALE_REPEAT_GROUPS, + ) + scale_hi_slots = pto.vmi.vload( + scale_ptr, + scale_off + _SCALE_LANES_PER_HALF, + size=_SCALE_LANES_PER_HALF, + stride=1, + group=_SCALE_REPEAT_GROUPS, + ) + scale_lo_f32 = pto.vmi.vbrc( + scale_lo_slots, + size=_HALF_ELEMS_PER_LOOP, + group=_SCALE_REPEAT_GROUPS, + ) + scale_hi_f32 = pto.vmi.vbrc( + scale_hi_slots, + size=_HALF_ELEMS_PER_LOOP, + group=_SCALE_REPEAT_GROUPS, + ) + y_lo_f32 = pto.vmi.vmul(x_lo_f32, scale_lo_f32, mask256) + y_hi_f32 = pto.vmi.vmul(x_hi_f32, scale_hi_f32, mask256) + + if pto.const_expr(DST_FMT == "f32"): + pto.vmi.vstore(y_lo_f32, y_ptr, y_off, mask256) + pto.vmi.vstore(y_hi_f32, y_ptr, y_off + _HALF_ELEMS_PER_LOOP, mask256) + else: + pto.vmi.vstore( + pto.vmi.vcvt(y_lo_f32, dst_dtype), + y_ptr, + y_off, + mask256, + ) + pto.vmi.vstore( + pto.vmi.vcvt(y_hi_f32, dst_dtype), + y_ptr, + y_off + _HALF_ELEMS_PER_LOOP, + mask256, + ) + + +def _runtime_entry( + x_gm: pto.ptr(pto.ui8, "gm"), + scale_gm: pto.ptr(pto.ui8, "gm"), + y_gm, + *, + SRC_FMT: pto.const_expr = "e4m3", + DST_FMT: pto.const_expr = "f32", + ROW_BLOCK_NUM: pto.const_expr = 4, + COL_BLOCK_NUM: pto.const_expr = 4, +): + effective_col_block_num = COL_BLOCK_NUM + (COL_BLOCK_NUM % 2) + total_scale_num = ROW_BLOCK_NUM * effective_col_block_num + total_block_num = total_scale_num + scale_loop_num = (total_scale_num + _BLOCKS_PER_DATA_LOOP - 1) // _BLOCKS_PER_DATA_LOOP + loop_num2vf = (total_block_num + _BLOCKS_PER_DATA_LOOP - 1) // _BLOCKS_PER_DATA_LOOP + padded_total_block_num = loop_num2vf * _BLOCKS_PER_DATA_LOOP + total_elems = padded_total_block_num * _BLOCK_SIZE + + x_ub_ptr = pto.castptr(pto.const(_X_BASE_ADDR, dtype=pto.ui64), pto.ptr(pto.ui8, "ub")) + scale_ub_ptr = pto.castptr(pto.const(_RAW_SCALE_BASE_ADDR, dtype=pto.ui64), pto.ptr(pto.ui8, "ub")) + x_bytes = total_elems + scale_bytes = scale_loop_num * _RAW_SCALE_BYTES_PER_LOOP + + if pto.const_expr(DST_FMT == "f32"): + y_ub_ptr = pto.castptr(pto.const(_Y_BASE_ADDR, dtype=pto.ui64), pto.ptr(pto.f32, "ub")) + y_bytes = total_elems * 4 + else: + y_ub_ptr = pto.castptr( + pto.const(_Y_BASE_ADDR, dtype=pto.ui64), + pto.ptr(pto.bf16 if pto.const_expr(DST_FMT == "bf16") else pto.f16, "ub"), + ) + y_bytes = total_elems * 2 + + pto.mte_gm_ub(x_gm, x_ub_ptr, 0, x_bytes, nburst=(1, x_bytes, x_bytes)) + pto.mte_gm_ub(scale_gm, scale_ub_ptr, 0, scale_bytes, nburst=(1, scale_bytes, scale_bytes)) + + pto.set_flag(pto.Pipe.MTE2, pto.Pipe.V, event_id=0) + pto.wait_flag(pto.Pipe.MTE2, pto.Pipe.V, event_id=0) + _compute_scale_ub( + pto.const(_RAW_SCALE_BASE_ADDR, dtype=pto.i64), + pto.const(_SCALE_BASE_ADDR, dtype=pto.i64), + DST_FMT="f32", + LOOP_NUM2VF=scale_loop_num, + ) + _compute_data_ub( + pto.const(_X_BASE_ADDR, dtype=pto.i64), + pto.const(_SCALE_BASE_ADDR, dtype=pto.i64), + pto.const(_Y_BASE_ADDR, dtype=pto.i64), + SRC_FMT=SRC_FMT, + DST_FMT=DST_FMT, + LOOP_NUM2VF=loop_num2vf, + ) + + pto.set_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=0) + pto.wait_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=0) + pto.mte_ub_gm(y_ub_ptr, y_gm, y_bytes, nburst=(1, y_bytes, y_bytes)) + pto.pipe_barrier(pto.Pipe.ALL) + + +@pto.jit( + name="anti_mx_quant_tail_axis_compute_scale_probe", + target="a5", + backend="vpto", + mode="explicit", + kernel_kind="vector", + insert_sync=False, +) +def anti_mx_quant_tail_axis_compute_scale_probe( + scale_local_addr: pto.i64 = _X_BASE_ADDR, + scale_buf_addr: pto.i64 = _SCALE_BASE_ADDR, + *, + DST_FMT: pto.const_expr = "f32", + LOOP_NUM2VF: pto.const_expr = 1, +): + _compute_scale_ub( + scale_local_addr, + scale_buf_addr, + DST_FMT=DST_FMT, + LOOP_NUM2VF=LOOP_NUM2VF, + ) + + +@pto.jit( + name="anti_mx_quant_tail_axis_compute_data_probe", + target="a5", + backend="vpto", + mode="explicit", + kernel_kind="vector", + insert_sync=False, +) +def anti_mx_quant_tail_axis_compute_data_probe( + x_local_addr: pto.i64 = _X_BASE_ADDR, + scale_buf_addr: pto.i64 = _SCALE_BASE_ADDR, + y_local_addr: pto.i64 = _Y_BASE_ADDR, + *, + SRC_FMT: pto.const_expr = "e4m3", + DST_FMT: pto.const_expr = "f32", + LOOP_NUM2VF: pto.const_expr = 1, +): + _compute_data_ub( + x_local_addr, + scale_buf_addr, + y_local_addr, + SRC_FMT=SRC_FMT, + DST_FMT=DST_FMT, + LOOP_NUM2VF=LOOP_NUM2VF, + ) + + +@pto.jit( + name="dequant_vmi_f32", + target="a5", + backend="vpto", + mode="explicit", + kernel_kind="vector", + insert_sync=False, +) +def dequant_vmi_f32( + x_gm: pto.ptr(pto.ui8, "gm"), + scale_gm: pto.ptr(pto.ui8, "gm"), + y_gm: pto.ptr(pto.f32, "gm"), + *, + SRC_FMT: pto.const_expr = "e4m3", + ROW_BLOCK_NUM: pto.const_expr = 4, + COL_BLOCK_NUM: pto.const_expr = 4, +): + _runtime_entry( + x_gm, + scale_gm, + y_gm, + SRC_FMT=SRC_FMT, + DST_FMT="f32", + ROW_BLOCK_NUM=ROW_BLOCK_NUM, + COL_BLOCK_NUM=COL_BLOCK_NUM, + ) + + +@pto.jit( + name="dequant_vmi_bf16", + target="a5", + backend="vpto", + mode="explicit", + kernel_kind="vector", + insert_sync=False, +) +def dequant_vmi_bf16( + x_gm: pto.ptr(pto.ui8, "gm"), + scale_gm: pto.ptr(pto.ui8, "gm"), + y_gm: pto.ptr(pto.bf16, "gm"), + *, + SRC_FMT: pto.const_expr = "e4m3", + ROW_BLOCK_NUM: pto.const_expr = 4, + COL_BLOCK_NUM: pto.const_expr = 4, +): + _runtime_entry( + x_gm, + scale_gm, + y_gm, + SRC_FMT=SRC_FMT, + DST_FMT="bf16", + ROW_BLOCK_NUM=ROW_BLOCK_NUM, + COL_BLOCK_NUM=COL_BLOCK_NUM, + ) + + +@pto.jit( + name="dequant_vmi_f16", + target="a5", + backend="vpto", + mode="explicit", + kernel_kind="vector", + insert_sync=False, +) +def dequant_vmi_f16( + x_gm: pto.ptr(pto.ui8, "gm"), + scale_gm: pto.ptr(pto.ui8, "gm"), + y_gm: pto.ptr(pto.f16, "gm"), + *, + SRC_FMT: pto.const_expr = "e4m3", + ROW_BLOCK_NUM: pto.const_expr = 4, + COL_BLOCK_NUM: pto.const_expr = 4, +): + _runtime_entry( + x_gm, + scale_gm, + y_gm, + SRC_FMT=SRC_FMT, + DST_FMT="f16", + ROW_BLOCK_NUM=ROW_BLOCK_NUM, + COL_BLOCK_NUM=COL_BLOCK_NUM, + ) diff --git a/test/kernel-test/kernels/dequant/vmi/backend.py b/test/kernel-test/kernels/dequant/vmi/backend.py new file mode 100644 index 0000000000..a10e7326fa --- /dev/null +++ b/test/kernel-test/kernels/dequant/vmi/backend.py @@ -0,0 +1,119 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""Runtime VMI backend for dequant.""" + +from __future__ import annotations + +import os +from pathlib import Path + +from kernel_test.backends import ArtifactPlan, RunPurpose +from kernel_test.npu_runtime import ensure_runtime, stream_ptr, sync +from kernel_test.pto_artifacts import materialize_artifact_plan + +from .anti_mx_quant_tail_axis_vmi import dequant_vmi_bf16, dequant_vmi_f16, dequant_vmi_f32 +from ..runtime import artifact_case_dir, prepare_compile_args, prepare_launch_args + +_VMI_ROOT = Path(__file__).resolve().parent +_KERNEL_ROOT = _VMI_ROOT.parent +_GENERATED_DIR = _KERNEL_ROOT / "generated" +_COMPILED: dict[tuple[str, str, int, int], object] = {} + + +def _kernel_for_dst(dst_fmt: str): + if dst_fmt == "f32": + return dequant_vmi_f32 + if dst_fmt == "bf16": + return dequant_vmi_bf16 + if dst_fmt == "f16": + return dequant_vmi_f16 + raise ValueError(f"unsupported dst_fmt: {dst_fmt}") + + +def _prepare_runtime_kernel(case: dict) -> object: + compile_args = prepare_compile_args(case) + cache_key = ( + compile_args.src_fmt, + compile_args.dst_fmt, + compile_args.row_block_num, + compile_args.col_block_num, + ) + compiled = _COMPILED.get(cache_key) + if compiled is None: + compiled = _kernel_for_dst(compile_args.dst_fmt).compile( + SRC_FMT=compile_args.src_fmt, + ROW_BLOCK_NUM=compile_args.row_block_num, + COL_BLOCK_NUM=compile_args.col_block_num, + ) + _COMPILED[cache_key] = compiled + return compiled + + +def _build_artifact_plan(case: dict[str, object]) -> ArtifactPlan: + compile_args = prepare_compile_args(case) + compiled = _prepare_runtime_kernel(case) + + case_dir = artifact_case_dir(_GENERATED_DIR, case) + return ArtifactPlan( + generated_dir=_GENERATED_DIR, + case_dir=case_dir, + vmi_text=compiled.mlir_text(), + alias_stem=compile_args.named_stem, + ) + + +class AntiMxQuantTailAxisVmiBackend: + """Runtime VMI backend for dequant.""" + + name = "vmi" + + def is_supported(self, case: object, *, purpose: RunPurpose) -> tuple[bool, str | None]: + if purpose == "cycle": + return False, "backend=vmi has correctness launch only; no cycle probe yet" + supported = case["src_fmt"] in {"e4m3", "e5m2"} and case["dst_fmt"] in {"f32", "bf16", "f16"} + if supported: + return True, None + return False, "backend=vmi not wired for this case" + + def launch(self, case: object, *, purpose: RunPurpose) -> object: + if purpose != "correctness": + raise ValueError(f"unsupported purpose for dequant vmi backend: {purpose}") + + ensure_runtime("dequant") + launch_args = prepare_launch_args(case) + compiled = _prepare_runtime_kernel(case) + artifacts = materialize_artifact_plan( + "correctness", + _build_artifact_plan(case), + root_alias=prepare_compile_args(case).default_alias, + ) + + compiled[1, stream_ptr()]( + launch_args.x.data_ptr(), + launch_args.scale.data_ptr(), + launch_args.y.data_ptr(), + ) + sync() + + return { + "y": launch_args.y, + **dict(artifacts.paths), + } + + def cache_tag(self) -> str: + backend_py = _VMI_ROOT / "backend.py" + kernel_py = _VMI_ROOT / "anti_mx_quant_tail_axis_vmi.py" + return ( + f"vmi:{backend_py}:{os.path.getmtime(backend_py):.0f}:" + f"{kernel_py}:{os.path.getmtime(kernel_py):.0f}" + ) + + def build_artifact_plan(self, case_id: str, case: object) -> ArtifactPlan: + del case_id + return _build_artifact_plan(case) diff --git a/test/kernel-test/kernels/rope/__init__.py b/test/kernel-test/kernels/rope/__init__.py new file mode 100644 index 0000000000..afec404a31 --- /dev/null +++ b/test/kernel-test/kernels/rope/__init__.py @@ -0,0 +1,31 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""Rope kernel adapter for the kernel-test framework.""" + +from __future__ import annotations + +from kernel_test.registry import OperatorSpec, make_operator_spec + +from .backends import create_backend +from .spec import cycle_fields, list_cases, verify_case + + +def get_operator_spec() -> OperatorSpec: + """Return the rope operator registration for the shared registry.""" + + return make_operator_spec( + name="rope", + default_backend="cce", + backend_names=("cce", "mi", "vmi"), + create_backend=create_backend, + list_cases=list_cases, + verify=verify_case, + cycle_fields=cycle_fields, + summary="VF sim rope kernel adapter", + ) diff --git a/test/kernel-test/kernels/rope/backends.py b/test/kernel-test/kernels/rope/backends.py new file mode 100644 index 0000000000..dac62a79d4 --- /dev/null +++ b/test/kernel-test/kernels/rope/backends.py @@ -0,0 +1,49 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""Backend factory for the rope kernel.""" + +from __future__ import annotations + +from kernel_test.backends import BackendAdapter + + +def create_backend(name: str) -> BackendAdapter: + """Create one rope backend adapter from the local backend packages.""" + + if name == "cce": + try: + from .cce.backend import RopeCceBackend + except ModuleNotFoundError as exc: + raise RuntimeError( + "rope cce backend requires the rope runtime dependencies; use the " + "`pto` conda environment or install numpy/torch/torch_npu first" + ) from exc + + return RopeCceBackend() + if name == "vmi": + try: + from .vmi.backend import RopeVmiBackend + except ModuleNotFoundError as exc: + raise RuntimeError( + "rope vmi backend requires the rope runtime dependencies; use the " + "`pto` conda environment or install ptodsl/torch/torch_npu first" + ) from exc + + return RopeVmiBackend() + if name == "mi": + try: + from .mi.backend import RopeMiBackend + except ModuleNotFoundError as exc: + raise RuntimeError( + "rope mi backend requires the rope runtime dependencies; use the " + "`pto` conda environment or install ptodsl/torch/torch_npu first" + ) from exc + + return RopeMiBackend() + raise ValueError(f"unknown rope backend: {name}") diff --git a/test/kernel-test/kernels/rope/cce/CMakeLists.txt b/test/kernel-test/kernels/rope/cce/CMakeLists.txt new file mode 100644 index 0000000000..80fda67d83 --- /dev/null +++ b/test/kernel-test/kernels/rope/cce/CMakeLists.txt @@ -0,0 +1,82 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +cmake_minimum_required(VERSION 3.20) + +if(NOT DEFINED ASCEND_HOME_PATH) + if(DEFINED ENV{ASCEND_HOME_PATH}) + set(ASCEND_HOME_PATH $ENV{ASCEND_HOME_PATH}) + elseif(DEFINED ENV{ASCEND_TOOLKIT_HOME}) + set(ASCEND_HOME_PATH $ENV{ASCEND_TOOLKIT_HOME}) + else() + message(FATAL_ERROR "ASCEND_HOME_PATH is not set") + endif() +endif() + +if(NOT DEFINED ASCEND_DRIVER_PATH) + if(DEFINED ENV{ASCEND_DRIVER_PATH}) + set(ASCEND_DRIVER_PATH $ENV{ASCEND_DRIVER_PATH}) + else() + set(ASCEND_DRIVER_PATH "/usr/local/Ascend/driver") + endif() +endif() + +find_program(BISHENG_EXECUTABLE bisheng HINTS "${ASCEND_HOME_PATH}/bin") +if(NOT BISHENG_EXECUTABLE) + message(FATAL_ERROR "bisheng compiler not found") +endif() + +set(CMAKE_CXX_COMPILER "${BISHENG_EXECUTABLE}" CACHE FILEPATH "" FORCE) + +project(rope_cce LANGUAGES CXX) + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_POSITION_INDEPENDENT_CODE ON) + +add_library(rope_cce SHARED rope_cce_kernel.cpp) + +target_include_directories( + rope_cce + PRIVATE + "${CMAKE_CURRENT_SOURCE_DIR}" + "${ASCEND_HOME_PATH}/include" + "${ASCEND_HOME_PATH}/include/cce" + "${ASCEND_HOME_PATH}/pkg_inc" + "${ASCEND_HOME_PATH}/pkg_inc/profiling" + "${ASCEND_HOME_PATH}/pkg_inc/runtime/runtime" + "${ASCEND_DRIVER_PATH}/kernel_inc" +) + +target_compile_options( + rope_cce + PRIVATE + -O3 + -std=gnu++17 + -Wno-macro-redefined + -Wno-ignored-attributes + -Wno-unknown-attributes + -fPIC + -xcce + -Xhost-start + -Xhost-end + "SHELL:-mllvm -cce-aicore-stack-size=0x8000" + "SHELL:-mllvm -cce-aicore-function-stack-size=0x8000" + "SHELL:-mllvm -cce-aicore-record-overflow=true" + "SHELL:-mllvm -cce-aicore-addr-transform" + "SHELL:-mllvm -cce-aicore-dcci-insert-for-scalar=false" + --cce-aicore-arch=dav-c310-vec + --cce-simd-vf-fusion=false + -DREGISTER_BASE +) + +target_link_options( + rope_cce + PRIVATE + --cce-fatobj-link + -Wl,-soname,librope_cce.so +) diff --git a/test/kernel-test/kernels/rope/cce/backend.py b/test/kernel-test/kernels/rope/cce/backend.py new file mode 100644 index 0000000000..d72724a348 --- /dev/null +++ b/test/kernel-test/kernels/rope/cce/backend.py @@ -0,0 +1,140 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""CCE backend for the rope kernel.""" + +from __future__ import annotations + +import ctypes +import os +import subprocess +from pathlib import Path + +from kernel_test.backends import RunPurpose +from kernel_test.npu_runtime import ensure_runtime, stream_ptr, sync + +from ..runtime import RopeLaunchArgs, prepare_launch_args +from ..tile_config import DTYPES, MODES, sim_fn_name + +_CCE_ROOT = Path(__file__).resolve().parent +_BUILD_DIR = _CCE_ROOT / "build" +_KERNEL_SOURCE = _CCE_ROOT / "rope_cce_kernel.cpp" +_LIB_PATH = _BUILD_DIR / "librope_cce.so" +_CMAKE_FILE = _CCE_ROOT / "CMakeLists.txt" +_SIM_SYMBOLS: tuple[str, ...] = tuple( + sim_fn_name(mode, dtype, cycle=cycle) + for cycle in (False, True) + for dtype in DTYPES + for mode in MODES +) +_LIB: ctypes.CDLL | None = None + + +def _ascend_home() -> Path: + home = os.environ.get("ASCEND_HOME_PATH") or os.environ.get("ASCEND_TOOLKIT_HOME") + if not home: + raise EnvironmentError("ASCEND_HOME_PATH is not set. Source CANN setenv.bash first.") + return Path(home) + + +def _run(cmd: list[str], cwd: Path) -> None: + subprocess.run(cmd, cwd=cwd, check=True) + + +def _build_lib(force: bool = False) -> Path: + _BUILD_DIR.mkdir(parents=True, exist_ok=True) + if _LIB_PATH.is_file() and not force: + return _LIB_PATH + + ascend = _ascend_home() + driver = os.environ.get("ASCEND_DRIVER_PATH", "/usr/local/Ascend/driver") + _run( + [ + "cmake", + "-S", + str(_CCE_ROOT), + "-B", + str(_BUILD_DIR), + f"-DASCEND_HOME_PATH={ascend}", + f"-DASCEND_DRIVER_PATH={driver}", + ], + cwd=_CCE_ROOT, + ) + _run(["cmake", "--build", str(_BUILD_DIR), "--target", "rope_cce"], cwd=_CCE_ROOT) + return _LIB_PATH + + +def _bind_lib(lib: ctypes.CDLL) -> None: + argtypes = [ctypes.c_void_p] * 6 + for name in _SIM_SYMBOLS: + fn = getattr(lib, name) + fn.argtypes = argtypes + fn.restype = None + + +def _load_lib() -> ctypes.CDLL: + global _LIB + if _LIB is None: + _LIB = ctypes.CDLL(str(_build_lib())) + _bind_lib(_LIB) + return _LIB + + +def _vp(t) -> ctypes.c_void_p: + return ctypes.c_void_p(t.data_ptr()) + +def _launch_cce(lib: ctypes.CDLL, launch_args: RopeLaunchArgs) -> object: + fn = getattr(lib, launch_args.fn_name) + fn( + stream_ptr(), + _vp(launch_args.x), + _vp(launch_args.cos), + _vp(launch_args.sin), + _vp(launch_args.y), + _vp(launch_args.params), + ) + sync() + return launch_args.y + + +def rope_f16(launch_args: RopeLaunchArgs) -> object: + """Launch the local rope f16 CCE kernel.""" + + return _launch_cce(_load_lib(), launch_args) + + +def rope_bf16(launch_args: RopeLaunchArgs) -> object: + """Launch the local rope bf16 CCE kernel.""" + + return _launch_cce(_load_lib(), launch_args) + + +def rope_f32(launch_args: RopeLaunchArgs) -> object: + """Launch the local rope f32 CCE kernel.""" + + return _launch_cce(_load_lib(), launch_args) + + +class RopeCceBackend: + """CCE rope backend implemented locally under kernel-test.""" + + name = "cce" + _launchers = { + "f16": rope_f16, + "bf16": rope_bf16, + "f32": rope_f32, + } + + def is_supported(self, case: object, *, purpose: RunPurpose) -> tuple[bool, str | None]: + del case, purpose + return True, None + + def launch(self, case: object, *, purpose: RunPurpose) -> object: + ensure_runtime("rope") + launch_args = prepare_launch_args(case, cycle=purpose == "cycle") + return self._launchers[launch_args.dtype](launch_args) diff --git a/test/kernel-test/kernels/rope/cce/rope_cce_compute.h b/test/kernel-test/kernels/rope/cce/rope_cce_compute.h new file mode 100644 index 0000000000..df05eb2ef0 --- /dev/null +++ b/test/kernel-test/kernels/rope/cce/rope_cce_compute.h @@ -0,0 +1,729 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// Upstream: cce/tile_kernels_port/rope/csrc/inc/rope_cce_compute.h +#ifndef ROPE_CCE_COMPUTE_H +#define ROPE_CCE_COMPUTE_H + +#include "rope_cce_shim.h" + +#if defined(__DAV_VEC__) + +namespace rope_cce { + +/*=========================================================================== + * + * ROPES — ROTARY POSITIONAL EMBEDDING FOR NPU CCE + * ================================================ + * + * This file contains the **in-register compute kernels** for RoPE + * (Rotary Positional Embedding), executed on a single AIV (vector core). + * + * **Mathematical model** (per complex pair `[x0, x1]`): + * + * y0 = x0 * cos(θ) − x1 * sin(θ) + * y1 = x0 * sin(θ) + x1 * cos(θ) + * + * where θ = position × frequency, and `(cos, sin)` are precomputed tables + * stored in the same UB as the input `x`. + * + * **Data layout on UB** (AB layout, `[B, S, N, D]`): + * + * For a given `(s, n, d)` the RoPE kernel operates on two halves of `D`: + * + * x_0 = x[s, n, d ∈ [0, D/2)] — "low" half + * x_1 = x[s, n, d ∈ [D/2, D)] — "high" half + * + * cos and sin are similarly stored as two halves, sharing the same `s` stride. + * + * **Two rotation modes**: + * + * Half (mode 0, NeoX-style): + * Pairs are `(x[d], x[d+D/2])` for each `d ∈ [0, D/2)`. + * Load two independent half-registers, apply the rotation, store back. + * Used by Llama, Qwen, and most transformer variants. + * + * Interleave (mode 1, GPT-J style): + * Pairs are `(x[2k], x[2k+1])` for each even index `2k`. + * Uses `vdintlv_x2` / `vintlv_x2` to form the partner pairs within + * one register. + * Used by GPT-J and PaLM. + * + * **Intrinsic call pattern** (CCE intrinsics, no AscendC API): + * + * - `vlds` — load from UB into vector register + * - `vcvt` — type conversion (b16↔f32, fp16↔fp32) + * - `vmul/vadd/vsub` — element-wise floating-point arithmetic + * - `vdintlv/vintlv` — deinterleave/interleave two registers + * - `vsts` — store from register to UB + * - `plt_b16/b32` — construct mask for `cnt` active lanes + * + * **Three compute variants**: + * + * - `ComputeF16` — fp16 throughout (fastest, lowest precision) + * - `ComputeBf16` — bf16 input → fp32 inner math → bf16 output (mid precision, same speed) + * - `ComputeF32` — fp32 throughout (slowest, bit-exact reference) + * + * All three variants follow the same outer loop structure: + * + * for each sequence slice `s`: + * for each half-d-block `rep`: + * load cos/sin half-registers (shared across heads) + * for each head `n`: + * load x half-registers (per head) + * compute y0 = cos * x0 - sin * x1 + * compute y1 = sin * x0 + cos * x1 + * store y0, y1 + * + * The cos/sin loads are hoisted outside the `n` loop because cos/sin are + * shared across all heads (`N` dimension). + * + *===========================================================================*/ +/*=========================================================================== + * ComputeF16 — fp16 RoPE computation (HALF and INTERLEAVE modes) + * + * INPUTS: x, cos, sin in UB, all fp16 (`__ubuf__ half *`) + * OUTPUTS: y in UB, fp16 + * PRECISION: full fp16 native — each vmul/vadd/vsub is an IEEE fp16 op. + * + * Data flow per (s, n): + * 1. Load cos_0, cos_1, sin_0, sin_1, x_0, x_1 from UB (vlds NORM). + * 2. y0 = cos_0 * x0 - sin_0 * x1 ; vmul, vmul, vsub + * 3. y1 = cos_1 * x1 + sin_1 * x0 ; vmul, vmul, vadd + * 4. Store y0, y1 to UB (vsts NORM_B16). + * + * Vector register width: 256 B = 128 fp16 elements (VL_F16 = 128). + * Mask granularity: b16 (one bit per fp16 element). + * + * PTO equivalents (for cross-reference): + * - vlds_norm_b16 → pto.vlds {dist = "DIST_NORM_B16"} + * - vmul_f16, vadd_f16 → pto.vmul / pto.vadd + * - vsub_f16 → pto.vsub + * - vsts_norm_b16 → pto.vsts {dist = "DIST_NORM_B16"} + * + * Register pressure: 10 vector_f16 registers (all 256 B each). + *===========================================================================*/ +ROPE_CCE_INTERNAL void ComputeF16( + __ubuf__ uint16_t *x_ub16, + __ubuf__ uint16_t *cos_ub16, + __ubuf__ uint16_t *sin_ub16, + __ubuf__ uint16_t *y_ub16, + int32_t sCount, int32_t nCount, + int32_t dLen, int32_t dAlign, + int32_t xSStep, int32_t xNStep, + int32_t csSStep, + int32_t ySStep, int32_t yNStep, + int32_t mode) +{ + __VEC_SCOPE__ + { + // Expected UB effect of this vector scope (local tensor view): + // x_ub16 -> x: fp16[sCount, nCount, dLen] with strides xSStep/xNStep + // cos_ub16 -> cos: fp16[sCount, dLen] with stride csSStep + // sin_ub16 -> sin: fp16[sCount, dLen] with stride csSStep + // y_ub16 -> y: fp16[sCount, nCount, dLen] with strides ySStep/yNStep + // Typical values from vf_sim tests: + // correctness default: sCount=15, nCount=32, dLen=dAlign=64 + // xNStep=yNStep=csSStep=64, xSStep=ySStep=nCount*64=2048 + // wall-time configs: (s,n)=(1,2),(15,4),(15,8),(15,16),(15,32) + // so xSStep/ySStep range from 128 to 2048 elements; compile-time caps + // are sCount<=15 and nCount<=32 in the standalone VF harness. + // + // Pseudocode after this scope completes: + // y[...] is overwritten; x/cos/sin are read-only. + // + // if mode == 0: # HALF / NeoX layout + // half = dLen // 2 + // for s in range(sCount): + // for n in range(nCount): + // for d in range(half): + // y[s,n,d] = fp16(x[s,n,d] * cos[s,d] + // - x[s,n,d+half] * sin[s,d]) + // y[s,n,d+half] = fp16(x[s,n,d+half] * cos[s,d+half] + // + x[s,n,d] * sin[s,d+half]) + // + // else: # INTERLEAVE / GPT-J layout + // rot = [-x1, x0, -x3, x2, ...] per (s,n) row + // y[s,n,:] = fp16(x[s,n,:] * cos[s,:] + rot * sin[s,:]) + __ubuf__ half *xH = (__ubuf__ half *)x_ub16; + __ubuf__ half *cosH = (__ubuf__ half *)cos_ub16; + __ubuf__ half *sinH = (__ubuf__ half *)sin_ub16; + __ubuf__ half *yH = (__ubuf__ half *)y_ub16; + + int32_t halfD = dLen / 2; + int32_t halfDAl = calign(halfD, (int32_t)BLOCK_BYTE_32 / 2); + + vector_f16 hx1, hx2, hc1, hc2, hs1, hs2, ht1, ht2, hout1, hout2; + + // ---- HALF mode (mode == 0, NeoX-style) ---- + // + // Loop structure: + // s ∈ [0, sCount): outer sequence slice + // rep ∈ [0, repeatTimes): inner half-D tile (128 fp16 per iter) + // [cos/sin loads hoisted here — shared across all n] + // n ∈ [0, nCount): per-head loop + // + // Register roles (each 256 B = 128 fp16): + // hc1 / hc2: cos values for (low half, high half) of D + // hs1 / hs2: sin values for (low half, high half) of D + // hx1 / hx2: x input for (low half, high half) of D (per head) + // ht1 / ht2: temp intermediates + // hout1 / hout2: y output for (low half, high half) + // + // repeatTimesF16 = ceil(halfD / 128) — how many 128-element chunks + // are needed to cover half of D. + if (mode == 0) { + int32_t repeatTimesF16 = cdiv(halfD, (int32_t)VL_F16); + for (uint16_t s = 0; s < (uint16_t)sCount; s++) { + int32_t csOff = s * csSStep; + for (uint16_t rep = 0; rep < (uint16_t)repeatTimesF16; rep++) { + uint16_t elemOff = (uint16_t)((uint32_t)rep * VL_F16); + uint32_t cnt = (uint32_t)(halfD - (int32_t)elemOff); + if (cnt > (uint32_t)VL_F16) cnt = (uint32_t)VL_F16; + // b16 granularity mask — one bit per fp16 element, `cnt` active. + MaskReg mask16 = simd_inlined::make_mask_b16(cnt); + + // ---- Hoisted cos/sin loads (shared across heads) ---- + // Load the 4 half-width cos/sin halves: + // hc1 = cos[s, d=elemOff .. elemOff+128] (low half) + // hc2 = cos[s, d=elemOff+halfDAl .. ] (high half) + // hs1 = sin[s, d=elemOff .. elemOff+128] (low half) + // hs2 = sin[s, d=elemOff+halfDAl .. ] (high half) + simd_inlined::vlds_norm_b16(hc1, cosH + csOff + elemOff, 0); + simd_inlined::vlds_norm_b16(hc2, cosH + csOff + elemOff + halfDAl, 0); + simd_inlined::vlds_norm_b16(hs1, sinH + csOff + elemOff, 0); + simd_inlined::vlds_norm_b16(hs2, sinH + csOff + elemOff + halfDAl, 0); + + for (uint16_t n = 0; n < (uint16_t)nCount; n++) { + int32_t xOff = s * xSStep + n * xNStep; + int32_t yOff = s * ySStep + n * yNStep; + + // ---- Per-head x loads ---- + // Load x halves for this head/position: + // hx1 = x[s, n, d=elemOff .. elemOff+128] (low half) + // hx2 = x[s, n, d=elemOff+halfDAl .. ] (high half) + simd_inlined::vlds_norm_b16(hx1, xH + xOff + elemOff, 0); + simd_inlined::vlds_norm_b16(hx2, xH + xOff + elemOff + halfDAl, 0); + + // ---- Compute y0 = cos_0 * x0 - sin_0 * x1 ---- + simd_inlined::vmul_f16(ht1, hc1, hx1, mask16); // ht1 = cos_0 * x0 + simd_inlined::vmul_f16(ht2, hs1, hx2, mask16); // ht2 = sin_0 * x1 + simd_inlined::vsub_f16(hout1, ht1, ht2, mask16); // y0 = cos_0*x0 - sin_0*x1 + + // ---- Compute y1 = cos_1 * x1 + sin_1 * x0 ---- + simd_inlined::vmul_f16(ht1, hc2, hx2, mask16); // ht1 = cos_1 * x1 + simd_inlined::vmul_f16(ht2, hs2, hx1, mask16); // ht2 = sin_1 * x0 + simd_inlined::vadd_f16(hout2, ht1, ht2, mask16); // y1 = cos_1*x1 + sin_1*x0 + + // ---- Store result (dense 128-fp16 → UB) ---- + simd_inlined::vsts_norm_b16(hout1, yH + yOff + elemOff, 0, mask16); + simd_inlined::vsts_norm_b16(hout2, yH + yOff + elemOff + halfDAl, 0, mask16); + } + } + } + } else { + // ---- INTERLEAVE mode (mode == 1, GPT-J-style) ---- + // + // Pairs are adjacent elements: (x[2k], x[2k+1]). + // Strategy: + // 1. Load 64 contiguous fp16 (x[0..63]). + // 2. `vdintlv_x2(even, odd, x, x)` splits into: + // even = x[0], x[2], ..., x[62] (32 elements) + // odd = x[1], x[3], ..., x[63] (32 elements) + // 3. Negate the odd part to form the "imaginary" partner: + // neg_odd = -x[1], -x[3], ..., -x[63] + // 4. `vintlv_x2(low, high, neg_odd, even)` rebuilds: + // new = [-x1, x0, -x3, x2, ..., -x63, x62] + // This is exactly the "rotated partner" vector needed for + // y = x * cos + rotated(x) * sin. + // + // blockSize is set to 64 (VL_F32) because each interleave step + // consumes 64 fp16 elements (32 even + 32 odd). 256 B register + // = 4 such blocks per iteration. + // + // Register roles (each 256 B = 128 fp16): + // xr : loaded x block + // cosr : loaded cos block (64 fp16, interleaved pattern) + // sinr : loaded sin block (same) + // heven : even-indexed elements after vdintlv + // hodd : odd-indexed elements after vdintlv + // hnegodd: -odd + // hxnew / hxnew_hi: rotated partner vector (low / high halves) + // hta / htb: arithmetic intermediates + // negOne : broadcast scalar -1.0f (all 128 lanes) + vector_f16 xr, cosr, sinr, heven, hodd, hnegodd, hxnew, hxnew_hi, hta, htb; + vector_f16 negOne; + simd_inlined::vbr_f16(negOne, (half)(-1.0f)); + + int32_t blockSize = (int32_t)VL_F32; // 64 fp16 per interleave block + int32_t dBlocks = cdiv(dLen, blockSize); + + for (uint16_t s = 0; s < (uint16_t)sCount; s++) { + int32_t csOff = s * csSStep; + for (uint16_t blk = 0; blk < (uint16_t)dBlocks; blk++) { + int32_t off = (int32_t)blk * blockSize; + int32_t remaining = dLen - off; + uint32_t cnt = (remaining > blockSize) ? (uint32_t)blockSize : (uint32_t)remaining; + // pairCnt = number of (even, odd) pairs in this block = cnt/2 (rounded up). + uint32_t pairCnt = (cnt + 1U) / 2U; + + MaskReg mask = simd_inlined::make_mask_b16(cnt); // full block + MaskReg maskPair = simd_inlined::make_mask_b16(pairCnt); // pair-only + + // ---- Hoisted cos/sin loads (shared across heads) ---- + simd_inlined::vlds_norm_b16(cosr, cosH + csOff + off, 0); + simd_inlined::vlds_norm_b16(sinr, sinH + csOff + off, 0); + + for (uint16_t n = 0; n < (uint16_t)nCount; n++) { + int32_t xOff = s * xSStep + n * xNStep; + int32_t yOff = s * ySStep + n * yNStep; + + // ---- Per-head x load ---- + // xr = x[s, n, d=off .. off+blockSize] (64 contiguous fp16) + simd_inlined::vlds_norm_b16(xr, xH + xOff + off, 0); + + // ---- Form the rotated partner vector ---- + // Step 1: split xr into even/odd half-index streams: + // heven = x[0], x[2], ..., x[62] + // hodd = x[1], x[3], ..., x[63] + simd_inlined::vdintlv_x2(heven, hodd, xr, xr); + // Step 2: negate the odd elements to match the rotation + // identity [-x1, -x3, ..., -x63]. + simd_inlined::vmul_f16(hnegodd, hodd, negOne, maskPair); + // Step 3: re-interleave (-odd, even) to form the partner: + // hxnew = [-x1, x0, -x3, x2, ..., -x63, x62] + // hxnew_hi = high half of the interleaved result (for cnt > 64) + simd_inlined::vintlv_x2(hxnew, hxnew_hi, hnegodd, heven); + + // ---- y = x * cos + rotated(x) * sin ---- + simd_inlined::vmul_f16(hta, xr, cosr, mask); // hta = x * cos + simd_inlined::vmul_f16(htb, hxnew, sinr, mask); // htb = rotated(x) * sin + simd_inlined::vadd_f16(htb, hta, htb, mask); // y = hta + htb + simd_inlined::vsts_norm_b16(htb, yH + yOff + off, 0, mask); + } + } + } + } + } +} + +/* + * ComputeBf16 — bf16 RoPE computation (HALF and INTERLEAVE modes) + * x/y are bf16 in GM, cos/sin are fp16 in GM. + * Internal compute is done in fp32: load bf16→fp32, load fp16→fp32, + * fp32 math, fp32→bf16 store. + * + * Load/conversion convention: + * All b16 loads in this function use UNPK_B16 mode, which places 64 + * valid elements at the EVEN halfword positions [0, 2, ..., 126] of + * the 128-lane register; the odd positions contain zero/padding. + * This is why only vcvt_*_to_fp32_even (PART_EVEN) is used to widen + * to fp32 — it recovers all 64 valid elements. vcvt_*_to_fp32_odd + * would extract the padding lanes and produce garbage. (The _odd + * wrapper exists in the shim for other kernels that fill all 128 + * lanes via denser loads, e.g. mx_quant with DINTLV_B16 x2.) + * + * HALF mode (mode==0): loads bf16 x and fp16 cos/sin via UNPK_B16, + * widens both to fp32 via PART_EVEN, does fp32 mul/add/sub, narrows + * back to bf16 via PART_EVEN + PK_B32. + * + * INTERLEAVE mode (mode==1): same load pattern, plus fp32-domain + * vdintlv/vintlv for the even/odd pair shuffle. Avoids the + * incorrect fp16 bit-reinterpretation approach. + */ +ROPE_CCE_INTERNAL void ComputeBf16( + __ubuf__ uint16_t *x_ub16, + __ubuf__ uint16_t *cos_ub16, + __ubuf__ uint16_t *sin_ub16, + __ubuf__ uint16_t *y_ub16, + int32_t sCount, int32_t nCount, + int32_t dLen, int32_t dAlign, + int32_t xSStep, int32_t xNStep, + int32_t csSStep, + int32_t ySStep, int32_t yNStep, + int32_t mode) +{ + __VEC_SCOPE__ + { + // Expected UB effect of this vector scope (local tensor view): + // x_ub16 -> x: bf16[sCount, nCount, dLen] with strides xSStep/xNStep + // cos_ub16 -> cos: fp16[sCount, dLen] with stride csSStep + // sin_ub16 -> sin: fp16[sCount, dLen] with stride csSStep + // y_ub16 -> y: bf16[sCount, nCount, dLen] with strides ySStep/yNStep + // Typical values from vf_sim tests: + // correctness default: sCount=15, nCount=32, dLen=dAlign=64 + // xNStep=yNStep=csSStep=64, xSStep=ySStep=nCount*64=2048 + // wall-time configs: (s,n)=(1,2),(15,4),(15,8),(15,16),(15,32) + // so xSStep/ySStep range from 128 to 2048 elements; compile-time caps + // are sCount<=15 and nCount<=32 in the standalone VF harness. + // + // Pseudocode after this scope completes: + // y[...] is overwritten; x/cos/sin are read-only. + // All arithmetic below is fp32; stores narrow back to bf16. + // + // if mode == 0: # HALF / NeoX layout + // half = dLen // 2 + // for s in range(sCount): + // for n in range(nCount): + // xf = x[s,n,:].astype(float32) + // cf = cos[s,:].astype(float32) + // sf = sin[s,:].astype(float32) + // for d in range(half): + // y[s,n,d] = bf16(xf[d] * cf[d] + // - xf[d+half] * sf[d]) + // y[s,n,d+half] = bf16(xf[d+half] * cf[d+half] + // + xf[d] * sf[d+half]) + // + // else: # INTERLEAVE / GPT-J layout + // xf = x[s,n,:].astype(float32) + // rot = [-xf[1], xf[0], -xf[3], xf[2], ...] + // y[s,n,:] = bf16(xf * cos[s,:].astype(float32) + // + rot * sin[s,:].astype(float32)) + // ---- Typed UB pointer setup ---- + // x/y are bf16 in UB; cos/sin are fp16 in UB. + // Casts only set the address-space / element-type view. + __ubuf__ bfloat16_t *xB = (__ubuf__ bfloat16_t *)x_ub16; + __ubuf__ bfloat16_t *yB = (__ubuf__ bfloat16_t *)y_ub16; + __ubuf__ half *cosH = (__ubuf__ half *)cos_ub16; + __ubuf__ half *sinH = (__ubuf__ half *)sin_ub16; + + // halfD = D/2 (the width of each rotatable half) + // halfDAl = halfD aligned up to (BLOCK_BYTE_32/2) = 16 elements, + // ensuring strided offsets land on 32-byte boundaries. + // repeatTimes = ceil(halfD / VL_F32) — each iteration processes + // VL_F32 = 64 fp32 elements (= one full vector register). + int32_t halfD = dLen / 2; + int32_t halfDAl = calign(halfD, (int32_t)BLOCK_BYTE_32 / 2); + int32_t repeatTimes = cdiv(halfD, (int32_t)VL_F32); + + // ---- Register allocation ---- + // fc, fc2, fs, fs2: cos/sin widened to fp32 (low and high halves) + // fx0, fx1: x halves widened to fp32 + // ft, ft2: arithmetic intermediates + // htmp0, htmp1: scratch 128-fp16 / 128-bf16 slots (used as the + // landing zone for UNPK_B16 loads and for PK_B32 stores) + vector_f32 fc, fc2, fs, fs2, fx0, fx1, ft, ft2; + vector_f16 htmp0, htmp1; + + if (mode == 0) { + // ==== ComputeBf16 HALF mode ==== + // Same outer loop structure as ComputeF16 HALF, but with + // widening load + narrowing store: + // vlds(UNPK_B16) → b16 at even register positions + // vcvt_*_to_fp32_even → fp32 (64 lanes, matching active mask) + // vmul / vadd / vsub → fp32 arithmetic + // vcvt_f32_to_bf16_narrow → bf16 at even register positions + // vsts(PK_B32) → dense bf16 written back to UB + for (uint16_t s = 0; s < (uint16_t)sCount; s++) { + int32_t csOff = s * csSStep; + for (uint16_t rep = 0; rep < (uint16_t)repeatTimes; rep++) { + uint32_t elemOff = (uint32_t)rep * (uint32_t)VL_F32; + int32_t elemOffH = (int32_t)elemOff; + uint32_t cnt = (uint32_t)(halfD - elemOffH); + if (cnt > (uint32_t)VL_F32) cnt = (uint32_t)VL_F32; + // 64-lane (b32) mask — `cnt` active fp32 lanes. + MaskReg mask32 = simd_inlined::make_mask(cnt); + + // ---- Hoisted cos/sin load-and-widen (shared across heads) ---- + // For both halves (low & high) of D: + // load 64 fp16 via UNPK_B16 (landing at even positions) + // widen to 64 fp32 via PART_EVEN. + simd_inlined::vlds_unpk_b16(htmp0, cosH + csOff + elemOffH, 0); + simd_inlined::vcvt_fp16_to_fp32_even(fc, htmp0, mask32); // fc = cos_low (fp32) + simd_inlined::vlds_unpk_b16(htmp0, sinH + csOff + elemOffH, 0); + simd_inlined::vcvt_fp16_to_fp32_even(fs, htmp0, mask32); // fs = sin_low (fp32) + + simd_inlined::vlds_unpk_b16(htmp0, cosH + csOff + elemOffH + halfDAl, 0); + simd_inlined::vcvt_fp16_to_fp32_even(fc2, htmp0, mask32); // fc2 = cos_high (fp32) + simd_inlined::vlds_unpk_b16(htmp0, sinH + csOff + elemOffH + halfDAl, 0); + simd_inlined::vcvt_fp16_to_fp32_even(fs2, htmp0, mask32); // fs2 = sin_high (fp32) + + for (uint16_t n = 0; n < (uint16_t)nCount; n++) { + int32_t xOff = s * xSStep + n * xNStep; + int32_t yOff = s * ySStep + n * yNStep; + + // ---- Per-head x load-and-widen ---- + // Load two bf16 halves via UNPK_B16, widen each to fp32. + simd_inlined::vlds_unpk_b16_bf16(htmp0, xB + xOff + elemOffH, 0); + simd_inlined::vcvt_bf16_to_fp32_even(fx0, htmp0, mask32); // fx0 = x_low (fp32) + simd_inlined::vlds_unpk_b16_bf16(htmp1, xB + xOff + elemOffH + halfDAl, 0); + simd_inlined::vcvt_bf16_to_fp32_even(fx1, htmp1, mask32); // fx1 = x_high (fp32) + + // ---- Compute y0 = cos_low * x_low - sin_low * x_high ---- + simd_inlined::vmul_f32(ft, fc, fx0, mask32); // ft = cos_low * x_low + simd_inlined::vmul_f32(ft2, fs, fx1, mask32); // ft2 = sin_low * x_high + simd_inlined::vsub_f32(ft, ft, ft2, mask32); // y0 = ft - ft2 + + // ---- Narrow y0 to bf16 and store (low half of y) ---- + simd_inlined::vcvt_f32_to_bf16_narrow(htmp0, ft, mask32); // fp32 → bf16 at even lanes + simd_inlined::vsts_pk_b32_bf16(htmp0, yB + yOff + elemOffH, 0, mask32); // dense bf16 → UB + + // ---- Compute y1 = cos_high * x_high + sin_high * x_low ---- + simd_inlined::vmul_f32(ft, fc2, fx1, mask32); // ft = cos_high * x_high + simd_inlined::vmul_f32(ft2, fs2, fx0, mask32); // ft2 = sin_high * x_low + simd_inlined::vadd_f32(ft, ft, ft2, mask32); // y1 = ft + ft2 + + // ---- Narrow y1 to bf16 and store (high half of y) ---- + simd_inlined::vcvt_f32_to_bf16_narrow(htmp0, ft, mask32); + simd_inlined::vsts_pk_b32_bf16(htmp0, yB + yOff + elemOffH + halfDAl, 0, mask32); + } + } + } + } else { + // ==== ComputeBf16 INTERLEAVE mode ==== + // Same strategy as ComputeF16 INTERLEAVE (vdintlv/vintlv), + // but performed in fp32 register space after widening the + // bf16 x via UNPK_B16 + PART_EVEN. + // + // negOne is a broadcast fp32 -1.0 to all 64 lanes — used + // to negate the odd-indexed x elements with a single vmul. + vector_f32 fx_even, fx_odd, negOne; + simd_inlined::vbr_f32(negOne, -1.0f); + + // blockSize = 64 fp32 lanes; each block covers 64 bf16 input elements. + int32_t blockSize = (int32_t)VL_F32; + int32_t dBlocks = cdiv(dLen, blockSize); + + for (uint16_t s = 0; s < (uint16_t)sCount; s++) { + int32_t csOff = s * csSStep; + for (uint16_t blk = 0; blk < (uint16_t)dBlocks; blk++) { + int32_t off = (int32_t)blk * blockSize; + int32_t remaining = dLen - off; + uint32_t cnt = (remaining > blockSize) ? (uint32_t)blockSize : (uint32_t)remaining; + // pairCnt = ceil(cnt/2) — number of even/odd pairs in this block. + uint32_t pairCnt = (cnt + 1U) / 2U; + + MaskReg mask32 = simd_inlined::make_mask(cnt); // full block (b32) + MaskReg maskHalf = simd_inlined::make_mask(pairCnt); // pairs only + + // ---- Hoisted cos/sin load-and-widen ---- + simd_inlined::vlds_unpk_b16(htmp0, cosH + csOff + off, 0); + simd_inlined::vcvt_fp16_to_fp32_even(fc, htmp0, mask32); // fc = cos (fp32) + simd_inlined::vlds_unpk_b16(htmp0, sinH + csOff + off, 0); + simd_inlined::vcvt_fp16_to_fp32_even(fs, htmp0, mask32); // fs = sin (fp32) + + for (uint16_t n = 0; n < (uint16_t)nCount; n++) { + int32_t xOff = s * xSStep + n * xNStep; + int32_t yOff = s * ySStep + n * yNStep; + + // ---- Per-head x load-and-widen ---- + simd_inlined::vlds_unpk_b16_bf16(htmp0, xB + xOff + off, 0); + simd_inlined::vcvt_bf16_to_fp32_even(fx0, htmp0, mask32); // fx0 = x (fp32) + + // ---- Form the rotated partner in fp32 space ---- + // Step 1: split into even/odd streams: + // fx_even = x[0], x[2], x[4], ... + // fx_odd = x[1], x[3], x[5], ... + simd_inlined::vdintlv_x2(fx_even, fx_odd, fx0, fx0); + // Step 2: negate the odd elements to -x[1], -x[3], ... + simd_inlined::vmul_f32(fx_odd, fx_odd, negOne, maskHalf); + // Step 3: re-interleave to form the rotated partner: + // ft = [-x1, x0, -x3, x2, ...] (low half) + // ft2 = high half of interleaved result + simd_inlined::vintlv_x2(ft, ft2, fx_odd, fx_even); + + // ---- y = x * cos + rotated(x) * sin ---- + // ft2 = x * cos (using the ORIGINAL widened `fx0`, not the + // deinterleaved-even copy — both hold the same values but + // `fx0` avoids aliasing concerns with the vintlv output). + // ft = rotated(x) * sin + // then add to get y. + simd_inlined::vmul_f32(ft2, fx0, fc, mask32); // ft2 = x * cos + simd_inlined::vmul_f32(ft, ft, fs, mask32); // ft = rotated(x) * sin + simd_inlined::vadd_f32(ft, ft2, ft, mask32); // y = x*cos + rotated*sin + + // ---- Narrow and store ---- + simd_inlined::vcvt_f32_to_bf16_narrow(htmp0, ft, mask32); + simd_inlined::vsts_pk_b32_bf16(htmp0, yB + yOff + off, 0, mask32); + } + } + } + } + } +} + +/*=========================================================================== + * ComputeF32 — fp32 RoPE computation (HALF and INTERLEAVE modes) + * + * INPUTS: x, cos, sin in UB, all fp32 (`__ubuf__ float *`) + * OUTPUTS: y in UB, fp32 + * PRECISION: native fp32 — bit-exact vs fp64 PyTorch reference. + * + * This variant exists primarily for validation (max_diff = 0.0 vs fp64 + * reference). The cost is 2× higher memory traffic (4 B/element) and + * therefore 2× lower bandwidth than fp16/bf16 for the same shape. + * + * Data flow per (s, n): + * Same as ComputeF16, but: + * - Loads via `vlds_norm_b32` (dense 64-fp32 per register). + * - All arithmetic in fp32 registers. + * - Stores via `vsts_norm_b32` (dense 64-fp32 back to UB). + * + * Vector register width: 256 B = 64 fp32 elements (VL_F32 = 64). + * Mask granularity: b32 (one bit per fp32 element). + *===========================================================================*/ +ROPE_CCE_INTERNAL void ComputeF32( + __ubuf__ uint16_t *x_ub16, + __ubuf__ uint16_t *cos_ub16, + __ubuf__ uint16_t *sin_ub16, + __ubuf__ uint16_t *y_ub16, + int32_t sCount, int32_t nCount, + int32_t dLen, int32_t dAlign, + int32_t xSStep, int32_t xNStep, + int32_t csSStep, + int32_t ySStep, int32_t yNStep, + int32_t mode) +{ + __VEC_SCOPE__ + { + // Expected UB effect of this vector scope (local tensor view): + // x_ub16 -> x: fp32[sCount, nCount, dLen] with strides xSStep/xNStep + // cos_ub16 -> cos: fp32[sCount, dLen] with stride csSStep + // sin_ub16 -> sin: fp32[sCount, dLen] with stride csSStep + // y_ub16 -> y: fp32[sCount, nCount, dLen] with strides ySStep/yNStep + // Typical values from vf_sim tests: + // correctness default: sCount=15, nCount=32, dLen=dAlign=64 + // xNStep=yNStep=csSStep=64, xSStep=ySStep=nCount*64=2048 + // wall-time configs: (s,n)=(1,2),(15,4),(15,8),(15,16),(15,32) + // so xSStep/ySStep range from 128 to 2048 elements; compile-time caps + // are sCount<=15 and nCount<=32 in the standalone VF harness. + // + // Pseudocode after this scope completes: + // y[...] is overwritten; x/cos/sin are read-only. + // + // if mode == 0: # HALF / NeoX layout + // half = dLen // 2 + // for s in range(sCount): + // for n in range(nCount): + // for d in range(half): + // y[s,n,d] = x[s,n,d] * cos[s,d] + // - x[s,n,d+half] * sin[s,d] + // y[s,n,d+half] = x[s,n,d+half] * cos[s,d+half] + // + x[s,n,d] * sin[s,d+half] + // + // else: # INTERLEAVE / GPT-J layout + // xe, xo = x[s,n,0::2], x[s,n,1::2] + // ce, co = cos[s,0::2], cos[s,1::2] + // se, so = sin[s,0::2], sin[s,1::2] + // y[s,n,0::2] = xe * ce - xo * se + // y[s,n,1::2] = xo * co + xe * so + __ubuf__ float *xF = (__ubuf__ float *)x_ub16; + __ubuf__ float *cosF = (__ubuf__ float *)cos_ub16; + __ubuf__ float *sinF = (__ubuf__ float *)sin_ub16; + __ubuf__ float *yF = (__ubuf__ float *)y_ub16; + + // halfDAl aligned to (BLOCK_BYTE_32 / 4) = 8 fp32 elements = 32 bytes. + int32_t halfD = dLen / 2; + int32_t halfDAl = calign(halfD, (int32_t)BLOCK_BYTE_32 / 4); + int32_t repeatTimes = cdiv(halfD, (int32_t)VL_F32); + + // Register roles (each 256 B = 64 fp32): + // fc0/fc1, fs0/fs1: cos/sin halves from both halves of D. + // fx0/fx1: x halves per head. + // ft0/ft1: arithmetic intermediates. + // fy0/fy1: y output halves. + vector_f32 fx0, fx1, fc0, fc1, fs0, fs1, ft0, ft1, fy0, fy1; + + // ---- HALF mode (mode == 0) ---- + // Identical dataflow to ComputeF16 HALF modulo element precision. + if (mode == 0) { + for (uint16_t s = 0; s < (uint16_t)sCount; s++) { + int32_t csOff = s * csSStep; + for (uint16_t rep = 0; rep < (uint16_t)repeatTimes; rep++) { + uint32_t elemOff = (uint32_t)rep * (uint32_t)VL_F32; + uint32_t cnt = (uint32_t)(halfD - (int32_t)elemOff); + if (cnt > (uint32_t)VL_F32) cnt = (uint32_t)VL_F32; + MaskReg mask32 = simd_inlined::make_mask(cnt); + + // ---- Hoisted cos/sin loads (shared across heads) ---- + // fc0 = cos_low, fc1 = cos_high + // fs0 = sin_low, fs1 = sin_high + simd_inlined::vlds_norm_b32(fc0, cosF + csOff + elemOff, 0); + simd_inlined::vlds_norm_b32(fc1, cosF + csOff + elemOff + halfDAl, 0); + simd_inlined::vlds_norm_b32(fs0, sinF + csOff + elemOff, 0); + simd_inlined::vlds_norm_b32(fs1, sinF + csOff + elemOff + halfDAl, 0); + + for (uint16_t n = 0; n < (uint16_t)nCount; n++) { + int32_t xOff = s * xSStep + n * xNStep; + int32_t yOff = s * ySStep + n * yNStep; + + // ---- Per-head x loads ---- + simd_inlined::vlds_norm_b32(fx0, xF + xOff + elemOff, 0); // x_low + simd_inlined::vlds_norm_b32(fx1, xF + xOff + elemOff + halfDAl, 0); // x_high + + // y0 = cos_low * x_low - sin_low * x_high + simd_inlined::vmul_f32(ft0, fc0, fx0, mask32); + simd_inlined::vmul_f32(ft1, fs0, fx1, mask32); + simd_inlined::vsub_f32(fy0, ft0, ft1, mask32); + // y1 = cos_high * x_high + sin_high * x_low + simd_inlined::vmul_f32(ft0, fc1, fx1, mask32); + simd_inlined::vmul_f32(ft1, fs1, fx0, mask32); + simd_inlined::vadd_f32(fy1, ft0, ft1, mask32); + + // y → UB + simd_inlined::vsts_norm_b32(fy0, yF + yOff + elemOff, 0, mask32); + simd_inlined::vsts_norm_b32(fy1, yF + yOff + elemOff + halfDAl, 0, mask32); + } + } + } + } else { + // ---- INTERLEAVE mode (mode == 1) ---- + // Same pattern as ComputeF16/ComputeBf16 INTERLEAVE (vdintlv/vintlv + // to form the rotated partner), in fp32 register space. + // Note: since fp32 registers have 64 lanes, blockSize = 64, so each + // block corresponds to 64 fp32 elements, i.e. 32 complex pairs. + vector_f32 heven, hodd, hnegOdd, hxnew, hxnew_hi, fta, ftb; + vector_f32 negOne; + simd_inlined::vbr_f32(negOne, -1.0f); + + int32_t blockSize = (int32_t)VL_F32; // 64 fp32 per block + int32_t dBlocks = cdiv(dLen, blockSize); + + for (uint16_t s = 0; s < (uint16_t)sCount; s++) { + int32_t csOff = s * csSStep; + for (uint16_t blk = 0; blk < (uint16_t)dBlocks; blk++) { + int32_t off = (int32_t)blk * blockSize; + int32_t remaining = dLen - off; + uint32_t cnt = (remaining > blockSize) ? (uint32_t)blockSize : (uint32_t)remaining; + + MaskReg mask = simd_inlined::make_mask(cnt); + MaskReg maskHalf = simd_inlined::make_mask((cnt + 1U) / 2U); + + // ---- Hoisted cos/sin loads ---- + simd_inlined::vlds_norm_b32(fc0, cosF + csOff + off, 0); + simd_inlined::vlds_norm_b32(fs0, sinF + csOff + off, 0); + + for (uint16_t n = 0; n < (uint16_t)nCount; n++) { + int32_t xOff = s * xSStep + n * xNStep; + int32_t yOff = s * ySStep + n * yNStep; + + // ---- Per-head x load ---- + simd_inlined::vlds_norm_b32(fx0, xF + xOff + off, 0); + + // ---- Form the rotated partner: [-x1, x0, -x3, x2, ...] ---- + simd_inlined::vdintlv_x2(heven, hodd, fx0, fx0); // split even/odd + simd_inlined::vmul_f32(hnegOdd, hodd, negOne, maskHalf); // negate odd + simd_inlined::vintlv_x2(hxnew, hxnew_hi, hnegOdd, heven); // rebuild partner + + // y = x * cos + rotated(x) * sin + simd_inlined::vmul_f32(fta, fx0, fc0, mask); // x * cos + simd_inlined::vmul_f32(ftb, hxnew, fs0, mask); // rotated(x) * sin + simd_inlined::vadd_f32(fy0, fta, ftb, mask); // y + simd_inlined::vsts_norm_b32(fy0, yF + yOff + off, 0, mask); + } + } + } + } + } +} + +} + +#endif + +#endif diff --git a/test/kernel-test/kernels/rope/cce/rope_cce_gm_io.h b/test/kernel-test/kernels/rope/cce/rope_cce_gm_io.h new file mode 100644 index 0000000000..f398c288ba --- /dev/null +++ b/test/kernel-test/kernels/rope/cce/rope_cce_gm_io.h @@ -0,0 +1,71 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#ifndef ROPE_CCE_GM_IO_H +#define ROPE_CCE_GM_IO_H + +#include "rope_cce_shim.h" + +namespace rope_cce { + +// UB caps sized for the production AB tile (ubFactorBS=15, ubFactorN=32, D=64). +// Runtime sCount/nCount passed to the kernel must stay within these limits. +constexpr int32_t kMaxS = 15; +constexpr int32_t kMaxN = 32; +constexpr int32_t kMaxD = 64; +constexpr int32_t kMaxDAlign = 64; + +constexpr int32_t kEsF16 = 2; +constexpr int32_t kEsF32 = 4; + +ROPE_CCE_INTERNAL void GmLoadContig( + __ubuf__ uint16_t *ub, + __gm__ uint16_t *gm, + uint32_t elemCount, + int32_t es) +{ + uint32_t bytes = elemCount * (uint32_t)es; + copy_gm_to_ubuf_align_v2( + ub, gm, + 0, 1, bytes, + 0, 0, false, 0, + bytes, bytes); +} + +ROPE_CCE_INTERNAL void GmStoreContig( + __gm__ uint16_t *gm, + __ubuf__ uint16_t *ub, + uint32_t elemCount, + int32_t es) +{ + uint32_t bytes = elemCount * (uint32_t)es; + copy_ubuf_to_gm_align_v2( + gm, ub, + 0, 1, bytes, + 0, bytes, bytes); +} + +ROPE_CCE_INTERNAL void GmLoadF16Contig( + __ubuf__ uint16_t *ub, + __gm__ uint16_t *gm, + uint32_t elemCount) +{ + GmLoadContig(ub, gm, elemCount, kEsF16); +} + +ROPE_CCE_INTERNAL void GmStoreF16Contig( + __gm__ uint16_t *gm, + __ubuf__ uint16_t *ub, + uint32_t elemCount) +{ + GmStoreContig(gm, ub, elemCount, kEsF16); +} + +} // namespace rope_cce + +#endif diff --git a/test/kernel-test/kernels/rope/cce/rope_cce_kernel.cpp b/test/kernel-test/kernels/rope/cce/rope_cce_kernel.cpp new file mode 100644 index 0000000000..82e0af1168 --- /dev/null +++ b/test/kernel-test/kernels/rope/cce/rope_cce_kernel.cpp @@ -0,0 +1,164 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +/** + * Standalone single-core rope CCE kernel. + * Wraps ComputeF16/Bf16/F32 (__VEC_SCOPE__) with trivial GM<->UB I/O. + */ +#include "rope_cce_compute.h" +#include "rope_cce_gm_io.h" + +using namespace rope_cce; + +namespace { + +ROPE_CCE_INTERNAL void ReadTileParams( + __gm__ int32_t *params_g, int32_t &sCount, int32_t &nCount) +{ + sCount = params_g[0]; + nCount = params_g[1]; +} + +template +ROPE_CCE_INTERNAL void RunRopeTile( + __gm__ uint16_t *x_g, + __gm__ uint16_t *cos_g, + __gm__ uint16_t *sin_g, + __gm__ uint16_t *y_g, + int32_t sCount, + int32_t nCount) +{ + int32_t es = (kDtypeMode == 2) ? kEsF32 : kEsF16; + int32_t cosSinElems = sCount * kMaxDAlign; + int32_t xyElems = sCount * nCount * kMaxDAlign; + int32_t cosSinBytes = cosSinElems * es; + int32_t xyBytes = xyElems * es; + + __ubuf__ uint16_t *ub_base = (__ubuf__ uint16_t *)0x00000; + __ubuf__ uint16_t *cos_ub = ub_base; + __ubuf__ uint16_t *sin_ub = (__ubuf__ uint16_t *)((uintptr_t)cos_ub + cosSinBytes); + __ubuf__ uint16_t *x_ub = (__ubuf__ uint16_t *)((uintptr_t)sin_ub + cosSinBytes); + __ubuf__ uint16_t *y_ub = (__ubuf__ uint16_t *)((uintptr_t)x_ub + xyBytes); + + GmLoadContig(cos_ub, cos_g, (uint32_t)cosSinElems, es); + GmLoadContig(sin_ub, sin_g, (uint32_t)cosSinElems, es); + GmLoadContig(x_ub, x_g, (uint32_t)xyElems, es); + + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + if constexpr (kDtypeMode == 0) { + ComputeF16( + x_ub, cos_ub, sin_ub, y_ub, + sCount, nCount, + kMaxD, kMaxDAlign, + nCount * kMaxDAlign, kMaxDAlign, + kMaxDAlign, + nCount * kMaxDAlign, kMaxDAlign, + kMode); + } else if constexpr (kDtypeMode == 1) { + ComputeBf16( + x_ub, cos_ub, sin_ub, y_ub, + sCount, nCount, + kMaxD, kMaxDAlign, + nCount * kMaxDAlign, kMaxDAlign, + kMaxDAlign, + nCount * kMaxDAlign, kMaxDAlign, + kMode); + } else { + ComputeF32( + x_ub, cos_ub, sin_ub, y_ub, + sCount, nCount, + kMaxD, kMaxDAlign, + nCount * kMaxDAlign, kMaxDAlign, + kMaxDAlign, + nCount * kMaxDAlign, kMaxDAlign, + kMode); + } + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + + GmStoreContig(y_g, y_ub, (uint32_t)xyElems, es); +} + +template +__global__ AICORE void rope_cce_sim_kernel( + __gm__ uint16_t *x_g, + __gm__ uint16_t *cos_g, + __gm__ uint16_t *sin_g, + __gm__ uint16_t *y_g, + __gm__ int32_t *params_g) +{ +#if defined(__DAV_VEC__) + int32_t sCount = 0; + int32_t nCount = 0; + ReadTileParams(params_g, sCount, nCount); + RunRopeTile(x_g, cos_g, sin_g, y_g, sCount, nCount); +#endif +} + +template +__global__ AICORE void rope_cce_cycle_kernel( + __gm__ uint16_t *x_g, + __gm__ uint16_t *cos_g, + __gm__ uint16_t *sin_g, + __gm__ uint16_t *y_g, + __gm__ int32_t *params_g) +{ +#if defined(__DAV_VEC__) + int32_t sCount = 0; + int32_t nCount = 0; + ReadTileParams(params_g, sCount, nCount); + RunRopeTile(x_g, cos_g, sin_g, y_g, sCount, nCount); +#endif +} + +} // namespace + +#define ROPE_CCE_LAUNCH_SIM(mode, dtype_mode, suffix) \ + void call_rope_cce_sim_##suffix( \ + void *stream, void *x, void *cos, void *sin, void *y, void *params) \ + { \ + rope_cce_sim_kernel<<<1, nullptr, stream>>>( \ + (__gm__ uint16_t *)x, \ + (__gm__ uint16_t *)cos, \ + (__gm__ uint16_t *)sin, \ + (__gm__ uint16_t *)y, \ + (__gm__ int32_t *)params); \ + } + +#define ROPE_CCE_LAUNCH_CYCLE(mode, dtype_mode, suffix) \ + void call_rope_cce_cycle_##suffix( \ + void *stream, void *x, void *cos, void *sin, void *y, void *params) \ + { \ + rope_cce_cycle_kernel<<<1, nullptr, stream>>>( \ + (__gm__ uint16_t *)x, \ + (__gm__ uint16_t *)cos, \ + (__gm__ uint16_t *)sin, \ + (__gm__ uint16_t *)y, \ + (__gm__ int32_t *)params); \ + } + +extern "C" { + +ROPE_CCE_LAUNCH_SIM(0, 0, half_f16) +ROPE_CCE_LAUNCH_SIM(1, 0, interleave_f16) +ROPE_CCE_LAUNCH_SIM(0, 1, half_bf16) +ROPE_CCE_LAUNCH_SIM(1, 1, interleave_bf16) +ROPE_CCE_LAUNCH_SIM(0, 2, half_f32) +ROPE_CCE_LAUNCH_SIM(1, 2, interleave_f32) + +ROPE_CCE_LAUNCH_CYCLE(0, 0, half_f16) +ROPE_CCE_LAUNCH_CYCLE(1, 0, interleave_f16) +ROPE_CCE_LAUNCH_CYCLE(0, 1, half_bf16) +ROPE_CCE_LAUNCH_CYCLE(1, 1, interleave_bf16) +ROPE_CCE_LAUNCH_CYCLE(0, 2, half_f32) +ROPE_CCE_LAUNCH_CYCLE(1, 2, interleave_f32) + +} diff --git a/test/kernel-test/kernels/rope/cce/rope_cce_shim.h b/test/kernel-test/kernels/rope/cce/rope_cce_shim.h new file mode 100644 index 0000000000..31adca0189 --- /dev/null +++ b/test/kernel-test/kernels/rope/cce/rope_cce_shim.h @@ -0,0 +1,361 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// Upstream: cce/tile_kernels_port/rope/csrc/inc/rope_cce_shim.h +#ifndef ROPE_CCE_SHIM_H +#define ROPE_CCE_SHIM_H + +#include + +/*=========================================================================== + * + * ROPE CCE SHIM — Typed wrappers around raw CCE intrinsics + * ======================================================== + * + * This header provides a small set of strongly-typed template wrappers + * around the raw CCE SIMD intrinsics. The wrappers use `__simd_callee__` + * + `always_inline` to keep the compiler happy while enabling bisheng's + * SIMD-inlining / fusion pass. + * + * Naming convention: + * + * vlds__ — load into register + * vsts__ — store from register + * vcvt__to__ — type conversion + * vmul/vadd/vsub_ — element-wise arithmetic + * vdintlv_x2 / vintlv_x2 — in-register deinterleave / interleave + * vbr_ — scalar broadcast + * make_mask[_b16](cnt) — construct partial predicate mask + * + * Each wrapper maps 1:1 to a PTO IR primitive. See the inline comments + * above each wrapper for the specific PTO equivalent, hardware semantics, + * and typical use cases within the RoPE kernel. + * + * Register width summary on Ascend 950DT (dav-c310-vec): + * + * vector register = 256 bytes = 2048 bits + * vector_f16 / vector_bf16 / vector_u16 = 128 elements + * vector_f32 / vector_u32 / vector_s32 = 64 elements + * + *===========================================================================*/ + +#ifndef __CPU_SIM +#define AICORE [aicore] +#else +#define AICORE +#endif + +#ifdef __CPU_SIM +#define ROPE_CCE_INTERNAL inline +#define ROPE_SIMD_FN inline +#else +// `__simd_callee__` marks the function as a candidate for bisheng's +// SIMD-inlining pass — the compiler will then merge adjacent simd calls +// into a single wide op when possible. `always_inline` ensures no call +// overhead even at -O0. +#define ROPE_CCE_INTERNAL AICORE inline __attribute__((always_inline)) +#define ROPE_SIMD_FN __attribute__((always_inline)) __simd_callee__ inline +#endif + +namespace rope_cce { + +// Vector lane counts (elements per 256-byte register). +constexpr uint16_t VL_F32 = 64; // fp32 / i32 +constexpr uint16_t VL_F16 = 128; // fp16 / bf16 / i16 +// CCE 32-byte hardware block size (used for aligning strided UB offsets). +constexpr uint16_t BLOCK_BYTE_32 = 32; + +// Integer ceiling helpers used throughout the kernel. +ROPE_CCE_INTERNAL constexpr int32_t cdiv(int32_t a, int32_t b) { return (a + b - 1) / b; } +ROPE_CCE_INTERNAL constexpr int32_t calign(int32_t a, int32_t b) { return (a + b - 1) / b * b; } + +// Predicate register alias — used as the `mask` parameter of every +// predicated intrinsic (vadd, vmul, vcvt, vsts, ...). +using MaskReg = vector_bool; + +namespace simd_inlined { + +// Load 64 b16 values from UB into a 128-lane vector register using +// UNPK_B16 mode. The hardware places the 64 valid elements at EVEN +// halfword positions [0, 2, 4, ..., 126] within the register; the odd +// positions [1, 3, 5, ..., 127] are filled with zero/padding. +// +// This matches vlds_unpk_b16_bf16 below and pairs exclusively with +// vcvt_*_to_fp32_*even (PART_EVEN). Using vcvt_odd after UNPK_B16 +// would extract the zero/padding lanes and produce garbage. +template +ROPE_SIMD_FN void vlds_unpk_b16(F16Dst &dst, U16Src src, int32_t off) +{ + vlds((vector_f16 &)dst, (__ubuf__ half *)src, off, UNPK_B16); +} + +// Load 128 b16 values densely (NORM mode). All 128 halfword positions +// [0, 1, 2, ..., 127] in the register are filled with valid data. +// This is the load used by ComputeF16, which stays entirely in fp16 +// arithmetic and does not widen to fp32 (no vcvt call needed). +template +ROPE_SIMD_FN void vlds_norm_b16(F16Dst &dst, U16Src src, int32_t off) +{ + vlds((vector_f16 &)dst, (__ubuf__ half *)src, off, NORM); +} + +// Load 64 bf16 values from UB into a 128-lane vector register using +// UNPK_B16 mode. Same even-lane-only placement as vlds_unpk_b16 above. +template +ROPE_SIMD_FN void vlds_unpk_b16_bf16(BF16Dst &dst, U16Src src, int32_t off) +{ + vlds((vector_bf16 &)dst, (__ubuf__ bfloat16_t *)src, off, UNPK_B16); +} + +// Extract fp32 values from ODD halfword positions [1, 3, 5, ..., 127]. +// UNUSED by this kernel: RoPE loads b16 data via UNPK_B16, which places +// valid elements at EVEN positions only, so PART_EVEN covers all 64 +// elements. PART_ODD would extract zero/padding from the odd lanes. +// +// Kept in the shim for completeness. A kernel that fills all 128 b16 +// lanes (e.g. via DINTLV_B16 x2 + vintlv x2 double-width loads, as in +// mx_quant) needs BOTH vcvt_even and vcvt_odd to widen the full register. +template +ROPE_SIMD_FN void vcvt_fp16_to_fp32_odd(F32Dst &dst, F16Src src, MaskReg mask) +{ + vcvt((vector_f32 &)dst, (vector_f16 &)src, mask, PART_ODD, MODE_ZEROING); +} + +// Extract fp32 values from EVEN halfword positions [0, 2, 4, ..., 126]. +// This is the widen half of the UNPK_B16 → PART_EVEN pair: after +// vlds_unpk_b16 places 64 valid b16 elements at even positions, +// vcvt_even recovers all 64 into a fp32 vector register. +template +ROPE_SIMD_FN void vcvt_fp16_to_fp32_even(F32Dst &dst, F16Src src, MaskReg mask) +{ + vcvt((vector_f32 &)dst, (vector_f16 &)src, mask, PART_EVEN, MODE_ZEROING); +} + +// Extract fp32 values from EVEN halfword positions of a bf16 register. +// Pairs with vlds_unpk_b16_bf16 (UNPK_B16 load). Same even-only +// semantics as vcvt_fp16_to_fp32_even; no _odd counterpart needed. +template +ROPE_SIMD_FN void vcvt_bf16_to_fp32_even(F32Dst &dst, BF16Src src, MaskReg mask) +{ + vcvt((vector_f32 &)dst, (vector_bf16 &)src, mask, PART_EVEN, MODE_ZEROING); +} + +// --- Narrowing conversions: fp32 → fp16 / fp32 → bf16 --- +// +// Both use ROUND_R (round-to-nearest-even) and RS_DISABLE (no +// saturation). PART_EVEN places the narrowed 16-bit result at EVEN +// halfword positions [0, 2, ..., 126] so that a later PK_B32 store +// can pack them densely into UB. +// +// After this call the output register has valid bf16/fp16 bits only at +// the even halfword positions; the odd lanes are zero. +// +// Note: `vcvt_fp32_to_fp16_narrow` is NOT USED in this kernel. +// ComputeF16 stays entirely in fp16 (no widening/narrowing at all) and +// ComputeF32 stays entirely in fp32. It is kept in the shim for +// completeness (useful when porting a fp32-accum → fp16-output kernel). +template +ROPE_SIMD_FN void vcvt_fp32_to_fp16_narrow(F16Dst &dst, F32Src src, MaskReg mask) +{ + vcvt((vector_f16 &)dst, (vector_f32 &)src, mask, ROUND_R, RS_DISABLE, PART_EVEN, MODE_ZEROING); +} + +template +ROPE_SIMD_FN void vcvt_f32_to_bf16_narrow(BF16Dst &dst, F32Src src, MaskReg mask) +{ + vcvt((vector_bf16 &)dst, (vector_f32 &)src, mask, ROUND_R, RS_DISABLE, PART_EVEN, MODE_ZEROING); +} + +// --- Arithmetic: element-wise vmul / vadd / vsub in fp16 or fp32 --- +// +// PTO IR: pto.vmul, pto.vadd, pto.vsub +// Hardware: 1-cycle throughput per instruction (fully pipelined). +// +// `mask` controls which lanes participate in the compute. +// `MODE_ZEROING` writes 0 to inactive lanes (vs. `MODE_MERGE` which +// keeps the prior register bits). For RoPE we always use `MODE_ZEROING`. +template +ROPE_SIMD_FN void vmul_f32(F32Dst &dst, F32SrcA a, F32SrcB b, MaskReg mask) +{ + vmul((vector_f32 &)dst, (vector_f32 &)a, (vector_f32 &)b, mask, MODE_ZEROING); +} + +template +ROPE_SIMD_FN void vmul_f16(F16Dst &dst, F16SrcA a, F16SrcB b, MaskReg mask) +{ + vmul((vector_f16 &)dst, (vector_f16 &)a, (vector_f16 &)b, mask, MODE_ZEROING); +} + +template +ROPE_SIMD_FN void vadd_f32(F32Dst &dst, F32SrcA a, F32SrcB b, MaskReg mask) +{ + vadd((vector_f32 &)dst, (vector_f32 &)a, (vector_f32 &)b, mask, MODE_ZEROING); +} + +template +ROPE_SIMD_FN void vadd_f16(F16Dst &dst, F16SrcA a, F16SrcB b, MaskReg mask) +{ + vadd((vector_f16 &)dst, (vector_f16 &)a, (vector_f16 &)b, mask, MODE_ZEROING); +} + +template +ROPE_SIMD_FN void vsub_f32(F32Dst &dst, F32SrcA a, F32SrcB b, MaskReg mask) +{ + vsub((vector_f32 &)dst, (vector_f32 &)a, (vector_f32 &)b, mask, MODE_ZEROING); +} + +template +ROPE_SIMD_FN void vsub_f16(F16Dst &dst, F16SrcA a, F16SrcB b, MaskReg mask) +{ + vsub((vector_f16 &)dst, (vector_f16 &)a, (vector_f16 &)b, mask, MODE_ZEROING); +} + +// --- In-register interleave / deinterleave --- +// +// PTO IR: pto.vintlv, pto.vdintlv +// +// vdintlv_x2(dst_even, dst_odd, s0, s1): +// Splits elements with EVEN and ODD indices from two source registers +// into two output registers: +// dst_even[i] = src[2i] (even indices: 0, 2, 4, ...) +// dst_odd[i] = src[2i + 1] (odd indices: 1, 3, 5, ...) +// Used by RoPE INTERLEAVE mode to separate the "real" from "imaginary" +// part of each (x[2k], x[2k+1]) complex pair. +// +// vintlv_x2(dst_low, dst_high, s0, s1): +// Merges two streams element-by-element: +// dst_low[i] = s0[i/2] if i even else s1[i/2] (interleaved first half) +// dst_high[i] = ... (second half, for cnt > lane width) +// In RoPE this is used to re-interleave (neg_odd, even) back into: +// [-x1, x0, -x3, x2, ...] — the "rotated partner" vector. +// +// The "_x2" suffix in these wrappers reflects that the underlying CCE +// `vintlv` / `vdintlv` intrinsics always produce CONSUMER two outputs. +template +ROPE_SIMD_FN void vdintlv_x2(Dst &d0, Dst &d1, Src s0, Src s1) +{ + vdintlv(d0, d1, s0, s1); +} + +template +ROPE_SIMD_FN void vintlv_x2(Dst &d0, Dst &d1, Src s0, Src s1) +{ + vintlv(d0, d1, s0, s1); +} + +// --- Broadcast: scalar → every lane of a vector register --- +// +// PTO IR: pto.vbr +// +// `vbr_f32(dst, val)` sets each of the 64 fp32 lanes to `val`. +// `vbr_f16(dst, val)` sets each of the 128 fp16 lanes to `val`. +// +// Used to materialise the `-1.0` constant for the odd-element negation +// in INTERLEAVE mode. A single broadcast is much more efficient than +// loading a pre-filled constant tensor from UB. +template +ROPE_SIMD_FN void vbr_f32(F32Dst &dst, float val) +{ + vbr((vector_f32 &)dst, val); +} + +template +ROPE_SIMD_FN void vbr_f16(F16Dst &dst, half val) +{ + vbr((vector_f16 &)dst, val); +} + +// --- Mask construction: plt_b32 / plt_b16 --- +// +// PTO IR: pto.plt_b32, pto.plt_b16 +// +// `plt_b32(cnt, POST_UPDATE)` sets the first `cnt` bits of a b32 mask +// (1 bit per fp32 lane, 64 lanes total), leaving the rest zero. +// `plt_b16(cnt, POST_UPDATE)` is the b16 variant (1 bit per fp16 lane, +// 128 lanes total). +// +// `POST_UPDATE` means that the predicate counter advances AFTER the +// current use — this is the standard choice for a one-shot partial +// mask (the counter doesn't carry state across calls in our code). +// +// Usage: `cnt = min(remaining_elements, VL)` for the last partial block +// of a D tile — otherwise full-block loads use cnt = VL directly. +ROPE_CCE_INTERNAL MaskReg make_mask(uint32_t cnt) +{ + return plt_b32(cnt, POST_UPDATE); +} + +ROPE_CCE_INTERNAL MaskReg make_mask_b16(uint32_t cnt) +{ + return plt_b16(cnt, POST_UPDATE); +} + +// --- Store operations: PK_B32 vs NORM_B16 vs NORM_B32 --- +// +// PTO IR: pto.vsts with the appropriate {dist = "..."}. +// +// Store modes: +// * PK_B32: PACK_B32 store. Used after `vcvt_*_to_*_even` +// conversions that placed 16-bit results at EVEN +// halfword positions [0, 2, ..., 126] of the register. +// PK_B32 extracts those even-halfword slots and writes +// them densely (N consecutive halfwords) to UB. +// The mask controls how many halfwords are written. +// * NORM_B16: Dense store of N fp16 elements from lanes [0 .. N). +// Used by ComputeF16 where each store writes the full +// 128-lane packed fp16 block. +// * NORM_B32: Dense store of N fp32 elements from lanes [0 .. N). +// Used by ComputeF32. +// +// `vsts_pk_b32` (fp16 variant): NOT USED in this kernel. +// ComputeF16 stores via `vsts_norm_b16` directly, and ComputeBf16 +// stores via `vsts_pk_b32_bf16`. Kept for completeness (useful when +// porting a fp32-accum → fp16-output kernel). +template +ROPE_SIMD_FN void vsts_pk_b32(F16Src src, U16Dst dst, int32_t off, MaskReg mask) +{ + vsts((vector_f16 &)src, (__ubuf__ half *)dst, off, PK_B32, mask); +} + +// `vsts_pk_b32_bf16`: USED by ComputeBf16 (narrowing fp32 → bf16 store). +template +ROPE_SIMD_FN void vsts_pk_b32_bf16(BF16Src src, U16Dst dst, int32_t off, MaskReg mask) +{ + vsts((vector_bf16 &)src, (__ubuf__ bfloat16_t *)dst, off, PK_B32, mask); +} + +template +ROPE_SIMD_FN void vsts_norm_b16(F16Src src, U16Dst dst, int32_t off, MaskReg mask) +{ + vsts((vector_f16 &)src, (__ubuf__ half *)dst, off, NORM_B16, mask); +} + +template +ROPE_SIMD_FN void vsts_norm_b32(F32Src src, U32Dst dst, int32_t off, MaskReg mask) +{ + vsts((vector_f32 &)src, (__ubuf__ float *)dst, off, NORM_B32, mask); +} + +// --- fp32 load: vlds_norm_b32 --- +// +// Load 64 fp32 elements densely from UB into a 64-lane fp32 vector +// register. Used by ComputeF32 (NORM load matches the NORM store it +// round-trips through). +// +// PTO IR: pto.vlds {dist = "DIST_NORM_B32"} +template +ROPE_SIMD_FN void vlds_norm_b32(F32Dst &dst, U32Src src, int32_t off) +{ + vlds((vector_f32 &)dst, (__ubuf__ float *)src, off, NORM); +} + +} + +} + +#endif diff --git a/test/kernel-test/kernels/rope/cycle_metrics.py b/test/kernel-test/kernels/rope/cycle_metrics.py new file mode 100644 index 0000000000..f4b5f0089d --- /dev/null +++ b/test/kernel-test/kernels/rope/cycle_metrics.py @@ -0,0 +1,68 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""Rope-specific cycle-report entrypoint built on shared kernel-test contracts.""" + +from __future__ import annotations + +import glob +import os + +from kernel_test.cycle_reporting import CycleReporterSpec, run_cycle_report + + +def default_cycle_out_dirs(sim_root: str | None = None) -> list[str]: + """Return kernel-test rope sim output dirs in canonical case order.""" + + root = sim_root or os.path.join( + os.path.dirname(__file__), + "..", + "..", + "sim_outputs", + "rope", + ) + try: + from .tile_config import DTYPES, MODES + + ordered = [ + os.path.join(root, backend, f"{dtype}_{mode}") + for backend in ("cce", "vmi", "mi") + for dtype in DTYPES + for mode in MODES + if os.path.isdir(os.path.join(root, backend, f"{dtype}_{mode}")) + ] + if ordered: + return ordered + except ImportError: + pass + return sorted( + path + for path in glob.glob(os.path.join(root, "*", "*")) + if os.path.isdir(path) + ) + + +CYCLE_REPORTER = CycleReporterSpec( + name="rope", + default_out_dirs=default_cycle_out_dirs, + missing_message="No rope cycle output dirs found. Run kernel-test/scripts/run_cycle.sh first.", +) + + +def get_cycle_reporter() -> CycleReporterSpec: + """Return the rope cycle reporter registration.""" + + return CYCLE_REPORTER + + +def main(argv: list[str] | None = None) -> int: + return run_cycle_report(CYCLE_REPORTER, argv) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/test/kernel-test/kernels/rope/mi/backend.py b/test/kernel-test/kernels/rope/mi/backend.py new file mode 100644 index 0000000000..82df2c265d --- /dev/null +++ b/test/kernel-test/kernels/rope/mi/backend.py @@ -0,0 +1,483 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""PTODSL MI backend for the rope kernel.""" + +import os +from pathlib import Path + +from ptodsl import pto + +from kernel_test.backends import ArtifactPlan, RunPurpose +from kernel_test.npu_runtime import ensure_runtime, stream_ptr, sync + +from ..runtime import RopeLaunchArgs, artifact_case_dir, prepare_launch_args +from ..tile_config import MAX_N, MAX_S, SIM_D + +_MI_ROOT = Path(__file__).resolve().parent +_GENERATED_DIR = _MI_ROOT.parent / "generated" +_MAX_XY_ROWS = MAX_S * MAX_N +_CS_ELEMS = MAX_S * SIM_D +_XY_ELEMS = _MAX_XY_ROWS * SIM_D +_UB_BASE_2B_COS = 0 +_UB_BASE_2B_SIN = _UB_BASE_2B_COS + _CS_ELEMS * 2 +_UB_BASE_2B_X = _UB_BASE_2B_SIN + _CS_ELEMS * 2 +_UB_BASE_2B_Y = _UB_BASE_2B_X + _XY_ELEMS * 2 +_UB_BASE_4B_COS = 0 +_UB_BASE_4B_SIN = _UB_BASE_4B_COS + _CS_ELEMS * 4 +_UB_BASE_4B_X = _UB_BASE_4B_SIN + _CS_ELEMS * 4 +_UB_BASE_4B_Y = _UB_BASE_4B_X + _XY_ELEMS * 4 + + +@pto.jit( + name="rope_mi_f16", + target="a5", + backend="vpto", + mode="explicit", + kernel_kind="vector", + insert_sync=False, +) +def rope_mi_f16( + x_gm: pto.ptr(pto.f16, "gm"), + cos_gm: pto.ptr(pto.f16, "gm"), + sin_gm: pto.ptr(pto.f16, "gm"), + y_gm: pto.ptr(pto.f16, "gm"), + s_count: pto.i32, + n_count: pto.i32, + *, + MODE: pto.const_expr = 0, +): + rows = s_count * n_count + row_bytes = SIM_D * 2 + cs_bytes = SIM_D * 2 + + cos_ptr = pto.castptr(pto.const(_UB_BASE_2B_COS, dtype=pto.ui64), pto.ptr(pto.f16, "ub")) + sin_ptr = pto.castptr(pto.const(_UB_BASE_2B_SIN, dtype=pto.ui64), pto.ptr(pto.f16, "ub")) + x_ptr = pto.castptr(pto.const(_UB_BASE_2B_X, dtype=pto.ui64), pto.ptr(pto.f16, "ub")) + y_ptr = pto.castptr(pto.const(_UB_BASE_2B_Y, dtype=pto.ui64), pto.ptr(pto.f16, "ub")) + + pto.mte_gm_ub(cos_gm, cos_ptr, 0, cs_bytes, nburst=(s_count, cs_bytes, cs_bytes)) + pto.mte_gm_ub(sin_gm, sin_ptr, 0, cs_bytes, nburst=(s_count, cs_bytes, cs_bytes)) + pto.mte_gm_ub(x_gm, x_ptr, 0, row_bytes, nburst=(rows, row_bytes, row_bytes)) + + pto.set_flag(pto.Pipe.MTE2, pto.Pipe.V, event_id=0) + pto.wait_flag(pto.Pipe.MTE2, pto.Pipe.V, event_id=0) + + half_mask = pto.pge_b16(pto.MaskPattern.VL32) + interleave_mask = pto.pge_b16(pto.MaskPattern.VL64) + x_s_step = n_count * SIM_D + + if MODE == 0: + for s in range(0, s_count, 1): + x_s_off = s * x_s_step + cs_off = s * SIM_D + cs_hi_off = cs_off + 32 + + cos_lo = pto.vlds(cos_ptr, cs_off) + cos_hi = pto.vlds(cos_ptr, cs_hi_off) + sin_lo = pto.vlds(sin_ptr, cs_off) + sin_hi = pto.vlds(sin_ptr, cs_hi_off) + + for n in range(0, n_count, 1): + row_off = x_s_off + n * SIM_D + x_hi_off = row_off + 32 + + x_lo = pto.vlds(x_ptr, row_off) + x_hi = pto.vlds(x_ptr, x_hi_off) + + y_lo = pto.vsub( + pto.vmul(cos_lo, x_lo, half_mask), + pto.vmul(sin_lo, x_hi, half_mask), + half_mask, + ) + y_hi = pto.vadd( + pto.vmul(cos_hi, x_hi, half_mask), + pto.vmul(sin_hi, x_lo, half_mask), + half_mask, + ) + + pto.vsts(y_lo, y_ptr, row_off, half_mask) + pto.vsts(y_hi, y_ptr, x_hi_off, half_mask) + else: + for s in range(0, s_count, 1): + x_s_off = s * x_s_step + cs_off = s * SIM_D + + cos_even, cos_odd = pto.vldsx2( + cos_ptr, + cs_off, + pto.DeinterleaveDist.DINTLV_B16, + ) + sin_even, sin_odd = pto.vldsx2( + sin_ptr, + cs_off, + pto.DeinterleaveDist.DINTLV_B16, + ) + + for n in range(0, n_count, 1): + row_off = x_s_off + n * SIM_D + x_even, x_odd = pto.vldsx2( + x_ptr, + row_off, + pto.DeinterleaveDist.DINTLV_B16, + ) + + y_even = pto.vsub( + pto.vmul(x_even, cos_even, half_mask), + pto.vmul(x_odd, sin_even, half_mask), + half_mask, + ) + y_odd = pto.vadd( + pto.vmul(x_odd, cos_odd, half_mask), + pto.vmul(x_even, sin_odd, half_mask), + half_mask, + ) + + pto.vstsx2( + y_even, + y_odd, + y_ptr, + row_off, + pto.InterleaveDist.INTLV_B16, + interleave_mask, + ) + + pto.set_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=0) + pto.wait_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=0) + pto.mte_ub_gm(y_ptr, y_gm, row_bytes, nburst=(rows, row_bytes, row_bytes)) + pto.pipe_barrier(pto.Pipe.ALL) + + +@pto.jit( + name="rope_mi_bf16", + target="a5", + backend="vpto", + mode="explicit", + kernel_kind="vector", + insert_sync=False, +) +def rope_mi_bf16( + x_gm: pto.ptr(pto.bf16, "gm"), + cos_gm: pto.ptr(pto.f16, "gm"), + sin_gm: pto.ptr(pto.f16, "gm"), + y_gm: pto.ptr(pto.bf16, "gm"), + s_count: pto.i32, + n_count: pto.i32, + *, + MODE: pto.const_expr = 0, +): + rows = s_count * n_count + row_bytes = SIM_D * 2 + cs_bytes = SIM_D * 2 + + cos_ptr = pto.castptr(pto.const(_UB_BASE_2B_COS, dtype=pto.ui64), pto.ptr(pto.f16, "ub")) + sin_ptr = pto.castptr(pto.const(_UB_BASE_2B_SIN, dtype=pto.ui64), pto.ptr(pto.f16, "ub")) + x_ptr = pto.castptr(pto.const(_UB_BASE_2B_X, dtype=pto.ui64), pto.ptr(pto.bf16, "ub")) + y_ptr = pto.castptr(pto.const(_UB_BASE_2B_Y, dtype=pto.ui64), pto.ptr(pto.bf16, "ub")) + + pto.mte_gm_ub(cos_gm, cos_ptr, 0, cs_bytes, nburst=(s_count, cs_bytes, cs_bytes)) + pto.mte_gm_ub(sin_gm, sin_ptr, 0, cs_bytes, nburst=(s_count, cs_bytes, cs_bytes)) + pto.mte_gm_ub(x_gm, x_ptr, 0, row_bytes, nburst=(rows, row_bytes, row_bytes)) + + pto.set_flag(pto.Pipe.MTE2, pto.Pipe.V, event_id=0) + pto.wait_flag(pto.Pipe.MTE2, pto.Pipe.V, event_id=0) + + mask16_all = pto.pset_b16(pto.MaskPattern.ALL) + half_mask = pto.pge_b32(pto.MaskPattern.VL32) + full_mask32 = pto.pset_b32(pto.MaskPattern.ALL) + x_s_step = n_count * SIM_D + + if MODE == 0: + for s in range(0, s_count, 1): + x_s_off = s * x_s_step + cs_off = s * SIM_D + cs_hi_off = cs_off + 32 + + cos_lo = pto.vcvt(pto.vlds(cos_ptr, cs_off, dist="UNPK_B16"), pto.f32, mask16_all, part="EVEN") + cos_hi = pto.vcvt(pto.vlds(cos_ptr, cs_hi_off, dist="UNPK_B16"), pto.f32, mask16_all, part="EVEN") + sin_lo = pto.vcvt(pto.vlds(sin_ptr, cs_off, dist="UNPK_B16"), pto.f32, mask16_all, part="EVEN") + sin_hi = pto.vcvt(pto.vlds(sin_ptr, cs_hi_off, dist="UNPK_B16"), pto.f32, mask16_all, part="EVEN") + + for n in range(0, n_count, 1): + row_off = x_s_off + n * SIM_D + x_hi_off = row_off + 32 + + x_lo = pto.vcvt(pto.vlds(x_ptr, row_off, dist="UNPK_B16"), pto.f32, mask16_all, part="EVEN") + x_hi = pto.vcvt(pto.vlds(x_ptr, x_hi_off, dist="UNPK_B16"), pto.f32, mask16_all, part="EVEN") + + y_lo = pto.vsub( + pto.vmul(cos_lo, x_lo, half_mask), + pto.vmul(sin_lo, x_hi, half_mask), + half_mask, + ) + y_hi = pto.vadd( + pto.vmul(cos_hi, x_hi, half_mask), + pto.vmul(sin_hi, x_lo, half_mask), + half_mask, + ) + + pto.vsts( + pto.vcvt(y_lo, pto.bf16, half_mask, rnd="R", sat="SAT", part="EVEN"), + y_ptr, + row_off, + half_mask, + dist=pto.VStoreDist.PK_B32, + ) + pto.vsts( + pto.vcvt(y_hi, pto.bf16, half_mask, rnd="R", sat="SAT", part="EVEN"), + y_ptr, + x_hi_off, + half_mask, + dist=pto.VStoreDist.PK_B32, + ) + else: + for s in range(0, s_count, 1): + x_s_off = s * x_s_step + cs_off = s * SIM_D + + cos = pto.vcvt(pto.vlds(cos_ptr, cs_off, dist="UNPK_B16"), pto.f32, mask16_all, part="EVEN") + sin = pto.vcvt(pto.vlds(sin_ptr, cs_off, dist="UNPK_B16"), pto.f32, mask16_all, part="EVEN") + + for n in range(0, n_count, 1): + row_off = x_s_off + n * SIM_D + x = pto.vcvt(pto.vlds(x_ptr, row_off, dist="UNPK_B16"), pto.f32, mask16_all, part="EVEN") + x_even, x_odd = pto.vdintlv(x, x) + rot, _ = pto.vintlv(pto.vneg(x_odd, half_mask), x_even) + y = pto.vadd( + pto.vmul(x, cos, full_mask32), + pto.vmul(rot, sin, full_mask32), + full_mask32, + ) + + pto.vsts( + pto.vcvt(y, pto.bf16, full_mask32, rnd="R", sat="SAT", part="EVEN"), + y_ptr, + row_off, + full_mask32, + dist=pto.VStoreDist.PK_B32, + ) + + pto.set_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=0) + pto.wait_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=0) + pto.mte_ub_gm(y_ptr, y_gm, row_bytes, nburst=(rows, row_bytes, row_bytes)) + pto.pipe_barrier(pto.Pipe.ALL) + + +@pto.jit( + name="rope_mi_f32", + target="a5", + backend="vpto", + mode="explicit", + kernel_kind="vector", + insert_sync=False, +) +def rope_mi_f32( + x_gm: pto.ptr(pto.f32, "gm"), + cos_gm: pto.ptr(pto.f32, "gm"), + sin_gm: pto.ptr(pto.f32, "gm"), + y_gm: pto.ptr(pto.f32, "gm"), + s_count: pto.i32, + n_count: pto.i32, + *, + MODE: pto.const_expr = 0, +): + rows = s_count * n_count + row_bytes = SIM_D * 4 + cs_bytes = SIM_D * 4 + + cos_ptr = pto.castptr(pto.const(_UB_BASE_4B_COS, dtype=pto.ui64), pto.ptr(pto.f32, "ub")) + sin_ptr = pto.castptr(pto.const(_UB_BASE_4B_SIN, dtype=pto.ui64), pto.ptr(pto.f32, "ub")) + x_ptr = pto.castptr(pto.const(_UB_BASE_4B_X, dtype=pto.ui64), pto.ptr(pto.f32, "ub")) + y_ptr = pto.castptr(pto.const(_UB_BASE_4B_Y, dtype=pto.ui64), pto.ptr(pto.f32, "ub")) + + pto.mte_gm_ub(cos_gm, cos_ptr, 0, cs_bytes, nburst=(s_count, cs_bytes, cs_bytes)) + pto.mte_gm_ub(sin_gm, sin_ptr, 0, cs_bytes, nburst=(s_count, cs_bytes, cs_bytes)) + pto.mte_gm_ub(x_gm, x_ptr, 0, row_bytes, nburst=(rows, row_bytes, row_bytes)) + + pto.set_flag(pto.Pipe.MTE2, pto.Pipe.V, event_id=0) + pto.wait_flag(pto.Pipe.MTE2, pto.Pipe.V, event_id=0) + + half_mask = pto.pge_b32(pto.MaskPattern.VL32) + interleave_mask = pto.pset_b32(pto.MaskPattern.ALL) + x_s_step = n_count * SIM_D + + if MODE == 0: + for s in range(0, s_count, 1): + x_s_off = s * x_s_step + cs_off = s * SIM_D + cs_hi_off = cs_off + 32 + + cos_lo = pto.vlds(cos_ptr, cs_off) + cos_hi = pto.vlds(cos_ptr, cs_hi_off) + sin_lo = pto.vlds(sin_ptr, cs_off) + sin_hi = pto.vlds(sin_ptr, cs_hi_off) + + for n in range(0, n_count, 1): + row_off = x_s_off + n * SIM_D + x_hi_off = row_off + 32 + + x_lo = pto.vlds(x_ptr, row_off) + x_hi = pto.vlds(x_ptr, x_hi_off) + + y_lo = pto.vsub( + pto.vmul(cos_lo, x_lo, half_mask), + pto.vmul(sin_lo, x_hi, half_mask), + half_mask, + ) + y_hi = pto.vadd( + pto.vmul(cos_hi, x_hi, half_mask), + pto.vmul(sin_hi, x_lo, half_mask), + half_mask, + ) + + pto.vsts(y_lo, y_ptr, row_off, half_mask) + pto.vsts(y_hi, y_ptr, x_hi_off, half_mask) + else: + for s in range(0, s_count, 1): + x_s_off = s * x_s_step + cs_off = s * SIM_D + + cos_even, cos_odd = pto.vldsx2( + cos_ptr, + cs_off, + pto.DeinterleaveDist.DINTLV_B32, + ) + sin_even, sin_odd = pto.vldsx2( + sin_ptr, + cs_off, + pto.DeinterleaveDist.DINTLV_B32, + ) + + for n in range(0, n_count, 1): + row_off = x_s_off + n * SIM_D + x_even, x_odd = pto.vldsx2( + x_ptr, + row_off, + pto.DeinterleaveDist.DINTLV_B32, + ) + + y_even = pto.vsub( + pto.vmul(x_even, cos_even, half_mask), + pto.vmul(x_odd, sin_even, half_mask), + half_mask, + ) + y_odd = pto.vadd( + pto.vmul(x_odd, cos_odd, half_mask), + pto.vmul(x_even, sin_odd, half_mask), + half_mask, + ) + + pto.vstsx2( + y_even, + y_odd, + y_ptr, + row_off, + pto.InterleaveDist.INTLV_B32, + interleave_mask, + ) + + pto.set_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=0) + pto.wait_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=0) + pto.mte_ub_gm(y_ptr, y_gm, row_bytes, nburst=(rows, row_bytes, row_bytes)) + pto.pipe_barrier(pto.Pipe.ALL) + + +_COMPILED: dict[tuple[str, int], object] = {} + + +def _kernel_for_dtype(dtype: str): + if dtype == "f16": + return "f16", rope_mi_f16 + if dtype == "bf16": + return "bf16", rope_mi_bf16 + if dtype == "f32": + return "f32", rope_mi_f32 + raise ValueError(f"unsupported mi dtype: {dtype}") + + +def _prepare(dtype: str, mode_value: int) -> object: + dtype_key, kernel = _kernel_for_dtype(dtype) + key = (dtype_key, mode_value) + compiled = _COMPILED.get(key) + if compiled is None: + compiled = kernel.compile(MODE=mode_value) + _COMPILED[key] = compiled + return compiled + + +def _launch(launch_args: RopeLaunchArgs): + compiled = _prepare(launch_args.dtype, launch_args.mode_value) + compiled[1, stream_ptr()]( + launch_args.x.data_ptr(), + launch_args.cos.data_ptr(), + launch_args.sin.data_ptr(), + launch_args.y.data_ptr(), + launch_args.s_count, + launch_args.n_count, + ) + sync() + return launch_args.y + + +def _build_artifact_plan(case: dict[str, object]) -> ArtifactPlan: + compiled = _prepare(case["dtype"], 0 if case["mode"] == "half" else 1) + case_dir = artifact_case_dir(_GENERATED_DIR, case, backend_name="mi") + return ArtifactPlan( + generated_dir=_GENERATED_DIR, + case_dir=case_dir, + mi_text=compiled.mlir_text(), + ) + + +def rope_f16(launch_args: RopeLaunchArgs): + """Launch the local rope f16 MI kernel.""" + + return _launch(launch_args) + + +def rope_bf16(launch_args: RopeLaunchArgs): + """Launch the local rope bf16 MI kernel.""" + + return _launch(launch_args) + + +def rope_f32(launch_args: RopeLaunchArgs): + """Launch the local rope f32 MI kernel.""" + + return _launch(launch_args) + + +class RopeMiBackend: + """Pure-PTODSL MI backend for rope.""" + + name = "mi" + _launchers = { + "f16": rope_f16, + "bf16": rope_bf16, + "f32": rope_f32, + } + + def is_supported(self, case: object, *, purpose: RunPurpose) -> tuple[bool, str | None]: + del purpose + supported = case["dtype"] in {"f16", "bf16", "f32"} and case["mode"] in {"half", "interleave"} + if supported: + return True, None + return False, "backend=mi not wired for this case" + + def launch(self, case: object, *, purpose: RunPurpose) -> object: + ensure_runtime("rope") + launch_args = prepare_launch_args(case, cycle=purpose == "cycle") + return self._launchers[launch_args.dtype](launch_args) + + def cache_tag(self) -> str: + backend_py = _MI_ROOT / "backend.py" + return f"mi:{backend_py}:{os.path.getmtime(backend_py):.0f}" + + def build_artifact_plan(self, case_id: str, case: object) -> ArtifactPlan: + del case_id + return _build_artifact_plan(case) diff --git a/test/kernel-test/kernels/rope/reference.py b/test/kernel-test/kernels/rope/reference.py new file mode 100644 index 0000000000..b3ed885aab --- /dev/null +++ b/test/kernel-test/kernels/rope/reference.py @@ -0,0 +1,92 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""CPU golden references for the rope kernel.""" + +from __future__ import annotations + +import numpy as np +import torch + +from .tile_config import DEFAULT_TILE, DTYPES, MODES, TileConfig + +SEED = 42 + + +def cpu_rotary_half(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + x1, x2 = torch.chunk(x, 2, -1) + x_new = torch.cat((-x2, x1), dim=-1) + cos_b = cos.unsqueeze(1) + sin_b = sin.unsqueeze(1) + return cos_b * x + sin_b * x_new + + +def cpu_rotary_interleave(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + x1 = x[..., ::2] + x2 = x[..., 1::2] + x_new = torch.stack((-x2, x1), dim=-1).reshape(x.shape) + cos_b = cos.unsqueeze(1) + sin_b = sin.unsqueeze(1) + return x * cos_b + x_new * sin_b + + +def _to_numpy(t: torch.Tensor) -> np.ndarray: + if t.dtype == torch.bfloat16: + return t.float().numpy() + return t.numpy() + + +def _torch_dtype(dtype: str) -> torch.dtype: + if dtype == "f16": + return torch.float16 + if dtype == "bf16": + return torch.bfloat16 + if dtype == "f32": + return torch.float32 + raise ValueError(f"unknown dtype: {dtype}") + + +def generate_case(mode: str, dtype: str, tile: TileConfig | None = None) -> dict: + if mode not in MODES: + raise ValueError(f"unknown mode: {mode}") + if dtype not in DTYPES: + raise ValueError(f"unknown dtype: {dtype}") + tile = tile or DEFAULT_TILE + + x_dtype = _torch_dtype(dtype) + cs_dtype = torch.float16 if dtype == "bf16" else x_dtype + + torch.manual_seed(SEED) + x = torch.randn(tile.x_shape, dtype=x_dtype) + cos = torch.randn(tile.cs_shape, dtype=cs_dtype).abs().clamp(0.1, 0.9) + sin = torch.randn(tile.cs_shape, dtype=cs_dtype).abs().clamp(0.1, 0.9) + + if mode == "half": + y = cpu_rotary_half(x, cos, sin) + else: + y = cpu_rotary_interleave(x, cos, sin) + + return { + "mode": mode, + "dtype": dtype, + "tile": tile, + "x": _to_numpy(x), + "cos": _to_numpy(cos), + "sin": _to_numpy(sin), + "y": _to_numpy(y), + "params": np.array([tile.s, tile.n], dtype=np.int32), + } + + +def generate_all(tile: TileConfig | None = None) -> dict[str, dict]: + tile = tile or DEFAULT_TILE + cases: dict[str, dict] = {} + for dtype in DTYPES: + for mode in MODES: + cases[f"{dtype}_{mode}"] = generate_case(mode, dtype, tile=tile) + return cases diff --git a/test/kernel-test/kernels/rope/runtime.py b/test/kernel-test/kernels/rope/runtime.py new file mode 100644 index 0000000000..7cdf0281f3 --- /dev/null +++ b/test/kernel-test/kernels/rope/runtime.py @@ -0,0 +1,106 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""Runtime data preparation for the rope kernel.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +from kernel_test.npu_runtime import device_str, empty_npu + +from .tile_config import sim_fn_name + +if TYPE_CHECKING: + import numpy as np + import torch + + +@dataclass(frozen=True) +class RopeLaunchArgs: + """Prepared runtime arguments shared by all rope backends.""" + + dtype: str + mode: str + mode_value: int + fn_name: str + x: torch.Tensor + cos: torch.Tensor + sin: torch.Tensor + y: torch.Tensor + params: torch.Tensor + s_count: int + n_count: int + + +def torch_dtype(dtype: str) -> torch.dtype: + import torch + + if dtype == "f16": + return torch.float16 + if dtype == "bf16": + return torch.bfloat16 + if dtype == "f32": + return torch.float32 + raise ValueError(f"unknown dtype: {dtype}") + + +def params_tensor(params: np.ndarray) -> torch.Tensor: + import numpy as np + import torch + + return torch.from_numpy(np.asarray(params, dtype=np.int32)).to(device_str()) + + +def artifact_case_dir_name(case: dict[str, object]) -> str: + """Build a stable per-case artifact directory name for rope.""" + + tile = case["tile"] + return f"{case['dtype']}_{case['mode']}_s{tile.s}_n{tile.n}" + + +def artifact_case_dir(root: Path, case: dict[str, object], *, backend_name: str) -> Path: + """Return the backend-specific rope artifact directory for one case.""" + + return root / backend_name / artifact_case_dir_name(case) + + +def prepare_launch_args(case: dict, *, cycle: bool = False) -> RopeLaunchArgs: + """Convert one rope case into prepared device tensors and launch metadata.""" + + import torch + + dtype = case["dtype"] + x_dtype = torch_dtype(dtype) + cs_dtype = torch_dtype("f16") if dtype == "bf16" else x_dtype + mode = case["mode"] + mode_value = 0 if mode == "half" else 1 + dev = device_str() + + x = torch.from_numpy(case["x"]).to(x_dtype).to(dev) + cos = torch.from_numpy(case["cos"]).to(cs_dtype).to(dev) + sin = torch.from_numpy(case["sin"]).to(cs_dtype).to(dev) + y = empty_npu(case["y"].shape, x_dtype) + params = params_tensor(case["params"]) + s_count, n_count = [int(value) for value in case["params"]] + + return RopeLaunchArgs( + dtype=dtype, + mode=mode, + mode_value=mode_value, + fn_name=sim_fn_name(mode, dtype, cycle=cycle), + x=x, + cos=cos, + sin=sin, + y=y, + params=params, + s_count=s_count, + n_count=n_count, + ) diff --git a/test/kernel-test/kernels/rope/spec.py b/test/kernel-test/kernels/rope/spec.py new file mode 100644 index 0000000000..f31152bbaf --- /dev/null +++ b/test/kernel-test/kernels/rope/spec.py @@ -0,0 +1,93 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""Case listing and verification for the rope kernel.""" + +from __future__ import annotations + +import os + +from kernel_test.results import CaseResult + +from .tile_config import DEFAULT_TILE, DTYPES, MODES, TOLERANCE, TileConfig + + +def tile_from_env() -> TileConfig | None: + """Build a rope tile override from environment variables if present.""" + + s = os.environ.get("ROPE_VF_S") + n = os.environ.get("ROPE_VF_N") + if s is None or n is None: + return None + return TileConfig(name="env", s=int(s), n=int(n)) + + +def default_tile_from_env() -> TileConfig: + """Return the configured tile, falling back to the default rope tile.""" + + return tile_from_env() or DEFAULT_TILE + + +def _lightweight_cases(tile: TileConfig) -> dict[str, object]: + return { + f"{dtype}_{mode}": {"mode": mode, "dtype": dtype, "tile": tile} + for dtype in DTYPES + for mode in MODES + } + + +def list_cases(workflow: str) -> dict[str, object]: + """Return the rope case matrix for the requested workflow.""" + + if workflow not in {"correctness", "cycle"}: + raise ValueError(f"unsupported rope workflow: {workflow}") + + tile = default_tile_from_env() + try: + from .reference import generate_all + except ModuleNotFoundError as exc: + if exc.name not in {"numpy", "torch"}: + raise + return _lightweight_cases(tile) + + return generate_all(tile=tile) + + +def verify_case(case_id: str, case: object, output: object) -> CaseResult: + """Verify a rope backend output against the generated golden tensors.""" + + y_host = output.cpu() + if case["dtype"] == "bf16": + got = y_host.float().numpy() + else: + got = y_host.numpy() + + max_diff = float(abs(got.astype("float32") - case["y"].astype("float32")).max()) + tile = case["tile"] + message = ( + f"{case['dtype']}/{case['mode']}: maxDiff={max_diff:.6f} " + f"tile=s{tile.s}_n{tile.n}" + ) + return CaseResult( + ok=max_diff < TOLERANCE[case["dtype"]], + message=message, + ) + + +def cycle_fields(case_id: str, case: object, backend: object) -> dict[str, object]: + """Build stable cycle marker fields for one rope case.""" + + del backend, case_id + tile = case["tile"] + return { + "mode": case["mode"], + "dtype": case["dtype"], + "s": tile.s, + "n": tile.n, + "vf_inner": tile.vf_inner_iters, + } diff --git a/test/kernel-test/kernels/rope/tile_config.py b/test/kernel-test/kernels/rope/tile_config.py new file mode 100644 index 0000000000..ea8b47b1c9 --- /dev/null +++ b/test/kernel-test/kernels/rope/tile_config.py @@ -0,0 +1,67 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""Shared tile configuration for the rope kernel.""" + +from __future__ import annotations + +from dataclasses import dataclass + +SIM_B = 1 +SIM_D = 64 +SIM_D_ALIGN = 64 + +DEFAULT_S = 15 +DEFAULT_N = 32 + +MAX_S = 15 +MAX_N = 32 + +MODES = ("half", "interleave") +DTYPES = ("f16", "bf16", "f32") + +TOLERANCE = { + "f16": 0.01, + "bf16": 0.07, + "f32": 1e-5, +} + + +@dataclass(frozen=True) +class TileConfig: + name: str + s: int + n: int + + @property + def vf_inner_iters(self) -> int: + return self.s * self.n + + @property + def x_shape(self) -> tuple[int, int, int]: + return (self.s, self.n, SIM_D) + + @property + def cs_shape(self) -> tuple[int, int]: + return (self.s, SIM_D) + + +WALLTIME_CONFIGS: tuple[TileConfig, ...] = ( + TileConfig("tiny", s=1, n=2), + TileConfig("n4", s=15, n=4), + TileConfig("n8", s=15, n=8), + TileConfig("n16", s=15, n=16), + TileConfig("prod", s=DEFAULT_S, n=DEFAULT_N), +) + +DEFAULT_TILE = TileConfig("prod", s=DEFAULT_S, n=DEFAULT_N) + + +def sim_fn_name(mode: str, dtype: str, cycle: bool = False) -> str: + prefix = "call_rope_cce_cycle" if cycle else "call_rope_cce_sim" + return f"{prefix}_{mode}_{dtype}" diff --git a/test/kernel-test/kernels/rope/vmi/backend.py b/test/kernel-test/kernels/rope/vmi/backend.py new file mode 100644 index 0000000000..08168cabdf --- /dev/null +++ b/test/kernel-test/kernels/rope/vmi/backend.py @@ -0,0 +1,448 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""VMI backend for the rope kernel.""" + +import os +from pathlib import Path + +from ptodsl import pto + +from kernel_test.backends import ArtifactPlan, RunPurpose +from kernel_test.npu_runtime import ensure_runtime, stream_ptr, sync + +from ..runtime import RopeLaunchArgs, artifact_case_dir, prepare_launch_args +from ..tile_config import MAX_N, MAX_S, SIM_D + +_VMI_ROOT = Path(__file__).resolve().parent +_GENERATED_DIR = _VMI_ROOT.parent / "generated" +_MAX_XY_ROWS = MAX_S * MAX_N +_CS_ELEMS = MAX_S * SIM_D +_XY_ELEMS = _MAX_XY_ROWS * SIM_D +_UB_BASE_2B_COS = 0 +_UB_BASE_2B_SIN = _UB_BASE_2B_COS + _CS_ELEMS * 2 +_UB_BASE_2B_X = _UB_BASE_2B_SIN + _CS_ELEMS * 2 +_UB_BASE_2B_Y = _UB_BASE_2B_X + _XY_ELEMS * 2 +_UB_BASE_4B_COS = 0 +_UB_BASE_4B_SIN = _UB_BASE_4B_COS + _CS_ELEMS * 4 +_UB_BASE_4B_X = _UB_BASE_4B_SIN + _CS_ELEMS * 4 +_UB_BASE_4B_Y = _UB_BASE_4B_X + _XY_ELEMS * 4 + + +@pto.jit( + name="rope_vmi_f16", + target="a5", + backend="vpto", + mode="explicit", + kernel_kind="vector", + insert_sync=False, +) +def rope_vmi_f16( + x_gm: pto.ptr(pto.f16, "gm"), + cos_gm: pto.ptr(pto.f16, "gm"), + sin_gm: pto.ptr(pto.f16, "gm"), + y_gm: pto.ptr(pto.f16, "gm"), + s_count: pto.i32, + n_count: pto.i32, + *, + MODE: pto.const_expr = 0, +): + rows = s_count * n_count + row_bytes = SIM_D * 2 + cs_bytes = SIM_D * 2 + + cos_ptr = pto.castptr(pto.const(_UB_BASE_2B_COS, dtype=pto.ui64), pto.ptr(pto.f16, "ub")) + sin_ptr = pto.castptr(pto.const(_UB_BASE_2B_SIN, dtype=pto.ui64), pto.ptr(pto.f16, "ub")) + x_ptr = pto.castptr(pto.const(_UB_BASE_2B_X, dtype=pto.ui64), pto.ptr(pto.f16, "ub")) + y_ptr = pto.castptr(pto.const(_UB_BASE_2B_Y, dtype=pto.ui64), pto.ptr(pto.f16, "ub")) + + pto.mte_gm_ub(cos_gm, cos_ptr, 0, cs_bytes, nburst=(s_count, cs_bytes, cs_bytes)) + pto.mte_gm_ub(sin_gm, sin_ptr, 0, cs_bytes, nburst=(s_count, cs_bytes, cs_bytes)) + pto.mte_gm_ub(x_gm, x_ptr, 0, row_bytes, nburst=(rows, row_bytes, row_bytes)) + + pto.set_flag(pto.Pipe.MTE2, pto.Pipe.V, event_id=0) + pto.wait_flag(pto.Pipe.MTE2, pto.Pipe.V, event_id=0) + + half_mask = pto.vmi.create_mask(32, size=64) + full_mask = pto.vmi.create_mask(64, size=64) + x_s_step = n_count * SIM_D + + if MODE == 0: + for s in range(0, s_count, 1): + x_s_off = s * x_s_step + cs_off = s * SIM_D + cs_hi_off = cs_off + 32 + + cos_lo = pto.vmi.vload(cos_ptr, cs_off, size=64) + cos_hi = pto.vmi.vload(cos_ptr, cs_hi_off, size=64) + sin_lo = pto.vmi.vload(sin_ptr, cs_off, size=64) + sin_hi = pto.vmi.vload(sin_ptr, cs_hi_off, size=64) + + for n in range(0, n_count, 1): + row_off = x_s_off + n * SIM_D + x_hi_off = row_off + 32 + + x_lo = pto.vmi.vload(x_ptr, row_off, size=64) + x_hi = pto.vmi.vload(x_ptr, x_hi_off, size=64) + + y_lo = pto.vmi.vsub( + pto.vmi.vmul(cos_lo, x_lo, half_mask), + pto.vmi.vmul(sin_lo, x_hi, half_mask), + half_mask, + ) + y_hi = pto.vmi.vadd( + pto.vmi.vmul(cos_hi, x_hi, half_mask), + pto.vmi.vmul(sin_hi, x_lo, half_mask), + half_mask, + ) + + pto.vmi.vstore(y_lo, y_ptr, row_off, half_mask) + pto.vmi.vstore(y_hi, y_ptr, x_hi_off, half_mask) + else: + for s in range(0, s_count, 1): + x_s_off = s * x_s_step + cs_off = s * SIM_D + + cos = pto.vmi.vload(cos_ptr, cs_off, size=64) + sin = pto.vmi.vload(sin_ptr, cs_off, size=64) + cos_even, cos_odd = pto.vmi.vdintlv(cos, cos, full_mask) + sin_even, sin_odd = pto.vmi.vdintlv(sin, sin, full_mask) + + for n in range(0, n_count, 1): + row_off = x_s_off + n * SIM_D + x = pto.vmi.vload(x_ptr, row_off, size=64) + x_even, x_odd = pto.vmi.vdintlv(x, x, full_mask) + + y_even = pto.vmi.vsub( + pto.vmi.vmul(x_even, cos_even, half_mask), + pto.vmi.vmul(x_odd, sin_even, half_mask), + half_mask, + ) + y_odd = pto.vmi.vadd( + pto.vmi.vmul(x_odd, cos_odd, half_mask), + pto.vmi.vmul(x_even, sin_odd, half_mask), + half_mask, + ) + + y, _ = pto.vmi.vintlv(y_even, y_odd, full_mask) + pto.vmi.vstore(y, y_ptr, row_off, full_mask) + + pto.set_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=0) + pto.wait_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=0) + pto.mte_ub_gm(y_ptr, y_gm, row_bytes, nburst=(rows, row_bytes, row_bytes)) + pto.pipe_barrier(pto.Pipe.ALL) + + +@pto.jit( + name="rope_vmi_bf16", + target="a5", + backend="vpto", + mode="explicit", + kernel_kind="vector", + insert_sync=False, +) +def rope_vmi_bf16( + x_gm: pto.ptr(pto.bf16, "gm"), + cos_gm: pto.ptr(pto.f16, "gm"), + sin_gm: pto.ptr(pto.f16, "gm"), + y_gm: pto.ptr(pto.bf16, "gm"), + s_count: pto.i32, + n_count: pto.i32, + *, + MODE: pto.const_expr = 0, +): + rows = s_count * n_count + row_bytes = SIM_D * 2 + cs_bytes = SIM_D * 2 + + cos_ptr = pto.castptr(pto.const(_UB_BASE_2B_COS, dtype=pto.ui64), pto.ptr(pto.f16, "ub")) + sin_ptr = pto.castptr(pto.const(_UB_BASE_2B_SIN, dtype=pto.ui64), pto.ptr(pto.f16, "ub")) + x_ptr = pto.castptr(pto.const(_UB_BASE_2B_X, dtype=pto.ui64), pto.ptr(pto.bf16, "ub")) + y_ptr = pto.castptr(pto.const(_UB_BASE_2B_Y, dtype=pto.ui64), pto.ptr(pto.bf16, "ub")) + + pto.mte_gm_ub(cos_gm, cos_ptr, 0, cs_bytes, nburst=(s_count, cs_bytes, cs_bytes)) + pto.mte_gm_ub(sin_gm, sin_ptr, 0, cs_bytes, nburst=(s_count, cs_bytes, cs_bytes)) + pto.mte_gm_ub(x_gm, x_ptr, 0, row_bytes, nburst=(rows, row_bytes, row_bytes)) + + pto.set_flag(pto.Pipe.MTE2, pto.Pipe.V, event_id=0) + pto.wait_flag(pto.Pipe.MTE2, pto.Pipe.V, event_id=0) + + half_mask = pto.vmi.create_mask(32, size=64) + full_mask = pto.vmi.create_mask(64, size=64) + x_s_step = n_count * SIM_D + + if MODE == 0: + for s in range(0, s_count, 1): + x_s_off = s * x_s_step + cs_off = s * SIM_D + cs_hi_off = cs_off + 32 + + cos_lo = pto.vmi.vcvt(pto.vmi.vload(cos_ptr, cs_off, size=64), pto.f32) + cos_hi = pto.vmi.vcvt(pto.vmi.vload(cos_ptr, cs_hi_off, size=64), pto.f32) + sin_lo = pto.vmi.vcvt(pto.vmi.vload(sin_ptr, cs_off, size=64), pto.f32) + sin_hi = pto.vmi.vcvt(pto.vmi.vload(sin_ptr, cs_hi_off, size=64), pto.f32) + + for n in range(0, n_count, 1): + row_off = x_s_off + n * SIM_D + x_hi_off = row_off + 32 + + x_lo = pto.vmi.vcvt(pto.vmi.vload(x_ptr, row_off, size=64), pto.f32) + x_hi = pto.vmi.vcvt(pto.vmi.vload(x_ptr, x_hi_off, size=64), pto.f32) + + y_lo_f32 = pto.vmi.vsub( + pto.vmi.vmul(cos_lo, x_lo, half_mask), + pto.vmi.vmul(sin_lo, x_hi, half_mask), + half_mask, + ) + y_hi_f32 = pto.vmi.vadd( + pto.vmi.vmul(cos_hi, x_hi, half_mask), + pto.vmi.vmul(sin_hi, x_lo, half_mask), + half_mask, + ) + + y_lo = pto.vmi.vcvt(y_lo_f32, pto.bf16) + y_hi = pto.vmi.vcvt(y_hi_f32, pto.bf16) + + pto.vmi.vstore(y_lo, y_ptr, row_off, half_mask) + pto.vmi.vstore(y_hi, y_ptr, x_hi_off, half_mask) + else: + for s in range(0, s_count, 1): + x_s_off = s * x_s_step + cs_off = s * SIM_D + + cos = pto.vmi.vcvt(pto.vmi.vload(cos_ptr, cs_off, size=64), pto.f32) + sin = pto.vmi.vcvt(pto.vmi.vload(sin_ptr, cs_off, size=64), pto.f32) + cos_even, cos_odd = pto.vmi.vdintlv(cos, cos, half_mask) + sin_even, sin_odd = pto.vmi.vdintlv(sin, sin, half_mask) + + for n in range(0, n_count, 1): + row_off = x_s_off + n * SIM_D + x = pto.vmi.vcvt(pto.vmi.vload(x_ptr, row_off, size=64), pto.f32) + x_even, x_odd = pto.vmi.vdintlv(x, x, half_mask) + + y_even_f32 = pto.vmi.vsub( + pto.vmi.vmul(x_even, cos_even, half_mask), + pto.vmi.vmul(x_odd, sin_even, half_mask), + half_mask, + ) + y_odd_f32 = pto.vmi.vadd( + pto.vmi.vmul(x_odd, cos_odd, half_mask), + pto.vmi.vmul(x_even, sin_odd, half_mask), + half_mask, + ) + + y_even = pto.vmi.vcvt(y_even_f32, pto.bf16) + y_odd = pto.vmi.vcvt(y_odd_f32, pto.bf16) + + y, _ = pto.vmi.vintlv(y_even, y_odd, half_mask) + pto.vmi.vstore(y, y_ptr, row_off, full_mask) + + pto.set_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=0) + pto.wait_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=0) + pto.mte_ub_gm(y_ptr, y_gm, row_bytes, nburst=(rows, row_bytes, row_bytes)) + pto.pipe_barrier(pto.Pipe.ALL) + + +@pto.jit( + name="rope_vmi_f32", + target="a5", + backend="vpto", + mode="explicit", + kernel_kind="vector", + insert_sync=False, +) +def rope_vmi_f32( + x_gm: pto.ptr(pto.f32, "gm"), + cos_gm: pto.ptr(pto.f32, "gm"), + sin_gm: pto.ptr(pto.f32, "gm"), + y_gm: pto.ptr(pto.f32, "gm"), + s_count: pto.i32, + n_count: pto.i32, + *, + MODE: pto.const_expr = 0, +): + rows = s_count * n_count + row_bytes = SIM_D * 4 + cs_bytes = SIM_D * 4 + + cos_ptr = pto.castptr(pto.const(_UB_BASE_4B_COS, dtype=pto.ui64), pto.ptr(pto.f32, "ub")) + sin_ptr = pto.castptr(pto.const(_UB_BASE_4B_SIN, dtype=pto.ui64), pto.ptr(pto.f32, "ub")) + x_ptr = pto.castptr(pto.const(_UB_BASE_4B_X, dtype=pto.ui64), pto.ptr(pto.f32, "ub")) + y_ptr = pto.castptr(pto.const(_UB_BASE_4B_Y, dtype=pto.ui64), pto.ptr(pto.f32, "ub")) + + pto.mte_gm_ub(cos_gm, cos_ptr, 0, cs_bytes, nburst=(s_count, cs_bytes, cs_bytes)) + pto.mte_gm_ub(sin_gm, sin_ptr, 0, cs_bytes, nburst=(s_count, cs_bytes, cs_bytes)) + pto.mte_gm_ub(x_gm, x_ptr, 0, row_bytes, nburst=(rows, row_bytes, row_bytes)) + + pto.set_flag(pto.Pipe.MTE2, pto.Pipe.V, event_id=0) + pto.wait_flag(pto.Pipe.MTE2, pto.Pipe.V, event_id=0) + + half_mask = pto.vmi.create_mask(32, size=64) + full_mask = pto.vmi.create_mask(64, size=64) + x_s_step = n_count * SIM_D + + if MODE == 0: + for s in range(0, s_count, 1): + x_s_off = s * x_s_step + cs_off = s * SIM_D + cs_hi_off = cs_off + 32 + + cos_lo = pto.vmi.vload(cos_ptr, cs_off, size=64) + cos_hi = pto.vmi.vload(cos_ptr, cs_hi_off, size=64) + sin_lo = pto.vmi.vload(sin_ptr, cs_off, size=64) + sin_hi = pto.vmi.vload(sin_ptr, cs_hi_off, size=64) + + for n in range(0, n_count, 1): + row_off = x_s_off + n * SIM_D + x_hi_off = row_off + 32 + + x_lo = pto.vmi.vload(x_ptr, row_off, size=64) + x_hi = pto.vmi.vload(x_ptr, x_hi_off, size=64) + + y_lo = pto.vmi.vsub( + pto.vmi.vmul(cos_lo, x_lo, half_mask), + pto.vmi.vmul(sin_lo, x_hi, half_mask), + half_mask, + ) + y_hi = pto.vmi.vadd( + pto.vmi.vmul(cos_hi, x_hi, half_mask), + pto.vmi.vmul(sin_hi, x_lo, half_mask), + half_mask, + ) + + pto.vmi.vstore(y_lo, y_ptr, row_off, half_mask) + pto.vmi.vstore(y_hi, y_ptr, x_hi_off, half_mask) + else: + for s in range(0, s_count, 1): + x_s_off = s * x_s_step + cs_off = s * SIM_D + + cos = pto.vmi.vload(cos_ptr, cs_off, size=64) + sin = pto.vmi.vload(sin_ptr, cs_off, size=64) + cos_even, cos_odd = pto.vmi.vdintlv(cos, cos, half_mask) + sin_even, sin_odd = pto.vmi.vdintlv(sin, sin, half_mask) + + for n in range(0, n_count, 1): + row_off = x_s_off + n * SIM_D + x = pto.vmi.vload(x_ptr, row_off, size=64) + x_even, x_odd = pto.vmi.vdintlv(x, x, half_mask) + + y_even = pto.vmi.vsub( + pto.vmi.vmul(x_even, cos_even, half_mask), + pto.vmi.vmul(x_odd, sin_even, half_mask), + half_mask, + ) + y_odd = pto.vmi.vadd( + pto.vmi.vmul(x_odd, cos_odd, half_mask), + pto.vmi.vmul(x_even, sin_odd, half_mask), + half_mask, + ) + + y, _ = pto.vmi.vintlv(y_even, y_odd, half_mask) + pto.vmi.vstore(y, y_ptr, row_off, full_mask) + + pto.set_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=0) + pto.wait_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=0) + pto.mte_ub_gm(y_ptr, y_gm, row_bytes, nburst=(rows, row_bytes, row_bytes)) + pto.pipe_barrier(pto.Pipe.ALL) + + +_COMPILED: dict[tuple[str, int], object] = {} + + +def _kernel_for_dtype(dtype: str): + if dtype == "f16": + return "f16", rope_vmi_f16 + if dtype == "bf16": + return "bf16", rope_vmi_bf16 + if dtype == "f32": + return "f32", rope_vmi_f32 + raise ValueError(f"unsupported vmi dtype: {dtype}") + + +def _prepare(dtype: str, mode_value: int) -> object: + key, kernel = _kernel_for_dtype(dtype) + cache_key = (key, mode_value) + compiled = _COMPILED.get(cache_key) + if compiled is None: + compiled = kernel.compile(MODE=mode_value) + _COMPILED[cache_key] = compiled + return compiled + + +def _launch(launch_args: RopeLaunchArgs): + compiled = _prepare(launch_args.dtype, launch_args.mode_value) + compiled[1, stream_ptr()]( + launch_args.x.data_ptr(), + launch_args.cos.data_ptr(), + launch_args.sin.data_ptr(), + launch_args.y.data_ptr(), + launch_args.s_count, + launch_args.n_count, + ) + sync() + return launch_args.y + + +def _build_artifact_plan(case: dict[str, object]) -> ArtifactPlan: + compiled = _prepare(case["dtype"], 0 if case["mode"] == "half" else 1) + case_dir = artifact_case_dir(_GENERATED_DIR, case, backend_name="vmi") + return ArtifactPlan( + generated_dir=_GENERATED_DIR, + case_dir=case_dir, + vmi_text=compiled.mlir_text(), + ) + + +def rope_f16(launch_args: RopeLaunchArgs): + """Launch the local rope f16 VMI kernel.""" + + return _launch(launch_args) + + +def rope_bf16(launch_args: RopeLaunchArgs): + """Launch the local rope bf16 VMI kernel.""" + + return _launch(launch_args) + + +def rope_f32(launch_args: RopeLaunchArgs): + """Launch the local rope f32 VMI kernel.""" + + return _launch(launch_args) + + +class RopeVmiBackend: + """Pure-PTODSL VMI backend for rope.""" + + name = "vmi" + _launchers = { + "f16": rope_f16, + "bf16": rope_bf16, + "f32": rope_f32, + } + + def is_supported(self, case: object, *, purpose: RunPurpose) -> tuple[bool, str | None]: + del purpose + supported = case["dtype"] in {"f16", "bf16", "f32"} and case["mode"] in {"half", "interleave"} + if supported: + return True, None + return False, "backend=vmi not wired for this case" + + def launch(self, case: object, *, purpose: RunPurpose) -> object: + ensure_runtime("rope") + launch_args = prepare_launch_args(case, cycle=purpose == "cycle") + return self._launchers[launch_args.dtype](launch_args) + + def cache_tag(self) -> str: + backend_py = _VMI_ROOT / "backend.py" + return f"vmi:{backend_py}:{os.path.getmtime(backend_py):.0f}" + + def build_artifact_plan(self, case_id: str, case: object) -> ArtifactPlan: + del case_id + return _build_artifact_plan(case) diff --git a/test/kernel-test/kernels/rope/vmi/rope_bf16.vmi.pto b/test/kernel-test/kernels/rope/vmi/rope_bf16.vmi.pto new file mode 100644 index 0000000000..2a6b23700a --- /dev/null +++ b/test/kernel-test/kernels/rope/vmi/rope_bf16.vmi.pto @@ -0,0 +1,364 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @rope_vmi_bf16( + %x_gm: !pto.ptr, + %cos_gm: !pto.ptr, + %sin_gm: !pto.ptr, + %y_gm: !pto.ptr, + %sCount: i32, + %nCount: i32, + %mode: i32) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c0_i16 = arith.constant 0 : i16 + %c1 = arith.constant 1 : index + %c1_i16 = arith.constant 1 : i16 + %c2_i32 = arith.constant 2 : i32 + %c15 = arith.constant 15 : index + %c16 = arith.constant 16 : index + %c63 = arith.constant 63 : index + %c1_i64 = arith.constant 1 : i64 + %c2_i64 = arith.constant 2 : i64 + %c31_i64 = arith.constant 31 : i64 + %c32_i64 = arith.constant 32 : i64 + %c32 = arith.constant 32 : index + %c64 = arith.constant 64 : index + %c64_i32 = arith.constant 64 : i32 + %c128_i64 = arith.constant 128 : i64 + %c0_i32 = arith.constant 0 : i32 + %c0_i64 = arith.constant 0 : i64 + + %s_count = arith.index_cast %sCount : i32 to index + %s_count_i16 = arith.index_cast %s_count : index to i16 + %n_count = arith.index_cast %nCount : i32 to index + %n_count_i16 = arith.index_cast %n_count : index to i16 + + %cos_elems = arith.muli %s_count, %c64 : index + %xy_rows = arith.muli %s_count, %n_count : index + %xy_elems = arith.muli %xy_rows, %c64 : index + + %cos_elems_i64 = arith.index_cast %cos_elems : index to i64 + %xy_elems_i64 = arith.index_cast %xy_elems : index to i64 + %cos_bytes = arith.muli %cos_elems_i64, %c2_i64 : i64 + %xy_bytes = arith.muli %xy_elems_i64, %c2_i64 : i64 + + %cos_lines = arith.addi %cos_bytes, %c31_i64 : i64 + %cos_lines_q = arith.divui %cos_lines, %c32_i64 : i64 + %cos_dma_bytes = arith.muli %cos_lines_q, %c32_i64 : i64 + + %xy_lines = arith.addi %xy_bytes, %c31_i64 : i64 + %xy_lines_q = arith.divui %xy_lines, %c32_i64 : i64 + %xy_dma_bytes = arith.muli %xy_lines_q, %c32_i64 : i64 + + %sin_ub_off = arith.addi %cos_dma_bytes, %c0_i64 : i64 + %x_ub_off = arith.addi %sin_ub_off, %cos_dma_bytes : i64 + %y_ub_off = arith.addi %x_ub_off, %xy_dma_bytes : i64 + + %cos_ub = pto.castptr %c0_i64 : i64 -> !pto.ptr + %sin_ub = pto.castptr %sin_ub_off : i64 -> !pto.ptr + %x_ub = pto.castptr %x_ub_off : i64 -> !pto.ptr + %y_ub = pto.castptr %y_ub_off : i64 -> !pto.ptr + + pto.mte_gm_ub %cos_gm, %cos_ub, %c0_i64, %cos_dma_bytes + nburst(%c1_i64, %cos_dma_bytes, %cos_dma_bytes) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + pto.mte_gm_ub %sin_gm, %sin_ub, %c0_i64, %cos_dma_bytes + nburst(%c1_i64, %cos_dma_bytes, %cos_dma_bytes) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + pto.mte_gm_ub %x_gm, %x_ub, %c0_i64, %xy_dma_bytes + nburst(%c1_i64, %xy_dma_bytes, %xy_dma_bytes) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + + pto.set_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + + %half_d_i32 = arith.divui %c64_i32, %c2_i32 : i32 + %half_d = arith.index_cast %half_d_i32 : i32 to index + %half_d_plus = arith.addi %half_d, %c15 : index + %half_d_blocks = arith.divui %half_d_plus, %c16 : index + %half_d_aligned = arith.muli %half_d_blocks, %c16 : index + + %half_repeat_num = arith.addi %half_d, %c63 : index + %half_repeats = arith.divui %half_repeat_num, %c64 : index + %half_repeats_i16 = arith.index_cast %half_repeats : index to i16 + + %x_s_step = arith.muli %n_count, %c64 : index + %x_n_step = arith.addi %c64, %c0 : index + %cs_s_step = arith.addi %c64, %c0 : index + %y_s_step = arith.muli %n_count, %c64 : index + %y_n_step = arith.addi %c64, %c0 : index + + %is_half_mode = arith.cmpi eq, %mode, %c0_i32 : i32 + + pto.vecscope { + // ======================================================================== + // VMI vs MI — RoPE bf16 向量计算路径对比概览 + // ======================================================================== + // 循环结构、数学公式与 MI 版本完全一致，差异在于表达层级： + // + // bf16 相比 f16 的关键差异：cos/sin 为 f16，x/y 为 bf16，需要统一提升到 + // f32 做乘加运算再截断回 bf16。MI 版本中这一过程暴露了大量硬件细节： + // + // ┌──────────────────────────┬──────────────────────────────────────┐ + // │ MI（硬件细节暴露） │ VMI（语义级抽象） │ + // ├──────────────────────────┼──────────────────────────────────────┤ + // │ vlds + UNPK_B16 解包 │ vmi.load（编译器决定物理排布） │ + // │ 有效数据仅在 EVEN 位置 │ continuous 逻辑连续向量 │ + // │ vcvt + PART_EVEN 取数据 │ vmi.extf（语义精度提升） │ + // │ 选错 part = 读到填充 0 │ 无需区分 EVEN/ODD，编译器自动处理 │ + // │ vcvt + PART_EVEN 窄化 │ vmi.truncf（语义精度截断） │ + // │ + rnd="R" sat="SAT" │ 无需手写舍入/饱和参数 │ + // │ vsts + PK_B32 打包写回 │ vmi.masked_store（编译器处理 pack） │ + // │ pset_b16 + pge_b32 多mask│ vmi.create_mask（统一 vmi.mask） │ + // │ vdintlv/vintlv 解交织 │ channel_split / channel_merge │ + // └──────────────────────────┴──────────────────────────────────────┘ + // + // 核心收益：MI 版本中 UNPK_B16 → PART_EVEN → f32 compute → PART_EVEN + // → PK_B32 这条"解包→取偶→算→放偶→打包"的硬件链路完全由编译器处理。 + // VMI 中开发者只需表达：load → extf → mulf/addf/subf → truncf → store。 + // ======================================================================== + // + // ======================================================================== + // Half 模式（contiguous-half layout）— bf16/f16→f32→bf16 混合精度计算 + // 将一个 64-d 行拆成 x1 和 x2，执行 RoPE 旋转： + // 所有计算在 f32 精度下进行，结果截断回 bf16。 + // ======================================================================== + scf.if %is_half_mode { + scf.for %s = %c0_i16 to %s_count_i16 step %c1_i16 : i16 { + %s_idx = arith.index_cast %s : i16 to index + %x_s_off = arith.muli %s_idx, %x_s_step : index + %cs_s_off = arith.muli %s_idx, %cs_s_step : index + %y_s_off = arith.muli %s_idx, %y_s_step : index + + %cs_off = arith.addi %cs_s_off, %c0 : index + scf.for %rep = %c0_i16 to %half_repeats_i16 step %c1_i16 : i16 { + %rep_idx = arith.index_cast %rep : i16 to index + %elem_off = arith.muli %rep_idx, %c64 : index + %rem = arith.subi %half_d, %elem_off : index + %lt_64 = arith.cmpi ult, %rem, %c64 : index + %active = arith.select %lt_64, %rem, %c64 : index + // VMI: create_mask 逻辑长度创建，统一 vmi.mask。 + // MI 对应: pset_b16 "PAT_ALL" → !pto.mask (load 用) + // pge_b32 "PAT_VL32" → !pto.mask (compute 用) + // 不同数据类型需要不同粒度的 mask——VMI 统一为一种。 + %mask = pto.vmi.create_mask %active : index -> !pto.vmi.mask<64xpred> + + %cos1_off = arith.addi %cs_off, %elem_off : index + %cos2_base = arith.addi %cs_off, %half_d_aligned : index + %cos2_off = arith.addi %cos2_base, %elem_off : index + %sin1_off = arith.addi %cs_off, %elem_off : index + %sin2_base = arith.addi %cs_off, %half_d_aligned : index + %sin2_off = arith.addi %sin2_base, %elem_off : index + + // VMI: vmi.load 语义——加载逻辑连续向量，cos/sin 为 f16。 + // MI 对应: pto.vlds %cos[%off] {dist = "UNPK_B16"} + // → !pto.vreg<128xf16>， + // UNPK_B16 把 64 个 f16 解包到 128xf16 的 EVEN 位置（ODD 填 0）， + // 后续 vcvt {part = "EVEN"} 才能取到有效数据。 + // VMI 屏蔽了这个隐含协议：load → extf 直接表达"加载并提升精度"。 + %cos1_16 = pto.vmi.load %cos_ub[%cos1_off] + : !pto.ptr -> !pto.vmi.vreg<64xf16> + %cos2_16 = pto.vmi.load %cos_ub[%cos2_off] + : !pto.ptr -> !pto.vmi.vreg<64xf16> + %sin1_16 = pto.vmi.load %sin_ub[%sin1_off] + : !pto.ptr -> !pto.vmi.vreg<64xf16> + %sin2_16 = pto.vmi.load %sin_ub[%sin2_off] + : !pto.ptr -> !pto.vmi.vreg<64xf16> + + // VMI: extf 将 f16/bf16 语义提升为 f32——无需 part 选择。 + // MI 对应: pto.vcvt %cos_lo_16, %mask16_all {part = "EVEN"} + // → 开发者必须知道 UNPK_B16 后数据在 EVEN，选 ODD 读到填充 0。 + // VMI 编译器追踪数据来源，自动选择正确的转换 part。 + %cos1 = pto.vmi.extf %cos1_16 + : !pto.vmi.vreg<64xf16> -> !pto.vmi.vreg<64xf32> + %cos2 = pto.vmi.extf %cos2_16 + : !pto.vmi.vreg<64xf16> -> !pto.vmi.vreg<64xf32> + %sin1 = pto.vmi.extf %sin1_16 + : !pto.vmi.vreg<64xf16> -> !pto.vmi.vreg<64xf32> + %sin2 = pto.vmi.extf %sin2_16 + : !pto.vmi.vreg<64xf16> -> !pto.vmi.vreg<64xf32> + + scf.for %n = %c0_i16 to %n_count_i16 step %c1_i16 : i16 { + %n_idx = arith.index_cast %n : i16 to index + %x_n_off = arith.muli %n_idx, %x_n_step : index + %y_n_off = arith.muli %n_idx, %y_n_step : index + + %x_off = arith.addi %x_s_off, %x_n_off : index + %y_off = arith.addi %y_s_off, %y_n_off : index + %x1_off = arith.addi %x_off, %elem_off : index + %x2_base = arith.addi %x_off, %half_d_aligned : index + %x2_off = arith.addi %x2_base, %elem_off : index + %y1_off = arith.addi %y_off, %elem_off : index + %y2_base = arith.addi %y_off, %half_d_aligned : index + %y2_off = arith.addi %y2_base, %elem_off : index + + // VMI: vmi.load 语义——加载 bf16 逻辑向量，自动处理物理排布。 + // MI 对应: pto.vlds %x[%off] {dist = "UNPK_B16"} + // → !pto.vreg<128xbf16>，64 个 bf16 解包到 128 宽的 EVEN 位置。 + %x1_16 = pto.vmi.load %x_ub[%x1_off] + : !pto.ptr -> !pto.vmi.vreg<64xbf16> + %x2_16 = pto.vmi.load %x_ub[%x2_off] + : !pto.ptr -> !pto.vmi.vreg<64xbf16> + + // VMI: extf 语义提升 bf16 → f32，无需 part。 + // MI 对应: pto.vcvt %x_lo_16, %mask16_all {part = "EVEN"} + // → !pto.vreg<64xf32>，必须指定 EVEN part。 + %x1 = pto.vmi.extf %x1_16 + : !pto.vmi.vreg<64xbf16> -> !pto.vmi.vreg<64xf32> + %x2 = pto.vmi.extf %x2_16 + : !pto.vmi.vreg<64xbf16> -> !pto.vmi.vreg<64xf32> + + // ======================================================================== + // Half 模式 RoPE 核心计算 (f32, 截断回 bf16) + // 公式: out1 = x1*cos - x2*sin, out2 = x2*cos + x1*sin + // + // VMI: mulf / subf / addf 无 mask 参数——编译器管理。 + // MI 对应: + // %t0 = pto.vmul %cos_lo, %x_lo, %mask32_half ← mask + // %t1 = pto.vmul %sin_lo, %x_hi, %mask32_half + // %y_lo_f32 = pto.vsub %t0, %t1, %mask32_half + // 每条指令都需要显式 mask，VMI 消除这一负担。 + // ======================================================================== + // y1 = x1 * cos - x2 * sin + %x1_cos = pto.vmi.mulf %x1, %cos1 + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xf32> + %x2_sin = pto.vmi.mulf %x2, %sin1 + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xf32> + %out1_f32 = pto.vmi.subf %x1_cos, %x2_sin + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xf32> + + // y2 = x2 * cos + x1 * sin + %x2_cos = pto.vmi.mulf %x2, %cos2 + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xf32> + %x1_sin = pto.vmi.mulf %x1, %sin2 + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xf32> + %out2_f32 = pto.vmi.addf %x2_cos, %x1_sin + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xf32> + + // VMI: truncf 语义截断 f32 → bf16，无需 rnd/sat 参数。 + // MI 对应: pto.vcvt %y_lo_f32, %mask32_half + // {part = "EVEN", rnd = "R", sat = "SAT"} + // → !pto.vreg<128xbf16>，然后还需 PK_B32 store 打包写回。 + // VMI 中 truncf + masked_store 直接表达"截断并写回"， + // 编译器自动填入 rnd/sat 并处理 pack。 + %out1 = pto.vmi.truncf %out1_f32 + : !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xbf16> + %out2 = pto.vmi.truncf %out2_f32 + : !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xbf16> + + // VMI: masked_store 统一掩码写回。 + // MI 对应: pto.vsts %y_lo_16, %y_ub[%y_off], %mask32_half + // {dist = "PK_B32"} + // PK_B32 告诉硬件把散在 EVEN 位置的有效 bf16 打包成内存中 + // 连续的 32 个元素——开发者需要理解何时用 PK、何时不用。 + // VMI 编译器根据数据布局自动选择打包策略。 + pto.vmi.masked_store %out1, %y_ub[%y1_off], %mask + : !pto.vmi.vreg<64xbf16>, !pto.ptr, !pto.vmi.mask<64xpred> + pto.vmi.masked_store %out2, %y_ub[%y2_off], %mask + : !pto.vmi.vreg<64xbf16>, !pto.ptr, !pto.vmi.mask<64xpred> + } + } + } + } else { + scf.for %s = %c0_i16 to %s_count_i16 step %c1_i16 : i16 { + %s_idx = arith.index_cast %s : i16 to index + %x_s_off = arith.muli %s_idx, %x_s_step : index + %cs_s_off = arith.muli %s_idx, %cs_s_step : index + %y_s_off = arith.muli %s_idx, %y_s_step : index + %cs_off = arith.addi %cs_s_off, %c0 : index + // ======================================================================== + // Interleave 模式（奇偶交织布局）— bf16/f16→f32→bf16 混合精度 + // 公式: y = x*cos + rot(x)*sin, rot(x) = [-x_odd, x_even] + // + // VMI: channel_split/merge 将奇偶分解表达为语义通道算子。 + // MI 对应: vdintlv(x,x) → even/odd → negate odd → vintlv(-odd, even) + // 这条"硬件解交织→手工negate→硬件交织"的链路被 VMI 的 + // channel_split → negf → channel_merge 取代为更直观的语义表达。 + // ======================================================================== + %mask = pto.vmi.create_mask %c64 : index -> !pto.vmi.mask<128xpred> + // VMI: 加载 f16 cos/sin，extf 提升到 f32。 + // MI 对应: vlds + UNPK_B16 → vcvt + PART_EVEN（隐式协议）， + // VMI 中 load + extf 直接表达"加载并提升"。 + %cos16 = pto.vmi.load %cos_ub[%cs_off] + : !pto.ptr -> !pto.vmi.vreg<128xf16> + %sin16 = pto.vmi.load %sin_ub[%cs_off] + : !pto.ptr -> !pto.vmi.vreg<128xf16> + + %cos = pto.vmi.extf %cos16 + : !pto.vmi.vreg<128xf16> -> !pto.vmi.vreg<128xf32> + %sin = pto.vmi.extf %sin16 + : !pto.vmi.vreg<128xf16> -> !pto.vmi.vreg<128xf32> + + scf.for %n = %c0_i16 to %n_count_i16 step %c1_i16 : i16 { + %n_idx = arith.index_cast %n : i16 to index + %x_n_off = arith.muli %n_idx, %x_n_step : index + %y_n_off = arith.muli %n_idx, %y_n_step : index + + %x_off = arith.addi %x_s_off, %x_n_off : index + %y_off = arith.addi %y_s_off, %y_n_off : index + + %x16 = pto.vmi.load %x_ub[%x_off] + : !pto.ptr -> !pto.vmi.vreg<128xbf16> + %x = pto.vmi.extf %x16 + : !pto.vmi.vreg<128xbf16> -> !pto.vmi.vreg<128xf32> + + // ======================================================================== + // VMI: channel_split 拆出 even/odd → negf odd → channel_merge + // 构造旋转伴侣向量 rot(x) = [-x_odd, x_even]。 + // + // MI 对应 (bf16 interleave): + // %x16 = vlds + UNPK_B16 → 128xbf16 (有效数据在 EVEN) + // %x = vcvt + PART_EVEN → 64xf32 (取 EVEN 位置的有效数据) + // %x_even, %x_odd = vdintlv(x, x) → 硬件解交织 + // ... (8条 vmul/vadd/vsub + mask) + // %y_pack, _ = vintlv(y_even, y_odd) → 硬件交织 + // %y_pack_16 = vcvt + PART_EVEN + rnd/sat → 128xbf16 + // vsts + PK_B32 → store + // + // VMI 简化: + // 1. load(128xbf16) 直接得到满宽向量（无需 UNPK_B16 半宽） + // 2. extf → 128xf32 全宽精度提升（无需 part 选择） + // 3. channel_split/merge 语义化奇偶分解（无需 vdintlv/vintlv） + // 4. mulf/addf 无 mask 参数 + // 5. truncf 自动处理 rnd/sat + // 6. masked_store 自动处理 PK 策略 + // ======================================================================== + %x_even, %x_odd = "pto.vmi.channel_split"(%x) + : (!pto.vmi.vreg<128xf32>) -> (!pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32>) + %neg_x_odd = pto.vmi.negf %x_odd + : !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xf32> + %rot = "pto.vmi.channel_merge"(%neg_x_odd, %x_even) + : (!pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32>) -> !pto.vmi.vreg<128xf32> + + %x_cos = pto.vmi.mulf %x, %cos + : !pto.vmi.vreg<128xf32>, !pto.vmi.vreg<128xf32> -> !pto.vmi.vreg<128xf32> + %rot_sin = pto.vmi.mulf %rot, %sin + : !pto.vmi.vreg<128xf32>, !pto.vmi.vreg<128xf32> -> !pto.vmi.vreg<128xf32> + %y_f32 = pto.vmi.addf %x_cos, %rot_sin + : !pto.vmi.vreg<128xf32>, !pto.vmi.vreg<128xf32> -> !pto.vmi.vreg<128xf32> + %y16 = pto.vmi.truncf %y_f32 + : !pto.vmi.vreg<128xf32> -> !pto.vmi.vreg<128xbf16> + + pto.vmi.masked_store %y16, %y_ub[%y_off], %mask + : !pto.vmi.vreg<128xbf16>, !pto.ptr, !pto.vmi.mask<128xpred> + } + } + } + } + + pto.set_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + pto.wait_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + + pto.mte_ub_gm %y_ub, %y_gm, %xy_dma_bytes + nburst(%c1_i64, %xy_dma_bytes, %xy_dma_bytes) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + pto.barrier #pto.pipe + return + } +} diff --git a/test/kernel-test/kernels/rope/vmi/rope_f16.vmi.pto b/test/kernel-test/kernels/rope/vmi/rope_f16.vmi.pto new file mode 100644 index 0000000000..8f5858e079 --- /dev/null +++ b/test/kernel-test/kernels/rope/vmi/rope_f16.vmi.pto @@ -0,0 +1,338 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @rope_vmi_f16( + %x_gm: !pto.ptr, + %cos_gm: !pto.ptr, + %sin_gm: !pto.ptr, + %y_gm: !pto.ptr, + %sCount: i32, + %nCount: i32, + %mode: i32) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c0_i16 = arith.constant 0 : i16 + %c1 = arith.constant 1 : index + %c1_i16 = arith.constant 1 : i16 + %c2_i32 = arith.constant 2 : i32 + %c127 = arith.constant 127 : index + %c15 = arith.constant 15 : index + %c16 = arith.constant 16 : index + %c1_i64 = arith.constant 1 : i64 + %c2_i64 = arith.constant 2 : i64 + %c31_i64 = arith.constant 31 : i64 + %c32_i64 = arith.constant 32 : i64 + %c64 = arith.constant 64 : index + %c64_i32 = arith.constant 64 : i32 + %c128 = arith.constant 128 : index + %c128_i64 = arith.constant 128 : i64 + %c0_i32 = arith.constant 0 : i32 + %c0_i64 = arith.constant 0 : i64 + %s_count = arith.index_cast %sCount : i32 to index + %s_count_i16 = arith.index_cast %s_count : index to i16 + %n_count = arith.index_cast %nCount : i32 to index + %n_count_i16 = arith.index_cast %n_count : index to i16 + + %cos_elems = arith.muli %s_count, %c64 : index + %xy_rows = arith.muli %s_count, %n_count : index + %xy_elems = arith.muli %xy_rows, %c64 : index + + %cos_elems_i64 = arith.index_cast %cos_elems : index to i64 + %xy_elems_i64 = arith.index_cast %xy_elems : index to i64 + %cos_bytes = arith.muli %cos_elems_i64, %c2_i64 : i64 + %xy_bytes = arith.muli %xy_elems_i64, %c2_i64 : i64 + + %cos_lines = arith.addi %cos_bytes, %c31_i64 : i64 + %cos_lines_q = arith.divui %cos_lines, %c32_i64 : i64 + %cos_dma_bytes = arith.muli %cos_lines_q, %c32_i64 : i64 + + %xy_lines = arith.addi %xy_bytes, %c31_i64 : i64 + %xy_lines_q = arith.divui %xy_lines, %c32_i64 : i64 + %xy_dma_bytes = arith.muli %xy_lines_q, %c32_i64 : i64 + + %sin_ub_off = arith.addi %cos_dma_bytes, %c0_i64 : i64 + %x_ub_off = arith.addi %sin_ub_off, %cos_dma_bytes : i64 + %y_ub_off = arith.addi %x_ub_off, %xy_dma_bytes : i64 + + %cos_ub = pto.castptr %c0_i64 : i64 -> !pto.ptr + %sin_ub = pto.castptr %sin_ub_off : i64 -> !pto.ptr + %x_ub = pto.castptr %x_ub_off : i64 -> !pto.ptr + %y_ub = pto.castptr %y_ub_off : i64 -> !pto.ptr + + pto.mte_gm_ub %cos_gm, %cos_ub, %c0_i64, %cos_dma_bytes + nburst(%c1_i64, %cos_dma_bytes, %cos_dma_bytes) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + pto.mte_gm_ub %sin_gm, %sin_ub, %c0_i64, %cos_dma_bytes + nburst(%c1_i64, %cos_dma_bytes, %cos_dma_bytes) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + pto.mte_gm_ub %x_gm, %x_ub, %c0_i64, %xy_dma_bytes + nburst(%c1_i64, %xy_dma_bytes, %xy_dma_bytes) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + + pto.set_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + + %half_d_i32 = arith.divui %c64_i32, %c2_i32 : i32 + %half_d = arith.index_cast %half_d_i32 : i32 to index + %half_d_plus = arith.addi %half_d, %c15 : index + %half_d_blocks = arith.divui %half_d_plus, %c16 : index + %half_d_aligned = arith.muli %half_d_blocks, %c16 : index + + %half_repeat_num = arith.addi %half_d, %c127 : index + %half_repeats = arith.divui %half_repeat_num, %c128 : index + %full_repeat_num = arith.addi %c64, %c127 : index + %full_repeats = arith.divui %full_repeat_num, %c128 : index + %half_repeats_i16 = arith.index_cast %half_repeats : index to i16 + %full_repeats_i16 = arith.index_cast %full_repeats : index to i16 + + %x_s_step = arith.muli %n_count, %c64 : index + %x_n_step = arith.addi %c64, %c0 : index + %cs_s_step = arith.addi %c64, %c0 : index + %y_s_step = arith.muli %n_count, %c64 : index + %y_n_step = arith.addi %c64, %c0 : index + + %is_half_mode = arith.cmpi eq, %mode, %c0_i32 : i32 + + pto.vecscope { + // ======================================================================== + // VMI vs MI — RoPE f16 向量计算路径对比概览 + // ======================================================================== + // 循环结构、数学公式与 MI 版本完全一致，差异在于表达层级： + // + // ┌──────────────────────┬────────────────────────────────────┐ + // │ MI（硬件细节暴露） │ VMI（语义级抽象） │ + // ├──────────────────────┼────────────────────────────────────┤ + // │ vlds + dist 模式 │ vmi.load（编译器决定物理排布） │ + // │ UNPK_B16 解包 │ continuous 逻辑连续向量 │ + // │ vcvt + part EVEN/ODD │ vmi.extf / vmi.truncf（纯语义转换） │ + // │ 显式 rnd="R" sat= │ 无需指定硬件舍入/饱和参数 │ + // │ vmul/vsub/vadd │ vmi.mulf / subf / addf │ + // │ + mask │ 无 mask 参数（编译器管理） │ + // │ vsts + PK_B32 打包 │ vmi.masked_store（编译器处理 pack） │ + // │ pge_b16/pset_b16 │ vmi.create_mask（逻辑长度创建） │ + // │ vdintlv/vintlv 解交织 │ channel_split / channel_merge │ + // │ 手工 negate 穿插 │ 语义化的旋转伴侣向量构造 │ + // └──────────────────────┴────────────────────────────────────┘ + // + // 核心收益：loop body 从 10+ 条物理指令压缩到 ~6 条语义指令， + // 代码直接反映算法流程（load → mul → sub/add → store）， + // 不再夹杂寄存器拆分、part 选择、pack/merge 等硬件噪音。 + // ======================================================================== + // + // ======================================================================== + // Half 模式（contiguous-half layout）— 纯 f16 向量计算路径 + // 将一个 64-d 行拆成 x1（前半）和 x2（后半），执行 RoPE 旋转： + // out1 = x1 * cos - x2 * sin + // out2 = x2 * cos + x1 * sin + // MI 版本中此模式需 f16 直接乘加（无需 f32 中间表示），但每次 vmul/vsub/vadd + // 都需携带 mask，且 load 需 UNPK_B16 解包、store 需 PK_B32 打包。 + // ======================================================================== + scf.if %is_half_mode { + scf.for %s = %c0_i16 to %s_count_i16 step %c1_i16 : i16 { + %s_idx = arith.index_cast %s : i16 to index + %x_s_off = arith.muli %s_idx, %x_s_step : index + %cs_s_off = arith.muli %s_idx, %cs_s_step : index + %y_s_off = arith.muli %s_idx, %y_s_step : index + + %cs_off = arith.addi %cs_s_off, %c0 : index + scf.for %rep = %c0_i16 to %half_repeats_i16 step %c1_i16 : i16 { + %rep_idx = arith.index_cast %rep : i16 to index + %elem_off = arith.muli %rep_idx, %c128 : index + %rem = arith.subi %half_d, %elem_off : index + %lt_128 = arith.cmpi ult, %rem, %c128 : index + %active = arith.select %lt_128, %rem, %c128 : index + // VMI: create_mask 用逻辑长度直接创建，统一 vmi.mask。 + // MI 对应: pge_b16 "PAT_VL32" → !pto.mask， + // 不同数据类型需要不同粒度的 mask（b16/b32/b8），VMI 统一为一种。 + %mask = pto.vmi.create_mask %active : index -> !pto.vmi.mask<128xpred> + + %cos1_off = arith.addi %cs_off, %elem_off : index + %cos2_base = arith.addi %cs_off, %half_d_aligned : index + %cos2_off = arith.addi %cos2_base, %elem_off : index + %sin1_off = arith.addi %cs_off, %elem_off : index + %sin2_base = arith.addi %cs_off, %half_d_aligned : index + %sin2_off = arith.addi %sin2_base, %elem_off : index + + // VMI: vmi.load 语义——加载逻辑连续向量 vreg<128xf16>。 + // MI 对应: pto.vlds %cos_ub[%cs_off] {dist = "UNPK_B16"} + // → !pto.vreg<128xf16>，需要开发者指定 UNPK_B16 解包模式， + // 且明白有效数据只占 EVEN 位置，ODD 是填充 0。 + // VMI 消除了 dist 模式选择——编译器根据上下文自动决定物理排布。 + %cos1_16 = pto.vmi.load %cos_ub[%cos1_off] + : !pto.ptr -> !pto.vmi.vreg<128xf16> + %cos2_16 = pto.vmi.load %cos_ub[%cos2_off] + : !pto.ptr -> !pto.vmi.vreg<128xf16> + %sin1_16 = pto.vmi.load %sin_ub[%sin1_off] + : !pto.ptr -> !pto.vmi.vreg<128xf16> + %sin2_16 = pto.vmi.load %sin_ub[%sin2_off] + : !pto.ptr -> !pto.vmi.vreg<128xf16> + + scf.for %n = %c0_i16 to %n_count_i16 step %c1_i16 : i16 { + %n_idx = arith.index_cast %n : i16 to index + %x_n_off = arith.muli %n_idx, %x_n_step : index + %y_n_off = arith.muli %n_idx, %y_n_step : index + + %x_off = arith.addi %x_s_off, %x_n_off : index + %y_off = arith.addi %y_s_off, %y_n_off : index + %x1_off = arith.addi %x_off, %elem_off : index + %x2_base = arith.addi %x_off, %half_d_aligned : index + %x2_off = arith.addi %x2_base, %elem_off : index + %y1_off = arith.addi %y_off, %elem_off : index + %y2_base = arith.addi %y_off, %half_d_aligned : index + %y2_off = arith.addi %y2_base, %elem_off : index + + %x1_16 = pto.vmi.load %x_ub[%x1_off] + : !pto.ptr -> !pto.vmi.vreg<128xf16> + %x2_16 = pto.vmi.load %x_ub[%x2_off] + : !pto.ptr -> !pto.vmi.vreg<128xf16> + + // ======================================================================== + // Half 模式 RoPE 核心计算 (f16, 无 f32 中间转换) + // 公式: out1 = x1*cos - x2*sin + // out2 = x2*cos + x1*sin + // + // VMI: vmi.mulf / subf / addf 直接表达向量乘减加，无需 mask 参数。 + // MI 对应 (f16 half 模式): + // %t0 = pto.vmul %cos_lo, %x_lo, %mask16_half ← 需 mask + // %t1 = pto.vmul %sin_lo, %x_hi, %mask16_half + // %y_lo = pto.vsub %t0, %t1, %mask16_half + // ... + // 每条算术指令都需要显式携带 mask（b16 for f16）， + // VMI 消除了这个负担——编译器自动管理 mask。 + // ======================================================================== + %x1_cos = pto.vmi.mulf %x1_16, %cos1_16 + : !pto.vmi.vreg<128xf16>, !pto.vmi.vreg<128xf16> -> !pto.vmi.vreg<128xf16> + %x2_sin = pto.vmi.mulf %x2_16, %sin1_16 + : !pto.vmi.vreg<128xf16>, !pto.vmi.vreg<128xf16> -> !pto.vmi.vreg<128xf16> + %out1 = pto.vmi.subf %x1_cos, %x2_sin + : !pto.vmi.vreg<128xf16>, !pto.vmi.vreg<128xf16> -> !pto.vmi.vreg<128xf16> + + %x2_cos = pto.vmi.mulf %x2_16, %cos2_16 + : !pto.vmi.vreg<128xf16>, !pto.vmi.vreg<128xf16> -> !pto.vmi.vreg<128xf16> + %x1_sin = pto.vmi.mulf %x1_16, %sin2_16 + : !pto.vmi.vreg<128xf16>, !pto.vmi.vreg<128xf16> -> !pto.vmi.vreg<128xf16> + %out2 = pto.vmi.addf %x2_cos, %x1_sin + : !pto.vmi.vreg<128xf16>, !pto.vmi.vreg<128xf16> -> !pto.vmi.vreg<128xf16> + + // VMI: masked_store 统一处理部分写入，一个 mask 管理活跃 lane。 + // MI 对应: pto.vsts %y_lo, %y_ub[%y_off], %mask16_half + // → !pto.vreg<128xf16> + !pto.mask， + // 且 f16 store 在 MI 中可能需要 {pk = "PK_B32"} 打包模式 + // （当数据从 f32 narrowing 后分散在 EVEN 位置时）， + // VMI 无需关心 pack 策略。 + pto.vmi.masked_store %out1, %y_ub[%y1_off], %mask + : !pto.vmi.vreg<128xf16>, !pto.ptr, !pto.vmi.mask<128xpred> + pto.vmi.masked_store %out2, %y_ub[%y2_off], %mask + : !pto.vmi.vreg<128xf16>, !pto.ptr, !pto.vmi.mask<128xpred> + } + } + } + } else { + scf.for %s = %c0_i16 to %s_count_i16 step %c1_i16 : i16 { + %s_idx = arith.index_cast %s : i16 to index + %x_s_off = arith.muli %s_idx, %x_s_step : index + %cs_s_off = arith.muli %s_idx, %cs_s_step : index + %y_s_off = arith.muli %s_idx, %y_s_step : index + + %cs_off = arith.addi %cs_s_off, %c0 : index + scf.for %rep = %c0_i16 to %full_repeats_i16 step %c1_i16 : i16 { + %rep_idx = arith.index_cast %rep : i16 to index + %elem_off = arith.muli %rep_idx, %c128 : index + %rem = arith.subi %c64, %elem_off : index + %lt_128 = arith.cmpi ult, %rem, %c128 : index + %active = arith.select %lt_128, %rem, %c128 : index + // ======================================================================== + // Interleave 模式（奇偶交织布局）— 纯 f16 计算路径 + // 公式: y = x * cos + rot(x) * sin + // 其中 rot(x) = [-x_odd, x_even] 是旋转伴侣向量 + // + // 与 MI/CCE 参考实现保持一致：奇偶拆分、乘加、合并都在 f16 上直接完成， + // 不再经过 extf → f32 compute → truncf 这条额外链路。 + // ======================================================================== + %mask = pto.vmi.create_mask %active : index -> !pto.vmi.mask<128xpred> + + %cs_elem_off = arith.addi %cs_off, %elem_off : index + %cos16 = pto.vmi.load %cos_ub[%cs_elem_off] + : !pto.ptr -> !pto.vmi.vreg<128xf16> + %sin16 = pto.vmi.load %sin_ub[%cs_elem_off] + : !pto.ptr -> !pto.vmi.vreg<128xf16> + %cos_even, %cos_odd = "pto.vmi.channel_split"(%cos16) + : (!pto.vmi.vreg<128xf16>) -> (!pto.vmi.vreg<64xf16>, !pto.vmi.vreg<64xf16>) + %sin_even, %sin_odd = "pto.vmi.channel_split"(%sin16) + : (!pto.vmi.vreg<128xf16>) -> (!pto.vmi.vreg<64xf16>, !pto.vmi.vreg<64xf16>) + + scf.for %n = %c0_i16 to %n_count_i16 step %c1_i16 : i16 { + %n_idx = arith.index_cast %n : i16 to index + %x_n_off = arith.muli %n_idx, %x_n_step : index + %y_n_off = arith.muli %n_idx, %y_n_step : index + + %x_off = arith.addi %x_s_off, %x_n_off : index + %y_off = arith.addi %y_s_off, %y_n_off : index + %x_elem_off = arith.addi %x_off, %elem_off : index + %y_elem_off = arith.addi %y_off, %elem_off : index + + %x16 = pto.vmi.load %x_ub[%x_elem_off] + : !pto.ptr -> !pto.vmi.vreg<128xf16> + + // ======================================================================== + // VMI: channel_split / channel_merge 是第一公民的通道算子， + // 语义化地将交织的 128xf16 拆成 [even: 64xf16, odd: 64xf16]， + // 先完成与 MI 一致的成对 f16 乘加，再 merge 回交织结果。 + // + // MI 对应（f16 interleave 模式，纯 f16 无 f32 中间）: + // %x_even, %x_odd = pto.vdintlv %x, %x ← 硬件解交织指令 + // %y_even = x_even*cos_even - x_odd*sin_even + // %y_odd = x_odd*cos_odd + x_even*sin_odd + // %y_pack, %y_pack_hi = pto.vintlv %y_even, %y_odd ← 硬件交织指令 + // + // VMI 的关键简化： + // 1. channel_split/merge 表达"拆奇偶/合奇偶"的语义意图， + // 而非 MI 的 vdintlv/vintlv 硬件指令（需要理解寄存器 lane 映射） + // 2. 直接在 even/odd 通道上表达成对 f16 乘加，避免额外的 f32 中间值 + // 3. 所有 mulf/addf/subf 无需 mask 参数——编译器自动管理 + // ======================================================================== + %x_even, %x_odd = "pto.vmi.channel_split"(%x16) + : (!pto.vmi.vreg<128xf16>) -> (!pto.vmi.vreg<64xf16>, !pto.vmi.vreg<64xf16>) + %x_even_cos = pto.vmi.mulf %x_even, %cos_even + : !pto.vmi.vreg<64xf16>, !pto.vmi.vreg<64xf16> -> !pto.vmi.vreg<64xf16> + %x_odd_sin = pto.vmi.mulf %x_odd, %sin_even + : !pto.vmi.vreg<64xf16>, !pto.vmi.vreg<64xf16> -> !pto.vmi.vreg<64xf16> + %y_even = pto.vmi.subf %x_even_cos, %x_odd_sin + : !pto.vmi.vreg<64xf16>, !pto.vmi.vreg<64xf16> -> !pto.vmi.vreg<64xf16> + + %x_odd_cos = pto.vmi.mulf %x_odd, %cos_odd + : !pto.vmi.vreg<64xf16>, !pto.vmi.vreg<64xf16> -> !pto.vmi.vreg<64xf16> + %x_even_sin = pto.vmi.mulf %x_even, %sin_odd + : !pto.vmi.vreg<64xf16>, !pto.vmi.vreg<64xf16> -> !pto.vmi.vreg<64xf16> + %y_odd = pto.vmi.addf %x_odd_cos, %x_even_sin + : !pto.vmi.vreg<64xf16>, !pto.vmi.vreg<64xf16> -> !pto.vmi.vreg<64xf16> + + %out = "pto.vmi.channel_merge"(%y_even, %y_odd) + : (!pto.vmi.vreg<64xf16>, !pto.vmi.vreg<64xf16>) -> !pto.vmi.vreg<128xf16> + + // VMI: masked_store 用统一的 vmi.mask<128xpred> 掩码写回。 + // MI 对应: pto.vsts %y_pack, %y_ub[%y_off], %mask16_full + // → 需要 !pto.mask + 可能需要 pk 模式。 + pto.vmi.masked_store %out, %y_ub[%y_elem_off], %mask + : !pto.vmi.vreg<128xf16>, !pto.ptr, !pto.vmi.mask<128xpred> + } + } + } + } + } + + pto.set_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + pto.wait_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + + pto.mte_ub_gm %y_ub, %y_gm, %xy_dma_bytes + nburst(%c1_i64, %xy_dma_bytes, %xy_dma_bytes) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + pto.barrier #pto.pipe + return + } +} diff --git a/test/kernel-test/kernels/rope/vmi/rope_f32.vmi.pto b/test/kernel-test/kernels/rope/vmi/rope_f32.vmi.pto new file mode 100644 index 0000000000..1b35bd15b8 --- /dev/null +++ b/test/kernel-test/kernels/rope/vmi/rope_f32.vmi.pto @@ -0,0 +1,331 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @rope_vmi_f32( + %x_gm: !pto.ptr, + %cos_gm: !pto.ptr, + %sin_gm: !pto.ptr, + %y_gm: !pto.ptr, + %sCount: i32, + %nCount: i32, + %mode: i32) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c0_i16 = arith.constant 0 : i16 + %c1_i16 = arith.constant 1 : i16 + %c2_i32 = arith.constant 2 : i32 + %c4_i64 = arith.constant 4 : i64 + %c7 = arith.constant 7 : index + %c8 = arith.constant 8 : index + %c31_i64 = arith.constant 31 : i64 + %c32_i64 = arith.constant 32 : i64 + %c63 = arith.constant 63 : index + %c64 = arith.constant 64 : index + %c64_i32 = arith.constant 64 : i32 + %c0_i32 = arith.constant 0 : i32 + %c1_i64 = arith.constant 1 : i64 + %c0_i64 = arith.constant 0 : i64 + %s_count = arith.index_cast %sCount : i32 to index + %n_count = arith.index_cast %nCount : i32 to index + %s_count_i16 = arith.index_cast %s_count : index to i16 + %n_count_i16 = arith.index_cast %n_count : index to i16 + + %cos_elems = arith.muli %s_count, %c64 : index + %xy_rows = arith.muli %s_count, %n_count : index + %xy_elems = arith.muli %xy_rows, %c64 : index + + %cos_elems_i64 = arith.index_cast %cos_elems : index to i64 + %xy_elems_i64 = arith.index_cast %xy_elems : index to i64 + %cos_bytes = arith.muli %cos_elems_i64, %c4_i64 : i64 + %xy_bytes = arith.muli %xy_elems_i64, %c4_i64 : i64 + + %cos_lines = arith.addi %cos_bytes, %c31_i64 : i64 + %cos_lines_q = arith.divui %cos_lines, %c32_i64 : i64 + %cos_dma_bytes = arith.muli %cos_lines_q, %c32_i64 : i64 + + %xy_lines = arith.addi %xy_bytes, %c31_i64 : i64 + %xy_lines_q = arith.divui %xy_lines, %c32_i64 : i64 + %xy_dma_bytes = arith.muli %xy_lines_q, %c32_i64 : i64 + + %sin_ub_off = arith.addi %cos_dma_bytes, %c0_i64 : i64 + %x_ub_off = arith.addi %sin_ub_off, %cos_dma_bytes : i64 + %y_ub_off = arith.addi %x_ub_off, %xy_dma_bytes : i64 + + %cos_ub = pto.castptr %c0_i64 : i64 -> !pto.ptr + %sin_ub = pto.castptr %sin_ub_off : i64 -> !pto.ptr + %x_ub = pto.castptr %x_ub_off : i64 -> !pto.ptr + %y_ub = pto.castptr %y_ub_off : i64 -> !pto.ptr + + pto.mte_gm_ub %cos_gm, %cos_ub, %c0_i64, %cos_dma_bytes + nburst(%c1_i64, %cos_dma_bytes, %cos_dma_bytes) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + pto.mte_gm_ub %sin_gm, %sin_ub, %c0_i64, %cos_dma_bytes + nburst(%c1_i64, %cos_dma_bytes, %cos_dma_bytes) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + pto.mte_gm_ub %x_gm, %x_ub, %c0_i64, %xy_dma_bytes + nburst(%c1_i64, %xy_dma_bytes, %xy_dma_bytes) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + + pto.set_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + + %half_d_i32 = arith.divui %c64_i32, %c2_i32 : i32 + %half_d = arith.index_cast %half_d_i32 : i32 to index + %half_d_plus = arith.addi %half_d, %c7 : index + %half_d_blocks = arith.divui %half_d_plus, %c8 : index + %half_d_aligned = arith.muli %half_d_blocks, %c8 : index + + %half_repeat_num = arith.addi %half_d, %c63 : index + %half_repeats = arith.divui %half_repeat_num, %c64 : index + %half_repeats_i16 = arith.index_cast %half_repeats : index to i16 + + %x_s_step = arith.muli %n_count, %c64 : index + %x_n_step = arith.addi %c64, %c0 : index + %cs_s_step = arith.addi %c64, %c0 : index + %y_s_step = arith.muli %n_count, %c64 : index + %y_n_step = arith.addi %c64, %c0 : index + + %is_half_mode = arith.cmpi eq, %mode, %c0_i32 : i32 + + pto.vecscope { + // ======================================================================== + // VMI vs MI — RoPE f32 向量计算路径对比概览 + // ======================================================================== + // f32 是全精度路径，无需 extf/truncf 类型转换，差异最纯粹： + // + // ┌──────────────────────────┬──────────────────────────────────────┐ + // │ MI（硬件细节暴露） │ VMI（语义级抽象） │ + // ├──────────────────────────┼──────────────────────────────────────┤ + // │ vlds（无 dist 模式） │ vmi.load（同样简洁，但类型系统统一） │ + // │ vmul/vsub/vadd │ vmi.mulf / subf / addf │ + // │ + mask 每条都带 │ 无 mask 参数（编译器管理） │ + // │ vsts + mask │ vmi.masked_store（统一 vmi.mask） │ + // │ vdintlv/vintlv 硬件交织 │ channel_split / channel_merge │ + // │ 错用会导致错位 │ 语义化奇偶分解，意图明确 │ + // │ pge_b32 "PAT_VL32" │ vmi.create_mask（逻辑长度创建） │ + // │ pset_b32 "PAT_ALL" │ 统一 vmi.mask 类型 │ + // └──────────────────────────┴──────────────────────────────────────┘ + // + // f32 的主要简化集中在两点： + // 1. 算术指令消除 mask 参数（MI 中每条 vmul/vadd/vsub 都需 mask） + // 2. channel_split/merge 替代 vdintlv/vintlv，将"硬件 lane 重排" + // 提升为"奇偶通道分解"的语义表达 + // ======================================================================== + // + // ======================================================================== + // Half 模式 — 纯 f32 向量计算 + // 将 64-d 行拆成两个 32 元素半区，执行 RoPE 旋转： + // ======================================================================== + scf.if %is_half_mode { + scf.for %s = %c0_i16 to %s_count_i16 step %c1_i16 : i16 { + %s_idx = arith.index_cast %s : i16 to index + %x_s_off = arith.muli %s_idx, %x_s_step : index + %cs_s_off = arith.muli %s_idx, %cs_s_step : index + %y_s_off = arith.muli %s_idx, %y_s_step : index + + %cs_off = arith.addi %cs_s_off, %c0 : index + scf.for %rep = %c0_i16 to %half_repeats_i16 step %c1_i16 : i16 { + %rep_idx = arith.index_cast %rep : i16 to index + %elem_off = arith.muli %rep_idx, %c64 : index + %rem = arith.subi %half_d, %elem_off : index + %lt_64 = arith.cmpi ult, %rem, %c64 : index + %active = arith.select %lt_64, %rem, %c64 : index + // VMI: create_mask 逻辑长度 → vmi.mask<64xpred>。 + // MI 对应: pge_b32 "PAT_VL32" → !pto.mask。 + // VMI 统一了 mask 类型，不再区分 b16/b32/b8。 + %mask = pto.vmi.create_mask %active : index -> !pto.vmi.mask<64xpred> + + %cos1_off = arith.addi %cs_off, %elem_off : index + %cos2_base = arith.addi %cs_off, %half_d_aligned : index + %cos2_off = arith.addi %cos2_base, %elem_off : index + %sin1_off = arith.addi %cs_off, %elem_off : index + %sin2_base = arith.addi %cs_off, %half_d_aligned : index + %sin2_off = arith.addi %sin2_base, %elem_off : index + + // VMI: vmi.load 语义——f32 直接加载，无需 dist 模式。 + // MI 对应: pto.vlds %cos_ub[%cs_off] → !pto.vreg<64xf32> + // 因为 f32 不存在 f16/bf16 的 UNPK/PK 问题，MI 的 load 也相对简洁。 + // 但 VMI 用统一的 vmi.load 接口覆盖所有数据类型，接口一致性更好。 + %cos1 = pto.vmi.load %cos_ub[%cos1_off] + : !pto.ptr -> !pto.vmi.vreg<64xf32> + %cos2 = pto.vmi.load %cos_ub[%cos2_off] + : !pto.ptr -> !pto.vmi.vreg<64xf32> + %sin1 = pto.vmi.load %sin_ub[%sin1_off] + : !pto.ptr -> !pto.vmi.vreg<64xf32> + %sin2 = pto.vmi.load %sin_ub[%sin2_off] + : !pto.ptr -> !pto.vmi.vreg<64xf32> + + scf.for %n = %c0_i16 to %n_count_i16 step %c1_i16 : i16 { + %n_idx = arith.index_cast %n : i16 to index + %x_n_off = arith.muli %n_idx, %x_n_step : index + %y_n_off = arith.muli %n_idx, %y_n_step : index + + %x_off = arith.addi %x_s_off, %x_n_off : index + %y_off = arith.addi %y_s_off, %y_n_off : index + %x1_off = arith.addi %x_off, %elem_off : index + %x2_base = arith.addi %x_off, %half_d_aligned : index + %x2_off = arith.addi %x2_base, %elem_off : index + %y1_off = arith.addi %y_off, %elem_off : index + %y2_base = arith.addi %y_off, %half_d_aligned : index + %y2_off = arith.addi %y2_base, %elem_off : index + + %x1 = pto.vmi.load %x_ub[%x1_off] + : !pto.ptr -> !pto.vmi.vreg<64xf32> + %x2 = pto.vmi.load %x_ub[%x2_off] + : !pto.ptr -> !pto.vmi.vreg<64xf32> + + // ======================================================================== + // Half 模式 RoPE 核心计算 (纯 f32) + // 公式: out1 = x1*cos - x2*sin, out2 = x2*cos + x1*sin + // + // VMI: mulf / subf / addf 直接表达向量运算，无 mask 参数。 + // MI 对应: + // %t0 = pto.vmul %cos_lo, %x_lo, %mask32_half ← mask + // %t1 = pto.vmul %sin_lo, %x_hi, %mask32_half + // %y_lo = pto.vsub %t0, %t1, %mask32_half + // 即使 f32 不需要类型转换，MI 的每条算术指令仍然绑定 mask。 + // VMI 编译器根据向量宽度自动生成正确的 mask。 + // ======================================================================== + // y1 = x1 * cos - x2 * sin + %x1_cos = pto.vmi.mulf %x1, %cos1 + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xf32> + %x2_sin = pto.vmi.mulf %x2, %sin1 + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xf32> + %out1 = pto.vmi.subf %x1_cos, %x2_sin + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xf32> + + // y2 = x2 * cos + x1 * sin + %x2_cos = pto.vmi.mulf %x2, %cos2 + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xf32> + %x1_sin = pto.vmi.mulf %x1, %sin2 + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xf32> + %out2 = pto.vmi.addf %x2_cos, %x1_sin + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xf32> + + // VMI: masked_store 统一掩码写回。 + // MI 对应: pto.vsts %y_lo, %y_ub[%y_off], %mask32_half + // → !pto.vreg<64xf32> + !pto.mask。 + pto.vmi.masked_store %out1, %y_ub[%y1_off], %mask + : !pto.vmi.vreg<64xf32>, !pto.ptr, !pto.vmi.mask<64xpred> + pto.vmi.masked_store %out2, %y_ub[%y2_off], %mask + : !pto.vmi.vreg<64xf32>, !pto.ptr, !pto.vmi.mask<64xpred> + } + } + } + } else { + scf.for %s = %c0_i16 to %s_count_i16 step %c1_i16 : i16 { + %s_idx = arith.index_cast %s : i16 to index + %x_s_off = arith.muli %s_idx, %x_s_step : index + %cs_s_off = arith.muli %s_idx, %cs_s_step : index + %y_s_off = arith.muli %s_idx, %y_s_step : index + + %cs_off = arith.addi %cs_s_off, %c0 : index + // ======================================================================== + // Interleave 模式（奇偶交织布局）— 纯 f32 偶奇成对 RoPE + // 公式: y_even = x_even*cos_even - x_odd*sin_even + // y_odd = x_odd *cos_odd + x_even*sin_odd + // + // 这是 f32 VMI interleave 与 f16/bf16 的关键差异： + // f16/bf16 使用 rot(x) = [-x_odd, x_even] 的旋转伴侣向量形式， + // f32 直接对 even/odd 通道分别做 RoPE 乘加，原因在于 f32 的 + // cos/sin 每个元素都需要成对使用（even 配 odd），先 split cos/sin + // 再分别乘加比构造 rot(x) 更自然。 + // + // MI 对应: + // vdintlv(cos, cos) → cos_even, cos_odd (硬件解交织) + // vdintlv(x, x) → x_even, x_odd + // 8条 vmul/vadd/vsub + mask + // vintlv(y_even, y_odd) → y_pack, _ (硬件交织回去) + // + // VMI: channel_split/merge 语义化表达奇偶分解/合并。 + // ======================================================================== + // y_even = x_even * cos_even - x_odd * sin_even + // y_odd = x_odd * cos_odd + x_even * sin_odd + %mask = pto.vmi.create_mask %c64 : index -> !pto.vmi.mask<128xpred> + %cos = pto.vmi.load %cos_ub[%cs_off] + : !pto.ptr -> !pto.vmi.vreg<128xf32> + %sin = pto.vmi.load %sin_ub[%cs_off] + : !pto.ptr -> !pto.vmi.vreg<128xf32> + + // VMI: channel_split 将交织的 [e0,o0, e1,o1, ...] 拆成 + // even=[e0,e1,...] 和 odd=[o0,o1,...] 两个独立通道， + // 使得后续 RoPE 成对计算可以用分离的向量直接表达。 + // MI 对应: pto.vdintlv %cos, %cos → !pto.vreg<64xf32>, !pto.vreg<64xf32> + // 硬件解交织指令，需要理解寄存器 lane 映射关系。 + %cos_even, %cos_odd = "pto.vmi.channel_split"(%cos) + : (!pto.vmi.vreg<128xf32>) -> (!pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32>) + %sin_even, %sin_odd = "pto.vmi.channel_split"(%sin) + : (!pto.vmi.vreg<128xf32>) -> (!pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32>) + + scf.for %n = %c0_i16 to %n_count_i16 step %c1_i16 : i16 { + %n_idx = arith.index_cast %n : i16 to index + %x_n_off = arith.muli %n_idx, %x_n_step : index + %y_n_off = arith.muli %n_idx, %y_n_step : index + + %x_off = arith.addi %x_s_off, %x_n_off : index + %y_off = arith.addi %y_s_off, %y_n_off : index + + %x = pto.vmi.load %x_ub[%x_off] + : !pto.ptr -> !pto.vmi.vreg<128xf32> + + // MI 对应: pto.vdintlv %x, %x → x_even(64xf32), x_odd(64xf32) + // 硬件解交织——VMI 用 channel_split 语义化表达同样的意图。 + %x_even, %x_odd = "pto.vmi.channel_split"(%x) + : (!pto.vmi.vreg<128xf32>) -> (!pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32>) + + // ======================================================================== + // 成对 RoPE 计算 (even/odd 通道独立) + // y_even = x_even*cos_even - x_odd*sin_even + // y_odd = x_odd *cos_odd + x_even*sin_odd + // + // VMI: mulf/subf/addf 无 mask——4条乘+2条加减表达完整的成对 RoPE。 + // MI 对应: 8条 vmul/vadd/vsub，每条都带 mask。 + // ======================================================================== + %x_even_cos = pto.vmi.mulf %x_even, %cos_even + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xf32> + %x_odd_sin = pto.vmi.mulf %x_odd, %sin_even + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xf32> + %y_even = pto.vmi.subf %x_even_cos, %x_odd_sin + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xf32> + + %x_odd_cos = pto.vmi.mulf %x_odd, %cos_odd + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xf32> + %x_even_sin = pto.vmi.mulf %x_even, %sin_odd + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xf32> + %y_odd = pto.vmi.addf %x_odd_cos, %x_even_sin + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xf32> + + // VMI: channel_merge 将两个独立通道合并回交织格式 [e0,o0, e1,o1, ...]。 + // MI 对应: pto.vintlv %y_even, %y_odd → y_pack(64xf32), y_pack_hi(64xf32) + // 硬件交织指令，产生两个输出（低半/高半），而 channel_merge 直接 + // 产出一个 vreg<128xf32>，更符合"逻辑向量"的思维模型。 + %out = "pto.vmi.channel_merge"(%y_even, %y_odd) + : (!pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32>) -> !pto.vmi.vreg<128xf32> + + // VMI: masked_store 统一掩码写回，无需区分 b16/b32 mask 类型。 + // MI 对应: pto.vsts %y_pack, %y_ub[%y_off], %mask32_full + // → 需要 !pto.mask，且 vintlv 产出的 y_pack_hi 在 MI 中 + // 需要额外的 store 指令——VMI 的单个 128xf32 向量更简洁。 + pto.vmi.masked_store %out, %y_ub[%y_off], %mask + : !pto.vmi.vreg<128xf32>, !pto.ptr, !pto.vmi.mask<128xpred> + } + } + } + } + + pto.set_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + pto.wait_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + + pto.mte_ub_gm %y_ub, %y_gm, %xy_dma_bytes + nburst(%c1_i64, %xy_dma_bytes, %xy_dma_bytes) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + pto.barrier #pto.pipe + return + } +} diff --git a/test/kernel-test/run.py b/test/kernel-test/run.py new file mode 100644 index 0000000000..d1eca00676 --- /dev/null +++ b/test/kernel-test/run.py @@ -0,0 +1,18 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""Unified CLI entry point for the kernel-test framework.""" + +from __future__ import annotations + +from kernel_test.cli import main + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/test/kernel-test/scripts/.gitkeep b/test/kernel-test/scripts/.gitkeep new file mode 100644 index 0000000000..8b13789179 --- /dev/null +++ b/test/kernel-test/scripts/.gitkeep @@ -0,0 +1 @@ + diff --git a/test/kernel-test/scripts/README.md b/test/kernel-test/scripts/README.md new file mode 100644 index 0000000000..bb96769ea0 --- /dev/null +++ b/test/kernel-test/scripts/README.md @@ -0,0 +1,35 @@ + + +# kernel-test scripts + +This directory is split into two groups: + +## User-facing entrypoints + +- `run_cycle.sh` + - Recommended command for cycle collection. + - Expands selected cases, invokes `kernel-test/run.py` once per case, and prints a cycle report. + - Accepts `--kernel-dir ` or `KERNEL_TEST_KERNEL_DIR` when kernels live outside the default `kernels/` root. +- `run_sim.sh` + - Generic `cannsim` transport for one Python entrypoint. + - Use this when you want to run a specific script under cannsim directly. +- `run_msprof.sh` + - Generic `msprof` transport for one Python entrypoint. + - Use this when you want direct `msprof` artifacts for a specific script. + +## Internal helpers + +- `helpers/common.sh` +- `helpers/run_sim_entry.sh` +- `helpers/report_cycles.py` + +These helper scripts are framework internals. They are called by the user-facing +entrypoints above and are not intended to be run directly. diff --git a/test/kernel-test/scripts/helpers/common.sh b/test/kernel-test/scripts/helpers/common.sh new file mode 100644 index 0000000000..e32444d384 --- /dev/null +++ b/test/kernel-test/scripts/helpers/common.sh @@ -0,0 +1,190 @@ +#!/usr/bin/env bash +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + + +# Internal helper library for kernel-test shell entrypoints. + +kt_resolve_paths() { + local invoked_path="${1:-${BASH_SOURCE[0]}}" + + KT_SCRIPT_PATH="$(realpath -- "${invoked_path}")" + KT_SCRIPT_DIR="$(cd -- "$(dirname -- "${KT_SCRIPT_PATH}")" && pwd)" + KT_ROOT="$(cd -- "${KT_SCRIPT_DIR}/.." && pwd)" + KT_REPO_ROOT="$(cd -- "${KT_ROOT}/.." && pwd)" +} + +kt_source_ascend_env() { + local ascend_home_path + ascend_home_path="${ASCEND_HOME_PATH:-${ASCEND_TOOLKIT_HOME:-/usr/local/Ascend/ascend-toolkit/latest}}" + # shellcheck disable=SC1090 + source "${ascend_home_path}/bin/setenv.bash" +} + +kt_default_python_cmd() { + printf '%s\n' "${KERNEL_TEST_PYTHON_CMD:-python}" +} + +kt_quote_words() { + local word="" + + for word in "$@"; do + printf '%q ' "${word}" + done +} + +kt_run_python_cmd() { + local python_cmd="$1" + shift + + local quoted_args="" + quoted_args="$(kt_quote_words "$@")" + eval "${python_cmd} ${quoted_args}" +} + +kt_exec_python_cmd() { + local python_cmd="$1" + shift + + local quoted_args="" + quoted_args="$(kt_quote_words "$@")" + eval "exec ${python_cmd} ${quoted_args}" +} + +kt_write_args_file() { + local work_dir="$1" + shift + + local args_file="" + args_file="$(mktemp "${work_dir}/.kernel_test_args.XXXXXX")" + printf '%s\0' "$@" > "${args_file}" + printf '%s\n' "${args_file}" +} + +kt_load_args_file() { + local args_file="${1:-${KERNEL_TEST_ARGS_FILE:-}}" + local arg="" + + KT_LOADED_ARGS=() + if [ -z "${args_file}" ]; then + return 0 + fi + if [ ! -f "${args_file}" ]; then + echo "args file not found: ${args_file}" >&2 + return 1 + fi + + while IFS= read -r -d '' arg; do + KT_LOADED_ARGS+=("${arg}") + done < "${args_file}" +} + +kt_prepare_cannsim_workspace() { + local work_dir="$1" + + mkdir -p "${work_dir}/log_ca" + rm -f "${work_dir}/instr.bin" + find "${work_dir}/log_ca" -mindepth 1 -delete 2>/dev/null || rm -rf "${work_dir}/log_ca"/* 2>/dev/null || true + cd "${work_dir}" +} + +kt_find_latest_log() { + local output_dir="$1" + + find "${output_dir}" -name cannsim.log -type f -printf '%T@ %p\n' 2>/dev/null | sort -nr | head -1 | cut -d' ' -f2- +} + +kt_log_has_cycle_completion() { + local log_path="$1" + + [ -f "${log_path}" ] || return 1 + grep -Eq '^CYCLE_(DONE|SKIP) ' "${log_path}" +} + +kt_descendant_pids() { + local root_pid="$1" + local child_pid="" + + pgrep -P "${root_pid}" 2>/dev/null || true + for child_pid in $(pgrep -P "${root_pid}" 2>/dev/null || true); do + kt_descendant_pids "${child_pid}" + done +} + +kt_terminate_descendants_matching() { + local root_pid="$1" + local pattern="$2" + local descendant="" + local cmdline="" + + while IFS= read -r descendant; do + [ -n "${descendant}" ] || continue + cmdline="$(ps -p "${descendant}" -o cmd= 2>/dev/null || true)" + if [[ "${cmdline}" =~ ${pattern} ]]; then + kill -TERM "${descendant}" 2>/dev/null || true + fi + done < <(kt_descendant_pids "${root_pid}" | sort -nr | uniq) +} + +kt_handle_cannsim_exit() { + local rc="$1" + local output_dir="$2" + local latest_log="" + + if [ "${rc}" -eq 0 ]; then + return 0 + fi + + latest_log="$(kt_find_latest_log "${output_dir}")" + if [ -n "${latest_log}" ]; then + if grep -q "All tests PASSED!" "${latest_log}"; then + echo "==> cannsim exited ${rc} but ${latest_log} reports success; treating as success" + return 0 + fi + if grep -Eq '^CYCLE_(DONE|SKIP) ' "${latest_log}"; then + echo "==> cannsim exited ${rc} after framework cycle completion markers in ${latest_log}; treating as success" + return 0 + fi + fi + + echo "==> cannsim exited ${rc}; inspect ${output_dir}" >&2 + return "${rc}" +} + +kt_collect_cannsim_artifacts() { + local work_dir="$1" + local out_dir="$2" + local cannsim_run="" + local dest="" + local tmp="" + + cannsim_run="$(find "${out_dir}" -maxdepth 1 -type d -name 'cannsim_*' 2>/dev/null | sort | tail -1 || true)" + if [ -z "${cannsim_run}" ]; then + return 0 + fi + + if [ -d "${work_dir}/log_ca" ] && [ ! -d "${cannsim_run}/log_ca" ]; then + cp -a "${work_dir}/log_ca" "${cannsim_run}/log_ca" + fi + + if [ -f "${work_dir}/instr.bin" ]; then + dest="${cannsim_run}/instr.bin" + if [ ! -f "${dest}" ] || [ -L "${dest}" ] || [ "${dest}" -ef "${work_dir}/instr.bin" ]; then + tmp="${cannsim_run}/.instr.bin.tmp" + cp -f "${work_dir}/instr.bin" "${tmp}" + mv -f "${tmp}" "${dest}" + fi + fi + + if [ -f "${cannsim_run}/instr.bin" ]; then + mkdir -p "${cannsim_run}/report" + set +e + cannsim report -e "${cannsim_run}" -o "${cannsim_run}/report" -n 0 + set -e + fi +} diff --git a/test/kernel-test/scripts/helpers/report_cycles.py b/test/kernel-test/scripts/helpers/report_cycles.py new file mode 100644 index 0000000000..e26b30bfb3 --- /dev/null +++ b/test/kernel-test/scripts/helpers/report_cycles.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +"""Internal helper: dispatch cycle reporting to one kernel-local analyzer.""" + +from __future__ import annotations + +import argparse +import os +import sys +from pathlib import Path + +KT_ROOT = Path(__file__).resolve().parents[2] +if str(KT_ROOT) not in sys.path: + sys.path.insert(0, str(KT_ROOT)) + +from kernel_test.registry import import_kernel_module + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description="Report cycle metrics for one kernel-test operator") + parser.add_argument( + "--kernel-dir", + help=( + "Kernel package root to discover. Accepts either a directory that contains " + "kernel subdirectories or one kernel package directory. " + "Defaults to test/kernel-test/kernels or $KERNEL_TEST_KERNEL_DIR." + ), + ) + parser.add_argument("--op", required=True, help="Kernel name") + parser.add_argument("--table", action="store_true", help="Print compact table output") + parser.add_argument("out_dirs", nargs="*", help="Per-case sim output directories") + args = parser.parse_args(argv) + kernel_dir = args.kernel_dir or os.environ.get("KERNEL_TEST_KERNEL_DIR") + + try: + module = import_kernel_module(args.op, kernel_dir=kernel_dir, submodule="cycle_metrics") + except ModuleNotFoundError as exc: + print(f"no cycle metrics analyzer for kernel {args.op}: {exc}", file=sys.stderr) + return 1 + + if hasattr(module, "get_cycle_reporter"): + from kernel_test.cycle_reporting import run_cycle_report + + reporter = module.get_cycle_reporter() + forwarded: list[str] = [] + if args.table: + forwarded.append("--table") + forwarded.extend(args.out_dirs) + return int(run_cycle_report(reporter, forwarded)) + + forwarded = [] + if args.table: + forwarded.append("--table") + forwarded.extend(args.out_dirs) + return int(module.main(forwarded)) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/test/kernel-test/scripts/helpers/run_sim_entry.sh b/test/kernel-test/scripts/helpers/run_sim_entry.sh new file mode 100644 index 0000000000..d3a4ae30fe --- /dev/null +++ b/test/kernel-test/scripts/helpers/run_sim_entry.sh @@ -0,0 +1,27 @@ +#!/usr/bin/env bash +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +set -euo pipefail + +# Internal cannsim entry shim. Not intended for direct user invocation. + +SCRIPT_PATH="${BASH_SOURCE[0]}" +# shellcheck disable=SC1091 +source "$(cd -- "$(dirname -- "$(realpath -- "${SCRIPT_PATH}")")" && pwd)/common.sh" + +kt_resolve_paths "${SCRIPT_PATH}" +kt_load_args_file "${KERNEL_TEST_ARGS_FILE:-}" +kt_prepare_cannsim_workspace "${KERNEL_TEST_WORK_DIR:?KERNEL_TEST_WORK_DIR is required}" + +if [ "$#" -lt 1 ]; then + echo "run_sim_entry.sh requires the python script path from cannsim -u" >&2 + exit 2 +fi + +kt_exec_python_cmd "$(kt_default_python_cmd)" "$1" "${KT_LOADED_ARGS[@]}" diff --git a/test/kernel-test/scripts/run_cycle.sh b/test/kernel-test/scripts/run_cycle.sh new file mode 100755 index 0000000000..c070911164 --- /dev/null +++ b/test/kernel-test/scripts/run_cycle.sh @@ -0,0 +1,285 @@ +#!/usr/bin/env bash +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +set -euo pipefail + +# User-facing entrypoint: expand selected cases and collect cycle metrics. + +SCRIPT_PATH="${BASH_SOURCE[0]}" +# shellcheck disable=SC1091 +SCRIPT_DIR="$(cd -- "$(dirname -- "$(realpath -- "${SCRIPT_PATH}")")" && pwd)" +source "${SCRIPT_DIR}/helpers/common.sh" + +kt_resolve_paths "${SCRIPT_PATH}" + +OP="" +BACKEND="" +KERNEL_DIR="${KERNEL_TEST_KERNEL_DIR:-}" +OUTPUT_ROOT="${KT_ROOT}/sim_outputs" +PYTHON_CMD="$(kt_default_python_cmd)" +CASE_FILTER="" +PARALLEL_SIM=0 +JOBS="" +ENGINE="msprof" +declare -a REQUESTED_CASES=() +declare -a SUCCESS_CASE_DIRS=() + +while [ $# -gt 0 ]; do + case "$1" in + --op) + [ $# -lt 2 ] && { echo "--op requires a value" >&2; exit 1; } + OP="$2" + shift 2 + ;; + --backend) + [ $# -lt 2 ] && { echo "--backend requires a value" >&2; exit 1; } + BACKEND="$2" + shift 2 + ;; + --kernel-dir) + [ $# -lt 2 ] && { echo "--kernel-dir requires a value" >&2; exit 1; } + KERNEL_DIR="$2" + shift 2 + ;; + --case) + [ $# -lt 2 ] && { echo "--case requires a value" >&2; exit 1; } + REQUESTED_CASES+=("$2") + shift 2 + ;; + --case-filter) + [ $# -lt 2 ] && { echo "--case-filter requires a value" >&2; exit 1; } + CASE_FILTER="$2" + shift 2 + ;; + --parallel-sim) + [ $# -lt 2 ] && { echo "--parallel-sim requires 0 or 1" >&2; exit 1; } + PARALLEL_SIM="$2" + shift 2 + ;; + --engine) + [ $# -lt 2 ] && { echo "--engine requires a value" >&2; exit 1; } + ENGINE="$2" + shift 2 + ;; + --jobs) + [ $# -lt 2 ] && { echo "--jobs requires a value" >&2; exit 1; } + JOBS="$2" + shift 2 + ;; + --output-root) + [ $# -lt 2 ] && { echo "--output-root requires a value" >&2; exit 1; } + OUTPUT_ROOT="$2" + shift 2 + ;; + --python-cmd) + [ $# -lt 2 ] && { echo "--python-cmd requires a value" >&2; exit 1; } + PYTHON_CMD="$2" + shift 2 + ;; + --help|-h) + cat < --backend [options] + +Recommended user-facing command for cycle measurement. + +Options: + --kernel-dir Kernel package root or one kernel package directory. + --case Select one case. Repeatable. + --case-filter Filter case ids by substring. + --engine Cycle runner: msprof or cannsim. Default: msprof. + --parallel-sim <0|1> Run cannsim jobs sequentially or in parallel. + --jobs Max parallel jobs when --parallel-sim=1. + --output-root Root directory for sim outputs. + --python-cmd Python launcher for cannsim runs. Default: $(kt_default_python_cmd) +EOF + exit 0 + ;; + *) + echo "unknown argument: $1" >&2 + exit 1 + ;; + esac +done + +[ -n "${OP}" ] || { echo "--op is required" >&2; exit 1; } +[ -n "${BACKEND}" ] || { echo "--backend is required" >&2; exit 1; } +[[ "${PARALLEL_SIM}" =~ ^[01]$ ]] || { echo "--parallel-sim must be 0 or 1" >&2; exit 1; } +[[ "${ENGINE}" =~ ^(msprof|cannsim)$ ]] || { echo "--engine must be msprof or cannsim" >&2; exit 1; } + +OUTPUT_ROOT="$(realpath -m -- "${OUTPUT_ROOT}")" +mkdir -p "${OUTPUT_ROOT}/${OP}/${BACKEND}" + +list_cases() { + local args=("${KT_ROOT}/run.py") + if [ -n "${KERNEL_DIR}" ]; then + args+=(--kernel-dir "${KERNEL_DIR}") + fi + args+=(--op "${OP}" --workflow cycle --list-cases) + kt_run_python_cmd "${PYTHON_CMD}" "${args[@]}" +} + +declare -a CASES=() +if [ "${#REQUESTED_CASES[@]}" -gt 0 ]; then + CASES=("${REQUESTED_CASES[@]}") +else + while IFS= read -r case_id; do + [ -n "${case_id}" ] || continue + CASES+=("${case_id}") + done < <(list_cases) +fi + +if [ -n "${CASE_FILTER}" ]; then + declare -a FILTERED_CASES=() + for case_id in "${CASES[@]}"; do + if [[ "${case_id}" == *"${CASE_FILTER}"* ]]; then + FILTERED_CASES+=("${case_id}") + fi + done + CASES=("${FILTERED_CASES[@]}") +fi + +[ "${#CASES[@]}" -gt 0 ] || { echo "no cases matched the requested selection" >&2; exit 1; } + +if [ "${PARALLEL_SIM}" -eq 1 ] && [ -z "${JOBS}" ]; then + JOBS="$(getconf _NPROCESSORS_ONLN 2>/dev/null || echo 4)" +fi +if [ -n "${JOBS}" ] && ! [[ "${JOBS}" =~ ^[1-9][0-9]*$ ]]; then + echo "--jobs must be a positive integer" >&2 + exit 1 +fi + +run_one_case() { + local case_id="$1" + local case_dir="${OUTPUT_ROOT}/${OP}/${BACKEND}/${case_id}" + local log_file="${case_dir}/driver.log" + local run_py="${KT_ROOT}/run.py" + + mkdir -p "${case_dir}" + echo "==> ${ENGINE} case ${case_id}" + local kernel_args=() + if [ -n "${KERNEL_DIR}" ]; then + kernel_args+=(--kernel-dir "${KERNEL_DIR}") + fi + if [ "${ENGINE}" = "msprof" ]; then + "${KT_SCRIPT_DIR}/run_msprof.sh" \ + --output "${case_dir}/msprof" \ + "${run_py}" \ + -- \ + "${kernel_args[@]}" \ + --op "${OP}" \ + --workflow cycle \ + --backend "${BACKEND}" \ + --case "${case_id}" \ + > "${log_file}" 2>&1 + else + "${KT_SCRIPT_DIR}/run_sim.sh" \ + --output "${case_dir}" \ + --python-cmd "${PYTHON_CMD}" \ + "${run_py}" \ + -- \ + "${kernel_args[@]}" \ + --op "${OP}" \ + --workflow cycle \ + --backend "${BACKEND}" \ + --case "${case_id}" \ + > "${log_file}" 2>&1 + fi +} + +report_cycles() { + if [ "${#SUCCESS_CASE_DIRS[@]}" -eq 0 ]; then + return 0 + fi + if [ ! -f "${KT_SCRIPT_DIR}/helpers/report_cycles.py" ]; then + return 0 + fi + + report_args=(--op "${OP}") + if [ -n "${KERNEL_DIR}" ]; then + report_args+=(--kernel-dir "${KERNEL_DIR}") + fi + if [ "${#SUCCESS_CASE_DIRS[@]}" -gt 1 ]; then + report_args+=(--table) + fi + report_args+=("${SUCCESS_CASE_DIRS[@]}") + + echo "==> analyze cycles" + kt_run_python_cmd "${PYTHON_CMD}" "${KT_SCRIPT_DIR}/helpers/report_cycles.py" "${report_args[@]}" +} + +FAILED=0 +declare -a FAILED_CASES=() + +if [ "${PARALLEL_SIM}" -eq 0 ]; then + for case_id in "${CASES[@]}"; do + if ! run_one_case "${case_id}"; then + FAILED=1 + FAILED_CASES+=("${case_id}") + echo "FAIL case=${case_id} log=${OUTPUT_ROOT}/${OP}/${BACKEND}/${case_id}/driver.log" >&2 + else + SUCCESS_CASE_DIRS+=("${OUTPUT_ROOT}/${OP}/${BACKEND}/${case_id}") + echo "PASS case=${case_id} log=${OUTPUT_ROOT}/${OP}/${BACKEND}/${case_id}/driver.log" + fi + done +else + declare -A PID_TO_CASE=() + active_jobs=0 + + handle_finished_job() { + local finished_pid="$1" + local job_rc="$2" + local case_id="${PID_TO_CASE[${finished_pid}]}" + + unset 'PID_TO_CASE[$finished_pid]' + active_jobs=$((active_jobs - 1)) + + if [ "${job_rc}" -ne 0 ]; then + FAILED=1 + FAILED_CASES+=("${case_id}") + echo "FAIL case=${case_id} log=${OUTPUT_ROOT}/${OP}/${BACKEND}/${case_id}/driver.log" >&2 + else + SUCCESS_CASE_DIRS+=("${OUTPUT_ROOT}/${OP}/${BACKEND}/${case_id}") + echo "PASS case=${case_id} log=${OUTPUT_ROOT}/${OP}/${BACKEND}/${case_id}/driver.log" + fi + } + + for case_id in "${CASES[@]}"; do + while [ "${active_jobs}" -ge "${JOBS}" ]; do + finished_pid="" + if wait -n -p finished_pid; then + job_rc=0 + else + job_rc=$? + fi + handle_finished_job "${finished_pid}" "${job_rc}" + done + + run_one_case "${case_id}" & + pid=$! + PID_TO_CASE["${pid}"]="${case_id}" + active_jobs=$((active_jobs + 1)) + done + + while [ "${#PID_TO_CASE[@]}" -gt 0 ]; do + finished_pid="" + if wait -n -p finished_pid; then + job_rc=0 + else + job_rc=$? + fi + handle_finished_job "${finished_pid}" "${job_rc}" + done +fi + +echo "SUMMARY op=${OP} backend=${BACKEND} total=${#CASES[@]} failed=${#FAILED_CASES[@]}" +report_cycles +if [ "${FAILED}" -ne 0 ]; then + printf 'FAILED_CASES %s\n' "${FAILED_CASES[*]}" >&2 + exit 1 +fi diff --git a/test/kernel-test/scripts/run_msprof.sh b/test/kernel-test/scripts/run_msprof.sh new file mode 100755 index 0000000000..a5ee7b2f70 --- /dev/null +++ b/test/kernel-test/scripts/run_msprof.sh @@ -0,0 +1,88 @@ +#!/usr/bin/env bash +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +set -euo pipefail + +# User-facing entrypoint: generic msprof transport for one Python entry. + +SCRIPT_PATH="${BASH_SOURCE[0]}" +# shellcheck disable=SC1091 +SCRIPT_DIR="$(cd -- "$(dirname -- "$(realpath -- "${SCRIPT_PATH}")")" && pwd)" +source "${SCRIPT_DIR}/helpers/common.sh" + +kt_resolve_paths "${SCRIPT_PATH}" + +OUTPUT_DIR="" +SCRIPT_FILE="" +declare -a SCRIPT_ARGS=() + +while [ $# -gt 0 ]; do + case "$1" in + --output) + [ $# -lt 2 ] && { echo "--output requires a value" >&2; exit 1; } + OUTPUT_DIR="$2" + shift 2 + ;; + --help|-h) + cat < [--