diff --git a/renderdoc/data/glsl/debuguniforms.h b/renderdoc/data/glsl/debuguniforms.h
index 85d1363b4..34477728d 100644
--- a/renderdoc/data/glsl/debuguniforms.h
+++ b/renderdoc/data/glsl/debuguniforms.h
@@ -70,6 +70,7 @@
 #define OPENGL 1
 
 #define HISTOGRAM_UBOS
+#define HEATMAP_UBO
 
 #ifdef GL_ES
 #define OPENGL_ES 1
@@ -146,24 +147,6 @@ BINDING(0) uniform TexDisplayUBOData
 }
 INST_NAME(texdisplay);
 
-#define HEATMAP_DISABLED 0
-#define HEATMAP_LINEAR 1
-#define HEATMAP_TRISIZE 2
-
-#define HEATMAP_RAMPSIZE 22
-
-BINDING(1) uniform HeatmapData
-{
-  int HeatmapMode;
-  int DummyA;
-  int DummyB;
-  int DummyC;
-
-  // must match size of colorRamp on C++ side
-  vec4 ColorRamp[22];
-}
-INST_NAME(heatmap);
-
 BINDING(0) uniform FontUBOData
 {
   vec2 TextPosition;
@@ -224,6 +207,28 @@ INST_NAME(str);
 
 #endif
 
+#define HEATMAP_DISABLED 0
+#define HEATMAP_LINEAR 1
+#define HEATMAP_TRISIZE 2
+
+#define HEATMAP_RAMPSIZE 22
+
+#if defined(HEATMAP_UBO) || defined(__cplusplus)
+
+BINDING(1) uniform HeatmapData
+{
+  int HeatmapMode;
+  int DummyA;
+  int DummyB;
+  int DummyC;
+
+  // must match size of colorRamp on C++ side
+  vec4 ColorRamp[22];
+}
+INST_NAME(heatmap);
+
+#endif
+
 #if defined(HISTOGRAM_UBOS) || defined(__cplusplus)
 
 BINDING(2) uniform HistogramUBOData
diff --git a/renderdoc/driver/vulkan/vk_shader_cache.cpp b/renderdoc/driver/vulkan/vk_shader_cache.cpp
index 75622d6a0..aa1d8ec49 100644
--- a/renderdoc/driver/vulkan/vk_shader_cache.cpp
+++ b/renderdoc/driver/vulkan/vk_shader_cache.cpp
@@ -170,9 +170,13 @@ VulkanShaderCache::VulkanShaderCache(WrappedVulkan *driver)
     if(driverVersion.TexelFetchBrokenDriver())
       defines += "#define NO_TEXEL_FETCH\n";
 
+    // add UBO definitions as needed, to workaround ARM bug
     if(config.builtin == BuiltinShader::TextVS || config.builtin == BuiltinShader::TextFS)
       defines += "#define FONT_UBOS\n";
 
+    if(config.builtin == BuiltinShader::TexDisplayFS)
+      defines += "#define HEATMAP_UBO\n";
+
     GenerateGLSLShader(sources, eShaderVulkan, defines, GetDynamicEmbeddedResource(config.resource),
                        430, config.uniforms);