[SystemZ] Fix 256-bit shifts when i128 is legal

When i128 is a legal type, SelectionDAG now attempts to use
SRL_PARTS etc. with type i128, which is not implemented.  Fix
by marking those as Expand, just like we do for i64.

Fixes #77132

Author: uweigand

llvm/lib/Target/SystemZ/SystemZISelLowering.cpp

@@ -340,6 +340,13 @@ SystemZTargetLowering::SystemZTargetLowering(const TargetMachine &TM,
   setLibcallName(RTLIB::SHL_I128, nullptr);
   setLibcallName(RTLIB::SRA_I128, nullptr);
 
+  // Also expand 256 bit shifts if i128 is a legal type.
+  if (isTypeLegal(MVT::i128)) {
+    setOperationAction(ISD::SRL_PARTS, MVT::i128, Expand);
+    setOperationAction(ISD::SHL_PARTS, MVT::i128, Expand);
+    setOperationAction(ISD::SRA_PARTS, MVT::i128, Expand);
+  }
+
   // Handle bitcast from fp128 to i128.
   if (!isTypeLegal(MVT::i128))
     setOperationAction(ISD::BITCAST, MVT::i128, Custom);

llvm/test/CodeGen/SystemZ/shift-16.ll

@@ -0,0 +1,132 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 2
+; Test that 256-bit shifts still work when i128 is a legal type
+;
+; RUN: llc < %s -mtriple=s390x-linux-gnu -mcpu=z13 | FileCheck %s
+
+; Shift left.
+define i256 @f1(i256 %a, i256 %sh) {
+; CHECK-LABEL: f1:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vl %v0, 0(%r3), 3
+; CHECK-NEXT:    vl %v1, 16(%r3), 3
+; CHECK-NEXT:    l %r0, 28(%r4)
+; CHECK-NEXT:    clijhe %r0, 128, .LBB0_2
+; CHECK-NEXT:  # %bb.1:
+; CHECK-NEXT:    lhi %r1, 128
+; CHECK-NEXT:    sr %r1, %r0
+; CHECK-NEXT:    vlvgp %v2, %r1, %r1
+; CHECK-NEXT:    vrepb %v2, %v2, 15
+; CHECK-NEXT:    vsrlb %v3, %v1, %v2
+; CHECK-NEXT:    vsrl %v2, %v3, %v2
+; CHECK-NEXT:    vlvgp %v3, %r0, %r0
+; CHECK-NEXT:    vrepb %v3, %v3, 15
+; CHECK-NEXT:    vslb %v4, %v0, %v3
+; CHECK-NEXT:    vslb %v1, %v1, %v3
+; CHECK-NEXT:    vsl %v4, %v4, %v3
+; CHECK-NEXT:    vo %v2, %v4, %v2
+; CHECK-NEXT:    vsl %v1, %v1, %v3
+; CHECK-NEXT:    cijlh %r0, 0, .LBB0_3
+; CHECK-NEXT:    j .LBB0_4
+; CHECK-NEXT:  .LBB0_2:
+; CHECK-NEXT:    ahik %r1, %r0, -128
+; CHECK-NEXT:    vlvgp %v2, %r1, %r1
+; CHECK-NEXT:    vrepb %v2, %v2, 15
+; CHECK-NEXT:    vslb %v1, %v1, %v2
+; CHECK-NEXT:    vsl %v2, %v1, %v2
+; CHECK-NEXT:    vgbm %v1, 0
+; CHECK-NEXT:    cije %r0, 0, .LBB0_4
+; CHECK-NEXT:  .LBB0_3:
+; CHECK-NEXT:    vlr %v0, %v2
+; CHECK-NEXT:  .LBB0_4:
+; CHECK-NEXT:    vst %v1, 16(%r2), 3
+; CHECK-NEXT:    vst %v0, 0(%r2), 3
+; CHECK-NEXT:    br %r14
+  %res = shl i256 %a, %sh
+  ret i256 %res
+}
+
+; Shift right logical.
+define i256 @f2(i256 %a, i256 %sh) {
+; CHECK-LABEL: f2:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vl %v0, 16(%r3), 3
+; CHECK-NEXT:    vl %v1, 0(%r3), 3
+; CHECK-NEXT:    l %r0, 28(%r4)
+; CHECK-NEXT:    clijhe %r0, 128, .LBB1_2
+; CHECK-NEXT:  # %bb.1:
+; CHECK-NEXT:    lhi %r1, 128
+; CHECK-NEXT:    sr %r1, %r0
+; CHECK-NEXT:    vlvgp %v2, %r1, %r1
+; CHECK-NEXT:    vrepb %v2, %v2, 15
+; CHECK-NEXT:    vslb %v3, %v1, %v2
+; CHECK-NEXT:    vsl %v2, %v3, %v2
+; CHECK-NEXT:    vlvgp %v3, %r0, %r0
+; CHECK-NEXT:    vrepb %v3, %v3, 15
+; CHECK-NEXT:    vsrlb %v4, %v0, %v3
+; CHECK-NEXT:    vsrlb %v1, %v1, %v3
+; CHECK-NEXT:    vsrl %v4, %v4, %v3
+; CHECK-NEXT:    vo %v2, %v4, %v2
+; CHECK-NEXT:    vsrl %v1, %v1, %v3
+; CHECK-NEXT:    cijlh %r0, 0, .LBB1_3
+; CHECK-NEXT:    j .LBB1_4
+; CHECK-NEXT:  .LBB1_2:
+; CHECK-NEXT:    ahik %r1, %r0, -128
+; CHECK-NEXT:    vlvgp %v2, %r1, %r1
+; CHECK-NEXT:    vrepb %v2, %v2, 15
+; CHECK-NEXT:    vsrlb %v1, %v1, %v2
+; CHECK-NEXT:    vsrl %v2, %v1, %v2
+; CHECK-NEXT:    vgbm %v1, 0
+; CHECK-NEXT:    cije %r0, 0, .LBB1_4
+; CHECK-NEXT:  .LBB1_3:
+; CHECK-NEXT:    vlr %v0, %v2
+; CHECK-NEXT:  .LBB1_4:
+; CHECK-NEXT:    vst %v1, 0(%r2), 3
+; CHECK-NEXT:    vst %v0, 16(%r2), 3
+; CHECK-NEXT:    br %r14
+  %res = lshr i256 %a, %sh
+  ret i256 %res
+}
+
+; Shift right arithmetic.
+define i256 @f3(i256 %a, i256 %sh) {
+; CHECK-LABEL: f3:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vl %v0, 16(%r3), 3
+; CHECK-NEXT:    vl %v2, 0(%r3), 3
+; CHECK-NEXT:    l %r0, 28(%r4)
+; CHECK-NEXT:    clijhe %r0, 128, .LBB2_2
+; CHECK-NEXT:  # %bb.1:
+; CHECK-NEXT:    lhi %r1, 128
+; CHECK-NEXT:    sr %r1, %r0
+; CHECK-NEXT:    vlvgp %v1, %r0, %r0
+; CHECK-NEXT:    vlvgp %v4, %r1, %r1
+; CHECK-NEXT:    vrepb %v3, %v1, 15
+; CHECK-NEXT:    vrepb %v4, %v4, 15
+; CHECK-NEXT:    vsrab %v1, %v2, %v3
+; CHECK-NEXT:    vslb %v2, %v2, %v4
+; CHECK-NEXT:    vsl %v2, %v2, %v4
+; CHECK-NEXT:    vsrlb %v4, %v0, %v3
+; CHECK-NEXT:    vsra %v1, %v1, %v3
+; CHECK-NEXT:    vsrl %v3, %v4, %v3
+; CHECK-NEXT:    vo %v2, %v3, %v2
+; CHECK-NEXT:    cijlh %r0, 0, .LBB2_3
+; CHECK-NEXT:    j .LBB2_4
+; CHECK-NEXT:  .LBB2_2:
+; CHECK-NEXT:    vrepib %v1, 127
+; CHECK-NEXT:    vsrab %v3, %v2, %v1
+; CHECK-NEXT:    ahik %r1, %r0, -128
+; CHECK-NEXT:    vsra %v1, %v3, %v1
+; CHECK-NEXT:    vlvgp %v3, %r1, %r1
+; CHECK-NEXT:    vrepb %v3, %v3, 15
+; CHECK-NEXT:    vsrab %v2, %v2, %v3
+; CHECK-NEXT:    vsra %v2, %v2, %v3
+; CHECK-NEXT:    cije %r0, 0, .LBB2_4
+; CHECK-NEXT:  .LBB2_3:
+; CHECK-NEXT:    vlr %v0, %v2
+; CHECK-NEXT:  .LBB2_4:
+; CHECK-NEXT:    vst %v1, 0(%r2), 3
+; CHECK-NEXT:    vst %v0, 16(%r2), 3
+; CHECK-NEXT:    br %r14
+  %res = ashr i256 %a, %sh
+  ret i256 %res
+}