# Project export: Operator Optimization This document was generated by HackStack to give an AI agent context about a hackathon project. Sections are labeled with their provenance; content marked as truncated was cut to keep this document small. ## Project metadata - Hackathon: TreeHacks 2025 - Tagline: Optimizing some core computational primitives used in the design space of Liquid AI models, as described in the opening presentation. - Devpost: https://devpost.com/software/operator-optimization - GitHub: https://github.com/kaloca/liquid_treehacks_challenge - Team: 2 GitHub contributor(s) — Zymrael (4 commits), Gabriel Noya (1 commits) ## Devpost submission (written by the team) ### Overview How to test pdes/ directory has .pde files. kernels/ directory has custom kernels. Need to be installed to environment by python setup.py. Issues Faced 1*Compilation Errors* The C++ extension failed to compile because: The out version of the function was missing or incorrectly defined. There were signature mismatches in function definitions. The meta function (fused_fftconv_meta) wasn’t properly registered. Template errors and ATen API changes caused unexpected failures. The out version of the function was missing or incorrectly defined. There were signature mismatches in function definitions. The meta function (fused_fftconv_meta) wasn’t properly registered. Template errors and ATen API changes caused unexpected failures. Fix: Defined fused_fftconv_out properly. Used resize_ and copy_ to ensure the output tensor was correctly shaped without modifying inputs. Registered both versions correctly: cpp TORCH_LIBRARY(myop, m) { m.def("fused_fftconv(Tensor input, Tensor filter) -> Tensor", fused_fftconv); m.def("fused_fftconv.out(Tensor input, Tensor filter, *, Tensor(a!) out) -> Tensor(a!)", fused_fftconv_out); } 2*XNNPACK Runtime Error: Shape Mismatch* -Error: Attempted to change the tensor rank which is immutable: old=3, new=2 The exported model expected a 2D input for the linear layer (e.g., (256, 4)) but received a 3D input (1, 4, 256). This happened because the partitioner isolated the linear operation in a subgraph with shape (B*L, C), while the overall model expected (B, C, L). This happened because the partitioner isolated the linear operation in a subgraph with shape (B*L, C), while the overall model expected (B, C, L). Fix Options: Option A (Python code change): Keep linear as a 3D operation and avoid flattening. Option B (C++ kernel change): Reshape the input inside C++ before convolution to match XNNPACK’s expected shape. Fix Options: Option A (Python code change): Keep linear as a 3D operation and avoid flattening. Option B (C++ kernel change): Reshape the input inside C++ before convolution to match XNNPACK’s expected shape. Most of the time, adjusting the Python side is simpler. 3*PyTorch Export Errors (torch._dynamo.exc.Unsupported: out= op was called where output tensor was non-contiguous)* The export function failed when tracing the model because out was not contiguous. Fix: Ensure out is contiguous before passing it to fused_fftconv_out. For example: python{out = torch.empty_like(x, memory_format=torch.contiguous_format) torch.ops.fftconv.fused_fftconv.out(x, filter, out=out)} Ensure out is contiguous before passing it to fused_fftconv_out. For example: python{out = torch.empty_like(x, memory_format=torch.contiguous_format) torch.ops.fftconv.fused_fftconv.out(x, filter, out=out)} Next Steps & Future Work If I had more time, I would: Optimize for Android: Use NEON intrinsics to improve performance on ARM-based devices. Improve Memory Efficiency: Avoid unnecessary copies and temporary tensors. Explore Other Backends: Possibly consider IPEX or other compilers for further performance gains. Summary Writing custom PyTorch C++ extensions. Exporting models to XNNPACK & ExecutorTorch. Debugging both compile-time and runtime errors. While I managed to get it all running, there wasn't enough time to work on actual optimizations for the hardware. ## README (from the GitHub repository) # Liquid AI: TreeHacks Challenges 🌲🌲 ## Challenge 2: Operator Optimization This challenges tests your ability to think cross-functionally across numerics and hardware. You will be optimizing some core computational primitives used in the design space of Liquid AI models, as described in the opening presentation. Throughout this challenge, [Executorch](https://github.com/pytorch/executorch) will be your friend. We will be using Executorch to transform PyTorch and / or C++ kernels into a representation that can be run on our target device: Samsung Galaxy S24 Ultra. It could be useful to keep the following resources open as you get started: * Tutorial on exporting models and operators: [tutorial](https://pytorch.org/executorch/stable/tutorials/export-to-executorch-tutorial.html) * Docs on kernel registration: [docs](https://pytorch.org/executorch/stable/kernel-library-custom-aten-kernel.html#custom-ops-api-best-practices) Any optimization is allowed as long as (a) the results are numerically correct (absolute tolerance of 1e-3 on random inputs) and (b) the custom operator can be run via our profiling [script](pte_android.py) directly on device (one of the Samsung Galaxy S24 Ultra phones provided by the team). See below for more details on how submissions will be scored. ## Getting Started Our build system supports both Linux and MacOS (no Windows!). ### Dependencies The dependencies for this repository could be split into the following groups, don't run anything yet, just read. 1. The git submodules * It is mainly the executorch repository + its submodules. * `make setup-submodules` pulls the git submodules but does not attempt to sync them or update. There are patches applied to the repositories which make the automated management complicated. In any case, you see a problem with a submodule, delete it and re-run that target. 2. Build tools * `make check-tools` verifies the basic tools required and suggest how to obtain those 3. Python dependencies * `make setup-python-deps` installs the python dependencies declared for that repo, including pytorch, **excluding executorch**. The purpose is to have a minimal set of tools for compilation-only needs. * Dependencies live in`requirements.txt` and `requirements-torch.txt` 4. Executorch as a python dependency * `make setup-executorch` installs executorch and the python dependencies. This calls `install_executorch.sh` (see official documentation). You may need to update your compiler versions and fix your paths. * If you are making changes to the executorch repository `./vendor/executorch` you would need to re-run that step to install them. * The executorch repo can't be installed as editable. ## One-step build With your preferred virtual environment (e.g., with conda `conda create -n treehacks python=3.11`), install the dependencies and build the project. ```bash make setup ``` If `make setup-submodules` fails for you, another option is to clone executorch and move it to `vendor/executorch` before calling `make setup` again. To customize your executorch build, refer to this [page](https://pytorch.org/executorch/main/getting-started-setup.html). Our setup step will automatically install Executorch with XNNPACK backend enabled (which could be useful to you). If you don't need the backend, you can also install the wheels with pip. ### Running PTEs The pte is a representation of your code that can be ingested by a runner. We provide pre-compiled pte runner binaries for Android (`runner/android-arm64-v8a`) and MacOS (`runner/macos-arm64`). You will need these to profile your code (check example usage in `quickstart.py`). These are stored using `git lfs`. We also provide pte runner binaries for Linux (`runner/linux`). If you are on a different OS and need to use the pte runner locally, or if you wish to customize it to your needs, you can recompile the Executorch runner following the official instructions. ## Numerical references We provide a simple mathematical description of the primitives under consideration (`references/math_ref.pdf`). There are multiple approaches to improve your implementation, including purely algorithmic tweaks (e.g., switching from FFT to direct convolution, or Winograd methods) and hardware-specific optimizations (e.g., minimizing overheads, using different data types). Convolutions and recurrences are a cornerstone of numerical computing, and have been implemented and optimized over the years on many different platforms. You can find excellent resources online to get ideas. ## Example custom kernel See `examples/fftconv` for an end-to-end example of how to register a custom operator with executorch and produce the corresponding pte file. To run the notebook, `python setup.py install` in the ops folder. To export, you will need to adapt to the "out" convention (see kernel registration docs in Executorch), linked above. ### Convolutions An example of convolution is provided under `examples/direct_conv`. A convolution implementation using FFT and custom kernel registration is provided under `examples/fftconv`. ### Linear recurrences Some examples using linear scan algorithmsare provided under `examples/linear_scan`. ## Profiling ### Local You can gain a lot of insight by measuring the latency and memory usage of your implementation locally (on your laptop, using CPUs). If you want accurate numbers, we also provide phones. To measure latency on your machine (assuming you have the right pte runner!), use `PYTHONPATH=profiling PTE_RUNNER_PATH=runner/macos-arm64/pte_runner python quickstart.py`. You should see something like this: ``` ProfilingResults(raw=[0.035606, 0.001095, 0.000502, 0.00056], p10=0.0005193999999999999, p50=0.0008275, p90=0.025252700000000003, min=0.000502, avg=0.00944075, max=0.035606) ``` these are the latency measurements. ### Phone `adb` is required to run on the phone. See the [docs](https://developer.android.com/tools/adb) and follow the steps. The phones have USB debugging mode already enabled. To measure on the phone, connect your device to one of the provided Samsung devices and run ` PYTHONPATH=profiling python pte_android.py --pte your_pte_file.pte --runner_path runner/android-arm64-v8a/pte_runner` ### Measurement protocol The objective is to minimize latency across pre-specified input shapes, with random inputs (we will measure and average over 100 random inputs, 10 runs, for each setting). Target input shapes (batch size, channels, sequence length) for both recurrences and convolutions: * 1, 512, 64 * 1, 512, 256 * 1, 512, 1024 * 1, 512, 2048 * 1, 512, 4096 * 1, 512, 16384 * 1, 512, 65536 * 1, 2048, 64 * 1, 2048, 256 * 1, 2048, 1024 * 1, 2048, 2048 * 1, 2048, 4096 * 1, 2048, 16384 * 1, 2048, 65536 In addition, for convolutions, we will measure filter sizes: 4, 32, 128, and as long as the input sequence (max 65536), resulting in 4 times the number of measurements (we will still average across filter size settings). You can provide a pte that accepts dynamic shapes, or different ptes for each shape, allowing you to optimize for the best performance. We will check for numerical correctness with absolute tolerance of 1e-3 on random and structured inputs (e.g., a tensor of ones). ### Issues #### Issues with data type support during compilation These are generally resolved by making sure `gcc` is up-to-date, or switching to a different toolchain, for example `llvm` on arm64. #### Exec format error `OSError: [Errno 8] Exec format error` This is generally due to a mismatch between the pte runner and the architecture of the machine you are running on. To measure latency locally, you can use compile and use Executorch's native runner. ## Detected evidence (automated analysis) Indexed codebase: 20 recognized source files, 64 KB. - C++ (language) — detected in the code - Hugging Face (technology) — detected in the code - Python (language) — detected in the code - PyTorch (technology) — detected in the code ## Codebase structure (from repository index) ### Files (49 of 49) ``` .gitmodules delegate.pte examples/direct_conv/short_conv.ipynb examples/fftconv/.ipynb_checkpoints/fftcconv-checkpoint.ipynb examples/fftconv/fftcconv.ipynb examples/fftconv/ops/__init__.py examples/fftconv/ops/fftconv.cpp examples/fftconv/ops/fused_fftconv.egg-info/dependency_links.txt examples/fftconv/ops/fused_fftconv.egg-info/PKG-INFO examples/fftconv/ops/fused_fftconv.egg-info/SOURCES.txt examples/fftconv/ops/fused_fftconv.egg-info/top_level.txt examples/fftconv/ops/setup.py examples/linear_scan/linear_scan.ipynb examples/linear_scan/ops/__init__.py examples/linear_scan/ops/linear_scan.cpp examples/linear_scan/ops/setup.py fftconvtest.py kernels/fftconv/ops/__init__.py kernels/fftconv/ops/fftconv.cpp kernels/fftconv/ops/fused_fftconv.egg-info/dependency_links.txt kernels/fftconv/ops/fused_fftconv.egg-info/PKG-INFO kernels/fftconv/ops/fused_fftconv.egg-info/SOURCES.txt kernels/fftconv/ops/fused_fftconv.egg-info/top_level.txt kernels/fftconv/ops/setup.py kernels/test_models.py kernels/winograd/ops/__init__.py kernels/winograd/ops/generic_winograd_conv.egg-info/dependency_links.txt kernels/winograd/ops/generic_winograd_conv.egg-info/PKG-INFO kernels/winograd/ops/generic_winograd_conv.egg-info/SOURCES.txt kernels/winograd/ops/generic_winograd_conv.egg-info/top_level.txt kernels/winograd/ops/setup.py kernels/winograd/ops/winograd.cpp Makefile profiling/android_profiler.py profiling/profiler.py pte_android.py pte_local.py ptes/fftconv.pte ptes/winograd.pte quickstart.py README.md requirements-cmake.txt requirements-torch.txt requirements.txt runner/android-arm64-v8a/pte_runner runner/linux/pte_runner runner/macos-arm64/pte_runner test.pte treehacks.ipynb ``` ### Dependencies - requirements.txt: blobfile@==3.0.0, boto3@==1.35.88, einops@==0.8.0, expecttest, flatbuffers, hydra-core@==1.3.2, hypothesis, ipykernel, lm_eval@==0.4.5, matplotlib, memory_profiler, numpy@<2.0, omegaconf, parameterized, pip@>=23, pycryptodome@==3.21.0, pytest, pytest-xdist, pyyaml@==6.0.2, rich@==13.9.4, ruamel.yaml, safetensors@==0.5.1, scipy@==1.13, sentencepiece@==0.2.0, setuptools@>=63, snakeviz@==2.2.2, tiktoken@==0.8.0, tokenizers@==0.20.3, tomli@==2.2.1, torch@==2.6.0.dev20250104, torchaudio@==2.6.0.dev20250104, torchsr, torchvision@==0.22.0.dev20250104, transformers@==4.46.1, wandb@==0.19.1, wheel, zstd@==1.5.6.1 ### Recent commits (newest first) - Update README.md - added custom kernels, made everything run - allow compilation even when executorch library is not defined - fix quickstart.py command typo - Update README.md - fix gitmodules - push math ref doc - fix pte_android.py - remove old profiler - upload profiling code - start challenge - fix typos - pre-release README ## Key source files (fetched from GitHub, selected and truncated for size) ### requirements.txt ``` # ----------------------------------------------------------------------------- # build tools # must be in sync with vendor/executorch/install_requirements.py # ----------------------------------------------------------------------------- -r requirements-cmake.txt pip>=23 pyyaml==6.0.2 setuptools>=63 tomli==2.2.1 wheel zstd==1.5.6.1 # ----------------------------------------------------------------------------- # torch and its deps # must be in sync with vendor/executorch/install_requirements.py # ----------------------------------------------------------------------------- transformers==4.46.1 torchsr # later index wins, liquid is later --extra-index-url https://download.pytorch.org/whl/nightly/cpu torch==2.6.0.dev20250104 torchvision==0.22.0.dev20250104 torchaudio==2.6.0.dev20250104 # ----------------------------------------------------------------------------- # export related # must be in sync with vendor/executorch/examples/models/llama/install_requirements.sh # Note: torchao is installed as a follow up step # ----------------------------------------------------------------------------- snakeviz==2.2.2 sentencepiece==0.2.0 lm_eval==0.4.5 tiktoken==0.8.0 blobfile==3.0.0 # Restore numpy if >= 2.0 numpy==1.21.3; python_version == '3.10' numpy<2.0; python_version >= '3.11' # to preserve the numpy version scipy==1.9; python_version == '3.10' scipy==1.13; python_version >= '3.11' # indirect deps, but important enough to be pinned safetensors==0.5.1 tokenizers==0.20.3 # ----------------------------------------------------------------------------- # v2 deps # ----------------------------------------------------------------------------- hydra-core==1.3.2 omegaconf pytest==8.3.4 einops==0.8.0 wandb==0.19.1 # ----------------------------------------------------------------------------- # encryption # ----------------------------------------------------------------------------- pycryptodome==3.21.0 boto3==1.35.88 # ----------------------------------------------------------------------------- # benchmark # ----------------------------------------------------------------------------- # formatting rich==13.9.4 # ----------------------------------------------------------------------------- # executorch wheel, from vendor/executorch/pyproject.toml # ----------------------------------------------------------------------------- expecttest flatbuffers hypothesis parameterized pytest pytest-xdist ruamel.yaml memory_profiler ipykernel matplotlib ``` ### fftconvtest.py ```python import torch import torch.nn as nn import fused_fftconv # Import your compiled extension class FFTConvModule(nn.Module): def __init__(self): super().__init__() def forward(self, x, filter): return fused_fftconv.fused_fftconv(x, filter) # Example input B, C, N, K = 4, 8, 128, 32 # Adjust as needed example_input = torch.randn(B, C, N, dtype=torch.float32) example_filter = torch.randn(C, K, dtype=torch.float32) # Instantiate module module = FFTConvModule() ``` ### pte_local.py ```python """ Local PTE Profiler This script provides functionality to profile PyTorch exported models locally. It measures execution time of forward method for the given PTE file. Usage: python pte_local.py --pte Example: python pte_local.py --pte model.pte """ import argparse from pathlib import Path import pprint import torch from quickstart import LocalPyProfiler def main(args): # prepare inputs inputs = (torch.ones(1, 5, dtype=torch.int64), torch.zeros(1, dtype=torch.int64)) # profile with LocalPyProfiler profiler = LocalPyProfiler(args.runner_path) profiling_result = profiler.profile(args.pte, inputs, repeats=4) pprint.pprint(profiling_result) if __name__ == "__main__": parser = argparse.ArgumentParser(description="Profile existing pte locally") parser.add_argument("--pte", type=Path, help="Path to pte file", required=True) parser.add_argument( "--runner_path", type=Path, help="path to pte_runner (for default is 'runner/linux/pte_runner')", required=True, ) args = parser.parse_args() main(args) ``` ### pte_android.py ```python #!/usr/bin/env python3 """ Profile PTE via Android. ## Local device: ``` PYTHONPATH=profiling \ --pte your_pte_file.pte \ --runner_path runner/android-arm64-v8a/pte_runner ``` """ import argparse import logging from pathlib import Path import pprint import torch from android_profiler import AndroidProfiler def main(args): # prepare inputs inputs = (torch.randn(128, dtype=torch.float32),) # instantiate profiler profiler = AndroidProfiler( runner_local_path=args.runner_path, device_id=args.device_id, device_work_dir=args.device_workdir or AndroidProfiler.device_work_dir, cpu_threads=args.cpu_threads, adb_port=args.adb_port, ignore_cpu_throttling=( args.ignore_cpu_throttling if args.ignore_cpu_throttling is not None else AndroidProfiler.ignore_cpu_throttling ), cpu_throttling_max_wait_secs=( args.cpu_throttling_max_wait_secs if args.cpu_throttling_max_wait_secs is not None else AndroidProfiler.cpu_throttling_max_wait_secs ), use_fixed_performance_mode=( args.use_fixed_performance_mode if args.use_fixed_performance_mode is not None else AndroidProfiler.use_fixed_performance_mode ), cpu_throttling_check_interval=( args.cpu_throttling_check_interval if args.cpu_throttling_check_interval is not None else AndroidProfiler.cpu_throttling_check_interval ), cpu_throttling_thermal_status_threshold=( args.cpu_throttling_thermal_status_threshold if args.cpu_throttling_thermal_status_threshold is not None else AndroidProfiler.cpu_throttling_thermal_status_threshold ), ) # profile profiling_result = profiler.profile(args.pte, inputs, repeats=args.repeats) # print results pprint.pprint(profiling_result) if __name__ == "__main__": parser = argparse.ArgumentParser(description="Profile existing pte via ssh") parser.add_argument("--pte", type=Path, help="Path to pte file", required=True) parser.add_argument( "--runner_path", type=Path, help="the pte runner local path", required=True ) parser.add_argument("--device_id", type=str, help="adb device id") parser.add_argument( "--device_workdir", type=Path, help="working directory on remote", ) parser.add_argument( "--cpu_threads", type=int, help="Concurrency level", default=AndroidProfiler.cpu_threads, ) parser.add_argument("--adb_port", type=int, help="Adb port if non-default required") parser.add_argument( "--repeats", type=int, help="Number of times to invoke the method", default=10 ) parser.add_argument( "--ignore_cpu_throttling", action="store_true", help="Disable throttling detection and cool down.", ) parser.add_argument( "--cpu_throttling_max_wait_secs", type=int, help="Max interval in seconds to wait for the device cool down.", ) parser.add_argument( "--use_fixed_performance_mode", action="store_true", help="Enable the fix performance mode", ) parser.add_argument( "--cpu_throttling_check_interval", type=int, help="Interval in seconds to check for the throttling conditions.", ) parser.add_argument( "--cpu_throttling_thermal_status_threshold", type=int, help="The thermal status, reaching which is considered throttling, use values from " "https://source.android.com/docs/core/power/thermal-mitigation#codes .", ) args = parser.parse_args() logging.basicConfig(level=logging.INFO) main(args) ``` ### quickstart.py ```python """ This provides a quickstart reference that shows how to export a module wrapping torch.sin to Edge Dialect, lower it to XNNPACK, and profile it using the local Python profiler. PYTHONPATH=profiling python quickstart.py """ from pathlib import Path from typing import Tuple, Union import os import subprocess import tempfile import torch import torch.nn as nn import numpy as np from torch.export import export, ExportedProgram from executorch.exir import EdgeProgramManager, to_edge from executorch.exir.backend.backend_api import LoweredBackendModule, to_backend from executorch.backends.xnnpack.utils.configs import get_xnnpack_edge_compile_config from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner from executorch.exir.backend.test.backend_with_compiler_demo import ( # noqa BackendWithCompilerDemo, ) from executorch.devtools import Inspector from profiling.profiler import Profiler, ProfilingResults, extract_stats class LocalPyProfiler(Profiler): """ Local Python profiler implementation for ExecutorTorch modules. """ def __init__( self, runner_path: Path, ): """ Initialize local profiler. Args: runner_path: Path to the profiler runner """ # Save arguments self._runner_path = runner_path def _load_pte(self, temp_dir: Path, pte: Union[bytes, Path]): if isinstance(pte, bytes): pte_path = temp_dir / "model.pte" with open(pte_path, "wb") as f: f.write(pte) return pte_path if isinstance(pte, Path): return pte elif isinstance(pte, str): pte_path = Path(pte) return pte_path raise ValueError(f"Invalid pte type: {type(pte)}, expected bytes or Path") def profile(self, pte: Union[bytes, Path], inputs: Tuple[torch.Tensor, ...], repeats: int) -> ProfilingResults: """ Invoke repeats times forward method of the pte and collect execution times. Args: pte: Compiled model in bytes or Path to the model file inputs: Tuple of input tensors for the model repeats: Number of times to repeat the profiling Returns: ProfilingResults containing statistical measures of the profiling run """ with tempfile.TemporaryDirectory() as temp_dir: temp_dir = Path(temp_dir) pte_path = self._load_pte(temp_dir, pte) inputs_path = temp_dir / "inputs.npz" np.savez(inputs_path, **{str(i): t.numpy() for i, t in enumerate(inputs)}) etdump_path = temp_dir / "etdump.bin" profile_cmd = [ self._runner_path, "-model_path", pte_path, "-etdump_path", etdump_path, "-inputs_npz_path", inputs_path, "-iter", f"{repeats}", ] try: subprocess.run(profile_cmd, capture_output=True, check=True, text=True) except subprocess.CalledProcessError as e: print(f"Command failed with exit code {e.returncode}") print(f"Error output:\n{e.stderr}") raise inspector = Inspector(etdump_path, debug_buffer_path=os.devnull) return extract_stats(inspector) class LowerableModule(nn.Module): def __init__(self): super().__init__() def forward(self, x): return torch.sin(x) import fused_fftconv # Import your compiled extension class FFTConvModule(nn.Module): def __init__(self): super().__init__() def forward(self, x, filter): out = torch.empty(x.shape, dtype=x.dtype, device=x.device) torch.ops.fftconv.fused_fftconv.out(x, filter, out=out) # Call the out= variant return out def run(): # Example input B, C, N, K = 4, 8, 128, 32 # Adjust as needed example_input = torch.randn(B, C, N, dtype=torch.float32) example_filter = torch.randn(C, K, dtype=torch.float32) from kernels.test_models import WinogradConvModule # Instantiate module module = FFTConvModule() B, C, N, K = 4, 8, 128, 32 # Adjust as needed example_input = torch.randn(B, C, N, dtype=torch.float32) example_filter = torch.randn(C, K, dtype=torch.float32) # Instantiate module module = WinogradConvModule() example_args = (example_input, example_filter) # Export and lower the module to Edge Dialect # example_args = (torch.randn(128, dtype=torch.float32),) example_args = (example_input, example_filter) # module = LowerableModule() aten_dialect_program: ExportedProgram = export(module, example_args) edge_config = get_xnnpack_edge_compile_config() edge_program: EdgeProgramManager = to_edge(aten_dialect_program, compile_config=edge_config) # Lower the module edge_manager_to_backend: LoweredBackendModule = edge_program.to_backend(XnnpackPartitioner()) print(edge_manager_to_backend) et_program = edge_manager_to_backend.to_executorch() print(et_program) # Serialize and save it to a file save_path = "test.pte" with open(save_path, "wb") as f: f.write(et_program.buffer) pte_runner_path = os.environ.get("PTE_RUNNER_PATH", 'runner/linux/pte_runner') profiler = LocalPyProfiler(pte_runner_path) profiling_result = profiler.profile(save_path, example_args, repeats=4) print(profiling_result) ``` ### kernels/test_models.py ```python import torch import torch.nn as nn import fused_fftconv # Import your compiled extension import generic_winograd_conv class FFTConvModule(nn.Module): def __init__(self): super().__init__() def forward(self, x, filter): out = torch.empty(x.shape, dtype=x.dtype, device=x.device) torch.ops.fftconv.fused_fftconv.out(x, filter, out=out) # Call the out= variant return out @torch._dynamo.disable() def opaque_winograd_conv(x, filt): B, C, N = x.shape K = filt.shape[1] out_len = N - K + 1 # Allocate output on the same device/dtype as x (force real allocation) out = torch.empty((B, C, out_len), dtype=x.dtype, device=x.device) # Call the custom op using its "out" variant. torch.ops.multiwinograd.generic_winograd_conv.out(x, filt, out=out) return out class WinogradConvModule(nn.Module): def __init__(self): super().__init__() def forward(self, x, filter): return opaque_winograd_conv(x, filter) ``` ### profiling/profiler.py ```python from abc import ABC, abstractmethod from dataclasses import dataclass from typing import List, Tuple, Union from pathlib import Path import torch from executorch.devtools import Inspector @dataclass(frozen=True) class ProfilingResults: """ Data class containing profiling results in milliseconds. """ raw: List[float] p10: float p50: float p90: float min: float avg: float max: float class Profiler(ABC): """ Abstract base class for implementing profilers. """ @abstractmethod def profile(self, pte: Union[bytes, Path], inputs: Tuple[torch.Tensor, ...], repeats: int) -> ProfilingResults: """ Profile the execution of a model with given inputs. Args: pte: Compiled model in bytes or Path to the compiled model file. inputs: Tuple of input tensors for the model. repeats: Number of times to repeat the profiling. Returns: ProfilingResults containing statistical measures of the profiling run. """ pass def extract_stats(inspector: Inspector) -> ProfilingResults: df = inspector.to_dataframe( include_units=False, include_delegate_debug_data=False, ) stats = df[df["event_name"] == "Method::execute"][["raw", "p10", "p50", "p90", "min", "avg", "max"]].to_dict( orient="records" )[0] if not stats: raise ValueError("No statistics found in inspector data") return ProfilingResults(**stats) ``` ### profiling/android_profiler.py ```python from pathlib import Path from typing import Optional, Tuple, Union import os import tempfile import logging import subprocess import contextlib import dataclasses import subprocess import hashlib from executorch.devtools import Inspector import numpy as np import torch from profiler import Profiler, ProfilingResults, extract_stats logger = logging.getLogger(__name__) @contextlib.contextmanager def pte_as_file(pte: Union[bytes, Path]): if isinstance(pte, bytes): with tempfile.NamedTemporaryFile(suffix=".pt") as tmp_file: tmp_file.write(pte) tmp_file.flush() yield Path(tmp_file.name) elif isinstance(pte, Path): yield pte else: raise ValueError(f"Invalid pte type: {type(pte)}, expected bytes or Path") @contextlib.contextmanager def input_as_npz_file(inputs: Tuple[torch.Tensor, ...] | Path): if isinstance(inputs, Path): yield inputs else: with tempfile.NamedTemporaryFile(suffix=".npz") as tmp_file: np.savez(tmp_file, **{str(i): t.numpy() for i, t in enumerate(inputs)}) tmp_file.flush() yield Path(tmp_file.name) @dataclasses.dataclass class AdbShell: device_id: Optional[str] = None port: Optional[str] = None port: Optional[str] = None def cmd(self, *args): r = ["adb"] if self.device_id: r.extend(["-s", str(self.device_id)]) if self.port: r.extend(["-P", str(self.port)]) r = ["adb"] if self.device_id: r.extend(["-s", str(self.device_id)]) if self.port: r.extend(["-P", str(self.port)]) r.extend(args) return r def calculate_adb_md5(self, path: Path): try: subprocess.check_output(self.cmd("shell", "stat", str(path))) except subprocess.CalledProcessError: logger.info(f"file {path} is not accessable on remote") return None # ba73f6793ff93db50abb82d1a0d1752a /data/local/tmp/foo out = subprocess.check_output(self.cmd("shell", "md5sum", str(path))) return out.split()[0].strip() def push_files(self, paths_mapping: dict[Path, Path]): for src, dst in paths_mapping.items(): def do_copy(): logger.info(f"pushing {src} -> {dst}") subprocess.check_call(self.cmd("shell", "mkdir", "-p", str(dst.parent))) subprocess.check_call(self.cmd("push", str(src), str(dst))) if src.stat().st_size < 1024**2: # copy small file do_copy() elif (remote_hash := self.calculate_adb_md5(dst)) is not None: if calculate_local_md5(src) != remote_hash: do_copy() else: logger.info(f"skip pusing identical {src} -> {dst}") else: do_copy() def pull_files(self, paths_mapping: dict[Path, Path]): for src, dst in paths_mapping.items(): logger.info(f"pulling {src} -> {dst}") os.makedirs(dst.parent, exist_ok=True) subprocess.check_call(self.cmd("pull", str(src), str(dst))) def delete_files(self, paths: list[Path]): for p in paths: logger.info(f"deleting {p}") subprocess.check_call(self.cmd("shell", "rm", str(p))) def make_executable(self, paths: list[Path]): for p in paths: logger.info(f"making executable {p}") subprocess.check_call(self.cmd("shell", "chmod", "+x", str(p))) def calculate_local_md5(file_path: Path): h = hashlib.md5() with file_path.open("rb") as f: for chunk in iter(lambda: f.read(4096), b""): h.update(chunk) return h.hexdigest().encode("utf-8") @dataclasses.dataclass class AndroidProfiler(Profiler): runner_local_path: Path device_id: Optional[str] = None device_work_dir: Path = Path("/data/local/tmp/profiling") cpu_threads: int = -1 adb_port: Optional[int] = None ignore_cpu_throttling: bool = False cpu_throttling_max_wait_secs: int = 60 * 5 use_fixed_performance_mode: bool = False cpu_throttling_check_interval: int = 1 cpu_throttling_thermal_status_threshold: int = 1 def profile( self, pte: Union[bytes, Path], inputs: Tuple[torch.Tensor, ...] | Path, repeats: int, ) -> ProfilingResults: assert self.device_work_dir adb = AdbShell(device_id=self.device_id, port=self.adb_port) # check the device connection subprocess.check_call( adb.cmd( "shell", " ".join( f'echo "{k}=$(getprop {k})";' for k in ( "ro.product.model", "ro.boot.serialno", "ro.vendor.build.fingerprint", ) ), ), ) with ( pte_as_file(pte) as pte_local_path, input_as_npz_file(inputs) as inputs_local_path, tempfile.NamedTemporaryFile() as etdump_local_file, ): runner_remote_path = self.device_work_dir / self.runner_local_path.name pte_remote_path = self.device_work_dir / pte_local_path.name inputs_remote_path = self.device_work_dir / inputs_local_path.name etdump_remote_path = self.device_work_dir / f"{pte_local_path.name}.etdump" etdump_local_path = Path(etdump_local_file.name) adb.push_files( { self.runner_local_path: runner_remote_path, pte_local_path: pte_remote_path, inputs_local_path: inputs_remote_path, } ) adb.make_executable([runner_remote_path]) pte_runner_args = [ str(runner_remote_path), "-model_path", str(pte_remote_path), [truncated — 1805 more characters] ``` ### examples/fftconv/ops/__init__.py ```python import torch from . import fused_fftconv __all__ = ['fused_fftconv'] ``` ### kernels/fftconv/ops/__init__.py ```python import torch from . import fused_fftconv __all__ = ['fused_fftconv'] ``` [10 more indexed source files omitted to keep this export small. The full file list is in the Codebase structure section above.]