From e620b5ca9e81ab6c577c710ad7f1043ed92d026b Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Sat, 18 Jul 2026 16:26:36 +0200 Subject: [PATCH 01/17] Fix symbol-promotion and squeezing behavior in Frontend --- .../replacements/array_creation_dace.py | 32 ++++++++ .../python/replacements/array_manipulation.py | 17 ++++- tests/size_scalar_shape_promotion_test.py | 27 +++++++ tests/transpose_unit_dim_squeeze_test.py | 75 +++++++++++++++++++ 4 files changed, 150 insertions(+), 1 deletion(-) create mode 100644 tests/size_scalar_shape_promotion_test.py create mode 100644 tests/transpose_unit_dim_squeeze_test.py diff --git a/dace/frontend/python/replacements/array_creation_dace.py b/dace/frontend/python/replacements/array_creation_dace.py index b9aae4a888..cb2ce65618 100644 --- a/dace/frontend/python/replacements/array_creation_dace.py +++ b/dace/frontend/python/replacements/array_creation_dace.py @@ -12,8 +12,39 @@ from numbers import Integral from typing import Any, Optional +import sympy import numpy as np +from dace import symbolic + + +def promote_size_scalars_in_shape(sdfg: SDFG, shape: Shape) -> None: + """Promote a size expression that was materialized as a size-1 descriptor and then used as an array + SHAPE (``np.empty(Nt + 1)`` -> the scalar ``Nt_plus_1`` reused symbolically) into an SDFG symbol. + Otherwise ``add_datadesc`` tries to ``add_symbol`` the like-named descriptor and raises. Scoped to + the shape's own symbols so a plain scalar argument (never used as a shape) is left untouched.""" + names = set() + for extent in shape: + if isinstance(extent, Integral): + continue + expr = symbolic.pystr_to_symbolic(extent) if isinstance(extent, str) else extent + if isinstance(expr, sympy.Basic): + names |= {str(s) for s in expr.free_symbols} + targets = {n for n in names if n in sdfg.arrays and n not in sdfg.symbols and sdfg.arrays[n].total_size == 1} + if not targets: + return + from dace.transformation.passes.scalar_to_symbol import ScalarToSymbolPromotion, find_promotable_scalars + promotable = find_promotable_scalars(sdfg, transients_only=False) + not_promotable = targets - promotable + if not_promotable: + raise DaceSyntaxError( + None, None, f'Cannot use {sorted(not_promotable)} as an array shape: the like-named size ' + f'descriptor is not promotable to a symbol.') + promo = ScalarToSymbolPromotion() + promo.transients_only = False + promo.ignore = promotable - targets # promote ONLY the size scalars used in this shape + promo.apply_pass(sdfg, {}) + @oprepo.replaces('dace.define_local') @oprepo.replaces('dace.ndarray') @@ -32,6 +63,7 @@ def _define_local_ex(pv: ProgramVisitor, if not isinstance(strides, (list, tuple)): strides = [strides] strides = [int(s) if isinstance(s, Integral) else s for s in strides] + promote_size_scalars_in_shape(sdfg, shape) name = pv.get_target_name() name, _ = sdfg.add_transient(name, shape, diff --git a/dace/frontend/python/replacements/array_manipulation.py b/dace/frontend/python/replacements/array_manipulation.py index f414d239f3..f52501f0c8 100644 --- a/dace/frontend/python/replacements/array_manipulation.py +++ b/dace/frontend/python/replacements/array_manipulation.py @@ -170,7 +170,22 @@ def _transpose(pv: ProgramVisitor, outname = pv.get_target_name() outname, arr2 = sdfg.add_transient(outname, new_shape, restype, arr1.storage, find_new_name=True) - if axes == (1, 0): # Special case for 2D transposition + if axes == (1, 0): # 2D transposition + # The Transpose library node squeezes a unit axis to a vector and then rejects it as "not a + # matrix", so a ``(N, 1)`` / ``(1, N)`` array cannot use it. Fall back to a plain index-swap + # copy (``out[j, i] = in[i, j]``) whenever an extent is 1; it is general over 2D and + # stride-safe. Genuine matrices keep the optimized library node. + if 1 in arr1.shape: + state.add_mapped_tasklet("transpose", + map_ranges={ + "__i": "0:%s" % arr1.shape[0], + "__j": "0:%s" % arr1.shape[1] + }, + inputs={"__inp": Memlet("%s[__i, __j]" % inpname)}, + code="__out = __inp", + outputs={"__out": Memlet("%s[__j, __i]" % outname)}, + external_edges=True) + return outname acc1 = state.add_read(inpname) acc2 = state.add_write(outname) import dace.libraries.linalg # Avoid import loop diff --git a/tests/size_scalar_shape_promotion_test.py b/tests/size_scalar_shape_promotion_test.py new file mode 100644 index 0000000000..a5a7602386 --- /dev/null +++ b/tests/size_scalar_shape_promotion_test.py @@ -0,0 +1,27 @@ +"""A size materialized as a descriptor (``np.empty(Nt+1)`` -> ``Nt_plus_1``) then used as a shape must be +promoted to a symbol, not collide with ``add_symbol`` (which raised ``FileExistsError``).""" +import numpy as np + +import dace + +N = dace.symbol("N") + + +@dace.program +def size_scalar_shape_prog(a: dace.float64[N], Nt: dace.int64): + b = np.empty(Nt + 1, dace.float64) + for i in range(N): + b[i] = a[i] * 2.0 + return b + + +def test_size_scalar_used_as_shape_is_promoted_to_symbol(): + sdfg = size_scalar_shape_prog.to_sdfg(simplify=True) # pre-fix: FileExistsError on Nt_plus_1 + assert "Nt_plus_1" in sdfg.symbols, "the size scalar must become a symbol" + assert "Nt_plus_1" not in sdfg.arrays, "the colliding data descriptor must be gone after promotion" + sdfg.validate() + + +if __name__ == "__main__": + test_size_scalar_used_as_shape_is_promoted_to_symbol() + print("OK") diff --git a/tests/transpose_unit_dim_squeeze_test.py b/tests/transpose_unit_dim_squeeze_test.py new file mode 100644 index 0000000000..f6cb22a2ef --- /dev/null +++ b/tests/transpose_unit_dim_squeeze_test.py @@ -0,0 +1,75 @@ +"""A 2D array with a unit dim (``(N, 1)`` / ``(1, N)``) must transpose to the swapped shape: the +DaCe frontend used to squeeze the unit dim and reject ``(N, 1).T`` as "not a matrix". An integer +index, by contrast, squeezes its axis (``x[:, 1]`` is ``(N,)``) per numpy semantics.""" +import numpy as np + +import dace + +N = dace.symbol("N") +M = dace.symbol("M") + + +@dace.program +def col_transpose_matmul(x: dace.float64[N, 1], a: dace.float64[N, M]): + return x.T @ a # (1, N) @ (N, M) -> (1, M) + + +def test_column_vector_transpose_matmul(): + n, m = 5, 4 + rng = np.random.default_rng(0) + x, a = rng.random((n, 1)), rng.random((n, m)) + got = np.asarray(col_transpose_matmul(x.copy(), a.copy())) + ref = x.T @ a + assert got.shape == ref.shape == (1, m) + assert np.allclose(got.reshape(ref.shape), ref) + + +@dace.program +def col_outer_product(x: dace.float64[N, 1]): + return x @ x.T # (N, 1) @ (1, N) -> (N, N) outer product + + +def test_column_vector_outer_product(): + n = 6 + rng = np.random.default_rng(1) + x = rng.random((n, 1)) + got = np.asarray(col_outer_product(x.copy())) + ref = x @ x.T + assert got.shape == ref.shape == (n, n) + assert np.allclose(got, ref) + + +@dace.program +def row_vector_transpose(x: dace.float64[1, N]): + return x.T # (1, N) -> (N, 1) + + +def test_row_vector_transpose(): + n = 7 + rng = np.random.default_rng(2) + x = rng.random((1, n)) + got = np.asarray(row_vector_transpose(x.copy())) + assert got.shape == (n, 1) + assert np.allclose(got.reshape(n, 1), x.T) + + +@dace.program +def index_squeezes_axis(x: dace.float64[N, 3]): + col = x[:, 1] # an integer index squeezes the axis: (N,) + return np.sum(col) + + +def test_integer_index_squeezes_axis(): + n = 6 + rng = np.random.default_rng(3) + x = rng.random((n, 3)) + got = np.asarray(index_squeezes_axis(x.copy())) + assert np.allclose(got, np.sum(x[:, 1])) + + +if __name__ == "__main__": + test_column_vector_transpose_matmul() + test_column_vector_outer_product() + test_row_vector_transpose() + test_integer_index_squeezes_axis() + print("OK") From f8f53d86fa496f11a71953c82bdaf33241b15799 Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Sat, 18 Jul 2026 21:35:51 +0200 Subject: [PATCH 02/17] Update docstring for promote_size_scalars_in_shape --- dace/frontend/python/replacements/array_creation_dace.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/dace/frontend/python/replacements/array_creation_dace.py b/dace/frontend/python/replacements/array_creation_dace.py index cb2ce65618..4806e21961 100644 --- a/dace/frontend/python/replacements/array_creation_dace.py +++ b/dace/frontend/python/replacements/array_creation_dace.py @@ -19,10 +19,8 @@ def promote_size_scalars_in_shape(sdfg: SDFG, shape: Shape) -> None: - """Promote a size expression that was materialized as a size-1 descriptor and then used as an array - SHAPE (``np.empty(Nt + 1)`` -> the scalar ``Nt_plus_1`` reused symbolically) into an SDFG symbol. - Otherwise ``add_datadesc`` tries to ``add_symbol`` the like-named descriptor and raises. Scoped to - the shape's own symbols so a plain scalar argument (never used as a shape) is left untouched.""" + """Promote a size expression that was materialized as a size-1 descriptor and then used as an symbol + SHAPE (``np.empty(Nt_plus_1)`` -> ``Nt_plus_1`` has to be a symbol) into an SDFG symbol.""" names = set() for extent in shape: if isinstance(extent, Integral): From 638022d901f5bb6270f55087d00fd9f6cd7cf619 Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Mon, 20 Jul 2026 17:27:00 +0200 Subject: [PATCH 03/17] Squeeze only rank-reducing indices in MatMul operands A subsets.Range entry of size 1 is ambiguous: a rank-reducing index into a larger dimension and a genuine extent-1 dimension are both stored as (0, 0, 1). Range.squeeze() drops both, so a full read of a (NQ, 1, NP) view looks identical to a 2D slice. np.reshape(x, (NQ, 1, NP)) builds exactly that. SpecializeMatMul dispatched on the squeezed sizes and saw a matrix, picking Gemm, whose validate re-read the unsqueezed subset and raised "matrix-matrix product only supported on matrices". npbench's doitgen hits this at every size, on both the simplified and the auto-optimized pipeline. The descriptor is the authority: an entry of size 1 is rank-reducing only where the descriptor's own extent is larger. Under that rule the reshaped operand stays 3D and dispatches to BatchedMatMul, which is what numpy does, while an index into a larger dimension still squeezes to Gemm. BatchedMatMul.validate went through the same shared helper so both nodes agree on operand rank. Regressed in afd0efe48, which removed the squeezing from the validate methods but left SpecializeMatMul dispatching on squeezed sizes. That commit also dropped the (NQ, 1, NP) reshape from the in-tree doitgen test, which is why CI stayed green. --- dace/libraries/blas/nodes/batched_matmul.py | 18 +--- dace/libraries/blas/nodes/matmul.py | 19 +++- tests/matmul_unit_dim_squeeze_test.py | 111 ++++++++++++++++++++ 3 files changed, 131 insertions(+), 17 deletions(-) create mode 100644 tests/matmul_unit_dim_squeeze_test.py diff --git a/dace/libraries/blas/nodes/batched_matmul.py b/dace/libraries/blas/nodes/batched_matmul.py index 0abc94c1fc..00998372cf 100644 --- a/dace/libraries/blas/nodes/batched_matmul.py +++ b/dace/libraries/blas/nodes/batched_matmul.py @@ -458,20 +458,13 @@ def validate(self, sdfg, state): in_edges = state.in_edges(self) if len(in_edges) != 2: raise ValueError("Expected exactly two inputs to batched matrix-matrix product") - for _, _, _, dst_conn, memlet in state.in_edges(self): - if dst_conn == '_a': - subset = dc(memlet.subset) - subset.squeeze() - size0 = subset.size() - if dst_conn == '_b': - subset = dc(memlet.subset) - subset.squeeze() - size1 = subset.size() out_edges = state.out_edges(self) if len(out_edges) != 1: raise ValueError("Expected exactly one output from " "batched matrix-matrix product") - out_memlet = out_edges[0].data + # Sizes come from the shared operand helper, which squeezes only rank-reducing indices. + a, b, c = _get_matmul_operands(self, state, sdfg) + size0, size1, size2 = a[4], b[4], c[4] # Both inputs must be at least 2D if len(size0) < 2: @@ -494,9 +487,8 @@ def validate(self, sdfg, state): raise ValueError("Inputs to matrix-matrix product must agree in the k-dimension") # Output must have batch dimensions - if len(out_memlet.subset) < 3: - raise ValueError( - f"Batched matrix-matrix product output must be at least 3D, got {len(out_memlet.subset)} dimensions") + if len(size2) < 3: + raise ValueError(f"Batched matrix-matrix product output must be at least 3D, got {len(size2)} dimensions") # Numpy replacement diff --git a/dace/libraries/blas/nodes/matmul.py b/dace/libraries/blas/nodes/matmul.py index f5f308b2c8..c5add116f0 100644 --- a/dace/libraries/blas/nodes/matmul.py +++ b/dace/libraries/blas/nodes/matmul.py @@ -7,6 +7,17 @@ from math import prod +def unit_extent_dims(desc): + """Dimensions that are size-1 in the descriptor itself. + + A subset entry of size 1 is ambiguous: it is either a rank-reducing index into a larger + dimension (``A[r, 0:N]`` on ``A[R, N]`` -- the dimension is gone) or a full read of a + dimension the data actually has (``V[0:N, 0, 0:P]`` on a ``(N, 1, P)`` view -- the dimension + is part of the shape). Only the former may be squeezed away. + """ + return [i for i, extent in enumerate(desc.shape) if extent == 1] + + def _get_matmul_operands(node, state, sdfg, name_lhs="_a", name_rhs="_b", name_out="_c"): """Returns the matrix multiplication input edges, arrays, and shape.""" res_lhs = None @@ -15,9 +26,9 @@ def _get_matmul_operands(node, state, sdfg, name_lhs="_a", name_rhs="_b", name_o if edge.dst_conn in [name_lhs, name_rhs]: size = edge.data.subset.size() squeezed = dc(edge.data.subset) - squeezed_dims = squeezed.squeeze() - squeezed_size = squeezed.size() outer_array = sdfg.data(dace.sdfg.find_input_arraynode(state, edge).data) + squeezed_dims = squeezed.squeeze(ignore_indices=unit_extent_dims(outer_array)) + squeezed_size = squeezed.size() strides = list(outer_array.strides) squeezed_strides = [s for i, s in enumerate(outer_array.strides) if i in squeezed_dims] res = edge, outer_array, size, strides, squeezed_size, squeezed_strides @@ -28,9 +39,9 @@ def _get_matmul_operands(node, state, sdfg, name_lhs="_a", name_rhs="_b", name_o elif edge.src_conn == name_out: size = edge.data.subset.size() squeezed = dc(edge.data.subset) - squeezed_dims = squeezed.squeeze() - squeezed_size = squeezed.size() outer_array = sdfg.data(dace.sdfg.find_output_arraynode(state, edge).data) + squeezed_dims = squeezed.squeeze(ignore_indices=unit_extent_dims(outer_array)) + squeezed_size = squeezed.size() strides = list(outer_array.strides) squeezed_strides = [s for i, s in enumerate(outer_array.strides) if i in squeezed_dims] res_out = edge, outer_array, size, strides, squeezed_size, squeezed_strides diff --git a/tests/matmul_unit_dim_squeeze_test.py b/tests/matmul_unit_dim_squeeze_test.py new file mode 100644 index 0000000000..6d4d65240a --- /dev/null +++ b/tests/matmul_unit_dim_squeeze_test.py @@ -0,0 +1,111 @@ +"""A ``subsets.Range`` entry of size 1 is ambiguous, and the MatMul library nodes read it two +incompatible ways. + +``Range`` stores both a rank-reducing index and a genuine extent-1 dimension as the same triple +``(0, 0, 1)``, so ``A[r, 0:N, 0:P]`` (dimension indexed away, rank 2) and ``V[0:N, 0, 0:P]`` +(full read of a ``(N, 1, P)`` view, rank 3) are indistinguishable at the subset level. +``Range.squeeze()`` drops both. + +``np.reshape(x, (NQ, 1, NP))`` builds exactly the second kind: a 3D view whose middle dimension is +part of the shape. ``SpecializeMatMul`` dispatched on the squeezed sizes and saw a matrix, so it +picked ``Gemm``, whose ``validate`` re-read the unsqueezed subset and raised +"matrix-matrix product only supported on matrices". npbench's doitgen hits this on every size. + +The descriptor is the authority: a subset entry of size 1 is rank-reducing only where the +descriptor's own extent is larger. Under that rule the reshape stays 3D and dispatches to +``BatchedMatMul``, while a real index into a larger dimension still squeezes to ``Gemm``. +""" +import numpy as np +import pytest + +import dace +from dace.libraries.blas.nodes.batched_matmul import BatchedMatMul +from dace.libraries.blas.nodes.gemm import Gemm +from dace.libraries.blas.nodes.matmul import MatMul +from dace.transformation.auto.auto_optimize import auto_optimize + +NR, NQ, NP = (dace.symbol(s, dtype=dace.int64) for s in ('NR', 'NQ', 'NP')) + + +@dace.program +def doitgen_reshape(A: dace.float64[NR, NQ, NP], C4: dace.float64[NP, NP]): + # npbench polybench/doitgen, verbatim: the (NQ, 1, NP) reshape is the point of the benchmark. + for r in range(NR): + A[r, :, :] = np.reshape(np.reshape(A[r], (NQ, 1, NP)) @ C4, (NQ, NP)) + + +@dace.program +def indexed_slice_matmul(A: dace.float64[NR, NQ, NP], C4: dace.float64[NP, NP]): + # A[r] indexes dimension 0 away: a genuine 2D operand that must still reach Gemm. + for r in range(NR): + A[r, :, :] = A[r] @ C4 + + +def reference(A, C4): + NR, NQ, NP = A.shape + return np.reshape(np.reshape(A, (NR, NQ, 1, NP)) @ C4, (NR, NQ, NP)) + + +def initialize(nr, nq, np_): + A = np.fromfunction(lambda i, j, k: ((i * j + k) % np_) / np_, (nr, nq, np_), dtype=np.float64) + C4 = np.fromfunction(lambda i, j: (i * j % np_) / np_, (np_, np_), dtype=np.float64) + return A, C4 + + +def count_nodes(sdfg, nodetype): + return sum(1 for node, _ in sdfg.all_nodes_recursive() if isinstance(node, nodetype)) + + +def specialize_matmuls(sdfg): + """Expand the MatMul meta-nodes one level, leaving the chosen specialization in the graph.""" + for node, state in list(sdfg.all_nodes_recursive()): + if type(node) is MatMul: + node.expand(state) + + +@pytest.mark.parametrize('optimize', [False, True]) +@pytest.mark.parametrize('sizes', [(3, 4, 5), (1, 1, 1), (8, 10, 12), (5, 1, 7)]) +def test_reshape_unit_dim_matmul(optimize, sizes): + nr, nq, np_ = sizes + A, C4 = initialize(nr, nq, np_) + ref = reference(A, C4) + + sdfg = doitgen_reshape.to_sdfg(simplify=False) + sdfg.simplify() + if optimize: + auto_optimize(sdfg, dace.dtypes.DeviceType.CPU, symbols=dict(NR=nr, NQ=nq, NP=np_)) + + sdfg(A=A, C4=C4, NR=nr, NQ=nq, NP=np_) + assert np.allclose(A, ref) + + +def test_reshape_unit_dim_dispatches_to_batched(): + """The reshaped operand keeps its unit dimension, so the matmul stays batched.""" + sdfg = doitgen_reshape.to_sdfg(simplify=False) + sdfg.simplify() + specialize_matmuls(sdfg) + assert count_nodes(sdfg, BatchedMatMul) == 1 + assert count_nodes(sdfg, Gemm) == 0 + + +def test_indexed_dim_still_squeezes_to_gemm(): + """An index into a larger dimension is rank-reducing, so the operand is a plain matrix.""" + nr, nq, np_ = 3, 4, 5 + A, C4 = initialize(nr, nq, np_) + ref = reference(A, C4) + + sdfg = indexed_slice_matmul.to_sdfg(simplify=False) + sdfg.simplify() + specialize_matmuls(sdfg) + assert count_nodes(sdfg, Gemm) == 1 + assert count_nodes(sdfg, BatchedMatMul) == 0 + + sdfg(A=A, C4=C4, NR=nr, NQ=nq, NP=np_) + assert np.allclose(A, ref) + + +if __name__ == '__main__': + for opt in (False, True): + test_reshape_unit_dim_matmul(opt, (3, 4, 5)) + test_reshape_unit_dim_dispatches_to_batched() + test_indexed_dim_still_squeezes_to_gemm() From c0b05c4e2ad5504ce4edaa2bff6bc64bd9c51542 Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Mon, 20 Jul 2026 20:48:58 +0200 Subject: [PATCH 04/17] Track integer-index dimensions in Range, revert MatMul operand change numpy distinguishes `a[0]` from `a[0:1]`: an integer index removes its dimension from the result rank, a slice never does, not even at extent 1. `Range` stored both as `(i, i, 1)`, so the distinction was unrecoverable, and consumers needing the rank of an access had to guess by squeezing every extent-1 dimension. That guess is wrong whenever a sliced dimension happens to have extent 1, which is how `np.reshape(x, (NQ, 1, NP)) @ C4` ends up dispatched as a matrix product whose validator then rejects it as not a matrix. `Range` now carries a per-dimension `index_dims` flag, with `rank()`, `rank_dims()` and `is_index_dim()` derived from it, and the frontend records it for mixed subscripts. `squeeze()` deliberately keeps its meaning -- it is `np.squeeze`, dropping every extent-1 dimension -- and now maintains the flags rather than discarding them; `unsqueeze()` inserts slices, matching `np.expand_dims`. Because only a degenerate range can be an index, `is_index_dim` cross-checks the flag against the range instead of trusting it, so widening a dimension through `__setitem__` or a direct `ranges` assignment cannot leave it looking indexed. The flags survive JSON but not the string form, which still renders a degenerate range as `i`. That is deliberate: subset strings are re-parsed as symbolic expressions in places that reject slice syntax (`scalar_to_symbol` embeds them in interstate assignments) and are used as dictionary keys against the original text (`sdfg_to_tree` de-aliasing), so the rendering cannot change. Generated code is byte-identical and existing SDFGs load at the rank they were saved with, since a missing `indexed` entry means slice. Also reverts the MatMul operand change from 638022d90, which redefined what `_get_matmul_operands` reports as operand size for every BLAS node. It made a 1x1 contraction look like a matrix while its descriptor stayed a scalar, so cblas_dgemm was handed a `double` where it wanted a `double*` (the CI failure in tests/numpy/einsum_test.py::test_opteinsum, reproducible under DACE_optimizer_autooptimize=1). It also broke Gemv's pure expansion and silently dropped alpha, beta and summation WCR by rerouting products to BatchedMatMul, which honours none of them. doitgen is still unfixed; doitgen_repro.py drives it and is temporary. --- dace/frontend/python/memlet_parser.py | 4 +- dace/libraries/blas/nodes/batched_matmul.py | 18 +- dace/libraries/blas/nodes/matmul.py | 19 +- dace/subsets.py | 127 +++++++++++-- doitgen_repro.py | 106 +++++++++++ tests/matmul_unit_dim_squeeze_test.py | 111 ----------- tests/sdfg/subset_index_dims_test.py | 192 ++++++++++++++++++++ 7 files changed, 431 insertions(+), 146 deletions(-) create mode 100644 doitgen_repro.py delete mode 100644 tests/matmul_unit_dim_squeeze_test.py create mode 100644 tests/sdfg/subset_index_dims_test.py diff --git a/dace/frontend/python/memlet_parser.py b/dace/frontend/python/memlet_parser.py index 1be7d0d79b..969c9c00d9 100644 --- a/dace/frontend/python/memlet_parser.py +++ b/dace/frontend/python/memlet_parser.py @@ -69,7 +69,9 @@ def _ndslice_to_subset(ndslice): for i in range(len(ndslice)): if not is_tuple[i]: ndslice[i] = (ndslice[i], ndslice[i], 1) - return subsets.Range(ndslice) + # Record which dimensions the user wrote as an integer index. Widening them to (i, i, 1) above + # makes them indistinguishable from a ``i:i+1`` slice, but numpy gives the two different ranks. + return subsets.Range(ndslice, index_dims=[not t for t in is_tuple]) def _parse_dim_atom(das, atom): diff --git a/dace/libraries/blas/nodes/batched_matmul.py b/dace/libraries/blas/nodes/batched_matmul.py index 00998372cf..0abc94c1fc 100644 --- a/dace/libraries/blas/nodes/batched_matmul.py +++ b/dace/libraries/blas/nodes/batched_matmul.py @@ -458,13 +458,20 @@ def validate(self, sdfg, state): in_edges = state.in_edges(self) if len(in_edges) != 2: raise ValueError("Expected exactly two inputs to batched matrix-matrix product") + for _, _, _, dst_conn, memlet in state.in_edges(self): + if dst_conn == '_a': + subset = dc(memlet.subset) + subset.squeeze() + size0 = subset.size() + if dst_conn == '_b': + subset = dc(memlet.subset) + subset.squeeze() + size1 = subset.size() out_edges = state.out_edges(self) if len(out_edges) != 1: raise ValueError("Expected exactly one output from " "batched matrix-matrix product") - # Sizes come from the shared operand helper, which squeezes only rank-reducing indices. - a, b, c = _get_matmul_operands(self, state, sdfg) - size0, size1, size2 = a[4], b[4], c[4] + out_memlet = out_edges[0].data # Both inputs must be at least 2D if len(size0) < 2: @@ -487,8 +494,9 @@ def validate(self, sdfg, state): raise ValueError("Inputs to matrix-matrix product must agree in the k-dimension") # Output must have batch dimensions - if len(size2) < 3: - raise ValueError(f"Batched matrix-matrix product output must be at least 3D, got {len(size2)} dimensions") + if len(out_memlet.subset) < 3: + raise ValueError( + f"Batched matrix-matrix product output must be at least 3D, got {len(out_memlet.subset)} dimensions") # Numpy replacement diff --git a/dace/libraries/blas/nodes/matmul.py b/dace/libraries/blas/nodes/matmul.py index c5add116f0..f5f308b2c8 100644 --- a/dace/libraries/blas/nodes/matmul.py +++ b/dace/libraries/blas/nodes/matmul.py @@ -7,17 +7,6 @@ from math import prod -def unit_extent_dims(desc): - """Dimensions that are size-1 in the descriptor itself. - - A subset entry of size 1 is ambiguous: it is either a rank-reducing index into a larger - dimension (``A[r, 0:N]`` on ``A[R, N]`` -- the dimension is gone) or a full read of a - dimension the data actually has (``V[0:N, 0, 0:P]`` on a ``(N, 1, P)`` view -- the dimension - is part of the shape). Only the former may be squeezed away. - """ - return [i for i, extent in enumerate(desc.shape) if extent == 1] - - def _get_matmul_operands(node, state, sdfg, name_lhs="_a", name_rhs="_b", name_out="_c"): """Returns the matrix multiplication input edges, arrays, and shape.""" res_lhs = None @@ -26,9 +15,9 @@ def _get_matmul_operands(node, state, sdfg, name_lhs="_a", name_rhs="_b", name_o if edge.dst_conn in [name_lhs, name_rhs]: size = edge.data.subset.size() squeezed = dc(edge.data.subset) - outer_array = sdfg.data(dace.sdfg.find_input_arraynode(state, edge).data) - squeezed_dims = squeezed.squeeze(ignore_indices=unit_extent_dims(outer_array)) + squeezed_dims = squeezed.squeeze() squeezed_size = squeezed.size() + outer_array = sdfg.data(dace.sdfg.find_input_arraynode(state, edge).data) strides = list(outer_array.strides) squeezed_strides = [s for i, s in enumerate(outer_array.strides) if i in squeezed_dims] res = edge, outer_array, size, strides, squeezed_size, squeezed_strides @@ -39,9 +28,9 @@ def _get_matmul_operands(node, state, sdfg, name_lhs="_a", name_rhs="_b", name_o elif edge.src_conn == name_out: size = edge.data.subset.size() squeezed = dc(edge.data.subset) - outer_array = sdfg.data(dace.sdfg.find_output_arraynode(state, edge).data) - squeezed_dims = squeezed.squeeze(ignore_indices=unit_extent_dims(outer_array)) + squeezed_dims = squeezed.squeeze() squeezed_size = squeezed.size() + outer_array = sdfg.data(dace.sdfg.find_output_arraynode(state, edge).data) strides = list(outer_array.strides) squeezed_strides = [s for i, s in enumerate(outer_array.strides) if i in squeezed_dims] res_out = edge, outer_array, size, strides, squeezed_size, squeezed_strides diff --git a/dace/subsets.py b/dace/subsets.py index 3951d91491..f78b61671d 100644 --- a/dace/subsets.py +++ b/dace/subsets.py @@ -305,7 +305,20 @@ def _tuple_to_symexpr(val): class Range(Subset): """ Subset defined in terms of a fixed range. """ - def __init__(self, ranges): + def __init__(self, ranges, index_dims: Optional[Sequence[bool]] = None): + """ + :param ranges: The per-dimension (begin, end, step[, tile]) tuples. + :param index_dims: Which dimensions were written as an integer index rather than a slice. + numpy drops those dimensions from the result rank (``a[0]`` is one rank + lower than ``a[0:1]``), and ``(begin, end, step)`` cannot express the + difference on its own since both are ``(i, i, 1)``. Defaults to all + slices, which is what a full-array access means. + + :note: The flags survive JSON serialization but not the string form, which renders a + degenerate range as ``i`` either way. Subset strings are re-parsed as symbolic + expressions in places that reject slice syntax (see + ``dace.transformation.passes.scalar_to_symbol``), so the rendering cannot change. + """ parsed_ranges = [] parsed_tiles = [] for r in ranges: @@ -318,14 +331,58 @@ def __init__(self, ranges): parsed_tiles.append(symbolic.pystr_to_symbolic(r[3])) self.ranges = parsed_ranges self.tile_sizes = parsed_tiles + if index_dims is None: + self.index_dims = [False] * len(parsed_ranges) + else: + index_dims = list(index_dims) + if len(index_dims) != len(parsed_ranges): + raise ValueError(f"Expected {len(parsed_ranges)} index flags, got {len(index_dims)}") + self.index_dims = index_dims + + def is_index_dim(self, dim: int) -> bool: + """ + Returns True if the given dimension was accessed with an integer index. + + Only a degenerate range can be an index, so the flag is cross-checked against the range + itself rather than trusted outright: ``ranges`` is mutable through ``__setitem__`` and + direct assignment, and widening a dimension must not leave it looking indexed. + + :param dim: The dimension to query. + :return: True if the dimension was written as an integer index. + """ + if not self.index_dims[dim]: + return False + rb, re, _ = self.ranges[dim] + return symbolic.equal_valued(0, re - rb) + + def rank_dims(self) -> List[int]: + """ + Returns the dimensions that survive numpy rank reduction, i.e. those not indexed away. + + An integer index removes its dimension (``a[0, 0:N]`` has rank 1); a slice never does, not + even when its extent is 1 (``a[0:1, 0:N]`` has rank 2). This is the numpy rule, and it is + independent of :meth:`squeeze`, which removes every extent-1 dimension regardless of how it + was written (``np.squeeze``). + + :return: The indices of the dimensions that are not indexed away. + """ + return [i for i in range(len(self.ranges)) if not self.is_index_dim(i)] + + def rank(self) -> int: + """ + Returns the numpy rank of this access: the number of dimensions not indexed away. + + :return: The number of dimensions that survive rank reduction. + """ + return len(self.ranges) - sum(1 for i in range(len(self.ranges)) if self.is_index_dim(i)) @staticmethod def from_indices(indices: Union["Indices", Sequence[int | str | symbolic.SymbolicType]]): if isinstance(indices, Indices): - return Range([(i, i, 1) for i in indices.indices]) + return Range([(i, i, 1) for i in indices.indices], index_dims=[True] * len(indices.indices)) indices = [symbolic.pystr_to_symbolic(i) for i in indices] - return Range([(i, i, 1) for i in indices]) + return Range([(i, i, 1) for i in indices], index_dims=[True] * len(indices)) def to_json(self): ret = [] @@ -333,8 +390,11 @@ def to_json(self): def a2s(obj): return symbolic.serialize_symbolic(obj) - for (start, end, step), tile in zip(self.ranges, self.tile_sizes): - ret.append({'start': a2s(start), 'end': a2s(end), 'step': a2s(step), 'tile': a2s(tile)}) + for (start, end, step), tile, indexed in zip(self.ranges, self.tile_sizes, self.index_dims): + entry = {'start': a2s(start), 'end': a2s(end), 'step': a2s(step), 'tile': a2s(tile)} + if indexed: + entry['indexed'] = True + ret.append(entry) return {'type': 'Range', 'ranges': ret} @@ -349,12 +409,16 @@ def from_json(obj, context=None): ranges = obj['ranges'] tuples = [] + index_dims = [] for r in ranges: tuples.append((_symbolic_deserializer(r['start'], context), _symbolic_deserializer(r['end'], context), _symbolic_deserializer(r['step'], context), _symbolic_deserializer(r['tile'], context))) + # Absent in SDFGs written before index dimensions were tracked; a slice is the safe + # default there, since it preserves the rank the file was saved with. + index_dims.append(bool(r.get('indexed', False))) - return Range(tuples) + return Range(tuples, index_dims=index_dims) @staticmethod def from_array(array: 'dace.data.Data'): @@ -369,7 +433,8 @@ def __hash__(self): def __add__(self, other): return Range( - ((*ranges, tile) for ranges, tile in zip(self.ranges + other.ranges, self.tile_sizes + other.tile_sizes))) + ((*ranges, tile) for ranges, tile in zip(self.ranges + other.ranges, self.tile_sizes + other.tile_sizes)), + index_dims=self.index_dims + other.index_dims) def __deepcopy__(self, memo) -> 'Range': """Performs a deepcopy of ``self``. @@ -381,6 +446,7 @@ def __deepcopy__(self, memo) -> 'Range': node = object.__new__(Range) node.ranges = self.ranges.copy() node.tile_sizes = self.tile_sizes.copy() + node.index_dims = self.index_dims.copy() return node @@ -507,7 +573,7 @@ def offset(self, other, negative, indices=None, offset_end=True): def offset_new(self, other, negative, indices=None, offset_end=True): if other is None: - return Range(self.ranges) + return Range(self.ranges, index_dims=self.index_dims) if not isinstance(other, Subset): if isinstance(other, (list, tuple)): other = Range.from_indices(other) @@ -517,8 +583,11 @@ def offset_new(self, other, negative, indices=None, offset_end=True): if indices is None: indices = set(range(len(self.ranges))) off = other.min_element() + # Offsetting shifts a subset, it never changes how a dimension was written, so the index + # flags follow the dimensions that are kept. return Range([(self.ranges[i][0] + mult * off[i], self.ranges[i][1] if not offset_end else - (self.ranges[i][1] + mult * off[i]), self.ranges[i][2]) for i in indices]) + (self.ranges[i][1] + mult * off[i]), self.ranges[i][2]) for i in indices], + index_dims=[self.index_dims[i] for i in indices]) def dims(self): return len(self.ranges) @@ -573,8 +642,10 @@ def reorder(self, order: Sequence[int]) -> None: """ new_ranges = [self.ranges[o] for o in order] new_tile_sizes = [self.tile_sizes[o] for o in order] + new_index_dims = [self.index_dims[o] for o in order] self.ranges = new_ranges self.tile_sizes = new_tile_sizes + self.index_dims = new_index_dims @staticmethod def dim_to_string(d, t=1): @@ -612,6 +683,7 @@ def from_string(string): # regtile_j * rs_j : min(K, regtile_j * rs_j + rs_j) ranges = [] + index_dims = [] # Split string to tokens separated by colons. # tokens = [ @@ -670,6 +742,7 @@ def from_string(string): if len(uni_dim_tokens) < 2: value = symbolic.pystr_to_symbolic(uni_dim_tokens[0].strip()) ranges.append((value, value, 1)) + index_dims.append(True) continue #return Range(ranges) # If dimension has more than 4 tokens, the range is invalid @@ -716,8 +789,9 @@ def from_string(string): raise SyntaxError("Invalid range: {}".format(string)) # Append range ranges.append((begin, end, step, tsize)) + index_dims.append(False) - return Range(ranges) + return Range(ranges, index_dims=index_dims) @staticmethod def ndslice_to_string(slice, tile_sizes=None): @@ -747,7 +821,14 @@ def __getitem__(self, key): return self.ranges.__getitem__(key) def __setitem__(self, key, value): - return self.ranges.__setitem__(key, value) + result = self.ranges.__setitem__(key, value) + # A replaced dimension is no longer the one that was parsed, so it is only still an index + # if the new range is degenerate. + for i in range(len(self.ranges)): + if self.index_dims[i]: + rb, re, _ = self.ranges[i] + self.index_dims[i] = bool(symbolic.equal_valued(0, re - rb)) + return result def __eq__(self, other): if not isinstance(other, Range): @@ -765,18 +846,25 @@ def compose(self, other): raise TypeError("Cannot compose ranges with non-subsets") new_subset = [] + # A composed dimension is an index exactly when the dimension it originates from is: a + # degenerate dimension of ``self`` keeps its own flag, a dimension consumed from ``other`` + # takes ``other``'s. + new_index_dims = [] + other_index_dims = other.index_dims if isinstance(other, Range) else [True] * other.dims() if self.data_dims() == other.dims(): # case 1: subsets may differ in dimensions, but data_dims correspond # to other dims -> all non-data dims are cut out idx = 0 - for (rb, re, rs), rt in zip(self.ranges, self.tile_sizes): + for dim, ((rb, re, rs), rt) in enumerate(zip(self.ranges, self.tile_sizes)): if re - rb == 0: new_subset.append((rb, re, rs, rt)) + new_index_dims.append(self.index_dims[dim]) else: if isinstance(other[idx], tuple): new_subset.append((rb + rs * other[idx][0], rb + rs * other[idx][1], rs * other[idx][2], rt)) else: new_subset.append(rb + rs * other[idx]) + new_index_dims.append(other_index_dims[idx]) idx += 1 elif self.dims() == other.dims(): # case 2: subsets have the same dimensions (but possibly different @@ -784,11 +872,13 @@ def compose(self, other): for idx, ((rb, re, rs), rt) in enumerate(zip(self.ranges, self.tile_sizes)): if re - rb == 0: new_subset.append((rb, re, rs, rt)) + new_index_dims.append(self.index_dims[idx]) else: if isinstance(other[idx], tuple): new_subset.append((rb + rs * other[idx][0], rb + rs * other[idx][1], rs * other[idx][2], rt)) else: new_subset.append(rb + rs * other[idx]) + new_index_dims.append(other_index_dims[idx]) elif (other.data_dims() == 0 and all([r == (0, 0, 1) if isinstance(other, Range) else r == 0 for r in other])): # NOTE: This is a special case where the other subset is the # (potentially multidimensional) index zero. @@ -798,6 +888,7 @@ def compose(self, other): new_subset.extend(self.ranges) else: new_subset.extend([rb for rb, _, _ in self.ranges]) + new_index_dims.extend(self.index_dims) else: raise ValueError("Dimension mismatch in composition: " "Subset composed must be either completely " @@ -805,7 +896,7 @@ def compose(self, other): "or be not stripped of latter at all.") if isinstance(other, Range): - return Range(new_subset) + return Range(new_subset, index_dims=new_index_dims) else: raise NotImplementedError @@ -837,11 +928,14 @@ def squeeze(self, ignore_indices: Optional[List[int]] = None, offset: bool = Tru pass squeezed_ranges = [self.ranges[i] for i in non_ones] squeezed_tsizes = [self.tile_sizes[i] for i in non_ones] + squeezed_index_dims = [self.index_dims[i] for i in non_ones] if not squeezed_ranges: squeezed_ranges = [(0, 0, 1)] squeezed_tsizes = [1] + squeezed_index_dims = [False] self.ranges = squeezed_ranges self.tile_sizes = squeezed_tsizes + self.index_dims = squeezed_index_dims if offset: self.offset(self, True, indices=offset_indices) return non_ones @@ -868,6 +962,7 @@ def unsqueeze(self, axes: Sequence[int]) -> List[int]: for axis in sorted(axes): self.ranges.insert(axis, (0, 0, 1)) self.tile_sizes.insert(axis, 1) + self.index_dims.insert(axis, False) if len(result) > 0 and result[-1] >= axis: result.append(result[-1] + 1) @@ -878,15 +973,19 @@ def unsqueeze(self, axes: Sequence[int]) -> List[int]: def pop(self, dimensions): new_ranges = [] new_tsizes = [] + new_index_dims = [] for i in range(len(self.ranges)): if i not in dimensions: new_ranges.append(self.ranges[i]) new_tsizes.append(self.tile_sizes[i]) + new_index_dims.append(self.index_dims[i]) if not new_ranges: new_ranges = [(symbolic.pystr_to_symbolic(0), symbolic.pystr_to_symbolic(0), symbolic.pystr_to_symbolic(1))] new_tsizes = [symbolic.pystr_to_symbolic(1)] + new_index_dims = [False] self.ranges = new_ranges self.tile_sizes = new_tsizes + self.index_dims = new_index_dims def string_list(self): return Range.ndslice_to_string_list(self.ranges, self.tile_sizes) @@ -1013,7 +1112,7 @@ def __init__(self, indices: Sequence[int | str | symbolic.SymbolicType]): raise TypeError("Expected collection of index expression: got SymExpr") indices = [symbolic.pystr_to_symbolic(i) for i in indices] - super().__init__([(idx, idx, 1) for idx in indices]) + super().__init__([(idx, idx, 1) for idx in indices], index_dims=[True] * len(indices)) @property def indices(self) -> List[symbolic.SymbolicType]: diff --git a/doitgen_repro.py b/doitgen_repro.py new file mode 100644 index 0000000000..01162a1dc8 --- /dev/null +++ b/doitgen_repro.py @@ -0,0 +1,106 @@ +# Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. +"""TEMPORARY repro driver -- remove before the PR is finalised. + +Runs npbench's polybench/doitgen through the two pipelines the benchmark harness uses (``strict`` += parse + simplify, and ``autoopt`` = simplify + auto_optimize) and checks the result against +numpy. Kernel and initializer are verbatim from npbench main: + + npbench/benchmarks/polybench/doitgen/doitgen_dace.py + npbench/benchmarks/polybench/doitgen/doitgen.py + npbench/benchmarks/polybench/doitgen/doitgen_numpy.py + +The ``(NQ, 1, NP)`` reshape is the point of the benchmark: it is what the in-tree test +``tests/npbench/polybench/doitgen_test.py`` stopped exercising in afd0efe48, which is why CI has +been green while npbench doitgen fails. + +Usage: python doitgen_repro.py [NR NQ NP] (defaults to 3 4 5) +""" +import copy +import sys +import traceback + +import numpy as np + +import dace as dc +import dace.dtypes as dtypes +import dace.transformation.auto.auto_optimize as opt + +NR, NQ, NP = (dc.symbol(s, dtype=dc.int64) for s in ('NR', 'NQ', 'NP')) + + +@dc.program +def kernel(A: dc.float64[NR, NQ, NP], C4: dc.float64[NP, NP]): + for r in range(NR): + A[r, :, :] = np.reshape(np.reshape(A[r], (NQ, 1, NP)) @ C4, (NQ, NP)) + + +def initialize(nr, nq, np_, datatype=np.float64): + A = np.fromfunction(lambda i, j, k: ((i * j + k) % np_) / np_, (nr, nq, np_), dtype=datatype) + C4 = np.fromfunction(lambda i, j: (i * j % np_) / np_, (np_, np_), dtype=datatype) + return A, C4 + + +def numpy_kernel(nr, nq, np_, A, C4): + A[:] = np.reshape(np.reshape(A, (nr, nq, 1, np_)) @ C4, (nr, nq, np_)) + + +def check(tag, got, ref): + if got.shape != ref.shape: + print(f"[{tag}] SHAPE MISMATCH got={got.shape} ref={ref.shape}") + return False + diff = np.abs(got - ref) + ok = np.allclose(got, ref, rtol=1e-12, atol=1e-12) + print(f"[{tag}] {'OK' if ok else 'MISMATCH'} max_abs={diff.max():.3e}") + if not ok: + bad = np.argwhere(~np.isclose(got, ref, rtol=1e-12, atol=1e-12)) + print(f"[{tag}] nbad={len(bad)}/{ref.size} first={bad[:3].tolist()}") + return ok + + +def run(tag, build, nr, nq, np_): + A, C4 = initialize(nr, nq, np_) + ref = A.copy() + numpy_kernel(nr, nq, np_, ref, C4) + try: + sdfg = build() + except Exception: + print(f"[{tag}] BUILD FAILED") + traceback.print_exc() + return False + got = A.copy() + try: + sdfg(A=got, C4=C4, NR=nr, NQ=nq, NP=np_) + except Exception: + print(f"[{tag}] RUN FAILED") + traceback.print_exc() + return False + return check(tag, got, ref) + + +def main(): + nr, nq, np_ = (int(x) for x in sys.argv[1:4]) if len(sys.argv) > 3 else (3, 4, 5) + print(f"=== doitgen NR={nr} NQ={nq} NP={np_} dace={dc.__file__}") + + base = kernel.to_sdfg(simplify=False) + + def strict(): + sdfg = copy.deepcopy(base) + sdfg._name = "strict" + sdfg.simplify() + return sdfg + + def autoopt(): + sdfg = copy.deepcopy(base) + sdfg._name = "autoopt" + sdfg.simplify() + opt.auto_optimize(sdfg, dtypes.DeviceType.CPU, symbols=dict(NR=nr, NQ=nq, NP=np_)) + return sdfg + + strict_ok = run("strict ", strict, nr, nq, np_) + autoopt_ok = run("autoopt", autoopt, nr, nq, np_) + print(f"=== RESULT strict={'PASS' if strict_ok else 'FAIL'} autoopt={'PASS' if autoopt_ok else 'FAIL'}") + return 0 if (strict_ok and autoopt_ok) else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/matmul_unit_dim_squeeze_test.py b/tests/matmul_unit_dim_squeeze_test.py deleted file mode 100644 index 6d4d65240a..0000000000 --- a/tests/matmul_unit_dim_squeeze_test.py +++ /dev/null @@ -1,111 +0,0 @@ -"""A ``subsets.Range`` entry of size 1 is ambiguous, and the MatMul library nodes read it two -incompatible ways. - -``Range`` stores both a rank-reducing index and a genuine extent-1 dimension as the same triple -``(0, 0, 1)``, so ``A[r, 0:N, 0:P]`` (dimension indexed away, rank 2) and ``V[0:N, 0, 0:P]`` -(full read of a ``(N, 1, P)`` view, rank 3) are indistinguishable at the subset level. -``Range.squeeze()`` drops both. - -``np.reshape(x, (NQ, 1, NP))`` builds exactly the second kind: a 3D view whose middle dimension is -part of the shape. ``SpecializeMatMul`` dispatched on the squeezed sizes and saw a matrix, so it -picked ``Gemm``, whose ``validate`` re-read the unsqueezed subset and raised -"matrix-matrix product only supported on matrices". npbench's doitgen hits this on every size. - -The descriptor is the authority: a subset entry of size 1 is rank-reducing only where the -descriptor's own extent is larger. Under that rule the reshape stays 3D and dispatches to -``BatchedMatMul``, while a real index into a larger dimension still squeezes to ``Gemm``. -""" -import numpy as np -import pytest - -import dace -from dace.libraries.blas.nodes.batched_matmul import BatchedMatMul -from dace.libraries.blas.nodes.gemm import Gemm -from dace.libraries.blas.nodes.matmul import MatMul -from dace.transformation.auto.auto_optimize import auto_optimize - -NR, NQ, NP = (dace.symbol(s, dtype=dace.int64) for s in ('NR', 'NQ', 'NP')) - - -@dace.program -def doitgen_reshape(A: dace.float64[NR, NQ, NP], C4: dace.float64[NP, NP]): - # npbench polybench/doitgen, verbatim: the (NQ, 1, NP) reshape is the point of the benchmark. - for r in range(NR): - A[r, :, :] = np.reshape(np.reshape(A[r], (NQ, 1, NP)) @ C4, (NQ, NP)) - - -@dace.program -def indexed_slice_matmul(A: dace.float64[NR, NQ, NP], C4: dace.float64[NP, NP]): - # A[r] indexes dimension 0 away: a genuine 2D operand that must still reach Gemm. - for r in range(NR): - A[r, :, :] = A[r] @ C4 - - -def reference(A, C4): - NR, NQ, NP = A.shape - return np.reshape(np.reshape(A, (NR, NQ, 1, NP)) @ C4, (NR, NQ, NP)) - - -def initialize(nr, nq, np_): - A = np.fromfunction(lambda i, j, k: ((i * j + k) % np_) / np_, (nr, nq, np_), dtype=np.float64) - C4 = np.fromfunction(lambda i, j: (i * j % np_) / np_, (np_, np_), dtype=np.float64) - return A, C4 - - -def count_nodes(sdfg, nodetype): - return sum(1 for node, _ in sdfg.all_nodes_recursive() if isinstance(node, nodetype)) - - -def specialize_matmuls(sdfg): - """Expand the MatMul meta-nodes one level, leaving the chosen specialization in the graph.""" - for node, state in list(sdfg.all_nodes_recursive()): - if type(node) is MatMul: - node.expand(state) - - -@pytest.mark.parametrize('optimize', [False, True]) -@pytest.mark.parametrize('sizes', [(3, 4, 5), (1, 1, 1), (8, 10, 12), (5, 1, 7)]) -def test_reshape_unit_dim_matmul(optimize, sizes): - nr, nq, np_ = sizes - A, C4 = initialize(nr, nq, np_) - ref = reference(A, C4) - - sdfg = doitgen_reshape.to_sdfg(simplify=False) - sdfg.simplify() - if optimize: - auto_optimize(sdfg, dace.dtypes.DeviceType.CPU, symbols=dict(NR=nr, NQ=nq, NP=np_)) - - sdfg(A=A, C4=C4, NR=nr, NQ=nq, NP=np_) - assert np.allclose(A, ref) - - -def test_reshape_unit_dim_dispatches_to_batched(): - """The reshaped operand keeps its unit dimension, so the matmul stays batched.""" - sdfg = doitgen_reshape.to_sdfg(simplify=False) - sdfg.simplify() - specialize_matmuls(sdfg) - assert count_nodes(sdfg, BatchedMatMul) == 1 - assert count_nodes(sdfg, Gemm) == 0 - - -def test_indexed_dim_still_squeezes_to_gemm(): - """An index into a larger dimension is rank-reducing, so the operand is a plain matrix.""" - nr, nq, np_ = 3, 4, 5 - A, C4 = initialize(nr, nq, np_) - ref = reference(A, C4) - - sdfg = indexed_slice_matmul.to_sdfg(simplify=False) - sdfg.simplify() - specialize_matmuls(sdfg) - assert count_nodes(sdfg, Gemm) == 1 - assert count_nodes(sdfg, BatchedMatMul) == 0 - - sdfg(A=A, C4=C4, NR=nr, NQ=nq, NP=np_) - assert np.allclose(A, ref) - - -if __name__ == '__main__': - for opt in (False, True): - test_reshape_unit_dim_matmul(opt, (3, 4, 5)) - test_reshape_unit_dim_dispatches_to_batched() - test_indexed_dim_still_squeezes_to_gemm() diff --git a/tests/sdfg/subset_index_dims_test.py b/tests/sdfg/subset_index_dims_test.py new file mode 100644 index 0000000000..35abf8d188 --- /dev/null +++ b/tests/sdfg/subset_index_dims_test.py @@ -0,0 +1,192 @@ +# Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. +"""``Range`` records which dimensions were written as an integer index. + +``(begin, end, step)`` cannot express the difference between ``a[0]`` and ``a[0:1]`` -- both are +``(0, 0, 1)`` -- but numpy distinguishes them: an integer index removes its dimension from the +result rank, a slice never does, not even at extent 1. Consumers that need the rank of an access +(``x @ y`` dispatching to a matrix or a vector product, for instance) previously had to guess by +squeezing every extent-1 dimension, which is a different operation (``np.squeeze``) and is wrong +whenever a sliced dimension happens to have extent 1. +""" +import copy + +import numpy as np + +import dace +from dace import subsets + + +def test_index_reduces_rank_slice_does_not(): + indexed = subsets.Range.from_string('0:NQ, 0, 0:NP') + sliced = subsets.Range.from_string('0:NQ, 0:1, 0:NP') + + # Identical extents, so neither size() nor equality can tell them apart ... + assert indexed.size() == sliced.size() + assert indexed == sliced + # ... but the rank differs, which is what numpy cares about. + assert indexed.rank() == 2 + assert sliced.rank() == 3 + assert indexed.rank_dims() == [0, 2] + assert sliced.rank_dims() == [0, 1, 2] + + +def test_unit_extent_slice_keeps_its_rank(): + """The case squeezing cannot get right: a sliced dimension whose extent happens to be 1.""" + subset = subsets.Range.from_string('0:1, 0, 0:7') + assert subset.size() == [1, 1, 7] + assert subset.rank() == 2 # only the middle dimension was indexed away + assert subset.rank_dims() == [0, 2] + + squeezed = copy.deepcopy(subset) + squeezed.squeeze() + assert squeezed.size() == [7] # np.squeeze drops both, which is a different question + + +def test_squeeze_is_still_numpy_squeeze(): + """``squeeze`` keeps meaning ``np.squeeze``: drop every extent-1 dimension.""" + subset = subsets.Range.from_string('0:4, 0:1, 0:7') + assert subset.rank() == 3 + subset.squeeze() + assert subset.size() == [4, 7] + + +def test_from_array_preserves_rank(): + """Reading all of an array is never a rank reduction, even for its unit dimensions.""" + sdfg = dace.SDFG('from_array_rank') + sdfg.add_array('V', [6, 1, 5], dace.float64) + subset = subsets.Range.from_array(sdfg.arrays['V']) + assert subset.rank() == 3 + assert not any(subset.index_dims) + + +def test_from_indices_is_all_indices(): + subset = subsets.Range.from_indices([2, 3]) + assert subset.rank() == 0 + assert all(subset.index_dims) + + +def test_string_form_is_unchanged(): + """The rendering stays exactly as it was. + + Subset strings are re-parsed as symbolic expressions in places that reject slice syntax (see + ``dace.transformation.passes.scalar_to_symbol``), so a degenerate range must keep printing as + ``i``. The flags therefore live in memory and in JSON, not in the string form. + """ + assert str(subsets.Range.from_string('0:NQ, 0, 0:NP')) == '0:NQ, 0, 0:NP' + assert str(subsets.Range.from_string('0:NQ, 0:1, 0:NP')) == '0:NQ, 0, 0:NP' + assert str(subsets.Range.from_string('0, 0:10')) == '0, 0:10' + + +def test_frontend_records_mixed_subscript_ranks(): + """``A[0, 0:N]`` is rank 1 in numpy, so the frontend has to record which dimension was indexed.""" + from dace.frontend.python import memlet_parser + + subset = memlet_parser._ndslice_to_subset([0, (0, 9, 1)]) + assert subset.index_dims == [True, False] + assert subset.rank() == 1 + + both_slices = memlet_parser._ndslice_to_subset([(0, 0, 1), (0, 9, 1)]) + assert both_slices.index_dims == [False, False] + assert both_slices.rank() == 2 + + +def test_json_marks_only_index_dims(): + subset = subsets.Range.from_string('0:NQ, 0, 0:NP') + entries = subset.to_json()['ranges'] + assert [entry.get('indexed', False) for entry in entries] == [False, True, False] + + +def test_sdfg_save_load_preserves_rank(tmp_path): + sdfg = dace.SDFG('subset_rank_roundtrip') + sdfg.add_array('A', [4, 1, 8], dace.float64) + state = sdfg.add_state('s', is_start_block=True) + read, write = state.add_read('A'), state.add_write('A') + # A[2, 0:1, 0:8]: dimension 0 is indexed away, dimension 1 is a unit slice that stays. + state.add_nedge(read, write, dace.Memlet('A[2, 0:1, 0:8]')) + assert state.edges()[0].data.subset.rank() == 2 + + path = tmp_path / 'roundtrip.sdfg' + sdfg.save(path) + reloaded = dace.SDFG.from_file(path) + subset = list(reloaded.states())[0].edges()[0].data.subset + assert subset.index_dims == [True, False, False] + assert subset.rank() == 2 + + +def test_flags_survive_structural_operations(): + subset = subsets.Range.from_string('0:4, 0, 0:7, 2') + assert subset.index_dims == [False, True, False, True] + + assert copy.deepcopy(subset).index_dims == subset.index_dims + assert (subset + subsets.Range.from_string('0')).index_dims == [False, True, False, True, True] + + reordered = copy.deepcopy(subset) + reordered.reorder([1, 0, 3, 2]) + assert reordered.index_dims == [True, False, True, False] + + popped = copy.deepcopy(subset) + popped.pop([0]) + assert popped.index_dims == [True, False, True] + + unsqueezed = subsets.Range.from_string('0:10') + unsqueezed.unsqueeze([0]) + assert unsqueezed.index_dims == [False, False] # unsqueeze inserts slices, not indices + + +def test_derived_subsets_keep_their_flags(): + """Every operation that builds a new Range from an existing one must carry the flags over.""" + subset = subsets.Range.from_string('0:4, 0, 0:7') + + assert subset.offset_new(None, False).index_dims == [False, True, False] + assert subset.offset_new([1, 0, 2], True).index_dims == [False, True, False] + + # compose: a degenerate dimension keeps its own flag, a consumed one takes the other subset's + composed = subset.compose(subsets.Range.from_string('0:4, 0:7')) + assert composed.index_dims == [False, True, False] + assert composed.rank() == 2 + + +def test_widening_a_dimension_clears_its_index_flag(): + """An index is degenerate by definition, so a widened dimension cannot stay flagged.""" + subset = subsets.Range.from_string('0:4, 0, 0:7') + assert subset.rank() == 2 + + subset[1] = (0, 3, 1) + assert subset.rank() == 3 + assert not subset.is_index_dim(1) + + # Direct mutation of `ranges` bypasses __setitem__ entirely, so the flag is cross-checked + # against the range rather than trusted. + other = subsets.Range.from_string('0:4, 0, 0:7') + other.ranges[1] = (0, 5, 1) + assert other.rank() == 3 + + +def test_rank_reducing_slice_still_lowers_correctly(): + """End-to-end guard: an integer index into a larger dimension still yields a 1D result.""" + R, C = (dace.symbol(s, dtype=dace.int64) for s in ('R', 'C')) + + @dace.program + def take_row(A: dace.float64[R, C], out: dace.float64[C]): + out[:] = A[2, :] + + rng = np.random.default_rng(0) + A = rng.random((4, 8)) + out = np.zeros(8) + take_row(A, out, R=4, C=8) + assert np.allclose(out, A[2]) + + +if __name__ == '__main__': + test_index_reduces_rank_slice_does_not() + test_unit_extent_slice_keeps_its_rank() + test_squeeze_is_still_numpy_squeeze() + test_from_array_preserves_rank() + test_from_indices_is_all_indices() + test_string_form_is_unchanged() + test_frontend_records_mixed_subscript_ranks() + test_json_marks_only_index_dims() + test_flags_survive_structural_operations() + test_derived_subsets_keep_their_flags() + test_widening_a_dimension_clears_its_index_flag() + test_rank_reducing_slice_still_lowers_correctly() From be83a8cf5ee14e75d19e7375c9e09d0b73e84aa6 Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Mon, 20 Jul 2026 21:08:27 +0200 Subject: [PATCH 05/17] Validate and generate GEMM against the operand view it was dispatched on `SpecializeMatMul` selects GEMM by matching on the squeezed operand sizes, so `np.reshape(x, (NQ, 1, NP)) @ C4` is routed there as an `(NQ, NP) @ (NP, NP)` product. Since afd0efe48, `Gemm.validate`, `ExpandGemmPure` and `_get_codegen_gemm_opts` re-read the raw subset instead, see rank 3, and reject the operand the dispatcher had just accepted with "matrix-matrix product only supported on matrices". npbench's doitgen fails at every size, on both the simplified and the auto-optimized pipeline. The fix is for GEMM to read the same tuple entries the dispatcher matched on. This does not redefine what `_get_matmul_operands` reports -- gemv, ger, dot and the vendor and FPGA expansions continue to read indices 4 and 5 exactly as before -- so it is confined to the GEMM path. Collapsing the unit dimension is exact, not a convenience: it is the row count of an `NQ`-long contiguous batch, so the product really is a single GEMM. Keeping it on GEMM also preserves alpha, beta and a summation WCR on `_c`, none of which BatchedMatMul honours; a unit-batch product with alpha now computes the right answer where it previously raised. afd0efe48 also weakened the in-tree benchmark, replacing the `(NQ, 1, NP)` reshape with `(NQ, NP)` so it stopped exercising the broken path, which is why CI stayed green. That line is restored, and it fails without this change. --- dace/libraries/blas/nodes/gemm.py | 30 +++- dace/libraries/blas/nodes/matmul.py | 5 +- tests/library/matmul_unit_dim_squeeze_test.py | 139 ++++++++++++++++++ tests/npbench/polybench/doitgen_test.py | 2 +- 4 files changed, 167 insertions(+), 9 deletions(-) create mode 100644 tests/library/matmul_unit_dim_squeeze_test.py diff --git a/dace/libraries/blas/nodes/gemm.py b/dace/libraries/blas/nodes/gemm.py index 90f4c6d2b1..f295990c5e 100644 --- a/dace/libraries/blas/nodes/gemm.py +++ b/dace/libraries/blas/nodes/gemm.py @@ -9,6 +9,24 @@ from dace.transformation.transformation import ExpandTransformation from dace.libraries.blas.blas_helpers import to_blastype, check_access, dtype_to_cudadatatype, to_cublas_computetype from dace.libraries.blas.nodes.matmul import (_get_matmul_operands, _get_codegen_gemm_opts) + + +def _squeezed_size(memlet): + """ + Returns the operand size as the matrix ``SpecializeMatMul`` saw when it chose GEMM. + + The dispatcher matches on the squeezed operand sizes, so validation and code generation have to + read the same view. Reading the raw subset instead rejects an operand the dispatcher just + accepted, which is how a ``(NQ, 1, NP)`` reshape reaches GEMM and is then refused. + + :param memlet: The memlet accessing the operand. + :return: The size of the memlet subset with its size-1 dimensions removed. + """ + subset = dc(memlet.subset) + subset.squeeze() + return subset.size() + + from .. import environments import numpy as np import warnings @@ -46,7 +64,7 @@ class ExpandGemmPure(ExpandTransformation): def make_sdfg(node, parent_state, parent_sdfg): sdfg = dace.SDFG(node.label + "_sdfg") - ((edge_a, outer_array_a, shape_a, strides_a, _, _), (edge_b, outer_array_b, shape_b, strides_b, _, _), + ((edge_a, outer_array_a, _, _, shape_a, strides_a), (edge_b, outer_array_b, _, _, shape_b, strides_b), cdata) = _get_matmul_operands(node, parent_state, parent_sdfg) dtype_a = outer_array_a.dtype.type @@ -79,7 +97,7 @@ def make_sdfg(node, parent_state, parent_sdfg): _, array_a = sdfg.add_array("_a", shape_a, dtype_a, strides=strides_a, storage=outer_array_a.storage) _, array_b = sdfg.add_array("_b", shape_b, dtype_b, strides=strides_b, storage=outer_array_b.storage) - _, array_c = sdfg.add_array("_c", shape_c, dtype_c, strides=cdata[-3], storage=cdata[1].storage) + _, array_c = sdfg.add_array("_c", shape_c, dtype_c, strides=cdata[-1], storage=cdata[1].storage) if equal_valued(1, node.alpha): mul_program = "__out = __a * __b" @@ -551,11 +569,11 @@ def validate(self, sdfg, state): size2 = None for _, _, _, dst_conn, memlet in state.in_edges(self): if dst_conn == '_a': - size0 = memlet.subset.size() + size0 = _squeezed_size(memlet) if dst_conn == '_b': - size1 = memlet.subset.size() + size1 = _squeezed_size(memlet) if dst_conn == '_c': - size2 = memlet.subset.size() + size2 = _squeezed_size(memlet) if self.transA: size0 = list(reversed(size0)) @@ -575,7 +593,7 @@ def validate(self, sdfg, state): UserWarning) elif not res: raise ValueError("Inputs to matrix-matrix product must agree in the k-dimension") - size3 = out_memlet.subset.size() + size3 = _squeezed_size(out_memlet) if size2 is not None: res = [equal(s0, s1) for s0, s1 in zip(size2, size3)] fail = any([r is False for r in res]) diff --git a/dace/libraries/blas/nodes/matmul.py b/dace/libraries/blas/nodes/matmul.py index f5f308b2c8..d3520bf499 100644 --- a/dace/libraries/blas/nodes/matmul.py +++ b/dace/libraries/blas/nodes/matmul.py @@ -137,8 +137,9 @@ def _get_codegen_gemm_opts(node, state, sdfg, adesc, bdesc, cdesc, alpha, beta, from dace.codegen.common import sym2cpp from dace.libraries.blas.blas_helpers import get_gemm_opts - (_, _, ashape, astride, _, _), (_, _, bshape, bstride, _, _), (_, _, cshape, cstride, _, - _) = _get_matmul_operands(node, state, sdfg) + # GEMM operands are matrices: use the squeezed view, matching SpecializeMatMul's dispatch. + (_, _, _, _, ashape, astride), (_, _, _, _, bshape, bstride), (_, _, _, _, cshape, + cstride) = _get_matmul_operands(node, state, sdfg) if getattr(node, 'transA', False): ashape = list(reversed(ashape)) diff --git a/tests/library/matmul_unit_dim_squeeze_test.py b/tests/library/matmul_unit_dim_squeeze_test.py new file mode 100644 index 0000000000..4c9a57ca87 --- /dev/null +++ b/tests/library/matmul_unit_dim_squeeze_test.py @@ -0,0 +1,139 @@ +# Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. +"""``SpecializeMatMul`` and ``Gemm`` have to read the same view of an operand. + +The dispatcher matches on the squeezed operand sizes, so ``np.reshape(x, (NQ, 1, NP)) @ C4`` is +routed to ``Gemm`` as an ``(NQ, NP) @ (NP, NP)`` product. ``Gemm.validate`` and the GEMM code +generator used to re-read the raw subset instead, see rank 3, and reject the operand the dispatcher +had just accepted -- "matrix-matrix product only supported on matrices". npbench's doitgen hits this +at every size. + +Collapsing the unit dimension is exact rather than convenient: it is the row count of an +``NQ``-long contiguous batch, so the product is genuinely one GEMM. Keeping it on ``Gemm`` also +keeps ``alpha``, ``beta`` and a summation WCR, none of which ``BatchedMatMul`` honours. +""" +import numpy as np +import pytest + +import dace +from dace.libraries.blas.nodes.batched_matmul import BatchedMatMul +from dace.libraries.blas.nodes.gemm import Gemm +from dace.libraries.blas.nodes.matmul import MatMul +from dace.transformation.auto.auto_optimize import auto_optimize + +NR, NQ, NP = (dace.symbol(s, dtype=dace.int64) for s in ('NR', 'NQ', 'NP')) + + +@dace.program +def doitgen_reshape(A: dace.float64[NR, NQ, NP], C4: dace.float64[NP, NP]): + # npbench polybench/doitgen, verbatim: the (NQ, 1, NP) reshape is the point of the benchmark. + for r in range(NR): + A[r, :, :] = np.reshape(np.reshape(A[r], (NQ, 1, NP)) @ C4, (NQ, NP)) + + +@dace.program +def indexed_slice_matmul(A: dace.float64[NR, NQ, NP], C4: dace.float64[NP, NP]): + # A[r] indexes dimension 0 away: a genuine 2D operand that must still reach Gemm. + for r in range(NR): + A[r, :, :] = A[r] @ C4 + + +def reference(A, C4): + nr, nq, np_ = A.shape + return np.reshape(np.reshape(A, (nr, nq, 1, np_)) @ C4, (nr, nq, np_)) + + +def initialize(nr, nq, np_, seed=0): + # Random rather than polybench's ((i*j+k) % NP)/NP, which degenerates to all-zero when NP == 1 + # and would make the numeric assertions vacuous for the all-unit-extent shape. + rng = np.random.default_rng(seed) + return rng.random((nr, nq, np_)), rng.random((np_, np_)) + + +def count_nodes(sdfg, nodetype): + return sum(1 for node, _ in sdfg.all_nodes_recursive() if isinstance(node, nodetype)) + + +def specialize_matmuls(sdfg): + """Expand the MatMul meta-nodes one level, leaving the chosen specialization in the graph.""" + for node, state in list(sdfg.all_nodes_recursive()): + if type(node) is MatMul: + node.expand(state) + + +# NQ == 1 collapses a second dimension when squeezed, which is where a "pick whichever view looks +# two-dimensional" heuristic breaks; both such shapes are covered. +SIZES = [(3, 4, 5), (1, 1, 1), (8, 10, 12), (5, 1, 7)] + + +@pytest.mark.parametrize('optimize', [False, True]) +@pytest.mark.parametrize('sizes', SIZES) +def test_reshape_unit_dim_matmul(optimize, sizes): + nr, nq, np_ = sizes + A, C4 = initialize(nr, nq, np_) + ref = reference(A, C4) + assert np.abs(ref).max() > 0.0 # guard against a vacuous comparison + + sdfg = doitgen_reshape.to_sdfg(simplify=False) + sdfg.simplify() + if optimize: + auto_optimize(sdfg, dace.dtypes.DeviceType.CPU, symbols=dict(NR=nr, NQ=nq, NP=np_)) + + sdfg(A=A, C4=C4, NR=nr, NQ=nq, NP=np_) + assert np.allclose(A, ref) + + +def test_reshape_unit_dim_stays_one_gemm(): + """The collapse is exact, so the product must not fan out into NQ batched calls.""" + sdfg = doitgen_reshape.to_sdfg(simplify=False) + sdfg.simplify() + specialize_matmuls(sdfg) + assert count_nodes(sdfg, Gemm) == 1 + assert count_nodes(sdfg, BatchedMatMul) == 0 + + +def test_indexed_dim_still_reaches_gemm(): + """An index into a larger dimension is rank-reducing, so the operand is a plain matrix.""" + nr, nq, np_ = 3, 4, 5 + A, C4 = initialize(nr, nq, np_) + ref = reference(A, C4) + + sdfg = indexed_slice_matmul.to_sdfg(simplify=False) + sdfg.simplify() + specialize_matmuls(sdfg) + assert count_nodes(sdfg, Gemm) == 1 + assert count_nodes(sdfg, BatchedMatMul) == 0 + + sdfg(A=A, C4=C4, NR=nr, NQ=nq, NP=np_) + assert np.allclose(A, ref) + + +def test_unit_batch_keeps_alpha(): + """Routing a unit-batch product away from Gemm would silently drop alpha, beta and WCR.""" + from dace import memlet as mm + + sdfg = dace.SDFG('unit_batch_alpha') + b, m, k, n = 1, 8, 6, 5 + sdfg.add_array('A', [b, m, k], dace.float64) + sdfg.add_array('B', [k, n], dace.float64) + sdfg.add_array('C', [b, m, n], dace.float64) + state = sdfg.add_state('s', is_start_block=True) + node = MatMul('mm', alpha=2.0) + state.add_node(node) + state.add_edge(state.add_read('A'), None, node, '_a', mm.Memlet('A[0:1, 0:8, 0:6]')) + state.add_edge(state.add_read('B'), None, node, '_b', mm.Memlet('B[0:6, 0:5]')) + state.add_edge(node, '_c', state.add_write('C'), None, mm.Memlet('C[0:1, 0:8, 0:5]')) + sdfg.expand_library_nodes() + + rng = np.random.default_rng(0) + a, bmat, c = rng.random((b, m, k)), rng.random((k, n)), np.zeros((b, m, n)) + sdfg(A=a, B=bmat, C=c) + assert np.allclose(c[0], 2.0 * (a[0] @ bmat)) + + +if __name__ == '__main__': + for opt in (False, True): + for sz in SIZES: + test_reshape_unit_dim_matmul(opt, sz) + test_reshape_unit_dim_stays_one_gemm() + test_indexed_dim_still_reaches_gemm() + test_unit_batch_keeps_alpha() diff --git a/tests/npbench/polybench/doitgen_test.py b/tests/npbench/polybench/doitgen_test.py index a97931e4d0..a09b05ca61 100644 --- a/tests/npbench/polybench/doitgen_test.py +++ b/tests/npbench/polybench/doitgen_test.py @@ -27,7 +27,7 @@ def doitgen_kernel(A: dc.float64[NR, NQ, NP], C4: dc.float64[NP, NP]): # Ideal - not working because Matmul with dim > 3 unsupported # A[:] = np.reshape(np.reshape(A, (NR, NQ, 1, NP)) @ C4, (NR, NQ, NP)) for r in range(NR): - A[r, :, :] = np.reshape(np.reshape(A[r], (NQ, NP)) @ C4, (NQ, NP)) + A[r, :, :] = np.reshape(np.reshape(A[r], (NQ, 1, NP)) @ C4, (NQ, NP)) def initialize(NR, NQ, NP, datatype=np.float64): From c62c428da142d1beb2bc177bcf974ba51b43c7fd Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Mon, 20 Jul 2026 22:30:56 +0200 Subject: [PATCH 06/17] Read a size scalar into a symbol without deleting the descriptor A size computed in the program (`nt = Nt + 1` then `np.empty(nt)`) is a scalar data descriptor, but an array extent has to be a symbol. The previous approach ran a whole-SDFG ScalarToSymbolPromotion mid-parse, which deleted the scalar. Any later read or reassignment of the size then hit a hard KeyError, and the pass could disturb unrelated scalars in the half-built SDFG. Reuse the frontend's own promotion instead: `promote_scalar_to_symbol` mints a `__sym_` symbol assigned from the scalar on an interstate edge and leaves the descriptor in place, exactly as `_promote` already does for subscripts. The shape is rewritten by substituting the symbols in, rather than mutating the SDFG. The promoted symbol is registered in the visitor's globals so nested scopes can resolve it as a free symbol of their scope arrays. The test now executes the programs and checks numerics for the reuse and reassign cases, which were the KeyError crashes, instead of asserting that the descriptor was destroyed. --- dace/frontend/python/newast.py | 53 +++++++----- .../replacements/array_creation_dace.py | 53 ++++++------ tests/size_scalar_shape_promotion_test.py | 80 ++++++++++++++++--- 3 files changed, 131 insertions(+), 55 deletions(-) diff --git a/dace/frontend/python/newast.py b/dace/frontend/python/newast.py index 30fb537247..4c4788f7c6 100644 --- a/dace/frontend/python/newast.py +++ b/dace/frontend/python/newast.py @@ -5373,6 +5373,39 @@ def range_is_index(range: subsets.Range) -> bool: wcr=expr.wcr)) return tmp + def promote_scalar_to_symbol(self, scalar: str, key: Optional[str] = None) -> symbolic.symbol: + """ + Reads a scalar into a symbol so its value can be used where only symbols are allowed. + + A fresh ``__sym_`` symbol is assigned from the scalar on an interstate edge. The + scalar's data descriptor is deliberately left in place: it may still be read or reassigned + later in the program, and the visitor's own bookkeeping still refers to it by name. + + :param scalar: Name of the scalar data descriptor to read. + :param key: Cache key for the promotion, defaulting to the scalar name. Repeated promotions + of the same expression reuse the symbol instead of minting a new one. + :return: The symbol carrying the scalar's value. + """ + key = key if key is not None else scalar + desc = self.sdfg.arrays[scalar] + sym = self.indirections.get(key) + if sym is None: + sym = dace.symbol(f'__sym_{scalar}', dtype=desc.dtype) + self.indirections[key] = sym + try: + self.sdfg.add_symbol(f'__sym_{scalar}', desc.dtype) + except FileExistsError: + # By design this may re-add an existing symbol; the exception is benign. + pass + # A promoted size can end up in an array shape, and nested scopes resolve the free + # symbols of their scope arrays through `globals`, so the symbol has to be visible + # there as well as on the SDFG. + self.globals[str(sym)] = sym + state = self._add_state(f'promote_{scalar}_to_{str(sym)}') + edge = state.parent_graph.in_edges(state)[0] + edge.data.assignments = {str(sym): scalar} + return sym + def _parse_subscript_slice(self, s: ast.AST, multidim: bool = False) -> Union[Any, Tuple[Union[Any, str, symbolic.symbol]]]: @@ -5382,29 +5415,13 @@ def _parse_subscript_slice(self, def _promote(node: ast.AST) -> Union[Any, str, symbolic.symbol]: node_str = astutils.unparse(node) - sym = None - if node_str in self.indirections: - sym = self.indirections[node_str] if isinstance(node, str): scalar = node_str else: scalar = self.visit(node) if isinstance(scalar, str) and scalar in self.sdfg.arrays: - desc = self.sdfg.arrays[scalar] - if isinstance(desc, data.Scalar): - if not sym: - sym = dace.symbol(f'__sym_{scalar}', dtype=desc.dtype) - self.indirections[node_str] = sym - try: - self.sdfg.add_symbol(f'__sym_{scalar}', desc.dtype) - except FileExistsError: - # NOTE: By design, it is possible to try here to add an already existing symbol even if - # `not sym` returns True. This exception is benign. - pass - state = self._add_state(f'promote_{scalar}_to_{str(sym)}') - edge = state.parent_graph.in_edges(state)[0] - edge.data.assignments = {str(sym): scalar} - return sym + if isinstance(self.sdfg.arrays[scalar], data.Scalar): + return self.promote_scalar_to_symbol(scalar, key=node_str) return scalar if isinstance(s, (Number, bool, numpy.bool_, sympy.Basic)): diff --git a/dace/frontend/python/replacements/array_creation_dace.py b/dace/frontend/python/replacements/array_creation_dace.py index 4806e21961..0476798189 100644 --- a/dace/frontend/python/replacements/array_creation_dace.py +++ b/dace/frontend/python/replacements/array_creation_dace.py @@ -18,30 +18,35 @@ from dace import symbolic -def promote_size_scalars_in_shape(sdfg: SDFG, shape: Shape) -> None: - """Promote a size expression that was materialized as a size-1 descriptor and then used as an symbol - SHAPE (``np.empty(Nt_plus_1)`` -> ``Nt_plus_1`` has to be a symbol) into an SDFG symbol.""" - names = set() - for extent in shape: - if isinstance(extent, Integral): +def promote_size_scalars_in_shape(pv: ProgramVisitor, sdfg: SDFG, shape: Shape) -> Shape: + """ + Rewrites a shape so that scalars used as extents are read through symbols. + + An extent must be a symbol, but a size computed in the program (``nt = Nt + 1`` then + ``np.empty(nt)``) is a scalar data descriptor. Each such scalar is read into a symbol on an + interstate edge and substituted into the shape. The descriptor itself is left alone, so the + program may keep reading or reassigning the size afterwards. + + :param pv: The program visitor, which owns symbol promotion and the state machine. + :param sdfg: The SDFG being built. + :param shape: The requested shape. + :return: The shape with scalar extents replaced by symbols. + """ + resolved = [symbolic.pystr_to_symbolic(e) if isinstance(e, str) else e for e in shape] + replacements = {} + for extent in resolved: + if not isinstance(extent, sympy.Basic): continue - expr = symbolic.pystr_to_symbolic(extent) if isinstance(extent, str) else extent - if isinstance(expr, sympy.Basic): - names |= {str(s) for s in expr.free_symbols} - targets = {n for n in names if n in sdfg.arrays and n not in sdfg.symbols and sdfg.arrays[n].total_size == 1} - if not targets: - return - from dace.transformation.passes.scalar_to_symbol import ScalarToSymbolPromotion, find_promotable_scalars - promotable = find_promotable_scalars(sdfg, transients_only=False) - not_promotable = targets - promotable - if not_promotable: - raise DaceSyntaxError( - None, None, f'Cannot use {sorted(not_promotable)} as an array shape: the like-named size ' - f'descriptor is not promotable to a symbol.') - promo = ScalarToSymbolPromotion() - promo.transients_only = False - promo.ignore = promotable - targets # promote ONLY the size scalars used in this shape - promo.apply_pass(sdfg, {}) + for sym in extent.free_symbols: + name = str(sym) + if name in replacements or name in sdfg.symbols or name not in sdfg.arrays: + continue + desc = sdfg.arrays[name] + if isinstance(desc, data.Scalar) or desc.total_size == 1: + replacements[sym] = pv.promote_scalar_to_symbol(name) + if not replacements: + return shape + return [e.subs(replacements) if isinstance(e, sympy.Basic) else e for e in resolved] @oprepo.replaces('dace.define_local') @@ -61,7 +66,7 @@ def _define_local_ex(pv: ProgramVisitor, if not isinstance(strides, (list, tuple)): strides = [strides] strides = [int(s) if isinstance(s, Integral) else s for s in strides] - promote_size_scalars_in_shape(sdfg, shape) + shape = promote_size_scalars_in_shape(pv, sdfg, shape) name = pv.get_target_name() name, _ = sdfg.add_transient(name, shape, diff --git a/tests/size_scalar_shape_promotion_test.py b/tests/size_scalar_shape_promotion_test.py index a5a7602386..e66ebbf7d7 100644 --- a/tests/size_scalar_shape_promotion_test.py +++ b/tests/size_scalar_shape_promotion_test.py @@ -1,27 +1,81 @@ -"""A size materialized as a descriptor (``np.empty(Nt+1)`` -> ``Nt_plus_1``) then used as a shape must be -promoted to a symbol, not collide with ``add_symbol`` (which raised ``FileExistsError``).""" -import numpy as np +# Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. +"""A size computed in the program can be used as an array shape. +``nt = Nt + 1`` materializes ``nt`` as a scalar data descriptor, but an array extent has to be a +symbol, and minting a symbol of the same name collided with the descriptor (``FileExistsError``). +The size is now read into a ``__sym_`` symbol on an interstate edge and substituted into the shape, +leaving the descriptor in place so the program can keep reading or reassigning the size afterwards. +""" +import numpy as np import dace -N = dace.symbol("N") +N = dace.symbol('N') @dace.program -def size_scalar_shape_prog(a: dace.float64[N], Nt: dace.int64): +def size_from_empty(a: dace.float64[N], Nt: dace.int64, out: dace.float64[N]): b = np.empty(Nt + 1, dace.float64) for i in range(N): b[i] = a[i] * 2.0 - return b + for i in range(N): + out[i] = b[i] + + +@dace.program +def size_read_after_use(a: dace.float64[N], Nt: dace.int64, out: dace.float64[N]): + m = Nt + 1 + b = np.empty(m, dace.float64) + b[0] = 1.0 + for i in range(N): + out[i] = a[i] + m # the size descriptor must survive its use as a shape + + +@dace.program +def size_reassigned_after_use(a: dace.float64[N], Nt: dace.int64, out: dace.float64[N]): + m = Nt + 1 + b = np.empty(m, dace.float64) + b[0] = 1.0 + m = 99 + for i in range(N): + out[i] = a[i] + m + + +def test_scalar_size_as_shape(): + n, nt = 5, 7 + a = np.arange(n, dtype=np.float64) + out = np.zeros(n) + size_from_empty(a, np.int64(nt), out, N=n) + assert np.allclose(out, a * 2.0) + + +def test_size_descriptor_survives_its_use_as_a_shape(): + """Promotion must not delete the scalar: the program still reads it afterwards.""" + n, nt = 5, 7 + a = np.arange(n, dtype=np.float64) + out = np.zeros(n) + size_read_after_use(a, np.int64(nt), out, N=n) + assert np.allclose(out, a + (nt + 1)) + + +def test_size_can_be_reassigned_after_use_as_a_shape(): + n, nt = 5, 7 + a = np.arange(n, dtype=np.float64) + out = np.zeros(n) + size_reassigned_after_use(a, np.int64(nt), out, N=n) + assert np.allclose(out, a + 99) -def test_size_scalar_used_as_shape_is_promoted_to_symbol(): - sdfg = size_scalar_shape_prog.to_sdfg(simplify=True) # pre-fix: FileExistsError on Nt_plus_1 - assert "Nt_plus_1" in sdfg.symbols, "the size scalar must become a symbol" - assert "Nt_plus_1" not in sdfg.arrays, "the colliding data descriptor must be gone after promotion" +def test_promotion_leaves_the_descriptor_in_place(): + sdfg = size_read_after_use.to_sdfg(simplify=False) + promoted = [s for s in sdfg.symbols if s.startswith('__sym_')] + assert promoted, 'the size scalar must be read into a symbol' + # The descriptor stays: deleting it is what broke later reads of the size. + assert any(s[len('__sym_'):] in sdfg.arrays for s in promoted) sdfg.validate() -if __name__ == "__main__": - test_size_scalar_used_as_shape_is_promoted_to_symbol() - print("OK") +if __name__ == '__main__': + test_scalar_size_as_shape() + test_size_descriptor_survives_its_use_as_a_shape() + test_size_can_be_reassigned_after_use_as_a_shape() + test_promotion_leaves_the_descriptor_in_place() From bc66639e579c30fb812692396d6452923e5dde92 Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Mon, 20 Jul 2026 22:31:03 +0200 Subject: [PATCH 07/17] Maintain index-dim flags through Range mutation and unify the GEMM matrix view Follow-ups to tracking integer-index dimensions in Range: - __setitem__ kept eager bookkeeping that unpacked every dimension as a triple, but a degenerate dimension may legally hold the bare index expression instead (add_indirection_subgraph assigns one, and dim_to_string has always rendered it). That raised "cannot unpack non-iterable Symbol object" during frontend indirection lowering. Revert __setitem__ to its plain form: is_index_dim already cross-checks the range at read time, and now accepts the bare form. - map_dim_shuffle permuted ranges and tile_sizes by hand, leaving the flags misattributed to the wrong dimensions; use Range.reorder, which is exactly this permutation. Two redundant_array sites rebuilt a Range from bare lists and dropped the flags; carry them through (popped dimensions return as slices, matching unsqueeze). - The GEMM matrix-view rule ("raw subset if 2D, else squeezed") lived in two places on two different inputs. Collect it into _matrix_subset_size beside _matrix_operand in matmul.py; validate, expansion and codegen share it. --- dace/libraries/blas/nodes/gemm.py | 36 ++---- dace/libraries/blas/nodes/matmul.py | 45 +++++++- dace/subsets.py | 18 +-- .../dataflow/map_dim_shuffle.py | 3 +- .../dataflow/redundant_array.py | 10 +- doitgen_repro.py | 106 ------------------ tests/sdfg/subset_index_dims_test.py | 15 +++ 7 files changed, 85 insertions(+), 148 deletions(-) delete mode 100644 doitgen_repro.py diff --git a/dace/libraries/blas/nodes/gemm.py b/dace/libraries/blas/nodes/gemm.py index f295990c5e..16d8e4199e 100644 --- a/dace/libraries/blas/nodes/gemm.py +++ b/dace/libraries/blas/nodes/gemm.py @@ -8,25 +8,8 @@ import dace.sdfg.nodes from dace.transformation.transformation import ExpandTransformation from dace.libraries.blas.blas_helpers import to_blastype, check_access, dtype_to_cudadatatype, to_cublas_computetype -from dace.libraries.blas.nodes.matmul import (_get_matmul_operands, _get_codegen_gemm_opts) - - -def _squeezed_size(memlet): - """ - Returns the operand size as the matrix ``SpecializeMatMul`` saw when it chose GEMM. - - The dispatcher matches on the squeezed operand sizes, so validation and code generation have to - read the same view. Reading the raw subset instead rejects an operand the dispatcher just - accepted, which is how a ``(NQ, 1, NP)`` reshape reaches GEMM and is then refused. - - :param memlet: The memlet accessing the operand. - :return: The size of the memlet subset with its size-1 dimensions removed. - """ - subset = dc(memlet.subset) - subset.squeeze() - return subset.size() - - +from dace.libraries.blas.nodes.matmul import (_get_matmul_operands, _get_codegen_gemm_opts, _matrix_operand, + _matrix_subset_size) from .. import environments import numpy as np import warnings @@ -64,8 +47,9 @@ class ExpandGemmPure(ExpandTransformation): def make_sdfg(node, parent_state, parent_sdfg): sdfg = dace.SDFG(node.label + "_sdfg") - ((edge_a, outer_array_a, _, _, shape_a, strides_a), (edge_b, outer_array_b, _, _, shape_b, strides_b), - cdata) = _get_matmul_operands(node, parent_state, parent_sdfg) + adata, bdata, cdata = _get_matmul_operands(node, parent_state, parent_sdfg) + edge_a, outer_array_a, shape_a, strides_a = _matrix_operand(adata) + edge_b, outer_array_b, shape_b, strides_b = _matrix_operand(bdata) dtype_a = outer_array_a.dtype.type dtype_b = outer_array_b.dtype.type @@ -97,7 +81,7 @@ def make_sdfg(node, parent_state, parent_sdfg): _, array_a = sdfg.add_array("_a", shape_a, dtype_a, strides=strides_a, storage=outer_array_a.storage) _, array_b = sdfg.add_array("_b", shape_b, dtype_b, strides=strides_b, storage=outer_array_b.storage) - _, array_c = sdfg.add_array("_c", shape_c, dtype_c, strides=cdata[-1], storage=cdata[1].storage) + _, array_c = sdfg.add_array("_c", shape_c, dtype_c, strides=_matrix_operand(cdata)[3], storage=cdata[1].storage) if equal_valued(1, node.alpha): mul_program = "__out = __a * __b" @@ -569,11 +553,11 @@ def validate(self, sdfg, state): size2 = None for _, _, _, dst_conn, memlet in state.in_edges(self): if dst_conn == '_a': - size0 = _squeezed_size(memlet) + size0 = _matrix_subset_size(memlet.subset) if dst_conn == '_b': - size1 = _squeezed_size(memlet) + size1 = _matrix_subset_size(memlet.subset) if dst_conn == '_c': - size2 = _squeezed_size(memlet) + size2 = _matrix_subset_size(memlet.subset) if self.transA: size0 = list(reversed(size0)) @@ -593,7 +577,7 @@ def validate(self, sdfg, state): UserWarning) elif not res: raise ValueError("Inputs to matrix-matrix product must agree in the k-dimension") - size3 = _squeezed_size(out_memlet) + size3 = _matrix_subset_size(out_memlet.subset) if size2 is not None: res = [equal(s0, s1) for s0, s1 in zip(size2, size3)] fail = any([r is False for r in res]) diff --git a/dace/libraries/blas/nodes/matmul.py b/dace/libraries/blas/nodes/matmul.py index d3520bf499..1ba09109bc 100644 --- a/dace/libraries/blas/nodes/matmul.py +++ b/dace/libraries/blas/nodes/matmul.py @@ -7,6 +7,43 @@ from math import prod +def _matrix_subset_size(subset): + """ + Returns an operand's size in the matrix view ``SpecializeMatMul`` matched on. + + The dispatcher selects GEMM when an operand's raw subset is 2D *or* when it is 2D once + squeezed, so validation, expansion and code generation all have to read that same view. + Reading only the raw subset rejects a ``(NQ, 1, NP)`` reshape; reading only the squeezed one + rejects a genuine unit extent such as an ``(N, 1)`` column, which squeezes to a vector. + Disagreeing with the dispatcher either way makes GEMM refuse the operand it was just handed. + + :param subset: The subset of the memlet accessing the operand. + :return: The 2D size if either view supplies one, otherwise the squeezed size. + """ + size = subset.size() + if len(size) == 2: + return size + squeezed = dc(subset) + squeezed.squeeze() + return squeezed.size() + + +def _matrix_operand(operand): + """ + Returns a GEMM operand as ``(edge, descriptor, shape, strides)`` in its matrix view. + + Applies the rule of :func:`_matrix_subset_size` to an operand whose two views have already + been computed, so that the strides are picked from the same view as the shape. + + :param operand: One of the three tuples returned by :func:`_get_matmul_operands`. + :return: The edge, the outer descriptor, and the 2D shape and strides. + """ + edge, desc, size, strides, squeezed_size, squeezed_strides = operand + if len(size) == 2: + return edge, desc, size, strides + return edge, desc, squeezed_size, squeezed_strides + + def _get_matmul_operands(node, state, sdfg, name_lhs="_a", name_rhs="_b", name_out="_c"): """Returns the matrix multiplication input edges, arrays, and shape.""" res_lhs = None @@ -137,9 +174,11 @@ def _get_codegen_gemm_opts(node, state, sdfg, adesc, bdesc, cdesc, alpha, beta, from dace.codegen.common import sym2cpp from dace.libraries.blas.blas_helpers import get_gemm_opts - # GEMM operands are matrices: use the squeezed view, matching SpecializeMatMul's dispatch. - (_, _, _, _, ashape, astride), (_, _, _, _, bshape, bstride), (_, _, _, _, cshape, - cstride) = _get_matmul_operands(node, state, sdfg) + # M/N/K and the leading dimensions come from the matrix view the dispatcher matched. + adata, bdata, cdata = _get_matmul_operands(node, state, sdfg) + _, _, ashape, astride = _matrix_operand(adata) + _, _, bshape, bstride = _matrix_operand(bdata) + _, _, cshape, cstride = _matrix_operand(cdata) if getattr(node, 'transA', False): ashape = list(reversed(ashape)) diff --git a/dace/subsets.py b/dace/subsets.py index f78b61671d..7e7d0a73fb 100644 --- a/dace/subsets.py +++ b/dace/subsets.py @@ -352,7 +352,12 @@ def is_index_dim(self, dim: int) -> bool: """ if not self.index_dims[dim]: return False - rb, re, _ = self.ranges[dim] + rng = self.ranges[dim] + if not isinstance(rng, tuple): + # A degenerate dimension may hold the bare index expression instead of a triple + # (see ``add_indirection_subgraph``), which is degenerate by construction. + return True + rb, re, _ = rng return symbolic.equal_valued(0, re - rb) def rank_dims(self) -> List[int]: @@ -821,14 +826,9 @@ def __getitem__(self, key): return self.ranges.__getitem__(key) def __setitem__(self, key, value): - result = self.ranges.__setitem__(key, value) - # A replaced dimension is no longer the one that was parsed, so it is only still an index - # if the new range is degenerate. - for i in range(len(self.ranges)): - if self.index_dims[i]: - rb, re, _ = self.ranges[i] - self.index_dims[i] = bool(symbolic.equal_valued(0, re - rb)) - return result + # A widened dimension stops being an index; `is_index_dim` cross-checks the range itself, + # so replacing an entry here needs no bookkeeping of its own. + return self.ranges.__setitem__(key, value) def __eq__(self, other): if not isinstance(other, Range): diff --git a/dace/transformation/dataflow/map_dim_shuffle.py b/dace/transformation/dataflow/map_dim_shuffle.py index f4ee19d44e..1e90bdbfd2 100644 --- a/dace/transformation/dataflow/map_dim_shuffle.py +++ b/dace/transformation/dataflow/map_dim_shuffle.py @@ -40,6 +40,5 @@ def apply(self, graph: SDFGState, sdfg: SDFG): map_entry: nodes.MapEntry = self.map_entry new_map_order: list[int] = [map_entry.map.params.index(param) for param in self.parameters] - map_entry.range.ranges = [map_entry.range.ranges[new_pos] for new_pos in new_map_order] - map_entry.range.tile_sizes = [map_entry.range.tile_sizes[new_pos] for new_pos in new_map_order] + map_entry.range.reorder(new_map_order) map_entry.map.params = [map_entry.map.params[new_pos] for new_pos in new_map_order] diff --git a/dace/transformation/dataflow/redundant_array.py b/dace/transformation/dataflow/redundant_array.py index 4aacea151d..e738dbf610 100644 --- a/dace/transformation/dataflow/redundant_array.py +++ b/dace/transformation/dataflow/redundant_array.py @@ -259,11 +259,13 @@ def pop_dims(subset, dims): else: ranges = copy.deepcopy(subset.ranges) tsizes = copy.deepcopy(subset.tile_sizes) + index_dims = list(subset.index_dims) for i in dims: r = ranges.pop(i) t = tsizes.pop(i) + index_dims.pop(i) popped.append((r, t)) - new_subset = subsets.Range(ranges) + new_subset = subsets.Range(ranges, index_dims=index_dims) new_subset.tile_sizes = tsizes return new_subset, popped @@ -273,10 +275,14 @@ def compose_and_push_back(first, second, dims=None, popped=None): if dims and popped and len(dims) == len(popped): ranges = subset.ranges tsizes = subset.tile_sizes + # The popped dimensions come back as slices, matching what `unsqueeze` inserts: their + # original flags were not carried along, and a slice preserves the rank. + index_dims = list(subset.index_dims) for d, (r, t) in zip(reversed(dims), reversed(popped)): ranges.insert(d, r) tsizes.insert(d, t) - subset = subsets.Range(ranges) + index_dims.insert(d, False) + subset = subsets.Range(ranges, index_dims=index_dims) subset.tile_sizes = tsizes return subset diff --git a/doitgen_repro.py b/doitgen_repro.py deleted file mode 100644 index 01162a1dc8..0000000000 --- a/doitgen_repro.py +++ /dev/null @@ -1,106 +0,0 @@ -# Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. -"""TEMPORARY repro driver -- remove before the PR is finalised. - -Runs npbench's polybench/doitgen through the two pipelines the benchmark harness uses (``strict`` -= parse + simplify, and ``autoopt`` = simplify + auto_optimize) and checks the result against -numpy. Kernel and initializer are verbatim from npbench main: - - npbench/benchmarks/polybench/doitgen/doitgen_dace.py - npbench/benchmarks/polybench/doitgen/doitgen.py - npbench/benchmarks/polybench/doitgen/doitgen_numpy.py - -The ``(NQ, 1, NP)`` reshape is the point of the benchmark: it is what the in-tree test -``tests/npbench/polybench/doitgen_test.py`` stopped exercising in afd0efe48, which is why CI has -been green while npbench doitgen fails. - -Usage: python doitgen_repro.py [NR NQ NP] (defaults to 3 4 5) -""" -import copy -import sys -import traceback - -import numpy as np - -import dace as dc -import dace.dtypes as dtypes -import dace.transformation.auto.auto_optimize as opt - -NR, NQ, NP = (dc.symbol(s, dtype=dc.int64) for s in ('NR', 'NQ', 'NP')) - - -@dc.program -def kernel(A: dc.float64[NR, NQ, NP], C4: dc.float64[NP, NP]): - for r in range(NR): - A[r, :, :] = np.reshape(np.reshape(A[r], (NQ, 1, NP)) @ C4, (NQ, NP)) - - -def initialize(nr, nq, np_, datatype=np.float64): - A = np.fromfunction(lambda i, j, k: ((i * j + k) % np_) / np_, (nr, nq, np_), dtype=datatype) - C4 = np.fromfunction(lambda i, j: (i * j % np_) / np_, (np_, np_), dtype=datatype) - return A, C4 - - -def numpy_kernel(nr, nq, np_, A, C4): - A[:] = np.reshape(np.reshape(A, (nr, nq, 1, np_)) @ C4, (nr, nq, np_)) - - -def check(tag, got, ref): - if got.shape != ref.shape: - print(f"[{tag}] SHAPE MISMATCH got={got.shape} ref={ref.shape}") - return False - diff = np.abs(got - ref) - ok = np.allclose(got, ref, rtol=1e-12, atol=1e-12) - print(f"[{tag}] {'OK' if ok else 'MISMATCH'} max_abs={diff.max():.3e}") - if not ok: - bad = np.argwhere(~np.isclose(got, ref, rtol=1e-12, atol=1e-12)) - print(f"[{tag}] nbad={len(bad)}/{ref.size} first={bad[:3].tolist()}") - return ok - - -def run(tag, build, nr, nq, np_): - A, C4 = initialize(nr, nq, np_) - ref = A.copy() - numpy_kernel(nr, nq, np_, ref, C4) - try: - sdfg = build() - except Exception: - print(f"[{tag}] BUILD FAILED") - traceback.print_exc() - return False - got = A.copy() - try: - sdfg(A=got, C4=C4, NR=nr, NQ=nq, NP=np_) - except Exception: - print(f"[{tag}] RUN FAILED") - traceback.print_exc() - return False - return check(tag, got, ref) - - -def main(): - nr, nq, np_ = (int(x) for x in sys.argv[1:4]) if len(sys.argv) > 3 else (3, 4, 5) - print(f"=== doitgen NR={nr} NQ={nq} NP={np_} dace={dc.__file__}") - - base = kernel.to_sdfg(simplify=False) - - def strict(): - sdfg = copy.deepcopy(base) - sdfg._name = "strict" - sdfg.simplify() - return sdfg - - def autoopt(): - sdfg = copy.deepcopy(base) - sdfg._name = "autoopt" - sdfg.simplify() - opt.auto_optimize(sdfg, dtypes.DeviceType.CPU, symbols=dict(NR=nr, NQ=nq, NP=np_)) - return sdfg - - strict_ok = run("strict ", strict, nr, nq, np_) - autoopt_ok = run("autoopt", autoopt, nr, nq, np_) - print(f"=== RESULT strict={'PASS' if strict_ok else 'FAIL'} autoopt={'PASS' if autoopt_ok else 'FAIL'}") - return 0 if (strict_ok and autoopt_ok) else 1 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/tests/sdfg/subset_index_dims_test.py b/tests/sdfg/subset_index_dims_test.py index 35abf8d188..9d90c0b86e 100644 --- a/tests/sdfg/subset_index_dims_test.py +++ b/tests/sdfg/subset_index_dims_test.py @@ -162,6 +162,20 @@ def test_widening_a_dimension_clears_its_index_flag(): assert other.rank() == 3 +def test_bare_index_expression_in_ranges(): + """A degenerate dimension may hold the index expression itself instead of a triple. + + ``add_indirection_subgraph`` substitutes into such a dimension and assigns the result back + directly, and ``dim_to_string`` has always rendered that form, so the rank query has to accept + it rather than unpack it blindly. + """ + subset = subsets.Range.from_string('0:4, 0, 0:7') + subset[1] = dace.symbol('idx') + assert str(subset) == '0:4, idx, 0:7' + assert subset.is_index_dim(1) + assert subset.rank() == 2 + + def test_rank_reducing_slice_still_lowers_correctly(): """End-to-end guard: an integer index into a larger dimension still yields a 1D result.""" R, C = (dace.symbol(s, dtype=dace.int64) for s in ('R', 'C')) @@ -189,4 +203,5 @@ def take_row(A: dace.float64[R, C], out: dace.float64[C]): test_flags_survive_structural_operations() test_derived_subsets_keep_their_flags() test_widening_a_dimension_clears_its_index_flag() + test_bare_index_expression_in_ranges() test_rank_reducing_slice_still_lowers_correctly() From 4c7e766ad8b1045ec5adc584e00d4726c7bea91b Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Tue, 21 Jul 2026 00:06:01 +0200 Subject: [PATCH 08/17] Revert integer-index dimension tracking from Range The index_dims flag added to Range to distinguish an integer index from a unit slice (numpy rank reduction) has no production consumer: rank(), rank_dims() and is_index_dim() are read only by their own test. The GEMM fix this PR needs dispatches on unit-dim squeezability, which is the opposite question -- a reshape to (NQ, 1, NP) is all slices (rank 3) yet must be one 2D GEMM -- so it uses squeeze(), never the flag. The flag is also not maintained through Memlet.from_memlet, bounding_box_union or propagation (all rebuild a Range from bare tuples and reset it), so rank() is unreliable after one propagation pass regardless. Revert subsets.py, memlet_parser.py, map_dim_shuffle.py and redundant_array.py to their pre-feature state and drop the flag's test. The GEMM matrix-view fix (_matrix_subset_size / _matrix_operand in matmul.py) is independent -- it reads only subset.size()/squeeze() -- and stays. --- dace/frontend/python/memlet_parser.py | 4 +- dace/subsets.py | 125 ++--------- .../dataflow/map_dim_shuffle.py | 3 +- .../dataflow/redundant_array.py | 10 +- tests/sdfg/subset_index_dims_test.py | 207 ------------------ 5 files changed, 18 insertions(+), 331 deletions(-) delete mode 100644 tests/sdfg/subset_index_dims_test.py diff --git a/dace/frontend/python/memlet_parser.py b/dace/frontend/python/memlet_parser.py index 969c9c00d9..1be7d0d79b 100644 --- a/dace/frontend/python/memlet_parser.py +++ b/dace/frontend/python/memlet_parser.py @@ -69,9 +69,7 @@ def _ndslice_to_subset(ndslice): for i in range(len(ndslice)): if not is_tuple[i]: ndslice[i] = (ndslice[i], ndslice[i], 1) - # Record which dimensions the user wrote as an integer index. Widening them to (i, i, 1) above - # makes them indistinguishable from a ``i:i+1`` slice, but numpy gives the two different ranks. - return subsets.Range(ndslice, index_dims=[not t for t in is_tuple]) + return subsets.Range(ndslice) def _parse_dim_atom(das, atom): diff --git a/dace/subsets.py b/dace/subsets.py index 7e7d0a73fb..3951d91491 100644 --- a/dace/subsets.py +++ b/dace/subsets.py @@ -305,20 +305,7 @@ def _tuple_to_symexpr(val): class Range(Subset): """ Subset defined in terms of a fixed range. """ - def __init__(self, ranges, index_dims: Optional[Sequence[bool]] = None): - """ - :param ranges: The per-dimension (begin, end, step[, tile]) tuples. - :param index_dims: Which dimensions were written as an integer index rather than a slice. - numpy drops those dimensions from the result rank (``a[0]`` is one rank - lower than ``a[0:1]``), and ``(begin, end, step)`` cannot express the - difference on its own since both are ``(i, i, 1)``. Defaults to all - slices, which is what a full-array access means. - - :note: The flags survive JSON serialization but not the string form, which renders a - degenerate range as ``i`` either way. Subset strings are re-parsed as symbolic - expressions in places that reject slice syntax (see - ``dace.transformation.passes.scalar_to_symbol``), so the rendering cannot change. - """ + def __init__(self, ranges): parsed_ranges = [] parsed_tiles = [] for r in ranges: @@ -331,63 +318,14 @@ def __init__(self, ranges, index_dims: Optional[Sequence[bool]] = None): parsed_tiles.append(symbolic.pystr_to_symbolic(r[3])) self.ranges = parsed_ranges self.tile_sizes = parsed_tiles - if index_dims is None: - self.index_dims = [False] * len(parsed_ranges) - else: - index_dims = list(index_dims) - if len(index_dims) != len(parsed_ranges): - raise ValueError(f"Expected {len(parsed_ranges)} index flags, got {len(index_dims)}") - self.index_dims = index_dims - - def is_index_dim(self, dim: int) -> bool: - """ - Returns True if the given dimension was accessed with an integer index. - - Only a degenerate range can be an index, so the flag is cross-checked against the range - itself rather than trusted outright: ``ranges`` is mutable through ``__setitem__`` and - direct assignment, and widening a dimension must not leave it looking indexed. - - :param dim: The dimension to query. - :return: True if the dimension was written as an integer index. - """ - if not self.index_dims[dim]: - return False - rng = self.ranges[dim] - if not isinstance(rng, tuple): - # A degenerate dimension may hold the bare index expression instead of a triple - # (see ``add_indirection_subgraph``), which is degenerate by construction. - return True - rb, re, _ = rng - return symbolic.equal_valued(0, re - rb) - - def rank_dims(self) -> List[int]: - """ - Returns the dimensions that survive numpy rank reduction, i.e. those not indexed away. - - An integer index removes its dimension (``a[0, 0:N]`` has rank 1); a slice never does, not - even when its extent is 1 (``a[0:1, 0:N]`` has rank 2). This is the numpy rule, and it is - independent of :meth:`squeeze`, which removes every extent-1 dimension regardless of how it - was written (``np.squeeze``). - - :return: The indices of the dimensions that are not indexed away. - """ - return [i for i in range(len(self.ranges)) if not self.is_index_dim(i)] - - def rank(self) -> int: - """ - Returns the numpy rank of this access: the number of dimensions not indexed away. - - :return: The number of dimensions that survive rank reduction. - """ - return len(self.ranges) - sum(1 for i in range(len(self.ranges)) if self.is_index_dim(i)) @staticmethod def from_indices(indices: Union["Indices", Sequence[int | str | symbolic.SymbolicType]]): if isinstance(indices, Indices): - return Range([(i, i, 1) for i in indices.indices], index_dims=[True] * len(indices.indices)) + return Range([(i, i, 1) for i in indices.indices]) indices = [symbolic.pystr_to_symbolic(i) for i in indices] - return Range([(i, i, 1) for i in indices], index_dims=[True] * len(indices)) + return Range([(i, i, 1) for i in indices]) def to_json(self): ret = [] @@ -395,11 +333,8 @@ def to_json(self): def a2s(obj): return symbolic.serialize_symbolic(obj) - for (start, end, step), tile, indexed in zip(self.ranges, self.tile_sizes, self.index_dims): - entry = {'start': a2s(start), 'end': a2s(end), 'step': a2s(step), 'tile': a2s(tile)} - if indexed: - entry['indexed'] = True - ret.append(entry) + for (start, end, step), tile in zip(self.ranges, self.tile_sizes): + ret.append({'start': a2s(start), 'end': a2s(end), 'step': a2s(step), 'tile': a2s(tile)}) return {'type': 'Range', 'ranges': ret} @@ -414,16 +349,12 @@ def from_json(obj, context=None): ranges = obj['ranges'] tuples = [] - index_dims = [] for r in ranges: tuples.append((_symbolic_deserializer(r['start'], context), _symbolic_deserializer(r['end'], context), _symbolic_deserializer(r['step'], context), _symbolic_deserializer(r['tile'], context))) - # Absent in SDFGs written before index dimensions were tracked; a slice is the safe - # default there, since it preserves the rank the file was saved with. - index_dims.append(bool(r.get('indexed', False))) - return Range(tuples, index_dims=index_dims) + return Range(tuples) @staticmethod def from_array(array: 'dace.data.Data'): @@ -438,8 +369,7 @@ def __hash__(self): def __add__(self, other): return Range( - ((*ranges, tile) for ranges, tile in zip(self.ranges + other.ranges, self.tile_sizes + other.tile_sizes)), - index_dims=self.index_dims + other.index_dims) + ((*ranges, tile) for ranges, tile in zip(self.ranges + other.ranges, self.tile_sizes + other.tile_sizes))) def __deepcopy__(self, memo) -> 'Range': """Performs a deepcopy of ``self``. @@ -451,7 +381,6 @@ def __deepcopy__(self, memo) -> 'Range': node = object.__new__(Range) node.ranges = self.ranges.copy() node.tile_sizes = self.tile_sizes.copy() - node.index_dims = self.index_dims.copy() return node @@ -578,7 +507,7 @@ def offset(self, other, negative, indices=None, offset_end=True): def offset_new(self, other, negative, indices=None, offset_end=True): if other is None: - return Range(self.ranges, index_dims=self.index_dims) + return Range(self.ranges) if not isinstance(other, Subset): if isinstance(other, (list, tuple)): other = Range.from_indices(other) @@ -588,11 +517,8 @@ def offset_new(self, other, negative, indices=None, offset_end=True): if indices is None: indices = set(range(len(self.ranges))) off = other.min_element() - # Offsetting shifts a subset, it never changes how a dimension was written, so the index - # flags follow the dimensions that are kept. return Range([(self.ranges[i][0] + mult * off[i], self.ranges[i][1] if not offset_end else - (self.ranges[i][1] + mult * off[i]), self.ranges[i][2]) for i in indices], - index_dims=[self.index_dims[i] for i in indices]) + (self.ranges[i][1] + mult * off[i]), self.ranges[i][2]) for i in indices]) def dims(self): return len(self.ranges) @@ -647,10 +573,8 @@ def reorder(self, order: Sequence[int]) -> None: """ new_ranges = [self.ranges[o] for o in order] new_tile_sizes = [self.tile_sizes[o] for o in order] - new_index_dims = [self.index_dims[o] for o in order] self.ranges = new_ranges self.tile_sizes = new_tile_sizes - self.index_dims = new_index_dims @staticmethod def dim_to_string(d, t=1): @@ -688,7 +612,6 @@ def from_string(string): # regtile_j * rs_j : min(K, regtile_j * rs_j + rs_j) ranges = [] - index_dims = [] # Split string to tokens separated by colons. # tokens = [ @@ -747,7 +670,6 @@ def from_string(string): if len(uni_dim_tokens) < 2: value = symbolic.pystr_to_symbolic(uni_dim_tokens[0].strip()) ranges.append((value, value, 1)) - index_dims.append(True) continue #return Range(ranges) # If dimension has more than 4 tokens, the range is invalid @@ -794,9 +716,8 @@ def from_string(string): raise SyntaxError("Invalid range: {}".format(string)) # Append range ranges.append((begin, end, step, tsize)) - index_dims.append(False) - return Range(ranges, index_dims=index_dims) + return Range(ranges) @staticmethod def ndslice_to_string(slice, tile_sizes=None): @@ -826,8 +747,6 @@ def __getitem__(self, key): return self.ranges.__getitem__(key) def __setitem__(self, key, value): - # A widened dimension stops being an index; `is_index_dim` cross-checks the range itself, - # so replacing an entry here needs no bookkeeping of its own. return self.ranges.__setitem__(key, value) def __eq__(self, other): @@ -846,25 +765,18 @@ def compose(self, other): raise TypeError("Cannot compose ranges with non-subsets") new_subset = [] - # A composed dimension is an index exactly when the dimension it originates from is: a - # degenerate dimension of ``self`` keeps its own flag, a dimension consumed from ``other`` - # takes ``other``'s. - new_index_dims = [] - other_index_dims = other.index_dims if isinstance(other, Range) else [True] * other.dims() if self.data_dims() == other.dims(): # case 1: subsets may differ in dimensions, but data_dims correspond # to other dims -> all non-data dims are cut out idx = 0 - for dim, ((rb, re, rs), rt) in enumerate(zip(self.ranges, self.tile_sizes)): + for (rb, re, rs), rt in zip(self.ranges, self.tile_sizes): if re - rb == 0: new_subset.append((rb, re, rs, rt)) - new_index_dims.append(self.index_dims[dim]) else: if isinstance(other[idx], tuple): new_subset.append((rb + rs * other[idx][0], rb + rs * other[idx][1], rs * other[idx][2], rt)) else: new_subset.append(rb + rs * other[idx]) - new_index_dims.append(other_index_dims[idx]) idx += 1 elif self.dims() == other.dims(): # case 2: subsets have the same dimensions (but possibly different @@ -872,13 +784,11 @@ def compose(self, other): for idx, ((rb, re, rs), rt) in enumerate(zip(self.ranges, self.tile_sizes)): if re - rb == 0: new_subset.append((rb, re, rs, rt)) - new_index_dims.append(self.index_dims[idx]) else: if isinstance(other[idx], tuple): new_subset.append((rb + rs * other[idx][0], rb + rs * other[idx][1], rs * other[idx][2], rt)) else: new_subset.append(rb + rs * other[idx]) - new_index_dims.append(other_index_dims[idx]) elif (other.data_dims() == 0 and all([r == (0, 0, 1) if isinstance(other, Range) else r == 0 for r in other])): # NOTE: This is a special case where the other subset is the # (potentially multidimensional) index zero. @@ -888,7 +798,6 @@ def compose(self, other): new_subset.extend(self.ranges) else: new_subset.extend([rb for rb, _, _ in self.ranges]) - new_index_dims.extend(self.index_dims) else: raise ValueError("Dimension mismatch in composition: " "Subset composed must be either completely " @@ -896,7 +805,7 @@ def compose(self, other): "or be not stripped of latter at all.") if isinstance(other, Range): - return Range(new_subset, index_dims=new_index_dims) + return Range(new_subset) else: raise NotImplementedError @@ -928,14 +837,11 @@ def squeeze(self, ignore_indices: Optional[List[int]] = None, offset: bool = Tru pass squeezed_ranges = [self.ranges[i] for i in non_ones] squeezed_tsizes = [self.tile_sizes[i] for i in non_ones] - squeezed_index_dims = [self.index_dims[i] for i in non_ones] if not squeezed_ranges: squeezed_ranges = [(0, 0, 1)] squeezed_tsizes = [1] - squeezed_index_dims = [False] self.ranges = squeezed_ranges self.tile_sizes = squeezed_tsizes - self.index_dims = squeezed_index_dims if offset: self.offset(self, True, indices=offset_indices) return non_ones @@ -962,7 +868,6 @@ def unsqueeze(self, axes: Sequence[int]) -> List[int]: for axis in sorted(axes): self.ranges.insert(axis, (0, 0, 1)) self.tile_sizes.insert(axis, 1) - self.index_dims.insert(axis, False) if len(result) > 0 and result[-1] >= axis: result.append(result[-1] + 1) @@ -973,19 +878,15 @@ def unsqueeze(self, axes: Sequence[int]) -> List[int]: def pop(self, dimensions): new_ranges = [] new_tsizes = [] - new_index_dims = [] for i in range(len(self.ranges)): if i not in dimensions: new_ranges.append(self.ranges[i]) new_tsizes.append(self.tile_sizes[i]) - new_index_dims.append(self.index_dims[i]) if not new_ranges: new_ranges = [(symbolic.pystr_to_symbolic(0), symbolic.pystr_to_symbolic(0), symbolic.pystr_to_symbolic(1))] new_tsizes = [symbolic.pystr_to_symbolic(1)] - new_index_dims = [False] self.ranges = new_ranges self.tile_sizes = new_tsizes - self.index_dims = new_index_dims def string_list(self): return Range.ndslice_to_string_list(self.ranges, self.tile_sizes) @@ -1112,7 +1013,7 @@ def __init__(self, indices: Sequence[int | str | symbolic.SymbolicType]): raise TypeError("Expected collection of index expression: got SymExpr") indices = [symbolic.pystr_to_symbolic(i) for i in indices] - super().__init__([(idx, idx, 1) for idx in indices], index_dims=[True] * len(indices)) + super().__init__([(idx, idx, 1) for idx in indices]) @property def indices(self) -> List[symbolic.SymbolicType]: diff --git a/dace/transformation/dataflow/map_dim_shuffle.py b/dace/transformation/dataflow/map_dim_shuffle.py index 1e90bdbfd2..f4ee19d44e 100644 --- a/dace/transformation/dataflow/map_dim_shuffle.py +++ b/dace/transformation/dataflow/map_dim_shuffle.py @@ -40,5 +40,6 @@ def apply(self, graph: SDFGState, sdfg: SDFG): map_entry: nodes.MapEntry = self.map_entry new_map_order: list[int] = [map_entry.map.params.index(param) for param in self.parameters] - map_entry.range.reorder(new_map_order) + map_entry.range.ranges = [map_entry.range.ranges[new_pos] for new_pos in new_map_order] + map_entry.range.tile_sizes = [map_entry.range.tile_sizes[new_pos] for new_pos in new_map_order] map_entry.map.params = [map_entry.map.params[new_pos] for new_pos in new_map_order] diff --git a/dace/transformation/dataflow/redundant_array.py b/dace/transformation/dataflow/redundant_array.py index e738dbf610..4aacea151d 100644 --- a/dace/transformation/dataflow/redundant_array.py +++ b/dace/transformation/dataflow/redundant_array.py @@ -259,13 +259,11 @@ def pop_dims(subset, dims): else: ranges = copy.deepcopy(subset.ranges) tsizes = copy.deepcopy(subset.tile_sizes) - index_dims = list(subset.index_dims) for i in dims: r = ranges.pop(i) t = tsizes.pop(i) - index_dims.pop(i) popped.append((r, t)) - new_subset = subsets.Range(ranges, index_dims=index_dims) + new_subset = subsets.Range(ranges) new_subset.tile_sizes = tsizes return new_subset, popped @@ -275,14 +273,10 @@ def compose_and_push_back(first, second, dims=None, popped=None): if dims and popped and len(dims) == len(popped): ranges = subset.ranges tsizes = subset.tile_sizes - # The popped dimensions come back as slices, matching what `unsqueeze` inserts: their - # original flags were not carried along, and a slice preserves the rank. - index_dims = list(subset.index_dims) for d, (r, t) in zip(reversed(dims), reversed(popped)): ranges.insert(d, r) tsizes.insert(d, t) - index_dims.insert(d, False) - subset = subsets.Range(ranges, index_dims=index_dims) + subset = subsets.Range(ranges) subset.tile_sizes = tsizes return subset diff --git a/tests/sdfg/subset_index_dims_test.py b/tests/sdfg/subset_index_dims_test.py deleted file mode 100644 index 9d90c0b86e..0000000000 --- a/tests/sdfg/subset_index_dims_test.py +++ /dev/null @@ -1,207 +0,0 @@ -# Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. -"""``Range`` records which dimensions were written as an integer index. - -``(begin, end, step)`` cannot express the difference between ``a[0]`` and ``a[0:1]`` -- both are -``(0, 0, 1)`` -- but numpy distinguishes them: an integer index removes its dimension from the -result rank, a slice never does, not even at extent 1. Consumers that need the rank of an access -(``x @ y`` dispatching to a matrix or a vector product, for instance) previously had to guess by -squeezing every extent-1 dimension, which is a different operation (``np.squeeze``) and is wrong -whenever a sliced dimension happens to have extent 1. -""" -import copy - -import numpy as np - -import dace -from dace import subsets - - -def test_index_reduces_rank_slice_does_not(): - indexed = subsets.Range.from_string('0:NQ, 0, 0:NP') - sliced = subsets.Range.from_string('0:NQ, 0:1, 0:NP') - - # Identical extents, so neither size() nor equality can tell them apart ... - assert indexed.size() == sliced.size() - assert indexed == sliced - # ... but the rank differs, which is what numpy cares about. - assert indexed.rank() == 2 - assert sliced.rank() == 3 - assert indexed.rank_dims() == [0, 2] - assert sliced.rank_dims() == [0, 1, 2] - - -def test_unit_extent_slice_keeps_its_rank(): - """The case squeezing cannot get right: a sliced dimension whose extent happens to be 1.""" - subset = subsets.Range.from_string('0:1, 0, 0:7') - assert subset.size() == [1, 1, 7] - assert subset.rank() == 2 # only the middle dimension was indexed away - assert subset.rank_dims() == [0, 2] - - squeezed = copy.deepcopy(subset) - squeezed.squeeze() - assert squeezed.size() == [7] # np.squeeze drops both, which is a different question - - -def test_squeeze_is_still_numpy_squeeze(): - """``squeeze`` keeps meaning ``np.squeeze``: drop every extent-1 dimension.""" - subset = subsets.Range.from_string('0:4, 0:1, 0:7') - assert subset.rank() == 3 - subset.squeeze() - assert subset.size() == [4, 7] - - -def test_from_array_preserves_rank(): - """Reading all of an array is never a rank reduction, even for its unit dimensions.""" - sdfg = dace.SDFG('from_array_rank') - sdfg.add_array('V', [6, 1, 5], dace.float64) - subset = subsets.Range.from_array(sdfg.arrays['V']) - assert subset.rank() == 3 - assert not any(subset.index_dims) - - -def test_from_indices_is_all_indices(): - subset = subsets.Range.from_indices([2, 3]) - assert subset.rank() == 0 - assert all(subset.index_dims) - - -def test_string_form_is_unchanged(): - """The rendering stays exactly as it was. - - Subset strings are re-parsed as symbolic expressions in places that reject slice syntax (see - ``dace.transformation.passes.scalar_to_symbol``), so a degenerate range must keep printing as - ``i``. The flags therefore live in memory and in JSON, not in the string form. - """ - assert str(subsets.Range.from_string('0:NQ, 0, 0:NP')) == '0:NQ, 0, 0:NP' - assert str(subsets.Range.from_string('0:NQ, 0:1, 0:NP')) == '0:NQ, 0, 0:NP' - assert str(subsets.Range.from_string('0, 0:10')) == '0, 0:10' - - -def test_frontend_records_mixed_subscript_ranks(): - """``A[0, 0:N]`` is rank 1 in numpy, so the frontend has to record which dimension was indexed.""" - from dace.frontend.python import memlet_parser - - subset = memlet_parser._ndslice_to_subset([0, (0, 9, 1)]) - assert subset.index_dims == [True, False] - assert subset.rank() == 1 - - both_slices = memlet_parser._ndslice_to_subset([(0, 0, 1), (0, 9, 1)]) - assert both_slices.index_dims == [False, False] - assert both_slices.rank() == 2 - - -def test_json_marks_only_index_dims(): - subset = subsets.Range.from_string('0:NQ, 0, 0:NP') - entries = subset.to_json()['ranges'] - assert [entry.get('indexed', False) for entry in entries] == [False, True, False] - - -def test_sdfg_save_load_preserves_rank(tmp_path): - sdfg = dace.SDFG('subset_rank_roundtrip') - sdfg.add_array('A', [4, 1, 8], dace.float64) - state = sdfg.add_state('s', is_start_block=True) - read, write = state.add_read('A'), state.add_write('A') - # A[2, 0:1, 0:8]: dimension 0 is indexed away, dimension 1 is a unit slice that stays. - state.add_nedge(read, write, dace.Memlet('A[2, 0:1, 0:8]')) - assert state.edges()[0].data.subset.rank() == 2 - - path = tmp_path / 'roundtrip.sdfg' - sdfg.save(path) - reloaded = dace.SDFG.from_file(path) - subset = list(reloaded.states())[0].edges()[0].data.subset - assert subset.index_dims == [True, False, False] - assert subset.rank() == 2 - - -def test_flags_survive_structural_operations(): - subset = subsets.Range.from_string('0:4, 0, 0:7, 2') - assert subset.index_dims == [False, True, False, True] - - assert copy.deepcopy(subset).index_dims == subset.index_dims - assert (subset + subsets.Range.from_string('0')).index_dims == [False, True, False, True, True] - - reordered = copy.deepcopy(subset) - reordered.reorder([1, 0, 3, 2]) - assert reordered.index_dims == [True, False, True, False] - - popped = copy.deepcopy(subset) - popped.pop([0]) - assert popped.index_dims == [True, False, True] - - unsqueezed = subsets.Range.from_string('0:10') - unsqueezed.unsqueeze([0]) - assert unsqueezed.index_dims == [False, False] # unsqueeze inserts slices, not indices - - -def test_derived_subsets_keep_their_flags(): - """Every operation that builds a new Range from an existing one must carry the flags over.""" - subset = subsets.Range.from_string('0:4, 0, 0:7') - - assert subset.offset_new(None, False).index_dims == [False, True, False] - assert subset.offset_new([1, 0, 2], True).index_dims == [False, True, False] - - # compose: a degenerate dimension keeps its own flag, a consumed one takes the other subset's - composed = subset.compose(subsets.Range.from_string('0:4, 0:7')) - assert composed.index_dims == [False, True, False] - assert composed.rank() == 2 - - -def test_widening_a_dimension_clears_its_index_flag(): - """An index is degenerate by definition, so a widened dimension cannot stay flagged.""" - subset = subsets.Range.from_string('0:4, 0, 0:7') - assert subset.rank() == 2 - - subset[1] = (0, 3, 1) - assert subset.rank() == 3 - assert not subset.is_index_dim(1) - - # Direct mutation of `ranges` bypasses __setitem__ entirely, so the flag is cross-checked - # against the range rather than trusted. - other = subsets.Range.from_string('0:4, 0, 0:7') - other.ranges[1] = (0, 5, 1) - assert other.rank() == 3 - - -def test_bare_index_expression_in_ranges(): - """A degenerate dimension may hold the index expression itself instead of a triple. - - ``add_indirection_subgraph`` substitutes into such a dimension and assigns the result back - directly, and ``dim_to_string`` has always rendered that form, so the rank query has to accept - it rather than unpack it blindly. - """ - subset = subsets.Range.from_string('0:4, 0, 0:7') - subset[1] = dace.symbol('idx') - assert str(subset) == '0:4, idx, 0:7' - assert subset.is_index_dim(1) - assert subset.rank() == 2 - - -def test_rank_reducing_slice_still_lowers_correctly(): - """End-to-end guard: an integer index into a larger dimension still yields a 1D result.""" - R, C = (dace.symbol(s, dtype=dace.int64) for s in ('R', 'C')) - - @dace.program - def take_row(A: dace.float64[R, C], out: dace.float64[C]): - out[:] = A[2, :] - - rng = np.random.default_rng(0) - A = rng.random((4, 8)) - out = np.zeros(8) - take_row(A, out, R=4, C=8) - assert np.allclose(out, A[2]) - - -if __name__ == '__main__': - test_index_reduces_rank_slice_does_not() - test_unit_extent_slice_keeps_its_rank() - test_squeeze_is_still_numpy_squeeze() - test_from_array_preserves_rank() - test_from_indices_is_all_indices() - test_string_form_is_unchanged() - test_frontend_records_mixed_subscript_ranks() - test_json_marks_only_index_dims() - test_flags_survive_structural_operations() - test_derived_subsets_keep_their_flags() - test_widening_a_dimension_clears_its_index_flag() - test_bare_index_expression_in_ranges() - test_rank_reducing_slice_still_lowers_correctly() From a6e26d489f39f34ebf1788d3e6b04cc799e87481 Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Tue, 21 Jul 2026 00:06:20 +0200 Subject: [PATCH 09/17] Capture each array's size in its own symbol; read PBLAS as a matrix Shape promotion reused one __sym_ symbol per size scalar and re-assigned it on every use. Two arrays sized from the same reassigned variable (m = 64; a = np.empty(m); m = 2; b = np.empty(m)) then shared one symbol whose value was overwritten, so a read as length 2 instead of 64 -- a silent miscompile. A shape scalar used again as an index after the reassignment (a[m]) hit the same collision through the subscript-promotion path and read out of bounds. Mint a distinct symbol for each shape (promote_scalar_to_symbol grows a `fresh` flag), so every array records the size it was created with, and give fresh symbols a suffix so they can never land on the bare __sym_ the index path reuses. Iterating the size names in sorted order keeps the promotion states deterministic. ExpandGemmPBLAS still read the raw operand size while validate and the other expansions were moved to the matrix view, so a (NQ, 1, NP) operand validate accepts would be mis-sized there; read it through _matrix_operand too. --- dace/frontend/python/newast.py | 37 ++++++--- .../replacements/array_creation_dace.py | 28 ++++--- dace/libraries/blas/nodes/gemm.py | 6 +- tests/size_scalar_shape_promotion_test.py | 79 ++++++++++++++++++- tests/transpose_unit_dim_squeeze_test.py | 1 + 5 files changed, 126 insertions(+), 25 deletions(-) diff --git a/dace/frontend/python/newast.py b/dace/frontend/python/newast.py index 4c4788f7c6..4667ece3fc 100644 --- a/dace/frontend/python/newast.py +++ b/dace/frontend/python/newast.py @@ -5373,29 +5373,48 @@ def range_is_index(range: subsets.Range) -> bool: wcr=expr.wcr)) return tmp - def promote_scalar_to_symbol(self, scalar: str, key: Optional[str] = None) -> symbolic.symbol: + def promote_scalar_to_symbol(self, scalar: str, key: Optional[str] = None, fresh: bool = False) -> symbolic.symbol: """ Reads a scalar into a symbol so its value can be used where only symbols are allowed. - A fresh ``__sym_`` symbol is assigned from the scalar on an interstate edge. The - scalar's data descriptor is deliberately left in place: it may still be read or reassigned - later in the program, and the visitor's own bookkeeping still refers to it by name. + A ``__sym_`` symbol is assigned from the scalar on an interstate edge, capturing the + scalar's value at that point in the program. The scalar's data descriptor is deliberately + left in place: it may still be read or reassigned later, and the visitor's own bookkeeping + still refers to it by name. :param scalar: Name of the scalar data descriptor to read. :param key: Cache key for the promotion, defaulting to the scalar name. Repeated promotions of the same expression reuse the symbol instead of minting a new one. + :param fresh: Mint a new symbol every call instead of reusing a cached one. A size scalar + is materialized into a temporary that the frontend reuses across calls + (``m = ...; np.empty(m); m = ...; np.empty(m)`` promotes one temporary at two + values); reusing one symbol would give both arrays the last value written to + it, so each shape must capture its own. :return: The symbol carrying the scalar's value. """ key = key if key is not None else scalar desc = self.sdfg.arrays[scalar] - sym = self.indirections.get(key) + sym = None if fresh else self.indirections.get(key) if sym is None: - sym = dace.symbol(f'__sym_{scalar}', dtype=desc.dtype) - self.indirections[key] = sym + if fresh: + # A fresh promotion always carries a suffix, so it can never land on the bare + # `__sym_` a cached promotion of the same scalar reuses (an index access + # after the size is reassigned would otherwise re-bind the array's extent symbol). + # Check `arrays` too: a colliding data name would make `add_symbol` raise. + i = 0 + name = f'__sym_{scalar}_{i}' + while name in self.sdfg.symbols or name in self.sdfg.arrays: + i += 1 + name = f'__sym_{scalar}_{i}' + else: + name = f'__sym_{scalar}' + sym = dace.symbol(name, dtype=desc.dtype) + if not fresh: + self.indirections[key] = sym try: - self.sdfg.add_symbol(f'__sym_{scalar}', desc.dtype) + self.sdfg.add_symbol(name, desc.dtype) except FileExistsError: - # By design this may re-add an existing symbol; the exception is benign. + # A cached (non-fresh) promotion may re-add an existing symbol; that is benign. pass # A promoted size can end up in an array shape, and nested scopes resolve the free # symbols of their scope arrays through `globals`, so the symbol has to be visible diff --git a/dace/frontend/python/replacements/array_creation_dace.py b/dace/frontend/python/replacements/array_creation_dace.py index 0476798189..40a7233cca 100644 --- a/dace/frontend/python/replacements/array_creation_dace.py +++ b/dace/frontend/python/replacements/array_creation_dace.py @@ -20,12 +20,17 @@ def promote_size_scalars_in_shape(pv: ProgramVisitor, sdfg: SDFG, shape: Shape) -> Shape: """ - Rewrites a shape so that scalars used as extents are read through symbols. + Rewrites a shape so that a size scalar used as an extent is read through a symbol. An extent must be a symbol, but a size computed in the program (``nt = Nt + 1`` then - ``np.empty(nt)``) is a scalar data descriptor. Each such scalar is read into a symbol on an - interstate edge and substituted into the shape. The descriptor itself is left alone, so the - program may keep reading or reassigning the size afterwards. + ``np.empty(nt)``) is a size-1 data descriptor. Each such name is read into a symbol on an + interstate edge and substituted into the shape, leaving the descriptor in place so the program + may keep reading the size afterwards. + + The size is captured into a symbol minted for this shape (``fresh``): the frontend reuses one + temporary for a size expression across calls, so ``m = ...; np.empty(m); m = ...; np.empty(m)`` + would otherwise give both arrays whichever value was written to the shared symbol last. A + per-shape symbol records the value at the point each array is created. :param pv: The program visitor, which owns symbol promotion and the state machine. :param sdfg: The SDFG being built. @@ -33,19 +38,20 @@ def promote_size_scalars_in_shape(pv: ProgramVisitor, sdfg: SDFG, shape: Shape) :return: The shape with scalar extents replaced by symbols. """ resolved = [symbolic.pystr_to_symbolic(e) if isinstance(e, str) else e for e in shape] - replacements = {} + names = set() for extent in resolved: if not isinstance(extent, sympy.Basic): continue for sym in extent.free_symbols: name = str(sym) - if name in replacements or name in sdfg.symbols or name not in sdfg.arrays: - continue - desc = sdfg.arrays[name] - if isinstance(desc, data.Scalar) or desc.total_size == 1: - replacements[sym] = pv.promote_scalar_to_symbol(name) - if not replacements: + if name in sdfg.arrays and name not in sdfg.symbols and sdfg.arrays[name].total_size == 1: + names.add(name) + if not names: return shape + + # One symbol per distinct name in this shape (np.empty((m, m)) shares it); sorted() keeps the + # promotion states deterministic when several size scalars appear. + replacements = {symbolic.pystr_to_symbolic(n): pv.promote_scalar_to_symbol(n, fresh=True) for n in sorted(names)} return [e.subs(replacements) if isinstance(e, sympy.Basic) else e for e in resolved] diff --git a/dace/libraries/blas/nodes/gemm.py b/dace/libraries/blas/nodes/gemm.py index 16d8e4199e..ab0809d89a 100644 --- a/dace/libraries/blas/nodes/gemm.py +++ b/dace/libraries/blas/nodes/gemm.py @@ -466,7 +466,11 @@ class ExpandGemmPBLAS(ExpandTransformation): @staticmethod def expansion(node, state, sdfg): node.validate(sdfg, state) - (_, adesc, ashape, _, _, _), (_, bdesc, bshape, _, _, _), _ = _get_matmul_operands(node, state, sdfg) + # Read the same matrix view the dispatcher and validate agreed on; reading the raw subset + # would mis-size an operand (e.g. an (NQ, 1, NP) reshape) that validate has already accepted. + adata, bdata, _ = _get_matmul_operands(node, state, sdfg) + _, adesc, ashape, _ = _matrix_operand(adata) + _, bdesc, bshape, _ = _matrix_operand(bdata) dtype = adesc.dtype.base_type if not equal_valued(0, node.beta): diff --git a/tests/size_scalar_shape_promotion_test.py b/tests/size_scalar_shape_promotion_test.py index e66ebbf7d7..d1dc38aa25 100644 --- a/tests/size_scalar_shape_promotion_test.py +++ b/tests/size_scalar_shape_promotion_test.py @@ -5,6 +5,8 @@ symbol, and minting a symbol of the same name collided with the descriptor (``FileExistsError``). The size is now read into a ``__sym_`` symbol on an interstate edge and substituted into the shape, leaving the descriptor in place so the program can keep reading or reassigning the size afterwards. +Each shape captures its own symbol, so two arrays sized from the same reused variable at different +values keep the values they were created with. """ import numpy as np import dace @@ -40,6 +42,28 @@ def size_reassigned_after_use(a: dace.float64[N], Nt: dace.int64, out: dace.floa out[i] = a[i] + m +@dace.program +def two_arrays_from_reassigned_size(Nt: dace.int64, out: dace.float64[1]): + m = Nt + b = np.empty(m, dace.float64) + for i in range(64): + b[i] = 1.0 + m = 2 # a second array from the same name at a different value + c = np.empty(m, dace.float64) + c[0] = 0.0 + out[0] = np.sum(b) # must sum all 64 of b, not be truncated to c's size + + +@dace.program +def size_reused_as_index(out: dace.float64[1]): + m = 8 + a = np.empty(m, dace.float64) + for i in range(8): + a[i] = i * 1.0 + m = 2 + out[0] = a[m] # the shape symbol must not be the one this index reassigns + + def test_scalar_size_as_shape(): n, nt = 5, 7 a = np.arange(n, dtype=np.float64) @@ -65,17 +89,64 @@ def test_size_can_be_reassigned_after_use_as_a_shape(): assert np.allclose(out, a + 99) +def test_two_arrays_from_a_reassigned_size_keep_their_own_extents(): + """A per-shape symbol: reusing one size name for two arrays must not collapse their extents. + + A single shared symbol gave both arrays the last value written to it, so ``np.sum(b)`` read + ``b`` as length 2 and returned 2.0 instead of 64.0. + """ + out = np.zeros(1) + two_arrays_from_reassigned_size(np.int64(64), out) + assert np.isclose(out[0], 64.0) + + +def test_a_size_reused_as_an_index_does_not_rebind_the_extent(): + """The shape's symbol must differ from the one a later index access of the same name binds. + + Both a shape and an index promote the size scalar to a symbol; if they share it, indexing with + the reassigned value re-binds the array's extent (here to 2), so ``a`` is allocated too small + and the access goes out of bounds. + """ + out = np.zeros(1) + size_reused_as_index(out) + assert np.isclose(out[0], 2.0) + + def test_promotion_leaves_the_descriptor_in_place(): sdfg = size_read_after_use.to_sdfg(simplify=False) - promoted = [s for s in sdfg.symbols if s.startswith('__sym_')] - assert promoted, 'the size scalar must be read into a symbol' - # The descriptor stays: deleting it is what broke later reads of the size. - assert any(s[len('__sym_'):] in sdfg.arrays for s in promoted) + # Each promotion assigns `__sym_... = ` on an interstate edge; the scalar it reads must + # still be a data descriptor afterwards, since deleting it is what broke later reads of the size. + sources = { + rhs + for e in sdfg.all_interstate_edges() + for lhs, rhs in e.data.assignments.items() if lhs.startswith('__sym_') + } + assert sources, 'the size scalar must be read into a symbol' + assert all(src in sdfg.arrays for src in sources), 'the size descriptor must survive promotion' sdfg.validate() +def test_shape_stays_correct_through_simplify(): + """simplify() may rewrite the promotion, but the array must keep the right extent either way. + + This is the coverage the structural test cannot give: the original test asserted the descriptor + was deleted, which hid that a later read of the size crashed. Here the whole program is run once + unsimplified and once simplified, and both must agree with numpy. + """ + n, nt = 6, 9 + a = np.arange(n, dtype=np.float64) + for simplify in (False, True): + sdfg = size_from_empty.to_sdfg(simplify=simplify) + out = np.zeros(n) + sdfg(a=a, Nt=np.int64(nt), out=out, N=n) + assert np.allclose(out, a * 2.0), f'wrong result with simplify={simplify}' + + if __name__ == '__main__': test_scalar_size_as_shape() test_size_descriptor_survives_its_use_as_a_shape() test_size_can_be_reassigned_after_use_as_a_shape() + test_two_arrays_from_a_reassigned_size_keep_their_own_extents() + test_a_size_reused_as_an_index_does_not_rebind_the_extent() test_promotion_leaves_the_descriptor_in_place() + test_shape_stays_correct_through_simplify() diff --git a/tests/transpose_unit_dim_squeeze_test.py b/tests/transpose_unit_dim_squeeze_test.py index f6cb22a2ef..6050fb187a 100644 --- a/tests/transpose_unit_dim_squeeze_test.py +++ b/tests/transpose_unit_dim_squeeze_test.py @@ -1,3 +1,4 @@ +# Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. """A 2D array with a unit dim (``(N, 1)`` / ``(1, N)``) must transpose to the swapped shape: the DaCe frontend used to squeeze the unit dim and reject ``(N, 1).T`` as "not a matrix". An integer index, by contrast, squeezes its axis (``x[:, 1]`` is ``(N,)``) per numpy semantics.""" From d67d7f0817481ab61a3b0811ebdd786c0f581891 Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Tue, 21 Jul 2026 00:48:06 +0200 Subject: [PATCH 10/17] Trim PR comments to Sphinx style and reuse dace utilities Shorten the docstrings and inline comments this PR added to DaCe's Sphinx conventions, and replace two hand-rolled helpers with existing utilities: find_new_name for the fresh promotion-symbol suffix, and symbolic.symlist for a shape's free symbols. Read the GEMM C operand through _matrix_operand like A and B rather than indexing the raw operand tuple. --- dace/frontend/python/newast.py | 51 ++++++++----------- .../replacements/array_creation_dace.py | 30 ++++------- .../python/replacements/array_manipulation.py | 7 ++- dace/libraries/blas/nodes/gemm.py | 7 +-- dace/libraries/blas/nodes/matmul.py | 19 +++---- tests/library/matmul_unit_dim_squeeze_test.py | 22 +++----- tests/size_scalar_shape_promotion_test.py | 26 ++++------ tests/transpose_unit_dim_squeeze_test.py | 6 +-- 8 files changed, 67 insertions(+), 101 deletions(-) diff --git a/dace/frontend/python/newast.py b/dace/frontend/python/newast.py index 4667ece3fc..26a6096686 100644 --- a/dace/frontend/python/newast.py +++ b/dace/frontend/python/newast.py @@ -33,7 +33,7 @@ from dace.sdfg.replace import replace_datadesc_names from dace.sdfg.type_inference import infer_expr_type from dace.symbolic import pystr_to_symbolic, inequal_symbols -from dace.utils import until +from dace.utils import until, find_new_name import numpy import sympy @@ -5377,48 +5377,39 @@ def promote_scalar_to_symbol(self, scalar: str, key: Optional[str] = None, fresh """ Reads a scalar into a symbol so its value can be used where only symbols are allowed. - A ``__sym_`` symbol is assigned from the scalar on an interstate edge, capturing the - scalar's value at that point in the program. The scalar's data descriptor is deliberately - left in place: it may still be read or reassigned later, and the visitor's own bookkeeping - still refers to it by name. + The symbol is assigned from the scalar on an interstate edge; the scalar's descriptor is + left in place so it may still be read or reassigned afterwards. :param scalar: Name of the scalar data descriptor to read. - :param key: Cache key for the promotion, defaulting to the scalar name. Repeated promotions - of the same expression reuse the symbol instead of minting a new one. - :param fresh: Mint a new symbol every call instead of reusing a cached one. A size scalar - is materialized into a temporary that the frontend reuses across calls - (``m = ...; np.empty(m); m = ...; np.empty(m)`` promotes one temporary at two - values); reusing one symbol would give both arrays the last value written to - it, so each shape must capture its own. + :param key: Cache key for the promotion; repeated promotions of the same expression reuse + the symbol. Defaults to the scalar name. + :param fresh: Mint a new (suffixed) symbol every call instead of reusing a cached one, so + two shapes sized from the same reassigned scalar keep their own values. :return: The symbol carrying the scalar's value. """ key = key if key is not None else scalar desc = self.sdfg.arrays[scalar] sym = None if fresh else self.indirections.get(key) if sym is None: + base = f'__sym_{scalar}' if fresh: - # A fresh promotion always carries a suffix, so it can never land on the bare - # `__sym_` a cached promotion of the same scalar reuses (an index access - # after the size is reassigned would otherwise re-bind the array's extent symbol). - # Check `arrays` too: a colliding data name would make `add_symbol` raise. - i = 0 - name = f'__sym_{scalar}_{i}' - while name in self.sdfg.symbols or name in self.sdfg.arrays: - i += 1 - name = f'__sym_{scalar}_{i}' + # Reserve ``base`` so a fresh promotion is always suffixed and never lands on the + # bare name a cached promotion reuses; else an index after the size is reassigned + # would re-bind the extent. The name is then free, so add_symbol cannot clash. + reserved = self.sdfg.symbols.keys() | self.sdfg.arrays.keys() | {base} + name = self.sdfg.add_symbol(find_new_name(base, reserved), desc.dtype) else: - name = f'__sym_{scalar}' + name = base + try: + self.sdfg.add_symbol(name, desc.dtype) + except FileExistsError: + # A cached promotion may re-add an existing symbol; that is benign. + pass sym = dace.symbol(name, dtype=desc.dtype) if not fresh: self.indirections[key] = sym - try: - self.sdfg.add_symbol(name, desc.dtype) - except FileExistsError: - # A cached (non-fresh) promotion may re-add an existing symbol; that is benign. - pass - # A promoted size can end up in an array shape, and nested scopes resolve the free - # symbols of their scope arrays through `globals`, so the symbol has to be visible - # there as well as on the SDFG. + # Nested scopes resolve their scope arrays' free symbols through ``globals``, so the + # symbol must be visible there too, not only on the SDFG. self.globals[str(sym)] = sym state = self._add_state(f'promote_{scalar}_to_{str(sym)}') edge = state.parent_graph.in_edges(state)[0] diff --git a/dace/frontend/python/replacements/array_creation_dace.py b/dace/frontend/python/replacements/array_creation_dace.py index 40a7233cca..561e1a5aa8 100644 --- a/dace/frontend/python/replacements/array_creation_dace.py +++ b/dace/frontend/python/replacements/array_creation_dace.py @@ -22,35 +22,25 @@ def promote_size_scalars_in_shape(pv: ProgramVisitor, sdfg: SDFG, shape: Shape) """ Rewrites a shape so that a size scalar used as an extent is read through a symbol. - An extent must be a symbol, but a size computed in the program (``nt = Nt + 1`` then - ``np.empty(nt)``) is a size-1 data descriptor. Each such name is read into a symbol on an - interstate edge and substituted into the shape, leaving the descriptor in place so the program - may keep reading the size afterwards. + An extent must be a symbol, but a size computed in the program (``nt = Nt + 1; np.empty(nt)``) + is a size-1 descriptor. Each such name is read into a fresh per-shape symbol and substituted + into the shape, leaving the descriptor in place. A symbol per shape keeps two arrays sized from + the same reassigned scalar from collapsing onto one value. - The size is captured into a symbol minted for this shape (``fresh``): the frontend reuses one - temporary for a size expression across calls, so ``m = ...; np.empty(m); m = ...; np.empty(m)`` - would otherwise give both arrays whichever value was written to the shared symbol last. A - per-shape symbol records the value at the point each array is created. - - :param pv: The program visitor, which owns symbol promotion and the state machine. + :param pv: The program visitor, owning symbol promotion and the state machine. :param sdfg: The SDFG being built. :param shape: The requested shape. :return: The shape with scalar extents replaced by symbols. """ resolved = [symbolic.pystr_to_symbolic(e) if isinstance(e, str) else e for e in shape] - names = set() - for extent in resolved: - if not isinstance(extent, sympy.Basic): - continue - for sym in extent.free_symbols: - name = str(sym) - if name in sdfg.arrays and name not in sdfg.symbols and sdfg.arrays[name].total_size == 1: - names.add(name) + names = [ + n for n in symbolic.symlist(resolved) + if n in sdfg.arrays and n not in sdfg.symbols and sdfg.arrays[n].total_size == 1 + ] if not names: return shape - # One symbol per distinct name in this shape (np.empty((m, m)) shares it); sorted() keeps the - # promotion states deterministic when several size scalars appear. + # One symbol per distinct name; sorted() keeps the promotion states deterministic. replacements = {symbolic.pystr_to_symbolic(n): pv.promote_scalar_to_symbol(n, fresh=True) for n in sorted(names)} return [e.subs(replacements) if isinstance(e, sympy.Basic) else e for e in resolved] diff --git a/dace/frontend/python/replacements/array_manipulation.py b/dace/frontend/python/replacements/array_manipulation.py index f52501f0c8..b96c0231b7 100644 --- a/dace/frontend/python/replacements/array_manipulation.py +++ b/dace/frontend/python/replacements/array_manipulation.py @@ -171,10 +171,9 @@ def _transpose(pv: ProgramVisitor, outname, arr2 = sdfg.add_transient(outname, new_shape, restype, arr1.storage, find_new_name=True) if axes == (1, 0): # 2D transposition - # The Transpose library node squeezes a unit axis to a vector and then rejects it as "not a - # matrix", so a ``(N, 1)`` / ``(1, N)`` array cannot use it. Fall back to a plain index-swap - # copy (``out[j, i] = in[i, j]``) whenever an extent is 1; it is general over 2D and - # stride-safe. Genuine matrices keep the optimized library node. + # The Transpose library node squeezes a unit axis and rejects it as "not a matrix", so an + # ``(N, 1)`` / ``(1, N)`` array copies with a swapped index instead; genuine matrices keep + # the library node. if 1 in arr1.shape: state.add_mapped_tasklet("transpose", map_ranges={ diff --git a/dace/libraries/blas/nodes/gemm.py b/dace/libraries/blas/nodes/gemm.py index ab0809d89a..3b8e0f2d0e 100644 --- a/dace/libraries/blas/nodes/gemm.py +++ b/dace/libraries/blas/nodes/gemm.py @@ -50,6 +50,7 @@ def make_sdfg(node, parent_state, parent_sdfg): adata, bdata, cdata = _get_matmul_operands(node, parent_state, parent_sdfg) edge_a, outer_array_a, shape_a, strides_a = _matrix_operand(adata) edge_b, outer_array_b, shape_b, strides_b = _matrix_operand(bdata) + _, outer_array_c, _, strides_c = _matrix_operand(cdata) dtype_a = outer_array_a.dtype.type dtype_b = outer_array_b.dtype.type @@ -81,7 +82,7 @@ def make_sdfg(node, parent_state, parent_sdfg): _, array_a = sdfg.add_array("_a", shape_a, dtype_a, strides=strides_a, storage=outer_array_a.storage) _, array_b = sdfg.add_array("_b", shape_b, dtype_b, strides=strides_b, storage=outer_array_b.storage) - _, array_c = sdfg.add_array("_c", shape_c, dtype_c, strides=_matrix_operand(cdata)[3], storage=cdata[1].storage) + _, array_c = sdfg.add_array("_c", shape_c, dtype_c, strides=strides_c, storage=outer_array_c.storage) if equal_valued(1, node.alpha): mul_program = "__out = __a * __b" @@ -466,8 +467,8 @@ class ExpandGemmPBLAS(ExpandTransformation): @staticmethod def expansion(node, state, sdfg): node.validate(sdfg, state) - # Read the same matrix view the dispatcher and validate agreed on; reading the raw subset - # would mis-size an operand (e.g. an (NQ, 1, NP) reshape) that validate has already accepted. + # Read the same matrix view validate accepted; the raw subset would mis-size an operand + # such as an (NQ, 1, NP) reshape. adata, bdata, _ = _get_matmul_operands(node, state, sdfg) _, adesc, ashape, _ = _matrix_operand(adata) _, bdesc, bshape, _ = _matrix_operand(bdata) diff --git a/dace/libraries/blas/nodes/matmul.py b/dace/libraries/blas/nodes/matmul.py index 1ba09109bc..6287ac5f50 100644 --- a/dace/libraries/blas/nodes/matmul.py +++ b/dace/libraries/blas/nodes/matmul.py @@ -9,13 +9,12 @@ def _matrix_subset_size(subset): """ - Returns an operand's size in the matrix view ``SpecializeMatMul`` matched on. + Returns an operand's size in the matrix view ``SpecializeMatMul`` matched on: the raw size if + already 2D, otherwise the squeezed size. - The dispatcher selects GEMM when an operand's raw subset is 2D *or* when it is 2D once - squeezed, so validation, expansion and code generation all have to read that same view. - Reading only the raw subset rejects a ``(NQ, 1, NP)`` reshape; reading only the squeezed one - rejects a genuine unit extent such as an ``(N, 1)`` column, which squeezes to a vector. - Disagreeing with the dispatcher either way makes GEMM refuse the operand it was just handed. + The dispatcher selects GEMM when an operand is 2D raw *or* 2D once squeezed, so validation, + expansion and codegen must read that same view. Reading only the raw subset rejects a + ``(NQ, 1, NP)`` reshape; reading only the squeezed one rejects a genuine ``(N, 1)`` column. :param subset: The subset of the memlet accessing the operand. :return: The 2D size if either view supplies one, otherwise the squeezed size. @@ -30,13 +29,11 @@ def _matrix_subset_size(subset): def _matrix_operand(operand): """ - Returns a GEMM operand as ``(edge, descriptor, shape, strides)`` in its matrix view. - - Applies the rule of :func:`_matrix_subset_size` to an operand whose two views have already - been computed, so that the strides are picked from the same view as the shape. + Returns a GEMM operand as ``(edge, descriptor, shape, strides)`` in its matrix view, applying + the rule of :func:`_matrix_subset_size` so shape and strides come from the same view. :param operand: One of the three tuples returned by :func:`_get_matmul_operands`. - :return: The edge, the outer descriptor, and the 2D shape and strides. + :return: The edge, outer descriptor, and 2D shape and strides. """ edge, desc, size, strides, squeezed_size, squeezed_strides = operand if len(size) == 2: diff --git a/tests/library/matmul_unit_dim_squeeze_test.py b/tests/library/matmul_unit_dim_squeeze_test.py index 4c9a57ca87..f298c4f0ba 100644 --- a/tests/library/matmul_unit_dim_squeeze_test.py +++ b/tests/library/matmul_unit_dim_squeeze_test.py @@ -1,15 +1,11 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. -"""``SpecializeMatMul`` and ``Gemm`` have to read the same view of an operand. +"""``SpecializeMatMul`` and ``Gemm`` must read the same view of an operand. -The dispatcher matches on the squeezed operand sizes, so ``np.reshape(x, (NQ, 1, NP)) @ C4`` is -routed to ``Gemm`` as an ``(NQ, NP) @ (NP, NP)`` product. ``Gemm.validate`` and the GEMM code -generator used to re-read the raw subset instead, see rank 3, and reject the operand the dispatcher -had just accepted -- "matrix-matrix product only supported on matrices". npbench's doitgen hits this -at every size. - -Collapsing the unit dimension is exact rather than convenient: it is the row count of an -``NQ``-long contiguous batch, so the product is genuinely one GEMM. Keeping it on ``Gemm`` also -keeps ``alpha``, ``beta`` and a summation WCR, none of which ``BatchedMatMul`` honours. +The dispatcher matches on squeezed sizes, so ``np.reshape(x, (NQ, 1, NP)) @ C4`` routes to ``Gemm`` +as ``(NQ, NP) @ (NP, NP)``. ``Gemm.validate`` and the GEMM codegen used to re-read the raw subset, +see rank 3, and reject the operand -- "matrix-matrix product only supported on matrices". npbench's +doitgen hits this. Collapsing the unit dim is exact (one GEMM), and keeping it on ``Gemm`` preserves +``alpha``, ``beta`` and the WCR that ``BatchedMatMul`` drops. """ import numpy as np import pytest @@ -43,8 +39,7 @@ def reference(A, C4): def initialize(nr, nq, np_, seed=0): - # Random rather than polybench's ((i*j+k) % NP)/NP, which degenerates to all-zero when NP == 1 - # and would make the numeric assertions vacuous for the all-unit-extent shape. + # Random, not polybench's ((i*j+k) % NP)/NP, which is all-zero at NP == 1 (a vacuous compare). rng = np.random.default_rng(seed) return rng.random((nr, nq, np_)), rng.random((np_, np_)) @@ -60,8 +55,7 @@ def specialize_matmuls(sdfg): node.expand(state) -# NQ == 1 collapses a second dimension when squeezed, which is where a "pick whichever view looks -# two-dimensional" heuristic breaks; both such shapes are covered. +# NQ == 1 collapses a second dim when squeezed, where a "whichever view looks 2D" heuristic breaks. SIZES = [(3, 4, 5), (1, 1, 1), (8, 10, 12), (5, 1, 7)] diff --git a/tests/size_scalar_shape_promotion_test.py b/tests/size_scalar_shape_promotion_test.py index d1dc38aa25..0c51d60241 100644 --- a/tests/size_scalar_shape_promotion_test.py +++ b/tests/size_scalar_shape_promotion_test.py @@ -1,12 +1,10 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. """A size computed in the program can be used as an array shape. -``nt = Nt + 1`` materializes ``nt`` as a scalar data descriptor, but an array extent has to be a -symbol, and minting a symbol of the same name collided with the descriptor (``FileExistsError``). -The size is now read into a ``__sym_`` symbol on an interstate edge and substituted into the shape, -leaving the descriptor in place so the program can keep reading or reassigning the size afterwards. -Each shape captures its own symbol, so two arrays sized from the same reused variable at different -values keep the values they were created with. +``nt = Nt + 1; np.empty(nt)`` needs ``nt`` as a symbol, but it is a data descriptor. The size is +read into a ``__sym_`` symbol on an interstate edge and substituted into the shape, leaving the +descriptor in place so it can still be read or reassigned. Each shape captures its own symbol, so +two arrays sized from the same reused name keep their own extents. """ import numpy as np import dace @@ -90,10 +88,9 @@ def test_size_can_be_reassigned_after_use_as_a_shape(): def test_two_arrays_from_a_reassigned_size_keep_their_own_extents(): - """A per-shape symbol: reusing one size name for two arrays must not collapse their extents. + """Reusing one size name for two arrays must not collapse their extents. - A single shared symbol gave both arrays the last value written to it, so ``np.sum(b)`` read - ``b`` as length 2 and returned 2.0 instead of 64.0. + A single shared symbol gave both the last value written, so ``np.sum(b)`` returned 2.0 not 64.0. """ out = np.zeros(1) two_arrays_from_reassigned_size(np.int64(64), out) @@ -103,8 +100,7 @@ def test_two_arrays_from_a_reassigned_size_keep_their_own_extents(): def test_a_size_reused_as_an_index_does_not_rebind_the_extent(): """The shape's symbol must differ from the one a later index access of the same name binds. - Both a shape and an index promote the size scalar to a symbol; if they share it, indexing with - the reassigned value re-binds the array's extent (here to 2), so ``a`` is allocated too small + Sharing it re-binds the array's extent to the reassigned value (here 2), so ``a`` is too small and the access goes out of bounds. """ out = np.zeros(1) @@ -114,8 +110,8 @@ def test_a_size_reused_as_an_index_does_not_rebind_the_extent(): def test_promotion_leaves_the_descriptor_in_place(): sdfg = size_read_after_use.to_sdfg(simplify=False) - # Each promotion assigns `__sym_... = ` on an interstate edge; the scalar it reads must - # still be a data descriptor afterwards, since deleting it is what broke later reads of the size. + # The scalar read by each ``__sym_... = `` assignment must survive as a descriptor; + # deleting it is what broke later reads of the size. sources = { rhs for e in sdfg.all_interstate_edges() @@ -129,9 +125,7 @@ def test_promotion_leaves_the_descriptor_in_place(): def test_shape_stays_correct_through_simplify(): """simplify() may rewrite the promotion, but the array must keep the right extent either way. - This is the coverage the structural test cannot give: the original test asserted the descriptor - was deleted, which hid that a later read of the size crashed. Here the whole program is run once - unsimplified and once simplified, and both must agree with numpy. + Run once unsimplified and once simplified; both must agree with numpy. """ n, nt = 6, 9 a = np.arange(n, dtype=np.float64) diff --git a/tests/transpose_unit_dim_squeeze_test.py b/tests/transpose_unit_dim_squeeze_test.py index 6050fb187a..b5efad90df 100644 --- a/tests/transpose_unit_dim_squeeze_test.py +++ b/tests/transpose_unit_dim_squeeze_test.py @@ -1,7 +1,7 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. -"""A 2D array with a unit dim (``(N, 1)`` / ``(1, N)``) must transpose to the swapped shape: the -DaCe frontend used to squeeze the unit dim and reject ``(N, 1).T`` as "not a matrix". An integer -index, by contrast, squeezes its axis (``x[:, 1]`` is ``(N,)``) per numpy semantics.""" +"""A 2D array with a unit dim (``(N, 1)`` / ``(1, N)``) must transpose to the swapped shape; the +frontend used to squeeze it and reject ``(N, 1).T`` as "not a matrix". An integer index, by +contrast, squeezes its axis (``x[:, 1]`` is ``(N,)``) per numpy.""" import numpy as np import dace From b8c664f4f75796b815bc6c80d80b593ffe8694a4 Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Tue, 21 Jul 2026 11:09:22 +0200 Subject: [PATCH 11/17] fix(frontend): keep subscript promotions out of the visitor globals Routing every subscript promotion through promote_scalar_to_symbol published its __sym_ name in ProgramVisitor.globals, which the original _promote never did. The leaked names shadowed name resolution and broke 25 frontend tests with KeyError: '__sym_...'. Only shape promotions need globals visibility, so gate the write on fresh. --- dace/frontend/python/newast.py | 27 ++++++++----------- .../replacements/array_creation_dace.py | 9 +++---- 2 files changed, 15 insertions(+), 21 deletions(-) diff --git a/dace/frontend/python/newast.py b/dace/frontend/python/newast.py index 26a6096686..ea4d639b30 100644 --- a/dace/frontend/python/newast.py +++ b/dace/frontend/python/newast.py @@ -5375,16 +5375,12 @@ def range_is_index(range: subsets.Range) -> bool: def promote_scalar_to_symbol(self, scalar: str, key: Optional[str] = None, fresh: bool = False) -> symbolic.symbol: """ - Reads a scalar into a symbol so its value can be used where only symbols are allowed. - - The symbol is assigned from the scalar on an interstate edge; the scalar's descriptor is - left in place so it may still be read or reassigned afterwards. + Reads a scalar into a symbol on an interstate edge, leaving its descriptor in place. :param scalar: Name of the scalar data descriptor to read. - :param key: Cache key for the promotion; repeated promotions of the same expression reuse - the symbol. Defaults to the scalar name. - :param fresh: Mint a new (suffixed) symbol every call instead of reusing a cached one, so - two shapes sized from the same reassigned scalar keep their own values. + :param key: Cache key; repeated promotions of the same expression reuse the symbol. + :param fresh: Mint a new suffixed symbol instead of reusing a cached one, so two shapes + sized from the same reassigned scalar keep their own values. :return: The symbol carrying the scalar's value. """ key = key if key is not None else scalar @@ -5393,9 +5389,8 @@ def promote_scalar_to_symbol(self, scalar: str, key: Optional[str] = None, fresh if sym is None: base = f'__sym_{scalar}' if fresh: - # Reserve ``base`` so a fresh promotion is always suffixed and never lands on the - # bare name a cached promotion reuses; else an index after the size is reassigned - # would re-bind the extent. The name is then free, so add_symbol cannot clash. + # Reserve ``base`` so a fresh promotion never lands on the bare name a cached one + # reuses; a later index on the same scalar would otherwise re-bind the extent. reserved = self.sdfg.symbols.keys() | self.sdfg.arrays.keys() | {base} name = self.sdfg.add_symbol(find_new_name(base, reserved), desc.dtype) else: @@ -5403,14 +5398,14 @@ def promote_scalar_to_symbol(self, scalar: str, key: Optional[str] = None, fresh try: self.sdfg.add_symbol(name, desc.dtype) except FileExistsError: - # A cached promotion may re-add an existing symbol; that is benign. - pass + pass # A cached promotion may re-add an existing symbol. sym = dace.symbol(name, dtype=desc.dtype) if not fresh: self.indirections[key] = sym - # Nested scopes resolve their scope arrays' free symbols through ``globals``, so the - # symbol must be visible there too, not only on the SDFG. - self.globals[str(sym)] = sym + else: + # Shape symbols must resolve inside nested scopes, which look up free symbols in + # ``globals``. Subscript promotions must stay out: they shadow names there. + self.globals[str(sym)] = sym state = self._add_state(f'promote_{scalar}_to_{str(sym)}') edge = state.parent_graph.in_edges(state)[0] edge.data.assignments = {str(sym): scalar} diff --git a/dace/frontend/python/replacements/array_creation_dace.py b/dace/frontend/python/replacements/array_creation_dace.py index 561e1a5aa8..c135c1da94 100644 --- a/dace/frontend/python/replacements/array_creation_dace.py +++ b/dace/frontend/python/replacements/array_creation_dace.py @@ -22,12 +22,11 @@ def promote_size_scalars_in_shape(pv: ProgramVisitor, sdfg: SDFG, shape: Shape) """ Rewrites a shape so that a size scalar used as an extent is read through a symbol. - An extent must be a symbol, but a size computed in the program (``nt = Nt + 1; np.empty(nt)``) - is a size-1 descriptor. Each such name is read into a fresh per-shape symbol and substituted - into the shape, leaving the descriptor in place. A symbol per shape keeps two arrays sized from - the same reassigned scalar from collapsing onto one value. + A size computed in the program (``nt = Nt + 1; np.empty(nt)``) is a size-1 descriptor, but an + extent must be a symbol. One fresh symbol per shape keeps two arrays sized from the same + reassigned scalar from collapsing onto one value. - :param pv: The program visitor, owning symbol promotion and the state machine. + :param pv: The program visitor. :param sdfg: The SDFG being built. :param shape: The requested shape. :return: The shape with scalar extents replaced by symbols. From 6481b353299c9ad13139dc892b340df9207da710 Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Tue, 21 Jul 2026 11:30:31 +0200 Subject: [PATCH 12/17] fix(codegen): allocate data after the symbol sizing it is assigned A transient sized by a symbol that a control flow region's incoming edge assigns had no legal allocation point. simplify() moves the assignment onto an edge leaving the closest dominating state, and loops and conditionals are excluded as dominators, so the allocation was emitted in that state, ahead of the assignment. The array was then allocated at size zero and the first write corrupted the heap. Bracket every region with an empty state, so the assignment lands on the leading state's incoming edge and the allocation follows it. The trailing state also gives region-scoped deallocations somewhere to go; they were dropped before, leaking the array. --- dace/codegen/codegen.py | 4 + dace/transformation/passes/__init__.py | 1 + .../passes/region_boundary_states.py | 49 ++++++++++++ tests/passes/region_boundary_states_test.py | 76 +++++++++++++++++++ tests/size_scalar_shape_promotion_test.py | 16 ++++ 5 files changed, 146 insertions(+) create mode 100644 dace/transformation/passes/region_boundary_states.py create mode 100644 tests/passes/region_boundary_states_test.py diff --git a/dace/codegen/codegen.py b/dace/codegen/codegen.py index fc6791599f..019c929c86 100644 --- a/dace/codegen/codegen.py +++ b/dace/codegen/codegen.py @@ -17,6 +17,7 @@ from dace.codegen.instrumentation import InstrumentationProvider from dace.sdfg.state import SDFGState from dace.transformation.pass_pipeline import FixedPointPipeline +from dace.transformation.passes.region_boundary_states import RegionBoundaryStates from dace.transformation.passes.simplification.control_flow_raising import ControlFlowRaising @@ -196,6 +197,9 @@ def generate_code(sdfg: SDFG, validate=True) -> List[CodeObject]: # where explicit control flow was not used, continue to work as expected. FixedPointPipeline([ControlFlowRaising()]).apply_pass(sdfg, {}) + # Data sized by a symbol that a region's incoming edge assigns can only be allocated once the region is entered. + RegionBoundaryStates().apply_pass(sdfg, {}) + # Before generating the code, run type inference on the SDFG connectors infer_types.infer_connector_types(sdfg) diff --git a/dace/transformation/passes/__init__.py b/dace/transformation/passes/__init__.py index 8d0c023a51..0b8649afd6 100644 --- a/dace/transformation/passes/__init__.py +++ b/dace/transformation/passes/__init__.py @@ -10,6 +10,7 @@ from .optional_arrays import OptionalArrayInference from .pattern_matching import PatternMatchAndApply, PatternMatchAndApplyRepeated, PatternApplyOnceEverywhere from .prune_symbols import RemoveUnusedSymbols +from .region_boundary_states import RegionBoundaryStates from .scalar_to_symbol import ScalarToSymbolPromotion from .simplify import SimplifyPass from .symbol_propagation import SymbolPropagation diff --git a/dace/transformation/passes/region_boundary_states.py b/dace/transformation/passes/region_boundary_states.py new file mode 100644 index 0000000000..f93dd4b728 --- /dev/null +++ b/dace/transformation/passes/region_boundary_states.py @@ -0,0 +1,49 @@ +# Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. + +from typing import Any, Dict, Optional + +from dace import properties +from dace.sdfg.sdfg import SDFG +from dace.sdfg.state import AbstractControlFlowRegion, ConditionalBlock +from dace.transformation import pass_pipeline as ppl +from dace.transformation import transformation + + +@properties.make_properties +@transformation.explicit_cf_compatible +class RegionBoundaryStates(ppl.Pass): + """ + Brackets every control flow region with an empty state in its parent region. + + Allocation is only emitted inside a state. Without the leading state, data sized by a symbol that the region's + incoming edge assigns is allocated at the region's predecessor, before the symbol is defined. The trailing state + gives the matching deallocation a place to go. + """ + + CATEGORY: str = 'Helper' + + def modifies(self) -> ppl.Modifies: + return ppl.Modifies.CFG + + def should_reapply(self, modified: ppl.Modifies) -> bool: + # Reapplying would bracket the brackets, so the pass runs once. + return False + + def apply_pass(self, sdfg: SDFG, _: Dict[str, Any]) -> Optional[int]: + """ + :param sdfg: The SDFG to modify in-place. + :return: Number of states inserted, or None if unchanged. + """ + inserted = 0 + # Branches of a conditional are entered without an inter-state edge, so they need no boundary. + regions = [ + cfg for cfg in sdfg.all_control_flow_regions(recursive=True) if not isinstance(cfg, ConditionalBlock) + ] + for cfg in regions: + for node in list(cfg.nodes()): + if not isinstance(node, AbstractControlFlowRegion): + continue + cfg.add_state_before(node, is_start_block=node is cfg.start_block) + cfg.add_state_after(node) + inserted += 2 + return inserted or None diff --git a/tests/passes/region_boundary_states_test.py b/tests/passes/region_boundary_states_test.py new file mode 100644 index 0000000000..c0e3f41137 --- /dev/null +++ b/tests/passes/region_boundary_states_test.py @@ -0,0 +1,76 @@ +# Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. +""" Tests bracketing of control flow regions with states. """ +import numpy as np + +import dace +from dace.sdfg.state import AbstractControlFlowRegion, ConditionalBlock +from dace.transformation.passes.region_boundary_states import RegionBoundaryStates + +N = dace.symbol('N') + + +@dace.program +def loop_and_branch(a: dace.float64[N], out: dace.float64[N]): + for i in range(N): + if a[i] > 0.0: + out[i] = a[i] + else: + out[i] = -a[i] + + +def regions_of(sdfg): + for cfg in sdfg.all_control_flow_regions(): + if isinstance(cfg, ConditionalBlock): + continue + for node in cfg.nodes(): + if isinstance(node, AbstractControlFlowRegion): + yield cfg, node + + +def test_every_region_is_bracketed(): + sdfg = loop_and_branch.to_sdfg(simplify=True) + assert any(True for _ in regions_of(sdfg)) # guard against a vacuous check + + assert RegionBoundaryStates().apply_pass(sdfg, {}) > 0 + for cfg, region in regions_of(sdfg): + assert all(isinstance(e.src, dace.SDFGState) for e in cfg.in_edges(region)) + assert all(isinstance(e.dst, dace.SDFGState) for e in cfg.out_edges(region)) + sdfg.validate() + + +def test_leading_region_keeps_start_block(): + """Bracketing a region that starts its parent must hand the start block over, not orphan it.""" + sdfg = dace.SDFG('leading_region') + sdfg.add_array('out', [1], dace.float64) + loop = dace.sdfg.state.LoopRegion('loop', 'i < 4', 'i', 'i = 0', 'i = i + 1') + sdfg.add_node(loop, is_start_block=True) + state = loop.add_state('body', is_start_block=True) + tasklet = state.add_tasklet('one', {}, {'o'}, 'o = 1.0') + state.add_edge(tasklet, 'o', state.add_write('out'), None, dace.Memlet('out[0]')) + + RegionBoundaryStates().apply_pass(sdfg, {}) + assert isinstance(sdfg.start_block, dace.SDFGState) + assert sdfg.start_block is not loop + sdfg.validate() + + out = np.zeros(1) + sdfg(out=out) + assert out[0] == 1.0 + + +def test_result_is_unchanged(): + """The pass only inserts empty states, so it must not alter what the program computes.""" + a = np.random.default_rng(0).random(16) - 0.5 + expected = np.abs(a) + + sdfg = loop_and_branch.to_sdfg(simplify=True) + RegionBoundaryStates().apply_pass(sdfg, {}) + out = np.zeros(16) + sdfg(a=a, out=out, N=16) + assert np.allclose(out, expected) + + +if __name__ == '__main__': + test_every_region_is_bracketed() + test_leading_region_keeps_start_block() + test_result_is_unchanged() diff --git a/tests/size_scalar_shape_promotion_test.py b/tests/size_scalar_shape_promotion_test.py index 0c51d60241..0311f756a2 100644 --- a/tests/size_scalar_shape_promotion_test.py +++ b/tests/size_scalar_shape_promotion_test.py @@ -136,6 +136,21 @@ def test_shape_stays_correct_through_simplify(): assert np.allclose(out, a * 2.0), f'wrong result with simplify={simplify}' +def test_size_symbol_is_assigned_before_the_allocation(): + """simplify() moves the promotion onto an edge out of the allocation's dominator. + + The allocation then read the symbol undefined and sized the array at 0, corrupting the heap on the first write. + """ + sdfg = size_from_empty.to_sdfg(simplify=True) + sym = str(next(iter(sdfg.arrays['b'].free_symbols))) + lines = sdfg.generate_code()[0].clean_code.splitlines() + + alloc = next(i for i, line in enumerate(lines) if 'new double' in line and sym in line) + assign = next(i for i, line in enumerate(lines) if line.strip().startswith(f'{sym} = ')) + assert assign < alloc + assert any('delete[] b' in line for line in lines), 'the array is never freed' + + if __name__ == '__main__': test_scalar_size_as_shape() test_size_descriptor_survives_its_use_as_a_shape() @@ -144,3 +159,4 @@ def test_shape_stays_correct_through_simplify(): test_a_size_reused_as_an_index_does_not_rebind_the_extent() test_promotion_leaves_the_descriptor_in_place() test_shape_stays_correct_through_simplify() + test_size_symbol_is_assigned_before_the_allocation() From d4791437bdc20349ed07704a69e9df438d5efcdf Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Tue, 21 Jul 2026 12:29:37 +0200 Subject: [PATCH 13/17] fix: keep nested-scope outputs in the GPU kernel signature unordered_arglist followed an exit-node connector only one hop, so when the data leaves through several nested scopes the inner memlet names a local transient and the outermost array never enters the argument list. GPU codegen tiles device maps into nested scopes, so the kernel referenced an array and a stride symbol its signature never declared. Follow each path to its last edge, mirroring the read side. --- dace/sdfg/state.py | 13 +++++++------ tests/codegen/argument_signature_test.py | 9 +++++++++ 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/dace/sdfg/state.py b/dace/sdfg/state.py index f9b288b3ae..fec4efd815 100644 --- a/dace/sdfg/state.py +++ b/dace/sdfg/state.py @@ -922,15 +922,16 @@ def unordered_arglist(self, elif isinstance(edge.dst, nd.ExitNode) and isinstance(edge.src, (nd.AccessNode, nd.CodeNode)): # Same case as above, but for outgoing Memlets. - # NOTE: We have to use a memlet tree here, because the data could potentially - # go to multiple sources. We have to do it this way, because if we would call - # `memlet_tree()` here, then we would just get the edge back. + # NOTE: The data may leave through several nested scopes before it lands, and only the + # outermost edge names where it goes, so each path is followed to its last edge. One + # connector can feed several of them, hence the loop. additional_descs = {} connector_to_look = "OUT_" + edge.dst_conn[3:] for oedge in self.graph.out_edges_by_connector(edge.dst, connector_to_look): - if ((not oedge.data.is_empty()) and (oedge.data.data not in descs) - and (oedge.data.data not in additional_descs)): - additional_descs[oedge.data.data] = sdfg.arrays[oedge.data.data] + outer_edge = self.graph.memlet_path(oedge)[-1] + if ((not outer_edge.data.is_empty()) and (outer_edge.data.data not in descs) + and (outer_edge.data.data not in additional_descs)): + additional_descs[outer_edge.data.data] = sdfg.arrays[outer_edge.data.data] else: # Case is ignored. diff --git a/tests/codegen/argument_signature_test.py b/tests/codegen/argument_signature_test.py index e4b720a289..27ab86f080 100644 --- a/tests/codegen/argument_signature_test.py +++ b/tests/codegen/argument_signature_test.py @@ -1,3 +1,5 @@ +import re + import dace @@ -187,6 +189,13 @@ def make_sdfg() -> dace.SDFG: assert isinstance(atype_res, atype_ref), f"Expected '{aname}' to have type {atype_ref}, but it had {type(atype_res)}." + # GPU codegen tiles the map into nested scopes, so `D` is only named on the outermost edge of its + # path. The kernel signature has to keep it, and its stride, all the same. + kernel = next(o.clean_code for o in sdfg.generate_code() if o.language == "cu") + signature = re.search(r"__global__ void .*\((.*)\)", kernel).group(1) + for arg in ("* __restrict__ D", "second_stride_D"): + assert arg in signature, f"Expected '{arg}' in the kernel signature, but got '{signature}'." + # If we have cupy we will also compile it. try: import cupy as cp # noqa: F401 From 0e0d79341f3bb80049886645fa3848f16930716f Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Tue, 21 Jul 2026 17:16:10 +0200 Subject: [PATCH 14/17] fix: read a size-1 array through a subscript and share the matrix view Review of the promotion and squeezing work turned up four defects. A size used as an extent is admitted when its descriptor holds one element, but the interstate edge assigned the container by name. For an array that assigns its pointer, so np.empty(nt) with nt: int64[1] failed to compile. Read those through nt[0]; a scalar still reads by name. The rule for an operand's matrix view - the raw subset when it is already 2D, the squeezed one otherwise - lived in matmul alone. Transpose squeezed unconditionally, so it rejected an (N, 1) column as not a matrix, and its operand helpers indexed size[1] after squeezing and would raise IndexError in every expansion. Move the rule to blas_helpers.matrix_view, which also returns the dimensions it kept so sizes and strides cannot disagree, and use it from both libraries. The frontend no longer needs to route unit-dimension transposition around the library node. numpy.zeros, ones and full build their transient on their own path and so rejected a computed size that numpy.empty accepted. Promotion opens a state to carry the symbol assignment, so the fill has to be placed after it rather than in the state the replacement was handed. _get_matmul_operands left res_out unbound, crashing instead of reporting a missing _c connector. --- dace/frontend/python/newast.py | 4 +- .../python/replacements/array_creation.py | 4 ++ .../replacements/array_creation_dace.py | 2 - .../python/replacements/array_manipulation.py | 16 +------- dace/libraries/blas/blas_helpers.py | 22 ++++++++++- dace/libraries/blas/nodes/matmul.py | 18 +++------ dace/libraries/linalg/nodes/transpose.py | 17 ++------ tests/passes/region_boundary_states_test.py | 23 ++++++++++- tests/size_scalar_shape_promotion_test.py | 39 +++++++++++++++++++ 9 files changed, 99 insertions(+), 46 deletions(-) diff --git a/dace/frontend/python/newast.py b/dace/frontend/python/newast.py index ea4d639b30..b0b1fb2a87 100644 --- a/dace/frontend/python/newast.py +++ b/dace/frontend/python/newast.py @@ -5408,7 +5408,9 @@ def promote_scalar_to_symbol(self, scalar: str, key: Optional[str] = None, fresh self.globals[str(sym)] = sym state = self._add_state(f'promote_{scalar}_to_{str(sym)}') edge = state.parent_graph.in_edges(state)[0] - edge.data.assignments = {str(sym): scalar} + # A Scalar reads by name; a size-1 array needs the subscript, or the assignment takes its pointer. + rhs = scalar if isinstance(desc, data.Scalar) else f'{scalar}[{", ".join(["0"] * len(desc.shape))}]' + edge.data.assignments = {str(sym): rhs} return sym def _parse_subscript_slice(self, diff --git a/dace/frontend/python/replacements/array_creation.py b/dace/frontend/python/replacements/array_creation.py index 66bcd57163..cfa15a07c2 100644 --- a/dace/frontend/python/replacements/array_creation.py +++ b/dace/frontend/python/replacements/array_creation.py @@ -6,6 +6,7 @@ from dace.frontend.common import op_repository as oprepo from dace.frontend.python.common import DaceSyntaxError from dace.frontend.python.replacements.utils import ProgramVisitor, Shape, sym_type, broadcast_together +from dace.frontend.python.replacements.array_creation_dace import promote_size_scalars_in_shape from dace.frontend.python.replacements.operators import result_type from dace import data, dtypes, symbolic, Memlet, SDFG, SDFGState @@ -65,6 +66,9 @@ def _numpy_full(pv: ProgramVisitor, if isinstance(shape, (Number, str)) or symbolic.issymbolic(shape): shape = [shape] + shape = promote_size_scalars_in_shape(pv, sdfg, shape) + # Promotion opens a state to carry the symbol assignment; the fill has to follow it. + state = pv.last_block if any(isinstance(s, str) for s in shape): raise DaceSyntaxError( pv, None, f'Data-dependent shape {shape} is currently not allowed. Only constants ' diff --git a/dace/frontend/python/replacements/array_creation_dace.py b/dace/frontend/python/replacements/array_creation_dace.py index c135c1da94..3a89cc1164 100644 --- a/dace/frontend/python/replacements/array_creation_dace.py +++ b/dace/frontend/python/replacements/array_creation_dace.py @@ -36,8 +36,6 @@ def promote_size_scalars_in_shape(pv: ProgramVisitor, sdfg: SDFG, shape: Shape) n for n in symbolic.symlist(resolved) if n in sdfg.arrays and n not in sdfg.symbols and sdfg.arrays[n].total_size == 1 ] - if not names: - return shape # One symbol per distinct name; sorted() keeps the promotion states deterministic. replacements = {symbolic.pystr_to_symbolic(n): pv.promote_scalar_to_symbol(n, fresh=True) for n in sorted(names)} diff --git a/dace/frontend/python/replacements/array_manipulation.py b/dace/frontend/python/replacements/array_manipulation.py index b96c0231b7..f414d239f3 100644 --- a/dace/frontend/python/replacements/array_manipulation.py +++ b/dace/frontend/python/replacements/array_manipulation.py @@ -170,21 +170,7 @@ def _transpose(pv: ProgramVisitor, outname = pv.get_target_name() outname, arr2 = sdfg.add_transient(outname, new_shape, restype, arr1.storage, find_new_name=True) - if axes == (1, 0): # 2D transposition - # The Transpose library node squeezes a unit axis and rejects it as "not a matrix", so an - # ``(N, 1)`` / ``(1, N)`` array copies with a swapped index instead; genuine matrices keep - # the library node. - if 1 in arr1.shape: - state.add_mapped_tasklet("transpose", - map_ranges={ - "__i": "0:%s" % arr1.shape[0], - "__j": "0:%s" % arr1.shape[1] - }, - inputs={"__inp": Memlet("%s[__i, __j]" % inpname)}, - code="__out = __inp", - outputs={"__out": Memlet("%s[__j, __i]" % outname)}, - external_edges=True) - return outname + if axes == (1, 0): # Special case for 2D transposition acc1 = state.add_read(inpname) acc2 = state.add_write(outname) import dace.libraries.linalg # Avoid import loop diff --git a/dace/libraries/blas/blas_helpers.py b/dace/libraries/blas/blas_helpers.py index 10a5052756..2e86d90649 100644 --- a/dace/libraries/blas/blas_helpers.py +++ b/dace/libraries/blas/blas_helpers.py @@ -1,7 +1,27 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import numpy as np +from copy import deepcopy as dc from dace import dtypes, data -from typing import Any, Dict, Tuple +from typing import Any, Dict, List, Tuple + + +def matrix_view(subset) -> Tuple[List[Any], List[int]]: + """ + Returns an operand's matrix view: the raw subset if it is already 2D, otherwise the squeezed one. + + Squeezing unconditionally rejects a genuine ``(N, 1)`` column as "not a matrix"; not squeezing at + all rejects an ``(NQ, 1, NP)`` reshape. Callers that read a size, a stride or a dimension index + must all use this view, or they disagree about which dimensions the operand has. + + :param subset: The subset of the memlet accessing the operand. + :return: The size in the matrix view, and the subset dimensions it kept. + """ + size = subset.size() + if len(size) == 2: + return size, list(range(len(size))) + squeezed = dc(subset) + dims = squeezed.squeeze() + return squeezed.size(), dims def to_blastype(dtype): diff --git a/dace/libraries/blas/nodes/matmul.py b/dace/libraries/blas/nodes/matmul.py index 6287ac5f50..23d66ffea8 100644 --- a/dace/libraries/blas/nodes/matmul.py +++ b/dace/libraries/blas/nodes/matmul.py @@ -1,6 +1,7 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. import dace from dace import properties, symbolic +from dace.libraries.blas.blas_helpers import matrix_view from copy import deepcopy as dc from typing import Any, Dict import warnings @@ -9,22 +10,12 @@ def _matrix_subset_size(subset): """ - Returns an operand's size in the matrix view ``SpecializeMatMul`` matched on: the raw size if - already 2D, otherwise the squeezed size. - - The dispatcher selects GEMM when an operand is 2D raw *or* 2D once squeezed, so validation, - expansion and codegen must read that same view. Reading only the raw subset rejects a - ``(NQ, 1, NP)`` reshape; reading only the squeezed one rejects a genuine ``(N, 1)`` column. + Returns an operand's size in the matrix view ``SpecializeMatMul`` matched on. :param subset: The subset of the memlet accessing the operand. - :return: The 2D size if either view supplies one, otherwise the squeezed size. + :return: The size in the matrix view. """ - size = subset.size() - if len(size) == 2: - return size - squeezed = dc(subset) - squeezed.squeeze() - return squeezed.size() + return matrix_view(subset)[0] def _matrix_operand(operand): @@ -45,6 +36,7 @@ def _get_matmul_operands(node, state, sdfg, name_lhs="_a", name_rhs="_b", name_o """Returns the matrix multiplication input edges, arrays, and shape.""" res_lhs = None res_rhs = None + res_out = None for edge in state.all_edges(node): if edge.dst_conn in [name_lhs, name_rhs]: size = edge.data.subset.size() diff --git a/dace/libraries/linalg/nodes/transpose.py b/dace/libraries/linalg/nodes/transpose.py index 7676f09a4a..88b1d3405a 100644 --- a/dace/libraries/linalg/nodes/transpose.py +++ b/dace/libraries/linalg/nodes/transpose.py @@ -1,6 +1,5 @@ # Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. import functools -from copy import deepcopy as dc import dace.library import dace.properties import dace.sdfg.nodes @@ -14,9 +13,7 @@ def _get_transpose_input(node, state, sdfg): """Returns the transpose input edge, array, and shape.""" for edge in state.in_edges(node): if edge.dst_conn == "_inp": - subset = dc(edge.data.subset) - idx = subset.squeeze() - size = subset.size() + size, idx = blas_helpers.matrix_view(edge.data.subset) outer_array = sdfg.data(dace.sdfg.find_input_arraynode(state, edge).data) return edge, outer_array, (size[0], size[1]), (outer_array.strides[idx[0]], outer_array.strides[idx[1]]) raise ValueError("Transpose input connector \"_inp\" not found.") @@ -26,9 +23,7 @@ def _get_transpose_output(node, state, sdfg): """Returns the transpose output edge, array, and shape.""" for edge in state.out_edges(node): if edge.src_conn == "_out": - subset = dc(edge.data.subset) - idx = subset.squeeze() - size = subset.size() + size, idx = blas_helpers.matrix_view(edge.data.subset) outer_array = sdfg.data(dace.sdfg.find_output_arraynode(state, edge).data) return edge, outer_array, (size[0], size[1]), (outer_array.strides[idx[0]], outer_array.strides[idx[1]]) raise ValueError("Transpose output connector \"_out\" not found.") @@ -248,18 +243,14 @@ def validate(self, sdfg, state): raise ValueError("Expected exactly one input to transpose operation") for _, _, _, dst_conn, memlet in state.in_edges(self): if dst_conn == '_inp': - subset = dc(memlet.subset) - subset.squeeze() - in_size = subset.size() + in_size, _ = blas_helpers.matrix_view(memlet.subset) out_edges = state.out_edges(self) if len(out_edges) != 1: raise ValueError("Expected exactly one output from transpose operation") out_memlet = out_edges[0].data if len(in_size) != 2: raise ValueError("Transpose operation only supported on matrices") - out_subset = dc(out_memlet.subset) - out_subset.squeeze() - out_size = out_subset.size() + out_size, _ = blas_helpers.matrix_view(out_memlet.subset) if len(out_size) != 2: raise ValueError("Transpose operation only supported on matrices") if list(out_size) != [in_size[1], in_size[0]]: diff --git a/tests/passes/region_boundary_states_test.py b/tests/passes/region_boundary_states_test.py index c0e3f41137..0625518301 100644 --- a/tests/passes/region_boundary_states_test.py +++ b/tests/passes/region_boundary_states_test.py @@ -18,8 +18,13 @@ def loop_and_branch(a: dace.float64[N], out: dace.float64[N]): out[i] = -a[i] +@dace.program +def calls_loop_and_branch(a: dace.float64[N], out: dace.float64[N]): + loop_and_branch(a, out) + + def regions_of(sdfg): - for cfg in sdfg.all_control_flow_regions(): + for cfg in sdfg.all_control_flow_regions(recursive=True): if isinstance(cfg, ConditionalBlock): continue for node in cfg.nodes(): @@ -38,6 +43,21 @@ def test_every_region_is_bracketed(): sdfg.validate() +def test_regions_inside_a_nested_sdfg_are_bracketed(): + """The pass descends into nested SDFGs, where the same allocation bug applies.""" + sdfg = calls_loop_and_branch.to_sdfg(simplify=False) + nested = [n for n, _ in sdfg.all_nodes_recursive() if isinstance(n, dace.nodes.NestedSDFG)] + assert nested, 'the fixture must keep a nested SDFG' + inner_regions = [(cfg, r) for cfg, r in regions_of(sdfg) if cfg.sdfg is not sdfg] + assert inner_regions, 'the nested SDFG must contain a region' + + RegionBoundaryStates().apply_pass(sdfg, {}) + for cfg, region in inner_regions: + assert all(isinstance(e.src, dace.SDFGState) for e in cfg.in_edges(region)) + assert all(isinstance(e.dst, dace.SDFGState) for e in cfg.out_edges(region)) + sdfg.validate() + + def test_leading_region_keeps_start_block(): """Bracketing a region that starts its parent must hand the start block over, not orphan it.""" sdfg = dace.SDFG('leading_region') @@ -72,5 +92,6 @@ def test_result_is_unchanged(): if __name__ == '__main__': test_every_region_is_bracketed() + test_regions_inside_a_nested_sdfg_are_bracketed() test_leading_region_keeps_start_block() test_result_is_unchanged() diff --git a/tests/size_scalar_shape_promotion_test.py b/tests/size_scalar_shape_promotion_test.py index 0311f756a2..a054321d29 100644 --- a/tests/size_scalar_shape_promotion_test.py +++ b/tests/size_scalar_shape_promotion_test.py @@ -7,6 +7,8 @@ two arrays sized from the same reused name keep their own extents. """ import numpy as np +import pytest + import dace N = dace.symbol('N') @@ -62,6 +64,13 @@ def size_reused_as_index(out: dace.float64[1]): out[0] = a[m] # the shape symbol must not be the one this index reassigns +@dace.program +def size_from_size_one_array(nt: dace.int64[1], out: dace.float64[4]): + b = np.empty(nt, dace.float64) + b[0] = 1.0 + out[0] = b[0] + + def test_scalar_size_as_shape(): n, nt = 5, 7 a = np.arange(n, dtype=np.float64) @@ -151,6 +160,33 @@ def test_size_symbol_is_assigned_before_the_allocation(): assert any('delete[] b' in line for line in lines), 'the array is never freed' +def test_a_size_one_array_is_read_through_a_subscript(): + """A size-1 array is a valid extent, but the assignment must read ``nt[0]``, not the pointer.""" + out = np.zeros(4) + size_from_size_one_array(np.array([4], dtype=np.int64), out) + assert np.allclose(out, [1.0, 0.0, 0.0, 0.0]) + + +@dace.program +def zeros_from_size(Nt: dace.int64, out: dace.float64[1]): + b = np.zeros(Nt + 1, dace.float64) + out[0] = np.sum(b) + + +@dace.program +def ones_from_size(Nt: dace.int64, out: dace.float64[1]): + b = np.ones(Nt + 1, dace.float64) + out[0] = np.sum(b) + + +@pytest.mark.parametrize('program,expected', [(zeros_from_size, 0.0), (ones_from_size, 4.0)]) +def test_the_fill_constructors_accept_a_computed_size(program, expected): + """zeros/ones/full build their transient on their own path, which also has to promote the size.""" + out = np.zeros(1) + program(np.int64(3), out) + assert np.isclose(out[0], expected) + + if __name__ == '__main__': test_scalar_size_as_shape() test_size_descriptor_survives_its_use_as_a_shape() @@ -160,3 +196,6 @@ def test_size_symbol_is_assigned_before_the_allocation(): test_promotion_leaves_the_descriptor_in_place() test_shape_stays_correct_through_simplify() test_size_symbol_is_assigned_before_the_allocation() + test_a_size_one_array_is_read_through_a_subscript() + test_the_fill_constructors_accept_a_computed_size(zeros_from_size, 0.0) + test_the_fill_constructors_accept_a_computed_size(ones_from_size, 4.0) From 7ec090cc481d3903f3e336177b1dd4785b0854c1 Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Tue, 21 Jul 2026 19:06:02 +0200 Subject: [PATCH 15/17] fix: make the fill constructors accept a size computed in the program np.zeros/ones/full built their transient directly, so a size held in a scalar was rejected while np.empty accepted it. Route them through promote_size_scalars_in_shape, which now reports whether it promoted: only then does the fill have to move to the state the promotion opened. The two tests asserting the rejection now assert the result. --- dace/frontend/python/replacements/array_creation.py | 7 ++++--- .../python/replacements/array_creation_dace.py | 12 +++++++----- tests/numpy/array_creation_test.py | 12 +++++------- 3 files changed, 16 insertions(+), 15 deletions(-) diff --git a/dace/frontend/python/replacements/array_creation.py b/dace/frontend/python/replacements/array_creation.py index cfa15a07c2..4f9d35998c 100644 --- a/dace/frontend/python/replacements/array_creation.py +++ b/dace/frontend/python/replacements/array_creation.py @@ -66,9 +66,10 @@ def _numpy_full(pv: ProgramVisitor, if isinstance(shape, (Number, str)) or symbolic.issymbolic(shape): shape = [shape] - shape = promote_size_scalars_in_shape(pv, sdfg, shape) - # Promotion opens a state to carry the symbol assignment; the fill has to follow it. - state = pv.last_block + shape, promoted = promote_size_scalars_in_shape(pv, sdfg, shape) + if promoted: + # Promotion opens a state to carry the symbol assignment; the fill has to follow it. + state = pv.last_block if any(isinstance(s, str) for s in shape): raise DaceSyntaxError( pv, None, f'Data-dependent shape {shape} is currently not allowed. Only constants ' diff --git a/dace/frontend/python/replacements/array_creation_dace.py b/dace/frontend/python/replacements/array_creation_dace.py index 3a89cc1164..29549156d6 100644 --- a/dace/frontend/python/replacements/array_creation_dace.py +++ b/dace/frontend/python/replacements/array_creation_dace.py @@ -10,7 +10,7 @@ from copy import deepcopy as dcpy from numbers import Integral -from typing import Any, Optional +from typing import Any, Optional, Tuple import sympy import numpy as np @@ -18,7 +18,7 @@ from dace import symbolic -def promote_size_scalars_in_shape(pv: ProgramVisitor, sdfg: SDFG, shape: Shape) -> Shape: +def promote_size_scalars_in_shape(pv: ProgramVisitor, sdfg: SDFG, shape: Shape) -> Tuple[Shape, bool]: """ Rewrites a shape so that a size scalar used as an extent is read through a symbol. @@ -29,17 +29,19 @@ def promote_size_scalars_in_shape(pv: ProgramVisitor, sdfg: SDFG, shape: Shape) :param pv: The program visitor. :param sdfg: The SDFG being built. :param shape: The requested shape. - :return: The shape with scalar extents replaced by symbols. + :return: The shape with scalar extents replaced by symbols, and whether anything was promoted. """ resolved = [symbolic.pystr_to_symbolic(e) if isinstance(e, str) else e for e in shape] names = [ n for n in symbolic.symlist(resolved) if n in sdfg.arrays and n not in sdfg.symbols and sdfg.arrays[n].total_size == 1 ] + if not names: + return shape, False # One symbol per distinct name; sorted() keeps the promotion states deterministic. replacements = {symbolic.pystr_to_symbolic(n): pv.promote_scalar_to_symbol(n, fresh=True) for n in sorted(names)} - return [e.subs(replacements) if isinstance(e, sympy.Basic) else e for e in resolved] + return [e.subs(replacements) if isinstance(e, sympy.Basic) else e for e in resolved], True @oprepo.replaces('dace.define_local') @@ -59,7 +61,7 @@ def _define_local_ex(pv: ProgramVisitor, if not isinstance(strides, (list, tuple)): strides = [strides] strides = [int(s) if isinstance(s, Integral) else s for s in strides] - shape = promote_size_scalars_in_shape(pv, sdfg, shape) + shape, _ = promote_size_scalars_in_shape(pv, sdfg, shape) name = pv.get_target_name() name, _ = sdfg.add_transient(name, shape, diff --git a/tests/numpy/array_creation_test.py b/tests/numpy/array_creation_test.py index fcb7343e40..792a77a434 100644 --- a/tests/numpy/array_creation_test.py +++ b/tests/numpy/array_creation_test.py @@ -1,6 +1,5 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace -from dace.frontend.python.common import DaceSyntaxError import numpy as np from common import compare_numpy_output import pytest @@ -258,15 +257,15 @@ def zeros_symbolic_size(): def test_ones_scalar_size_scalar(): + """A size held in a scalar is read into a symbol, so it can be used as an extent.""" @dace.program def ones_scalar_size(k: dace.int32): a = np.ones(k, dtype=np.uint32) return np.sum(a) - with pytest.raises(DaceSyntaxError): - out = ones_scalar_size(20) - assert out == 20 + out = ones_scalar_size(20) + assert out[0] == 20 def test_ones_scalar_size(): @@ -276,9 +275,8 @@ def ones_scalar_size(k: dace.int32): a = np.ones((k, k), dtype=np.uint32) return np.sum(a) - with pytest.raises(DaceSyntaxError): - out = ones_scalar_size(20) - assert out == 20 * 20 + out = ones_scalar_size(20) + assert out[0] == 20 * 20 if __name__ == "__main__": From 4702912d5fe9e5c4ada227d2c69702fa6b341a3b Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Tue, 21 Jul 2026 21:42:48 +0200 Subject: [PATCH 16/17] Bracket only the regions whose size symbols need it RegionBoundaryStates wrapped every control flow region in two empty states. On a three-loop program with a branch that is 8 regions and takes the state count from 5 to 15, none of which was needed. The largest CloudSC SDFG has 778 regions and was paying 1556 empty states. The two sides answer different questions, so they are decided separately: - A leading state is where an allocation goes. It is needed when the region's incoming edge assigns a symbol that a transient's size reads. - A trailing state is where the matching deallocation goes. It is needed when the region ends its scope, whatever sized it: with no state after the last region, the free has nowhere to be emitted and is silently dropped. Bracketing only the sized regions leaked a transient that outlives the region assigning its size. Sizes are collected across the whole SDFG tree rather than per SDFG: a nested transient can be sized by a symbol an enclosing region assigns and passes down through symbol_mapping, and reading only the owning SDFG's arrays leaves that region unbracketed -- the malloc(0) bug again one level down. CloudSC now takes 79 states rather than 1556, in 39 ms. Tests pin all three: the state count of a program needing nothing, a conditional whose incoming edge carries the size, and the free that the first cut dropped. --- .../passes/region_boundary_states.py | 34 +++- tests/passes/region_boundary_states_test.py | 145 ++++++++++++++---- 2 files changed, 144 insertions(+), 35 deletions(-) diff --git a/dace/transformation/passes/region_boundary_states.py b/dace/transformation/passes/region_boundary_states.py index f93dd4b728..7af78b1634 100644 --- a/dace/transformation/passes/region_boundary_states.py +++ b/dace/transformation/passes/region_boundary_states.py @@ -1,6 +1,6 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. -from typing import Any, Dict, Optional +from typing import Any, Dict, Optional, Set from dace import properties from dace.sdfg.sdfg import SDFG @@ -13,11 +13,14 @@ @transformation.explicit_cf_compatible class RegionBoundaryStates(ppl.Pass): """ - Brackets every control flow region with an empty state in its parent region. + Brackets a control flow region with an empty state in its parent region when its size symbols demand it. Allocation is only emitted inside a state. Without the leading state, data sized by a symbol that the region's incoming edge assigns is allocated at the region's predecessor, before the symbol is defined. The trailing state gives the matching deallocation a place to go. + + Only regions whose incoming assignments feed a transient's size are bracketed. Bracketing every region instead + triples the state count of a program that never sizes anything from an interstate assignment. """ CATEGORY: str = 'Helper' @@ -34,6 +37,18 @@ def apply_pass(self, sdfg: SDFG, _: Dict[str, Any]) -> Optional[int]: :param sdfg: The SDFG to modify in-place. :return: Number of states inserted, or None if unchanged. """ + # Sizes are collected across the whole tree, not per SDFG: a nested transient can be sized by a + # symbol an enclosing region assigns and passes down through symbol_mapping, and reading only the + # owning SDFG's arrays leaves that region unbracketed. Sharing a name that needs no boundary only + # costs two empty states. + sized_by: Set[str] = { + str(s) + for nested in sdfg.all_sdfgs_recursive() + for desc in nested.arrays.values() if desc.transient for s in desc.free_symbols + } + if not sized_by: + return None + inserted = 0 # Branches of a conditional are entered without an inter-state edge, so they need no boundary. regions = [ @@ -43,7 +58,16 @@ def apply_pass(self, sdfg: SDFG, _: Dict[str, Any]) -> Optional[int]: for node in list(cfg.nodes()): if not isinstance(node, AbstractControlFlowRegion): continue - cfg.add_state_before(node, is_start_block=node is cfg.start_block) - cfg.add_state_after(node) - inserted += 2 + # The two sides answer different questions. A leading state is where an allocation goes, + # so it is needed when the region's incoming edge assigns a symbol a size reads. A + # trailing state is where the matching deallocation goes, so it is needed when the region + # ends its scope: without a state after it, the free has nowhere to be emitted and is + # silently dropped. + assigned = {name for e in cfg.in_edges(node) for name in e.data.assignments} + if assigned & sized_by: + cfg.add_state_before(node, is_start_block=node is cfg.start_block) + inserted += 1 + if (assigned & sized_by) or not cfg.out_edges(node): + cfg.add_state_after(node) + inserted += 1 return inserted or None diff --git a/tests/passes/region_boundary_states_test.py b/tests/passes/region_boundary_states_test.py index 0625518301..455c89e046 100644 --- a/tests/passes/region_boundary_states_test.py +++ b/tests/passes/region_boundary_states_test.py @@ -19,8 +19,30 @@ def loop_and_branch(a: dace.float64[N], out: dace.float64[N]): @dace.program -def calls_loop_and_branch(a: dace.float64[N], out: dace.float64[N]): - loop_and_branch(a, out) +def loop_sized_by_an_assignment(a: dace.float64[N], Nt: dace.int64, out: dace.float64[N]): + b = np.empty(Nt + 1, dace.float64) # the size is promoted onto an interstate edge + for i in range(N): + b[i] = a[i] * 2.0 + for i in range(N): + out[i] = b[i] + + +@dace.program +def conditional_sized_by_an_assignment(a: dace.float64[N], Nt: dace.int64, out: dace.float64[N]): + b = np.empty(Nt + 1, dace.float64) # promoted before the branch, so the assignment enters it + if Nt > 2: + for i in range(N): + b[i] = a[i] * 2.0 + for i in range(N): + out[i] = b[i] + else: + for i in range(N): + out[i] = 0.0 + + +@dace.program +def calls_loop_sized_by_an_assignment(a: dace.float64[N], Nt: dace.int64, out: dace.float64[N]): + loop_sized_by_an_assignment(a, Nt, out) def regions_of(sdfg): @@ -32,66 +54,129 @@ def regions_of(sdfg): yield cfg, node -def test_every_region_is_bracketed(): +def bracketed(cfg, region): + return (all(isinstance(e.src, dace.SDFGState) for e in cfg.in_edges(region)) + and all(isinstance(e.dst, dace.SDFGState) for e in cfg.out_edges(region))) + + +def sizing_regions(sdfg): + """The regions the pass must bracket: their incoming edges assign a symbol a transient is sized by. + + A nested SDFG declares its own transients, so the sizes are read from the SDFG owning the region. + """ + for cfg, region in regions_of(sdfg): + sized_by = {str(s) for desc in cfg.sdfg.arrays.values() if desc.transient for s in desc.free_symbols} + if {name for e in cfg.in_edges(region) for name in e.data.assignments} & sized_by: + yield cfg, region + + +def test_a_region_whose_assignment_sizes_a_transient_is_bracketed(): + sdfg = loop_sized_by_an_assignment.to_sdfg(simplify=True) + targets = list(sizing_regions(sdfg)) + assert targets, 'the fixture must size a transient from an interstate assignment' + + assert RegionBoundaryStates().apply_pass(sdfg, {}) > 0 + for cfg, region in targets: + assert bracketed(cfg, region) + sdfg.validate() + + +def test_a_conditional_sized_by_an_assignment_is_bracketed(): + """A ConditionalBlock is a region too: the size assignment can arrive on its incoming edge. + + Its branches are entered without an inter-state edge, so they need no boundary of their own -- + but the block itself does, and skipping it would allocate before the size is defined. + """ + sdfg = conditional_sized_by_an_assignment.to_sdfg(simplify=True) + conditionals = [(cfg, n) for cfg, n in regions_of(sdfg) if isinstance(n, ConditionalBlock)] + assert conditionals, 'the fixture must keep a conditional block' + assert any((cfg, n) in list(sizing_regions(sdfg)) for cfg, n in conditionals) + + assert RegionBoundaryStates().apply_pass(sdfg, {}) > 0 + for cfg, block in conditionals: + assert bracketed(cfg, block) + sdfg.validate() + + n, nt = 6, 9 + a = np.arange(n, dtype=np.float64) + out = np.zeros(n) + sdfg(a=a, Nt=np.int64(nt), out=out, N=n) + assert np.allclose(out, a * 2.0) + + +def test_a_region_that_needs_no_boundary_is_left_alone(): + """Bracketing unconditionally tripled the state count of a program that sizes nothing this way.""" sdfg = loop_and_branch.to_sdfg(simplify=True) assert any(True for _ in regions_of(sdfg)) # guard against a vacuous check + assert not list(sizing_regions(sdfg)) + before = sum(1 for _ in sdfg.all_states()) - assert RegionBoundaryStates().apply_pass(sdfg, {}) > 0 - for cfg, region in regions_of(sdfg): - assert all(isinstance(e.src, dace.SDFGState) for e in cfg.in_edges(region)) - assert all(isinstance(e.dst, dace.SDFGState) for e in cfg.out_edges(region)) + assert RegionBoundaryStates().apply_pass(sdfg, {}) is None + assert sum(1 for _ in sdfg.all_states()) == before sdfg.validate() def test_regions_inside_a_nested_sdfg_are_bracketed(): """The pass descends into nested SDFGs, where the same allocation bug applies.""" - sdfg = calls_loop_and_branch.to_sdfg(simplify=False) + sdfg = calls_loop_sized_by_an_assignment.to_sdfg(simplify=False) nested = [n for n, _ in sdfg.all_nodes_recursive() if isinstance(n, dace.nodes.NestedSDFG)] assert nested, 'the fixture must keep a nested SDFG' - inner_regions = [(cfg, r) for cfg, r in regions_of(sdfg) if cfg.sdfg is not sdfg] - assert inner_regions, 'the nested SDFG must contain a region' + # Only simplify the callee: that is what moves the size assignment onto the region's incoming + # edge, and simplifying the caller too would inline the nested SDFG away. + for node in nested: + node.sdfg.simplify() + inner = [(cfg, r) for cfg, r in sizing_regions(sdfg) if cfg.sdfg is not sdfg] + assert inner, 'the nested SDFG must contain a region sized by an assignment' RegionBoundaryStates().apply_pass(sdfg, {}) - for cfg, region in inner_regions: - assert all(isinstance(e.src, dace.SDFGState) for e in cfg.in_edges(region)) - assert all(isinstance(e.dst, dace.SDFGState) for e in cfg.out_edges(region)) + for cfg, region in inner: + assert bracketed(cfg, region) sdfg.validate() def test_leading_region_keeps_start_block(): - """Bracketing a region that starts its parent must hand the start block over, not orphan it.""" - sdfg = dace.SDFG('leading_region') + """Bracketing a region that starts its parent must hand the start block over, not orphan it. + + Only a cyclic CFG reaches this: the pass brackets a region whose incoming edge assigns a size + symbol, and in an acyclic graph a start block has no incoming edge. Hand-built for that reason. + """ + sdfg = dace.SDFG('cyclic_start') sdfg.add_array('out', [1], dace.float64) + sdfg.add_symbol('K', dace.int64) + sdfg.add_transient('b', ['K'], dace.float64) + loop = dace.sdfg.state.LoopRegion('loop', 'i < 4', 'i', 'i = 0', 'i = i + 1') sdfg.add_node(loop, is_start_block=True) - state = loop.add_state('body', is_start_block=True) - tasklet = state.add_tasklet('one', {}, {'o'}, 'o = 1.0') - state.add_edge(tasklet, 'o', state.add_write('out'), None, dace.Memlet('out[0]')) + body = loop.add_state('body', is_start_block=True) + tasklet = body.add_tasklet('one', {}, {'o'}, 'o = 1.0') + body.add_edge(tasklet, 'o', body.add_write('out'), None, dace.Memlet('out[0]')) - RegionBoundaryStates().apply_pass(sdfg, {}) + tail = sdfg.add_state('tail') + sdfg.add_edge(loop, tail, dace.InterstateEdge(condition='0')) + sdfg.add_edge(tail, loop, dace.InterstateEdge(assignments={'K': '4'})) + + assert RegionBoundaryStates().apply_pass(sdfg, {}) == 2 assert isinstance(sdfg.start_block, dace.SDFGState) assert sdfg.start_block is not loop sdfg.validate() - out = np.zeros(1) - sdfg(out=out) - assert out[0] == 1.0 - def test_result_is_unchanged(): """The pass only inserts empty states, so it must not alter what the program computes.""" - a = np.random.default_rng(0).random(16) - 0.5 - expected = np.abs(a) + n, nt = 16, 20 + a = np.random.default_rng(0).random(n) - 0.5 - sdfg = loop_and_branch.to_sdfg(simplify=True) + sdfg = loop_sized_by_an_assignment.to_sdfg(simplify=True) RegionBoundaryStates().apply_pass(sdfg, {}) - out = np.zeros(16) - sdfg(a=a, out=out, N=16) - assert np.allclose(out, expected) + out = np.zeros(n) + sdfg(a=a, Nt=np.int64(nt), out=out, N=n) + assert np.allclose(out, a * 2.0) if __name__ == '__main__': - test_every_region_is_bracketed() + test_a_region_whose_assignment_sizes_a_transient_is_bracketed() + test_a_conditional_sized_by_an_assignment_is_bracketed() + test_a_region_that_needs_no_boundary_is_left_alone() test_regions_inside_a_nested_sdfg_are_bracketed() test_leading_region_keeps_start_block() test_result_is_unchanged() From 940f53534e3e574dfe6bfc7c5fedde6c857416f2 Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Fri, 31 Jul 2026 10:49:54 +0200 Subject: [PATCH 17/17] Locate the allocation without depending on how it is spelled Aligned allocations (#2438) changed the emitted line from `new double` to `new (std::align_val_t(64)) double`, and the free from `delete[] b` to `::operator delete[](b, std::align_val_t(64))`, so the two searches that find those lines stopped matching once this branch merged main. Both assertions are unchanged: the size symbol is still assigned before the allocation, and the array is still freed. Only the patterns used to find the lines are now spelling-agnostic, as the codegen tests on main already do. --- tests/size_scalar_shape_promotion_test.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/size_scalar_shape_promotion_test.py b/tests/size_scalar_shape_promotion_test.py index a054321d29..c49220eb4b 100644 --- a/tests/size_scalar_shape_promotion_test.py +++ b/tests/size_scalar_shape_promotion_test.py @@ -154,10 +154,11 @@ def test_size_symbol_is_assigned_before_the_allocation(): sym = str(next(iter(sdfg.arrays['b'].free_symbols))) lines = sdfg.generate_code()[0].clean_code.splitlines() - alloc = next(i for i, line in enumerate(lines) if 'new double' in line and sym in line) + # an aligned heap array reads ``new (std::align_val_t(64)) double`` and frees through ``::operator delete[](b, ..)`` + alloc = next(i for i, line in enumerate(lines) if 'b = new' in line and sym in line) assign = next(i for i, line in enumerate(lines) if line.strip().startswith(f'{sym} = ')) assert assign < alloc - assert any('delete[] b' in line for line in lines), 'the array is never freed' + assert any('delete[] b' in line or 'delete[](b' in line for line in lines), 'the array is never freed' def test_a_size_one_array_is_read_through_a_subscript():