fix nanobind differences (and others) (#110)

Author: makslevental

.github/workflows/test.yml

@@ -34,18 +34,12 @@ jobs:
       fail-fast: false
       matrix:
         os: [ ubuntu-22.04, macos-13, macos-14, windows-2022 ]
-        py_version: [ "3.8", "3.9", "3.10", "3.11", "3.12" ]
+        py_version: [ "3.9", "3.10", "3.11", "3.12", "3.13" ]
 
         exclude:
-          - os: macos-13
-            py_version: "3.8"
-
           - os: macos-13
             py_version: "3.9"
 
-          - os: macos-14
-            py_version: "3.8"
-
           - os: macos-14
             py_version: "3.9"
 
@@ -174,7 +168,7 @@ jobs:
       fail-fast: false
       matrix:
         os: [ ubuntu-22.04 ]
-        py_version: [ "3.8", "3.9", "3.10", "3.11", "3.12" ]
+        py_version: [ "3.9", "3.10", "3.11", "3.12", "3.13" ]
 
     steps:
       - name: Checkout
@@ -189,7 +183,7 @@ jobs:
           install: |
             
             apt-get update -q -y
-            apt-get install -y wget build-essential
+            apt-get install -y wget build-essential git
             
             mkdir -p ~/miniconda3
             wget -q https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-aarch64.sh -O miniconda.sh

examples/mwe.py

@@ -131,6 +131,8 @@ def pats():
     .finalize_memref_to_llvm()
     # Convert Func to LLVM (always needed).
     .convert_func_to_llvm()
+    .convert_arith_to_llvm()
+    .convert_cf_to_llvm()
     # Convert Index to LLVM (always needed).
     .convert_index_to_llvm()
     # Convert remaining unrealized_casts (always needed).

examples/vectorization_e2e.ipynb

@@ -424,6 +424,8 @@
     "    .finalize_memref_to_llvm()\n",
     "    # Convert Func to LLVM (always needed).\n",
     "    .convert_func_to_llvm()\n",
+    "    .convert_arith_to_llvm()\n",
+    "    .convert_cf_to_llvm()\n",
     "    # Convert Index to LLVM (always needed).\n",
     "    .convert_index_to_llvm()\n",
     "    # Convert remaining unrealized_casts (always needed).\n",

mlir/extras/ast/canonicalize.py

@@ -98,7 +98,7 @@ def transform_ast(
     module_code_o = compile(module, f.__code__.co_filename, "exec")
     new_f_code_o = find_func_in_code_object(module_code_o, f.__name__)
     n_lines = len(inspect.getsource(f).splitlines())
-    line_starts = list(findlinestarts(new_f_code_o))
+    line_starts = list(filter(lambda el: el[1], findlinestarts(new_f_code_o)))
     if (
         max([l for _, l in line_starts]) - min([l for _, l in line_starts]) + 1
         > n_lines

mlir/extras/dialects/ext/_shaped_value.py

@@ -9,7 +9,7 @@
 
 
 # mixin that requires `is_constant`
-class ShapedValue:
+def ShapedValue(cls):
     @cached_property
     def literal_value(self) -> np.ndarray:
         if not self.is_constant:
@@ -42,3 +42,22 @@ def n_elements(self) -> int:
     @cached_property
     def dtype(self) -> Type:
         return self._shaped_type.element_type
+
+    setattr(cls, "literal_value", literal_value)
+    cls.literal_value.__set_name__(None, "literal_value")
+    setattr(cls, "_shaped_type", _shaped_type)
+    cls._shaped_type.__set_name__(None, "_shaped_type")
+
+    setattr(cls, "has_static_shape", has_static_shape)
+    setattr(cls, "has_rank", has_rank)
+
+    setattr(cls, "rank", rank)
+    cls.rank.__set_name__(None, "rank")
+    setattr(cls, "shape", shape)
+    cls.shape.__set_name__(None, "shape")
+    setattr(cls, "n_elements", n_elements)
+    cls.n_elements.__set_name__(None, "n_elements")
+    setattr(cls, "dtype", dtype)
+    cls.dtype.__set_name__(None, "dtype")
+
+    return cls

mlir/extras/dialects/ext/arith.py

@@ -7,6 +7,7 @@
 from typing import Optional, Tuple, Union
 
 from bytecode import ConcreteBytecode
+from einspect.structs import PyTypeObject
 
 from ...ast.canonicalize import StrictTransformer, Canonicalizer, BytecodePatcher
 from ...ast.util import ast_call
@@ -138,7 +139,13 @@ def index_cast(
     )
 
 
-class ArithValueMeta(type(Value)):
+nb_meta_cls = type(Value)
+
+_Py_TPFLAGS_BASETYPE = 1 << 10
+PyTypeObject.from_object(nb_meta_cls).tp_flags |= _Py_TPFLAGS_BASETYPE
+
+
+class ArithValueMeta(nb_meta_cls):
     """Metaclass that orchestrates the Python object protocol
     (i.e., calling __new__ and __init__) for Indexing dialect extension values
     (created using `mlir_value_subclass`).

mlir/extras/dialects/ext/func.py

@@ -192,7 +192,9 @@ def __init__(
 
     def _is_decl(self):
         # magic constant found from looking at the code for an empty fn
-        if sys.version_info.minor == 12:
+        if sys.version_info.minor == 13:
+            return self.body_builder.__code__.co_code == b"\x95\x00g\x00"
+        elif sys.version_info.minor == 12:
             return self.body_builder.__code__.co_code == b"\x97\x00y\x00"
         elif sys.version_info.minor == 11:
             return self.body_builder.__code__.co_code == b"\x97\x00d\x00S\x00"

mlir/extras/dialects/ext/memref.py

@@ -129,7 +129,8 @@ def store(
 
 
 @register_value_caster(MemRefType.static_typeid)
-class MemRef(Value, ShapedValue):
+@ShapedValue
+class MemRef(Value):
     def __str__(self):
         return f"{self.__class__.__name__}({self.get_name()}, {self.type})"
 

mlir/extras/dialects/ext/tensor.py

@@ -109,7 +109,8 @@ def insert_slice(
 
 # TODO(max): unify vector/memref/tensor
 @register_value_caster(RankedTensorType.static_typeid)
-class Tensor(ShapedValue, ArithValue):
+@ShapedValue
+class Tensor(ArithValue):
     def __getitem__(self, idx: tuple) -> "Tensor":
         loc = get_user_code_loc()
 

mlir/extras/dialects/ext/vector.py

@@ -14,7 +14,8 @@
 
 
 @register_value_caster(VectorType.static_typeid)
-class Vector(ShapedValue, ArithValue):
+@ShapedValue
+class Vector(ArithValue):
     def __getitem__(self, idx: tuple) -> "Vector":
         loc = get_user_code_loc()
 
@@ -105,7 +106,7 @@ def transfer_read(
     if isinstance(padding, int):
         padding = constant(padding, type=source.type.element_type)
     if in_bounds is None:
-        in_bounds = [None] * len(permutation_map.results)
+        raise ValueError("in_bounds cannot be None")
 
     return _transfer_read(
         vector=vector_t,