Browse Source

infer processor from ms_context inside function CreateCNode

pull/14237/head
looop5 5 years ago
parent
commit
dd43972839
2 changed files with 16 additions and 1 deletions
  1. +15
    -1
      mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_helper.cc
  2. +1
    -0
      mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_helper.h

+ 15
- 1
mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_helper.cc View File

@@ -33,6 +33,7 @@
#include "pipeline/jit/parse/python_adapter.h"
#include "pipeline/jit/action.h"
#include "vm/segment_runner.h"
#include "utils/ms_context.h"
#if ENABLE_GPU
#include "runtime/device/gpu/kernel_info_setter.h"
#endif
@@ -796,6 +797,19 @@ std::vector<int64_t> GetReduceAxis(const AnfNodePtr &node) {
return axis;
}

kernel::Processor GetProcessorFromContext() {
kernel::Processor processor = kernel::Processor::UNKNOWN;
auto context_ptr = MsContext::GetInstance();
MS_EXCEPTION_IF_NULL(context_ptr);
auto device_info = context_ptr->get_param<std::string>(MS_CTX_DEVICE_TARGET);
if (device_info == kGPUDevice) {
processor = kernel::Processor::CUDA;
} else if (device_info == kAscendDevice) {
processor = kernel::Processor::AICORE;
}
return processor;
}

CNodePtr CreateCNode(const std::vector<AnfNodePtr> &inputs, const FuncGraphPtr &func_graph, const DataInfo &out_info) {
// Limitation: 1. Node's attributes should be set out of this function; 2. only one output.
MS_EXCEPTION_IF_NULL(out_info.type);
@@ -852,7 +866,7 @@ CNodePtr CreateCNode(const std::vector<AnfNodePtr> &inputs, const FuncGraphPtr &
info_builder.SetInputsDeviceType(input_types);
info_builder.SetOutputsFormat(output_formats);
info_builder.SetOutputsDeviceType(output_types);
info_builder.SetProcessor(kernel::Processor::CUDA);
info_builder.SetProcessor(GetProcessorFromContext());
info_builder.SetKernelType(KernelType::AKG_KERNEL);
info_builder.SetFusionType(kernel::FusionType::OPAQUE);
auto selected_info = info_builder.Build();


+ 1
- 0
mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_helper.h View File

@@ -87,6 +87,7 @@ TypePtr GetType(const AnfNodePtr &node);
ShapeVector GetShape(const AnfNodePtr &node);
ShapeVector GetDeviceShape(const AnfNodePtr &node);
std::vector<int64_t> GetReduceAxis(const AnfNodePtr &node);
kernel::Processor GetProcessorFromContext();

CNodePtr CreateCNode(const std::vector<AnfNodePtr> &inputs, const FuncGraphPtr &func_graph, const DataInfo &out_info);
void SetNodeAttrSafely(const std::string &key, const ValuePtr &value, const AnfNodePtr &node);


Loading…
Cancel
Save