def optimize(self, graph: Graph) -> Tuple[Graph, bool]: flag_changed = False for op in traverse.filter_nodes(traverse.listup_operators(graph), Sgemm): # type: Sgemm A = op.inputs["A"] B = op.inputs["B"] M = op.M N = op.N K = op.K transpose_A = op.transpose_A transpose_B = op.transpose_B if transpose_A: if TextureShape.get(A) != [M, K]: flag_changed = True TextureShape.set(A, width=K, height=M) else: if TextureShape.get(A) != [K, M]: flag_changed = True TextureShape.set(A, width=M, height=K) if transpose_B: if TextureShape.get(B) != [K, N]: flag_changed = True TextureShape.set(B, width=N, height=K) else: if TextureShape.get(B) != [N, K]: flag_changed = True TextureShape.set(B, width=K, height=N) return graph, flag_changed
def optimize(self, graph: Graph) -> Tuple[Graph, bool]: flag_changed = False for op in traverse.filter_nodes(traverse.listup_operators(graph), Tensordot): # type: Tensordot A = op.inputs["A"] B = op.inputs["B"] axes = op.axes K = mul(A.shape_dict[a] for a in axes[0]) M = A.size // K N = B.size // K if all([self.optimize_channel_mode, K % 4 == 0]): if ChannelMode.get(A) != ChannelModeEnum.RGBA: flag_changed = True ChannelMode.set(A, ChannelModeEnum.RGBA) if ChannelMode.get(B) != ChannelModeEnum.RGBA: flag_changed = True ChannelMode.set(B, ChannelModeEnum.RGBA) texture_shape_A = [M, K // 4] texture_shape_B = [N, K // 4] else: if ChannelMode.get(A) != ChannelModeEnum.R: flag_changed = True ChannelMode.set(A, ChannelModeEnum.R) if ChannelMode.get(B) != ChannelModeEnum.R: flag_changed = True ChannelMode.set(B, ChannelModeEnum.R) texture_shape_A = [M, K] texture_shape_B = [N, K] if TextureShape.get(A) != texture_shape_A: flag_changed = True TextureShape.set(A, height=texture_shape_A[0], width=texture_shape_A[1]) if TextureShape.get(B) != texture_shape_B: flag_changed = True TextureShape.set(B, height=texture_shape_B[0], width=texture_shape_B[1]) return graph, flag_changed
def update(self) -> bool: base = self.base # type: Concat y = base.outputs["y"] workspace = base.inputs["workspace"] flag_changed = False if y.order != workspace.order: flag_changed = True workspace.change_order(base.outputs["y"].order) if TextureShape.get(y) != TextureShape.get(workspace): flag_changed = True width, height = TextureShape.get(y) TextureShape.set(workspace, height, width) return flag_changed
def _replace_input(op: Operator, var_name: str, target: ChannelModeEnum): """ before) v -{op}- after) v -{conversion}- v' -{op}- """ v = op.inputs[var_name] if ChannelMode.get(v) == target: return False if target == ChannelModeEnum.RGBA: v_new = convert_r_to_rgba(v) else: v_new = convert_rgba_to_r(v) TextureShape.set(v_new, height=TextureShape.get(v)[0], width=TextureShape.get(v)[1]) op.replace_input(v, v_new) return True
def optimize(self, graph: Graph) -> Tuple[Graph, bool]: flag_changed = False for op in traverse.filter_nodes(traverse.listup_operators(graph), Tensordot): A = op.inputs["A"] B = op.inputs["B"] axes = op.axes K = mul(A.shape_dict[a] for a in axes[0]) M = A.size // K N = B.size // K if K % 4 == 0: if ChannelMode.get(A) != ChannelModeEnum.RGBA: flag_changed = True ChannelMode.set(A, ChannelModeEnum.RGBA) if ChannelMode.get(B) != ChannelModeEnum.RGBA: flag_changed = True ChannelMode.set(B, ChannelModeEnum.RGBA) else: if ChannelMode.get(A) != ChannelModeEnum.R: flag_changed = True ChannelMode.set(A, ChannelModeEnum.R) if ChannelMode.get(B) != ChannelModeEnum.R: flag_changed = True ChannelMode.set(B, ChannelModeEnum.R) if TextureShape.get(A) != (M, K): flag_changed = True TextureShape.set(A, height=M, width=K) if TextureShape.get(B) != (N, K): flag_changed = True TextureShape.set(B, height=N, width=K) return graph, flag_changed
def optimize(self, graph: Graph) -> Tuple[Graph, bool]: flag_changed = False for sgemm in traverse.filter_nodes(traverse.listup_operators(graph), Sgemm): # type: Sgemm A = sgemm.inputs["A"] B = sgemm.inputs["B"] M = sgemm.M N = sgemm.N K = sgemm.K transpose_A = sgemm.transpose_A transpose_B = sgemm.transpose_B if all([ self.optimize_channel_mode, K % 4 == 0, isinstance(A, ConstantVariable) or transpose_A == True, isinstance(B, ConstantVariable) or transpose_B == False ]): if transpose_A != True: assert isinstance(A, ConstantVariable) flag_changed = True old_A = A A = ConstantVariable( A.data.reshape([K, M]).transpose(), Order([Axis(None), Axis(None)])) ChannelMode.set(A, ChannelMode.get(old_A)) sgemm.replace_input(old_A, A, with_assert=False) sgemm.parameters["transpose_A"] = transpose_A = True if transpose_B != False: assert isinstance(B, ConstantVariable) flag_changed = True old_B = B B = ConstantVariable( B.data.reshape([K, N]).transpose(), Order([Axis(None), Axis(None)])) ChannelMode.set(B, ChannelMode.get(old_B)) sgemm.replace_input(old_B, B, with_assert=False) sgemm.parameters["transpose_B"] = transpose_B = False if ChannelMode.get(A) != ChannelModeEnum.RGBA: flag_changed = True ChannelMode.set(A, ChannelModeEnum.RGBA) if ChannelMode.get(B) != ChannelModeEnum.RGBA: flag_changed = True ChannelMode.set(B, ChannelModeEnum.RGBA) texture_shape_A = [M, K // 4] if transpose_A else [K // 4, M] texture_shape_B = [K // 4, N] if transpose_B else [N, K // 4] else: if ChannelMode.get(A) != ChannelModeEnum.R: flag_changed = True ChannelMode.set(A, ChannelModeEnum.R) if ChannelMode.get(B) != ChannelModeEnum.R: flag_changed = True ChannelMode.set(B, ChannelModeEnum.R) texture_shape_A = [M, K] if transpose_A else [K, M] texture_shape_B = [K, N] if transpose_B else [N, K] if TextureShape.get(A) != texture_shape_A: flag_changed = True TextureShape.set(A, height=texture_shape_A[0], width=texture_shape_A[1]) if TextureShape.get(B) != texture_shape_B: flag_changed = True TextureShape.set(B, height=texture_shape_B[0], width=texture_shape_B[1]) if flag_changed: graph, _ = ConstantFolding().optimize(graph) return graph, flag_changed