Skip to content

Commit

Permalink
Improved MLF to contain workspace info
Browse files Browse the repository at this point in the history 
* addressing comments - 1

Change-Id: I729158eda564f0a878226702fbbcd1dbf7f014f7
  • Loading branch information
@manupak
manupak committed May 5, 2021
1 parent cdbc0a9 commit b458a89
Show file tree
Hide file tree
Showing 2 changed files with 60 additions and 50 deletions.
13 changes: 12 additions & 1 deletion src/relay/backend/graph_executor_codegen.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -227,11 +227,16 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
* \param func The main function that contains calls to relay primitive functions
*/
void UpdateMainWorkspaceSize(const Function& func) {
// This is a Map<device,Map<storage_id, size>>
std::unordered_map<int, std::unordered_map<int, int>> sid_workspace;
// This is a Map<device, workspace_size>
std::unordered_map<int, int> device_workspace;
// This is a Map<device, size_of_inputs_and_outputs>
std::unordered_map<int, int> device_io;
// This is a Map<device, size_of_constants>
std::unordered_map<int, int> device_consts;

// Initialize the maps to zero
for (const auto& kv : storage_device_map_) {
auto sids = kv.second[0];
auto devices = kv.second[1];
Expand All @@ -240,9 +245,11 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
sid_workspace[devices[i]][sids[i]] = 0;
device_io[devices[i]] = 0;
device_consts[devices[i]] = 0;
device_workspace[devices[i]] = 0;
}
}

// Collect sizes of tensors
for (const auto& kv : storage_device_map_) {
auto size_bytes = CalculateRelayExprSizeBytes(kv.first->checked_type());
auto sids = kv.second[0];
Expand All @@ -259,20 +266,24 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
continue;
}
for (uint32_t i = 0; i < sids.size(); i++) {
// Here we record the largest size of the tensor
// that share the same storage id, because storage_id will
// be shared between multiple tensors that are not live simultaneously.
if (size_bytes > sid_workspace[devices[i]][sids[i]]) {
sid_workspace[devices[i]][sids[i]] = size_bytes;
}
}
}

// Once we know the sizes of sids, we need to accumulate per device
for (const auto& dev_sid_size : sid_workspace) {
auto dev = dev_sid_size.first;
device_workspace[dev] = 0;
for (const auto& sid_size : dev_sid_size.second) {
device_workspace[dev] += sid_size.second;
}
}

