[flang-commits] [flang] [flang][Evaluate] Implement rewriting framework for evaluate::Expr (PR #153037)

Krzysztof Parzyszek via flang-commits flang-commits at lists.llvm.org
Mon Aug 11 10:51:20 PDT 2025


https://github.com/kparzysz updated https://github.com/llvm/llvm-project/pull/153037

>From 59ca403b295c706dbd15d4291b7fd34591403164 Mon Sep 17 00:00:00 2001
From: Krzysztof Parzyszek <Krzysztof.Parzyszek at amd.com>
Date: Fri, 8 Aug 2025 14:46:40 -0500
Subject: [PATCH 1/3] [flang][OpenMP] Refactor creating atomic analysis, NFC

Turn it into a class that combines the information and generates the
analysis instead of having independent functions do it.
---
 flang/lib/Semantics/check-omp-atomic.cpp | 161 ++++++++++++++---------
 1 file changed, 97 insertions(+), 64 deletions(-)

diff --git a/flang/lib/Semantics/check-omp-atomic.cpp b/flang/lib/Semantics/check-omp-atomic.cpp
index fcb0f9ad1e25d..a5fe820b1069b 100644
--- a/flang/lib/Semantics/check-omp-atomic.cpp
+++ b/flang/lib/Semantics/check-omp-atomic.cpp
@@ -222,47 +222,77 @@ static void SetAssignment(parser::AssignmentStmt::TypedAssignment &assign,
   }
 }
 
-static parser::OpenMPAtomicConstruct::Analysis::Op MakeAtomicAnalysisOp(
-    int what,
-    const std::optional<evaluate::Assignment> &maybeAssign = std::nullopt) {
-  parser::OpenMPAtomicConstruct::Analysis::Op operation;
-  operation.what = what;
-  SetAssignment(operation.assign, maybeAssign);
-  return operation;
-}
+namespace {
+struct AtomicAnalysis {
+  AtomicAnalysis(const SomeExpr &atom, const MaybeExpr &cond = std::nullopt)
+      : atom_(atom), cond_(cond) {}
+
+  AtomicAnalysis &addOp0(int what,
+      const std::optional<evaluate::Assignment> &maybeAssign = std::nullopt) {
+    return addOp(op0_, what, maybeAssign);
+  }
+  AtomicAnalysis &addOp1(int what,
+      const std::optional<evaluate::Assignment> &maybeAssign = std::nullopt) {
+    return addOp(op1_, what, maybeAssign);
+  }
 
-static parser::OpenMPAtomicConstruct::Analysis MakeAtomicAnalysis(
-    const SomeExpr &atom, const MaybeExpr &cond,
-    parser::OpenMPAtomicConstruct::Analysis::Op &&op0,
-    parser::OpenMPAtomicConstruct::Analysis::Op &&op1) {
-  // Defined in flang/include/flang/Parser/parse-tree.h
-  //
-  // struct Analysis {
-  //   struct Kind {
-  //     static constexpr int None = 0;
-  //     static constexpr int Read = 1;
-  //     static constexpr int Write = 2;
-  //     static constexpr int Update = Read | Write;
-  //     static constexpr int Action = 3; // Bits containing N, R, W, U
-  //     static constexpr int IfTrue = 4;
-  //     static constexpr int IfFalse = 8;
-  //     static constexpr int Condition = 12; // Bits containing IfTrue, IfFalse
-  //   };
-  //   struct Op {
-  //     int what;
-  //     TypedAssignment assign;
-  //   };
-  //   TypedExpr atom, cond;
-  //   Op op0, op1;
-  // };
-
-  parser::OpenMPAtomicConstruct::Analysis an;
-  SetExpr(an.atom, atom);
-  SetExpr(an.cond, cond);
-  an.op0 = std::move(op0);
-  an.op1 = std::move(op1);
-  return an;
-}
+  operator parser::OpenMPAtomicConstruct::Analysis() const {
+    // Defined in flang/include/flang/Parser/parse-tree.h
+    //
+    // struct Analysis {
+    //   struct Kind {
+    //     static constexpr int None = 0;
+    //     static constexpr int Read = 1;
+    //     static constexpr int Write = 2;
+    //     static constexpr int Update = Read | Write;
+    //     static constexpr int Action = 3; // Bits containing None, Read,
+    //                                      // Write, Update
+    //     static constexpr int IfTrue = 4;
+    //     static constexpr int IfFalse = 8;
+    //     static constexpr int Condition = 12; // Bits containing IfTrue,
+    //                                          // IfFalse
+    //   };
+    //   struct Op {
+    //     int what;
+    //     TypedAssignment assign;
+    //   };
+    //   TypedExpr atom, cond;
+    //   Op op0, op1;
+    // };
+
+    parser::OpenMPAtomicConstruct::Analysis an;
+    SetExpr(an.atom, atom_);
+    SetExpr(an.cond, cond_);
+    an.op0 = std::move(op0_);
+    an.op1 = std::move(op1_);
+    return an;
+  }
+
+private:
+  struct Op {
+    operator parser::OpenMPAtomicConstruct::Analysis::Op() const {
+      parser::OpenMPAtomicConstruct::Analysis::Op op;
+      op.what = what;
+      SetAssignment(op.assign, assign);
+      return op;
+    }
+
+    int what;
+    std::optional<evaluate::Assignment> assign;
+  };
+
+  AtomicAnalysis &addOp(Op &op, int what,
+      const std::optional<evaluate::Assignment> &maybeAssign) {
+    op.what = what;
+    op.assign = maybeAssign;
+    return *this;
+  }
+
+  const SomeExpr &atom_;
+  const MaybeExpr &cond_;
+  Op op0_, op1_;
+};
+} // namespace
 
 /// Check if `expr` satisfies the following conditions for x and v:
 ///
