* struct iterate_module_ctx - HW module iterator
  * @fn: function to apply to each HW module instance
  * @data: optional internal data to the function iterator
+ * @rc: return code for optional use of iterator/iterator-caller
  */
 struct iterate_module_ctx {
        /*
         * @inst: HW module instance within the block
         * @offset: current HW module instance offset from the 1-st HW module instance
         *          in the 1-st block
-        * @data: function specific data
+        * @ctx: the iterator context.
         */
-       void (*fn)(struct hl_device *hdev, int block, int inst, u32 offset, void *data);
+       void (*fn)(struct hl_device *hdev, int block, int inst, u32 offset,
+                       struct iterate_module_ctx *ctx);
        void *data;
+       int rc;
 };
 
 struct hl_block_glbl_sec {
 
        int dcore, inst, tpc_seq;
        u32 offset;
 
+       /* init the return code */
+       ctx->rc = 0;
+
        for (dcore = 0; dcore < NUM_OF_DCORES; dcore++) {
                for (inst = 0; inst < NUM_OF_TPC_PER_DCORE; inst++) {
                        tpc_seq = dcore * NUM_OF_TPC_PER_DCORE + inst;
 
                        offset = (DCORE_OFFSET * dcore) + (DCORE_TPC_OFFSET * inst);
 
-                       ctx->fn(hdev, dcore, inst, offset, ctx->data);
+                       ctx->fn(hdev, dcore, inst, offset, ctx);
+                       if (ctx->rc) {
+                               dev_err(hdev->dev, "TPC iterator failed for DCORE%d TPC%d\n",
+                                                       dcore, inst);
+                               return;
+                       }
                }
        }
 
 
        /* special check for PCI TPC (DCORE0_TPC6) */
        offset = DCORE_TPC_OFFSET * (NUM_DCORE0_TPC - 1);
-       ctx->fn(hdev, 0, NUM_DCORE0_TPC - 1, offset, ctx->data);
+       ctx->fn(hdev, 0, NUM_DCORE0_TPC - 1, offset, ctx);
+       if (ctx->rc)
+               dev_err(hdev->dev, "TPC iterator failed for DCORE0 TPC6\n");
 }
 
 static bool gaudi2_host_phys_addr_valid(u64 addr)
 };
 
 static void gaudi2_init_tpc_config(struct hl_device *hdev, int dcore, int inst,
-                                                       u32 offset, void *data)
+                                       u32 offset, struct iterate_module_ctx *ctx)
 {
        struct gaudi2_device *gaudi2 = hdev->asic_specific;
-       struct gaudi2_tpc_init_cfg_data *cfg_data = data;
+       struct gaudi2_tpc_init_cfg_data *cfg_data = ctx->data;
        u32 queue_id_base;
        u8 seq;
 
 }
 
 static void gaudi2_is_tpc_engine_idle(struct hl_device *hdev, int dcore, int inst, u32 offset,
-                                       void *data)
+                                       struct iterate_module_ctx *ctx)
 {
-       struct gaudi2_tpc_idle_data *idle_data = (struct gaudi2_tpc_idle_data *)data;
+       struct gaudi2_tpc_idle_data *idle_data = ctx->data;
        u32 tpc_cfg_sts, qm_glbl_sts0, qm_glbl_sts1, qm_cgm_sts;
        bool is_eng_idle;
        int engine_idx;
 }
 
 static void gaudi2_tpc_mmu_prepare(struct hl_device *hdev, int dcore, int inst,        u32 offset,
-                                       void *data)
+                                       struct iterate_module_ctx *ctx)
 {
-       struct gaudi2_tpc_mmu_data *mmu_data = (struct gaudi2_tpc_mmu_data *)data;
+       struct gaudi2_tpc_mmu_data *mmu_data = ctx->data;
 
        WREG32(mmDCORE0_TPC0_CFG_AXUSER_HB_MMU_BP + offset, 0);
        WREG32(mmDCORE0_TPC0_CFG_AXUSER_HB_ASID + offset, mmu_data->rw_asid);
 
 };
 
 static void gaudi2_config_tpcs_glbl_sec(struct hl_device *hdev, int dcore, int inst, u32 offset,
-                                       void *data)
+                                               struct iterate_module_ctx *ctx)
 {
-       struct gaudi2_tpc_pb_data *pb_data = (struct gaudi2_tpc_pb_data *)data;
+       struct gaudi2_tpc_pb_data *pb_data = ctx->data;
 
        hl_config_glbl_sec(hdev, gaudi2_pb_dcr0_tpc0, pb_data->glbl_sec,
                                        offset, pb_data->block_array_size);
 struct gaudi2_tpc_arc_pb_data {
        u32 unsecured_regs_arr_size;
        u32 arc_regs_arr_size;
-       int rc;
 };
 
 static void gaudi2_config_tpcs_pb_ranges(struct hl_device *hdev, int dcore, int inst, u32 offset,
-                                       void *data)
+                                               struct iterate_module_ctx *ctx)
 {
-       struct gaudi2_tpc_arc_pb_data *pb_data = (struct gaudi2_tpc_arc_pb_data *)data;
+       struct gaudi2_tpc_arc_pb_data *pb_data = ctx->data;
 
-       pb_data->rc |= hl_init_pb_ranges(hdev, HL_PB_SHARED, HL_PB_NA, 1,
+       ctx->rc = hl_init_pb_ranges(hdev, HL_PB_SHARED, HL_PB_NA, 1,
                                        offset, gaudi2_pb_dcr0_tpc0_arc,
                                        pb_data->arc_regs_arr_size,
                                        gaudi2_pb_dcr0_tpc0_arc_unsecured_regs,
        tpc_arc_pb_data.arc_regs_arr_size = ARRAY_SIZE(gaudi2_pb_dcr0_tpc0_arc);
        tpc_arc_pb_data.unsecured_regs_arr_size =
                        ARRAY_SIZE(gaudi2_pb_dcr0_tpc0_arc_unsecured_regs);
-       tpc_arc_pb_data.rc = 0;
+
        tpc_iter.fn = &gaudi2_config_tpcs_pb_ranges;
        tpc_iter.data = &tpc_arc_pb_data;
        gaudi2_iterate_tpcs(hdev, &tpc_iter);
 
-       return tpc_arc_pb_data.rc;
+       return tpc_iter.rc;
 }
 
 static int gaudi2_init_pb_sm_objs(struct hl_device *hdev)
 };
 
 static void gaudi2_ack_pb_tpc_config(struct hl_device *hdev, int dcore, int inst, u32 offset,
-                                       void *data)
+                                       struct iterate_module_ctx *ctx)
 {
-       struct gaudi2_ack_pb_tpc_data *pb_data = (struct gaudi2_ack_pb_tpc_data *)data;
+       struct gaudi2_ack_pb_tpc_data *pb_data = ctx->data;
 
        hl_ack_pb_single_dcore(hdev, offset, HL_PB_SINGLE_INSTANCE, HL_PB_NA,
                                gaudi2_pb_dcr0_tpc0, pb_data->tpc_regs_array_size);