Add SsaContext: Template to access generic IR functionality

[nhaehnle]

Summary

Some functionality cannot be accessed directly from templates in a
straightforward generic way at all, and other functionality could be
accessed but only inefficiently.

Examples are:

• Given an SSA value, determining the basic block in which it is defined (if any) isn't possible in MachineIR without access to MachineRegisterInfo.

• Printing an llvm::Value properly requires a ModuleSlotTracker whose creation is expensive. Printing large functions can become prohibitively slow unless the ModuleSlotTracker is cached.

Each supported IR provides an implementation of the SsaContext
"concept" defined in this patch through which such functionality can be
exposed.

Change-Id: I7ed2b5c7698e6f3784db6201bcdb842173e23379

llvm/include/llvm/CodeGen/MachineSsaContext.h

+//===- MachineSsaContext.h - SsaContext for Machine IR ----------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+/// \file
+///
+/// This file defines the MachineSsaContext to allow generic algorithms to
+/// operate on MachineIR in SSA form.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CODEGEN_MACHINESSACONTEXT_H
+#define LLVM_CODEGEN_MACHINESSACONTEXT_H
+
+#include "llvm/CodeGen/MachineBasicBlock.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineInstr.h"
+#include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/Support/Handle.h"
+#include "llvm/Support/SsaContext.h"
+
+namespace llvm {
+
+namespace handle_detail {
+
+/// Specialization of \ref WrapperImpl for values (llvm::Register).
+template <>
+class WrapperImpl<SsaValueHandle, Register>
+    : public WrapperImplBase<SsaValueHandleTag> {
+public:
+  static SsaValueHandle wrapRef(Register value) {
+    // Physical registers are unsupported by design.
+    assert(!value.isValid() || value.isVirtual());
+    uintptr_t wrapped = value.id();
+    assert((wrapped != 0) == value.isValid());
+
+    // Guard against producing values reserved for DenseMap markers. This is de
+    // facto impossible, because it'd require 2^31 virtual registers to be in
+    // use on a 32-bit architecture.
+    assert(wrapped != (uintptr_t)-1 && wrapped != (uintptr_t)-2);
+
+    return make(reinterpret_cast<void *>(wrapped));
+  }
+
+  static Register unwrapRef(SsaValueHandle value) {
+    uintptr_t wrapped = reinterpret_cast<uintptr_t>(get(value));
+    return Register(wrapped);
+  }
+};
+
+} // namespace handle_detail
+
+/// \brief SsaContext for Machine IR in SSA form.
+class MachineSsaContext {
+public:
+  using BlockRef = MachineBasicBlock *;
+  using InstructionRef = MachineInstr *;
+  using ValueRef = Register;
+  using Wrapper = HandleWrapper<BlockHandle, BlockRef,
+                                InstructionHandle, InstructionRef,
+                                SsaValueHandle, ValueRef>;
+
+  explicit MachineSsaContext(MachineBasicBlock *block)
+      : m_regInfo(&block->getParent()->getRegInfo()) {}
+  explicit MachineSsaContext(MachineInstr *instr)
+      : m_regInfo(&instr->getParent()->getParent()->getRegInfo()) {}
+
+  /// Get the defining block of a value.
+  MachineBasicBlock *getDefBlock(ValueRef value) const {
+    if (!value)
+      return nullptr;
+    return m_regInfo->getVRegDef(value)->getParent();
+  }
+
+  Printable printableName(MachineBasicBlock *block) const;
+  Printable printable(MachineInstr *instruction) const;
+  Printable printable(Register value) const;
+
+private:
+  MachineRegisterInfo *m_regInfo;
+};
+
+template <> struct SsaContextForImpl<MachineBasicBlock *> {
+  using Context = MachineSsaContext;
+};
+template <> struct SsaContextForImpl<MachineInstr *> {
+  using Context = MachineSsaContext;
+};
+template <> struct SsaContextForImpl<Register> {
+  using Context = MachineSsaContext;
+};
+
+} // namespace llvm
+
+#endif // LLVM_CODEGEN_MACHINESSACONTEXT_H

llvm/include/llvm/IR/CFG.h

 //===- CFG.h ----------------------------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 /// \file

 #define LLVM_IR_CFG_H
 
 #include "llvm/ADT/GraphTraits.h"
 #include "llvm/ADT/iterator.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Value.h"
 #include "llvm/Support/Casting.h"
