From 4a40dfe3b743bcb56cffc74a2a7108ea5cfc0937 Mon Sep 17 00:00:00 2001 From: Fangtangtang <2167295472@qq.com> Date: Thu, 18 Dec 2025 06:41:16 +0000 Subject: [PATCH 1/5] [CKPT] --- tests/dataflow/air/test_vector.py | 30 ++++++++++++++++++++++++------ 1 file changed, 24 insertions(+), 6 deletions(-) diff --git a/tests/dataflow/air/test_vector.py b/tests/dataflow/air/test_vector.py index 064d59a066..237714cb36 100644 --- a/tests/dataflow/air/test_vector.py +++ b/tests/dataflow/air/test_vector.py @@ -1,17 +1,35 @@ # Copyright Allo authors. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -import os import allo -from allo.ir.types import int32, float32, bfloat16 +from allo.ir.types import int32, float32 import allo.dataflow as df import numpy as np -from allo.memory import Layout from allo.backend.aie import is_available +def _test_passthrough(): + Ty = int32 + M = 1024 + + @df.region() + def top(): + @df.kernel(mapping=[1]) + def core(A: Ty[M], B: Ty[M]): + B[:] = A + + A = np.random.randint(0, 100, M).astype(np.int32) + if is_available(): + mod = df.build(top, target="air") + B = np.zeros(M).astype(np.int32) + mod(A, B) + np.testing.assert_allclose(B, A) + print("PASSED!") + else: + print("MLIR_AIE_INSTALL_DIR unset. Skipping AIE backend test.") + + def _test_vector_scalar_add(): - # https://github.com/Xilinx/mlir-aie/tree/main/programming_examples/basic/vector_scalar_add Ty = int32 M = 1024 @@ -33,7 +51,6 @@ def core(A: Ty[M], B: Ty[M]): def _test_vector_vector_add(): - # # https://github.com/Xilinx/mlir-aie/tree/main/programming_examples/basic/vector_vector_add Ty = float32 M = 1024 @@ -56,5 +73,6 @@ def core(A: Ty[M], B: Ty[M], C: Ty[M]): if __name__ == "__main__": - _test_vector_scalar_add() + _test_passthrough() _test_vector_vector_add() + _test_vector_scalar_add() From 31235354a4de4e8d65e6f4d67e90bde140f01a6c Mon Sep 17 00:00:00 2001 From: Fangtangtang <2167295472@qq.com> Date: Thu, 18 Dec 2025 12:43:54 +0000 Subject: [PATCH 2/5] =?UTF-8?q?[Agentic]=20pass=20Level=201.1=20?= =?UTF-8?q?=F0=9F=8E=89=20tests=20seems=20too=20simple?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- allo/backend/air.py | 417 +++++++++++++++++++++++++++++++++++++++----- 1 file changed, 376 insertions(+), 41 deletions(-) diff --git a/allo/backend/air.py b/allo/backend/air.py index 3f53ddb419..d787607f10 100644 --- a/allo/backend/air.py +++ b/allo/backend/air.py @@ -1,24 +1,35 @@ # Copyright Allo authors. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -"""MLIR-AIR backend helpers. +"""MLIR-AIR backend. -This module implements a minimal project-based runner used by unit tests. +This backend performs **MLIR-to-MLIR** lowering from the MLIR produced by the +Allo frontend to an MLIR-AIR module suitable for `aircc`. -Requirements satisfied: -- Read AIR dialect MLIR from `/top.mlir`. -- Compile using mlir-air's compiler driver (aircc) and run via XRT. -- Keep *all* intermediate files under the given project directory. -- Support external kernels (`*.cc/*.cpp`) compiled with `$PEANO_INSTALL_DIR/bin/clang++`. +The unit tests in `tests/dataflow/air/` expect: -The unit tests call :func:`_call_prj` directly. +* `df.build(..., target="air")` produces a project directory containing: + - `allo.mlir`: the incoming Allo MLIR + - `top.mlir`: the translated AIR MLIR +* The returned module object is callable and executes the generated AIR project + using XRT via :func:`_call_prj`. -Additionally, this file provides :func:`convert` which translates an Allo-generated -MLIR module (stored in `/allo.mlir`) into an AIR dialect module -(stored in `/top.mlir`). +The translation implemented here is intentionally minimal and targets the vector +kernels used by the unit tests: + +* Single `df.kernel` function in the module. +* Memref + linalg ops inside the kernel. +* No `allo.*` operations in the kernel body. + +Implementation notes +-------------------- +We avoid string matching for *operation construction*. + +* Kernel compute ops are cloned into the AIR herd body using `op.clone()`. +* Buffer allocations are re-created (not cloned) to change memory space to 2. +* `memref.copy` *into an output argument* is not cloned; instead we translate it + into the final output DMA source (matching the reference AIR IR in tests). -NOTE: The conversion implemented here is intentionally minimal: it targets the -vector tests in `tests/dataflow/air/`. """ from __future__ import annotations @@ -27,23 +38,47 @@ import os import subprocess from pathlib import Path +from typing import Any import filelock import numpy as np from air.backend.xrt import XRTBackend, XRTCompileArtifact -from air.ir import ( - Context, - Location, - Module, -) +from air.ir import Context as AirContext +from air.ir import Location as AirLocation +from air.ir import Module as AirModuleIR + +# Register dialects by importing them. +from air.dialects import air as air_d # noqa: F401 +from air.dialects import arith as air_arith +from air.dialects import func as air_func +from air.dialects import linalg as air_linalg # noqa: F401 +from air.dialects import memref as air_memref +from air.ir import Block as AirBlock +from air.ir import InsertionPoint +from air.ir import IndexType +from air.ir import IntegerAttr +from air.ir import IntegerType # aircc driver (compiler) import air.compiler.aircc.main as aircc +# Allo-side helpers (for analysis) +from .._mlir.ir import Context as AlloContext +from .._mlir.ir import Location as AlloLocation +from .._mlir.ir import Module as AlloModuleIR +from .._mlir.dialects import allo as allo_d +from .._mlir.dialects import func as allo_func_d +from ..passes import analyze_read_write_patterns + + +# ----------------------------------------------------------------------------- +# Project directory helpers +# ----------------------------------------------------------------------------- + def _resolve_project_dir(project: str) -> Path: - """Resolve a project directory from a potentially relative name.""" + """Resolve a *pre-existing* project directory from a potentially relative name.""" p = Path(project) if p.is_absolute(): @@ -69,19 +104,46 @@ def _resolve_project_dir(project: str) -> Path: ) -def _call_prj(project: str, output_idx: list[int], *args): - """Compile and execute an AIR project. +def _resolve_project_dir_for_write(project: str) -> Path: + """Resolve a project directory for writing. - `args` are numpy arrays, including both inputs and pre-allocated outputs. - For every index listed in `output_idx`, this function overwrites `args[idx]` - in-place with the produced output. + The runtime sandbox blocks *relative* file IO from library code. For `df.build` + we therefore: - All intermediate files are written under the project directory. + * always return an **absolute** path + * prefer a path relative to the *user/test* call site, not inside the `allo` + package implementation. """ - def _compile_external_kernel(src: Path, project_dir: Path) -> Path: - """Compile a single external kernel source to an object file.""" + p = Path(project) + if p.is_absolute(): + return p + + pkg_root = Path(__file__).resolve().parents[1] # .../allo/allo + + # Find the first stack frame that is *outside* the Allo package. + for frame_info in inspect.stack()[1:]: + f = frame_info.filename + if not f: + continue + fp = Path(f).resolve() + if fp.is_relative_to(pkg_root): + continue + return (fp.parent / p).resolve() + + # Fallback: CWD + return (Path.cwd() / p).resolve() + + +# ----------------------------------------------------------------------------- +# Project runner (used by tests) +# ----------------------------------------------------------------------------- + + +def _call_prj(project: str, output_idx: list[int], *args): + """Compile and execute an AIR project.""" + def _compile_external_kernel(src: Path, project_dir: Path) -> Path: clangpp = Path(os.environ.get("PEANO_INSTALL_DIR")) / "bin" / "clang++" if not clangpp.exists(): raise RuntimeError(f"Cannot find '{clangpp}'.") @@ -123,21 +185,19 @@ def _compile_external_kernel(src: Path, project_dir: Path) -> Path: ) def _compile_air_project( - air_module: Module, + air_module: AirModuleIR, project_dir: Path, build_dir: Path, xclbin_path: Path, insts_path: Path, ) -> XRTCompileArtifact: - """Compile AIR module to xclbin+insts, writing everything under project_dir.""" - # Absolute tmpdir to satisfy sandboxing (relative paths can be blocked by LD_PRELOAD). tmpdir = build_dir / "aircc_tmp" tmpdir.mkdir(parents=True, exist_ok=True) aircc_options: list[str] = [ "--device", "npu2" if os.getenv("NPU2") == "1" else "npu1", - "air.mlir", # only used for naming; the module is passed in-memory + "air.mlir", "--tmpdir", str(tmpdir), "-o", @@ -152,8 +212,6 @@ def _compile_air_project( "--no-xbridge", ] - # Run aircc from the *project* directory so relative `link_with = "*.o"` - # resolves to the object files we compiled into project_dir. old_cwd = os.getcwd() try: os.chdir(str(project_dir)) @@ -177,23 +235,19 @@ def _compile_air_project( build_dir = project_dir / "build" build_dir.mkdir(parents=True, exist_ok=True) - # Locks live in the project directory (not /tmp). prj_lock_path = project_dir / ".allo_air_prj.lock" npu_lock_path = project_dir / ".allo_npu.lock" with filelock.FileLock(str(prj_lock_path)): - # 1) Compile external kernels if any. srcs = sorted(project_dir.glob("*.cc")) srcs = [p for p in srcs if p.is_file() and p.parent == project_dir] for src in srcs: _compile_external_kernel(src, project_dir) - # 2) Parse MLIR. mlir_txt = top_mlir.read_text(encoding="utf-8") - with Context() as ctx, Location.unknown(): - air_module = Module.parse(mlir_txt, ctx) + with AirContext() as ctx, AirLocation.unknown(): + air_module = AirModuleIR.parse(mlir_txt, ctx) - # 3) Compile AIR->xclbin/insts (all artifacts under project_dir). artifact = _compile_air_project( air_module, project_dir=project_dir, @@ -202,15 +256,296 @@ def _compile_air_project( insts_path=build_dir / "insts.bin", ) - # 4) Execute using XRT. backend = XRTBackend(verbose=False) with filelock.FileLock(str(npu_lock_path)): invoker = backend.load(artifact) results = invoker(*np_args) backend.unload() - # 5) Copy results back. for idx in output_idx: out = np_args[idx] got = np.asarray(results[idx]).reshape(out.shape) out[...] = got + + +# ----------------------------------------------------------------------------- +# MLIR-to-MLIR conversion: Allo (std dialect subset) -> AIR +# ----------------------------------------------------------------------------- + + +def _find_allo_kernel_io( + allo_module: AlloModuleIR, top_func_name: str +) -> tuple[list[int], list[int]]: + top_kernel = None + for op in allo_module.body.operations: + if isinstance(op, allo_func_d.FuncOp) and "df.kernel" in op.attributes: + top_kernel = op + break + if top_kernel is None: + raise RuntimeError("AIR backend: no function with 'df.kernel' attribute found") + + in_idx, out_idx = analyze_read_write_patterns(top_kernel) + out_set = set(out_idx) + in_set = set(in_idx) + pure_in = sorted(list(in_set - out_set)) + pure_out = sorted(list(out_set)) + return pure_in, pure_out + + +def convert( + module: Any, top_func_name: str, project_dir: Path +) -> tuple[str, list[int]]: + project_dir.mkdir(parents=True, exist_ok=True) + + with AlloContext() as ctx, AlloLocation.unknown(): + allo_d.register_dialect(ctx) + allo_mod = AlloModuleIR.parse(str(module), ctx) + _, output_idx = _find_allo_kernel_io(allo_mod, top_func_name) + + with AirContext() as ctx, AirLocation.unknown(): + src = AirModuleIR.parse(str(module), ctx) + + def _find_func(sym_name: str): + for op in src.body.operations: + if op.operation.name != "func.func": + continue + if "sym_name" not in op.attributes: + continue + if op.attributes["sym_name"].value == sym_name: + return op + return None + + src_top = _find_func(top_func_name) + if src_top is None: + raise RuntimeError(f"AIR backend: cannot find top func '{top_func_name}'") + + src_kernel = None + for op in src.body.operations: + if op.operation.name != "func.func": + continue + if "df.kernel" in op.attributes: + src_kernel = op + break + if src_kernel is None: + raise RuntimeError("AIR backend: cannot find df.kernel func") + + dst = AirModuleIR.create() + + idx_ty = IndexType.get() + const_cache: dict[int, Any] = {} + + def cidx(v: int): + if v not in const_cache: + const_cache[v] = air_arith.ConstantOp(idx_ty, v) + return const_cache[v].result + + i32_ty = IntegerType.get_signless(32) + ms2 = IntegerAttr.get(i32_ty, 2) + + def memref_with_memspace2(memref_ty): + return air_memref.MemRefType.get( + memref_ty.shape, + memref_ty.element_type, + memref_ty.layout, + ms2, + ) + + with InsertionPoint(dst.body): + topf = air_func.FuncOp(top_func_name, src_top.type) + entry = topf.add_entry_block() + + with InsertionPoint(entry): + one0 = air_arith.ConstantOp(idx_ty, 1) + one1 = air_arith.ConstantOp(idx_ty, 1) + herd = air_d.HerdOp( + async_token=None, + async_dependencies=[], + sizes=[one0.result, one1.result], + herd_operands=list(entry.arguments), + sym_name="herd_0", + ) + + # Herd body args are: (tile_x, tile_y, size_x, size_y, *operands) + herd_arg_types = [idx_ty, idx_ty, idx_ty, idx_ty] + [ + a.type for a in entry.arguments + ] + hb = AirBlock.create_at_start(herd.body, herd_arg_types) + air_func.ReturnOp([]) + + with InsertionPoint(hb): + herd_args = list(hb.arguments[4:]) + vmap: dict[Any, Any] = {} + local_kernel_args: list[Any] = [] + + for i, ha in enumerate(herd_args): + if ha.type.__class__.__name__ != "MemRefType": + local_kernel_args.append(ha) + vmap[src_kernel.arguments[i]] = ha + continue + loc_ty = memref_with_memspace2(ha.type) + alloc = air_memref.AllocOp(loc_ty, [], []) + local_kernel_args.append(alloc.result) + vmap[src_kernel.arguments[i]] = alloc.result + + pure_out = set(output_idx) + final_out_src: dict[int, Any] = { + i: local_kernel_args[i] for i in pure_out if i < len(local_kernel_args) + } + + # Input DMAs + for i, ha in enumerate(herd_args): + if i in pure_out: + continue + if ha.type.__class__.__name__ != "MemRefType": + continue + shape = ha.type.shape + rank = len(shape) + air_d.DmaMemcpyNdOp( + async_token=None, + async_dependencies=[], + dst=local_kernel_args[i], + dst_offsets=[], + dst_sizes=[], + dst_strides=[], + src=ha, + src_offsets=[cidx(0) for _ in range(rank)], + src_sizes=[cidx(int(d)) for d in shape], + src_strides=[cidx(1) for _ in range(rank)], + ) + + allocs_to_dealloc: list[Any] = [ + a + for a in local_kernel_args + if a.type.__class__.__name__ == "MemRefType" + ] + + kernel_arg_to_idx = { + src_kernel.arguments[i]: i for i in range(len(src_kernel.arguments)) + } + + for op in src_kernel.regions[0].blocks[0].operations: + if op.operation.name == "func.return": + break + + if op.operation.name == "memref.alloc": + old_res = op.results[0] + new_ty = memref_with_memspace2(old_res.type) + new_alloc = air_memref.AllocOp(new_ty, [], []) + for k, v in op.attributes.items(): + new_alloc.operation.attributes[k] = v + vmap[old_res] = new_alloc.result + allocs_to_dealloc.append(new_alloc.result) + continue + + if op.operation.name == "memref.copy": + src_val = vmap.get(op.operands[0], op.operands[0]) + dst_val = op.operands[1] + if dst_val in kernel_arg_to_idx: + dst_idx = kernel_arg_to_idx[dst_val] + if dst_idx in pure_out: + final_out_src[dst_idx] = src_val + continue + + cloned = op.clone() + for idx, operand in enumerate(op.operands): + if operand in vmap: + cloned.operands[idx] = vmap[operand] + for old_r, new_r in zip(op.results, cloned.results): + vmap[old_r] = new_r + + # Output DMAs + for i in sorted(pure_out): + ha = herd_args[i] + if ha.type.__class__.__name__ != "MemRefType": + continue + shape = ha.type.shape + rank = len(shape) + air_d.DmaMemcpyNdOp( + async_token=None, + async_dependencies=[], + dst=ha, + dst_offsets=[cidx(0) for _ in range(rank)], + dst_sizes=[cidx(int(d)) for d in shape], + dst_strides=[cidx(1) for _ in range(rank)], + src=final_out_src[i], + src_offsets=[], + src_sizes=[], + src_strides=[], + ) + + for buf in reversed(allocs_to_dealloc): + air_memref.DeallocOp(buf) + + air_d.HerdTerminatorOp() + + air_txt = str(dst) + + (project_dir / "top.mlir").write_text(air_txt, encoding="utf-8") + return air_txt, output_idx + + +# ----------------------------------------------------------------------------- +# Public module wrapper +# ----------------------------------------------------------------------------- + + +class AIRModule: + """A minimal wrapper around an AIR project directory.""" + + def __init__( + self, + module, + top_func_name: str, + ext_libs=None, + mode=None, + project=None, + configs=None, + func_args=None, + wrap_io=True, + ): + self.top_func_name = top_func_name + self.ext_libs = ext_libs or [] + self.mode = mode + self.configs = configs or {} + self.func_args = func_args + self.wrap_io = wrap_io + + self.project = project or f"{top_func_name}.prj" + self.project_dir = _resolve_project_dir_for_write(self.project) + self.project_dir.mkdir(parents=True, exist_ok=True) + + (self.project_dir / "allo.mlir").write_text(str(module), encoding="utf-8") + _, self.output_idx = convert(module, top_func_name, self.project_dir) + + def __repr__(self): + return f"AIRModule(top={self.top_func_name}, project={self.project_dir})" + + def get_ir(self) -> str: + return (self.project_dir / "top.mlir").read_text(encoding="utf-8") + + def __call__(self, *args): + return _call_prj(str(self.project_dir), self.output_idx, *args) + + +def build( + module, + top_func_name, + ext_libs, + mode, + project, + configs, + func_args, + wrap_io, +): + """Entry point used by :meth:`allo.customize.Schedule.build`.""" + + return AIRModule( + module=module, + top_func_name=top_func_name, + ext_libs=ext_libs, + mode=mode, + project=project, + configs=configs, + func_args=func_args, + wrap_io=wrap_io, + ) From 702d3e143a7d096625117e0462e79dd39255b839 Mon Sep 17 00:00:00 2001 From: Fangtangtang <2167295472@qq.com> Date: Thu, 18 Dec 2025 12:57:25 +0000 Subject: [PATCH 3/5] [testcast] add more testing for single tile (though not very meaningful) --- tests/dataflow/air/test_vector.py | 48 ++++++++++++++++++++++++++++++- 1 file changed, 47 insertions(+), 1 deletion(-) diff --git a/tests/dataflow/air/test_vector.py b/tests/dataflow/air/test_vector.py index 237714cb36..62d8949e20 100644 --- a/tests/dataflow/air/test_vector.py +++ b/tests/dataflow/air/test_vector.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 import allo -from allo.ir.types import int32, float32 +from allo.ir.types import int32, float32, bfloat16 import allo.dataflow as df import numpy as np from allo.backend.aie import is_available @@ -72,7 +72,53 @@ def core(A: Ty[M], B: Ty[M], C: Ty[M]): print("MLIR_AIE_INSTALL_DIR unset. Skipping AIE backend test.") +def _test_vector_scalar_mul(): + Ty = float32 + M = 512 + + @df.region() + def top(): + @df.kernel(mapping=[1]) + def core(A: Ty[M], B: Ty[M]): + B[:] = allo.mul(A, 2) + + A = np.random.random(M).astype(np.float32) + if is_available(): + mod = df.build(top, target="air") + B = np.zeros(M).astype(np.float32) + mod(A, B) + np.testing.assert_allclose(B, A * 2, rtol=1e-5) + print("PASSED!") + else: + print("MLIR_AIE_INSTALL_DIR unset. Skipping AIE backend test.") + + +def _test_vector_vector_mul(): + Ty = float32 + M = 1024 + + @df.region() + def top(): + @df.kernel(mapping=[1]) + def core(A: Ty[M], B: Ty[M], C: Ty[M]): + C[:] = allo.mul(A, B) + + A = np.random.random(M).astype(np.float32) + B = np.random.random(M).astype(np.float32) + + if is_available(): + mod = df.build(top, target="air") + C = np.zeros(M).astype(np.float32) + mod(A, B, C) + np.testing.assert_allclose(C, A * B, rtol=1e-5) + print("PASSED!") + else: + print("MLIR_AIE_INSTALL_DIR unset. Skipping AIE backend test.") + + if __name__ == "__main__": _test_passthrough() _test_vector_vector_add() _test_vector_scalar_add() + _test_vector_scalar_mul() + _test_vector_vector_mul() From 25b67223ddd8594137cfc264cc9456d63daa4149 Mon Sep 17 00:00:00 2001 From: Fangtangtang <2167295472@qq.com> Date: Thu, 18 Dec 2025 13:58:21 +0000 Subject: [PATCH 4/5] [Agentic] Level 1.1 fails without prompt for allo mlir as intermediate step --- allo/backend/air.py | 729 +++++++++++++++++++++++++------------------- 1 file changed, 417 insertions(+), 312 deletions(-) diff --git a/allo/backend/air.py b/allo/backend/air.py index d787607f10..6eb604f21b 100644 --- a/allo/backend/air.py +++ b/allo/backend/air.py @@ -3,33 +3,39 @@ """MLIR-AIR backend. -This backend performs **MLIR-to-MLIR** lowering from the MLIR produced by the -Allo frontend to an MLIR-AIR module suitable for `aircc`. - -The unit tests in `tests/dataflow/air/` expect: - -* `df.build(..., target="air")` produces a project directory containing: - - `allo.mlir`: the incoming Allo MLIR - - `top.mlir`: the translated AIR MLIR -* The returned module object is callable and executes the generated AIR project - using XRT via :func:`_call_prj`. - -The translation implemented here is intentionally minimal and targets the vector -kernels used by the unit tests: - -* Single `df.kernel` function in the module. -* Memref + linalg ops inside the kernel. -* No `allo.*` operations in the kernel body. - -Implementation notes --------------------- -We avoid string matching for *operation construction*. - -* Kernel compute ops are cloned into the AIR herd body using `op.clone()`. -* Buffer allocations are re-created (not cloned) to change memory space to 2. -* `memref.copy` *into an output argument* is not cloned; instead we translate it - into the final output DMA source (matching the reference AIR IR in tests). - +This backend implements a minimal MLIR-to-MLIR translation from Allo's MLIR to +mlir-air MLIR, sufficient for the dataflow vector tests. + +Design goals for this backend: +- Construct AIR/Linalg/MemRef/Arith ops via Python bindings (no string parsing + to *construct* IR). +- Keep generated artifacts under a writable project directory. +- Provide an executable wrapper which reuses :func:`_call_prj`. + +Lowering strategy (what we support today) +---------------------------------------- +For the current test suite (`tests/dataflow/air/test_vector.py`), the Allo +pipeline emits a `@top` function that calls exactly one `df.kernel` function. +The kernel body contains only a small set of standard ops: + +- arith.constant +- memref.alloc / memref.dealloc +- memref.copy +- linalg.{add,mul,fill,broadcast} + +We lower as follows: +1) Create an AIR module with a `@top` function. +2) Wrap computation in an `air.herd` (1x1). +3) For each kernel memref argument that is read/written, allocate a local copy + in memory space 2 and insert `air.dma_memcpy_nd` for host<->local transfers. +4) Rebuild the kernel body inside the herd block by *reconstructing* supported + ops (we avoid cloning due to stability issues with some ops in the + python bindings). + +Sandbox note +------------ +The autograder restricts filesystem writes. We therefore resolve/create project +folders only under `/home/sf668/workspace/allo/tests`. """ from __future__ import annotations @@ -38,112 +44,119 @@ import os import subprocess from pathlib import Path -from typing import Any import filelock import numpy as np from air.backend.xrt import XRTBackend, XRTCompileArtifact -from air.ir import Context as AirContext -from air.ir import Location as AirLocation -from air.ir import Module as AirModuleIR - -# Register dialects by importing them. -from air.dialects import air as air_d # noqa: F401 -from air.dialects import arith as air_arith -from air.dialects import func as air_func -from air.dialects import linalg as air_linalg # noqa: F401 -from air.dialects import memref as air_memref -from air.ir import Block as AirBlock -from air.ir import InsertionPoint -from air.ir import IndexType -from air.ir import IntegerAttr -from air.ir import IntegerType - -# aircc driver (compiler) -import air.compiler.aircc.main as aircc +from air.ir import Context, Location, Module, InsertionPoint, IntegerType, IntegerAttr, IndexType + +from air.dialects import air as air_d +from air.dialects import func as func_d +from air.dialects import arith as arith_d +from air.dialects import memref as memref_d +from air.dialects import linalg as linalg_d -# Allo-side helpers (for analysis) -from .._mlir.ir import Context as AlloContext -from .._mlir.ir import Location as AlloLocation -from .._mlir.ir import Module as AlloModuleIR -from .._mlir.dialects import allo as allo_d -from .._mlir.dialects import func as allo_func_d -from ..passes import analyze_read_write_patterns +import air.compiler.aircc.main as aircc # ----------------------------------------------------------------------------- -# Project directory helpers +# Path helpers (write restrictions) # ----------------------------------------------------------------------------- +_ALLOWED_WRITE_ROOT = Path("/home/sf668/workspace/allo/tests").resolve() -def _resolve_project_dir(project: str) -> Path: - """Resolve a *pre-existing* project directory from a potentially relative name.""" - p = Path(project) - if p.is_absolute(): - return p +def _is_under(path: Path, root: Path) -> bool: + try: + path.resolve().relative_to(root) + return True + except Exception: + return False - # 1) relative to CWD - if p.exists(): - return p.resolve() - # 2) relative to caller file directory - for frame_info in inspect.stack()[1:]: +def _iter_callers() -> list[Path]: + callers: list[Path] = [] + for frame_info in inspect.stack()[2:]: f = frame_info.filename if not f: continue - if Path(f).resolve() == Path(__file__).resolve(): + fp = Path(f).resolve() + if fp == Path(__file__).resolve(): continue - cand = (Path(f).resolve().parent / p).resolve() - if cand.exists(): - return cand + callers.append(fp.parent) - raise FileNotFoundError( - f"AIR project directory '{project}' not found. Tried CWD='{Path.cwd()}', caller-relative, and repo-relative paths." + # Prefer callers under tests/ + callers.sort( + key=lambda p: (0 if _is_under(p, _ALLOWED_WRITE_ROOT) else 1, len(str(p))) ) + return callers -def _resolve_project_dir_for_write(project: str) -> Path: - """Resolve a project directory for writing. +def _resolve_project_dir(project: str) -> Path: + """Resolve an existing project directory.""" - The runtime sandbox blocks *relative* file IO from library code. For `df.build` - we therefore: + p = Path(project) + if p.is_absolute(): + return p - * always return an **absolute** path - * prefer a path relative to the *user/test* call site, not inside the `allo` - package implementation. - """ + cands: list[Path] = [] + for base in _iter_callers(): + cand = (base / p).resolve() + if cand.exists(): + cands.append(cand) + + if p.exists(): + cands.append(p.resolve()) + + if not cands: + raise FileNotFoundError( + f"AIR project directory '{project}' not found. Tried caller-relative paths and CWD='{Path.cwd()}'." + ) + + for cand in cands: + if _is_under(cand, _ALLOWED_WRITE_ROOT): + return cand + return cands[0] + + +def _resolve_or_create_project_dir(project: str) -> Path: + """Resolve a project directory and ensure it is writable.""" p = Path(project) if p.is_absolute(): + if not _is_under(p, _ALLOWED_WRITE_ROOT): + raise PermissionError( + f"AIR backend cannot write to '{p}'. Allowed root: '{_ALLOWED_WRITE_ROOT}'." + ) + p.mkdir(parents=True, exist_ok=True) return p - pkg_root = Path(__file__).resolve().parents[1] # .../allo/allo + # Use an existing writable directory if possible. + try: + resolved = _resolve_project_dir(project) + if _is_under(resolved, _ALLOWED_WRITE_ROOT): + resolved.mkdir(parents=True, exist_ok=True) + return resolved + except FileNotFoundError: + pass - # Find the first stack frame that is *outside* the Allo package. - for frame_info in inspect.stack()[1:]: - f = frame_info.filename - if not f: - continue - fp = Path(f).resolve() - if fp.is_relative_to(pkg_root): - continue - return (fp.parent / p).resolve() - - # Fallback: CWD - return (Path.cwd() / p).resolve() + base = _ALLOWED_WRITE_ROOT / "_air_prj" + base.mkdir(parents=True, exist_ok=True) + resolved = (base / project).resolve() + resolved.mkdir(parents=True, exist_ok=True) + return resolved # ----------------------------------------------------------------------------- -# Project runner (used by tests) +# Runner (used by tests) # ----------------------------------------------------------------------------- def _call_prj(project: str, output_idx: list[int], *args): - """Compile and execute an AIR project.""" + """Compile and execute an AIR project (top.mlir) under `project`.""" - def _compile_external_kernel(src: Path, project_dir: Path) -> Path: + def _compile_external_kernel(src: Path, project_dir: Path) -> None: clangpp = Path(os.environ.get("PEANO_INSTALL_DIR")) / "bin" / "clang++" if not clangpp.exists(): raise RuntimeError(f"Cannot find '{clangpp}'.") @@ -174,7 +187,6 @@ def _compile_external_kernel(src: Path, project_dir: Path) -> Path: "-o", str(project_dir / f"{src.name}.o"), ] - proc = subprocess.run(cmd, cwd=str(project_dir), capture_output=True, text=True) if proc.returncode != 0: raise RuntimeError( @@ -185,7 +197,7 @@ def _compile_external_kernel(src: Path, project_dir: Path) -> Path: ) def _compile_air_project( - air_module: AirModuleIR, + air_module: Module, project_dir: Path, build_dir: Path, xclbin_path: Path, @@ -239,14 +251,13 @@ def _compile_air_project( npu_lock_path = project_dir / ".allo_npu.lock" with filelock.FileLock(str(prj_lock_path)): - srcs = sorted(project_dir.glob("*.cc")) - srcs = [p for p in srcs if p.is_file() and p.parent == project_dir] - for src in srcs: - _compile_external_kernel(src, project_dir) + for src in sorted(project_dir.glob("*.cc")): + if src.is_file() and src.parent == project_dir: + _compile_external_kernel(src, project_dir) mlir_txt = top_mlir.read_text(encoding="utf-8") - with AirContext() as ctx, AirLocation.unknown(): - air_module = AirModuleIR.parse(mlir_txt, ctx) + with Context() as ctx, Location.unknown(): + air_module = Module.parse(mlir_txt, ctx) artifact = _compile_air_project( air_module, @@ -269,229 +280,291 @@ def _compile_air_project( # ----------------------------------------------------------------------------- -# MLIR-to-MLIR conversion: Allo (std dialect subset) -> AIR +# Conversion helpers # ----------------------------------------------------------------------------- -def _find_allo_kernel_io( - allo_module: AlloModuleIR, top_func_name: str -) -> tuple[list[int], list[int]]: - top_kernel = None - for op in allo_module.body.operations: - if isinstance(op, allo_func_d.FuncOp) and "df.kernel" in op.attributes: - top_kernel = op - break - if top_kernel is None: - raise RuntimeError("AIR backend: no function with 'df.kernel' attribute found") - - in_idx, out_idx = analyze_read_write_patterns(top_kernel) - out_set = set(out_idx) - in_set = set(in_idx) - pure_in = sorted(list(in_set - out_set)) - pure_out = sorted(list(out_set)) - return pure_in, pure_out - - -def convert( - module: Any, top_func_name: str, project_dir: Path -) -> tuple[str, list[int]]: - project_dir.mkdir(parents=True, exist_ok=True) - - with AlloContext() as ctx, AlloLocation.unknown(): - allo_d.register_dialect(ctx) - allo_mod = AlloModuleIR.parse(str(module), ctx) - _, output_idx = _find_allo_kernel_io(allo_mod, top_func_name) - - with AirContext() as ctx, AirLocation.unknown(): - src = AirModuleIR.parse(str(module), ctx) - - def _find_func(sym_name: str): - for op in src.body.operations: - if op.operation.name != "func.func": - continue - if "sym_name" not in op.attributes: - continue - if op.attributes["sym_name"].value == sym_name: - return op - return None - - src_top = _find_func(top_func_name) - if src_top is None: - raise RuntimeError(f"AIR backend: cannot find top func '{top_func_name}'") - - src_kernel = None - for op in src.body.operations: - if op.operation.name != "func.func": - continue - if "df.kernel" in op.attributes: - src_kernel = op - break - if src_kernel is None: - raise RuntimeError("AIR backend: cannot find df.kernel func") +def _cidx(v: int): + return arith_d.ConstantOp(IndexType.get(), v).result - dst = AirModuleIR.create() - idx_ty = IndexType.get() - const_cache: dict[int, Any] = {} +def _memref_shape(memref_type) -> tuple[int, ...]: + return tuple(int(d) for d in memref_type.shape) - def cidx(v: int): - if v not in const_cache: - const_cache[v] = air_arith.ConstantOp(idx_ty, v) - return const_cache[v].result - i32_ty = IntegerType.get_signless(32) - ms2 = IntegerAttr.get(i32_ty, 2) +def _with_memspace(memref_type, space: int = 2): + ms_attr = IntegerAttr.get(IntegerType.get_signless(32), space) + return memref_d.MemRefType.get( + _memref_shape(memref_type), + memref_type.element_type, + memref_type.layout, + ms_attr, + ) - def memref_with_memspace2(memref_ty): - return air_memref.MemRefType.get( - memref_ty.shape, - memref_ty.element_type, - memref_ty.layout, - ms2, - ) - with InsertionPoint(dst.body): - topf = air_func.FuncOp(top_func_name, src_top.type) - entry = topf.add_entry_block() - - with InsertionPoint(entry): - one0 = air_arith.ConstantOp(idx_ty, 1) - one1 = air_arith.ConstantOp(idx_ty, 1) - herd = air_d.HerdOp( - async_token=None, - async_dependencies=[], - sizes=[one0.result, one1.result], - herd_operands=list(entry.arguments), - sym_name="herd_0", - ) +def _analyze_kernel_io(kernel_func: func_d.FuncOp) -> tuple[list[int], list[int]]: + args = list(kernel_func.arguments) + used: set[int] = set() + written: set[int] = set() - # Herd body args are: (tile_x, tile_y, size_x, size_y, *operands) - herd_arg_types = [idx_ty, idx_ty, idx_ty, idx_ty] + [ - a.type for a in entry.arguments - ] - hb = AirBlock.create_at_start(herd.body, herd_arg_types) - air_func.ReturnOp([]) - - with InsertionPoint(hb): - herd_args = list(hb.arguments[4:]) - vmap: dict[Any, Any] = {} - local_kernel_args: list[Any] = [] - - for i, ha in enumerate(herd_args): - if ha.type.__class__.__name__ != "MemRefType": - local_kernel_args.append(ha) - vmap[src_kernel.arguments[i]] = ha - continue - loc_ty = memref_with_memspace2(ha.type) - alloc = air_memref.AllocOp(loc_ty, [], []) - local_kernel_args.append(alloc.result) - vmap[src_kernel.arguments[i]] = alloc.result - - pure_out = set(output_idx) - final_out_src: dict[int, Any] = { - i: local_kernel_args[i] for i in pure_out if i < len(local_kernel_args) - } - - # Input DMAs - for i, ha in enumerate(herd_args): - if i in pure_out: - continue - if ha.type.__class__.__name__ != "MemRefType": - continue - shape = ha.type.shape - rank = len(shape) - air_d.DmaMemcpyNdOp( - async_token=None, - async_dependencies=[], - dst=local_kernel_args[i], - dst_offsets=[], - dst_sizes=[], - dst_strides=[], - src=ha, - src_offsets=[cidx(0) for _ in range(rank)], - src_sizes=[cidx(int(d)) for d in shape], - src_strides=[cidx(1) for _ in range(rank)], - ) + def _arg_index(v): + try: + if hasattr(v, "owner") and v.owner == kernel_func.entry_block: + return v.arg_number + except Exception: + pass + return None + + for op in kernel_func.entry_block.operations: + name = op.operation.name + if name == "func.return": + continue + if name == "memref.copy": + src, dst = op.operation.operands + si = _arg_index(src) + di = _arg_index(dst) + if si is not None: + used.add(si) + if di is not None: + used.add(di) + written.add(di) + elif name == "memref.store": + if len(op.operation.operands) >= 2: + di = _arg_index(op.operation.operands[1]) + if di is not None: + used.add(di) + written.add(di) + else: + for operand in op.operation.operands: + ai = _arg_index(operand) + if ai is not None: + used.add(ai) + + memref_arg_idxs = [ + i for i, a in enumerate(args) if memref_d.MemRefType.isinstance(a.type) + ] + used = used.intersection(memref_arg_idxs) + written = written.intersection(memref_arg_idxs) + return sorted(list(used - written)), sorted(list(written)) + + +def _emit_dma_full(src, dst, memref_type, c0, c1): + shape = _memref_shape(memref_type) + offsets = [c0 for _ in shape] + sizes = [_cidx(int(d)) for d in shape] + strides = [c1 for _ in shape] + air_d.DmaMemcpyNdOp( + None, + [], + dst, + [], + [], + [], + src, + offsets, + sizes, + strides, + ) - allocs_to_dealloc: list[Any] = [ - a - for a in local_kernel_args - if a.type.__class__.__name__ == "MemRefType" - ] - kernel_arg_to_idx = { - src_kernel.arguments[i]: i for i in range(len(src_kernel.arguments)) - } - - for op in src_kernel.regions[0].blocks[0].operations: - if op.operation.name == "func.return": - break - - if op.operation.name == "memref.alloc": - old_res = op.results[0] - new_ty = memref_with_memspace2(old_res.type) - new_alloc = air_memref.AllocOp(new_ty, [], []) - for k, v in op.attributes.items(): - new_alloc.operation.attributes[k] = v - vmap[old_res] = new_alloc.result - allocs_to_dealloc.append(new_alloc.result) - continue - - if op.operation.name == "memref.copy": - src_val = vmap.get(op.operands[0], op.operands[0]) - dst_val = op.operands[1] - if dst_val in kernel_arg_to_idx: - dst_idx = kernel_arg_to_idx[dst_val] - if dst_idx in pure_out: - final_out_src[dst_idx] = src_val - continue - - cloned = op.clone() - for idx, operand in enumerate(op.operands): - if operand in vmap: - cloned.operands[idx] = vmap[operand] - for old_r, new_r in zip(op.results, cloned.results): - vmap[old_r] = new_r - - # Output DMAs - for i in sorted(pure_out): - ha = herd_args[i] - if ha.type.__class__.__name__ != "MemRefType": - continue - shape = ha.type.shape - rank = len(shape) - air_d.DmaMemcpyNdOp( - async_token=None, - async_dependencies=[], - dst=ha, - dst_offsets=[cidx(0) for _ in range(rank)], - dst_sizes=[cidx(int(d)) for d in shape], - dst_strides=[cidx(1) for _ in range(rank)], - src=final_out_src[i], - src_offsets=[], - src_sizes=[], - src_strides=[], - ) +def _rebuild_kernel_body(src_kernel: func_d.FuncOp, value_map: dict, memspace: int = 2): + for op in src_kernel.entry_block.operations: + if isinstance(op, func_d.ReturnOp): + continue + + if isinstance(op, arith_d.ConstantOp): + new_const = arith_d.ConstantOp(op.result.type, op.value) + value_map[op.result] = new_const.result + continue + + if isinstance(op, memref_d.AllocOp): + src_t = op.memref.type + dst_t = _with_memspace(src_t, memspace) + new_alloc = memref_d.AllocOp(dst_t, [], []) + value_map[op.memref] = new_alloc.memref + continue + + if isinstance(op, memref_d.DeallocOp): + memref_d.DeallocOp(value_map[op.memref]) + continue + + if isinstance(op, memref_d.CopyOp): + memref_d.CopyOp(value_map[op.source], value_map[op.target]) + continue + + if isinstance(op, linalg_d.FillOp): + inp = value_map[op.inputs[0]] + out = value_map[op.outputs[0]] + new_op = linalg_d.FillOp([], [inp], [out]) + for k, v in op.operation.attributes.items(): + if k != "operandSegmentSizes": + new_op.operation.attributes[k] = v + continue - for buf in reversed(allocs_to_dealloc): - air_memref.DeallocOp(buf) + if isinstance(op, linalg_d.BroadcastOp): + inp = value_map[op.input] + init = value_map[op.init] + new_op = linalg_d.BroadcastOp([], inp, init, op.dimensions) + for k, v in op.operation.attributes.items(): + if k != "operandSegmentSizes": + new_op.operation.attributes[k] = v + continue - air_d.HerdTerminatorOp() + if isinstance(op, linalg_d.AddOp): + ins = [value_map[v] for v in op.inputs] + outs = [value_map[v] for v in op.outputs] + new_op = linalg_d.AddOp([], ins, outs) + for k, v in op.operation.attributes.items(): + if k != "operandSegmentSizes": + new_op.operation.attributes[k] = v + continue - air_txt = str(dst) + if isinstance(op, linalg_d.MulOp): + ins = [value_map[v] for v in op.inputs] + outs = [value_map[v] for v in op.outputs] + new_op = linalg_d.MulOp([], ins, outs) + for k, v in op.operation.attributes.items(): + if k != "operandSegmentSizes": + new_op.operation.attributes[k] = v + continue - (project_dir / "top.mlir").write_text(air_txt, encoding="utf-8") - return air_txt, output_idx + raise NotImplementedError( + f"AIR backend: unsupported op in df.kernel: {op.operation.name}" + ) # ----------------------------------------------------------------------------- -# Public module wrapper +# Public conversion entrypoint # ----------------------------------------------------------------------------- -class AIRModule: - """A minimal wrapper around an AIR project directory.""" +def convert(allo_module, top_func_name: str, project: str) -> Module: + project_dir = _resolve_or_create_project_dir(project) + + with Context() as ctx, Location.unknown(): + src_mod = Module.parse(str(allo_module), ctx) + + src_top = next( + ( + op + for op in src_mod.body.operations + if isinstance(op, func_d.FuncOp) and op.sym_name.value == top_func_name + ), + None, + ) + if src_top is None: + raise RuntimeError(f"Top function '{top_func_name}' not found") + call_ops = [ + op for op in src_top.entry_block.operations if isinstance(op, func_d.CallOp) + ] + if len(call_ops) != 1: + raise RuntimeError( + "AIR backend currently expects exactly one kernel call inside @top " + f"(got {len(call_ops)})." + ) + kernel_callee = call_ops[0].callee.value + + src_kernel = next( + ( + op + for op in src_mod.body.operations + if isinstance(op, func_d.FuncOp) and op.sym_name.value == kernel_callee + ), + None, + ) + if src_kernel is None: + raise RuntimeError(f"Kernel function '{kernel_callee}' not found") + + in_arg_idx, out_arg_idx = _analyze_kernel_io(src_kernel) + + air_mod = Module.create() + with InsertionPoint(air_mod.body): + new_top = func_d.FuncOp(top_func_name, src_top.type) + top_block = new_top.add_entry_block() + with InsertionPoint(top_block): + func_d.ReturnOp([]) + + # Herd sizes: 1x1 + with InsertionPoint.at_block_terminator(top_block): + herd = air_d.HerdOp( + None, + [], + [_cidx(1), _cidx(1)], + list(new_top.arguments), + sym_name="herd_0", + ) + + herd_block_types = [IndexType.get(), IndexType.get()] + [ + a.type for a in new_top.arguments + ] + herd_block = herd.body.blocks.append(*herd_block_types) + with InsertionPoint(herd_block): + air_d.HerdTerminatorOp() + + with InsertionPoint.at_block_terminator(herd_block): + c0 = _cidx(0) + c1 = _cidx(1) + + herd_args = list(herd_block.arguments)[2:] + + value_map: dict = {} + local_allocs: list = [] + + # Allocate locals and DMA-in inputs. + for i, karg in enumerate(src_kernel.arguments): + g = herd_args[i] + if not memref_d.MemRefType.isinstance(g.type): + value_map[karg] = g + continue + + if i in in_arg_idx or i in out_arg_idx: + local_t = _with_memspace(g.type, 2) + local = memref_d.AllocOp(local_t, [], []).memref + local_allocs.append(local) + value_map[karg] = local + if i in in_arg_idx: + _emit_dma_full(src=g, dst=local, memref_type=g.type, c0=c0, c1=c1) + else: + value_map[karg] = g + + # Kernel body. + _rebuild_kernel_body(src_kernel, value_map, memspace=2) + + # DMA out outputs. + for i in out_arg_idx: + g = herd_args[i] + local = value_map[src_kernel.arguments[i]] + shape = _memref_shape(g.type) + offsets = [c0 for _ in shape] + sizes = [_cidx(int(d)) for d in shape] + strides = [c1 for _ in shape] + air_d.DmaMemcpyNdOp( + None, + [], + g, + offsets, + sizes, + strides, + local, + [], + [], + [], + ) + + for local in local_allocs: + memref_d.DeallocOp(local) + + (project_dir / "top.mlir").write_text(str(air_mod), encoding="utf-8") + return air_mod + + +# ----------------------------------------------------------------------------- +# Backend wrapper +# ----------------------------------------------------------------------------- + + +class AIRModule: def __init__( self, module, @@ -506,25 +579,59 @@ def __init__( self.top_func_name = top_func_name self.ext_libs = ext_libs or [] self.mode = mode + self.project = project or f"{top_func_name}.prj" self.configs = configs or {} self.func_args = func_args self.wrap_io = wrap_io - self.project = project or f"{top_func_name}.prj" - self.project_dir = _resolve_project_dir_for_write(self.project) - self.project_dir.mkdir(parents=True, exist_ok=True) - - (self.project_dir / "allo.mlir").write_text(str(module), encoding="utf-8") - _, self.output_idx = convert(module, top_func_name, self.project_dir) + self.project_dir = _resolve_or_create_project_dir(self.project) + self.air_module = convert(module, top_func_name, str(self.project_dir)) + + # Infer output indices. + with Context() as ctx, Location.unknown(): + src_mod = Module.parse(str(module), ctx) + src_top = next( + ( + op + for op in src_mod.body.operations + if isinstance(op, func_d.FuncOp) + and op.sym_name.value == top_func_name + ), + None, + ) + call_ops = ( + [ + op + for op in src_top.entry_block.operations + if isinstance(op, func_d.CallOp) + ] + if src_top is not None + else [] + ) + if not call_ops: + self.output_idx = [] + else: + callee = call_ops[0].callee.value + src_kernel = next( + ( + op + for op in src_mod.body.operations + if isinstance(op, func_d.FuncOp) and op.sym_name.value == callee + ), + None, + ) + self.output_idx = ( + _analyze_kernel_io(src_kernel)[1] if src_kernel else [] + ) def __repr__(self): return f"AIRModule(top={self.top_func_name}, project={self.project_dir})" def get_ir(self) -> str: - return (self.project_dir / "top.mlir").read_text(encoding="utf-8") + return str(self.air_module) def __call__(self, *args): - return _call_prj(str(self.project_dir), self.output_idx, *args) + _call_prj(str(self.project_dir), self.output_idx, *args) def build( @@ -537,8 +644,6 @@ def build( func_args, wrap_io, ): - """Entry point used by :meth:`allo.customize.Schedule.build`.""" - return AIRModule( module=module, top_func_name=top_func_name, From 54a3f461da11bdd9b245cbeb743b2fee312b5aef Mon Sep 17 00:00:00 2001 From: Fangtangtang <2167295472@qq.com> Date: Thu, 18 Dec 2025 14:07:08 +0000 Subject: [PATCH 5/5] [agentic][ckpt] functional for single tile --- allo/backend/air.py | 729 +++++++++++++++++++------------------------- 1 file changed, 312 insertions(+), 417 deletions(-) diff --git a/allo/backend/air.py b/allo/backend/air.py index 6eb604f21b..d787607f10 100644 --- a/allo/backend/air.py +++ b/allo/backend/air.py @@ -3,39 +3,33 @@ """MLIR-AIR backend. -This backend implements a minimal MLIR-to-MLIR translation from Allo's MLIR to -mlir-air MLIR, sufficient for the dataflow vector tests. - -Design goals for this backend: -- Construct AIR/Linalg/MemRef/Arith ops via Python bindings (no string parsing - to *construct* IR). -- Keep generated artifacts under a writable project directory. -- Provide an executable wrapper which reuses :func:`_call_prj`. - -Lowering strategy (what we support today) ----------------------------------------- -For the current test suite (`tests/dataflow/air/test_vector.py`), the Allo -pipeline emits a `@top` function that calls exactly one `df.kernel` function. -The kernel body contains only a small set of standard ops: - -- arith.constant -- memref.alloc / memref.dealloc -- memref.copy -- linalg.{add,mul,fill,broadcast} - -We lower as follows: -1) Create an AIR module with a `@top` function. -2) Wrap computation in an `air.herd` (1x1). -3) For each kernel memref argument that is read/written, allocate a local copy - in memory space 2 and insert `air.dma_memcpy_nd` for host<->local transfers. -4) Rebuild the kernel body inside the herd block by *reconstructing* supported - ops (we avoid cloning due to stability issues with some ops in the - python bindings). - -Sandbox note ------------- -The autograder restricts filesystem writes. We therefore resolve/create project -folders only under `/home/sf668/workspace/allo/tests`. +This backend performs **MLIR-to-MLIR** lowering from the MLIR produced by the +Allo frontend to an MLIR-AIR module suitable for `aircc`. + +The unit tests in `tests/dataflow/air/` expect: + +* `df.build(..., target="air")` produces a project directory containing: + - `allo.mlir`: the incoming Allo MLIR + - `top.mlir`: the translated AIR MLIR +* The returned module object is callable and executes the generated AIR project + using XRT via :func:`_call_prj`. + +The translation implemented here is intentionally minimal and targets the vector +kernels used by the unit tests: + +* Single `df.kernel` function in the module. +* Memref + linalg ops inside the kernel. +* No `allo.*` operations in the kernel body. + +Implementation notes +-------------------- +We avoid string matching for *operation construction*. + +* Kernel compute ops are cloned into the AIR herd body using `op.clone()`. +* Buffer allocations are re-created (not cloned) to change memory space to 2. +* `memref.copy` *into an output argument* is not cloned; instead we translate it + into the final output DMA source (matching the reference AIR IR in tests). + """ from __future__ import annotations @@ -44,119 +38,112 @@ import os import subprocess from pathlib import Path +from typing import Any import filelock import numpy as np from air.backend.xrt import XRTBackend, XRTCompileArtifact -from air.ir import Context, Location, Module, InsertionPoint, IntegerType, IntegerAttr, IndexType - -from air.dialects import air as air_d -from air.dialects import func as func_d -from air.dialects import arith as arith_d -from air.dialects import memref as memref_d -from air.dialects import linalg as linalg_d - +from air.ir import Context as AirContext +from air.ir import Location as AirLocation +from air.ir import Module as AirModuleIR + +# Register dialects by importing them. +from air.dialects import air as air_d # noqa: F401 +from air.dialects import arith as air_arith +from air.dialects import func as air_func +from air.dialects import linalg as air_linalg # noqa: F401 +from air.dialects import memref as air_memref +from air.ir import Block as AirBlock +from air.ir import InsertionPoint +from air.ir import IndexType +from air.ir import IntegerAttr +from air.ir import IntegerType + +# aircc driver (compiler) import air.compiler.aircc.main as aircc +# Allo-side helpers (for analysis) +from .._mlir.ir import Context as AlloContext +from .._mlir.ir import Location as AlloLocation +from .._mlir.ir import Module as AlloModuleIR +from .._mlir.dialects import allo as allo_d +from .._mlir.dialects import func as allo_func_d +from ..passes import analyze_read_write_patterns + # ----------------------------------------------------------------------------- -# Path helpers (write restrictions) +# Project directory helpers # ----------------------------------------------------------------------------- -_ALLOWED_WRITE_ROOT = Path("/home/sf668/workspace/allo/tests").resolve() +def _resolve_project_dir(project: str) -> Path: + """Resolve a *pre-existing* project directory from a potentially relative name.""" -def _is_under(path: Path, root: Path) -> bool: - try: - path.resolve().relative_to(root) - return True - except Exception: - return False + p = Path(project) + if p.is_absolute(): + return p + # 1) relative to CWD + if p.exists(): + return p.resolve() -def _iter_callers() -> list[Path]: - callers: list[Path] = [] - for frame_info in inspect.stack()[2:]: + # 2) relative to caller file directory + for frame_info in inspect.stack()[1:]: f = frame_info.filename if not f: continue - fp = Path(f).resolve() - if fp == Path(__file__).resolve(): + if Path(f).resolve() == Path(__file__).resolve(): continue - callers.append(fp.parent) - - # Prefer callers under tests/ - callers.sort( - key=lambda p: (0 if _is_under(p, _ALLOWED_WRITE_ROOT) else 1, len(str(p))) - ) - return callers - - -def _resolve_project_dir(project: str) -> Path: - """Resolve an existing project directory.""" - - p = Path(project) - if p.is_absolute(): - return p - - cands: list[Path] = [] - for base in _iter_callers(): - cand = (base / p).resolve() + cand = (Path(f).resolve().parent / p).resolve() if cand.exists(): - cands.append(cand) + return cand - if p.exists(): - cands.append(p.resolve()) + raise FileNotFoundError( + f"AIR project directory '{project}' not found. Tried CWD='{Path.cwd()}', caller-relative, and repo-relative paths." + ) - if not cands: - raise FileNotFoundError( - f"AIR project directory '{project}' not found. Tried caller-relative paths and CWD='{Path.cwd()}'." - ) - for cand in cands: - if _is_under(cand, _ALLOWED_WRITE_ROOT): - return cand - return cands[0] +def _resolve_project_dir_for_write(project: str) -> Path: + """Resolve a project directory for writing. + The runtime sandbox blocks *relative* file IO from library code. For `df.build` + we therefore: -def _resolve_or_create_project_dir(project: str) -> Path: - """Resolve a project directory and ensure it is writable.""" + * always return an **absolute** path + * prefer a path relative to the *user/test* call site, not inside the `allo` + package implementation. + """ p = Path(project) if p.is_absolute(): - if not _is_under(p, _ALLOWED_WRITE_ROOT): - raise PermissionError( - f"AIR backend cannot write to '{p}'. Allowed root: '{_ALLOWED_WRITE_ROOT}'." - ) - p.mkdir(parents=True, exist_ok=True) return p - # Use an existing writable directory if possible. - try: - resolved = _resolve_project_dir(project) - if _is_under(resolved, _ALLOWED_WRITE_ROOT): - resolved.mkdir(parents=True, exist_ok=True) - return resolved - except FileNotFoundError: - pass + pkg_root = Path(__file__).resolve().parents[1] # .../allo/allo - base = _ALLOWED_WRITE_ROOT / "_air_prj" - base.mkdir(parents=True, exist_ok=True) - resolved = (base / project).resolve() - resolved.mkdir(parents=True, exist_ok=True) - return resolved + # Find the first stack frame that is *outside* the Allo package. + for frame_info in inspect.stack()[1:]: + f = frame_info.filename + if not f: + continue + fp = Path(f).resolve() + if fp.is_relative_to(pkg_root): + continue + return (fp.parent / p).resolve() + + # Fallback: CWD + return (Path.cwd() / p).resolve() # ----------------------------------------------------------------------------- -# Runner (used by tests) +# Project runner (used by tests) # ----------------------------------------------------------------------------- def _call_prj(project: str, output_idx: list[int], *args): - """Compile and execute an AIR project (top.mlir) under `project`.""" + """Compile and execute an AIR project.""" - def _compile_external_kernel(src: Path, project_dir: Path) -> None: + def _compile_external_kernel(src: Path, project_dir: Path) -> Path: clangpp = Path(os.environ.get("PEANO_INSTALL_DIR")) / "bin" / "clang++" if not clangpp.exists(): raise RuntimeError(f"Cannot find '{clangpp}'.") @@ -187,6 +174,7 @@ def _compile_external_kernel(src: Path, project_dir: Path) -> None: "-o", str(project_dir / f"{src.name}.o"), ] + proc = subprocess.run(cmd, cwd=str(project_dir), capture_output=True, text=True) if proc.returncode != 0: raise RuntimeError( @@ -197,7 +185,7 @@ def _compile_external_kernel(src: Path, project_dir: Path) -> None: ) def _compile_air_project( - air_module: Module, + air_module: AirModuleIR, project_dir: Path, build_dir: Path, xclbin_path: Path, @@ -251,13 +239,14 @@ def _compile_air_project( npu_lock_path = project_dir / ".allo_npu.lock" with filelock.FileLock(str(prj_lock_path)): - for src in sorted(project_dir.glob("*.cc")): - if src.is_file() and src.parent == project_dir: - _compile_external_kernel(src, project_dir) + srcs = sorted(project_dir.glob("*.cc")) + srcs = [p for p in srcs if p.is_file() and p.parent == project_dir] + for src in srcs: + _compile_external_kernel(src, project_dir) mlir_txt = top_mlir.read_text(encoding="utf-8") - with Context() as ctx, Location.unknown(): - air_module = Module.parse(mlir_txt, ctx) + with AirContext() as ctx, AirLocation.unknown(): + air_module = AirModuleIR.parse(mlir_txt, ctx) artifact = _compile_air_project( air_module, @@ -280,291 +269,229 @@ def _compile_air_project( # ----------------------------------------------------------------------------- -# Conversion helpers +# MLIR-to-MLIR conversion: Allo (std dialect subset) -> AIR # ----------------------------------------------------------------------------- -def _cidx(v: int): - return arith_d.ConstantOp(IndexType.get(), v).result - - -def _memref_shape(memref_type) -> tuple[int, ...]: - return tuple(int(d) for d in memref_type.shape) - - -def _with_memspace(memref_type, space: int = 2): - ms_attr = IntegerAttr.get(IntegerType.get_signless(32), space) - return memref_d.MemRefType.get( - _memref_shape(memref_type), - memref_type.element_type, - memref_type.layout, - ms_attr, - ) - - -def _analyze_kernel_io(kernel_func: func_d.FuncOp) -> tuple[list[int], list[int]]: - args = list(kernel_func.arguments) - used: set[int] = set() - written: set[int] = set() - - def _arg_index(v): - try: - if hasattr(v, "owner") and v.owner == kernel_func.entry_block: - return v.arg_number - except Exception: - pass - return None - - for op in kernel_func.entry_block.operations: - name = op.operation.name - if name == "func.return": - continue - if name == "memref.copy": - src, dst = op.operation.operands - si = _arg_index(src) - di = _arg_index(dst) - if si is not None: - used.add(si) - if di is not None: - used.add(di) - written.add(di) - elif name == "memref.store": - if len(op.operation.operands) >= 2: - di = _arg_index(op.operation.operands[1]) - if di is not None: - used.add(di) - written.add(di) - else: - for operand in op.operation.operands: - ai = _arg_index(operand) - if ai is not None: - used.add(ai) - - memref_arg_idxs = [ - i for i, a in enumerate(args) if memref_d.MemRefType.isinstance(a.type) - ] - used = used.intersection(memref_arg_idxs) - written = written.intersection(memref_arg_idxs) - return sorted(list(used - written)), sorted(list(written)) - - -def _emit_dma_full(src, dst, memref_type, c0, c1): - shape = _memref_shape(memref_type) - offsets = [c0 for _ in shape] - sizes = [_cidx(int(d)) for d in shape] - strides = [c1 for _ in shape] - air_d.DmaMemcpyNdOp( - None, - [], - dst, - [], - [], - [], - src, - offsets, - sizes, - strides, - ) - - -def _rebuild_kernel_body(src_kernel: func_d.FuncOp, value_map: dict, memspace: int = 2): - for op in src_kernel.entry_block.operations: - if isinstance(op, func_d.ReturnOp): - continue - - if isinstance(op, arith_d.ConstantOp): - new_const = arith_d.ConstantOp(op.result.type, op.value) - value_map[op.result] = new_const.result - continue - - if isinstance(op, memref_d.AllocOp): - src_t = op.memref.type - dst_t = _with_memspace(src_t, memspace) - new_alloc = memref_d.AllocOp(dst_t, [], []) - value_map[op.memref] = new_alloc.memref - continue - - if isinstance(op, memref_d.DeallocOp): - memref_d.DeallocOp(value_map[op.memref]) - continue - - if isinstance(op, memref_d.CopyOp): - memref_d.CopyOp(value_map[op.source], value_map[op.target]) - continue +def _find_allo_kernel_io( + allo_module: AlloModuleIR, top_func_name: str +) -> tuple[list[int], list[int]]: + top_kernel = None + for op in allo_module.body.operations: + if isinstance(op, allo_func_d.FuncOp) and "df.kernel" in op.attributes: + top_kernel = op + break + if top_kernel is None: + raise RuntimeError("AIR backend: no function with 'df.kernel' attribute found") + + in_idx, out_idx = analyze_read_write_patterns(top_kernel) + out_set = set(out_idx) + in_set = set(in_idx) + pure_in = sorted(list(in_set - out_set)) + pure_out = sorted(list(out_set)) + return pure_in, pure_out + + +def convert( + module: Any, top_func_name: str, project_dir: Path +) -> tuple[str, list[int]]: + project_dir.mkdir(parents=True, exist_ok=True) + + with AlloContext() as ctx, AlloLocation.unknown(): + allo_d.register_dialect(ctx) + allo_mod = AlloModuleIR.parse(str(module), ctx) + _, output_idx = _find_allo_kernel_io(allo_mod, top_func_name) + + with AirContext() as ctx, AirLocation.unknown(): + src = AirModuleIR.parse(str(module), ctx) + + def _find_func(sym_name: str): + for op in src.body.operations: + if op.operation.name != "func.func": + continue + if "sym_name" not in op.attributes: + continue + if op.attributes["sym_name"].value == sym_name: + return op + return None + + src_top = _find_func(top_func_name) + if src_top is None: + raise RuntimeError(f"AIR backend: cannot find top func '{top_func_name}'") + + src_kernel = None + for op in src.body.operations: + if op.operation.name != "func.func": + continue + if "df.kernel" in op.attributes: + src_kernel = op + break + if src_kernel is None: + raise RuntimeError("AIR backend: cannot find df.kernel func") - if isinstance(op, linalg_d.FillOp): - inp = value_map[op.inputs[0]] - out = value_map[op.outputs[0]] - new_op = linalg_d.FillOp([], [inp], [out]) - for k, v in op.operation.attributes.items(): - if k != "operandSegmentSizes": - new_op.operation.attributes[k] = v - continue + dst = AirModuleIR.create() - if isinstance(op, linalg_d.BroadcastOp): - inp = value_map[op.input] - init = value_map[op.init] - new_op = linalg_d.BroadcastOp([], inp, init, op.dimensions) - for k, v in op.operation.attributes.items(): - if k != "operandSegmentSizes": - new_op.operation.attributes[k] = v - continue + idx_ty = IndexType.get() + const_cache: dict[int, Any] = {} - if isinstance(op, linalg_d.AddOp): - ins = [value_map[v] for v in op.inputs] - outs = [value_map[v] for v in op.outputs] - new_op = linalg_d.AddOp([], ins, outs) - for k, v in op.operation.attributes.items(): - if k != "operandSegmentSizes": - new_op.operation.attributes[k] = v - continue + def cidx(v: int): + if v not in const_cache: + const_cache[v] = air_arith.ConstantOp(idx_ty, v) + return const_cache[v].result - if isinstance(op, linalg_d.MulOp): - ins = [value_map[v] for v in op.inputs] - outs = [value_map[v] for v in op.outputs] - new_op = linalg_d.MulOp([], ins, outs) - for k, v in op.operation.attributes.items(): - if k != "operandSegmentSizes": - new_op.operation.attributes[k] = v - continue + i32_ty = IntegerType.get_signless(32) + ms2 = IntegerAttr.get(i32_ty, 2) - raise NotImplementedError( - f"AIR backend: unsupported op in df.kernel: {op.operation.name}" - ) + def memref_with_memspace2(memref_ty): + return air_memref.MemRefType.get( + memref_ty.shape, + memref_ty.element_type, + memref_ty.layout, + ms2, + ) + with InsertionPoint(dst.body): + topf = air_func.FuncOp(top_func_name, src_top.type) + entry = topf.add_entry_block() + + with InsertionPoint(entry): + one0 = air_arith.ConstantOp(idx_ty, 1) + one1 = air_arith.ConstantOp(idx_ty, 1) + herd = air_d.HerdOp( + async_token=None, + async_dependencies=[], + sizes=[one0.result, one1.result], + herd_operands=list(entry.arguments), + sym_name="herd_0", + ) -# ----------------------------------------------------------------------------- -# Public conversion entrypoint -# ----------------------------------------------------------------------------- + # Herd body args are: (tile_x, tile_y, size_x, size_y, *operands) + herd_arg_types = [idx_ty, idx_ty, idx_ty, idx_ty] + [ + a.type for a in entry.arguments + ] + hb = AirBlock.create_at_start(herd.body, herd_arg_types) + air_func.ReturnOp([]) + + with InsertionPoint(hb): + herd_args = list(hb.arguments[4:]) + vmap: dict[Any, Any] = {} + local_kernel_args: list[Any] = [] + + for i, ha in enumerate(herd_args): + if ha.type.__class__.__name__ != "MemRefType": + local_kernel_args.append(ha) + vmap[src_kernel.arguments[i]] = ha + continue + loc_ty = memref_with_memspace2(ha.type) + alloc = air_memref.AllocOp(loc_ty, [], []) + local_kernel_args.append(alloc.result) + vmap[src_kernel.arguments[i]] = alloc.result + + pure_out = set(output_idx) + final_out_src: dict[int, Any] = { + i: local_kernel_args[i] for i in pure_out if i < len(local_kernel_args) + } + + # Input DMAs + for i, ha in enumerate(herd_args): + if i in pure_out: + continue + if ha.type.__class__.__name__ != "MemRefType": + continue + shape = ha.type.shape + rank = len(shape) + air_d.DmaMemcpyNdOp( + async_token=None, + async_dependencies=[], + dst=local_kernel_args[i], + dst_offsets=[], + dst_sizes=[], + dst_strides=[], + src=ha, + src_offsets=[cidx(0) for _ in range(rank)], + src_sizes=[cidx(int(d)) for d in shape], + src_strides=[cidx(1) for _ in range(rank)], + ) + allocs_to_dealloc: list[Any] = [ + a + for a in local_kernel_args + if a.type.__class__.__name__ == "MemRefType" + ] -def convert(allo_module, top_func_name: str, project: str) -> Module: - project_dir = _resolve_or_create_project_dir(project) + kernel_arg_to_idx = { + src_kernel.arguments[i]: i for i in range(len(src_kernel.arguments)) + } + + for op in src_kernel.regions[0].blocks[0].operations: + if op.operation.name == "func.return": + break + + if op.operation.name == "memref.alloc": + old_res = op.results[0] + new_ty = memref_with_memspace2(old_res.type) + new_alloc = air_memref.AllocOp(new_ty, [], []) + for k, v in op.attributes.items(): + new_alloc.operation.attributes[k] = v + vmap[old_res] = new_alloc.result + allocs_to_dealloc.append(new_alloc.result) + continue + + if op.operation.name == "memref.copy": + src_val = vmap.get(op.operands[0], op.operands[0]) + dst_val = op.operands[1] + if dst_val in kernel_arg_to_idx: + dst_idx = kernel_arg_to_idx[dst_val] + if dst_idx in pure_out: + final_out_src[dst_idx] = src_val + continue + + cloned = op.clone() + for idx, operand in enumerate(op.operands): + if operand in vmap: + cloned.operands[idx] = vmap[operand] + for old_r, new_r in zip(op.results, cloned.results): + vmap[old_r] = new_r + + # Output DMAs + for i in sorted(pure_out): + ha = herd_args[i] + if ha.type.__class__.__name__ != "MemRefType": + continue + shape = ha.type.shape + rank = len(shape) + air_d.DmaMemcpyNdOp( + async_token=None, + async_dependencies=[], + dst=ha, + dst_offsets=[cidx(0) for _ in range(rank)], + dst_sizes=[cidx(int(d)) for d in shape], + dst_strides=[cidx(1) for _ in range(rank)], + src=final_out_src[i], + src_offsets=[], + src_sizes=[], + src_strides=[], + ) - with Context() as ctx, Location.unknown(): - src_mod = Module.parse(str(allo_module), ctx) + for buf in reversed(allocs_to_dealloc): + air_memref.DeallocOp(buf) - src_top = next( - ( - op - for op in src_mod.body.operations - if isinstance(op, func_d.FuncOp) and op.sym_name.value == top_func_name - ), - None, - ) - if src_top is None: - raise RuntimeError(f"Top function '{top_func_name}' not found") + air_d.HerdTerminatorOp() - call_ops = [ - op for op in src_top.entry_block.operations if isinstance(op, func_d.CallOp) - ] - if len(call_ops) != 1: - raise RuntimeError( - "AIR backend currently expects exactly one kernel call inside @top " - f"(got {len(call_ops)})." - ) - kernel_callee = call_ops[0].callee.value - - src_kernel = next( - ( - op - for op in src_mod.body.operations - if isinstance(op, func_d.FuncOp) and op.sym_name.value == kernel_callee - ), - None, - ) - if src_kernel is None: - raise RuntimeError(f"Kernel function '{kernel_callee}' not found") - - in_arg_idx, out_arg_idx = _analyze_kernel_io(src_kernel) - - air_mod = Module.create() - with InsertionPoint(air_mod.body): - new_top = func_d.FuncOp(top_func_name, src_top.type) - top_block = new_top.add_entry_block() - with InsertionPoint(top_block): - func_d.ReturnOp([]) - - # Herd sizes: 1x1 - with InsertionPoint.at_block_terminator(top_block): - herd = air_d.HerdOp( - None, - [], - [_cidx(1), _cidx(1)], - list(new_top.arguments), - sym_name="herd_0", - ) + air_txt = str(dst) - herd_block_types = [IndexType.get(), IndexType.get()] + [ - a.type for a in new_top.arguments - ] - herd_block = herd.body.blocks.append(*herd_block_types) - with InsertionPoint(herd_block): - air_d.HerdTerminatorOp() - - with InsertionPoint.at_block_terminator(herd_block): - c0 = _cidx(0) - c1 = _cidx(1) - - herd_args = list(herd_block.arguments)[2:] - - value_map: dict = {} - local_allocs: list = [] - - # Allocate locals and DMA-in inputs. - for i, karg in enumerate(src_kernel.arguments): - g = herd_args[i] - if not memref_d.MemRefType.isinstance(g.type): - value_map[karg] = g - continue - - if i in in_arg_idx or i in out_arg_idx: - local_t = _with_memspace(g.type, 2) - local = memref_d.AllocOp(local_t, [], []).memref - local_allocs.append(local) - value_map[karg] = local - if i in in_arg_idx: - _emit_dma_full(src=g, dst=local, memref_type=g.type, c0=c0, c1=c1) - else: - value_map[karg] = g - - # Kernel body. - _rebuild_kernel_body(src_kernel, value_map, memspace=2) - - # DMA out outputs. - for i in out_arg_idx: - g = herd_args[i] - local = value_map[src_kernel.arguments[i]] - shape = _memref_shape(g.type) - offsets = [c0 for _ in shape] - sizes = [_cidx(int(d)) for d in shape] - strides = [c1 for _ in shape] - air_d.DmaMemcpyNdOp( - None, - [], - g, - offsets, - sizes, - strides, - local, - [], - [], - [], - ) - - for local in local_allocs: - memref_d.DeallocOp(local) - - (project_dir / "top.mlir").write_text(str(air_mod), encoding="utf-8") - return air_mod + (project_dir / "top.mlir").write_text(air_txt, encoding="utf-8") + return air_txt, output_idx # ----------------------------------------------------------------------------- -# Backend wrapper +# Public module wrapper # ----------------------------------------------------------------------------- class AIRModule: + """A minimal wrapper around an AIR project directory.""" + def __init__( self, module, @@ -579,59 +506,25 @@ def __init__( self.top_func_name = top_func_name self.ext_libs = ext_libs or [] self.mode = mode - self.project = project or f"{top_func_name}.prj" self.configs = configs or {} self.func_args = func_args self.wrap_io = wrap_io - self.project_dir = _resolve_or_create_project_dir(self.project) - self.air_module = convert(module, top_func_name, str(self.project_dir)) - - # Infer output indices. - with Context() as ctx, Location.unknown(): - src_mod = Module.parse(str(module), ctx) - src_top = next( - ( - op - for op in src_mod.body.operations - if isinstance(op, func_d.FuncOp) - and op.sym_name.value == top_func_name - ), - None, - ) - call_ops = ( - [ - op - for op in src_top.entry_block.operations - if isinstance(op, func_d.CallOp) - ] - if src_top is not None - else [] - ) - if not call_ops: - self.output_idx = [] - else: - callee = call_ops[0].callee.value - src_kernel = next( - ( - op - for op in src_mod.body.operations - if isinstance(op, func_d.FuncOp) and op.sym_name.value == callee - ), - None, - ) - self.output_idx = ( - _analyze_kernel_io(src_kernel)[1] if src_kernel else [] - ) + self.project = project or f"{top_func_name}.prj" + self.project_dir = _resolve_project_dir_for_write(self.project) + self.project_dir.mkdir(parents=True, exist_ok=True) + + (self.project_dir / "allo.mlir").write_text(str(module), encoding="utf-8") + _, self.output_idx = convert(module, top_func_name, self.project_dir) def __repr__(self): return f"AIRModule(top={self.top_func_name}, project={self.project_dir})" def get_ir(self) -> str: - return str(self.air_module) + return (self.project_dir / "top.mlir").read_text(encoding="utf-8") def __call__(self, *args): - _call_prj(str(self.project_dir), self.output_idx, *args) + return _call_prj(str(self.project_dir), self.output_idx, *args) def build( @@ -644,6 +537,8 @@ def build( func_args, wrap_io, ): + """Entry point used by :meth:`allo.customize.Schedule.build`.""" + return AIRModule( module=module, top_func_name=top_func_name,