spirv-reader: support sample_id builtin variable

TODO: support sample_id declared with signed integer store
type, and then having the pointer passed to a helper function.

Bug: tint:471
Change-Id: Iac303ff6118b2d2d518e5070a8d589dcd3616f39
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/40020
Commit-Queue: Ben Clayton <bclayton@google.com>
Reviewed-by: Ben Clayton <bclayton@google.com>
Auto-Submit: David Neto <dneto@google.com>

src/reader/spirv/enum_converter.cc

@@ -86,6 +86,8 @@ ast::Builtin EnumConverter::ToBuiltin(SpvBuiltIn b) {
       return ast::Builtin::kLocalInvocationIndex;
     case SpvBuiltInGlobalInvocationId:
       return ast::Builtin::kGlobalInvocationId;
+    case SpvBuiltInSampleId:
+      return ast::Builtin::kSampleId;
     default:
       break;
   }

src/reader/spirv/function.cc

@@ -991,7 +991,7 @@ bool FunctionEmitter::EmitBody() {
     return false;
   }
 
-  if (!RegisterIgnoredBuiltInVariables()) {
+  if (!RegisterSpecialBuiltInVariables()) {
     return false;
   }
   if (!RegisterLocallyDefinedValues()) {
@@ -2023,6 +2023,10 @@ TypedExpression FunctionEmitter::MakeExpression(uint32_t id) {
       Fail() << "unhandled use of a pointer to the PointSize builtin, with ID: "
              << id;
       return {};
+    case SkipReason::kSampleIdBuiltinPointer:
+      Fail() << "unhandled use of a pointer to the SampleId builtin, with ID: "
+             << id;
+      return {};
   }
   if (identifier_values_.count(id) || parser_impl_.IsScalarSpecConstant(id)) {
     auto name = namer_.Name(id);
@@ -2999,9 +3003,24 @@ bool FunctionEmitter::EmitStatement(const spvtools::opt::Instruction& inst) {
       // Memory accesses must be issued in SPIR-V program order.
       // So represent a load by a new const definition.
       const auto ptr_id = inst.GetSingleWordInOperand(0);
-      if (GetSkipReason(ptr_id) == SkipReason::kPointSizeBuiltinPointer) {
+      switch (GetSkipReason(ptr_id)) {
-        GetDefInfo(inst.result_id())->skip = SkipReason::kPointSizeBuiltinValue;
+        case SkipReason::kPointSizeBuiltinPointer:
-        return true;
+          GetDefInfo(inst.result_id())->skip =
+              SkipReason::kPointSizeBuiltinValue;
+          return true;
+        case SkipReason::kSampleIdBuiltinPointer: {
+          // The SPIR-V variable is i32, but WGSL requires u32.
+          auto var_id = parser_impl_.IdForSpecialBuiltIn(SpvBuiltInSampleId);
+          auto name = namer_.Name(var_id);
+          ast::Expression* id_expr = create<ast::IdentifierExpression>(
+              Source{}, builder_.Symbols().Register(name));
+          auto expr = TypedExpression{
+              i32_, create<ast::TypeConstructorExpression>(
+                        Source{}, i32_, ast::ExpressionList{id_expr})};
+          return EmitConstDefinition(inst, expr);
+        }
+        default:
+          break;
       }
       auto expr = MakeExpression(ptr_id);
       // The load result type is the pointee type of its operand.
@@ -3626,11 +3645,11 @@ TypedExpression FunctionEmitter::MakeVectorShuffle(
           create<ast::TypeConstructorExpression>(source, result_type, values)};
 }
 
-bool FunctionEmitter::RegisterIgnoredBuiltInVariables() {
+bool FunctionEmitter::RegisterSpecialBuiltInVariables() {
   size_t index = def_info_.size();
-  for (auto& ignored_var : parser_impl_.ignored_builtins()) {
+  for (auto& special_var : parser_impl_.special_builtins()) {
-    const auto id = ignored_var.first;
+    const auto id = special_var.first;
-    const auto builtin = ignored_var.second;
+    const auto builtin = special_var.second;
     def_info_[id] =
         std::make_unique<DefInfo>(*(def_use_mgr_->GetDef(id)), 0, index);
     ++index;
@@ -3639,8 +3658,11 @@ bool FunctionEmitter::RegisterIgnoredBuiltInVariables() {
       case SpvBuiltInPointSize:
         def->skip = SkipReason::kPointSizeBuiltinPointer;
         break;
+      case SpvBuiltInSampleId:
+        def->skip = SkipReason::kSampleIdBuiltinPointer;
+        break;
       default:
-        return Fail() << "unrecognized ignored builtin: " << int(builtin);
+        return Fail() << "unrecognized special builtin: " << int(builtin);
     }
   }
   return true;

src/reader/spirv/function.h

@@ -223,6 +223,10 @@ enum class SkipReason {
   /// PointSize builtin. Use 1.0f instead, because that's the only value
   /// supported by WebGPU.
   kPointSizeBuiltinValue,
+
+  /// `kSampleIdBuiltinPointer`: the value is a pointer to the SampleId builtin
+  /// variable.  Don't generate its address.
+  kSampleIdBuiltinPointer,
 };
 
 /// Bookkeeping info for a SPIR-V ID defined in the function, or some
@@ -328,6 +332,9 @@ inline std::ostream& operator<<(std::ostream& o, const DefInfo& di) {
     case SkipReason::kPointSizeBuiltinValue:
       o << " skip:pointsize_value";
       break;
+    case SkipReason::kSampleIdBuiltinPointer:
+      o << " skip:sampleid_pointer";
+      break;
   }
   o << "}";
   return o;
@@ -467,10 +474,11 @@ class FunctionEmitter {
   bool FindIfSelectionInternalHeaders();
 
   /// Creates a DefInfo record for each module-scope builtin variable
-  /// that should be ignored.
+  /// that should be handled specially.  Either it's ignored, or its store
+  /// type is converted on load.
   /// Populates the `def_info_` mapping for such IDs.
   /// @returns false on failure
-  bool RegisterIgnoredBuiltInVariables();
+  bool RegisterSpecialBuiltInVariables();
 
   /// Creates a DefInfo record for each locally defined SPIR-V ID.
   /// Populates the `def_info_` mapping with basic results for such IDs.

src/reader/spirv/parser_impl.cc

@@ -1265,8 +1265,18 @@ ast::Variable* ParserImpl::MakeVariable(
       const auto spv_builtin = static_cast<SpvBuiltIn>(deco[1]);
       switch (spv_builtin) {
         case SpvBuiltInPointSize:
-          ignored_builtins_[id] = spv_builtin;
+          special_builtins_[id] = spv_builtin;
           return nullptr;
+        case SpvBuiltInSampleId:
+          // The SPIR-V variable might is likely to be signed (because GLSL
+          // requires signed), but WGSL requires unsigned.  Handle specially
+          // so we always perform the conversion at load and store.
+          if (auto* forced_type = unsigned_type_for_[type]) {
+            // Requires conversion and special handling in code generation.
+            special_builtins_[id] = spv_builtin;
+            type = forced_type;
+          }
+          break;
         default:
           break;
       }

src/reader/spirv/parser_impl.h

@@ -477,9 +477,22 @@ class ParserImpl : Reader {
   /// A map of SPIR-V identifiers to builtins
   using BuiltInsMap = std::unordered_map<uint32_t, SpvBuiltIn>;
 
-  /// @returns a map of builtins that should be ignored as they do not exist in
+  /// @returns a map of builtins that should be handled specially by code
-  /// WGSL.
+  /// generation. Either the builtin does not exist in WGSL, or a type
-  const BuiltInsMap& ignored_builtins() const { return ignored_builtins_; }
+  /// conversion must be implemented on load and store.
+  const BuiltInsMap& special_builtins() const { return special_builtins_; }
+
+  /// @param builtin the SPIR-V builtin variable kind
+  /// @returns the SPIR-V ID for the variable defining the given builtin, or 0
+  uint32_t IdForSpecialBuiltIn(SpvBuiltIn builtin) const {
+    // Do a linear search.
+    for (const auto& entry : special_builtins_) {
+      if (entry.second == builtin) {
+        return entry.first;
+      }
+    }
+    return 0;
+  }
 
  private:
   /// Converts a specific SPIR-V type to a Tint type. Integer case
@@ -634,8 +647,10 @@ class ParserImpl : Reader {
       handle_type_;
 
   /// Maps the SPIR-V ID of a module-scope builtin variable that should be
-  /// ignored, to its builtin kind.
+  /// ignored or type-converted, to its builtin kind.
-  BuiltInsMap ignored_builtins_;
+  /// See also BuiltInPositionInfo which is a separate mechanism for a more
+  /// complex case of replacing an entire structure.
+  BuiltInsMap special_builtins_;
 };
 
 }  // namespace spirv

src/reader/spirv/parser_impl_module_var_test.cc

@@ -2021,6 +2021,388 @@ TEST_F(SpvModuleScopeVarParserTest, ScalarSpecConstant_UsedInFunction) {
       << ToString(program, fe.ast_body());
 }
 
+// Returns the start of a shader for testing SampleId,
+// parameterized by store type of %int or %uint
+std::string SampleIdPreamble(std::string store_type) {
+  return R"(
+    OpCapability Shader
+    OpCapability SampleRateShading
+    OpMemoryModel Logical Simple
+    OpEntryPoint Fragment %main "main" %1
+    OpExecutionMode %main OriginUpperLeft
+    OpDecorate %1 BuiltIn SampleId
+    %void = OpTypeVoid
+    %voidfn = OpTypeFunction %void
+    %float = OpTypeFloat 32
+    %uint = OpTypeInt 32 0
+    %int = OpTypeInt 32 1
+    %ptr_ty = OpTypePointer Input )" +
+         store_type + R"(
+    %1 = OpVariable %ptr_ty Input
+)";
+}
+
+TEST_F(SpvModuleScopeVarParserTest, SampleId_I32_Load_Direct) {
+  const std::string assembly = SampleIdPreamble("%int") + R"(
+    %main = OpFunction %void None %voidfn
+    %entry = OpLabel
+    %2 = OpLoad %int %1
+    OpReturn
+    OpFunctionEnd
+ )";
+  auto p = parser(test::Assemble(assembly));
+  ASSERT_TRUE(p->BuildAndParseInternalModule()) << p->error() << assembly;
+  EXPECT_TRUE(p->error().empty());
+  const auto module_str = p->program().to_str();
+  // Correct declaration
+  EXPECT_THAT(module_str, HasSubstr(R"(
+  Variable{
+    Decorations{
+      BuiltinDecoration{sample_id}
+    }
+    x_1
+    in
+    __u32
+  })"));
+
+  // Correct body
+  EXPECT_THAT(module_str, HasSubstr(R"(
+    VariableDeclStatement{
+      VariableConst{
+        x_2
+        none
+        __i32
+        {
+          TypeConstructor[not set]{
+            __i32
+            Identifier[not set]{x_1}
+          }
+        }
+      }
+    })"))
+      << module_str;
+}
+
+TEST_F(SpvModuleScopeVarParserTest, SampleId_I32_Load_CopyObject) {
+  const std::string assembly = SampleIdPreamble("%int") + R"(
+    %main = OpFunction %void None %voidfn
+    %entry = OpLabel
+    %copy_ptr = OpCopyObject %ptr_ty %1
+    %2 = OpLoad %int %copy_ptr
+    OpReturn
+    OpFunctionEnd
+ )";
+  auto p = parser(test::Assemble(assembly));
+  ASSERT_TRUE(p->BuildAndParseInternalModule()) << p->error() << assembly;
+  EXPECT_TRUE(p->error().empty());
+  const auto module_str = p->program().to_str();
+  // Correct declaration
+  EXPECT_THAT(module_str, HasSubstr(R"(
+  Variable{
+    Decorations{
+      BuiltinDecoration{sample_id}
+    }
+    x_1
+    in
+    __u32
+  })"));
+
+  // Correct body
+  EXPECT_THAT(module_str, HasSubstr(R"(
+    VariableDeclStatement{
+      VariableConst{
+        x_2
+        none
+        __i32
+        {
+          TypeConstructor[not set]{
+            __i32
+            Identifier[not set]{x_1}
+          }
+        }
+      }
+    })"))
+      << module_str;
+}
+
+TEST_F(SpvModuleScopeVarParserTest, SampleId_I32_Load_AccessChain) {
+  const std::string assembly = SampleIdPreamble("%int") + R"(
+    %main = OpFunction %void None %voidfn
+    %entry = OpLabel
+    %copy_ptr = OpAccessChain %ptr_ty %1
+    %2 = OpLoad %int %copy_ptr
+    OpReturn
+    OpFunctionEnd
+ )";
+  auto p = parser(test::Assemble(assembly));
+  ASSERT_TRUE(p->BuildAndParseInternalModule()) << p->error() << assembly;
+  EXPECT_TRUE(p->error().empty());
+  const auto module_str = p->program().to_str();
+  // Correct declaration
+  EXPECT_THAT(module_str, HasSubstr(R"(
+  Variable{
+    Decorations{
+      BuiltinDecoration{sample_id}
+    }
+    x_1
+    in
+    __u32
+  })"));
+
+  // Correct body
+  EXPECT_THAT(module_str, HasSubstr(R"(
+    VariableDeclStatement{
+      VariableConst{
+        x_2
+        none
+        __i32
+        {
+          TypeConstructor[not set]{
+            __i32
+            Identifier[not set]{x_1}
+          }
+        }
+      }
+    })"))
+      << module_str;
+}
+
+TEST_F(SpvModuleScopeVarParserTest, SampleId_I32_FunctParam) {
+  const std::string assembly = SampleIdPreamble("%int") + R"(
+    %helper_ty = OpTypeFunction %int %ptr_ty
+    %helper = OpFunction %int None %helper_ty
+    %param = OpFunctionParameter %ptr_ty
+    %helper_entry = OpLabel
+    %3 = OpLoad %int %param
+    OpReturnValue %3
+    OpFunctionEnd
+
+    %main = OpFunction %void None %voidfn
+    %entry = OpLabel
+    %result = OpFunctionCall %int %helper %1
+    OpReturn
+    OpFunctionEnd
+ )";
+  auto p = parser(test::Assemble(assembly));
+  // TODO(dneto): We can handle this if we make a shadow variable and mutate
+  // the parameter type.
+  ASSERT_FALSE(p->BuildAndParseInternalModule());
+  EXPECT_THAT(
+      p->error(),
+      HasSubstr(
+          "unhandled use of a pointer to the SampleId builtin, with ID: 1"));
+}
+
+TEST_F(SpvModuleScopeVarParserTest, SampleId_U32_Load_Direct) {
+  const std::string assembly = SampleIdPreamble("%uint") + R"(
+    %main = OpFunction %void None %voidfn
+    %entry = OpLabel
+    %2 = OpLoad %uint %1
+    OpReturn
+    OpFunctionEnd
+ )";
+  auto p = parser(test::Assemble(assembly));
+  ASSERT_TRUE(p->BuildAndParseInternalModule()) << p->error() << assembly;
+  EXPECT_TRUE(p->error().empty());
+  const auto module_str = p->program().to_str();
+  // Correct declaration
+  EXPECT_THAT(module_str, HasSubstr(R"(
+  Variable{
+    Decorations{
+      BuiltinDecoration{sample_id}
+    }
+    x_1
+    in
+    __u32
+  })"));
+
+  // Correct body
+  EXPECT_THAT(module_str, HasSubstr(R"(
+    VariableDeclStatement{
+      VariableConst{
+        x_2
+        none
+        __u32
+        {
+          Identifier[not set]{x_1}
+        }
+      }
+    })"))
+      << module_str;
+}
+
+TEST_F(SpvModuleScopeVarParserTest, SampleId_U32_Load_CopyObject) {
+  const std::string assembly = SampleIdPreamble("%uint") + R"(
+    %main = OpFunction %void None %voidfn
+    %entry = OpLabel
+    %copy_ptr = OpCopyObject %ptr_ty %1
+    %2 = OpLoad %uint %copy_ptr
+    OpReturn
+    OpFunctionEnd
+ )";
+  auto p = parser(test::Assemble(assembly));
+  ASSERT_TRUE(p->BuildAndParseInternalModule()) << p->error() << assembly;
+  EXPECT_TRUE(p->error().empty());
+  const auto module_str = p->program().to_str();
+  // Correct declaration
+  EXPECT_THAT(module_str, HasSubstr(R"(
+  Variable{
+    Decorations{
+      BuiltinDecoration{sample_id}
+    }
+    x_1
+    in
+    __u32
+  })"));
+
+  // Correct body
+  EXPECT_THAT(module_str, HasSubstr(R"(
+    VariableDeclStatement{
+      VariableConst{
+        x_11
+        none
+        __ptr_in__u32
+        {
+          Identifier[not set]{x_1}
+        }
+      }
+    }
+    VariableDeclStatement{
+      VariableConst{
+        x_2
+        none
+        __u32
+        {
+          Identifier[not set]{x_11}
+        }
+      }
+    })"))
+      << module_str;
+}
+
+TEST_F(SpvModuleScopeVarParserTest, SampleId_U32_Load_AccessChain) {
+  const std::string assembly = SampleIdPreamble("%uint") + R"(
+    %main = OpFunction %void None %voidfn
+    %entry = OpLabel
+    %copy_ptr = OpAccessChain %ptr_ty %1
+    %2 = OpLoad %uint %copy_ptr
+    OpReturn
+    OpFunctionEnd
+ )";
+  auto p = parser(test::Assemble(assembly));
+  ASSERT_TRUE(p->BuildAndParseInternalModule()) << p->error() << assembly;
+  EXPECT_TRUE(p->error().empty());
+  const auto module_str = p->program().to_str();
+  // Correct declaration
+  EXPECT_THAT(module_str, HasSubstr(R"(
+  Variable{
+    Decorations{
+      BuiltinDecoration{sample_id}
+    }
+    x_1
+    in
+    __u32
+  })"));
+
+  // Correct body
+  EXPECT_THAT(module_str, HasSubstr(R"(
+    VariableDeclStatement{
+      VariableConst{
+        x_2
+        none
+        __u32
+        {
+          Identifier[not set]{x_1}
+        }
+      }
+    })"))
+      << module_str;
+}
+
+TEST_F(SpvModuleScopeVarParserTest, SampleId_U32_FunctParam) {
+  const std::string assembly = SampleIdPreamble("%uint") + R"(
+    %helper_ty = OpTypeFunction %uint %ptr_ty
+    %helper = OpFunction %uint None %helper_ty
+    %param = OpFunctionParameter %ptr_ty
+    %helper_entry = OpLabel
+    %3 = OpLoad %uint %param
+    OpReturnValue %3
+    OpFunctionEnd
+
+    %main = OpFunction %void None %voidfn
+    %entry = OpLabel
+    %result = OpFunctionCall %uint %helper %1
+    OpReturn
+    OpFunctionEnd
+ )";
+  auto p = parser(test::Assemble(assembly));
+  // TODO(dneto): We can handle this if we make a shadow variable and mutate
+  // the parameter type.
+  ASSERT_TRUE(p->BuildAndParseInternalModule()) << p->error() << assembly;
+  EXPECT_TRUE(p->error().empty());
+  const auto module_str = p->program().to_str();
+  // Correct declaration
+  EXPECT_THAT(module_str, HasSubstr(R"(
+  Variable{
+    Decorations{
+      BuiltinDecoration{sample_id}
+    }
+    x_1
+    in
+    __u32
+  })"));
+
+  // Correct bodies
+  EXPECT_THAT(module_str, HasSubstr(R"(
+  Function x_11 -> __u32
+  (
+    VariableConst{
+      x_12
+      none
+      __ptr_in__u32
+    }
+  )
+  {
+    VariableDeclStatement{
+      VariableConst{
+        x_3
+        none
+        __u32
+        {
+          Identifier[not set]{x_12}
+        }
+      }
+    }
+    Return{
+      {
+        Identifier[not set]{x_3}
+      }
+    }
+  }
+  Function main -> __void
+  StageDecoration{fragment}
+  ()
+  {
+    VariableDeclStatement{
+      VariableConst{
+        x_15
+        none
+        __u32
+        {
+          Call[not set]{
+            Identifier[not set]{x_11}
+            (
+              Identifier[not set]{x_1}
+            )
+          }
+        }
+      }
+    }
+    Return{}
+  }
+})")) << module_str;
+}
+
 }  // namespace
 }  // namespace spirv
 }  // namespace reader