[WIP][Performance Remarks] Support `MLIR_ENABLE_REMARK` for command line tools by sfzhu93 · Pull Request #5250 · triton-lang/triton

sfzhu93 · 2024-11-25T08:14:16Z

This PR enhances Triton's command-line tools, such as triton-opt, by adding support for the MLIR_ENABLE_REMARK option. It introduces a singleton class to manage the emission of performance warnings, leveraging the MLIR infrastructure for diagnostic information (e.g., op->emitWarning() and op->emitRemark()). This approach replaces the previous approach in #3922 and #4350 by using a macro based on the new wrapper class instead of directly emitting remarks.

Initially, we considered introducing a switch to disable emitRemark in MLIR. However, similar implementations, like the -Rpass and -W switches in Clang, directly control information from specific passes. Therefore, we opted to create a wrapper class for the MLIR APIs rather than modifying upstream MLIR to support disabling certain diagnostic information.

This PR is still ongoing. 1. We may need to update related Python code to utilize this new wrapper class. This work is still in progress. 2. Also, unit test is necessary to check if warnings are really disabled.

New contributor declaration

I am not making a trivial change, such as fixing a typo in a comment.
I have written a PR description following these
rules.
I have run pre-commit run --from-ref origin/main --to-ref HEAD.
Select one of the following.
- I have added tests.
  - /test for lit tests
  - /unittest for C++ tests
  - /python/test for end-to-end tests
- This PR does not need a test because FILL THIS IN.
Select one of the following.
- I have not added any lit tests.
- The lit tests I have added follow these best practices,
  including the "tests should be minimal" section. (Usually running Python code
  and using the instructions it generates is not minimal.)

ThomasRaoux

I wonder why we can't handle this by setting the right checks in SourceMgrErrorDiagnosticHandler maybe @Mogball has an idea.

ThomasRaoux · 2024-11-25T08:19:48Z

