[ShrinkWrap] Allow shrinkwrapping past memory accesses to jump tables

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.

Differential Revision: https://reviews.llvm.org/D149886

Author: jonathonpenix

llvm/lib/CodeGen/ShrinkWrap.cpp

@@ -288,8 +288,8 @@ 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.
   ///
@@ -302,6 +302,8 @@ bool ShrinkWrap::useOrDefCSROrFI(const MachineInstr &MI,
         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

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

@@ -0,0 +1,82 @@
+; 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
+}