[mlir][ArmSVE] Add basic arithmetic operations

[jsetoain]

Summary

While we figure out how to best add Standard support for scalable
vectors, these instructions provide a workaround for basic arithmetic
between scalable vectors.

[nicolasvasilache]

Is there an integration test running on some simulator that could be added in the future (similar to what we have for AVX512 and AMX) ?

[ftynse]

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

311–315

We shouldn't be needing this if assemblyFormat is used.

512–518

Nit: I wonder if it isn't simpler to define a ScalableFPOp class and factor out all common parts (arguments, results, assembly format, traits) into it so that concrete classes only need name + doc.

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

49–55 ↗

(On Diff #338798)

Please add tests for user-visible error messages.

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

48

Could we also check the converted type at least once?

[aartbik]

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

354

sub commutative?

378

same

450

div commutative?

474

same

498

same

[jsetoain]

Address comments

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

354

Yes, I noticed while refactoring. Too much auto-pilot :-/

512–518

Is this alright or did I over do it? I'm happy to move the summary and description back to each definition if this is too much.

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

49–55 ↗

(On Diff #338798)

I've dropped this part of the change, it's going to require more than this to avoid the assert, I'd rather move it to its own patch.

[jsetoain]

In D100837#2711320, @nicolasvasilache wrote:

Is there an integration test running on some simulator that could be added in the future (similar to what we have for AVX512 and AMX) ?

I believe we can use qemu, but I wouldn't know how to write the CI test. I'll ask around just in case it already exists, otherwise I'll need some time to look into it.

[jsetoain]

I can't land revisions myself, if nobody else has any issues with the patch, please do that when ready. Thanks!

[kiranchandramohan]

@jsetoain I can land this for you.

[ftynse]

In D100837#2738134, @kiranchandramohan wrote:

@jsetoain I can land this for you.

I've done this already.

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

class ArmSVE_IntrBinaryOverloadedOp<string mnemonic,
   LLVM_IntrOpBase</*Dialect dialect=*/ArmSVE_Dialect,
                   /*string opName=*/"intr." # mnemonic,
                   /*string enumName=*/"aarch64_sve_" # !subst(".", "_", mnemonic),
                   /*list<int> overloadedResults=*/[0],
                   /*list<int> overloadedOperands=*/[], // defined by result overload
                   /*list<OpTrait> traits=*/traits,
                   /*int numResults=*/1>;
 
+class ScalableFOp<string mnemonic, string op_description,
+                  list<OpTrait> traits = []> :
+  ArmSVE_Op<mnemonic, !listconcat(traits,
+                       [AllTypesMatch<["src1", "src2", "dst"]>])> {
+  let summary = op_description # " for scalable vectors of floats";
+  let description = [{
+    The `arm_sve.}] # mnemonic # [{` operations takes two scalable vectors and
+    returns one scalable vector with the result of the }] # op_description # [{.
+  }];
+  let arguments = (ins
+          ScalableVectorOf<[AnyFloat]>:$src1,
+          ScalableVectorOf<[AnyFloat]>:$src2
+  );
+  let results = (outs ScalableVectorOf<[AnyFloat]>:$dst);
+  let assemblyFormat =
+    "$src1 `,` $src2 attr-dict `:` type($src1)";
+}
+
+class ScalableIOp<string mnemonic, string op_description,
+                  list<OpTrait> traits = []> :
+  ArmSVE_Op<mnemonic, !listconcat(traits,
+                       [AllTypesMatch<["src1", "src2", "dst"]>])> {
+  let summary = op_description # " for scalable vectors of integers";
+  let description = [{
+    The `arm_sve.}] # mnemonic # [{` operation takes two scalable vectors and
+    returns one scalable vector with the result of the }] # op_description # [{.
+  }];
+  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)";
+}
+
 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.
let description = [{
     times, once on each <4xi32> segment of the scalable vector. The vector_scale
     op can be used to calculate the step in vector-length agnostic (VLA) loops.
   }];
   let results = (outs Index:$res);
   let assemblyFormat =
     "attr-dict `:` type($res)";
 }
 
