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)