[OPTIMIZER] Fix MoveOpAfterLayoutConversion to bail out when there is a Select Op#2678
Conversation
| 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)) |
There was a problem hiding this comment.
It doesn't seem like the right solution to me, since select is still a layout preserving op.
There was a problem hiding this comment.
@ThomasRaoux If the number of operands is >=2, should we still move forward the cvtOp? I'm concerning about the cost of extra cvts
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Sounds good to me. Let's see if @manman-ren would like to take a stab on fixing the problem.
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
@Jokeren @ThomasRaoux Added a commit on top of the stack, let me how it looks. Thanks!
There was a problem hiding this comment.
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.
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.
… 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>
… 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>
… 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>
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.