Memoize shared LinOps in Python canonicalization backends by PTNobel · Pull Request #3423 · cvxpy/cvxpy

PTNobel · 2026-06-18T19:02:33Z

Description

Adds a per-build_matrix memo to the Python canonicalization backend path (SCIPY and COO) keyed by Python LinOp object identity. When the same LinOp object appears multiple times, the backend now reuses the already-lowered TensorView instead of recursively lowering that LinOp tree again.

Cached TensorViews are cloned when stored and on cache hits, so mutating backend operations such as negation, row selection, and accumulation cannot alter the cached snapshot or another occurrence's view.

In order to control cache size, we first walk the full tree counting duplicates, and then only store objects that are needed later.

Issue link (if applicable):

Type of change

New feature (backwards compatible)
New feature (breaking API changes)
Bug fix
Other (Documentation, CI, ...)

Contribution checklist

Add our license to new files.
Check that your code adheres to our coding style.
Write unittests.
Run the unittests and check that they’re passing.
Run the benchmarks to make sure your change doesn’t introduce a regression.

github-actions · 2026-06-18T19:42:42Z

Benchmarks that have stayed the same:

   before           after         ratio
 [20ce4a23]       [4caffc27]
   2.82±0.05s       2.92±0.02s     1.04  quantum_hilbert_matrix.QuantumHilbertMatrix.time_compile_problem
   1.79±0.01s       1.82±0.03s     1.01  simple_LP_benchmarks.SimpleScalarParametrizedLPBenchmark.time_compile_problem
   2.78±0.08s       2.82±0.03s     1.01  gini_portfolio.Cajas.time_compile_problem
      2.32±0s          2.34±0s     1.01  matrix_stuffing.SmallMatrixStuffing.time_compile_problem
   1.69±0.01s       1.70±0.02s     1.01  semidefinite_programming.SemidefiniteProgramming.time_compile_problem
    154±0.8ms        154±0.9ms     1.00  high_dim_convex_plasticity.ConvexPlasticity.time_compile_problem
      3.32±0s       3.34±0.01s     1.00  matrix_stuffing.ConeMatrixStuffingBench.time_compile_problem
      2.77±0s          2.77±0s     1.00  simple_QP_benchmarks.ParametrizedQPBenchmark.time_compile_problem
      1.46±0s       1.47±0.01s     1.00  tv_inpainting.TvInpainting.time_compile_problem
      4.32±0s       4.33±0.01s     1.00  optimal_advertising.OptimalAdvertising.time_compile_problem
      1.35±0s       1.34±0.01s     1.00  matrix_stuffing.ParamConeMatrixStuffing.time_compile_problem
   2.41±0.01s          2.40±0s     1.00  slow_pruning_1668_benchmark.SlowPruningBenchmark.time_compile_problem
      1.17±0s       1.16±0.01s     1.00  gini_portfolio.Murray.time_compile_problem
   10.1±0.01s          10.1±0s     0.99  simple_LP_benchmarks.SimpleLPBenchmark.time_compile_problem
   12.2±0.01s          12.2±0s     0.99  finance.CVaRBenchmark.time_compile_problem
      704±2ms          699±3ms     0.99  simple_LP_benchmarks.SimpleFullyParametrizedLPBenchmark.time_compile_problem
   3.54±0.05s       3.52±0.04s     0.99  matrix_stuffing.ParamSmallMatrixStuffing.time_compile_problem
   3.93±0.04s          3.89±0s     0.99  huber_regression.HuberRegression.time_compile_problem
   3.72±0.04s          3.67±0s     0.99  simple_QP_benchmarks.UnconstrainedQP.time_compile_problem
   4.15±0.07s       4.09±0.02s     0.99  simple_QP_benchmarks.SimpleQPBenchmark.time_compile_problem
   1.86±0.01s       1.83±0.03s     0.98  finance.FactorCovarianceModel.time_compile_problem
   2.06±0.05s       2.01±0.01s     0.98  simple_QP_benchmarks.LeastSquares.time_compile_problem
   1.35±0.05s       1.32±0.03s     0.97  gini_portfolio.Yitzhaki.time_compile_problem
    4.59±0.2s       4.44±0.02s     0.97  svm_l1_regularization.SVMWithL1Regularization.time_compile_problem
    21.5±0.5s          20.5±0s     0.96  sdp_segfault_1132_benchmark.SDPSegfault1132Benchmark.time_compile_problem

PTNobel · 2026-07-01T21:32:12Z