// Populate FunctionInfo
auto fi_node = make_object<FunctionInfoNode>();
for (const auto& dev_and_size : device_workspace) {
auto tgt = GetTargetFromInteger(dev_and_size.first);
Expand Down
97 changes: 48 additions & 49 deletions tests/python/unittest/test_micro_model_library_format.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -201,59 +201,58 @@ def @main(%a : Tensor[(1, 2), uint8], %b : Tensor[(1, 2), float32], %c : Tensor[

@tvm.testing.requires_micro
def test_export_model_library_format_workspace():
with utils.TempDirectory.set_keep_for_debug(True):
target = tvm.target.target.micro("host")
with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
relay_mod = tvm.parser.fromtext(
"""
#[version = "0.0.5"]
def @main(%p0: Tensor[(1, 56, 56, 128), int16], %p1: Tensor[(3, 3, 128, 1), int16], %p2: Tensor[(1, 1, 1, 128), int32]){
%0 = nn.conv2d(%p0, %p1, padding=[1, 1, 1, 1], groups=128, channels=128, kernel_size=[3, 3], data_layout="NHWC", kernel_layout="HWOI", out_dtype="int32") /* ty=Tensor[(1, 56, 56, 128), int32] */;
%1 = add(%0, %p2) /* ty=Tensor[(1, 56, 56, 128), int32] */;
%2 = fixed_point_multiply(%1, multiplier=2080045879, shift=-4) /* ty=Tensor[(1, 56, 56, 128), int32] */;
%3 = clip(%2, a_min=0f, a_max=255f) /* ty=Tensor[(1, 56, 56, 128), int32] */;
cast(%3, dtype="uint8") /* ty=Tensor[(1, 56, 56, 128), uint8] */
}
"""
)
factory = tvm.relay.build(relay_mod, target, target_host=target, mod_name="qnn_conv2d")
target = tvm.target.target.micro("host")
with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
relay_mod = tvm.parser.fromtext(
"""
#[version = "0.0.5"]
def @main(%p0: Tensor[(1, 56, 56, 128), int16], %p1: Tensor[(3, 3, 128, 1), int16], %p2: Tensor[(1, 1, 1, 128), int32]){
%0 = nn.conv2d(%p0, %p1, padding=[1, 1, 1, 1], groups=128, channels=128, kernel_size=[3, 3], data_layout="NHWC", kernel_layout="HWOI", out_dtype="int32") /* ty=Tensor[(1, 56, 56, 128), int32] */;
%1 = add(%0, %p2) /* ty=Tensor[(1, 56, 56, 128), int32] */;
%2 = fixed_point_multiply(%1, multiplier=2080045879, shift=-4) /* ty=Tensor[(1, 56, 56, 128), int32] */;
%3 = clip(%2, a_min=0f, a_max=255f) /* ty=Tensor[(1, 56, 56, 128), int32] */;
cast(%3, dtype="uint8") /* ty=Tensor[(1, 56, 56, 128), uint8] */
}
"""
)
factory = tvm.relay.build(relay_mod, target, target_host=target, mod_name="qnn_conv2d")

temp_dir = utils.tempdir()
mlf_tar_path = temp_dir.relpath("lib.tar")
import tvm.micro as micro
temp_dir = utils.tempdir()
mlf_tar_path = temp_dir.relpath("lib.tar")
import tvm.micro as micro

micro.export_model_library_format(factory, mlf_tar_path)
tf = tarfile.open(mlf_tar_path)
micro.export_model_library_format(factory, mlf_tar_path)
tf = tarfile.open(mlf_tar_path)

extract_dir = temp_dir.relpath("extract")
os.mkdir(extract_dir)
tf.extractall(extract_dir)
extract_dir = temp_dir.relpath("extract")
os.mkdir(extract_dir)
tf.extractall(extract_dir)

with open(os.path.join(extract_dir, "metadata.json")) as json_f:
metadata = json.load(json_f)
assert metadata["version"] == 2
assert metadata["model_name"] == "qnn_conv2d"
export_datetime = datetime.datetime.strptime(
metadata["export_datetime"], "%Y-%m-%d %H:%M:%SZ"
)
assert (datetime.datetime.now() - export_datetime) < datetime.timedelta(seconds=60 * 5)
assert metadata["target"] == {"1": str(target)}
assert metadata["memory"]["functions"] == {
"main_function": [
{
"constants_size_bytes": 0,
"device": 1,
"io_size_bytes": 1207040,
"workspace_size_bytes": 2466816,
}
],
"operator_functions": [
{
"function_name": "fused_nn_conv2d_add_fixed_point_multiply_clip_cast",
"workspace": [{"device": 1, "workspace_size_bytes": 2466816}],
}
],
}
with open(os.path.join(extract_dir, "metadata.json")) as json_f:
metadata = json.load(json_f)
assert metadata["version"] == 2
assert metadata["model_name"] == "qnn_conv2d"
export_datetime = datetime.datetime.strptime(
metadata["export_datetime"], "%Y-%m-%d %H:%M:%SZ"
)
assert (datetime.datetime.now() - export_datetime) < datetime.timedelta(seconds=60 * 5)
assert metadata["target"] == {"1": str(target)}
assert metadata["memory"]["functions"] == {
"main_function": [
{
"constants_size_bytes": 0,
"device": 1,
"io_size_bytes": 1207040,
"workspace_size_bytes": 2466816,
}
],
"operator_functions": [
{
"function_name": "fused_nn_conv2d_add_fixed_point_multiply_clip_cast",
"workspace": [{"device": 1, "workspace_size_bytes": 2466816}],
}
],
}


@tvm.testing.requires_micro
Expand Down

0 comments on commit b458a89

Please sign in to comment.