*                           dma_free_coherent(). This is ASIC function because
  *                           its implementation is not trivial when the driver
  *                           is loaded in simulation mode (not upstreamed).
- * @scrub_device_mem: Scrub device memory given an address and size
+ * @scrub_device_mem: Scrub the entire SRAM and DRAM.
  * @scrub_device_dram: Scrub the dram memory of the device.
  * @get_int_queue_base: get the internal queue base address.
  * @test_queues: run simple test on all queues for sanity check.
                                        dma_addr_t *dma_handle, gfp_t flag);
        void (*asic_dma_free_coherent)(struct hl_device *hdev, size_t size,
                                        void *cpu_addr, dma_addr_t dma_handle);
-       int (*scrub_device_mem)(struct hl_device *hdev, u64 addr, u64 size);
+       int (*scrub_device_mem)(struct hl_device *hdev);
        int (*scrub_device_dram)(struct hl_device *hdev, u64 val);
        void* (*get_int_queue_base)(struct hl_device *hdev, u32 queue_id,
                                dma_addr_t *dma_handle, u16 *queue_len);
 
        }
 
        /* Scrub both SRAM and DRAM */
-       rc = hdev->asic_funcs->scrub_device_mem(hdev, 0, 0);
+       rc = hdev->asic_funcs->scrub_device_mem(hdev);
        if (rc)
                goto disable_pci_access;
 
        return 0;
 }
 
-static int gaudi_scrub_device_mem(struct hl_device *hdev, u64 addr, u64 size)
+static int gaudi_scrub_device_mem(struct hl_device *hdev)
 {
        struct asic_fixed_properties *prop = &hdev->asic_prop;
+       u64 addr, size, dummy_val;
        int rc = 0;
        u64 val = 0;
 
        if (!hdev->memory_scrub)
                return 0;
 
-       if (!addr && !size) {
-               /* Wait till device is idle */
-               rc = hl_poll_timeout(
-                               hdev,
-                               mmDMA0_CORE_STS0/* dummy */,
-                               val/* dummy */,
-                               (hdev->asic_funcs->is_device_idle(hdev, NULL,
-                                               0, NULL)),
-                                               1000,
-                                               HBM_SCRUBBING_TIMEOUT_US);
-               if (rc) {
-                       dev_err(hdev->dev, "waiting for idle timeout\n");
-                       return -EIO;
-               }
+       /* Wait till device is idle */
+       rc = hl_poll_timeout(hdev,
+                       mmDMA0_CORE_STS0 /* dummy */,
+                       dummy_val /* dummy */,
+                       (hdev->asic_funcs->is_device_idle(hdev, NULL, 0, NULL)),
+                       1000,
+                       HBM_SCRUBBING_TIMEOUT_US);
+       if (rc) {
+               dev_err(hdev->dev, "waiting for idle timeout\n");
+               return -EIO;
+       }
 
-               /* Scrub SRAM */
-               addr = prop->sram_user_base_address;
-               size = hdev->pldm ? 0x10000 :
-                               (prop->sram_size - SRAM_USER_BASE_OFFSET);
-               val = 0x7777777777777777ull;
+       /* Scrub SRAM */
+       addr = prop->sram_user_base_address;
+       size = hdev->pldm ? 0x10000 : prop->sram_size - SRAM_USER_BASE_OFFSET;
+       val = 0x7777777777777777ull;
 
-               rc = gaudi_memset_device_memory(hdev, addr, size, val);
-               if (rc) {
-                       dev_err(hdev->dev,
-                               "Failed to clear SRAM in mem scrub all\n");
-                       return rc;
-               }
+       dev_dbg(hdev->dev, "Scrubing SRAM: 0x%09llx - 0x%09llx val: 0x%llx\n",
+                       addr, addr + size, val);
+       rc = gaudi_memset_device_memory(hdev, addr, size, val);
+       if (rc) {
+               dev_err(hdev->dev, "Failed to clear SRAM (%d)\n", rc);
+               return rc;
+       }
 
-               /* Scrub HBM using all DMA channels in parallel */
-               rc = gaudi_scrub_device_dram(hdev, 0xdeadbeaf);
-               if (rc)
-                       dev_err(hdev->dev,
-                               "Failed to clear HBM in mem scrub all\n");
+       /* Scrub HBM using all DMA channels in parallel */
+       rc = gaudi_scrub_device_dram(hdev, 0xdeadbeaf);
+       if (rc) {
+               dev_err(hdev->dev, "Failed to clear HBM (%d)\n", rc);
+               return rc;
        }
 
-       return rc;
+       return 0;
 }
 
 static void *gaudi_get_int_queue_base(struct hl_device *hdev,
 
                                u32 len, u32 original_len, u64 cq_addr, u32 cq_val,
                                u32 msix_vec, bool eb);
 int goya_cs_parser(struct hl_device *hdev, struct hl_cs_parser *parser);
-int goya_scrub_device_mem(struct hl_device *hdev, u64 addr, u64 size);
+int goya_scrub_device_mem(struct hl_device *hdev);
 void *goya_get_int_queue_base(struct hl_device *hdev, u32 queue_id,
                                dma_addr_t *dma_handle, u16 *queue_len);
 u32 goya_get_dma_desc_list_size(struct hl_device *hdev, struct sg_table *sgt);