[llvm] [LLVM][CodeGen][SVE] Add isel patterns for abs-diff partial reductions. (PR #212800)

[Benjamin Maxwell via llvm-commits]

================
@@ -82,23 +123,63 @@ define <16 x i8> @uabs_v16i8(<16 x i8> %a, <16 x i8> %b) {
   ret <16 x i8> %uabs
 }
 
-; TODO: This case could be lowered to a uabal[bt] pair.
-define <vscale x 4 x i32> @uabs_nxv16i8_wide_add(<vscale x 4 x i32> %acc, <vscale x 8 x i16> %a, <vscale x 8 x i16> %b) {
-; CHECK-LABEL: uabs_nxv16i8_wide_add:
-; CHECK:       // %bb.0:
-; CHECK-NEXT:    ptrue p0.h
-; CHECK-NEXT:    uabd z1.h, p0/m, z1.h, z2.h
-; CHECK-NEXT:    uaddwb z0.s, z0.s, z1.h
-; CHECK-NEXT:    uaddwt z0.s, z0.s, z1.h
-; CHECK-NEXT:    ret
-  %umax = tail call <vscale x 8 x i16> @llvm.umax.nxv8i16(<vscale x 8 x i16> %a, <vscale x 8 x i16> %b)
-  %umin = tail call <vscale x 8 x i16> @llvm.umin.nxv8i16(<vscale x 8 x i16> %a, <vscale x 8 x i16> %b)
+define <vscale x 8 x i16> @uabs_nxv16i8_wide_add(<vscale x 8 x i16> %acc, <vscale x 16 x i8> %a, <vscale x 16 x i8> %b) {
+; SVE2-LABEL: uabs_nxv16i8_wide_add:
+; SVE2:       // %bb.0:
+; SVE2-NEXT:    uabalb z0.h, z1.b, z2.b
+; SVE2-NEXT:    uabalt z0.h, z1.b, z2.b
+; SVE2-NEXT:    ret
+;
+; SVE2p3-LABEL: uabs_nxv16i8_wide_add:
+; SVE2p3:       // %bb.0:
+; SVE2p3-NEXT:    uabal z0.h, z1.b, z2.b
+; SVE2p3-NEXT:    ret
+  %umax = tail call <vscale x 16 x i8> @llvm.umax(<vscale x 16 x i8> %a, <vscale x 16 x i8> %b)
+  %umin = tail call <vscale x 16 x i8> @llvm.umin(<vscale x 16 x i8> %a, <vscale x 16 x i8> %b)
+  %uabs = sub <vscale x 16 x i8> %umax, %umin
+  %ext = zext <vscale x 16 x i8> %uabs to <vscale x 16 x i16>
+  %reduce = call <vscale x 8 x i16> @llvm.vector.partial.reduce.add(<vscale x 8 x i16> %acc, <vscale x 16 x i16> %ext)
+  ret <vscale x 8 x i16> %reduce
+}
+
+define <vscale x 4 x i32> @uabs_nxv8i16_wide_add(<vscale x 4 x i32> %acc, <vscale x 8 x i16> %a, <vscale x 8 x i16> %b) {
+; SVE2-LABEL: uabs_nxv8i16_wide_add:
+; SVE2:       // %bb.0:
+; SVE2-NEXT:    uabalb z0.s, z1.h, z2.h
+; SVE2-NEXT:    uabalt z0.s, z1.h, z2.h
+; SVE2-NEXT:    ret
+;
+; SVE2p3-LABEL: uabs_nxv8i16_wide_add:
+; SVE2p3:       // %bb.0:
+; SVE2p3-NEXT:    uabal z0.s, z1.h, z2.h
+; SVE2p3-NEXT:    ret
+  %umax = tail call <vscale x 8 x i16> @llvm.umax(<vscale x 8 x i16> %a, <vscale x 8 x i16> %b)
+  %umin = tail call <vscale x 8 x i16> @llvm.umin(<vscale x 8 x i16> %a, <vscale x 8 x i16> %b)
   %uabs = sub <vscale x 8 x i16> %umax, %umin
   %ext = zext <vscale x 8 x i16> %uabs to <vscale x 8 x i32>
-  %reduce = call <vscale x 4 x i32> @llvm.vector.partial.reduce.add.v4i32.v16i32(<vscale x 4 x i32> %acc, <vscale x 8 x i32> %ext)
+  %reduce = call <vscale x 4 x i32> @llvm.vector.partial.reduce.add(<vscale x 4 x i32> %acc, <vscale x 8 x i32> %ext)
   ret <vscale x 4 x i32> %reduce
 }
 
+define <vscale x 2 x i64> @uabs_nxv4i32_wide_add(<vscale x 2 x i64> %acc, <vscale x 4 x i32> %a, <vscale x 4 x i32> %b) {
+; SVE2-LABEL: uabs_nxv4i32_wide_add:
+; SVE2:       // %bb.0:
+; SVE2-NEXT:    uabalb z0.d, z1.s, z2.s
+; SVE2-NEXT:    uabalt z0.d, z1.s, z2.s
+; SVE2-NEXT:    ret
+;
+; SVE2p3-LABEL: uabs_nxv4i32_wide_add:
+; SVE2p3:       // %bb.0:
+; SVE2p3-NEXT:    uabal z0.d, z1.s, z2.s
+; SVE2p3-NEXT:    ret
+  %umax = tail call <vscale x 4 x i32> @llvm.umax(<vscale x 4 x i32> %a, <vscale x 4 x i32> %b)
+  %umin = tail call <vscale x 4 x i32> @llvm.umin(<vscale x 4 x i32> %a, <vscale x 4 x i32> %b)
+  %uabs = sub <vscale x 4 x i32> %umax, %umin
+  %ext = zext <vscale x 4 x i32> %uabs to <vscale x 4 x i64>
+  %reduce = call <vscale x 2 x i64> @llvm.vector.partial.reduce.add(<vscale x 2 x i64> %acc, <vscale x 4 x i64> %ext)
+  ret <vscale x 2 x i64> %reduce
+}
+
----------------
MacDue wrote:

> To further this point, would we then need to repeatedly document what _mt and _m1 mean?

I would not think so, as for these predicates the predication needs to be handled in the pattern (simply ignoring likely would be incorrect), so it's less surprising.

I think there's ways to make this pattern more obvious, for example we could have a `(SVEUndefPredication)` complex pattern, that could 'any predicate', possibly with an possibly with an additional assert/check that it's only used with `_PRED` nodes.

There's also many uses of `_p` nodes with 'AllActive' predicate matches, maybe that's always intended. But it could be that 'undef' meaning may have been missed in cases (possible as the name alone does not explain that behavior unlike the other nodes, `_PRED` vs something like `_PRED_UNDEF`). 

Anyway, I digress. Yes there's good docs, but I think names are important too.  


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