diff --git a/ptodsl/docs/user_guide/04-type-system-and-buffer.md b/ptodsl/docs/user_guide/04-type-system-and-buffer.md index 9ae8265fb2..3aafb9c6a5 100644 --- a/ptodsl/docs/user_guide/04-type-system-and-buffer.md +++ b/ptodsl/docs/user_guide/04-type-system-and-buffer.md @@ -365,7 +365,7 @@ scalar.store(y4, ptr, offset) ## 4.10 Explicit scratch buffers -A scratch buffer is a lane-local temporary buffer inside SIMT code. Allocate scratch buffers with `pto.alloc_buffer`: +A scratch buffer is a local temporary buffer inside an explicit kernel or helper body. Allocate scratch buffers with `pto.alloc_buffer`. A buffer allocated outside dimensioned inline SIMT scopes, or passed to SIMT helpers, can be accessed by those scopes or helpers in the enclosing explicit body: ```text pto.alloc_buffer(shape, dtype) diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index ede6b5b8dc..fbf0d02526 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -2337,21 +2337,20 @@ def _tile_transfer_partition(tv, tile, *, offsets=None, sizes=None, context: str def alloc_buffer(shape, dtype, **kwargs): """ - Allocate SIMT lane-local scratch storage and return an address-like value. + Allocate explicit-body scratch storage and return an address-like value. - The allocation emits an LLVM stack allocation in the surrounding SIMT - helper. UB scratch uses explicit ``pto.castptr`` / ``pto.addptr`` pointer - authoring and an appropriate host launch wrapper. + The allocation emits an LLVM stack allocation in the surrounding explicit + kernel or helper body. UB scratch uses explicit ``pto.castptr`` / + ``pto.addptr`` pointer authoring and an appropriate host launch wrapper. """ if kwargs: unexpected = ", ".join(sorted(kwargs)) raise TypeError( f"pto.alloc_buffer(...) does not accept keyword argument(s): {unexpected}. " - "It only allocates SIMT local buffers; author UB scratch explicitly with " + "It only allocates explicit-body local buffers; author UB scratch explicitly with " "pto.castptr/pto.addptr and pass the dynamic UB byte count at launch." ) _require_explicit_mode("pto.alloc_buffer(...)") - _require_simt_subkernel("pto.alloc_buffer(...)") element_type = _resolve(dtype) element_count = _static_alloc_buffer_element_count(shape) elem_bytes = _element_bytewidth(element_type) diff --git a/ptodsl/ptodsl/_surface_values.py b/ptodsl/ptodsl/_surface_values.py index d0fdbf42bc..5f9eb8f35b 100644 --- a/ptodsl/ptodsl/_surface_values.py +++ b/ptodsl/ptodsl/_surface_values.py @@ -623,6 +623,8 @@ def wrap_like_surface_value(template, value): for dim in valid_shape ) return TileValue(value, **metadata) + if isinstance(template, AllocatedBufferValue): + return AllocatedBufferValue(value, **template.surface_metadata) if isinstance(template, AddressValue): return AddressValue(value) return wrap_surface_value(value) diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index 79f7108b78..a5f8e96d55 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -4359,11 +4359,8 @@ def fake_run_ptoas_cmd(cmd, *, cwd=None): is not None, "alloc_buffer probe should keep allocation inside the SIMT helper body", ) - expect_raises( - RuntimeError, - alloc_buffer_outside_simt_probe.compile, - "pto.alloc_buffer(...) may only be used inside a @pto.simt helper or inline pto.simt() scope", - ) + alloc_buffer_top_level_text = alloc_buffer_outside_simt_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify(alloc_buffer_top_level_text, "alloc_buffer top-level specialization") rmsnorm_alloc_buffer_text = rmsnorm_alloc_buffer_layout_probe.compile().mlir_text() expect_parse_roundtrip_and_verify(rmsnorm_alloc_buffer_text, "RMSNorm hand-authored UB layout specialization") for expected_offset in (4096, 4224, 12416, 20608):