Skip to content

[OPTIMIZER] Fix MoveOpAfterLayoutConversion to bail out when there is a Select Op#2678

Merged
Jokeren merged 4 commits into
triton-lang:mainfrom
manman-ren:fix-2655
Nov 21, 2023
Merged

[OPTIMIZER] Fix MoveOpAfterLayoutConversion to bail out when there is a Select Op#2678
Jokeren merged 4 commits into
triton-lang:mainfrom
manman-ren:fix-2655

Conversation

@manman-ren
Copy link
Copy Markdown
Collaborator

Summary: Fix #2655 When there is a Select Op after a Load Op, the type of the operands for the Select Op will be different, we can't use the same newCvtTy for all the operands when creating ConvertLayoutOp.

mlir::TypeID::get<arith::ArithDialect>())
return mlir::failure();
// Also bail out if there exists an op after Load that is Select.
if (isa<arith::SelectOp>(currOp))
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It doesn't seem like the right solution to me, since select is still a layout preserving op.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ThomasRaoux If the number of operands is >=2, should we still move forward the cvtOp? I'm concerning about the cost of extra cvts

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It doesn't seem like the right solution to me, since select is still a layout preserving op.

Yes I'm don't see why select is different than other ops here? That seems arbitrary.

@ThomasRaoux If the number of operands is >=2, should we still move forward the cvtOp? I'm concerning about the cost of extra cvts

It's hard to tell as in general the data should be constant or coming a pipelined load so there wouldn't necessarily be a cvt. If we want more complex heuristic we should probably deal with this as a more global heuristic but without real life examples it is hard to tell.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It feels like this is just working around the problem that the code below assumes all the operand types are the same. I would prefer fixing this code even if we later want to restrict the propagation for heuristic reasons.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sounds good to me. Let's see if @manman-ren would like to take a stab on fixing the problem.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Given that select is preserving the layout, the issue is when generating the ConvertLayoutOp, we use the same newCvtTy while select takes different operand types: i1 vs. another type. We may consider bailing out for select due to heuristic reasons, but for now we should fix the issue by using different newCvtTy for the i1 operand. I will try to fix it.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Jokeren @ThomasRaoux Added a commit on top of the stack, let me how it looks. Thanks!

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM now

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @Jokeren! Also fixed a clang-format issue. It seems to indicate that "Merging can be performed automatically with 1 approving review.", I guess once it is approved, and there are no issues with ci, the PR will be merged. I don't think I have commit access as this is my first PR for Triton.

manman-ren and others added 4 commits November 20, 2023 12:06
Summary: Fix triton-lang#2655
When there is a Select Op after a Load Op, the type of the operands for
the Select Op will be different, we can't use the same newCvtTy for all
the operands when creating ConvertLayoutOp.

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
…and.

This patch uses the corresponding newCvtTy for each operand, instead of using
the same newCvtTy which is from the first operand.
@Jokeren Jokeren changed the title Fix MoveOpAfterLayoutConversion to bail out when there is a Select Op [OPTIMIZER] Fix MoveOpAfterLayoutConversion to bail out when there is a Select Op Nov 21, 2023
@Jokeren Jokeren merged commit e2d3d13 into triton-lang:main Nov 21, 2023
qcolombet pushed a commit to qcolombet/triton that referenced this pull request Nov 24, 2023
… a Select Op (triton-lang#2678)

Summary: Fix triton-lang#2655 When there is
a Select Op after a Load Op, the type of the operands for the Select Op
will be different, we can't use the same newCvtTy for all the operands
when creating ConvertLayoutOp.

---------

Co-authored-by: Manman Ren <mren@meta.com>
Co-authored-by: Manman Ren <mren@fb.com>
@manman-ren manman-ren deleted the fix-2655 branch January 23, 2024 18:57
feihugis pushed a commit to feihugis/triton that referenced this pull request Feb 13, 2024
… a Select Op (triton-lang#2678)

Summary: Fix triton-lang#2655 When there is
a Select Op after a Load Op, the type of the operands for the Select Op
will be different, we can't use the same newCvtTy for all the operands
when creating ConvertLayoutOp.

---------

Co-authored-by: Manman Ren <mren@meta.com>
Co-authored-by: Manman Ren <mren@fb.com>
pingzhuu pushed a commit to siliconflow/triton that referenced this pull request Apr 2, 2024
… a Select Op (triton-lang#2678)

Summary: Fix triton-lang#2655 When there is
a Select Op after a Load Op, the type of the operands for the Select Op
will be different, we can't use the same newCvtTy for all the operands
when creating ConvertLayoutOp.

---------

Co-authored-by: Manman Ren <mren@meta.com>
Co-authored-by: Manman Ren <mren@fb.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

tritongpu-optimize-dot-operands crashes with some arith.select op

4 participants