Implement support for AArch64ISD::MOVI in computeKnownBits

[adriantong1024]

Summary

This helps simplify a USHR+ORR into USRA on AArch64

[dmgreen]

Sounds like a good use of isTargetCanonicalConstantNode.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

1977

cast<> will already assert.

1979

Does CN->getAPIntValue() work, or does it need to change the width of the APInt?

[foad]

llvm/test/CodeGen/AArch64/shift-accumulate.ll

124

Pre-commit this test and rebase the current patch on top of it?

[tschuett]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

1976

Could you you change it to dyn_cast with if to get the correct assert and error message?

[adriantong1024]

Address comments from reviewers.

Thanks !

[adriantong1024]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

1979

I think its needs to be the same width as Known. Otherwise Known bitwidth will be changed and cause trouble later.

llvm/test/CodeGen/AArch64/shift-accumulate.ll

124

Hi Jay

Thank you for the comment. Do you mean I should change this test and commit it as a separate patch before the actual change (i.e. the implementation of the MOVI) ?

[foad]

llvm/test/CodeGen/AArch64/shift-accumulate.ll

124

Yes. Commit the test in a form that will pass on trunk, and then when you rebase this patch, it'll show the difference in the codegen caused by your patch. That makes it much easier for reviewers to see the effect of your patch.

(I'm not an AArch64 reviewer but this is pretty common practice for the backends I've worked on.)

[adriantong1024]

llvm/test/CodeGen/AArch64/shift-accumulate.ll

124

Makes sense !. I will do so once this patch is accepted. Thanks for teaching me this.

[tschuett]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

1976

I still believe that you need and if to check in Release mode whether it really is a constant node.

[foad]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

1976

No, just use cast<> and remove the assert.

[adriantong1024]

Address comments from reviewers. Thanks !

[adriantong1024]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

1976

Thanks for the comment. I went back to check "The cast<> operator is a “checked cast” operation. It converts a pointer or reference from a base class to a derived class, causing an assertion failure if it is not really an instance of the right type." This seems to be what we want.

Let me know if there are other concerns!

[dmgreen]

There are other MOVI nodes created in ISel, have you considered adding support for those too?

MOVI,
MOVIshift,
MOVIedit,
MOVImsl,
FMOV,
MVNIshift,
MVNImsl,

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

1976

cast will do the assert for you. I would think it would be OK to always use case and drop the assert.

The operand of a MOVI will always be a ConstantSDNode. There is no need to check for it in release builds, it should be correct by construction.

llvm/test/CodeGen/AArch64/shift-accumulate.ll

124

Can you expand this to other types too. For example this, as well as i64 variants:

define <4 x i32> @usra_with_movi16(<8 x i16> %0, <8 x i16> %1) {
; CHECK-LABEL: usra_with_movi:
; CHECK:       // %bb.0:
; CHECK-NEXT:    movi v2.16b, #1
; CHECK-NEXT:    cmeq v0.16b, v0.16b, v1.16b
; CHECK-NEXT:    and v0.16b, v0.16b, v2.16b
; CHECK-NEXT:    usra v0.8h, v0.8h, #7
; CHECK-NEXT:    ret
  %3 = icmp eq <8 x i16> %0, %1
  %4 = zext <8 x i1> %3 to <8 x i16>
  %5 = bitcast <8 x i16> %4 to <4 x i32>
  %6 = lshr <4 x i32> %5, <i32 7, i32 7, i32 7, i32 7>
  %7 = or <4 x i32> %6, %5
  ret <4 x i32> %7
}

[adriantong1024]

In D137108#3899142, @dmgreen wrote:

There are other MOVI nodes created in ISel, have you considered adding support for those too?

MOVI,
MOVIshift,
MOVIedit,
MOVImsl,
FMOV,
MVNIshift,
MVNImsl,

This patch comes out of a workload we have inside the company. Only MOVI needs to be fixed here. I will certainly implement support for others if there are use cases for them.

[adriantong1024]

Address test case comment from dmgreen. Thanks !

[dmgreen]

It may be worth implementing them in separate patches to keep the changes small. If you can add the test cases for different sizes, this LGTM.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

1979

Makes sense. The constant should always be an i32, even though the typesize is i8.

[dmgreen]

Thanks. LGTM

[adriantong1024]

Rebase.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

void AArch64TargetLowering::computeKnownBitsForTargetNode(
   }
   case AArch64ISD::VASHR: {
     KnownBits Known2;
     Known = DAG.computeKnownBits(Op->getOperand(0), Depth + 1);
     Known2 = DAG.computeKnownBits(Op->getOperand(1), Depth + 1);
     Known = KnownBits::ashr(Known, Known2);
     break;
   }
+  case AArch64ISD::MOVI: {
+    ConstantSDNode *CN = cast<ConstantSDNode>(Op->getOperand(0));
+    Known =
+        KnownBits::makeConstant(APInt(Known.getBitWidth(), CN->getZExtValue()));
+    break;
+  }
   case AArch64ISD::LOADgot:
   case AArch64ISD::ADDlow: {
     if (!Subtarget->isTargetILP32())
       break;
     // In ILP32 mode all valid pointers are in the low 4GB of the address-space.
     Known.Zero = APInt::getHighBitsSet(64, 32);
     break;
   }
