Files
unity-application/Packages/com.unity.barracuda/Runtime/Core/BarracudaArray.cs
2023-03-18 19:53:17 +00:00

546 lines
20 KiB
C#

using System;
using System.Runtime.ConstrainedExecution;
using System.Runtime.InteropServices;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using Unity.Mathematics;
using UnityEngine;
using UnityEngine.Assertions;
namespace Unity.Barracuda
{
///see https://referencesource.microsoft.com/#mscorlib/system/runtime/interopservices/safehandle.cs
internal class NativeMemorySafeHandle : SafeHandle
{
public readonly Allocator m_AllocatorLabel;
[ReliabilityContract(Consistency.WillNotCorruptState, Cer.MayFail)]
public unsafe NativeMemorySafeHandle(long size, int alignment, Allocator allocator) : base(IntPtr.Zero, true)
{
m_AllocatorLabel = allocator;
if (size > 0)
SetHandle((IntPtr)UnsafeUtility.Malloc(size, alignment, allocator));
}
public override bool IsInvalid {
get { return handle == IntPtr.Zero; }
}
[ReliabilityContract(Consistency.WillNotCorruptState, Cer.Success)]
protected override unsafe bool ReleaseHandle()
{
UnsafeUtility.Free((void*)handle, m_AllocatorLabel);
return true;
}
}
internal class PinnedMemorySafeHandle : SafeHandle
{
private readonly GCHandle m_GCHandle;
[ReliabilityContract(Consistency.WillNotCorruptState, Cer.MayFail)]
public PinnedMemorySafeHandle(object managedObject) : base(IntPtr.Zero, true)
{
m_GCHandle = GCHandle.Alloc(managedObject, GCHandleType.Pinned);
IntPtr pinnedPtr = m_GCHandle.AddrOfPinnedObject();
SetHandle(pinnedPtr);
}
public override bool IsInvalid {
get { return handle == IntPtr.Zero; }
}
[ReliabilityContract(Consistency.WillNotCorruptState, Cer.Success)]
protected override bool ReleaseHandle()
{
m_GCHandle.Free();
return true;
}
}
/// <summary>
/// A BarracudaArrayFromManagedArray exposes a buffer of managed memory as if it was native memory (by pinning it).
/// </summary>
public class BarracudaArrayFromManagedArray : BarracudaArray
{
private readonly int m_PinnedMemoryByteOffset;
public BarracudaArrayFromManagedArray(float[] srcData, int srcOffset = 0) : this(srcData, srcOffset, sizeof(float), DataType.Float, srcData.Length-srcOffset)
{
}
public BarracudaArrayFromManagedArray(byte[] srcData, int srcOffset, DataType destType, int numDestElement) : this(srcData, srcOffset, sizeof(byte), destType, numDestElement)
{
}
private unsafe BarracudaArrayFromManagedArray(Array srcData, int srcElementOffset, int srcElementSize, DataType destElementType, int numDestElement) : base(new PinnedMemorySafeHandle(srcData), destElementType, numDestElement)
{
m_PinnedMemoryByteOffset = srcElementSize * srcElementOffset;
//Safety checks
int requiredAlignment = DataAlignmentSize(destElementType);
int srcLenghtInByte = (srcData.Length - srcElementOffset) * srcElementSize;
int dstLenghtInByte = numDestElement * DataItemSize(destElementType);
IntPtr pinnedPtrWithOffset = (IntPtr) base.RawPtr + m_PinnedMemoryByteOffset;
if (srcElementOffset > srcData.Length)
throw new ArgumentOutOfRangeException(nameof (srcElementOffset), "SrcElementOffset must be <= srcData.Length");
if (dstLenghtInByte > srcLenghtInByte)
throw new ArgumentOutOfRangeException(nameof (numDestElement), "NumDestElement too big for srcData and srcElementOffset");
if (pinnedPtrWithOffset.ToInt64() % requiredAlignment != 0)
throw new InvalidOperationException($"The BarracudaArrayFromManagedArray source ptr (including offset) need to be aligned on {requiredAlignment} bytes for the data to be express as {destElementType}.");
var neededSrcPaddedLengthInByte = LengthWithPaddingForGPUCopy(destElementType, numDestElement) * DataItemSize(destElementType);
if (srcLenghtInByte < neededSrcPaddedLengthInByte)
throw new InvalidOperationException($"The BarracudaArrayFromManagedArray source ptr (including offset) is to small to account for extra padding needing for type {destElementType}.");
}
public override unsafe void* RawPtr => (byte*) base.RawPtr + m_PinnedMemoryByteOffset;
}
public enum DataType
{
Float,
Half
}
/// <summary>
/// A BarracudaArray exposes a buffer of native memory to managed code.
/// </summary>
public class BarracudaArray : IDisposable
{
protected readonly SafeHandle m_SafeHandle;
private readonly Allocator m_Allocator;
private readonly int m_Length;
private readonly DataType m_DataType;
#region helpers
public static int DataItemSize(DataType dataType)
{
if (dataType == DataType.Float)
return UnsafeUtility.SizeOf<float>();
if (dataType == DataType.Half)
return UnsafeUtility.SizeOf<half>();
throw new NotImplementedException($"Type {dataType} not supported.");
}
public static int DataAlignmentSize(DataType dataType)
{
if (dataType == DataType.Float)
return UnsafeUtility.AlignOf<float>();
if (dataType == DataType.Half)
return UnsafeUtility.AlignOf<uint>();
throw new NotImplementedException($"Type {dataType} not supported.");
}
public static int LengthWithPaddingForGPUCopy(DataType dataType, int length)
{
if (dataType == DataType.Float)
return length;
if (dataType == DataType.Half)
return length + (length % 2);
throw new NotImplementedException($"Type {dataType} not supported.");
}
private void CheckElementAccess(DataType dataType, long index)
{
//Disabled by default for performance reasons.
#if ENABLE_BARRACUDA_DEBUG
if (Disposed)
throw new InvalidOperationException("The BarracudaArray was disposed.");
if (index <0 || index >= m_Length)
throw new IndexOutOfRangeException($"Accessing BarracudaArray of length {m_Length} at index {index}, data type is {m_DataType}.");
if (dataType != m_DataType)
throw new InvalidOperationException($"Accessing BarracudaArray of data type {m_DataType} as if it was {dataType}.");
#endif
}
#endregion
protected BarracudaArray(SafeHandle safeHandle, DataType dataType, int dataLength)
{
m_DataType = dataType;
m_Length = dataLength;
m_SafeHandle = safeHandle;
m_Allocator = Allocator.Persistent;
}
public BarracudaArray(int length, DataType dataType = DataType.Float, Allocator allocator = Allocator.Persistent)
{
if (!UnsafeUtility.IsValidAllocator(allocator))
throw new InvalidOperationException("The BarracudaArray should use a valid allocator.");
if (length < 0)
throw new ArgumentOutOfRangeException(nameof (length), "Length must be >= 0");
m_DataType = dataType;
m_Length = length;
m_SafeHandle = new NativeMemorySafeHandle(LengthWithPaddingForGPUCopy(m_DataType, m_Length) * DataItemSize(dataType), DataAlignmentSize(dataType), allocator);
m_Allocator = allocator;
}
public unsafe void ZeroMemory()
{
var numByteToClear = LengthWithPaddingForGPUCopy(m_DataType, m_Length) * DataItemSize(m_DataType);
UnsafeUtility.MemClear(RawPtr, numByteToClear);
}
public virtual void Dispose()
{
m_SafeHandle.Dispose();
}
#region properties
public DataType Type => m_DataType;
public int SizeOfType => DataItemSize(m_DataType);
public int Length => m_Length;
public long LongLength => m_Length;
public virtual unsafe void* RawPtr
{
get
{
if (Disposed)
throw new InvalidOperationException("The BarracudaArray was disposed.");
return (void*)m_SafeHandle.DangerousGetHandle();
}
}
public bool Disposed => m_SafeHandle.IsClosed;
#endregion
#region indexers and single access accessor
public unsafe float* AddressAt(long index)
{
Assert.AreEqual(DataType.Float, m_DataType);
return (float*) RawPtr + index;
}
public unsafe half* HalfAddressAt(long index)
{
Assert.AreEqual(DataType.Half, m_DataType);
return (half*) RawPtr + index;
}
public unsafe void* RawAddressAt(long index)
{
if (m_DataType == DataType.Half)
return HalfAddressAt(index);
else
return AddressAt(index);
}
public float this[long index]
{
get => this[(int)index];
set => this[(int)index] = value;
}
public float this[int index]
{
get
{
switch (m_DataType)
{
case DataType.Float:
return GetFloat(index);
default:
return GetHalf(index);
}
}
set
{
switch (m_DataType)
{
case DataType.Float:
SetFloat(index, value);
break;
default:
SetHalf(index, (half) value);
break;
}
}
}
public unsafe float GetFloat(int index)
{
CheckElementAccess(DataType.Float, index);
return UnsafeUtility.ReadArrayElement<float>(RawPtr, index);
}
public unsafe half GetHalf(int index)
{
CheckElementAccess(DataType.Half, index);
return UnsafeUtility.ReadArrayElement<half>(RawPtr, index);
}
public unsafe void SetFloat(int index, float value)
{
CheckElementAccess(DataType.Float, index);
UnsafeUtility.WriteArrayElement<float>(RawPtr, index, value);
}
public unsafe void SetHalf(int index, half value)
{
CheckElementAccess(DataType.Half, index);
UnsafeUtility.WriteArrayElement<half>(RawPtr, index, value);
}
#endregion
#region copy to other memory containers
public void UploadToComputeBuffer(ComputeBuffer buffer)
{
UploadToComputeBuffer(buffer, 0, 0, m_Length);
}
public unsafe void UploadToComputeBuffer(ComputeBuffer buffer, int elementStartIndex, int computeBufferStartIndex, int numElementToCopy)
{
if (numElementToCopy == 0)
return;
if (m_DataType == DataType.Float)
{
NativeArray<float> nativeArray = NativeArrayUnsafeUtility.ConvertExistingDataToNativeArray<float>(RawPtr, m_Length, m_Allocator);
#if ENABLE_UNITY_COLLECTIONS_CHECKS
NativeArrayUnsafeUtility.SetAtomicSafetyHandle(ref nativeArray, AtomicSafetyHandle.Create());
#endif
buffer.SetData(nativeArray, elementStartIndex, computeBufferStartIndex, numElementToCopy);
}
else if (m_DataType == DataType.Half)
{
if (elementStartIndex % 2 == 1 || computeBufferStartIndex % 2 == 1)
throw new ArgumentException($"For half buffer type nativeBufferStartIndex and computeBufferStartIndex should be modulo of 2.");
numElementToCopy += numElementToCopy % 2;
int uintBufferViewLength = LengthWithPaddingForGPUCopy(m_DataType, m_Length) / 2;
NativeArray<uint> nativeArray = NativeArrayUnsafeUtility.ConvertExistingDataToNativeArray<uint>(RawPtr, uintBufferViewLength, m_Allocator);
#if ENABLE_UNITY_COLLECTIONS_CHECKS
NativeArrayUnsafeUtility.SetAtomicSafetyHandle(ref nativeArray, AtomicSafetyHandle.Create());
#endif
//TODO fp16 should computeBufferStartIndex be expressed in half or uint? For now in half
buffer.SetData(nativeArray, elementStartIndex/2, computeBufferStartIndex/2, numElementToCopy/2);
}
else
{
throw new NotImplementedException($"Type {m_DataType} not supported.");
}
}
/// <summary>
/// Warning, this return a copy! Do not use to modify a BarracudaArray
/// </summary>
public static implicit operator float[](BarracudaArray barracudaArray)
{
var floatArray = new float[barracudaArray.Length];
Copy(barracudaArray, 0, floatArray, 0, barracudaArray.Length);
return floatArray;
}
public void CopyTo(BarracudaArray dst, int dstOffset)
{
Copy(this, 0, dst, dstOffset, Length);
}
public void CopyTo(BarracudaArray dst, long dstOffset)
{
Copy(this, 0, dst, (int)dstOffset, Length);
}
public static void Copy(BarracudaArray sourceArray, BarracudaArray destinationArray, int length = -1)
{
Copy(sourceArray, 0, destinationArray, 0, length);
}
public static void Copy(float[] sourceArray, BarracudaArray destinationArray, int length = -1)
{
Copy(sourceArray, 0, destinationArray, 0, length);
}
public static unsafe void Copy(
BarracudaArray sourceArray,
int sourceIndex,
BarracudaArray destinationArray,
int destinationIndex,
int length)
{
if (length < 0)
length = sourceArray.Length;
if (length == 0)
return;
if (sourceIndex+length > sourceArray.Length)
throw new ArgumentException($"Cannot copy {length} element from sourceIndex {sourceIndex} and Barracuda array of length {sourceArray.Length}.");
if (destinationIndex+length > destinationArray.Length)
throw new ArgumentException($"Cannot copy {length} element to sourceIndex {destinationIndex} and Barracuda array of length {destinationArray.Length}.");
//Same type we can do a memcopy
if (sourceArray.m_DataType == destinationArray.m_DataType)
{
int itemSize = DataItemSize(sourceArray.m_DataType);
void* srcPtr = (byte*)sourceArray.RawPtr + sourceIndex * itemSize;
void* dstPtr = (byte*)destinationArray.RawPtr + destinationIndex * itemSize;
UnsafeUtility.MemCpy(dstPtr, srcPtr, length * itemSize);
}
else//different type, we need to iterate and cast
{
for (var i=0; i < length; ++i)
{
//this will use float as intermediate/common representation
destinationArray[destinationIndex+i] = sourceArray[sourceIndex+i];
}
}
}
public static unsafe void Copy(
BarracudaArray sourceArray,
int sourceIndex,
float[] destinationArray,
int destinationIndex,
int length)
{
if (length < 0)
length = sourceArray.Length;
if (length == 0)
return;
if (sourceIndex+length > sourceArray.Length)
throw new ArgumentException($"Cannot copy {length} element from sourceIndex {sourceIndex} and Barracuda array of length {sourceArray.Length}.");
if (destinationIndex+length > destinationArray.Length)
throw new ArgumentException($"Cannot copy {length} element to sourceIndex {destinationIndex} and array of length {destinationArray.Length}.");
//Same type we can do a memcopy
if (sourceArray.m_DataType == DataType.Float)
{
fixed (void* dstPtr = &destinationArray[destinationIndex])
{
int itemSize = DataItemSize(sourceArray.m_DataType);
void* srcPtr = (byte*)sourceArray.RawPtr + sourceIndex * itemSize;
UnsafeUtility.MemCpy(dstPtr, srcPtr, length * itemSize);
}
}
else//different type, we need to iterate and cast
{
for (var i=0; i < length; ++i)
{
//this will use float as intermediate/common representation
destinationArray[destinationIndex+i] = sourceArray[sourceIndex+i];
}
}
}
public static unsafe void BlockCopy(
BarracudaArray sourceArray,
int sourceByteOffset,
byte[] destinationArray,
int destinationByteOffset,
int lengthInBytes)
{
int itemSize = sourceArray.SizeOfType;
int srcLengthBytes = sourceArray.Length * itemSize;
if (lengthInBytes == 0)
return;
if (lengthInBytes < 0)
lengthInBytes = srcLengthBytes;
if (sourceByteOffset+lengthInBytes > srcLengthBytes)
throw new ArgumentException($"Cannot copy {lengthInBytes} bytes from sourceByteOffset {sourceByteOffset} and BarracudaArray of {srcLengthBytes} num bytes.");
if (destinationByteOffset+lengthInBytes > destinationArray.Length)
throw new ArgumentException($"Cannot copy {lengthInBytes} bytes to destinationByteOffset {destinationByteOffset} and byte[] array of {destinationArray.Length} num bytes.");
fixed (void* dstPtr = &destinationArray[destinationByteOffset])
{
void* srcPtr = (byte*)sourceArray.RawPtr + sourceByteOffset;
UnsafeUtility.MemCpy(dstPtr, srcPtr, lengthInBytes);
}
}
public static unsafe void BlockCopy(
byte[] sourceArray,
int sourceByteOffset,
BarracudaArray destinationArray,
int destinationByteOffset,
int lengthInBytes)
{
if (lengthInBytes == 0)
return;
if (lengthInBytes < 0)
lengthInBytes = sourceArray.Length;
if (sourceByteOffset+lengthInBytes > sourceArray.Length)
throw new ArgumentException($"Cannot copy {lengthInBytes} bytes from sourceByteOffset {sourceByteOffset} and byte[] array of {sourceArray.Length} num bytes.");
var fullDestPaddedSizeInByte = LengthWithPaddingForGPUCopy(destinationArray.Type, destinationArray.Length) * DataItemSize(destinationArray.Type);
if (destinationByteOffset+lengthInBytes > fullDestPaddedSizeInByte)
throw new ArgumentException($"Cannot copy {lengthInBytes} bytes to destinationByteOffset {destinationByteOffset} and byte[] array of {destinationArray.Length} num bytes.");
void* dstPtr = (byte*)destinationArray.RawPtr + destinationByteOffset;
fixed (void* srcPtr = &sourceArray[sourceByteOffset])
{
UnsafeUtility.MemCpy(dstPtr, srcPtr, lengthInBytes);
}
}
public static void Copy(
float[] sourceArray,
int sourceIndex,
BarracudaArray destinationArray,
long destinationIndex,
int length)
{
Copy(sourceArray, sourceIndex, destinationArray, (int)destinationIndex, length);
}
public static unsafe void Copy(
float[] sourceArray,
int sourceIndex,
BarracudaArray destinationArray,
int destinationIndex,
int length)
{
if (length < 0)
length = sourceArray.Length;
if (length == 0)
return;
if (sourceIndex+length > sourceArray.Length)
throw new ArgumentException($"Cannot copy {length} element from sourceIndex {sourceIndex} and Barracuda array of length {sourceArray.Length}.");
if (destinationIndex+length > destinationArray.Length)
throw new ArgumentException($"Cannot copy {length} element to sourceIndex {destinationIndex} and Barracuda array of length {destinationArray.Length}.");
//Same type we can do a memcopy
if (destinationArray.m_DataType == DataType.Float)
{
fixed (void* srcPtr = &sourceArray[sourceIndex])
{
int itemSize = DataItemSize(destinationArray.m_DataType);
void* dstPtr = (byte*)destinationArray.RawPtr + destinationIndex * itemSize;
UnsafeUtility.MemCpy(dstPtr, srcPtr, length * itemSize);
}
}
else//different type, we need to iterate and cast
{
for (var i=0; i < length; ++i)
{
//this will use float as intermediate/common representation
destinationArray[destinationIndex+i] = sourceArray[sourceIndex+i];
}
}
}
#endregion
}
static class BarracudaArrayExtensionHelper
{
public static void CopyToBarracudaArray(this float[] sourceArray, BarracudaArray destinationArray, int destinationIndex)
{
BarracudaArray.Copy(sourceArray, 0, destinationArray, destinationIndex, sourceArray.Length);
}
public static void CopyToBarracudaArray(this float[] sourceArray, BarracudaArray destinationArray, long destinationIndex)
{
BarracudaArray.Copy(sourceArray, 0, destinationArray, (int)destinationIndex, sourceArray.Length);
}
}
} // namespace Barracuda