[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>

Author: mcinally

clang/include/clang/Driver/Options.td

@@ -5480,7 +5480,7 @@ def mrecip_EQ : CommaJoined<["-"], "mrecip=">, Group<m_Group>,
            "<value> = ( ['!'] ['vec-'] ('rcp'|'sqrt') [('h'|'s'|'d')] [':'<n>] ) | 'all' | 'default' | 'none'">,
   MarshallingInfoStringVector<CodeGenOpts<"Reciprocals">>;
 def mprefer_vector_width_EQ : Joined<["-"], "mprefer-vector-width=">, Group<m_Group>,
-  Visibility<[ClangOption, CC1Option]>,
+  Visibility<[ClangOption, CC1Option, FlangOption, FC1Option]>,
   HelpText<"Specifies preferred vector width for auto-vectorization. Defaults to 'none' which allows target specific decisions.">,
   MarshallingInfoString<CodeGenOpts<"PreferVectorWidth">>;
 def mstack_protector_guard_EQ : Joined<["-"], "mstack-protector-guard=">, Group<m_Group>,

flang/include/flang/Frontend/CodeGenOptions.h

@@ -53,6 +53,9 @@ class CodeGenOptions : public CodeGenOptionsBase {
   /// 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;

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",

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 =

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))

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;

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>(

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");
 }

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'

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

@@ -1893,6 +1893,7 @@ def LLVM_LLVMFuncOp : LLVM_Op<"func", [
     OptionalAttr<FramePointerKindAttr>:$frame_pointer,
     OptionalAttr<StrAttr>:$target_cpu,
     OptionalAttr<StrAttr>:$tune_cpu,
+    OptionalAttr<StrAttr>:$prefer_vector_width,
     OptionalAttr<LLVM_TargetFeaturesAttr>:$target_features,
     OptionalAttr<BoolAttr>:$unsafe_fp_math,
     OptionalAttr<BoolAttr>:$no_infs_fp_math,

mlir/lib/Target/LLVMIR/ModuleImport.cpp

@@ -2636,6 +2636,10 @@ void ModuleImport::processFunctionAttributes(llvm::Function *func,
     funcOp.setTargetFeaturesAttr(
         LLVM::TargetFeaturesAttr::get(context, attr.getValueAsString()));
 
+  if (llvm::Attribute attr = func->getFnAttribute("prefer-vector-width");
+      attr.isStringAttribute())
+    funcOp.setPreferVectorWidth(attr.getValueAsString());
+
   if (llvm::Attribute attr = func->getFnAttribute("unsafe-fp-math");
       attr.isStringAttribute())
     funcOp.setUnsafeFpMath(attr.getValueAsBool());

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

@@ -1549,6 +1549,9 @@ LogicalResult ModuleTranslation::convertOneFunction(LLVMFuncOp func) {
   if (auto tuneCpu = func.getTuneCpu())
     llvmFunc->addFnAttr("tune-cpu", *tuneCpu);
 
+  if (auto preferVectorWidth = func.getPreferVectorWidth())
+    llvmFunc->addFnAttr("prefer-vector-width", *preferVectorWidth);
+
   if (auto attr = func.getVscaleRange())
     llvmFunc->addFnAttr(llvm::Attribute::getWithVScaleRangeArgs(
         getLLVMContext(), attr->getMinRange().getInt(),

mlir/test/Target/LLVMIR/Import/prefer-vector-width.ll

@@ -0,0 +1,9 @@
+; RUN: mlir-translate -import-llvm -split-input-file %s | FileCheck %s
+
+; CHECK-LABEL: llvm.func @prefer_vector_width()
+; CHECK-SAME: prefer_vector_width = "128"
+define void @prefer_vector_width() #0 {
+  ret void
+}
+
+attributes #0 = { "prefer-vector-width"="128" }

mlir/test/Target/LLVMIR/prefer-vector-width.mlir

@@ -0,0 +1,8 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// CHECK: define void @prefer_vector_width() #[[ATTRS:.*]] {
+// CHECK: attributes #[[ATTRS]] = { "prefer-vector-width"="128" }
+
+llvm.func @prefer_vector_width() attributes {prefer_vector_width = "128"} {
+  llvm.return
+}