[llvm] [InstCombine][X86] Fold add(vpmadd52(0, a, b), x) -> vpmadd52(x, a, b) (PR #195937)

[via llvm-commits]

================
@@ -0,0 +1,127 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -passes=instcombine -mtriple=x86_64-- -mattr=+avx512ifma -S | FileCheck %s
+
+declare <2 x i64> @llvm.x86.avx512.vpmadd52h.uq.128(<2 x i64>, <2 x i64>, <2 x i64>)
+declare <2 x i64> @llvm.x86.avx512.vpmadd52l.uq.128(<2 x i64>, <2 x i64>, <2 x i64>)
+declare <4 x i64> @llvm.x86.avx512.vpmadd52h.uq.256(<4 x i64>, <4 x i64>, <4 x i64>)
+declare <4 x i64> @llvm.x86.avx512.vpmadd52l.uq.256(<4 x i64>, <4 x i64>, <4 x i64>)
+declare <8 x i64> @llvm.x86.avx512.vpmadd52h.uq.512(<8 x i64>, <8 x i64>, <8 x i64>)
+declare <8 x i64> @llvm.x86.avx512.vpmadd52l.uq.512(<8 x i64>, <8 x i64>, <8 x i64>)
+
+; ---- Positive cases should fold ----
+
+define <2 x i64> @fold_h_128(<2 x i64> %x, <2 x i64> %a, <2 x i64> %b) {
+; CHECK-LABEL: @fold_h_128(
+; CHECK-NEXT:    [[R:%.*]] = call <2 x i64> @llvm.x86.avx512.vpmadd52h.uq.128(<2 x i64> [[X:%.*]], <2 x i64> [[A:%.*]], <2 x i64> [[B:%.*]])
+; CHECK-NEXT:    ret <2 x i64> [[R]]
+;
+  %m = call <2 x i64> @llvm.x86.avx512.vpmadd52h.uq.128(<2 x i64> zeroinitializer, <2 x i64> %a, <2 x i64> %b)
+  %r = add <2 x i64> %m, %x
+  ret <2 x i64> %r
+}
+
+define <2 x i64> @fold_l_128(<2 x i64> %x, <2 x i64> %a, <2 x i64> %b) {
+; CHECK-LABEL: @fold_l_128(
+; CHECK-NEXT:    [[R:%.*]] = call <2 x i64> @llvm.x86.avx512.vpmadd52l.uq.128(<2 x i64> [[X:%.*]], <2 x i64> [[A:%.*]], <2 x i64> [[B:%.*]])
+; CHECK-NEXT:    ret <2 x i64> [[R]]
+;
+  %m = call <2 x i64> @llvm.x86.avx512.vpmadd52l.uq.128(<2 x i64> zeroinitializer, <2 x i64> %a, <2 x i64> %b)
+  %r = add <2 x i64> %m, %x
+  ret <2 x i64> %r
+}
+
+define <4 x i64> @fold_h_256(<4 x i64> %x, <4 x i64> %a, <4 x i64> %b) {
+; CHECK-LABEL: @fold_h_256(
+; CHECK-NEXT:    [[R:%.*]] = call <4 x i64> @llvm.x86.avx512.vpmadd52h.uq.256(<4 x i64> [[X:%.*]], <4 x i64> [[A:%.*]], <4 x i64> [[B:%.*]])
+; CHECK-NEXT:    ret <4 x i64> [[R]]
+;
+  %m = call <4 x i64> @llvm.x86.avx512.vpmadd52h.uq.256(<4 x i64> zeroinitializer, <4 x i64> %a, <4 x i64> %b)
+  %r = add <4 x i64> %m, %x
+  ret <4 x i64> %r
+}
+
+define <4 x i64> @fold_l_256(<4 x i64> %x, <4 x i64> %a, <4 x i64> %b) {
+; CHECK-LABEL: @fold_l_256(
+; CHECK-NEXT:    [[R:%.*]] = call <4 x i64> @llvm.x86.avx512.vpmadd52l.uq.256(<4 x i64> [[X:%.*]], <4 x i64> [[A:%.*]], <4 x i64> [[B:%.*]])
+; CHECK-NEXT:    ret <4 x i64> [[R]]
+;
+  %m = call <4 x i64> @llvm.x86.avx512.vpmadd52l.uq.256(<4 x i64> zeroinitializer, <4 x i64> %a, <4 x i64> %b)
+  %r = add <4 x i64> %m, %x
+  ret <4 x i64> %r
+}
+
+define <8 x i64> @fold_h_512(<8 x i64> %x, <8 x i64> %a, <8 x i64> %b) {
+; CHECK-LABEL: @fold_h_512(
+; CHECK-NEXT:    [[R:%.*]] = call <8 x i64> @llvm.x86.avx512.vpmadd52h.uq.512(<8 x i64> [[X:%.*]], <8 x i64> [[A:%.*]], <8 x i64> [[B:%.*]])
+; CHECK-NEXT:    ret <8 x i64> [[R]]
+;
+  %m = call <8 x i64> @llvm.x86.avx512.vpmadd52h.uq.512(<8 x i64> zeroinitializer, <8 x i64> %a, <8 x i64> %b)
+  %r = add <8 x i64> %m, %x
+  ret <8 x i64> %r
+}
+
+define <8 x i64> @fold_l_512(<8 x i64> %x, <8 x i64> %a, <8 x i64> %b) {
+; CHECK-LABEL: @fold_l_512(
+; CHECK-NEXT:    [[R:%.*]] = call <8 x i64> @llvm.x86.avx512.vpmadd52l.uq.512(<8 x i64> [[X:%.*]], <8 x i64> [[A:%.*]], <8 x i64> [[B:%.*]])
+; CHECK-NEXT:    ret <8 x i64> [[R]]
+;
+  %m = call <8 x i64> @llvm.x86.avx512.vpmadd52l.uq.512(<8 x i64> zeroinitializer, <8 x i64> %a, <8 x i64> %b)
+  %r = add <8 x i64> %m, %x
+  ret <8 x i64> %r
+}
+
+define <4 x i64> @fold_commuted(<4 x i64> %x, <4 x i64> %a, <4 x i64> %b) {
+; CHECK-LABEL: @fold_commuted(
+; CHECK-NEXT:    [[R:%.*]] = call <4 x i64> @llvm.x86.avx512.vpmadd52h.uq.256(<4 x i64> [[X:%.*]], <4 x i64> [[A:%.*]], <4 x i64> [[B:%.*]])
+; CHECK-NEXT:    ret <4 x i64> [[R]]
+;
+  %m = call <4 x i64> @llvm.x86.avx512.vpmadd52h.uq.256(<4 x i64> zeroinitializer, <4 x i64> %a, <4 x i64> %b)
+  %r = add <4 x i64> %x, %m
+  ret <4 x i64> %r
+}
+
+; ---- Negative cases should NOT fold ----
+
+define <4 x i64> @nofold_nonzero_acc(<4 x i64> %x, <4 x i64> %acc, <4 x i64> %a, <4 x i64> %b) {
+; CHECK-LABEL: @nofold_nonzero_acc(
+; CHECK-NEXT:    [[M:%.*]] = call <4 x i64> @llvm.x86.avx512.vpmadd52h.uq.256(<4 x i64> [[ACC:%.*]], <4 x i64> [[A:%.*]], <4 x i64> [[B:%.*]])
+; CHECK-NEXT:    [[R:%.*]] = add <4 x i64> [[M]], [[X:%.*]]
+; CHECK-NEXT:    ret <4 x i64> [[R]]
+;
+  %m = call <4 x i64> @llvm.x86.avx512.vpmadd52h.uq.256(<4 x i64> %acc, <4 x i64> %a, <4 x i64> %b)
+  %r = add <4 x i64> %m, %x
+  ret <4 x i64> %r
+}
----------------
as4230 wrote:

I just pushed another commit with a test that hits the fold with that. The test I added before reused the same madd52 call to make `add <4 x i64> %m, %m` which avoids the check you want because a previous fold does `add x, x -> shl x, 1` first. If I make two equivalent madd52 calls to create `add <4 x i64> %m1, %m2` then the fold fires and exposes a issue. The new madd52 was being inserted at II's location but the X (the second madd52) is defined after it. I set it before the Add instead so that doesn't happen. I made a test case with 2 and 3 calls for this.


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