Revert tile sizes, apply collapse on MPI.

devito/operator/xdsl_operator.py

@@ -53,9 +53,9 @@
 # gpu-launch-sink-index-computations seemed to have no impact
 MLIR_GPU_PIPELINE = lambda block_sizes: f'"builtin.module(test-math-algebraic-simplification,scf-parallel-loop-tiling{{parallel-loop-tile-sizes={block_sizes}}},func.func(gpu-map-parallel-loops),convert-parallel-loops-to-gpu,lower-affine, canonicalize,cse, fold-memref-alias-ops, gpu-launch-sink-index-computations, gpu-kernel-outlining, canonicalize{{region-simplify}},cse,fold-memref-alias-ops,expand-strided-metadata,lower-affine,canonicalize,cse,func.func(gpu-async-region),canonicalize,cse,convert-arith-to-llvm{{index-bitwidth=64}},convert-scf-to-cf,convert-cf-to-llvm{{index-bitwidth=64}},canonicalize,cse,convert-func-to-llvm{{use-bare-ptr-memref-call-conv}},gpu.module(convert-gpu-to-nvvm,reconcile-unrealized-casts,canonicalize,gpu-to-cubin),gpu-to-llvm,canonicalize,cse)"'
 
-XDSL_CPU_PIPELINE = lambda tile_sizes, collapse: f'"stencil-shape-inference,convert-stencil-to-ll-mlir{{{"tile-sizes="+",".join(["128"]*tile_sizes) if tile_sizes > 0 else ""} collapse={collapse}}},printf-to-llvm"'
+XDSL_CPU_PIPELINE = lambda tile_sizes, collapse: f'"stencil-shape-inference,convert-stencil-to-ll-mlir{{{"tile-sizes="+",".join(["64"]*tile_sizes) if tile_sizes > 0 else ""}}},printf-to-llvm"'
 XDSL_GPU_PIPELINE = "stencil-shape-inference,convert-stencil-to-ll-mlir{target=gpu},reconcile-unrealized-casts,printf-to-llvm"
-XDSL_MPI_PIPELINE = lambda decomp, nb_tiled_dims: f'"dmp-decompose{decomp},canonicalize-dmp,convert-stencil-to-ll-mlir{{tile-sizes={",".join(["64"]*nb_tiled_dims)}}},dmp-to-mpi{{mpi_init=false}},lower-mpi,printf-to-llvm"'
+XDSL_MPI_PIPELINE = lambda decomp, nb_tiled_dims, collapse: f'"dmp-decompose-2d{decomp},canonicalize-dmp,convert-stencil-to-ll-mlir{{tile-sizes={",".join(["64"]*nb_tiled_dims)} collapse={collapse}}},dmp-to-mpi{{mpi_init=false}},lower-mpi,printf-to-llvm"'
 
 
 class XDSLOperator(Operator):
@@ -136,7 +136,7 @@ def _jit_compile(self):
                 decomp = "2d-grid" if len(shape) == 2 else "3d-grid"
 
                 decomp = f"{{strategy={decomp} slices={slices} restrict_domain=false}}"
-                xdsl_pipeline = XDSL_MPI_PIPELINE(decomp, to_tile)
+                xdsl_pipeline = XDSL_MPI_PIPELINE(decomp, to_tile, to_tile)
             elif is_gpu:
                 xdsl_pipeline = XDSL_GPU_PIPELINE
                 mlir_pipeline = MLIR_GPU_PIPELINE(block_sizes)