+
+def ScalableAddIOp : ScalableIOp<"addi", "addition", [Commutative]>;
+
+def ScalableAddFOp : ScalableFOp<"addf", "addition", [Commutative]>;
+
+def ScalableSubIOp : ScalableIOp<"subi", "subtraction">;
+
+def ScalableSubFOp : ScalableFOp<"subf", "subtraction">;
+
+def ScalableMulIOp : ScalableIOp<"muli", "multiplication", [Commutative]>;
+
+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 UmmlaIntrOp :
   ArmSVE_IntrBinaryOverloadedOp<"ummla">,
   Arguments<(ins LLVM_AnyVector, LLVM_AnyVector, LLVM_AnyVector)>;
 
 def SmmlaIntrOp :
   ArmSVE_IntrBinaryOverloadedOp<"smmla">,
   Arguments<(ins LLVM_AnyVector, LLVM_AnyVector, LLVM_AnyVector)>;
 
 def SdotIntrOp :
   ArmSVE_IntrBinaryOverloadedOp<"sdot">,
   Arguments<(ins LLVM_AnyVector, LLVM_AnyVector, LLVM_AnyVector)>;
 
 def UdotIntrOp :
   ArmSVE_IntrBinaryOverloadedOp<"udot">,
   Arguments<(ins LLVM_AnyVector, LLVM_AnyVector, LLVM_AnyVector)>;
 
 def VectorScaleIntrOp:
   ArmSVE_NonSVEIntrUnaryOverloadedOp<"vscale">;
 
 #endif // ARMSVE_OPS

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>;
 
+static void
+populateBasicSVEArithmeticExportPatterns(LLVMTypeConverter &converter,
+                                         OwningRewritePatternList &patterns) {
+  // clang-format off
+  patterns.add<OneToOneConvertToLLVMPattern<ScalableAddIOp, LLVM::AddOp>,
+               OneToOneConvertToLLVMPattern<ScalableAddFOp, LLVM::FAddOp>,
+               OneToOneConvertToLLVMPattern<ScalableSubIOp, LLVM::SubOp>,
+               OneToOneConvertToLLVMPattern<ScalableSubFOp, LLVM::FSubOp>,
+               OneToOneConvertToLLVMPattern<ScalableMulIOp, LLVM::MulOp>,
+               OneToOneConvertToLLVMPattern<ScalableMulFOp, LLVM::FMulOp>,
+               OneToOneConvertToLLVMPattern<ScalableSDivIOp, LLVM::SDivOp>,
+               OneToOneConvertToLLVMPattern<ScalableUDivIOp, LLVM::UDivOp>,
+               OneToOneConvertToLLVMPattern<ScalableDivFOp, LLVM::FDivOp>
+              >(converter);
+  // clang-format on
+}
+
+static void
+configureBasicSVEArithmeticLegalizations(LLVMConversionTarget &target) {
+  // clang-format off
+  target.addIllegalOp<ScalableAddIOp,
+                      ScalableAddFOp,
+                      ScalableSubIOp,
+                      ScalableSubFOp,
+                      ScalableMulIOp,
+                      ScalableMulFOp,
+                      ScalableSDivIOp,
+                      ScalableUDivIOp,
+                      ScalableDivFOp>();
+  // 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);
   });
   converter.addSourceMaterialization(addUnrealizedCast);
 
   // clang-format off
   patterns.add<ForwardOperands<CallOp>,
                ForwardOperands<CallIndirectOp>,
                ForwardOperands<ReturnOp>>(converter,
                                           &converter.getContext());
   patterns.add<SdotOpLowering,
                SmmlaOpLowering,
                UdotOpLowering,
                UmmlaOpLowering,
                VectorScaleOpLowering>(converter);
   // clang-format on
+  populateBasicSVEArithmeticExportPatterns(converter, patterns);
 }
 
 void mlir::configureArmSVELegalizeForExportTarget(
     LLVMConversionTarget &target) {
-  target.addLegalOp<SdotIntrOp>();
-  target.addIllegalOp<SdotOp>();
-  target.addLegalOp<SmmlaIntrOp>();
-  target.addIllegalOp<SmmlaOp>();
-  target.addLegalOp<UdotIntrOp>();
-  target.addIllegalOp<UdotOp>();
-  target.addLegalOp<UmmlaIntrOp>();
-  target.addIllegalOp<UmmlaOp>();
-  target.addLegalOp<VectorScaleIntrOp>();
-  target.addIllegalOp<VectorScaleOp>();
+  target.addLegalOp<SdotIntrOp, SmmlaIntrOp, UdotIntrOp, UmmlaIntrOp,
+                    VectorScaleIntrOp>();
+  target.addIllegalOp<SdotOp, SmmlaOp, UdotOp, UmmlaOp, VectorScaleOp>();
   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()) &&
            !hasScalableVectorType(op.getType().getResults());
   });
   target.addDynamicallyLegalOp<CallOp, CallIndirectOp, ReturnOp>(
       [hasScalableVectorType](Operation *op) {
         return !hasScalableVectorType(op->getOperandTypes()) &&
                !hasScalableVectorType(op->getResultTypes());
       });
