slyubomirsky commented on code in PR #16721:
URL: https://github.com/apache/tvm/pull/16721#discussion_r1525640493
##########
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:
It might be worth discussing whether this is the specification we would
want. Perhaps it would be good to have a helper function to assert/convert
tensors into the expected format. Right now, tensors that don't conform to the
expectations about stride/offset are completely "invisible" to Relax other than
using these newly added functions
##########
src/tir/transforms/plan_update_buffer_allocation_location.cc:
##########
@@ -32,21 +32,21 @@
namespace tvm {
namespace tir {
-class CollectUnmanagedAllocations : public StmtExprVisitor {
+class CollectManagedAllocations : public StmtExprVisitor {
Review Comment:
Do these changes relate to the stated purpose of the PR?
##########
src/tir/transforms/lower_tvm_builtin.cc:
##########
@@ -38,6 +38,19 @@ namespace tir {
// These information are needed during codegen.
class BuiltinLower : public StmtExprMutator {
public:
+ static PrimFunc Build(PrimFunc func) {
Review Comment:
Is this also related to the stated purpose of the PR?
--
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]
