[mlir][spirv] Add support for Ldexp operation in glsl extended instruction set

[Weiwei-2021]

Summary

Add support for Ldexp operation in glsl extended instruction set for SPIR-V dialect.

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

[Weiwei-2021]

Forget to mention it in the commit msg (too quick to click the button:)). I am not 100% sure how to refer an operand in the error msg (e.g., x, exp in this case). Therefore, I keep the words the same as the SPIR-V Extended Instruction Set Spec.

[antiagainst]

Nice, thanks for the contribution!

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

1102

This does not match with the assembly format. Actually with assembly format defined you can omit this grammar.

1108

These need to be updated to match the assembly format.

1119

This should be named as y to match with the description?

1123

Could you define it as

attr-dict $x  `:` type($x) `,` $exp `:` type($exp) `->` type($result)

so that the type is directly associated with the operand? This makes it easy to read. (Right now that the majority of the ops aren't defined this way but I'm considering to refactoring them to use this scheme.)

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

3593

Nit: just use Type instead of Ty. Two more characters makes things more readable and I would typically expect tab to complete it so you won't need to even type those two characters for most of the cases. :)

3597

This can be autogenerated by adding AllTypesMatch<["x", "result"]> to the trait. Typically we want to autogen when possible.

3601

What about

if (significantType.isa<FloatType>() != exponentType.isa<IntegerType>())
  return ldexpOp.emitOpError("operands must both be scalars or vectors");

auto getNumElements = [](Type type) -> unsigned {
  if (auto vectorType = type.dyn_cast<VectorType>())
    return vectorType.getNumElements;
  return 1;
};

if (getNumElements(...) != getNumElements(...))
  return ldexpOp.emitOpError("operands must have the same number of elements");

?

[antiagainst]

In D97228#2580084, @Weiwei-2021 wrote:

Forget to mention it in the commit msg (too quick to click the button:)). I am not 100% sure how to refer an operand in the error msg (e.g., x, exp in this case). Therefore, I keep the words the same as the SPIR-V Extended Instruction Set Spec.

Yeah using the symbols defined in ODS is the way to go.

[Weiwei-2021]

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

1102

good to know. Thanks!

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

3597

Thanks and good to know this tip. Do we have any document for things like this? My main reference is this link: https://mlir.llvm.org/docs/OpDefinitions/, but it is not that directly to get tips like this.

3601

LGTM, thanks for writing it.

[antiagainst]

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

3597

I typically directly read OpBase.td to discover them. Don't worry about not finding every one of them for now. You'll learn gradually. Code reviews are also for this purpose. :)

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

def SPV_GLSLFrexpStructOp : SPV_GLSLOp<"FrexpStruct", 52, [NoSideEffect]> {
 
   let assemblyFormat = [{
     attr-dict $operand `:` type($operand) `->` type($result)
   }];
 
   let verifier = [{ return ::verifyGLSLFrexpStructOp(*this); }];
 }
 
+def SPV_GLSLLdexpOp : SPV_GLSLOp<"Ldexp", 53, [NoSideEffect]> {
+  let summary = "Builds y such that y = significand * 2^exponent";
+
+  let description = [{
+    Builds a floating-point number from x and the corresponding
+    integral exponent of two in exp:
+
+    significand * 2^exponent
+
+    If this product is too large to be represented in the floating-point
+    type, the resulting value is undefined. If exp is greater than +128
+    (single precision) or +1024 (double precision), the resulting value is
+    undefined. If exp is less than -126 (single precision) or -1022 (double precision),
+    the result may be flushed to zero. Additionally, splitting the value
+    into a significand and exponent using frexp and then reconstructing a
+    floating-point value using ldexp should yield the original input for
+    zero and all finite non-denormalized values.
+
+    The operand x must be a scalar or vector whose component type is floating-point.
+
+    The exp operand must be a scalar or vector with integer component type.
+    The number of components in x and exp must be the same.
+
+    Result Type must be the same type as the type of x. Results are computed per
+    component.
+
+    <!-- End of AutoGen section -->
+    ```
+    float-scalar-vector-type ::= float-type |
+                                 `vector<` integer-literal `x` float-type `>`
+    ldexp-op ::= ssa-id `=` `spv.GLSL.Ldexp` ssa-use `:`
+                              float-scalar-vector-type `,` integer-type
+    ```
+    #### Example:
+
+    ```mlir
+    %2 = spv.GLSL.Ldexp %0, %1 : f32, i32 -> f32
+    %3 = spv.GLSL.Ldexp %0, %1 : vector<3xf32>, vector<3xi32> -> vector<3xf32>
+    ```
+  }];
+
+  let arguments = (ins
+    SPV_ScalarOrVectorOf<SPV_Float>:$x,
+    SPV_ScalarOrVectorOf<SPV_Integer>:$exp
+  );
+
+  let results = (outs
+    SPV_ScalarOrVectorOf<SPV_Float>:$result
+  );
+
+  let assemblyFormat = [{
+    attr-dict $x `,` $exp `:` type($x) `,` type($exp) `->` type($result)
+  }];
+
+  let verifier = [{ return ::verify(*this); }];
+}
+
+
 #endif // MLIR_DIALECT_SPIRV_IR_GLSL_OPS

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