@@ -805,9 +835,9 @@ void OmpStructureChecker::CheckAtomicUpdateOnly(
       CheckAtomicUpdateAssignment(*maybeUpdate, action.source);
 
       using Analysis = parser::OpenMPAtomicConstruct::Analysis;
-      x.analysis = MakeAtomicAnalysis(atom, std::nullopt,
-          MakeAtomicAnalysisOp(Analysis::Update, maybeUpdate),
-          MakeAtomicAnalysisOp(Analysis::None));
+      x.analysis = AtomicAnalysis(atom)
+                       .addOp0(Analysis::Update, maybeUpdate)
+                       .addOp1(Analysis::None);
     } else if (!IsAssignment(action.stmt)) {
       context_.Say(
           source, "ATOMIC UPDATE operation should be an assignment"_err_en_US);
@@ -889,9 +919,11 @@ void OmpStructureChecker::CheckAtomicConditionalUpdate(
   }
 
   using Analysis = parser::OpenMPAtomicConstruct::Analysis;
-  x.analysis = MakeAtomicAnalysis(assign.lhs, update.cond,
-      MakeAtomicAnalysisOp(Analysis::Update | Analysis::IfTrue, assign),
-      MakeAtomicAnalysisOp(Analysis::None));
+  const SomeExpr &atom{assign.lhs};
+
+  x.analysis = AtomicAnalysis(atom, update.cond)
+                   .addOp0(Analysis::Update | Analysis::IfTrue, assign)
+                   .addOp1(Analysis::None);
 }
 
 void OmpStructureChecker::CheckAtomicUpdateCapture(
@@ -936,13 +968,13 @@ void OmpStructureChecker::CheckAtomicUpdateCapture(
   }
 
   if (GetActionStmt(&body.front()).stmt == uact.stmt) {
-    x.analysis = MakeAtomicAnalysis(atom, std::nullopt,
-        MakeAtomicAnalysisOp(action, update),
-        MakeAtomicAnalysisOp(Analysis::Read, capture));
+    x.analysis = AtomicAnalysis(atom)
+                     .addOp0(action, update)
+                     .addOp1(Analysis::Read, capture);
   } else {
-    x.analysis = MakeAtomicAnalysis(atom, std::nullopt,
-        MakeAtomicAnalysisOp(Analysis::Read, capture),
-        MakeAtomicAnalysisOp(action, update));
+    x.analysis = AtomicAnalysis(atom)
+                     .addOp0(Analysis::Read, capture)
+                     .addOp1(action, update);
   }
 }
 