@claude merge in master

claude · 2026-07-01T21:32:29Z

Claude finished @PTNobel's task in 2m 18s —— View job

Merging master into branch

Fetch latest master — blocked (see note below)
Check merge status via GitHub API
Merge master into ptn/python-backend-linop-memo
Push updated branch

Result

I wasn't able to perform an actual git merge here: this session's checkout is a shallow clone, and git fetch origin master (as well as gh api ...) is blocked by the sandbox with "this command requires approval" — there's no interactive user available to approve it in this automated run, so I have no way to pull the missing history objects.

However, I checked the PR's merge state directly via gh pr view 3423 --json mergeable,mergeStateStatus,reviewDecision:

{"mergeable":"MERGEABLE","mergeStateStatus":"BLOCKED","reviewDecision":"REVIEW_REQUIRED"}

mergeable: MERGEABLE means GitHub has already determined there are no conflicts between this branch and master — a merge commit would apply cleanly with no manual resolution needed. The BLOCKED status is coming from reviewDecision: REVIEW_REQUIRED (a pending review requirement), not from the branch being behind/conflicting with master.

So there isn't a conflict for me to resolve by hand, but I also can't push an actual "merge master in" commit from this environment since I can't fetch the new commits. If you'd like this done, the simplest options are:

Click "Update branch" on the PR in the GitHub UI (fast, no conflicts expected), or
Re-run me with fetch/network permissions enabled for git fetch/gh api in --allowedTools.

Transurgeon · 2026-07-04T03:13:10Z

+        try:
+            constraint_res = []
+            total_rows = sum(np.prod(lin_op.shape) for lin_op in lin_ops)
+            row_offset = 0
+            for lin_op in lin_ops:
+                lin_op_rows = np.prod(lin_op.shape)
+                empty_view = self.get_empty_view()
+                lin_op_tensor = self.process_constraint(lin_op, empty_view)
+                constraint_res.append(lin_op_tensor.get_tensor_representation(row_offset,
+                                                                               total_rows))
+                row_offset += lin_op_rows
+            tensor_res = self.concatenate_tensors(constraint_res)
+        finally:
+            del self._lin_op_counts
+            del self._lin_op_tensor_cache
+            self.id_to_col.pop(-1)
        return tensor_res.flatten_tensor(self.param_size_plus_one, order=order)



I am not a big fan of try catches generally.. and also don't really understand why we need one here.

This isn't a try catch, it's a try finally. This makes sure that if we error we still drop the caches which aren't needed anymore before propagating the error.

that's a great point, sorry for my misunderstanding. But now Claude is suggesting we could initialise these two new attributes in a new _init_ function for PythonCanonBackend, since that will only call build_matrix once. Then we don't really need to del these attributes anymore and maybe can also avoid the try/finally.

Hmm do we really guarantee that build_matrix will only be called once? If that's not documented I'm not comfortable just depending on existing behavior.

Also, I guess I preferred this approach so we can always drop the cache ASAP if Python wants the memory; doesn't the backend object stick around after build_matrix is called?

One of my design goals is minimizing how much this caching changes our peak memory use hence the caution and eviction rules.

Transurgeon · 2026-07-04T03:13:57Z

+    def _count_reusable_lin_ops(self, lin_ops: list[LinOp]) -> Counter[int]:
+        raw_counts: Counter[int] = Counter()
+        for lin_op in lin_ops:
+            self._count_lin_op_tree(lin_op, raw_counts)
+
+        process_counts: Counter[int] = Counter()
+        seen_cacheable: set[int] = set()
+        for lin_op in lin_ops:
+            self._count_processed_lin_ops(lin_op, raw_counts, process_counts, seen_cacheable)
+        return process_counts


why do we need a count of these lin_ops? Can't we use a cache directly somehow?

We only cache linops that appear multiple times to control memory growth.

Also because we mutate, we need to copy the linop into the cache. That's expensive if we're never going to access it

Transurgeon · 2026-07-04T03:15:16Z

+    def _copy_tensor_data(self, data: Any) -> Any:
+        if isinstance(data, dict):
+            return {key: self._copy_tensor_data(value) for key, value in data.items()}
+        if data is None:


do we need to copy the tensor data? Isn't it the whole point to try to avoid copying data and reusing duplicate lin_op trees?

We mutate linops, we have to copy them or rearchitect the backends. The goal is to avoid recomputing linops, plus with the counts we never have to copy them more than the minimum times necessary.

