[llvm] [Attributor] Teach HeapToStack about conservative GC allocators (PR #113299)

via llvm-commits llvm-commits at lists.llvm.org
Tue Oct 22 04:32:32 PDT 2024


https://github.com/aykevl created https://github.com/llvm/llvm-project/pull/113299

When a conservative GC allocates memory, and this memory is flagged as never explicitly freed, the memory can be stack allocated even if a pointer is passed to functions that don't have 'nofree' set.

This is very useful for a conservative GC where memory is known to never be explicitly freed.

---

Context: I plan on using this optimization in TinyGo, so that we don't need a custom heap-to-stack transformation pass but can instead rely on the Attributor pass instead (which will likely do a better job anyway).

I am not entirely sure about the "nofree" `allockind`. This means that memory returned by the allocator will _never_ be explicitly freed. Instead, it might be better to use the `nofree` parameter attribute on the return value (which is currently not allowed). This way, the attribute can be added to a call site and each call site can opt into the "no explicit free" behavior if the compiler knows this particular allocation is never explicitly freed. What do you think?

>From 624e186ff8be9f06a5ef264f0ecf9941976c22c4 Mon Sep 17 00:00:00 2001
From: Ayke van Laethem <aykevanlaethem at gmail.com>
Date: Tue, 22 Oct 2024 13:16:24 +0200
Subject: [PATCH] [Attributor] Teach HeapToStack about conservative GC
 allocators

When a conservative GC allocates memory, and this memory is flagged as
never explicitly freed, the memory can be stack allocated even if a
pointer is passed to functions that don't have 'nofree' set.

This is very useful for a conservative GC where memory is known to be
never explicitly freed.
---
 llvm/docs/LangRef.rst                         |  2 +
 llvm/include/llvm/Analysis/MemoryBuiltins.h   |  5 ++
 llvm/include/llvm/IR/Attributes.h             |  2 +
 llvm/lib/Analysis/MemoryBuiltins.cpp          |  4 ++
 llvm/lib/AsmParser/LLParser.cpp               |  2 +
 llvm/lib/IR/Attributes.cpp                    |  2 +
 llvm/lib/IR/Verifier.cpp                      |  3 +-
 .../Transforms/IPO/AttributorAttributes.cpp   | 13 +++--
 .../Transforms/Attributor/heap_to_stack.ll    | 52 +++++++++++++++----
 9 files changed, 71 insertions(+), 14 deletions(-)

diff --git a/llvm/docs/LangRef.rst b/llvm/docs/LangRef.rst
index b83675c6ed97aa..7edeb090bce959 100644
--- a/llvm/docs/LangRef.rst
+++ b/llvm/docs/LangRef.rst
@@ -1887,6 +1887,8 @@ example:
       zeroed.
     * "aligned": the function returns memory aligned according to the
       ``allocalign`` parameter.
+    * "nofree": the block of memory is not explicitly freed (but might be freed
+      by a conservative GC when unreferenced).
 
     The first three options are mutually exclusive, and the remaining options
     describe more details of how the function behaves. The remaining options
diff --git a/llvm/include/llvm/Analysis/MemoryBuiltins.h b/llvm/include/llvm/Analysis/MemoryBuiltins.h
index 7b48844cc9e8e9..c03e5cadecb8ff 100644
--- a/llvm/include/llvm/Analysis/MemoryBuiltins.h
+++ b/llvm/include/llvm/Analysis/MemoryBuiltins.h
@@ -100,6 +100,11 @@ Value *getFreedOperand(const CallBase *CB, const TargetLibraryInfo *TLI);
 /// insertion or speculative execution of allocation routines.
 bool isRemovableAlloc(const CallBase *V, const TargetLibraryInfo *TLI);
 
+// Whether this is a function that allocates memory that will never be
+// explicitly freed. The memory might be freed in the background by a GC when
+// unreferenced.
+bool isNoFreeAllocFunction(const CallBase *CB);
+
 /// Gets the alignment argument for an aligned_alloc-like function, using either
 /// built-in knowledge based on fuction names/signatures or allocalign
 /// attributes. Note: the Value returned may not indicate a valid alignment, per
diff --git a/llvm/include/llvm/IR/Attributes.h b/llvm/include/llvm/IR/Attributes.h
index 2755ced404dddb..c7eb8189c13f04 100644
--- a/llvm/include/llvm/IR/Attributes.h
+++ b/llvm/include/llvm/IR/Attributes.h
@@ -55,6 +55,8 @@ enum class AllocFnKind : uint64_t {
   Zeroed = 1 << 4,        // Allocator function returns zeroed memory
   Aligned = 1 << 5,       // Allocator function aligns allocations per the
                           // `allocalign` argument
+  NoFree = 1 << 6,        // Allocator function returns memory that's never
+                          // freed
   LLVM_MARK_AS_BITMASK_ENUM(/* LargestValue = */ Aligned)
 };
 
diff --git a/llvm/lib/Analysis/MemoryBuiltins.cpp b/llvm/lib/Analysis/MemoryBuiltins.cpp
index dc2dc4c1733b5e..4ec0db03ade21f 100644
--- a/llvm/lib/Analysis/MemoryBuiltins.cpp
+++ b/llvm/lib/Analysis/MemoryBuiltins.cpp
@@ -339,6 +339,10 @@ bool llvm::isRemovableAlloc(const CallBase *CB, const TargetLibraryInfo *TLI) {
   return isAllocLikeFn(CB, TLI);
 }
 
+bool llvm::isNoFreeAllocFunction(const CallBase *CB) {
+  return checkFnAllocKind(CB, AllocFnKind::NoFree);
+}
+
 Value *llvm::getAllocAlignment(const CallBase *V,
                                const TargetLibraryInfo *TLI) {
   const std::optional<AllocFnsTy> FnData = getAllocationData(V, AnyAlloc, TLI);
diff --git a/llvm/lib/AsmParser/LLParser.cpp b/llvm/lib/AsmParser/LLParser.cpp
index 6a2372c9751408..7fcabb64f17965 100644
--- a/llvm/lib/AsmParser/LLParser.cpp
+++ b/llvm/lib/AsmParser/LLParser.cpp
@@ -2470,6 +2470,8 @@ bool LLParser::parseAllocKind(AllocFnKind &Kind) {
       Kind |= AllocFnKind::Zeroed;
     } else if (A == "aligned") {
       Kind |= AllocFnKind::Aligned;
+    } else if (A == "nofree") {
+      Kind |= AllocFnKind::NoFree;
     } else {
       return error(KindLoc, Twine("unknown allockind ") + A);
     }
diff --git a/llvm/lib/IR/Attributes.cpp b/llvm/lib/IR/Attributes.cpp
index e9daa01b899e8f..c63375edc2dd40 100644
--- a/llvm/lib/IR/Attributes.cpp
+++ b/llvm/lib/IR/Attributes.cpp
@@ -600,6 +600,8 @@ std::string Attribute::getAsString(bool InAttrGrp) const {
       parts.push_back("zeroed");
     if ((Kind & AllocFnKind::Aligned) != AllocFnKind::Unknown)
       parts.push_back("aligned");
+    if ((Kind & AllocFnKind::NoFree) != AllocFnKind::Unknown)
+      parts.push_back("nofree");
     return ("allockind(\"" +
             Twine(llvm::join(parts.begin(), parts.end(), ",")) + "\")")
         .str();
diff --git a/llvm/lib/IR/Verifier.cpp b/llvm/lib/IR/Verifier.cpp
index f34fe7594c8602..19666f4cc43bc4 100644
--- a/llvm/lib/IR/Verifier.cpp
+++ b/llvm/lib/IR/Verifier.cpp
@@ -2318,7 +2318,8 @@ void Verifier::verifyFunctionAttrs(FunctionType *FT, AttributeList Attrs,
           "'allockind()' requires exactly one of alloc, realloc, and free");
     if ((Type == AllocFnKind::Free) &&
         ((K & (AllocFnKind::Uninitialized | AllocFnKind::Zeroed |
-               AllocFnKind::Aligned)) != AllocFnKind::Unknown))
+               AllocFnKind::Aligned | AllocFnKind::NoFree)) !=
+         AllocFnKind::Unknown))
       CheckFailed("'allockind(\"free\")' doesn't allow uninitialized, zeroed, "
                   "or aligned modifiers.");
     AllocFnKind ZeroedUninit = AllocFnKind::Uninitialized | AllocFnKind::Zeroed;
