Program Listing for File computation.cc
↰ Return to documentation for file (/WorkSpace/CINN/cinn/frontend/computation.cc)
// Copyright (c) 2021 CINN Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "cinn/frontend/computation.h"
#include "cinn/frontend/program_pass.h"
#include "cinn/hlir/framework/graph.h"
#include "cinn/hlir/framework/graph_compiler.h"
#include "cinn/hlir/framework/pass.h"
#include "cinn/hlir/framework/scope.h"
namespace cinn {
namespace frontend {
struct ComputationContext {
Target target;
void *stream;
std::shared_ptr<hlir::framework::Graph> graph;
std::shared_ptr<hlir::framework::Scope> scope;
std::shared_ptr<hlir::framework::Program> program;
std::shared_ptr<hlir::framework::GraphCompiler> graph_compiler;
CinnComputation::CompileOptions compile_options;
std::vector<hlir::framework::Tensor> inputs;
std::vector<hlir::framework::Tensor> outputs;
std::unordered_map<std::string, Variable> varmap;
std::unordered_map<std::string, std::string> varmap_paddle2program;
};
std::shared_ptr<ComputationContext> CompileProgram(const Target &target,
Program &program,
const std::vector<Variable> &outputs,
std::shared_ptr<hlir::framework::Scope> scope,
const CinnComputation::CompileOptions &options,
void *stream) {
std::shared_ptr<ComputationContext> ctx(new ComputationContext());
ctx->stream = stream;
ctx->target = target;
ctx->compile_options = options;
if (ctx->compile_options.use_decomposer) {
ProgramPass::Apply(&program, target, {"Decomposer"});
}
ctx->graph.reset(new hlir::framework::Graph(program, target));
if (ctx->compile_options.use_default_passes) {
hlir::framework::ApplyPass(ctx->graph.get(), "InferShape");
#ifndef CINN_WITH_CUDA
if (target.arch == Target::Arch::X86) {
hlir::framework::ApplyPass(ctx->graph.get(), "AlterLayout");
}
#endif
hlir::framework::ApplyPass(ctx->graph.get(), "ConstPropagate");
hlir::framework::ApplyPass(ctx->graph.get(), "OpFusion");
}
for (auto &pass_name : ctx->compile_options.passes) {
hlir::framework::ApplyPass(ctx->graph.get(), pass_name);
}
ctx->scope = hlir::framework::BuildScope(target, ctx->graph, scope);
ctx->graph_compiler.reset(new hlir::framework::GraphCompiler(target, ctx->scope, ctx->graph));
std::unordered_set<std::string> fetch_var_ids;
for (auto &out : outputs) {
fetch_var_ids.insert(out->id);
}
ctx->program = ctx->graph_compiler->Build(options, std::move(fetch_var_ids)).runtime_program;
if (ctx->compile_options.do_prerun) {
ctx->program->PreRun();
}
for (auto &in_v : program.GetInputs()) {
hlir::framework::Tensor t = ctx->scope->GetTensor(in_v->id);
ctx->inputs.push_back(t);
}
for (auto &out_v : outputs) {
hlir::framework::Tensor t = ctx->scope->GetTensor(out_v->id);
ctx->outputs.push_back(t);
}
return ctx;
}
std::vector<std::string> CinnComputation::GetAllTensorNames() {
std::vector<std::string> res;
for (auto &v : context_->scope->var_names()) {
res.push_back(std::string(v));
}
return res;
}
std::shared_ptr<CinnComputation> CinnComputation::CompilePaddleModel(
const Target &target,
const std::string &model_path,
const std::vector<std::string> &input_names,
const std::vector<hlir::framework::shape_t> &input_shapes,
bool params_combined,
const CompileOptions &options,
void *stream) {
CHECK(input_names.size() == input_shapes.size());
auto scope = std::make_shared<hlir::framework::Scope>();
auto loadedProgram = LoadPaddleProgram(model_path, scope.get(), params_combined, target);
auto &program = std::get<0>(loadedProgram);
auto &varmap = std::get<1>(loadedProgram);
auto &varmap_paddle2program = std::get<2>(loadedProgram);
auto &fetch_names = std::get<3>(loadedProgram);
std::vector<Variable> input_vars;
std::vector<Variable> output_vars;
for (int i = 0; i < input_names.size(); i++) {
auto &name = input_names[i];
auto &var = varmap.at(name);
var->shape = input_shapes[i];
input_vars.push_back(var);
}
program->SetInputs({input_vars});
program->Validate();
VLOG(3) << "program:\n" << *program;
for (auto &name : fetch_names) {
output_vars.push_back(varmap.at(name));
}
std::shared_ptr<ComputationContext> ctx = CompileProgram(target, *program, output_vars, scope, options, stream);
for (auto &v : varmap) {
ctx->varmap[v.first] = v.second;
}
for (auto &v : varmap_paddle2program) {
ctx->varmap_paddle2program[v.first] = v.second;
}
auto computation = std::make_shared<CinnComputation>();
computation->context_ = std::move(ctx);
return computation;
}
std::shared_ptr<CinnComputation> CinnComputation::BuildAndCompile(const Target &target,
BaseBuilder &builder,
const CompileOptions &options,
const std::vector<Variable> &outputs,
void *stream) {
auto program = builder.Build();
return Compile(target, program, options, outputs, stream);
}
std::shared_ptr<CinnComputation> CinnComputation::Compile(const Target &target,
Program &program,
const CompileOptions &options,
const std::vector<Variable> &outputs,
void *stream) {
std::vector<Variable> output_vars = outputs;
if (output_vars.empty()) {
output_vars.push_back(program[program.size() - 1].GetOutput(0));
}
std::shared_ptr<ComputationContext> ctx = CompileProgram(target, program, output_vars, nullptr, options, stream);
auto computation = std::make_shared<CinnComputation>();
computation->context_ = std::move(ctx);
return computation;
}
void CinnComputation::SetTensorData(const std::string &tname, void *data, size_t size) {
hlir::framework::Tensor t = GetTensor(tname);
SetTensorData(t, data, size);
}
void CinnComputation::SetTensorData(hlir::framework::Tensor &t, void *data, size_t size) {
void *tdata = reinterpret_cast<void *>(t->mutable_data<float>(context_->target));
CHECK_EQ(size, t->shape().numel() * sizeof(float));
if (context_->target.arch == Target::Arch::NVGPU) {
#ifdef CINN_WITH_CUDA
CUDA_CALL(cudaMemcpy(tdata, data, size, cudaMemcpyHostToDevice));
#else
CINN_NOT_IMPLEMENTED
#endif
} else if (context_->target.arch == Target::Arch::X86) {
memcpy(tdata, data, size);
} else {
CINN_NOT_IMPLEMENTED
}
}
void CinnComputation::GetTensorData(hlir::framework::Tensor &t, void *data, size_t size) {
void *tdata = reinterpret_cast<void *>(t->mutable_data<float>(context_->target));
CHECK_EQ(size, t->shape().numel() * sizeof(float));
if (context_->target.arch == Target::Arch::NVGPU) {
#ifdef CINN_WITH_CUDA
CUDA_CALL(cudaMemcpy(data, tdata, size, cudaMemcpyDeviceToHost));
#else
CINN_NOT_IMPLEMENTED
#endif
} else if (context_->target.arch == Target::Arch::X86) {
memcpy(data, tdata, size);
} else {
CINN_NOT_IMPLEMENTED
}
}
void CinnComputation::GetTensorData(const std::string &tname, void *data, size_t size) {
hlir::framework::Tensor t = GetTensor(tname);
GetTensorData(t, data, size);
}
std::vector<hlir::framework::Tensor> CinnComputation::GetInputTensors() { return context_->inputs; }
std::vector<hlir::framework::Tensor> CinnComputation::GetOutputTensors() { return context_->outputs; }
hlir::framework::Tensor CinnComputation::GetTensor(const std::string &tname) {
if (context_->scope->FindVar(tname)) {
return context_->scope->GetTensor(tname);
}
auto it = context_->varmap_paddle2program.find(tname);
if (it == context_->varmap_paddle2program.end()) {
LOG(FATAL) << "No variable called [" << tname
<< "] found in computation\nThe existing vars: " << utils::Join(context_->scope->var_names(), ", ");
}
return context_->scope->GetTensor(it->second);
}
void CinnComputation::Execute(const std::map<std::string, cinn_pod_value_t> *name2podargs) {
context_->program->Execute(name2podargs, context_->stream);
}
} // namespace frontend
} // namespace cinn