From 4afa5f896b07c8e72562d0e3a4a3394e2e8a78e2 Mon Sep 17 00:00:00 2001
From: Ahmed Harmouche <ahmedharmouche92@gmail.com>
Date: Thu, 19 Sep 2024 08:55:48 +0200
Subject: [PATCH] Implement dot

---
 lib/CodeGen/MLIRCodeGen.cpp         | 4 ++++
 test/CodeGen/functions_builtin.glsl | 9 +++++++++
 2 files changed, 13 insertions(+)

diff --git a/lib/CodeGen/MLIRCodeGen.cpp b/lib/CodeGen/MLIRCodeGen.cpp
index 3ddf309..2a2264f 100644
--- a/lib/CodeGen/MLIRCodeGen.cpp
+++ b/lib/CodeGen/MLIRCodeGen.cpp
@@ -73,6 +73,10 @@ void MLIRCodeGen::initBuiltinFuncMap() {
     {"cosh", [](mlir::MLIRContext &context, mlir::OpBuilder &builder, mlir::ValueRange operands) {
       return builder.create<spirv::GLCoshOp>(builder.getUnknownLoc(), operands[0]);
     }},
+    {"dot", [](mlir::MLIRContext &context, mlir::OpBuilder &builder, mlir::ValueRange operands) {
+      auto op0ElementType = operands[0].getType().dyn_cast<mlir::VectorType>().getElementType();
+      return builder.create<spirv::DotOp>(builder.getUnknownLoc(), op0ElementType, operands[0], operands[1]);
+    }},
     {"exp", [](mlir::MLIRContext &context, mlir::OpBuilder &builder, mlir::ValueRange operands) {
       return builder.create<spirv::GLExpOp>(builder.getUnknownLoc(), operands[0]);
     }},
diff --git a/test/CodeGen/functions_builtin.glsl b/test/CodeGen/functions_builtin.glsl
index b9527b1..18519de 100644
--- a/test/CodeGen/functions_builtin.glsl
+++ b/test/CodeGen/functions_builtin.glsl
@@ -100,4 +100,13 @@ void main() {
 
     // CHECK: %34 = spirv.GL.SSign %33 : si32
     b = sign(-1);
+
+    vec3 dotTestA = vec3(1.0, 0.0, 1.2);
+    dvec3 dotTestB = dvec3(0.1lf, 0.2lf, 0.3lf);
+
+    // CHECK: %41 = spirv.Dot %39, %40 : vector<3xf32> -> f32
+    float floatDotRes = dot(dotTestA, dotTestA);
+
+    // CHECK: %45 = spirv.Dot %43, %44 : vector<3xf64> -> f64
+    double doubleDotRes = dot(dotTestB, dotTestB);
 }