[wgsl-writer] Add GenerateEntryPoint

This Cl adds the preliminary GenerateEntryPoint method to the WGSL
writer.

Bug: tint:211
Change-Id: Ib414ff66d482179f10eeeb890f6127bc585cd664
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/28045
Reviewed-by: David Neto <dneto@google.com>
Commit-Queue: David Neto <dneto@google.com>

include/tint/tint.h

@@ -18,6 +18,7 @@
 // TODO(tint:88): When implementing support for an install target, all of these
 //                headers will need to be moved to include/tint/.
 
+#include "src/ast/pipeline_stage.h"
 #include "src/context.h"
 #include "src/reader/reader.h"
 #include "src/type_determiner.h"

samples/main.cc

@@ -46,6 +46,10 @@ struct Options {
   bool dump_ast = false;
 
   Format format = Format::kNone;
+
+  bool emit_single_entry_point = false;
+  tint::ast::PipelineStage stage;
+  std::string ep_name;
 };
 
 const char kUsage[] = R"(Usage: tint [options] <input-file>
@@ -60,6 +64,7 @@ const char kUsage[] = R"(Usage: tint [options] <input-file>
                                    .metal  -> msl
                                    .hlsl   -> hlsl
                                If none matches, then default to SPIR-V assembly.
+  -ep <compute|fragment|vertex> <name>  -- Output single entry point
   --output-file <name>      -- Output file name.  Use "-" for standard output
   -o <name>                 -- Output file name.  Use "-" for standard output
   --parse-only              -- Stop after parsing the input
@@ -138,6 +143,19 @@ Format infer_format(const std::string& filename) {
   return Format::kNone;
 }
 
