[Attributor] Pack out arguments into a struct #119267

elhewaty · 2024-12-09T20:41:54Z

[Attributor] Add pre-commit tests
[Attributor] Pack out arguments into a struct

PR Summary:
Transforms functions with write-only pointer arguments into functions that return a struct containing the original return value and pointer-written values, improving GPU performance, and maybe x86 too.

Example:
we need to convert this:

__device__ bool foo(int *dst) {
    *dst = threadIdx.x;
    return true;
}

into

struct Out { bool ret; int dst; };

__device__ Out foo() {
    Out result;
    result.dst = threadIdx.x;
    result.ret = true;
    return result;
}

llvmbot · 2024-12-09T20:42:29Z

@llvm/pr-subscribers-llvm-transforms

Author: None (elhewaty)

Changes

[Attributor] Add pre-commit tests
[Attributor] Pack out arguments into a struct

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

2 Files Affected:

(modified) llvm/lib/Transforms/IPO/Attributor.cpp (+113)
(added) llvm/test/Transforms/Attributor/remove_out_args.ll (+16)

diff --git a/llvm/lib/Transforms/IPO/Attributor.cpp b/llvm/lib/Transforms/IPO/Attributor.cpp
index 116f419129a239..90eecf45892ee6 100644
--- a/llvm/lib/Transforms/IPO/Attributor.cpp
+++ b/llvm/lib/Transforms/IPO/Attributor.cpp
@@ -2963,6 +2963,119 @@ bool Attributor::shouldSeedAttribute(AbstractAttribute &AA) {
   return Result;
 }
 
