[llvm] [VectorUtils] Add helper to get list of metadata to propagate (NFC). (PR #135003)

[Florian Hahn via llvm-commits]

https://github.com/fhahn created https://github.com/llvm/llvm-project/pull/135003

Split off the logic to filter out metadata that should not be propagated to a helper that populates a list with metadata kinds that should be preserved.

The current version just uses is_contained on an array to check if a metadata kind is supported. Given that most instructions will only have a small number of metadata kinds to start with, this shouldn't be worse than iterating over all kinds once and querying the instruction for individual metadata kinds, which involves quite a bit of indirection.

I plan ot use this utility in a follow-up patch in other places as well.

>From 5b32fe46694e32cfbc2cbc91858060031d77cc84 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Wed, 9 Apr 2025 13:44:53 +0100
Subject: [PATCH] [VectorUtils] Add helper to get list of metadata to propagate
 (NFC).

Split off the logic to filter out metadata that should not be propagated
to a helper that populates a list with metadata kinds that should be
preserved.

The current version just uses is_contained on an array to check if a
metadata kind is supported. Given that most instructions will only have
a small number of metadata kinds to start with, this shouldn't be worse
than iterating over all kinds once and querying the instruction for
individual metadata kinds, which involves quite a bit of indirection.

I plan ot use this utility in a follow-up patch in other places as well.
---
 llvm/lib/Analysis/VectorUtils.cpp | 37 +++++++++++++++++++++++--------
 1 file changed, 28 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/Analysis/VectorUtils.cpp b/llvm/lib/Analysis/VectorUtils.cpp
index 46f588f4c6705..1c0ea87990a3a 100644
--- a/llvm/lib/Analysis/VectorUtils.cpp
+++ b/llvm/lib/Analysis/VectorUtils.cpp
@@ -984,19 +984,38 @@ MDNode *llvm::intersectAccessGroups(const Instruction *Inst1,
   return MDNode::get(Ctx, Intersection);
 }
 
+/// Add metadata from \p Inst to \p Metadata, if it can be preserved after
+/// vectorization.
+static void getMetadataToPropagate(
+    Instruction *Inst,
+    SmallVectorImpl<std::pair<unsigned, MDNode *>> &Metadata) {
+  Inst->getAllMetadataOtherThanDebugLoc(Metadata);
+  unsigned SupportedIDs[] = {
+      LLVMContext::MD_tbaa,         LLVMContext::MD_alias_scope,
+      LLVMContext::MD_noalias,      LLVMContext::MD_fpmath,
+      LLVMContext::MD_nontemporal,  LLVMContext::MD_invariant_load,
+      LLVMContext::MD_access_group, LLVMContext::MD_mmra};
+
+  // Remove any unsupported metadata kinds from Metadata.
+  for (unsigned Idx = 0; Idx != Metadata.size();) {
+    if (is_contained(SupportedIDs, Metadata[Idx].first))
+      Idx++;
+    else {
+      // Swap element to end and remove it.
+      std::swap(Metadata[Idx], Metadata.back());
+      Metadata.pop_back();
+    }
+  }
+}
+
 /// \returns \p I after propagating metadata from \p VL.
 Instruction *llvm::propagateMetadata(Instruction *Inst, ArrayRef<Value *> VL) {
   if (VL.empty())
     return Inst;
-  Instruction *I0 = cast<Instruction>(VL[0]);
-  SmallVector<std::pair<unsigned, MDNode *>, 4> Metadata;
-  I0->getAllMetadataOtherThanDebugLoc(Metadata);
-
-  for (auto Kind : {LLVMContext::MD_tbaa, LLVMContext::MD_alias_scope,
-                    LLVMContext::MD_noalias, LLVMContext::MD_fpmath,
-                    LLVMContext::MD_nontemporal, LLVMContext::MD_invariant_load,
-                    LLVMContext::MD_access_group, LLVMContext::MD_mmra}) {
-    MDNode *MD = I0->getMetadata(Kind);
+  SmallVector<std::pair<unsigned, MDNode *>> Metadata;
+  getMetadataToPropagate(cast<Instruction>(VL[0]), Metadata);
+
+  for (auto &[Kind, MD] : Metadata) {
     for (int J = 1, E = VL.size(); MD && J != E; ++J) {
       const Instruction *IJ = cast<Instruction>(VL[J]);
       MDNode *IMD = IJ->getMetadata(Kind);