gemini-code-assist[bot] commented on code in PR #431: URL: https://github.com/apache/tvm-ffi/pull/431#discussion_r2773467241
########## examples/kernel_library/scale_kernel.cu: ########## @@ -0,0 +1,65 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +#include "tvm_ffi_utils.h" + +// [cuda_kernel.begin] +template <typename T> +__global__ void ScaleKernel(T* out, const T* in, T factor, int64_t n) { + int64_t i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) { + out[i] = in[i] * factor; + } +} +// [cuda_kernel.end] + +// [function.begin] +void Scale(TensorView output, TensorView input, double factor) { + // --- 1. Validate inputs --- + CHECK_INPUT(input); + CHECK_INPUT(output); + CHECK_DIM(1, input); + CHECK_DEVICE(input, output); + TVM_FFI_CHECK(input.numel() == output.numel(), ValueError) << "input/output size mismatch"; Review Comment:  The validation logic checks that the input and output tensors are on the same device and have the same number of elements, but it doesn't verify that they have the same `dtype`. Since the `ScaleKernel` is templated on a single type `T` for both input and output, a `dtype` mismatch would lead to incorrect behavior or memory corruption. It's important to add a check to ensure `input.dtype() == output.dtype()`. ``` CHECK_DEVICE(input, output); TVM_FFI_CHECK(input.dtype() == output.dtype(), ValueError) << "input/output dtype mismatch"; TVM_FFI_CHECK(input.numel() == output.numel(), ValueError) << "input/output size mismatch"; ``` ########## examples/kernel_library/tvm_ffi_utils.h: ########## @@ -0,0 +1,71 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +#ifndef KERNEL_LIBRARY_TVM_FFI_UTILS_H_ +#define KERNEL_LIBRARY_TVM_FFI_UTILS_H_ + +#include <tvm/ffi/extra/c_env_api.h> +#include <tvm/ffi/extra/cuda/device_guard.h> +#include <tvm/ffi/tvm_ffi.h> + +namespace ffi = tvm::ffi; +using ffi::Optional; +using ffi::Tensor; +using ffi::TensorView; + +// [check_macros.begin] +// --- Reusable validation macros --- +#define CHECK_CUDA(x) \ + TVM_FFI_CHECK((x).device().device_type == kDLCUDA, ValueError) << #x " must be a CUDA tensor" +#define CHECK_CONTIGUOUS(x) \ + TVM_FFI_CHECK((x).IsContiguous(), ValueError) << #x " must be contiguous" +#define CHECK_INPUT(x) \ + CHECK_CUDA(x); \ + CHECK_CONTIGUOUS(x) Review Comment:  This macro expands to two separate statements. This can lead to unexpected behavior when used in control flow statements without braces, like `if (condition) CHECK_INPUT(x);`. In this case, only the first check (`CHECK_CUDA`) would be conditional. To make the macro robust, it's a standard C++ practice to wrap multi-statement macros in a `do { ... } while(0)` block, which turns them into a single statement. ```suggestion #define CHECK_INPUT(x) \ do { \ CHECK_CUDA(x); \ CHECK_CONTIGUOUS(x); \ } while (0) ``` ########## examples/kernel_library/tvm_ffi_utils.h: ########## @@ -0,0 +1,71 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +#ifndef KERNEL_LIBRARY_TVM_FFI_UTILS_H_ +#define KERNEL_LIBRARY_TVM_FFI_UTILS_H_ + +#include <tvm/ffi/extra/c_env_api.h> +#include <tvm/ffi/extra/cuda/device_guard.h> +#include <tvm/ffi/tvm_ffi.h> + +namespace ffi = tvm::ffi; +using ffi::Optional; +using ffi::Tensor; +using ffi::TensorView; + +// [check_macros.begin] +// --- Reusable validation macros --- +#define CHECK_CUDA(x) \ + TVM_FFI_CHECK((x).device().device_type == kDLCUDA, ValueError) << #x " must be a CUDA tensor" +#define CHECK_CONTIGUOUS(x) \ + TVM_FFI_CHECK((x).IsContiguous(), ValueError) << #x " must be contiguous" +#define CHECK_INPUT(x) \ + CHECK_CUDA(x); \ + CHECK_CONTIGUOUS(x) +#define CHECK_DIM(d, x) \ + TVM_FFI_CHECK((x).ndim() == (d), ValueError) << #x " must be a " #d "D tensor" +#define CHECK_DEVICE(a, b) \ + TVM_FFI_CHECK((a).device().device_type == (b).device().device_type, ValueError) \ + << #a " and " #b " must be on the same device type"; \ + TVM_FFI_CHECK((a).device().device_id == (b).device().device_id, ValueError) \ + << #a " and " #b " must be on the same device" Review Comment:  Similar to `CHECK_INPUT`, this macro expands to multiple statements, which can be problematic inside control structures without braces (e.g., `if`). Additionally, there's an inconsistent use of semicolons within the macro definition. To ensure it behaves as a single, safe statement, it should be wrapped in a `do { ... } while(0)` block, and the internal semicolon should be removed for consistency. ```suggestion #define CHECK_DEVICE(a, b) \ do { \ TVM_FFI_CHECK((a).device().device_type == (b).device().device_type, ValueError) \ << #a " and " #b " must be on the same device type"; \ TVM_FFI_CHECK((a).device().device_id == (b).device_id, ValueError) \ << #a " and " #b " must be on the same device"; \ } while (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]
