[Ref Mode] PyTorch reference mode (eager only) #339
Conversation
stack-info: PR: #339, branch: yf225/stack/34
meta-cla
bot
added
the
CLA Signed
helion/ref/hl_patch.py
Outdated
| ) | ||
|
|
||
| # Step 3: Handle block_size (in ref mode, full dim size is always used as block_size) | ||
| block_size_list = [None] * len(end_list) |
There was a problem hiding this comment.
Always use full dim size in ref modes regardless of block_size value
examples/concatenate.py
Outdated
| x_part = hl.load( | ||
| x, [tile0, tile1], extra_mask=(tile1.index < x.size(1))[None, :] |
There was a problem hiding this comment.
Since we are treating tile as a Python slice object in ref mode, tile.index no longer works and we have to use hl.tile_index().
To make the UX better, in a follow-up PR I am thinking of adding a RefTile class that Dynamo can understand, and support tile APIs like .index / .begin / .end in that class.
There was a problem hiding this comment.
We should make tile.index work. Maybe rather than changing the examples we should skip tests in reference mode.
test/test_ref_compile.py
Outdated
| class TestExamplesRefCompile(test_examples.TestExamples): | ||
| """Run all TestExamples tests in reference torch.compile mode via HELION_REF_COMPILE=1.""" | ||
|
|
||
| # NOTE: All tests in TestExamples are run in ref torch.compile(fullgraph=True) mode by default in this test file. |
There was a problem hiding this comment.
Currently all examples in TestExamples pass with ref eager mode and ref compile mode.
Planning to add more ref mode unit tests to cover test_reduce.py, test_associative_scan.py etc. in the next PR.
yf225
changed the title
yf225
changed the title
Stacked PRs:
[Ref Mode] PyTorch reference mode (eager only)
Part of #77.
Please see inline code comments on the PR.