add test for mem[v:]

Author: makslevental

mlir/extras/dialects/ext/_shaped_value.py

@@ -163,14 +163,13 @@ def _indices_to_indexer(
         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:
+            if 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)
+                    isinstance(stop, int)
                     and in_shape[in_axis] != ShapedType.get_dynamic_size()
                     and stop >= in_shape[in_axis]
                 ):
@@ -205,6 +204,7 @@ def _indices_to_indexer(
                     step = 1
                 if stop is None:
                     stop = in_shape[in_axis]
+
                 indices[in_axis] = slice(start, stop, step)
 
                 out_axis += 1
@@ -307,20 +307,55 @@ def _has_index_type(e: Any) -> bool:
 
 def _is_constant_index(e: Any) -> bool:
     return (
-        isinstance(e, Scalar)
-        and e.is_constant()
+        (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)
+        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()
+        (isinstance(e, Scalar) and e.is_constant())
         or (isinstance(e, (int, float, bool)) and e != ShapedType.get_dynamic_size())
         or e is None
     )
+
+
+def _maybe_compute_size(start, stop, step):
+    from ....dialects import arith
+
+    # TODO(max): figure out how to use actual canonicalizers
+    if (
+        isinstance(start, Value)
+        and isinstance(stop, Value)
+        and stop.owner.operands[0]._eq(start)
+        and stop.owner.operands[1].is_constant()
+    ):
+        return stop.owner.operands[1].literal_value
+    elif (
+        isinstance(start, Value)
+        and isinstance(start.owner.opview, arith.MulIOp)
+        and isinstance(stop, Value)
+        and isinstance(stop.owner.opview, arith.MulIOp)
+        and isinstance(stop.owner.operands[0].owner.opview, arith.AddIOp)
+        and start.owner.operands[0] == stop.owner.operands[0].owner.operands[0]
+        and stop.owner.operands[1].is_constant()
+        and isinstance(step, int)
+        or (isinstance(step, Scalar) and step.is_constant())
+    ):
+        # looks like this
+        # l = lambda l: l * D
+        # r = lambda r: (r + 1) * D
+        # a, b, c = (
+        #     A[l(i) : r(i), l(j) : r(j)],
+        #     B[l(i) : r(i), l(j) : r(j)],
+        #     C[l(i) : r(i), l(j) : r(j)],
+        # )
+        return stop.owner.operands[1]
+    else:
+        return stop - start

mlir/extras/dialects/ext/memref.py

@@ -5,7 +5,7 @@
 
 import numpy as np
 
-from ._shaped_value import ShapedValue, _indices_to_indexer
+from ._shaped_value import ShapedValue, _indices_to_indexer, _maybe_compute_size
 from .arith import Scalar, constant, index_cast
 from .tensor import compute_result_shape_reassoc_list
 from .vector import Vector
@@ -26,8 +26,8 @@
 )
 from ....dialects.memref import (
     _is_static_int_like,
-    _infer_memref_subview_result_type,
     _generated_subview,
+    _is_constant_int_like,
 )
 from ....dialects.memref import *
 from ....ir import (
@@ -283,37 +283,44 @@ def expand_shape(
     )
 
 
-def _maybe_compute_size(start, stop, step):
-    # TODO(max): figure out how to use actual canonicalizers
-    if (
-        isinstance(start, Value)
-        and isinstance(stop, Value)
-        and stop.owner.operands[0]._eq(start)
-        and stop.owner.operands[1].is_constant()
-    ):
-        return stop.owner.operands[1].literal_value
-    elif (
-        isinstance(start, Value)
-        and isinstance(start.owner.opview, arith.MulIOp)
-        and isinstance(stop, Value)
-        and isinstance(stop.owner.opview, arith.MulIOp)
-        and isinstance(stop.owner.operands[0].owner.opview, arith.AddIOp)
-        and start.owner.operands[0] == stop.owner.operands[0].owner.operands[0]
-        and stop.owner.operands[1].is_constant()
-        and isinstance(step, int)
-        or (isinstance(step, Scalar) and step.is_constant())
-    ):
-        # looks like this
-        # l = lambda l: l * D
-        # r = lambda r: (r + 1) * D
-        # a, b, c = (
-        #     A[l(i) : r(i), l(j) : r(j)],
-        #     B[l(i) : r(i), l(j) : r(j)],
-        #     C[l(i) : r(i), l(j) : r(j)],
-        # )
-        return stop.owner.operands[1]
+def _infer_memref_subview_result_type(source_memref_type, offsets, sizes, strides):
+    source_strides, source_offset = source_memref_type.get_strides_and_offset()
+    # "canonicalize" from tuple|list -> list
+    offsets, sizes, strides, source_strides = map(
+        list, (offsets, sizes, strides, source_strides)
+    )
+
+    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
+        for offset, target_stride in zip(offsets, source_strides):
+            target_offset += offset * target_stride
+
+    target_strides = []
+    for source_stride, static_stride in zip(source_strides, strides):
+        target_strides.append(source_stride * static_stride)
+
+    if all(isinstance(s, int) for s in sizes):
+        # If default striding then no need to complicate things for downstream ops (e.g., expand_shape).
+        default_strides = list(accumulate(sizes[1:][::-1], operator.mul))[::-1] + [1]
     else:
-        return stop - start
+        default_strides = None
+        sizes = [
+            s if isinstance(s, int) else ShapedType.get_dynamic_size() for s in sizes
+        ]
+
+    if target_strides == default_strides and target_offset == 0:
+        layout = None
+    else:
+        layout = StridedLayoutAttr.get(target_offset, target_strides)
+
+    return MemRefType.get(
+        sizes,
+        source_memref_type.element_type,
+        layout,
+        source_memref_type.memory_space,
+    )
 
 
 def subview(
@@ -333,21 +340,17 @@ def subview(
         sizes = []
     if strides is None:
         strides = []
-    source_strides, source_offset = source.type.get_strides_and_offset()
-    if result_type is None and all(
-        all(_is_static_int_like(i) for i in s) for s in [sizes, strides, source_strides]
-    ):
+
+    for s in [offsets, sizes, strides]:
+        for idx, i in enumerate(s):
+            if _is_constant_int_like(i):
+                s[idx] = i.owner.opview.literal_value
+
+    if result_type is None:
         # If any are arith.constant results then this will canonicalize to python int
         # (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)
-    elif result_type is None:
-        raise ValueError(
-            "mixed static/dynamic offset/sizes/strides requires explicit result type."
+        result_type = _infer_memref_subview_result_type(
+            source.type, offsets, sizes, strides
         )
 
     offsets, _packed_offsets, static_offsets = _dispatch_mixed_values(offsets)

tests/test_gpu.py

@@ -1088,7 +1088,7 @@ def test_amdgpu_bank_conflicts(ctx: MLIRContext):
 
     set_container_module(ctx.module)
 
-    M = 1024
+    M = 128
 
     @gpu_func
     def no_bank_conflicts(A: T.memref(M, M, T.f32()), B: T.memref(M, M, T.f32())):

tests/test_memref.py

@@ -62,12 +62,21 @@ def test_simple_slicing(ctx: MLIRContext):
     w = mem[5:]
     w = mem[:5]
 
+    one = constant(1, index=True) * 2
+    w = mem[one:]
+
     correct = dedent(
         """\
     module {
       %alloc = memref.alloc() : memref<10xi32>
       %subview = memref.subview %alloc[5] [5] [1] : memref<10xi32> to memref<5xi32, strided<[1], offset: 5>>
       %subview_0 = memref.subview %alloc[0] [5] [1] : memref<10xi32> to memref<5xi32>
+      %c1 = arith.constant 1 : index
+      %c2 = arith.constant 2 : index
+      %0 = arith.muli %c1, %c2 : index
+      %c10 = arith.constant 10 : index
+      %1 = arith.subi %c10, %0 : index
+      %subview_1 = memref.subview %alloc[%0] [%1] [1] : memref<10xi32> to memref<?xi32, strided<[1], offset: ?>>
     }
     """
     )
@@ -146,14 +155,7 @@ def test_ellipsis_and_full_slice_plus_coordinate_1(ctx: MLIRContext):
 
     one = constant(1, index=True) * 2
     w = mem[one, :, :, ...]
-
-    try:
-        w = mem[1, :, :, :, :]
-    except IndexError as e:
-        assert (
-            str(e)
-            == "Too many indices for shaped type with rank: 5 non-None/Ellipsis indices for dim 4."
-        )
+    w = mem[one:, :, :, ...]
 
     correct = dedent(
         f"""\
@@ -169,12 +171,23 @@ def test_ellipsis_and_full_slice_plus_coordinate_1(ctx: MLIRContext):
       %c2 = arith.constant 2 : index
       %0 = arith.muli %c1_4, %c2 : index
       %subview_5 = memref.subview %alloc[%0, 0, 0, 0] [1, 22, 333, 4444] [1, 1, 1, 1] : memref<10x22x333x4444xi32> to memref<1x22x333x4444xi32, strided<{golden_w_3_strides}, offset: ?>>
-      %c1_6 = arith.constant 1 : index
+      %c10 = arith.constant 10 : index
+      %1 = arith.subi %c10, %0 : index
+      %subview_6 = memref.subview %alloc[%0, 0, 0, 0] [%1, 22, 333, 4444] [1, 1, 1, 1] : memref<10x22x333x4444xi32> to memref<?x22x333x4444xi32, strided<{golden_w_3_strides}, offset: ?>>
     }}
     """
     )
     filecheck(correct, ctx.module)
 
+    try:
+        w = mem[1, :, :, :, :]
+    except IndexError as e:
+        assert (
+                str(e)
+                == "Too many indices for shaped type with rank: 5 non-None/Ellipsis indices for dim 4."
+        )
+
+
 
 def test_ellipsis_and_full_slice_plus_coordinate_2(ctx: MLIRContext):
     sizes = (10, 22, 333, 4444)