+#include "llvm/Support/Handle.h"
+#include "llvm/Support/SsaContext.h"
 #include <cassert>
 #include <cstddef>
 #include <iterator>
 
 namespace llvm {
 
 class BasicBlock;
 class Instruction;

 };
 template <> struct GraphTraits<Inverse<const Function*>> :
   public GraphTraits<Inverse<const BasicBlock*>> {
   static NodeRef getEntryNode(Inverse<const Function *> G) {
     return &G.Graph->getEntryBlock();
   }
 };
 
+//===----------------------------------------------------------------------===//
+// LLVM IR SsaContext
+//===----------------------------------------------------------------------===//
+
+/// \brief SsaContext for LLVM IR.
+class IrSsaContext {
+public:
+  using BlockRef = BasicBlock *;
+  using InstructionRef = Instruction *;
+  using ValueRef = Value *;
+  using Wrapper = HandleWrapper<BlockHandle, BlockRef,
+                                InstructionHandle, InstructionRef,
+                                SsaValueHandle, ValueRef>;
+
+  explicit IrSsaContext(BasicBlock *);
+  explicit IrSsaContext(Instruction *);
+  ~IrSsaContext();
+
+  BasicBlock *getDefBlock(Value *value) const {
+    if (auto *instruction = dyn_cast<Instruction>(value))
+      return instruction->getParent();
+    if (auto *argument = dyn_cast<Argument>(value))
+      return &argument->getParent()->getEntryBlock();
+    return nullptr;
+  }
+
+  Printable printableName(BasicBlock *block) const;
+  Printable printable(Value *value) const;
+
+private:
+  mutable std::unique_ptr<ModuleSlotTracker> m_moduleSlotTracker;
+
+  void ensureModuleSlotTracker(const Function &function) const;
+};
+
+template <> struct SsaContextForImpl<BasicBlock *> {
+  using Context = IrSsaContext;
+};
+template <> struct SsaContextForImpl<Instruction *> {
+  using Context = IrSsaContext;
+};
+template <> struct SsaContextForImpl<Value *> { using Context = IrSsaContext; };
+
 } // end namespace llvm
 
 #endif // LLVM_IR_CFG_H

llvm/include/llvm/Support/SsaContext.h

+//===- SsaContext.h ---------------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+/// \file
+///
+/// This file defines a set of classes around the SsaContext concept that allows
+/// code to abstract over some of the differences between SSA-form IRs beyond
+/// what is possible from templates via context-free functions like
+/// \ref llvm::GraphTraits<T>::child_begin()/child_end().
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_SUPPORT_SSACONTEXT_H
+#define LLVM_SUPPORT_SSACONTEXT_H
+
+#include "llvm/Support/Printable.h"
+
+namespace llvm {
+
+class Printable;
+
+/// Users of SsaContext should use \ref SsaContextFor instead.
+///
+/// This template is specialized by every supported IR. The specialization
+/// should exist for:
+///  - block "references", e.g. `llvm::BasicBlock *`
+///  - instruction "references", e.g. `llvm::MachineInstr *`
+///  - value "references", e.g. `mlir::Value`
+template <typename RefType> struct SsaContextForImpl {
+  // Context is the type implementing the SsaContext concept. It must be the
+  // same for every RefType of an IR.
+
+  // If you get a compiler error on the following line, it is because you're
+  // (indirectly?) using SsaContextFor<T> where T is not a supported IR object
+  // reference type, or the necessary specialization of this template is not
+  // in scope (probably due to a missing #include).
+  using Context = typename RefType::MissingSpecializationOfSsaContextForImpl;
+};
+
+/// \brief Discover the SsaContext implementation for an IR.
+///
+/// The template parameter must be an IR object reference type, e.g.
+/// llvm::BasicBlock* (for LLVM IR basic block) or mlir::Value (for MLIR
+/// values).
+///
+/// The context class must provide the following:
+///
+///   class TheSsaContext {
+///   public:
+///     using BlockRef = ...;          // e.g., llvm::BasicBlock *
+///     using InstructionRef = ...;    // e.g., llvm::Instruction *
+///     using ValueRef = ...;          // e.g., llvm::Value *
+///     using Wrapper = ...;           // suitable llvm::HandleWrapper
+///
+///     explicit TheSsaContext(BlockRef block);
+///     explicit TheSsaContext(InstructionRef instruction);
+///
+///     // Get the block in which a given value is defined. Returns a null-like
+///     // BlockRef if the value is not defined in a block (e.g. it is a
+///     // constant). Function arguments are defined in the function entry
+///     // block.
+///     BlockRef getDefBlock(ValueRef value) const;
+///
+///     Printable printableName(BlockRef ref) const;
+///     Printable printable(InstructionRef ref) const;
+///     Printable printable(ValueRef ref) const;
+///   };
+///
+template <typename RefType>
+using SsaContextFor = typename SsaContextForImpl<RefType>::Context;
+
+} // namespace llvm
+
+#endif // LLVM_SUPPORT_SSACONTEXT_H

