diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 3454d6c..cf55a6c 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -2,30 +2,192 @@ name: tests on: push: + branches: + - main pull_request: + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +env: + BUILD_TYPE: Release + LLVM_COMMIT: 6146a88f60492b520a36f8f8f3231e15f3cc6082 + LLVM_BUILD_DIR: ${{ github.workspace }}/llvm-project/build + AMOEBA_BUILD_DIR: ${{ github.workspace }}/build/amoeba + CCACHE_DIR: ${{ github.workspace }}/.ccache + CCACHE_COMPRESS: "true" + CCACHE_MAXSIZE: 4G jobs: - pytest: + python-unit-tests: runs-on: ubuntu-latest - strategy: - matrix: - python-version: ["3.10", "3.11"] - steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v7 - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: - python-version: ${{ matrix.python-version }} + python-version: 3.11.13 - - name: Install package + - name: Install Synapse and test dependencies run: | python -m pip install --upgrade pip python -m pip install -e . python -m pip install pytest - - name: Run tests - run: pytest -q \ No newline at end of file + - name: Run frontend and language tests + run: python -m pytest -q tests/python/frontend tests/python/language + + compiler-integration: + runs-on: ubuntu-22.04 + timeout-minutes: 240 + + steps: + - name: Checkout Synapse + uses: actions/checkout@v7 + + - name: Initialize Amoeba and Neura + run: | + git submodule update --init mlir/amoeba + git -C mlir/amoeba config \ + submodule.thirdparty/neura.url \ + https://github.com/coredac/neura.git + git -C mlir/amoeba submodule update --init thirdparty/neura + + - name: Verify compiler submodules + run: | + test "$(git -C mlir/amoeba rev-parse HEAD)" = "$(git rev-parse HEAD:mlir/amoeba)" + test "$(git -C mlir/amoeba/thirdparty/neura rev-parse HEAD)" = "$(git -C mlir/amoeba rev-parse HEAD:thirdparty/neura)" + + - name: Install build dependencies + run: | + sudo apt-get update + sudo apt-get install --yes ccache clang lld ninja-build + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: 3.11.13 + + - name: Restore Python package cache + uses: actions/cache@v6 + with: + path: ~/.cache/pip + key: ${{ runner.os }}-pip-python-3.11.13-synapse + + - name: Install Python build and test dependencies + run: | + python -m pip install --upgrade pip + python -m pip install pybind11==2.13.6 nanobind==2.15.0 pytest + python -m pip install -e . + + - name: Restore ccache + id: ccache + uses: actions/cache/restore@v6 + with: + path: ${{ env.CCACHE_DIR }} + key: ${{ runner.os }}-ccache-${{ env.LLVM_COMMIT }}-${{ github.sha }} + restore-keys: | + ${{ runner.os }}-ccache-${{ env.LLVM_COMMIT }}- + + - name: Restore LLVM build + id: llvm-cache + uses: actions/cache/restore@v6 + with: + path: llvm-project + key: ${{ runner.os }}-llvm-python-3.11.13-${{ env.LLVM_COMMIT }}-${{ env.BUILD_TYPE }}-v2 + + - name: Build LLVM and MLIR + if: steps.llvm-cache.outputs.cache-hit != 'true' + run: | + mkdir -p "${CCACHE_DIR}" + git init llvm-project + git -C llvm-project remote add origin https://github.com/llvm/llvm-project.git + git -C llvm-project fetch --depth=1 --filter=blob:none origin "${LLVM_COMMIT}" + git -C llvm-project checkout --detach FETCH_HEAD + + cmake -G Ninja \ + -S llvm-project/llvm \ + -B llvm-project/build \ + -DLLVM_ENABLE_PROJECTS="mlir;clang" \ + -DLLVM_BUILD_EXAMPLES=OFF \ + -DLLVM_TARGETS_TO_BUILD=Native \ + -DCMAKE_BUILD_TYPE="${BUILD_TYPE}" \ + -DLLVM_ENABLE_ASSERTIONS=ON \ + -DCMAKE_C_COMPILER=clang \ + -DCMAKE_CXX_COMPILER=clang++ \ + -DCMAKE_CXX_STANDARD=17 \ + -DCMAKE_CXX_FLAGS="-std=c++17 -frtti" \ + -DLLVM_ENABLE_LLD=ON \ + -DMLIR_INSTALL_AGGREGATE_OBJECTS=ON \ + -DLLVM_ENABLE_RTTI=ON \ + -DLLVM_CCACHE_BUILD=ON \ + -DMLIR_ENABLE_BINDINGS_PYTHON=ON \ + -DMLIR_BINDINGS_PYTHON_NB_DOMAIN=mlir \ + -DPython3_EXECUTABLE="$(which python)" \ + -DPython_EXECUTABLE="$(which python)" \ + -DCMAKE_C_COMPILER_LAUNCHER=ccache \ + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache + + cmake --build llvm-project/build --parallel 2 + + - name: Verify LLVM build + run: | + test -f llvm-project/build/lib/cmake/llvm/LLVMConfig.cmake + test -f llvm-project/build/lib/cmake/mlir/MLIRConfig.cmake + test -x llvm-project/build/bin/llvm-lit + + - name: Save LLVM build + if: steps.llvm-cache.outputs.cache-hit != 'true' + uses: actions/cache/save@v6 + with: + path: llvm-project + key: ${{ steps.llvm-cache.outputs.cache-primary-key }} + + - name: Configure Amoeba + run: | + cmake -G Ninja \ + -S mlir/amoeba \ + -B "${AMOEBA_BUILD_DIR}" \ + -DCMAKE_BUILD_TYPE="${BUILD_TYPE}" \ + -DPython3_EXECUTABLE="$(which python)" \ + -DPython_EXECUTABLE="$(which python)" \ + -DCMAKE_C_COMPILER_LAUNCHER=ccache \ + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache + + - name: Build Amoeba compiler and Python bindings + run: | + cmake --build "${AMOEBA_BUILD_DIR}" \ + --target mlir-amoeba-opt AmoebaPythonModules \ + --parallel 2 + + - name: Verify Amoeba build + run: | + test -x build/amoeba/tools/mlir-amoeba-opt/mlir-amoeba-opt + test -d build/amoeba/python_packages/amoeba_core/taskflow_mlir + + - name: Run compiler integration tests + run: python -m pytest -q tests/python/compiler + + - name: Save ccache + if: steps.ccache.outputs.cache-hit != 'true' + uses: actions/cache/save@v6 + with: + path: ${{ env.CCACHE_DIR }} + key: ${{ steps.ccache.outputs.cache-primary-key }} + + - name: Show ccache statistics + if: always() + run: | + if command -v ccache >/dev/null 2>&1; then + ccache --show-stats + else + echo "ccache was not installed because setup did not complete." + fi diff --git a/.gitmodules b/.gitmodules index 0fd26b1..7e9f9ca 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,3 +1,3 @@ -[submodule "mlir/neura"] - path = mlir/neura - url = https://github.com/coredac/neura +[submodule "mlir/amoeba"] + path = mlir/amoeba + url = https://github.com/coredac/amoeba diff --git a/README.md b/README.md index b03062d..ac83b8b 100644 --- a/README.md +++ b/README.md @@ -1,53 +1,189 @@ # SYNAPSE -SYNAPSE is a programming model and compiler frontend for -task-level dataflow systems. +SYNAPSE is a Python programming model and compiler frontend for spatial +dataflow systems. Its first backend target is a multi-CGRA architecture in +which each task may contain a program explicitly placed on a CGRA tile array. + +The current single-task compiler path lowers a Python TileArray program into: + +```text +func.func + taskflow.task + neura.kernel + placed Neura operations +``` + +The Neura backend then legalizes predicated values, inserts data movement, and +maps the placed operations to tiles, links, registers, and time steps. ## Repository Layout ```text -python/synapse/ Python package source -python/synapse/frontend Source capture and frontend parser utilities +python/synapse/ + language/ User-facing spatial and TileArray language APIs + frontend/ Python program capture and lowering to MLIR + compiler/ Backend compiler orchestration examples/ Small frontend examples -tests/ Pytest tests -mlir/neura/ Downstream NEURA / Taskflow compiler stack +tests/python/ Python unit and compiler-integration tests +mlir/amoeba/ Pinned Amoeba compiler dependency + thirdparty/neura/ Pinned Neura backend dependency managed by Amoeba ``` -## Setup +## Requirements + +SYNAPSE currently requires Python 3.11. + +The compiler-integration path additionally requires: + +- CMake, Ninja, Clang, LLD, and ccache; +- `pybind11==2.13.6` and `nanobind==2.15.0`; and +- LLVM/MLIR at commit + [`6146a88f60492b520a36f8f8f3231e15f3cc6082`](https://github.com/llvm/llvm-project/commit/6146a88f60492b520a36f8f8f3231e15f3cc6082). + +This is the same LLVM revision and Python binding configuration used by the +pinned Amoeba workflow. + +## Checkout + +After cloning SYNAPSE, initialize Amoeba and its direct Neura dependency: + +```bash +git submodule update --init mlir/amoeba +git -C mlir/amoeba submodule update --init thirdparty/neura +``` + +These commands intentionally avoid downloading Neura's nested benchmark +submodules, which are not required to build the SYNAPSE compiler path. + +## Python Setup Create or activate a Python environment, then install SYNAPSE in editable mode from the repository root: ```bash -cd $PROJECT_PATH/synapse +python -m pip install --upgrade pip python -m pip install -e . +python -m pip install pytest ``` Editable install only needs to be done once per environment. After that, changes under `python/synapse/` are picked up directly. -For tests, install `pytest`: +## Build LLVM and MLIR + +Install the Python dependencies into the same Python 3.11 environment that +will configure LLVM and Amoeba: ```bash -python -m pip install pytest +python -m pip install pybind11==2.13.6 nanobind==2.15.0 ``` -## Run The GEMM Parser Example +Choose a location for LLVM, clone it, and check out the pinned revision: ```bash -cd $PROJECT_PATH/synapse -python examples/gemm.py +export SYNAPSE_LLVM_PROJECT=/absolute/path/to/llvm-project + +git clone https://github.com/llvm/llvm-project.git "${SYNAPSE_LLVM_PROJECT}" +git -C "${SYNAPSE_LLVM_PROJECT}" checkout \ + 6146a88f60492b520a36f8f8f3231e15f3cc6082 +``` + +Configure and build LLVM/MLIR with Python bindings enabled: + +```bash +cmake -G Ninja \ + -S "${SYNAPSE_LLVM_PROJECT}/llvm" \ + -B "${SYNAPSE_LLVM_PROJECT}/build" \ + -DLLVM_ENABLE_PROJECTS="mlir;clang" \ + -DLLVM_BUILD_EXAMPLES=OFF \ + -DLLVM_TARGETS_TO_BUILD=Native \ + -DCMAKE_BUILD_TYPE=Release \ + -DLLVM_ENABLE_ASSERTIONS=ON \ + -DCMAKE_C_COMPILER=clang \ + -DCMAKE_CXX_COMPILER=clang++ \ + -DCMAKE_CXX_STANDARD=17 \ + -DCMAKE_CXX_FLAGS="-std=c++17 -frtti" \ + -DLLVM_ENABLE_LLD=ON \ + -DMLIR_INSTALL_AGGREGATE_OBJECTS=ON \ + -DLLVM_ENABLE_RTTI=ON \ + -DLLVM_CCACHE_BUILD=ON \ + -DMLIR_ENABLE_BINDINGS_PYTHON=ON \ + -DMLIR_BINDINGS_PYTHON_NB_DOMAIN=mlir \ + -DPython3_EXECUTABLE="$(which python)" \ + -DPython_EXECUTABLE="$(which python)" \ + -DCMAKE_C_COMPILER_LAUNCHER=ccache \ + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache + +cmake --build "${SYNAPSE_LLVM_PROJECT}/build" --parallel 2 +``` + +## Build Amoeba and Neura + +From the SYNAPSE repository root, point Amoeba at the LLVM build and create the +compiler artifacts under `build/amoeba`: + +```bash +export LLVM_BUILD_DIR="${SYNAPSE_LLVM_PROJECT}/build" + +cmake -G Ninja \ + -S mlir/amoeba \ + -B build/amoeba \ + -DCMAKE_BUILD_TYPE=Release \ + -DPython3_EXECUTABLE="$(which python)" \ + -DPython_EXECUTABLE="$(which python)" + +cmake --build build/amoeba \ + --target mlir-amoeba-opt AmoebaPythonModules \ + --parallel 2 ``` -The example parses a plain Python GEMM function and prints its Python AST. At -this stage, the GEMM function is not executed; SYNAPSE only reads its source. +The build produces: + +```text +build/amoeba/tools/mlir-amoeba-opt/mlir-amoeba-opt +build/amoeba/python_packages/amoeba_core/taskflow_mlir/ +``` ## Run Tests +Run the frontend and language tests without building LLVM: + +```bash +python -m pytest -q tests/python/frontend tests/python/language +``` + +After building Amoeba, run the compiler-integration tests: + ```bash -cd $PROJECT_PATH/synapse -pytest -q +python -m pytest -q tests/python/compiler ``` -The current tests check that the frontend parser can capture a Python GEMM -function and expose the expected AST nodes. +Or run the complete Python suite: + +```bash +python -m pytest -q tests/python +``` + +## Run the GEMM Parser Example + +```bash +python examples/gemm.py +``` + +This example captures a plain Python GEMM function and prints its Python AST. +The TileArray compiler path is exercised by the tests under +`tests/python/compiler`. + +## Continuous Integration + +The GitHub Actions workflow contains two layers: + +- `python-unit-tests` runs frontend and language tests on Python 3.11 without + building LLVM. +- `compiler-integration` uses Python 3.11, checks out the pinned Amoeba and + Neura revisions, downloads the pinned LLVM commit, builds LLVM/MLIR with + Python bindings, builds `mlir-amoeba-opt` and `AmoebaPythonModules`, and runs + the compiler tests. + +LLVM and ccache artifacts are cached using the pinned LLVM revision so later +workflow runs do not rebuild the entire dependency stack unnecessarily. diff --git a/examples/gemm.py b/examples/gemm.py index b0bd60a..6fbba93 100644 --- a/examples/gemm.py +++ b/examples/gemm.py @@ -1,5 +1,6 @@ from synapse.frontend.parser import dump_ast + def gemm(A, B, C): for i in range(128): for j in range(128): @@ -8,5 +9,6 @@ def gemm(A, B, C): acc += A[i][k] * B[k][j] C[i][j] = acc + if __name__ == "__main__": print(dump_ast(gemm)) diff --git a/mlir/README.md b/mlir/README.md index e49c9ba..a7e2ecc 100644 --- a/mlir/README.md +++ b/mlir/README.md @@ -1,23 +1,56 @@ -# SYNAPSE MLIR Submodules +# SYNAPSE MLIR Dependencies -This directory is reserved for MLIR/compiler submodules used by SYNAPSE. +This directory contains the compiler projects used by SYNAPSE after its Python +frontend has captured a program. -Current layout: +## Dependency layout ```text mlir/ - neura/ # git submodule: https://github.com/coredac/neura + amoeba/ # git submodule: https://github.com/coredac/amoeba + thirdparty/neura/ # nested submodule managed by Amoeba ``` -SYNAPSE itself only defines the programming model, frontend, internal graph IR, -and lowering to Taskflow IR. Compilation from Taskflow IR to NEURA or lower -hardware/compiler targets belongs to the NEURA/Taskflow compiler stack in this -directory. +SYNAPSE owns the programming model, Python frontend, and lowering of a captured +program into compiler IR. Amoeba provides the backend-neutral Taskflow dialect +and backend integration. Neura is Amoeba's CGRA backend and provides the +single-CGRA dialect, mapping, routing, register allocation, and code generation. -The submodule is registered in `.gitmodules`: +SYNAPSE therefore depends on Amoeba rather than carrying a second, independent +Neura checkout. This keeps the Taskflow and Neura interfaces on the revisions +tested together by Amoeba. + +## Initial single-task compilation boundary + +The first CGRA path intentionally handles one task only. The frontend treats the +whole captured program as an implicit task and lowers it to the following +container hierarchy: + +```text +taskflow.task + neura.kernel + placed Neura operations with ordinary MLIR value types +``` + +The frontend fixes the spatial operation placement selected by the TileArray +program, but it does not construct Neura's predicated value type or final +mapping metadata. The backend compilation flow performs: ```text -[submodule "mlir/neura"] - path = mlir/neura - url = https://github.com/coredac/neura +--leverage-predicated-value + -> --insert-data-mov + -> --map-to-accelerator="mapping-strategy=template mapping-mode=spatial-only" +``` + +The final mapped Neura IR contains tile coordinates, time steps, links, and +register information. The single-task path does not yet define user-facing +task syntax or perform inter-task allocation, placement, scheduling, +replication, or communication. + +## Checkout + +Initialize Amoeba and its Neura dependency recursively: + +```sh +git submodule update --init --recursive ``` diff --git a/mlir/amoeba b/mlir/amoeba new file mode 160000 index 0000000..784cb87 --- /dev/null +++ b/mlir/amoeba @@ -0,0 +1 @@ +Subproject commit 784cb87a45f51574cb45088da25bd674bed03abb diff --git a/mlir/neura b/mlir/neura deleted file mode 160000 index 5763d85..0000000 --- a/mlir/neura +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 5763d850d2618392e14fe4bc4714d189ad906290 diff --git a/pyproject.toml b/pyproject.toml index e861013..8c7c081 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,10 +5,16 @@ build-backend = "setuptools.build_meta" [project] name = "synapse" version = "0.0.0" -requires-python = ">=3.10" +requires-python = ">=3.11" [tool.setuptools] package-dir = {"" = "python"} [tool.setuptools.packages.find] -where = ["python"] \ No newline at end of file +where = ["python"] + +[tool.pytest.ini_options] +pythonpath = [ + "python", + "build/amoeba/python_packages/amoeba_core", +] diff --git a/python/synapse/__init__.py b/python/synapse/__init__.py index e69de29..236470b 100644 --- a/python/synapse/__init__.py +++ b/python/synapse/__init__.py @@ -0,0 +1,3 @@ +from .compiler import compile + +__all__ = ["compile"] diff --git a/python/synapse/compiler/__init__.py b/python/synapse/compiler/__init__.py new file mode 100644 index 0000000..236470b --- /dev/null +++ b/python/synapse/compiler/__init__.py @@ -0,0 +1,3 @@ +from .compiler import compile + +__all__ = ["compile"] diff --git a/python/synapse/compiler/compiler.py b/python/synapse/compiler/compiler.py new file mode 100644 index 0000000..1213525 --- /dev/null +++ b/python/synapse/compiler/compiler.py @@ -0,0 +1,70 @@ +"""Top-Level Synapse Compilation Flow.""" + +import subprocess +from collections.abc import Callable +from pathlib import Path +from tempfile import TemporaryDirectory + +from synapse.frontend.lowering import lower + + +def compile(program: Callable, *, target: str) -> str: + """Compile a Synapse program for the selected backend.""" + + # We only support the Neura backend for now, so we raise an error if the user tries to compile for any other target. + if target != "neura": + raise ValueError(f"unsupported compilation target: {target}") + # TODO: Support the amoeba backend. + + neura_ir = lower(program) + return _run_neura_backend(neura_ir) + + +def _run_neura_backend(neura_ir: str) -> str: + """Legalize Neura values, insert data movement, and run template mapping.""" + + repository_root = Path(__file__).resolve().parents[3] + amoeba_opt = ( + repository_root + / "build" + / "amoeba" + / "tools" + / "mlir-amoeba-opt" + / "mlir-amoeba-opt" + ) + + if not amoeba_opt.is_file(): + raise FileNotFoundError(f"Amoeba compiler is not built: {amoeba_opt}") + + with TemporaryDirectory(prefix="synapse-") as temporary_directory: + output_path = Path(temporary_directory) / "mapped.mlir" + + command = [ + str(amoeba_opt), + "--leverage-predicated-value", + "--insert-data-mov", + ( + "--map-to-accelerator=" + "mapping-strategy=template " + "mapping-mode=spatial-only" + ), + "-o", + str(output_path), + ] + + completed = subprocess.run( + command, + input=neura_ir, + capture_output=True, + text=True, + check=False, + ) + + if completed.returncode != 0: + diagnostics = "\n".join( + output for output in (completed.stdout, completed.stderr) if output + ) + + raise RuntimeError(f"Neura backend compilation failed:\n{diagnostics}") + + return output_path.read_text(encoding="utf-8") diff --git a/python/synapse/frontend/lowering.py b/python/synapse/frontend/lowering.py new file mode 100644 index 0000000..74b3cfb --- /dev/null +++ b/python/synapse/frontend/lowering.py @@ -0,0 +1,187 @@ +"""Lower Synapse Python programs to compiler input IR.""" + +from collections.abc import Callable + +from synapse.language.spatial import Tile +from synapse.language.tile_array_program import ( + AddOp, + ConstantOp, + TileArrayBuilder, + TileArrayProgram, + TileArrayScalarType, +) + + +def lower(program_fn: Callable) -> str: + """Lower one tile-array program to pre-mapping Taskflow and Neura IR.""" + + builder = TileArrayBuilder() + + # Execute the user's tile-array DSL while recording its operations. + with builder: + program_fn() + + program = builder.build() + + return _lower_tile_array_program( + program_name=program_fn.__name__, + program=program, + ) + + +def _lower_tile_array_program( + *, + program_name: str, + program: TileArrayProgram, +) -> str: + """Convert a TileArrayProgram into pre-mapping Taskflow and Neura IR. + + MLIR imports remain local so users can import ``synapse.language`` without + requiring the compiled Amoeba Python bindings. + """ + + from taskflow_mlir.dialects import func, neura, taskflow + from taskflow_mlir.ir import ( + Context, + DictAttr, + F32Type, + FloatAttr, + InsertionPoint, + IntegerAttr, + IntegerType, + Location, + Module, + StringAttr, + ) + + with Context(), Location.unknown(): + taskflow.register_dialect() + neura.register_dialect() + + i32 = IntegerType.get_signless(32) + + def get_mlir_type(dtype: TileArrayScalarType): + """Translate a frontend scalar type into an MLIR type.""" + + if dtype == TileArrayScalarType.I32: + return i32 + + if dtype == TileArrayScalarType.F32: + return F32Type.get() + + raise NotImplementedError( + f"unsupported tile-array scalar type: {dtype.value}" + ) + + def get_constant_attribute( + operation: ConstantOp, + result_type, + ): + """Build the typed MLIR attribute for a constant value.""" + + if operation.result.dtype == TileArrayScalarType.I32: + if not isinstance(operation.value, int): + raise TypeError("an i32 constant requires an integer value") + return IntegerAttr.get(result_type, operation.value) + + if operation.result.dtype == TileArrayScalarType.F32: + if not isinstance(operation.value, (float, int)): + raise TypeError("an f32 constant requires a numeric value") + return FloatAttr.get(result_type, float(operation.value)) + + raise NotImplementedError( + f"unsupported constant type: {operation.result.dtype.value}" + ) + + def get_placement(tile: Tile) -> DictAttr: + """Build Neura placement directly from a Tile coordinate.""" + + return DictAttr.get( + { + "x": IntegerAttr.get(i32, tile.x), + "y": IntegerAttr.get(i32, tile.y), + } + ) + + module = Module.create() + + # This milestone lowers one Python function into one task containing + # one manually placed Neura kernel. + with InsertionPoint(module.body): + function = func.FuncOp(program_name, ([], [])) + function_block = function.add_entry_block() + + with InsertionPoint(function_block): + task = taskflow.TaskflowTaskOp( + done_reads=[], + done_writes=[], + value_outputs=[], + will_reads=[], + will_writes=[], + value_inputs=[], + task_name=program_name, + original_read_memrefs=[], + original_write_memrefs=[], + ) + task_block = task.body.blocks.append() + + func.ReturnOp([]) + + with InsertionPoint(task_block): + kernel = neura.KernelOp( + outputs=[], + inputs=[], + iter_args_init=[], + accelerator=StringAttr.get("neura"), + ) + kernel_block = kernel.body.blocks.append() + + taskflow.TaskflowYieldOp( + done_reads=[], + done_writes=[], + value_results=[], + ) + + # Map frontend value IDs to the MLIR SSA values produced while + # lowering the recorded operations. + values_by_id = {} + + with InsertionPoint(kernel_block): + for operation in program.operations: + result_type = get_mlir_type(operation.result.dtype) + + if isinstance(operation, ConstantOp): + mlir_operation = neura.ConstantOp( + result_type, + get_constant_attribute(operation, result_type), + ) + + elif isinstance(operation, AddOp): + lhs = values_by_id[operation.lhs.id] + rhs = values_by_id[operation.rhs.id] + + mlir_operation = neura.AddOp( + result_type, + lhs, + rhs=rhs, + ) + + else: + raise NotImplementedError( + f"unsupported tile-array operation: {type(operation).__name__}" + ) + + mlir_operation.operation.attributes["placement"] = get_placement( + operation.tile + ) + values_by_id[operation.result.id] = mlir_operation.result + + neura.YieldOp( + iter_args_next=[], + results_=[], + ) + + if not module.operation.verify(): + raise RuntimeError("generated Taskflow/Neura module is invalid") + + return str(module) diff --git a/python/synapse/language/__init__.py b/python/synapse/language/__init__.py new file mode 100644 index 0000000..eca34d4 --- /dev/null +++ b/python/synapse/language/__init__.py @@ -0,0 +1,9 @@ +"""Public Synapse language API.""" + +from .spatial import TileArray +from .tile_array_program import TileArrayScalarType, add, constant + +i32 = TileArrayScalarType.I32 +f32 = TileArrayScalarType.F32 + +__all__ = ["TileArray", "TileArrayScalarType", "add", "constant", "f32", "i32"] diff --git a/python/synapse/language/spatial.py b/python/synapse/language/spatial.py new file mode 100644 index 0000000..a18c886 --- /dev/null +++ b/python/synapse/language/spatial.py @@ -0,0 +1,62 @@ +"""Hardware spatial features exposed by the Synapse language. +This file exposes spatial structures that programmers can use to organize +computation and data movement according to the target hardware hierarchy. + +TileArray currently exposes the two-dimensional tile array of a CGRA. + +Future spatial abstractions may expose inter-core structures, such as the +core array of a multi-CGRA, AMD AIE/NPU, or Tenstorrent. +""" + + +class Tile: + """A hardware tile in a CGRA TileArray. + + A Tile is owned by a TileArray. Its ``x`` and ``y`` coordinates identify + the corresponding position in the target CGRA tile array. + """ + + def __init__(self, x: int, y: int, array: "TileArray"): + self.array = array + self.x = x + self.y = y + + +class TileArray: + """A parameterized two-dimensional tile array. + + In the initial implementation, ``x_tiles`` and ``y_tiles`` must match the + dimensions of the target CGRA. A tile accessed as ``array[x, y]`` + corresponds directly to the tile at that hardware coordinate. + + Coordinates follow Neura's convention: increasing ``x`` moves east + (right), and increasing ``y`` moves north (up). + + Example: + array = TileArray(x_tiles=4, y_tiles=4) + tile = array[1, 2] + """ + + def __init__(self, x_tiles: int, y_tiles: int): + self.x_tiles = x_tiles + self.y_tiles = y_tiles + + self._tiles = [ + [Tile(x=x, y=y, array=self) for x in range(x_tiles)] + for y in range(y_tiles) + ] + + def tiles(self): + """Iterate over all tiles in the array. + + The iteration order is a Python programming convenience and + does not specify sequential hardware execution. + """ + for y_row in self._tiles: + yield from y_row + + def __getitem__(self, coordinate: tuple[int, int]) -> Tile: + """Return the tile at the given ``(x, y)`` coordinate.""" + + x, y = coordinate + return self._tiles[y][x] diff --git a/python/synapse/language/tile_array_program.py b/python/synapse/language/tile_array_program.py new file mode 100644 index 0000000..12f0ff8 --- /dev/null +++ b/python/synapse/language/tile_array_program.py @@ -0,0 +1,261 @@ +"""Programming model for computations placed on a TileArray. + +This module defines the typed frontend representation of a tile-array program. +It records computation independently of MLIR. Compiler lowering later converts +the recorded program into Taskflow and Neura operations. +""" + +from __future__ import annotations + +from contextvars import ContextVar +from dataclasses import dataclass, field +from enum import Enum +from typing import TypeAlias + +from .spatial import Tile, TileArray + + +class TileArrayScalarType(str, Enum): + """The scalar types supported by the TileArray programming model. + This is intentionally independent of MLIR types. The lowering converts + these frontend types into the corresponding MLIR types. + + Additional scalar types can be added here as the language grows. + """ + + I32 = "i32" + F32 = "f32" + + +@dataclass(frozen=True) +class TileArrayValue: + """A typed value produced by one tile-array operation.""" + + id: int + dtype: TileArrayScalarType + _builder: TileArrayBuilder = field(repr=False) + + +# --------------------------------------------------------------- +# Typed tile-array operations +# --------------------------------------------------------------- + + +@dataclass(frozen=True) +class ConstantOp: + """A scalar constant produced by a tile-array operation.""" + + result: TileArrayValue + value: int | float + tile: Tile + + +@dataclass(frozen=True) +class AddOp: + """A scalar addition executed by a tile-array operation.""" + + result: TileArrayValue + lhs: TileArrayValue + rhs: TileArrayValue + tile: Tile + + +# This union explicitly lists every operation currently supported by the +# tile-array frontend. Future operations such as MacOp and GatherOp should be +# added here. +TileArrayOp: TypeAlias = ConstantOp | AddOp + + +@dataclass(frozen=True) +class TileArrayProgram: + """A tile-array program produced by TileArrayBuilder.""" + + array: TileArray + operations: tuple[TileArrayOp, ...] + + +# --------------------------------------------------------------- +# Internal program builder +# --------------------------------------------------------------- +class TileArrayBuilder: + """A tile-array program builder. + + The builder records typed operations in user-program order. Once build() + is called, it returns a TileArrayProgram. + """ + + def __init__(self): + self._array: TileArray | None = None + self._operations: list[TileArrayOp] = [] + self._next_value_id = 0 + self._token = None + self._is_built = False + + def __enter__(self): + """Make this builder active for tile-array DSL calls.""" + if _active_builder.get() is not None: + raise RuntimeError("Cannot enter a nested TileArrayBuilder context") + if self._token is not None: + raise RuntimeError("TileArrayBuilder context is already active") + self._token = _active_builder.set(self) + return self + + def __exit__(self, exc_type, exc_value, traceback) -> None: + """Restore the previously active builder.""" + if self._token is None: + raise RuntimeError("TileArrayBuilder is not active") + + _active_builder.reset(self._token) + self._token = None + + def _new_value(self, *, dtype: TileArrayScalarType) -> TileArrayValue: + """Allocate the next value identifier.""" + result = TileArrayValue(id=self._next_value_id, dtype=dtype, _builder=self) + self._next_value_id += 1 + return result + + def _bind_tile_array(self, tile: Tile) -> None: + """Bind the program to one TileArray.""" + if not isinstance(tile, Tile): + raise TypeError("tile must be a Tile") + if self._array is None: + self._array = tile.array + return + if tile.array is not self._array: + raise ValueError( + "all operations in a TileArrayProgram must use tiles from the same TileArray" + ) + + def _check_operand(self, operand: TileArrayValue) -> None: + """Verify that an operand was produced by this builder.""" + if not isinstance(operand, TileArrayValue): + raise TypeError("operation operand must be a TileArrayValue") + + if operand._builder is not self: + raise ValueError( + "operation operand belongs to a different TileArrayProgram" + ) + + def _check_can_emit(self) -> None: + """Reject operations emitted after the program has been finalized.""" + if self._is_built: + raise RuntimeError( + "cannot emit operations after building a TileArrayProgram" + ) + + def emit_constant( + self, value: int | float, *, dtype: TileArrayScalarType, tile: Tile + ) -> TileArrayValue: + """Record one scalar constant operation.""" + self._check_can_emit() + self._bind_tile_array(tile) + result = self._new_value(dtype=dtype) + + self._operations.append(ConstantOp(result=result, value=value, tile=tile)) + return result + + def emit_add( + self, lhs: TileArrayValue, rhs: TileArrayValue, *, tile: Tile + ) -> TileArrayValue: + """Record one scalar addition operation.""" + self._check_can_emit() + self._bind_tile_array(tile) + self._check_operand(lhs) + self._check_operand(rhs) + + if lhs.dtype != rhs.dtype: + raise TypeError("add operands must have the same tile-array value type") + + result = self._new_value(dtype=lhs.dtype) + + self._operations.append(AddOp(result=result, lhs=lhs, rhs=rhs, tile=tile)) + return result + + def build(self) -> TileArrayProgram: + """Finish recording and return a program.""" + if self._token is not None: + raise RuntimeError( + "cannot build a TileArrayProgram while its builder is active" + ) + if self._array is None: + raise RuntimeError( + "cannot build an empty TileArrayProgram without a TileArray" + ) + self._is_built = True + return TileArrayProgram(array=self._array, operations=tuple(self._operations)) + + +# The active builder is compiler-internal state. Public DSL calls use it to +# find the builder created by frontend lowering. +_active_builder: ContextVar[TileArrayBuilder | None] = ContextVar( + "active_tile_array_builder", + default=None, +) + + +def _require_active_builder() -> TileArrayBuilder: + """Return the active builder.""" + builder = _active_builder.get() + + if builder is None: + raise RuntimeError( + "tile-array operations must be called while lowering a Synapse program" + ) + + return builder + + +# --------------------------------------------------------------- +# User-facing tile-array program DSL +# --------------------------------------------------------------- +def constant( + value: int | float, *, tile: Tile, dtype: TileArrayScalarType | None = None +) -> TileArrayValue: + """Create a scalar constant on one hardware tile. + + Integer literals default to i32. Floating-point literals default to f32. + Use an explicit dtype when a different representation is required: + constant(1.0, tile=tile, dtype=TileArrayScalarType.F32) + """ + if isinstance(value, bool): + raise TypeError("boolean constants are not supported yet") + + if dtype is None: + if isinstance(value, int): + dtype = TileArrayScalarType.I32 + elif isinstance(value, float): + dtype = TileArrayScalarType.F32 + else: + raise TypeError( + "constant currently supports integer and floating-point values" + ) + + if not isinstance(dtype, TileArrayScalarType): + raise TypeError("dtype must be a TileArrayScalarType") + + if dtype == TileArrayScalarType.I32 and not isinstance(value, int): + raise TypeError("an i32 constant requires an integer value") + + if dtype in (TileArrayScalarType.F32,) and not isinstance(value, (int, float)): + raise TypeError("a floating-point constant requires a numeric value") + + return _require_active_builder().emit_constant( + value, + dtype=dtype, + tile=tile, + ) + + +def add( + lhs: TileArrayValue, + rhs: TileArrayValue, + *, + tile: Tile, +) -> TileArrayValue: + """Create a scalar addition on one hardware tile.""" + + return _require_active_builder().emit_add( + lhs, + rhs, + tile=tile, + ) diff --git a/tests/python/compiler/test_add_constant_kernel.py b/tests/python/compiler/test_add_constant_kernel.py new file mode 100644 index 0000000..db551a9 --- /dev/null +++ b/tests/python/compiler/test_add_constant_kernel.py @@ -0,0 +1,61 @@ +import synapse +import synapse.language as synl +from synapse.frontend.lowering import lower + +PRE_MAPPING_IR = """ +module { + func.func @add_constant() { + taskflow.task @add_constant : () -> () { + neura.kernel attributes {accelerator = "neura"} { + %0 = "neura.constant"() <{value = 1 : i32}> {placement = {x = 0 : i32, y = 0 : i32}} : () -> i32 + %1 = "neura.constant"() <{value = 2 : i32}> {placement = {x = 2 : i32, y = 0 : i32}} : () -> i32 + %2 = "neura.add"(%0, %1) {placement = {x = 1 : i32, y = 0 : i32}} : (i32, i32) -> i32 + neura.yield + } + taskflow.yield + } + return + } +} +""".strip() + + +MAPPED_IR = """ +module { + func.func @add_constant() { + taskflow.task @add_constant : () -> () { + neura.kernel attributes {accelerator = "neura", mapping_info = {compiled_ii = 1 : i32, mapping_mode = "spatial-only", mapping_strategy = "template", rec_mii = 1 : i32, res_mii = 1 : i32, x_tiles = 4 : i32, y_tiles = 4 : i32}} { + %0 = "neura.constant"() <{value = 1 : i32}> {dfg_id = 0 : i32, mapping_locs = [{id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 0 : i32, x = 0 : i32, y = 0 : i32}]} : () -> !neura.data + %1 = "neura.constant"() <{value = 2 : i32}> {dfg_id = 1 : i32, mapping_locs = [{id = 2 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "tile", time_step = 0 : i32, x = 2 : i32, y = 0 : i32}]} : () -> !neura.data + %2 = "neura.data_mov"(%0) {dfg_id = 3 : i32, mapping_locs = [{id = 0 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 0 : i32}]} : (!neura.data) -> !neura.data + %3 = "neura.data_mov"(%1) {dfg_id = 4 : i32, mapping_locs = [{id = 5 : i32, index_per_ii = 0 : i32, invalid_iterations = 0 : i32, resource = "link", time_step = 0 : i32}]} : (!neura.data) -> !neura.data + %4 = "neura.add"(%2, %3) {dfg_id = 5 : i32, mapping_locs = [{id = 1 : i32, index_per_ii = 0 : i32, invalid_iterations = 1 : i32, resource = "tile", time_step = 1 : i32, x = 1 : i32, y = 0 : i32}]} : (!neura.data, !neura.data) -> !neura.data + neura.yield {dfg_id = 2 : i32} + } + taskflow.yield + } + return + } +} +""".strip() + + +def add_constant(): + array = synl.TileArray(x_tiles=4, y_tiles=4) + + lhs = synl.constant(1, tile=array[0, 0]) + rhs = synl.constant(2, tile=array[2, 0]) + + synl.add(lhs, rhs, tile=array[1, 0]) + + +def test_add_constant_lowers_to_placed_neura_ir(): + actual = lower(add_constant) + + assert actual.strip() == PRE_MAPPING_IR + + +def test_add_constant_compiles_to_mapped_neura_ir(): + actual = synapse.compile(add_constant, target="neura") + + assert actual.strip() == MAPPED_IR diff --git a/tests/frontend/test_parser.py b/tests/python/frontend/test_parser.py similarity index 100% rename from tests/frontend/test_parser.py rename to tests/python/frontend/test_parser.py diff --git a/tests/python/language/test_spatial.py b/tests/python/language/test_spatial.py new file mode 100644 index 0000000..1cab673 --- /dev/null +++ b/tests/python/language/test_spatial.py @@ -0,0 +1,39 @@ +import synapse.language as synl + + +def test_tile_array_is_parameterized(): + array_2x3 = synl.TileArray(x_tiles=2, y_tiles=3) + array_4x4 = synl.TileArray(x_tiles=4, y_tiles=4) + + assert array_2x3.x_tiles == 2 + assert array_2x3.y_tiles == 3 + + assert array_4x4.x_tiles == 4 + assert array_4x4.y_tiles == 4 + + +def test_tile_array_exposes_tiles(): + array = synl.TileArray(x_tiles=2, y_tiles=3) + + coordinates = {(tile.x, tile.y) for tile in array.tiles()} + + assert coordinates == {(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)} + + +def test_access_tile_array_by_coordinate(): + array = synl.TileArray(x_tiles=2, y_tiles=3) + tile = array[1, 2] + + assert tile.array is array + assert tile.x == 1 + assert tile.y == 2 + + assert tile is array[1, 2] + + enumerated_tile = next( + candidate + for candidate in array.tiles() + if candidate.x == 1 and candidate.y == 2 + ) + + assert tile is enumerated_tile diff --git a/tests/python/language/test_tile_array_program.py b/tests/python/language/test_tile_array_program.py new file mode 100644 index 0000000..847e7cb --- /dev/null +++ b/tests/python/language/test_tile_array_program.py @@ -0,0 +1,48 @@ +import synapse.language as synl + + +def test_records_tile_array_program(): + array = synl.TileArray(x_tiles=4, y_tiles=4) + builder = synl.tile_array_program.TileArrayBuilder() + + with builder: + lhs = synl.constant(1, tile=array[0, 0]) + rhs = synl.constant(2, tile=array[0, 2]) + result = synl.add(lhs, rhs, tile=array[0, 1]) + + program = builder.build() + lhs_op, rhs_op, add_op = program.operations + + assert program.array is array + assert isinstance(lhs_op, synl.tile_array_program.ConstantOp) + assert isinstance(rhs_op, synl.tile_array_program.ConstantOp) + assert isinstance(add_op, synl.tile_array_program.AddOp) + assert [(value.id, value.dtype) for value in (lhs, rhs, result)] == [ + (0, synl.i32), + (1, synl.i32), + (2, synl.i32), + ] + assert [op.tile for op in program.operations] == [ + array[0, 0], + array[0, 2], + array[0, 1], + ] + assert (add_op.lhs, add_op.rhs) == (lhs, rhs) + + +def test_infers_supported_scalar_types(): + array = synl.TileArray(x_tiles=1, y_tiles=3) + builder = synl.tile_array_program.TileArrayBuilder() + + with builder: + integer = synl.constant(1, tile=array[0, 0]) + floating = synl.constant(1.0, tile=array[0, 1]) + explicit_f32 = synl.constant( + 1, + tile=array[0, 2], + dtype=synl.f32, + ) + + assert integer.dtype == synl.i32 + assert floating.dtype == synl.f32 + assert explicit_f32.dtype == synl.f32