Support for CUDA heap allocations

lib/type/DIFinder.cpp

@@ -35,6 +35,7 @@
 
 #include <cassert>
 #include <iterator>
+#include <optional>
 
 namespace llvm {
 class DbgVariableIntrinsic;
@@ -51,6 +52,15 @@ llvm::Value* get_alloca_for(const DbgVar* dbg_var) {
   return dbg_var->getVariableLocationOp(0);
 #endif
 }
+
+template <typename DbgVar>
+llvm::DIExpression* get_expr_for(const DbgVar* dbg_var) {
+#if LLVM_VERSION_MAJOR < 13
+  return dbg_var->getExpression();
+#else
+  return dbg_var->getExpression();
+#endif
+}
 }  // namespace compat
 
 #if LLVM_VERSION_MAJOR < 19
@@ -129,4 +139,18 @@ std::optional<llvm::DILocation*> find_location(const llvm::Instruction* inst) {
   return {};
 }
 
+std::optional<LocalAccessData> get_array_access_assignment(const llvm::CallBase* call) {
+  auto intrinsic = find_intrinsic(call);
+  if (intrinsic) {
+    auto* expr = compat::get_expr_for(intrinsic.value());
+    if (auto* array = llvm::dyn_cast<llvm::DIExpression>(expr)) {
+      if (array->isFragment()) {
+        return {LocalAccessData{intrinsic.value()->getVariable(), array}};
+      }
+    }
+    return {LocalAccessData{intrinsic.value()->getVariable(), {}}};
+  }
+  return {};
+}
+
 }  // namespace dimeta::difinder

lib/type/DIFinder.h

@@ -8,6 +8,7 @@
 #ifndef DIMETA_DIFINDER_H
 #define DIMETA_DIFINDER_H
 
+#include <llvm/IR/DebugInfoMetadata.h>
 #include <optional>
 
 namespace llvm {
@@ -17,14 +18,22 @@ class DILocalVariable;
 class DILocation;
 class CallBase;
 class Instruction;
+class DIExpression;
 }  // namespace llvm
 
 namespace dimeta::difinder {
 
+struct LocalAccessData {
+  llvm::DILocalVariable* var;
+  std::optional<llvm::DIExpression*> array_access;
+};
+
 std::optional<llvm::DILocalVariable*> find_local_variable(const llvm::Instruction* ai, bool bitcast_search = false);
 
 std::optional<llvm::DILocation*> find_location(const llvm::Instruction* inst);
 
+std::optional<LocalAccessData> get_array_access_assignment(const llvm::CallBase* call);
+
 }  // namespace dimeta::difinder
 
 #endif

lib/type/DIParser.cpp

@@ -11,8 +11,11 @@
 #include "DimetaData.h"
 #include "support/Logger.h"
 
+#include <cassert>
 #include <cstdlib>
 #include <llvm/BinaryFormat/Dwarf.h>
