diff --git a/docs/spec/core-ir.md b/docs/spec/core-ir.md index 4080726e0..fc01567d0 100644 --- a/docs/spec/core-ir.md +++ b/docs/spec/core-ir.md @@ -759,6 +759,10 @@ implicitly. Predicates on a surrounding `ComposedLayoutPattern` instead check the whole matched layout. Execution-mesh declarations may admit leading modes with `StarPattern`; operand mesh declarations keep their exact axis structure because shard attrs refer to those axes. +A wildcard in `MeshPattern.topologies` selects the innermost execution level +and binds its one-name tuple. Reusing that wildcard in an operand mesh requires +the same level; valid topology names and extents come from the module's Target +Facts, as specified in [target §11](./target.md#11-target-facts-projection). `ShardLayoutPattern` names the same `layout`, `attrs`, and `mesh` fields as `ShardLayout`: `layout` and `mesh` are nested patterns, while `attrs` remains an exact structural value. Its mesh pattern may use the bare/composed identity diff --git a/docs/spec/target.md b/docs/spec/target.md index a880a4a72..10e31e022 100644 --- a/docs/spec/target.md +++ b/docs/spec/target.md @@ -56,6 +56,7 @@ def registered_targets() -> Mapping[str, type[Target]]: ... by inheriting a backend-specific analysis base class. A missing Facts projection MUST fail when the selected analyzer requests it. - `facts_result`, `TargetFactsError`, `TopologyFacts`, `TopologyLevelFacts`, + `MemoryExecutionFacts`, `MemoryHierarchyFacts`, `ThroughputFacts`, and `PerformanceServiceFacts` MUST be importable from `tilefoundry.target` for provider implementations. - `name` MUST be a non-empty class variable declared directly by every @@ -583,6 +584,9 @@ class AppleM2Pro: unified_memory_bandwidth_bytes_per_second: int _unit_flops: tuple[tuple[str, tuple[tuple[DType, int], ...]], ...] + performance_amx_f16: int | None + performance_host_lane_events: int | None + def throughput_for(self, unit: str, dtype: DType) -> int: ... ``` @@ -608,6 +612,17 @@ class AppleM2Pro: - `throughput_for` MUST return a measured per-unit throughput recorded in the installed document, and MUST raise an actionable error for a unit or compute DType with no measured entry rather than return an estimate. + - `performance_amx_f16` and `performance_host_lane_events` MUST read their + available or unavailable entries from the installed document. Estimates MUST + state their sources and assumptions separately from measured throughput and + MUST NOT populate `throughput_for`'s measured-rate map. + - AMX `PerformanceServiceFacts` MUST use the measured NEON FP32 rate for host + arithmetic. FP16 service MUST divide the documented per-AMX-unit estimate + across core streams using the AMX unit count and performance-core count. + Non-FLOP work MUST use the documented per-core lane-event estimate. + Global-memory bandwidth MUST be shared across the same core streams. + Unrecorded DTypes and services MUST remain absent. + ## 9. `AmxTarget` @@ -822,9 +837,19 @@ class TopologyFacts: def level(self, name: str | None) -> TopologyLevelFacts | None: ... def parallel(self) -> TopologyLevelFacts | None: ... + + +class MemoryExecutionFacts: + topologies: tuple[str, ...] ``` - constraints: + - `MemoryExecutionFacts.topologies` MUST list the Target's execution levels + admitted for synchronous Copy and Fill. Each name MUST occur in + `TopologyFacts.topologies`. CUDA admits its thread level; AMX admits its + core and coprocessor levels. An empty list admits no memory execution scope. + Shared instruction matching MUST specialize the memory mesh wildcard from + this projection before checking operands. - `max_logical_units` is the program-side declaration limit. `max_physical_units` is the machine-side count of positions the level divides over. For CUDA CTAs it counts one CTA per SM for this model; it is diff --git a/docs/spec/tir.md b/docs/spec/tir.md index 69494d81b..80a0ec15e 100644 --- a/docs/spec/tir.md +++ b/docs/spec/tir.md @@ -582,6 +582,30 @@ class TensorView(Op): - `T.ptr_of` MAY point at an allocated `ShardTensor`; the view rebuilds over the engine pointer rather than reusing the existing shard layout. +Apple AMX contracts are in `tilefoundry.ir.tir.amx.instructions`: +`T.amx_load_x` and `T.amx_load_y` transfer 32 FP16 values from global memory +to registers; `T.amx_fma16` accumulates their outer product into a 32-by-32 +FP32 tile; `T.amx_store` transfers that tile to global memory. The accumulator +uses `Layout((32, (16, 2)), (32, (1, 16)))` to describe physical Z row order. +These operations execute on the `amx` mesh and require `amx_outer_product`. +Native wrappers for the corresponding loads, FMA and physical Z-row stores, +plus AMX state enable/disable, live in `include/tilefoundry/runtime/amx/instructions.h`. + +CPU tensor kernels in `tilefoundry.ir.tir.cpu.tensor` explicitly select host +implementations for the tensor operations used by hybrid CPU/AMX schedules. +They execute in a single `core` participant scope, read tensor operands and +write an explicit global-memory destination. Attributes and output types are +checked against the corresponding HIR operation. Access relations, evaluation +and arithmetic costs retain that operation's semantics; these are software +kernel contracts rather than AMX ISA operations. +These contracts require the Target's `cpu_tensor` capability; they MUST NOT +appear among CUDA or topology-free CPU target instructions. + +Copy and Fill use a shared mesh wildcard, specialized from the Target's +`MemoryExecutionFacts`. They require a contiguous run and sharded register +operands on that same level. The module's Target Facts validate topology names +and extents; memory declarations do not enumerate backend topology names. + ##### Copy ```python class Copy(Op): diff --git a/examples/qwen_prefill-amx/.gitignore b/examples/qwen_prefill-amx/.gitignore new file mode 100644 index 000000000..30399290e --- /dev/null +++ b/examples/qwen_prefill-amx/.gitignore @@ -0,0 +1,5 @@ +/artifacts/ +/prepared/ +/kernels/.build/ +__pycache__/ +*.py[cod] diff --git a/examples/qwen_prefill-amx/README.md b/examples/qwen_prefill-amx/README.md new file mode 100644 index 000000000..e012bc715 --- /dev/null +++ b/examples/qwen_prefill-amx/README.md @@ -0,0 +1,205 @@ +# Qwen text prefill on TileFoundry — Apple AMX kernels + +Runtime execution and CLI integration rechecked on **2026-10-19**, in the +development checkout based on `3cf660e`. + +Qwen3-1.7B and Qwen3.5-4B text prefill on one M2 Pro: token IDs in, last-position +logits out. Batch one, no history. The large projections use FP16 AMX outer +products with FP32 accumulation; attention, GDN and the other host operations +remain in the supplied native llama.cpp runtime. The model, its projection schedules and its runtime are all present +in this directory. + +Section 4 compares this AMX runtime with local llama.cpp Q8_0 and F16 execution. + +--- + +## 1. Environment + +This is the environment used for the current execution checks. The entry point +uses the installed packages and builds its FFI entry on first use. + +| | | +|---|---| +| Machine | Apple M2 Pro, 32 GiB unified memory | +| Execution | CPU and Apple AMX, eight logical workers | +| Python | 3.12.13 | +| `tilefoundry` | this development checkout, commit `3cf660e` | +| Torch | 2.13.0, for resources and FFI tensors | +| Native model runtime | the supplied experimental llama.cpp checkout, selected by `LLAMA_CPP` | +| Other packages | `gguf` 0.19.0, `safetensors` 0.8.0, `apache-tvm-ffi` 0.1.12 | +| Native toolchain | CMake and Xcode Command Line Tools; `clang++` and SDK located through `xcrun` | +| Weights | Qwen3-1.7B or Qwen3.5-4B, Q8_0 GGUF | + +Q8_0 weights are dequantized to FP16 during preparation. That preserves the +checkpoint's quantization error and changes the projection arithmetic from +the native Q8 path. Prepared weights take more memory than the GGUF file. + +## 2. How this was produced + +The task for this example, stated as a prompt: + +```text +Get Qwen3-1.7B and Qwen3.5-4B text prefill running on TileFoundry. +Weights: the Q8_0 GGUF checkpoints. +Backend: Apple AMX on M2 Pro, eight workers. +Given: the existing user-provided AMX runtime, including its native +instruction wrappers and projection implementation. Use it as the +starting point for the C++ runtime and AMX instruction contracts. +Preserve its existing model execution path and host operations. +Scope: batch one, token IDs in, last-position logits out. + +Deliver model.py (whole-model HIR), scheduled.py (scheduled HIR), +runtime_model.py (Python runtime modules), kernels/ (C++ AMX kernels), +and run.py (the entry point). + +Make the weight layouts, AMX instruction choices and worker partition +explicit in the scheduled HIR. Put the AMX instruction support in +TileFoundry itself. Both HIR modules must be usable through the CLI. + +The runtime module should load prepared weights and execute directly. +Build the native FFI entry on first use against the supplied runtime. Verify execution without rerunning +performance experiments or adding tests. +``` + +## 3. How to use it + +From the repository root: + +```sh +export LLAMA_CPP=/path/to/the/supplied/llama.cpp +python examples/qwen_prefill-amx/run.py --model 1.7b --checkpoint /path/to/Qwen3-1.7B-Q8_0.gguf +python examples/qwen_prefill-amx/run.py --model 4b --checkpoint /path/to/Qwen3.5-4B-Q8_0.gguf +``` + +The first run prepares packed projection weights under `prepared//` +beside `run.py`, builds the supplied native libraries and FFI entry under +`kernels/.build/`, and loads the model. Subsequent runs reuse the outputs. +`LLAMA_CPP` must contain the provided AMX implementation and its +`tf_amx_load_prepared` / `tf_amx_clear_prepared` exports; a standard upstream +checkout alone does not provide that execution path. `LLAMA_BUILD` can select +an existing CPU/Accelerate build. No source checkout is modified automatically. +The projection kernel is in the supplied `ggml/src/ggml-blas/shared_fp16.hpp`; +`kernels/prefill.cpp` provides its FFI and native model-driver bridge. + +`--prepared` can reuse the original experiment's packed `.bin` directory +for the same checkpoint; its metadata does not need a separate manifest. + +| Flag | Meaning | +|---|---| +| `--model 1.7b\|4b` | Which model to construct; required | +| `--checkpoint PATH` | Source Q8_0 GGUF; required | +| `--prepared DIR` | Prepared resource directory; defaults to this example's `prepared//` | +| `--sequence N` | Static prompt length, a multiple of 32 from 128 to 8192; default 128 | + +`run.py` feeds a fixed valid token sequence through one prefill and prints the +logits shape and next-token ID. It stops there: the example covers text prefill, +with no decode loop, MTP or vision path. + +```text +model.py the whole-model HIR and checkpoint weight converters +scheduled.py the same model with explicit AMX projection schedules +runtime_model.py RuntimeModule adapter for the original native model driver +kernels/prefill.cpp the original TVM FFI model and projection entries +kernels/__init__.py first-use native build and load +weights.py GGUF reading and checkpoint-to-module weight aliases +config.py the two models' architecture parameters +run.py preparation, loading and one runtime call +``` + +The runtime also accepts caller-supplied token IDs. After preparing the weights, +run this from the example directory: + +```python +import torch +from config import config +from runtime_model import QwenPrefill +from weights import PreparedModelResource + +resource = PreparedModelResource("/path/to/Qwen3-1.7B-Q8_0.gguf", "prepared/1.7b") +with QwenPrefill(config("1.7b"), sequence=128) as runtime: + runtime.load(resource) + token_ids = torch.arange(128, dtype=torch.int32) + logits = runtime(token_ids) +``` + +`model.build(c, 128)` constructs the authored HIR; +`scheduled.make_scheduled(c, 128)` constructs the scheduled HIR used by +`runtime.module` for inspection. Execution calls the original native model +graph rather than evaluating the IR. Use one checkpoint per process, as in +the supplied runtime. The HIR keeps weights on the root module and passes them +explicitly to shared child definitions, retaining the root's weight converters. +Each source file also exposes `Qwen3_1_7B` and `Qwen3_5_4B` as +CLI-selectable modules with a 128-token prompt: + +```sh +tilefoundry analyze examples/qwen_prefill-amx/model.py:Qwen3_1_7B /tmp/qwen-hir.txt --compute-cost --memory +tilefoundry schedule candidates examples/qwen_prefill-amx/model.py:Qwen3_1_7B /tmp/qwen-candidates.txt +tilefoundry analyze examples/qwen_prefill-amx/scheduled.py:Qwen3_1_7B /tmp/qwen-scheduled.txt --compute-cost --memory +tilefoundry schedule finalize examples/qwen_prefill-amx/scheduled.py:Qwen3_1_7B /tmp/qwen-tir.py +``` + +To generate the modeled performance timeline: + +```sh +tilefoundry analyze examples/qwen_prefill-amx/scheduled.py:Qwen3_1_7B /tmp/qwen-performance.txt --performance +``` + +Use `:Qwen3_5_4B` for the hybrid model. Analysis and finalization run on the +model descriptions and do not require weights or the native extension. + +## 4. Where it stands + +Measured on **2026-10-19**, Apple M2 Pro, 32 GiB unified memory, CPU-only, eight +threads/workers. Our implementation uses the supplied packed-FP16 AMX runtime +for eligible large projections and retains its native C++ model execution. +The AMX column measures that modified llama.cpp execution path; the model +driver and host operations still come from llama.cpp. +llama.cpp is the local `e5819f0dd` working tree (build 2, existing BLAS/AMX edits), +rebuilt in Release with AppleClang 21.0.0. Accelerate, BLAS, native ARM +instructions, CPU repacking and Flash Attention are enabled; Metal and KleidiAI +are disabled. Native baselines unset `TF_AMX_MODE` / `TF_AMX_CHECKPOINT`, allow +BLAS dispatch (`-nopo 0`) and disable thread-pool polling (`--poll 0`). + +CPU builds with and without OpenMP were checked. One configuration was selected +per model and weight format by the lower sum of medians across all four lengths, +then retained for that entire group. Qwen3-1.7B Q8_0 uses the native thread pool; +the other native groups enable OpenMP in the CPU and BLAS backends. The local +CMake configuration required an explicit `GGML_USE_OPENMP` definition and +OpenMP linkage on the BLAS target to activate its existing parallel-conversion +branch; the llama.cpp C++ source was unchanged. OpenMP uses eight threads, +passive waiting and `KMP_BLOCKTIME=0`; all runners set +`VECLIB_MAXIMUM_THREADS=8`. + +Batch-one text prefill, zero history, last-position logits, batch/microbatch +2048 and F16 K/V cache. Each length has one full-length warmup and five timed +calls, run serially. All latencies below are **median milliseconds**. Preparation, +compilation and loading are excluded. Our timer includes state reset, FFI and +final-logit copy; llama-bench clears state before its timer. No decode, MTP or +vision is timed. CPU affinity and clocks were not locked. + +| Model | Prompt tokens | Our AMX (ms) | llama.cpp Q8_0 (ms) | llama.cpp F16 (ms) | vs Q8_0 | vs F16 | +|---|---:|---:|---:|---:|---:|---:| +| Qwen3-1.7B | 128 | 192.90 | 220.18 | 368.14 | 1.14× | 1.91× | +| Qwen3-1.7B | 512 | 741.72 | 901.61 | 988.67 | 1.22× | 1.33× | +| Qwen3-1.7B | 1024 | 1655.07 | 2025.08 | 1827.90 | 1.22× | 1.10× | +| Qwen3-1.7B | 2048 | 3711.20 | 4742.70 | 3732.24 | 1.28× | 1.01× | +| Qwen3.5-4B | 128 | 525.31 | 684.32 | 1107.83 | 1.30× | 2.11× | +| Qwen3.5-4B | 512 | 2085.52 | 2743.58 | 2727.69 | 1.32× | 1.31× | +| Qwen3.5-4B | 1024 | 4126.70 | 5899.10 | 5241.54 | 1.43× | 1.27× | +| Qwen3.5-4B | 2048 | 9653.82 | 11808.71 | 10596.87 | 1.22× | 1.10× | + +Ratios are **llama.cpp latency / our AMX latency**: above 1× means our runtime +is faster. Results close to 1× indicate similar performance within the limits +of this desktop measurement. + +All paths start from the same Q8_0 GGUF checkpoints. Our runtime prepares FP16 +projection weights; native F16 files are expanded offline using +`llama-quantize --allow-requantize ... F16 8`, preserving source quantization +error and auxiliary F32 tensors. Native BLAS-selected F16 matmuls convert +weights to F32 and call `cblas_sgemm`. + +The measured entry dispatches through the supplied AMX backend: each prefill +has 193 AMX matrix calls for 1.7B and 197 for 4B. Across all eight measured shapes, +its logits match the archived experiment outputs exactly (maximum absolute +difference 0). The table uses these current-entry measurements. First-use +weight preparation and execution were also checked on 1.7B. diff --git a/examples/qwen_prefill-amx/config.py b/examples/qwen_prefill-amx/config.py new file mode 100644 index 000000000..92a79866f --- /dev/null +++ b/examples/qwen_prefill-amx/config.py @@ -0,0 +1,76 @@ +"""Architecture parameters for the two Qwen text prefill examples.""" + +CONFIGS = { + "1.7b": { + "hidden_size": 2048, + "intermediate_size": 6144, + "vocab_size": 151936, + "num_attention_heads": 16, + "num_key_value_heads": 8, + "head_dim": 128, + "rms_norm_eps": 1e-06, + "num_hidden_layers": 28, + "rope_theta": 1000000, + }, + "4b": { + "hidden_size": 2560, + "intermediate_size": 9216, + "vocab_size": 248320, + "num_attention_heads": 16, + "num_key_value_heads": 4, + "head_dim": 256, + "rms_norm_eps": 1e-06, + "num_hidden_layers": 32, + "attn_output_gate": True, + "layer_types": [ + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + ], + "linear_num_key_heads": 16, + "linear_num_value_heads": 32, + "linear_key_head_dim": 128, + "linear_value_head_dim": 128, + "linear_conv_kernel_dim": 4, + "rope_parameters": { + "mrope_interleaved": True, + "mrope_section": [11, 11, 10], + "rope_type": "default", + "rope_theta": 10000000, + "partial_rotary_factor": 0.25, + }, + }, +} + + +def config(model): + return dict(CONFIGS[model]) diff --git a/examples/qwen_prefill-amx/kernels/__init__.py b/examples/qwen_prefill-amx/kernels/__init__.py new file mode 100644 index 000000000..bce35f28f --- /dev/null +++ b/examples/qwen_prefill-amx/kernels/__init__.py @@ -0,0 +1,118 @@ +"""Build the FFI entry against the supplied experimental llama.cpp runtime.""" + +import os +from pathlib import Path +import platform +import shutil +import subprocess +import sys +import tvm_ffi +from tvm_ffi import libinfo + +_module = None + + +def ops(): + global _module + if _module is not None: + return _module + if platform.system() != "Darwin" or platform.machine() != "arm64": + raise RuntimeError("This runtime requires Apple Silicon macOS") + source = os.environ.get("LLAMA_CPP") + if not source: + raise RuntimeError( + "Set LLAMA_CPP to the supplied experimental llama.cpp checkout" + ) + source = Path(source).expanduser().resolve() + native = source / "ggml/src/ggml-blas" + if not (native / "shared_fp16.hpp").is_file(): + raise RuntimeError( + "LLAMA_CPP must contain the supplied AMX runtime (shared_fp16.hpp)" + ) + root = Path(__file__).resolve().parent + build = ( + Path(os.environ.get("LLAMA_BUILD", root / ".build/llama")) + .expanduser() + .resolve() + ) + cmake = shutil.which("cmake") or str(Path(sys.executable).parent / "cmake") + subprocess.run( + [ + cmake, + "-S", + str(source), + "-B", + str(build), + "-DCMAKE_BUILD_TYPE=Release", + "-DGGML_METAL=OFF", + "-DGGML_BLAS=ON", + "-DGGML_BLAS_VENDOR=Apple", + "-DLLAMA_BUILD_TESTS=OFF", + "-DLLAMA_BUILD_SERVER=OFF", + ], + check=True, + stdout=subprocess.DEVNULL, + ) + subprocess.run( + [cmake, "--build", str(build), "--target", "llama", "ggml-blas", "-j", "8"], + check=True, + stdout=subprocess.DEVNULL, + ) + library = build / "bin/libggml-blas.dylib" + symbols = subprocess.check_output(["nm", "-gU", str(library)], text=True) + if ( + "_tf_amx_load_prepared" not in symbols + or "_tf_amx_clear_prepared" not in symbols + ): + raise RuntimeError( + "The supplied runtime must export tf_amx_load_prepared and tf_amx_clear_prepared" + ) + out = root / ".build/prefill.dylib" + out.parent.mkdir(parents=True, exist_ok=True) + ffi = Path(libinfo.find_libtvm_ffi()) + dependencies = ( + root / "prefill.cpp", + ffi, + library, + build / "bin/libllama.dylib", + *native.glob("*.hpp"), + ) + if not out.exists() or any( + p.stat().st_mtime > out.stat().st_mtime for p in dependencies + ): + clang = subprocess.check_output( + ["xcrun", "--find", "clang++"], text=True + ).strip() + sdk = subprocess.check_output(["xcrun", "--show-sdk-path"], text=True).strip() + subprocess.run( + [ + clang, + "-std=c++17", + "-O3", + "-shared", + "-fPIC", + "-isysroot", + sdk, + "-framework", + "Accelerate", + "-I" + str(Path(tvm_ffi.__file__).parent / "include"), + "-I" + str(source / "include"), + "-I" + str(source / "ggml/include"), + "-I" + str(native), + str(root / "prefill.cpp"), + str(ffi), + "-L" + str(build / "bin"), + "-lllama", + "-lggml-cpu", + "-lggml-base", + "-lggml", + "-lggml-blas", + "-Wl,-rpath," + str(build / "bin"), + "-Wl,-rpath," + str(ffi.parent), + "-o", + str(out), + ], + check=True, + ) + _module = tvm_ffi.load_module(str(out)) + return _module diff --git a/examples/qwen_prefill-amx/kernels/prefill.cpp b/examples/qwen_prefill-amx/kernels/prefill.cpp new file mode 100644 index 000000000..2bb75ed49 --- /dev/null +++ b/examples/qwen_prefill-amx/kernels/prefill.cpp @@ -0,0 +1,196 @@ +/** TVM FFI entry for the supplied CPU/AMX model runtime. */ +#include +#include +#include +#include "llama.h" +#include "ggml-cpu.h" +#include "ggml-backend.h" +#include "shared_fp16.hpp" +#include +#include +#include +#include +#include + +extern "C" void tf_amx_load_prepared(const char *, const char *, int, int); +extern "C" void tf_amx_clear_prepared(); +using tvm::ffi::String; +using tvm::ffi::TensorView; +static std::mutex mutex; +struct Session { + llama_model *model = nullptr; + llama_context *ctx = nullptr; + ggml_threadpool_t workers = nullptr; + int capacity = 0; + ~Session() { + if (ctx) { + llama_detach_threadpool(ctx); + llama_free(ctx); + } + if (workers) + ggml_threadpool_free(workers); + if (model) + llama_model_free(model); + } +}; +static std::unique_ptr session; +static std::string process_checkpoint; +static int64_t process_threads = 0; +static void tensor(const TensorView &t, int ndim, int code, int bits) { + auto dtype = t.dtype(); + if (t.device().device_type != kDLCPU || t.ndim() != ndim || + dtype.code != code || dtype.bits != bits || dtype.lanes != 1 || + !t.IsContiguous()) + throw std::invalid_argument("Wrong CPU tensor rank/dtype/strides"); +} +static void *ptr(const TensorView &t) { + return (char *)t.data_ptr() + t.byte_offset(); +} +int64_t load_model(String path, int64_t threads, int64_t capacity, + int64_t mode) { + std::lock_guard guard(mutex); + if (session) + throw std::runtime_error( + "Close the loaded model before loading another"); + if (threads < 1 || threads > 12 || capacity < 128 || capacity > 8192 || + (mode != 0 && mode != 7)) + throw std::invalid_argument("Invalid model options"); + std::string filename(path), t = std::to_string(threads), + m = std::to_string(mode); + if (!process_checkpoint.empty() && + (filename != process_checkpoint || threads != process_threads)) + throw std::runtime_error( + "Use a separate process for another checkpoint/worker count"); + setenv("TF_AMX_THREADS", t.c_str(), 1); + setenv("TF_AMX_MODE", m.c_str(), 1); + setenv("TF_AMX_CHECKPOINT", filename.c_str(), 1); + unsetenv("TF_SHARED_RESOURCE"); + ggml_backend_load_all(); + llama_backend_init(); + auto value = std::make_unique(); + value->capacity = capacity; + auto mp = llama_model_default_params(); + mp.n_gpu_layers = 0; + mp.use_extra_bufts = true; + value->model = llama_model_load_from_file(filename.c_str(), mp); + if (!value->model) + throw std::runtime_error("Cannot load GGUF"); + auto cp = llama_context_default_params(); + cp.n_ctx = capacity; + cp.n_batch = capacity; + cp.n_ubatch = capacity; + cp.n_threads = threads; + cp.n_threads_batch = threads; + cp.type_k = GGML_TYPE_F16; + cp.type_v = GGML_TYPE_F16; + cp.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED; + value->ctx = llama_init_from_model(value->model, cp); + if (!value->ctx) + throw std::runtime_error("Cannot create model context"); + auto tp = ggml_threadpool_params_default(threads); + tp.poll = 0; + value->workers = ggml_threadpool_new(&tp); + if (!value->workers) + throw std::runtime_error("Cannot create worker pool"); + llama_attach_threadpool(value->ctx, value->workers, value->workers); + int64_t vocab = llama_vocab_n_tokens(llama_model_get_vocab(value->model)); + session = std::move(value); + process_checkpoint = filename; + process_threads = threads; + return vocab; +} +void load_weight(String name, String path, int64_t n, int64_t k) { + std::lock_guard guard(mutex); + if (!session) + throw std::runtime_error("Load the model first"); + if (n <= 0 || k <= 0 || n % 64 || k % 32 || n > INT32_MAX || + k > INT32_MAX || n * k > INT32_MAX) + throw std::invalid_argument("Invalid weight dimensions"); + std::string a(name), b(path); + tf_amx_load_prepared(a.c_str(), b.c_str(), n, k); +} +void close_model() { + std::lock_guard guard(mutex); + session.reset(); + tf_amx_clear_prepared(); +} +void prefill(TensorView tokens, TensorView logits) { + std::lock_guard guard(mutex); + if (!session) + throw std::runtime_error("Model not loaded"); + tensor(tokens, 1, kDLInt, 32); + tensor(logits, 1, kDLFloat, 32); + int64_t length = tokens.size(0); + if (length < 1 || length > session->capacity) + throw std::invalid_argument("Token length out of bounds"); + int n = length, + v = llama_vocab_n_tokens(llama_model_get_vocab(session->model)); + if (n < 1 || n > session->capacity || logits.size(0) != v) + throw std::invalid_argument("Token/logit shape out of bounds"); + auto *ids = (int32_t *)ptr(tokens); + for (int i = 0; i < n; i++) + if (ids[i] < 0 || ids[i] >= v) + throw std::invalid_argument("Invalid token ID"); + llama_memory_clear(llama_get_memory(session->ctx), true); + auto *batch = llama_batch_ext_init(session->ctx); + for (int i = 0; i < n; i++) { + int index = llama_batch_ext_add_token(batch, 0, ids[i]); + llama_pos p = i; + llama_batch_ext_set_pos(batch, index, &p); + } + llama_batch_ext_set_output_logits(batch, n - 1, true); + int status = llama_process(session->ctx, LLAMA_PROCESS_TYPE_DECODE, batch); + llama_batch_ext_free(batch); + if (status) + throw std::runtime_error("Prefill failed: " + std::to_string(status)); + std::memcpy(ptr(logits), llama_get_logits_ith(session->ctx, -1), + v * sizeof(float)); +} +void linear(TensorView x, TensorView w, TensorView out) { + std::lock_guard guard(mutex); + tensor(x, 2, kDLFloat, 32); + tensor(w, 3, kDLFloat, 16); + tensor(out, 2, kDLFloat, 32); + if (w.size(0) <= 0 || w.size(0) > INT32_MAX / 64) + throw std::invalid_argument("Invalid N panels"); + int64_t m = x.size(0), k = x.size(1), n = w.size(0) * w.size(2); + if (m <= 0 || k <= 0 || m > INT32_MAX - 31 || k % 32 || w.size(1) != k || + (w.size(2) != 64 && w.size(2) != 32) || out.size(0) != m || + out.size(1) != n || n * k > INT32_MAX || m * n > INT32_MAX || + ((m + 31) / 32 * 32) * k > INT32_MAX) + throw std::invalid_argument("Invalid linear shape"); + shared_fp16::Weight weight((const _Float16 *)ptr(w), n, k, w.size(2)); + shared_fp16::prefill((float *)ptr(x), m, weight, (float *)ptr(out), + w.size(2), true); +} +int64_t token_count(String text) { + std::lock_guard guard(mutex); + if (!session) + throw std::runtime_error("Model not loaded"); + std::string str(text); + if (str.size() > INT32_MAX) + throw std::invalid_argument("Text too long"); + return -llama_tokenize(llama_model_get_vocab(session->model), str.data(), + str.size(), nullptr, 0, true, true); +} +void tokenize(String text, TensorView out) { + std::lock_guard guard(mutex); + if (!session) + throw std::runtime_error("Model not loaded"); + tensor(out, 1, kDLInt, 32); + std::string str(text); + if (str.size() > INT32_MAX || out.size(0) > INT32_MAX) + throw std::invalid_argument("Text too long"); + int actual = llama_tokenize(llama_model_get_vocab(session->model), + str.data(), str.size(), (llama_token *)ptr(out), + out.size(0), true, true); + if (actual != out.size(0)) + throw std::invalid_argument("Token output length mismatch"); +} +TVM_FFI_DLL_EXPORT_TYPED_FUNC(token_count, token_count); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(tokenize, tokenize); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(load_model, load_model); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(load_weight, load_weight); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(close_model, close_model); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(prefill, prefill); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(linear, linear); diff --git a/examples/qwen_prefill-amx/model.py b/examples/qwen_prefill-amx/model.py new file mode 100644 index 000000000..ba527ecfb --- /dev/null +++ b/examples/qwen_prefill-amx/model.py @@ -0,0 +1,453 @@ +"""Whole text-model Prefill HIR, with all layers and checkpoint-owned weights. + +Every compute and composition step is an authored @func, including the whole +Prefill entry. The entry lists static instance calls; layer bodies remain shared. GGUF's tiled GDN head order, +materialized RMS scales and negative exp(A_log) are retained explicitly. +""" + +import math +from tilefoundry import module, func +from tilefoundry.dsl import Tensor, ConstTensor, ModuleList, Topology, tf +from tilefoundry.target import AmxTarget + + +def projection(M, N, K): + @module(entry="forward") + class Linear: + @func + def forward(x: Tensor[(M, K), "f32"], w: ConstTensor[(K, N), "f16"]): + return tf.matmul(tf.cast(x, dtype="f16"), w, out_dtype="f32") + + @forward.converter("w") + def _(raw: ConstTensor[(N, K), "f16"]): + return tf.transpose(raw, perm=(1, 0)) + + return Linear + + +def build(config, S, *, linear=projection): + H = config["hidden_size"] + F = config["intermediate_size"] + V = config["vocab_size"] + HQ = config["num_attention_heads"] + HK = config["num_key_value_heads"] + D = config["head_dim"] + EPS = config["rms_norm_eps"] + LAYERS = config["num_hidden_layers"] + GATED = config.get("attn_output_gate", False) + R = int(D * config.get("rope_parameters", {}).get("partial_rotary_factor", 1)) + Q = HQ * D + KV = HK * D + REPEAT = HQ // HK + ATTN_SCALE = math.sqrt(D) + + @module(entry="forward") + class Norm: + @func + def forward(x: Tensor[(S, H), "f32"], gamma: ConstTensor[(H,), "f32"]): + return tf.rms_norm(x, gamma, eps=EPS) + + @module(entry="forward") + class MLP: + gate = ModuleList(linear(S, F, H)) + up = ModuleList(linear(S, F, H)) + down = ModuleList(linear(S, H, F)) + + @func + def combine(g: Tensor[(S, F), "f32"], u: Tensor[(S, F), "f32"]): + return tf.silu(g) * u + + @func + def forward(x: Tensor[(S, H), "f32"]): + return down(combine(gate(x), up(x))) + + @module() + class AttentionMath: + @func + def query_norm( + q: Tensor[(1, S, HQ, D), "f32"], gamma_q: ConstTensor[(D,), "f32"] + ): + return tf.rms_norm(q, gamma_q, eps=EPS) + + @func + def key_norm( + k: Tensor[(1, S, HK, D), "f32"], gamma_k: ConstTensor[(D,), "f32"] + ): + return tf.rms_norm(k, gamma_k, eps=EPS) + + if R == D: + + @func + def query_rope( + q: Tensor[(1, S, HQ, D), "f32"], + cos: Tensor[(S, R), "f32"], + sin: Tensor[(S, R), "f32"], + ): + rot = q[:, :, :, :R] + turned = tf.concat( + [-rot[:, :, :, R // 2 :], rot[:, :, :, : R // 2]], axis=-1 + ) + head = rot * tf.reshape( + cos, new_shape=(1, S, 1, R) + ) + turned * tf.reshape(sin, new_shape=(1, S, 1, R)) + return head + + @func + def key_rope( + k: Tensor[(1, S, HK, D), "f32"], + cos: Tensor[(S, R), "f32"], + sin: Tensor[(S, R), "f32"], + ): + rot = k[:, :, :, :R] + turned = tf.concat( + [-rot[:, :, :, R // 2 :], rot[:, :, :, : R // 2]], axis=-1 + ) + head = rot * tf.reshape( + cos, new_shape=(1, S, 1, R) + ) + turned * tf.reshape(sin, new_shape=(1, S, 1, R)) + return head + + else: + + @func + def query_rope( + q: Tensor[(1, S, HQ, D), "f32"], + cos: Tensor[(S, R), "f32"], + sin: Tensor[(S, R), "f32"], + ): + rot = q[:, :, :, :R] + turned = tf.concat( + [-rot[:, :, :, R // 2 :], rot[:, :, :, : R // 2]], axis=-1 + ) + head = rot * tf.reshape( + cos, new_shape=(1, S, 1, R) + ) + turned * tf.reshape(sin, new_shape=(1, S, 1, R)) + return tf.concat([head, q[:, :, :, R:]], axis=-1) + + @func + def key_rope( + k: Tensor[(1, S, HK, D), "f32"], + cos: Tensor[(S, R), "f32"], + sin: Tensor[(S, R), "f32"], + ): + rot = k[:, :, :, :R] + turned = tf.concat( + [-rot[:, :, :, R // 2 :], rot[:, :, :, : R // 2]], axis=-1 + ) + head = rot * tf.reshape( + cos, new_shape=(1, S, 1, R) + ) + turned * tf.reshape(sin, new_shape=(1, S, 1, R)) + return tf.concat([head, k[:, :, :, R:]], axis=-1) + + @func + def attend( + q: Tensor[(1, S, HQ, D), "f32"], + k: Tensor[(1, S, HK, D), "f32"], + v: Tensor[(1, S, HK, D), "f32"], + mask: Tensor[(S, S), "f32"], + ): + qe = tf.transpose(q, perm=(0, 2, 1, 3)) + ke = tf.transpose( + tf.repeat_interleave(k, repeats=REPEAT, axis=2), perm=(0, 2, 3, 1) + ) + ve = tf.transpose( + tf.repeat_interleave(v, repeats=REPEAT, axis=2), perm=(0, 2, 1, 3) + ) + score = tf.matmul(qe, ke, out_dtype="f32") / ATTN_SCALE + mask + prob = tf.softmax(score, axis=-1) + return tf.reshape( + tf.transpose(tf.matmul(prob, ve, out_dtype="f32"), perm=(0, 2, 1, 3)), + new_shape=(S, Q), + ) + + @module(entry="query_norm") + class QueryNorm: + query_norm = AttentionMath.query_norm + + @module(entry="key_norm") + class KeyNorm: + key_norm = AttentionMath.key_norm + + @module(entry="forward") + class FullAttention: + query = ModuleList(linear(S, Q * (2 if GATED else 1), H)) + key = ModuleList(linear(S, KV, H)) + value = ModuleList(linear(S, KV, H)) + output = ModuleList(linear(S, H, Q)) + query_norm = ModuleList(QueryNorm) + key_norm = ModuleList(KeyNorm) + query_rope = AttentionMath.query_rope + key_rope = AttentionMath.key_rope + attend = AttentionMath.attend + if GATED: + + @func + def split_query(x: Tensor[(S, 2 * Q), "f32"]): + packed = tf.reshape(x, new_shape=(S, HQ, 2 * D)) + indices = tf.arange(Tensor[(D,), "i64"], start=0) + return tf.reshape( + tf.index_select(packed, indices, dim=2), new_shape=(1, S, HQ, D) + ) + + @func + def split_gate(x: Tensor[(S, 2 * Q), "f32"]): + packed = tf.reshape(x, new_shape=(S, HQ, 2 * D)) + indices = tf.arange(Tensor[(D,), "i64"], start=D) + return tf.reshape( + tf.index_select(packed, indices, dim=2), new_shape=(S, Q) + ) + + @func + def gate_output(x: Tensor[(S, Q), "f32"], gate: Tensor[(S, Q), "f32"]): + return x * tf.sigmoid(gate) + else: + + @func + def split_query(x: Tensor[(S, Q), "f32"]): + return tf.reshape(x, new_shape=(1, S, HQ, D)) + + @func + def split_kv(x: Tensor[(S, KV), "f32"]): + return tf.reshape(x, new_shape=(1, S, HK, D)) + + if GATED: + + @func + def forward( + x: Tensor[(S, H), "f32"], + cos: Tensor[(S, R), "f32"], + sin: Tensor[(S, R), "f32"], + mask: Tensor[(S, S), "f32"], + ): + qg = query(x) + q = query_rope(query_norm(split_query(qg)), cos, sin) + k = key_rope(key_norm(split_kv(key(x))), cos, sin) + v = split_kv(value(x)) + out = attend(q, k, v, mask) + return output(gate_output(out, split_gate(qg))) + else: + + @func + def forward( + x: Tensor[(S, H), "f32"], + cos: Tensor[(S, R), "f32"], + sin: Tensor[(S, R), "f32"], + mask: Tensor[(S, S), "f32"], + ): + qg = query(x) + q = query_rope(query_norm(split_query(qg)), cos, sin) + k = key_rope(key_norm(split_kv(key(x))), cos, sin) + v = split_kv(value(x)) + out = attend(q, k, v, mask) + return output(out) + + if "linear_num_key_heads" in config: + NK = config["linear_num_key_heads"] + NV = config["linear_num_value_heads"] + DK = config["linear_key_head_dim"] + DV = config["linear_value_head_dim"] + CK = config["linear_conv_kernel_dim"] + KD = NK * DK + VD = NV * DV + CV = 2 * KD + VD + VR = NV // NK + GDN_SCALE = math.sqrt(DK) + + @module(entry="forward") + class GDNMath: + @func + def forward( + qkv: Tensor[(S, CV), "f32"], + z: Tensor[(S, VD), "f32"], + alpha: Tensor[(S, NV), "f32"], + beta_raw: Tensor[(S, NV), "f32"], + conv: ConstTensor[(CV, CK), "f32"], + decay: ConstTensor[(NV,), "f32"], + dt_bias: ConstTensor[(NV,), "f32"], + gamma: ConstTensor[(DV,), "f32"], + ): + padded = tf.concat([tf.zeros(Tensor[(CK - 1, CV), "f32"]), qkv], axis=0) + mixed = tf.zeros(Tensor[(S, CV), "f32"]) + for t in range(S): + window = padded[t : t + CK, :] + row = tf.silu( + tf.reduce( + window * tf.transpose(conv, perm=(1, 0)), + axes=(0,), + keepdim=True, + kind="sum", + ) + ) + mixed = tf.insert_slice(mixed, row, (t, 0)) + q = tf.reshape(mixed[:, :KD], new_shape=(S, NK, DK)) + k = tf.reshape(mixed[:, KD : 2 * KD], new_shape=(S, NK, DK)) + # GGUF uses tiled V-head order. Repeat the K-head bank, rather + # than duplicating each K head consecutively (HF's ordering). + q = tf.reshape( + tf.repeat_interleave( + tf.reshape(q, new_shape=(S, 1, NK, DK)), repeats=VR, axis=1 + ), + new_shape=(S, NV, DK), + ) + k = tf.reshape( + tf.repeat_interleave( + tf.reshape(k, new_shape=(S, 1, NK, DK)), repeats=VR, axis=1 + ), + new_shape=(S, NV, DK), + ) + q = q * tf.rsqrt( + tf.reduce(tf.square(q), axes=(-1,), keepdim=True, kind="sum") + 1e-6 + ) + k = k * tf.rsqrt( + tf.reduce(tf.square(k), axes=(-1,), keepdim=True, kind="sum") + 1e-6 + ) + v = tf.reshape(mixed[:, 2 * KD :], new_shape=(S, NV, DV)) + g = decay * tf.softplus(alpha + dt_bias) + beta = tf.sigmoid(beta_raw) + state = tf.zeros(Tensor[(NV, DK, DV), "f32"]) + out = tf.zeros(Tensor[(S, NV, DV), "f32"]) + for t in range(S): + state = state * tf.reshape( + tf.exp(g[t : t + 1, :]), new_shape=(NV, 1, 1) + ) + key = tf.reshape(k[t : t + 1, :, :], new_shape=(NV, DK, 1)) + recall = tf.reduce( + state * key, axes=(1,), keepdim=False, kind="sum" + ) + delta = ( + tf.reshape(v[t : t + 1, :, :], new_shape=(NV, DV)) - recall + ) * tf.reshape(beta[t : t + 1, :], new_shape=(NV, 1)) + state = state + key * tf.reshape(delta, new_shape=(NV, 1, DV)) + query = ( + tf.reshape(q[t : t + 1, :, :], new_shape=(NV, DK, 1)) + / GDN_SCALE + ) + read = tf.reshape( + tf.reduce(state * query, axes=(1,), keepdim=False, kind="sum"), + new_shape=(1, NV, DV), + ) + out = tf.insert_slice(out, read, (t, 0, 0)) + gated = tf.rms_norm(out, gamma, eps=EPS) * tf.silu( + tf.reshape(z, new_shape=(S, NV, DV)) + ) + return tf.reshape(gated, new_shape=(S, VD)) + + @module(entry="forward") + class LinearAttention: + qkv = ModuleList(linear(S, CV, H)) + gate = ModuleList(linear(S, VD, H)) + alpha = ModuleList(linear(S, NV, H)) + beta = ModuleList(linear(S, NV, H)) + output = ModuleList(linear(S, H, VD)) + math = ModuleList(GDNMath) + + @func + def forward( + x: Tensor[(S, H), "f32"], + cos: Tensor[(S, R), "f32"], + sin: Tensor[(S, R), "f32"], + mask: Tensor[(S, S), "f32"], + ): + return output(math(qkv(x), gate(x), alpha(x), beta(x))) + + def layer(kind): + mixer_type = FullAttention if kind == "full_attention" else LinearAttention + + @module(entry="forward") + class Layer: + norm_in = ModuleList(Norm) + mixer = ModuleList(mixer_type) + norm_ffn = ModuleList(Norm) + mlp = ModuleList(MLP) + + @func + def add_mixer(x: Tensor[(S, H), "f32"], y: Tensor[(S, H), "f32"]): + return x + y + + @func + def add_mlp(x: Tensor[(S, H), "f32"], y: Tensor[(S, H), "f32"]): + return x + y + + @func + def forward( + x: Tensor[(S, H), "f32"], + cos: Tensor[(S, R), "f32"], + sin: Tensor[(S, R), "f32"], + mask: Tensor[(S, S), "f32"], + ): + x = add_mixer(x, mixer(norm_in(x), cos, sin, mask)) + return add_mlp(x, mlp(norm_ffn(x))) + + return Layer + + kinds = config.get("layer_types", ["full_attention"] * LAYERS) + + definitions = {kind: layer(kind) for kind in set(kinds)} + + @module(target=AmxTarget(), topologies=(Topology("core", 8), Topology("amx", 1))) + class Base: + layers = ModuleList(tuple(definitions[kind] for kind in kinds)) + norm = ModuleList(Norm) + head = ModuleList(linear(1, V, H)) + + @func + def embed(tokens: Tensor[(S,), "i64"], table: ConstTensor[(V, H), "f16"]): + return tf.cast(tf.index_select(table, tokens, dim=0), dtype="f32") + + @func + def last(x: Tensor[(S, H), "f32"]): + return x[S - 1 : S, :] + + # Static slot indices support the heterogeneous GDN/attention stack. + # Definitions stay shared: this emits calls, never cloned layer bodies. + from pathlib import Path + import hashlib + + lines = [ + "@module(entry='prefill',target=AmxTarget(),topologies=(Topology('core',8),Topology('amx',1)))", + "class QwenPrefill:", + " layers=ModuleList(DEFINITIONS)", + " norm=ModuleList(NORM)", + " head=ModuleList(HEAD)", + " embed=BASE.embed", + " last=BASE.last", + " @func", + " def prefill(tokens:Tensor[(S,),'i64'],cos:Tensor[(S,R),'f32'],sin:Tensor[(S,R),'f32'],mask:Tensor[(S,S),'f32'],table:ConstTensor[(V,H),'f16']):", + " hidden=embed(tokens,table)", + ] + lines += [f" hidden=layers[{i}](hidden,cos,sin,mask)" for i in range(LAYERS)] + lines += [" return head(last(norm(hidden)))"] + source = "\n".join(lines) + "\n" + dest = Path(__file__).parent / "artifacts/authored" + dest.mkdir(parents=True, exist_ok=True) + path = dest / ( + hashlib.sha256((str(config) + source).encode()).hexdigest()[:16] + ".py" + ) + path.write_text(source) + namespace = dict( + BASE=Base, + DEFINITIONS=tuple(definitions[k] for k in kinds), + NORM=Norm, + HEAD=Base.head, + ModuleList=ModuleList, + module=module, + func=func, + Tensor=Tensor, + ConstTensor=ConstTensor, + AmxTarget=AmxTarget, + Topology=Topology, + S=S, + R=R, + H=H, + V=V, + ) + exec(compile(source, str(path), "exec"), namespace) + from weights import explicit_weights + + return explicit_weights(namespace["QwenPrefill"]) + + +from config import config + +Qwen3_1_7B = build(config("1.7b"), 128) +Qwen3_5_4B = build(config("4b"), 128) diff --git a/examples/qwen_prefill-amx/run.py b/examples/qwen_prefill-amx/run.py new file mode 100644 index 000000000..0f0e2041e --- /dev/null +++ b/examples/qwen_prefill-amx/run.py @@ -0,0 +1,36 @@ +"""Load the original CPU/AMX runtime and execute one text prefill.""" + +import argparse +from pathlib import Path +import torch +from config import config +from runtime_model import QwenPrefill +from weights import PreparedModelResource + + +def run(model, checkpoint, sequence=128, prepared=None): + torch.set_num_threads(8) + source = Path(prepared) if prepared else Path(__file__).parent / "prepared" / model + resource = PreparedModelResource(checkpoint, source) + with QwenPrefill(config(model), sequence=sequence) as runtime: + runtime.load(resource) + tokens = (torch.arange(sequence, dtype=torch.int32) * 17 + 100) % min( + 30000, runtime.vocab + ) + logits = runtime(tokens) + if not torch.isfinite(logits).all(): + raise RuntimeError("Prefill produced non-finite logits") + print( + f"{model}: logits {tuple(logits.shape)}, next token {int(logits.argmax())}" + ) + return logits + + +if __name__ == "__main__": + p = argparse.ArgumentParser() + p.add_argument("--model", choices=("1.7b", "4b"), required=True) + p.add_argument("--checkpoint", required=True) + p.add_argument("--sequence", type=int, default=128) + p.add_argument("--prepared") + a = p.parse_args() + run(a.model, a.checkpoint, a.sequence, a.prepared) diff --git a/examples/qwen_prefill-amx/runtime_model.py b/examples/qwen_prefill-amx/runtime_model.py new file mode 100644 index 000000000..41c0b572b --- /dev/null +++ b/examples/qwen_prefill-amx/runtime_model.py @@ -0,0 +1,78 @@ +"""Python adapter for the original supplied C++/AMX prefill runtime.""" + +import torch +from tilefoundry.runtime import RuntimeModule +from kernels import ops +from scheduled import make_scheduled + + +class QwenPrefill(RuntimeModule): + """Execute one model through its original native prefill driver. + + The scheduled HIR remains available as ``module`` for inspection. Execution + uses the supplied native graph, including its existing host operations. + """ + + def __init__(self, config, *, sequence=128, threads=8, capacity=2048): + super().__init__("QwenPrefill", entry="prefill") + if sequence < 128 or sequence > 8192 or sequence % 32: + raise ValueError("sequence must be a multiple of 32 between 128 and 8192") + self._semantic = make_scheduled(config, sequence) + self.threads = threads + self.capacity = max(capacity, sequence) + self.sequence = sequence + self.vocab = config["vocab_size"] + self.ffi = None + self.loaded = False + + @property + def module(self): + return self._semantic + + def load(self, resource, *, placement=None): + if self.loaded: + raise RuntimeError("Close the loaded model before loading another") + self.ffi = ops() + resource.prepare() + try: + vocab = int( + self.ffi.load_model( + str(resource.checkpoint), self.threads, self.capacity, 7 + ) + ) + self.loaded = True + if vocab != self.vocab: + raise ValueError( + "The checkpoint does not match the selected model vocabulary" + ) + for name, n, k, path in resource.entries(): + self.ffi.load_weight(name, str(path), n, k) + except BaseException: + self.close() + raise + + def forward(self, tokens): + if not self.loaded: + raise RuntimeError("Load the model resource first") + if ( + tokens.device.type != "cpu" + or tokens.dtype != torch.int32 + or tokens.ndim != 1 + ): + raise ValueError("Expected a CPU int32 token vector") + if tokens.numel() != self.sequence: + raise ValueError("Token length must equal the module's sequence length") + logits = torch.empty(self.vocab, dtype=torch.float32) + self.ffi.prefill(tokens.contiguous(), logits) + return logits.reshape(1, self.vocab) + + def close(self): + if self.loaded: + self.ffi.close_model() + self.loaded = False + + def __enter__(self): + return self + + def __exit__(self, *exc): + self.close() diff --git a/examples/qwen_prefill-amx/scheduled.py b/examples/qwen_prefill-amx/scheduled.py new file mode 100644 index 000000000..b04a528b5 --- /dev/null +++ b/examples/qwen_prefill-amx/scheduled.py @@ -0,0 +1,225 @@ +"""Qwen HIR with explicit AMX projection layouts and worker schedules.""" + +from dataclasses import replace +from tilefoundry import module, func +from tilefoundry.dsl import Tensor, ConstTensor, Mesh, T, Topology, tf +from tilefoundry.ir.types import Layout, DType +from tilefoundry.ir.tir.amx.instructions import Z_LAYOUT +from model import build +from tilefoundry.ir.core import Call, Var +from tilefoundry.ir.visitor import collect_exprs, ExprCloner, BindingSubstitutionCloner +from tilefoundry.ir.hir.schedule import ScheduleOp +from tilefoundry.ir.hir.mesh_region import MeshRegion +from tilefoundry.ir.types import Mesh as MeshValue +from tilefoundry.ir.tir.cpu.tensor import instruction_for, INSTRUCTIONS + +BM = 32 +WORKERS = 8 + + +def scheduled_projection(M, N, K): + if K <= 0 or N <= 0 or K % 32 or N % 32: + raise ValueError( + "The native AMX runtime requires positive K and N that are multiples of 32" + ) + B = 64 + if M != 1 and (M <= 0 or M % 32): + raise ValueError( + "This schedule requires a multiple of 32 prompt tokens, or M=1 for the last-token head" + ) + if M < 128 or N % 64: + + @module(entry="forward", topologies=(Topology("core", 8), Topology("amx", 1))) + class Head: + @func + def forward(x: Tensor[(M, K), "f32"], w: ConstTensor[(K, N), "f16"]): + return tf.matmul(tf.cast(x, dtype="f16"), w, out_dtype="f32") + + @forward.converter("w") + def _(raw: ConstTensor[(N, K), "f16"]): + return tf.transpose(raw, perm=(1, 0)) + + return cpu_schedule(Head) + + @module(entry="forward", topologies=(Topology("core", 8), Topology("amx", 1))) + class Linear: + @func + def forward(x: Tensor[(M, K), "f32"], w: ConstTensor[(N // B, K, B), "f16"]): + with Mesh(("core",), layout=(1,), names=("packer",)) as _packer: + # One CPU pack, outside all K issues. The physical 32-row + # panel makes each AMX Y load a contiguous 64-byte vector. + half = tf.schedule( + (x, tf.empty(Tensor[(M, K), "f16", None, "gmem"])), + op=T.cpu_cast(dtype=DType.f16), + ) + view = tf.transpose( + tf.reshape(half, new_shape=(M // BM, BM, K)), perm=(0, 2, 1) + ) + packed = tf.schedule( + ( + view, + tf.empty( + Tensor[ + (M // BM, K, BM), + "f16", + Layout((M // BM, K, BM), (K * BM, BM, 1)), + "gmem", + ] + ), + ), + op=T.copy(), + ) + out = tf.zeros(Tensor[(M, N), "f32"]) + with Mesh(("core",), layout=(WORKERS,), names=("worker",)) as cores: + for panel in range(cores.worker, N // B, WORKERS): + for half_n in range(B // 32): + for mb in range(M // BM): + with Mesh(("amx",), layout=(1,), names=("unit",)) as _amx: + acc = tf.zeros( + Tensor[(BM, 32), "f32", Z_LAYOUT, "rmem"] + ) + for k in range(K): + a = tf.schedule( + ( + tf.reshape( + packed[mb : mb + 1, k : k + 1, :], + new_shape=(BM, 1), + ), + tf.empty( + Tensor[(BM, 1), "f16", None, "rmem"] + ), + ), + op=T.amx_load_y(), + ) + b = tf.schedule( + ( + tf.reshape( + w[ + panel : panel + 1, + k : k + 1, + half_n * 32 : half_n * 32 + 32, + ], + new_shape=(1, 32), + ), + tf.empty( + Tensor[(1, 32), "f16", None, "rmem"] + ), + ), + op=T.amx_load_x(), + ) + acc = tf.schedule( + (acc, a, b), + op=T.amx_fma16(), + repeat=(1, 1, 1), + order=(0, 1, 2), + ) + tile = tf.schedule( + ( + acc, + tf.empty(Tensor[(BM, 32), "f32", None, "gmem"]), + ), + op=T.amx_store(), + ) + out = tf.insert_slice( + out, tile, (mb * BM, panel * B + half_n * 32) + ) + return out + + @forward.converter("w") + def _(raw: ConstTensor[(N, K), "f16"]): + return tf.transpose( + tf.reshape(raw, new_shape=(N // B, B, K)), perm=(0, 2, 1) + ) + + return Linear + + +class CPUChoices(ExprCloner): + """Record explicit software-kernel choices without wrapping whole functions.""" + + def visit_Call(self, call, ctx=None): + args = tuple(self.visit(arg, ctx) for arg in call.args) + if (type(call.target), len(args)) in INSTRUCTIONS and not ( + call.type.shape == () and call.type.dtype in (DType.i64, DType.i32) + ): + op = instruction_for(call.target, len(args), call.type.dtype) + from tilefoundry.ir.hir.tensor.empty import Empty + from tilefoundry.ir.types import StorageKind + from tilefoundry.ir.types.stride import compact_row_major + + ty = replace( + call.type, + storage=StorageKind.GMEM, + layout=Layout( + call.type.shape, tuple(compact_row_major(call.type.shape)) + ), + ) + dst = Call(Empty(type=ty), (), type=ty) + return Call(ScheduleOp(op=op), (*args, dst), type=ty) + return replace(call, args=args) + + +def cpu_schedule(module, memo=None): + from tilefoundry.dsl import ModuleList + from tilefoundry.ir.core import Op + + memo = {} if memo is None else memo + if id(module) in memo: + return memo[id(module)] + functions = [] + for fn in module.functions: + if fn.body is None: + raise ValueError("No dispatch prototypes") + if isinstance(fn.body, MeshRegion) or any( + isinstance(v, Call) and not isinstance(v.target, Op) + for v in collect_exprs(fn.body) + ): + functions.append(fn) + continue + body = CPUChoices().visit(fn.body) + mesh = MeshValue((Topology("core", 8),), Layout((1,), (1,)), ("worker",)) + params = tuple(Var(p.name, type=p.type) for p in fn.params) + body = BindingSubstitutionCloner().visit( + body, {id(p): v for p, v in zip(fn.params, params)} + ) + body = MeshRegion( + mesh, body, params=params, args=fn.params, type=fn.return_type + ) + functions.append(replace(fn, body=body)) + mapping = {id(old): new for old, new in zip(module.functions, functions)} + + class Retarget(ExprCloner): + def visit_Call(self, call, ctx=None): + result = replace( + call, + args=tuple(self.visit(arg, ctx) for arg in call.args), + target=self.visit(call.target, ctx) + if isinstance(call.target, Call) + else call.target, + ) + return replace(result, target=mapping.get(id(result.target), result.target)) + + for fn in functions: + fn.body = Retarget().visit(fn.body) + children = { + key: ModuleList(tuple(cpu_schedule(child, memo) for child in slot.elements)) + if slot.indexed + else ModuleList(cpu_schedule(slot.elements[0], memo)) + for key, slot in module.children.items() + } + result = replace(module, functions=tuple(functions), children=children) + memo[id(module)] = result + return result + + +def make_scheduled(config, sequence=128): + return replace( + cpu_schedule(build(config, sequence, linear=scheduled_projection)), + topologies=(Topology("core", 8), Topology("amx", 1)), + ) + + +from config import config + +Qwen3_1_7B = make_scheduled(config("1.7b"), 128) +Qwen3_5_4B = make_scheduled(config("4b"), 128) diff --git a/examples/qwen_prefill-amx/weights.py b/examples/qwen_prefill-amx/weights.py new file mode 100644 index 000000000..b7320f202 --- /dev/null +++ b/examples/qwen_prefill-amx/weights.py @@ -0,0 +1,317 @@ +"""Canonical GGUF resource aliases for the whole Prefill HIR tree.""" + +from pathlib import Path +import numpy as np +import torch +from gguf import GGUFReader, GGMLQuantizationType + + +def explicit_weights(root): + """Keep model weights on the root and pass them through shared definitions.""" + from dataclasses import replace + from tilefoundry.ir.core import Call, Var + from tilefoundry.ir.core.module import ModuleList + from tilefoundry.ir.hir.function import Function + from tilefoundry.ir.visitor import BindingSubstitutionCloner, collect_exprs + from tilefoundry.visitor_registry.contexts import ValuationContext + + cache = {} + + def lift(mod): + if id(mod) in cache: + return cache[id(mod)] + children = {} + specs = {name: (ty, None, ()) for name, ty in mod.weights.items()} + for fn in mod.functions: + for name, converter in fn.converters: + specs[name] = (specs[name][0], converter, ()) + for key, slot in mod.children.items(): + definitions = [] + for i, child in enumerate(slot.elements): + definition, child_specs, _ = lift(child) + definitions.append(definition) + prefix = f"{key}__{i}__" if slot.indexed else f"{key}__" + for name, (ty, converter, path) in child_specs.items(): + specs[prefix + name] = (ty, converter, (prefix, *path)) + children[key] = ModuleList( + tuple(definitions) if slot.indexed else definitions[0] + ) + needs, functions = {}, {} + + def callee(call): + frame = ValuationContext(module=mod)._callee_context(call) + if frame is None: + return None + if isinstance(call.target, Function): + return call.target, "" + slot = ( + call.target + if isinstance(call.target, ModuleList) + else call.target.args[0] + ) + index = ( + 0 if isinstance(call.target, ModuleList) else call.target.args[1].value + ) + prefix = f"{slot.key}__{index}__" if slot.indexed else f"{slot.key}__" + return frame.function, prefix + + def needed(fn): + if id(fn) not in needs: + names = [p.name for p in fn.params if p.is_const] + for call in collect_exprs(fn.body): + if not isinstance(call, Call) or (item := callee(call)) is None: + continue + target, prefix = item + target_needs = ( + needed(target) + if not prefix + else lift( + ValuationContext(module=mod, function=fn)._callee_context(call).module + )[2][id(target)] + ) + names.extend(prefix + name for name in target_needs) + needs[id(fn)] = tuple(dict.fromkeys(names)) + return needs[id(fn)] + + def rewrite(fn): + if id(fn) in functions: + return functions[id(fn)] + inputs = tuple( + Var(p.name, type=p.type) for p in fn.params if not p.is_const + ) + weights = {name: Var(name, type=specs[name][0]) for name in needed(fn)} + input_iter = iter(inputs) + bindings = { + id(p): weights[p.name] if p.is_const else next(input_iter) + for p in fn.params + } + + class Rewrite(BindingSubstitutionCloner): + def visit_ModuleList(self, slot, ctx=None): + return children[slot.key] + + def visit_Call(self, call, ctx=None): + args = tuple(self.visit(arg, ctx) for arg in call.args) + item = callee(call) + if item is None: + return replace(call, args=args) + target, prefix = item + if prefix: + child = ValuationContext(module=mod, function=fn)._callee_context(call).module + child_needs = lift(child)[2][id(target)] + target = self.visit(call.target, ctx) + else: + args = tuple( + arg + for param, arg in zip(target.params, args) + if not param.is_const + ) + child_needs = needed(target) + target = rewrite(target) + return replace( + call, + target=target, + args=(*args, *(weights[prefix + name] for name in child_needs)), + ) + + result = Function.build( + name=fn.name, + params=(*inputs, *weights.values()), + body=Rewrite().visit(fn.body, bindings), + return_type=fn.return_type, + ) + functions[id(fn)] = result + return result + + result = replace( + mod, functions=tuple(rewrite(fn) for fn in mod.functions), children=children + ) + cache[id(mod)] = result, specs, needs + return cache[id(mod)] + + result, specs, needs = lift(root) + entry = result.entry_function() + names = needs[id(root.entry_function())] + weights = set(names) + params = tuple( + Var(p.name, type=p.type, is_const=p.name in weights) for p in entry.params + ) + converters = [] + for name in names: + _, converter, path = specs[name] + if converter is None: + continue + prefix = "".join(path) + raw = tuple( + Var(prefix + p.name, type=p.type, is_const=True) for p in converter.params + ) + converted = Function.build( + name=prefix + converter.name, + params=raw, + body=BindingSubstitutionCloner().visit( + converter.body, {id(p): v for p, v in zip(converter.params, raw)} + ), + return_type=converter.return_type, + ) + converters.append((name, converted)) + entry = Function.build( + name=entry.name, + params=params, + body=BindingSubstitutionCloner().visit( + entry.body, {id(p): v for p, v in zip(entry.params, params)} + ), + return_type=entry.return_type, + converters=tuple(converters), + ) + return replace( + result, + functions=tuple( + entry if fn.name == entry.name else fn for fn in result.functions + ), + ) + + +def aliases(config): + out = { + "table": "token_embd.weight", + "norm.gamma": "output_norm.weight", + "head.raw": "token_embd.weight", + } + for i, kind in enumerate( + config.get("layer_types", ["full_attention"] * config["num_hidden_layers"]) + ): + p = f"layers.{i}" + b = f"blk.{i}" + out[p + ".norm_in.gamma"] = b + ".attn_norm.weight" + out[p + ".norm_ffn.gamma"] = b + ( + ".post_attention_norm.weight" + if "linear_num_key_heads" in config + else ".ffn_norm.weight" + ) + for name in ("gate", "up", "down"): + out[p + ".mlp." + name + ".raw"] = b + ".ffn_" + name + ".weight" + if kind == "full_attention": + for name, tensor in ( + ("query", "attn_q"), + ("key", "attn_k"), + ("value", "attn_v"), + ("output", "attn_output"), + ): + out[p + ".mixer." + name + ".raw"] = b + "." + tensor + ".weight" + out[p + ".mixer.query_norm.gamma_q"] = b + ".attn_q_norm.weight" + out[p + ".mixer.key_norm.gamma_k"] = b + ".attn_k_norm.weight" + else: + for name, tensor in ( + ("qkv", "attn_qkv"), + ("gate", "attn_gate"), + ("alpha", "ssm_alpha"), + ("beta", "ssm_beta"), + ("output", "ssm_out"), + ): + out[p + ".mixer." + name + ".raw"] = b + "." + tensor + ".weight" + for name, tensor in ( + ("conv", "ssm_conv1d.weight"), + ("decay", "ssm_a"), + ("dt_bias", "ssm_dt.bias"), + ("gamma", "ssm_norm.weight"), + ): + out[p + ".mixer.math." + name] = b + "." + tensor + return {key.replace(".", "__"): value for key, value in out.items()} + + +class GGUFResource: + def __init__(self, path, alias, *, prefix="", shared=None): + self.path = Path(path) + self.alias = alias + self.prefix = prefix + self.shared = ( + shared + if shared is not None + else {"reader": GGUFReader(self.path), "cache": {}} + ) + self.tensors = {t.name: t for t in self.shared["reader"].tensors} + + def load(self, name): + raw = self.alias[self.prefix + name] + cache = self.shared["cache"] + if raw in cache: + return cache[raw] + t = self.tensors[raw] + if t.tensor_type == GGMLQuantizationType.Q8_0: + k, n = map(int, t.shape) + blocks = np.asarray(t.data).reshape(n, k // 32, 34) + codes = blocks[:, :, 2:].copy().view(np.int8).astype(np.float32) + scale = ( + blocks[:, :, :2] + .copy() + .view(np.float16) + .reshape(n, k // 32, 1) + .astype(np.float32) + ) + value = torch.from_numpy((codes * scale).reshape(n, k).astype(np.float16)) + elif t.tensor_type == GGMLQuantizationType.F32: + value = torch.from_numpy( + np.asarray(t.data).copy().reshape(tuple(map(int, t.shape[::-1]))) + ).float() + else: + raise ValueError(f"Unsupported source dtype {t.name}: {t.tensor_type}") + cache.clear() + cache[raw] = value + return value + + def load_group(self, name): + return None + + def subtree(self, seg): + return GGUFResource( + self.path, self.alias, prefix=self.prefix + seg + ".", shared=self.shared + ) + + +class PreparedModelResource: + """GGUF checkpoint with the original N/64,K,64 prepared AMX weights.""" + + def __init__(self, checkpoint, directory): + self.checkpoint = Path(checkpoint).expanduser().resolve() + self.directory = Path(directory).expanduser().resolve() + self.reader = GGUFReader(str(self.checkpoint)) + + def entries(self): + for tensor in self.reader.tensors: + if ( + tensor.tensor_type != GGMLQuantizationType.Q8_0 + or len(tensor.shape) != 2 + ): + continue + k, n = map(int, tensor.shape) + if n % 64 or k % 32: + continue + if not tensor.name or any( + c + not in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_." + for c in tensor.name + ): + raise ValueError("Invalid GGUF tensor name") + yield tensor.name, n, k, self.directory / (tensor.name + ".bin") + + def prepare(self): + from scheduled import scheduled_projection + from tilefoundry.runtime import DictResource, SafetensorsResource + + self.directory.mkdir(parents=True, exist_ok=True) + raw = GGUFResource( + self.checkpoint, {t.name: t.name for t in self.reader.tensors} + ) + for name, n, k, path in self.entries(): + if path.exists(): + if path.stat().st_size != n * k * 2: + raise ValueError(f"Invalid prepared weight size: {path}") + continue + stage = self.directory / name + scheduled_projection(128, n, k).prepare( + DictResource({"raw": raw.load(name)}), str(stage), device="cpu" + ) + packed = SafetensorsResource(str(stage), device="cpu").load("w") + temporary = path.with_suffix(".bin.tmp") + packed.contiguous().numpy().tofile(temporary) + temporary.replace(path) diff --git a/include/tilefoundry/runtime/amx/instructions.h b/include/tilefoundry/runtime/amx/instructions.h new file mode 100644 index 000000000..8a2b2fdfa --- /dev/null +++ b/include/tilefoundry/runtime/amx/instructions.h @@ -0,0 +1,39 @@ +#pragma once +#include + +namespace tilefoundry::amx { +/** Enable the current thread's Apple AMX state. */ +inline void start() { + asm volatile("nop\nnop\nnop\n.word 0x201220" ::: "memory"); +} +/** Release the current thread's Apple AMX state. */ +inline void stop() { + asm volatile("nop\nnop\nnop\n.word 0x201221" ::: "memory"); +} +/** Load one 64-byte vector into X register zero. */ +inline void load_x(const void *ptr) { + auto operand = reinterpret_cast(ptr) & ((1ULL << 57) - 1); + asm volatile(".word (0x201000 + 0%0 - ((0%0 >> 4) * 6))" ::"r"(operand) + : "memory"); +} +/** Load one 64-byte vector into Y register zero. */ +inline void load_y(const void *ptr) { + auto operand = reinterpret_cast(ptr) & ((1ULL << 57) - 1); + asm volatile(".word (0x201020 + 0%0 - ((0%0 >> 4) * 6))" ::"r"(operand) + : "memory"); +} +/** Accumulate the FP16 outer product into FP32 Z, optionally clearing Z first. + */ +inline void fma16(bool clear = false) { + uint64_t operand = (1ULL << 62) | (uint64_t(clear) << 27); + asm volatile(".word (0x2011e0 + 0%0 - ((0%0 >> 4) * 6))" ::"r"(operand) + : "memory"); +} +/** Store one physical 64-byte row of Z. */ +inline void store_z(void *ptr, uint64_t row) { + auto operand = + (reinterpret_cast(ptr) & ((1ULL << 57) - 1)) | (row << 56); + asm volatile(".word (0x2010a0 + 0%0 - ((0%0 >> 4) * 6))" ::"r"(operand) + : "memory"); +} +} diff --git a/pyproject.toml b/pyproject.toml index 85afe9c35..3f90b0bf4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -65,6 +65,7 @@ where = ["src"] "tilefoundry.target.amx.hardware" = ["*.toml"] [tool.setuptools.data-files] +"include/tilefoundry/runtime/amx" = ["include/tilefoundry/runtime/amx/*.h"] "share/tilefoundry/spec" = ["docs/spec/*.md"] # Tutorial pages have distinct filenames, so one destination is sufficient. "share/tilefoundry/tutorial" = ["docs/tutorial/*.md"] diff --git a/src/tilefoundry/inspection/pattern_printer.py b/src/tilefoundry/inspection/pattern_printer.py index 319907816..07686df21 100644 --- a/src/tilefoundry/inspection/pattern_printer.py +++ b/src/tilefoundry/inspection/pattern_printer.py @@ -429,7 +429,12 @@ def rules_ComposedLayoutPattern(self, pattern, name) -> tuple[str, ...]: return self._rules_of((pattern.inner, pattern.offset, pattern.outer)) def visit_MeshPattern(self, pattern, name) -> str: - return f"Mesh({pattern.topologies!r}, {self.written(pattern.layout)})" + topologies = ( + repr(pattern.topologies) + if isinstance(pattern.topologies, tuple) + else self.written(pattern.topologies) + ) + return f"Mesh({topologies}, {self.written(pattern.layout)})" def rules_MeshPattern(self, pattern, name) -> tuple[str, ...]: return self.rules(pattern.layout) diff --git a/src/tilefoundry/ir/hir/schedule.py b/src/tilefoundry/ir/hir/schedule.py index ee7e925af..729dd2f21 100644 --- a/src/tilefoundry/ir/hir/schedule.py +++ b/src/tilefoundry/ir/hir/schedule.py @@ -47,6 +47,7 @@ operand_relations, project, ) +from tilefoundry.visitor_registry.registries import cost_evaluator_registry @dataclass(frozen=True) @@ -170,6 +171,9 @@ def _schedule_access_relation(call: Call, ctx) -> tuple[AccessRelation, ...]: @register_cost_evaluator(ScheduleOp) def _schedule_cost(call: Call, ctx) -> Cost: + evaluator = cost_evaluator_registry.lookup(type(call.target.op)) + if evaluator is not None: + return evaluator(Call(call.target.op, call.args, type=call.type), ctx) stated = relations_of(call, ctx) local = projected(stated, call, ctx) types = (*(ctx.local_type_of(arg) for arg in call.args), ctx.local_output_type(call)) @@ -216,7 +220,7 @@ def match_instruction(op_type, attributes, operand_types, mesh, *, frame, repeat raise ScheduleRefusal(f"{op_type.__name__} schedule operands must be tensors", op=op) atom = getattr(op, "atom", None) if not isinstance(atom, type): - admit_fields(op, operand_types, mesh, PatternMatcher(dict(getattr(atom, "bindings", {})))) + admit_fields(op, operand_types, mesh, PatternMatcher(dict(getattr(atom, "bindings", {}))), target=None if frame is None else frame.target) tiles = project(op, operand_types, mesh, repeat=repeat, order=order) bindings, _tiles = check_call( op, operand_types, mesh, frame=frame, repeat=repeat, order=order, projection=tiles @@ -241,7 +245,7 @@ def match_instruction(op_type, attributes, operand_types, mesh, *, frame, repeat for seed in partial: matcher = PatternMatcher(seed) try: - admit_fields(op, operand_types, mesh, matcher) + admit_fields(op, operand_types, mesh, matcher, target=None if frame is None else frame.target) except ScheduleRefusal as error: refusal = refusal or error continue @@ -316,6 +320,7 @@ def _infer_schedule(call: Call, ctx) -> TensorType: tuple(ctx.type_of(arg) for arg in call.args), ctx.current_mesh, PatternMatcher(dict(getattr(getattr(op, "atom", None), "bindings", {}))), + target=ctx.target, ) tiles = project( op, diff --git a/src/tilefoundry/ir/pattern/match.py b/src/tilefoundry/ir/pattern/match.py index 90c14460e..d8de1c956 100644 --- a/src/tilefoundry/ir/pattern/match.py +++ b/src/tilefoundry/ir/pattern/match.py @@ -95,7 +95,16 @@ def match(self, pattern, subject) -> bool: if isinstance(subject, ShardLayout) and isinstance( pattern, _shard_layout_pattern_type() ): - subject = participant_layout(subject, pattern.mesh.topologies) + topologies = pattern.mesh.topologies + if not isinstance(topologies, tuple): + topologies = self.bindings.get( + topologies.name, + tuple( + getattr(level, "name", level) + for level in subject.mesh.topologies[-1:] + ), + ) + subject = participant_layout(subject, topologies) except ValueError: found = self._fail(pattern, subject) else: @@ -405,9 +414,18 @@ def visit_ComposedLayoutPattern(self, pattern, subject) -> bool: def visit_MeshPattern(self, pattern, subject) -> bool: if not isinstance(subject, Mesh): return self._fail(pattern, subject) + levels = separate(subject) + if not isinstance(pattern.topologies, tuple): + if not levels: + return self._fail(pattern, subject) + selected = levels[-1] + names = tuple(getattr(level, "name", level) for level in selected.topologies) + return self.match(pattern.topologies, names) and self.match( + pattern.layout, selected.layout + ) picked = tuple( level - for level in separate(subject) + for level in levels if getattr(level.topologies[0], "name", level.topologies[0]) in pattern.topologies ) found = tuple(getattr(level.topologies[0], "name", level.topologies[0]) for level in picked) @@ -443,7 +461,7 @@ def visit_ShardLayoutPattern(self, pattern, subject) -> bool: subject_names = tuple( getattr(topology, "name", topology) for topology in subject.mesh.topologies ) - if subject_names != pattern.mesh.topologies: + if not self.match(pattern.mesh.topologies, subject_names): return self._fail(pattern, subject) return all( self.match(place, value) diff --git a/src/tilefoundry/ir/pattern/pattern.py b/src/tilefoundry/ir/pattern/pattern.py index 686f52856..acfefa2a4 100644 --- a/src/tilefoundry/ir/pattern/pattern.py +++ b/src/tilefoundry/ir/pattern/pattern.py @@ -331,13 +331,13 @@ def fixed(self): @dataclass(frozen=True) class MeshPattern(Pattern): - """Match named mesh levels after separating and recomposing them.""" + """Match named levels, or bind the innermost level with a wildcard.""" - topologies: tuple[str, ...] + topologies: tuple[str, ...] | WildcardPattern layout: object def __post_init__(self): - if ( + if not isinstance(self.topologies, WildcardPattern) and ( not self.topologies or any(not isinstance(name, str) or not name for name in self.topologies) or len(set(self.topologies)) != len(self.topologies) diff --git a/src/tilefoundry/ir/pattern/utils.py b/src/tilefoundry/ir/pattern/utils.py index bac916f3e..bff557af9 100644 --- a/src/tilefoundry/ir/pattern/utils.py +++ b/src/tilefoundry/ir/pattern/utils.py @@ -114,8 +114,13 @@ def thread_execution_mesh() -> MeshPattern: return MeshPattern(("thread",), _THREAD_EXECUTION_LAYOUT) -def declared_execution_mesh(op_type) -> MeshPattern | OrPattern: - """Return one op's reserved execution-mesh declaration or reject it.""" +def memory_execution_mesh() -> MeshPattern: + """Bind a contiguous run at the current innermost topology level.""" + return MeshPattern(WildcardPattern(), _THREAD_EXECUTION_LAYOUT) + + +def declared_execution_mesh(op_type, target=None) -> MeshPattern | OrPattern: + """Resolve one execution declaration, specializing memory levels from Target Facts.""" declaration = op_type if isinstance(op_type, type) else type(op_type) factory = getattr(op_type, "execution_mesh_pattern", None) if callable(factory): @@ -129,6 +134,18 @@ def declared_execution_mesh(op_type) -> MeshPattern | OrPattern: raise ValueError( f"{name} execution_mesh must be a constant, ParamDef, or operand mesh reference" ) + if target is not None and any(isinstance(choice.topologies, WildcardPattern) for choice in choices): + from tilefoundry.target.facts import MemoryExecutionFacts # noqa: PLC0415 + + levels = target.get_facts(MemoryExecutionFacts).topologies + expanded = [] + for choice in choices: + if isinstance(choice.topologies, WildcardPattern): + expanded.extend(replace(choice, topologies=(level,)) for level in levels) + else: + expanded.append(choice) + choices = tuple(expanded) + return choices[0] if len(choices) == 1 else OrPattern(*choices) return declared diff --git a/src/tilefoundry/ir/tir/amx/__init__.py b/src/tilefoundry/ir/tir/amx/__init__.py new file mode 100644 index 000000000..997a62372 --- /dev/null +++ b/src/tilefoundry/ir/tir/amx/__init__.py @@ -0,0 +1 @@ +"""Apple AMX instruction contracts.""" diff --git a/src/tilefoundry/ir/tir/amx/instructions.py b/src/tilefoundry/ir/tir/amx/instructions.py new file mode 100644 index 000000000..4215fbfcc --- /dev/null +++ b/src/tilefoundry/ir/tir/amx/instructions.py @@ -0,0 +1,142 @@ +"""Apple AMX FP16 outer-product and memory-transfer contracts. + +X holds 32 weight columns, Y holds 32 activation rows, and Z accumulates +32x32 FP32 results. The native module implements the same explicit contract. +""" + +from tilefoundry.evaluator.registry import register_schedule_eval +from tilefoundry.evaluator.value import TensorValue +from tilefoundry.ir.core import Op, OpCapability +from tilefoundry.ir.core.param_def import MemoryEffect, ParamDef +from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import ( + ComposedLayoutPattern, + LayoutPattern, + MeshPattern, + TensorPattern, + WildcardPattern, +) +from tilefoundry.ir.types import ( + DType, + Layout, + Mesh, + StorageKind, + Topology, + UnitType, +) +from tilefoundry.visitor_registry import register_typeinfer +from tilefoundry.visitor_registry.access_relation import ( + identity_relations, + matmul_relations, + register_access_relation, +) + +Z_LAYOUT = Layout((32, (16, 2)), (32, (1, 16))) +AMX = Mesh((Topology("amx", 1),), Layout((1,), (1,)), ("unit",)) + + +def mesh_pattern(name): + return MeshPattern( + (name,), + ComposedLayoutPattern( + inner=None, + offset=WildcardPattern(name + "_offset"), + outer=LayoutPattern(((1,),), ((1,),)), + ), + ) + + +AMX_PATTERN = mesh_pattern("amx") +R = MemoryEffect.READ +W = MemoryEffect.WRITE + + +@register_op(dialect="T", category="amx", name="amx_fma16") +class Fma16(Op): + capability = OpCapability("amx_outer_product") + resource = "amx_engine" + execution_mesh = AMX_PATTERN + acc = ParamDef( + kind="input", + effect=R | W, + pattern=TensorPattern( + shape=(32, 32), + dtype=DType.f32, + storage=StorageKind.RMEM, + layout=LayoutPattern.from_layout(Z_LAYOUT), + ), + ) + lhs = ParamDef( + kind="input", + effect=R, + pattern=TensorPattern(shape=(32, 1), dtype=DType.f16, storage=StorageKind.RMEM), + ) + rhs = ParamDef( + kind="input", + effect=R, + pattern=TensorPattern(shape=(1, 32), dtype=DType.f16, storage=StorageKind.RMEM), + ) + + +@register_access_relation(Fma16) +def relation(call, ctx): + _, a, b = (ctx.type_of(v) for v in call.args) + *operands, out = matmul_relations(a.shape, b.shape, (-2, -1, -1, -2)) + return (out, *operands, out) + + +@register_schedule_eval(Fma16) +def evaluate(ctx): + c, a, b = (v.data for v in ctx.args) + return TensorValue(data=c + a.float() @ b.float(), type=ctx.result_type) + + +@register_typeinfer(Fma16) +def infer(call, ctx): + return UnitType() + + +def transfer(name, shape, dtype, storage, execution): + cls = register_op(dialect="T", category="amx", name=name)( + type( + name, + (Op,), + dict( + __module__=__name__, + capability=OpCapability("amx_outer_product"), + resource="amx_engine", + execution_mesh=execution, + src=ParamDef( + kind="input", + effect=R, + pattern=TensorPattern( + shape=shape, + dtype=dtype, + storage=StorageKind.RMEM + if storage is StorageKind.GMEM + else StorageKind.GMEM, + layout=LayoutPattern.from_layout(Z_LAYOUT) + if storage is StorageKind.GMEM + else None, + ), + ), + dst=ParamDef( + kind="input", + effect=W, + pattern=TensorPattern(shape=shape, dtype=dtype, storage=storage), + ), + ), + ) + ) + register_typeinfer(cls)(lambda call, ctx: UnitType()) + register_access_relation(cls)(identity_relations) + register_schedule_eval(cls)( + lambda ctx: TensorValue(data=ctx.args[0].data.clone(), type=ctx.result_type) + ) + globals()[name] = cls + return cls + + +LoadX = transfer("amx_load_x", (1, 32), DType.f16, StorageKind.RMEM, AMX_PATTERN) +LoadY = transfer("amx_load_y", (32, 1), DType.f16, StorageKind.RMEM, AMX_PATTERN) +StoreZ = transfer("amx_store", (32, 32), DType.f32, StorageKind.GMEM, AMX_PATTERN) diff --git a/src/tilefoundry/ir/tir/cpu/__init__.py b/src/tilefoundry/ir/tir/cpu/__init__.py new file mode 100644 index 000000000..50f0db158 --- /dev/null +++ b/src/tilefoundry/ir/tir/cpu/__init__.py @@ -0,0 +1 @@ +"""CPU tensor-kernel contracts.""" diff --git a/src/tilefoundry/ir/tir/cpu/tensor.py b/src/tilefoundry/ir/tir/cpu/tensor.py new file mode 100644 index 000000000..5124eab35 --- /dev/null +++ b/src/tilefoundry/ir/tir/cpu/tensor.py @@ -0,0 +1,203 @@ +"""Explicit CPU tensor-kernel contracts for hybrid CPU/AMX schedules.""" + +from dataclasses import replace +from enum import Enum + +import isl + +from tilefoundry.evaluator.registry import eval_registry, register_schedule_eval +from tilefoundry.ir.core import Call, Op, OpCapability +from tilefoundry.ir.core.param_def import MemoryEffect, ParamDef, collect_param_defs +from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.hir.math.binary import Binary +from tilefoundry.ir.hir.math.softplus import Softplus +from tilefoundry.ir.hir.math.unary import Unary +from tilefoundry.ir.hir.nn.matmul import MatMul +from tilefoundry.ir.hir.nn.rms_norm import RMSNorm +from tilefoundry.ir.hir.nn.sigmoid import Sigmoid +from tilefoundry.ir.hir.nn.silu import Silu +from tilefoundry.ir.hir.nn.softmax import SoftMax +from tilefoundry.ir.hir.tensor.arange import Arange +from tilefoundry.ir.hir.tensor.cast import Cast +from tilefoundry.ir.hir.tensor.concat import Concat +from tilefoundry.ir.hir.tensor.full_like import FullLike +from tilefoundry.ir.hir.tensor.index_select import IndexSelect +from tilefoundry.ir.hir.tensor.reduce import Reduce +from tilefoundry.ir.hir.tensor.repeat_interleave import RepeatInterleave +from tilefoundry.ir.hir.tensor.reshape import Reshape +from tilefoundry.ir.hir.tensor.transpose import Transpose +from tilefoundry.ir.pattern import ( + ComposedLayoutPattern, + LayoutPattern, + MeshPattern, + TensorPattern, + WildcardPattern, +) +from tilefoundry.ir.types import DType, StorageKind, UnitType +from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt +from tilefoundry.visitor_registry.access_relation import ( + access_relation_registry, + register_access_relation, +) +from tilefoundry.visitor_registry.contexts import TrafficBytes, TypeInferContext +from tilefoundry.visitor_registry.registries import ( + cost_evaluator_registry, + register_cost_evaluator, + typeinfer_registry, +) + +CORE = MeshPattern( + ("core",), + ComposedLayoutPattern( + inner=None, offset=WildcardPattern("core_offset"), outer=LayoutPattern(((1,),), ((1,),)) + ), +) +INSTRUCTIONS = {} + + +def _semantic_type(ty): + """Treat scalar-by-value inputs as uniform HIR operands.""" + return replace(ty, storage=StorageKind.UMAT) if getattr(ty, "shape", None) == () else ty + + +def _register(hir, arity): + original_attrs = tuple(p for p in collect_param_defs(hir) if p.kind == "attribute") + attributes = {p.name: replace(p, effect=None) for p in original_attrs} + if "dtype" not in attributes: + attributes["dtype"] = ParamDef(kind="attribute", annotation=DType, default=DType.f32) + + def constructor(self, **bound): + for name, value in tuple(bound.items()): + if isinstance(value, str) and name in ("dtype", "out_dtype"): + bound[name] = DType.from_name(value) + for parameter in original_attrs: + value = bound.get(parameter.name) + annotation = parameter.annotation + if ( + isinstance(value, str) + and isinstance(annotation, type) + and issubclass(annotation, Enum) + ): + bound[parameter.name] = annotation(value) + Op.__init__(self, **bound) + + attrs = { + "__init__": constructor, + "__module__": __name__, + "capability": OpCapability("cpu_tensor"), + "resource": "cpu_engine", + "execution_mesh": CORE, + **attributes, + } + for index in range(arity): + attrs[f"arg{index}"] = ParamDef( + kind="input", effect=MemoryEffect.READ, pattern=TensorPattern() + ) + attrs["dst"] = ParamDef( + kind="input", + effect=MemoryEffect.WRITE, + pattern=TensorPattern(storage=StorageKind.GMEM), + ) + name = "cpu_" + hir._op_schema.name + (f"_{arity}" if hir is Concat else "") + cls = register_op(dialect="T", category="cpu", name=name)( + type("Cpu" + hir.__name__ + str(arity), (Op,), attrs) + ) + + def reconstruct(op): + return hir(**{p.name: getattr(op, p.name) for p in original_attrs}) + + @register_typeinfer(cls) + def infer(call, ctx): + return UnitType() + + @register_access_relation(cls) + def relation(call, ctx): + original = Call(reconstruct(call.target), call.args[:arity], type=UnitType()) + local = TypeInferContext( + memo={id(arg): (arg, _semantic_type(ctx.type_of(arg))) for arg in original.args} + ) + original.type = typeinfer_registry.lookup(hir)(original, local).type + local.memo[id(original)] = (original, original.type) + boundaries = access_relation_registry.lookup(hir)(original, local) + output = boundaries[-1] + if not output.relation.is_single_valued(): + relation = output.relation + n = relation.dim(isl.dim_type.IN) + rank = relation.dim(isl.dim_type.OUT) + write = relation.move_dims(isl.dim_type.IN, n, isl.dim_type.OUT, 0, rank).add_dims( + isl.dim_type.OUT, rank + ) + for axis in range(rank): + write = write.equate(isl.dim_type.IN, n + axis, isl.dim_type.OUT, axis) + domain = write.domain() + boundaries = tuple( + replace( + item, + relation=item.relation.add_dims(isl.dim_type.IN, rank).intersect_domain(domain), + ) + for item in boundaries[:-1] + ) + (replace(output, relation=write),) + return (*boundaries, boundaries[-1]) + + @register_verify_stmt(cls) + def verify(call, ctx): + target = ctx.resolve_target() + if target is not None and target.name not in ("cpu", "amx"): + ctx.error(call, "Host tensor kernels require a CPU/AMX target") + original = Call(reconstruct(call.target), call.args[:arity], type=UnitType()) + local = TypeInferContext( + memo={id(arg): (arg, _semantic_type(ctx.type_of(arg))) for arg in original.args} + ) + expected = typeinfer_registry.lookup(hir)(original, local).type + actual = ctx.type_of(call.args[-1]) + if expected.shape != actual.shape or expected.dtype != actual.dtype: + ctx.error(call, f"CPU tensor kernel writes {actual}, expected {expected}") + + @register_cost_evaluator(cls) + def cost(call, ctx): + original = Call(reconstruct(call.target), call.args[:arity], type=call.type) + result = cost_evaluator_registry.lookup(hir)(original, ctx) + return replace(result, traffic=(*result.traffic, TrafficBytes())) + + @register_schedule_eval(cls) + def evaluate(ctx): + op = reconstruct(ctx.op) + result = eval_registry.lookup(hir)(ctx.for_op(op, ctx.args[:arity], ctx.result_type)) + return result + + INSTRUCTIONS[hir, arity] = cls + globals()[cls.__name__] = cls + + +for hir, arity in ( + (Reshape, 1), + (Transpose, 1), + (Softplus, 1), + (Binary, 2), + (Unary, 1), + (MatMul, 2), + (RMSNorm, 2), + (SoftMax, 1), + (Silu, 1), + (Sigmoid, 1), + (Cast, 1), + (Reduce, 1), + (IndexSelect, 2), + (RepeatInterleave, 1), + (FullLike, 1), + (Arange, 0), + (Concat, 2), +): + _register(hir, arity) + + +def instruction_for(operation, arity, dtype): + cls = INSTRUCTIONS[type(operation), arity] + attrs = { + p.name: getattr(operation, p.name) + for p in collect_param_defs(type(operation)) + if p.kind == "attribute" + } + if "dtype" in attrs: + attrs["dtype"] = dtype + return cls(**attrs) diff --git a/src/tilefoundry/ir/tir/memory/copy.py b/src/tilefoundry/ir/tir/memory/copy.py index 929af1e16..5b88407c0 100644 --- a/src/tilefoundry/ir/tir/memory/copy.py +++ b/src/tilefoundry/ir/tir/memory/copy.py @@ -28,7 +28,7 @@ class Copy(Op): capability = OpCapability(None) - execution_mesh = utils.thread_execution_mesh() + execution_mesh = utils.memory_execution_mesh() src = ParamDef( kind="input", diff --git a/src/tilefoundry/ir/tir/memory/fill.py b/src/tilefoundry/ir/tir/memory/fill.py index 31a9aad1f..c4ee41ba1 100644 --- a/src/tilefoundry/ir/tir/memory/fill.py +++ b/src/tilefoundry/ir/tir/memory/fill.py @@ -22,7 +22,7 @@ class Fill(Op): """Fills ``tensor`` element-wise with ``value`` (rank-0 scalar).""" - execution_mesh = utils.thread_execution_mesh() + execution_mesh = utils.memory_execution_mesh() tensor = ParamDef( kind="input", diff --git a/src/tilefoundry/passes/transforms/convert_hir_to_tir.py b/src/tilefoundry/passes/transforms/convert_hir_to_tir.py index 2876c31ec..2d7da3fc9 100644 --- a/src/tilefoundry/passes/transforms/convert_hir_to_tir.py +++ b/src/tilefoundry/passes/transforms/convert_hir_to_tir.py @@ -1619,7 +1619,12 @@ def instruction(self, call: Call, op, operands, ctx) -> Expr: attributes = { param.name: getattr(op, param.name) for param in type(op)._op_schema.signature - if param.kind == "attribute" and (not param.has_default or param.name == "fill") + if param.kind == "attribute" + and ( + not param.has_default + or getattr(op, param.name) != param.default + or param.name == "fill" + ) } ctx.builder.add(Evaluate(type(op)(**attributes), tuple(issued))) result_param = next(param for param in params if param.effect & MemoryEffect.WRITE) diff --git a/src/tilefoundry/schedule/facts.py b/src/tilefoundry/schedule/facts.py index 6e7e62708..825612540 100644 --- a/src/tilefoundry/schedule/facts.py +++ b/src/tilefoundry/schedule/facts.py @@ -78,13 +78,13 @@ def one(target: Target, wanted: str | type) -> dict[str, Any]: branches[0].attribute, { op_identifier(item.declaration): declared_execution_mesh( - op_type(**{item.attribute: item.declaration}) + op_type(**{item.attribute: item.declaration}), target ) for item in branches }, ) else: - execution_mesh = declared_execution_mesh(op_type) + execution_mesh = declared_execution_mesh(op_type, target) sections = PatternPrinter().declaration(op_type) return { "target": target.identity, diff --git a/src/tilefoundry/target/__init__.py b/src/tilefoundry/target/__init__.py index b08fda7ff..df170b4e9 100644 --- a/src/tilefoundry/target/__init__.py +++ b/src/tilefoundry/target/__init__.py @@ -15,6 +15,7 @@ from tilefoundry.target.cuda import CudaArchitecture, CudaDevice, CudaTarget from tilefoundry.target.cuda.spec import H200_SXM_ID from tilefoundry.target.facts import ( + MemoryExecutionFacts, TargetFactsError, TopologyFacts, TopologyLevelFacts, @@ -75,6 +76,7 @@ def validate_cuda_topology_levels(target: Target, names) -> None: "CudaTarget", "Device", "MemoryHierarchyFacts", + "MemoryExecutionFacts", "PerformanceServiceFacts", "Target", "TargetFactsError", diff --git a/src/tilefoundry/target/amx/device.py b/src/tilefoundry/target/amx/device.py index 6eaa6a72f..8f18c4537 100644 --- a/src/tilefoundry/target/amx/device.py +++ b/src/tilefoundry/target/amx/device.py @@ -34,6 +34,8 @@ class AppleM2Pro(Device): _unit_flops: tuple[tuple[str, tuple[tuple[DType, int], ...]], ...] + performance_amx_f16: int | None + performance_host_lane_events: int | None def _python_import_module(self) -> str: if type(self) is AppleM2Pro: diff --git a/src/tilefoundry/target/amx/facts.py b/src/tilefoundry/target/amx/facts.py index 6c93df61a..7c6f112f7 100644 --- a/src/tilefoundry/target/amx/facts.py +++ b/src/tilefoundry/target/amx/facts.py @@ -13,6 +13,7 @@ MemoryHierarchyFacts, MemoryLevelRelation, MemoryRelationKind, + PerformanceServiceFacts, ThroughputFacts, ) from tilefoundry.ir.types import DType @@ -99,7 +100,35 @@ def throughput(target: AmxTarget, query: object = None) -> ThroughputFacts: ) +def performance_service(target: AmxTarget, query: object = None) -> PerformanceServiceFacts: + """Project the documented M2 Pro compute and host-service model. + + FP16 work shares AMX clusters; FP32 host work uses the measured NEON rate. + Non-FLOP services use the document's explicitly estimated lane-event rate. + """ + unit = "core" if query is None else str(query) + if unit not in ("core", "amx"): + raise ValueError(f"AMX performance services have no topology level {unit!r}") + device = target.device + share = device.performance_core_count + rates = dict(device.unit_flops_per_second.get("neon", {})) + if device.performance_amx_f16 is not None: + rates[DType.f16] = device.performance_amx_f16 * device.sm_count // share + lane_events = device.performance_host_lane_events + return PerformanceServiceFacts( + unit_flops=tuple(sorted(rates.items(), key=lambda item: item[0].name)), + unit_ops=tuple( + (kind, lane_events) + for kind in ("integer", "predicate", "select", "local_move", "special") + if lane_events is not None + ), + unit_bandwidth=(("gmem", device.unified_memory_bandwidth_bytes_per_second // share),), + unit=unit, + ) + + __all__ = [ "memory_hierarchy", "throughput", + "performance_service", ] diff --git a/src/tilefoundry/target/amx/hardware/apple_amx.toml b/src/tilefoundry/target/amx/hardware/apple_amx.toml index a5e34149e..07bf6c532 100644 --- a/src/tilefoundry/target/amx/hardware/apple_amx.toml +++ b/src/tilefoundry/target/amx/hardware/apple_amx.toml @@ -11,11 +11,11 @@ source = "https://github.com/corsix/amx" conditions = "Apple publishes no AMX documentation, so the architecture is described from a reverse-engineered reference" [facts.capabilities] -value = ["amx_outer_product", "amx_resident_accumulator"] +value = ["amx_outer_product", "amx_resident_accumulator", "cpu_tensor"] unit = "name" origin = "reference" source = "https://github.com/corsix/amx" -conditions = "instruction families the AMX code path may select" +conditions = "instruction families the AMX code path may select; cpu_tensor identifies the host software contracts in src/tilefoundry/ir/tir/cpu/tensor.py" [facts.instruction.compute_dtypes] value = ["f16", "f32"] diff --git a/src/tilefoundry/target/amx/hardware/apple_m2_pro.toml b/src/tilefoundry/target/amx/hardware/apple_m2_pro.toml index 083c4097e..d52db042d 100644 --- a/src/tilefoundry/target/amx/hardware/apple_m2_pro.toml +++ b/src/tilefoundry/target/amx/hardware/apple_m2_pro.toml @@ -110,3 +110,18 @@ unit = "element" origin = "derived" source = "Arm A64 instruction set, FMLA (by element)" conditions = "four FMLA vD.4s, vN.4s, vM.s[i] over four accumulator registers are the smallest whole f32 outer product NEON issues" + + +[facts.performance.amx.f16] +value = 504869155110 +unit = "flop/s" +origin = "estimated" +source = "throughput.amx.f32 in this document; https://github.com/corsix/amx/issues/6#issuecomment-1477016531" +conditions = "M2 Pro nominal FP16-input/FP32-accumulator service estimate, anchored to this host's measured 504.9 GFLOP/s AMX F32 GEMM rate. The reference reports near-equal full-Z rates for mixed FP16 (one 4096-byte accumulator, 1636.8 GFLOP/s) and FP32 (four accumulators totaling 4096 bytes, 1636.9 GFLOP/s). Apply only that ratio: floor(504900000000 * 16368 / 16369). This is not a new M2 Pro FP16 measurement or a transfer of another chip's absolute throughput. Cluster sharing is applied by the Facts projection." + +[facts.performance.host.lane_events] +value = 53850000000 +unit = "op/s" +origin = "estimated" +source = "throughput.neon.f32 in this document; https://github.com/corsix/amx/issues/7#issuecomment-1475996279" +conditions = "M2 Pro optimistic CPU service model, not measured integer or transcendental throughput. Divide this host's measured 107.7 GFLOP/s NEON F32 FMA rate by two FLOPs per lane FMA: 53.85 billion lane events/s per performance core. Assign one lane event per integer, predicate, select, local-move or special-function result. The reference explains NEON lane/FMA counting, not these services' throughput. Polynomial expansion, instruction mix and Torch overhead are outside this nominal model." diff --git a/src/tilefoundry/target/amx/spec.py b/src/tilefoundry/target/amx/spec.py index a60c402dc..ff95273d7 100644 --- a/src/tilefoundry/target/amx/spec.py +++ b/src/tilefoundry/target/amx/spec.py @@ -97,6 +97,10 @@ def build_apple_m2_pro(document: HardwareDocument) -> AppleM2Pro: "memory.unified.bandwidth", unit="byte/s" ), _unit_flops=unit_flops, + performance_amx_f16=reader.optional_integer("performance.amx.f16", unit="flop/s"), + performance_host_lane_events=reader.optional_integer( + "performance.host.lane_events", unit="op/s" + ), ) reader.text("geometry.neon_f32_outer_product") reader.close() diff --git a/src/tilefoundry/target/amx/target.py b/src/tilefoundry/target/amx/target.py index eb8aaa0b3..0f3bec270 100644 --- a/src/tilefoundry/target/amx/target.py +++ b/src/tilefoundry/target/amx/target.py @@ -26,6 +26,7 @@ select, ) from tilefoundry.target.facts import ( + MemoryExecutionFacts, TopologyFacts, TopologyLevelFacts, facts_result, @@ -134,6 +135,11 @@ def get_facts(self, facts_type: type, query: object | None = None): """Project AMX hardware through the facts this Target owns.""" if facts_type is TopologyFacts and query is None: return facts_result(self, facts_type, self._topology_facts()) + if facts_type is MemoryExecutionFacts and query is None: + return facts_result( + self, facts_type, + MemoryExecutionFacts(tuple(level.name for level in self._topology_facts().topologies)), + ) if facts_type is TopologyLevelFacts: level = self._topology_facts().level( query if isinstance(query, str) else None @@ -144,15 +150,19 @@ def get_facts(self, facts_type: type, query: object | None = None): from tilefoundry.analysis.facts import ( # noqa: PLC0415 MemoryHierarchyFacts, + PerformanceServiceFacts, ThroughputFacts, ) from tilefoundry.target.amx.facts import ( # noqa: PLC0415 memory_hierarchy, + performance_service, throughput, ) if facts_type is MemoryHierarchyFacts: return facts_result(self, facts_type, memory_hierarchy(self, query)) + if facts_type is PerformanceServiceFacts: + return facts_result(self, facts_type, performance_service(self, query)) if facts_type is ThroughputFacts: return facts_result(self, facts_type, throughput(self, query)) return super().get_facts(facts_type, query) diff --git a/src/tilefoundry/target/base.py b/src/tilefoundry/target/base.py index 1b5e77df0..af520399d 100644 --- a/src/tilefoundry/target/base.py +++ b/src/tilefoundry/target/base.py @@ -11,7 +11,7 @@ from typing import Any, ClassVar, Mapping, TypeVar from tilefoundry.ir.types import DType, Topology -from tilefoundry.target.facts import TopologyFacts, facts_result +from tilefoundry.target.facts import MemoryExecutionFacts, TopologyFacts, facts_result from tilefoundry.target.hardware.envelope import ( DuplicateRegistrationError, HardwareDocument, @@ -243,6 +243,8 @@ def get_facts( """ if facts_type is TopologyFacts and query is None: return facts_result(self, facts_type, TopologyFacts(())) + if facts_type is MemoryExecutionFacts and query is None: + return facts_result(self, facts_type, MemoryExecutionFacts(())) raise UnsupportedCapabilityError( f"{_target_summary(self)}: no Facts projection for " f"{getattr(facts_type, '__name__', facts_type)!r}" diff --git a/src/tilefoundry/target/cuda/target.py b/src/tilefoundry/target/cuda/target.py index c2f1cf39c..dbb3b8b32 100644 --- a/src/tilefoundry/target/cuda/target.py +++ b/src/tilefoundry/target/cuda/target.py @@ -26,6 +26,7 @@ build_cuda_device, ) from tilefoundry.target.facts import ( + MemoryExecutionFacts, TopologyFacts, TopologyLevelFacts, facts_result, @@ -155,6 +156,10 @@ def get_facts(self, facts_type: type, query: object | None = None): """Project CUDA hardware through the facts this Target owns.""" if facts_type is TopologyFacts and query is None: return facts_result(self, facts_type, self._topology_facts()) + if facts_type is MemoryExecutionFacts and query is None: + return facts_result( + self, facts_type, MemoryExecutionFacts((self._topology_facts().topologies[-1].name,)) + ) if facts_type is TopologyLevelFacts: level = self._topology_facts().level( query if isinstance(query, str) else None diff --git a/src/tilefoundry/target/facts.py b/src/tilefoundry/target/facts.py index 965444e8c..680784b9b 100644 --- a/src/tilefoundry/target/facts.py +++ b/src/tilefoundry/target/facts.py @@ -56,6 +56,13 @@ def parallel(self) -> TopologyLevelFacts | None: return self.level(self.parallel_level) +@dataclasses.dataclass(frozen=True) +class MemoryExecutionFacts: + """Topology levels that may execute synchronous memory operations.""" + + topologies: tuple[str, ...] + + class TargetFactsError(Exception): """A Target failed to provide the requested immutable Facts aggregate.""" @@ -88,5 +95,6 @@ def facts_result( "TargetFactsError", "TopologyFacts", "TopologyLevelFacts", + "MemoryExecutionFacts", "facts_result", ] diff --git a/src/tilefoundry/visitor_registry/candidates/amx.py b/src/tilefoundry/visitor_registry/candidates/amx.py new file mode 100644 index 000000000..6488a37bc --- /dev/null +++ b/src/tilefoundry/visitor_registry/candidates/amx.py @@ -0,0 +1,8 @@ +"""Apple AMX matrix instruction candidates.""" + +from tilefoundry.ir.hir.nn.matmul import MatMul +from tilefoundry.ir.tir.amx.instructions import Fma16 + +from . import register_candidates + +register_candidates(MatMul, (Fma16,)) diff --git a/src/tilefoundry/visitor_registry/instruction.py b/src/tilefoundry/visitor_registry/instruction.py index 671129617..ca8f20b1b 100644 --- a/src/tilefoundry/visitor_registry/instruction.py +++ b/src/tilefoundry/visitor_registry/instruction.py @@ -313,7 +313,7 @@ def _pattern_refusal(matcher: PatternMatcher, op, parameter=None) -> ScheduleRef return ScheduleRefusal(PatternPrinter().refusal_reason(refusal, parameter), path=path, op=op) -def admit_fields(op, operands, mesh, matcher): +def admit_fields(op, operands, mesh, matcher, *, target=None): """Admit attributes, execution mesh and operand fields before tile projection.""" for param in type(op)._op_schema.signature: if param.kind != "attribute" or param.name == "execution_mesh": @@ -336,7 +336,7 @@ def admit_fields(op, operands, mesh, matcher): raise ScheduleRefusal(f"{param.name} must be int, got {type(value).__name__}", op=op) if param.pattern is not None and not matcher.match(param.pattern, value): raise _pattern_refusal(matcher, op, param.name) - if not matcher.match(declared_execution_mesh(op), mesh): + if not matcher.match(declared_execution_mesh(op, target), mesh): raise _pattern_refusal(matcher, op, "mesh") for param, type_ in zip(call_params(op, len(operands)), operands, strict=True): pattern = _operand_pattern(param, op, matcher.bindings) @@ -375,7 +375,7 @@ def check_call(op, operands, mesh, *, frame, repeat=None, order=None, projection ) atom = getattr(op, "atom", None) matcher = PatternMatcher(dict(getattr(atom, "bindings", {}))) - admit_fields(op, whole, mesh, matcher) + admit_fields(op, whole, mesh, matcher, target=None if frame is None else frame.target) try: tiles = ( projection diff --git a/tests/schedule/test_fixtures.py b/tests/schedule/test_fixtures.py index 10595b381..4254880af 100644 --- a/tests/schedule/test_fixtures.py +++ b/tests/schedule/test_fixtures.py @@ -230,6 +230,40 @@ T.relu all targets T.reduce all targets """ +_AMX_FACTS_TEXT = """\ +target apple.m2_pro +instructions + T.amx_fma16 amx_outer_product + T.amx_load_x amx_outer_product + T.amx_load_y amx_outer_product + T.amx_store amx_outer_product + T.arange all targets + T.binary all targets + T.cast all targets + T.clamp all targets + T.unary all targets + T.where all targets + T.cpu_arange cpu_tensor + T.cpu_binary cpu_tensor + T.cpu_cast cpu_tensor + T.cpu_concat_2 cpu_tensor + T.cpu_full_like cpu_tensor + T.cpu_index_select cpu_tensor + T.cpu_matmul cpu_tensor + T.cpu_reduce cpu_tensor + T.cpu_repeat_interleave cpu_tensor + T.cpu_reshape cpu_tensor + T.cpu_rms_norm cpu_tensor + T.cpu_sigmoid cpu_tensor + T.cpu_silu cpu_tensor + T.cpu_softmax cpu_tensor + T.cpu_softplus cpu_tensor + T.cpu_transpose cpu_tensor + T.cpu_unary cpu_tensor + T.copy all targets + T.relu all targets + T.reduce all targets +""" CANDIDATE_GOLDEN = ( Path(__file__).parents[1] / "fixtures" @@ -1480,10 +1514,10 @@ def _reshape_tensor(expected: bool = False) -> torch.Tensor: _TARGET_NEUTRAL_FACTS.format(target="cpu"), ), ( - "Apple facts list only target neutral instructions", + "Apple facts list AMX and CPU instructions", "apple.m2_pro", (), - _TARGET_NEUTRAL_FACTS.format(target="apple.m2_pro"), + _AMX_FACTS_TEXT, ), ) ),