[WGSL] Initial support for override default values

https://bugs.webkit.org/show_bug.cgi?id=262622 rdar://116466820 Reviewed by Mike Wyrzykowski. Pass the initializer expression to the API, which can later be evaluated with the evaluation function introduced in 268833@main. * Source/WebGPU/WGSL/GlobalVariableRewriter.cpp: (WGSL::RewriteGlobalVariables::usesOverride): * Source/WebGPU/WGSL/WGSL.h: * Source/WebGPU/WebGPU/Pipeline.mm: (WebGPU::createConstantValues): Canonical link: https://commits.webkit.org/268905@main
WebKit · Oct 5, 2023 · aba1459 · aba1459
1 parent e8585e2
commit aba1459
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 5 deletions.
diff --git a/Source/WebGPU/WGSL/GlobalVariableRewriter.cpp b/Source/WebGPU/WGSL/GlobalVariableRewriter.cpp
@@ -847,7 +847,7 @@ void RewriteGlobalVariables::usesOverride(AST::Variable& variable)
         }
     }
 
-    m_entryPointInformation->specializationConstants.add(originalName, Reflection::SpecializationConstant { variable.name(), constantType });
+    m_entryPointInformation->specializationConstants.add(originalName, Reflection::SpecializationConstant { variable.name(), constantType, variable.maybeInitializer() });
 }
 
 void RewriteGlobalVariables::insertStructs(const UsedResources& usedResources)

diff --git a/Source/WebGPU/WGSL/WGSL.h b/Source/WebGPU/WGSL/WGSL.h
@@ -204,6 +204,7 @@ enum class SpecializationConstantType : uint8_t {
 struct SpecializationConstant {
     String mangledName;
     SpecializationConstantType type;
+    AST::Expression* defaultValue;
 };
 
 struct EntryPointInformation {

diff --git a/Source/WebGPU/WebGPU/Pipeline.mm b/Source/WebGPU/WebGPU/Pipeline.mm
@@ -68,6 +68,39 @@
         return { };
 
     auto constantValues = [MTLFunctionConstantValues new];
+    for (auto& kvp : entryPointInformation.specializationConstants) {
+        auto& specializationConstant = kvp.value;
+        if (!specializationConstant.defaultValue)
+            continue;
+
+        auto constantValue = WGSL::evaluate(*kvp.value.defaultValue, wgslConstantValues);
+        auto addResult = wgslConstantValues.add(specializationConstant.mangledName, constantValue);
+        ASSERT_UNUSED(addResult, addResult.isNewEntry);
+
+        switch (specializationConstant.type) {
+        case WGSL::Reflection::SpecializationConstantType::Boolean: {
+            bool value = constantValue.toBool();
+            [constantValues setConstantValue:&value type:MTLDataTypeBool withName:specializationConstant.mangledName];
+            break;
+        }
+        case WGSL::Reflection::SpecializationConstantType::Float: {
+            float value = constantValue.toDouble();
+            [constantValues setConstantValue:&value type:MTLDataTypeFloat withName:specializationConstant.mangledName];
+            break;
+        }
+        case WGSL::Reflection::SpecializationConstantType::Int: {
+            int value = constantValue.toInt();
+            [constantValues setConstantValue:&value type:MTLDataTypeInt withName:specializationConstant.mangledName];
+            break;
+        }
+        case WGSL::Reflection::SpecializationConstantType::Unsigned: {
+            unsigned value = constantValue.toInt();
+            [constantValues setConstantValue:&value type:MTLDataTypeUInt withName:specializationConstant.mangledName];
+            break;
+        }
+        }
+    }
+
     for (uint32_t i = 0; i < constantCount; ++i) {
         const auto& entry = constants[i];
         auto indexIterator = entryPointInformation.specializationConstants.find(fromAPI(entry.key));
@@ -77,25 +110,25 @@
         switch (specializationConstant.type) {
         case WGSL::Reflection::SpecializationConstantType::Boolean: {
             bool value = entry.value;
-            wgslConstantValues.add(fromAPI(entry.key), value);
+            wgslConstantValues.set(fromAPI(entry.key), value);
             [constantValues setConstantValue:&value type:MTLDataTypeBool withName:specializationConstant.mangledName];
             break;
         }
         case WGSL::Reflection::SpecializationConstantType::Float: {
             float value = entry.value;
-            wgslConstantValues.add(fromAPI(entry.key), value);
+            wgslConstantValues.set(fromAPI(entry.key), value);
             [constantValues setConstantValue:&value type:MTLDataTypeFloat withName:specializationConstant.mangledName];
             break;
         }
         case WGSL::Reflection::SpecializationConstantType::Int: {
             int value = entry.value;
-            wgslConstantValues.add(fromAPI(entry.key), value);
+            wgslConstantValues.set(fromAPI(entry.key), value);
             [constantValues setConstantValue:&value type:MTLDataTypeInt withName:specializationConstant.mangledName];
             break;
         }
         case WGSL::Reflection::SpecializationConstantType::Unsigned: {
             unsigned value = entry.value;
-            wgslConstantValues.add(fromAPI(entry.key), value);
+            wgslConstantValues.set(fromAPI(entry.key), value);
             [constantValues setConstantValue:&value type:MTLDataTypeUInt withName:specializationConstant.mangledName];
             break;
         }