[WGSL] Implement arithmetic operator overload and codegen

https://bugs.webkit.org/show_bug.cgi?id=255323 rdar://problem/107923724 Reviewed by Tadeu Zagallo. Implement overload resolution and codegen for +, -, *, /, %. * Source/WebGPU/WGSL/Metal/MetalFunctionWriter.cpp: (WGSL::Metal::FunctionDefinitionWriter::visit): * Source/WebGPU/WGSL/TypeDeclarations.rb: * Source/WebGPU/WGSL/generator/main.rb: * Source/WebGPU/WGSL/tests/valid/overload.wgsl: Canonical link: https://commits.webkit.org/262908@main
WebKit · Apr 13, 2023 · 9dcc43f · 9dcc43f
1 parent 2dd8ae9
commit 9dcc43f
Show file tree

Hide file tree

Showing 4 changed files with 75 additions and 13 deletions.
diff --git a/Source/WebGPU/WGSL/Metal/MetalFunctionWriter.cpp b/Source/WebGPU/WGSL/Metal/MetalFunctionWriter.cpp
@@ -511,13 +511,12 @@ void FunctionDefinitionWriter::visit(AST::BinaryExpression& binary)
     case AST::BinaryOperation::Multiply:
         m_stringBuilder.append(" * ");
         break;
-
     case AST::BinaryOperation::Divide:
+        m_stringBuilder.append(" / ");
+        break;
     case AST::BinaryOperation::Modulo:
-        // FIXME: Implement these
-        RELEASE_ASSERT_NOT_REACHED();
+        m_stringBuilder.append(" % ");
         break;
-
     case AST::BinaryOperation::And:
         m_stringBuilder.append(" & ");
         break;

diff --git a/Source/WebGPU/WGSL/TypeDeclarations.rb b/Source/WebGPU/WGSL/TypeDeclarations.rb
@@ -1,13 +1,31 @@
 # FIXME: add all the missing type declarations here
 operator :+, {
     [T < Number].(T, T) => T,
+
+    # vector addition
     [T < Number, N].(Vector[T, N], T) => Vector[T, N],
     [T < Number, N].(T, Vector[T, N]) => Vector[T, N],
     [T < Number, N].(Vector[T, N], Vector[T, N]) => Vector[T, N],
+
+    # matrix-matrix addition
     [T < Float, C, R].(Matrix[T, C, R], Matrix[T, C, R]) => Matrix[T, C, R],
 }
 
+operator :-, {
+    # unary
+    [T < SignedNumber].(T) => T,
+    [T < SignedNumber, N].(Vector[T, N]) => Vector[T, N],
+
+    # binary
+    [T < Number].(T, T) => T,
+    [T < Number, N].(Vector[T, N], T) => Vector[T, N],
+    [T < Number, N].(T, Vector[T, N]) => Vector[T, N],
+    [T < Number, N].(Vector[T, N], Vector[T, N]) => Vector[T, N],
+}
+
 operator :*, {
+    [T < Number].(T, T) => T,
+
     # vector scaling
     [T < Number, N].(Vector[T, N], T) => Vector[T, N],
     [T < Number, N].(T, Vector[T, N]) => Vector[T, N],
@@ -16,17 +34,27 @@
     # matrix-vector multiplication
     [T < Float, C, R].(Matrix[T, C, R], Vector[T, C]) => Vector[T, R],
     [T < Float, C, R].(Vector[T, R], Matrix[T, C, R]) => Vector[T, C],
+
+    # matrix-matrix multiplication
+    [T < Float, C, R, K].(Matrix[T, K, R], Matrix[T, C, K]) => Matrix[T, C, R],
 }
 
-operator :-, {
-  # unary
-  [T < SignedNumber].(T) => T,
-  [T < SignedNumber, N].(Vector[T, N]) => Vector[T, N],
-
-  # binary
-  [T < Number, N].(Vector[T, N], T) => Vector[T, N],
-  [T < Number, N].(T, Vector[T, N]) => Vector[T, N],
-  [T < Number, N].(Vector[T, N], Vector[T, N]) => Vector[T, N],
+operator :/, {
+    [T < Number].(T, T) => T,
+
+    # vector scaling
+    [T < Number, N].(Vector[T, N], T) => Vector[T, N],
+    [T < Number, N].(T, Vector[T, N]) => Vector[T, N],
+    [T < Number, N].(Vector[T, N], Vector[T, N]) => Vector[T, N],
+}
+
+operator :%, {
+    [T < Number].(T, T) => T,
+
+    # vector scaling
+    [T < Number, N].(Vector[T, N], T) => Vector[T, N],
+    [T < Number, N].(T, Vector[T, N]) => Vector[T, N],
+    [T < Number, N].(Vector[T, N], Vector[T, N]) => Vector[T, N],
 }
 
 # Comparison operations

diff --git a/Source/WebGPU/WGSL/generator/main.rb b/Source/WebGPU/WGSL/generator/main.rb
@@ -286,6 +286,7 @@ def self.prologue
         N = Variable.new(:N, @NumericVariable)
         C = Variable.new(:C, @NumericVariable)
         R = Variable.new(:R, @NumericVariable)
+        K = Variable.new(:K, @NumericVariable)
 
         Number = Constraint.new(:Number)
         Integer = Constraint.new(:Integer)

diff --git a/Source/WebGPU/WGSL/tests/valid/overload.wgsl b/Source/WebGPU/WGSL/tests/valid/overload.wgsl
@@ -37,6 +37,14 @@ fn testAdd() {
 }
 
 fn testMultiply() {
+  {
+    let x = 0 * 0;
+    let x1 = 0i * 0i;
+    let x2 = 0u * 0u;
+    let x3 = 0.0 * 0.0;
+    let x4 = 0.0f * 0.0f;
+  }
+
   let v2 = vec2<f32>(0, 0);
   let v4 = vec4<f32>(0, 0, 0, 0);
   let m = mat2x4<f32>(0, 0, 0, 0, 0, 0, 0, 0);
@@ -45,6 +53,32 @@ fn testMultiply() {
   let r3 = vec2(1, 1) * 1;
   let r4 = 1 * vec2(1, 1);
   let r5 = vec2(1, 1) * vec2(1, 1);
+
+  let x0 = mat2x2(0, 0, 0, 0) * mat2x2(0, 0, 0, 0);
+  let x1 = mat2x2(0, 0, 0, 0) * mat3x2(0, 0, 0, 0, 0, 0);
+  let x2 = mat2x2(0, 0, 0, 0) * mat4x2(0, 0, 0, 0, 0, 0, 0, 0);
+}
+
+fn testDivision() {
+   let x = 0 / 1;
+   let x1 = 0i / 1i;
+   let x2 = 0u / 1u;
+   let x3 = 0.0 / 1.0;
+   let x4 = 0.0f / 1.0f;
+   let x5 = vec2(0.0) / 1.0;
+   let x6 = 0.0 / vec2(1.0, 1.0);
+   let x7 = vec2(0.0, 0.0) / vec2(1.0, 1.0);
+}
+
+fn testModulo() {
+   let x = 0 % 1;
+   let x1 = 0i % 1i;
+   let x2 = 0u % 1u;
+   let x3 = 0.0 % 1.0;
+   let x4 = 0.0f % 1.0f;
+   let x5 = vec2(0.0) % 1.0;
+   let x6 = 0.0 % vec2(1.0, 1.0);
+   let x7 = vec2(0.0, 0.0) % vec2(1.0, 1.0);
 }
 
 fn testTextureSample() {