memoryChecks.h

Go to the documentation of this file.

 
#pragma once
 
#include <tvm/api.h>
 
#include <Eigen/Core>
 
namespace tvm::utils
{
 
#ifdef NDEBUG
inline bool is_malloc_allowed() { return true; }
inline bool set_is_malloc_allowed(bool b) { return b; }
inline void override_is_malloc_allowed(bool) {}
inline void restore_is_malloc_allowed() {}
#else
namespace internal
{
TVM_DLLAPI bool is_malloc_allowed_();
TVM_DLLAPI bool set_is_malloc_allowed_(bool allow);
TVM_DLLAPI void override_is_malloc_allowed_(bool v);
TVM_DLLAPI void restore_is_malloc_allowed_();
 
inline void check_malloc_coherency()
{
#  ifdef EIGEN_RUNTIME_NO_MALLOC
  assert(Eigen::internal::is_malloc_allowed() == is_malloc_allowed_()
         && "Eigen::internal::is_malloc_allowed() at the caller site and tvm::utils::is_malloc_allowed are differing. "
            "This might be due to the fact that Eigen::internal::set_is_alloc_allowed was called independently/instead "
            "of tvm::utils::set_is_malloc_allowed in the current program/library.");
#  endif
}
 
inline void enforce_malloc_coherency()
{
#  ifdef EIGEN_RUNTIME_NO_MALLOC
  Eigen::internal::set_is_malloc_allowed(is_malloc_allowed_());
#  endif
}
} // namespace internal
 
inline bool is_malloc_allowed()
{
  internal::check_malloc_coherency();
  return internal::is_malloc_allowed_();
}
 
inline bool set_is_malloc_allowed(bool allow)
{
  internal::check_malloc_coherency();
  internal::set_is_malloc_allowed_(allow);
  internal::enforce_malloc_coherency();
  return allow;
}
 
inline void override_is_malloc_allowed(bool allow)
{
  internal::check_malloc_coherency();
  internal::override_is_malloc_allowed_(allow);
  internal::enforce_malloc_coherency();
}
 
inline void restore_is_malloc_allowed()
{
  internal::check_malloc_coherency();
  internal::restore_is_malloc_allowed_();
  internal::enforce_malloc_coherency();
}
#endif
 
#define TVM_TEMPORARY_ALLOW_EIGEN_MALLOC(x)     \
  tvm::utils::override_is_malloc_allowed(true); \
  x;                                            \
  tvm::utils::restore_is_malloc_allowed();
 
} // namespace tvm::utils