Format some Python files with black. NFC

mlir/test/lit.cfg.py

@@ -18,54 +18,68 @@
 # Configuration file for the 'lit' test runner.
 
 # name: The name of this test suite.
-config.name = 'AIRMLIR'
+config.name = "AIRMLIR"
 
 config.test_format = lit.formats.ShTest(not llvm_config.use_lit_shell)
 
 # suffixes: A list of file extensions to treat as test files.
-config.suffixes = ['.mlir']
+config.suffixes = [".mlir"]
 
 # test_source_root: The root path where tests are located.
 config.test_source_root = os.path.dirname(__file__)
 
 # test_exec_root: The root path where tests should be run.
-config.test_exec_root = os.path.join(config.air_obj_root, 'test')
-air_runtime_lib = os.path.join(config.air_obj_root, "runtime_lib", config.runtime_test_target)
-
-config.substitutions.append(('%PATH%', config.environment['PATH']))
-config.substitutions.append(('%shlibext', config.llvm_shlib_ext))
-config.substitutions.append(('%CLANG', "clang"))
-config.substitutions.append(('%LLC', config.llvm_tools_dir + "/llc"))
-config.substitutions.append(('%OPT', config.llvm_tools_dir + "/opt"))
-config.substitutions.append(("%airhost_inc", "-I" + air_runtime_lib + "/airhost/include"))
-config.substitutions.append(("%aircpu_lib", "-L" + air_runtime_lib + "/aircpu -laircpu"))
-config.substitutions.append(("%mlir_async_lib", "-L" + config.llvm_obj_root + "/lib -lmlir_async_runtime"))
-config.substitutions.append(("%ld_lib_path", "LD_LIBRARY_PATH=" + air_runtime_lib + "/aircpu:" + config.llvm_obj_root + "/lib"))
-
-llvm_config.with_system_environment(
-    ['HOME', 'INCLUDE', 'LIB', 'TMP', 'TEMP'])
+config.test_exec_root = os.path.join(config.air_obj_root, "test")
+air_runtime_lib = os.path.join(
+    config.air_obj_root, "runtime_lib", config.runtime_test_target
+)
+
+config.substitutions.append(("%PATH%", config.environment["PATH"]))
+config.substitutions.append(("%shlibext", config.llvm_shlib_ext))
+config.substitutions.append(("%CLANG", "clang"))
+config.substitutions.append(("%LLC", config.llvm_tools_dir + "/llc"))
+config.substitutions.append(("%OPT", config.llvm_tools_dir + "/opt"))
+config.substitutions.append(
+    ("%airhost_inc", "-I" + air_runtime_lib + "/airhost/include")
+)
+config.substitutions.append(
+    ("%aircpu_lib", "-L" + air_runtime_lib + "/aircpu -laircpu")
+)
+config.substitutions.append(
+    ("%mlir_async_lib", "-L" + config.llvm_obj_root + "/lib -lmlir_async_runtime")
+)
+config.substitutions.append(
+    (
+        "%ld_lib_path",
+        "LD_LIBRARY_PATH="
+        + air_runtime_lib
+        + "/aircpu:"
+        + config.llvm_obj_root
+        + "/lib",
+    )
+)
+
+llvm_config.with_system_environment(["HOME", "INCLUDE", "LIB", "TMP", "TEMP"])
 
 llvm_config.use_default_substitutions()
 
 # excludes: A list of directories to exclude from the testsuite. The 'Inputs'
 # subdirectories contain auxiliary inputs for various tests in their parent
 # directories.
-config.excludes = ['Inputs', 'Examples', 'CMakeLists.txt', 'README.txt', 'LICENSE.txt']
+config.excludes = ["Inputs", "Examples", "CMakeLists.txt", "README.txt", "LICENSE.txt"]
 
 # test_source_root: The root path where tests are located.
 config.test_source_root = os.path.dirname(__file__)
 
 # test_exec_root: The root path where tests should be run.
-config.test_exec_root = os.path.join(config.air_obj_root, 'test')
-config.air_tools_dir = os.path.join(config.air_obj_root, 'bin')
+config.test_exec_root = os.path.join(config.air_obj_root, "test")
+config.air_tools_dir = os.path.join(config.air_obj_root, "bin")
 
 # Tweak the PATH to include the tools dir.
