feat(python): add NumPy ArrayMethod float64 loops

codegen/src/targets/python/extension.rs

@@ -3,11 +3,8 @@ use anyhow::Result;
 use crate::{
     model::{CompiledFunction, Output},
     output::GeneratedFile,
-    render::{
-        Writer,
-        c::{self, Dialect},
-    },
-    targets::group_by_source,
+    render::Writer,
+    targets::{group_by_source, native::c_result_name},
 };
 
 use super::C_HEADER;
@@ -19,15 +16,17 @@ pub(crate) fn render(functions: &[CompiledFunction]) -> Result<Vec<GeneratedFile
         let register = format!("ptfkit_register_{slug}");
         let mut source = Writer::new();
         source.write(C_HEADER);
-        source.write("#include \"ufunc.h\"\n\n");
+        source.write(format_args!(
+            "#include <ptfkit/{slug}.h>\n#include \"ufunc.h\"\n\n"
+        ));
         for function in functions {
             source.write(ufunc(function)?);
         }
         source.line(format_args!("int {register}(PyObject *module) {{"));
         source.indented(|writer| {
             for function in functions {
                 writer.line(format_args!(
-                    "if (ptfkit_add_ufunc(module, \"{name}\", {name}_functions, {name}_types, {nin}, {nout}) < 0) return -1;",
+                    "if (ptfkit_add_ufunc(module, \"{name}\", {nin}, {nout}, &{name}_spec) < 0) return -1;",
                     name = function.core.name,
                     nin = function.core.inputs.len(),
                     nout = output_count(&function.core.output),
@@ -46,6 +45,7 @@ pub(crate) fn render(functions: &[CompiledFunction]) -> Result<Vec<GeneratedFile
     entry.write(
         r#"#define PY_SSIZE_T_CLEAN
 #define PY_ARRAY_UNIQUE_SYMBOL PTFKIT_ARRAY_API
+#define NPY_TARGET_VERSION NPY_2_0_API_VERSION
 #include <Python.h>
 #include <numpy/arrayobject.h>
 #include <numpy/ufuncobject.h>"#,
@@ -88,50 +88,102 @@ fn ufunc(function: &CompiledFunction) -> Result<String> {
         ],
         Output::Struct(fields) => fields.clone(),
     };
-    let types = std::iter::repeat_n("NPY_DOUBLE", inputs.len() + values.len())
-        .collect::<Vec<_>>()
-        .join(", ");
     let mut writer = Writer::new();
-    writer.line(format_args!("static void {name}_loop(char **args, const npy_intp *dimensions, const npy_intp *steps, void *data) {{"));
+    writer.line(format_args!(
+        "static int {name}_contiguous_loop(PyArrayMethod_Context *context, char *const *data, const npy_intp *dimensions, const npy_intp *strides, NpyAuxData *transferdata) {{"
+    ));
     writer.indented(|writer| {
-        writer.line("npy_intp index;");
-        writer.line("for (index = 0; index < dimensions[0]; index++) {");
+        writer.line("(void)context;");
+        writer.line("(void)strides;");
+        writer.line("(void)transferdata;");
+        for (index, input) in inputs.iter().enumerate() {
+            writer.line(format_args!(
+                "const double *in_{input} = (const double *)data[{index}];"
+            ));
+        }
+        for (index, value) in values.iter().enumerate() {
+            writer.line(format_args!(
+                "double *out_{value} = (double *)data[{}];",
+                inputs.len() + index
+            ));
+        }
+        writer.line("for (npy_intp index = 0; index < dimensions[0]; index++) {");
         writer.indented(|writer| {
-            for (index, input) in inputs.iter().enumerate() {
-                writer.line(format_args!(
-                    "const double {input} = *(const double *)args[{index}];"
-                ));
-            }
-            for variable in &function.ir.variables {
-                writer.write(format_args!("const double {} = ", variable.name));
-                writer.write(c::expression(
-                    &variable.expression,
-                    inputs,
-                    &function.ir.variables,
-                    Dialect::C,
-                ));
-                writer.line(";");
+            for input in inputs {
+                writer.line(format_args!("const double {input} = in_{input}[index];"));
             }
-            for (index, value) in values.iter().enumerate() {
+            render_kernel_call(writer, function, inputs, &values, Some("[index]"));
+        });
+        writer.line("}");
+        writer.line("return 0;");
+    });
+    writer.line("}");
+    writer.blank_line();
+    writer.line(format_args!(
+        "static int {name}_strided_loop(PyArrayMethod_Context *context, char *const *data, const npy_intp *dimensions, const npy_intp *strides, NpyAuxData *transferdata) {{"
+    ));
+    writer.indented(|writer| {
+        writer.line("(void)context;");
+        writer.line("(void)transferdata;");
+        writer.line("for (npy_intp index = 0; index < dimensions[0]; index++) {");
+        writer.indented(|writer| {
+            for (index, input) in inputs.iter().enumerate() {
                 writer.line(format_args!(
-                    "*(double *)args[{}] = {value};",
-                    inputs.len() + index
+                    "const double {input} = *(const double *)(data[{index}] + index * strides[{index}]);"
                 ));
             }
-            writer.line(format_args!(
-                "for (int arg = 0; arg < {}; arg++) args[arg] += steps[arg];",
-                inputs.len() + values.len()
-            ));
+            render_kernel_call(writer, function, inputs, &values, None);
         });
         writer.line("}");
+        writer.line("return 0;");
     });
     writer.line("}");
     writer.write(format_args!(
-        "static PyUFuncGenericFunction {name}_functions[] = {{ {name}_loop }};\nstatic char {name}_types[] = {{ {types} }};\n\n"
+        "static PyType_Slot {name}_slots[] = {{\n    {{NPY_METH_strided_loop, {name}_strided_loop}},\n    {{NPY_METH_contiguous_loop, {name}_contiguous_loop}},\n    {{0, NULL}},\n}};\nstatic PyArrayMethod_Spec {name}_spec = {{\n    .name = \"{name}\",\n    .nin = {},\n    .nout = {},\n    .casting = NPY_SAME_KIND_CASTING,\n    .slots = {name}_slots,\n}};\n\n",
+        inputs.len(),
+        values.len(),
     ));
     Ok(writer.into_string())
 }
 
+fn render_kernel_call(
+    writer: &mut Writer,
+    function: &CompiledFunction,
+    inputs: &[String],
+    values: &[String],
+    output_index: Option<&str>,
+) {
+    let arguments = inputs.join(", ");
+    let result = match &function.core.output {
+        Output::Scalar => "double".to_owned(),
+        Output::Struct(_) => c_result_name(
+            function.entry.spec.functions[function.function_index]
+                .result_class()
+                .expect("record output has a result class"),
+        ),
+    };
+    writer.line(format_args!(
+        "const {result} ptfkit_result = {}({arguments});",
+        function.core.name
+    ));
+    for (index, value) in values.iter().enumerate() {
+        let result_value = match &function.core.output {
+            Output::Scalar => "ptfkit_result".to_owned(),
+            Output::Struct(_) => format!("ptfkit_result.{value}"),
+        };
+        match output_index {
+            None => writer.line(format_args!(
+                "*(double *)(data[{}] + index * strides[{}]) = {result_value};",
+                inputs.len() + index,
+                inputs.len() + index
+            )),
+            Some(output_index) => {
+                writer.line(format_args!("out_{value}{output_index} = {result_value};"))
+            }
+        }
+    }
+}
+
 fn output_count(output: &Output) -> usize {
     match output {
         Output::Scalar => 1,

codegen/src/targets/python/wrapper.rs

@@ -140,6 +140,7 @@ fn module_source(
     module.write(format_args!(
         "from __future__ import annotations\n\nfrom typing import {typing_imports}\n\n"
     ));
+    module.line("from ptfkit._dispatch import call as _call");
     module.block(
         "from ptfkit._ptfkit import (",
         |writer| {
@@ -204,16 +205,6 @@ fn render_function(module: &mut Module, function: &PythonFunction<'_>) {
         .result_class
         .map(|class| format!("{class}[NDArray[floating]]"))
         .unwrap_or_else(|| "NDArray[floating]".into());
-    let out = if function.result_class.is_some() {
-        "tuple(out)"
-    } else {
-        "out"
-    };
-    let result = if let Some(result_class) = function.result_class {
-        format!("return {result_class}(*values)")
-    } else {
-        "return values".into()
-    };
     module.line("@overload");
     module.line(format_args!(
         "def {}(*, {}) -> {scalar_result}: ...",
@@ -242,22 +233,23 @@ fn render_function(module: &mut Module, function: &PythonFunction<'_>) {
     module.line(format_args!(") -> {scalar_result} | {array_result}:"));
     module.indented(|writer| {
         render_docstring(writer, &function.docstring, 4);
-        writer.line("if out is None:");
-        writer.indented(|writer| {
-            writer.line(format_args!(
-                "values = _{}({})",
-                function.rust_name, function.parameters
-            ));
-        });
-        writer.line("else:");
+        if function.result_class.is_some() {
+            writer.line("values = _call(");
+        } else {
+            writer.line("return _call(");
+        }
         writer.indented(|writer| {
-            writer.line(format_args!(
-                "values = _{}({}, out={out})",
-                function.rust_name, function.parameters
-            ));
+            writer.line(format_args!("_{},", function.rust_name));
+            for parameter in function.parameters.split(", ") {
+                writer.line(format_args!("{parameter},"));
+            }
+            writer.line("out=out,");
         });
-        writer.blank_line();
-        writer.line(result);
+        writer.line(")");
+        if let Some(result_class) = function.result_class {
+            writer.blank_line();
+            writer.line(format_args!("return {result_class}(*values)"));
+        }
     });
 }
 

targets/ptfkit-py/CMakeLists.txt

@@ -4,7 +4,12 @@ project(ptfkit LANGUAGES C)
 
 find_package(Python REQUIRED COMPONENTS Interpreter Development.Module NumPy)
 
+set(PTFKIT_NATIVE_INCLUDE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/vendor/ptfkit-native/include")
+if(NOT EXISTS "${PTFKIT_NATIVE_INCLUDE_DIR}/ptfkit/ptfkit.h")
+    set(PTFKIT_NATIVE_INCLUDE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../ptfkit-native/include")
+endif()
+
 Python_add_library(_ptfkit MODULE WITH_SOABI src/ptfkit/ptfkit.c)
 target_link_libraries(_ptfkit PRIVATE Python::NumPy)
-target_include_directories(_ptfkit PRIVATE src/ptfkit)
+target_include_directories(_ptfkit PRIVATE src/ptfkit "${PTFKIT_NATIVE_INCLUDE_DIR}")
 install(TARGETS _ptfkit DESTINATION ptfkit)

targets/ptfkit-py/pyproject.toml

@@ -62,3 +62,6 @@ testpaths = ["tests"]
 [tool.scikit-build]
 minimum-version = "build-system.requires"
 wheel.exclude = ["*.c", "*.h"]
+
+[tool.scikit-build.sdist.force-include]
+"../ptfkit-native/include" = "vendor/ptfkit-native/include"

targets/ptfkit-py/src/ptfkit/_dispatch.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+from typing import Any
+
+from numpy import asarray, float64, ufunc
+
+
+def call(function: ufunc, *inputs: object, out: object) -> Any:  # noqa: ANN401
+    inputs = tuple(asarray(value, dtype=float64) for value in inputs)
+    if out is None:
+        return function(*inputs)
+    if isinstance(out, tuple):
+        out = tuple(out)
+    return function(*inputs, out=out)

targets/ptfkit-py/src/ptfkit/ahuja1984.c

@@ -1,28 +1,61 @@
 /* @generated by ptfkit-codegen; DO NOT EDIT. */
+#include <ptfkit/ahuja1984.h>
 #include "ufunc.h"
 
-static void calc_ptf_ahuja1984_loop(char **args, const npy_intp *dimensions, const npy_intp *steps,
-                                    void *data) {
-    npy_intp index;
-    for (index = 0; index < dimensions[0]; index++) {
-        const double total_porosity = *(const double *)args[0];
-        const double theta_33 = *(const double *)args[1];
-        const double coefficient_b = *(const double *)args[2];
-        const double exponent_n = *(const double *)args[3];
-        const double effective_porosity = total_porosity - theta_33;
-        const double k_sat = coefficient_b * pow(effective_porosity, exponent_n);
-        *(double *)args[4] = k_sat;
-        for (int arg = 0; arg < 5; arg++)
-            args[arg] += steps[arg];
+static int calc_ptf_ahuja1984_contiguous_loop(PyArrayMethod_Context *context, char *const *data,
+                                              const npy_intp *dimensions, const npy_intp *strides,
+                                              NpyAuxData *transferdata) {
+    (void)context;
+    (void)strides;
+    (void)transferdata;
+    const double *in_total_porosity = (const double *)data[0];
+    const double *in_theta_33 = (const double *)data[1];
+    const double *in_coefficient_b = (const double *)data[2];
+    const double *in_exponent_n = (const double *)data[3];
+    double *out_k_sat = (double *)data[4];
+    for (npy_intp index = 0; index < dimensions[0]; index++) {
+        const double total_porosity = in_total_porosity[index];
+        const double theta_33 = in_theta_33[index];
+        const double coefficient_b = in_coefficient_b[index];
+        const double exponent_n = in_exponent_n[index];
+        const double ptfkit_result =
+            calc_ptf_ahuja1984(total_porosity, theta_33, coefficient_b, exponent_n);
+        out_k_sat[index] = ptfkit_result;
     }
+    return 0;
+}
+
+static int calc_ptf_ahuja1984_strided_loop(PyArrayMethod_Context *context, char *const *data,
+                                           const npy_intp *dimensions, const npy_intp *strides,
+                                           NpyAuxData *transferdata) {
+    (void)context;
+    (void)transferdata;
+    for (npy_intp index = 0; index < dimensions[0]; index++) {
+        const double total_porosity = *(const double *)(data[0] + index * strides[0]);
+        const double theta_33 = *(const double *)(data[1] + index * strides[1]);
+        const double coefficient_b = *(const double *)(data[2] + index * strides[2]);
+        const double exponent_n = *(const double *)(data[3] + index * strides[3]);
+        const double ptfkit_result =
+            calc_ptf_ahuja1984(total_porosity, theta_33, coefficient_b, exponent_n);
+        *(double *)(data[4] + index * strides[4]) = ptfkit_result;
+    }
+    return 0;
 }
-static PyUFuncGenericFunction calc_ptf_ahuja1984_functions[] = {calc_ptf_ahuja1984_loop};
-static char calc_ptf_ahuja1984_types[] = {NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE,
-                                          NPY_DOUBLE};
+static PyType_Slot calc_ptf_ahuja1984_slots[] = {
+    {NPY_METH_strided_loop, calc_ptf_ahuja1984_strided_loop},
+    {NPY_METH_contiguous_loop, calc_ptf_ahuja1984_contiguous_loop},
+    {0, NULL},
+};
+static PyArrayMethod_Spec calc_ptf_ahuja1984_spec = {
+    .name = "calc_ptf_ahuja1984",
+    .nin = 4,
+    .nout = 1,
+    .casting = NPY_SAME_KIND_CASTING,
+    .slots = calc_ptf_ahuja1984_slots,
+};
 
 int ptfkit_register_ahuja1984(PyObject *module) {
-    if (ptfkit_add_ufunc(module, "calc_ptf_ahuja1984", calc_ptf_ahuja1984_functions,
-                         calc_ptf_ahuja1984_types, 4, 1) < 0)
+    if (ptfkit_add_ufunc(module, "calc_ptf_ahuja1984", 4, 1, &calc_ptf_ahuja1984_spec) < 0)
         return -1;
     return 0;
 }

targets/ptfkit-py/src/ptfkit/ahuja1984.py

@@ -23,6 +23,7 @@
 
 from typing import TYPE_CHECKING, overload
 
+from ptfkit._dispatch import call as _call
 from ptfkit._ptfkit import (
     calc_ptf_ahuja1984 as _calc_ptf_ahuja1984,
 )
@@ -100,9 +101,11 @@ def calc_ptf_ahuja1984(
         total_porosity must be greater than or equal to theta_33.
 
     """
-    if out is None:
-        values = _calc_ptf_ahuja1984(total_porosity, theta_33, coefficient_b, exponent_n)
-    else:
-        values = _calc_ptf_ahuja1984(total_porosity, theta_33, coefficient_b, exponent_n, out=out)
-
-    return values
+    return _call(
+        _calc_ptf_ahuja1984,
+        total_porosity,
+        theta_33,
+        coefficient_b,
+        exponent_n,
+        out=out,
+    )