incorporate comments

Author: makslevental

mlir/include/mlir-c/BuiltinTypes.h

@@ -411,9 +411,8 @@ MLIR_CAPI_EXPORTED MlirAttribute mlirMemRefTypeGetMemorySpace(MlirType type);
 /// Returns the strides of the MemRef if the layout map is in strided form.
 /// Both strides and offset are out params. strides must point to pre-allocated
 /// memory of length equal to the rank of the memref.
-MLIR_CAPI_EXPORTED void mlirMemRefTypeGetStridesAndOffset(MlirType type,
-                                                          int64_t *strides,
-                                                          int64_t *offset);
+MLIR_CAPI_EXPORTED MlirLogicalResult mlirMemRefTypeGetStridesAndOffset(
+    MlirType type, int64_t *strides, int64_t *offset);
 
 /// Returns the memory spcae of the given Unranked MemRef type.
 MLIR_CAPI_EXPORTED MlirAttribute

mlir/lib/Bindings/Python/IRTypes.cpp

@@ -12,6 +12,8 @@
 
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/BuiltinTypes.h"
+#include "mlir-c/Support.h"
+
 #include <optional>
 
 namespace py = pybind11;
@@ -618,12 +620,15 @@ class PyMemRefType : public PyConcreteType<PyMemRefType, PyShapedType> {
               return mlirMemRefTypeGetLayout(self);
             },
             "The layout of the MemRef type.")
