Harden assertBuffersHaveSameSize to check shapes. #3531
Conversation
| // and a MatmulOp, and StmtSort::getExprs only traverse via the first | ||
| // registered definition (i.e. the Select). This sounds like a bug -- I wonder | ||
| // how nvFuser resets the TensorView uses of a kir::Kernel, also non-SSA. | ||
| hic->addOutput(tva_j_unsqueezed); |
There was a problem hiding this comment.
The introduction of tva_j_unsqueezed triggered a weird problem that @samnordmann is probably aware of. I added more explanation and wonder what @naoyam think about this.
There was a problem hiding this comment.
We could make TensorView live even if no output depends on it for HostIR. Not sure if that would solve the issue, though, as I'm still not entirely clear what the issue is.
There was a problem hiding this comment.
I am aware we artificially need to add the matmul's output as a fusion output to fix the data dependency, that is why tvc_j was added in the first place. However, I was not aware of the other bug you're mentioning -- that we only traverse the first registered producing Expr.
Would the program break if you only let hic->addOutput(tvc_j); ?
There was a problem hiding this comment.
Would the program break if you only let hic->addOutput(tvc_j);?
Yes as I commented at https://github.com/NVIDIA/Fuser/pull/3531/files#diff-30df6421558f87ef0024b01f11752c35d3d68b80a9e6e0ec0fd49de535acb91aR917
There was a problem hiding this comment.
Ok but I am not sure to fully understand the reason why it breaks. Even if the visitor only traverses through the first definition, i.e., the SelectOp, then tvc_j should still be invalidated because the SelectOp consumes the index j
There was a problem hiding this comment.
Yes, tvcj will be invalidated but tvaj unsqueeze won't be. As a result it holds always hold the first iteration value
|
!test |
| // We could have added `tvc_j` instead as an output, which transitively | ||
| // consumes `tva_j_unsqueezed`. However, `tvc_j` has two definitions, a Select | ||
| // and a MatmulOp, and StmtSort::getExprs only traverse via the first | ||
| // registered definition (i.e. the Select). This sounds like a bug -- I wonder |
There was a problem hiding this comment.
StmtSort and other stuff in iter_visiter.h assume the SSA property of Fusion.
There was a problem hiding this comment.
Does that mean TVs in a kir::Kernel (also non-SSA) get wrong Val::uses(), which should be avoided using?
There was a problem hiding this comment.
It may be actually used, but non-SSA definitions in the Kernel IR are pretty limited so far, so we may not encounter any problems. But in general, it isn't a well ironed out use scenario.
There was a problem hiding this comment.
Got it. I'll have to revisit how we evaluate for-loops. One potential approach is to only invalidate loop-index-dependent scalars and let TensorView ops in the loop body run unconditionally.
| // and a MatmulOp, and StmtSort::getExprs only traverse via the first | ||
| // registered definition (i.e. the Select). This sounds like a bug -- I wonder | ||
| // how nvFuser resets the TensorView uses of a kir::Kernel, also non-SSA. | ||
| hic->addOutput(tva_j_unsqueezed); |
There was a problem hiding this comment.
I am aware we artificially need to add the matmul's output as a fusion output to fix the data dependency, that is why tvc_j was added in the first place. However, I was not aware of the other bug you're mentioning -- that we only traverse the first registered producing Expr.
Would the program break if you only let hic->addOutput(tvc_j); ?
I wrote this to make the allgather-related issue discovered in #3284 (comment) easier to expose. And it seems a good runtime check to have in extra, because
_allgather_basetreats I/O tensors as flat buffers and ignores the shapes.