[VPlan] Implement printing VPIRMetadata. (#168385)

fhahn · web-flow · commit 0e11a92447df · 2025-12-04T10:56:16.000Z
mplement printing for VPIRMetadata, using generic dyn_cast to VPIRMetadata. Depends on #166245 PR: #168385
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -1020,6 +1020,11 @@ class VPIRMetadata {
         find_if(Metadata, [Kind](const auto &P) { return P.first == Kind; });
     return It != Metadata.end() ? It->second : nullptr;
   }
+
+#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
+  /// Print metadata with node IDs.
+  void print(raw_ostream &O, VPSlotTracker &SlotTracker) const;
+#endif
 };
 
 /// This is a concrete Recipe that models a single VPlan-level instruction.
diff --git a/llvm/lib/Transforms/Vectorize/VPlanHelpers.h b/llvm/lib/Transforms/Vectorize/VPlanHelpers.h
@@ -394,21 +394,40 @@ class VPSlotTracker {
   /// require slot tracking.
   std::unique_ptr<ModuleSlotTracker> MST;
 
+  /// Cached metadata kind names from the Module's LLVMContext.
+  SmallVector<StringRef> MDNames;
+
+  /// Cached Module pointer for printing metadata.
+  const Module *M = nullptr;
+
   void assignName(const VPValue *V);
   LLVM_ABI_FOR_TEST void assignNames(const VPlan &Plan);
   void assignNames(const VPBasicBlock *VPBB);
   std::string getName(const Value *V);
 
 public:
   VPSlotTracker(const VPlan *Plan = nullptr) {
-    if (Plan)
+    if (Plan) {
       assignNames(*Plan);
+      if (auto *ScalarHeader = Plan->getScalarHeader())
+        M = ScalarHeader->getIRBasicBlock()->getModule();
+    }
   }
 
   /// Returns the name assigned to \p V, if there is one, otherwise try to
   /// construct one from the underlying value, if there's one; else return
   /// <badref>.
   std::string getOrCreateName(const VPValue *V) const;
+
+  /// Returns the cached metadata kind names.
+  ArrayRef<StringRef> getMDNames() {
+    if (MDNames.empty() && M)
+      M->getContext().getMDKindNames(MDNames);
+    return MDNames;
+  }
+
+  /// Returns the cached Module pointer.
+  const Module *getModule() const { return M; }
 };
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
diff --git a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
@@ -377,6 +377,9 @@ void VPRecipeBase::print(raw_ostream &O, const Twine &Indent,
     O << ", !dbg ";
     DL.print(O);
   }
+
+  if (auto *Metadata = dyn_cast<VPIRMetadata>(this))
+    Metadata->print(O, SlotTracker);
 }
 #endif
 
@@ -1599,6 +1602,25 @@ void VPIRMetadata::intersect(const VPIRMetadata &Other) {
   Metadata = std::move(MetadataIntersection);
 }
 
