[mlir][spirv] Add spv.GLSL.FrexpStruct

co-authored-by: Alan Liu <alanliu.yf@gmail.com>

Reviewed By: antiagainst

Differential Revision: https://reviews.llvm.org/D96527

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVGLSLOps.td

@@ -1012,4 +1012,61 @@ def SPV_GLSLFmaOp : SPV_GLSLTernaryArithmeticOp<"Fma", 50, SPV_Float> {
   }];
 }
 
+// ----
+
+def SPV_GLSLFrexpStructOp : SPV_GLSLOp<"FrexpStruct", 52, [NoSideEffect]> {
+  let summary = "Splits x into two components such that x = significand * 2^exponent";
+
+  let description = [{
+    Result is a structure containing x split into a floating-point significand
+    in the range (-1.0, 0.5] or [0.5, 1.0) and an integral exponent of 2, such that:
+
+    x = significand * 2^exponent
+
+    If x is a zero, the exponent is 0.0. If x is an infinity or a NaN, the
+    exponent is undefined. If x is 0.0, the significand is 0.0. If x is -0.0,
+    the significand is -0.0
+
+    Result Type must be an OpTypeStruct with two members. Member 0 must have
+    the same type as the type of x. Member 0 holds the significand. Member 1
+    must be a scalar or vector with integer component type, with 32-bit
+    component width. Member 1 holds the exponent. These two members and x must
+    have the same number of components.
+
+    The operand x must be a scalar or vector whose component type is
+    floating-point.
+
+    <!-- End of AutoGen section -->
+    ```
+    float-scalar-vector-type ::= float-type |
+                                 `vector<` integer-literal `x` float-type `>`
+    integer-scalar-vector-type ::= integer-type |
+                                 `vector<` integer-literal `x` integer-type `>`
+    frexpstruct-op ::= ssa-id `=` `spv.GLSL.FrexpStruct` ssa-use `:`
+                                  `!spv.struct<` float-scalar-vector-type `,`
+                                                  integer-scalar-vector-type `>`
+    ```
+    #### Example:
+
+    ```mlir
+    %2 = spv.GLSL.FrexpStruct %0 : f32 -> !spv.struct<f32, i32>
+    %3 = spv.GLSL.FrexpStruct %0 : vector<3xf32> -> !spv.struct<vector<3xf32>, vector<3xi32>>
+    ```
+  }];
+
+  let arguments = (ins
+    SPV_ScalarOrVectorOf<SPV_Float>:$operand
+  );
+
+  let results = (outs
+    SPV_AnyStruct:$result
+  );
+
+  let assemblyFormat = [{
+    attr-dict $operand `:` type($operand) `->` type($result)
+  }];
+
+  let verifier = [{ return ::verifyGLSLFrexpStructOp(*this); }];
+}
+
 #endif // MLIR_DIALECT_SPIRV_IR_GLSL_OPS

mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp

