[mlir][ArmSVE] Add masked arithmetic operations

[jsetoain]

Summary

These instructions map to SVE-specific instrinsics that accept a
predicate operand to support control flow in vector code.

[jsetoain]

Add stronger type checking

[jsetoain]

This comment was removed by jsetoain.

[jsetoain]

Use classes to reduce redundancy in operation descriptions

[jsetoain]

Add type checks

[ftynse]

LGTM with comments addressed.

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

387

Shouldn't this be commutative?

468

Nit: we can have type constraints stricter than AnyVector here

mlir/lib/Dialect/ArmSVE/IR/ArmSVEDialect.cpp

69

Nit: please add doc comment

74

I'd rather return null or assert here. We don't use NoneType to denote error state.

[jsetoain]

Address comments

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

468

Since the LLVM Dialect already supports Scalable Vectors, this quick fix could go directly into LLVM dialect. To keep the patch contained within ArmSVE, and given that nobody else is using this yet, I think it's better to leave it here for now. Once the scalable vector has a more defined shape outside of ArmSVE, I myself will take care of adapting and moving these LLVM checks to where they belong. Alternatively, if this is not acceptable, I can create a different patch for LLVM dialect adding the extra type constraint, and make this patch depend on that one. That way, LLVM dialect changes don't get mixed up with hw-specific dialect changes.

[ftynse]

You can follow this https://llvm.org/docs/DeveloperPolicy.html#obtaining-commit-access to get commit access for yourself.

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

int64_t getNumElements() const {
       for (auto dim : shape)
         num *= dim;
       return num;
     }
   }];
 }
 
 //===----------------------------------------------------------------------===//
+// Additional LLVM type constraints
+//===----------------------------------------------------------------------===//
+def LLVMScalableVectorType :
+  Type<CPred<"$_self.isa<::mlir::LLVM::LLVMScalableVectorType>()">,
+       "LLVM dialect scalable vector type">;
+
+//===----------------------------------------------------------------------===//
 // ArmSVE op definitions
 //===----------------------------------------------------------------------===//
 
 class ArmSVE_Op<string mnemonic, list<OpTrait> traits = []> :
   Op<ArmSVE_Dialect, mnemonic, traits> {}
 
 class ArmSVE_NonSVEIntrUnaryOverloadedOp<string mnemonic,
                                          list<OpTrait> traits =[]> :
let arguments = (ins
           ScalableVectorOf<[I8, I16, I32, I64]>:$src1,
           ScalableVectorOf<[I8, I16, I32, I64]>:$src2
   );
   let results = (outs ScalableVectorOf<[I8, I16, I32, I64]>:$dst);
   let assemblyFormat =
     "$src1 `,` $src2 attr-dict `:` type($src1)";
 }
 