+  configureBasicSVEArithmeticLegalizations(target);
 }

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

func @arm_sve_ummla(%a: !arm_sve.vector<16xi8>,
                     %c: !arm_sve.vector<4xi32>)
     -> !arm_sve.vector<4xi32> {
   // CHECK: arm_sve.intr.ummla
   %0 = arm_sve.ummla %c, %a, %b :
                !arm_sve.vector<16xi8> to !arm_sve.vector<4xi32>
   return %0 : !arm_sve.vector<4xi32>
 }
 
+func @arm_sve_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>) -> !arm_sve.vector<4xi32> {
+  // CHECK: llvm.mul {{.*}}: !llvm.vec<? x 4 x i32>
+  %0 = arm_sve.muli %a, %b : !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>
+}
+
+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 @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_ummla(%a: !arm_sve.vector<16xi8>,
                     %b: !arm_sve.vector<16xi8>,
                     %c: !arm_sve.vector<4xi32>) -> !arm_sve.vector<4xi32> {
   // CHECK: arm_sve.ummla {{.*}}: !arm_sve.vector<16xi8> to !arm_sve.vector<4xi3
   %0 = arm_sve.ummla %c, %a, %b :
              !arm_sve.vector<16xi8> to !arm_sve.vector<4xi32>
   return %0 : !arm_sve.vector<4xi32>
 }
 
+func @arm_sve_arithi(%a: !arm_sve.vector<4xi32>,
+                     %b: !arm_sve.vector<4xi32>,
+                     %c: !arm_sve.vector<4xi32>) -> !arm_sve.vector<4xi32> {
+  // CHECK: arm_sve.muli {{.*}}: !arm_sve.vector<4xi32>
+  %0 = arm_sve.muli %a, %b : !arm_sve.vector<4xi32>
+  // CHECK: arm_sve.addi {{.*}}: !arm_sve.vector<4xi32>
+  %1 = arm_sve.addi %0, %c : !arm_sve.vector<4xi32>
+  return %1 : !arm_sve.vector<4xi32>
+}
+
+func @arm_sve_arithf(%a: !arm_sve.vector<4xf32>,
+                     %b: !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 @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_ummla(%arg0: !llvm.vec<?x16 x i8>,
                          -> !llvm.vec<?x4 x i32> {
   // CHECK: call <vscale x 4 x i32> @llvm.aarch64.sve.ummla.nxv4i32(<vscale x 4
   %0 = "arm_sve.intr.ummla"(%arg2, %arg0, %arg1) :
     (!llvm.vec<?x4 x i32>, !llvm.vec<?x16 x i8>, !llvm.vec<?x16 x i8>)
         -> !llvm.vec<?x4 x i32>
   llvm.return %0 : !llvm.vec<?x4 x i32>
 }
 
+// CHECK-LABEL: define <vscale x 4 x i32> @arm_sve_arithi
+llvm.func @arm_sve_arithi(%arg0: !llvm.vec<? x 4 x i32>,
+                          %arg1: !llvm.vec<? x 4 x i32>,
+                          %arg2: !llvm.vec<? x 4 x i32>)
+                          -> !llvm.vec<? x 4 x i32> {
+  // CHECK: mul <vscale x 4 x i32>
+  %0 = llvm.mul %arg0, %arg1 : !llvm.vec<? x 4 x i32>
+  // CHECK: add <vscale x 4 x i32>
+  %1 = llvm.add %0, %arg2 : !llvm.vec<? x 4 x i32>
+  llvm.return %1 : !llvm.vec<? x 4 x i32>
+}
+
+// CHECK-LABEL: define <vscale x 4 x float> @arm_sve_arithf
+llvm.func @arm_sve_arithf(%arg0: !llvm.vec<? x 4 x f32>,
+                          %arg1: !llvm.vec<? x 4 x f32>,
+                          %arg2: !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 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
 }