diff --git a/runtime/hsa-runtime/core/inc/amd_gpu_agent.h b/runtime/hsa-runtime/core/inc/amd_gpu_agent.h
index 1653823823..f514dda78c 100644
--- a/runtime/hsa-runtime/core/inc/amd_gpu_agent.h
+++ b/runtime/hsa-runtime/core/inc/amd_gpu_agent.h
@@ -151,7 +151,7 @@ class GpuAgentInt : public core::Agent {
    virtual uint64_t TranslateTime(uint64_t tick) = 0;
 
    // @brief Invalidate caches on the agent which may hold code object data.
-   virtual void InvalidateCodeCaches() = 0;
+   virtual void InvalidateCodeCaches(void *ptr, size_t size) = 0;
 
    // @brief Sets the coherency type of this agent.
    //
@@ -353,7 +353,7 @@ class GpuAgent : public GpuAgentInt {
   uint64_t TranslateTime(uint64_t tick) override;
 
   // @brief Override from AMD::GpuAgentInt.
-  void InvalidateCodeCaches() override;
+  void InvalidateCodeCaches(void* ptr, size_t size) override;
 
   // @brief Override from AMD::GpuAgentInt.
   bool current_coherency_type(hsa_amd_coherency_type_t type) override;
diff --git a/runtime/hsa-runtime/core/inc/amd_gpu_pm4.h b/runtime/hsa-runtime/core/inc/amd_gpu_pm4.h
index 0be269681c..2e871067cb 100644
--- a/runtime/hsa-runtime/core/inc/amd_gpu_pm4.h
+++ b/runtime/hsa-runtime/core/inc/amd_gpu_pm4.h
@@ -81,6 +81,8 @@
 #  define PM4_ACQUIRE_MEM_COHER_CNTL_SH_ICACHE_ACTION_ENA  (1 << 29)
 #define PM4_ACQUIRE_MEM_DW2_COHER_SIZE(x)                  (((x) & 0xFFFFFFFF) << 0)
 #define PM4_ACQUIRE_MEM_DW3_COHER_SIZE_HI(x)               (((x) & 0xFF) << 0)
+#define PM4_ACQUIRE_MEM_DW4_COHER_BASE(x)                  ((x >> 8) & 0xFFFFFFFF)
+#define PM4_ACQUIRE_MEM_DW4_COHER_BASE_HI(x)               ((x >> 40) & 0xFFFFFF)
 #define PM4_ACQUIRE_MEM_DW7_GCR_CNTL(x)                    (((x) & 0x7FFFF) << 0)
 #  define PM4_ACQUIRE_MEM_GCR_CNTL_GLI_INV(x)              (((x) & 0x3) << 0)
 #  define PM4_ACQUIRE_MEM_GCR_CNTL_GLK_INV                 (1 << 7)
diff --git a/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp b/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp
index f7e0fcd51b..e772f22b60 100644
--- a/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp
+++ b/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp
@@ -2310,7 +2310,7 @@ void GpuAgent::BindTrapHandler() {
   assert(err == HSAKMT_STATUS_SUCCESS && "hsaKmtSetTrapHandler() failed");
 }
 
-void GpuAgent::InvalidateCodeCaches() {
+void GpuAgent::InvalidateCodeCaches(void *ptr, size_t size) {
   // Check for microcode cache invalidation support.
   // This is deprecated in later microcode builds.
   if (isa_->GetMajorVersion() == 7) {
@@ -2352,8 +2352,17 @@ void GpuAgent::InvalidateCodeCaches() {
 
   cache_inv[0] = PM4_HDR(PM4_HDR_IT_OPCODE_ACQUIRE_MEM, cache_inv_size_dw,
              isa_->GetMajorVersion());
-  cache_inv[2] = PM4_ACQUIRE_MEM_DW2_COHER_SIZE(0xFFFFFFFF);
-  cache_inv[3] = PM4_ACQUIRE_MEM_DW3_COHER_SIZE_HI(0xFF);
+
+  if (ptr) {
+    size_t size_granule = (size + 0xFF) >> 8;
+    cache_inv[2] = PM4_ACQUIRE_MEM_DW2_COHER_SIZE(size_granule);
+    cache_inv[3] = PM4_ACQUIRE_MEM_DW3_COHER_SIZE_HI(size_granule >> 32);
+    cache_inv[4] = PM4_ACQUIRE_MEM_DW4_COHER_BASE((uint64_t)ptr);
+    cache_inv[5] = PM4_ACQUIRE_MEM_DW4_COHER_BASE_HI((uint64_t)ptr);
+  } else {
+    cache_inv[2] = PM4_ACQUIRE_MEM_DW2_COHER_SIZE(0xFFFFFFFF);
+    cache_inv[3] = PM4_ACQUIRE_MEM_DW3_COHER_SIZE_HI(0xFF);
+  }
 
   // Submit the command to the utility queue and wait for it to complete.
   queues_[QueueUtility]->ExecutePM4(cache_inv, cache_inv_size_dw * sizeof(uint32_t));
diff --git a/runtime/hsa-runtime/core/runtime/amd_loader_context.cpp b/runtime/hsa-runtime/core/runtime/amd_loader_context.cpp
index 0b027d77ff..94f21eb9af 100644
--- a/runtime/hsa-runtime/core/runtime/amd_loader_context.cpp
+++ b/runtime/hsa-runtime/core/runtime/amd_loader_context.cpp
@@ -360,7 +360,7 @@ bool RegionMemory::Freeze() {
 
   // Invalidate agent caches which may hold lines of the new allocation.
   if (is_code_ && (region_->owner()->device_type() == core::Agent::kAmdGpuDevice))
-    ((AMD::GpuAgent*)region_->owner())->InvalidateCodeCaches();
+    ((AMD::GpuAgent*)region_->owner())->InvalidateCodeCaches(ptr_, size_);
 
   return true;
 }