[METAL] Split kernels and compile them separately#7980

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tqchen merged 5 commits into from
May 25, 2021
Merged

[METAL] Split kernels and compile them separately #7980

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
b0b51fd
[METAL] Split kernels and compile them separately
echuraev May 5, 2021
4938d70
Add backward compatible behaviour
echuraev May 6, 2021
070bba4
Move SplitKernels to a common file
echuraev May 12, 2021
8041a29
Remove using GetSource
echuraev May 12, 2021
df55c33
Update the bundles
echuraev May 13, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions apps/android_camera/app/src/main/jni/tvm_runtime.h
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -57,6 +57,7 @@
#ifdef TVM_OPENCL_RUNTIME
#include "../src/runtime/opencl/opencl_device_api.cc"
#include "../src/runtime/opencl/opencl_module.cc"
#include "../src/runtime/source_utils.cc"
#endif

#ifdef TVM_VULKAN_RUNTIME
Expand Down
1 change: 1 addition & 0 deletions apps/android_rpc/app/src/main/jni/tvm_runtime.h
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -62,6 +62,7 @@
#ifdef TVM_OPENCL_RUNTIME
#include "../src/runtime/opencl/opencl_device_api.cc"
#include "../src/runtime/opencl/opencl_module.cc"
#include "../src/runtime/source_utils.cc"
#endif

#ifdef TVM_VULKAN_RUNTIME
Expand Down
1 change: 1 addition & 0 deletions golang/src/tvm_runtime_pack.cc
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -68,3 +68,4 @@
// Uncomment the following lines to enable OpenCL
// #include "../../src/runtime/opencl/opencl_device_api.cc"
// #include "../../src/runtime/opencl/opencl_module.cc"
// #include "../src/runtime/source_utils.cc"
79 changes: 45 additions & 34 deletions src/runtime/metal/metal_module.mm
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -30,6 +30,7 @@
#include "../file_utils.h"
#include "../meta_data.h"
#include "../pack_args.h"
#include "../source_utils.h"
#include "../thread_storage_scope.h"
#include "metal_common.h"

Expand All@@ -43,7 +44,9 @@
public:
explicit MetalModuleNode(std::string data, std::string fmt,
std::unordered_map<std::string, FunctionInfo> fmap, std::string source)
: data_(data), fmt_(fmt), fmap_(fmap), source_(source) {}
: data_(data), fmt_(fmt), fmap_(fmap), source_(source) {
parsed_kernels_ = SplitKernels(data);
}
const char* type_key() const final { return "metal"; }

PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) final;
Expand DownExpand Up@@ -71,6 +74,7 @@ void SaveToBinary(dmlc::Stream* stream) final {
return "";
}
}

