pymc-devs / pymc-devs/pytensor

An Elemwise feeding `dest[idx].set(...)` / `.inc(...)` isn't fused into the write (extra temp + copy)

Open
#2,192 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement graph rewriting memory opt numba performance
Dominant language
Python
Stars
644
Forks
208
Avg merge
2d 14h
Merged PRs (30d)
16

Description

Description

When an Elemwise result is written into a subtensor via .set(...)/.inc(...), the Elemwise output is materialized in its own buffer and then copied into the destination, even when the write itself is in-place. The Elemwise could instead write directly into the destination region.

Repro (Numba backend; destination intermediate so the write is in-place):

import pytensor, pytensor.tensor as pt
x  = pt.vector("x")
o0 = pt.vector("o0")
o  = o0 + 1.0
res = o[1:].set(pt.exp(x))
fn = pytensor.function([x, o0], res, mode="NUMBA")
pytensor.dprint(fn, print_memory_map=True)
SetSubtensor{start:}  d={0: [0]}    # in-place on o
 ├─ Add (o0 + 1)
 ├─ Exp                             # own buffer (inplace_pattern={}, destroy_map={})
 │  └─ x
 └─ 1

So per call we allocate exp(x) in its own buffer and then copy it into o[1:]. The in-place SetSubtensor only avoids copying the whole o, not the Elemwise temp.

Elemwise inplace can only destroy one of its own inputs; the write destination o[1:] is not an input to Exp, so there's no path for the Elemwise to write its result directly into the destination.

The advanced-indexing fusion (IndexedElemwise) already avoids this (after #2015) — its inner graph is AdvancedIncSubtensor1{inplace,set}(buffer, Exp(...), idx), so the Elemwise is fused into the scatter and the result lands in the destination with no temp. The basic Subtensor/IncSubtensor path has no equivalent rewrite.

We should generalize the IndexedElemwise fuse-elemwise-into-write machinery to also absorb basic Subtensor/IncSubtensor, unifying the slice and advanced paths.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the provided Numba reproduction and memory map to confirm the extra Elemwise buffer and copy. Then trace the IndexedElemwise fusion machinery described as following #2015 and compare it with the basic Subtensor/IncSubtensor path. Done means the basic slice write absorbs the Elemwise result so it lands directly in the destination without the temporary buffer or copy.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
performance
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.