Merge branch 'master' into dependabot/pip/sympy-gte-1.9-and-lt-1.14

.github/workflows/ci-lit.yml

@@ -27,7 +27,7 @@ jobs:
 
     - name: Install requirements and xDSL
       run: |
-        pip install git+https://github.com/xdslproject/xdsl@f8bb935880276cf077e0a80f1905105d0a98eb33
+        pip install git+https://github.com/xdslproject/xdsl@48b530615ab877e980d1f6339c18e63970011311
         pip install -e .[tests]
 
     - name: Execute lit tests

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

@@ -36,7 +36,7 @@ jobs:
       run: |
         pip install -e .[tests]
         pip install mpi4py
-        pip install git+https://github.com/xdslproject/xdsl@f8bb935880276cf077e0a80f1905105d0a98eb33
+        pip install git+https://github.com/xdslproject/xdsl@48b530615ab877e980d1f6339c18e63970011311
 
     - name: Test with MPI + openmp
       run: |

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

@@ -36,7 +36,7 @@ jobs:
       run: |
         pip install -e .[tests]
         pip install mpi4py
-        pip install git+https://github.com/xdslproject/xdsl@f8bb935880276cf077e0a80f1905105d0a98eb33
+        pip install git+https://github.com/xdslproject/xdsl@48b530615ab877e980d1f6339c18e63970011311
 
     - name: Test with MPI - no Openmp
       run: |

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

@@ -36,7 +36,7 @@ jobs:
       run: |
         pip install -e .[tests]
         pip install mpi4py
-        pip install git+https://github.com/xdslproject/xdsl@f8bb935880276cf077e0a80f1905105d0a98eb33
+        pip install git+https://github.com/xdslproject/xdsl@48b530615ab877e980d1f6339c18e63970011311
 
     - name: Test no-MPI, Openmp
       run: |

.github/workflows/ci-mlir.yml

@@ -35,7 +35,7 @@ jobs:
     - name: Install requirements and xDSL
       run: |
         pip install -e .[tests]
-        pip install git+https://github.com/xdslproject/xdsl@f8bb935880276cf077e0a80f1905105d0a98eb33
+        pip install git+https://github.com/xdslproject/xdsl@48b530615ab877e980d1f6339c18e63970011311
         
     - name: Test no-MPI, no-Openmp
       run: |

devito/ir/xdsl_iet/cluster_to_ssa.py

