[llvm] f14cb49 - [ORC] Add "wrap" and "unwrap" steps to ExecutorAddr toPtr/fromPtr.

[Lang Hames via llvm-commits]

Author: Lang Hames
Date: 2022-08-26T12:32:44-07:00
New Revision: f14cb494a34db6df9853d713c6027a476f030dbf

URL: https://github.com/llvm/llvm-project/commit/f14cb494a34db6df9853d713c6027a476f030dbf
DIFF: https://github.com/llvm/llvm-project/commit/f14cb494a34db6df9853d713c6027a476f030dbf.diff

LOG: [ORC] Add "wrap" and "unwrap" steps to ExecutorAddr toPtr/fromPtr.

The wrap/unwrap operations are applied to pointers after/before conversion to/from
raw addresses. They can be used to tag, untag, sign, or strip signing from
pointers. They currently default to 'rawPtr' (identity) on all platforms, but it
is expected that the default will be set based on the host architecture, e.g.
they would default to signing/stripping for arm64e.

Added: 
    

Modified: 
    llvm/include/llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h
    llvm/unittests/ExecutionEngine/Orc/ExecutorAddressTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h b/llvm/include/llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h
index 5d545f8abdb99..1264531e8f0df 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h
@@ -14,6 +14,7 @@
 #define LLVM_EXECUTIONENGINE_ORC_SHARED_EXECUTORADDRESS_H
 
 #include "llvm/ADT/DenseMapInfo.h"
+#include "llvm/ADT/identity.h"
 #include "llvm/ExecutionEngine/Orc/Shared/SimplePackedSerialization.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/raw_ostream.h"
@@ -29,6 +30,45 @@ using ExecutorAddrDiff = uint64_t;
 /// Represents an address in the executor process.
 class ExecutorAddr {
 public:
+  /// A wrap/unwrap function that leaves pointers unmodified.
+  template <typename T> using rawPtr = llvm::identity<T>;
+
+  /// Default wrap function to use on this host.
+  template <typename T> using defaultWrap = rawPtr<T>;
+
+  /// Default unwrap function to use on this host.
+  template <typename T> using defaultUnwrap = rawPtr<T>;
+
+  /// Merges a tag into the raw address value:
+  ///   P' = P | (TagValue << TagOffset).
+  class Tag {
+  public:
+    constexpr Tag(uintptr_t TagValue, uintptr_t TagOffset)
+        : TagMask(TagValue << TagOffset) {}
+
+    template <typename T> constexpr T *operator()(T *P) {
+      return reinterpret_cast<T*>(reinterpret_cast<uintptr_t>(P) | TagMask);
+    }
+
+  private:
+    uintptr_t TagMask;
+  };
+
+  /// Strips a tag of the given length from the given offset within the pointer:
+  /// P' = P & ~(((1 << TagLen) -1) << TagOffset)
+  class Untag {
+  public:
+    constexpr Untag(uintptr_t TagLen, uintptr_t TagOffset)
+        : UntagMask(~(((1 << TagLen) - 1) << TagOffset)) {}
+
+    template <typename T> constexpr T *operator()(T *P) {
+      return reinterpret_cast<T *>(reinterpret_cast<uintptr_t>(P) & UntagMask);
+    }
+
+  private:
+    uintptr_t UntagMask;
+  };
+
   ExecutorAddr() = default;
 
   /// Create an ExecutorAddr from the given value.
@@ -36,27 +76,30 @@ class ExecutorAddr {
 
   /// Create an ExecutorAddr from the given pointer.
   /// Warning: This should only be used when JITing in-process.
-  template <typename T> static ExecutorAddr fromPtr(T *Value) {
+  template <typename T, typename UnwrapFn = defaultUnwrap<T *>>
+  static ExecutorAddr fromPtr(T *Ptr, UnwrapFn &&Unwrap = UnwrapFn()) {
     return ExecutorAddr(
-        static_cast<uint64_t>(reinterpret_cast<uintptr_t>(Value)));
+        static_cast<uint64_t>(reinterpret_cast<uintptr_t>(Unwrap(Ptr))));
   }
 
   /// Cast this ExecutorAddr to a pointer of the given type.
   /// Warning: This should only be used when JITing in-process.
-  template <typename T>
-  std::enable_if_t<std::is_pointer<T>::value, T> toPtr() const {
+  template <typename T, typename WrapFn = defaultWrap<T>>
+  std::enable_if_t<std::is_pointer<T>::value, T>
+  toPtr(WrapFn &&Wrap = WrapFn()) const {
     uintptr_t IntPtr = static_cast<uintptr_t>(Addr);
     assert(IntPtr == Addr && "ExecutorAddr value out of range for uintptr_t");
-    return reinterpret_cast<T>(IntPtr);
+    return Wrap(reinterpret_cast<T>(IntPtr));
   }
 
   /// Cast this ExecutorAddr to a pointer of the given function type.
   /// Warning: This should only be used when JITing in-process.
-  template <typename T>
-  std::enable_if_t<std::is_function<T>::value, T *> toPtr() const {
+  template <typename T, typename WrapFn = defaultWrap<T *>>
+  std::enable_if_t<std::is_function<T>::value, T *>
+  toPtr(WrapFn &&Wrap = WrapFn()) const {
     uintptr_t IntPtr = static_cast<uintptr_t>(Addr);
     assert(IntPtr == Addr && "ExecutorAddr value out of range for uintptr_t");
-    return reinterpret_cast<T *>(IntPtr);
+    return Wrap(reinterpret_cast<T *>(IntPtr));
   }
 
   uint64_t getValue() const { return Addr; }

diff  --git a/llvm/unittests/ExecutionEngine/Orc/ExecutorAddressTest.cpp b/llvm/unittests/ExecutionEngine/Orc/ExecutorAddressTest.cpp
index 7cbfdffd49bcf..f829de6d6cb3c 100644
--- a/llvm/unittests/ExecutionEngine/Orc/ExecutorAddressTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/ExecutorAddressTest.cpp
@@ -58,6 +58,26 @@ TEST(ExecutorAddrTest, PtrConversionWithFunctionType) {
   EXPECT_EQ(FPtr, &F);
 }
 
+TEST(ExecutorAddrTest, WrappingAndUnwrapping) {
+  constexpr uintptr_t RawAddr = 0x123456;
+  int *RawPtr = (int *)RawAddr;
+
+  constexpr uintptr_t TagOffset = 8 * (sizeof(uintptr_t) - 1);
+  uintptr_t TagVal = 0xA5;
+  uintptr_t TagBits = TagVal << TagOffset;
+  void *TaggedPtr = (void *)((uintptr_t)RawPtr | TagBits);
+
+  ExecutorAddr EA =
+      ExecutorAddr::fromPtr(TaggedPtr, ExecutorAddr::Untag(8, TagOffset));
+
+  EXPECT_EQ(EA.getValue(), RawAddr);
+
+  void *ReconstitutedTaggedPtr =
+      EA.toPtr<void *>(ExecutorAddr::Tag(TagVal, TagOffset));
+
+  EXPECT_EQ(TaggedPtr, ReconstitutedTaggedPtr);
+}
+
 TEST(ExecutorAddrTest, AddrRanges) {
   ExecutorAddr A0(0), A1(1), A2(2), A3(3);
   ExecutorAddrRange R0(A0, A1), R1(A1, A2), R2(A2, A3), R3(A0, A2), R4(A1, A3);