diff --git a/llvm/projects/keras/frontend/promise_translator.py b/llvm/projects/keras/frontend/promise_translator.py index a4843c529136980df82f22e16d9e4dc2db878c5d..ab8008293d39ccc1278e0a45bc51a1ed974a1439 100644 --- a/llvm/projects/keras/frontend/promise_translator.py +++ b/llvm/projects/keras/frontend/promise_translator.py @@ -42,7 +42,14 @@ class State: def getConvOp(self): for i in range(self.num_ops): - if self.ops[i].layer_type == "Conv2D": + if self.ops[i].layer_type == "Conv2D" or self.ops[i].layer_type == "DepthwiseConv2D": + return self.ops[i] + return None # Should be unreachable + + + def getDepthwiseConvOp(self): + for i in range(self.num_ops): + if self.ops[i].layer_type == "DepthwiseConv2D": return self.ops[i] return None # Should be unreachable @@ -90,6 +97,12 @@ class State: return True return False + def isDepthwiseConv(self): + if "depthwise" in self.op_string: + return True + return False + + def isPool(self): if "pool" in self.op_string and self.num_ops == 1: return True @@ -373,7 +386,7 @@ class PromiseRtTranslator: if input_node_name in self.output_map: input_var_name = self.output_map[input_node_name] else: - print ("Input Var not found - Aborting....") + print ("Input Var not found - Aborting....", input_node_name, "\n") sys.exit(0) return input_var_name @@ -554,6 +567,52 @@ class PromiseRtTranslator: state.clear() + + def genDepthwiseConvLayer(self, state): + print ("\n\n Layer = ", state.op_string, "\n\n") + + conv_op = state.getDepthwiseConvOp() + first_op = state.getFirstOp() + last_op = state.getLastOp() + + input_var = self.getSingleInputName(first_op) + output_var = self.getVariableName(last_op) + + w_name, b_name = self.getWeightNames(conv_op) + + activation_id = state.getActivationID() + padding = state.getPadding() + pool_id, pool_size = state.getPoolInfo() + strides = state.getStrides() + + promise_layer_str = "void* " + output_var + " = " + promise_layer_str += "tensorConvolution(" + input_var + ", " + promise_layer_str += w_name + ", " + promise_layer_str += str(padding) + ", " + promise_layer_str += str(padding) + ", " + promise_layer_str += str(strides[0]) + ", " + promise_layer_str += str(strides[1]) + ", " + promise_layer_str += "1, " + + C = conv_op.weights.shape[2] + promise_layer_str += str(C) + "); \n" + + # FIX: ADD code for TensorAdd and ACTIVATION + # TODO: ADD code for TensorAdd and ACTIVATION + + print (promise_layer_str) + + self.program_str += promise_layer_str + + + + self.appendLayerString("DepthwiseConv", state) + + state.clear() + + + + def genSoftmaxLayer(self, state): print ("\n\n Layer = ", state.op_string, "\n\n") @@ -666,7 +725,10 @@ class PromiseRtTranslator: elif state.isConv(): self.genConvLayer(state) - + + elif state.isDepthwiseConv(): + self.genDepthwiseConvLayer(state) + elif state.isPool(): self.genPoolLayer(state) @@ -674,6 +736,8 @@ class PromiseRtTranslator: self.genActivationLayer(state) + + def shouldVisit(self, cur_node): layer_name = cur_node.layer_name # NOTE: visit a node if not already visited and all predecessors are visited @@ -715,6 +779,23 @@ class PromiseRtTranslator: self.traverseSuccessors(cur_node, state) + + + def handle_depthwise_conv(self, cur_node, state): + if not self.shouldVisit(cur_node): + return + + layer_name = cur_node.layer_name + print ("handle_depthwise_conv", layer_name) + self.visited_nodes[layer_name] = True + + self.genPreviousLayer(state) + + # Appending depthwise_conv to state + state.add(cur_node, "depthwise") + + self.traverseSuccessors(cur_node, state) + def handle_add(self, cur_node, state): @@ -807,6 +888,9 @@ class PromiseRtTranslator: if layer_type == "Conv2D": self.handle_conv(output_node, state) + if layer_type == "DepthwiseConv2D": + self.handle_depthwise_conv(output_node, state) + if layer_type == "Dense": self.handle_dense(output_node, state)