ipmmu_ctx_write(domain->mmu->root, domain->context_id, reg, data);
 }
 
+static u32 ipmmu_utlb_reg(struct ipmmu_vmsa_device *mmu, unsigned int reg)
+{
+       return reg;
+}
+
+static void ipmmu_imuasid_write(struct ipmmu_vmsa_device *mmu,
+                               unsigned int utlb, u32 data)
+{
+       ipmmu_write(mmu, ipmmu_utlb_reg(mmu, IMUASID(utlb)), data);
+}
+
+static void ipmmu_imuctr_write(struct ipmmu_vmsa_device *mmu,
+                              unsigned int utlb, u32 data)
+{
+       ipmmu_write(mmu, ipmmu_utlb_reg(mmu, IMUCTR(utlb)), data);
+}
+
 /* -----------------------------------------------------------------------------
  * TLB and microTLB Management
  */
         */
 
        /* TODO: What should we set the ASID to ? */
-       ipmmu_write(mmu, IMUASID(utlb), 0);
+       ipmmu_imuasid_write(mmu, utlb, 0);
        /* TODO: Do we need to flush the microTLB ? */
-       ipmmu_write(mmu, IMUCTR(utlb),
-                   IMUCTR_TTSEL_MMU(domain->context_id) | IMUCTR_FLUSH |
-                   IMUCTR_MMUEN);
+       ipmmu_imuctr_write(mmu, utlb, IMUCTR_TTSEL_MMU(domain->context_id) |
+                                     IMUCTR_FLUSH | IMUCTR_MMUEN);
        mmu->utlb_ctx[utlb] = domain->context_id;
 }
 
 {
        struct ipmmu_vmsa_device *mmu = domain->mmu;
 
-       ipmmu_write(mmu, IMUCTR(utlb), 0);
+       ipmmu_imuctr_write(mmu, utlb, 0);
        mmu->utlb_ctx[utlb] = IPMMU_CTX_INVALID;
 }