more compat with py38

Author: makslevental

mlir/extras/ast/canonicalize.py

@@ -8,9 +8,10 @@
 from dis import findlinestarts
 from opcode import opmap
 from types import CodeType
+from typing import List, Union
 
+import astunparse
 from bytecode import ConcreteBytecode
-from typing import List, Union
 
 from ..ast.util import get_module_cst, copy_func
 
@@ -33,14 +34,14 @@ def transform_func(f, *transformer_ctors: type(Transformer)):
     module = get_module_cst(f)
     context = types.SimpleNamespace()
     for transformer_ctor in transformer_ctors:
-        orig_code = ast.unparse(module)
+        orig_code = astunparse.unparse(module)
         func_node = module.body[0]
         replace = transformer_ctor(
             context=context, first_lineno=f.__code__.co_firstlineno - 1
         )
         logger.debug("[transformer] %s", replace.__class__.__name__)
         func_node = replace.generic_visit(func_node)
-        new_code = ast.unparse(func_node)
+        new_code = astunparse.unparse(func_node)
 
         diff = list(
             difflib.unified_diff(

mlir/extras/dialects/ext/func.py

@@ -163,7 +163,7 @@ def _is_decl(self):
             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"
-        elif sys.version_info.minor == 10:
+        elif sys.version_info.minor in {8, 9, 10}:
             return self.body_builder.__code__.co_code == b"d\x00S\x00"
         else:
             raise NotImplementedError(f"{sys.version_info.minor} not supported.")

mlir/extras/dialects/ext/scf.py

@@ -455,7 +455,7 @@ class InsertEmptyYield(StrictTransformer):
     def visit_If(self, updated_node: ast.If) -> ast.If:
         updated_node = self.generic_visit(updated_node)
 
-        new_yield = ast.Expr(ast.Yield())
+        new_yield = ast.Expr(ast.Yield(value=None))
         if not is_yield(updated_node.body[-1]):
             updated_node.body = append_hidden_node(
                 updated_node.body, deepcopy(new_yield)
@@ -469,7 +469,7 @@ def visit_If(self, updated_node: ast.If) -> ast.If:
 
     def visit_For(self, updated_node: ast.For) -> ast.For:
         updated_node = self.generic_visit(updated_node)
-        new_yield = ast.Expr(ast.Yield())
+        new_yield = ast.Expr(ast.Yield(value=None))
         if not is_yield(updated_node.body[-1]):
             updated_node.body = append_hidden_node(updated_node.body, new_yield)
         return updated_node
@@ -536,7 +536,7 @@ def visit_While(self, updated_node: ast.While) -> List[ast.AST]:
             next.__name__,
             [
                 ast.Name(f"w_{updated_node.lineno}", ctx=ast.Load()),
-                ast.Constant(False),
+                ast.Constant(False, kind="bool"),
             ],
         )
         next_ = ast.fix_missing_locations(ast.copy_location(next_, updated_node))
@@ -603,7 +603,8 @@ def visit_If(self, updated_node: ast.If) -> Union[ast.With, List[ast.With]]:
         if updated_node.orelse:
             if_op_name = ast.Name(f"__if_op__{updated_node.lineno}", ctx=ast.Load())
             withitem = ast.withitem(
-                context_expr=ast_call(else_ctx_manager.__name__, args=[if_op_name])
+                context_expr=ast_call(else_ctx_manager.__name__, args=[if_op_name]),
+                optional_vars=None,
             )
             else_with = ast.With(items=[withitem])
             if is_elif:

mlir/extras/testing/testing.py

@@ -24,7 +24,9 @@ def filecheck(correct: str, module):
     # try to find using which
     if not filecheck_path.exists():
         filecheck_path = shutil.which(filecheck_name)
-    assert Path(filecheck_path).exists() is not None, "couldn't find FileCheck"
+    assert (
+        filecheck_path is not None and Path(filecheck_path).exists() is not None
+    ), "couldn't find FileCheck"
 
     correct = "\n".join(filter(None, correct.splitlines()))
     correct = dedent(correct)

mlir/extras/util.py

@@ -39,6 +39,10 @@
     TypeID = object
 
 
+def is_relative_to(self, other):
+    return other == self or other in self.parents
+
+
 def get_user_code_loc(user_base: Optional[Path] = None):
     from .. import extras
 
@@ -52,20 +56,18 @@ def get_user_code_loc(user_base: Optional[Path] = None):
         user_base = Path(prev_frame.f_code.co_filename)
 
     while prev_frame.f_back and (
-        Path(prev_frame.f_code.co_filename).is_relative_to(mlir_extras_root_path)
-        or Path(prev_frame.f_code.co_filename).is_relative_to(sys.prefix)
-        or Path(prev_frame.f_code.co_filename).is_relative_to(user_base)
+        is_relative_to(Path(prev_frame.f_code.co_filename), mlir_extras_root_path)
+        or is_relative_to(Path(prev_frame.f_code.co_filename), sys.prefix)
+        or is_relative_to(Path(prev_frame.f_code.co_filename), user_base)
     ):
         prev_frame = prev_frame.f_back
     frame_info = inspect.getframeinfo(prev_frame)
     if sys.version_info.minor >= 11:
         return Location.file(
             frame_info.filename, frame_info.lineno, frame_info.positions.col_offset
         )
-    elif sys.version_info.minor == 10:
-        return Location.file(frame_info.filename, frame_info.lineno, col=0)
     else:
-        raise NotImplementedError(f"{sys.version_info.minor} not supported.")
+        return Location.file(frame_info.filename, frame_info.lineno, col=0)
 
 
 @contextlib.contextmanager

requirements.txt

@@ -3,4 +3,5 @@ PyYAML
 black
 bytecode
 inflection
-numpy
+numpy
+astunparse

tests/test_func.py

@@ -46,7 +46,7 @@ def demo_fun1():
         assert demo_fun1.__code__.co_code == b"\x97\x00y\x00"
     elif sys.version_info.minor == 11:
         assert demo_fun1.__code__.co_code == b"\x97\x00d\x00S\x00"
-    elif sys.version_info.minor == 10:
+    elif sys.version_info.minor in {8, 9, 10}:
         assert demo_fun1.__code__.co_code == b"d\x00S\x00"
     else:
         raise NotImplementedError(f"{sys.version_info.minor} not supported.")

tests/test_gpu.py

@@ -1,3 +1,4 @@
+import sys
 from textwrap import dedent
 
 import pytest
@@ -138,6 +139,7 @@ def mat_product_kernel(
     filecheck(correct, ctx.module)
 
 
+@pytest.mark.skipif(sys.version_info < (3, 10), reason="requires python3.10 or higher")
 def test_class_call(ctx: MLIRContext):
     scale = 1
     M, N, K = 4 * scale, 16 * scale, 8 * scale
@@ -163,7 +165,10 @@ def mat_product_kernel(
     b = alloc((N, K), T.f32())
     c = alloc((M, K), T.f32())
 
-    MyClass1.mat_product_kernel[grid_size:= [4, 4, 1], block_size:= [1, 1, 1]](a, b, c)
+    # this is to avoid python 3.8 parser
+    eval(
+        "MyClass1.mat_product_kernel[grid_size:= [4, 4, 1], block_size:= [1, 1, 1]](a, b, c)"
+    )
 
     correct = dedent(
         """\
@@ -196,6 +201,7 @@ def mat_product_kernel(
     filecheck(correct, ctx.module)
 
 
+@pytest.mark.skipif(sys.version_info < (3, 10), reason="requires python3.10 or higher")
 def test_class_call_from_func(ctx: MLIRContext):
     scale = 1
     M, N, K = 4 * scale, 16 * scale, 8 * scale
@@ -227,8 +233,9 @@ def main():
         b = alloc((N, K), T.f32())
         c = alloc((M, K), T.f32())
 
-        MyClass1.mat_product_kernel[grid_size:= [4, 4, 1], block_size:= [1, 1, 1]](
-            a, b, c
+        MyClass1
+        eval(
+            "MyClass1.mat_product_kernel[grid_size:= [4, 4, 1], block_size:= [1, 1, 1]](a, b, c)"
         )
 
     ctx.module.operation.verify()
@@ -267,6 +274,7 @@ def main():
     filecheck(correct, ctx.module)
 
 
+@pytest.mark.skipif(sys.version_info < (3, 10), reason="requires python3.10 or higher")
 def test_async_object(ctx: MLIRContext):
     scale = 1
     M, N, K = 4 * scale, 16 * scale, 8 * scale
@@ -300,12 +308,9 @@ def main():
 
         w = wait()
         stream = mlir_zero(llvm_ptr_t())
-        MyClass1.mat_product_kernel[grid_size:= [4, 4, 1], block_size:= [1, 1, 1]](
-            a,
-            b,
-            c,
-            async_dependencies=[w],
-            stream=stream,
+        MyClass1
+        eval(
+            "MyClass1.mat_product_kernel[grid_size:= [4, 4, 1], block_size:= [1, 1, 1]](a, b, c, async_dependencies=[w], stream=stream)"
         )
 
     correct = dedent(

tests/test_transformers.py

@@ -1,6 +1,7 @@
 import ast
 import sys
 from textwrap import dedent
+from typing import Tuple
 
 import astpretty
 import pytest
@@ -22,7 +23,7 @@
 pytest.mark.usefixtures("ctx")
 
 
-def _fields(n: ast.AST, show_offsets: bool = True) -> tuple[str, ...]:
+def _fields(n: ast.AST, show_offsets: bool = True) -> Tuple[str, ...]:
     strip = {"type_ignores", "decorator_list", "type_comment", "ctx", "kind"}
     fields = tuple(f for f in n._fields if f not in strip)
     attributes = ("lineno",) if "lineno" in n._attributes else ()