[RISCV] Add support for call returns to RISCVSExtWRemoval.

[craig.topper]

Summary

The implementation is inspired by code from PPCInstrInfo.

We look for a copy from X10(A0) preceded by an ADJCALLSTACKUP and
a PseudoCall. We use the PseudoCall to look up the IR function being
called to find it's return attributes.

[asb]

Left a few comments based on an initial shallow review.

Is it possible to test the zext case, or is it too fiddly to construct a reasonable test case?

llvm/lib/Target/RISCV/RISCVSExtWRemoval.cpp

370

I think the note in your commit message "We use the PseudoCall to look up the IR function being called to find it's return attributes." is helpful to comprehension.

How about "For a method return value, we check the ZExt/SExt attributes, using the PseudoCALL instruction to retrieve the IR function."?

376

auto?

388

Maybe auto? We're not overly consistent on this in LLVM and it's short anyway. But it would match the use of auto a few lines down.

[craig.topper]

Address review comments. Add zext test.

[asb]

LGTM.

llvm/lib/Target/RISCV/RISCVSExtWRemoval.cpp

 //===-------------- RISCVSExtWRemoval.cpp - MI sext.w Removal -------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===---------------------------------------------------------------------===//
 //
case RISCV::COPY: {
       // If this is the entry block and the register is livein, see if we know
       // it is sign extended.
       if (MI->getParent() == &MF->front()) {
         Register VReg = MI->getOperand(0).getReg();
         if (MF->getRegInfo().isLiveIn(VReg) && RVFI->isSExt32Register(VReg))
           continue;
       }
 
-      // TODO: Handle returns from calls?
+      Register CopySrcReg = MI->getOperand(1).getReg();
+      if (CopySrcReg == RISCV::X10) {
+        // For a method return value, we check the ZExt/SExt flags in attribute.
+        // We assume the following code sequence for method call.
+        // PseudoCALL @bar, ...
+        // ADJCALLSTACKUP 0, 0, implicit-def dead $x2, implicit $x2
+        // %0:gpr = COPY $x10
+        //
+        // We use the PseudoCall to look up the IR function being called to find
+        // its return attributes.
+        const MachineBasicBlock *MBB = MI->getParent();
+        auto II = MI->getIterator();
+        if (II == MBB->instr_begin() ||
+            (--II)->getOpcode() != RISCV::ADJCALLSTACKUP)
+          return false;
 
-      if (!AddRegDefToWorkList(MI->getOperand(1).getReg()))
+        const MachineInstr &CallMI = *(--II);
+        if (!CallMI.isCall() || !CallMI.getOperand(0).isGlobal())
+          return false;
+
+        auto *CalleeFn =
+            dyn_cast_if_present<Function>(CallMI.getOperand(0).getGlobal());
+        if (!CalleeFn)
+          return false;
+
+        auto *IntTy = dyn_cast<IntegerType>(CalleeFn->getReturnType());
+        if (!IntTy)
+          return false;
+
+        const AttributeSet &Attrs = CalleeFn->getAttributes().getRetAttrs();
+        unsigned BitWidth = IntTy->getBitWidth();
+        return (BitWidth <= 32 && Attrs.hasAttribute(Attribute::SExt)) ||
+               (BitWidth < 32 && Attrs.hasAttribute(Attribute::ZExt));
+      }
+
+      if (!AddRegDefToWorkList(CopySrcReg))
         return false;
 
       break;
     }
 
     // For these, we just need to check if the 1st operand is sign extended.
     case RISCV::BCLRI:
     case RISCV::BINVI:

llvm/test/CodeGen/RISCV/sextw-removal.ll

if.end:                                           ; preds = %entry
   %x.addr.3.lobit.not = xor i32 %x.addr.3.lobit, -1
   %r.4 = add nsw i32 %r.3, %x.addr.3.lobit.not
   br label %cleanup
 
 cleanup:                                          ; preds = %entry, %if.end
   %retval.0 = phi i32 [ %r.4, %if.end ], [ 0, %entry ]
   ret i32 %retval.0
 }
