[ShrinkWrap] Allow shrinkwrapping past memory accesses to jump tables

[jpenix-quic]

Summary

This patch adds a check for whether the memory operand is known to be
a jump table and, if so, allows shrinkwrapping to continue. In the
case that we are looking at a jump table, I believe it is safe to
assume that the access will not be to the stack (but please correct me
if I am wrong here).

In the test attached, this is helpful in that we are able to generate
only one instruction for each non-default case in the original switch
statement.

[jpenix-quic]

Make if statement slightly more tidy.

[thegameg]

LGTM

llvm/lib/CodeGen/ShrinkWrap.cpp

 INITIALIZE_PASS_DEPENDENCY(MachinePostDominatorTree)
 INITIALIZE_PASS_DEPENDENCY(MachineLoopInfo)
 INITIALIZE_PASS_DEPENDENCY(MachineOptimizationRemarkEmitterPass)
 INITIALIZE_PASS_END(ShrinkWrap, DEBUG_TYPE, "Shrink Wrap Pass", false, false)
 
 bool ShrinkWrap::useOrDefCSROrFI(const MachineInstr &MI,
                                  RegScavenger *RS) const {
   /// Check if \p Op is known to access an address not on the function's stack .
-  /// At the moment, accesses where the underlying object is a global or a
-  /// function argument are considered non-stack accesses. Note that the
+  /// At the moment, accesses where the underlying object is a global, function
+  /// argument, or jump table are considered non-stack accesses. Note that the
   /// caller's stack may get accessed when passing an argument via the stack,
   /// but not the stack of the current function.
   ///
   auto IsKnownNonStackPtr = [](MachineMemOperand *Op) {
     if (Op->getValue()) {
       const Value *UO = getUnderlyingObject(Op->getValue());
       if (!UO)
         return false;
       if (auto *Arg = dyn_cast<Argument>(UO))
         return !Arg->hasPassPointeeByValueCopyAttr();
       return isa<GlobalValue>(UO);
     }
+    if (const PseudoSourceValue *PSV = Op->getPseudoValue())
+      return PSV->isJumpTable();
     return false;
   };
   // This prevents premature stack popping when occurs a indirect stack
   // access.  It is overly aggressive for the moment.
   // TODO:
   //       - Further, data dependency and alias analysis can validate
   //         that load and stores never derive from the stack pointer.
   if (MI.mayLoadOrStore() && (MI.isCall() || MI.hasUnmodeledSideEffects() ||

llvm/test/CodeGen/RISCV/shrinkwrap-jump-table.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 2
+; RUN: llc -mtriple riscv64 < %s | FileCheck %s
+
+declare signext i32 @func1(ptr noundef) local_unnamed_addr
+declare signext i32 @func2(ptr noundef) local_unnamed_addr
+declare signext i32 @func3(ptr noundef) local_unnamed_addr
+declare signext i32 @func4(ptr noundef) local_unnamed_addr
+declare signext i32 @func5(ptr noundef) local_unnamed_addr
+declare signext i32 @default_func(ptr noundef) local_unnamed_addr
+
+define dso_local signext i32 @test_shrinkwrap_jump_table(ptr noundef %m) local_unnamed_addr {
+; CHECK-LABEL: test_shrinkwrap_jump_table:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    lw a1, 0(a0)
+; CHECK-NEXT:    addi a1, a1, -1
+; CHECK-NEXT:    li a2, 4
+; CHECK-NEXT:    bltu a2, a1, .LBB0_3
+; CHECK-NEXT:  # %bb.1: # %entry
+; CHECK-NEXT:    slli a1, a1, 2
+; CHECK-NEXT:    lui a2, %hi(.LJTI0_0)
+; CHECK-NEXT:    addi a2, a2, %lo(.LJTI0_0)
+; CHECK-NEXT:    add a1, a1, a2
+; CHECK-NEXT:    lw a1, 0(a1)
+; CHECK-NEXT:    jr a1
+; CHECK-NEXT:  .LBB0_2: # %sw.bb
+; CHECK-NEXT:    tail func1@plt
+; CHECK-NEXT:  .LBB0_3: # %sw.default
+; CHECK-NEXT:    addi sp, sp, -16
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    sd ra, 8(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    .cfi_offset ra, -8
+; CHECK-NEXT:    call default_func@plt
+; CHECK-NEXT:    li a0, 0
+; CHECK-NEXT:    ld ra, 8(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    addi sp, sp, 16
+; CHECK-NEXT:    ret
+; CHECK-NEXT:  .LBB0_4: # %sw.bb1
+; CHECK-NEXT:    tail func2@plt
+; CHECK-NEXT:  .LBB0_5: # %sw.bb3
+; CHECK-NEXT:    tail func3@plt
+; CHECK-NEXT:  .LBB0_6: # %sw.bb5
+; CHECK-NEXT:    tail func4@plt
+; CHECK-NEXT:  .LBB0_7: # %sw.bb7
+; CHECK-NEXT:    tail func5@plt
+entry:
+  %0 = load i32, ptr %m, align 4
+  switch i32 %0, label %sw.default [
+    i32 1, label %sw.bb
+    i32 2, label %sw.bb1
+    i32 3, label %sw.bb3
+    i32 4, label %sw.bb5
+    i32 5, label %sw.bb7
+  ]
+
+sw.bb:
+  %call = tail call signext i32 @func1(ptr noundef nonnull %m)
+  br label %sw.epilog
+
+sw.bb1:
+  %call2 = tail call signext i32 @func2(ptr noundef nonnull %m)
+  br label %sw.epilog
+
+sw.bb3:
+  %call4 = tail call signext i32 @func3(ptr noundef nonnull %m)
+  br label %sw.epilog
+
+sw.bb5:
+  %call6 = tail call signext i32 @func4(ptr noundef nonnull %m)
+  br label %sw.epilog
+
+sw.bb7:
+  %call8 = tail call signext i32 @func5(ptr noundef nonnull %m)
+  br label %sw.epilog
+
+sw.default:
+  %call9 = tail call signext i32 @default_func(ptr noundef nonnull %m)
+  br label %sw.epilog
+
+sw.epilog:
+  %ret.0 = phi i32 [ 0, %sw.default ], [ %call8, %sw.bb7 ], [ %call6, %sw.bb5 ], [ %call4, %sw.bb3 ], [ %call2, %sw.bb1 ], [ %call, %sw.bb ]
+  ret i32 %ret.0
+}