[mlir][spirv] Add support for FrexpStruct in GLSL extended instruction set

[Weiwei-2021]

Summary

add support for FrexpStruct in glsl extended instruction set.

co-authored-by: Alan Liu <[e-mail]>

FYI, I tried to use define_inst.sh to autogen FrexpStruct in this branch but cannot make work. It seems that this script only supports spirv core instructions. Alan did this part before, I will try to sync up with him on how he makes it. But whether wonder anyone else know how to do it.

[ergawy]

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

1021

Look like the description is taken from an old revision of the spec. Given that Frexp is not supported here (and deprecated), it would be great to use the latest revision: https://www.khronos.org/registry/spir-v/specs/unified1/GLSL.std.450.html

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

3555

nit: I think you mean to convey the negation of this message, something like:

(loc, "expected result structure type to consist of two members, but provided") << type;

3576

We need to make sure the struct has 2 members before accessing them.

3576

nit: Here and below, can we change the names from: memberZero... to significand... and memberOne... to exponent...?

3585

nit: No need for the comment, the condition and error message are both quite clear.

3594

nit: The convention is not to use {} if there is a single statement in the scope. I consistently forget this too :D.

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

394

I think this test exercises 2 different error paths at once: (1) the argument is not a floating point type, and (2) the arity of the argument and the second struct member must match (i.e. either both scalars or both vectors). Can you split this to 2 different cases as such?

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

31

One more test for the vector case would also be great.

[antiagainst]

Nice, thanks for working on this!

FYI, I tried to use define_inst.sh to autogen FrexpStruct in this branch but cannot make work. It seems that this script only supports spirv core instructions. Alan did this part before, I will try to sync up with him on how he makes it. But whether wonder anyone else know how to do it.

Yeah I think the script is not extended to support GLSL ops yet. Would be great to make it work there. But it's also fine for me to just manually add for now given we have been doing that for a while.

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

1021

+1

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

3540

We could use

let assemblyFormat = "attr-dict $operand `:` type($operand) `->` type($result)

in the ODS to autogenerate the parser and printer.

Here I see you are validating the result type. There is no need to do it here: validation should be deferred to the validator where possible, which is auto partially autogenerated by the ODS spec (like SPV_AnyStruct would guarantee that the result type is a struct). In general, only validation needed for parsing itself should be checked in the parser.

3551

s/SPIRV/SPIR-V/

SPIR-V is the official name. We cannot use it in the code but in comment or strings we should use the official name. :)

3573

No need to duplicate the check here given this is already guaranteed by SPV_AnyStruct. The manual validation only needs to cover what's not covered by ODS.

3587

Could you add a test case for this error?

3588

... must be the same type ...

[Weiwei-2021]

I tried to use define_inst.sh to create autogen GLSL

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

1021

Thanks for clarifying the generation of glsl extension. I was confused by the fact that all ops containing "<!-- End of AutoGen section-->". We added this instruction months ago, and I didn't notice the out-of-date problem. I will update with the latest version.

A side question is that if knronos update the spirv version's, those documents needs to get update or we haven't done things like this?

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

3540

Thanks. Remove parse and print. Add in ODS instead.

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

394

Actually it is just a lazy copy and past. :)

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

31

Seems that this whole test only tests scalar type although those operations support vector. Let us leave it like this considering that the test in Dialect is checking the validation of vector types.

[antiagainst]

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

1021

A side question is that if knronos update the spirv version's, those documents needs to get update or we haven't done things like this?

Yeah. But I won't worry that too much. The spec is not updated very frequently. Changes will be reflected when we add new ops. There is a delay there but it's fine to me.

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

let description = [{
 
     ```mlir
     %0 = spv.GLSL.Fma %a, %b, %c : f32
     %1 = spv.GLSL.Fma %a, %b, %c : vector<3xf16>
     ```
   }];
 }
 
+// ----
+
+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

for (auto operand : enclosedOp.getOperands())
     if (!isa<spirv::ConstantOp, spirv::ReferenceOfOp,
              spirv::SpecConstantOperationOp>(operand.getDefiningOp()))
       return constOp.emitOpError(
           "invalid operand, must be defined by a constant operation");
 
   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 {
 
 // TableGen'erated operation interfaces for querying versions, extensions, and
 // capabilities.
 #include "mlir/Dialect/SPIRV/IR/SPIRVAvailability.cpp.inc"
 } // namespace spirv
 } // namespace mlir
 
 // TablenGen'erated operation definitions.
 #define GET_OP_CLASSES
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.cpp.inc"
 
 namespace mlir {
 namespace spirv {
 // TableGen'erated operation availability interface implementations.
 #include "mlir/Dialect/SPIRV/IR/SPIRVOpAvailabilityImpl.inc"
 
 } // namespace spirv
 } // namespace mlir

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

 
 // -----
 
 func @fma(%a : vector<3xf32>, %b : vector<3xf32>, %c : vector<3xf32>) -> () {
   // CHECK: spv.GLSL.Fma {{%[^,]*}}, {{%[^,]*}}, {{%[^,]*}} : 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

spv.func @fmul(%arg0 : f32, %arg1 : f32) "None" {
     // CHECK: {{%.*}} = spv.GLSL.Sinh {{%.*}} : f32
     %9 = spv.GLSL.Sinh %arg0 : f32
     // CHECK: {{%.*}} = spv.GLSL.Cosh {{%.*}} : f32
     %10 = spv.GLSL.Cosh %arg0 : f32
     // CHECK: {{%.*}} = spv.GLSL.Pow {{%.*}} : f32
     %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
   }
 
   spv.func @fclamp(%arg0 : f32, %arg1 : f32, %arg2 : f32) "None" {
     // CHECK: spv.GLSL.FClamp {{%[^,]*}}, {{%[^,]*}}, {{%[^,]*}} : f32
     %13 = spv.GLSL.FClamp %arg0, %arg1, %arg2 : f32
     spv.Return
   }