+
+define void @test16(i32 signext %arg, i32 signext %arg1) nounwind {
+; CHECK-LABEL: test16:
+; CHECK:       # %bb.0: # %bb
+; CHECK-NEXT:    addi sp, sp, -32
+; CHECK-NEXT:    sd ra, 24(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    sd s0, 16(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    sd s1, 8(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    mv s0, a1
+; CHECK-NEXT:    call bar@plt
+; CHECK-NEXT:    mv s1, a0
+; CHECK-NEXT:  .LBB19_1: # %bb2
+; CHECK-NEXT:    # =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    mv a0, s1
+; CHECK-NEXT:    call bar@plt
+; CHECK-NEXT:    sllw s1, s1, s0
+; CHECK-NEXT:    bnez a0, .LBB19_1
+; CHECK-NEXT:  # %bb.2: # %bb7
+; CHECK-NEXT:    ld ra, 24(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    ld s0, 16(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    ld s1, 8(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    addi sp, sp, 32
+; CHECK-NEXT:    ret
+;
+; NOREMOVAL-LABEL: test16:
+; NOREMOVAL:       # %bb.0: # %bb
+; NOREMOVAL-NEXT:    addi sp, sp, -32
+; NOREMOVAL-NEXT:    sd ra, 24(sp) # 8-byte Folded Spill
+; NOREMOVAL-NEXT:    sd s0, 16(sp) # 8-byte Folded Spill
+; NOREMOVAL-NEXT:    sd s1, 8(sp) # 8-byte Folded Spill
+; NOREMOVAL-NEXT:    mv s0, a1
+; NOREMOVAL-NEXT:    call bar@plt
+; NOREMOVAL-NEXT:    mv s1, a0
+; NOREMOVAL-NEXT:  .LBB19_1: # %bb2
+; NOREMOVAL-NEXT:    # =>This Inner Loop Header: Depth=1
+; NOREMOVAL-NEXT:    sext.w a0, s1
+; NOREMOVAL-NEXT:    call bar@plt
+; NOREMOVAL-NEXT:    sllw s1, s1, s0
+; NOREMOVAL-NEXT:    bnez a0, .LBB19_1
+; NOREMOVAL-NEXT:  # %bb.2: # %bb7
+; NOREMOVAL-NEXT:    ld ra, 24(sp) # 8-byte Folded Reload
+; NOREMOVAL-NEXT:    ld s0, 16(sp) # 8-byte Folded Reload
+; NOREMOVAL-NEXT:    ld s1, 8(sp) # 8-byte Folded Reload
+; NOREMOVAL-NEXT:    addi sp, sp, 32
+; NOREMOVAL-NEXT:    ret
+bb:
+  %i = call signext i32 @bar(i32 signext %arg)
+  br label %bb2
+
+bb2:                                              ; preds = %bb2, %bb
+  %i3 = phi i32 [ %i, %bb ], [ %i5, %bb2 ]
+  %i4 = tail call signext i32 @bar(i32 signext %i3)
+  %i5 = shl i32 %i3, %arg1
+  %i6 = icmp eq i32 %i4, 0
+  br i1 %i6, label %bb7, label %bb2
+
+bb7:                                              ; preds = %bb2
+  ret void
+}
+
+define void @test17(i32 signext %arg, i32 signext %arg1) nounwind {
+; CHECK-LABEL: test17:
+; CHECK:       # %bb.0: # %bb
+; CHECK-NEXT:    addi sp, sp, -32
+; CHECK-NEXT:    sd ra, 24(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    sd s0, 16(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    sd s1, 8(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    mv s0, a1
+; CHECK-NEXT:    call bat@plt
+; CHECK-NEXT:    mv s1, a0
+; CHECK-NEXT:  .LBB20_1: # %bb2
+; CHECK-NEXT:    # =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    mv a0, s1
+; CHECK-NEXT:    call bar@plt
+; CHECK-NEXT:    sllw s1, s1, s0
+; CHECK-NEXT:    bnez a0, .LBB20_1
+; CHECK-NEXT:  # %bb.2: # %bb7
+; CHECK-NEXT:    ld ra, 24(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    ld s0, 16(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    ld s1, 8(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    addi sp, sp, 32
+; CHECK-NEXT:    ret
+;
+; NOREMOVAL-LABEL: test17:
+; NOREMOVAL:       # %bb.0: # %bb
+; NOREMOVAL-NEXT:    addi sp, sp, -32
+; NOREMOVAL-NEXT:    sd ra, 24(sp) # 8-byte Folded Spill
+; NOREMOVAL-NEXT:    sd s0, 16(sp) # 8-byte Folded Spill
+; NOREMOVAL-NEXT:    sd s1, 8(sp) # 8-byte Folded Spill
+; NOREMOVAL-NEXT:    mv s0, a1
+; NOREMOVAL-NEXT:    call bat@plt
+; NOREMOVAL-NEXT:    mv s1, a0
+; NOREMOVAL-NEXT:  .LBB20_1: # %bb2
+; NOREMOVAL-NEXT:    # =>This Inner Loop Header: Depth=1
+; NOREMOVAL-NEXT:    sext.w a0, s1
+; NOREMOVAL-NEXT:    call bar@plt
+; NOREMOVAL-NEXT:    sllw s1, s1, s0
+; NOREMOVAL-NEXT:    bnez a0, .LBB20_1
+; NOREMOVAL-NEXT:  # %bb.2: # %bb7
+; NOREMOVAL-NEXT:    ld ra, 24(sp) # 8-byte Folded Reload
+; NOREMOVAL-NEXT:    ld s0, 16(sp) # 8-byte Folded Reload
+; NOREMOVAL-NEXT:    ld s1, 8(sp) # 8-byte Folded Reload
+; NOREMOVAL-NEXT:    addi sp, sp, 32
+; NOREMOVAL-NEXT:    ret
+bb:
+  %i = call zeroext i16 @bat(i32 signext %arg)
+  %zext = zext i16 %i to i32
+  br label %bb2
+
+bb2:                                              ; preds = %bb2, %bb
+  %i3 = phi i32 [ %zext, %bb ], [ %i5, %bb2 ]
+  %i4 = tail call signext i32 @bar(i32 signext %i3)
+  %i5 = shl i32 %i3, %arg1
+  %i6 = icmp eq i32 %i4, 0
+  br i1 %i6, label %bb7, label %bb2
+
+bb7:                                              ; preds = %bb2
+  ret void
+}
+declare zeroext i16 @bat(i32 signext)