public MemoryAllocation GetOrAllocateMainMemory(OutputConnector output)
{
if (!MainMemoryMap.TryGetValue(output, out var alloc))
{
uint elementSize;
if (output.Owner.GetType().Name.StartsWith("Quantized"))
{
elementSize = 1;
}
else
{
switch (output.Owner)
{
case K210Conv2d _:
case K210AddPadding _:
case K210RemovePadding _:
case K210Upload _:
case layers.Quantize _:
case layers.Requantize _:
elementSize = 1;
break;
default:
elementSize = 4;
break;
}
}
var dimensions = output.Dimensions;
alloc = new MemoryAllocation(MainMemoryAllocator.Allocate((uint)dimensions.GetSize() * elementSize));
MainMemoryMap.Add(output, alloc);
}
else
{
alloc.Node.AddRef();
}
return(alloc);
}
public MemoryAllocation GetOrAllocateKPUMemory(OutputConnector output)
{
if (!KPUMemoryMap.TryGetValue(output, out var alloc))
{
var dimensions = output.Dimensions;
(var groups, var rowLength, _) = K210Helper.GetRowLayout(dimensions[3]);
var oneLineChannels = Math.Min(dimensions[1], groups);
var blocks = (int)Math.Ceiling(dimensions[1] / (double)oneLineChannels);
if (dimensions[1] == 921)
{
}
var size = rowLength * dimensions[2] * blocks;
alloc = new MemoryAllocation(KPUMemoryAllocator.Allocate((uint)size));
KPUMemoryMap.Add(output, alloc);
}
else
{
alloc.Node.AddRef();
}
return(alloc);
}