bugfix torch.xpu.empty_cache (#3550) (#3552)

Co-authored-by: Jinghui <[email protected]>
(cherry picked from commit 51465f8)

Author: guangyey

csrc/gpu/runtime/CachingDeviceAllocator.cpp

@@ -380,6 +380,20 @@ void CachingDeviceAllocator::recordQueue(void* buffer, sycl::queue* queue) {
 void CachingDeviceAllocator::emptyCache() {
   std::lock_guard<std::recursive_mutex> lock(mutex);
   synchronize_and_free_events(std::nullopt);
+
+  /*
+   * See Note [Safe to Free Blocks on BlockPool]
+   *
+   * torch.xpu.empty_cache will release all unoccupied cached memory currently
+   * held on all the GPUs. So we have to do a device-level synchronization on
+   * all GPUs.
+   */
+  int count = 0;
+  AT_DPCPP_CHECK(dpcppGetDeviceCount(&count));
+  for (auto i = 0; i < count; i++) {
+    xpu::dpcpp::deviceSynchronize(i);
+  }
+
   free_blocks(large_blocks, large_blocks.begin(), large_blocks.end());
   free_blocks(small_blocks, small_blocks.begin(), small_blocks.end());
 }
@@ -545,6 +559,15 @@ size_t CachingDeviceAllocator::try_merge_blocks(
   return subsumed_size;
 }
 
+/**
+ * Note [Safe to Free Blocks on BlockPool]
+ *
+ * Callers must ensure that all accesses to the block, whose raw pointer is
+ * allocated by SYCL APIs, have been completed before invoking sycl::free.
+ *
+ * We have to do a device-level synchronization before free these blocks to
+ * guarantee that all kernels can access to the blocks have finished.
+ */
 void CachingDeviceAllocator::free_blocks(
     BlockPool& blocks,
     BlockPool::iterator it,
@@ -579,6 +602,11 @@ void CachingDeviceAllocator::free_cached_blocks(DeviceId di) {
   Block lower_bound(di, nullptr, 0);
   Block upper_bound(di + 1, nullptr, 0);
 
+  /*
+   * See Note [Safe to Free Blocks on BlockPool]
+   */
+  xpu::dpcpp::deviceSynchronize(di);
+
   free_blocks(
       large_blocks,
       large_blocks.lower_bound(&lower_bound),

intel_extension_for_pytorch/transformers/models/xpu/optimize_transformers/modules/Functions.py

@@ -655,7 +655,6 @@ def _ipex_beam_search(
     # IPEXTransformerAtten.release_all_static_cached_resources()
     reserved_mem = round(torch.xpu.memory_reserved() / 1024**3, 3)
     if reserved_mem > 50:
-        torch.xpu.synchronize()
         torch.xpu.empty_cache()
     if hasattr(self, "token_latency") and self.token_latency:
         return out, latency_list

intel_extension_for_pytorch/transformers/models/xpu/optimize_transformers/modules/transformer_modules/BaseAttention.py

@@ -114,7 +114,6 @@ def end_of_attention(self):
                 not self.is_beam_search()
                 and IPEXTransformerAttn.timestamp % self.runtime_cache_size == 0
             ):
-                torch.xpu.synchronize()
                 torch.xpu.empty_cache()
             IPEXTransformerAttn.timestamp += 1
 

intel_extension_for_pytorch/transformers/models/xpu/optimize_transformers/modules/transformer_modules/Decoderblock.py

@@ -39,7 +39,6 @@ def port_all_parameters_to_new_module(self):
         self.port_mlp_parameter()
         self.port_norm_parameter()
         self.port_module_specific_parameter()
-        torch.xpu.synchronize()
         torch.xpu.empty_cache()
         # for debug
         # self.print_all_paramter_with_name

tests/gpu/regression/test_fill.py

@@ -9,7 +9,6 @@ def test_fill(self):
         Regression desc:
           fill_ may set values to part of large-size tensor.
         """
-        torch.xpu.synchronize()
         torch.xpu.empty_cache()
 
         output_cpu = torch.zeros([2, 8, 256, 512, 224])