[flang] Add support for -mrecip[=<list>]

Cameron McInally · Cameron McInally · commit e7c60252b021 · 2025-05-30T08:37:10.000-07:00
This patch adds support for the -mrecip command line option. The parsing of
this options is equivalent to Clang's and it is implemented by setting the
"reciprocal-estimates" function attribute.
diff --git a/clang/include/clang/Driver/Options.td b/clang/include/clang/Driver/Options.td
@@ -5472,9 +5472,10 @@ def mno_implicit_float : Flag<["-"], "mno-implicit-float">, Group<m_Group>,
   HelpText<"Don't generate implicit floating point or vector instructions">;
 def mimplicit_float : Flag<["-"], "mimplicit-float">, Group<m_Group>;
 def mrecip : Flag<["-"], "mrecip">, Group<m_Group>,
+  Visibility<[ClangOption, CC1Option, FlangOption, FC1Option]>,
   HelpText<"Equivalent to '-mrecip=all'">;
 def mrecip_EQ : CommaJoined<["-"], "mrecip=">, Group<m_Group>,
-  Visibility<[ClangOption, CC1Option]>,
+  Visibility<[ClangOption, CC1Option, FlangOption, FC1Option]>,
   HelpText<"Control use of approximate reciprocal and reciprocal square root instructions followed by <n> iterations of "
            "Newton-Raphson refinement. "
            "<value> = ( ['!'] ['vec-'] ('rcp'|'sqrt') [('h'|'s'|'d')] [':'<n>] ) | 'all' | 'default' | 'none'">,
