Create an MLIR environment with matrix multiplication #652
Conversation
Co-authored-by: kyle <kyle.w.herndon@gmail.com>
facebook-github-bot
added
the
CLA Signed
Codecov Report
@@ Coverage Diff @@
## development #652 +/- ##
===============================================
- Coverage 88.42% 82.13% -6.30%
===============================================
Files 119 120 +1
Lines 7475 7550 +75
===============================================
- Hits 6610 6201 -409
- Misses 865 1349 +484
Continue to review full report at Codecov.
|
|
One thing to note is that the CI from the MLIR env is not setup yet. |
|
Woohoo! Here comes the big PR! 😄 There's a lot to unpack here. I'll make a start today, but it might take a couple of rounds for me to go through it all. Great work @sogartar et al! Cheers, |
|
For the CI, could you split the existing Cheers, |
compiler_gym/envs/mlir/__init__.py
Outdated
| # TODO(boian): Make better config | ||
| # observation_spaces = {"runtime": "Runtime"} | ||
| # reward_spaces = {"runtime": "Runtime"} | ||
| # for reward_space in reward_spaces: | ||
| # register( | ||
| # id=f"mlir-{reward_space}-v0", | ||
| # entry_point="compiler_gym.envs.mlir:MlirEnv", | ||
| # kwargs={ | ||
| # "service": MLIR_SERVICE_BINARY, | ||
| # "reward_space": reward_spaces[reward_space], | ||
| # }, | ||
| # ) | ||
| # | ||
| # for observation_space, reward_space in product(observation_spaces, reward_spaces): | ||
| # register( | ||
| # id=f"mlir-{observation_space}-{reward_space}-v0", | ||
| # entry_point="compiler_gym.envs.mlir:MlirEnv", | ||
| # kwargs={ | ||
| # "service": MLIR_SERVICE_BINARY, | ||
| # "observation_space": observation_spaces[observation_space], | ||
| # "reward_space": reward_spaces[reward_space], | ||
| # }, | ||
| # ) |
There was a problem hiding this comment.
Should this be committed?
There was a problem hiding this comment.
@KyleHerndon is working on this, but it is a big change and it will be in another PR.
| def benchmark_uris(self) -> Iterable[str]: | ||
| return (self.name_from_size(mnk) for mnk in matmul_sizes) | ||
|
|
||
| def benchmark(self, uri: str) -> MatmulBenchmark: |
There was a problem hiding this comment.
I believe you could use benchmark_from_parsed_uri() instead of benchmark() to save you the string parsing here
There was a problem hiding this comment.
This part is going to be changed in another PR.
|
Hey, Sorry for the delay in the full review, I just moved to a new office and am setting up a new dev environment. I'm hitting this error, not sure if it's related to your changes or my own fault: Have you seen this before? I started reading through the code and left a bunch of minor comments, but I think I'll hold off until I get a chance to build and test the new environment before I go any further. I don't want to spam you with nitpicks until I understand the big picture 🙂 Note to self: I got as far as Cheers, |
|
@ChrisCummins, I added the external dependencies of the MLIR env to the python packaging. It required the clang compiler, Google benchmark and |
compiler_gym/wrappers/core.py
Outdated
| @@ -371,3 +376,70 @@ def multistep( | |||
| def convert_reward(self, reward): | |||
| """Translate a reward to the new space.""" | |||
| raise NotImplementedError | |||
|
|
|||
|
|
|||
| class ConversionWrapperEnv(CompilerEnvWrapper): | |||
There was a problem hiding this comment.
I'm not sure about this. I like that the wrappers in base gym operate only on action/observation/reward at a time. Could you add the space conversion methods to the relevant wrapper and use them instead? E.g. merge convert_observation_space into ObservationWrapper. You could then have:
env = MlirObservationWrapper(MlirActionWrapper(MlirRewardWrapper(env)))
|
Okay, I have completed a full pass of the PR. Minor nitpicks left inline, but in general, this looks excellent. Many thanks @sogartar and @KyleHerndon! While this isn't merge blocking, I'd like to discuss with you folks your goals / timeline is for making the MLIR environment a part of the "core" installable package. In particular, I want to avoid this kind of thing: if config.enable_llvm_env:
...
if config.enable_mlir_env:
...As I understand, it would be totally possible to build a combined LLVM + MLIR environment package, it would just require a couple of passes of the cmake build and merging the generated runfiles, right? Is the idea to split the cmake build into separate projects for each environment so that they can have mutually exclusive dependencies? If for some reason there are unavoidable extra dependencies that MLIR requires that complicates the build, could we not still have the python class definition available, and defer the >>> compiler_gym.make("mlir-v0")
EnvironmentNotSupported("Please install compiler_gym[mlir] to use this environment")Again, this doesn't need to be resolved in this PR, assuming that you're okay with the pip wheel continuing to be LLVM-only for the time being. One final request that is blocking is for you to add a new page in the documentation describing this environment. Take a look at envs/gcc for an example. Since it is still early days, it would be fine to say that its under construction and that the doc can be updated as more features are added, but it's important to have the basic there. Cheers, |
…pported action spaces
I added the
I am also not happy with the logic of the conditional inclusion of the LLVM and MLIR environments in the package. It needs refactoring. Having multiple subpackages that can be separately built and installed with pip seems a good idea. I think it would be beneficial to have also a development package that has all stuff related to creating new environments. Like the C++ code in The packages could look like that @ChrisCummins, I will let you know in a few days when we are going to be able to work on that. |
|
Thanks for the latest iteration of changes, it's coming together great.
Ah, what I meant was a new rst doc in
Yep, I'll continue discussion in #669. Cheers, |
I added some documentation under |
|
@ChrisCummins, did you got a chance to see the documentation? |
This release adds a new compiler environment, new APIs, and a suite of backend improvements to improve the flexibility of CompilerGym environments. Many thanks to code contributors: @sogartar, @KyleHerndon, @SoumyajitKarmakar, @uduse, and @anthony0727! Highlights of this release include: - [mlir] Began work on a new environment for matrix multiplication using MLIR ([#652](#652), thanks @KyleHerndon and @sogartar!). Note this environment is not yet included in the pypi package and must be [compiled from source](https://github.com/facebookresearch/CompilerGym/blob/development/INSTALL.md#building-from-source-with-cmake). - [llvm] Added a new `env.benchmark_from_clang_invocation()` method ([#577](#577)) that can be used for constructing LLVM environment automatically from C/C++ compiler invocations. This makes it much easier to integrate CompilerGym with your existing build scripts. - Added three new wrapper classes: `Counter`, that provides op counts for analysis ([#683](#683)); `SynchronousSqliteLogger`, that provides logging of environment interactions to a relational database ([#679](#679)), and `ForkOnStep` that provides an `undo()` operation ([#682](#682)). - Added `reward_space` and `observation_space` parameters to `env.reset()` ([#659](#659), thanks @SoumyajitKarmakar!) This release includes a number of improvements to the backend APIs that make it easier to write new CompilerGym environments: - Refactored the backend to make `CompilerEnv` an abstract interface, and `ClientServiceCompilerEnv` the concrete implementation of this interface. This enables new environments to be implemented without using gRPC ([#633](#633), thanks @sogartar!). - Extended the support for different types of action and observation spaces ([#641](#641), [#643](#643), thanks @sogartar!), including new `Permutation` and `SpaceSequence` spaces ([#645](#645), thanks @sogartar!).. - Added a new `disk/` subdirectory to compiler service's working directories, which is symlinked to an on-disk location for devices which support in-memory working directories. This fixes a bug with leftover temporary directories from LLVM ([#672](#672)). This release also includes numerous bug fixes and improvements, many of which were reported or fixed by the community. For example, fixing a bug in cache file locations ([#656](#656), thanks @uduse!), and a missing flag definition in example code ([#684](#684), thanks @anthony0727!). **Full Changelog**: v0.2.3...v0.2.4 This release brings in deprecating changes to the core `env.step()` routine, and lays the groundwork for enabling new types of compiler optimizations to be exposed through CompilerGym. Many thanks to code contributors: @mostafaelhoushi, @sogartar, @KyleHerndon, @uduse, @parthchadha, and @xtremey! Highlights of this release include: - Added a new `TextSizeInBytes` observation space for LLVM ([#575](#575)). * Added a new PPO leaderboard entry ([#580](#580). Thanks @xtremey! - Fixed a bug in which temporary directories created by the LLVM environment were not cleaned up ([#592](#592)). - **[Backend]** The function `createAndRunCompilerGymService` now returns an int, which is the exit return code ([#592](#592)). - Improvements to the examples documentation ([#548](#548)) and FAQ ([#586](#586)) Deprecations and breaking changes: - `CompilerEnv.step` no longer accepts a list of actions ([#627](#627)). A new method, `CompilerEnv.multistep` provides this functionality. This is to provide compatibility with environments whose action spaces are lists. To update your code, replace any calls to `env.step()` which take a list of actions to use `env.multistep()`. Thanks @sogartar! - The arguments `observations` and `rewards` to `step()` have been renamed `observation_spaces` and `reward_spaces`, respectively ([#627](#627)). - `Reward.id` has been renamed `Reward.name` ([#565](#565), [#612](#612)). Thanks @parthchadha! * The backend protocol buffer schema has been updated to natively support more types of observation and action, and to support nested spaces ([#531](#531)). Thanks @sogartar!
This release adds a new compiler environment, new APIs, and a suite of backend improvements to improve the flexibility of CompilerGym environments. Many thanks to code contributors: @sogartar, @KyleHerndon, @SoumyajitKarmakar, @uduse, and @anthony0727! Highlights of this release include: - [mlir] Began work on a new environment for matrix multiplication using MLIR ([#652](#652), thanks @KyleHerndon and @sogartar!). Note this environment is not yet included in the pypi package and must be [compiled from source](https://github.com/facebookresearch/CompilerGym/blob/development/INSTALL.md#building-from-source-with-cmake). - [llvm] Added a new `env.benchmark_from_clang_invocation()` method ([#577](#577)) that can be used for constructing LLVM environment automatically from C/C++ compiler invocations. This makes it much easier to integrate CompilerGym with your existing build scripts. - Added three new wrapper classes: `Counter`, that provides op counts for analysis ([#683](#683)); `SynchronousSqliteLogger`, that provides logging of environment interactions to a relational database ([#679](#679)), and `ForkOnStep` that provides an `undo()` operation ([#682](#682)). - Added `reward_space` and `observation_space` parameters to `env.reset()` ([#659](#659), thanks @SoumyajitKarmakar!) This release includes a number of improvements to the backend APIs that make it easier to write new CompilerGym environments: - Refactored the backend to make `CompilerEnv` an abstract interface, and `ClientServiceCompilerEnv` the concrete implementation of this interface. This enables new environments to be implemented without using gRPC ([#633](#633), thanks @sogartar!). - Extended the support for different types of action and observation spaces ([#641](#641), [#643](#643), thanks @sogartar!), including new `Permutation` and `SpaceSequence` spaces ([#645](#645), thanks @sogartar!).. - Added a new `disk/` subdirectory to compiler service's working directories, which is symlinked to an on-disk location for devices which support in-memory working directories. This fixes a bug with leftover temporary directories from LLVM ([#672](#672)). This release also includes numerous bug fixes and improvements, many of which were reported or fixed by the community. For example, fixing a bug in cache file locations ([#656](#656), thanks @uduse!), and a missing flag definition in example code ([#684](#684), thanks @anthony0727!). **Full Changelog**: v0.2.3...v0.2.4 This release brings in deprecating changes to the core `env.step()` routine, and lays the groundwork for enabling new types of compiler optimizations to be exposed through CompilerGym. Many thanks to code contributors: @mostafaelhoushi, @sogartar, @KyleHerndon, @uduse, @parthchadha, and @xtremey! Highlights of this release include: - Added a new `TextSizeInBytes` observation space for LLVM ([#575](#575)). * Added a new PPO leaderboard entry ([#580](#580). Thanks @xtremey! - Fixed a bug in which temporary directories created by the LLVM environment were not cleaned up ([#592](#592)). - **[Backend]** The function `createAndRunCompilerGymService` now returns an int, which is the exit return code ([#592](#592)). - Improvements to the examples documentation ([#548](#548)) and FAQ ([#586](#586)) Deprecations and breaking changes: - `CompilerEnv.step` no longer accepts a list of actions ([#627](#627)). A new method, `CompilerEnv.multistep` provides this functionality. This is to provide compatibility with environments whose action spaces are lists. To update your code, replace any calls to `env.step()` which take a list of actions to use `env.multistep()`. Thanks @sogartar! - The arguments `observations` and `rewards` to `step()` have been renamed `observation_spaces` and `reward_spaces`, respectively ([#627](#627)). - `Reward.id` has been renamed `Reward.name` ([#565](#565), [#612](#612)). Thanks @parthchadha! * The backend protocol buffer schema has been updated to natively support more types of observation and action, and to support nested spaces ([#531](#531)). Thanks @sogartar!
This release adds a new compiler environment, new APIs, and a suite of backend improvements to improve the flexibility of CompilerGym environments. Many thanks to code contributors: @sogartar, @KyleHerndon, @SoumyajitKarmakar, @uduse, and @anthony0727! Highlights of this release include: - [mlir] Began work on a new environment for matrix multiplication using MLIR ([#652](#652), thanks @KyleHerndon and @sogartar!). Note this environment is not yet included in the pypi package and must be [compiled from source](https://github.com/facebookresearch/CompilerGym/blob/development/INSTALL.md#building-from-source-with-cmake). - [llvm] Added a new `env.benchmark_from_clang_invocation()` method ([#577](#577)) that can be used for constructing LLVM environment automatically from C/C++ compiler invocations. This makes it much easier to integrate CompilerGym with your existing build scripts. - Added three new wrapper classes: `Counter`, that provides op counts for analysis ([#683](#683)); `SynchronousSqliteLogger`, that provides logging of environment interactions to a relational database ([#679](#679)), and `ForkOnStep` that provides an `undo()` operation ([#682](#682)). - Added `reward_space` and `observation_space` parameters to `env.reset()` ([#659](#659), thanks @SoumyajitKarmakar!) This release includes a number of improvements to the backend APIs that make it easier to write new CompilerGym environments: - Refactored the backend to make `CompilerEnv` an abstract interface, and `ClientServiceCompilerEnv` the concrete implementation of this interface. This enables new environments to be implemented without using gRPC ([#633](#633), thanks @sogartar!). - Extended the support for different types of action and observation spaces ([#641](#641), [#643](#643), thanks @sogartar!), including new `Permutation` and `SpaceSequence` spaces ([#645](#645), thanks @sogartar!).. - Added a new `disk/` subdirectory to compiler service's working directories, which is symlinked to an on-disk location for devices which support in-memory working directories. This fixes a bug with leftover temporary directories from LLVM ([#672](#672)). This release also includes numerous bug fixes and improvements, many of which were reported or fixed by the community. For example, fixing a bug in cache file locations ([#656](#656), thanks @uduse!), and a missing flag definition in example code ([#684](#684), thanks @anthony0727!). **Full Changelog**: v0.2.3...v0.2.4 This release brings in deprecating changes to the core `env.step()` routine, and lays the groundwork for enabling new types of compiler optimizations to be exposed through CompilerGym. Many thanks to code contributors: @mostafaelhoushi, @sogartar, @KyleHerndon, @uduse, @parthchadha, and @xtremey! Highlights of this release include: - Added a new `TextSizeInBytes` observation space for LLVM ([#575](#575)). * Added a new PPO leaderboard entry ([#580](#580). Thanks @xtremey! - Fixed a bug in which temporary directories created by the LLVM environment were not cleaned up ([#592](#592)). - **[Backend]** The function `createAndRunCompilerGymService` now returns an int, which is the exit return code ([#592](#592)). - Improvements to the examples documentation ([#548](#548)) and FAQ ([#586](#586)) Deprecations and breaking changes: - `CompilerEnv.step` no longer accepts a list of actions ([#627](#627)). A new method, `CompilerEnv.multistep` provides this functionality. This is to provide compatibility with environments whose action spaces are lists. To update your code, replace any calls to `env.step()` which take a list of actions to use `env.multistep()`. Thanks @sogartar! - The arguments `observations` and `rewards` to `step()` have been renamed `observation_spaces` and `reward_spaces`, respectively ([#627](#627)). - `Reward.id` has been renamed `Reward.name` ([#565](#565), [#612](#612)). Thanks @parthchadha! * The backend protocol buffer schema has been updated to natively support more types of observation and action, and to support nested spaces ([#531](#531)). Thanks @sogartar!
This also adds config flags to enable the MLIR and LLVM environments. They are mutually exclusive due to their dependence on different LLVM versions and the problem of their coexistence in one CMake configuration.
Enabling them is available through CMake with the flags
The LLVM env is enabled by default.
These flags are propagated to python where sub-modules are included conditionally.