gemini-code-assist[bot] commented on code in PR #19762:
URL: https://github.com/apache/tvm/pull/19762#discussion_r3408996283
##########
src/runtime/extra/contrib/coreml/coreml_runtime.mm:
##########
@@ -154,32 +155,19 @@
NSDictionary* json = [NSJSONSerialization JSONObjectWithData:data
options:NSJSONReadingAllowFragments
error:nil];
- NSffi::Array<NSString*>* input_names = json[@"inputs"];
+ NSArray<NSString*>* input_names = json[@"inputs"];
// Copy input tensors to corresponding data entries.
for (auto i = 0; i < args.size() - 1; ++i) {
- TVM_FFI_ICHECK(args[i].type_code() == kTVMDLTensorHandle ||
- args[i].type_code() == kTVMTensorHandle)
- << "Expect Tensor or DLTensor as inputs\n";
- if (args[i].type_code() == kTVMDLTensorHandle || args[i].type_code()
== kTVMTensorHandle) {
- model_->SetInput([input_names[i] UTF8String], args[i]);
- } else {
- LOG(FATAL) << "Not implemented";
- }
+ model_->SetInput([input_names[i] UTF8String],
args[i].cast<DLTensor*>());
}
Review Comment:
If the metadata JSON is invalid or missing the `inputs` key, `json` or
`input_names` will be `nil`. In Objective-C, calling `UTF8String` on a `nil`
object returns `nil`, which when converted to `std::string` will cause a crash.
Additionally, if `args.size()` is 0, `args.size() - 1` will underflow (since
`size()` is unsigned), leading to an out-of-bounds loop. We should validate
`json`, `input_names`, and the argument count to ensure safety.
```
NSDictionary* json = [NSJSONSerialization JSONObjectWithData:data
options:NSJSONReadingAllowFragments
error:nil];
TVM_FFI_ICHECK(json != nil) << "Failed to parse metadata JSON";
NSArray<NSString*>* input_names = json[@"inputs"];
TVM_FFI_ICHECK(input_names != nil) << "Metadata JSON does not contain
'inputs'";
TVM_FFI_ICHECK_EQ(args.size(), [input_names count] + 1)
<< "Expect " << [input_names count] + 1 << " arguments, but got "
<< args.size();
// Copy input tensors to corresponding data entries.
for (auto i = 0; i < args.size() - 1; ++i) {
model_->SetInput([input_names[i] UTF8String],
args[i].cast<DLTensor*>());
}
```
##########
src/runtime/extra/contrib/coreml/coreml_runtime.mm:
##########
@@ -196,7 +184,7 @@
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef().def_packed("tvm.coreml_runtime.create", [](ffi::PackedArgs
args, ffi::Any* rv) {
- *rv = CoreMLRuntimeCreate(args[0], args[1]);
+ *rv = CoreMLRuntimeCreate(args[0].cast<std::string>(),
args[1].cast<std::string>());
});
Review Comment:
To prevent potential out-of-bounds access or crashes when the frontend
passes an incorrect number of arguments, we should validate that `args.size()`
is exactly 2 before accessing `args[0]` and `args[1]`.
```
refl::GlobalDef().def_packed("tvm.coreml_runtime.create",
[](ffi::PackedArgs args, ffi::Any* rv) {
TVM_FFI_ICHECK_EQ(args.size(), 2) << "Expect 2 arguments for
tvm.coreml_runtime.create";
*rv = CoreMLRuntimeCreate(args[0].cast<std::string>(),
args[1].cast<std::string>());
});
```
##########
src/runtime/extra/contrib/coreml/coreml_runtime.mm:
##########
@@ -136,12 +137,12 @@
return ffi::Function([this](ffi::PackedArgs args, ffi::Any* rv) {
model_->Invoke(); });
} else if (name == "set_input") {
return ffi::Function([this](ffi::PackedArgs args, ffi::Any* rv) {
- const auto& input_name = args[0].operator std::string();
- model_->SetInput(input_name, args[1]);
+ model_->SetInput(args[0].cast<std::string>(), args[1].cast<DLTensor*>());
});
} else if (name == "get_output") {
- return ffi::Function(
- [this](ffi::PackedArgs args, ffi::Any* rv) { *rv =
model_->GetOutput(args[0]); });
+ return ffi::Function([this](ffi::PackedArgs args, ffi::Any* rv) {
+ *rv = model_->GetOutput(args[0].cast<int>());
+ });
Review Comment:
To prevent potential out-of-bounds access or crashes when the frontend
passes an incorrect number of arguments, we should validate that `args.size()`
is exactly 1 before accessing `args[0]`.
```
return ffi::Function([this](ffi::PackedArgs args, ffi::Any* rv) {
TVM_FFI_ICHECK_EQ(args.size(), 1) << "Expect 1 argument for
get_output";
*rv = model_->GetOutput(args[0].cast<int>());
});
```
##########
src/runtime/extra/contrib/coreml/coreml_runtime.mm:
##########
@@ -136,12 +137,12 @@
return ffi::Function([this](ffi::PackedArgs args, ffi::Any* rv) {
model_->Invoke(); });
} else if (name == "set_input") {
return ffi::Function([this](ffi::PackedArgs args, ffi::Any* rv) {
- const auto& input_name = args[0].operator std::string();
- model_->SetInput(input_name, args[1]);
+ model_->SetInput(args[0].cast<std::string>(), args[1].cast<DLTensor*>());
});
Review Comment:
To prevent potential out-of-bounds access or crashes when the frontend
passes an incorrect number of arguments, we should validate that `args.size()`
is exactly 2 before accessing `args[0]` and `args[1]`.
```
return ffi::Function([this](ffi::PackedArgs args, ffi::Any* rv) {
TVM_FFI_ICHECK_EQ(args.size(), 2) << "Expect 2 arguments for
set_input";
model_->SetInput(args[0].cast<std::string>(),
args[1].cast<DLTensor*>());
});
```
--
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]
