[llvm] [AArch64][SVE] Use ADD/ADR instead of MUL/MLA for x*N (PR #198566)

[Paul Walker via llvm-commits]

================
@@ -1948,6 +1948,75 @@ let Predicates = [HasSVE] in {
   defm : adrShiftPat<nxv4i32, nxv4i1, i32, ADR_LSL_ZZZ_S_2, 2>;
   defm : adrShiftPat<nxv4i32, nxv4i1, i32, ADR_LSL_ZZZ_S_3, 3>;
 
+  // Avoid MUL/MLA for small constants where ADD/ADR forms are available.
+  multiclass sveMulImmAddAdrPat<ValueType Ty, ValueType PredTy,
+                                ValueType ScalarTy, Instruction Add,
+                                Instruction Adr1, Instruction Adr2,
+                                Instruction Adr3> {
+    def : Pat<(Ty (AArch64mul_p (PredTy (SVEAnyPredicate)), Ty:$Op,
+                                (Ty (splat_vector (ScalarTy 2))))),
+              (Add $Op, $Op)>;
+    def : Pat<(Ty (AArch64mul_p (PredTy (SVEAnyPredicate)), Ty:$Op,
+                                (Ty (splat_vector (ScalarTy 3))))),
+              (Adr1 $Op, $Op)>;
+    def : Pat<(Ty (AArch64mul_p (PredTy (SVEAnyPredicate)), Ty:$Op,
+                                (Ty (splat_vector (ScalarTy 5))))),
+              (Adr2 $Op, $Op)>;
+    def : Pat<(Ty (AArch64mul_p (PredTy (SVEAnyPredicate)), Ty:$Op,
+                                (Ty (splat_vector (ScalarTy 9))))),
+              (Adr3 $Op, $Op)>;
+
+    def : Pat<(Ty (AArch64mul_p (PredTy (SVEAnyPredicate)),
+                                (Ty (splat_vector (ScalarTy 2))), Ty:$Op)),
+              (Add $Op, $Op)>;
+    def : Pat<(Ty (AArch64mul_p (PredTy (SVEAnyPredicate)),
+                                (Ty (splat_vector (ScalarTy 3))), Ty:$Op)),
+              (Adr1 $Op, $Op)>;
+    def : Pat<(Ty (AArch64mul_p (PredTy (SVEAnyPredicate)),
+                                (Ty (splat_vector (ScalarTy 5))), Ty:$Op)),
+              (Adr2 $Op, $Op)>;
+    def : Pat<(Ty (AArch64mul_p (PredTy (SVEAnyPredicate)),
+                                (Ty (splat_vector (ScalarTy 9))), Ty:$Op)),
+              (Adr3 $Op, $Op)>;
+  }
+
+  multiclass sveMlaImmAdrPat<ValueType Ty, ValueType PredTy,
+                             ValueType ScalarTy, Instruction Adr1,
+                             Instruction Adr2, Instruction Adr3> {
+    def : Pat<(Ty (AArch64mla_p (PredTy (SVEAnyPredicate)), Ty:$Acc, Ty:$Op,
+                                (Ty (splat_vector (ScalarTy 2))))),
+              (Adr1 $Acc, $Op)>;
+    def : Pat<(Ty (AArch64mla_p (PredTy (SVEAnyPredicate)), Ty:$Acc, Ty:$Op,
+                                (Ty (splat_vector (ScalarTy 4))))),
+              (Adr2 $Acc, $Op)>;
+    def : Pat<(Ty (AArch64mla_p (PredTy (SVEAnyPredicate)), Ty:$Acc, Ty:$Op,
+                                (Ty (splat_vector (ScalarTy 8))))),
+              (Adr3 $Acc, $Op)>;
+
+    def : Pat<(Ty (AArch64mla_p (PredTy (SVEAnyPredicate)), Ty:$Acc,
+                                (Ty (splat_vector (ScalarTy 2))), Ty:$Op)),
+              (Adr1 $Acc, $Op)>;
+    def : Pat<(Ty (AArch64mla_p (PredTy (SVEAnyPredicate)), Ty:$Acc,
+                                (Ty (splat_vector (ScalarTy 4))), Ty:$Op)),
+              (Adr2 $Acc, $Op)>;
+    def : Pat<(Ty (AArch64mla_p (PredTy (SVEAnyPredicate)), Ty:$Acc,
+                                (Ty (splat_vector (ScalarTy 8))), Ty:$Op)),
+              (Adr3 $Acc, $Op)>;
+  }
+
+  let AddedComplexity = 10 in {
+    defm : sveMulImmAddAdrPat<nxv2i64, nxv2i1, i64, ADD_ZZZ_D,
+                              ADR_LSL_ZZZ_D_1, ADR_LSL_ZZZ_D_2,
+                              ADR_LSL_ZZZ_D_3>;
+    defm : sveMulImmAddAdrPat<nxv4i32, nxv4i1, i32, ADD_ZZZ_S,
+                              ADR_LSL_ZZZ_S_1, ADR_LSL_ZZZ_S_2,
+                              ADR_LSL_ZZZ_S_3>;
+    defm : sveMlaImmAdrPat<nxv2i64, nxv2i1, i64, ADR_LSL_ZZZ_D_1,
+                           ADR_LSL_ZZZ_D_2, ADR_LSL_ZZZ_D_3>;
+    defm : sveMlaImmAdrPat<nxv4i32, nxv4i1, i32, ADR_LSL_ZZZ_S_1,
+                           ADR_LSL_ZZZ_S_2, ADR_LSL_ZZZ_S_3>;
----------------
paulwalker-arm wrote:

Given they take the same parameter is there any reason not to merge sveMulImmAddAdrPat and sveMlaImmAdrPat into a single class? You might also be able to merge in the logic from `adrShiftPat` so there's just a single class covering all the ADR use cases.

https://github.com/llvm/llvm-project/pull/198566