Skip to content
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.

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

Plumb automatic schedulers into Python interface #1145

Open
jacobhinkle opened this issue Oct 25, 2023 · 3 comments
Open

Plumb automatic schedulers into Python interface #1145

jacobhinkle opened this issue Oct 25, 2023 · 3 comments
Labels
enhancement New feature or request

Comments

@jacobhinkle
Copy link
Collaborator

jacobhinkle commented Oct 25, 2023

For autotuning, and for developing heuristics, it would be convenient to be able to schedule complicated fusions from the Python side. Currently, doing so requires us to write the whole schedule from the primitives like split, parallelize, etc. However, it would be great to be able to create in Python some surrogate of, say a PointwiseParams or MatmulParams object and call schedulePointwise or scheduleMatmul with it. The key is that this would let us manipulate the heuristic in order to explore different regimes of our automatic schedulers, without needing to reproduce the complicated logic in the schedule* functions.

Each param object inherits from HeuristicParams but appends a number of custom attributes, some of which are compound types. For example, MatmulParams contains, among other things, a MatMulTileOptions struct that contains multiple GemmTile structs each holding a triple of ints. We could represent this with a dict in Python, but we will need to maintain a translation function for each type heuristic parameter object.

Instead we could try to devise a system where we could automatically convert from a nested dictionary of simple types. For example, each of the structs would declare names and types for their members using a macro that would generate the conversion code.

struct GemmTile : public FromDict {
 FROM_DICT_INT_ATTR(m);
 FROM_DICT_INT_ATTR(n);
 FROM_DICT_INT_ATTR(k);
  // rest of definition unmodified 
}

struct MatMulTileOptions : public FromDict {
  FROM_DICT_CLASS_ATTR(GemmTile, cta_tile, 128, 128, 32);
  FROM_DICT_CLASS_ATTR(GemmTile, warp_tile, 64, 64, 32);
  FROM_DICT_CLASS_ATTR(GemmTile, instruction_tile, 16, 8, 16);
  // ...
}

// HeuristicParams would inherit from FromDict
struct MatmulParams : public HeuristicParams {
  // ...
  FROM_DICT_BOOL_ATTR(async_gmem_load_operands, false);
  FROM_DICT_CLASS_ATTR(MatMulTileOptions, tile_sizes);
  // ...
}

// Then we'd automatically generate a static method
// MatmulParams MatmulParams::fromDict(const std::unordered_map<std::string, std::variant<...>>& d);

Plumbing these in to the python frontend would then be very simple, and presumably this would also make serde simple for the heuristic info.

Is it possible to create such a class FromDict and associated macros? Is this pattern already implemented somewhere?

@jacobhinkle jacobhinkle added the enhancement New feature or request label Oct 25, 2023
@jacobhinkle
Copy link
Collaborator Author

cc @kevinstephano

@rdspring1 rdspring1 pinned this issue Nov 5, 2023
@rdspring1 rdspring1 unpinned this issue Nov 5, 2023
@jacobhinkle
Copy link
Collaborator Author

See also the more recent issue #2418

@jacobhinkle
Copy link
Collaborator Author

Fixed by #3106

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

No branches or pull requests

1 participant