Add Python bindings for override mechanism. #1526

vcanicTT · 2024-12-06T15:01:23Z

Add initial version of Python lib for optimizer overrides. This library is just pybind wrapper around C++ class for optimizer overrides in MLIR.

nobradovictt · 2024-12-06T15:07:49Z

include/ttmlir/Bindings/Python/TTMLIRModule.h

 void populatePassesModule(py::module &m);
+


Nit Extra line

nobradovictt · 2024-12-06T15:08:43Z

include/ttmlir/Dialect/TTNN/Utils/OptimizerOverrides.h

 private:
  // Options for the TTIR to TTNN backend pipeline,
  // we use them to extract the names and the deafulat values.
-  TTIRToTTNNBackendPipelineOptions pipelineOptions;
+  // TTIRToTTNNBackendPipelineOptions pipelineOptions;


Removed. Done.

nobradovictt · 2024-12-06T15:09:00Z

include/ttmlir/Dialect/TTNN/Utils/PassOverrides.h

@@ -13,6 +13,19 @@

 namespace mlir::tt::ttnn {

+struct OptionNames {
+


Nit Extra line

nobradovictt · 2024-12-06T15:12:46Z

include/ttmlir/Dialect/TTNN/Utils/PassOverrides.h

@@ -13,6 +13,19 @@

 namespace mlir::tt::ttnn {

+struct OptionNames {


Why did you switch to this helper struct for names? Names are already defined in pipeline options, previous solution was thus better it didn't require redefining in multiple places in code when adding new option?

We can't bind LLVM structures (class) and their virtual methods. LLVM has separate RTTI. Done.

nobradovictt · 2024-12-06T15:14:14Z

lib/Dialect/TTNN/Utils/PassOverrides.cpp

@@ -6,6 +6,18 @@

 namespace mlir::tt::ttnn {

+const std::string OptionNames::optimizerPassEnabled = "enable-optimizer";


See my other comment above for these new constants.

nobradovictt · 2024-12-06T15:19:25Z

lib/Dialect/TTNN/Utils/PassOverrides.cpp

+const std::string OptionNames::overrideOutputLayout = "override-output-layout";
+const std::string OptionNames::memoryLayoutAnalysisEnabled =
+    "memory-layout-analysis-enabled";
+const std::string OptionNames::memReconfigEnabled = "memreconfig-enabled";


odjuricicTT

Looks good!

odjuricicTT · 2024-12-06T15:29:53Z

include/ttmlir/Dialect/TTNN/Utils/OptimizerOverrides.h

@@ -65,10 +68,22 @@ class OptimizerOverridesHandler {
                               TensorMemoryLayout, tt::ttnn::Layout,
                               tt::DataType);

+  // Wrapper methods we use to expose the adders to the python bindings
+  std::unordered_map<std::string, InputLayoutOverrideParams>
+  getInputLayoutOverridesPybindWrapper() const;


What is the reason to create wrappers instead of changing the existing methods?

C++ maps bind to Python dict, it's easier to bind and deal with C++ types instead of LLVM. I use wrappers only for bindings, within the MLIR C++ you can use both of them.

odjuricicTT · 2024-12-06T15:32:58Z

lib/Dialect/TTNN/Utils/PassOverrides.cpp

@@ -6,6 +6,18 @@

 namespace mlir::tt::ttnn {

+const std::string OptionNames::optimizerPassEnabled = "enable-optimizer";


I'm guessing that there is a problem when reading these from pipeline options?

Yes. We can't bind class TTIRToTTNNBackendPipelineOptions, we build python lib with RTTI, and if the structure has or inherits some virtual methods, something wrong happens with the linker. We can talk more about this with Stefi and Topalko, about what we can do in the future, but I think we can't do too much. If we want to expose something from MLIR to Python in this way, we always need to have some simple(C++) structure, or to extend pybind and LLVM. There is also, llvmpy lib, but often they have the same issues with the linking.

odjuricicTT · 2024-12-06T15:33:30Z

lib/Dialect/TTNN/Utils/PassOverrides.cpp

@@ -6,6 +6,18 @@

 namespace mlir::tt::ttnn {

+const std::string OptionNames::optimizerPassEnabled = "enable-optimizer";
+const std::string OptionNames::overrideInputLayout = "insert-memreconfig";
+const std::string OptionNames::overrideOutputLayout = "override-output-layout";


Nit: Might be cleaner just to have these set in the .h file?

vprajapati-tt

Bindings LGTM, glad to see you got around the issues :) Unsure if we want to hold both overrides and optimizer_overrides, there isn't anything planned for the overrides module at the moment.

vcanicTT · 2024-12-08T09:27:47Z

Bindings LGTM, glad to see you got around the issues :) Unsure if we want to hold both overrides and optimizer_overrides, there isn't anything planned for the overrides module at the moment.

I'm not sure about that module, you can remove it.

Optimizer teams want to have separate module for all kinds of overrides we use through the optimizer. If we want to extend override mechanism, we need to design it better and write documentation before we start the implementation.

nobradovictt

Thanks for addressing feedback promptly! Enjoy your vacation!

* Add Python bindings for override mechanism. * Add Python bindings for override mechanism. * Add Python bindings for override mechanism.

vcanicTT requested review from nsmithtt, tapspatel, vprajapati-tt, odjuricicTT, sdjordjevicTT, svuckovicTT, mtopalovicTT, rpavlovicTT, nobradovictt, jserbedzijaTT and jnie-TT as code owners December 6, 2024 15:01

vcanicTT force-pushed the vcanic/override_mechanism_pybind branch from 1ef0a35 to 10b7547 Compare December 6, 2024 15:06

nobradovictt suggested changes Dec 6, 2024

View reviewed changes

nobradovictt reviewed Dec 6, 2024

View reviewed changes

odjuricicTT reviewed Dec 6, 2024

View reviewed changes

vprajapati-tt approved these changes Dec 6, 2024

View reviewed changes

vcanicTT added 3 commits December 8, 2024 09:31

Add Python bindings for override mechanism.

22c6c4e

Add Python bindings for override mechanism.

dc2fa71

Add Python bindings for override mechanism.

458f00a

vcanicTT force-pushed the vcanic/override_mechanism_pybind branch from a95db1d to 458f00a Compare December 8, 2024 09:31

nobradovictt approved these changes Dec 8, 2024

View reviewed changes

nsmithtt requested a review from vprajapati-tt December 9, 2024 02:04

vcanicTT merged commit 223d244 into main Dec 9, 2024
19 checks passed

azecevicTT pushed a commit that referenced this pull request Dec 17, 2024

Add Python bindings for override mechanism. (#1526)

57a0eb5

* Add Python bindings for override mechanism. * Add Python bindings for override mechanism. * Add Python bindings for override mechanism.

mtopalovicTT deleted the vcanic/override_mechanism_pybind branch December 31, 2024 08:45

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add Python bindings for override mechanism. #1526

Add Python bindings for override mechanism. #1526

vcanicTT commented Dec 6, 2024

nobradovictt Dec 6, 2024

vcanicTT Dec 8, 2024

nobradovictt Dec 6, 2024

vcanicTT Dec 8, 2024

nobradovictt Dec 6, 2024

vcanicTT Dec 8, 2024

nobradovictt Dec 6, 2024

vcanicTT Dec 8, 2024

nobradovictt Dec 6, 2024

vcanicTT Dec 8, 2024

nobradovictt Dec 6, 2024

vcanicTT Dec 8, 2024

odjuricicTT left a comment

odjuricicTT Dec 6, 2024

vcanicTT Dec 8, 2024

odjuricicTT Dec 6, 2024

vcanicTT Dec 8, 2024

odjuricicTT Dec 6, 2024

vcanicTT Dec 8, 2024

vprajapati-tt left a comment

vcanicTT commented Dec 8, 2024

nobradovictt left a comment

		@@ -13,6 +13,19 @@

		namespace mlir::tt::ttnn {

		struct OptionNames {

		@@ -6,6 +6,18 @@

		namespace mlir::tt::ttnn {

		const std::string OptionNames::optimizerPassEnabled = "enable-optimizer";

Add Python bindings for override mechanism. #1526

Add Python bindings for override mechanism. #1526

Conversation

vcanicTT commented Dec 6, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

odjuricicTT left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

vprajapati-tt left a comment

Choose a reason for hiding this comment

vcanicTT commented Dec 8, 2024

nobradovictt left a comment

Choose a reason for hiding this comment