src/tint/lang/spirv/reader/lower/shader_io.cc - dawn - Git at Google

 // Copyright 2024 The Dawn & Tint Authors
 //
 // Redistribution and use in source and binary forms, with or without
 // modification, are permitted provided that the following conditions are met:
 //
 // 1. Redistributions of source code must retain the above copyright notice, this
 //    list of conditions and the following disclaimer.
 //
 // 2. Redistributions in binary form must reproduce the above copyright notice,
 //    this list of conditions and the following disclaimer in the documentation
 //    and/or other materials provided with the distribution.
 //
 // 3. Neither the name of the copyright holder nor the names of its
 //    contributors may be used to endorse or promote products derived from
 //    this software without specific prior written permission.
 //
 // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 // AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 // DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
 // FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 // DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 // SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 // CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

 #include "src/tint/lang/spirv/reader/lower/shader_io.h"

 #include <utility>

 #include "src/tint/lang/core/ir/builder.h"
 #include "src/tint/lang/core/ir/module.h"
 #include "src/tint/lang/core/ir/transform/common/referenced_module_vars.h"
 #include "src/tint/lang/core/ir/validator.h"

 namespace tint::spirv::reader::lower {

 namespace {

 using namespace tint::core::fluent_types;  // NOLINT

 /// PIMPL state for the transform.
 struct State {
     /// The IR module.
     core::ir::Module& ir;

     /// The IR builder.
     core::ir::Builder b{ir};

     /// The type manager.
     core::type::Manager& ty{ir.Types()};

     /// A map from block to its containing function.
     Hashmap<core::ir::Block*, core::ir::Function*, 64> block_to_function{};

     /// A map from each function to a map from input variable to parameter.
     Hashmap<core::ir::Function*, Hashmap<core::ir::Var*, core::ir::Value*, 4>, 8>
         function_parameter_map{};

     /// The set of output variables that have been processed.
     Hashset<core::ir::Var*, 4> output_variables{};

     /// The mapping from functions to their transitively referenced output variables.
     core::ir::ReferencedModuleVars referenced_output_vars{
         ir, [](const core::ir::Var* var) {
             auto* view = var->Result(0)->Type()->As<core::type::MemoryView>();
             return view && view->AddressSpace() == core::AddressSpace::kOut;
         }};

     /// Process the module.
     void Process() {
         // Process outputs first, as that may introduce new functions that input variables need to
         // be propagated through.
         ProcessOutputs();
         ProcessInputs();
     }

     /// Process output variables.
     /// Changes output variables to the `private` address space and wraps entry points that produce
     /// outputs with new functions that copy the outputs from the private variables to the return
     /// value.
     void ProcessOutputs() {
         // Update entry point functions to return their outputs, using a wrapper function.
         // Use a worklist as `ProcessEntryPointOutputs()` will add new functions.
         Vector<core::ir::Function*, 4> entry_points;
         for (auto& func : ir.functions) {
             if (func->Stage() != core::ir::Function::PipelineStage::kUndefined) {
                 entry_points.Push(func);
             }
         }
         for (auto& ep : entry_points) {
             ProcessEntryPointOutputs(ep);
         }

         // Remove attributes from all of the original structs and module-scope output variables.
         // This is done last as we need to copy attributes during `ProcessEntryPointOutputs()`.
         for (auto& var : output_variables) {
             var->SetAttributes({});
             if (auto* str = var->Result(0)->Type()->UnwrapPtr()->As<core::type::Struct>()) {
                 for (auto* member : str->Members()) {
                     // TODO(crbug.com/tint/745): Remove the const_cast.
                     const_cast<core::type::StructMember*>(member)->SetAttributes({});
                 }
             }
         }
     }

     /// Process input variables.
     /// Pass inputs down the call stack as parameters to any functions that need them.
     void ProcessInputs() {
         // Seed the block-to-function map with the function entry blocks.
         for (auto& func : ir.functions) {
             block_to_function.Add(func->Block(), func);
         }

         // Gather the list of all module-scope input variables.
         Vector<core::ir::Var*, 4> inputs;
         for (auto* global : *ir.root_block) {
             if (auto* var = global->As<core::ir::Var>()) {
                 auto addrspace = var->Result(0)->Type()->As<core::type::Pointer>()->AddressSpace();
                 if (addrspace == core::AddressSpace::kIn) {
                     inputs.Push(var);
                 }
             }
         }

         // Replace the input variables with function parameters.
         for (auto* var : inputs) {
             ReplaceInputPointerUses(var, var->Result(0));
             var->Destroy();
         }
     }

     /// Replace an output pointer address space to make it `private`.
     /// @param value the output variable
     void ReplaceOutputPointerAddressSpace(core::ir::InstructionResult* value) {
         // Change the address space to `private`.
         auto* old_ptr_type = value->Type();
         auto* new_ptr_type = ty.ptr(core::AddressSpace::kPrivate, old_ptr_type->UnwrapPtr());
         value->SetType(new_ptr_type);

         // Update all uses of the module-scope variable.
         value->ForEachUse([&](core::ir::Usage use) {
             if (auto* access = use.instruction->As<core::ir::Access>()) {
                 ReplaceOutputPointerAddressSpace(access->Result(0));
             } else if (!use.instruction->IsAnyOf<core::ir::Load, core::ir::LoadVectorElement,
                                                  core::ir::Store, core::ir::StoreVectorElement>()) {
                 TINT_UNREACHABLE()
                     << "unexpected instruction: " << use.instruction->TypeInfo().name;
             }
         });
     }

     /// Process the outputs of an entry point function, adding a wrapper function to forward outputs
     /// through the return value.
     /// @param ep the entry point
     void ProcessEntryPointOutputs(core::ir::Function* ep) {
         const auto& referenced_outputs = referenced_output_vars.TransitiveReferences(ep);
         if (referenced_outputs.IsEmpty()) {
             return;
         }

         // Add a wrapper function to return either a single value or a struct.
         auto* wrapper = b.Function(ty.void_(), ep->Stage());
         if (auto name = ir.NameOf(ep)) {
             ir.SetName(ep, name.Name() + "_inner");
             ir.SetName(wrapper, name);
         }

         // Call the original entry point and make it a regular function.
         ep->SetStage(core::ir::Function::PipelineStage::kUndefined);
         b.Append(wrapper->Block(), [&] {  //
             b.Call(ep);
         });

         // Collect all outputs into a list of struct member declarations.
         // Also add instructions to load their final values in the wrapper function.
         Vector<core::ir::Value*, 4> results;
         Vector<core::type::Manager::StructMemberDesc, 4> output_descriptors;
         auto add_output = [&](Symbol name, const core::type::Type* type,
                               core::type::StructMemberAttributes attributes) {
             if (!name) {
                 name = ir.symbols.New();
             }
             output_descriptors.Push(core::type::Manager::StructMemberDesc{name, type, attributes});
         };
         for (auto* var : referenced_outputs) {
             // Change the address space of the variable to private and update its uses, if we
             // haven't already seen this variable.
             if (output_variables.Add(var)) {
                 ReplaceOutputPointerAddressSpace(var->Result(0));
             }

             // Copy the variable attributes to the struct member.
             const auto& original_attributes = var->Attributes();
             core::type::StructMemberAttributes var_attributes;
             var_attributes.invariant = original_attributes.invariant;
             var_attributes.builtin = original_attributes.builtin;
             var_attributes.location = original_attributes.location;
             var_attributes.interpolation = original_attributes.interpolation;

             auto var_type = var->Result(0)->Type()->UnwrapPtr();
             if (auto* str = var_type->As<core::type::Struct>()) {
                 // Add an output for each member of the struct.
                 for (auto* member : str->Members()) {
                     // Use the base variable attributes if not specified directly on the member.
                     auto member_attributes = member->Attributes();
                     if (auto base_loc = var_attributes.location) {
                         // Location values increment from the base location value on the variable.
                         member_attributes.location = base_loc.value() + member->Index();
                     }
                     if (!member_attributes.interpolation) {
                         member_attributes.interpolation = var_attributes.interpolation;
                     }

                     add_output(member->Name(), member->Type(), std::move(member_attributes));

                     // Load the final result from the member of the original struct variable.
                     b.Append(wrapper->Block(), [&] {  //
                         auto* access =
                             b.Access(ty.ptr<private_>(member->Type()), var, u32(member->Index()));
                         results.Push(b.Load(access)->Result(0));
                     });
                 }
             } else {
                 // Load the final result from the original variable.
                 b.Append(wrapper->Block(), [&] {
                     results.Push(b.Load(var)->Result(0));

                     // If we're dealing with sample_mask, extract the scalar from the array.
                     if (var_attributes.builtin == core::BuiltinValue::kSampleMask) {
                         var_type = ty.u32();
                         results.Back() = b.Access(ty.u32(), results.Back(), u32(0))->Result(0);
                     }
                 });
                 add_output(ir.NameOf(var), var_type, std::move(var_attributes));
             }
         }

         if (output_descriptors.Length() == 1) {
             // Copy the output attributes to the function return.
             const auto& attributes = output_descriptors[0].attributes;
             wrapper->SetReturnInvariant(attributes.invariant);
             if (attributes.builtin) {
                 wrapper->SetReturnBuiltin(attributes.builtin.value());
             } else if (attributes.location) {
                 core::ir::Location loc;
                 loc.value = attributes.location.value();
                 loc.interpolation = attributes.interpolation;
                 wrapper->SetReturnLocation(std::move(loc));
             }

             // Return the output from the wrapper function.
             wrapper->SetReturnType(output_descriptors[0].type);
             b.Append(wrapper->Block(), [&] {  //
                 b.Return(wrapper, results[0]);
             });
         } else {
             // Create a struct to hold all of the output values.
             auto* str = ty.Struct(ir.symbols.New(), std::move(output_descriptors));
             wrapper->SetReturnType(str);

             // Collect the output values and return them from the wrapper function.
             b.Append(wrapper->Block(), [&] {  //
                 b.Return(wrapper, b.Construct(str, std::move(results)));
             });
         }
     }

     /// Replace a use of an input pointer value.
     /// @param var the originating input variable
     /// @param value the input pointer value
     void ReplaceInputPointerUses(core::ir::Var* var, core::ir::Value* value) {
         Vector<core::ir::Instruction*, 8> to_destroy;
         value->ForEachUse([&](core::ir::Usage use) {
             auto* object = value;
             if (object->Type()->Is<core::type::Pointer>()) {
                 // Get (or create) the function parameter that will replace the variable.
                 auto* func = ContainingFunction(use.instruction);
                 object = GetParameter(func, var);
             }

             Switch(
                 use.instruction,
                 [&](core::ir::Load* l) {
                     // Fold the load away and replace its uses with the new parameter.
                     l->Result(0)->ReplaceAllUsesWith(object);
                     to_destroy.Push(l);
                 },
                 [&](core::ir::LoadVectorElement* lve) {
                     // Replace the vector element load with an access instruction.
                     auto* access = b.AccessWithResult(lve->DetachResult(), object, lve->Index());
                     access->InsertBefore(lve);
                     to_destroy.Push(lve);
                 },
                 [&](core::ir::Access* a) {
                     if (!a->Indices().IsEmpty()) {
                         // Remove the pointer from the source and destination type.
                         a->SetOperand(core::ir::Access::kObjectOperandOffset, object);
                         a->Result(0)->SetType(a->Result(0)->Type()->UnwrapPtr());
                         ReplaceInputPointerUses(var, a->Result(0));
                     } else {
                         // Fold the access away and replace its uses.
                         ReplaceInputPointerUses(var, a->Result(0));
                         to_destroy.Push(a);
                     }
                 },
                 TINT_ICE_ON_NO_MATCH);
         });

         // Clean up orphaned instructions.
         for (auto* inst : to_destroy) {
             inst->Destroy();
         }
     }

     /// Get the function that contains an instruction.
     /// @param inst the instruction
     /// @returns the function
     core::ir::Function* ContainingFunction(core::ir::Instruction* inst) {
         return block_to_function.GetOrAdd(inst->Block(), [&] {  //
             return ContainingFunction(inst->Block()->Parent());
         });
     }

     /// Get or create a function parameter to replace a module-scope variable.
     /// @param func the function
     /// @param var the module-scope variable
     /// @returns the function parameter
     core::ir::Value* GetParameter(core::ir::Function* func, core::ir::Var* var) {
         return function_parameter_map.GetOrAddZero(func).GetOrAdd(var, [&] {
             const bool entry_point = func->Stage() != core::ir::Function::PipelineStage::kUndefined;
             auto* var_type = var->Result(0)->Type()->UnwrapPtr();

             // Use a scalar u32 for sample_mask builtins for entry point parameters.
             if (entry_point && var->Attributes().builtin == core::BuiltinValue::kSampleMask) {
                 TINT_ASSERT(var_type->Is<core::type::Array>());
                 TINT_ASSERT(var_type->As<core::type::Array>()->ConstantCount() == 1u);
                 var_type = ty.u32();
             }

             // Create a new function parameter for the input.
             auto* param = b.FunctionParam(var_type);
             func->AppendParam(param);
             if (auto name = ir.NameOf(var)) {
                 ir.SetName(param, name);
             }

             // Add attributes to the parameter if this is an entry point function.
             if (entry_point) {
                 AddEntryPointParameterAttributes(param, var->Attributes());
             }

             // Update the callsites of this function.
             func->ForEachUse([&](core::ir::Usage use) {
                 if (auto* call = use.instruction->As<core::ir::UserCall>()) {
                     // Recurse into the calling function.
                     auto* caller = ContainingFunction(call);
                     call->AppendArg(GetParameter(caller, var));
                 } else if (!use.instruction->Is<core::ir::Return>()) {
                     TINT_UNREACHABLE()
                         << "unexpected instruction: " << use.instruction->TypeInfo().name;
                 }
             });

             core::ir::Value* result = param;
             if (entry_point && var->Attributes().builtin == core::BuiltinValue::kSampleMask) {
                 // Construct an array from the scalar sample_mask builtin value for entry points.
                 b.Prepend(func->Block(), [&] {  //
                     result = b.Construct(var->Result(0)->Type()->UnwrapPtr(), param)->Result(0);
                 });
             }
             return result;
         });
     }

     /// Add attributes to an entry point function parameter.
     /// @param param the parameter
     /// @param attributes the attributes
     void AddEntryPointParameterAttributes(core::ir::FunctionParam* param,
                                           const core::ir::IOAttributes& attributes) {
         if (auto* str = param->Type()->UnwrapPtr()->As<core::type::Struct>()) {
             for (auto* member : str->Members()) {
                 // Use the base variable attributes if not specified directly on the member.
                 auto member_attributes = member->Attributes();
                 if (auto base_loc = attributes.location) {
                     // Location values increment from the base location value on the variable.
                     member_attributes.location = base_loc.value() + member->Index();
                 }
                 if (!member_attributes.interpolation) {
                     member_attributes.interpolation = attributes.interpolation;
                 }
                 // TODO(crbug.com/tint/745): Remove the const_cast.
                 const_cast<core::type::StructMember*>(member)->SetAttributes(
                     std::move(member_attributes));
             }
         } else {
             // Set attributes directly on the function parameter.
             param->SetInvariant(attributes.invariant);
             if (attributes.builtin) {
                 param->SetBuiltin(attributes.builtin.value());
             } else if (attributes.location) {
                 core::ir::Location loc;
                 loc.value = attributes.location.value();
                 loc.interpolation = attributes.interpolation;
                 param->SetLocation(std::move(loc));
             }
         }
     }
 };

 }  // namespace

 Result<SuccessType> ShaderIO(core::ir::Module& ir) {
     auto result = ValidateAndDumpIfNeeded(ir, "ShaderIO transform");
     if (result != Success) {
         return result.Failure();
     }

     State{ir}.Process();

     return Success;
 }

 }  // namespace tint::spirv::reader::lower
	// Copyright 2024 The Dawn & Tint Authors
	//
	// Redistribution and use in source and binary forms, with or without
	// modification, are permitted provided that the following conditions are met:
	//
	// 1. Redistributions of source code must retain the above copyright notice, this
	// list of conditions and the following disclaimer.
	//
	// 2. Redistributions in binary form must reproduce the above copyright notice,
	// this list of conditions and the following disclaimer in the documentation
	// and/or other materials provided with the distribution.
	//
	// 3. Neither the name of the copyright holder nor the names of its
	// contributors may be used to endorse or promote products derived from
	// this software without specific prior written permission.
	//
	// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
	// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
	// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
	// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
	// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
	// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
	// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
	// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
	// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
	// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

	#include "src/tint/lang/spirv/reader/lower/shader_io.h"

	#include <utility>

	#include "src/tint/lang/core/ir/builder.h"
	#include "src/tint/lang/core/ir/module.h"
	#include "src/tint/lang/core/ir/transform/common/referenced_module_vars.h"
	#include "src/tint/lang/core/ir/validator.h"

	namespace tint::spirv::reader::lower {

	namespace {

	using namespace tint::core::fluent_types; // NOLINT

	/// PIMPL state for the transform.
	struct State {
	/// The IR module.
	core::ir::Module& ir;

	/// The IR builder.
	core::ir::Builder b{ir};

	/// The type manager.
	core::type::Manager& ty{ir.Types()};

	/// A map from block to its containing function.
	Hashmap<core::ir::Block, core::ir::Function, 64> block_to_function{};

	/// A map from each function to a map from input variable to parameter.
	Hashmap<core::ir::Function, Hashmap<core::ir::Var, core::ir::Value*, 4>, 8>
	function_parameter_map{};

	/// The set of output variables that have been processed.
	Hashset<core::ir::Var*, 4> output_variables{};

	/// The mapping from functions to their transitively referenced output variables.
	core::ir::ReferencedModuleVars referenced_output_vars{
	ir, [](const core::ir::Var* var) {
	auto* view = var->Result(0)->Type()->As<core::type::MemoryView>();
	return view && view->AddressSpace() == core::AddressSpace::kOut;
	}};

	/// Process the module.
	void Process() {
	// Process outputs first, as that may introduce new functions that input variables need to
	// be propagated through.
	ProcessOutputs();
	ProcessInputs();
	}

	/// Process output variables.
	/// Changes output variables to the `private` address space and wraps entry points that produce
	/// outputs with new functions that copy the outputs from the private variables to the return
	/// value.
	void ProcessOutputs() {
	// Update entry point functions to return their outputs, using a wrapper function.
	// Use a worklist as `ProcessEntryPointOutputs()` will add new functions.
	Vector<core::ir::Function*, 4> entry_points;
	for (auto& func : ir.functions) {
	if (func->Stage() != core::ir::Function::PipelineStage::kUndefined) {
	entry_points.Push(func);
	}
	}
	for (auto& ep : entry_points) {
	ProcessEntryPointOutputs(ep);
	}

	// Remove attributes from all of the original structs and module-scope output variables.
	// This is done last as we need to copy attributes during `ProcessEntryPointOutputs()`.
	for (auto& var : output_variables) {
	var->SetAttributes({});
	if (auto* str = var->Result(0)->Type()->UnwrapPtr()->As<core::type::Struct>()) {
	for (auto* member : str->Members()) {
	// TODO(crbug.com/tint/745): Remove the const_cast.
	const_cast<core::type::StructMember*>(member)->SetAttributes({});
	}
	}
	}
	}

	/// Process input variables.
	/// Pass inputs down the call stack as parameters to any functions that need them.
	void ProcessInputs() {
	// Seed the block-to-function map with the function entry blocks.
	for (auto& func : ir.functions) {
	block_to_function.Add(func->Block(), func);
	}

	// Gather the list of all module-scope input variables.
	Vector<core::ir::Var*, 4> inputs;
	for (auto* global : *ir.root_block) {
	if (auto* var = global->As<core::ir::Var>()) {
	auto addrspace = var->Result(0)->Type()->As<core::type::Pointer>()->AddressSpace();
	if (addrspace == core::AddressSpace::kIn) {
	inputs.Push(var);
	}
	}
	}

	// Replace the input variables with function parameters.
	for (auto* var : inputs) {
	ReplaceInputPointerUses(var, var->Result(0));
	var->Destroy();
	}
	}

	/// Replace an output pointer address space to make it `private`.
	/// @param value the output variable
	void ReplaceOutputPointerAddressSpace(core::ir::InstructionResult* value) {
	// Change the address space to `private`.
	auto* old_ptr_type = value->Type();
	auto* new_ptr_type = ty.ptr(core::AddressSpace::kPrivate, old_ptr_type->UnwrapPtr());
	value->SetType(new_ptr_type);

	// Update all uses of the module-scope variable.
	value->ForEachUse([&](core::ir::Usage use) {
	if (auto* access = use.instruction->As<core::ir::Access>()) {
	ReplaceOutputPointerAddressSpace(access->Result(0));
	} else if (!use.instruction->IsAnyOf<core::ir::Load, core::ir::LoadVectorElement,
	core::ir::Store, core::ir::StoreVectorElement>()) {
	TINT_UNREACHABLE()
	<< "unexpected instruction: " << use.instruction->TypeInfo().name;
	}
	});
	}

	/// Process the outputs of an entry point function, adding a wrapper function to forward outputs
	/// through the return value.
	/// @param ep the entry point
	void ProcessEntryPointOutputs(core::ir::Function* ep) {
	const auto& referenced_outputs = referenced_output_vars.TransitiveReferences(ep);
	if (referenced_outputs.IsEmpty()) {
	return;
	}

	// Add a wrapper function to return either a single value or a struct.
	auto* wrapper = b.Function(ty.void_(), ep->Stage());
	if (auto name = ir.NameOf(ep)) {
	ir.SetName(ep, name.Name() + "_inner");
	ir.SetName(wrapper, name);
	}

	// Call the original entry point and make it a regular function.
	ep->SetStage(core::ir::Function::PipelineStage::kUndefined);
	b.Append(wrapper->Block(), [&] { //
	b.Call(ep);
	});

	// Collect all outputs into a list of struct member declarations.
	// Also add instructions to load their final values in the wrapper function.
	Vector<core::ir::Value*, 4> results;
	Vector<core::type::Manager::StructMemberDesc, 4> output_descriptors;
	auto add_output = [&](Symbol name, const core::type::Type* type,
	core::type::StructMemberAttributes attributes) {
	if (!name) {
	name = ir.symbols.New();
	}
	output_descriptors.Push(core::type::Manager::StructMemberDesc{name, type, attributes});
	};
	for (auto* var : referenced_outputs) {
	// Change the address space of the variable to private and update its uses, if we
	// haven't already seen this variable.
	if (output_variables.Add(var)) {
	ReplaceOutputPointerAddressSpace(var->Result(0));
	}

	// Copy the variable attributes to the struct member.
	const auto& original_attributes = var->Attributes();
	core::type::StructMemberAttributes var_attributes;
	var_attributes.invariant = original_attributes.invariant;
	var_attributes.builtin = original_attributes.builtin;
	var_attributes.location = original_attributes.location;
	var_attributes.interpolation = original_attributes.interpolation;

	auto var_type = var->Result(0)->Type()->UnwrapPtr();
	if (auto* str = var_type->As<core::type::Struct>()) {
	// Add an output for each member of the struct.
	for (auto* member : str->Members()) {
	// Use the base variable attributes if not specified directly on the member.
	auto member_attributes = member->Attributes();
	if (auto base_loc = var_attributes.location) {
	// Location values increment from the base location value on the variable.
	member_attributes.location = base_loc.value() + member->Index();
	}
	if (!member_attributes.interpolation) {
	member_attributes.interpolation = var_attributes.interpolation;
	}

	add_output(member->Name(), member->Type(), std::move(member_attributes));

	// Load the final result from the member of the original struct variable.
	b.Append(wrapper->Block(), [&] { //
	auto* access =
	b.Access(ty.ptr<private_>(member->Type()), var, u32(member->Index()));
	results.Push(b.Load(access)->Result(0));
	});
	}
	} else {
	// Load the final result from the original variable.
	b.Append(wrapper->Block(), [&] {
	results.Push(b.Load(var)->Result(0));

	// If we're dealing with sample_mask, extract the scalar from the array.
	if (var_attributes.builtin == core::BuiltinValue::kSampleMask) {
	var_type = ty.u32();
	results.Back() = b.Access(ty.u32(), results.Back(), u32(0))->Result(0);
	}
	});
	add_output(ir.NameOf(var), var_type, std::move(var_attributes));
	}
	}

	if (output_descriptors.Length() == 1) {
	// Copy the output attributes to the function return.
	const auto& attributes = output_descriptors[0].attributes;
	wrapper->SetReturnInvariant(attributes.invariant);
	if (attributes.builtin) {
	wrapper->SetReturnBuiltin(attributes.builtin.value());
	} else if (attributes.location) {
	core::ir::Location loc;
	loc.value = attributes.location.value();
	loc.interpolation = attributes.interpolation;
	wrapper->SetReturnLocation(std::move(loc));
	}

	// Return the output from the wrapper function.
	wrapper->SetReturnType(output_descriptors[0].type);
	b.Append(wrapper->Block(), [&] { //
	b.Return(wrapper, results[0]);
	});
	} else {
	// Create a struct to hold all of the output values.
	auto* str = ty.Struct(ir.symbols.New(), std::move(output_descriptors));
	wrapper->SetReturnType(str);

	// Collect the output values and return them from the wrapper function.
	b.Append(wrapper->Block(), [&] { //
	b.Return(wrapper, b.Construct(str, std::move(results)));
	});
	}
	}

	/// Replace a use of an input pointer value.
	/// @param var the originating input variable
	/// @param value the input pointer value
	void ReplaceInputPointerUses(core::ir::Var* var, core::ir::Value* value) {
	Vector<core::ir::Instruction*, 8> to_destroy;
	value->ForEachUse([&](core::ir::Usage use) {
	auto* object = value;
	if (object->Type()->Is<core::type::Pointer>()) {
	// Get (or create) the function parameter that will replace the variable.
	auto* func = ContainingFunction(use.instruction);
	object = GetParameter(func, var);
	}

	Switch(
	use.instruction,
	[&](core::ir::Load* l) {
	// Fold the load away and replace its uses with the new parameter.
	l->Result(0)->ReplaceAllUsesWith(object);
	to_destroy.Push(l);
	},
	[&](core::ir::LoadVectorElement* lve) {
	// Replace the vector element load with an access instruction.
	auto* access = b.AccessWithResult(lve->DetachResult(), object, lve->Index());
	access->InsertBefore(lve);
	to_destroy.Push(lve);
	},
	[&](core::ir::Access* a) {
	if (!a->Indices().IsEmpty()) {
	// Remove the pointer from the source and destination type.
	a->SetOperand(core::ir::Access::kObjectOperandOffset, object);
	a->Result(0)->SetType(a->Result(0)->Type()->UnwrapPtr());
	ReplaceInputPointerUses(var, a->Result(0));
	} else {
	// Fold the access away and replace its uses.
	ReplaceInputPointerUses(var, a->Result(0));
	to_destroy.Push(a);
	}
	},
	TINT_ICE_ON_NO_MATCH);
	});

	// Clean up orphaned instructions.
	for (auto* inst : to_destroy) {
	inst->Destroy();
	}
	}

	/// Get the function that contains an instruction.
	/// @param inst the instruction
	/// @returns the function
	core::ir::Function* ContainingFunction(core::ir::Instruction* inst) {
	return block_to_function.GetOrAdd(inst->Block(), [&] { //
	return ContainingFunction(inst->Block()->Parent());
	});
	}

	/// Get or create a function parameter to replace a module-scope variable.
	/// @param func the function
	/// @param var the module-scope variable
	/// @returns the function parameter
	core::ir::Value* GetParameter(core::ir::Function* func, core::ir::Var* var) {
	return function_parameter_map.GetOrAddZero(func).GetOrAdd(var, [&] {
	const bool entry_point = func->Stage() != core::ir::Function::PipelineStage::kUndefined;
	auto* var_type = var->Result(0)->Type()->UnwrapPtr();

	// Use a scalar u32 for sample_mask builtins for entry point parameters.
	if (entry_point && var->Attributes().builtin == core::BuiltinValue::kSampleMask) {
	TINT_ASSERT(var_type->Is<core::type::Array>());
	TINT_ASSERT(var_type->As<core::type::Array>()->ConstantCount() == 1u);
	var_type = ty.u32();
	}

	// Create a new function parameter for the input.
	auto* param = b.FunctionParam(var_type);
	func->AppendParam(param);
	if (auto name = ir.NameOf(var)) {
	ir.SetName(param, name);
	}

	// Add attributes to the parameter if this is an entry point function.
	if (entry_point) {
	AddEntryPointParameterAttributes(param, var->Attributes());
	}

	// Update the callsites of this function.
	func->ForEachUse([&](core::ir::Usage use) {
	if (auto* call = use.instruction->As<core::ir::UserCall>()) {
	// Recurse into the calling function.
	auto* caller = ContainingFunction(call);
	call->AppendArg(GetParameter(caller, var));
	} else if (!use.instruction->Is<core::ir::Return>()) {
	TINT_UNREACHABLE()
	<< "unexpected instruction: " << use.instruction->TypeInfo().name;
	}
	});

	core::ir::Value* result = param;
	if (entry_point && var->Attributes().builtin == core::BuiltinValue::kSampleMask) {
	// Construct an array from the scalar sample_mask builtin value for entry points.
	b.Prepend(func->Block(), [&] { //
	result = b.Construct(var->Result(0)->Type()->UnwrapPtr(), param)->Result(0);
	});
	}
	return result;
	});
	}

	/// Add attributes to an entry point function parameter.
	/// @param param the parameter
	/// @param attributes the attributes
	void AddEntryPointParameterAttributes(core::ir::FunctionParam* param,
	const core::ir::IOAttributes& attributes) {
	if (auto* str = param->Type()->UnwrapPtr()->As<core::type::Struct>()) {
	for (auto* member : str->Members()) {
	// Use the base variable attributes if not specified directly on the member.
	auto member_attributes = member->Attributes();
	if (auto base_loc = attributes.location) {
	// Location values increment from the base location value on the variable.
	member_attributes.location = base_loc.value() + member->Index();
	}
	if (!member_attributes.interpolation) {
	member_attributes.interpolation = attributes.interpolation;
	}
	// TODO(crbug.com/tint/745): Remove the const_cast.
	const_cast<core::type::StructMember*>(member)->SetAttributes(
	std::move(member_attributes));
	}
	} else {
	// Set attributes directly on the function parameter.
	param->SetInvariant(attributes.invariant);
	if (attributes.builtin) {
	param->SetBuiltin(attributes.builtin.value());
	} else if (attributes.location) {
	core::ir::Location loc;
	loc.value = attributes.location.value();
	loc.interpolation = attributes.interpolation;
	param->SetLocation(std::move(loc));
	}
	}
	}
	};

	} // namespace

	Result<SuccessType> ShaderIO(core::ir::Module& ir) {
	auto result = ValidateAndDumpIfNeeded(ir, "ShaderIO transform");
	if (result != Success) {
	return result.Failure();
	}

	State{ir}.Process();

	return Success;
	}

	} // namespace tint::spirv::reader::lower