@@ -1087,15 +1119,16 @@ void OmpStructureChecker::CheckAtomicConditionalUpdateCapture(
 
   evaluate::Assignment updAssign{*GetEvaluateAssignment(update.ift.stmt)};
   evaluate::Assignment capAssign{*GetEvaluateAssignment(capture.stmt)};
+  const SomeExpr &atom{updAssign.lhs};
 
   if (captureFirst) {
-    x.analysis = MakeAtomicAnalysis(updAssign.lhs, update.cond,
-        MakeAtomicAnalysisOp(Analysis::Read | captureWhen, capAssign),
-        MakeAtomicAnalysisOp(Analysis::Write | updateWhen, updAssign));
+    x.analysis = AtomicAnalysis(atom, update.cond)
+                     .addOp0(Analysis::Read | captureWhen, capAssign)
+                     .addOp1(Analysis::Write | updateWhen, updAssign);
   } else {
-    x.analysis = MakeAtomicAnalysis(updAssign.lhs, update.cond,
-        MakeAtomicAnalysisOp(Analysis::Write | updateWhen, updAssign),
-        MakeAtomicAnalysisOp(Analysis::Read | captureWhen, capAssign));
+    x.analysis = AtomicAnalysis(atom, update.cond)
+                     .addOp0(Analysis::Write | updateWhen, updAssign)
+                     .addOp1(Analysis::Read | captureWhen, capAssign);
   }
 }
 
