torch/csrc/jit/passes/to_batch.cpp
// replace aten operator node with BatchTensor operator graph
replace aten operator node with BatchTensor operator graph
replace aten operator node with BatchTensor operator graph
void ToBatch::visitAten(Node* n, Block* block, Block* res_block){
  auto res_graph = res_block->owningGraph();
  auto func_name = std::string(n->kind().toUnqualString());
  std::vector<Value*> new_inputs;
  for(Value *input : n->inputs()){
    if(rn_env.find(input) == rn_env.end()){  
      auto new_input = batch_map.at(input);
      new_inputs.insert(new_inputs.end(), new_input.begin(), new_input.end());
    }
    else{  
      new_inputs.push_back(rn_env.at(input));
    }
  }
  for(size_t i = 0; i < new_inputs.size(); i++){
    auto input = new_inputs[i];
    if(input->type() == IntType::get() || input->type() == FloatType::get()){
      auto to_tensor_node = res_graph->createNumToTensor(input);
      res_graph->insertNode(to_tensor_node);
      new_inputs[i] = to_tensor_node->output();
    }
  }
  auto batch_graph = getBatchOperator(func_name, new_inputs.size());
  auto outputs = script::inlineCallTo(*res_block->owningGraph(), *batch_graph, new_inputs);
  if(outputs.size() == 1){
    if(n->outputs()[0]->type() != outputs[0]->type()){
      Node* to_scalar_node;
      if(n->outputs()[0]->type() == IntType::get()){
        to_scalar_node = res_graph->createTensorToNum(IntType::get(), outputs[0]);
      }
      else if(n->outputs()[0]->type() == FloatType::get()){
        to_scalar_node = res_graph->createTensorToNum(FloatType::get(), outputs[0]);
      }
      else{
        throw std::runtime_error("NYI: scalar type other than int, float is not supported yet");
      }
      res_graph->insertNode(to_scalar_node);
      rn_env[n->outputs()[0]] = to_scalar_node->output();
    }
    else
      rn_env[n->outputs()[0]] = outputs[0];
  }
  else{
    for(size_t i = 0; i < n->outputs().size(); i++){
      auto output = n->outputs()[i];
      batch_map[output] = std::vector<Value*>(outputs.begin() + i * EXP_BTENSOR_SIZE, outputs.begin() + i * EXP_BTENSOR_SIZE + EXP_BTENSOR_SIZE);
    }
  }
}
void ToBatch :: visitAten ( Node * n , Block * block , Block * res_block ) { auto res_graph = res_block -> owningGraph ( ) ; auto func_name = std :: string ( n -> kind ( ) . toUnqualString ( ) ) ; std :: vector < Value * > new_inputs ; for ( Value * input : n -> inputs ( ) ) { if ( rn_env . find ( input ) == rn_env . end ( ) ) { auto new_input = batch_map . at ( input ) ; new_inputs . insert ( new_inputs . end ( ) , new_input . begin ( ) , new_input . end ( ) ) ; } else { new_inputs . push_back ( rn_env . at ( input ) ) ; } } for ( size_t i = 0 ; i < new_inputs . size ( ) ; i ++ ) { auto input = new_inputs [ i ] ; if ( input -> type ( ) == IntType :: get ( ) || input -> type ( ) == FloatType :: get ( ) ) { auto to_tensor_node = res_graph -> createNumToTensor ( input ) ; res_graph -> insertNode ( to_tensor_node ) ; new_inputs [ i ] = to_tensor_node -> output ( ) ; } } auto batch_graph = getBatchOperator ( func_name , new_inputs . size ( ) ) ; auto outputs = script :: inlineCallTo ( * res_block -> owningGraph ( ) , * batch_graph , new_inputs ) ; if ( outputs . size ( ) == 1 ) { if ( n -> outputs ( ) [ 0 ] -> type ( ) != outputs [ 0 ] -> type ( ) ) { Node * to_scalar_node ; if ( n -> outputs ( ) [ 0 ] -> type ( ) == IntType :: get ( ) ) { to_scalar_node = res_graph -> createTensorToNum ( IntType :: get ( ) , outputs [ 0 ] ) ; } else if ( n -> outputs ( ) [ 0 ] -> type ( ) == FloatType :: get ( ) ) { to_scalar_node = res_graph -> createTensorToNum ( FloatType :: get ( ) , outputs [ 0 ] ) ; } else { throw std :: runtime_error ( " " ) ; } res_graph -> insertNode ( to_scalar_node ) ; rn_env [ n -> outputs ( ) [ 0 ] ] = to_scalar_node -> output ( ) ; } else rn_env [ n -> outputs ( ) [ 0 ] ] = outputs [ 0 ] ; } else { for ( size_t i = 0 ; i < n -> outputs ( ) . size ( ) ; i ++ ) { auto output = n -> outputs ( ) [ i ] ; batch_map [ output ] = std :: vector < Value * > ( outputs . begin ( ) + i * EXP_BTENSOR_SIZE , outputs . begin ( ) + i * EXP_BTENSOR_SIZE + EXP_BTENSOR_SIZE ) ; } } }
replace aten operator node with BatchTensor operator graph
replace aten operator node with BatchTensor operator graph
// non-tensor input // batched tensor input // transform scalar to tensor before pass to batch operator script // Assume all outputs from inlined operator implementation are in the triple form batched tensor or just a single non-tensor. // if previous output is scalar, transform new output back to scalar from dynamic