tiling examples

Author: makslevental

.github/workflows/test.yml

@@ -49,7 +49,7 @@ jobs:
           if [ ${{ matrix.os }} == 'windows-2022' ]; then
             pytest -s --ignore-glob=*test_smoke* tests
           else
-            pytest -s --ignore-glob=*test_smoke* tests
+            pytest --capture=tee-sys --ignore-glob=*test_smoke* tests
           fi
 
   test-against-torch-mlir-bindings:
@@ -86,9 +86,9 @@ jobs:
         shell: bash
         run: |
           if [ ${{ matrix.os }} == 'windows-2022' ]; then
-            pytest -s tests/test_smoke.py
+            pytest --capture=tee-sys tests/test_smoke.py
           else
-            pytest -s tests/test_smoke.py
+            pytest --capture=tee-sys tests/test_smoke.py
           fi
 
 
@@ -126,7 +126,7 @@ jobs:
         shell: bash
         run: |
           if [ ${{ matrix.os }} == 'windows-2022' ]; then
-            pytest -s tests/test_smoke.py
+            pytest --capture=tee-sys tests/test_smoke.py
           else
-            pytest -s tests/test_smoke.py
+            pytest --capture=tee-sys tests/test_smoke.py
           fi

mlir_utils/ast/canonicalize.py

