[flang] Set KIND in compiler generated COUNT for SIZE(PACK) (#79801)

Compiler was rewriting SIZE(PACK(x, MASK)) to COUNT(MASK). It was
wrapping the COUNT call without a KIND argument (leading to INTEGER(4)
result in the characteristics) in an Expr<ExtentType> (implying
INTEGER(8) result), this lead to inconsistencies that later hit verifier
errors in lowering.

Set the KIND argument to the KIND of ExtentType to ensure the built
expression is consistent.

This requires giving access to some safe place where the "kind" name can
be saved and turned into a CharBlock (count has a DIM argument that
require using the KIND keyword here). For the FoldingContext that belong
to SemanticsContext, this is the same string set as the one used by
SemanticsContext for similar purposes.

Author: jeanPerier

flang/include/flang/Evaluate/common.h

@@ -20,6 +20,7 @@
 #include "flang/Parser/message.h"
 #include <cinttypes>
 #include <map>
+#include <set>
 #include <string>
 
 namespace Fortran::semantics {
@@ -217,26 +218,30 @@ class FoldingContext {
 public:
   FoldingContext(const common::IntrinsicTypeDefaultKinds &d,
       const IntrinsicProcTable &t, const TargetCharacteristics &c,
-      const common::LanguageFeatureControl &lfc)
+      const common::LanguageFeatureControl &lfc,
+      std::set<std::string> &tempNames)
       : defaults_{d}, intrinsics_{t}, targetCharacteristics_{c},
-        languageFeatures_{lfc} {}
+        languageFeatures_{lfc}, tempNames_{tempNames} {}
   FoldingContext(const parser::ContextualMessages &m,
       const common::IntrinsicTypeDefaultKinds &d, const IntrinsicProcTable &t,
-      const TargetCharacteristics &c, const common::LanguageFeatureControl &lfc)
+      const TargetCharacteristics &c, const common::LanguageFeatureControl &lfc,
+      std::set<std::string> &tempNames)
       : messages_{m}, defaults_{d}, intrinsics_{t}, targetCharacteristics_{c},
-        languageFeatures_{lfc} {}
+        languageFeatures_{lfc}, tempNames_{tempNames} {}
   FoldingContext(const FoldingContext &that)
       : messages_{that.messages_}, defaults_{that.defaults_},
         intrinsics_{that.intrinsics_},
         targetCharacteristics_{that.targetCharacteristics_},
         pdtInstance_{that.pdtInstance_}, impliedDos_{that.impliedDos_},
-        languageFeatures_{that.languageFeatures_} {}
+        languageFeatures_{that.languageFeatures_}, tempNames_{that.tempNames_} {
+  }
   FoldingContext(
       const FoldingContext &that, const parser::ContextualMessages &m)
       : messages_{m}, defaults_{that.defaults_}, intrinsics_{that.intrinsics_},
         targetCharacteristics_{that.targetCharacteristics_},
         pdtInstance_{that.pdtInstance_}, impliedDos_{that.impliedDos_},
-        languageFeatures_{that.languageFeatures_} {}
+        languageFeatures_{that.languageFeatures_}, tempNames_{that.tempNames_} {
+  }
 
   parser::ContextualMessages &messages() { return messages_; }
   const parser::ContextualMessages &messages() const { return messages_; }
@@ -273,6 +278,10 @@ class FoldingContext {
     return common::ScopedSet(pdtInstance_, nullptr);
   }
 
+  parser::CharBlock SaveTempName(std::string &&name) {
+    return {*tempNames_.emplace(std::move(name)).first};
+  }
+
 private:
   parser::ContextualMessages messages_;
   const common::IntrinsicTypeDefaultKinds &defaults_;
@@ -282,6 +291,7 @@ class FoldingContext {
   bool inModuleFile_{false};
   std::map<parser::CharBlock, ConstantSubscript> impliedDos_;
   const common::LanguageFeatureControl &languageFeatures_;
+  std::set<std::string> &tempNames_;
 };
 
 void RealFlagWarnings(FoldingContext &, const RealFlags &, const char *op);

flang/include/flang/Lower/Bridge.h

@@ -21,6 +21,7 @@
 #include "flang/Optimizer/Builder/FIRBuilder.h"
 #include "flang/Optimizer/Dialect/Support/KindMapping.h"
 #include "mlir/IR/BuiltinOps.h"
+#include <set>
 
 namespace llvm {
 class DataLayout;
@@ -111,7 +112,7 @@ class LoweringBridge {
   }
 
   /// Create a folding context. Careful: this is very expensive.
-  Fortran::evaluate::FoldingContext createFoldingContext() const;
+  Fortran::evaluate::FoldingContext createFoldingContext();
 
   Fortran::semantics::SemanticsContext &getSemanticsContext() const {
     return semanticsContext;
@@ -164,6 +165,7 @@ class LoweringBridge {
   const Fortran::lower::LoweringOptions &loweringOptions;
   const std::vector<Fortran::lower::EnvironmentDefault> &envDefaults;
   const Fortran::common::LanguageFeatureControl &languageFeatures;
+  std::set<std::string> tempNames;
 };
 
 } // namespace lower

flang/lib/Evaluate/shape.cpp

@@ -973,9 +973,14 @@ auto GetShapeHelper::operator()(const ProcedureRef &call) const -> Result {
               }
             }
           } else {
-            // Non-scalar MASK= -> [COUNT(mask)]
-            ActualArguments toCount{ActualArgument{common::Clone(
-                DEREF(call.arguments().at(1).value().UnwrapExpr()))}};
+            // Non-scalar MASK= -> [COUNT(mask, KIND=extent_kind)]
+            ActualArgument kindArg{
+                AsGenericExpr(Constant<ExtentType>{ExtentType::kind})};
+            kindArg.set_keyword(context_->SaveTempName("kind"));
+            ActualArguments toCount{
+                ActualArgument{common::Clone(
+                    DEREF(call.arguments().at(1).value().UnwrapExpr()))},
+                std::move(kindArg)};
             auto specific{context_->intrinsics().Probe(
                 CallCharacteristics{"count"}, toCount, *context_)};
             CHECK(specific);

flang/lib/Lower/Bridge.cpp

@@ -5058,9 +5058,9 @@ class FirConverter : public Fortran::lower::AbstractConverter {
 } // namespace
 
 Fortran::evaluate::FoldingContext
-Fortran::lower::LoweringBridge::createFoldingContext() const {
+Fortran::lower::LoweringBridge::createFoldingContext() {
   return {getDefaultKinds(), getIntrinsicTable(), getTargetCharacteristics(),
-          getLanguageFeatures()};
+          getLanguageFeatures(), tempNames};
 }
 
 void Fortran::lower::LoweringBridge::lower(

flang/lib/Semantics/semantics.cpp

@@ -313,7 +313,7 @@ SemanticsContext::SemanticsContext(
       globalScope_{*this}, intrinsicModulesScope_{globalScope_.MakeScope(
                                Scope::Kind::IntrinsicModules, nullptr)},
       foldingContext_{parser::ContextualMessages{&messages_}, defaultKinds_,
-          intrinsics_, targetCharacteristics_, languageFeatures_} {}
+          intrinsics_, targetCharacteristics_, languageFeatures_, tempNames_} {}
 
 SemanticsContext::~SemanticsContext() {}
 

flang/test/Evaluate/rewrite07.f90

@@ -0,0 +1,8 @@
+! RUN: %flang_fc1 -fdebug-unparse %s 2>&1 | FileCheck %s
+
+subroutine test_pack_size_rewrite(x, mask)
+    real :: x(:)
+    logical, intent(in) :: mask(:)
+    ! CHECK: CALL test(count(mask,kind=8_8))
+    call test(size(pack(x, mask), dim=1, kind=8))
+end subroutine

flang/unittests/Evaluate/expression.cpp

@@ -23,8 +23,9 @@ int main() {
   auto intrinsics{Fortran::evaluate::IntrinsicProcTable::Configure(defaults)};
   TargetCharacteristics targetCharacteristics;
   Fortran::common::LanguageFeatureControl languageFeatures;
+  std::set<std::string> tempNames;
   FoldingContext context{Fortran::parser::ContextualMessages{nullptr}, defaults,
-      intrinsics, targetCharacteristics, languageFeatures};
+      intrinsics, targetCharacteristics, languageFeatures, tempNames};
   ex1 = Fold(context, std::move(ex1));
   MATCH("-10_4", ex1.AsFortran());
   MATCH("1_4/2_4", (DefaultIntegerExpr{1} / DefaultIntegerExpr{2}).AsFortran());

flang/unittests/Evaluate/folding.cpp

@@ -50,10 +50,11 @@ void TestHostRuntimeSubnormalFlushing() {
     TargetCharacteristics noFlushingTargetCharacteristics;
     noFlushingTargetCharacteristics.set_areSubnormalsFlushedToZero(false);
     Fortran::common::LanguageFeatureControl languageFeatures;
+    std::set<std::string> tempNames;
     FoldingContext flushingContext{messages, defaults, intrinsics,
-        flushingTargetCharacteristics, languageFeatures};
+        flushingTargetCharacteristics, languageFeatures, tempNames};
     FoldingContext noFlushingContext{messages, defaults, intrinsics,
-        noFlushingTargetCharacteristics, languageFeatures};
+        noFlushingTargetCharacteristics, languageFeatures, tempNames};
 
     DynamicType r4{R4{}.GetType()};
     // Test subnormal argument flushing

flang/unittests/Evaluate/intrinsics.cpp

@@ -106,8 +106,8 @@ struct TestCall {
     auto messages{strings.Messages(buffer)};
     TargetCharacteristics targetCharacteristics;
     common::LanguageFeatureControl languageFeatures;
-    FoldingContext context{
-        messages, defaults, table, targetCharacteristics, languageFeatures};
+    FoldingContext context{messages, defaults, table, targetCharacteristics,
+        languageFeatures, tempNames};
     std::optional<SpecificCall> si{table.Probe(call, args, context)};
     if (resultType.has_value()) {
       TEST(si.has_value());
@@ -142,6 +142,7 @@ struct TestCall {
   ActualArguments args;
   std::string name;
   std::vector<std::string> keywords;
+  std::set<std::string> tempNames;
 };
 
 void TestIntrinsics() {