[flang] Add support for -mprefer-vector-width=<value> (#142073)

This patch adds support for the -mprefer-vector-width= command line
option. The parsing of this options is equivalent to Clang's and it is
implemented by setting the "prefer-vector-width" function attribute.

Co-authored-by: Cameron McInally <cmcinally@nvidia.com>

8 files changed, 47 insertions, 1 deletions

diff --git a/flang/include/flang/Frontend/CodeGenOptions.h b/flang/include/flang/Frontend/CodeGenOptions.h
index 2b4e823..61e56e5 100644
--- a/flang/include/flang/Frontend/CodeGenOptions.h
+++ b/flang/include/flang/Frontend/CodeGenOptions.h
@@ -53,6 +53,9 @@ public:
   /// The paths to the pass plugins that were registered using -fpass-plugin.
   std::vector<std::string> LLVMPassPlugins;
 
+  // The prefered vector width, if requested by -mprefer-vector-width.
+  std::string PreferVectorWidth;
+
   /// 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
index 7149353..1b19704 100644
--- a/flang/include/flang/Optimizer/Transforms/Passes.td
+++ b/flang/include/flang/Optimizer/Transforms/Passes.td
@@ -429,6 +429,10 @@ 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<"preferVectorWidth", "prefer-vector-width", "std::string",
+              /*default=*/"",
+              "Set the prefer-vector-width attribute on functions in the "
+              "module.">,
        Option<"tuneCPU", "tune-cpu", "std::string", /*default=*/"",
               "Set the tune-cpu attribute on functions in the module.">,
        Option<"setNoCapture", "set-nocapture", "bool", /*default=*/"false",
diff --git a/flang/include/flang/Tools/CrossToolHelpers.h b/flang/include/flang/Tools/CrossToolHelpers.h
index 118695b..058024a 100644
--- 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();
+    PreferVectorWidth = opts.PreferVectorWidth;
     if (opts.InstrumentFunctions) {
       InstrumentFunctionEntry = "__cyg_profile_func_enter";
       InstrumentFunctionExit = "__cyg_profile_func_exit";
@@ -126,6 +127,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 PreferVectorWidth = ""; ///< Set prefer-vector-width attribute for
+                                      ///< functions.
   bool NSWOnLoopVarInc = true; ///< Add nsw flag to loop variable increments.
   bool EnableOpenMP = false; ///< Enable OpenMP lowering.
   std::string InstrumentFunctionEntry =
diff --git a/flang/lib/Frontend/CompilerInvocation.cpp b/flang/lib/Frontend/CompilerInvocation.cpp
index ba25318..90a0029 100644
--- a/flang/lib/Frontend/CompilerInvocation.cpp
+++ b/flang/lib/Frontend/CompilerInvocation.cpp
@@ -309,6 +309,20 @@ static void parseCodeGenArgs(Fortran::frontend::CodeGenOptions &opts,
   for (auto *a : args.filtered(clang::driver::options::OPT_fpass_plugin_EQ))
     opts.LLVMPassPlugins.push_back(a->getValue());
 
+  // -mprefer_vector_width option
+  if (const llvm::opt::Arg *a = args.getLastArg(
+          clang::driver::options::OPT_mprefer_vector_width_EQ)) {
+    llvm::StringRef s = a->getValue();
+    unsigned width;
+    if (s == "none")
+      opts.PreferVectorWidth = "none";
+    else if (s.getAsInteger(10, width))
+      diags.Report(clang::diag::err_drv_invalid_value)
+          << a->getAsString(args) << a->getValue();
+    else
+      opts.PreferVectorWidth = s.str();
+  }
+
   // -fembed-offload-object option
   for (auto *a :
        args.filtered(clang::driver::options::OPT_fembed_offload_object_EQ))
diff --git a/flang/lib/Frontend/FrontendActions.cpp b/flang/lib/Frontend/FrontendActions.cpp
index 38dfaad..012d0fd 100644
--- a/flang/lib/Frontend/FrontendActions.cpp
+++ b/flang/lib/Frontend/FrontendActions.cpp
@@ -741,6 +741,8 @@ void CodeGenAction::generateLLVMIR() {
     config.VScaleMax = vsr->second;
   }
 
+  config.PreferVectorWidth = opts.PreferVectorWidth;
+
   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
index 378913f..0c774ee 100644
--- a/flang/lib/Optimizer/Passes/Pipelines.cpp
+++ b/flang/lib/Optimizer/Passes/Pipelines.cpp
@@ -358,7 +358,7 @@ void createDefaultFIRCodeGenPassPipeline(mlir::PassManager &pm,
       {framePointerKind, config.InstrumentFunctionEntry,
        config.InstrumentFunctionExit, config.NoInfsFPMath, config.NoNaNsFPMath,
        config.ApproxFuncFPMath, config.NoSignedZerosFPMath, config.UnsafeFPMath,
-       /*tuneCPU=*/"", setNoCapture, setNoAlias}));
+       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
index c8cdba0..041aa87 100644
--- 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 (!preferVectorWidth.empty())
+    func->setAttr(
+        mlir::LLVM::LLVMFuncOp::getPreferVectorWidthAttrName(llvmFuncOpName),
+        mlir::StringAttr::get(context, preferVectorWidth));
 
   LLVM_DEBUG(llvm::dbgs() << "=== End " DEBUG_TYPE " ===\n");
 }
diff --git a/flang/test/Driver/prefer-vector-width.f90 b/flang/test/Driver/prefer-vector-width.f90
new file mode 100644
index 0000000..d0f5fd2
--- /dev/null
+++ b/flang/test/Driver/prefer-vector-width.f90
@@ -0,0 +1,16 @@
+! Test that -mprefer-vector-width works as expected.
+
+! RUN: %flang_fc1 -emit-llvm -o - %s 2>&1| FileCheck %s --check-prefix=CHECK-DEF
+! RUN: %flang_fc1 -mprefer-vector-width=none -emit-llvm -o - %s 2>&1| FileCheck %s --check-prefix=CHECK-NONE
+! RUN: %flang_fc1 -mprefer-vector-width=128 -emit-llvm -o - %s 2>&1| FileCheck %s --check-prefix=CHECK-128
+! RUN: %flang_fc1 -mprefer-vector-width=256 -emit-llvm -o - %s 2>&1| FileCheck %s --check-prefix=CHECK-256
+! RUN: not %flang_fc1 -mprefer-vector-width=xxx  -emit-llvm -o - %s 2>&1| FileCheck %s --check-prefix=CHECK-INVALID
+
+subroutine func
+end subroutine func
+
+! CHECK-DEF-NOT: attributes #0 = { "prefer-vector-width"={{.*}} }
+! CHECK-NONE: attributes #0 = { "prefer-vector-width"="none" }
+! CHECK-128: attributes #0 = { "prefer-vector-width"="128" }
+! CHECK-256: attributes #0 = { "prefer-vector-width"="256" }
+! CHECK-INVALID:error: invalid value 'xxx' in '-mprefer-vector-width=xxx'