for (cap = 1; cap <= cap_count; cap++) {
                void __iomem *register_block;
-               u32 hdr;
-               int decoder_cnt;
+               struct cxl_reg_map *rmap;
                u16 cap_id, offset;
-               u32 length;
+               u32 length, hdr;
 
                hdr = readl(base + cap * 0x4);
 
                cap_id = FIELD_GET(CXL_CM_CAP_HDR_ID_MASK, hdr);
                offset = FIELD_GET(CXL_CM_CAP_PTR_MASK, hdr);
                register_block = base + offset;
+               hdr = readl(register_block);
 
+               rmap = NULL;
                switch (cap_id) {
-               case CXL_CM_CAP_CAP_ID_HDM:
+               case CXL_CM_CAP_CAP_ID_HDM: {
+                       int decoder_cnt;
+
                        dev_dbg(dev, "found HDM decoder capability (0x%x)\n",
                                offset);
 
-                       hdr = readl(register_block);
-
                        decoder_cnt = cxl_hdm_decoder_count(hdr);
                        length = 0x20 * decoder_cnt + 0x10;
-
-                       map->hdm_decoder.valid = true;
-                       map->hdm_decoder.offset = CXL_CM_OFFSET + offset;
-                       map->hdm_decoder.size = length;
+                       rmap = &map->hdm_decoder;
                        break;
+               }
                default:
                        dev_dbg(dev, "Unknown CM cap ID: %d (0x%x)\n", cap_id,
                                offset);
                        break;
                }
+
+               if (!rmap)
+                       continue;
+               rmap->valid = true;
+               rmap->offset = CXL_CM_OFFSET + offset;
+               rmap->size = length;
        }
 }
 EXPORT_SYMBOL_NS_GPL(cxl_probe_component_regs, CXL);
        cap_count = FIELD_GET(CXLDEV_CAP_ARRAY_COUNT_MASK, cap_array);
 
        for (cap = 1; cap <= cap_count; cap++) {
+               struct cxl_reg_map *rmap;
                u32 offset, length;
                u16 cap_id;
 
                offset = readl(base + cap * 0x10 + 0x4);
                length = readl(base + cap * 0x10 + 0x8);
 
+               rmap = NULL;
                switch (cap_id) {
                case CXLDEV_CAP_CAP_ID_DEVICE_STATUS:
                        dev_dbg(dev, "found Status capability (0x%x)\n", offset);
-
-                       map->status.valid = true;
-                       map->status.offset = offset;
-                       map->status.size = length;
+                       rmap = &map->status;
                        break;
                case CXLDEV_CAP_CAP_ID_PRIMARY_MAILBOX:
                        dev_dbg(dev, "found Mailbox capability (0x%x)\n", offset);
-                       map->mbox.valid = true;
-                       map->mbox.offset = offset;
-                       map->mbox.size = length;
+                       rmap = &map->mbox;
                        break;
                case CXLDEV_CAP_CAP_ID_SECONDARY_MAILBOX:
                        dev_dbg(dev, "found Secondary Mailbox capability (0x%x)\n", offset);
                        break;
                case CXLDEV_CAP_CAP_ID_MEMDEV:
                        dev_dbg(dev, "found Memory Device capability (0x%x)\n", offset);
-                       map->memdev.valid = true;
-                       map->memdev.offset = offset;
-                       map->memdev.size = length;
+                       rmap = &map->memdev;
                        break;
                default:
                        if (cap_id >= 0x8000)
                                dev_dbg(dev, "Unknown cap ID: %#x offset: %#x\n", cap_id, offset);
                        break;
                }
+
+               if (!rmap)
+                       continue;
+               rmap->valid = true;
+               rmap->offset = offset;
+               rmap->size = length;
        }
 }
 EXPORT_SYMBOL_NS_GPL(cxl_probe_device_regs, CXL);