diff --git a/OnnxBridge/LLAMA/sytorchBackendRep.py b/OnnxBridge/LLAMA/sytorchBackendRep.py index 36ae78e9..5a237651 100644 --- a/OnnxBridge/LLAMA/sytorchBackendRep.py +++ b/OnnxBridge/LLAMA/sytorchBackendRep.py @@ -22,6 +22,7 @@ def func_call(node, value_info): "GlobalAveragePool": "GlobalAvgPool2D", "Add": "add", "ConvTranspose": "ConvTranspose3D", + "Slice": "Slice", } return func_map[node.op_type] @@ -65,6 +66,7 @@ def inputs_to_take(node): "BatchNormalization": 1, "GlobalAveragePool": 1, "ConvTranspose": 1, + "Slice": 1, } return tmp_dict[node] @@ -261,7 +263,7 @@ def prepare_export(program, var_dict, value_info, mode, scale, bitlength, backen # Check nodes for assertions and modifications for node in program: func = getattr(OnnxNode, node.op_type) - func(node) + func(node, value_info) # Start CPP program number_of_nodes = 0 diff --git a/OnnxBridge/LLAMA/sytorch_func_calls.py b/OnnxBridge/LLAMA/sytorch_func_calls.py index e71d20e6..091068f8 100644 --- a/OnnxBridge/LLAMA/sytorch_func_calls.py +++ b/OnnxBridge/LLAMA/sytorch_func_calls.py @@ -226,6 +226,16 @@ def Flatten(cls, attributes, inputs, outputs, value_info, var_dict, mode, indent logger.debug("Inside Flatten function call.") return str(f"{' ' * indent}new Flatten();") + @classmethod + def Slice(cls, attributes, inputs, outputs, value_info, var_dict, mode, indent): + logger.debug("Inside Slice function call.") + starts = attributes["starts"] + ends = attributes["ends"] + axes = attributes["axes"] + steps = attributes["steps"] + + return str(f"{' ' * indent}new Slice({starts}, {ends}, {axes}, {steps});") + @classmethod def Reshape(cls, attributes, inputs, outputs, value_info, var_dict, mode, indent): logger.debug("Inside Reshape function call.") diff --git a/OnnxBridge/backend.py b/OnnxBridge/backend.py index 0f5781d8..544c3e80 100644 --- a/OnnxBridge/backend.py +++ b/OnnxBridge/backend.py @@ -133,6 +133,7 @@ def is_compatible(cls, model, backend, device: str = "2PC", **kwargs): "GlobalAveragePool", "Add", "ConvTranspose", + "Slice", ] implemented_secfloat = [ "Relu", diff --git a/OnnxBridge/utils/onnx_nodes.py b/OnnxBridge/utils/onnx_nodes.py index 5842482f..76001c2b 100644 --- a/OnnxBridge/utils/onnx_nodes.py +++ b/OnnxBridge/utils/onnx_nodes.py @@ -8,22 +8,22 @@ class OnnxNode: """ @classmethod - def input(cls, node): + def input(cls, node, value_info): assert isinstance(node, Input) logger.debug("Input is OK!") @classmethod - def output(cls, node): + def output(cls, node, value_info): assert isinstance(node, Output) logger.debug("Output is OK!") @classmethod - def Cast(cls, node): + def Cast(cls, node, value_info): assert len(node.inputs) == 1 logger.debug("Cast is OK!") @classmethod - def Pad(cls, node): + def Pad(cls, node, value_info): assert node.attrs["mode"] == "constant" if node.opset_version >= 11: @@ -43,52 +43,52 @@ def Pad(cls, node): logger.debug("Pad is OK! (with possible modifications)") @classmethod - def Concat(cls, node): + def Concat(cls, node, value_info): # Nothing to assert yet logger.debug("Concat is OK! (No assertions)") @classmethod - def Sigmoid(cls, node): + def Sigmoid(cls, node, value_info): assert len(node.inputs) == 1 logger.debug("Sigmoid is OK!") @classmethod - def HardSigmoid(cls, node): + def HardSigmoid(cls, node, value_info): assert len(node.inputs) == 1 logger.debug("Hard Sigmoid is OK!") @classmethod - def Relu(cls, node): + def Relu(cls, node, value_info): assert len(node.inputs) == 1 logger.debug("Relu is OK!") @classmethod - def Div(cls, node): + def Div(cls, node, value_info): # todo is div there or not? in athos it takes one input pass @classmethod - def Add(cls, node): + def Add(cls, node, value_info): assert len(node.inputs) == 2 logger.debug("Add is OK!") @classmethod - def Sub(cls, node): + def Sub(cls, node, value_info): assert len(node.inputs) == 2 logger.debug("Sub is OK!") @classmethod - def Mul(cls, node): + def Mul(cls, node, value_info): assert len(node.inputs) == 2 logger.debug("Mul is OK!") @classmethod - def Gather(cls, node): + def Gather(cls, node, value_info): # Nothing to assert yet logger.debug("Concat is OK! (No assertions)") @classmethod - def Gemm(cls, node): + def Gemm(cls, node, value_info): # todo transpose done separately in gemm if "alpha" not in node.attrs: node.attrs["alpha"] = 1.0 @@ -100,22 +100,22 @@ def Gemm(cls, node): node.attrs["transB"] = 0 @classmethod - def Constant(cls, node): + def Constant(cls, node, value_info): # Nothing to assert yet logger.debug("Concat is OK! (No assertions)") @classmethod - def Transpose(cls, node): + def Transpose(cls, node, value_info): assert len(node.inputs) == 1 logger.debug("Transpose is OK!") @classmethod - def Split(cls, node): + def Split(cls, node, value_info): node.inputs = node.inputs[:1] logger.debug("Split is OK! (with possible modifications)") @classmethod - def ReduceMean(cls, node): + def ReduceMean(cls, node, value_info): keepdims = node.attrs["keepdims"] axes = node.attrs["axes"] @@ -128,56 +128,79 @@ def ReduceMean(cls, node): logger.debug("ReduceMean is OK! (with possible modifications)") @classmethod - def MatMul(cls, node): + def MatMul(cls, node, value_info): # todo transpose and mul assert len(node.inputs) == 2 logger.debug("MatMul is OK!") @classmethod - def BatchNormalization(cls, node): + def BatchNormalization(cls, node, value_info): assert len(node.inputs) == 5 node.inputs = node.inputs[:3] logger.debug("Batch Normalization is OK! (with possible modifications)") @classmethod - def Unsqueeze(cls, node): + def Unsqueeze(cls, node, value_info): pass @classmethod - def Reshape(cls, node): + def Reshape(cls, node, value_info): pass @classmethod - def Flatten(cls, node): + def Flatten(cls, node, value_info): assert len(node.inputs) == 1 @classmethod - def Conv(cls, node): + def Slice(cls, node, value_info): + # input: data, starts, ends, axes, steps + assert len(node.inputs) >= 3 + starts = value_info[node.inputs[1]][1][0] + ends = value_info[node.inputs[2]][1][0] + assert starts == ends + axes = 0 + steps = 0 + if len(node.inputs) > 3: + axes = value_info[node.inputs[3]][1][0] + assert axes == starts + if len(node.inputs) > 4: + steps = value_info[node.inputs[4]][1][0] + assert steps == starts + + node.attrs["starts"] = starts + node.attrs["ends"] = ends + node.attrs["axes"] = axes + node.attrs["steps"] = steps + print(node) + logger.debug("Slice is OK!") + + @classmethod + def Conv(cls, node, value_info): pass @classmethod - def MaxPool(cls, node): + def MaxPool(cls, node, value_info): pass @classmethod - def AveragePool(cls, node): + def AveragePool(cls, node, value_info): pass @classmethod - def GlobalAveragePool(cls, node): + def GlobalAveragePool(cls, node, value_info): pass @classmethod - def ConvTranspose(cls, node): + def ConvTranspose(cls, node, value_info): pass @classmethod - def LeakyRelu(cls, node): + def LeakyRelu(cls, node, value_info): if "alpha" not in node.attributes: node.attributes["alpha"] = 0.01 @classmethod - def Tanh(cls, node): + def Tanh(cls, node, value_info): assert len(node.inputs) == 1 assert len(node.outputs) == 1 # we can print the node at this step and get info on all node parameters diff --git a/sytorch/include/sytorch/layers/layers.h b/sytorch/include/sytorch/layers/layers.h index 72f9a181..fb646634 100644 --- a/sytorch/include/sytorch/layers/layers.h +++ b/sytorch/include/sytorch/layers/layers.h @@ -109,6 +109,7 @@ class Layer { virtual TensorRef getweights() { return TensorRef(nullptr, 0); }; virtual TensorRef getbias() { return TensorRef(nullptr, 0); }; + virtual std::vector> get_params() { return {TensorRef(nullptr, 0)}; }; virtual std::vector get_output_dims(const std::vector> &inShapes) = 0; virtual void setBackend(Backend *b) { @@ -438,6 +439,335 @@ class Flatten : public Layer { } }; +template +class Slice : public Layer +{ +public: + Tensor1D starts, ends, axes, steps; + i64 params; + + Slice(i64 start, i64 end, i64 axis, i64 step) : Layer("Slice"), starts(start), ends(end), axes(axis), steps(step) + { + params = start; + } + + void _resize(const std::vector> &shapes) + { + always_assert(shapes.size() == 1); + auto &shape = shapes[0]; + always_assert(shape.size() >= 2); + } + + void _forward(Tensor &a) + { + std::unordered_map axis_map; + for (u64 i = 0; i < axes.size(); i++) + { + axis_map[axes.data[i] / (1 << this->scale)] = i; + } + + if (a.shape.size() == 2) + { + auto a_2d = a.as_2d(); + auto out_2d = this->activation.as_2d(); + u64 d1 = a.shape[0]; + u64 d2 = a.shape[1]; + for (u64 i = 0, i_o = 0; i < d1;) + { + if (axis_map.find(0) != axis_map.end()) + { + i64 start = starts.data[axis_map[0]] / (1 << this->scale); + i64 end = ends.data[axis_map[0]] / (1 << this->scale); + i64 step = steps.data[axis_map[0]] / (1 << this->scale); + if (i < start or i >= end or (i - start) % step != 0) + { + i++; + continue; + } + } + + for (u64 j = 0, j_o = 0; j < d2;) + { + if (axis_map.find(1) != axis_map.end()) + { + i64 start = starts.data[axis_map[1]] / (1 << this->scale); + i64 end = ends.data[axis_map[1]] / (1 << this->scale); + i64 step = steps.data[axis_map[1]] / (1 << this->scale); + if (j < start or j >= end or (j - start) % step != 0) + { + j++; + continue; + } + } + out_2d(i_o, j_o) = a_2d(i, j); + j++; + j_o++; + } + i++; + i_o++; + } + } + else if (a.shape.size() == 4) + { + auto a_4d = a.as_4d(); + auto out_4d = this->activation.as_4d(); + u64 d1 = a.shape[0]; + u64 d2 = a.shape[1]; + u64 d3 = a.shape[2]; + u64 d4 = a.shape[3]; + + for (u64 i = 0, i_o = 0; i < d1;) + { + if (axis_map.find(0) != axis_map.end()) + { + i64 start = starts.data[axis_map[0]] / (1 << this->scale); + i64 end = ends.data[axis_map[0]] / (1 << this->scale); + i64 step = steps.data[axis_map[0]] / (1 << this->scale); + if (i < start or i >= end or (i - start) % step != 0) + { + i++; + continue; + } + } + for (u64 j = 0, j_o = 0; j < d2;) + { + if (axis_map.find(1) != axis_map.end()) + { + i64 start = starts.data[axis_map[1]] / (1 << this->scale); + i64 end = ends.data[axis_map[1]] / (1 << this->scale); + i64 step = steps.data[axis_map[1]] / (1 << this->scale); + if (j < start or j >= end or (j - start) % step != 0) + { + j++; + continue; + } + } + for (u64 k = 0, k_o = 0; k < d3;) + { + if (axis_map.find(2) != axis_map.end()) + { + i64 start = starts.data[axis_map[2]] / (1 << this->scale); + i64 end = ends.data[axis_map[2]] / (1 << this->scale); + i64 step = steps.data[axis_map[2]] / (1 << this->scale); + if (k < start or k >= end or (k - start) % step != 0) + { + k++; + continue; + } + } + for (u64 l = 0, l_o = 0; l < d4;) + { + if (axis_map.find(3) != axis_map.end()) + { + i64 start = starts.data[axis_map[3]] / (1 << this->scale); + i64 end = ends.data[axis_map[3]] / (1 << this->scale); + i64 step = steps.data[axis_map[3]] / (1 << this->scale); + if (l < start or l >= end or (l - start) % step != 0) + { + l++; + continue; + } + } + out_4d(i_o, j_o, k_o, l_o) = a_4d(i, j, k, l); + l++; + l_o++; + } + k++; + k_o++; + } + j++; + j_o++; + } + i++; + i_o++; + } + } + else if (a.shape.size() == 5) + { + auto a_5d = a.as_5d(); + auto out_5d = this->activation.as_5d(); + u64 d1 = a.shape[0]; + u64 d2 = a.shape[1]; + u64 d3 = a.shape[2]; + u64 d4 = a.shape[3]; + u64 d5 = a.shape[4]; + + for (u64 i = 0, i_o = 0; i < d1;) + { + if (axis_map.find(0) != axis_map.end()) + { + i64 start = starts.data[axis_map[0]] / (1 << this->scale); + i64 end = ends.data[axis_map[0]] / (1 << this->scale); + i64 step = steps.data[axis_map[0]] / (1 << this->scale); + if (i < start or i >= end or (i - start) % step != 0) + { + i++; + continue; + } + } + for (u64 j = 0, j_o = 0; j < d2;) + { + if (axis_map.find(1) != axis_map.end()) + { + i64 start = starts.data[axis_map[1]] / (1 << this->scale); + i64 end = ends.data[axis_map[1]] / (1 << this->scale); + i64 step = steps.data[axis_map[1]] / (1 << this->scale); + if (j < start or j >= end or (j - start) % step != 0) + { + j++; + continue; + } + } + for (u64 k = 0, k_o = 0; k < d3;) + { + if (axis_map.find(2) != axis_map.end()) + { + i64 start = starts.data[axis_map[2]] / (1 << this->scale); + i64 end = ends.data[axis_map[2]] / (1 << this->scale); + i64 step = steps.data[axis_map[2]] / (1 << this->scale); + if (k < start or k >= end or (k - start) % step != 0) + { + k++; + continue; + } + } + for (u64 l = 0, l_o = 0; l < d4;) + { + if (axis_map.find(3) != axis_map.end()) + { + i64 start = starts.data[axis_map[3]] / (1 << this->scale); + i64 end = ends.data[axis_map[3]] / (1 << this->scale); + i64 step = steps.data[axis_map[3]] / (1 << this->scale); + if (l < start or l >= end or (l - start) % step != 0) + { + l++; + continue; + } + } + for (u64 m = 0, m_o = 0; m < d5;) + { + if (axis_map.find(4) != axis_map.end()) + { + i64 start = starts.data[axis_map[4]] / (1 << this->scale); + i64 end = ends.data[axis_map[4]] / (1 << this->scale); + i64 step = steps.data[axis_map[4]] / (1 << this->scale); + if (m < start or m >= end or (m - start) % step != 0) + { + m++; + continue; + } + } + out_5d(i_o, j_o, k_o, l_o, m_o) = a_5d(i, j, k, l, m); + m_o++; + m++; + } + l_o++; + l++; + } + k_o++; + k++; + } + j_o++; + j++; + } + i_o++; + i++; + } + } + else + { + throw std::runtime_error("Slice operator supported only for 2d, 4d and 5d tensors"); + } + } + + std::vector> get_params() + { + return {starts.ref(), ends.ref(), axes.ref(), steps.ref()}; + } + + std::vector get_output_dims(const std::vector> &inShapes) + { + always_assert(inShapes.size() == 1); + auto &inShape = inShapes[0]; + std::vector output_shape(inShape.size()); + // make copy of inShape + std::vector inShapeCopy(inShape.size()); + std::copy(inShape.begin(), inShape.end(), inShapeCopy.begin()); + + for (u64 i = 0; i < params; i++) + { + i64 start = starts.data[i] / (1 << this->scale); + i64 end = ends.data[i] / (1 << this->scale); + if (axes.size() == 0) + { + axes.resize(params); + for (u64 j = 0; j < params; j++) + { + axes.data[j] = j * (1 << this->scale); + } + } + if (steps.size() == 0) + { + steps.resize(params); + steps.fill(1 * (1 << this->scale)); + } + i64 axis = axes.data[i] / (1 << this->scale); + i64 step = steps.data[i] / (1 << this->scale); + assert(step != 0); + + // modifying axis as per nhwc format + if (axis >= 2) + { + axis = axis - 1; + axes.data[i] = axis * (1 << this->scale); + } + else if (axis == 1) + { + axis = inShapeCopy.size() - 1; + axes.data[i] = axis * (1 << this->scale); + } + + // handling case where start/end/axes/step is -ve + if (axis < 0) + { + axis = inShapeCopy.size() + axis; + axes.data[i] = axis * (1 << this->scale); + } + if (start < 0) + { + start = inShapeCopy[axis] + start; + starts.data[i] = start * (1 << this->scale); + } + if (end < 0) + { + end = inShapeCopy[axis] + end; + ends.data[i] = end * (1 << this->scale); + } + if (step < 0) + { + u64 lost = (end - start - 1) % (-step); + start = start + lost; + starts.data[i] = start * (1 << this->scale); + step = -step; + steps.data[i] = step * (1 << this->scale); + } + for (int i = 0; i < inShapeCopy.size(); i++) + { + if (i == axis) + { + output_shape[i] = ceil((end - start) / (double)step); + } + else + { + output_shape[i] = inShapeCopy[i]; + } + } + inShapeCopy = output_shape; + } + return output_shape; + } +}; + template class FC : public Layer { public: diff --git a/sytorch/include/sytorch/module.h b/sytorch/include/sytorch/module.h index fe26a08a..97459e4a 100644 --- a/sytorch/include/sytorch/module.h +++ b/sytorch/include/sytorch/module.h @@ -169,6 +169,18 @@ class SytorchModule { } wIdx += 4 * channel; } + else if (layer->name == "Slice") + { + auto params = layer->get_params(); + for (u64 i = 0; i < params.size(); i++) + { + for (u64 j = 0; j < params[i].size; j++) + { + params[i].data[j] = i64(floatWeights[wIdx + j] * (1LL << scale)); + } + wIdx += params[i].size; + } + } else { auto weights = layer->getweights(); for (u64 j = 0; j < weights.size; j++) { diff --git a/sytorch/include/sytorch/tensor.h b/sytorch/include/sytorch/tensor.h index fc2ac7bd..6892497c 100644 --- a/sytorch/include/sytorch/tensor.h +++ b/sytorch/include/sytorch/tensor.h @@ -349,6 +349,17 @@ class Tensor1D { } } + void resize(u64 s) + { + if (this->d1 == s) + { + return; + } + delete[] this->data; + this->d1 = s; + this->data = new T[s]; + } + ~Tensor1D() { delete[] this->data; }