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); }