Inside TensorFlow: Graph rewriting (Macros, not functions)

Key Takeaways

my name is Eugene I'm working on grappa and performance but today I'm going to talk mostly about how we do graduating for functions so functions is very important part of TF 2.0 it's very important for the end-user how user thinks about the flow computation so it's not graph anymore it's a composition of functions but that creates a bunch of problems for us at runtime and so basically we have to rewrite autograph and all the factions to make it executable and to avoid some specific problems that comes with the functions so this is the on the left of this button code so this kind of v1 Python that people usually write so you define your variables some constants you do assignment and then you want to return update variables so if you don't add explicit control dependency you're going to read the initial values so you have to think what are you dependencies and have to make everything very explicit and so on the right this is some textual graph graph representation so there's not enough space to move to the full graph def and the prota so this is roughly how it will look like if you do to the back string internal flow if you get a graph death due to debug shrink you can get some short representation of the graph something like on the right so variables actually will handle up this is single kernel that returns DT a source data type constants just simple controls with the values and then when it was an ad it is an assignment as a fulcrum that takes resource control from our handle up and the constant and assigns and add a value and when you do it it is another emptiness or Colonel read the variable up so you get turn throw the th or scandal a and C and because I have explicit control dependencies this series depends on update a and bacey both of them and you fetch with a NGC but explicitly defining control dependencies is very annoying so intestine flow to point all the same graph will look something like this so we have your variables defined outside of a function scope and you have a tasteful functional notation I didn't get so I want to add two different constants two variables a and C and get back our data results so the Python looks something like on the left so you don't have to explicitly add control dependencies which make life easier and you don't have to think about program order so TF function traces the function and it is all the necessary control dependencies so when you execute a graph you will get the program order you would expect from there a co-op item code so where function deff so the graph representation looks something like on the right so we have a graph def and the graph def is basically a sequence of no deaths and we also have a function deff it's a way to group together some sub graph with input arguments and output arguments so this function on the left from the pattern will be transferred to a function deff that will be represented in something like this so we have two input parameters over source type a and B and the return type of the function is pair of ins so there update a assign add we are signing the new value to input resource and then we do to reach Biofreeze and we have control dependencies from the previous separation touching the same resource and this is edited automatically by there's a 42.0 tamati control dependencies so we don't have to do anything and then we have a return values so we have functions at graph level at function deff and graph def interval they have two types of returns you might return the data type so you return two tenders read a and rich C but also it has a special kind of donal notation control written so it doesn't in term any data but you might group some of the herbs inside the function deff and you can specify that these herbs have a name and they must always run and when you have a control dependency on the function call runtime verifies that every hope that is part of controllers on set will run and the main graph will look something like this so we have a partition call it is dev to point all function call mechanism there does partitioning and multi-device function invocation and graph optimisation and all the other details that are required for the runtime and read operations is just the identity they treats the first and the second output of function left so any questions does it clear like it's annotation on the right side so in this example in particular like the read operations already depend on the update so in this case the control dependency there is not that required but I guess in other case if you don't return the read values in those know though very depends only on the input a and C but it also it has a control dependency on the update yeah because this control dependency is it automatically by DF 2.0 at their function when it trace the function and if you have multiple TF ops touching the same resource automatic depends which I think will add control depends if I'm saying like in in terms of the return values like in this case see if you have the really de andrade see they automatically capture the if you try to fetch really and read C or tomorrow you make sure that the updates are runs in this is the control dependencies are not that useful yeah the control returns oh yeah well yes I mean just just getting out that in the if it would be more useful if we don't return the read values then in this case yeah a control test would be very useful yeah yeah so if you have some variable update for example you have a function that doesn't have returns that you are a function only for the side-effects updates some counters much of counter statistics in this case you don't have your eight outputs so this is why you have control job but right now chief 2.0 will air to these controller Jones to both places this is needed because to separate the data path and control path because when you in line that function and it turns out that no one is using output one so you might end up that you will not execute your stay tablets inside a function but it is guaranteed that everything on the control path on the controller zone will be executed under some conditions little mention later if you two variable are in a small sign in front control dependency I mean this one yeah so the data dependency goes by name so a is regulated dependency and this one is a control dependency so we don't read the actual output it just means that update a must execute before we can execute read a internal representation of control dependencies in the graph so yeah if you're up has two inputs you will have a B as inputs and then you can add as many as you want to control dependencies at a graphics equation time run time will make sure that everything in control dependencies will be executed before the kernel little special notation so here's another example so we have variables a and B vectors of length 10 and variable C it's just a counter all types are integers and we want to get a strata slice from variable a and B and then we also want to increment the counter every time we take a slice we want to know how many slices did you take for whatever reason so this graph does make much sense but anyway and we only interested in slice one and the graph will look something like this we have three variables work and elopes constant assign add for the counter and then we read variable this one so we read variables from this should be B so we read variables from the resource from the handle and then we take a strata slice and then we fetch slice one and because we fetch slice one we can remove slice zero it is not needed full slice one we can also remove from the graph a and the soul but you see that the nice property we get that slice zero is invalid if you try to take slice from 20 to 30 in a variable of size 10 you'll get fail at runtime but because you don't really need slice 0 everything's just fine as long as you don't fetch y 0 this graph is taught to the village but if you do the same graph but you put your slices inside the function so the same kind of thing you have three variables a b c c is the counter you do a sign ad inside the function so nice you don't have to add explicit control dependency you get the account updated automatically then you take slice 0 and slice 1 and the don't Apple and you invoke your function you don't you are not instance in slice 0 you get you want only to get the second value so you just keep it and you're trying to print it so at a graph depth level and function there it will look something of this you'll have a good slice function that takes three resources a B and C it will have an increment for the counter it will be the values of a B take a slices and it translates 0 and slice 1 and will have a control shown for increment automatically added by to metric depends the tracking so slices do not have any control dependencies because slices they do not depend on variable C they depends on different variables and then you have a function call partition call to the get slice a and B and you don't mean the first slice you take the second one my function call output one and you're trying to fetch it or print it out seven and the problem right now that this code will fail because the function semantics internal currently is strict with regards to its inputs and outputs so before function return can function can return it must finish the execution of all the output parameters and before function can start executing it must finish executing of all the inputs yeah so it's not the property so all inputs must be evaluated and must be not dead the function can start running and not net it means that if you have inputs from switches is sutures are used to represent control flow into the flow so there is special notion of that tensor so it so basically when you around your graph at runtime and you have dead tensors you don't execute any currency you just update the flag and propagate this deadness across the graph so we can have a large graph and if note that is represented a switch and based on the runtime value you execute only half of the graph so another property of the function is that all outputs must be related and also not dead so for example you have a function with the control flow and you want to take a gradient of that function if you have to return all intermediates from the function and some of them might be dead because they might be from intermediate computations for that from untaken branch and basic view your function can be dead but sometimes that includes okay so there is special way to avoid that and I think might not be used right now but before that when we generated gradient functions to the photo point all marked all the functions for the fourth pass that they can return that inputs and instead of dead inputs runtime would return empty tensor and move just continue so but the most confusing thing if you do this code without the F function notation just a regular Python function get slice a B and C the same inputs you have doing increment but then you have to add explicit control dependency because you're not using TF 2.0 to the automatic control dependency tracking so it will get graph on the right and it is exactly the same graph as you would get if you don't use functions at all if you just write it line by line one by one and it will go out in the four run time to prune all the notes that are not used and your graph is valid again but the only difference is that don't have TF dot function notation and even if I completely remove this counter and I don't want to count anything in C so the only difference would be the tip the function notation and it is a huge difference at run time one graph fails and another graph is just fine and this is the real example so Alex gave presentation embowered functions in terms of flow and that's how some people use it in reinforced running so you have a function that returns some environment state does action and does environment sure that you get ABEC and you just run the loop on action and in the end you're on your asset function so if you would annotate this function with TF dot function notation on every loop iteration you will get the state you will do the action and you would do there as that which is not fine another example from terms of flow probability it was kind of so you have TF dot function that has three inputs you have like the first output comes from some constant inside the function constant and some temptations then you take first second and third inputs and do some computation and output at position two and three and you pass it to exactly the same function and then you pass to the next function so we have the first input of the first function is some failed computation it might be placeholder that you didn't define it be a strata slice with innovative bounce or fail to read from an old file system but if you trace this fish we are trying to feature the second value if you trace across the function boundaries this one this one yeah we ended up in this constant and it's totally valid if we try if you try to fetch the last output of the last function and we'll trace it we'll get to the second input of the first function and it's totally valid so as long as we don't try to fetch this computation or this computation we don't really need the first input and a lot of code in terms of all probability relies on this property of the function that the function is not really a function call is just some Python scoping language construct that allows you to separate your graph construction and when people moving to terms of flow eager and they annotate all their functions with table function it all breaks because function at runtime level are strict but that's not what it was designed how people wanted to do functions so if you open graph protocol committee when functions were added to the total flow and there's very important thing so the column might start execution as soon as sum of inputs are ad and if you want to enforce that all your inputs ready you must use topple because runtime is allowed to start executing when the first snippet is ready it doesn't have to wait for all the inputs and the consumer of the function may start executing earlier is that function value is ready so if you want to ensure that your function has a strict semantics you must use tuples otherwise runtime is allowed to do motivate ones and start with the kidding as early as it once happened to be that implementation was strict and people didn't notice that and they're out of court International that relies on strict semantics but we don't really want strict functions we'd really love to get lazy functions so nothing is evaluated until it's really nice say and in tensorflow executive implementation necessary means that we want to execute the primitive gentle for colonel so primitive terms of whole kernel is any op like add multiply convolution might move and composite kernels are right now currently only functions i think there are some proposals to add more composite core support but it isn't know how to look like but currently we love functions to be lazy and ideally would laugh to make switch lazy as well so if you have switch your conditional execution right now we all have to execute all the inputs to the switch even if you don't use it later but is very hard to write to the current executor and also non stick semantics is very important for performance so imagine that you have to print all we have a function ResNet and we have a huge model that does then resonate influences and each ResNet has and no 256 input parameters so if these parameters are living on remote parameter servers you can't start educating your function before you fetch all these 256 tenders to your local machine and start running and then you have well you will take a lot of memory to keep them at the same time on the host on the device or whatever and it will take out of time until to get them on the network if even if we used in the 4.0 and functions do not get tensors as inputs and they just get resources and they do read variable hope we still have kind of the same program because we can't start executing the function before the previous update to that your source finished so before you can start running key or as near to 50 function even over the resources as inputs you have to wait for all the control dependencies so for all the previous updates to all the previous variables so even if don't get word of network traffic you still wait on parameter server completing that updates or maybe something else so we really related functions and run them as soon as something close ready start fishing parameters running the first layer and fetch the parameters for the second layer only when we completed the first layer so we didn't get lazy functions originally the cost functions were added to the tenth floor runtime much later than original executor and engraft dev was created and tensorflow executor and runtime is super strict so people often think that channel flow is lazy so you fetch way of a and then you get back and pull only whatever you need well that's not actually how it works so terms of runtime is very strict and greedy and easy using pushes and this received lazy evaluation it's just a product of pruning so we remove just before execution we remove all the notes that we don't need to compute the output and then we start looking at the notes one by one that idea they don't have inputs and run them and update all the consumers that you are on and this is the fundamental property of how executor CC intensive flow works and it's almost impossible to do anything for that and you can't really touch it it's super complex it's super performance critical and adding clays evasion is just impossible so that's why we end up with strict functions but originally even in design documents from 2015 people discussed this problem and people wanted to have laser functions already semantics because people at that time thought that if you have some else Leia in other than layer stem layer as a function you don't really want to wait for all inputs it is too critical for performance but that was too hard to implement at that time and we ended up with six functions but no one used functions at v1 it was a very rare thing and it was not a problem but in TIFF 2.0 with Devon presentation we wrap all the functions into function depths and now we have hundreds thousands of functions at runtime and now have all these problems with strict lazy semantics and console semantics is not really defined so different people think different things what should be the semantics and what is the right semantics so right now we kind of between strict and lazy and we can but sometimes it's still strict so it's a little bit of a mess so to get back our lazy semantics one way is to fix execution model runtime which is impossible without rewriting from scratch so we might do it later but this nonsensical can do so the easiest way to do that is just to inline all the functions into the graph so instead of graph def with a partition call to the function call you will get one single graph which is executable and function 1004 runtime can start executing notes as any impurity so you don't have any dysfunction boundary anymore but you still have to do it very carefully because if you inline it without because if you if you just inline the function body inside the graph you might get completely different semantics because some people rely on control dependencies that edge of the function all the strictness and all the side effects visibility into the 42.0 when it adds traumatic control planes is checking it has some assumptions about like what is the program execution order and if function aligning would violate that it would be very very sad situation so to the fourth point all Python content has some function semantics and graph construction semantics that helps a lot to get this program all the semantics so if the function or the mutable state is represented as resource tensors so their source is just a handle to device and some memory location so you can imagine this is a pointer on GPU or pointer to the buffer on CPU and we pass the sources as inputs to the functions and if the function has an input for versa a it will have an income control edge from the last op in graph or the last function that has with touch the chess awesome it has in the same resource as input so if you have assigned to variable and then you pass theirselves for the same variable to a function you will have a control dependency from that assign and if anyone else is touching the same resource after the function call or any other of TF 2.0 will add outgoing control edge to the next op so if we pass a variable to a function and then you outside your function you have a read variable up you have to point all will edit control dependencies from your function call to the treat so if you do the read you will observe all the changes that were made to that variable inside the function body so the most important thing I guess from T of the function notation that it does automatic control the prints tracking so when you write your python function and you have some idea what should be your program or the execution semantics so you add one to a then you add one to be you think that the earth to be should happen before after a and if you add one to a then you read a you expect that you see and then you value of the variable a it was not the case in TF one in tf2 when you wait to have the function notation it will add all the necessary control dependences so it will add control dependencies between all the herbs that have the same resource as input and it will also add control dependencies between all stage who ops inside function body so all your stateful operations will be always executed in the same order and it will give you some synergy so if you have multiple prints inside a function body you should every time observe that brings in the same order in tf2 point v1 you can observe that prints in any order and that's confusing and all the stage show ops basic aside effects ops that can have side effects they will be added to control the the controller output so function and functional lining must respect that so we don't want to lose any of the side effects so updates to the state and these are some rules so all the side effects to the resources that happened before function call must be visible to all the notes inside function call and all the side effects to the resources kept inside function body must be visible to every op function that using the same resource after a function completed currently it's implemented too so you have a control dependencies from the OP that made some side effect to the function call and the function call has a control dependencies to the next stop that might need that side effects so to enforce that semantics function aligning has a special rules so it will add a special note input control note so if your function call has any control dependencies that input control note will all that control dependencies will be forward to that input control node and all the function inputs will have a control dependencies from that node so it basically means that you can't start executing your function before your control dependencies satisfy it all the notes are executed it also means that if your function call doesn't have any control dependencies total for runtime is free to start running your function body as soon as anything is ready also it will add up control node so this node who have control edges from all their side effects from all the control outputs of a function so if for some reason you have a side-effect inside a function and this side effect is not connected to one of the control outputs when function will be aligned terms of flow is free to prune the side effect oh so you might end up with some partial observed state updates but that should not happen in practice in therefore 2.0 because all side effects when you can start a graph from bite and all the side effects should be properly connected to control outputs but that's not the case of scares models I think caris doesn't use what went automatic control the planes taking in some cases oh it was not using some time ago so that might be violated by Karras in some models and I hope that doesn't happen right now and also there is a sum from if the function call does not have an outgoing control edge it means that no one long doesn't care about what's happening inside the function what are the side effects so if you have an ongoing data edge someone needs from the data output but if function code doesn't have outgoing control edge it means that function might do whatever it wants inside function body updates any variables counters print anything send any data over the network and no one cares about what it does so that might lead to some troubles but I think in 2.0 that also never happens in practice because a traumatic dependency tracking for nested function calls will add all the required control the parentheses and when you have when it executes the top level function you should also add all the control dependences but again it might happen with some model so they do not use automatic control dependencies checking so that's how the function like function look like after landing so this is a function from previous example so the function takes three resources as inputs and reads variables a B increments the counter and does the Stratos slice one of which is invalid and returns both slices and control output is increment so a function call and this will be after aligning the graph def so we no longer have function deff and we know what we have function call we have just one graph with multiple notes so we have incoming input control note this is not all so the function call now do you have any kind of dependencies so no op has empty control dependencies we read variables from in and we depend on input control not so any reads from the variables from the inputs will happen after the input control note executed then we have an increment increment the counter then we do two studies slices and then we have two identical for function returns and we slice 0 for the first drawn value and slice one for the second and we have an output control node and output control node has control dependencies from the side-effects inside function and this function has only one side effect a sign a sign at counter increment so read variable hope marked stay flop but it is not a side effect because read variable op can't modify can have side effects you just observe the states so interior we could add variable up to this up control node but in practice there are many stateful ops that just reading the state and read variable up just one of them and slice so P so here the function call person cannot get slice and slice is identical the reads the second output now we don't have function corner anymore and slice is just an attention out the trees directly from the inlined function return anywhere so we read available from the counter and it has automatically control dependencies to the output control node so every time we read the counter we must observe all the side effects that happen to the counter inside function body and now we can do pruning again so we don't use return 0 we can prove it we don't use slide 0 that is invalid 22:13 we can prune it and we don't need the value of the variable so we can prove read the variable up and so again back to the graph of that is valid and can be executed at some time without exceptions so there are few more problems so when you have a function and you in line function body and function but it does not have any device annotations you have to decide on what device to place that note so before tier 2.0 we also had single device function so if the function kernel is placed on CPU all the notes inside function body would be executed on CPU in EF 2.0 it's to limit it because if you have ResNet in your function and that function call is on CPU you don't want to run your net on CPU or you might want to use multiple CPUs GPUs or have it function the transfer multiple devices so there are multiple strategies how to place notes of the length function so you can do nothing at all and rely on placer you can make force function to be single device and we do some time for you one graphs for computability mode primarily oh the right strategy for multi device functions in TM 2.0 that all the nodes inside function body must be placed on the same job replica and task but the device might be empty and then will lie on the place to place them on the correct device because imagine that you have a function call and stupid runtime and your function call happened to be on another machine and then when you would execute that function call on that machine at runtime it would be able to choose only from the devices available on this machine so if we were doing learning and don't add any device annotation place in my computer my sub device placements so if the user place a function call like ResNet function call on machine 1 and that doesn't function inside doesn't have any device annotation and Devine line it and then we run place huh all that ResNet knows my replaced on completely different machine in GPUs and it will break assumptions of the user so when we know why functions currently we override job tasks and replicas and then end on touch device so the placer will pick the right device for the note even after aligning also have a bunch of functions created internal v1 that don't use control outputs and don't use control the pins inside function body at all so after in line such functions you might end up with completely different execution semantics and what's of not multi-pronged that's another fun part of current runtime that current function runtime it does not prove any stay flops and it is very different from execution semantics of the graph because if you have stateful ops variable days inside your graph and you don't have control dependencies runtime will prove much but if you have exactly the same graph inside function runtime will execute all the stateful ops and this mismatch is very difficult to handle when you inline your function graph and sign on the outside gravid F because you have different notions of what should be pruned and when so you have to point all always in lines all the functions because if you have to point always guaranteed to have correct control dependences and function control outputs but we don't inline functions by default in legacy graphs grappler is responsible for in landing functions in legacy graphs but the wrapper does a lot of additional checks that all the stageful ops has a path to one of the output that we don't have any side effect for ops inside function body that are not connected to anything because they will be pruned and there are bunch of other also in TF one you can get function with mismatch deadness which should not be possible but that happens so the upper is very conservative so it in lines only if it can prove that it is safe and it does not change the semantics and this function ilanic is a heat source of bugs sometimes people think that function aligning to the problem and often it is as the problem and it because it is very complicated and semantics was never defined properly for the function so I mostly had to come up with semantics to make all the tests pass and a bunch of hockey workarounds for different tests I hope that we might be able to get better semantics at some point but at now it's super my son also have much of other function like functional Ops so for example we have a functional if so this is basically predicate and you have attributes with two functions then and else function and another of is functional while so in v1 we have switch next iteration and to exit nodes to represent wire loops in v2 we have a functional while so you define a function for the body for the condition and then you just run it at runtime as a special op also have a case something like if with multiple branches currently we lower all these functional ops 2 v1 control flow so basically the if node if node becomes a switch node and to function color node for them and else functions and these functions are then aligned does any other function call so we do that primarily for the stick semantics so if you run your while or if is the open side graph it will have strict semantics and some models just expect lazy semantics from you if it also were limiting for concurrency so we want control flow while loops for example you can run multiple parallel loop iterations at a time you can't tell where alterations if you do if you try to do that with the functional while it's just impossible you have to wait forever iteration to completely finish before you can start the next direction so a lot of people want to move to functional ops function control flow button practice it's very difficult primarily because of the performance what it makes a lot of analysis easy because we often have to reconstruct what was the wire loop from the graph graph from all the switches next iteration and tracks it notes and it is very back prominent have a lot of troubles with it so if you would have at the graphic immunization level functional ops that would help and then at the late stage we can just lower all of them to switching measures to get the good performance so here's example of how functional F looks like in a graph so we have a flag just a boolean variable we read the flag with the read variable up we have some constant zero so we have two functions plus one that adds to interest in put one in the transept immediately and plus two so we read the flag boolean flag then we have zero and result is if plus one plus two so if the flag is true we add plus one if the flag is false we add plus 2 to 0 so that will be 1 or 2 depending on the flag so when we lower this functional if to be one control flow constructs we will have the level up for the flag which will have the 0 for the cost so we have a switch node based on the flag and 0 so the switch node has two outputs it will output the value 0 on the output 0 if the flag is true and you output the value 0 on the output 1 if the flag is false and allow a new start will be dead so to basically prevent execution from one of these nodes so then function is as function call with a switch out to 0 and else function is another function call with the switch node output 1 so either flag is false this node will be dead and will not be executed and then we have a result as a merge so we merging results of their function and else function and one of them is going to be dead and another one will be alive and this will be the final result so after we lower so this is after we lower if you note to function calls and then function learning kicks in and we get rid of the function call notes in the basic you have then function return just add from to each one and else functions at this should be 2 and we merging the return values all the functions yeah basically it so that's how we get rid of all the functions so we have functions as a mental model for the end user how to you think about your total flow graph how you built a program so you no longer think in terms of graph and support notes you think in terms of functions but when we get these functions in the runtime we still can order them to the soup of notes because that's what we have to do for performance and sometimes for correctness because there is a kind of a promise of TF but function notation into the floor eager mode if you want your view one semantics back just annotate your functional with DF that function and you'll get back a graph but that's not completely true because if you have nested function calls annotated with stab the function you would have multiple function there's multiple function call notes and you'll get strict semantics and this not to what was the semantics of you want so the only way to provide users what's promised is to inline all the functions and then we'll get back to the single graph with all a semantics with the pruning and all the Knights Pro Prius of page fault graph [Music] you