Skip to content

Fix visitCall in deviceIR. Always visit argument nodes #178

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed

Fix visitCall in deviceIR. Always visit argument nodes #178

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 0 additions & 3 deletions helion/_compiler/device_ir.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -422,9 +422,6 @@ def _body(self, body: list[ast.stmt]) -> None:

def _to_proxy(self, node: ast.AST) -> object:
assert isinstance(node, ExtendedAST)
type_info = node._type_info
if not type_info.contains_tensor():
return type_info.proxy()
return self.visit(node)
Comment on lines 423 to 425
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This function has become the same as self.visit. Let's delete it and update the callers to call self.visit directly.


def visit_BinOp(self, node: ast.BinOp) -> object:
Expand Down
21 changes: 21 additions & 0 deletions test/test_atomic_add.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,15 @@ def atomic_add_float_kernel(x: torch.Tensor, indices: torch.Tensor) -> torch.Ten
return x


@helion.kernel()
def atomic_add_w_tile_attr(x: torch.Tensor) -> torch.Tensor:
"""Test atomic_add where the index is a symbolic int"""
y = torch.zeros_like(x, device=x.device, dtype=torch.int32)
for tile in hl.tile(x.size(0)):
hl.atomic_add(y, [tile.begin], 1)
return y


class TestAtomicOperations(TestCase):
maxDiff = 16384

Expand Down Expand Up @@ -203,6 +212,18 @@ def bad_atomic_add_kernel(x: torch.Tensor, y: torch.Tensor):
)
self.assertIn("Invalid memory semantic 'ERROR'", str(ctx.exception))

def test_atomic_add_w_tile_attr(self):
"""Test atomic_add where the index is a symbolic int"""
x = torch.randn(20, device=DEVICE)
code, result = code_and_output(
atomic_add_w_tile_attr,
(x,),
block_sizes=[2],
)

expected = torch.tensor([1, 0], device=DEVICE, dtype=torch.int32).repeat(10)
torch.testing.assert_close(result, expected)


if __name__ == "__main__":
unittest.main()
Loading