[IROutliner] Ensure that phi values that are passed in as arguments are remapped as arguments

[AndrewLitteken]

Summary

Issue: https://github.com/llvm/llvm-project/issues/54430

For incoming values of phi nodes added to an outlined function to accommodate different exit paths in the function, when a value is a constant that is passed into the outlined function as an argument, we find the corresponding value in the first extracted function used to fill the overall outlined function. When this value is an argument, the corresponding value used will be the old value, prior to outlining. This patch maintains a mapping from these values to arguments, and uses this mapping to update the added phi node accordingly.

[paquette]

LGTM with nit on a comment

llvm/test/Transforms/IROutliner/exit-phi-nodes-incoming-value-constant-argument.ll

5

This sentence doesn't make sense?

llvm/include/llvm/Transforms/IPO/IROutliner.h

struct OutlinableRegion {
   /// Mapping the extracted argument number to the argument number in the
   /// overall function.  Since there will be inputs, such as elevated constants
   /// that are not the same in each region in a SimilarityGroup, or values that
   /// cannot be sunk into the extracted section in every region, we must keep
   /// track of which extracted argument maps to which overall argument.
   DenseMap<unsigned, unsigned> ExtractedArgToAgg;
   DenseMap<unsigned, unsigned> AggArgToExtracted;
 
+  /// Values in the outlined functions will often be replaced by arguments. When
+  /// finding corresponding values from one region to another, the found value
+  /// will be the value the argument previously replaced.  This structure maps
+  /// any replaced values for the region to the aggregate aggregate argument
+  /// in the overall function.
+  DenseMap<Value *, Value *> RemappedArguments;
+
   /// Marks whether we need to change the order of the arguments when mapping
   /// the old extracted function call to the new aggregate outlined function
   /// call.
   bool ChangedArgOrder = false;
 
   /// Marks whether this region ends in a branch, there is special handling
   /// required for the following basic blocks in this case.
   bool EndsInBranch = false;

llvm/lib/Transforms/IPO/IROutliner.cpp

if (Argument *A = dyn_cast<Argument>(IncomingVal)) {
       NewPN->setIncomingValue(Idx, Val);
       continue;
     }
     
     // Find the corresponding value in the overall function.
     IncomingVal = findOutputMapping(OutputMappings, IncomingVal);
     Value *Val = Region.findCorrespondingValueIn(*FirstRegion, IncomingVal);
     assert(Val && "Value is nullptr?");
+    DenseMap<Value *, Value *>::iterator RemappedIt =
+        FirstRegion->RemappedArguments.find(Val);
+    if (RemappedIt != FirstRegion->RemappedArguments.end())
+      Val = RemappedIt->second;
     NewPN->setIncomingValue(Idx, Val);
   }
   return NewPN;
 }
 
 // Within an extracted function, replace the argument uses of the extracted
 // region with the arguments of the function for an OutlinableGroup.
 //
