[Relax] Express dynamic arguments of strided_slice as arguments #16826
Conversation
src/relax/op/tensor/index.cc
Outdated
| // PrimExpr bounds_offset = tvm::if_then_else(stride < 0, -1, 0); | ||
| // index = tvm::min(tvm::max(index, 0 + bounds_offset), extent + bounds_offset); |
There was a problem hiding this comment.
Probably should just be deleted, I assume
There was a problem hiding this comment.
Thank you, and deleted!
| // TODO(Lunderberg): Implement this check using `IsBaseOf`. Doing | ||
| // so will require a way to represent a `relax::TupleStructInfo` of | ||
| // unknown length, where each element has the same `StructInfo`. |
There was a problem hiding this comment.
Hm, I think the idea of having a list type has come up before. I think we've had lists via Objects and PackedFuncs before, probably for gradient. Reifying it in the type system could be feasible if it's a common enough use case.
|
@Archermmt Can you look into the failing MSC unit tests? It seems that it is a limitation in the Relax to MSC conversion, that it doesn't correctly handle in-line leaf nodes that occur within a |
Prior to this commit, `relax.op.strided_slice` stored the `axes`, `begin`, `end`, and `strides` in the `CallNode::attrs`. However, the attributes are only intended to store static values. The indices used used for `relax.op.strided_slice` must frequently be in terms of symbolic shape variables, which should not be stored in the attributes. While some utilities have special handling for `relax.op.strided_slice` (e.g. `tvm::relax::Bind`), many do not (e.g. `tvm::relax::WellFormed` and `tvm::relax::FreeSymbolicVars`). As a result, the symbolic expressions in `relax.op.strided_slice` will fail to be updated in generic utilities, and will fail to trigger safeguards when this occurs. This commit changes the representation of `relax.op.strided_slice` to store all arguments in the `relax::CallNode::args`, rather than the `relax::CallNode::attrs`. As mentioned in a comment from apache#13987, which initially implemented `relax.op.strided_slice`, this was an intended refactor once `relax::PrimValue` was fully supported.
Lunderberg
force-pushed
the
relax_use_primvalue_for_strided_slice
branch
from
8e4dd20 to
f596467
Compare
|
Merged from main into PR branch to resolve conflict. |
|
CI is passing, and this PR is approved. I'm going to give one final merge from main into this PR, as there have been a couple of recent changes in |
|
CI passed with main, and a quick smoke test with MLC-LLM passed as well, so it's time to merge. |
Prior to this commit,
relax.op.strided_slicestored theaxes,begin,end, andstridesin theCallNode::attrs. However, the attributes are only intended to store static values. The indices used used forrelax.op.strided_slicemust frequently be in terms of symbolic shape variables, which should not be stored in the attributes. While some utilities have special handling forrelax.op.strided_slice(e.g.tvm::relax::Bind), many do not (e.g.tvm::relax::WellFormedandtvm::relax::FreeSymbolicVars). As a result, the symbolic expressions inrelax.op.strided_slicewill fail to be updated in generic utilities, and will fail to trigger safeguards when this occurs.This commit changes the representation of
relax.op.strided_sliceto store all arguments in therelax::CallNode::args, rather than therelax::CallNode::attrs. As mentioned in a comment from #13987, which initially implementedrelax.op.strided_slice, this was an intended refactor oncerelax::PrimValuewas fully supported.