+tint::ast::PipelineStage convert_to_pipeline_stage(const std::string& name) {
+  if (name == "compute") {
+    return tint::ast::PipelineStage::kCompute;
+  }
+  if (name == "fragment") {
+    return tint::ast::PipelineStage::kFragment;
+  }
+  if (name == "vertex") {
+    return tint::ast::PipelineStage::kVertex;
+  }
+  return tint::ast::PipelineStage::kNone;
+}
+
 bool ParseArgs(const std::vector<std::string>& args, Options* opts) {
   for (size_t i = 1; i < args.size(); ++i) {
     const std::string& arg = args[i];
@@ -153,6 +171,22 @@ bool ParseArgs(const std::vector<std::string>& args, Options* opts) {
         std::cerr << "Unknown output format: " << args[i] << std::endl;
         return false;
       }
+    } else if (arg == "-ep") {
+      if (i + 2 >= args.size()) {
+        std::cerr << "Missing values for -ep" << std::endl;
+        return false;
+      }
+      i++;
+      opts->stage = convert_to_pipeline_stage(args[i]);
+      if (opts->stage == tint::ast::PipelineStage::kNone) {
+        std::cerr << "Invalid pipeline stage: " << args[i] << std::endl;
+        return false;
+      }
+
+      i++;
+      opts->ep_name = args[i];
+      opts->emit_single_entry_point = true;
+
     } else if (arg == "-o" || arg == "--output-name") {
       ++i;
       if (i >= args.size()) {
@@ -471,10 +505,17 @@ int main(int argc, const char** argv) {
     return 1;
   }
 
+  if (options.emit_single_entry_point) {
+    if (!writer->GenerateEntryPoint(options.stage, options.ep_name)) {
+      std::cerr << "Failed to generate: " << writer->error() << std::endl;
+      return 1;
+    }
+  } else {
     if (!writer->Generate()) {
       std::cerr << "Failed to generate: " << writer->error() << std::endl;
       return 1;
     }
+  }
 
 #if TINT_BUILD_SPV_WRITER
   if (options.format == Format::kSpvAsm) {

src/ast/function.cc

@@ -156,6 +156,15 @@ void Function::add_ancestor_entry_point(const std::string& ep) {
   ancestor_entry_points_.push_back(ep);
 }
 
+bool Function::HasAncestorEntryPoint(const std::string& name) const {
+  for (const auto& point : ancestor_entry_points_) {
+    if (point == name) {
+      return true;
+    }
+  }
+  return false;
+}
+
 const Statement* Function::get_last_statement() const {
   return body_->last();
 }

src/ast/function.h

@@ -115,6 +115,10 @@ class Function : public Node {
   const std::vector<std::string>& ancestor_entry_points() const {
     return ancestor_entry_points_;
   }
+  /// Checks if the given entry point is an ancestor
+  /// @param name the entry point name
+  /// @returns true if |name| is an ancestor entry point of this function
+  bool HasAncestorEntryPoint(const std::string& name) const;
 
   /// Sets the return type of the function
   /// @param type the return type

src/writer/wgsl/generator.cc

@@ -38,8 +38,13 @@ bool Generator::Generate() {
   return ret;
 }
 
-bool Generator::GenerateEntryPoint(ast::PipelineStage, const std::string&) {
-  return false;
+bool Generator::GenerateEntryPoint(ast::PipelineStage stage,
+                                   const std::string& name) {
+  auto ret = impl_->GenerateEntryPoint(module_, stage, name);
+  if (!ret) {
+    error_ = impl_->error();
+  }
+  return ret;
 }
 
 std::string Generator::result() const {

src/writer/wgsl/generator_impl.cc

@@ -112,6 +112,101 @@ bool GeneratorImpl::Generate(const ast::Module& module) {
   return true;
 }
 
+bool GeneratorImpl::GenerateEntryPoint(const ast::Module& module,
+                                       ast::PipelineStage stage,
+                                       const std::string& name) {
+  // TODO(dsinclair): We're always emitting imports even if they aren't needed.
+  for (const auto& import : module.imports()) {
+    if (!EmitImport(import.get())) {
+      return false;
+    }
+  }
+  if (!module.imports().empty()) {
+    out_ << std::endl;
+  }
+
+  bool found_entry_point = false;
+  std::string ep_function_name = "";
+  for (const auto& ep : module.entry_points()) {
+    std::string ep_name = ep->name();
+    if (ep_name.empty()) {
+      ep_name = ep->function_name();
+    }
+    ep_function_name = ep->function_name();
+
+    if (ep->stage() != stage || ep_name != name) {
+      continue;
+    }
+    if (!EmitEntryPoint(ep.get())) {
+      return false;
+    }
+    found_entry_point = true;
+    break;
+  }
+  out_ << std::endl;
+
+  if (!found_entry_point) {
+    error_ = "Unable to find requested entry point: " + name;
+    return false;
+  }
+
+  // TODO(dsinclair): We always emit aliases even if they aren't strictly needed
+  for (auto* const alias : module.alias_types()) {
+    if (!EmitAliasType(alias)) {
+      return false;
+    }
+  }
+  if (!module.alias_types().empty()) {
+    out_ << std::endl;
+  }
+
+  // TODO(dsinclair): This should be smarter and only emit needed const
+  // variables
+  for (const auto& var : module.global_variables()) {
+    if (!var->is_const()) {
+      continue;
+    }
+    if (!EmitVariable(var.get())) {
+      return false;
+    }
+  }
+
+  auto* func = module.FindFunctionByName(ep_function_name);
+  if (!func) {
+    error_ = "Unable to find entry point function: " + ep_function_name;
+    return false;
+  }
+
+  bool found_func_variable = false;
+  for (auto* var : func->referenced_module_variables()) {
+    if (!EmitVariable(var)) {
+      return false;
+    }
+    found_func_variable = true;
+  }
+  if (found_func_variable) {
+    out_ << std::endl;
+  }
+
+  for (const auto& f : module.functions()) {
+    if (!f->HasAncestorEntryPoint(name)) {
+      continue;
+    }
+
+    if (!EmitFunction(f.get())) {
+      return false;
+    }
+    out_ << std::endl;
+  }
+
+  if (!EmitFunction(func)) {
+    return false;
+  }
+  out_ << std::endl;
+
+  return true;
+}
+
 bool GeneratorImpl::EmitAliasType(const ast::type::AliasType* alias) {
   make_indent();
   out_ << "type " << alias->name() << " = ";

src/writer/wgsl/generator_impl.h

@@ -47,6 +47,14 @@ class GeneratorImpl : public TextGenerator {
   /// @returns true on successful generation; false otherwise
   bool Generate(const ast::Module& module);
 
+  /// Generates a single entry point
+  /// @param module the module to generate from
+  /// @param stage the pipeline stage
+  /// @param name the entry point name
+  bool GenerateEntryPoint(const ast::Module& module,
+                          ast::PipelineStage stage,
+                          const std::string& name);
+
   /// Handles generating an alias
   /// @param alias the alias to generate
   /// @returns true if the alias was emitted