1 //
2 // Copyright 2024 The ANGLE Project Authors. All rights reserved.
3 // Use of this source code is governed by a BSD-style license that can be
4 // found in the LICENSE file.
5 //
6
7 #include "libANGLE/renderer/wgpu/wgpu_wgsl_util.h"
8
9 #include <sstream>
10
11 #include "common/PackedEnums.h"
12 #include "common/PackedGLEnums_autogen.h"
13 #include "libANGLE/Program.h"
14 #include "libANGLE/ProgramExecutable.h"
15
16 namespace rx
17 {
18 namespace webgpu
19 {
20
21 namespace
22 {
23 const bool kOutputReplacements = false;
24
WgslReplaceLocationMarkers(const std::string & shaderSource,std::map<std::string,int> varNameToLocation)25 std::string WgslReplaceLocationMarkers(const std::string &shaderSource,
26 std::map<std::string, int> varNameToLocation)
27 {
28 const char *marker = "@location(@@@@@@) ";
29 const char *endOfName = " : ";
30
31 std::string newSource;
32 newSource.reserve(shaderSource.size());
33
34 size_t currPos = 0;
35 while (true)
36 {
37 size_t nextMarker = shaderSource.find(marker, currPos, strlen(marker));
38 if (nextMarker == std::string::npos)
39 {
40 // Copy the rest of the shader and end the loop.
41 newSource.append(shaderSource, currPos);
42 break;
43 }
44 else
45 {
46 // Copy up to the next marker
47 newSource.append(shaderSource, currPos, nextMarker - currPos);
48
49 // Extract name from something like `@location(@@@@@@) NAME : TYPE`.
50 size_t startOfNamePos = nextMarker + strlen(marker);
51 size_t endOfNamePos = shaderSource.find(endOfName, startOfNamePos, strlen(endOfName));
52 std::string name(shaderSource.c_str() + startOfNamePos, endOfNamePos - startOfNamePos);
53
54 // Use the shader variable's name to get the assigned location
55 auto locationIter = varNameToLocation.find(name);
56 if (locationIter == varNameToLocation.end())
57 {
58 ASSERT(false);
59 return "";
60 }
61
62 // TODO(anglebug.com/42267100): if the GLSL input is a matrix there should be multiple
63 // WGSL input variables (multiple vectors representing the columns of the matrix).
64 int location = locationIter->second;
65 std::ostringstream locationReplacementStream;
66 locationReplacementStream << "@location(" << location << ") " << name;
67
68 if (kOutputReplacements)
69 {
70 std::cout << "Replace \"" << marker << name << "\" with \""
71 << locationReplacementStream.str() << "\"" << std::endl;
72 }
73
74 // Append the new `@location(N) name` and then continue from the ` : type`.
75 newSource.append(locationReplacementStream.str());
76 currPos = endOfNamePos;
77 }
78 }
79 return newSource;
80 }
81
82 } // namespace
83
84 template <typename T>
WgslAssignLocations(const std::string & shaderSource,const std::vector<T> shaderVars,const gl::ProgramMergedVaryings & mergedVaryings,gl::ShaderType shaderType)85 std::string WgslAssignLocations(const std::string &shaderSource,
86 const std::vector<T> shaderVars,
87 const gl::ProgramMergedVaryings &mergedVaryings,
88 gl::ShaderType shaderType)
89 {
90 std::map<std::string, int> varNameToLocation;
91 for (const T &shaderVar : shaderVars)
92 {
93 if (shaderVar.isBuiltIn())
94 {
95 continue;
96 }
97 varNameToLocation[shaderVar.name] = shaderVar.getLocation();
98 }
99
100 int currLocMarker = 0;
101 for (const gl::ProgramVaryingRef &linkedVarying : mergedVaryings)
102 {
103 gl::ShaderBitSet supportedShaderStages =
104 gl::ShaderBitSet({gl::ShaderType::Vertex, gl::ShaderType::Fragment});
105 ASSERT(linkedVarying.frontShaderStage == gl::ShaderType::InvalidEnum ||
106 supportedShaderStages.test(linkedVarying.frontShaderStage));
107 ASSERT(linkedVarying.backShaderStage == gl::ShaderType::InvalidEnum ||
108 supportedShaderStages.test(linkedVarying.backShaderStage));
109 if (!linkedVarying.frontShader && !linkedVarying.backShader)
110 {
111 continue;
112 }
113 const sh::ShaderVariable *shaderVar = shaderType == gl::ShaderType::Vertex
114 ? linkedVarying.frontShader
115 : linkedVarying.backShader;
116 if (shaderVar)
117 {
118 if (shaderVar->isBuiltIn())
119 {
120 continue;
121 }
122 ASSERT(varNameToLocation.find(shaderVar->name) == varNameToLocation.end());
123 varNameToLocation[shaderVar->name] = currLocMarker++;
124 }
125 else
126 {
127 const sh::ShaderVariable *otherShaderVar = shaderType == gl::ShaderType::Vertex
128 ? linkedVarying.backShader
129 : linkedVarying.frontShader;
130 if (!otherShaderVar->isBuiltIn())
131 {
132 // Increment `currLockMarker` to keep locations in sync with the WGSL source
133 // generated for the other shader stage, which will also have incremented
134 // `currLocMarker` when seeing this variable.
135 currLocMarker++;
136 }
137 }
138 }
139
140 return WgslReplaceLocationMarkers(shaderSource, varNameToLocation);
141 }
142
143 template std::string WgslAssignLocations<gl::ProgramInput>(
144 const std::string &shaderSource,
145 const std::vector<gl::ProgramInput> shaderVars,
146 const gl::ProgramMergedVaryings &mergedVaryings,
147 gl::ShaderType shaderType);
148
149 template std::string WgslAssignLocations<gl::ProgramOutput>(
150 const std::string &shaderSource,
151 const std::vector<gl::ProgramOutput> shaderVars,
152 const gl::ProgramMergedVaryings &mergedVaryings,
153 gl::ShaderType shaderType);
154
155 } // namespace webgpu
156 } // namespace rx
157