Address comments

Author: larryliu0820

backends/aoti/common_shims.cpp

@@ -50,36 +50,14 @@ AOTITorchError aoti_torch_get_storage_offset(
 }
 
 AOTITorchError aoti_torch_get_strides(Tensor* tensor, int64_t** ret_strides) {
-  auto it = internal::tensor_to_strides.find(tensor);
-  bool needs_update = false;
-
-  if (it == internal::tensor_to_strides.end()) {
-    needs_update = true;
-  } else {
-    // Check if cached values are still valid
-    auto tensor_strides = tensor->strides();
-    if (it->second.size() != static_cast<size_t>(tensor->dim())) {
-      needs_update = true;
-    } else {
-      for (int i = 0; i < tensor->dim(); i++) {
-        if (it->second[i] != tensor_strides[i]) {
-          needs_update = true;
-          break;
-        }
-      }
-    }
-  }
-
-  if (needs_update) {
-    std::vector<int64_t> strides(tensor->dim());
-    auto tensor_strides = tensor->strides();
-    for (int i = 0; i < tensor->dim(); i++) {
-      strides[i] = tensor_strides[i];
-    }
-    it =
-        internal::tensor_to_strides.insert_or_assign(tensor, std::move(strides))
-            .first;
+  std::vector<int64_t> strides(tensor->dim());
+  auto tensor_strides = tensor->strides();
+  for (ssize_t i = 0; i < tensor->dim(); i++) {
+    strides[i] = static_cast<int64_t>(tensor_strides[i]);
   }
+  auto it =
+      internal::tensor_to_strides.insert_or_assign(tensor, std::move(strides))
+          .first;
 
   // For 0D tensors, data() returns nullptr on empty vectors, but we need to
   // return a valid pointer
@@ -100,35 +78,13 @@ AOTITorchError aoti_torch_get_dtype(Tensor* tensor, int32_t* ret_dtype) {
 }
 
 AOTITorchError aoti_torch_get_sizes(Tensor* tensor, int64_t** ret_sizes) {
-  auto it = internal::tensor_to_sizes.find(tensor);
-  bool needs_update = false;
-
-  if (it == internal::tensor_to_sizes.end()) {
-    needs_update = true;
-  } else {
-    // Check if cached values are still valid
-    auto tensor_sizes = tensor->sizes();
-    if (it->second.size() != static_cast<size_t>(tensor->dim())) {
-      needs_update = true;
-    } else {
-      for (int i = 0; i < tensor->dim(); i++) {
-        if (it->second[i] != tensor_sizes[i]) {
-          needs_update = true;
-          break;
-        }
-      }
-    }
-  }
-
-  if (needs_update) {
-    std::vector<int64_t> sizes(tensor->dim());
-    auto tensor_sizes = tensor->sizes();
-    for (int i = 0; i < tensor->dim(); i++) {
-      sizes[i] = tensor_sizes[i];
-    }
-    it = internal::tensor_to_sizes.insert_or_assign(tensor, std::move(sizes))
-             .first;
+  std::vector<int64_t> sizes(tensor->dim());
+  auto tensor_sizes = tensor->sizes();
+  for (ssize_t i = 0; i < tensor->dim(); i++) {
+    sizes[i] = static_cast<int64_t>(tensor_sizes[i]);
   }
+  auto it = internal::tensor_to_sizes.insert_or_assign(tensor, std::move(sizes))
+                .first;
 
   // For 0D tensors, data() returns nullptr on empty vectors, but we need to
   // return a valid pointer

backends/cuda/runtime/cuda_backend.cpp

@@ -170,7 +170,8 @@ class ET_EXPERIMENTAL CudaBackend final
     // static/singleton across the whole process. When we share multiple methods
     // (meaning multiple so_handle) in the same process, we need to re-register
     // the symbols from the so_handle that is being used in this execution.
-    register_shared_library_functions(handle->so_handle);
+    ET_CHECK_OK_OR_RETURN_ERROR(
+        register_shared_library_functions(handle->so_handle));
 
     size_t n_inputs;
     AOTInductorModelContainerGetNumInputs(handle->container_handle, &n_inputs);

extension/llm/runner/multimodal_prefiller.cpp

@@ -67,11 +67,11 @@ Result<uint64_t> MultimodalPrefiller::prefill(
           InvalidArgument,
           "Model expects uint8_t image data, but image has float data.");
     } else {
-      ET_LOG(
-          Error,
+      ET_CHECK_OR_RETURN_ERROR(
+          false,
+          NotSupported,
           "Unsupported image encoder input dtype: %s",
           ::executorch::runtime::toString(expected_dtype));
-      return ::executorch::runtime::Error::NotSupported;
     }
 
     // The model might expect a 4D tensor (NCHW), but toTensor() returns a 3D
@@ -119,12 +119,12 @@ Result<uint64_t> MultimodalPrefiller::prefill(
             convert_to_bfloat16(audio_tensor),
             "Failed to convert audio tensor to bfloat16");
       } else {
-        ET_LOG(
-            Error,
+        ET_CHECK_OR_RETURN_ERROR(
+            false,
+            NotSupported,
             "Unsupported audio encoder input dtype: %s. Expecting %s",
             ::executorch::runtime::toString(audio_tensor->scalar_type()),
             ::executorch::runtime::toString(expected_dtype));
-        return ::executorch::runtime::Error::NotSupported;
       }
     }
 

extension/llm/runner/util.h

@@ -145,7 +145,7 @@ inline runtime::Result<TensorPtr> populate_start_pos_or_cache_position(
  * Helper function to convert a float tensor to bfloat16.
  * Creates a new tensor with bfloat16 dtype and copies/converts the data.
  */
-::executorch::runtime::Result<::executorch::extension::TensorPtr>
+inline ::executorch::runtime::Result<::executorch::extension::TensorPtr>
 convert_to_bfloat16(const ::executorch::extension::TensorPtr& src_tensor) {
   ET_CHECK_OR_RETURN_ERROR(
       src_tensor->scalar_type() == ::executorch::aten::ScalarType::Float,