Skip to content

Commit 0abce1e

Browse files
committed
Add reallocation wihtout TruffleBoundary
1 parent 637f43a commit 0abce1e

File tree

5 files changed

+155
-100
lines changed

5 files changed

+155
-100
lines changed

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/common/SequenceStorageNodes.java

Lines changed: 15 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,7 @@
112112
import com.oracle.graal.python.runtime.PythonContext;
113113
import com.oracle.graal.python.runtime.PythonOptions;
114114
import com.oracle.graal.python.runtime.exception.PException;
115+
import com.oracle.graal.python.runtime.native_memory.NativeBuffer;
115116
import com.oracle.graal.python.runtime.object.PythonObjectFactory;
116117
import com.oracle.graal.python.runtime.sequence.PSequence;
117118
import com.oracle.graal.python.runtime.sequence.storage.ArrayBasedSequenceStorage;
@@ -787,9 +788,8 @@ protected static SequenceStorage doMroSequenceStorage(MroSequenceStorage storage
787788
@Specialization
788789
protected static SequenceStorage doNativeInt(NativeIntSequenceStorage storage, int start, int stop, int step, int length,
789790
@Bind("this") Node node) {
790-
var valueBufferAddr = doNativePrimitiveSliceInBound(node, start, step, length, storage);
791-
long sizeInBytes = length * storage.getItemSize();
792-
return PythonContext.get(node).nativeBufferContext.wrapToIntStorage(valueBufferAddr, sizeInBytes, length);
791+
NativeBuffer sliceValueBuffer = doNativePrimitiveSliceInBound(PythonContext.get(node), start, step, length, storage);
792+
return PythonContext.get(node).nativeBufferContext.createNativeIntStorage(sliceValueBuffer, length);
793793
}
794794

795795
@Specialization
@@ -814,25 +814,25 @@ protected static SequenceStorage doNativeObject(NativeObjectSequenceStorage stor
814814
return new ObjectSequenceStorage(newArray);
815815
}
816816

817-
private static long doNativePrimitiveSliceInBound(Node node, int start, int step, int sliceLength, NativePrimitiveSequenceStorage storage) {
817+
private static NativeBuffer doNativePrimitiveSliceInBound(PythonContext pythonCtx, int start, int step, int sliceLength, NativePrimitiveSequenceStorage storage) {
818+
var unsafe = pythonCtx.getUnsafe();
818819
long itemSize = storage.getItemSize();
819-
var nativeContext = PythonContext.get(node).nativeBufferContext;
820820
long sizeInBytes = sliceLength * itemSize;
821-
long toAddr = nativeContext.allocateNativeMemory(sizeInBytes);
821+
NativeBuffer sliceBuffer = NativeBuffer.allocateNew(sizeInBytes);
822822

823823
if (step == 1) {
824824
var startAddress = storage.getValueBufferAddr() + (start * itemSize);
825-
nativeContext.copyMemory(startAddress, toAddr, sizeInBytes);
826-
return toAddr;
825+
unsafe.copyMemory(startAddress, sliceBuffer.getMemoryAddress(), sizeInBytes);
826+
return sliceBuffer;
827827
}
828828

829829
var stepInBytes = step * itemSize;
830-
for (long srcAddr = storage.getValueBufferAddr() + (start * itemSize), destAddr = toAddr,
830+
for (long srcAddr = storage.getValueBufferAddr() + (start * itemSize), destAddr = sliceBuffer.getMemoryAddress(),
831831
j = 0; j < sliceLength; srcAddr += stepInBytes, destAddr += itemSize, j++) {
832-
nativeContext.copyMemory(srcAddr, destAddr, itemSize);
832+
unsafe.copyMemory(srcAddr, destAddr, itemSize);
833833
}
834834

835-
return toAddr;
835+
return sliceBuffer;
836836
}
837837

838838
@NeverDefault
@@ -3144,14 +3144,9 @@ static void doObject(ObjectSequenceStorage storage, int cap) {
31443144
}
31453145

31463146
@Specialization
3147-
static void doNativePrimitive(Node inliningTarget, NativePrimitiveSequenceStorage storage, int cap) {
3147+
static void doNativePrimitive(NativePrimitiveSequenceStorage storage, int cap) {
31483148
if (CompilerDirectives.injectBranchProbability(CompilerDirectives.UNLIKELY_PROBABILITY, cap > storage.getCapacity())) {
3149-
var nativeContext = getContext(inliningTarget).nativeBufferContext;
3150-
long newCapacityInBytes = storage.getItemSize() * cap;
3151-
var oldAddr = storage.getValueBufferAddr();
3152-
var newAddr = nativeContext.allocateNativeMemory(newCapacityInBytes);
3153-
nativeContext.copyMemory(oldAddr, newAddr, storage.getCapacityInBytes());
3154-
nativeContext.setNewValueAddrToStorage(storage, newAddr, newCapacityInBytes);
3149+
storage.reallocate(cap);
31553150
}
31563151
}
31573152

@@ -3298,11 +3293,9 @@ static SequenceStorage doMro(MroSequenceStorage storage) {
32983293
@Specialization
32993294
static SequenceStorage doNativeInt(Node inliningTarget, NativeIntSequenceStorage storage) {
33003295
var nativeContext = PythonContext.get(inliningTarget).getContext().nativeBufferContext;
3301-
var capacityInBytes = storage.getCapacityInBytes();
3302-
var toAddr = nativeContext.allocateNativeMemory(capacityInBytes);
3303-
nativeContext.copyMemory(storage.getValueBufferAddr(), toAddr, capacityInBytes);
3296+
var copiedBuffer = storage.getValueBuffer().copy();
33043297

3305-
return nativeContext.wrapToIntStorage(toAddr, capacityInBytes, storage.length());
3298+
return nativeContext.createNativeIntStorage(copiedBuffer, storage.length());
33063299
}
33073300

33083301
@Specialization

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/runtime/NativeBufferContext.java

Lines changed: 17 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@
4242

4343
import com.oracle.graal.python.PythonLanguage;
4444
import com.oracle.graal.python.builtins.objects.cext.hpy.GraalHPyContext;
45+
import com.oracle.graal.python.runtime.native_memory.NativeBuffer;
4546
import com.oracle.graal.python.runtime.sequence.storage.NativeIntSequenceStorage;
4647
import com.oracle.graal.python.runtime.sequence.storage.NativePrimitiveSequenceStorage;
4748
import com.oracle.graal.python.util.PythonUtils;
@@ -59,10 +60,10 @@
5960
public class NativeBufferContext {
6061
private static final Unsafe unsafe = PythonUtils.initUnsafe();
6162

62-
@CompilationFinal private ReferenceQueue<NativePrimitiveSequenceStorage> referenceQueue;
63+
private ReferenceQueue<NativePrimitiveSequenceStorage> referenceQueue;
6364
// We need to keep around the references since Phantom References are one way bound
6465
// Key: MemoryAddr, Value: PhantomReference
65-
private ConcurrentHashMap<Long, NativePrimitiveReference> phantomReferences;
66+
private ConcurrentHashMap<NativePrimitiveReference, NativePrimitiveReference> phantomReferences;
6667

6768
private Thread nativeBufferReferenceCleanerThread;
6869

@@ -81,37 +82,11 @@ public void initReferenceQueue() {
8182
}
8283
}
8384

84-
public long allocateNativeMemory(long capacityInBytes) {
85-
assert capacityInBytes >= 0;
86-
return unsafe.allocateMemory(capacityInBytes);
87-
}
88-
89-
public void copyMemory(long fromAddress, long toAddress, long capacityInBytes) {
90-
assert capacityInBytes >= 0;
91-
unsafe.copyMemory(fromAddress, toAddress, capacityInBytes);
92-
}
93-
9485
@TruffleBoundary
95-
public void releaseMemory(long memoryAddr) {
96-
NativePrimitiveReference ref = phantomReferences.remove(memoryAddr);
97-
ref.release(unsafe);
98-
}
99-
100-
@TruffleBoundary
101-
public void setNewValueAddrToStorage(NativePrimitiveSequenceStorage storage, long valueAddr, long capacityInBytes) {
102-
assert capacityInBytes >= 0;
103-
releaseMemory(storage.getValueBufferAddr());
104-
storage.setValueBufferAddr(valueAddr, capacityInBytes);
86+
public NativeIntSequenceStorage createNativeIntStorage(NativeBuffer valueBuffer, int length) {
87+
var storage = new NativeIntSequenceStorage(valueBuffer, length);
10588
var phantomRef = new NativePrimitiveReference(storage, getReferenceQueue());
106-
phantomReferences.put(valueAddr, phantomRef);
107-
}
108-
109-
@TruffleBoundary
110-
public NativeIntSequenceStorage wrapToIntStorage(long valueBufferAddr, long capacityInBytes, int length) {
111-
assert capacityInBytes >= 0;
112-
var storage = new NativeIntSequenceStorage(valueBufferAddr, capacityInBytes, length);
113-
var phantomRef = new NativePrimitiveReference(storage, getReferenceQueue());
114-
phantomReferences.put(storage.getValueBufferAddr(), phantomRef);
89+
phantomReferences.put(phantomRef, phantomRef);
11590

11691
return storage;
11792
}
@@ -132,10 +107,10 @@ public void finalizeContext() {
132107

133108
public NativeIntSequenceStorage toNativeIntStorage(int[] arr) {
134109
long sizeInBytes = (long) arr.length * (long) Integer.BYTES;
135-
long addr = allocateNativeMemory(sizeInBytes);
110+
NativeBuffer nativeBuffer = NativeBuffer.allocateNew(sizeInBytes);
136111

137-
unsafe.copyMemory(arr, Unsafe.ARRAY_INT_BASE_OFFSET, null, addr, sizeInBytes);
138-
return wrapToIntStorage(addr, sizeInBytes, arr.length);
112+
unsafe.copyMemory(arr, Unsafe.ARRAY_INT_BASE_OFFSET, null, nativeBuffer.getMemoryAddress(), sizeInBytes);
113+
return createNativeIntStorage(nativeBuffer, arr.length);
139114
}
140115

141116
private ReferenceQueue<NativePrimitiveSequenceStorage> getReferenceQueue() {
@@ -151,13 +126,11 @@ private ReferenceQueue<NativePrimitiveSequenceStorage> getReferenceQueue() {
151126
static final class NativeBufferDeallocatorRunnable implements Runnable {
152127
private static final TruffleLogger LOGGER = GraalHPyContext.getLogger(NativeBufferDeallocatorRunnable.class);
153128

154-
private static final Unsafe unsafe = PythonUtils.initUnsafe();
155-
156129
private final ReferenceQueue<NativePrimitiveSequenceStorage> referenceQueue;
157-
private final ConcurrentHashMap<Long, NativePrimitiveReference> references;
130+
private final ConcurrentHashMap<NativePrimitiveReference, NativePrimitiveReference> references;
158131

159132
public NativeBufferDeallocatorRunnable(ReferenceQueue<NativePrimitiveSequenceStorage> referenceQueue,
160-
ConcurrentHashMap<Long, NativePrimitiveReference> references) {
133+
ConcurrentHashMap<NativePrimitiveReference, NativePrimitiveReference> references) {
161134
this.referenceQueue = referenceQueue;
162135
this.references = references;
163136
}
@@ -170,8 +143,8 @@ public void run() {
170143
while (!pythonContext.getThreadState(language).isShuttingDown()) {
171144
try {
172145
NativePrimitiveReference phantomRef = (NativePrimitiveReference) referenceQueue.remove();
173-
phantomRef.release(unsafe);
174-
references.remove(phantomRef.getMemoryAddress());
146+
phantomRef.release();
147+
references.remove(phantomRef);
175148
} catch (InterruptedException e) {
176149
Thread.currentThread().interrupt();
177150
LOGGER.fine("Native buffer reference cleaner thread was interrupted and is exiting");
@@ -184,23 +157,15 @@ public void run() {
184157

185158
static final class NativePrimitiveReference extends PhantomReference<NativePrimitiveSequenceStorage> {
186159

187-
private final long memoryAddress;
188-
private boolean released = false;
160+
private final NativeBuffer buffer;
189161

190162
public NativePrimitiveReference(NativePrimitiveSequenceStorage referent, ReferenceQueue<NativePrimitiveSequenceStorage> q) {
191163
super(referent, q);
192-
this.memoryAddress = referent.getValueBufferAddr();
193-
}
194-
195-
public void release(Unsafe unsafe) {
196-
if (!released) {
197-
unsafe.freeMemory(memoryAddress);
198-
released = true;
199-
}
164+
this.buffer = referent.getValueBuffer();
200165
}
201166

202-
public long getMemoryAddress() {
203-
return memoryAddress;
167+
public void release() {
168+
buffer.release();
204169
}
205170
}
206171
}
Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
/*
2+
* Copyright (c) 2024, Oracle and/or its affiliates. All rights reserved.
3+
* DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
4+
*
5+
* The Universal Permissive License (UPL), Version 1.0
6+
*
7+
* Subject to the condition set forth below, permission is hereby granted to any
8+
* person obtaining a copy of this software, associated documentation and/or
9+
* data (collectively the "Software"), free of charge and under any and all
10+
* copyright rights in the Software, and any and all patent rights owned or
11+
* freely licensable by each licensor hereunder covering either (i) the
12+
* unmodified Software as contributed to or provided by such licensor, or (ii)
13+
* the Larger Works (as defined below), to deal in both
14+
*
15+
* (a) the Software, and
16+
*
17+
* (b) any piece of software and/or hardware listed in the lrgrwrks.txt file if
18+
* one is included with the Software each a "Larger Work" to which the Software
19+
* is contributed by such licensors),
20+
*
21+
* without restriction, including without limitation the rights to copy, create
22+
* derivative works of, display, perform, and distribute the Software and make,
23+
* use, sell, offer for sale, import, export, have made, and have sold the
24+
* Software and the Larger Work(s), and to sublicense the foregoing rights on
25+
* either these or other terms.
26+
*
27+
* This license is subject to the following condition:
28+
*
29+
* The above copyright notice and either this complete permission notice or at a
30+
* minimum a reference to the UPL must be included in all copies or substantial
31+
* portions of the Software.
32+
*
33+
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
34+
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
35+
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
36+
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
37+
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
38+
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
39+
* SOFTWARE.
40+
*/
41+
package com.oracle.graal.python.runtime.native_memory;
42+
43+
import com.oracle.graal.python.util.PythonUtils;
44+
import sun.misc.Unsafe;
45+
46+
public class NativeBuffer {
47+
48+
private static final Unsafe unsafe = PythonUtils.initUnsafe();
49+
private long memoryAddress;
50+
private long capacityInBytes;
51+
52+
private NativeBuffer(long memoryAddress, long capacityInBytes) {
53+
this.memoryAddress = memoryAddress;
54+
this.capacityInBytes = capacityInBytes;
55+
}
56+
57+
public void reallocate(long newCapacityInBytes) {
58+
assert newCapacityInBytes >= 0;
59+
long newMemoryAddr = unsafe.allocateMemory(newCapacityInBytes);
60+
unsafe.copyMemory(memoryAddress, newMemoryAddr, newCapacityInBytes);
61+
unsafe.freeMemory(memoryAddress);
62+
memoryAddress = newMemoryAddr;
63+
capacityInBytes = newCapacityInBytes;
64+
}
65+
66+
public NativeBuffer copy() {
67+
long newAddr = unsafe.allocateMemory(capacityInBytes);
68+
unsafe.copyMemory(memoryAddress, newAddr, capacityInBytes);
69+
70+
return new NativeBuffer(newAddr, capacityInBytes);
71+
}
72+
73+
public NativeBuffer copy(long newCapacityInBytes) {
74+
assert newCapacityInBytes >= 0;
75+
long newAddr = unsafe.allocateMemory(newCapacityInBytes);
76+
unsafe.copyMemory(memoryAddress, newAddr, capacityInBytes);
77+
78+
return new NativeBuffer(newAddr, newCapacityInBytes);
79+
}
80+
81+
public void release() {
82+
unsafe.freeMemory(memoryAddress);
83+
}
84+
85+
public static NativeBuffer allocateNew(long capacityInBytes) {
86+
assert capacityInBytes >= 0;
87+
long adr = unsafe.allocateMemory(capacityInBytes);
88+
return new NativeBuffer(adr, capacityInBytes);
89+
}
90+
91+
public long getMemoryAddress() {
92+
return memoryAddress;
93+
}
94+
95+
public long getCapacityInBytes() {
96+
return capacityInBytes;
97+
}
98+
}

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/runtime/sequence/storage/NativeIntSequenceStorage.java

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -40,10 +40,12 @@
4040
*/
4141
package com.oracle.graal.python.runtime.sequence.storage;
4242

43+
import com.oracle.graal.python.runtime.native_memory.NativeBuffer;
44+
4345
public class NativeIntSequenceStorage extends NativePrimitiveSequenceStorage {
4446

45-
public NativeIntSequenceStorage(long valueBufferAddr, long bufferCapacityInBytes, int length) {
46-
super(valueBufferAddr, bufferCapacityInBytes, length, Integer.BYTES);
47+
public NativeIntSequenceStorage(NativeBuffer valueBuffer, int length) {
48+
super(valueBuffer, length, Integer.BYTES);
4749
}
4850

4951
@Override
@@ -57,12 +59,12 @@ public Object getIndicativeValue() {
5759
}
5860

5961
public int getIntItemNormalized(int idx) {
60-
long indexInBytes = (long) idx * this.itemSize;
61-
return unsafe.getInt(this.valueBufferAddr + indexInBytes);
62+
long indexInBytes = (long) idx * Integer.BYTES;
63+
return unsafe.getInt(getValueBufferAddr() + indexInBytes);
6264
}
6365

6466
public void setIntItemNormalized(int idx, int value) {
65-
long indexInBytes = (long) idx * this.itemSize;
66-
unsafe.putInt(this.valueBufferAddr + indexInBytes, value);
67+
long indexInBytes = (long) idx * Integer.BYTES;
68+
unsafe.putInt(getValueBufferAddr() + indexInBytes, value);
6769
}
6870
}

0 commit comments

Comments
 (0)