auto wrapper generation

Author: makslevental

.github/workflows/test.yml

@@ -47,9 +47,9 @@ jobs:
         shell: bash
         run: |
           if [ ${{ matrix.os }} == 'windows-2022' ]; then
-            pytest --ignore-glob=*test_smoke* tests
+            pytest -s --ignore-glob=*test_smoke* tests
           else
-            pytest --ignore-glob=*test_smoke* tests
+            pytest -s --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 tests/test_smoke.py
+            pytest -s tests/test_smoke.py
           else
-            pytest tests/test_smoke.py
+            pytest -s tests/test_smoke.py
           fi
 
 
@@ -126,7 +126,7 @@ jobs:
         shell: bash
         run: |
           if [ ${{ matrix.os }} == 'windows-2022' ]; then
-            pytest tests/test_smoke.py
+            pytest -s tests/test_smoke.py
           else
-            pytest tests/test_smoke.py
+            pytest -s tests/test_smoke.py
           fi

mlir_utils/context.py

@@ -37,3 +37,4 @@ def mlir_mod_ctx(
         ip = mlir.ir.InsertionPoint(module.body)
         stack.enter_context(ip)
         yield MLIRContext(context, module)
+    context._clear_live_operations()

mlir_utils/dialects/ext/tensor.py

@@ -144,12 +144,15 @@ def __getitem__(self, idx: tuple) -> "Tensor":
         if not self.has_rank():
             raise ValueError("only ranked tensor slicing/indexing supported")
 
+        if idx is None:
+            return expand_dims(self, (0,), loc=loc)
         if idx == Ellipsis or idx == slice(None):
             return self
         if isinstance(idx, tuple) and all(i == slice(None) for i in idx):
             return self
-        if idx is None:
-            return expand_dims(self, (0,), loc=loc)
+        if isinstance(idx, tuple) and all(i == slice(None) or i is None for i in idx):
+            nones = [i for i, n in enumerate(idx) if n is None]
+            return expand_dims(self, nones, loc=loc)
 
         idx = list((idx,) if isinstance(idx, int) else idx)
         for i, d in enumerate(idx):
@@ -237,6 +240,13 @@ class _Indexer:
     def is_constant(self):
         return all(_is_constant_index(i) for i in self.indices)
 
+    def is_full(self):
+        return all(
+            isinstance(idx, slice)
+            and len(range(*idx.indices(self.in_shape[i]))) == self.in_shape[i]
+            for i, idx in enumerate(self.indices)
+        )
+
     # waiting on hashable slices in 3.12 https://stackoverflow.com/a/76562346
     # @lru_cache(maxsize=1)
     def static_offsets(self):
@@ -551,7 +561,9 @@ def _extract_slice(
     indexer = _indices_to_indexer(idx, ten.shape)
     out = ten
 
-    if indexer.is_constant():
+    if indexer.is_full():
+        out = out
+    elif indexer.is_constant():
         out = extract_slice(
             out,
             static_offsets=indexer.static_offsets(),

mlir_utils/runtime/passes.py

@@ -1,10 +1,82 @@
-from __future__ import annotations
-
 import logging
+import os
+import sys
+import tempfile
+from contextlib import ExitStack
+from io import StringIO
+from typing import Optional
+
+from mlir.ir import StringAttr
+from mlir.passmanager import PassManager
+
+from mlir_utils.util import disable_multithreading
 
 logger = logging.getLogger(__name__)
 
 
+class MlirCompilerError(Exception):
+    pass
+
+
+def get_module_name_for_debug_dump(module):
+    if "debug_module_name" not in module.operation.attributes:
+        return "UnnammedModule"
+    return StringAttr(module.operation.attributes["debug_module_name"]).value
+
+
+def run_pipeline(
+    module,
+    pipeline: str,
+    description: Optional[str] = None,
+    enable_ir_printing=False,
+    print_pipeline=False,
+):
+    """Runs `pipeline` on `module`, with a nice repro report if it fails."""
+    module_name = get_module_name_for_debug_dump(module)
+    try:
+        original_stderr = sys.stderr
+        sys.stderr = StringIO()
+        # Lower module in place to make it ready for compiler backends.
+        with ExitStack() as stack:
+            stack.enter_context(module.context)
+            asm_for_error_report = module.operation.get_asm(
+                large_elements_limit=10,
+                enable_debug_info=True,
+            )
+            pm = PassManager.parse(pipeline)
+            if print_pipeline:
+                print(pm)
+            if enable_ir_printing:
+                stack.enter_context(disable_multithreading())
+                pm.enable_ir_printing()
+
+            pm.run(module.operation)
+    except Exception as e:
+        print(e, file=sys.stderr)
+        filename = os.path.join(tempfile.gettempdir(), module_name + ".mlir")
+        with open(filename, "w") as f:
+            f.write(asm_for_error_report)
+        debug_options = "-mlir-print-ir-after-all -mlir-disable-threading"
+        description = description or f"{module_name} compile"
+
+        message = f"""\
+            {description} failed with the following diagnostics:
+
+            {'*' * 80}
+            {sys.stderr.getvalue().strip()}
+            {'*' * 80}
+
+            For developers, the error can be reproduced with:
+            $ mlir-opt {debug_options} -pass-pipeline='{pipeline}' {filename}
+            """
+        trimmed_message = "\n".join([m.lstrip() for m in message.split("\n")])
+        raise MlirCompilerError(trimmed_message)
+    finally:
+        sys.stderr = original_stderr
+
+    return module
+
+
 class Pipeline:
     _pipeline: list[str] = []
 
@@ -13,17 +85,17 @@ def __init__(self, pipeline=None, wrapper=None):
             pipeline = []
         self._pipeline = pipeline
 
-    def Func(self, p: Pipeline):
+    def Func(self, p: "Pipeline"):
         assert isinstance(p, Pipeline)
         self._pipeline.append(f"func.func({p.materialize(module=False)})")
         return self
 
-    def Spirv(self, p: Pipeline):
+    def Spirv(self, p: "Pipeline"):
         assert isinstance(p, Pipeline)
         self._pipeline.append(f"spirv.module({p.materialize(module=False)})")
         return self
 
-    def Gpu(self, p: Pipeline):
+    def Gpu(self, p: "Pipeline"):
         assert isinstance(p, Pipeline)
         self._pipeline.append(f"gpu.module({p.materialize(module=False)})")
         return self
@@ -38,13 +110,6 @@ def materialize(self, module=True):
     def __str__(self):
         return self.materialize()
 
-    def __add__(self, other: Pipeline):
-        return Pipeline(self._pipeline + other._pipeline)
-
-    def __iadd__(self, other: Pipeline):
-        self._pipeline += other._pipeline
-        return self
-
     def add_pass(self, pass_name, **kwargs):
         kwargs = {
             k.replace("_", "-"): int(v) if isinstance(v, bool) else v
@@ -57,6 +122,7 @@ def add_pass(self, pass_name, **kwargs):
         else:
             pass_str = f"{pass_name}"
         self._pipeline.append(pass_str)
+        return self
 
     def lower_to_llvm_(self):
         return any(["to-llvm" in p for p in self._pipeline])