Add pattern to rewrite chains of MLA instructions into a tree for increased ILP

Author: jcohen-apple

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

@@ -6709,6 +6709,7 @@ bool AArch64InstrInfo::isAccumulationOpcode(unsigned Opcode) const {
   case AArch64::SABAv4i32:
   case AArch64::SABAv8i16:
   case AArch64::SABAv8i8:
+  case AArch64::MLAv8i8:
     return true;
   }
 
@@ -6720,6 +6721,8 @@ std::optional<unsigned> AArch64InstrInfo::getAccumulationStartOpcode(
   switch (AccumulationOpcode) {
   default:
     llvm_unreachable("Unknown accumulator opcode");
+  case AArch64::MLAv8i8:
+    return AArch64::MULv8i8;
   case AArch64::UABALB_ZZZ_D:
     return AArch64::UABDLB_ZZZ_D;
   case AArch64::UABALB_ZZZ_H:
@@ -7593,6 +7596,7 @@ std::optional<unsigned> AArch64InstrInfo::getReduceOpcodeForAccumulator(
     return AArch64::ADDv2i32;
   case AArch64::UABAv8i8:
   case AArch64::SABAv8i8:
+  case AArch64::MLAv8i8:
     return AArch64::ADDv8i8;
   default:
     llvm_unreachable("Unknown accumulator opcode");

llvm/test/CodeGen/AArch64/aarch64-reassociate-accumulators.ll

@@ -1,5 +1,36 @@
-; RUN: opt -passes=loop-unroll %s -o - | llc -O3 - -mtriple=arm64e-apple-darwin -o - | Filecheck %s
+; RUN: opt -passes=loop-unroll %s -o - | llc -O3 - -mtriple=arm64e-apple-darwin -machine-combiner-recurse -o - | Filecheck %s
 
+define i8 @mla_i8_accumulation(ptr %ptr1, ptr %ptr2) {
+entry:
+  br label %loop
+loop:
+  %i = phi i32 [ 0, %entry ], [ %next_i, %loop ]
+  %acc_phi = phi <8 x i8> [ zeroinitializer, %entry ], [ %acc_next, %loop ]
+  %ptr1_i = getelementptr i8, ptr %ptr1, i32 %i
+  %ptr2_i = getelementptr i8, ptr %ptr2, i32 %i
+  %a = load <8 x i8>, <8 x i8>* %ptr1_i, align 1
+  %b = load <8 x i8>, <8 x i8>* %ptr2_i, align 1
+  %mul = mul <8 x i8> %a, %b
+  %acc_next = add <8 x i8> %acc_phi, %mul
+  %next_i = add i32 %i, 8
+  %cmp = icmp slt i32 %next_i, 64
+  br i1 %cmp, label %loop, label %exit
+exit:
+  %reduce = call i8 @llvm.vector.reduce.add.v8i8(<8 x i8> %acc_next)
+  ret i8 %reduce
+}
+; CHECK-LABEL: mla_i8_accumulation
+; CHECK: mul.8b v1
+; CHECK: mul.8b v0
+; CHECK: mul.8b v2
+; CHECK: mla.8b v1
+; CHECK: mla.8b v0
+; CHECK: mla.8b v2
+; CHECK: mla.8b v1
+; CHECK: mla.8b v0
+; CHECK: add.8b v1, v2, v1
+; CHECK: add.8b v0, v1, v0
+; CHECK: addv.8b
 
 define i16 @sabal_i8_to_i16_accumulation(ptr %ptr1, ptr %ptr2) {
 entry: