| 27 | #include <tvm/ffi/function.h> |
| 28 | |
| 29 | void add_one_cpu(tvm::ffi::TensorView x, tvm::ffi::TensorView y) { |
| 30 | TVM_FFI_ICHECK(x.ndim() == 1) << "x must be a 1D tensor"; |
| 31 | DLDataType f32_dtype{kDLFloat, 32, 1}; |
| 32 | TVM_FFI_ICHECK(x.dtype() == f32_dtype) << "x must be a float tensor"; |
| 33 | TVM_FFI_ICHECK(y.ndim() == 1) << "y must be a 1D tensor"; |
| 34 | TVM_FFI_ICHECK(y.dtype() == f32_dtype) << "y must be a float tensor"; |
| 35 | TVM_FFI_ICHECK(x.size(0) == y.size(0)) << "x and y must have the same shape"; |
| 36 | for (int i = 0; i < x.size(0); ++i) { |
| 37 | static_cast<float*>(y.data_ptr())[i] = static_cast<float*>(x.data_ptr())[i] + 1; |
| 38 | } |
| 39 | } |
| 40 | |
| 41 | TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_one_cpu, add_one_cpu) |