[llvm] llvm-reduce: Add values to return reduction (PR #132686)

[Nikita Popov via llvm-commits]

================
@@ -0,0 +1,244 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Try to reduce a function by inserting new return instructions. Try to insert
+// an early return for each instruction value at that point. This requires
+// mutating the return type, or finding instructions with a compatible type.
+//
+//===----------------------------------------------------------------------===//
+
+#define DEBUG_TYPE "llvm-reduce"
+
+#include "ReduceValuesToReturn.h"
+
+#include "Delta.h"
+#include "Utils.h"
+#include "llvm/IR/CFG.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/Cloning.h"
+
+using namespace llvm;
+
+/// Return true if it is legal to emit a copy of the function with a non-void
+/// return type.
+static bool canUseNonVoidReturnType(const Function &F) {
+  // Functions with sret arguments must return void.
+  return !F.hasStructRetAttr() &&
+         CallingConv::supportsNonVoidReturnType(F.getCallingConv());
+}
+
+/// Return true if it's legal to replace a function return type to use \p Ty.
+static bool isReallyValidReturnType(Type *Ty) {
+  return FunctionType::isValidReturnType(Ty) && !Ty->isTokenTy() &&
+         Ty->isFirstClassType();
+}
+
+/// Insert a ret inst after \p NewRetValue, which returns the value it produces.
+static void rewriteFuncWithReturnType(Function &OldF, Value *NewRetValue) {
+  Type *NewRetTy = NewRetValue->getType();
+  FunctionType *OldFuncTy = OldF.getFunctionType();
+
+  FunctionType *NewFuncTy =
+      FunctionType::get(NewRetTy, OldFuncTy->params(), OldFuncTy->isVarArg());
+
+  LLVMContext &Ctx = OldF.getContext();
+  Instruction *NewRetI = cast<Instruction>(NewRetValue);
+  BasicBlock *NewRetBlock = NewRetI->getParent();
+
+  BasicBlock::iterator NewValIt = NewRetI->getIterator();
+
+  // Hack up any return values in other blocks, we can't leave them as ret void.
+  if (OldFuncTy->getReturnType()->isVoidTy()) {
+    for (BasicBlock &OtherRetBB : OldF) {
+      if (&OtherRetBB != NewRetBlock) {
+        auto *OrigRI = dyn_cast<ReturnInst>(OtherRetBB.getTerminator());
+        if (!OrigRI)
+          continue;
+
+        OrigRI->eraseFromParent();
+        ReturnInst::Create(Ctx, getDefaultValue(NewRetTy), &OtherRetBB);
+      }
+    }
+  }
+
+  // Now prune any CFG edges we have to deal with.
+  //
+  // Use KeepOneInputPHIs in case the instruction we are using for the return is
+  // that phi.
+  // TODO: Could avoid this with fancier iterator management.
+  for (BasicBlock *Succ : successors(NewRetBlock))
+    Succ->removePredecessor(NewRetBlock, /*KeepOneInputPHIs=*/true);
+
+  // Now delete the tail of this block, in reverse to delete uses before defs.
+  for (Instruction &I : make_early_inc_range(
+           make_range(NewRetBlock->rbegin(), NewValIt.getReverse()))) {
+    Value *Replacement = getDefaultValue(I.getType());
+    I.replaceAllUsesWith(Replacement);
+    I.eraseFromParent();
+  }
+
+  ReturnInst::Create(Ctx, NewRetValue, NewRetBlock);
+
+  // TODO: We may be eliminating blocks that were originally unreachable. We
+  // probably ought to only be pruning blocks that became dead directly as a
+  // result of our pruning here.
+  EliminateUnreachableBlocks(OldF);
+
+  Function *NewF =
+      Function::Create(NewFuncTy, OldF.getLinkage(), OldF.getAddressSpace(), "",
+                       OldF.getParent());
+
+  NewF->removeFromParent();
+  OldF.getParent()->getFunctionList().insertAfter(OldF.getIterator(), NewF);
+  NewF->takeName(&OldF);
+  NewF->copyAttributesFrom(&OldF);
+
+  // Adjust the callsite uses to the new return type. We pre-filtered cases
+  // where the original call type was incorrectly non-void.
+  for (User *U : make_early_inc_range(OldF.users())) {
+    if (auto *CB = dyn_cast<CallBase>(U);
+        CB && CB->getCalledOperand() == &OldF) {
+      if (CB->getType()->isVoidTy()) {
+        FunctionType *CallType = CB->getFunctionType();
+
+        // The callsite may not match the new function type, in an undefined
+        // behavior way. Only mutate the local return type.
+        FunctionType *NewCallType = FunctionType::get(
+            NewRetTy, CallType->params(), CallType->isVarArg());
+
+        CB->mutateType(NewRetTy);
+        CB->setCalledFunction(NewCallType, NewF);
+      } else {
+        assert(CB->getType() == NewRetTy &&
+               "only handle exact return type match with non-void returns");
+      }
+    }
+  }
+
+  // Preserve the parameters of OldF.
+  ValueToValueMapTy VMap;
+  for (auto Z : zip_first(OldF.args(), NewF->args())) {
+    Argument &OldArg = std::get<0>(Z);
+    Argument &NewArg = std::get<1>(Z);
+
+    NewArg.setName(OldArg.getName()); // Copy the name over...
+    VMap[&OldArg] = &NewArg;          // Add mapping to VMap
+  }
+
+  SmallVector<ReturnInst *, 8> Returns; // Ignore returns cloned.
+  CloneFunctionInto(NewF, &OldF, VMap,
+                    CloneFunctionChangeType::LocalChangesOnly, Returns, "",
+                    /*CodeInfo=*/nullptr);
----------------
nikic wrote:

> No, Function doesn't have a way to mutate the FunctionType. The NumArgs are also an allocated array from the original signature, so it would need some new non-trivial API to hack up a function with a new signature

I was only thinking about the particular case where where we're not changing arguments. But yeah, looks like an API for that currently doesn't exist, though possibly it should. E.g. we recently relaxed this for globals, and you can now change their value type using replaceInitializer().

But in any case, I still think that this should not be using cloning and instead moving the instructions. This is the pattern I had in mind: https://github.com/llvm/llvm-project/blob/1004fae222efeee215780c4bb4e64eb82b07fb4f/llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp#L225-L239

(Really, we should have a utility for changing a functions signature, as this comes up in quite a few places, and every one really shouldn't be coming up with its own bespoke way of doing things.)

https://github.com/llvm/llvm-project/pull/132686