void CorrectConstantsForBroadCast(ref Model nhwc)
{
List <Layer> correctedConstants = new List <Layer>();
for (int l = 0; l < nhwc.layers.Count; l++)
{
Layer layer = nhwc.layers[l];
for (int i = 0; i < layer.inputs.Length; i++)
{
var input = layer.inputs[i];
if (!ModelAnalyzer.IsLayerBroacastable(layer))
{
continue;
}
if (!m_RanksByName.ContainsKey(input) || !m_RanksByName.ContainsKey(layer.name))
{
continue;
}
Layer inputLayer;
bool found = ModelAnalyzer.FindLayerByName(nhwc, input, out inputLayer);
if (!found)
{
continue;
}
if (!ModelOptimizer.IsLayerConstant(inputLayer))
{
continue;
}
if (m_RanksByName[input] < 1 || m_RanksByName[input] == m_RanksByName[layer.name])
{
continue;
}
if (inputLayer.weights.Length == 1)
{
continue;
}
if (m_RanksByName[input] > m_RanksByName[layer.name])
{
throw new Exception($"constant must be lower rank than input for broadcast to work, TODO add transpose before input");
}
Layer correctedConstLayer = new Layer("c_" + inputLayer.name + "For_" + layer.name, Layer.Type.Load);
// transpose dataset
correctedConstLayer.datasets = new Layer.DataSet[1];
Array.Copy(inputLayer.datasets, correctedConstLayer.datasets, inputLayer.datasets.Length);
correctedConstLayer.datasets[0].name = correctedConstLayer.name;
correctedConstLayer.weights = new float[inputLayer.weights.Length];
var X = inputLayer.DataSetToTensor(0);
int[] permutations = new[] { 0, 1, 2, 3 };
var rank = m_RanksByName[layer.name].Value;
switch (rank)
{
case 2:
// ONNX: 5,7 + 7
// Barracuda: 5,_,_,7 + 7,_,_,- => _,_,_,7
permutations = new[] { 1, 2, 3, 0 };
break;
case 3:
// ONNX: 5,7,3 + 3
// Barracuda: 5,_,3,7 + 3,_,_,_ => _,_,3,_
if (m_RanksByName[input] == 1)
{
permutations = new[] { 1, 2, 0, 3 }
}
;
// ONNX: 5,7,3 + 7,3
// Barracuda: 5,_,3,7 + 7,_,_,3 => _,_,3,7
else if (m_RanksByName[input] == 2)
{
permutations = new[] { 1, 2, 3, 0 }
}
;
break;
case 4:
// ONNX: 2,5,7,3 + 3
// Barracuda: 2,7,3,5 + 3,_,_,_ => _,_,3,_
if (m_RanksByName[input] == 1)
{
permutations = new[] { 1, 2, 0, 3 }
}
;
// ONNX: 2,5,7,3 + 7,3
// Barracuda: 2,7,3,5 + 7,_,_,3 => _,7,3,_
else if (m_RanksByName[input] == 2)
{
permutations = new[] { 2, 0, 1, 3 }
}
;
// ONNX: 2,5,7,3 + 5,7,3
// Barracuda: 2,7,3,5 + 5,_,3,7 => _,7,3,5
else if (m_RanksByName[input] == 3)
{
permutations = new[] { 1, 3, 2, 0 }
}
;
break;
}
if (m_isModelExportedFromNCHW && (m_layersChannelOrder[layer.name] == LayoutTransposeRemovalHelper.ChannelsOrder.NHWC))
{
switch (rank)
{
case 2:
// ONNX: 5,7 + 7
// Barracuda: 5,_,_,7 + 7,_,_,- => _,_,_,7
permutations = new[] { 1, 2, 3, 0 };
break;
case 3:
// ONNX: 5,7,3 + 3
// Barracuda: 5,_,7,3 + 3,_,_,_ => _,_,_,3
if (m_RanksByName[input] == 1)
{
permutations = new[] { 1, 2, 3, 0 }
}
;
// ONNX: 5,7,3 + 7,3
// Barracuda: 5,_,7,3 + 7,_,_,3 => _,_,7,3
else if (m_RanksByName[input] == 2)
{
permutations = new[] { 2, 3, 0, 1 }
}
;
break;
case 4:
// ONNX: 2,5,7,3 + 3
// Barracuda: 2,5,7,3 + 3,_,_,_ => _,_,_,3
if (m_RanksByName[input] == 1)
{
permutations = new[] { 1, 2, 3, 0 }
}
;
// ONNX: 2,5,7,3 + 7,3
// Barracuda: 2,5,7,3 + 7,_,_,3 => _,_,7,3,
else if (m_RanksByName[input] == 2)
{
permutations = new[] { 2, 3, 0, 1 }
}
;
// ONNX: 2,5,7,3 + 5,7,3
// Barracuda: 2,5,7,3 + 5,_,7,3 => _,5,7,3
else if (m_RanksByName[input] == 3)
{
permutations = new[] { 1, 0, 2, 3 }
}
;
break;
}
}
var O = m_Ops.Transpose(X, permutations);
correctedConstLayer.ApplyTensorToDataSet(O, 0);
correctedConstants.Add(correctedConstLayer);
layer.inputs[i] = correctedConstLayer.name;
}
nhwc.layers[l] = layer;
}
foreach (var l in correctedConstants)
{
nhwc.layers.Insert(0, l);
}
}
}
}
// works on IRModel
public bool InferAllLayersChannelOrder(Model model, out Dictionary <string, ChannelsOrder> layerChannelOrder)
{
layerChannelOrder = new Dictionary <string, ChannelsOrder>();
IDictionary <string, TensorShape?> shapesByName = new Dictionary <string, TensorShape?>();
IDictionary <string, int?> ranksByName = new Dictionary <string, int?>();
foreach (var i in model.inputs)
{
ranksByName[i.name] = i.rank;
if (!ModelAnalyzer.IsInputShapeAcceptablyKnowForShapeInference(i))
{
continue;
}
shapesByName[i.name] = new TensorShape(i.shape);
}
IRShapeInferenceAndConstantFusing shapeInferencePass = new IRShapeInferenceAndConstantFusing();
shapeInferencePass.InferAllShapes(model, ref shapesByName, ref ranksByName);
// flood-fill approach: NCHW layout is propagated from NCHW ops
// * onnx-nchw ops are flagged as being native nchw
// * nchw layout is propagated to upstream and downstream nodes
// foreach node:
// take layout being propagated to
// if T or T-1 flip layout depending on upstream/downstream direction
// - stop if layout is the same as previously propagated
// - native nchw layout has priority
Queue <(string, ChannelsOrder, FlowDirection)> layersToInferLayout = new Queue <(string, ChannelsOrder, FlowDirection)>();
for (int l = 0; l < model.layers.Count; l++)
{
var layer = model.layers[l];
if (!IsLayerNecessarilyNCHWOnnx(layer))
{
continue;
}
layersToInferLayout.Enqueue((layer.name, ChannelsOrder.NativeNCHW, FlowDirection.Seed));
}
while (layersToInferLayout.Any())
{
(string, ChannelsOrder, FlowDirection)layerData = layersToInferLayout.Dequeue();
string name = layerData.Item1;
ChannelsOrder deducedChannelOrder = layerData.Item2;
// 0: in-place native
// 1: downstream
// 2: upstream
FlowDirection flowDirection = layerData.Item3;
if (!layerChannelOrder.ContainsKey(name))
{
layerChannelOrder[name] = deducedChannelOrder;
}
else if (deducedChannelOrder == layerChannelOrder[name])
{
continue;
}
else if (layerChannelOrder[name] == ChannelsOrder.NativeNCHW)
{
continue;
}
// heuristic to stop ping-pong loop, prioritize NHWC over NCHW as it implies less transposes
// if incoming is NativeNCHW always propagate that
// TODO: count # of transpose swaps
else if (layerChannelOrder[name] == ChannelsOrder.NHWC && deducedChannelOrder != ChannelsOrder.NativeNCHW)
{
continue;
}
Layer layer;
bool found = ModelAnalyzer.FindLayerByName(model, name, out layer);
if (IsLayerChangingLayoutToNHWC(layer, shapesByName, ranksByName))
{
// NCHW -> T -> NHWC
if (((deducedChannelOrder == ChannelsOrder.NCHW) || (deducedChannelOrder == ChannelsOrder.NativeNCHW)) && (flowDirection == FlowDirection.Downstream))
{
deducedChannelOrder = ChannelsOrder.TransposeToNHWC;
}
// NCHW <- T <- NHWC
else if ((deducedChannelOrder == ChannelsOrder.NHWC) && (flowDirection == FlowDirection.Upstream))
{
deducedChannelOrder = ChannelsOrder.TransposeToNHWC;
}
}
else if (IsLayerChangingLayoutToNCHW(layer, shapesByName, ranksByName))
{
// NHWC -> T-1 -> NCHW
if ((deducedChannelOrder == ChannelsOrder.NHWC) && (flowDirection == FlowDirection.Downstream))
{
deducedChannelOrder = ChannelsOrder.TransposeToNCHW;
}
// NHWC <- T-1 <- NCHW
else if (((deducedChannelOrder == ChannelsOrder.NCHW) || (deducedChannelOrder == ChannelsOrder.NativeNCHW)) && (flowDirection == FlowDirection.Upstream))
{
deducedChannelOrder = ChannelsOrder.TransposeToNCHW;
}
}
if ((deducedChannelOrder == ChannelsOrder.TransposeToNCHW || deducedChannelOrder == ChannelsOrder.TransposeToNHWC) && (deducedChannelOrder == layerChannelOrder[name]))
{
continue;
}
layerChannelOrder[name] = deducedChannelOrder;
foreach (var input in layer.inputs)
{
if (deducedChannelOrder == ChannelsOrder.TransposeToNCHW)
{
layersToInferLayout.Enqueue((input, ChannelsOrder.NHWC, FlowDirection.Upstream));
}
else if (deducedChannelOrder == ChannelsOrder.TransposeToNHWC)
{
layersToInferLayout.Enqueue((input, ChannelsOrder.NCHW, FlowDirection.Upstream));
}
else
{
layersToInferLayout.Enqueue((input, deducedChannelOrder, FlowDirection.Upstream));
}
}
var outputs = ModelAnalyzer.FindLayerOutputs(model, layer.name);
foreach (var output in outputs)
{
if (deducedChannelOrder == ChannelsOrder.TransposeToNCHW)
{
layersToInferLayout.Enqueue((output, ChannelsOrder.NCHW, FlowDirection.Downstream));
}
else if (deducedChannelOrder == ChannelsOrder.TransposeToNHWC)
{
layersToInferLayout.Enqueue((output, ChannelsOrder.NHWC, FlowDirection.Downstream));
}
else
{
layersToInferLayout.Enqueue((output, deducedChannelOrder, FlowDirection.Downstream));
}
}
}
bool modelExportedASNHWC = false;
foreach (string key in layerChannelOrder.Keys.ToList())
{
var value = layerChannelOrder[key];
if (value == ChannelsOrder.NativeNCHW)
{
layerChannelOrder[key] = ChannelsOrder.NCHW;
}
if (value == ChannelsOrder.NHWC)
{
modelExportedASNHWC = true;
}
}
return(modelExportedASNHWC);
}
int[] GetPermutationForBroadcast(int targetRank, int rank, bool isNHWC = false)
{
int[] permutations = new[] { 0, 1, 2, 3 };
if (rank == 0 || targetRank == 1)
{
return(permutations);
}
switch (targetRank)
{
case 2:
// ONNX: 5,7 + 7
// Barracuda: 5,_,_,7 + 7,_,_,- => _,_,_,7
permutations = new[] { 1, 2, 3, 0 };
break;
case 3:
// ONNX: 5,7,3 + 3
// Barracuda: 5,_,3,7 + 3,_,_,_ => _,_,3,_
if (rank == 1)
{
permutations = new[] { 1, 2, 0, 3 }
}
;
// ONNX: 5,7,3 + 7,3
// Barracuda: 5,_,3,7 + 7,_,_,3 => _,_,3,7
else if (rank == 2)
{
permutations = new[] { 1, 2, 3, 0 }
}
;
break;
case 4:
// ONNX: 2,5,7,3 + 3
// Barracuda: 2,7,3,5 + 3,_,_,_ => _,_,3,_
if (rank == 1)
{
permutations = new[] { 1, 2, 0, 3 }
}
;
// ONNX: 2,5,7,3 + 7,3
// Barracuda: 2,7,3,5 + 7,_,_,3 => _,7,3,_
else if (rank == 2)
{
permutations = new[] { 1, 0, 3, 2 }
}
;
// ONNX: 2,5,7,3 + 5,7,3
// Barracuda: 2,7,3,5 + 5,_,3,7 => _,7,3,5
else if (rank == 3)
{
permutations = new[] { 1, 3, 2, 0 }
}
;
break;
}
if (isNHWC)
{
switch (targetRank)
{
case 2:
// ONNX: 5,7 + 7
// Barracuda: 5,_,_,7 + 7,_,_,- => _,_,_,7
permutations = new[] { 1, 2, 3, 0 };
break;
case 3:
// ONNX: 5,7,3 + 3
// Barracuda: 5,_,7,3 + 3,_,_,_ => _,_,_,3
if (rank == 1)
{
permutations = new[] { 1, 2, 3, 0 }
}
;
// ONNX: 5,7,3 + 7,3
// Barracuda: 5,_,7,3 + 7,_,_,3 => _,_,7,3
else if (rank == 2)
{
permutations = new[] { 1, 2, 0, 3 }
}
;
break;
case 4:
// ONNX: 2,5,7,3 + 3
// Barracuda: 2,5,7,3 + 3,_,_,_ => _,_,_,3
if (rank == 1)
{
permutations = new[] { 1, 2, 3, 0 }
}
;
// ONNX: 2,5,7,3 + 7,3
// Barracuda: 2,5,7,3 + 7,_,_,3 => _,_,7,3,
else if (rank == 2)
{
permutations = new[] { 1, 2, 0, 3 }
}
;
// ONNX: 2,5,7,3 + 5,7,3
// Barracuda: 2,5,7,3 + 5,_,7,3 => _,5,7,3
else if (rank == 3)
{
permutations = new[] { 1, 0, 2, 3 }
}
;
break;
}
}
return(permutations);
}
void CorrectConstantsForBroadCast(ref Model nhwc)
{
List <Layer> correctedConstants = new List <Layer>();
for (int l = 0; l < nhwc.layers.Count; l++)
{
Layer layer = nhwc.layers[l];
for (int i = 0; i < layer.inputs.Length; i++)
{
var input = layer.inputs[i];
if (!ModelAnalyzer.IsLayerBroacastable(layer))
{
continue;
}
if (!m_RanksByName.ContainsKey(input) || !m_RanksByName.ContainsKey(layer.name))
{
continue;
}
Layer inputLayer;
bool found = ModelAnalyzer.FindLayerByName(nhwc, input, out inputLayer);
if (!found)
{
continue;
}
if (!ModelOptimizer.IsLayerConstant(inputLayer))
{
continue;
}
if (m_RanksByName[input] < 1 || m_RanksByName[input] == m_RanksByName[layer.name])
{
continue;
}
if (inputLayer.weights.Length == 1)
{
continue;
}
if (m_RanksByName[input] > m_RanksByName[layer.name])
{
throw new Exception($"constant must be lower rank than input for broadcast to work, TODO add transpose before input");
}
Layer correctedConstLayer = new Layer("c_" + inputLayer.name + "For_" + layer.name, Layer.Type.Load);
// transpose dataset
correctedConstLayer.datasets = new Layer.DataSet[1];
Array.Copy(inputLayer.datasets, correctedConstLayer.datasets, inputLayer.datasets.Length);
correctedConstLayer.datasets[0].name = correctedConstLayer.name;
correctedConstLayer.weights = new BarracudaArray(inputLayer.weights.Length);
var X = inputLayer.DataSetToTensor(0);
var rank = m_RanksByName[layer.name].Value;
var inputRank = m_RanksByName[input].Value;
int[] permutations = GetPermutationForBroadcast(rank, inputRank, (m_isModelExportedFromNHWC && (m_layersChannelOrder[layer.name] == LayoutTransposeRemovalHelper.ChannelsOrder.NHWC)));
var O = m_Ops.Transpose(X, permutations);
correctedConstLayer.ApplyTensorToDataSet(O, 0);
O.Dispose();
X.Dispose();
correctedConstants.Add(correctedConstLayer);
layer.inputs[i] = correctedConstLayer.name;
}
nhwc.layers[l] = layer;
}
foreach (var l in correctedConstants)
{
nhwc.layers.Insert(0, l);
}
}
void CorrectDynamicInputsForBroadCast(ref Model nhwc)
{
// insert transposes before broadcastalbe ops
for (int l = 0; l < nhwc.layers.Count; l++)
{
Layer layer = nhwc.layers[l];
if (!ModelAnalyzer.IsLayerBroacastable(layer))
{
continue;
}
if (!m_RanksByName.ContainsKey(layer.name) || m_RanksByName[layer.name] == null)
{
continue;
}
int maxRank = m_RanksByName[layer.name].Value;
if (maxRank <= 1)
{
continue;
}
for (int i = 0; i < layer.inputs.Length; i++)
{
string input = layer.inputs[i];
if (!m_RanksByName.ContainsKey(input) || m_RanksByName[input] == null)
{
continue;
}
int inputRank = m_RanksByName[input].Value;
if (inputRank < 1 || inputRank == maxRank)
{
continue;
}
int[] permutations = GetPermutationForBroadcast(maxRank, inputRank, (m_isModelExportedFromNHWC && (m_layersChannelOrder[layer.name] == LayoutTransposeRemovalHelper.ChannelsOrder.NHWC)));
Layer transpose = new Layer("transpose_forbroadcast_" + layer.name + "_" + input, Layer.Type.Transpose);
transpose.inputs = new[] { input };
transpose.pool = permutations;
nhwc.layers[l].inputs[i] = transpose.name;
nhwc.layers.Insert(l, transpose);
l += 1;
}
}
}
}
}