[mlir][spirv] Define spv.GLSL.Fma and add lowerings

[hanchung]

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

39–40

I think this is not part of the patch? It's better to have a cleanup in another patch, or at least mention it in commit message.

[antiagainst]

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

39–40

This is lightweight enough so I think it's fine to combine. Will update the commit message.

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

let description = [{
 
     ```mlir
     %2 = spv.GLSL.SClamp %x, %min, %max : si32
     %3 = spv.GLSL.SClamp %x, %min, %max : vector<3xsi16>
     ```
   }];
 }
 
+// -----
+
+def SPV_GLSLFmaOp : SPV_GLSLTernaryArithmeticOp<"Fma", 50, SPV_Float> {
+  let summary = "Computes a * b + c.";
+
+  let description = [{
+    In uses where this operation is decorated with NoContraction:
+
+    - fma is considered a single operation, whereas the expression a * b + c
+      is considered two operations.
+    - The precision of fma can differ from the precision of the expression
+      a * b + c.
+    - fma will be computed with the same precision as any other fma decorated
+      with NoContraction, giving invariant results for the same input values
+      of a, b, and c.
+
+    Otherwise, in the absence of a NoContraction decoration, there are no
+    special constraints on the number of operations or difference in precision
+    between fma and the expression a * b +c.
+
+    The operands must all be a scalar or vector whose component type is
+    floating-point.
+
+    Result Type and the type of all operands must be the same type. Results
+    are computed per component.
+
+    <!-- End of AutoGen section -->
+    ```
+    fma-op ::= ssa-id `=` `spv.GLSL.Fma` ssa-use, ssa-use, ssa-use `:`
+               float-scalar-vector-type
+    ```
+    #### Example:
+
+    ```mlir
+    %0 = spv.GLSL.Fma %a, %b, %c : f32
+    %1 = spv.GLSL.Fma %a, %b, %c : vector<3xf16>
+    ```
+  }];
+}
+
 #endif // MLIR_DIALECT_SPIRV_IR_GLSL_OPS

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

struct VectorBroadcastConvert final
   matchAndRewrite(vector::BroadcastOp broadcastOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     if (broadcastOp.source().getType().isa<VectorType>() ||
         !spirv::CompositeType::isValid(broadcastOp.getVectorType()))
       return failure();
     vector::BroadcastOp::Adaptor adaptor(operands);
     SmallVector<Value, 4> source(broadcastOp.getVectorType().getNumElements(),
                                  adaptor.source());
-    Value construct = rewriter.create<spirv::CompositeConstructOp>(
-        broadcastOp.getLoc(), broadcastOp.getVectorType(), source);
-    rewriter.replaceOp(broadcastOp, construct);
+    rewriter.replaceOpWithNewOp<spirv::CompositeConstructOp>(
+        broadcastOp, broadcastOp.getVectorType(), source);
     return success();
   }
 };
 
 struct VectorExtractOpConvert final
     : public OpConversionPattern<vector::ExtractOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(vector::ExtractOp extractOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     if (extractOp.getType().isa<VectorType>() ||
         !spirv::CompositeType::isValid(extractOp.getVectorType()))
       return failure();
     vector::ExtractOp::Adaptor adaptor(operands);
     int32_t id = extractOp.position().begin()->cast<IntegerAttr>().getInt();
-    Value newExtract = rewriter.create<spirv::CompositeExtractOp>(
-        extractOp.getLoc(), adaptor.vector(), id);
-    rewriter.replaceOp(extractOp, newExtract);
+    rewriter.replaceOpWithNewOp<spirv::CompositeExtractOp>(
+        extractOp, adaptor.vector(), id);
+    return success();
+  }
+};
+
+struct VectorFmaOpConvert final : public OpConversionPattern<vector::FMAOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(vector::FMAOp fmaOp, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    if (!spirv::CompositeType::isValid(fmaOp.getVectorType()))
+      return failure();
+    vector::FMAOp::Adaptor adaptor(operands);
+    rewriter.replaceOpWithNewOp<spirv::GLSLFmaOp>(
+        fmaOp, fmaOp.getType(), adaptor.lhs(), adaptor.rhs(), adaptor.acc());
     return success();
   }
 };
 
 struct VectorInsertOpConvert final
     : public OpConversionPattern<vector::InsertOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(vector::InsertOp insertOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     if (insertOp.getSourceType().isa<VectorType>() ||
         !spirv::CompositeType::isValid(insertOp.getDestVectorType()))
       return failure();
     vector::InsertOp::Adaptor adaptor(operands);
     int32_t id = insertOp.position().begin()->cast<IntegerAttr>().getInt();
-    Value newInsert = rewriter.create<spirv::CompositeInsertOp>(
-        insertOp.getLoc(), adaptor.source(), adaptor.dest(), id);
-    rewriter.replaceOp(insertOp, newInsert);
+    rewriter.replaceOpWithNewOp<spirv::CompositeInsertOp>(
+        insertOp, adaptor.source(), adaptor.dest(), id);
     return success();
   }
 };
 
 struct VectorExtractElementOpConvert final
     : public OpConversionPattern<vector::ExtractElementOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(vector::ExtractElementOp extractElementOp,
                   ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     if (!spirv::CompositeType::isValid(extractElementOp.getVectorType()))
       return failure();
     vector::ExtractElementOp::Adaptor adaptor(operands);
-    Value newExtractElement = rewriter.create<spirv::VectorExtractDynamicOp>(
-        extractElementOp.getLoc(), extractElementOp.getType(), adaptor.vector(),
+    rewriter.replaceOpWithNewOp<spirv::VectorExtractDynamicOp>(
+        extractElementOp, extractElementOp.getType(), adaptor.vector(),
         extractElementOp.position());
-    rewriter.replaceOp(extractElementOp, newExtractElement);
     return success();
   }
 };
 
 struct VectorInsertElementOpConvert final
     : public OpConversionPattern<vector::InsertElementOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(vector::InsertElementOp insertElementOp,
                   ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     if (!spirv::CompositeType::isValid(insertElementOp.getDestVectorType()))
       return failure();
     vector::InsertElementOp::Adaptor adaptor(operands);
-    Value newInsertElement = rewriter.create<spirv::VectorInsertDynamicOp>(
-        insertElementOp.getLoc(), insertElementOp.getType(),
-        insertElementOp.dest(), adaptor.source(), insertElementOp.position());
-    rewriter.replaceOp(insertElementOp, newInsertElement);
+    rewriter.replaceOpWithNewOp<spirv::VectorInsertDynamicOp>(
+        insertElementOp, insertElementOp.getType(), insertElementOp.dest(),
+        adaptor.source(), insertElementOp.position());
     return success();
   }
 };
 
 } // namespace
 
 void mlir::populateVectorToSPIRVPatterns(MLIRContext *context,
                                          SPIRVTypeConverter &typeConverter,
                                          OwningRewritePatternList &patterns) {
-  patterns.insert<VectorBroadcastConvert, VectorExtractOpConvert,
-                  VectorInsertOpConvert, VectorExtractElementOpConvert,
-                  VectorInsertElementOpConvert>(typeConverter, context);
+  patterns.insert<VectorBroadcastConvert, VectorExtractElementOpConvert,
+                  VectorExtractOpConvert, VectorFmaOpConvert,
+                  VectorInsertOpConvert, VectorInsertElementOpConvert>(
+      typeConverter, context);
 }

