From 45a574f18fadcde90a1e872334903912925d6c5c Mon Sep 17 00:00:00 2001 From: Emilien Bauer Date: Sat, 21 Oct 2023 20:23:37 +0100 Subject: [PATCH] Integrate stencil collapse in MPI pipeline. --- devito/operator/xdsl_operator.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/devito/operator/xdsl_operator.py b/devito/operator/xdsl_operator.py index 6dbc750c86..211ee0f804 100644 --- a/devito/operator/xdsl_operator.py +++ b/devito/operator/xdsl_operator.py @@ -55,7 +55,7 @@ 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_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-2d{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, tile_sizes, collapse: f'"dmp-decompose-2d{decomp},canonicalize-dmp,convert-stencil-to-ll-mlir{{{"tile-sizes="+",".join(["128"]*tile_sizes) if tile_sizes > 0 else ""} 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)