aparler/PyNorch
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

sub broadcasted cuda

Browse source
This commit is contained in:
lucasdelimanogueira 2024-05-17 13:15:33 -03:00
parent 64fb3f9edb
commit de72cd67ba
14 changed files with 128 additions and 30 deletions
Download patch file Download diff file Expand all files Collapse all files

8
norch/csrc/cpu.cpp
View file

@ -58,7 +58,7 @@ void sub_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) {
}
}
void sub_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape) {
void sub_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size) {
int max_ndim = tensor1->ndim > tensor2->ndim ? tensor1->ndim : tensor2->ndim;
// Calculate strides for broadcasting
@ -75,12 +75,12 @@ void sub_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_
int dim2 = i < tensor2->ndim ? tensor2->shape[tensor2->ndim - max_ndim + i] : 1;
strides1[i] = dim1 == broadcasted_shape[i] ? stride1 : 0;
strides2[i] = dim2 == broadcasted_shape[i] ? stride2 : 0;
stride1 *= broadcasted_shape[i];
stride2 *= broadcasted_shape[i];
stride1 *= (dim1 == broadcasted_shape[i]) ? dim1 : 1;
stride2 *= (dim2 == broadcasted_shape[i]) ? dim2 : 1;
}
// Perform element-wise addition with broadcasting
for (int i = 0; i < tensor1->size; i++) {
for (int i = 0; i < broadcasted_size; i++) {
int index1 = 0, index2 = 0;
int linear_index = i;
for (int j = max_ndim - 1; j >= 0; j--) {