From 78604eca8cae4d00bb10fd5d6c648ef6caa5e707 Mon Sep 17 00:00:00 2001 From: Dominik Adamski Date: Tue, 23 Jul 2024 06:48:53 -0500 Subject: [PATCH 1/2] [Flang][OpenMP] Erase trip count for generic kernels We determine the type of the kernel based on the MLIR code, which changes during the lowering phase. Some kernels, such as those with multiple workshare loops, are initially classified as SPMD kernels, but are later recognized as generic kernels during PFT lowering. In such cases, we need to identify the change in type and clear the trip count if it was previously set. --- flang/lib/Lower/OpenMP/OpenMP.cpp | 34 ++++++++++++++++++------------- 1 file changed, 20 insertions(+), 14 deletions(-) diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp index 291e282c8b0b55..696b68f5547a0f 100644 --- a/flang/lib/Lower/OpenMP/OpenMP.cpp +++ b/flang/lib/Lower/OpenMP/OpenMP.cpp @@ -1678,21 +1678,27 @@ genLoopNestOp(lower::AbstractConverter &converter, lower::SymMap &symTable, firOpBuilder.getModule().getOperation()); auto targetOp = loopNestOp->getParentOfType(); - if (offloadMod && targetOp && !offloadMod.getIsTargetDevice() && - targetOp.isTargetSPMDLoop()) { - // Lower loop bounds and step, and process collapsing again, putting lowered - // values outside of omp.target this time. This enables calculating and - // accessing the trip count in the host, which is needed when lowering to - // LLVM IR via the OMPIRBuilder. - HostClausesInsertionGuard guard(firOpBuilder); - mlir::omp::CollapseClauseOps collapseClauseOps; - llvm::SmallVector iv; - ClauseProcessor cp(converter, semaCtx, item->clauses); - cp.processCollapse(loc, eval, collapseClauseOps, iv); - targetOp.getTripCountMutable().assign(calculateTripCount( - converter.getFirOpBuilder(), loc, collapseClauseOps)); + if (offloadMod && targetOp && !offloadMod.getIsTargetDevice()) { + if (targetOp.isTargetSPMDLoop()) { + // Lower loop bounds and step, and process collapsing again, putting + // lowered values outside of omp.target this time. This enables + // calculating and accessing the trip count in the host, which is needed + // when lowering to LLVM IR via the OMPIRBuilder. + HostClausesInsertionGuard guard(firOpBuilder); + mlir::omp::CollapseClauseOps collapseClauseOps; + llvm::SmallVector iv; + ClauseProcessor cp(converter, semaCtx, item->clauses); + cp.processCollapse(loc, eval, collapseClauseOps, iv); + targetOp.getTripCountMutable().assign(calculateTripCount( + converter.getFirOpBuilder(), loc, collapseClauseOps)); + } else { + // The MLIR target operation was updated during PFT lowering, + // and it is no longer an SPMD kernel. Erase the trip count because + // as it is now invalid. + if (targetOp.getTripCountMutable().size()) + targetOp.getTripCountMutable().erase(0); + } } - return loopNestOp; } From 0e8e02e6cee861ef568fcbebd18c855b8af1331a Mon Sep 17 00:00:00 2001 From: Dominik Adamski Date: Tue, 30 Jul 2024 21:16:40 +0200 Subject: [PATCH 2/2] Apply Sergio's remark --- flang/lib/Lower/OpenMP/OpenMP.cpp | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp index 696b68f5547a0f..c0d02dd4c5edb2 100644 --- a/flang/lib/Lower/OpenMP/OpenMP.cpp +++ b/flang/lib/Lower/OpenMP/OpenMP.cpp @@ -1691,12 +1691,11 @@ genLoopNestOp(lower::AbstractConverter &converter, lower::SymMap &symTable, cp.processCollapse(loc, eval, collapseClauseOps, iv); targetOp.getTripCountMutable().assign(calculateTripCount( converter.getFirOpBuilder(), loc, collapseClauseOps)); - } else { + } else if (targetOp.getTripCountMutable().size()) { // The MLIR target operation was updated during PFT lowering, // and it is no longer an SPMD kernel. Erase the trip count because // as it is now invalid. - if (targetOp.getTripCountMutable().size()) - targetOp.getTripCountMutable().erase(0); + targetOp.getTripCountMutable().erase(0); } } return loopNestOp;