[MLIR] Share TypeTree with MLIR as GenericTypeTree

enzyme/Enzyme/MLIR/Analysis/ActivityAnalysis.cpp

@@ -935,15 +935,17 @@
     // Even if returning a pointer, this instruction is considered inactive
     // since the instruction doesn't prop gradients. Thus, so long as we don't
     // return an object containing a float, this instruction is inactive
-    // if (!TR.intType(1, I, /*errifNotFound*/ false).isPossibleFloat()) {
-    //   if (EnzymePrintActivity)
-    //     llvm::errs()
-    //         << " constant instruction from known non-float non-writing "
-    //            "instruction "
-    //         << *I << "\n";
-    //   InsertConstantInstruction(TR, I);
-    //   return true;
-    // }
+    if (I->getNumResults() != 0 && llvm::all_of(I->getResults(), [&](Value v) {
+          return !TR.intType(1, v, /*errIfNotFound*/ false).isPossibleFloat();
+        })) {
+      if (EnzymePrintActivity)
+        llvm::errs()
+            << " constant instruction from known non-float non-writing "
+               "instruction "
+            << *I << "\n";
+      InsertConstantOperation(TR, I);
+      return true;
+    }
 
     // If all returned values constant otherwise, the operation is inactive
     if (llvm::all_of(I->getResults(),
@@ -1599,20 +1601,17 @@
     return true;
   }
 
-#if 0
-  // This value is certainly a pointer to an integer (and only and integer, not
+  // This value is certainly a pointer to an integer (and only an integer, not
   // a pointer or float). Therefore its value is constant
-  // TODO use typeInfo for more aggressive activity analysis
-  if (val->getType()->isPointerTy() &&
-      cast<PointerType>(val->getType())->isIntOrIntVectorTy() &&
-      TR.firstPointer(1, val, /*errifnotfound*/ false).isIntegral()) {
+  auto mutableTy = dyn_cast<AutoDiffTypeInterface>(Val.getType());
+  if (mutableTy && mutableTy.isMutable() &&
+      TR.firstPointer(1, Val, /*errIfNotFound*/ false).isIntegral()) {
     if (EnzymePrintActivity)
       llvm::errs() << " Value const as integral pointer" << (int)directions
-                   << " " << *val << "\n";
+                   << " " << Val << "\n";
     InsertConstantValue(TR, Val);
     return true;
   }
-#endif
 
   // TODO: since in MLIR globals are operations like others, this should be
   // handled via interfaces.
@@ -3184,16 +3183,16 @@
           }
           if (shouldContinue) {
             if (EnzymePrintActivity)
               llvm::errs() << "      -- continuing indirect store from " << val
                            << " into:\n";
             done.insert(std::make_tuple(SI.getOperation(), SI.getValue(), UA));
             for (Value TmpOrig : newAllocaSet) {
 
               for (Operation *a : TmpOrig.getUsers()) {
                 todo.push_back(std::make_tuple(a, TmpOrig, UA));
                 if (EnzymePrintActivity)
                   llvm::errs() << "         ** " << *a << "\n";
               }
               AllocaSet.insert(TmpOrig);
               shouldContinue = true;
             }

enzyme/Enzyme/MLIR/Interfaces/CMakeLists.txt

@@ -19,6 +19,7 @@ add_mlir_library(
   TransformUtils.cpp
   EnzymeLogic.cpp
   EnzymeLogicReverse.cpp
+  MTypeResults.cpp
   Utils.cpp
   DEPENDS
   MLIROffsetViewInterfaceIncGen

enzyme/Enzyme/MLIR/Interfaces/EnzymeLogic.h

@@ -4,9 +4,12 @@
 #include "mlir/IR/IRMapping.h"
 #include "mlir/Interfaces/FunctionInterfaces.h"
 
-#include "../../TypeAnalysis/TypeAnalysis.h"
+#include "../../TypeAnalysis/TypeTree.h"
 #include "../../Utils.h"
 #include <functional>
+#include <map>
+#include <string>
+#include <vector>
 
 namespace mlir {
 namespace enzyme {
@@ -27,21 +30,65 @@ class MTypeAnalysis {
   }
 };
 
+} // namespace enzyme
+} // namespace mlir
+
+template <> struct ConcreteTypeTraits<mlir::Type> {
+  static bool isFloatingPoint(mlir::Type T) { return static_cast<bool>(T); }
+
+  static void print(llvm::raw_ostream &OS, mlir::Type T) { OS << T; }
+
+  static std::string floatName(mlir::Type T) {
+    std::string S;
+    llvm::raw_string_ostream OS(S);
+    OS << T;
+    return S;
+  }
+
+  static bool less(mlir::Type A, mlir::Type B) {
+    return A.getAsOpaquePointer() < B.getAsOpaquePointer();
+  }
+};
+
+template <> struct TypeTreeTraits<mlir::Type> {
+  using Origin = mlir::Operation *;
+
+  // Mirror the defaults of -enzyme-max-type-depth and -enzyme-max-type-offset.
+  static unsigned maxDepth() { return 6; }
+
+  static int maxOffset() { return 500; }
+
+  template <typename TreeT>
+  static void depthLimitExceeded(const TreeT &Tree, Origin Orig,
+                                 const llvm::Twine &Detail) {
+    std::string Msg = ("not handling more than " + llvm::Twine(maxDepth()) +
+                       " pointer lookups deep dt:" + Tree.str() + Detail)
+                          .str();
+    if (Orig)
+      Orig->emitWarning(Msg);
+    else
+      llvm::errs() << Msg << "\n";
+  }
+};
+
+namespace mlir {
+namespace enzyme {
+
+using MConcreteType = GenericConcreteType<mlir::Type>;
+
+using MTypeTree = GenericTypeTree<mlir::Type>;
+
+// The MLIR counterpart of TypeResults. There is no type analysis behind it.
 class MTypeResults {
 public:
-  // TODO
-  TypeTree getReturnAnalysis() { return TypeTree(); }
-  TypeTree query(Value) const { return TypeTree(); }
-  ConcreteType intType(size_t num, Value val, bool errIfNotFound = true,
-                       bool pointerIntSame = false) const {
-    if (isa<IntegerType, IndexType>(val.getType())) {
-      return BaseType::Integer;
-    }
-    if (errIfNotFound) {
-      llvm_unreachable("something happened");
-    }
-    return BaseType::Unknown;
-  }
+  MTypeTree query(Value val) const;
+
+  /// What the value itself is.
+  MConcreteType intType(size_t num, Value val, bool errIfNotFound = true,
+                        bool pointerIntSame = false) const;
+
+  MConcreteType firstPointer(size_t num, Value val, bool errIfNotFound = true,
+                             bool pointerIntSame = false) const;
 };
 
 class MEnzymeLogic {

enzyme/Enzyme/MLIR/Interfaces/MTypeResults.cpp

@@ -0,0 +1,44 @@
+//===- MTypeResults.cpp - Type facts for MLIR activity analysis -----------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// MTypeResults answers the activity analysis's questions about what a value
+// holds. Nothing is inferred: every answer is read off the value's type, which
+// is exact where the type says what a value holds and Unknown elsewhere.
+//
+//===----------------------------------------------------------------------===//
+
+#include "EnzymeLogic.h"
+
+using namespace mlir;
+using namespace mlir::enzyme;
+
+MTypeTree MTypeResults::query(Value val) const {
+  // TODO: cache the tree per type in a DenseMap<Type, MTypeTree> once
+  // building it is more than a few isa checks.
+  MTypeTree tree;
+  return tree;
+}
+
+MConcreteType MTypeResults::intType(size_t num, Value val, bool errIfNotFound,
+                                    bool pointerIntSame) const {
+  if (isa<IntegerType, IndexType>(val.getType()))
+    return BaseType::Integer;
+  if (errIfNotFound)
+    llvm_unreachable("MTypeResults::intType: could not deduce the type of the "
+                     "value from its MLIR type");
+  return BaseType::Unknown;
+}
+
+MConcreteType MTypeResults::firstPointer(size_t num, Value val,
+                                         bool errIfNotFound,
+                                         bool pointerIntSame) const {
+  if (errIfNotFound)
+    llvm_unreachable("MTypeResults::firstPointer: could not deduce the pointee "
+                     "type of the value from its MLIR type");
+  return BaseType::Unknown;
+}