[LV][LAA] Vectorize math lib calls with mem write-only attribute

Teach LAA to consider safe specific math lib calls which are known to
have set the memory write-only attribute. Those attributes are set to
calls by inferNonMandatoryLibFuncAttrs, in BuildLibCalls.cpp, and the
current ones are modf/modff and frexp/frexpf.

This happens only when the calls are found through TLI to have
vectorized counterparts.

Author: paschalis-mpeis

clang/test/CodeGen/aarch64-veclib-function-calls-linear-ptrs.c

@@ -17,7 +17,7 @@ vectorize.
 
 // CHECK-LABEL: define dso_local void @frexp_f64(
 // CHECK-SAME: ptr nocapture noundef readonly [[IN:%.*]], ptr nocapture noundef writeonly [[OUT1:%.*]], ptr nocapture noundef writeonly [[OUT2:%.*]], i32 noundef [[N:%.*]]) local_unnamed_addr #[[ATTR0:[0-9]+]] {
-// CHECK:    [[CALL:%.*]] = tail call double @frexp(double noundef [[TMP0:%.*]], ptr noundef [[ADD_PTR:%.*]]) #[[ATTR2:[0-9]+]]
+// CHECK:    [[CALL:%.*]] = tail call double @frexp(double noundef [[TMP0:%.*]], ptr noundef [[ADD_PTR:%.*]]) #[[ATTR5:[0-9]+]]
 //
 void frexp_f64(double *in, double *out1, int *out2, int N) {
   for (int i = 0; i < N; ++i)
@@ -26,30 +26,27 @@ void frexp_f64(double *in, double *out1, int *out2, int N) {
 
 // CHECK-LABEL: define dso_local void @frexp_f32(
 // CHECK-SAME: ptr nocapture noundef readonly [[IN:%.*]], ptr nocapture noundef writeonly [[OUT1:%.*]], ptr nocapture noundef writeonly [[OUT2:%.*]], i32 noundef [[N:%.*]]) local_unnamed_addr #[[ATTR0]] {
-// CHECK:    [[CALL:%.*]] = tail call float @frexpf(float noundef [[TMP0:%.*]], ptr noundef [[ADD_PTR:%.*]]) #[[ATTR2]]
+// CHECK:    [[CALL:%.*]] = tail call float @frexpf(float noundef [[TMP0:%.*]], ptr noundef [[ADD_PTR:%.*]]) #[[ATTR5]]
 //
 void frexp_f32(float *in, float *out1, int *out2, int N) {
   for (int i = 0; i < N; ++i)
     *out1 = frexpf(in[i], out2+i);
 }
 
-
-// TODO: LAA must allow vectorization.
-
 // CHECK-LABEL: define dso_local void @modf_f64(
 // CHECK-SAME: ptr nocapture noundef readonly [[IN:%.*]], ptr nocapture noundef writeonly [[OUT1:%.*]], ptr nocapture noundef writeonly [[OUT2:%.*]], i32 noundef [[N:%.*]]) local_unnamed_addr #[[ATTR0]] {
-// CHECK:    [[CALL:%.*]] = tail call double @modf(double noundef [[TMP0:%.*]], ptr noundef [[ADD_PTR:%.*]]) #[[ATTR3:[0-9]+]]
+// CHECK:    [[TMP11:%.*]] = tail call <vscale x 2 x double> @armpl_svmodf_f64_x(<vscale x 2 x double> [[WIDE_MASKED_LOAD:%.*]], ptr [[TMP10:%.*]], <vscale x 2 x i1> [[ACTIVE_LANE_MASK:%.*]])
+// CHECK:    [[CALL:%.*]] = tail call double @modf(double noundef [[TMP14:%.*]], ptr noundef [[ADD_PTR:%.*]]) #[[ATTR6:[0-9]+]]
 //
 void modf_f64(double *in, double *out1, double *out2, int N) {
   for (int i = 0; i < N; ++i)
       out1[i] = modf(in[i], out2+i);
 }
 
-// TODO: LAA must allow vectorization.
-
 // CHECK-LABEL: define dso_local void @modf_f32(
 // CHECK-SAME: ptr nocapture noundef readonly [[IN:%.*]], ptr nocapture noundef writeonly [[OUT1:%.*]], ptr nocapture noundef writeonly [[OUT2:%.*]], i32 noundef [[N:%.*]]) local_unnamed_addr #[[ATTR0]] {
-// CHECK:    [[CALL:%.*]] = tail call float @modff(float noundef [[TMP0:%.*]], ptr noundef [[ADD_PTR:%.*]]) #[[ATTR4:[0-9]+]]
+// CHECK:    [[TMP11:%.*]] = tail call <vscale x 4 x float> @armpl_svmodf_f32_x(<vscale x 4 x float> [[WIDE_MASKED_LOAD:%.*]], ptr [[TMP10:%.*]], <vscale x 4 x i1> [[ACTIVE_LANE_MASK:%.*]])
+// CHECK:    [[CALL:%.*]] = tail call float @modff(float noundef [[TMP14:%.*]], ptr noundef [[ADD_PTR:%.*]]) #[[ATTR7:[0-9]+]]
 //
 void modf_f32(float *in, float *out1, float *out2, int N) {
   for (int i = 0; i < N; ++i)

llvm/lib/Analysis/LoopAccessAnalysis.cpp

@@ -2309,6 +2309,20 @@ bool LoopAccessInfo::canAnalyzeLoop() {
   return true;
 }
 
+/// Returns whether \p I is a known math library call that has memory write-only
+/// attribute set.
+static bool isMathLibCallMemWriteOnly(const TargetLibraryInfo *TLI,
+                                      const Instruction &I) {
+  auto *Call = dyn_cast<CallInst>(&I);
+  if (!Call)
+    return false;
+
+  LibFunc Func;
+  TLI->getLibFunc(*Call, Func);
+  return Func == LibFunc::LibFunc_modf || Func == LibFunc::LibFunc_modff ||
+         Func == LibFunc::LibFunc_frexp || Func == LibFunc::LibFunc_frexpf;
+}
+
 void LoopAccessInfo::analyzeLoop(AAResults *AA, LoopInfo *LI,
                                  const TargetLibraryInfo *TLI,
                                  DominatorTree *DT) {
@@ -2405,6 +2419,11 @@ void LoopAccessInfo::analyzeLoop(AAResults *AA, LoopInfo *LI,
 
       // Save 'store' instructions. Abort if other instructions write to memory.
       if (I.mayWriteToMemory()) {
+        // We can safety handle math functions that have vectorized
+        // counterparts and have the memory write-only attribute set.
+        if (isMathLibCallMemWriteOnly(TLI, I))
+          continue;
+
         auto *St = dyn_cast<StoreInst>(&I);
         if (!St) {
           recordAnalysis("CantVectorizeInstruction", St)