[LegalizeDAG] Truncate condition operand of ISD::SELECT

[evgeny777]

Summary

I'm facing a crash when compiling large in-house project for AArch64 using clang, which happens because condition argument of ISD::SELECT can't be expanded (for my case required expansion is i128 -> i64). This patch adds required support.

[efriedma]

It seems a little dubious to create a SELECT where the condition isn't either an i1 or a naturally promoted i1... maybe it would make sense to restrict that instead? (And fix the VSELECT expansion to just truncate the condition instead of trying to get fancy.)

lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

3118 ↗

(On Diff #132172)

Probably want to clarify that you're intentionally throwing away the high bits because the operand is required to conform to getBooleanContents.

[evgeny777]

Addressed review comments. The new version also handles scalar types which size is not power of 2.

[evgeny777]

Any comments on this? @efriedma ?

[delena]

lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

368 ↗

(On Diff #132335)

Can the condition be wider than getSetCCResultType?

test/CodeGen/AArch64/expand-select.ll

4 ↗

(On Diff #132335)

I think you should put here regular CHECKs, like in all other llc tests.

[evgeny777]

lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

368 ↗

(On Diff #132335)

Yes, if this wider type is legal for target architecture. However see @efriedma comment above.

If boolean type is not legal then you either crash when trying to expand condition or when trying to promote it (if type bit width is not power of 2). Here I chose to truncate it to size of getSetCCResultType to avoid promoting it back in PromoteTargetBoolean.

test/CodeGen/AArch64/expand-select.ll

4 ↗

(On Diff #132335)

Any suggestion on what should be checked? Probably debug output from -debug-only=selectiondag,isel?

[delena]

lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

368 ↗

(On Diff #132335)

My question was incorrect. I wanted to ask whether the getSetCCResultType may be wider than Cond. Do you need to extend the Cond in this case?

test/CodeGen/AArch64/expand-select.ll

4 ↗

(On Diff #132335)

just use utils/update_llc_test_checks.py to generate "CHECKS":

utils/update_llc_test_checks.py --llc-binary=<path>/llc test/CodeGen/AArch64/expand-select.ll

[evgeny777]

lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

368 ↗

(On Diff #132335)

I don't think so. PromoteIntOp_SELECT will do this for me. It calls PromoteTargetBoolean which in turn zero extends condition to the size of getSetCCResultType

test/CodeGen/AArch64/expand-select.ll

4 ↗

(On Diff #132335)

Thanks, I'll give it a try.

[evgeny777]

Added checks to test case

llvm/trunk/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

if (TLI.getBooleanContents(false, false) !=
     if (Cond->getOpcode() == ISD::SETCC) {
       EVT OpVT = Cond->getOperand(0).getValueType();
       ScalarBool = TLI.getBooleanContents(OpVT.getScalarType());
       VecBool = TLI.getBooleanContents(OpVT);
     } else
       ScalarBool = TargetLowering::UndefinedBooleanContent;
   }
 
-  if (ScalarBool != VecBool) {
-    EVT CondVT = Cond.getValueType();
+  if (ScalarBool != VecBool) {
     switch (ScalarBool) {
       case TargetLowering::UndefinedBooleanContent:
         break;
       case TargetLowering::ZeroOrOneBooleanContent:
         assert(VecBool == TargetLowering::UndefinedBooleanContent ||
                VecBool == TargetLowering::ZeroOrNegativeOneBooleanContent);
         // Vector read from all ones, scalar expects a single 1 so mask.
         Cond = DAG.getNode(ISD::AND, SDLoc(N), CondVT,
                            Cond, DAG.getConstant(1, SDLoc(N), CondVT));
         break;
       case TargetLowering::ZeroOrNegativeOneBooleanContent:
         assert(VecBool == TargetLowering::UndefinedBooleanContent ||
                VecBool == TargetLowering::ZeroOrOneBooleanContent);
         // Vector reads from a one, scalar from all ones so sign extend.
         Cond = DAG.getNode(ISD::SIGN_EXTEND_INREG, SDLoc(N), CondVT,
                            Cond, DAG.getValueType(MVT::i1));
         break;
     }
   }
 
+  // Truncate the condition if needed
+  auto BoolVT = getSetCCResultType(CondVT);
+  if (BoolVT.bitsLT(CondVT))
+    Cond = DAG.getNode(ISD::TRUNCATE, SDLoc(N), BoolVT, Cond);
+
   return DAG.getSelect(SDLoc(N),
                        LHS.getValueType(), Cond, LHS,
                        GetScalarizedVector(N->getOperand(2)));
 }
 
 SDValue DAGTypeLegalizer::ScalarizeVecRes_SELECT(SDNode *N) {
   SDValue LHS = GetScalarizedVector(N->getOperand(1));
   return DAG.getSelect(SDLoc(N),

llvm/trunk/test/CodeGen/AArch64/expand-select.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64-unknown-linux-gnu -O3 %s -o - | FileCheck %s
+
+define void @foo(i32 %In1, <2 x i128> %In2, <2 x i128> %In3, <2 x i128> *%Out) {
+; CHECK-LABEL: foo:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    and w9, w0, #0x1
+; CHECK-NEXT:    fmov s0, wzr
+; CHECK-NEXT:    ldp x10, x8, [sp, #8]
+; CHECK-NEXT:    fmov s1, w9
+; CHECK-NEXT:    ldr x9, [sp]
+; CHECK-NEXT:    cmeq v0.4s, v1.4s, v0.4s
+; CHECK-NEXT:    fmov w11, s0
+; CHECK-NEXT:    tst w11, #0x1
+; CHECK-NEXT:    csel x11, x2, x6, ne
+; CHECK-NEXT:    csel x12, x3, x7, ne
+; CHECK-NEXT:    csel x9, x4, x9, ne
+; CHECK-NEXT:    csel x10, x5, x10, ne
+; CHECK-NEXT:    stp x9, x10, [x8, #16]
+; CHECK-NEXT:    stp x11, x12, [x8]
+; CHECK-NEXT:    ret
+  %cond = and i32 %In1, 1
+  %cbool = icmp eq i32 %cond, 0
+  %res = select i1 %cbool, <2 x i128> %In2, <2 x i128> %In3
+  store <2 x i128> %res, <2 x i128> *%Out
+
+  ret void
+}
+
+; Check case when scalar size is not power of 2.
+define void @bar(i32 %In1, <2 x i96> %In2, <2 x i96> %In3, <2 x i96> *%Out) {
+; CHECK-LABEL: bar:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    and w10, w0, #0x1
+; CHECK-NEXT:    fmov s0, wzr
+; CHECK-NEXT:    fmov s1, w10
+; CHECK-NEXT:    cmeq v0.4s, v1.4s, v0.4s
+; CHECK-NEXT:    ldp x11, x8, [sp, #8]
+; CHECK-NEXT:    ldr x9, [sp]
+; CHECK-NEXT:    dup v1.4s, v0.s[0]
+; CHECK-NEXT:    mov x10, v1.d[1]
+; CHECK-NEXT:    lsr x10, x10, #32
+; CHECK-NEXT:    tst w10, #0x1
+; CHECK-NEXT:    fmov w10, s0
+; CHECK-NEXT:    csel x11, x5, x11, ne
+; CHECK-NEXT:    csel x9, x4, x9, ne
+; CHECK-NEXT:    tst w10, #0x1
+; CHECK-NEXT:    csel x10, x3, x7, ne
+; CHECK-NEXT:    csel x12, x2, x6, ne
+; CHECK-NEXT:    stur x9, [x8, #12]
+; CHECK-NEXT:    str x12, [x8]
+; CHECK-NEXT:    str w10, [x8, #8]
+; CHECK-NEXT:    str w11, [x8, #20]
+; CHECK-NEXT:    ret
+  %cond = and i32 %In1, 1
+  %cbool = icmp eq i32 %cond, 0
+  %res = select i1 %cbool, <2 x i96> %In2, <2 x i96> %In3
+  store <2 x i96> %res, <2 x i96> *%Out
+
+  ret void
+}