[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 = [{
+    Same semantics as in Frexp, except that the entire result is in the
+    instruction’s result; there is not a pointer operand to write through.
+
+    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 exponent. These two members
+    must have the same number of components. This structure type must be
+    explicitly declared by the module.
+
+    <!-- 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 : !spv.struct<f32, i32>
+    %3 = spv.GLSL.FrexpStruct %0 : !spv.struct<vector<3xf32>, vector<3xi32>>
+    ```
+  }];
+
+  let arguments = (ins
+    SPV_ScalarOrVectorOf<SPV_Float>:$operand
+  );
+
+  let results = (outs
+    SPV_AnyStruct:$result
+  );
+
+  let parser = [{ return ::parseGLSLFrexpStructOp(parser, result); }];
+
+  let printer = [{ return ::printGLSLFrexpStructOp(*this, p); }];
+
+  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 ParseResult parseGLSLFrexpStructOp(OpAsmParser &parser,
+                                          OperationState &state) {
+  OpAsmParser::OperandType operand;
+  Type type;
+  auto loc = parser.getCurrentLocation();
+  if (parser.parseOperand(operand) || parser.parseColonType(type))
+    return failure();
+
+  spirv::StructType resultType = type.dyn_cast<spirv::StructType>();
+  if (!resultType)
+    return parser.emitError(
+               loc, "result type must be a SPIRV structure type, but provided ")
+           << type;
+
+  if (resultType.getNumElements() != 2)
+    return parser.emitError(loc, "result structure type with two members");
+
+  state.addTypes(type);
+  return parser.resolveOperand(operand, resultType.getElementType(0),
+                               state.operands);
+}
+
+static void printGLSLFrexpStructOp(spirv::GLSLFrexpStructOp frexpStructOp,
+                                   OpAsmPrinter &printer) {
+  printer << spirv::GLSLFrexpStructOp::getOperationName() << " "
+          << frexpStructOp.operand() << " : "
+          << frexpStructOp.result().getType();
+}
+
+static LogicalResult
+verifyGLSLFrexpStructOp(spirv::GLSLFrexpStructOp frexpStructOp) {
+  spirv::StructType structTy =
+      frexpStructOp.result().getType().dyn_cast<spirv::StructType>();
+  if (!structTy)
+    return frexpStructOp.emitError("result type must be struct type");
+
+  Type memberZeroTy = structTy.getElementType(0);
+  Type memberOneTy = structTy.getElementType(1);
+  VectorType memberOneVecTy = memberOneTy.dyn_cast<VectorType>();
+  IntegerType memberOneIntTy = memberOneTy.dyn_cast<IntegerType>();
+
+  Type operandTy = frexpStructOp.operand().getType();
+  VectorType operandVecTy = operandTy.dyn_cast<VectorType>();
+  FloatType operandFTy = operandTy.dyn_cast<FloatType>();
+
+  // Check that member zero has the same type as the operand
+  if (memberZeroTy != operandTy)
+    return frexpStructOp.emitError("member zero of the resulting struct type "
+                                   "must be of the same type as the operand");
+
+  // Check that member one must be scalar or vector of 32bit integer type
+  if (memberOneVecTy) {
+    IntegerType componentIntTy =
+        memberOneVecTy.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 (!(memberOneIntTy && memberOneIntTy.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 && memberOneVecTy &&
+      (memberOneVecTy.getNumElements() == operandVecTy.getNumElements()))
+    return success();
+
+  if (operandFTy && memberOneIntTy)
+    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

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 {{%.*}} : !spv.struct<(f32, i32)>
+  %2 = spv.GLSL.FrexpStruct %arg0 : !spv.struct<(f32, i32)>
+  return
+}
+
+func @frexp_struct_64(%arg0 : f64) -> () {
+  // CHECK: spv.GLSL.FrexpStruct {{%.*}} : !spv.struct<(f64, i32)>
+  %2 = spv.GLSL.FrexpStruct %arg0 : !spv.struct<(f64, i32)>
+  return
+}
+
+func @frexp_struct_vec(%arg0 : vector<3xf32>) -> () {
+  // CHECK: spv.GLSL.FrexpStruct {{%.*}} : !spv.struct<(vector<3xf32>, vector<3xi32>)>
+  %2 = spv.GLSL.FrexpStruct %arg0 : !spv.struct<(vector<3xf32>, vector<3xi32>)>
+  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 : !spv.struct<(i32, vector<3xi32>)>
+  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 : !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 : !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 {{%.*}} : !spv.struct<(f32, i32)>
+    %13 = spv.GLSL.FrexpStruct %arg0 : !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
   }