fix striding calcs for memref (depends on llvm/llvm-project#79393)

Author: makslevental

mlir/extras/dialects/ext/arith.py

@@ -5,17 +5,20 @@
 from typing import Union, Optional, Tuple
 
 import numpy as np
+
+from ...util import get_user_code_loc, infer_mlir_type, mlir_type_to_np_dtype
+from ...._mlir_libs._mlir import register_value_caster
 from ....dialects import arith as arith_dialect
-from ....dialects.arith import *
 from ....dialects import complex as complex_dialect
 from ....dialects._arith_enum_gen import (
     _arith_cmpfpredicateattr,
     CmpFPredicate,
     CmpIPredicate,
     _arith_cmpipredicateattr,
 )
-from ....dialects.arith import _is_integer_like_type
 from ....dialects._ods_common import get_op_result_or_value, get_op_result_or_op_results
+from ....dialects.arith import *
+from ....dialects.arith import _is_integer_like_type
 from ....dialects.linalg.opdsl.lang.emitter import (
     _is_floating_point_type,
     _is_integer_type,
@@ -44,9 +47,6 @@
     FloatAttr,
 )
 
-from ...util import get_user_code_loc, infer_mlir_type, mlir_type_to_np_dtype
-from ...._mlir_libs._mlir import register_value_caster
-
 
 def constant(
     value: Union[int, float, bool, np.ndarray],

mlir/extras/dialects/ext/cf.py

@@ -1,16 +1,12 @@
 from typing import Union, List
 
-from ....dialects.cf import *
+from ...util import get_user_code_loc, Successor
 from ....dialects._cf_ops_gen import _Dialect
 from ....dialects._ods_common import (
-    get_op_result_or_value,
-    get_op_results_or_values,
-    get_default_loc_context,
-    segmented_accessor,
     _cext,
 )
-from ....ir import Value, InsertionPoint, Block, OpView
-from ...util import get_user_code_loc, Successor
+from ....dialects.cf import *
+from ....ir import Value, InsertionPoint, Block
 
 
 @_cext.register_operation(_Dialect, replace=True)

mlir/extras/dialects/ext/func.py

@@ -6,12 +6,10 @@
 from ...meta import op_region_builder
 from ...util import get_user_code_loc, make_maybe_no_args_decorator
 from ....dialects.func import *
-from ....extras import types as T
 from ....ir import (
     FlatSymbolRefAttr,
     FunctionType,
     InsertionPoint,
-    StringAttr,
     Type,
     TypeAttr,
     Value,

mlir/extras/dialects/ext/gpu.py

@@ -2,9 +2,17 @@
 from functools import partial
 from typing import Optional, Any, List, Tuple
 
+from .arith import constant
+from .func import FuncBase
+from ... import types as T
+from ...meta import (
+    region_op,
+)
+from ...util import get_user_code_loc, make_maybe_no_args_decorator, ModuleMeta
+from ....dialects._gpu_ops_gen import _Dialect
 from ....dialects._ods_common import get_default_loc_context, _cext
+from ....dialects._ods_common import get_op_result_or_op_results
 from ....dialects.gpu import *
-from ....dialects._gpu_ops_gen import _Dialect
 from ....ir import (
     Type,
     Attribute,
@@ -17,15 +25,6 @@
     Value,
 )
 
-from ... import types as T
-from .arith import constant
-from .func import FuncBase
-from ...meta import (
-    region_op,
-)
-from ....dialects._ods_common import get_op_result_or_op_results
-from ...util import get_user_code_loc, make_maybe_no_args_decorator, ModuleMeta
-
 
 def block_id_x():
     return block_id("x")

mlir/extras/dialects/ext/linalg.py

@@ -1,5 +1,5 @@
-from ....dialects import linalg
 from ...util import get_user_code_loc
+from ....dialects import linalg
 
 
 def abs(I, O, *, loc=None, ip=None):

mlir/extras/dialects/ext/llvm.py

@@ -1,5 +1,5 @@
-from ....ir import Type
 from ....dialects.llvm import *
+from ....ir import Type
 
 
 def llvm_ptr_t():

mlir/extras/dialects/ext/memref.py

@@ -1,18 +1,16 @@
-import re
 from functools import cached_property, reduce
-from typing import Tuple, Sequence, Optional, Union
+from typing import Tuple, Sequence, Union
 
-from ....ir import Type, Value, MemRefType, ShapedType, MLIRError
-
-from ... import types as T
-from ....dialects.memref import *
-from ....dialects import memref, arith
 from .arith import Scalar, constant
 from .tensor import _indices_to_indexer, compute_result_shape_reassoc_list
+from ... import types as T
 from ...meta import region_op
-from ...._mlir_libs._mlir import register_value_caster
 from ...util import get_user_code_loc
+from ...._mlir_libs._mlir import register_value_caster
+from ....dialects import memref, arith
 from ....dialects._ods_common import get_op_result_or_op_results
+from ....dialects.memref import *
+from ....ir import Type, Value, MemRefType, ShapedType
 
 S = ShapedType.get_dynamic_size()
 
@@ -70,71 +68,6 @@ def store(
     return get_op_result_or_op_results(StoreOp(value, mem, indices, loc=loc, ip=ip))
 
 
-def subview(
-    source: "MemRef",
-    offsets: Optional[Sequence[Value]] = None,
-    strides: Optional[Sequence[Value]] = None,
-    static_offsets: Optional[Sequence[int]] = None,
-    static_sizes: Optional[Sequence[int]] = None,
-    static_strides: Optional[Sequence[int]] = None,
-    *,
-    loc=None,
-    ip=None,
-):
-    if loc is None:
-        loc = get_user_code_loc()
-    if offsets is None:
-        offsets = []
-    if static_offsets is None:
-        static_offsets = []
-    if strides is None:
-        strides = []
-    if static_strides is None:
-        static_strides = []
-    assert static_sizes, f"this convenience method only handles static sizes"
-    sizes = []
-    wrong_type = T.memref(*static_sizes, source.dtype)
-    if offsets and static_offsets:
-        assert all(s == S for s in static_offsets)
-    if strides and static_strides:
-        assert all(s == S for s in static_strides)
-    val = memref.subview(
-        wrong_type,
-        source,
-        offsets,
-        sizes,
-        strides,
-        static_offsets,
-        static_sizes,
-        static_strides,
-        loc=loc,
-        ip=ip,
-    )
-    # dumbest hack ever - the default builder doesn't connect to inferReturnTypes
-    # but the diag message does
-    try:
-        val.owner.verify()
-        return val
-    except MLIRError as e:
-        diag = str(e.error_diagnostics[0])
-        correct_type = re.findall(r"'memref<(.*)>'", diag)
-        assert len(correct_type) == 1
-        correct_type = Type.parse(f"memref<{correct_type[0]}>")
-        val.owner.erase()
-        return memref.subview(
-            correct_type,
-            source,
-            offsets,
-            sizes,
-            strides,
-            static_offsets,
-            static_sizes,
-            static_strides,
-            loc=loc,
-            ip=ip,
-        )
-
-
 @register_value_caster(MemRefType.static_typeid)
 class MemRef(Value):
     def __str__(self):
@@ -266,16 +199,15 @@ def _subview(
     if indexer.is_constant():
         out = subview(
             out,
-            static_offsets=indexer.static_offsets(),
-            static_sizes=indexer.static_sizes(),
-            static_strides=indexer.static_strides(),
+            offsets=indexer.static_offsets(),
+            sizes=indexer.static_sizes(),
+            strides=indexer.static_strides(),
             loc=loc,
             ip=ip,
         )
     else:
         # special tile case
         offsets = [None] * len(indexer.in_shape)
-        static_offsets = [None] * len(indexer.in_shape)
         static_sizes = [None] * len(indexer.in_shape)
         static_strides = [None] * len(indexer.in_shape)
         for i, ind in enumerate(indexer.indices):
@@ -292,15 +224,13 @@ def _subview(
                 and ind.step.is_constant()
             ):
                 offsets[i] = ind.start
-                static_offsets[i] = S
                 static_sizes[i] = maybe_size.literal_value
                 static_strides[i] = (
                     ind.step.literal_value if isinstance(ind.step, Scalar) else ind.step
                 )
             else:
                 raise RuntimeError(f"indexing not supported {indexer.indices}")
         offsets = list(filter(None, offsets))
-        static_offsets = list(filter(None, static_offsets))
         static_sizes = list(filter(None, static_sizes))
         static_strides = list(filter(None, static_strides))
         assert (
@@ -312,9 +242,8 @@ def _subview(
         out = subview(
             out,
             offsets=offsets,
-            static_offsets=static_offsets,
-            static_sizes=static_sizes,
-            static_strides=static_strides,
+            sizes=static_sizes,
+            strides=static_strides,
             loc=loc,
             ip=ip,
         )

mlir/extras/dialects/ext/nvgpu.py

@@ -6,7 +6,7 @@
     TensorMapOOBKind,
     TensorMapInterleaveKind,
 )
-from ....ir import Attribute, Type
+from ....ir import Type
 
 
 def tensormap_descriptor(

mlir/extras/testing/testing.py

@@ -12,9 +12,12 @@
 from ..context import MLIRContext, mlir_mod_ctx
 from .generate_test_checks import main
 from ..runtime.refbackend import LLVMJITBackend
+from ...ir import Module
 
 
 def filecheck(correct: str, module):
+    if isinstance(module, Module):
+        assert module.operation.verify()
     filecheck_name = "FileCheck"
     if platform.system() == "Windows":
         filecheck_name += ".exe"