[mlir][ArmSVE] Add basic mask generation operation

[jsetoain]

Summary

These arm_sve.cmp functions are needed to generate scalable vector
masks as long as scalable vectors are not part of the standard types.
Once in standard, these can be removed and std.cmp can be used
instead.

[Matt]

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

446

(Likely) Minor typo: "unordered not equal" (dropping "or").

[Matt]

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

533

Minor edit: Use integer comparison operation cmpi (instead of cmpf) in the example for cmpi.

[jsetoain]

Amended documentation.
Fixed bug in parseCmpOp (assuming a vector must be LLVMFixed or LLVMScalable)

[jsetoain]

Added missing empty line

[jsetoain]

Added a slightly more interesting test

[ftynse]

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

406–432

This will create name clashes should you ever need to include ArmSVE.td and StandardOps/Ops.td in the same file, e.g. for DRR patterns. Have you considered exposing the predicates from StandardOps/Ops.td to StandardOps/StandardOpBase.td and reusing them instead of copying?

[jsetoain]

Moved Cmp[F|I]Predicate from Ops.td to StandardOpsBase.td and reuse them in ArmSVE

[jsetoain]

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

406–432

Something like this?

[jsetoain]

This patch breaks mlir-nvidia, and a couple of flang-related builders:

• https://lab.llvm.org/buildbot/#/builders/61/builds/10754
• https://lab.llvm.org/buildbot/#/builders/160/builds/2534
• https://lab.llvm.org/buildbot/#/builders/161/builds/2594
• https://lab.llvm.org/buildbot/#/builders/134/builds/6180

[awarzynski]

Hi @jsetoain.

Thank you for working on this! Sadly, this change is causing multiple BuildBot workers to fail:

• https://lab.llvm.org/buildbot/#/builders/134/builds/6180
• https://lab.llvm.org/buildbot/#/builders/160/builds/2534
• https://lab.llvm.org/buildbot/#/builders/161/builds/2594

I suspect that this change was only tested with static lib builds. Could you please test with BUILD_SHARED_LIBS set to On? Thank you!

*EDIT*
Apologies, I've only noticed your comment *after* posting this.

[jsetoain]

Sorry it took me so long to revert, and thank you for the instructions!

[jsetoain]

