diff --git a/src/TensorFlowNET.Core/Util/UnmanagedExtensions.cs b/src/TensorFlowNET.Core/Util/UnmanagedExtensions.cs new file mode 100644 index 00000000..02b8bb73 --- /dev/null +++ b/src/TensorFlowNET.Core/Util/UnmanagedExtensions.cs @@ -0,0 +1,94 @@ +using System; +using System.IO; +using System.Runtime.CompilerServices; +using System.Runtime.InteropServices; +using NumSharp.Backends.Unmanaged; + +namespace Tensorflow.Util +{ + public static class UnmanagedExtensions + { + //internally UnmanagedMemoryStream can't construct with null address. + private static readonly unsafe byte* _empty = (byte*) Marshal.AllocHGlobal(1); + + /// + /// Creates a memory stream based on given . + /// + /// The block to stream. Can be default/null. + /// There is no need to dispose the returned + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static UnmanagedMemoryStream Stream(this UnmanagedMemoryBlock block) + { + unsafe + { + if (block.Address == null) + return new UnmanagedMemoryStream(_empty, 0); + return new UnmanagedMemoryStream(block.Address, block.BytesCount); + } + } + + /// + /// Creates a memory stream based on given . + /// + /// The block to stream. Can be default/null. + /// Offset from the start of the block. + /// There is no need to dispose the returned + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static UnmanagedMemoryStream Stream(this UnmanagedMemoryBlock block, long offset) + { + if (block.BytesCount - offset <= 0) + throw new ArgumentOutOfRangeException(nameof(offset)); + + unsafe + { + if (block.Address == null) + return new UnmanagedMemoryStream(_empty, 0); + return new UnmanagedMemoryStream(block.Address + offset, block.BytesCount - offset); + } + } + + /// + /// Creates a memory stream based on given . + /// + /// The block to stream. Can be IntPtr.Zero. + /// The length of the block in bytes. + /// There is no need to dispose the returned + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static UnmanagedMemoryStream Stream(this IntPtr address, long length) + { + if (length <= 0) + throw new ArgumentOutOfRangeException(nameof(length)); + + unsafe + { + if (address == IntPtr.Zero) + return new UnmanagedMemoryStream(_empty, 0); + + // ReSharper disable once AssignNullToNotNullAttribute + return new UnmanagedMemoryStream((byte*) address, length); + } + } + + /// + /// Creates a memory stream based on given . + /// + /// The block to stream. Can be IntPtr.Zero. + /// Offset from the start of the block. + /// The length of the block in bytes. + /// There is no need to dispose the returned + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static UnmanagedMemoryStream Stream(this IntPtr address, long offset, long length) + { + if (length <= 0) + throw new ArgumentOutOfRangeException(nameof(length)); + + unsafe + { + if (address == IntPtr.Zero) + return new UnmanagedMemoryStream(_empty, 0); + + return new UnmanagedMemoryStream((byte*) address + offset, length); + } + } + } +} \ No newline at end of file