[RFC][llvm] Added llvm.loop.vectorize.reassociation.enable metadata.

This metadata allows unsafe reassociations of computations during
the loop vectorization. For example, it allows vectorizing loops
with floating-point reductions without the need to compile the whole
function/program with -fassociative-math.

Author: vzakhari

llvm/docs/LangRef.rst

@@ -7593,6 +7593,22 @@ Note that setting ``llvm.loop.interleave.count`` to 1 disables interleaving
 multiple iterations of the loop. If ``llvm.loop.interleave.count`` is set to 0
 then the interleave count will be determined automatically.
 
+'``llvm.loop.vectorize.reassociation.enable``' Metadata
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+This metadata selectively allows or disallows reassociating computations,
+which otherwise may be unsafe to reassociate, during the loop vectorization.
+For example, a floating point ``ADD`` reduction without ``reassoc`` fast-math
+flags may be vectorized provided that this metadata allows it. The first
+operand is the string ``llvm.loop.vectorize.reassociation.enable``
+and the second operand is a bit. If the bit operand value is 1 unsafe
+reassociations aqre enabled. A value of 0 disables unsafe reassociations.
+
+.. code-block:: llvm
+
+   !0 = !{!"llvm.loop.vectorize.reassociation.enable", i1 0}
+   !1 = !{!"llvm.loop.vectorize.reassociation.enable", i1 1}
+
 '``llvm.loop.vectorize.enable``' Metadata
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 

llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h

@@ -64,7 +64,8 @@ class LoopVectorizeHints {
     HK_FORCE,
     HK_ISVECTORIZED,
     HK_PREDICATE,
-    HK_SCALABLE
+    HK_SCALABLE,
+    HK_REASSOCIATE,
   };
 
   /// Hint - associates name and validation with the hint value.
@@ -97,6 +98,10 @@ class LoopVectorizeHints {
   /// Says whether we should use fixed width or scalable vectorization.
   Hint Scalable;
 
+  /// Says whether unsafe reassociation of computations is allowed
+  /// during the loop vectorization.
+  Hint Reassociate;
+
   /// Return the loop metadata prefix.
   static StringRef Prefix() { return "llvm.loop."; }
 
@@ -162,6 +167,13 @@ class LoopVectorizeHints {
     return (ScalableForceKind)Scalable.Value == SK_FixedWidthOnly;
   }
 
+  enum ForceKind getReassociate() const {
+    if ((ForceKind)Reassociate.Value == FK_Undefined &&
+        hasDisableAllTransformsHint(TheLoop))
+      return FK_Disabled;
+    return (ForceKind)Reassociate.Value;
+  }
+
   /// If hints are provided that force vectorization, use the AlwaysPrint
   /// pass name to force the frontend to print the diagnostic.
   const char *vectorizeAnalysisPassName() const;

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

@@ -97,6 +97,7 @@ bool LoopVectorizeHints::Hint::validate(unsigned Val) {
   case HK_ISVECTORIZED:
   case HK_PREDICATE:
   case HK_SCALABLE:
+  case HK_REASSOCIATE:
     return (Val == 0 || Val == 1);
   }
   return false;
@@ -112,6 +113,8 @@ LoopVectorizeHints::LoopVectorizeHints(const Loop *L,
       IsVectorized("isvectorized", 0, HK_ISVECTORIZED),
       Predicate("vectorize.predicate.enable", FK_Undefined, HK_PREDICATE),
       Scalable("vectorize.scalable.enable", SK_Unspecified, HK_SCALABLE),
+      Reassociate("vectorize.reassociation.enable", FK_Undefined,
+                  HK_REASSOCIATE),
       TheLoop(L), ORE(ORE) {
   // Populate values with existing loop metadata.
   getHintsFromMetadata();
@@ -251,6 +254,7 @@ bool LoopVectorizeHints::allowReordering() const {
   ElementCount EC = getWidth();
   return HintsAllowReordering &&
          (getForce() == LoopVectorizeHints::FK_Enabled ||
+          getReassociate() == LoopVectorizeHints::FK_Enabled ||
           EC.getKnownMinValue() > 1);
 }
 
@@ -300,8 +304,8 @@ void LoopVectorizeHints::setHint(StringRef Name, Metadata *Arg) {
     return;
   unsigned Val = C->getZExtValue();
 
-  Hint *Hints[] = {&Width,        &Interleave, &Force,
-                   &IsVectorized, &Predicate,  &Scalable};
+  Hint *Hints[] = {&Width,     &Interleave, &Force,      &IsVectorized,
+                   &Predicate, &Scalable,   &Reassociate};
   for (auto *H : Hints) {
     if (Name == H->Name) {
       if (H->validate(Val))

llvm/test/Transforms/LoopVectorize/reduction-reassociate.ll

@@ -0,0 +1,47 @@
+; Check that the loop with a floating-point reduction is vectorized
+; due to llvm.loop.vectorize.reassociation.enable metadata.
+; RUN: opt -passes=loop-vectorize -S < %s 2>&1 | FileCheck %s
+
+source_filename = "FIRModule"
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+; Function Attrs: nofree norecurse nosync nounwind memory(argmem: readwrite)
+define void @test_(ptr captures(none) %0, ptr readonly captures(none) %1) local_unnamed_addr #0 {
+; CHECK-LABEL: define void @test_(
+; CHECK-NEXT:    fadd contract <4 x float> {{.*}}
+; CHECK-NEXT:    call contract float @llvm.vector.reduce.fadd.v4f32(float -0.000000e+00, <4 x float> {{.*}})
+;
+  %invariant.gep = getelementptr i8, ptr %1, i64 -4
+  %.promoted = load float, ptr %0, align 4
+  br label %3
+
+3:                                                ; preds = %2, %3
+  %indvars.iv = phi i64 [ 1, %2 ], [ %indvars.iv.next, %3 ]
+  %4 = phi float [ %.promoted, %2 ], [ %6, %3 ]
+  %gep = getelementptr float, ptr %invariant.gep, i64 %indvars.iv
+  %5 = load float, ptr %gep, align 4
+  %6 = fadd contract float %4, %5
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %exitcond.not = icmp eq i64 %indvars.iv.next, 1001
+  br i1 %exitcond.not, label %7, label %3, !llvm.loop !2
+
+7:                                                ; preds = %3
+  %.lcssa = phi float [ %6, %3 ]
+  store float %.lcssa, ptr %0, align 4
+  ret void
+}
+
+attributes #0 = { nofree norecurse nosync nounwind memory(argmem: readwrite) "target-cpu"="x86-64" }
+
+!llvm.ident = !{!0}
+!llvm.module.flags = !{!1}
+
+!0 = !{!"flang version 21.0.0"}
+!1 = !{i32 2, !"Debug Info Version", i32 3}
+!2 = distinct !{!2, !3}
+!3 = !{!"llvm.loop.vectorize.reassociation.enable", i1 true}
+
+; CHECK-NOT: llvm.loop.vectorize.reassociation.enable
+; CHECK: [[META3]] = !{!"llvm.loop.isvectorized", i32 1}
+; CHECK: [[META4]] = !{!"llvm.loop.unroll.runtime.disable"}