diff --git a/drivers/pci/msi.c b/drivers/pci/msi.c
index d4497141d083a71d5d5206496fee58fbddc5cb13..4a7da3c3e0353c3c746e9b10be9a093c8d085920 100644
--- a/drivers/pci/msi.c
+++ b/drivers/pci/msi.c
@@ -1243,6 +1243,10 @@ static void pci_msi_domain_update_chip_ops(struct msi_domain_info *info)
 	BUG_ON(!chip);
 	if (!chip->irq_write_msi_msg)
 		chip->irq_write_msi_msg = pci_msi_domain_write_msg;
+	if (!chip->irq_mask)
+		chip->irq_mask = pci_msi_mask_irq;
+	if (!chip->irq_unmask)
+		chip->irq_unmask = pci_msi_unmask_irq;
 }
 
 /**
diff --git a/kernel/irq/msi.c b/kernel/irq/msi.c
index 7e6512b9dc1ff2682394cdd0fea9a8c6d01cf6e9..be9149f62eb86e63ac06194d4eeaa4065823ea62 100644
--- a/kernel/irq/msi.c
+++ b/kernel/irq/msi.c
@@ -228,11 +228,7 @@ static void msi_domain_update_chip_ops(struct msi_domain_info *info)
 {
 	struct irq_chip *chip = info->chip;
 
-	BUG_ON(!chip);
-	if (!chip->irq_mask)
-		chip->irq_mask = pci_msi_mask_irq;
-	if (!chip->irq_unmask)
-		chip->irq_unmask = pci_msi_unmask_irq;
+	BUG_ON(!chip || !chip->irq_mask || !chip->irq_unmask);
 	if (!chip->irq_set_affinity)
 		chip->irq_set_affinity = msi_domain_set_affinity;
 }