Lunderberg commented on code in PR #16721:
URL: https://github.com/apache/tvm/pull/16721#discussion_r1526266543
##########
src/relax/op/tensor/inspect.cc:
##########
@@ -346,6 +356,116 @@ TVM_REGISTER_OP("relax.inspect.tensor_shape_i")
.set_attr<FNormalize>("FNormalize", NormalizeToKnownPrimValue)
.set_attr<Bool>("FPurity", Bool(true));
+//// relax.tensor_stride_i
+
+Expr tensor_stride_i(Expr expr) {
+ static const Op& op = Op::Get("relax.inspect.tensor_stride_i");
+ return Call(op, {expr});
+}
+
+StructInfo InferStructInfoTensorStride(const Call& call, const BlockBuilder&) {
+ auto dlpack_type = DataType::Int(64);
+
+ auto [tensor_sinfo, axis_sinfo] = GetTensorArgInfoWithIndex(call);
+
+ auto opt_tensor_shape = tensor_sinfo->GetShape();
+ auto int_imm_axis = axis_sinfo->value.as<IntImmNode>();
+
+ if (int_imm_axis && opt_tensor_shape.defined()) {
+ // As of 2024-03-14, Relax does not have an explicit
+ // representation for striding in `TensorStructInfo`. The
+ // `FLegalize` function for most operators is implemented in terms
+ // of `topi`, and is then converted from TE to `tir::PrimFunc`
+ // using `tvm::tir::CreatePrimFunc`. The `te::Tensor` is
+ // converted to a `tir::Buffer` in `RewriteStageToBlock`, and uses
+ // the default empty list for the strides. The empty strides
+ // represent a compact data array.
+ //
+ // Therefore, while Relax does not explicitly represent the
+ // striding of a tensor, it implicitly requires compact striding
+ // for any legalizable Tensor.
Review Comment:
Agreed. This was more to state what the current implementation requires.
The upgrade path that I'd imagine would be with a sequence of incremental
changes:
1. Provide ability to read the strides/offset.
2. Insert assert statements (probably while lowering `R.match_cast`) that
the user-provided buffers meet the implicit requirements.
3. Update the `TensorStructInfo` to specify the byte offset and striding.
(I'd default to a contiguous striding, but an arbitrary offset. The offset is
the more useful one, and could be applied before delegating to an external
kernel.)
4. Update the assert statements to match the more flexible
`TensorStructInfo`.
This PR would be step (1) in this sequence.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