diff --git a/flang/include/flang/Frontend/CodeGenOptions.h b/flang/include/flang/Frontend/CodeGenOptions.h
@@ -56,6 +56,9 @@ class CodeGenOptions : public CodeGenOptionsBase {
   // The prefered vector width, if requested by -mprefer-vector-width.
   std::string PreferVectorWidth;
 
+  // List of reciprocal estimate sub-options.
+  std::string Reciprocals;
+
   /// List of filenames passed in using the -fembed-offload-object option. These
   /// are offloading binaries containing device images and metadata.
   std::vector<std::string> OffloadObjects;
diff --git a/flang/include/flang/Optimizer/Transforms/Passes.td b/flang/include/flang/Optimizer/Transforms/Passes.td
@@ -429,6 +429,9 @@ def FunctionAttr : Pass<"function-attr", "mlir::func::FuncOp"> {
               "module.">,
        Option<"unsafeFPMath", "unsafe-fp-math", "bool", /*default=*/"false",
               "Set the unsafe-fp-math attribute on functions in the module.">,
+       Option<"reciprocals", "mrecip", "std::string", /*default=*/"",
+              "Set the reciprocal-estimates attribute on functions in the "
+              "module.">,
        Option<"preferVectorWidth", "prefer-vector-width", "std::string",
               /*default=*/"",
               "Set the prefer-vector-width attribute on functions in the "
diff --git a/flang/include/flang/Tools/CrossToolHelpers.h b/flang/include/flang/Tools/CrossToolHelpers.h
@@ -102,6 +102,7 @@ struct MLIRToLLVMPassPipelineConfig : public FlangEPCallBacks {
     UnsafeFPMath = mathOpts.getAssociativeMath() &&
         mathOpts.getReciprocalMath() && NoSignedZerosFPMath &&
         ApproxFuncFPMath && mathOpts.getFPContractEnabled();
+    Reciprocals = opts.Reciprocals;
     PreferVectorWidth = opts.PreferVectorWidth;
     if (opts.InstrumentFunctions) {
       InstrumentFunctionEntry = "__cyg_profile_func_enter";
@@ -127,6 +128,8 @@ struct MLIRToLLVMPassPipelineConfig : public FlangEPCallBacks {
   bool NoSignedZerosFPMath =
       false; ///< Set no-signed-zeros-fp-math attribute for functions.
   bool UnsafeFPMath = false; ///< Set unsafe-fp-math attribute for functions.
+  std::string Reciprocals = ""; ///< Set reciprocal-estimate attribute for
+                                ///< functions.
   std::string PreferVectorWidth = ""; ///< Set prefer-vector-width attribute for
                                       ///< functions.
   bool NSWOnLoopVarInc = true; ///< Add nsw flag to loop variable increments.
diff --git a/flang/lib/Frontend/CompilerInvocation.cpp b/flang/lib/Frontend/CompilerInvocation.cpp
@@ -1251,6 +1251,164 @@ static bool parseIntegerOverflowArgs(CompilerInvocation &invoc,
   return true;
 }
 
+/// This is a helper function for validating the optional refinement step
+/// parameter in reciprocal argument strings. Return false if there is an error
+/// parsing the refinement step. Otherwise, return true and set the Position
+/// of the refinement step in the input string.
+///
+/// \param [in] in The input string
+/// \param [in] a The compiler invocation arguments to parse
+/// \param [out] position The position of the refinement step in input string
+/// \param [out] diags DiagnosticsEngine to report erros with
+static bool getRefinementStep(llvm::StringRef in, const llvm::opt::Arg &a,
+                              size_t &position,
+                              clang::DiagnosticsEngine &diags) {
+  const char refinementStepToken = ':';
+  position = in.find(refinementStepToken);
+  if (position != llvm::StringRef::npos) {
+    llvm::StringRef option = a.getOption().getName();
+    llvm::StringRef refStep = in.substr(position + 1);
+    // Allow exactly one numeric character for the additional refinement
+    // step parameter. This is reasonable for all currently-supported
+    // operations and architectures because we would expect that a larger value
+    // of refinement steps would cause the estimate "optimization" to
+    // under-perform the native operation. Also, if the estimate does not
+    // converge quickly, it probably will not ever converge, so further
+    // refinement steps will not produce a better answer.
+    if (refStep.size() != 1) {
+      diags.Report(clang::diag::err_drv_invalid_value) << option << refStep;
+      return false;
+    }
+    char refStepChar = refStep[0];
+    if (refStepChar < '0' || refStepChar > '9') {
+      diags.Report(clang::diag::err_drv_invalid_value) << option << refStep;
+      return false;
+    }
+  }
+  return true;
+}
+
+/// Parses all -mrecip=<list> arguments and populates the
+/// CompilerInvocation accordingly. Returns false if new errors are generated.
+///
+/// \param [out] invoc Stores the processed arguments
+/// \param [in] args The compiler invocation arguments to parse
+/// \param [out] diags DiagnosticsEngine to report erros with
+static bool parseMRecip(CompilerInvocation &invoc, llvm::opt::ArgList &args,
+                        clang::DiagnosticsEngine &diags) {
+  llvm::StringRef disabledPrefixIn = "!";
+  llvm::StringRef disabledPrefixOut = "!";
+  llvm::StringRef enabledPrefixOut = "";
+  llvm::StringRef out = "";
+  Fortran::frontend::CodeGenOptions &opts = invoc.getCodeGenOpts();
+
+  const llvm::opt::Arg *a =
+      args.getLastArg(clang::driver::options::OPT_mrecip,
+                      clang::driver::options::OPT_mrecip_EQ);
+  if (!a)
+    return true;
+
+  unsigned numOptions = a->getNumValues();
+  if (numOptions == 0) {
+    // No option is the same as "all".
+    opts.Reciprocals = "all";
+    return true;
+  }
+
+  // Pass through "all", "none", or "default" with an optional refinement step.
+  if (numOptions == 1) {
+    llvm::StringRef val = a->getValue(0);
+    size_t refStepLoc;
+    if (!getRefinementStep(val, *a, refStepLoc, diags))
+      return false;
+    llvm::StringRef valBase = val.slice(0, refStepLoc);
+    if (valBase == "all" || valBase == "none" || valBase == "default") {
+      opts.Reciprocals = args.MakeArgString(val);
+      return true;
+    }
+  }
+
+  // Each reciprocal type may be enabled or disabled individually.
+  // Check each input value for validity, concatenate them all back together,
+  // and pass through.
+
+  llvm::StringMap<bool> optionStrings;
+  optionStrings.insert(std::make_pair("divd", false));
+  optionStrings.insert(std::make_pair("divf", false));
+  optionStrings.insert(std::make_pair("divh", false));
+  optionStrings.insert(std::make_pair("vec-divd", false));
+  optionStrings.insert(std::make_pair("vec-divf", false));
+  optionStrings.insert(std::make_pair("vec-divh", false));
+  optionStrings.insert(std::make_pair("sqrtd", false));
+  optionStrings.insert(std::make_pair("sqrtf", false));
+  optionStrings.insert(std::make_pair("sqrth", false));
+  optionStrings.insert(std::make_pair("vec-sqrtd", false));
+  optionStrings.insert(std::make_pair("vec-sqrtf", false));
+  optionStrings.insert(std::make_pair("vec-sqrth", false));
+
+  for (unsigned i = 0; i != numOptions; ++i) {
+    llvm::StringRef val = a->getValue(i);
+
+    bool isDisabled = val.starts_with(disabledPrefixIn);
+    // Ignore the disablement token for string matching.
+    if (isDisabled)
+      val = val.substr(1);
+
+    size_t refStep;
+    if (!getRefinementStep(val, *a, refStep, diags))
+      return false;
+
+    llvm::StringRef valBase = val.slice(0, refStep);
+    llvm::StringMap<bool>::iterator optionIter = optionStrings.find(valBase);
+    if (optionIter == optionStrings.end()) {
+      // Try again specifying float suffix.
+      optionIter = optionStrings.find(valBase.str() + 'f');
+      if (optionIter == optionStrings.end()) {
+        // The input name did not match any known option string.
+        diags.Report(clang::diag::err_drv_invalid_value)
+            << a->getOption().getName() << val;
+        return false;
+      }
+      // The option was specified without a half or float or double suffix.
+      // Make sure that the double or half entry was not already specified.
+      // The float entry will be checked below.
+      if (optionStrings[valBase.str() + 'd'] ||
+          optionStrings[valBase.str() + 'h']) {
+        diags.Report(clang::diag::err_drv_invalid_value)
+            << a->getOption().getName() << val;
+        return false;
+      }
+    }
+
+    if (optionIter->second == true) {
+      // Duplicate option specified.
+      diags.Report(clang::diag::err_drv_invalid_value)
+          << a->getOption().getName() << val;
+      return false;
+    }
+
+    // Mark the matched option as found. Do not allow duplicate specifiers.
+    optionIter->second = true;
+
+    // If the precision was not specified, also mark the double and half entry
+    // as found.
+    if (valBase.back() != 'f' && valBase.back() != 'd' &&
+        valBase.back() != 'h') {
+      optionStrings[valBase.str() + 'd'] = true;
+      optionStrings[valBase.str() + 'h'] = true;
+    }
+
+    // Build the output string.
+    llvm::StringRef prefix = isDisabled ? disabledPrefixOut : enabledPrefixOut;
+    out = args.MakeArgString(out + prefix + val);
+    if (i != numOptions - 1)
+      out = args.MakeArgString(out + ",");
+  }
+
+  opts.Reciprocals = args.MakeArgString(out); // Handle the rest.
+  return true;
+}
+
 /// Parses all floating point related arguments and populates the
 /// CompilerInvocation accordingly.
 /// Returns false if new errors are generated.
@@ -1398,6 +1556,7 @@ static bool parseLangOptionsArgs(CompilerInvocation &invoc,
 
   success &= parseIntegerOverflowArgs(invoc, args, diags);
   success &= parseFloatingPointArgs(invoc, args, diags);
+  success &= parseMRecip(invoc, args, diags);
   success &= parseVScaleArgs(invoc, args, diags);
 
   return success;
diff --git a/flang/lib/Frontend/FrontendActions.cpp b/flang/lib/Frontend/FrontendActions.cpp
@@ -743,6 +743,8 @@ void CodeGenAction::generateLLVMIR() {
 
   config.PreferVectorWidth = opts.PreferVectorWidth;
 
+  config.Reciprocals = opts.Reciprocals;
+
   if (ci.getInvocation().getFrontendOpts().features.IsEnabled(
           Fortran::common::LanguageFeature::OpenMP))
     config.EnableOpenMP = true;
diff --git a/flang/lib/Optimizer/Passes/Pipelines.cpp b/flang/lib/Optimizer/Passes/Pipelines.cpp
@@ -358,7 +358,8 @@ void createDefaultFIRCodeGenPassPipeline(mlir::PassManager &pm,
       {framePointerKind, config.InstrumentFunctionEntry,
        config.InstrumentFunctionExit, config.NoInfsFPMath, config.NoNaNsFPMath,
        config.ApproxFuncFPMath, config.NoSignedZerosFPMath, config.UnsafeFPMath,
-       config.PreferVectorWidth, /*tuneCPU=*/"", setNoCapture, setNoAlias}));
+       config.Reciprocals, config.PreferVectorWidth, /*tuneCPU=*/"",
+       setNoCapture, setNoAlias}));
 
   if (config.EnableOpenMP) {
     pm.addNestedPass<mlir::func::FuncOp>(
diff --git a/flang/lib/Optimizer/Transforms/FunctionAttr.cpp b/flang/lib/Optimizer/Transforms/FunctionAttr.cpp
@@ -107,6 +107,10 @@ void FunctionAttrPass::runOnOperation() {
     func->setAttr(
         mlir::LLVM::LLVMFuncOp::getUnsafeFpMathAttrName(llvmFuncOpName),
         mlir::BoolAttr::get(context, true));
+  if (!reciprocals.empty())
+    func->setAttr(
+        mlir::LLVM::LLVMFuncOp::getReciprocalEstimatesAttrName(llvmFuncOpName),
+        mlir::StringAttr::get(context, reciprocals));
   if (!preferVectorWidth.empty())
     func->setAttr(
         mlir::LLVM::LLVMFuncOp::getPreferVectorWidthAttrName(llvmFuncOpName),
diff --git a/flang/test/Driver/mrecip.f90 b/flang/test/Driver/mrecip.f90
@@ -0,0 +1,27 @@
+! Test that -mrecip[=<list>] works as expected.
+
+! RUN: %flang_fc1 -emit-llvm -o - %s 2>&1| FileCheck %s --check-prefix=CHECK-OMIT
+! RUN: %flang_fc1 -mrecip -emit-llvm -o - %s 2>&1| FileCheck %s --check-prefix=CHECK-NOARG
+! RUN: %flang_fc1 -mrecip=all -emit-llvm -o - %s 2>&1| FileCheck %s --check-prefix=CHECK-ALL
+! RUN: %flang_fc1 -mrecip=none -emit-llvm -o - %s 2>&1| FileCheck %s --check-prefix=CHECK-NONE
+! RUN: %flang_fc1 -mrecip=default -emit-llvm -o - %s 2>&1| FileCheck %s --check-prefix=CHECK-DEF
+! RUN: %flang_fc1 -mrecip=divd,divf,divh,vec-divd,vec-divf,vec-divh,sqrtd,sqrtf,sqrth,vec-sqrtd,vec-sqrtf,vec-sqrth -emit-llvm -o - %s 2>&1| FileCheck %s --check-prefix=CHECK-POS
+! RUN: %flang_fc1 -mrecip=!divd,!divf,!divh,!vec-divd,!vec-divf,!vec-divh,!sqrtd,!sqrtf,!sqrth,!vec-sqrtd,!vec-sqrtf,!vec-sqrth
+! -emit-llvm -o - %s 2>&1| FileCheck %s --check-prefix=CHECK-NEG
+! RUN: %flang_fc1 -mrecip=!divd,divf,!divh,sqrtd,!sqrtf,sqrth -emit-llvm -o - %s 2>&1| FileCheck %s --check-prefix=CHECK-MIX
+! RUN: not %flang_fc1 -mrecip=xxx  -emit-llvm -o - %s 2>&1| FileCheck %s --check-prefix=CHECK-INV
+! RUN: not %flang_fc1 -mrecip=divd,divd  -emit-llvm -o - %s 2>&1| FileCheck %s --check-prefix=CHECK-DUP
+
+subroutine func
+end subroutine func
+
+! CHECK-OMIT-NOT: attributes #0 = { "reciprocal-estimates"={{.*}} }
+! CHECK-NOARG: attributes #0 = { "reciprocal-estimates"="all" }
+! CHECK-ALL: attributes #0 = { "reciprocal-estimates"="all" }
+! CHECK-NONE: attributes #0 = { "reciprocal-estimates"="none" }
+! CHECK-DEF: attributes #0 = { "reciprocal-estimates"="default" }
+! CHECK-POS: attributes #0 = { "reciprocal-estimates"="divd,divf,divh,vec-divd,vec-divf,vec-divh,sqrtd,sqrtf,sqrth,vec-sqrtd,vec-sqrtf,vec-sqrth" }
+! CHECK-NEG: attributes #0 = { "reciprocal-estimates"="!divd,!divf,!divh,!vec-divd,!vec-divf,!vec-divh,!sqrtd,!sqrtf,!sqrth,!vec-sqrtd,!vec-sqrtf,!vec-sqrth" }
+! CHECK-MIX: attributes #0 = { "reciprocal-estimates"="!divd,divf,!divh,sqrtd,!sqrtf,sqrth" }
+! CHECK-INV: error: invalid value 'xxx' in 'mrecip='
+! CHECK-DUP: error: invalid value 'divd' in 'mrecip='
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
@@ -1895,6 +1895,7 @@ def LLVM_LLVMFuncOp : LLVM_Op<"func", [
     OptionalAttr<StrAttr>:$tune_cpu,
     OptionalAttr<StrAttr>:$prefer_vector_width,
     OptionalAttr<LLVM_TargetFeaturesAttr>:$target_features,
+    OptionalAttr<StrAttr>:$reciprocal_estimates,
     OptionalAttr<BoolAttr>:$unsafe_fp_math,
     OptionalAttr<BoolAttr>:$no_infs_fp_math,
     OptionalAttr<BoolAttr>:$no_nans_fp_math,
diff --git a/mlir/lib/Target/LLVMIR/ModuleImport.cpp b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
@@ -2640,6 +2640,11 @@ void ModuleImport::processFunctionAttributes(llvm::Function *func,
       attr.isStringAttribute())
     funcOp.setPreferVectorWidth(attr.getValueAsString());
 
+  if (llvm::Attribute attr = func->getFnAttribute("reciprocal-estimates");
+      attr.isStringAttribute())
+    funcOp.setReciprocalEstimatesAttr(
+        StringAttr::get(context, attr.getValueAsString()));
+
   if (llvm::Attribute attr = func->getFnAttribute("unsafe-fp-math");
       attr.isStringAttribute())
     funcOp.setUnsafeFpMath(attr.getValueAsBool());
diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -1557,6 +1557,9 @@ LogicalResult ModuleTranslation::convertOneFunction(LLVMFuncOp func) {
         getLLVMContext(), attr->getMinRange().getInt(),
         attr->getMaxRange().getInt()));
 
+  if (auto reciprocalEstimates = func.getReciprocalEstimates())
+    llvmFunc->addFnAttr("reciprocal-estimates", *reciprocalEstimates);
+
   if (auto unsafeFpMath = func.getUnsafeFpMath())
     llvmFunc->addFnAttr("unsafe-fp-math", llvm::toStringRef(*unsafeFpMath));
 
diff --git a/mlir/test/Target/LLVMIR/Import/mrecip.ll b/mlir/test/Target/LLVMIR/Import/mrecip.ll
@@ -0,0 +1,9 @@
+; RUN: mlir-translate -import-llvm -split-input-file %s | FileCheck %s
+
+; CHECK-LABEL: llvm.func @reciprocal_estimates()
+; CHECK-SAME: reciprocal_estimates = "all"
+define void @reciprocal_estimates() #0 {
+  ret void
+}
+
+attributes #0 = { "reciprocal-estimates" = "all" }
diff --git a/mlir/test/Target/LLVMIR/mrecip.mlir b/mlir/test/Target/LLVMIR/mrecip.mlir
@@ -0,0 +1,8 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// CHECK: define void @reciprocal_estimates() #[[ATTRS:.*]] {
+// CHECK: attributes #[[ATTRS]] = { "reciprocal-estimates"="all" }
+
+llvm.func @reciprocal_estimates() attributes {reciprocal_estimates = "all"} {
+  llvm.return
+}