@@ -75,7 +75,8 @@ def transform_ast(
     n_lines = len(inspect.getsource(f).splitlines())
     line_starts = list(findlinestarts(new_f_code_o))
     assert (
-        line_starts[-1][1] - line_starts[0][1] == n_lines - 1
+        max([l for _, l in line_starts]) - min([l for _, l in line_starts]) + 1
+        == n_lines
     ), f"something went wrong with the line numbers for the rewritten/canonicalized function"
     return copy_func(f, new_f_code_o)
 

mlir_utils/dialects/ext/func.py

@@ -247,7 +247,7 @@ def func(
 ) -> FuncBase:
     if loc is None:
         loc = get_user_code_loc()
-    return FuncBase(
+    func = FuncBase(
         body_builder=f,
         func_op_ctor=FuncOp.__base__,
         return_op_ctor=ReturnOp,
@@ -259,3 +259,5 @@ def func(
         loc=loc,
         ip=ip,
     )
+    func.__name__ = f.__name__
+    return func

mlir_utils/dialects/ext/memref.py

@@ -4,13 +4,13 @@
 
 from mlir.ir import Type, Value, MemRefType, ShapedType, MLIRError
 
-from mlir_utils.dialects import memref
+import mlir_utils.types as T
+from mlir_utils.dialects import memref, arith
 from mlir_utils.dialects.ext.arith import Scalar, constant
 from mlir_utils.dialects.ext.tensor import (
     _indices_to_indexer,
     compute_result_shape_reassoc_list,
 )
-import mlir_utils.types as T
 from mlir_utils.util import (
     register_value_caster,
     get_user_code_loc,
@@ -88,6 +88,8 @@ def store(
 
 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,
@@ -97,11 +99,23 @@ def subview(
 ):
     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"
-    offsets = sizes = strides = []
-    result = T.memref(*static_sizes, source.dtype)
+    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(
-        result,
+        wrong_type,
         source,
         offsets,
         sizes,
@@ -270,7 +284,51 @@ def _subview(
             ip=ip,
         )
     else:
-        raise ValueError(f"non-constant indices not supported {indexer}")
+        # 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):
+            maybe_size = maybe_cast(ind.stop.owner.operands[1])
+            if (
+                isinstance(ind.start.owner.opview, arith.MulIOp)
+                and isinstance(ind.stop.owner.opview, arith.MulIOp)
+                and isinstance(ind.stop.owner.operands[0].owner.opview, arith.AddIOp)
+                and ind.start.owner.operands[0]
+                == ind.stop.owner.operands[0].owner.operands[0]
+                and maybe_size.is_constant()
+                and isinstance(ind.step, int)
+                or isinstance(ind.step, Scalar)
+                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 (
+            len(offsets)
+            == len(static_sizes)
+            == len(static_strides)
+            == len(indexer.in_shape)
+        ), f"not each slice is statically known: {indexer.indices}"
+        out = subview(
+            out,
+            offsets=offsets,
+            static_offsets=static_offsets,
+            static_sizes=static_sizes,
+            static_strides=static_strides,
+            loc=loc,
+            ip=ip,
+        )
 
     # This adds newaxis/None dimensions.
     return expand_shape(out, indexer.newaxis_dims, loc=loc, ip=ip)

mlir_utils/dialects/ext/tensor.py

@@ -268,10 +268,14 @@ def static_sizes(self):
                 sizes.append(1)
             elif isinstance(i, slice):
                 start, stop, step = map(int, (i.start, i.stop, i.step))
-                s = ((stop - start) // step) + 1
-                if (stop - start) % step == 0:
-                    s -= 1
-                sizes.append(s)
+                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)
@@ -496,12 +500,12 @@ 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 step == 1:
+            if step is None or isinstance(step, int) and step == 1:
                 step = None
             if step is None:
-                if start is None or start == 0:
+                if start is None or isinstance(start, int) and start == 0:
                     start = None
-                if stop is None or stop >= in_shape[in_axis]:
+                if stop is None or isinstance(stop, int) and stop >= in_shape[in_axis]:
                     stop = None
             # Handle slice(None) and slice(None, None, -1)
             if (
@@ -529,6 +533,8 @@ def _indices_to_indexer(
                     )
                     raise IndexError(msg)
 
+                if step is None:
+                    step = 1
                 indices[in_axis] = slice(start, stop, step)
 
                 out_axis += 1

mlir_utils/runtime/refbackend.py

@@ -1,25 +1,35 @@
 import ctypes
 import logging
 import os
+import platform
 import warnings
 from pathlib import Path
 
 import numpy as np
 from mlir import _mlir_libs
 from mlir.dialects.func import FuncOp, CallOp
-from mlir.execution_engine import ExecutionEngine
 from mlir.ir import UnitAttr, Module, MemRefType
-from mlir.runtime import (
-    UnrankedMemRefDescriptor,
-    get_ranked_memref_descriptor,
-    unranked_memref_to_numpy,
-    get_unranked_memref_descriptor,
-)
 
+try:
+    from mlir.execution_engine import ExecutionEngine
+    from mlir.runtime import (
+        UnrankedMemRefDescriptor,
+        get_ranked_memref_descriptor,
+        unranked_memref_to_numpy,
+        get_unranked_memref_descriptor,
+    )
+except:
+    warnings.warn("no execution engine in mlir bindings; refbackend won't work")
+
+
+import mlir_utils.types as T
 from mlir_utils.dialects.memref import cast
 from mlir_utils.runtime.passes import Pipeline, run_pipeline
-from mlir_utils.types import memref_type_to_np_dtype, mlir_type_to_ctype, np_dtype_to_mlir_type
-import mlir_utils.types as T
+from mlir_utils.types import (
+    memref_type_to_np_dtype,
+    mlir_type_to_ctype,
+    np_dtype_to_mlir_type,
+)
 from mlir_utils.util import shlib_ext, find_ops, shlib_prefix
 
 logger = logging.getLogger(__name__)
@@ -86,8 +96,8 @@ def convert_arg_to_ctype(arg, unranked=True):
     if isinstance(arg, CData) or isinstance(arg, (int, float, bool)):
         return arg
     elif isinstance(arg, np.ndarray):
-        assert (
-            np_dtype_to_mlir_type(arg.dtype.type)
+        assert np_dtype_to_mlir_type(
+            arg.dtype.type
         ), f"unsupported numpy array type {arg.dtype}"
         if unranked:
             return ctypes.pointer(ctypes.pointer(get_unranked_memref_descriptor(arg)))
@@ -145,6 +155,8 @@ def invoke(*args):
             self.ee.invoke(
                 function_name, *[convert_arg_to_ctype(a, unranked=False) for a in args]
             )
+            if self.results is not None and len(self.results) == 1:
+                return self.results[0]
             return self.results
 
         return invoke
@@ -193,7 +205,7 @@ def __init__(
         self,
         shared_lib_paths=None,
     ):
-        if shared_lib_paths is None:
+        if shared_lib_paths is None and platform.system() != "Windows":
             shared_lib_paths = [
                 ASYNC_RUNTIME_LIB_PATH,
                 C_RUNNER_UTILS_LIB_PATH,

mlir_utils/types.py

@@ -1,4 +1,3 @@
-import sys
 import ctypes
 from functools import partial
 from typing import Union
@@ -24,6 +23,7 @@
     UnrankedMemRefType,
     UnrankedTensorType,
     VectorType,
+    StridedLayoutAttr,
 )
 
 _index = lambda: IndexType.get()
@@ -228,9 +228,17 @@ def tensor(*args, element_type: Type = None):
         )
 
 
-def memref(*args, element_type: Type = None, memory_space: int = None):
+def memref(
+    *args,
+    element_type: Type = None,
+    memory_space: int = None,
+    layout: tuple[tuple[int, ...], int] = None,
+):
     if memory_space is None:
         memory_space = 0
+    if layout is not None:
+        strides, offset = layout
+        layout = StridedLayoutAttr.get(offset, strides)
     memory_space = Attribute.parse(str(memory_space))
     if not len(args) or len(args) == 1 and isinstance(args[-1], Type):
         return shaped(
@@ -242,7 +250,9 @@ def memref(*args, element_type: Type = None, memory_space: int = None):
         return shaped(
             *args,
             element_type=element_type,
-            type_constructor=partial(MemRefType.get, memory_space=memory_space),
+            type_constructor=partial(
+                MemRefType.get, memory_space=memory_space, layout=layout
+            ),
         )
 
 

pyproject.toml

@@ -11,7 +11,6 @@ dependencies = [
     "PyYAML",
     "inflection",
     "bytecode",
-    "libcst",
 ]
 
 [project.optional-dependencies]