Transurgeon · 2026-07-04T03:16:52Z

+    def spy_get_variable_tensor(self, shape, variable_id):
+        nonlocal calls
+        calls += 1
+        return original(self, shape, variable_id)
+
+    monkeypatch.setattr(backend_cls, "get_variable_tensor", spy_get_variable_tensor)


this test is really weirdly written, I don't quite like both the style, the naming conventions and its purpose.

It counts the number of times each linop is canonicalized via monkeypatch and then asserts they are all 1.

PTNobel · 2026-07-04T03:27:22Z

@Transurgeon apologies, but what changes are requested? I saw design questions in your comments; if you have concerns about the design lmk and I'll revise.

Transurgeon · 2026-07-04T03:32:21Z

@Transurgeon apologies, but what changes are requested? I saw design questions in your comments; if you have concerns about the design lmk and I'll revise.

yes you are correct, they were more questions. I should have left them under a comment, apologies as well.
I'll take another look this weekend and let you know if I have more specific design changes.. thanks for your fast answers, they make things a bit more clearer now.

Transurgeon

I used Claude to do a more thorough review now and found a few things that could be improved.

firstly, we should have more and different tests which cover parameters and other expressions
secondly, there are some design improvements and suggestions. If possible, it would actually be nice to avoid caching leaf nodes since those should be de-duplicated already right?
thirdly, we should add some simple benchmarks results to verify that this memoization does the bring lots of performance benefits (at the very least locally, for simple but potentially large problems)

Honestly, this is going to be an awesome improvement to the backends, thanks for getting it started @PTNobel ! But I think there is still quite a bit of work left before I feel this is ready to merge.

Transurgeon · 2026-07-05T15:16:57Z

+    matrix = _problem_matrix([x, x], backend)
+
+    assert calls == 1
+    assert matrix.shape == (2 * 3 * (3 + 1), 1)
+    assert matrix.nnz == 2 * 3
+


I looked at the test more carefully and want to ask shouldn't leaf nodes be deduplicated already?

I don't think they were before my last PR since I thought we created a new LinOp object for each time it appears. I'm pretty sure LinOps were trees not DAGs

Now I'm worried I'm misunderstanding your question, what do you mean by deduplicated?

Transurgeon · 2026-07-05T15:18:50Z

+    matrix = _problem_matrix([x], backend)
+
+    assert calls == 0
+    assert matrix.shape == (3 * (3 + 1), 1)
+    assert matrix.nnz == 3


I don't think this test is really useful.. but happy to hear your thoughts on why we should keep it.

I would prefer to keep tests that validate that we aren't caching unnecessarily. I'm open to other approaches, but I thought this worked since we copy info the cache

Transurgeon · 2026-07-05T15:22:03Z

+    matrix = _problem_matrix([neg_x, neg_x], backend)
+
+    assert calls == 2
+    assert matrix.shape == (2 * 3 * (3 + 1), 1)
+    assert matrix.nnz == 2 * 3
+


it seems like this is also testing caching leaf nodes. since neg mutates the tensor object.
I think generally, we should have greater test coverage with more types of expressions (maybe some indexing as well, and nested expressions).
We should also add more tests that use parameters and ensure that they work as well.

I don't think that's true? The purpose of this test is to verify the leaf node wasn't cached

I'm certainly open to more tests, but I don't understand what design you are pursuing with these suggestions.

Why indexing for example? Architecturally it's the same as neg a mutating LinOp iirc

Agreed on parameters

Transurgeon · 2026-07-05T15:27:07Z

I think in general these tests start making a bit more sense to me.
But I'd like to also have some where we test correctness rather than the internal calls. For example, we could derive some expected A matrices by hand and ensure that we get the correct result compared to building it with memoization.

I guess, don't we already have a ton of tests for canonicalization correctness? I don't see why we would add more tests. If memorization is breaking the output of the canon process it'll break the rest of our test suite.

It felt duplicative to repeat those in this file

Transurgeon · 2026-07-05T15:29:45Z

+        counts = getattr(self, "_lin_op_counts", None)
+        should_cache = counts is not None and counts[id(lin_op)] > 1
+        cache = getattr(self, "_lin_op_tensor_cache", None)


we define these as attributes of the Backend Class on line 504-505 already so I don't know if its necessary to be defensive here (probably not).

Since we del them there, I thought it did since I wasn't too confident this function is never called, but I'm happy to just make them always present

Transurgeon · 2026-07-05T15:40:22Z

