From bbe1e5b593e4fbdb63b9d1c6e8ac6fc31e058844 Mon Sep 17 00:00:00 2001
From: lkdvos <lukas.devos@ugent.be>
Date: Fri, 5 Jul 2024 17:11:32 +0200
Subject: [PATCH] Start allocator implementations

---
 Project.toml                      |  4 ++
 src/TensorOperations.jl           |  1 +
 src/implementation/allocator.jl   | 81 ++++++++++++++++++++++++++++++-
 src/implementation/strided.jl     |  2 +-
 src/indexnotation/tensormacros.jl | 23 +++++++++
 5 files changed, 109 insertions(+), 2 deletions(-)

diff --git a/Project.toml b/Project.toml
index 371c9bc..f900838 100644
--- a/Project.toml
+++ b/Project.toml
@@ -4,11 +4,13 @@ authors = ["Lukas Devos <lukas.devos@ugent.be>", "Maarten Van Damme <maartenvd19
 version = "5.0.0"
 
 [deps]
+Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 LRUCache = "8ac3fa9e-de4c-5943-b1dc-09c6b5f20637"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
+PtrArrays = "43287f4e-b6f4-7ad1-bb20-aadabca52c3d"
 Strided = "5e0ebb24-38b0-5f93-81fe-25c709ecae67"
 StridedViews = "4db3bf67-4bd7-4b4e-b153-31dc3fb37143"
 TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
@@ -26,6 +28,7 @@ TensorOperationscuTENSORExt = ["cuTENSOR", "CUDA"]
 
 [compat]
 Aqua = "0.6, 0.7, 0.8"
+Bumper = "0.6"
 CUDA = "5.4.0"
 ChainRulesCore = "1"
 ChainRulesTestUtils = "1"
@@ -34,6 +37,7 @@ LRUCache = "1"
 LinearAlgebra = "1.6"
 Logging = "1.6"
 PackageExtensionCompat = "1"
+PtrArrays = "1.2"
 Random = "1"
 Strided = "2.0.4"
 StridedViews = "0.3"
diff --git a/src/TensorOperations.jl b/src/TensorOperations.jl
index a15f7bc..1c85c92 100644
--- a/src/TensorOperations.jl
+++ b/src/TensorOperations.jl
@@ -7,6 +7,7 @@ using LinearAlgebra
 using LinearAlgebra: mul!, BlasFloat
 using Strided
 using StridedViews: isstrided
+using PtrArrays, Bumper
 using LRUCache
 
 using Base.Meta: isexpr
diff --git a/src/implementation/allocator.jl b/src/implementation/allocator.jl
index 6d77470..005e164 100644
--- a/src/implementation/allocator.jl
+++ b/src/implementation/allocator.jl
@@ -1,5 +1,5 @@
 # ------------------------------------------------------------------------------------------
-# DefaultAllocator
+# Allocator backends
 # ------------------------------------------------------------------------------------------
 """
     DefaultAllocator()
@@ -30,6 +30,17 @@ parameters `Min`, `Mout`, `Mtemp`` can be any of the CUDA.jl memory types, i.e.
 """
 struct CUDAAllocator{Mout,Min,Mtemp} end
 
+"""
+    ManualAllocator()
+
+Allocator that bypasses Julia's memory management for temporary tensors by leveraging `Libc.malloc`
+and `Libc.free` directly. This can be useful for reducing the pressure on the garbage collector.
+This backend will allocate using `DefaultAllocator` for output tensors that escape the `@tensor`
+block, which will thus still be managed using Julia's GC. The other tensors will be backed by
+`PtrArray` instances, from `PtrArrays.jl`, thus requiring compatibility with that interface.
+"""
+struct ManualAllocator end
+
 # ------------------------------------------------------------------------------------------
 # Generic implementation
 # ------------------------------------------------------------------------------------------
@@ -148,3 +159,71 @@ end
 function tensorfree!(C, allocator=DefaultAllocator())
     return nothing
 end
+
+# ------------------------------------------------------------------------------------------
+# ManualAllocator implementation
+# ------------------------------------------------------------------------------------------
+Base.@constprop :aggressive function tensoralloc_add(TC, A, pA::Index2Tuple, conjA::Bool,
+                                                     istemp::Bool,
+                                                     ::ManualAllocator)
+    structure = tensoradd_structure(A, pA, conjA)
+    if istemp
+        return malloc(TC, structure...)
+    else
+        ttype = tensoradd_type(TC, A, pA, conjA)
+        return tensoralloc(ttype, structure, istemp)::ttype
+    end
+end
+
+Base.@constprop :aggressive function tensoralloc_contract(TC,
+                                                          A, pA::Index2Tuple, conjA::Bool,
+                                                          B, pB::Index2Tuple, conjB::Bool,
+                                                          pAB::Index2Tuple, istemp::Bool,
+                                                          ::ManualAllocator)
+    structure = tensorcontract_structure(A, pA, conjA, B, pB, conjB, pAB)
+    if istemp
+        return malloc(TC, structure...)
+    else
+        ttype = tensorcontract_type(TC, A, pA, conjA, B, pB, conjB, pAB)
+        return tensoralloc(ttype, structure, istemp)::ttype
+    end
+end
+
+function tensorfree!(C, ::ManualAllocator)
+    free(C)
+    return nothing
+end
+
+# ------------------------------------------------------------------------------------------
+# BumperAllocator implementation
+# ------------------------------------------------------------------------------------------
+
+Base.@constprop :aggressive function tensoralloc_add(TC, A::AbstractArray, pA::Index2Tuple,
+                                                     conjA::Bool,
+                                                     istemp::Bool,
+                                                     buf::Union{SlabBuffer,AllocBuffer})
+    structure = tensoradd_structure(A, pA, conjA)
+    if istemp
+        return Bumper.alloc!(buf, TC, structure...)
+    else
+        ttype = tensoradd_type(TC, A, pA, conjA)
+        return tensoralloc(ttype, structure, istemp)::ttype
+    end
+end
+
+Base.@constprop :aggressive function tensoralloc_contract(TC,
+                                                          A::AbstractArray, pA::Index2Tuple,
+                                                          conjA::Bool,
+                                                          B::AbstractArray, pB::Index2Tuple,
+                                                          conjB::Bool,
+                                                          pAB::Index2Tuple, istemp::Bool,
+                                                          buf::Union{SlabBuffer,
+                                                                     AllocBuffer})
+    structure = tensorcontract_structure(A, pA, conjA, B, pB, conjB, pAB)
+    if istemp
+        return Bumper.alloc!(buf, TC, structure...)
+    else
+        ttype = tensorcontract_type(TC, A, pA, conjA, B, pB, conjB, pAB)
+        return tensoralloc(ttype, structure, istemp)::ttype
+    end
+end
diff --git a/src/implementation/strided.jl b/src/implementation/strided.jl
index 89f4714..c79e75d 100644
--- a/src/implementation/strided.jl
+++ b/src/implementation/strided.jl
@@ -1,4 +1,4 @@
-const HostStridedView{T,N,A<:Array{T}} = StridedView{T,N,A}
+const HostStridedView{T,N,A<:DenseArray{T}} = StridedView{T,N,A}
 const HostStridedViewOrDiagonal{T} = Union{HostStridedView{T},Diagonal{T,Vector{T}}}
 const StridedViewOrDiagonal = Union{StridedView,Diagonal}
 
diff --git a/src/indexnotation/tensormacros.jl b/src/indexnotation/tensormacros.jl
index b1f7b8b..d2499bf 100644
--- a/src/indexnotation/tensormacros.jl
+++ b/src/indexnotation/tensormacros.jl
@@ -316,3 +316,26 @@ macro cutensor(ex::Expr)
                                 Expr(:call, GlobalRef(TensorOperations, :CUDAAllocator))))
     return esc(parser(ex))
 end
+
+"""
+    @butensor tensor_expr
+
+Use Bumper.jl to handle allocation of temporary tensors. This macro will use the default
+buffer and automatically reset it after the tensor expression has been evaluated. This macro
+is equivalent to `@no_escape @tensor tensor_expr` with all temporary allocations handled by
+Bumper.jl.
+"""
+macro butensor(ex...)
+    buf_sym = gensym("buffer")
+    cp_sym = gensym("checkpoint")
+    res_sym = gensym("result")
+    return esc(quote
+                   $buf_sym = $(Expr(:call, GlobalRef(Bumper, :default_buffer)))
+                   $cp_sym = $(Expr(:call, GlobalRef(Bumper, :checkpoint_save), buf_sym))
+                   $res_sym = $(Expr(:macrocall,
+                                     GlobalRef(TensorOperations, Symbol("@tensor")),
+                                     __source__, :(allocator = $buf_sym), ex...))
+                   $(Expr(:call, GlobalRef(Bumper, :checkpoint_restore!), cp_sym))
+                   $res_sym
+               end)
+end