torch/csrc/jit/passes/specialize_undef.cpp
// propagate undefined information through a gradient graph and
// remove grad_of blocks if present.
// Note: this is a very limited pass. It only propagates undefines for
// operations generated by the symbolic autodiff code and cleans up
// AutogradAdds when possible. Outputs of other nodes are conservatively
// marked Unknown and not optimized.
propagate undefined information through a gradient graph and
remove grad_of blocks if present.
Note: this is a very limited pass. It only propagates undefines for
operations generated by the symbolic autodiff code and cleans up
AutogradAdds when possible. Outputs of other nodes are conservatively
marked Unknown and not optimized.
propagate undefined information through a gradient graph and remove grad_of blocks if present . Note : this is a very limited pass . It only propagates undefines for operations generated by the symbolic autodiff code and cleans up AutogradAdds when possible . Outputs of other nodes are conservatively marked Unknown and not optimized .
void specializeUndef(Graph & g, const std::vector<bool>& defined) {
  enum class State { Defined, Undefined, Unknown };
  std::unordered_map<Value*, State> state;
  for(size_t i = 0; i < defined.size(); i++) {
    state[g.inputs()[i]] = defined.at(i) ? State::Defined : State::Undefined;
  }
  for(auto it = g.nodes().begin(); it != g.nodes().end(); ++it) {
    auto n = *it;
    switch(n->kind()) {
      case prim::GradOf: {
        auto all_undefined =
            std::all_of(n->inputs().begin(), n->inputs().end(), [&](Value* v) {
              return state[v] == State::Undefined;
            });
        if(all_undefined) {
          auto undef = g.createUndefined()->insertAfter(n)->output();
          for(auto o : n->outputs()) {
            o->replaceAllUsesWith(undef);
          }
        } else {
          auto body = n->blocks().at(0);
          for(auto input : n->inputs()){
            JIT_ASSERT(state[input] != State::Unknown);
          }
          for(auto it = body->nodes().begin(); it != body->nodes().end();) {
            auto block_node = *it++;
            block_node->moveBefore(n);
          }
          for(size_t i = 0; i < n->outputs().size(); ++i)
            n->outputs().at(i)->replaceAllUsesWith(body->outputs().at(i));
        }
        it.destroyCurrent();
      } break;
      case prim::AutogradAdd: {
        auto a = n->input(0);
        auto b = n->input(1);
        if(state[a] == State::Undefined) {
          n->output()->replaceAllUsesWith(b);
          it.destroyCurrent();
        } else if(state[b] == State::Undefined) {
          n->output()->replaceAllUsesWith(a);
          it.destroyCurrent();
        } else if(state[a] == State::Defined && state[b] == State::Defined) {
          WithInsertPoint guard(n);
          Value* new_add = toVar(a) + toVar(b);
          state[new_add] = State::Defined;
          n->output()->replaceAllUsesWith(new_add);
          it.destroyCurrent();
        } else {
          state[n->output()] = State::Unknown;
        }
      } break;
      case prim::Undefined: {
        state[n->output()] = State::Undefined;
      } break;
      default:
        for(auto o : n->outputs()) {
          state[o] = State::Unknown;
        }
        break;
    }
  }
}
void specializeUndef ( Graph & g , const std :: vector < bool > & defined ) { enum class State { Defined , Undefined , Unknown } ; std :: unordered_map < Value * , State > state ; for ( size_t i = 0 ; i < defined . size ( ) ; i ++ ) { state [ g . inputs ( ) [ i ] ] = defined . at ( i ) ? State :: Defined : State :: Undefined ; } for ( auto it = g . nodes ( ) . begin ( ) ; it != g . nodes ( ) . end ( ) ; ++ it ) { auto n = * it ; switch ( n -> kind ( ) ) { case prim :: GradOf : { auto all_undefined = std :: all_of ( n -> inputs ( ) . begin ( ) , n -> inputs ( ) . end ( ) , [ & ] ( Value * v ) { return state [ v ] == State :: Undefined ; } ) ; if ( all_undefined ) { auto undef = g . createUndefined ( ) -> insertAfter ( n ) -> output ( ) ; for ( auto o : n -> outputs ( ) ) { o -> replaceAllUsesWith ( undef ) ; } } else { auto body = n -> blocks ( ) . at ( 0 ) ; for ( auto input : n -> inputs ( ) ) { JIT_ASSERT ( state [ input ] != State :: Unknown ) ; } for ( auto it = body -> nodes ( ) . begin ( ) ; it != body -> nodes ( ) . end ( ) ; ) { auto block_node = * it ++ ; block_node -> moveBefore ( n ) ; } for ( size_t i = 0 ; i < n -> outputs ( ) . size ( ) ; ++ i ) n -> outputs ( ) . at ( i ) -> replaceAllUsesWith ( body -> outputs ( ) . at ( i ) ) ; } it . destroyCurrent ( ) ; } break ; case prim :: AutogradAdd : { auto a = n -> input ( 0 ) ; auto b = n -> input ( 1 ) ; if ( state [ a ] == State :: Undefined ) { n -> output ( ) -> replaceAllUsesWith ( b ) ; it . destroyCurrent ( ) ; } else if ( state [ b ] == State :: Undefined ) { n -> output ( ) -> replaceAllUsesWith ( a ) ; it . destroyCurrent ( ) ; } else if ( state [ a ] == State :: Defined && state [ b ] == State :: Defined ) { WithInsertPoint guard ( n ) ; Value * new_add = toVar ( a ) + toVar ( b ) ; state [ new_add ] = State :: Defined ; n -> output ( ) -> replaceAllUsesWith ( new_add ) ; it . destroyCurrent ( ) ; } else { state [ n -> output ( ) ] = State :: Unknown ; } } break ; case prim :: Undefined : { state [ n -> output ( ) ] = State :: Undefined ; } break ; default : for ( auto o : n -> outputs ( ) ) { state [ o ] = State :: Unknown ; } break ; } } }
propagate undefined information through a gradient graph and
remove grad_of blocks if present.
propagate undefined information through a gradient graph and remove grad_of blocks if present .
// Property 1: if all the gradInputs to the GradOf are undefined // then the gradOutputs are also zero and will be represented as undefined nodes // Property 2: GradOfs are required to correctly handle combinations // of defined and undefined inputs. They are expected to produce defined // output tensors in this case. // Remove the GradOf, splicing its body back into the surrounding block // we should never get into a situation when specializing a GradOf // where we do not know if a value is defined since at the top level // a gradient graph is composed of Linear nodes and AutogradAdds // and LinearNodes only appear in these graphs // hoist the nodes in the GradOf body to be before the linear block // if one is undefined, we can just drop the add // Undef + b == b // a + Undef == a // when both are defined, we can use a normal, optimizable add instruction // otherwise we have conditionally-defined things, and we need // to actually run an AutogradAdd which will guard for undefs // so we leave the op as is