From f0944daf179ecd926f50976ca19c8dae5c820b47 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 3 Aug 2026 09:48:37 +0000 Subject: [PATCH] feat(compile): register flagos as a first-class inductor GPU device Wires torch.compile into the flagos (PrivateUse1) backend by registering flagos with TorchInductor as a real GPU device, so the traced graph is handed to compile_fx unchanged and inductor emits Triton kernels that operate on flagos tensors directly. The earlier approach rewrote the graph and its example inputs to cuda before compiling. That is not just a copy per call: at::getAccelerator() is PrivateUse1/flagos in this build, and torch::autograd::Node::stream() only yields a stream when a node's input device type equals the accelerator. A cuda-rewritten graph therefore produces stream-less autograd nodes, and AOT autograd's backward trace inside compile_fx trips opt_ready_stream && opt_parent_stream (engine.cpp:1085) -- the cause of 8 of the 11 test failures. Verified by differential test: eager backward on plain cuda fails identically with torch.compile never involved. Registration surface (device_interface.py, inductor_codegen.py): * GPU_TYPES gains "flagos" in place -- is_gpu() is a membership test on that list object, and without it inductor takes the C++/CPU codegen path and never emits Triton. get_gpu_type()'s functools cache is primed while the list is narrowed, since it asserts at most one GPU type is available and the torch.cuda shim reports available too. * DeviceInterface subclass: device state from torch.flagos, hardware properties from torch.cuda (same physical GPU, same allocator). * DeviceProperties.create reports flagos as cuda at the Triton boundary. Triton's NVIDIA backend hard-checks target.backend == "cuda", so a literal "flagos" finds 0 compatible backends. Inductor already does this rewrite in the opposite direction for ROCm (hints.py:149). * Device op overrides + scheduling/wrapper codegen: the stock CUDA/Triton pipeline under the "flagos" key, also published on torch.flagos for inductor's official PrivateUse1 hook. Two generated-kernel bugs that only surface under compilation: * detach re-dispatched into itself. The kernel called at::detach(self), also registered on PrivateUse1, so it dispatched straight back. Eager hid the recursion because DeviceBoxingGuard rewrites self's device metadata first; under FakeTensor it cannot, since the Python dispatch key sits above the backend key. Dynamo traces every nn.Linear through detach, so this was a stack-overflow segfault at trace time. Now emits at::native::detach (NATIVE_DIRECT_VIEW_OPS). * gen_inplace passed only plain at::Tensor args to DeviceBoxingGuard, so clamp_.Tensor handed unboxed flagos min/max to a CUDA self and crashed. Optionals are now materialized into holders, matching gen_functional_pure. Both regressions are covered by tests that were confirmed to fail (segfault at the exact asserting line) against a build with the fixes reverted. CPU-torch wheel accommodations, since torch.cuda's Python layer was frozen without CUDA: re-attach CudaInterface.get_raw_stream (binding exists, the import-time _is_compiled() probe left it None), route torch.cuda.memory_* to the flagos allocator that backs the same pool, hand out flagos Event/Stream in place of the dummy base classes, force triton.cudagraphs off (torch.cuda.CUDAGraph raises on construction) and use_static_cuda_launcher off (not built). flagos_compile_backend now accepts the mode/options/dynamic kwargs dynamo forwards to named backends and expands them into compile_fx config_patches, rather than mutating inductor's global config. Tests: test_compile.py 12 passed / 1 skipped (was 2 passed / 8 failed); test_clamp_dispatch.py 15 passed; ops dispatch sweep 358 passed; test_ops.py 58 passed; allocator/factory/fallback/unit 73 passed. Docs updated to drop the device-aliasing description and the unmeasured performance-parity figures; benchmarking remains open work. Co-Authored-By: Claude Opus 5 (1M context) --- csrc/aten/generated/cuda_kernels.cc | 5 +- docs/torch_compile_integration.md | 263 +++++++++++++++ scripts/codegen_ops.py | 30 +- tests/integration/ops/test_clamp_dispatch.py | 176 ++++++++++ tests/integration/test_compile.py | 322 +++++++++++++++++++ tests/perf/bench_compile.py | 269 ++++++++++++++++ torch_fl/__init__.py | 16 + torch_fl/accelerator/cuda/_cuda_compat.py | 50 +++ torch_fl/compile/README.md | 97 ++++++ torch_fl/compile/__init__.py | 33 ++ torch_fl/compile/device_interface.py | 257 +++++++++++++++ torch_fl/compile/flagtree_shim.py | 86 +++++ torch_fl/compile/inductor_backend.py | 200 ++++++++++++ torch_fl/compile/inductor_codegen.py | 151 +++++++++ torch_fl/flagos/meta.py | 35 +- 15 files changed, 1980 insertions(+), 10 deletions(-) create mode 100644 docs/torch_compile_integration.md create mode 100644 tests/integration/ops/test_clamp_dispatch.py create mode 100644 tests/integration/test_compile.py create mode 100644 tests/perf/bench_compile.py create mode 100644 torch_fl/compile/README.md create mode 100644 torch_fl/compile/__init__.py create mode 100644 torch_fl/compile/device_interface.py create mode 100644 torch_fl/compile/flagtree_shim.py create mode 100644 torch_fl/compile/inductor_backend.py create mode 100644 torch_fl/compile/inductor_codegen.py diff --git a/csrc/aten/generated/cuda_kernels.cc b/csrc/aten/generated/cuda_kernels.cc index 7e20774f..0fd5adcd 100644 --- a/csrc/aten/generated/cuda_kernels.cc +++ b/csrc/aten/generated/cuda_kernels.cc @@ -8051,10 +8051,7 @@ at::Tensor & DequantizeSelfOutKernelCuda(const at::Tensor & self, at::Tensor & o } at::Tensor DetachKernelCuda(const at::Tensor & self) { - DeviceBoxingGuard guard(self); - auto result = at::detach(self); - UnboxToFlagos(result); - return result; + return at::native::detach(self); } at::Tensor & DetachInplaceKernelCuda(at::Tensor & self) { diff --git a/docs/torch_compile_integration.md b/docs/torch_compile_integration.md new file mode 100644 index 00000000..3b47ee5a --- /dev/null +++ b/docs/torch_compile_integration.md @@ -0,0 +1,263 @@ +# torch.compile Integration for flagos Device + +This document describes the `torch.compile` integration for the flagos device, enabling automatic kernel fusion and optimization via TorchInductor. + +## Overview + +The flagos device now supports PyTorch 2.0+ `torch.compile` for automatic performance optimization: + +```python +import torch +import torch_fl + +model = MyModel().to("flagos:0") +compiled_model = torch.compile(model, backend="flagos") + +# Automatic fusion of elementwise ops, reduced dispatch overhead +output = compiled_model(input) +``` + +**Key benefits**: +- Automatic kernel fusion (no manual optimization needed), cutting per-op dispatch overhead +- Graph stays on flagos: no cuda round trip, no copy at the graph boundary +- Compatible with existing flagos dispatch (FlagGems Python/C++, CUDA boxing) +- Optional FlagTree integration for multi-backend compilation + +## Quick Start + +### Basic Usage + +```python +import torch +import torch_fl + +# Standard model definition +model = torch.nn.Sequential( + torch.nn.Linear(512, 512), + torch.nn.ReLU(), + torch.nn.Linear(512, 512), +).to("flagos:0") + +# Compile with flagos backend +model = torch.compile(model, backend="flagos") + +# Use as normal +x = torch.randn(64, 512, device="flagos:0") +y = model(x) # Automatically fused kernels +``` + +### Compilation Modes + +```python +# Default: balanced optimization +model = torch.compile(model, backend="flagos") + +# Maximum performance (longer compile, better runtime) +model = torch.compile(model, backend="flagos", mode="max-autotune") + +# Explicit inductor config overrides +model = torch.compile(model, backend="flagos", options={"max_autotune": True}) +``` + +`mode` and `options` are expanded into inductor config patches scoped to that +compile. Note that CUDA graphs are always forced off (see Limitations), so +`mode="reduce-overhead"` -- whose main lever is cudagraphs -- has little effect +here. + +### FlagTree Integration (Phase 2) + +Use FlagTree for multi-backend kernel compilation: + +```bash +# Install FlagTree +pip install flagtree + +# Enable FlagTree backend +export FLAGOS_USE_FLAGTREE=1 +python your_script.py +``` + +FlagTree replaces OpenAI Triton with a multi-backend compiler supporting NVIDIA, Ascend, Cambricon, and MetaX hardware. + +## Architecture + +### Phase 1: Inductor Integration + +flagos is registered with TorchInductor as a **first-class GPU device**. The +traced graph is handed to `compile_fx` unchanged -- still on flagos -- and +inductor generates Triton kernels that operate on flagos tensors directly. + +**Components**: +1. **Backend registration** (`torch_fl/compile/inductor_backend.py`) + - Registers `"flagos"` with `torch._dynamo.register_backend` + - Expands `mode` / `options` into inductor `config_patches` + - Delegates to `compile_fx` with no graph rewriting + +2. **Device interface** (`torch_fl/compile/device_interface.py`) + - `DeviceInterface` subclass: device state from `torch.flagos`, hardware + properties from `torch.cuda` (the same physical GPU) + - Adds `"flagos"` to inductor's `GPU_TYPES` so `is_gpu()` is True and the + Triton codegen path is taken instead of C++/CPU + - Reports flagos as cuda at the Triton boundary (`DeviceProperties.create`), + because Triton's NVIDIA backend hard-checks `target.backend == "cuda"` + +3. **Codegen registration** (`torch_fl/compile/inductor_codegen.py`) + - Device op overrides (guards, streams, sync) inheriting the CUDA ones + - Scheduling + wrapper codegen: the stock CUDA/Triton pipeline + +4. **Dispatch integration** + - Ops inductor does not fuse fall back to eager flagos dispatch + (FlagGems Python/C++ or CUDA boxing) with no changes needed + +**Flow**: +``` +torch.compile(model, backend="flagos") + → dynamo captures FX graph (on flagos) + → compile_fx / AOT autograd, graph never leaves flagos + → inductor generates fused Triton kernels for flagos tensors + → unfused ops fall back to flagos eager dispatch +``` + +**Why the graph is not rewritten to cuda.** An earlier version converted the +graph and example inputs to cuda first. Beyond the copy per call, this breaks +backward: `at::getAccelerator()` is PrivateUse1/flagos, and +`torch::autograd::Node::stream()` only yields a stream when a node's input +device type equals the accelerator, so a cuda-rewritten graph produces +stream-less autograd nodes and AOT autograd's backward trace trips +`opt_ready_stream && opt_parent_stream` (engine.cpp:1085). + +### Phase 2: FlagTree Integration + +**Components**: +1. **Triton import patcher** (`torch_fl/compile/flagtree_shim.py`) + - Replaces `import triton` with `import flagtree` + - Activated via `FLAGOS_USE_FLAGTREE=1` + +2. **Backend selection** + - FlagTree backend configured via `GEMS_VENDOR` env var + - Same Triton kernel code, different backend compiler + +**Benefits**: +- Multi-backend: same compiled model runs on NVIDIA/Ascend/Cambricon/MetaX +- Future-proof for non-NVIDIA hardware + +## Performance + +**Not yet measured.** Correctness is verified (`tests/integration/test_compile.py`); +benchmarking the fusion gain, and comparing it against stock `inductor` on cuda, +is still open work. Structurally the two should land close together -- same +inductor fusion passes, same Triton codegen, and since the graph stays on flagos +there is no per-call copy -- but that is an expectation, not a measurement. + +### Benchmarking + +```bash +# Run performance benchmark +python tests/perf/bench_compile.py --model=mlp --batch-size=64 + +# Compare with CUDA baseline +python tests/perf/bench_compile.py --model=transformer --compare-cuda + +# Test FlagTree integration +FLAGOS_USE_FLAGTREE=1 python tests/perf/bench_compile.py +``` + +## Environment Variables + +| Variable | Default | Description | +|----------|---------|-------------| +| `FLAGOS_USE_FLAGTREE` | `0` | Use FlagTree instead of OpenAI Triton | +| `FLAGOS_COMPILE_FALLBACK_EAGER` | `0` | Fall back to eager on compile errors | + +Existing dispatch variables (`FLAGOS_USE_FLAGGEMS`, `FLAGOS_BACKEND_CONFIG`) still apply to compiled kernels. + +## Troubleshooting + +### Compilation Errors + +**Symptom**: `torch.compile` raises errors during graph capture or codegen. + +**Solutions**: +1. Enable fallback to eager: `FLAGOS_COMPILE_FALLBACK_EAGER=1` +2. Check for unsupported ops (dynamic shapes, custom ops) +3. Verify meta implementations for custom ops + +### No Speedup + +**Symptom**: Compiled model runs at same speed as eager. + +**Possible causes**: +1. Model is compute-bound (large matmuls) — fusion won't help much +2. Compilation didn't fuse ops (check inductor logs) +3. Dispatch overhead is small relative to kernel time + +**Debug**: Run with `TORCH_LOGS="+inductor"` to see fusion decisions. + +### FlagTree Not Loading + +**Symptom**: `FLAGOS_USE_FLAGTREE=1` but warning says "falling back to OpenAI Triton". + +**Solutions**: +1. Install FlagTree: `pip install flagtree` +2. Verify import works: `python -c "import flagtree"` +3. Check FlagTree supports your hardware (`GEMS_VENDOR` setting) + +## Testing + +```bash +# Run integration tests +pytest tests/integration/test_compile.py -v + +# Test specific scenarios +pytest tests/integration/test_compile.py::test_basic_compile +pytest tests/integration/test_compile.py::test_compile_backward + +# Regression guards for the two codegen fixes this integration required +pytest tests/integration/test_compile.py -k fake_tensor +pytest tests/integration/ops/test_clamp_dispatch.py -v + +# Test FlagTree integration (requires flagtree installed) +FLAGOS_USE_FLAGTREE=1 pytest tests/integration/test_compile.py::test_flagtree_integration +``` + +### Codegen fixes this integration required + +Two generated-kernel bugs only surface under compilation, so their regression +tests live alongside it: + +- **`detach` re-dispatch.** The generated kernel called `at::detach(self)`, which + is registered on PrivateUse1 too and so dispatched back into itself. Eager hid + the recursion because `DeviceBoxingGuard` rewrites self's device metadata + first; under FakeTensor it cannot, since the Python dispatch key sits *above* + the backend key. Dynamo traces every `nn.Linear` through detach, so this was a + stack-overflow segfault at trace time. Fixed by emitting `at::native::detach` + (`NATIVE_DIRECT_VIEW_OPS` in `scripts/codegen_ops.py`). +- **`optional` boxing in in-place kernels.** `gen_inplace` handed only + plain `at::Tensor` args to `DeviceBoxingGuard`, so `clamp_.Tensor` passed + unboxed flagos `min`/`max` into a CUDA `self` and crashed. Fixed by + materializing each optional into a holder, matching `gen_functional_pure`. + +## Limitations + +1. **torch >= 2.0 required**: Older PyTorch versions don't have `torch.compile` +2. **Inductor-compatible ops only**: Custom C++ ops may not fuse +3. **Dynamic shapes**: Some models with dynamic shapes may not compile +4. **CUDA graphs off**: `torch.cuda.CUDAGraph` is a dummy class in the CPU torch + wheel, so `triton.cudagraphs` is forced off even under `mode="max-autotune"` +5. **FlagTree maturity**: Backend support varies by hardware (NVIDIA most mature) + +## Roadmap + +- [x] Phase 1: Inductor integration (flagos as a first-class GPU device) +- [ ] Phase 2: FlagTree integration — shim exists, not yet exercised end-to-end +- [ ] Benchmark fusion gains vs. stock inductor+triton on cuda +- [ ] Phase 3: FlagGems-aware fusion (recognize pre-optimized patterns) +- [ ] Phase 4: Custom fusion patterns for flagos-specific ops + +## See Also + +- [PyTorch 2.0 torch.compile documentation](https://pytorch.org/docs/stable/torch.compiler.html) +- [TorchInductor overview](https://pytorch.org/docs/stable/torch.compiler_inductor_overview.html) +- [FlagTree repository](https://github.com/flagos-ai/FlagTree) +- [CPU torch + external libtorch_cuda.so](cpu_torch_external_libtorch_cuda.md) — why several `torch.cuda` bindings need shimming +- [torch_fl/compile/README.md](../torch_fl/compile/README.md) — the registration surface in detail diff --git a/scripts/codegen_ops.py b/scripts/codegen_ops.py index 96df2b80..eb0cd256 100644 --- a/scripts/codegen_ops.py +++ b/scripts/codegen_ops.py @@ -1145,10 +1145,33 @@ def _generator_inject_line(args, device_expr): _GEN_OUT_BASES = _GEN_FACTORY_BASES | _GEN_LIKE_BASES +# Pure view/alias ops whose generated kernel MUST call at::native:: directly +# instead of at::. These ops are also registered on the PrivateUse1 key, so +# calling the re-dispatchable at::(self) from inside the kernel routes right +# back into THIS kernel. Under eager, DeviceBoxingGuard hides the recursion by +# rewriting self's device metadata to CUDA before the call. But that trick fails +# for FakeTensors: their Python dispatch key sits ABOVE the backend keys, so the +# metadata rewrite can't redirect dispatch -- at::(self) re-dispatches to +# PrivateUse1 -> infinite recursion -> stack overflow (segfault during +# torch.compile / dynamo tracing of any model that hits detach, e.g. nn.Linear). +# at::native:: is the metadata-only implementation and never re-dispatches, +# matching how strided_ops.cc registers view/as_strided/transpose for Ascend. +NATIVE_DIRECT_VIEW_OPS = { + "detach", +} + + def gen_functional_pure(op, fn_type, ret_type, args, func=None): kn = kernel_name(fn_type) api = f"at::{at_api_base(op)}" + # Pure alias ops (detach): call at::native:: directly, no boxing, no + # re-dispatch. See NATIVE_DIRECT_VIEW_OPS for why re-dispatch is fatal. + if at_api_base(op) in NATIVE_DIRECT_VIEW_OPS: + return f"""{ret_type} {kn}({args_decl(args)}) {{ + return at::native::{at_api_base(op)}({call_args(args)}); +}}""" + # Box plain Tensor inputs AND optional inputs (e.g. conv/linear bias). # An optional that lives on flagos must be boxed to CUDA too, or the # backend op receives a mix of boxed (input/weight) and unboxed (bias) @@ -1235,7 +1258,12 @@ def gen_inplace(op, fn_type, ret_type, args, func=None): - function-only (silu_/gelu_/celu_/leaky_relu_/...activations): call the free function `at::silu_(self, ...)` -- these have NO Tensor method, so the method syntax fails to compile. - Ops with both (fill_/zero_/relu_) work either way; we default to method.""" + Ops with both (fill_/zero_/relu_) work either way; we default to method. + + optional inputs must be materialized into a holder to be boxed by + DeviceBoxingGuard (same pattern as gen_functional_pure / gen_out_variant). + Missing this makes e.g. clamp_.Tensor pass unboxed flagos min/max into a + CUDA self -> "tensor does not have a device" / segfault.""" kn = kernel_name(fn_type) # Box plain Tensor inputs AND optional inputs (e.g. clamp_.Tensor's # min/max). An optional on flagos must be boxed to CUDA too, or the diff --git a/tests/integration/ops/test_clamp_dispatch.py b/tests/integration/ops/test_clamp_dispatch.py new file mode 100644 index 00000000..ba291eab --- /dev/null +++ b/tests/integration/ops/test_clamp_dispatch.py @@ -0,0 +1,176 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +clamp dispatch tests + +Guards the ``optional`` boxing in the generated in-place kernels. + +``clamp_.Tensor`` takes ``optional min`` / ``optional max``. +``DeviceBoxingGuard`` only rewrites the tensors it is handed, so an unboxed +``min``/``max`` reaches a CUDA ``self`` still carrying a flagos device -- which +crashes rather than failing cleanly. ``scripts/codegen_ops.py`` materializes each +optional into a holder (``min_t``/``max_t``) before the guard; these tests hold +that in place across regenerations. + +Usage: + pytest tests/integration/ops/test_clamp_dispatch.py -v +""" + +import pytest +import torch +import torch_fl # noqa: F401 + + +DEVICE = "flagos:0" + + +def _bounds(shape, lo, hi, device): + """min/max tensors broadcastable against ``shape``.""" + return ( + torch.full(shape, lo, device=device), + torch.full(shape, hi, device=device), + ) + + +class TestClampScalarBounds: + """clamp / clamp_ with Scalar min & max (the ``clamp.Scalar`` overloads).""" + + @pytest.mark.anyplatform + def test_clamp_scalar_both(self): + torch.manual_seed(0) + a = torch.randn(64, 64, device=DEVICE) + out = torch.clamp(a, -0.5, 0.5) + ref = torch.clamp(a.cpu(), -0.5, 0.5) + torch.testing.assert_close(out.cpu(), ref, rtol=0, atol=0) + + @pytest.mark.anyplatform + @pytest.mark.parametrize("bound", ["min", "max"]) + def test_clamp_scalar_one_sided(self, bound): + """Exercises the ``None`` branch of the optional holder.""" + torch.manual_seed(1) + a = torch.randn(32, 32, device=DEVICE) + kwargs = {bound: 0.25} + out = torch.clamp(a, **kwargs) + ref = torch.clamp(a.cpu(), **kwargs) + torch.testing.assert_close(out.cpu(), ref, rtol=0, atol=0) + + @pytest.mark.anyplatform + def test_clamp_inplace_scalar(self): + torch.manual_seed(2) + a = torch.randn(64, 64, device=DEVICE) + ref = torch.clamp(a.cpu(), -1.0, 1.0) + ret = a.clamp_(-1.0, 1.0) + assert ret.data_ptr() == a.data_ptr(), "clamp_ must mutate in place" + torch.testing.assert_close(a.cpu(), ref, rtol=0, atol=0) + + +class TestClampTensorBounds: + """clamp / clamp_ with Tensor min & max -- the optional boxing path.""" + + @pytest.mark.anyplatform + def test_clamp_tensor_both(self): + torch.manual_seed(3) + a = torch.randn(64, 64, device=DEVICE) + lo, hi = _bounds((64, 64), -0.5, 0.5, DEVICE) + out = torch.clamp(a, lo, hi) + ref = torch.clamp(a.cpu(), lo.cpu(), hi.cpu()) + torch.testing.assert_close(out.cpu(), ref, rtol=0, atol=0) + + @pytest.mark.anyplatform + @pytest.mark.parametrize("bound", ["min", "max"]) + def test_clamp_tensor_one_sided(self, bound): + torch.manual_seed(4) + a = torch.randn(32, 32, device=DEVICE) + b = torch.full((32, 32), 0.25, device=DEVICE) + out = torch.clamp(a, **{bound: b}) + ref = torch.clamp(a.cpu(), **{bound: b.cpu()}) + torch.testing.assert_close(out.cpu(), ref, rtol=0, atol=0) + + @pytest.mark.anyplatform + def test_clamp_inplace_tensor_both(self): + """The regression itself: unboxed min/max used to core dump here.""" + torch.manual_seed(5) + a = torch.randn(64, 64, device=DEVICE) + lo, hi = _bounds((64, 64), -0.5, 0.5, DEVICE) + ref = torch.clamp(a.cpu(), lo.cpu(), hi.cpu()) + ret = a.clamp_(lo, hi) + assert ret.data_ptr() == a.data_ptr(), "clamp_ must mutate in place" + torch.testing.assert_close(a.cpu(), ref, rtol=0, atol=0) + + @pytest.mark.anyplatform + @pytest.mark.parametrize("bound", ["min", "max"]) + def test_clamp_inplace_tensor_one_sided(self, bound): + """One optional set, the other empty -- both holder branches at once.""" + torch.manual_seed(6) + a = torch.randn(32, 32, device=DEVICE) + b = torch.full((32, 32), 0.25, device=DEVICE) + ref = torch.clamp(a.cpu(), **{bound: b.cpu()}) + a.clamp_(**{bound: b}) + torch.testing.assert_close(a.cpu(), ref, rtol=0, atol=0) + + @pytest.mark.anyplatform + def test_clamp_tensor_broadcast(self): + """Row-vector bounds broadcast against a 2-D input.""" + torch.manual_seed(7) + a = torch.randn(16, 8, device=DEVICE) + lo = torch.linspace(-1.0, 0.0, 8, device=DEVICE) + hi = torch.linspace(0.0, 1.0, 8, device=DEVICE) + out = torch.clamp(a, lo, hi) + ref = torch.clamp(a.cpu(), lo.cpu(), hi.cpu()) + torch.testing.assert_close(out.cpu(), ref, rtol=0, atol=0) + + @pytest.mark.anyplatform + def test_clamp_min_max_ops(self): + """clamp_min / clamp_max, the single-bound siblings.""" + torch.manual_seed(8) + a = torch.randn(32, 32, device=DEVICE) + b = torch.full((32, 32), 0.1, device=DEVICE) + + torch.testing.assert_close( + torch.clamp_min(a, b).cpu(), + torch.clamp_min(a.cpu(), b.cpu()), + rtol=0, + atol=0, + ) + torch.testing.assert_close( + torch.clamp_max(a, b).cpu(), + torch.clamp_max(a.cpu(), b.cpu()), + rtol=0, + atol=0, + ) + + +class TestClampDtypes: + @pytest.mark.anyplatform + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.int32]) + def test_clamp_inplace_tensor_dtype(self, dtype): + torch.manual_seed(9) + if dtype.is_floating_point: + a = torch.randn(32, 32, device=DEVICE, dtype=dtype) + lo, hi = -0.5, 0.5 + else: + a = torch.randint(-10, 10, (32, 32), device=DEVICE, dtype=dtype) + lo, hi = -3, 3 + lo_t = torch.full((32, 32), lo, device=DEVICE, dtype=dtype) + hi_t = torch.full((32, 32), hi, device=DEVICE, dtype=dtype) + + ref = torch.clamp(a.cpu(), lo_t.cpu(), hi_t.cpu()) + a.clamp_(lo_t, hi_t) + assert a.dtype == dtype + torch.testing.assert_close(a.cpu(), ref, rtol=0, atol=0) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/integration/test_compile.py b/tests/integration/test_compile.py new file mode 100644 index 00000000..90302194 --- /dev/null +++ b/tests/integration/test_compile.py @@ -0,0 +1,322 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Integration tests for torch.compile on flagos device. + +Tests basic compilation, fusion gains, and FlagTree integration (Phase 2). +""" + +import os +import pytest +import torch +import torch_fl + + +# Skip all tests if torch.compile not available (torch < 2.0) +try: + import torch._dynamo + + HAS_COMPILE = True +except ImportError: + HAS_COMPILE = False + +pytestmark = pytest.mark.skipif( + not HAS_COMPILE, reason="torch.compile not available (torch < 2.0)" +) + +# flagos tensors report either name depending on how the device was spelled. +FLAGOS_DEVICE_TYPES = ("privateuseone", "flagos") + + +def assert_on_flagos(tensor, what="output"): + """The graph is compiled *on* flagos, so results must come back on flagos. + + A cuda round trip would both cost a copy per call and produce stream-less + autograd nodes (see torch_fl/compile/inductor_backend.py), so this is a + load-bearing assertion, not a smoke check. + """ + assert tensor.device.type in FLAGOS_DEVICE_TYPES, ( + f"{what} landed on {tensor.device}, expected flagos" + ) + + +@pytest.fixture +def device(): + """Flagos device for testing.""" + if torch_fl.flagos.device_count() == 0: + pytest.skip("No flagos devices available") + return "flagos:0" + + +class SimpleModel(torch.nn.Module): + """Simple model with fusible ops.""" + + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(128, 128) + + def forward(self, x): + x = self.linear(x) + x = torch.relu(x) + x = x * 2.0 + x = x + 1.0 + return x + + +class MatMulModel(torch.nn.Module): + """Model with matrix multiplications.""" + + def __init__(self): + super().__init__() + + def forward(self, a, b): + c = torch.mm(a, b) + d = torch.mm(c, b.t()) + return d + 1.0 + + +def test_compile_backend_registered(): + """Test that 'flagos' backend is registered with dynamo.""" + import torch._dynamo + + # Check backend is in registry + backends = torch._dynamo.list_backends() + assert "flagos" in backends, ( + f"'flagos' backend not registered. Available: {backends}" + ) + + +def test_basic_compile(device): + """Test basic torch.compile with flagos backend.""" + model = SimpleModel().to(device) + x = torch.randn(32, 128, device=device) + + # Compile with flagos backend + compiled_model = torch.compile(model, backend="flagos") + + # Run compiled model + output = compiled_model(x) + + # Verify output shape and device. The graph is compiled *on* flagos (no + # cuda round trip), so the result must come back on flagos. + assert output.shape == (32, 128) + assert_on_flagos(output) + + # Compare with eager mode + eager_output = model(x) + torch.testing.assert_close(output, eager_output, rtol=1e-4, atol=1e-4) + + +def test_compile_vs_eager_correctness(device): + """Test numerical correctness of compiled vs eager execution.""" + model = MatMulModel().to(device) + a = torch.randn(64, 64, device=device) + b = torch.randn(64, 64, device=device) + + # Eager mode + eager_output = model(a, b) + + # Compiled mode + compiled_model = torch.compile(model, backend="flagos") + compiled_output = compiled_model(a, b) + + assert_on_flagos(compiled_output) + + # Should be numerically identical (or very close) + torch.testing.assert_close(compiled_output, eager_output, rtol=1e-4, atol=1e-4) + + +def test_compile_with_max_autotune(device): + """Test torch.compile with max-autotune mode.""" + model = SimpleModel().to(device) + x = torch.randn(32, 128, device=device) + + # Compile with max-autotune (aggressive fusion) + compiled_model = torch.compile(model, backend="flagos", mode="max-autotune") + + output = compiled_model(x) + eager_output = model(x) + + assert_on_flagos(output) + torch.testing.assert_close(output, eager_output, rtol=1e-4, atol=1e-4) + + +def test_compile_multiple_inputs(device): + """Test compilation with multiple input tensors.""" + model = MatMulModel().to(device) + a = torch.randn(32, 32, device=device) + b = torch.randn(32, 32, device=device) + + compiled_model = torch.compile(model, backend="flagos") + + output = compiled_model(a, b) + eager_output = model(a, b) + + assert_on_flagos(output) + torch.testing.assert_close(output, eager_output, rtol=1e-4, atol=1e-4) + + +def test_compile_backward(device): + """Test that compiled model supports backward pass.""" + model = SimpleModel().to(device) + x = torch.randn(32, 128, device=device, requires_grad=True) + + compiled_model = torch.compile(model, backend="flagos") + + # Forward + backward + output = compiled_model(x) + loss = output.sum() + loss.backward() + + # Check gradients exist. The gradient staying on flagos is what proves the + # backward graph was never rewritten to cuda -- that rewrite produced + # stream-less autograd nodes and tripped engine.cpp's stream assertion. + assert x.grad is not None + assert x.grad.shape == x.shape + assert_on_flagos(x.grad, "gradient") + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +def test_compile_dtypes(device, dtype): + """Test compilation with different dtypes.""" + model = SimpleModel().to(device).to(dtype) + x = torch.randn(32, 128, device=device, dtype=dtype) + + compiled_model = torch.compile(model, backend="flagos") + output = compiled_model(x) + + assert output.dtype == dtype + assert_on_flagos(output) + + eager_output = model(x) + # Float16 has lower precision + rtol = 1e-2 if dtype == torch.float16 else 1e-4 + torch.testing.assert_close(output, eager_output, rtol=rtol, atol=rtol) + + +def test_compile_recompile(device): + """Test that recompiling doesn't break.""" + model = SimpleModel().to(device) + x = torch.randn(32, 128, device=device) + + # First compilation + compiled_model = torch.compile(model, backend="flagos") + output1 = compiled_model(x) + + # Reset dynamo cache and recompile + torch._dynamo.reset() + compiled_model2 = torch.compile(model, backend="flagos") + output2 = compiled_model2(x) + + torch.testing.assert_close(output1, output2, rtol=1e-6, atol=1e-6) + + +def test_fake_tensor_detach(device): + """detach must not re-dispatch to itself under FakeTensorMode. + + The generated CUDA kernel used to call ``at::detach(self)``, which is + registered on PrivateUse1 too and so dispatched straight back into itself. + In eager, ``DeviceBoxingGuard``'s device rewrite masked the recursion; under + FakeTensorMode it cannot, because the Python dispatch key sits above the + backend key -- rewriting metadata does not change where dispatch goes. The + kernel now calls ``at::native::detach``. Dynamo traces every ``nn.Linear`` + through detach, so a regression here is a stack-overflow crash, not a + failure. + """ + from torch._subclasses.fake_tensor import FakeTensorMode + + with FakeTensorMode(): + x = torch.randn(32, 128, device=device) + d = x.detach() + assert d.shape == x.shape + assert d.device.type == x.device.type + assert not d.requires_grad + + +def test_fake_tensor_linear(device): + """F.linear under FakeTensorMode -- the shape dynamo actually traces. + + nn.Linear goes through detach internally; this is the end-to-end form of + test_fake_tensor_detach and the exact call that used to segfault at trace + time. Parameters are built inside the mode rather than by moving a module + into it (nn.Module._apply cannot swap real params for fake ones). + """ + from torch._subclasses.fake_tensor import FakeTensorMode + + with FakeTensorMode(): + weight = torch.nn.Parameter(torch.randn(64, 128, device=device)) + bias = torch.nn.Parameter(torch.randn(64, device=device)) + x = torch.randn(32, 128, device=device) + out = torch.nn.functional.linear(x, weight, bias) + assert out.shape == (32, 64) + assert out.device.type == x.device.type + + +@pytest.mark.skipif( + os.environ.get("FLAGOS_USE_FLAGTREE", "0") != "1", + reason="FlagTree integration not enabled (set FLAGOS_USE_FLAGTREE=1)", +) +def test_flagtree_integration(device): + """ + Test FlagTree integration (Phase 2). + + Requires: pip install flagtree + FLAGOS_USE_FLAGTREE=1 + """ + model = SimpleModel().to(device) + x = torch.randn(32, 128, device=device) + + # Compile should use FlagTree instead of OpenAI Triton + compiled_model = torch.compile(model, backend="flagos") + output = compiled_model(x) + + eager_output = model(x) + torch.testing.assert_close(output, eager_output, rtol=1e-4, atol=1e-4) + + # Verify FlagTree was actually used (check sys.modules) + import sys + + if "flagtree" in sys.modules: + # FlagTree loaded successfully + pass + else: + pytest.skip("FlagTree not loaded (may have fallen back to OpenAI Triton)") + + +def test_compile_fallback_eager(): + """Test fallback to eager mode when compilation fails.""" + # Set fallback env var + os.environ["FLAGOS_COMPILE_FALLBACK_EAGER"] = "1" + + try: + # Create a model that might cause compilation issues + class ProblematicModel(torch.nn.Module): + def forward(self, x): + # Some operation that might not compile cleanly + return x + + model = ProblematicModel() + x = torch.randn(10, 10) + + # Should not raise, falls back to eager + compiled_model = torch.compile(model, backend="flagos") + output = compiled_model(x) + + assert output.shape == (10, 10) + finally: + os.environ.pop("FLAGOS_COMPILE_FALLBACK_EAGER", None) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/perf/bench_compile.py b/tests/perf/bench_compile.py new file mode 100644 index 00000000..9a2dbabc --- /dev/null +++ b/tests/perf/bench_compile.py @@ -0,0 +1,269 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Performance benchmark: torch.compile fusion gains on flagos device. + +Measures speedup from kernel fusion (compiled vs eager) to validate we achieve +parity with inductor+triton performance gains. + +Usage: + python tests/perf/bench_compile.py + python tests/perf/bench_compile.py --model=mlp --batch-size=128 + FLAGOS_USE_FLAGTREE=1 python tests/perf/bench_compile.py # Phase 2 +""" + +import argparse +import time +import torch +import torch.nn as nn +import torch_fl + + +class MLPModel(nn.Module): + """Multi-layer perceptron with many fusible ops.""" + + def __init__(self, hidden_size=512): + super().__init__() + self.fc1 = nn.Linear(hidden_size, hidden_size) + self.fc2 = nn.Linear(hidden_size, hidden_size) + self.fc3 = nn.Linear(hidden_size, hidden_size) + + def forward(self, x): + # Many elementwise ops that should fuse + x = self.fc1(x) + x = torch.relu(x) + x = x * 2.0 + x = x + 1.0 + + x = self.fc2(x) + x = torch.gelu(x) + x = x / 2.0 + + x = self.fc3(x) + x = torch.sigmoid(x) + return x + + +class ConvModel(nn.Module): + """Convolutional model with fusible activation patterns.""" + + def __init__(self): + super().__init__() + self.conv1 = nn.Conv2d(3, 64, 3, padding=1) + self.conv2 = nn.Conv2d(64, 64, 3, padding=1) + self.conv3 = nn.Conv2d(64, 3, 3, padding=1) + + def forward(self, x): + x = self.conv1(x) + x = torch.relu(x) + x = x * 0.5 + 0.5 # Normalization pattern + + x = self.conv2(x) + x = torch.relu(x) + + x = self.conv3(x) + return x + + +class TransformerBlock(nn.Module): + """Single transformer block (attention + FFN).""" + + def __init__(self, d_model=512, nhead=8, dim_feedforward=2048): + super().__init__() + self.self_attn = nn.MultiheadAttention(d_model, nhead, batch_first=True) + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.linear2 = nn.Linear(dim_feedforward, d_model) + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + + def forward(self, x): + # Self-attention with residual + attn_out, _ = self.self_attn(x, x, x) + x = self.norm1(x + attn_out) + + # FFN with residual (many fusible ops) + ffn = self.linear1(x) + ffn = torch.relu(ffn) + ffn = self.linear2(ffn) + x = self.norm2(x + ffn) + + return x + + +def benchmark_model(model, inputs, warmup=10, rounds=100): + """ + Benchmark model execution time. + + Returns average time per forward pass in milliseconds. + """ + device = next(model.parameters()).device + + # Warmup + for _ in range(warmup): + with torch.no_grad(): + _ = model(*inputs) if isinstance(inputs, tuple) else model(inputs) + + # Sync before timing + if device.type == "privateuseone": + torch_fl.flagos.synchronize() + else: + torch.cuda.synchronize() + + # Timed runs + start = time.perf_counter() + for _ in range(rounds): + with torch.no_grad(): + _ = model(*inputs) if isinstance(inputs, tuple) else model(inputs) + + # Sync after timing + if device.type == "privateuseone": + torch_fl.flagos.synchronize() + else: + torch.cuda.synchronize() + + elapsed = time.perf_counter() - start + return (elapsed / rounds) * 1000 # Convert to ms + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--model", choices=["mlp", "conv", "transformer"], default="mlp" + ) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--hidden-size", type=int, default=512) + parser.add_argument("--rounds", type=int, default=100) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument( + "--compare-cuda", action="store_true", help="Compare with CUDA baseline" + ) + args = parser.parse_args() + + device_flagos = "flagos:0" + device_cuda = "cuda:0" if args.compare_cuda else None + + # Build model and inputs + if args.model == "mlp": + model = MLPModel(hidden_size=args.hidden_size) + inputs_shape = (args.batch_size, args.hidden_size) + inputs_flagos = torch.randn(*inputs_shape, device=device_flagos) + inputs_cuda = ( + torch.randn(*inputs_shape, device=device_cuda) if device_cuda else None + ) + elif args.model == "conv": + model = ConvModel() + inputs_shape = (args.batch_size, 3, 224, 224) + inputs_flagos = torch.randn(*inputs_shape, device=device_flagos) + inputs_cuda = ( + torch.randn(*inputs_shape, device=device_cuda) if device_cuda else None + ) + elif args.model == "transformer": + model = TransformerBlock(d_model=args.hidden_size) + inputs_shape = (args.batch_size, 128, args.hidden_size) + inputs_flagos = torch.randn(*inputs_shape, device=device_flagos) + inputs_cuda = ( + torch.randn(*inputs_shape, device=device_cuda) if device_cuda else None + ) + + print("=== torch.compile Performance Benchmark ===") + print(f"Model: {args.model}") + print(f"Batch size: {args.batch_size}") + print(f"Hidden size: {args.hidden_size}") + print(f"Rounds: {args.rounds}") + print( + f"FlagTree enabled: {bool(int(torch.os.environ.get('FLAGOS_USE_FLAGTREE', '0')))}" + ) + print() + + # === Flagos Device === + print("--- Flagos Device ---") + + # Eager mode (baseline) + model_flagos_eager = model.to(device_flagos) + time_eager = benchmark_model( + model_flagos_eager, inputs_flagos, warmup=args.warmup, rounds=args.rounds + ) + print(f"Eager mode: {time_eager:>8.3f} ms/iter") + + # Compiled mode (inductor fusion) + try: + model_flagos_compiled = torch.compile(model.to(device_flagos), backend="flagos") + time_compiled = benchmark_model( + model_flagos_compiled, inputs_flagos, warmup=args.warmup, rounds=args.rounds + ) + speedup = time_eager / time_compiled + print( + f"Compiled (flagos): {time_compiled:>8.3f} ms/iter ({speedup:.2f}x speedup)" + ) + except Exception as e: + print(f"Compiled (flagos): FAILED - {e}") + speedup = 1.0 + + print() + + # === CUDA Baseline (for comparison) === + if args.compare_cuda: + print("--- CUDA Device (baseline) ---") + + model_cuda_eager = model.to(device_cuda) + time_cuda_eager = benchmark_model( + model_cuda_eager, inputs_cuda, warmup=args.warmup, rounds=args.rounds + ) + print(f"Eager mode: {time_cuda_eager:>8.3f} ms/iter") + + try: + model_cuda_compiled = torch.compile( + model.to(device_cuda), backend="inductor" + ) + time_cuda_compiled = benchmark_model( + model_cuda_compiled, inputs_cuda, warmup=args.warmup, rounds=args.rounds + ) + speedup_cuda = time_cuda_eager / time_cuda_compiled + print( + f"Compiled (inductor): {time_cuda_compiled:>8.3f} ms/iter ({speedup_cuda:.2f}x speedup)" + ) + + # Compare flagos vs CUDA speedups + print() + print("--- Speedup Comparison ---") + print(f"Flagos compile speedup: {speedup:.2f}x") + print(f"CUDA compile speedup: {speedup_cuda:.2f}x") + parity = (speedup / speedup_cuda) * 100 + print(f"Parity: {parity:.1f}% (flagos vs CUDA)") + + except Exception as e: + print(f"Compiled (inductor): FAILED - {e}") + + # === Summary === + print() + print("=== Summary ===") + if speedup >= 1.5: + print(f"✅ Fusion gain achieved: {speedup:.2f}x speedup") + elif speedup >= 1.1: + print(f"⚠️ Modest fusion gain: {speedup:.2f}x speedup") + else: + print(f"❌ No significant fusion gain: {speedup:.2f}x") + + if args.compare_cuda and speedup_cuda > 0: + if parity >= 80: + print(f"✅ Parity with inductor+triton: {parity:.1f}%") + elif parity >= 60: + print(f"⚠️ Approaching parity: {parity:.1f}%") + else: + print(f"❌ Below parity: {parity:.1f}%") + + +if __name__ == "__main__": + main() diff --git a/torch_fl/__init__.py b/torch_fl/__init__.py index 472dc9b4..ad50f54b 100644 --- a/torch_fl/__init__.py +++ b/torch_fl/__init__.py @@ -803,6 +803,22 @@ def _accum_grad_hook(param, *, ddp_model=self): _patch_ddp_for_flagos() + +# Register torch.compile backend for flagos device (torch 2.0+) +def _register_compile_backend(): + """Register the 'flagos' backend with torch._dynamo if available.""" + try: + from torch_fl.compile.inductor_backend import register_backend + + register_backend() + except (ImportError, AttributeError): + # torch._dynamo not available (torch < 2.0) or inductor missing + pass + + +_register_compile_backend() + + __all__ = [ "flagos", "distributed", diff --git a/torch_fl/accelerator/cuda/_cuda_compat.py b/torch_fl/accelerator/cuda/_cuda_compat.py index e196d592..dce0c293 100644 --- a/torch_fl/accelerator/cuda/_cuda_compat.py +++ b/torch_fl/accelerator/cuda/_cuda_compat.py @@ -392,6 +392,56 @@ def _exchange_device(idx): torch.cuda.current_stream = lambda device=None: _StreamShim(_device_index(device)) torch.cuda.default_stream = lambda device=None: _StreamShim(_device_index(device)) + # torch.cuda.Event/Stream are dummy base classes in the CPU wheel and raise + # on construction. flagos ships working ones over the same physical GPU + # (its Event does real elapsed_time), so hand those out instead. inductor's + # kernel benchmarking constructs torch.cuda.Event(enable_timing=True). + if _flagos is not None: + if getattr(_flagos, "Event", None) is not None: + torch.cuda.Event = _flagos.Event + if getattr(_flagos, "Stream", None) is not None: + torch.cuda.Stream = _flagos.Stream + + # Memory stats. Every torch.cuda.memory_* query goes through + # torch._C._cuda_memoryStats, which the CPU wheel does not build, so they all + # raise AttributeError. flagos delegates its allocator to + # c10::cuda::CUDACachingAllocator, so its own stats describe the very same + # pool -- route the CUDA queries there. inductor's autotuner needs these to + # size its benchmark scratch budget (copy_args_to_cpu_if_needed). + if _flagos is not None: + torch.cuda.memory_allocated = lambda device=None: _flagos.memory_allocated( + _device_index(device) + ) + torch.cuda.memory_reserved = lambda device=None: _flagos.memory_reserved( + _device_index(device) + ) + + def _memory_stats(device=None): + """flagos stats under the nested keys torch.cuda callers expect. + + flagos reports flat names (``peak_allocated_bytes``); torch.cuda's + schema is ``allocated_bytes.all.peak``. Emit both so either style of + lookup resolves. + """ + stats = dict(_flagos.memory_stats(_device_index(device))) + for flat, nested in ( + ("allocated_bytes", "allocated_bytes.all.current"), + ("peak_allocated_bytes", "allocated_bytes.all.peak"), + ("reserved_bytes", "reserved_bytes.all.current"), + ("peak_reserved_bytes", "reserved_bytes.all.peak"), + ): + if flat in stats: + stats[nested] = stats[flat] + return stats + + torch.cuda.memory_stats = _memory_stats + torch.cuda.max_memory_allocated = lambda device=None: _memory_stats(device).get( + "peak_allocated_bytes", 0 + ) + torch.cuda.max_memory_reserved = lambda device=None: _memory_stats(device).get( + "peak_reserved_bytes", 0 + ) + # triton reads torch._C._cuda_getCurrentRawStream(idx) -> raw handle. try: torch._C._cuda_getCurrentRawStream = lambda idx=0: 0 diff --git a/torch_fl/compile/README.md b/torch_fl/compile/README.md new file mode 100644 index 00000000..a8bd00ca --- /dev/null +++ b/torch_fl/compile/README.md @@ -0,0 +1,97 @@ +# torch.compile Integration for FlagOS + +This directory contains the torch.compile backend integration for the flagos device. + +## Overview + +The `flagos` backend registers flagos with TorchInductor as a **first-class GPU +device**, then hands the traced graph to `compile_fx` unchanged. Inductor +generates Triton kernels that operate on flagos tensors directly -- there is no +conversion to cuda and no copy at the graph boundary. + +This works because flagos runs on the physical GPU that `torch.cuda` describes: +its allocator delegates to `c10::cuda::CUDACachingAllocator`, so a flagos +tensor's storage already *is* CUDA memory, and device indices line up +(`flagos.set_device(i)` moves the CUDA current device). + +## Usage + +```python +import torch +import torch_fl + +def my_model(x): + z = x + 1.0 + z = torch.nn.functional.relu(z) + z = z * 2.0 + return z + +x = torch.randn(4096, 4096, device='flagos:0') + +# Compile with flagos backend +compiled_model = torch.compile(my_model, backend='flagos') +result = compiled_model(x) + +# mode / options are forwarded to inductor as config patches +compiled_model = torch.compile(my_model, backend='flagos', mode='max-autotune') +``` + +## Implementation Notes + +### Why the graph stays on flagos + +An earlier version rewrote the graph and its example inputs to cuda before +calling `compile_fx`. That is not merely a copy cost -- it breaks backward. +`at::getAccelerator()` is PrivateUse1/flagos in this build, and +`torch::autograd::Node::stream()` only yields a stream when a node's input +device type equals the accelerator. A cuda-rewritten graph therefore produces +stream-less autograd nodes, and AOT autograd's backward trace inside +`compile_fx` trips `opt_ready_stream && opt_parent_stream` (engine.cpp:1085). + +Keeping the graph on flagos avoids that, and removes a copy-in/copy-out per call. + +### Registration surface (`device_interface.py`, `inductor_codegen.py`) + +| What | Why | +|---|---| +| `GPU_TYPES.append("flagos")` | `is_gpu()` is a membership test on this list; without it inductor picks the C++/CPU codegen path and never emits Triton. Must be in place -- callers captured the list object at import. | +| Prime `get_gpu_type()`'s cache | It asserts at most one GPU type is available, and the torch.cuda shim reports available alongside flagos. | +| `register_interface_for_device` | Inductor's `DeviceInterface`: device state from `torch.flagos`, hardware properties from `torch.cuda` (same GPU). | +| `DeviceProperties.create` wrap | Reports flagos as cuda at the Triton boundary. Triton's NVIDIA backend hard-checks `target.backend == "cuda"`, so a literal `"flagos"` finds zero compatible backends. Inductor already does this in the opposite direction for ROCm (`hints.py:149`). | +| `register_device_op_overrides` | Device guard / stream / synchronize snippets spliced into generated code. Inherits the CUDA ones; only Python-level device manipulation routes through `torch.flagos`. | +| `register_backend_for_device` | Scheduling + wrapper codegen -- the stock CUDA/Triton pipeline under the `"flagos"` key. | + +The four codegen classes are also published on `torch.flagos` so inductor's +official PrivateUse1 hook (`init_backend_registration`, `codegen/common.py:578`) +can register flagos on its own. + +### CPU-torch wheel accommodations + +This build pairs a CPU-only pip torch with an externally supplied +`libtorch_cuda.so`, so several `torch.cuda` Python bindings are missing. The +backend compensates: + +- `use_static_cuda_launcher = False` -- `torch._C._StaticCudaLauncher` is not built. +- `triton.cudagraphs = False` -- `torch.cuda.CUDAGraph` is a dummy base class + that raises on construction; `mode="max-autotune"` would otherwise enable it. +- `CudaInterface.get_raw_stream` is re-attached -- the binding exists, but the + import-time `torch.cuda._is_compiled()` probe left it at `None`. + +See `torch_fl/accelerator/cuda/_cuda_compat.py` for the memory-stats and +Event/Stream shims that inductor's autotuner needs. + +## Environment Variables + +- `FLAGOS_USE_FLAGTREE=1` - Reserved for FlagTree integration (Phase 2) +- `FLAGOS_COMPILE_FALLBACK_EAGER=1` - Fall back to eager mode on compile errors + +## Limitations + +1. Single device - multi-GPU compilation not yet exercised +2. FlagTree integration is a stub (Phase 2) + +## Future Work + +- [ ] FlagTree integration to replace OpenAI Triton +- [ ] Benchmark fusion gains against stock inductor+triton on cuda +- [ ] Multi-GPU compilation support diff --git a/torch_fl/compile/__init__.py b/torch_fl/compile/__init__.py new file mode 100644 index 00000000..3a7540ed --- /dev/null +++ b/torch_fl/compile/__init__.py @@ -0,0 +1,33 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +torch.compile backend for flagos device. + +Enables inductor-based kernel fusion for models running on the flagos device: + model = torch.compile(model, backend="flagos") + +The backend internally uses TorchInductor for graph optimization and kernel +fusion, with device context patched to target the flagos device. Generated +Triton kernels dispatch through the existing flagos routing infrastructure +(kFlagOsPython/kFlagOs/cuda boxing). + +Optional FlagTree integration (Phase 2): + FLAGOS_USE_FLAGTREE=1 model = torch.compile(model, backend="flagos") +This replaces OpenAI Triton with FlagTree for multi-backend kernel compilation. +""" + +from torch_fl.compile.inductor_backend import flagos_compile_backend + +__all__ = ["flagos_compile_backend"] diff --git a/torch_fl/compile/device_interface.py b/torch_fl/compile/device_interface.py new file mode 100644 index 00000000..b582452a --- /dev/null +++ b/torch_fl/compile/device_interface.py @@ -0,0 +1,257 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Device runtime interface that lets TorchInductor treat flagos as a GPU. + +Inductor asks a `DeviceInterface` for everything it needs at codegen and +autotune time: current device, streams, hardware properties, Triton capability. +Registering one for flagos is what makes inductor generate Triton kernels for +flagos tensors *directly*, instead of us rewriting the graph to cuda first. + +That rewrite is not a cosmetic difference. `at::getAccelerator()` returns +PrivateUse1 (flagos) in this build, and `torch::autograd::Node::stream()` only +returns a stream when a node's input device type equals the accelerator. So a +graph whose inputs were rewritten to cuda produces autograd nodes with no +stream, and the engine trips `opt_ready_stream && opt_parent_stream` (see +engine.cpp:1085) as soon as AOT autograd traces the backward. Keeping the graph +on flagos avoids that entirely -- and avoids a copy-in/copy-out per call. + +Hardware queries proxy to torch.cuda: flagos runs on the same physical GPU and +its allocator delegates to c10::cuda::CUDACachingAllocator, so device +properties, compute capability and raw streams are the CUDA ones. Device +indices line up too (flagos.set_device(i) moves the CUDA current device). +""" + +from typing import Any, Optional, Union + +import torch + +from torch._dynamo.device_interface import ( + DeviceInterface, + caching_worker_current_devices, + caching_worker_device_properties, +) + + +DEVICE_TYPE = "flagos" + + +def _device_index(device: Any) -> Optional[int]: + """Normalize str / torch.device / int into a plain device index.""" + if device is None: + return None + if isinstance(device, str): + device = torch.device(device) + if isinstance(device, torch.device): + return device.index + return int(device) + + +class FlagOSDeviceInterface(DeviceInterface): + """Inductor's device runtime interface, backed by flagos + torch.cuda. + + Mirrors torch._dynamo.device_interface.CudaInterface. Anything touching + *hardware* goes to torch.cuda (same GPU); anything touching *device state* + goes to torch.flagos so the two stay in sync. + """ + + device = torch.flagos.device # type: ignore[assignment] + + # Inductor captures these through dynamo; flagos ships its own shims that + # proxy the CUDA streams/events of the same physical GPU. + Event = torch.flagos.Event # type: ignore[assignment] + Stream = torch.flagos.Stream # type: ignore[assignment] + + class Worker: + """Property queries that must work in forked compile workers. + + Workers cannot touch the GPU, so properties are recorded in the parent + process and read from the cache here (same contract as CudaInterface). + """ + + @staticmethod + def set_device(device: int) -> None: + caching_worker_current_devices[DEVICE_TYPE] = device + + @staticmethod + def current_device() -> int: + if DEVICE_TYPE in caching_worker_current_devices: + return caching_worker_current_devices[DEVICE_TYPE] + return torch.flagos.current_device() + + @staticmethod + def get_device_properties(device: Any = None) -> Any: + idx = _device_index(device) + if idx is None: + idx = FlagOSDeviceInterface.Worker.current_device() + + if DEVICE_TYPE not in caching_worker_device_properties: + caching_worker_device_properties[DEVICE_TYPE] = [ + torch.cuda.get_device_properties(i) + for i in range(torch.cuda.device_count()) + ] + + return caching_worker_device_properties[DEVICE_TYPE][idx] + + # --- device state: flagos ------------------------------------------------ + current_device = staticmethod(torch.flagos.current_device) + set_device = staticmethod(torch.flagos.set_device) + device_count = staticmethod(torch.flagos.device_count) + synchronize = staticmethod(torch.flagos.synchronize) + + # --- streams: flagos shims proxy the CUDA stream of the same GPU --------- + stream = staticmethod(torch.flagos.stream) # type: ignore[assignment] + current_stream = staticmethod(torch.flagos.current_stream) # type: ignore[assignment] + set_stream = staticmethod(torch.cuda.set_stream) # type: ignore[assignment] + _set_stream_by_id = staticmethod(torch.cuda._set_stream_by_id) # type: ignore[assignment] + # Generated Triton launch code passes this raw stream handle to the kernel. + get_raw_stream = staticmethod(torch._C._cuda_getCurrentRawStream) # type: ignore[assignment] + + # --- hardware: same physical GPU as cuda -------------------------------- + get_device_properties = staticmethod(torch.cuda.get_device_properties) # type: ignore[assignment] + memory_allocated = staticmethod(torch.flagos.memory_allocated) + exchange_device = staticmethod(torch.cuda._exchange_device) # type: ignore[assignment] + maybe_exchange_device = staticmethod(torch.cuda._maybe_exchange_device) # type: ignore[assignment] + + @staticmethod + def is_available() -> bool: + return torch.flagos.device_count() > 0 + + @staticmethod + def is_bf16_supported(including_emulation: bool = True) -> bool: + return torch.cuda.is_bf16_supported() + + @staticmethod + def get_compute_capability(device: Any = None) -> Union[int, str]: + major, minor = torch.cuda.get_device_capability(_device_index(device)) + return major * 10 + minor + + @staticmethod + def is_triton_capable(device: Any = None) -> bool: + return torch.cuda.get_device_properties(_device_index(device)).major >= 7 + + @staticmethod + def raise_if_triton_unavailable(device: Any = None) -> None: + from torch._inductor.exc import GPUTooOldForTriton + import inspect + + if not FlagOSDeviceInterface.is_triton_capable(device): + raise GPUTooOldForTriton( + torch.cuda.get_device_properties(_device_index(device)), + inspect.currentframe(), + ) + + import triton.backends + + if "nvidia" not in triton.backends.backends: + raise RuntimeError("triton not built with the 'nvidia' backend") + + +def _register_gpu_type() -> None: + """Teach inductor that flagos is a GPU, not a CPU-like device. + + Two separate things read GPU_TYPES: + + * `is_gpu()` -- a membership test. Without flagos in the list inductor picks + the C++/CPU codegen path and never emits Triton. The append must be in + place; callers captured this exact list object at import time. + + * `get_gpu_type()` -- picks *the* single GPU type, asserting at most one + entry of GPU_TYPES is available. torch_fl's torch.cuda shim reports + available alongside flagos, so that assert would fire (post_grad's + ConstructorMoverPass hits it). It is `functools.cache`d, so we prime the + cache with flagos while the list is temporarily narrowed, and every later + caller gets the memoized answer. + """ + from torch._inductor import utils as inductor_utils + + if DEVICE_TYPE in inductor_utils.GPU_TYPES: + return + + saved = list(inductor_utils.GPU_TYPES) + try: + inductor_utils.GPU_TYPES[:] = [DEVICE_TYPE] + inductor_utils.get_gpu_type() + finally: + inductor_utils.GPU_TYPES[:] = saved + inductor_utils.GPU_TYPES.append(DEVICE_TYPE) + + +def _patch_device_properties() -> None: + """Report flagos devices to the Triton layer as cuda. + + `DeviceProperties.type` is what inductor forwards to Triton as + `GPUTarget.backend` (hints.py -> triton_heuristics.py:718 -> triton's + make_backend). Triton's NVIDIA backend hard-checks `target.backend == "cuda"`, + so a literal "flagos" finds zero compatible backends. The same read also + drives cubin vs hsaco selection and the launcher's interface lookup -- all of + which we want to be the CUDA ones, because that is the hardware. + + Rewriting the device type at this boundary is what inductor already does for + ROCm in the opposite direction (`hints.py:149`, cuda -> hip). We wrap + `create` rather than editing the read sites so the functools cache, and every + other field, keep working unchanged. + """ + from torch._inductor.runtime.hints import DeviceProperties + + if getattr(DeviceProperties.create, "_flagos_patched", False): + return + + original = DeviceProperties.create + + def create(device: Any) -> Any: + if device is not None and getattr(device, "type", None) == DEVICE_TYPE: + device = torch.device("cuda", device.index or 0) + return original(device) + + create._flagos_patched = True # type: ignore[attr-defined] + DeviceProperties.create = create # type: ignore[method-assign, assignment] + + +def _repair_cuda_interface_raw_stream() -> None: + """Give the stock CudaInterface its raw-stream getter back. + + Because DeviceProperties reports flagos as cuda (see _patch_device_properties), + inductor's *runtime* launcher resolves the stock `CudaInterface`, and its + autotuner calls `get_raw_stream(current_device())`. That attribute is bound at + import time from `torch._C._cuda_getCurrentRawStream`, but only when + `torch.cuda._is_compiled()` -- which is False for the CPU torch wheel, so it + lands on None and the autotuner raises "'NoneType' object is not callable". + + The binding itself is present: torch_fl's torch.cuda shim installs it + (accelerator/cuda/_cuda_compat.py). So this just re-attaches what the import + time probe missed. + """ + import torch._dynamo.device_interface as di + + getter = getattr(torch._C, "_cuda_getCurrentRawStream", None) + if getter is None or di.CudaInterface.get_raw_stream is not None: + return + + di.get_cuda_stream = getter + di.CudaInterface.get_raw_stream = staticmethod(getter) # type: ignore[assignment] + + +def register_flagos_device_interface() -> None: + """Register the flagos device interface + GPU type with inductor. + + Idempotent; called by the compile backend before each compile_fx. + """ + from torch._dynamo.device_interface import register_interface_for_device + + _register_gpu_type() + _patch_device_properties() + _repair_cuda_interface_raw_stream() + register_interface_for_device(DEVICE_TYPE, FlagOSDeviceInterface) diff --git a/torch_fl/compile/flagtree_shim.py b/torch_fl/compile/flagtree_shim.py new file mode 100644 index 00000000..9a320aff --- /dev/null +++ b/torch_fl/compile/flagtree_shim.py @@ -0,0 +1,86 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +FlagTree integration for torch.compile (Phase 2). + +Patches inductor's Triton imports to use FlagTree instead of OpenAI Triton, +enabling multi-backend kernel compilation (NVIDIA/Ascend/Cambricon/MetaX). + +Activated via FLAGOS_USE_FLAGTREE=1 environment variable. +""" + +import sys + + +_original_triton = None +_patched = False + + +def patch_inductor_triton(): + """ + Replace inductor's triton imports with flagtree. + + FlagTree is API-compatible with OpenAI Triton (it's a fork), so this is + a drop-in replacement. Inductor generates Triton kernel code; we just + swap which compiler JITs it. + + This must be called before inductor imports triton (i.e., before the + first torch.compile call that uses inductor). + """ + global _patched, _original_triton + + if _patched: + return + + try: + import flagtree as triton + except ImportError: + raise ImportError( + "FLAGOS_USE_FLAGTREE=1 but 'flagtree' package not installed. " + "Install with: pip install flagtree" + ) + + # Save original triton if already imported + if "triton" in sys.modules: + _original_triton = sys.modules["triton"] + + # Replace triton in sys.modules with flagtree + sys.modules["triton"] = triton + + # Also replace triton.language (inductor imports this) + if hasattr(triton, "language"): + sys.modules["triton.language"] = triton.language + + _patched = True + + +def unpatch_inductor_triton(): + """ + Restore original OpenAI Triton (for testing/cleanup). + """ + global _patched, _original_triton + + if not _patched: + return + + if _original_triton is not None: + sys.modules["triton"] = _original_triton + if hasattr(_original_triton, "language"): + sys.modules["triton.language"] = _original_triton.language + else: + sys.modules.pop("triton", None) + sys.modules.pop("triton.language", None) + + _patched = False diff --git a/torch_fl/compile/inductor_backend.py b/torch_fl/compile/inductor_backend.py new file mode 100644 index 00000000..778aea54 --- /dev/null +++ b/torch_fl/compile/inductor_backend.py @@ -0,0 +1,200 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Inductor-based compile backend for the flagos device. + +flagos is registered with inductor as a first-class GPU device (see +device_interface.py and inductor_codegen.py), so the traced graph is handed to +`compile_fx` *as is* -- still on flagos. Inductor generates Triton kernels for +it directly. + +Why not rewrite the graph to cuda (as an earlier version did): `at::getAccelerator()` +is PrivateUse1/flagos here, and `torch::autograd::Node::stream()` only yields a +stream when a node's input device type matches the accelerator. A cuda-rewritten +graph therefore produces stream-less autograd nodes, and AOT autograd's backward +trace inside compile_fx trips `opt_ready_stream && opt_parent_stream` +(engine.cpp:1085). Staying on flagos also removes a copy-in/copy-out per call. +""" + +import os +from typing import Any, Callable, Dict, List, Optional + +import torch +import torch.fx +import torch.cuda + + +def _patch_cuda_rng_for_cpu_torch(): + """ + Workaround for CPU torch + external libtorch_cuda.so setup. + + dynamo tries to capture torch.cuda.get_rng_state() during tracing, but + CPU torch doesn't have torch._C._cuda_getDevice() binding. We patch + torch.cuda to provide stub implementations that prevent the crash. + + Only applied when torch._C lacks CUDA bindings (CPU torch build). + """ + import torch as torch_module + + if hasattr(torch_module._C, "_cuda_getDevice"): + return # Native CUDA torch, no patch needed + + # CPU torch detected - patch CUDA RNG functions + import torch.cuda as cuda_module + + # Stub implementations that won't be called (dynamo just needs them callable) + def _stub_get_rng_state(device=None): + # Return empty tensor as placeholder (dynamo won't execute this) + return torch_module.tensor([], dtype=torch_module.uint8) + + def _stub_set_rng_state(new_state, device=None): + pass # No-op + + cuda_module.get_rng_state = _stub_get_rng_state + cuda_module.set_rng_state = _stub_set_rng_state + + +# Apply CPU torch workaround at module load time +_patch_cuda_rng_for_cpu_torch() + + +def _resolve_config_patches( + mode: Optional[str], + options: Optional[Dict[str, Any]], + dynamic: Optional[bool], +) -> Dict[str, Any]: + """Turn torch.compile's mode/options into an inductor config patch dict. + + Same expansion `_TorchCompileInductorWrapper` does, plus the flagos-specific + overrides this build needs. Passing these to compile_fx as `config_patches` + scopes them to this compile, instead of mutating inductor's global config. + """ + patches: Dict[str, Any] = {} + + if mode and mode != "default": + from torch._inductor import list_mode_options + + patches.update(list_mode_options(mode, dynamic)) + if options: + patches.update({k.replace("-", "_"): v for k, v in options.items()}) + + # CUDA graphs need torch.cuda.CUDAGraph, a dummy base class in the CPU torch + # wheel ("Tried to instantiate dummy base class CUDAGraph"). mode= + # "max-autotune" turns them on, so force them back off. + patches["triton.cudagraphs"] = False + + # The static launcher needs torch._C._StaticCudaLauncher, which the CPU + # torch wheel does not build (we supply libtorch_cuda.so externally). Fall + # back to the regular Triton launch path. + if not hasattr(torch._C, "_StaticCudaLauncher"): + patches["use_static_cuda_launcher"] = False + + return patches + + +def flagos_compile_backend( + gm: torch.fx.GraphModule, + example_inputs: List[torch.Tensor], + mode: Optional[str] = None, + options: Optional[Dict[str, Any]] = None, + dynamic: Optional[bool] = None, +) -> Callable: + """ + torch.compile backend for flagos device. + + Registers flagos as an inductor GPU device, then delegates the graph to + inductor unchanged. Generated Triton kernels run on flagos tensors directly. + + `mode` / `options` arrive as kwargs from dynamo's `_TorchCompileWrapper` + (torch/__init__.py) whenever torch.compile is called with them on a named + backend; they are expanded into inductor config patches. + + Usage: + model = torch.compile(model, backend="flagos") + # or + model = torch.compile(model, backend="flagos", mode="max-autotune") + + Environment: + FLAGOS_USE_FLAGTREE=1 : Use FlagTree instead of OpenAI Triton (Phase 2) + FLAGOS_COMPILE_FALLBACK_EAGER=1 : Fall back to eager on compile errors + """ + # Import inductor lazily (not all torch builds have it) + try: + from torch._inductor.compile_fx import compile_fx + except ImportError as e: + if os.environ.get("FLAGOS_COMPILE_FALLBACK_EAGER", "0") == "1": + return gm.forward + raise RuntimeError( + "torch._inductor not available. Install torch with inductor support " + "or set FLAGOS_COMPILE_FALLBACK_EAGER=1 to fall back to eager." + ) from e + + config_patches = _resolve_config_patches(mode, options, dynamic) + + # Phase 2 hook: FlagTree integration + use_flagtree = os.environ.get("FLAGOS_USE_FLAGTREE", "0") == "1" + if use_flagtree: + try: + from torch_fl.compile.flagtree_shim import patch_inductor_triton + + patch_inductor_triton() + except ImportError: + import warnings + + warnings.warn( + "FLAGOS_USE_FLAGTREE=1 but flagtree_shim not available. " + "Falling back to OpenAI Triton." + ) + + # Make inductor treat flagos as a GPU device. Order matters: is_gpu() must + # answer True and the device interface must be resolvable before the + # codegen backend registration reads them. + from torch_fl.compile.device_interface import register_flagos_device_interface + from torch_fl.compile.inductor_codegen import ( + publish_codegen_on_device_module, + register_flagos_codegen, + ) + + register_flagos_device_interface() + publish_codegen_on_device_module() + register_flagos_codegen() + + # Hand the graph to inductor untouched -- it is on flagos and stays there. + try: + return compile_fx(gm, example_inputs, config_patches=config_patches) + except Exception as e: + if os.environ.get("FLAGOS_COMPILE_FALLBACK_EAGER", "0") == "1": + import warnings + + warnings.warn(f"Inductor compilation failed: {e}. Falling back to eager.") + return gm.forward + raise + + +def register_backend(): + """ + Register the flagos backend with torch._dynamo. + + Called automatically on import torch_fl if torch 2.0+ detected. + """ + try: + import torch._dynamo + + torch._dynamo.register_backend( + name="flagos", compiler_fn=flagos_compile_backend + ) + except (ImportError, AttributeError): + # torch._dynamo not available (torch < 2.0) + pass diff --git a/torch_fl/compile/inductor_codegen.py b/torch_fl/compile/inductor_codegen.py new file mode 100644 index 00000000..e180407e --- /dev/null +++ b/torch_fl/compile/inductor_codegen.py @@ -0,0 +1,151 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Inductor codegen backend registration for the flagos device. + +Two things get registered here: + +1. `DeviceOpOverrides` -- the snippets inductor splices into generated code for + device guards, stream lookup and synchronization. flagos reuses the CUDA + ones almost verbatim: a flagos tensor's storage *is* CUDA memory (the flagos + allocator delegates to c10::cuda::CUDACachingAllocator), and torch_fl ships a + torch.cuda shim over the same physical GPU. Only the Python-level device + guard and set_device switch to `torch.flagos`, so generated code moves the + flagos current device rather than desyncing the two. + +2. The scheduling + wrapper codegen classes. These are the stock CUDA ones -- + flagos wants exactly the Triton/CUDA pipeline -- registered under the + "flagos" device key. + +Registration goes through inductor's official PrivateUse1 path where possible: +`init_backend_registration()` looks up `Scheduling`, `PythonWrapperCodegen`, +`CppWrapperCodegen` and `WrapperFxCodegen` on the device module named by +`torch._C._get_privateuse1_backend_name()` (i.e. `torch.flagos`). We publish +those four names *and* call `register_backend_for_device` directly, since +`init_backend_registration` may already have run before torch_fl was imported. +""" + +from torch._inductor.codegen.common import ( + get_scheduling_for_device, + register_backend_for_device, + register_device_op_overrides, +) +from torch._inductor.codegen.cuda.device_op_overrides import CUDADeviceOpOverrides + + +DEVICE_TYPE = "flagos" + + +class FlagOSDeviceOpOverrides(CUDADeviceOpOverrides): + """Code snippets inductor emits into generated flagos kernels. + + Inherits the CUDA implementation for everything C++/driver related (kernel + headers, stream types, TMA helpers, AOTI guards) because those operate on + the underlying CUDA runtime, which is exactly what flagos runs on. Only the + Python-level device manipulation is overridden, to go through torch.flagos + so generated code doesn't desync the flagos and CUDA current device. + """ + + def set_device(self, device_idx: int) -> str: + return f"torch.flagos.set_device({device_idx})" + + def device_guard(self, device_idx: int) -> str: + # torch.flagos has no _DeviceGuard; torch.flagos.device(idx) is the + # equivalent context manager and keeps flagos/cuda current device in + # sync (flagos.set_device moves the CUDA device too). + return f"torch.flagos.device({device_idx})" + + def synchronize(self) -> str: + return "torch.flagos.synchronize()" + + def import_get_raw_stream_as(self, name: str) -> str: + # flagos streams are the CUDA streams of the same GPU, and this handle + # is what gets handed to the Triton kernel launcher. + return f"from torch._C import _cuda_getCurrentRawStream as {name}" + + +def register_flagos_codegen() -> None: + """Register flagos scheduling/wrapper codegen + device op overrides. + + Idempotent: safe to call before every compile_fx. + """ + from torch._inductor.codegen.common import ( + device_op_overrides_dict, + init_backend_registration, + ) + + # Populate the built-in registrations first, so our checks below see the + # real state and so the CUDA overrides module is imported. + init_backend_registration() + + if DEVICE_TYPE not in device_op_overrides_dict: + register_device_op_overrides(DEVICE_TYPE, FlagOSDeviceOpOverrides()) + + if get_scheduling_for_device(DEVICE_TYPE) is None: + from torch._inductor import config + from torch._inductor.codegen.cpp_wrapper_gpu import CppWrapperGpu + from torch._inductor.codegen.cuda_combined_scheduling import ( + CUDACombinedScheduling, + ) + from torch._inductor.codegen.halide import HalideScheduling + from torch._inductor.codegen.simd import SIMDScheduling + from torch._inductor.codegen.wrapper import PythonWrapperCodegen + from torch._inductor.codegen.wrapper_fxir import WrapperFxCodegen + + # Same table CUDA uses; flagos wants the Triton pipeline. + backends = { + "triton": CUDACombinedScheduling, + "halide": HalideScheduling, + } + register_backend_for_device( + DEVICE_TYPE, + lambda scheduling: backends.get(config.cuda_backend, SIMDScheduling)( + scheduling + ), + PythonWrapperCodegen, + CppWrapperGpu, + WrapperFxCodegen, + ) + + +def publish_codegen_on_device_module() -> None: + """Expose the four codegen classes on `torch.flagos`. + + This is inductor's sanctioned PrivateUse1 hook: `init_backend_registration` + reads `Scheduling` / `PythonWrapperCodegen` / `CppWrapperCodegen` / + `WrapperFxCodegen` off the device module. Setting them means a fresh + inductor process registers flagos on its own, without our backend having to + run first. + """ + import torch + + from torch._inductor.codegen.cpp_wrapper_gpu import CppWrapperGpu + from torch._inductor.codegen.cuda_combined_scheduling import ( + CUDACombinedScheduling, + ) + from torch._inductor.codegen.wrapper import PythonWrapperCodegen + from torch._inductor.codegen.wrapper_fxir import WrapperFxCodegen + + mod = getattr(torch, DEVICE_TYPE, None) + if mod is None: + return + for name, cls in ( + ("Scheduling", CUDACombinedScheduling), + ("PythonWrapperCodegen", PythonWrapperCodegen), + ("CppWrapperCodegen", CppWrapperGpu), + ("WrapperFxCodegen", WrapperFxCodegen), + ): + if not hasattr(mod, name): + setattr(mod, name, cls) diff --git a/torch_fl/flagos/meta.py b/torch_fl/flagos/meta.py index 5856c239..3e54f6f1 100644 --- a/torch_fl/flagos/meta.py +++ b/torch_fl/flagos/meta.py @@ -15,9 +15,34 @@ # Meta functions for FlagOS device # These are used by torch.compile and other meta-dispatch mechanisms +""" +Meta tensor implementations for shape inference during torch.compile tracing. -# You can add meta implementations here if needed -# For example: -# @impl("aten::some_op", "Meta") -# def some_op_meta(self): -# return torch.empty_like(self) +When torch.compile traces a model with flagos tensors, it needs to infer output +shapes without executing kernels. We register meta implementations that compute +output shapes/dtypes for ops that don't have default meta kernels. + +Most ops inherit meta kernels from their CPU/CUDA implementations. We only need +to register meta kernels for: +1. Custom ops specific to flagos +2. Ops where the default meta kernel is incorrect for our backend +3. Ops that torch.compile explicitly requires but are missing + +Start with a minimal set and expand as needed based on compile errors. +""" + + +# Meta kernels are registered via torch.library.impl +# Format: @torch.library.impl("aten::op_name", "Meta") + +# Example: if we had a custom flagos-specific op +# @torch.library.impl("flagos::custom_op", "Meta") +# def custom_op_meta(input: Tensor, alpha: float) -> Tensor: +# return torch.empty_like(input) + + +# Most standard ops (mm, add, conv, etc.) already have meta kernels from +# torch's default registrations. We inherit those automatically. + +# Placeholder for future custom meta implementations as needed +# (torch.compile will error with specific op names if meta kernels are missing)