[IRSim][IROutliner] Ignoring Musttail Function

Musttail calls require extra handling to properly propagate the calling convention information and tail call information. The outliner does not currently do this, so we ignore call instructions that utilize the swifttailcc and tailcc calling convention as well as functions marked with the attribute musttail.

Reviewers: paquette, aschwaighofer

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

llvm/include/llvm/Analysis/IRSimilarityIdentifier.h

@@ -560,6 +560,18 @@ struct IRInstructionMapper {
         return Illegal;
       if (!F && !IsIndirectCall)
         return Illegal;
+      // Functions marked with the swifttailcc and tailcc calling conventions
+      // require special handling when outlining musttail functions.  The
+      // calling convention must be passed down to the outlined function as
+      // well. Further, there is special handling for musttail calls as well,
+      // requiring a return call directly after.  For now, the outliner does not
+      // support this, so we do not handle matching this case either.
+      if ((CI.getCallingConv() == CallingConv::SwiftTail ||
+           CI.getCallingConv() == CallingConv::Tail) &&
+          !EnableMustTailCalls)
+        return Illegal;
+      if (CI.isMustTailCall() && !EnableMustTailCalls)
+        return Illegal;
       return Legal;
     }
     // TODO: We do not current handle similarity that changes the control flow.
@@ -581,6 +593,10 @@ struct IRInstructionMapper {
     // Flag that lets the classifier know whether we should allow intrinsics to
     // be checked for similarity.
     bool EnableIntrinsics = false;
+  
+    // Flag that lets the classifier know whether we should allow tail calls to
+    // be checked for similarity.
+    bool EnableMustTailCalls = false;
   };
 
   /// Maps an Instruction to a member of InstrType.
@@ -968,11 +984,13 @@ public:
   IRSimilarityIdentifier(bool MatchBranches = true,
                          bool MatchIndirectCalls = true,
                          bool MatchCallsWithName = false,
-                         bool MatchIntrinsics = true)
+                         bool MatchIntrinsics = true,
+                         bool MatchMustTailCalls = true)
       : Mapper(&InstDataAllocator, &InstDataListAllocator),
         EnableBranches(MatchBranches), EnableIndirectCalls(MatchIndirectCalls),
         EnableMatchingCallsByName(MatchCallsWithName),
-        EnableIntrinsics(MatchIntrinsics) {}
+        EnableIntrinsics(MatchIntrinsics),
+        EnableMustTailCalls(MatchMustTailCalls) {}
 
 private:
   /// Map the instructions in the module to unsigned integers, using mapping
@@ -1065,6 +1083,10 @@ private:
   /// similarity.
   bool EnableIntrinsics = true;
 
+  // The flag variable that marks whether we should allow tailcalls
+  // to be checked for similarity.
+  bool EnableMustTailCalls = false;
+
   /// The SimilarityGroups found with the most recent run of \ref
   /// findSimilarity. None if there is no recent run.
   Optional<SimilarityGroupList> SimilarityCandidates;

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

@@ -51,6 +51,7 @@
 struct OutlinableGroup;
 
 namespace llvm {
+using namespace CallingConv;
 using namespace IRSimilarity;
 
 class Module;
@@ -372,6 +373,25 @@ private:
       // the call in outlined functions.
       if (CI.canReturnTwice())
         return false;
+      // TODO: Update the outliner to capture whether the outlined function
+      // needs these extra attributes.
+
+      // Functions marked with the swifttailcc and tailcc calling conventions
+      // require special handling when outlining musttail functions.  The
+      // calling convention must be passed down to the outlined function as
+      // well. Further, there is special handling for musttail calls as well,
+      // requiring a return call directly after.  For now, the outliner does not
+      // support this.
+      bool IsTailCC = CI.getCallingConv() == CallingConv::SwiftTail ||
+                      CI.getCallingConv() == CallingConv::Tail;
+      if (IsTailCC && !EnableMustTailCalls)
+        return false;
+      if (CI.isMustTailCall() && !EnableMustTailCalls)
+        return false;
+      // The outliner can only handle musttail items if it is also accompanied
+      // by the tailcc or swifttailcc calling convention.
+      if (CI.isMustTailCall() && !IsTailCC)
+        return false;
       return true;
     }
     // TODO: Handle FreezeInsts.  Since a frozen value could be frozen inside
@@ -397,6 +417,9 @@ private:
     // The flag variable that marks whether we should allow intrinsics
     // instructions to be outlined.
     bool EnableIntrinsics = false;
+
+    // The flag variable that marks whether we should allow musttail calls.
+    bool EnableMustTailCalls = false;
   };
 
   /// A InstVisitor used to exclude certain instructions from being outlined.

llvm/lib/Analysis/IRSimilarityIdentifier.cpp

