[MetaSchedule][M3b] Database

[junrushao]

This PR is part of the meta schedule project (#8473) that adds a workload registry, as well as
a generic database of tuning records. This feature is future-compatible with dynamic shape
auto-tuning work by @ZihengJiang and @ArmageddonKnight.

Depend on PR #9059. The PR is ready for review now!

Co-authored-by: Xiyou Zhou <[email protected]>
Co-authored-by: Bohan Hou <[email protected]>
Co-authored-by: Ruihang Lai <[email protected]>
Co-authored-by: Hongyi Jin <[email protected]>
Co-authored-by: Wuwei Lin <[email protected]>
Co-authored-by: Siyuan Feng <[email protected]>

[junrushao]

CC @comaniac @mbrookhart @jroesch @ZihengJiang

[comaniac]

Will try to find time reviewing this PR this afternoon or tomorrow.

[comaniac]

Took a fast pass and had an offline discussion with @junrushao1994. Two major points:

1. The APIs are bit confusing (e.g.,lookup_or_add and add). It might be better to use something like commit_workload and commit_record. Meanwhile, since we may add more database implementations in the future (e.g., sqlite), this general API can also make sure we won't need to refine it again in AutoTIR.
2. Better to remove workload registry accordingly to make the design simple.

[junrushao]

Per discussion with @comaniac.

The naming change:

• lookup_or_add seems pretty random, and we should rename it as commit_workload which is more consistent with database terminologies
• add seems okay but commit_tuning_record sounds better and more consistent with commit_workload

Minor change in the code structure:

• Consolidate WorkloadRegistry into JSONFile database, so that there is only one unified database API, and this could mimic the two-table database behavior

[mbrookhart]

First kind of stream of consciousness pass through the code. Mostly documentation/comment requests. It seems the direct use of the Database is all wrapped in a JSON-related file object, which kind of confuses me, would you mind providing a brief explanation of the class structure somewhere?

[junrushao]

It seems the direct use of the Database is all wrapped in a JSON-related file object, which kind of confuses me, would you mind providing a brief explanation of the class structure somewhere?

Hey thanks @mbrookhart for asking! Let's focus on the database API and ignore the workload registry part (which will be removed according to previous discussion with Cody in this thread).

• The base abstract class is Database, which has 4 pure virtual methods: commit_workload, commit_tuning_record, get_top_k, size
• The system interacts with any subclass of Database by calling these 4 methods
• To implement a concrete subclass of Database: in pure C++, we can just inherit Database and do so; in pure python, we provide PyDatabase to inherit from, where we plays some tricks, so that everything can be implemented in pure python
• JSONFile is a simple default implementation of Database for open source
• Developers can provide other implementations too by inheriting Database or PyDatabase, for example, database backed by SQLite, etc, as long as the class implements the 4 methods

[mbrookhart]

The first three points and the last one make sense, but I think I'm a little confused by JSONFile being the default user-facing database. JSON means a very particular thing, you convert to/from it, but it doesn't have class methods or a user interface, per say. I would say all of the JSON functionality is important for conversion and serialization, but I'm not sure it makes sense for it to be the primary interface point?

[junrushao]

@mbrookhart Yeah indeed JSONFile can be very confusing given JSON means very particular thing. Would love to hear about your proposals

[junrushao]

@mbrookhart @comaniac I updated the patch according to your comments. Major changes:

• Remove WorkloadRegistry, WorkloadToken and the annoying token_id_
• Rename lookup_or_add => commit_workload
• Rename add => commit_tuning_record
• Various refactors to consolidate more stuff

Let me know what you guys think :-)

[junrushao]

BTW, here is the convention we are using in TensorIR and meta schedule project:

• An abstract base class that describe how the interface looks like, e.g. Schedule, Database, Builder, Runner
• Concrete implementation as subclasses are hidden in src/ folder, and mostly hidden in cc files, which are not directly visible to developers. For example, PyDatabase, ConcreteSchedule
• The abstract class provides factory methods to create these concrete implementations, whose name is consistent with the names of the subclasses. For example:

class Database : ... {
 public:
  static Database JSONFile(...); // TBD: Rename 
  static Database SQLiteDB(...);
};

And then another problem here is what name should we change JSONFile for

[comaniac]

The refactored version is much cleaner :D
Left some comments mostly regarding to the naming. Otherwise LGTM.

[mbrookhart]

I think I agree with @comaniac here, definitely much cleaner now, and I think the only thing I'm still concerned about is naming. Maybe we could just call it JSONDatabase? I think that's a little less ambiguous but still succinct.

Otherwise, LGTM

[junrushao]

Hey thanks guys for the discussion! Let's go with the name JSONDatabase then 😄

[junrushao]

PR updated according to our discussion 👍 @mbrookhart @comaniac please take another look. Thanks guys for the discussion!

[comaniac]

LGTM. Thanks

[mbrookhart]

Thanks, Junru!

[junrushao]

Thanks everybody for the review :-)