Adds `UnitAttribute` to root ops for tuner #19345

nithinsubbiah · 2024-12-02T23:04:25Z

Tags rootOp with a UnitAttribute for tuner to pick up. Please note that this attribute isn't used by the compiler at any level, it's only intended to help the tuner.

Signed-off-by: nithinsubbiah <[email protected]>

kuhar

Nice. Can you also add a test case for the strip-compilation-info pass to make sure we don't accidentally remove it there?

hanhanW

Why do we have changes on CPU side? Can we avoid that? My concern is that people need to think about the attribute even they do not use tuners. The other concern is how do you define root op if there is op decomposition happening?

IMO, it is a layering/pluggable issue. People should not have to interpret "root_op" attribute when it is not needed.

kuhar · 2024-12-03T00:17:42Z

Why do we have changes on CPU side? Can we avoid that? My concern is that people need to think about the attribute even they do not use tuners. The other concern is how do you define root op if there is op decomposition happening?

We have a PR in review to add tuning for one of the CPU pipelines as well.

IMO, it is a layering/pluggable issue. People should not have to interpret "root_op" attribute when it is not needed.

What do you mean by 'interpret'? This is a hint for other tooling only, the compiler is not meant to use it any way.

hanhanW · 2024-12-03T00:28:48Z

What do you mean by 'interpret'? This is a hint for other tooling only, the compiler is not meant to use it any way.

By 'interpret', I meant that you're exposing unnecessary information for users who are on the default path. People could start asking like (1) what does it mean? (2) should we take it into account when we decompose an operation? etc. They are new overheads for people who do not have the tooling context (and they probably don't care).

If it is a hint for other tooling only, why do we need to set it in the compiler default path? I think my question is more like what is the rule of attaching "random" attribute to an operation?

nithinsubbiah · 2024-12-03T01:42:19Z

I think the discussion boils down to the question - can we have op attribute that doesn't necessarily get used in the compilation pipeline?

Also given we're building out the tuner infrastructure, which supplements the compiler, this is a good example to hash out how the information exchange can happen between the two. @hanhanW Would love to hear your ideas on that

kuhar · 2024-12-03T01:53:43Z

If it is a hint for other tooling only, why do we need to set it in the compiler default path?

I think this is a practical solution considering the possible alternatives:

Reimplementing the root op identification logic in the tuner. This is a lot of code that has to be kept in sync with IREE.
Exposing root op identification as a public C/python API. This seems less bad to me, but exposes compiler internals that probably shouldn't be exposed.
Annotating root ops. This is the least ambiguous solution and hides any compiler implementation details except for one unit attribute.

We could put it under a flag, but flags add complexity on their own. Adding a new flag would also have to have some justification, and I think in this case the whole system is simpler by adding this unit attr by default. (Compiler flag is even more visible than an attribute and also requires documentation/explanation.)

I think my question is more like what is the rule of attaching "random" attribute to an operation?

This is a good question that I don't have a good answer for. For now, it seems like tuner is the only external tool that we deal with, and after working on it for > 6 months it seems to me like the root op identification is the only piece of information that we miss from configured executables to make the tuner work. My proposal would be to allow the attribute in this case a broader compromise (very low cost to the compiler, very high usefulness to the tuner), and carefully evaluate if we discover we need more than one attribute.

kuhar · 2024-12-03T02:04:30Z

compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_x86_64_lowering_strategy.mlir