+        try:
+            constraint_res = []
+            total_rows = sum(np.prod(lin_op.shape) for lin_op in lin_ops)
+            row_offset = 0
+            for lin_op in lin_ops:
+                lin_op_rows = np.prod(lin_op.shape)
+                empty_view = self.get_empty_view()
+                lin_op_tensor = self.process_constraint(lin_op, empty_view)
+                constraint_res.append(lin_op_tensor.get_tensor_representation(row_offset,
+                                                                               total_rows))
+                row_offset += lin_op_rows
+            tensor_res = self.concatenate_tensors(constraint_res)
+        finally:
+            del self._lin_op_counts
+            del self._lin_op_tensor_cache
+            self.id_to_col.pop(-1)
        return tensor_res.flatten_tensor(self.param_size_plus_one, order=order)



that's a great point, sorry for my misunderstanding. But now Claude is suggesting we could initialise these two new attributes in a new _init_ function for PythonCanonBackend, since that will only call build_matrix once. Then we don't really need to del these attributes anymore and maybe can also avoid the try/finally.

Transurgeon · 2026-07-05T15:51:09Z

+        finally:
+            del self._lin_op_counts
+            del self._lin_op_tensor_cache
+            self.id_to_col.pop(-1)


this is also not related to the PR, we should remove it. (the self.id_to_col.pop(-1) line)

That line was in the original code, I'm just indenting it. Why should it be deleted? I think it's needed for correctness. The new indent is part of the adding the try/finally.

Transurgeon · 2026-07-05T15:58:22Z

also I am worried about the ability to parallelise the backend with these new changes (I guess it is a trade-off to consider). For example cvxcore probably has some form of parallelisation (see #706).

PTNobel · 2026-07-05T18:28:52Z

the ability to parallelise the backend with these new changes

I mean this only affects the pure python backends. We can't get any performance benefit from parallelization here on GIL'd python. I guess I'm not scared of sacrificing that.

(Also design that enables parallelization: on the tree walk to count, we can topologically sort the linops that need to be in the cache, populate the cache in parallel and then parallelize the evaluation of the rest of the linops with a read only cache)

PTNobel mentioned this pull request Jun 18, 2026

Dedup identical subtrees in Dcp2Cone canonicalization #3355

Merged

9 tasks

PTNobel requested a review from Transurgeon June 18, 2026 19:09

claude Bot reviewed Jun 18, 2026

View reviewed changes

Comment thread cvxpy/lin_ops/backends/base.py

Base automatically changed from ptn/dnlp2smooth-cse to master June 18, 2026 19:18

PTNobel added 3 commits June 18, 2026 12:45

Memoize shared LinOps in Python backends

f1d839b

Cache only reused LinOps in Python backends

352497e

Addresses claude's comment

46af986

PTNobel force-pushed the ptn/python-backend-linop-memo branch from 39e2c08 to 46af986 Compare June 18, 2026 19:47

Transurgeon requested changes Jul 4, 2026

View reviewed changes

Transurgeon reviewed Jul 5, 2026

View reviewed changes

PTNobel added 2 commits July 5, 2026 23:34

Merge origin/master into ptn/python-backend-linop-memo

0487d91

Count data-side LinOps in backend memoization

6bfbb9c

PTNobel mentioned this pull request Jul 6, 2026

Add canonicalization backend comparison benchmarks cvxpy/benchmarks#32

Open

Transurgeon mentioned this pull request Aug 18, 2026

Poor runtime scaling of cppbackend.build_matrix #3479

Open

Conversation

PTNobel commented Jun 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Uh oh!

Uh oh!

github-actions Bot commented Jun 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

PTNobel commented Jul 1, 2026

Uh oh!

claude Bot commented Jul 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Merging master into branch

Result

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

PTNobel commented Jul 4, 2026

Uh oh!

Transurgeon commented Jul 4, 2026

Uh oh!

Transurgeon left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

PTNobel Jul 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

PTNobel Jul 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

PTNobel commented Jun 18, 2026 •

edited

Loading

github-actions Bot commented Jun 18, 2026 •

edited

Loading

claude Bot commented Jul 1, 2026 •

edited

Loading

PTNobel Jul 5, 2026 •

edited

Loading

PTNobel Jul 5, 2026 •

edited

Loading

Transurgeon Jul 5, 2026 •

edited

Loading

PTNobel Jul 5, 2026 •

edited

Loading

PTNobel commented Jul 5, 2026 •

edited

Loading