gemini-code-assist[bot] commented on code in PR #284:
URL: https://github.com/apache/tvm-ffi/pull/284#discussion_r2555739588
##########
include/tvm/ffi/container/shape.h:
##########
@@ -330,11 +330,17 @@ inline constexpr bool use_default_type_traits_v<Shape> =
false;
// Allow auto conversion from Array<int64_t> to Shape, but not from Shape to
Array<int64_t>
template <>
-struct TypeTraits<Shape> : public ObjectRefWithFallbackTraitsBase<Shape,
Array<int64_t>> {
+struct TypeTraits<Shape> : public ObjectRefWithFallbackTraitsBase<Shape,
Array<int64_t>, DLTensor*> {
static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIShape;
+
TVM_FFI_INLINE static Shape ConvertFallbackValue(Array<int64_t> src) {
return Shape(std::move(src));
}
+
+ TVM_FFI_INLINE static Shape ConvertFallbackValue(DLTensor* src) {
+ std::vector<int64_t> shape_vec(src->shape, src->shape + src->ndim);
+ return Shape(Array<int64_t>(shape_vec));
Review Comment:
The current implementation creates an intermediate `std::vector` and then an
`Array<int64_t>`, resulting in multiple memory allocations and data copies.
This can be optimized by directly constructing the `Shape` from the
`DLTensor`'s shape data, which avoids these intermediate steps.
```c
return Shape(src->shape, src->shape + src->ndim);
```
##########
tests/cpp/test_shape.cc:
##########
@@ -93,4 +93,18 @@ TEST(Shape, ShapeView) {
EXPECT_EQ(view_from_init.Product(), 7 * 8 * 9);
}
+TEST(Shape, DLTensor) {
+ int64_t shape_data[] = {1, 2, 3};
+ DLTensor dltensor;
+ dltensor.ndim = 3;
+ dltensor.shape = shape_data;
+
+ AnyView view_dl = &dltensor;
+ auto shape = view_dl.cast<Shape>();
+ EXPECT_EQ(shape.size(), 3);
+ EXPECT_EQ(shape[0], 1);
+ EXPECT_EQ(shape[1], 2);
+ EXPECT_EQ(shape[2], 3);
+}
Review Comment:
The added test case covers the conversion for a non-empty shape, which is
great. To make the test more comprehensive, I suggest also adding a test case
for an empty shape (when `ndim` is 0). This will ensure the conversion logic
correctly handles this important edge case.
```c
TEST(Shape, DLTensor) {
// Test with a non-empty shape
int64_t shape_data[] = {1, 2, 3};
DLTensor dltensor;
dltensor.ndim = 3;
dltensor.shape = shape_data;
AnyView view_dl = &dltensor;
auto shape = view_dl.cast<Shape>();
EXPECT_EQ(shape.size(), 3);
EXPECT_EQ(shape[0], 1);
EXPECT_EQ(shape[1], 2);
EXPECT_EQ(shape[2], 3);
// Test with an empty shape (ndim = 0)
DLTensor dltensor_empty;
dltensor_empty.ndim = 0;
dltensor_empty.shape = nullptr;
AnyView view_dl_empty = &dltensor_empty;
auto shape_empty = view_dl_empty.cast<Shape>();
EXPECT_EQ(shape_empty.size(), 0);
}
```
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
