diff --git a/dace/codegen/codegen.py b/dace/codegen/codegen.py index 78a9499f50..7183a76b73 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 @@ -198,6 +199,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/frontend/python/newast.py b/dace/frontend/python/newast.py index a4e1eaa2a8..12ae9a9f5a 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 @@ -5370,6 +5370,46 @@ 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, fresh: bool = False) -> symbolic.symbol: + """ + 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; 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 + desc = self.sdfg.arrays[scalar] + sym = None if fresh else self.indirections.get(key) + if sym is None: + base = f'__sym_{scalar}' + if fresh: + # 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: + name = base + try: + self.sdfg.add_symbol(name, desc.dtype) + except FileExistsError: + pass # A cached promotion may re-add an existing symbol. + sym = dace.symbol(name, dtype=desc.dtype) + if not fresh: + self.indirections[key] = 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] + # 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, s: ast.AST, multidim: bool = False) -> Union[Any, Tuple[Union[Any, str, symbolic.symbol]]]: @@ -5379,29 +5419,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.py b/dace/frontend/python/replacements/array_creation.py index 66bcd57163..4f9d35998c 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,10 @@ def _numpy_full(pv: ProgramVisitor, if isinstance(shape, (Number, str)) or symbolic.issymbolic(shape): shape = [shape] + 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 b9aae4a888..29549156d6 100644 --- a/dace/frontend/python/replacements/array_creation_dace.py +++ b/dace/frontend/python/replacements/array_creation_dace.py @@ -10,10 +10,39 @@ 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 +from dace import symbolic + + +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. + + 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. + :param sdfg: The SDFG being built. + :param shape: The requested shape. + :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], True + @oprepo.replaces('dace.define_local') @oprepo.replaces('dace.ndarray') @@ -32,6 +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) name = pv.get_target_name() name, _ = sdfg.add_transient(name, shape, 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/gemm.py b/dace/libraries/blas/nodes/gemm.py index c78bcdf03b..9180f45e37 100644 --- a/dace/libraries/blas/nodes/gemm.py +++ b/dace/libraries/blas/nodes/gemm.py @@ -8,7 +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) +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 @@ -46,8 +47,10 @@ 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) + _, outer_array_c, _, strides_c = _matrix_operand(cdata) dtype_a = outer_array_a.dtype.type dtype_b = outer_array_b.dtype.type @@ -79,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=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" @@ -464,7 +467,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 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) dtype = adesc.dtype.base_type if not equal_valued(0, node.beta): @@ -551,11 +558,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 = _matrix_subset_size(memlet.subset) if dst_conn == '_b': - size1 = memlet.subset.size() + size1 = _matrix_subset_size(memlet.subset) if dst_conn == '_c': - size2 = memlet.subset.size() + size2 = _matrix_subset_size(memlet.subset) if self.transA: size0 = list(reversed(size0)) @@ -575,7 +582,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 = _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 0d07014f5e..ade20fecca 100644 --- a/dace/libraries/blas/nodes/matmul.py +++ b/dace/libraries/blas/nodes/matmul.py @@ -1,16 +1,42 @@ # 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 from math import prod +def _matrix_subset_size(subset): + """ + Returns an operand's size in the matrix view ``SpecializeMatMul`` matched on. + + :param subset: The subset of the memlet accessing the operand. + :return: The size in the matrix view. + """ + return matrix_view(subset)[0] + + +def _matrix_operand(operand): + """ + 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, outer descriptor, and 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 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() @@ -137,8 +163,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 - (_, _, 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 node.transA: ashape = list(reversed(ashape)) 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/dace/sdfg/state.py b/dace/sdfg/state.py index 93c10bc539..a976fff754 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/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..7af78b1634 --- /dev/null +++ b/dace/transformation/passes/region_boundary_states.py @@ -0,0 +1,73 @@ +# Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. + +from typing import Any, Dict, Optional, Set + +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 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' + + 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. + """ + # 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 = [ + 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 + # 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/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 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..f298c4f0ba --- /dev/null +++ b/tests/library/matmul_unit_dim_squeeze_test.py @@ -0,0 +1,133 @@ +# Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. +"""``SpecializeMatMul`` and ``Gemm`` must read the same view of an operand. + +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 + +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, 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_)) + + +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 dim when squeezed, where a "whichever view looks 2D" heuristic breaks. +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): 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__": diff --git a/tests/passes/region_boundary_states_test.py b/tests/passes/region_boundary_states_test.py new file mode 100644 index 0000000000..455c89e046 --- /dev/null +++ b/tests/passes/region_boundary_states_test.py @@ -0,0 +1,182 @@ +# 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] + + +@dace.program +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): + for cfg in sdfg.all_control_flow_regions(recursive=True): + if isinstance(cfg, ConditionalBlock): + continue + for node in cfg.nodes(): + if isinstance(node, AbstractControlFlowRegion): + yield cfg, node + + +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, {}) 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_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' + # 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: + 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. + + 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) + 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]')) + + 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() + + +def test_result_is_unchanged(): + """The pass only inserts empty states, so it must not alter what the program computes.""" + n, nt = 16, 20 + a = np.random.default_rng(0).random(n) - 0.5 + + sdfg = loop_sized_by_an_assignment.to_sdfg(simplify=True) + RegionBoundaryStates().apply_pass(sdfg, {}) + 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_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() diff --git a/tests/size_scalar_shape_promotion_test.py b/tests/size_scalar_shape_promotion_test.py new file mode 100644 index 0000000000..c49220eb4b --- /dev/null +++ b/tests/size_scalar_shape_promotion_test.py @@ -0,0 +1,202 @@ +# 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; 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 pytest + +import dace + +N = dace.symbol('N') + + +@dace.program +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 + 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 + + +@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 + + +@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) + 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_two_arrays_from_a_reassigned_size_keep_their_own_extents(): + """Reusing one size name for two arrays must not collapse their extents. + + 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) + 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. + + 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) + 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) + # 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() + 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. + + Run once unsimplified and once simplified; 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}' + + +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() + + # 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 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(): + """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() + 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() + 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) diff --git a/tests/transpose_unit_dim_squeeze_test.py b/tests/transpose_unit_dim_squeeze_test.py new file mode 100644 index 0000000000..b5efad90df --- /dev/null +++ b/tests/transpose_unit_dim_squeeze_test.py @@ -0,0 +1,76 @@ +# 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 +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 + +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")