+#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
+void VPIRMetadata::print(raw_ostream &O, VPSlotTracker &SlotTracker) const {
+  const Module *M = SlotTracker.getModule();
+  if (Metadata.empty() || !M)
+    return;
+
+  ArrayRef<StringRef> MDNames = SlotTracker.getMDNames();
+  O << " (";
+  interleaveComma(Metadata, O, [&](const auto &KindNodePair) {
+    auto [Kind, Node] = KindNodePair;
+    assert(Kind < MDNames.size() && !MDNames[Kind].empty() &&
+           "Unexpected unnamed metadata kind");
+    O << "!" << MDNames[Kind] << " ";
+    Node->printAsOperand(O, M);
+  });
+  O << ")";
+}
+#endif
+
 void VPWidenCallRecipe::execute(VPTransformState &State) {
   assert(State.VF.isVector() && "not widening");
   assert(Variant != nullptr && "Can't create vector function.");
diff --git a/llvm/test/Transforms/LoopVectorize/vplan-printing-metadata.ll b/llvm/test/Transforms/LoopVectorize/vplan-printing-metadata.ll
@@ -7,11 +7,16 @@ define void @test_widen_metadata(ptr noalias %A, ptr noalias %B, i32 %n) {
 ; CHECK:      VPlan 'Initial VPlan for VF={4},UF>=1' {
 ; CHECK:      <x1> vector loop: {
 ; CHECK:        vector.body:
-; CHECK:          WIDEN ir<%lv> = load vp<{{.*}}>
-; CHECK:          WIDEN-CAST ir<%conv> = sitofp ir<%lv> to float
-; CHECK:          WIDEN ir<%mul> = fmul ir<%conv>, ir<2.000000e+00>
+; CHECK:          WIDEN ir<%lv> = load vp<{{.*}}> (!tbaa ![[TBAA:[0-9]+]])
+; CHECK:          WIDEN-CAST ir<%conv> = sitofp ir<%lv> to float (!fpmath ![[FPMATH:[0-9]+]])
+; CHECK:          WIDEN ir<%mul> = fmul ir<%conv>, ir<2.000000e+00> (!fpmath ![[FPMATH]])
 ; CHECK:          WIDEN-CAST ir<%conv.back> = fptosi ir<%mul> to i32
-; CHECK:          WIDEN store vp<{{.*}}>, ir<%conv.back>
+; CHECK:          WIDEN store vp<{{.*}}>, ir<%conv.back> (!tbaa ![[TBAA]])
+; CHECK:      ir-bb<loop>:
+; CHECK:        IR   %lv = load i32, ptr %gep.A, align 4, !tbaa ![[TBAA]]
+; CHECK:        IR   %conv = sitofp i32 %lv to float, !fpmath ![[FPMATH]]
+; CHECK:        IR   %mul = fmul float %conv, 2.000000e+00, !fpmath ![[FPMATH]]
+; CHECK:        IR   store i32 %conv.back, ptr %gep.B, align 4, !tbaa ![[TBAA]]
 ;
 entry:
   br label %loop
@@ -40,9 +45,13 @@ define void @test_intrinsic_with_metadata(ptr noalias %A, ptr noalias %B, i32 %n
 ; CHECK:      VPlan 'Initial VPlan for VF={4},UF>=1' {
 ; CHECK:      <x1> vector loop: {
 ; CHECK:        vector.body:
-; CHECK:          WIDEN ir<%lv> = load vp<{{.*}}>
-; CHECK:          WIDEN-INTRINSIC ir<%sqrt> = call llvm.sqrt(ir<%lv>)
-; CHECK:          WIDEN store vp<{{.*}}>, ir<%sqrt>
+; CHECK:          WIDEN ir<%lv> = load vp<{{.*}}> (!tbaa ![[TBAA2:[0-9]+]])
+; CHECK:          WIDEN-INTRINSIC ir<%sqrt> = call llvm.sqrt(ir<%lv>) (!fpmath ![[FPMATH2:[0-9]+]])
+; CHECK:          WIDEN store vp<{{.*}}>, ir<%sqrt> (!tbaa ![[TBAA2]])
+; CHECK:      ir-bb<loop>:
+; CHECK:        IR   %lv = load float, ptr %gep.A, align 4, !tbaa ![[TBAA2]]
+; CHECK:        IR   %sqrt = call float @llvm.sqrt.f32(float %lv), !fpmath ![[FPMATH2]]
+; CHECK:        IR   store float %sqrt, ptr %gep.B, align 4, !tbaa ![[TBAA2]]
 ;
 entry:
   br label %loop
@@ -67,11 +76,14 @@ define void @test_widen_with_multiple_metadata(ptr noalias %A, ptr noalias %B, i
 ; CHECK:      VPlan 'Initial VPlan for VF={4},UF>=1' {
 ; CHECK:      <x1> vector loop: {
 ; CHECK:        vector.body:
-; CHECK:          WIDEN ir<%lv> = load vp<{{.*}}>
+; CHECK:          WIDEN ir<%lv> = load vp<{{.*}}> (!tbaa ![[TBAA3:[0-9]+]])
 ; CHECK:          WIDEN-CAST ir<%conv> = sitofp ir<%lv> to float
 ; CHECK:          WIDEN ir<%mul> = fmul ir<%conv>, ir<2.000000e+00>
 ; CHECK:          WIDEN-CAST ir<%conv.back> = fptosi ir<%mul> to i32
-; CHECK:          WIDEN store vp<{{.*}}>, ir<%conv.back>
+; CHECK:          WIDEN store vp<{{.*}}>, ir<%conv.back> (!tbaa ![[TBAA3]])
+; CHECK:      ir-bb<loop>:
+; CHECK:        IR   %lv = load i32, ptr %gep.A, align 4, !tbaa ![[TBAA3]]
+; CHECK:        IR   store i32 %conv.back, ptr %gep.B, align 4, !tbaa ![[TBAA3]]
 ;
 entry:
   br label %loop
diff --git a/llvm/unittests/Transforms/Vectorize/VPlanTestBase.h b/llvm/unittests/Transforms/Vectorize/VPlanTestBase.h
@@ -86,15 +86,21 @@ class VPlanTestIRBase : public testing::Test {
 class VPlanTestBase : public testing::Test {
 protected:
   LLVMContext C;
-  std::unique_ptr<BasicBlock> ScalarHeader;
+  std::unique_ptr<Module> M;
+  Function *F;
+  BasicBlock *ScalarHeader;
   SmallVector<std::unique_ptr<VPlan>> Plans;
 
-  VPlanTestBase() : ScalarHeader(BasicBlock::Create(C, "scalar.header")) {
-    BranchInst::Create(&*ScalarHeader, &*ScalarHeader);
+  VPlanTestBase() {
+    M = std::make_unique<Module>("VPlanTestModule", C);
+    FunctionType *FTy = FunctionType::get(Type::getVoidTy(C), false);
+    F = Function::Create(FTy, GlobalValue::ExternalLinkage, "f", M.get());
+    ScalarHeader = BasicBlock::Create(C, "scalar.header", F);
+    BranchInst::Create(ScalarHeader, ScalarHeader);
   }
 
   VPlan &getPlan() {
-    Plans.push_back(std::make_unique<VPlan>(&*ScalarHeader));
+    Plans.push_back(std::make_unique<VPlan>(ScalarHeader));
     VPlan &Plan = *Plans.back();
     VPValue *DefaultTC = Plan.getConstantInt(32, 1024);
     Plan.setTripCount(DefaultTC);
