Lunderberg commented on code in PR #17145:
URL: https://github.com/apache/tvm/pull/17145#discussion_r1672412396
##########
python/tvm/relax/op/memory/view.py:
##########
@@ -92,3 +92,20 @@ def _normalize(expr, relax_cls):
relative_byte_offset = _normalize(relative_byte_offset, PrimValue)
return _ffi_api.view(data, shape, dtype, relative_byte_offset) # type:
ignore
+
+
+def ensure_aligned(data: Expr) -> Expr:
Review Comment:
Can we rename this to `ensure_zero_offset`? The alignment for a view
depends on whether `byte_offset % desired_alignment == 0`, not whether
`byte_offset==0`. That is, a buffer with a non-zero `byte_offset` can be
aligned.
##########
src/relax/op/memory/view.cc:
##########
@@ -355,5 +354,28 @@ TVM_REGISTER_OP("relax.memory.view")
.set_attr<FLegalize>("FLegalize", LegalizeView)
.set_attr<Bool>("FPurity", Bool(true));
+Expr ensure_aligned(const Expr& x) {
Review Comment:
Same name change request here, to replace `ensure_aligned` with
`ensure_zero_offset`.
##########
src/runtime/relax_vm/builtin.cc:
##########
@@ -545,6 +545,19 @@
TVM_REGISTER_GLOBAL("vm.builtin.tensor_to_shape").set_body_typed([](NDArray data
return ShapeTuple(out_shape);
});
+TVM_REGISTER_GLOBAL("vm.builtin.ensure_aligned").set_body_typed([](NDArray
data) {
+ if (data->byte_offset == 0) {
+ return data;
+ }
+ DLManagedTensor* dl_tensor = data.ToDLPack();
+ dl_tensor->dl_tensor.data =
+ reinterpret_cast<char*>(dl_tensor->dl_tensor.data) +
dl_tensor->dl_tensor.byte_offset;
+ dl_tensor->dl_tensor.byte_offset = 0;
+ // For platforms that does not support pointer arithmetic, we need to copy
the data to a new
+ // buffer.
+ return NDArray::FromDLPack(dl_tensor);
Review Comment:
This constructs a new `NDArray`, but does not actually copy the data.
Instead, `NDArray::FromDLPack` takes ownership of the `DLManagedTensor*`. The
pointer arithmetic performed earlier would still be invalid.
##########
src/relax/backend/vm/vm_builtin_lower.cc:
##########
@@ -124,6 +128,19 @@ class VMBuiltinLowerMutator : public ExprMutator {
}
}
+ Expr View(const Call& view_node) {
Review Comment:
The handling of `R.memory.view` is already handled by
`relax.transform.LegalizeOps`, and doesn't require any additional handling here.
##########
src/relax/transform/static_plan_block_memory.cc:
##########
@@ -286,8 +286,13 @@ class TokenAllocator1D {
std::vector<StorageToken> full_pool_;
};
-/*! \brief Check if the input op is "relax.reshape". */
-bool IsReshape(const Expr& op) { return op.same_as(Op::Get("relax.reshape")); }
+/*! \brief Check if the input op is a memory op that return the same buffer as
the input buffer. */
Review Comment:
Nitpick: Replace "a memory op that return the same buffer" with "a memory op
that **may** return the same buffer", since it isn't guaranteed whether
`ensure_zero_offset` will return a copy or an alias.
##########
src/runtime/relax_vm/builtin.cc:
##########
@@ -545,6 +545,19 @@
TVM_REGISTER_GLOBAL("vm.builtin.tensor_to_shape").set_body_typed([](NDArray data
return ShapeTuple(out_shape);
});
+TVM_REGISTER_GLOBAL("vm.builtin.ensure_aligned").set_body_typed([](NDArray
data) {
+ if (data->byte_offset == 0) {
+ return data;
+ }
+ DLManagedTensor* dl_tensor = data.ToDLPack();
+ dl_tensor->dl_tensor.data =
Review Comment:
This pointer addition is not valid on all targets. Can we add a method to
the `DeviceAPI` that returns a boolean, indicating whether the host may perform
pointer arithmetic on a `DLTensor::data` pointer?
##########
src/relax/op/memory/view.cc:
##########
@@ -334,13 +334,12 @@ Expr LegalizeView(const BlockBuilder& bb, const Call&
call) {
relative_byte_offset = relax::PrimValue::Int64(0);
}
- StructInfoDeriveFunc infer_sinfo_env_func;
- infer_sinfo_env_func =
EnvFunc::Get("tvm.relax.struct_info.infer_view_sinfo");
- auto runtime_view_sinfo = FuncStructInfo::OpaqueFunc(infer_sinfo_env_func,
true);
-
- ExternFunc runtime_view_func("runtime.TVMArrayCreateView",
runtime_view_sinfo);
+ if (shape.same_as(call->args[1]) && dtype.same_as(call->args[2]) &&
+ relative_byte_offset.same_as(call->args[3])) {
+ return call;
+ }
- return Call(runtime_view_func, {data, shape, dtype, relative_byte_offset});
+ return Call(call->op, {data, shape, dtype, relative_byte_offset});
Review Comment:
This change means that `R.memory.view` is still present in the output of
`LegalizeOps`, but a legalization function should replace the operator with a
lowered form.
##########
tests/python/relax/test_op_view.py:
##########
@@ -483,18 +483,7 @@ def main(A: R.Tensor([4096], "float32")):
class Expected:
@R.function
def main(A: R.Tensor([4096], "float32")):
- B = R.ExternFunc(
- "runtime.TVMArrayCreateView",
- R.Callable(
- derive_func="tvm.relax.struct_info.infer_view_sinfo",
- purity=True,
- ),
- )(
- A,
- R.shape([64, 64]),
- R.dtype("float32"),
- R.prim_value(0),
- )
+ B = R.memory.view(A, shape=R.shape([64, 64]), dtype="float32",
relative_byte_offset=0)
Review Comment:
I don't think this is the desired behavior for `R.memory.view`, as the
operator hasn't actually been legalized.
--
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]
