[BackendBench] Add Helion DSL kernel generation support - #200
Conversation
stack-info: PR: #200, branch: karthickai/stack/1
5e5f618 to
9ad8ddc
Compare
meta-cla
Bot
added
the
CLA Signed
|
Thank you for adding Helion support! I have a question about how Helion uses PyTorch operators to generate Triton kernels. Specifically, I'm concerned about the potential for infinite loops in BackendBench. If the Helion kernel implementation (e.g., add__Tensor_kernel_impl) calls torch.add, and BackendBench registers the Helion kernel to replace the PyTorch eager op (torch.ops.aten.add.Tensor), wouldn't that result in a recursive loop? |
It will not cause infinite recursion because Helion uses for example import torch
import helion
import helion.language as hl
os.environ['HELION_AUTOTUNE_EFFORT'] = 'none'
# llm generated kernel I modified to torch.add instead of +
@helion.kernel()
def my_add_kernel(x, y):
x, y = torch.broadcast_tensors(x, y)
out = torch.empty(x.shape, dtype=torch.promote_types(x.dtype, y.dtype), device=x.device)
out_flat = out.view(-1)
x_flat = x.contiguous().view(-1)
y_flat = y.contiguous().view(-1)
for tile in hl.tile(out_flat.size()):
out_flat[tile] = torch.add(x_flat[tile], y_flat[tile])
return out
x = torch.randn(10, device='cuda')
y = torch.randn(10, device='cuda')
lib = torch.library.Library("aten", "IMPL")
lib.impl("add.Tensor", my_add_kernel, "CUDA")
out = torch.add(x, y)output tensor([-1.1052, -0.9066, -0.1886, -2.5633, 2.8374, -0.8006, 0.7526, 2.2208,
0.2794, 3.7469], device='cuda:0') |
|
Thank you! Just a small request. Could you please change the name of the wrapper function from add_kernel_impl to add__Tensor_kernel_impl? This would align better with our current naming standard. |
I believe it is generating correctly import torch
import helion
import helion.language as hl
@helion.kernel()
def add___Tensor_helion_kernel(input_tensor: torch.Tensor, other: torch.Tensor) -> torch.Tensor:
input_tensor, other = torch.broadcast_tensors(input_tensor, other)
out = torch.empty(
input_tensor.shape,
dtype=torch.promote_types(input_tensor.dtype, other.dtype),
device=input_tensor.device
)
input_flat = input_tensor.contiguous().view(-1)
other_flat = other.contiguous().view(-1)
out_flat = out.view(-1)
for tile in hl.tile(out_flat.size()):
out_flat[tile] = torch.add(input_flat[tile], other_flat[tile])
return out
def add___Tensor_kernel_impl(*args, **kwargs) -> torch.Tensor:
if len(args) >= 2:
input_tensor = args[0]
other = args[1]
elif len(args) == 1 and 'other' in kwargs:
input_tensor = args[0]
other = kwargs['other']
elif 'input' in kwargs and 'other' in kwargs:
input_tensor = kwargs['input']
other = kwargs['other']
else:
raise ValueError("add_ requires 'input' and 'other' arguments")
if 'alpha' in kwargs:
alpha = kwargs['alpha']
other = other * alpha
original_input = args[0]
original_device = input_tensor.device
if not input_tensor.is_cuda:
if not torch.cuda.is_available():
raise RuntimeError("CUDA is not available")
input_tensor = input_tensor.cuda()
if torch.is_tensor(other) and not other.is_cuda:
if not torch.cuda.is_available():
raise RuntimeError("CUDA is not available")
other = other.cuda()
result = add___Tensor_helion_kernel(input_tensor, other)
if original_device.type == 'cpu':
result = result.cpu()
if result.shape == original_input.shape:
original_input.copy_(result)
else:
try:
original_input.copy_(result)
except RuntimeError as e:
raise RuntimeError(f"Cannot perform in-place addition with broadcasting: "
f"input shape {original_input.shape} vs result shape {result.shape}. "
f"Original error: {e}")
return original_input |
| def exp(x: torch.Tensor) -> torch.Tensor: | ||
| out = torch.empty_like(x) | ||
| for tile in hl.tile(x.size()): | ||
| out[tile] = torch.exp(x[tile]) |
There was a problem hiding this comment.
Is this a typical pattern in Helion kernels? In BackendBench, the goal is to avoid falling back to PyTorch, since that would allow the LLM to bypass generating the actual kernel code. However, I'm not sure if Helion handles PyTorch operations differently.
| return out | ||
|
|
||
| # WRAPPER TEMPLATE - Helion kernels are called like Python functions | ||
| def add_kernel_impl(*args, **kwargs) -> torch.Tensor: |
There was a problem hiding this comment.
We are using overload names. Please use something like add__Tensor_kernel_impl here.
2 participants
Stacked PRs:
This PR adds support for Helion DSL to BackendBench and fixes #31
High-level view: Helion is a higher-level abstraction over Triton that:
hl.tile())Added
--dsl helionflag it enables helion DSL for kernel generation.Added
--daemon/--no-daemonflag to control worker process daemon mode (defaultdaemon=True).--no-daemonmust be set for helion because helion autotuning spawns subprocesses.helion autotuning can take significant time so set
HELION_AUTOTUNE_EFFORT=noneenv variable to disable autotuning for faster execution.