// get a from primary context in device_id
id<MTLComputePipelineState> GetPipelineState(size_t device_id, const std::string& func_name) {
metal::MetalWorkspace* w = metal::MetalWorkspace::Global();
Expand All@@ -85,44 +89,52 @@ void SaveToBinary(dmlc::Stream* stream) final {
if (it != e.smap.end()) return it->second;
// compile
NSError* err_msg = nil;
if (e.lib == nil) {
if (fmt_ == "metal") {
MTLCompileOptions* opts = [MTLCompileOptions alloc];
opts.languageVersion = MTLLanguageVersion2_3;
opts.fastMathEnabled = YES;
// opts = nil;
e.lib = [w->devices[device_id]
newLibraryWithSource:[NSString stringWithUTF8String:data_.c_str()]
options:opts
error:&err_msg];
[opts dealloc];
if (e.lib == nil) {
LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
}
if (err_msg != nil) {
LOG(INFO) << "Warning: " << [[err_msg localizedDescription] UTF8String];
}
} else {
// Build from library.
auto q = dispatch_queue_create("q", DISPATCH_QUEUE_SERIAL);
auto data = dispatch_data_create(data_.c_str(), data_.length(), q,
^{
});
e.lib = [w->devices[device_id] newLibraryWithData:data error:&err_msg];
if (err_msg != nil || e.lib == nil) {
LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
}
id<MTLLibrary> lib = nil;
std::string source;
auto kernel = parsed_kernels_.find(func_name);
// If we cannot find this kernel in parsed_kernels_, it means that all kernels going together
// without explicit separator. In this case we use data_ with all kernels. It done for backward
// compatibility.
if (kernel != parsed_kernels_.end())

@elvin-nelvin-nMay 11, 2021•
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would add here a comment that this behaviour is done for backward compatibility when we had all kernels going together without explicit separators, but looks good

@echuraevechuraevMay 12, 2021

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done

@csullivancsullivanMay 13, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we have metal deployments that require the old behavior that can't be rebuilt easily? Otherwise we don't need to preserve backward compatibility.

@elvin-nelvin-nMay 13, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I cannot say for sure. I can imagine some use cases which might have issues with recompilation, but this is just assumptions.

If we do not declare backward compatibility for now, we can skip my comment and return original behaviour and rise an exception if we do not find kernel in parsed_kernels_ container.

@csullivancsullivanMay 13, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In that case I personally would prefer the metal backward compat be specific to the metal backend and the general SplitKernel util to contain the existence ICHECK on the function delimiter.

@elvin-nelvin-nMay 14, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tqchen do we care about backward compatibility or not? If not, it might be ok so far, but we need to decide when we will start to care and what use cases should work and be regularly verified in the future. It must not be a deal only metal backend

@csullivan what bahavior do you propose in common SplitKernel? Metal backend uses it, it cannot just abort execution. Warnings in the runtime are also useles. It might be an error code, but empty returned map detects the absent of Function delimiter with the same efficiency as return code

source = kernel->second;
else
source = data_;
if (fmt_ == "metal") {
MTLCompileOptions* opts = [MTLCompileOptions alloc];
opts.languageVersion = MTLLanguageVersion2_3;
opts.fastMathEnabled = YES;
// opts = nil;
lib =
[w->devices[device_id] newLibraryWithSource:[NSString stringWithUTF8String:source.c_str()]
options:opts
error:&err_msg];
[opts dealloc];
if (lib == nil) {
LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
}
if (err_msg != nil) {
LOG(INFO) << "Warning: " << [[err_msg localizedDescription] UTF8String];
}
} else {
// Build from library.
auto q = dispatch_queue_create("q", DISPATCH_QUEUE_SERIAL);
auto data = dispatch_data_create(source.c_str(), source.length(), q,
^{
});
lib = [w->devices[device_id] newLibraryWithData:data error:&err_msg];
if (err_msg != nil || lib == nil) {
LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
}
}
id<MTLFunction> f =
[e.lib newFunctionWithName:[NSString stringWithUTF8String:func_name.c_str()]];
id<MTLFunction> f = [lib newFunctionWithName:[NSString stringWithUTF8String:func_name.c_str()]];
ICHECK(f != nil) << "cannot find function " << func_name;
id<MTLComputePipelineState> state =
[w->devices[device_id] newComputePipelineStateWithFunction:f error:&err_msg];
ICHECK(state != nil) << "cannot get state:"
<< " for function " << func_name
<< [[err_msg localizedDescription] UTF8String];
[f release];
[lib release];
// The state.threadExecutionWidth can change dynamically according
// to the resource constraint in kernel, so it is not strictly hold
// Turn of warp aware optimziation for now.
Expand All@@ -135,13 +147,10 @@ void SaveToBinary(dmlc::Stream* stream) final {
private:
// device specific entry
struct DeviceEntry {
// library
id<MTLLibrary> lib = nil;
// state cache;
std::unordered_map<std::string, id<MTLComputePipelineState>> smap;
std::unordered_map<std::string, id<MTLComputePipelineState>> smap;

~DeviceEntry() {
if (lib != nil) [lib release];
for (auto&& kv : smap) {
[kv.second release];
}
Expand All@@ -159,6 +168,8 @@ void SaveToBinary(dmlc::Stream* stream) final {
std::vector<DeviceEntry> finfo_;
// internal mutex when updating the module
std::mutex mutex_;
// parsed kernel data
std::unordered_map<std::string, std::string> parsed_kernels_;
};

// a wrapped function class to get packed func.
Expand Down
8 changes: 0 additions & 8 deletions src/runtime/opencl/opencl_common.h
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -326,14 +326,6 @@ class OpenCLModuleNode : public ModuleNode {
cl_kernel InstallKernel(cl::OpenCLWorkspace* w, cl::OpenCLThreadEntry* t,
const std::string& func_name, const KTRefEntry& e);

/*
* \brief Splits the provided serialized source file into separate
* source for each kernel primitive.
* \param source The serialized program source file (fmt: cl)
* \return Mapping from primitive name to kernel source
*/
std::unordered_map<std::string, std::string> SplitKernels(std::string source) const;

private:
// The workspace, need to keep reference to use it in destructor.
// In case of static destruction order problem.
Expand Down
39 changes: 6 additions & 33 deletions src/runtime/opencl/opencl_module.cc
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -29,6 +29,7 @@
#include <unordered_map>
#include <vector>

#include "../source_utils.h"
#include "opencl_common.h"

namespace tvm {
Expand DownExpand Up@@ -188,6 +189,11 @@ void OpenCLModuleNode::Init() {

// split into source artifacts for each kernel
parsed_kernels_ = SplitKernels(GetSource("cl"));
ICHECK(!parsed_kernels_.empty()) << "The OpenCL module expects a kernel delimited "
<< "source from code generation, but no kernel "
<< "delimiter was found.";
ICHECK_EQ(workspace_->num_registered_kernels, parsed_kernels_.size())
<< "The number of registered kernels does not match number of parsed kernel sources";
// zero initialize cl_program pointers for each device kernel
for (auto& kv : parsed_kernels_) {
programs_.insert({kv.first, std::vector<cl_program>(workspace_->devices.size(), nullptr)});
Expand DownExpand Up@@ -242,39 +248,6 @@ cl_kernel OpenCLModuleNode::InstallKernel(cl::OpenCLWorkspace* w, cl::OpenCLThre
return kernel;
}

std::unordered_map<std::string, std::string> OpenCLModuleNode::SplitKernels(
std::string source) const {
std::unordered_map<std::string, std::string> split_kernels;
if (source.size()) {
std::string del{"// Function: "};
size_t end;
size_t begin = source.find(del);
ICHECK(begin != std::string::npos) << "The OpenCL module expects a kernel delimited "
<< "source from code generation, but no kernel "
<< "delimiter was found.";
for (size_t num_kernels = 0; num_kernels < workspace_->num_registered_kernels; num_kernels++) {
begin += del.size();
end = source.find('\n', begin);
std::string func_name = source.substr(begin, end - begin);
begin = ++end;
// std::string::substr returns either start of next kernel
// or std::string::npos, in the latter case substr returns
// all characters until the end of the source string.
end = source.find(del, begin);
std::string func_source =
source.substr(begin, (end == std::string::npos) ? end : end - begin);
split_kernels.insert({func_name, func_source});
begin = end;
if (end == std::string::npos) {
break;
}
}
}
ICHECK_EQ(workspace_->num_registered_kernels, split_kernels.size())
<< "The number of registered kernels does not match number of parsed kernel sources";
return split_kernels;
}

Module OpenCLModuleCreate(std::string data, std::string fmt,
std::unordered_map<std::string, FunctionInfo> fmap, std::string source) {
auto n = make_object<OpenCLModuleNode>(data, fmt, fmap, source);
Expand Down
49 changes: 49 additions & 0 deletions src/runtime/source_utils.cc
View file
Open in desktop
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you 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.
*/

/*!
* \file source_utils.cc
*/
#include "source_utils.h"

namespace tvm {
namespace runtime {

std::unordered_map<std::string, std::string> SplitKernels(std::string source,
Comment thread
echuraev marked this conversation as resolved.
std::string delimiter) {
std::unordered_map<std::string, std::string> split_kernels;
if (source.size()) {
size_t begin = source.find(delimiter);

@csullivancsullivanMay 12, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I noticed you deleted ICHECK here, probably because the error message referred to OpenCL specifically, but there still should be an

 ICHECK(begin != std::string::npos)

on the first search for a kernel delimiter. If none are found, then it is better to catch an error in the code generation early.

@echuraevechuraevMay 13, 2021

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I did it to keep backward compatibility as @elvin-nnov requested. In case when the Metal library was generated before this change, we won't find any delimiters in the source code.
But I added some checks to opencl module:
https://github.com/apache/tvm/pull/7980/files#diff-9d619c5f05402f93afa552f93675dea84d789f8f6b0c553845aa183128f38494R192-R196

Also, I added check in metal module were kernels found or not: https://github.com/apache/tvm/pull/7980/files#diff-dc30040e2ec938d8dd0e6e87d6517870f037093775d4b52e0f249fdbebfdb9a3R95-R101

@csullivancsullivanMay 13, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would prefer the metal backward compat be specific to the metal backend and the general SplitKernel util to contain the existence ICHECK on the function delimiter.

size_t end = begin;
while (end != std::string::npos) {
begin += delimiter.size();
end = source.find('\n', begin);
std::string func_name = source.substr(begin, end - begin);
begin = ++end;
end = source.find(delimiter, begin);
std::string func_source =
source.substr(begin, (end == std::string::npos) ? end : end - begin);
split_kernels.insert({func_name, func_source});
begin = end;
}
}
return split_kernels;
}
} // namespace runtime
} // namespace tvm
44 changes: 44 additions & 0 deletions src/runtime/source_utils.h
View file
Open in desktop
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you 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.
*/

/*!
* \file source_utils.h
* \brief Minimum source manipulation utils for runtime.
*/

#ifndef TVM_RUNTIME_SOURCE_UTILS_H_
#define TVM_RUNTIME_SOURCE_UTILS_H_

#include <string>
#include <unordered_map>

namespace tvm {
namespace runtime {
/*!
* \brief Split the source file on separate kernels by specified delimiter.
* \param source The source code of the kernels.
* \param delimiter The delimiter which is using for splitting kernels.
* \return Mapping from primitive name to kernel source
*/
std::unordered_map<std::string, std::string> SplitKernels(std::string source,
std::string delimiter = "// Function: ");
} // namespace runtime
} // namespace tvm

#endif // TVM_RUNTIME_SOURCE_UTILS_H_
25 changes: 14 additions & 11 deletions src/target/source/codegen_metal.cc
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -305,27 +305,30 @@ void CodeGenMetal::VisitExpr_(const CallNode* op, std::ostream& os) { // NOLINT
runtime::Module BuildMetal(IRModule mod, Target target) {
using tvm::runtime::Registry;
bool output_ssa = false;
CodeGenMetal cg;
cg.Init(output_ssa);

std::stringstream code;
std::stringstream source;
std::string fmt = "metal";
for (auto kv : mod->functions) {
ICHECK(kv.second->IsInstance<PrimFuncNode>()) << "CodeGenMetal: Can only take PrimFunc";
code << "// Function: " << kv.first->name_hint << std::endl;
CodeGenMetal cg;
cg.Init(output_ssa);
auto f = Downcast<PrimFunc>(kv.second);
auto calling_conv = f->GetAttr<Integer>(tvm::attr::kCallingConv);
ICHECK(calling_conv == CallingConv::kDeviceKernelLaunch)
<< "CodeGenMetal: expect calling_conv equals CallingConv::kDeviceKernelLaunch";
cg.AddFunction(f);
std::string fsource = cg.Finish();
if (const auto* f = Registry::Get("tvm_callback_metal_compile")) {
source << fsource;
fsource = (*f)(fsource).operator std::string();
fmt = "metallib";
}
code << fsource;
}

std::string code = cg.Finish();
std::string fmt = "metal";
std::string source = "";
if (const auto* f = Registry::Get("tvm_callback_metal_compile")) {
source = code;
code = (*f)(code).operator std::string();
fmt = "metallib";
}
return MetalModuleCreate(code, fmt, ExtractFuncInfo(mod), source);
return MetalModuleCreate(code.str(), fmt, ExtractFuncInfo(mod), source.str());
}

TVM_REGISTER_GLOBAL("target.build.metal").set_body_typed(BuildMetal);
Expand Down
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all 
 blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks");
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '

Skip to content
[METAL] Split kernels and compile them separately#7980

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

tqchen
merged 5 commits into

from

May 25, 2021
 Merged


[METAL] Split kernels and compile them separately
#7980
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


5 commits


Select commit
Hold shift + click to select a range
b0b51fd

[METAL] Split kernels and compile them separately


echuraev May 5, 2021
4938d70

Add backward compatible behaviour


echuraev May 6, 2021
070bba4

Move SplitKernels to a common file


echuraev May 12, 2021
8041a29

Remove using GetSource


echuraev May 12, 2021
df55c33

Update the bundles


echuraev May 13, 2021

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


1 change: 1 addition & 0 deletions
apps/android_camera/app/src/main/jni/tvm_runtime.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -57,6 +57,7 @@
#ifdef TVM_OPENCL_RUNTIME
#include "../src/runtime/opencl/opencl_device_api.cc"
#include "../src/runtime/opencl/opencl_module.cc"
#include "../src/runtime/source_utils.cc"
#endif

#ifdef TVM_VULKAN_RUNTIME
Expand Down

1 change: 1 addition & 0 deletions
apps/android_rpc/app/src/main/jni/tvm_runtime.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -62,6 +62,7 @@
#ifdef TVM_OPENCL_RUNTIME
#include "../src/runtime/opencl/opencl_device_api.cc"
#include "../src/runtime/opencl/opencl_module.cc"
#include "../src/runtime/source_utils.cc"
#endif

#ifdef TVM_VULKAN_RUNTIME
Expand Down

1 change: 1 addition & 0 deletions
golang/src/tvm_runtime_pack.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -68,3 +68,4 @@
// Uncomment the following lines to enable OpenCL
// #include "../../src/runtime/opencl/opencl_device_api.cc"
// #include "../../src/runtime/opencl/opencl_module.cc"
// #include "../src/runtime/source_utils.cc"

79 changes: 45 additions & 34 deletions
src/runtime/metal/metal_module.mm

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -30,6 +30,7 @@
#include "../file_utils.h"
#include "../meta_data.h"
#include "../pack_args.h"
#include "../source_utils.h"
#include "../thread_storage_scope.h"
#include "metal_common.h"

Expand All@@ -43,7 +44,9 @@
 public:
 explicit MetalModuleNode(std::string data, std::string fmt,
 std::unordered_map<std::string, FunctionInfo> fmap, std::string source)
 : data_(data), fmt_(fmt), fmap_(fmap), source_(source) {}
 : data_(data), fmt_(fmt), fmap_(fmap), source_(source) {
 parsed_kernels_ = SplitKernels(data);
 }
 const char* type_key() const final { return "metal"; }

 PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) final;
Expand DownExpand Up@@ -71,6 +74,7 @@ void SaveToBinary(dmlc::Stream* stream) final {
 return "";
 }
 }

 // get a from primary context in device_id
 id<MTLComputePipelineState> GetPipelineState(size_t device_id, const std::string& func_name) {
 metal::MetalWorkspace* w = metal::MetalWorkspace::Global();
Expand All@@ -85,44 +89,52 @@ void SaveToBinary(dmlc::Stream* stream) final {
 if (it != e.smap.end()) return it->second;
 // compile
 NSError* err_msg = nil;
 if (e.lib == nil) {
 if (fmt_ == "metal") {
 MTLCompileOptions* opts = [MTLCompileOptions alloc];
 opts.languageVersion = MTLLanguageVersion2_3;
 opts.fastMathEnabled = YES;
 // opts = nil;
 e.lib = [w->devices[device_id]
 newLibraryWithSource:[NSString stringWithUTF8String:data_.c_str()]
 options:opts
 error:&err_msg];
 [opts dealloc];
 if (e.lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 if (err_msg != nil) {
 LOG(INFO) << "Warning: " << [[err_msg localizedDescription] UTF8String];
 }
 } else {
 // Build from library.
 auto q = dispatch_queue_create("q", DISPATCH_QUEUE_SERIAL);
 auto data = dispatch_data_create(data_.c_str(), data_.length(), q,
 ^{
 });
 e.lib = [w->devices[device_id] newLibraryWithData:data error:&err_msg];
 if (err_msg != nil || e.lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 id<MTLLibrary> lib = nil;
 std::string source;
 auto kernel = parsed_kernels_.find(func_name);
 // If we cannot find this kernel in parsed_kernels_, it means that all kernels going together
 // without explicit separator. In this case we use data_ with all kernels. It done for backward
 // compatibility.
 if (kernel != parsed_kernels_.end())
@elvin-nelvin-nMay 11, 2021•

edited

Loading

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would add here a comment that this behaviour is done for backward compatibility when we had all kernels going together without explicit separators, but looks good
@echuraevechuraevMay 12, 2021

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done
@csullivancsullivanMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we have metal deployments that require the old behavior that can't be rebuilt easily? Otherwise we don't need to preserve backward compatibility.
@elvin-nelvin-nMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I cannot say for sure. I can imagine some use cases which might have issues with recompilation, but this is just assumptions.
If we do not declare backward compatibility for now, we can skip my comment and return original behaviour and rise an exception if we do not find kernel in parsed_kernels_ container.
@csullivancsullivanMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In that case I personally would prefer the metal backward compat be specific to the metal backend and the general SplitKernel util to contain the existence ICHECK on the function delimiter.
@elvin-nelvin-nMay 14, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tqchen do we care about backward compatibility or not? If not, it might be ok so far, but we need to decide when we will start to care and what use cases should work and be regularly verified in the future. It must not be a deal only metal backend
@csullivan what bahavior do you propose in common SplitKernel? Metal backend uses it, it cannot just abort execution. Warnings in the runtime are also useles. It might be an error code, but empty returned map detects the absent of Function delimiter with the same efficiency as return code
 source = kernel->second;
 else
 source = data_;
 if (fmt_ == "metal") {
 MTLCompileOptions* opts = [MTLCompileOptions alloc];
 opts.languageVersion = MTLLanguageVersion2_3;
 opts.fastMathEnabled = YES;
 // opts = nil;
 lib =
 [w->devices[device_id] newLibraryWithSource:[NSString stringWithUTF8String:source.c_str()]
 options:opts
 error:&err_msg];
 [opts dealloc];
 if (lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 if (err_msg != nil) {
 LOG(INFO) << "Warning: " << [[err_msg localizedDescription] UTF8String];
 }
 } else {
 // Build from library.
 auto q = dispatch_queue_create("q", DISPATCH_QUEUE_SERIAL);
 auto data = dispatch_data_create(source.c_str(), source.length(), q,
 ^{
 });
 lib = [w->devices[device_id] newLibraryWithData:data error:&err_msg];
 if (err_msg != nil || lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 }
 id<MTLFunction> f =
 [e.lib newFunctionWithName:[NSString stringWithUTF8String:func_name.c_str()]];
 id<MTLFunction> f = [lib newFunctionWithName:[NSString stringWithUTF8String:func_name.c_str()]];
 ICHECK(f != nil) << "cannot find function " << func_name;
 id<MTLComputePipelineState> state =
 [w->devices[device_id] newComputePipelineStateWithFunction:f error:&err_msg];
 ICHECK(state != nil) << "cannot get state:"
 << " for function " << func_name
 << [[err_msg localizedDescription] UTF8String];
 [f release];
 [lib release];
 // The state.threadExecutionWidth can change dynamically according
 // to the resource constraint in kernel, so it is not strictly hold
 // Turn of warp aware optimziation for now.
Expand All@@ -135,13 +147,10 @@ void SaveToBinary(dmlc::Stream* stream) final {
 private:
 // device specific entry
 struct DeviceEntry {
 // library
 id<MTLLibrary> lib = nil;
 // state cache;
 std::unordered_map<std::string, id<MTLComputePipelineState>> smap;
 std::unordered_map<std::string, id<MTLComputePipelineState>> smap;

 ~DeviceEntry() {
 if (lib != nil) [lib release];
 for (auto&& kv : smap) {
 [kv.second release];
 }
Expand All@@ -159,6 +168,8 @@ void SaveToBinary(dmlc::Stream* stream) final {
 std::vector<DeviceEntry> finfo_;
 // internal mutex when updating the module
 std::mutex mutex_;
 // parsed kernel data
 std::unordered_map<std::string, std::string> parsed_kernels_;
};

// a wrapped function class to get packed func.
Expand Down

8 changes: 0 additions & 8 deletions
src/runtime/opencl/opencl_common.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -326,14 +326,6 @@ class OpenCLModuleNode : public ModuleNode {
 cl_kernel InstallKernel(cl::OpenCLWorkspace* w, cl::OpenCLThreadEntry* t,
 const std::string& func_name, const KTRefEntry& e);

 /*
 * \brief Splits the provided serialized source file into separate
 * source for each kernel primitive.
 * \param source The serialized program source file (fmt: cl)
 * \return Mapping from primitive name to kernel source
 */
 std::unordered_map<std::string, std::string> SplitKernels(std::string source) const;

 private:
 // The workspace, need to keep reference to use it in destructor.
 // In case of static destruction order problem.
Expand Down

39 changes: 6 additions & 33 deletions
src/runtime/opencl/opencl_module.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -29,6 +29,7 @@
#include <unordered_map>
#include <vector>

#include "../source_utils.h"
#include "opencl_common.h"

namespace tvm {
Expand DownExpand Up@@ -188,6 +189,11 @@ void OpenCLModuleNode::Init() {

 // split into source artifacts for each kernel
 parsed_kernels_ = SplitKernels(GetSource("cl"));
 ICHECK(!parsed_kernels_.empty()) << "The OpenCL module expects a kernel delimited "
 << "source from code generation, but no kernel "
 << "delimiter was found.";
 ICHECK_EQ(workspace_->num_registered_kernels, parsed_kernels_.size())
 << "The number of registered kernels does not match number of parsed kernel sources";
 // zero initialize cl_program pointers for each device kernel
 for (auto& kv : parsed_kernels_) {
 programs_.insert({kv.first, std::vector<cl_program>(workspace_->devices.size(), nullptr)});
Expand DownExpand Up@@ -242,39 +248,6 @@ cl_kernel OpenCLModuleNode::InstallKernel(cl::OpenCLWorkspace* w, cl::OpenCLThre
 return kernel;
}

std::unordered_map<std::string, std::string> OpenCLModuleNode::SplitKernels(
 std::string source) const {
 std::unordered_map<std::string, std::string> split_kernels;
 if (source.size()) {
 std::string del{"// Function: "};
 size_t end;
 size_t begin = source.find(del);
 ICHECK(begin != std::string::npos) << "The OpenCL module expects a kernel delimited "
 << "source from code generation, but no kernel "
 << "delimiter was found.";
 for (size_t num_kernels = 0; num_kernels < workspace_->num_registered_kernels; num_kernels++) {
 begin += del.size();
 end = source.find('\n', begin);
 std::string func_name = source.substr(begin, end - begin);
 begin = ++end;
 // std::string::substr returns either start of next kernel
 // or std::string::npos, in the latter case substr returns
 // all characters until the end of the source string.
 end = source.find(del, begin);
 std::string func_source =
 source.substr(begin, (end == std::string::npos) ? end : end - begin);
 split_kernels.insert({func_name, func_source});
 begin = end;
 if (end == std::string::npos) {
 break;
 }
 }
 }
 ICHECK_EQ(workspace_->num_registered_kernels, split_kernels.size())
 << "The number of registered kernels does not match number of parsed kernel sources";
 return split_kernels;
}

Module OpenCLModuleCreate(std::string data, std::string fmt,
 std::unordered_map<std::string, FunctionInfo> fmap, std::string source) {
 auto n = make_object<OpenCLModuleNode>(data, fmt, fmap, source);
Expand Down

49 changes: 49 additions & 0 deletions
src/runtime/source_utils.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you 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.
 */

/*!
 * \file source_utils.cc
 */
#include "source_utils.h"

namespace tvm {
namespace runtime {

std::unordered_map<std::string, std::string> SplitKernels(std::string source,
Comment thread
echuraev marked this conversation as resolved.

 std::string delimiter) {
 std::unordered_map<std::string, std::string> split_kernels;
 if (source.size()) {
 size_t begin = source.find(delimiter);
@csullivancsullivanMay 12, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I noticed you deleted ICHECK here, probably because the error message referred to OpenCL specifically, but there still should be an
 ICHECK(begin != std::string::npos) 
on the first search for a kernel delimiter. If none are found, then it is better to catch an error in the code generation early.
@echuraevechuraevMay 13, 2021

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I did it to keep backward compatibility as @elvin-nnov requested. In case when the Metal library was generated before this change, we won't find any delimiters in the source code.

But I added some checks to opencl module:
https://github.com/apache/tvm/pull/7980/files#diff-9d619c5f05402f93afa552f93675dea84d789f8f6b0c553845aa183128f38494R192-R196
Also, I added check in metal module were kernels found or not: https://github.com/apache/tvm/pull/7980/files#diff-dc30040e2ec938d8dd0e6e87d6517870f037093775d4b52e0f249fdbebfdb9a3R95-R101
@csullivancsullivanMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would prefer the metal backward compat be specific to the metal backend and the general SplitKernel util to contain the existence ICHECK on the function delimiter.
 size_t end = begin;
 while (end != std::string::npos) {
 begin += delimiter.size();
 end = source.find('\n', begin);
 std::string func_name = source.substr(begin, end - begin);
 begin = ++end;
 end = source.find(delimiter, begin);
 std::string func_source =
 source.substr(begin, (end == std::string::npos) ? end : end - begin);
 split_kernels.insert({func_name, func_source});
 begin = end;
 }
 }
 return split_kernels;
}
} // namespace runtime
} // namespace tvm

44 changes: 44 additions & 0 deletions
src/runtime/source_utils.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you 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.
 */

/*!
 * \file source_utils.h
 * \brief Minimum source manipulation utils for runtime.
 */

#ifndef TVM_RUNTIME_SOURCE_UTILS_H_
#define TVM_RUNTIME_SOURCE_UTILS_H_

#include <string>
#include <unordered_map>

namespace tvm {
namespace runtime {
/*!
 * \brief Split the source file on separate kernels by specified delimiter.
 * \param source The source code of the kernels.
 * \param delimiter The delimiter which is using for splitting kernels.
 * \return Mapping from primitive name to kernel source
 */
std::unordered_map<std::string, std::string> SplitKernels(std::string source,
 std::string delimiter = "// Function: ");
} // namespace runtime
} // namespace tvm

#endif // TVM_RUNTIME_SOURCE_UTILS_H_

25 changes: 14 additions & 11 deletions
src/target/source/codegen_metal.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -305,27 +305,30 @@ void CodeGenMetal::VisitExpr_(const CallNode* op, std::ostream& os) { // NOLINT
runtime::Module BuildMetal(IRModule mod, Target target) {
 using tvm::runtime::Registry;
 bool output_ssa = false;
 CodeGenMetal cg;
 cg.Init(output_ssa);

 std::stringstream code;
 std::stringstream source;
 std::string fmt = "metal";
 for (auto kv : mod->functions) {
 ICHECK(kv.second->IsInstance<PrimFuncNode>()) << "CodeGenMetal: Can only take PrimFunc";
 code << "// Function: " << kv.first->name_hint << std::endl;
 CodeGenMetal cg;
 cg.Init(output_ssa);
 auto f = Downcast<PrimFunc>(kv.second);
 auto calling_conv = f->GetAttr<Integer>(tvm::attr::kCallingConv);
 ICHECK(calling_conv == CallingConv::kDeviceKernelLaunch)
 << "CodeGenMetal: expect calling_conv equals CallingConv::kDeviceKernelLaunch";
 cg.AddFunction(f);
 std::string fsource = cg.Finish();
 if (const auto* f = Registry::Get("tvm_callback_metal_compile")) {
 source << fsource;
 fsource = (*f)(fsource).operator std::string();
 fmt = "metallib";
 }
 code << fsource;
 }

 std::string code = cg.Finish();
 std::string fmt = "metal";
 std::string source = "";
 if (const auto* f = Registry::Get("tvm_callback_metal_compile")) {
 source = code;
 code = (*f)(code).operator std::string();
 fmt = "metallib";
 }
 return MetalModuleCreate(code, fmt, ExtractFuncInfo(mod), source);
 return MetalModuleCreate(code.str(), fmt, ExtractFuncInfo(mod), source.str());
}

TVM_REGISTER_GLOBAL("target.build.metal").set_body_typed(BuildMetal);
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix");
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
[METAL] Split kernels and compile them separately#7980

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

tqchen
merged 5 commits into

from

May 25, 2021
 Merged


[METAL] Split kernels and compile them separately
#7980
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


5 commits


Select commit
Hold shift + click to select a range
b0b51fd

[METAL] Split kernels and compile them separately


echuraev May 5, 2021
4938d70

Add backward compatible behaviour


echuraev May 6, 2021
070bba4

Move SplitKernels to a common file


echuraev May 12, 2021
8041a29

Remove using GetSource


echuraev May 12, 2021
df55c33

Update the bundles


echuraev May 13, 2021

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


1 change: 1 addition & 0 deletions
apps/android_camera/app/src/main/jni/tvm_runtime.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -57,6 +57,7 @@
#ifdef TVM_OPENCL_RUNTIME
#include "../src/runtime/opencl/opencl_device_api.cc"
#include "../src/runtime/opencl/opencl_module.cc"
#include "../src/runtime/source_utils.cc"
#endif

#ifdef TVM_VULKAN_RUNTIME
Expand Down

1 change: 1 addition & 0 deletions
apps/android_rpc/app/src/main/jni/tvm_runtime.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -62,6 +62,7 @@
#ifdef TVM_OPENCL_RUNTIME
#include "../src/runtime/opencl/opencl_device_api.cc"
#include "../src/runtime/opencl/opencl_module.cc"
#include "../src/runtime/source_utils.cc"
#endif

#ifdef TVM_VULKAN_RUNTIME
Expand Down

1 change: 1 addition & 0 deletions
golang/src/tvm_runtime_pack.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -68,3 +68,4 @@
// Uncomment the following lines to enable OpenCL
// #include "../../src/runtime/opencl/opencl_device_api.cc"
// #include "../../src/runtime/opencl/opencl_module.cc"
// #include "../src/runtime/source_utils.cc"

79 changes: 45 additions & 34 deletions
src/runtime/metal/metal_module.mm

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -30,6 +30,7 @@
#include "../file_utils.h"
#include "../meta_data.h"
#include "../pack_args.h"
#include "../source_utils.h"
#include "../thread_storage_scope.h"
#include "metal_common.h"

Expand All@@ -43,7 +44,9 @@
 public:
 explicit MetalModuleNode(std::string data, std::string fmt,
 std::unordered_map<std::string, FunctionInfo> fmap, std::string source)
 : data_(data), fmt_(fmt), fmap_(fmap), source_(source) {}
 : data_(data), fmt_(fmt), fmap_(fmap), source_(source) {
 parsed_kernels_ = SplitKernels(data);
 }
 const char* type_key() const final { return "metal"; }

 PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) final;
Expand DownExpand Up@@ -71,6 +74,7 @@ void SaveToBinary(dmlc::Stream* stream) final {
 return "";
 }
 }

 // get a from primary context in device_id
 id<MTLComputePipelineState> GetPipelineState(size_t device_id, const std::string& func_name) {
 metal::MetalWorkspace* w = metal::MetalWorkspace::Global();
Expand All@@ -85,44 +89,52 @@ void SaveToBinary(dmlc::Stream* stream) final {
 if (it != e.smap.end()) return it->second;
 // compile
 NSError* err_msg = nil;
 if (e.lib == nil) {
 if (fmt_ == "metal") {
 MTLCompileOptions* opts = [MTLCompileOptions alloc];
 opts.languageVersion = MTLLanguageVersion2_3;
 opts.fastMathEnabled = YES;
 // opts = nil;
 e.lib = [w->devices[device_id]
 newLibraryWithSource:[NSString stringWithUTF8String:data_.c_str()]
 options:opts
 error:&err_msg];
 [opts dealloc];
 if (e.lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 if (err_msg != nil) {
 LOG(INFO) << "Warning: " << [[err_msg localizedDescription] UTF8String];
 }
 } else {
 // Build from library.
 auto q = dispatch_queue_create("q", DISPATCH_QUEUE_SERIAL);
 auto data = dispatch_data_create(data_.c_str(), data_.length(), q,
 ^{
 });
 e.lib = [w->devices[device_id] newLibraryWithData:data error:&err_msg];
 if (err_msg != nil || e.lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 id<MTLLibrary> lib = nil;
 std::string source;
 auto kernel = parsed_kernels_.find(func_name);
 // If we cannot find this kernel in parsed_kernels_, it means that all kernels going together
 // without explicit separator. In this case we use data_ with all kernels. It done for backward
 // compatibility.
 if (kernel != parsed_kernels_.end())
@elvin-nelvin-nMay 11, 2021•

edited

Loading

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would add here a comment that this behaviour is done for backward compatibility when we had all kernels going together without explicit separators, but looks good
@echuraevechuraevMay 12, 2021

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done
@csullivancsullivanMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we have metal deployments that require the old behavior that can't be rebuilt easily? Otherwise we don't need to preserve backward compatibility.
@elvin-nelvin-nMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I cannot say for sure. I can imagine some use cases which might have issues with recompilation, but this is just assumptions.
If we do not declare backward compatibility for now, we can skip my comment and return original behaviour and rise an exception if we do not find kernel in parsed_kernels_ container.
@csullivancsullivanMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In that case I personally would prefer the metal backward compat be specific to the metal backend and the general SplitKernel util to contain the existence ICHECK on the function delimiter.
@elvin-nelvin-nMay 14, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tqchen do we care about backward compatibility or not? If not, it might be ok so far, but we need to decide when we will start to care and what use cases should work and be regularly verified in the future. It must not be a deal only metal backend
@csullivan what bahavior do you propose in common SplitKernel? Metal backend uses it, it cannot just abort execution. Warnings in the runtime are also useles. It might be an error code, but empty returned map detects the absent of Function delimiter with the same efficiency as return code
 source = kernel->second;
 else
 source = data_;
 if (fmt_ == "metal") {
 MTLCompileOptions* opts = [MTLCompileOptions alloc];
 opts.languageVersion = MTLLanguageVersion2_3;
 opts.fastMathEnabled = YES;
 // opts = nil;
 lib =
 [w->devices[device_id] newLibraryWithSource:[NSString stringWithUTF8String:source.c_str()]
 options:opts
 error:&err_msg];
 [opts dealloc];
 if (lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 if (err_msg != nil) {
 LOG(INFO) << "Warning: " << [[err_msg localizedDescription] UTF8String];
 }
 } else {
 // Build from library.
 auto q = dispatch_queue_create("q", DISPATCH_QUEUE_SERIAL);
 auto data = dispatch_data_create(source.c_str(), source.length(), q,
 ^{
 });
 lib = [w->devices[device_id] newLibraryWithData:data error:&err_msg];
 if (err_msg != nil || lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 }
 id<MTLFunction> f =
 [e.lib newFunctionWithName:[NSString stringWithUTF8String:func_name.c_str()]];
 id<MTLFunction> f = [lib newFunctionWithName:[NSString stringWithUTF8String:func_name.c_str()]];
 ICHECK(f != nil) << "cannot find function " << func_name;
 id<MTLComputePipelineState> state =
 [w->devices[device_id] newComputePipelineStateWithFunction:f error:&err_msg];
 ICHECK(state != nil) << "cannot get state:"
 << " for function " << func_name
 << [[err_msg localizedDescription] UTF8String];
 [f release];
 [lib release];
 // The state.threadExecutionWidth can change dynamically according
 // to the resource constraint in kernel, so it is not strictly hold
 // Turn of warp aware optimziation for now.
Expand All@@ -135,13 +147,10 @@ void SaveToBinary(dmlc::Stream* stream) final {
 private:
 // device specific entry
 struct DeviceEntry {
 // library
 id<MTLLibrary> lib = nil;
 // state cache;
 std::unordered_map<std::string, id<MTLComputePipelineState>> smap;
 std::unordered_map<std::string, id<MTLComputePipelineState>> smap;

 ~DeviceEntry() {
 if (lib != nil) [lib release];
 for (auto&& kv : smap) {
 [kv.second release];
 }
Expand All@@ -159,6 +168,8 @@ void SaveToBinary(dmlc::Stream* stream) final {
 std::vector<DeviceEntry> finfo_;
 // internal mutex when updating the module
 std::mutex mutex_;
 // parsed kernel data
 std::unordered_map<std::string, std::string> parsed_kernels_;
};

// a wrapped function class to get packed func.
Expand Down

8 changes: 0 additions & 8 deletions
src/runtime/opencl/opencl_common.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -326,14 +326,6 @@ class OpenCLModuleNode : public ModuleNode {
 cl_kernel InstallKernel(cl::OpenCLWorkspace* w, cl::OpenCLThreadEntry* t,
 const std::string& func_name, const KTRefEntry& e);

 /*
 * \brief Splits the provided serialized source file into separate
 * source for each kernel primitive.
 * \param source The serialized program source file (fmt: cl)
 * \return Mapping from primitive name to kernel source
 */
 std::unordered_map<std::string, std::string> SplitKernels(std::string source) const;

 private:
 // The workspace, need to keep reference to use it in destructor.
 // In case of static destruction order problem.
Expand Down

39 changes: 6 additions & 33 deletions
src/runtime/opencl/opencl_module.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -29,6 +29,7 @@
#include <unordered_map>
#include <vector>

#include "../source_utils.h"
#include "opencl_common.h"

namespace tvm {
Expand DownExpand Up@@ -188,6 +189,11 @@ void OpenCLModuleNode::Init() {

 // split into source artifacts for each kernel
 parsed_kernels_ = SplitKernels(GetSource("cl"));
 ICHECK(!parsed_kernels_.empty()) << "The OpenCL module expects a kernel delimited "
 << "source from code generation, but no kernel "
 << "delimiter was found.";
 ICHECK_EQ(workspace_->num_registered_kernels, parsed_kernels_.size())
 << "The number of registered kernels does not match number of parsed kernel sources";
 // zero initialize cl_program pointers for each device kernel
 for (auto& kv : parsed_kernels_) {
 programs_.insert({kv.first, std::vector<cl_program>(workspace_->devices.size(), nullptr)});
Expand DownExpand Up@@ -242,39 +248,6 @@ cl_kernel OpenCLModuleNode::InstallKernel(cl::OpenCLWorkspace* w, cl::OpenCLThre
 return kernel;
}

std::unordered_map<std::string, std::string> OpenCLModuleNode::SplitKernels(
 std::string source) const {
 std::unordered_map<std::string, std::string> split_kernels;
 if (source.size()) {
 std::string del{"// Function: "};
 size_t end;
 size_t begin = source.find(del);
 ICHECK(begin != std::string::npos) << "The OpenCL module expects a kernel delimited "
 << "source from code generation, but no kernel "
 << "delimiter was found.";
 for (size_t num_kernels = 0; num_kernels < workspace_->num_registered_kernels; num_kernels++) {
 begin += del.size();
 end = source.find('\n', begin);
 std::string func_name = source.substr(begin, end - begin);
 begin = ++end;
 // std::string::substr returns either start of next kernel
 // or std::string::npos, in the latter case substr returns
 // all characters until the end of the source string.
 end = source.find(del, begin);
 std::string func_source =
 source.substr(begin, (end == std::string::npos) ? end : end - begin);
 split_kernels.insert({func_name, func_source});
 begin = end;
 if (end == std::string::npos) {
 break;
 }
 }
 }
 ICHECK_EQ(workspace_->num_registered_kernels, split_kernels.size())
 << "The number of registered kernels does not match number of parsed kernel sources";
 return split_kernels;
}

Module OpenCLModuleCreate(std::string data, std::string fmt,
 std::unordered_map<std::string, FunctionInfo> fmap, std::string source) {
 auto n = make_object<OpenCLModuleNode>(data, fmt, fmap, source);
Expand Down

49 changes: 49 additions & 0 deletions
src/runtime/source_utils.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you 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.
 */

/*!
 * \file source_utils.cc
 */
#include "source_utils.h"

namespace tvm {
namespace runtime {

std::unordered_map<std::string, std::string> SplitKernels(std::string source,
Comment thread
echuraev marked this conversation as resolved.

 std::string delimiter) {
 std::unordered_map<std::string, std::string> split_kernels;
 if (source.size()) {
 size_t begin = source.find(delimiter);
@csullivancsullivanMay 12, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I noticed you deleted ICHECK here, probably because the error message referred to OpenCL specifically, but there still should be an
 ICHECK(begin != std::string::npos) 
on the first search for a kernel delimiter. If none are found, then it is better to catch an error in the code generation early.
@echuraevechuraevMay 13, 2021

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I did it to keep backward compatibility as @elvin-nnov requested. In case when the Metal library was generated before this change, we won't find any delimiters in the source code.

But I added some checks to opencl module:
https://github.com/apache/tvm/pull/7980/files#diff-9d619c5f05402f93afa552f93675dea84d789f8f6b0c553845aa183128f38494R192-R196
Also, I added check in metal module were kernels found or not: https://github.com/apache/tvm/pull/7980/files#diff-dc30040e2ec938d8dd0e6e87d6517870f037093775d4b52e0f249fdbebfdb9a3R95-R101
@csullivancsullivanMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would prefer the metal backward compat be specific to the metal backend and the general SplitKernel util to contain the existence ICHECK on the function delimiter.
 size_t end = begin;
 while (end != std::string::npos) {
 begin += delimiter.size();
 end = source.find('\n', begin);
 std::string func_name = source.substr(begin, end - begin);
 begin = ++end;
 end = source.find(delimiter, begin);
 std::string func_source =
 source.substr(begin, (end == std::string::npos) ? end : end - begin);
 split_kernels.insert({func_name, func_source});
 begin = end;
 }
 }
 return split_kernels;
}
} // namespace runtime
} // namespace tvm

44 changes: 44 additions & 0 deletions
src/runtime/source_utils.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you 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.
 */

/*!
 * \file source_utils.h
 * \brief Minimum source manipulation utils for runtime.
 */

#ifndef TVM_RUNTIME_SOURCE_UTILS_H_
#define TVM_RUNTIME_SOURCE_UTILS_H_

#include <string>
#include <unordered_map>

namespace tvm {
namespace runtime {
/*!
 * \brief Split the source file on separate kernels by specified delimiter.
 * \param source The source code of the kernels.
 * \param delimiter The delimiter which is using for splitting kernels.
 * \return Mapping from primitive name to kernel source
 */
std::unordered_map<std::string, std::string> SplitKernels(std::string source,
 std::string delimiter = "// Function: ");
} // namespace runtime
} // namespace tvm

#endif // TVM_RUNTIME_SOURCE_UTILS_H_

25 changes: 14 additions & 11 deletions
src/target/source/codegen_metal.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -305,27 +305,30 @@ void CodeGenMetal::VisitExpr_(const CallNode* op, std::ostream& os) { // NOLINT
runtime::Module BuildMetal(IRModule mod, Target target) {
 using tvm::runtime::Registry;
 bool output_ssa = false;
 CodeGenMetal cg;
 cg.Init(output_ssa);

 std::stringstream code;
 std::stringstream source;
 std::string fmt = "metal";
 for (auto kv : mod->functions) {
 ICHECK(kv.second->IsInstance<PrimFuncNode>()) << "CodeGenMetal: Can only take PrimFunc";
 code << "// Function: " << kv.first->name_hint << std::endl;
 CodeGenMetal cg;
 cg.Init(output_ssa);
 auto f = Downcast<PrimFunc>(kv.second);
 auto calling_conv = f->GetAttr<Integer>(tvm::attr::kCallingConv);
 ICHECK(calling_conv == CallingConv::kDeviceKernelLaunch)
 << "CodeGenMetal: expect calling_conv equals CallingConv::kDeviceKernelLaunch";
 cg.AddFunction(f);
 std::string fsource = cg.Finish();
 if (const auto* f = Registry::Get("tvm_callback_metal_compile")) {
 source << fsource;
 fsource = (*f)(fsource).operator std::string();
 fmt = "metallib";
 }
 code << fsource;
 }

 std::string code = cg.Finish();
 std::string fmt = "metal";
 std::string source = "";
 if (const auto* f = Registry::Get("tvm_callback_metal_compile")) {
 source = code;
 code = (*f)(code).operator std::string();
 fmt = "metallib";
 }
 return MetalModuleCreate(code, fmt, ExtractFuncInfo(mod), source);
 return MetalModuleCreate(code.str(), fmt, ExtractFuncInfo(mod), source.str());
}

TVM_REGISTER_GLOBAL("target.build.metal").set_body_typed(BuildMetal);
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length > 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms");
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
[METAL] Split kernels and compile them separately#7980

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

tqchen
merged 5 commits into

from

May 25, 2021
 Merged


[METAL] Split kernels and compile them separately
#7980
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


5 commits


Select commit
Hold shift + click to select a range
b0b51fd

[METAL] Split kernels and compile them separately


echuraev May 5, 2021
4938d70

Add backward compatible behaviour


echuraev May 6, 2021
070bba4

Move SplitKernels to a common file


echuraev May 12, 2021
8041a29

Remove using GetSource


echuraev May 12, 2021
df55c33

Update the bundles


echuraev May 13, 2021

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


1 change: 1 addition & 0 deletions
apps/android_camera/app/src/main/jni/tvm_runtime.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -57,6 +57,7 @@
#ifdef TVM_OPENCL_RUNTIME
#include "../src/runtime/opencl/opencl_device_api.cc"
#include "../src/runtime/opencl/opencl_module.cc"
#include "../src/runtime/source_utils.cc"
#endif

#ifdef TVM_VULKAN_RUNTIME
Expand Down

1 change: 1 addition & 0 deletions
apps/android_rpc/app/src/main/jni/tvm_runtime.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -62,6 +62,7 @@
#ifdef TVM_OPENCL_RUNTIME
#include "../src/runtime/opencl/opencl_device_api.cc"
#include "../src/runtime/opencl/opencl_module.cc"
#include "../src/runtime/source_utils.cc"
#endif

#ifdef TVM_VULKAN_RUNTIME
Expand Down

1 change: 1 addition & 0 deletions
golang/src/tvm_runtime_pack.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -68,3 +68,4 @@
// Uncomment the following lines to enable OpenCL
// #include "../../src/runtime/opencl/opencl_device_api.cc"
// #include "../../src/runtime/opencl/opencl_module.cc"
// #include "../src/runtime/source_utils.cc"

79 changes: 45 additions & 34 deletions
src/runtime/metal/metal_module.mm

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -30,6 +30,7 @@
#include "../file_utils.h"
#include "../meta_data.h"
#include "../pack_args.h"
#include "../source_utils.h"
#include "../thread_storage_scope.h"
#include "metal_common.h"

Expand All@@ -43,7 +44,9 @@
 public:
 explicit MetalModuleNode(std::string data, std::string fmt,
 std::unordered_map<std::string, FunctionInfo> fmap, std::string source)
 : data_(data), fmt_(fmt), fmap_(fmap), source_(source) {}
 : data_(data), fmt_(fmt), fmap_(fmap), source_(source) {
 parsed_kernels_ = SplitKernels(data);
 }
 const char* type_key() const final { return "metal"; }

 PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) final;
Expand DownExpand Up@@ -71,6 +74,7 @@ void SaveToBinary(dmlc::Stream* stream) final {
 return "";
 }
 }

 // get a from primary context in device_id
 id<MTLComputePipelineState> GetPipelineState(size_t device_id, const std::string& func_name) {
 metal::MetalWorkspace* w = metal::MetalWorkspace::Global();
Expand All@@ -85,44 +89,52 @@ void SaveToBinary(dmlc::Stream* stream) final {
 if (it != e.smap.end()) return it->second;
 // compile
 NSError* err_msg = nil;
 if (e.lib == nil) {
 if (fmt_ == "metal") {
 MTLCompileOptions* opts = [MTLCompileOptions alloc];
 opts.languageVersion = MTLLanguageVersion2_3;
 opts.fastMathEnabled = YES;
 // opts = nil;
 e.lib = [w->devices[device_id]
 newLibraryWithSource:[NSString stringWithUTF8String:data_.c_str()]
 options:opts
 error:&err_msg];
 [opts dealloc];
 if (e.lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 if (err_msg != nil) {
 LOG(INFO) << "Warning: " << [[err_msg localizedDescription] UTF8String];
 }
 } else {
 // Build from library.
 auto q = dispatch_queue_create("q", DISPATCH_QUEUE_SERIAL);
 auto data = dispatch_data_create(data_.c_str(), data_.length(), q,
 ^{
 });
 e.lib = [w->devices[device_id] newLibraryWithData:data error:&err_msg];
 if (err_msg != nil || e.lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 id<MTLLibrary> lib = nil;
 std::string source;
 auto kernel = parsed_kernels_.find(func_name);
 // If we cannot find this kernel in parsed_kernels_, it means that all kernels going together
 // without explicit separator. In this case we use data_ with all kernels. It done for backward
 // compatibility.
 if (kernel != parsed_kernels_.end())
@elvin-nelvin-nMay 11, 2021•

edited

Loading

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would add here a comment that this behaviour is done for backward compatibility when we had all kernels going together without explicit separators, but looks good
@echuraevechuraevMay 12, 2021

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done
@csullivancsullivanMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we have metal deployments that require the old behavior that can't be rebuilt easily? Otherwise we don't need to preserve backward compatibility.
@elvin-nelvin-nMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I cannot say for sure. I can imagine some use cases which might have issues with recompilation, but this is just assumptions.
If we do not declare backward compatibility for now, we can skip my comment and return original behaviour and rise an exception if we do not find kernel in parsed_kernels_ container.
@csullivancsullivanMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In that case I personally would prefer the metal backward compat be specific to the metal backend and the general SplitKernel util to contain the existence ICHECK on the function delimiter.
@elvin-nelvin-nMay 14, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tqchen do we care about backward compatibility or not? If not, it might be ok so far, but we need to decide when we will start to care and what use cases should work and be regularly verified in the future. It must not be a deal only metal backend
@csullivan what bahavior do you propose in common SplitKernel? Metal backend uses it, it cannot just abort execution. Warnings in the runtime are also useles. It might be an error code, but empty returned map detects the absent of Function delimiter with the same efficiency as return code
 source = kernel->second;
 else
 source = data_;
 if (fmt_ == "metal") {
 MTLCompileOptions* opts = [MTLCompileOptions alloc];
 opts.languageVersion = MTLLanguageVersion2_3;
 opts.fastMathEnabled = YES;
 // opts = nil;
 lib =
 [w->devices[device_id] newLibraryWithSource:[NSString stringWithUTF8String:source.c_str()]
 options:opts
 error:&err_msg];
 [opts dealloc];
 if (lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 if (err_msg != nil) {
 LOG(INFO) << "Warning: " << [[err_msg localizedDescription] UTF8String];
 }
 } else {
 // Build from library.
 auto q = dispatch_queue_create("q", DISPATCH_QUEUE_SERIAL);
 auto data = dispatch_data_create(source.c_str(), source.length(), q,
 ^{
 });
 lib = [w->devices[device_id] newLibraryWithData:data error:&err_msg];
 if (err_msg != nil || lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 }
 id<MTLFunction> f =
 [e.lib newFunctionWithName:[NSString stringWithUTF8String:func_name.c_str()]];
 id<MTLFunction> f = [lib newFunctionWithName:[NSString stringWithUTF8String:func_name.c_str()]];
 ICHECK(f != nil) << "cannot find function " << func_name;
 id<MTLComputePipelineState> state =
 [w->devices[device_id] newComputePipelineStateWithFunction:f error:&err_msg];
 ICHECK(state != nil) << "cannot get state:"
 << " for function " << func_name
 << [[err_msg localizedDescription] UTF8String];
 [f release];
 [lib release];
 // The state.threadExecutionWidth can change dynamically according
 // to the resource constraint in kernel, so it is not strictly hold
 // Turn of warp aware optimziation for now.
Expand All@@ -135,13 +147,10 @@ void SaveToBinary(dmlc::Stream* stream) final {
 private:
 // device specific entry
 struct DeviceEntry {
 // library
 id<MTLLibrary> lib = nil;
 // state cache;
 std::unordered_map<std::string, id<MTLComputePipelineState>> smap;
 std::unordered_map<std::string, id<MTLComputePipelineState>> smap;

 ~DeviceEntry() {
 if (lib != nil) [lib release];
 for (auto&& kv : smap) {
 [kv.second release];
 }
Expand All@@ -159,6 +168,8 @@ void SaveToBinary(dmlc::Stream* stream) final {
 std::vector<DeviceEntry> finfo_;
 // internal mutex when updating the module
 std::mutex mutex_;
 // parsed kernel data
 std::unordered_map<std::string, std::string> parsed_kernels_;
};

// a wrapped function class to get packed func.
Expand Down

8 changes: 0 additions & 8 deletions
src/runtime/opencl/opencl_common.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -326,14 +326,6 @@ class OpenCLModuleNode : public ModuleNode {
 cl_kernel InstallKernel(cl::OpenCLWorkspace* w, cl::OpenCLThreadEntry* t,
 const std::string& func_name, const KTRefEntry& e);

 /*
 * \brief Splits the provided serialized source file into separate
 * source for each kernel primitive.
 * \param source The serialized program source file (fmt: cl)
 * \return Mapping from primitive name to kernel source
 */
 std::unordered_map<std::string, std::string> SplitKernels(std::string source) const;

 private:
 // The workspace, need to keep reference to use it in destructor.
 // In case of static destruction order problem.
Expand Down

39 changes: 6 additions & 33 deletions
src/runtime/opencl/opencl_module.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -29,6 +29,7 @@
#include <unordered_map>
#include <vector>

#include "../source_utils.h"
#include "opencl_common.h"

namespace tvm {
Expand DownExpand Up@@ -188,6 +189,11 @@ void OpenCLModuleNode::Init() {

 // split into source artifacts for each kernel
 parsed_kernels_ = SplitKernels(GetSource("cl"));
 ICHECK(!parsed_kernels_.empty()) << "The OpenCL module expects a kernel delimited "
 << "source from code generation, but no kernel "
 << "delimiter was found.";
 ICHECK_EQ(workspace_->num_registered_kernels, parsed_kernels_.size())
 << "The number of registered kernels does not match number of parsed kernel sources";
 // zero initialize cl_program pointers for each device kernel
 for (auto& kv : parsed_kernels_) {
 programs_.insert({kv.first, std::vector<cl_program>(workspace_->devices.size(), nullptr)});
Expand DownExpand Up@@ -242,39 +248,6 @@ cl_kernel OpenCLModuleNode::InstallKernel(cl::OpenCLWorkspace* w, cl::OpenCLThre
 return kernel;
}

std::unordered_map<std::string, std::string> OpenCLModuleNode::SplitKernels(
 std::string source) const {
 std::unordered_map<std::string, std::string> split_kernels;
 if (source.size()) {
 std::string del{"// Function: "};
 size_t end;
 size_t begin = source.find(del);
 ICHECK(begin != std::string::npos) << "The OpenCL module expects a kernel delimited "
 << "source from code generation, but no kernel "
 << "delimiter was found.";
 for (size_t num_kernels = 0; num_kernels < workspace_->num_registered_kernels; num_kernels++) {
 begin += del.size();
 end = source.find('\n', begin);
 std::string func_name = source.substr(begin, end - begin);
 begin = ++end;
 // std::string::substr returns either start of next kernel
 // or std::string::npos, in the latter case substr returns
 // all characters until the end of the source string.
 end = source.find(del, begin);
 std::string func_source =
 source.substr(begin, (end == std::string::npos) ? end : end - begin);
 split_kernels.insert({func_name, func_source});
 begin = end;
 if (end == std::string::npos) {
 break;
 }
 }
 }
 ICHECK_EQ(workspace_->num_registered_kernels, split_kernels.size())
 << "The number of registered kernels does not match number of parsed kernel sources";
 return split_kernels;
}

Module OpenCLModuleCreate(std::string data, std::string fmt,
 std::unordered_map<std::string, FunctionInfo> fmap, std::string source) {
 auto n = make_object<OpenCLModuleNode>(data, fmt, fmap, source);
Expand Down

49 changes: 49 additions & 0 deletions
src/runtime/source_utils.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you 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.
 */

/*!
 * \file source_utils.cc
 */
#include "source_utils.h"

namespace tvm {
namespace runtime {

std::unordered_map<std::string, std::string> SplitKernels(std::string source,
Comment thread
echuraev marked this conversation as resolved.

 std::string delimiter) {
 std::unordered_map<std::string, std::string> split_kernels;
 if (source.size()) {
 size_t begin = source.find(delimiter);
@csullivancsullivanMay 12, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I noticed you deleted ICHECK here, probably because the error message referred to OpenCL specifically, but there still should be an
 ICHECK(begin != std::string::npos) 
on the first search for a kernel delimiter. If none are found, then it is better to catch an error in the code generation early.
@echuraevechuraevMay 13, 2021

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I did it to keep backward compatibility as @elvin-nnov requested. In case when the Metal library was generated before this change, we won't find any delimiters in the source code.

But I added some checks to opencl module:
https://github.com/apache/tvm/pull/7980/files#diff-9d619c5f05402f93afa552f93675dea84d789f8f6b0c553845aa183128f38494R192-R196
Also, I added check in metal module were kernels found or not: https://github.com/apache/tvm/pull/7980/files#diff-dc30040e2ec938d8dd0e6e87d6517870f037093775d4b52e0f249fdbebfdb9a3R95-R101
@csullivancsullivanMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would prefer the metal backward compat be specific to the metal backend and the general SplitKernel util to contain the existence ICHECK on the function delimiter.
 size_t end = begin;
 while (end != std::string::npos) {
 begin += delimiter.size();
 end = source.find('\n', begin);
 std::string func_name = source.substr(begin, end - begin);
 begin = ++end;
 end = source.find(delimiter, begin);
 std::string func_source =
 source.substr(begin, (end == std::string::npos) ? end : end - begin);
 split_kernels.insert({func_name, func_source});
 begin = end;
 }
 }
 return split_kernels;
}
} // namespace runtime
} // namespace tvm

44 changes: 44 additions & 0 deletions
src/runtime/source_utils.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you 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.
 */

/*!
 * \file source_utils.h
 * \brief Minimum source manipulation utils for runtime.
 */

#ifndef TVM_RUNTIME_SOURCE_UTILS_H_
#define TVM_RUNTIME_SOURCE_UTILS_H_

#include <string>
#include <unordered_map>

namespace tvm {
namespace runtime {
/*!
 * \brief Split the source file on separate kernels by specified delimiter.
 * \param source The source code of the kernels.
 * \param delimiter The delimiter which is using for splitting kernels.
 * \return Mapping from primitive name to kernel source
 */
std::unordered_map<std::string, std::string> SplitKernels(std::string source,
 std::string delimiter = "// Function: ");
} // namespace runtime
} // namespace tvm

#endif // TVM_RUNTIME_SOURCE_UTILS_H_

25 changes: 14 additions & 11 deletions
src/target/source/codegen_metal.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -305,27 +305,30 @@ void CodeGenMetal::VisitExpr_(const CallNode* op, std::ostream& os) { // NOLINT
runtime::Module BuildMetal(IRModule mod, Target target) {
 using tvm::runtime::Registry;
 bool output_ssa = false;
 CodeGenMetal cg;
 cg.Init(output_ssa);

 std::stringstream code;
 std::stringstream source;
 std::string fmt = "metal";
 for (auto kv : mod->functions) {
 ICHECK(kv.second->IsInstance<PrimFuncNode>()) << "CodeGenMetal: Can only take PrimFunc";
 code << "// Function: " << kv.first->name_hint << std::endl;
 CodeGenMetal cg;
 cg.Init(output_ssa);
 auto f = Downcast<PrimFunc>(kv.second);
 auto calling_conv = f->GetAttr<Integer>(tvm::attr::kCallingConv);
 ICHECK(calling_conv == CallingConv::kDeviceKernelLaunch)
 << "CodeGenMetal: expect calling_conv equals CallingConv::kDeviceKernelLaunch";
 cg.AddFunction(f);
 std::string fsource = cg.Finish();
 if (const auto* f = Registry::Get("tvm_callback_metal_compile")) {
 source << fsource;
 fsource = (*f)(fsource).operator std::string();
 fmt = "metallib";
 }
 code << fsource;
 }

 std::string code = cg.Finish();
 std::string fmt = "metal";
 std::string source = "";
 if (const auto* f = Registry::Get("tvm_callback_metal_compile")) {
 source = code;
 code = (*f)(code).operator std::string();
 fmt = "metallib";
 }
 return MetalModuleCreate(code, fmt, ExtractFuncInfo(mod), source);
 return MetalModuleCreate(code.str(), fmt, ExtractFuncInfo(mod), source.str());
}

TVM_REGISTER_GLOBAL("target.build.metal").set_body_typed(BuildMetal);
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links");
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '

Skip to content
[METAL] Split kernels and compile them separately#7980

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

tqchen
merged 5 commits into

from

May 25, 2021
 Merged


[METAL] Split kernels and compile them separately
#7980
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


5 commits


Select commit
Hold shift + click to select a range
b0b51fd

[METAL] Split kernels and compile them separately


echuraev May 5, 2021
4938d70

Add backward compatible behaviour


echuraev May 6, 2021
070bba4

Move SplitKernels to a common file


echuraev May 12, 2021
8041a29

Remove using GetSource


echuraev May 12, 2021
df55c33

Update the bundles


echuraev May 13, 2021

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


1 change: 1 addition & 0 deletions
apps/android_camera/app/src/main/jni/tvm_runtime.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -57,6 +57,7 @@
#ifdef TVM_OPENCL_RUNTIME
#include "../src/runtime/opencl/opencl_device_api.cc"
#include "../src/runtime/opencl/opencl_module.cc"
#include "../src/runtime/source_utils.cc"
#endif

#ifdef TVM_VULKAN_RUNTIME
Expand Down

1 change: 1 addition & 0 deletions
apps/android_rpc/app/src/main/jni/tvm_runtime.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -62,6 +62,7 @@
#ifdef TVM_OPENCL_RUNTIME
#include "../src/runtime/opencl/opencl_device_api.cc"
#include "../src/runtime/opencl/opencl_module.cc"
#include "../src/runtime/source_utils.cc"
#endif

#ifdef TVM_VULKAN_RUNTIME
Expand Down

1 change: 1 addition & 0 deletions
golang/src/tvm_runtime_pack.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -68,3 +68,4 @@
// Uncomment the following lines to enable OpenCL
// #include "../../src/runtime/opencl/opencl_device_api.cc"
// #include "../../src/runtime/opencl/opencl_module.cc"
// #include "../src/runtime/source_utils.cc"

79 changes: 45 additions & 34 deletions
src/runtime/metal/metal_module.mm

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -30,6 +30,7 @@
#include "../file_utils.h"
#include "../meta_data.h"
#include "../pack_args.h"
#include "../source_utils.h"
#include "../thread_storage_scope.h"
#include "metal_common.h"

Expand All@@ -43,7 +44,9 @@
 public:
 explicit MetalModuleNode(std::string data, std::string fmt,
 std::unordered_map<std::string, FunctionInfo> fmap, std::string source)
 : data_(data), fmt_(fmt), fmap_(fmap), source_(source) {}
 : data_(data), fmt_(fmt), fmap_(fmap), source_(source) {
 parsed_kernels_ = SplitKernels(data);
 }
 const char* type_key() const final { return "metal"; }

 PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) final;
Expand DownExpand Up@@ -71,6 +74,7 @@ void SaveToBinary(dmlc::Stream* stream) final {
 return "";
 }
 }

 // get a from primary context in device_id
 id<MTLComputePipelineState> GetPipelineState(size_t device_id, const std::string& func_name) {
 metal::MetalWorkspace* w = metal::MetalWorkspace::Global();
Expand All@@ -85,44 +89,52 @@ void SaveToBinary(dmlc::Stream* stream) final {
 if (it != e.smap.end()) return it->second;
 // compile
 NSError* err_msg = nil;
 if (e.lib == nil) {
 if (fmt_ == "metal") {
 MTLCompileOptions* opts = [MTLCompileOptions alloc];
 opts.languageVersion = MTLLanguageVersion2_3;
 opts.fastMathEnabled = YES;
 // opts = nil;
 e.lib = [w->devices[device_id]
 newLibraryWithSource:[NSString stringWithUTF8String:data_.c_str()]
 options:opts
 error:&err_msg];
 [opts dealloc];
 if (e.lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 if (err_msg != nil) {
 LOG(INFO) << "Warning: " << [[err_msg localizedDescription] UTF8String];
 }
 } else {
 // Build from library.
 auto q = dispatch_queue_create("q", DISPATCH_QUEUE_SERIAL);
 auto data = dispatch_data_create(data_.c_str(), data_.length(), q,
 ^{
 });
 e.lib = [w->devices[device_id] newLibraryWithData:data error:&err_msg];
 if (err_msg != nil || e.lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 id<MTLLibrary> lib = nil;
 std::string source;
 auto kernel = parsed_kernels_.find(func_name);
 // If we cannot find this kernel in parsed_kernels_, it means that all kernels going together
 // without explicit separator. In this case we use data_ with all kernels. It done for backward
 // compatibility.
 if (kernel != parsed_kernels_.end())
@elvin-nelvin-nMay 11, 2021•

edited

Loading

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would add here a comment that this behaviour is done for backward compatibility when we had all kernels going together without explicit separators, but looks good
@echuraevechuraevMay 12, 2021

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done
@csullivancsullivanMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we have metal deployments that require the old behavior that can't be rebuilt easily? Otherwise we don't need to preserve backward compatibility.
@elvin-nelvin-nMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I cannot say for sure. I can imagine some use cases which might have issues with recompilation, but this is just assumptions.
If we do not declare backward compatibility for now, we can skip my comment and return original behaviour and rise an exception if we do not find kernel in parsed_kernels_ container.
@csullivancsullivanMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In that case I personally would prefer the metal backward compat be specific to the metal backend and the general SplitKernel util to contain the existence ICHECK on the function delimiter.
@elvin-nelvin-nMay 14, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tqchen do we care about backward compatibility or not? If not, it might be ok so far, but we need to decide when we will start to care and what use cases should work and be regularly verified in the future. It must not be a deal only metal backend
@csullivan what bahavior do you propose in common SplitKernel? Metal backend uses it, it cannot just abort execution. Warnings in the runtime are also useles. It might be an error code, but empty returned map detects the absent of Function delimiter with the same efficiency as return code
 source = kernel->second;
 else
 source = data_;
 if (fmt_ == "metal") {
 MTLCompileOptions* opts = [MTLCompileOptions alloc];
 opts.languageVersion = MTLLanguageVersion2_3;
 opts.fastMathEnabled = YES;
 // opts = nil;
 lib =
 [w->devices[device_id] newLibraryWithSource:[NSString stringWithUTF8String:source.c_str()]
 options:opts
 error:&err_msg];
 [opts dealloc];
 if (lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 if (err_msg != nil) {
 LOG(INFO) << "Warning: " << [[err_msg localizedDescription] UTF8String];
 }
 } else {
 // Build from library.
 auto q = dispatch_queue_create("q", DISPATCH_QUEUE_SERIAL);
 auto data = dispatch_data_create(source.c_str(), source.length(), q,
 ^{
 });
 lib = [w->devices[device_id] newLibraryWithData:data error:&err_msg];
 if (err_msg != nil || lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 }
 id<MTLFunction> f =
 [e.lib newFunctionWithName:[NSString stringWithUTF8String:func_name.c_str()]];
 id<MTLFunction> f = [lib newFunctionWithName:[NSString stringWithUTF8String:func_name.c_str()]];
 ICHECK(f != nil) << "cannot find function " << func_name;
 id<MTLComputePipelineState> state =
 [w->devices[device_id] newComputePipelineStateWithFunction:f error:&err_msg];
 ICHECK(state != nil) << "cannot get state:"
 << " for function " << func_name
 << [[err_msg localizedDescription] UTF8String];
 [f release];
 [lib release];
 // The state.threadExecutionWidth can change dynamically according
 // to the resource constraint in kernel, so it is not strictly hold
 // Turn of warp aware optimziation for now.
Expand All@@ -135,13 +147,10 @@ void SaveToBinary(dmlc::Stream* stream) final {
 private:
 // device specific entry
 struct DeviceEntry {
 // library
 id<MTLLibrary> lib = nil;
 // state cache;
 std::unordered_map<std::string, id<MTLComputePipelineState>> smap;
 std::unordered_map<std::string, id<MTLComputePipelineState>> smap;

 ~DeviceEntry() {
 if (lib != nil) [lib release];
 for (auto&& kv : smap) {
 [kv.second release];
 }
Expand All@@ -159,6 +168,8 @@ void SaveToBinary(dmlc::Stream* stream) final {
 std::vector<DeviceEntry> finfo_;
 // internal mutex when updating the module
 std::mutex mutex_;
 // parsed kernel data
 std::unordered_map<std::string, std::string> parsed_kernels_;
};

// a wrapped function class to get packed func.
Expand Down

8 changes: 0 additions & 8 deletions
src/runtime/opencl/opencl_common.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -326,14 +326,6 @@ class OpenCLModuleNode : public ModuleNode {
 cl_kernel InstallKernel(cl::OpenCLWorkspace* w, cl::OpenCLThreadEntry* t,
 const std::string& func_name, const KTRefEntry& e);

 /*
 * \brief Splits the provided serialized source file into separate
 * source for each kernel primitive.
 * \param source The serialized program source file (fmt: cl)
 * \return Mapping from primitive name to kernel source
 */
 std::unordered_map<std::string, std::string> SplitKernels(std::string source) const;

 private:
 // The workspace, need to keep reference to use it in destructor.
 // In case of static destruction order problem.
Expand Down

39 changes: 6 additions & 33 deletions
src/runtime/opencl/opencl_module.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -29,6 +29,7 @@
#include <unordered_map>
#include <vector>

#include "../source_utils.h"
#include "opencl_common.h"

namespace tvm {
Expand DownExpand Up@@ -188,6 +189,11 @@ void OpenCLModuleNode::Init() {

 // split into source artifacts for each kernel
 parsed_kernels_ = SplitKernels(GetSource("cl"));
 ICHECK(!parsed_kernels_.empty()) << "The OpenCL module expects a kernel delimited "
 << "source from code generation, but no kernel "
 << "delimiter was found.";
 ICHECK_EQ(workspace_->num_registered_kernels, parsed_kernels_.size())
 << "The number of registered kernels does not match number of parsed kernel sources";
 // zero initialize cl_program pointers for each device kernel
 for (auto& kv : parsed_kernels_) {
 programs_.insert({kv.first, std::vector<cl_program>(workspace_->devices.size(), nullptr)});
Expand DownExpand Up@@ -242,39 +248,6 @@ cl_kernel OpenCLModuleNode::InstallKernel(cl::OpenCLWorkspace* w, cl::OpenCLThre
 return kernel;
}

std::unordered_map<std::string, std::string> OpenCLModuleNode::SplitKernels(
 std::string source) const {
 std::unordered_map<std::string, std::string> split_kernels;
 if (source.size()) {
 std::string del{"// Function: "};
 size_t end;
 size_t begin = source.find(del);
 ICHECK(begin != std::string::npos) << "The OpenCL module expects a kernel delimited "
 << "source from code generation, but no kernel "
 << "delimiter was found.";
 for (size_t num_kernels = 0; num_kernels < workspace_->num_registered_kernels; num_kernels++) {
 begin += del.size();
 end = source.find('\n', begin);
 std::string func_name = source.substr(begin, end - begin);
 begin = ++end;
 // std::string::substr returns either start of next kernel
 // or std::string::npos, in the latter case substr returns
 // all characters until the end of the source string.
 end = source.find(del, begin);
 std::string func_source =
 source.substr(begin, (end == std::string::npos) ? end : end - begin);
 split_kernels.insert({func_name, func_source});
 begin = end;
 if (end == std::string::npos) {
 break;
 }
 }
 }
 ICHECK_EQ(workspace_->num_registered_kernels, split_kernels.size())
 << "The number of registered kernels does not match number of parsed kernel sources";
 return split_kernels;
}

Module OpenCLModuleCreate(std::string data, std::string fmt,
 std::unordered_map<std::string, FunctionInfo> fmap, std::string source) {
 auto n = make_object<OpenCLModuleNode>(data, fmt, fmap, source);
Expand Down

49 changes: 49 additions & 0 deletions
src/runtime/source_utils.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you 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.
 */

/*!
 * \file source_utils.cc
 */
#include "source_utils.h"

namespace tvm {
namespace runtime {

std::unordered_map<std::string, std::string> SplitKernels(std::string source,
Comment thread
echuraev marked this conversation as resolved.

 std::string delimiter) {
 std::unordered_map<std::string, std::string> split_kernels;
 if (source.size()) {
 size_t begin = source.find(delimiter);
@csullivancsullivanMay 12, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I noticed you deleted ICHECK here, probably because the error message referred to OpenCL specifically, but there still should be an
 ICHECK(begin != std::string::npos) 
on the first search for a kernel delimiter. If none are found, then it is better to catch an error in the code generation early.
@echuraevechuraevMay 13, 2021

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I did it to keep backward compatibility as @elvin-nnov requested. In case when the Metal library was generated before this change, we won't find any delimiters in the source code.

But I added some checks to opencl module:
https://github.com/apache/tvm/pull/7980/files#diff-9d619c5f05402f93afa552f93675dea84d789f8f6b0c553845aa183128f38494R192-R196
Also, I added check in metal module were kernels found or not: https://github.com/apache/tvm/pull/7980/files#diff-dc30040e2ec938d8dd0e6e87d6517870f037093775d4b52e0f249fdbebfdb9a3R95-R101
@csullivancsullivanMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would prefer the metal backward compat be specific to the metal backend and the general SplitKernel util to contain the existence ICHECK on the function delimiter.
 size_t end = begin;
 while (end != std::string::npos) {
 begin += delimiter.size();
 end = source.find('\n', begin);
 std::string func_name = source.substr(begin, end - begin);
 begin = ++end;
 end = source.find(delimiter, begin);
 std::string func_source =
 source.substr(begin, (end == std::string::npos) ? end : end - begin);
 split_kernels.insert({func_name, func_source});
 begin = end;
 }
 }
 return split_kernels;
}
} // namespace runtime
} // namespace tvm

44 changes: 44 additions & 0 deletions
src/runtime/source_utils.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you 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.
 */

/*!
 * \file source_utils.h
 * \brief Minimum source manipulation utils for runtime.
 */

#ifndef TVM_RUNTIME_SOURCE_UTILS_H_
#define TVM_RUNTIME_SOURCE_UTILS_H_

#include <string>
#include <unordered_map>

namespace tvm {
namespace runtime {
/*!
 * \brief Split the source file on separate kernels by specified delimiter.
 * \param source The source code of the kernels.
 * \param delimiter The delimiter which is using for splitting kernels.
 * \return Mapping from primitive name to kernel source
 */
std::unordered_map<std::string, std::string> SplitKernels(std::string source,
 std::string delimiter = "// Function: ");
} // namespace runtime
} // namespace tvm

#endif // TVM_RUNTIME_SOURCE_UTILS_H_

25 changes: 14 additions & 11 deletions
src/target/source/codegen_metal.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -305,27 +305,30 @@ void CodeGenMetal::VisitExpr_(const CallNode* op, std::ostream& os) { // NOLINT
runtime::Module BuildMetal(IRModule mod, Target target) {
 using tvm::runtime::Registry;
 bool output_ssa = false;
 CodeGenMetal cg;
 cg.Init(output_ssa);

 std::stringstream code;
 std::stringstream source;
 std::string fmt = "metal";
 for (auto kv : mod->functions) {
 ICHECK(kv.second->IsInstance<PrimFuncNode>()) << "CodeGenMetal: Can only take PrimFunc";
 code << "// Function: " << kv.first->name_hint << std::endl;
 CodeGenMetal cg;
 cg.Init(output_ssa);
 auto f = Downcast<PrimFunc>(kv.second);
 auto calling_conv = f->GetAttr<Integer>(tvm::attr::kCallingConv);
 ICHECK(calling_conv == CallingConv::kDeviceKernelLaunch)
 << "CodeGenMetal: expect calling_conv equals CallingConv::kDeviceKernelLaunch";
 cg.AddFunction(f);
 std::string fsource = cg.Finish();
 if (const auto* f = Registry::Get("tvm_callback_metal_compile")) {
 source << fsource;
 fsource = (*f)(fsource).operator std::string();
 fmt = "metallib";
 }
 code << fsource;
 }

 std::string code = cg.Finish();
 std::string fmt = "metal";
 std::string source = "";
 if (const auto* f = Registry::Get("tvm_callback_metal_compile")) {
 source = code;
 code = (*f)(code).operator std::string();
 fmt = "metallib";
 }
 return MetalModuleCreate(code, fmt, ExtractFuncInfo(mod), source);
 return MetalModuleCreate(code.str(), fmt, ExtractFuncInfo(mod), source.str());
}

TVM_REGISTER_GLOBAL("target.build.metal").set_body_typed(BuildMetal);
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Auto-enable theater mode on YouTube\n(function() {\n function tryTheater() {\n var btn = document.querySelector('button[aria-label=\"Theater mode\"], ytd-player #player button[title=\"Theater mode\"]');\n if (btn && !btn.classList.contains('activated')) {\n btn.click();\n }\n }\n \n // Try immediately\n tryTheater();\n \n // Try after navigation (SPA)\n var lastUrl = location.href;\n setInterval(function() {\n if (location.href !== lastUrl) {\n lastUrl = location.href;\n setTimeout(tryTheater, 500);\n }\n }, 1000);\n \n // Also try on player load\n var observer = new MutationObserver(tryTheater);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "YouTube Theater Mode Default");
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
[METAL] Split kernels and compile them separately#7980

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

tqchen
merged 5 commits into

from

May 25, 2021
 Merged


[METAL] Split kernels and compile them separately
#7980
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


5 commits


Select commit
Hold shift + click to select a range
b0b51fd

[METAL] Split kernels and compile them separately


echuraev May 5, 2021
4938d70

Add backward compatible behaviour


echuraev May 6, 2021
070bba4

Move SplitKernels to a common file


echuraev May 12, 2021
8041a29

Remove using GetSource


echuraev May 12, 2021
df55c33

Update the bundles


echuraev May 13, 2021

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


1 change: 1 addition & 0 deletions
apps/android_camera/app/src/main/jni/tvm_runtime.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -57,6 +57,7 @@
#ifdef TVM_OPENCL_RUNTIME
#include "../src/runtime/opencl/opencl_device_api.cc"
#include "../src/runtime/opencl/opencl_module.cc"
#include "../src/runtime/source_utils.cc"
#endif

#ifdef TVM_VULKAN_RUNTIME
Expand Down

1 change: 1 addition & 0 deletions
apps/android_rpc/app/src/main/jni/tvm_runtime.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -62,6 +62,7 @@
#ifdef TVM_OPENCL_RUNTIME
#include "../src/runtime/opencl/opencl_device_api.cc"
#include "../src/runtime/opencl/opencl_module.cc"
#include "../src/runtime/source_utils.cc"
#endif

#ifdef TVM_VULKAN_RUNTIME
Expand Down

1 change: 1 addition & 0 deletions
golang/src/tvm_runtime_pack.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -68,3 +68,4 @@
// Uncomment the following lines to enable OpenCL
// #include "../../src/runtime/opencl/opencl_device_api.cc"
// #include "../../src/runtime/opencl/opencl_module.cc"
// #include "../src/runtime/source_utils.cc"

79 changes: 45 additions & 34 deletions
src/runtime/metal/metal_module.mm

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -30,6 +30,7 @@
#include "../file_utils.h"
#include "../meta_data.h"
#include "../pack_args.h"
#include "../source_utils.h"
#include "../thread_storage_scope.h"
#include "metal_common.h"

Expand All@@ -43,7 +44,9 @@
 public:
 explicit MetalModuleNode(std::string data, std::string fmt,
 std::unordered_map<std::string, FunctionInfo> fmap, std::string source)
 : data_(data), fmt_(fmt), fmap_(fmap), source_(source) {}
 : data_(data), fmt_(fmt), fmap_(fmap), source_(source) {
 parsed_kernels_ = SplitKernels(data);
 }
 const char* type_key() const final { return "metal"; }

 PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) final;
Expand DownExpand Up@@ -71,6 +74,7 @@ void SaveToBinary(dmlc::Stream* stream) final {
 return "";
 }
 }

 // get a from primary context in device_id
 id<MTLComputePipelineState> GetPipelineState(size_t device_id, const std::string& func_name) {
 metal::MetalWorkspace* w = metal::MetalWorkspace::Global();
Expand All@@ -85,44 +89,52 @@ void SaveToBinary(dmlc::Stream* stream) final {
 if (it != e.smap.end()) return it->second;
 // compile
 NSError* err_msg = nil;
 if (e.lib == nil) {
 if (fmt_ == "metal") {
 MTLCompileOptions* opts = [MTLCompileOptions alloc];
 opts.languageVersion = MTLLanguageVersion2_3;
 opts.fastMathEnabled = YES;
 // opts = nil;
 e.lib = [w->devices[device_id]
 newLibraryWithSource:[NSString stringWithUTF8String:data_.c_str()]
 options:opts
 error:&err_msg];
 [opts dealloc];
 if (e.lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 if (err_msg != nil) {
 LOG(INFO) << "Warning: " << [[err_msg localizedDescription] UTF8String];
 }
 } else {
 // Build from library.
 auto q = dispatch_queue_create("q", DISPATCH_QUEUE_SERIAL);
 auto data = dispatch_data_create(data_.c_str(), data_.length(), q,
 ^{
 });
 e.lib = [w->devices[device_id] newLibraryWithData:data error:&err_msg];
 if (err_msg != nil || e.lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 id<MTLLibrary> lib = nil;
 std::string source;
 auto kernel = parsed_kernels_.find(func_name);
 // If we cannot find this kernel in parsed_kernels_, it means that all kernels going together
 // without explicit separator. In this case we use data_ with all kernels. It done for backward
 // compatibility.
 if (kernel != parsed_kernels_.end())
@elvin-nelvin-nMay 11, 2021•

edited

Loading

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would add here a comment that this behaviour is done for backward compatibility when we had all kernels going together without explicit separators, but looks good
@echuraevechuraevMay 12, 2021

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done
@csullivancsullivanMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we have metal deployments that require the old behavior that can't be rebuilt easily? Otherwise we don't need to preserve backward compatibility.
@elvin-nelvin-nMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I cannot say for sure. I can imagine some use cases which might have issues with recompilation, but this is just assumptions.
If we do not declare backward compatibility for now, we can skip my comment and return original behaviour and rise an exception if we do not find kernel in parsed_kernels_ container.
@csullivancsullivanMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In that case I personally would prefer the metal backward compat be specific to the metal backend and the general SplitKernel util to contain the existence ICHECK on the function delimiter.
@elvin-nelvin-nMay 14, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tqchen do we care about backward compatibility or not? If not, it might be ok so far, but we need to decide when we will start to care and what use cases should work and be regularly verified in the future. It must not be a deal only metal backend
@csullivan what bahavior do you propose in common SplitKernel? Metal backend uses it, it cannot just abort execution. Warnings in the runtime are also useles. It might be an error code, but empty returned map detects the absent of Function delimiter with the same efficiency as return code
 source = kernel->second;
 else
 source = data_;
 if (fmt_ == "metal") {
 MTLCompileOptions* opts = [MTLCompileOptions alloc];
 opts.languageVersion = MTLLanguageVersion2_3;
 opts.fastMathEnabled = YES;
 // opts = nil;
 lib =
 [w->devices[device_id] newLibraryWithSource:[NSString stringWithUTF8String:source.c_str()]
 options:opts
 error:&err_msg];
 [opts dealloc];
 if (lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 if (err_msg != nil) {
 LOG(INFO) << "Warning: " << [[err_msg localizedDescription] UTF8String];
 }
 } else {
 // Build from library.
 auto q = dispatch_queue_create("q", DISPATCH_QUEUE_SERIAL);
 auto data = dispatch_data_create(source.c_str(), source.length(), q,
 ^{
 });
 lib = [w->devices[device_id] newLibraryWithData:data error:&err_msg];
 if (err_msg != nil || lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 }
 id<MTLFunction> f =
 [e.lib newFunctionWithName:[NSString stringWithUTF8String:func_name.c_str()]];
 id<MTLFunction> f = [lib newFunctionWithName:[NSString stringWithUTF8String:func_name.c_str()]];
 ICHECK(f != nil) << "cannot find function " << func_name;
 id<MTLComputePipelineState> state =
 [w->devices[device_id] newComputePipelineStateWithFunction:f error:&err_msg];
 ICHECK(state != nil) << "cannot get state:"
 << " for function " << func_name
 << [[err_msg localizedDescription] UTF8String];
 [f release];
 [lib release];
 // The state.threadExecutionWidth can change dynamically according
 // to the resource constraint in kernel, so it is not strictly hold
 // Turn of warp aware optimziation for now.
Expand All@@ -135,13 +147,10 @@ void SaveToBinary(dmlc::Stream* stream) final {
 private:
 // device specific entry
 struct DeviceEntry {
 // library
 id<MTLLibrary> lib = nil;
 // state cache;
 std::unordered_map<std::string, id<MTLComputePipelineState>> smap;
 std::unordered_map<std::string, id<MTLComputePipelineState>> smap;

 ~DeviceEntry() {
 if (lib != nil) [lib release];
 for (auto&& kv : smap) {
 [kv.second release];
 }
Expand All@@ -159,6 +168,8 @@ void SaveToBinary(dmlc::Stream* stream) final {
 std::vector<DeviceEntry> finfo_;
 // internal mutex when updating the module
 std::mutex mutex_;
 // parsed kernel data
 std::unordered_map<std::string, std::string> parsed_kernels_;
};

// a wrapped function class to get packed func.
Expand Down

8 changes: 0 additions & 8 deletions
src/runtime/opencl/opencl_common.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -326,14 +326,6 @@ class OpenCLModuleNode : public ModuleNode {
 cl_kernel InstallKernel(cl::OpenCLWorkspace* w, cl::OpenCLThreadEntry* t,
 const std::string& func_name, const KTRefEntry& e);

 /*
 * \brief Splits the provided serialized source file into separate
 * source for each kernel primitive.
 * \param source The serialized program source file (fmt: cl)
 * \return Mapping from primitive name to kernel source
 */
 std::unordered_map<std::string, std::string> SplitKernels(std::string source) const;

 private:
 // The workspace, need to keep reference to use it in destructor.
 // In case of static destruction order problem.
Expand Down

39 changes: 6 additions & 33 deletions
src/runtime/opencl/opencl_module.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -29,6 +29,7 @@
#include <unordered_map>
#include <vector>

#include "../source_utils.h"
#include "opencl_common.h"

namespace tvm {
Expand DownExpand Up@@ -188,6 +189,11 @@ void OpenCLModuleNode::Init() {

 // split into source artifacts for each kernel
 parsed_kernels_ = SplitKernels(GetSource("cl"));
 ICHECK(!parsed_kernels_.empty()) << "The OpenCL module expects a kernel delimited "
 << "source from code generation, but no kernel "
 << "delimiter was found.";
 ICHECK_EQ(workspace_->num_registered_kernels, parsed_kernels_.size())
 << "The number of registered kernels does not match number of parsed kernel sources";
 // zero initialize cl_program pointers for each device kernel
 for (auto& kv : parsed_kernels_) {
 programs_.insert({kv.first, std::vector<cl_program>(workspace_->devices.size(), nullptr)});
Expand DownExpand Up@@ -242,39 +248,6 @@ cl_kernel OpenCLModuleNode::InstallKernel(cl::OpenCLWorkspace* w, cl::OpenCLThre
 return kernel;
}

std::unordered_map<std::string, std::string> OpenCLModuleNode::SplitKernels(
 std::string source) const {
 std::unordered_map<std::string, std::string> split_kernels;
 if (source.size()) {
 std::string del{"// Function: "};
 size_t end;
 size_t begin = source.find(del);
 ICHECK(begin != std::string::npos) << "The OpenCL module expects a kernel delimited "
 << "source from code generation, but no kernel "
 << "delimiter was found.";
 for (size_t num_kernels = 0; num_kernels < workspace_->num_registered_kernels; num_kernels++) {
 begin += del.size();
 end = source.find('\n', begin);
 std::string func_name = source.substr(begin, end - begin);
 begin = ++end;
 // std::string::substr returns either start of next kernel
 // or std::string::npos, in the latter case substr returns
 // all characters until the end of the source string.
 end = source.find(del, begin);
 std::string func_source =
 source.substr(begin, (end == std::string::npos) ? end : end - begin);
 split_kernels.insert({func_name, func_source});
 begin = end;
 if (end == std::string::npos) {
 break;
 }
 }
 }
 ICHECK_EQ(workspace_->num_registered_kernels, split_kernels.size())
 << "The number of registered kernels does not match number of parsed kernel sources";
 return split_kernels;
}

Module OpenCLModuleCreate(std::string data, std::string fmt,
 std::unordered_map<std::string, FunctionInfo> fmap, std::string source) {
 auto n = make_object<OpenCLModuleNode>(data, fmt, fmap, source);
Expand Down

49 changes: 49 additions & 0 deletions
src/runtime/source_utils.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you 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.
 */

/*!
 * \file source_utils.cc
 */
#include "source_utils.h"

namespace tvm {
namespace runtime {

std::unordered_map<std::string, std::string> SplitKernels(std::string source,
Comment thread
echuraev marked this conversation as resolved.

 std::string delimiter) {
 std::unordered_map<std::string, std::string> split_kernels;
 if (source.size()) {
 size_t begin = source.find(delimiter);
@csullivancsullivanMay 12, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I noticed you deleted ICHECK here, probably because the error message referred to OpenCL specifically, but there still should be an
 ICHECK(begin != std::string::npos) 
on the first search for a kernel delimiter. If none are found, then it is better to catch an error in the code generation early.
@echuraevechuraevMay 13, 2021

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I did it to keep backward compatibility as @elvin-nnov requested. In case when the Metal library was generated before this change, we won't find any delimiters in the source code.

But I added some checks to opencl module:
https://github.com/apache/tvm/pull/7980/files#diff-9d619c5f05402f93afa552f93675dea84d789f8f6b0c553845aa183128f38494R192-R196
Also, I added check in metal module were kernels found or not: https://github.com/apache/tvm/pull/7980/files#diff-dc30040e2ec938d8dd0e6e87d6517870f037093775d4b52e0f249fdbebfdb9a3R95-R101
@csullivancsullivanMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would prefer the metal backward compat be specific to the metal backend and the general SplitKernel util to contain the existence ICHECK on the function delimiter.
 size_t end = begin;
 while (end != std::string::npos) {
 begin += delimiter.size();
 end = source.find('\n', begin);
 std::string func_name = source.substr(begin, end - begin);
 begin = ++end;
 end = source.find(delimiter, begin);
 std::string func_source =
 source.substr(begin, (end == std::string::npos) ? end : end - begin);
 split_kernels.insert({func_name, func_source});
 begin = end;
 }
 }
 return split_kernels;
}
} // namespace runtime
} // namespace tvm

44 changes: 44 additions & 0 deletions
src/runtime/source_utils.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you 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.
 */

/*!
 * \file source_utils.h
 * \brief Minimum source manipulation utils for runtime.
 */

#ifndef TVM_RUNTIME_SOURCE_UTILS_H_
#define TVM_RUNTIME_SOURCE_UTILS_H_

#include <string>
#include <unordered_map>

namespace tvm {
namespace runtime {
/*!
 * \brief Split the source file on separate kernels by specified delimiter.
 * \param source The source code of the kernels.
 * \param delimiter The delimiter which is using for splitting kernels.
 * \return Mapping from primitive name to kernel source
 */
std::unordered_map<std::string, std::string> SplitKernels(std::string source,
 std::string delimiter = "// Function: ");
} // namespace runtime
} // namespace tvm

#endif // TVM_RUNTIME_SOURCE_UTILS_H_

25 changes: 14 additions & 11 deletions
src/target/source/codegen_metal.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -305,27 +305,30 @@ void CodeGenMetal::VisitExpr_(const CallNode* op, std::ostream& os) { // NOLINT
runtime::Module BuildMetal(IRModule mod, Target target) {
 using tvm::runtime::Registry;
 bool output_ssa = false;
 CodeGenMetal cg;
 cg.Init(output_ssa);

 std::stringstream code;
 std::stringstream source;
 std::string fmt = "metal";
 for (auto kv : mod->functions) {
 ICHECK(kv.second->IsInstance<PrimFuncNode>()) << "CodeGenMetal: Can only take PrimFunc";
 code << "// Function: " << kv.first->name_hint << std::endl;
 CodeGenMetal cg;
 cg.Init(output_ssa);
 auto f = Downcast<PrimFunc>(kv.second);
 auto calling_conv = f->GetAttr<Integer>(tvm::attr::kCallingConv);
 ICHECK(calling_conv == CallingConv::kDeviceKernelLaunch)
 << "CodeGenMetal: expect calling_conv equals CallingConv::kDeviceKernelLaunch";
 cg.AddFunction(f);
 std::string fsource = cg.Finish();
 if (const auto* f = Registry::Get("tvm_callback_metal_compile")) {
 source << fsource;
 fsource = (*f)(fsource).operator std::string();
 fmt = "metallib";
 }
 code << fsource;
 }

 std::string code = cg.Finish();
 std::string fmt = "metal";
 std::string source = "";
 if (const auto* f = Registry::Get("tvm_callback_metal_compile")) {
 source = code;
 code = (*f)(code).operator std::string();
 fmt = "metallib";
 }
 return MetalModuleCreate(code, fmt, ExtractFuncInfo(mod), source);
 return MetalModuleCreate(code.str(), fmt, ExtractFuncInfo(mod), source.str());
}

TVM_REGISTER_GLOBAL("target.build.metal").set_body_typed(BuildMetal);
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers");
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
[METAL] Split kernels and compile them separately#7980

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

tqchen
merged 5 commits into

from

May 25, 2021
 Merged


[METAL] Split kernels and compile them separately
#7980
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


5 commits


Select commit
Hold shift + click to select a range
b0b51fd

[METAL] Split kernels and compile them separately


echuraev May 5, 2021
4938d70

Add backward compatible behaviour


echuraev May 6, 2021
070bba4

Move SplitKernels to a common file


echuraev May 12, 2021
8041a29

Remove using GetSource


echuraev May 12, 2021
df55c33

Update the bundles


echuraev May 13, 2021

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


1 change: 1 addition & 0 deletions
apps/android_camera/app/src/main/jni/tvm_runtime.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -57,6 +57,7 @@
#ifdef TVM_OPENCL_RUNTIME
#include "../src/runtime/opencl/opencl_device_api.cc"
#include "../src/runtime/opencl/opencl_module.cc"
#include "../src/runtime/source_utils.cc"
#endif

#ifdef TVM_VULKAN_RUNTIME
Expand Down

1 change: 1 addition & 0 deletions
apps/android_rpc/app/src/main/jni/tvm_runtime.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -62,6 +62,7 @@
#ifdef TVM_OPENCL_RUNTIME
#include "../src/runtime/opencl/opencl_device_api.cc"
#include "../src/runtime/opencl/opencl_module.cc"
#include "../src/runtime/source_utils.cc"
#endif

#ifdef TVM_VULKAN_RUNTIME
Expand Down

1 change: 1 addition & 0 deletions
golang/src/tvm_runtime_pack.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -68,3 +68,4 @@
// Uncomment the following lines to enable OpenCL
// #include "../../src/runtime/opencl/opencl_device_api.cc"
// #include "../../src/runtime/opencl/opencl_module.cc"
// #include "../src/runtime/source_utils.cc"

79 changes: 45 additions & 34 deletions
src/runtime/metal/metal_module.mm

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -30,6 +30,7 @@
#include "../file_utils.h"
#include "../meta_data.h"
#include "../pack_args.h"
#include "../source_utils.h"
#include "../thread_storage_scope.h"
#include "metal_common.h"

Expand All@@ -43,7 +44,9 @@
 public:
 explicit MetalModuleNode(std::string data, std::string fmt,
 std::unordered_map<std::string, FunctionInfo> fmap, std::string source)
 : data_(data), fmt_(fmt), fmap_(fmap), source_(source) {}
 : data_(data), fmt_(fmt), fmap_(fmap), source_(source) {
 parsed_kernels_ = SplitKernels(data);
 }
 const char* type_key() const final { return "metal"; }

 PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) final;
Expand DownExpand Up@@ -71,6 +74,7 @@ void SaveToBinary(dmlc::Stream* stream) final {
 return "";
 }
 }

 // get a from primary context in device_id
 id<MTLComputePipelineState> GetPipelineState(size_t device_id, const std::string& func_name) {
 metal::MetalWorkspace* w = metal::MetalWorkspace::Global();
Expand All@@ -85,44 +89,52 @@ void SaveToBinary(dmlc::Stream* stream) final {
 if (it != e.smap.end()) return it->second;
 // compile
 NSError* err_msg = nil;
 if (e.lib == nil) {
 if (fmt_ == "metal") {
 MTLCompileOptions* opts = [MTLCompileOptions alloc];
 opts.languageVersion = MTLLanguageVersion2_3;
 opts.fastMathEnabled = YES;
 // opts = nil;
 e.lib = [w->devices[device_id]
 newLibraryWithSource:[NSString stringWithUTF8String:data_.c_str()]
 options:opts
 error:&err_msg];
 [opts dealloc];
 if (e.lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 if (err_msg != nil) {
 LOG(INFO) << "Warning: " << [[err_msg localizedDescription] UTF8String];
 }
 } else {
 // Build from library.
 auto q = dispatch_queue_create("q", DISPATCH_QUEUE_SERIAL);
 auto data = dispatch_data_create(data_.c_str(), data_.length(), q,
 ^{
 });
 e.lib = [w->devices[device_id] newLibraryWithData:data error:&err_msg];
 if (err_msg != nil || e.lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 id<MTLLibrary> lib = nil;
 std::string source;
 auto kernel = parsed_kernels_.find(func_name);
 // If we cannot find this kernel in parsed_kernels_, it means that all kernels going together
 // without explicit separator. In this case we use data_ with all kernels. It done for backward
 // compatibility.
 if (kernel != parsed_kernels_.end())
@elvin-nelvin-nMay 11, 2021•

edited

Loading

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would add here a comment that this behaviour is done for backward compatibility when we had all kernels going together without explicit separators, but looks good
@echuraevechuraevMay 12, 2021

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done
@csullivancsullivanMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we have metal deployments that require the old behavior that can't be rebuilt easily? Otherwise we don't need to preserve backward compatibility.
@elvin-nelvin-nMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I cannot say for sure. I can imagine some use cases which might have issues with recompilation, but this is just assumptions.
If we do not declare backward compatibility for now, we can skip my comment and return original behaviour and rise an exception if we do not find kernel in parsed_kernels_ container.
@csullivancsullivanMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In that case I personally would prefer the metal backward compat be specific to the metal backend and the general SplitKernel util to contain the existence ICHECK on the function delimiter.
@elvin-nelvin-nMay 14, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tqchen do we care about backward compatibility or not? If not, it might be ok so far, but we need to decide when we will start to care and what use cases should work and be regularly verified in the future. It must not be a deal only metal backend
@csullivan what bahavior do you propose in common SplitKernel? Metal backend uses it, it cannot just abort execution. Warnings in the runtime are also useles. It might be an error code, but empty returned map detects the absent of Function delimiter with the same efficiency as return code
 source = kernel->second;
 else
 source = data_;
 if (fmt_ == "metal") {
 MTLCompileOptions* opts = [MTLCompileOptions alloc];
 opts.languageVersion = MTLLanguageVersion2_3;
 opts.fastMathEnabled = YES;
 // opts = nil;
 lib =
 [w->devices[device_id] newLibraryWithSource:[NSString stringWithUTF8String:source.c_str()]
 options:opts
 error:&err_msg];
 [opts dealloc];
 if (lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 if (err_msg != nil) {
 LOG(INFO) << "Warning: " << [[err_msg localizedDescription] UTF8String];
 }
 } else {
 // Build from library.
 auto q = dispatch_queue_create("q", DISPATCH_QUEUE_SERIAL);
 auto data = dispatch_data_create(source.c_str(), source.length(), q,
 ^{
 });
 lib = [w->devices[device_id] newLibraryWithData:data error:&err_msg];
 if (err_msg != nil || lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 }
 id<MTLFunction> f =
 [e.lib newFunctionWithName:[NSString stringWithUTF8String:func_name.c_str()]];
 id<MTLFunction> f = [lib newFunctionWithName:[NSString stringWithUTF8String:func_name.c_str()]];
 ICHECK(f != nil) << "cannot find function " << func_name;
 id<MTLComputePipelineState> state =
 [w->devices[device_id] newComputePipelineStateWithFunction:f error:&err_msg];
 ICHECK(state != nil) << "cannot get state:"
 << " for function " << func_name
 << [[err_msg localizedDescription] UTF8String];
 [f release];
 [lib release];
 // The state.threadExecutionWidth can change dynamically according
 // to the resource constraint in kernel, so it is not strictly hold
 // Turn of warp aware optimziation for now.
Expand All@@ -135,13 +147,10 @@ void SaveToBinary(dmlc::Stream* stream) final {
 private:
 // device specific entry
 struct DeviceEntry {
 // library
 id<MTLLibrary> lib = nil;
 // state cache;
 std::unordered_map<std::string, id<MTLComputePipelineState>> smap;
 std::unordered_map<std::string, id<MTLComputePipelineState>> smap;

 ~DeviceEntry() {
 if (lib != nil) [lib release];
 for (auto&& kv : smap) {
 [kv.second release];
 }
Expand All@@ -159,6 +168,8 @@ void SaveToBinary(dmlc::Stream* stream) final {
 std::vector<DeviceEntry> finfo_;
 // internal mutex when updating the module
 std::mutex mutex_;
 // parsed kernel data
 std::unordered_map<std::string, std::string> parsed_kernels_;
};

// a wrapped function class to get packed func.
Expand Down

8 changes: 0 additions & 8 deletions
src/runtime/opencl/opencl_common.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -326,14 +326,6 @@ class OpenCLModuleNode : public ModuleNode {
 cl_kernel InstallKernel(cl::OpenCLWorkspace* w, cl::OpenCLThreadEntry* t,
 const std::string& func_name, const KTRefEntry& e);

 /*
 * \brief Splits the provided serialized source file into separate
 * source for each kernel primitive.
 * \param source The serialized program source file (fmt: cl)
 * \return Mapping from primitive name to kernel source
 */
 std::unordered_map<std::string, std::string> SplitKernels(std::string source) const;

 private:
 // The workspace, need to keep reference to use it in destructor.
 // In case of static destruction order problem.
Expand Down

39 changes: 6 additions & 33 deletions
src/runtime/opencl/opencl_module.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -29,6 +29,7 @@
#include <unordered_map>
#include <vector>

#include "../source_utils.h"
#include "opencl_common.h"

namespace tvm {
Expand DownExpand Up@@ -188,6 +189,11 @@ void OpenCLModuleNode::Init() {

 // split into source artifacts for each kernel
 parsed_kernels_ = SplitKernels(GetSource("cl"));
 ICHECK(!parsed_kernels_.empty()) << "The OpenCL module expects a kernel delimited "
 << "source from code generation, but no kernel "
 << "delimiter was found.";
 ICHECK_EQ(workspace_->num_registered_kernels, parsed_kernels_.size())
 << "The number of registered kernels does not match number of parsed kernel sources";
 // zero initialize cl_program pointers for each device kernel
 for (auto& kv : parsed_kernels_) {
 programs_.insert({kv.first, std::vector<cl_program>(workspace_->devices.size(), nullptr)});
Expand DownExpand Up@@ -242,39 +248,6 @@ cl_kernel OpenCLModuleNode::InstallKernel(cl::OpenCLWorkspace* w, cl::OpenCLThre
 return kernel;
}

std::unordered_map<std::string, std::string> OpenCLModuleNode::SplitKernels(
 std::string source) const {
 std::unordered_map<std::string, std::string> split_kernels;
 if (source.size()) {
 std::string del{"// Function: "};
 size_t end;
 size_t begin = source.find(del);
 ICHECK(begin != std::string::npos) << "The OpenCL module expects a kernel delimited "
 << "source from code generation, but no kernel "
 << "delimiter was found.";
 for (size_t num_kernels = 0; num_kernels < workspace_->num_registered_kernels; num_kernels++) {
 begin += del.size();
 end = source.find('\n', begin);
 std::string func_name = source.substr(begin, end - begin);
 begin = ++end;
 // std::string::substr returns either start of next kernel
 // or std::string::npos, in the latter case substr returns
 // all characters until the end of the source string.
 end = source.find(del, begin);
 std::string func_source =
 source.substr(begin, (end == std::string::npos) ? end : end - begin);
 split_kernels.insert({func_name, func_source});
 begin = end;
 if (end == std::string::npos) {
 break;
 }
 }
 }
 ICHECK_EQ(workspace_->num_registered_kernels, split_kernels.size())
 << "The number of registered kernels does not match number of parsed kernel sources";
 return split_kernels;
}

Module OpenCLModuleCreate(std::string data, std::string fmt,
 std::unordered_map<std::string, FunctionInfo> fmap, std::string source) {
 auto n = make_object<OpenCLModuleNode>(data, fmt, fmap, source);
Expand Down

49 changes: 49 additions & 0 deletions
src/runtime/source_utils.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you 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.
 */

/*!
 * \file source_utils.cc
 */
#include "source_utils.h"

namespace tvm {
namespace runtime {

std::unordered_map<std::string, std::string> SplitKernels(std::string source,
Comment thread
echuraev marked this conversation as resolved.

 std::string delimiter) {
 std::unordered_map<std::string, std::string> split_kernels;
 if (source.size()) {
 size_t begin = source.find(delimiter);
@csullivancsullivanMay 12, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I noticed you deleted ICHECK here, probably because the error message referred to OpenCL specifically, but there still should be an
 ICHECK(begin != std::string::npos) 
on the first search for a kernel delimiter. If none are found, then it is better to catch an error in the code generation early.
@echuraevechuraevMay 13, 2021

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I did it to keep backward compatibility as @elvin-nnov requested. In case when the Metal library was generated before this change, we won't find any delimiters in the source code.

But I added some checks to opencl module:
https://github.com/apache/tvm/pull/7980/files#diff-9d619c5f05402f93afa552f93675dea84d789f8f6b0c553845aa183128f38494R192-R196
Also, I added check in metal module were kernels found or not: https://github.com/apache/tvm/pull/7980/files#diff-dc30040e2ec938d8dd0e6e87d6517870f037093775d4b52e0f249fdbebfdb9a3R95-R101
@csullivancsullivanMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would prefer the metal backward compat be specific to the metal backend and the general SplitKernel util to contain the existence ICHECK on the function delimiter.
 size_t end = begin;
 while (end != std::string::npos) {
 begin += delimiter.size();
 end = source.find('\n', begin);
 std::string func_name = source.substr(begin, end - begin);
 begin = ++end;
 end = source.find(delimiter, begin);
 std::string func_source =
 source.substr(begin, (end == std::string::npos) ? end : end - begin);
 split_kernels.insert({func_name, func_source});
 begin = end;
 }
 }
 return split_kernels;
}
} // namespace runtime
} // namespace tvm

44 changes: 44 additions & 0 deletions
src/runtime/source_utils.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you 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.
 */

/*!
 * \file source_utils.h
 * \brief Minimum source manipulation utils for runtime.
 */

#ifndef TVM_RUNTIME_SOURCE_UTILS_H_
#define TVM_RUNTIME_SOURCE_UTILS_H_

#include <string>
#include <unordered_map>

namespace tvm {
namespace runtime {
/*!
 * \brief Split the source file on separate kernels by specified delimiter.
 * \param source The source code of the kernels.
 * \param delimiter The delimiter which is using for splitting kernels.
 * \return Mapping from primitive name to kernel source
 */
std::unordered_map<std::string, std::string> SplitKernels(std::string source,
 std::string delimiter = "// Function: ");
} // namespace runtime
} // namespace tvm

#endif // TVM_RUNTIME_SOURCE_UTILS_H_

25 changes: 14 additions & 11 deletions
src/target/source/codegen_metal.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -305,27 +305,30 @@ void CodeGenMetal::VisitExpr_(const CallNode* op, std::ostream& os) { // NOLINT
runtime::Module BuildMetal(IRModule mod, Target target) {
 using tvm::runtime::Registry;
 bool output_ssa = false;
 CodeGenMetal cg;
 cg.Init(output_ssa);

 std::stringstream code;
 std::stringstream source;
 std::string fmt = "metal";
 for (auto kv : mod->functions) {
 ICHECK(kv.second->IsInstance<PrimFuncNode>()) << "CodeGenMetal: Can only take PrimFunc";
 code << "// Function: " << kv.first->name_hint << std::endl;
 CodeGenMetal cg;
 cg.Init(output_ssa);
 auto f = Downcast<PrimFunc>(kv.second);
 auto calling_conv = f->GetAttr<Integer>(tvm::attr::kCallingConv);
 ICHECK(calling_conv == CallingConv::kDeviceKernelLaunch)
 << "CodeGenMetal: expect calling_conv equals CallingConv::kDeviceKernelLaunch";
 cg.AddFunction(f);
 std::string fsource = cg.Finish();
 if (const auto* f = Registry::Get("tvm_callback_metal_compile")) {
 source << fsource;
 fsource = (*f)(fsource).operator std::string();
 fmt = "metallib";
 }
 code << fsource;
 }

 std::string code = cg.Finish();
 std::string fmt = "metal";
 std::string source = "";
 if (const auto* f = Registry::Get("tvm_callback_metal_compile")) {
 source = code;
 code = (*f)(code).operator std::string();
 fmt = "metallib";
 }
 return MetalModuleCreate(code, fmt, ExtractFuncInfo(mod), source);
 return MetalModuleCreate(code.str(), fmt, ExtractFuncInfo(mod), source.str());
}

TVM_REGISTER_GLOBAL("target.build.metal").set_body_typed(BuildMetal);
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode");
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();

Skip to content
[METAL] Split kernels and compile them separately#7980

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

tqchen
merged 5 commits into

from

May 25, 2021
 Merged


[METAL] Split kernels and compile them separately
#7980
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


5 commits


Select commit
Hold shift + click to select a range
b0b51fd

[METAL] Split kernels and compile them separately


echuraev May 5, 2021
4938d70

Add backward compatible behaviour


echuraev May 6, 2021
070bba4

Move SplitKernels to a common file


echuraev May 12, 2021
8041a29

Remove using GetSource


echuraev May 12, 2021
df55c33

Update the bundles


echuraev May 13, 2021

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


1 change: 1 addition & 0 deletions
apps/android_camera/app/src/main/jni/tvm_runtime.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -57,6 +57,7 @@
#ifdef TVM_OPENCL_RUNTIME
#include "../src/runtime/opencl/opencl_device_api.cc"
#include "../src/runtime/opencl/opencl_module.cc"
#include "../src/runtime/source_utils.cc"
#endif

#ifdef TVM_VULKAN_RUNTIME
Expand Down

1 change: 1 addition & 0 deletions
apps/android_rpc/app/src/main/jni/tvm_runtime.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -62,6 +62,7 @@
#ifdef TVM_OPENCL_RUNTIME
#include "../src/runtime/opencl/opencl_device_api.cc"
#include "../src/runtime/opencl/opencl_module.cc"
#include "../src/runtime/source_utils.cc"
#endif

#ifdef TVM_VULKAN_RUNTIME
Expand Down

1 change: 1 addition & 0 deletions
golang/src/tvm_runtime_pack.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -68,3 +68,4 @@
// Uncomment the following lines to enable OpenCL
// #include "../../src/runtime/opencl/opencl_device_api.cc"
// #include "../../src/runtime/opencl/opencl_module.cc"
// #include "../src/runtime/source_utils.cc"

79 changes: 45 additions & 34 deletions
src/runtime/metal/metal_module.mm

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -30,6 +30,7 @@
#include "../file_utils.h"
#include "../meta_data.h"
#include "../pack_args.h"
#include "../source_utils.h"
#include "../thread_storage_scope.h"
#include "metal_common.h"

Expand All@@ -43,7 +44,9 @@
 public:
 explicit MetalModuleNode(std::string data, std::string fmt,
 std::unordered_map<std::string, FunctionInfo> fmap, std::string source)
 : data_(data), fmt_(fmt), fmap_(fmap), source_(source) {}
 : data_(data), fmt_(fmt), fmap_(fmap), source_(source) {
 parsed_kernels_ = SplitKernels(data);
 }
 const char* type_key() const final { return "metal"; }

 PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) final;
Expand DownExpand Up@@ -71,6 +74,7 @@ void SaveToBinary(dmlc::Stream* stream) final {
 return "";
 }
 }

 // get a from primary context in device_id
 id<MTLComputePipelineState> GetPipelineState(size_t device_id, const std::string& func_name) {
 metal::MetalWorkspace* w = metal::MetalWorkspace::Global();
Expand All@@ -85,44 +89,52 @@ void SaveToBinary(dmlc::Stream* stream) final {
 if (it != e.smap.end()) return it->second;
 // compile
 NSError* err_msg = nil;
 if (e.lib == nil) {
 if (fmt_ == "metal") {
 MTLCompileOptions* opts = [MTLCompileOptions alloc];
 opts.languageVersion = MTLLanguageVersion2_3;
 opts.fastMathEnabled = YES;
 // opts = nil;
 e.lib = [w->devices[device_id]
 newLibraryWithSource:[NSString stringWithUTF8String:data_.c_str()]
 options:opts
 error:&err_msg];
 [opts dealloc];
 if (e.lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 if (err_msg != nil) {
 LOG(INFO) << "Warning: " << [[err_msg localizedDescription] UTF8String];
 }
 } else {
 // Build from library.
 auto q = dispatch_queue_create("q", DISPATCH_QUEUE_SERIAL);
 auto data = dispatch_data_create(data_.c_str(), data_.length(), q,
 ^{
 });
 e.lib = [w->devices[device_id] newLibraryWithData:data error:&err_msg];
 if (err_msg != nil || e.lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 id<MTLLibrary> lib = nil;
 std::string source;
 auto kernel = parsed_kernels_.find(func_name);
 // If we cannot find this kernel in parsed_kernels_, it means that all kernels going together
 // without explicit separator. In this case we use data_ with all kernels. It done for backward
 // compatibility.
 if (kernel != parsed_kernels_.end())
@elvin-nelvin-nMay 11, 2021•

edited

Loading

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would add here a comment that this behaviour is done for backward compatibility when we had all kernels going together without explicit separators, but looks good
@echuraevechuraevMay 12, 2021

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done
@csullivancsullivanMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we have metal deployments that require the old behavior that can't be rebuilt easily? Otherwise we don't need to preserve backward compatibility.
@elvin-nelvin-nMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I cannot say for sure. I can imagine some use cases which might have issues with recompilation, but this is just assumptions.
If we do not declare backward compatibility for now, we can skip my comment and return original behaviour and rise an exception if we do not find kernel in parsed_kernels_ container.
@csullivancsullivanMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In that case I personally would prefer the metal backward compat be specific to the metal backend and the general SplitKernel util to contain the existence ICHECK on the function delimiter.
@elvin-nelvin-nMay 14, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tqchen do we care about backward compatibility or not? If not, it might be ok so far, but we need to decide when we will start to care and what use cases should work and be regularly verified in the future. It must not be a deal only metal backend
@csullivan what bahavior do you propose in common SplitKernel? Metal backend uses it, it cannot just abort execution. Warnings in the runtime are also useles. It might be an error code, but empty returned map detects the absent of Function delimiter with the same efficiency as return code
 source = kernel->second;
 else
 source = data_;
 if (fmt_ == "metal") {
 MTLCompileOptions* opts = [MTLCompileOptions alloc];
 opts.languageVersion = MTLLanguageVersion2_3;
 opts.fastMathEnabled = YES;
 // opts = nil;
 lib =
 [w->devices[device_id] newLibraryWithSource:[NSString stringWithUTF8String:source.c_str()]
 options:opts
 error:&err_msg];
 [opts dealloc];
 if (lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 if (err_msg != nil) {
 LOG(INFO) << "Warning: " << [[err_msg localizedDescription] UTF8String];
 }
 } else {
 // Build from library.
 auto q = dispatch_queue_create("q", DISPATCH_QUEUE_SERIAL);
 auto data = dispatch_data_create(source.c_str(), source.length(), q,
 ^{
 });
 lib = [w->devices[device_id] newLibraryWithData:data error:&err_msg];
 if (err_msg != nil || lib == nil) {
 LOG(FATAL) << "Fail to compile metal lib:" << [[err_msg localizedDescription] UTF8String];
 }
 }
 id<MTLFunction> f =
 [e.lib newFunctionWithName:[NSString stringWithUTF8String:func_name.c_str()]];
 id<MTLFunction> f = [lib newFunctionWithName:[NSString stringWithUTF8String:func_name.c_str()]];
 ICHECK(f != nil) << "cannot find function " << func_name;
 id<MTLComputePipelineState> state =
 [w->devices[device_id] newComputePipelineStateWithFunction:f error:&err_msg];
 ICHECK(state != nil) << "cannot get state:"
 << " for function " << func_name
 << [[err_msg localizedDescription] UTF8String];
 [f release];
 [lib release];
 // The state.threadExecutionWidth can change dynamically according
 // to the resource constraint in kernel, so it is not strictly hold
 // Turn of warp aware optimziation for now.
Expand All@@ -135,13 +147,10 @@ void SaveToBinary(dmlc::Stream* stream) final {
 private:
 // device specific entry
 struct DeviceEntry {
 // library
 id<MTLLibrary> lib = nil;
 // state cache;
 std::unordered_map<std::string, id<MTLComputePipelineState>> smap;
 std::unordered_map<std::string, id<MTLComputePipelineState>> smap;

 ~DeviceEntry() {
 if (lib != nil) [lib release];
 for (auto&& kv : smap) {
 [kv.second release];
 }
Expand All@@ -159,6 +168,8 @@ void SaveToBinary(dmlc::Stream* stream) final {
 std::vector<DeviceEntry> finfo_;
 // internal mutex when updating the module
 std::mutex mutex_;
 // parsed kernel data
 std::unordered_map<std::string, std::string> parsed_kernels_;
};

// a wrapped function class to get packed func.
Expand Down

8 changes: 0 additions & 8 deletions
src/runtime/opencl/opencl_common.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -326,14 +326,6 @@ class OpenCLModuleNode : public ModuleNode {
 cl_kernel InstallKernel(cl::OpenCLWorkspace* w, cl::OpenCLThreadEntry* t,
 const std::string& func_name, const KTRefEntry& e);

 /*
 * \brief Splits the provided serialized source file into separate
 * source for each kernel primitive.
 * \param source The serialized program source file (fmt: cl)
 * \return Mapping from primitive name to kernel source
 */
 std::unordered_map<std::string, std::string> SplitKernels(std::string source) const;

 private:
 // The workspace, need to keep reference to use it in destructor.
 // In case of static destruction order problem.
Expand Down

39 changes: 6 additions & 33 deletions
src/runtime/opencl/opencl_module.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -29,6 +29,7 @@
#include <unordered_map>
#include <vector>

#include "../source_utils.h"
#include "opencl_common.h"

namespace tvm {
Expand DownExpand Up@@ -188,6 +189,11 @@ void OpenCLModuleNode::Init() {

 // split into source artifacts for each kernel
 parsed_kernels_ = SplitKernels(GetSource("cl"));
 ICHECK(!parsed_kernels_.empty()) << "The OpenCL module expects a kernel delimited "
 << "source from code generation, but no kernel "
 << "delimiter was found.";
 ICHECK_EQ(workspace_->num_registered_kernels, parsed_kernels_.size())
 << "The number of registered kernels does not match number of parsed kernel sources";
 // zero initialize cl_program pointers for each device kernel
 for (auto& kv : parsed_kernels_) {
 programs_.insert({kv.first, std::vector<cl_program>(workspace_->devices.size(), nullptr)});
Expand DownExpand Up@@ -242,39 +248,6 @@ cl_kernel OpenCLModuleNode::InstallKernel(cl::OpenCLWorkspace* w, cl::OpenCLThre
 return kernel;
}

std::unordered_map<std::string, std::string> OpenCLModuleNode::SplitKernels(
 std::string source) const {
 std::unordered_map<std::string, std::string> split_kernels;
 if (source.size()) {
 std::string del{"// Function: "};
 size_t end;
 size_t begin = source.find(del);
 ICHECK(begin != std::string::npos) << "The OpenCL module expects a kernel delimited "
 << "source from code generation, but no kernel "
 << "delimiter was found.";
 for (size_t num_kernels = 0; num_kernels < workspace_->num_registered_kernels; num_kernels++) {
 begin += del.size();
 end = source.find('\n', begin);
 std::string func_name = source.substr(begin, end - begin);
 begin = ++end;
 // std::string::substr returns either start of next kernel
 // or std::string::npos, in the latter case substr returns
 // all characters until the end of the source string.
 end = source.find(del, begin);
 std::string func_source =
 source.substr(begin, (end == std::string::npos) ? end : end - begin);
 split_kernels.insert({func_name, func_source});
 begin = end;
 if (end == std::string::npos) {
 break;
 }
 }
 }
 ICHECK_EQ(workspace_->num_registered_kernels, split_kernels.size())
 << "The number of registered kernels does not match number of parsed kernel sources";
 return split_kernels;
}

Module OpenCLModuleCreate(std::string data, std::string fmt,
 std::unordered_map<std::string, FunctionInfo> fmap, std::string source) {
 auto n = make_object<OpenCLModuleNode>(data, fmt, fmap, source);
Expand Down

49 changes: 49 additions & 0 deletions
src/runtime/source_utils.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you 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.
 */

/*!
 * \file source_utils.cc
 */
#include "source_utils.h"

namespace tvm {
namespace runtime {

std::unordered_map<std::string, std::string> SplitKernels(std::string source,
Comment thread
echuraev marked this conversation as resolved.

 std::string delimiter) {
 std::unordered_map<std::string, std::string> split_kernels;
 if (source.size()) {
 size_t begin = source.find(delimiter);
@csullivancsullivanMay 12, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I noticed you deleted ICHECK here, probably because the error message referred to OpenCL specifically, but there still should be an
 ICHECK(begin != std::string::npos) 
on the first search for a kernel delimiter. If none are found, then it is better to catch an error in the code generation early.
@echuraevechuraevMay 13, 2021

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I did it to keep backward compatibility as @elvin-nnov requested. In case when the Metal library was generated before this change, we won't find any delimiters in the source code.

But I added some checks to opencl module:
https://github.com/apache/tvm/pull/7980/files#diff-9d619c5f05402f93afa552f93675dea84d789f8f6b0c553845aa183128f38494R192-R196
Also, I added check in metal module were kernels found or not: https://github.com/apache/tvm/pull/7980/files#diff-dc30040e2ec938d8dd0e6e87d6517870f037093775d4b52e0f249fdbebfdb9a3R95-R101
@csullivancsullivanMay 13, 2021

Copy link


Copy Markdown

Contributor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would prefer the metal backward compat be specific to the metal backend and the general SplitKernel util to contain the existence ICHECK on the function delimiter.
 size_t end = begin;
 while (end != std::string::npos) {
 begin += delimiter.size();
 end = source.find('\n', begin);
 std::string func_name = source.substr(begin, end - begin);
 begin = ++end;
 end = source.find(delimiter, begin);
 std::string func_source =
 source.substr(begin, (end == std::string::npos) ? end : end - begin);
 split_kernels.insert({func_name, func_source});
 begin = end;
 }
 }
 return split_kernels;
}
} // namespace runtime
} // namespace tvm

44 changes: 44 additions & 0 deletions
src/runtime/source_utils.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you 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.
 */

/*!
 * \file source_utils.h
 * \brief Minimum source manipulation utils for runtime.
 */

#ifndef TVM_RUNTIME_SOURCE_UTILS_H_
#define TVM_RUNTIME_SOURCE_UTILS_H_

#include <string>
#include <unordered_map>

namespace tvm {
namespace runtime {
/*!
 * \brief Split the source file on separate kernels by specified delimiter.
 * \param source The source code of the kernels.
 * \param delimiter The delimiter which is using for splitting kernels.
 * \return Mapping from primitive name to kernel source
 */
std::unordered_map<std::string, std::string> SplitKernels(std::string source,
 std::string delimiter = "// Function: ");
} // namespace runtime
} // namespace tvm

#endif // TVM_RUNTIME_SOURCE_UTILS_H_

25 changes: 14 additions & 11 deletions
src/target/source/codegen_metal.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -305,27 +305,30 @@ void CodeGenMetal::VisitExpr_(const CallNode* op, std::ostream& os) { // NOLINT
runtime::Module BuildMetal(IRModule mod, Target target) {
 using tvm::runtime::Registry;
 bool output_ssa = false;
 CodeGenMetal cg;
 cg.Init(output_ssa);

 std::stringstream code;
 std::stringstream source;
 std::string fmt = "metal";
 for (auto kv : mod->functions) {
 ICHECK(kv.second->IsInstance<PrimFuncNode>()) << "CodeGenMetal: Can only take PrimFunc";
 code << "// Function: " << kv.first->name_hint << std::endl;
 CodeGenMetal cg;
 cg.Init(output_ssa);
 auto f = Downcast<PrimFunc>(kv.second);
 auto calling_conv = f->GetAttr<Integer>(tvm::attr::kCallingConv);
 ICHECK(calling_conv == CallingConv::kDeviceKernelLaunch)
 << "CodeGenMetal: expect calling_conv equals CallingConv::kDeviceKernelLaunch";
 cg.AddFunction(f);
 std::string fsource = cg.Finish();
 if (const auto* f = Registry::Get("tvm_callback_metal_compile")) {
 source << fsource;
 fsource = (*f)(fsource).operator std::string();
 fmt = "metallib";
 }
 code << fsource;
 }

 std::string code = cg.Finish();
 std::string fmt = "metal";
 std::string source = "";
 if (const auto* f = Registry::Get("tvm_callback_metal_compile")) {
 source = code;
 code = (*f)(code).operator std::string();
 fmt = "metallib";
 }
 return MetalModuleCreate(code, fmt, ExtractFuncInfo(mod), source);
 return MetalModuleCreate(code.str(), fmt, ExtractFuncInfo(mod), source.str());
}

TVM_REGISTER_GLOBAL("target.build.metal").set_body_typed(BuildMetal);
Expand Down