[flang-commits] [flang] [flang][OpenMP] Simplify check for DISTRIBUTE/LINEAR restriction, NFC (PR #204813)

[Krzysztof Parzyszek via flang-commits]

https://github.com/kparzysz created https://github.com/llvm/llvm-project/pull/204813

Use `CollectAffectedDoLoops` instead of traversing the loop nest by hand.

>From b0e2a24b767c181aaf55e8b653fb80b1c1cf09df Mon Sep 17 00:00:00 2001
From: Krzysztof Parzyszek <Krzysztof.Parzyszek at amd.com>
Date: Fri, 19 Jun 2026 07:31:43 -0500
Subject: [PATCH] [flang][OpenMP] Simplify check for DISTRIBUTE/LINEAR
 restriction, NFC

Use `CollectAffectedDoLoops` instead of traversing the loop nest by hand.
---
 flang/lib/Semantics/check-omp-loop.cpp | 38 +++++---------------------
 1 file changed, 7 insertions(+), 31 deletions(-)

diff --git a/flang/lib/Semantics/check-omp-loop.cpp b/flang/lib/Semantics/check-omp-loop.cpp
index 689e3a0da89ca..7fa223a020c89 100644
--- a/flang/lib/Semantics/check-omp-loop.cpp
+++ b/flang/lib/Semantics/check-omp-loop.cpp
@@ -582,48 +582,24 @@ void OmpStructureChecker::CheckAssociatedLoopConstraints(
 
 void OmpStructureChecker::CheckDistLinear(
     const parser::OpenMPLoopConstruct &x) {
-  const parser::OmpClauseList &clauses{x.BeginDir().Clauses()};
-
+  unsigned version{context_.langOptions().OpenMPVersion};
   SymbolSourceMap indexVars;
 
   // Collect symbols of all the variables from linear clauses
-  for (auto &clause : clauses.v) {
+  for (auto &clause : x.BeginDir().Clauses().v) {
     if (std::get_if<parser::OmpClause::Linear>(&clause.u)) {
       GetSymbolsInObjectList(*parser::omp::GetOmpObjectList(clause), indexVars);
     }
   }
 
   if (!indexVars.empty()) {
-    // Get collapse level, if given, to find which loops are "associated."
-    std::int64_t collapseVal{GetOrdCollapseLevel(x)};
-    // Include the top loop if no collapse is specified
-    if (collapseVal == 0) {
-      collapseVal = 1;
-    }
-
-    // Match the loop index variables with the collected symbols from linear
-    // clauses.
-    for (auto &construct : std::get<parser::Block>(x.t)) {
-      std::int64_t curCollapseVal{collapseVal};
-      for (const parser::DoConstruct *loop{
-               parser::omp::GetDoConstruct(construct)};
-          loop;) {
-        if (loop->IsDoNormal()) {
-          const parser::Name &itrVal{GetLoopIndex(loop)};
-          if (itrVal.symbol) {
-            // Remove the symbol from the collected set
-            indexVars.erase(&itrVal.symbol->GetUltimate());
-          }
-          curCollapseVal--;
-          if (curCollapseVal == 0) {
-            break;
+    if (auto doLoops{CollectAffectedDoLoops(x, version, &context_)}) {
+      for (const parser::DoConstruct *loop : *doLoops) {
+        for (const omp::LoopControl &control : GetLoopControls(*loop)) {
+          if (control.iv.symbol) {
+            indexVars.erase(&control.iv.symbol->GetUltimate());
           }
         }
-        // Get the next DoConstruct if block is not empty.
-        const auto &block{std::get<parser::Block>(loop->t)};
-        const auto it{block.begin()};
-        loop = it != block.end() ? parser::Unwrap<parser::DoConstruct>(*it)
-                                 : nullptr;
       }
     }