[flang-commits] [flang] [Flang][OpenMP] Move assert for wrapper syms and block args to genLoopNestOp (PR #103731)

[Sergio Afonso via flang-commits]

https://github.com/skatrak created https://github.com/llvm/llvm-project/pull/103731

This patch adds an assert to `genLoopNestClauses` to ensure the number of symbols and corresponding loop wrapper entry block arguments have the same size. This is checked by some of the callers, but it makes more sense moving it into the function itself and avoid having to replicate it.

>From ebfe3e7a233a000b814c851a047f978d28a438c3 Mon Sep 17 00:00:00 2001
From: Sergio Afonso <safonsof at amd.com>
Date: Wed, 14 Aug 2024 10:58:14 +0100
Subject: [PATCH] [Flang][OpenMP] Move assert for wrapper syms and block args
 to genLoopNestOp

This patch adds an assert to `genLoopNestClauses` to ensure the number of
symbols and corresponding loop wrapper entry block arguments have the same
size. This is checked by some of the callers, but it makes more sense moving
it into the function itself and avoid having to replicate it.
---
 flang/lib/Lower/OpenMP/OpenMP.cpp | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index 26825468df9b1d..97539d6b331068 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -1372,6 +1372,9 @@ genLoopNestOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
               llvm::ArrayRef<const semantics::Symbol *> wrapperSyms,
               llvm::ArrayRef<mlir::BlockArgument> wrapperArgs,
               llvm::omp::Directive directive, DataSharingProcessor &dsp) {
+  assert(wrapperSyms.size() == wrapperArgs.size() &&
+         "Number of symbols and wrapper block arguments must match");
+
   auto ivCallback = [&](mlir::Operation *op) {
     genLoopVars(op, converter, loc, iv, wrapperSyms, wrapperArgs);
     return llvm::SmallVector<const semantics::Symbol *>(iv);
@@ -2084,8 +2087,6 @@ static void genCompositeDistributeSimd(
       llvm::concat<mlir::BlockArgument>(distributeOp.getRegion().getArguments(),
                                         simdOp.getRegion().getArguments()));
 
-  assert(wrapperArgs.empty() &&
-         "Block args for omp.simd and omp.distribute currently not expected");
   genLoopNestOp(converter, symTable, semaCtx, eval, loc, queue, item,
                 loopNestClauseOps, iv, /*wrapperSyms=*/{}, wrapperArgs,
                 llvm::omp::Directive::OMPD_distribute_simd, dsp);
@@ -2133,8 +2134,6 @@ static void genCompositeDoSimd(lower::AbstractConverter &converter,
   auto wrapperArgs = llvm::to_vector(llvm::concat<mlir::BlockArgument>(
       wsloopOp.getRegion().getArguments(), simdOp.getRegion().getArguments()));
 
-  assert(wsloopReductionSyms.size() == wrapperArgs.size() &&
-         "Number of symbols and wrapper block arguments must match");
   genLoopNestOp(converter, symTable, semaCtx, eval, loc, queue, item,
                 loopNestClauseOps, iv, wsloopReductionSyms, wrapperArgs,
                 llvm::omp::Directive::OMPD_do_simd, dsp);