[TVM] Fixed SPIR-V codegen incorrectly not declaring the interface for the entry point (#1400)

src/codegen/spirv/codegen_spirv.cc

@@ -35,7 +35,7 @@ std::vector<uint32_t> CodeGenSPIRV::BuildFunction(const LoweredFunc& f) {
       pod_args.push_back(arg);
     }
   }
-  spirv::Value func_ptr = builder_->DeclareKenrelFunction(f->name);
+  spirv::Value func_ptr = builder_->NewFunction();
   builder_->StartFunction(func_ptr);
 
   // All the POD arguments are passed in through PushConstant
@@ -56,6 +56,8 @@ std::vector<uint32_t> CodeGenSPIRV::BuildFunction(const LoweredFunc& f) {
   builder_->MakeInst(spv::OpReturn);
   builder_->MakeInst(spv::OpFunctionEnd);
 
+  builder_->CommitKernelFunction(func_ptr, f->name);
+
   return builder_->Finalize();
 }
 

src/codegen/spirv/ir_builder.cc

@@ -222,12 +222,21 @@ Value IRBuilder::GetPushConstant(
   return this->MakeValue(spv::OpLoad, v_type, ptr);
 }
 
-Value IRBuilder::DeclareKenrelFunction(const std::string& name) {
-  Value val = NewValue(t_void_func_, kFunction);
+Value IRBuilder::NewFunction() {
+  return NewValue(t_void_func_, kFunction);
+}
+
+void IRBuilder::CommitKernelFunction(const Value& func, const std::string& name) {
+  CHECK_EQ(func.flag, kFunction);
   ib_.Begin(spv::OpEntryPoint)
-      .AddSeq(spv::ExecutionModelGLCompute, val, name)
-      .Commit(&entry_);
-  return val;
+    .AddSeq(spv::ExecutionModelGLCompute, func, name);
+  if (workgroup_id_.id != 0) {
+    ib_.Add(workgroup_id_);
+  }
+  if (local_id_.id != 0) {
+    ib_.Add(local_id_);
+  }
+  ib_.Commit(&entry_);
 }
 
 void IRBuilder::StartFunction(const Value& func) {

src/codegen/spirv/ir_builder.h

@@ -485,11 +485,18 @@ class IRBuilder {
    */
   Value GetPushConstant(Value ptr_push_const, const SType& v_type, uint32_t index);
   /*!
-   * \brief Declare a kernel function
-   * \param name Name of the entry point.
+   * \brief Declare a new function
    * \return The created function ID.
    */
-  Value DeclareKenrelFunction(const std::string& name);
+  Value NewFunction();
+  /*!
+   * \brief Declare the entry point for a kernel function. This should be
+   * invoked after building the function so the builder is aware of which
+   * variables to declare as part of the function's interface.
+   * \param func The previously declared function.
+   * \param name Name of the entry point.
+   */
+  void CommitKernelFunction(const Value& func, const std::string& name);
   /*!
    * \brief Start function scope.
    * \param func function to be started.