Allow alloc_buffer outside SIMT scopes

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)

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)

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)

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):