Implement dot

wpmed92 · Sep 19, 2024 · 4afa5f8 · 4afa5f8
1 parent 38534f6
commit 4afa5f8
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git 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
@@ -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);
 }