habanalabs: set dma mask from fw once fw done iatu config

When setting "DMA mask from FW" we are reading PSOC_GLOBAL_CONF register
which is allowed only once FW has done it's iATU configuration.
Signed-off-by: Ohad Sharabi <osharabi@habana.ai>
Reviewed-by: Oded Gabbay <ogabbay@kernel.org>
Signed-off-by: Oded Gabbay <ogabbay@kernel.org>

drivers/misc/habanalabs/common/pci/pci.c

@@ -421,6 +421,12 @@ int hl_pci_init(struct hl_device *hdev)
 		goto unmap_pci_bars;
 	}
 
+	/* Driver must sleep in order for FW to finish the iATU configuration */
+	if (hdev->asic_prop.iatu_done_by_fw) {
+		usleep_range(2000, 3000);
+		hdev->asic_funcs->set_dma_mask_from_fw(hdev);
+	}
+
 	rc = dma_set_mask_and_coherent(&pdev->dev,
 					DMA_BIT_MASK(hdev->dma_mask));
 	if (rc) {
@@ -430,10 +436,6 @@ int hl_pci_init(struct hl_device *hdev)
 		goto unmap_pci_bars;
 	}
 
-	/* Driver must sleep in order for FW to finish the iATU configuration */
-	if (hdev->asic_prop.iatu_done_by_fw)
-		usleep_range(2000, 3000);
-
 	return 0;
 
 unmap_pci_bars:

drivers/misc/habanalabs/gaudi/gaudi.c

@@ -600,10 +600,8 @@ static int gaudi_init_iatu(struct hl_device *hdev)
 	struct hl_outbound_pci_region outbound_region;
 	int rc;
 
-	if (hdev->asic_prop.iatu_done_by_fw) {
-		hdev->asic_funcs->set_dma_mask_from_fw(hdev);
+	if (hdev->asic_prop.iatu_done_by_fw)
 		return 0;
-	}
 
 	/* Inbound Region 0 - Bar 0 - Point to SRAM + CFG */
 	inbound_region.mode = PCI_BAR_MATCH_MODE;

drivers/misc/habanalabs/goya/goya.c

@@ -532,10 +532,8 @@ static int goya_init_iatu(struct hl_device *hdev)
 	struct hl_outbound_pci_region outbound_region;
 	int rc;
 
-	if (hdev->asic_prop.iatu_done_by_fw) {
-		hdev->asic_funcs->set_dma_mask_from_fw(hdev);
+	if (hdev->asic_prop.iatu_done_by_fw)
 		return 0;
-	}
 
 	/* Inbound Region 0 - Bar 0 - Point to SRAM and CFG */
 	inbound_region.mode = PCI_BAR_MATCH_MODE;