3 files changed, 29 insertions, 14 deletions

diff --git a/src/gallium/drivers/iris/iris_context.h b/src/gallium/drivers/iris/iris_context.h
index ac5bdaf26e3..596b195b5a8 100644
--- a/src/gallium/drivers/iris/iris_context.h
+++ b/src/gallium/drivers/iris/iris_context.h
@@ -628,9 +628,9 @@ const struct shader_info *iris_get_shader_info(const struct iris_context *ice,
                                                gl_shader_stage stage);
 unsigned iris_get_shader_num_ubos(const struct iris_context *ice,
                                   gl_shader_stage stage);
-uint32_t iris_get_scratch_space(struct iris_context *ice,
-                                unsigned per_thread_scratch,
-                                gl_shader_stage stage);
+struct iris_bo *iris_get_scratch_space(struct iris_context *ice,
+                                       unsigned per_thread_scratch,
+                                       gl_shader_stage stage);
 
 /* iris_program_cache.c */
 
diff --git a/src/gallium/drivers/iris/iris_program.c b/src/gallium/drivers/iris/iris_program.c
index 8a3a77ddb05..ba8b60d3e6d 100644
--- a/src/gallium/drivers/iris/iris_program.c
+++ b/src/gallium/drivers/iris/iris_program.c
@@ -1481,10 +1481,8 @@ iris_fill_cs_push_const_buffer(struct brw_cs_prog_data *cs_prog_data,
 
 /**
  * Allocate scratch BOs as needed for the given per-thread size and stage.
- *
- * Returns the 32-bit "Scratch Space Base Pointer" value.
  */
-uint32_t
+struct iris_bo *
 iris_get_scratch_space(struct iris_context *ice,
                        unsigned per_thread_scratch,
                        gl_shader_stage stage)
@@ -1526,7 +1524,7 @@ iris_get_scratch_space(struct iris_context *ice,
       *bop = iris_bo_alloc(bufmgr, "scratch", size, IRIS_MEMZONE_SHADER);
    }
 
-   return (*bop)->gtt_offset;
+   return *bop;
 }
 
 void
diff --git a/src/gallium/drivers/iris/iris_state.c b/src/gallium/drivers/iris/iris_state.c
index e4e42f50d82..4c95d6d90e0 100644
--- a/src/gallium/drivers/iris/iris_state.c
+++ b/src/gallium/drivers/iris/iris_state.c
@@ -3113,8 +3113,9 @@ KSP(const struct iris_compiled_shader *shader)
    pkt.Enable           = true;                                           \
                                                                           \
    if (prog_data->total_scratch) {                                        \
-      uint32_t scratch_addr =                                             \
+      struct iris_bo *bo =                                                \
          iris_get_scratch_space(ice, prog_data->total_scratch, stage);    \
+      uint32_t scratch_addr = bo->gtt_offset;                             \
       pkt.PerThreadScratchSpace = ffs(prog_data->total_scratch) - 11;     \
       pkt.ScratchSpaceBasePointer = rw_bo(NULL, scratch_addr);            \
    }
@@ -3308,9 +3309,10 @@ iris_store_fs_state(struct iris_context *ice,
          KSP(shader) + brw_wm_prog_data_prog_offset(wm_prog_data, ps, 2);
 
       if (prog_data->total_scratch) {
-         uint32_t scratch_addr =
+         struct iris_bo *bo =
             iris_get_scratch_space(ice, prog_data->total_scratch,
                                    MESA_SHADER_FRAGMENT);
+         uint32_t scratch_addr = bo->gtt_offset;
          ps.PerThreadScratchSpace = ffs(prog_data->total_scratch) - 11;
          ps.ScratchSpaceBasePointer = rw_bo(NULL, scratch_addr);
       }
@@ -3804,12 +3806,19 @@ iris_restore_render_saved_bos(struct iris_context *ice,
    for (int stage = 0; stage <= MESA_SHADER_FRAGMENT; stage++) {
       if (clean & (IRIS_DIRTY_VS << stage)) {
          struct iris_compiled_shader *shader = ice->shaders.prog[stage];
+
          if (shader) {
             struct iris_bo *bo = iris_resource_bo(shader->assembly.res);
             iris_use_pinned_bo(batch, bo, false);
-         }
 
-         // XXX: scratch buffer
+            struct brw_stage_prog_data *prog_data = shader->prog_data;
+
+            if (prog_data->total_scratch > 0) {
+               struct iris_bo *bo =
+                  iris_get_scratch_space(ice, prog_data->total_scratch, stage);
+               iris_use_pinned_bo(batch, bo, true);
+            }
+         }
       }
    }
 
@@ -3886,12 +3895,19 @@ iris_restore_compute_saved_bos(struct iris_context *ice,
 
    if (clean & IRIS_DIRTY_CS) {
       struct iris_compiled_shader *shader = ice->shaders.prog[stage];
+
       if (shader) {
          struct iris_bo *bo = iris_resource_bo(shader->assembly.res);
          iris_use_pinned_bo(batch, bo, false);
-      }
 
-      // XXX: scratch buffer
+         struct brw_stage_prog_data *prog_data = shader->prog_data;
+
+         if (prog_data->total_scratch > 0) {
+            struct iris_bo *bo =
+               iris_get_scratch_space(ice, prog_data->total_scratch, stage);
+            iris_use_pinned_bo(batch, bo, true);
+         }
+      }
    }
 }
 
@@ -4695,9 +4711,10 @@ iris_upload_compute_state(struct iris_context *ice,
 
       iris_emit_cmd(batch, GENX(MEDIA_VFE_STATE), vfe) {
          if (prog_data->total_scratch) {
-            uint32_t scratch_addr =
+            struct iris_bo *bo =
                iris_get_scratch_space(ice, prog_data->total_scratch,
                                       MESA_SHADER_COMPUTE);
+            uint32_t scratch_addr = bo->gtt_offset;
             vfe.PerThreadScratchSpace = ffs(prog_data->total_scratch) - 11;
             vfe.ScratchSpaceBasePointer = rw_bo(NULL, scratch_addr);
          }