import functools import logging import math from typing import Dict, Iterable, Union import sympy import torch from torch.utils._sympy.value_ranges import ValueRangeAnalysis, ValueRanges from .ir import FloorDiv, InterpreterShim, LoopBody, ModularIndexing from .utils import sympy_subs from .virtualized import V log = logging.getLogger(__name__) def dominated_nodes( initial_queue: Union[torch.fx.Node, Iterable[torch.fx.Node]], skip_filter=None ): """Returns the set of nodes whose values depend on those within initial_queue""" if isinstance(initial_queue, torch.fx.Node): initial_queue = [initial_queue] dominated_set = set(initial_queue) while initial_queue: node = initial_queue.pop() for user in node.users: if skip_filter and skip_filter(user): continue if user not in dominated_set: dominated_set.add(user) initial_queue.append(user) return dominated_set def val_expressable_in_32_bits(val): if hasattr(val, "is_Boolean") and val.is_Boolean: return True if isinstance(val, sympy.Expr): assert val.is_constant() if val.is_Integer or val.is_Boolean: val = int(val) else: val = float(val) # bound within mantissa if isinstance(val, float): return val <= (2**24) and val >= -(2**24) if isinstance(val, int): iinfo = torch.iinfo(torch.int32) return val <= iinfo.max and val >= iinfo.min raise Exception(f"Unexpected value {val}") def range_expressable_in_32_bits(range): return val_expressable_in_32_bits(range.lower) and val_expressable_in_32_bits( range.upper ) class OptimizeIndexing: """ Performs Value Range Analysis on LoopBody's fx graph to reduce precision of intermediaries from int64 to int32. This is an important optimization for indexing kernels such as Upsample and Interpolate. """ def __init__( self, loop_body: LoopBody, indices_ranges: Dict[sympy.Symbol, int], indexing_exprs: Dict[str, sympy.Expr], ): self.loop_body = loop_body self.indices_range = indices_ranges self.indexing_exprs = indexing_exprs self.replacement_vals = {} self.interp_env = {} self.submodules = self.swap_submodules(dict(loop_body.submodules)) indirect_var_set = set(loop_body.indirect_vars) self.index_indirect_dependecies = { index: expr.free_symbols & indirect_var_set for index, expr in indexing_exprs.items() } self.all_graphs = [loop_body.root_block.graph] + [ block.graph for block in loop_body.subblocks.values() ] for k, v in indices_ranges.items(): self.replace_indirect(k, ValueRanges(0, v)) # avoid computing these values, pessimistically assume that they are unbounded self.tensor_values_set = dominated_nodes( [ node for node in self.all_nodes if node.target in ["load", "reduction"] or "masked_subblock" in node.target ] ) def run(self): """Compute Value Ranges and try reduce precision of 'to_dtype' nodes to int32 where possible""" int64_dtype_nodes = [ node for node in self.all_nodes if ( node.target == "to_dtype" and node.args[2] == torch.int64 and node not in self.tensor_values_set ) ] if not int64_dtype_nodes: return for node in self.tensor_values_set: # we need to evaluate masked_subblock to recurse, and we need to set indirect values if ( "masked_subblock" not in node.target and "set_indirect" not in node.target ): self.interp_env[node] = torch._inductor.optimize_indexing.ValueRanges( -math.inf, math.inf ) interpreter = InterpreterShim(self.loop_body.root_block.graph, self.submodules) interpreter.run(V.get_ops_handler(), initial_env=self.interp_env) # TODO - if dominated node of one to_dtype is not expressible in int32, # we should short circuit another to_dtype node if that node also dominates for node in int64_dtype_nodes: self.try_to_reduce_precision(node) def try_to_reduce_precision(self, node): # if a downstream use of a node explicitly converts to int32, or float16/float32/float64, # then it's precision is set for that chain of uses, and we don't need to consider those # dominated values def skip_filter(node): return node.target == "to_dtype" and node.args[2] in ( torch.int32, torch.float32, torch.float64, ) # TODO - there are dominated uses whose dtype does not depend on whether # we reduce the precision here, e.g. add(int64, int64) one of the args can be reduced to # int32 without changing the output precision of the node. this case hasn't shown up for dominated in dominated_nodes(node, skip_filter): if dominated.target in ["store", "output"]: continue if "set_indirect" in dominated.target: idx = int(dominated.target[len("set_indirect") :]) indirect_var = self.loop_body.indirect_vars[idx] for index, indirect_vals in self.index_indirect_dependecies.items(): if indirect_var in indirect_vals: index_val = self.replacement_vals[index] if math.isinf(index_val.lower) or math.isinf(index_val.upper): return # all indices are integers, so make sure that we # use the bounds of integers instead of floats. # TODO - not sure if we should be doing int/float casts while tracing, # might interfere with sympy. index_val_int = ValueRanges( int(index_val.lower), int(index_val.upper) ) if not range_expressable_in_32_bits(index_val_int): return if not range_expressable_in_32_bits(self.interp_env[dominated]): return args = list(node.args) args[2] = torch.int32 node.args = tuple(args) @property def all_nodes(self): for graph in self.all_graphs: for node in graph.nodes: yield node def swap_submodules(self, submodules): keys = list(submodules.keys()) for key in keys: if key == "get_index": submodules[key] = self.get_index elif "masked_subblock" in key: subblock = self.loop_body.subblocks[key] submodules[key] = functools.partial( self.masked_subblock, subblock, self.interp_env ) else: assert "set_indirect" in key idx = int(key[len("set_indirect") :]) var = self.loop_body.indirect_vars[idx] indirect = functools.partial(self.set_indirect, var) submodules[key] = indirect return submodules def masked_subblock(self, subblock, env, mask, value): interp = InterpreterShim(subblock.graph, self.submodules) interp.run(V.get_ops_handler(), initial_env=env) output = [node for node in subblock.graph.nodes if node.target == "output"] assert len(output) == 1 # dont bother unioning with value since the load from buffer will be # pessimistically assumed to be inf anyway return interp.env[output[0]] def set_indirect(self, var, new_var): self.replace_indirect(var, new_var) return new_var def replace_indirect(self, old, new): """Swap in a variable used in indirect indexing""" assert isinstance(new, ValueRanges) self.replacement_vals[old] = new def get_index(self, name): if name in self.replacement_vals: return self.replacement_vals[name] out = self._get_index_impl(name) self.replacement_vals[name] = out return out def _get_index_impl(self, name): expr = self.indexing_exprs[name] free_symbols = list(expr.free_symbols) if len(free_symbols) == 0: return ValueRanges(expr, expr) if expr in self.replacement_vals: return self.replacement_vals[expr] def replace_symbols_for_deriv(expr, ignore_mod=False): # for the purposes of finding local, minimum, maximum, assume smoothness def mod_indexing_rep(x, y, z): if z.is_constant(): return x / y # never really happens, we'll bail on optimizing return (x / y) % z def indexing_div_rep(x, y): return x / y return expr.replace(ModularIndexing, mod_indexing_rep).replace( FloorDiv, indexing_div_rep ) symbols = expr.free_symbols monotonic_increasing = [] monotonic_decreasing = [] other_symbols = [] expr_for_deriv = replace_symbols_for_deriv(expr, True) for symbol in symbols: diff = sympy.diff(expr_for_deriv, symbol) if diff.is_positive: monotonic_increasing.append(symbol) elif diff.is_positive is False: # can return None monotonic_decreasing.append(symbol) else: other_symbols.append(symbol) if not other_symbols: max_val = sympy_subs( expr, { k: (v.upper if k in monotonic_increasing else v.lower) for k, v in self.replacement_vals.items() }, ) min_val = sympy_subs( expr, { k: (v.lower if k in monotonic_increasing else v.upper) for k, v in self.replacement_vals.items() }, ) return ValueRanges(min_val, max_val) else: # bail on optimizing, have not run into this yet return ValueRanges(-math.inf, math.inf) def indexing_dtype_strength_reduction(loop_body: LoopBody): """ Performs Value Range Analysis on LoopBody's fx graph to reduce precision of intermediaries from int64 to int32 """ indices = dict(loop_body.var_ranges) indexing = dict(loop_body.indexing_exprs) with V.set_ops_handler(ValueRangeAnalysis()): OptimizeIndexing(loop_body, indices, indexing).run()