[MetaSchedule] Consolidate module hashing and equality testing #13050

masahi · 2022-10-12T00:39:35Z

Currently, MS uses StructuralEqual/Hash in task extraction / evo search / database. Sometimes, we want to use different hashing and equality testing methods, for example (1) to ignore NDArray (#12706) or (2) to enable anchor-op only tuning (identify conv2d and conv2d -> add subgraphs as equal).

To enable such flexibility, this PR consolidate raw calls to StructuralEqual/Hash into one place, which for now is named ModuleEquality. Since hashing is also done for equality testing, I think it is appropriate to call the component responsible for hashing / equality test that way. But other suggestions are welcome.

Importantly, task extraction and database are now using the same hashing / equal method based on TIR mod, while previously task extraction was using a cache key-ed on relay mod.

cc @junrushao @zxybazh @tqchen

tvm-bot · 2022-10-12T00:39:37Z

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

cc @Hzfengsy, @junrushao _{See #10317 for details}
Built docs for commit 95c01b7 can be found here.

_{Generated by tvm-bot}

masahi · 2022-10-12T00:45:20Z

src/meta_schedule/database/database.cc

  ObjectPtr<WorkloadNode> n = runtime::make_object<WorkloadNode>();
-  n->shash = tvm::StructuralHash()(mod);
  n->mod = mod;
+  n->shash = ModuleEquality::Create("structural")->Hash(mod);


Since this constructor is not used anywhere, I'm hard coding the instance here to avoid passing the mod_eq_name string. I want to remove this constructor if there is no objection.

masahi · 2022-10-12T00:50:00Z

src/meta_schedule/database/json_database.cc

 /*! \brief The default database implementation, which mimics two database tables with two files. */
 class JSONDatabaseNode : public DatabaseNode {
 public:
+  explicit JSONDatabaseNode(String mod_eq_name = "structural")


The default value is required, since we need to be able to default-construct JSONDatabaseNode and other database nodes due to the use of TVM_REGISTER_NODE_TYPE macro.

Is it possible if we bypass it via some tricks like unique ptr? (Sorry I don’t think very deeply so feel free to object)

I didn't get what you meant here.

just noticed that DatabaseNode::mod_eq_ has been already a unique_ptr - so we don't need an explicit default constructor DatabaseNode::DatabaseNode, but instead we could do inline initialization:

- std::unique_ptr<ModuleEquality> mod_eq_; + std::unique_ptr<ModuleEquality> mod_eq_{nullptr};

BTW, just to nitpick, shall we change mod_eq_ from private to protected so it could be accessed by derived classes?

Since JSONDatabaseNode is constructed by Database class,

tvm/src/meta_schedule/database/json_database.cc

Lines 147 to 150 in eecb7fd

Database Database::JSONDatabase(String path_workload, String path_tuning_record,

bool allow_missing) {

int num_threads = std::thread::hardware_concurrency();

ObjectPtr<JSONDatabaseNode> n = make_object<JSONDatabaseNode>();

and Database is not a subclass of DatabaseNode, setting mod_eq after JSONDatabaseNode is constructed requires making mod_eq a public member or introducing a setter method in DatabaseNode.

Moreover, it would also introduce complication to initializing workloads2idx_ (see #13050 (comment)).

Just tried: Removing the construction of workloads2idx_, workloads2idx_(/*bucket_count*/ 0, WorkloadHash(), WorkloadEqual(GetModuleEquality())) leads to compilation failure since workloads2idx_ cannot be default-constructed (due to WorkloadEqual which requires mod_eq to be initialized).

I see. Then I'm fine with the current implementation

masahi · 2022-10-12T00:52:37Z

src/meta_schedule/database/json_database.cc

 public:
+  explicit JSONDatabaseNode(String mod_eq_name = "structural")
+      : DatabaseNode(mod_eq_name),
+        workloads2idx_(/*bucket_count*/ 0, WorkloadHash(), WorkloadEqual(GetModuleEquality())) {}


Since WorkloadEqual needs to be aware of the module hashing / equality testing being used, we need to construct this unordered_map after the ModuleEquality instance is constructed in DatabaseNode by L72.

src/relay/backend/task_extraction.cc

masahi · 2022-10-12T01:06:50Z

src/meta_schedule/database/database.cc

-    shash = tvm::StructuralHash()(mod);
-    String recalc_shash = SHash2Str(shash);
-    CHECK_EQ(recalc_shash, str_shash) << "ValueError: Structural hash changed. Given: " << str_shash
-                                      << "; Recalculated: " << recalc_shash;


Since ModuleEquality information is not written to json (Workload doesn't have ModuleEquality instance), we cannot determine the appropriate hashing method here. So I had to remove this verification.

Perhaps we could leave a TODO here in case of future changes

I think the removal make sense. On the other hand, we are always using structural hash as the shash. I think we may not really need the functionality to specify the shash direcly in constructor.

What about we construct it from the mod and mod_eq_name string so that we can obtain the customized hash result?

Maybe we can store the mod_eq_name in Workload so that when we parse it we can still check the shash results?

What do you think?

To keep things clean, I want to minimize the number of unique ModuleEquality instances scattered throughout MS infra. Currently, only task extraction and Database "own" a ModuleEquality instance.

Given this goal, I think having Database compute a hash of Workload using its ModuleEquality is more natural than passing mod_eq_name to Workload and let it compute hash. Moreover, all Workload would end up having the same mod_eq_name, so it feels redundant.

One way to restore hash verification is to do verification inside DatabaseNode after Workload::FromJSON is called, since the database does know what the ModuleEquality instance is supposed to be.

Restored hash verification in JSONDatabase, the only place where Workload::FromJSON is used currently.

junrushao

Otherwise LGTM!

zxybazh

Thanks for sending the new feature! It's very exciting to see such improvements!

Generally it's looking good. I have some minor suggestions over the Workload usage and question towards the hash map of workloads.

Other than these, currently I can see the only option of module equality functions defined on the c++ side, do you think we can also allow easier customization of the hash and equal function via providing FFI functions or by overridding the class from python side (i.e., providing a PyModuleEquality class)?

src/relay/backend/task_extraction.cc

src/meta_schedule/search_strategy/evolutionary_search.cc

src/relay/backend/task_extraction.cc

include/tvm/meta_schedule/database.h

zxybazh · 2022-10-14T05:45:28Z

src/meta_schedule/database/database.cc

-    shash = tvm::StructuralHash()(mod);
-    String recalc_shash = SHash2Str(shash);
-    CHECK_EQ(recalc_shash, str_shash) << "ValueError: Structural hash changed. Given: " << str_shash
-                                      << "; Recalculated: " << recalc_shash;


I think the removal make sense. On the other hand, we are always using structural hash as the shash. I think we may not really need the functionality to specify the shash direcly in constructor.

What about we construct it from the mod and mod_eq_name string so that we can obtain the customized hash result?

Maybe we can store the mod_eq_name in Workload so that when we parse it we can still check the shash results?

What do you think?

src/meta_schedule/database/database.cc

include/tvm/meta_schedule/database.h

masahi · 2022-10-14T07:55:45Z

currently I can see the only option of module equality functions defined on the c++ side, do you think we can also allow easier customization of the hash and equal function via providing FFI functions or by overridding the class from python side (i.e., providing a PyModuleEquality class)?

@zxybazh This can be considered in future, but I don't expect that people would want to add a custom ModuleEquality often. So for now, ModuleEquality is not exposed to python, this lets me spare from boilerplates and make ModuleEquality creation as simple as

std::unique_ptr<ModuleEquality> ModuleEquality::Create(const std::string& mod_eq_name) {
  if (mod_eq_name == "structural") {
    return std::make_unique<ModuleEqualityStructural>();
  }
  LOG(FATAL) << "Unknown module equality " << mod_eq_name;
  return nullptr;
}

…Hash

zxybazh

Thanks @masahi for addressing my questions and adding the verification, the PR looks good to me! Looking forward to the introduction of new module equality functions!

junrushao

LGTM

…ng NDArray raw data (#13091) A follow up to #13050, also builds on #13001. This PR enables the functionality in #12706 without changing the existing `StructuralEqual/Hash`. A question for discussion: Should this be the default ModuleEquality used by MS? It has no effect for the `link-params = False` case, and it simplifies the MS tuning API usage for the `link-params = True` case (Hexagon etc).

…e#13050) Currently, MS uses `StructuralEqual/Hash` in task extraction / evo search / database. Sometimes, we want to use different hashing and equality testing methods, for example (1) to ignore NDArray (apache#12706) or (2) to enable anchor-op only tuning (identify `conv2d` and `conv2d -> add` subgraphs as equal). To enable such flexibility, this PR consolidate raw calls to `StructuralEqual/Hash` into one place, which for now is named `ModuleEquality`. Since hashing is also done for equality testing, I think it is appropriate to call the component responsible for hashing / equality test that way. But other suggestions are welcome. Importantly, task extraction and database are now using the same hashing / equal method based on TIR mod, while previously task extraction was using a cache key-ed on relay mod.

…ng NDArray raw data (apache#13091) A follow up to apache#13050, also builds on apache#13001. This PR enables the functionality in apache#12706 without changing the existing `StructuralEqual/Hash`. A question for discussion: Should this be the default ModuleEquality used by MS? It has no effect for the `link-params = False` case, and it simplifies the MS tuning API usage for the `link-params = True` case (Hexagon etc).

…e#13050) Currently, MS uses `StructuralEqual/Hash` in task extraction / evo search / database. Sometimes, we want to use different hashing and equality testing methods, for example (1) to ignore NDArray (apache#12706) or (2) to enable anchor-op only tuning (identify `conv2d` and `conv2d -> add` subgraphs as equal). To enable such flexibility, this PR consolidate raw calls to `StructuralEqual/Hash` into one place, which for now is named `ModuleEquality`. Since hashing is also done for equality testing, I think it is appropriate to call the component responsible for hashing / equality test that way. But other suggestions are welcome. Importantly, task extraction and database are now using the same hashing / equal method based on TIR mod, while previously task extraction was using a cache key-ed on relay mod.

…ng NDArray raw data (apache#13091) A follow up to apache#13050, also builds on apache#13001. This PR enables the functionality in apache#12706 without changing the existing `StructuralEqual/Hash`. A question for discussion: Should this be the default ModuleEquality used by MS? It has no effect for the `link-params = False` case, and it simplifies the MS tuning API usage for the `link-params = True` case (Hexagon etc).

masahi commented Oct 12, 2022

View reviewed changes

src/relay/backend/task_extraction.cc Show resolved Hide resolved

masahi force-pushed the metasch-refactor-structural branch from 3c79ad5 to b71ca1b Compare October 12, 2022 00:55

github-actions bot requested review from junrushao, tqchen and zxybazh October 12, 2022 01:01

masahi commented Oct 12, 2022

View reviewed changes

masahi force-pushed the metasch-refactor-structural branch from b71ca1b to 0c8ee7f Compare October 12, 2022 01:13

junrushao approved these changes Oct 14, 2022

View reviewed changes

zxybazh reviewed Oct 14, 2022

View reviewed changes

masahi added 7 commits October 15, 2022 04:07

[MetaSchedule] Introduce ModuleEqual/Hash to replace StructuralEqual/…

025d10b

…Hash

moving away from global ModuleEqality object

f4f8cb7

move module_equality header to src

16134f7

add doc

20e41f0

pass module_equality to extract_tasks in tune_relay

f4d8a96

restored module hash verification

364e478

doc fix

95c01b7

masahi force-pushed the metasch-refactor-structural branch from 3058a68 to 95c01b7 Compare October 14, 2022 19:07

zxybazh approved these changes Oct 15, 2022

View reviewed changes

junrushao approved these changes Oct 17, 2022

View reviewed changes

junrushao merged commit cbca28d into apache:main Oct 17, 2022

masahi mentioned this pull request Oct 17, 2022

[MetaSchedule] Introduce a variant of ModuleEquality to enable ignoring NDArray raw data #13091

Merged

masahi mentioned this pull request Oct 26, 2022

[MetaSchedule] Enable anchor-block tuning #13206

Merged

leandron mentioned this pull request Feb 1, 2023

TVM v0.11.0 Release Candidate Notes #13899

Closed

	Database Database::JSONDatabase(String path_workload, String path_tuning_record,
	bool allow_missing) {
	int num_threads = std::thread::hardware_concurrency();
	ObjectPtr<JSONDatabaseNode> n = make_object<JSONDatabaseNode>();

[MetaSchedule] Consolidate module hashing and equality testing #13050

[MetaSchedule] Consolidate module hashing and equality testing #13050

Uh oh!

Conversation

masahi commented Oct 12, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tvm-bot commented Oct 12, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

masahi Oct 12, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

masahi Oct 14, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

junrushao left a comment

Choose a reason for hiding this comment

Uh oh!

zxybazh left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

masahi commented Oct 14, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

zxybazh left a comment

Choose a reason for hiding this comment

Uh oh!

junrushao left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

masahi commented Oct 12, 2022 •

edited

Loading

tvm-bot commented Oct 12, 2022 •

edited

Loading

masahi Oct 12, 2022 •

edited

Loading

masahi Oct 14, 2022 •

edited

Loading

masahi commented Oct 14, 2022 •

edited

Loading