+class DiagnosticEmitter {
+  // singleton pattern
+private:
+  inline static DiagnosticEmitter *instance{nullptr};
+  bool shouldEmitPerfWarning;
+  DiagnosticEmitter() : shouldEmitPerfWarning(false){};
+  ~DiagnosticEmitter() = default;
+
+public:
+  DiagnosticEmitter(const DiagnosticEmitter &) = delete;
+  DiagnosticEmitter &operator=(const DiagnosticEmitter &) = delete;
+
+  static DiagnosticEmitter *getInstance() {
+    if (!instance) {
+      instance = new DiagnosticEmitter();
+      if (tools::getBoolEnv("MLIR_ENABLE_REMARK")) {
+        instance->shouldEmitPerfWarning = true;
+      }
+    }
+    return instance;
+  }
+
+  static void setShouldEmitPerfWarning(bool shouldEmit) {
+    DiagnosticEmitter::getInstance()->shouldEmitPerfWarning = shouldEmit;
+  }
+
+  static std::optional<InFlightDiagnostic>
+  getPerfWarningStream(const OpState &op) {
+    if (DiagnosticEmitter::getInstance()->shouldEmitPerfWarning) {
+      return op->emitRemark();
+    } else {
+      return std::nullopt;
+    }
+  }


This is not going to be thread safe. We should stay away from globals or singleton

ThomasRaoux · 2024-11-25T08:19:59Z

+#define EMIT_PERF_WARNING(op, message)                                         \
+  if (auto out = mlir::triton::DiagnosticEmitter::getPerfWarningStream(op)) {  \
+    *out << message;                                                           \
+  }


why does it need to be a macro?

was trying to make an API supporting emitPerfWarn(op) << "foo" << "bar" but find it difficult to make it print nothing when the switch is off. Making an empty InFlightDiagnostic triggers an assertion in MLIR when it attempts to attach notes.

sfzhu93 · 2024-11-25T08:38:54Z

I wonder why we can't handle this by setting the right checks in SourceMgrErrorDiagnosticHandler maybe @Mogball has an idea.

We did so in the compiler.py. Personally I prefer this checks as well. I have thought about introducing a customized MLIR main function following what's in https://github.com/llvm/llvm-project/blob/main/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp in triton-opt. My main concern is that the main function would probably have 50-100 more LOC and most code simply repeats what's in MlirOptMain.cpp. So I introduced a wrapper class which could be useful for future customizations as well.

I can move on customizing the main function for triton-opt if you think that is better.

Mogball · 2024-11-25T15:45:55Z

If this is just meant for CLI tools driven by MlirOptMain, it might be easier to fix this upstream: we can add a config option to MlirOptMainConfig that controls suppressing warnings, remarks, etc.

IMO a macro is quite awkward and it's easy to forget to use the macro instead of the normal MLIR API. Forking MlirOptMain.cpp is also an option -- most of the code in there probably isn't useful for Triton.

sfzhu93 · 2024-11-25T19:05:13Z

@Mogball Thanks for the suggestion! I'll look into making an update to MLIR upstream to address this.

Forking MlirOptMain.cpp is also an option -- most of the code in there probably isn't useful for Triton.

Could you elaborate on what you mean by forking MlirOptMain.cpp? Are you suggesting that we maintain a separate/partial copy of it within Triton?

Mogball · 2024-11-25T19:06:44Z

@Mogball Thanks for the suggestion! I'll look into making an update to MLIR upstream to address this.

Forking MlirOptMain.cpp is also an option -- most of the code in there probably isn't useful for Triton.

Could you elaborate on what you mean by forking MlirOptMain.cpp? Are you suggesting that we maintain a separate/partial copy of it within Triton?

Yeah, a partial copy.

If you have an upstream PR, please share the link here and I'll TAL

sfzhu93 · 2024-11-26T05:26:17Z

llvm/llvm-project#117669 @Mogball
tested locally, both with and without the argument. This test gives a lot of vectorization warnings and now we can suppress.

./triton-opt ~/triton/test/Conversion/tritongpu_to_llvm.mlir -split-input-file --allocate-shared-memory --convert-triton-gpu-to-llvm --mlir-print-op-on-diagnostic=false --mlir-disable-diagnostic > ~/tmp.txt

This PR adds a command line argument `--mlir-disable-diagnostic` for disabling diagnostic information for mlir-opt. When debugging with mlir-opt, some developers would like to disable the diagnostic information and focus specifically on the dumped IR. For example, triton-lang/triton#5250

sfzhu93 · 2024-11-28T02:55:50Z

Closing this PR, as llvm/llvm-project#117669 is merging. After that, we should be able to update the upstream LLVM commit hash to address the problem.

init commit

80c9586

ThomasRaoux reviewed Nov 25, 2024

View reviewed changes

sfzhu93 mentioned this pull request Nov 26, 2024

[MLIR][mlir-opt] add support for disabling diagnostics llvm/llvm-project#117669

Merged

sfzhu93 closed this Nov 28, 2024

sfzhu93 deleted the triton-opt-add-cmd-arg branch December 3, 2024 02:25

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[WIP][Performance Remarks] Support `MLIR_ENABLE_REMARK` for command line tools#5250

[WIP][Performance Remarks] Support `MLIR_ENABLE_REMARK` for command line tools#5250
sfzhu93 wants to merge 1 commit intotriton-lang:mainfrom
sfzhu93:triton-opt-add-cmd-arg

sfzhu93 commented Nov 25, 2024

Uh oh!

ThomasRaoux left a comment

Uh oh!

ThomasRaoux Nov 25, 2024

Uh oh!

ThomasRaoux Nov 25, 2024

Uh oh!

sfzhu93 Nov 25, 2024

Uh oh!

sfzhu93 commented Nov 25, 2024 •

edited

Loading

Uh oh!

Mogball commented Nov 25, 2024

Uh oh!

sfzhu93 commented Nov 25, 2024

Uh oh!

Mogball commented Nov 25, 2024

Uh oh!

sfzhu93 commented Nov 26, 2024 •

edited

Loading

Uh oh!

sfzhu93 commented Nov 28, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

sfzhu93 commented Nov 25, 2024

New contributor declaration

Uh oh!

ThomasRaoux left a comment

Choose a reason for hiding this comment

Uh oh!

ThomasRaoux Nov 25, 2024

Choose a reason for hiding this comment

Uh oh!

ThomasRaoux Nov 25, 2024

Choose a reason for hiding this comment

Uh oh!

sfzhu93 Nov 25, 2024

Choose a reason for hiding this comment

Uh oh!

sfzhu93 commented Nov 25, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Mogball commented Nov 25, 2024

Uh oh!

sfzhu93 commented Nov 25, 2024

Uh oh!

Mogball commented Nov 25, 2024

Uh oh!

sfzhu93 commented Nov 26, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

sfzhu93 commented Nov 28, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

sfzhu93 commented Nov 25, 2024 •

edited

Loading

sfzhu93 commented Nov 26, 2024 •

edited

Loading