+#include <llvm/IR/DebugInfoMetadata.h>
+#include <llvm/Support/Casting.h>
 
 namespace dimeta::diparser {
 
@@ -39,6 +42,7 @@ class DIEventVisitor : public visitor::DINodeVisitor<DIEventVisitor> {
   bool visitCompositeType(const llvm::DICompositeType*);
 
   bool visitNode(const llvm::DINode* node);
+  void leaveNode(const llvm::DINode* node);
 
   bool visitRecurringCompositeType(const llvm::DICompositeType*);
 
@@ -88,6 +92,8 @@ bool DIEventVisitor::visitDerivedType(const llvm::DIDerivedType* derived_type) {
     current_.is_member_static = derived->isStaticMember();
   };
 
+  current_.derived_size = derived_type->getSizeInBits() / 8;
+
   const auto tag = derived_type->getTag();
   switch (tag) {
     case DW_TAG_member: {
@@ -168,6 +174,10 @@ bool DIEventVisitor::visitNode(const llvm::DINode* node) {
       }
       // LOG_FATAL(range_count);
     }
+  } else if (const auto* sub_routine = llvm::dyn_cast<llvm::DISubroutineType>(node)) {
+    // LOG_FATAL(sub_routine);
+    current_.type = const_cast<llvm::DISubroutineType*>(sub_routine);
+    events_.make_function_ptr(current_);
   }
   return true;
 }
@@ -243,6 +253,12 @@ void DIEventVisitor::leaveBasicType(const llvm::DIBasicType*) {
   current_ = state::MetaData{};
 }
 
+void DIEventVisitor::leaveNode(const llvm::DINode* node) {
+  if (llvm::isa<llvm::DISubroutineType>(node)) {
+    current_ = state::MetaData{};
+  }
+}
+
 bool DIEventVisitor::visitRecurringCompositeType(const llvm::DICompositeType* recurring_composite) {
   current_.is_recurring = true;
   return this->visitCompositeType(recurring_composite);

lib/type/DIParser.h

@@ -58,6 +58,8 @@ struct MetaData {
   bool is_void_ptr{false};
   bool is_vector{false};
   bool is_forward_decl{false};
+
+  Extent derived_size{0};
 };
 
 using MetaStack = llvm::SmallVector<MetaData, 4>;
@@ -67,6 +69,7 @@ using MetaStack = llvm::SmallVector<MetaData, 4>;
 class DIParseEvents {
  public:
   virtual void make_fundamental(const state::MetaData&)   = 0;
+  virtual void make_function_ptr(const state::MetaData&)  = 0;
   virtual void make_void_ptr(const state::MetaData&)      = 0;
   virtual void make_vtable(const state::MetaData&)        = 0;
   virtual void make_enum_member(const state::MetaData&)   = 0;

lib/type/DIRootType.cpp

@@ -121,6 +121,26 @@ std::optional<llvm::DIType*> type_of_call_argument(const dataflow::ValuePath& pa
   return {};
 }
 
+std::optional<llvm::DIType*> type_of_argument(const llvm::Argument& argument) {
+  if (auto* subprogram = argument.getParent()->getSubprogram(); subprogram != nullptr) {
+    const auto type_array = subprogram->getType()->getTypeArray();
+    const auto arg_pos    = [&](const auto arg_num) {
+      if (argument.hasStructRetAttr()) {
+        // return value is passed as argument at this point
+        return arg_num;  // see test cpp/heap_lhs_function_opt_nofwd.cpp
+      }
+      return arg_num + 1;
+    }(argument.getArgNo());
+
+    LOG_DEBUG(log::ditype_str(subprogram) << " -> " << argument)
+    LOG_DEBUG("Arg data: " << argument.getArgNo() << " Type num operands: " << type_array->getNumOperands())
+    assert(arg_pos < type_array.size() && "Arg position greater than DI type array of subprogram!");
+    return type_array[arg_pos];
+  }
+
+  return {};
+}
+
 }  // namespace helper
 
 std::optional<llvm::DIType*> find_type_root(const dataflow::CallValuePath& call_path) {
@@ -149,6 +169,14 @@ std::optional<llvm::DIType*> find_type_root(const dataflow::CallValuePath& call_
       return local_di_var.value()->getType();
     }
 
+    for (auto user : alloca->users()) {
+      if (auto store = llvm::dyn_cast<llvm::StoreInst>(user)) {
+        if (const auto* argument = llvm::dyn_cast<llvm::Argument>(store->getValueOperand())) {
+          return helper::type_of_argument(*argument);
+        }
+      }
+    }
+
     // see test heap_case_inheritance.cpp (e.g., returns several objects as base class pointer):
     // TODO: check if that ever applies to C, should probably only execute for C++ codes.
     LOG_DEBUG("Dataflow analysis of alloca")
@@ -225,23 +253,7 @@ std::optional<llvm::DIType*> find_type_root(const dataflow::CallValuePath& call_
   }
 
   if (const auto* argument = llvm::dyn_cast<llvm::Argument>(root_value)) {
-    if (auto* subprogram = argument->getParent()->getSubprogram(); subprogram != nullptr) {
-      const auto type_array = subprogram->getType()->getTypeArray();
-      const auto arg_pos    = [&](const auto arg_num) {
-        if (argument->hasStructRetAttr()) {
-          // return value is passed as argument at this point
-          return arg_num;  // see test cpp/heap_lhs_function_opt_nofwd.cpp
-        }
-        return arg_num + 1;
-      }(argument->getArgNo());
-
-      LOG_DEBUG(log::ditype_str(subprogram) << " -> " << *argument)
-      LOG_DEBUG("Arg data: " << argument->getArgNo() << " Type num operands: " << type_array->getNumOperands())
-      assert(arg_pos < type_array.size() && "Arg position greater than DI type array of subprogram!");
-      return type_array[arg_pos];
-    }
-
-    return {};
+    return helper::type_of_argument(*argument);
   }
 
   if (const auto* const_expr = llvm::dyn_cast<llvm::ConstantExpr>(root_value)) {

lib/type/DITypeExtractor.cpp

@@ -83,7 +83,7 @@ bool load_to(const llvm::LoadInst* load) {
   return detail::get_operand_to<T>(load).has_value();
 }
 
-std::optional<llvm::DIType*> reset_load_related_basic(const dataflow::ValuePath&, llvm::DIType* type_to_reset,
+std::optional<llvm::DIType*> reset_load_related_basic(const dataflow::ValuePath& path, llvm::DIType* type_to_reset,
                                                       const llvm::LoadInst* load) {
   auto* type = type_to_reset;
 
@@ -101,6 +101,25 @@ std::optional<llvm::DIType*> reset_load_related_basic(const dataflow::ValuePath&
     return type;
   }
 
+  // a (last?) load to a GEP of a composite likely loads the first member in an optimized context:
+  const bool last_load = path.start_value().value_or(nullptr) == load;
+  if (last_load && load_to<llvm::GetElementPtrInst>(load)) {
+    if (auto* may_be_member_type = llvm::dyn_cast<llvm::DIDerivedType>(type)) {
+      LOG_DEBUG("Load on GEP, return basetype " << log::ditype_str(may_be_member_type->getBaseType()))
+      if (di::util::is_member(*may_be_member_type)) {
+        auto type_of_member = may_be_member_type->getBaseType();
+        if (auto member = llvm::dyn_cast<llvm::DIDerivedType>(type_of_member)) {
+          may_be_member_type = member;
+          if (auto* member_composite_type = llvm::dyn_cast<llvm::DICompositeType>(member->getBaseType())) {
+            auto members_of_composite_type = di::util::get_composite_members(*member_composite_type);
+            assert(members_of_composite_type.size() > 0 && "Load to composite expects at least one member");
+            return (*members_of_composite_type.begin());
+          }
+        }
+      }
+    }
+  }
+
   if (auto* maybe_ptr_to_type = llvm::dyn_cast<llvm::DIDerivedType>(type)) {
     if (di::util::is_pointer(*maybe_ptr_to_type)) {
       LOG_DEBUG("Load of pointer-like " << log::ditype_str(maybe_ptr_to_type))

lib/type/DIUtil.cpp

@@ -173,6 +173,12 @@ bool is_non_static_member(const llvm::DINode& elem) {
          llvm::cast<llvm::DIType>(elem).getFlags() != llvm::DINode::DIFlags::FlagStaticMember;
 }
 
+bool is_member(const llvm::DINode& elem) {
+  const auto* type = llvm::dyn_cast<llvm::DIType>(&elem);
+  return elem.getTag() == llvm::dwarf::DW_TAG_member ||
+         ((type != nullptr) && (type->getFlags() == llvm::DINode::DIFlags::FlagStaticMember));
+}
+
 size_t get_num_composite_members(const llvm::DICompositeType& composite) {
   const auto num_members =
       llvm::count_if(composite.getElements(), [&](const auto* node) { return is_non_static_member(*node); });

lib/type/DIUtil.h

@@ -36,6 +36,7 @@ std::optional<StructMember> resolve_byte_offset_to_member_of(const llvm::DICompo
 bool is_pointer(const llvm::DIType& di_type);
 bool is_pointer_like(const llvm::DIType& di_type);
 bool is_non_static_member(const llvm::DINode& elem);
+bool is_member(const llvm::DINode& elem);
 size_t get_num_composite_members(const llvm::DICompositeType& composite);
 llvm::SmallVector<llvm::DIDerivedType*, 4> get_composite_members(const llvm::DICompositeType& composite);
 

lib/type/DIVisitor.h

@@ -47,8 +47,8 @@ class DINodeVisitor {
            invoke_if<DILocalVariable>(&DINodeVisitor::traverseVariable, std::forward<T>(type)) ||
            invoke_if<DIGlobalVariable>(&DINodeVisitor::traverseVariable, std::forward<T>(type)) ||
            invoke_if<DIEnumerator>(&DINodeVisitor::traverseNode, std::forward<T>(type)) ||
-           invoke_if<DISubrange>(&DINodeVisitor::traverseNode, std::forward<T>(type));
-    ;
+           invoke_if<DISubrange>(&DINodeVisitor::traverseNode, std::forward<T>(type)) ||
+           invoke_if<DISubroutineType>(&DINodeVisitor::traverseNode, std::forward<T>(type));
   }
 
  protected:
@@ -74,14 +74,13 @@ class DINodeVisitor {
   }
 
   bool traverseNode(const llvm::DINode* node) {
-    //    assert(
-    //        (llvm::isa<llvm::DIVariable>(node) || llvm::isa<llvm::DIType>(node) ||
-    //        llvm::isa<llvm::DIEnumerator>(node)) && "Can only visit variables or types");
-    if (const auto* type = llvm::dyn_cast<llvm::DIType>(node)) {
-      return traverseType(type);
-    }
-    if (const auto* var = llvm::dyn_cast<llvm::DIVariable>(node)) {
-      return traverseVariable(var);
+    if (!llvm::isa<llvm::DISubroutineType>(node)) {  // TODO give subroutine types their own iteration funcs
+      if (const auto* type = llvm::dyn_cast<llvm::DIType>(node)) {
+        return traverseType(type);
+      }
+      if (const auto* var = llvm::dyn_cast<llvm::DIVariable>(node)) {
+        return traverseVariable(var);
+      }
     }
 
     const auto exit = dimeta::util::create_scope_exit([&]() { get().leaveNode(node); });
@@ -198,6 +197,7 @@ class DINodeVisitor {
     ++depth_composite_;
     const auto exit = dimeta::util::create_scope_exit([&]() {
       get().leaveCompositeType(composite_type);
+      visited_dinodes_.erase(composite_type);
       assert(depth_composite_ > 0);
       --depth_composite_;
     });