for (unsigned ArgIdx = 0; ArgIdx < Region.ExtractedFunction->arg_size();
     Argument *Arg = Region.ExtractedFunction->getArg(ArgIdx);
     // The argument is an input, so we can simply replace it with the overall
     // argument value
     if (ArgIdx < Region.NumExtractedInputs) {
       LLVM_DEBUG(dbgs() << "Replacing uses of input " << *Arg << " in function "
                         << *Region.ExtractedFunction << " with " << *AggArg
                         << " in function " << *Group.OutlinedFunction << "\n");
       Arg->replaceAllUsesWith(AggArg);
+      Value *V = Region.Call->getArgOperand(ArgIdx);
+      Region.RemappedArguments.insert(std::make_pair(V, AggArg));
       continue;
     }
 
     // If we are replacing an output, we place the store value in its own
     // block inside the overall function before replacing the use of the output
     // in the function.
     assert(Arg->hasOneUse() && "Output argument can only have one use");
     User *InstAsUser = Arg->user_back();

llvm/test/Transforms/IROutliner/exit-phi-nodes-incoming-value-constant-argument.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --include-generated-funcs
+; RUN: opt -S -verify -iroutliner -ir-outlining-no-cost < %s | FileCheck %s
+
+; When consolidating PHINodes, the outliner replaces the incoming value with
+; a corresponding value from the first outlined section.  When this replaced
+; value is passed in as an argument, the corresponding value is found outside
+; of the outlined region, and must be replaced with an argument to avoid
+; dominating value errors. This checks that we use the argument to replace
+; the incoming value.
+
+define void @func1(i32 %0, i32 %1) local_unnamed_addr #0 {
+bb1:
+  br label %bb5s
+
+bb2:
+  %a = add i32 %0, %1
+  %b = add i32 %0, %1
+  %c = icmp eq i32 %b, %a
+  br i1 %c, label %bb5, label %bb3
+
+bb3:
+  %d = add i32 %0, %1
+  br label %bb5
+
+bb4:
+  %e = sub i32 %0, %1
+  br label %bb2
+
+bb5:
+  ret void
+}
+
+define void @func2(i32 %0, i32 %1) local_unnamed_addr #0 {
+bb1:
+  br label %bb5
+
+bb2:
+  %a = sub i32 %0, %1
+  %b = add i32 %0, %1
+  %c = icmp eq i32 %b, 1
+  br i1 %c, label %bb5, label %bb3
+
+bb3:
+  %d = add i32 %0, %1
+  br label %bb5
+
+bb4:
+  %e = add i32 %0, %1
+  br label %bb2
+
+bb5:
+  %f = phi i32 [ 0, %bb1 ], [ 1, %bb2 ], [ 1, %bb3 ]
+  ret void
+}
+; CHECK-LABEL: @func1(
+; CHECK-NEXT:  bb1:
+; CHECK-NEXT:    br label [[BB5:%.*]]
+; CHECK:       bb2:
+; CHECK-NEXT:    [[A:%.*]] = add i32 [[TMP0:%.*]], [[TMP1:%.*]]
+; CHECK-NEXT:    call void @outlined_ir_func_0(i32 [[TMP0]], i32 [[TMP1]], i32 [[A]], i32* null, i32 -1)
+; CHECK-NEXT:    br label [[BB5]]
+; CHECK:       bb4:
+; CHECK-NEXT:    [[E:%.*]] = sub i32 [[TMP0]], [[TMP1]]
+; CHECK-NEXT:    br label [[BB2:%.*]]
+; CHECK:       bb5:
+; CHECK-NEXT:    ret void
+;
+;
+; CHECK-LABEL: @func2(
+; CHECK-NEXT:  bb1:
+; CHECK-NEXT:    [[F_CE_LOC:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    br label [[BB5:%.*]]
+; CHECK:       bb2:
+; CHECK-NEXT:    [[A:%.*]] = sub i32 [[TMP0:%.*]], [[TMP1:%.*]]
+; CHECK-NEXT:    [[LT_CAST:%.*]] = bitcast i32* [[F_CE_LOC]] to i8*
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 -1, i8* [[LT_CAST]])
+; CHECK-NEXT:    call void @outlined_ir_func_0(i32 [[TMP0]], i32 [[TMP1]], i32 1, i32* [[F_CE_LOC]], i32 0)
+; CHECK-NEXT:    [[F_CE_RELOAD:%.*]] = load i32, i32* [[F_CE_LOC]], align 4
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 -1, i8* [[LT_CAST]])
+; CHECK-NEXT:    br label [[BB5]]
+; CHECK:       bb4:
+; CHECK-NEXT:    [[E:%.*]] = add i32 [[TMP0]], [[TMP1]]
+; CHECK-NEXT:    br label [[BB2:%.*]]
+; CHECK:       bb5:
+; CHECK-NEXT:    [[F:%.*]] = phi i32 [ 0, [[BB1:%.*]] ], [ [[F_CE_RELOAD]], [[BB2]] ]
+; CHECK-NEXT:    ret void
+;
+;
+; CHECK-LABEL: define internal void @outlined_ir_func_0(
+; CHECK-NEXT:  newFuncRoot:
+; CHECK-NEXT:    br label [[BB2_TO_OUTLINE:%.*]]
+; CHECK:       bb2_to_outline:
+; CHECK-NEXT:    [[B:%.*]] = add i32 [[TMP0:%.*]], [[TMP1:%.*]]
+; CHECK-NEXT:    [[C:%.*]] = icmp eq i32 [[B]], [[TMP2:%.*]]
+; CHECK-NEXT:    br i1 [[C]], label [[PHI_BLOCK:%.*]], label [[BB3:%.*]]
+; CHECK:       bb3:
+; CHECK-NEXT:    [[D:%.*]] = add i32 [[TMP0]], [[TMP1]]
+; CHECK-NEXT:    br label [[PHI_BLOCK]]
+; CHECK:       bb5.exitStub:
+; CHECK-NEXT:    switch i32 [[TMP4:%.*]], label [[FINAL_BLOCK_0:%.*]] [
+; CHECK-NEXT:    i32 0, label [[OUTPUT_BLOCK_1_0:%.*]]
+; CHECK-NEXT:    ]
+; CHECK:       output_block_1_0:
+; CHECK-NEXT:    store i32 [[TMP5:%.*]], i32* [[TMP3:%.*]], align 4
+; CHECK-NEXT:    br label [[FINAL_BLOCK_0]]
+; CHECK:       phi_block:
+; CHECK-NEXT:    [[TMP5]] = phi i32 [ [[TMP2]], [[BB2_TO_OUTLINE]] ], [ [[TMP2]], [[BB3]] ]
+; CHECK-NEXT:    br label [[BB5_EXITSTUB:%.*]]
+; CHECK:       final_block_0:
+; CHECK-NEXT:    ret void
+;