@@ -198,7 +198,7 @@ def _visit_math_nodes(self, dim: SteppingDimension, node: Expr,
             if output_indexed is not None:
                 space_offsets = ([node.indices[d] - output_indexed.indices[d]
                                  for d in node.function.space_dimensions])
-                temp = self.function_values[(node.function, time_offset)]
+                temp = self.apply_temps[(node.function, time_offset)]
                 access = stencil.AccessOp.get(temp, space_offsets)
                 return access.res
             # Otherwise, generate a load op
@@ -287,7 +287,7 @@ def _visit_math_nodes(self, dim: SteppingDimension, node: Expr,
             SSAargs = (self._visit_math_nodes(dim, arg, output_indexed)
                        for arg in node.args)
             return reduce(lambda x, y : arith.AndI(x, y).result, SSAargs)
-        
+
         # Trigonometric functions
         elif isinstance(node, sin):
             assert len(node.args) == 1, "Expected single argument for sin."
@@ -298,13 +298,13 @@ def _visit_math_nodes(self, dim: SteppingDimension, node: Expr,
             assert len(node.args) == 1, "Expected single argument for cos."           
             return math.CosOp(self._visit_math_nodes(dim, node.args[0],
                               output_indexed)).result
-        
+
         elif isinstance(node, tan):
             assert len(node.args) == 1, "Expected single argument for TanOp."
-            
+
             return math.TanOp(self._visit_math_nodes(dim, node.args[0],
                               output_indexed)).result
-                   
+
         elif isinstance(node, Relational):
             if isinstance(node, GreaterThan):
                 mnemonic = "sge"
@@ -378,11 +378,22 @@ def build_stencil_step(self, dim: SteppingDimension, eq: LoweredEq) -> None:
             apply_arg.name_hint = apply_op.name_hint.replace("temp", "blk")
 
         self.apply_temps = {k: v for k, v in zip(read_functions, apply.region.block.args)}
-        # Update the function values with the new temps
-        self.function_values |= self.apply_temps
 
         with ImplicitBuilder(apply.region.block):
-            stencil.ReturnOp.get([self._visit_math_nodes(dim, eq.rhs, eq.lhs)])
+            result = self._visit_math_nodes(dim, eq.rhs, eq.lhs)
+            expected_type = apply.res[0].type.get_element_type()
+            match expected_type:
+                case result.type:
+                    pass
+                case builtin.f32:
+                    if result.type == IndexType():
+                        result = arith.IndexCastOp(result, builtin.i64).result
+                    result = arith.SIToFPOp(result, builtin.f32).result
+                case builtin.IndexType:
+                    result = arith.IndexCastOp(result, IndexType()).result
+                case _:
+                    raise Exception(f"Unexpected result type {type(result)}")
+            stencil.ReturnOp.get([result])
 
         lb = stencil.IndexAttr.get(*([0] * len(shape)))
         ub = stencil.IndexAttr.get(*shape)
@@ -391,12 +402,10 @@ def build_stencil_step(self, dim: SteppingDimension, eq: LoweredEq) -> None:
             apply.res[0],
             self.function_values[self.out_time_buffer],
             stencil.StencilBoundsAttr(zip(lb, ub)),
-            stencil.TempType(len(shape),
-                             element_type=dtype_to_xdsltype(write_function.dtype))
         )
-
-        store.temp_with_halo.name_hint = f"{write_function.name}_t{self.out_time_buffer[1]}_temp"  # noqa
-        self.temps[self.out_time_buffer] = store.temp_with_halo
+        load = stencil.LoadOp.get(self.function_values[self.out_time_buffer])
+        load.res.name_hint = f"{write_function.name}_t{self.out_time_buffer[1]}_temp"  # noqa
+        self.temps[self.out_time_buffer] = load.res
 
     def build_generic_step_expression(self, dim: SteppingDimension, eq: LoweredEq):
         # Sources
@@ -439,7 +448,6 @@ def build_condition(self, dim: SteppingDimension, eq: BooleanFunction):
             self.build_generic_step_expression(dim, eq)
             scf.Yield()
 
-
     def build_time_loop(
         self, eqs: list[Any], step_dim: SteppingDimension, **kwargs
     ):
@@ -450,7 +458,7 @@ def build_time_loop(
         ub = iet_ssa.LoadSymbolic.get(
             step_dim.symbolic_max._C_name, IndexType()
         )
-        
+
         one = arith.Constant.from_int_and_width(1, IndexType())
 
         # Devito iterates from time_m to time_M *inclusive*, MLIR only takes
@@ -497,7 +505,7 @@ def build_time_loop(
             for i, (f, t) in enumerate(self.time_buffers)
         }
         self.function_values |= self.block_args
-        
+
         # Name the block argument for debugging
         for (f, t), arg in self.block_args.items():
             arg.name_hint = f"{f.name}_t{t}"
@@ -513,8 +521,7 @@ def build_time_loop(
 
     def lower_devito_Eqs(self, eqs: list[Any], **kwargs):
         # Lower devito Equations to xDSL
-
-
+
         for eq in eqs:
             lowered = self.operator._lower_exprs(as_tuple(eq), **kwargs)
             if isinstance(eq, Eq):
@@ -546,7 +553,7 @@ def _lower_injection(self, eqs: list[LoweredEq]):
                 lb = arith.Constant.from_int_and_width(int(lower), IndexType())
             else:
                 raise NotImplementedError(f"Lower bound of type {type(lower)} not supported")
-            
+
             try:
                 name = interval.dim.symbolic_min.name
             except:
@@ -633,7 +640,7 @@ def convert(self, eqs: Iterable[Eq], **kwargs) -> ModuleOp:
         # Instantiate the module.
         self.function_values: dict[tuple[Function, int], SSAValue] = {}
         self.symbol_values: dict[str, SSAValue] = {}
-        
+
         module = ModuleOp(Region([block := Block([])]))
         with ImplicitBuilder(block):
             # Get all functions used in the equations
@@ -647,7 +654,7 @@ def convert(self, eqs: Iterable[Eq], **kwargs) -> ModuleOp:
                         functions.add(f.function)
 
                 elif isinstance(eq, Injection):
-                    
+
                     functions.add(eq.field.function)
                     for f in retrieve_functions(eq.expr):
                         if isinstance(f, PointSource):

devito/xdsl_core/xdsl_cpu.py

@@ -597,7 +597,10 @@ def generate_MLIR_OPENMP_PIPELINE(kwargs):
 
 def generate_XDSL_CPU_PIPELINE(nb_tiled_dims):
     passes = [
-        "stencil-shape-inference",
+        "canonicalize",
+        "cse",
+        "shape-inference",
+        "stencil-bufferize",
         "convert-stencil-to-ll-mlir",
         f"scf-parallel-loop-tiling{{{generate_tiling_arg(nb_tiled_dims)}}}",
         "printf-to-llvm",
@@ -609,7 +612,10 @@ def generate_XDSL_CPU_PIPELINE(nb_tiled_dims):
 
 def generate_XDSL_CPU_noop_PIPELINE():
     passes = [
-        "stencil-shape-inference",
+        "canonicalize",
+        "cse",
+        "shape-inference",
+        "stencil-bufferize",
         "convert-stencil-to-ll-mlir",
         "printf-to-llvm"
     ]
@@ -619,11 +625,15 @@ def generate_XDSL_CPU_noop_PIPELINE():
 
 def generate_XDSL_MPI_PIPELINE(decomp, nb_tiled_dims):
     passes = [
+        "canonicalize",
+        "cse",
         f"distribute-stencil{decomp}",
+        "shape-inference",
         "canonicalize-dmp",
+        "stencil-bufferize",
+        "dmp-to-mpi{mpi_init=false}",
         "convert-stencil-to-ll-mlir",
         f"scf-parallel-loop-tiling{{{generate_tiling_arg(nb_tiled_dims)}}}",
-        "dmp-to-mpi{mpi_init=false}",
         "lower-mpi",
         "printf-to-llvm",
         "canonicalize"

devito/xdsl_core/xdsl_gpu.py

@@ -140,7 +140,10 @@ def _jit_compile(self):
 
 def generate_XDSL_GPU_PIPELINE():
     passes = [
-        "stencil-shape-inference",
+        "canonicalize",
+        "cse",
+        "shape-inference",
+        "stencil-bufferize",
         "convert-stencil-to-ll-mlir",
         "reconcile-unrealized-casts",
         "printf-to-llvm",

tests/filecheck/.lit_test_times.txt

@@ -0,0 +1,9 @@
+1.568606e-01 shape_inference.mlir
+8.218265e-02 version.mlir
+2.304106e-01 xdsl_mpi_pipeline.mlir
+2.372591e-01 xdsl_mpi_pipeline_b.mlir
+7.415709e-01 xdsl_mpi_pipeline_c.mlir
+7.533531e-01 xdsl_mpi_pipeline_d.mlir
+7.714086e-01 xdsl_mpi_pipeline_e.mlir
+2.611284e-01 xdsl_pipeline.mlir
+1.698225e-01 xdsl_pipeline_openmp.mlir

tests/filecheck/shape_inference.mlir

@@ -1,4 +1,4 @@
-// RUN: xdsl-opt -p stencil-shape-inference %s | filecheck %s
+// RUN: xdsl-opt -p shape-inference %s | filecheck %s
 
 builtin.module {
   func.func @Kernel(%f2_vec0 : !stencil.field<[-2,5]x[-2,5]xf32>, %f2_vec1 : !stencil.field<[-2,5]x[-2,5]xf32>, %timers : !llvm.ptr) {
@@ -64,7 +64,7 @@ builtin.module {
         %47 = arith.mulf %46, %dt_1 : f32
         stencil.return %47 : f32
       }
-      %f2_t1_temp_1 = stencil.store %f2_t1_temp to %f2_t1 ([0, 0] : [3, 3]) : !stencil.temp<?x?xf32> to !stencil.field<[-2,5]x[-2,5]xf32> with_halo : !stencil.temp<?x?xf32>
+      stencil.store %f2_t1_temp to %f2_t1(<[0, 0], [3, 3]>)  : !stencil.temp<?x?xf32> to !stencil.field<[-2,5]x[-2,5]xf32>
       scf.yield %f2_t1, %f2_t0 : !stencil.field<[-2,5]x[-2,5]xf32>, !stencil.field<[-2,5]x[-2,5]xf32>
     }
     %5 = func.call @timer_end(%0) : (f64) -> f64
@@ -139,7 +139,7 @@ builtin.module {
 // CHECK-NEXT:          %47 = arith.mulf %46, %dt_1 : f32
 // CHECK-NEXT:          stencil.return %47 : f32
 // CHECK-NEXT:        }
-// CHECK-NEXT:        %f2_t1_temp_1 = stencil.store %f2_t1_temp to %f2_t1 ([0, 0] : [3, 3]) : !stencil.temp<[0,3]x[0,3]xf32> to !stencil.field<[-2,5]x[-2,5]xf32> with_halo : !stencil.temp<?x?xf32>
+// CHECK-NEXT:        stencil.store %f2_t1_temp to %f2_t1(<[0, 0], [3, 3]>)  : !stencil.temp<[0,3]x[0,3]xf32> to !stencil.field<[-2,5]x[-2,5]xf32>
 // CHECK-NEXT:        scf.yield %f2_t1, %f2_t0 : !stencil.field<[-2,5]x[-2,5]xf32>, !stencil.field<[-2,5]x[-2,5]xf32>
 // CHECK-NEXT:      }
 // CHECK-NEXT:      %5 = func.call @timer_end(%0) : (f64) -> f64