spirv-reader: Sink pointer-to-vector-component

WGSL does not support pointer-to-vector-component, so the SPIR-V
reader needs to sink these pointers into their use.

Bug: tint:491
Change-Id: Ib5ae87d2f6bbac13280314ba11369d7ced505b56
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/68241
Kokoro: Kokoro <noreply+kokoro@google.com>
Reviewed-by: David Neto <dneto@google.com>

src/reader/spirv/function.cc

@@ -2541,6 +2541,12 @@ TypedExpression FunctionEmitter::MakeExpression(uint32_t id) {
       Fail() << "internal error: unhandled use of opaque object with ID: "
              << id;
       return {};
+    case SkipReason::kSinkPointerIntoUse: {
+      // Replace the pointer with its source reference expression.
+      auto source_expr = GetDefInfo(id)->sink_pointer_source_expr;
+      TINT_ASSERT(Reader, source_expr.type->Is<Reference>());
+      return source_expr;
+    }
     case SkipReason::kPointSizeBuiltinValue: {
       return {ty_.F32(),
               create<ast::ScalarConstructorExpression>(
@@ -3470,6 +3476,12 @@ bool FunctionEmitter::EmitConstDefinition(
   if (!expr) {
     return false;
   }
+
+  // Do not generate pointers that we want to sink.
+  if (GetDefInfo(inst.result_id())->skip == SkipReason::kSinkPointerIntoUse) {
+    return true;
+  }
+
   expr = AddressOfIfNeeded(expr, &inst);
   auto* ast_const = parser_impl_.MakeVariable(
       inst.result_id(), ast::StorageClass::kNone, expr.type, true, expr.expr,
@@ -3735,6 +3747,8 @@ bool FunctionEmitter::EmitStatement(const spvtools::opt::Instruction& inst) {
       const auto skip = GetSkipReason(value_id);
       if (skip != SkipReason::kDontSkip) {
         GetDefInfo(inst.result_id())->skip = skip;
+        GetDefInfo(inst.result_id())->sink_pointer_source_expr =
+            GetDefInfo(value_id)->sink_pointer_source_expr;
         return true;
       }
       auto expr = AddressOfIfNeeded(MakeExpression(value_id), &inst);
@@ -4214,6 +4228,8 @@ TypedExpression FunctionEmitter::MakeAccessChain(
   if (base_skip != SkipReason::kDontSkip) {
     // This can occur for AccessChain with no indices.
     GetDefInfo(inst.result_id())->skip = base_skip;
+    GetDefInfo(inst.result_id())->sink_pointer_source_expr =
+        GetDefInfo(base_id)->sink_pointer_source_expr;
     return {};
   }
 
@@ -4221,6 +4237,7 @@ TypedExpression FunctionEmitter::MakeAccessChain(
   uint32_t first_index = 1;
   const auto num_in_operands = inst.NumInOperands();
 
+  bool sink_pointer = false;
   TypedExpression current_expr;
 
   // If the variable was originally gl_PerVertex, then in the AST we
@@ -4352,6 +4369,8 @@ TypedExpression FunctionEmitter::MakeAccessChain(
         }
         // All vector components are the same type.
         pointee_type_id = pointee_type_inst->GetSingleWordInOperand(0);
+        // Sink pointers to vector components.
+        sink_pointer = true;
         break;
       case SpvOpTypeMatrix:
         // Use array syntax.
@@ -4407,6 +4426,13 @@ TypedExpression FunctionEmitter::MakeAccessChain(
     TINT_ASSERT(Reader, type && type->Is<Reference>());
     current_expr = TypedExpression{type, next_expr};
   }
+
+  if (sink_pointer) {
+    // Capture the reference so that we can sink it into the point of use.
+    GetDefInfo(inst.result_id())->skip = SkipReason::kSinkPointerIntoUse;
+    GetDefInfo(inst.result_id())->sink_pointer_source_expr = current_expr;
+  }
+
   return current_expr;
 }
 

src/reader/spirv/function.h

@@ -205,6 +205,12 @@ enum class SkipReason {
   /// function parameter).
   kOpaqueObject,
 
+  /// `kSinkPointerIntoUse`: used to avoid emitting certain pointer expressions,
+  /// by instead generating their reference expression directly at the point of
+  /// use. For example, we apply this to OpAccessChain when indexing into a
+  /// vector, to avoid generating address-of vector component expressions.
+  kSinkPointerIntoUse,
+
   /// `kPointSizeBuiltinPointer`: the value is a pointer to the Position builtin
   /// variable.  Don't generate its address.  Avoid generating stores to this
   /// pointer.
@@ -296,6 +302,11 @@ struct DefInfo {
   /// This is kInvalid for non-pointers.
   ast::StorageClass storage_class = ast::StorageClass::kInvalid;
 
+  /// The expression to use when sinking pointers into their use.
+  /// When encountering a use of this instruction, we will emit this expression
+  /// instead.
+  TypedExpression sink_pointer_source_expr = {};
+
   /// The reason, if any, that this value should be ignored.
   /// Normally no values are ignored.  This field can be updated while
   /// generating code because sometimes we only discover necessary facts
@@ -320,6 +331,9 @@ inline std::ostream& operator<<(std::ostream& o, const DefInfo& di) {
     case SkipReason::kOpaqueObject:
       o << " skip:opaque";
       break;
+    case SkipReason::kSinkPointerIntoUse:
+      o << " skip:sink_pointer";
+      break;
     case SkipReason::kPointSizeBuiltinPointer:
       o << " skip:pointsize_pointer";
       break;

src/reader/spirv/function_memory_test.cc

@@ -390,6 +390,119 @@ TEST_F(SpvParserMemoryTest, EmitStatement_AccessChain_VectorNonConstIndex) {
               HasSubstr("myvar[a_dynamic_index] = 42u;"));
 }
 
+TEST_F(SpvParserMemoryTest,
+       EmitStatement_AccessChain_VectorComponent_MultiUse) {
+  // WGSL does not support pointer-to-vector-component, so test that we sink
+  // these pointers into the point of use.
+  const std::string assembly = Preamble() + R"(
+     OpName %1 "myvar"
+     %void = OpTypeVoid
+     %voidfn = OpTypeFunction %void
+     %uint = OpTypeInt 32 0
+     %store_ty = OpTypeVector %uint 4
+     %uint_2 = OpConstant %uint 2
+     %uint_42 = OpConstant %uint 42
+     %elem_ty = OpTypePointer Private %uint
+     %var_ty = OpTypePointer Private %store_ty
+     %1 = OpVariable %var_ty Private
+     %100 = OpFunction %void None %voidfn
+     %entry = OpLabel
+     %ptr = OpAccessChain %elem_ty %1 %uint_2
+     %load = OpLoad %uint %ptr
+     %result = OpIAdd %uint %load %uint_2
+     OpStore %ptr %result
+     OpReturn
+     OpFunctionEnd
+  )";
+  auto p = parser(test::Assemble(assembly));
+  ASSERT_TRUE(p->BuildAndParseInternalModuleExceptFunctions())
+      << assembly << p->error();
+  auto fe = p->function_emitter(100);
+  EXPECT_TRUE(fe.EmitBody()) << p->error();
+  auto ast_body = fe.ast_body();
+  auto wgsl = test::ToString(p->program(), ast_body);
+  EXPECT_THAT(wgsl, Not(HasSubstr("&")));
+  EXPECT_THAT(wgsl, HasSubstr(" = myvar.z;"));
+  EXPECT_THAT(wgsl, HasSubstr("myvar.z = "));
+}
+
+TEST_F(SpvParserMemoryTest,
+       EmitStatement_AccessChain_VectorComponent_MultiUse_NonConstIndex) {
+  // WGSL does not support pointer-to-vector-component, so test that we sink
+  // these pointers into the point of use.
+  const std::string assembly = Preamble() + R"(
+     OpName %1 "myvar"
+     %void = OpTypeVoid
+     %voidfn = OpTypeFunction %void
+     %uint = OpTypeInt 32 0
+     %store_ty = OpTypeVector %uint 4
+     %uint_2 = OpConstant %uint 2
+     %uint_42 = OpConstant %uint 42
+     %elem_ty = OpTypePointer Private %uint
+     %var_ty = OpTypePointer Private %store_ty
+     %1 = OpVariable %var_ty Private
+     %2 = OpVariable %elem_ty Private
+     %100 = OpFunction %void None %voidfn
+     %entry = OpLabel
+     %idx = OpLoad %uint %2
+     %ptr = OpAccessChain %elem_ty %1 %idx
+     %load = OpLoad %uint %ptr
+     %result = OpIAdd %uint %load %uint_2
+     OpStore %ptr %result
+     OpReturn
+     OpFunctionEnd
+  )";
+  auto p = parser(test::Assemble(assembly));
+  ASSERT_TRUE(p->BuildAndParseInternalModuleExceptFunctions())
+      << assembly << p->error();
+  auto fe = p->function_emitter(100);
+  EXPECT_TRUE(fe.EmitBody()) << p->error();
+  auto ast_body = fe.ast_body();
+  auto wgsl = test::ToString(p->program(), ast_body);
+  EXPECT_THAT(wgsl, Not(HasSubstr("&")));
+  EXPECT_THAT(wgsl, HasSubstr(" = myvar[x_12];"));
+  EXPECT_THAT(wgsl, HasSubstr("myvar[x_12] = "));
+}
+
+TEST_F(SpvParserMemoryTest,
+       EmitStatement_AccessChain_VectorComponent_SinkThroughChain) {
+  // Test that we can sink a pointer-to-vector-component through a chain of
+  // instructions that propagate it.
+  const std::string assembly = Preamble() + R"(
+     OpName %1 "myvar"
+     %void = OpTypeVoid
+     %voidfn = OpTypeFunction %void
+     %uint = OpTypeInt 32 0
+     %store_ty = OpTypeVector %uint 4
+     %uint_2 = OpConstant %uint 2
+     %uint_42 = OpConstant %uint 42
+     %elem_ty = OpTypePointer Private %uint
+     %var_ty = OpTypePointer Private %store_ty
+     %1 = OpVariable %var_ty Private
+     %100 = OpFunction %void None %voidfn
+     %entry = OpLabel
+     %ptr = OpAccessChain %elem_ty %1 %uint_2
+     %ptr2 = OpCopyObject %elem_ty %ptr
+     %ptr3 = OpInBoundsAccessChain %elem_ty %ptr2
+     %ptr4 = OpAccessChain %elem_ty %ptr3
+     %load = OpLoad %uint %ptr3
+     %result = OpIAdd %uint %load %uint_2
+     OpStore %ptr4 %result
+     OpReturn
+     OpFunctionEnd
+  )";
+  auto p = parser(test::Assemble(assembly));
+  ASSERT_TRUE(p->BuildAndParseInternalModuleExceptFunctions())
+      << assembly << p->error();
+  auto fe = p->function_emitter(100);
+  EXPECT_TRUE(fe.EmitBody()) << p->error();
+  auto ast_body = fe.ast_body();
+  auto wgsl = test::ToString(p->program(), ast_body);
+  EXPECT_THAT(wgsl, Not(HasSubstr("&")));
+  EXPECT_THAT(wgsl, HasSubstr(" = myvar.z;"));
+  EXPECT_THAT(wgsl, HasSubstr("myvar.z = "));
+}
+
 TEST_F(SpvParserMemoryTest, EmitStatement_AccessChain_Matrix) {
   const std::string assembly = Preamble() + R"(
      OpName %1 "myvar"