Skip to content

[Type] Allow none return typing, for kernels and funcs#18

Merged
hughperkins merged 7 commits intomainfrom
hp/allow-none-return-typing
Jun 10, 2025
Merged

[Type] Allow none return typing, for kernels and funcs#18
hughperkins merged 7 commits intomainfrom
hp/allow-none-return-typing

Conversation

@hughperkins
Copy link
Copy Markdown
Collaborator

@hughperkins hughperkins commented Jun 7, 2025

Issue: #

Brief Summary

Allow none return typing, for kernels and funcs

This is now possible:

@ti.kernel
def foo() -> None:
    x[0] += 1

This is now possible

@ti.func
def foo() -> None:
    x[0] += 1

copilot:summary

Walkthrough

copilot:walkthrough

@hughperkins
Copy link
Copy Markdown
Collaborator Author

hughperkins commented Jun 10, 2025

Thanks!

@hughperkins hughperkins enabled auto-merge (squash) June 10, 2025 13:57
@hughperkins hughperkins merged commit 041d4a9 into main Jun 10, 2025
16 of 17 checks passed
@hughperkins hughperkins deleted the hp/allow-none-return-typing branch June 10, 2025 22:59
@hughperkins hughperkins mentioned this pull request Apr 22, 2026
Merged
hughperkins added a commit that referenced this pull request Apr 26, 2026
@hughperkins
[Interop] Add _ZerocopyCache and new helpers in _interop.py
d340072 
Introduce the helper machinery that the per-class to_torch / to_numpy methods
will migrate to in subsequent commits. Existing public symbols (can_zerocopy,
dlpack_to_torch, invalidate_zerocopy_cache, current_arch_is_cpu) are preserved
as deprecated shims so the in-tree pre-rework callers continue to work; they
will be removed once every call site is migrated.

New surface:

- _ZerocopyCache: per-instance container with two independent slots (torch
  tensor + numpy ndarray), each filled lazily on first access via
  torch.utils.dlpack.from_dlpack and numpy.from_dlpack respectively. Numpy
  zero-copy now bypasses torch entirely (closes review #6).

- make_zerocopy_cache_if_supported(owner, ...): constructs a cache when
  zero-copy is supported and registers `owner` with `pyquadrants.cache_holders`
  so invalidation is wired automatically (closes review #18).

- get_zerocopy_torch / get_zerocopy_numpy: thin entry points that implement
  the always-zerocopy-then-clone semantic (closes review #15, #16, #21) and
  the Apple Metal double-sync (qd.sync() on read AND torch.mps.synchronize()
  after .clone()/.to(); closes review #1, #22, #23).

Also applies the small lints from the review:
- Module-level constant for the torch>2.9.1 MPS bytes_offset probe; drops the
  pointless lru_cache wrapper around a zero-arg helper (closes review #2).
- ASCII '...' instead of Unicode horizontal ellipsis '\u2026' in the docstring
  (closes review #3).
- Top-level imports for numpy and torch (try/except for the no-torch CI
  case); no per-call lazy imports in the new code path (closes review #7,
  #9). The deprecated shim still does what the existing per-class methods
  expect; the new helpers are torch-clean.

cache_holders is still empty until the next commits register Ndarray /
ScalarField / MatrixField; this commit alone is no-op behaviourally.
hughperkins added a commit that referenced this pull request Apr 26, 2026
@hughperkins
[Interop] Migrate Ndarray to new _ZerocopyCache helpers
96b8bf7 
Convert Ndarray.to_torch / _ndarray_to_numpy / _ndarray_matrix_to_numpy
to thin calls into _interop.get_zerocopy_torch / get_zerocopy_numpy.

Behaviour changes vs. the previous PR state:

- to_torch(copy=True) now ALSO uses the zerocopy export and clones the
  result, instead of skipping zerocopy entirely. Zerocopy DLPack is
  cheaper than the kernel-copy fallback even when followed by a clone
  (closes review #15, #16, #21).

- to_torch on Apple Metal now does both syncs: qd.sync() before the read
  AND torch.mps.synchronize() after the .clone() so the cloned buffer
  is actually populated by the time the call returns (closes review
  #1, #22, #23).

- to_numpy(copy=False) now goes through numpy.from_dlpack directly --
  no torch round-trip, no torch import required (closes review #6).

- to_numpy default (copy=None) keeps the current "independent copy"
  semantics: numpy arrays conventionally outlive their source, so the
  default is the safer one. Zero-copy is opt-in via copy=False.

Cache invalidation switches from the legacy _qd_dlpack_tc attribute to
the per-class _zerocopy_cache (cached_property) + _invalidate_zerocopy_cache
method. Registration in pyquadrants.cache_holders happens automatically
on first cache access (closes review #18).

The existing _reset hook drops its zerocopy invalidation call: that work
is now handled by the cache_holders loop in impl.reset() BEFORE C++
teardown, eliminating the use-after-free that motivated the Genesis
to_numpy() always-copy workaround.
hughperkins added a commit that referenced this pull request Apr 26, 2026
@hughperkins
[Interop] Migrate ScalarField to new _ZerocopyCache helpers
ddbfd9a 
Convert ScalarField.to_torch / to_numpy to thin calls into
_interop.get_zerocopy_torch / get_zerocopy_numpy, mirroring the Ndarray
migration in the previous commit.

- Adds Field._invalidate_zerocopy_cache method (used by all subclasses)
  so the impl.reset() cache_holders loop has a single contract to call.
- Adds ScalarField._zerocopy_cache as a cached_property; constructed once
  per instance instead of re-checking can_zerocopy on every call (closes
  review #17).
- Registers ScalarField in pyquadrants.cache_holders on first cache
  access (closes review #18).
- to_torch(copy=True) now zerocopies and clones rather than skipping
  zerocopy entirely (closes review #15, #16, #21); to_torch on Metal
  does both syncs (closes review #1, #22, #23).
- to_numpy(copy=False) goes through numpy.from_dlpack directly, no
  torch round-trip and no torch import required (closes review #6).

Default to_numpy semantic (copy=None) remains "independent copy" for
lifetime safety; zero-copy is opt-in via copy=False.
hughperkins added a commit that referenced this pull request Apr 26, 2026
@hughperkins
[Interop] Migrate MatrixField to new _ZerocopyCache helpers
2985e0e 
Convert MatrixField.to_torch / to_numpy to thin calls into
_interop.get_zerocopy_torch / get_zerocopy_numpy.

Adds a single _matrix_view_shape() helper that returns (expected_shape,
as_vector); both to_torch and to_numpy now go through it instead of
duplicating the n / m / keep_dims branching (closes review #13, #14).

Same wins as the ScalarField migration:
- _zerocopy_cache as cached_property (closes review #17)
- automatic registration with pyquadrants.cache_holders (closes #18)
- copy=True now zerocopies + clones (closes #15, #16, #21)
- to_numpy(copy=False) goes through numpy.from_dlpack directly, no torch
  round-trip (closes #6)
- Apple Metal double-sync via the centralised helpers (closes #1, #22, #23)

The redundant reshape after zerocopy stays (DLPack returns the field's flat
n-d shape; numpy/torch users expect the matrix dims appended) but it's now
in one place per method instead of two.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@CharlesMasson CharlesMasson CharlesMasson approved these changes

@duburcqa duburcqa duburcqa approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@hughperkins @CharlesMasson @duburcqa