diff --git a/llvm/lib/Transforms/IPO/AttributorAttributes.cpp b/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
index 5d17ffa61272ee..3c21e5599a0742 100644
--- a/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
+++ b/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
@@ -7063,10 +7063,15 @@ ChangeStatus AAHeapToStackFunction::updateImpl(Attributor &A) {
         bool IsAssumedNoCapture = AA::hasAssumedIRAttr<Attribute::NoCapture>(
             A, this, CBIRP, DepClassTy::OPTIONAL, IsKnownNoCapture);
 
-        // If a call site argument use is nofree, we are fine.
-        bool IsKnownNoFree;
-        bool IsAssumedNoFree = AA::hasAssumedIRAttr<Attribute::NoFree>(
-            A, this, CBIRP, DepClassTy::OPTIONAL, IsKnownNoFree);
+        // Check for potentially calls only when the allocator returns memory
+        // that may be freed explicitly.
+        bool IsAssumedNoFree = true;
+        if (!isNoFreeAllocFunction(AI.CB)) {
+          // If a call site argument use is nofree, we are fine.
+          bool IsKnownNoFree;
+          IsAssumedNoFree = AA::hasAssumedIRAttr<Attribute::NoFree>(
+              A, this, CBIRP, DepClassTy::OPTIONAL, IsKnownNoFree);
+        }
 
         if (!IsAssumedNoCapture ||
             (AI.LibraryFunctionId != LibFunc___kmpc_alloc_shared &&
diff --git a/llvm/test/Transforms/Attributor/heap_to_stack.ll b/llvm/test/Transforms/Attributor/heap_to_stack.ll
index 33ac066e43d093..d1f387ddc5a1e7 100644
--- a/llvm/test/Transforms/Attributor/heap_to_stack.ll
+++ b/llvm/test/Transforms/Attributor/heap_to_stack.ll
@@ -46,13 +46,13 @@ define void @h2s_value_simplify_interaction(i1 %c, ptr %A) {
 ; CHECK:       f2:
 ; CHECK-NEXT:    [[L:%.*]] = load i8, ptr [[M]], align 16
 ; CHECK-NEXT:    call void @usei8(i8 [[L]])
-; CHECK-NEXT:    call void @no_sync_func(ptr noalias nocapture nofree noundef nonnull align 16 dereferenceable(1) [[M]]) #[[ATTR11:[0-9]+]]
+; CHECK-NEXT:    call void @no_sync_func(ptr noalias nocapture nofree noundef nonnull align 16 dereferenceable(1) [[M]]) #[[ATTR12:[0-9]+]]
 ; CHECK-NEXT:    br label [[J]]
 ; CHECK:       dead:
 ; CHECK-NEXT:    unreachable
 ; CHECK:       j:
 ; CHECK-NEXT:    [[PHI:%.*]] = phi ptr [ [[M]], [[F]] ], [ null, [[F2]] ]
-; CHECK-NEXT:    tail call void @no_sync_func(ptr nocapture nofree noundef align 16 [[PHI]]) #[[ATTR11]]
+; CHECK-NEXT:    tail call void @no_sync_func(ptr nocapture nofree noundef align 16 [[PHI]]) #[[ATTR12]]
 ; CHECK-NEXT:    ret void
 ;
 entry:
@@ -359,7 +359,7 @@ define void @test9() {
 ; CHECK-NEXT:    [[I:%.*]] = tail call noalias ptr @malloc(i64 noundef 4)
 ; CHECK-NEXT:    tail call void @no_sync_func(ptr nocapture nofree [[I]])
 ; CHECK-NEXT:    store i32 10, ptr [[I]], align 4
-; CHECK-NEXT:    tail call void @foo_nounw(ptr nofree nonnull align 4 dereferenceable(4) [[I]]) #[[ATTR11]]
+; CHECK-NEXT:    tail call void @foo_nounw(ptr nofree nonnull align 4 dereferenceable(4) [[I]]) #[[ATTR12]]
 ; CHECK-NEXT:    tail call void @free(ptr nocapture nonnull align 4 dereferenceable(4) [[I]])
 ; CHECK-NEXT:    ret void
 ;
@@ -419,7 +419,7 @@ define void @test11() {
 ; CHECK-LABEL: define {{[^@]+}}@test11() {
 ; CHECK-NEXT:  bb:
 ; CHECK-NEXT:    [[I_H2S:%.*]] = alloca i8, i64 4, align 1
-; CHECK-NEXT:    tail call void @sync_will_return(ptr [[I_H2S]]) #[[ATTR11]]
+; CHECK-NEXT:    tail call void @sync_will_return(ptr [[I_H2S]]) #[[ATTR12]]
 ; CHECK-NEXT:    ret void
 ;
 bb:
@@ -670,7 +670,7 @@ define void @test16c(i8 %v, ptr %P) {
 ; CHECK-NEXT:  bb:
 ; CHECK-NEXT:    [[I_H2S:%.*]] = alloca i8, i64 4, align 1
 ; CHECK-NEXT:    store ptr [[I_H2S]], ptr [[P]], align 8
-; CHECK-NEXT:    tail call void @no_sync_func(ptr nocapture nofree [[I_H2S]]) #[[ATTR11]]
+; CHECK-NEXT:    tail call void @no_sync_func(ptr nocapture nofree [[I_H2S]]) #[[ATTR12]]
 ; CHECK-NEXT:    ret void
 ;
 bb:
@@ -703,7 +703,7 @@ define void @test16e(i8 %v) norecurse {
 ; CHECK-NEXT:  bb:
 ; CHECK-NEXT:    [[I_H2S:%.*]] = alloca i8, i64 4, align 1
 ; CHECK-NEXT:    store ptr [[I_H2S]], ptr @G, align 8
-; CHECK-NEXT:    call void @usei8p(ptr nocapture nofree [[I_H2S]]) #[[ATTR12:[0-9]+]]
+; CHECK-NEXT:    call void @usei8p(ptr nocapture nofree [[I_H2S]]) #[[ATTR13:[0-9]+]]
 ; CHECK-NEXT:    ret void
 ;
 bb:
@@ -715,6 +715,39 @@ bb:
   ret void
 }
 
+declare noalias ptr @gc_alloc(i64) allockind("alloc,zeroed,nofree") allocsize(0)
+
+; Check that a heap allocated object that is not captured and not explicitly
+; freed can be converted to a stack object. This is often true for GCs.
+define void @test_alloc_nofree_nocapture() {
+; CHECK-LABEL: define {{[^@]+}}@test_alloc_nofree_nocapture() {
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    [[I_H2S:%.*]] = alloca i8, i64 4, align 1
+; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr [[I_H2S]], i8 0, i64 4, i1 false)
+; CHECK-NEXT:    tail call void @nocapture_func_frees_pointer(ptr noalias nocapture [[I_H2S]])
+; CHECK-NEXT:    ret void
+;
+bb:
+  %i = tail call noalias ptr @gc_alloc(i64 4)
+  tail call void @nocapture_func_frees_pointer(ptr %i)
+  ret void
+}
+
+; Check that a nofree heap allocated object that is captured is not converted to
+; a stack object.
+define void @test_alloc_nofree_captured() {
+; CHECK-LABEL: define {{[^@]+}}@test_alloc_nofree_captured() {
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    [[I:%.*]] = tail call noalias ptr @gc_alloc(i64 noundef 4)
+; CHECK-NEXT:    tail call void @foo(ptr [[I]])
+; CHECK-NEXT:    ret void
+;
+bb:
+  %i = tail call noalias ptr @gc_alloc(i64 4)
+  tail call void @foo(ptr %i)
+  ret void
+}
+
 ;.
 ; CHECK: attributes #[[ATTR0:[0-9]+]] = { allockind("alloc,uninitialized") allocsize(0) }
 ; CHECK: attributes #[[ATTR1:[0-9]+]] = { nounwind willreturn }
@@ -726,9 +759,10 @@ bb:
 ; CHECK: attributes #[[ATTR7:[0-9]+]] = { allockind("alloc,uninitialized,aligned") allocsize(1) }
 ; CHECK: attributes #[[ATTR8:[0-9]+]] = { allockind("alloc,zeroed") allocsize(0,1) }
 ; CHECK: attributes #[[ATTR9]] = { norecurse }
-; CHECK: attributes #[[ATTR10:[0-9]+]] = { nocallback nofree nounwind willreturn memory(argmem: write) }
-; CHECK: attributes #[[ATTR11]] = { nounwind }
-; CHECK: attributes #[[ATTR12]] = { nocallback nosync nounwind willreturn }
+; CHECK: attributes #[[ATTR10:[0-9]+]] = { allockind("alloc,zeroed,nofree") allocsize(0) }
+; CHECK: attributes #[[ATTR11:[0-9]+]] = { nocallback nofree nounwind willreturn memory(argmem: write) }
+; CHECK: attributes #[[ATTR12]] = { nounwind }
+; CHECK: attributes #[[ATTR13]] = { nocallback nosync nounwind willreturn }
 ;.
 ;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
 ; CGSCC: {{.*}}



More information about the llvm-commits mailing list