[mlir][spirv] Add OpenCL fma op and lowering

[Hardcode84]

Summary

Also, it seems Khronos has chanhed html spec format so small adjustment to script was needed.
Base op parsing is also probably broken.

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

 // return type matches.
 class SPV_OCLBinaryArithmeticOp<string mnemonic, int opcode, Type type,
                                 list<Trait> traits = []> :
   SPV_OCLBinaryOp<mnemonic, opcode, type, type,
                   traits # [SameOperandsAndResultType]> {                    
   let assemblyFormat = "operands attr-dict `:` type($result)";
 }
 
+// Base class for OpenCL binary ops.
+class SPV_OCLTernaryOp<string mnemonic, int opcode, Type resultType,
+                      Type operandType, list<Trait> traits = []> :
+  SPV_OCLOp<mnemonic, opcode, !listconcat([NoSideEffect], traits)> {
+
+  let arguments = (ins
+    SPV_ScalarOrVectorOf<operandType>:$x,
+    SPV_ScalarOrVectorOf<operandType>:$y,
+    SPV_ScalarOrVectorOf<operandType>:$z
+  );
+
+  let results = (outs
+    SPV_ScalarOrVectorOf<resultType>:$result
+  );
+
+  let hasVerifier = 0;
+}
+
+// Base class for OpenCL Ternary arithmetic ops where operand types and
+// return type matches.
+class SPV_OCLTernaryArithmeticOp<string mnemonic, int opcode, Type type,
+                                list<Trait> traits = []> :
+  SPV_OCLTernaryOp<mnemonic, opcode, type, type,
+                  traits # [SameOperandsAndResultType]> {
+  let assemblyFormat = "operands attr-dict `:` type($result)";
+}
+
+
 // -----
 
-def SPV_OCLErfOp : SPV_OCLUnaryArithmeticOp<"erf", 18, SPV_Float> {
+def SPV_OCLFmaOp : SPV_OCLTernaryArithmeticOp<"fma", 26, SPV_Float> {
   let summary = [{
-    Error function of x encountered in integrating the normal distribution.
+    Compute the correctly rounded floating-point representation of the sum
+    of c with the infinitely precise product of a and b. Rounding of
+    intermediate products shall not occur. Edge case results are per the
+    IEEE 754-2008 standard.
   }];
 
   let description = [{
-    Result Type and x must be floating-point or vector(2,3,4,8,16) of
+    Result Type, a, b and c must be floating-point or vector(2,3,4,8,16) of
     floating-point values.
 
     All of the operands, including the Result Type operand, must be of the
     same type.
 
     <!-- End of AutoGen section -->
 
     ```
-    float-scalar-vector-type ::= float-type |
-                                 `vector<` integer-literal `x` float-type `>`
-    erf-op ::= ssa-id `=` `spv.OCL.erf` ssa-use `:`
+    fma-op ::= ssa-id `=` `spv.OCL.fma` ssa-use, ssa-use, ssa-use `:`
                float-scalar-vector-type
     ```mlir
 
-    #### Example:
-
     ```
-    %2 = spv.OCL.erf %0 : f32
-    %3 = spv.OCL.erf %1 : vector<3xf16>
+    %0 = spv.OCL.fma %a, %b, %c : f32
+    %1 = spv.OCL.fma %a, %b, %c : vector<3xf16>
     ```
   }];
 }
 
 // -----
 
 def SPV_OCLCeilOp : SPV_OCLUnaryArithmeticOp<"ceil", 12, SPV_Float> {
   let summary = [{
let description = [{
     %2 = spv.OCL.cos %0 : f32
     %3 = spv.OCL.cos %1 : vector<3xf16>
     ```
   }];
 }
 
 // -----
 
+def SPV_OCLErfOp : SPV_OCLUnaryArithmeticOp<"erf", 18, SPV_Float> {
+  let summary = [{
+    Error function of x encountered in integrating the normal distribution.
+  }];
+
+  let description = [{
+    Result Type and x must be floating-point or vector(2,3,4,8,16) of
+    floating-point values.
+
+    All of the operands, including the Result Type operand, must be of the
+    same type.
+
+    <!-- End of AutoGen section -->
+
+    ```
+    float-scalar-vector-type ::= float-type |
+                                 `vector<` integer-literal `x` float-type `>`
+    erf-op ::= ssa-id `=` `spv.OCL.erf` ssa-use `:`
+               float-scalar-vector-type
+    ```mlir
+
+    #### Example:
+
+    ```
+    %2 = spv.OCL.erf %0 : f32
+    %3 = spv.OCL.erf %1 : vector<3xf16>
+    ```
+  }];
+}
+
+// -----
+
 def SPV_OCLExpOp : SPV_OCLUnaryArithmeticOp<"exp", 19, SPV_Float> {
   let summary = "Exponentiation of Operand 1";
 
   let description = [{
     Compute the base-e exponential of x. (i.e. ex)
 
     Result Type and x must be floating-point or vector(2,3,4,8,16) of
     floating-point values.

mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp

void populateMathToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
   // GLSL patterns
   patterns
       .add<Log1pOpPattern<spirv::GLSLLogOp>, ExpM1OpPattern<spirv::GLSLExpOp>,
            spirv::ElementwiseOpPattern<math::AbsOp, spirv::GLSLFAbsOp>,
            spirv::ElementwiseOpPattern<math::CeilOp, spirv::GLSLCeilOp>,
            spirv::ElementwiseOpPattern<math::CosOp, spirv::GLSLCosOp>,
            spirv::ElementwiseOpPattern<math::ExpOp, spirv::GLSLExpOp>,
            spirv::ElementwiseOpPattern<math::FloorOp, spirv::GLSLFloorOp>,
+           spirv::ElementwiseOpPattern<math::FmaOp, spirv::GLSLFmaOp>,
            spirv::ElementwiseOpPattern<math::LogOp, spirv::GLSLLogOp>,
            spirv::ElementwiseOpPattern<math::PowFOp, spirv::GLSLPowOp>,
            spirv::ElementwiseOpPattern<math::RsqrtOp, spirv::GLSLInverseSqrtOp>,
            spirv::ElementwiseOpPattern<math::SinOp, spirv::GLSLSinOp>,
            spirv::ElementwiseOpPattern<math::SqrtOp, spirv::GLSLSqrtOp>,
-           spirv::ElementwiseOpPattern<math::TanhOp, spirv::GLSLTanhOp>,
-           spirv::ElementwiseOpPattern<math::FmaOp, spirv::GLSLFmaOp>>(
+           spirv::ElementwiseOpPattern<math::TanhOp, spirv::GLSLTanhOp>>(
           typeConverter, patterns.getContext());
 
   // OpenCL patterns
   patterns.add<Log1pOpPattern<spirv::OCLLogOp>, ExpM1OpPattern<spirv::OCLExpOp>,
                spirv::ElementwiseOpPattern<math::AbsOp, spirv::OCLFAbsOp>,
                spirv::ElementwiseOpPattern<math::CeilOp, spirv::OCLCeilOp>,
                spirv::ElementwiseOpPattern<math::CosOp, spirv::OCLCosOp>,
                spirv::ElementwiseOpPattern<math::ErfOp, spirv::OCLErfOp>,
                spirv::ElementwiseOpPattern<math::ExpOp, spirv::OCLExpOp>,
                spirv::ElementwiseOpPattern<math::FloorOp, spirv::OCLFloorOp>,
+               spirv::ElementwiseOpPattern<math::FmaOp, spirv::OCLFmaOp>,
                spirv::ElementwiseOpPattern<math::LogOp, spirv::OCLLogOp>,
                spirv::ElementwiseOpPattern<math::PowFOp, spirv::OCLPowOp>,
                spirv::ElementwiseOpPattern<math::RsqrtOp, spirv::OCLRsqrtOp>,
                spirv::ElementwiseOpPattern<math::SinOp, spirv::OCLSinOp>,
                spirv::ElementwiseOpPattern<math::SqrtOp, spirv::OCLSqrtOp>,
                spirv::ElementwiseOpPattern<math::TanhOp, spirv::OCLTanhOp>>(
       typeConverter, patterns.getContext());
 }
 
 } // namespace mlir

mlir/test/Conversion/MathToSPIRV/math-to-glsl-spirv.mlir

 
 // CHECK-LABEL: @float32_binary_vector
 func @float32_binary_vector(%lhs: vector<4xf32>, %rhs: vector<4xf32>) {
   // CHECK: spv.GLSL.Pow %{{.*}}: vector<4xf32>
   %0 = math.powf %lhs, %rhs : vector<4xf32>
   return
 }
 
-  // CHECK-LABEL: @float32_ternary_scalar
 func @float32_ternary_scalar(%a: f32, %b: f32, %c: f32) {
   // CHECK: spv.GLSL.Fma %{{.*}}: f32
   %0 = math.fma %a, %b, %c : f32
   return
 }
 
 // CHECK-LABEL: @float32_ternary_vector
 func @float32_ternary_vector(%a: vector<4xf32>, %b: vector<4xf32>,
                             %c: vector<4xf32>) {
   // CHECK: spv.GLSL.Fma %{{.*}}: vector<4xf32>
   %0 = math.fma %a, %b, %c : vector<4xf32>
   return
 }
 
 } // end module

mlir/test/Conversion/MathToSPIRV/math-to-opencl-spirv.mlir

 
 // CHECK-LABEL: @float32_binary_vector
 func @float32_binary_vector(%lhs: vector<4xf32>, %rhs: vector<4xf32>) {
   // CHECK: spv.OCL.pow %{{.*}}: vector<4xf32>
   %0 = math.powf %lhs, %rhs : vector<4xf32>
   return
 }
 
+// CHECK-LABEL: @float32_ternary_scalar
+func @float32_ternary_scalar(%a: f32, %b: f32, %c: f32) {
+  // CHECK: spv.OCL.fma %{{.*}}: f32
+  %0 = math.fma %a, %b, %c : f32
+  return
+}
+
+// CHECK-LABEL: @float32_ternary_vector
+func @float32_ternary_vector(%a: vector<4xf32>, %b: vector<4xf32>,
+                            %c: vector<4xf32>) {
+  // CHECK: spv.OCL.fma %{{.*}}: vector<4xf32>
+  %0 = math.fma %a, %b, %c : vector<4xf32>
+  return
+}
+
 } // end module

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

 // -----
 
 func @sabs(%arg0 : i32) -> () {
   // expected-error @+2 {{expected non-function type}}
   %2 = spv.OCL.s_abs %arg0 :
   return
 }
 
+// -----
+
+//===----------------------------------------------------------------------===//
+// spv.OCL.fma
+//===----------------------------------------------------------------------===//
+
+func @fma(%a : f32, %b : f32, %c : f32) -> () {
+  // CHECK: spv.OCL.fma {{%[^,]*}}, {{%[^,]*}}, {{%[^,]*}} : f32
+  %2 = spv.OCL.fma %a, %b, %c : f32
+  return
+}
+
+// -----
+
+func @fma(%a : vector<3xf32>, %b : vector<3xf32>, %c : vector<3xf32>) -> () {
+  // CHECK: spv.OCL.fma {{%[^,]*}}, {{%[^,]*}}, {{%[^,]*}} : vector<3xf32>
+  %2 = spv.OCL.fma %a, %b, %c : vector<3xf32>
+  return
+}

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

spv.func @integer_insts(%arg0 : i32) "None" {
     spv.Return
   }
   
   spv.func @vector_size16(%arg0 : vector<16xf32>) "None" {
     // CHECK: {{%.*}} = spv.OCL.fabs {{%.*}} : vector<16xf32>
     %0 = spv.OCL.fabs %arg0 : vector<16xf32>
     spv.Return
   }
+
+  spv.func @fma(%arg0 : f32, %arg1 : f32, %arg2 : f32) "None" {
+    // CHECK: spv.OCL.fma {{%[^,]*}}, {{%[^,]*}}, {{%[^,]*}} : f32
+    %13 = spv.OCL.fma %arg0, %arg1, %arg2 : f32
+    spv.Return
+  }
 }

mlir/utils/spirv/gen_spirv_dialect.py

def get_spirv_doc_from_html_spec(url, settings):
   spec = response.content
 
   from bs4 import BeautifulSoup
   spirv = BeautifulSoup(spec, 'html.parser')
 
   doc = {}
 
   if settings.gen_ocl_ops:
-    section_anchor = spirv.find('h2', {'id': '_a_id_binary_a_binary_form'})
+    section_anchor = spirv.find('h2', {'id': '_binary_form'})
     for section in section_anchor.parent.find_all('div', {'class': 'sect2'}):
       for table in section.find_all('table'):
         inst_html = table.tbody.tr.td
         opname = inst_html.a['id']
         # Ignore the first line, which is just the opname.
         doc[opname] = inst_html.text.split('\n', 1)[1].strip()
   else:
     section_anchor = spirv.find('h3', {'id': '_a_id_instructions_a_instructions'})