Optimize particle primitive shader

3 files changed, 29 insertions, 26 deletions

diff --git a/source/blender/draw/modes/object_mode.c b/source/blender/draw/modes/object_mode.c
index 05c9488ba82..f213102b0f0 100644
--- a/source/blender/draw/modes/object_mode.c
+++ b/source/blender/draw/modes/object_mode.c
@@ -204,6 +204,7 @@ static struct {
 	GPUShader *grid_sh;
 	GPUShader *part_dot_sh;
 	GPUShader *part_prim_sh;
+	GPUShader *part_axis_sh;
 	float camera_pos[3];
 	float grid_settings[5];
 	float grid_mat[4][4];
@@ -293,6 +294,10 @@ static void OBJECT_engine_init(void *vedata)
 		e_data.part_prim_sh = DRW_shader_create(datatoc_object_particle_prim_vert_glsl, NULL, datatoc_object_particle_prim_frag_glsl, NULL);
 	}
 
+	if (!e_data.part_axis_sh) {
+		e_data.part_axis_sh = DRW_shader_create(datatoc_object_particle_prim_vert_glsl, NULL, datatoc_object_particle_prim_frag_glsl, "#define USE_AXIS\n");
+	}
+
 	if (!e_data.part_dot_sh) {
 		e_data.part_dot_sh = DRW_shader_create(datatoc_object_particle_dot_vert_glsl, NULL, datatoc_object_particle_dot_frag_glsl, NULL);
 	}
@@ -442,6 +447,7 @@ static void OBJECT_engine_free(void)
 	DRW_SHADER_FREE_SAFE(e_data.object_empty_image_wire_sh);
 	DRW_SHADER_FREE_SAFE(e_data.grid_sh);
 	DRW_SHADER_FREE_SAFE(e_data.part_prim_sh);
+	DRW_SHADER_FREE_SAFE(e_data.part_axis_sh);
 	DRW_SHADER_FREE_SAFE(e_data.part_dot_sh);
 }
 
@@ -1421,14 +1427,18 @@ static void OBJECT_cache_populate(void *vedata, Object *ob)
 								break;
 							case PART_DRAW_CROSS:
 								shgrp = DRW_shgroup_instance_create(e_data.part_prim_sh, psl->particle, DRW_cache_particles_get_prim(PART_DRAW_CROSS));
+								DRW_shgroup_uniform_texture(shgrp, "ramp", globals_ramp);
+								DRW_shgroup_uniform_vec3(shgrp, "color", &ma->r, 1);
 								DRW_shgroup_uniform_int(shgrp, "screen_space", &screen_space[0], 1);
 								break;
 							case PART_DRAW_CIRC:
 								shgrp = DRW_shgroup_instance_create(e_data.part_prim_sh, psl->particle, DRW_cache_particles_get_prim(PART_DRAW_CIRC));
+								DRW_shgroup_uniform_texture(shgrp, "ramp", globals_ramp);
+								DRW_shgroup_uniform_vec3(shgrp, "color", &ma->r, 1);
 								DRW_shgroup_uniform_int(shgrp, "screen_space", &screen_space[1], 1);
 								break;
 							case PART_DRAW_AXIS:
-								shgrp = DRW_shgroup_instance_create(e_data.part_prim_sh, psl->particle, DRW_cache_particles_get_prim(PART_DRAW_AXIS));
+								shgrp = DRW_shgroup_instance_create(e_data.part_axis_sh, psl->particle, DRW_cache_particles_get_prim(PART_DRAW_AXIS));
 								DRW_shgroup_uniform_int(shgrp, "screen_space", &screen_space[0], 1);
 								break;
 							default:
@@ -1439,10 +1449,8 @@ static void OBJECT_cache_populate(void *vedata, Object *ob)
 							DRW_shgroup_attrib_float(shgrp, "pos", 3);
 							DRW_shgroup_attrib_float(shgrp, "rot", 4);
 							DRW_shgroup_attrib_float(shgrp, "val", 1);
-							DRW_shgroup_uniform_vec3(shgrp, "color", &ma->r, 1);
 							DRW_shgroup_uniform_short_to_int(shgrp, "draw_size", &part->draw_size, 1);
 							DRW_shgroup_uniform_float(shgrp, "pixel_size", DRW_viewport_pixelsize_get(), 1);
-							DRW_shgroup_uniform_texture(shgrp, "ramp", globals_ramp);
 							DRW_shgroup_instance_batch(shgrp, geom);
 						}
 					}
diff --git a/source/blender/draw/modes/shaders/object_particle_prim_frag.glsl b/source/blender/draw/modes/shaders/object_particle_prim_frag.glsl
index 0e58d70b69b..aa455a85cf0 100644
--- a/source/blender/draw/modes/shaders/object_particle_prim_frag.glsl
+++ b/source/blender/draw/modes/shaders/object_particle_prim_frag.glsl
@@ -1,26 +1,9 @@
 
-uniform vec3 color;
-uniform sampler1D ramp;
-
-flat in int finalAxis;
-flat in float finalVal;
+flat in vec4 finalColor;
 
 out vec4 fragColor;
 
 void main()
 {
-	if (finalAxis == -1) {
-		if (finalVal < 0.0) {
-			fragColor.rgb = color;
-		}
-		else {
-			fragColor.rgb = texture(ramp, finalVal).rgb;
-		}
-	}
-	else {
-		fragColor.rgb = vec3(0.0);
-		fragColor[finalAxis] = 1.0;
-	}
-
-	fragColor.a = 1.0;
+	fragColor = finalColor;
 }
diff --git a/source/blender/draw/modes/shaders/object_particle_prim_vert.glsl b/source/blender/draw/modes/shaders/object_particle_prim_vert.glsl
index 7b328c73de3..55b183b6933 100644
--- a/source/blender/draw/modes/shaders/object_particle_prim_vert.glsl
+++ b/source/blender/draw/modes/shaders/object_particle_prim_vert.glsl
@@ -6,6 +6,8 @@ uniform mat4 ProjectionMatrix;
 uniform int screen_space;
 uniform float pixel_size;
 uniform int draw_size;
+uniform vec3 color;
+uniform sampler1D ramp;
 
 in vec3 pos;
 in vec4 rot;
@@ -13,8 +15,7 @@ in float val;
 in vec3 inst_pos;
 in int axis;
 
-flat out int finalAxis;
-flat out float finalVal;
+flat out vec4 finalColor;
 
 vec3 rotate(vec3 vec, vec4 quat)
 {
@@ -47,6 +48,17 @@ void main()
 		gl_Position = ModelViewProjectionMatrix * vec4(pos + rotate(inst_pos * pix_size * size, rot), 1.0);
 	}
 
-	finalAxis = axis;
-	finalVal = val;
+#ifdef USE_AXIS
+	finalColor.rgb = vec3(0.0);
+	finalColor[axis] = 1.0;
+#else
+	if (val < 0.0) {
+		finalColor.rgb = color;
+	}
+	else {
+		finalColor.rgb = texture(ramp, val).rgb;
+	}
+#endif
+
+	finalColor.a = 1.0;
 }