src/main/java/org/apache/sysds/runtime/lineage/LineageDedupUtils.java - systemds - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you 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.
  */

 package org.apache.sysds.runtime.lineage;

 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.Stack;

 import org.apache.sysds.runtime.controlprogram.BasicProgramBlock;
 import org.apache.sysds.runtime.controlprogram.ForProgramBlock;
 import org.apache.sysds.runtime.controlprogram.FunctionProgramBlock;
 import org.apache.sysds.runtime.controlprogram.IfProgramBlock;
 import org.apache.sysds.runtime.controlprogram.ProgramBlock;
 import org.apache.sysds.runtime.controlprogram.WhileProgramBlock;
 import org.apache.sysds.runtime.controlprogram.context.ExecutionContext;
 import org.apache.sysds.utils.Explain;

 public class LineageDedupUtils {
 	public static final String DEDUP_DELIM = "_";
 	private static Lineage _tmpLineage = null;
 	private static Lineage _mainLineage = null;
 	private static ArrayList<Long> _numDistinctPaths = new ArrayList<>();
 	private static long _maxNumPaths = 0;
 	private static int _numPaths = 0;

 	public static boolean isValidDedupBlock(ProgramBlock pb, boolean inLoop) {
 		// Only the last level loop-body in nested loop structure is valid for deduplication
 		boolean ret = true; //basic program block
 		if (pb instanceof FunctionProgramBlock) {
 			FunctionProgramBlock fsb = (FunctionProgramBlock)pb;
 			for (ProgramBlock cpb : fsb.getChildBlocks())
 				ret &= isValidDedupBlock(cpb, inLoop);
 		}
 		else if (pb instanceof WhileProgramBlock) {
 			if( inLoop ) return false;
 			WhileProgramBlock wpb = (WhileProgramBlock) pb;
 			for (ProgramBlock cpb : wpb.getChildBlocks())
 				ret &= isValidDedupBlock(cpb, true);
 		}
 		else if (pb instanceof IfProgramBlock) {
 			IfProgramBlock ipb = (IfProgramBlock) pb;
 			for (ProgramBlock cpb : ipb.getChildBlocksIfBody())
 				ret &= isValidDedupBlock(cpb, inLoop);
 			for (ProgramBlock cpb : ipb.getChildBlocksElseBody())
 				ret &= isValidDedupBlock(cpb, inLoop);
 		}
 		else if (pb instanceof ForProgramBlock) { //incl parfor
 			if( inLoop ) return false;
 			ForProgramBlock fpb = (ForProgramBlock) pb;
 			for (ProgramBlock cpb : fpb.getChildBlocks())
 				ret &= isValidDedupBlock(cpb, true);
 		}
 		return ret;
 	}

 	public static LineageDedupBlock computeDedupBlock(ProgramBlock fpb, ExecutionContext ec) {
 		LineageDedupBlock ldb = new LineageDedupBlock();
 		ec.getLineage().setInitDedupBlock(ldb);
 		ldb.traceProgramBlocks(fpb.getChildBlocks(), ec);
 		ec.getLineage().setInitDedupBlock(null);
 		return ldb;
 	}

 	public static LineageDedupBlock initializeDedupBlock(ProgramBlock fpb, ExecutionContext ec) {
 		LineageDedupBlock ldb = new LineageDedupBlock();
 		ec.getLineage().setInitDedupBlock(ldb);
 		// create/reuse a lineage object to trace the loop iterations
 		initLocalLineage(ec);
 		// save the original lineage object
 		_mainLineage = ec.getLineage();
 		// count and save the number of distinct paths
 		ldb.setNumPathsInPBs(fpb.getChildBlocks(), ec);
 		ec.getLineage().setInitDedupBlock(null);
 		return ldb;
 	}

 	public static void setNewDedupPatch(LineageDedupBlock ldb, ProgramBlock fpb, ExecutionContext ec) {
 		// No need to trace anymore if all the paths are taken,
 		// instead reuse the stored maps for this and future interations
 		// But, tracing is required if reuse is enabled, even though the traced
 		// lineage map is discarded after each iteration (if all paths are taken)
 		if (ldb.isAllPathsTaken() && LineageCacheConfig.ReuseCacheType.isNone())
 			return;

 		// copy the input LineageItems of the loop-body
 		initLocalLineage(ec);
 		ArrayList<String> inputnames = fpb.getStatementBlock().getInputstoSB();
 		LineageItem[] liinputs = LineageItemUtils.getLineageItemInputstoSB(inputnames, ec);
 		// TODO: find the inputs from the ProgramBlock instead of StatementBlock
 		String ph = LineageItemUtils.LPLACEHOLDER;
 		int i = 0;
 		for (String input : inputnames) {
 			// Skip empty variables to correctly map lineage items to live variables
 			if (ec.getVariable(input) == null)
 				continue;
 			// Wrap the inputs with order-preserving placeholders.
 			LineageItem phInput = new LineageItem(ph+String.valueOf(i), new LineageItem[] {liinputs[i]});
 			_tmpLineage.set(input, phInput);
 			i++;
 		}
 		// also copy the dedupblock to trace the taken path (bitset)
 		_tmpLineage.setDedupBlock(ldb);
 		// attach the lineage object to the execution context
 		ec.setLineage(_tmpLineage);
 	}

 	public static void replaceLineage(ExecutionContext ec) {
 		// replace the local lineage with the original one
 		ec.setLineage(_mainLineage);
 	}

 	public static Map<String, Integer> setDedupMap(LineageDedupBlock ldb, long takenPath) {
 		// if this iteration took a new path, store the corresponding map
 		if (ldb.getMap(takenPath) == null) {
 			LineageMap patchMap = _tmpLineage.getLineageMap();
 			// Clean unused variables, and cut the DAGs at placeholders
 			cleanDedupMap(patchMap);
 			ldb.setMap(takenPath, patchMap);
 		}
 		// Copy and return the hash values of all the roots
 		if (!LineageCacheConfig.ReuseCacheType.isNone()) {
 			Map<String, Integer> dedupPatchHashList = new HashMap<>();
 			LineageMap patchMap = _tmpLineage.getLineageMap();
 			for (Map.Entry<String, LineageItem> litem : patchMap.getTraces().entrySet()) {
 				if (!litem.getValue().isPlaceholder())
 					dedupPatchHashList.put(litem.getKey(), litem.getValue().hashCode());
 			}
 			return dedupPatchHashList;
 		}
 		return null;
 	}

 	private static void initLocalLineage(ExecutionContext ec) {
 		_mainLineage = ec.getLineage();
 		_tmpLineage = _tmpLineage == null ? new Lineage() : _tmpLineage;
 		_tmpLineage.clearLineageMap();
 		_tmpLineage.clearDedupBlock();
 	}

 	public static String mergeExplainDedupBlocks(ExecutionContext ec) {
 		Map<ProgramBlock, LineageDedupBlock> dedupBlocks = ec.getLineage().getDedupBlocks();
 		StringBuilder sb = new StringBuilder();
 		// Gather all the DAG roots of all the paths in all the loops.
 		for (Map.Entry<ProgramBlock, LineageDedupBlock> dblock : dedupBlocks.entrySet()) {
 			if (dblock.getValue() != null) {
 				String forKey = dblock.getKey().getStatementBlock().getName();
 				LineageDedupBlock dedup = dblock.getValue();
 				for (Map.Entry<Long, LineageMap> patch : dedup.getPathMaps().entrySet()) {
 					for (Map.Entry<String, LineageItem> root : patch.getValue().getTraces().entrySet()) {
 						// Encode all the information in the headers that're
 						// needed by the deserialization logic.
 						sb.append("patch");
 						sb.append(DEDUP_DELIM);
 						sb.append(root.getKey());
 						sb.append(DEDUP_DELIM);
 						sb.append(forKey);
 						sb.append(DEDUP_DELIM);
 						sb.append(patch.getKey());
 						sb.append("\n");
 						sb.append(Explain.explain(root.getValue()));
 						sb.append("\n");

 					}
 				}
 			}
 		}
 		return sb.toString();
 	}

 	private static void cleanDedupMap(LineageMap lmap) {
 		// Gather all the DAG roots and cut each at placeholder
 		Set<String> emptyRoots = new HashSet<>();
 		for (Map.Entry<String, LineageItem> litem : lmap.getTraces().entrySet()) {
 			LineageItem root = litem.getValue();
 			// Clean empty DAG roots such as iterator i
 			if (root.isPlaceholder()) {
 				emptyRoots.add(litem.getKey());
 				continue;
 			}
 			root.resetVisitStatusNR();
 			cutAtPlaceholder(root);
 			if (!LineageCacheConfig.ReuseCacheType.isNone())
 				// These chopped DAGs can lead to incorrect reuse.
 				// FIXME: This logic removes only the live variables from lineage cache.
 				//        Need a way to remove all entries that are cached in this iteration.
 				LineageCache.removeEntry(root);
 		}
 		lmap.getTraces().keySet().removeAll(emptyRoots);
 	}

 	private static void cutAtPlaceholder(LineageItem root) {
 		Stack<LineageItem> q = new Stack<>();
 		q.push(root);
 		while (!q.empty()) {
 			LineageItem tmp = q.pop();
 			if (tmp.isVisited())
 				continue;

 			if (tmp.isPlaceholder()) {
 				// set inputs to null
 				tmp.resetInputs();
 				tmp.setVisited();
 				continue;
 			}

 			if (tmp.getInputs() != null)
 				for (int i=0; i<tmp.getInputs().length; i++) {
 					LineageItem li = tmp.getInputs()[i];
 					q.push(li);
 				}
 			tmp.setVisited();
 		}
 	}

 	//------------------------------------------------------------------------------
 	/* The below static functions help to compute the number of distinct paths
 	 * in any program block, and are used for diagnostic purposes. These will
 	 * be removed in future.
 	 */

 	public static long computeNumPaths(ProgramBlock fpb, ExecutionContext ec) {
 		if (fpb == null || fpb.getChildBlocks() == null)
 			return 0;
 		_numDistinctPaths.clear();
 		long n = numPathsInPBs(fpb.getChildBlocks(), ec);
 		if (n > _maxNumPaths) {
 			_maxNumPaths = n;
 			System.out.println("\nmax no of paths : " + _maxNumPaths + "\n");
 		}
 		return n;
 	}

 	public static long numPathsInPBs (ArrayList<ProgramBlock> pbs, ExecutionContext ec) {
 		if (_numDistinctPaths.size() == 0)
 			_numDistinctPaths.add(0L);
 		for (ProgramBlock pb : pbs)
 			numPathsInPB(pb, ec, _numDistinctPaths);
 		return _numDistinctPaths.size();
 	}

 	private static void numPathsInPB(ProgramBlock pb, ExecutionContext ec, ArrayList<Long> paths) {
 		if (pb instanceof IfProgramBlock)
 			numPathsInIfPB((IfProgramBlock)pb, ec, paths);
 		else if (pb instanceof BasicProgramBlock)
 			return;
 		else
 			return;
 	}

 	private static void numPathsInIfPB(IfProgramBlock ipb, ExecutionContext ec, ArrayList<Long> paths) {
 		ipb.setLineageDedupPathPos(_numPaths++);
 		ArrayList<Long> rep = new ArrayList<>();
 		int pathKey = 1 << (_numPaths-1);
 		for (long p : paths) {
 			long pathIndex = p | pathKey;
 			rep.add(pathIndex);
 		}
 		_numDistinctPaths.addAll(rep);
 		for (ProgramBlock pb : ipb.getChildBlocksIfBody())
 			numPathsInPB(pb, ec, rep);
 		for (ProgramBlock pb : ipb.getChildBlocksElseBody())
 			numPathsInPB(pb, ec, paths);
 	}
 }
	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you 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.
	*/

	package org.apache.sysds.runtime.lineage;

	import java.util.ArrayList;
	import java.util.HashMap;
	import java.util.HashSet;
	import java.util.Map;
	import java.util.Set;
	import java.util.Stack;

	import org.apache.sysds.runtime.controlprogram.BasicProgramBlock;
	import org.apache.sysds.runtime.controlprogram.ForProgramBlock;
	import org.apache.sysds.runtime.controlprogram.FunctionProgramBlock;
	import org.apache.sysds.runtime.controlprogram.IfProgramBlock;
	import org.apache.sysds.runtime.controlprogram.ProgramBlock;
	import org.apache.sysds.runtime.controlprogram.WhileProgramBlock;
	import org.apache.sysds.runtime.controlprogram.context.ExecutionContext;
	import org.apache.sysds.utils.Explain;

	public class LineageDedupUtils {
	public static final String DEDUP_DELIM = "_";
	private static Lineage _tmpLineage = null;
	private static Lineage _mainLineage = null;
	private static ArrayList<Long> _numDistinctPaths = new ArrayList<>();
	private static long _maxNumPaths = 0;
	private static int _numPaths = 0;

	public static boolean isValidDedupBlock(ProgramBlock pb, boolean inLoop) {
	// Only the last level loop-body in nested loop structure is valid for deduplication
	boolean ret = true; //basic program block
	if (pb instanceof FunctionProgramBlock) {
	FunctionProgramBlock fsb = (FunctionProgramBlock)pb;
	for (ProgramBlock cpb : fsb.getChildBlocks())
	ret &= isValidDedupBlock(cpb, inLoop);
	}
	else if (pb instanceof WhileProgramBlock) {
	if( inLoop ) return false;
	WhileProgramBlock wpb = (WhileProgramBlock) pb;
	for (ProgramBlock cpb : wpb.getChildBlocks())
	ret &= isValidDedupBlock(cpb, true);
	}
	else if (pb instanceof IfProgramBlock) {
	IfProgramBlock ipb = (IfProgramBlock) pb;
	for (ProgramBlock cpb : ipb.getChildBlocksIfBody())
	ret &= isValidDedupBlock(cpb, inLoop);
	for (ProgramBlock cpb : ipb.getChildBlocksElseBody())
	ret &= isValidDedupBlock(cpb, inLoop);
	}
	else if (pb instanceof ForProgramBlock) { //incl parfor
	if( inLoop ) return false;
	ForProgramBlock fpb = (ForProgramBlock) pb;
	for (ProgramBlock cpb : fpb.getChildBlocks())
	ret &= isValidDedupBlock(cpb, true);
	}
	return ret;
	}

	public static LineageDedupBlock computeDedupBlock(ProgramBlock fpb, ExecutionContext ec) {
	LineageDedupBlock ldb = new LineageDedupBlock();
	ec.getLineage().setInitDedupBlock(ldb);
	ldb.traceProgramBlocks(fpb.getChildBlocks(), ec);
	ec.getLineage().setInitDedupBlock(null);
	return ldb;
	}

	public static LineageDedupBlock initializeDedupBlock(ProgramBlock fpb, ExecutionContext ec) {
	LineageDedupBlock ldb = new LineageDedupBlock();
	ec.getLineage().setInitDedupBlock(ldb);
	// create/reuse a lineage object to trace the loop iterations
	initLocalLineage(ec);
	// save the original lineage object
	_mainLineage = ec.getLineage();
	// count and save the number of distinct paths
	ldb.setNumPathsInPBs(fpb.getChildBlocks(), ec);
	ec.getLineage().setInitDedupBlock(null);
	return ldb;
	}

	public static void setNewDedupPatch(LineageDedupBlock ldb, ProgramBlock fpb, ExecutionContext ec) {
	// No need to trace anymore if all the paths are taken,
	// instead reuse the stored maps for this and future interations
	// But, tracing is required if reuse is enabled, even though the traced
	// lineage map is discarded after each iteration (if all paths are taken)
	if (ldb.isAllPathsTaken() && LineageCacheConfig.ReuseCacheType.isNone())
	return;

	// copy the input LineageItems of the loop-body
	initLocalLineage(ec);
	ArrayList<String> inputnames = fpb.getStatementBlock().getInputstoSB();
	LineageItem[] liinputs = LineageItemUtils.getLineageItemInputstoSB(inputnames, ec);
	// TODO: find the inputs from the ProgramBlock instead of StatementBlock
	String ph = LineageItemUtils.LPLACEHOLDER;
	int i = 0;
	for (String input : inputnames) {
	// Skip empty variables to correctly map lineage items to live variables
	if (ec.getVariable(input) == null)
	continue;
	// Wrap the inputs with order-preserving placeholders.
	LineageItem phInput = new LineageItem(ph+String.valueOf(i), new LineageItem[] {liinputs[i]});
	_tmpLineage.set(input, phInput);
	i++;
	}
	// also copy the dedupblock to trace the taken path (bitset)
	_tmpLineage.setDedupBlock(ldb);
	// attach the lineage object to the execution context
	ec.setLineage(_tmpLineage);
	}

	public static void replaceLineage(ExecutionContext ec) {
	// replace the local lineage with the original one
	ec.setLineage(_mainLineage);
	}

	public static Map<String, Integer> setDedupMap(LineageDedupBlock ldb, long takenPath) {
	// if this iteration took a new path, store the corresponding map
	if (ldb.getMap(takenPath) == null) {
	LineageMap patchMap = _tmpLineage.getLineageMap();
	// Clean unused variables, and cut the DAGs at placeholders
	cleanDedupMap(patchMap);
	ldb.setMap(takenPath, patchMap);
	}
	// Copy and return the hash values of all the roots
	if (!LineageCacheConfig.ReuseCacheType.isNone()) {
	Map<String, Integer> dedupPatchHashList = new HashMap<>();
	LineageMap patchMap = _tmpLineage.getLineageMap();
	for (Map.Entry<String, LineageItem> litem : patchMap.getTraces().entrySet()) {
	if (!litem.getValue().isPlaceholder())
	dedupPatchHashList.put(litem.getKey(), litem.getValue().hashCode());
	}
	return dedupPatchHashList;
	}
	return null;
	}

	private static void initLocalLineage(ExecutionContext ec) {
	_mainLineage = ec.getLineage();
	_tmpLineage = _tmpLineage == null ? new Lineage() : _tmpLineage;
	_tmpLineage.clearLineageMap();
	_tmpLineage.clearDedupBlock();
	}

	public static String mergeExplainDedupBlocks(ExecutionContext ec) {
	Map<ProgramBlock, LineageDedupBlock> dedupBlocks = ec.getLineage().getDedupBlocks();
	StringBuilder sb = new StringBuilder();
	// Gather all the DAG roots of all the paths in all the loops.
	for (Map.Entry<ProgramBlock, LineageDedupBlock> dblock : dedupBlocks.entrySet()) {
	if (dblock.getValue() != null) {
	String forKey = dblock.getKey().getStatementBlock().getName();
	LineageDedupBlock dedup = dblock.getValue();
	for (Map.Entry<Long, LineageMap> patch : dedup.getPathMaps().entrySet()) {
	for (Map.Entry<String, LineageItem> root : patch.getValue().getTraces().entrySet()) {
	// Encode all the information in the headers that're
	// needed by the deserialization logic.
	sb.append("patch");
	sb.append(DEDUP_DELIM);
	sb.append(root.getKey());
	sb.append(DEDUP_DELIM);
	sb.append(forKey);
	sb.append(DEDUP_DELIM);
	sb.append(patch.getKey());
	sb.append("\n");
	sb.append(Explain.explain(root.getValue()));
	sb.append("\n");

	}
	}
	}
	}
	return sb.toString();
	}

	private static void cleanDedupMap(LineageMap lmap) {
	// Gather all the DAG roots and cut each at placeholder
	Set<String> emptyRoots = new HashSet<>();
	for (Map.Entry<String, LineageItem> litem : lmap.getTraces().entrySet()) {
	LineageItem root = litem.getValue();
	// Clean empty DAG roots such as iterator i
	if (root.isPlaceholder()) {
	emptyRoots.add(litem.getKey());
	continue;
	}
	root.resetVisitStatusNR();
	cutAtPlaceholder(root);
	if (!LineageCacheConfig.ReuseCacheType.isNone())
	// These chopped DAGs can lead to incorrect reuse.
	// FIXME: This logic removes only the live variables from lineage cache.
	// Need a way to remove all entries that are cached in this iteration.
	LineageCache.removeEntry(root);
	}
	lmap.getTraces().keySet().removeAll(emptyRoots);
	}

	private static void cutAtPlaceholder(LineageItem root) {
	Stack<LineageItem> q = new Stack<>();
	q.push(root);
	while (!q.empty()) {
	LineageItem tmp = q.pop();
	if (tmp.isVisited())
	continue;

	if (tmp.isPlaceholder()) {
	// set inputs to null
	tmp.resetInputs();
	tmp.setVisited();
	continue;
	}

	if (tmp.getInputs() != null)
	for (int i=0; i<tmp.getInputs().length; i++) {
	LineageItem li = tmp.getInputs()[i];
	q.push(li);
	}
	tmp.setVisited();
	}
	}

	//------------------------------------------------------------------------------
	/* The below static functions help to compute the number of distinct paths
	* in any program block, and are used for diagnostic purposes. These will
	* be removed in future.
	*/

	public static long computeNumPaths(ProgramBlock fpb, ExecutionContext ec) {
	if (fpb == null \|\| fpb.getChildBlocks() == null)
	return 0;
	_numDistinctPaths.clear();
	long n = numPathsInPBs(fpb.getChildBlocks(), ec);
	if (n > _maxNumPaths) {
	_maxNumPaths = n;
	System.out.println("\nmax no of paths : " + _maxNumPaths + "\n");
	}
	return n;
	}

	public static long numPathsInPBs (ArrayList<ProgramBlock> pbs, ExecutionContext ec) {
	if (_numDistinctPaths.size() == 0)
	_numDistinctPaths.add(0L);
	for (ProgramBlock pb : pbs)
	numPathsInPB(pb, ec, _numDistinctPaths);
	return _numDistinctPaths.size();
	}

	private static void numPathsInPB(ProgramBlock pb, ExecutionContext ec, ArrayList<Long> paths) {
	if (pb instanceof IfProgramBlock)
	numPathsInIfPB((IfProgramBlock)pb, ec, paths);
	else if (pb instanceof BasicProgramBlock)
	return;
	else
	return;
	}

	private static void numPathsInIfPB(IfProgramBlock ipb, ExecutionContext ec, ArrayList<Long> paths) {
	ipb.setLineageDedupPathPos(_numPaths++);
	ArrayList<Long> rep = new ArrayList<>();
	int pathKey = 1 << (_numPaths-1);
	for (long p : paths) {
	long pathIndex = p \| pathKey;
	rep.add(pathIndex);
	}
	_numDistinctPaths.addAll(rep);
	for (ProgramBlock pb : ipb.getChildBlocksIfBody())
	numPathsInPB(pb, ec, rep);
	for (ProgramBlock pb : ipb.getChildBlocksElseBody())
	numPathsInPB(pb, ec, paths);
	}
	}