+class ScalableMaskedFOp<string mnemonic, string op_description,
+                        list<OpTrait> traits = []> :
+  ArmSVE_Op<mnemonic, !listconcat(traits,
+                       [AllTypesMatch<["src1", "src2", "res"]>,
+                        TypesMatchWith<
+                          "mask has i1 element type and same shape as operands",
+                          "src1", "mask", "getI1SameShape($_self)">])> {
+  let summary = "masked " # op_description # " for scalable vectors of floats";
+  let description = [{
+    The `arm_sve.}] # mnemonic # [{` operation takes one scalable vector mask
+    and two scalable vector operands, and perform floating point }] #
+    op_description # [{ on active lanes. Inactive lanes will keep the value of
+    the first operand.}];
+  let arguments = (ins
+          ScalableVectorOf<[I1]>:$mask,
+          ScalableVectorOf<[AnyFloat]>:$src1,
+          ScalableVectorOf<[AnyFloat]>:$src2
+  );
+  let results = (outs ScalableVectorOf<[AnyFloat]>:$res);
+  let assemblyFormat =
+    "$mask `,` $src1 `,` $src2 attr-dict `:` type($mask) `,` type($res)";
+}
+
+class ScalableMaskedIOp<string mnemonic, string op_description,
+                        list<OpTrait> traits = []> :
+  ArmSVE_Op<mnemonic, !listconcat(traits,
+                       [AllTypesMatch<["src1", "src2", "res"]>,
+                        TypesMatchWith<
+                          "mask has i1 element type and same shape as operands",
+                          "src1", "mask", "getI1SameShape($_self)">])> {
+  let summary = "masked " # op_description # " for scalable vectors of integers";
+  let description = [{
+    The `arm_sve.}] # mnemonic # [{` operation takes one scalable vector mask
+    and two scalable vector operands, and perform integer }] #
+    op_description # [{ on active lanes. Inactive lanes will keep the value of
+    the first operand.}];
+  let arguments = (ins
+          ScalableVectorOf<[I1]>:$mask,
+          ScalableVectorOf<[I8, I16, I32, I64]>:$src1,
+          ScalableVectorOf<[I8, I16, I32, I64]>:$src2
+  );
+  let results = (outs ScalableVectorOf<[I8, I16, I32, I64]>:$res);
+  let assemblyFormat =
+    "$mask `,` $src1 `,` $src2 attr-dict `:` type($mask) `,` type($res)";
+}
+
 def SdotOp : ArmSVE_Op<"sdot",
                [NoSideEffect,
                AllTypesMatch<["src1", "src2"]>,
                AllTypesMatch<["acc", "dst"]>,
              ]> {
   let summary = "Vector-vector dot product and accumulate op";
   let description = [{
     SDOT: Signed integer addition of dot product.

 def ScalableMulFOp : ScalableFOp<"mulf", "multiplication", [Commutative]>;
 
 def ScalableSDivIOp : ScalableIOp<"divi_signed", "signed division">;
 
 def ScalableUDivIOp : ScalableIOp<"divi_unsigned", "unsigned division">;
 
 def ScalableDivFOp : ScalableFOp<"divf", "division">;
 
+def ScalableMaskedAddIOp : ScalableMaskedIOp<"masked.addi", "addition",
+                                             [Commutative]>;
+
+def ScalableMaskedAddFOp : ScalableMaskedFOp<"masked.addf", "addition",
+                            [Commutative]>;
+
+def ScalableMaskedSubIOp : ScalableMaskedIOp<"masked.subi", "subtraction">;
+
+def ScalableMaskedSubFOp : ScalableMaskedFOp<"masked.subf", "subtraction">;
+
+def ScalableMaskedMulIOp : ScalableMaskedIOp<"masked.muli", "multiplication",
+                            [Commutative]>;
+
+def ScalableMaskedMulFOp : ScalableMaskedFOp<"masked.mulf", "multiplication",
+                            [Commutative]>;
+
+def ScalableMaskedSDivIOp : ScalableMaskedIOp<"masked.divi_signed",
+                                              "signed division">;
+
+def ScalableMaskedUDivIOp : ScalableMaskedIOp<"masked.divi_unsigned",
+                                              "unsigned division">;
+
+def ScalableMaskedDivFOp : ScalableMaskedFOp<"masked.divf", "division">;
+
 def UmmlaIntrOp :
   ArmSVE_IntrBinaryOverloadedOp<"ummla">,
-  Arguments<(ins LLVM_AnyVector, LLVM_AnyVector, LLVM_AnyVector)>;
+  Arguments<(ins LLVMScalableVectorType, LLVMScalableVectorType,
+             LLVMScalableVectorType)>;
 
 def SmmlaIntrOp :
   ArmSVE_IntrBinaryOverloadedOp<"smmla">,
-  Arguments<(ins LLVM_AnyVector, LLVM_AnyVector, LLVM_AnyVector)>;
+  Arguments<(ins LLVMScalableVectorType, LLVMScalableVectorType,
+             LLVMScalableVectorType)>;
 
 def SdotIntrOp :
   ArmSVE_IntrBinaryOverloadedOp<"sdot">,
-  Arguments<(ins LLVM_AnyVector, LLVM_AnyVector, LLVM_AnyVector)>;
+  Arguments<(ins LLVMScalableVectorType, LLVMScalableVectorType,
+             LLVMScalableVectorType)>;
 
 def UdotIntrOp :
   ArmSVE_IntrBinaryOverloadedOp<"udot">,
-  Arguments<(ins LLVM_AnyVector, LLVM_AnyVector, LLVM_AnyVector)>;
+  Arguments<(ins LLVMScalableVectorType, LLVMScalableVectorType,
+             LLVMScalableVectorType)>;
+
+def ScalableMaskedAddIIntrOp :
+  ArmSVE_IntrBinaryOverloadedOp<"add">,
+  Arguments<(ins LLVMScalableVectorType, LLVMScalableVectorType,
+             LLVMScalableVectorType)>;
+
+def ScalableMaskedAddFIntrOp :
+  ArmSVE_IntrBinaryOverloadedOp<"fadd">,
+  Arguments<(ins LLVMScalableVectorType, LLVMScalableVectorType,
+             LLVMScalableVectorType)>;
+
+def ScalableMaskedMulIIntrOp :
+  ArmSVE_IntrBinaryOverloadedOp<"mul">,
+  Arguments<(ins LLVMScalableVectorType, LLVMScalableVectorType,
+             LLVMScalableVectorType)>;
+
+def ScalableMaskedMulFIntrOp :
+  ArmSVE_IntrBinaryOverloadedOp<"fmul">,
+  Arguments<(ins LLVMScalableVectorType, LLVMScalableVectorType,
+             LLVMScalableVectorType)>;
+
+def ScalableMaskedSubIIntrOp :
+  ArmSVE_IntrBinaryOverloadedOp<"sub">,
+  Arguments<(ins LLVMScalableVectorType, LLVMScalableVectorType,
+             LLVMScalableVectorType)>;
+
+def ScalableMaskedSubFIntrOp :
+  ArmSVE_IntrBinaryOverloadedOp<"fsub">,
+  Arguments<(ins LLVMScalableVectorType, LLVMScalableVectorType,
+             LLVMScalableVectorType)>;
+
+def ScalableMaskedSDivIIntrOp :
+  ArmSVE_IntrBinaryOverloadedOp<"sdiv">,
+  Arguments<(ins LLVMScalableVectorType, LLVMScalableVectorType,
+             LLVMScalableVectorType)>;
+
+def ScalableMaskedUDivIIntrOp :
+  ArmSVE_IntrBinaryOverloadedOp<"udiv">,
+  Arguments<(ins LLVMScalableVectorType, LLVMScalableVectorType,
+             LLVMScalableVectorType)>;
+
+def ScalableMaskedDivFIntrOp :
+  ArmSVE_IntrBinaryOverloadedOp<"fdiv">,
+  Arguments<(ins LLVMScalableVectorType, LLVMScalableVectorType,
+             LLVMScalableVectorType)>;
 
 def VectorScaleIntrOp:
   ArmSVE_NonSVEIntrUnaryOverloadedOp<"vscale">;
 
 #endif // ARMSVE_OPS

mlir/lib/Dialect/ArmSVE/IR/ArmSVEDialect.cpp

 #include "mlir/IR/Builders.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 
+static Type getI1SameShape(Type type);
+
 #define GET_OP_CLASSES
 #include "mlir/Dialect/ArmSVE/ArmSVE.cpp.inc"
 
 #define GET_TYPEDEF_CLASSES
 #include "mlir/Dialect/ArmSVE/ArmSVETypes.cpp.inc"
 
 void arm_sve::ArmSVEDialect::initialize() {
   addOperations<
Type arm_sve::ArmSVEDialect::parseType(DialectAsmParser &parser) const {
   parser.emitError(typeLoc, "unknown type in ArmSVE dialect");
   return Type();
 }
 
 void arm_sve::ArmSVEDialect::printType(Type type, DialectAsmPrinter &os) const {
   if (failed(generatedTypePrinter(type, os)))
     llvm_unreachable("unexpected 'arm_sve' type kind");
 }
+
+//===----------------------------------------------------------------------===//
+// ScalableVector versions of general helpers for comparison ops
+//===----------------------------------------------------------------------===//
+
+// Return the scalable vector of the same shape and containing i1.
+static Type getI1SameShape(Type type) {
+  auto i1Type = IntegerType::get(type.getContext(), 1);
+  if (auto sVectorType = type.dyn_cast<arm_sve::ScalableVectorType>())
+    return arm_sve::ScalableVectorType::get(type.getContext(),
+                                            sVectorType.getShape(), i1Type);
+  return nullptr;
+}

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

 }
 
 using SdotOpLowering = OneToOneConvertToLLVMPattern<SdotOp, SdotIntrOp>;
 using SmmlaOpLowering = OneToOneConvertToLLVMPattern<SmmlaOp, SmmlaIntrOp>;
 using UdotOpLowering = OneToOneConvertToLLVMPattern<UdotOp, UdotIntrOp>;
 using UmmlaOpLowering = OneToOneConvertToLLVMPattern<UmmlaOp, UmmlaIntrOp>;
 using VectorScaleOpLowering =
     OneToOneConvertToLLVMPattern<VectorScaleOp, VectorScaleIntrOp>;
+using ScalableMaskedAddIOpLowering =
+    OneToOneConvertToLLVMPattern<ScalableMaskedAddIOp,
+                                 ScalableMaskedAddIIntrOp>;
+using ScalableMaskedAddFOpLowering =
+    OneToOneConvertToLLVMPattern<ScalableMaskedAddFOp,
+                                 ScalableMaskedAddFIntrOp>;
+using ScalableMaskedSubIOpLowering =
+    OneToOneConvertToLLVMPattern<ScalableMaskedSubIOp,
+                                 ScalableMaskedSubIIntrOp>;
+using ScalableMaskedSubFOpLowering =
+    OneToOneConvertToLLVMPattern<ScalableMaskedSubFOp,
+                                 ScalableMaskedSubFIntrOp>;
+using ScalableMaskedMulIOpLowering =
+    OneToOneConvertToLLVMPattern<ScalableMaskedMulIOp,
+                                 ScalableMaskedMulIIntrOp>;
+using ScalableMaskedMulFOpLowering =
+    OneToOneConvertToLLVMPattern<ScalableMaskedMulFOp,
+                                 ScalableMaskedMulFIntrOp>;
+using ScalableMaskedSDivIOpLowering =
+    OneToOneConvertToLLVMPattern<ScalableMaskedSDivIOp,
+                                 ScalableMaskedSDivIIntrOp>;
+using ScalableMaskedUDivIOpLowering =
+    OneToOneConvertToLLVMPattern<ScalableMaskedUDivIOp,
+                                 ScalableMaskedUDivIIntrOp>;
+using ScalableMaskedDivFOpLowering =
+    OneToOneConvertToLLVMPattern<ScalableMaskedDivFOp,
+                                 ScalableMaskedDivFIntrOp>;
 
 static void
 populateBasicSVEArithmeticExportPatterns(LLVMTypeConverter &converter,
                                          OwningRewritePatternList &patterns) {
   // clang-format off
   patterns.add<OneToOneConvertToLLVMPattern<ScalableAddIOp, LLVM::AddOp>,
                OneToOneConvertToLLVMPattern<ScalableAddFOp, LLVM::FAddOp>,
                OneToOneConvertToLLVMPattern<ScalableSubIOp, LLVM::SubOp>,
void mlir::populateArmSVELegalizeForLLVMExportPatterns(
   patterns.add<ForwardOperands<CallOp>,
                ForwardOperands<CallIndirectOp>,
                ForwardOperands<ReturnOp>>(converter,
                                           &converter.getContext());
   patterns.add<SdotOpLowering,
                SmmlaOpLowering,
                UdotOpLowering,
                UmmlaOpLowering,
-               VectorScaleOpLowering>(converter);
+               VectorScaleOpLowering,
+               ScalableMaskedAddIOpLowering,
+               ScalableMaskedAddFOpLowering,
+               ScalableMaskedSubIOpLowering,
+               ScalableMaskedSubFOpLowering,
+               ScalableMaskedMulIOpLowering,
+               ScalableMaskedMulFOpLowering,
+               ScalableMaskedSDivIOpLowering,
+               ScalableMaskedUDivIOpLowering,
+               ScalableMaskedDivFOpLowering>(converter);
   // clang-format on
   populateBasicSVEArithmeticExportPatterns(converter, patterns);
 }
 
 void mlir::configureArmSVELegalizeForExportTarget(
     LLVMConversionTarget &target) {
-  target.addLegalOp<SdotIntrOp, SmmlaIntrOp, UdotIntrOp, UmmlaIntrOp,
-                    VectorScaleIntrOp>();
-  target.addIllegalOp<SdotOp, SmmlaOp, UdotOp, UmmlaOp, VectorScaleOp>();
+  // clang-format off
+  target.addLegalOp<SdotIntrOp,
+                    SmmlaIntrOp,
+                    UdotIntrOp,
+                    UmmlaIntrOp,
+                    VectorScaleIntrOp,
+                    ScalableMaskedAddIIntrOp,
+                    ScalableMaskedAddFIntrOp,
+                    ScalableMaskedSubIIntrOp,
+                    ScalableMaskedSubFIntrOp,
+                    ScalableMaskedMulIIntrOp,
+                    ScalableMaskedMulFIntrOp,
+                    ScalableMaskedSDivIIntrOp,
+                    ScalableMaskedUDivIIntrOp,
+                    ScalableMaskedDivFIntrOp>();
+  target.addIllegalOp<SdotOp,
+                      SmmlaOp,
+                      UdotOp,
+                      UmmlaOp,
+                      VectorScaleOp,
+                      ScalableMaskedAddIOp,
+                      ScalableMaskedAddFOp,
+                      ScalableMaskedSubIOp,
+                      ScalableMaskedSubFOp,
+                      ScalableMaskedMulIOp,
+                      ScalableMaskedMulFOp,
+                      ScalableMaskedSDivIOp,
+                      ScalableMaskedUDivIOp,
+                      ScalableMaskedDivFOp>();
+  // clang-format on
   auto hasScalableVectorType = [](TypeRange types) {
     for (Type type : types)
       if (type.isa<arm_sve::ScalableVectorType>())
         return true;
     return false;
   };
   target.addDynamicallyLegalOp<FuncOp>([hasScalableVectorType](FuncOp op) {
     return !hasScalableVectorType(op.getType().getInputs()) &&

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

func @arm_sve_arithi(%a: !arm_sve.vector<4xi32>,
   // CHECK: llvm.add {{.*}}: !llvm.vec<? x 4 x i32>
   %1 = arm_sve.addi %0, %c : !arm_sve.vector<4xi32>
   // CHECK: llvm.sub {{.*}}: !llvm.vec<? x 4 x i32>
   %2 = arm_sve.subi %1, %d : !arm_sve.vector<4xi32>
   // CHECK: llvm.sdiv {{.*}}: !llvm.vec<? x 4 x i32>
   %3 = arm_sve.divi_signed %2, %e : !arm_sve.vector<4xi32>
   // CHECK: llvm.udiv {{.*}}: !llvm.vec<? x 4 x i32>
   %4 = arm_sve.divi_unsigned %2, %e : !arm_sve.vector<4xi32>
-  return %3 : !arm_sve.vector<4xi32>
+  return %4 : !arm_sve.vector<4xi32>
 }
 
 func @arm_sve_arithf(%a: !arm_sve.vector<4xf32>,
                      %b: !arm_sve.vector<4xf32>,
                      %c: !arm_sve.vector<4xf32>,
                      %d: !arm_sve.vector<4xf32>,
                      %e: !arm_sve.vector<4xf32>) -> !arm_sve.vector<4xf32> {
   // CHECK: llvm.fmul {{.*}}: !llvm.vec<? x 4 x f32>
   %0 = arm_sve.mulf %a, %b : !arm_sve.vector<4xf32>
   // CHECK: llvm.fadd {{.*}}: !llvm.vec<? x 4 x f32>
   %1 = arm_sve.addf %0, %c : !arm_sve.vector<4xf32>
   // CHECK: llvm.fsub {{.*}}: !llvm.vec<? x 4 x f32>
   %2 = arm_sve.subf %1, %d : !arm_sve.vector<4xf32>
   // CHECK: llvm.fdiv {{.*}}: !llvm.vec<? x 4 x f32>
   %3 = arm_sve.divf %2, %e : !arm_sve.vector<4xf32>
   return %3 : !arm_sve.vector<4xf32>
 }
 
+func @arm_sve_arithi_masked(%a: !arm_sve.vector<4xi32>,
+                            %b: !arm_sve.vector<4xi32>,
+                            %c: !arm_sve.vector<4xi32>,
+                            %d: !arm_sve.vector<4xi32>,
+                            %e: !arm_sve.vector<4xi32>,
+                            %mask: !arm_sve.vector<4xi1>
+                            ) -> !arm_sve.vector<4xi32> {
+  // CHECK: arm_sve.intr.add{{.*}}: (!llvm.vec<? x 4 x i1>, !llvm.vec<? x 4 x i32>, !llvm.vec<? x 4 x i32>) -> !llvm.vec<? x 4 x i32>
+  %0 = arm_sve.masked.addi %mask, %a, %b : !arm_sve.vector<4xi1>,
+                                           !arm_sve.vector<4xi32>
+  // CHECK: arm_sve.intr.sub{{.*}}: (!llvm.vec<? x 4 x i1>, !llvm.vec<? x 4 x i32>, !llvm.vec<? x 4 x i32>) -> !llvm.vec<? x 4 x i32>
+  %1 = arm_sve.masked.subi %mask, %0, %c : !arm_sve.vector<4xi1>,
+                                           !arm_sve.vector<4xi32>
+  // CHECK: arm_sve.intr.mul{{.*}}: (!llvm.vec<? x 4 x i1>, !llvm.vec<? x 4 x i32>, !llvm.vec<? x 4 x i32>) -> !llvm.vec<? x 4 x i32>
+  %2 = arm_sve.masked.muli %mask, %1, %d : !arm_sve.vector<4xi1>,
+                                           !arm_sve.vector<4xi32>
+  // CHECK: arm_sve.intr.sdiv{{.*}}: (!llvm.vec<? x 4 x i1>, !llvm.vec<? x 4 x i32>, !llvm.vec<? x 4 x i32>) -> !llvm.vec<? x 4 x i32>
+  %3 = arm_sve.masked.divi_signed %mask, %2, %e : !arm_sve.vector<4xi1>,
+                                                  !arm_sve.vector<4xi32>
+  // CHECK: arm_sve.intr.udiv{{.*}}: (!llvm.vec<? x 4 x i1>, !llvm.vec<? x 4 x i32>, !llvm.vec<? x 4 x i32>) -> !llvm.vec<? x 4 x i32>
+  %4 = arm_sve.masked.divi_unsigned %mask, %3, %e : !arm_sve.vector<4xi1>,
+                                                    !arm_sve.vector<4xi32>
+  return %4 : !arm_sve.vector<4xi32>
+}
+
+func @arm_sve_arithf_masked(%a: !arm_sve.vector<4xf32>,
+                            %b: !arm_sve.vector<4xf32>,
+                            %c: !arm_sve.vector<4xf32>,
+                            %d: !arm_sve.vector<4xf32>,
+                            %e: !arm_sve.vector<4xf32>,
+                            %mask: !arm_sve.vector<4xi1>
+                            ) -> !arm_sve.vector<4xf32> {
+  // CHECK: arm_sve.intr.fadd{{.*}}: (!llvm.vec<? x 4 x i1>, !llvm.vec<? x 4 x f32>, !llvm.vec<? x 4 x f32>) -> !llvm.vec<? x 4 x f32>
+  %0 = arm_sve.masked.addf %mask, %a, %b : !arm_sve.vector<4xi1>,
+                                           !arm_sve.vector<4xf32>
+  // CHECK: arm_sve.intr.fsub{{.*}}: (!llvm.vec<? x 4 x i1>, !llvm.vec<? x 4 x f32>, !llvm.vec<? x 4 x f32>) -> !llvm.vec<? x 4 x f32>
+  %1 = arm_sve.masked.subf %mask, %0, %c : !arm_sve.vector<4xi1>,
+                                           !arm_sve.vector<4xf32>
+  // CHECK: arm_sve.intr.fmul{{.*}}: (!llvm.vec<? x 4 x i1>, !llvm.vec<? x 4 x f32>, !llvm.vec<? x 4 x f32>) -> !llvm.vec<? x 4 x f32>
+  %2 = arm_sve.masked.mulf %mask, %1, %d : !arm_sve.vector<4xi1>,
+                                           !arm_sve.vector<4xf32>
+  // CHECK: arm_sve.intr.fdiv{{.*}}: (!llvm.vec<? x 4 x i1>, !llvm.vec<? x 4 x f32>, !llvm.vec<? x 4 x f32>) -> !llvm.vec<? x 4 x f32>
+  %3 = arm_sve.masked.divf %mask, %2, %e : !arm_sve.vector<4xi1>,
+                                           !arm_sve.vector<4xf32>
+  return %3 : !arm_sve.vector<4xf32>
+}
+
 func @get_vector_scale() -> index {
   // CHECK: arm_sve.vscale
   %0 = arm_sve.vector_scale : index
   return %0 : index
 }

mlir/test/Dialect/ArmSVE/roundtrip.mlir

func @arm_sve_arithf(%a: !arm_sve.vector<4xf32>,
                      %c: !arm_sve.vector<4xf32>) -> !arm_sve.vector<4xf32> {
   // CHECK: arm_sve.mulf {{.*}}: !arm_sve.vector<4xf32>
   %0 = arm_sve.mulf %a, %b : !arm_sve.vector<4xf32>
   // CHECK: arm_sve.addf {{.*}}: !arm_sve.vector<4xf32>
   %1 = arm_sve.addf %0, %c : !arm_sve.vector<4xf32>
   return %1 : !arm_sve.vector<4xf32>
 }
 
+func @arm_sve_masked_arithi(%a: !arm_sve.vector<4xi32>,
+                            %b: !arm_sve.vector<4xi32>,
+                            %c: !arm_sve.vector<4xi32>,
+                            %d: !arm_sve.vector<4xi32>,
+                            %e: !arm_sve.vector<4xi32>,
+                            %mask: !arm_sve.vector<4xi1>)
+                            -> !arm_sve.vector<4xi32> {
+  // CHECK: arm_sve.masked.muli {{.*}}: !arm_sve.vector<4xi1>, !arm_sve.vector
+  %0 = arm_sve.masked.muli %mask, %a, %b : !arm_sve.vector<4xi1>,
+                                           !arm_sve.vector<4xi32>
+  // CHECK: arm_sve.masked.addi {{.*}}: !arm_sve.vector<4xi1>, !arm_sve.vector
+  %1 = arm_sve.masked.addi %mask, %0, %c : !arm_sve.vector<4xi1>,
+                                           !arm_sve.vector<4xi32>
+  // CHECK: arm_sve.masked.subi {{.*}}: !arm_sve.vector<4xi1>, !arm_sve.vector
+  %2 = arm_sve.masked.subi %mask, %1, %d : !arm_sve.vector<4xi1>,
+                                           !arm_sve.vector<4xi32>
+  // CHECK: arm_sve.masked.divi_signed
+  %3 = arm_sve.masked.divi_signed %mask, %2, %e : !arm_sve.vector<4xi1>,
+                                                  !arm_sve.vector<4xi32>
+  // CHECK: arm_sve.masked.divi_unsigned
+  %4 = arm_sve.masked.divi_unsigned %mask, %3, %e : !arm_sve.vector<4xi1>,
+                                                    !arm_sve.vector<4xi32>
+  return %2 : !arm_sve.vector<4xi32>
+}
+
+func @arm_sve_masked_arithf(%a: !arm_sve.vector<4xf32>,
+                            %b: !arm_sve.vector<4xf32>,
+                            %c: !arm_sve.vector<4xf32>,
+                            %d: !arm_sve.vector<4xf32>,
+                            %e: !arm_sve.vector<4xf32>,
+                            %mask: !arm_sve.vector<4xi1>)
+                            -> !arm_sve.vector<4xf32> {
+  // CHECK: arm_sve.masked.mulf {{.*}}: !arm_sve.vector<4xi1>, !arm_sve.vector
+  %0 = arm_sve.masked.mulf %mask, %a, %b : !arm_sve.vector<4xi1>,
+                                           !arm_sve.vector<4xf32>
+  // CHECK: arm_sve.masked.addf {{.*}}: !arm_sve.vector<4xi1>, !arm_sve.vector
+  %1 = arm_sve.masked.addf %mask, %0, %c : !arm_sve.vector<4xi1>,
+                                           !arm_sve.vector<4xf32>
+  // CHECK: arm_sve.masked.subf {{.*}}: !arm_sve.vector<4xi1>, !arm_sve.vector
+  %2 = arm_sve.masked.subf %mask, %1, %d : !arm_sve.vector<4xi1>,
+                                           !arm_sve.vector<4xf32>
+  // CHECK: arm_sve.masked.divf {{.*}}: !arm_sve.vector<4xi1>, !arm_sve.vector
+  %3 = arm_sve.masked.divf %mask, %2, %e : !arm_sve.vector<4xi1>,
+                                           !arm_sve.vector<4xf32>
+  return %3 : !arm_sve.vector<4xf32>
+}
+
 func @get_vector_scale() -> index {
   // CHECK: arm_sve.vector_scale : index
   %0 = arm_sve.vector_scale : index
   return %0 : index
 }

mlir/test/Target/LLVMIR/arm-sve.mlir

llvm.func @arm_sve_arithf(%arg0: !llvm.vec<? x 4 x f32>,
                           -> !llvm.vec<? x 4 x f32> {
   // CHECK: fmul <vscale x 4 x float>
   %0 = llvm.fmul %arg0, %arg1 : !llvm.vec<? x 4 x f32>
   // CHECK: fadd <vscale x 4 x float>
   %1 = llvm.fadd %0, %arg2 : !llvm.vec<? x 4 x f32>
   llvm.return %1 : !llvm.vec<? x 4 x f32>
 }
 
+// CHECK-LABEL: define <vscale x 4 x i32> @arm_sve_arithi_masked
+llvm.func @arm_sve_arithi_masked(%arg0: !llvm.vec<? x 4 x i32>,
+                                 %arg1: !llvm.vec<? x 4 x i32>,
+                                 %arg2: !llvm.vec<? x 4 x i32>,
+                                 %arg3: !llvm.vec<? x 4 x i32>,
+                                 %arg4: !llvm.vec<? x 4 x i32>,
+                                 %arg5: !llvm.vec<? x 4 x i1>)
+                                 -> !llvm.vec<? x 4 x i32> {
+  // CHECK: call <vscale x 4 x i32> @llvm.aarch64.sve.add.nxv4i32
+  %0 = "arm_sve.intr.add"(%arg5, %arg0, %arg1) : (!llvm.vec<? x 4 x i1>,
+                                                  !llvm.vec<? x 4 x i32>,
+                                                  !llvm.vec<? x 4 x i32>)
+                                                  -> !llvm.vec<? x 4 x i32>
+  // CHECK: call <vscale x 4 x i32> @llvm.aarch64.sve.sub.nxv4i32
+  %1 = "arm_sve.intr.sub"(%arg5, %0, %arg1) : (!llvm.vec<? x 4 x i1>,
+                                               !llvm.vec<? x 4 x i32>,
+                                               !llvm.vec<? x 4 x i32>)
+                                               -> !llvm.vec<? x 4 x i32>
+  // CHECK: call <vscale x 4 x i32> @llvm.aarch64.sve.mul.nxv4i32
+  %2 = "arm_sve.intr.mul"(%arg5, %1, %arg3) : (!llvm.vec<? x 4 x i1>,
+                                               !llvm.vec<? x 4 x i32>,
+                                               !llvm.vec<? x 4 x i32>)
+                                               -> !llvm.vec<? x 4 x i32>
+  // CHECK: call <vscale x 4 x i32> @llvm.aarch64.sve.sdiv.nxv4i32
+  %3 = "arm_sve.intr.sdiv"(%arg5, %2, %arg4) : (!llvm.vec<? x 4 x i1>,
+                                               !llvm.vec<? x 4 x i32>,
+                                               !llvm.vec<? x 4 x i32>)
+                                               -> !llvm.vec<? x 4 x i32>
+  // CHECK: call <vscale x 4 x i32> @llvm.aarch64.sve.udiv.nxv4i32
+  %4 = "arm_sve.intr.udiv"(%arg5, %3, %arg4) : (!llvm.vec<? x 4 x i1>,
+                                               !llvm.vec<? x 4 x i32>,
+                                               !llvm.vec<? x 4 x i32>)
+                                               -> !llvm.vec<? x 4 x i32>
+  llvm.return %4 : !llvm.vec<? x 4 x i32>
+}
+
+// CHECK-LABEL: define <vscale x 4 x float> @arm_sve_arithf_masked
+llvm.func @arm_sve_arithf_masked(%arg0: !llvm.vec<? x 4 x f32>,
+                                 %arg1: !llvm.vec<? x 4 x f32>,
+                                 %arg2: !llvm.vec<? x 4 x f32>,
+                                 %arg3: !llvm.vec<? x 4 x f32>,
+                                 %arg4: !llvm.vec<? x 4 x f32>,
+                                 %arg5: !llvm.vec<? x 4 x i1>)
+                                 -> !llvm.vec<? x 4 x f32> {
+  // CHECK: call <vscale x 4 x float> @llvm.aarch64.sve.fadd.nxv4f32
+  %0 = "arm_sve.intr.fadd"(%arg5, %arg0, %arg1) : (!llvm.vec<? x 4 x i1>,
+                                                   !llvm.vec<? x 4 x f32>,
+                                                   !llvm.vec<? x 4 x f32>)
+                                                   -> !llvm.vec<? x 4 x f32>
+  // CHECK: call <vscale x 4 x float> @llvm.aarch64.sve.fsub.nxv4f32
+  %1 = "arm_sve.intr.fsub"(%arg5, %0, %arg2) : (!llvm.vec<? x 4 x i1>,
+                                                !llvm.vec<? x 4 x f32>,
+                                                !llvm.vec<? x 4 x f32>)
+                                                -> !llvm.vec<? x 4 x f32>
+  // CHECK: call <vscale x 4 x float> @llvm.aarch64.sve.fmul.nxv4f32
+  %2 = "arm_sve.intr.fmul"(%arg5, %1, %arg3) : (!llvm.vec<? x 4 x i1>,
+                                                !llvm.vec<? x 4 x f32>,
+                                                !llvm.vec<? x 4 x f32>)
+                                                -> !llvm.vec<? x 4 x f32>
+  // CHECK: call <vscale x 4 x float> @llvm.aarch64.sve.fdiv.nxv4f32
+  %3 = "arm_sve.intr.fdiv"(%arg5, %2, %arg4) : (!llvm.vec<? x 4 x i1>,
+                                                !llvm.vec<? x 4 x f32>,
+                                                !llvm.vec<? x 4 x f32>)
+                                                -> !llvm.vec<? x 4 x f32>
+  llvm.return %3 : !llvm.vec<? x 4 x f32>
+}
+
 // CHECK-LABEL: define i64 @get_vector_scale()
 llvm.func @get_vector_scale() -> i64 {
   // CHECK: call i64 @llvm.vscale.i64()
   %0 = "arm_sve.vscale"() : () -> i64
   llvm.return %0 : i64
 }