deal with workspace

deal with workspace on cuda

mindspore/_extends/graph_kernel/parallel_estimate.py

@@ -21,7 +21,7 @@ from . import model
 
 
 def estimate_ops(json_str: str):
-    """Call costmodel to estimate ops."""
+    """Call cost model to estimate ops."""
     try:
         json_obj = json.loads(json_str)
         graph_descs = json_obj["graph_desc"]
@@ -38,7 +38,7 @@ def estimate_ops(json_str: str):
 
 
 def estimate_calulation_amount(json_str: str):
-    """Call costmodel to estimate calculation amount of op."""
+    """Call cost model to estimate calculation amount of op."""
     try:
         graph_desc = json.loads(json_str)
         comp = model.load_composite(graph_desc)

mindspore/_extends/graph_kernel/splitter.py

@@ -24,7 +24,7 @@ from . import utils
 
 
 def split_with_json(json_str, flags_str):
-    """Call costmodel to split GraphKernel"""
+    """Call cost model to split GraphKernel"""
     try:
         graph_desc = json.loads(json_str)
         flags = json.loads(flags_str)

mindspore/ccsrc/backend/kernel_compiler/akg/ascend/akg_ascend_kernel_build.cc

@@ -44,8 +44,10 @@ KernelPackPtr AkgAscendKernelBuilder::AkgInsertCache(const std::string &kernel_n
 void AkgAscendKernelBuilder::AkgSetKernelMod(const KernelPackPtr &kernel_pack,
                                              const AkgKernelJsonGenerator &json_generator, const AnfNodePtr &anf_node) {
   auto kernel_mod_ptr = std::make_shared<AkgKernelMod>(kernel_pack);
+  auto kernel_json_info = kernel_pack->kernel_json_info();
   kernel_mod_ptr->SetInputSizeList(json_generator.input_size_list());
   kernel_mod_ptr->SetOutputSizeList(json_generator.output_size_list());
+  kernel_mod_ptr->SetWorkspaceSizeList(kernel_json_info.workspaces);
   AnfAlgo::SetKernelMod(kernel_mod_ptr, anf_node.get());
 }
 

mindspore/ccsrc/backend/kernel_compiler/akg/ascend/akg_ascend_kernel_mod.cc

@@ -49,7 +49,7 @@ const std::vector<size_t> &AkgKernelMod::GetOutputSizeList() const { return outp
 
 const std::vector<size_t> &AkgKernelMod::GetWorkspaceSizeList() const { return workspace_size_list_; }
 
-bool AkgKernelMod::Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &,
+bool AkgKernelMod::Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
                           const std::vector<AddressPtr> &outputs, void *stream_ptr) {
   if (stream_ptr == nullptr) {
     MS_LOG(ERROR) << "stream_ptr should not be nullptr.";
@@ -74,6 +74,10 @@ bool AkgKernelMod::Launch(const std::vector<AddressPtr> &inputs, const std::vect
                        [](const AddressPtr &input) -> void * { return input->addr; });
   (void)std::transform(std::begin(outputs), std::end(outputs), std::back_inserter(runtime_args),
                        [](const AddressPtr &output) -> void * { return output->addr; });
+  if (!workspace.empty()) {
+    (void)std::transform(std::begin(workspace), std::end(workspace), std::back_inserter(runtime_args),
+                         [](const AddressPtr &addr) -> void * { return addr->addr; });
+  }
 
   rtL2Ctrl_t *l2ctrl = nullptr;
   auto stream = static_cast<rtStream_t *>(stream_ptr);
@@ -86,7 +90,8 @@ bool AkgKernelMod::Launch(const std::vector<AddressPtr> &inputs, const std::vect
   return true;
 }
 
-std::vector<TaskInfoPtr> AkgKernelMod::GenTask(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &,
+std::vector<TaskInfoPtr> AkgKernelMod::GenTask(const std::vector<AddressPtr> &inputs,
+                                               const std::vector<AddressPtr> &workspace,
                                                const std::vector<AddressPtr> &outputs, uint32_t stream_id) {
   if (kernel_pack_ == nullptr) {
     MS_LOG(EXCEPTION) << "kernel pack should not be nullptr.";
@@ -107,6 +112,10 @@ std::vector<TaskInfoPtr> AkgKernelMod::GenTask(const std::vector<AddressPtr> &in
                        [](const AddressPtr &input) -> void * { return input->addr; });
   (void)std::transform(std::begin(outputs), std::end(outputs), std::back_inserter(output_data_addrs),
                        [](const AddressPtr &output) -> void * { return output->addr; });
+  if (!workspace.empty()) {
+    (void)std::transform(std::begin(workspace), std::end(workspace), std::back_inserter(workspace_addrs),
+                         [](const AddressPtr &workspace) -> void * { return workspace->addr; });
+  }
 
   uint32_t block_dim = DEFAULT_BLOCK_DIM;  // default blockdim equal to 1.
   auto func_stub = KernelManager::GenFuncStub(*kernel_pack_, false, &block_dim);

mindspore/ccsrc/backend/kernel_compiler/akg/gpu/akg_gpu_kernel_build.cc

@@ -39,8 +39,10 @@ KernelPackPtr AkgGpuKernelBuilder::AkgInsertCache(const std::string &kernel_name
 void AkgGpuKernelBuilder::AkgSetKernelMod(const KernelPackPtr &kernel_pack,
                                           const AkgKernelJsonGenerator &json_generator, const AnfNodePtr &anf_node) {
   auto kernel_mod_ptr = std::make_shared<GpuKernelMod>(kernel_pack);
+  auto kernel_json_info = kernel_pack->kernel_json_info();
   kernel_mod_ptr->SetInputSizeList(json_generator.input_size_list());
   kernel_mod_ptr->SetOutputSizeList(json_generator.output_size_list());
+  kernel_mod_ptr->SetWorkspaceSizeList(kernel_json_info.workspaces);
   AnfAlgo::SetKernelMod(kernel_mod_ptr, anf_node.get());
 }
 

mindspore/ccsrc/backend/kernel_compiler/akg/gpu/akg_gpu_kernel_mod.cc

@@ -92,13 +92,15 @@ void GpuKernelMod::SetInputSizeList(const std::vector<size_t> &size_list) { inpu
 
 void GpuKernelMod::SetOutputSizeList(const std::vector<size_t> &size_list) { output_size_list_ = size_list; }
 
+void GpuKernelMod::SetWorkspaceSizeList(const std::vector<size_t> &size_list) { workspace_size_list_ = size_list; }
+
 const std::vector<size_t> &GpuKernelMod::GetInputSizeList() const { return input_size_list_; }
 
 const std::vector<size_t> &GpuKernelMod::GetOutputSizeList() const { return output_size_list_; }
 
 const std::vector<size_t> &GpuKernelMod::GetWorkspaceSizeList() const { return workspace_size_list_; }
 
-bool GpuKernelMod::Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &,
+bool GpuKernelMod::Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
                           const std::vector<AddressPtr> &outputs, void *stream_ptr) {
   if (stream_ptr == 0) {
     MS_LOG(ERROR) << "stream_ptr should not be nullptr.";
@@ -122,6 +124,10 @@ bool GpuKernelMod::Launch(const std::vector<AddressPtr> &inputs, const std::vect
                        [](const AddressPtr &input) -> void * { return reinterpret_cast<void *>(&(input->addr)); });
   (void)std::transform(std::begin(outputs), std::end(outputs), std::back_inserter(runtimeargs),
                        [](const AddressPtr &output) -> void * { return reinterpret_cast<void *>(&(output->addr)); });
+  if (!workspace.empty()) {
+    (void)std::transform(std::begin(workspace), std::end(workspace), std::back_inserter(runtimeargs),
+                         [](const AddressPtr &addr) -> void * { return addr->addr; });
+  }
   result = cuLaunchKernel(kernel_addr, thread_info[0], thread_info[1], thread_info[2], thread_info[3], thread_info[4],
                           thread_info[5], 0, reinterpret_cast<CUstream>(stream_ptr),
                           reinterpret_cast<void **>(&runtimeargs[0]), 0);

mindspore/ccsrc/backend/kernel_compiler/akg/gpu/akg_gpu_kernel_mod.h

@@ -60,6 +60,7 @@ class GpuKernelMod : public KernelMod {
 
   void SetInputSizeList(const std::vector<size_t> &size_list);
   void SetOutputSizeList(const std::vector<size_t> &size_list);
+  void SetWorkspaceSizeList(const std::vector<size_t> &size_list);
   const std::vector<size_t> &GetInputSizeList() const override;
   const std::vector<size_t> &GetOutputSizeList() const override;
   const std::vector<size_t> &GetWorkspaceSizeList() const override;

mindspore/ccsrc/backend/kernel_compiler/kash/kernel_pack.cc

@@ -118,6 +118,16 @@ bool KernelPack::ReadFromJsonFile(const std::string &json_f, const std::string &
     if (!CheckHash(json_f, bin_f, js)) {
       return false;
     }
+
+    // cuda json file may have workspace information
+    if (js.find("workspace") != js.end()) {
+      auto workspace = js.at("workspace");
+      std::vector<size_t> sizes = workspace.at("size");
+      for (auto size : sizes) {
+        kernel_json_info_.workspaces.push_back(size);
+      }
+    }
+
     return true;
   }