[DAG][AArch64] Fix truncated vscale constant types

[dmgreen]

Summary

It appears that vscale values truncated to i1 causes mismatches in the constant types when created in getNode. https://godbolt.org/z/TaaTo86ne.

[dmgreen]

Just trunc the constant, as opposed to being opinionated about the sign.

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5754–5756

SelectionDAG::getVScale does the right thing by truncating the APInt value, but it only does so after the assert:

assert(MulImm.getSignificantBits() <= VT.getSizeInBits() &&
          "Immediate does not fit VT");

which fails because for MulImm=1, MulImm.getSignificantBits() results in 2. This happens because it adds the sign bit, which for a 1-bit value is nonsensical. I think the better thing to do is to fix the assert itself.

[DavidSpickett]

I have confirmed this fixes the assertion failure seen on the SVE VLA bots.

[dmgreen]

Thanks for checking!

This version adds a special case for i1 to the assert.

[dmgreen]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5754–5756

I'm not sure it wouldn't be better for getVScale to only accept APInt of the correct bitwidth, to be more precise. I couldn't see any other code that looked like the APInt and the VT could be a mismatch.

The new version just adds an override for i1 to the assert. Let me know what you think.

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

1945

I guess this still needs an additional check that for MVT::i1 the immediate fits, i.e. ... || (VT == MVT::i1 && MulImm <= 1) ?

[dmgreen]

I was thinking about it, and I think it would be better to not be sloppy about the APInt passed to getVScale. Similar to existing methods in DAG like getConstant, we should be asserting that the size of the APInt and the type match. Otherwise it should be up to the caller to make them match in a way that is appropriate.

It appears from looking through the calls to getVScale that the version in getNode when truncating the value is the only one where they can be a mismatch.

[dmgreen]

Assert the sizes of the type and APInt match in getVScale.

[awarzynski]

Thanks for fixing this @dmgreen !

Any chance to land this soon? We (folks working on SVE/SSVE/SME support in MLIR) rely heavily on clang-aarch64-sve-vla that's currently failing.

[sdesmalen]

In D155626#4516126, @dmgreen wrote:

I was thinking about it, and I think it would be better to not be sloppy about the APInt passed to getVScale. Similar to existing methods in DAG like getConstant, we should be asserting that the size of the APInt and the type match. Otherwise it should be up to the caller to make them match in a way that is appropriate.

It appears from looking through the calls to getVScale that the version in getNode when truncating the value is the only one where they can be a mismatch.

That makes sense. I thought there would have been more instances that needed fixing, but glad to see that wasn't the case.

[dmgreen]

Thanks. I'll commit this now, just after my machine has finished rerunning the tests. Hopefully I didn't miss any cases for getVScale.

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

if (!CondCodeNodes[Cond]) {
     InsertNode(N);
   }
 
   return SDValue(CondCodeNodes[Cond], 0);
 }
 
 SDValue SelectionDAG::getVScale(const SDLoc &DL, EVT VT, APInt MulImm,
                                 bool ConstantFold) {
   assert(MulImm.getSignificantBits() <= VT.getSizeInBits() &&
          "Immediate does not fit VT");
 
   MulImm = MulImm.sextOrTrunc(VT.getSizeInBits());
 
   if (ConstantFold) {
     const MachineFunction &MF = getMachineFunction();
     auto Attr = MF.getFunction().getFnAttribute(Attribute::VScaleRange);
     if (Attr.isValid()) {
if (OpOpcode == ISD::ZERO_EXTEND || OpOpcode == ISD::SIGN_EXTEND ||
         return getNode(OpOpcode, DL, VT, N1.getOperand(0));
       if (N1.getOperand(0).getValueType().bitsGT(VT))
         return getNode(ISD::TRUNCATE, DL, VT, N1.getOperand(0));
       return N1.getOperand(0);
     }
     if (OpOpcode == ISD::UNDEF)
       return getUNDEF(VT);
     if (OpOpcode == ISD::VSCALE && !NewNodesMustHaveLegalTypes)
-      return getVScale(DL, VT, N1.getConstantOperandAPInt(0));
+      return getVScale(
+          DL, VT,
+          N1.getConstantOperandAPInt(0).zextOrTrunc(VT.getScalarSizeInBits()));
     break;
   case ISD::ANY_EXTEND_VECTOR_INREG:
   case ISD::ZERO_EXTEND_VECTOR_INREG:
   case ISD::SIGN_EXTEND_VECTOR_INREG:
     assert(VT.isVector() && "This DAG node is restricted to vector types.");
     assert(N1.getValueType().bitsLE(VT) &&
            "The input must be the same size or smaller than the result.");
     assert(VT.getVectorMinNumElements() <

llvm/test/CodeGen/AArch64/sve-vscale.ll

 ; CHECK:       rdvl x0, #31
 ; CHECK-NEXT:  ret
 define i32 @rdvl_max() nounwind {
   %vscale = call i32 @llvm.vscale.i32()
   %1 = mul nsw i32 %vscale, 496
   ret i32 %1
 }
 
+define i1 @rdvl_i1() {
+; CHECK-LABEL: rdvl_i1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    rdvl x8, #-1
+; CHECK-NEXT:    asr x8, x8, #4
+; CHECK-NEXT:    and w0, w8, #0x1
+; CHECK-NEXT:    ret
+  %a = tail call i64 @llvm.vscale.i64()
+  %b = trunc i64 %a to i1
+  ret i1 %b
+}
+
 ;
 ; CNTH
 ;
 
 ; CHECK-LABEL: cnth:
 ; CHECK:       cnth x0{{$}}
 ; CHECK-NEXT:  ret
 define i32 @cnth() nounwind {