[HIP] Implement workaround for hipMemset2D

source/adapters/hip/enqueue.cpp

@@ -769,29 +769,19 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueMemBufferCopyRect(
   return Result;
 }
 
-// HIP has no memset functions that allow setting values more than 4 bytes. UR
-// API lets you pass an arbitrary "pattern" to the buffer fill, which can be
-// more than 4 bytes. We must break up the pattern into 1 byte values, and set
-// the buffer using multiple strided calls.  The first 4 patterns are set using
-// hipMemsetD32Async then all subsequent 1 byte patterns are set using
-// hipMemset2DAsync which is called for each pattern.
-ur_result_t commonMemSetLargePattern(hipStream_t Stream, uint32_t PatternSize,
-                                     size_t Size, const void *pPattern,
-                                     hipDeviceptr_t Ptr) {
-  // Calculate the number of patterns, stride, number of times the pattern
-  // needs to be applied, and the number of times the first 32 bit pattern
-  // needs to be applied.
+static inline void memsetRemainPattern(hipStream_t Stream, uint32_t PatternSize,
+                                       size_t Size, const void *pPattern,
+                                       hipDeviceptr_t Ptr) {
+
+  // Calculate the number of patterns, stride and the number of times the
+  // pattern needs to be applied.
   auto NumberOfSteps = PatternSize / sizeof(uint8_t);
   auto Pitch = NumberOfSteps * sizeof(uint8_t);
   auto Height = Size / NumberOfSteps;
-  auto Count32 = Size / sizeof(uint32_t);
 
-  // Get 4-byte chunk of the pattern and call hipMemsetD32Async
-  auto Value = *(static_cast<const uint32_t *>(pPattern));
-  UR_CHECK_ERROR(hipMemsetD32Async(Ptr, Value, Count32, Stream));
   for (auto step = 4u; step < NumberOfSteps; ++step) {
     // take 1 byte of the pattern
-    Value = *(static_cast<const uint8_t *>(pPattern) + step);
+    auto Value = *(static_cast<const uint8_t *>(pPattern) + step);
 
     // offset the pointer to the part of the buffer we want to write to
     auto OffsetPtr = reinterpret_cast<void *>(reinterpret_cast<uint8_t *>(Ptr) +
@@ -801,6 +791,55 @@ ur_result_t commonMemSetLargePattern(hipStream_t Stream, uint32_t PatternSize,
     UR_CHECK_ERROR(hipMemset2DAsync(OffsetPtr, Pitch, Value, sizeof(uint8_t),
                                     Height, Stream));
   }
+}
+
+// HIP has no memset functions that allow setting values more than 4 bytes. UR
+// API lets you pass an arbitrary "pattern" to the buffer fill, which can be
+// more than 4 bytes. We must break up the pattern into 1 byte values, and set
+// the buffer using multiple strided calls.  The first 4 patterns are set using
+// hipMemsetD32Async then all subsequent 1 byte patterns are set using
+// hipMemset2DAsync which is called for each pattern.
+ur_result_t commonMemSetLargePattern(hipStream_t Stream, uint32_t PatternSize,
+                                     size_t Size, const void *pPattern,
+                                     hipDeviceptr_t Ptr) {
+
+  // Get 4-byte chunk of the pattern and call hipMemsetD32Async
+  auto Count32 = Size / sizeof(uint32_t);
+  auto Value = *(static_cast<const uint32_t *>(pPattern));
+  UR_CHECK_ERROR(hipMemsetD32Async(Ptr, Value, Count32, Stream));
+
+  // There is a bug in ROCm prior to 6.0.0 version which causes hipMemset2D
+  // to behave incorrectly when acting on host pinned memory.
+  // In such a case, the memset operation is partially emulated with memcpy.
+#if HIP_VERSION_MAJOR < 6
+  hipPointerAttribute_t ptrAttribs{};
+  UR_CHECK_ERROR(hipPointerGetAttributes(&ptrAttribs, (const void *)Ptr));
+
+  // The hostPointer attribute is non-null also for shared memory allocations.
+  // To make sure that this workaround only executes for host pinned memory, we
+  // need to check that isManaged attribute is false.
+  if (ptrAttribs.hostPointer && !ptrAttribs.isManaged) {
+    const auto NumOfCopySteps = Size / PatternSize;
+    const auto Offset = sizeof(uint32_t);
+    const auto LeftPatternSize = PatternSize - Offset;
+    const auto OffsetPatternPtr = reinterpret_cast<const void *>(
+        reinterpret_cast<const uint8_t *>(pPattern) + Offset);
+
+    // Loop through the memory area to memset, advancing each time by the
+    // PatternSize and memcpy the left over pattern bits.
+    for (uint32_t i = 0; i < NumOfCopySteps; ++i) {
+      auto OffsetDstPtr = reinterpret_cast<void *>(
+          reinterpret_cast<uint8_t *>(Ptr) + Offset + i * PatternSize);
+      UR_CHECK_ERROR(hipMemcpyAsync(OffsetDstPtr, OffsetPatternPtr,
+                                    LeftPatternSize, hipMemcpyHostToHost,
+                                    Stream));
+    }
+  } else {
+    memsetRemainPattern(Stream, PatternSize, Size, pPattern, Ptr);
+  }
+#else
+  memsetRemainPattern(Stream, PatternSize, Size, pPattern, Ptr);
+#endif
   return UR_RESULT_SUCCESS;
 }