[transform] Add first index offset transform

Adding new transform to workaround D3D's vertex/instance_index always
starting from 0

Bug: dawn:548
Change-Id: I048f39e76e236570f3ce337a77d804384ee659a4
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/34600
Commit-Queue: Enrico Galli <enrico.galli@intel.com>
Reviewed-by: dan sinclair <dsinclair@chromium.org>

BUILD.gn

@@ -415,6 +415,8 @@ source_set("libtint_core_src") {
     "src/transform/bound_array_accessors.h",
     "src/transform/emit_vertex_point_size.cc",
     "src/transform/emit_vertex_point_size.h",
+    "src/transform/first_index_offset.cc",
+    "src/transform/first_index_offset.h",
     "src/transform/manager.cc",
     "src/transform/manager.h",
     "src/transform/transform.cc",
@@ -821,6 +823,7 @@ source_set("tint_unittests_core_src") {
     "src/scope_stack_test.cc",
     "src/transform/bound_array_accessors_test.cc",
     "src/transform/emit_vertex_point_size_test.cc",
+    "src/transform/first_index_offset_test.cc",
     "src/transform/vertex_pulling_test.cc",
     "src/type_determiner_test.cc",
     "src/validator/validator_control_block_test.cc",

include/tint/tint.h

@@ -26,6 +26,7 @@
 #include "src/reader/reader.h"
 #include "src/transform/bound_array_accessors.h"
 #include "src/transform/emit_vertex_point_size.h"
+#include "src/transform/first_index_offset.h"
 #include "src/transform/manager.h"
 #include "src/transform/vertex_pulling.h"
 #include "src/type_determiner.h"

samples/main.cc

@@ -75,6 +75,7 @@ const char kUsage[] = R"(Usage: tint [options] <input-file>
                                Available transforms:
                                 bound_array_accessors
                                 emit_vertex_point_size
+                                first_index_offset
   --parse-only              -- Stop after parsing the input
   --dump-ast                -- Dump the generated AST to stdout
   --dawn-validation         -- SPIRV outputs are validated with the same flags
@@ -519,6 +520,9 @@ int main(int argc, const char** argv) {
     } else if (name == "emit_vertex_point_size") {
       transform_manager.append(
           std::make_unique<tint::transform::EmitVertexPointSize>());
+    } else if (name == "first_index_offset") {
+      transform_manager.append(
+          std::make_unique<tint::transform::FirstIndexOffset>(0, 0));
     } else {
       std::cerr << "Unknown transform name: " << name << std::endl;
       return 1;

src/CMakeLists.txt

@@ -236,6 +236,8 @@ set(TINT_LIB_SRCS
   transform/emit_vertex_point_size.h
   transform/bound_array_accessors.cc
   transform/bound_array_accessors.h
+  transform/first_index_offset.cc
+  transform/first_index_offset.h
   transform/manager.cc
   transform/manager.h
   transform/transform.cc
@@ -431,6 +433,7 @@ set(TINT_TEST_SRCS
   scope_stack_test.cc
   transform/emit_vertex_point_size_test.cc
   transform/bound_array_accessors_test.cc
+  transform/first_index_offset_test.cc
   transform/vertex_pulling_test.cc
   type_determiner_test.cc
   validator/validator_control_block_test.cc

src/ast/function.cc

@@ -74,6 +74,15 @@ void Function::add_referenced_module_variable(Variable* var) {
   referenced_module_vars_.push_back(var);
 }
 
+void Function::add_local_referenced_module_variable(Variable* var) {
+  for (const auto* v : local_referenced_module_vars_) {
+    if (v->name() == var->name()) {
+      return;
+    }
+  }
+  local_referenced_module_vars_.push_back(var);
+}
+
 const std::vector<std::pair<Variable*, LocationDecoration*>>
 Function::referenced_location_variables() const {
   std::vector<std::pair<Variable*, LocationDecoration*>> ret;
@@ -186,6 +195,23 @@ Function::referenced_multisampled_texture_variables() const {
   return ReferencedSampledTextureVariablesImpl(true);
 }
 
+const std::vector<std::pair<Variable*, BuiltinDecoration*>>
+Function::local_referenced_builtin_variables() const {
+  std::vector<std::pair<Variable*, BuiltinDecoration*>> ret;
+
+  for (auto* var : local_referenced_module_variables()) {
+    if (auto* decorated = var->As<DecoratedVariable>()) {
+      for (auto* deco : decorated->decorations()) {
+        if (auto* builtin = deco->As<BuiltinDecoration>()) {
+          ret.push_back({var, builtin});
+          break;
+        }
+      }
+    }
+  }
+  return ret;
+}
+
 void Function::add_ancestor_entry_point(const std::string& ep) {
   for (const auto& point : ancestor_entry_points_) {
     if (point == ep) {

src/ast/function.h

@@ -90,12 +90,20 @@ class Function : public Castable<Function, Node> {
   /// is not already included.
   /// @param var the module variable to add
   void add_referenced_module_variable(Variable* var);
+  /// Adds the given variable to the list of locally referenced module variables
+  /// if it is not already included.
+  /// @param var the module variable to add
+  void add_local_referenced_module_variable(Variable* var);
   /// Note: If this function calls other functions, the return will also include
   /// all of the referenced variables from the callees.
   /// @returns the referenced module variables
   const std::vector<Variable*>& referenced_module_variables() const {
     return referenced_module_vars_;
   }
+  /// @returns the locally referenced module variables
+  const std::vector<Variable*>& local_referenced_module_variables() const {
+    return local_referenced_module_vars_;
+  }
   /// Retrieves any referenced location variables
   /// @returns the <variable, decoration> pair.
   const std::vector<std::pair<Variable*, LocationDecoration*>>
@@ -135,6 +143,11 @@ class Function : public Castable<Function, Node> {
   const std::vector<std::pair<Variable*, Function::BindingInfo>>
   referenced_multisampled_texture_variables() const;
 
+  /// Retrieves any locally referenced builtin variables
+  /// @returns the <variable, decoration> pairs.
+  const std::vector<std::pair<Variable*, BuiltinDecoration*>>
+  local_referenced_builtin_variables() const;
+
   /// Adds an ancestor entry point
   /// @param ep the entry point ancestor
   void add_ancestor_entry_point(const std::string& ep);
@@ -189,6 +202,7 @@ class Function : public Castable<Function, Node> {
   type::Type* return_type_ = nullptr;
   BlockStatement* body_ = nullptr;
   std::vector<Variable*> referenced_module_vars_;
+  std::vector<Variable*> local_referenced_module_vars_;
   std::vector<std::string> ancestor_entry_points_;
   FunctionDecorationList decorations_;
 };

src/transform/first_index_offset.cc

@@ -0,0 +1,225 @@
+// Copyright 2020 The Tint Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "src/transform/first_index_offset.h"
+
+#include <cassert>
+#include <utility>
+
+#include "src/ast/array_accessor_expression.h"
+#include "src/ast/assignment_statement.h"
+#include "src/ast/binary_expression.h"
+#include "src/ast/bitcast_expression.h"
+#include "src/ast/break_statement.h"
+#include "src/ast/builtin_decoration.h"
+#include "src/ast/call_statement.h"
+#include "src/ast/case_statement.h"
+#include "src/ast/constructor_expression.h"
+#include "src/ast/decorated_variable.h"
+#include "src/ast/else_statement.h"
+#include "src/ast/expression.h"
+#include "src/ast/fallthrough_statement.h"
+#include "src/ast/identifier_expression.h"
+#include "src/ast/if_statement.h"
+#include "src/ast/loop_statement.h"
+#include "src/ast/member_accessor_expression.h"
+#include "src/ast/return_statement.h"
+#include "src/ast/scalar_constructor_expression.h"
+#include "src/ast/struct.h"
+#include "src/ast/struct_block_decoration.h"
+#include "src/ast/struct_decoration.h"
+#include "src/ast/struct_member.h"
+#include "src/ast/struct_member_offset_decoration.h"
+#include "src/ast/switch_statement.h"
+#include "src/ast/type/struct_type.h"
+#include "src/ast/type/u32_type.h"
+#include "src/ast/type_constructor_expression.h"
+#include "src/ast/unary_op_expression.h"
+#include "src/ast/variable.h"
+#include "src/ast/variable_decl_statement.h"
+#include "src/ast/variable_decoration.h"
+#include "src/type_determiner.h"
+
+namespace tint {
+namespace transform {
+namespace {
+
+constexpr char kStructName[] = "TintFirstIndexOffsetData";
+constexpr char kBufferName[] = "tint_first_index_data";
+constexpr char kFirstVertexName[] = "tint_first_vertex_index";
+constexpr char kFirstInstanceName[] = "tint_first_instance_index";
+constexpr char kIndexOffsetPrefix[] = "tint_first_index_offset_";
+
+}  // namespace
+
+FirstIndexOffset::FirstIndexOffset(uint32_t binding, uint32_t set)
+    : binding_(binding), set_(set) {}
+
+FirstIndexOffset::~FirstIndexOffset() = default;
+
+Transform::Output FirstIndexOffset::Run(ast::Module* in) {
+  Output out;
+  out.module = in->Clone();
+  auto* mod = &out.module;
+
+  // Running TypeDeterminer as we require local_referenced_builtin_variables()
+  // to be populated
+  TypeDeterminer td(mod);
+  if (!td.Determine()) {
+    diag::Diagnostic err;
+    err.severity = diag::Severity::Error;
+    err.message = td.error();
+    out.diagnostics.add(std::move(err));
+    return out;
+  }
+
+  std::string vertex_index_name;
+  std::string instance_index_name;
+
+  for (ast::Variable* var : mod->global_variables()) {
+    if (auto* dec_var = var->As<ast::DecoratedVariable>()) {
+      if (dec_var->name() == kBufferName) {
+        diag::Diagnostic err;
+        err.message = "First index offset transform has already been applied.";
+        err.severity = diag::Severity::Error;
+        out.diagnostics.add(std::move(err));
+        return out;
+      }
+
+      for (ast::VariableDecoration* dec : dec_var->decorations()) {
+        if (auto* blt_dec = dec->As<ast::BuiltinDecoration>()) {
+          ast::Builtin blt_type = blt_dec->value();
+          if (blt_type == ast::Builtin::kVertexIdx) {
+            vertex_index_name = var->name();
+            var->set_name(kIndexOffsetPrefix + var->name());
+            has_vertex_index_ = true;
+          } else if (blt_type == ast::Builtin::kInstanceIdx) {
+            instance_index_name = var->name();
+            var->set_name(kIndexOffsetPrefix + var->name());
+            has_instance_index_ = true;
+          }
+        }
+      }
+    }
+  }
+
+  if (!has_vertex_index_ && !has_instance_index_) {
+    return out;
+  }
+
+  ast::Variable* buffer_var = AddUniformBuffer(mod);
+
+  for (ast::Function* func : mod->functions()) {
+    for (const auto& data : func->local_referenced_builtin_variables()) {
+      if (data.second->value() == ast::Builtin::kVertexIdx) {
+        AddFirstIndexOffset(vertex_index_name, kFirstVertexName, buffer_var,
+                            func, mod);
+      } else if (data.second->value() == ast::Builtin::kInstanceIdx) {
+        AddFirstIndexOffset(instance_index_name, kFirstInstanceName, buffer_var,
+                            func, mod);
+      }
+    }
+  }
+
+  return out;
+}
+
+bool FirstIndexOffset::HasVertexIndex() {
+  return has_vertex_index_;
+}
+
+bool FirstIndexOffset::HasInstanceIndex() {
+  return has_instance_index_;
+}
+
+uint32_t FirstIndexOffset::GetFirstVertexOffset() {
+  assert(has_vertex_index_);
+  return vertex_index_offset_;
+}
+
+uint32_t FirstIndexOffset::GetFirstInstanceOffset() {
+  assert(has_instance_index_);
+  return instance_index_offset_;
+}
+
+ast::Variable* FirstIndexOffset::AddUniformBuffer(ast::Module* mod) {
+  auto* u32_type = mod->create<ast::type::U32>();
+  ast::StructMemberList members;
+  uint32_t offset = 0;
+  if (has_vertex_index_) {
+    ast::StructMemberDecorationList member_dec;
+    member_dec.push_back(
+        mod->create<ast::StructMemberOffsetDecoration>(offset, Source{}));
+    members.push_back(mod->create<ast::StructMember>(kFirstVertexName, u32_type,
+                                                     std::move(member_dec)));
+    vertex_index_offset_ = offset;
+    offset += 4;
+  }
+
+  if (has_instance_index_) {
+    ast::StructMemberDecorationList member_dec;
+    member_dec.push_back(
+        mod->create<ast::StructMemberOffsetDecoration>(offset, Source{}));
+    members.push_back(mod->create<ast::StructMember>(
+        kFirstInstanceName, u32_type, std::move(member_dec)));
+    instance_index_offset_ = offset;
+    offset += 4;
+  }
+
+  ast::StructDecorationList decos;
+  decos.push_back(mod->create<ast::StructBlockDecoration>(Source{}));
+
+  auto* struct_type = mod->create<ast::type::Struct>(
+      kStructName,
+      mod->create<ast::Struct>(std::move(decos), std::move(members)));
+
+  auto* idx_var =
+      mod->create<ast::DecoratedVariable>(mod->create<ast::Variable>(
+          Source{}, kBufferName, ast::StorageClass::kUniform, struct_type));
+
+  ast::VariableDecorationList decorations;
+  decorations.push_back(
+      mod->create<ast::BindingDecoration>(binding_, Source{}));
+  decorations.push_back(mod->create<ast::SetDecoration>(set_, Source{}));
+  idx_var->set_decorations(std::move(decorations));
+
+  mod->AddGlobalVariable(idx_var);
+
+  mod->AddConstructedType(struct_type);
+
+  return idx_var;
+}
+
+void FirstIndexOffset::AddFirstIndexOffset(const std::string& original_name,
+                                           const std::string& field_name,
+                                           ast::Variable* buffer_var,
+                                           ast::Function* func,
+                                           ast::Module* mod) {
+  auto* buffer = mod->create<ast::IdentifierExpression>(buffer_var->name());
+  auto* var = mod->create<ast::Variable>(Source{}, original_name,
+                                         ast::StorageClass::kNone,
+                                         mod->create<ast::type::U32>());
+
+  var->set_is_const(true);
+  var->set_constructor(mod->create<ast::BinaryExpression>(
+      ast::BinaryOp::kAdd,
+      mod->create<ast::IdentifierExpression>(kIndexOffsetPrefix + var->name()),
+      mod->create<ast::MemberAccessorExpression>(
+          buffer, mod->create<ast::IdentifierExpression>(field_name))));
+  func->body()->insert(0,
+                       mod->create<ast::VariableDeclStatement>(std::move(var)));
+}
+
+}  // namespace transform
+}  // namespace tint

src/transform/first_index_offset.h

@@ -0,0 +1,117 @@
+// Copyright 2020 The Tint Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef SRC_TRANSFORM_FIRST_INDEX_OFFSET_H_
+#define SRC_TRANSFORM_FIRST_INDEX_OFFSET_H_
+
+#include <string>
+
+#include "src/ast/module.h"
+#include "src/transform/transform.h"
+
+namespace tint {
+namespace transform {
+
+/// Adds firstVertex/Instance (injected via root constants) to
+/// vertex/instance_idx builtins.
+///
+/// This transform assumes that Name transform has been run before.
+///
+/// Unlike other APIs, D3D always starts vertex and instance numbering at 0,
+/// regardless of the firstVertex/Instance value specified. This transformer
+/// adds the value of firstVertex/Instance to each builtin. This action is
+/// performed by adding a new constant equal to original builtin +
+/// firstVertex/Instance to each function that references one of these builtins.
+///
+/// Note that D3D does not have any semantics for firstVertex/Instance.
+/// Therefore, these values must by passed to the shader.
+///
+/// Before:
+///   [[builtin(vertex_index)]] var<in> vert_idx : u32;
+///   fn func() -> u32 {
+///     return vert_idx;
+///   }
+///
+/// After:
+///   [[block]]
+///   struct TintFirstIndexOffsetData {
+///     [[offset(0)]] tint_first_vertex_index : u32;
+///     [[offset(4)]] tint_first_instance_index : u32;
+///   };
+///   [[builtin(vertex_index)]] var<in> tint_first_index_offset_vert_idx : u32;
+///   [[binding(N), set(M)]] var<uniform> tint_first_index_data :
+///                                                    TintFirstIndexOffsetData;
+///   fn func() -> u32 {
+///     const vert_idx = (tint_first_index_offset_vert_idx +
+///                       tint_first_index_data.tint_first_vertex_index);
+///     return vert_idx;
+///   }
+///
+class FirstIndexOffset : public Transform {
+ public:
+  /// Constructor
+  /// @param binding the binding() for firstVertex/Instance uniform
+  /// @param set the set() for firstVertex/Instance uniform
+  FirstIndexOffset(uint32_t binding, uint32_t set);
+  ~FirstIndexOffset() override;
+
+  /// Runs the transform on `module`, returning the transformation result.
+  /// @note Users of Tint should register the transform with transform manager
+  /// and invoke its Run(), instead of directly calling the transform's Run().
+  /// Calling Run() directly does not perform module state cleanup operations.
+  /// @param module the source module to transform
+  /// @returns the transformation result
+  Output Run(ast::Module* module) override;
+
+  /// @returns whether shader uses vertex_index
+  bool HasVertexIndex();
+
+  /// @returns whether shader uses instance_index
+  bool HasInstanceIndex();
+
+  /// @returns offset of firstVertex into constant buffer
+  uint32_t GetFirstVertexOffset();
+
+  /// @returns offset of firstInstance into constant buffer
+  uint32_t GetFirstInstanceOffset();
+
+ private:
+  /// Adds uniform buffer with firstVertex/Instance to module
+  /// @returns variable of new uniform buffer
+  ast::Variable* AddUniformBuffer(ast::Module* mod);
+  /// Adds constant with modified original_name builtin to func
+  /// @param original_name the name of the original builtin used in function
+  /// @param field_name name of field in firstVertex/Instance buffer
+  /// @param buffer_var variable of firstVertex/Instance buffer
+  /// @param func function to modify
+  /// @returns true if it was able to add constant to function
+  void AddFirstIndexOffset(const std::string& original_name,
+                           const std::string& field_name,
+                           ast::Variable* buffer_var,
+                           ast::Function* func,
+                           ast::Module* module);
+
+  uint32_t binding_;
+  uint32_t set_;
+
+  bool has_vertex_index_ = false;
+  bool has_instance_index_ = false;
+
+  uint32_t vertex_index_offset_ = 0;
+  uint32_t instance_index_offset_ = 0;
+};
+}  // namespace transform
+}  // namespace tint
+
+#endif  // SRC_TRANSFORM_FIRST_INDEX_OFFSET_H_

src/transform/first_index_offset_test.cc

@@ -0,0 +1,392 @@
+// Copyright 2020 The Tint Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "src/transform/first_index_offset.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "gtest/gtest.h"
+#include "src/ast/block_statement.h"
+#include "src/ast/builder.h"
+#include "src/ast/builtin.h"
+#include "src/ast/builtin_decoration.h"
+#include "src/ast/call_expression.h"
+#include "src/ast/call_statement.h"
+#include "src/ast/decorated_variable.h"
+#include "src/ast/function.h"
+#include "src/ast/identifier_expression.h"
+#include "src/ast/module.h"
+#include "src/ast/return_statement.h"
+#include "src/ast/storage_class.h"
+#include "src/ast/type/u32_type.h"
+#include "src/ast/variable.h"
+#include "src/ast/variable_decoration.h"
+#include "src/diagnostic/formatter.h"
+#include "src/source.h"
+#include "src/transform/manager.h"
+
+namespace tint {
+namespace transform {
+namespace {
+
+class FirstIndexOffsetTest : public testing::Test {};
+
+struct ModuleBuilder : public ast::BuilderWithModule {
+  ast::Module Module() {
+    Build();
+    return std::move(*mod);
+  }
+
+ protected:
+  void AddBuiltinInput(const std::string& name, ast::Builtin builtin) {
+    auto* var = Var(name, ast::StorageClass::kInput, ty.u32);
+    auto* dec_var = create<ast::DecoratedVariable>(var);
+    ast::VariableDecorationList decs;
+    decs.push_back(create<ast::BuiltinDecoration>(builtin, Source{}));
+    dec_var->set_decorations(std::move(decs));
+    mod->AddGlobalVariable(dec_var);
+  }
+
+  ast::Function* AddFunction(const std::string& name,
+                             ast::VariableList params = {}) {
+    auto* func = create<ast::Function>(Source{}, name, std::move(params),
+                                       ty.u32, create<ast::BlockStatement>(),
+                                       ast::FunctionDecorationList());
+    mod->AddFunction(func);
+    return func;
+  }
+
+  virtual void Build() = 0;
+};
+
+TEST_F(FirstIndexOffsetTest, Error_AlreadyTransformed) {
+  struct Builder : public ModuleBuilder {
+    void Build() override {
+      AddBuiltinInput("vert_idx", ast::Builtin::kVertexIdx);
+    }
+  };
+
+  Manager manager;
+  manager.append(std::make_unique<FirstIndexOffset>(0, 0));
+  manager.append(std::make_unique<FirstIndexOffset>(1, 1));
+
+  auto module = Builder{}.Module();
+  auto result = manager.Run(&module);
+
+  // Release the source module to ensure there's no uncloned data in result
+  { auto tmp = std::move(module); }
+
+  ASSERT_EQ(diag::Formatter().format(result.diagnostics),
+            "error: First index offset transform has already been applied.");
+}
+
+TEST_F(FirstIndexOffsetTest, EmptyModule) {
+  Manager manager;
+  manager.append(std::make_unique<FirstIndexOffset>(0, 0));
+
+  ast::Module module;
+  auto result = manager.Run(&module);
+
+  // Release the source module to ensure there's no uncloned data in result
+  { auto tmp = std::move(module); }
+
+  ASSERT_FALSE(result.diagnostics.contains_errors())
+      << diag::Formatter().format(result.diagnostics);
+
+  EXPECT_EQ("Module{\n}\n", result.module.to_str());
+}
+
+TEST_F(FirstIndexOffsetTest, BasicModuleVertexIndex) {
+  struct Builder : public ModuleBuilder {
+    void Build() override {
+      AddBuiltinInput("vert_idx", ast::Builtin::kVertexIdx);
+      ast::Function* func = AddFunction("test");
+      func->body()->append(create<ast::ReturnStatement>(
+          Source{}, create<ast::IdentifierExpression>("vert_idx")));
+    }
+  };
+
+  Manager manager;
+  manager.append(std::make_unique<FirstIndexOffset>(1, 2));
+
+  auto module = Builder{}.Module();
+  auto result = manager.Run(&module);
+
+  // Release the source module to ensure there's no uncloned data in result
+  { auto tmp = std::move(module); }
+
+  ASSERT_FALSE(result.diagnostics.contains_errors())
+      << diag::Formatter().format(result.diagnostics);
+
+  EXPECT_EQ(R"(Module{
+  TintFirstIndexOffsetData Struct{
+    [[block]]
+    StructMember{[[ offset 0 ]] tint_first_vertex_index: __u32}
+  }
+  DecoratedVariable{
+    Decorations{
+      BuiltinDecoration{vertex_idx}
+    }
+    tint_first_index_offset_vert_idx
+    in
+    __u32
+  }
+  DecoratedVariable{
+    Decorations{
+      BindingDecoration{1}
+      SetDecoration{2}
+    }
+    tint_first_index_data
+    uniform
+    __struct_TintFirstIndexOffsetData
+  }
+  Function test -> __u32
+  ()
+  {
+    VariableDeclStatement{
+      VariableConst{
+        vert_idx
+        none
+        __u32
+        {
+          Binary[__u32]{
+            Identifier[__ptr_in__u32]{tint_first_index_offset_vert_idx}
+            add
+            MemberAccessor[__ptr_uniform__u32]{
+              Identifier[__ptr_uniform__struct_TintFirstIndexOffsetData]{tint_first_index_data}
+              Identifier[not set]{tint_first_vertex_index}
+            }
+          }
+        }
+      }
+    }
+    Return{
+      {
+        Identifier[__u32]{vert_idx}
+      }
+    }
+  }
+}
+)",
+            result.module.to_str());
+}
+
+TEST_F(FirstIndexOffsetTest, BasicModuleInstanceIndex) {
+  struct Builder : public ModuleBuilder {
+    void Build() override {
+      AddBuiltinInput("inst_idx", ast::Builtin::kInstanceIdx);
+    }
+  };
+
+  Manager manager;
+  manager.append(std::make_unique<FirstIndexOffset>(1, 7));
+
+  auto module = Builder{}.Module();
+  auto result = manager.Run(&module);
+
+  // Release the source module to ensure there's no uncloned data in result
+  { auto tmp = std::move(module); }
+
+  ASSERT_FALSE(result.diagnostics.contains_errors())
+      << diag::Formatter().format(result.diagnostics);
+  EXPECT_EQ(R"(Module{
+  TintFirstIndexOffsetData Struct{
+    [[block]]
+    StructMember{[[ offset 0 ]] tint_first_instance_index: __u32}
+  }
+  DecoratedVariable{
+    Decorations{
+      BuiltinDecoration{instance_idx}
+    }
+    tint_first_index_offset_inst_idx
+    in
+    __u32
+  }
+  DecoratedVariable{
+    Decorations{
+      BindingDecoration{1}
+      SetDecoration{7}
+    }
+    tint_first_index_data
+    uniform
+    __struct_TintFirstIndexOffsetData
+  }
+}
+)",
+            result.module.to_str());
+}
+
+TEST_F(FirstIndexOffsetTest, BasicModuleBothIndex) {
+  struct Builder : public ModuleBuilder {
+    void Build() override {
+      AddBuiltinInput("inst_idx", ast::Builtin::kInstanceIdx);
+      AddBuiltinInput("vert_idx", ast::Builtin::kVertexIdx);
+    }
+  };
+
+  auto transform = std::make_unique<FirstIndexOffset>(1, 7);
+  auto* transform_ptr = transform.get();
+
+  Manager manager;
+  manager.append(std::move(transform));
+
+  auto module = Builder{}.Module();
+  auto result = manager.Run(&module);
+
+  // Release the source module to ensure there's no uncloned data in result
+  { auto tmp = std::move(module); }
+
+  ASSERT_FALSE(result.diagnostics.contains_errors())
+      << diag::Formatter().format(result.diagnostics);
+  EXPECT_EQ(R"(Module{
+  TintFirstIndexOffsetData Struct{
+    [[block]]
+    StructMember{[[ offset 0 ]] tint_first_vertex_index: __u32}
+    StructMember{[[ offset 4 ]] tint_first_instance_index: __u32}
+  }
+  DecoratedVariable{
+    Decorations{
+      BuiltinDecoration{instance_idx}
+    }
+    tint_first_index_offset_inst_idx
+    in
+    __u32
+  }
+  DecoratedVariable{
+    Decorations{
+      BuiltinDecoration{vertex_idx}
+    }
+    tint_first_index_offset_vert_idx
+    in
+    __u32
+  }
+  DecoratedVariable{
+    Decorations{
+      BindingDecoration{1}
+      SetDecoration{7}
+    }
+    tint_first_index_data
+    uniform
+    __struct_TintFirstIndexOffsetData
+  }
+}
+)",
+            result.module.to_str());
+
+  EXPECT_TRUE(transform_ptr->HasVertexIndex());
+  EXPECT_EQ(transform_ptr->GetFirstVertexOffset(), 0u);
+
+  EXPECT_TRUE(transform_ptr->HasInstanceIndex());
+  EXPECT_EQ(transform_ptr->GetFirstInstanceOffset(), 4u);
+}
+
+TEST_F(FirstIndexOffsetTest, NestedCalls) {
+  struct Builder : public ModuleBuilder {
+    void Build() override {
+      AddBuiltinInput("vert_idx", ast::Builtin::kVertexIdx);
+      ast::Function* func1 = AddFunction("func1");
+      func1->body()->append(create<ast::ReturnStatement>(
+          Source{}, create<ast::IdentifierExpression>("vert_idx")));
+      ast::Function* func2 = AddFunction("func2");
+      func2->body()->append(create<ast::ReturnStatement>(
+          Source{}, create<ast::CallExpression>(
+                        create<ast::IdentifierExpression>("func1"),
+                        ast::ExpressionList{})));
+    }
+  };
+
+  auto transform = std::make_unique<FirstIndexOffset>(2, 2);
+
+  Manager manager;
+  manager.append(std::move(transform));
+
+  auto module = Builder{}.Module();
+  auto result = manager.Run(&module);
+
+  // Release the source module to ensure there's no uncloned data in result
+  { auto tmp = std::move(module); }
+
+  ASSERT_FALSE(result.diagnostics.contains_errors())
+      << diag::Formatter().format(result.diagnostics);
+  EXPECT_EQ(R"(Module{
+  TintFirstIndexOffsetData Struct{
+    [[block]]
+    StructMember{[[ offset 0 ]] tint_first_vertex_index: __u32}
+  }
+  DecoratedVariable{
+    Decorations{
+      BuiltinDecoration{vertex_idx}
+    }
+    tint_first_index_offset_vert_idx
+    in
+    __u32
+  }
+  DecoratedVariable{
+    Decorations{
+      BindingDecoration{2}
+      SetDecoration{2}
+    }
+    tint_first_index_data
+    uniform
+    __struct_TintFirstIndexOffsetData
+  }
+  Function func1 -> __u32
+  ()
+  {
+    VariableDeclStatement{
+      VariableConst{
+        vert_idx
+        none
+        __u32
+        {
+          Binary[__u32]{
+            Identifier[__ptr_in__u32]{tint_first_index_offset_vert_idx}
+            add
+            MemberAccessor[__ptr_uniform__u32]{
+              Identifier[__ptr_uniform__struct_TintFirstIndexOffsetData]{tint_first_index_data}
+              Identifier[not set]{tint_first_vertex_index}
+            }
+          }
+        }
+      }
+    }
+    Return{
+      {
+        Identifier[__u32]{vert_idx}
+      }
+    }
+  }
+  Function func2 -> __u32
+  ()
+  {
+    Return{
+      {
+        Call[__u32]{
+          Identifier[__u32]{func1}
+          (
+          )
+        }
+      }
+    }
+  }
+}
+)",
+            result.module.to_str());
+}
+
+}  // namespace
+}  // namespace transform
+}  // namespace tint

src/type_determiner.cc

@@ -71,8 +71,8 @@ void TypeDeterminer::set_error(const Source& src, const std::string& msg) {
   error_ += msg;
 }
 
-void TypeDeterminer::set_referenced_from_function_if_needed(
+void TypeDeterminer::set_referenced_from_function_if_needed(ast::Variable* var,
-    ast::Variable* var) {
+                                                            bool local) {
   if (current_function_ == nullptr) {
     return;
   }
@@ -82,6 +82,9 @@ void TypeDeterminer::set_referenced_from_function_if_needed(
   }
 
   current_function_->add_referenced_module_variable(var);
+  if (local) {
+    current_function_->add_local_referenced_module_variable(var);
+  }
 }
 
 bool TypeDeterminer::Determine() {
@@ -394,7 +397,7 @@ bool TypeDeterminer::DetermineCall(ast::CallExpression* expr) {
 
         // We inherit any referenced variables from the callee.
         for (auto* var : callee_func->referenced_module_variables()) {
-          set_referenced_from_function_if_needed(var);
+          set_referenced_from_function_if_needed(var, false);
         }
       }
 
@@ -849,7 +852,7 @@ bool TypeDeterminer::DetermineIdentifier(ast::IdentifierExpression* expr) {
           mod_->create<ast::type::Pointer>(var->type(), var->storage_class()));
     }
 
-    set_referenced_from_function_if_needed(var);
+    set_referenced_from_function_if_needed(var, true);
     return true;
   }
 

src/type_determiner.h

@@ -112,7 +112,7 @@ class TypeDeterminer {
 
  private:
   void set_error(const Source& src, const std::string& msg);
-  void set_referenced_from_function_if_needed(ast::Variable* var);
+  void set_referenced_from_function_if_needed(ast::Variable* var, bool local);
   void set_entry_points(const std::string& fn_name, const std::string& ep_name);
 
   bool DetermineArrayAccessor(ast::ArrayAccessorExpression* expr);