bool AArch64TargetLowering::SimplifyDemandedBitsForTargetNode(
   }
 
   return TargetLowering::SimplifyDemandedBitsForTargetNode(
       Op, OriginalDemandedBits, OriginalDemandedElts, Known, TLO, Depth);
 }
 
 bool AArch64TargetLowering::isTargetCanonicalConstantNode(SDValue Op) const {
   return Op.getOpcode() == AArch64ISD::DUP ||
+         Op.getOpcode() == AArch64ISD::MOVI ||
          (Op.getOpcode() == ISD::EXTRACT_SUBVECTOR &&
           Op.getOperand(0).getOpcode() == AArch64ISD::DUP) ||
          TargetLowering::isTargetCanonicalConstantNode(Op);
 }
 
 bool AArch64TargetLowering::isConstantUnsignedBitfieldExtractLegal(
     unsigned Opc, LLT Ty1, LLT Ty2) const {
   return Ty1 == Ty2 && (Ty1 == LLT::scalar(32) || Ty1 == LLT::scalar(64));
 }

llvm/test/CodeGen/AArch64/fast-isel-cmp-vec.ll

bb2:
   %z = zext <16 x i1> %c to <16 x i8>
   ret <16 x i8> %z
 }
 
 define <16 x i8> @icmp_constfold_v16i8(<16 x i8> %a) {
 ; CHECK-LABEL: icmp_constfold_v16i8:
 ; CHECK:       ; %bb.0:
 ; CHECK-NEXT:    movi.16b v0, #1
-; CHECK-NEXT:    and.16b v0, v0, v0
 ; CHECK-NEXT:    ret
   %1 = icmp eq <16 x i8> %a, %a
   br label %bb2
 bb2:
   %2 = zext <16 x i1> %1 to <16 x i8>
   ret <16 x i8> %2
 }

llvm/test/CodeGen/AArch64/shift-accumulate.ll

; CHECK-NEXT:    ret
   %2 = and <4 x i32> %0, <i32 3221225471, i32 3221225471,i32 3221225471,i32 3221225471>
   %3 = bitcast <4 x i32> %2 to <2 x i64>
   %4 = lshr <2 x i64> %3, <i64 63, i64 63>
   %5 = ashr <2 x i64> %3, <i64 62, i64 62>
   %6 = or <2 x i64> %4, %5
   ret <2 x i64> %6
 }
 
 ; Expected to be able to deduce movi is generate a vector of integer
 ; and turn USHR+ORR into USRA.
 define <8 x i16> @usra_with_movi_v8i16(<16 x i8> %0, <16 x i8> %1) {
 ; CHECK-LABEL: usra_with_movi_v8i16:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    movi v2.16b, #1
 ; CHECK-NEXT:    cmeq v0.16b, v0.16b, v1.16b
 ; CHECK-NEXT:    and v0.16b, v0.16b, v2.16b
-; CHECK-NEXT:    ushr v1.8h, v0.8h, #7
-; CHECK-NEXT:    orr v0.16b, v1.16b, v0.16b
+; CHECK-NEXT:    usra v0.8h, v0.8h, #7
 ; CHECK-NEXT:    ret
   %3 = icmp eq <16 x i8> %0, %1
   %4 = zext <16 x i1> %3 to <16 x i8>
   %5 = bitcast <16 x i8> %4 to <8 x i16>
   %6 = lshr <8 x i16> %5, <i16 7, i16 7, i16 7, i16 7, i16 7, i16 7, i16 7, i16 7>
   %7 = or <8 x i16> %6, %5
   ret <8 x i16> %7
 }
 
 ; Expected to be able to deduce movi is generate a vector of integer
 ; and turn USHR+ORR into USRA.
 define <4 x i32> @usra_with_movi_v4i32(<16 x i8> %0, <16 x i8> %1) {
 ; CHECK-LABEL: usra_with_movi_v4i32:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    movi v2.16b, #1
 ; CHECK-NEXT:    cmeq v0.16b, v0.16b, v1.16b
 ; CHECK-NEXT:    and v0.16b, v0.16b, v2.16b
-; CHECK-NEXT:    ushr v1.4s, v0.4s, #15
-; CHECK-NEXT:    orr v0.16b, v1.16b, v0.16b
+; CHECK-NEXT:    usra v0.4s, v0.4s, #15
 ; CHECK-NEXT:    ret
   %3 = icmp eq <16 x i8> %0, %1
   %4 = zext <16 x i1> %3 to <16 x i8>
   %5 = bitcast <16 x i8> %4 to <4 x i32>
   %6 = lshr <4 x i32> %5, <i32 15, i32 15, i32 15, i32 15>
   %7 = or <4 x i32> %6, %5
   ret <4 x i32> %7
 }
 
 ; Expected to be able to deduce movi is generate a vector of integer
 ; and turn USHR+ORR into USRA.
 define <2 x i64> @usra_with_movi_v2i64(<16 x i8> %0, <16 x i8> %1) {
 ; CHECK-LABEL: usra_with_movi_v2i64:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    movi v2.16b, #1
 ; CHECK-NEXT:    cmeq v0.16b, v0.16b, v1.16b
 ; CHECK-NEXT:    and v0.16b, v0.16b, v2.16b
-; CHECK-NEXT:    ushr v1.2d, v0.2d, #31
-; CHECK-NEXT:    orr v0.16b, v1.16b, v0.16b
+; CHECK-NEXT:    usra v0.2d, v0.2d, #31
 ; CHECK-NEXT:    ret
   %3 = icmp eq <16 x i8> %0, %1
   %4 = zext <16 x i1> %3 to <16 x i8>
   %5 = bitcast <16 x i8> %4 to <2 x i64>
   %6 = lshr <2 x i64> %5, <i64 31, i64 31>
   %7 = or <2 x i64> %6, %5
   ret <2 x i64> %7
 }