From 59575c87b0e492e2a4945616a4993608d2ed20a3 Mon Sep 17 00:00:00 2001
From: lujiao <lujiao@amd.com>
Date: Fri, 20 Mar 2026 18:28:44 +0800
Subject: [PATCH] Add class member function support for addSet

Add class member function support
---
 example/ExampleMain.cpp                 | 22 ++++++++++++++--------
 include/llvm-dialects/Dialect/Visitor.h | 15 +++++++++++++++
 test/example/visitor-basic.ll           |  4 +++-
 3 files changed, 32 insertions(+), 9 deletions(-)

diff --git a/example/ExampleMain.cpp b/example/ExampleMain.cpp
index 763a9ab..6dca04e 100644
--- a/example/ExampleMain.cpp
+++ b/example/ExampleMain.cpp
@@ -173,6 +173,10 @@ struct VisitorNest {
   void visitBinaryOperator(BinaryOperator &inst) {
     *out << "visiting BinaryOperator: " << inst << '\n';
   }
+  void visitInst(Instruction &inst) {
+    *out << "visiting llvm instruction: " << inst << '\n';
+  }
+
   VisitorResult visitUnaryInstruction(UnaryInstruction &inst) {
     *out << "visiting UnaryInstruction (pre): " << inst << '\n';
     return isa<LoadInst>(inst) ? VisitorResult::Stop : VisitorResult::Continue;
@@ -212,10 +216,10 @@ template <bool rpot> const Visitor<VisitorContainer> &getExampleVisitor() {
             });
             b.add(&VisitorNest::visitUnaryInstruction);
             b.add<xd::cpp::SetReadOp>([](VisitorNest &self, xd::cpp::SetReadOp &op) {
-              *self.out << "visiting SetReadOp: " << op << '\n';
-              return op.getType()->isIntegerTy(1) ? VisitorResult::Stop
-                                                  : VisitorResult::Continue;
-            });
+                  *self.out << "visiting SetReadOp: " << op << '\n';
+                  return op.getType()->isIntegerTy(1) ? VisitorResult::Stop
+                                                      : VisitorResult::Continue;
+                });
             b.addSet<xd::cpp::SetReadOp, xd::cpp::SetWriteOp>(
                 [](VisitorNest &self, llvm::Instruction &op) {
                   if (isa<xd::cpp::SetReadOp>(op)) {
@@ -224,11 +228,13 @@ template <bool rpot> const Visitor<VisitorContainer> &getExampleVisitor() {
                     *self.out << "visiting SetWriteOp (set): " << op << '\n';
                   }
                 });
+            b.addSet<xd::cpp::SetReadOp, xd::cpp::SetWriteOp>(
+                &VisitorNest::visitInst);
             b.addSet(complexSet, [](VisitorNest &self, llvm::Instruction &op) {
               assert((op.getOpcode() == Instruction::Ret ||
                       (isa<IntrinsicInst>(&op) &&
-                          cast<IntrinsicInst>(&op)->getIntrinsicID() ==
-                              Intrinsic::umin)) &&
+                       cast<IntrinsicInst>(&op)->getIntrinsicID() ==
+                           Intrinsic::umin)) &&
                      "Unexpected operation detected while visiting OpSet!");
 
               if (op.getOpcode() == Instruction::Ret) {
@@ -246,8 +252,8 @@ template <bool rpot> const Visitor<VisitorContainer> &getExampleVisitor() {
             b.add(&VisitorNest::visitBinaryOperator);
             b.nest<raw_ostream>([](VisitorBuilder<raw_ostream> &b) {
               b.add<xd::cpp::WriteOp>([](raw_ostream &out, xd::cpp::WriteOp &op) {
-                out << "visiting WriteOp: " << op << '\n';
-              });
+                    out << "visiting WriteOp: " << op << '\n';
+                  });
               b.add<xd::cpp::WriteVarArgOp>(
                   [](raw_ostream &out, xd::cpp::WriteVarArgOp &op) {
                     out << "visiting WriteVarArgOp: " << op << ":\n";
diff --git a/include/llvm-dialects/Dialect/Visitor.h b/include/llvm-dialects/Dialect/Visitor.h
index fd6e731..41568ac 100644
--- a/include/llvm-dialects/Dialect/Visitor.h
+++ b/include/llvm-dialects/Dialect/Visitor.h
@@ -414,6 +414,12 @@ class VisitorBuilder : private detail::VisitorBuilderBase {
     return *this;
   }
 
+  template <typename... OpTs>
+  VisitorBuilder &addSet(void (PayloadT::*fn)(llvm::Instruction &)) {
+    addSetMemberCase(detail::VisitorKey::opSet<OpTs...>(), fn);
+    return *this;
+  }
+
   VisitorBuilder &addSet(const OpSet &opSet,
                          VisitorResult (*fn)(PayloadT &,
                                              llvm::Instruction &I)) {
@@ -512,6 +518,15 @@ class VisitorBuilder : private detail::VisitorBuilderBase {
     VisitorBuilderBase::add(key, &VisitorBuilder::setForwarder<ReturnT>, data);
   }
 
+  template <typename ReturnT>
+  void addSetMemberCase(detail::VisitorKey key,
+                        ReturnT (PayloadT::*fn)(llvm::Instruction &)) {
+    detail::VisitorCallbackData data{};
+    static_assert(sizeof(fn) <= sizeof(data.data));
+    memcpy(&data.data, &fn, sizeof(fn));
+    VisitorBuilderBase::add(key, &VisitorBuilder::setForwarder<ReturnT>, data);
+  }
+
   template <typename OpT, typename ReturnT>
   void addMemberFnCase(detail::VisitorKey key, ReturnT (PayloadT::*fn)(OpT &)) {
     detail::VisitorCallbackData data{};
diff --git a/test/example/visitor-basic.ll b/test/example/visitor-basic.ll
index 091d73b..bdc1975 100644
--- a/test/example/visitor-basic.ll
+++ b/test/example/visitor-basic.ll
@@ -28,13 +28,15 @@
 ; DEFAULT-NEXT: visiting SetReadOp: %v.0 = call i1 @xd.ir.set.read__i1()
 ; DEFAULT-NEXT: visiting SetReadOp: %v.1 = call i32 @xd.ir.set.read__i32()
 ; DEFAULT-NEXT: visiting SetReadOp (set): %v.1 = call i32 @xd.ir.set.read__i32()
+; DEFAULT-NEXT: visiting llvm instruction:   %v.1 = call i32 @xd.ir.set.read__i32()
 ; DEFAULT-NEXT: visiting UnaryInstruction (pre): %v.2 = trunc i32 %v.1 to i8
 ; DEFAULT-NEXT: visiting UnaryInstruction: %v.2 = trunc i32 %v.1 to i8
 ; DEFAULT-NEXT: visiting SetWriteOp (set): call void (...) @xd.ir.set.write(i8 %v.2)
+; DEFAULT-NEXT: visiting llvm instruction:   call void (...) @xd.ir.set.write(i8 %v.2)
 ; DEFAULT-NEXT: visiting WriteVarArgOp: call void (...) @xd.ir.write.vararg(i8 %t, i32 %v2, i32 %q)
 ; DEFAULT-NEXT:   %v2 =
 ; DEFAULT-NEXT:   %q =
-; DEFAULT-NEXT: visiting umin (set): %vm = call i32 @llvm.umin.i32(i32 %v1, i32 %q) 
+; DEFAULT-NEXT: visiting umin (set): %vm = call i32 @llvm.umin.i32(i32 %v1, i32 %q)
 ; DEFAULT-NEXT: visiting StringAttrOp: Hello world!
 ; DEFAULT-NEXT: visiting CallInst:   %0 = call i32 @op.func(i32 %v1, i32 %q)
 ; DEFAULT-NEXT: visiting CallBrInst:   callbr void @callee()