@@ -1163,6 +1163,7 @@ SimilarityGroupList &IRSimilarityIdentifier::findSimilarity(
   Mapper.InstClassifier.EnableIndirectCalls = EnableIndirectCalls;
   Mapper.EnableMatchCallsByName = EnableMatchingCallsByName;
   Mapper.InstClassifier.EnableIntrinsics = EnableIntrinsics;
+  Mapper.InstClassifier.EnableMustTailCalls = EnableMustTailCalls;
 
   populateMapper(Modules, InstrList, IntegerMapping);
   findCandidates(InstrList, IntegerMapping);
@@ -1176,6 +1177,7 @@ SimilarityGroupList &IRSimilarityIdentifier::findSimilarity(Module &M) {
   Mapper.InstClassifier.EnableIndirectCalls = EnableIndirectCalls;
   Mapper.EnableMatchCallsByName = EnableMatchingCallsByName;
   Mapper.InstClassifier.EnableIntrinsics = EnableIntrinsics;
+  Mapper.InstClassifier.EnableMustTailCalls = EnableMustTailCalls;
 
   std::vector<IRInstructionData *> InstrList;
   std::vector<unsigned> IntegerMapping;
@@ -1197,7 +1199,8 @@ IRSimilarityIdentifierWrapperPass::IRSimilarityIdentifierWrapperPass()
 
 bool IRSimilarityIdentifierWrapperPass::doInitialization(Module &M) {
   IRSI.reset(new IRSimilarityIdentifier(!DisableBranches, !DisableIndirectCalls,
-                                        MatchCallsByName, !DisableIntrinsics));
+                                        MatchCallsByName, !DisableIntrinsics,
+                                        false));
   return false;
 }
 
@@ -1215,7 +1218,8 @@ AnalysisKey IRSimilarityAnalysis::Key;
 IRSimilarityIdentifier IRSimilarityAnalysis::run(Module &M,
                                                  ModuleAnalysisManager &) {
   auto IRSI = IRSimilarityIdentifier(!DisableBranches, !DisableIndirectCalls,
-                                     MatchCallsByName, !DisableIntrinsics);
+                                     MatchCallsByName, !DisableIntrinsics,
+                                     false);
   IRSI.findSimilarity(M);
   return IRSI;
 }

llvm/test/Transforms/IROutliner/outlining-musttail.ll

@@ -0,0 +1,34 @@
+; 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
+
+; Check that we do not outline musttail when swifttaill cc or tailcc
+; is not present.
+
+declare void @musttail()
+
+define void @f1() {
+  %a = alloca i32, align 4
+  store i32 2, i32* %a, align 4
+  musttail call void @musttail()
+  ret void
+}
+
+define void @f2() {
+  %a = alloca i32, align 4
+  store i32 2, i32* %a, align 4
+  musttail call void @musttail()
+  ret void
+}
+; CHECK-LABEL: @f1(
+; CHECK-NEXT:    [[A:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    store i32 2, i32* [[A]], align 4
+; CHECK-NEXT:    musttail call void @musttail()
+; CHECK-NEXT:    ret void
+;
+;
+; CHECK-LABEL: @f2(
+; CHECK-NEXT:    [[A:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    store i32 2, i32* [[A]], align 4
+; CHECK-NEXT:    musttail call void @musttail()
+; CHECK-NEXT:    ret void
+;

llvm/test/Transforms/IROutliner/outlining-swifttailcc.ll

@@ -0,0 +1,33 @@
+; 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
+
+; Check that we do not outline musttail calls when swifttailcc is present.
+
+declare swifttailcc void @musttail()
+
+define swifttailcc void @f1() {
+  %a = alloca i32, align 4
+  store i32 2, i32* %a, align 4
+  musttail call swifttailcc void @musttail()
+  ret void
+}
+
+define swifttailcc void @f2() {
+  %a = alloca i32, align 4
+  store i32 2, i32* %a, align 4
+  musttail call swifttailcc void @musttail()
+  ret void
+}
+; CHECK-LABEL: @f1(
+; CHECK-NEXT:    [[A:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    store i32 2, i32* [[A]], align 4
+; CHECK-NEXT:    musttail call swifttailcc void @musttail()
+; CHECK-NEXT:    ret void
+;
+;
+; CHECK-LABEL: @f2(
+; CHECK-NEXT:    [[A:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    store i32 2, i32* [[A]], align 4
+; CHECK-NEXT:    musttail call swifttailcc void @musttail()
+; CHECK-NEXT:    ret void
+;

llvm/test/Transforms/IROutliner/outlining-tailcc.ll

@@ -0,0 +1,33 @@
+; 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
+
+; Check that we not do outline musttail calls when tailcc is present.
+
+declare tailcc void @musttail()
+
+define tailcc void @f1() {
+  %a = alloca i32, align 4
+  store i32 2, i32* %a, align 4
+  musttail call tailcc void @musttail()
+  ret void
+}
+
+define tailcc void @f2() {
+  %a = alloca i32, align 4
+  store i32 2, i32* %a, align 4
+  musttail call tailcc void @musttail()
+  ret void
+}
+; CHECK-LABEL: @f1(
+; CHECK-NEXT:    [[A:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    store i32 2, i32* [[A]], align 4
+; CHECK-NEXT:    musttail call tailcc void @musttail()
+; CHECK-NEXT:    ret void
+;
+;
+; CHECK-LABEL: @f2(
+; CHECK-NEXT:    [[A:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    store i32 2, i32* [[A]], align 4
+; CHECK-NEXT:    musttail call tailcc void @musttail()
+; CHECK-NEXT:    ret void
+;