uint32_t ODM_MEM_PWR_CTRL3;
        uint32_t DMU_MEM_PWR_CNTL;
        uint32_t MMHUBBUB_MEM_PWR_CNTL;
+       uint32_t DCHUBBUB_ARB_HOSTVM_CNTL;
 };
  /* set field name */
 #define HWS_SF(blk_name, reg_name, field_name, post_fix)\
        type DOMAIN_POWER_FORCEON;\
        type DOMAIN_POWER_GATE;\
        type DOMAIN_PGFSM_PWR_STATUS;\
-       type HPO_HDMISTREAMCLK_G_GATE_DIS;
+       type HPO_HDMISTREAMCLK_G_GATE_DIS;\
+       type DISABLE_HOSTVM_FORCE_ALLOW_PSTATE;
 
 struct dce_hwseq_shift {
        HWSEQ_REG_FIELD_LIST(uint8_t)
 
 #include "dce/dmub_outbox.h"
 #include "dc_link_dp.h"
 #include "inc/link_dpcd.h"
+#include "dcn10/dcn10_hw_sequencer.h"
 
 #define DC_LOGGER_INIT(logger)
 
        }
        return false;
 }
+
+static void apply_riommu_invalidation_wa(struct dc *dc)
+{
+       struct dce_hwseq *hws = dc->hwseq;
+
+       if (!hws->wa.early_riommu_invalidation)
+               return;
+
+       REG_UPDATE(DCHUBBUB_ARB_HOSTVM_CNTL, DISABLE_HOSTVM_FORCE_ALLOW_PSTATE, 0);
+}
+
+void dcn31_init_pipes(struct dc *dc, struct dc_state *context)
+{
+       dcn10_init_pipes(dc, context);
+       apply_riommu_invalidation_wa(dc);
+
+}
 
                struct dc_state *context);
 bool dcn31_is_abm_supported(struct dc *dc,
                struct dc_state *context, struct dc_stream_state *stream);
+void dcn31_init_pipes(struct dc *dc, struct dc_state *context);
 
 #endif /* __DC_HWSS_DCN31_H__ */
 
 };
 
 static const struct hwseq_private_funcs dcn31_private_funcs = {
-       .init_pipes = dcn10_init_pipes,
+       .init_pipes = dcn31_init_pipes,
        .update_plane_addr = dcn20_update_plane_addr,
        .plane_atomic_disconnect = dcn10_plane_atomic_disconnect,
        .update_mpcc = dcn20_update_mpcc,
 
 
 #define HWSEQ_DCN31_REG_LIST()\
        SR(DCHUBBUB_GLOBAL_TIMER_CNTL), \
+       SR(DCHUBBUB_ARB_HOSTVM_CNTL), \
        SR(DIO_MEM_PWR_CTRL), \
        SR(ODM_MEM_PWR_CTRL3), \
        SR(DMU_MEM_PWR_CNTL), \
 #define HWSEQ_DCN31_MASK_SH_LIST(mask_sh)\
        HWSEQ_DCN_MASK_SH_LIST(mask_sh), \
        HWS_SF(, DCHUBBUB_GLOBAL_TIMER_CNTL, DCHUBBUB_GLOBAL_TIMER_REFDIV, mask_sh), \
+       HWS_SF(, DCHUBBUB_ARB_HOSTVM_CNTL, DISABLE_HOSTVM_FORCE_ALLOW_PSTATE, mask_sh), \
        HWS_SF(, DOMAIN0_PG_CONFIG, DOMAIN_POWER_FORCEON, mask_sh), \
        HWS_SF(, DOMAIN0_PG_CONFIG, DOMAIN_POWER_GATE, mask_sh), \
        HWS_SF(, DOMAIN1_PG_CONFIG, DOMAIN_POWER_FORCEON, mask_sh), \
                hws->regs = &hwseq_reg;
                hws->shifts = &hwseq_shift;
                hws->masks = &hwseq_mask;
+               hws->wa.early_riommu_invalidation = true;
        }
        return hws;
 }
 
        bool DEGVIDCN10_254;
        bool DEGVIDCN21;
        bool disallow_self_refresh_during_multi_plane_transition;
+       bool early_riommu_invalidation;
 };
 
 struct hwseq_wa_state {