[llvm-branch-commits] [llvm] [Inline][Cloning] Drop incompatible attributes from `NewFunc` (PR #91350)

[via llvm-branch-commits]

https://github.com/AtariDreams created https://github.com/llvm/llvm-project/pull/91350

Performing `instSimplify` while cloning is unsafe due to incomplete remapping (as reported in #87534). Ideally, `instSimplify` ought to reason on the updated newly-cloned function, after returns have been rewritten and callee entry basic block / call-site have been fixed up. This is in contrast to `CloneAndPruneIntoFromInst` behaviour, which is inherently expected to clone basic blocks, with pruning on top of – if any –, and not actually fixing up returns / CFG, which should be up to the Inliner. We may solve this by letting `instSimplify` work on the newly-cloned function, while maintaining old function attributes, so as to avoid inconsistencies between the yet-to-be-solved return type, and new function ret type attributes.

(cherry picked from commit 1bb929833b18db4a26a4d145d7270597cb5d48ce)

>From 8d591584cbecb9a2c4a2b9e1ace24e13974ad7f5 Mon Sep 17 00:00:00 2001
From: Antonio Frighetto <me at antoniofrighetto.com>
Date: Mon, 29 Apr 2024 17:53:29 +0200
Subject: [PATCH] [Inline][Cloning] Drop incompatible attributes from `NewFunc`
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Performing `instSimplify` while cloning is unsafe due to incomplete
remapping (as reported in #87534). Ideally, `instSimplify` ought to
reason on the updated newly-cloned function, after returns have been
rewritten and callee entry basic block / call-site have been fixed up.
This is in contrast to `CloneAndPruneIntoFromInst` behaviour, which
is inherently expected to clone basic blocks, with pruning on top of
– if any –, and not actually fixing up returns / CFG, which should be
up to the Inliner. We may solve this by letting `instSimplify` work on
the newly-cloned function, while maintaining old function attributes,
so as to avoid inconsistencies between the yet-to-be-solved return
type, and new function ret type attributes.

(cherry picked from commit 1bb929833b18db4a26a4d145d7270597cb5d48ce)
---
 llvm/lib/Transforms/Utils/CloneFunction.cpp   | 21 ++++++++----
 .../Inline/inline-drop-attributes.ll          | 32 +++++++++++++++++++
 2 files changed, 46 insertions(+), 7 deletions(-)
 create mode 100644 llvm/test/Transforms/Inline/inline-drop-attributes.ll

diff --git a/llvm/lib/Transforms/Utils/CloneFunction.cpp b/llvm/lib/Transforms/Utils/CloneFunction.cpp
index c0f333364fa58..a6821c9df8a92 100644
--- a/llvm/lib/Transforms/Utils/CloneFunction.cpp
+++ b/llvm/lib/Transforms/Utils/CloneFunction.cpp
@@ -17,6 +17,7 @@
 #include "llvm/Analysis/DomTreeUpdater.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/LoopInfo.h"
+#include "llvm/IR/AttributeMask.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DebugInfo.h"
@@ -823,13 +824,16 @@ void llvm::CloneAndPruneIntoFromInst(Function *NewFunc, const Function *OldFunc,
     }
   }
 
-  // Make a second pass over the PHINodes now that all of them have been
-  // remapped into the new function, simplifying the PHINode and performing any
-  // recursive simplifications exposed. This will transparently update the
-  // WeakTrackingVH in the VMap. Notably, we rely on that so that if we coalesce
-  // two PHINodes, the iteration over the old PHIs remains valid, and the
-  // mapping will just map us to the new node (which may not even be a PHI
-  // node).
+  // Drop all incompatible return attributes that cannot be applied to NewFunc
+  // during cloning, so as to allow instruction simplification to reason on the
+  // old state of the function. The original attributes are restored later.
+  AttributeMask IncompatibleAttrs =
+      AttributeFuncs::typeIncompatible(OldFunc->getReturnType());
+  AttributeList Attrs = NewFunc->getAttributes();
+  NewFunc->removeRetAttrs(IncompatibleAttrs);
+
+  // As phi-nodes have been now remapped, allow incremental simplification of
+  // newly-cloned instructions.
   const DataLayout &DL = NewFunc->getParent()->getDataLayout();
   SmallSetVector<const Value *, 8> Worklist;
   for (unsigned Idx = 0, Size = PHIToResolve.size(); Idx != Size; ++Idx)
@@ -871,6 +875,9 @@ void llvm::CloneAndPruneIntoFromInst(Function *NewFunc, const Function *OldFunc,
       VMap[OrigV] = I;
   }
 
+  // Restore attributes.
+  NewFunc->setAttributes(Attrs);
+
   // Remap debug intrinsic operands now that all values have been mapped.
   // Doing this now (late) preserves use-before-defs in debug intrinsics. If
   // we didn't do this, ValueAsMetadata(use-before-def) operands would be
diff --git a/llvm/test/Transforms/Inline/inline-drop-attributes.ll b/llvm/test/Transforms/Inline/inline-drop-attributes.ll
new file mode 100644
index 0000000000000..9a451f4b8699a
--- /dev/null
+++ b/llvm/test/Transforms/Inline/inline-drop-attributes.ll
@@ -0,0 +1,32 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
+; RUN: opt < %s -passes=inline -S | FileCheck %s
+; RUN: opt < %s -passes='cgscc(inline)' -S | FileCheck %s
+
+define void @callee() {
+; CHECK-LABEL: define void @callee() {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[VAL_PTR:%.*]] = load ptr, ptr null, align 8
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq ptr [[VAL_PTR]], null
+; CHECK-NEXT:    [[VAL:%.*]] = load i64, ptr null, align 8
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i64 undef, i64 [[VAL]]
+; CHECK-NEXT:    ret void
+;
+entry:
+  %val_ptr = load ptr, ptr null, align 8
+  %cmp = icmp eq ptr %val_ptr, null
+  %val = load i64, ptr null, align 8
+  %sel = select i1 %cmp, i64 undef, i64 %val
+  ret void
+}
+
+define noundef i1 @caller() {
+; CHECK-LABEL: define noundef i1 @caller() {
+; CHECK-NEXT:    [[VAL_PTR_I:%.*]] = load ptr, ptr null, align 8
+; CHECK-NEXT:    [[CMP_I:%.*]] = icmp eq ptr [[VAL_PTR_I]], null
+; CHECK-NEXT:    [[VAL_I:%.*]] = load i64, ptr null, align 8
+; CHECK-NEXT:    [[SEL_I:%.*]] = select i1 [[CMP_I]], i64 undef, i64 [[VAL_I]]
+; CHECK-NEXT:    ret i1 false
+;
+  call void @callee()
+  ret i1 false
+}