diff --git a/pytensor/compile/debug/debugmode.py b/pytensor/compile/debug/debugmode.py index 135d8ecd83..7131ef7fec 100644 --- a/pytensor/compile/debug/debugmode.py +++ b/pytensor/compile/debug/debugmode.py @@ -1316,7 +1316,10 @@ def printstuff(self): # List of default version of make thunk. # This is needed to know if the user overrode it. -default_make_thunk = [get_unbound_function(COp.make_thunk)] +default_make_thunk = [ + get_unbound_function(Op.make_thunk), + get_unbound_function(COp.make_thunk), +] # Debug mode cheats and initializes the linker in a different way in @@ -1369,6 +1372,7 @@ def make_all( # can't import at toplevel because of circular import TODO: # don't do this ugly hacky way of setting the # filter_checks_isfinite + from pytensor.link.c.dispatch.basic import c_thunk_from_dispatch fgraph = self.fgraph input_storage_ = input_storage @@ -1421,15 +1425,10 @@ def make_all( debug = hasattr(node.op, "debug_perform") try: - if ( - not self.maker.mode.check_c_code - or debug - or not isinstance(node.op, COp) - ): + if not self.maker.mode.check_c_code or debug: raise MethodNotDefined() - node.op.prepare_node(node, storage_map, compute_map, "c") - thunk = node.op.make_c_thunk( + thunk = c_thunk_from_dispatch( node, storage_map, compute_map, no_recycling ) thunks_c.append(thunk) @@ -1451,19 +1450,18 @@ def make_all( else: thunks_py.append(None) - if ( - not self.maker.mode.check_c_code - and thunks_py[-1] is None - and isinstance(node.op, COp) - ): - _logger.warning( - f"Op {node.op} doesn't have a perform, forcing check of the C code" - ) - node.op.prepare_node(node, storage_map, compute_map, "c") - thunk = node.op.make_c_thunk( - node, storage_map, compute_map, no_recycling - ) - thunks_c[-1] = thunk + if not self.maker.mode.check_c_code and thunks_py[-1] is None: + try: + thunk = c_thunk_from_dispatch( + node, storage_map, compute_map, no_recycling + ) + except (NotImplementedError, MethodNotDefined): + pass + else: + _logger.warning( + f"Op {node.op} doesn't have a perform, forcing check of the C code" + ) + thunks_c[-1] = thunk # If the op defined its own make_thunk, use the generated thunk if thunk_other is not None: diff --git a/pytensor/link/c/basic.py b/pytensor/link/c/basic.py index 1a3a2c3ff8..eda35ee40d 100644 --- a/pytensor/link/c/basic.py +++ b/pytensor/link/c/basic.py @@ -13,6 +13,7 @@ from pytensor.compile.compilelock import lock_ctx from pytensor.configdefaults import config from pytensor.graph.basic import ( + Apply, AtomicVariable, Constant, ) @@ -553,6 +554,7 @@ class CLinker(Linker): def __init__(self, schedule=None): self.fgraph = None + self._node_impls: dict[Apply, CLinkerOp] = {} super().__init__(scheduler=schedule) def accept( @@ -573,6 +575,23 @@ def accept( self.no_recycling = no_recycling return self + def _impl_for(self, node) -> CLinkerOp: + """Return `node`'s C implementation, resolved and memoized via `c_funcify`.""" + try: + return self._node_impls[node] + except KeyError: + pass + # Imported lazily so `import pytensor` does not load the dispatch + # registrations; the import triggers them on first compilation. + from pytensor.link.c.dispatch.basic import c_funcify + + try: + impl = c_funcify(node.op, node=node) + except NotImplementedError as exc: + raise NotImplementedError(f"{node.op} cannot produce C code") from exc + self._node_impls[node] = impl + return impl + def fetch_variables(self): """Fills the inputs, outputs, variables, orphans, temps and node_order fields.""" fgraph = self.fgraph @@ -591,10 +610,14 @@ def fetch_variables(self): # that needs it self.node_params = dict() for node in self.node_order: - if not isinstance(node.op, CLinkerOp): + try: + impl = self._impl_for(node) + except NotImplementedError: + # No C implementation; code_gen will raise if this node is + # actually compiled. continue try: - params = node.op.get_params(node) + params = impl.get_params(node) except MethodNotDefined: params = NoParams if params is not NoParams: @@ -602,10 +625,10 @@ def fetch_variables(self): # same params. if params in self.node_params: var = self.node_params[params] - assert var.type == node.params_type + assert var.type == impl.params_type fgraph.clients[var].append((node, "params")) else: - var = Constant(node.params_type, params) + var = Constant(impl.params_type, params) fgraph.clients[var] = [(node, "params")] self.node_params[params] = var self.variables.append(var) @@ -775,10 +798,7 @@ def code_gen(self): id += 2 for node_num, node in enumerate(self.node_order): - op = node.op - - if not isinstance(op, CLinkerOp): - raise NotImplementedError(f"{op} cannot produce C code") + op = self._impl_for(node) sub = dict(failure_var=failure_var) @@ -905,7 +925,9 @@ def support_code(self): """ ) # generic support code - for x in [y.type for y in self.variables] + [y.op for y in self.node_order]: + for x in [y.type for y in self.variables] + [ + self._impl_for(y) for y in self.node_order + ]: support_code = x.c_support_code() if isinstance(support_code, list): ret.extend(support_code) @@ -941,7 +963,9 @@ def compile_args(self): c_compiler = self.c_compiler() - for x in [y.type for y in self.variables] + [y.op for y in self.node_order]: + for x in [y.type for y in self.variables] + [ + self._impl_for(y) for y in self.node_order + ]: if isinstance(x, CLinkerObject): ret += x.c_compile_args(c_compiler=c_compiler) @@ -949,7 +973,9 @@ def compile_args(self): # The args set by the compiler include the user flags. We do not want # to reorder them ret += c_compiler.compile_args() - for x in [y.type for y in self.variables] + [y.op for y in self.node_order]: + for x in [y.type for y in self.variables] + [ + self._impl_for(y) for y in self.node_order + ]: if isinstance(x, CLinkerObject): no_comp = x.c_no_compile_args(c_compiler=c_compiler) @@ -970,7 +996,9 @@ def headers(self): """ ret = [] c_compiler = self.c_compiler() - for x in [y.type for y in self.variables] + [y.op for y in self.node_order]: + for x in [y.type for y in self.variables] + [ + self._impl_for(y) for y in self.node_order + ]: if isinstance(x, CLinkerObject): ret += x.c_headers(c_compiler=c_compiler) return uniq(ret) @@ -984,14 +1012,18 @@ def init_code(self): """ ret = [] - for x in [y.type for y in self.variables] + [y.op for y in self.node_order]: + for x in [y.type for y in self.variables] + [ + self._impl_for(y) for y in self.node_order + ]: if isinstance(x, CLinkerObject): ret += x.c_init_code() return uniq(ret) def c_compiler(self): c_compiler = None - for x in [y.type for y in self.variables] + [y.op for y in self.node_order]: + for x in [y.type for y in self.variables] + [ + self._impl_for(y) for y in self.node_order + ]: # FIXME: Why would a `Type` have a `c_compiler` field?! if hasattr(x, "c_compiler"): x_compiler = x.c_compiler() @@ -1021,7 +1053,9 @@ def header_dirs(self): """ ret = [] c_compiler = self.c_compiler() - for x in [y.type for y in self.variables] + [y.op for y in self.node_order]: + for x in [y.type for y in self.variables] + [ + self._impl_for(y) for y in self.node_order + ]: if isinstance(x, CLinkerObject): ret += x.c_header_dirs(c_compiler=c_compiler) # filter out empty strings/None @@ -1037,7 +1071,9 @@ def libraries(self): """ ret = [] c_compiler = self.c_compiler() - for x in [y.type for y in self.variables] + [y.op for y in self.node_order]: + for x in [y.type for y in self.variables] + [ + self._impl_for(y) for y in self.node_order + ]: if isinstance(x, CLinkerObject): ret += x.c_libraries(c_compiler=c_compiler) return uniq(ret) @@ -1052,7 +1088,9 @@ def lib_dirs(self): """ ret = [] c_compiler = self.c_compiler() - for x in [y.type for y in self.variables] + [y.op for y in self.node_order]: + for x in [y.type for y in self.variables] + [ + self._impl_for(y) for y in self.node_order + ]: if isinstance(x, CLinkerObject): ret += x.c_lib_dirs(c_compiler=c_compiler) # filter out empty strings/None @@ -1433,8 +1471,14 @@ def in_sig(i, topological_pos, i_idx): version = [] for node_pos, node in enumerate(order): - if hasattr(node.op, "c_code_cache_version_apply"): - version.append(node.op.c_code_cache_version_apply(node)) + try: + impl = self._impl_for(node) + except NotImplementedError: + # No C implementation: contributes no version entry (code_gen + # raises later if this graph is compiled). + pass + else: + version.append(impl.c_code_cache_version_apply(node)) props = getattr(node.op, "__props__", None) @@ -1819,6 +1863,8 @@ def accept(self, fgraph, no_recycling=None, profile=None): def make_all( self, profiler=None, input_storage=None, output_storage=None, storage_map=None ): + from pytensor.link.c.dispatch.basic import make_node_thunk_with_c_dispatch + fgraph = self.fgraph order = self.schedule(fgraph) no_recycling = self.no_recycling @@ -1838,9 +1884,16 @@ def make_all( thunks = [] for node in order: - # make_thunk will try by default C code, otherwise - # it fall back to python. - thunks += [node.op.make_thunk(node, storage_map, compute_map, no_recycling)] + # Try the C dispatch first, otherwise fall back to Python. + thunks += [ + make_node_thunk_with_c_dispatch( + node, + storage_map, + compute_map, + no_recycling, + try_c=bool(config.cxx), + ) + ] thunks[-1].inputs = [storage_map[v] for v in node.inputs] thunks[-1].outputs = [storage_map[v] for v in node.outputs] diff --git a/pytensor/link/c/dispatch/__init__.py b/pytensor/link/c/dispatch/__init__.py new file mode 100644 index 0000000000..cb5d74c71f --- /dev/null +++ b/pytensor/link/c/dispatch/__init__.py @@ -0,0 +1,6 @@ +# isort: off +from pytensor.link.c.dispatch.basic import c_funcify + +import pytensor.link.c.dispatch.blas + +# isort: on diff --git a/pytensor/link/c/dispatch/basic.py b/pytensor/link/c/dispatch/basic.py new file mode 100644 index 0000000000..c8abbf7fcb --- /dev/null +++ b/pytensor/link/c/dispatch/basic.py @@ -0,0 +1,203 @@ +import warnings +from collections.abc import Collection +from functools import singledispatch +from typing import NoReturn + +from pytensor.graph.basic import Apply, Variable +from pytensor.graph.fg import FunctionGraph +from pytensor.graph.op import ComputeMapType, Op, StorageMapType, ThunkType +from pytensor.graph.utils import MethodNotDefined +from pytensor.link.c.interface import CLinkerOp +from pytensor.link.c.op import ( + COp, + CThunkWrapperType, + is_cthunk_wrapper_type, +) + + +@singledispatch +def c_funcify(op: Op, node: Apply | None = None, **kwargs) -> CLinkerOp: + """Return the C implementation of `op` at `node`. + + By default an op implementing `CLinkerOp` (every `COp`) is its own + implementation; otherwise raise `NotImplementedError` and let the caller fall + back to the Python thunk. + """ + if isinstance(op, CLinkerOp): + return op + raise NotImplementedError(f"No C implementation registered for {type(op).__name__}") + + +def _hashable_aliasing_map(aliasing_map: dict[int, list[int]]) -> tuple: + return tuple(sorted((idx, tuple(vals)) for idx, vals in aliasing_map.items())) + + +class CImpl(CLinkerOp): + """A C implementation of an `Op`, detached from the op. + + Returned by `c_funcify`; never a graph op. Subclasses that add configuration + must extend `_impl_props`, which backs equality, hashing, and the cache key. + """ + + # `Apply.clone_with_new_inputs` reads this off the node's op when the + # single-node graph is cloned for compilation; impl outputs never depend on + # input values (the graph op already fixed the output types). + _output_type_depends_on_input_value = False + + def __init__( + self, + op: Op, + *, + destroy_map: dict[int, list[int]] | None = None, + view_map: dict[int, list[int]] | None = None, + ): + self.op = op + if destroy_map is None: + destroy_map = getattr(op, "destroy_map", {}) + if view_map is None: + view_map = getattr(op, "view_map", {}) + self.destroy_map = destroy_map + self.view_map = view_map + + def _impl_props(self) -> tuple: + return ( + self.op, + _hashable_aliasing_map(self.destroy_map), + _hashable_aliasing_map(self.view_map), + ) + + def __eq__(self, other) -> bool: + return type(self) is type(other) and self._impl_props() == other._impl_props() + + def __hash__(self) -> int: + return hash((type(self), *self._impl_props())) + + def __str__(self) -> str: + return f"{type(self).__name__}{{{self.op}}}" + + def make_node(self, *inputs) -> NoReturn: + raise RuntimeError( + f"{type(self).__name__} is a C implementation, not a graph op." + ) + + def prepare_node( + self, + node: Apply, + storage_map: StorageMapType, + compute_map: ComputeMapType | None, + impl: str | None, + ) -> None: + """No-op: C preparation happens when `c_funcify` constructs the impl.""" + + +def c_thunk_from_dispatch( + node: Apply, + storage_map: StorageMapType, + compute_map: ComputeMapType | None, + no_recycling: Collection[Variable], +) -> CThunkWrapperType: + """Compile a C thunk for `node`, taking its implementation from `c_funcify`. + + Raises + ------ + NotImplementedError + If `node.op` has no C implementation, or has float16 inputs/outputs. + MethodNotDefined + If the implementation declines this node (e.g. an unsupported dtype). + + Callers fall back to a Python thunk on either. + """ + # Imported here to avoid an import cycle. + import pytensor.link.c.basic + + # Resolve eagerly so an unimplemented op raises before prepare_node runs and + # before any compilation work; CLinker re-resolves (memoized) during codegen. + c_funcify(node.op, node=node) + + node.op.prepare_node( + node, storage_map=storage_map, compute_map=compute_map, impl="c" + ) + + node_input_storage = [storage_map[r] for r in node.inputs] + node_output_storage = [storage_map[r] for r in node.outputs] + + fgraph = FunctionGraph(node.inputs, node.outputs) + fgraph_no_recycling = [ + new_o + for (new_o, old_o) in zip(fgraph.outputs, node.outputs, strict=True) + if old_o in no_recycling + ] + cl = pytensor.link.c.basic.CLinker().accept( + fgraph, no_recycling=fgraph_no_recycling + ) + + # float16 gets special treatment since running unprepared C code will get bad + # results. + if not getattr(node.op, "_f16_ok", False): + + def is_f16(t): + return getattr(t, "dtype", "") == "float16" + + if any(is_f16(i.type) for i in node.inputs) or any( + is_f16(o.type) for o in node.outputs + ): + # get_dynamic_module just tries to build the C code; it raises for + # impls without C code, in which case we don't want to warn. + cl.get_dynamic_module() + warnings.warn(f"Disabling C code for {node.op} due to unsupported float16") + raise NotImplementedError("float16") + + outputs = cl.make_thunk( + input_storage=node_input_storage, output_storage=node_output_storage + ) + thunk, _node_input_filters, _node_output_filters = outputs + + if compute_map is None: + rval = is_cthunk_wrapper_type(thunk) + else: + cm_entries = [compute_map[o] for o in node.outputs] + + @is_cthunk_wrapper_type + def rval(thunk=thunk, cm_entries=cm_entries): + thunk() + for entry in cm_entries: + entry[0] = True + + rval.thunk = thunk + rval.cthunk = thunk.cthunk + rval.inputs = node_input_storage + rval.outputs = node_output_storage + rval.lazy = False + return rval + + +# Ops whose `make_thunk` is one of these run through the dispatch; anything else +# overrode `make_thunk` (e.g. `IfElse`, `Scan`) and keeps its custom path. +_DEFAULT_MAKE_THUNKS = (Op.make_thunk, COp.make_thunk) + + +def make_node_thunk_with_c_dispatch( + node: Apply, + storage_map: StorageMapType, + compute_map: ComputeMapType | None, + no_recycling: Collection[Variable], + *, + try_c: bool, + fallback_impl: str | None = None, +) -> ThunkType: + """Make a thunk for `node`, trying the C dispatch first when `try_c`. + + When the C attempt fails (no implementation, or the implementation declines + the node) the fallback passes ``impl="py"`` so `COp.make_thunk` does not + retry the C path. + """ + if try_c and type(node.op).make_thunk in _DEFAULT_MAKE_THUNKS: + try: + return c_thunk_from_dispatch(node, storage_map, compute_map, no_recycling) + except (NotImplementedError, MethodNotDefined): + fallback_impl = "py" + # Op.make_thunk is untyped upstream; pin the result to its real type. + thunk: ThunkType = node.op.make_thunk( + node, storage_map, compute_map, no_recycling, impl=fallback_impl + ) + return thunk diff --git a/pytensor/link/c/dispatch/blas.py b/pytensor/link/c/dispatch/blas.py new file mode 100644 index 0000000000..f331ad56fe --- /dev/null +++ b/pytensor/link/c/dispatch/blas.py @@ -0,0 +1,129 @@ +from collections.abc import Hashable + +from pytensor.configdefaults import config +from pytensor.graph.basic import Apply +from pytensor.graph.utils import MethodNotDefined +from pytensor.link.c.dispatch.basic import CImpl, c_funcify +from pytensor.link.c.params_type import Params, ParamsType +from pytensor.scalar import bool as bool_t +from pytensor.tensor.blas._core import ldflags, must_initialize_y_gemv +from pytensor.tensor.blas.c_code.blas_headers import ( + blas_header_text, + blas_header_version, +) +from pytensor.tensor.blas.c_code.codegen import gemm_c_code, gemv_c_code, ger_c_code +from pytensor.tensor.blas.gemm import Gemm +from pytensor.tensor.blas.gemv import Gemv +from pytensor.tensor.blas.ger import Ger + + +class BlasImpl(CImpl): + """Base for the BLAS C implementations: link flags, headers, and the ``inplace`` param. + + The generated code reads ``inplace`` out of a `ParamsType` rather than baking it in, so one + compiled module serves both forms of an op. + """ + + params_type = ParamsType(inplace=bool_t) + + def get_params(self, node: Apply) -> Params: + return self.params_type.get_params(self.op) + + def c_support_code(self, **kwargs) -> str: + return blas_header_text() + + def c_headers(self, **kwargs) -> list[str]: + return [] + + def c_libraries(self, **kwargs) -> list[str]: + return ldflags() + + def c_compile_args(self, **kwargs) -> list[str]: + return ldflags(libs=False, flags=True) + + def c_lib_dirs(self, **kwargs) -> list[str]: + return ldflags(libs=False, libs_dir=True) + + def c_header_dirs(self, **kwargs) -> list[str]: + return ldflags(libs=False, include_dir=True) + + +class GemmImpl(BlasImpl): + """C implementation of `Gemm`.""" + + op: Gemm + + def c_support_code(self, **kwargs) -> str: + # BLAS declarations plus the MOD macro and compute_strides helper the GEMM templates + # in codegen.py expect. + mod_str = """ + #ifndef MOD + #define MOD % + #endif + void compute_strides(npy_intp *shape, int N_shape, int type_size, npy_intp *res) { + int s; + res[N_shape - 1] = type_size; + for (int i = N_shape - 1; i > 0; i--) { + s = shape[i]; + res[i - 1] = res[i] * (s > 0 ? s : 1); + } + } + """ + return blas_header_text() + mod_str + + def c_code_cache_version(self) -> tuple[Hashable, ...]: + return (8, 14, blas_header_version()) + + def c_code(self, node, name, inp, out, sub) -> str: + if node.inputs[0].type.dtype.startswith("complex"): + raise MethodNotDefined("GemmImpl.c_code") + return gemm_c_code(node, name, inp, out, sub) + + +class GemvImpl(BlasImpl): + """C implementation of `Gemv`.""" + + op: Gemv + + def c_code_cache_version(self) -> tuple[Hashable, ...]: + return (18, blas_header_version(), must_initialize_y_gemv()) + + def c_code(self, node, name, inp, out, sub) -> str: + # No `blas__ldflags` guard: the fallback header supplies its own [sd]gemv_ and [sd]dot_. + if node.outputs[0].dtype not in ("float32", "float64"): + raise MethodNotDefined("GemvImpl.c_code") + return gemv_c_code(node, name, inp, out, sub) + + +class GerImpl(BlasImpl): + """C implementation of `Ger`.""" + + op: Ger + + def c_code_cache_version(self) -> tuple[Hashable, ...]: + return (11, blas_header_version()) + + def c_code(self, node, name, inp, out, sub) -> str: + # Unlike gemv, the fallback header defines no [sd]ger_, so without link flags this code + # would not link. + if not config.blas__ldflags or node.outputs[0].dtype not in ( + "float32", + "float64", + ): + raise MethodNotDefined("GerImpl.c_code") + return ger_c_code(node, name, inp, out, sub) + + +@c_funcify.register(Gemm) +def c_funcify_Gemm(op, node=None, **kwargs) -> GemmImpl: + return GemmImpl(op) + + +@c_funcify.register(Gemv) +def c_funcify_Gemv(op, node=None, **kwargs) -> GemvImpl: + return GemvImpl(op) + + +@c_funcify.register(Ger) +def c_funcify_Ger(op, node=None, **kwargs) -> GerImpl: + return GerImpl(op) diff --git a/pytensor/link/jax/dispatch/blas.py b/pytensor/link/jax/dispatch/blas.py index a0d0faeabb..c67b742bdd 100644 --- a/pytensor/link/jax/dispatch/blas.py +++ b/pytensor/link/jax/dispatch/blas.py @@ -1,7 +1,7 @@ import jax.numpy as jnp from pytensor.link.jax.dispatch import jax_funcify -from pytensor.tensor.blas import BatchedDot +from pytensor.tensor.blas import BatchedDot, Gemm, Gemv, Ger @jax_funcify.register(BatchedDot) @@ -12,3 +12,31 @@ def batched_dot(a, b): return jnp.matmul(a, b) return batched_dot + + +@jax_funcify.register(Gemm) +def jax_funcify_Gemm(op, **kwargs): + def gemm(z, alpha, x, y, beta): + # Written out rather than fused by hand: XLA contracts this to a single dot with + # alpha and beta folded into it. + return beta * z + alpha * jnp.matmul(x, y) + + return gemm + + +@jax_funcify.register(Gemv) +def jax_funcify_Gemv(op, **kwargs): + def gemv(y, alpha, A, x, beta): + # As with Gemm above, XLA folds the scalars into the contraction itself. + return beta * y + alpha * jnp.matmul(A, x) + + return gemv + + +@jax_funcify.register(Ger) +def jax_funcify_Ger(op, **kwargs): + def ger(A, alpha, x, y): + # As with Gemm above, there is no fused primitive to call and none is needed. + return A + alpha * jnp.outer(x, y) + + return ger diff --git a/pytensor/link/mlx/dispatch/blas.py b/pytensor/link/mlx/dispatch/blas.py index 9181d8604d..f7213e9e58 100644 --- a/pytensor/link/mlx/dispatch/blas.py +++ b/pytensor/link/mlx/dispatch/blas.py @@ -1,8 +1,8 @@ import mlx.core as mx -from pytensor.graph.basic import Constant from pytensor.link.mlx.dispatch import mlx_funcify -from pytensor.tensor.blas import BatchedDot, Gemv, Ger +from pytensor.link.utils import get_static_scalar +from pytensor.tensor.blas import BatchedDot, Gemm, Gemv, Ger @mlx_funcify.register(BatchedDot) @@ -15,10 +15,28 @@ def batched_dot(a, b): return batched_dot +@mlx_funcify.register(Gemm) +def mlx_funcify_Gemm(op, node=None, **kwargs): + static_alpha = get_static_scalar(node, 1) + static_beta = get_static_scalar(node, 4) + + if static_alpha is not None and static_beta is not None: + + def gemm(z, alpha, x, y, beta): + return mx.addmm(z, x, y, alpha=static_alpha, beta=static_beta) + + else: + + def gemm(z, alpha, x, y, beta): + return beta * z + alpha * mx.matmul(x, y) + + return gemm + + @mlx_funcify.register(Gemv) def mlx_funcify_Gemv(op, node=None, **kwargs): - static_alpha = _as_float_constant(node.inputs[1]) if node is not None else None - static_beta = _as_float_constant(node.inputs[4]) if node is not None else None + static_alpha = get_static_scalar(node, 1) + static_beta = get_static_scalar(node, 4) if static_alpha is not None and static_beta is not None: @@ -35,7 +53,7 @@ def gemv(y, alpha, A, x, beta): @mlx_funcify.register(Ger) def mlx_funcify_Ger(op, node=None, **kwargs): - static_alpha = _as_float_constant(node.inputs[1]) if node is not None else None + static_alpha = get_static_scalar(node, 1) if static_alpha is not None: @@ -52,12 +70,3 @@ def ger(A, alpha, x, y): return A + alpha * mx.outer(x, y) return ger - - -def _as_float_constant(var): - if not isinstance(var, Constant): - return None - try: - return float(var.data) - except (TypeError, ValueError): - return None diff --git a/pytensor/link/numba/dispatch/__init__.py b/pytensor/link/numba/dispatch/__init__.py index e8c4aab8db..a8eb8bea23 100644 --- a/pytensor/link/numba/dispatch/__init__.py +++ b/pytensor/link/numba/dispatch/__init__.py @@ -5,6 +5,7 @@ from pytensor.link.numba.dispatch.basic import numba_funcify, numba_typify # Load dispatch specializations +import pytensor.link.numba.dispatch.blas import pytensor.link.numba.dispatch.blockwise import pytensor.link.numba.dispatch.compile_ops import pytensor.link.numba.dispatch.elemwise diff --git a/pytensor/link/numba/dispatch/blas.py b/pytensor/link/numba/dispatch/blas.py new file mode 100644 index 0000000000..1b07c23243 --- /dev/null +++ b/pytensor/link/numba/dispatch/blas.py @@ -0,0 +1,116 @@ +import numpy as np + +from pytensor.link.numba.dispatch import basic as numba_basic +from pytensor.link.numba.dispatch.basic import register_funcify_default_op_cache_key +from pytensor.link.numba.dispatch.linalg.products import _gemm, _ger +from pytensor.tensor.blas import Gemm, Gemv, Ger + + +@register_funcify_default_op_cache_key(Gemm) +def numba_funcify_Gemm(op, node, **kwargs): + """Dispatch ``Gemm`` to one BLAS call, with its scalars carried as gemm's own alpha and beta.""" + dtype = node.outputs[0].type.numpy_dtype + + if op.inplace: + + @numba_basic.numba_njit + def gemm(Z, alpha, X, Y, beta): + return _gemm(X, Y, Z, False, False, alpha.item(), beta.item()) + + else: + + @numba_basic.numba_njit + def gemm(Z, alpha, X, Y, beta): + # `Z` is only broadcast against the product, so the accumulator gemm writes into takes + # the product's shape rather than `Z`'s. Copying also leaves `Z` intact, which is the + # whole difference between this op and its inplace form. + out = np.empty((X.shape[0], Y.shape[1]), dtype=dtype) + _gemm(X, Y, out, False, False, alpha.item(), 0.0) + b = beta.item() + if b == 1.0: + out += Z + elif b != 0.0: + out += b * Z + return out + + cache_version = 3 + return gemm, cache_version + + +@register_funcify_default_op_cache_key(Gemv) +def numba_funcify_Gemv(op, node, **kwargs): + """Dispatch ``Gemv`` to a single BLAS call, with its scalars carried as gemm's own + alpha and beta.""" + # The vectors reach `_gemm` as one-column matrices rather than through a separate + # gemv binding: BLAS reads the same buffers either way, and gemm already resolves + # each operand's memory order without copying. + if op.inplace: + + @numba_basic.numba_njit + def gemv(y, alpha, A, x, beta): + _gemm( + A, + np.expand_dims(x, 1), + np.expand_dims(y, 1), + False, + False, + alpha.item(), + beta.item(), + ) + return y + + else: + + @numba_basic.numba_njit + def gemv(y, alpha, A, x, beta): + # Accumulating into a copy leaves `y` intact, which is the whole difference + # between this op and its inplace form. + out = y.copy() + _gemm( + A, + np.expand_dims(x, 1), + np.expand_dims(out, 1), + False, + False, + alpha.item(), + beta.item(), + ) + return out + + cache_version = 1 + return gemv, cache_version + + +@register_funcify_default_op_cache_key(Ger) +def numba_funcify_Ger(op, node, **kwargs): + """Dispatch ``Ger`` to one BLAS rank-1 update.""" + dtype = node.outputs[0].type.numpy_dtype + + if op.inplace: + + @numba_basic.numba_njit + def ger(A, alpha, x, y): + return _ger(alpha.item(), x, y, A) + + else: + + @numba_basic.numba_njit + def ger(A, alpha, x, y): + # Writing `A` and the update together keeps this to one pass over the + # output; copying `A` in and letting BLAS accumulate on top would touch it + # twice. Leaving `A` itself alone is the whole difference between this op + # and its inplace form. + rows = x.shape[0] + cols = y.shape[0] + out = np.empty((rows, cols), dtype=dtype) + a = alpha.item() + for i in range(rows): + scaled = a * x[i] + for j in range(cols): + out[i, j] = A[i, j] + scaled * y[j] + return out + + # Bump whenever `_ger` changes: it is inlined here, so its source is not part of + # this key. + cache_version = 4 + return ger, cache_version diff --git a/pytensor/link/numba/dispatch/elemwise.py b/pytensor/link/numba/dispatch/elemwise.py index 029b7e42e0..d6982b6e50 100644 --- a/pytensor/link/numba/dispatch/elemwise.py +++ b/pytensor/link/numba/dispatch/elemwise.py @@ -22,6 +22,7 @@ register_funcify_and_cache_key, register_funcify_default_op_cache_key, ) +from pytensor.link.numba.dispatch.linalg.products import _gemm from pytensor.link.numba.dispatch.string_codegen import ( CODE_TOKEN, build_source_code, @@ -1256,6 +1257,11 @@ def argmax(x): return argmax, cache_version +_GEMM_DTYPES = frozenset( + np.dtype(name) for name in ("float32", "float64", "complex64", "complex128") +) + + @register_funcify_default_op_cache_key(Dot) def numba_funcify_Dot(op, node, **kwargs): # Numba's `np.dot` does not support integer dtypes, so we need to cast to float. @@ -1280,7 +1286,55 @@ def numba_funcify_Dot(op, node, **kwargs): f"{x_dtype=}, {y_dtype=}, {out_dtype=}, {numba_dot_dtype=}" ) - if x_dtype == numba_dot_dtype and y_dtype == numba_dot_dtype: + cast_x = x_dtype != numba_dot_dtype + cast_y = y_dtype != numba_dot_dtype + + if numba_dot_dtype in _GEMM_DTYPES: + # `gemm` reads each operand's memory order as a transpose flag, so an + # operand that reaches here transposed costs nothing, where `np.dot` would + # have to be handed a contiguous copy of it. + if not cast_x and not cast_y: + + @numba_basic.numba_njit + def dot(x, y, out=None): + if out is None: + out = np.empty((x.shape[0], y.shape[1]), dtype=numba_dot_dtype) + return _gemm(x, y, out, False, False, 1.0, 0.0) + + elif cast_x and not cast_y: + + @numba_basic.numba_njit + def dot(x, y, out=None): + if out is None: + out = np.empty((x.shape[0], y.shape[1]), dtype=numba_dot_dtype) + return _gemm(x.astype(numba_dot_dtype), y, out, False, False, 1.0, 0.0) + + elif not cast_x and cast_y: + + @numba_basic.numba_njit + def dot(x, y, out=None): + if out is None: + out = np.empty((x.shape[0], y.shape[1]), dtype=numba_dot_dtype) + return _gemm(x, y.astype(numba_dot_dtype), out, False, False, 1.0, 0.0) + + else: + + @numba_basic.numba_njit + def dot(x, y, out=None): + if out is None: + out = np.empty((x.shape[0], y.shape[1]), dtype=numba_dot_dtype) + return _gemm( + x.astype(numba_dot_dtype), + y.astype(numba_dot_dtype), + out, + False, + False, + 1.0, + 0.0, + ) + + elif not cast_x and not cast_y: + # A dtype BLAS has no kind for, float16 being the one that reaches here. @numba_basic.numba_njit def dot(x, y, out=None): @@ -1289,22 +1343,22 @@ def dot(x, y, out=None): np.dot(x, y, out) return out - elif x_dtype == numba_dot_dtype and y_dtype != numba_dot_dtype: + elif cast_x and not cast_y: @numba_basic.numba_njit def dot(x, y, out=None): if out is None: - return np.asarray(np.dot(x, y.astype(numba_dot_dtype))) - np.dot(x, y.astype(numba_dot_dtype), out) + return np.asarray(np.dot(x.astype(numba_dot_dtype), y)) + np.dot(x.astype(numba_dot_dtype), y, out) return out - elif x_dtype != numba_dot_dtype and y_dtype == numba_dot_dtype: + elif not cast_x and cast_y: @numba_basic.numba_njit def dot(x, y, out=None): if out is None: - return np.asarray(np.dot(x.astype(numba_dot_dtype), y)) - np.dot(x.astype(numba_dot_dtype), y, out) + return np.asarray(np.dot(x, y.astype(numba_dot_dtype))) + np.dot(x, y.astype(numba_dot_dtype), out) return out else: @@ -1318,15 +1372,17 @@ def dot(x, y, out=None): np.dot(x.astype(numba_dot_dtype), y.astype(numba_dot_dtype), out) return out - cache_version = 2 + # Bump whenever `_gemm` changes: it is inlined here, so its source is not + # part of this key. + cache_version = 6 if out_dtype == numba_dot_dtype: - # np.dot can write straight into the pre-allocated batch output slice. + # The product can be written straight into the pre-allocated batch output slice. dot.handles_out = True return dot, cache_version else: - # Output needs a dtype cast np.dot can't do in place, so fall back to + # Output needs a dtype cast the product can't do in place, so fall back to # the copying store_core_outputs wrapper. @numba_basic.numba_njit def dot_with_cast(x, y): @@ -1338,18 +1394,34 @@ def dot_with_cast(x, y): @register_funcify_default_op_cache_key(BatchedDot) def numba_funcify_BatchedDot(op, node, **kwargs): dtype = node.outputs[0].type.numpy_dtype + x, y = node.inputs + + # Numba does not support 3D matmul (numba#3804), so either path loops over the batch. + if x.type.numpy_dtype == y.type.numpy_dtype == dtype and dtype in _GEMM_DTYPES: - @numba_basic.numba_njit - def batched_dot(x, y, out=None): - # Numba does not support 3D matmul - # https://github.com/numba/numba/issues/3804 - if out is None: - shape = x.shape[:-1] + y.shape[2:] - out = np.empty(shape, dtype=dtype) - for i in range(out.shape[0]): - out[i] = np.dot(x[i], y[i]) + @numba_basic.numba_njit + def batched_dot(x, y, out=None): + if out is None: + shape = x.shape[:-1] + y.shape[2:] + out = np.empty(shape, dtype=dtype) + for i in range(out.shape[0]): + # Each slice keeps whatever layout its batch carries, and gemm reads that as a + # transpose flag rather than copying the slice to satisfy `np.dot`. + _gemm(x[i], y[i], out[i], False, False, 1.0, 0.0) - return out + return out + + else: + + @numba_basic.numba_njit + def batched_dot(x, y, out=None): + if out is None: + shape = x.shape[:-1] + y.shape[2:] + out = np.empty(shape, dtype=dtype) + for i in range(out.shape[0]): + out[i] = np.dot(x[i], y[i]) + + return out batched_dot.handles_out = True - return batched_dot, 1 + return batched_dot, 2 diff --git a/pytensor/link/numba/dispatch/linalg/_BLAS.py b/pytensor/link/numba/dispatch/linalg/_BLAS.py index faa558f330..0d02cfc598 100644 --- a/pytensor/link/numba/dispatch/linalg/_BLAS.py +++ b/pytensor/link/numba/dispatch/linalg/_BLAS.py @@ -72,3 +72,107 @@ def trsm(SIDE, UPLO, TRANSA, DIAG, M, N, ALPHA, A, LDA, B, LDB): fn(SIDE, UPLO, TRANSA, DIAG, M, N, ALPHA, A, LDA, B, LDB) return trsm + + @classmethod + def numba_xgemm(cls, dtype) -> CPUDispatcher: + r""" + Compute a general matrix-matrix product, overwriting :math:`C`. + + .. math:: + + C \leftarrow \alpha \, op(A) \, op(B) + \beta C + + where :math:`op(X)` is :math:`X`, :math:`X^T` or :math:`X^H` according to + ``TRANSA`` and ``TRANSB``. Taking the transposes as flags is the point of + binding this directly: BLAS reads either operand in transposed order at no + cost, so a caller never has to materialize one. + """ + + kind = get_blas_kind(dtype) + float_ptr = _get_nb_float_from_dtype(kind) + unique_func_name = f"scipy.blas.{kind}gemm" + + @numba_basic.numba_njit + def get_gemm_pointer(): + with numba.objmode(ptr=types.intp): + ptr = get_blas_ptr(dtype, "gemm") + return ptr + + gemm_function_type = types.FunctionType( + types.void( + nb_i32p, # TRANSA + nb_i32p, # TRANSB + nb_i32p, # M + nb_i32p, # N + nb_i32p, # K + float_ptr, # ALPHA + float_ptr, # A + nb_i32p, # LDA + float_ptr, # B + nb_i32p, # LDB + float_ptr, # BETA + float_ptr, # C + nb_i32p, # LDC + ) + ) + + @numba_basic.numba_njit + def gemm(TRANSA, TRANSB, M, N, K, ALPHA, A, LDA, B, LDB, BETA, C, LDC): + fn = _call_cached_ptr( + get_ptr_func=get_gemm_pointer, + func_type_ref=gemm_function_type, + unique_func_name_lit=unique_func_name, + ) + fn(TRANSA, TRANSB, M, N, K, ALPHA, A, LDA, B, LDB, BETA, C, LDC) + + return gemm + + @classmethod + def numba_xger(cls, dtype) -> CPUDispatcher: + r""" + Add a rank-1 update to a general matrix, overwriting :math:`A`. + + .. math:: + + A \leftarrow \alpha \, x \, y^T + A + + BLAS names this ``ger`` only for the real kinds; the complex ones split it + into an unconjugated ``geru`` and a conjugated ``gerc``, and this binds the + unconjugated one so every dtype computes the same :math:`x y^T`. + """ + + kind = get_blas_kind(dtype) + float_ptr = _get_nb_float_from_dtype(kind) + name = "ger" if kind in "sd" else "geru" + unique_func_name = f"scipy.blas.{kind}{name}" + + @numba_basic.numba_njit + def get_ger_pointer(): + with numba.objmode(ptr=types.intp): + ptr = get_blas_ptr(dtype, name) + return ptr + + ger_function_type = types.FunctionType( + types.void( + nb_i32p, # M + nb_i32p, # N + float_ptr, # ALPHA + float_ptr, # X + nb_i32p, # INCX + float_ptr, # Y + nb_i32p, # INCY + float_ptr, # A + nb_i32p, # LDA + ) + ) + + @numba_basic.numba_njit + def ger(M, N, ALPHA, X, INCX, Y, INCY, A, LDA): + fn = _call_cached_ptr( + get_ptr_func=get_ger_pointer, + func_type_ref=ger_function_type, + unique_func_name_lit=unique_func_name, + ) + fn(M, N, ALPHA, X, INCX, Y, INCY, A, LDA) + + return ger diff --git a/pytensor/link/numba/dispatch/linalg/products.py b/pytensor/link/numba/dispatch/linalg/products.py index 7c2561d6fb..41c0247a97 100644 --- a/pytensor/link/numba/dispatch/linalg/products.py +++ b/pytensor/link/numba/dispatch/linalg/products.py @@ -1,12 +1,13 @@ import numpy as np from numba.core.extending import overload from numba.core.types import Complex, Float -from numba.np.linalg import _copy_to_fortran_order, ensure_lapack +from numba.np.linalg import _copy_to_fortran_order, ensure_blas, ensure_lapack from scipy import linalg from pytensor import config from pytensor.link.numba.dispatch import basic as numba_basic from pytensor.link.numba.dispatch.basic import register_funcify_default_op_cache_key +from pytensor.link.numba.dispatch.linalg._BLAS import _BLAS from pytensor.link.numba.dispatch.linalg._LAPACK import ( _LAPACK, _get_underlying_float, @@ -328,3 +329,228 @@ def expm(a): cache_version = 1 return expm, cache_version + + +def _gemm(A, B, C, transa=False, transb=False, alpha=1.0, beta=0.0): + r""" + Compute a general matrix-matrix product, overwriting ``C``. + + .. math:: + + C \leftarrow \alpha \, op(A) \, op(B) + \beta C + + Parameters + ---------- + A, B : ndarray + The operands, each 2-d and of the same dtype as ``C``. + C : ndarray + The output, overwritten in place and also returned. Either memory order + works; a row-major ``C`` is handled by computing the transposed product + rather than by copying. + transa, transb : bool, optional + Transpose the corresponding operand. The transpose is a flag BLAS reads, not + an array that gets built, so neither operand is ever copied to apply one. + Both default to False. + alpha, beta : scalar, optional + Scale the product and the incoming ``C`` respectively. ``alpha`` defaults to + 1.0 and ``beta`` to 0.0, which overwrites ``C`` rather than accumulating into + it. + + Returns + ------- + C : ndarray + The same array passed in. + """ + op_a = A.T if transa else A + op_b = B.T if transb else B + product = alpha * (op_a @ op_b) + # BLAS leaves C unread when beta is zero, so callers are free to pass a buffer they never + # initialized. Scaling that buffer by zero would turn stray inf or nan bytes into nan here. + C[:] = product if beta == 0 else product + beta * C + return C + + +@overload(_gemm) +def _gemm_impl(A, B, C, transa, transb, alpha, beta): + ensure_blas() + _check_linalg_matrix(A, ndim=2, dtype=(Float, Complex), func_name="gemm") + _check_linalg_matrix(B, ndim=2, dtype=(Float, Complex), func_name="gemm") + _check_linalg_matrix(C, ndim=2, dtype=(Float, Complex), func_name="gemm") + + numba_gemm = _BLAS().numba_xgemm(A.dtype) + dtype = A.dtype + + def impl(A, B, C, transa, transb, alpha, beta): + # BLAS reads column-major, so a C-ordered array's buffer already *is* its + # transpose. Folding that into the flag each operand carries keeps every case + # copy-free: only an array that is neither C- nor F-contiguous is materialized. + if A.flags.f_contiguous: + A_work = A + A_trans = transa + LDA = np.int32(max(1, A.shape[0])) + elif A.flags.c_contiguous: + A_work = A + A_trans = not transa + LDA = np.int32(max(1, A.shape[1])) + else: + A_work = _copy_to_fortran_order(A) + A_trans = transa + LDA = np.int32(max(1, A.shape[0])) + + if B.flags.f_contiguous: + B_work = B + B_trans = transb + LDB = np.int32(max(1, B.shape[0])) + elif B.flags.c_contiguous: + B_work = B + B_trans = not transb + LDB = np.int32(max(1, B.shape[1])) + else: + B_work = _copy_to_fortran_order(B) + B_trans = transb + LDB = np.int32(max(1, B.shape[0])) + + # M, N and K describe the logical product, so they come from the caller's + # transposes rather than the layout-adjusted ones. + M = np.int32(A.shape[1] if transa else A.shape[0]) + K = np.int32(A.shape[0] if transa else A.shape[1]) + N = np.int32(B.shape[0] if transb else B.shape[1]) + + # BLAS trusts the extents it is handed, so a mismatch here reads past the end + # of an operand rather than failing. numpy raises for these inputs, so + # must this. + if (B.shape[1] if transb else B.shape[0]) != K: + raise ValueError("gemm: operands have mismatched contraction dimensions") + if C.shape[0] != M or C.shape[1] != N: + raise ValueError("gemm: output shape does not match the product") + + ALPHA = np.full(1, alpha, dtype=dtype) + BETA = np.full(1, beta, dtype=dtype) + + if C.flags.f_contiguous: + numba_gemm( + val_to_int_ptr(ord("T") if A_trans else ord("N")), + val_to_int_ptr(ord("T") if B_trans else ord("N")), + val_to_int_ptr(M), + val_to_int_ptr(N), + val_to_int_ptr(K), + ALPHA.ctypes, + A_work.ctypes, + val_to_int_ptr(LDA), + B_work.ctypes, + val_to_int_ptr(LDB), + BETA.ctypes, + C.ctypes, + val_to_int_ptr(np.int32(max(1, C.shape[0]))), + ) + else: + # A row-major C is C^T to BLAS, and C^T = op(B)^T op(A)^T, so the operands + # swap places and each flag flips. This is the case pytensor hits, since + # its outputs are C-ordered. A C that is neither C- nor F-contiguous has no + # leading dimension BLAS can address, so it goes through a copy. + needs_writeback = not C.flags.c_contiguous + C_work = np.ascontiguousarray(C) if needs_writeback else C + numba_gemm( + val_to_int_ptr(ord("N") if B_trans else ord("T")), + val_to_int_ptr(ord("N") if A_trans else ord("T")), + val_to_int_ptr(N), + val_to_int_ptr(M), + val_to_int_ptr(K), + ALPHA.ctypes, + B_work.ctypes, + val_to_int_ptr(LDB), + A_work.ctypes, + val_to_int_ptr(LDA), + BETA.ctypes, + C_work.ctypes, + val_to_int_ptr(np.int32(max(1, C_work.shape[1]))), + ) + if needs_writeback: + C[:] = C_work + return C + + return impl + + +def _ger(alpha, x, y, A): + r""" + Add a rank-1 update to a general matrix, overwriting ``A``. + + .. math:: + + A \leftarrow \alpha \, x \, y^T + A + + Parameters + ---------- + alpha : scalar + Scale the outer product. + x, y : ndarray + The 1-d factors, of lengths matching ``A``'s rows and columns respectively. + A : ndarray + The matrix updated in place and also returned. + + Returns + ------- + A : ndarray + The same array passed in. + """ + A[:] = alpha * np.outer(x, y) + A + return A + + +@overload(_ger) +def _ger_impl(alpha, x, y, A): + ensure_blas() + _check_linalg_matrix(A, ndim=2, dtype=(Float, Complex), func_name="ger") + + numba_ger = _BLAS().numba_xger(A.dtype) + dtype = A.dtype + + def impl(alpha, x, y, A): + ALPHA = np.full(1, alpha, dtype=dtype) + INC = val_to_int_ptr(np.int32(1)) + + # ger walks each vector by a fixed increment of one element, so a strided view + # reads the wrong entries unless it is made contiguous first. The vectors are + # O(m + n) against the update's O(m n), so copying one costs less than carrying + # its own increment would, and it keeps a negative stride -- which BLAS addresses + # from the far end of the buffer -- out of the picture entirely. + x_work = x if x.flags.c_contiguous else np.ascontiguousarray(x) + y_work = y if y.flags.c_contiguous else np.ascontiguousarray(y) + + # ger has no transpose flag and does not need one: a C-ordered A's buffer is + # A^T, and transposing the identity gives A^T <- alpha y x^T + A^T. Swapping + # the vectors and the extents computes that, updating a row-major A in place. + if A.flags.f_contiguous: + numba_ger( + val_to_int_ptr(np.int32(A.shape[0])), + val_to_int_ptr(np.int32(A.shape[1])), + ALPHA.ctypes, + x_work.ctypes, + INC, + y_work.ctypes, + INC, + A.ctypes, + val_to_int_ptr(np.int32(max(1, A.shape[0]))), + ) + else: + # An array that is neither C- nor F-contiguous has to be updated through a + # copy, so the result is written back to keep the in-place contract. + needs_writeback = not A.flags.c_contiguous + A_work = np.ascontiguousarray(A) if needs_writeback else A + numba_ger( + val_to_int_ptr(np.int32(A_work.shape[1])), + val_to_int_ptr(np.int32(A_work.shape[0])), + ALPHA.ctypes, + y_work.ctypes, + INC, + x_work.ctypes, + INC, + A_work.ctypes, + val_to_int_ptr(np.int32(max(1, A_work.shape[1]))), + ) + if needs_writeback: + A[:] = A_work + return A + + return impl diff --git a/pytensor/link/pytorch/dispatch/blas.py b/pytensor/link/pytorch/dispatch/blas.py index 5691551998..389286a609 100644 --- a/pytensor/link/pytorch/dispatch/blas.py +++ b/pytensor/link/pytorch/dispatch/blas.py @@ -1,7 +1,8 @@ import torch from pytensor.link.pytorch.dispatch import pytorch_funcify -from pytensor.tensor.blas import BatchedDot +from pytensor.link.utils import get_static_scalar +from pytensor.tensor.blas import BatchedDot, Gemm, Gemv, Ger @pytorch_funcify.register(BatchedDot) @@ -12,3 +13,56 @@ def batched_dot(a, b): return torch.bmm(a, b) return batched_dot + + +@pytorch_funcify.register(Gemm) +def pytorch_funcify_Gemm(op, node=None, **kwargs): + static_alpha = get_static_scalar(node, 1) + static_beta = get_static_scalar(node, 4) + + if static_alpha is not None and static_beta is not None: + + def gemm(z, alpha, x, y, beta): + return torch.addmm(z, x, y, beta=static_beta, alpha=static_alpha) + + else: + + def gemm(z, alpha, x, y, beta): + return beta * z + alpha * torch.matmul(x, y) + + return gemm + + +@pytorch_funcify.register(Gemv) +def pytorch_funcify_Gemv(op, node=None, **kwargs): + static_alpha = get_static_scalar(node, 1) + static_beta = get_static_scalar(node, 4) + + if static_alpha is not None and static_beta is not None: + + def gemv(y, alpha, A, x, beta): + return torch.addmv(y, A, x, beta=static_beta, alpha=static_alpha) + + else: + + def gemv(y, alpha, A, x, beta): + return beta * y + alpha * torch.matmul(A, x) + + return gemv + + +@pytorch_funcify.register(Ger) +def pytorch_funcify_Ger(op, node=None, **kwargs): + static_alpha = get_static_scalar(node, 1) + + if static_alpha is not None: + + def ger(A, alpha, x, y): + return torch.addr(A, x, y, beta=1.0, alpha=static_alpha) + + else: + + def ger(A, alpha, x, y): + return A + alpha * torch.outer(x, y) + + return ger diff --git a/pytensor/link/utils.py b/pytensor/link/utils.py index 624d53969b..184b5394a4 100644 --- a/pytensor/link/utils.py +++ b/pytensor/link/utils.py @@ -845,3 +845,39 @@ def get_destroy_dependencies(fgraph: FunctionGraph) -> dict[Apply, list[Variable for prereq in order.get(node, []): destroy_dependencies[node].extend(prereq.outputs) return destroy_dependencies + + +def get_static_scalar(node: Apply | None, input_index: int) -> float | None: + """Return one of a node's inputs as a Python float, when it is constant at compile time. + + Backends use this to decide whether a scalar can be baked into a fused kernel call, + several of which take their scaling factors as plain floats rather than as arrays. + + Parameters + ---------- + node : Apply or None + The node being dispatched. None when the backend was given no node. + input_index : int + Position of the input to resolve. + + Returns + ------- + float or None + The value, or None when that input is not a compile-time scalar or cannot be + expressed as a float. + """ + from pytensor.tensor.basic import get_underlying_scalar_constant_value + + if node is None: + return None + + value = get_underlying_scalar_constant_value( + node.inputs[input_index], raise_not_constant=False + ) + if not isinstance(value, np.ndarray): + return None + try: + return float(value) + except TypeError: + # a complex constant, which no caller can pass on as a float + return None diff --git a/pytensor/link/vm.py b/pytensor/link/vm.py index 239f73df80..a2dfe91f43 100644 --- a/pytensor/link/vm.py +++ b/pytensor/link/vm.py @@ -1210,6 +1210,8 @@ def make_all( output_storage=None, storage_map=None, ): + from pytensor.link.c.dispatch.basic import make_node_thunk_with_c_dispatch + fgraph = self.fgraph order = self.schedule(fgraph) @@ -1227,6 +1229,7 @@ def make_all( impl = None if self.c_thunks is False: impl = "py" + use_c_dispatch = self.c_thunks is not False and bool(config.cxx) for node in order: try: thunk_start = time.perf_counter() @@ -1234,7 +1237,14 @@ def make_all( # no need to cause duplicate c code by passing # no_recycling here. thunks.append( - node.op.make_thunk(node, storage_map, compute_map, [], impl=impl) + make_node_thunk_with_c_dispatch( + node, + storage_map, + compute_map, + [], + try_c=use_c_dispatch, + fallback_impl=impl, + ) ) linker_make_thunk_time[node] = time.perf_counter() - thunk_start if not hasattr(thunks[-1], "lazy"): diff --git a/pytensor/tensor/basic.py b/pytensor/tensor/basic.py index 52f63d9e34..eebfaded8b 100644 --- a/pytensor/tensor/basic.py +++ b/pytensor/tensor/basic.py @@ -1763,7 +1763,7 @@ def do_constant_folding(self, fgraph, node): if not clients: return False - from pytensor.tensor.blas import CGemv, CGer, Gemv, Ger + from pytensor.tensor.blas import Gemv, Ger from pytensor.tensor.subtensor import ( AdvancedIncSubtensor, IncSubtensor, @@ -1796,7 +1796,7 @@ def do_constant_folding(self, fgraph, node): idx == 0 and isinstance( client_op, - IncSubtensor | AdvancedIncSubtensor | Gemv | CGemv | Ger | CGer, + IncSubtensor | AdvancedIncSubtensor | Gemv | Ger, ) ): # Ops that will work inplace on the Alloc. So if they diff --git a/pytensor/tensor/blas/__init__.py b/pytensor/tensor/blas/__init__.py index 7b014d6b04..254532c42b 100644 --- a/pytensor/tensor/blas/__init__.py +++ b/pytensor/tensor/blas/__init__.py @@ -11,17 +11,8 @@ Where there is a discrepancy between how things do work and how they *should* work, both aspects should be documented. -There are four kinds of BLAS Ops in PyTensor: - - Python implementations (this file) - - SciPy-based (blas_scipy) - - C-based (blas_c) - -Notes ------ -Unfortunately (because it's confusing) this file currently contains Ops -that contain both Python and C versions. I think it would be better to -move the C implementations to blas_c so that this file is pure Python. --JB +The Ops here are pure Python; their SciPy and C implementations live in the +corresponding backend dispatch registries. Ops @@ -90,18 +81,6 @@ BatchedDot, _batched_dot, ) -from pytensor.tensor.blas.blas_c import ( - BaseBLAS, - CGemv, - CGer, - cgemv_inplace, - cgemv_no_inplace, - cger_inplace, - cger_no_inplace, -) -from pytensor.tensor.blas.blas_c import ( - must_initialize_y_gemv as must_initialize_y_gemv_c, -) from pytensor.tensor.blas.c_code.blas_headers import ( blas_header_text, blas_header_version, @@ -115,19 +94,13 @@ GemmRelated, _dot22, _dot22scalar, - gemm, - gemm_inplace, - gemm_no_inplace, ) -from pytensor.tensor.blas.gemv import Gemv, gemv, gemv_inplace, gemv_no_inplace -from pytensor.tensor.blas.ger import Ger, ger, ger_destructive +from pytensor.tensor.blas.gemv import Gemv +from pytensor.tensor.blas.ger import Ger __all__ = [ - "BaseBLAS", "BatchedDot", - "CGemv", - "CGer", "Dot22", "Dot22Scalar", "Gemm", @@ -141,22 +114,9 @@ "_logger", "blas_header_text", "blas_header_version", - "cgemv_inplace", - "cgemv_no_inplace", - "cger_inplace", - "cger_no_inplace", - "gemm", - "gemm_inplace", - "gemm_no_inplace", - "gemv", - "gemv_inplace", - "gemv_no_inplace", - "ger", - "ger_destructive", "ldflags", "mkl_threads_text", "must_initialize_y_gemv", - "must_initialize_y_gemv_c", "openblas_threads_text", "view_roots", ] diff --git a/pytensor/tensor/blas/_core.py b/pytensor/tensor/blas/_core.py index 08aeff7766..ebea31012c 100644 --- a/pytensor/tensor/blas/_core.py +++ b/pytensor/tensor/blas/_core.py @@ -48,7 +48,7 @@ def must_initialize_y_gemv(): must_initialize_y_gemv._result = None # type: ignore -def ldflags(libs=True, flags=False, libs_dir=False, include_dir=False): +def ldflags(libs=True, flags=False, libs_dir=False, include_dir=False) -> list[str]: """Extract a list of compilation flags from config.blas__ldflags. Depending on the options, different type of flags will be kept. diff --git a/pytensor/tensor/blas/blas_c.py b/pytensor/tensor/blas/blas_c.py deleted file mode 100644 index bf9eaac88a..0000000000 --- a/pytensor/tensor/blas/blas_c.py +++ /dev/null @@ -1,134 +0,0 @@ -from pytensor.link.c.op import COp -from pytensor.link.c.params_type import ParamsType -from pytensor.scalar import bool as bool_t -from pytensor.tensor.blas._core import ldflags -from pytensor.tensor.blas.c_code.blas_headers import ( - blas_header_text, - blas_header_version, -) -from pytensor.tensor.blas.c_code.codegen import gemv_c_code, ger_c_code -from pytensor.tensor.blas.gemv import Gemv -from pytensor.tensor.blas.ger import Ger - - -class BaseBLAS(COp): - def c_libraries(self, **kwargs): - return ldflags() - - def c_compile_args(self, **kwargs): - return ldflags(libs=False, flags=True) - - def c_lib_dirs(self, **kwargs): - return ldflags(libs=False, libs_dir=True) - - def c_header_dirs(self, **kwargs): - return ldflags(libs=False, include_dir=True) - - def c_support_code(self, **kwargs): - return blas_header_text() - - -# ##### ####### ####### -# GER -# ##### ####### ####### - - -class CGer(BaseBLAS, Ger): - """C implementation of GER (rank-1 update): Z = A + alpha * outer(x, y).""" - - params_type = ParamsType( - destructive=bool_t, - ) - - def c_code(self, node, name, inp, out, sub): - return ger_c_code(node, name, inp, out, sub) - - def c_code_cache_version(self): - return (11, blas_header_version()) - - -cger_inplace = CGer(True) -cger_no_inplace = CGer(False) - - -# ##### ####### ####### -# GEMV -# ##### ####### ####### - - -class CGemv(BaseBLAS, Gemv): - params_type = ParamsType( - inplace=bool_t, - ) - - def __init__(self, inplace): - super().__init__(inplace) - - def c_code(self, node, name, inp, out, sub): - return gemv_c_code(node, name, inp, out, sub) - - def c_code_cache_version(self): - return (18, blas_header_version(), must_initialize_y_gemv()) - - -cgemv_inplace = CGemv(inplace=True) -cgemv_no_inplace = CGemv(inplace=False) - - -def must_initialize_y_gemv(): - if must_initialize_y_gemv._force_init_beta is None: - from pytensor.link.c.cmodule import GCC_compiler - - """ - Test issue 1569. - Namely when evaluating - - beta*y + alpha*dot(A, x) - - where we set y * beta = zeros of the correct dimensions we - do not actually set y = zeros and instead let the BLAS - perform beta*y with uninitialized memory for - speed. Occasionally the memory contains values that are - equivalent to NaN in which case the product beta*y contains - NaN's for correctly implemented BLAS libraries. In this - situation, since we are introducing the NaN's, we need to test - whether the BLAS performs correctly. If it *does*, i.e. it - actually performs the multiplication beta*y which will result - in NaN's in the result, then we need initialize the memory to - zeros. - """ - test_code = """ -#include -extern "C" void dgemv_(char*, const int*, const int*, const double *, const double *, const int*, const double *, const int*, const double *, double *, const int *); -int main() { - double A[2][2] = {{1., 1.}, {1., 1.}}; - double x[2] = {1., 1.}; - double y[2] = {NAN, NAN}; - const int s = 2; - const int inc = 1; - const double alpha = 1.0; - const double beta = 0.0; - - dgemv_("T", &s, &s, &alpha, A, &s, x, &inc, &beta, &y, &inc); - - return (isnan(y[0]) || isnan(y[1]) ? 1 : 0; -} -""" - res = GCC_compiler.try_compile_tmp( - test_code, - tmp_prefix="check_beta_", - flags=ldflags(libs=True, flags=True, libs_dir=True), - try_run=True, - ) - if res: - if res[0]: - must_initialize_y_gemv._force_init_beta = res[1] - else: - must_initialize_y_gemv._force_init_beta = False - else: - must_initialize_y_gemv._force_init_beta = False - - return must_initialize_y_gemv._force_init_beta - - -must_initialize_y_gemv._force_init_beta = None # type: ignore[attr-defined] diff --git a/pytensor/tensor/blas/c_code/blas_headers.py b/pytensor/tensor/blas/c_code/blas_headers.py index 973c9feccb..13d692f442 100644 --- a/pytensor/tensor/blas/c_code/blas_headers.py +++ b/pytensor/tensor/blas/c_code/blas_headers.py @@ -740,7 +740,7 @@ def cblas_header_text(): """ -def blas_header_text(): +def blas_header_text() -> str: """C header for the fortran blas interface""" blas_code = "" @@ -1054,9 +1054,9 @@ def openblas_threads_text(): return header -def blas_header_version(): +def blas_header_version() -> tuple[int, ...]: # Version for the base header - version = (10,) + version: tuple[int, ...] = (10,) if detect_macos_sdot_bug(): if detect_macos_sdot_bug.fix_works: # Version with fix diff --git a/pytensor/tensor/blas/c_code/codegen.py b/pytensor/tensor/blas/c_code/codegen.py index 2952f36282..a79ecfde2b 100644 --- a/pytensor/tensor/blas/c_code/codegen.py +++ b/pytensor/tensor/blas/c_code/codegen.py @@ -12,6 +12,9 @@ # silently get stale compiled binaries. +from pytensor.tensor.blas._core import must_initialize_y_gemv + + # ##### ####### ####### # GEMM family (Gemm, Dot22, Dot22Scalar) # ##### ####### ####### @@ -462,7 +465,7 @@ def _assemble_gemm_call( *, setup_z, check_ab, broadcast_xy, ab_constants_float, ab_constants_double -): +) -> str: """Concatenate the GEMM template fragments in execution order.""" return "".join( ( @@ -488,7 +491,7 @@ def _assemble_gemm_call( ) -def gemm_c_code(node, name, inputs, outputs, sub): +def gemm_c_code(node, name, inputs, outputs, sub) -> str: r"""C code for ``Gemm``: :math:`z \leftarrow b\,z + a\,xy` (in/out-of-place).""" _z, _a, _x, _y, _b = inputs (_zout,) = outputs @@ -506,7 +509,7 @@ def gemm_c_code(node, name, inputs, outputs, sub): return code % dict(_z=_z, _a=_a, _x=_x, _y=_y, _b=_b, _zout=_zout, **sub) -def dot22_c_code(node, name, inputs, outputs, sub): +def dot22_c_code(node, name, inputs, outputs, sub) -> str: r"""C code for ``Dot22``: :math:`z \leftarrow xy`, allocating a fresh output.""" _x, _y = inputs (_zout,) = outputs @@ -539,15 +542,12 @@ def dot22scalar_c_code(node, name, inputs, outputs, sub): # ##### ####### ####### -def gemv_c_code(node, name, inputs, outputs, sub): - r"""C code for ``CGemv``: :math:`z \leftarrow \beta\,y + \alpha\,Ax`. +def gemv_c_code(node, name, inputs, outputs, sub) -> str: + r"""C code for ``Gemv``: :math:`z \leftarrow \beta\,y + \alpha\,Ax`. ``z`` aliases ``y`` when inplace, otherwise a fresh copy; :math:`A` is a matrix and :math:`x`, :math:`y` are vectors. """ - # Imported lazily to avoid an import cycle (blas_c imports this module). - from pytensor.tensor.blas.blas_c import must_initialize_y_gemv - y, alpha, A, x, beta = inputs (z,) = outputs must_initialize_y = must_initialize_y_gemv() @@ -804,8 +804,8 @@ def gemv_c_code(node, name, inputs, outputs, sub): # ##### ####### ####### -def ger_c_code(node, name, inputs, outputs, sub): - r"""C code for ``CGer``: rank-1 update :math:`Z = A + \alpha\,x y^{\top}`.""" +def ger_c_code(node, name, inputs, outputs, sub) -> str: + r"""C code for ``Ger``: rank-1 update :math:`Z = A + \alpha\,x y^{\top}`.""" A, a, x, y = inputs (Z,) = outputs fail = sub["fail"] @@ -845,12 +845,12 @@ def ger_c_code(node, name, inputs, outputs, sub): else if (PyArray_DESCR({A})->type_num == NPY_FLOAT) {{ elemsize = 4;}} else {{ - PyErr_SetString(PyExc_NotImplementedError, "complex CGer"); + PyErr_SetString(PyExc_NotImplementedError, "complex Ger"); {fail}; }} - // copy A if !self.destructive or A is fully strided - if (!{params}->destructive + // copy A if !self.inplace or A is fully strided + if (!{params}->inplace || (PyArray_STRIDES({A})[0] < 0) || (PyArray_STRIDES({A})[1] < 0) || ((PyArray_STRIDES({A})[0] != elemsize) diff --git a/pytensor/tensor/blas/gemm.py b/pytensor/tensor/blas/gemm.py index adfeedffd1..1f82e77f75 100644 --- a/pytensor/tensor/blas/gemm.py +++ b/pytensor/tensor/blas/gemm.py @@ -2,11 +2,10 @@ import pytensor.scalar from pytensor.graph.basic import Apply +from pytensor.graph.op import Op from pytensor.graph.utils import InconsistencyError, MethodNotDefined from pytensor.link.c.op import COp -from pytensor.link.c.params_type import ParamsType from pytensor.printing import FunctionPrinter, pprint -from pytensor.scalar import bool as bool_t from pytensor.tensor.basic import as_tensor_variable from pytensor.tensor.blas._core import ( ldflags, @@ -19,7 +18,6 @@ from pytensor.tensor.blas.c_code.codegen import ( dot22_c_code, dot22scalar_c_code, - gemm_c_code, ) from pytensor.tensor.type import DenseTensorType, tensor @@ -73,27 +71,16 @@ def build_gemm_version(self): return (14, blas_header_version()) -class Gemm(GemmRelated): - """In-place version of matrix-matrix multiplication (with accumulation). +class Gemm(Op): + r"""Matrix-matrix product with accumulation. - When a and b are scalars and x, y, and z are matrices, then + .. math:: - gemm(z,a,x,y,b) - - is similar to - - b*z + a*dot(x,y) - - The difference between the two is that the top form is destructive - on z, whereas the bottom form is not. Gemm works in-place on the - storage associated with z, and the L{Variable} returned by Gemm - has a storage that will be aliased to the storage of the z - argument. Because of this in-place computation, an L{Apply} of - this op will destroy the L{Variable} z on which it operates. (See - L{DestructiveOps} for an explanation of what destroying means in - the context of pytensor graphs. See L{BlasLapackSupport} for more - optimized linear algebra operations.) + Z \leftarrow \beta Z + \alpha X Y + for matrices :math:`X`, :math:`Y` and :math:`Z` and scalars :math:`\alpha` and + :math:`\beta`. Constructed with ``inplace=True``, the output aliases ``Z``'s storage + and the op destroys it; otherwise ``Z`` is left untouched. """ E_rank = "gemm only works for rank 2" @@ -103,10 +90,8 @@ class Gemm(GemmRelated): E_float = "gemm requires floating-point dtypes" __props__ = ("inplace",) - params_type = ParamsType( - inplace=bool_t, - ) check_input = False + gufunc_signature = "(m,n),(),(m,k),(k,n),()->(m,n)" def __init__(self, inplace): self.inplace = inplace @@ -224,25 +209,15 @@ def infer_shape(self, node, input_shapes): ) ] - def c_code(self, node, name, inp, out, sub): - if node.inputs[0].type.dtype.startswith("complex"): - raise MethodNotDefined(f"{self.__class__.__name__}.c_code") - return gemm_c_code(node, name, inp, out, sub) - - def c_code_cache_version(self): - gv = self.build_gemm_version() - if gv: - return (8, *gv) - else: - return gv + def inplace_on_inputs(self, allowed_inplace_inputs: list[int]) -> Op: + """``Gemm`` accumulates into its first input, so that is the only one it can destroy.""" + if 0 in allowed_inplace_inputs: + return type(self)(inplace=True) + return self -gemm_inplace = Gemm(inplace=True) -gemm_no_inplace = Gemm(inplace=False) -# For the user interface. PyTensor optimization will make them inplace -gemm = gemm_no_inplace -pprint.assign(gemm_inplace, FunctionPrinter(["gemm_inplace"])) -pprint.assign(gemm_no_inplace, FunctionPrinter(["gemm_no_inplace"])) +pprint.assign(Gemm(inplace=True), FunctionPrinter(["gemm_inplace"])) +pprint.assign(Gemm(inplace=False), FunctionPrinter(["gemm_no_inplace"])) class Dot22(GemmRelated): diff --git a/pytensor/tensor/blas/gemv.py b/pytensor/tensor/blas/gemv.py index 44027d47f5..4f874a0e4d 100644 --- a/pytensor/tensor/blas/gemv.py +++ b/pytensor/tensor/blas/gemv.py @@ -1,8 +1,3 @@ -"""BLAS GEMV operation: matrix-vector multiply with accumulation. - -Computes: beta * y + alpha * dot(A, x) -""" - import numpy as np from pytensor.graph.basic import Apply @@ -14,23 +9,31 @@ class Gemv(Op): - """ - expression is beta * y + alpha * A x + r"""Matrix-vector product with accumulation. + + .. math:: - A is matrix - x, y are vectors - alpha, beta are scalars - output is a vector that can be inplace on y + y \leftarrow \beta y + \alpha A x + for matrix :math:`A`, vectors :math:`x` and :math:`y` and scalars :math:`\alpha` and + :math:`\beta`. Constructed with ``inplace=True``, the output aliases ``y``'s storage + and the op destroys it; otherwise ``y`` is left untouched. """ __props__ = ("inplace",) + gufunc_signature = "(m),(),(m,n),(n),()->(m)" def __init__(self, inplace): self.inplace = inplace if inplace: self.destroy_map = {0: [0]} + def inplace_on_inputs(self, allowed_inplace_inputs: list[int]) -> Op: + """``Gemv`` accumulates into ``y``, so that is the only input it can destroy.""" + if 0 in allowed_inplace_inputs: + return type(self)(inplace=True) + return self + def __str__(self): if self.inplace: return f"{self.__class__.__name__}{{inplace}}" @@ -109,9 +112,3 @@ def perform(self, node, inputs, out_storage): def infer_shape(self, node, input_shapes): return [input_shapes[0]] - - -gemv_no_inplace = Gemv(inplace=False) -gemv_inplace = Gemv(inplace=True) -# For the user interface. Opt will make them inplace later -gemv = gemv_no_inplace diff --git a/pytensor/tensor/blas/ger.py b/pytensor/tensor/blas/ger.py index dedea9d287..b5f27abeeb 100644 --- a/pytensor/tensor/blas/ger.py +++ b/pytensor/tensor/blas/ger.py @@ -6,28 +6,36 @@ class Ger(Op): - """ - BLAS defines general rank-1 update GER as A <- A + alpha x y' + r"""Rank-1 update of a matrix. - for matrix A, scalar alpha, vectors x and y. + .. math:: - This interface to GER allows non-destructive operation on A via the - `destructive` argument to the constructor. + A \leftarrow A + \alpha x y^{\top} + for matrix :math:`A`, scalar :math:`\alpha` and vectors :math:`x` and :math:`y`. + Constructed with ``inplace=True``, the output aliases ``A``'s storage and the op + destroys it; otherwise ``A`` is left untouched. """ - __props__ = ("destructive",) + __props__ = ("inplace",) + gufunc_signature = "(m,n),(),(m),(n)->(m,n)" - def __init__(self, destructive): - self.destructive = destructive - if destructive: + def __init__(self, inplace): + self.inplace = inplace + if inplace: self.destroy_map = {0: [0]} + def inplace_on_inputs(self, allowed_inplace_inputs: list[int]) -> Op: + """``Ger`` updates ``A`` in place, so that is the only input it can destroy.""" + if 0 in allowed_inplace_inputs: + return type(self)(inplace=True) + return self + def __str__(self): - if self.destructive: - return f"{self.__class__.__name__}{{destructive}}" + if self.inplace: + return f"{self.__class__.__name__}{{inplace}}" else: - return f"{self.__class__.__name__}{{non-destructive}}" + return f"{self.__class__.__name__}{{no_inplace}}" def make_node(self, A, alpha, x, y): A = as_tensor_variable(A) @@ -63,14 +71,10 @@ def perform(self, node, inputs, output_storage): ger_func = scipy_linalg.get_blas_funcs("ger", dtype=A.dtype) if A.flags["C_CONTIGUOUS"]: # Work on transposed system to avoid copying - A = ger_func(alpha, y, x, a=A.T, overwrite_a=self.destructive).T + A = ger_func(alpha, y, x, a=A.T, overwrite_a=self.inplace).T else: - A = ger_func(alpha, x, y, a=A, overwrite_a=self.destructive) + A = ger_func(alpha, x, y, a=A, overwrite_a=self.inplace) output_storage[0][0] = A def infer_shape(self, node, input_shapes): return [input_shapes[0]] - - -ger = Ger(destructive=False) -ger_destructive = Ger(destructive=True) diff --git a/pytensor/tensor/math.py b/pytensor/tensor/math.py index bf76ab97a1..2d0477f1bf 100644 --- a/pytensor/tensor/math.py +++ b/pytensor/tensor/math.py @@ -3046,9 +3046,9 @@ class Dot(Op): ----- Matrix-matrix products are sometimes optimized to Dot22 or Gemm ops (see tensor.blas). - Vector-vector products are sometimes optimized to Ger or CGer (see + Vector-vector products are sometimes optimized to Ger (see tensor.blas). - Matrix-vector products are sometimes optimized to Gemv, CGemv (see + Matrix-vector products are sometimes optimized to Gemv (see tensor.blas). """ @@ -3197,9 +3197,9 @@ def dense_dot(a, b): ----- Matrix-matrix products are sometimes optimized to Dot22 or Gemm ops (see tensor.blas). - Vector-vector products are sometimes optimized to Ger or CGer (see + Vector-vector products are sometimes optimized to Ger (see tensor.blas). - Matrix-vector products are sometimes optimized to Gemv, CGemv (see + Matrix-vector products are sometimes optimized to Gemv (see tensor.blas). """ diff --git a/pytensor/tensor/rewriting/__init__.py b/pytensor/tensor/rewriting/__init__.py index 095d3b74a3..fadd5b813a 100644 --- a/pytensor/tensor/rewriting/__init__.py +++ b/pytensor/tensor/rewriting/__init__.py @@ -1,7 +1,6 @@ import pytensor.tensor.rewriting.assumptions import pytensor.tensor.rewriting.basic import pytensor.tensor.rewriting.blas -import pytensor.tensor.rewriting.blas_c import pytensor.tensor.rewriting.blockwise import pytensor.tensor.rewriting.einsum import pytensor.tensor.rewriting.elemwise diff --git a/pytensor/tensor/rewriting/blas.py b/pytensor/tensor/rewriting/blas.py index 03a1e8b0ab..b9bdb84544 100644 --- a/pytensor/tensor/rewriting/blas.py +++ b/pytensor/tensor/rewriting/blas.py @@ -85,16 +85,14 @@ from pytensor.tensor import basic as ptb from pytensor.tensor.blas import ( Dot22, + Gemm, + Gemv, + Ger, _batched_dot, _dot22, _dot22scalar, - gemm_inplace, - gemm_no_inplace, - gemv_inplace, - gemv_no_inplace, - ger, - ger_destructive, ) +from pytensor.tensor.blockwise import Blockwise from pytensor.tensor.elemwise import DimShuffle, Elemwise from pytensor.tensor.exceptions import NotScalarConstantError from pytensor.tensor.math import ( @@ -106,9 +104,12 @@ sub, variadic_add, ) +from pytensor.tensor.rewriting.basic import elemwise_of +from pytensor.tensor.rewriting.blockwise import blockwise_of from pytensor.tensor.rewriting.elemwise import local_dimshuffle_lift from pytensor.tensor.type import ( TensorType, + float_dtypes, integer_dtypes, values_eq_approx_remove_inf_nan, ) @@ -177,7 +178,7 @@ def _beta_L_plus_alpha_M(fgraph, beta, L, alpha, M, recurse_flip=True): # if res_is_a(M, _dot22, 1): if M.owner and M.owner.op == _dot22: Ml, Mr = M.owner.inputs - rval = [gemm_no_inplace(L, alpha, Ml, Mr, beta)] + rval = [Blockwise(Gemm(inplace=False))(L, alpha, Ml, Mr, beta)] return rval, M # it also might be the case that there is a dimshuffle between the + @@ -192,19 +193,25 @@ def _beta_L_plus_alpha_M(fgraph, beta, L, alpha, M, recurse_flip=True): if M.owner.op.new_order == (0,): # it is making a column MM into a vector MMl, MMr = MM.owner.inputs - g = gemm_no_inplace(L.dimshuffle(0, "x"), alpha, MMl, MMr, beta) + g = Blockwise(Gemm(inplace=False))( + L.dimshuffle(0, "x"), alpha, MMl, MMr, beta + ) rval = [g.dimshuffle(0)] return rval, MM if M.owner.op.new_order == (1,): # it is making a row MM into a vector MMl, MMr = MM.owner.inputs - g = gemm_no_inplace(L.dimshuffle("x", 0), alpha, MMl, MMr, beta) + g = Blockwise(Gemm(inplace=False))( + L.dimshuffle("x", 0), alpha, MMl, MMr, beta + ) rval = [g.dimshuffle(1)] return rval, MM if len(M.owner.op.new_order) == 0: # it is making a row MM into a vector MMl, MMr = MM.owner.inputs - g = gemm_no_inplace(L.dimshuffle("x", "x"), alpha, MMl, MMr, beta) + g = Blockwise(Gemm(inplace=False))( + L.dimshuffle("x", "x"), alpha, MMl, MMr, beta + ) rval = [g.dimshuffle()] return rval, MM @@ -595,39 +602,21 @@ def local_dot_to_dot22(fgraph, node): _logger.info(f"Not optimizing dot with inputs {x} {y} {x.type} {y.type}") -@node_rewriter([gemm_no_inplace], inplace=True) -def local_inplace_gemm(fgraph, node): - if node.op == gemm_no_inplace: - new_out = [gemm_inplace(*node.inputs)] - copy_stack_trace(node.outputs, new_out) - return new_out - - -@node_rewriter([gemv_no_inplace], inplace=True) -def local_inplace_gemv(fgraph, node): - if node.op == gemv_no_inplace: - new_out = [gemv_inplace(*node.inputs)] - copy_stack_trace(node.outputs, new_out) - return new_out - - -@node_rewriter([ger], inplace=True) -def local_inplace_ger(fgraph, node): - if node.op == ger: - new_out = [ger_destructive(*node.inputs)] - copy_stack_trace(node.outputs, new_out) - return new_out - - -@node_rewriter([gemm_no_inplace]) +@node_rewriter([Gemm, blockwise_of(Gemm)]) def local_gemm_to_gemv(fgraph, node): """GEMM acting on row or column matrices -> GEMV.""" + # These trackers match any `Gemm`; an inplace one must not be traded for a `Gemv` + # or `Ger` that would no longer destroy `z`. + core_op = node.op.core_op if isinstance(node.op, Blockwise) else node.op + if core_op.inplace: + return None + z, a, x, y, b = node.inputs if z.broadcastable == x.broadcastable == (True, False): - r = gemv_no_inplace(z.dimshuffle(1), a, y.T, x.dimshuffle(1), b) + r = Blockwise(Gemv(inplace=False))(z.dimshuffle(1), a, y.T, x.dimshuffle(1), b) new_out = [r.dimshuffle("x", 0)] elif z.broadcastable == y.broadcastable == (False, True): - r = gemv_no_inplace(z.dimshuffle(0), a, x, y.dimshuffle(0), b) + r = Blockwise(Gemv(inplace=False))(z.dimshuffle(0), a, x, y.dimshuffle(0), b) new_out = [r.dimshuffle(0, "x")] else: return @@ -635,9 +624,15 @@ def local_gemm_to_gemv(fgraph, node): return new_out -@node_rewriter([gemm_no_inplace]) +@node_rewriter([Gemm, blockwise_of(Gemm)]) def local_gemm_to_ger(fgraph, node): """GEMM computing an outer-product -> GER.""" + # These trackers match any `Gemm`; an inplace one must not be traded for a `Gemv` + # or `Ger` that would no longer destroy `z`. + core_op = node.op.core_op if isinstance(node.op, Blockwise) else node.op + if core_op.inplace: + return None + z, a, x, y, b = node.inputs if x.broadcastable[1] and y.broadcastable[0]: # x and y are both vectors so this might qualifies for a GER @@ -650,11 +645,11 @@ def local_gemm_to_ger(fgraph, node): return if bval == 1: # best case a natural GER - rval = ger(z, a, xv, yv) + rval = Blockwise(Ger(inplace=False))(z, a, xv, yv) new_out = [rval] elif bval == 0: # GER on zeros_like should be faster than GEMM zeros = ptb.zeros([x.shape[0], y.shape[1]], x.dtype) - rval = ger(zeros, a, xv, yv) + rval = Blockwise(Ger(inplace=False))(zeros, a, xv, yv) new_out = [rval] else: # if bval is another constant, then z is being usefully @@ -678,26 +673,26 @@ def local_dot22_to_ger_or_gemv(fgraph, node): xv = x.dimshuffle(0) yv = y.dimshuffle(1) zeros = ptb.zeros([x.shape[0], y.shape[1]], dtype=x.dtype) - rval = ger(zeros, one, xv, yv) + rval = Blockwise(Ger(inplace=False))(zeros, one, xv, yv) new_out = [rval] elif xb[0] and yb[1]: # x and y are both vectors so this qualifies for a sdot / ddot - # PyTensor's CGemv will call sdot/ddot at runtime, the Scipy Gemv may not + # PyTensor's C Gemv will call sdot/ddot at runtime, the Scipy Gemv may not xv = x.dimshuffle(1) zeros = ptb.AllocEmpty(x.dtype)(1) - rval = gemv_no_inplace(zeros, one, y.T, xv, zero) + rval = Blockwise(Gemv(inplace=False))(zeros, one, y.T, xv, zero) new_out = [rval.dimshuffle("x", 0)] elif xb[0] and not yb[0] and not yb[1]: # x is vector, y is matrix so try gemv xv = x.dimshuffle(1) zeros = ptb.AllocEmpty(x.dtype)(y.shape[1]) - rval = gemv_no_inplace(zeros, one, y.T, xv, zero) + rval = Blockwise(Gemv(inplace=False))(zeros, one, y.T, xv, zero) new_out = [rval.dimshuffle("x", 0)] elif not xb[0] and not xb[1] and yb[1]: # x is matrix, y is vector, try gemv yv = y.dimshuffle(0) zeros = ptb.AllocEmpty(x.dtype)(x.shape[0]) - rval = gemv_no_inplace(zeros, one, x, yv, zero) + rval = Blockwise(Gemv(inplace=False))(zeros, one, x, yv, zero) new_out = [rval.dimshuffle(0, "x")] else: return @@ -744,20 +739,6 @@ def local_dot22_to_ger_or_gemv(fgraph, node): ) -blas_opt_inplace = dfs_rewriter( - local_inplace_gemm, local_inplace_gemv, local_inplace_ger, name="blas_opt_inplace" -) -optdb.register( - "InplaceBlasOpt", - blas_opt_inplace, - "fast_run", - "inplace", - "blas_opt_inplace", - # Before we try to make elemwise things inplace (70.5) - position=50.2, -) - - @node_rewriter([mul]) def local_dot22_to_dot22scalar(fgraph, node): """ @@ -923,3 +904,154 @@ def specialize_matmul_to_batched_dot(fgraph, node): copy_stack_trace(node.outputs, [new_out]) return [new_out] + + +def _split_scalar_factor(var, dtype): + """ + Split ``alpha * X`` into the scalar and the rest. + + Parameters + ---------- + var : TensorVariable + The term to inspect. + dtype : str + The dtype the scalar has to be usable as. + + Returns + ------- + alpha : TensorVariable or None + The scalar factor, or None when there is none to peel. + rest : TensorVariable + ``var`` with the scalar removed, or ``var`` itself when ``alpha`` is None. + """ + node = var.owner + if node is None or not isinstance(node.op, Elemwise): + return None, var + if not isinstance(node.op.scalar_op, pytensor.scalar.Mul): + return None, var + + scalars, others = [], [] + for term in node.inputs: + scalar = _as_scalar(term, dtype=dtype) + if scalar is None: + others.append(term) + else: + scalars.append(scalar) + + # Anything but a single non-scalar term is a product this rewrite cannot read. + if not scalars or len(others) != 1: + return None, var + + alpha = scalars[0] if len(scalars) == 1 else mul(*scalars) + return alpha, others[0] + + +def _as_matrix_product(fgraph, var): + """ + Return the two matrices of a rank-2 matrix product, or None if ``var`` is not one. + + A product read by more than one client is rejected: folding it into a `Gemm` + would compute it a second time for whoever else consumes it. + """ + node = var.owner + if node is None: + return None + + op = node.op + core_op = op.core_op if isinstance(op, Blockwise) else op + if not isinstance(core_op, Dot): + return None + if len(fgraph.clients[var]) > 1: + return None + + x, y = node.inputs + if x.type.ndim != 2 or y.type.ndim != 2: + return None + return x, y + + +@register_specialize +@node_rewriter([elemwise_of(pytensor.scalar.Add)]) +def local_add_dot_to_gemm(fgraph, node): + r""" + Rewrite :math:`\beta C + \alpha A B` as a single `Gemm`. + + Either scalar may be absent, so this also covers :math:`C + AB`, + :math:`C + \alpha AB` and :math:`\beta C + AB`. The matrix product is matched + through `Blockwise` as well as bare, because ``@`` builds a `Blockwise` of `Dot` + that is only unwrapped much later. A matrix-vector product is matched too, and is + handed to `Gemm` as a column or row matrix so that `local_gemm_to_gemv` can turn it + into a `Gemv`. + + An accumulator that only broadcasts against the product is left alone: BLAS writes + into the output buffer itself, so folding one in would cost the pass it saves. + """ + if len(node.inputs) != 2: + return None + + [out] = node.outputs + dtype = out.type.dtype + if dtype not in float_dtypes or out.type.ndim not in (1, 2): + return None + + for index, term in enumerate(node.inputs): + alpha, product = _split_scalar_factor(term, dtype) + + # `Gemm` computes into a buffer shaped like the product, so a product that only + # broadcasts against the sum would hand the replacement a narrower type than the + # node it replaces. + if product.type.broadcastable != out.type.broadcastable: + continue + + # `matmul` and `dot` both promote a vector operand to a matrix and squeeze the + # rank-2 product back down, so a rank-1 sum reaches this rewrite with the squeeze + # still on it. Take it off, build the `Gemm` at rank 2, and squeeze the result. + squeeze_order = None + if product.owner is not None and isinstance(product.owner.op, DimShuffle): + if out.type.ndim != 1 or len(fgraph.clients[product]) > 1: + continue + squeeze_order = product.owner.op.new_order + if squeeze_order not in ((0,), (1,)): + continue + [product] = product.owner.inputs + matrices = _as_matrix_product(fgraph, product) + if matrices is None: + continue + x, y = matrices + + beta, z = _split_scalar_factor(node.inputs[1 - index], dtype) + if z.type.ndim != out.type.ndim: + continue + # BLAS accumulates into the output buffer itself, so an accumulator that only + # broadcasts against the product would have to be materialized to the product's + # shape first -- exactly the pass this rewrite exists to remove. + if z.type.broadcastable != out.type.broadcastable: + continue + + # Gemm reads one dtype across all five inputs and will not upcast for us. + one = ptb.constant(np.asarray(1.0, dtype=dtype)) + alpha = one if alpha is None else alpha + beta = one if beta is None else beta + if not all(term.type.dtype == dtype for term in (z, x, y, alpha, beta)): + continue + + if squeeze_order == (0,): + z = z.dimshuffle(0, "x") + elif squeeze_order == (1,): + z = z.dimshuffle("x", 0) + + # Emitted through `Blockwise` so that `InplaceBlockwiseOptimizer` can ask the op itself + # whether it may destroy `z`; the wrapper is unwrapped again once inplace has been decided. + gemm = Blockwise(Gemm(inplace=False))(z, alpha, x, y, beta) + new_out = [gemm if squeeze_order is None else gemm.dimshuffle(*squeeze_order)] + copy_stack_trace(node.outputs, new_out) + return new_out + + return None + + +# Both also run inside `blas_optdb`, which sits at optdb position 1.7, ahead of +# `specialize` at 2.0. The `Gemm` above is created after that has already run, so they +# are registered here as well to catch the rank-1 and row/column-matrix products. +register_specialize(local_gemm_to_ger, name="local_gemm_to_ger_after_add_dot") +register_specialize(local_gemm_to_gemv, name="local_gemm_to_gemv_after_add_dot") diff --git a/pytensor/tensor/rewriting/blas_c.py b/pytensor/tensor/rewriting/blas_c.py deleted file mode 100644 index eaf4dd89e4..0000000000 --- a/pytensor/tensor/rewriting/blas_c.py +++ /dev/null @@ -1,76 +0,0 @@ -from pytensor.configdefaults import config -from pytensor.graph.rewriting.basic import dfs_rewriter -from pytensor.tensor import basic as ptb -from pytensor.tensor.blas import gemv_inplace, gemv_no_inplace, ger, ger_destructive -from pytensor.tensor.blas.blas_c import ( - CGemv, - CGer, - cgemv_inplace, - cgemv_no_inplace, - cger_inplace, -) -from pytensor.tensor.rewriting.blas import blas_optdb, node_rewriter, optdb - - -@node_rewriter([ger, ger_destructive]) -def use_c_ger(fgraph, node): - if not config.blas__ldflags: - return - # Only float32 and float64 are supported for now. - if node.op == ger and node.outputs[0].dtype in ("float32", "float64"): - return [CGer(False)(*node.inputs)] - if node.op == ger_destructive and node.outputs[0].dtype in ("float32", "float64"): - return [CGer(True)(*node.inputs)] - - -@node_rewriter([CGer(False)]) -def make_c_ger_destructive(fgraph, node): - return [cger_inplace(*node.inputs)] - - -@node_rewriter([gemv_inplace, gemv_no_inplace]) -def use_c_gemv(fgraph, node): - if not config.blas__ldflags: - return - # Only float32 and float64 are supported for now. - if node.op == gemv_no_inplace and node.outputs[0].dtype in ("float32", "float64"): - return [cgemv_no_inplace(*node.inputs)] - if node.op == gemv_inplace and node.outputs[0].dtype in ("float32", "float64"): - return [cgemv_inplace(*node.inputs)] - - -@node_rewriter([CGemv(inplace=False)]) -def make_c_gemv_destructive(fgraph, node): - inputs = list(node.inputs) - dest = inputs[0] - if ( - dest.owner - and isinstance(dest.owner.op, ptb.AllocEmpty) - and len(fgraph.clients[dest]) > 1 - ): - inputs[0] = ptb.AllocEmpty(dest.dtype)(*dest.owner.inputs) - - return [cgemv_inplace(*inputs)] - - -blas_optdb.register( - "use_c_blas", - dfs_rewriter(use_c_ger, use_c_gemv), - "fast_run", - "c_blas", - "cxx_only", - position=20, -) - -# this matches the InplaceBlasOpt defined in blas.py -optdb.register( - "c_blas_destructive", - dfs_rewriter( - make_c_ger_destructive, make_c_gemv_destructive, name="c_blas_destructive" - ), - "fast_run", - "inplace", - "c_blas", - "cxx_only", - position=70.0, -) diff --git a/pytensor/tensor/rewriting/blockwise.py b/pytensor/tensor/rewriting/blockwise.py index 45daae925b..d0fa684ca0 100644 --- a/pytensor/tensor/rewriting/blockwise.py +++ b/pytensor/tensor/rewriting/blockwise.py @@ -1,6 +1,7 @@ from pytensor.compile.mode import optdb from pytensor.graph import Constant, Op, node_rewriter from pytensor.graph.destroyhandler import inplace_candidates +from pytensor.graph.fg import Output from pytensor.graph.replace import vectorize_graph from pytensor.graph.rewriting.basic import copy_stack_trace, dfs_rewriter from pytensor.graph.rewriting.unify import OpPattern, OpPatternOpTypeType @@ -344,6 +345,61 @@ def create_inplace_node(self, node, inplace_pattern): return inplace_blockwise_op.make_node(*node.inputs) +def _accepts_inplace_on(node, input_index): + """Return whether ``node``'s op would destroy the input at ``input_index`` if allowed.""" + op = node.op + core_op = op.core_op if isinstance(op, Blockwise) else op + inplace_op = core_op.inplace_on_inputs(allowed_inplace_inputs=[input_index]) + return any( + input_index in destroyed for destroyed in inplace_op.destroy_map.values() + ) + + +@node_rewriter([AllocEmpty]) +def local_split_alloc_empty_clients(fgraph, node): + """Give each client that wants to write into an `AllocEmpty` a buffer of its own. + + `AllocEmpty` produces uninitialized memory, so no client can depend on what another + left in it. Sharing one buffer only stops every client but the first from claiming it + for an inplace operation. + """ + [out] = node.outputs + clients = fgraph.clients[out] + if len(clients) < 2: + return None + + # The first client keeps the original buffer; the rest get their own. Kept in client + # order so the rewrite emits the same graph on every run. + destroyers = dict.fromkeys( + client + for client, input_index in clients[1:] + if not isinstance(client.op, Output) + and _accepts_inplace_on(client, input_index) + ) + if not destroyers: + return None + + replacements = {} + for client in destroyers: + new_inputs = [ + node.op(*node.inputs) if inp is out else inp for inp in client.inputs + ] + new_client = client.clone_with_new_inputs(new_inputs) + copy_stack_trace(client.outputs, new_client.outputs) + replacements.update(zip(client.outputs, new_client.outputs, strict=True)) + return replacements + + +optdb.register( + "local_split_alloc_empty_clients", + dfs_rewriter(local_split_alloc_empty_clients), + "fast_run", + "inplace", + # After the last merge pass, before any inplace rewrite claims the shared buffer. + position=49.4, +) + + optdb.register( "blockwise_inplace", InplaceBlockwiseOptimizer(), diff --git a/tests/benchmarks/test_blas.py b/tests/benchmarks/test_blas.py index 17ae5573ce..4d1726cbb7 100644 --- a/tests/benchmarks/test_blas.py +++ b/tests/benchmarks/test_blas.py @@ -3,7 +3,7 @@ from pytensor import In, function from pytensor.tensor import dot, empty, matrix, outer, scalar, tensor, vector -from pytensor.tensor.blas.blas_c import CGemv +from pytensor.tensor.blas import Gemv @pytest.mark.parametrize("dtype", ("float64", "float32", "mixed")) @@ -62,7 +62,7 @@ def test_cgemv_vector_dot_benchmark(benchmark): a = vector("A", shape=(n,)) b = vector("x", shape=(n,)) - out = CGemv(inplace=True)( + out = Gemv(inplace=True)( empty((1,)), 1.0, a[None], @@ -97,7 +97,7 @@ def test_cgemv_negative_strides_benchmark( x = vector("x", shape=(A.type.shape[-1],)) y = vector("y", shape=(A.type.shape[0],)) - out = CGemv(inplace=False)( + out = Gemv(inplace=False)( y, 1.0, A, diff --git a/tests/link/c/test_dispatch.py b/tests/link/c/test_dispatch.py new file mode 100644 index 0000000000..3afbb5b5bb --- /dev/null +++ b/tests/link/c/test_dispatch.py @@ -0,0 +1,266 @@ +import numpy as np +import pytest + +import pytensor +import pytensor.scalar as ps +import pytensor.tensor as pt +from pytensor.compile.debug.debugmode import BadThunkOutput, DebugMode +from pytensor.compile.mode import Mode +from pytensor.configdefaults import config +from pytensor.graph.basic import Apply +from pytensor.graph.fg import FunctionGraph +from pytensor.graph.op import Op +from pytensor.graph.utils import MethodNotDefined +from pytensor.link.c.basic import CLinker +from pytensor.link.c.dispatch.basic import ( + CImpl, + c_funcify, + c_thunk_from_dispatch, +) +from pytensor.link.vm import VMLinker +from pytensor.tensor.shape import Shape, Shape_i + + +pytestmark = pytest.mark.skipif( + not config.cxx, reason="A C compiler is required to test the C dispatch" +) + +CVM_MODE = Mode(linker="cvm", optimizer=None) +PY_MODE = Mode(linker="py", optimizer=None) + + +class ScalarOpBase(Op): + """A pure scalar op: only `make_node` and `perform`.""" + + __props__ = () + increment = 1.0 + + def make_node(self, x): + x = ps.as_scalar(x) + return Apply(self, [x], [x.type()]) + + def perform(self, node, inputs, output_storage): + (x,) = inputs + output_storage[0][0] = np.dtype(node.outputs[0].dtype).type(x + self.increment) + + +class IncOne(ScalarOpBase): + pass + + +class IncOneNoImpl(ScalarOpBase): + pass + + +class IncOneDeclining(ScalarOpBase): + pass + + +class IncOneImpl(CImpl): + def c_code(self, node, name, inputs, outputs, sub): + (x,) = inputs + (z,) = outputs + return f"{z} = {x} + 1;" + + def c_code_cache_version(self): + return (1,) + + +class DecliningImpl(CImpl): + def c_code(self, node, name, inputs, outputs, sub): + raise MethodNotDefined("c_code") + + def c_code_cache_version(self): + return (1,) + + +@c_funcify.register(IncOne) +def c_funcify_inc_one(op, node=None, **kwargs): + return IncOneImpl(op) + + +@c_funcify.register(IncOneDeclining) +def c_funcify_declining(op, node=None, **kwargs): + return DecliningImpl(op) + + +def make_thunk_for(op, x_value=2.0, dtype="float64"): + x = ps.ScalarType(dtype)("x") + out = op(x) + node = out.owner + storage_map = {x: [np.dtype(dtype).type(x_value)], out: [None]} + compute_map = {x: [True], out: [False]} + thunk = c_thunk_from_dispatch(node, storage_map, compute_map, []) + return thunk, storage_map, compute_map, out + + +def test_pure_op_gains_c_thunk(): + thunk, storage_map, compute_map, out = make_thunk_for(IncOne()) + + assert hasattr(thunk, "cthunk") + assert thunk.lazy is False + assert thunk.inputs == [storage_map[out.owner.inputs[0]]] + assert thunk.outputs == [storage_map[out]] + + thunk() + assert storage_map[out][0] == 3.0 + assert compute_map[out][0] is True + + +def test_pure_op_cvm_function_matches_perform(): + x = ps.float64("x") + out = IncOne()(x) + + f_c = pytensor.function([x], out, mode=CVM_MODE) + f_py = pytensor.function([x], out, mode=PY_MODE) + assert f_c(2.0) == f_py(2.0) == 3.0 + + +def test_unregistered_pure_op_falls_back(): + op = IncOneNoImpl() + with pytest.raises(NotImplementedError, match="No C implementation registered"): + c_funcify(op) + + x = ps.float64("x") + f = pytensor.function([x], op(x), mode=CVM_MODE) + assert f(2.0) == 3.0 + + +def test_declining_impl_falls_back(): + op = IncOneDeclining() + with pytest.raises(MethodNotDefined): + make_thunk_for(op) + + x = ps.float64("x") + f = pytensor.function([x], op(x), mode=CVM_MODE) + assert f(2.0) == 3.0 + + +def test_cop_is_its_own_impl(): + op = Shape() + assert c_funcify(op) is op + + +def test_float16_guard_falls_back(): + # The guard raises NotImplementedError either way; the warning is only + # emitted when the impl's C code builds for f16, but ScalarType's own C + # support rejects f16 first. + op = IncOne() + with pytest.raises(NotImplementedError): + make_thunk_for(op, dtype="float16") + + x = ps.ScalarType("float16")("x") + f = pytensor.function([x], op(x), mode=CVM_MODE) + assert f(np.float16(2.0)) == np.float16(3.0) + + +def test_vm_without_c_thunks_skips_dispatch(monkeypatch): + def fail_dispatch(*args, **kwargs): + raise AssertionError("dispatch should not run when c_thunks=False") + + monkeypatch.setattr( + "pytensor.link.c.dispatch.basic.c_thunk_from_dispatch", fail_dispatch + ) + + x = ps.float64("x") + mode = Mode(linker=VMLinker(use_cloop=False, c_thunks=False), optimizer=None) + f = pytensor.function([x], IncOne()(x), mode=mode) + assert f(2.0) == 3.0 + + +@pytest.mark.parametrize("linker", ["c", "c|py"]) +def test_whole_graph_linkers_use_dispatch(linker): + x = ps.float64("x") + f = pytensor.function([x], IncOne()(x), mode=Mode(linker=linker, optimizer=None)) + assert f(2.0) == 3.0 + + +def test_whole_graph_c_linker_unregistered_raises(): + x = ps.float64("x") + with pytest.raises(NotImplementedError, match="cannot produce C code"): + pytensor.function([x], IncOneNoImpl()(x), mode=Mode(linker="c", optimizer=None)) + + +def test_cmodule_key_stable_and_versioned(): + def key_for_fresh_graph(): + x = ps.float64("x") + out = IncOne()(x) + fgraph = FunctionGraph([x], [out]) + return CLinker().accept(fgraph).cmodule_key() + + key_a = key_for_fresh_graph() + key_b = key_for_fresh_graph() + assert key_a == key_b + + version, _sig = key_a + # The registered impl's cache version makes the module versioned (cacheable + # across processes), even though the graph op itself has no C methods. + assert version != () + assert IncOneImpl(IncOne()).c_code_cache_version() in version + + +def test_params_constants_deduplicated_across_nodes(): + x = pt.matrix("x") + y = pt.matrix("y") + out = Shape_i(0)(x) + Shape_i(0)(y) + + fgraph = FunctionGraph([x, y], [out]) + cl = CLinker().accept(fgraph) + shape_i_nodes = [n for n in cl.node_order if isinstance(n.op, Shape_i)] + assert len(shape_i_nodes) == 2 + # Both Shape_i(0) nodes share one params Constant. + assert len(cl.node_params) == 1 + + f = pytensor.function([x, y], out, mode=Mode(linker="c", optimizer=None)) + assert f(np.ones((3, 2)), np.ones((5, 2))) == 8 + + +class IncOneWrongImpl(ScalarOpBase): + pass + + +class WrongImpl(CImpl): + def c_code(self, node, name, inputs, outputs, sub): + (x,) = inputs + (z,) = outputs + return f"{z} = {x} + 2;" # disagrees with perform on purpose + + def c_code_cache_version(self): + return (1,) + + +@c_funcify.register(IncOneWrongImpl) +def c_funcify_wrong(op, node=None, **kwargs): + return WrongImpl(op) + + +def test_cop_graph_resolves_to_identity(): + # The parity guarantee: every COp node resolves to itself, so CLinker calls + # the op's own c_code/cache-version methods and produces byte-identical + # source and cache keys. + x = pt.matrix("x") + out = (x.T + 1.0).sum(axis=0) + fgraph = FunctionGraph([x], [out]) + cl = CLinker().accept(fgraph) + + for node in cl.node_order: + assert cl._impl_for(node) is node.op + + # Source generation works and the module is versioned (cacheable). + assert isinstance(cl.get_src_code(), str) + version, _sig = cl.cmodule_key() + assert version != () + + +def test_debugmode_cross_checks_dispatch_impl(): + x = ps.float64("x") + f = pytensor.function( + [x], IncOne()(x), mode=DebugMode(optimizer=None, check_py_code=True) + ) + assert f(2.0) == 3.0 + + f_wrong = pytensor.function( + [x], IncOneWrongImpl()(x), mode=DebugMode(optimizer=None, check_py_code=True) + ) + with pytest.raises(BadThunkOutput): + f_wrong(2.0) diff --git a/tests/link/jax/test_blas.py b/tests/link/jax/test_blas.py index a738b6701e..1fe6424d1a 100644 --- a/tests/link/jax/test_blas.py +++ b/tests/link/jax/test_blas.py @@ -1,6 +1,7 @@ import numpy as np import pytest +from pytensor import tensor as pt from pytensor.compile.maker import function from pytensor.compile.mode import Mode from pytensor.configdefaults import config @@ -31,3 +32,92 @@ def test_jax_BatchedDot(): pytensor_jax_fn = function([a, b], [out], mode=jax_mode) with pytest.raises(TypeError): pytensor_jax_fn(*inputs) + + +@pytest.mark.parametrize( + "alpha, beta", + [(None, None), (2.0, None), (None, 3.0), (2.0, 3.0)], + ids=["plain", "alpha", "beta", "alpha_beta"], +) +def test_jax_Gemm(alpha, beta): + # Gemm is what local_add_dot_to_gemm folds `beta * C + alpha * (A @ B)` into, so the + # scales arrive as constants. JAX has no fused kernel to reach; XLA folds them + # into the dot itself. + z = pt.matrix("z", dtype=config.floatX) + x = pt.matrix("x", dtype=config.floatX) + y = pt.matrix("y", dtype=config.floatX) + one = np.asarray(1.0, dtype=config.floatX) + + out = pt_blas.Gemm(inplace=False)( + z, + one if alpha is None else np.asarray(alpha, dtype=config.floatX), + x, + y, + one if beta is None else np.asarray(beta, dtype=config.floatX), + ) + + rng = np.random.default_rng(418) + test_values = [ + rng.normal(size=shape).astype(config.floatX) + for shape in ((4, 6), (4, 5), (5, 6)) + ] + compare_jax_and_py([z, x, y], [out], test_values) + + +def test_jax_Gemm_runtime_scales(): + z = pt.matrix("z", dtype=config.floatX) + x = pt.matrix("x", dtype=config.floatX) + y = pt.matrix("y", dtype=config.floatX) + alpha = pt.scalar("alpha", dtype=config.floatX) + beta = pt.scalar("beta", dtype=config.floatX) + + out = pt_blas.Gemm(inplace=False)(z, alpha, x, y, beta) + + rng = np.random.default_rng(418) + test_values = [ + rng.normal(size=shape).astype(config.floatX) + for shape in ((4, 6), (4, 5), (5, 6)) + ] + compare_jax_and_py( + [z, alpha, x, y, beta], + [out], + [ + test_values[0], + np.asarray(2.0, dtype=config.floatX), + test_values[1], + test_values[2], + np.asarray(3.0, dtype=config.floatX), + ], + ) + + +def test_jax_Ger(): + # local_gemm_to_ger folds a rank-1 `A + alpha * outer(x, y)` into Ger. + A = pt.matrix("A", dtype=config.floatX) + x = pt.vector("x", dtype=config.floatX) + y = pt.vector("y", dtype=config.floatX) + out = pt_blas.Ger(inplace=False)(A, np.asarray(2.0, dtype=config.floatX), x, y) + + rng = np.random.default_rng(418) + test_values = [ + rng.normal(size=shape).astype(config.floatX) for shape in ((4, 6), (4,), (6,)) + ] + compare_jax_and_py([A, x, y], [out], test_values) + + +def test_jax_Ger_runtime_alpha(): + A = pt.matrix("A", dtype=config.floatX) + x = pt.vector("x", dtype=config.floatX) + y = pt.vector("y", dtype=config.floatX) + alpha = pt.scalar("alpha", dtype=config.floatX) + out = pt_blas.Ger(inplace=False)(A, alpha, x, y) + + rng = np.random.default_rng(418) + A_val, x_val, y_val = ( + rng.normal(size=shape).astype(config.floatX) for shape in ((4, 6), (4,), (6,)) + ) + compare_jax_and_py( + [A, alpha, x, y], + [out], + [A_val, np.asarray(2.0, dtype=config.floatX), x_val, y_val], + ) diff --git a/tests/link/mlx/test_blas.py b/tests/link/mlx/test_blas.py index 240ec3fedf..0df5f7e804 100644 --- a/tests/link/mlx/test_blas.py +++ b/tests/link/mlx/test_blas.py @@ -19,7 +19,7 @@ def test_mlx_Gemv_static_scales(): A = pt.matrix("A", dtype=config.floatX) x = pt.vector("x", dtype=config.floatX) - out = pt_blas.gemv_no_inplace( + out = pt_blas.Gemv(inplace=False)( y, np.asarray(0.5, dtype=config.floatX), A, @@ -46,7 +46,7 @@ def test_mlx_Gemv_symbolic_scales(): alpha = pt.scalar("alpha", dtype=config.floatX) beta = pt.scalar("beta", dtype=config.floatX) - out = pt_blas.gemv_no_inplace(y, alpha, A, x, beta) + out = pt_blas.Gemv(inplace=False)(y, alpha, A, x, beta) rng = np.random.default_rng(sum(map(ord, "test_mlx_Gemv_symbolic_scales"))) y_test = rng.normal(size=(3,)).astype(config.floatX) @@ -67,7 +67,7 @@ def test_mlx_Ger_static_scale(): x = pt.vector("x", dtype=config.floatX) y = pt.vector("y", dtype=config.floatX) - out = pt_blas.ger(A, np.asarray(0.5, dtype=config.floatX), x, y) + out = pt_blas.Ger(inplace=False)(A, np.asarray(0.5, dtype=config.floatX), x, y) rng = np.random.default_rng(sum(map(ord, "test_mlx_Ger_static_scale"))) A_test = rng.normal(size=(3, 2)).astype(config.floatX) @@ -87,7 +87,7 @@ def test_mlx_Ger_symbolic_scale(): y = pt.vector("y", dtype=config.floatX) alpha = pt.scalar("alpha", dtype=config.floatX) - out = pt_blas.ger(A, alpha, x, y) + out = pt_blas.Ger(inplace=False)(A, alpha, x, y) rng = np.random.default_rng(sum(map(ord, "test_mlx_Ger_symbolic_scale"))) A_test = rng.normal(size=(3, 2)).astype(config.floatX) @@ -124,3 +124,93 @@ def test_mlx_BatchedDot(): inputs = [a_test_value[:-1], b_test_value] with pytest.raises(TypeError): pytensor_mlx_fn(*inputs) + + +@pytest.mark.parametrize( + "alpha, beta", + [(None, None), (2.0, None), (None, 3.0), (2.0, 3.0)], + ids=["plain", "alpha", "beta", "alpha_beta"], +) +def test_mlx_Gemm(alpha, beta): + # Gemm is what local_add_dot_to_gemm folds `beta * C + alpha * (A @ B)` into, so the + # scales arrive as constants. They reach MLX's fused kernel only in that form. + z = pt.matrix("z", dtype=config.floatX) + x = pt.matrix("x", dtype=config.floatX) + y = pt.matrix("y", dtype=config.floatX) + one = np.asarray(1.0, dtype=config.floatX) + + out = pt_blas.Gemm(inplace=False)( + z, + one if alpha is None else np.asarray(alpha, dtype=config.floatX), + x, + y, + one if beta is None else np.asarray(beta, dtype=config.floatX), + ) + + rng = np.random.default_rng(418) + test_values = [ + rng.normal(size=shape).astype(config.floatX) + for shape in ((4, 6), (4, 5), (5, 6)) + ] + compare_mlx_and_py([z, x, y], [out], test_values) + + +def test_mlx_Gemm_runtime_scales(): + # With alpha and beta as graph variables there is no constant to hand the fused + # kernel, so this covers the arithmetic path the dispatch falls back to. + z = pt.matrix("z", dtype=config.floatX) + x = pt.matrix("x", dtype=config.floatX) + y = pt.matrix("y", dtype=config.floatX) + alpha = pt.scalar("alpha", dtype=config.floatX) + beta = pt.scalar("beta", dtype=config.floatX) + + out = pt_blas.Gemm(inplace=False)(z, alpha, x, y, beta) + + rng = np.random.default_rng(418) + test_values = [ + rng.normal(size=shape).astype(config.floatX) + for shape in ((4, 6), (4, 5), (5, 6)) + ] + compare_mlx_and_py( + [z, alpha, x, y, beta], + [out], + [ + test_values[0], + np.asarray(2.0, dtype=config.floatX), + test_values[1], + test_values[2], + np.asarray(3.0, dtype=config.floatX), + ], + ) + + +def test_mlx_Ger(): + # local_gemm_to_ger folds a rank-1 `A + alpha * outer(x, y)` into Ger. + A = pt.matrix("A", dtype=config.floatX) + x = pt.vector("x", dtype=config.floatX) + y = pt.vector("y", dtype=config.floatX) + out = pt_blas.Ger(inplace=False)(A, np.asarray(2.0, dtype=config.floatX), x, y) + + rng = np.random.default_rng(418) + test_values = [ + rng.normal(size=shape).astype(config.floatX) for shape in ((4, 6), (4,), (6,)) + ] + compare_mlx_and_py([A, x, y], [out], test_values) + + +def test_mlx_Ger_runtime_alpha(): + A = pt.matrix("A", dtype=config.floatX) + x = pt.vector("x", dtype=config.floatX) + y = pt.vector("y", dtype=config.floatX) + alpha = pt.scalar("alpha", dtype=config.floatX) + out = pt_blas.Ger(inplace=False)(A, alpha, x, y) + + rng = np.random.default_rng(418) + A_val, x_val, y_val = ( + rng.normal(size=shape).astype(config.floatX) for shape in ((4, 6), (4,), (6,)) + ) + compare_mlx_and_py( + [A, alpha, x, y], + [out], + [A_val, np.asarray(2.0, dtype=config.floatX), x_val, y_val], + ) diff --git a/tests/link/numba/linalg/test_products.py b/tests/link/numba/linalg/test_products.py index 3e09b96cea..6b687175f2 100644 --- a/tests/link/numba/linalg/test_products.py +++ b/tests/link/numba/linalg/test_products.py @@ -3,6 +3,8 @@ import pytensor.tensor as pt from pytensor import In, config +from pytensor.link.numba.dispatch.basic import numba_njit +from pytensor.link.numba.dispatch.linalg.products import _gemm from pytensor.tensor.linalg.products import Expm, expm from tests.link.numba.test_basic import compare_numba_and_py, numba_inplace_mode @@ -86,3 +88,25 @@ def test_expm_integer_input(self): _, res = compare_numba_and_py([A], [y], [val]) np.testing.assert_array_equal(val, original) assert res[0].dtype == np.float64 + + +@numba_njit(final_function=True) +def _gemm_jit(A, B, C, transa, transb, alpha, beta): + return _gemm(A, B, C, transa, transb, alpha, beta) + + +@pytest.mark.parametrize("poison", [np.inf, np.nan], ids=["inf", "nan"]) +def test_gemm_ignores_C_when_beta_is_zero(poison): + """`numba_funcify_Dot` allocates C with `np.empty` and passes beta=0, so C holds whatever the + allocator returned. BLAS does not read C in that case and neither may the reference, or stray + non-finite bytes would multiply by zero into nan.""" + A = np.ascontiguousarray(rng.normal(size=(3, 2))) + B = np.ascontiguousarray(rng.normal(size=(2, 4))) + uninitialized = np.full((3, 4), poison) + + np.testing.assert_allclose( + _gemm(A, B, uninitialized.copy(), False, False, 1.0, 0.0), A @ B + ) + np.testing.assert_allclose( + _gemm_jit(A, B, uninitialized.copy(), False, False, 1.0, 0.0), A @ B + ) diff --git a/tests/link/numba/test_blas.py b/tests/link/numba/test_blas.py new file mode 100644 index 0000000000..0eded7ef50 --- /dev/null +++ b/tests/link/numba/test_blas.py @@ -0,0 +1,194 @@ +import numpy as np +import pytest + +import pytensor +import pytensor.tensor as pt +from pytensor import config +from pytensor.tensor.blas import Gemm, Gemv, Ger + + +pytestmark = pytest.mark.filterwarnings("error") + +pytest.importorskip("numba") + +floatX = config.floatX + + +@pytest.mark.parametrize( + "z_shape", [(6, 5), (1, 5), (6, 1), (1, 1)], ids=["full", "row", "column", "scalar"] +) +def test_gemm_broadcasts_its_accumulator(z_shape): + """``Gemm`` takes a ``z`` that is only broadcast against the product -- a bias row is the common + case -- so the buffer it accumulates into is the product's shape, not ``z``'s.""" + z = pt.tensor("z", shape=z_shape, dtype=floatX) + x = pt.tensor("x", shape=(6, 3), dtype=floatX) + y = pt.tensor("y", shape=(3, 5), dtype=floatX) + alpha, beta = pt.scalar("alpha", dtype=floatX), pt.scalar("beta", dtype=floatX) + + rng = np.random.default_rng(sum(map(ord, f"gemm_broadcasts {z_shape}"))) + z_np = rng.normal(size=z_shape).astype(floatX) + x_np = rng.normal(size=(6, 3)).astype(floatX) + y_np = rng.normal(size=(3, 5)).astype(floatX) + + fn = pytensor.function( + [z, alpha, x, y, beta], Gemm(inplace=False)(z, alpha, x, y, beta), mode="NUMBA" + ) + np.testing.assert_allclose( + fn(z_np, 2.0, x_np, y_np, 0.5), 0.5 * z_np + 2.0 * (x_np @ y_np), rtol=1e-5 + ) + + +def test_gemm_no_inplace_leaves_its_accumulator_alone(): + """The non-inplace form must not write through ``z``; only the inplace form may.""" + z = pt.tensor("z", shape=(6, 5), dtype=floatX) + x = pt.tensor("x", shape=(6, 3), dtype=floatX) + y = pt.tensor("y", shape=(3, 5), dtype=floatX) + alpha, beta = pt.scalar("alpha", dtype=floatX), pt.scalar("beta", dtype=floatX) + + rng = np.random.default_rng(sum(map(ord, "gemm_no_inplace"))) + z_np = rng.normal(size=(6, 5)).astype(floatX) + original = z_np.copy() + x_np = rng.normal(size=(6, 3)).astype(floatX) + y_np = rng.normal(size=(3, 5)).astype(floatX) + + fn = pytensor.function( + [z, alpha, x, y, beta], Gemm(inplace=False)(z, alpha, x, y, beta), mode="NUMBA" + ) + result = fn(z_np, 2.0, x_np, y_np, 0.5) + np.testing.assert_allclose(result, 0.5 * original + 2.0 * (x_np @ y_np), rtol=1e-5) + np.testing.assert_array_equal(z_np, original) + + +def test_gemm_inplace_writes_through_its_accumulator(): + """The inplace form is chosen precisely to avoid the copy, so it has to destroy ``z``.""" + rng = np.random.default_rng(sum(map(ord, "gemm_inplace"))) + z_np = rng.normal(size=(6, 5)).astype(floatX) + x_np = rng.normal(size=(6, 3)).astype(floatX) + y_np = rng.normal(size=(3, 5)).astype(floatX) + expected = 0.5 * z_np + 2.0 * (x_np @ y_np) + + z = pytensor.shared(z_np) + x = pt.tensor("x", shape=(6, 3), dtype=floatX) + y = pt.tensor("y", shape=(3, 5), dtype=floatX) + alpha, beta = pt.scalar("alpha", dtype=floatX), pt.scalar("beta", dtype=floatX) + + # An inplace op cannot be built into a graph directly; only the rewrite that introduces one + # may, so the test has to say it knows what it is asking for. + fn = pytensor.function( + [alpha, x, y, beta], + Gemm(inplace=True)(z, alpha, x, y, beta), + mode="NUMBA", + accept_inplace=True, + ) + np.testing.assert_allclose(fn(2.0, x_np, y_np, 0.5), expected, rtol=1e-5) + np.testing.assert_allclose(z.get_value(), expected, rtol=1e-5) + + +def test_gemv_no_inplace_leaves_its_accumulator_alone(): + """The non-inplace form must not write through ``y``; only the inplace form may.""" + y = pt.tensor("y", shape=(6,), dtype=floatX) + A = pt.tensor("A", shape=(6, 5), dtype=floatX) + x = pt.tensor("x", shape=(5,), dtype=floatX) + alpha, beta = pt.scalar("alpha", dtype=floatX), pt.scalar("beta", dtype=floatX) + + rng = np.random.default_rng(sum(map(ord, "gemv_no_inplace"))) + y_np = rng.normal(size=6).astype(floatX) + original = y_np.copy() + A_np = rng.normal(size=(6, 5)).astype(floatX) + x_np = rng.normal(size=5).astype(floatX) + + fn = pytensor.function( + [y, alpha, A, x, beta], Gemv(inplace=False)(y, alpha, A, x, beta), mode="NUMBA" + ) + result = fn(y_np, 2.0, A_np, x_np, 0.5) + np.testing.assert_allclose(result, 0.5 * original + 2.0 * (A_np @ x_np), rtol=1e-5) + np.testing.assert_array_equal(y_np, original) + + +def test_gemv_inplace_writes_through_its_accumulator(): + """The inplace form is chosen precisely to avoid the copy, so it has to destroy ``y``.""" + rng = np.random.default_rng(sum(map(ord, "gemv_inplace"))) + y_np = rng.normal(size=6).astype(floatX) + A_np = rng.normal(size=(6, 5)).astype(floatX) + x_np = rng.normal(size=5).astype(floatX) + expected = 0.5 * y_np + 2.0 * (A_np @ x_np) + + y = pytensor.shared(y_np) + A = pt.tensor("A", shape=(6, 5), dtype=floatX) + x = pt.tensor("x", shape=(5,), dtype=floatX) + alpha, beta = pt.scalar("alpha", dtype=floatX), pt.scalar("beta", dtype=floatX) + + fn = pytensor.function( + [alpha, A, x, beta], + Gemv(inplace=True)(y, alpha, A, x, beta), + mode="NUMBA", + accept_inplace=True, + ) + np.testing.assert_allclose(fn(2.0, A_np, x_np, 0.5), expected, rtol=1e-5) + np.testing.assert_allclose(y.get_value(), expected, rtol=1e-5) + + +def test_ger_no_inplace_leaves_its_accumulator_alone(): + """The non-inplace form must not write through ``A``; only the inplace form may.""" + A = pt.tensor("A", shape=(6, 5), dtype=floatX) + x = pt.tensor("x", shape=(6,), dtype=floatX) + y = pt.tensor("y", shape=(5,), dtype=floatX) + alpha = pt.scalar("alpha", dtype=floatX) + + rng = np.random.default_rng(sum(map(ord, "ger_no_inplace"))) + A_np = rng.normal(size=(6, 5)).astype(floatX) + original = A_np.copy() + x_np = rng.normal(size=6).astype(floatX) + y_np = rng.normal(size=5).astype(floatX) + + fn = pytensor.function( + [A, alpha, x, y], Ger(inplace=False)(A, alpha, x, y), mode="NUMBA" + ) + result = fn(A_np, 2.0, x_np, y_np) + np.testing.assert_allclose(result, original + 2.0 * np.outer(x_np, y_np), rtol=1e-5) + np.testing.assert_array_equal(A_np, original) + + +def test_ger_inplace_writes_through_its_accumulator(): + """The inplace form is chosen precisely to avoid the copy, so it has to destroy ``A``.""" + rng = np.random.default_rng(sum(map(ord, "ger_inplace"))) + A_np = rng.normal(size=(6, 5)).astype(floatX) + x_np = rng.normal(size=6).astype(floatX) + y_np = rng.normal(size=5).astype(floatX) + expected = A_np + 2.0 * np.outer(x_np, y_np) + + A = pytensor.shared(A_np) + x = pt.tensor("x", shape=(6,), dtype=floatX) + y = pt.tensor("y", shape=(5,), dtype=floatX) + alpha = pt.scalar("alpha", dtype=floatX) + + fn = pytensor.function( + [alpha, x, y], + Ger(inplace=True)(A, alpha, x, y), + mode="NUMBA", + accept_inplace=True, + ) + np.testing.assert_allclose(fn(2.0, x_np, y_np), expected, rtol=1e-5) + np.testing.assert_allclose(A.get_value(), expected, rtol=1e-5) + + +def test_ger_inplace_reads_strided_vectors(): + """``ger`` walks each vector by a fixed increment, so a strided view has to be copied.""" + rng = np.random.default_rng(sum(map(ord, "ger_strided"))) + A_np = rng.normal(size=(6, 5)).astype(floatX) + x_np = rng.normal(size=12).astype(floatX) + y_np = rng.normal(size=10).astype(floatX) + expected = A_np + 2.0 * np.outer(x_np[::2], y_np[::2]) + + A = pytensor.shared(A_np) + x = pt.tensor("x", shape=(12,), dtype=floatX) + y = pt.tensor("y", shape=(10,), dtype=floatX) + alpha = pt.scalar("alpha", dtype=floatX) + + fn = pytensor.function( + [alpha, x, y], + Ger(inplace=True)(A, alpha, x[::2], y[::2]), + mode="NUMBA", + accept_inplace=True, + ) + np.testing.assert_allclose(fn(2.0, x_np, y_np), expected, rtol=1e-5) diff --git a/tests/link/pytorch/test_blas.py b/tests/link/pytorch/test_blas.py index 4b9fc4d55f..beec5a92e7 100644 --- a/tests/link/pytorch/test_blas.py +++ b/tests/link/pytorch/test_blas.py @@ -1,6 +1,7 @@ import numpy as np import pytest +from pytensor import tensor as pt from pytensor.configdefaults import config from pytensor.tensor import blas as pt_blas from pytensor.tensor.type import tensor3 @@ -21,3 +22,93 @@ def test_pytorch_BatchedDot(): inputs = [a_test[:-1], b_test] with pytest.raises(TypeError): pytensor_pytorch_fn(*inputs) + + +@pytest.mark.parametrize( + "alpha, beta", + [(None, None), (2.0, None), (None, 3.0), (2.0, 3.0)], + ids=["plain", "alpha", "beta", "alpha_beta"], +) +def test_pytorch_Gemm(alpha, beta): + # Gemm is what local_add_dot_to_gemm folds `beta * C + alpha * (A @ B)` into, so the + # scales arrive as constants. They reach PyTorch's fused kernel only in that form. + z = pt.matrix("z", dtype=config.floatX) + x = pt.matrix("x", dtype=config.floatX) + y = pt.matrix("y", dtype=config.floatX) + one = np.asarray(1.0, dtype=config.floatX) + + out = pt_blas.Gemm(inplace=False)( + z, + one if alpha is None else np.asarray(alpha, dtype=config.floatX), + x, + y, + one if beta is None else np.asarray(beta, dtype=config.floatX), + ) + + rng = np.random.default_rng(418) + test_values = [ + rng.normal(size=shape).astype(config.floatX) + for shape in ((4, 6), (4, 5), (5, 6)) + ] + compare_pytorch_and_py([z, x, y], [out], test_values) + + +def test_pytorch_Gemm_runtime_scales(): + # With alpha and beta as graph variables there is no constant to hand the fused + # kernel, so this covers the arithmetic path the dispatch falls back to. + z = pt.matrix("z", dtype=config.floatX) + x = pt.matrix("x", dtype=config.floatX) + y = pt.matrix("y", dtype=config.floatX) + alpha = pt.scalar("alpha", dtype=config.floatX) + beta = pt.scalar("beta", dtype=config.floatX) + + out = pt_blas.Gemm(inplace=False)(z, alpha, x, y, beta) + + rng = np.random.default_rng(418) + test_values = [ + rng.normal(size=shape).astype(config.floatX) + for shape in ((4, 6), (4, 5), (5, 6)) + ] + compare_pytorch_and_py( + [z, alpha, x, y, beta], + [out], + [ + test_values[0], + np.asarray(2.0, dtype=config.floatX), + test_values[1], + test_values[2], + np.asarray(3.0, dtype=config.floatX), + ], + ) + + +def test_pytorch_Ger(): + # local_gemm_to_ger folds a rank-1 `A + alpha * outer(x, y)` into Ger. + A = pt.matrix("A", dtype=config.floatX) + x = pt.vector("x", dtype=config.floatX) + y = pt.vector("y", dtype=config.floatX) + out = pt_blas.Ger(inplace=False)(A, np.asarray(2.0, dtype=config.floatX), x, y) + + rng = np.random.default_rng(418) + test_values = [ + rng.normal(size=shape).astype(config.floatX) for shape in ((4, 6), (4,), (6,)) + ] + compare_pytorch_and_py([A, x, y], [out], test_values) + + +def test_pytorch_Ger_runtime_alpha(): + A = pt.matrix("A", dtype=config.floatX) + x = pt.vector("x", dtype=config.floatX) + y = pt.vector("y", dtype=config.floatX) + alpha = pt.scalar("alpha", dtype=config.floatX) + out = pt_blas.Ger(inplace=False)(A, alpha, x, y) + + rng = np.random.default_rng(418) + A_val, x_val, y_val = ( + rng.normal(size=shape).astype(config.floatX) for shape in ((4, 6), (4,), (6,)) + ) + compare_pytorch_and_py( + [A, alpha, x, y], + [out], + [A_val, np.asarray(2.0, dtype=config.floatX), x_val, y_val], + ) diff --git a/tests/tensor/rewriting/test_blockwise.py b/tests/tensor/rewriting/test_blockwise.py index c511e21f66..5b19445c8e 100644 --- a/tests/tensor/rewriting/test_blockwise.py +++ b/tests/tensor/rewriting/test_blockwise.py @@ -9,6 +9,8 @@ from pytensor.graph.traversal import apply_ancestors from pytensor.scalar import log as scalar_log from pytensor.tensor import add, alloc, iscalar, matrix, scalar, tensor, tensor3 +from pytensor.tensor.basic import AllocEmpty +from pytensor.tensor.blas import Gemv from pytensor.tensor.blockwise import Blockwise, BlockwiseWithCoreShape from pytensor.tensor.elemwise import Elemwise from pytensor.tensor.linalg.inverse import MatrixPinv @@ -186,3 +188,38 @@ def test_blockwise_reshape(): new_y.eval({"x": test_x}, mode=no_rewrites), rewritten_y.eval({"x": test_x}, mode=no_rewrites), ) + + +def test_split_alloc_empty_clients_enables_inplace(): + """Two destructive clients of one `AllocEmpty` each get a buffer they can destroy.""" + x = matrix("x") + y = tensor("y", shape=(None,)) + z = tensor("z", shape=(None,)) + + f = function([x, y, z], [y @ x, z @ x], mode="cvm") + nodes = f.maker.fgraph.apply_nodes + gemvs = [n.op for n in nodes if isinstance(n.op, Gemv)] + assert len(gemvs) == 2 + assert all(op.inplace for op in gemvs), gemvs + # One buffer per destroyer, rather than one shared between them. + assert len([n for n in nodes if isinstance(n.op, AllocEmpty)]) == 2 + + rng = np.random.default_rng(sum(map(ord, "split_alloc_empty"))) + x_val = rng.normal(size=(3, 3)) + y_val = rng.normal(size=(3,)) + z_val = rng.normal(size=(3,)) + out_y, out_z = f(x_val, y_val, z_val) + np.testing.assert_allclose(out_y, y_val @ x_val) + np.testing.assert_allclose(out_z, z_val @ x_val) + + +def test_split_alloc_empty_clients_leaves_readers_alone(): + """An `AllocEmpty` shared by clients that cannot destroy it stays a single buffer.""" + shape = iscalar("shape") + buffer = AllocEmpty(config.floatX)(shape) + out = add(buffer * 2, buffer * 3) + + fg = FunctionGraph([shape], [out]) + rewrite_graph(fg, include=("fast_run", "inplace")) + allocs = [n for n in fg.apply_nodes if isinstance(n.op, AllocEmpty)] + assert len(allocs) == 1, allocs diff --git a/tests/tensor/rewriting/test_math.py b/tests/tensor/rewriting/test_math.py index f6683b621d..19aed8c93a 100644 --- a/tests/tensor/rewriting/test_math.py +++ b/tests/tensor/rewriting/test_math.py @@ -33,7 +33,6 @@ from pytensor.scalar import PolyGamma, Psi, TriGamma from pytensor.tensor.basic import Alloc, constant, join, second, switch from pytensor.tensor.blas import Dot22, Gemv -from pytensor.tensor.blas.blas_c import CGemv from pytensor.tensor.blockwise import Blockwise from pytensor.tensor.elemwise import CAReduce, DimShuffle, Elemwise from pytensor.tensor.linalg.constructors import BlockDiagonal @@ -4034,7 +4033,7 @@ def test_local_sumsqr2dot(): assert any( isinstance( n.op, - Dot | Dot22 | Gemv | CGemv, + Dot | Dot22 | Gemv, ) for n in f.maker.fgraph.toposort() ) diff --git a/tests/tensor/rewriting/test_subtensor_lift.py b/tests/tensor/rewriting/test_subtensor_lift.py index 76f6d65308..ca62e7a9eb 100644 --- a/tests/tensor/rewriting/test_subtensor_lift.py +++ b/tests/tensor/rewriting/test_subtensor_lift.py @@ -46,7 +46,6 @@ make_vector, ) from pytensor.tensor.blas import Dot22, Gemv -from pytensor.tensor.blas.blas_c import CGemv from pytensor.tensor.blockwise import Blockwise from pytensor.tensor.elemwise import DimShuffle, Elemwise from pytensor.tensor.math import Dot @@ -386,7 +385,7 @@ def test_equality(a, b): topo = f.maker.fgraph.toposort() assert test_equality(f(d1, d2), np.dot(d1, d2)[1]) # DimShuffle happen in FAST_COMPILE - assert isinstance(topo[-1].op, CGemv | Gemv | DimShuffle) + assert isinstance(topo[-1].op, Gemv | DimShuffle) # slice f = function([m1, m2], pt.dot(m1, m2)[1:2], mode=mode) diff --git a/tests/tensor/test_blas.py b/tests/tensor/test_blas.py index 184c90f453..0e8b5a3b8f 100644 --- a/tests/tensor/test_blas.py +++ b/tests/tensor/test_blas.py @@ -28,16 +28,8 @@ _batched_dot, _dot22, _dot22scalar, - gemm, - gemm_inplace, - gemm_no_inplace, - gemv, - gemv_inplace, - gemv_no_inplace, - ger, - ger_destructive, ) -from pytensor.tensor.elemwise import DimShuffle +from pytensor.tensor.elemwise import DimShuffle, Elemwise from pytensor.tensor.math import Dot, dot, mean, mul, outer, sigmoid from pytensor.tensor.rewriting.blas import local_dot22_to_dot22scalar, local_gemm_to_ger from pytensor.tensor.type import ( @@ -81,7 +73,7 @@ def sharedX(x, name): class TestGemm: """ - This test suite is supposed to establish that gemm works as it is supposed to. + This test suite is supposed to establish that Gemm works as it is supposed to. """ def setup_method(self): @@ -110,7 +102,7 @@ def cmp_linker(z, a, x, y, b, l): f = inplace_func( [tz, ta, tx, ty, tb], - gemm_inplace(tz, ta, tx, ty, tb), + Gemm(inplace=True)(tz, ta, tx, ty, tb), mode=Mode(optimizer=None, linker=l), ) f(z, a, x, y, b) @@ -136,7 +128,7 @@ def cmp_linker(z, a, x, y, b, l): def test_basic(self): Gemm.debug = True with pytest.raises(TypeError, match=Gemm.E_rank): - gemm_no_inplace([1.0], 1.0, [1.0], [1.0], 1.0) + Gemm(inplace=False)([1.0], 1.0, [1.0], [1.0], 1.0) def test_basic_1(self): with pytest.raises(TypeError, match=Gemm.E_rank): @@ -199,7 +191,7 @@ def test_factorised_scalar(self): lr2 = pt.constant(2).astype(config.floatX) l2_reg = pt.constant(0.0001).astype(config.floatX) - # test constant merge with gemm + # test constant merge with Gemm f = function( [a, b], updates=[(s, lr1 * dot(a, b) + l2_reg * lr2 * s)], @@ -209,7 +201,7 @@ def test_factorised_scalar(self): # , , # 2e-06)] assert len(f) == 1 - assert f[0].op == gemm_inplace + assert f[0].op == Gemm(inplace=True) # test factored scalar with merge f = function( @@ -221,7 +213,7 @@ def test_factorised_scalar(self): # , , # -2e-06)] assert len(f) == 1 - assert f[0].op == gemm_inplace + assert f[0].op == Gemm(inplace=True) # test factored scalar with merge and neg f = function( @@ -233,14 +225,14 @@ def test_factorised_scalar(self): # , , # 0.999998)] assert len(f) == 1 - assert f[0].op == gemm_inplace + assert f[0].op == Gemm(inplace=True) def test_destroy_map0(self): # test that only first input can be overwritten. rng = np.random.default_rng(seed=utt.fetch_seed()) Z = as_tensor_variable(rng.random((2, 2))) with pytest.raises(InconsistencyError, match=Gemm.E_z_uniq): - gemm_inplace(Z, 1.0, Z, Z, 1.0) + Gemm(inplace=True)(Z, 1.0, Z, Z, 1.0) def test_destroy_map1(self): # test that only first input can be overwritten. @@ -250,7 +242,7 @@ def test_destroy_map1(self): Zt = Z.transpose() assert isinstance(Zt.owner.op, DimShuffle) and Zt.owner.op.view_map == {0: [0]} with pytest.raises(InconsistencyError, match=Gemm.E_z_uniq): - gemm_inplace(Z, 1.0, A, Zt, 1.0) + Gemm(inplace=True)(Z, 1.0, A, Zt, 1.0) def test_destroy_map2(self): # test that only first input can be overwritten. @@ -260,7 +252,7 @@ def test_destroy_map2(self): Zt = Z.transpose() assert isinstance(Zt.owner.op, DimShuffle) and Zt.owner.op.view_map == {0: [0]} with pytest.raises(InconsistencyError, match=Gemm.E_z_uniq): - gemm_inplace(Z, 1.0, Zt, A, 1.0) + Gemm(inplace=True)(Z, 1.0, Zt, A, 1.0) def test_destroy_map3(self): # test that only first input can be overwritten @@ -268,7 +260,7 @@ def test_destroy_map3(self): Z = as_tensor_variable(rng.random((2, 2))) A = as_tensor_variable(rng.random((2, 2))) with pytest.raises(InconsistencyError, match=Gemm.E_z_uniq): - gemm_inplace(Z, 1.0, Z, A, 1.0) + Gemm(inplace=True)(Z, 1.0, Z, A, 1.0) def test_destroy_map4(self): # test that dot args can be aliased @@ -276,10 +268,10 @@ def test_destroy_map4(self): Z = shared(rng.random((2, 2)), name="Z") A = shared(rng.random((2, 2)), name="A") one = pt.constant(1.0).astype(Z.dtype) - f = inplace_func([], gemm_inplace(Z, one, A, A, one)) + f = inplace_func([], Gemm(inplace=True)(Z, one, A, A, one)) # TODO FIXME: This is a bad test f() - f = inplace_func([], gemm_inplace(Z, one, A, A.T, one)) + f = inplace_func([], Gemm(inplace=True)(Z, one, A, A.T, one)) # TODO FIXME: This is a bad test f() @@ -292,17 +284,13 @@ def test_transposes(self): def t(z, x, y, a=1.0, b=0.0, l="c|py", dt="float64"): z, a, x, y, b = (np.asarray(p, dtype=dt) for p in (z, a, x, y, b)) - # z_orig = z.copy() z_after = self._gemm(z, a, x, y, b) tz, ta, tx, ty, tb = (shared(p) for p in (z, a, x, y, b)) - # f = inplace_func([tz,ta,tx,ty,tb], gemm_inplace(tz,ta,tx,ty,tb), - # mode = Mode(optimizer = None, linker=l)) - # f(z, a, x, y, b) f = inplace_func( [], - gemm_inplace(tz, ta, tx, ty, tb), + Gemm(inplace=True)(tz, ta, tx, ty, tb), mode=Mode(optimizer=None, linker=l), ) f() @@ -360,7 +348,7 @@ def t(z, x, y, a=1.0, b=0.0, l="c|py", dt="float64"): for i in range(3): f_i = inplace_func( [], - gemm_inplace(tz[:, :, i], ta, tx[:, :, i], ty[:, :, i], tb), + Gemm(inplace=True)(tz[:, :, i], ta, tx[:, :, i], ty[:, :, i], tb), mode=Mode(optimizer=None, linker=l), ) for j in range(3): @@ -375,7 +363,9 @@ def t(z, x, y, a=1.0, b=0.0, l="c|py", dt="float64"): z_after[:, :, i], tz.get_value(borrow=True)[:, :, i] ) - tz_i = gemm_no_inplace(tz[:, :, i], ta, tx[:, :, i], ty[:, :, i], tb) + tz_i = Gemm(inplace=False)( + tz[:, :, i], ta, tx[:, :, i], ty[:, :, i], tb + ) g_i = function( [], tz_i, @@ -404,7 +394,7 @@ def t(z, x, y, a=1.0, b=0.0, l="c|py", dt="float64"): class TestGemmNoFlags: - gemm = gemm_no_inplace + gemm = Gemm(inplace=False) M = 4 N = 5 K = 6 @@ -582,10 +572,10 @@ def just_gemm(i, o, ishapes=None, max_graphlen=0, expected_nb_gemm=1): nb_gemm = 0 for node in f.maker.fgraph.apply_nodes: assert not isinstance(node.op, Dot), ( - "_dot22 not changed to gemm_inplace in graph" + "_dot22 not changed to an inplace Gemm in graph" ) assert node.op != _dot22 - if node.op == gemm_inplace: + if node.op == Gemm(inplace=True): nb_gemm += 1 assert nb_gemm == expected_nb_gemm, (nb_gemm, expected_nb_gemm) g = inplace_func( @@ -596,7 +586,7 @@ def just_gemm(i, o, ishapes=None, max_graphlen=0, expected_nb_gemm=1): on_unused_input="ignore", ) for node in g.maker.fgraph.apply_nodes: - assert node.op != gemm_inplace, "gemm_inplace in original graph" + assert node.op != Gemm(inplace=True), "Gemm(inplace=True) in original graph" graphlen = len(f.maker.fgraph.toposort()) assert not (max_graphlen and (graphlen <= max_graphlen)), ( @@ -667,7 +657,7 @@ def test_gemm_opt_double_gemm(): o = [ ( a * dot(X, Y) - + gemm_inplace(Z, b, S.T, R.T, pt.constant(1.0).astype(config.floatX)) + + Gemm(inplace=True)(Z, b, S.T, R.T, pt.constant(1.0).astype(config.floatX)) ) ] f = inplace_func( @@ -812,15 +802,15 @@ def test_gemm_opt_vector_stuff(): u, v = vector(), vector() f = inplace_func([a, u, v], a + dot(u, v), mode=mode_blas_opt) - assert gemm_inplace not in [n.op for n in f.maker.fgraph.apply_nodes] + assert Gemm(inplace=True) not in [n.op for n in f.maker.fgraph.apply_nodes] f = inplace_func([a, u, X, Y], a * u + dot(X, Y), mode=mode_blas_opt) - assert gemm_inplace not in [n.op for n in f.maker.fgraph.apply_nodes] + assert Gemm(inplace=True) not in [n.op for n in f.maker.fgraph.apply_nodes] def test_gemm_unrolled(): - # This test that the gemm optimizer remove the dot22 that was - # present in the graph. Otherwise, this add a gemm, but still + # This test that the Gemm optimizer remove the dot22 that was + # present in the graph. Otherwise, this add a Gemm, but still # compute the dot22. # This was not always the case in the with this the following code. @@ -868,7 +858,7 @@ def update_H(cur_V): def test_inplace0(): - # should fail to insert gemm_inplace because gemm_inplace would + # should fail to insert an inplace Gemm because it would # create cycles X, Y, Z, a, b = ( matrix("X"), @@ -880,16 +870,16 @@ def test_inplace0(): R, S, c = matrix("R"), matrix("S"), scalar("c") f = inplace_func([Z, b, R, S], [Z * (Z + b * dot(R, S).T)], mode=mode_blas_opt) - assert gemm_inplace not in [n.op for n in f.maker.fgraph.apply_nodes] - assert gemm_no_inplace in [n.op for n in f.maker.fgraph.apply_nodes] + assert Gemm(inplace=True) not in [n.op for n in f.maker.fgraph.apply_nodes] + assert Gemm(inplace=False) in [n.op for n in f.maker.fgraph.apply_nodes] - # gemm_inplace should be inserted here, to work in-place on Z*c + # An inplace Gemm should be inserted here, to work in-place on Z*c f = inplace_func( [X, Y, Z, a, b, R, S, c], [Z * (c * Z + a * dot(X, Y) + b * dot(R, S).T)], mode=mode_blas_opt, ) - assert gemm_inplace in [n.op for n in f.maker.fgraph.apply_nodes] + assert Gemm(inplace=True) in [n.op for n in f.maker.fgraph.apply_nodes] def test_inplace1(): @@ -898,7 +888,7 @@ def test_inplace1(): f = inplace_func([X, Y, Z], [Z + Z + dot(X, Y)], mode=mode_blas_opt) # pytensor.printing.debugprint(f) # it doesn't work inplace because we didn't mark Z as mutable input - assert [n.op for n in f.maker.fgraph.apply_nodes] == [gemm_no_inplace] + assert [n.op for n in f.maker.fgraph.apply_nodes] == [Gemm(inplace=False)] @pytest.mark.parametrize("linker", ("py", "cvm")) @@ -909,13 +899,13 @@ def test_gemm_broadcasting(inplace, linker): mode = Mode(linker=linker) if inplace: - out = gemm_inplace(z, a, x, y, b) + out = Gemm(inplace=True)(z, a, x, y, b) f = pytensor.function([z, x, y, a, b], out, accept_inplace=True, mode=mode) - assert [node.op for node in f.maker.fgraph.toposort()] == [gemm_inplace] + assert [node.op for node in f.maker.fgraph.toposort()] == [Gemm(inplace=True)] else: - out = gemm_no_inplace(z, a, x, y, b) + out = Gemm(inplace=False)(z, a, x, y, b) f = pytensor.function([z, x, y, a, b], out, mode=mode) - assert [node.op for node in f.maker.fgraph.toposort()] == [gemm_no_inplace] + assert [node.op for node in f.maker.fgraph.toposort()] == [Gemm(inplace=False)] shapes_z = [(5, 3), (1, 3), (5, 1), (1, 1)] shapes_x = [(5, 4), (1, 4)] @@ -943,7 +933,7 @@ def test_gemm_static_shape(): z = matrix("z", shape=(1, 1)) x = matrix("x", shape=(5, 4)) y = matrix("y", shape=(4, 3)) - assert gemm_no_inplace(z, a, x, y, b).type.shape == (5, 3) + assert Gemm(inplace=False)(z, a, x, y, b).type.shape == (5, 3) def test_dot22(): @@ -1169,7 +1159,7 @@ def test_local_dot22_to_dot22scalar(): def test_dot_w_self(): # This can trigger problems in the optimization because what would - # normally be a gemm must not be because the output is aliased to + # normally be a Gemm must not be because the output is aliased to # one of the inputs. A = shared(value=np.ones((2, 2))) @@ -1191,7 +1181,7 @@ def test_dot_w_self(): class TestGemv(unittest_tools.OptimizationTestMixin): def test_dot_vv(self): - # Currently we generate a gemv for that case + # Currently we generate a Gemv for that case rng = np.random.default_rng(unittest_tools.fetch_seed()) v = shared(np.array(rng.uniform(size=(2,)), dtype="float32")) w = shared(np.array(rng.uniform(size=(2,)), dtype="float32")) @@ -1329,7 +1319,7 @@ def test_gemv2(self): ) def test_gemv_broadcast(self): - # test gemv with some broadcasted input + # test Gemv with some broadcasted input rng = np.random.default_rng(unittest_tools.fetch_seed()) v1 = shared(np.array(rng.uniform(size=(2,)), dtype="float32")) v2_orig = np.array(rng.uniform(size=(1,)), dtype="float32") @@ -1346,8 +1336,8 @@ def test_gemv_broadcast(self): topo = f.maker.fgraph.toposort() assert sum(isinstance(node.op, Gemv) for node in topo) == 1 - # call gemv directly for mixed broadcast pattern. - o = gemv_no_inplace(v2, 0.5, m, v1, 0.25) + # call Gemv directly for mixed broadcast pattern. + o = Gemv(inplace=False)(v2, 0.5, m, v1, 0.25) f = function([], o, mode=mode_blas_opt) assert np.allclose( f(), 0.5 * np.dot(m.get_value(), v1.get_value()) + 0.25 * v2.get_value() @@ -1382,8 +1372,8 @@ def test_gemv_dimensions(self): f(A_val, ones_4, ones_6) -# The following gemv tests were added in March 2011 by Ian Goodfellow -# and are based on the gemv tests from scipy +# The following Gemv tests were added in March 2011 by Ian Goodfellow +# and are based on the Gemv tests from scipy # http://projects.scipy.org/scipy/browser/trunk/scipy/linalg/tests/test_fblas.py?rev=6803 # NOTE: At the time these tests were written, pytensor did not have a # conjugate function. If such a thing is ever added, the tests involving @@ -1603,7 +1593,7 @@ def test_upcasting_scalar_nogemv(self): rval = dot(a, x) * alpha + y f = function([alpha], rval, mode=self.mode) - # this function is currently optimized so that the gemv is + # this function is currently optimized so that the Gemv is # done inplace on a temporarily allocated-buffer, which is # then scaled by alpha and to t with a fused elemwise. n_gemvs = 0 @@ -1620,14 +1610,14 @@ def test_upcasting_scalar_nogemv(self): class TestSgemv(BaseGemv, unittest_tools.OptimizationTestMixin): dtype = np.float32 - gemv = gemv_no_inplace - gemv_inplace = gemv_inplace + gemv = Gemv(inplace=False) + gemv_inplace = Gemv(inplace=True) class TestDgemv(BaseGemv, unittest_tools.OptimizationTestMixin): dtype = np.float64 - gemv = gemv_no_inplace - gemv_inplace = gemv_inplace + gemv = Gemv(inplace=False) + gemv_inplace = Gemv(inplace=True) # The optimization to put Gemv don't work for complex type for now. @@ -1671,58 +1661,73 @@ def setup_method(self): self.za = zscalar() def test_works_on_all_valid_dtypes(self): - assert self.fm.type == ger(self.fm, self.fa, self.fv, self.fv_2).type - assert self.fm.type == ger(self.fm, self.fa, self.fv, self.fv_2).type - assert self.fm.type == ger(self.fm, self.fa, self.fv, self.fv_2).type - assert self.fm.type == ger(self.fm, self.fa, self.fv, self.fv_2).type + assert ( + self.fm.type + == Ger(inplace=False)(self.fm, self.fa, self.fv, self.fv_2).type + ) + assert ( + self.fm.type + == Ger(inplace=False)(self.fm, self.fa, self.fv, self.fv_2).type + ) + assert ( + self.fm.type + == Ger(inplace=False)(self.fm, self.fa, self.fv, self.fv_2).type + ) + assert ( + self.fm.type + == Ger(inplace=False)(self.fm, self.fa, self.fv, self.fv_2).type + ) def test_fails_on_invalid_dtypes(self): with pytest.raises(TypeError): - ger(imatrix(), iscalar(), ivector(), ivector()) + Ger(inplace=False)(imatrix(), iscalar(), ivector(), ivector()) def test_fails_for_nonscalar_alpha(self): with pytest.raises(TypeError): - ger(self.fm, self.fm, self.fv, self.fv_2) + Ger(inplace=False)(self.fm, self.fm, self.fv, self.fv_2) # boundary case - fv1 has the right dtype and could be dimshuffled to a # scalar, but that's not make_node's job. with pytest.raises(TypeError): - ger(self.fm, self.fv1, self.fv, self.fv_2) + Ger(inplace=False)(self.fm, self.fv1, self.fv, self.fv_2) # actually doing the aforementioned dimshuffle makes it work assert ( - self.fm.type == ger(self.fm, self.fv1.dimshuffle(), self.fv, self.fv_2).type + self.fm.type + == Ger(inplace=False)( + self.fm, self.fv1.dimshuffle(), self.fv, self.fv_2 + ).type ) def test_fails_for_nonmatrix_A(self): with pytest.raises(TypeError): - ger(self.fv, self.fa, self.fv, self.fv_2) + Ger(inplace=False)(self.fv, self.fa, self.fv, self.fv_2) def test_fails_for_nonvector_x_or_y(self): with pytest.raises(TypeError): - ger(self.fm, self.fa, self.fv.dimshuffle("x", 0), self.fv_2) + Ger(inplace=False)(self.fm, self.fa, self.fv.dimshuffle("x", 0), self.fv_2) with pytest.raises(TypeError): - ger(self.fm, self.fa, self.fv, self.fv_2.dimshuffle("x", 0)) + Ger(inplace=False)(self.fm, self.fa, self.fv, self.fv_2.dimshuffle("x", 0)) def test_fails_for_mixed_dtypes(self): with pytest.raises(TypeError): - ger(self.dm, self.fa, self.fv, self.fv_2) + Ger(inplace=False)(self.dm, self.fa, self.fv, self.fv_2) with pytest.raises(TypeError): - ger(self.fm, self.da, self.fv, self.fv_2) + Ger(inplace=False)(self.fm, self.da, self.fv, self.fv_2) with pytest.raises(TypeError): - ger(self.fm, self.fa, self.dv, self.fv_2) + Ger(inplace=False)(self.fm, self.fa, self.dv, self.fv_2) with pytest.raises(TypeError): - ger(self.fm, self.fa, self.fv, self.dv_2) + Ger(inplace=False)(self.fm, self.fa, self.fv, self.dv_2) with pytest.raises(TypeError): - ger(self.cm, self.fa, self.fv, self.dv_2) + Ger(inplace=False)(self.cm, self.fa, self.fv, self.dv_2) with pytest.raises(TypeError): - ger(self.cm, self.fa, self.fv, self.zv_2) + Ger(inplace=False)(self.cm, self.fa, self.fv, self.zv_2) class TestGerOpContract(unittest_tools.OpContractTestMixin): def setup_method(self): - self.ops = [ger, ger_destructive] + self.ops = [Ger(inplace=False), Ger(inplace=True)] def clone(self, op): - return Ger(op.destructive) + return Ger(op.inplace) class TestGer(unittest_tools.OptimizationTestMixin): @@ -1735,9 +1740,9 @@ def setup_method(self): self.a = tensor(dtype=dtype, shape=()) self.x = tensor(dtype=dtype, shape=(None,)) self.y = tensor(dtype=dtype, shape=(None,)) - self.ger = ger - self.ger_destructive = ger_destructive - self.gemm = gemm_no_inplace + self.ger = Ger(inplace=False) + self.ger_inplace = Ger(inplace=True) + self.gemm = Gemm(inplace=False) def function(self, inputs, outputs, updates=None): if updates is None: @@ -1751,7 +1756,7 @@ def test_b_0_triggers_ger(self): # test local_gemm_to_ger opt assert local_gemm_to_ger.transform( None, - gemm_no_inplace( + Gemm(inplace=False)( self.A, self.a, self.x.dimshuffle(0, "x"), @@ -1764,7 +1769,7 @@ def test_b_1_triggers_ger(self): # test local_gemm_to_ger opt assert local_gemm_to_ger.transform( None, - gemm_no_inplace( + Gemm(inplace=False)( self.A, self.a, self.x.dimshuffle(0, "x"), @@ -1777,7 +1782,7 @@ def test_b_other_does_not_triggers_ger(self): # test local_gemm_to_ger opt assert not local_gemm_to_ger.transform( None, - gemm_no_inplace( + Gemm(inplace=False)( self.A, self.a, self.x.dimshuffle(0, "x"), @@ -1790,7 +1795,7 @@ def test_b_nonconst_does_not_triggers_ger(self): # test local_gemm_to_ger opt assert not local_gemm_to_ger.transform( None, - gemm_no_inplace( + Gemm(inplace=False)( self.A, self.a, self.x.dimshuffle(0, "x"), @@ -1802,7 +1807,7 @@ def test_b_nonconst_does_not_triggers_ger(self): def test_outer(self): rng = np.random.default_rng(unittest_tools.fetch_seed()) f = self.function([self.x, self.y], outer(self.x, self.y)) - self.assertFunctionContains(f, self.ger_destructive) + self.assertFunctionContains(f, self.ger_inplace) f( rng.random(5).astype(self.dtype), rng.random(4).astype(self.dtype), @@ -1847,7 +1852,7 @@ def test_scaled_A_plus_scaled_outer(self): np.asarray(0.2, self.dtype) * self.A + np.asarray(0.1, self.dtype) * outer(self.x, self.y), ) - # Why gemm? This make the graph simpler did we test that it + # Why Gemm? This make the graph simpler did we test that it # make it faster? self.assertFunctionContains(f, self.gemm) f( @@ -1861,7 +1866,7 @@ def test_scaled_A_plus_scaled_outer(self): rng.random(4).astype(self.dtype), ).shape == (5, 4) - def given_dtype(self, dtype, M, N, *, destructive=True): + def given_dtype(self, dtype, M, N, *, inplace=True): # test corner case shape and dtype rng = np.random.default_rng(unittest_tools.fetch_seed()) @@ -1870,9 +1875,7 @@ def given_dtype(self, dtype, M, N, *, destructive=True): y = tensor(dtype=dtype, shape=(None,)) f = self.function([A, x, y], A + 0.1 * outer(x, y)) - self.assertFunctionContains( - f, self.ger_destructive if destructive else self.ger - ) + self.assertFunctionContains(f, self.ger_inplace if inplace else self.ger) f( rng.random((M, N)).astype(dtype), rng.random(M).astype(dtype), @@ -1885,28 +1888,28 @@ def given_dtype(self, dtype, M, N, *, destructive=True): ).shape == (5, 4) def test_f32_0_0(self): - return self.given_dtype("float32", 0, 0, destructive=config.floatX != "float32") + return self.given_dtype("float32", 0, 0, inplace=config.floatX != "float32") def test_f32_1_0(self): - return self.given_dtype("float32", 1, 0, destructive=config.floatX != "float32") + return self.given_dtype("float32", 1, 0, inplace=config.floatX != "float32") def test_f32_0_1(self): - return self.given_dtype("float32", 0, 1, destructive=config.floatX != "float32") + return self.given_dtype("float32", 0, 1, inplace=config.floatX != "float32") def test_f32_1_1(self): - return self.given_dtype("float32", 1, 1, destructive=config.floatX != "float32") + return self.given_dtype("float32", 1, 1, inplace=config.floatX != "float32") def test_f32_4_4(self): - return self.given_dtype("float32", 4, 4, destructive=config.floatX != "float32") + return self.given_dtype("float32", 4, 4, inplace=config.floatX != "float32") def test_f32_7_1(self): - return self.given_dtype("float32", 7, 1, destructive=config.floatX != "float32") + return self.given_dtype("float32", 7, 1, inplace=config.floatX != "float32") def test_f32_1_2(self): - return self.given_dtype("float32", 1, 2, destructive=config.floatX != "float32") + return self.given_dtype("float32", 1, 2, inplace=config.floatX != "float32") def test_f64_4_5(self): - return self.given_dtype("float64", 4, 5, destructive=False) + return self.given_dtype("float64", 4, 5, inplace=False) def test_c64_7_1(self): return self.given_dtype("complex64", 7, 1) @@ -1924,7 +1927,7 @@ def test_inplace(self): (A, A + pt.constant(0.1, dtype=self.dtype) * outer(self.x, self.y)) ], ) - self.assertFunctionContains(f, self.ger_destructive) + self.assertFunctionContains(f, self.ger_inplace) # TODO: Test something about the updated value of `A` f( rng.random(4).astype(self.dtype), @@ -2364,7 +2367,7 @@ def test_gemm(self): b = scalar("b") self._compile_and_check( [x, y, a, z, b], - [gemm(z, a, x, y, b)], + [Gemm(inplace=False)(z, a, x, y, b)], [ rng.random((2, 3)).astype(config.floatX), rng.random((3, 4)).astype(config.floatX), @@ -2384,7 +2387,7 @@ def test_gemm_broadcast(self): # Broadcast Z self._compile_and_check( [x, y, a, z, b], - [gemm(z, a, x, y, b)], + [Gemm(inplace=False)(z, a, x, y, b)], [ rng.random((2, 3)).astype(config.floatX), rng.random((3, 4)).astype(config.floatX), @@ -2398,7 +2401,7 @@ def test_gemm_broadcast(self): # Broadcast dot(X, Y) self._compile_and_check( [x, y, a, z, b], - [gemm(z, a, x, y, b)], + [Gemm(inplace=False)(z, a, x, y, b)], [ rng.random((1, 3)).astype(config.floatX), rng.random((3, 4)).astype(config.floatX), @@ -2417,7 +2420,7 @@ def test_gemv(self): b = scalar("b") self._compile_and_check( [y, a, A, x, b], - [gemv(y, a, A, x, b)], + [Gemv(inplace=False)(y, a, A, x, b)], [ rng.random((2,)).astype(config.floatX), np.asarray(0.5, dtype=config.floatX), @@ -2435,7 +2438,7 @@ def test_ger(self): a = scalar("a") self._compile_and_check( [A, a, x, y], - [ger(A, a, x, y)], + [Ger(inplace=False)(A, a, x, y)], [ rng.random((2, 3)).astype(config.floatX), np.asarray(0.5, dtype=config.floatX), @@ -2534,3 +2537,162 @@ def test_batched_dot_blas_flags(): [batched_dot_thunk] = fn.vm.thunks assert not hasattr(batched_dot_thunk, "cthunk") np.testing.assert_allclose(fn(x_test, y_test), x_test @ y_test) + + +@pytest.mark.parametrize( + "z_shape", [(1, 3), (4, 1), (1, 1)], ids=["row", "column", "scalar"] +) +def test_broadcast_accumulator_is_left_out_of_gemm(z_shape): + """BLAS accumulates into the output buffer, so an accumulator that only broadcasts + against the product would have to be materialized to its shape first.""" + mode = Mode(linker="py", optimizer="fast_run").excluding("BlasOpt") + alpha = pt.scalar("alpha", dtype="float64") + x = pt.matrix("x", dtype="float64") + y = pt.matrix("y", dtype="float64") + z = pt.tensor("z", shape=z_shape, dtype="float64") + + fn = pytensor.function([alpha, x, y, z], z + alpha * (x @ y), mode=mode) + assert not [ap for ap in fn.maker.fgraph.toposort() if isinstance(ap.op, Gemm)] + + rng = np.random.default_rng(8) + xv = rng.standard_normal((4, 5)) + yv = rng.standard_normal((5, 3)) + zv = rng.standard_normal(z_shape) + result = fn(2.5, xv, yv, zv) + np.testing.assert_allclose(result, zv + 2.5 * (xv @ yv)) + # A static shape contradicting the result would mean the accumulator reached an op + # whose output type is the accumulator's own. + declared = fn.maker.fgraph.outputs[0].type.shape + assert all( + length is None or length == actual + for length, actual in zip(declared, result.shape, strict=True) + ) + + +def test_broadcast_product_is_left_out_of_gemm(): + """`Gemm` computes into a buffer shaped like the product, so a product that only + broadcasts against the sum cannot stand in for it.""" + mode = Mode(linker="py", optimizer="fast_run").excluding("BlasOpt") + a = pt.tensor("a", shape=(1, 4), dtype="float64") + b = pt.tensor("b", shape=(4, 3), dtype="float64") + x = pt.tensor("x", shape=(None, 5), dtype="float64") + y = pt.tensor("y", shape=(5, 3), dtype="float64") + + fn = pytensor.function([a, b, x, y], (a @ b) + (x @ y), mode=mode) + assert not [ap for ap in fn.maker.fgraph.toposort() if isinstance(ap.op, Gemm)] + + rng = np.random.default_rng(9) + av = rng.standard_normal((1, 4)) + bv = rng.standard_normal((4, 3)) + xv = rng.standard_normal((6, 5)) + yv = rng.standard_normal((5, 3)) + result = fn(av, bv, xv, yv) + np.testing.assert_allclose(result, av @ bv + xv @ yv) + declared = fn.maker.fgraph.outputs[0].type.shape + assert all( + length is None or length == actual + for length, actual in zip(declared, result.shape, strict=True) + ) + + +class TestGeneralBlasLowering: + """`Gemm`, `Gemv` and `Ger` built outside `blas_optdb`. + + These rewrites run in `specialize` and match `Dot` directly rather than `Dot22`, so + they are the only BLAS lowering the backends that exclude `blas_optdb` ever see. + """ + + mode = Mode(linker="py", optimizer="fast_run").excluding("BlasOpt") + + @staticmethod + def _apply_of(fn, op_type): + return [ap for ap in fn.maker.fgraph.toposort() if isinstance(ap.op, op_type)] + + @staticmethod + def _inputs(rng, *shapes): + return [rng.standard_normal(shape) for shape in shapes] + + def test_matrix_product_to_gemm(self): + alpha = pt.scalar("alpha", dtype="float64") + x = pt.matrix("x", dtype="float64") + y = pt.matrix("y", dtype="float64") + c = pt.matrix("c", dtype="float64") + + fn = pytensor.function([alpha, x, y, c], c + alpha * (x @ y), mode=self.mode) + assert len(self._apply_of(fn, Gemm)) == 1 + # The sum is carried by Gemm's own beta rather than left as a separate add. + adds = [ + ap + for ap in self._apply_of(fn, Elemwise) + if isinstance(ap.op.scalar_op, ps.Add) + ] + assert not adds + + rng = np.random.default_rng(2) + xv, yv, cv = self._inputs(rng, (4, 5), (5, 3), (4, 3)) + np.testing.assert_allclose(fn(2.5, xv, yv, cv), cv + 2.5 * (xv @ yv)) + + def test_outer_product_to_ger(self): + alpha = pt.scalar("alpha", dtype="float64") + u = pt.vector("u", dtype="float64") + v = pt.vector("v", dtype="float64") + c = pt.matrix("c", dtype="float64") + + fn = pytensor.function( + [alpha, u, v, c], c + alpha * pt.outer(u, v), mode=self.mode + ) + assert len(self._apply_of(fn, Ger)) == 1 + + rng = np.random.default_rng(3) + uv, vv, cv = self._inputs(rng, 4, 3, (4, 3)) + np.testing.assert_allclose(fn(2.5, uv, vv, cv), cv + 2.5 * np.outer(uv, vv)) + + @pytest.mark.parametrize("build", [lambda x, v: x @ v, dot], ids=["matmul", "dot"]) + def test_matrix_vector_to_gemv(self, build): + """``@`` promotes the vector and squeezes the result; `dot` builds the rank-1 + operand directly. Both reach `Gemv` through `_as_matrix_product`.""" + alpha = pt.scalar("alpha", dtype="float64") + x = pt.matrix("x", dtype="float64") + v = pt.vector("v", dtype="float64") + w = pt.vector("w", dtype="float64") + + fn = pytensor.function( + [alpha, x, v, w], w + alpha * build(x, v), mode=self.mode + ) + assert len(self._apply_of(fn, Gemv)) == 1 + + rng = np.random.default_rng(4) + xv, vv, wv = self._inputs(rng, (4, 5), 5, 4) + np.testing.assert_allclose(fn(2.5, xv, vv, wv), wv + 2.5 * (xv @ vv)) + + def test_vector_matrix_to_gemv(self): + alpha = pt.scalar("alpha", dtype="float64") + v = pt.vector("v", dtype="float64") + y = pt.matrix("y", dtype="float64") + w = pt.vector("w", dtype="float64") + + fn = pytensor.function([alpha, v, y, w], w + alpha * (v @ y), mode=self.mode) + assert len(self._apply_of(fn, Gemv)) == 1 + + rng = np.random.default_rng(5) + vv, yv, wv = self._inputs(rng, 5, (5, 3), 3) + np.testing.assert_allclose(fn(2.5, vv, yv, wv), wv + 2.5 * (vv @ yv)) + + def test_shared_product_is_left_alone(self): + """Folding a product read twice would compute it once for each consumer.""" + alpha = pt.scalar("alpha", dtype="float64") + x = pt.matrix("x", dtype="float64") + y = pt.matrix("y", dtype="float64") + c = pt.matrix("c", dtype="float64") + + product = x @ y + fn = pytensor.function( + [alpha, x, y, c], [c + alpha * product, product], mode=self.mode + ) + assert not self._apply_of(fn, Gemm) + + rng = np.random.default_rng(7) + xv, yv, cv = self._inputs(rng, (4, 5), (5, 3), (4, 3)) + summed, plain = fn(2.5, xv, yv, cv) + np.testing.assert_allclose(plain, xv @ yv) + np.testing.assert_allclose(summed, cv + 2.5 * (xv @ yv)) diff --git a/tests/tensor/test_blas_c.py b/tests/tensor/test_blas_c.py index 5b3da5cff9..f6cda0b2b2 100644 --- a/tests/tensor/test_blas_c.py +++ b/tests/tensor/test_blas_c.py @@ -1,14 +1,13 @@ -from warnings import warn - import numpy as np import pytest import pytensor import pytensor.tensor as pt from pytensor.compile import get_mode +from pytensor.graph.utils import MethodNotDefined +from pytensor.link.c.dispatch.basic import c_funcify from pytensor.tensor.basic import AllocEmpty -from pytensor.tensor.blas import Ger -from pytensor.tensor.blas.blas_c import CGemv, CGer, must_initialize_y_gemv +from pytensor.tensor.blas import Gemv, Ger from pytensor.tensor.type import ( dmatrix, dscalar, @@ -38,7 +37,7 @@ def skip_if_blas_ldflags_empty(*functions_detected): ) -class TestCGer(OptimizationTestMixin): +class TestGerC(OptimizationTestMixin): def setup_method(self): self.manual_setup_method() @@ -69,46 +68,34 @@ def run_f(self, f): def b(self, bval): return pt.as_tensor_variable(np.asarray(bval, dtype=self.dtype)) - def test_eq(self): - assert CGer(True) == CGer(True) - assert CGer(False) == CGer(False) - assert CGer(False) != CGer(True) - - assert CGer(True) != Ger(True) - assert CGer(False) != Ger(False) - - # assert that eq works for non-CGer instances - assert CGer(False) is not None - assert CGer(True) is not None - - def test_hash(self): - assert hash(CGer(True)) == hash(CGer(True)) - assert hash(CGer(False)) == hash(CGer(False)) - assert hash(CGer(False)) != hash(CGer(True)) - def test_optimization_pipeline(self): skip_if_blas_ldflags_empty() f = self.function([self.x, self.y], pt.outer(self.x, self.y)) - self.assertFunctionContains(f, CGer(destructive=True)) + self.assertFunctionContains(f, Ger(inplace=True)) f(self.xval, self.yval) # DebugMode tests correctness def test_optimization_pipeline_float(self): skip_if_blas_ldflags_empty() self.manual_setup_method("float32") f = self.function([self.x, self.y], pt.outer(self.x, self.y)) - self.assertFunctionContains(f, CGer(destructive=True)) + self.assertFunctionContains(f, Ger(inplace=True)) f(self.xval, self.yval) # DebugMode tests correctness - def test_int_fails(self): + def test_integer_outer_product_is_not_routed_to_ger(self): + # `Ger.make_node` rejects integer dtypes, so a rewrite that routed this to `Ger` + # would raise while rewriting rather than fall back. self.manual_setup_method("int32") f = self.function([self.x, self.y], pt.outer(self.x, self.y)) - self.assertFunctionContains0(f, CGer(destructive=True)) - self.assertFunctionContains0(f, CGer(destructive=False)) + self.assertFunctionContains0(f, Ger(inplace=True)) + self.assertFunctionContains0(f, Ger(inplace=False)) + np.testing.assert_array_equal( + f(self.xval, self.yval), np.outer(self.xval, self.yval) + ) def test_A_plus_outer(self): skip_if_blas_ldflags_empty() f = self.function([self.A, self.x, self.y], self.A + pt.outer(self.x, self.y)) - self.assertFunctionContains(f, CGer(destructive=False)) + self.assertFunctionContains(f, Ger(inplace=False)) self.run_f(f) # DebugMode tests correctness def test_A_plus_scaled_outer(self): @@ -116,16 +103,16 @@ def test_A_plus_scaled_outer(self): f = self.function( [self.A, self.x, self.y], self.A + 0.1 * pt.outer(self.x, self.y) ) - self.assertFunctionContains(f, CGer(destructive=False)) + self.assertFunctionContains(f, Ger(inplace=False)) self.run_f(f) # DebugMode tests correctness -class TestCGemv(OptimizationTestMixin): +class TestGemvC(OptimizationTestMixin): """ - Tests of CGemv specifically. + Tests of the C implementation of `Gemv`. Generic tests of Gemv-compatibility, including both dtypes are - done below in TestCGemvFloat32 and TestCGemvFloat64 + done below in TestGemvCFloat32 and TestGemvCFloat64 """ def setup_method(self): @@ -156,7 +143,7 @@ def test_nan_beta_0(self, inplace): mode=mode, ) [node] = f.maker.fgraph.apply_nodes - assert isinstance(node.op, CGemv) and node.op.inplace == inplace + assert isinstance(node.op, Gemv) and node.op.inplace == inplace for rows in (3, 1, 0): for cols in (1, 0): Aval = np.ones((rows, cols), dtype=self.dtype) @@ -172,7 +159,7 @@ def test_optimizations_vm(self): # Assert that the dot was optimized somehow self.assertFunctionContains0(f, pt.dot) - self.assertFunctionContains1(f, CGemv(inplace=True)) + self.assertFunctionContains1(f, Gemv(inplace=True)) # Assert they produce the same output assert np.allclose(f(self.xval, self.Aval), np.dot(self.xval, self.Aval)) @@ -190,7 +177,7 @@ def test_optimizations_mv(self): # Assert that the dot was optimized somehow self.assertFunctionContains0(f, pt.dot) - self.assertFunctionContains1(f, CGemv(inplace=True)) + self.assertFunctionContains1(f, Gemv(inplace=True)) # Assert they produce the same output assert np.allclose(f(self.Aval, self.yval), np.dot(self.Aval, self.yval)) @@ -200,16 +187,6 @@ def test_optimizations_mv(self): np.dot(self.Aval[::-1, ::-1], self.yval), ) - def test_must_initialize_y_gemv(self): - if must_initialize_y_gemv(): - # FIME: This warn should be emitted by the function if we find it relevant - # Not in a test that doesn't care about the outcome either way - warn( - "WARNING: The current BLAS requires PyTensor to initialize" - " memory for some GEMV calls which will result in a minor" - " degradation in performance for such calls." - ) - def t_gemv1(self, m_shp): """test vector2 + dot(matrix, vector1)""" rng = np.random.default_rng(unittest_tools.fetch_seed()) @@ -223,7 +200,7 @@ def t_gemv1(self, m_shp): # Assert they produce the same output assert np.allclose(f(), np.dot(m.get_value(), v1.get_value()) + v2_orig) topo = [n.op for n in f.maker.fgraph.toposort()] - assert topo == [CGemv(inplace=False)], topo + assert topo == [Gemv(inplace=False)], topo # test the inplace version g = pytensor.function( @@ -236,7 +213,7 @@ def t_gemv1(self, m_shp): v2.get_value(), np.dot(m.get_value(), v1.get_value()) + v2_orig ) topo = [n.op for n in g.maker.fgraph.toposort()] - assert topo == [CGemv(inplace=True)] + assert topo == [Gemv(inplace=True)] # Do the same tests with a matrix with strides in both dimensions m.set_value(m.get_value(borrow=True)[::-1, ::-1], borrow=True) @@ -308,7 +285,7 @@ def test_empty_A(self): y = dvector("y") alpha = 1.0 beta = dscalar("beta") - gemv = CGemv(inplace=True)(y, alpha, A, x, beta) + gemv = Gemv(inplace=True)(y, alpha, A, x, beta) fn = pytensor.function( [A, x, y, beta], gemv, @@ -323,29 +300,29 @@ def test_empty_A(self): np.testing.assert_allclose(out, expected) -class TestCGemvFloat32(BaseGemv, OptimizationTestMixin): +class TestGemvCFloat32(BaseGemv, OptimizationTestMixin): mode = mode_blas_opt dtype = "float32" - gemv = CGemv(inplace=False) - gemv_inplace = CGemv(inplace=True) + gemv = Gemv(inplace=False) + gemv_inplace = Gemv(inplace=True) def setup_method(self): skip_if_blas_ldflags_empty() -class TestCGemvFloat64(BaseGemv, OptimizationTestMixin): +class TestGemvCFloat64(BaseGemv, OptimizationTestMixin): mode = mode_blas_opt dtype = "float64" - gemv = CGemv(inplace=False) - gemv_inplace = CGemv(inplace=True) + gemv = Gemv(inplace=False) + gemv_inplace = Gemv(inplace=True) def setup_method(self): skip_if_blas_ldflags_empty() -class TestCGemvNoFlags: +class TestGemvCNoFlags: mode = mode_blas_opt - gemv = CGemv(inplace=False) + gemv = Gemv(inplace=False) M = 4 N = 5 slice_step = 3 @@ -407,14 +384,14 @@ def compute_ref(self, alpha, A, x, beta, y, transpose_A, slice_tensors): return ref_val @pytensor.config.change_flags(blas__ldflags="") - def run_cgemv(self, dtype, ALPHA, BETA, transpose_A, slice_tensors): + def run_gemv(self, dtype, ALPHA, BETA, transpose_A, slice_tensors): f = self.get_function( dtype, transpose_A=transpose_A, slice_tensors=slice_tensors ) values = self.get_data( dtype, ALPHA, BETA, transpose_A=transpose_A, slice_tensors=slice_tensors ) - assert any(isinstance(node.op, CGemv) for node in f.maker.fgraph.apply_nodes) + assert any(isinstance(node.op, Gemv) for node in f.maker.fgraph.apply_nodes) z_val = f(*values) assert z_val.dtype == dtype assert z_val.ndim == 1 @@ -422,13 +399,13 @@ def run_cgemv(self, dtype, ALPHA, BETA, transpose_A, slice_tensors): ref_val = self.compute_ref(*((*values, transpose_A, slice_tensors))) unittest_tools.assert_allclose(ref_val, z_val) - def test_cgemv(self): + def test_gemv(self): for dtype in ("float32", "float64"): for alpha in (0, 1, -2): for beta in (0, 1, -2): for transpose_A in (False, True): for slice_tensors in (False, True): - self.run_cgemv( + self.run_gemv( dtype, alpha, beta, @@ -437,9 +414,53 @@ def test_cgemv(self): ) -class TestSdotNoFlags(TestCGemvNoFlags): +class TestSdotNoFlags(TestGemvCNoFlags): M = 1 class TestBlasStridesC(TestBlasStrides): mode = mode_blas_opt + + +class TestGerCNoFlags: + """`Ger` has no C implementation to fall back on when BLAS is not linked. + + The bundled no-BLAS header supplies ``[sd]gemm_`` and ``[sd]gemv_`` but no ``[sd]ger_``, + so the C implementation declines the node and the Python ``perform`` runs instead. + """ + + @pytensor.config.change_flags(blas__ldflags="") + def test_c_implementation_declines(self): + A = matrix(dtype="float64") + x = vector(dtype="float64") + y = vector(dtype="float64") + alpha = scalar(dtype="float64") + node = Ger(inplace=False)(A, alpha, x, y).owner + + impl = c_funcify(node.op, node=node) + # Emitting the code anyway would reach the linker and fail there instead. + with pytest.raises(MethodNotDefined): + impl.c_code( + node, "n", ["A", "a", "x", "y"], ["Z"], {"fail": "", "params": "p"} + ) + + @pytensor.config.change_flags(blas__ldflags="") + @pytest.mark.parametrize("dtype", ["float32", "float64"]) + def test_ger_without_blas(self, dtype): + A = matrix(dtype=dtype) + x = vector(dtype=dtype) + y = vector(dtype=dtype) + alpha = scalar(dtype=dtype) + + f = pytensor.function( + [A, alpha, x, y], Ger(inplace=False)(A, alpha, x, y), mode=mode_blas_opt + ) + + rng = np.random.default_rng(sum(map(ord, f"ger_without_blas {dtype}"))) + A_val = rng.random((4, 5)).astype(dtype) + x_val = rng.random(4).astype(dtype) + y_val = rng.random(5).astype(dtype) + + out = f(A_val, np.asarray(2.0, dtype=dtype), x_val, y_val) + assert out.dtype == dtype + unittest_tools.assert_allclose(A_val + 2.0 * np.outer(x_val, y_val), out) diff --git a/tests/tensor/test_math.py b/tests/tensor/test_math.py index f3f7257b3a..c31df77759 100644 --- a/tests/tensor/test_math.py +++ b/tests/tensor/test_math.py @@ -2840,7 +2840,7 @@ def test_Dot(self): [advec, bdvec], [dot(advec, bdvec)], [advec_val, bdvec_val], - (Dot, blas.Dot22, blas.Gemv, blas.CGemv), + (Dot, blas.Dot22, blas.Gemv), ) # mat/mat @@ -2861,7 +2861,7 @@ def test_Dot(self): [advec, bdmat], [dot(advec, bdmat)], [advec_val, bdmat_val], - (Dot, blas.Dot22, blas.Gemv, blas.CGemv), + (Dot, blas.Dot22, blas.Gemv), ) # mat/vec @@ -2870,7 +2870,7 @@ def test_Dot(self): [admat, bdvec], [dot(admat, bdvec)], [admat_val, bdvec_val], - (Dot, blas.Dot22, blas.Gemv, blas.CGemv), + (Dot, blas.Dot22, blas.Gemv), ) diff --git a/tests/tensor/test_sharedvar.py b/tests/tensor/test_sharedvar.py index df506223c5..e96207c587 100644 --- a/tests/tensor/test_sharedvar.py +++ b/tests/tensor/test_sharedvar.py @@ -527,11 +527,11 @@ def test_specify_shape_inplace(self): == 1 ) assert all( - node.op == pytensor.tensor.blas.gemm_inplace + node.op == pytensor.tensor.blas.Gemm(inplace=True) for node in topo if isinstance(node.op, pytensor.tensor.blas.Gemm) ) - # Their is no inplace gemm for sparse + # There is no inplace Gemm for sparse # assert all(node.op.inplace for node in topo if node.op.__class__.__name__ == "StructuredDot") s_shared_specify = specify_shape( s_shared, s_shared.get_value(borrow=True).shape @@ -561,7 +561,7 @@ def test_specify_shape_inplace(self): == 1 ) assert all( - node.op == pytensor.tensor.blas.gemm_inplace + node.op == pytensor.tensor.blas.Gemm(inplace=True) for node in topo if isinstance(node.op, pytensor.tensor.blas.Gemm) ) @@ -593,7 +593,7 @@ def test_specify_shape_inplace(self): == 1 ) assert all( - node.op == pytensor.tensor.blas.gemm_inplace + node.op == pytensor.tensor.blas.Gemm(inplace=True) for node in topo if isinstance(node.op, pytensor.tensor.blas.Gemm) ) diff --git a/tests/test_printing.py b/tests/test_printing.py index 55476ecb34..2804868099 100644 --- a/tests/test_printing.py +++ b/tests/test_printing.py @@ -307,7 +307,7 @@ def test_debugprint(): print_view_map=True, ) s = s.getvalue() - Gemv_op_name = "CGemv" if pytensor.config.blas__ldflags else "Gemv" + Gemv_op_name = "Gemv" exp_res = dedent( r""" Composite{(i0 + (i1 - i2))} 4