utils.py

# Copyright 2021-2022 NVIDIA Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
from __future__ import annotations

import traceback
from functools import reduce
from string import ascii_lowercase, ascii_uppercase
from types import FrameType
from typing import Any, Callable, List, Sequence, Tuple, Union, cast

import legate.core.types as ty
import numpy as np
import pyarrow as pa

from .types import NdShape

SUPPORTED_DTYPES = {
    bool: ty.bool_,
    np.bool_: ty.bool_,
    np.int8: ty.int8,
    np.int16: ty.int16,
    np.int32: ty.int32,
    int: ty.int64,  # np.int is int
    np.int64: ty.int64,
    np.uint8: ty.uint8,
    np.uint16: ty.uint16,
    np.uint32: ty.uint32,
    np.uint64: ty.uint64,  # np.uint is np.uint64
    np.float16: ty.float16,
    np.float32: ty.float32,
    float: ty.float64,
    np.float64: ty.float64,
    np.complex64: ty.complex64,
    np.complex128: ty.complex128,
    complex: ty.complex128,
}

CUNUMERIC_TYPE_MAP = {
    bool: ty.bool_,
    int: ty.int64,
    float: ty.float64,
    complex: ty.complex128,
    pa.bool_: ty.bool_,
    pa.int8: ty.int8,
    pa.int16: ty.int16,
    pa.int32: ty.int32,
    pa.int64: ty.int64,  # np.int is int
    pa.uint8: ty.uint8,
    pa.uint16: ty.uint16,
    pa.uint32: ty.uint32,
    pa.uint64: ty.uint64,  # np.uint is np.uint64
    pa.float16: ty.float16,
    pa.float32: ty.float32,
    pa.float64: ty.float64,
}


def is_advanced_indexing(key: Any) -> bool:
    if key is Ellipsis or key is None:  # np.newdim case
        return False
    if np.isscalar(key):
        return False
    if isinstance(key, slice):
        return False
    if isinstance(key, tuple):
        return any(is_advanced_indexing(k) for k in key)
    # Any other kind of thing leads to advanced indexing
    return True


def find_last_user_stacklevel() -> int:
    stacklevel = 1
    for frame, _ in traceback.walk_stack(None):
        if not frame.f_globals["__name__"].startswith("cunumeric"):
            break
        stacklevel += 1
    return stacklevel


def get_line_number_from_frame(frame: FrameType) -> str:
    return f"{frame.f_code.co_filename}:{frame.f_lineno}"


def find_last_user_frames(top_only: bool = True) -> str:
    for last, _ in traceback.walk_stack(None):
        if "__name__" not in last.f_globals:
            continue
        name = last.f_globals["__name__"]
        if not any(name.startswith(pkg) for pkg in ("cunumeric", "legate")):
            break

    if top_only:
        return get_line_number_from_frame(last)

    frames: list[FrameType] = []
    curr: Union[FrameType, None] = last
    while curr is not None:
        if "legion_top.py" in curr.f_code.co_filename:
            break
        frames.append(curr)
        curr = curr.f_back
    return "|".join(get_line_number_from_frame(f) for f in frames)


def is_supported_dtype(dtype: Any) -> bool:
    if not isinstance(dtype, np.dtype):
        raise TypeError("expected a NumPy dtype")
    return dtype.type in SUPPORTED_DTYPES


def convert_to_cunumeric_dtype(dtype: Any) -> Any:
    if dtype in CUNUMERIC_TYPE_MAP:
        return CUNUMERIC_TYPE_MAP[dtype]
    raise TypeError("dtype is not supported")


def calculate_volume(shape: NdShape) -> int:
    if len(shape) == 0:
        return 0
    return reduce(lambda x, y: x * y, shape)


def get_arg_dtype(dtype: np.dtype[Any]) -> np.dtype[Any]:
    return np.dtype(
        [("arg", np.int64), ("arg_value", dtype)],
        align=True,
    )


def get_arg_value_dtype(dtype: np.dtype[Any]) -> np.dtype[Any]:
    dt = dtype.fields["arg_value"][0].type  # type: ignore [index]
    return cast(np.dtype[Any], dt)


Modes = Tuple[List[str], List[str], List[str]]


def dot_modes(a_ndim: int, b_ndim: int) -> Modes:
    a_modes = list(ascii_lowercase[:a_ndim])
    b_modes = list(ascii_uppercase[:b_ndim])
    if a_ndim == 0 or b_ndim == 0:
        out_modes = a_modes + b_modes
    elif b_ndim == 1:
        b_modes[-1] = a_modes[-1]
        out_modes = a_modes[:-1]
    else:
        b_modes[-2] = a_modes[-1]
        out_modes = a_modes[:-1] + b_modes[:-2] + [b_modes[-1]]
    return (a_modes, b_modes, out_modes)