-llvm_config.with_environment('PATH', config.llvm_tools_dir, append_path=True)
-llvm_config.with_environment('PATH', config.aie_tools_dir, append_path=True)
+llvm_config.with_environment("PATH", config.llvm_tools_dir, append_path=True)
+llvm_config.with_environment("PATH", config.aie_tools_dir, append_path=True)
 
 tool_dirs = [config.air_tools_dir, config.aie_tools_dir, config.llvm_tools_dir]
-tools = [
-    'air-opt', 'air-translate', 'air-runner', 'aie-opt'
-]
+tools = ["air-opt", "air-translate", "air-runner", "aie-opt"]
 
 llvm_config.add_tool_substitutions(tools, tool_dirs)

python/air/backend/__init__.py

@@ -3,4 +3,3 @@
 # Copyright (C) 2022, Xilinx Inc.
 # Copyright (C) 2022, Advanced Micro Devices, Inc.
 # SPDX-License-Identifier: MIT
-

python/air/backend/abc.py

@@ -14,12 +14,12 @@
 # A type shared between the result of `AirBackend.compile` and the
 # input to `AirBackend.load`. Each backend will likely have a
 # different definition of this type.
-CompiledArtifact = TypeVar('CompiledArtifact')
+CompiledArtifact = TypeVar("CompiledArtifact")
 
 # A wrapper around a backend-specific loaded program representation
 # that uniformly translates the `x.method(...)` interface expected of
 # Torch modules into appropriate lower-level operations.
-Invoker = TypeVar('Invoker')
+Invoker = TypeVar("Invoker")
 
 
 class AirBackend(abc.ABC):
@@ -28,6 +28,7 @@ class AirBackend(abc.ABC):
     Backends are recommended to raise meaningful exceptions in case of error,
     ideally with easy reproduction instructions.
     """
+
     @abc.abstractmethod
     def compile(self, module: Module) -> CompiledArtifact:
         """Compile the provided MLIR module into a compiled artifact.

python/air/backend/cpu_backend.py

@@ -10,7 +10,9 @@
 import air.mlir.ir
 import air.mlir.passmanager
 
-from torch_mlir_e2e_test.linalg_on_tensors_backends.refbackend import RefBackendLinalgOnTensorsBackend
+from torch_mlir_e2e_test.linalg_on_tensors_backends.refbackend import (
+    RefBackendLinalgOnTensorsBackend,
+)
 
 from .abc import AirBackend
 
@@ -24,80 +26,89 @@
 
 path = Path(air.backend.__file__).resolve().parent
 ctypes.CDLL(f"{path}/../../../runtime_lib/aircpu/libaircpu.so", mode=ctypes.RTLD_GLOBAL)
