Visitor: allow visiting core instructions and intrinsics

This is mostly built on an extension of OpDescription.

- comment out references to dbg.addr, which is from fairly recent LLVM
- use std::variant and ArrayRef in OpDescription

Author: nhaehnle

example/ExampleMain.cpp

@@ -30,6 +30,8 @@
 
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/IR/DerivedTypes.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IRPrinter/IRPrintingPasses.h"
 #include "llvm/Support/CommandLine.h"
@@ -173,10 +175,24 @@ template <bool rpot> const Visitor<VisitorContainer> &getExampleVisitor() {
             b.add<xd::ReadOp>([](VisitorNest &self, xd::ReadOp &op) {
               *self.out << "visiting ReadOp: " << op << '\n';
             });
+            b.add<UnaryInstruction>(
+                [](VisitorNest &self, UnaryInstruction &inst) {
+                  *self.out << "visiting UnaryInstruction: " << inst << '\n';
+                });
+            b.add<BinaryOperator>([](VisitorNest &self, BinaryOperator &inst) {
+              *self.out << "visiting BinaryOperator: " << inst << '\n';
+            });
             b.nest<raw_ostream>([](VisitorBuilder<raw_ostream> &b) {
               b.add<xd::WriteOp>([](raw_ostream &out, xd::WriteOp &op) {
                 out << "visiting WriteOp: " << op << '\n';
               });
+              b.add<ReturnInst>([](raw_ostream &out, ReturnInst &ret) {
+                out << "visiting ReturnInst: " << ret << '\n';
+              });
+              b.addIntrinsic(
+                  Intrinsic::umax, [](raw_ostream &out, IntrinsicInst &umax) {
+                    out << "visiting umax intrinsic: " << umax << '\n';
+                  });
             });
             b.nest<VisitorInnermost>([](VisitorBuilder<VisitorInnermost> &b) {
               b.add<xd::ITruncOp>([](VisitorInnermost &inner,

include/llvm-dialects/Dialect/OpDescription.h

@@ -16,30 +16,59 @@
 
 #pragma once
 
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/StringRef.h"
 
+#include <vector>
+
 namespace llvm {
 class Function;
 class Instruction;
 } // namespace llvm
 
 namespace llvm_dialects {
 
-/// @brief Reflect an operation defined by a dialect
+/// @brief Reflect an operation.
 class OpDescription {
+public:
+  enum class Kind {
+    Core,
+    Dialect,
+    DialectWithOverloads,
+    Intrinsic,
+  };
+
 public:
   OpDescription(bool hasOverloads, llvm::StringRef mnemonic)
-      : m_hasOverloads(hasOverloads), m_mnemonic(mnemonic) {}
+      : m_kind(hasOverloads ? Kind::DialectWithOverloads : Kind::Dialect),
+        m_op(mnemonic) {}
+  OpDescription(Kind kind, unsigned opcode) : m_kind(kind), m_op(opcode) {}
+  OpDescription(Kind kind, llvm::MutableArrayRef<unsigned> opcodes);
 
   template <typename OpT>
   static const OpDescription& get();
 
+  Kind getKind() const { return m_kind; }
+  llvm::ArrayRef<unsigned> getOpcodes() const;
+
   bool matchInstruction(llvm::Instruction &inst) const;
   bool matchDeclaration(llvm::Function &decl) const;
 
+  bool canMatchDeclaration() const {
+    return m_kind == Kind::Dialect || m_kind == Kind::DialectWithOverloads ||
+           m_kind == Kind::Intrinsic;
+  }
+
 private:
-  bool m_hasOverloads;
-  llvm::StringRef m_mnemonic;
+  bool matchIntrinsic(unsigned intrinsicId) const;
+
+  Kind m_kind;
+
+  // Holds one of:
+  //  - core instruction opcode or intrinsic ID
+  //  - sorted array of opcodes or intrinsic IDs
+  //  - mnemonic
+  std::variant<unsigned, llvm::ArrayRef<unsigned>, llvm::StringRef> m_op;
 };
 
 } // namespace llvm_dialects

include/llvm-dialects/Dialect/Visitor.h

@@ -24,6 +24,7 @@
 namespace llvm {
 class Function;
 class Instruction;
+class IntrinsicInst;
 class Module;
 } // namespace llvm
 
@@ -90,6 +91,43 @@ namespace detail {
 
 class VisitorBase;
 
+/// @brief Key describing the condition of a visitor case
+class VisitorKey {
+  friend class VisitorBase;
+
+public:
+  template <typename OpT> static VisitorKey op() {
+    VisitorKey key{Kind::OpDescription};
+    key.m_description = &OpDescription::get<OpT>();
+    return key;
+  }
+
+  static VisitorKey intrinsic(unsigned id) {
+    VisitorKey key{Kind::Intrinsic};
+    key.m_intrinsicId = id;
+    return key;
+  }
+
+  bool matchInstruction(llvm::Instruction &inst) const;
+  bool matchDeclaration(llvm::Function &decl) const;
+
+  bool canMatchDeclaration() const {
+    return m_kind == Kind::Intrinsic || m_description->canMatchDeclaration();
+  }
+
+private:
+  enum class Kind {
+    OpDescription,
+    Intrinsic,
+  };
+
+  VisitorKey(Kind kind) : m_kind(kind) {}
+
+  Kind m_kind;
+  const OpDescription *m_description = nullptr;
+  unsigned m_intrinsicId = 0;
+};
+
 using VisitorCallback = void (void *, void *, llvm::Instruction *);
 using PayloadProjectionCallback = void *(void *);
 
@@ -101,13 +139,15 @@ struct PayloadProjection {
 };
 
 struct VisitorCase {
-  const OpDescription *description = nullptr;
+  VisitorKey key;
   VisitorCallback *callback = nullptr;
   void *callbackData = nullptr;
 
   // If non-negative, a byte offset to apply to the payload. If negative,
   // a shifted index into the projections vector.
   ssize_t projection = 0;
+
+  explicit VisitorCase(VisitorKey key) : key(key) {}
 };
 
 template <typename PayloadT, typename NestedPayloadT>
@@ -124,7 +164,7 @@ class VisitorBuilderBase {
 
   void setStrategy(VisitorStrategy strategy);
 
-  void add(const OpDescription &desc, void *extra, VisitorCallback *fn);
+  void add(VisitorKey key, void *extra, VisitorCallback *fn);
 
 public:
   PayloadProjectionCallback *m_projection = nullptr;
@@ -223,13 +263,13 @@ class VisitorBuilder : private detail::VisitorBuilderBase {
   Visitor<PayloadT> build() { return std::move(*this); }
 
   template <typename OpT> VisitorBuilder &add(void (*fn)(PayloadT &, OpT &)) {
-    VisitorBuilderBase::add(
-        OpDescription::get<OpT>(),
-        (void *)fn,
-        [](void *extra, void *payload, llvm::Instruction *op) {
-          auto fn = (void (*)(PayloadT &, OpT &))extra;
-          fn(*static_cast<PayloadT *>(payload), *llvm::cast<OpT>(op));
-        });
+    addCase<OpT>(detail::VisitorKey::op<OpT>(), fn);
+    return *this;
+  }
+
+  VisitorBuilder &addIntrinsic(unsigned id,
+                               void (*fn)(PayloadT &, llvm::IntrinsicInst &)) {
+    addCase<llvm::IntrinsicInst>(detail::VisitorKey::intrinsic(id), fn);
     return *this;
   }
 
@@ -258,6 +298,17 @@ class VisitorBuilder : private detail::VisitorBuilderBase {
 private:
   explicit VisitorBuilder(VisitorBuilderBase *parent)
       : VisitorBuilderBase(parent) {}
+
+  template <typename OpT>
+  void addCase(detail::VisitorKey key, void (*fn)(PayloadT &, OpT &)) {
+    VisitorBuilderBase::add(key, (void *)fn, &VisitorBuilder::forwarder<OpT>);
+  }
+
+  template <typename OpT>
+  static void forwarder(void *extra, void *payload, llvm::Instruction *op) {
+    auto fn = (void (*)(PayloadT &, OpT &))extra;
+    fn(*static_cast<PayloadT *>(payload), *llvm::cast<OpT>(op));
+  }
 };
 
 } // namespace llvm_dialects