Merge pull request #118 from xdslproject/emilien/coming-soon

Unhack stencil.store result through bufferization

.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@bd700e9665c040d478a4f3bfa286ede66216b5ed
         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@bd700e9665c040d478a4f3bfa286ede66216b5ed
 
     - 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@bd700e9665c040d478a4f3bfa286ede66216b5ed
 
     - 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@bd700e9665c040d478a4f3bfa286ede66216b5ed
 
     - 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@bd700e9665c040d478a4f3bfa286ede66216b5ed
         
     - 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
@@ -378,8 +378,6 @@ 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):
             result = self._visit_math_nodes(dim, eq.rhs, eq.lhs)
@@ -404,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

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