Add Typetree to/from md

enzyme/Enzyme/CApi.cpp

@@ -711,6 +711,29 @@ void EnzymeExtractReturnInfo(EnzymeAugmentedReturnPtr ret, int64_t *data,
   }
 }
 
+static MDNode *extractMDNode(MetadataAsValue *MAV) {
+  Metadata *MD = MAV->getMetadata();
+  assert((isa<MDNode>(MD) || isa<ConstantAsMetadata>(MD)) &&
+         "Expected a metadata node or a canonicalized constant");
+
+  if (MDNode *N = dyn_cast<MDNode>(MD))
+    return N;
+
+  return MDNode::get(MAV->getContext(), MD);
+}
+
+CTypeTreeRef EnzymeTypeTreeFromMD(LLVMValueRef Val) {
+  TypeTree *Ret = new TypeTree();
+  MDNode *N = Val ? extractMDNode(unwrap<MetadataAsValue>(Val)) : nullptr;
+  Ret->insertFromMD(N);
+  return (CTypeTreeRef)N;
+}
+
+LLVMValueRef EnzymeTypeTreeToMD(CTypeTreeRef CTR, LLVMContextRef ctx) {
+  auto MD = ((TypeTree *)CTR)->toMD(*unwrap(ctx));
+  return wrap(MetadataAsValue::get(MD->getContext(), MD));
+}
+
 CTypeTreeRef EnzymeNewTypeTree() { return (CTypeTreeRef)(new TypeTree()); }
 CTypeTreeRef EnzymeNewTypeTreeCT(CConcreteType CT, LLVMContextRef ctx) {
   return (CTypeTreeRef)(new TypeTree(eunwrap(CT, *unwrap(ctx))));
@@ -850,7 +873,7 @@ void EnzymeSetStringMD(LLVMValueRef Inst, const char *Kind, LLVMValueRef Val) {
 LLVMValueRef EnzymeGetStringMD(LLVMValueRef Inst, const char *Kind) {
   auto *I = unwrap<Instruction>(Inst);
   assert(I && "Expected instruction");
-  if (auto *MD = I->getMetadata(KindID))
+  if (auto *MD = I->getMetadata(Kind))
     return wrap(MetadataAsValue::get(I->getContext(), MD));
   return nullptr;
 }

enzyme/Enzyme/FunctionUtils.cpp

@@ -494,6 +494,10 @@ UpgradeAllocasToMallocs(Function *NewF, DerivativeMode mode,
                     {ConstantAsMetadata::get(ConstantInt::get(
                         IntegerType::get(AI->getContext(), 64), align))}));
 
+    for (auto MD : {"enzyme_active", "enzyme_inactive", "enzyme_type"})
+      if (auto M = AI->getMetadata(MD))
+        CI->setMetadata(MD, M);
+
     if (rep != CI) {
       cast<Instruction>(rep)->setMetadata("enzyme_caststack",
                                           MDNode::get(CI->getContext(), {}));

enzyme/Enzyme/TypeAnalysis/BaseType.h

@@ -25,6 +25,7 @@
 #ifndef ENZYME_TYPE_ANALYSIS_BASE_TYPE_H
 #define ENZYME_TYPE_ANALYSIS_BASE_TYPE_H 1
 
+#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/ErrorHandling.h"
 #include <string>
 
@@ -60,7 +61,7 @@ static inline std::string to_string(BaseType t) {
 }
 
 /// Convert string to BaseType
-static inline BaseType parseBaseType(std::string str) {
+static inline BaseType parseBaseType(llvm::StringRef str) {
   if (str == "Integer")
     return BaseType::Integer;
   if (str == "Float")
@@ -73,4 +74,4 @@ static inline BaseType parseBaseType(std::string str) {
     return BaseType::Unknown;
   llvm_unreachable("Unknown BaseType string");
 }
-#endif
+#endif

enzyme/Enzyme/TypeAnalysis/ConcreteType.h

@@ -65,9 +65,9 @@ class ConcreteType {
   /// Construct a ConcreteType from a string
   ///  A Concrete Type's string representation is given by the string of the
   ///  enum If it is a floating point it is given by Float@<specific_type>
-  ConcreteType(std::string Str, llvm::LLVMContext &C) {
+  ConcreteType(llvm::StringRef Str, llvm::LLVMContext &C) {
     auto Sep = Str.find('@');
-    if (Sep != std::string::npos) {
+    if (Sep != llvm::StringRef::npos) {
       SubTypeEnum = BaseType::Float;
       assert(Str.substr(0, Sep) == "Float");
       auto SubName = Str.substr(Sep + 1);

enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp

@@ -788,6 +788,9 @@ void TypeAnalyzer::considerTBAA() {
 
   for (BasicBlock &BB : *fntypeinfo.Function) {
     for (Instruction &I : BB) {
+      if (auto MD = I.getMetadata("enzyme_type")) {
+        updateAnalysis(&I, TypeTree::fromMD(MD), &I);
+      }
 
       if (CallBase *call = dyn_cast<CallBase>(&I)) {
         Function *F = call->getCalledFunction();

enzyme/Enzyme/TypeAnalysis/TypeTree.h

@@ -1207,6 +1207,52 @@ class TypeTree : public std::enable_shared_from_this<TypeTree> {
     out += "}";
     return out;
   }
+
+  llvm::MDNode *toMD(llvm::LLVMContext &ctx) {
+    llvm::SmallVector<llvm::Metadata *, 1> subMD;
+    std::map<int, TypeTree> todo;
+    ConcreteType base(BaseType::Unknown);
+    for (auto &pair : mapping) {
+      if (pair.first.size() == 0) {
+        base = pair.second;
+        continue;
+      }
+      auto next(pair.first);
+      next.erase(next.begin());
+      todo[pair.first[0]].mapping.insert(std::make_pair(next, pair.second));
+    }
+    subMD.push_back(llvm::MDString::get(ctx, base.str()));
+    for (auto pair : todo) {
+      subMD.push_back(llvm::ConstantAsMetadata::get(
+          llvm::ConstantInt::get(llvm::IntegerType::get(ctx, 32), pair.first)));
+      subMD.push_back(pair.second.toMD(ctx));
+    }
+    return llvm::MDNode::get(ctx, subMD);
+  };
+
+  void insertFromMD(llvm::MDNode *md, const std::vector<int> &prev = {}) {
+    ConcreteType base(
+        llvm::cast<llvm::MDString>(md->getOperand(0))->getString(),
+        md->getContext());
+    if (base != BaseType::Unknown)
+      mapping.insert(std::make_pair(prev, base));
+    for (size_t i = 1; i < md->getNumOperands(); i += 2) {
+      auto off = llvm::cast<llvm::ConstantInt>(
+                     llvm::cast<llvm::ConstantAsMetadata>(md->getOperand(i))
+                         ->getValue())
+                     ->getSExtValue();
+      auto next(prev);
+      next.push_back((int)off);
+      insertFromMD(llvm::cast<llvm::MDNode>(md->getOperand(i + 1)), next);
+    }
+  }
+
+  static TypeTree fromMD(llvm::MDNode *md) {
+    TypeTree ret;
+    std::vector<int> off;
+    ret.insertFromMD(md, off);
+    return ret;
+  }
 };
 
 #endif