gemini-code-assist[bot] commented on code in PR #308:
URL: https://github.com/apache/tvm-ffi/pull/308#discussion_r2586140944
##########
tests/python/test_function.py:
##########
@@ -383,3 +383,22 @@ def __tvm_ffi_float__(self) -> float:
y = fecho(FloatProtocol(10))
assert isinstance(y, float)
assert y == 10
+
+
+def test_function_with_value_protocol() -> None:
+ class ValueProtocol:
+ def __init__(self, value: Any) -> None:
+ self.value = value
+
+ def __tvm_ffi_value__(self) -> Any:
+ return self.value
+
+ fecho = tvm_ffi.get_global_func("testing.echo")
+ assert fecho(ValueProtocol(10)) == 10
+ assert tuple(fecho(ValueProtocol([1, 2, 3]))) == (1, 2, 3)
+ assert tuple(fecho(ValueProtocol([1, 2, ValueProtocol(3)]))) == (1, 2, 3)
+ nested_value_protocol = ValueProtocol(ValueProtocol(ValueProtocol(10)))
+ fecho(nested_value_protocol) == 10
Review Comment:
This line is missing an `assert`. The result of the equality check is not
used, so this doesn't actually test anything.
```suggestion
assert fecho(nested_value_protocol) == 10
```
##########
python/tvm_ffi/cython/tvm_ffi_python_helpers.h:
##########
@@ -394,20 +396,45 @@ class TVMFFIPyCallManager {
return -1;
}
}
+
+ /*!
+ * \brief Set an py_arg to out using the __tvm_ffi_value__ protocol.
+ * \param setter_factory The factory function to create the setter
+ * \param ctx The call context
+ * \param py_arg The python argument to be set
+ * \param out The output argument
+ * \return 0 on success, -1 on failure
+ */
+ TVM_FFI_INLINE int SetArgumentFFIValueProtocol(TVMFFIPyArgSetterFactory
setter_factory,
+ TVMFFIPyCallContext* ctx,
PyObject* py_arg,
+ TVMFFIAny* out) {
+ ++ctx->nested_tvm_ffi_value_count;
+ int ret_code = SetArgument(setter_factory, ctx, py_arg, out);
+ --ctx->nested_tvm_ffi_value_count;
+ // retain the static object if the return value only if it is the
outermost nested level
+ // NOTE: we still keep the invariance that each ArgSetter can have at most
one temporary FFI
+ // object so it ensures that we won't overflow the temporary FFI object
stack
+ if (ret_code == 0 && ctx->nested_tvm_ffi_value_count == 0 &&
+ out->type_index >= kTVMFFIStaticObjectBegin) {
+ if (TVMFFIObjectIncRef(out->v_ptr) != 0) {
+ PyErr_SetString(PyExc_RuntimeError, "Failed to retain temp object");
+ return -1;
+ }
+ // TVMFFIPyPushTempFFIObject(ctx, out->v_ptr);
+ // push the temporary FFI object to the call context
+ ctx->temp_ffi_objects[ctx->num_temp_ffi_objects++] = out->v_ptr;
+ return 0;
+ }
+ return ret_code;
Review Comment:
There's a potential memory corruption bug here due to double-freeing a
temporary object. The `SetArgument` call might invoke a setter (e.g., for
`list` or `tuple`) that creates a temporary FFI object and pushes it to
`ctx->temp_ffi_objects`. This logic then unconditionally increments the
reference count and pushes the same object handle again if it's the outermost
call. This will lead to the object being de-referenced twice when the call
stack is cleaned up, likely causing a crash.
To fix this, you should check if a temporary object has already been pushed
by `SetArgument` before attempting to retain it again. You can do this by
checking the number of temporary objects before and after the `SetArgument`
call. I've also taken the liberty to clean up the code by using
`TVMFFIPyPushTempFFIObject`.
```c
++ctx->nested_tvm_ffi_value_count;
int old_num_temp_ffi_objects = ctx->num_temp_ffi_objects;
int ret_code = SetArgument(setter_factory, ctx, py_arg, out);
--ctx->nested_tvm_ffi_value_count;
// retain the static object if the return value only if it is the
outermost nested level
// NOTE: we still keep the invariance that each ArgSetter can have at
most one temporary FFI
// object so it ensures that we won't overflow the temporary FFI object
stack
if (ret_code == 0 && ctx->nested_tvm_ffi_value_count == 0 &&
out->type_index >= kTVMFFIStaticObjectBegin) {
bool already_retained = ctx->num_temp_ffi_objects >
old_num_temp_ffi_objects;
if (!already_retained) {
if (TVMFFIObjectIncRef(out->v_ptr) != 0) {
PyErr_SetString(PyExc_RuntimeError, "Failed to retain temp
object");
return -1;
}
TVMFFIPyPushTempFFIObject(ctx, out->v_ptr);
}
return 0;
}
return ret_code;
```
--
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]
