Drive the MLIR lowering through xDSL.

.github/workflows/ci-mlir-mpi.yml

@@ -38,10 +38,10 @@ jobs:
       run: |
         pip install -e .[tests]
         pip install mpi4py
-        pip install git+https://github.com/xdslproject/xdsl@5500ff6d82d1a920b369615292ba507ecbf92fc9
+        pip install git+https://github.com/xdslproject/xdsl@46a20b6c346a25172fa2f27aac01eb2a6b579659
 
     - name: Test with MPI
       run: |
         # Add mlir-opt to the path
         export PATH=/xdsl-sc/llvm-project/build/bin/:$PATH
-        pytest -m "parallel" -k "not adjoint" tests/test_xdsl_*
+        pytest -m "parallel" -k "not adjoint" tests/test_xdsl_* -vvv

.github/workflows/ci-mlir.yml

@@ -38,20 +38,20 @@ jobs:
       run: |
         pip install -e .[tests]
         pip install mpi4py
-        pip install git+https://github.com/xdslproject/xdsl@5500ff6d82d1a920b369615292ba507ecbf92fc9
+        pip install git+https://github.com/xdslproject/xdsl@46a20b6c346a25172fa2f27aac01eb2a6b579659
 
     - name: Test no-MPI, no-Openmp
       run: |
         export DEVITO_MPI=0
         export DEVITO_LANGUAGE=C
         # Add mlir-opt to the path
         export PATH=/xdsl-sc/llvm-project/build/bin/:$PATH
-        pytest -m "not parallel" -k "not adjoint" tests/test_xdsl_*
+        pytest -m "not parallel" -k "not adjoint" tests/test_xdsl_* -vvv
 
     - name: Test no-MPI, Openmp
       run: |
           export DEVITO_MPI=0
           export DEVITO_LANGUAGE=openmp
           # Add mlir-opt to the path
           export PATH=/xdsl-sc/llvm-project/build/bin/:$PATH
-          pytest -m "not parallel" -k "not adjoint" tests/test_xdsl_*
+          pytest -m "not parallel" -k "not adjoint" tests/test_xdsl_* -vvv

devito/core/cpu.py

@@ -9,6 +9,7 @@
 from collections import OrderedDict
 
 from functools import partial
+from typing import Iterable
 
 from devito.core.operator import CoreOperator, CustomOperator, ParTile
 from devito.exceptions import InvalidOperator
@@ -372,18 +373,24 @@ def _jit_compile(self):
 
                 # xdsl-opt, get xDSL IR
                 # TODO: Remove quotes in pipeline; currently workaround with [1:-1]
-                xdsl = xDSLOptMain(args=[source_name, "-p", xdsl_pipeline[1:-1]])
+                xdsl_args = [source_name,
+                             "--allow-unregistered-dialect",
+                             "-p",
+                             xdsl_pipeline[1:-1]+','+mlir_pipeline]
+                xdsl = xDSLOptMain(args=xdsl_args)
                 out = io.StringIO()
+                perf("-----------------")
+                perf(f"xdsl-opt {' '.join(xdsl_args)}")
                 with redirect_stdout(out):
                     xdsl.run()
 
                 # mlir-opt
-                mlir_cmd = f'mlir-opt -p {mlir_pipeline}'
-                out = self.compile(mlir_cmd, out.getvalue())
-                #  Printer().print(out)
+                # mlir_cmd = f'mlir-opt -p {mlir_pipeline}'
+                # out = self.compile(mlir_cmd, out.getvalue())
+                # #  Printer().print(out)
 
                 mlir_translate_cmd = 'mlir-translate --mlir-to-llvmir'
-                out = self.compile(mlir_translate_cmd, out)
+                out = self.compile(mlir_translate_cmd, out.getvalue())
                 # Printer().print(out)
 
                 # Compile with clang and get LLVM-IR
@@ -634,19 +641,25 @@ def _jit_compile(self):
 
                 # xdsl-opt, get xDSL IR
                 # TODO: Remove quotes in pipeline; currently workaround with [1:-1]
-                xdsl = xDSLOptMain(args=[source_name, "-p", xdsl_pipeline[1:-1]])
+                xdsl_args = [source_name,
+                             "--allow-unregistered-dialect",
+                             "-p",
+                             xdsl_pipeline[1:-1]+','+mlir_pipeline]
+                xdsl = xDSLOptMain(args=xdsl_args)
                 out = io.StringIO()
+                perf("-----------------")
+                perf(f"xdsl-opt {' '.join(xdsl_args)}")
                 with redirect_stdout(out):
                     xdsl.run()
 
                 # mlir-opt
-                mlir_cmd = f'mlir-opt -p {mlir_pipeline}'
-                out = self.compile(mlir_cmd, out.getvalue())
+                # mlir_cmd = f'mlir-opt -p {mlir_pipeline}'
+                # out = self.compile(mlir_cmd, out.getvalue())
 
                 # Printer().print(out)
 
                 mlir_translate_cmd = 'mlir-translate --mlir-to-llvmir'
-                out = self.compile(mlir_translate_cmd, out)
+                out = self.compile(mlir_translate_cmd, out.getvalue())
                 # Printer().print(out)
 
                 # Compile with clang and get LLVM-IR