@@ -1638,7 +1638,7 @@ func.func @pad_only() attributes {hal.executable.target = #executable_target_emb
 //  CHECK-SAME:     translation_info = #[[TRANSLATION]]
 //      CHECK:   tensor.pad {{.+}} {
 //      CHECK:     tensor.yield
-// CHECK-NEXT:   } {lowering_config = #[[CONFIG]]}


or alternatively, only match lowering_config = #[[COFIG]] in this test to allow for other attributes

hanhanW · 2024-12-03T02:27:15Z

I did not join any design discussion, so I did not put a blocker on the PR. Feel free to ignore my comments if people have considered the case in discussion. 😅

I agree that it is the practical solution, and it is very valuable to see IREE + tuner working. I think we should be careful on default behavior changes, and it is why I chime in. People could start building the components without being aware of that it could impact default path. We could spend months to recover the state or get rid of tech debts. Transform matcher is a good example. The MaterializeHomogeneousEncodings pass is the other example that I'm trying to fix.

The builtin ukernel story is successful to me because we did not mess up any other default path; we figured a way to turn it on by default in IREE. Backends that do not care ukernels do not have to deal with any ukernel context. They even don't see the ukernel attribute floating in the target configuration.

We could put it under a flag, but flags add complexity on their own. Adding a new flag would also have to have some justification, and I think in this case the whole system is simpler by adding this unit attr by default. (Compiler flag is even more visible than an attribute and also requires documentation/explanation.)

I'm +1 on putting it under a flag. I honestly think that "having some justification" is good in this case. It explains to people who do not have context about it. How about having an iree-codegen-annotate-root-op flag (with default false value)? And we only set the unit attribute when the flag is turned on. We're having this kind of flag (being false by default) in the IREE project (e.g., IREE::Stream::createAnnotateAffinitiesPass), and I think we can follow the same idea. This way people understand that the attribute is mainly for tuner development because you have to add a documentation to the flag; other contributors can continue their tuning work. There is no hurt to turn it on by default when we find it is the only solution.

In the comment/documentation, you can say that there are no behavioral changes on the default path, and it is mainly used for tuner project development. What do you think?

iree/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.cpp

Lines 84 to 89 in 006c5d8

    
           // Annotate all ops/resources with the analyzed affinities. 
        
           // This should have no behavioral changes during conversion but allows for 
        
           // debugging of analysis errors in end-user tooling. 
        
           if (clAnnotateInputAffinities) { 
        
             passManager.addPass(IREE::Stream::createAnnotateAffinitiesPass()); 
        
           }

compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/test/root_op_attribute.mlir

MaheshRavishankar

I'd add some warning that this is only used for tuner and is not intended for any other use. Relying on root op for codegen is a bit of a footgun I think.

Signed-off-by: nithinsubbiah <[email protected]>

kuhar

LGTM. If this gets any more complicated in the future, I think we should guard it with a flag and/or rethink the overall approach. But in the current state, I think this is fine as-is.

hanhanW

Thanks, @nithinsubbiah, for ping me to review it. Again, I'm not a blocker because I don't actively work on tuning area.

I'm not approving because I'm -1 on not guarding it over a flag in the first place. Please see the inline comment for the reason. Basically, I want the attribute to evolve in our control and I want to prevent adding overheads to other contributors and reviewers in advance.

hanhanW · 2024-12-03T22:07:05Z

compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.cpp

@@ -568,4 +569,13 @@ void eraseCompilationInfo(Operation *op) {
  op->removeAttr(kCompilationInfoAttrName);
 }

+//===----------------------------------------------------------------------===//
+// Helpers for setting `iree_codegen.root_op` attribute on root operations for


There is no iree_codegen.root_op attribute. What you're doing is adding a unit attribute to the op. Perhaps rephrase it like "Helpers for setting attributes for tuner".

hanhanW · 2024-12-03T22:18:25Z

compiler/src/iree/compiler/Codegen/LLVMGPU/test/root_op_attribute.mlir

I was going to suggest updating the checks in other lit files, then I found that the structure is quite different from LLVMCPU. In LLVMGPU we group lit tests with operation categories, and all the file names start with config_. Maybe we can rename it to config_root_op_attribute.mlir?

hanhanW · 2024-12-03T22:27:49Z

compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h

+/// Sets an attribute to identify the rootOp and adds any information needed for
+/// the tuner from compiler. Currently, only sets a `UnitAttr`. Note that this
+/// attribute is not used by the compiler at any level and is only intended for
+/// tuner use.
+void setRootOpInfo(Operation *op);


It is a deeply hidden warning to me. It could be outdated soon if other IREE contributors start relying on the attribute to do transformations. We could stop them, but it is like overheads++. Because reviewers need to know the information, and the contributors waste some time on the patch.

My expectation was guarding it over a flag in the first place, so it won't evolve without controls. If we start with guarding it over a flag, people who want to use the attribute will to be aware of this information.

hanhanW · 2024-12-03T22:36:20Z

compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h

 /// Convenience function that sets the lowering configuration on the operation
 /// and translation info.
 inline LogicalResult setOpConfigAndEntryPointFnTranslation(
    mlir::FunctionOpInterface entryPointFn, Operation *op,
    IREE::Codegen::LoweringConfigAttrInterface config,
    IREE::Codegen::TranslationInfoAttr translationInfo) {
+  setRootOpInfo(op);


To be more specific -- If it were me, I'd add a flag in this file and wrap this statement with the flag.

MaheshRavishankar · 2024-12-04T00:13:29Z

I think adding a flag that adds the root op annotation is useful. We can leave the flag off by default, but use it when calling from the tuner. Once the tuner is done, it doesnt need to add the root op marker on the final configuration. This will limit the use case to just the tuner. Does that work @kuhar ?

kuhar · 2024-12-04T02:24:12Z

I think adding a flag that adds the root op annotation is useful. We can leave the flag off by default, but use it when calling from the tuner. Once the tuner is done, it doesnt need to add the root op marker on the final configuration. This will limit the use case to just the tuner. Does that work @kuhar ?

OK, I live with a new flag. Yes, the flow is like what you described.

MaheshRavishankar · 2024-12-04T06:34:53Z

This just might be an indication to let the compiler know that it is being involved by the tuner. There might be other places this information would be useful.

Signed-off-by: nithinsubbiah <[email protected]>

hanhanW

LGTM, just two nit. Thanks!

compiler/src/iree/compiler/Codegen/LLVMGPU/test/BUILD.bazel

compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h

Signed-off-by: nithinsubbiah <[email protected]>

nithinsubbiah requested a review from MaheshRavishankar as a code owner December 2, 2024 23:04

nithinsubbiah requested review from kuhar and removed request for MaheshRavishankar December 2, 2024 23:04

nithinsubbiah force-pushed the annotate_rootOp branch from e5d92a8 to 14e9c96 Compare December 2, 2024 23:06

nithinsubbiah and others added 3 commits December 2, 2024 15:08

Add UnitAttribute to root ops for tuner

14e9c96

Signed-off-by: nithinsubbiah <[email protected]>

Add lit test

b1e8726

Signed-off-by: nithinsubbiah <[email protected]>

Fix whitespace

01ff091

nithinsubbiah requested a review from hanhanW as a code owner December 3, 2024 00:01

Update LLVMCPU lit test

8654394

Signed-off-by: nithinsubbiah <[email protected]>

nithinsubbiah mentioned this pull request Dec 3, 2024

[tuner] Reduce maintenance burden and prepare for more codegen pipelines nod-ai/shark-ai#453

Closed

12 tasks

kuhar reviewed Dec 3, 2024

View reviewed changes

hanhanW reviewed Dec 3, 2024

View reviewed changes

hanhanW requested a review from MaheshRavishankar December 3, 2024 00:18

kuhar reviewed Dec 3, 2024

View reviewed changes

hanhanW reviewed Dec 3, 2024

View reviewed changes

compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/test/root_op_attribute.mlir Outdated Show resolved Hide resolved

MaheshRavishankar reviewed Dec 3, 2024

View reviewed changes

nithinsubbiah requested review from qedawkins and Groverkss as code owners December 3, 2024 17:59

Add test case in strip-compilation-info pass

53c15ff

Signed-off-by: nithinsubbiah <[email protected]>

nithinsubbiah requested review from kuhar, MaheshRavishankar and hanhanW December 3, 2024 18:03

kuhar approved these changes Dec 3, 2024

View reviewed changes

nithinsubbiah enabled auto-merge (squash) December 3, 2024 21:47

kuhar disabled auto-merge December 3, 2024 21:48

hanhanW reviewed Dec 3, 2024

View reviewed changes

nithinsubbiah force-pushed the annotate_rootOp branch from 2ed31a0 to 9d91a9d Compare December 4, 2024 20:37

Guard tuner attributes with flag

9d91a9d

Signed-off-by: nithinsubbiah <[email protected]>

nithinsubbiah requested review from kuhar and hanhanW December 4, 2024 20:42

nithinsubbiah force-pushed the annotate_rootOp branch from 47498b2 to 721a9e7 Compare December 4, 2024 20:46

Merge branch 'main' into annotate_rootOp

721a9e7

Signed-off-by: nithinsubbiah <[email protected]>

hanhanW approved these changes Dec 4, 2024

View reviewed changes

compiler/src/iree/compiler/Codegen/LLVMGPU/test/BUILD.bazel Outdated Show resolved Hide resolved

compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h Outdated Show resolved Hide resolved

nithinsubbiah enabled auto-merge (squash) December 4, 2024 21:34

Update test name

2ae4c0a

Signed-off-by: nithinsubbiah <[email protected]>

nithinsubbiah merged commit df34911 into iree-org:main Dec 4, 2024
38 checks passed

kuhar mentioned this pull request Dec 12, 2024

[tuner] Add direct TD spec generation for candidates nod-ai/shark-ai#606

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Adds `UnitAttribute` to root ops for tuner #19345

Adds `UnitAttribute` to root ops for tuner #19345

nithinsubbiah commented Dec 2, 2024 •

edited

Loading

kuhar left a comment

hanhanW left a comment

kuhar commented Dec 3, 2024

hanhanW commented Dec 3, 2024

nithinsubbiah commented Dec 3, 2024 •

edited

Loading

kuhar commented Dec 3, 2024 •

edited

Loading

kuhar Dec 3, 2024

hanhanW commented Dec 3, 2024

MaheshRavishankar left a comment

kuhar left a comment

hanhanW left a comment

hanhanW Dec 3, 2024

hanhanW Dec 3, 2024

hanhanW Dec 3, 2024

hanhanW Dec 3, 2024

MaheshRavishankar commented Dec 4, 2024

kuhar commented Dec 4, 2024

MaheshRavishankar commented Dec 4, 2024

hanhanW left a comment

Adds UnitAttribute to root ops for tuner #19345

Adds UnitAttribute to root ops for tuner #19345

Conversation

nithinsubbiah commented Dec 2, 2024 • edited Loading

kuhar left a comment

Choose a reason for hiding this comment

hanhanW left a comment

Choose a reason for hiding this comment

kuhar commented Dec 3, 2024

hanhanW commented Dec 3, 2024

nithinsubbiah commented Dec 3, 2024 • edited Loading

kuhar commented Dec 3, 2024 • edited Loading

kuhar Dec 3, 2024

Choose a reason for hiding this comment

hanhanW commented Dec 3, 2024

MaheshRavishankar left a comment

Choose a reason for hiding this comment

kuhar left a comment

Choose a reason for hiding this comment

hanhanW left a comment

Choose a reason for hiding this comment

hanhanW Dec 3, 2024

Choose a reason for hiding this comment

hanhanW Dec 3, 2024

Choose a reason for hiding this comment

hanhanW Dec 3, 2024

Choose a reason for hiding this comment

hanhanW Dec 3, 2024

Choose a reason for hiding this comment

MaheshRavishankar commented Dec 4, 2024

kuhar commented Dec 4, 2024

MaheshRavishankar commented Dec 4, 2024

hanhanW left a comment

Choose a reason for hiding this comment

Adds `UnitAttribute` to root ops for tuner #19345

Adds `UnitAttribute` to root ops for tuner #19345

nithinsubbiah commented Dec 2, 2024 •

edited

Loading

nithinsubbiah commented Dec 3, 2024 •

edited

Loading

kuhar commented Dec 3, 2024 •

edited

Loading