/*
* Based on the implementation of the THTensor_(indexSelect) in torch7
*/
static void THCudaTensor_indexSelect(THCudaTensor *tensor, THCudaTensor *src, int dim, THLongTensor *index)
{
long i, numel;
THLongStorage *newSize;
THCudaTensor *tSlice, *sSlice;
long *index_data;
THArgCheck(index->nDimension == 1, 3, "Index is supposed to be a vector");
THArgCheck(dim < src->nDimension,4,"Indexing dim is out of bounds");
THArgCheck(src->nDimension > 0,2,"Source tensor is empty");
numel = THLongTensor_nElement(index);
newSize = THLongStorage_newWithSize(src->nDimension);
THLongStorage_rawCopy(newSize,src->size);
newSize->data[dim] = numel;
THCudaTensor_resize(tensor,newSize,NULL);
THLongStorage_free(newSize);
index = THLongTensor_newContiguous(index);
index_data = THLongTensor_data(index);
for (i=0; i<numel; i++)
{
if (src->nDimension > 1)
{
tSlice = THCudaTensor_new();
sSlice = THCudaTensor_new();
THCudaTensor_select(tSlice, tensor, dim, i);
THCudaTensor_select(sSlice, src, dim, index_data[i]-1);
THCudaTensor_copy(tSlice, sSlice);
THCudaTensor_free(tSlice);
THCudaTensor_free(sSlice);
}
else
{ // It's faster to copy a float from an address in the device to another address in the device than
// retrieving it to the host memory and recopy it to the device memory
THCudaCheck(cudaMemcpy(tensor->storage->data + tensor->storageOffset + i,\
src->storage->data + src->storageOffset + index_data[i]-1, sizeof(float), cudaMemcpyDeviceToDevice));
}
}
THLongTensor_free(index);
}
THLongStorage *THCudaTensor_newStrideOf(THCState *state, THCudaTensor *self)
{
THLongStorage *stride = THLongStorage_newWithSize(self->nDimension);
THLongStorage_rawCopy(stride, self->stride);
return stride;
}
THLongStorage *THZTensor_(newStrideOf)(THZTensor *self)
{
THLongStorage *stride = THLongStorage_newWithSize(self->nDimension);
THLongStorage_rawCopy(stride, self->stride);
return stride;
}
THLongStorage *THZTensor_(newSizeOf)(THZTensor *self)
{
THLongStorage *size = THLongStorage_newWithSize(self->nDimension);
THLongStorage_rawCopy(size, self->size);
return size;
}
THLongStorage *THTensor_(newStrideOf)(THTensor *self)
{
THLongStorage *stride = THLongStorage_newWithSize(self->_dim());
THLongStorage_rawCopy(stride, self->stride);
return stride;
}
THLongStorage *THTensor_(newSizeOf)(THTensor *self)
{
THLongStorage *size = THLongStorage_newWithSize(self->_dim());
THLongStorage_rawCopy(size, self->size);
return size;
}
THLongStorage *THCSTensor_(newSizeOf)(THCState *state, THCSTensor *self)
{
THLongStorage *size = THLongStorage_newWithSize(self->nDimensionI + self->nDimensionV);
THLongStorage_rawCopy(size, self->size);
return size;
}
THLongStorage *THClTensor_newSizeOf(THClState *state, THClTensor *self)
{
THLongStorage *size = THLongStorage_newWithSize(self->nDimension);
THLongStorage_rawCopy(size, self->size);
return size;
}