[spirv] Adding sqrt op in the GLSL extension.

PiperOrigin-RevId: 283769736

diff --git a/mlir/include/mlir/Dialect/SPIRV/SPIRVGLSLOps.td b/mlir/include/mlir/Dialect/SPIRV/SPIRVGLSLOps.td
index 217b6d9..2a1e8f3 100644 (file)
--- a/mlir/include/mlir/Dialect/SPIRV/SPIRVGLSLOps.td
+++ b/mlir/include/mlir/Dialect/SPIRV/SPIRVGLSLOps.td
@@ -516,6 +516,36 @@ def SPV_GLSLSSignOp : SPV_GLSLUnaryArithmeticOp<"SSign", 7, SPV_Integer> {
 
 // -----
 
+def SPV_GLSLSqrtOp : SPV_GLSLUnaryArithmeticOp<"Sqrt", 31, SPV_Float> {
+  let summary = "Returns the square root of the operand";
+
+  let description = [{
+    Result is the square root of x. Result is undefined if x < 0.
+
+    The operand x must be a scalar or vector whose component type is
+    floating-point.
+
+    Result Type and the type of x must be the same type. Results are computed
+    per component.
+
+    ### Custom assembly format
+    ``` {.ebnf}
+    float-scalar-vector-type ::= float-type |
+                                 `vector<` integer-literal `x` float-type `>`
+    sqrt-op ::= ssa-id `=` `spv.GLSL.Sqrt` ssa-use `:`
+                float-scalar-vector-type
+    ```
+    For example:
+
+    ```
+    %2 = spv.GLSL.Sqrt %0 : f32
+    %3 = spv.GLSL.Sqrt %1 : vector<3xf16>
+    ```
+  }];
+}
+
+// -----
+
 def SPV_GLSLTanhOp : SPV_GLSLUnaryArithmeticOp<"Tanh", 21, SPV_Float16or32> {
   let summary = "Hyperbolic tangent of operand in radians";
 

diff --git a/mlir/test/Dialect/SPIRV/Serialization/glsl-ops.mlir b/mlir/test/Dialect/SPIRV/Serialization/glsl-ops.mlir
index adc87bd..2ebe188 100644 (file)
--- a/mlir/test/Dialect/SPIRV/Serialization/glsl-ops.mlir
+++ b/mlir/test/Dialect/SPIRV/Serialization/glsl-ops.mlir
@@ -6,6 +6,8 @@ spv.module "Logical" "GLSL450" {
     %0 = spv.GLSL.Exp %arg0 : f32
     // CHECK: {{%.*}} = spv.GLSL.FMax {{%.*}}, {{%.*}} : f32
     %1 = spv.GLSL.FMax %arg0, %arg1 : f32
+    // CHECK: {{%.*}} = spv.GLSL.Sqrt {{%.*}} : f32
+    %2 = spv.GLSL.Sqrt %arg0 : f32
     spv.Return
   }
-}
\ No newline at end of file
+}

diff --git a/mlir/test/Dialect/SPIRV/glslops.mlir b/mlir/test/Dialect/SPIRV/glslops.mlir
index 5b78b52..e3a3ca5 100644 (file)
--- a/mlir/test/Dialect/SPIRV/glslops.mlir
+++ b/mlir/test/Dialect/SPIRV/glslops.mlir
@@ -71,3 +71,39 @@ func @fmaxf64(%arg0 : f64, %arg1 : f64) -> () {
   %2 = spv.GLSL.FMax %arg0, %arg1 : f64
   return
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spv.GLSL.InverseSqrt
+//===----------------------------------------------------------------------===//
+
+func @inversesqrt(%arg0 : f32) -> () {
+  // CHECK: spv.GLSL.InverseSqrt {{%.*}} : f32
+  %2 = spv.GLSL.InverseSqrt %arg0 : f32
+  return
+}
+
+func @inversesqrtvec(%arg0 : vector<3xf16>) -> () {
+  // CHECK: spv.GLSL.InverseSqrt {{%.*}} : vector<3xf16>
+  %2 = spv.GLSL.InverseSqrt %arg0 : vector<3xf16>
+  return
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spv.GLSL.Sqrt
+//===----------------------------------------------------------------------===//
+
+func @sqrt(%arg0 : f32) -> () {
+  // CHECK: spv.GLSL.Sqrt {{%.*}} : f32
+  %2 = spv.GLSL.Sqrt %arg0 : f32
+  return
+}
+
+func @sqrtvec(%arg0 : vector<3xf16>) -> () {
+  // CHECK: spv.GLSL.Sqrt {{%.*}} : vector<3xf16>
+  %2 = spv.GLSL.Sqrt %arg0 : vector<3xf16>
+  return
+}