Passing in make_fx input into helion

[mitkotak]

Hey was wondering whether I can pass in a make_fx()'d graph into helion to leverage its autotuner and spit out triton code.

import torch

import helion
import helion.language as hl
from torch.fx.experimental.proxy_tensor import make_fx

class F(torch.nn.Module):
    def forward(self, x: torch.Tensor):
        return x + 1


# Need to pass device args otherwise it complains
@helion.kernel()
def F_kernel(func_m, device):
    return func_m

x = torch.randn(1000).to(device='cuda')
m = make_fx(F(), tracing_mode="symbolic", _allow_non_fake_inputs=True, _error_on_data_dependent_ops=True)(x)

torch.testing.assert_allclose(m(x), F()(x))

F_kernel(lambda x: m(x), x)

Traceback (most recent call last):
  File "/home/mkotak/atomic_architects/projects/helion_playground/fx_helion.py", line 24, in <module>
    F_kernel(lambda x: F()(x), x)
  File "/home/mkotak/atomic_architects/projects/helion_playground/.venv/lib/python3.11/site-packages/helion/runtime/kernel.py", line 272, in __call__
    return self.bind(args)(*args)
           ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mkotak/atomic_architects/projects/helion_playground/.venv/lib/python3.11/site-packages/helion/runtime/kernel.py", line 581, in __call__
    self.autotune(args)
  File "/home/mkotak/atomic_architects/projects/helion_playground/.venv/lib/python3.11/site-packages/helion/runtime/kernel.py", line 473, in autotune
    config = self.settings.autotuner_fn(self, args, **kwargs).autotune()
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mkotak/atomic_architects/projects/helion_playground/.venv/lib/python3.11/site-packages/helion/runtime/settings.py", line 68, in default_autotuner_fn
    return LocalAutotuneCache(DifferentialEvolutionSearch(bound_kernel, args, **kwargs))  # pyright: ignore[reportArgumentType]
                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mkotak/atomic_architects/projects/helion_playground/.venv/lib/python3.11/site-packages/helion/autotuner/differential_evolution.py", line 34, in __init__
    super().__init__(kernel, args)
  File "/home/mkotak/atomic_architects/projects/helion_playground/.venv/lib/python3.11/site-packages/helion/autotuner/base_search.py", line 319, in __init__
    self.config_gen: ConfigGeneration = ConfigGeneration(self.config_spec)
                                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mkotak/atomic_architects/projects/helion_playground/.venv/lib/python3.11/site-packages/helion/autotuner/config_generation.py", line 55, in __init__
    self.min_block_size: int = max(
                               ^^^^
ValueError: max() arg is an empty sequence

[jansel]

That example is returning func_m not calling it, which doesn't seem right.

I'd expect this to work (haven't tested it though) if you have an hl.tile loop and wrap the FX graph in a function. Something like:

@helion.kernel()
def F_kernel(func_m, x):
    out = torch.empty_like(x)
    for tile in hl.tile(out.size())
       out[tile] = func_m(x[tile])
    return m

m = make_fx(...)
# need to wrap it in a function since helion doesn't understand GraphModule args
m_wrapped = lambda x: m(x)
F_kernel(m_wrapped, x)

Because function calls are traced with make_fx.

It would be straightforward to support taking m directly (deleting m_wrapped) by adding a case in arg processing. I can do that in the next few days.

Would something like that meet your needs?

[mitkotak]

Yup this example fits my workload.

@helion.kernel()
def F_kernel(func_m, x):
    out = torch.empty_like(x)
    for tile in hl.tile(out.size()):
       out[tile] = func_m(x[tile])
    return out

x = torch.randn(1000).to(device='cuda')
m = make_fx(F(), tracing_mode="symbolic", _allow_non_fake_inputs=True, _error_on_data_dependent_ops=True)(x)

m_wrapped = lambda x: m(x)
F_kernel(m_wrapped, x)