def inner_modes(a_ndim: int, b_ndim: int) -> Modes:
    a_modes = list(ascii_lowercase[:a_ndim])
    b_modes = list(ascii_uppercase[:b_ndim])
    if a_ndim == 0 or b_ndim == 0:
        out_modes = a_modes + b_modes
    else:
        b_modes[-1] = a_modes[-1]
        out_modes = a_modes[:-1] + b_modes[:-1]
    return (a_modes, b_modes, out_modes)


def matmul_modes(a_ndim: int, b_ndim: int) -> Modes:
    if a_ndim == 0 or b_ndim == 0:
        raise ValueError("Scalars not allowed in matmul")
    a_modes = list(ascii_lowercase[-a_ndim:])
    b_modes = list(ascii_lowercase[-b_ndim:])
    if b_ndim >= 2:
        a_modes[-1] = "A"
        b_modes[-2] = "A"
    if b_ndim == 1:
        out_modes = a_modes[:-1]
    elif a_ndim == 1:
        out_modes = b_modes[:-2] + [b_modes[-1]]
    else:
        out_modes = (
            list(ascii_lowercase[-max(a_ndim, b_ndim) : -2])
            + [a_modes[-2]]
            + [b_modes[-1]]
        )
    return (a_modes, b_modes, out_modes)


Axes = Sequence[int]
AxesPair = Tuple[Axes, Axes]
AxesPairLikeTuple = Union[
    Tuple[int, int],
    Tuple[int, Axes],
    Tuple[Axes, int],
    Tuple[Axes, Axes],
]
AxesPairLike = Union[int, AxesPairLikeTuple]


def tensordot_modes(a_ndim: int, b_ndim: int, axes: AxesPairLike) -> Modes:
    def convert_int_axes(axes: int) -> AxesPair:
        return list(range(a_ndim - axes, a_ndim)), list(range(axes))

    def convert_seq_axes(axes: AxesPairLikeTuple) -> AxesPair:
        a_axes, b_axes = axes
        return (
            [a_axes] if isinstance(a_axes, int) else list(a_axes),
            [b_axes] if isinstance(b_axes, int) else list(b_axes),
        )

    def convert_axes(axes: AxesPairLike) -> AxesPair:
        if isinstance(axes, int):
            a_axes, b_axes = convert_int_axes(axes)
        else:
            a_axes, b_axes = convert_seq_axes(axes)

        return (
            [ax + a_ndim if ax < 0 else ax for ax in a_axes],
            [ax + b_ndim if ax < 0 else ax for ax in b_axes],
        )

    def check_axes(a_axes: Axes, b_axes: Axes) -> None:
        if (
            len(a_axes) != len(b_axes)
            or len(a_axes) > a_ndim
            or len(b_axes) > b_ndim
            or len(a_axes) != len(set(a_axes))
            or len(b_axes) != len(set(b_axes))
            or any(ax < 0 for ax in a_axes)
            or any(ax < 0 for ax in b_axes)
            or any(ax >= a_ndim for ax in a_axes)
            or any(ax >= b_ndim for ax in b_axes)
        ):
            raise ValueError("Invalid axes argument")

    a_axes, b_axes = convert_axes(axes)

    check_axes(a_axes, b_axes)

    a_modes = list(ascii_lowercase[:a_ndim])
    b_modes = list(ascii_uppercase[:b_ndim])
    for a_i, b_i in zip(a_axes, b_axes):
        b_modes[b_i] = a_modes[a_i]
    a_out = [a_modes[a_i] for a_i in sorted(set(range(a_ndim)) - set(a_axes))]
    b_out = [b_modes[b_i] for b_i in sorted(set(range(b_ndim)) - set(b_axes))]

    return (a_modes, b_modes, a_out + b_out)


def deep_apply(obj: Any, func: Callable[[Any], Any]) -> Any:
    """
    Apply the provided function to objects contained at any depth within a data
    structure.

    This function will recurse over arbitrary nestings of lists, tuples and
    dicts. This recursion logic is rather limited, but this function is
    primarily meant to be used for arguments of NumPy API calls, which
    shouldn't nest their arrays very deep.
    """
    if type(obj) == list:
        return [deep_apply(x, func) for x in obj]
    elif type(obj) == tuple:
        return tuple(deep_apply(x, func) for x in obj)
    elif type(obj) == dict:
        return {k: deep_apply(v, func) for k, v in obj.items()}
    else:
        return func(obj)