iris: Init aux map state for compute engine

We need to write the aux map address during the batch initialization for
compute engine as well otherwise we would run into gpu hang with
compression enabled.

Closes: https://gitlab.freedesktop.org/mesa/mesa/-/issues/10194

Signed-off-by: Sagar Ghuge <sagar.ghuge@intel.com>
Reviewed-by: Tapani Pälli <tapani.palli@intel.com>
Reviewed-by: José Roberto de Souza <jose.souza@intel.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/26342>

src/gallium/drivers/iris/iris_state.c

@@ -6153,14 +6153,22 @@ static void
 init_aux_map_state(struct iris_batch *batch)
 {
    struct iris_screen *screen = batch->screen;
+   const struct intel_device_info *devinfo = screen->devinfo;
    void *aux_map_ctx = iris_bufmgr_get_aux_map_context(screen->bufmgr);
    if (!aux_map_ctx)
       return;
 
    uint64_t base_addr = intel_aux_map_get_base(aux_map_ctx);
    assert(base_addr != 0 && align64(base_addr, 32 * 1024) == base_addr);
-   iris_load_register_imm64(batch, GENX(GFX_AUX_TABLE_BASE_ADDR_num),
-                            base_addr);
+
+   bool use_compute_reg = batch->name == IRIS_BATCH_COMPUTE &&
+                          devinfo->has_compute_engine &&
+                          debug_get_bool_option("INTEL_COMPUTE_CLASS", false);
+
+   uint32_t reg = use_compute_reg ? GENX(COMPCS0_AUX_TABLE_BASE_ADDR_num) :
+                                    GENX(GFX_AUX_TABLE_BASE_ADDR_num);
+
+   iris_load_register_imm64(batch, reg, base_addr);
 }
 #endif