[llvm-branch-commits] [llvm] [RISCV] Support i64 accumulator for Zvdot4a8i partial reductions (PR #215517)

[Philip Reames via llvm-branch-commits]

================
@@ -0,0 +1,899 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=riscv32 -mattr=+v -verify-machineinstrs < %s | FileCheck %s --check-prefixes=CHECK,NODOT,NODOT32
+; RUN: llc -mtriple=riscv64 -mattr=+v -verify-machineinstrs < %s | FileCheck %s --check-prefixes=CHECK,NODOT,NODOT64
+; RUN: llc -mtriple=riscv32 -mattr=+v,+experimental-zvdot4a8i -verify-machineinstrs < %s | FileCheck %s --check-prefixes=CHECK,DOT,DOT32
+; RUN: llc -mtriple=riscv64 -mattr=+v,+experimental-zvdot4a8i -verify-machineinstrs < %s | FileCheck %s --check-prefixes=CHECK,DOT,DOT64
+
+; Tests for partial.reduce.add into an i64 accumulator with i8 inputs. When
+; Zvdot4a8i is available these are lowered to a vdot4a* (i8 -> i32) followed by
+; a widening (i32 -> i64) accumulate, mirroring the AArch64 sdot+sadalp idiom.
+
+define <vscale x 1 x i64> @vdot4a_i64(<vscale x 1 x i64> %acc, <vscale x 8 x i8> %a, <vscale x 8 x i8> %b) {
+; NODOT32-LABEL: vdot4a_i64:
+; NODOT32:       # %bb.0: # %entry
+; NODOT32-NEXT:    li a0, 1
+; NODOT32-NEXT:    vsetvli a1, zero, e32, m4, ta, ma
+; NODOT32-NEXT:    vsext.vf4 v12, v9
+; NODOT32-NEXT:    vsext.vf4 v24, v10
+; NODOT32-NEXT:    vwmul.vv v16, v12, v24
+; NODOT32-NEXT:    vsetvli zero, zero, e64, m8, ta, ma
+; NODOT32-NEXT:    vmul.vx v16, v16, a0
+; NODOT32-NEXT:    vsetvli a0, zero, e64, m1, ta, ma
+; NODOT32-NEXT:    vadd.vv v9, v19, v20
+; NODOT32-NEXT:    vadd.vv v8, v8, v16
+; NODOT32-NEXT:    vadd.vv v10, v17, v18
+; NODOT32-NEXT:    vadd.vv v8, v23, v8
+; NODOT32-NEXT:    vadd.vv v11, v21, v22
+; NODOT32-NEXT:    vadd.vv v9, v10, v9
+; NODOT32-NEXT:    vadd.vv v8, v11, v8
+; NODOT32-NEXT:    vadd.vv v8, v9, v8
+; NODOT32-NEXT:    ret
+;
+; NODOT64-LABEL: vdot4a_i64:
+; NODOT64:       # %bb.0: # %entry
+; NODOT64-NEXT:    vsetvli a0, zero, e32, m4, ta, ma
+; NODOT64-NEXT:    vsext.vf4 v12, v9
+; NODOT64-NEXT:    vsext.vf4 v24, v10
+; NODOT64-NEXT:    vwmul.vv v16, v12, v24
+; NODOT64-NEXT:    vsetvli a0, zero, e64, m1, ta, ma
+; NODOT64-NEXT:    vadd.vv v9, v19, v20
+; NODOT64-NEXT:    vadd.vv v8, v8, v16
+; NODOT64-NEXT:    vadd.vv v10, v17, v18
+; NODOT64-NEXT:    vadd.vv v8, v23, v8
+; NODOT64-NEXT:    vadd.vv v11, v21, v22
+; NODOT64-NEXT:    vadd.vv v9, v10, v9
+; NODOT64-NEXT:    vadd.vv v8, v11, v8
+; NODOT64-NEXT:    vadd.vv v8, v9, v8
+; NODOT64-NEXT:    ret
+;
+; DOT-LABEL: vdot4a_i64:
+; DOT:       # %bb.0: # %entry
+; DOT-NEXT:    vsetvli a0, zero, e32, m1, ta, ma
+; DOT-NEXT:    vmv.v.i v11, 0
+; DOT-NEXT:    vdot4a.vv v11, v9, v10
+; DOT-NEXT:    csrr a0, vlenb
+; DOT-NEXT:    srli a0, a0, 3
+; DOT-NEXT:    vslidedown.vx v9, v11, a0
----------------
preames wrote:

Hm, this is exactly what I asked for, but I hadn't considered the implications at m1 and the need for the high mf2 extract.  That's a bit ugly.  However, the m2 and above is a definite improvement since the high half extract folds into the reduce change.

Possibly we want an alternate lowering for m1 vs m2 and above?  This is fine to explore in a follow up patch if you want.  One point worth considering is that vwredsum(u).vs is technically a partial reduction if the other lanes are zeroed.  You'd need a wide add for the accumulator.  



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