Initial push

mcimadamore · Jun 6, 2024 · 2092836 · 2092836
1 parent 3cbdf8d
commit 2092836
Show file tree

Hide file tree

Showing 6 changed files with 147 additions and 7 deletions.
diff --git a/src/java.base/share/classes/java/lang/foreign/SegmentAllocator.java b/src/java.base/share/classes/java/lang/foreign/SegmentAllocator.java
@@ -33,7 +33,7 @@
 import jdk.internal.foreign.ArenaImpl;
 import jdk.internal.foreign.SlicingAllocator;
 import jdk.internal.foreign.StringSupport;
-import jdk.internal.foreign.Utils;
+import jdk.internal.foreign.ZeroingAllocator;
 import jdk.internal.vm.annotation.ForceInline;
 
 /**
@@ -712,6 +712,20 @@ static SegmentAllocator prefixAllocator(MemorySegment segment) {
         return (AbstractMemorySegmentImpl)segment;
     }
 
+    /**
+     * {@return a segment allocator zero-initializes the segments obtained by the provided allocator}
+     * <p>
+     * Equivalent to (but likely more efficient than) the following code:
+     * {@snippet lang=java :
+     * SegmentAllocator allocator = ...
+     * SegmentAllocator zeroingAllocator = (size, align) -> allocator.allocate(size, align).fill((byte)0);
+     * }
+     * @param allocator the allocator the returned allocator delegates to
+     */
+    static SegmentAllocator zeroingAllocator(SegmentAllocator allocator) {
+        return ZeroingAllocator.of(allocator);
+    }
+
     private static void assertWritable(MemorySegment segment) {
         // Implicit null check
         if (segment.isReadOnly()) {

diff --git a/src/java.base/share/classes/jdk/internal/foreign/ArenaImpl.java b/src/java.base/share/classes/jdk/internal/foreign/ArenaImpl.java
@@ -31,7 +31,7 @@
 import java.lang.foreign.MemorySegment.Scope;
 import java.util.Objects;
 
-public final class ArenaImpl implements Arena {
+public final class ArenaImpl implements Arena, ZeroingAllocator {
 
     private final MemorySessionImpl session;
     private final boolean shouldReserveMemory;

diff --git a/src/java.base/share/classes/jdk/internal/foreign/ZeroingAllocator.java b/src/java.base/share/classes/jdk/internal/foreign/ZeroingAllocator.java
@@ -0,0 +1,55 @@
+package jdk.internal.foreign;
+
+import jdk.internal.vm.annotation.Stable;
+
+import java.lang.foreign.MemorySegment;
+import java.lang.foreign.SegmentAllocator;
+import java.lang.foreign.ValueLayout;
+
+/**
+ * Marker interface for allocators that provide zeroing semantics.
+ */
+public interface ZeroingAllocator extends SegmentAllocator {
+
+    static ZeroingAllocator of(SegmentAllocator allocator) {
+        if (allocator instanceof ZeroingAllocator zeroingAllocator) {
+            // already a zeroing allocator
+            return zeroingAllocator;
+        } else if (allocator instanceof MemorySegment ||
+                   allocator instanceof SlicingAllocator) {
+            // prefix and slicing allocators do not zero
+            return (byteSize, byteAlignment) -> allocator.allocate(byteSize, byteAlignment).fill((byte)0);
+        } else {
+            return new ZeroingWrapper(allocator);
+        }
+    }
+
+    class ZeroingWrapper implements ZeroingAllocator {
+        static final MemorySegment ZERO = MemorySegment.ofArray(new long[1024]);
+        @Stable
+        static final ValueLayout[] ALIGNED_LAYOUTS;
+
+        final SegmentAllocator allocator;
+
+        ZeroingWrapper(SegmentAllocator allocator) {
+            this.allocator = allocator;
+        }
+
+        static {
+            ALIGNED_LAYOUTS = new ValueLayout[8];
+            ALIGNED_LAYOUTS[0] = ValueLayout.JAVA_BYTE;
+            ALIGNED_LAYOUTS[1] = ValueLayout.JAVA_SHORT;
+            ALIGNED_LAYOUTS[3] = ValueLayout.JAVA_INT;
+            ALIGNED_LAYOUTS[7] = ValueLayout.JAVA_LONG;
+        }
+
+        public MemorySegment allocate(long size, long align) {
+            if (align > 8 || size > ZERO.byteSize()) {
+                // slow path (might perform double zeroing)
+                return allocator.allocate(size, align).fill((byte) 0);
+            } else {
+                return allocator.allocateFrom(ALIGNED_LAYOUTS[(int)align - 1], ZERO, ALIGNED_LAYOUTS[(int)align - 1], 0, size);
+            }
+        };
+    }
+}
diff --git a/test/jdk/java/foreign/TestSegmentAllocators.java b/test/jdk/java/foreign/TestSegmentAllocators.java
@@ -478,8 +478,13 @@ interface OfDoubleArray extends AllocationFunction<double[], ValueLayout.OfDoubl
     }
 
     enum AllocationFactory {
+        ARENA(false, (_, arena) -> arena),
+        ZEROING(false, (_, arena) -> SegmentAllocator.zeroingAllocator(arena)),
         SLICING(true, (size, arena) -> {
             return SegmentAllocator.slicingAllocator(arena.allocate(size, 1));
+        }),
+        ZERO_SLICING(true, (size, arena) -> {
+            return SegmentAllocator.zeroingAllocator(SLICING.factory.apply(size, arena));
         });
 
         private final boolean isBound;
@@ -611,8 +616,15 @@ public double[] toArray(MemorySegment segment, ValueLayout layout) {
 
     @DataProvider(name = "allocators")
     static Object[][] allocators() {
+        SegmentAllocator prefix = SegmentAllocator.prefixAllocator(Arena.global().allocate(10, 1));
+        SegmentAllocator slicing = SegmentAllocator.slicingAllocator(Arena.global().allocate(10, 1));
         return new Object[][] {
-                { SegmentAllocator.prefixAllocator(Arena.global().allocate(10, 1)) },
+                { Arena.global() },
+                { SegmentAllocator.zeroingAllocator(Arena.global()) },
+                { prefix },
+                { SegmentAllocator.zeroingAllocator(prefix) },
+                { slicing },
+                { SegmentAllocator.zeroingAllocator(slicing) },
         };
     }
 }
diff --git a/test/micro/org/openjdk/bench/java/lang/foreign/AllocFromTest.java b/test/micro/org/openjdk/bench/java/lang/foreign/AllocFromTest.java
@@ -87,6 +87,13 @@ public MemorySegment alloc_unsafe_arena() {
         }
     }
 
+    @Benchmark
+    public MemorySegment alloc_zeroing_arena() {
+        try (ZeroingArena arena = new ZeroingArena()) {
+            return arena.allocateFrom(ValueLayout.JAVA_BYTE, arr);
+        }
+    }
+
     @Benchmark
     public MemorySegment alloc_pool_arena() {
         try (Arena arena = pool.acquire()) {
@@ -169,4 +176,29 @@ public MemorySegment allocate(long byteSize, long byteAlignment) {
                     .reinterpret(byteSize, arena, ms -> Utils.unsafe.freeMemory(ms.address()));
         }
     }
+
+    public static class ZeroingArena implements Arena {
+        final Arena arena = Arena.ofConfined();
+        final SegmentAllocator allocator = SegmentAllocator.zeroingAllocator(arena);
+
+        @Override
+        public Scope scope() {
+            return arena.scope();
+        }
+
+        @Override
+        public void close() {
+            arena.close();
+        }
+
+        @Override
+        public MemorySegment allocate(long byteSize, long byteAlignment) {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public MemorySegment allocateFrom(ValueLayout.OfByte layout, byte... elements) {
+            return allocator.allocateFrom(layout, elements);
+        }
+    }
 }
diff --git a/test/micro/org/openjdk/bench/java/lang/foreign/AllocTest.java b/test/micro/org/openjdk/bench/java/lang/foreign/AllocTest.java
@@ -71,16 +71,23 @@ public MemorySegment alloc_confined() {
     }
 
     @Benchmark
-    public long alloc_calloc_arena() {
+    public MemorySegment alloc_calloc_arena() {
         try (CallocArena arena = new CallocArena()) {
-            return arena.allocate(size).address();
+            return arena.allocate(size);
         }
     }
 
     @Benchmark
-    public long alloc_unsafe_arena() {
+    public MemorySegment alloc_unsafe_arena() {
         try (UnsafeArena arena = new UnsafeArena()) {
-            return arena.allocate(size).address();
+            return arena.allocate(size);
+        }
+    }
+
+    @Benchmark
+    public MemorySegment alloc_zeroing_arena() {
+        try (ZeroingArena arena = new ZeroingArena()) {
+            return arena.allocate(size);
         }
     }
 
@@ -139,4 +146,24 @@ public MemorySegment allocate(long byteSize, long byteAlignment) {
             return segment.reinterpret(byteSize, arena, ms -> Utils.unsafe.freeMemory(segment.address()));
         }
     }
+
+    public static class ZeroingArena implements Arena {
+        final Arena arena = Arena.ofConfined();
+        final SegmentAllocator allocator = SegmentAllocator.zeroingAllocator(arena);
+
+        @Override
+        public Scope scope() {
+            return arena.scope();
+        }
+
+        @Override
+        public void close() {
+            arena.close();
+        }
+
+        @Override
+        public MemorySegment allocate(long byteSize, long byteAlignment) {
+            return allocator.allocate(byteSize, byteAlignment);
+        }
+    }
 }