[VMVX] Added math.fma to vm.fma lowering (#24915)

There was no lowering from math.fma to vm.fma. This patch adds one.

Vector types for ternary conversions need to be supported as well:
https://github.com/iree-org/iree/issues/24937

`Assisted-by: Claude Opus 5`

Signed-off-by: Ayokunle Amodu
[ayokunle321@gmail.com](mailto:ayokunle321@gmail.com)

---------

Signed-off-by: Ayokunle Amodu <ayokunle321@gmail.com>

compiler/src/iree/compiler/Dialect/VM/Conversion/MathToVM/Patterns.cpp

diff --git a/compiler/src/iree/compiler/Dialect/VM/Conversion/MathToVM/Patterns.cpp b/compiler/src/iree/compiler/Dialect/VM/Conversion/MathToVM/Patterns.cpp
index 78ff73d..798fcda 100644
--- a/compiler/src/iree/compiler/Dialect/VM/Conversion/MathToVM/Patterns.cpp
+++ b/compiler/src/iree/compiler/Dialect/VM/Conversion/MathToVM/Patterns.cpp
@@ -80,6 +80,34 @@
   }
 };
 
+template <typename SrcOpTy, typename Dst32OpTy, typename Dst64OpTy>
+class TernaryArithmeticOpConversion : public OpConversionPattern<SrcOpTy> {
+  using OpConversionPattern<SrcOpTy>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(SrcOpTy srcOp, typename SrcOpTy::Adaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    if (isa<VectorType>(srcOp.getResult().getType())) {
+      return failure();
+    }
+
+    Type type = adaptor.getA().getType();
+    switch (type.getIntOrFloatBitWidth()) {
+    case 32:
+      rewriter.replaceOpWithNewOp<Dst32OpTy>(srcOp, type, adaptor.getA(),
+                                             adaptor.getB(), adaptor.getC());
+      break;
+    case 64:
+      rewriter.replaceOpWithNewOp<Dst64OpTy>(srcOp, type, adaptor.getA(),
+                                             adaptor.getB(), adaptor.getC());
+      break;
+    default:
+      return rewriter.notifyMatchFailure(srcOp, "unsupported bit width");
+    }
+    return success();
+  }
+};
+
 } // namespace
 
 void populateMathToVMPatterns(MLIRContext *context,
@@ -127,6 +155,8 @@
       // so far. So a check of lhs (the float value) is sufficient.
       BinaryArithmeticOpConversion<math::FPowIOp, IREE::VM::FPowI32Op,
                                    IREE::VM::FPowI64Op>,
+      TernaryArithmeticOpConversion<math::FmaOp, IREE::VM::FMAF32Op,
+                                    IREE::VM::FMAF64Op>,
       UnaryArithmeticOpConversion<math::RsqrtOp, IREE::VM::RsqrtF32Op,
                                   IREE::VM::RsqrtF64Op>,
       UnaryArithmeticOpConversion<math::SqrtOp, IREE::VM::SqrtF32Op,

compiler/src/iree/compiler/Dialect/VM/Conversion/MathToVM/test/arithmetic_ops.mlir

diff --git a/compiler/src/iree/compiler/Dialect/VM/Conversion/MathToVM/test/arithmetic_ops.mlir b/compiler/src/iree/compiler/Dialect/VM/Conversion/MathToVM/test/arithmetic_ops.mlir
index 97ec9bc..12eb82a 100644
--- a/compiler/src/iree/compiler/Dialect/VM/Conversion/MathToVM/test/arithmetic_ops.mlir
+++ b/compiler/src/iree/compiler/Dialect/VM/Conversion/MathToVM/test/arithmetic_ops.mlir
@@ -76,22 +76,25 @@
   // CHECK: vm.erf.f32
   %15 = math.erf %14 : f32
 
+  // CHECK: vm.fma.f32
+  %16 = math.fma %arg0, %arg0, %15 : f32
+
   // CHECK: vm.abs.f32
-  %16 = math.absf %15 : f32
+  %17 = math.absf %16 : f32
 
   // CHECK: vm.ceil.f32
-  %17 = math.ceil %16 : f32
+  %18 = math.ceil %17 : f32
 
   // CHECK: vm.floor.f32
-  %18 = math.floor %17 : f32
+  %19 = math.floor %18 : f32
 
   // CHECK: vm.round.f32
-  %19 = math.round %18 : f32
+  %20 = math.round %19 : f32
 
   // CHECK: vm.round.f32.even
-  %20 = math.roundeven %19 : f32
+  %21 = math.roundeven %20 : f32
 
-  return %20 : f32
+  return %21 : f32
 }
 
 // -----