llvm/lib/CodeGen/CMakeLists.txt

add_llvm_component_library(LLVMCodeGen
   MachinePipeliner.cpp
   MachinePostDominators.cpp
   MachineRegionInfo.cpp
   MachineRegisterInfo.cpp
   MachineScheduler.cpp
   MachineSink.cpp
   MachineSizeOpts.cpp
   MachineSSAUpdater.cpp
+  MachineSsaContext.cpp
   MachineStripDebug.cpp
   MachineTraceMetrics.cpp
   MachineVerifier.cpp
   ModuloSchedule.cpp
   MultiHazardRecognizer.cpp
   PatchableFunction.cpp
   MBFIWrapper.cpp
   MIRPrinter.cpp

llvm/lib/CodeGen/MachineSsaContext.cpp

+//===- MachineSsaContext.cpp ------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGen/MachineSsaContext.h"
+
+#include "llvm/IR/BasicBlock.h"
+
+using namespace llvm;
+
+Printable MachineSsaContext::printableName(MachineBasicBlock *block) const {
+  return Printable([block](raw_ostream &out) { block->printName(out); });
+}
+
+Printable MachineSsaContext::printable(MachineInstr *instruction) const {
+  return Printable(
+      [instruction](raw_ostream &out) { instruction->print(out); });
+}
+
+Printable MachineSsaContext::printable(Register value) const {
+  MachineRegisterInfo *mri = m_regInfo;
+  return Printable([mri, value](raw_ostream &out) {
+    out << printReg(value, mri->getTargetRegisterInfo(), 0, mri);
+
+    if (value) {
+      out << ": ";
+
+      MachineInstr *instr = mri->getUniqueVRegDef(value);
+      instr->print(out);
+    }
+  });
+}

llvm/lib/IR/CFG.cpp

+//===- CFG.cpp --------------------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/IR/CFG.h"
+
+#include "llvm/IR/ModuleSlotTracker.h"
+
+using namespace llvm;
+
+IrSsaContext::IrSsaContext(BasicBlock *) {}
+IrSsaContext::IrSsaContext(Instruction *) {}
+IrSsaContext::~IrSsaContext() {}
+
+Printable IrSsaContext::printable(Value *value) const {
+  if (!m_moduleSlotTracker) {
+    const Function *function = nullptr;
+
+    if (auto *instruction = dyn_cast<Instruction>(value)) {
+      function = instruction->getParent()->getParent();
+    } else if (auto *argument = dyn_cast<Argument>(value)) {
+      function = argument->getParent();
+    }
+
+    if (function)
+      ensureModuleSlotTracker(*function);
+  }
+
+  if (m_moduleSlotTracker) {
+    return Printable([&mst = *m_moduleSlotTracker, value](raw_ostream &out) {
+      value->print(out, mst, true);
+    });
+  } else {
+    return Printable([value](raw_ostream &out) { value->print(out, true); });
+  }
+}
+
+Printable IrSsaContext::printableName(BasicBlock *block) const {
+  if (block->hasName()) {
+    return Printable([block](raw_ostream &out) { out << block->getName(); });
+  } else {
+    ensureModuleSlotTracker(*block->getParent());
+    return Printable([&mst = *m_moduleSlotTracker, block](raw_ostream &out) {
+      out << mst.getLocalSlot(block);
+    });
+  }
+}
+
+void IrSsaContext::ensureModuleSlotTracker(const Function &function) const {
+  if (!m_moduleSlotTracker) {
+    m_moduleSlotTracker =
+        std::make_unique<ModuleSlotTracker>(function.getParent(), false);
+    m_moduleSlotTracker->incorporateFunction(function);
+  }
+}

llvm/lib/IR/CMakeLists.txt

 add_llvm_component_library(LLVMCore
   AbstractCallSite.cpp
   AsmWriter.cpp
   Attributes.cpp
   AutoUpgrade.cpp
   BasicBlock.cpp
+  CFG.cpp
   Comdat.cpp
   ConstantFold.cpp
   ConstantRange.cpp
   Constants.cpp
   Core.cpp
   DIBuilder.cpp
   DataLayout.cpp
   DebugInfo.cpp