2 files changed, 10 insertions, 4 deletions
diff --git a/drivers/accel/habanalabs/common/firmware_if.c b/drivers/accel/habanalabs/common/firmware_if.c
index eb000e035026..ef228087ef55 100644
--- a/drivers/accel/habanalabs/common/firmware_if.c
+++ b/drivers/accel/habanalabs/common/firmware_if.c
@@ -2647,7 +2647,7 @@ static int hl_fw_dynamic_init_cpu(struct hl_device *hdev,
 				fw_loader->dynamic_loader.comm_desc.cur_fw_ver);
 
 		if (rc)
-			goto out;
+			return rc;
 
 		/* read binning info from preboot */
 		if (hdev->support_preboot_binning) {
@@ -2660,15 +2660,19 @@ static int hl_fw_dynamic_init_cpu(struct hl_device *hdev,
 
 			rc = hdev->asic_funcs->set_dram_properties(hdev);
 			if (rc)
-				goto out;
+				return rc;
+
+			rc = hdev->asic_funcs->set_binning_masks(hdev);
+			if (rc)
+				return rc;
 
 			dev_dbg(hdev->dev,
 				"Read binning masks: tpc: 0x%llx, dram: 0x%llx, edma: 0x%x, dec: 0x%x, rot:0x%x\n",
 				hdev->tpc_binning, hdev->dram_binning, hdev->edma_binning,
 				hdev->decoder_binning, hdev->rotator_binning);
 		}
-out:
-		return rc;
+
+		return 0;
 	}
 
 	/* load boot fit to FW */
diff --git a/drivers/accel/habanalabs/common/habanalabs.h b/drivers/accel/habanalabs/common/habanalabs.h
index cd474422163d..0b7fe4afd92d 100644
--- a/drivers/accel/habanalabs/common/habanalabs.h
+++ b/drivers/accel/habanalabs/common/habanalabs.h
@@ -1549,6 +1549,7 @@ struct engines_data {
  * @set_engine_cores: set a config command to engine cores
  * @send_device_activity: indication to FW about device availability
  * @set_dram_properties: set DRAM related properties.
+ * @set_binning_masks: set binning/enable masks for all relevant components.
  */
 struct hl_asic_funcs {
 	int (*early_init)(struct hl_device *hdev);
@@ -1687,6 +1688,7 @@ struct hl_asic_funcs {
 					u32 num_cores, u32 core_command);
 	int (*send_device_activity)(struct hl_device *hdev, bool open);
 	int (*set_dram_properties)(struct hl_device *hdev);
+	int (*set_binning_masks)(struct hl_device *hdev);
 };