[Issue] `torch.compile` fails with FlyDSL kernels

[amd-cgilli]

Summary

torch.compile(fullgraph=True) doesn't work with FlyDSL compiled kernels.

Two distinct failure modes are reproduced:

1. Compiling the kernel inside torch.compile
2. Calling a pre-compiled kernel inside torch.compile

System configuration

• GPU: MI355X
• OS: Ubuntu 24.04
• PyTorch: 2.10.0+rocm7.2.2.git40d237bf

Repro

Code available here.

Both scripts work correctly in eager mode and only fail at the torch.compile step.

Case 1: compilation inside torch.compile

FLYDSL_RUNTIME_ENABLE_CACHE=0 python tests/kernels/test_gemm_torch_compile.py 

Case 2: invocation of pre-compiled kernel inside torch.compile

FLYDSL_RUNTIME_ENABLE_CACHE=0 python tests/kernels/test_gemm_torch_compile_wrapper.py

[coderfeli]

I reproduced both failure modes locally.

A few findings:

• Compiling a FlyDSL kernel inside torch.compile(fullgraph=True) is not currently traceable. Dynamo first runs into Python-side arch detection (subprocess / threading.Lock), and after bypassing that it continues into FlyDSL/MLIR Python reflection and binding code such as inspect.getsource() / ir.InsertionPoint.
• Calling a flyc.compile(...) result directly inside torch.compile(fullgraph=True) also fails because Dynamo traces into the Python/ctypes argument packing path (memoryview / struct.Struct.pack_into).
• torch.compiler.disable is not a workaround for fullgraph=True; it produces a graph-break error.

The path that does work is to keep both FlyDSL compilation steps outside the Dynamo graph, then expose the precompiled executable as a PyTorch custom op (torch.library.custom_op) with a fake/meta implementation. With that shape, Dynamo sees an opaque op instead of tracing through FlyDSL dispatch. I verified the reported GEMM wrapper pattern with torch.compile(fullgraph=True, backend="eager") and backend="aot_eager" and the output passed the existing accuracy check.

So the likely fix is to add a small FlyDSL helper around CompiledFunction for registering a precompiled executable as a PyTorch custom op, rather than trying to make flyc.compile or CompiledFunction.__call__ itself traceable.

[amd-cgilli]

@amd-asalykov

[amd-asalykov]

@coderfeli thank you for your feedback! In fact, I did some experiments with precompiled exe + torch.library.custom_op and I can confirm it works, but I observed perf. degradation when integrating into a model and running e2e benchmarks. Let me prepare a minimal repro code