[llvm] [IROutliner] do not outline non-uniform immarg constants (PR #195450)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Arda Serdar Pektezol (pektezol)

<details>
<summary>Changes</summary>

Fixes #<!-- -->194733 

`IROutliner` can outline structurally similar regions by lifting non-uniform constants into arguments of the outlined function. This is invalid for operands passed to callee parameters marked `immarg`, because those operands must remain immediate constants at the call site.

A fuzzer found a case where `llvm.memcpy` calls differed only in the `i1 immarg` operand. The outliner replaced that constant with an outlined-function argument, producing invalid IR: `call void @<!-- -->llvm.memcpy.p0.p0.i64(ptr null, ptr null, i64 0, i1 %0)`

Fix this by detecting when a non-uniform constant operand corresponds to an `immarg` call parameter and rejecting the outlining group. Identical `immarg` constants are still outlineable, because they remain constants in the outlined function. Added a regression test covering this case, and also tested the existing `outline-memcpy.ll` regression to make sure the fix doesn't break existing behavior.

---
Full diff: https://github.com/llvm/llvm-project/pull/195450.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/IPO/IROutliner.cpp (+43) 
- (added) llvm/test/Transforms/IROutliner/illegal-different-immarg.ll (+28) 


``````````diff
diff --git a/llvm/lib/Transforms/IPO/IROutliner.cpp b/llvm/lib/Transforms/IPO/IROutliner.cpp
index c1640f3d0e2a7..157fc51a9ab2d 100644
--- a/llvm/lib/Transforms/IPO/IROutliner.cpp
+++ b/llvm/lib/Transforms/IPO/IROutliner.cpp
@@ -579,11 +579,54 @@ collectRegionsConstants(OutlinableRegion &Region,
   return ConstantsTheSame;
 }
 
+/// Check whether \p Region contains a non-uniform constant operand that must
+/// remain an immediate argument.
+///
+/// \param Region - The region to check for non-uniform immarg constants.
+/// \param NotSame - The set of global value numbers that do not have the same
+/// constant in each region.
+/// \returns true if \p Region contains a non-uniform immarg constant, and false
+/// otherwise.
+static bool containsNonUniformImmArgConstant(OutlinableRegion &Region,
+                                             DenseSet<unsigned> &NotSame) {
+  IRSimilarityCandidate &C = *Region.Candidate;
+
+  for (IRInstructionData &ID : C) {
+    auto *CB = dyn_cast<CallBase>(ID.Inst);
+    if (!CB)
+      continue;
+
+    for (unsigned ArgIdx = 0, ArgEnd = CB->arg_size(); ArgIdx != ArgEnd;
+         ++ArgIdx) {
+      if (!CB->paramHasAttr(ArgIdx, Attribute::ImmArg))
+        continue;
+
+      Value *Arg = CB->getArgOperand(ArgIdx);
+      if (!isa<Constant>(Arg))
+        continue;
+
+      std::optional<unsigned> GVN = C.getGVN(Arg);
+      assert(GVN && "Expected a GVN for immarg operand?");
+      if (NotSame.contains(*GVN))
+        return true;
+    }
+  }
+
+  return false;
+}
+
 void OutlinableGroup::findSameConstants(DenseSet<unsigned> &NotSame) {
   DenseMap<unsigned, Constant *> GVNToConstant;
 
   for (OutlinableRegion *Region : Regions)
     collectRegionsConstants(*Region, GVNToConstant, NotSame);
+
+  for (OutlinableRegion *Region : Regions) {
+    if (!containsNonUniformImmArgConstant(*Region, NotSame))
+      continue;
+    IgnoreGroup = true;
+    return;
+  }
 }
 
 void OutlinableGroup::collectGVNStoreSets(Module &M) {
diff --git a/llvm/test/Transforms/IROutliner/illegal-different-immarg.ll b/llvm/test/Transforms/IROutliner/illegal-different-immarg.ll
new file mode 100644
index 0000000000000..04f1ef7165869
--- /dev/null
+++ b/llvm/test/Transforms/IROutliner/illegal-different-immarg.ll
@@ -0,0 +1,28 @@
+; RUN: opt -S -passes=verify,iroutliner -ir-outlining-no-cost < %s | FileCheck %s
+
+; Do not outline structurally similar regions when doing so would replace a
+; non-uniform immarg constant with an outlined-function argument.
+
+declare void @llvm.memcpy.p0.p0.i64(ptr noalias writeonly captures(none), ptr noalias readonly captures(none), i64, i1 immarg)
+
+define ptr @PR194733() {
+entry:
+  br i1 false, label %then, label %else
+
+then:
+  %0 = load ptr, ptr null, align 8
+  call void @llvm.memcpy.p0.p0.i64(ptr null, ptr null, i64 0, i1 false)
+  ret ptr null
+
+else:
+  %1 = load ptr, ptr null, align 8
+  call void @llvm.memcpy.p0.p0.i64(ptr null, ptr null, i64 0, i1 true)
+  %2 = load ptr, ptr null, align 8
+  call void @llvm.memcpy.p0.p0.i64(ptr null, ptr null, i64 0, i1 false)
+  ret ptr null
+}
+; CHECK-LABEL: define ptr @PR194733(
+; CHECK-NOT: outlined_ir_func
+; CHECK: call void @llvm.memcpy.p0.p0.i64(ptr null, ptr null, i64 0, i1 false)
+; CHECK: call void @llvm.memcpy.p0.p0.i64(ptr null, ptr null, i64 0, i1 true)
+; CHECK: call void @llvm.memcpy.p0.p0.i64(ptr null, ptr null, i64 0, i1 false)

``````````

</details>


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