+// For now: argument can be put in the struct if it's write only and
+// has no aliases.
+static bool canBeComapctedInAStruct(const Argument &Arg, Attributor &A, const AbstractAttribute &QueryingAA) {
+  IRPosition ArgPosition = IRPosition::argument(Arg);
+  // Check if Arg has no alias.
+  auto *AAliasInfo = A.getAAFor<AANoAlias>(QueryingAA, ArgPosition, DepClassTy::NONE);
+  if (!AAliasInfo || !AAliasInfo->isKnownNoAlias())
+    return false;
+
+  // Check if Arg is write-only.
+  const auto *MemBehaviorAA =
+      A.getAAFor<AAMemoryBehavior>(QueryingAA, ArgPosition, DepClassTy::NONE);
+  if (!MemBehaviorAA || !MemBehaviorAA->isKnownWriteOnly())
+    return false;
+
+  return true;
+}
+
+static void replaceArgRetWithStructRetCalls(Function &OldFunction, Function &NewFunction) {
+  for (auto UseItr = OldFunction.use_begin(); UseItr != OldFunction.use_end(); ++UseItr) {
+    CallBase *Call = dyn_cast<CallBase>(UseItr->getUser());
+    if (!Call)
+      continue;
+
+    IRBuilder<> Builder(Call);
+    SmallVector<Value *, 8> NewArgs;
+    for (unsigned ArgIdx = 0; ArgIdx < Call->arg_size(); ++ArgIdx)
+      if (std::find_if(OldFunction.arg_begin(), OldFunction.arg_end(),
+          [&](Argument &Arg) { return &Arg == Call->getArgOperand(ArgIdx); }) == OldFunction.arg_end())
+        NewArgs.push_back(Call->getArgOperand(ArgIdx));
+
+    CallInst *NewCall = Builder.CreateCall(&NewFunction, NewArgs);
+    Call->replaceAllUsesWith(NewCall);
+    Call->eraseFromParent();
+  }
+}
+
+static bool convertOutArgsToRetStruct(Function &F, Attributor &A, AbstractAttribute &QueryingAA) {
+  // Get valid ptr args.
+  DenseMap<Argument *, Type *>  PtrToType;
+  for (unsigned ArgIdx = 0; ArgIdx < F.arg_size(); ++ArgIdx) {
+    Argument *Arg = F.getArg(ArgIdx);
+    if (Arg->getType()->isPointerTy() && canBeComapctedInAStruct(*Arg, A, QueryingAA)) {
+      // Get the the type of the pointer through its users
+      for (auto UseItr = Arg->use_begin(); UseItr != Arg->use_end(); ++UseItr) {
+        auto *Store = dyn_cast<StoreInst>(UseItr->getUser());
+        if (Store)
+          PtrToType[Arg] = Store->getValueOperand()->getType();
+      }
+    }
+  }
+
+  // If there is no valid candidates then return false.
+  if (PtrToType.empty())
+    return false;
+
+  // Create the new struct return type.
+  SmallVector<Type *, 4> OutStructElements;
+  if (auto *OriginalFuncTy = F.getReturnType(); !OriginalFuncTy->isVoidTy())
+    OutStructElements.push_back(OriginalFuncTy);
+
+  for (const auto &[Arg, Type] : PtrToType)
+    OutStructElements.push_back(Type);
+
+  auto *ReturnStructType = StructType::create(F.getContext(), OutStructElements, (F.getName() + "Out").str());
+
+  // Get the new Args.
+  SmallVector<Type *, 4> NewParamTypes;
+  for (unsigned ArgIdx = 0; ArgIdx < F.arg_size(); ++ArgIdx)
+    if (!PtrToType.count(F.getArg(ArgIdx)))
+      NewParamTypes.push_back(F.getArg(ArgIdx)->getType());
+
+  auto *NewFunctionType = FunctionType::get(ReturnStructType, NewParamTypes, F.isVarArg());
+  auto *NewFunction = Function::Create(NewFunctionType, F.getLinkage(), F.getAddressSpace(), F.getName());
+
+  // Map old args to new args.
+  ValueToValueMapTy VMap;
+  auto *NewArgIt = NewFunction->arg_begin();
+  for (Argument &OldArg : F.args())
+    if (!PtrToType.count(F.getArg(OldArg.getArgNo())))
+      VMap[&OldArg] = &(*NewArgIt++);
+
+  // Clone the old function into the new one.
+  SmallVector<ReturnInst *, 8> Returns;
+  CloneFunctionInto(NewFunction, &F, VMap, CloneFunctionChangeType::LocalChangesOnly, Returns);
+
+  // Update the return values (make it struct).
+  for (ReturnInst *Ret : Returns) {
+    IRBuilder<> Builder(Ret);
+    SmallVector<Value *, 4> StructValues;
+    // Include original return type, if any
+    if (auto *OriginalFuncTy = F.getReturnType(); !OriginalFuncTy->isVoidTy())
+      StructValues.push_back(Ret->getReturnValue());
+
+    // Create a load instruction to fill the struct element.
+    for (const auto &[Arg, Ty] : PtrToType) {
+      Value *OutVal = Builder.CreateLoad(Ty, VMap[Arg]);
+      StructValues.push_back(OutVal);
+    }
+
+    // Build the return struct incrementally.
+    Value *StructRetVal = UndefValue::get(ReturnStructType);
+    for (unsigned i = 0; i < StructValues.size(); ++i)
+      StructRetVal = Builder.CreateInsertValue(StructRetVal, StructValues[i], i);
+
+    Builder.CreateRet(StructRetVal);
+    Ret->eraseFromParent();
+  }
+
+  replaceArgRetWithStructRetCalls(F, *NewFunction);
+  F.eraseFromParent();
+}
+
 ChangeStatus Attributor::rewriteFunctionSignatures(
     SmallSetVector<Function *, 8> &ModifiedFns) {
   ChangeStatus Changed = ChangeStatus::UNCHANGED;
diff --git a/llvm/test/Transforms/Attributor/remove_out_args.ll b/llvm/test/Transforms/Attributor/remove_out_args.ll
new file mode 100644
index 00000000000000..bd52bf5d80656c
--- /dev/null
+++ b/llvm/test/Transforms/Attributor/remove_out_args.ll
@@ -0,0 +1,16 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes=attributor < %s | FileCheck %s
+
+
+
+define i1 @foo(ptr %dst) {
+; CHECK-LABEL: define noundef i1 @foo(
+; CHECK-SAME: ptr nocapture nofree noundef nonnull writeonly align 4 dereferenceable(4) [[DST:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    store i32 42, ptr [[DST]], align 4
+; CHECK-NEXT:    ret i1 true
+;
+entry:
+  store i32 42, ptr %dst
+  ret i1 1
+}

github-actions · 2024-12-09T20:45:29Z

✅ With the latest revision this PR passed the C/C++ code formatter.

elhewaty · 2024-12-09T20:49:14Z

I know it's not complete. but I wonder if I am on the right track as I am confused.
another thing is where to call convertOutArgsToRetStruct 😖?

llvm/lib/Transforms/IPO/Attributor.cpp

llvm/test/Transforms/Attributor/remove_out_args.ll

llvm/lib/Transforms/IPO/Attributor.cpp

arsenm · 2024-12-11T02:47:47Z

llvm/test/Transforms/Attributor/remove_out_args.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes=attributor < %s | FileCheck %s
+
+


This needs about 100 more tests llvm/test/CodeGen/AMDGPU/rewrite-out-arguments.ll and llvm/test/CodeGen/AMDGPU/rewrite-out-arguments-address-space.ll would be better starting points

I will. BTW, I don't have an AMD GPU, I have a nvidia one, should J add the tests in NVPTX?

It doesn't matter what GPU you have, it doesn't matter when working on the compiler. NVPTX will be less useful for IR testing since it doesn't use a non-0 address space for stack objects

Still need all of these tests, using AMDGPU

@arsenm, I will test using AMDGPU, but I don't know why my changes don't change this simple test? Do I need to call the some function somewhere? the optimizer dosn't see my changes or do my changes don't affect the input?

This is an externally visible function, so it's not allowed to change the signature. It would need to be be internal or similar linkage, with a call use

llvm/lib/Transforms/IPO/Attributor.cpp

github-actions · 2024-12-21T21:44:25Z

⚠️ undef deprecator found issues in your code. ⚠️

You can test this locally with the following command:

git diff -U0 --pickaxe-regex -S '([^a-zA-Z0-9#_-]undef[^a-zA-Z0-9_-]|UndefValue::get)' e6ced4da4499007a366ada31cfb07e0b4fbf2393 7e34621d2aca606ec1836ef761a1cca590249e45 llvm/test/Transforms/Attributor/remove_out_args.ll llvm/include/llvm/Transforms/IPO/Attributor.h llvm/lib/Transforms/IPO/Attributor.cpp llvm/lib/Transforms/IPO/AttributorAttributes.cpp

The following files introduce new uses of undef:

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

Undef is now deprecated and should only be used in the rare cases where no replacement is possible. For example, a load of uninitialized memory yields undef. You should use poison values for placeholders instead.

In tests, avoid using undef and having tests that trigger undefined behavior. If you need an operand with some unimportant value, you can add a new argument to the function and use that instead.

For example, this is considered a bad practice:

define void @fn() {
  ...
  br i1 undef, ...
}

Please use the following instead:

define void @fn(i1 %cond) {
  ...
  br i1 %cond, ...
}

Please refer to the Undefined Behavior Manual for more information.

elhewaty · 2024-12-21T21:48:50Z

@arsenm @shiltian @vidsinghal @jdoerfert @jhuber6
I am sorry but this is my first time with the attributor. I keep getting compilation errors like:

error: invalid new-expression of abstract class type ‘{anonymous}::AAConvertOutArgumentFunction’
13207 |     AA = new (A.Allocator) CLASS##SUFFIX(IRP, A);

But I need you to share your opinions so far, I will write tests to cover all possible cases, but at least I need this to work with the simple test function written before.

What did I mess up while adding the attributor?
I think I need to change the whole thing, but this is my best for now!

arsenm · 2024-12-23T07:40:06Z

@arsenm @shiltian @vidsinghal @jdoerfert @jhuber6 I am sorry but this is my first time with the attributor. I keep getting compilation errors like:
error: invalid new-expression of abstract class type ‘{anonymous}::AAConvertOutArgumentFunction’
13207 |     AA = new (A.Allocator) CLASS##SUFFIX(IRP, A);

You didn't implement a virtual method of the base class. If you look further along in the error message it should tell you which ones

elhewaty · 2024-12-23T10:10:50Z

@arsenm @shiltian @vidsinghal @jdoerfert @jhuber6 I am sorry but this is my first time with the attributor. I keep getting compilation errors like:
error: invalid new-expression of abstract class type ‘{anonymous}::AAConvertOutArgumentFunction’
13207 |     AA = new (A.Allocator) CLASS##SUFFIX(IRP, A);
You didn't implement a virtual method of the base class. If you look further along in the error message it should tell you which ones

@arsenm, Thanks.
Would you take a look? how to use AAConvertOutArgumentFunction It's not changing the sample function in the tests?

elhewaty · 2024-12-24T12:12:44Z

@arsenm, Thanks. Would you take a look? how to use AAConvertOutArgumentFunction It's not changing the sample function in the tests?

@shiltian can you take a look?

arsenm · 2024-12-26T12:55:31Z

llvm/include/llvm/Transforms/IPO/Attributor.h

+    const IRPosition &ArgPos = IRPosition::argument(Arg);
+    auto *AAMem = A.getAAFor<AAMemoryBehavior>(AA, ArgPos, DepClassTy::NONE);
+
+    return Arg.hasNoAliasAttr() && AAMem && AAMem->isKnownWriteOnly() &&


Should this be querying from AANoAlias instead of directly checking the attribute is present on the argument?

@arsenm, in the following example does ptr %dts have any aliases in @foo? as AANoAlias->isKnownNoAlias() returns false.

define internal i1 @foo(ptr %dst) { entry: store i32 42, ptr %dst ret i1 true } define i1 @fee(i32 %x, i32 %y) { %ptr = alloca i32 %a = call i1 @foo(ptr %ptr, i32 %y) %b = load i32, ptr %ptr %c = icmp sle i32 %b, %x %xor = xor i1 %a, %c ret i1 %xor }

This callsite has the wrong type, so this test is UB and it probably doesn't bother trying to optimize it correctly

llvm/include/llvm/Transforms/IPO/Attributor.h

arsenm · 2024-12-26T12:56:23Z

llvm/lib/Transforms/IPO/Attributor.cpp

@@ -3412,6 +3412,9 @@ void Attributor::identifyDefaultAbstractAttributes(Function &F) {
    // Every function can track active assumptions.
    getOrCreateAAFor<AAAssumptionInfo>(FPos);

+    // Every function can have out arguments.


Only functions with pointer arguments, not sure if it's worth checking if there are any

arsenm · 2024-12-26T12:56:47Z