-ctypes.CDLL(f"/FIXME/PATH/TO/llvm/lib/libmlir_async_runtime.so.17git", mode=ctypes.RTLD_GLOBAL)
-
-__all__ = [
-    "AirCpuBackend",
-    "make_dynamo_backend",
-    "DEFAULT_PIPELINE"
-]
-
-DEFAULT_PIPELINE = "builtin.module("+",".join([
-    "air-to-async",
-    "canonicalize",
-    "cse"
-])+")"
-
-ASYNC_TO_LLVM_PIPELINE = "builtin.module("+",".join([
-    "func.func(buffer-deallocation)",
-    "async-to-async-runtime",
-    "async-runtime-ref-counting",
-    "async-runtime-ref-counting-opt",
-    "convert-async-to-llvm",
-    "canonicalize",
-    "cse"
-])+")"
-
-REF_BACKEND_LOWERING_PIPELINE = "builtin.module(" + ",".join([
-    "func.func(refback-generalize-tensor-pad)",
-    # Apply some optimizations. It would be great if MLIR had more useful
-    # optimizations that worked out of the box here.
-    # Note: When measured, this doesn't seem to actually help that much
-    # for the linalg-on-tensors backend.
-    # This is likely because if things are naturally fusable we usually already
-    # emit things in that form from the high level (e.g. single linalg-generic).
-    # Other backends are likely to benefit more.
-    "func.func(linalg-fuse-elementwise-ops)",
-    # Bufferize.
-    "func.func(scf-bufferize)",
-    "func.func(tm-tensor-bufferize)",
-    "func.func(empty-tensor-to-alloc-tensor)",
-    "func.func(linalg-bufferize)",
-    "func-bufferize",
-    "arith-bufferize",
-    "refback-mlprogram-bufferize",
-    "func.func(tensor-bufferize)",
-    "func.func(finalizing-bufferize)",
-    #"func.func(buffer-deallocation)",
-    # Munge to make it ExecutionEngine compatible.
-    # Specifically, we rewrite calling convention boundaries to be in terms
-    # of unranked memref, and we rewrite the return to actually be a
-    # callback that consumes the return (the final munged function always
-    # returns void at the C level -- we get the return value by providing the
-    # callback).
-    "refback-munge-calling-conventions",
-    # Insert global variable and instruction sequence for getting the next
-    # global seed used in stateful rng.
-    # Lower to LLVM
-    "func.func(tm-tensor-to-loops)",
-    "func.func(refback-munge-memref-copy)",
-    "func.func(convert-linalg-to-loops)",
-    "func.func(lower-affine)",
-    "convert-scf-to-cf",
-    "func.func(refback-expand-ops-for-llvm)",
-    "func.func(arith-expand)",
-    "func.func(convert-math-to-llvm)",
-    # Handle some complex mlir::math ops (e.g. atan2)
-    "convert-math-to-libm",
-    "convert-linalg-to-llvm",
-    "expand-strided-metadata",
-    "finalize-memref-to-llvm",
-    "lower-affine",
-    "func.func(convert-arith-to-llvm)",
-    "convert-func-to-llvm",
-    "convert-cf-to-llvm",
-    "reconcile-unrealized-casts",
-]) + ")"
+ctypes.CDLL(
+    f"/FIXME/PATH/TO/llvm/lib/libmlir_async_runtime.so.17git", mode=ctypes.RTLD_GLOBAL
+)
+
+__all__ = ["AirCpuBackend", "make_dynamo_backend", "DEFAULT_PIPELINE"]
+
+DEFAULT_PIPELINE = (
+    "builtin.module(" + ",".join(["air-to-async", "canonicalize", "cse"]) + ")"
+)
+
+ASYNC_TO_LLVM_PIPELINE = (
+    "builtin.module("
+    + ",".join(
+        [
+            "func.func(buffer-deallocation)",
+            "async-to-async-runtime",
+            "async-runtime-ref-counting",
+            "async-runtime-ref-counting-opt",
+            "convert-async-to-llvm",
+            "canonicalize",
+            "cse",
+        ]
+    )
+    + ")"
+)
+
+REF_BACKEND_LOWERING_PIPELINE = (
+    "builtin.module("
+    + ",".join(
+        [
+            "func.func(refback-generalize-tensor-pad)",
+            # Apply some optimizations. It would be great if MLIR had more useful
+            # optimizations that worked out of the box here.
+            # Note: When measured, this doesn't seem to actually help that much
+            # for the linalg-on-tensors backend.
+            # This is likely because if things are naturally fusable we usually already
+            # emit things in that form from the high level (e.g. single linalg-generic).
+            # Other backends are likely to benefit more.
+            "func.func(linalg-fuse-elementwise-ops)",
+            # Bufferize.
+            "func.func(scf-bufferize)",
+            "func.func(tm-tensor-bufferize)",
+            "func.func(empty-tensor-to-alloc-tensor)",
+            "func.func(linalg-bufferize)",
+            "func-bufferize",
+            "arith-bufferize",
+            "refback-mlprogram-bufferize",
+            "func.func(tensor-bufferize)",
+            "func.func(finalizing-bufferize)",
+            # "func.func(buffer-deallocation)",
+            # Munge to make it ExecutionEngine compatible.
+            # Specifically, we rewrite calling convention boundaries to be in terms
+            # of unranked memref, and we rewrite the return to actually be a
+            # callback that consumes the return (the final munged function always
+            # returns void at the C level -- we get the return value by providing the
+            # callback).
+            "refback-munge-calling-conventions",
+            # Insert global variable and instruction sequence for getting the next
+            # global seed used in stateful rng.
+            # Lower to LLVM
+            "func.func(tm-tensor-to-loops)",
+            "func.func(refback-munge-memref-copy)",
+            "func.func(convert-linalg-to-loops)",
+            "func.func(lower-affine)",
+            "convert-scf-to-cf",
+            "func.func(refback-expand-ops-for-llvm)",
+            "func.func(arith-expand)",
+            "func.func(convert-math-to-llvm)",
+            # Handle some complex mlir::math ops (e.g. atan2)
+            "convert-math-to-libm",
+            "convert-linalg-to-llvm",
+            "expand-strided-metadata",
+            "finalize-memref-to-llvm",
+            "lower-affine",
+            "func.func(convert-arith-to-llvm)",
+            "convert-func-to-llvm",
+            "convert-cf-to-llvm",
+            "reconcile-unrealized-casts",
+        ]
+    )
+    + ")"
+)
+
 
 class AirCpuBackend(AirBackend):
     """Main entry-point for the AIR CPU backend.
@@ -106,6 +117,7 @@ class AirCpuBackend(AirBackend):
     for JIT execution.
 
     """
