[Util] Add int narrowing pattern for arith.select#23928
[Util] Add int narrowing pattern for arith.select#23928Max191 wants to merge 1 commit intoiree-org:mainfrom
Conversation
Narrow arith.select through index_cast ops when both true/false operands are either index_cast from a narrower type or arith.constant. This handles the common pattern from boolean mask computation: %t = arith.index_cast %a : i32 to index %f = arith.constant 0 : index %s = arith.select %cond, %t, %f : index -> %f_narrow = arith.constant 0 : i32 %s_narrow = arith.select %cond, %a, %f_narrow : i32 %s = arith.index_cast %s_narrow : i32 to index Signed-off-by: Max Dawkins <max.dawkins@gmail.com>
| // Constant with a value that fits in the narrow type. | ||
| auto constOp = dyn_cast<arith::ConstantOp>(defOp); | ||
| if (!constOp) { | ||
| return false; | ||
| } |
There was a problem hiding this comment.
Can we use ConstantLikeInterface maybe?
There was a problem hiding this comment.
I don't think an interface exists for it, but I think I can use the ConstantLike trait. I'll update the implementation.
There was a problem hiding this comment.
There are m_Constant-type matchers.
krzysz00
left a comment
krzysz00
left a comment
There was a problem hiding this comment.
I've got some implementation comments
| if (!trueDef || !falseDef) { | ||
| return failure(); | ||
| } | ||
| Operation *castOp = |
There was a problem hiding this comment.
This could probably work on extsi/extui also, and I colud've sworn there's a nice helper for this.
Also, what you're looking for is isa<IndexCastOp, IndexCastUIOp>(trueDef) and similar for falseDef - if it isa one of these, you've found the castOp.
| narrowElemTy; | ||
| } | ||
| // Constant with a value that fits in the narrow type. | ||
| auto constOp = dyn_cast<arith::ConstantOp>(defOp); |
There was a problem hiding this comment.
You could check for the range being constant here?
There was a problem hiding this comment.
Yeah, I think the check is if the min == max of this value fits in the narrower type.
Also, I think there's code elsewhere around here that handles similar tests for for lops.
| if (defOp->getName() == castOp->getName()) { | ||
| return defOp->getOperand(0); | ||
| } | ||
| auto intAttr = cast<IntegerAttr>(cast<arith::ConstantOp>(defOp).getValue()); |
There was a problem hiding this comment.
I think for this case I'd either follow existing logic for materializing constants in a generic way or I'd just createOrFold the relevant truncation to get the new constant.
|
(my comment is drive-by, non need to wait for me) |
3 participants
Narrow arith.select through index_cast ops when both true/false operands are either index_cast from a narrower type or arith.constant. For example: