[fixup] Skip the two-stage LLVM IR generation, map the op directly to the LLVM IR intrinsic

Author: momchil-velikov

mlir/include/mlir/Dialect/ArmSVE/IR/ArmSVE.td

@@ -293,10 +293,10 @@ def UsmmlaOp : ArmSVE_Op<"usmmla", [Pure,
     "$acc `,` $src1 `,` $src2 attr-dict `:` type($src1) `to` type($dst)";
 }
 
-
-def BfmmlaOp : ArmSVE_Op<"bfmmla", [Pure,
-                                    AllTypesMatch<["src1", "src2"]>,
-                                    AllTypesMatch<["acc", "dst"]>]> {
+def BfmmlaOp : ArmSVE_IntrOp<"bfmmla", [Pure,
+                                        AllTypesMatch<["src1", "src2"]>,
+                                        AllTypesMatch<["acc", "res"]>,
+                                        ]> {
   let summary = "BFloat16 matrix multiply-accumulate";
   let description = [{
     BFMMLA: BFloat16 matrix multiply-accumulate into 2×2 matrices";
@@ -317,9 +317,9 @@ def BfmmlaOp : ArmSVE_Op<"bfmmla", [Pure,
           ScalableVectorOfLengthAndType<[8], [BF16]>:$src1,
           ScalableVectorOfLengthAndType<[8], [BF16]>:$src2
   );
-  let results = (outs ScalableVectorOfLengthAndType<[4], [F32]>:$dst);
+  let results = (outs ScalableVectorOfLengthAndType<[4], [F32]>:$res);
   let assemblyFormat =
-    "$acc `,` $src1 `,` $src2 attr-dict `:` type($src1) `to` type($dst)";
+    "$acc `,` $src1 `,` $src2 attr-dict `:` type($src1) `to` type($res)";
 }
 
 class SvboolTypeConstraint<string lhsArg, string rhsArg> : TypesMatchWith<
@@ -619,12 +619,6 @@ def UsmmlaIntrOp :
   ArmSVE_IntrBinaryOverloadedOp<"usmmla">,
   Arguments<(ins AnyScalableVectorOfAnyRank, AnyScalableVectorOfAnyRank, AnyScalableVectorOfAnyRank)>;
 
-def BfmmlaIntrOp :
-  ArmSVE_IntrOp<"bfmmla", [Pure, TypeIs<"res", ScalableVectorOfLengthAndType<[4], [F32]>>]>,
-  Arguments<(ins Arg<ScalableVectorOfLengthAndType<[4], [F32]>, "acc">:$acc,
-                 Arg<ScalableVectorOfLengthAndType<[8], [BF16]>, "lhs">:$lhs,
-                 Arg<ScalableVectorOfLengthAndType<[8], [BF16]>, "rhs">:$rhs)>;
-
 def SdotIntrOp :
   ArmSVE_IntrBinaryOverloadedOp<"sdot">,
   Arguments<(ins AnyScalableVectorOfAnyRank, AnyScalableVectorOfAnyRank, AnyScalableVectorOfAnyRank)>;

mlir/lib/Dialect/ArmSVE/Transforms/LegalizeForLLVMExport.cpp

@@ -25,7 +25,6 @@ using SmmlaOpLowering = OneToOneConvertToLLVMPattern<SmmlaOp, SmmlaIntrOp>;
 using UdotOpLowering = OneToOneConvertToLLVMPattern<UdotOp, UdotIntrOp>;
 using UmmlaOpLowering = OneToOneConvertToLLVMPattern<UmmlaOp, UmmlaIntrOp>;
 using UsmmlaOpLowering = OneToOneConvertToLLVMPattern<UsmmlaOp, UsmmlaIntrOp>;
-using BfmmlaOpLowering = OneToOneConvertToLLVMPattern<BfmmlaOp, BfmmlaIntrOp>;
 using DupQLaneLowering =
     OneToOneConvertToLLVMPattern<DupQLaneOp, DupQLaneIntrOp>;
 using ScalableMaskedAddIOpLowering =
@@ -192,8 +191,7 @@ void mlir::populateArmSVELegalizeForLLVMExportPatterns(
   // Populate conversion patterns
 
   // clang-format off
-  patterns.add<BfmmlaOpLowering,
-               ConvertFromSvboolOpLowering,
+  patterns.add<ConvertFromSvboolOpLowering,
                ConvertToSvboolOpLowering,
                DupQLaneLowering,
                PselOpLowering,
@@ -222,7 +220,7 @@ void mlir::populateArmSVELegalizeForLLVMExportPatterns(
 void mlir::configureArmSVELegalizeForExportTarget(
     LLVMConversionTarget &target) {
   // clang-format off
-  target.addLegalOp<BfmmlaIntrOp,
+  target.addLegalOp<BfmmlaOp,
                     ConvertFromSvboolIntrOp,
                     ConvertToSvboolIntrOp,
                     DupQLaneIntrOp,
@@ -244,8 +242,7 @@ void mlir::configureArmSVELegalizeForExportTarget(
                     ZipX2IntrOp,
                     ZipX4IntrOp,
                     SdotIntrOp>();
-  target.addIllegalOp<BfmmlaOp,
-                      ConvertFromSvboolOp,
+  target.addIllegalOp<ConvertFromSvboolOp,
                       ConvertToSvboolOp,
                       DupQLaneOp,
                       PselOp,

mlir/test/Dialect/ArmSVE/legalize-for-llvm.mlir

@@ -60,15 +60,6 @@ func.func @arm_sve_usmmla(%a: vector<[16]xi8>,
 
 // -----
 
-func.func @arm_sve_bfmmla(%a: vector<[8]xbf16>,
-                          %b: vector<[8]xbf16>,
-                          %c: vector<[4]xf32>) -> vector<[4]xf32> {
-  // CHECK: arm_sve.intr.bfmmla
-  %0 = arm_sve.bfmmla %c, %a, %b : vector<[8]xbf16> to vector<[4]xf32>
-  return %0 : vector<[4]xf32>
-}
-// -----
-
 func.func @arm_sve_arithi_masked(%a: vector<[4]xi32>,
                             %b: vector<[4]xi32>,
                             %c: vector<[4]xi32>,

mlir/test/Dialect/ArmSVE/roundtrip.mlir

@@ -58,8 +58,8 @@ func.func @arm_sve_usmmla(%a: vector<[16]xi8>,
 func.func @arm_sve_bfmmla(%a: vector<[8]xbf16>,
                           %b: vector<[8]xbf16>,
                           %c: vector<[4]xf32>) -> vector<[4]xf32> {
-  // CHECK: arm_sve.bfmmla {{.*}}: vector<[8]xbf16> to vector<[4]xf32>
-  %0 = arm_sve.bfmmla %c, %a, %b : vector<[8]xbf16> to vector<[4]xf32>
+  // CHECK: arm_sve.intr.bfmmla {{.*}}: vector<[8]xbf16> to vector<[4]xf32>
+  %0 = arm_sve.intr.bfmmla %c, %a, %b : vector<[8]xbf16> to vector<[4]xf32>
   return %0 : vector<[4]xf32>
 }