[Vulkan] Some fixes of Vulkan codegen (#309)

- add exp2 dispatch
- handle thread invariant op
- disable smem merge for vulkan

Co-authored-by: spectrometerHBH <[email protected]>

Author: junrushao

src/target/spirv/codegen_spirv.cc

@@ -509,6 +509,8 @@ spirv::Value CodeGenSPIRV::VisitExpr_(const CallNode* op) {
     spirv::SType ptr_type = builder_->GetPointerType(ele_stype, buffer_val.stype.storage_class);
     ICHECK(var_map_.count(buffer_node));
     return builder_->StructArrayAccess(ptr_type, var_map_[buffer_node], MakeValue(index));
+  } else if (op->op.same_as(builtin::tvm_thread_invariant())) {
+    return MakeValue(op->args[0]);
   } else {
     LOG(FATAL) << "Unresolved call  " << op->op;
   }

src/target/spirv/intrin_rule_spirv.cc

@@ -82,6 +82,9 @@ TVM_REGISTER_OP("tir.fabs")
 TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>("vulkan.FLowerIntrinsic",
                                                      DispatchGLSLPureIntrin<GLSLstd450Exp>);
 
+TVM_REGISTER_OP("tir.exp2")
+    .set_attr<FLowerIntrinsic>("vulkan.FLowerIntrinsic", DispatchGLSLPureIntrin<GLSLstd450Exp2>);
+
 TVM_REGISTER_OP("tir.sin").set_attr<FLowerIntrinsic>("vulkan.FLowerIntrinsic",
                                                      DispatchGLSLPureIntrin<GLSLstd450Sin>);
 

src/tir/transforms/merge_shared_memory_allocations.cc

@@ -662,6 +662,11 @@ namespace transform {
 Pass MergeSharedMemoryAllocations() {
   auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) {
     bool merge_static_smem = ctx->GetConfig<Bool>("tir.merge_static_smem", Bool(false)).value();
+    // disable this pass for Vulkan
+    auto target = Target::Current(true);
+    if (target.defined() && target->kind->name == "vulkan") {
+      return f;
+    }
     auto* n = f.CopyOnWrite();
     n->body = MergeSharedMemoryAllocations(std::move(n->body), merge_static_smem);
     return f;

src/tir/transforms/storage_rewrite.cc

@@ -1705,8 +1705,13 @@ namespace transform {
 Pass StorageRewrite() {
   auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) {
     bool merge_static_smem = ctx->GetConfig<Bool>("tir.merge_static_smem", Bool(false)).value();
+    // disable merge_static_smem for Vulkan
+    auto target = Target::Current(true);
+    if (target.defined() && target->kind->name == "vulkan") {
+      merge_static_smem = false;
+    }
     auto* n = f.CopyOnWrite();
-    n->body = StoragePlanRewriter().Rewrite(std::move(n->body), true, !merge_static_smem);
+    n->body = StoragePlanRewriter().Rewrite(std::move(n->body), true, merge_static_smem);
     // Parameters may not be rewritten, but internal allocations may.
     // Vectorization of AllocateConst is currently disabled, as it has
     // indexing issues for types that include padding (e.g. int8x3