From 73f6ec9b8d0d67256e39ffc3d6a79456dc176e25 Mon Sep 17 00:00:00 2001 From: QiHang Zheng <597323109@qq.com> Date: Mon, 19 Oct 2026 02:00:51 +0800 Subject: [PATCH 1/8] fix(analysis,schedule): preserve instance footprints and instruction choices --- src/tilefoundry/analysis/footprint.py | 20 +++--- .../passes/transforms/convert_hir_to_tir.py | 7 +- tests/analysis/test_analyze_cross_module.py | 70 ++++++++++++++++++- tests/schedule/test_attribute_preservation.py | 40 +++++++++++ 4 files changed, 123 insertions(+), 14 deletions(-) create mode 100644 tests/schedule/test_attribute_preservation.py diff --git a/src/tilefoundry/analysis/footprint.py b/src/tilefoundry/analysis/footprint.py index d649e569..9390beac 100644 --- a/src/tilefoundry/analysis/footprint.py +++ b/src/tilefoundry/analysis/footprint.py @@ -1187,27 +1187,23 @@ def resolve_footprint(footprint, ctx): continue def project_image(image): + relation = isl.map.from_range(image) for axis in reversed(range(len(item.ownership))): - image = image.set_dim_name(isl.dim_type.SET, axis, f"__tf_instance_{axis}") - image = image.move_dims( - isl.dim_type.PARAM, image.dim(isl.dim_type.PARAM), isl.dim_type.SET, axis, 1 - ) + # Instance coordinates are relation dimensions, not + # external parameters with an unknown SSA binding. Keep + # their correlation with addresses through projection. + relation = relation.move_dims(isl.dim_type.IN, 0, isl.dim_type.OUT, axis, 1) source = BufferAlias( item.target, - AccessRelation(isl.map.from_range(image), item.values), + AccessRelation(relation, item.values), AliasKind.FORWARD, ) projected, precision = binding.project(source) relation = projected.relation.relation values = dict(projected.relation.values) - relation = relation.project_out(isl.dim_type.IN, 0, depth) for axis in range(len(item.ownership)): - parameter = relation.find_dim_by_name( - isl.dim_type.PARAM, f"__tf_instance_{axis}" - ) - relation = relation.move_dims( - isl.dim_type.OUT, axis, isl.dim_type.PARAM, parameter, 1 - ) + relation = relation.move_dims(isl.dim_type.OUT, axis, isl.dim_type.IN, depth, 1) + relation = relation.project_out(isl.dim_type.IN, 0, depth) image = relation.range() return projected.target, image, values, precision diff --git a/src/tilefoundry/passes/transforms/convert_hir_to_tir.py b/src/tilefoundry/passes/transforms/convert_hir_to_tir.py index b95f657f..0e1d9ded 100644 --- a/src/tilefoundry/passes/transforms/convert_hir_to_tir.py +++ b/src/tilefoundry/passes/transforms/convert_hir_to_tir.py @@ -1618,7 +1618,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/tests/analysis/test_analyze_cross_module.py b/tests/analysis/test_analyze_cross_module.py index 610750c3..ada58c4f 100644 --- a/tests/analysis/test_analyze_cross_module.py +++ b/tests/analysis/test_analyze_cross_module.py @@ -11,8 +11,10 @@ from __future__ import annotations from dataclasses import dataclass, replace +from types import SimpleNamespace from typing import ClassVar +import isl import pytest from tests.fixtures.logical.hir_composition import REFERENCE_PROGRAMS, CrossModule, Expert @@ -20,10 +22,20 @@ from tests.fixtures.shapes.scaled_modules import FusedScaledParent, ScaledChild, StackedScaledParent from tests.models.qwen3_1_7b.case import CASE as QWEN_CASE from tests.models.qwen3_1_7b.model import Qwen3_1_7B +from tilefoundry.analysis.access import SSAValue from tilefoundry.analysis.api import analyze from tilefoundry.analysis.check import check_program from tilefoundry.analysis.errors import AnalysisError -from tilefoundry.analysis.metadata import ComputeCostMetadata, MemoryMetadata, RegionMemoryMetadata +from tilefoundry.analysis.footprint import ReachedAddresses, resolve_footprint +from tilefoundry.analysis.metadata import ( + BindingValuation, + ComputeCostMetadata, + Footprint, + LoopNestValuation, + MemoryMetadata, + RegionMemoryMetadata, +) +from tilefoundry.analysis.precision import AnalysisPrecision from tilefoundry.analysis.registry import Analyzer from tilefoundry.dsl import * from tilefoundry.ir.core import Call, Constant, Tuple, Var, attach_metadata, get_metadata @@ -46,6 +58,7 @@ from tilefoundry.ir.types.utils import tensor_types from tilefoundry.ir.visitor import collect_exprs from tilefoundry.target import CudaTarget, Target +from tilefoundry.visitor_registry.access_relation import AccessRelation from tilefoundry.visitor_registry.call_frame import CallFrame, CallFrameVisitor _H200 = "nvidia.h200_sxm" @@ -541,3 +554,58 @@ def test_repeated_callee_placement_compares_alias_even_at_the_same_offset(): ] assert len(buffers) == 2 assert all(len(record.offsets) == 1 for record in buffers) + + +@pytest.mark.parametrize( + "ownership,coordinates", + [ + (("layers",), "{ [layer, addr] : 0 <= layer < 3 and addr = 2*layer }"), + ( + ("layers", "experts"), + "{ [layer, expert, addr] : 0 <= layer < 2 and 0 <= expert < 3 and addr = 10*layer + expert }", + ), + ], +) +def test_nested_weight_footprint_projection_keeps_instance_address_correlation( + ownership, coordinates +): + + weight = Var("w", type=TensorType((32,), DType.f32, None, StorageKind.GMEM), is_const=True) + item = ReachedAddresses( + SSAValue(weight), + None, + DType.f32, + isl.set(coordinates), + AnalysisPrecision.EXACT, + ownership=ownership, + ) + footprint = Footprint((((), (item,)),)) + + def project(source): + shifted = source.relation.relation.apply_range(isl.map("{ [i] -> [j] : j=i+4 }")) + return replace( + source, relation=AccessRelation(shifted, source.relation.values) + ), AnalysisPrecision.EXACT + + values = { + BindingValuation: SimpleNamespace(project=project), + LoopNestValuation: SimpleNamespace(domain=lambda: isl.set("{ [] }")), + } + result = resolve_footprint( + footprint, SimpleNamespace(request_valuation=lambda kind: values[kind]) + ) + actual = result.groups[0][1][0] + rank = len(ownership) + # Shift only the address coordinate; instance coordinates remain correlated. + offset = isl.map( + "{ [" + + ",".join(f"i{k}" for k in range(rank + 1)) + + "] -> [" + + ",".join([*(f"i{k}" for k in range(rank)), f"i{rank}+4"]) + + "] }" + ) + expected = item.reached.apply(offset) + assert actual.reached.is_equal(expected) + assert actual.must.is_equal(expected) + assert actual.ownership == ownership + assert not actual.values diff --git a/tests/schedule/test_attribute_preservation.py b/tests/schedule/test_attribute_preservation.py new file mode 100644 index 00000000..29a4c1a1 --- /dev/null +++ b/tests/schedule/test_attribute_preservation.py @@ -0,0 +1,40 @@ +"""Lowering must retain authored choices that differ from instruction defaults.""" + +import torch + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, T, Tensor, Topology, tf +from tilefoundry.ir.tir.arange import Arange +from tilefoundry.ir.visitor import StmtVisitor, walk_prim_function +from tilefoundry.schedule import finalize +from tilefoundry.target import CudaTarget + + +@module( + entry="run", + target=CudaTarget("nvidia.h200_sxm"), + topologies=(Topology("cta", 1), Topology("thread", 32)), +) +class NondefaultRange: + @func + def run(): + with Mesh(("thread",), layout=(1,), names=("lane",)) as _thread: + output = tf.empty(Tensor[(4,), "i64", None, "rmem"]) + return tf.schedule((output,), op=T.arange(start=5, step=2)) + + +def test_nondefault_instruction_attributes_survive_lowering(): + torch.testing.assert_close(NondefaultRange(), torch.tensor([5, 7, 9, 11])) + + class Find(StmtVisitor): + def __init__(self): + self.operations = [] + + def visit_Evaluate(self, statement): + if isinstance(statement.callable, Arange): + self.operations.append(statement.callable) + + visitor = Find() + lowered = finalize(NondefaultRange) + walk_prim_function(visitor, lowered.entry_function()) + assert [(op.start, op.step) for op in visitor.operations] == [(5, 2)] From 79565d2f22f41f1168e18baa01dbc633ddc4177c Mon Sep 17 00:00:00 2001 From: QiHang Zheng <597323109@qq.com> Date: Mon, 19 Oct 2026 02:00:52 +0800 Subject: [PATCH 2/8] feat(amx): admit memory operations on Apple executor scopes --- docs/spec/tir.md | 5 +++++ src/tilefoundry/ir/pattern/utils.py | 20 +++++++++++++++----- src/tilefoundry/ir/tir/memory/copy.py | 8 ++++++-- src/tilefoundry/ir/tir/memory/fill.py | 6 +++++- tests/target/test_amx_target.py | 22 ++++++++++++++++++++++ 5 files changed, 53 insertions(+), 8 deletions(-) diff --git a/docs/spec/tir.md b/docs/spec/tir.md index 69494d81..6c76d0b3 100644 --- a/docs/spec/tir.md +++ b/docs/spec/tir.md @@ -582,6 +582,11 @@ 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. +Copy and Fill admit CUDA `thread`, CPU `core`, and Apple `amx` execution +scopes through the shared memory execution pattern. A sharded register operand +must name the selected scope. Their canonical CUDA `execution_mesh` descriptor +remains available; admission uses `execution_mesh_pattern()`. + ##### Copy ```python class Copy(Op): diff --git a/src/tilefoundry/ir/pattern/utils.py b/src/tilefoundry/ir/pattern/utils.py index bac916f3..692c6682 100644 --- a/src/tilefoundry/ir/pattern/utils.py +++ b/src/tilefoundry/ir/pattern/utils.py @@ -114,6 +114,15 @@ def thread_execution_mesh() -> MeshPattern: return MeshPattern(("thread",), _THREAD_EXECUTION_LAYOUT) +def memory_execution_mesh() -> OrPattern: + """Memory statements may execute on CUDA threads or Apple CPU/AMX units.""" + return OrPattern( + thread_execution_mesh(), + MeshPattern(("core",), _THREAD_EXECUTION_LAYOUT), + MeshPattern(("amx",), _THREAD_EXECUTION_LAYOUT), + ) + + def declared_execution_mesh(op_type) -> MeshPattern | OrPattern: """Return one op's reserved execution-mesh declaration or reject it.""" declaration = op_type if isinstance(op_type, type) else type(op_type) @@ -132,13 +141,14 @@ def declared_execution_mesh(op_type) -> MeshPattern | OrPattern: return declared -def _sharded_on(execution_mesh: MeshPattern) -> OrPattern: +def _sharded_on(execution_mesh: MeshPattern | OrPattern) -> OrPattern: """A lexical tile, or a shard explicitly held by *execution_mesh*.""" return OrPattern( - ShardLayoutPattern( - layout=WildcardPattern(), - attrs=WildcardPattern(), - mesh=execution_mesh, + *( + ShardLayoutPattern(layout=WildcardPattern(), attrs=WildcardPattern(), mesh=mesh) + for mesh in ( + execution_mesh.patterns if isinstance(execution_mesh, OrPattern) else (execution_mesh,) + ) ), LayoutPattern(), ) diff --git a/src/tilefoundry/ir/tir/memory/copy.py b/src/tilefoundry/ir/tir/memory/copy.py index 929af1e1..2d78b454 100644 --- a/src/tilefoundry/ir/tir/memory/copy.py +++ b/src/tilefoundry/ir/tir/memory/copy.py @@ -30,15 +30,19 @@ class Copy(Op): execution_mesh = utils.thread_execution_mesh() + @classmethod + def execution_mesh_pattern(cls): + return utils.memory_execution_mesh() + src = ParamDef( kind="input", effect=MemoryEffect.READ, - pattern=utils.operand_tile(0, execution_mesh=execution_mesh), + pattern=utils.operand_tile(0, execution_mesh=utils.memory_execution_mesh()), ) dst = ParamDef( kind="input", effect=MemoryEffect.WRITE, - pattern=utils.operand_tile(1, execution_mesh=execution_mesh), + pattern=utils.operand_tile(1, execution_mesh=utils.memory_execution_mesh()), ) between = (SameConstraint("shape", "src", "dst"), SameConstraint("dtype", "src", "dst")) diff --git a/src/tilefoundry/ir/tir/memory/fill.py b/src/tilefoundry/ir/tir/memory/fill.py index 31a9aad1..030ae454 100644 --- a/src/tilefoundry/ir/tir/memory/fill.py +++ b/src/tilefoundry/ir/tir/memory/fill.py @@ -24,10 +24,14 @@ class Fill(Op): execution_mesh = utils.thread_execution_mesh() + @classmethod + def execution_mesh_pattern(cls): + return utils.memory_execution_mesh() + tensor = ParamDef( kind="input", effect=MemoryEffect.WRITE, - pattern=utils.operand_tile(0, execution_mesh=execution_mesh), + pattern=utils.operand_tile(0, execution_mesh=utils.memory_execution_mesh()), ) value = ParamDef(kind="input", effect=MemoryEffect.READ, pattern=is_scalar_tensor()) diff --git a/tests/target/test_amx_target.py b/tests/target/test_amx_target.py index 898f0176..3438937a 100644 --- a/tests/target/test_amx_target.py +++ b/tests/target/test_amx_target.py @@ -9,8 +9,12 @@ import pytest +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.inspection import as_script from tilefoundry.ir.types import DType, Topology from tilefoundry.ir.types.dim import DimVar +from tilefoundry.schedule import finalize from tilefoundry.target import ( AmxTarget, TopologyFacts, @@ -110,3 +114,21 @@ def test_unmeasured_units_and_dtypes_have_no_throughput_entry(): target.device.throughput_for("amx", DType.f16) with pytest.raises(ValueError, match="no measured compute-throughput"): target.device.throughput_for("sme", DType.f32) + + +@pytest.mark.parametrize("level", ["core", "amx"]) +def test_memory_lowering_admits_apple_executor_scopes(level): + """Copy/fill remain valid when scheduled in Apple executor scopes.""" + + @module(entry="run", target=AmxTarget(), topologies=(Topology("core", 8), Topology("amx", 1))) + class AppleMemory: + @func + def run(): + with Mesh((level,), layout=(1,), names=("unit",)) as _unit: + x = tf.zeros(Tensor[(2, 3), "f32", None, "gmem"]) + return tf.transpose(x, perm=(1, 0)) + + # finalize verifies concrete statements as well as authored HIR admission. + result = finalize(AppleMemory) + text = as_script(result) + assert "T.fill" in text and "T.copy" in text From ce7ddc6b6efbf922f150b2615119c09e851d4158 Mon Sep 17 00:00:00 2001 From: QiHang Zheng <597323109@qq.com> Date: Mon, 19 Oct 2026 02:02:22 +0800 Subject: [PATCH 3/8] feat(examples): migrate Qwen AMX prefill to shared definitions --- examples/qwen_prefill-amx/.gitignore | 3 + examples/qwen_prefill-amx/README.md | 96 +++ examples/qwen_prefill-amx/build.py | 35 ++ examples/qwen_prefill-amx/check.py | 92 +++ examples/qwen_prefill-amx/config.py | 11 + examples/qwen_prefill-amx/configs/1.7b.json | 30 + examples/qwen_prefill-amx/configs/4b.json | 104 ++++ examples/qwen_prefill-amx/export_ir.py | 53 ++ .../qwen_prefill-amx/host_instructions.py | 214 +++++++ examples/qwen_prefill-amx/instructions.py | 164 +++++ examples/qwen_prefill-amx/model.py | 444 ++++++++++++++ examples/qwen_prefill-amx/native/amx.hpp | 560 ++++++++++++++++++ .../qwen_prefill-amx/native/prefill_ffi.cpp | 38 ++ examples/qwen_prefill-amx/native/runtime.hpp | 124 ++++ examples/qwen_prefill-amx/prepare.py | 38 ++ examples/qwen_prefill-amx/run.py | 43 ++ examples/qwen_prefill-amx/runtime_model.py | 176 ++++++ examples/qwen_prefill-amx/scheduled.py | 221 +++++++ examples/qwen_prefill-amx/verify_ir.py | 34 ++ examples/qwen_prefill-amx/weights.py | 103 ++++ 20 files changed, 2583 insertions(+) create mode 100644 examples/qwen_prefill-amx/.gitignore create mode 100644 examples/qwen_prefill-amx/README.md create mode 100644 examples/qwen_prefill-amx/build.py create mode 100644 examples/qwen_prefill-amx/check.py create mode 100644 examples/qwen_prefill-amx/config.py create mode 100644 examples/qwen_prefill-amx/configs/1.7b.json create mode 100644 examples/qwen_prefill-amx/configs/4b.json create mode 100644 examples/qwen_prefill-amx/export_ir.py create mode 100644 examples/qwen_prefill-amx/host_instructions.py create mode 100644 examples/qwen_prefill-amx/instructions.py create mode 100644 examples/qwen_prefill-amx/model.py create mode 100644 examples/qwen_prefill-amx/native/amx.hpp create mode 100644 examples/qwen_prefill-amx/native/prefill_ffi.cpp create mode 100644 examples/qwen_prefill-amx/native/runtime.hpp create mode 100644 examples/qwen_prefill-amx/prepare.py create mode 100644 examples/qwen_prefill-amx/run.py create mode 100644 examples/qwen_prefill-amx/runtime_model.py create mode 100644 examples/qwen_prefill-amx/scheduled.py create mode 100644 examples/qwen_prefill-amx/verify_ir.py create mode 100644 examples/qwen_prefill-amx/weights.py diff --git a/examples/qwen_prefill-amx/.gitignore b/examples/qwen_prefill-amx/.gitignore new file mode 100644 index 00000000..d048eb85 --- /dev/null +++ b/examples/qwen_prefill-amx/.gitignore @@ -0,0 +1,3 @@ +/artifacts/ +__pycache__/ +*.py[cod] diff --git a/examples/qwen_prefill-amx/README.md b/examples/qwen_prefill-amx/README.md new file mode 100644 index 00000000..75a50b67 --- /dev/null +++ b/examples/qwen_prefill-amx/README.md @@ -0,0 +1,96 @@ +# Qwen text Prefill on Apple AMX + +Qwen3-1.7B and Qwen3.5-4B, batch one, text Prefill only. The example provides +whole-model HIR, explicit AMX schedules, retained-call TIR and a hand-written +TVM FFI runtime. No Decode, MTP, vision or compiler code generation is involved. + +The model uses `ModuleList` instance slots: the dense model's 28 layers share +one layer definition; the hybrid model's 32 layers share two. Each instance +retains its own resources. Static calls select the heterogeneous layer slots; +the compiler does not clone or inline their bodies. + +## Files + +| File | Role | +|---|---| +| `model.py` | Dense/hybrid architecture and whole Prefill HIR | +| `scheduled.py` | Explicit destinations, activation/weight layout and worker schedule | +| `instructions.py` | Example-local X/Y load, FMA16, Z store and packing contracts | +| `host_instructions.py` | Typed host software kernels for non-AMX operations | +| `runtime_model.py` | One runtime twin per shared HIR definition | +| `weights.py` | GGUF aliases and Q8_0 decoding | +| `native/` | Standalone hand-written AMX kernel, thread pool and reverse-engineered instruction header | +| `prepare.py`, `run.py`, `export_ir.py`, `check.py` | Reproduction entry points | + +AMX contracts are local extensions in this example, not claims that an AMX +code generator has been added to TileFoundry. Host operations are software +kernels implemented using Torch; they are not AMX ISA instructions. The native +projection kernel and its FFI module have no llama.cpp dependency. + +## Schedule and numerical contract + +Weights arrive as Q8_0 GGUF and are dequantized to FP16. The original quantization +error remains. `Module.prepare` packs each projection from `[N,K]` to `[N/B,K,B]`, +where B is 64, or 32 for narrow projections. Activations are converted to FP16 +and packed into contiguous 32-row panels before accumulation. + +Eight logical workers distribute N panels cyclically. Each holds a 32x32 FP32 +Z accumulator, advances K by loading 32 weights into X and 32 activations into +Y, and issues FP16 outer products with FP32 accumulation. Z stays resident until +the tile is complete, then its interleaved rows are unpacked. The final M=1 +vocabulary projection uses a host kernel; it is not a Decode workload. + +The reverse-engineered AMX header is retained from the original user-provided +runtime. This path is validated on M2 Pro. Worker identity is logical; the +runtime does not promise OS thread affinity. + +## Run + +Use an Apple Silicon Mac with Command Line Tools and the existing `tf` Conda +environment. Dependencies are TileFoundry from this checkout, Torch, `gguf`, +`safetensors` and `tvm_ffi`. Use the same Python for building and running. + +```sh +conda activate tf +python examples/qwen_prefill-amx/build.py +python examples/qwen_prefill-amx/check.py + +python examples/qwen_prefill-amx/prepare.py --model 1.7b --checkpoint /path/to/Qwen3-1.7B-Q8_0.gguf +python examples/qwen_prefill-amx/run.py --model 1.7b + +python examples/qwen_prefill-amx/prepare.py --model 4b --checkpoint /path/to/Qwen3.5-4B-Q8_0.gguf +python examples/qwen_prefill-amx/run.py --model 4b +``` + +Config snapshots are in `configs/`. Checkpoint paths must be supplied explicitly; +there are no user-specific absolute paths or downloads during a run. `prepare` +uses the authored converters and writes canonical weights under +`artifacts/prepared-`. `run` loads only those prepared resources and +executes one 128-token Prefill over a fixed valid token pattern. It checks the +logits shape and finiteness, prints the first-token argmax and actual AMX call +counts, and saves `smoke.json`. It has no warmup, repetition or speedup table. + +Prompt lengths must be positive multiples of 32; `--sequence` selects the static +shape. The prepare/load/run stages are separate, and packed weight layouts do +not depend on prompt length. Large prepared weights and generated binaries are +ignored by Git. + +## Inspect the compiler path + +```sh +python examples/qwen_prefill-amx/export_ir.py --model 1.7b +python examples/qwen_prefill-amx/export_ir.py --model 4b +python examples/qwen_prefill-amx/verify_ir.py --model 1.7b +python examples/qwen_prefill-amx/verify_ir.py --model 4b +``` + +The three source files are emitted into `artifacts/ir-/`. Finalization +uses the current API and retained calls, with no old node-budget, recursion-limit +or no-traffic workaround. The runtime is hand-written; emitting TIR does not +compile that TIR into the executed native kernel. + +`check.py` is a small dense/hybrid migration check: actual converters, loaded +resources and native AMX are compared with the unscheduled whole-model evaluator. +The full-model smoke confirms execution, not generation quality. This migration +does not rerun performance experiments and does not claim the old llama.cpp +bridge's speedups for this runtime. diff --git a/examples/qwen_prefill-amx/build.py b/examples/qwen_prefill-amx/build.py new file mode 100644 index 00000000..7a5e9910 --- /dev/null +++ b/examples/qwen_prefill-amx/build.py @@ -0,0 +1,35 @@ +"""Build the standalone hand-written AMX TVM FFI module.""" + +from pathlib import Path +import platform, subprocess +import tvm_ffi +from tvm_ffi import libinfo + +ROOT = Path(__file__).resolve().parent +if __name__ == "__main__": + if platform.system() != "Darwin" or platform.machine() != "arm64": + raise SystemExit("This AMX runtime requires Apple Silicon macOS") + out = ROOT / "artifacts" + out.mkdir(exist_ok=True) + clang = subprocess.check_output(["xcrun", "--find", "clang++"], text=True).strip() + sdk = subprocess.check_output(["xcrun", "--show-sdk-path"], text=True).strip() + ffi = Path(libinfo.find_libtvm_ffi()) + subprocess.run( + [ + clang, + "-std=c++17", + "-O3", + "-shared", + "-fPIC", + "-isysroot", + sdk, + "-I" + str(Path(tvm_ffi.__file__).parent / "include"), + str(ROOT / "native/prefill_ffi.cpp"), + str(ffi), + "-Wl,-rpath," + str(ffi.parent), + "-o", + str(out / "prefill.dylib"), + ], + check=True, + ) + print(out / "prefill.dylib") diff --git a/examples/qwen_prefill-amx/check.py b/examples/qwen_prefill-amx/check.py new file mode 100644 index 00000000..4d1e289e --- /dev/null +++ b/examples/qwen_prefill-amx/check.py @@ -0,0 +1,92 @@ +"""Tiny migration smoke: prepare/load/native run against the model HIR.""" + +from pathlib import Path +import torch +from tilefoundry.runtime import DictResource, SafetensorsResource +from model import build +from scheduled import make_scheduled +from runtime_model import make_runtime, inputs + +ROOT = Path(__file__).resolve().parent +CONFIG = dict( + hidden_size=64, + intermediate_size=128, + vocab_size=64, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + rms_norm_eps=1e-6, + num_hidden_layers=2, + rope_theta=10000, +) + + +def raw_weights(mod, prefix="", out=None): + out = {} if out is None else out + for fn in mod.functions: + converters = dict(fn.converters) + for param in fn.params: + if not param.is_const: + continue + for p in ( + converters[param.name].params if param.name in converters else (param,) + ): + key = prefix + p.name + if key in out: + continue + x = torch.randn(p.type.shape) * 0.025 + if p.name.startswith("gamma"): + x = 1 + x + if p.name == "decay": + x = -x.abs() - 0.01 + out[key] = x.half() if p.type.dtype.name == "f16" else x.float() + for key, slot in mod.children.items(): + for i, ch in enumerate(slot.elements): + raw_weights( + ch, prefix + key + "." + (str(i) + "." if slot.indexed else ""), out + ) + return out + + +if __name__ == "__main__": + torch.set_num_threads(8) + torch.manual_seed(17) + for name in ("dense", "hybrid"): + c = dict(CONFIG) + if name == "hybrid": + c.update( + attn_output_gate=True, + layer_types=["linear_attention", "full_attention"], + linear_num_key_heads=16, + linear_num_value_heads=32, + linear_key_head_dim=4, + linear_value_head_dim=4, + linear_conv_kernel_dim=4, + rope_parameters={"partial_rotary_factor": 0.5, "rope_theta": 10000}, + ) + plain = build(c, 32) + scheduled = make_scheduled(c, 32) + raw = DictResource(raw_weights(plain)) + outputs = [] + args = inputs(c, 32, torch.arange(32) % 64) + for label, sem in [("plain", plain), ("scheduled", scheduled)]: + p = ROOT / "artifacts" / ("migration-check-" + name + "-" + label) + sem.prepare(raw, str(p), device="cpu") + resource = SafetensorsResource(str(p), device="cpu") + if label == "plain": + result = sem.load(resource)(*args) + else: + runtime = make_runtime(sem, c)() + runtime.load(resource) + result = runtime(*args) + assert runtime.stats["amx_calls"] > 0 + assert torch.isfinite(result).all() + outputs.append(result) + torch.testing.assert_close(outputs[1], outputs[0], rtol=2e-4, atol=2e-4) + print( + "PASS", + name, + "native runtime vs whole-model HIR", + float((outputs[1] - outputs[0]).abs().max()), + flush=True, + ) diff --git a/examples/qwen_prefill-amx/config.py b/examples/qwen_prefill-amx/config.py new file mode 100644 index 00000000..9d40686c --- /dev/null +++ b/examples/qwen_prefill-amx/config.py @@ -0,0 +1,11 @@ +"""Model config snapshots and local resource paths.""" + +import json +from pathlib import Path + +ROOT = Path(__file__).resolve().parent + + +def config(model): + data = json.loads((ROOT / "configs" / f"{model}.json").read_text()) + return data.get("text_config", data) diff --git a/examples/qwen_prefill-amx/configs/1.7b.json b/examples/qwen_prefill-amx/configs/1.7b.json new file mode 100644 index 00000000..044a86ec --- /dev/null +++ b/examples/qwen_prefill-amx/configs/1.7b.json @@ -0,0 +1,30 @@ +{ + "architectures": [ + "Qwen3ForCausalLM" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "bos_token_id": 151643, + "eos_token_id": 151645, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 2048, + "initializer_range": 0.02, + "intermediate_size": 6144, + "max_position_embeddings": 40960, + "max_window_layers": 28, + "model_type": "qwen3", + "num_attention_heads": 16, + "num_hidden_layers": 28, + "num_key_value_heads": 8, + "rms_norm_eps": 1e-06, + "rope_scaling": null, + "rope_theta": 1000000, + "sliding_window": null, + "tie_word_embeddings": true, + "torch_dtype": "bfloat16", + "transformers_version": "4.51.0", + "use_cache": true, + "use_sliding_window": false, + "vocab_size": 151936 +} \ No newline at end of file diff --git a/examples/qwen_prefill-amx/configs/4b.json b/examples/qwen_prefill-amx/configs/4b.json new file mode 100644 index 00000000..557d961b --- /dev/null +++ b/examples/qwen_prefill-amx/configs/4b.json @@ -0,0 +1,104 @@ +{ + "architectures": [ + "Qwen3_5ForConditionalGeneration" + ], + "image_token_id": 248056, + "model_type": "qwen3_5", + "text_config": { + "attention_bias": false, + "attention_dropout": 0.0, + "attn_output_gate": true, + "dtype": "bfloat16", + "eos_token_id": 248044, + "full_attention_interval": 4, + "head_dim": 256, + "hidden_act": "silu", + "hidden_size": 2560, + "initializer_range": 0.02, + "intermediate_size": 9216, + "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_conv_kernel_dim": 4, + "linear_key_head_dim": 128, + "linear_num_key_heads": 16, + "linear_num_value_heads": 32, + "linear_value_head_dim": 128, + "max_position_embeddings": 262144, + "mlp_only_layers": [], + "model_type": "qwen3_5_text", + "mtp_num_hidden_layers": 1, + "mtp_use_dedicated_embeddings": false, + "num_attention_heads": 16, + "num_hidden_layers": 32, + "num_key_value_heads": 4, + "rms_norm_eps": 1e-06, + "tie_word_embeddings": true, + "use_cache": true, + "vocab_size": 248320, + "mamba_ssm_dtype": "float32", + "rope_parameters": { + "mrope_interleaved": true, + "mrope_section": [ + 11, + 11, + 10 + ], + "rope_type": "default", + "rope_theta": 10000000, + "partial_rotary_factor": 0.25 + } + }, + "tie_word_embeddings": true, + "transformers_version": "4.57.0.dev0", + "video_token_id": 248057, + "vision_config": { + "deepstack_visual_indexes": [], + "depth": 24, + "hidden_act": "gelu_pytorch_tanh", + "hidden_size": 1024, + "in_channels": 3, + "initializer_range": 0.02, + "intermediate_size": 4096, + "model_type": "qwen3_5", + "num_heads": 16, + "num_position_embeddings": 2304, + "out_hidden_size": 2560, + "patch_size": 16, + "spatial_merge_size": 2, + "temporal_patch_size": 2 + }, + "vision_end_token_id": 248054, + "vision_start_token_id": 248053 +} \ No newline at end of file diff --git a/examples/qwen_prefill-amx/export_ir.py b/examples/qwen_prefill-amx/export_ir.py new file mode 100644 index 00000000..8bef2b68 --- /dev/null +++ b/examples/qwen_prefill-amx/export_ir.py @@ -0,0 +1,53 @@ +"""Export whole-model HIR, scheduled HIR and retained-call TIR.""" + +import argparse, json +from pathlib import Path +from tilefoundry.inspection import as_script +from tilefoundry.schedule import finalize +from tilefoundry.analysis.check import check_program +from config import config, ROOT +from model import build +from scheduled import make_scheduled + + +def export(model, sequence=128, output=None): + out = Path(output) if output else ROOT / "artifacts" / ("ir-" + model) + out.mkdir(parents=True, exist_ok=True) + c = config(model) + plain = build(c, sequence) + plain, _ = check_program(plain, plain.entry_function()) + scheduled = make_scheduled(c, sequence) + + def registrations(source): + future = "from __future__ import annotations\n" + imports = "import instructions, host_instructions\n" + return ( + source.replace(future, future + imports, 1) + if future in source + else imports + source + ) + + (out / "model.py").write_text(as_script(plain)) + (out / "scheduled.py").write_text(registrations(as_script(scheduled))) + lowered = finalize(scheduled) + (out / "finalized.py").write_text(registrations(as_script(lowered))) + report = dict( + model=model, + sequence=sequence, + layers=len(scheduled.layers), + shared_layer_definitions=len({id(layer) for layer in scheduled.layers}), + codegen=False, + performance_measured=False, + ) + (out / "manifest.json").write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(report), flush=True) + return out + + +if __name__ == "__main__": + p = argparse.ArgumentParser() + p.add_argument("--model", choices=("1.7b", "4b"), required=True) + p.add_argument("--sequence", type=int, default=128) + p.add_argument("--output") + a = p.parse_args() + export(a.model, a.sequence, a.output) diff --git a/examples/qwen_prefill-amx/host_instructions.py b/examples/qwen_prefill-amx/host_instructions.py new file mode 100644 index 00000000..2a1c3d6e --- /dev/null +++ b/examples/qwen_prefill-amx/host_instructions.py @@ -0,0 +1,214 @@ +"""Typed CPU tensor kernels, not ISA atoms or an opaque whole-model callback. + +Each kernel preserves its HIR operation's attributes and coordinate relations. +The schedule chooses these kernels explicitly for the CPU fallback portions of +an AMX program. No CPU code generation is provided by this module. +""" + +import isl + + +from dataclasses import replace +from enum import Enum + +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.transpose import Transpose +from tilefoundry.ir.hir.tensor.reshape import Reshape +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.where import Where +from tilefoundry.ir.pattern import TensorPattern +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 TypeInferContext +from tilefoundry.visitor_registry.registries import typeinfer_registry + +from instructions import CORE_PATTERN + +CORE = CORE_PATTERN +INSTRUCTIONS = {} + + +def _semantic_type(ty): + # A host tensor-kernel call can pass a scalar register by value. The HIR + # rule treats scalar literals as uniform rather than tensor residency. + 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(None), + "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 = "amx_host_" + 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(): + # A software row kernel writes each result cell exactly once. + # Expose output coordinates in the logical issue domain so the + # explicit destination/result relation is the identity, while + # each result may still read its entire normalization row. + 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.frame.target if ctx.frame is not None else None + 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_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), + (Where, 3), + (Arange, 0), + (Concat, 2), + (Concat, 3), +): + _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/examples/qwen_prefill-amx/instructions.py b/examples/qwen_prefill-amx/instructions.py new file mode 100644 index 00000000..ac60a744 --- /dev/null +++ b/examples/qwen_prefill-amx/instructions.py @@ -0,0 +1,164 @@ +"""Example-local AMX contracts; no machine-code generation is involved. + +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.ir.core import Op, OpCapability +from tilefoundry.ir.core.param_def import ParamDef, MemoryEffect +from tilefoundry.ir.core.register import register_op +from tilefoundry.ir.pattern import ( + TensorPattern, + LayoutPattern, + MeshPattern, + ComposedLayoutPattern, + WildcardPattern, +) +from tilefoundry.ir.types import ( + DType, + Layout, + Mesh, + Topology, + StorageKind, + UnitType, + LayoutBase, +) +from tilefoundry.visitor_registry import register_typeinfer +from tilefoundry.visitor_registry.access_relation import ( + register_access_relation, + identity_relations, + matmul_relations, +) +from tilefoundry.evaluator.registry import register_schedule_eval +from tilefoundry.evaluator.value import TensorValue + +Z_LAYOUT = Layout((32, (16, 2)), (32, (1, 16))) +AMX = Mesh((Topology("amx", 1),), Layout((1,), (1,)), ("unit",)) +CORE = Mesh((Topology("core", 8),), Layout((1,), (1,)), ("worker",)) + + +def mesh_pattern(name): + return MeshPattern( + (name,), + ComposedLayoutPattern( + inner=None, + offset=WildcardPattern(name + "_offset"), + outer=LayoutPattern(((1,),), ((1,),)), + ), + ) + + +AMX_PATTERN = mesh_pattern("amx") +CORE_PATTERN = mesh_pattern("core") +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) + + +@register_op(dialect="T", category="amx", name="amx_pack") +class Pack(Op): + capability = OpCapability("amx_outer_product") + resource = "cpu_engine" + execution_mesh = CORE_PATTERN + src = ParamDef(kind="input", effect=R, pattern=TensorPattern()) + dst = ParamDef( + kind="input", effect=W, pattern=TensorPattern(storage=StorageKind.GMEM) + ) + gmem_layout = ParamDef(kind="attribute", annotation=LayoutBase) + + +register_typeinfer(Pack)(lambda call, ctx: UnitType()) +register_access_relation(Pack)(identity_relations) +register_schedule_eval(Pack)( + lambda ctx: TensorValue(data=ctx.args[0].data.contiguous(), type=ctx.result_type) +) diff --git a/examples/qwen_prefill-amx/model.py b/examples/qwen_prefill-amx/model.py new file mode 100644 index 00000000..04bf2fbc --- /dev/null +++ b/examples/qwen_prefill-amx/model.py @@ -0,0 +1,444 @@ +"""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, 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()) + 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())", + "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, + S=S, + R=R, + H=H, + V=V, + ) + exec(compile(source, str(path), "exec"), namespace) + return namespace["QwenPrefill"] diff --git a/examples/qwen_prefill-amx/native/amx.hpp b/examples/qwen_prefill-amx/native/amx.hpp new file mode 100644 index 00000000..9fc36d26 --- /dev/null +++ b/examples/qwen_prefill-amx/native/amx.hpp @@ -0,0 +1,560 @@ +#pragma once +#include + +#ifdef __cplusplus +extern "C" { +#endif + +#define AMX_NOP_OP_IMM5(op, imm5) \ + __asm("nop\nnop\nnop\n.word (0x201000 + (%0 << 5) + %1)" \ + : \ + : "i"(op), "i"(imm5) \ + : "memory") + +#define AMX_OP_GPR(op, gpr) \ + __asm(".word (0x201000 + (%0 << 5) + 0%1 - ((0%1 >> 4) * 6))" \ + : \ + : "i"(op), "r"((uint64_t)(gpr)) \ + : "memory") + +#define AMX_LDX(gpr) AMX_OP_GPR(0, gpr) +#define AMX_LDY(gpr) AMX_OP_GPR(1, gpr) +#define AMX_STX(gpr) AMX_OP_GPR(2, gpr) +#define AMX_STY(gpr) AMX_OP_GPR(3, gpr) +#define AMX_LDZ(gpr) AMX_OP_GPR(4, gpr) +#define AMX_STZ(gpr) AMX_OP_GPR(5, gpr) +#define AMX_LDZI(gpr) AMX_OP_GPR(6, gpr) +#define AMX_STZI(gpr) AMX_OP_GPR(7, gpr) +#define AMX_EXTRX(gpr) AMX_OP_GPR(8, gpr) +#define AMX_EXTRY(gpr) AMX_OP_GPR(9, gpr) +#define AMX_FMA64(gpr) AMX_OP_GPR(10, gpr) +#define AMX_FMS64(gpr) AMX_OP_GPR(11, gpr) +#define AMX_FMA32(gpr) AMX_OP_GPR(12, gpr) +#define AMX_FMS32(gpr) AMX_OP_GPR(13, gpr) +#define AMX_MAC16(gpr) AMX_OP_GPR(14, gpr) +#define AMX_FMA16(gpr) AMX_OP_GPR(15, gpr) +#define AMX_FMS16(gpr) AMX_OP_GPR(16, gpr) +#define AMX_SET() AMX_NOP_OP_IMM5(17, 0) +#define AMX_CLR() AMX_NOP_OP_IMM5(17, 1) +#define AMX_VECINT(gpr) AMX_OP_GPR(18, gpr) +#define AMX_VECFP(gpr) AMX_OP_GPR(19, gpr) +#define AMX_MATINT(gpr) AMX_OP_GPR(20, gpr) +#define AMX_MATFP(gpr) AMX_OP_GPR(21, gpr) +#define AMX_GENLUT(gpr) AMX_OP_GPR(22, gpr) + +#ifdef __cplusplus +} +#endif + +template struct operand { + inline constexpr uint64_t bind(void *ptr) const noexcept { + return val | (((1ull << 57) - 1) & (uint64_t)ptr); + } + + inline constexpr operator uint64_t() const noexcept { return val; } + + template + inline constexpr T set_field(const uint64_t value) const noexcept { + constexpr uint64_t clr_mask = ~(((1ull << Width) - 1ull) << OffSet); + uint64_t value_mask = (value << OffSet); + return T{(val & clr_mask) | value_mask}; + } + + uint64_t val; +}; + +struct ldxy : operand { + /** + * @brief load multiple register + * + * factor = multiple ? 2 : 1 + * + * factor2 = multiple_means_four ? 2 : 1 + * + * so total load factor * factor2 * 64 byte + */ + constexpr auto multiple() const { return set_field<62, 1>(1); } + constexpr auto non_consecutive() const { return set_field<61, 1>(1); } + constexpr auto multiple_four() const { return set_field<60, 1>(1); } + /** + * @brief choice reg id, 0~5 + */ + constexpr auto register_index(const uint64_t index) const { + return set_field<56, 3>(index); + } +}; + +struct stxy : operand { + /** + * @brief store (multiple ? 2 : 1) * 64 bytes. + */ + constexpr auto multiple() const { return set_field<62, 1>(1); } + /** + * @brief choice reg id, 0~5 + */ + constexpr auto register_index(const uint64_t index) const { + return set_field<56, 3>(index); + } +}; + +struct ldstz : operand { + /** + * @brief load or store (multiple ? 2 : 1) * register. + */ + constexpr auto multiple() const { return set_field<62, 1>(1); } + /** + * @brief choice row id, 0 ~ 63 + */ + constexpr auto row_index(const uint64_t index) const { + return set_field<56, 6>(index); + } +}; + +struct ldstzi : operand { + /** + * @brief interleave load one row to two row. + * right: x[0:16:2] => z[row,8:16], x[1:16:2] => z[row+1,8:16] + * left: x[0:16:2] => z[row,0:8], x[1:16:2] => z[row+1,0:8] + */ + constexpr auto right() const { return set_field<56, 1>(1); } + /** + * @brief choice row id, 0 ~ 63 + */ + constexpr auto row_index(const uint64_t index) const { + return set_field<57, 5>(index); + } +}; + +/** + * @brief + * matrix mode: + * the x, y is one register(64 bytes). + * z[j][z_row][i] += x[i] * y[j] + * note 0 <= z_row < 64 / lane, 0 <= j < 64 / z_row. + * vector mode: z[_][i] += x[i] * y[i] + */ +template struct fma_base : operand { + constexpr auto vector_mode() const { + return this->template set_field<63, 1>(1); + } + /** + * @brief x_mode and x_value is combined options. + * 0: Enable all lanes (x_value==0), or odd lanes only (x_value==1), or even + * lanes only (x_value==2), or no lanes (x_value==anything else) + * + * 1: Only enable lane x_value + * + * 2: Only enable the first x_value lanes, or all lanes when x_value is zero + * + * 3: Only enable the last x_value lanes, or all lanes when x_value is zero + */ + constexpr auto x_mode(uint64_t v) const { + return this->template set_field<46, 2>(v); + } + constexpr auto x_value(uint64_t v) const { + return this->template set_field<41, 5>(v); + } + /** + * @brief same like x_mode and x_value + */ + constexpr auto y_mode(uint64_t v) const { + return this->template set_field<37, 2>(v); + } + /** + * @brief Ignored in vector mode + */ + constexpr auto y_value(uint64_t v) const { + return this->template set_field<32, 5>(v); + } + constexpr auto skip_x() const { return this->template set_field<29, 1>(1); } + constexpr auto skip_y() const { return this->template set_field<28, 1>(1); } + constexpr auto skip_z() const { return this->template set_field<27, 1>(1); } + + /** + * @brief 0 <= z_row < 64 / lane, eg. + * fma64 : 0 <= z_row < 8 (lane = 8) + * fma32 : 0 <= z_row < 4 (lane = 16) + * fma16 : 0 <= z_row < 2 (lane = 32) + */ + constexpr auto z_row(uint64_t v) const { + return this->template set_field<20, 6>(v); + } + /** + * @brief x_offset in bytes + */ + constexpr auto x_offset(uint64_t v) const { + return this->template set_field<10, 9>(v); + } + /** + * @brief y_offset in bytes + */ + constexpr auto y_offset(uint64_t v) const { + return this->template set_field<0, 9>(v); + } +}; + +struct fma16 : fma_base { + constexpr auto f32_z() const { return set_field<62, 1>(1); } +}; + +struct fma32 : fma_base { + constexpr auto f16_x() const { return set_field<61, 1>(1); } + constexpr auto f16_y() const { return set_field<60, 1>(1); } +}; + +struct fma64 : fma_base { + constexpr auto f16_x() const { return set_field<61, 1>(1); } + constexpr auto f16_y() const { return set_field<60, 1>(1); } +}; + +/** + * @brief fms is just doing mul and sub. + */ +template struct fms_base : operand { + constexpr auto vector_mode() const { + return this->template set_field<63, 1>(1); + } + /** + * @brief x_mode and x_value is combined options. + * 0: Enable all lanes (x_value==0), or odd lanes only (x_value==1), or even + * lanes only (x_value==2), or no lanes (x_value==anything else) + * + * 1: Only enable lane x_value + * + * 2: Only enable the first x_value lanes, or all lanes when x_value is zero + * + * 3: Only enable the last x_value lanes, or all lanes when x_value is zero + */ + constexpr auto x_mode(uint64_t v) const { + return this->template set_field<46, 2>(v); + } + constexpr auto x_value(uint64_t v) const { + return this->template set_field<41, 5>(v); + } + /** + * @brief same like x_mode and x_value + */ + constexpr auto y_mode(uint64_t v) const { + return this->template set_field<37, 2>(v); + } + /** + * @brief Ignored in vector mode + */ + constexpr auto y_value(uint64_t v) const { + return this->template set_field<32, 5>(v); + } + constexpr auto skip_x() const { return this->template set_field<29, 1>(1); } + constexpr auto skip_y() const { return this->template set_field<28, 1>(1); } + constexpr auto skip_z() const { return this->template set_field<27, 1>(1); } + + /** + * @brief 0 <= z_row < 64 / lane, eg. + * fma64 : 0 <= z_row < 8 (lane = 8) + * fma32 : 0 <= z_row < 4 (lane = 16) + * fma16 : 0 <= z_row < 2 (lane = 32) + */ + constexpr auto z_row(uint64_t v) const { + return this->template set_field<20, 6>(v); + } + /** + * @brief x_offset in bytes + */ + constexpr auto x_offset(uint64_t v) const { + return this->template set_field<10, 9>(v); + } + /** + * @brief y_offset in bytes + */ + constexpr auto y_offset(uint64_t v) const { + return this->template set_field<0, 9>(v); + } +}; + +struct fms16 : fms_base { + constexpr auto f32_z() const { return set_field<62, 1>(1); } +}; + +struct fms32 : fms_base { + constexpr auto f16_x() const { return set_field<61, 1>(1); } + constexpr auto f16_y() const { return set_field<60, 1>(1); } +}; + +struct fms64 : fms_base { + constexpr auto f16_x() const { return set_field<61, 1>(1); } + constexpr auto f16_y() const { return set_field<60, 1>(1); } +}; + +enum matfp_alu_mode : uint64_t { + /** + * @brief z + x*y + * + */ + kfma = 0, + /** + * @brief z - x*y + */ + kfms = 1, + /** + * @brief x <= 0 ? 0 : y + */ + cond = 4, +}; + +enum matfp_dtype_t : uint64_t { + /** + * @brief xy: bf16 z: bf16 (one row from each two) M2 only + */ + bf16bf16 = 0, + /** + * @brief xy: bf16 z: f32 (all rows, interleaved pairs) M2 only + */ + bf16f32 = 1, + /** + * @brief xy: f16 z: f16 (one row from each two) anything else + */ + f16f16 = 2, + /** + * @brief xy: f16 z: f32 (all rows, interleaved pairs) + */ + f16f32 = 3, + /** + * @brief xy: f32 z: f32 (one row from each four) + */ + f32f32 = 4, + /** + * @brief xy: f64 z: f64 (one row from each eight) + */ + f64f64 = 7, +}; + +enum matfp_enable_mode : uint64_t { + /** + * @brief config mode by enable value. + * 0 : Enable all lanes + * 1 : odd lanes only + * 2 : even lanes only + * 3 : enable all lanes but override the ALU operation to 0.0 + * 4 or 5 : enable all lanes but override their value to 0.0 + * others : no lanes enabled + */ + config = 0, + // Only enable lane #N + kindex = 1, + // Only enable the first N lanes, or all lanes when N is zero + take = 2, + // Only enable the last N lanes, or all lanes when N is zero + take_last = 3, + // Only enable the first N lanes (no lanes when N is zero) + take_or_non = 4, + // Only enable the last N lanes (no lanes when N is zero) + take_last_or_non = 5, + // No lanes enabled + disable = 6, + // No lanes enabled + // disable2 = 7, +}; + +enum shuffle_mode : uint64_t { + /** + * @brief no shuffle. + * + */ + s0 = 0, + /** + * @brief pack half, eg. f32[16] view as f32[2,8] then pack axis 0. + */ + s1 = 0, + /** + * @brief pack quarter, eg. f32[16] view as f32[4,4] then pack axis 0. + */ + s2 = 0, + /** + * @brief pack one eighth, eg. f32[16] view as f32[8,2] then pack axis 0. + */ + s3 = 0, +}; + +/** + * @brief matfp : z[j][i] ยฑ= f(x[i], y[j]) + * support shuffle x/y/z and indexed load. + */ +struct matfp : operand { + + /** + * @brief depend on the y_mode + */ + constexpr auto y_value(uint64_t v) const { return set_field<57, 6>(v); } + + constexpr auto index_load() const { return set_field<53, 1>(1); } + + /** + * @brief only enabled when index_load() == 1. + */ + constexpr auto index_register(uint64_t v) const { + return set_field<49, 3>(v); + } + /** + * @brief indices is 4 bit, otherwise is 2 bit + */ + constexpr auto index_4bit() const { return set_field<48, 1>(1); } + + /** + * @brief load to y, otherwise load to x. + */ + constexpr auto indexed_to_y() const { return set_field<47, 1>(1); } + + /** + * @brief only enabled when index_load() == 0 + */ + constexpr auto alu_mode(matfp_alu_mode v) const { + return set_field<47, 6>(v); + } + + /** + * @brief Lane width mode + */ + constexpr auto dtype_mode(matfp_dtype_t v) const { return set_field<42, 4>(v); } + constexpr auto x_mode(matfp_enable_mode v) const { + return set_field<38, 3>(v); + } + constexpr auto x_value(uint64_t v) const { return set_field<32, 5>(v); } + constexpr auto x_shuffle(shuffle_mode v) const { return set_field<29, 2>(v); } + constexpr auto y_shuffle(shuffle_mode v) const { return set_field<27, 2>(v); } + constexpr auto y_mode(matfp_enable_mode v) const { + return set_field<23, 3>(v); + } + constexpr auto z_row(uint64_t v) const { return set_field<20, 3>(v); } + constexpr auto x_offset(uint64_t v) const { return set_field<10, 9>(v); } + constexpr auto y_offset(uint64_t v) const { return set_field<0, 9>(v); } +}; + +/** + * @brief Data type (DT) Index type (IT) Lane count (LC) Direction mode + */ +enum genlut_mode : uint64_t { + /** @brief + * f32 u4 16 Generate `0` + */ + g_f32_u4_16 = 0, + /** @brief + * bf16 or f16 u5 32 Generate `1` M2 + * required for bf16. See bit 30 + */ + g_f16_u5_32, + /** @brief + * f64 u4 8 Generate `2` High + * bit of index always generated as `0` + */ + g_f64_u4_8, + /** @brief + * i32 u4 16 Generate `3` + */ + g_i32_u4_16, + /** @brief + * i16 u5 32 Generate `4` + */ + g_i16_u5_32, + /** @brief + * u32 u4 16 Generate `5` + */ + g_u32_u4_16, + /** @brief + * u16 u5 32 Generate `6` + */ + g_u16_u5_32, + /** @brief + * any 32-bit u2 16 Lookup `7` + */ + lc_a32_u2_16, + /** @brief + * any 16-bit u2 32 Lookup `8` + */ + lc_a16_u2_32, + /** @brief + * any 8-bit u2 64 Lookup `9` + */ + lc_a8_u2_64, + /** @brief + * any 64-bit u4 8 Lookup `10` High bit + * of index ignored + */ + lc_a64_u4_8, + /** @brief + * any 32-bit u4 16 Lookup `11` + */ + lc_a32_u4_16, + /** @brief + * any 16-bit u4 32 Lookup `12` + */ + lc_a16_u4_32, + /** @brief + * any 8-bit u4 64 Lookup `13` + */ + lc_a8_u4_64, + /** @brief + * any 16-bit u5 32 Lookup `14` + */ + lc_a16_u5_32, + /** @brief + * any 8-bit u5 64 Lookup `15` + */ + lc_a8_u5_64, +}; + +struct genlut : operand { + + /** + * @brief the table register index + */ + constexpr auto register_table(uint64_t v) const { + return set_field<60, 3>(v); + } + + /** + * @brief Table is from Y or from X + * + */ + constexpr auto table_from_y() const { return set_field<59, 1>(1); } + + /** + * @brief Table is from Y or from X + * @table mode table + */ + constexpr auto mode(genlut_mode v) const { return set_field<53, 4>(v); } + + /** + * @brief Data type is bf16 (1) or f16 (0), enabled when mode==1 and m2 only. + */ + constexpr auto bf16() const { return set_field<30, 1>(1); } + + /** + * @brief Destination is Z (1) or is X or Y (0), Ignored in generate modes; + * they always write to X or Y. + */ + constexpr auto dest_z() const { return set_field<26, 1>(1); } + + /** + * @brief enabled when dest_z == 0, Destination is Y (1) or is X (0) + */ + constexpr auto dest_y() const { return set_field<25, 1>(1); } + + /** + * @brief enabled when dest_z ==1, Destination register index (hi) + */ + constexpr auto register_hi(uint64_t v) const { return set_field<23, 3>(v); } + + /** + * @brief Destination register index (lo) + */ + constexpr auto register_lo(uint64_t v) const { return set_field<20, 3>(v); } + + /** + * @brief Source is from Y (1) or from X (0) + */ + constexpr auto src_from_y() const { return set_field<10, 1>(1); } + + /** + * @brief Source offset (in bytes) + */ + constexpr auto src_offset(uint64_t v) const { return set_field<0, 9>(v); } +}; diff --git a/examples/qwen_prefill-amx/native/prefill_ffi.cpp b/examples/qwen_prefill-amx/native/prefill_ffi.cpp new file mode 100644 index 00000000..6f9f3af3 --- /dev/null +++ b/examples/qwen_prefill-amx/native/prefill_ffi.cpp @@ -0,0 +1,38 @@ +// Hand-written AMX projection runtime; independent of llama.cpp. +#include +#include +#include "runtime.hpp" +#include +#include +#include +using tvm::ffi::TensorView; +static std::mutex mutex; +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(); +} +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); +} + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(linear, linear); diff --git a/examples/qwen_prefill-amx/native/runtime.hpp b/examples/qwen_prefill-amx/native/runtime.hpp new file mode 100644 index 00000000..07d699bd --- /dev/null +++ b/examples/qwen_prefill-amx/native/runtime.hpp @@ -0,0 +1,124 @@ +#pragma once +#include "amx.hpp" +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace local_amx { +class Pool { + std::vector threads; + std::mutex mutex; + std::condition_variable ready, done; + std::atomic next{0}; + int count = 0, generation = 0, remaining = 0; + bool stop = false; + std::function task; + + public: + explicit Pool(int n) { + for (int i = 0; i < n; i++) + threads.emplace_back([this] { + int seen = 0; + for (;;) { + std::unique_lock lock(mutex); + ready.wait(lock, + [&] { return stop || generation != seen; }); + if (stop) + return; + seen = generation; + lock.unlock(); + for (int i = next.fetch_add(1); i < count; + i = next.fetch_add(1)) + task(i); + lock.lock(); + if (--remaining == 0) + done.notify_one(); + } + }); + } + ~Pool() { + { + std::lock_guard lock(mutex); + stop = true; + } + ready.notify_all(); + for (auto &thread : threads) + thread.join(); + } + void run(int n, std::function fn) { + std::unique_lock lock(mutex); + task = std::move(fn); + count = n; + next = 0; + remaining = threads.size(); + generation++; + ready.notify_all(); + done.wait(lock, [&] { return remaining == 0; }); + } +}; +inline Pool &pool() { + static Pool value( + std::getenv("TF_AMX_THREADS") + ? std::max(1, std::atoi(std::getenv("TF_AMX_THREADS"))) + : 8); + return value; +} +inline std::vector<_Float16> pack_input(const float *a, int m, int k) { + std::vector<_Float16> out(((m + 31) / 32) * k * 32, 0); + for (int m0 = 0; m0 < m; m0 += 32) + for (int p = 0; p < k; p++) + for (int j = 0; j < std::min(32, m - m0); j++) + out[(m0 / 32 * k + p) * 32 + j] = (_Float16)a[(m0 + j) * k + p]; + return out; +} +} +namespace shared_fp16 { +struct Weight { + int n, k, b; + const _Float16 *data; + Weight(const _Float16 *packed, int nn, int kk, int bb) + : n(nn), k(kk), b(bb), data(packed) {} + const _Float16 *at(int n0, int p) const { + return data + (n0 / b * k + p) * b + n0 % b; + } +}; +inline void prefill(const float *x, int m, const Weight &w, float *out, + int taskcols, bool cyclic = false) { + auto a = local_amx::pack_input(x, m, w.k); + int panels = (w.n + taskcols - 1) / taskcols; + local_amx::pool().run(cyclic ? 8 : panels, [&](int task) { + alignas(128) float z[32]; + AMX_SET(); + for (int panel = task; panel < panels; panel += cyclic ? 8 : panels) + for (int n0 = panel * taskcols; + n0 < std::min(w.n, (panel + 1) * taskcols); n0 += 32) + for (int m0 = 0; m0 < m; m0 += 32) { + for (int p = 0; p < w.k; p++) { + AMX_LDX(ldxy{}.bind((void *)w.at(n0, p))); + AMX_LDY( + ldxy{}.bind(a.data() + (m0 / 32 * w.k + p) * 32)); + auto op = fma16{}.f32_z(); + AMX_FMA16(p == 0 ? op.skip_z() : op); + } + for (int row = 0; row < std::min(32, m - m0); row++) { + AMX_STZ(ldstz{}.row_index(row * 2).bind(z)); + AMX_STZ(ldstz{}.row_index(row * 2 + 1).bind(z + 16)); + for (int j = 0; j < 32; j += 8) { + auto pair = vzipq_f32(vld1q_f32(z + j / 2), + vld1q_f32(z + 16 + j / 2)); + vst1q_f32(out + (m0 + row) * w.n + n0 + j, + pair.val[0]); + vst1q_f32(out + (m0 + row) * w.n + n0 + j + 4, + pair.val[1]); + } + } + } + AMX_CLR(); + }); +} +} diff --git a/examples/qwen_prefill-amx/prepare.py b/examples/qwen_prefill-amx/prepare.py new file mode 100644 index 00000000..cdef06a2 --- /dev/null +++ b/examples/qwen_prefill-amx/prepare.py @@ -0,0 +1,38 @@ +"""Canonical GGUF weights and offline layout conversion via Module.prepare.""" + +from pathlib import Path +import argparse, json +from scheduled import make_scheduled +from config import config, ROOT +from weights import aliases, GGUFResource + + +def prepare(model, checkpoint, output=None, sequence=128): + c = config(model) + program = make_scheduled(c, sequence) + out = Path(output) if output else ROOT / "artifacts" / ("prepared-" + model) + program.prepare(GGUFResource(Path(checkpoint), aliases(c)), str(out), device="cpu") + (out / "manifest.json").write_text( + json.dumps( + dict( + model=model, + sequence=sequence, + checkpoint=str(Path(checkpoint).resolve()), + format="module-list-v1", + ), + indent=2, + ) + + "\n" + ) + print(out, flush=True) + return out + + +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("--output") + p.add_argument("--sequence", type=int, default=128) + a = p.parse_args() + prepare(a.model, a.checkpoint, a.output, a.sequence) diff --git a/examples/qwen_prefill-amx/run.py b/examples/qwen_prefill-amx/run.py new file mode 100644 index 00000000..7d6ffe26 --- /dev/null +++ b/examples/qwen_prefill-amx/run.py @@ -0,0 +1,43 @@ +"""One Prefill smoke run; no warmups, repetitions or performance comparison.""" + +import argparse, json +from pathlib import Path +import torch +from tilefoundry.runtime import SafetensorsResource +from config import config, ROOT +from scheduled import make_scheduled +from runtime_model import make_runtime, inputs + + +def run(model, sequence=128, prepared=None): + torch.set_num_threads(8) + c = config(model) + semantic = make_scheduled(c, sequence) + runtime = make_runtime(semantic, c)() + source = Path(prepared) if prepared else ROOT / "artifacts" / ("prepared-" + model) + runtime.load(SafetensorsResource(str(source), device="cpu")) + tokens = (torch.arange(sequence) * 17 + 100) % min(30000, c["vocab_size"]) + logits = runtime(*inputs(c, sequence, tokens)) + assert tuple(logits.shape) == (1, c["vocab_size"]) and torch.isfinite(logits).all() + report = dict( + model=model, + sequence=sequence, + shape=list(logits.shape), + argmax=int(logits.argmax()), + finite=True, + stats=runtime.stats, + smoke_only=True, + codegen=False, + ) + (source / "smoke.json").write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(report), flush=True) + return report + + +if __name__ == "__main__": + p = argparse.ArgumentParser() + p.add_argument("--model", choices=("1.7b", "4b"), required=True) + p.add_argument("--sequence", type=int, default=128) + p.add_argument("--prepared") + a = p.parse_args() + run(a.model, 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 00000000..858dd834 --- /dev/null +++ b/examples/qwen_prefill-amx/runtime_model.py @@ -0,0 +1,176 @@ +"""One-to-one runtime twins of the whole authored/scheduled model tree. + +AMX linear leaves launch the existing hand-written native kernel through TVM +FFI. All other leaves execute their typed tensor program or a checked GDN +implementation. Hand-written runtime twins follow the HIR function call tree. +There is no llama.cpp whole-model callback in this path and no codegen. +""" + +import os +from pathlib import Path +import torch, tvm_ffi +from tilefoundry.evaluator import evaluate +from tilefoundry.runtime import runtime_func, runtime_module +from tilefoundry.dsl import ModuleList +from tilefoundry.ir.core import Call +from tilefoundry.ir.visitor import collect_exprs +from tilefoundry.ir.hir.schedule import ScheduleOp +from instructions import Fma16 + +R = Path(__file__).resolve().parent + + +def gdn(qkv, z, alpha, beta_raw, conv, decay, dt_bias, gamma, *, nk, nv, dk, dv, eps): + import torch.nn.functional as F + + s = qkv.shape[0] + kd = nk * dk + vd = nv * dv + mixed = ( + F.silu( + F.conv1d( + qkv.T.unsqueeze(0), + conv.unsqueeze(1), + padding=conv.shape[1] - 1, + groups=qkv.shape[1], + )[:, :, :s] + ) + .squeeze(0) + .T + ) + q = mixed[:, :kd].reshape(s, nk, dk).repeat(1, nv // nk, 1) + k = mixed[:, kd : 2 * kd].reshape(s, nk, dk).repeat(1, nv // nk, 1) + q = q * torch.rsqrt((q * q).sum(-1, keepdim=True) + 1e-6) + k = k * torch.rsqrt((k * k).sum(-1, keepdim=True) + 1e-6) + v = mixed[:, 2 * kd :].reshape(s, nv, dv) + g = decay * F.softplus(alpha + dt_bias) + beta = torch.sigmoid(beta_raw) + state = torch.zeros(nv, dk, dv) + out = torch.empty(s, nv, dv) + for t in range(s): + state = state * torch.exp(g[t]).reshape(nv, 1, 1) + key = k[t].reshape(nv, dk, 1) + delta = (v[t] - (state * key).sum(1)) * beta[t].reshape(nv, 1) + state = state + key * delta.reshape(nv, 1, dv) + out[t] = (state * (q[t].reshape(nv, dk, 1) / (dk**0.5))).sum(1) + out = out * torch.rsqrt((out * out).mean(-1, keepdim=True) + eps) * gamma + return (out * F.silu(z.reshape(s, nv, dv))).reshape(s, vd) + + +def make_runtime(semantic, config, *, native=True): + os.environ["TF_AMX_THREADS"] = "8" + ffi = tvm_ffi.load_module(str(R / "artifacts/prefill.dylib")) if native else None + stats = {"amx_calls": 0, "tensor_calls": 0} + + classes = {} + + def twin(mod): + if id(mod) in classes: + return classes[id(mod)] + namespace = { + key: ModuleList(tuple(twin(child) for child in slot.elements)) + if slot.indexed + else twin(slot.elements[0]) + for key, slot in mod.children.items() + } + namespace["stats"] = stats + for fn in mod.functions: + is_amx = ( + len(fn.params) == 2 + and fn.params[1].name == "w" + and len(fn.params[1].type.shape) == 3 + and any( + isinstance(v, Call) + and isinstance(v.target, ScheduleOp) + and isinstance(v.target.op, Fma16) + for v in collect_exprs(fn.body) + ) + ) + + def body(self, *args, fn=fn, is_amx=is_amx): + if fn.name == "prefill": + tokens, cos, sin, mask = args[:4] + hidden = self.embed(tokens) + for child in self.layers: + hidden = child(hidden, cos, sin, mask) + return self.head(self.last(self.norm(hidden))) + if fn.name == "forward" and self.children: + if hasattr(self, "norm_in"): + x, cos, sin, mask = args + x = self.add_mixer( + x, self.mixer(self.norm_in(x), cos, sin, mask) + ) + return self.add_mlp(x, self.mlp(self.norm_ffn(x))) + if hasattr(self, "down"): + x = args[0] + return self.down(self.combine(self.gate(x), self.up(x))) + if hasattr(self, "qkv"): + x = args[0] + return self.output( + self.math( + self.qkv(x), self.gate(x), self.alpha(x), self.beta(x) + ) + ) + if hasattr(self, "query"): + x, cos, sin, mask = args + qg = self.query(x) + q = self.query_rope( + self.query_norm(self.split_query(qg)), cos, sin + ) + k = self.key_rope( + self.key_norm(self.split_kv(self.key(x))), cos, sin + ) + v = self.split_kv(self.value(x)) + out = self.attend(q, k, v, mask) + if config.get("attn_output_gate", False): + out = self.gate_output(out, self.split_gate(qg)) + return self.output(out) + if is_amx: + x, w = args + M, K = x.shape + N = w.shape[0] * w.shape[2] + if native: + stats["amx_calls"] += 1 + y = torch.empty(M, N, dtype=torch.float32) + ffi.linear(x.contiguous(), w, y) + return y + raw = w.permute(0, 2, 1).reshape(N, K) + return x.half().float() @ raw.float().T + if ( + any(p.name == "conv" for p in fn.params) + and len(args) == 8 + and "linear_num_key_heads" in config + ): + return gdn( + *args, + nk=config["linear_num_key_heads"], + nv=config["linear_num_value_heads"], + dk=config["linear_key_head_dim"], + dv=config["linear_value_head_dim"], + eps=config["rms_norm_eps"], + ) + stats["tensor_calls"] += 1 + return evaluate(fn, *args) + + namespace[fn.name] = runtime_func(body) + classes[id(mod)] = runtime_module(mod)( + type(mod.name + "Runtime", (), namespace) + ) + return classes[id(mod)] + + return twin(semantic) + + +def inputs(config, sequence, tokens): + r = int( + config["head_dim"] + * config.get("rope_parameters", {}).get("partial_rotary_factor", 1) + ) + theta = config.get("rope_parameters", {}).get( + "rope_theta", config.get("rope_theta", 1000000) + ) + inv = 1 / (theta ** (torch.arange(0, r, 2, dtype=torch.float32) / r)) + freq = torch.arange(sequence, dtype=torch.float32)[:, None] * inv[None, :] + freq = torch.cat([freq, freq], dim=-1) + mask = torch.full((sequence, sequence), float("-inf")).triu(1) + return tokens.long(), freq.cos(), freq.sin(), mask diff --git a/examples/qwen_prefill-amx/scheduled.py b/examples/qwen_prefill-amx/scheduled.py new file mode 100644 index 00000000..01264af1 --- /dev/null +++ b/examples/qwen_prefill-amx/scheduled.py @@ -0,0 +1,221 @@ +"""Explicit AMX projection schedule, plus typed CPU-kernel fallback choices.""" + +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 instructions import Z_LAYOUT +from tilefoundry.ir.core import Call, Var +from tilefoundry.ir.visitor import collect_exprs +from tilefoundry.ir.hir.schedule import ScheduleOp +from tilefoundry.ir.hir.mesh_region import MeshRegion +from tilefoundry.ir.visitor import ExprCloner, BindingSubstitutionCloner +from host_instructions import instruction_for, INSTRUCTIONS +from tilefoundry.ir.types import Mesh as MeshValue +from model import build + +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 N % 64 == 0 else 32 + 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 == 1: + + @module(entry="forward", topologies=(Topology("core", 8), Topology("amx", 1))) + class Head: + @func + def forward(x: Tensor[(1, 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.amx_host_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.amx_pack( + gmem_layout=Layout((M // BM, K, BM), (K * BM, BM, 1)) + ), + ) + 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) + # Integer coordinate arithmetic stays index arithmetic, not a kernel. + 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): + root = cpu_schedule(build(config, sequence, linear=scheduled_projection)) + return replace(root, topologies=(Topology("core", 8), Topology("amx", 1))) diff --git a/examples/qwen_prefill-amx/verify_ir.py b/examples/qwen_prefill-amx/verify_ir.py new file mode 100644 index 00000000..6003d118 --- /dev/null +++ b/examples/qwen_prefill-amx/verify_ir.py @@ -0,0 +1,34 @@ +"""Reload exported scheduled/TIR sources, without compiling machine code.""" + +import argparse, importlib.util +from config import ROOT +from tilefoundry.ir.core.module import Module +from tilefoundry.ir.tir.prim_function import PrimFunction + + +def verify(model): + path = ROOT / "artifacts" / ("ir-" + model) + for name in ("model", "scheduled", "finalized"): + spec = importlib.util.spec_from_file_location( + "exported_" + name, path / (name + ".py") + ) + loaded = importlib.util.module_from_spec(spec) + spec.loader.exec_module(loaded) + root = loaded.QwenPrefill + assert isinstance(root, Module) + if name == "finalized": + assert isinstance(root.entry_function(), PrimFunction) + print( + "PASS", + model, + name, + "shared layers:", + len({id(m) for m in root.layers}), + flush=True, + ) + + +if __name__ == "__main__": + p = argparse.ArgumentParser() + p.add_argument("--model", choices=("1.7b", "4b"), required=True) + verify(p.parse_args().model) diff --git a/examples/qwen_prefill-amx/weights.py b/examples/qwen_prefill-amx/weights.py new file mode 100644 index 00000000..fdbeead2 --- /dev/null +++ b/examples/qwen_prefill-amx/weights.py @@ -0,0 +1,103 @@ +"""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 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 out + + +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 + ) From bf4d61121bfcd986726d5911fd65cd8f837ce06f Mon Sep 17 00:00:00 2001 From: QiHang Zheng <597323109@qq.com> Date: Mon, 19 Oct 2026 11:05:07 +0800 Subject: [PATCH 4/8] docs(examples): record AMX execution checks and satisfy source hygiene --- examples/qwen_prefill-amx/README.md | 15 +++++++++++++++ src/tilefoundry/analysis/footprint.py | 4 +--- tests/analysis/test_analyze_cross_module.py | 1 - tests/target/test_amx_target.py | 1 - 4 files changed, 16 insertions(+), 5 deletions(-) diff --git a/examples/qwen_prefill-amx/README.md b/examples/qwen_prefill-amx/README.md index 75a50b67..5c0aa7c7 100644 --- a/examples/qwen_prefill-amx/README.md +++ b/examples/qwen_prefill-amx/README.md @@ -94,3 +94,18 @@ resources and native AMX are compared with the unscheduled whole-model evaluator The full-model smoke confirms execution, not generation quality. This migration does not rerun performance experiments and does not claim the old llama.cpp bridge's speedups for this runtime. + +## Execution validation + +Rechecked on M2 Pro on 2026-10-19 using the commands above and existing prepared +Q8_0 resources. Both full-model runs used 128 tokens; all logits were finite. + +| Model | Logits shape | Native AMX calls | First-token argmax | +|---|---|---:|---:| +| Qwen3-1.7B | `[1, 151936]` | 196 | 1813 | +| Qwen3.5-4B | `[1, 248320]` | 248 | 306 | + +The dense/hybrid migration checks also passed against the whole-model HIR +evaluator, with maximum absolute errors of 0 and 5.74e-5 respectively. These +checks establish execution and small-model numerical agreement; the fixed-token +full-model smoke does not establish generation quality or performance. diff --git a/src/tilefoundry/analysis/footprint.py b/src/tilefoundry/analysis/footprint.py index 9390beac..77340591 100644 --- a/src/tilefoundry/analysis/footprint.py +++ b/src/tilefoundry/analysis/footprint.py @@ -1187,11 +1187,9 @@ def resolve_footprint(footprint, ctx): continue def project_image(image): + """Keep instance coordinates correlated with addresses during projection.""" relation = isl.map.from_range(image) for axis in reversed(range(len(item.ownership))): - # Instance coordinates are relation dimensions, not - # external parameters with an unknown SSA binding. Keep - # their correlation with addresses through projection. relation = relation.move_dims(isl.dim_type.IN, 0, isl.dim_type.OUT, axis, 1) source = BufferAlias( item.target, diff --git a/tests/analysis/test_analyze_cross_module.py b/tests/analysis/test_analyze_cross_module.py index ada58c4f..8502d6d9 100644 --- a/tests/analysis/test_analyze_cross_module.py +++ b/tests/analysis/test_analyze_cross_module.py @@ -596,7 +596,6 @@ def project(source): ) actual = result.groups[0][1][0] rank = len(ownership) - # Shift only the address coordinate; instance coordinates remain correlated. offset = isl.map( "{ [" + ",".join(f"i{k}" for k in range(rank + 1)) diff --git a/tests/target/test_amx_target.py b/tests/target/test_amx_target.py index 3438937a..d6e5bd65 100644 --- a/tests/target/test_amx_target.py +++ b/tests/target/test_amx_target.py @@ -128,7 +128,6 @@ def run(): x = tf.zeros(Tensor[(2, 3), "f32", None, "gmem"]) return tf.transpose(x, perm=(1, 0)) - # finalize verifies concrete statements as well as authored HIR admission. result = finalize(AppleMemory) text = as_script(result) assert "T.fill" in text and "T.copy" in text From 3cf660e3d4611f856e0222609ca7f3363a7a7fd7 Mon Sep 17 00:00:00 2001 From: QiHang Zheng <597323109@qq.com> Date: Mon, 19 Oct 2026 11:27:11 +0800 Subject: [PATCH 5/8] refactor(examples): simplify Qwen runtime and move AMX atoms into core --- docs/spec/tir.md | 9 + examples/qwen_prefill-amx/.gitignore | 2 + examples/qwen_prefill-amx/README.md | 135 ++--- examples/qwen_prefill-amx/build.py | 35 -- examples/qwen_prefill-amx/check.py | 92 --- examples/qwen_prefill-amx/config.py | 79 ++- examples/qwen_prefill-amx/configs/1.7b.json | 30 - examples/qwen_prefill-amx/configs/4b.json | 104 ---- examples/qwen_prefill-amx/export_ir.py | 53 -- .../qwen_prefill-amx/host_instructions.py | 214 ------- examples/qwen_prefill-amx/kernels/__init__.py | 60 ++ .../prefill_ffi.cpp => kernels/prefill.cpp} | 2 +- .../{native => kernels}/runtime.hpp | 19 +- examples/qwen_prefill-amx/native/amx.hpp | 560 ------------------ examples/qwen_prefill-amx/prepare.py | 38 -- examples/qwen_prefill-amx/run.py | 49 +- examples/qwen_prefill-amx/runtime_model.py | 332 +++++++---- examples/qwen_prefill-amx/scheduled.py | 107 +--- examples/qwen_prefill-amx/verify_ir.py | 34 -- .../tilefoundry/runtime/amx/instructions.h | 39 ++ pyproject.toml | 1 + src/tilefoundry/analysis/footprint.py | 18 +- src/tilefoundry/ir/tir/amx/__init__.py | 1 + .../tilefoundry/ir/tir/amx}/instructions.py | 38 +- .../passes/transforms/convert_hir_to_tir.py | 7 +- tests/analysis/test_analyze_cross_module.py | 69 +-- tests/schedule/test_attribute_preservation.py | 40 -- tests/target/test_amx_target.py | 21 - 28 files changed, 511 insertions(+), 1677 deletions(-) delete mode 100644 examples/qwen_prefill-amx/build.py delete mode 100644 examples/qwen_prefill-amx/check.py delete mode 100644 examples/qwen_prefill-amx/configs/1.7b.json delete mode 100644 examples/qwen_prefill-amx/configs/4b.json delete mode 100644 examples/qwen_prefill-amx/export_ir.py delete mode 100644 examples/qwen_prefill-amx/host_instructions.py create mode 100644 examples/qwen_prefill-amx/kernels/__init__.py rename examples/qwen_prefill-amx/{native/prefill_ffi.cpp => kernels/prefill.cpp} (95%) rename examples/qwen_prefill-amx/{native => kernels}/runtime.hpp (88%) delete mode 100644 examples/qwen_prefill-amx/native/amx.hpp delete mode 100644 examples/qwen_prefill-amx/prepare.py delete mode 100644 examples/qwen_prefill-amx/verify_ir.py create mode 100644 include/tilefoundry/runtime/amx/instructions.h create mode 100644 src/tilefoundry/ir/tir/amx/__init__.py rename {examples/qwen_prefill-amx => src/tilefoundry/ir/tir/amx}/instructions.py (82%) delete mode 100644 tests/schedule/test_attribute_preservation.py diff --git a/docs/spec/tir.md b/docs/spec/tir.md index 6c76d0b3..9503212c 100644 --- a/docs/spec/tir.md +++ b/docs/spec/tir.md @@ -582,6 +582,15 @@ 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`. + Copy and Fill admit CUDA `thread`, CPU `core`, and Apple `amx` execution scopes through the shared memory execution pattern. A sharded register operand must name the selected scope. Their canonical CUDA `execution_mesh` descriptor diff --git a/examples/qwen_prefill-amx/.gitignore b/examples/qwen_prefill-amx/.gitignore index d048eb85..30399290 100644 --- a/examples/qwen_prefill-amx/.gitignore +++ b/examples/qwen_prefill-amx/.gitignore @@ -1,3 +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 index 5c0aa7c7..1f8eeef5 100644 --- a/examples/qwen_prefill-amx/README.md +++ b/examples/qwen_prefill-amx/README.md @@ -1,111 +1,64 @@ -# Qwen text Prefill on Apple AMX +# Qwen text prefill on Apple AMX -Qwen3-1.7B and Qwen3.5-4B, batch one, text Prefill only. The example provides -whole-model HIR, explicit AMX schedules, retained-call TIR and a hand-written -TVM FFI runtime. No Decode, MTP, vision or compiler code generation is involved. - -The model uses `ModuleList` instance slots: the dense model's 28 layers share -one layer definition; the hybrid model's 32 layers share two. Each instance -retains its own resources. Static calls select the heterogeneous layer slots; -the compiler does not clone or inline their bodies. +Qwen3-1.7B and Qwen3.5-4B text prefill, batch one, on Apple Silicon. The HIR +expresses the whole model; the scheduled HIR assigns its projections to AMX. +A Python runtime module implements the same module tree and calls C++ AMX +kernels through TVM FFI. Other operations run on the CPU using Torch. ## Files -| File | Role | +| File | Purpose | |---|---| -| `model.py` | Dense/hybrid architecture and whole Prefill HIR | -| `scheduled.py` | Explicit destinations, activation/weight layout and worker schedule | -| `instructions.py` | Example-local X/Y load, FMA16, Z store and packing contracts | -| `host_instructions.py` | Typed host software kernels for non-AMX operations | -| `runtime_model.py` | One runtime twin per shared HIR definition | -| `weights.py` | GGUF aliases and Q8_0 decoding | -| `native/` | Standalone hand-written AMX kernel, thread pool and reverse-engineered instruction header | -| `prepare.py`, `run.py`, `export_ir.py`, `check.py` | Reproduction entry points | - -AMX contracts are local extensions in this example, not claims that an AMX -code generator has been added to TileFoundry. Host operations are software -kernels implemented using Torch; they are not AMX ISA instructions. The native -projection kernel and its FFI module have no llama.cpp dependency. - -## Schedule and numerical contract - -Weights arrive as Q8_0 GGUF and are dequantized to FP16. The original quantization -error remains. `Module.prepare` packs each projection from `[N,K]` to `[N/B,K,B]`, -where B is 64, or 32 for narrow projections. Activations are converted to FP16 -and packed into contiguous 32-row panels before accumulation. - -Eight logical workers distribute N panels cyclically. Each holds a 32x32 FP32 -Z accumulator, advances K by loading 32 weights into X and 32 activations into -Y, and issues FP16 outer products with FP32 accumulation. Z stays resident until -the tile is complete, then its interleaved rows are unpacked. The final M=1 -vocabulary projection uses a host kernel; it is not a Decode workload. - -The reverse-engineered AMX header is retained from the original user-provided -runtime. This path is validated on M2 Pro. Worker identity is logical; the -runtime does not promise OS thread affinity. +| `model.py` | Whole-model HIR, shared layer definitions and weight converters | +| `scheduled.py` | AMX projection schedule, packed layouts and worker assignment | +| `runtime_model.py` | Explicit Python runtime modules for projections, attention, GDN and layers | +| `kernels/prefill.cpp`, `kernels/runtime.hpp` | Native projection kernel and FFI entry | +| `kernels/__init__.py` | Build and load the native kernel on first use | +| `weights.py` | Q8_0 GGUF reading and resource aliases | +| `config.py` | Architecture parameters for the two models | +| `run.py` | Weight preparation, loading and one prefill call | ## Run -Use an Apple Silicon Mac with Command Line Tools and the existing `tf` Conda -environment. Dependencies are TileFoundry from this checkout, Torch, `gguf`, -`safetensors` and `tvm_ffi`. Use the same Python for building and running. +Requires an Apple Silicon Mac, Xcode Command Line Tools, Python 3.12+, this +TileFoundry revision, Torch, `gguf`, `safetensors` and `apache-tvm-ffi`. ```sh -conda activate tf -python examples/qwen_prefill-amx/build.py -python examples/qwen_prefill-amx/check.py - -python examples/qwen_prefill-amx/prepare.py --model 1.7b --checkpoint /path/to/Qwen3-1.7B-Q8_0.gguf -python examples/qwen_prefill-amx/run.py --model 1.7b - -python examples/qwen_prefill-amx/prepare.py --model 4b --checkpoint /path/to/Qwen3.5-4B-Q8_0.gguf -python examples/qwen_prefill-amx/run.py --model 4b +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 ``` -Config snapshots are in `configs/`. Checkpoint paths must be supplied explicitly; -there are no user-specific absolute paths or downloads during a run. `prepare` -uses the authored converters and writes canonical weights under -`artifacts/prepared-`. `run` loads only those prepared resources and -executes one 128-token Prefill over a fixed valid token pattern. It checks the -logits shape and finiteness, prints the first-token argmax and actual AMX call -counts, and saves `smoke.json`. It has no warmup, repetition or speedup table. +The first run prepares the GGUF weights under `prepared//` and compiles +the C++ kernel under `kernels/.build/`. Subsequent runs reuse both. Use +`--prepared /path/to/prepared-weights` to choose another resource directory. +`--sequence` selects a positive multiple of 32; the default is 128. The entry +runs a fixed token sequence and prints the logits shape and next-token ID. -Prompt lengths must be positive multiples of 32; `--sequence` selects the static -shape. The prepare/load/run stages are separate, and packed weight layouts do -not depend on prompt length. Large prepared weights and generated binaries are -ignored by Git. +The runtime module can also be called directly with your own token IDs: -## Inspect the compiler path +```python +from config import config +from runtime_model import make_runtime, inputs +from tilefoundry.runtime import SafetensorsResource -```sh -python examples/qwen_prefill-amx/export_ir.py --model 1.7b -python examples/qwen_prefill-amx/export_ir.py --model 4b -python examples/qwen_prefill-amx/verify_ir.py --model 1.7b -python examples/qwen_prefill-amx/verify_ir.py --model 4b +c = config("1.7b") +runtime = make_runtime(c, sequence=128)() +runtime.load(SafetensorsResource("prepared/1.7b", device="cpu")) +logits = runtime(*inputs(c, 128, token_ids)) # token_ids: CPU int64 tensor [128] ``` -The three source files are emitted into `artifacts/ir-/`. Finalization -uses the current API and retained calls, with no old node-budget, recursion-limit -or no-traffic workaround. The runtime is hand-written; emitting TIR does not -compile that TIR into the executed native kernel. - -`check.py` is a small dense/hybrid migration check: actual converters, loaded -resources and native AMX are compared with the unscheduled whole-model evaluator. -The full-model smoke confirms execution, not generation quality. This migration -does not rerun performance experiments and does not claim the old llama.cpp -bridge's speedups for this runtime. - -## Execution validation +Run this snippet from the example directory. `runtime.module` is the scheduled +HIR module; `model.build(c, 128)` constructs its unscheduled counterpart. -Rechecked on M2 Pro on 2026-10-19 using the commands above and existing prepared -Q8_0 resources. Both full-model runs used 128 tokens; all logits were finite. +## Execution arrangement -| Model | Logits shape | Native AMX calls | First-token argmax | -|---|---|---:|---:| -| Qwen3-1.7B | `[1, 151936]` | 196 | 1813 | -| Qwen3.5-4B | `[1, 248320]` | 248 | 306 | +Q8_0 weights are dequantized to FP16 and packed from `[N,K]` to `[N/B,K,B]`, +with B=64, or 32 for narrow projections. Eight logical workers divide N panels +cyclically. Each AMX unit accumulates a 32ร—32 FP32 tile from FP16 outer products; +activations are packed into contiguous 32-row panels. The final one-row vocabulary +projection uses Torch. GDN, normalization, rotary embeddings and attention also +execute in the Python runtime using Torch. -The dense/hybrid migration checks also passed against the whole-model HIR -evaluator, with maximum absolute errors of 0 and 5.74e-5 respectively. These -checks establish execution and small-model numerical agreement; the fixed-token -full-model smoke does not establish generation quality or performance. +AMX instruction contracts and the native instruction wrappers are provided by +TileFoundry. The example supplies the model-specific schedule and hand-written +runtime. It covers text prefill through logits, with no decode, MTP or vision. diff --git a/examples/qwen_prefill-amx/build.py b/examples/qwen_prefill-amx/build.py deleted file mode 100644 index 7a5e9910..00000000 --- a/examples/qwen_prefill-amx/build.py +++ /dev/null @@ -1,35 +0,0 @@ -"""Build the standalone hand-written AMX TVM FFI module.""" - -from pathlib import Path -import platform, subprocess -import tvm_ffi -from tvm_ffi import libinfo - -ROOT = Path(__file__).resolve().parent -if __name__ == "__main__": - if platform.system() != "Darwin" or platform.machine() != "arm64": - raise SystemExit("This AMX runtime requires Apple Silicon macOS") - out = ROOT / "artifacts" - out.mkdir(exist_ok=True) - clang = subprocess.check_output(["xcrun", "--find", "clang++"], text=True).strip() - sdk = subprocess.check_output(["xcrun", "--show-sdk-path"], text=True).strip() - ffi = Path(libinfo.find_libtvm_ffi()) - subprocess.run( - [ - clang, - "-std=c++17", - "-O3", - "-shared", - "-fPIC", - "-isysroot", - sdk, - "-I" + str(Path(tvm_ffi.__file__).parent / "include"), - str(ROOT / "native/prefill_ffi.cpp"), - str(ffi), - "-Wl,-rpath," + str(ffi.parent), - "-o", - str(out / "prefill.dylib"), - ], - check=True, - ) - print(out / "prefill.dylib") diff --git a/examples/qwen_prefill-amx/check.py b/examples/qwen_prefill-amx/check.py deleted file mode 100644 index 4d1e289e..00000000 --- a/examples/qwen_prefill-amx/check.py +++ /dev/null @@ -1,92 +0,0 @@ -"""Tiny migration smoke: prepare/load/native run against the model HIR.""" - -from pathlib import Path -import torch -from tilefoundry.runtime import DictResource, SafetensorsResource -from model import build -from scheduled import make_scheduled -from runtime_model import make_runtime, inputs - -ROOT = Path(__file__).resolve().parent -CONFIG = dict( - hidden_size=64, - intermediate_size=128, - vocab_size=64, - num_attention_heads=4, - num_key_value_heads=2, - head_dim=16, - rms_norm_eps=1e-6, - num_hidden_layers=2, - rope_theta=10000, -) - - -def raw_weights(mod, prefix="", out=None): - out = {} if out is None else out - for fn in mod.functions: - converters = dict(fn.converters) - for param in fn.params: - if not param.is_const: - continue - for p in ( - converters[param.name].params if param.name in converters else (param,) - ): - key = prefix + p.name - if key in out: - continue - x = torch.randn(p.type.shape) * 0.025 - if p.name.startswith("gamma"): - x = 1 + x - if p.name == "decay": - x = -x.abs() - 0.01 - out[key] = x.half() if p.type.dtype.name == "f16" else x.float() - for key, slot in mod.children.items(): - for i, ch in enumerate(slot.elements): - raw_weights( - ch, prefix + key + "." + (str(i) + "." if slot.indexed else ""), out - ) - return out - - -if __name__ == "__main__": - torch.set_num_threads(8) - torch.manual_seed(17) - for name in ("dense", "hybrid"): - c = dict(CONFIG) - if name == "hybrid": - c.update( - attn_output_gate=True, - layer_types=["linear_attention", "full_attention"], - linear_num_key_heads=16, - linear_num_value_heads=32, - linear_key_head_dim=4, - linear_value_head_dim=4, - linear_conv_kernel_dim=4, - rope_parameters={"partial_rotary_factor": 0.5, "rope_theta": 10000}, - ) - plain = build(c, 32) - scheduled = make_scheduled(c, 32) - raw = DictResource(raw_weights(plain)) - outputs = [] - args = inputs(c, 32, torch.arange(32) % 64) - for label, sem in [("plain", plain), ("scheduled", scheduled)]: - p = ROOT / "artifacts" / ("migration-check-" + name + "-" + label) - sem.prepare(raw, str(p), device="cpu") - resource = SafetensorsResource(str(p), device="cpu") - if label == "plain": - result = sem.load(resource)(*args) - else: - runtime = make_runtime(sem, c)() - runtime.load(resource) - result = runtime(*args) - assert runtime.stats["amx_calls"] > 0 - assert torch.isfinite(result).all() - outputs.append(result) - torch.testing.assert_close(outputs[1], outputs[0], rtol=2e-4, atol=2e-4) - print( - "PASS", - name, - "native runtime vs whole-model HIR", - float((outputs[1] - outputs[0]).abs().max()), - flush=True, - ) diff --git a/examples/qwen_prefill-amx/config.py b/examples/qwen_prefill-amx/config.py index 9d40686c..92a79866 100644 --- a/examples/qwen_prefill-amx/config.py +++ b/examples/qwen_prefill-amx/config.py @@ -1,11 +1,76 @@ -"""Model config snapshots and local resource paths.""" +"""Architecture parameters for the two Qwen text prefill examples.""" -import json -from pathlib import Path - -ROOT = Path(__file__).resolve().parent +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): - data = json.loads((ROOT / "configs" / f"{model}.json").read_text()) - return data.get("text_config", data) + return dict(CONFIGS[model]) diff --git a/examples/qwen_prefill-amx/configs/1.7b.json b/examples/qwen_prefill-amx/configs/1.7b.json deleted file mode 100644 index 044a86ec..00000000 --- a/examples/qwen_prefill-amx/configs/1.7b.json +++ /dev/null @@ -1,30 +0,0 @@ -{ - "architectures": [ - "Qwen3ForCausalLM" - ], - "attention_bias": false, - "attention_dropout": 0.0, - "bos_token_id": 151643, - "eos_token_id": 151645, - "head_dim": 128, - "hidden_act": "silu", - "hidden_size": 2048, - "initializer_range": 0.02, - "intermediate_size": 6144, - "max_position_embeddings": 40960, - "max_window_layers": 28, - "model_type": "qwen3", - "num_attention_heads": 16, - "num_hidden_layers": 28, - "num_key_value_heads": 8, - "rms_norm_eps": 1e-06, - "rope_scaling": null, - "rope_theta": 1000000, - "sliding_window": null, - "tie_word_embeddings": true, - "torch_dtype": "bfloat16", - "transformers_version": "4.51.0", - "use_cache": true, - "use_sliding_window": false, - "vocab_size": 151936 -} \ No newline at end of file diff --git a/examples/qwen_prefill-amx/configs/4b.json b/examples/qwen_prefill-amx/configs/4b.json deleted file mode 100644 index 557d961b..00000000 --- a/examples/qwen_prefill-amx/configs/4b.json +++ /dev/null @@ -1,104 +0,0 @@ -{ - "architectures": [ - "Qwen3_5ForConditionalGeneration" - ], - "image_token_id": 248056, - "model_type": "qwen3_5", - "text_config": { - "attention_bias": false, - "attention_dropout": 0.0, - "attn_output_gate": true, - "dtype": "bfloat16", - "eos_token_id": 248044, - "full_attention_interval": 4, - "head_dim": 256, - "hidden_act": "silu", - "hidden_size": 2560, - "initializer_range": 0.02, - "intermediate_size": 9216, - "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_conv_kernel_dim": 4, - "linear_key_head_dim": 128, - "linear_num_key_heads": 16, - "linear_num_value_heads": 32, - "linear_value_head_dim": 128, - "max_position_embeddings": 262144, - "mlp_only_layers": [], - "model_type": "qwen3_5_text", - "mtp_num_hidden_layers": 1, - "mtp_use_dedicated_embeddings": false, - "num_attention_heads": 16, - "num_hidden_layers": 32, - "num_key_value_heads": 4, - "rms_norm_eps": 1e-06, - "tie_word_embeddings": true, - "use_cache": true, - "vocab_size": 248320, - "mamba_ssm_dtype": "float32", - "rope_parameters": { - "mrope_interleaved": true, - "mrope_section": [ - 11, - 11, - 10 - ], - "rope_type": "default", - "rope_theta": 10000000, - "partial_rotary_factor": 0.25 - } - }, - "tie_word_embeddings": true, - "transformers_version": "4.57.0.dev0", - "video_token_id": 248057, - "vision_config": { - "deepstack_visual_indexes": [], - "depth": 24, - "hidden_act": "gelu_pytorch_tanh", - "hidden_size": 1024, - "in_channels": 3, - "initializer_range": 0.02, - "intermediate_size": 4096, - "model_type": "qwen3_5", - "num_heads": 16, - "num_position_embeddings": 2304, - "out_hidden_size": 2560, - "patch_size": 16, - "spatial_merge_size": 2, - "temporal_patch_size": 2 - }, - "vision_end_token_id": 248054, - "vision_start_token_id": 248053 -} \ No newline at end of file diff --git a/examples/qwen_prefill-amx/export_ir.py b/examples/qwen_prefill-amx/export_ir.py deleted file mode 100644 index 8bef2b68..00000000 --- a/examples/qwen_prefill-amx/export_ir.py +++ /dev/null @@ -1,53 +0,0 @@ -"""Export whole-model HIR, scheduled HIR and retained-call TIR.""" - -import argparse, json -from pathlib import Path -from tilefoundry.inspection import as_script -from tilefoundry.schedule import finalize -from tilefoundry.analysis.check import check_program -from config import config, ROOT -from model import build -from scheduled import make_scheduled - - -def export(model, sequence=128, output=None): - out = Path(output) if output else ROOT / "artifacts" / ("ir-" + model) - out.mkdir(parents=True, exist_ok=True) - c = config(model) - plain = build(c, sequence) - plain, _ = check_program(plain, plain.entry_function()) - scheduled = make_scheduled(c, sequence) - - def registrations(source): - future = "from __future__ import annotations\n" - imports = "import instructions, host_instructions\n" - return ( - source.replace(future, future + imports, 1) - if future in source - else imports + source - ) - - (out / "model.py").write_text(as_script(plain)) - (out / "scheduled.py").write_text(registrations(as_script(scheduled))) - lowered = finalize(scheduled) - (out / "finalized.py").write_text(registrations(as_script(lowered))) - report = dict( - model=model, - sequence=sequence, - layers=len(scheduled.layers), - shared_layer_definitions=len({id(layer) for layer in scheduled.layers}), - codegen=False, - performance_measured=False, - ) - (out / "manifest.json").write_text(json.dumps(report, indent=2) + "\n") - print(json.dumps(report), flush=True) - return out - - -if __name__ == "__main__": - p = argparse.ArgumentParser() - p.add_argument("--model", choices=("1.7b", "4b"), required=True) - p.add_argument("--sequence", type=int, default=128) - p.add_argument("--output") - a = p.parse_args() - export(a.model, a.sequence, a.output) diff --git a/examples/qwen_prefill-amx/host_instructions.py b/examples/qwen_prefill-amx/host_instructions.py deleted file mode 100644 index 2a1c3d6e..00000000 --- a/examples/qwen_prefill-amx/host_instructions.py +++ /dev/null @@ -1,214 +0,0 @@ -"""Typed CPU tensor kernels, not ISA atoms or an opaque whole-model callback. - -Each kernel preserves its HIR operation's attributes and coordinate relations. -The schedule chooses these kernels explicitly for the CPU fallback portions of -an AMX program. No CPU code generation is provided by this module. -""" - -import isl - - -from dataclasses import replace -from enum import Enum - -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.transpose import Transpose -from tilefoundry.ir.hir.tensor.reshape import Reshape -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.where import Where -from tilefoundry.ir.pattern import TensorPattern -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 TypeInferContext -from tilefoundry.visitor_registry.registries import typeinfer_registry - -from instructions import CORE_PATTERN - -CORE = CORE_PATTERN -INSTRUCTIONS = {} - - -def _semantic_type(ty): - # A host tensor-kernel call can pass a scalar register by value. The HIR - # rule treats scalar literals as uniform rather than tensor residency. - 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(None), - "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 = "amx_host_" + 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(): - # A software row kernel writes each result cell exactly once. - # Expose output coordinates in the logical issue domain so the - # explicit destination/result relation is the identity, while - # each result may still read its entire normalization row. - 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.frame.target if ctx.frame is not None else None - 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_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), - (Where, 3), - (Arange, 0), - (Concat, 2), - (Concat, 3), -): - _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/examples/qwen_prefill-amx/kernels/__init__.py b/examples/qwen_prefill-amx/kernels/__init__.py new file mode 100644 index 00000000..8d8224ff --- /dev/null +++ b/examples/qwen_prefill-amx/kernels/__init__.py @@ -0,0 +1,60 @@ +"""Build and cache the Apple Silicon prefill extension on first use.""" + +from pathlib import Path +import platform +import subprocess +import tvm_ffi +from tvm_ffi import libinfo +import sys +import tilefoundry + +_module = None + + +def ops(): + global _module + if _module is None: + if platform.system() != "Darwin" or platform.machine() != "arm64": + raise RuntimeError("The AMX example requires Apple Silicon macOS") + root = Path(__file__).resolve().parent + include = Path(sys.prefix) / "include" + if not (include / "tilefoundry/runtime/amx/instructions.h").exists(): + include = Path(tilefoundry.__file__).resolve().parents[2] / "include" + sources = ( + root / "prefill.cpp", + root / "runtime.hpp", + include / "tilefoundry/runtime/amx/instructions.h", + ) + out = root / ".build/prefill.dylib" + ffi = Path(libinfo.find_libtvm_ffi()) + if not out.exists() or any( + p.stat().st_mtime > out.stat().st_mtime for p in sources + ): + out.parent.mkdir(exist_ok=True) + 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, + "-I" + str(Path(tvm_ffi.__file__).parent / "include"), + "-I" + str(include), + str(root / "prefill.cpp"), + str(ffi), + "-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/native/prefill_ffi.cpp b/examples/qwen_prefill-amx/kernels/prefill.cpp similarity index 95% rename from examples/qwen_prefill-amx/native/prefill_ffi.cpp rename to examples/qwen_prefill-amx/kernels/prefill.cpp index 6f9f3af3..545d1def 100644 --- a/examples/qwen_prefill-amx/native/prefill_ffi.cpp +++ b/examples/qwen_prefill-amx/kernels/prefill.cpp @@ -1,4 +1,4 @@ -// Hand-written AMX projection runtime; independent of llama.cpp. +/** Qwen prefill projection and TVM FFI entry. */ #include #include #include "runtime.hpp" diff --git a/examples/qwen_prefill-amx/native/runtime.hpp b/examples/qwen_prefill-amx/kernels/runtime.hpp similarity index 88% rename from examples/qwen_prefill-amx/native/runtime.hpp rename to examples/qwen_prefill-amx/kernels/runtime.hpp index 07d699bd..f4460941 100644 --- a/examples/qwen_prefill-amx/native/runtime.hpp +++ b/examples/qwen_prefill-amx/kernels/runtime.hpp @@ -1,5 +1,5 @@ #pragma once -#include "amx.hpp" +#include #include #include #include @@ -93,21 +93,20 @@ inline void prefill(const float *x, int m, const Weight &w, float *out, int panels = (w.n + taskcols - 1) / taskcols; local_amx::pool().run(cyclic ? 8 : panels, [&](int task) { alignas(128) float z[32]; - AMX_SET(); + tilefoundry::amx::start(); for (int panel = task; panel < panels; panel += cyclic ? 8 : panels) for (int n0 = panel * taskcols; n0 < std::min(w.n, (panel + 1) * taskcols); n0 += 32) for (int m0 = 0; m0 < m; m0 += 32) { for (int p = 0; p < w.k; p++) { - AMX_LDX(ldxy{}.bind((void *)w.at(n0, p))); - AMX_LDY( - ldxy{}.bind(a.data() + (m0 / 32 * w.k + p) * 32)); - auto op = fma16{}.f32_z(); - AMX_FMA16(p == 0 ? op.skip_z() : op); + tilefoundry::amx::load_x(w.at(n0, p)); + tilefoundry::amx::load_y(a.data() + + (m0 / 32 * w.k + p) * 32); + tilefoundry::amx::fma16(p == 0); } for (int row = 0; row < std::min(32, m - m0); row++) { - AMX_STZ(ldstz{}.row_index(row * 2).bind(z)); - AMX_STZ(ldstz{}.row_index(row * 2 + 1).bind(z + 16)); + tilefoundry::amx::store_z(z, row * 2); + tilefoundry::amx::store_z(z + 16, row * 2 + 1); for (int j = 0; j < 32; j += 8) { auto pair = vzipq_f32(vld1q_f32(z + j / 2), vld1q_f32(z + 16 + j / 2)); @@ -118,7 +117,7 @@ inline void prefill(const float *x, int m, const Weight &w, float *out, } } } - AMX_CLR(); + tilefoundry::amx::stop(); }); } } diff --git a/examples/qwen_prefill-amx/native/amx.hpp b/examples/qwen_prefill-amx/native/amx.hpp deleted file mode 100644 index 9fc36d26..00000000 --- a/examples/qwen_prefill-amx/native/amx.hpp +++ /dev/null @@ -1,560 +0,0 @@ -#pragma once -#include - -#ifdef __cplusplus -extern "C" { -#endif - -#define AMX_NOP_OP_IMM5(op, imm5) \ - __asm("nop\nnop\nnop\n.word (0x201000 + (%0 << 5) + %1)" \ - : \ - : "i"(op), "i"(imm5) \ - : "memory") - -#define AMX_OP_GPR(op, gpr) \ - __asm(".word (0x201000 + (%0 << 5) + 0%1 - ((0%1 >> 4) * 6))" \ - : \ - : "i"(op), "r"((uint64_t)(gpr)) \ - : "memory") - -#define AMX_LDX(gpr) AMX_OP_GPR(0, gpr) -#define AMX_LDY(gpr) AMX_OP_GPR(1, gpr) -#define AMX_STX(gpr) AMX_OP_GPR(2, gpr) -#define AMX_STY(gpr) AMX_OP_GPR(3, gpr) -#define AMX_LDZ(gpr) AMX_OP_GPR(4, gpr) -#define AMX_STZ(gpr) AMX_OP_GPR(5, gpr) -#define AMX_LDZI(gpr) AMX_OP_GPR(6, gpr) -#define AMX_STZI(gpr) AMX_OP_GPR(7, gpr) -#define AMX_EXTRX(gpr) AMX_OP_GPR(8, gpr) -#define AMX_EXTRY(gpr) AMX_OP_GPR(9, gpr) -#define AMX_FMA64(gpr) AMX_OP_GPR(10, gpr) -#define AMX_FMS64(gpr) AMX_OP_GPR(11, gpr) -#define AMX_FMA32(gpr) AMX_OP_GPR(12, gpr) -#define AMX_FMS32(gpr) AMX_OP_GPR(13, gpr) -#define AMX_MAC16(gpr) AMX_OP_GPR(14, gpr) -#define AMX_FMA16(gpr) AMX_OP_GPR(15, gpr) -#define AMX_FMS16(gpr) AMX_OP_GPR(16, gpr) -#define AMX_SET() AMX_NOP_OP_IMM5(17, 0) -#define AMX_CLR() AMX_NOP_OP_IMM5(17, 1) -#define AMX_VECINT(gpr) AMX_OP_GPR(18, gpr) -#define AMX_VECFP(gpr) AMX_OP_GPR(19, gpr) -#define AMX_MATINT(gpr) AMX_OP_GPR(20, gpr) -#define AMX_MATFP(gpr) AMX_OP_GPR(21, gpr) -#define AMX_GENLUT(gpr) AMX_OP_GPR(22, gpr) - -#ifdef __cplusplus -} -#endif - -template struct operand { - inline constexpr uint64_t bind(void *ptr) const noexcept { - return val | (((1ull << 57) - 1) & (uint64_t)ptr); - } - - inline constexpr operator uint64_t() const noexcept { return val; } - - template - inline constexpr T set_field(const uint64_t value) const noexcept { - constexpr uint64_t clr_mask = ~(((1ull << Width) - 1ull) << OffSet); - uint64_t value_mask = (value << OffSet); - return T{(val & clr_mask) | value_mask}; - } - - uint64_t val; -}; - -struct ldxy : operand { - /** - * @brief load multiple register - * - * factor = multiple ? 2 : 1 - * - * factor2 = multiple_means_four ? 2 : 1 - * - * so total load factor * factor2 * 64 byte - */ - constexpr auto multiple() const { return set_field<62, 1>(1); } - constexpr auto non_consecutive() const { return set_field<61, 1>(1); } - constexpr auto multiple_four() const { return set_field<60, 1>(1); } - /** - * @brief choice reg id, 0~5 - */ - constexpr auto register_index(const uint64_t index) const { - return set_field<56, 3>(index); - } -}; - -struct stxy : operand { - /** - * @brief store (multiple ? 2 : 1) * 64 bytes. - */ - constexpr auto multiple() const { return set_field<62, 1>(1); } - /** - * @brief choice reg id, 0~5 - */ - constexpr auto register_index(const uint64_t index) const { - return set_field<56, 3>(index); - } -}; - -struct ldstz : operand { - /** - * @brief load or store (multiple ? 2 : 1) * register. - */ - constexpr auto multiple() const { return set_field<62, 1>(1); } - /** - * @brief choice row id, 0 ~ 63 - */ - constexpr auto row_index(const uint64_t index) const { - return set_field<56, 6>(index); - } -}; - -struct ldstzi : operand { - /** - * @brief interleave load one row to two row. - * right: x[0:16:2] => z[row,8:16], x[1:16:2] => z[row+1,8:16] - * left: x[0:16:2] => z[row,0:8], x[1:16:2] => z[row+1,0:8] - */ - constexpr auto right() const { return set_field<56, 1>(1); } - /** - * @brief choice row id, 0 ~ 63 - */ - constexpr auto row_index(const uint64_t index) const { - return set_field<57, 5>(index); - } -}; - -/** - * @brief - * matrix mode: - * the x, y is one register(64 bytes). - * z[j][z_row][i] += x[i] * y[j] - * note 0 <= z_row < 64 / lane, 0 <= j < 64 / z_row. - * vector mode: z[_][i] += x[i] * y[i] - */ -template struct fma_base : operand { - constexpr auto vector_mode() const { - return this->template set_field<63, 1>(1); - } - /** - * @brief x_mode and x_value is combined options. - * 0: Enable all lanes (x_value==0), or odd lanes only (x_value==1), or even - * lanes only (x_value==2), or no lanes (x_value==anything else) - * - * 1: Only enable lane x_value - * - * 2: Only enable the first x_value lanes, or all lanes when x_value is zero - * - * 3: Only enable the last x_value lanes, or all lanes when x_value is zero - */ - constexpr auto x_mode(uint64_t v) const { - return this->template set_field<46, 2>(v); - } - constexpr auto x_value(uint64_t v) const { - return this->template set_field<41, 5>(v); - } - /** - * @brief same like x_mode and x_value - */ - constexpr auto y_mode(uint64_t v) const { - return this->template set_field<37, 2>(v); - } - /** - * @brief Ignored in vector mode - */ - constexpr auto y_value(uint64_t v) const { - return this->template set_field<32, 5>(v); - } - constexpr auto skip_x() const { return this->template set_field<29, 1>(1); } - constexpr auto skip_y() const { return this->template set_field<28, 1>(1); } - constexpr auto skip_z() const { return this->template set_field<27, 1>(1); } - - /** - * @brief 0 <= z_row < 64 / lane, eg. - * fma64 : 0 <= z_row < 8 (lane = 8) - * fma32 : 0 <= z_row < 4 (lane = 16) - * fma16 : 0 <= z_row < 2 (lane = 32) - */ - constexpr auto z_row(uint64_t v) const { - return this->template set_field<20, 6>(v); - } - /** - * @brief x_offset in bytes - */ - constexpr auto x_offset(uint64_t v) const { - return this->template set_field<10, 9>(v); - } - /** - * @brief y_offset in bytes - */ - constexpr auto y_offset(uint64_t v) const { - return this->template set_field<0, 9>(v); - } -}; - -struct fma16 : fma_base { - constexpr auto f32_z() const { return set_field<62, 1>(1); } -}; - -struct fma32 : fma_base { - constexpr auto f16_x() const { return set_field<61, 1>(1); } - constexpr auto f16_y() const { return set_field<60, 1>(1); } -}; - -struct fma64 : fma_base { - constexpr auto f16_x() const { return set_field<61, 1>(1); } - constexpr auto f16_y() const { return set_field<60, 1>(1); } -}; - -/** - * @brief fms is just doing mul and sub. - */ -template struct fms_base : operand { - constexpr auto vector_mode() const { - return this->template set_field<63, 1>(1); - } - /** - * @brief x_mode and x_value is combined options. - * 0: Enable all lanes (x_value==0), or odd lanes only (x_value==1), or even - * lanes only (x_value==2), or no lanes (x_value==anything else) - * - * 1: Only enable lane x_value - * - * 2: Only enable the first x_value lanes, or all lanes when x_value is zero - * - * 3: Only enable the last x_value lanes, or all lanes when x_value is zero - */ - constexpr auto x_mode(uint64_t v) const { - return this->template set_field<46, 2>(v); - } - constexpr auto x_value(uint64_t v) const { - return this->template set_field<41, 5>(v); - } - /** - * @brief same like x_mode and x_value - */ - constexpr auto y_mode(uint64_t v) const { - return this->template set_field<37, 2>(v); - } - /** - * @brief Ignored in vector mode - */ - constexpr auto y_value(uint64_t v) const { - return this->template set_field<32, 5>(v); - } - constexpr auto skip_x() const { return this->template set_field<29, 1>(1); } - constexpr auto skip_y() const { return this->template set_field<28, 1>(1); } - constexpr auto skip_z() const { return this->template set_field<27, 1>(1); } - - /** - * @brief 0 <= z_row < 64 / lane, eg. - * fma64 : 0 <= z_row < 8 (lane = 8) - * fma32 : 0 <= z_row < 4 (lane = 16) - * fma16 : 0 <= z_row < 2 (lane = 32) - */ - constexpr auto z_row(uint64_t v) const { - return this->template set_field<20, 6>(v); - } - /** - * @brief x_offset in bytes - */ - constexpr auto x_offset(uint64_t v) const { - return this->template set_field<10, 9>(v); - } - /** - * @brief y_offset in bytes - */ - constexpr auto y_offset(uint64_t v) const { - return this->template set_field<0, 9>(v); - } -}; - -struct fms16 : fms_base { - constexpr auto f32_z() const { return set_field<62, 1>(1); } -}; - -struct fms32 : fms_base { - constexpr auto f16_x() const { return set_field<61, 1>(1); } - constexpr auto f16_y() const { return set_field<60, 1>(1); } -}; - -struct fms64 : fms_base { - constexpr auto f16_x() const { return set_field<61, 1>(1); } - constexpr auto f16_y() const { return set_field<60, 1>(1); } -}; - -enum matfp_alu_mode : uint64_t { - /** - * @brief z + x*y - * - */ - kfma = 0, - /** - * @brief z - x*y - */ - kfms = 1, - /** - * @brief x <= 0 ? 0 : y - */ - cond = 4, -}; - -enum matfp_dtype_t : uint64_t { - /** - * @brief xy: bf16 z: bf16 (one row from each two) M2 only - */ - bf16bf16 = 0, - /** - * @brief xy: bf16 z: f32 (all rows, interleaved pairs) M2 only - */ - bf16f32 = 1, - /** - * @brief xy: f16 z: f16 (one row from each two) anything else - */ - f16f16 = 2, - /** - * @brief xy: f16 z: f32 (all rows, interleaved pairs) - */ - f16f32 = 3, - /** - * @brief xy: f32 z: f32 (one row from each four) - */ - f32f32 = 4, - /** - * @brief xy: f64 z: f64 (one row from each eight) - */ - f64f64 = 7, -}; - -enum matfp_enable_mode : uint64_t { - /** - * @brief config mode by enable value. - * 0 : Enable all lanes - * 1 : odd lanes only - * 2 : even lanes only - * 3 : enable all lanes but override the ALU operation to 0.0 - * 4 or 5 : enable all lanes but override their value to 0.0 - * others : no lanes enabled - */ - config = 0, - // Only enable lane #N - kindex = 1, - // Only enable the first N lanes, or all lanes when N is zero - take = 2, - // Only enable the last N lanes, or all lanes when N is zero - take_last = 3, - // Only enable the first N lanes (no lanes when N is zero) - take_or_non = 4, - // Only enable the last N lanes (no lanes when N is zero) - take_last_or_non = 5, - // No lanes enabled - disable = 6, - // No lanes enabled - // disable2 = 7, -}; - -enum shuffle_mode : uint64_t { - /** - * @brief no shuffle. - * - */ - s0 = 0, - /** - * @brief pack half, eg. f32[16] view as f32[2,8] then pack axis 0. - */ - s1 = 0, - /** - * @brief pack quarter, eg. f32[16] view as f32[4,4] then pack axis 0. - */ - s2 = 0, - /** - * @brief pack one eighth, eg. f32[16] view as f32[8,2] then pack axis 0. - */ - s3 = 0, -}; - -/** - * @brief matfp : z[j][i] ยฑ= f(x[i], y[j]) - * support shuffle x/y/z and indexed load. - */ -struct matfp : operand { - - /** - * @brief depend on the y_mode - */ - constexpr auto y_value(uint64_t v) const { return set_field<57, 6>(v); } - - constexpr auto index_load() const { return set_field<53, 1>(1); } - - /** - * @brief only enabled when index_load() == 1. - */ - constexpr auto index_register(uint64_t v) const { - return set_field<49, 3>(v); - } - /** - * @brief indices is 4 bit, otherwise is 2 bit - */ - constexpr auto index_4bit() const { return set_field<48, 1>(1); } - - /** - * @brief load to y, otherwise load to x. - */ - constexpr auto indexed_to_y() const { return set_field<47, 1>(1); } - - /** - * @brief only enabled when index_load() == 0 - */ - constexpr auto alu_mode(matfp_alu_mode v) const { - return set_field<47, 6>(v); - } - - /** - * @brief Lane width mode - */ - constexpr auto dtype_mode(matfp_dtype_t v) const { return set_field<42, 4>(v); } - constexpr auto x_mode(matfp_enable_mode v) const { - return set_field<38, 3>(v); - } - constexpr auto x_value(uint64_t v) const { return set_field<32, 5>(v); } - constexpr auto x_shuffle(shuffle_mode v) const { return set_field<29, 2>(v); } - constexpr auto y_shuffle(shuffle_mode v) const { return set_field<27, 2>(v); } - constexpr auto y_mode(matfp_enable_mode v) const { - return set_field<23, 3>(v); - } - constexpr auto z_row(uint64_t v) const { return set_field<20, 3>(v); } - constexpr auto x_offset(uint64_t v) const { return set_field<10, 9>(v); } - constexpr auto y_offset(uint64_t v) const { return set_field<0, 9>(v); } -}; - -/** - * @brief Data type (DT) Index type (IT) Lane count (LC) Direction mode - */ -enum genlut_mode : uint64_t { - /** @brief - * f32 u4 16 Generate `0` - */ - g_f32_u4_16 = 0, - /** @brief - * bf16 or f16 u5 32 Generate `1` M2 - * required for bf16. See bit 30 - */ - g_f16_u5_32, - /** @brief - * f64 u4 8 Generate `2` High - * bit of index always generated as `0` - */ - g_f64_u4_8, - /** @brief - * i32 u4 16 Generate `3` - */ - g_i32_u4_16, - /** @brief - * i16 u5 32 Generate `4` - */ - g_i16_u5_32, - /** @brief - * u32 u4 16 Generate `5` - */ - g_u32_u4_16, - /** @brief - * u16 u5 32 Generate `6` - */ - g_u16_u5_32, - /** @brief - * any 32-bit u2 16 Lookup `7` - */ - lc_a32_u2_16, - /** @brief - * any 16-bit u2 32 Lookup `8` - */ - lc_a16_u2_32, - /** @brief - * any 8-bit u2 64 Lookup `9` - */ - lc_a8_u2_64, - /** @brief - * any 64-bit u4 8 Lookup `10` High bit - * of index ignored - */ - lc_a64_u4_8, - /** @brief - * any 32-bit u4 16 Lookup `11` - */ - lc_a32_u4_16, - /** @brief - * any 16-bit u4 32 Lookup `12` - */ - lc_a16_u4_32, - /** @brief - * any 8-bit u4 64 Lookup `13` - */ - lc_a8_u4_64, - /** @brief - * any 16-bit u5 32 Lookup `14` - */ - lc_a16_u5_32, - /** @brief - * any 8-bit u5 64 Lookup `15` - */ - lc_a8_u5_64, -}; - -struct genlut : operand { - - /** - * @brief the table register index - */ - constexpr auto register_table(uint64_t v) const { - return set_field<60, 3>(v); - } - - /** - * @brief Table is from Y or from X - * - */ - constexpr auto table_from_y() const { return set_field<59, 1>(1); } - - /** - * @brief Table is from Y or from X - * @table mode table - */ - constexpr auto mode(genlut_mode v) const { return set_field<53, 4>(v); } - - /** - * @brief Data type is bf16 (1) or f16 (0), enabled when mode==1 and m2 only. - */ - constexpr auto bf16() const { return set_field<30, 1>(1); } - - /** - * @brief Destination is Z (1) or is X or Y (0), Ignored in generate modes; - * they always write to X or Y. - */ - constexpr auto dest_z() const { return set_field<26, 1>(1); } - - /** - * @brief enabled when dest_z == 0, Destination is Y (1) or is X (0) - */ - constexpr auto dest_y() const { return set_field<25, 1>(1); } - - /** - * @brief enabled when dest_z ==1, Destination register index (hi) - */ - constexpr auto register_hi(uint64_t v) const { return set_field<23, 3>(v); } - - /** - * @brief Destination register index (lo) - */ - constexpr auto register_lo(uint64_t v) const { return set_field<20, 3>(v); } - - /** - * @brief Source is from Y (1) or from X (0) - */ - constexpr auto src_from_y() const { return set_field<10, 1>(1); } - - /** - * @brief Source offset (in bytes) - */ - constexpr auto src_offset(uint64_t v) const { return set_field<0, 9>(v); } -}; diff --git a/examples/qwen_prefill-amx/prepare.py b/examples/qwen_prefill-amx/prepare.py deleted file mode 100644 index cdef06a2..00000000 --- a/examples/qwen_prefill-amx/prepare.py +++ /dev/null @@ -1,38 +0,0 @@ -"""Canonical GGUF weights and offline layout conversion via Module.prepare.""" - -from pathlib import Path -import argparse, json -from scheduled import make_scheduled -from config import config, ROOT -from weights import aliases, GGUFResource - - -def prepare(model, checkpoint, output=None, sequence=128): - c = config(model) - program = make_scheduled(c, sequence) - out = Path(output) if output else ROOT / "artifacts" / ("prepared-" + model) - program.prepare(GGUFResource(Path(checkpoint), aliases(c)), str(out), device="cpu") - (out / "manifest.json").write_text( - json.dumps( - dict( - model=model, - sequence=sequence, - checkpoint=str(Path(checkpoint).resolve()), - format="module-list-v1", - ), - indent=2, - ) - + "\n" - ) - print(out, flush=True) - return out - - -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("--output") - p.add_argument("--sequence", type=int, default=128) - a = p.parse_args() - prepare(a.model, a.checkpoint, a.output, a.sequence) diff --git a/examples/qwen_prefill-amx/run.py b/examples/qwen_prefill-amx/run.py index 7d6ffe26..fbd86d81 100644 --- a/examples/qwen_prefill-amx/run.py +++ b/examples/qwen_prefill-amx/run.py @@ -1,37 +1,37 @@ -"""One Prefill smoke run; no warmups, repetitions or performance comparison.""" +"""Load Qwen weights and execute one text prefill through its runtime module.""" -import argparse, json +import argparse from pathlib import Path import torch from tilefoundry.runtime import SafetensorsResource -from config import config, ROOT -from scheduled import make_scheduled +from config import config from runtime_model import make_runtime, inputs -def run(model, sequence=128, prepared=None): +def run(model, sequence=128, prepared=None, checkpoint=None): torch.set_num_threads(8) c = config(model) - semantic = make_scheduled(c, sequence) - runtime = make_runtime(semantic, c)() - source = Path(prepared) if prepared else ROOT / "artifacts" / ("prepared-" + model) + runtime = make_runtime(c, sequence)() + source = Path(prepared) if prepared else Path(__file__).parent / "prepared" / model + if not source.exists(): + if checkpoint is None: + raise ValueError( + "Supply --checkpoint to prepare weights, or --prepared for existing weights" + ) + from weights import aliases, GGUFResource + + runtime.module.prepare( + GGUFResource(checkpoint, aliases(c)), str(source), device="cpu" + ) runtime.load(SafetensorsResource(str(source), device="cpu")) tokens = (torch.arange(sequence) * 17 + 100) % min(30000, c["vocab_size"]) - logits = runtime(*inputs(c, sequence, tokens)) - assert tuple(logits.shape) == (1, c["vocab_size"]) and torch.isfinite(logits).all() - report = dict( - model=model, - sequence=sequence, - shape=list(logits.shape), - argmax=int(logits.argmax()), - finite=True, - stats=runtime.stats, - smoke_only=True, - codegen=False, - ) - (source / "smoke.json").write_text(json.dumps(report, indent=2) + "\n") - print(json.dumps(report), flush=True) - return report + with torch.inference_mode(): + logits = runtime(*inputs(c, sequence, tokens)) + if not torch.isfinite(logits).all(): + raise RuntimeError("Prefill produced non-finite logits") + token = int(logits.argmax()) + print(f"{model}: logits {tuple(logits.shape)}, next token {token}") + return logits if __name__ == "__main__": @@ -39,5 +39,6 @@ def run(model, sequence=128, prepared=None): p.add_argument("--model", choices=("1.7b", "4b"), required=True) p.add_argument("--sequence", type=int, default=128) p.add_argument("--prepared") + p.add_argument("--checkpoint") a = p.parse_args() - run(a.model, a.sequence, a.prepared) + run(a.model, a.sequence, a.prepared, a.checkpoint) diff --git a/examples/qwen_prefill-amx/runtime_model.py b/examples/qwen_prefill-amx/runtime_model.py index 858dd834..b2ebcbd9 100644 --- a/examples/qwen_prefill-amx/runtime_model.py +++ b/examples/qwen_prefill-amx/runtime_model.py @@ -1,23 +1,11 @@ -"""One-to-one runtime twins of the whole authored/scheduled model tree. - -AMX linear leaves launch the existing hand-written native kernel through TVM -FFI. All other leaves execute their typed tensor program or a checked GDN -implementation. Hand-written runtime twins follow the HIR function call tree. -There is no llama.cpp whole-model callback in this path and no codegen. -""" - -import os -from pathlib import Path -import torch, tvm_ffi -from tilefoundry.evaluator import evaluate +"""Python runtime twins with native AMX projections and Torch host operations.""" + +import torch +import torch.nn.functional as F from tilefoundry.runtime import runtime_func, runtime_module from tilefoundry.dsl import ModuleList -from tilefoundry.ir.core import Call -from tilefoundry.ir.visitor import collect_exprs -from tilefoundry.ir.hir.schedule import ScheduleOp -from instructions import Fma16 - -R = Path(__file__).resolve().parent +from kernels import ops +from scheduled import make_scheduled def gdn(qkv, z, alpha, beta_raw, conv, decay, dt_bias, gamma, *, nk, nv, dk, dv, eps): @@ -57,108 +45,222 @@ def gdn(qkv, z, alpha, beta_raw, conv, decay, dt_bias, gamma, *, nk, nv, dk, dv, return (out * F.silu(z.reshape(s, nv, dv))).reshape(s, vd) -def make_runtime(semantic, config, *, native=True): - os.environ["TF_AMX_THREADS"] = "8" - ffi = tvm_ffi.load_module(str(R / "artifacts/prefill.dylib")) if native else None - stats = {"amx_calls": 0, "tensor_calls": 0} - - classes = {} - - def twin(mod): - if id(mod) in classes: - return classes[id(mod)] - namespace = { - key: ModuleList(tuple(twin(child) for child in slot.elements)) - if slot.indexed - else twin(slot.elements[0]) - for key, slot in mod.children.items() - } - namespace["stats"] = stats - for fn in mod.functions: - is_amx = ( - len(fn.params) == 2 - and fn.params[1].name == "w" - and len(fn.params[1].type.shape) == 3 - and any( - isinstance(v, Call) - and isinstance(v.target, ScheduleOp) - and isinstance(v.target.op, Fma16) - for v in collect_exprs(fn.body) +def make_runtime(config, sequence=128): + semantic = make_scheduled(config, sequence) + eps = config["rms_norm_eps"] + hq = config["num_attention_heads"] + hk = config["num_key_value_heads"] + d = config["head_dim"] + gated = config.get("attn_output_gate", False) + + def linear(sem, *, head=False): + @runtime_module(sem) + class Linear: + @runtime_func + def forward(self, x, w): + if head: + return x.half().float() @ w.float() + out = torch.empty(x.shape[0], w.shape[0] * w.shape[2]) + ops().linear(x.contiguous(), w, out) + return out + + return Linear + + def make_norm(sem): + @runtime_module(sem) + class Norm: + @runtime_func + def forward(self, x, gamma): + return F.rms_norm(x, (x.shape[-1],), gamma, eps) + + return Norm + + def make_mlp(sem): + @runtime_module(sem) + class MLP: + gate = linear(sem.gate) + up = linear(sem.up) + down = linear(sem.down) + + @runtime_func + def combine(self, g, u): + return F.silu(g) * u + + @runtime_func + def forward(self, x): + return self.down(self.combine(self.gate(x), self.up(x))) + + return MLP + + def attention(sem): + @runtime_module(sem.query_norm) + class QueryNorm: + @runtime_func + def query_norm(self, q, gamma_q): + return F.rms_norm(q, (d,), gamma_q, eps) + + @runtime_module(sem.key_norm) + class KeyNorm: + @runtime_func + def key_norm(self, k, gamma_k): + return F.rms_norm(k, (d,), gamma_k, eps) + + def rope(x, cos, sin): + r = cos.shape[-1] + part = x[..., :r] + turned = torch.cat((-part[..., r // 2 :], part[..., : r // 2]), -1) + result = part * cos[None, :, None, :] + turned * sin[None, :, None, :] + return torch.cat((result, x[..., r:]), -1) + + @runtime_module(sem) + class Attention: + query = linear(sem.query) + key = linear(sem.key) + value = linear(sem.value) + output = linear(sem.output) + query_norm = QueryNorm + key_norm = KeyNorm + + @runtime_func + def query_rope(self, q, cos, sin): + return rope(q, cos, sin) + + @runtime_func + def key_rope(self, k, cos, sin): + return rope(k, cos, sin) + + @runtime_func + def split_query(self, x): + return x.reshape(1, sequence, hq, -1)[..., :d] + + @runtime_func + def split_kv(self, x): + return x.reshape(1, sequence, hk, d) + + if gated: + + @runtime_func + def split_gate(self, x): + return x.reshape(sequence, hq, 2 * d)[..., d:].reshape( + sequence, hq * d + ) + + @runtime_func + def gate_output(self, x, gate): + return x * torch.sigmoid(gate) + + @runtime_func + def attend(self, q, k, v, mask): + q = q.transpose(1, 2) + k = k.repeat_interleave(hq // hk, 2).permute(0, 2, 3, 1) + v = v.repeat_interleave(hq // hk, 2).transpose(1, 2) + p = torch.softmax(q @ k / (d**0.5) + mask, -1) + return (p @ v).transpose(1, 2).reshape(sequence, hq * d) + + @runtime_func + def forward(self, x, cos, sin, mask): + qg = self.query(x) + q = self.query_rope(self.query_norm(self.split_query(qg)), cos, sin) + k = self.key_rope(self.key_norm(self.split_kv(self.key(x))), cos, sin) + out = self.attend(q, k, self.split_kv(self.value(x)), mask) + if gated: + out = self.gate_output(out, self.split_gate(qg)) + return self.output(out) + + return Attention + + def recurrent(sem): + @runtime_module(sem.math) + class GDN: + @runtime_func + def forward(self, qkv, z, alpha, beta_raw, conv, decay, dt_bias, gamma): + return gdn( + qkv, + z, + alpha, + beta_raw, + conv, + decay, + dt_bias, + gamma, + nk=config["linear_num_key_heads"], + nv=config["linear_num_value_heads"], + dk=config["linear_key_head_dim"], + dv=config["linear_value_head_dim"], + eps=eps, ) + + @runtime_module(sem) + class Recurrent: + qkv = linear(sem.qkv) + gate = linear(sem.gate) + alpha = linear(sem.alpha) + beta = linear(sem.beta) + output = linear(sem.output) + math = GDN + + @runtime_func + def forward(self, x, cos, sin, mask): + return self.output( + self.math(self.qkv(x), self.gate(x), self.alpha(x), self.beta(x)) + ) + + return Recurrent + + def layer(sem, kind): + @runtime_module(sem) + class Layer: + norm_in = make_norm(sem.norm_in) + norm_ffn = make_norm(sem.norm_ffn) + mlp = make_mlp(sem.mlp) + mixer = ( + attention(sem.mixer) + if kind == "full_attention" + else recurrent(sem.mixer) ) - def body(self, *args, fn=fn, is_amx=is_amx): - if fn.name == "prefill": - tokens, cos, sin, mask = args[:4] - hidden = self.embed(tokens) - for child in self.layers: - hidden = child(hidden, cos, sin, mask) - return self.head(self.last(self.norm(hidden))) - if fn.name == "forward" and self.children: - if hasattr(self, "norm_in"): - x, cos, sin, mask = args - x = self.add_mixer( - x, self.mixer(self.norm_in(x), cos, sin, mask) - ) - return self.add_mlp(x, self.mlp(self.norm_ffn(x))) - if hasattr(self, "down"): - x = args[0] - return self.down(self.combine(self.gate(x), self.up(x))) - if hasattr(self, "qkv"): - x = args[0] - return self.output( - self.math( - self.qkv(x), self.gate(x), self.alpha(x), self.beta(x) - ) - ) - if hasattr(self, "query"): - x, cos, sin, mask = args - qg = self.query(x) - q = self.query_rope( - self.query_norm(self.split_query(qg)), cos, sin - ) - k = self.key_rope( - self.key_norm(self.split_kv(self.key(x))), cos, sin - ) - v = self.split_kv(self.value(x)) - out = self.attend(q, k, v, mask) - if config.get("attn_output_gate", False): - out = self.gate_output(out, self.split_gate(qg)) - return self.output(out) - if is_amx: - x, w = args - M, K = x.shape - N = w.shape[0] * w.shape[2] - if native: - stats["amx_calls"] += 1 - y = torch.empty(M, N, dtype=torch.float32) - ffi.linear(x.contiguous(), w, y) - return y - raw = w.permute(0, 2, 1).reshape(N, K) - return x.half().float() @ raw.float().T - if ( - any(p.name == "conv" for p in fn.params) - and len(args) == 8 - and "linear_num_key_heads" in config - ): - return gdn( - *args, - nk=config["linear_num_key_heads"], - nv=config["linear_num_value_heads"], - dk=config["linear_key_head_dim"], - dv=config["linear_value_head_dim"], - eps=config["rms_norm_eps"], - ) - stats["tensor_calls"] += 1 - return evaluate(fn, *args) + @runtime_func + def add_mixer(self, x, y): + return x + y - namespace[fn.name] = runtime_func(body) - classes[id(mod)] = runtime_module(mod)( - type(mod.name + "Runtime", (), namespace) - ) - return classes[id(mod)] + @runtime_func + def add_mlp(self, x, y): + return x + y + + @runtime_func + def forward(self, x, cos, sin, mask): + x = self.add_mixer(x, self.mixer(self.norm_in(x), cos, sin, mask)) + return self.add_mlp(x, self.mlp(self.norm_ffn(x))) + + return Layer + + kinds = config.get("layer_types", ["full_attention"] * config["num_hidden_layers"]) + definitions = { + kind: layer(semantic.layers[kinds.index(kind)], kind) for kind in set(kinds) + } + + @runtime_module(semantic) + class QwenPrefill: + layers = ModuleList(tuple(definitions[kind] for kind in kinds)) + norm = make_norm(semantic.norm) + head = linear(semantic.head, head=True) + + @runtime_func + def embed(self, tokens, table): + return table[tokens].float() + + @runtime_func + def last(self, x): + return x[-1:] + + @runtime_func + def prefill(self, tokens, cos, sin, mask, table): + hidden = self.embed(tokens) + for child in self.layers: + hidden = child(hidden, cos, sin, mask) + return self.head(self.last(self.norm(hidden))) - return twin(semantic) + return QwenPrefill def inputs(config, sequence, tokens): diff --git a/examples/qwen_prefill-amx/scheduled.py b/examples/qwen_prefill-amx/scheduled.py index 01264af1..5bbe1da9 100644 --- a/examples/qwen_prefill-amx/scheduled.py +++ b/examples/qwen_prefill-amx/scheduled.py @@ -1,17 +1,10 @@ -"""Explicit AMX projection schedule, plus typed CPU-kernel fallback choices.""" +"""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 instructions import Z_LAYOUT -from tilefoundry.ir.core import Call, Var -from tilefoundry.ir.visitor import collect_exprs -from tilefoundry.ir.hir.schedule import ScheduleOp -from tilefoundry.ir.hir.mesh_region import MeshRegion -from tilefoundry.ir.visitor import ExprCloner, BindingSubstitutionCloner -from host_instructions import instruction_for, INSTRUCTIONS -from tilefoundry.ir.types import Mesh as MeshValue +from tilefoundry.ir.tir.amx.instructions import Z_LAYOUT from model import build BM = 32 @@ -40,7 +33,7 @@ def forward(x: Tensor[(1, K), "f32"], w: ConstTensor[(K, N), "f16"]): def _(raw: ConstTensor[(N, K), "f16"]): return tf.transpose(raw, perm=(1, 0)) - return cpu_schedule(Head) + return Head @module(entry="forward", topologies=(Topology("core", 8), Topology("amx", 1))) class Linear: @@ -49,10 +42,7 @@ 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.amx_host_cast(dtype=DType.f16), - ) + half = tf.cast(x, dtype="f16") view = tf.transpose( tf.reshape(half, new_shape=(M // BM, BM, K)), perm=(0, 2, 1) ) @@ -68,9 +58,7 @@ def forward(x: Tensor[(M, K), "f32"], w: ConstTensor[(N // B, K, B), "f16"]): ] ), ), - op=T.amx_pack( - gmem_layout=Layout((M // BM, K, BM), (K * BM, BM, 1)) - ), + op=T.copy(), ) out = tf.zeros(Tensor[(M, N), "f32"]) with Mesh(("core",), layout=(WORKERS,), names=("worker",)) as cores: @@ -137,85 +125,8 @@ def _(raw: ConstTensor[(N, K), "f16"]): 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) - # Integer coordinate arithmetic stays index arithmetic, not a kernel. - 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): - root = cpu_schedule(build(config, sequence, linear=scheduled_projection)) - return replace(root, topologies=(Topology("core", 8), Topology("amx", 1))) + return replace( + build(config, sequence, linear=scheduled_projection), + topologies=(Topology("core", 8), Topology("amx", 1)), + ) diff --git a/examples/qwen_prefill-amx/verify_ir.py b/examples/qwen_prefill-amx/verify_ir.py deleted file mode 100644 index 6003d118..00000000 --- a/examples/qwen_prefill-amx/verify_ir.py +++ /dev/null @@ -1,34 +0,0 @@ -"""Reload exported scheduled/TIR sources, without compiling machine code.""" - -import argparse, importlib.util -from config import ROOT -from tilefoundry.ir.core.module import Module -from tilefoundry.ir.tir.prim_function import PrimFunction - - -def verify(model): - path = ROOT / "artifacts" / ("ir-" + model) - for name in ("model", "scheduled", "finalized"): - spec = importlib.util.spec_from_file_location( - "exported_" + name, path / (name + ".py") - ) - loaded = importlib.util.module_from_spec(spec) - spec.loader.exec_module(loaded) - root = loaded.QwenPrefill - assert isinstance(root, Module) - if name == "finalized": - assert isinstance(root.entry_function(), PrimFunction) - print( - "PASS", - model, - name, - "shared layers:", - len({id(m) for m in root.layers}), - flush=True, - ) - - -if __name__ == "__main__": - p = argparse.ArgumentParser() - p.add_argument("--model", choices=("1.7b", "4b"), required=True) - verify(p.parse_args().model) diff --git a/include/tilefoundry/runtime/amx/instructions.h b/include/tilefoundry/runtime/amx/instructions.h new file mode 100644 index 00000000..8a2b2fdf --- /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 85afe9c3..3f90b0bf 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/analysis/footprint.py b/src/tilefoundry/analysis/footprint.py index 77340591..d649e569 100644 --- a/src/tilefoundry/analysis/footprint.py +++ b/src/tilefoundry/analysis/footprint.py @@ -1187,21 +1187,27 @@ def resolve_footprint(footprint, ctx): continue def project_image(image): - """Keep instance coordinates correlated with addresses during projection.""" - relation = isl.map.from_range(image) for axis in reversed(range(len(item.ownership))): - relation = relation.move_dims(isl.dim_type.IN, 0, isl.dim_type.OUT, axis, 1) + image = image.set_dim_name(isl.dim_type.SET, axis, f"__tf_instance_{axis}") + image = image.move_dims( + isl.dim_type.PARAM, image.dim(isl.dim_type.PARAM), isl.dim_type.SET, axis, 1 + ) source = BufferAlias( item.target, - AccessRelation(relation, item.values), + AccessRelation(isl.map.from_range(image), item.values), AliasKind.FORWARD, ) projected, precision = binding.project(source) relation = projected.relation.relation values = dict(projected.relation.values) - for axis in range(len(item.ownership)): - relation = relation.move_dims(isl.dim_type.OUT, axis, isl.dim_type.IN, depth, 1) relation = relation.project_out(isl.dim_type.IN, 0, depth) + for axis in range(len(item.ownership)): + parameter = relation.find_dim_by_name( + isl.dim_type.PARAM, f"__tf_instance_{axis}" + ) + relation = relation.move_dims( + isl.dim_type.OUT, axis, isl.dim_type.PARAM, parameter, 1 + ) image = relation.range() return projected.target, image, values, precision diff --git a/src/tilefoundry/ir/tir/amx/__init__.py b/src/tilefoundry/ir/tir/amx/__init__.py new file mode 100644 index 00000000..997a6237 --- /dev/null +++ b/src/tilefoundry/ir/tir/amx/__init__.py @@ -0,0 +1 @@ +"""Apple AMX instruction contracts.""" diff --git a/examples/qwen_prefill-amx/instructions.py b/src/tilefoundry/ir/tir/amx/instructions.py similarity index 82% rename from examples/qwen_prefill-amx/instructions.py rename to src/tilefoundry/ir/tir/amx/instructions.py index ac60a744..4215fbfc 100644 --- a/examples/qwen_prefill-amx/instructions.py +++ b/src/tilefoundry/ir/tir/amx/instructions.py @@ -1,40 +1,38 @@ -"""Example-local AMX contracts; no machine-code generation is involved. +"""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 ParamDef, MemoryEffect +from tilefoundry.ir.core.param_def import MemoryEffect, ParamDef from tilefoundry.ir.core.register import register_op from tilefoundry.ir.pattern import ( - TensorPattern, + ComposedLayoutPattern, LayoutPattern, MeshPattern, - ComposedLayoutPattern, + TensorPattern, WildcardPattern, ) from tilefoundry.ir.types import ( DType, Layout, Mesh, - Topology, StorageKind, + Topology, UnitType, - LayoutBase, ) from tilefoundry.visitor_registry import register_typeinfer from tilefoundry.visitor_registry.access_relation import ( - register_access_relation, identity_relations, matmul_relations, + register_access_relation, ) -from tilefoundry.evaluator.registry import register_schedule_eval -from tilefoundry.evaluator.value import TensorValue Z_LAYOUT = Layout((32, (16, 2)), (32, (1, 16))) AMX = Mesh((Topology("amx", 1),), Layout((1,), (1,)), ("unit",)) -CORE = Mesh((Topology("core", 8),), Layout((1,), (1,)), ("worker",)) def mesh_pattern(name): @@ -49,7 +47,6 @@ def mesh_pattern(name): AMX_PATTERN = mesh_pattern("amx") -CORE_PATTERN = mesh_pattern("core") R = MemoryEffect.READ W = MemoryEffect.WRITE @@ -143,22 +140,3 @@ def transfer(name, shape, dtype, storage, execution): 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) - - -@register_op(dialect="T", category="amx", name="amx_pack") -class Pack(Op): - capability = OpCapability("amx_outer_product") - resource = "cpu_engine" - execution_mesh = CORE_PATTERN - src = ParamDef(kind="input", effect=R, pattern=TensorPattern()) - dst = ParamDef( - kind="input", effect=W, pattern=TensorPattern(storage=StorageKind.GMEM) - ) - gmem_layout = ParamDef(kind="attribute", annotation=LayoutBase) - - -register_typeinfer(Pack)(lambda call, ctx: UnitType()) -register_access_relation(Pack)(identity_relations) -register_schedule_eval(Pack)( - lambda ctx: TensorValue(data=ctx.args[0].data.contiguous(), type=ctx.result_type) -) diff --git a/src/tilefoundry/passes/transforms/convert_hir_to_tir.py b/src/tilefoundry/passes/transforms/convert_hir_to_tir.py index 0e1d9ded..b95f657f 100644 --- a/src/tilefoundry/passes/transforms/convert_hir_to_tir.py +++ b/src/tilefoundry/passes/transforms/convert_hir_to_tir.py @@ -1618,12 +1618,7 @@ 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 getattr(op, param.name) != param.default - or param.name == "fill" - ) + if param.kind == "attribute" and (not param.has_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/tests/analysis/test_analyze_cross_module.py b/tests/analysis/test_analyze_cross_module.py index 8502d6d9..610750c3 100644 --- a/tests/analysis/test_analyze_cross_module.py +++ b/tests/analysis/test_analyze_cross_module.py @@ -11,10 +11,8 @@ from __future__ import annotations from dataclasses import dataclass, replace -from types import SimpleNamespace from typing import ClassVar -import isl import pytest from tests.fixtures.logical.hir_composition import REFERENCE_PROGRAMS, CrossModule, Expert @@ -22,20 +20,10 @@ from tests.fixtures.shapes.scaled_modules import FusedScaledParent, ScaledChild, StackedScaledParent from tests.models.qwen3_1_7b.case import CASE as QWEN_CASE from tests.models.qwen3_1_7b.model import Qwen3_1_7B -from tilefoundry.analysis.access import SSAValue from tilefoundry.analysis.api import analyze from tilefoundry.analysis.check import check_program from tilefoundry.analysis.errors import AnalysisError -from tilefoundry.analysis.footprint import ReachedAddresses, resolve_footprint -from tilefoundry.analysis.metadata import ( - BindingValuation, - ComputeCostMetadata, - Footprint, - LoopNestValuation, - MemoryMetadata, - RegionMemoryMetadata, -) -from tilefoundry.analysis.precision import AnalysisPrecision +from tilefoundry.analysis.metadata import ComputeCostMetadata, MemoryMetadata, RegionMemoryMetadata from tilefoundry.analysis.registry import Analyzer from tilefoundry.dsl import * from tilefoundry.ir.core import Call, Constant, Tuple, Var, attach_metadata, get_metadata @@ -58,7 +46,6 @@ from tilefoundry.ir.types.utils import tensor_types from tilefoundry.ir.visitor import collect_exprs from tilefoundry.target import CudaTarget, Target -from tilefoundry.visitor_registry.access_relation import AccessRelation from tilefoundry.visitor_registry.call_frame import CallFrame, CallFrameVisitor _H200 = "nvidia.h200_sxm" @@ -554,57 +541,3 @@ def test_repeated_callee_placement_compares_alias_even_at_the_same_offset(): ] assert len(buffers) == 2 assert all(len(record.offsets) == 1 for record in buffers) - - -@pytest.mark.parametrize( - "ownership,coordinates", - [ - (("layers",), "{ [layer, addr] : 0 <= layer < 3 and addr = 2*layer }"), - ( - ("layers", "experts"), - "{ [layer, expert, addr] : 0 <= layer < 2 and 0 <= expert < 3 and addr = 10*layer + expert }", - ), - ], -) -def test_nested_weight_footprint_projection_keeps_instance_address_correlation( - ownership, coordinates -): - - weight = Var("w", type=TensorType((32,), DType.f32, None, StorageKind.GMEM), is_const=True) - item = ReachedAddresses( - SSAValue(weight), - None, - DType.f32, - isl.set(coordinates), - AnalysisPrecision.EXACT, - ownership=ownership, - ) - footprint = Footprint((((), (item,)),)) - - def project(source): - shifted = source.relation.relation.apply_range(isl.map("{ [i] -> [j] : j=i+4 }")) - return replace( - source, relation=AccessRelation(shifted, source.relation.values) - ), AnalysisPrecision.EXACT - - values = { - BindingValuation: SimpleNamespace(project=project), - LoopNestValuation: SimpleNamespace(domain=lambda: isl.set("{ [] }")), - } - result = resolve_footprint( - footprint, SimpleNamespace(request_valuation=lambda kind: values[kind]) - ) - actual = result.groups[0][1][0] - rank = len(ownership) - offset = isl.map( - "{ [" - + ",".join(f"i{k}" for k in range(rank + 1)) - + "] -> [" - + ",".join([*(f"i{k}" for k in range(rank)), f"i{rank}+4"]) - + "] }" - ) - expected = item.reached.apply(offset) - assert actual.reached.is_equal(expected) - assert actual.must.is_equal(expected) - assert actual.ownership == ownership - assert not actual.values diff --git a/tests/schedule/test_attribute_preservation.py b/tests/schedule/test_attribute_preservation.py deleted file mode 100644 index 29a4c1a1..00000000 --- a/tests/schedule/test_attribute_preservation.py +++ /dev/null @@ -1,40 +0,0 @@ -"""Lowering must retain authored choices that differ from instruction defaults.""" - -import torch - -from tilefoundry import func, module -from tilefoundry.dsl import Mesh, T, Tensor, Topology, tf -from tilefoundry.ir.tir.arange import Arange -from tilefoundry.ir.visitor import StmtVisitor, walk_prim_function -from tilefoundry.schedule import finalize -from tilefoundry.target import CudaTarget - - -@module( - entry="run", - target=CudaTarget("nvidia.h200_sxm"), - topologies=(Topology("cta", 1), Topology("thread", 32)), -) -class NondefaultRange: - @func - def run(): - with Mesh(("thread",), layout=(1,), names=("lane",)) as _thread: - output = tf.empty(Tensor[(4,), "i64", None, "rmem"]) - return tf.schedule((output,), op=T.arange(start=5, step=2)) - - -def test_nondefault_instruction_attributes_survive_lowering(): - torch.testing.assert_close(NondefaultRange(), torch.tensor([5, 7, 9, 11])) - - class Find(StmtVisitor): - def __init__(self): - self.operations = [] - - def visit_Evaluate(self, statement): - if isinstance(statement.callable, Arange): - self.operations.append(statement.callable) - - visitor = Find() - lowered = finalize(NondefaultRange) - walk_prim_function(visitor, lowered.entry_function()) - assert [(op.start, op.step) for op in visitor.operations] == [(5, 2)] diff --git a/tests/target/test_amx_target.py b/tests/target/test_amx_target.py index d6e5bd65..898f0176 100644 --- a/tests/target/test_amx_target.py +++ b/tests/target/test_amx_target.py @@ -9,12 +9,8 @@ import pytest -from tilefoundry import func, module -from tilefoundry.dsl import Mesh, Tensor, tf -from tilefoundry.inspection import as_script from tilefoundry.ir.types import DType, Topology from tilefoundry.ir.types.dim import DimVar -from tilefoundry.schedule import finalize from tilefoundry.target import ( AmxTarget, TopologyFacts, @@ -114,20 +110,3 @@ def test_unmeasured_units_and_dtypes_have_no_throughput_entry(): target.device.throughput_for("amx", DType.f16) with pytest.raises(ValueError, match="no measured compute-throughput"): target.device.throughput_for("sme", DType.f32) - - -@pytest.mark.parametrize("level", ["core", "amx"]) -def test_memory_lowering_admits_apple_executor_scopes(level): - """Copy/fill remain valid when scheduled in Apple executor scopes.""" - - @module(entry="run", target=AmxTarget(), topologies=(Topology("core", 8), Topology("amx", 1))) - class AppleMemory: - @func - def run(): - with Mesh((level,), layout=(1,), names=("unit",)) as _unit: - x = tf.zeros(Tensor[(2, 3), "f32", None, "gmem"]) - return tf.transpose(x, perm=(1, 0)) - - result = finalize(AppleMemory) - text = as_script(result) - assert "T.fill" in text and "T.copy" in text From 34eff555fe7d3ace6188f820110e1f1d8035442e Mon Sep 17 00:00:00 2001 From: QiHang Zheng <597323109@qq.com> Date: Mon, 19 Oct 2026 14:15:09 +0800 Subject: [PATCH 6/8] fix(amx,examples): restore native prefill and complete CLI support --- docs/spec/target.md | 14 + docs/spec/tir.md | 8 + examples/qwen_prefill-amx/README.md | 219 ++++++++--- examples/qwen_prefill-amx/kernels/__init__.py | 150 +++++--- examples/qwen_prefill-amx/kernels/prefill.cpp | 166 ++++++++- examples/qwen_prefill-amx/kernels/runtime.hpp | 123 ------- examples/qwen_prefill-amx/model.py | 13 +- examples/qwen_prefill-amx/run.py | 44 +-- examples/qwen_prefill-amx/runtime_model.py | 340 ++++-------------- examples/qwen_prefill-amx/scheduled.py | 105 +++++- examples/qwen_prefill-amx/weights.py | 49 +++ src/tilefoundry/analysis/footprint.py | 18 +- src/tilefoundry/ir/hir/schedule.py | 4 + src/tilefoundry/ir/tir/cpu/__init__.py | 1 + src/tilefoundry/ir/tir/cpu/tensor.py | 203 +++++++++++ .../passes/transforms/convert_hir_to_tir.py | 7 +- src/tilefoundry/target/amx/device.py | 2 + src/tilefoundry/target/amx/facts.py | 29 ++ .../target/amx/hardware/apple_m2_pro.toml | 15 + src/tilefoundry/target/amx/spec.py | 4 + src/tilefoundry/target/amx/target.py | 4 + .../visitor_registry/candidates/amx.py | 8 + 22 files changed, 993 insertions(+), 533 deletions(-) delete mode 100644 examples/qwen_prefill-amx/kernels/runtime.hpp create mode 100644 src/tilefoundry/ir/tir/cpu/__init__.py create mode 100644 src/tilefoundry/ir/tir/cpu/tensor.py create mode 100644 src/tilefoundry/visitor_registry/candidates/amx.py diff --git a/docs/spec/target.md b/docs/spec/target.md index a880a4a7..e623e7c2 100644 --- a/docs/spec/target.md +++ b/docs/spec/target.md @@ -583,6 +583,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 +611,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` diff --git a/docs/spec/tir.md b/docs/spec/tir.md index 9503212c..eeae9b22 100644 --- a/docs/spec/tir.md +++ b/docs/spec/tir.md @@ -591,6 +591,14 @@ 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. + Copy and Fill admit CUDA `thread`, CPU `core`, and Apple `amx` execution scopes through the shared memory execution pattern. A sharded register operand must name the selected scope. Their canonical CUDA `execution_mesh` descriptor diff --git a/examples/qwen_prefill-amx/README.md b/examples/qwen_prefill-amx/README.md index 1f8eeef5..33db2422 100644 --- a/examples/qwen_prefill-amx/README.md +++ b/examples/qwen_prefill-amx/README.md @@ -1,64 +1,199 @@ -# Qwen text prefill on Apple AMX +# Qwen text prefill on TileFoundry โ€” Apple AMX kernels -Qwen3-1.7B and Qwen3.5-4B text prefill, batch one, on Apple Silicon. The HIR -expresses the whole model; the scheduled HIR assigns its projections to AMX. -A Python runtime module implements the same module tree and calls C++ AMX -kernels through TVM FFI. Other operations run on the CPU using Torch. +Runtime execution and CLI integration rechecked on **2026-10-19**, in the +development checkout based on `3cf660e`. -## Files +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. -| File | Purpose | +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. + +| | | |---|---| -| `model.py` | Whole-model HIR, shared layer definitions and weight converters | -| `scheduled.py` | AMX projection schedule, packed layouts and worker assignment | -| `runtime_model.py` | Explicit Python runtime modules for projections, attention, GDN and layers | -| `kernels/prefill.cpp`, `kernels/runtime.hpp` | Native projection kernel and FFI entry | -| `kernels/__init__.py` | Build and load the native kernel on first use | -| `weights.py` | Q8_0 GGUF reading and resource aliases | -| `config.py` | Architecture parameters for the two models | -| `run.py` | Weight preparation, loading and one prefill call | +| 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. -## Run +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). -Requires an Apple Silicon Mac, Xcode Command Line Tools, Python 3.12+, this -TileFoundry revision, Torch, `gguf`, `safetensors` and `apache-tvm-ffi`. +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 the GGUF weights under `prepared//` and compiles -the C++ kernel under `kernels/.build/`. Subsequent runs reuse both. Use -`--prepared /path/to/prepared-weights` to choose another resource directory. -`--sequence` selects a positive multiple of 32; the default is 128. The entry -runs a fixed token sequence and prints the logits shape and next-token ID. +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 runtime module can also be called directly with your own token IDs: +`--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 make_runtime, inputs -from tilefoundry.runtime import SafetensorsResource +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) +``` -c = config("1.7b") -runtime = make_runtime(c, sequence=128)() -runtime.load(SafetensorsResource("prepared/1.7b", device="cpu")) -logits = runtime(*inputs(c, 128, token_ids)) # token_ids: CPU int64 tensor [128] +`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. 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 ``` -Run this snippet from the example directory. `runtime.module` is the scheduled -HIR module; `model.build(c, 128)` constructs its unscheduled counterpart. +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. +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ร— | -## Execution arrangement +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. -Q8_0 weights are dequantized to FP16 and packed from `[N,K]` to `[N/B,K,B]`, -with B=64, or 32 for narrow projections. Eight logical workers divide N panels -cyclically. Each AMX unit accumulates a 32ร—32 FP32 tile from FP16 outer products; -activations are packed into contiguous 32-row panels. The final one-row vocabulary -projection uses Torch. GDN, normalization, rotary embeddings and attention also -execute in the Python runtime using Torch. +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`. -AMX instruction contracts and the native instruction wrappers are provided by -TileFoundry. The example supplies the model-specific schedule and hand-written -runtime. It covers text prefill through logits, with no decode, MTP or vision. +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/kernels/__init__.py b/examples/qwen_prefill-amx/kernels/__init__.py index 8d8224ff..bce35f28 100644 --- a/examples/qwen_prefill-amx/kernels/__init__.py +++ b/examples/qwen_prefill-amx/kernels/__init__.py @@ -1,60 +1,118 @@ -"""Build and cache the Apple Silicon prefill extension on first use.""" +"""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 -import sys -import tilefoundry _module = None def ops(): global _module - if _module is None: - if platform.system() != "Darwin" or platform.machine() != "arm64": - raise RuntimeError("The AMX example requires Apple Silicon macOS") - root = Path(__file__).resolve().parent - include = Path(sys.prefix) / "include" - if not (include / "tilefoundry/runtime/amx/instructions.h").exists(): - include = Path(tilefoundry.__file__).resolve().parents[2] / "include" - sources = ( - root / "prefill.cpp", - root / "runtime.hpp", - include / "tilefoundry/runtime/amx/instructions.h", + 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, ) - out = root / ".build/prefill.dylib" - ffi = Path(libinfo.find_libtvm_ffi()) - if not out.exists() or any( - p.stat().st_mtime > out.stat().st_mtime for p in sources - ): - out.parent.mkdir(exist_ok=True) - 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, - "-I" + str(Path(tvm_ffi.__file__).parent / "include"), - "-I" + str(include), - str(root / "prefill.cpp"), - str(ffi), - "-Wl,-rpath," + str(ffi.parent), - "-o", - str(out), - ], - check=True, - ) - _module = tvm_ffi.load_module(str(out)) + _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 index 545d1def..2bb75ed4 100644 --- a/examples/qwen_prefill-amx/kernels/prefill.cpp +++ b/examples/qwen_prefill-amx/kernels/prefill.cpp @@ -1,12 +1,41 @@ -/** Qwen prefill projection and TVM FFI entry. */ +/** TVM FFI entry for the supplied CPU/AMX model runtime. */ #include #include -#include "runtime.hpp" -#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 || @@ -17,6 +46,106 @@ static void tensor(const TensorView &t, int ndim, int code, int bits) { 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); @@ -34,5 +163,34 @@ void linear(TensorView x, TensorView w, TensorView out) { 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/kernels/runtime.hpp b/examples/qwen_prefill-amx/kernels/runtime.hpp deleted file mode 100644 index f4460941..00000000 --- a/examples/qwen_prefill-amx/kernels/runtime.hpp +++ /dev/null @@ -1,123 +0,0 @@ -#pragma once -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -namespace local_amx { -class Pool { - std::vector threads; - std::mutex mutex; - std::condition_variable ready, done; - std::atomic next{0}; - int count = 0, generation = 0, remaining = 0; - bool stop = false; - std::function task; - - public: - explicit Pool(int n) { - for (int i = 0; i < n; i++) - threads.emplace_back([this] { - int seen = 0; - for (;;) { - std::unique_lock lock(mutex); - ready.wait(lock, - [&] { return stop || generation != seen; }); - if (stop) - return; - seen = generation; - lock.unlock(); - for (int i = next.fetch_add(1); i < count; - i = next.fetch_add(1)) - task(i); - lock.lock(); - if (--remaining == 0) - done.notify_one(); - } - }); - } - ~Pool() { - { - std::lock_guard lock(mutex); - stop = true; - } - ready.notify_all(); - for (auto &thread : threads) - thread.join(); - } - void run(int n, std::function fn) { - std::unique_lock lock(mutex); - task = std::move(fn); - count = n; - next = 0; - remaining = threads.size(); - generation++; - ready.notify_all(); - done.wait(lock, [&] { return remaining == 0; }); - } -}; -inline Pool &pool() { - static Pool value( - std::getenv("TF_AMX_THREADS") - ? std::max(1, std::atoi(std::getenv("TF_AMX_THREADS"))) - : 8); - return value; -} -inline std::vector<_Float16> pack_input(const float *a, int m, int k) { - std::vector<_Float16> out(((m + 31) / 32) * k * 32, 0); - for (int m0 = 0; m0 < m; m0 += 32) - for (int p = 0; p < k; p++) - for (int j = 0; j < std::min(32, m - m0); j++) - out[(m0 / 32 * k + p) * 32 + j] = (_Float16)a[(m0 + j) * k + p]; - return out; -} -} -namespace shared_fp16 { -struct Weight { - int n, k, b; - const _Float16 *data; - Weight(const _Float16 *packed, int nn, int kk, int bb) - : n(nn), k(kk), b(bb), data(packed) {} - const _Float16 *at(int n0, int p) const { - return data + (n0 / b * k + p) * b + n0 % b; - } -}; -inline void prefill(const float *x, int m, const Weight &w, float *out, - int taskcols, bool cyclic = false) { - auto a = local_amx::pack_input(x, m, w.k); - int panels = (w.n + taskcols - 1) / taskcols; - local_amx::pool().run(cyclic ? 8 : panels, [&](int task) { - alignas(128) float z[32]; - tilefoundry::amx::start(); - for (int panel = task; panel < panels; panel += cyclic ? 8 : panels) - for (int n0 = panel * taskcols; - n0 < std::min(w.n, (panel + 1) * taskcols); n0 += 32) - for (int m0 = 0; m0 < m; m0 += 32) { - for (int p = 0; p < w.k; p++) { - tilefoundry::amx::load_x(w.at(n0, p)); - tilefoundry::amx::load_y(a.data() + - (m0 / 32 * w.k + p) * 32); - tilefoundry::amx::fma16(p == 0); - } - for (int row = 0; row < std::min(32, m - m0); row++) { - tilefoundry::amx::store_z(z, row * 2); - tilefoundry::amx::store_z(z + 16, row * 2 + 1); - for (int j = 0; j < 32; j += 8) { - auto pair = vzipq_f32(vld1q_f32(z + j / 2), - vld1q_f32(z + 16 + j / 2)); - vst1q_f32(out + (m0 + row) * w.n + n0 + j, - pair.val[0]); - vst1q_f32(out + (m0 + row) * w.n + n0 + j + 4, - pair.val[1]); - } - } - } - tilefoundry::amx::stop(); - }); -} -} diff --git a/examples/qwen_prefill-amx/model.py b/examples/qwen_prefill-amx/model.py index 04bf2fbc..17d80d40 100644 --- a/examples/qwen_prefill-amx/model.py +++ b/examples/qwen_prefill-amx/model.py @@ -7,7 +7,7 @@ import math from tilefoundry import module, func -from tilefoundry.dsl import Tensor, ConstTensor, ModuleList, tf +from tilefoundry.dsl import Tensor, ConstTensor, ModuleList, Topology, tf from tilefoundry.target import AmxTarget @@ -384,7 +384,7 @@ def forward( definitions = {kind: layer(kind) for kind in set(kinds)} - @module(target=AmxTarget()) + @module(target=AmxTarget(), topologies=(Topology("core", 8), Topology("amx", 1))) class Base: layers = ModuleList(tuple(definitions[kind] for kind in kinds)) norm = ModuleList(Norm) @@ -404,7 +404,7 @@ def last(x: Tensor[(S, H), "f32"]): import hashlib lines = [ - "@module(entry='prefill',target=AmxTarget())", + "@module(entry='prefill',target=AmxTarget(),topologies=(Topology('core',8),Topology('amx',1)))", "class QwenPrefill:", " layers=ModuleList(DEFINITIONS)", " norm=ModuleList(NORM)", @@ -435,6 +435,7 @@ def last(x: Tensor[(S, H), "f32"]): Tensor=Tensor, ConstTensor=ConstTensor, AmxTarget=AmxTarget, + Topology=Topology, S=S, R=R, H=H, @@ -442,3 +443,9 @@ def last(x: Tensor[(S, H), "f32"]): ) exec(compile(source, str(path), "exec"), namespace) return 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 index fbd86d81..0f0e2041 100644 --- a/examples/qwen_prefill-amx/run.py +++ b/examples/qwen_prefill-amx/run.py @@ -1,44 +1,36 @@ -"""Load Qwen weights and execute one text prefill through its runtime module.""" +"""Load the original CPU/AMX runtime and execute one text prefill.""" import argparse from pathlib import Path import torch -from tilefoundry.runtime import SafetensorsResource from config import config -from runtime_model import make_runtime, inputs +from runtime_model import QwenPrefill +from weights import PreparedModelResource -def run(model, sequence=128, prepared=None, checkpoint=None): +def run(model, checkpoint, sequence=128, prepared=None): torch.set_num_threads(8) - c = config(model) - runtime = make_runtime(c, sequence)() source = Path(prepared) if prepared else Path(__file__).parent / "prepared" / model - if not source.exists(): - if checkpoint is None: - raise ValueError( - "Supply --checkpoint to prepare weights, or --prepared for existing weights" - ) - from weights import aliases, GGUFResource - - runtime.module.prepare( - GGUFResource(checkpoint, aliases(c)), str(source), device="cpu" + 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())}" ) - runtime.load(SafetensorsResource(str(source), device="cpu")) - tokens = (torch.arange(sequence) * 17 + 100) % min(30000, c["vocab_size"]) - with torch.inference_mode(): - logits = runtime(*inputs(c, sequence, tokens)) - if not torch.isfinite(logits).all(): - raise RuntimeError("Prefill produced non-finite logits") - token = int(logits.argmax()) - print(f"{model}: logits {tuple(logits.shape)}, next token {token}") - return logits + 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") - p.add_argument("--checkpoint") a = p.parse_args() - run(a.model, a.sequence, a.prepared, a.checkpoint) + 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 index b2ebcbd9..41c0b572 100644 --- a/examples/qwen_prefill-amx/runtime_model.py +++ b/examples/qwen_prefill-amx/runtime_model.py @@ -1,278 +1,78 @@ -"""Python runtime twins with native AMX projections and Torch host operations.""" +"""Python adapter for the original supplied C++/AMX prefill runtime.""" import torch -import torch.nn.functional as F -from tilefoundry.runtime import runtime_func, runtime_module -from tilefoundry.dsl import ModuleList +from tilefoundry.runtime import RuntimeModule from kernels import ops from scheduled import make_scheduled -def gdn(qkv, z, alpha, beta_raw, conv, decay, dt_bias, gamma, *, nk, nv, dk, dv, eps): - import torch.nn.functional as F - - s = qkv.shape[0] - kd = nk * dk - vd = nv * dv - mixed = ( - F.silu( - F.conv1d( - qkv.T.unsqueeze(0), - conv.unsqueeze(1), - padding=conv.shape[1] - 1, - groups=qkv.shape[1], - )[:, :, :s] - ) - .squeeze(0) - .T - ) - q = mixed[:, :kd].reshape(s, nk, dk).repeat(1, nv // nk, 1) - k = mixed[:, kd : 2 * kd].reshape(s, nk, dk).repeat(1, nv // nk, 1) - q = q * torch.rsqrt((q * q).sum(-1, keepdim=True) + 1e-6) - k = k * torch.rsqrt((k * k).sum(-1, keepdim=True) + 1e-6) - v = mixed[:, 2 * kd :].reshape(s, nv, dv) - g = decay * F.softplus(alpha + dt_bias) - beta = torch.sigmoid(beta_raw) - state = torch.zeros(nv, dk, dv) - out = torch.empty(s, nv, dv) - for t in range(s): - state = state * torch.exp(g[t]).reshape(nv, 1, 1) - key = k[t].reshape(nv, dk, 1) - delta = (v[t] - (state * key).sum(1)) * beta[t].reshape(nv, 1) - state = state + key * delta.reshape(nv, 1, dv) - out[t] = (state * (q[t].reshape(nv, dk, 1) / (dk**0.5))).sum(1) - out = out * torch.rsqrt((out * out).mean(-1, keepdim=True) + eps) * gamma - return (out * F.silu(z.reshape(s, nv, dv))).reshape(s, vd) - - -def make_runtime(config, sequence=128): - semantic = make_scheduled(config, sequence) - eps = config["rms_norm_eps"] - hq = config["num_attention_heads"] - hk = config["num_key_value_heads"] - d = config["head_dim"] - gated = config.get("attn_output_gate", False) - - def linear(sem, *, head=False): - @runtime_module(sem) - class Linear: - @runtime_func - def forward(self, x, w): - if head: - return x.half().float() @ w.float() - out = torch.empty(x.shape[0], w.shape[0] * w.shape[2]) - ops().linear(x.contiguous(), w, out) - return out - - return Linear - - def make_norm(sem): - @runtime_module(sem) - class Norm: - @runtime_func - def forward(self, x, gamma): - return F.rms_norm(x, (x.shape[-1],), gamma, eps) - - return Norm - - def make_mlp(sem): - @runtime_module(sem) - class MLP: - gate = linear(sem.gate) - up = linear(sem.up) - down = linear(sem.down) - - @runtime_func - def combine(self, g, u): - return F.silu(g) * u - - @runtime_func - def forward(self, x): - return self.down(self.combine(self.gate(x), self.up(x))) - - return MLP - - def attention(sem): - @runtime_module(sem.query_norm) - class QueryNorm: - @runtime_func - def query_norm(self, q, gamma_q): - return F.rms_norm(q, (d,), gamma_q, eps) - - @runtime_module(sem.key_norm) - class KeyNorm: - @runtime_func - def key_norm(self, k, gamma_k): - return F.rms_norm(k, (d,), gamma_k, eps) - - def rope(x, cos, sin): - r = cos.shape[-1] - part = x[..., :r] - turned = torch.cat((-part[..., r // 2 :], part[..., : r // 2]), -1) - result = part * cos[None, :, None, :] + turned * sin[None, :, None, :] - return torch.cat((result, x[..., r:]), -1) - - @runtime_module(sem) - class Attention: - query = linear(sem.query) - key = linear(sem.key) - value = linear(sem.value) - output = linear(sem.output) - query_norm = QueryNorm - key_norm = KeyNorm - - @runtime_func - def query_rope(self, q, cos, sin): - return rope(q, cos, sin) - - @runtime_func - def key_rope(self, k, cos, sin): - return rope(k, cos, sin) - - @runtime_func - def split_query(self, x): - return x.reshape(1, sequence, hq, -1)[..., :d] - - @runtime_func - def split_kv(self, x): - return x.reshape(1, sequence, hk, d) - - if gated: - - @runtime_func - def split_gate(self, x): - return x.reshape(sequence, hq, 2 * d)[..., d:].reshape( - sequence, hq * d - ) - - @runtime_func - def gate_output(self, x, gate): - return x * torch.sigmoid(gate) - - @runtime_func - def attend(self, q, k, v, mask): - q = q.transpose(1, 2) - k = k.repeat_interleave(hq // hk, 2).permute(0, 2, 3, 1) - v = v.repeat_interleave(hq // hk, 2).transpose(1, 2) - p = torch.softmax(q @ k / (d**0.5) + mask, -1) - return (p @ v).transpose(1, 2).reshape(sequence, hq * d) - - @runtime_func - def forward(self, x, cos, sin, mask): - qg = self.query(x) - q = self.query_rope(self.query_norm(self.split_query(qg)), cos, sin) - k = self.key_rope(self.key_norm(self.split_kv(self.key(x))), cos, sin) - out = self.attend(q, k, self.split_kv(self.value(x)), mask) - if gated: - out = self.gate_output(out, self.split_gate(qg)) - return self.output(out) - - return Attention - - def recurrent(sem): - @runtime_module(sem.math) - class GDN: - @runtime_func - def forward(self, qkv, z, alpha, beta_raw, conv, decay, dt_bias, gamma): - return gdn( - qkv, - z, - alpha, - beta_raw, - conv, - decay, - dt_bias, - gamma, - nk=config["linear_num_key_heads"], - nv=config["linear_num_value_heads"], - dk=config["linear_key_head_dim"], - dv=config["linear_value_head_dim"], - eps=eps, +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 ) - - @runtime_module(sem) - class Recurrent: - qkv = linear(sem.qkv) - gate = linear(sem.gate) - alpha = linear(sem.alpha) - beta = linear(sem.beta) - output = linear(sem.output) - math = GDN - - @runtime_func - def forward(self, x, cos, sin, mask): - return self.output( - self.math(self.qkv(x), self.gate(x), self.alpha(x), self.beta(x)) - ) - - return Recurrent - - def layer(sem, kind): - @runtime_module(sem) - class Layer: - norm_in = make_norm(sem.norm_in) - norm_ffn = make_norm(sem.norm_ffn) - mlp = make_mlp(sem.mlp) - mixer = ( - attention(sem.mixer) - if kind == "full_attention" - else recurrent(sem.mixer) ) - - @runtime_func - def add_mixer(self, x, y): - return x + y - - @runtime_func - def add_mlp(self, x, y): - return x + y - - @runtime_func - def forward(self, x, cos, sin, mask): - x = self.add_mixer(x, self.mixer(self.norm_in(x), cos, sin, mask)) - return self.add_mlp(x, self.mlp(self.norm_ffn(x))) - - return Layer - - kinds = config.get("layer_types", ["full_attention"] * config["num_hidden_layers"]) - definitions = { - kind: layer(semantic.layers[kinds.index(kind)], kind) for kind in set(kinds) - } - - @runtime_module(semantic) - class QwenPrefill: - layers = ModuleList(tuple(definitions[kind] for kind in kinds)) - norm = make_norm(semantic.norm) - head = linear(semantic.head, head=True) - - @runtime_func - def embed(self, tokens, table): - return table[tokens].float() - - @runtime_func - def last(self, x): - return x[-1:] - - @runtime_func - def prefill(self, tokens, cos, sin, mask, table): - hidden = self.embed(tokens) - for child in self.layers: - hidden = child(hidden, cos, sin, mask) - return self.head(self.last(self.norm(hidden))) - - return QwenPrefill - - -def inputs(config, sequence, tokens): - r = int( - config["head_dim"] - * config.get("rope_parameters", {}).get("partial_rotary_factor", 1) - ) - theta = config.get("rope_parameters", {}).get( - "rope_theta", config.get("rope_theta", 1000000) - ) - inv = 1 / (theta ** (torch.arange(0, r, 2, dtype=torch.float32) / r)) - freq = torch.arange(sequence, dtype=torch.float32)[:, None] * inv[None, :] - freq = torch.cat([freq, freq], dim=-1) - mask = torch.full((sequence, sequence), float("-inf")).triu(1) - return tokens.long(), freq.cos(), freq.sin(), mask + 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 index 5bbe1da9..b04a528b 100644 --- a/examples/qwen_prefill-amx/scheduled.py +++ b/examples/qwen_prefill-amx/scheduled.py @@ -6,6 +6,12 @@ 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 @@ -16,24 +22,24 @@ def scheduled_projection(M, N, K): raise ValueError( "The native AMX runtime requires positive K and N that are multiples of 32" ) - B = 64 if N % 64 == 0 else 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 == 1: + if M < 128 or N % 64: @module(entry="forward", topologies=(Topology("core", 8), Topology("amx", 1))) class Head: @func - def forward(x: Tensor[(1, K), "f32"], w: ConstTensor[(K, N), "f16"]): + 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 Head + return cpu_schedule(Head) @module(entry="forward", topologies=(Topology("core", 8), Topology("amx", 1))) class Linear: @@ -42,7 +48,10 @@ 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.cast(x, dtype="f16") + 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) ) @@ -125,8 +134,92 @@ def _(raw: ConstTensor[(N, K), "f16"]): 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( - build(config, sequence, linear=scheduled_projection), + 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 index fdbeead2..4f5af7a1 100644 --- a/examples/qwen_prefill-amx/weights.py +++ b/examples/qwen_prefill-amx/weights.py @@ -101,3 +101,52 @@ 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/src/tilefoundry/analysis/footprint.py b/src/tilefoundry/analysis/footprint.py index d649e569..77340591 100644 --- a/src/tilefoundry/analysis/footprint.py +++ b/src/tilefoundry/analysis/footprint.py @@ -1187,27 +1187,21 @@ def resolve_footprint(footprint, ctx): continue def project_image(image): + """Keep instance coordinates correlated with addresses during projection.""" + relation = isl.map.from_range(image) for axis in reversed(range(len(item.ownership))): - image = image.set_dim_name(isl.dim_type.SET, axis, f"__tf_instance_{axis}") - image = image.move_dims( - isl.dim_type.PARAM, image.dim(isl.dim_type.PARAM), isl.dim_type.SET, axis, 1 - ) + relation = relation.move_dims(isl.dim_type.IN, 0, isl.dim_type.OUT, axis, 1) source = BufferAlias( item.target, - AccessRelation(isl.map.from_range(image), item.values), + AccessRelation(relation, item.values), AliasKind.FORWARD, ) projected, precision = binding.project(source) relation = projected.relation.relation values = dict(projected.relation.values) - relation = relation.project_out(isl.dim_type.IN, 0, depth) for axis in range(len(item.ownership)): - parameter = relation.find_dim_by_name( - isl.dim_type.PARAM, f"__tf_instance_{axis}" - ) - relation = relation.move_dims( - isl.dim_type.OUT, axis, isl.dim_type.PARAM, parameter, 1 - ) + relation = relation.move_dims(isl.dim_type.OUT, axis, isl.dim_type.IN, depth, 1) + relation = relation.project_out(isl.dim_type.IN, 0, depth) image = relation.range() return projected.target, image, values, precision diff --git a/src/tilefoundry/ir/hir/schedule.py b/src/tilefoundry/ir/hir/schedule.py index 2eee3ecb..9e53f3bd 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)) diff --git a/src/tilefoundry/ir/tir/cpu/__init__.py b/src/tilefoundry/ir/tir/cpu/__init__.py new file mode 100644 index 00000000..50f0db15 --- /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 00000000..e1124f9b --- /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(None), + "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.frame.target if ctx.frame is not None else None + 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/passes/transforms/convert_hir_to_tir.py b/src/tilefoundry/passes/transforms/convert_hir_to_tir.py index b95f657f..0e1d9ded 100644 --- a/src/tilefoundry/passes/transforms/convert_hir_to_tir.py +++ b/src/tilefoundry/passes/transforms/convert_hir_to_tir.py @@ -1618,7 +1618,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/target/amx/device.py b/src/tilefoundry/target/amx/device.py index 6eaa6a72..8f18c453 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 6c93df61..7c6f112f 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_m2_pro.toml b/src/tilefoundry/target/amx/hardware/apple_m2_pro.toml index 083c4097..d52db042 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 a60c402d..ff95273d 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 eb8aaa0b..5c0cc37b 100644 --- a/src/tilefoundry/target/amx/target.py +++ b/src/tilefoundry/target/amx/target.py @@ -144,15 +144,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/visitor_registry/candidates/amx.py b/src/tilefoundry/visitor_registry/candidates/amx.py new file mode 100644 index 00000000..6488a37b --- /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,)) From a9b7ae0d76c89667dbe60f874e9a525c96fcb3bb Mon Sep 17 00:00:00 2001 From: QiHang Zheng <597323109@qq.com> Date: Mon, 19 Oct 2026 15:00:12 +0800 Subject: [PATCH 7/8] fix(amx,examples): address review on analysis and memory scopes --- docs/spec/core-ir.md | 4 + docs/spec/tir.md | 8 +- examples/qwen_prefill-amx/README.md | 8 +- examples/qwen_prefill-amx/model.py | 4 +- examples/qwen_prefill-amx/weights.py | 171 +++++++++++++++++- src/tilefoundry/analysis/footprint.py | 18 +- src/tilefoundry/inspection/pattern_printer.py | 7 +- src/tilefoundry/ir/pattern/match.py | 24 ++- src/tilefoundry/ir/pattern/pattern.py | 6 +- src/tilefoundry/ir/pattern/utils.py | 21 +-- src/tilefoundry/ir/tir/memory/copy.py | 10 +- src/tilefoundry/ir/tir/memory/fill.py | 8 +- 12 files changed, 243 insertions(+), 46 deletions(-) diff --git a/docs/spec/core-ir.md b/docs/spec/core-ir.md index 585e4041..dd499f57 100644 --- a/docs/spec/core-ir.md +++ b/docs/spec/core-ir.md @@ -765,6 +765,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/tir.md b/docs/spec/tir.md index eeae9b22..1e8ccf7d 100644 --- a/docs/spec/tir.md +++ b/docs/spec/tir.md @@ -599,10 +599,10 @@ 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. -Copy and Fill admit CUDA `thread`, CPU `core`, and Apple `amx` execution -scopes through the shared memory execution pattern. A sharded register operand -must name the selected scope. Their canonical CUDA `execution_mesh` descriptor -remains available; admission uses `execution_mesh_pattern()`. +Copy and Fill bind the current innermost execution level through a shared mesh +wildcard. 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 diff --git a/examples/qwen_prefill-amx/README.md b/examples/qwen_prefill-amx/README.md index 33db2422..e012bc71 100644 --- a/examples/qwen_prefill-amx/README.md +++ b/examples/qwen_prefill-amx/README.md @@ -78,6 +78,8 @@ beside `run.py`, builds the supplied native libraries and FFI entry under `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. @@ -124,7 +126,9 @@ with QwenPrefill(config("1.7b"), sequence=128) as runtime: `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. Each source file also exposes `Qwen3_1_7B` and `Qwen3_5_4B` as +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 @@ -148,6 +152,8 @@ model descriptions and do not require weights or the native extension. 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 diff --git a/examples/qwen_prefill-amx/model.py b/examples/qwen_prefill-amx/model.py index 17d80d40..ba527ecf 100644 --- a/examples/qwen_prefill-amx/model.py +++ b/examples/qwen_prefill-amx/model.py @@ -442,7 +442,9 @@ def last(x: Tensor[(S, H), "f32"]): V=V, ) exec(compile(source, str(path), "exec"), namespace) - return namespace["QwenPrefill"] + from weights import explicit_weights + + return explicit_weights(namespace["QwenPrefill"]) from config import config diff --git a/examples/qwen_prefill-amx/weights.py b/examples/qwen_prefill-amx/weights.py index 4f5af7a1..22625f7e 100644 --- a/examples/qwen_prefill-amx/weights.py +++ b/examples/qwen_prefill-amx/weights.py @@ -6,6 +6,175 @@ 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.call_frame import CallFrame, CallFrameVisitor + + 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 = CallFrameVisitor.callee_frame(call, CallFrame(mod, None)) + 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( + CallFrameVisitor.callee_frame( + call, CallFrame(mod, fn) + ).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 = CallFrameVisitor.callee_frame( + call, CallFrame(mod, fn) + ).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", @@ -51,7 +220,7 @@ def aliases(config): ("gamma", "ssm_norm.weight"), ): out[p + ".mixer.math." + name] = b + "." + tensor - return out + return {key.replace(".", "__"): value for key, value in out.items()} class GGUFResource: diff --git a/src/tilefoundry/analysis/footprint.py b/src/tilefoundry/analysis/footprint.py index 77340591..d649e569 100644 --- a/src/tilefoundry/analysis/footprint.py +++ b/src/tilefoundry/analysis/footprint.py @@ -1187,21 +1187,27 @@ def resolve_footprint(footprint, ctx): continue def project_image(image): - """Keep instance coordinates correlated with addresses during projection.""" - relation = isl.map.from_range(image) for axis in reversed(range(len(item.ownership))): - relation = relation.move_dims(isl.dim_type.IN, 0, isl.dim_type.OUT, axis, 1) + image = image.set_dim_name(isl.dim_type.SET, axis, f"__tf_instance_{axis}") + image = image.move_dims( + isl.dim_type.PARAM, image.dim(isl.dim_type.PARAM), isl.dim_type.SET, axis, 1 + ) source = BufferAlias( item.target, - AccessRelation(relation, item.values), + AccessRelation(isl.map.from_range(image), item.values), AliasKind.FORWARD, ) projected, precision = binding.project(source) relation = projected.relation.relation values = dict(projected.relation.values) - for axis in range(len(item.ownership)): - relation = relation.move_dims(isl.dim_type.OUT, axis, isl.dim_type.IN, depth, 1) relation = relation.project_out(isl.dim_type.IN, 0, depth) + for axis in range(len(item.ownership)): + parameter = relation.find_dim_by_name( + isl.dim_type.PARAM, f"__tf_instance_{axis}" + ) + relation = relation.move_dims( + isl.dim_type.OUT, axis, isl.dim_type.PARAM, parameter, 1 + ) image = relation.range() return projected.target, image, values, precision diff --git a/src/tilefoundry/inspection/pattern_printer.py b/src/tilefoundry/inspection/pattern_printer.py index 31990781..07686df2 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/pattern/match.py b/src/tilefoundry/ir/pattern/match.py index 90c14460..d8de1c95 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 686f5285..acfefa2a 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 692c6682..09d219d1 100644 --- a/src/tilefoundry/ir/pattern/utils.py +++ b/src/tilefoundry/ir/pattern/utils.py @@ -114,13 +114,9 @@ def thread_execution_mesh() -> MeshPattern: return MeshPattern(("thread",), _THREAD_EXECUTION_LAYOUT) -def memory_execution_mesh() -> OrPattern: - """Memory statements may execute on CUDA threads or Apple CPU/AMX units.""" - return OrPattern( - thread_execution_mesh(), - MeshPattern(("core",), _THREAD_EXECUTION_LAYOUT), - MeshPattern(("amx",), _THREAD_EXECUTION_LAYOUT), - ) +def memory_execution_mesh() -> MeshPattern: + """Bind a contiguous run at the current innermost topology level.""" + return MeshPattern(WildcardPattern("execution_topologies"), _THREAD_EXECUTION_LAYOUT) def declared_execution_mesh(op_type) -> MeshPattern | OrPattern: @@ -141,14 +137,13 @@ def declared_execution_mesh(op_type) -> MeshPattern | OrPattern: return declared -def _sharded_on(execution_mesh: MeshPattern | OrPattern) -> OrPattern: +def _sharded_on(execution_mesh: MeshPattern) -> OrPattern: """A lexical tile, or a shard explicitly held by *execution_mesh*.""" return OrPattern( - *( - ShardLayoutPattern(layout=WildcardPattern(), attrs=WildcardPattern(), mesh=mesh) - for mesh in ( - execution_mesh.patterns if isinstance(execution_mesh, OrPattern) else (execution_mesh,) - ) + ShardLayoutPattern( + layout=WildcardPattern(), + attrs=WildcardPattern(), + mesh=execution_mesh, ), LayoutPattern(), ) diff --git a/src/tilefoundry/ir/tir/memory/copy.py b/src/tilefoundry/ir/tir/memory/copy.py index 2d78b454..5b88407c 100644 --- a/src/tilefoundry/ir/tir/memory/copy.py +++ b/src/tilefoundry/ir/tir/memory/copy.py @@ -28,21 +28,17 @@ class Copy(Op): capability = OpCapability(None) - execution_mesh = utils.thread_execution_mesh() - - @classmethod - def execution_mesh_pattern(cls): - return utils.memory_execution_mesh() + execution_mesh = utils.memory_execution_mesh() src = ParamDef( kind="input", effect=MemoryEffect.READ, - pattern=utils.operand_tile(0, execution_mesh=utils.memory_execution_mesh()), + pattern=utils.operand_tile(0, execution_mesh=execution_mesh), ) dst = ParamDef( kind="input", effect=MemoryEffect.WRITE, - pattern=utils.operand_tile(1, execution_mesh=utils.memory_execution_mesh()), + pattern=utils.operand_tile(1, execution_mesh=execution_mesh), ) between = (SameConstraint("shape", "src", "dst"), SameConstraint("dtype", "src", "dst")) diff --git a/src/tilefoundry/ir/tir/memory/fill.py b/src/tilefoundry/ir/tir/memory/fill.py index 030ae454..c4ee41ba 100644 --- a/src/tilefoundry/ir/tir/memory/fill.py +++ b/src/tilefoundry/ir/tir/memory/fill.py @@ -22,16 +22,12 @@ class Fill(Op): """Fills ``tensor`` element-wise with ``value`` (rank-0 scalar).""" - execution_mesh = utils.thread_execution_mesh() - - @classmethod - def execution_mesh_pattern(cls): - return utils.memory_execution_mesh() + execution_mesh = utils.memory_execution_mesh() tensor = ParamDef( kind="input", effect=MemoryEffect.WRITE, - pattern=utils.operand_tile(0, execution_mesh=utils.memory_execution_mesh()), + pattern=utils.operand_tile(0, execution_mesh=execution_mesh), ) value = ParamDef(kind="input", effect=MemoryEffect.READ, pattern=is_scalar_tensor()) From 89f208a11af756ac0940fe8ae00df8d697df6378 Mon Sep 17 00:00:00 2001 From: QiHang Zheng <597323109@qq.com> Date: Mon, 19 Oct 2026 15:48:31 +0800 Subject: [PATCH 8/8] fix(schedule): scope memory admission and CPU kernel capabilities --- docs/spec/target.md | 11 ++++++ docs/spec/tir.md | 10 +++-- examples/qwen_prefill-amx/weights.py | 12 ++---- src/tilefoundry/ir/hir/schedule.py | 5 ++- src/tilefoundry/ir/pattern/utils.py | 18 +++++++-- src/tilefoundry/ir/tir/cpu/tensor.py | 4 +- src/tilefoundry/schedule/facts.py | 4 +- src/tilefoundry/target/__init__.py | 2 + .../target/amx/hardware/apple_amx.toml | 4 +- src/tilefoundry/target/amx/target.py | 6 +++ src/tilefoundry/target/base.py | 4 +- src/tilefoundry/target/cuda/target.py | 5 +++ src/tilefoundry/target/facts.py | 8 ++++ .../visitor_registry/instruction.py | 6 +-- tests/schedule/test_fixtures.py | 38 ++++++++++++++++++- 15 files changed, 108 insertions(+), 29 deletions(-) diff --git a/docs/spec/target.md b/docs/spec/target.md index e623e7c2..10e31e02 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 @@ -836,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 1e8ccf7d..80a0ec15 100644 --- a/docs/spec/tir.md +++ b/docs/spec/tir.md @@ -598,11 +598,13 @@ 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 bind the current innermost execution level through a shared mesh -wildcard. 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 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 diff --git a/examples/qwen_prefill-amx/weights.py b/examples/qwen_prefill-amx/weights.py index 22625f7e..b7320f20 100644 --- a/examples/qwen_prefill-amx/weights.py +++ b/examples/qwen_prefill-amx/weights.py @@ -13,7 +13,7 @@ def explicit_weights(root): 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.call_frame import CallFrame, CallFrameVisitor + from tilefoundry.visitor_registry.contexts import ValuationContext cache = {} @@ -39,7 +39,7 @@ def lift(mod): needs, functions = {}, {} def callee(call): - frame = CallFrameVisitor.callee_frame(call, CallFrame(mod, None)) + frame = ValuationContext(module=mod)._callee_context(call) if frame is None: return None if isinstance(call.target, Function): @@ -66,9 +66,7 @@ def needed(fn): needed(target) if not prefix else lift( - CallFrameVisitor.callee_frame( - call, CallFrame(mod, fn) - ).module + ValuationContext(module=mod, function=fn)._callee_context(call).module )[2][id(target)] ) names.extend(prefix + name for name in target_needs) @@ -99,9 +97,7 @@ def visit_Call(self, call, ctx=None): return replace(call, args=args) target, prefix = item if prefix: - child = CallFrameVisitor.callee_frame( - call, CallFrame(mod, fn) - ).module + child = ValuationContext(module=mod, function=fn)._callee_context(call).module child_needs = lift(child)[2][id(target)] target = self.visit(call.target, ctx) else: diff --git a/src/tilefoundry/ir/hir/schedule.py b/src/tilefoundry/ir/hir/schedule.py index 5d55f93c..729dd2f2 100644 --- a/src/tilefoundry/ir/hir/schedule.py +++ b/src/tilefoundry/ir/hir/schedule.py @@ -220,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 @@ -245,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 @@ -320,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/utils.py b/src/tilefoundry/ir/pattern/utils.py index 09d219d1..bff557af 100644 --- a/src/tilefoundry/ir/pattern/utils.py +++ b/src/tilefoundry/ir/pattern/utils.py @@ -116,11 +116,11 @@ def thread_execution_mesh() -> MeshPattern: def memory_execution_mesh() -> MeshPattern: """Bind a contiguous run at the current innermost topology level.""" - return MeshPattern(WildcardPattern("execution_topologies"), _THREAD_EXECUTION_LAYOUT) + return MeshPattern(WildcardPattern(), _THREAD_EXECUTION_LAYOUT) -def declared_execution_mesh(op_type) -> MeshPattern | OrPattern: - """Return one op's reserved execution-mesh declaration or reject it.""" +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): @@ -134,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/cpu/tensor.py b/src/tilefoundry/ir/tir/cpu/tensor.py index e1124f9b..5124eab3 100644 --- a/src/tilefoundry/ir/tir/cpu/tensor.py +++ b/src/tilefoundry/ir/tir/cpu/tensor.py @@ -84,7 +84,7 @@ def constructor(self, **bound): attrs = { "__init__": constructor, "__module__": __name__, - "capability": OpCapability(None), + "capability": OpCapability("cpu_tensor"), "resource": "cpu_engine", "execution_mesh": CORE, **attributes, @@ -141,7 +141,7 @@ def relation(call, ctx): @register_verify_stmt(cls) def verify(call, ctx): - target = ctx.frame.target if ctx.frame is not None else None + 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()) diff --git a/src/tilefoundry/schedule/facts.py b/src/tilefoundry/schedule/facts.py index 6e7e6270..82561254 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 b08fda7f..df170b4e 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/hardware/apple_amx.toml b/src/tilefoundry/target/amx/hardware/apple_amx.toml index a5e34149..07bf6c53 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/target.py b/src/tilefoundry/target/amx/target.py index 5c0cc37b..0f3bec27 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 diff --git a/src/tilefoundry/target/base.py b/src/tilefoundry/target/base.py index 1b5e77df..af520399 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 c2f1cf39..dbb3b8b3 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 965444e8..680784b9 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/instruction.py b/src/tilefoundry/visitor_registry/instruction.py index 67112961..ca8f20b1 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 10595b38..4254880a 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, ), ) ),