From d3c829ea74fdee8a4e2707809cd017caa186fe9f Mon Sep 17 00:00:00 2001 From: da-gazzi <51782106+da-gazzi@users.noreply.github.com> Date: Thu, 20 Jul 2023 15:01:05 +0200 Subject: [PATCH 1/4] Devel (#8) * remove 'graphs' editing package as it's not used and buggy * remove remaining references to old 'graphs' package --- backends/cutie/__init__.py | 1 - backends/cutie/grrules/__init__.py | 23 - backends/cutie/grrules/ana/__init__.py | 28 - backends/cutie/grrules/ana/dporules.py | 1110 ------------------ backends/cutie/grrules/ana/folding.py | 127 -- backends/cutie/grrules/ana/lutactivation.py | 50 - backends/twn_accelerator/grrules/__init__.py | 23 - backends/twn_accelerator/grrules/dporules.py | 719 ------------ backends/twn_accelerator/grrules/folding.py | 178 --- backends/twn_accelerator/twn_accelerator.py | 4 +- editing/__init__.py | 1 - editing/graphs/__init__.py | 34 - editing/graphs/editor/__init__.py | 23 - editing/graphs/editor/editor.py | 237 ---- editing/graphs/graphs/__init__.py | 26 - editing/graphs/graphs/graphs.py | 276 ----- editing/graphs/graphs/modules.py | 113 -- editing/graphs/graphs/nodes.py | 99 -- editing/graphs/grrules/__init__.py | 42 - editing/graphs/grrules/dporules.py | 276 ----- editing/graphs/grrules/hlprrules.py | 417 ------- editing/graphs/grrules/seeker.py | 84 -- editing/graphs/traces/__init__.py | 50 - editing/graphs/traces/trace.py | 182 --- editing/graphs/utils/__init__.py | 23 - editing/graphs/utils/draw.py | 61 - 26 files changed, 2 insertions(+), 4205 deletions(-) delete mode 100644 backends/cutie/grrules/__init__.py delete mode 100644 backends/cutie/grrules/ana/__init__.py delete mode 100644 backends/cutie/grrules/ana/dporules.py delete mode 100644 backends/cutie/grrules/ana/folding.py delete mode 100644 backends/cutie/grrules/ana/lutactivation.py delete mode 100644 backends/twn_accelerator/grrules/__init__.py delete mode 100644 backends/twn_accelerator/grrules/dporules.py delete mode 100644 backends/twn_accelerator/grrules/folding.py delete mode 100755 editing/graphs/__init__.py delete mode 100644 editing/graphs/editor/__init__.py delete mode 100644 editing/graphs/editor/editor.py delete mode 100644 editing/graphs/graphs/__init__.py delete mode 100644 editing/graphs/graphs/graphs.py delete mode 100644 editing/graphs/graphs/modules.py delete mode 100644 editing/graphs/graphs/nodes.py delete mode 100644 editing/graphs/grrules/__init__.py delete mode 100644 editing/graphs/grrules/dporules.py delete mode 100644 editing/graphs/grrules/hlprrules.py delete mode 100644 editing/graphs/grrules/seeker.py delete mode 100644 editing/graphs/traces/__init__.py delete mode 100644 editing/graphs/traces/trace.py delete mode 100644 editing/graphs/utils/__init__.py delete mode 100644 editing/graphs/utils/draw.py diff --git a/backends/cutie/__init__.py b/backends/cutie/__init__.py index da596cd..d7c877e 100644 --- a/backends/cutie/__init__.py +++ b/backends/cutie/__init__.py @@ -19,5 +19,4 @@ # limitations under the License. # -from . import grrules from .cutie_export import convert_net, export_net diff --git a/backends/cutie/grrules/__init__.py b/backends/cutie/grrules/__init__.py deleted file mode 100644 index 1bc7168..0000000 --- a/backends/cutie/grrules/__init__.py +++ /dev/null @@ -1,23 +0,0 @@ -# -# __init__.py -# -# Author(s): -# Matteo Spallanzani -# -# Copyright (c) 2020-2021 ETH Zurich. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -from . import ana - diff --git a/backends/cutie/grrules/ana/__init__.py b/backends/cutie/grrules/ana/__init__.py deleted file mode 100644 index 587203a..0000000 --- a/backends/cutie/grrules/ana/__init__.py +++ /dev/null @@ -1,28 +0,0 @@ -# -# __init__.py -# -# Author(s): -# Matteo Spallanzani -# -# Copyright (c) 2020-2021 ETH Zurich. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -from .dporules import FoldANAConvBNANAActRule -from .dporules import FoldANAActANAConvBNANAActTypeARule -from .dporules import FoldANAActANAConvBNANAActTypeBRule -from .dporules import FoldANAActANALinearBNANAActTypeARule -from .dporules import FoldANAActANALinearBNANAActTypeBRule -from .dporules import FoldANAActLinearRule - diff --git a/backends/cutie/grrules/ana/dporules.py b/backends/cutie/grrules/ana/dporules.py deleted file mode 100644 index 133747d..0000000 --- a/backends/cutie/grrules/ana/dporules.py +++ /dev/null @@ -1,1110 +0,0 @@ -# -# dporules.py -# -# Author(s): -# Matteo Spallanzani -# -# Copyright (c) 2020-2021 ETH Zurich. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -import itertools -from collections import OrderedDict -import torch -import torch.nn as nn -import networkx as nx - -from .lutactivation import LUTActivation - -from .folding import fold_anaact_anaconv2d_bn2d_anaact, fold_anaact_analinear_bn1d_anaact -from quantlib.editing.graphs.graphs import Bipartite, PyTorchNode, __NODE_ID_FORMAT__ -from quantlib.editing.graphs.grrules.dporules import DPORule -from quantlib.editing.graphs.grrules import Seeker -import quantlib.editing.graphs as qg - -import quantlib.algorithms as qa - - -class FoldANAConvBNANAActRule(DPORule): - - def __init__(self, lut_entry_bits=16): - - self._lut_entry_bits = lut_entry_bits - - # Nodes of the interface - K_types = OrderedDict() - K_types.update({'HPin': qg.graphs.HelperInput.__name__}) - K_types.update({'HPTin': qg.graphs.HelperInputPrecisionTunnel.__name__}) - K_types = OrderedDict([('/'.join(['K-term', k]), v) for k, v in K_types.items()]) - - # Nodes in the core template graph - LK_types = OrderedDict() - LK_types.update({'ANAConv': qa.ana.ANAConv2d.__name__}) - LK_types.update({'BatchNorm': nn.BatchNorm2d.__name__}) - LK_types.update({'ANAActout': qa.ana.ANAActivation.__name__}) - LK_types = OrderedDict([('/'.join(['L-term', k]), v) for k, v in LK_types.items()]) - - # Nodes in the core replacement graph - RK_types = OrderedDict() - RK_types.update({'TWConv': nn.Conv2d.__name__}) - RK_types.update({'LUTAct': LUTActivation.__name__}) - RK_types = OrderedDict([('/'.join(['R-term', k]), v) for k, v in RK_types.items()]) - - K_node_IDs = list(K_types.keys()) - LK_node_IDs = list(LK_types.keys()) - RK_node_IDs = list(RK_types.keys()) - - # define the template graph L [L-term] - L_node_IDs = [K_node_IDs[0]] + LK_node_IDs + [K_node_IDs[-1]] - self.L = nx.DiGraph() - # Define arcs between nodes in full template graph - self.L.add_edges_from({(u, v) for u, v in zip(L_node_IDs[:-1], L_node_IDs[1:])}) - - # Here, graph is only operation nodes - # Necessary for seeker - nx.set_node_attributes(self.L, {vL: Bipartite.KERNEL for vL in set(self.L.nodes)}, 'bipartite') - nx.set_node_attributes(self.L, {**K_types, **LK_types}, 'type') - - # define the context (sub-)graph K [K-term] - VK = set(K_node_IDs) # precision tunnel nodes define the context graph - self.K = self.L.subgraph(VK) - - # define the template (sub-)graph L\K - VLK = set(self.L.nodes).difference(set(self.K.nodes)) - self.LK = self.L.subgraph(VLK) - - # define the replacement (sub-)graph R\K ["gluing" R\K to K yields the graph R, i.e., the R-term] - self.RK = nx.DiGraph() - self.RK.add_edges_from({(u, v) for u, v in zip(RK_node_IDs[:-1], RK_node_IDs[1:])}) - nx.set_node_attributes(self.RK, {vRK: Bipartite.KERNEL for vRK in set(self.RK.nodes)}, 'bipartite') - nx.set_node_attributes(self.RK, RK_types, 'type') - - # define the arcs that go from the vertices of K to those of R\K, and viceversa - E_K2RK = {(K_node_IDs[0], RK_node_IDs[0])} - E_RK2K = {(RK_node_IDs[-1], K_node_IDs[-1])} - E_K2RK2K = E_K2RK | E_RK2K - # disintegrate `E_K2RK` and `E_RK2K` along fibres to speed up rule application - # A fibre is kind of like fixing one argument of a two input one output function and looking at all possible outputs - self.F_K2RK = {vK: set(arc for arc in E_K2RK if arc[0] == vK) for vK in set(self.K.nodes)} - self.F_RK2K = {vK: set(arc for arc in E_RK2K if arc[1] == vK) for vK in set(self.K.nodes)} - - # since the GRR's L-term has been modified, rebuild the seeker - self.seeker = Seeker(self.L) - - # this machinery can generate always-new identifiers for different rule applications - self._counter = itertools.count() - - def _get_rule_count(self): - rule_count = ''.join(['FConvBNANA', __NODE_ID_FORMAT__.format(next(self._counter))]) - return rule_count - - def core(self, HI, g, nodes_dict): - - # generate the substitute (sub-)graph J\I - rule_count = self._get_rule_count() - g_RK2JI = {vRK: '_'.join([rule_count, vRK.replace('R-term/', '')]) for vRK in set(self.RK.nodes)} - JI = nx.relabel_nodes(self.RK, g_RK2JI, copy=True) - - # get pointers to the old modules; - # these pointers will enable two actions: - # 1. extracting the arguments required to perform the folding - # 2. extracting the parameters to instantiate the new modules - g_L2H = {vL: vH for vH, vL in g.items()} - mconv2d = nodes_dict[g_L2H['/'.join(['L-term', 'ANAConv'])]].nobj - mbn2d = nodes_dict[g_L2H['/'.join(['L-term', 'BatchNorm'])]].nobj - manaout = nodes_dict[g_L2H['/'.join(['L-term', 'ANAActout'])]].nobj - - # fold - tau, weight = fold_anaact_anaconv2d_bn2d_anaact(torch.Tensor([1.0]), - mconv2d.eps, mconv2d.weight_maybe_quant, - mbn2d.running_mean, mbn2d.running_var, mbn2d.eps, mbn2d.weight, - mbn2d.bias, - manaout.eps, - manaout.thresholds, - ceiltau=False) - - # build the new modules - mtwconv = nn.Conv2d(mconv2d.in_channels, mconv2d.out_channels, mconv2d.kernel_size, - stride=mconv2d.stride, padding=mconv2d.padding, dilation=mconv2d.dilation, - groups=mconv2d.groups, - bias=mconv2d.bias is not None).to(torch.device('cpu')) - mtwconv.weight.data = weight - - mlutact = LUTActivation(tau, manaout.quant_levels) - - # register the newly created nodes - vJI_2_ptnode = {} - vJI_2_ptnode[g_RK2JI['/'.join(['R-term', 'TWConv'])]] = PyTorchNode(mtwconv) - vJI_2_ptnode[g_RK2JI['/'.join(['R-term', 'LUTAct'])]] = PyTorchNode(mlutact) - - return JI, vJI_2_ptnode - - # G: Full/original graph - # nodes_dict: Mapping between node identifiers of G and actual underlying objects - # g: One instance of all occurences of the template in G, i.e. one application point for the replacement rule -> one morphism - def apply(self, G, nodes_dict, g): - - # create new containers - G = G.copy() - # Dictionary mapping of node identifiers to a payload - # keys in nodes_dict should be the same as G.nodes - nodes_dict = {**nodes_dict} - - # characterise the match graph H - # Occurence of template in the graph - # SPMATTEO: Some assumptions to discuss - VI = {vH for vH, vL in g.items() if vL in set(self.K.nodes)} # Occurence of context - VHI = {vH for vH, vL in g.items() if vL not in set(self.K.nodes)} # Occurence of core template - HI = G.subgraph(VHI) # HI is the subgraph induced by the set of nodes VHI - - # generate the substitute (sub-)graph J\I (completely detached from G) - # Instantiate blueprint of the replacement graph - JI, vJI_2_ptnode = self.core(HI, g, nodes_dict) - - # add the substitute (sub-)graph J\I to the main graph G - G = nx.compose(G, JI) # G now has two connected but 'independent' subgraphs - nodes_dict.update(vJI_2_ptnode) # Add new payloads from substitute graph - - # glue the substitute (sub-)graph J\I to the interface (sub-)graph I - JI2RK_morphisms = Seeker(self.RK).get_morphisms(JI) - assert len(JI2RK_morphisms) == 1 - g_JI2RK = JI2RK_morphisms[0] - g_RK2JI = {vRK: vJI for vJI, vRK in g_JI2RK.items()} - for vI in VI: # for each node in the interface subgraph of G - vK = g[vI] - G.add_edges_from({(vI, g_RK2JI[vRK]) for (_, vRK) in - self.F_K2RK[vK]}) # incoming interface connections from G to substitute graph - G.add_edges_from({(g_RK2JI[vRK], vI) for (vRK, _) in - self.F_RK2K[vK]}) # outcoming interface connections from substitute graph to G - # the new modules are fully integerized, so the precision tunnel should not embed integer numbers in floating point numbers - # Specific to integer arithmetic transformation -> No relation to graph editing, per-se - if nodes_dict[vI].ntype == qg.graphs.HelperInput.__name__: - pass - elif nodes_dict[vI].ntype == qg.graphs.HelperInputPrecisionTunnel.__name__: - nodes_dict[vI] = PyTorchNode(qg.graphs.HelperInputPrecisionTunnel(1.0)) - else: - raise TypeError # interface nodes should be objects of class `qg.graphs.HelperPrecisionTunnel` only - - # discard the match (sub-)graph H\I - # Assumption: removing a node also removes all arcs pointing to or from that node - G.remove_nodes_from(set(HI.nodes)) - - # Remove the payload, i.e. underying objects, accordingly - for vHI in VHI: - del nodes_dict[vHI] - - return G, nodes_dict - - def seek(self, G, nodes_dict): - gs = self.seeker.get_morphisms(G) - return gs - - -class FoldANAActANAConvBNANAActTypeARule(DPORule): # w/o max pooling - - def __init__(self, lut_entry_bits=16): - - self._lut_entry_bits = lut_entry_bits - - # Nodes of the interface - K_types = OrderedDict() - K_types.update({'HPTout': qg.graphs.HelperOutputPrecisionTunnel.__name__}) - K_types.update({'HPTin': qg.graphs.HelperInputPrecisionTunnel.__name__}) - K_types = OrderedDict([('/'.join(['K-term', k]), v) for k, v in K_types.items()]) - - # Nodes in the core template graph - LK_types = OrderedDict() - LK_types.update({'ANAActin': qa.ana.ANAActivation.__name__}) - LK_types.update({'ANAConv': qa.ana.ANAConv2d.__name__}) - LK_types.update({'BatchNorm': nn.BatchNorm2d.__name__}) - LK_types.update({'ANAActout': qa.ana.ANAActivation.__name__}) - LK_types = OrderedDict([('/'.join(['L-term', k]), v) for k, v in LK_types.items()]) - - # Nodes in the core replacement graph - RK_types = OrderedDict() - RK_types.update({'TWConv': nn.Conv2d.__name__}) - RK_types.update({'LUTAct': LUTActivation.__name__}) - RK_types = OrderedDict([('/'.join(['R-term', k]), v) for k, v in RK_types.items()]) - - K_node_IDs = list(K_types.keys()) - LK_node_IDs = list(LK_types.keys()) - RK_node_IDs = list(RK_types.keys()) - - # define the template graph L [L-term] - L_node_IDs = [K_node_IDs[0]] + LK_node_IDs + [K_node_IDs[-1]] - self.L = nx.DiGraph() - # Define arcs between nodes in full template graph - self.L.add_edges_from({(u, v) for u, v in zip(L_node_IDs[:-1], L_node_IDs[1:])}) - - # Here, graph is only operation nodes - # Necessary for seeker - nx.set_node_attributes(self.L, {vL: Bipartite.KERNEL for vL in set(self.L.nodes)}, 'bipartite') - nx.set_node_attributes(self.L, {**K_types, **LK_types}, 'type') - - # define the context (sub-)graph K [K-term] - VK = set(K_node_IDs) # precision tunnel nodes define the context graph - self.K = self.L.subgraph(VK) - - # define the template (sub-)graph L\K - VLK = set(self.L.nodes).difference(set(self.K.nodes)) - self.LK = self.L.subgraph(VLK) - - # define the replacement (sub-)graph R\K ["gluing" R\K to K yields the graph R, i.e., the R-term] - self.RK = nx.DiGraph() - self.RK.add_edges_from({(u, v) for u, v in zip(RK_node_IDs[:-1], RK_node_IDs[1:])}) - nx.set_node_attributes(self.RK, {vRK: Bipartite.KERNEL for vRK in set(self.RK.nodes)}, 'bipartite') - nx.set_node_attributes(self.RK, RK_types, 'type') - - # define the arcs that go from the vertices of K to those of R\K, and viceversa - E_K2RK = {(K_node_IDs[0], RK_node_IDs[0])} - E_RK2K = {(RK_node_IDs[-1], K_node_IDs[-1])} - E_K2RK2K = E_K2RK | E_RK2K - # disintegrate `E_K2RK` and `E_RK2K` along fibres to speed up rule application - # A fibre is kind of like fixing one argument of a two input one output function and looking at all possible outputs - self.F_K2RK = {vK: set(arc for arc in E_K2RK if arc[0] == vK) for vK in set(self.K.nodes)} - self.F_RK2K = {vK: set(arc for arc in E_RK2K if arc[1] == vK) for vK in set(self.K.nodes)} - - # # glue together the (sub-)graphs L\K and R\K along the vertices of K - # self.S = nx.compose(self.L, self.RK) - # self.S.add_edges_from(E_K2RK2K) - - # since the GRR's L-term has been modified, rebuild the seeker - self.seeker = Seeker(self.L) - - # this machinery can generate always-new identifiers for different rule applications - self._counter = itertools.count() - - def _get_rule_count(self): - rule_count = ''.join(['FANABNANATA', __NODE_ID_FORMAT__.format(next(self._counter))]) - return rule_count - - def core(self, HI, g, nodes_dict): - - # generate the substitute (sub-)graph J\I - rule_count = self._get_rule_count() - g_RK2JI = {vRK: '_'.join([rule_count, vRK.replace('R-term/', '')]) for vRK in set(self.RK.nodes)} - JI = nx.relabel_nodes(self.RK, g_RK2JI, copy=True) - - # get pointers to the old modules; - # these pointers will enable two actions: - # 1. extracting the arguments required to perform the folding - # 2. extracting the parameters to instantiate the new modules - g_L2H = {vL: vH for vH, vL in g.items()} - manain = nodes_dict[g_L2H['/'.join(['L-term', 'ANAActin'])]].nobj - mconv2d = nodes_dict[g_L2H['/'.join(['L-term', 'ANAConv'])]].nobj - mbn2d = nodes_dict[g_L2H['/'.join(['L-term', 'BatchNorm'])]].nobj - manaout = nodes_dict[g_L2H['/'.join(['L-term', 'ANAActout'])]].nobj - - # fold - tau, weight = fold_anaact_anaconv2d_bn2d_anaact(manain.eps, - mconv2d.eps, mconv2d.weight_maybe_quant, - mbn2d.running_mean, mbn2d.running_var, mbn2d.eps, mbn2d.weight, - mbn2d.bias, - manaout.eps, - manaout.thresholds) - - # build the new modules - mtwconv = nn.Conv2d(mconv2d.in_channels, mconv2d.out_channels, mconv2d.kernel_size, - stride=mconv2d.stride, padding=mconv2d.padding, dilation=mconv2d.dilation, - groups=mconv2d.groups, - bias=mconv2d.bias is not None).to(torch.device('cpu')) - mtwconv.weight.data = weight - - mlutact = LUTActivation(tau, manaout.quant_levels) - - # register the newly created nodes - vJI_2_ptnode = {} - vJI_2_ptnode[g_RK2JI['/'.join(['R-term', 'TWConv'])]] = PyTorchNode(mtwconv) - vJI_2_ptnode[g_RK2JI['/'.join(['R-term', 'LUTAct'])]] = PyTorchNode(mlutact) - - return JI, vJI_2_ptnode - - # G: Full/original graph - # nodes_dict: Mapping between node identifiers of G and actual underlying objects - # g: One instance of all occurences of the template in G, i.e. one application point for the replacement rule -> one morphism - def apply(self, G, nodes_dict, g): - - # create new containers - G = G.copy() - # Dictionary mapping of node identifiers to a payload - # keys in nodes_dict should be the same as G.nodes - nodes_dict = {**nodes_dict} - - # characterise the match graph H - # Occurence of template in the graph - # SPMATTEO: Some assumptions to discuss - VI = {vH for vH, vL in g.items() if vL in set(self.K.nodes)} # Occurence of context - VHI = {vH for vH, vL in g.items() if vL not in set(self.K.nodes)} # Occurence of core template - HI = G.subgraph(VHI) # HI is the subgraph induced by the set of nodes VHI - - # generate the substitute (sub-)graph J\I (completely detached from G) - # Instantiate blueprint of the replacement graph - JI, vJI_2_ptnode = self.core(HI, g, nodes_dict) - - # add the substitute (sub-)graph J\I to the main graph G - G = nx.compose(G, JI) # G now has two connected but 'independent' subgraphs - nodes_dict.update(vJI_2_ptnode) # Add new payloads from substitute graph - - # glue the substitute (sub-)graph J\I to the interface (sub-)graph I - JI2RK_morphisms = Seeker(self.RK).get_morphisms(JI) - assert len(JI2RK_morphisms) == 1 - g_JI2RK = JI2RK_morphisms[0] - g_RK2JI = {vRK: vJI for vJI, vRK in g_JI2RK.items()} - for vI in VI: # for each node in the interface subgraph of G - vK = g[vI] - G.add_edges_from({(vI, g_RK2JI[vRK]) for (_, vRK) in - self.F_K2RK[vK]}) # incoming interface connections from G to substitute graph - G.add_edges_from({(g_RK2JI[vRK], vI) for (vRK, _) in - self.F_RK2K[vK]}) # outcoming interface connections from substitute graph to G - # the new modules are fully integerized, so the precision tunnel should not embed integer numbers in floating point numbers - # Specific to integer arithmetic transformation -> No relation to graph editing, per-se - if nodes_dict[vI].ntype == qg.graphs.HelperOutputPrecisionTunnel.__name__: - nodes_dict[vI] = PyTorchNode(qg.graphs.HelperOutputPrecisionTunnel(1.0)) - elif nodes_dict[vI].ntype == qg.graphs.HelperInputPrecisionTunnel.__name__: - nodes_dict[vI] = PyTorchNode(qg.graphs.HelperInputPrecisionTunnel(1.0)) - else: - raise TypeError # interface nodes should be objects of class `qg.graphs.HelperPrecisionTunnel` only - - # discard the match (sub-)graph H\I - # Assumption: removing a node also removes all arcs pointing to or from that node - G.remove_nodes_from(set(HI.nodes)) - - # Remove the payload, i.e. underying objects, accordingly - for vHI in VHI: - del nodes_dict[vHI] - - return G, nodes_dict - - def seek(self, G, nodes_dict): - gs = self.seeker.get_morphisms(G) - return gs - - -class FoldANAActANAConvBNANAActTypeBRule(DPORule): # w/ max pooling - - def __init__(self, lut_entry_bits=16): - - self._lut_entry_bits = lut_entry_bits - - # Nodes of the interface - K_types = OrderedDict() - K_types.update({'HPTout': qg.graphs.HelperOutputPrecisionTunnel.__name__}) - K_types.update({'HPTin': qg.graphs.HelperInputPrecisionTunnel.__name__}) - K_types = OrderedDict([('/'.join(['K-term', k]), v) for k, v in K_types.items()]) - - # Nodes in the core template graph - LK_types = OrderedDict() - LK_types.update({'ANAActin': qa.ana.ANAActivation.__name__}) - LK_types.update({'MaxPool': nn.MaxPool2d.__name__}) - LK_types.update({'ANAConv': qa.ana.ANAConv2d.__name__}) - LK_types.update({'BatchNorm': nn.BatchNorm2d.__name__}) - LK_types.update({'ANAActout': qa.ana.ANAActivation.__name__}) - LK_types = OrderedDict([('/'.join(['L-term', k]), v) for k, v in LK_types.items()]) - - # Nodes in the core replacement graph - RK_types = OrderedDict() - RK_types.update({'MaxPool': nn.MaxPool2d.__name__}) - RK_types.update({'TWConv': nn.Conv2d.__name__}) - RK_types.update({'LUTAct': LUTActivation.__name__}) - RK_types = OrderedDict([('/'.join(['R-term', k]), v) for k, v in RK_types.items()]) - - K_node_IDs = list(K_types.keys()) - LK_node_IDs = list(LK_types.keys()) - RK_node_IDs = list(RK_types.keys()) - - # define the template graph L [L-term] - L_node_IDs = [K_node_IDs[0]] + LK_node_IDs + [K_node_IDs[-1]] - self.L = nx.DiGraph() - # Define arcs between nodes in full template graph - self.L.add_edges_from({(u, v) for u, v in zip(L_node_IDs[:-1], L_node_IDs[1:])}) - - # Here, graph is only operation nodes - # Necessary for seeker - nx.set_node_attributes(self.L, {vL: Bipartite.KERNEL for vL in set(self.L.nodes)}, 'bipartite') - nx.set_node_attributes(self.L, {**K_types, **LK_types}, 'type') - - # define the context (sub-)graph K [K-term] - VK = set(K_node_IDs) # precision tunnel nodes define the context graph - self.K = self.L.subgraph(VK) - - # define the template (sub-)graph L\K - VLK = set(self.L.nodes).difference(set(self.K.nodes)) - self.LK = self.L.subgraph(VLK) - - # define the replacement (sub-)graph R\K ["gluing" R\K to K yields the graph R, i.e., the R-term] - self.RK = nx.DiGraph() - self.RK.add_edges_from({(u, v) for u, v in zip(RK_node_IDs[:-1], RK_node_IDs[1:])}) - nx.set_node_attributes(self.RK, {vRK: Bipartite.KERNEL for vRK in set(self.RK.nodes)}, 'bipartite') - nx.set_node_attributes(self.RK, RK_types, 'type') - - # define the arcs that go from the vertices of K to those of R\K, and viceversa - E_K2RK = {(K_node_IDs[0], RK_node_IDs[0])} - E_RK2K = {(RK_node_IDs[-1], K_node_IDs[-1])} - E_K2RK2K = E_K2RK | E_RK2K - # disintegrate `E_K2RK` and `E_RK2K` along fibres to speed up rule application - # A fibre is kind of like fixing one argument of a two input one output function and looking at all possible outputs - self.F_K2RK = {vK: set(arc for arc in E_K2RK if arc[0] == vK) for vK in set(self.K.nodes)} - self.F_RK2K = {vK: set(arc for arc in E_RK2K if arc[1] == vK) for vK in set(self.K.nodes)} - - # # glue together the (sub-)graphs L\K and R\K along the vertices of K - # self.S = nx.compose(self.L, self.RK) - # self.S.add_edges_from(E_K2RK2K) - - # since the GRR's L-term has been modified, rebuild the seeker - self.seeker = Seeker(self.L) - - # this machinery can generate always-new identifiers for different rule applications - self._counter = itertools.count() - - def _get_rule_count(self): - rule_count = ''.join(['FANABNANATB', __NODE_ID_FORMAT__.format(next(self._counter))]) - return rule_count - - def core(self, HI, g, nodes_dict): - - # generate the substitute (sub-)graph J\I - rule_count = self._get_rule_count() - g_RK2JI = {vRK: '_'.join([rule_count, vRK.replace('R-term/', '')]) for vRK in set(self.RK.nodes)} - JI = nx.relabel_nodes(self.RK, g_RK2JI, copy=True) - - # get pointers to the old modules; - # these pointers will enable two actions: - # 1. extracting the arguments required to perform the folding - # 2. extracting the parameters to instantiate the new modules - g_L2H = {vL: vH for vH, vL in g.items()} - manain = nodes_dict[g_L2H['/'.join(['L-term', 'ANAActin'])]].nobj - mmxpold = nodes_dict[g_L2H['/'.join(['L-term', 'MaxPool'])]].nobj - mconv2d = nodes_dict[g_L2H['/'.join(['L-term', 'ANAConv'])]].nobj - mbn2d = nodes_dict[g_L2H['/'.join(['L-term', 'BatchNorm'])]].nobj - manaout = nodes_dict[g_L2H['/'.join(['L-term', 'ANAActout'])]].nobj - - # fold - tau, weight = fold_anaact_anaconv2d_bn2d_anaact(manain.eps, - mconv2d.eps, mconv2d.weight_maybe_quant, - mbn2d.running_mean, mbn2d.running_var, mbn2d.eps, mbn2d.weight, - mbn2d.bias, - manaout.eps, - manaout.thresholds) - - # build the new modules - mmxpnew = nn.MaxPool2d(kernel_size=mmxpold.kernel_size, stride=mmxpold.stride, padding=mmxpold.padding) - - mtwconv = nn.Conv2d(mconv2d.in_channels, mconv2d.out_channels, mconv2d.kernel_size, - stride=mconv2d.stride, padding=mconv2d.padding, dilation=mconv2d.dilation, - groups=mconv2d.groups, - bias=mconv2d.bias is not None).to(torch.device('cpu')) - mtwconv.weight.data = weight - - mlutact = LUTActivation(tau, manaout.quant_levels) - - # register the newly created nodes - vJI_2_ptnode = {} - vJI_2_ptnode[g_RK2JI['/'.join(['R-term', 'MaxPool'])]] = PyTorchNode(mmxpnew) - vJI_2_ptnode[g_RK2JI['/'.join(['R-term', 'TWConv'])]] = PyTorchNode(mtwconv) - vJI_2_ptnode[g_RK2JI['/'.join(['R-term', 'LUTAct'])]] = PyTorchNode(mlutact) - - return JI, vJI_2_ptnode - - # G: Full/original graph - # nodes_dict: Mapping between node identifiers of G and actual underlying objects - # g: One instance of all occurences of the template in G, i.e. one application point for the replacement rule -> one morphism - def apply(self, G, nodes_dict, g): - - # create new containers - G = G.copy() - # Dictionary mapping of node identifiers to a payload - # keys in nodes_dict should be the same as G.nodes - nodes_dict = {**nodes_dict} - - # characterise the match graph H - # Occurence of template in the graph - # SPMATTEO: Some assumptions to discuss - VI = {vH for vH, vL in g.items() if vL in set(self.K.nodes)} # Occurence of context - VHI = {vH for vH, vL in g.items() if vL not in set(self.K.nodes)} # Occurence of core template - HI = G.subgraph(VHI) # HI is the subgraph induced by the set of nodes VHI - - # generate the substitute (sub-)graph J\I (completely detached from G) - # Instantiate blueprint of the replacement graph - JI, vJI_2_ptnode = self.core(HI, g, nodes_dict) - - # add the substitute (sub-)graph J\I to the main graph G - G = nx.compose(G, JI) # G now has two connected but 'independent' subgraphs - nodes_dict.update(vJI_2_ptnode) # Add new payloads from substitute graph - - # glue the substitute (sub-)graph J\I to the interface (sub-)graph I - JI2RK_morphisms = Seeker(self.RK).get_morphisms(JI) - assert len(JI2RK_morphisms) == 1 - g_JI2RK = JI2RK_morphisms[0] - g_RK2JI = {vRK: vJI for vJI, vRK in g_JI2RK.items()} - for vI in VI: # for each node in the interface subgraph of G - vK = g[vI] - G.add_edges_from({(vI, g_RK2JI[vRK]) for (_, vRK) in - self.F_K2RK[vK]}) # incoming interface connections from G to substitute graph - G.add_edges_from({(g_RK2JI[vRK], vI) for (vRK, _) in - self.F_RK2K[vK]}) # outcoming interface connections from substitute graph to G - # the new modules are fully integerized, so the precision tunnel should not embed integer numbers in floating point numbers - # Specific to integer arithmetic transformation -> No relation to graph editing, per-se - if nodes_dict[vI].ntype == qg.graphs.HelperOutputPrecisionTunnel.__name__: - nodes_dict[vI] = PyTorchNode(qg.graphs.HelperOutputPrecisionTunnel(1.0)) - elif nodes_dict[vI].ntype == qg.graphs.HelperInputPrecisionTunnel.__name__: - nodes_dict[vI] = PyTorchNode(qg.graphs.HelperInputPrecisionTunnel(1.0)) - else: - raise TypeError # interface nodes should be objects of class `qg.graphs.HelperPrecisionTunnel` only - - # discard the match (sub-)graph H\I - # Assumption: removing a node also removes all arcs pointing to or from that node - G.remove_nodes_from(set(HI.nodes)) - - # Remove the payload, i.e. underying objects, accordingly - for vHI in VHI: - del nodes_dict[vHI] - - return G, nodes_dict - - def seek(self, G, nodes_dict): - gs = self.seeker.get_morphisms(G) - return gs - - -class FoldANAActANALinearBNANAActTypeARule(DPORule): # w/o pooling layers - - def __init__(self, lut_entry_bits=16): - - self._lut_entry_bits = lut_entry_bits - - # Nodes of the interface - K_types = OrderedDict() - K_types.update({'HPTout': qg.graphs.HelperOutputPrecisionTunnel.__name__}) - K_types.update({'HPTin': qg.graphs.HelperInputPrecisionTunnel.__name__}) - K_types = OrderedDict([('/'.join(['K-term', k]), v) for k, v in K_types.items()]) - - # Nodes in the core template graph - LK_types = OrderedDict() - LK_types.update({'ANAActin': qa.ana.ANAActivation.__name__}) - LK_types.update({'ANALinear': qa.ana.ANALinear.__name__}) - LK_types.update({'BatchNorm': nn.BatchNorm1d.__name__}) - LK_types.update({'ANAActout': qa.ana.ANAActivation.__name__}) - LK_types = OrderedDict([('/'.join(['L-term', k]), v) for k, v in LK_types.items()]) - - # Nodes in the core replacement graph - RK_types = OrderedDict() - RK_types.update({'TWLinear': nn.Linear.__name__}) - RK_types.update({'LUTAct': LUTActivation.__name__}) - RK_types = OrderedDict([('/'.join(['R-term', k]), v) for k, v in RK_types.items()]) - - K_node_IDs = list(K_types.keys()) - LK_node_IDs = list(LK_types.keys()) - RK_node_IDs = list(RK_types.keys()) - - # define the template graph L [L-term] - L_node_IDs = [K_node_IDs[0]] + LK_node_IDs + [K_node_IDs[-1]] - self.L = nx.DiGraph() - # Define arcs between nodes in full template graph - self.L.add_edges_from({(u, v) for u, v in zip(L_node_IDs[:-1], L_node_IDs[1:])}) - - # Here, graph is only operation nodes - # Necessary for seeker - nx.set_node_attributes(self.L, {vL: Bipartite.KERNEL for vL in set(self.L.nodes)}, 'bipartite') - nx.set_node_attributes(self.L, {**K_types, **LK_types}, 'type') - - # define the context (sub-)graph K [K-term] - VK = set(K_node_IDs) # precision tunnel nodes define the context graph - self.K = self.L.subgraph(VK) - - # define the template (sub-)graph L\K - VLK = set(self.L.nodes).difference(set(self.K.nodes)) - self.LK = self.L.subgraph(VLK) - - # define the replacement (sub-)graph R\K ["gluing" R\K to K yields the graph R, i.e., the R-term] - self.RK = nx.DiGraph() - self.RK.add_edges_from({(u, v) for u, v in zip(RK_node_IDs[:-1], RK_node_IDs[1:])}) - nx.set_node_attributes(self.RK, {vRK: Bipartite.KERNEL for vRK in set(self.RK.nodes)}, 'bipartite') - nx.set_node_attributes(self.RK, RK_types, 'type') - - # define the arcs that go from the vertices of K to those of R\K, and viceversa - E_K2RK = {(K_node_IDs[0], RK_node_IDs[0])} - E_RK2K = {(RK_node_IDs[-1], K_node_IDs[-1])} - E_K2RK2K = E_K2RK | E_RK2K - # disintegrate `E_K2RK` and `E_RK2K` along fibres to speed up rule application - # A fibre is kind of like fixing one argument of a two input one output function and looking at all possible outputs - self.F_K2RK = {vK: set(arc for arc in E_K2RK if arc[0] == vK) for vK in set(self.K.nodes)} - self.F_RK2K = {vK: set(arc for arc in E_RK2K if arc[1] == vK) for vK in set(self.K.nodes)} - - # # glue together the (sub-)graphs L\K and R\K along the vertices of K - # self.S = nx.compose(self.L, self.RK) - # self.S.add_edges_from(E_K2RK2K) - - # since the GRR's L-term has been modified, rebuild the seeker - self.seeker = Seeker(self.L) - - # this machinery can generate always-new identifiers for different rule applications - self._counter = itertools.count() - - def _get_rule_count(self): - rule_count = ''.join(['FANABNANALinTA', __NODE_ID_FORMAT__.format(next(self._counter))]) - return rule_count - - def core(self, HI, g, nodes_dict): - - # generate the substitute (sub-)graph J\I - rule_count = self._get_rule_count() - g_RK2JI = {vRK: '_'.join([rule_count, vRK.replace('R-term/', '')]) for vRK in set(self.RK.nodes)} - JI = nx.relabel_nodes(self.RK, g_RK2JI, copy=True) - - # get pointers to the old modules; - # these pointers will enable two actions: - # 1. extracting the arguments required to perform the folding - # 2. extracting the parameters to instantiate the new modules - g_L2H = {vL: vH for vH, vL in g.items()} - manain = nodes_dict[g_L2H['/'.join(['L-term', 'ANAActin'])]].nobj - mlinear = nodes_dict[g_L2H['/'.join(['L-term', 'ANALinear'])]].nobj - mbn1d = nodes_dict[g_L2H['/'.join(['L-term', 'BatchNorm'])]].nobj - manaout = nodes_dict[g_L2H['/'.join(['L-term', 'ANAActout'])]].nobj - - # fold - tau, weight = fold_anaact_analinear_bn1d_anaact(manain.eps, - mlinear.eps, mlinear.weight_maybe_quant, - mbn1d.running_mean, mbn1d.running_var, mbn1d.eps, mbn1d.weight, - mbn1d.bias, - manaout.eps, - manaout.thresholds) - - # build the new modules - mtwlinear = nn.Linear(mlinear.in_features, mlinear.out_features, - bias=mlinear.bias is not None).to(torch.device('cpu')) - mtwlinear.weight.data = weight - - mlutact = LUTActivation(tau, manaout.quant_levels) - - # register the newly created nodes - vJI_2_ptnode = {} - vJI_2_ptnode[g_RK2JI['/'.join(['R-term', 'TWLinear'])]] = PyTorchNode(mtwlinear) - vJI_2_ptnode[g_RK2JI['/'.join(['R-term', 'LUTAct'])]] = PyTorchNode(mlutact) - - return JI, vJI_2_ptnode - - # G: Full/original graph - # nodes_dict: Mapping between node identifiers of G and actual underlying objects - # g: One instance of all occurences of the template in G, i.e. one application point for the replacement rule -> one morphism - def apply(self, G, nodes_dict, g): - - # create new containers - G = G.copy() - # Dictionary mapping of node identifiers to a payload - # keys in nodes_dict should be the same as G.nodes - nodes_dict = {**nodes_dict} - - # characterise the match graph H - # Occurence of template in the graph - # SPMATTEO: Some assumptions to discuss - VI = {vH for vH, vL in g.items() if vL in set(self.K.nodes)} # Occurence of context - VHI = {vH for vH, vL in g.items() if vL not in set(self.K.nodes)} # Occurence of core template - HI = G.subgraph(VHI) # HI is the subgraph induced by the set of nodes VHI - - # generate the substitute (sub-)graph J\I (completely detached from G) - # Instantiate blueprint of the replacement graph - JI, vJI_2_ptnode = self.core(HI, g, nodes_dict) - - # add the substitute (sub-)graph J\I to the main graph G - G = nx.compose(G, JI) # G now has two connected but 'independent' subgraphs - nodes_dict.update(vJI_2_ptnode) # Add new payloads from substitute graph - - # glue the substitute (sub-)graph J\I to the interface (sub-)graph I - JI2RK_morphisms = Seeker(self.RK).get_morphisms(JI) - assert len(JI2RK_morphisms) == 1 - g_JI2RK = JI2RK_morphisms[0] - g_RK2JI = {vRK: vJI for vJI, vRK in g_JI2RK.items()} - for vI in VI: # for each node in the interface subgraph of G - vK = g[vI] - G.add_edges_from({(vI, g_RK2JI[vRK]) for (_, vRK) in - self.F_K2RK[vK]}) # incoming interface connections from G to substitute graph - G.add_edges_from({(g_RK2JI[vRK], vI) for (vRK, _) in - self.F_RK2K[vK]}) # outcoming interface connections from substitute graph to G - # the new modules are fully integerized, so the precision tunnel should not embed integer numbers in floating point numbers - # Specific to integer arithmetic transformation -> No relation to graph editing, per-se - if nodes_dict[vI].ntype == qg.graphs.HelperOutputPrecisionTunnel.__name__: - nodes_dict[vI] = PyTorchNode(qg.graphs.HelperOutputPrecisionTunnel(1.0)) - elif nodes_dict[vI].ntype == qg.graphs.HelperInputPrecisionTunnel.__name__: - nodes_dict[vI] = PyTorchNode(qg.graphs.HelperInputPrecisionTunnel(1.0)) - else: - raise TypeError # interface nodes should be objects of class `qg.graphs.HelperPrecisionTunnel` only - - # discard the match (sub-)graph H\I - # Assumption: removing a node also removes all arcs pointing to or from that node - G.remove_nodes_from(set(HI.nodes)) - - # Remove the payload, i.e. underying objects, accordingly - for vHI in VHI: - del nodes_dict[vHI] - - return G, nodes_dict - - def seek(self, G, nodes_dict): - gs = self.seeker.get_morphisms(G) - return gs - - -class FoldANAActANALinearBNANAActTypeBRule(DPORule): # w/ pooling layers - - def __init__(self, lut_entry_bits=16): - - self._lut_entry_bits = lut_entry_bits - - # Nodes of the interface - K_types = OrderedDict() - K_types.update({'HPTout': qg.graphs.HelperOutputPrecisionTunnel.__name__}) - K_types.update({'HPTin': qg.graphs.HelperInputPrecisionTunnel.__name__}) - K_types = OrderedDict([('/'.join(['K-term', k]), v) for k, v in K_types.items()]) - - # Nodes in the core template graph - LK_types = OrderedDict() - LK_types.update({'ANAActin': qa.ana.ANAActivation.__name__}) - LK_types.update({'MaxPool': nn.MaxPool2d.__name__}) - LK_types.update({'AvgPool': nn.AdaptiveAvgPool2d.__name__}) - LK_types.update({'ViewFlattenNd': qg.graphs.modules.ViewFlattenNd.__name__}) - LK_types.update({'ANALinear': qa.ana.ANALinear.__name__}) - LK_types.update({'BatchNorm': nn.BatchNorm1d.__name__}) - LK_types.update({'ANAActout': qa.ana.ANAActivation.__name__}) - LK_types = OrderedDict([('/'.join(['L-term', k]), v) for k, v in LK_types.items()]) - - # Nodes in the core replacement graph - RK_types = OrderedDict() - RK_types.update({'MaxPool': nn.MaxPool2d.__name__}) - RK_types.update({'TWLinear': nn.Linear.__name__}) - RK_types.update({'LUTAct': LUTActivation.__name__}) - RK_types = OrderedDict([('/'.join(['R-term', k]), v) for k, v in RK_types.items()]) - - K_node_IDs = list(K_types.keys()) - LK_node_IDs = list(LK_types.keys()) - RK_node_IDs = list(RK_types.keys()) - - # define the template graph L [L-term] - L_node_IDs = [K_node_IDs[0]] + LK_node_IDs + [K_node_IDs[-1]] - self.L = nx.DiGraph() - # Define arcs between nodes in full template graph - self.L.add_edges_from({(u, v) for u, v in zip(L_node_IDs[:-1], L_node_IDs[1:])}) - - # Here, graph is only operation nodes - # Necessary for seeker - nx.set_node_attributes(self.L, {vL: Bipartite.KERNEL for vL in set(self.L.nodes)}, 'bipartite') - nx.set_node_attributes(self.L, {**K_types, **LK_types}, 'type') - - # define the context (sub-)graph K [K-term] - VK = set(K_node_IDs) # precision tunnel nodes define the context graph - self.K = self.L.subgraph(VK) - - # define the template (sub-)graph L\K - VLK = set(self.L.nodes).difference(set(self.K.nodes)) - self.LK = self.L.subgraph(VLK) - - # define the replacement (sub-)graph R\K ["gluing" R\K to K yields the graph R, i.e., the R-term] - self.RK = nx.DiGraph() - self.RK.add_edges_from({(u, v) for u, v in zip(RK_node_IDs[:-1], RK_node_IDs[1:])}) - nx.set_node_attributes(self.RK, {vRK: Bipartite.KERNEL for vRK in set(self.RK.nodes)}, 'bipartite') - nx.set_node_attributes(self.RK, RK_types, 'type') - - # define the arcs that go from the vertices of K to those of R\K, and viceversa - E_K2RK = {(K_node_IDs[0], RK_node_IDs[0])} - E_RK2K = {(RK_node_IDs[-1], K_node_IDs[-1])} - E_K2RK2K = E_K2RK | E_RK2K - # disintegrate `E_K2RK` and `E_RK2K` along fibres to speed up rule application - # A fibre is kind of like fixing one argument of a two input one output function and looking at all possible outputs - self.F_K2RK = {vK: set(arc for arc in E_K2RK if arc[0] == vK) for vK in set(self.K.nodes)} - self.F_RK2K = {vK: set(arc for arc in E_RK2K if arc[1] == vK) for vK in set(self.K.nodes)} - - # # glue together the (sub-)graphs L\K and R\K along the vertices of K - # self.S = nx.compose(self.L, self.RK) - # self.S.add_edges_from(E_K2RK2K) - - # since the GRR's L-term has been modified, rebuild the seeker - self.seeker = Seeker(self.L) - - # this machinery can generate always-new identifiers for different rule applications - self._counter = itertools.count() - - def _get_rule_count(self): - rule_count = ''.join(['FANABNANALinTB', __NODE_ID_FORMAT__.format(next(self._counter))]) - return rule_count - - def core(self, HI, g, nodes_dict): - - # generate the substitute (sub-)graph J\I - rule_count = self._get_rule_count() - g_RK2JI = {vRK: '_'.join([rule_count, vRK.replace('R-term/', '')]) for vRK in set(self.RK.nodes)} - JI = nx.relabel_nodes(self.RK, g_RK2JI, copy=True) - - # get pointers to the old modules; - # these pointers will enable two actions: - # 1. extracting the arguments required to perform the folding - # 2. extracting the parameters to instantiate the new modules - g_L2H = {vL: vH for vH, vL in g.items()} - manain = nodes_dict[g_L2H['/'.join(['L-term', 'ANAActin'])]].nobj - mmxpold = nodes_dict[g_L2H['/'.join(['L-term', 'MaxPool'])]].nobj - mlinear = nodes_dict[g_L2H['/'.join(['L-term', 'ANALinear'])]].nobj - mbn1d = nodes_dict[g_L2H['/'.join(['L-term', 'BatchNorm'])]].nobj - manaout = nodes_dict[g_L2H['/'.join(['L-term', 'ANAActout'])]].nobj - - # fold - tau, weight = fold_anaact_analinear_bn1d_anaact(manain.eps, - mlinear.eps, mlinear.weight_maybe_quant, - mbn1d.running_mean, mbn1d.running_var, mbn1d.eps, mbn1d.weight, - mbn1d.bias, - manaout.eps, - manaout.thresholds) - - # build the new modules - mmxpnew = nn.MaxPool2d(kernel_size=mmxpold.kernel_size, stride=mmxpold.stride, padding=mmxpold.padding) - - mtwlinear = nn.Linear(mlinear.in_features, mlinear.out_features, - bias=mlinear.bias is not None).to(torch.device('cpu')) - mtwlinear.weight.data = weight - - mlutact = LUTActivation(tau, manaout.quant_levels) - - # register the newly created nodes - vJI_2_ptnode = {} - vJI_2_ptnode[g_RK2JI['/'.join(['R-term', 'MaxPool'])]] = PyTorchNode(mmxpnew) - vJI_2_ptnode[g_RK2JI['/'.join(['R-term', 'TWLinear'])]] = PyTorchNode(mtwlinear) - vJI_2_ptnode[g_RK2JI['/'.join(['R-term', 'LUTAct'])]] = PyTorchNode(mlutact) - - return JI, vJI_2_ptnode - - # G: Full/original graph - # nodes_dict: Mapping between node identifiers of G and actual underlying objects - # g: One instance of all occurences of the template in G, i.e. one application point for the replacement rule -> one morphism - def apply(self, G, nodes_dict, g): - - # create new containers - G = G.copy() - # Dictionary mapping of node identifiers to a payload - # keys in nodes_dict should be the same as G.nodes - nodes_dict = {**nodes_dict} - - # characterise the match graph H - # Occurence of template in the graph - # SPMATTEO: Some assumptions to discuss - VI = {vH for vH, vL in g.items() if vL in set(self.K.nodes)} # Occurence of context - VHI = {vH for vH, vL in g.items() if vL not in set(self.K.nodes)} # Occurence of core template - HI = G.subgraph(VHI) # HI is the subgraph induced by the set of nodes VHI - - # generate the substitute (sub-)graph J\I (completely detached from G) - # Instantiate blueprint of the replacement graph - JI, vJI_2_ptnode = self.core(HI, g, nodes_dict) - - # add the substitute (sub-)graph J\I to the main graph G - G = nx.compose(G, JI) # G now has two connected but 'independent' subgraphs - nodes_dict.update(vJI_2_ptnode) # Add new payloads from substitute graph - - # glue the substitute (sub-)graph J\I to the interface (sub-)graph I - JI2RK_morphisms = Seeker(self.RK).get_morphisms(JI) - assert len(JI2RK_morphisms) == 1 - g_JI2RK = JI2RK_morphisms[0] - g_RK2JI = {vRK: vJI for vJI, vRK in g_JI2RK.items()} - for vI in VI: # for each node in the interface subgraph of G - vK = g[vI] - G.add_edges_from({(vI, g_RK2JI[vRK]) for (_, vRK) in - self.F_K2RK[vK]}) # incoming interface connections from G to substitute graph - G.add_edges_from({(g_RK2JI[vRK], vI) for (vRK, _) in - self.F_RK2K[vK]}) # outcoming interface connections from substitute graph to G - # the new modules are fully integerized, so the precision tunnel should not embed integer numbers in floating point numbers - # Specific to integer arithmetic transformation -> No relation to graph editing, per-se - if nodes_dict[vI].ntype == qg.graphs.HelperOutputPrecisionTunnel.__name__: - nodes_dict[vI] = PyTorchNode(qg.graphs.HelperOutputPrecisionTunnel(1.0)) - elif nodes_dict[vI].ntype == qg.graphs.HelperInputPrecisionTunnel.__name__: - nodes_dict[vI] = PyTorchNode(qg.graphs.HelperInputPrecisionTunnel(1.0)) - else: - raise TypeError # interface nodes should be objects of class `qg.graphs.HelperPrecisionTunnel` only - - # discard the match (sub-)graph H\I - # Assumption: removing a node also removes all arcs pointing to or from that node - G.remove_nodes_from(set(HI.nodes)) - - # Remove the payload, i.e. underying objects, accordingly - for vHI in VHI: - del nodes_dict[vHI] - - return G, nodes_dict - - def seek(self, G, nodes_dict): - gs = self.seeker.get_morphisms(G) - return gs - - -class FoldANAActLinearRule(DPORule): - - def __init__(self, lut_entry_bits=16): - - self._lut_entry_bits = lut_entry_bits - - # Nodes of the interface - K_types = OrderedDict() - K_types.update({'HPTout': qg.graphs.HelperOutputPrecisionTunnel.__name__}) - K_types.update({'HPout': qg.graphs.HelperOutput.__name__}) - K_types = OrderedDict([('/'.join(['K-term', k]), v) for k, v in K_types.items()]) - - # Nodes in the core template graph - LK_types = OrderedDict() - LK_types.update({'ANAAct': qa.ana.ANAActivation.__name__}) - LK_types.update({'Linear': nn.Linear.__name__}) - LK_types = OrderedDict([('/'.join(['L-term', k]), v) for k, v in LK_types.items()]) - - # Nodes in the core replacement graph - RK_types = OrderedDict() - RK_types.update({'Linear': nn.Linear.__name__}) - RK_types = OrderedDict([('/'.join(['R-term', k]), v) for k, v in RK_types.items()]) - - K_node_IDs = list(K_types.keys()) - LK_node_IDs = list(LK_types.keys()) - RK_node_IDs = list(RK_types.keys()) - - # define the template graph L [L-term] - L_node_IDs = [K_node_IDs[0]] + LK_node_IDs + [K_node_IDs[-1]] - self.L = nx.DiGraph() - # Define arcs between nodes in full template graph - self.L.add_edges_from({(u, v) for u, v in zip(L_node_IDs[:-1], L_node_IDs[1:])}) - - # Here, graph is only operation nodes - # Necessary for seeker - nx.set_node_attributes(self.L, {vL: Bipartite.KERNEL for vL in set(self.L.nodes)}, 'bipartite') - nx.set_node_attributes(self.L, {**K_types, **LK_types}, 'type') - - # define the context (sub-)graph K [K-term] - VK = set(K_node_IDs) # precision tunnel nodes define the context graph - self.K = self.L.subgraph(VK) - - # define the template (sub-)graph L\K - VLK = set(self.L.nodes).difference(set(self.K.nodes)) - self.LK = self.L.subgraph(VLK) - - # define the replacement (sub-)graph R\K ["gluing" R\K to K yields the graph R, i.e., the R-term] - self.RK = nx.DiGraph() - ### WARNING! if R\K has only one node, this initialisation will fail! - self.RK.add_nodes_from(RK_node_IDs) - self.RK.add_edges_from({(u, v) for u, v in zip(RK_node_IDs[:-1], RK_node_IDs[1:])}) - nx.set_node_attributes(self.RK, {vRK: Bipartite.KERNEL for vRK in set(self.RK.nodes)}, 'bipartite') - nx.set_node_attributes(self.RK, RK_types, 'type') - - # define the arcs that go from the vertices of K to those of R\K, and viceversa - E_K2RK = {(K_node_IDs[0], RK_node_IDs[0])} - E_RK2K = {(RK_node_IDs[-1], K_node_IDs[-1])} - E_K2RK2K = E_K2RK | E_RK2K - # disintegrate `E_K2RK` and `E_RK2K` along fibres to speed up rule application - # A fibre is kind of like fixing one argument of a two input one output function and looking at all possible outputs - self.F_K2RK = {vK: set(arc for arc in E_K2RK if arc[0] == vK) for vK in set(self.K.nodes)} - self.F_RK2K = {vK: set(arc for arc in E_RK2K if arc[1] == vK) for vK in set(self.K.nodes)} - - # # glue together the (sub-)graphs L\K and R\K along the vertices of K - # self.S = nx.compose(self.L, self.RK) - # self.S.add_edges_from(E_K2RK2K) - - # since the GRR's L-term has been modified, rebuild the seeker - self.seeker = Seeker(self.L) - - # this machinery can generate always-new identifiers for different rule applications - self._counter = itertools.count() - - def _get_rule_count(self): - rule_count = ''.join(['FANAActLinear', __NODE_ID_FORMAT__.format(next(self._counter))]) - return rule_count - - def core(self, HI, g, nodes_dict): - - # generate the substitute (sub-)graph J\I - rule_count = self._get_rule_count() - g_RK2JI = {vRK: '_'.join([rule_count, vRK.replace('R-term/', '')]) for vRK in set(self.RK.nodes)} - JI = nx.relabel_nodes(self.RK, g_RK2JI, copy=True) - - # get pointers to the old modules; - # these pointers will enable two actions: - # 1. extracting the arguments required to perform the folding - # 2. extracting the parameters to instantiate the new modules - g_L2H = {vL: vH for vH, vL in g.items()} - manain = nodes_dict[g_L2H['/'.join(['L-term', 'ANAAct'])]].nobj - mlinearold = nodes_dict[g_L2H['/'.join(['L-term', 'Linear'])]].nobj - - # fold - weight = manain.eps.item() * mlinearold.weight.data - - # build the new modules - mlinearnew = nn.Linear(mlinearold.in_features, mlinearold.out_features, - bias=mlinearold.bias is not None).to(torch.device('cpu')) - mlinearnew.weight.data = weight - mlinearnew.bias.data = mlinearold.bias.data - - # register the newly created nodes - vJI_2_ptnode = {} - vJI_2_ptnode[g_RK2JI['/'.join(['R-term', 'Linear'])]] = PyTorchNode(mlinearnew) - - return JI, vJI_2_ptnode - - # G: Full/original graph - # nodes_dict: Mapping between node identifiers of G and actual underlying objects - # g: One instance of all occurences of the template in G, i.e. one application point for the replacement rule -> one morphism - def apply(self, G, nodes_dict, g): - - # create new containers - G = G.copy() - # Dictionary mapping of node identifiers to a payload - # keys in nodes_dict should be the same as G.nodes - nodes_dict = {**nodes_dict} - - # characterise the match graph H - # Occurence of template in the graph - # SPMATTEO: Some assumptions to discuss - VI = {vH for vH, vL in g.items() if vL in set(self.K.nodes)} # Occurence of context - VHI = {vH for vH, vL in g.items() if vL not in set(self.K.nodes)} # Occurence of core template - HI = G.subgraph(VHI) # HI is the subgraph induced by the set of nodes VHI - - # generate the substitute (sub-)graph J\I (completely detached from G) - # Instantiate blueprint of the replacement graph - JI, vJI_2_ptnode = self.core(HI, g, nodes_dict) - - # add the substitute (sub-)graph J\I to the main graph G - G = nx.compose(G, JI) # G now has two connected but 'independent' subgraphs - nodes_dict.update(vJI_2_ptnode) # Add new payloads from substitute graph - - # glue the substitute (sub-)graph J\I to the interface (sub-)graph I - JI2RK_morphisms = Seeker(self.RK).get_morphisms(JI) - assert len(JI2RK_morphisms) == 1 - g_JI2RK = JI2RK_morphisms[0] - g_RK2JI = {vRK: vJI for vJI, vRK in g_JI2RK.items()} - for vI in VI: # for each node in the interface subgraph of G - vK = g[vI] - G.add_edges_from({(vI, g_RK2JI[vRK]) for (_, vRK) in - self.F_K2RK[vK]}) # incoming interface connections from G to substitute graph - G.add_edges_from({(g_RK2JI[vRK], vI) for (vRK, _) in - self.F_RK2K[vK]}) # outcoming interface connections from substitute graph to G - # the new modules are fully integerized, so the precision tunnel should not embed integer numbers in floating point numbers - # Specific to integer arithmetic transformation -> No relation to graph editing, per-se - if nodes_dict[vI].ntype == qg.graphs.HelperOutputPrecisionTunnel.__name__: - nodes_dict[vI] = PyTorchNode(qg.graphs.HelperOutputPrecisionTunnel(1.0)) - elif nodes_dict[vI].ntype == qg.graphs.HelperOutput.__name__: - pass - else: - raise TypeError # interface nodes should be objects of class `qg.graphs.HelperPrecisionTunnel` only - - # discard the match (sub-)graph H\I - # Assumption: removing a node also removes all arcs pointing to or from that node - G.remove_nodes_from(set(HI.nodes)) - - # Remove the payload, i.e. underying objects, accordingly - for vHI in VHI: - del nodes_dict[vHI] - - return G, nodes_dict - - def seek(self, G, nodes_dict): - gs = self.seeker.get_morphisms(G) - return gs - diff --git a/backends/cutie/grrules/ana/folding.py b/backends/cutie/grrules/ana/folding.py deleted file mode 100644 index dac91f0..0000000 --- a/backends/cutie/grrules/ana/folding.py +++ /dev/null @@ -1,127 +0,0 @@ -# -# folding.py -# -# Author(s): -# Matteo Spallanzani -# -# Copyright (c) 2020-2021 ETH Zurich. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -import numpy as np -import torch - - -def fold_anaact_anaconv2d_bn2d_anaact(eps_x: torch.Tensor, - eps_w: torch.Tensor, weight: torch.Tensor, - mi: torch.Tensor, sigma: torch.Tensor, bn_eps: torch.Tensor, gamma: torch.Tensor, beta: torch.Tensor, - eps_s: torch.Tensor, theta: torch.Tensor, ceiltau: bool = True): - - def torch2numpyfp64(x): - return x.detach().cpu().numpy().astype(np.float64) - - eps_x = torch2numpyfp64(eps_x) - eps_w = torch2numpyfp64(eps_w) - weight = torch2numpyfp64(weight) - mi = torch2numpyfp64(mi) - sigma = torch2numpyfp64(sigma) - gamma = torch2numpyfp64(gamma) - beta = torch2numpyfp64(beta) - eps_s = torch2numpyfp64(eps_s) - theta = torch2numpyfp64(theta) - - # compensate for negative gammas - flip = np.sign(gamma) - w_tmp = weight.transpose(1, 2, 3, 0) - w_tmp *= flip - weight = w_tmp.transpose(3, 0, 1, 2) - - # https://github.com/pytorch/pytorch/blob/b5e832111e5e4bb3dd66d716d398b81fe70c6af0/torch/csrc/jit/tensorexpr/kernel.cpp#L2015 - sigma = np.sqrt(sigma + bn_eps) - - # folding - xi = gamma * (sigma ** -1) - zeta = beta - mi * xi - - gammaprime = flip * (xi * (eps_x * eps_w) / eps_s) - betaprime = zeta / eps_s - - # prepare for broadcasting - gammaprime = np.expand_dims(gammaprime, axis=0) - betaprime = np.expand_dims(betaprime, axis=0) - theta = np.expand_dims(theta, axis=-1) - - # absorb folded parameters into thresholds - tau = (theta - betaprime) / gammaprime - - assert np.all(tau[0] < tau[1]) - - def numpy2torchfp64(x): - return torch.from_numpy(x.astype(np.float64)) - - if ceiltau: - return numpy2torchfp64(np.ceil(tau)).float(), numpy2torchfp64(weight).float() - else: - return numpy2torchfp64(tau).float(), numpy2torchfp64(weight).float() - - -def fold_anaact_analinear_bn1d_anaact(eps_x: torch.Tensor, - eps_w: torch.Tensor, weight: torch.Tensor, - mi: torch.Tensor, sigma: torch.Tensor, bn_eps: torch.Tensor, gamma: torch.Tensor, beta: torch.Tensor, - eps_s: torch.Tensor, theta: torch.Tensor): - - def torch2numpyfp64(x): - return x.detach().cpu().numpy().astype(np.float64) - - eps_x = torch2numpyfp64(eps_x) - eps_w = torch2numpyfp64(eps_w) - weight = torch2numpyfp64(weight) - mi = torch2numpyfp64(mi) - sigma = torch2numpyfp64(sigma) - gamma = torch2numpyfp64(gamma) - beta = torch2numpyfp64(beta) - eps_s = torch2numpyfp64(eps_s) - theta = torch2numpyfp64(theta) - - # compensate for negative gammas - flip = np.sign(gamma) - w_tmp = weight.transpose(1, 0) - w_tmp *= flip - weight = w_tmp.transpose(1, 0) - - # https://github.com/pytorch/pytorch/blob/b5e832111e5e4bb3dd66d716d398b81fe70c6af0/torch/csrc/jit/tensorexpr/kernel.cpp#L2015 - sigma = np.sqrt(sigma + bn_eps) - - # folding - xi = gamma * (sigma ** -1) - zeta = beta - mi * xi - - gammaprime = flip * (xi * (eps_x * eps_w) / eps_s) - betaprime = zeta / eps_s - - # prepare for broadcasting - gammaprime = np.expand_dims(gammaprime, axis=0) - betaprime = np.expand_dims(betaprime, axis=0) - theta = np.expand_dims(theta, axis=-1) - - # absorb folded parameters into thresholds - tau = (theta - betaprime) / gammaprime - - assert np.all(tau[0] < tau[1]) - - def numpy2torchfp64(x): - return torch.from_numpy(x.astype(np.float64)) - - return numpy2torchfp64(np.ceil(tau)).float(), numpy2torchfp64(weight).float() - diff --git a/backends/cutie/grrules/ana/lutactivation.py b/backends/cutie/grrules/ana/lutactivation.py deleted file mode 100644 index 218c07b..0000000 --- a/backends/cutie/grrules/ana/lutactivation.py +++ /dev/null @@ -1,50 +0,0 @@ -# -# lutactivation.py -# -# Author(s): -# Matteo Spallanzani -# -# Copyright (c) 2020-2021 ETH Zurich. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -import torch -import torch.nn as nn - - -class LUTActivation(nn.Module): - - def __init__(self, tau, quant_levels): - super(LUTActivation, self).__init__() - self.setup_parameters(self, tau, quant_levels) - self.unfolded_tau = False - - @staticmethod - def setup_parameters(lutmod, tau, quant_levels): - lutmod.register_parameter('tau', nn.Parameter(tau, requires_grad=False)) - lutmod.register_parameter('quant_levels', nn.Parameter(quant_levels, requires_grad=False)) - lutmod.register_parameter('q0', nn.Parameter(quant_levels[0], requires_grad=False)) - lutmod.register_parameter('jumps', nn.Parameter(quant_levels[1:] - quant_levels[:-1], requires_grad=False)) - - def forward(self, x): - - if not self.unfolded_tau: - self.tau.data = self.tau[(...,) + (None,) * (x.dim() - 2)] - self.unfolded_tau = True - - x = x.unsqueeze(1) - cdf = (x - self.tau >= 0.0).float() - - y = self.q0 + torch.sum(self.jumps[(...,) + (None,) * (cdf.dim() - 2)] * cdf, 1) - - return y diff --git a/backends/twn_accelerator/grrules/__init__.py b/backends/twn_accelerator/grrules/__init__.py deleted file mode 100644 index d90d303..0000000 --- a/backends/twn_accelerator/grrules/__init__.py +++ /dev/null @@ -1,23 +0,0 @@ -# -# __init__.py -# -# Author(s): -# Matteo Spallanzani -# -# Copyright (c) 2020-2021 ETH Zurich. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -from .dporules import * - diff --git a/backends/twn_accelerator/grrules/dporules.py b/backends/twn_accelerator/grrules/dporules.py deleted file mode 100644 index 5fcb3be..0000000 --- a/backends/twn_accelerator/grrules/dporules.py +++ /dev/null @@ -1,719 +0,0 @@ -# -# dporules.py -# -# Author(s): -# Matteo Spallanzani -# -# Copyright (c) 2020-2021 ETH Zurich. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -import networkx as nx -from collections import OrderedDict -import itertools -import math -import torch -import torch.nn as nn - -import quantlib.editing.graphs as qg -from quantlib.editing.graphs.grrules.dporules import DPORule -from quantlib.editing.graphs.grrules import Seeker -from quantlib.editing.graphs.graphs.nodes import Bipartite, __NODE_ID_FORMAT__, PyTorchNode - -import quantlib.algorithms as qa - -from .folding import foldsteinqconvbnste, foldconvbnste, foldsteinqconvbn - - -__all__ = [ - 'FoldSTEINQConvBNSTETypeARule', - 'FoldSTEINQConvBNSTETypeBRule', - 'FoldConvBNSTERule', - 'FoldSTEINQConvBNRule', -] - - -class FoldSTEINQConvBNSTETypeARule(DPORule): # w/o max pooling - - def __init__(self, gamma_int_bits=10, gamma_frac_bits=17, beta_int_bits=8, beta_frac_bits=0): - - self._gamma_int_bits = gamma_int_bits - self._gamma_frac_bits = gamma_frac_bits - self._beta_int_bits = beta_int_bits - self._beta_frac_bits = beta_frac_bits - - # Nodes of the interface - K_types = OrderedDict() - K_types.update({'HPTout': qg.graphs.HelperOutputPrecisionTunnel.__name__}) - K_types.update({'HPTin': qg.graphs.HelperInputPrecisionTunnel.__name__}) - K_types = OrderedDict([('/'.join(['K-term', k]), v) for k, v in K_types.items()]) - - # Nodes in the core template graph - LK_types = OrderedDict() - LK_types.update({'STEin': qa.ste.STEActivation.__name__}) - LK_types.update({'Conv': qa.inq.INQConv2d.__name__}) - LK_types.update({'BatchNorm': nn.BatchNorm2d.__name__}) - LK_types.update({'ReLU': nn.ReLU.__name__}) - LK_types.update({'STEout': qa.ste.STEActivation.__name__}) - LK_types = OrderedDict([('/'.join(['L-term', k]), v) for k, v in LK_types.items()]) - - # Nodes in the core replacement graph - RK_types = OrderedDict() - RK_types.update({'TWConv': nn.Conv2d.__name__}) - RK_types.update({'XPAffine': nn.Conv2d.__name__}) - RK_types.update({'S&C': qg.graphs.ShiftAndClip.__name__}) - RK_types = OrderedDict([('/'.join(['R-term', k]), v) for k, v in RK_types.items()]) - - K_node_IDs = list(K_types.keys()) - LK_node_IDs = list(LK_types.keys()) - RK_node_IDs = list(RK_types.keys()) - - # define the template graph L [L-term] - L_node_IDs = [K_node_IDs[0]] + LK_node_IDs + [K_node_IDs[-1]] - self.L = nx.DiGraph() - # Define arcs between nodes in full template graph - self.L.add_edges_from({(u, v) for u, v in zip(L_node_IDs[:-1], L_node_IDs[1:])}) - - # Here, graph is only operation nodes - # Necessary for seeker - nx.set_node_attributes(self.L, {vL: Bipartite.KERNEL for vL in set(self.L.nodes)}, 'bipartite') - nx.set_node_attributes(self.L, {**K_types, **LK_types}, 'type') - - # define the context (sub-)graph K [K-term] - VK = set(K_node_IDs) # precision tunnel nodes define the context graph - self.K = self.L.subgraph(VK) - - # define the template (sub-)graph L\K - VLK = set(self.L.nodes).difference(set(self.K.nodes)) - self.LK = self.L.subgraph(VLK) - - # define the replacement (sub-)graph R\K ["gluing" R\K to K yields the graph R, i.e., the R-term] - self.RK = nx.DiGraph() - self.RK.add_edges_from({(u, v) for u, v in zip(RK_node_IDs[:-1], RK_node_IDs[1:])}) - nx.set_node_attributes(self.RK, {vRK: Bipartite.KERNEL for vRK in set(self.RK.nodes)}, 'bipartite') - nx.set_node_attributes(self.RK, RK_types, 'type') - - # define the arcs that go from the vertices of K to those of R\K, and viceversa - E_K2RK = {(K_node_IDs[0], RK_node_IDs[0])} - E_RK2K = {(RK_node_IDs[-1], K_node_IDs[-1])} - E_K2RK2K = E_K2RK | E_RK2K - # disintegrate `E_K2RK` and `E_RK2K` along fibres to speed up rule application - # A fibre is kind of like fixing one argument of a two input one output function and looking at all possible outputs - self.F_K2RK = {vK: set(arc for arc in E_K2RK if arc[0] == vK) for vK in set(self.K.nodes)} - self.F_RK2K = {vK: set(arc for arc in E_RK2K if arc[1] == vK) for vK in set(self.K.nodes)} - - # # glue together the (sub-)graphs L\K and R\K along the vertices of K - # self.S = nx.compose(self.L, self.RK) - # self.S.add_edges_from(E_K2RK2K) - - # since the GRR's L-term has been modified, rebuild the seeker - self.seeker = Seeker(self.L) - - # this machinery can generate always-new identifiers for different rule applications - self._counter = itertools.count() - - def _get_rule_count(self): - rule_count = ''.join(['FINQBNSTETA', __NODE_ID_FORMAT__.format(next(self._counter))]) - return rule_count - - def core(self, HI, g, nodes_dict): - - # generate the substitute (sub-)graph J\I - rule_count = self._get_rule_count() - g_RK2JI = {vRK: '_'.join([rule_count, vRK.replace('R-term/', '')]) for vRK in set(self.RK.nodes)} - JI = nx.relabel_nodes(self.RK, g_RK2JI, copy=True) - - # get pointers to the old modules; - # these pointers will enable two actions: - # 1. extracting the arguments required to perform the folding - # 2. extracting the parameters to instantiate the new modules - g_L2H = {vL: vH for vH, vL in g.items()} - mstein = nodes_dict[g_L2H['/'.join(['L-term', 'STEin'])]].nobj - minq2d = nodes_dict[g_L2H['/'.join(['L-term', 'Conv'])]].nobj - mbn2d = nodes_dict[g_L2H['/'.join(['L-term', 'BatchNorm'])]].nobj - msteout = nodes_dict[g_L2H['/'.join(['L-term', 'STEout'])]].nobj - - # fold - weight, gamma, beta = foldsteinqconvbnste(mstein.num_levels, mstein.abs_max_value, - minq2d.weight_frozen, - mbn2d.running_mean, mbn2d.running_var, mbn2d.eps, mbn2d.weight, mbn2d.bias, - msteout.num_levels, msteout.abs_max_value, - gamma_int_bits=self._gamma_int_bits, gamma_frac_bits=self._gamma_frac_bits, - beta_int_bits=self._beta_int_bits, beta_frac_bits=self._beta_frac_bits) - - # build the new modules - mtwconv = nn.Conv2d(minq2d.in_channels, minq2d.out_channels, minq2d.kernel_size, - stride=minq2d.stride, padding=minq2d.padding, dilation=minq2d.dilation, groups=minq2d.groups, - bias=minq2d.bias is not None).to(torch.device('cpu')) - mtwconv.weight.data = weight - - mxpaffine = nn.Conv2d(minq2d.out_channels, minq2d.out_channels, 1, - stride=1, padding=0, groups=minq2d.out_channels, - bias=True).to(torch.device('cpu')) - mxpaffine.weight.data = gamma - mxpaffine.bias.data = beta - - msandc = qg.graphs.ShiftAndClip(n_bits=math.ceil(math.log(msteout.num_levels, 2)), - shift=self._gamma_frac_bits, - signed=True, only_positive=True).to(torch.device('cpu')) - - # register the newly created nodes - vJI_2_ptnode = {} - vJI_2_ptnode[g_RK2JI['/'.join(['R-term', 'TWConv'])]] = PyTorchNode(mtwconv) - vJI_2_ptnode[g_RK2JI['/'.join(['R-term', 'XPAffine'])]] = PyTorchNode(mxpaffine) - vJI_2_ptnode[g_RK2JI['/'.join(['R-term', 'S&C'])]] = PyTorchNode(msandc) - - return JI, vJI_2_ptnode - - # G: Full/original graph - # nodes_dict: Mapping between node identifiers of G and actual underlying objects - # g: One instance of all occurences of the template in G, i.e. one application point for the replacement rule -> one morphism - def apply(self, G, nodes_dict, g): - - # create new containers - G = G.copy() - # Dictionary mapping of node identifiers to a payload - # keys in nodes_dict should be the same as G.nodes - nodes_dict = {**nodes_dict} - - # characterise the match graph H - # Occurence of template in the graph - # SPMATTEO: Some assumptions to discuss - VI = {vH for vH, vL in g.items() if vL in set(self.K.nodes)} # Occurence of context - VHI = {vH for vH, vL in g.items() if vL not in set(self.K.nodes)} # Occurence of core template - HI = G.subgraph(VHI) # HI is the subgraph induced by the set of nodes VHI - - # generate the substitute (sub-)graph J\I (completely detached from G) - # Instantiate blueprint of the replacement graph - JI, vJI_2_ptnode = self.core(HI, g, nodes_dict) - - # add the substitute (sub-)graph J\I to the main graph G - G = nx.compose(G, JI) # G now has two connected but 'independent' subgraphs - nodes_dict.update(vJI_2_ptnode) # Add new payloads from substitute graph - - # glue the substitute (sub-)graph J\I to the interface (sub-)graph I - JI2RK_morphisms = Seeker(self.RK).get_morphisms(JI) - assert len(JI2RK_morphisms) == 1 - g_JI2RK = JI2RK_morphisms[0] - g_RK2JI = {vRK: vJI for vJI, vRK in g_JI2RK.items()} - for vI in VI: # for each node in the interface subgraph of G - vK = g[vI] - G.add_edges_from({(vI, g_RK2JI[vRK]) for (_, vRK) in self.F_K2RK[vK]}) # incoming interface connections from G to substitute graph - G.add_edges_from({(g_RK2JI[vRK], vI) for (vRK, _) in self.F_RK2K[vK]}) # outcoming interface connections from substitute graph to G - # the new modules are fully integerized, so the precision tunnel should not embed integer numbers in floating point numbers - # Specific to integer arithmetic transformation -> No relation to graph editing, per-se - if nodes_dict[vI].ntype == qg.graphs.HelperOutputPrecisionTunnel.__name__: - nodes_dict[vI] = PyTorchNode(qg.graphs.HelperOutputPrecisionTunnel(1.0)) - elif nodes_dict[vI].ntype == qg.graphs.HelperInputPrecisionTunnel.__name__: - nodes_dict[vI] = PyTorchNode(qg.graphs.HelperInputPrecisionTunnel(1.0)) - else: - raise TypeError # interface nodes should be objects of class `qg.graphs.HelperPrecisionTunnel` only - - # discard the match (sub-)graph H\I - # Assumption: removing a node also removes all arcs pointing to or from that node - G.remove_nodes_from(set(HI.nodes)) - - # Remove the payload, i.e. underying objects, accordingly - for vHI in VHI: - del nodes_dict[vHI] - - return G, nodes_dict - - def seek(self, G, nodes_dict): - gs = self.seeker.get_morphisms(G) - return gs - - -class FoldSTEINQConvBNSTETypeBRule(DPORule): # w/o max pooling - - def __init__(self, gamma_int_bits=10, gamma_frac_bits=17, beta_int_bits=8, beta_frac_bits=0): - - self._gamma_int_bits = gamma_int_bits - self._gamma_frac_bits = gamma_frac_bits - self._beta_int_bits = beta_int_bits - self._beta_frac_bits = beta_frac_bits - - K_types = OrderedDict() - K_types.update({'HPTout': qg.graphs.HelperOutputPrecisionTunnel.__name__}) - K_types.update({'HPTin': qg.graphs.HelperInputPrecisionTunnel.__name__}) - K_types = OrderedDict([('/'.join(['K-term', k]), v) for k, v in K_types.items()]) - - LK_types = OrderedDict() - LK_types.update({'STEin': qa.ste.STEActivation.__name__}) - LK_types.update({'Conv': qa.inq.INQConv2d.__name__}) - LK_types.update({'BatchNorm': nn.BatchNorm2d.__name__}) - LK_types.update({'ReLU': nn.ReLU.__name__}) - LK_types.update({'MaxPool': nn.MaxPool2d.__name__}) - LK_types.update({'STEout': qa.ste.STEActivation.__name__}) - LK_types = OrderedDict([('/'.join(['L-term', k]), v) for k, v in LK_types.items()]) - - RK_types = OrderedDict() - RK_types.update({'TWConv': nn.Conv2d.__name__}) - RK_types.update({'XPAffine': nn.Conv2d.__name__}) - RK_types.update({'S&C': qg.graphs.ShiftAndClip.__name__}) - RK_types.update({'MaxPool': nn.MaxPool2d.__name__}) - RK_types = OrderedDict([('/'.join(['R-term', k]), v) for k, v in RK_types.items()]) - - K_node_IDs = list(K_types.keys()) - LK_node_IDs = list(LK_types.keys()) - RK_node_IDs = list(RK_types.keys()) - - # define the template graph L [L-term] - L_node_IDs = [K_node_IDs[0]] + LK_node_IDs + [K_node_IDs[-1]] - self.L = nx.DiGraph() - self.L.add_edges_from({(u, v) for u, v in zip(L_node_IDs[:-1], L_node_IDs[1:])}) - nx.set_node_attributes(self.L, {vL: Bipartite.KERNEL for vL in set(self.L.nodes)}, 'bipartite') - nx.set_node_attributes(self.L, {**K_types, **LK_types}, 'type') - - # define the context (sub-)graph K [K-term] - VK = set(K_node_IDs) # precision tunnel nodes define the context graph - self.K = self.L.subgraph(VK) - - # define the template (sub-)graph L\K - VLK = set(self.L.nodes).difference(set(self.K.nodes)) - self.LK = self.L.subgraph(VLK) - - # define the replacement (sub-)graph R\K ["gluing" R\K to K yields the graph R, i.e., the R-term] - self.RK = nx.DiGraph() - self.RK.add_edges_from({(u, v) for u, v in zip(RK_node_IDs[:-1], RK_node_IDs[1:])}) - nx.set_node_attributes(self.RK, {vRK: Bipartite.KERNEL for vRK in set(self.RK.nodes)}, 'bipartite') - nx.set_node_attributes(self.RK, RK_types, 'type') - - # define the arcs that go from the vertices of K to those of R\K, and viceversa - E_K2RK = {(K_node_IDs[0], RK_node_IDs[0])} - E_RK2K = {(RK_node_IDs[-1], K_node_IDs[-1])} - E_K2RK2K = E_K2RK | E_RK2K - # disintegrate `E_K2RK` and `E_RK2K` along fibres to speed up rule application - self.F_K2RK = {vK: set(arc for arc in E_K2RK if arc[0] == vK) for vK in set(self.K.nodes)} - self.F_RK2K = {vK: set(arc for arc in E_RK2K if arc[1] == vK) for vK in set(self.K.nodes)} - - # # glue together the (sub-)graphs L\K and R\K along the vertices of K - # self.S = nx.compose(self.L, self.RK) - # self.S.add_edges_from(E_K2RK2K) - - # since the GRR's L-term has been modified, rebuild the seeker - self.seeker = Seeker(self.L) - - # this machinery can generate always-new identifiers for different rule applications - self._counter = itertools.count() - - def _get_rule_count(self): - rule_count = ''.join(['FINQBNSTETB', __NODE_ID_FORMAT__.format(next(self._counter))]) - return rule_count - - def core(self, HI, g, nodes_dict): - - # generate the substitute (sub-)graph J\I - rule_count = self._get_rule_count() - g_RK2JI = {vRK: '_'.join([rule_count, vRK.replace('R-term/', '')]) for vRK in set(self.RK.nodes)} - JI = nx.relabel_nodes(self.RK, g_RK2JI, copy=True) - - # get pointers to the old modules; - # these pointers will enable two actions: - # 1. extracting the arguments required to perform the folding - # 2. extracting the parameters to instantiate the new modules - g_L2H = {vL: vH for vH, vL in g.items()} - mstein = nodes_dict[g_L2H['/'.join(['L-term', 'STEin'])]].nobj - minq2d = nodes_dict[g_L2H['/'.join(['L-term', 'Conv'])]].nobj - mbn2d = nodes_dict[g_L2H['/'.join(['L-term', 'BatchNorm'])]].nobj - msteout = nodes_dict[g_L2H['/'.join(['L-term', 'STEout'])]].nobj - mmxpold = nodes_dict[g_L2H['/'.join(['L-term', 'MaxPool'])]].nobj - - # fold - weight, gamma, beta = foldsteinqconvbnste(mstein.num_levels, mstein.abs_max_value, - minq2d.weight_frozen, - mbn2d.running_mean, mbn2d.running_var, mbn2d.eps, mbn2d.weight, mbn2d.bias, - msteout.num_levels, msteout.abs_max_value, - gamma_int_bits=self._gamma_int_bits, gamma_frac_bits=self._gamma_frac_bits, - beta_int_bits=self._beta_int_bits, beta_frac_bits=self._beta_frac_bits) - - # build the new modules - mtwconv = nn.Conv2d(minq2d.in_channels, minq2d.out_channels, minq2d.kernel_size, - stride=minq2d.stride, padding=minq2d.padding, dilation=minq2d.dilation, groups=minq2d.groups, - bias=minq2d.bias is not None).to(torch.device('cpu')) - mtwconv.weight.data = weight - - mxpaffine = nn.Conv2d(minq2d.out_channels, minq2d.out_channels, 1, - stride=1, padding=0, groups=minq2d.out_channels, - bias=True).to(torch.device('cpu')) - mxpaffine.weight.data = gamma - mxpaffine.bias.data = beta - - msandc = qg.graphs.ShiftAndClip(n_bits=math.ceil(math.log(msteout.num_levels, 2)), - shift=self._gamma_frac_bits, - signed=True, only_positive=True).to(torch.device('cpu')) - - mmxpnew = nn.MaxPool2d(kernel_size=mmxpold.kernel_size, stride=mmxpold.stride, padding=mmxpold.padding) - - # register the newly created nodes - vJI_2_ptnode = {} - vJI_2_ptnode[g_RK2JI['/'.join(['R-term', 'TWConv'])]] = PyTorchNode(mtwconv) - vJI_2_ptnode[g_RK2JI['/'.join(['R-term', 'XPAffine'])]] = PyTorchNode(mxpaffine) - vJI_2_ptnode[g_RK2JI['/'.join(['R-term', 'S&C'])]] = PyTorchNode(msandc) - vJI_2_ptnode[g_RK2JI['/'.join(['R-term', 'MaxPool'])]] = PyTorchNode(mmxpnew) - - return JI, vJI_2_ptnode - - def apply(self, G, nodes_dict, g): - - # create new containers - G = G.copy() - nodes_dict = {**nodes_dict} - - # characterise the match graph H - VI = {vH for vH, vL in g.items() if vL in set(self.K.nodes)} - VHI = {vH for vH, vL in g.items() if vL not in set(self.K.nodes)} - HI = G.subgraph(VHI) - - # generate the substitute (sub-)graph J\I - JI, vJI_2_ptnode = self.core(HI, g, nodes_dict) - - # add the substitute (sub-)graph J\I to the main graph G - G = nx.compose(G, JI) - nodes_dict.update(vJI_2_ptnode) - - # glue the substitute (sub-)graph J\I to the interface (sub-)graph I - JI2RK_morphisms = Seeker(self.RK).get_morphisms(JI) - assert len(JI2RK_morphisms) == 1 - g_JI2RK = JI2RK_morphisms[0] - g_RK2JI = {vRK: vJI for vJI, vRK in g_JI2RK.items()} - for vI in VI: - vK = g[vI] - G.add_edges_from({(vI, g_RK2JI[vRK]) for (_, vRK) in self.F_K2RK[vK]}) - G.add_edges_from({(g_RK2JI[vRK], vI) for (vRK, _) in self.F_RK2K[vK]}) - # the new modules are fully integerized, so the precision tunnel should not embed integer numbers in floating point numbers - if nodes_dict[vI].ntype == qg.graphs.HelperOutputPrecisionTunnel.__name__: - nodes_dict[vI] = PyTorchNode(qg.graphs.HelperOutputPrecisionTunnel(1.0)) - elif nodes_dict[vI].ntype == qg.graphs.HelperInputPrecisionTunnel.__name__: - nodes_dict[vI] = PyTorchNode(qg.graphs.HelperInputPrecisionTunnel(1.0)) - else: - raise TypeError # interface nodes should be objects of class `qg.graphs.HelperPrecisionTunnel` only - - # discard the match (sub-)graph H\I - G.remove_nodes_from(set(HI.nodes)) - for vHI in VHI: - del nodes_dict[vHI] - - return G, nodes_dict - - def seek(self, G, nodes_dict): - gs = self.seeker.get_morphisms(G) - return gs - - -class FoldConvBNSTERule(DPORule): - - def __init__(self): - - K_types = OrderedDict() - K_types.update({'HI': qg.graphs.HelperInput.__name__}) - K_types.update({'HPTin': qg.graphs.HelperInputPrecisionTunnel.__name__}) - K_types = OrderedDict([('/'.join(['K-term', k]), v) for k, v in K_types.items()]) - - LK_types = OrderedDict() - LK_types.update({'Conv': nn.Conv2d.__name__}) - LK_types.update({'BatchNorm': nn.BatchNorm2d.__name__}) - LK_types.update({'ReLU': nn.ReLU.__name__}) - LK_types.update({'STE': qa.ste.STEActivation.__name__}) - LK_types = OrderedDict([('/'.join(['L-term', k]), v) for k, v in LK_types.items()]) - - RK_types = OrderedDict() - RK_types.update({'Conv': nn.Conv2d.__name__}) - RK_types.update({'F&C': qg.graphs.FloorAndClip.__name__}) - RK_types = OrderedDict([('/'.join(['R-term', k]), v) for k, v in RK_types.items()]) - - K_node_IDs = list(K_types.keys()) - LK_node_IDs = list(LK_types.keys()) - RK_node_IDs = list(RK_types.keys()) - - # define the template graph L [L-term] - L_node_IDs = [K_node_IDs[0]] + LK_node_IDs + [K_node_IDs[-1]] - self.L = nx.DiGraph() - self.L.add_edges_from({(u, v) for u, v in zip(L_node_IDs[:-1], L_node_IDs[1:])}) - nx.set_node_attributes(self.L, {vL: Bipartite.KERNEL for vL in set(self.L.nodes)}, 'bipartite') - nx.set_node_attributes(self.L, {**K_types, **LK_types}, 'type') - - # define the context (sub-)graph K [K-term] - VK = set(K_node_IDs) # precision tunnel nodes define the context graph - self.K = self.L.subgraph(VK) - - # define the template (sub-)graph L\K - VLK = set(self.L.nodes).difference(set(self.K.nodes)) - self.LK = self.L.subgraph(VLK) - - # define the replacement (sub-)graph R\K ["gluing" R\K to K yields the graph R, i.e., the R-term] - self.RK = nx.DiGraph() - self.RK.add_edges_from({(u, v) for u, v in zip(RK_node_IDs[:-1], RK_node_IDs[1:])}) - nx.set_node_attributes(self.RK, {vRK: Bipartite.KERNEL for vRK in set(self.RK.nodes)}, 'bipartite') - nx.set_node_attributes(self.RK, RK_types, 'type') - - # define the arcs that go from the vertices of K to those of R\K, and viceversa - E_K2RK = {(K_node_IDs[0], RK_node_IDs[0])} - E_RK2K = {(RK_node_IDs[-1], K_node_IDs[-1])} - E_K2RK2K = E_K2RK | E_RK2K - # disintegrate `E_K2RK` and `E_RK2K` along fibres to speed up rule application - self.F_K2RK = {vK: set(arc for arc in E_K2RK if arc[0] == vK) for vK in set(self.K.nodes)} - self.F_RK2K = {vK: set(arc for arc in E_RK2K if arc[1] == vK) for vK in set(self.K.nodes)} - - # # glue together the (sub-)graphs L\K and R\K along the vertices of K - # self.S = nx.compose(self.L, self.RK) - # self.S.add_edges_from(E_K2RK2K) - - # since the GRR's L-term has been modified, rebuild the seeker - self.seeker = Seeker(self.L) - - # this machinery can generate always-new identifiers for different rule applications - self._counter = itertools.count() - - def _get_rule_count(self): - rule_count = ''.join(['FCBNSTE', __NODE_ID_FORMAT__.format(next(self._counter))]) - return rule_count - - def core(self, HI, g, nodes_dict): - - # generate the substitute (sub-)graph J\I - rule_count = self._get_rule_count() - g_RK2JI = {vRK: '_'.join([rule_count, vRK.replace('R-term/', '')]) for vRK in set(self.RK.nodes)} - JI = nx.relabel_nodes(self.RK, g_RK2JI, copy=True) - - # get pointers to the old modules; - # these pointers will enable two actions: - # 1. extracting the arguments required to perform the folding - # 2. extracting the parameters to instantiate the new modules - g_L2H = {vL: vH for vH, vL in g.items()} - mconvold = nodes_dict[g_L2H['/'.join(['L-term', 'Conv'])]].nobj - mbn2d = nodes_dict[g_L2H['/'.join(['L-term', 'BatchNorm'])]].nobj - mste = nodes_dict[g_L2H['/'.join(['L-term', 'STE'])]].nobj - - # fold - weight, bias = foldconvbnste(mconvold.weight, - mbn2d.running_mean, mbn2d.running_var, mbn2d.eps, mbn2d.weight, mbn2d.bias, - mste.num_levels, mste.abs_max_value) - - # build the new modules - mconvnew = nn.Conv2d(mconvold.in_channels, mconvold.out_channels, mconvold.kernel_size, - stride=mconvold.stride, padding=mconvold.padding, dilation=mconvold.dilation, groups=mconvold.groups, - bias=True).to(torch.device('cpu')) - mconvnew.weight.data = weight - mconvnew.bias.data = bias - - mfandc = qg.graphs.FloorAndClip(n_bits=math.ceil(math.log(mste.num_levels, 2)), - signed=True, only_positive=True).to(torch.device('cpu')) - - # register the newly created nodes - vJI_2_ptnode = {} - vJI_2_ptnode[g_RK2JI['/'.join(['R-term', 'Conv'])]] = PyTorchNode(mconvnew) - vJI_2_ptnode[g_RK2JI['/'.join(['R-term', 'F&C'])]] = PyTorchNode(mfandc) - - return JI, vJI_2_ptnode - - def apply(self, G, nodes_dict, g): - - # create new containers - G = G.copy() - nodes_dict = {**nodes_dict} - - # characterise the match graph H - VI = {vH for vH, vL in g.items() if vL in set(self.K.nodes)} - VHI = {vH for vH, vL in g.items() if vL not in set(self.K.nodes)} - HI = G.subgraph(VHI) - - # generate the substitute (sub-)graph J\I - JI, vJI_2_ptnode = self.core(HI, g, nodes_dict) - - # add the substitute (sub-)graph J\I to the main graph G - G = nx.compose(G, JI) - nodes_dict.update(vJI_2_ptnode) - - # glue the substitute (sub-)graph J\I to the interface (sub-)graph I - JI2RK_morphisms = Seeker(self.RK).get_morphisms(JI) - assert len(JI2RK_morphisms) == 1 - g_JI2RK = JI2RK_morphisms[0] - g_RK2JI = {vRK: vJI for vJI, vRK in g_JI2RK.items()} - for vI in VI: - vK = g[vI] - G.add_edges_from({(vI, g_RK2JI[vRK]) for (_, vRK) in self.F_K2RK[vK]}) - G.add_edges_from({(g_RK2JI[vRK], vI) for (vRK, _) in self.F_RK2K[vK]}) - # the new modules are fully integerized, so the precision tunnel should not embed integer numbers in floating point numbers - if nodes_dict[vI].ntype == qg.graphs.HelperInput.__name__: - pass - elif nodes_dict[vI].ntype == qg.graphs.HelperInputPrecisionTunnel.__name__: - nodes_dict[vI] = PyTorchNode(qg.graphs.HelperInputPrecisionTunnel(1.0)) - else: - raise TypeError # interface nodes should be objects of class `qg.graphs.HelperPrecisionTunnel` only - - # discard the match (sub-)graph H\I - G.remove_nodes_from(set(HI.nodes)) - for vHI in VHI: - del nodes_dict[vHI] - - return G, nodes_dict - - def seek(self, G, nodes_dict): - gs = self.seeker.get_morphisms(G) - return gs - - -class FoldSTEINQConvBNRule(DPORule): - - def __init__(self): - - K_types = OrderedDict() - K_types.update({'HI': qg.graphs.HelperOutputPrecisionTunnel.__name__}) - K_types.update({'MaxPool': nn.MaxPool2d.__name__}) - K_types = OrderedDict([('/'.join(['K-term', k]), v) for k, v in K_types.items()]) - - LK_types = OrderedDict() - LK_types.update({'STE': qa.ste.STEActivation.__name__}) - LK_types.update({'INQConv': qa.inq.INQConv2d.__name__}) - LK_types.update({'BatchNorm': nn.BatchNorm2d.__name__}) - LK_types.update({'ReLU': nn.ReLU.__name__}) - LK_types = OrderedDict([('/'.join(['L-term', k]), v) for k, v in LK_types.items()]) - - RK_types = OrderedDict() - RK_types.update({'Conv': nn.Conv2d.__name__}) - RK_types.update({'ReLU': nn.ReLU.__name__}) - RK_types = OrderedDict([('/'.join(['R-term', k]), v) for k, v in RK_types.items()]) - - K_node_IDs = list(K_types.keys()) - LK_node_IDs = list(LK_types.keys()) - RK_node_IDs = list(RK_types.keys()) - - # define the template graph L [L-term] - L_node_IDs = [K_node_IDs[0]] + LK_node_IDs + [K_node_IDs[-1]] - self.L = nx.DiGraph() - self.L.add_edges_from({(u, v) for u, v in zip(L_node_IDs[:-1], L_node_IDs[1:])}) - nx.set_node_attributes(self.L, {vL: Bipartite.KERNEL for vL in set(self.L.nodes)}, 'bipartite') - nx.set_node_attributes(self.L, {**K_types, **LK_types}, 'type') - - # define the context (sub-)graph K [K-term] - VK = set(K_node_IDs) # precision tunnel nodes define the context graph - self.K = self.L.subgraph(VK) - - # define the template (sub-)graph L\K - VLK = set(self.L.nodes).difference(set(self.K.nodes)) - self.LK = self.L.subgraph(VLK) - - # define the replacement (sub-)graph R\K ["gluing" R\K to K yields the graph R, i.e., the R-term] - self.RK = nx.DiGraph() - self.RK.add_edges_from({(u, v) for u, v in zip(RK_node_IDs[:-1], RK_node_IDs[1:])}) - nx.set_node_attributes(self.RK, {vRK: Bipartite.KERNEL for vRK in set(self.RK.nodes)}, 'bipartite') - nx.set_node_attributes(self.RK, RK_types, 'type') - - # define the arcs that go from the vertices of K to those of R\K, and viceversa - E_K2RK = {(K_node_IDs[0], RK_node_IDs[0])} - E_RK2K = {(RK_node_IDs[-1], K_node_IDs[-1])} - E_K2RK2K = E_K2RK | E_RK2K - # disintegrate `E_K2RK` and `E_RK2K` along fibres to speed up rule application - self.F_K2RK = {vK: set(arc for arc in E_K2RK if arc[0] == vK) for vK in set(self.K.nodes)} - self.F_RK2K = {vK: set(arc for arc in E_RK2K if arc[1] == vK) for vK in set(self.K.nodes)} - - # # glue together the (sub-)graphs L\K and R\K along the vertices of K - # self.S = nx.compose(self.L, self.RK) - # self.S.add_edges_from(E_K2RK2K) - - # since the GRR's L-term has been modified, rebuild the seeker - self.seeker = Seeker(self.L) - - # this machinery can generate always-new identifiers for different rule applications - self._counter = itertools.count() - - def _get_rule_count(self): - rule_count = ''.join(['FSTEINQBN', __NODE_ID_FORMAT__.format(next(self._counter))]) - return rule_count - - def core(self, HI, g, nodes_dict): - - # generate the substitute (sub-)graph J\I - rule_count = self._get_rule_count() - g_RK2JI = {vRK: '_'.join([rule_count, vRK.replace('R-term/', '')]) for vRK in set(self.RK.nodes)} - JI = nx.relabel_nodes(self.RK, g_RK2JI, copy=True) - - # get pointers to the old modules; - # these pointers will enable two actions: - # 1. extracting the arguments required to perform the folding - # 2. extracting the parameters to instantiate the new modules - g_L2H = {vL: vH for vH, vL in g.items()} - mste = nodes_dict[g_L2H['/'.join(['L-term', 'STE'])]].nobj - minq2d = nodes_dict[g_L2H['/'.join(['L-term', 'INQConv'])]].nobj - mbn2d = nodes_dict[g_L2H['/'.join(['L-term', 'BatchNorm'])]].nobj - mreluold = nodes_dict[g_L2H['/'.join(['L-term', 'ReLU'])]].nobj - - # fold - weight, bias = foldsteinqconvbn(mste.num_levels, mste.abs_max_value, - minq2d.weight_frozen, - mbn2d.running_mean, mbn2d.running_var, mbn2d.eps, mbn2d.weight, mbn2d.bias) - - # build the new modules - mconv = nn.Conv2d(minq2d.in_channels, minq2d.out_channels, minq2d.kernel_size, - stride=minq2d.stride, padding=minq2d.padding, dilation=minq2d.dilation, groups=minq2d.groups, - bias=True).to(torch.device('cpu')) - mconv.weight.data = weight - mconv.bias.data = bias - - mrelunew = nn.ReLU(inplace=True) - - # register the newly created nodes - vJI_2_ptnode = {} - vJI_2_ptnode[g_RK2JI['/'.join(['R-term', 'Conv'])]] = PyTorchNode(mconv) - vJI_2_ptnode[g_RK2JI['/'.join(['R-term', 'ReLU'])]] = PyTorchNode(mrelunew) - - return JI, vJI_2_ptnode - - def apply(self, G, nodes_dict, g): - - # create new containers - G = G.copy() - nodes_dict = {**nodes_dict} - - # characterise the match graph H - VI = {vH for vH, vL in g.items() if vL in set(self.K.nodes)} - VHI = {vH for vH, vL in g.items() if vL not in set(self.K.nodes)} - HI = G.subgraph(VHI) - - # generate the substitute (sub-)graph J\I - JI, vJI_2_ptnode = self.core(HI, g, nodes_dict) - - # add the substitute (sub-)graph J\I to the main graph G - G = nx.compose(G, JI) - nodes_dict.update(vJI_2_ptnode) - - # glue the substitute (sub-)graph J\I to the interface (sub-)graph I - JI2RK_morphisms = Seeker(self.RK).get_morphisms(JI) - assert len(JI2RK_morphisms) == 1 - g_JI2RK = JI2RK_morphisms[0] - g_RK2JI = {vRK: vJI for vJI, vRK in g_JI2RK.items()} - for vI in VI: - vK = g[vI] - G.add_edges_from({(vI, g_RK2JI[vRK]) for (_, vRK) in self.F_K2RK[vK]}) - G.add_edges_from({(g_RK2JI[vRK], vI) for (vRK, _) in self.F_RK2K[vK]}) - # the new modules are fully integerized, so the precision tunnel should not embed integer numbers in floating point numbers - if nodes_dict[vI].ntype == qg.graphs.HelperOutputPrecisionTunnel.__name__: - nodes_dict[vI] = PyTorchNode(qg.graphs.HelperOutputPrecisionTunnel(1.0)) - elif nodes_dict[vI].ntype == nn.MaxPool2d.__name__: - pass - else: - raise TypeError # interface nodes should be objects of class `qg.graphs.HelperPrecisionTunnel` only - - # discard the match (sub-)graph H\I - G.remove_nodes_from(set(HI.nodes)) - for vHI in VHI: - del nodes_dict[vHI] - - return G, nodes_dict - - def seek(self, G, nodes_dict): - gs = self.seeker.get_morphisms(G) - return gs - diff --git a/backends/twn_accelerator/grrules/folding.py b/backends/twn_accelerator/grrules/folding.py deleted file mode 100644 index a0c32f0..0000000 --- a/backends/twn_accelerator/grrules/folding.py +++ /dev/null @@ -1,178 +0,0 @@ -# -# folding.py -# -# Author(s): -# Matteo Spallanzani -# -# Copyright (c) 2020-2021 ETH Zurich. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -import numpy as np -import torch - - -def cast_gamma(gamma_fp64, int_bits=10, frac_bits=17): - - quantum = 2**(-frac_bits) - gamma_fp64 /= quantum - gamma_uint32 = gamma_fp64.astype('= 0.5).astype(np.int32) # round to the closest value (should minimise the proability of errors under the assumption that the inputs are uniformly distributed) - # beta_int32 += (probs > 0.).astype(np.int32) # round to ceiling - - assert np.all(np.logical_and(np.abs(beta_int32) <= (2**(int_bits + beta_frac_bits) / 2), beta_int32 != (2**(int_bits + beta_frac_bits) / 2))) # each beta is a SIGNED integer with `int_bits + beta_frac_bits` precision - - beta_int32 *= 2**n_truncated_bits - beta_int32_into_fp64 = beta_int32.astype(np.float64) - - return beta_int32_into_fp64 - - -def foldsteinqconvbnste(n_in, m_in, weight, mu, sigma, eps, gamma, beta, n_out, m_out, gamma_int_bits=10, gamma_frac_bits=17, beta_int_bits=8, beta_frac_bits=0): - - def torch2numpyfp64(x): - return x.detach().cpu().numpy().astype(np.float64) - - m_in = torch2numpyfp64(m_in) - weight = torch2numpyfp64(weight) - mu = torch2numpyfp64(mu) - sigma = torch2numpyfp64(sigma) - gamma = torch2numpyfp64(gamma) - beta = torch2numpyfp64(beta) - m_out = torch2numpyfp64(m_out) - - # see batch normalisation docs in PyTorch - sigma = np.sqrt(sigma + eps) - # STE quanta - eps_in = (2 * m_in) / (n_in - 1) - eps_out = (2 * m_out) / (n_out - 1) - - # REORDER - - # compensate for negative gammas - flip = np.sign(gamma) - w_tmp = weight.transpose(1, 2, 3, 0) - w_tmp *= flip - fweight = w_tmp.transpose(3, 0, 1, 2) - - # fold gamma - fgamma = (eps_in * gamma) / (eps_out * sigma) - fgamma *= flip - fgamma = fgamma.reshape(-1, 1, 1, 1) - - # fold beta - # # reordering rule when `x` is replaced by `x + C` - # fwsum = w.reshape((w.shape[0], -1)).sum(axis=1) - # fbeta = ((((C * fwsum) - mu) * gamma / sigma) + beta) / eps_out # "rounding" version - # fbeta = ((((C * fwsum) - mu) * gamma / sigma) + beta) / eps_out + 0.5 # "flooring" version - # standard reordering rule - # fbeta = (((-mu * gamma) / sigma) + beta) / eps_out # "rounding version - fbeta = (((-mu * gamma) / sigma) + beta) / eps_out + 0.5 # "flooring" version - - # CAST - - fgamma = cast_gamma(fgamma, int_bits=gamma_int_bits, frac_bits=gamma_frac_bits) - fbeta = cast_beta(fbeta, int_bits=beta_int_bits, gamma_frac_bits=gamma_frac_bits, beta_frac_bits=beta_frac_bits) - - def numpy2torchfp64(x): - return torch.from_numpy(x.astype(np.float64)) - - return numpy2torchfp64(fweight), numpy2torchfp64(fgamma), numpy2torchfp64(fbeta) - - -def foldconvbnste(weight, mu, sigma, eps, gamma, beta, n_out, m_out): - - def torch2numpyfp64(x): - return x.detach().cpu().numpy().astype(np.float64) - - weight = torch2numpyfp64(weight) - mu = torch2numpyfp64(mu) - sigma = torch2numpyfp64(sigma) - gamma = torch2numpyfp64(gamma) - beta = torch2numpyfp64(beta) - m_out = torch2numpyfp64(m_out) - - # see batch normalisation docs in PyTorch - sigma = np.sqrt(sigma + eps) - # STE quantum - eps_out = (2 * m_out) / (n_out - 1) - - # REORDER - - fweight = (weight * (gamma / sigma).reshape(-1, 1, 1, 1)) / eps_out - fbias = ((-mu * gamma) / sigma + beta) / eps_out + 0.5 - - def numpy2torchfp64(x): - return torch.from_numpy(x.astype(np.float64)) - - return numpy2torchfp64(fweight), numpy2torchfp64(fbias) - - -def foldsteinqconvbn(n_in, m_in, weight, mu, sigma, eps, gamma, beta): - - def torch2numpyfp64(x): - return x.detach().cpu().numpy().astype(np.float64) - - m_in = torch2numpyfp64(m_in) - weight = torch2numpyfp64(weight) - mu = torch2numpyfp64(mu) - sigma = torch2numpyfp64(sigma) - gamma = torch2numpyfp64(gamma) - beta = torch2numpyfp64(beta) - - # see batch normalisation docs in PyTorch - sigma = np.sqrt(sigma + eps) - # STE quantum - eps_in = (2 * m_in) / (n_in - 1) - - # REORDER - - fweight = weight * ((eps_in * gamma) / sigma).reshape(-1, 1, 1, 1) - fbias = (-mu * gamma) / sigma + beta + 0.5 - - def numpy2torchfp64(x): - return torch.from_numpy(x.astype(np.float64)) - - return numpy2torchfp64(fweight), numpy2torchfp64(fbias) - diff --git a/backends/twn_accelerator/twn_accelerator.py b/backends/twn_accelerator/twn_accelerator.py index 42de16f..df55c81 100644 --- a/backends/twn_accelerator/twn_accelerator.py +++ b/backends/twn_accelerator/twn_accelerator.py @@ -28,7 +28,7 @@ from quantlib.algorithms.inq import INQConv2d, INQConv1d from quantlib.algorithms.ste import STEActivation -from quantlib.graphs.analyse import Node +from quantlib.editing.lightweight import LightweightNode from .layers import layer_has_modules from mako.template import Template @@ -110,7 +110,7 @@ class TWNLayer: # for a specific TWN accelerator layer def __init__(self, layer_nodes : list, name : str, params : TWNAccelParams): self.layer_name = name - self.layer_nodes = [m.module if isinstance(m, Node) else m for m in layer_nodes] + self.layer_nodes = [m.module if isinstance(m, LightweightNode) else m for m in layer_nodes] self.params = params @property diff --git a/editing/__init__.py b/editing/__init__.py index a10edea..3b29b44 100755 --- a/editing/__init__.py +++ b/editing/__init__.py @@ -19,6 +19,5 @@ # limitations under the License. # -from . import graphs from . import lightweight from . import fx diff --git a/editing/graphs/__init__.py b/editing/graphs/__init__.py deleted file mode 100755 index 72a8a7c..0000000 --- a/editing/graphs/__init__.py +++ /dev/null @@ -1,34 +0,0 @@ -# -# __init__.py -# -# Author(s): -# Matteo Spallanzani -# -# Copyright (c) 2020-2021 ETH Zurich. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -from .editor import Editor -from . import graphs -from . import grrules -from . import traces -from . import utils - -__all__ = [ - 'editor', - 'graphs', - 'grrules', - 'utils', -] - diff --git a/editing/graphs/editor/__init__.py b/editing/graphs/editor/__init__.py deleted file mode 100644 index e1f5c40..0000000 --- a/editing/graphs/editor/__init__.py +++ /dev/null @@ -1,23 +0,0 @@ -# -# __init__.py -# -# Author(s): -# Matteo Spallanzani -# -# Copyright (c) 2020-2021 ETH Zurich. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -from .editor import Editor - diff --git a/editing/graphs/editor/editor.py b/editing/graphs/editor/editor.py deleted file mode 100644 index 209b43c..0000000 --- a/editing/graphs/editor/editor.py +++ /dev/null @@ -1,237 +0,0 @@ -# -# editor.py -# -# Author(s): -# Matteo Spallanzani -# -# Copyright (c) 2020-2021 ETH Zurich. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -import itertools -import tempfile -from collections import namedtuple -from datetime import datetime -from networkx.algorithms import bipartite -import networkx as nx - -# import graphs -# import utils - -from .. import graphs -from .. import grrules -from .. import utils - - -__FAILURE__ = False -__SUCCESS__ = True - - -Commit = namedtuple('Commit', ['rho', 'g', 'Gprime', 'nodes_dict']) - - -class History(object): - - def __init__(self, nx_graph, nodes_dict): - self._nx_graph = nx_graph # keep track of the original object - self._nodes_dict = nodes_dict - self._undo = [] - self._redo = [] - - def show(self): - print("-- History --") - for i, commit in enumerate(self._undo): - print(i, commit) - - def push(self, commit): - self._undo.append(commit) - self._redo.clear() - - def undo(self, n=1): - for i in range(0, n): - try: - self._redo.append(self._undo.pop()) - except IndexError: - print("Tried to undo {} steps, but history contained just {}. 'Undo' stack has been cleared.".format(n, i)) - break - - def redo(self, n=1): - for i in range(0, n): - try: - self._undo.append(self._redo.pop()) - except IndexError: - print("Tried to redo {} steps, but history contained just {}. 'Redo' stack has been cleared.".format(n, i)) - break - - def clear(self, force=False): - - if not force: - confirmation = input("This action is not reversible. Are you sure that you want to delete all the history? [yes/NO]") - force = confirmation.lower() == 'yes' - - if force: - self._undo.clear() - self._redo.clear() - - -class Editor(object): - - def __init__(self, qlgraph, onlykernel=False, graphviz=False): - - self.qlgraph = qlgraph - - self._input_nodes = None - self._output_nodes = None - - input_datanodes = {n for n, dp in nx.get_node_attributes(self.qlgraph.nx_graph, 'data_partition').items() if dp == graphs.DataPartition.INPUT} - output_datanodes = {n for n, dp in nx.get_node_attributes(self.qlgraph.nx_graph, 'data_partition').items() if dp == graphs.DataPartition.OUTPUT} - - if onlykernel: - - input_opnodes = set(itertools.chain.from_iterable([set([s for s in self.qlgraph.nx_graph.successors(n)]) for n in input_datanodes])) - output_opnodes = set(itertools.chain.from_iterable([set([p for p in self.qlgraph.nx_graph.predecessors(n)]) for n in output_datanodes])) - - self._input_nodes = input_opnodes - self._output_nodes = output_opnodes - - G = bipartite.projected_graph(self.qlgraph.nx_graph, {n for n in self.qlgraph.nx_graph.nodes if self.qlgraph.nx_graph.nodes[n]['bipartite'] == graphs.Bipartite.KERNEL}) - - else: - - self._input_nodes = input_datanodes - self._output_nodes = output_datanodes - - G = self.qlgraph.nx_graph - - nodes_dict = {k: v for k, v in self.qlgraph.nodes_dict.items() if k in G.nodes} - - self._history = History(G, nodes_dict) - self._in_session = False # put a lock on the history by preventing editing actions - self._rho = None # current GRR - self._graphviz = graphviz - self._cache_dir = None - - @property - def G(self): - try: - G = self._history._undo[-1].Gprime - except IndexError: - G = self._history._nx_graph - return G - - @property - def nodes_dict(self): - try: - nodes_dict = self._history._undo[-1].nodes_dict - except IndexError: - nodes_dict = self._history._nodes_dict - return nodes_dict - - def startup(self): - self._cache_dir = tempfile.TemporaryDirectory() - import os - print("Temporary cache directory created at {}".format(os.path.abspath(self._cache_dir.name))) - self._in_session = True - - def pause(self): - self._in_session = False - - def resume(self): - self._in_session = True - - def shutdown(self): - self._rho = None - self._in_session = False - self._apply_changes_to_graph() - self._cache_dir.cleanup() - self._history.clear(force=True) - - def set_grr(self, rho): - self._rho = rho - - def seek(self, **kwargs): - - if self._rho: - gs = self._rho.seek(self.G, self.nodes_dict, **kwargs) - else: - gs = None - print("No rule defined.") - - return gs - - def edit(self, gs=None, **kwargs): - - if self._rho and self._in_session: - - if gs is None: - gs = self.seek(**kwargs) - - for g in gs: - - try: - G_new, nodes_dict_new = self._rho.apply(self.G, self.nodes_dict, g) # derivation - self._history.push(Commit(self._rho, g, G_new, nodes_dict_new)) - status = __SUCCESS__ - - except Exception as e: - print("An issue arose while applying rule {} to graph <{}> at point: ".format(type(self._rho), self.G)) - for vH, vL in g.items(): - print("\t", vH, vL) - print(e) - status = __FAILURE__ - - if (status == __SUCCESS__) and self._graphviz: - self._take_snapshot() - - else: - if self._rho is None: - print("No rule defined for editor object <{}>.".format(self)) - else: - print("Editor object <{}> is not in an editing session.".format(self)) - - def add_io_handles(self): - - if isinstance(self.qlgraph, graphs.PyTorchGraph): - self.startup() - self.set_grr(grrules.AddInputNodeRule()) - self.edit(gs=self.seek(VIs=[[n] for n in self._input_nodes])) - self.set_grr(grrules.AddOutputNodeRule()) - self.edit(gs=self.seek(VIs=[[n] for n in self._output_nodes])) - self.pause() - - def _apply_changes_to_graph(self): - - self.qlgraph.nx_graph = self.G - self.qlgraph.nodes_dict = self.nodes_dict - - def _take_snapshot(self): - filename = datetime.now().strftime("%H:%M:%S_{}_{}".format(len(self._history._undo), type(self._history._undo[-1].rho))) - utils.draw_graph(self.G, self._cache_dir.name, filename) # take a snapshot of the edited graph - -# 1. label graph nodes (node label is usually computed as the aggregation of 1.partition and 2.type, but see COMMENT below) -# 2. define a graph rewriting rule (GRR) -# 3. 'discover' possible application points for the rules -# 4. 'filter' the sequence of application points (possibly NOT automatic) -# 5. 'apply' the rule to the filtered sequence of application points -# - each pair (rule, application_point) is called a 'transform', and the resulting graph is called a 'derivation' -# 6. 'generate_code' for the transformed graph -# 7. 'import_network' from the transformed graph's file - -# [COMMENT] Steps 1 and 2 are usually designed in reversed order: -# - the user first thinks to the rule -# - then decides which "pieces" should be in the label -# which "ingredients" did I use in the past to generate these labels? (my personal "database/record" of use cases) -# - ONNX op type -# - node scope - diff --git a/editing/graphs/graphs/__init__.py b/editing/graphs/graphs/__init__.py deleted file mode 100644 index 87e463d..0000000 --- a/editing/graphs/graphs/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -# -# __init__.py -# -# Author(s): -# Matteo Spallanzani -# -# Copyright (c) 2020-2021 ETH Zurich. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -from .modules import * -from .nodes import Bipartite, DataPartition, __NODE_ID_FORMAT__ -from .nodes import ONNXNode, PyTorchNode -from .graphs import ONNXGraph, PyTorchGraph - diff --git a/editing/graphs/graphs/graphs.py b/editing/graphs/graphs/graphs.py deleted file mode 100644 index a79ae3c..0000000 --- a/editing/graphs/graphs/graphs.py +++ /dev/null @@ -1,276 +0,0 @@ -# -# graphs.py -# -# Author(s): -# Matteo Spallanzani -# -# Copyright (c) 2020-2021 ETH Zurich. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -from packaging import version -import torch.onnx.utils -import torch -import itertools -import networkx as nx -from networkx.algorithms import dag - -from .. import graphs - - -class scope_name_workaround(object): - # this is a necessary "context manager" object for PyTorch >= 1.4 - # (see https://github.com/pytorch/pytorch/issues/33463#issuecomment-606399944) - - def __init__(self): - self.backup = None # store pointer to the definition of the "native" '_slow_forward' - - def __enter__(self): - - def _tracing_name(self_, tracing_state): - - if not tracing_state._traced_module_stack: - return None - - module = tracing_state._traced_module_stack[-1] - for name, child in module.named_children(): - if child is self_: - return name - - return None - - def _slow_forward(self_, *input, **kwargs): - - tracing_state = torch._C._get_tracing_state() - - if not tracing_state or isinstance(self_.forward, torch._C.ScriptMethod): - return self_.forward(*input, **kwargs) # no need to wrap and trace - - if not hasattr(tracing_state, '_traced_module_stack'): - tracing_state._traced_module_stack = [] - - name = _tracing_name(self_, tracing_state) - if name: - tracing_state.push_scope('{}[{}]'.format(self_._get_name(), name)) - else: - tracing_state.push_scope(self_._get_name()) - tracing_state._traced_module_stack.append(self_) - - try: - result = self_.forward(*input, **kwargs) - finally: - tracing_state.pop_scope() - tracing_state._traced_module_stack.pop() - - return result - - self.backup = torch.nn.Module._slow_forward - setattr(torch.nn.Module, '_slow_forward', _slow_forward) # replace '_slow_forward' with the version defined by this context manager - - def __exit__(self, type, value, tb): - setattr(torch.nn.Module, '_slow_forward', self.backup) # restore "native" '_slow_forward' method - - -class ONNXGraph(object): - - def __init__(self, net, dummy_input): - - if not isinstance(dummy_input, tuple): - dummy_input = (dummy_input,) - - assert version.parse(torch.__version__) >= version.parse('1.9.0') - from torch.onnx.symbolic_helper import _set_opset_version - _set_opset_version(11) # opset_version9 does not support `round` ONNX operator (even though docs for PyTorch 1.5.0 suggests so) - - with scope_name_workaround(): - # self.jit_graph, _, _ = torch.onnx.utils._model_to_graph(net, dummy_input, propagate=True, _retain_param_name=True) - self.jit_graph, _, _ = torch.onnx.utils._model_to_graph(net, dummy_input, _retain_param_name=True) - - # At this point, I have a handle on a `torch._C.Graph` object; its - # components are `torch._C.Node` objects, which are abstractions for - # operations; the "data pools" where operands (i.e., inputs) are read - # from and results (i.e., outputs) are written to are `torch._C.Value` - # objects. The definitions of these objects can be found in the file - # "torch/csrc/jit/ir/ir.h" in PyTorch's codebase: - # - # https://github.com/pytorch/pytorch . - # - # More specifically, look for the following definitions: - # - 'struct Value'; - # - 'struct TORCH_API Node'; - # - 'struct Graph'. - # - # These structures are exposed in Python via the 'pybind11' module. - - # build computational graph - opnodes_dict = {} - datanodes_dict = {} - arcs = [] - - datanode_2_onnx_id = dict() # data nodes will be discovered via tracing (I do not know in advance who they are) - datanode_id_gen = itertools.count() - - for i_op, opnode in enumerate(self.jit_graph.nodes()): - - # populate kernel partition of the computational graph - opnode_id = 'O' + graphs.__NODE_ID_FORMAT__.format(i_op) - opnodes_dict[opnode_id] = graphs.nodes.ONNXNode(opnode) - - # populate memory partition of the computational graph - # I might encouter the same data node ('torch._C.Value') again in - # other iterations of the loop on op nodes. I am trusting the fact - # that the object will have the same `debugName`; this seems - # reasonable, if nobody (or nothing) modifies the object - # in-between iterations. - for in_datanode in opnode.inputs(): - datanode_name = in_datanode.debugName() - try: # the data node has already been discovered - datanode_id = datanode_2_onnx_id[datanode_name] - except KeyError: - datanode_id = 'D' + graphs.__NODE_ID_FORMAT__.format(next(datanode_id_gen)) - datanode_2_onnx_id[datanode_name] = datanode_id - datanodes_dict[datanode_id] = graphs.nodes.ONNXNode(in_datanode) - arcs.append((datanode_id, opnode_id)) - - for out_datanode in opnode.outputs(): - datanode_name = out_datanode.debugName() - try: # the data node has already been discovered - datanode_id = datanode_2_onnx_id[datanode_name] - except KeyError: - datanode_id = 'D' + graphs.__NODE_ID_FORMAT__.format(next(datanode_id_gen)) - datanode_2_onnx_id[datanode_name] = datanode_id - datanodes_dict[datanode_id] = graphs.nodes.ONNXNode(out_datanode) # get_datanode_attributes(out_datanode) - arcs.append((opnode_id, datanode_id)) - - self.nx_graph = nx.DiGraph() - self.nx_graph.add_nodes_from(set(opnodes_dict.keys()), bipartite=graphs.Bipartite.KERNEL) - self.nx_graph.add_nodes_from(set(datanodes_dict.keys()), bipartite=graphs.Bipartite.MEMORY) - self.nx_graph.add_edges_from(set(arcs)) - - self.nodes_dict = {**opnodes_dict, **datanodes_dict} - - nx.set_node_attributes(self.nx_graph, {k: v.ntype for k, v in self.nodes_dict.items()}, 'type') - nx.set_node_attributes(self.nx_graph, {k: v.nscope for k, v in self.nodes_dict.items()}, 'scope') - - -class PyTorchGraph(object): - - def __init__(self, net, onnxgraph): - - # map ONNX graph to PyTorch graph - G = onnxgraph.nx_graph - assert '' not in {G.nodes[n]['scope'] for n in G.nodes}, "Argument graph {} has unscoped nodes.".format(G) - - g = nx.get_node_attributes(G, 'scope') - opnodes = {scope for n, scope in g.items() if G.nodes[n]['bipartite'] == graphs.Bipartite.KERNEL} - datanodes = {scope for n, scope in g.items() if G.nodes[n]['bipartite'] == graphs.Bipartite.MEMORY} - arcs = {arc for arc in map(lambda a: (g[a[0]], g[a[-1]]), G.edges)} - - self.nx_graph = nx.DiGraph() - self.nx_graph.add_nodes_from(opnodes, bipartite=graphs.Bipartite.KERNEL) - self.nx_graph.add_nodes_from(datanodes, bipartite=graphs.Bipartite.MEMORY) - self.nx_graph.add_edges_from(arcs) - - # remove data nodes which are used internaly to a PyTorch `nn.Module` (i.e., as "working memory"); - # beware: this operation is not reversible! - self.nx_graph.remove_nodes_from(PyTorchGraph.find_internal_datanodes(self.nx_graph)) - - # reassign IDs to nodes based on their topological sorting (I assume the graph is a DAG) - __NODE_ID_FORMAT__ = '{:06d}' - opnode_id_gen = itertools.count() - datanode_id_gen = itertools.count() - - onnx_scope_2_pytorch_id = {} - for n in dag.topological_sort(self.nx_graph): - if self.nx_graph.nodes[n]['bipartite'] == graphs.nodes.Bipartite.KERNEL: - node_id = 'O' + __NODE_ID_FORMAT__.format(next(opnode_id_gen)) - elif self.nx_graph.nodes[n]['bipartite'] == graphs.Bipartite.MEMORY: - node_id = 'D' + __NODE_ID_FORMAT__.format(next(datanode_id_gen)) - onnx_scope_2_pytorch_id[n] = node_id - - nx.relabel_nodes(self.nx_graph, onnx_scope_2_pytorch_id, copy=False) - self.pytorch_id_2_onnx_scope = {v: k for k, v in onnx_scope_2_pytorch_id.items()} - - # assign type and scope attributes to PyTorch graph nodes - opnodes_dict = {} - datanodes_dict = {} - - # populate kernel partition of the computational graph - for n in {n for n in self.nx_graph if self.nx_graph.nodes[n]['bipartite'] == graphs.Bipartite.KERNEL}: - onnx_scope = self.pytorch_id_2_onnx_scope[n] - if 'torch.view' in onnx_scope: # TODO: mind quantlib.editing.graphs/grrules/__init__.py:L16 - obj = graphs.ViewFlattenNd() - else: - obj = PyTorchGraph.get_pytorch_module_by_name(net, onnx_scope) - opnodes_dict[n] = graphs.PyTorchNode(obj) - - # populate memory partition of the computational graph - onnx_scope_2_onnx_id = {v: k for k, v in nx.get_node_attributes(G, 'scope').items()} - for n in {n for n in self.nx_graph.nodes if self.nx_graph.nodes[n]['bipartite'] == graphs.Bipartite.MEMORY}: - onnx_scope = self.pytorch_id_2_onnx_scope[n] - obj = onnxgraph.nodes_dict[onnx_scope_2_onnx_id[onnx_scope]].nobj - datanodes_dict[n] = graphs.PyTorchNode(obj) - - self.nodes_dict = {**opnodes_dict, **datanodes_dict} - nx.set_node_attributes(self.nx_graph, {k: v.ntype for k, v in self.nodes_dict.items()}, 'type') - nx.set_node_attributes(self.nx_graph, {k: v.nscope for k, v in self.nodes_dict.items()}, 'scope') - - # SCHEREMO: the traced graph treats all inputs to the compute graph as inputs and doesn't distinguish - # between user inputs and trainable parameters - we do this here by using the requires_grad attribute, - # which we tested to still be there after setting the model in eval. - # It's possible this breaks under some conditions, though. - python_id_2_datanode = {id(node.nobj): n for n, node in self.nodes_dict.items() if self.nx_graph.nodes[n]['bipartite'] == graphs.Bipartite.MEMORY} - inputs_python_ids = set(map(lambda n: id(n), onnxgraph.jit_graph.inputs())) - parameters_python_ids = set(filter(lambda id_: onnxgraph.nodes_dict[python_id_2_datanode[id_]].nobj.requires_grad(), inputs_python_ids)) - inputs_python_ids = inputs_python_ids.difference(parameters_python_ids) - outputs_python_ids = set(map(lambda n: id(n), onnxgraph.jit_graph.outputs())) - - # SCHEREMO: identify node attributes - each data node is either input (in the sense of the network input, not the compute graph inputs), output, parameter or other - # Most probably other is equivalent to intermediate feature map, but this requires further checking - data_partition_dict = {} - data_partition_dict.update({python_id_2_datanode[id_]: graphs.DataPartition.INPUT for id_ in inputs_python_ids}) - data_partition_dict.update({python_id_2_datanode[id_]: graphs.DataPartition.OUTPUT for id_ in outputs_python_ids}) - data_partition_dict.update({python_id_2_datanode[id_]: graphs.DataPartition.PARAMETER for id_ in parameters_python_ids}) - nx.set_node_attributes(self.nx_graph, data_partition_dict, 'data_partition') - - @staticmethod - def find_internal_datanodes(G): - - opnodes = {n for n in G.nodes if G.nodes[n]['bipartite'] == graphs.Bipartite.KERNEL} - datanodes = {n for n in G if G.nodes[n]['bipartite'] == graphs.Bipartite.MEMORY} - - internal_datanodes = [] - for datanode in datanodes: - - A = set(G.predecessors(datanode)) - B = set(G.successors(datanode)) - assert B.issubset(opnodes), "Data node {} has another data node as neighbour.".format(datanode) - - if A.issubset(B) and B.issubset(A) and len(A) == 1: - internal_datanodes.append(datanode) - - return internal_datanodes - - @staticmethod - def get_pytorch_module_by_name(module, target_name): - # this function is recursive: beware pitfalls related to inheritance! (https://stackoverflow.com/questions/13183501/staticmethod-and-recursion) - - for name, child in module.named_children(): - if name == target_name: - return child - elif name == target_name.split('.', 1)[0]: - return PyTorchGraph.get_pytorch_module_by_name(child, target_name.split('.', 1)[-1]) - - return module - diff --git a/editing/graphs/graphs/modules.py b/editing/graphs/graphs/modules.py deleted file mode 100644 index aa137f7..0000000 --- a/editing/graphs/graphs/modules.py +++ /dev/null @@ -1,113 +0,0 @@ -# -# modules.py -# -# Author(s): -# Matteo Spallanzani -# -# Copyright (c) 2020-2021 ETH Zurich. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -import torch -import torch.nn as nn - - -__all__ = [ - 'ViewFlattenNd', - 'ShiftAndClip', - 'FloorAndClip', - 'HelperInput', - 'HelperOutput', - 'HelperPrecisionTunnel', - 'HelperInputPrecisionTunnel', - 'HelperOutputPrecisionTunnel', -] - - -class ViewFlattenNd(nn.Module): - - def __init__(self): - super(ViewFlattenNd, self).__init__() - - def forward(self, x): - return x.view(x.size(0), -1) - - -class ShiftAndClip(nn.Module): - - def __init__(self, n_bits=8, shift=17, signed=True, only_positive=False): - super(ShiftAndClip, self).__init__() - self._n_bits = n_bits - self.shift = shift - self.min = -2**(self._n_bits - 1) if signed else 0 - self.min = 0 if only_positive else self.min - self.max = 2**(self._n_bits - 1) - 1 if signed else 2**self._n_bits - 1 - - def forward(self, x): - x //= 2**self.shift - x = torch.clamp(x, self.min, self.max) - return x - - -class FloorAndClip(nn.Module): - - def __init__(self, n_bits=8, signed=True, only_positive=False): - super(FloorAndClip, self).__init__() - self._n_bits = n_bits - self.min = -2**(self._n_bits - 1) if signed else 0 - self.min = 0 if only_positive else self.min - self.max = 2**(self._n_bits - 1) - 1 if signed else 2**self._n_bits - 1 - - def forward(self, x): - x = torch.floor(x) - x = torch.clamp(x, self.min, self.max) - return x - - -class HelperInput(nn.Identity): - - def __init__(self): - super(HelperInput, self).__init__() - - -class HelperOutput(nn.Identity): - - def __init__(self): - super(HelperOutput, self).__init__() - - -class HelperPrecisionTunnel(nn.Module): - - def __init__(self, eps_in, eps_out): - super(HelperPrecisionTunnel, self).__init__() - self.eps_in = eps_in - self.eps_out = eps_out - - def forward(self, x): - x = torch.div(x, self.eps_in) - x = torch.mul(x, self.eps_out) - return x - - -class HelperInputPrecisionTunnel(HelperPrecisionTunnel): - - def __init__(self, eps_in): - super(HelperInputPrecisionTunnel, self).__init__(eps_in, 1.0) - - -class HelperOutputPrecisionTunnel(HelperPrecisionTunnel): - - def __init__(self, eps_out): - super(HelperOutputPrecisionTunnel, self).__init__(1.0, eps_out) - diff --git a/editing/graphs/graphs/nodes.py b/editing/graphs/graphs/nodes.py deleted file mode 100644 index f0b6be3..0000000 --- a/editing/graphs/graphs/nodes.py +++ /dev/null @@ -1,99 +0,0 @@ -# -# nodes.py -# -# Author(s): -# Matteo Spallanzani -# -# Copyright (c) 2020-2021 ETH Zurich. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -import re -import torch -from enum import IntEnum, unique - - -__NODE_ID_FORMAT__ = '{:06d}' - - -@unique -class Bipartite(IntEnum): - KERNEL = 0 - MEMORY = 1 - CONTXT = 2 - - -@unique -class DataPartition(IntEnum): - INPUT = 0 - OUTPUT = 1 - PARAMETER = 2 - OTHER = 3 - - -class QuantLabNode(object): - - def __init__(self, obj): - self.nobj = obj - - -class ONNXNode(QuantLabNode): - - def __init__(self, obj): - super(ONNXNode, self).__init__(obj) - - @staticmethod - def onnx_scope_2_pytorch_scope(onnx_scope): - module_name_parts = re.findall('\[.*?\]', onnx_scope) - pytorch_scope = '.'.join([mn[1:-1] for mn in module_name_parts]) - return pytorch_scope - - @property - def ntype(self): - if isinstance(self.nobj, torch._C.Node): - ntype = self.nobj.kind() - elif isinstance(self.nobj, torch._C.Value): - ntype = '*' # data nodes are untyped ('onnx::Tensor'?) - return ntype - - @property - def nscope(self): - if isinstance(self.nobj, torch._C.Node): - nscope = ONNXNode.onnx_scope_2_pytorch_scope(self.nobj.scopeName()) - elif isinstance(self.nobj, torch._C.Value): - nscope = self.nobj.debugName() - return nscope - - -class PyTorchNode(QuantLabNode): - - def __init__(self, obj): - super(PyTorchNode, self).__init__(obj) - - @property - def ntype(self): - if isinstance(self.nobj, torch.nn.Module): - ntype = self.nobj.__class__.__name__ - elif isinstance(self.nobj, torch._C.Value): - ntype = '*' # data nodes are untyped ('onnx::Tensor'?) - return ntype - - @property - def nscope(self): - if isinstance(self.nobj, torch.nn.Module): - nscope = '' # the scope of `nn.Module`s usually depends on the "view" that the network's coder had of it at implementation time; we leave op nodes unscoped - elif isinstance(self.nobj, torch._C.Value): - nscope = self.nobj.debugName() - return nscope - diff --git a/editing/graphs/grrules/__init__.py b/editing/graphs/grrules/__init__.py deleted file mode 100644 index e2ea498..0000000 --- a/editing/graphs/grrules/__init__.py +++ /dev/null @@ -1,42 +0,0 @@ -# -# __init__.py -# -# Author(s): -# Matteo Spallanzani -# -# Copyright (c) 2020-2021 ETH Zurich. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -from collections import OrderedDict - -from .seeker import Seeker -from .dporules import * -from .hlprrules import * -from .. import traces - - -def load_rescoping_rules(modules=None): - - libtraces = traces.load_traces_library(modules=modules) - - librules = OrderedDict() - for mod_name, (L, K) in libtraces.items(): - if mod_name == 'ViewFlattenNd': - librules[mod_name] = ManualRescopingRule(L, K, 'torch.view') # TODO: mind quantlib.editing.graphs/graphs/graphs.py:L205 - else: - librules[mod_name] = AutoRescopingRule(L, K) - - return librules - diff --git a/editing/graphs/grrules/dporules.py b/editing/graphs/grrules/dporules.py deleted file mode 100644 index 33f377a..0000000 --- a/editing/graphs/grrules/dporules.py +++ /dev/null @@ -1,276 +0,0 @@ -# -# dporules.py -# -# Author(s): -# Matteo Spallanzani -# -# Copyright (c) 2020-2021 ETH Zurich. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -import networkx as nx -import itertools - -from .seeker import Seeker -from .. import graphs - - -__all__ = [ - 'ManualRescopingRule', - 'AutoRescopingRule', -] - - -class DPORule(object): - """This object represents an abstract *graph rewriting rule* (GRR). - - This process is similar to biological reactions between enzymes and other - proteins or ensembles of proteins (the so-called "substrates"). Indeed, - when an enzyme (the `Rule` object) encounters a substrate (an `nx.DiGraph` - object) it looks for suitable binding sites (the application points of the - rule) where it can catalyse a chemical reaction (an application of the - rule yielding a transformed `nx.DiGraph`). - """ - def __init__(self): - # 1. define the template graph L [L-term] - # 2. define the context (sub-)graph K [K-term] - # 3. define the template (sub-)graph L\K - # 4. define the replacement (sub-)graph R\K - # 5. define the connections between the context (sub-)graph K and the replacement (sub-)graph R\K to obtain the replacement graph R [R-term] - # 6. build the `Seeker` object; since the L-term of the GRR is usually defined or at least modified in the body of a GRR's constructor method, it is safe to build the seeker at the end - raise NotImplementedError - - def core(self, HI): - # transform the match (sub-)graph H\I into the substitute (sub-)graph J\I - raise NotImplementedError - - def apply(self, G, nodes_dict, g): - # 1. create copies of G and nodes_dict - # 2. characterise the match graph H: separate the interface (sub-)graph I from the match (sub-)graph H\I - # 3. transform the match (sub-)graph H\I into the substitute (sub-)graph J\I - # 4. glue the substitute (sub-)graph J\I to the main graph G via the interface (sub-)graph I - # 5. discard the match (sub-)graph H\I (including all the arcs between its nodes and the nodes of the interface (sub-)graph I) - # return G, nodes_dict # should be copies! - raise NotImplementedError - - def seek(self, G, nodes_dict): - # return gs (a `list` of `dictionaries` whose keys are nodes of G and values are nodes of L - raise NotImplementedError - - -class ManualRescopingRule(DPORule): - - def __init__(self, L, K, new_scope): - - # define the template graph L [L-term] - VLK = set(L.nodes).difference(set(K.nodes)) - self.L = nx.relabel_nodes(L, {vLK: '/'.join(['L-term', vLK]) for vLK in VLK}, copy=True) - - # define the context (sub-)graph K [K-term] - VK = set(self.L.nodes).intersection(set(K.nodes)) - nx.relabel_nodes(self.L, {vK: '/'.join(['K-term', vK]) for vK in VK}, copy=False) - VK = {vL for vL in set(self.L.nodes) if vL.startswith('K-term')} - self.K = self.L.subgraph(VK) - - # define the template (sub-)graph L\K - VLK = set(self.L.nodes).difference(set(self.K.nodes)) - self.LK = self.L.subgraph(VLK) - - # get the arcs that go from the vertices of K to those of L\K, and viceversa - E_K2LK2K = {arc for arc in set(self.L.edges).difference(set(self.LK.edges) | set(self.K.edges))} - E_K2LK = {arc for arc in E_K2LK2K if arc[0] in set(self.K.nodes)} - E_LK2K = {arc for arc in E_K2LK2K if arc[1] in set(self.K.nodes)} - - # define the replacement (sub-)graph R\K ["gluing" R\K to K yields the graph R, i.e., the R-term] -- rescoping rules implement one-to-one mappings - self.RK = nx.relabel_nodes(self.LK, {vLK: vLK.replace('L-term', 'R-term') for vLK in set(self.LK.nodes)}, copy=True) - - # define the arcs that go from the vertices of K to those of R\K, and viceversa - LK2RK_morphisms = Seeker(self.RK).get_morphisms(self.LK) - assert len(LK2RK_morphisms) == 1 - g_LK2RK = LK2RK_morphisms[0] - E_K2RK = {(u, g_LK2RK[v]) for u, v in E_K2LK} - E_RK2K = {(g_LK2RK[u], v) for u, v in E_LK2K} - E_K2RK2K = E_K2RK | E_RK2K - # disintegrate `E_K2RK` and `E_RK2K` along fibres to speed up rule application - self.F_K2RK = {vK: set(arc for arc in E_K2RK if arc[0] == vK) for vK in set(self.K.nodes)} - self.F_RK2K = {vK: set(arc for arc in E_RK2K if arc[1] == vK) for vK in set(self.K.nodes)} - - # # glue together the (sub-)graphs L\K and R\K along the vertices of K - # self.S = nx.compose(self.L, self.RK) - # self.S.add_edges_from(E_K2RK2K) - - # since the GRR's L-term will not be modified from now on, I can safely build the seeker - self.seeker = Seeker(self.L) - - # this machinery can generate always-new scope labels - self._scope = new_scope - self._counter = itertools.count() - - def _get_new_scope(self): - new_scope = '.'.join([self._scope, '{:03d}'.format(next(self._counter))]) - return new_scope - - def core(self, HI): - - # generate a new scope - new_scope = self._get_new_scope() - - # create a copy of the match (sub-)graph, but whose nodes have a new scope; its nodes are assigned different IDs to avoid conflicting IDs when gluing to G - JI = nx.relabel_nodes(HI, {vHI: vHI.replace('__tmp__', '') for vHI in set(HI.nodes)}, copy=True) - nx.set_node_attributes(JI, {vJI: new_scope for vJI in set(JI.nodes) if (JI.nodes[vJI]['bipartite'] == graphs.Bipartite.KERNEL)}, 'scope') - - return JI - - def apply(self, G, nodes_dict, g): - - # create new containers - G = G.copy() - nodes_dict = {**nodes_dict} - - # mark the to-be-rescoped nodes' IDs as obsolete - gkeys_2_tmpkeys = {vH: vH + '__tmp__' for vH, vL in g.items() if vL not in set(self.K.nodes)} - nx.relabel_nodes(G, gkeys_2_tmpkeys, copy=False) - # characterise the match graph H - VI = {vH for vH, vL in g.items() if vL in set(self.K.nodes)} - I = G.subgraph(VI) - VHI = {gkeys_2_tmpkeys[vH] for vH, vL in g.items() if vL not in set(self.K.nodes)} - HI = G.subgraph(VHI) - - # generate the substitute (sub-)graph J\I - JI = self.core(HI) - # add the substitute (sub-)graph J\I to the main graph G - G = nx.compose(G, JI) - - # compute the morphism 'g_{(J \setminus I) \to (R \setminus K)}': I need it to glue J\I to I - JI2RK_morphisms = Seeker(self.RK).get_morphisms(JI) - assert len(JI2RK_morphisms) == 1 - g_JI2RK = JI2RK_morphisms[0] - g_RK2JI = {vRK: vJI for vJI, vRK in g_JI2RK.items()} - # glue the substitue (sub-)graph J\I to the main graph G - for vI in set(I.nodes): - vK = g[vI] - G.add_edges_from({(vI, g_RK2JI[vRK]) for (_, vRK) in self.F_K2RK[vK]}) - G.add_edges_from({(g_RK2JI[vRK], vI) for (vRK, _) in self.F_RK2K[vK]}) - - # discard the match (sub-)graph H\I; arcs between H\I and I are deleted automatically - G.remove_nodes_from(set(HI.nodes)) - - return G, nodes_dict - - def seek(self, G, nodes_dict): - gs = self.seeker.get_morphisms(G) - return gs - - -class AutoRescopingRule(DPORule): - - def __init__(self, L, K): - - # define the graph L [L-term] - VL = set(L.nodes).difference(set(K.nodes)) - self.L = nx.relabel_nodes(L, {vL: '/'.join(['L-term', vL]) for vL in VL}, copy=True) - - # define the (sub-)graph K [K-term] - VK = set(self.L.nodes).intersection(set(K.nodes)) - nx.relabel_nodes(self.L, {vK: '/'.join(['K-term', vK]) for vK in VK}, copy=False) - VK = {vL for vL in set(self.L.nodes) if vL.startswith('K-term')} - self.K = self.L.subgraph(VK) - - # define the (sub-)graph L\K - VLK = set(self.L.nodes).difference(set(self.K.nodes)) - self.LK = self.L.subgraph(VLK) - - # get the arcs that go from the vertices of K to those of L\K, and viceversa - E_K2LK2K = {arc for arc in set(self.L.edges).difference(set(self.LK.edges) | set(self.K.edges))} - E_K2LK = {arc for arc in E_K2LK2K if arc[0] in set(self.K.nodes)} - E_LK2K = {arc for arc in E_K2LK2K if arc[1] in set(self.K.nodes)} - - # define the (sub-)graph R\K ["gluing" R\K to K yields the graph R, i.e., the R-term] -- rescoping rules implement one-to-one mappings - self.RK = nx.relabel_nodes(self.LK, {vLK: vLK.replace('L-term', 'R-term') for vLK in set(self.LK.nodes)}, copy=True) - - # define the arcs that go from the vertices of K to those of R\K, and viceversa - LK2RK_morphisms = Seeker(self.RK).get_morphisms(self.LK) - assert len(LK2RK_morphisms) == 1 - g_LK2RK = LK2RK_morphisms[0] - E_K2RK = {(u, g_LK2RK[v]) for u, v in E_K2LK} - E_RK2K = {(g_LK2RK[u], v) for u, v in E_LK2K} - E_K2RK2K = E_K2RK | E_RK2K - # disintegrate `E_K2RK` and `E_RK2K` along fibres to speed up rule application - self.F_K2RK = {vK: set(arc for arc in E_K2RK if arc[0] == vK) for vK in set(self.K.nodes)} - self.F_RK2K = {vK: set(arc for arc in E_RK2K if arc[1] == vK) for vK in set(self.K.nodes)} - - # # glue together the (sub-)graphs L\K and R\K along the vertices of K - # self.S = nx.compose(self.L, self.RK) - # self.S.add_edges_from(E_K2RK2K) - - # since the GRR's L-term will not be modified from now on, I can safely build the seeker - self.seeker = Seeker(self.L) - - def core(self, HI): - - # automatically detect the scope of the operations involved (should be unique!) - scopes = {HI.nodes[vHI]['scope'] for vHI in set(HI.nodes) if (HI.nodes[vHI]['bipartite'] == graphs.Bipartite.KERNEL)} - try: - scopes.remove('') - except KeyError: - pass - assert len(scopes) == 1 # up to now, quantlib.editing.s `nn.Module`s traces have included at least one correctly scoped operation... maybe we could suggest the user to apply a `ManualRescopingRule` when this does not happen? - new_scope = list(scopes)[0] - - # create a copy of the match (sub-)graph, but whose nodes have a new scope; its nodes are assigned different IDs to avoid conflicting IDs when gluing to G - JI = nx.relabel_nodes(HI, {vHI: vHI.replace('__tmp__', '') for vHI in set(HI.nodes)}, copy=True) - nx.set_node_attributes(JI, {vJI: new_scope for vJI in JI.nodes if (JI.nodes[vJI]['bipartite'] == graphs.Bipartite.KERNEL)}, 'scope') - - return JI - - def apply(self, G, nodes_dict, g): - - # create new containers - G = G.copy() - nodes_dict = {**nodes_dict} - - # mark the to-be-rescoped nodes' IDs as obsolete - gkeys_2_tmpkeys = {vH: vH + '__tmp__' for vH, vL in g.items() if vL not in set(self.K.nodes)} - nx.relabel_nodes(G, gkeys_2_tmpkeys, copy=False) - # characterise the match graph H - VI = {vH for vH, vL in g.items() if vL in set(self.K.nodes)} - I = G.subgraph(VI) - VHI = {gkeys_2_tmpkeys[vH] for vH, vL in g.items() if vL not in set(self.K.nodes)} - HI = G.subgraph(VHI) - - # generate the substitute (sub-)graph J\I - JI = self.core(HI) - # add the substitute (sub-)graph J\I to the main graph G - G = nx.compose(G, JI) - - # compute the morphism 'g_{(J \setminus I) \to (R \setminus K)}': I need it to glue J\I to I - JI2RK_morphisms = Seeker(self.RK).get_morphisms(JI) - assert len(JI2RK_morphisms) == 1 - g_JI2RK = JI2RK_morphisms[0] - g_RK2JI = {vRK: vJI for vJI, vRK in g_JI2RK.items()} - # glue the substitue (sub-)graph J\I to the main graph G - for vI in set(I.nodes): - vK = g[vI] - G.add_edges_from({(vI, g_RK2JI[vRK]) for (_, vRK) in self.F_K2RK[vK]}) - G.add_edges_from({(g_RK2JI[vRK], vI) for (vRK, _) in self.F_RK2K[vK]}) - - # discard the match (sub-)graph H\I; arcs between H\I and I are deleted automatically - G.remove_nodes_from(set(HI.nodes)) - - return G, nodes_dict - - def seek(self, G, nodes_dict): - gs = self.seeker.get_morphisms(G) - return gs - diff --git a/editing/graphs/grrules/hlprrules.py b/editing/graphs/grrules/hlprrules.py deleted file mode 100644 index b3fa646..0000000 --- a/editing/graphs/grrules/hlprrules.py +++ /dev/null @@ -1,417 +0,0 @@ -# -# hlprrules.py -# -# Author(s): -# Matteo Spallanzani -# -# Copyright (c) 2020-2021 ETH Zurich. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -import networkx as nx -from networkx.classes.filters import show_nodes, hide_edges -import itertools - -from .seeker import Seeker -#import graphs -from .. import graphs - -__all__ = [ - 'AddInputNodeRule', - 'AddOutputNodeRule', - 'RemoveInputNodeRule', - 'RemoveOutputNodeRule', - 'AddPrecisionTunnelRule', - 'RemovePrecisionTunnelRule', -] - - -class HelperRule(object): - """A GRR that prepares the graph for the 'core' editing. - - This GRR inserts nodes into the computational graph that are propedeutics - to the application of other GRRs; or, after all the 'core' GRRs have been - applied, it can remove specific subgraphs that could be replaced by an - identity operation or by groups of identity operations, and are therefore - redundant for computational purposes. - - This GRR still follows the algebraic approach to graph rewriting, but the - application points are computed 'on-the-fly'. In this sense, the rule - actually implements a vast (possibly infinite) set of GRRs. - """ - def __init__(self): - raise NotImplementedError - - def core(self): - raise NotImplementedError - - def apply(self, G, nodes_dict, g): - # return G, nodes_dict - raise NotImplementedError - - def seek(self, G, nodes_dict): - # return gs - raise NotImplementedError - - -###################### -## I/O HELPER NODES ## -###################### - -class AddIONodeRule(HelperRule): - - def __init__(self, io): - - self._io = io # either 'I' or 'O' - type = graphs.HelperInput.__name__ if self._io == 'I' else graphs.HelperOutput.__name__ - - self.RK = nx.DiGraph() - self.RK.add_nodes_from(set([''.join(['R-term/', 'H', self._io])]), bipartite=graphs.Bipartite.KERNEL, type=type) - - self._counter = itertools.count() - - def core(self): - - vJI = ''.join(['H', self._io, graphs.__NODE_ID_FORMAT__.format(next(self._counter))]) - JI = nx.relabel_nodes(self.RK, {vRK: vJI for vRK in set(self.RK.nodes)}, copy=True) - - m = graphs.HelperInput() if self._io == 'I' else graphs.HelperOutput() - vJI_2_ptnode = {vJI: graphs.PyTorchNode(m)} - - nx.set_node_attributes(JI, {vJI: '' for vJI in set(JI.nodes) if (JI.nodes[vJI]['bipartite'] == graphs.Bipartite.KERNEL)}, 'scope') - - return JI, vJI_2_ptnode - - def apply(self, G, nodes_dict, g): - - G = G.copy() - - VI = set(g.keys()) - I = G.subgraph(VI) - - JI, vJI_2_ptnode = self.core() - - G = nx.compose(G, JI) - if self._io == 'I': - E_JI2I = set(itertools.product(set(JI.nodes), set(I.nodes))) - G.add_edges_from(E_JI2I) - elif self._io == 'O': - E_I2JI = set(itertools.product(set(I.nodes), set(JI.nodes))) - G.add_edges_from(E_I2JI) - - nodes_dict = {**nodes_dict, **vJI_2_ptnode} - - return G, nodes_dict - - def seek(self, G, nodes_dict, VIs): - - if self._io == 'I': - VIs = list(filter(lambda VI: len(set(itertools.chain.from_iterable([set(G.predecessors(vI)) for vI in VI]))) == 0, VIs)) - if self._io == 'O': - assert len(set(itertools.chain.from_iterable(VIs))) == sum([len(VI) for VI in VIs]) # I assume that an operation can't write to multiple output nodes - VIs = list(filter(lambda VI: len(set(itertools.chain.from_iterable([set(G.successors(vI)) for vI in VI]))) == 0, VIs)) - - gs = [] - for VI in VIs: - g = {vI: None for vI in VI} # there is no fixed context term (K-term) for this rule! - gs.append(g) - - return gs - - -class AddInputNodeRule(AddIONodeRule): - - def __init__(self): - super(AddInputNodeRule, self).__init__('I') - - -class AddOutputNodeRule(AddIONodeRule): - - def __init__(self): - super(AddOutputNodeRule, self).__init__('O') - - -class RemoveIONodeRule(HelperRule): - - def __init__(self, io): - - self._io = io # either 'I' or 'O' - type = graphs.HelperInput.__name__ if self._io == 'I' else graphs.HelperOutput.__name__ - - # the I/O operation will serve as an "anchor"; from it, I will be able to (implicitly) generate and apply the graph rewriting rule on-the-fly - self.LK = nx.DiGraph() - self.LK.add_nodes_from(set([''.join(['L-term/', 'H', self._io])]), bipartite=graphs.Bipartite.KERNEL, type=type) - - self.seeker = Seeker(self.LK) - - def core(self): - pass - - def apply(self, G, nodes_dict, g): - - # create new containers - G = G.copy() - nodes_dict = {**nodes_dict} - - # characterise the match (sub-)graph H\I - VHI = set(g.keys()) - HI = G.subgraph(VHI) - - # delete the I/O nodes - G.remove_nodes_from(set(HI.nodes)) - for n in set(HI.nodes): - del nodes_dict[n] - - return G, nodes_dict - - def seek(self, G, nodes_dict): - - def is_valid_application_point(g): - - VJI = set(g.keys()) - assert len(VJI) == 1 - - if self._io == 'I': - is_ok = len(set(G.predecessors(next(iter(VJI))))) == 0 # an input node does not read the output of any other node - elif self._io == 'O': - is_ok = len(set(G.successors(next(iter(VJI))))) == 0 # an output node's output won't be read by any node - - return is_ok - - gs = self.seeker.get_morphisms(G) - gs = list(filter(is_valid_application_point, gs)) - - return gs - - -class RemoveInputNodeRule(RemoveIONodeRule): - - def __init__(self): - super(RemoveInputNodeRule, self).__init__('I') - - -class RemoveOutputNodeRule(RemoveIONodeRule): - - def __init__(self): - super(RemoveOutputNodeRule, self).__init__('O') - - -############################ -## PRECISION HELPER NODES ## -############################ - -class AddPrecisionTunnelRule(HelperRule): - - def __init__(self, type): - """Insert a 'precision tunnel' after idempotent operations. - - Imagine having a function composition :math:`o \circ f \circ i`, - which needs to be transformed into a second composition - :math:`h \circ g`. Suppose that the information required to compute - :math:`h` is contained in :math:`o` and (partly) in :math:`f`, i.e., - there exists a transformation :math:`T_{h} \,|\, h = T_{h}(o, f)`; - similarly, the information required to compute :math:`g` is contained - (partly) in :math:`f` and in :math:`i`, i.e., there exists a transform - :math:`T_{g} \,|\, g = T_{g}(f, i)`. We assume that :math:`T_{h}` and - :math:`T_{g}` can be applied in any order, but must be executed - sequentially; also, we assume that after the application of a - transform its inputs will be destroyed. We see then that there is no - valid order, since each of them will destroy :math:`f`, preventing the - application of the second transformation. - - If we suppose that :math:`f` is idempotent (i.e., it is such that - :math:`f \circ f = f`), we can rewrite the original term as - :math:`o \circ f \circ f \circ i` before applying the transformation - rules :math:`T_{h}` and :math:`T_{g}`. In this case, we can derive the - desired form :math:`T_{h}(o, f) \circ T_{g}(f, i)` without any issue. - - In particular, we focus on the case where :math:`f` is a quantization - operator, i.e., an activation function of the form - :math:`f = e \circ r_{p} \circ e^{-1}`, where :math:`r_{p}` is a - rounding operation at precision :math:`p`, and :math:`e` is an - element-wise multiplication by a positive number (possibly represented - in floating point). We assume that the 'anchor' of each application - point (i.e., the quantization operator :math:`f`) will receive data - from just one operation :math:`i`, but its outputs will be read by a - positive number of operations :math:`o^{(k)}, k = 0, \dots, K-1`, for - some positive integer :math:`K`. Then, we can rewrite each sequence - :math:`o^{(k)} \circ f \circ i` as - :math:`o^{(k)} \circ e \circ r_{p} \circ e^{-1} \circ e \circ r_{p} \circ e^{-1} \circ i`. - In this way, we will be able to apply :math:`T_{h}` :math:`K` times, - returning :math:`h^{(k)} = T_{h}(o^{(k)}, f)`, and :math:`T_{g}` just - once, returning :math:`g = T_{g}(f, i)`. - """ - - # the idempotent operation will serve as an "anchor"; from it, I will be able to (implicitly) generate and apply the graph rewriting rule on-the-fly - self.Kin = nx.DiGraph() - self.Kin.add_nodes_from(set([''.join(['K-term/', 'HPTin'])]), bipartite=graphs.Bipartite.KERNEL, type=type) - - self.seeker = Seeker(self.Kin) - - def core(self, H, Iin, Iout, eps, nodes_dict): - - # `HelperPrecisionTunnel` nodes are meant to serve as 'stubs' when applying full-fledged GRRs; - # this graph will ensure the correct connectivity between the 'pieces' of the full graph that will be derived by applying full-fledged GRRs - vH_2_vJI_PTin = {vH: vH.replace('O', 'HPTin') for vH in set(H.nodes).intersection(set(Iin.nodes))} - vH_2_vJI_PTout = {vH: vH.replace('O', 'HPTout') for vH in set(H.nodes).intersection(set(Iout.nodes))} - JI = nx.relabel_nodes(H, {**vH_2_vJI_PTin, **vH_2_vJI_PTout}, copy=True) - nx.set_node_attributes(JI, {vJI: graphs.HelperInputPrecisionTunnel.__name__ for vJI in set(vH_2_vJI_PTin.values())}, 'type') - nx.set_node_attributes(JI, {vJI: graphs.HelperOutputPrecisionTunnel.__name__ for vJI in set(vH_2_vJI_PTout.values())}, 'type') - - # replicate the "anchor" idempotent operation along each connection - vJI_PTout_2_vJI_PTclone = {vJI: vJI.replace('HPTout', 'HPTclone') for vJI in set(vH_2_vJI_PTout.values())} - for u, v in vJI_PTout_2_vJI_PTclone.items(): - Iin_clone = nx.relabel_nodes(H.subgraph(Iin), {next(iter(set(Iin.nodes))): v}, copy=True) - JI = nx.compose(JI, Iin_clone) - JI.add_edge(u, v) - - nx.set_node_attributes(JI, {vJI: '' for vJI in set(JI.nodes) if (JI.nodes[vJI]['bipartite'] == graphs.Bipartite.KERNEL)}, 'scope') - - # compute the connections of the new nodes to the old nodes - E_I2JI = {(vI, vJI) for vI, vJI in vH_2_vJI_PTin.items()} - vJI_PTclone_2_vJI_PTout = {v: k for k, v in vJI_PTout_2_vJI_PTclone.items()} - vJI_PTout_2_vH = {v: k for k, v in vH_2_vJI_PTout.items()} - E_JI2I = {(vJI, vJI_PTout_2_vH[vJI_PTclone_2_vJI_PTout[vJI]]) for vJI in set(vJI_PTclone_2_vJI_PTout.keys())} - - # register the technical specs of the new ops - vJI_2_ptnode = {} - for vJI in set(JI.nodes): - if JI.nodes[vJI]['type'] == graphs.HelperInputPrecisionTunnel.__name__: - ptnode = graphs.PyTorchNode(graphs.HelperInputPrecisionTunnel(eps)) - elif JI.nodes[vJI]['type'] == graphs.HelperOutputPrecisionTunnel.__name__: - ptnode = graphs.PyTorchNode(graphs.HelperOutputPrecisionTunnel(eps)) - else: - ptnode = nodes_dict[next(iter(set(Iin.nodes)))] # since the idempotent operation already exists, I just need a pointer to it - vJI_2_ptnode[vJI] = ptnode - - return JI, vJI_2_ptnode, E_I2JI, E_JI2I - - def apply(self, G, nodes_dict, g): - - # create new containers - G = G.copy() - nodes_dict = {**nodes_dict} - - # compute the match graph H on-the-fly - VIin = set(g.keys()) - VIout = set(G.successors(next(iter(VIin)))) - VI = VIin | VIout - VH = VI - H = G.subgraph(VH) - I = nx.subgraph_view(G, filter_node=show_nodes(VI), filter_edge=hide_edges(set(H.edges))) - Iin = I.subgraph(VIin) - Iout = I.subgraph(VIout) - - # create the precision tunnel - eps = nodes_dict[next(iter(VIin))].nobj.eps - # n = nodes_dict[next(iter(VIin))].nobj.num_levels - # m = nodes_dict[next(iter(VIin))].nobj.abs_max_value.item() # TODO: only `STEActivation` nodes have `abs_max_value` attribute! try to homogenise this in the future - # eps = (2 * m) / (n - 1) - JI, vJI_2_ptnode, E_I2JI, E_JI2I = self.core(H, Iin, Iout, eps, nodes_dict) - - # link the substitute (sub-)graph J\I to the interface (sub-)graph I - G = nx.compose(G, JI) - E_I2JI2I = E_I2JI | E_JI2I - G.add_edges_from(E_I2JI2I) - nodes_dict.update(vJI_2_ptnode) - - # delete H \ I (this it is NOT the match (sub-)graph H\I, but the difference between the match graph H and the interface (sub-)graph I) - G.remove_edges_from(set(H.edges)) - - return G, nodes_dict - - def seek(self, G, nodes_dict): - - def is_valid_application_point(g): - - VIin = set(g.keys()) - assert len(VIin) == 1 - VIout = set(G.successors(next(iter(VIin)))) - - is_ok = len(VIout) > 0 # adding a precision tunnel makes sense just in case the node has at least one output - return is_ok - - gs = self.seeker.get_morphisms(G) - gs = list(filter(is_valid_application_point, gs)) - - return gs - - -class RemovePrecisionTunnelRule(HelperRule): - - def __init__(self): - """Delete the `HelperPrecisionTunnel` nodes in the graph. - - This GRR is not mean to act as a full inverse of the - `AddPrecisionTunnelRule` GRR. It will only remove nodes whose - corresponding `nn.Module`s are of type `HelperPrecisionTunnel`; it - will not take care of 'reabsorbing' the copies of the idempotent - operations generated by applications of the `AddPrecisionTunnelRule`. - In fact, this GRR assumes that all those copies will be consumed by - full-fledged GRRs. Since `HelperPrecisionTunnel` modules are meant to - serve as 'stubs' when applying full-fledged GRRs, this rule should be - applied only after all such copies will have been absorbed. In - summary, this GRR is meant to 'clean up' the computational graph once - all the 'core' GRRs will have been applied. - """ - - # the input to the precision tunnel will serve as an "anchor"; once I locate such a node, I will be able to (implicitly) generate and apply the graph rewriting rule on-the-fly - self.LK = nx.DiGraph() - self.LK.add_nodes_from(set([''.join(['L-term/', 'HPTin'])]), bipartite=graphs.Bipartite.KERNEL, type=graphs.HelperInputPrecisionTunnel(1.0).__class__.__name__) - - self.seeker = Seeker(self.LK) - - def apply(self, G, nodes_dict, g): - - # create new containers - G = G.copy() - nodes_dict = {**nodes_dict} - - # characterise the match graph H - VHIin = set(g.keys()) - VHIout = set(G.successors(next(iter(VHIin)))) - VHI = VHIin | VHIout - VIin = set(G.predecessors(next(iter(VHIin)))) - VIout = set(itertools.chain.from_iterable([set(G.successors(vH)) for vH in VHIout])) - - # add J \ I; (this it is NOT the substitute (sub-)graph J\I, but the difference between the substitute graph J and the interface (sub-)graph I) - G.add_edges_from(list(itertools.product(VIin, VIout))) - - # delete the 'precision tunnel' nodes - G.remove_nodes_from(VHI) - for n in VHI: - del nodes_dict[n] - - return G, nodes_dict - - def seek(self, G, nodes_dict): - - def is_valid_application_point(g): - - VHin = set(g.keys()) - assert len(VHin) == 1 - VHout = set(G.successors(next(iter(VHin)))) - assert all([G.nodes[vH]['type'] == graphs.HelperOutputPrecisionTunnel.__name__ for vH in VHout]) - - epss_in = {nodes_dict[next(iter(VHin))].nobj.eps_in} - epss_out = {nodes_dict[vH].nobj.eps_out for vH in VHout} - is_ok = len(epss_in | epss_out) == 1 # all the output quanta must agree with the input quantum - - return is_ok - - gs = self.seeker.get_morphisms(G) - gs = list(filter(is_valid_application_point, gs)) - - return gs - diff --git a/editing/graphs/grrules/seeker.py b/editing/graphs/grrules/seeker.py deleted file mode 100644 index 6cd1387..0000000 --- a/editing/graphs/grrules/seeker.py +++ /dev/null @@ -1,84 +0,0 @@ -# -# seeker.py -# -# Author(s): -# Matteo Spallanzani -# -# Copyright (c) 2020-2021 ETH Zurich. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -from networkx.algorithms import isomorphism - - -class Seeker(object): - """This object looks for the application points of a graph rewriting rule. - """ - def __init__(self, T): - self.T = T - - @staticmethod - def is_morphism(T, G, g): - - is_ok = False - - for vH, vT in g.items(): - - is_same_partition = G.nodes[vH]['bipartite'] == T.nodes[vT]['bipartite'] - is_same_type = G.nodes[vH]['type'] == T.nodes[vT]['type'] - is_ok = is_same_partition and is_same_type # computational graphs are node-labelled graphs, where node types act as labels - - if not is_ok: - break - - return is_ok - - def get_morphisms(self, G): - - # In principle, morphisms do not need to be isomorphisms: this is a - # restriction that I chose to simplify the work on QNNs conversion (it - # makes solving ambiguities much easier). - # In particular, candidate matchings will be induced subgraph - # isomorphisms, not "spurious" monomorphisms: - # - # https://github.com/networkx/networkx/blob/master/networkx/algorithms/isomorphism/isomorphvf2.py . - # - - # G is the whole network graph, T is the template to be matched - matcher = isomorphism.DiGraphMatcher(G, self.T) - - # Return list of all occurences of the non-attributed template - # This means only the number of nodes + interconnections are matched - isomorphisms = list(matcher.subgraph_isomorphisms_iter()) - - # check the second morphism condition (label consistency) - # This checks that node types / classes (technically, string labels representing the class name) are as in the template - morphisms = [g for g in isomorphisms if Seeker.is_morphism(self.T, G, g)] - - # remove duplicate morphisms - # SCHEREMO: Check this a bit carefully - # SPMATTEO: Check what is filtered, pipe to logs - unique_VHs = set(frozenset(g.keys()) for g in morphisms) - VHs_2_morphisms = {VH: [g for g in morphisms if frozenset(g.keys()) == VH] for VH in unique_VHs} - morphisms = [v[0] for v in VHs_2_morphisms.values()] - - # List of dictionaries, mapping node identifiers to node identifiers - # Node identifier as in the networkx package - - # We map the node identifiers that are absolute/unique to the original graph G - # to node identifiers that are a bsolute/unique to the template graph self.T - # This assures there is no ambiguity in the mapping (<- This is the assumption here) - - return morphisms - diff --git a/editing/graphs/traces/__init__.py b/editing/graphs/traces/__init__.py deleted file mode 100644 index 2742e61..0000000 --- a/editing/graphs/traces/__init__.py +++ /dev/null @@ -1,50 +0,0 @@ -# -# __init__.py -# -# Author(s): -# Matteo Spallanzani -# -# Copyright (c) 2020-2021 ETH Zurich. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -import os -from collections import OrderedDict -import networkx as nx - -from .trace import __TRACES_LIBRARY__ -from .. import graphs - -def load_traces_library(modules=None): - - mod_2_trace_dir = {} - for root, dirs, files in os.walk(__TRACES_LIBRARY__): - if len(dirs) == 0: # terminal directories contain only trace files (graphviz, networkx) - mod_2_trace_dir[os.path.basename(root)] = root - - if modules is None: - modules = list(mod_2_trace_dir.keys()) # beware: there is no guarantee on the order in which the rescoping rules will be returned! - - libtraces = OrderedDict() - for mod_name in modules: - - L = nx.read_gpickle(os.path.join(mod_2_trace_dir[mod_name], 'networkx')) - VK = {n for n in L.nodes if L.nodes[n]['partition'] == graphs.Bipartite.CONTXT} - for n in L.nodes: - del L.nodes[n]['partition'] - K = L.subgraph(VK) - - libtraces[mod_name] = (L, K) - - return libtraces diff --git a/editing/graphs/traces/trace.py b/editing/graphs/traces/trace.py deleted file mode 100644 index a464acf..0000000 --- a/editing/graphs/traces/trace.py +++ /dev/null @@ -1,182 +0,0 @@ -# -# trace.py -# -# Author(s): -# Matteo Spallanzani -# -# Copyright (c) 2020-2021 ETH Zurich. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -import networkx as nx -import os - -# import graphs -# import quantlib.editing.graphs.utils - -from quantlib.editing.graphs import graphs -from quantlib.editing.graphs import utils - -__TRACES_LIBRARY__ = os.path.join(os.path.dirname(os.path.abspath(__file__)), '.libtraces') -# __TRACES_LIBRARY__ = os.path.join(os.path.expanduser('~'), 'Desktop', 'QuantLab', 'quantlib', 'editing', 'graphs', 'traces', '.libtraces') - -def trace_module(library, algorithm, mod, dummy_input): - - # trace graph - mod.eval() - onnxgraph = graphs.ONNXGraph(mod, dummy_input) - G = onnxgraph.nx_graph - - # locate interface nodes - node_2_partition = nx.get_node_attributes(G, 'bipartite') - for n in {n for n, onnxnode in onnxgraph.nodes_dict.items() if onnxnode.nobj in set(onnxgraph.jit_graph.inputs()) | set(onnxgraph.jit_graph.outputs())}: - node_2_partition[n] = graphs.Bipartite.CONTXT - nx.set_node_attributes(G, node_2_partition, 'partition') - - # store traces and graph picture - trace_dir = os.path.join(__TRACES_LIBRARY__, library, algorithm, mod.__class__.__name__) - if not os.path.isdir(trace_dir): - os.makedirs(trace_dir, exist_ok=True) - nx.write_gpickle(G, os.path.join(trace_dir, 'networkx')) - #utils.draw_graph(G, trace_dir, 'graphviz') - - -#################################### -## GENERIC PARAMETERS FOR TRACING ## -#################################### - -_batch_size = 1 -_n_input_channels = 8 -_n_output_channels = 8 -_dim1 = 32 -_dim2 = 32 -_dim3 = 32 -_kernel_size = 3 -_stride = 1 -_padding = 1 - - -############################# -## PYTORCH MODULES TRACING ## -############################# - -def trace_pytorch_modules(): - - import torch - import torch.nn as nn - - library = 'PyTorch' - - ##################### - ## AdaptiveAvgPool ## - ##################### - algorithm = 'AdaptiveAvgPool' - - mod_AdaptiveAvgPool1d = nn.AdaptiveAvgPool1d((int(_dim1 / 4))) - dummy_input = torch.ones(_batch_size, _n_input_channels, _dim1) - trace_module(library, algorithm, mod_AdaptiveAvgPool1d, dummy_input) - - mod_AdaptiveAvgPool2d = nn.AdaptiveAvgPool2d((int(_dim1 / 4), int(_dim2 / 4))) - dummy_input = torch.ones(_batch_size, _n_input_channels, _dim1, _dim2) - trace_module(library, algorithm, mod_AdaptiveAvgPool2d, dummy_input) - - mod_AdaptiveAvgPool3d = nn.AdaptiveAvgPool3d((int(_dim1 / 4), int(_dim2 / 4), int(_dim3 / 4))) - dummy_input = torch.ones(_batch_size, _n_input_channels, _dim1, _dim2, _dim3) - trace_module(library, algorithm, mod_AdaptiveAvgPool3d, dummy_input) - - ################ - ## torch.view ## - ################ - algorithm = 'ViewFlatten' - - mod_ViewFlattenNd = graphs.ViewFlattenNd() - dummy_input = torch.ones(_batch_size, _n_input_channels, _dim1, _dim2, _dim3) - trace_module(library, algorithm, mod_ViewFlattenNd, dummy_input) - - -############################## -## QUANTLIB MODULES TRACING ## -############################## - -def trace_quantlib_modules(): - - import torch - import quantlib.algorithms as qa - - library = 'QuantLab' - - ######### - ## STE ## - ######### - algorithm = 'STE' - - # TODO: upgrading to PyTorch 1.9.0 breaks tracing of STE activation for "some reason": we must look into this! - # mod_STEActivation = qa.ste.STEActivation() - # dummy_input = torch.ones((_batch_size, _n_input_channels)) - # trace_module(library, algorithm, mod_STEActivation, dummy_input) - - ######### - ## INQ ## - ######### - algorithm = 'INQ' - - mod_INQConv1d = qa.inq.INQConv1d(_n_input_channels, _n_output_channels, kernel_size=_kernel_size, stride=_stride, padding=_padding, bias=False) - dummy_inpyut = torch.ones((_batch_size, _n_input_channels, _dim1)) - trace_module(library, algorithm, mod_INQConv1d, dummy_inpyut) - - mod_INQConv2d = qa.inq.INQConv2d(_n_input_channels, _n_output_channels, kernel_size=_kernel_size, stride=_stride, padding=_padding, bias=False) - dummy_input = torch.ones((_batch_size, _n_input_channels, _dim1, _dim2)) - trace_module(library, algorithm, mod_INQConv2d, dummy_input) - - # mod_INQConv3d = qa.inq.INQConv3d(_n_input_channels, _n_output_channels, kernel_size=_kernel_size, stride=_stride, padding=_padding, bias=False) - # dummy_input = torch.ones((_batch_size, _n_input_channels, _dim1, _dim2, _dim3)) - # trace_module(library, algorithm, mod_INQConv3d, dummy_input) - - ######### - ## ANA ## - ######### - algorithm = 'ANA' - - quantizer_spec = {'nbits': 2, 'signed': True, 'balanced': True, 'eps': 1.0} - noise_type = 'uniform' - strategy = 'expectation' - - mod_ANAActivation = qa.ana.ANAActivation(quantizer_spec, noise_type, strategy) - dummy_input = torch.ones((_batch_size, _n_input_channels)) - trace_module(library, algorithm, mod_ANAActivation, dummy_input) - - mod_ANALinear = qa.ana.ANALinear(quantizer_spec, noise_type, strategy, _n_input_channels, _n_output_channels, bias=False) - dummy_input = torch.ones((_batch_size, _n_input_channels)) - trace_module(library, algorithm, mod_ANALinear, dummy_input) - - mod_ANAConv1d = qa.ana.ANAConv1d(quantizer_spec, noise_type, strategy, _n_input_channels, _n_output_channels, kernel_size=_kernel_size, stride=_stride, padding=_padding, bias=False) - dummy_input = torch.ones((_batch_size, _n_input_channels, _dim1)) - trace_module(library, algorithm, mod_ANAConv1d, dummy_input) - - mod_ANAConv2d = qa.ana.ANAConv2d(quantizer_spec, noise_type, strategy, _n_input_channels, _n_output_channels, kernel_size=_kernel_size, stride=_stride, padding=_padding, bias=False) - dummy_input = torch.ones((_batch_size, _n_input_channels, _dim1, _dim2)) - trace_module(library, algorithm, mod_ANAConv2d, dummy_input) - - mod_ANAConv3d = qa.ana.ANAConv3d(quantizer_spec, noise_type, strategy, _n_input_channels, _n_output_channels, kernel_size=_kernel_size, stride=_stride, padding=_padding, bias=False) - dummy_input = torch.ones((_batch_size, _n_input_channels, _dim1, _dim2, _dim3)) - trace_module(library, algorithm, mod_ANAConv3d, dummy_input) - - -if __name__ == '__main__': - - if not os.path.isdir(__TRACES_LIBRARY__): - os.mkdir(__TRACES_LIBRARY__) - - trace_pytorch_modules() - trace_quantlib_modules() diff --git a/editing/graphs/utils/__init__.py b/editing/graphs/utils/__init__.py deleted file mode 100644 index 7c0768d..0000000 --- a/editing/graphs/utils/__init__.py +++ /dev/null @@ -1,23 +0,0 @@ -# -# __init__.py -# -# Author(s): -# Matteo Spallanzani -# -# Copyright (c) 2020-2021 ETH Zurich. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -from .draw import draw_graph - diff --git a/editing/graphs/utils/draw.py b/editing/graphs/utils/draw.py deleted file mode 100644 index 682a02a..0000000 --- a/editing/graphs/utils/draw.py +++ /dev/null @@ -1,61 +0,0 @@ -# -# draw.py -# -# Author(s): -# Matteo Spallanzani -# -# Copyright (c) 2020-2021 ETH Zurich. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -from collections import namedtuple -import networkx as nx -import graphviz as gv - -from ...graphs import graphs - - -GVNodeAppearance = namedtuple('GVNodeAppearance', ['fontsize', 'shape', 'height', 'width', 'color', 'fillcolor']) - - -def _get_styles(): - - _styles = { - graphs.Bipartite.KERNEL.name: - GVNodeAppearance(fontsize='8', shape='circle', height='2.0', width='2.0', - color='cornflowerblue', fillcolor='cornflowerblue'), - graphs.Bipartite.MEMORY.name: - GVNodeAppearance(fontsize='8', shape='square', height='1.2', width='1.2', - color='brown2', fillcolor='brown2') - } - - return _styles - - -def draw_graph(G, save_dir, filename, node_2_label=None): - - # map nodes to labels and graphic styles - partition_2_style = _get_styles() - if (node_2_label is None) or (set(G.nodes) != set(node_2_label.keys())): - node_2_label = {n: G.nodes[n]['type'] for n in G.nodes} - - # build GraphViz graph - gvG = gv.Digraph(comment=filename) - for n, p in nx.get_node_attributes(G, 'bipartite').items(): - gvG.node(n, node_2_label[n], **partition_2_style[p.name]._asdict(), style='filled') - for e in G.edges: - gvG.edge(e[0], e[1]) - - gvG.render(directory=save_dir, filename=filename) - From f42a99e8b2cb9bfec2e4ae5eb1bfffda80b13f83 Mon Sep 17 00:00:00 2001 From: Luka Macan Date: Fri, 1 Dec 2023 16:10:47 +0100 Subject: [PATCH 2/4] Fix quant being copied even if it doesn't exist --- editing/fx/passes/pact/harmonize.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/editing/fx/passes/pact/harmonize.py b/editing/fx/passes/pact/harmonize.py index 9722051..134c4ad 100644 --- a/editing/fx/passes/pact/harmonize.py +++ b/editing/fx/passes/pact/harmonize.py @@ -170,7 +170,8 @@ def run_pass(self, gm : fx.GraphModule): #now, we assume the output of the new node has the same "quant" #meta properties as the output of the tree that is being #replaced - new_node.meta['quant'] = deepcopy(tree.end_node.meta['quant']) + if hasattr(tree.end_node, 'meta') and 'quant' in tree.end_node.meta.keys(): + new_node.meta['quant'] = deepcopy(tree.end_node.meta['quant']) # finally, delete the nodes in the tree for node in tree.nodes: gm.graph.erase_node(node) From 6a47ffbebaba4d4628c4a9ed0d82a2a388dc63f4 Mon Sep 17 00:00:00 2001 From: Luka Macan Date: Fri, 1 Dec 2023 16:33:43 +0100 Subject: [PATCH 3/4] Revert pytorch requant fix because it breaks dory --- algorithms/pact/pact_ops.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/algorithms/pact/pact_ops.py b/algorithms/pact/pact_ops.py index 9652634..3482909 100644 --- a/algorithms/pact/pact_ops.py +++ b/algorithms/pact/pact_ops.py @@ -112,8 +112,8 @@ def forward(ctx, x, mul, add, div, signed, n_levels_out, cmsis_requant): # division. Division is with flooring. else: y = x * mul + add - # Avoid round to even behaviour, friggin pytorch - y = torch.floor((y / div) + 0.5) + # LMACAN: Dory doesn't like the `+ 0.5` fix + y = torch.floor(y / div) if not signed: # if unsigned: clip y to interval (0, n_levels-1) From 9975c7d2e1b7dc5078e8fb1b4f9a868a5ef078ec Mon Sep 17 00:00:00 2001 From: Luka Macan Date: Tue, 9 Jan 2024 17:36:08 +0100 Subject: [PATCH 4/4] Fix Dory's addition operation clips the output --- backends/dory/dory_passes.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/backends/dory/dory_passes.py b/backends/dory/dory_passes.py index 68e3915..c43924d 100644 --- a/backends/dory/dory_passes.py +++ b/backends/dory/dory_passes.py @@ -116,7 +116,7 @@ def __init__(self): class DORYAdder(nn.Module): class DORYAdderFun(torch.autograd.Function): @staticmethod - def forward(ctx, x1, rq1, x2, rq2, rq_out): + def forward(ctx, x1, rq1, x2, rq2, rq_out, out_n_levels): if rq1: x1 = rq1(x1) if rq2: @@ -127,10 +127,20 @@ def forward(ctx, x1, rq1, x2, rq2, rq_out): if rq_out: x_sum = rq_out(x_sum) + out_signed = rq_out.signed if rq_out else False + if out_signed: + out_min = -(out_n_levels // 2) + out_max = (out_n_levels - 1) // 2 + else: + out_min = 0 + out_max = out_n_levels - 1 + + x_sum = torch.clamp(x_sum, out_min, out_max) + return x_sum @staticmethod - def symbolic(g, x1, rq1, x2, rq2, rq_out): + def symbolic(g, x1, rq1, x2, rq2, rq_out, out_n_levels): params = {} out_signed_inferred = False @@ -146,7 +156,10 @@ def symbolic(g, x1, rq1, x2, rq2, rq_out): mul = 1 add = 0 shift = 0 - n_l = 256 + if name == "out": + n_l = out_n_levels + else: + n_l = 256 requant = 0 if name == "out": if module: @@ -175,7 +188,7 @@ def __init__(self, in1_requant : Optional[nn.Module], in2_requant : Optional[nn. def forward(self, x1, x2): - return self.DORYAdderFun.apply(x1, self.in1_requant, x2, self.in2_requant, self.out_requant) + return self.DORYAdderFun.apply(x1, self.in1_requant, x2, self.in2_requant, self.out_requant, self.out_n_levels) class DORYReplaceAddersPass(OpTreeReplacementPass):