fix shaders

Author: bobcao3

python/taichi/shaders/SetImage_vk.frag

@@ -9,11 +9,16 @@ layout(location = 0) out vec4 out_color;
 layout(binding = 1) uniform UBO {
   float x_factor;
   float y_factor;
+  int transposed;
 }
 ubo;
 
 void main() {
-  vec2 coord = frag_texcoord.yx * vec2(ubo.y_factor,ubo.x_factor);
+  vec2 coord = vec2(ubo.y_factor,ubo.x_factor);
+  if (ubo.transposed != 0) {
+    coord *= frag_texcoord.yx;
+  } else {
+    coord *= frag_texcoord.xy;
+  }
   out_color = texture(texSampler, coord);
-  // out_color = vec4(frag_texcoord.xy,0,1);
 }

python/taichi/shaders/SetImage_vk_frag.spv

@@ -2372,14 +2372,14 @@ void KernelCodegen::run(TaichiKernelAttributes &kernel_attribs,
       std::vector<uint32_t> &spirv = success ? optimized_spv : task_res.spirv_code;
 
       std::string spirv_asm;
-      spirv_tools_->Disassemble(optimized_spv, &spirv_asm);
+      spirv_tools_->Disassemble(spirv, &spirv_asm);
       auto kernel_name = tp.ti_kernel_name;
       TI_WARN("SPIR-V Assembly dump for {} :\n{}\n\n", kernel_name, spirv_asm);
 
       std::ofstream fout(kernel_name + ".spv",
                          std::ios::binary | std::ios::out);
-      fout.write(reinterpret_cast<const char *>(optimized_spv.data()),
-                 optimized_spv.size() * sizeof(uint32_t));
+      fout.write(reinterpret_cast<const char *>(spirv.data()),
+                 spirv.size() * sizeof(uint32_t));
       fout.close();
     }
 #endif