@@ -1125,9 +1158,9 @@ void OmpStructureChecker::CheckAtomicRead(
       if (auto maybe{GetConvertInput(maybeRead->rhs)}) {
         const SomeExpr &atom{*maybe};
         using Analysis = parser::OpenMPAtomicConstruct::Analysis;
-        x.analysis = MakeAtomicAnalysis(atom, std::nullopt,
-            MakeAtomicAnalysisOp(Analysis::Read, maybeRead),
-            MakeAtomicAnalysisOp(Analysis::None));
+        x.analysis = AtomicAnalysis(atom)
+                         .addOp0(Analysis::Read, maybeRead)
+                         .addOp1(Analysis::None);
       }
     } else if (!IsAssignment(action.stmt)) {
       context_.Say(
@@ -1159,9 +1192,9 @@ void OmpStructureChecker::CheckAtomicWrite(
       CheckAtomicWriteAssignment(*maybeWrite, action.source);
 
       using Analysis = parser::OpenMPAtomicConstruct::Analysis;
-      x.analysis = MakeAtomicAnalysis(atom, std::nullopt,
-          MakeAtomicAnalysisOp(Analysis::Write, maybeWrite),
-          MakeAtomicAnalysisOp(Analysis::None));
+      x.analysis = AtomicAnalysis(atom)
+                       .addOp0(Analysis::Write, maybeWrite)
+                       .addOp1(Analysis::None);
     } else if (!IsAssignment(action.stmt)) {
       context_.Say(
           x.source, "ATOMIC WRITE operation should be an assignment"_err_en_US);

>From 2120125ad74caa41402f782a93a5a153520fdc1c Mon Sep 17 00:00:00 2001
From: Krzysztof Parzyszek <Krzysztof.Parzyszek at amd.com>
Date: Tue, 29 Jul 2025 16:17:29 -0500
Subject: [PATCH 2/3] [flang][Evaluate] Implement rewriting framework for
 evaluate::Expr

The structure of evaluate::Expr is highly customized for the specific
operation or entity that it represents. The different cases are expressed
with different types, which makes the traversal and modifications somewhat
complicated. There exists a framework for read-only traversal (traverse.h),
but there is nothing that helps with modifying evaluate::Expr.

It's rare that evaluate::Expr needs to be modified, but for the cases
where it needs to be, this code will make it easier.
---
 flang/include/flang/Evaluate/rewrite.h | 160 +++++++++++++++++++++++++
 1 file changed, 160 insertions(+)
 create mode 100644 flang/include/flang/Evaluate/rewrite.h

diff --git a/flang/include/flang/Evaluate/rewrite.h b/flang/include/flang/Evaluate/rewrite.h
new file mode 100644
index 0000000000000..034b1efa21977
--- /dev/null
+++ b/flang/include/flang/Evaluate/rewrite.h
@@ -0,0 +1,160 @@
+//===-- include/flang/Evaluate/rewrite.h ------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#ifndef FORTRAN_EVALUATE_REWRITE_H_
+#define FORTRAN_EVALUATE_REWRITE_H_
+
+#include "flang/Common/visit.h"
+#include "flang/Evaluate/expression.h"
+#include "flang/Support/Fortran.h"
+#include "llvm/ADT/STLExtras.h"
+
+#include <tuple>
+#include <type_traits>
+#include <utility>
+#include <variant>
+
+namespace Fortran::evaluate {
+namespace rewrite {
+namespace detail {
+template <typename, typename = void> //
+struct IsOperation {
+  static constexpr bool value{false};
+};
+
+template <typename T>
+struct IsOperation<T, std::void_t<decltype(T::operands)>> {
+  static constexpr bool value{true};
+};
+} // namespace detail
+
+template <typename T>
+constexpr bool is_operation_v{detail::IsOperation<T>::value};
+
+/// Individual Expr<T> rewriter that simply constructs an expression that is
+/// identical to the input. This is a suitable base class for all user-defined
+/// rewriters.
+struct Identity {
+  template <typename T, typename U>
+  Expr<T> operator()(Expr<T> &&x, const U &op) {
+    return std::move(x);
+  }
+};
+
+/// Bottom-up Expr<T> rewriter.
+///
+/// The Mutator traverses and reconstructs given Expr<T>. Going bottom-up,
+/// whenever the traversal visits a sub-node of type Expr<U> (for some U),
+/// it will invoke the user-provided rewriter via the () operator.
+///
+/// If x is of type Expr<U>, it will call (in pseudo-code):
+///   rewriter_(x, active_member_of(x.u))
+/// The second parameter is there to make it easier to overload the () operator
+/// for specific operations in Expr<...>.
+///
+/// The user rewriter is only invoked for Expr<U>, not for Operation, nor any
+/// other subobject.
+template <typename Rewriter> struct Mutator {
+  Mutator(Rewriter &rewriter) : rewriter_(rewriter) {}
+
+  template <typename T, typename U = llvm::remove_cvref_t<T>>
+  U operator()(T &&x) {
+    if constexpr (std::is_lvalue_reference_v<T>) {
+      return Mutate(U(x));
+    } else {
+      return Mutate(std::move(x));
+    }
+  }
+
+private:
+  template <typename T> struct LambdaWithRvalueCapture {
+    LambdaWithRvalueCapture(Rewriter &r, Expr<T> &&c)
+        : rewriter_(r), capture_(std::move(c)) {}
+    template <typename S> Expr<T> operator()(const S &s) {
+      return rewriter_(std::move(capture_), s);
+    }
+
+  private:
+    Rewriter &rewriter_;
+    Expr<T> &&capture_;
+  };
+
+  template <typename T, typename = std::enable_if_t<!is_operation_v<T>>>
+  T Mutate(T &&x) const {
+    return std::move(x);
+  }
+
+  template <typename D, typename = std::enable_if_t<is_operation_v<D>>>
+  D Mutate(D &&op, std::make_index_sequence<D::operands> t = {}) const {
+    return MutateOp(std::move(op), t);
+  }
+
+  template <typename T> //
+  Expr<T> Mutate(Expr<T> &&x) const {
+    // First construct the new expression with the rewritten op.
+    Expr<T> n{common::visit(
+        [&](auto &&s) { //
+          return Expr<T>(Mutate(std::move(s)));
+        },
+        std::move(x.u))};
+    // Return the rewritten expression. The second visit it to make sure
+    // that the second argument in the call to the rewriter is a part of
+    // the Expr<T> passed to it.
+    return common::visit(
+        LambdaWithRvalueCapture<T>(rewriter_, std::move(n)), std::move(n.u));
+  }
+
+  template <typename... Ts>
+  std::variant<Ts...> Mutate(std::variant<Ts...> &&u) const {
+    return common::visit(
+        [this](auto &&s) { return Mutate(std::move(s)); }, std::move(u));
+  }
+
+  template <typename... Ts>
+  std::tuple<Ts...> Mutate(std::tuple<Ts...> &&t) const {
+    return MutateTuple(std::move(t), std::index_sequence_for<Ts...>{});
+  }
+
+  template <typename... Ts, size_t... Is>
+  std::tuple<Ts...> MutateTuple(
+      std::tuple<Ts...> &&t, std::index_sequence<Is...>) const {
+    return std::make_tuple(Mutate(std::move(std::get<Is>(t))...));
+  }
+
+  template <typename D, size_t... Is>
+  D MutateOp(D &&op, std::index_sequence<Is...>) const {
+    return D(Mutate(std::move(op.template operand<Is>()))...);
+  }
+
+  template <typename T, size_t... Is>
+  Extremum<T> MutateOp(Extremum<T> &&op, std::index_sequence<Is...>) const {
+    return Extremum<T>(
+        op.ordering, Mutate(std::move(op.template operand<Is>()))...);
+  }
+
+  template <int K, size_t... Is>
+  ComplexComponent<K> MutateOp(
+      ComplexComponent<K> &&op, std::index_sequence<Is...>) const {
+    return ComplexComponent<K>(
+        op.isImaginaryPart, Mutate(std::move(op.template operand<Is>()))...);
+  }
+
+  template <int K, size_t... Is>
+  LogicalOperation<K> MutateOp(
+      LogicalOperation<K> &&op, std::index_sequence<Is...>) const {
+    return LogicalOperation<K>(
+        op.logicalOperator, Mutate(std::move(op.template operand<Is>()))...);
+  }
+
+  Rewriter &rewriter_;
+};
+
+template <typename Rewriter> Mutator(Rewriter &) -> Mutator<Rewriter>;
+} // namespace rewrite
+} // namespace Fortran::evaluate
+
+#endif // FORTRAN_EVALUATE_REWRITE_H_

>From d2ecbf084c01b991155092d03f1b892869f8a445 Mon Sep 17 00:00:00 2001
From: Krzysztof Parzyszek <Krzysztof.Parzyszek at amd.com>
Date: Mon, 11 Aug 2025 12:51:12 -0500
Subject: [PATCH 3/3] Update flang/include/flang/Evaluate/rewrite.h

Co-authored-by: Tom Eccles <tom.eccles at arm.com>
---
 flang/include/flang/Evaluate/rewrite.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/flang/include/flang/Evaluate/rewrite.h b/flang/include/flang/Evaluate/rewrite.h
index 034b1efa21977..50259cc0959f4 100644
--- a/flang/include/flang/Evaluate/rewrite.h
+++ b/flang/include/flang/Evaluate/rewrite.h
@@ -101,7 +101,7 @@ template <typename Rewriter> struct Mutator {
           return Expr<T>(Mutate(std::move(s)));
         },
         std::move(x.u))};
-    // Return the rewritten expression. The second visit it to make sure
+    // Return the rewritten expression. The second visit is to make sure
     // that the second argument in the call to the rewriter is a part of
     // the Expr<T> passed to it.
     return common::visit(



More information about the flang-commits mailing list