@@ -791,7 +804,7 @@ class Cpu64FsgOmpOperator(Cpu64FsgOperator):
 
 def generate_MLIR_CPU_PIPELINE():
     passes = [
-        "builtin.module(canonicalize",
+        "canonicalize",
         "cse",
         "loop-invariant-code-motion",
         "canonicalize",
@@ -808,15 +821,15 @@ def generate_MLIR_CPU_PIPELINE():
         "convert-func-to-llvm{use-bare-ptr-memref-call-conv}",
         "finalize-memref-to-llvm",
         "canonicalize",
-        "cse)"
+        "cse"
     ]
 
-    return generate_pipeline(passes)
+    return generate_mlir_pipeline(passes)
 
 
 def generate_MLIR_CPU_noop_PIPELINE():
     passes = [
-        "builtin.module(canonicalize",
+        "canonicalize",
         "cse",
         # "remove-dead-values",
         "canonicalize",
@@ -825,15 +838,15 @@ def generate_MLIR_CPU_noop_PIPELINE():
         "convert-math-to-llvm",
         "convert-func-to-llvm{use-bare-ptr-memref-call-conv}",
         "finalize-memref-to-llvm",
-        "canonicalize)",
+        "canonicalize",
     ]
 
-    return generate_pipeline(passes)
+    return generate_mlir_pipeline(passes)
 
 
 def generate_MLIR_OPENMP_PIPELINE():
     passes = [
-        "builtin.module(canonicalize",
+        "canonicalize",
         "cse",
         "loop-invariant-code-motion",
         "canonicalize",
@@ -858,10 +871,10 @@ def generate_MLIR_OPENMP_PIPELINE():
         # "reconcile-unrealized-casts",
         "canonicalize",
         # "print-ir",
-        "cse)"
+        "cse"
     ]
 
-    return generate_pipeline(passes)
+    return generate_mlir_pipeline(passes)
 
 
 def generate_XDSL_CPU_PIPELINE(nb_tiled_dims):
@@ -899,11 +912,21 @@ def generate_XDSL_MPI_PIPELINE(decomp, nb_tiled_dims):
     return generate_pipeline(passes)
 
 
-def generate_pipeline(passes):
+def generate_pipeline(passes: Iterable[str]):
     passes_string = ",".join(passes)
     return f'"{passes_string}"'
 
 
+def generate_mlir_pipeline(passes: Iterable[str]):
+    passes_string = ",".join(passes)
+    return 'mlir-opt{arguments='\
+           '"--mlir-print-op-generic",'\
+           '"--allow-unregistered-dialect"'\
+           ','\
+           f'"-p","builtin.module({passes_string})"'\
+           '}'
+
+
 # small interop shim script for stuff that we don't want to implement in mlir-ir
 _INTEROP_C = """
 #include <time.h>

devito/core/gpu.py

@@ -9,7 +9,7 @@
 
 from devito.core.operator import CoreOperator, CustomOperator, ParTile
 
-from devito.core.cpu import XdslAdvOperator, generate_pipeline
+from devito.core.cpu import XdslAdvOperator, generate_mlir_pipeline, generate_pipeline
 
 from devito.exceptions import InvalidOperator
 from devito.operator.operator import rcompile
@@ -439,19 +439,25 @@ def _jit_compile(self):
 
                 # xdsl-opt, get xDSL IR
                 # TODO: Remove quotes in pipeline; currently workaround with [1:-1]
-                xdsl = xDSLOptMain(args=[source_name, "-p", xdsl_pipeline[1:-1]])
+                xdsl_args = [source_name,
+                             "--allow-unregistered-dialect",
+                             "-p",
+                             xdsl_pipeline[1:-1]+','+mlir_pipeline]
+                xdsl = xDSLOptMain(args=xdsl_args)
                 out = io.StringIO()
+                perf("-----------------")
+                perf(f"xdsl-opt {' '.join(xdsl_args)}")
                 with redirect_stdout(out):
                     xdsl.run()
 
                 # mlir-opt
-                mlir_cmd = f'mlir-opt -p {mlir_pipeline}'
-                out = self.compile(mlir_cmd, out.getvalue())
+                # mlir_cmd = f'mlir-opt -p {mlir_pipeline}'
+                # out = self.compile(mlir_cmd, out.getvalue())
 
                 # Printer().print(out)
 
                 mlir_translate_cmd = 'mlir-translate --mlir-to-llvmir'
-                out = self.compile(mlir_translate_cmd, out)
+                out = self.compile(mlir_translate_cmd, out.getvalue())
                 # Printer().print(out)
 
                 # Compile with clang and get LLVM-IR
@@ -562,7 +568,7 @@ def generate_XDSL_GPU_PIPELINE():
 # gpu-launch-sink-index-computations seemed to have no impact
 def generate_MLIR_GPU_PIPELINE(block_sizes):
     passes = [
-        "builtin.module(test-math-algebraic-simplification",
+        "test-math-algebraic-simplification",
         f"scf-parallel-loop-tiling{{parallel-loop-tile-sizes={block_sizes}}}",
         "func.func(gpu-map-parallel-loops)",
         "convert-parallel-loops-to-gpu",
@@ -593,7 +599,7 @@ def generate_MLIR_GPU_PIPELINE(block_sizes):
         "gpu-to-llvm",
         "gpu-module-to-binary",
         "canonicalize",
-        "cse)"
+        "cse"
     ]
 
-    return generate_pipeline(passes)
+    return generate_mlir_pipeline(passes)