diff --git a/doc/extending/creating_a_c_op.rst b/doc/extending/creating_a_c_op.rst index 12105faa8d..33a80ae27b 100644 --- a/doc/extending/creating_a_c_op.rst +++ b/doc/extending/creating_a_c_op.rst @@ -713,322 +713,6 @@ C code. return c_code % locals() -Alternate way of defining C :class:`Op`\s -========================================= - -The two previous examples have covered the standard way of implementing C :class:`Op`\s -in PyTensor by inheriting from the class :class:`Op`. This process is mostly -simple but it still involves defining many methods as well as mixing, in the -same file, both Python and C code which tends to make the result less -readable. - -To help with this, PyTensor defines a class, `ExternalCOp`, from which new C :class:`Op`\s -can inherit. The class `ExternalCOp` aims to simplify the process of implementing -C :class:`Op`\s by doing the following : - -* It allows you to define the C implementation of your :class:`Op` in a distinct - C code file. This makes it easier to keep your Python and C code - readable and well indented. - -* It can automatically handle all the methods that return C code, - in addition to :meth:`Op.c_code_cache_version` based on the - provided external C implementation. - -To illustrate how much simpler the class `ExternalCOp` makes the process of defining -a new :class:`Op` with a C implementation, let's revisit the second example of this -tutorial, the `VectorTimesVector`\ :class:`Op`. In that example, we implemented an :class:`Op` -to perform the task of element-wise vector-vector multiplication. The two -following blocks of code illustrate what the :class:`Op` would look like if it was -implemented using the `ExternalCOp` class. - -The new :class:`Op` is defined inside a Python file with the following code : - -.. testcode:: - - import pytensor - from pytensor.link.c.op import ExternalCOp - - class VectorTimesVector(ExternalCOp): - __props__ = () - - func_file = "./vectorTimesVector.c" - func_name = "APPLY_SPECIFIC(vector_times_vector)" - - def __init__(self): - super().__init__(self.func_file, self.func_name) - - def make_node(self, x, y): - # Validate the inputs' type - if x.type.ndim != 1: - raise TypeError('x must be a 1-d vector') - if y.type.ndim != 1: - raise TypeError('y must be a 1-d vector') - - # Create an output variable of the same type as x - output_var = pytensor.tensor.type.TensorType( - dtype=pytensor.scalar.upcast(x.dtype, y.dtype), - shape=(None,))() - - return Apply(self, [x, y], [output_var]) - -And the following is the C implementation of the :class:`Op`, defined in an external -C file named ``vectorTimesVector.c``: - -.. code-block:: c - - #section support_code - - // Support code function - bool vector_same_shape(PyArrayObject* arr1, PyArrayObject* arr2) - { - return (PyArray_DIMS(arr1)[0] == PyArray_DIMS(arr2)[0]); - } - - - #section support_code_apply - - // Apply-specific support function - void APPLY_SPECIFIC(vector_elemwise_mult)( - DTYPE_INPUT_0* x_ptr, int x_str, - DTYPE_INPUT_1* y_ptr, int y_str, - DTYPE_OUTPUT_0* z_ptr, int z_str, int nbElements) - { - for (int i=0; i < nbElements; i++){ - z_ptr[i * z_str] = x_ptr[i * x_str] * y_ptr[i * y_str]; - } - } - - // Apply-specific main function - int APPLY_SPECIFIC(vector_times_vector)(PyArrayObject* input0, - PyArrayObject* input1, - PyArrayObject** output0) - { - // Validate that the inputs have the same shape - if ( !vector_same_shape(input0, input1)) - { - PyErr_Format(PyExc_ValueError, "Shape mismatch : " - "input0.shape[0] and input1.shape[0] should " - "match but x.shape[0] == %i and " - "y.shape[0] == %i", - PyArray_DIMS(input0)[0], PyArray_DIMS(input1)[0]); - return 1; - } - - // Validate that the output storage exists and has the same - // dimension as x. - if (NULL == *output0 || !(vector_same_shape(input0, *output0))) - { - /* Reference received to invalid output variable. - Decrease received reference's ref count and allocate new - output variable */ - Py_XDECREF(*output0); - *output0 = (PyArrayObject*)PyArray_EMPTY(1, - PyArray_DIMS(input0), - TYPENUM_OUTPUT_0, - 0); - - if (!*output0) { - PyErr_Format(PyExc_ValueError, - "Could not allocate output storage"); - return 1; - } - } - - // Perform the actual vector-vector multiplication - APPLY_SPECIFIC(vector_elemwise_mult)( - (DTYPE_INPUT_0*)PyArray_DATA(input0), - PyArray_STRIDES(input0)[0] / ITEMSIZE_INPUT_0, - (DTYPE_INPUT_1*)PyArray_DATA(input1), - PyArray_STRIDES(input1)[0] / ITEMSIZE_INPUT_1, - (DTYPE_OUTPUT_0*)PyArray_DATA(*output0), - PyArray_STRIDES(*output0)[0] / ITEMSIZE_OUTPUT_0, - PyArray_DIMS(input0)[0]); - - return 0; - } - -As you can see from this example, the Python and C implementations are nicely -decoupled which makes them much more readable than when they were intertwined -in the same file and the C code contained string formatting markers. - -Now that we have motivated the `ExternalCOp` class, we can have a more precise look at -what it does for us. For this, we go through the various elements that make up -this new version of the `VectorTimesVector`\ `Op` : - -* Parent class : instead of inheriting from the class :class:`Op`, - VectorTimesVector inherits from the class `ExternalCOp`. - -* Constructor : in our new `COp`, the :meth:`COp.__init__` method has an - important use; to inform the constructor of the `ExternalCOp` class - of the location, on the filesystem of the C implementation of - this `COp`. To do this, it gives a list of file paths containing - the C code for this `COp`. To auto-generate the c_code method - with a function call you can specify the function name as the - second parameter. The paths should be given as a relative - path from the folder where the descendant of the `ExternalCOp` class - is defined. - -* :meth:`ExternalCOp.make_node` : this method is absolutely - identical to the one in our old example. Using the `ExternalCOp` - class doesn't change anything here. - -* External C code : the external C code implements the various - functions associated with the `COp`. Writing this C code - involves a few subtleties which deserve their own respective - sections. - -Main function -------------- - -If you pass a function name to :meth:`ExternalCOp.__init___`, it must respect -the following constraints: - -* It must return an int. The value of that int indicates whether - the `Op` could perform its task or not. A value of 0 indicates - success while any non-zero value will interrupt the execution - of the PyTensor function. When returning non-zero the function - must set a python exception indicating the details of the - problem. - -* It must receive one argument for each input to the `Op` followed - by one pointer to an argument for each output of the `Op`. The - types for the argument is dependent on the Types (that is - pytensor Types) of your inputs and outputs. - -* You can specify the number of inputs and outputs for your `Op` - by setting the ``_cop_num_inputs`` and ``_cop_num_outputs`` - attributes on your `COp`. The main function will always be - called with that number of arguments, using NULL to fill in - for missing values at the end. This can be used if your `COp` - has a variable number of inputs or outputs, but with a fixed - maximum. - -For example, the main C function of an `COp` that takes two TensorTypes -(which has ``PyArrayObject *`` as its C type) as inputs and returns -both their sum and the difference between them would have four -parameters (two for the `COp`'s inputs and two for its outputs) and it's -signature would look something like this : - -.. code-block:: c - - int sumAndDiffOfScalars(PyArrayObject* in0, PyArrayObject* in1, - PyArrayObject** out0, PyArrayObject** out1) - -Macros ------- - -For certain section tags, your C code can benefit from a number of -pre-defined macros. These section tags have no macros: ``init_code``, -``support_code``. All other tags will have the support macros -discussed below. - -* ``APPLY_SPECIFIC(str)`` which will automatically append a name - unique to the :ref:`apply` node that applies the `Op` at the end - of the provided ``str``. The use of this macro is discussed - further below. - -For every input which has a :attr:`dtype` attribute (this means -Tensors), the following macros will be -defined unless your `Op` class has an :attr:`Op.check_input` attribute -defined to False. In these descriptions 'i' refers to the position -(indexed from 0) in the input array. - -* ``DTYPE_INPUT_{i}`` : NumPy dtype of the data in the array. - This is the variable type corresponding to the NumPy dtype, not the - string representation of the NumPy dtype. For instance, if the `Op`'s - first input is a float32 :class:`ndarray`, then the macro ``DTYPE_INPUT_0`` - corresponds to ``npy_float32`` and can directly be used to declare a - new variable of the same dtype as the data in the array : - - .. code-block:: c - - DTYPE_INPUT_0 myVar = someValue; - -* ``TYPENUM_INPUT_{i}`` : Typenum of the data in the array - -* ``ITEMSIZE_INPUT_{i}`` : Size, in bytes, of the elements in - the array. - -In the same way, the macros ``DTYPE_OUTPUT_{i}``, -``ITEMSIZE_OUTPUT_{i}`` and ``TYPENUM_OUTPUT_{i}`` are defined for -every output 'i' of the `Op`. - -In addition to these macros, the ``init_code_struct``, ``code``, and -``code_cleanup`` section tags also have the following macros: - -* ``FAIL`` : Code to insert at error points. A python exception - should be set prior to this code. An invocation look like this: - - .. code-block:: c - - if (error) { - // Set python exception - FAIL - } - - You can add a semicolon after the macro if it makes your editor - happy. - -* ``PARAMS`` : Name of the params variable for this node. (only - for `Op`\s which have params, which is discussed elsewhere) - -Finally the tag ``code`` and ``code_cleanup`` have macros to -pass the inputs and output names. These are name ``INPUT_{i}`` and -``OUTPUT_{i}`` where `i` is the 0-based index position in the input -and output arrays respectively. - -Support code ------------- - -Certain section are limited in what you can place in them due to -semantic and syntactic restrictions of the C++ language. Most of -these restrictions apply to the tags that end in ``_struct``. - -When we defined the ``VectorTimesVector`` `Op` without using the ``ExternalCOp`` -class, we had to make a distinction between two types of support_code -: the support code that was apply-specific and the support code that -wasn't. The apply-specific code was defined in the -``c_support_code_apply`` method and the elements defined in that -code (global variables and functions) had to include the name of the -Apply node in their own names to avoid conflicts between the different -versions of the apply-specific code. The code that wasn't -apply-specific was simply defined in the ``c_support_code`` method. - -To make indentifiers that include the :ref:`apply` node name use the -``APPLY_SPECIFIC(str)`` macro. In the above example, this macro is -used when defining the functions ``vector_elemwise_mult`` and -``vector_times_vector`` as well as when calling function -``vector_elemwise_mult`` from inside ``vector_times_vector``. - -When using the ``ExternalCOp`` class, we still have to make the distinction -between C code for each of the methods of a C class. These sections of -code are separated by ``#section `` markers. The tag determines -the name of the method this C code applies to with the rule that -```` applies to `c_`. Unknown tags are an error and will be -reported. Duplicate tags will be merged together in the order the -appear in the C files. - -The rules for knowing if where a piece of code should be put can be -sometimes tricky. The key thing to remember is that things that can -be shared between instances of the `Op` should be apply-agnostic and go -into a section which does not end in ``_apply`` or ``_struct``. The -distinction of ``_apply`` and ``_struct`` mostly hinghes on how you -want to manage the lifetime of the object. Note that to use an -apply-specific object, you have to be in a apply-specific section, so -some portions of the code that might seem apply-agnostic may still be -apply-specific because of the data they use (this does not include -arguments). - -In the above example, the ``function vector_same_shape`` is -apply-agnostic because it uses none of the macros defined by the class -``ExternalCOp`` and it doesn't rely on any apply-specific code. The function -``vector_elemwise_mult`` is apply-specific because it uses the -macros defined by ``ExternalCOp``. Finally, the function -``vector_times_vector`` is apply-specific because it uses those same -macros and also because it calls ``vector_elemwise_mult`` which is -an apply-specific function. - - Using GDB to debug :class:`COp`'s C code ======================================== diff --git a/pytensor/link/c/op.py b/pytensor/link/c/op.py index b67fb9048e..01804b21a2 100644 --- a/pytensor/link/c/op.py +++ b/pytensor/link/c/op.py @@ -256,6 +256,10 @@ def get_io_macros(inputs: list[str], outputs: list[str]) -> tuple[str, str]: class ExternalCOp(COp): """Class for an `Op` with an external C implementation. + .. deprecated:: + Inherit from `COp` and return the C code directly from `c_code` and + related methods instead. + One can inherit from this class, provide its constructor with a path to an external C source file and the name of a function within it, and define an `Op` for said function. @@ -306,6 +310,13 @@ def __init__( files overriding sections in previous files. """ + warnings.warn( + "ExternalCOp is deprecated and will be removed in a future release. " + "Inherit from COp and return the C code directly from `c_code` and " + "related methods instead.", + FutureWarning, + stacklevel=2, + ) if not isinstance(func_files, list): self.func_files = [Path(func_files)] else: @@ -633,16 +644,3 @@ class _NoPythonCOp(COp): def perform(self, node, inputs, output_storage): raise NotImplementedError("No Python implementation is provided by this COp.") - - -class _NoPythonExternalCOp(ExternalCOp): - """A class used to indicate that an `ExternalCOp` does not provide a Python implementation. - - XXX: Do not use this class; it's only for tracking bad implementations internally. - - """ - - def perform(self, node, inputs, output_storage): - raise NotImplementedError( - "No Python implementation is provided by this ExternalCOp." - ) diff --git a/pytensor/tensor/c_code/dimshuffle.c b/pytensor/tensor/c_code/dimshuffle.c deleted file mode 100644 index 0bfc5df3bb..0000000000 --- a/pytensor/tensor/c_code/dimshuffle.c +++ /dev/null @@ -1,86 +0,0 @@ -#section support_code_apply - -int APPLY_SPECIFIC(cpu_dimshuffle)(PyArrayObject *input, PyArrayObject **res, PARAMS_TYPE *params) { - npy_int64* new_order; - npy_intp nd_in; - npy_intp nd_out; - npy_intp* dimensions; - npy_intp* strides; - - if (!PyArray_IS_C_CONTIGUOUS(params->_new_order)) { - PyErr_SetString(PyExc_RuntimeError, "DimShuffle: param _new_order must be C-contiguous."); - return 1; - } - new_order = (npy_int64*) PyArray_DATA(params->_new_order); - nd_in = (npy_intp)(params->input_ndim); - nd_out = PyArray_SIZE(params->_new_order); - - if (PyArray_NDIM(input) != nd_in) { - PyErr_SetString(PyExc_ValueError, "DimShuffle: Input has less dimensions than expected."); - return 1; - } - - // Compute new dimensions and strides - dimensions = (npy_intp*) malloc(nd_out * sizeof(npy_intp)); - strides = (npy_intp*) malloc(nd_out * sizeof(npy_intp)); - if (dimensions == NULL || strides == NULL) { - PyErr_NoMemory(); - free(dimensions); - free(strides); - return 1; - }; - - npy_intp original_size = PyArray_SIZE(input); - npy_intp new_size = 1; - for (npy_intp i = 0; i < nd_out; ++i) { - // We set the strides of length 1 dimensions to PyArray_ITEMSIZE(input). - // The value is arbitrary, because there is never a next element. - // np.expand_dims(x, 0) and x[None] do different things here. - // I would prefer zero, but there are some poorly implemented BLAS operations - // That don't handle zero strides correctly. At least they won't fail because of DimShuffle. - if (new_order[i] != -1) { - dimensions[i] = PyArray_DIMS(input)[new_order[i]]; - strides[i] = PyArray_DIMS(input)[new_order[i]] == 1 ? PyArray_ITEMSIZE(input) : PyArray_STRIDES(input)[new_order[i]]; - } else { - dimensions[i] = 1; - strides[i] = PyArray_ITEMSIZE(input); - } - new_size *= dimensions[i]; - } - - if (original_size != new_size) { - PyErr_SetString(PyExc_ValueError, "DimShuffle: Attempting to squeeze axes with size not equal to one."); - free(dimensions); - free(strides); - return 1; - } - - if (*res) - Py_XDECREF(*res); - - // Create the new array. - *res = (PyArrayObject*)PyArray_New(&PyArray_Type, nd_out, dimensions, - PyArray_TYPE(input), strides, - PyArray_DATA(input), PyArray_ITEMSIZE(input), - // borrow only the writable flag from the base - // the NPY_OWNDATA flag will default to 0. - (NPY_ARRAY_WRITEABLE * PyArray_ISWRITEABLE(input)), - NULL); - - if (*res == NULL) { - free(dimensions); - free(strides); - return 1; - } - - // Declare it a view of the original input - Py_INCREF((PyObject*)input); - PyArray_SetBaseObject(*res, (PyObject*)input); - - // recalculate flags: CONTIGUOUS, FORTRAN, ALIGNED - PyArray_UpdateFlags(*res, NPY_ARRAY_UPDATE_ALL); - - free(strides); - free(dimensions); - return 0; -} \ No newline at end of file diff --git a/pytensor/tensor/elemwise.py b/pytensor/tensor/elemwise.py index 2a3de0a167..17cd0fba63 100644 --- a/pytensor/tensor/elemwise.py +++ b/pytensor/tensor/elemwise.py @@ -14,13 +14,12 @@ from pytensor.graph.replace import _vectorize_node, _vectorize_not_needed from pytensor.graph.utils import MethodNotDefined from pytensor.link.c.basic import failure_code -from pytensor.link.c.op import COp, ExternalCOp, OpenMPOp -from pytensor.link.c.params_type import ParamsType +from pytensor.link.c.op import COp, OpenMPOp from pytensor.misc.frozendict import frozendict from pytensor.printing import Printer, pprint from pytensor.scalar import get_scalar_type from pytensor.scalar.basic import identity as scalar_identity -from pytensor.scalar.basic import int64, upcast +from pytensor.scalar.basic import upcast from pytensor.tensor import elemwise_cgen as cgen from pytensor.tensor import get_vector_length from pytensor.tensor.basic import _get_vector_length, as_tensor_variable @@ -29,7 +28,6 @@ continuous_dtypes, discrete_dtypes, float_dtypes, - lvector, ) from pytensor.tensor.utils import ( broadcast_static_dim_lengths, @@ -40,7 +38,7 @@ from pytensor.utils import uniq, unzip -class DimShuffle(ExternalCOp): +class DimShuffle(COp): """ Allows to reorder the dimensions of a tensor or insert or remove broadcastable dimensions. @@ -114,20 +112,9 @@ class DimShuffle(ExternalCOp): _f16_ok = True check_input = False __props__ = ("input_ndim", "new_order") - c_func_file = "c_code/dimshuffle.c" - c_func_name = "APPLY_SPECIFIC(cpu_dimshuffle)" view_map = {0: [0]} - @property - def params_type(self): - return ParamsType( - _new_order=lvector, - input_ndim=int64, - ) - def __init__(self, *, input_ndim: int, new_order: Sequence[int | Literal["x"]]): - super().__init__([self.c_func_file], self.c_func_name) - if not isinstance(input_ndim, int): raise TypeError(f"input_ndim must be an integer, got {type(int)}") @@ -187,11 +174,86 @@ def __init__(self, *, input_ndim: int, new_order: Sequence[int | Literal["x"]]): self.is_matrix_transpose = not augment and is_left_expanded_matrix_transpose def __setstate__(self, state): + # Old pickles carry ExternalCOp attributes (func_files, ...); drop them, + # the C code is now emitted inline by `c_code`. + for key in ("func_files", "func_codes", "func_name", "code_sections"): + state.pop(key, None) self.__dict__.update(state) - if not hasattr(self, "func_files"): - # Perhaps we are loading an old `Op` version of DimShuffle. - # Let's just build the ExternalCOp. - super().__init__([self.c_func_file], self.c_func_name) + + def c_code_cache_version(self): + return (2,) + + def c_code(self, node, name, inputs, outputs, sub): + """Emit a straight-line view construction specialized on the static + permutation; dropped axes are guarded by runtime squeeze checks.""" + (inp,) = inputs + (out,) = outputs + fail = sub["fail"] + new_order = self._new_order + nd_out = len(new_order) + + guards = "\n".join( + f""" + if (PyArray_DIMS({inp})[{d}] != 1) {{ + PyErr_SetString(PyExc_ValueError, + "DimShuffle: cannot drop axis {d} with length not equal to one."); + {fail} + }}""" + for d in self.drop + ) + + assigns = [] + for i, j in enumerate(new_order): + if j == -1: + # An augmented (broadcast) axis. The length-1 stride is set to the + # itemsize rather than zero: the value is never dereferenced, but + # some BLAS implementations mishandle a zero stride. + assigns.append(f"dimensions[{i}] = 1;") + assigns.append(f"strides[{i}] = itemsize;") + else: + assigns.append(f"dimensions[{i}] = PyArray_DIMS({inp})[{j}];") + # Normalize length-1 strides to the itemsize for the same reason. + assigns.append( + f"strides[{i}] = PyArray_DIMS({inp})[{j}] == 1 ? " + f"itemsize : PyArray_STRIDES({inp})[{j}];" + ) + + if nd_out: + shape_block = ( + f"npy_intp dimensions[{nd_out}];\n" + f"npy_intp strides[{nd_out}];\n" + "\n".join(assigns) + ) + dims_ptr = "dimensions" + strides_ptr = "strides" + else: + shape_block = "" + dims_ptr = "NULL" + strides_ptr = "NULL" + + return f""" + {{ + npy_intp itemsize = PyArray_ITEMSIZE({inp}); + {guards} + {shape_block} + + Py_XDECREF({out}); + // Borrow only the writable flag from the input; NPY_OWNDATA stays 0. + {out} = (PyArrayObject*)PyArray_New( + &PyArray_Type, {nd_out}, {dims_ptr}, + PyArray_TYPE({inp}), {strides_ptr}, + PyArray_DATA({inp}), itemsize, + (NPY_ARRAY_WRITEABLE * PyArray_ISWRITEABLE({inp})), + NULL); + if ({out} == NULL) {{ + {fail} + }} + + // Declare the result a view of the input and recompute its flags. + Py_INCREF((PyObject*){inp}); + PyArray_SetBaseObject({out}, (PyObject*){inp}); + PyArray_UpdateFlags({out}, NPY_ARRAY_UPDATE_ALL); + }} + """ def make_node(self, inp): input = as_tensor_variable(inp) diff --git a/tests/link/c/test_params_type.py b/tests/link/c/test_params_type.py index d8bd2b754a..2e6d123020 100644 --- a/tests/link/c/test_params_type.py +++ b/tests/link/c/test_params_type.py @@ -338,7 +338,8 @@ def test_op_params(self): a, b, c = 2, 3, -7 x = matrix(dtype="float64") y1 = QuadraticOpFunc(a, b, c)(x) - y2 = QuadraticCOpFunc(a, b, c)(x) + with pytest.warns(FutureWarning, match="ExternalCOp is deprecated"): + y2 = QuadraticCOpFunc(a, b, c)(x) f1 = pytensor.function([x], y1, mode="CVM") f2 = pytensor.function([x], y2, mode="CVM") shape = (100, 100)