@@ -3534,6 +3534,57 @@ static LogicalResult verify(spirv::SpecConstantOperationOp constOp) {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// spv.GLSL.FrexpStruct
+//===----------------------------------------------------------------------===//
+static LogicalResult
+verifyGLSLFrexpStructOp(spirv::GLSLFrexpStructOp frexpStructOp) {
+  spirv::StructType structTy =
+      frexpStructOp.result().getType().dyn_cast<spirv::StructType>();
+
+  if (structTy.getNumElements() != 2)
+    return frexpStructOp.emitError("result type must be a struct  type "
+                                   "with two memebers");
+
+  Type significandTy = structTy.getElementType(0);
+  Type exponentTy = structTy.getElementType(1);
+  VectorType exponentVecTy = exponentTy.dyn_cast<VectorType>();
+  IntegerType exponentIntTy = exponentTy.dyn_cast<IntegerType>();
+
+  Type operandTy = frexpStructOp.operand().getType();
+  VectorType operandVecTy = operandTy.dyn_cast<VectorType>();
+  FloatType operandFTy = operandTy.dyn_cast<FloatType>();
+
+  if (significandTy != operandTy)
+    return frexpStructOp.emitError("member zero of the resulting struct type "
+                                   "must be the same type as the operand");
+
+  if (exponentVecTy) {
+    IntegerType componentIntTy =
+        exponentVecTy.getElementType().dyn_cast<IntegerType>();
+    if (!(componentIntTy && componentIntTy.getWidth() == 32))
+      return frexpStructOp.emitError(
+          "member one of the resulting struct type must"
+          "be a scalar or vector of 32 bit integer type");
+  } else if (!(exponentIntTy && exponentIntTy.getWidth() == 32)) {
+    return frexpStructOp.emitError(
+        "member one of the resulting struct type "
+        "must be a scalar or vector of 32 bit integer type");
+  }
+
+  // Check that the two member types have the same number of components
+  if (operandVecTy && exponentVecTy &&
+      (exponentVecTy.getNumElements() == operandVecTy.getNumElements()))
+    return success();
+
+  if (operandFTy && exponentIntTy)
+    return success();
+
+  return frexpStructOp.emitError(
+      "member one of the resulting struct type "
+      "must have the same number of components as the operand type");
+}
+
 namespace mlir {
 namespace spirv {
 

mlir/test/Dialect/SPIRV/IR/glsl-ops.mlir

@@ -365,3 +365,58 @@ func @fma(%a : vector<3xf32>, %b : vector<3xf32>, %c : vector<3xf32>) -> () {
   %2 = spv.GLSL.Fma %a, %b, %c : vector<3xf32>
   return
 }
+// -----
+
+//===----------------------------------------------------------------------===//
+// spv.GLSL.FrexpStruct
+//===----------------------------------------------------------------------===//
+
+func @frexp_struct(%arg0 : f32) -> () {
+  // CHECK: spv.GLSL.FrexpStruct {{%.*}} : f32 -> !spv.struct<(f32, i32)>
+  %2 = spv.GLSL.FrexpStruct %arg0 : f32 -> !spv.struct<(f32, i32)>
+  return
+}
+
+func @frexp_struct_64(%arg0 : f64) -> () {
+  // CHECK: spv.GLSL.FrexpStruct {{%.*}} : f64 -> !spv.struct<(f64, i32)>
+  %2 = spv.GLSL.FrexpStruct %arg0 : f64 -> !spv.struct<(f64, i32)>
+  return
+}
+
+func @frexp_struct_vec(%arg0 : vector<3xf32>) -> () {
+  // CHECK: spv.GLSL.FrexpStruct {{%.*}} : vector<3xf32> -> !spv.struct<(vector<3xf32>, vector<3xi32>)>
+  %2 = spv.GLSL.FrexpStruct %arg0 : vector<3xf32> -> !spv.struct<(vector<3xf32>, vector<3xi32>)>
+  return
+}
+
+// -----
+
+func @frexp_struct_mismatch_type(%arg0 : f32) -> () {
+  // expected-error @+1 {{member zero of the resulting struct type must be the same type as the operand}}
+  %2 = spv.GLSL.FrexpStruct %arg0 : f32 -> !spv.struct<(vector<3xf32>, i32)>
+  return
+}
+
+// -----
+
+func @frexp_struct_wrong_type(%arg0 : i32) -> () {
+  // expected-error @+1 {{op operand #0 must be 16/32/64-bit float or vector of 16/32/64-bit float values}}
+  %2 = spv.GLSL.FrexpStruct %arg0 : i32 -> !spv.struct<(i32, i32)>
+  return
+}
+
+// -----
+
+func @frexp_struct_mismatch_num_components(%arg0 : vector<3xf32>) -> () {
+  // expected-error @+1 {{member one of the resulting struct type must have the same number of components as the operand type}}
+  %2 = spv.GLSL.FrexpStruct %arg0 : vector<3xf32> -> !spv.struct<(vector<3xf32>, vector<2xi32>)>
+  return
+}
+
+// -----
+
+func @frexp_struct_not_i32(%arg0 : f32) -> () {
+  // expected-error @+1 {{member one of the resulting struct type must be a scalar or vector of 32 bit integer type}}
+  %2 = spv.GLSL.FrexpStruct %arg0 : f32 -> !spv.struct<(f32, i64)>
+  return
+}

mlir/test/Target/SPIRV/glsl-ops.mlir

@@ -28,6 +28,8 @@ spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
     %11 = spv.GLSL.Pow %arg0, %arg1 : f32
     // CHECK: {{%.*}} = spv.GLSL.Round {{%.*}} : f32
     %12 = spv.GLSL.Round %arg0 : f32
+    // CHECK: {{%.*}} = spv.GLSL.FrexpStruct {{%.*}} : f32 -> !spv.struct<(f32, i32)>
+    %13 = spv.GLSL.FrexpStruct %arg0 : f32 -> !spv.struct<(f32, i32)>
     spv.Return
   }