Method __trace

mllm/nn/Module.cpp:144–202

}

std::vector<Tensor> Module::__trace(const std::vector<Tensor>& inputs, const std::vector<AnyValue>& args) {
  auto ir_ctx = Context::instance().thisThread()->ir_context;

  // Generate unique name for this module instance to avoid conflicts with same-name modules
  std::string unique_module_name = ir_ctx->getUniqueModuleName(impl_->getAbsoluteName());
  // rename the module impl to unique name
  impl_->setAbsoluteName(unique_module_name);

  // Create call graph.
  auto call_op = ir_ctx->create<ir::graph::CallGraphOp>(ir_ctx->create<ir::SymbolAttr>(unique_module_name));

  // Create subgraph under ModuleOp
  ir::graph::SubGraphOp::ptr_t this_graph_ir = nullptr;
  {
    auto guard = ir::IRWriterGuard(ir_ctx, ir_ctx->topLevelOp()->cast_<ir::ModuleOp>()->getTopRegion());
    this_graph_ir = ir_ctx->create<ir::graph::SubGraphOp>(ir_ctx->create<ir::SymbolAttr>(unique_module_name), impl_);
  }
  this_graph_ir->setDevice(impl_->getDevice());

  // Wrap the inputs to tensor ir.
  std::vector<ir::tensor::TensorValue::ptr_t> inputs_ir = ir::tensor::wrapTensors2TensorIR(ir_ctx.get(), inputs);

  // Link inputs to subgraph.
  for (const auto& input_ir : inputs_ir) {
    (*input_ir)-- > this_graph_ir;
    (*input_ir)-- > call_op;
    this_graph_ir->getTopRegion()->inputs().push_back(input_ir);
  }

  // Forward
  std::vector<Tensor> outputs;
  {
    ir_ctx->setDevice(impl_->getDevice());
    auto guard = ir::IRWriterGuard(ir_ctx, this_graph_ir->getTopRegion());
    outputs = forward(inputs, args);
  }

  // wrap the outputs to tensor ir.
  std::vector<std::shared_ptr<ir::tensor::TensorValue>> outputs_ir = ir::tensor::wrapTensors2TensorIR(ir_ctx.get(), outputs);

  // link outputs to subgraph.
  for (const auto& output_ir : outputs_ir) {
    (*this_graph_ir)-- > output_ir;
    (*call_op)-- > output_ir;
    this_graph_ir->getTopRegion()->outputs().push_back(output_ir);
  }

  // create return op
  {
    auto guard = ir::IRWriterGuard(ir_ctx, this_graph_ir->getTopRegion());
    std::vector<ir::val_ptr_t> vals;
    vals.reserve(outputs_ir.size());
    for (auto& o : outputs_ir) vals.push_back(o);
    ir_ctx->create<ir::cf::ReturnOp>(vals);
  }

  // return the outputs.
  return outputs;