Fixed shared library dependency

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

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef ARMSVE_OPS
 #define ARMSVE_OPS
 
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
+include "mlir/Dialect/StandardOps/IR/StandardOpsBase.td"
 include "mlir/Dialect/ArmSVE/ArmSVEOpBase.td"
 
 //===----------------------------------------------------------------------===//
 // ArmSVE dialect definition
 //===----------------------------------------------------------------------===//
 
 def ArmSVE_Dialect : Dialect {
   let name = "arm_sve";

 def ScalableMaskedSDivIOp : ScalableMaskedIOp<"masked.divi_signed",
                                               "signed division">;
 
 def ScalableMaskedUDivIOp : ScalableMaskedIOp<"masked.divi_unsigned",
                                               "unsigned division">;
 
 def ScalableMaskedDivFOp : ScalableMaskedFOp<"masked.divf", "division">;
 
+//===----------------------------------------------------------------------===//
+// ScalableCmpFOp
+//===----------------------------------------------------------------------===//
+
+def ScalableCmpFOp : ArmSVE_Op<"cmpf", [NoSideEffect, SameTypeOperands,
+    TypesMatchWith<"result type has i1 element type and same shape as operands",
+    "lhs", "result", "getI1SameShape($_self)">]> {
+  let summary = "floating-point comparison operation for scalable vectors";
+  let description = [{
+    The `arm_sve.cmpf` operation compares two scalable vectors of floating point
+    elements according to the float comparison rules and the predicate specified
+    by the respective attribute. The predicate defines the type of comparison:
+    (un)orderedness, (in)equality and signed less/greater than (or equal to) as
+    well as predicates that are always true or false. The result is a scalable
+    vector of i1 elements. Unlike `arm_sve.cmpi`, the operands are always
+    treated as signed. The u prefix indicates *unordered* comparison, not
+    unsigned comparison, so "une" means unordered not equal. For the sake of
+    readability by humans, custom assembly form for the operation uses a
+    string-typed attribute for the predicate.  The value of this attribute
+    corresponds to lower-cased name of the predicate constant, e.g., "one" means
+    "ordered not equal".  The string representation of the attribute is merely a
+    syntactic sugar and is converted to an integer attribute by the parser.
+
+    Example:
+
+    ```mlir
+    %r = arm_sve.cmpf oeq, %0, %1 : !arm_sve.vector<4xf32>
+    ```
+  }];
+  let arguments = (ins
+    CmpFPredicateAttr:$predicate,
+    ScalableVectorOf<[AnyFloat]>:$lhs,
+    ScalableVectorOf<[AnyFloat]>:$rhs // TODO: This should support a simple scalar
+  );
+  let results = (outs ScalableVectorOf<[I1]>:$result);
+
+  let builders = [
+    OpBuilder<(ins "CmpFPredicate":$predicate, "Value":$lhs,
+                  "Value":$rhs), [{
+      buildScalableCmpFOp($_builder, $_state, predicate, lhs, rhs);
+    }]>];
+
+  let extraClassDeclaration = [{
+    static StringRef getPredicateAttrName() { return "predicate"; }
+    static CmpFPredicate getPredicateByName(StringRef name);
+
+    CmpFPredicate getPredicate() {
+      return (CmpFPredicate)(*this)->getAttrOfType<IntegerAttr>(
+          getPredicateAttrName()).getInt();
+    }
+  }];
+
+  let verifier = [{ return success(); }];
+
+  let assemblyFormat = "$predicate `,` $lhs `,` $rhs attr-dict `:` type($lhs)";
+}
+
+//===----------------------------------------------------------------------===//
+// ScalableCmpIOp
+//===----------------------------------------------------------------------===//
+
+def ScalableCmpIOp : ArmSVE_Op<"cmpi", [NoSideEffect, SameTypeOperands,
+    TypesMatchWith<"result type has i1 element type and same shape as operands",
+    "lhs", "result", "getI1SameShape($_self)">]> {
+  let summary = "integer comparison operation for scalable vectors";
+  let description = [{
+    The `arm_sve.cmpi` operation compares two scalable vectors of integer
+    elements according to the predicate specified by the respective attribute.
+
+    The predicate defines the type of comparison:
+
+    -   equal (mnemonic: `"eq"`; integer value: `0`)
+    -   not equal (mnemonic: `"ne"`; integer value: `1`)
+    -   signed less than (mnemonic: `"slt"`; integer value: `2`)
+    -   signed less than or equal (mnemonic: `"sle"`; integer value: `3`)
+    -   signed greater than (mnemonic: `"sgt"`; integer value: `4`)
+    -   signed greater than or equal (mnemonic: `"sge"`; integer value: `5`)
+    -   unsigned less than (mnemonic: `"ult"`; integer value: `6`)
+    -   unsigned less than or equal (mnemonic: `"ule"`; integer value: `7`)
+    -   unsigned greater than (mnemonic: `"ugt"`; integer value: `8`)
+    -   unsigned greater than or equal (mnemonic: `"uge"`; integer value: `9`)
+
+    Example:
+
+    ```mlir
+    %r = arm_sve.cmpi uge, %0, %1 : !arm_sve.vector<4xi32>
+    ```
+  }];
+
+  let arguments = (ins
+      CmpIPredicateAttr:$predicate,
+      ScalableVectorOf<[I8, I16, I32, I64]>:$lhs,
+      ScalableVectorOf<[I8, I16, I32, I64]>:$rhs
+  );
+  let results = (outs ScalableVectorOf<[I1]>:$result);
+
+  let builders = [
+    OpBuilder<(ins "CmpIPredicate":$predicate, "Value":$lhs,
+                 "Value":$rhs), [{
+      buildScalableCmpIOp($_builder, $_state, predicate, lhs, rhs);
+    }]>];
+
+  let extraClassDeclaration = [{
+    static StringRef getPredicateAttrName() { return "predicate"; }
+    static CmpIPredicate getPredicateByName(StringRef name);
+
+    CmpIPredicate getPredicate() {
+      return (CmpIPredicate)(*this)->getAttrOfType<IntegerAttr>(
+          getPredicateAttrName()).getInt();
+    }
+  }];
+
+  let verifier = [{ return success(); }];
+
+  let assemblyFormat = "$predicate `,` $lhs `,` $rhs attr-dict `:` type($lhs)";
+}
+
 def UmmlaIntrOp :
   ArmSVE_IntrBinaryOverloadedOp<"ummla">,
   Arguments<(ins LLVMScalableVectorType, LLVMScalableVectorType,
              LLVMScalableVectorType)>;
 
 def SmmlaIntrOp :
   ArmSVE_IntrBinaryOverloadedOp<"smmla">,
   Arguments<(ins LLVMScalableVectorType, LLVMScalableVectorType,
              LLVMScalableVectorType)>;
 
 def SdotIntrOp :
   ArmSVE_IntrBinaryOverloadedOp<"sdot">,
   Arguments<(ins LLVMScalableVectorType, LLVMScalableVectorType,
              LLVMScalableVectorType)>;
 
 def UdotIntrOp :
   ArmSVE_IntrBinaryOverloadedOp<"udot">,
   Arguments<(ins LLVMScalableVectorType, LLVMScalableVectorType,
              LLVMScalableVectorType)>;
 
 def ScalableMaskedAddIIntrOp :
   ArmSVE_IntrBinaryOverloadedOp<"add">,
   Arguments<(ins LLVMScalableVectorType, LLVMScalableVectorType,

mlir/include/mlir/Dialect/ArmSVE/ArmSVEDialect.h

 #define MLIR_DIALECT_ARMSVE_ARMSVEDIALECT_H
 
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 #include "mlir/Dialect/ArmSVE/ArmSVEDialect.h.inc"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
 
 #define GET_TYPEDEF_CLASSES
 #include "mlir/Dialect/ArmSVE/ArmSVETypes.h.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/ArmSVE/ArmSVE.h.inc"
 
 #endif // MLIR_DIALECT_ARMSVE_ARMSVEDIALECT_H

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

let description = [{
     ```
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // CmpFOp
 //===----------------------------------------------------------------------===//
 
-// The predicate indicates the type of the comparison to perform:
-// (un)orderedness, (in)equality and less/greater than (or equal to) as
-// well as predicates that are always true or false.
-def CMPF_P_FALSE   : I64EnumAttrCase<"AlwaysFalse", 0, "false">;
-def CMPF_P_OEQ     : I64EnumAttrCase<"OEQ", 1, "oeq">;
-def CMPF_P_OGT     : I64EnumAttrCase<"OGT", 2, "ogt">;
-def CMPF_P_OGE     : I64EnumAttrCase<"OGE", 3, "oge">;
-def CMPF_P_OLT     : I64EnumAttrCase<"OLT", 4, "olt">;
-def CMPF_P_OLE     : I64EnumAttrCase<"OLE", 5, "ole">;
-def CMPF_P_ONE     : I64EnumAttrCase<"ONE", 6, "one">;
-def CMPF_P_ORD     : I64EnumAttrCase<"ORD", 7, "ord">;
-def CMPF_P_UEQ     : I64EnumAttrCase<"UEQ", 8, "ueq">;
-def CMPF_P_UGT     : I64EnumAttrCase<"UGT", 9, "ugt">;
-def CMPF_P_UGE     : I64EnumAttrCase<"UGE", 10, "uge">;
-def CMPF_P_ULT     : I64EnumAttrCase<"ULT", 11, "ult">;
-def CMPF_P_ULE     : I64EnumAttrCase<"ULE", 12, "ule">;
-def CMPF_P_UNE     : I64EnumAttrCase<"UNE", 13, "une">;
-def CMPF_P_UNO     : I64EnumAttrCase<"UNO", 14, "uno">;
-def CMPF_P_TRUE    : I64EnumAttrCase<"AlwaysTrue", 15, "true">;
-
-def CmpFPredicateAttr : I64EnumAttr<
-    "CmpFPredicate", "",
-    [CMPF_P_FALSE, CMPF_P_OEQ, CMPF_P_OGT, CMPF_P_OGE, CMPF_P_OLT, CMPF_P_OLE,
-     CMPF_P_ONE, CMPF_P_ORD, CMPF_P_UEQ, CMPF_P_UGT, CMPF_P_UGE, CMPF_P_ULT,
-     CMPF_P_ULE, CMPF_P_UNE, CMPF_P_UNO, CMPF_P_TRUE]> {
-  let cppNamespace = "::mlir";
-}
-
 def CmpFOp : Std_Op<"cmpf", [NoSideEffect, SameTypeOperands,
     DeclareOpInterfaceMethods<VectorUnrollOpInterface>, TypesMatchWith<
     "result type has i1 element type and same shape as operands",
     "lhs", "result", "getI1SameShape($_self)">] # ElementwiseMappable.traits> {
   let summary = "floating-point comparison operation";
   let description = [{
     The `cmpf` operation compares its two operands according to the float
     comparison rules and the predicate specified by the respective attribute.
def CmpFOp : Std_Op<"cmpf", [NoSideEffect, SameTypeOperands,
 
   let assemblyFormat = "$predicate `,` $lhs `,` $rhs attr-dict `:` type($lhs)";
 }
 
 //===----------------------------------------------------------------------===//
 // CmpIOp
 //===----------------------------------------------------------------------===//
 
-def CMPI_P_EQ  : I64EnumAttrCase<"eq", 0>;
-def CMPI_P_NE  : I64EnumAttrCase<"ne", 1>;
-def CMPI_P_SLT : I64EnumAttrCase<"slt", 2>;
-def CMPI_P_SLE : I64EnumAttrCase<"sle", 3>;
-def CMPI_P_SGT : I64EnumAttrCase<"sgt", 4>;
-def CMPI_P_SGE : I64EnumAttrCase<"sge", 5>;
-def CMPI_P_ULT : I64EnumAttrCase<"ult", 6>;
-def CMPI_P_ULE : I64EnumAttrCase<"ule", 7>;
-def CMPI_P_UGT : I64EnumAttrCase<"ugt", 8>;
-def CMPI_P_UGE : I64EnumAttrCase<"uge", 9>;
-
-def CmpIPredicateAttr : I64EnumAttr<
-    "CmpIPredicate", "",
-    [CMPI_P_EQ, CMPI_P_NE, CMPI_P_SLT, CMPI_P_SLE, CMPI_P_SGT,
-     CMPI_P_SGE, CMPI_P_ULT, CMPI_P_ULE, CMPI_P_UGT, CMPI_P_UGE]> {
-  let cppNamespace = "::mlir";
-}
-
 def CmpIOp : Std_Op<"cmpi", [NoSideEffect, SameTypeOperands,
     DeclareOpInterfaceMethods<VectorUnrollOpInterface>, TypesMatchWith<
     "result type has i1 element type and same shape as operands",
     "lhs", "result", "getI1SameShape($_self)">] # ElementwiseMappable.traits> {
   let summary = "integer comparison operation";
   let description = [{
     The `cmpi` operation is a generic comparison for integer-like types. Its two
     arguments can be integers, vectors or tensors thereof as long as their types

mlir/include/mlir/Dialect/StandardOps/IR/StandardOpsBase.td

def AtomicRMWKindAttr : I64EnumAttr<
     "AtomicRMWKind", "",
     [ATOMIC_RMW_KIND_ADDF, ATOMIC_RMW_KIND_ADDI, ATOMIC_RMW_KIND_ASSIGN,
      ATOMIC_RMW_KIND_MAXF, ATOMIC_RMW_KIND_MAXS, ATOMIC_RMW_KIND_MAXU,
      ATOMIC_RMW_KIND_MINF, ATOMIC_RMW_KIND_MINS, ATOMIC_RMW_KIND_MINU,
      ATOMIC_RMW_KIND_MULF, ATOMIC_RMW_KIND_MULI]> {
   let cppNamespace = "::mlir";
 }
 
+// The predicate indicates the type of the comparison to perform:
+// (un)orderedness, (in)equality and less/greater than (or equal to) as
+// well as predicates that are always true or false.
+def CMPF_P_FALSE   : I64EnumAttrCase<"AlwaysFalse", 0, "false">;
+def CMPF_P_OEQ     : I64EnumAttrCase<"OEQ", 1, "oeq">;
+def CMPF_P_OGT     : I64EnumAttrCase<"OGT", 2, "ogt">;
+def CMPF_P_OGE     : I64EnumAttrCase<"OGE", 3, "oge">;
+def CMPF_P_OLT     : I64EnumAttrCase<"OLT", 4, "olt">;
+def CMPF_P_OLE     : I64EnumAttrCase<"OLE", 5, "ole">;
+def CMPF_P_ONE     : I64EnumAttrCase<"ONE", 6, "one">;
+def CMPF_P_ORD     : I64EnumAttrCase<"ORD", 7, "ord">;
+def CMPF_P_UEQ     : I64EnumAttrCase<"UEQ", 8, "ueq">;
+def CMPF_P_UGT     : I64EnumAttrCase<"UGT", 9, "ugt">;
+def CMPF_P_UGE     : I64EnumAttrCase<"UGE", 10, "uge">;
+def CMPF_P_ULT     : I64EnumAttrCase<"ULT", 11, "ult">;
+def CMPF_P_ULE     : I64EnumAttrCase<"ULE", 12, "ule">;
+def CMPF_P_UNE     : I64EnumAttrCase<"UNE", 13, "une">;
+def CMPF_P_UNO     : I64EnumAttrCase<"UNO", 14, "uno">;
+def CMPF_P_TRUE    : I64EnumAttrCase<"AlwaysTrue", 15, "true">;
+
+def CmpFPredicateAttr : I64EnumAttr<
+    "CmpFPredicate", "",
+    [CMPF_P_FALSE, CMPF_P_OEQ, CMPF_P_OGT, CMPF_P_OGE, CMPF_P_OLT, CMPF_P_OLE,
+     CMPF_P_ONE, CMPF_P_ORD, CMPF_P_UEQ, CMPF_P_UGT, CMPF_P_UGE, CMPF_P_ULT,
+     CMPF_P_ULE, CMPF_P_UNE, CMPF_P_UNO, CMPF_P_TRUE]> {
+  let cppNamespace = "::mlir";
+}
+
+def CMPI_P_EQ  : I64EnumAttrCase<"eq", 0>;
+def CMPI_P_NE  : I64EnumAttrCase<"ne", 1>;
+def CMPI_P_SLT : I64EnumAttrCase<"slt", 2>;
+def CMPI_P_SLE : I64EnumAttrCase<"sle", 3>;
+def CMPI_P_SGT : I64EnumAttrCase<"sgt", 4>;
+def CMPI_P_SGE : I64EnumAttrCase<"sge", 5>;
+def CMPI_P_ULT : I64EnumAttrCase<"ult", 6>;
+def CMPI_P_ULE : I64EnumAttrCase<"ule", 7>;
+def CMPI_P_UGT : I64EnumAttrCase<"ugt", 8>;
+def CMPI_P_UGE : I64EnumAttrCase<"uge", 9>;
+
+def CmpIPredicateAttr : I64EnumAttr<
+    "CmpIPredicate", "",
+    [CMPI_P_EQ, CMPI_P_NE, CMPI_P_SLT, CMPI_P_SLE, CMPI_P_SGT,
+     CMPI_P_SGE, CMPI_P_ULT, CMPI_P_ULE, CMPI_P_UGT, CMPI_P_UGE]> {
+  let cppNamespace = "::mlir";
+}
+
 #endif // STANDARD_OPS_BASE

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

 #include "mlir/Dialect/Vector/VectorOps.h"
 #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;
+using namespace arm_sve;
 
 static Type getI1SameShape(Type type);
+static void buildScalableCmpIOp(OpBuilder &build, OperationState &result,
+                                CmpIPredicate predicate, Value lhs, Value rhs);
+static void buildScalableCmpFOp(OpBuilder &build, OperationState &result,
+                                CmpFPredicate predicate, Value lhs, Value rhs);
 
 #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() {
+void ArmSVEDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/ArmSVE/ArmSVE.cpp.inc"
       >();
   addTypes<
 #define GET_TYPEDEF_LIST
 #include "mlir/Dialect/ArmSVE/ArmSVETypes.cpp.inc"
       >();
 }
 
 //===----------------------------------------------------------------------===//
 // ScalableVectorType
 //===----------------------------------------------------------------------===//
 
-Type arm_sve::ArmSVEDialect::parseType(DialectAsmParser &parser) const {
+Type ArmSVEDialect::parseType(DialectAsmParser &parser) const {
   llvm::SMLoc typeLoc = parser.getCurrentLocation();
   {
     Type genType;
     auto parseResult = generatedTypeParser(parser.getBuilder().getContext(),
                                            parser, "vector", genType);
     if (parseResult.hasValue())
       return genType;
   }
   parser.emitError(typeLoc, "unknown type in ArmSVE dialect");
   return Type();
 }
 
-void arm_sve::ArmSVEDialect::printType(Type type, DialectAsmPrinter &os) const {
+void 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);
+  if (auto sVectorType = type.dyn_cast<ScalableVectorType>())
+    return ScalableVectorType::get(type.getContext(), sVectorType.getShape(),
+                                   i1Type);
   return nullptr;
 }
+
+//===----------------------------------------------------------------------===//
+// CmpFOp
+//===----------------------------------------------------------------------===//
+
+static void buildScalableCmpFOp(OpBuilder &build, OperationState &result,
+                                CmpFPredicate predicate, Value lhs, Value rhs) {
+  result.addOperands({lhs, rhs});
+  result.types.push_back(getI1SameShape(lhs.getType()));
+  result.addAttribute(ScalableCmpFOp::getPredicateAttrName(),
+                      build.getI64IntegerAttr(static_cast<int64_t>(predicate)));
+}
+
+static void buildScalableCmpIOp(OpBuilder &build, OperationState &result,
+                                CmpIPredicate predicate, Value lhs, Value rhs) {
+  result.addOperands({lhs, rhs});
+  result.types.push_back(getI1SameShape(lhs.getType()));
+  result.addAttribute(ScalableCmpIOp::getPredicateAttrName(),
+                      build.getI64IntegerAttr(static_cast<int64_t>(predicate)));
+}

mlir/lib/Dialect/ArmSVE/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRArmSVE
   ArmSVEDialect.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/ArmSVE
 
   DEPENDS
   MLIRArmSVEIncGen
 
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRLLVMIR
+  MLIRStandard
   MLIRSideEffectInterfaces
   )

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

target.addIllegalOp<ScalableAddIOp,
                       ScalableMulIOp,
                       ScalableMulFOp,
                       ScalableSDivIOp,
                       ScalableUDivIOp,
                       ScalableDivFOp>();
   // clang-format on
 }
 
+static void
+populateSVEMaskGenerationExportPatterns(LLVMTypeConverter &converter,
+                                        OwningRewritePatternList &patterns) {
+  // clang-format off
+  patterns.add<OneToOneConvertToLLVMPattern<ScalableCmpFOp, LLVM::FCmpOp>,
+               OneToOneConvertToLLVMPattern<ScalableCmpIOp, LLVM::ICmpOp>
+              >(converter);
+  // clang-format on
+}
+
+static void
+configureSVEMaskGenerationLegalizations(LLVMConversionTarget &target) {
+  // clang-format off
+  target.addIllegalOp<ScalableCmpFOp,
+                      ScalableCmpIOp>();
+  // clang-format on
+}
+
 /// Populate the given list with patterns that convert from ArmSVE to LLVM.
 void mlir::populateArmSVELegalizeForLLVMExportPatterns(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns) {
   // Populate conversion patterns
   // Remove any ArmSVE-specific types from function signatures and results.
   populateFuncOpTypeConversionPattern(patterns, converter);
   converter.addConversion([&converter](ScalableVectorType svType) {
     return convertScalableVectorTypeToLLVM(svType, converter);
patterns.add<SdotOpLowering,
                ScalableMaskedSubFOpLowering,
                ScalableMaskedMulIOpLowering,
                ScalableMaskedMulFOpLowering,
                ScalableMaskedSDivIOpLowering,
                ScalableMaskedUDivIOpLowering,
                ScalableMaskedDivFOpLowering>(converter);
   // clang-format on
   populateBasicSVEArithmeticExportPatterns(converter, patterns);
+  populateSVEMaskGenerationExportPatterns(converter, patterns);
 }
 
 void mlir::configureArmSVELegalizeForExportTarget(
     LLVMConversionTarget &target) {
   // clang-format off
   target.addLegalOp<SdotIntrOp,
                     SmmlaIntrOp,
                     UdotIntrOp,
return !hasScalableVectorType(op.getType().getInputs()) &&
            !hasScalableVectorType(op.getType().getResults());
   });
   target.addDynamicallyLegalOp<CallOp, CallIndirectOp, ReturnOp>(
       [hasScalableVectorType](Operation *op) {
         return !hasScalableVectorType(op->getOperandTypes()) &&
                !hasScalableVectorType(op->getResultTypes());
       });
   configureBasicSVEArithmeticLegalizations(target);
+  configureSVEMaskGenerationLegalizations(target);
 }

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

result.attributes.set("predicate",
                         parser.getBuilder().getI64IntegerAttr(predicateValue));
 
   // The result type is either i1 or a vector type <? x i1> if the inputs are
   // vectors.
   Type resultType = IntegerType::get(builder.getContext(), 1);
   if (!isCompatibleType(type))
     return parser.emitError(trailingTypeLoc,
                             "expected LLVM dialect-compatible type");
-  if (LLVM::isCompatibleVectorType(type))
+  if (LLVM::isCompatibleVectorType(type)) {
+    if (type.isa<LLVM::LLVMScalableVectorType>()) {
+      resultType = LLVM::LLVMScalableVectorType::get(
+          resultType, LLVM::getVectorNumElements(type).getKnownMinValue());
+    } else {
       resultType = LLVM::getFixedVectorType(
           resultType, LLVM::getVectorNumElements(type).getFixedValue());
-  assert(!type.isa<LLVM::LLVMScalableVectorType>() &&
-         "unhandled scalable vector");
+    }
+  }
 
   result.addTypes({resultType});
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Printing/parsing for LLVM::AllocaOp.
 //===----------------------------------------------------------------------===//

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

func @arm_sve_arithf_masked(%a: !arm_sve.vector<4xf32>,
   %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 @arm_sve_mask_genf(%a: !arm_sve.vector<4xf32>,
+                        %b: !arm_sve.vector<4xf32>)
+                        -> !arm_sve.vector<4xi1> {
+  // CHECK: llvm.fcmp "oeq" {{.*}}: !llvm.vec<? x 4 x f32>
+  %0 = arm_sve.cmpf oeq, %a, %b : !arm_sve.vector<4xf32>
+  return %0 : !arm_sve.vector<4xi1>
+}
+
+func @arm_sve_mask_geni(%a: !arm_sve.vector<4xi32>,
+                        %b: !arm_sve.vector<4xi32>)
+                        -> !arm_sve.vector<4xi1> {
+  // CHECK: llvm.icmp "uge" {{.*}}: !llvm.vec<? x 4 x i32>
+  %0 = arm_sve.cmpi uge, %a, %b : !arm_sve.vector<4xi32>
+  return %0 : !arm_sve.vector<4xi1>
+}
+
+func @arm_sve_abs_diff(%a: !arm_sve.vector<4xi32>,
+                       %b: !arm_sve.vector<4xi32>)
+                       -> !arm_sve.vector<4xi32> {
+  // CHECK: llvm.sub {{.*}}: !llvm.vec<? x 4 x i32>
+  %z = arm_sve.subi %a, %a : !arm_sve.vector<4xi32>
+  // CHECK: llvm.icmp "sge" {{.*}}: !llvm.vec<? x 4 x i32>
+  %agb = arm_sve.cmpi sge, %a, %b : !arm_sve.vector<4xi32>
+  // CHECK: llvm.icmp "slt" {{.*}}: !llvm.vec<? x 4 x i32>
+  %bga = arm_sve.cmpi slt, %a, %b : !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>
+  %0 = arm_sve.masked.subi %agb, %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 %bga, %b, %a : !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>
+  %2 = arm_sve.masked.addi %agb, %z, %0 : !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>
+  %3 = arm_sve.masked.addi %bga, %2, %1 : !arm_sve.vector<4xi1>,
+                                          !arm_sve.vector<4xi32>
+  return %3 : !arm_sve.vector<4xi32>
+}
+
 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_masked_arithf(%a: !arm_sve.vector<4xf32>,
   %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 @arm_sve_mask_genf(%a: !arm_sve.vector<4xf32>,
+                        %b: !arm_sve.vector<4xf32>)
+                        -> !arm_sve.vector<4xi1> {
+  // CHECK: arm_sve.cmpf oeq, {{.*}}: !arm_sve.vector<4xf32>
+  %0 = arm_sve.cmpf oeq, %a, %b : !arm_sve.vector<4xf32>
+  return %0 : !arm_sve.vector<4xi1>
+}
+
+func @arm_sve_mask_geni(%a: !arm_sve.vector<4xi32>,
+                        %b: !arm_sve.vector<4xi32>)
+                        -> !arm_sve.vector<4xi1> {
+  // CHECK: arm_sve.cmpi uge, {{.*}}: !arm_sve.vector<4xi32>
+  %0 = arm_sve.cmpi uge, %a, %b : !arm_sve.vector<4xi32>
+  return %0 : !arm_sve.vector<4xi1>
+}
+
 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_masked(%arg0: !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 <vscale x 4 x i1> @arm_sve_mask_genf
+llvm.func @arm_sve_mask_genf(%arg0: !llvm.vec<? x 4 x f32>,
+                             %arg1: !llvm.vec<? x 4 x f32>)
+                             -> !llvm.vec<? x 4 x i1> {
+  // CHECK: fcmp oeq <vscale x 4 x float>
+  %0 = llvm.fcmp "oeq" %arg0, %arg1 : !llvm.vec<? x 4 x f32>
+  llvm.return %0 : !llvm.vec<? x 4 x i1>
+}
+
+// CHECK-LABEL: define <vscale x 4 x i1> @arm_sve_mask_geni
+llvm.func @arm_sve_mask_geni(%arg0: !llvm.vec<? x 4 x i32>,
+                             %arg1: !llvm.vec<? x 4 x i32>)
+                             -> !llvm.vec<? x 4 x i1> {
+  // CHECK: icmp uge <vscale x 4 x i32>
+  %0 = llvm.icmp "uge" %arg0, %arg1 : !llvm.vec<? x 4 x i32>
+  llvm.return %0 : !llvm.vec<? x 4 x i1>
+}
+
+// CHECK-LABEL: define <vscale x 4 x i32> @arm_sve_abs_diff
+llvm.func @arm_sve_abs_diff(%arg0: !llvm.vec<? x 4 x i32>,
+                            %arg1: !llvm.vec<? x 4 x i32>)
+                            -> !llvm.vec<? x 4 x i32> {
+  // CHECK: sub <vscale x 4 x i32>
+  %0 = llvm.sub %arg0, %arg0  : !llvm.vec<? x 4 x i32>
+  // CHECK: icmp sge <vscale x 4 x i32>
+  %1 = llvm.icmp "sge" %arg0, %arg1 : !llvm.vec<? x 4 x i32>
+  // CHECK: icmp slt <vscale x 4 x i32>
+  %2 = llvm.icmp "slt" %arg0, %arg1 : !llvm.vec<? x 4 x i32>
+  // CHECK: call <vscale x 4 x i32> @llvm.aarch64.sve.sub.nxv4i32
+  %3 = "arm_sve.intr.sub"(%1, %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
+  %4 = "arm_sve.intr.sub"(%2, %arg1, %arg0) : (!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.add.nxv4i32
+  %5 = "arm_sve.intr.add"(%1, %0, %3) : (!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.add.nxv4i32
+  %6 = "arm_sve.intr.add"(%2, %5, %4) : (!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 %6 : !llvm.vec<? x 4 x i32>
+}
+
 // 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
 }