31121514451

Committed 06 Aug 2026 04:55PM UTC coverage: 65.036% (+0.02%) from 65.021%

Build # 31121514451

Build Type

Pull #940

github

Committed by

web-flow

Commit Message

Merge 50b8a37e8 into 7d5b198bd

Pull Request Pull Request #940: Compile time reporting

Coverage Stats

32 of 34 new or added lines in 3 files covered. (94.12%)

1 existing line in 1 file now uncovered.

46293 of 71181 relevant lines covered (65.04%)

717.51 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

64.43

/python/docc/python/python_program.py

import inspect
import json
import shutil
import textwrap
import ast
import os
import getpass
import hashlib
import ml_dtypes
import numpy as np
from typing import Annotated, get_origin, get_args, Any, Optional
import time

from docc.sdfg import (
    Scalar,
    PrimitiveType,
    Pointer,
    Structure,
    Array,
    Type,
    Tensor,
    StructuredSDFG,
    StructuredSDFGBuilder,
    DoccMetrics,
)
from docc.compiler.docc_program import DoccProgram
from docc.compiler.compiled_sdfg import CompiledSDFG
from docc.python.ast_parser import ASTParser
from docc.python.types import element_type_from_sdfg_type


def _compile_wrapper(self, output_folder=None):
    """Wrapper to allow StructuredSDFG.compile() to return a CompiledSDFG."""
    lib_path = self._compile(output_folder)
    return CompiledSDFG(lib_path, self)


# Monkey-patch StructuredSDFG to add compile method
StructuredSDFG.compile = _compile_wrapper


def _map_python_type(dtype):
    """Map Python/numpy types to SDFG types."""
    # If it is already a sdfg Type, return it
    if isinstance(dtype, Type):
        return dtype

    # Handle Annotated for Arrays
    if get_origin(dtype) is Annotated:
        args = get_args(dtype)
        base_type = args[0]
        metadata = args[1:]

        if base_type is np.ndarray:
            # Convention: Annotated[np.ndarray, shape, dtype]
            shape = metadata[0]
            elem_type = Scalar(PrimitiveType.Double)  # Default

            if len(metadata) > 1:
                possible_dtype = metadata[1]
                elem_type = _map_python_type(possible_dtype)

            return Pointer(elem_type)

    # Handle numpy.ndarray[Shape, DType]
    if get_origin(dtype) is np.ndarray:
        args = get_args(dtype)
        # args[0] is shape, args[1] is dtype
        if len(args) >= 2:
            elem_type = _map_python_type(args[1])
            return Pointer(elem_type)

    # Handle a parametrized numpy dtype generic, e.g. numpy.dtype[numpy.float64]
    # (produced by npt.NDArray[RealT] -> ndarray[Any, dtype[float64]]).
    if get_origin(dtype) is np.dtype:
        inner = get_args(dtype)
        if inner:
            return _map_python_type(inner[0])

    # Simple mapping for python types
    if dtype is float or dtype is np.float64:
        return Scalar(PrimitiveType.Double)
    elif dtype is np.float32:
        return Scalar(PrimitiveType.Float)
    elif dtype is bool or dtype is np.bool_:
        return Scalar(PrimitiveType.Bool)
    elif dtype is int or dtype is np.int64:
        return Scalar(PrimitiveType.Int64)
    elif dtype is np.int32:
        return Scalar(PrimitiveType.Int32)
    elif dtype is np.int16:
        return Scalar(PrimitiveType.Int16)
    elif dtype is np.int8:
        return Scalar(PrimitiveType.Int8)
    elif dtype is np.uint64:
        return Scalar(PrimitiveType.UInt64)
    elif dtype is np.uint32:
        return Scalar(PrimitiveType.UInt32)
    elif dtype is np.uint16:
        return Scalar(PrimitiveType.UInt16)
    elif dtype is np.uint8:
        return Scalar(PrimitiveType.UInt8)

    # Handle Python classes - map to Structure type
    if inspect.isclass(dtype):
        # Use the class name as the structure name
        return Pointer(Structure(dtype.__name__))

    return dtype