-        .def_property_readonly(
-            "strides_and_offset",
+        .def(
+            "get_strides_and_offset",
             [](PyMemRefType &self) -> std::pair<std::vector<int64_t>, int64_t> {
               std::vector<int64_t> strides(mlirShapedTypeGetRank(self));
               int64_t offset;
-              mlirMemRefTypeGetStridesAndOffset(self, strides.data(), &offset);
+              if (mlirLogicalResultIsFailure(mlirMemRefTypeGetStridesAndOffset(
+                      self, strides.data(), &offset)))
+                throw std::runtime_error(
+                    "failed to extract strides and offset from memref");
               return {strides, offset};
             },
             "The strides and offset of the MemRef type.")

mlir/lib/CAPI/IR/BuiltinTypes.cpp

@@ -9,12 +9,14 @@
 #include "mlir-c/BuiltinTypes.h"
 #include "mlir-c/AffineMap.h"
 #include "mlir-c/IR.h"
+#include "mlir-c/Support.h"
 #include "mlir/CAPI/AffineMap.h"
 #include "mlir/CAPI/IR.h"
 #include "mlir/CAPI/Support.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Types.h"
+#include "mlir/Support/LogicalResult.h"
 
 #include <algorithm>
 
@@ -428,16 +430,16 @@ MlirAttribute mlirMemRefTypeGetMemorySpace(MlirType type) {
   return wrap(llvm::cast<MemRefType>(unwrap(type)).getMemorySpace());
 }
 
-void mlirMemRefTypeGetStridesAndOffset(MlirType type, int64_t *strides,
-                                       int64_t *offset) {
+MlirLogicalResult mlirMemRefTypeGetStridesAndOffset(MlirType type,
+                                                    int64_t *strides,
+                                                    int64_t *offset) {
   MemRefType memrefType = llvm::cast<MemRefType>(unwrap(type));
-  std::pair<SmallVector<int64_t>, int64_t> stridesOffsets =
-      getStridesAndOffset(memrefType);
-  assert(stridesOffsets.first.size() == memrefType.getRank() &&
-         "Strides and rank don't match for memref");
-  (void)std::copy(stridesOffsets.first.begin(), stridesOffsets.first.end(),
-                  strides);
-  *offset = stridesOffsets.second;
+  SmallVector<int64_t> strides_;
+  if (failed(getStridesAndOffset(memrefType, strides_, *offset)))
+    return mlirLogicalResultFailure();
+
+  (void)std::copy(strides_.begin(), strides_.end(), strides);
+  return mlirLogicalResultSuccess();
 }
 
 MlirTypeID mlirUnrankedMemRefTypeGetTypeID() {

mlir/python/mlir/dialects/_ods_common.py

@@ -2,16 +2,30 @@
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-# Provide a convenient name for sub-packages to resolve the main C-extension
-# with a relative import.
-from .._mlir_libs import _mlir as _cext
 from typing import (
+    List as _List,
+    Optional as _Optional,
     Sequence as _Sequence,
+    Tuple as _Tuple,
     Type as _Type,
     TypeVar as _TypeVar,
     Union as _Union,
 )
 
+from .._mlir_libs import _mlir as _cext
+from ..ir import (
+    ArrayAttr,
+    Attribute,
+    BoolAttr,
+    DenseI64ArrayAttr,
+    IntegerAttr,
+    IntegerType,
+    OpView,
+    Operation,
+    ShapedType,
+    Value,
+)
+
 __all__ = [
     "equally_sized_accessor",
     "get_default_loc_context",
@@ -138,3 +152,157 @@ def get_op_result_or_op_results(
 ResultValueTypeTuple = _cext.ir.Operation, _cext.ir.OpView, _cext.ir.Value
 ResultValueT = _Union[ResultValueTypeTuple]
 VariadicResultValueT = _Union[ResultValueT, _Sequence[ResultValueT]]
+
+StaticIntLike = _Union[int, IntegerAttr]
+ValueLike = _Union[Operation, OpView, Value]
+MixedInt = _Union[StaticIntLike, ValueLike]
+
+IntOrAttrList = _Sequence[_Union[IntegerAttr, int]]
+OptionalIntList = _Optional[_Union[ArrayAttr, IntOrAttrList]]
+
+BoolOrAttrList = _Sequence[_Union[BoolAttr, bool]]
+OptionalBoolList = _Optional[_Union[ArrayAttr, BoolOrAttrList]]
+
+MixedValues = _Union[_Sequence[_Union[StaticIntLike, ValueLike]], ArrayAttr, ValueLike]
+
+DynamicIndexList = _Sequence[_Union[MixedInt, _Sequence[MixedInt]]]
+
+
+def _dispatch_dynamic_index_list(
+    indices: _Union[DynamicIndexList, ArrayAttr],
+) -> _Tuple[_List[ValueLike], _Union[_List[int], ArrayAttr], _List[bool]]:
+    """Dispatches a list of indices to the appropriate form.
+
+    This is similar to the custom `DynamicIndexList` directive upstream:
+    provided indices may be in the form of dynamic SSA values or static values,
+    and they may be scalable (i.e., as a singleton list) or not. This function
+    dispatches each index into its respective form. It also extracts the SSA
+    values and static indices from various similar structures, respectively.
+    """
+    dynamic_indices = []
+    static_indices = [ShapedType.get_dynamic_size()] * len(indices)
+    scalable_indices = [False] * len(indices)
+
+    # ArrayAttr: Extract index values.
+    if isinstance(indices, ArrayAttr):
+        indices = [idx for idx in indices]
+
+    def process_nonscalable_index(i, index):
+        """Processes any form of non-scalable index.
+
+        Returns False if the given index was scalable and thus remains
+        unprocessed; True otherwise.
+        """
+        if isinstance(index, int):
+            static_indices[i] = index
+        elif isinstance(index, IntegerAttr):
+            static_indices[i] = index.value  # pytype: disable=attribute-error
+        elif isinstance(index, (Operation, Value, OpView)):
+            dynamic_indices.append(index)
+        else:
+            return False
+        return True
+
+    # Process each index at a time.
+    for i, index in enumerate(indices):
+        if not process_nonscalable_index(i, index):
+            # If it wasn't processed, it must be a scalable index, which is
+            # provided as a _Sequence of one value, so extract and process that.
+            scalable_indices[i] = True
+            assert len(index) == 1
+            ret = process_nonscalable_index(i, index[0])
+            assert ret
+
+    return dynamic_indices, static_indices, scalable_indices
+
+
+# Dispatches `MixedValues` that all represents integers in various forms into
+# the following three categories:
+#   - `dynamic_values`: a list of `Value`s, potentially from op results;
+#   - `packed_values`: a value handle, potentially from an op result, associated
+#                      to one or more payload operations of integer type;
+#   - `static_values`: an `ArrayAttr` of `i64`s with static values, from Python
+#                      `int`s, from `IntegerAttr`s, or from an `ArrayAttr`.
+# The input is in the form for `packed_values`, only that result is set and the
+# other two are empty. Otherwise, the input can be a mix of the other two forms,
+# and for each dynamic value, a special value is added to the `static_values`.
+def _dispatch_mixed_values(
+    values: MixedValues,
+) -> _Tuple[_List[Value], _Union[Operation, Value, OpView], DenseI64ArrayAttr]:
+    dynamic_values = []
+    packed_values = None
+    static_values = None
+    if isinstance(values, ArrayAttr):
+        static_values = values
+    elif isinstance(values, (Operation, Value, OpView)):
+        packed_values = values
+    else:
+        static_values = []
+        for size in values or []:
+            if isinstance(size, int):
+                static_values.append(size)
+            else:
+                static_values.append(ShapedType.get_dynamic_size())
+                dynamic_values.append(size)
+        static_values = DenseI64ArrayAttr.get(static_values)
+
+    return (dynamic_values, packed_values, static_values)
+
+
+def _get_value_or_attribute_value(
+    value_or_attr: _Union[any, Attribute, ArrayAttr]
+) -> any:
+    if isinstance(value_or_attr, Attribute) and hasattr(value_or_attr, "value"):
+        return value_or_attr.value
+    if isinstance(value_or_attr, ArrayAttr):
+        return _get_value_list(value_or_attr)
+    return value_or_attr
+
+
+def _get_value_list(
+    sequence_or_array_attr: _Union[_Sequence[any], ArrayAttr]
+) -> _Sequence[any]:
+    return [_get_value_or_attribute_value(v) for v in sequence_or_array_attr]
+
+
+def _get_int_array_attr(
+    values: _Optional[_Union[ArrayAttr, IntOrAttrList]]
+) -> ArrayAttr:
+    if values is None:
+        return None
+
+    # Turn into a Python list of Python ints.
+    values = _get_value_list(values)
+
+    # Make an ArrayAttr of IntegerAttrs out of it.
+    return ArrayAttr.get(
+        [IntegerAttr.get(IntegerType.get_signless(64), v) for v in values]
+    )
+
+
+def _get_int_array_array_attr(
+    values: _Optional[_Union[ArrayAttr, _Sequence[_Union[ArrayAttr, IntOrAttrList]]]]
+) -> ArrayAttr:
+    """Creates an ArrayAttr of ArrayAttrs of IntegerAttrs.
+
+    The input has to be a collection of a collection of integers, where any
+    Python _Sequence and ArrayAttr are admissible collections and Python ints and
+    any IntegerAttr are admissible integers. Both levels of collections are
+    turned into ArrayAttr; the inner level is turned into IntegerAttrs of i64s.
+    If the input is None, an empty ArrayAttr is returned.
+    """
+    if values is None:
+        return None
+
+    # Make sure the outer level is a list.
+    values = _get_value_list(values)
+
+    # The inner level is now either invalid or a mixed sequence of ArrayAttrs and
+    # Sequences. Make sure the nested values are all lists.
+    values = [_get_value_list(nested) for nested in values]
+
+    # Turn each nested list into an ArrayAttr.
+    values = [_get_int_array_attr(nested) for nested in values]
+
+    # Turn the outer list into an ArrayAttr.
+    return ArrayAttr.get(values)

mlir/python/mlir/dialects/memref.py

@@ -6,43 +6,52 @@
 from typing import Optional
 
 from ._memref_ops_gen import *
-from .arith import ConstantOp
-from .transform.structured import _dispatch_mixed_values, MixedValues
+from ._ods_common import _dispatch_mixed_values, MixedValues
+from .arith import ConstantOp, _is_integer_like_type
 from ..ir import Value, MemRefType, StridedLayoutAttr, ShapedType
 
 
-def _is_constant(i):
-    return isinstance(i, Value) and isinstance(i.owner.opview, ConstantOp)
+def _is_constant_int_like(i):
+    return (
+        isinstance(i, Value)
+        and isinstance(i.owner.opview, ConstantOp)
+        and _is_integer_like_type(i.type)
+    )
 
 
-def _is_static(i):
-    return (isinstance(i, int) and not ShapedType.is_dynamic_size(i)) or _is_constant(i)
+def _is_static_int_like(i):
+    return (
+        isinstance(i, int) and not ShapedType.is_dynamic_size(i)
+    ) or _is_constant_int_like(i)
 
 
 def _infer_memref_subview_result_type(
     source_memref_type, offsets, static_sizes, static_strides
 ):
-    source_strides, source_offset = source_memref_type.strides_and_offset
+    source_strides, source_offset = source_memref_type.get_strides_and_offset()
     # "canonicalize" from tuple|list -> list
     offsets, static_sizes, static_strides, source_strides = map(
         list, (offsets, static_sizes, static_strides, source_strides)
     )
 
-    assert all(
-        all(_is_static(i) for i in s)
+    if not all(
+        all(_is_static_int_like(i) for i in s)
         for s in [
             static_sizes,
             static_strides,
             source_strides,
         ]
-    ), f"Only inferring from python or mlir integer constant is supported"
+    ):
+        raise ValueError(
+            "Only inferring from python or mlir integer constant is supported."
+        )
 
     for s in [offsets, static_sizes, static_strides]:
         for idx, i in enumerate(s):
-            if _is_constant(i):
+            if _is_constant_int_like(i):
                 s[idx] = i.owner.opview.literal_value
 
-    if any(not _is_static(i) for i in offsets + [source_offset]):
+    if any(not _is_static_int_like(i) for i in offsets + [source_offset]):
         target_offset = ShapedType.get_dynamic_size()
     else:
         target_offset = source_offset
@@ -91,22 +100,22 @@ def subview(
         sizes = []
     if strides is None:
         strides = []
-    source_strides, source_offset = source.type.strides_and_offset
+    source_strides, source_offset = source.type.get_strides_and_offset()
     if result_type is None and all(
-        all(_is_static(i) for i in s) for s in [sizes, strides, source_strides]
+        all(_is_static_int_like(i) for i in s) for s in [sizes, strides, source_strides]
     ):
         # If any are arith.constant results then this will canonicalize to python int
-        # (which can then be used to fully specific the subview).
+        # (which can then be used to fully specify the subview).
         (
             offsets,
             sizes,
             strides,
             result_type,
         ) = _infer_memref_subview_result_type(source.type, offsets, sizes, strides)
-    else:
-        assert (
-            result_type is not None
-        ), "mixed static/dynamic offset/sizes/strides requires explicit result type"
+    elif result_type is None:
+        raise ValueError(
+            "mixed static/dynamic offset/sizes/strides requires explicit result type."
+        )
 
     offsets, _packed_offsets, static_offsets = _dispatch_mixed_values(offsets)
     sizes, _packed_sizes, static_sizes = _dispatch_mixed_values(sizes)