+
     def __init__(self):
         super().__init__()
         self.handle = None
@@ -114,8 +126,14 @@ def __init__(self):
     def __del__(self):
         self.unload()
 
-    def compile(self, air_module: air.mlir.ir.Module, pipeline=None,
-                verbose=False, segment_offset=None, segment_size=None):
+    def compile(
+        self,
+        air_module: air.mlir.ir.Module,
+        pipeline=None,
+        verbose=False,
+        segment_offset=None,
+        segment_size=None,
+    ):
         """Compiles an imported module, with a flat list of functions.
 
         The module is expected to be AIR dialect.
@@ -170,6 +188,7 @@ def unload(self):
         """Unload any loaded module and release resources."""
         pass
 
+
 def make_dynamo_backend(pipeline=None, verbose=False):
     """Make a PyTorch dynamo backend using AirCpuBackend.
 
@@ -185,21 +204,27 @@ def make_dynamo_backend(pipeline=None, verbose=False):
         A PyTorch dynamo backend
     """
     backend = AirCpuBackend()
+
     @make_simple_dynamo_backend
-    def air_backend(fx_graph: torch.fx.GraphModule,
-                    example_inputs: List[torch.Tensor]):
-
+    def air_backend(fx_graph: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):
+
         # get the linalg mlir of the model from torch_mlir
         mlir_module = torch_mlir.compile(
-            fx_graph, example_inputs,
-            output_type=torch_mlir.OutputType.LINALG_ON_TENSORS)
+            fx_graph,
+            example_inputs,
+            output_type=torch_mlir.OutputType.LINALG_ON_TENSORS,
+        )
 
         with air.mlir.ir.Context():
             air_module = air.mlir.ir.Module.parse(str(mlir_module))
-            pm = air.mlir.passmanager.PassManager.parse(air.compiler.util.LINALG_TENSOR_TO_MEMREF_PIPELINE)
+            pm = air.mlir.passmanager.PassManager.parse(
+                air.compiler.util.LINALG_TENSOR_TO_MEMREF_PIPELINE
+            )
             pm.run(air_module.operation)
             if pipeline is None:
-                pm = air.mlir.passmanager.PassManager.parse(linalg_on_tensors.LINALG_MEMREF_TO_AIR_PIPELINE)
+                pm = air.mlir.passmanager.PassManager.parse(
+                    linalg_on_tensors.LINALG_MEMREF_TO_AIR_PIPELINE
+                )
                 pm.run(air_module.operation)
             elif callable(pipeline):
                 air_module = pipeline(air_module)
@@ -219,5 +244,7 @@ def compiled_callable(*inputs):
             loaded = backend.load(compiled)
             result = loaded.forward(*inputs)
             return torch.from_numpy(result)
+
         return compiled_callable
-    return air_backend
+
+    return air_backend