test memref<vector> (#131)

Author: makslevental

.github/workflows/test.yml

@@ -66,15 +66,13 @@ jobs:
       - name: Test
         shell: bash
         run: |
-          if [ ${{ matrix.os }} == 'windows-2022' ]; then
-            pytest -s tests
-          else
-            pytest --capture=tee-sys tests
-          fi
+          
+          pytest tests
 
       - name: Test mwe
         shell: bash
         run: |
+          
           python examples/mwe.py
 
   test-other-host-bindings:

mlir/extras/dialects/ext/arith.py

@@ -6,6 +6,7 @@
 from functools import cached_property, partialmethod
 from typing import Optional, Tuple, Union
 
+import numpy as np
 from bytecode import ConcreteBytecode
 from einspect.structs import PyTypeObject
 
@@ -124,18 +125,21 @@ def constant(
 
 
 def index_cast(
-    value: Value,
+    in_: Value,
     *,
     to: Type = None,
+    out: Type = None,
     loc: Location = None,
     ip: InsertionPoint = None,
 ) -> Value:
     if loc is None:
         loc = get_user_code_loc()
-    if to is None:
-        to = IndexType.get()
+    assert bool(to) != bool(out), "either `to` or `out` must be set"
+    res_type = out or to
+    if res_type is None:
+        res_type = IndexType.get()
     return get_op_result_or_op_results(
-        arith_dialect.IndexCastOp(to, value, loc=loc, ip=ip)
+        arith_dialect.IndexCastOp(res_type, in_, loc=loc, ip=ip)
     )
 
 
@@ -234,26 +238,26 @@ def __call__(cls, *args, **kwargs):
 @register_attribute_builder("Arith_CmpIPredicateAttr", replace=True)
 def _arith_CmpIPredicateAttr(predicate: Union[str, Attribute], context: Context):
     predicates = {
-        "eq": CmpIPredicate.eq,
-        "ne": CmpIPredicate.ne,
-        "slt": CmpIPredicate.slt,
-        "sle": CmpIPredicate.sle,
-        "sgt": CmpIPredicate.sgt,
-        "sge": CmpIPredicate.sge,
-        "ult": CmpIPredicate.ult,
-        "ule": CmpIPredicate.ule,
-        "ugt": CmpIPredicate.ugt,
-        "uge": CmpIPredicate.uge,
-        0: CmpIPredicate.eq,
-        1: CmpIPredicate.ne,
-        2: CmpIPredicate.slt,
-        3: CmpIPredicate.sle,
-        4: CmpIPredicate.sgt,
-        5: CmpIPredicate.sge,
-        6: CmpIPredicate.ult,
-        7: CmpIPredicate.ule,
-        8: CmpIPredicate.ugt,
-        9: CmpIPredicate.uge,
+        "eq": arith_dialect.CmpIPredicate.eq,
+        "ne": arith_dialect.CmpIPredicate.ne,
+        "slt": arith_dialect.CmpIPredicate.slt,
+        "sle": arith_dialect.CmpIPredicate.sle,
+        "sgt": arith_dialect.CmpIPredicate.sgt,
+        "sge": arith_dialect.CmpIPredicate.sge,
+        "ult": arith_dialect.CmpIPredicate.ult,
+        "ule": arith_dialect.CmpIPredicate.ule,
+        "ugt": arith_dialect.CmpIPredicate.ugt,
+        "uge": arith_dialect.CmpIPredicate.uge,
+        0: arith_dialect.CmpIPredicate.eq,
+        1: arith_dialect.CmpIPredicate.ne,
+        2: arith_dialect.CmpIPredicate.slt,
+        3: arith_dialect.CmpIPredicate.sle,
+        4: arith_dialect.CmpIPredicate.sgt,
+        5: arith_dialect.CmpIPredicate.sge,
+        6: arith_dialect.CmpIPredicate.ult,
+        7: arith_dialect.CmpIPredicate.ule,
+        8: arith_dialect.CmpIPredicate.ugt,
+        9: arith_dialect.CmpIPredicate.uge,
     }
     if isinstance(predicate, Attribute):
         return predicate
@@ -264,29 +268,29 @@ def _arith_CmpIPredicateAttr(predicate: Union[str, Attribute], context: Context)
 @register_attribute_builder("Arith_CmpFPredicateAttr", replace=True)
 def _arith_CmpFPredicateAttr(predicate: Union[str, Attribute], context: Context):
     predicates = {
-        "false": CmpFPredicate.AlwaysFalse,
+        "false": arith_dialect.CmpFPredicate.AlwaysFalse,
         # ordered comparison
         # An ordered comparison checks if neither operand is NaN.
-        "oeq": CmpFPredicate.OEQ,
-        "ogt": CmpFPredicate.OGT,
-        "oge": CmpFPredicate.OGE,
-        "olt": CmpFPredicate.OLT,
-        "ole": CmpFPredicate.OLE,
-        "one": CmpFPredicate.ONE,
+        "oeq": arith_dialect.CmpFPredicate.OEQ,
+        "ogt": arith_dialect.CmpFPredicate.OGT,
+        "oge": arith_dialect.CmpFPredicate.OGE,
+        "olt": arith_dialect.CmpFPredicate.OLT,
+        "ole": arith_dialect.CmpFPredicate.OLE,
+        "one": arith_dialect.CmpFPredicate.ONE,
         # no clue what this one is
-        "ord": CmpFPredicate.ORD,
+        "ord": arith_dialect.CmpFPredicate.ORD,
         # unordered comparison
         # Conversely, an unordered comparison checks if either operand is a NaN.
-        "ueq": CmpFPredicate.UEQ,
-        "ugt": CmpFPredicate.UGT,
-        "uge": CmpFPredicate.UGE,
-        "ult": CmpFPredicate.ULT,
-        "ule": CmpFPredicate.ULE,
-        "une": CmpFPredicate.UNE,
+        "ueq": arith_dialect.CmpFPredicate.UEQ,
+        "ugt": arith_dialect.CmpFPredicate.UGT,
+        "uge": arith_dialect.CmpFPredicate.UGE,
+        "ult": arith_dialect.CmpFPredicate.ULT,
+        "ule": arith_dialect.CmpFPredicate.ULE,
+        "une": arith_dialect.CmpFPredicate.UNE,
         # no clue what this one is
-        "uno": CmpFPredicate.UNO,
+        "uno": arith_dialect.CmpFPredicate.UNO,
         # return always true
-        "true": CmpFPredicate.AlwaysTrue,
+        "true": arith_dialect.CmpFPredicate.AlwaysTrue,
     }
     if isinstance(predicate, Attribute):
         return predicate

mlir/extras/dialects/ext/func.py

@@ -24,6 +24,9 @@
 )
 
 
+_call = call
+
+
 def call(
     callee_or_results: Union[FuncOp, List[Type]],
     arguments_or_callee: Union[List[Value], FlatSymbolRefAttr, str],

mlir/extras/dialects/ext/memref.py

@@ -6,6 +6,7 @@
 from ._shaped_value import ShapedValue
 from .arith import Scalar, constant
 from .tensor import _indices_to_indexer, compute_result_shape_reassoc_list
+from .vector import Vector
 from ... import types as T
 from ...meta import region_op
 from ...util import (
@@ -28,7 +29,7 @@
 S = ShapedType.get_dynamic_size()
 
 
-def _alloc(
+def __alloc(
     op_ctor,
     sizes: Sequence[Union[int, Value]],
     element_type: Type,
@@ -64,6 +65,9 @@ def _alloc(
     )
 
 
+_alloc = alloc
+
+
 def alloc(
     sizes: Union[int, Value],
     element_type: Type = None,
@@ -74,7 +78,7 @@ def alloc(
 ):
     if loc is None:
         loc = get_user_code_loc()
-    return _alloc(
+    return __alloc(
         AllocOp,
         sizes,
         element_type,
@@ -85,6 +89,9 @@ def alloc(
     )
 
 
+_alloca = alloca
+
+
 def alloca(
     sizes: Union[int, Value],
     element_type: Type = None,
@@ -95,7 +102,7 @@ def alloca(
 ):
     if loc is None:
         loc = get_user_code_loc()
-    return _alloc(
+    return __alloc(
         AllocaOp,
         sizes,
         element_type,
@@ -106,26 +113,31 @@ def alloca(
     )
 
 
-def load(mem: Value, indices: Sequence[Union[Value, int]], *, loc=None, ip=None):
+def load(memref: Value, indices: Sequence[Union[Value, int]], *, loc=None, ip=None):
     if loc is None:
         loc = get_user_code_loc()
     indices = list(indices)
     for idx, i in enumerate(indices):
         if isinstance(i, int):
             indices[idx] = constant(i, index=True)
-    return get_op_result_or_op_results(LoadOp(mem, indices, loc=loc, ip=ip))
+    return get_op_result_or_op_results(LoadOp(memref, indices, loc=loc, ip=ip))
 
 
 def store(
-    value: Value, mem: Value, indices: Sequence[Union[Value, int]], *, loc=None, ip=None
+    value: Value,
+    memref: Value,
+    indices: Sequence[Union[Value, int]],
+    *,
+    loc=None,
+    ip=None,
 ):
     if loc is None:
         loc = get_user_code_loc()
     indices = list(indices)
     for idx, i in enumerate(indices):
         if isinstance(i, int):
             indices[idx] = constant(i, index=True)
-    return get_op_result_or_op_results(StoreOp(value, mem, indices, loc=loc, ip=ip))
+    return get_op_result_or_op_results(StoreOp(value, memref, indices, loc=loc, ip=ip))
 
 
 @register_value_caster(MemRefType.static_typeid)
@@ -176,7 +188,9 @@ def __setitem__(self, idx, val):
             if isinstance(val, (int, float)):
                 # TODO: this is an unchecked conversion
                 val = Scalar(val, dtype=self.dtype)
-            assert isinstance(val, Scalar), "coordinate insert requires scalar element"
+            assert isinstance(
+                val, (Scalar, Vector)
+            ), "coordinate insert requires scalar element"
             store(val, self, idx, loc=loc)
         else:
             _copy_to_subview(self, val, tuple(idx), loc=loc)

mlir/extras/dialects/ext/scf.py

@@ -105,7 +105,7 @@ def placeholder_opaque_t():
     return opaque("scf", "placeholder")
 
 
-for_ = region_op(_build_for, terminator=yield__)
+for__ = region_op(_build_for, terminator=yield__)
 
 
 @_cext.register_operation(_Dialect, replace=True)
@@ -365,7 +365,11 @@ def another_reduce(reduce_op):
             return r
 
 
-def yield_(*args):
+def yield_(*args, results_=None):
+    if len(args):
+        assert results_ is None, "must provide results_ or args"
+    if results_ is not None:
+        args = results_
     if len(args) == 1 and isinstance(args[0], (list, OpResultList)):
         args = list(args[0])
     y = yield__(args)

mlir/extras/dialects/ext/vector.py

@@ -142,11 +142,13 @@ def extract(vector, position, *, loc=None, ip=None):
 _insert = insert
 
 
-def insert(vector, val, position, *, loc=None, ip=None):
+def insert(vector, val, positions, *, loc=None, ip=None):
     if loc is None:
         loc = get_user_code_loc()
+    if len(positions) == 0:
+        raise ValueError("positions cannot be empty")
     dynamic_position, _packed_position, static_position = _dispatch_mixed_values(
-        position
+        positions
     )
     return _insert(
         val,

mlir/extras/runtime/passes.py

@@ -135,35 +135,50 @@ def add_pass(self, pass_name, **kwargs):
         self._pipeline.append(pass_str)
         return self
 
-    def lower_to_llvm_(self):
-        return any(["to-llvm" in p for p in self._pipeline])
-
-    def bufferize(self):
-        return (
-            self.Func(Pipeline().empty_tensor_to_alloc_tensor())
-            .one_shot_bufferize()
-            .Func(Pipeline().buffer_deallocation_simplification())
-        )
-
     def lower_to_llvm(self):
+        # https://github.com/makslevental/llvm-project/blob/f6643263631bcb0d191ef923963ac1a5ca9ac5fd/mlir/test/lib/Dialect/LLVM/TestLowerToLLVM.cpp#L44
         return (
-            self.cse()
-            .Func(Pipeline().lower_affine().arith_expand().convert_math_to_llvm())
-            .convert_math_to_libm()
-            .expand_strided_metadata()
-            .finalize_memref_to_llvm()
+            self.Func(
+                Pipeline()
+                # Blanket-convert any remaining high-level vector ops to loops if any remain.
+                .convert_vector_to_scf()
+                # Blanket-convert any remaining linalg ops to loops if any remain.
+                .convert_linalg_to_loops()
+            )
+            # Blanket-convert any remaining affine ops if any remain.
+            .lower_affine()
+            # Convert SCF to CF (always needed).
             .convert_scf_to_cf()
-            .convert_cf_to_llvm()
+            # Sprinkle some cleanups.
+            .canonicalize()
             .cse()
+            # Convert vector to LLVM (always needed).
+            .convert_vector_to_llvm(force_32bit_vector_indices=True)
+            # Convert Math to LLVM (always needed).
+            .Func(Pipeline().convert_math_to_llvm())
+            # Expand complicated MemRef operations before lowering them.
+            .expand_strided_metadata()
+            # The expansion may create affine expressions. Get rid of them.
             .lower_affine()
-            .Func(Pipeline().convert_arith_to_llvm())
+            # Convert MemRef to LLVM (always needed).
+            .finalize_memref_to_llvm()
+            # Convert Func to LLVM (always needed).
             .convert_func_to_llvm()
-            .canonicalize()
-            .convert_openmp_to_llvm()
-            .cse()
+            .convert_arith_to_llvm()
+            .convert_cf_to_llvm()
+            # Convert Index to LLVM (always needed).
+            .convert_index_to_llvm()
+            # Convert remaining unrealized_casts (always needed).
             .reconcile_unrealized_casts()
         )
 
+    def bufferize(self):
+        return (
+            self.Func(Pipeline().empty_tensor_to_alloc_tensor())
+            .one_shot_bufferize()
+            .Func(Pipeline().buffer_deallocation_simplification())
+        )
+
     def lower_to_openmp(self):
         return self.convert_scf_to_openmp().Func(Pipeline().lower_affine())
 

mlir/extras/runtime/refbackend.py

@@ -238,6 +238,8 @@ def cb(op):
                 return False
 
         kernel_func = find_ops(module.operation, cb, single=True)
+        if isinstance(kernel_func, list) and len(kernel_func) == 0:
+            raise ValueError(f"couldn't find kernel_func {kernel_name=}")
         if len(kernel_func.function_type.value.results) and generate_return_consumer:
             with InsertionPoint(module.body):
                 return_consumer = make_return_consumer(kernel_func)

mlir/extras/util.py

@@ -437,3 +437,4 @@ def __getitem__(self, item):
 
         # f is not a bound method since it was decorated...
         return self.f(self.instance, item, **kwargs)
+