Properly annotate kernels

src/backend/OpenCL/CHIPBackendOpenCL.cc

@@ -246,6 +246,8 @@ annotateIndirectPointers(const CHIPContextOpenCL &Ctx,
     return nullptr;
 
   cl_kernel_exec_info PtrListName;
+  cl_kernel_exec_info AllocType;
+
   switch (Ctx.getAllocStrategy()) {
   default:
     assert(!"Unexpected allocation strategy!");
@@ -273,6 +275,18 @@ annotateIndirectPointers(const CHIPContextOpenCL &Ctx,
     for (std::shared_ptr<void> Ptr : Ctx.getAllocPointers()) {
       AnnotationList.push_back(Ptr.get());
       AllocKeepAlives->push_back(Ptr);
+
+      switch (Ctx.MemManager_.getPointerMemoryType(Ptr.get())) {
+      case hipMemoryTypeHost:
+        AllocType = CL_KERNEL_EXEC_INFO_INDIRECT_HOST_ACCESS_INTEL;
+        break;
+      case hipMemoryTypeDevice:
+        AllocType = CL_KERNEL_EXEC_INFO_INDIRECT_DEVICE_ACCESS_INTEL;
+        break;
+      case hipMemoryTypeUnified:
+        AllocType = CL_KERNEL_EXEC_INFO_INDIRECT_SHARED_ACCESS_INTEL;
+        break;
+      }
     }
 
     // TODO: Optimization. Don't call this function again if we know the
@@ -281,6 +295,13 @@ annotateIndirectPointers(const CHIPContextOpenCL &Ctx,
                                    AnnotationList.size() * sizeof(void *),
                                    AnnotationList.data());
     CHIPERR_CHECK_LOG_AND_THROW_TABLE(clSetKernelExecInfo);
+
+    if (Ctx.getAllocStrategy() == AllocationStrategy::IntelUSM) {
+      cl_bool param = CL_TRUE;
+      clStatus = clSetKernelExecInfo(KernelAPIHandle, AllocType,
+                                     sizeof(cl_bool), &param);
+      CHIPERR_CHECK_LOG_AND_THROW_TABLE(clSetKernelExecInfo);
+    }
   }
 
   return AllocKeepAlives;
@@ -1259,13 +1280,8 @@ CHIPQueueOpenCL::launchImpl(chipstar::ExecItem *ExecItem) {
   LOCK(Backend->DubiousLockOpenCL);
 #endif
 
-  std::unique_ptr<std::vector<std::shared_ptr<void>>> AllocationsToKeepAlive;
-  // Disable annotation for CPU device. Runtime errror.
-  // https://community.intel.com/t5/OpenCL-for-CPU/OpenCL-Using-USM-on-Intel-CPU-Runtime-along-with/m-p/1621136#M7350
-  if (ChipEnvVars.getDevice().getType() != DeviceType::CPU) {
-    AllocationsToKeepAlive = annotateIndirectPointers(
-        *OclContext, Kernel->getModule()->getInfo(), KernelHandle);
-  }
+  auto AllocationsToKeepAlive = annotateIndirectPointers(
+      *OclContext, Kernel->getModule()->getInfo(), KernelHandle);
 
   auto [EventsToWait, EventLocks] = getSyncQueuesLastEvents(LaunchEvent, false);
   std::vector<cl_event> SyncQueuesEventHandles = getOpenCLHandles(EventsToWait);

src/backend/OpenCL/CHIPBackendOpenCL.hh

@@ -155,8 +155,8 @@ struct CHIPContextUSMExts {
   clMemFreeINTEL_fn clMemFreeINTEL;
 };
 
-using const_alloc_iterator = ConstMapKeyIterator<
-    std::map<std::shared_ptr<void>, size_t, PointerCmp<void>>>;
+using const_alloc_iterator = ConstMapKeyIterator<std::map<
+    std::shared_ptr<void>, std::pair<size_t, hipMemoryType>, PointerCmp<void>>>;
 
 /// Used to select a strategy for managing HIP allocations.
 enum class AllocationStrategy {
@@ -173,7 +173,9 @@ enum class AllocationStrategy {
 class MemoryManager {
 private:
   // ContextMutex should be enough
-  std::map<std::shared_ptr<void>, size_t, PointerCmp<void>> Allocations_;
+  std::map<std::shared_ptr<void>, std::pair<size_t, hipMemoryType>,
+           PointerCmp<void>>
+      Allocations_;
   cl::Context Context_;
   cl::Device Device_;
 
@@ -195,7 +197,8 @@ public:
   bool free(void *P);
   bool hasPointer(const void *Ptr);
   bool pointerSize(void *Ptr, size_t *Size);
-  bool pointerInfo(void *Ptr, void **Base, size_t *Size);
+  bool pointerInfo(void *Ptr, void **Base, size_t *Size,
+                   hipMemoryType *MemType);
   int memCopy(void *Dst, const void *Src, size_t Size, cl::CommandQueue &Queue);
   int memFill(void *Dst, size_t Size, const void *Pattern, size_t PatternSize,
               cl::CommandQueue &Queue);
@@ -212,6 +215,8 @@ public:
     return AllocStrategy_;
   }
 
+  hipMemoryType getPointerMemoryType(const void *Ptr) const;
+
   std::pair<cl_mem, size_t> translateDevPtrToBuffer(const void *DevPtr) const;
 
 private:

src/backend/OpenCL/MemoryManager.cc

@@ -135,6 +135,10 @@ std::shared_ptr<void> MemoryManager::allocateSVM(size_t Size, size_t Alignment,
   return Result;
 }
 
+hipMemoryType MemoryManager::getPointerMemoryType(const void *Ptr) const {
+  return hipMemoryTypeHost;
+}
+
 std::shared_ptr<void> MemoryManager::allocateUSM(size_t Size, size_t Alignment,
                                                  hipMemoryType MemType) {
   std::shared_ptr<void> Result;
@@ -240,12 +244,12 @@ void *MemoryManager::allocate(size_t Size, size_t Alignment,
 
   logTrace("Memory allocated: {} / {}\n", Ptr.get(), Size);
   assert(Allocations_.find(Ptr) == Allocations_.end());
-  Allocations_.emplace(Ptr, Size);
+  Allocations_.emplace(Ptr, std::make_pair(Size, MemType));
   return Ptr.get();
 }
 
 bool MemoryManager::free(void *Ptr) {
-  auto I = Allocations_.find(Ptr);
+  auto I = Allocations_.find(std::shared_ptr<void>(Ptr, [](void *) {}));
   if (I != Allocations_.end())
     Allocations_.erase(I);
 
@@ -257,33 +261,33 @@ bool MemoryManager::free(void *Ptr) {
 
 bool MemoryManager::hasPointer(const void *Ptr) {
   logTrace("hasPointer on: {}\n", Ptr);
-  return (Allocations_.find((void *)Ptr) != Allocations_.end());
+  return (Allocations_.find(std::shared_ptr<void>(
+              const_cast<void *>(Ptr), [](void *) {})) != Allocations_.end());
 }
 
 bool MemoryManager::pointerSize(void *Ptr, size_t *Size) {
   logTrace("pointerSize on: {}\n", Ptr);
-  auto I = Allocations_.find(Ptr);
+  auto I = Allocations_.find(std::shared_ptr<void>(Ptr, [](void *) {}));
   if (I != Allocations_.end()) {
-    *Size = I->second;
+    *Size = I->second.first;
     return true;
   } else {
     return false;
   }
 }
 
-bool MemoryManager::pointerInfo(void *Ptr, void **Base, size_t *Size) {
+bool MemoryManager::pointerInfo(void *Ptr, void **Base, size_t *Size,
+                                hipMemoryType *MemType) {
   logTrace("pointerInfo on: {}\n", Ptr);
-  for (auto I : Allocations_) {
-    if ((I.first.get() <= Ptr) &&
-        (Ptr < ((const char *)I.first.get() + I.second))) {
-      if (Base)
-        *Base = I.first.get();
-      if (Size)
-        *Size = I.second;
-      return true;
-    }
+  auto I = Allocations_.find(std::shared_ptr<void>(Ptr, [](void *) {}));
+  if (I != Allocations_.end()) {
+    *Base = I->first.get();
+    *Size = I->second.first;
+    *MemType = I->second.second;
+    return true;
+  } else {
+    return false;
   }
-  return false;
 }
 
 void MemoryManager::clear() {
@@ -302,22 +306,21 @@ std::pair<cl_mem, size_t>
 MemoryManager::translateDevPtrToBuffer(const void *DevPtr) const {
   assert(getAllocStrategy() == AllocationStrategy::BufferDevAddr);
 
-  // Find entry with largest key so that 'key <= DevPtr'.
-  auto UpperIt = Allocations_.upper_bound(DevPtr);
+  auto UpperIt = Allocations_.upper_bound(
+      std::shared_ptr<void>(const_cast<void *>(DevPtr), [](void *) {}));
   auto CandIt = UpperIt == Allocations_.begin() ? UpperIt : std::prev(UpperIt);
   if (CandIt == Allocations_.end() || DevPtr < CandIt->first.get())
     return {nullptr, 0};
 
   const char *BasePtr = static_cast<const char *>(CandIt->first.get());
   size_t Offset = 0;
 
-  // Handle offseted pointer, check it is within the bounds of the allocation.
   if (BasePtr != DevPtr) {
     const char *CandPtr = static_cast<const char *>(DevPtr);
-    if (CandPtr >= (BasePtr + CandIt->second))
+    if (CandPtr >= (BasePtr + CandIt->second.first))
       return {nullptr, 0};
     Offset = CandPtr - BasePtr;
   }
 
   return {DevPtrToBuffer_.at(BasePtr), Offset};
-}
+}