[compiler] strip execution context affinities in const eval #18663
Conversation
During compile-time constant evaluation in pass iree-consteval-jit-globals it does not make sense to assign device/queue affinities. We will be compiling and executing it on the compilation host. The JITed IR is first stripped of all attributes that assign affinities and all flow.tensor.transfer ops. Signed-off-by: Boian Petkantchin <[email protected]>
Signed-off-by: Boian Petkantchin <[email protected]>
sogartar
force-pushed
the
jit-globals-strip-device-affinities
branch
from
0f6be21
to
5a294e4
Compare
|
Adding @ScottTodd to reviewers because he reviewed those device transfer codes before. |
ScottTodd
added
the
compiler/dialects
| if (failed(stripExecutionContextAffinities( | ||
| programBuilder.getTargetModule()))) { | ||
| return signalPassFailure(); | ||
| } |
There was a problem hiding this comment.
Seems like this code could instead go in importInitializer(). Do we have a preference for making the import function take on more responsibility or treating this as a cleanup step?
…tInitializer Signed-off-by: Boian Petkantchin <[email protected]>
sogartar
force-pushed
the
jit-globals-strip-device-affinities
branch
from
be72e43
to
5b0910f
Compare
|
|
||
| LogicalResult matchAndRewrite(IREE::Stream::AffinityOpInterface op, | ||
| PatternRewriter &rewriter) const override { | ||
| // Shouldn't we reject ops for which `op.requiresAffinity() == true`? |
There was a problem hiding this comment.
if you're confused when you're writing this then someone coming across this in the future will be doubly confused
There was a problem hiding this comment.
I removed the comment. I was not sure if removing this attribute when it is required may result in an illegal operation.
There was a problem hiding this comment.
I guess since we are attaching the interface to Flow ops, having the attribute removed, must not lead to illegality.
… numerics (#237) Verifies the IREE module numerical accuracy compared to execution with PyTorch. The prefill step result has very low absolute accuracy of around `1e-2` for FP32. The resulting cache state of prefill is way off. The decode step accuracy is also completely off. This test is market as skipped until iree-org/iree#18663 is merged. Without it the IREE compilation will crash.
|
@benvanik, will you be able to review again? |
There was a problem hiding this comment.
This should not be relying on specific ops and instead needs to walk attributes and replace in there. Any op in the IR can have an affinity assigned to it. You can use AttrTypeReplacer to walk and replace any AffinityAttr on any ops. Don't use pattern matching - just walk the ops and then use recursivelyReplaceElementsIn. Also make sure to call it on the function itself.
|
@benvanik, will you be able to review this again? |
During compile-time constant evaluation in pass
iree-consteval-jit-globals it does not make sense to assign device/queue affinities. We will be compiling and executing it on the compilation host.
The JITed IR is first stripped of all attributes that assign affinities and all flow.tensor.transfer ops.