add more amdgpu examples (#133)

Author: makslevental

mlir/extras/ast/canonicalize.py

@@ -7,7 +7,7 @@
 from abc import ABC, abstractmethod
 from dis import findlinestarts
 from opcode import opmap
-from typing import List, Union, Sequence
+from typing import List, Union, Sequence, get_type_hints
 
 import astunparse
 from bytecode import ConcreteBytecode
@@ -29,6 +29,43 @@ def visit_FunctionDef(self, node: ast.FunctionDef):
         return node
 
 
+# https://stackoverflow.com/a/66582895/9045206
+class AnnotationsCollector(ast.NodeVisitor):
+    def __init__(self):
+        self.annotations = {}
+
+    def visit_AnnAssign(self, node):
+        if node.simple:
+            # 'simple' == a single name, not an attribute or subscription.
+            # we can therefore count on `node.target.id` to exist. This is
+            # the same criteria used for module and class-level variable
+            # annotations.
+            self.annotations[node.target.id] = node.annotation
+
+
+def function_local_annotations(func):
+    """Return a mapping of name to string annotations for function locals
+
+    Python does not retain PEP 526 "variable: annotation" variable annotations
+    within a function body, as local variables do not have a lifetime beyond
+    the local namespace. This function extracts the mapping from functions that
+    have source code available.
+
+    """
+    source = inspect.getsource(func)
+    mod = ast.parse(source)
+    assert mod.body and isinstance(mod.body[0], (ast.FunctionDef, ast.AsyncFunctionDef))
+    collector = AnnotationsCollector()
+    collector.visit(mod.body[0])
+    func.__annotations__.update(
+        {
+            name: ast.get_source_segment(source, node)
+            for name, node in collector.annotations.items()
+        }
+    )
+    return get_type_hints(func)
+
+
 def transform_func(f, *transformer_ctors: type(Transformer)):
     module = get_module_cst(f)
     context = types.SimpleNamespace()

mlir/extras/dialects/ext/_shaped_value.py

@@ -1,9 +1,11 @@
+from dataclasses import dataclass
 from functools import cached_property, reduce
-from typing import Tuple
-
 import numpy as np
+from typing import Tuple, Union, List, Any
 
-from ....ir import DenseElementsAttr, ShapedType, Type
+from ....dialects.linalg.opdsl.lang.emitter import _is_index_type
+from .arith import Scalar
+from ....ir import DenseElementsAttr, ShapedType, Type, Value, RankedTensorType
 
 S = ShapedType.get_dynamic_size()
 
@@ -61,3 +63,262 @@ def dtype(self) -> Type:
     cls.dtype.__set_name__(None, "dtype")
 
     return cls
+
+
+@dataclass(frozen=True)
+class _Indexer:
+    indices: Tuple[Union[int, Scalar, slice, "Ellipsis", None]]
+    newaxis_dims: Tuple[int, "Ellipsis"]
+    in_shape: Tuple[Union[Value, int]]
+
+    def is_constant(self):
+        return all(_is_constant_index(i) for i in self.indices)
+
+    def is_full(self):
+        return all(
+            isinstance(idx, slice)
+            # TODO(max): could also work for constant Scalar
+            and all([isinstance(x, int) for x in [idx.start, idx.stop, idx.step]])
+            and len(range(*idx.indices(self.in_shape[i]))) == self.in_shape[i]
+            for i, idx in enumerate(self.indices)
+        )
+
+    # waiting on hashable slices in 3.12 https://stackoverflow.com/a/76562346
+    # @lru_cache(maxsize=1)
+    def static_offsets(self):
+        offsets = []
+        for i in self.indices:
+            if isinstance(i, (int, Scalar)):
+                offsets.append(int(i))
+            elif isinstance(i, slice):
+                offsets.append(int(i.start))
+            else:
+                raise ValueError(f"idx {i} not supported with static offsets")
+        return tuple(offsets)
+
+    # @lru_cache(maxsize=1)
+    def static_sizes(self):
+        sizes = []
+        for i in self.indices:
+            if isinstance(i, (int, Scalar)):
+                sizes.append(1)
+            elif isinstance(i, slice):
+                start, stop, step = map(int, (i.start, i.stop, i.step))
+                if all(isinstance(j, int) for j in (start, stop, step)):
+                    s = ((stop - start) // step) + 1
+                    if (stop - start) % step == 0:
+                        s -= 1
+                    sizes.append(s)
+                else:
+                    raise ValueError(f"idx {i} not supported with static sizes")
+
+            else:
+                raise ValueError(f"idx {i} not supported with static sizes")
+        return tuple(sizes)
+
+    # @lru_cache(maxsize=1)
+    def static_strides(self):
+        strides = []
+        for i in self.indices:
+            if isinstance(i, (int, Scalar)):
+                strides.append(1)
+            elif isinstance(i, slice):
+                strides.append(int(i.step))
+            else:
+                raise ValueError(f"idx {i} not supported with static strides")
+        return tuple(strides)
+
+
+def _indices_to_indexer(
+    idx: Tuple[Union[Scalar, slice, "Ellipsis", None]], in_shape: Tuple[int]
+) -> _Indexer:
+    """Processes sequence of index objects and constructs _Indexer with
+    corresponding indexing tensor and collapse dims (i.e., scatter/gather dims).
+
+    Args:
+      idx: Sequence (list or tuple) of slices, ellipses, Scalar, or Tensors.
+      in_shape: The shape of the tensor being indexed into.
+
+    Returns:
+      _Indexer object.
+
+    """
+    idx = _canonicalize_tuple_index(idx, len(in_shape))
+
+    in_axis = 0  # Current axis in input.
+    out_axis = 0  # Current axis in output.
+    indices: List[Union[Scalar, slice, Ellipsis, None]] = [slice(None)] * len(in_shape)
+    newaxis_dims: List[int] = []
+
+    # nb: idx_e <-> idx_element
+    for idx_i, idx_e in enumerate(idx):
+        if _is_scalar(idx_e) and _has_index_type(idx_e):
+            # Handle basic Scalar indexes.
+            indices[in_axis] = idx_e
+            in_axis += 1
+        # Handle newaxis (None)
+        elif idx_e is None:
+            newaxis_dims.append(out_axis)
+            out_axis += 1
+        elif isinstance(idx_e, slice):
+            # Normalize the slice to use None when possible
+            start, stop, step = idx_e.start, idx_e.stop, idx_e.step
+            if step is None or isinstance(step, int) and step == 1:
+                step = None
+            if step is None:
+                if start is None or isinstance(start, int) and start == 0:
+                    start = None
+                if (
+                    stop is None
+                    or isinstance(stop, int)
+                    and in_shape[in_axis] != ShapedType.get_dynamic_size()
+                    and stop >= in_shape[in_axis]
+                ):
+                    stop = None
+            # Handle slice(None) and slice(None, None, -1)
+            if (
+                start is None
+                and stop is None
+                and (step is None or isinstance(step, int) and step == -1)
+            ):
+                if step == -1:
+                    raise IndexError(
+                        f"Negative step indexing mode not yet supported:\n{idx}"
+                    )
+                indices[in_axis] = slice(None)
+                out_axis += 1
+                in_axis += 1
+
+            # Handle slice index (only static shape supported)
+            else:
+                if (
+                    not isinstance(in_shape[in_axis], int)
+                    or in_shape[in_axis] == ShapedType.get_dynamic_size()
+                ):
+                    msg = (
+                        "Cannot use NumPy slice indexing on an array dimension whose "
+                        f"size is not statically known ({in_shape[in_axis]}). "
+                    )
+                    raise IndexError(msg)
+
+                if step is None:
+                    step = 1
+                indices[in_axis] = slice(start, stop, step)
+
+                out_axis += 1
+                in_axis += 1
+        else:
+            raise IndexError(f"Indexing mode not yet supported:\n{idx}")
+
+    for i, idx in enumerate(indices):
+        if _is_constant_index(idx) and _is_constant_scalar(in_shape[i]):
+            if isinstance(idx, slice):
+                indices[i] = slice(*idx.indices(int(in_shape[i])))
+            elif isinstance(idx, Scalar):
+                indices[i] = int(idx)
+
+    return _Indexer(
+        newaxis_dims=tuple(newaxis_dims), indices=tuple(indices), in_shape=in_shape
+    )
+
+
+def _canonicalize_tuple_index(idx: Tuple[Any], rank: int):
+    """Helper to
+    1. remove Ellipsis and replace with implicit trailing slice(None)s.
+    2. cast Python lists of lists or numpy arrays to index Tensors
+
+    Args:
+      rank: Rank of tensor.
+      idx: Index object (Scalar, Tensor, slice, Ellipse, or None).
+
+    Returns:
+      Tuple of index objects with no ellipses.
+    """
+
+    len_without_none = 0
+    for e in idx:
+        if e is None or e is Ellipsis:
+            continue
+        else:
+            len_without_none += 1
+
+    if len_without_none > rank:
+        raise IndexError(
+            f"Too many indices for shaped type with rank: {len_without_none} "
+            f"non-None/Ellipsis indices for dim {rank}."
+        )
+    ellipses = (i for i, elt in enumerate(idx) if elt is Ellipsis)
+    ellipsis_index = next(ellipses, None)
+    if ellipsis_index is not None:
+        if next(ellipses, None) is not None:
+            raise IndexError(
+                f"Multiple ellipses (...) not supported: {list(map(type, idx))}."
+            )
+        colons = (slice(None),) * (rank - len_without_none)
+        idx = idx[:ellipsis_index] + colons + idx[ellipsis_index + 1 :]
+    elif len_without_none < rank:
+        colons = (slice(None),) * (rank - len_without_none)
+        idx = tuple(idx) + colons
+    return idx
+
+
+def _is_int_arraylike(x):
+    """Returns True if x is array-like with integer dtype, False otherwise.
+
+    Positive (i.e., return True) examples are e.g., [[0], [1]], [[0, 1]],
+    [[[0, 1]], [[0, 1]]].
+    """
+    return (
+        isinstance(x, int)
+        and not isinstance(x, bool)
+        or isinstance(x, (list, tuple))
+        and all(_is_int_arraylike(e) for e in x)
+    )
+
+
+def _is_scalar(e: Any) -> bool:
+    """Checks whether e is a Scalar or can be used to construct a Scalar.
+
+    Args:
+      e: Anything
+    """
+    return isinstance(e, Scalar) or isinstance(e, (int, float, bool))
+
+
+def _has_index_type(e: Any) -> bool:
+    """Checks whether e has MLIR index type or a Python value that can be used
+    to construct an index type.
+
+    Args:
+      e: Anything
+    """
+    return (
+        isinstance(e, int)
+        or isinstance(e, np.ndarray)
+        and e.dtype in {np.intp}
+        or isinstance(e, Value)
+        and _is_index_type(e.type)
+        or isinstance(e.type, RankedTensorType)
+        and _is_index_type(e.type.element_type)
+    )
+
+
+def _is_constant_index(e: Any) -> bool:
+    return (
+        isinstance(e, Scalar)
+        and e.is_constant()
+        or isinstance(e, (int, float, bool))
+        or isinstance(e, slice)
+        and _is_constant_scalar(e.start)
+        and _is_constant_scalar(e.stop)
+        and _is_constant_scalar(e.step)
+    )
+
+
+def _is_constant_scalar(e: Any) -> bool:
+    return (
+        isinstance(e, Scalar)
+        and e.is_constant()
+        or (isinstance(e, (int, float, bool)) and e != ShapedType.get_dynamic_size())
+        or e is None
+    )

mlir/extras/dialects/ext/memref.py

@@ -3,9 +3,9 @@
 
 import numpy as np
 
-from ._shaped_value import ShapedValue
+from ._shaped_value import ShapedValue, _indices_to_indexer
 from .arith import Scalar, constant
-from .tensor import _indices_to_indexer, compute_result_shape_reassoc_list
+from .tensor import compute_result_shape_reassoc_list
 from .vector import Vector
 from ... import types as T
 from ...meta import region_op
@@ -15,7 +15,7 @@
     infer_mlir_type,
 )
 from ...._mlir_libs._mlir import register_value_caster
-from ....dialects import memref, arith
+from ....dialects import memref, arith, vector
 from ....dialects._ods_common import get_op_result_or_op_results
 from ....dialects.memref import *
 from ....ir import (
@@ -191,7 +191,15 @@ def __setitem__(self, idx, val):
             assert isinstance(
                 val, (Scalar, Vector)
             ), "coordinate insert requires scalar element"
-            store(val, self, idx, loc=loc)
+            if isinstance(val, Scalar):
+                store(val, self, idx, loc=loc)
+            elif isinstance(val, Vector):
+                return vector.StoreOp(
+                    valueToStore=val,
+                    base=self,
+                    indices=idx,
+                    loc=loc,
+                )
         else:
             _copy_to_subview(self, val, tuple(idx), loc=loc)