verifyGLSLFrexpStructOp(spirv::GLSLFrexpStructOp frexpStructOp) {
   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");
 }
 
+//===----------------------------------------------------------------------===//
+// spv.GLSL.Ldexp
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verify(spirv::GLSLLdexpOp ldexpOp) {
+  Type resultTy = ldexpOp.result().getType();
+  Type significandTy = ldexpOp.x().getType();
+  Type exponentTy = ldexpOp.exp().getType();
+
+  if (resultTy != significandTy)
+    return ldexpOp.emitError(
+        "result type must be the same type as the type of the x operand");
+
+  if (significandTy.isa<FloatType>()) {
+    if (!exponentTy.isa<IntegerType>())
+      return ldexpOp.emitError("the number of components in x and exp "
+                               "must be the same");
+  } else if (VectorType significandVecTy =
+                 significandTy.dyn_cast<VectorType>()) {
+    if (!exponentTy.isa<VectorType>())
+      return ldexpOp.emitError("the number of components in x and exp "
+                               "must be the same");
+
+    VectorType exponentVecTy = exponentTy.cast<VectorType>();
+    if (exponentVecTy.getNumElements() != significandVecTy.getNumElements())
+      return ldexpOp.emitError("the number of components in x and exp "
+                               "must be the same");
+  }
+  return success();
+}
+
 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 @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
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spv.GLSL.Ldexp
+//===----------------------------------------------------------------------===//
+
+func @ldexp(%arg0 : f32, %arg1 : i32) -> () {
+  // CHECK: {{%.*}} = spv.GLSL.Ldexp {{%.*}}, {{%.*}} : f32, i32 -> f32
+  %0 = spv.GLSL.Ldexp %arg0, %arg1 : f32, i32 -> f32
+  return
+}
+
+// -----
+func @ldexp_vec(%arg0 : vector<3xf32>, %arg1 : vector<3xi32>) -> () {
+  // CHECK: {{%.*}} = spv.GLSL.Ldexp {{%.*}}, {{%.*}} : vector<3xf32>, vector<3xi32> -> vector<3xf32>
+  %0 = spv.GLSL.Ldexp %arg0, %arg1 : vector<3xf32>, vector<3xi32> -> vector<3xf32>
+  return
+}
+
+// -----
+
+func @ldexp_wrong_result_type(%arg0 : vector<3xf32>, %arg1 : vector<3xi32>) -> () {
+  // expected-error @+1 {{result type must be the same type as the type of the x operand}}
+  %0 = spv.GLSL.Ldexp %arg0, %arg1 : vector<3xf32>, vector<3xi32> -> vector<2xf32>
+  return
+}
+
+// -----
+
+func @ldexp_wrong_type_scalar(%arg0 : f32, %arg1 : vector<2xi32>) -> () {
+  // expected-error @+1 {{the number of components in x and exp must be the same}}
+  %0 = spv.GLSL.Ldexp %arg0, %arg1 : f32, vector<2xi32> -> f32
+  return
+}
+
+// -----
+
+func @ldexp_wrong_type_vec_1(%arg0 : vector<3xf32>, %arg1 : i32) -> () {
+  // expected-error @+1 {{the number of components in x and exp must be the same}}
+  %0 = spv.GLSL.Ldexp %arg0, %arg1 : vector<3xf32>, i32 -> vector<3xf32>
+  return
+}
+
+// -----
+
+func @ldexp_wrong_type_vec_2(%arg0 : vector<3xf32>, %arg1 : vector<2xi32>) -> () {
+  // expected-error @+1 {{the number of components in x and exp must be the same}}
+  %0 = spv.GLSL.Ldexp %arg0, %arg1 : vector<3xf32>, vector<2xi32> -> vector<3xf32>
+  return
+}
 No newline at end of file

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

 // RUN: mlir-translate -test-spirv-roundtrip %s | FileCheck %s
 
 spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
-  spv.func @fmul(%arg0 : f32, %arg1 : f32) "None" {
+  spv.func @fmul(%arg0 : f32, %arg1 : f32, %arg2 : i32) "None" {
     // CHECK: {{%.*}} = spv.GLSL.Exp {{%.*}} : f32
     %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
     // CHECK: {{%.*}} = spv.GLSL.Cos {{%.*}} : f32
     %3 = spv.GLSL.Cos %arg0 : f32
spv.func @fmul(%arg0 : f32, %arg1 : f32, %arg2 : i32) "None" {
     // 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)>
+    // CHECK: {{%.*}} = spv.GLSL.Ldexp {{%.*}}, {{%.*}} : f32, i32 -> f32
+    %14 = spv.GLSL.Ldexp %arg0, %arg2 : f32, i32 -> f32
     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
   }