Traceback (most recent call last):
  File "/home/mkotak/atomic_architects/projects/helion_playground/.venv/lib/python3.11/site-packages/helion/_compiler/type_propagation.py", line 765, in propagate_call
    _CheckForIndexCalls.retry_call(
  File "/home/mkotak/atomic_architects/projects/helion_playground/.venv/lib/python3.11/site-packages/helion/language/tile_proxy.py", line 122, in retry_call
    return fn(*proxy_args, **proxy_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mkotak/atomic_architects/projects/helion_playground/.venv/lib/python3.11/site-packages/helion/_compiler/lift_closures.py", line 56, in wrapper
    result = new_func(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mkotak/atomic_architects/projects/helion_playground/fx_helion.py", line 23, in <lambda>
    m_wrapped = lambda x: m(x)
                          ^
  File "/home/mkotak/atomic_architects/projects/helion_playground/.venv/lib/python3.11/site-packages/helion/_compiler/lift_closures.py", line 22, in __getitem__
    return HostFunction.current().register_fake(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mkotak/atomic_architects/projects/helion_playground/.venv/lib/python3.11/site-packages/helion/_compiler/host_function.py", line 166, in register_fake
    value = CompileEnvironment.current().to_fake(obj, origin)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mkotak/atomic_architects/projects/helion_playground/.venv/lib/python3.11/site-packages/helion/_compiler/compile_environment.py", line 260, in to_fake
    raise TypeError(f"unsupported argument type {type(obj)} ({origin})")
TypeError: unsupported argument type <class 'torch.fx.graph_module.GraphModule.__new__.<locals>.GraphModuleImpl'> (AttributeOrigin(value=GlobalOrigin(name='_source_module'), key='m'))

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/mkotak/atomic_architects/projects/helion_playground/fx_helion.py", line 24, in <module>
    F_kernel(m_wrapped, x)
  File "/home/mkotak/atomic_architects/projects/helion_playground/.venv/lib/python3.11/site-packages/helion/runtime/kernel.py", line 272, in __call__
    return self.bind(args)(*args)
           ^^^^^^^^^^^^^^^
  File "/home/mkotak/atomic_architects/projects/helion_playground/.venv/lib/python3.11/site-packages/helion/runtime/kernel.py", line 158, in bind
    bound_kernel = BoundKernel(self, args)
                   ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mkotak/atomic_architects/projects/helion_playground/.venv/lib/python3.11/site-packages/helion/runtime/kernel.py", line 338, in __init__
    self.host_function: HostFunction = HostFunction(
                                       ^^^^^^^^^^^^^
  File "/home/mkotak/atomic_architects/projects/helion_playground/.venv/lib/python3.11/site-packages/helion/_compiler/host_function.py", line 108, in __init__
    propagate_types(self, fake_args)
  File "/home/mkotak/atomic_architects/projects/helion_playground/.venv/lib/python3.11/site-packages/helion/_compiler/type_propagation.py", line 2163, in propagate_types
    prop.visit(stmt)
  File "/home/mkotak/atomic_architects/projects/helion_playground/.venv/lib/python3.11/site-packages/helion/_compiler/type_propagation.py", line 1492, in visit
    type_info = visitor(node)
                ^^^^^^^^^^^^^
  File "/home/mkotak/atomic_architects/projects/helion_playground/.venv/lib/python3.11/site-packages/helion/_compiler/type_propagation.py", line 1998, in visit_For
    body = self._loop_body(node.body)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mkotak/atomic_architects/projects/helion_playground/.venv/lib/python3.11/site-packages/helion/_compiler/type_propagation.py", line 1962, in _loop_body
    self.visit(stmt)
  File "/home/mkotak/atomic_architects/projects/helion_playground/.venv/lib/python3.11/site-packages/helion/_compiler/type_propagation.py", line 1492, in visit
    type_info = visitor(node)
                ^^^^^^^^^^^^^
  File "/home/mkotak/atomic_architects/projects/helion_playground/.venv/lib/python3.11/site-packages/helion/_compiler/type_propagation.py", line 1886, in visit_Assign
    type_info = self.visit(node.value)
                ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mkotak/atomic_architects/projects/helion_playground/.venv/lib/python3.11/site-packages/helion/_compiler/type_propagation.py", line 1492, in visit
    type_info = visitor(node)
                ^^^^^^^^^^^^^
  File "/home/mkotak/atomic_architects/projects/helion_playground/.venv/lib/python3.11/site-packages/helion/_compiler/type_propagation.py", line 1828, in visit_Call
    return func.propagate_call(tuple(args), kwargs, self.origin())  # pyright: ignore[reportReturnType]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mkotak/atomic_architects/projects/helion_playground/.venv/lib/python3.11/site-packages/helion/_compiler/type_propagation.py", line 799, in propagate_call
    raise exc.TorchOpTracingError(e) from e
helion.exc.TorchOpTracingError: TypeError: unsupported argument type <class 'torch.fx.graph_module.GraphModule.__new__.<locals>.GraphModuleImpl'> (AttributeOrigin(value=GlobalOrigin(name='_source_module'), key='m'))
While processing:
  File "/home/mkotak/atomic_architects/projects/helion_playground/fx_helion.py", line 16, in F_kernel
    out[tile] = func_m(x[tile])
                ^^^^^^^^^^^^^^^

[jansel]

I'll fix this next week.

[jansel]

#488 adds support for this but requires the GraphModule to be stateless (no parameters). If you need stateful graph modules you should lift the params into inputs the way AOTAutograd does.