Gather op implementation [#1015] #1151
Conversation
ddilbazTT
requested review from
jnie-TT,
kmabeeTT,
AleksKnezevic,
pilkicTT,
sdjordjevicTT,
svuckovicTT,
mtopalovicTT,
rpavlovicTT,
jserbedzijaTT,
nobradovictt,
nsmithtt,
mrakitaTT and
tapspatel
as code owners
lib/Conversion/TTIRToTTIRDecomposition/TTIRToTTIRDecomposition.cpp
Outdated
Show resolved
Hide resolved
ddilbazTT
force-pushed
the
ddilbaz/gather
branch
2 times, most recently
from
c05933b
to
b833e79
Compare
| : public OpConversionPattern<ttir::GatherOp> { | ||
| using OpConversionPattern<ttir::GatherOp>::OpConversionPattern; | ||
|
|
||
| LogicalResult checkBasicLegality(ttir::GatherOp op, |
There was a problem hiding this comment.
Is this something that operation verifiers can't catch?
There was a problem hiding this comment.
We should use verifiers for these scenarios. The only place I think it makes sense to do it outside of the verifiers is in dialects that we don't control, like StableHLO - I don't know if there's a better way to do it though.
There was a problem hiding this comment.
I think the intention was to keep stablehlo to ttir conversion layer as thin as possible. Whilst using verifiers, we would mark valid stablehlo gather ops as invalid since they cannot be lowered into embedding. We followed the TTIR to TTIR Pattern Matching approach. I'm not sure if one approach is better than the other.
There was a problem hiding this comment.
@sdjordjevicTT / @svuckovicTT, this is during conversion and is checking to see if the matchAndRewrite (see below) is possible, i.e. we're able to legally convert the op, which is separate from verification.
There was a problem hiding this comment.
This one is especially complicated because TTIR gather is a legal op in TTIR, but this rewrite pattern is trying to decomp into embedding op. This is because stablehlo does not have embedding, it's using the more general op gather.
lib/Conversion/TTIRToTTIRDecomposition/TTIRToTTIRDecomposition.cpp
Outdated
Show resolved
Hide resolved
lib/Conversion/TTIRToTTIRDecomposition/TTIRToTTIRDecomposition.cpp
Outdated
Show resolved
Hide resolved
ddilbazTT
force-pushed
the
ddilbaz/gather
branch
from
f406619
to
fddcfcd
Compare
ddilbazTT
force-pushed
the
ddilbaz/gather
branch
from
fddcfcd
to
7dcf489
Compare
| : public OpConversionPattern<ttir::GatherOp> { | ||
| using OpConversionPattern<ttir::GatherOp>::OpConversionPattern; | ||
|
|
||
| LogicalResult checkBasicLegality(ttir::GatherOp op, |
There was a problem hiding this comment.
@sdjordjevicTT / @svuckovicTT, this is during conversion and is checking to see if the matchAndRewrite (see below) is possible, i.e. we're able to legally convert the op, which is separate from verification.
| : public OpConversionPattern<ttir::GatherOp> { | ||
| using OpConversionPattern<ttir::GatherOp>::OpConversionPattern; | ||
|
|
||
| LogicalResult checkBasicLegality(ttir::GatherOp op, |
There was a problem hiding this comment.
This one is especially complicated because TTIR gather is a legal op in TTIR, but this rewrite pattern is trying to decomp into embedding op. This is because stablehlo does not have embedding, it's using the more general op gather.
Gather op is lowered into embedding. Used TTIR pass 38a4a46. Used Jackson's embedding fixes from e798a17 PR. Blocked by tt-metal issue 14584.