mlir/test/Conversion/VectorToSPIRV/simple.mlir

 
 // -----
 
 func @insert_element_negative(%val: f32, %arg0 : vector<5xf32>, %id : i32) {
 // expected-error @+1 {{failed to legalize operation 'vector.insertelement'}}
   %0 = vector.insertelement %val, %arg0[%id : i32] : vector<5xf32>
   spv.Return
 }
+
+// -----
+
+// CHECK-LABEL: func @fma
+//  CHECK-SAME: %[[A:.*]]: vector<4xf32>, %[[B:.*]]: vector<4xf32>, %[[C:.*]]: vector<4xf32>
+//       CHECK:   spv.GLSL.Fma %[[A]], %[[B]], %[[C]] : vector<4xf32>
+func @fma(%a: vector<4xf32>, %b: vector<4xf32>, %c: vector<4xf32>) {
+  %0 = vector.fma %a, %b, %c: vector<4xf32>
+  spv.Return
+}

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

 
 // -----
 
 func @fclamp(%arg0 : i32, %min : i32, %max : i32) -> () {
   // expected-error @+1 {{must be 8/16/32/64-bit signed integer or vector}}
   %2 = spv.GLSL.SClamp %arg0, %min, %max : i32
   return
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spv.GLSL.Fma
+//===----------------------------------------------------------------------===//
+
+func @fma(%a : f32, %b : f32, %c : f32) -> () {
+  // CHECK: spv.GLSL.Fma {{%[^,]*}}, {{%[^,]*}}, {{%[^,]*}} : f32
+  %2 = spv.GLSL.Fma %a, %b, %c : f32
+  return
+}
+
+// -----
+
+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
+}

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

spv.func @uclamp(%arg0 : ui32, %arg1 : ui32, %arg2 : ui32) "None" {
     spv.Return
   }
 
   spv.func @sclamp(%arg0 : si32, %arg1 : si32, %arg2 : si32) "None" {
     // CHECK: spv.GLSL.SClamp {{%[^,]*}}, {{%[^,]*}}, {{%[^,]*}} : si32
     %13 = spv.GLSL.SClamp %arg0, %arg1, %arg2 : si32
     spv.Return
   }
+
+  spv.func @fma(%arg0 : f32, %arg1 : f32, %arg2 : f32) "None" {
+    // CHECK: spv.GLSL.Fma {{%[^,]*}}, {{%[^,]*}}, {{%[^,]*}} : f32
+    %13 = spv.GLSL.Fma %arg0, %arg1, %arg2 : f32
+    spv.Return
+  }
 }