3 files changed, 24 insertions, 17 deletions

diff --git a/src/gallium/drivers/radeonsi/si_state.c b/src/gallium/drivers/radeonsi/si_state.c
index 54151eb0d19..c845df17009 100644
--- a/src/gallium/drivers/radeonsi/si_state.c
+++ b/src/gallium/drivers/radeonsi/si_state.c
@@ -2283,6 +2283,7 @@ int si_shader_select(struct pipe_context *ctx,
 			FREE(shader);
 			return r;
 		}
+		si_shader_init_pm4_state(shader);
 		sel->num_shaders++;
 	}
 
diff --git a/src/gallium/drivers/radeonsi/si_state.h b/src/gallium/drivers/radeonsi/si_state.h
index f70bddfb8ec..4f5140ca188 100644
--- a/src/gallium/drivers/radeonsi/si_state.h
+++ b/src/gallium/drivers/radeonsi/si_state.h
@@ -31,6 +31,7 @@
 #include "radeon/r600_pipe_common.h"
 
 struct si_screen;
+struct si_shader;
 
 struct si_state_blend {
 	struct si_pm4_state	pm4;
@@ -270,6 +271,7 @@ unsigned si_tile_mode_index(struct r600_texture *rtex, unsigned level, bool sten
 /* si_state_draw.c */
 extern const struct r600_atom si_atom_cache_flush;
 extern const struct r600_atom si_atom_msaa_config;
+void si_shader_init_pm4_state(struct si_shader *shader);
 void si_emit_cache_flush(struct r600_common_context *sctx, struct r600_atom *atom);
 void si_draw_vbo(struct pipe_context *ctx, const struct pipe_draw_info *dinfo);
 
diff --git a/src/gallium/drivers/radeonsi/si_state_draw.c b/src/gallium/drivers/radeonsi/si_state_draw.c
index eb7ed403f1c..707e2343494 100644
--- a/src/gallium/drivers/radeonsi/si_state_draw.c
+++ b/src/gallium/drivers/radeonsi/si_state_draw.c
@@ -300,6 +300,27 @@ static void si_shader_ps(struct si_shader *shader)
 		       S_00B02C_USER_SGPR(num_user_sgprs));
 }
 
+void si_shader_init_pm4_state(struct si_shader *shader)
+{
+	switch (shader->selector->type) {
+	case PIPE_SHADER_VERTEX:
+		if (shader->key.vs.as_es)
+			si_shader_es(shader);
+		else
+			si_shader_vs(shader);
+		break;
+	case PIPE_SHADER_GEOMETRY:
+		si_shader_gs(shader);
+		si_shader_vs(shader->gs_copy_shader);
+		break;
+	case PIPE_SHADER_FRAGMENT:
+		si_shader_ps(shader);
+		break;
+	default:
+		assert(0);
+	}
+}
+
 /*
  * Drawing
  */
@@ -598,22 +619,12 @@ static void si_update_derived_state(struct si_context *sctx)
 
 	if (sctx->gs_shader) {
 		si_shader_select(ctx, sctx->gs_shader);
-
-		if (!sctx->gs_shader->current->pm4) {
-			si_shader_gs(sctx->gs_shader->current);
-			si_shader_vs(sctx->gs_shader->current->gs_copy_shader);
-		}
-
 		si_pm4_bind_state(sctx, gs, sctx->gs_shader->current->pm4);
 		si_pm4_bind_state(sctx, vs, sctx->gs_shader->current->gs_copy_shader->pm4);
 
 		sctx->b.streamout.stride_in_dw = sctx->gs_shader->so.stride;
 
 		si_shader_select(ctx, sctx->vs_shader);
-
-		if (!sctx->vs_shader->current->pm4)
-			si_shader_es(sctx->vs_shader->current);
-
 		si_pm4_bind_state(sctx, es, sctx->vs_shader->current->pm4);
 
 		if (!sctx->gs_rings)
@@ -639,10 +650,6 @@ static void si_update_derived_state(struct si_context *sctx)
 		si_pm4_bind_state(sctx, gs_onoff, sctx->gs_on);
 	} else {
 		si_shader_select(ctx, sctx->vs_shader);
-
-		if (!sctx->vs_shader->current->pm4)
-			si_shader_vs(sctx->vs_shader->current);
-
 		si_pm4_bind_state(sctx, vs, sctx->vs_shader->current->pm4);
 
 		sctx->b.streamout.stride_in_dw = sctx->vs_shader->so.stride;
@@ -671,9 +678,6 @@ static void si_update_derived_state(struct si_context *sctx)
 		sctx->ps_shader->current = sel->current;
 	}
 
-	if (!sctx->ps_shader->current->pm4)
-		si_shader_ps(sctx->ps_shader->current);
-
 	si_pm4_bind_state(sctx, ps, sctx->ps_shader->current->pm4);
 
 	if (si_pm4_state_changed(sctx, ps) || si_pm4_state_changed(sctx, vs)) {