[AArch64]Fix an assertion failure about concating two build_vector in DAG Combiner

[• HaoLiu]

Summary

Hi Tim and other reviewers,

In DAG Combiner function visitCONCAT_VECTORS, it tries to do following optimization:

fold (concat_vectors (BUILD_VECTOR A, B, ...), (BUILD_VECTOR C, D, ...))
    -> (BUILD_VECTOR A, B, ..., C, D, ...)

But this optimization doesn't check whether A/B and C/D are different types. The test case in this patch shows that two SHL may be optimized into a v4i16 BUILD_VECTOR from four i32 constant 0.
There are two ways to fix this bug:
(1) Check whether two BUILD_VECTOR are from different types, if not, truncate the larger type.
(2) When combine 2 SHL, try to generate v4i16 BUILD_VECTOR from v4i16 constant 0. But we can't make sure no other optimization may generate such v4i16 from i32. We choose the first way.

Review please.

Thanks
-Hao

[t.p.northover]

Hi Hao,

I think you picked the right place to fix this. Even if you did patch AArch64, the bug would still be latent in the DAGCombiner.

I've just got one comment on the logic, really:

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

10377–10378

I think the logic here would be simpler if you defined an "EVT MinTy" and then truncated all inputs down to that in the loop. The trivial truncate operations will be automatically discarded in SelectionDAG::getNode.

[• HaoLiu]

Hi Tim,

Your suggestion is good to simplify the logic.
I've refactored the patch.

Thanks,
-Hao

[t.p.northover]

Hi Hao,

I think this looks good now. Thanks for updating the patch!

Tim.

[• HaoLiu]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

     SmallVector<SDValue, 8> Opnds;
     unsigned BuildVecNumElts =  N0.getNumOperands();
 
-    for (unsigned i = 0; i != BuildVecNumElts; ++i)
-      Opnds.push_back(N0.getOperand(i));
-    for (unsigned i = 0; i != BuildVecNumElts; ++i)
-      Opnds.push_back(N1.getOperand(i));
+    EVT SclTy0 = N0.getOperand(0)->getValueType(0);
+    EVT SclTy1 = N1.getOperand(0)->getValueType(0);
+    if (SclTy0.isFloatingPoint()) {
+      for (unsigned i = 0; i != BuildVecNumElts; ++i)
+        Opnds.push_back(N0.getOperand(i));
+      for (unsigned i = 0; i != BuildVecNumElts; ++i)
+        Opnds.push_back(N1.getOperand(i));
+    } else {
+      // If BUILD_VECTOR are from built from integer, they may have different
+      // operand types. Get the smaller type and truncate all operands to it.
+      EVT MinTy = SclTy0.bitsLE(SclTy1) ? SclTy0 : SclTy1;
+      for (unsigned i = 0; i != BuildVecNumElts; ++i)
+        Opnds.push_back(DAG.getNode(ISD::TRUNCATE, SDLoc(N), MinTy,
+                        N0.getOperand(i)));
+      for (unsigned i = 0; i != BuildVecNumElts; ++i)
+        Opnds.push_back(DAG.getNode(ISD::TRUNCATE, SDLoc(N), MinTy,
+                        N1.getOperand(i)));
+    }
 
     return DAG.getNode(ISD::BUILD_VECTOR, SDLoc(N), VT, Opnds);
   }

test/CodeGen/AArch64/arm64-build-vector.ll

   %b = add <8 x i16> %a, <i16 -32768, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef>
   %c = mul <8 x i16> %b, <i16 -20864, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef, i16 undef>
   ret <8 x i16> %c
-  }
+
+; There is an optimization in DAG Combiner as following:
+;   fold (concat_vectors (BUILD_VECTOR A, B, ...), (BUILD_VECTOR C, D, ...))
+;        -> (BUILD_VECTOR A, B, ..., C, D, ...)
+; This case checks when A,B and C,D are different types, there should be no
+; assertion failure.
+define <8 x i16> @concat_2_build_vector(<4 x i16> %in0) {
+; CHECK-LABEL: concat_2_build_vector:
+; CHECK: movi
+  %vshl_n = shl <4 x i16> %in0, <i16 8, i16 8, i16 8, i16 8>
+  %vshl_n2 = shl <4 x i16> %vshl_n, <i16 9, i16 9, i16 9, i16 9>
+  %shuffle.i = shufflevector <4 x i16> %vshl_n2, <4 x i16> zeroinitializer, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
+  ret <8 x i16> %shuffle.i
+}
 No newline at end of file