[VectorCombine] Enable scalarizeBinopOrCmp for scalable vectors

[MattDevereau]

Summary

[VectorCombine] Enable scalarizeBinopOrCmp for scalable vectors

This reverts a change to exclude scalarizeBinopOrCmp in VectorCombine for
scalable vectors which caused poor scalable Binop codegen.

[spatel]

LGTM - see inline for a couple of additions.

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

1724

Add a comment here to make the logic difference explicit:

// This transform works with scalable and fixed vectors.

Should there be a TODO comment about allowing more scalable transforms?

llvm/test/Transforms/VectorCombine/AArch64/scalarize-scalable.ll

5

This shows a series of folds, but it would be good to have a couple of minimum tests too. Something like this should give us coverage for integer types and compares?

define <vscale x 4 x i32> @scalarize_scalable_udiv(i32 %x, i32 %y) {
  %splatx = insertelement <vscale x 4 x i32> poison, i32 %x, i64 0
  %splaty = insertelement <vscale x 4 x i32> poison, i32 %y, i64 0
  %r = udiv <vscale x 4 x i32> %splatx, %splaty
  ret <vscale x 4 x i32> %r
}

define <vscale x 4 x i1> @scalarize_scalable_icmp(i32 %x, i32 %y) {
  %splatx = insertelement <vscale x 4 x i32> poison, i32 %x, i64 0
  %splaty = insertelement <vscale x 4 x i32> poison, i32 %y, i64 0
  %r = icmp sgt <vscale x 4 x i32> %splatx, %splaty
  ret <vscale x 4 x i1> %r
}

[MattDevereau]

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

1724

I'll add the comment on push.
Regarding the TODO: I think it's ok to put in to signal that there's more scalable optimization work possible here, but to be explicit on why this patch exists, we currently just want to fix this single regression case we've identified and are not actively seeking more scalable optimizations in this area just yet.

llvm/test/Transforms/VectorCombine/AArch64/scalarize-scalable.ll

5

Sure, I'm perfectly happy to throw in some more tests

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

if (IsFixedVectorType) {
         break;
       case Instruction::ShuffleVector:
         MadeChange |= widenSubvectorLoad(I);
         break;
       case Instruction::Load:
         MadeChange |= scalarizeLoadExtract(I);
         break;
       default:
-        MadeChange |= scalarizeBinopOrCmp(I);
         break;
       }
     }
+
+    // This transform works with scalable and fixed vectors
+    // TODO: Identify and allow other scalable transforms
+    if (isa<VectorType>(I.getType()))
+      MadeChange |= scalarizeBinopOrCmp(I);
+
     if (Opcode == Instruction::Store)
       MadeChange |= foldSingleElementStore(I);
 
 
     // If this is an early pipeline invocation of this pass, we are done.
     if (TryEarlyFoldsOnly)
       return;
 

llvm/test/Transforms/VectorCombine/AArch64/scalarize-scalable.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes=vector-combine -S %s | FileCheck %s
+
+define <vscale x 4 x float> @scalarize_scalable(float %0, float %1, float %2, float %3) {
+; CHECK-LABEL: @scalarize_scalable(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[DOTSCALAR:%.*]] = fadd fast float [[TMP1:%.*]], [[TMP0:%.*]]
+; CHECK-NEXT:    [[DOTSCALAR1:%.*]] = fadd fast float [[TMP2:%.*]], [[DOTSCALAR]]
+; CHECK-NEXT:    [[DOTSCALAR2:%.*]] = fadd fast float [[TMP3:%.*]], [[DOTSCALAR1]]
+; CHECK-NEXT:    [[TMP4:%.*]] = insertelement <vscale x 4 x float> poison, float [[DOTSCALAR2]], i64 0
+; CHECK-NEXT:    [[TMP5:%.*]] = shufflevector <vscale x 4 x float> [[TMP4]], <vscale x 4 x float> poison, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT:    ret <vscale x 4 x float> [[TMP5]]
+;
+entry:
+  %broadcast.splatinsert = insertelement <vscale x 4 x float> poison, float %0, i64 0
+  %broadcast.splatinsert28 = insertelement <vscale x 4 x float> poison, float %1, i64 0
+  %broadcast.splatinsert30 = insertelement <vscale x 4 x float> poison, float %2, i64 0
+  %broadcast.splatinsert32 = insertelement <vscale x 4 x float> poison, float %3, i64 0
+  %4 = fadd fast <vscale x 4 x float> %broadcast.splatinsert28, %broadcast.splatinsert
+  %5 = fadd fast <vscale x 4 x float> %broadcast.splatinsert30, %4
+  %6 = fadd fast <vscale x 4 x float> %broadcast.splatinsert32, %5
+  %7 = shufflevector <vscale x 4 x float> %6, <vscale x 4 x float> poison, <vscale x 4 x i32> zeroinitializer
+  ret <vscale x 4 x float> %7
+}
+
+define <vscale x 4 x i32> @scalarize_scalable_udiv(i32 %x, i32 %y) {
+; CHECK-LABEL: @scalarize_scalable_udiv(
+; CHECK-NEXT:    [[R_SCALAR:%.*]] = udiv i32 [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = insertelement <vscale x 4 x i32> poison, i32 [[R_SCALAR]], i64 0
+; CHECK-NEXT:    ret <vscale x 4 x i32> [[R]]
+;
+  %splatx = insertelement <vscale x 4 x i32> poison, i32 %x, i64 0
+  %splaty = insertelement <vscale x 4 x i32> poison, i32 %y, i64 0
+  %r = udiv <vscale x 4 x i32> %splatx, %splaty
+  ret <vscale x 4 x i32> %r
+}
+
+define <vscale x 4 x i1> @scalarize_scalable_icmp(i32 %x, i32 %y) {
+; CHECK-LABEL: @scalarize_scalable_icmp(
+; CHECK-NEXT:    [[R_SCALAR:%.*]] = icmp sgt i32 [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = insertelement <vscale x 4 x i1> poison, i1 [[R_SCALAR]], i64 0
+; CHECK-NEXT:    ret <vscale x 4 x i1> [[R]]
+;
+  %splatx = insertelement <vscale x 4 x i32> poison, i32 %x, i64 0
+  %splaty = insertelement <vscale x 4 x i32> poison, i32 %y, i64 0
+  %r = icmp sgt <vscale x 4 x i32> %splatx, %splaty
+  ret <vscale x 4 x i1> %r
+}