[TLE] Add alias support for tle.gpu.alloc to enable shared memory aliasing #826
+582
−4
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
Collaborator
There was a problem hiding this comment. The test file cannot be placed here. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,62 @@ | ||
| """E2E test: verify alloc alias shares the same physical smem via modify-and-observe.""" | ||
| import torch | ||
| import triton | ||
| import triton.language as tl | ||
| import triton.experimental.tle.language as tle | ||
|
|
||
| BLOCK = 128 | ||
|
|
||
|
|
||
| @triton.jit | ||
| def alias_e2e_kernel(in_ptr, out_ptr, N, BLOCK: tl.constexpr): | ||
| pid = tl.program_id(0) | ||
| offs = pid * BLOCK + tl.arange(0, BLOCK) | ||
| mask = offs < N | ||
|
|
||
| # Phase 1: allocate v_smem, write the first batch of data (original input) | ||
| v_smem = tle.gpu.alloc([BLOCK], dtype=tl.float32, scope=tle.gpu.smem) | ||
| tle.gpu.copy(in_ptr + offs, v_smem, [BLOCK]) | ||
|
|
||
| # Phase 2: alias — o_smem reuses v_smem's physical memory | ||
| o_smem = tle.gpu.alloc( | ||
| [BLOCK], dtype=tl.float32, scope=tle.gpu.smem, | ||
| alias=v_smem, alias_offset_bytes=0, | ||
| ) | ||
|
|
||
| # Phase 3: critical — write second batch of data via v_smem (input + BLOCK offset), | ||
| # overwriting the same physical memory. If alias is correct, o_smem | ||
| # should see the overwritten values. | ||
| offs2 = offs + BLOCK | ||
| mask2 = offs2 < N | ||
| tle.gpu.copy(in_ptr + offs2, v_smem, [BLOCK]) | ||
|
|
||
| # Phase 4: read from o_smem | ||
| tle.gpu.copy(o_smem, out_ptr + offs, [BLOCK]) | ||
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| N = 2048 # enough for two BLOCK-sized chunks | ||
| x = torch.randn(N, device="cuda", dtype=torch.float32) | ||
| y = torch.zeros(1024, device="cuda", dtype=torch.float32) # only first 1024 | ||
|
|
||
| grid = (triton.cdiv(1024, BLOCK),) | ||
| alias_e2e_kernel[grid](x, y, 1024, BLOCK=BLOCK) | ||
| torch.cuda.synchronize() | ||
|
|
||
| # Verify: y should equal x[128:1152] (second batch), not x[0:1024] (first batch). | ||
| # This confirms that v_smem's overwrite was observed by o_smem. | ||
| expected = x[128:1152] | ||
| max_diff_alias = (expected - y).abs().max().item() | ||
|
|
||
| # Control: y should NOT equal the first batch | ||
| first_batch = x[:1024] | ||
| max_diff_original = (first_batch - y).abs().max().item() | ||
|
|
||
| if max_diff_alias < 1e-5 and max_diff_original > 1e-5: | ||
| print("PASSED: alias E2E verified") | ||
| print(f" - o_smem matches overwritten data (second batch): max diff = {max_diff_alias:.2e}") | ||
| print(f" - o_smem differs from original data (first batch): max diff = {max_diff_original:.2e}") | ||
| else: | ||
| print(f"FAILED:") | ||
| print(f" diff vs overwritten (should be 0): {max_diff_alias:.2e}") | ||
| print(f" diff vs original (should be >0): {max_diff_original:.2e}") |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,130 @@ | ||
| # flagtree tle | ||
| """TLE alloc alias unit tests — mock-based, no GPU required.""" | ||
|
|
||
| import pytest | ||
| import triton.language as tl | ||
| import triton.experimental.tle.language as tle | ||
|
|
||
|
|
||
| class TestAllocAlias: | ||
|
|
||
| class _FakeTensor: | ||
| def __init__(self, handle, ty): | ||
| self.handle = handle | ||
| self.type = ty | ||
|
|
||
| class _FakeBuilder: | ||
|
|
||
| def __init__(self): | ||
| self.memdesc_type_args = None | ||
| self.memdesc_alias_args = None | ||
| self.swizzled_encoding_args = None | ||
|
|
||
| def get_half_ty(self): | ||
| return "fp16" | ||
|
|
||
| def make_swizzled_shared_encoding_attr(self, vector_size, per_phase, max_phase, order, ctas_per_cga, | ||
| cta_split_num, cta_order): | ||
| self.swizzled_encoding_args = ( | ||
| vector_size, per_phase, max_phase, | ||
| list(order), list(ctas_per_cga), list(cta_split_num), list(cta_order), | ||
| ) | ||
| return "fake_layout" | ||
|
|
||
| def get_memdesc_type(self, shape, element_ty, layout, space, alloc_shape=None): | ||
| self.memdesc_type_args = (list(shape), element_ty, layout, space, alloc_shape) | ||
| return ("memdesc", tuple(shape), element_ty, layout, space, | ||
| None if alloc_shape is None else tuple(alloc_shape)) | ||
|
|
||
| def create_local_alloc(self, *args): | ||
| return "alloc_handle" | ||
|
|
||
| def create_memdesc_alias(self, result_ty, src, offset_bytes): | ||
| self.memdesc_alias_args = (result_ty, src, offset_bytes) | ||
| return "alias_handle" | ||
|
|
||
| class _FakeSemantic: | ||
|
|
||
| def __init__(self): | ||
| self.builder = TestAllocAlias._FakeBuilder() | ||
|
|
||
| def to_tensor(self, value): | ||
| if isinstance(value, TestAllocAlias._FakeTensor): | ||
| return value | ||
| if isinstance(value, bool): | ||
| return TestAllocAlias._FakeTensor(f"pred_{value}", tl.int1) | ||
| if isinstance(value, int): | ||
| return TestAllocAlias._FakeTensor(f"stage_{value}", tl.int32) | ||
| raise TypeError(f"unsupported fake tensor input: {value!r}") | ||
|
|
||
| def _make_buffer(self, shape): | ||
| semantic = self._FakeSemantic() | ||
| layout = tle.gpu.swizzled_shared_layout.make_default(len(shape)) | ||
| return ( | ||
| tle.gpu.buffered_tensor("base", tl.float16, shape, tle.gpu.smem, layout, semantic), | ||
| semantic, | ||
| ) | ||
|
|
||
| def test_alloc_alias_creates_typed_memdesc_alias_view(self): | ||
| """alloc(alias=...) returns a typed view without creating a new allocation.""" | ||
| buffer, semantic = self._make_buffer([4, 16, 32]) | ||
| alias = tle.gpu.alloc( | ||
| (2, 16, 16), | ||
| tl.float16, | ||
| layout=buffer.type.layout, | ||
| alias=buffer, | ||
| alias_offset_bytes=64, | ||
| _semantic=semantic, | ||
| ) | ||
|
|
||
| assert isinstance(alias, tle.gpu.buffered_tensor) | ||
| assert alias.handle == "alias_handle" | ||
| assert alias.shape == [2, 16, 16] | ||
| assert alias.dtype == tl.float16 | ||
| assert alias.type.storage is tle.gpu.smem | ||
| assert semantic.builder.memdesc_alias_args == ( | ||
| ("memdesc", (2, 16, 16), "fp16", "fake_layout", "smem", None), | ||
| "base", | ||
| 64, | ||
| ) | ||
|
|
||
| def test_alloc_alias_rejects_init_value(self): | ||
| buffer, semantic = self._make_buffer([4, 16, 32]) | ||
| init = self._FakeTensor("init", tl.float16) | ||
|
|
||
| with pytest.raises(ValueError, match="alias mode cannot be combined"): | ||
| tle.gpu.alloc( | ||
| (2, 16, 16), | ||
| tl.float16, | ||
| layout=buffer.type.layout, | ||
| init_value=init, | ||
| alias=buffer, | ||
| _semantic=semantic, | ||
| ) | ||
|
|
||
| def test_alloc_alias_rejects_non_smem_buffer(self): | ||
| """alias source must be a shared-memory buffered_tensor.""" | ||
| semantic = self._FakeSemantic() | ||
| fake_buffer = self._FakeTensor("tmem_buf", tl.float16) | ||
|
|
||
| with pytest.raises(ValueError, match="tle.buffered_tensor"): | ||
| tle.gpu.alloc( | ||
| (2, 16, 16), | ||
| tl.float16, | ||
| alias=fake_buffer, | ||
| _semantic=semantic, | ||
| ) | ||
|
|
||
| def test_alloc_alias_invalid_offset_type(self): | ||
| """Bytes are a compile-time argument.""" | ||
| buffer, semantic = self._make_buffer([4, 16, 32]) | ||
|
|
||
| with pytest.raises(ValueError, match="compile-time integer"): | ||
| tle.gpu.alloc( | ||
| (2, 16, 16), | ||
| tl.float16, | ||
| layout=buffer.type.layout, | ||
| alias=buffer, | ||
| alias_offset_bytes=b"not_an_int", | ||
| _semantic=semantic, | ||
| ) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
#ifdef __TLE__