class PythonProgram(DoccProgram):

    def __init__(
        self,
        func,
        target: str = "none",
        category: str = "server",
        instrumentation_mode: Optional[str] = None,
        capture_args: Optional[bool] = None,
        remote_tuning: bool = False,
    ):
        super().__init__(
            name=func.__name__,
            target=target,
            category=category,
            instrumentation_mode=instrumentation_mode,
            capture_args=capture_args,
            remote_tuning=remote_tuning,
        )
        self.func = func
        self._last_structure_member_info = {}

    def __call__(self, *args: Any) -> Any:
        # JIT compile and run. CompiledSDFG validates the call mode (numpy /
        # cupy / torch) and rejects GPU arrays on non-device-resident artifacts.
        compiled = self.compile(*args)
        res = compiled(*args)

        # Handle return value conversion based on annotation
        sig = inspect.signature(self.func)
        ret_annotation = sig.return_annotation

        if ret_annotation is not inspect.Signature.empty:
            if get_origin(ret_annotation) is Annotated:
                type_args = get_args(ret_annotation)
                if len(type_args) >= 1 and type_args[0] is np.ndarray:
                    shape = None
                    if len(type_args) >= 2:
                        shape = type_args[1]

                    if shape is not None:
                        try:
                            return np.ctypeslib.as_array(res, shape=shape)
                        except Exception:
                            pass

        # Try to infer return shape from metadata
        if hasattr(compiled, "get_return_shape"):
            shape = compiled.get_return_shape(*args)
            if shape is not None:
                try:
                    return np.ctypeslib.as_array(res, shape=shape)
                except Exception:
                    pass

        return res

    def compile(
        self,
        *args: Any,
        output_folder: Optional[str] = None,
        instrumentation_mode: Optional[str] = None,
        capture_args: Optional[bool] = None,
        remote_tuning: Optional[bool] = None,
    ) -> CompiledSDFG:
        original_output_folder = output_folder

        metrics = DoccMetrics()
        compile_start_time = time.perf_counter()
        metrics.add_metric("function", self.name, "source")
        metrics.add_frontend_source_info("python")

        # Resolve options
        instrumentation_mode, capture_args, remote_tuning = (
            self._resolve_compile_options(
                instrumentation_mode, capture_args, remote_tuning
            )
        )

        # When binary reuse is requested, the build run must persist the
        # normalized SDFG (py4.norm.json) so a later run can reload it without
        # re-parsing/recompiling. Force the dump if instrumentation/capture
        # would not already produce it.
        docc_reuse_binaries = os.environ.get("DOCC_REUSE_BINARIES")
        if docc_reuse_binaries and not self.debug_dump:
            self.debug_dump = True

        # 1. Analyze arguments and shapes
        arg_types = []
        shape_values = []  # List of unique shape values found
        shape_sources = []  # List of (arg_idx, dim_idx) for each unique shape value

        # Mapping from (arg_idx, dim_idx) -> unique_shape_idx
        arg_shape_mapping = {}

        # First pass: collect scalar integer arguments and their values
        sig = inspect.signature(self.func)
        params = list(sig.parameters.items())
        scalar_int_params = {}  # Maps value -> parameter name (first one wins)
        for i, ((name, param), arg) in enumerate(zip(params, args)):
            if isinstance(arg, (int, np.integer)) and not isinstance(
                arg, (bool, np.bool_)
            ):
                val = int(arg)
                if val not in scalar_int_params:
                    scalar_int_params[val] = name

        for i, arg in enumerate(args):
            t = self._infer_type(arg)
            arg_types.append(t)

            if isinstance(arg, np.ndarray):
                for dim_idx, dim_val in enumerate(arg.shape):
                    # Check if we've seen this value
                    if dim_val in shape_values:
                        # Reuse
                        u_idx = shape_values.index(dim_val)
                    else:
                        # New
                        u_idx = len(shape_values)
                        shape_values.append(dim_val)
                        shape_sources.append((i, dim_idx))

                    arg_shape_mapping[(i, dim_idx)] = u_idx

        # 2. Signature - include scalar-shape equivalences for correct caching
        mapping_sig = sorted(arg_shape_mapping.items())
        type_sig = ", ".join(self._type_to_str(t) for t in arg_types)
        signature = f"{type_sig}|{mapping_sig}"

        # In-memory cache key: the structural signature plus the resolved compile
        # options, so repeated in-process compiles with different
        # instrumentation/arg-capture/remote-tuning do not alias to the first
        # built binary (the on-disk hash already accounts for these options).
        mem_cache_key = (
            f"{signature}|{capture_args}|{instrumentation_mode}|{remote_tuning}"
        )

        if output_folder is None:
            source_path = inspect.getsourcefile(self.func)
            hash_input = f"{source_path}|{self.name}|{self.target}|{self.category}|{capture_args}|{instrumentation_mode}|{remote_tuning}|{signature}".encode(
                "utf-8"
            )
            stable_id = hashlib.sha256(hash_input).hexdigest()[:16]
            filename = os.path.basename(inspect.getsourcefile(self.func))

            docc_tmp = os.environ.get("DOCC_TMP")
            if docc_tmp:
                output_folder = (
                    f"{docc_tmp}/{filename}-{self.name}-{self.target}-{stable_id}"
                )
            else:
                user = os.getenv("USER")
                if not user:
                    user = getpass.getuser()
                output_folder = f"/tmp/{user}/DOCC/{self.name}-{stable_id}"

        if original_output_folder is None and mem_cache_key in self.cache:
            return self.cache[mem_cache_key]

        # 3. Reuse a previously built binary if requested and available.
        # Structure arguments need per-member layout info that is only produced
        # while parsing the kernel, so reuse is limited to plain array/scalar
        # kernels; anything else falls through to a full rebuild.
        has_struct_args = any(
            isinstance(t, Pointer)
            and t.has_pointee_type()
            and isinstance(t.pointee_type, Structure)
            for t in arg_types
        )
        if docc_reuse_binaries and not has_struct_args:
            reused = self._try_reuse_binary(output_folder, shape_sources)
            if reused is not None:
                if original_output_folder is None:
                    self.cache[mem_cache_key] = reused

                metrics.capture_env_vars()
                metrics.append_to(output_folder)
                return reused

        # 4. Build SDFG
        if os.path.exists(output_folder):
            # Multiple python processes running the same code?
            shutil.rmtree(output_folder)
        sdfg, out_args, out_shapes, out_strides = self._build_sdfg(
            arg_types, args, arg_shape_mapping, shape_values
        )

        lib_path = self.sdfg_pipe(
            sdfg,
            output_folder,
            instrumentation_mode,
            capture_args,
            remote_tuning,
            metrics=metrics,
        )

        # Persist the return-value layout so a later DOCC_REUSE_BINARIES run can
        # rebuild the CompiledSDFG without re-parsing the kernel.
        if output_folder:
            self._persist_return_layout(output_folder, sdfg, out_shapes, out_strides)

        # 5. Create CompiledSDFG
        compiled = CompiledSDFG(
            lib_path,
            sdfg,
            shape_sources,
            self._last_structure_member_info,
            out_args,
            out_shapes,
            out_strides,
            device_resident=self._device_resident,
            device_backend=self._device_backend,
            target=self.target,
        )

        # Cache if using default output folder
        if original_output_folder is None:
            self.cache[mem_cache_key] = compiled

        compile_time_ms = round((time.perf_counter() - compile_start_time) * 1000)
        metrics.add_metric("compile_time_ms", compile_time_ms, "compile")
        metrics.capture_env_vars()
        metrics.append_to(output_folder)

        return compiled

    def _persist_return_layout(
        self, output_folder: str, sdfg: StructuredSDFG, out_shapes, out_strides
    ) -> None:
        """Stamp the return-value layout into the persisted SDFG metadata.

        The return shapes/strides are discovered while parsing the kernel and
        are not otherwise recoverable from the SDFG structure. Persisting them
        (into the same ``py4.norm.json`` the reuse path loads) lets a later
        ``DOCC_REUSE_BINARIES`` run reconstruct the CompiledSDFG without
        re-parsing/recompiling.
        """
        json_path = os.path.join(output_folder, f"{sdfg.name}.py4.norm.json")
        if not os.path.exists(json_path):
            return
        try:
            with open(json_path) as f:
                data = json.load(f)
            metadata = data.setdefault("metadata", {})
            metadata["output_shapes"] = json.dumps(out_shapes)
            metadata["output_strides"] = json.dumps(out_strides)
            with open(json_path, "w") as f:
                json.dump(data, f)
        except (OSError, ValueError):
            pass

    def _try_reuse_binary(
        self, output_folder: Optional[str], shape_sources
    ) -> Optional[CompiledSDFG]:
        """Reload a cached ``.so`` + normalized SDFG instead of recompiling.

        Mirrors the strictness of the pytorch/mlir binary-reuse path: when the
        cache directory does not exist yet this returns ``None`` so the caller
        performs a first build (no error); but when the directory *does* exist
        and a required artifact is missing, it raises ``ValueError`` so a broken
        or stale cache surfaces loudly instead of silently recompiling. The
        calling convention (device residency) and return-value layout are
        restored from the persisted SDFG metadata so arguments are marshalled
        exactly as they were at build time.
        """
        if not output_folder:
            return None

        # Cache directory absent -> first build. Let the caller build it; this
        # is the one case the mlir/pytorch frontends also treat as non-fatal.
        if not os.path.exists(output_folder):
            return None

        sdfg_name = f"{self.name}_sdfg"
        lib_path = os.path.join(output_folder, f"lib{sdfg_name}.so")
        json_path = os.path.join(output_folder, f"{sdfg_name}.py4.norm.json")
        if not os.path.exists(lib_path):
            raise ValueError(f"Tried reusing binary '{lib_path}' but does not exist")
        if not os.path.exists(json_path):
            raise ValueError(f"Tried loading SDFG '{json_path}' but does not exist")

        sdfg = StructuredSDFG.from_file(json_path)

        # Return arguments are recoverable directly from the SDFG signature.
        out_args = [name for name in sdfg.arguments if name.startswith("_docc_ret_")]

        # Return-value layout was persisted into the SDFG metadata at build time.
        out_shapes = {}
        out_strides = {}
        shapes_meta = sdfg.metadata("output_shapes")
        strides_meta = sdfg.metadata("output_strides")
        if shapes_meta:
            try:
                out_shapes = json.loads(shapes_meta)
            except ValueError:
                out_shapes = {}
        if strides_meta:
            try:
                out_strides = json.loads(strides_meta)
            except ValueError:
                out_strides = {}

        # Restore the device-residency calling convention chosen at compile time;
        # otherwise a device-resident binary would be fed host pointers.
        self._device_resident = sdfg.metadata("device_resident") == "1"
        backend = sdfg.metadata("device_backend")
        self._device_backend = backend or None

        return CompiledSDFG(
            lib_path,
            sdfg,
            shape_sources,
            {},
            out_args,
            out_shapes,
            out_strides,
            device_resident=self._device_resident,
            device_backend=self._device_backend,
            target=self.target,
        )

    def to_sdfg(self, *args: Any) -> StructuredSDFG:
        arg_types = [self._infer_type(arg) for arg in args]

        # Build shape mapping
        shape_values = []
        shape_sources = []
        arg_shape_mapping = {}

        sig = inspect.signature(self.func)
        params = list(sig.parameters.items())
        scalar_int_params = {}
        for i, ((name, param), arg) in enumerate(zip(params, args)):
            if isinstance(arg, (int, np.integer)) and not isinstance(
                arg, (bool, np.bool_)
            ):
                val = int(arg)
                if val not in scalar_int_params:
                    scalar_int_params[val] = name

        for i, arg in enumerate(args):
            if isinstance(arg, np.ndarray):
                for dim_idx, dim_val in enumerate(arg.shape):
                    if dim_val in shape_values:
                        u_idx = shape_values.index(dim_val)
                    else:
                        u_idx = len(shape_values)
                        shape_values.append(dim_val)
                        shape_sources.append((i, dim_idx))
                    arg_shape_mapping[(i, dim_idx)] = u_idx

        sdfg, _, _, _ = self._build_sdfg(
            arg_types, args, arg_shape_mapping, shape_values
        )
        return sdfg

    def _convert_inputs(self, args: tuple) -> tuple:
        return args

    def _convert_outputs(self, result: Any, original_args: tuple) -> Any:
        return result

    def _get_signature(self, arg_types):
        return ", ".join(self._type_to_str(t) for t in arg_types)

    def _type_to_str(self, t):
        if isinstance(t, Scalar):
            return f"Scalar({t.primitive_type})"
        elif isinstance(t, Array):
            return f"Array({self._type_to_str(t.element_type)}, {t.num_elements})"
        elif isinstance(t, Pointer):
            return f"Pointer({self._type_to_str(t.pointee_type)})"
        elif isinstance(t, Structure):
            return f"Structure({t.name})"
        return str(t)

    def _infer_type(self, arg):
        if isinstance(arg, (float, np.float64)):
            return Scalar(PrimitiveType.Double)
        elif isinstance(arg, np.float32):
            return Scalar(PrimitiveType.Float)
        elif isinstance(arg, (bool, np.bool_)):
            return Scalar(PrimitiveType.Bool)
        elif isinstance(arg, (int, np.int64)):
            return Scalar(PrimitiveType.Int64)
        elif isinstance(arg, np.int32):
            return Scalar(PrimitiveType.Int32)
        elif isinstance(arg, np.int16):
            return Scalar(PrimitiveType.Int16)
        elif isinstance(arg, np.int8):
            return Scalar(PrimitiveType.Int8)
        elif isinstance(arg, np.uint64):
            return Scalar(PrimitiveType.UInt64)
        elif isinstance(arg, np.uint32):
            return Scalar(PrimitiveType.UInt32)
        elif isinstance(arg, np.uint16):
            return Scalar(PrimitiveType.UInt16)
        elif isinstance(arg, np.uint8):
            return Scalar(PrimitiveType.UInt8)
        elif isinstance(arg, np.ndarray):
            # Map dtype
            if arg.dtype == np.float64:
                elem_type = Scalar(PrimitiveType.Double)
            elif arg.dtype == np.float32:
                elem_type = Scalar(PrimitiveType.Float)
            elif arg.dtype == np.float16:
                elem_type = Scalar(PrimitiveType.Half)
            elif arg.dtype == ml_dtypes.bfloat16:
                elem_type = Scalar(PrimitiveType.BFloat)
            elif arg.dtype == np.bool_:
                elem_type = Scalar(PrimitiveType.Bool)
            elif arg.dtype == np.int64:
                elem_type = Scalar(PrimitiveType.Int64)
            elif arg.dtype == np.int32:
                elem_type = Scalar(PrimitiveType.Int32)
            elif arg.dtype == np.int16:
                elem_type = Scalar(PrimitiveType.Int16)
            elif arg.dtype == np.int8:
                elem_type = Scalar(PrimitiveType.Int8)
            elif arg.dtype == np.uint64:
                elem_type = Scalar(PrimitiveType.UInt64)
            elif arg.dtype == np.uint32:
                elem_type = Scalar(PrimitiveType.UInt32)
            elif arg.dtype == np.uint16:
                elem_type = Scalar(PrimitiveType.UInt16)
            elif arg.dtype == np.uint8:
                elem_type = Scalar(PrimitiveType.UInt8)
            else:
                raise ValueError(f"Unsupported numpy dtype: {arg.dtype}")

            return Pointer(elem_type)
        elif isinstance(arg, str):
            # Explicitly reject strings - they are not supported
            raise ValueError(f"Unsupported argument type: {type(arg)}")
        else:
            # Check if it's a class instance
            if hasattr(arg, "__class__") and not isinstance(arg, type):
                # It's an instance of a class, return pointer to Structure
                return Pointer(Structure(arg.__class__.__name__))
            raise ValueError(f"Unsupported argument type: {type(arg)}")

    def _build_sdfg(
        self,
        arg_types,
        args,
        arg_shape_mapping,
        shape_values,
    ):
        sig = inspect.signature(self.func)

        # Handle return type - always void for SDFG, output args used for returns
        return_type = Scalar(PrimitiveType.Void)
        infer_return_type = True

        # Parse return annotation to determine output arguments if possible
        explicit_returns = []
        if sig.return_annotation is not inspect.Signature.empty:
            infer_return_type = False

            # Helper to normalize annotation to list of types
            def normalize_annotation(ann):
                # Handle Tuple[type, ...]
                origin = get_origin(ann)
                if origin is tuple:
                    type_args = get_args(ann)
                    # Tuple[()] or Tuple w/o args
                    if not type_args:
                        return []
                    # Tuple[int, float]
                    if len(type_args) > 0 and type_args[-1] is not Ellipsis:
                        return [_map_python_type(t) for t in type_args]
                    # Tuple[int, ...] - not supported for fixed number of returns yet?
                    # For now assume fixed tuple
                    return [_map_python_type(t) for t in type_args]
                else:
                    return [_map_python_type(ann)]

            explicit_returns = normalize_annotation(sig.return_annotation)
            for rt in explicit_returns:
                if not isinstance(rt, Type):
                    # Fallback if map failed (e.g. invalid annotation)
                    infer_return_type = True
                    explicit_returns = []
                    break

        builder = StructuredSDFGBuilder(f"{self.name}_sdfg", return_type)

        # Add pre-defined return arguments if we know them
        if not infer_return_type:
            for i, dtype in enumerate(explicit_returns):
                # Scalar -> Pointer(Scalar)
                # Array -> Already Pointer(Scalar). Keep it.
                arg_type = dtype
                if isinstance(dtype, Scalar):
                    arg_type = Pointer(dtype)

                builder.add_container(f"_docc_ret_{i}", arg_type, is_argument=True)

        # Register structure types for any class arguments
        # Also track member name to index mapping for each structure
        structures_to_register = {}
        structure_member_info = {}  # Maps struct_name -> {member_name: (index, type)}
        for i, (arg, dtype) in enumerate(zip(args, arg_types)):
            if isinstance(dtype, Pointer) and dtype.has_pointee_type():
                pointee = dtype.pointee_type
                if isinstance(pointee, Structure):
                    struct_name = pointee.name
                    if struct_name not in structures_to_register:
                        # Get class from arg to introspect members
                        if hasattr(arg, "__dict__"):
                            # Use __dict__ to get only instance attributes
                            # Sort by name to ensure consistent ordering
                            # Note: This alphabetical ordering is used to define the
                            # structure layout and must match the order expected by
                            # the backend code generation
                            member_types = []
                            member_names = []
                            member_shapes = []
                            for attr_name, attr_value in sorted(arg.__dict__.items()):
                                if not attr_name.startswith("_"):
                                    # Infer member type from instance attribute
                                    # Check bool before int since bool is subclass of int
                                    member_type = None
                                    member_shape = None
                                    if isinstance(attr_value, bool):
                                        member_type = Scalar(PrimitiveType.Bool)
                                    elif isinstance(attr_value, (int, np.int64)):
                                        member_type = Scalar(PrimitiveType.Int64)
                                    elif isinstance(attr_value, (float, np.float64)):
                                        member_type = Scalar(PrimitiveType.Double)
                                    elif isinstance(attr_value, np.int32):
                                        member_type = Scalar(PrimitiveType.Int32)
                                    elif isinstance(attr_value, np.float32):
                                        member_type = Scalar(PrimitiveType.Float)
                                    elif isinstance(attr_value, np.ndarray):
                                        # Array member: stored as a pointer field
                                        # (struct-of-arrays). Record the concrete
                                        # shape so attribute access can build a
                                        # tensor view over the member pointer.
                                        member_type = self._infer_type(attr_value)
                                        member_shape = [
                                            str(int(s)) for s in attr_value.shape
                                        ]
                                    # TODO: Consider using np.integer and np.floating abstract types
                                    # for more comprehensive numpy type coverage
                                    # TODO: Add support for nested structures

                                    if member_type is not None:
                                        member_types.append(member_type)
                                        member_names.append(attr_name)
                                        member_shapes.append(member_shape)

                            if member_types:
                                structures_to_register[struct_name] = member_types
                                # Build member name to (index, type, shape) mapping.
                                # shape is None for scalar members and a list of
                                # dimension-size strings for array members.
                                structure_member_info[struct_name] = {
                                    name: (idx, mtype, shape)
                                    for idx, (name, mtype, shape) in enumerate(
                                        zip(member_names, member_types, member_shapes)
                                    )
                                }

        # Store structure_member_info for later use in CompiledSDFG
        self._last_structure_member_info = structure_member_info

        # Register all discovered structures with the builder
        for struct_name, member_types in structures_to_register.items():
            builder.add_structure(struct_name, member_types)

        # Register arguments
        params = list(sig.parameters.items())
        if len(params) != len(arg_types):
            raise ValueError(
                f"Argument count mismatch: expected {len(params)}, got {len(arg_types)}"
            )

        # Add regular arguments
        tensor_table = {}
        for i, ((name, param), dtype, arg) in enumerate(zip(params, arg_types, args)):
            builder.add_container(name, dtype, is_argument=True)

            # Store layout information for arrays
            if isinstance(arg, np.ndarray):
                element_type = element_type_from_sdfg_type(dtype)

                shapes = []
                for dim_idx in range(arg.ndim):
                    dim_val = arg.shape[dim_idx]
                    if dim_val == 1:
                        # Always use literal "1" for size-1 dimensions to enable
                        # proper broadcasting detection
                        shapes.append("1")
                    else:
                        u_idx = arg_shape_mapping[(i, dim_idx)]
                        shapes.append(f"_s{u_idx}")

                strides = []
                if arg.flags["C_CONTIGUOUS"]:
                    # Row-major: stride[i] = product of shapes[i+1:]
                    for dim_idx in range(arg.ndim):
                        if dim_idx == arg.ndim - 1:
                            strides.append("1")
                        else:
                            suffix_shapes = shapes[dim_idx + 1 :]
                            if len(suffix_shapes) == 1:
                                strides.append(suffix_shapes[0])
                            else:
                                strides.append("(" + " * ".join(suffix_shapes) + ")")
                elif arg.flags["F_CONTIGUOUS"]:
                    # Column-major: stride[i] = product of shapes[:i]
                    for dim_idx in range(arg.ndim):
                        if dim_idx == 0:
                            strides.append("1")
                        else:
                            prefix_shapes = shapes[:dim_idx]
                            if len(prefix_shapes) == 1:
                                strides.append(prefix_shapes[0])
                            else:
                                strides.append("(" + " * ".join(prefix_shapes) + ")")
                else:
                    # Non-contiguous: use actual stride values
                    for dim_idx in range(arg.ndim):
                        stride_val = arg.strides[dim_idx] // arg.itemsize
                        strides.append(f"{stride_val}")

                offset = "0"
                tensor_table[name] = Tensor(element_type, shapes, strides, offset)

            elif isinstance(arg, np.generic):
                # NumPy scalar types (np.float64, np.int32, etc.) should be treated
                # as 0-d arrays for type promotion purposes - they trigger full
                # promotion, unlike Python literals which adapt to the array dtype
                element_type = element_type_from_sdfg_type(dtype)
                tensor_table[name] = Tensor(element_type, [], [], "0")

        # Add unified shape arguments only for shapes without scalar equivalents
        # and skip size-1 dimensions (they use literal "1" instead)
        for i in range(len(shape_values)):
            if shape_values[i] != 1:
                builder.add_container(
                    f"_s{i}", Scalar(PrimitiveType.Int64), is_argument=True
                )
                builder.add_assumption_lb(f"_s{i}", "1")  # Shapes must be positive
                builder.add_assumption_const(f"_s{i}", True)  # Shapes are constant

        # Create symbol table for parser
        container_table = {}
        for i, ((name, param), dtype, arg) in enumerate(zip(params, arg_types, args)):
            container_table[name] = dtype

        for i in range(len(shape_values)):
            if shape_values[i] != 1:
                container_table[f"_s{i}"] = Scalar(PrimitiveType.Int64)

        # Parse AST
        source_lines, start_line = inspect.getsourcelines(self.func)
        source = textwrap.dedent("".join(source_lines))
        tree = ast.parse(source)
        ast.increment_lineno(tree, start_line - 1)
        func_def = tree.body[0]

        filename = inspect.getsourcefile(self.func)
        function_name = self.func.__name__

        # Combine globals with closure variables (closure takes precedence)
        combined_globals = dict(self.func.__globals__)
        if self.func.__closure__ is not None and self.func.__code__.co_freevars:
            for name, cell in zip(
                self.func.__code__.co_freevars, self.func.__closure__
            ):
                combined_globals[name] = cell.cell_contents

        parser = ASTParser(
            builder,
            tensor_table,
            container_table,
            filename,
            function_name,
            infer_return_type=infer_return_type,
            globals_dict=combined_globals,
            structure_member_info=structure_member_info,
        )
        for node in func_def.body:
            parser.visit(node)

        # Emit hoisted allocations at function entry
        parser.memory_handler.emit_allocations()

        sdfg = builder.move()
        # Mark return arguments metadata
        out_args = []
        for name in sdfg.arguments:
            if name.startswith("_docc_ret_"):
                out_args.append(name)

        return (
            sdfg,
            out_args,
            parser.captured_return_shapes,
            parser.captured_return_strides,
        )


def native(
    func=None,
    *,
    target="none",
    category="server",
    instrumentation_mode=None,
    capture_args=None,
    remote_tuning=False,
):
    """Decorator to create a PythonProgram from a Python function.

    Example:
        @native
        def my_function(x: np.ndarray) -> np.ndarray:
            return x * 2

        result = my_function(np.array([1.0, 2.0, 3.0]))
    """
    if func is None:
        return lambda f: PythonProgram(
            f,
            target=target,
            category=category,
            instrumentation_mode=instrumentation_mode,
            capture_args=capture_args,
            remote_tuning=remote_tuning,
        )
    return PythonProgram(
        func,
        target=target,
        category=category,
        instrumentation_mode=instrumentation_mode,
        capture_args=capture_args,
        remote_tuning=remote_tuning,
    )

daisytuner / docc / 31121514451

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous