jiebi/CodeConvo huggingface dataset:
(Addcdiv,       (),                 ((S, S), (S, S), torch.rand(S, S) + 1e-2)   ), (Addcdiv,       (0.6,),             ((S, S), (S, S), torch.rand(S, S) + 1e-2), 'scale'), (IndexAdd,      (0,),               ((S, S), index_variable(2, S), (2, S))      ), <del> (IndexCopy,     (0,),               ((S, S), index_variable(2, S), (2, S))      ), </del> <ins> # (IndexCopy,     (0,),               ((S, S), index_variable(2, S), (2, S))      ), </ins> (IndexFill,     (0, 2),             ((S, S), index_variable(2, S))              ), (IndexSelect,   (0,),               ((S, S), index_variable(2, S))              ), (Gather,        (0,),               ((M, S), gather_variable((S, S), 1, M))     ),
<del> import torch from torch.autograd.function import Function from torch._thnn import type2backend import torch.backends.cudnn as cudnn class BatchNorm(Function): def __init__(self, running_mean, running_var, training, momentum, eps): super(BatchNorm, self).__init__() self.running_mean = running_mean self.running_var = running_var self.training = training self.momentum = momentum self.eps = eps def forward(self, input, weight=None, bias=None): self.save_for_backward(input, weight, bias) # don't use cuDNN for half inputs because cuDNN requires the weight and # bias tensors to be floats, unlike THCUNN which requires half tensors. self.use_cudnn = (cudnn.is_acceptable(input) and cudnn.version() > 5110 and weight is not None and bias is not None and not isinstance(input, torch.cuda.HalfTensor)) # temporary buffers used in forward and backward num_features = input.size(1) _save_mean = input.new(num_features) _save_std = input.new(num_features) output = input.new(input.size()) if self.use_cudnn: torch._C._cudnn_batch_norm_forward( input, output, weight, bias, self.running_mean, self.running_var, _save_mean, _save_std, self.training, self.momentum, self.eps) else: backend = type2backend[type(input)] backend.BatchNormalization_updateOutput( backend.library_state, input, output, weight, bias, self.running_mean, self.running_var, _save_mean, _save_std, self.training, self.momentum, self.eps) if self.requires_grad: self._save_mean = _save_mean self._save_std = _save_std return output def backward(self, grad_output): input, weight, bias = self.saved_tensors grad_input, grad_weight, grad_bias = None, None, None if self.needs_input_grad[0] or self.use_cudnn: grad_input = input.new(input.size()) if (len(self.needs_input_grad) > 1 and self.needs_input_grad[1]) or self.use_cudnn: grad_weight = weight.new(weight.size()).zero_() if (len(self.needs_input_grad) > 1 and self.needs_input_grad[2]) or self.use_cudnn: grad_bias = bias.new(bias.size()).zero_() if self.use_cudnn and self.training: # cudnn does not support backward in evaluate mode torch._C._cudnn_batch_norm_backward( input, grad_output, grad_input, grad_weight, grad_bias, weight, self.running_mean, self.running_var, self._save_mean, self._save_std, self.training, self.eps) else: grad_output = grad_output.contiguous() backend = type2backend[type(input)] backend.BatchNormalization_backward( backend.library_state, input, grad_output, grad_input, grad_weight, grad_bias, weight, self.running_mean, self.running_var, self._save_mean, self._save_std, self.training, 1.0, self.eps) return grad_input, grad_weight, grad_bias </del>
// most fields THPFunction* self = (THPFunction*)obj; new (&self->cdata) torch::autograd::PyFunction(obj); <ins> new (&self->cdata_ptr) std::weak_ptr<torch::autograd::PyFunction>(); </ins> self->cdata.num_inputs = -1; self->cdata.is_stochastic = PyObject_IsInstance(obj, THPStochasticFunctionClass); return obj;
} }; <ins> // Similar to shared_from_this. There's a problem that the Python object // and its cdata depend on each other being alive, so we can't keep // shared_ptrs as members, but we'd like to be able to manage the lifetime of // the objects using shared_ptrs in the C++ graph. The only way to get a new // shared_ptr that references them is through THPFunction_asFunction. When // called for the first time it will allocate a new shared_ptr and save a // weak_ptr in cdata_ptr attr. Later, when we try to take another reference, // we'll try to lock cdata_ptr and return its value if successful. Otherwise it // means that all shared_ptrs returned previously have been freed, so we can // create a new one. This ensures that this object is managed by at most one // shared_ptr control block at any time - a guarantee we depend on in other places // (e.g. we use weak_ptrs in SavedVariable because we know it won't go out of scope). </ins> std::shared_ptr<PyFunction> THPFunction_asFunction(THPFunction* self) { if (!self) { return std::shared_ptr<PyFunction>(); } Py_INCREF((PyObject*)self); <del> return std::shared_ptr<PyFunction>(&self->cdata, Decref()); </del> <ins> auto ptr = self->cdata_ptr.lock(); if (ptr) return ptr; ptr = std::shared_ptr<PyFunction>(&self->cdata, Decref()); self->cdata_ptr = ptr; return ptr; </ins> }
// should have saved the grad accumulator. Even if the Variable no longer // alive, the accumulator should be kept alive by the references in the graph). if (requires_grad && !grad_fn && weak_grad_fn.expired() && grad_accumulator.expired()) <del> throw std::logic_error("No grad accumulator for a saved leaf!"); </del> <ins> throw std::logic_error("No grad accumulator for a saved leaf!"); </ins> new_var->grad_accumulator = grad_accumulator; return new_var;
LDFLAGS="-L$INSTALL_DIR/lib " LD_POSTFIX=".so.1" LD_POSTFIX_UNVERSIONED=".so" <del> if [[ $(uname) == 'Darwin' ]]; then </del> <ins> if [[ $(uname) == 'Darwin' ]]; then </ins> LDFLAGS="$LDFLAGS -Qunused-arguments -Wl,-rpath,@loader_path" LD_POSTFIX=".1.dylib" LD_POSTFIX_UNVERSIONED=".dylib"
-DCMAKE_CXX_FLAGS="$C_FLAGS $CPP_FLAGS" make install cp "lib/libnccl.so.1" "${INSTALL_DIR}/lib/libnccl.so.1" <del> ln -s "${INSTALL_DIR}/lib/libnccl.so.1" "${INSTALL_DIR}/lib/libnccl.so" </del> <ins> if [ ! -f "${INSTALL_DIR}/lib/libnccl.so" ]; then ln -s "${INSTALL_DIR}/lib/libnccl.so.1" "${INSTALL_DIR}/lib/libnccl.so" fi </ins> cd ../.. }
if hasattr(lib, 'cudnnGetErrorString'): lib.cudnnGetErrorString.restype = ctypes.c_char_p __cudnn_version = lib.cudnnGetVersion() <ins> compile_version = torch._C._cudnn_version() # Check that cuDNN major and minor versions match if (__cudnn_version // 100) != (compile_version // 100): raise RuntimeError( 'cuDNN version mismatch: PyTorch was compiled against {} ' 'but linked against {}'.format(compile_version, __cudnn_version)) </ins> else: lib = None return lib
torch.mm(a, b, out=self.output) else: if self.transA: <del> a = a.transpose(2, 3) </del> <ins> a = a.transpose(1, 2) </ins> if self.transB: <del> b = b.transpose(2, 3) </del> <ins> b = b.transpose(1, 2) </ins> self.output.resize_(a.size(0), a.size(1), b.size(2)) torch.bmm(a, b, out=self.output)
for i in range(10): Variable(torch.randn(10, 10), _grad_fn=CollectOnDelete()) <del> @unittest.skipIf(not torch.cuda.is_available() or torch.cuda.device_count() < 2, "CUDA not available or <2 GPUs detected") </del> <ins> @unittest.skipIf(torch.cuda.device_count() < 2, "no multi-GPU") </ins> def test_unused_output_gpu(self): from torch.nn.parallel._functions import Broadcast x = Variable(torch.randn(5, 5).float().cuda(), requires_grad=True)
y.sum().backward() self.assertEqual(x.grad.data, torch.ones(5, 5) * 2) <ins> @unittest.skipIf(torch.cuda.device_count() < 2, "no multi-GPU") def test_backward_device(self): # check that current device matches the variable's device device = [None] class Identity(torch.autograd.Function): @staticmethod def forward(ctx, x): return x.clone() @staticmethod def backward(ctx, grad_output): device[0] = torch.cuda.current_device() return grad_output.clone() v = Variable(torch.randn(1).cuda(1), requires_grad=True) Identity.apply(v).backward() self.assertEqual(device[0], 1) </ins> def test_detach(self): x = Variable(torch.randn(10, 10), requires_grad=True) y = x + 2
// This Engine's ReadyQueues and their corresponding threads are leaked here Engine::~Engine() = default; <del> auto Engine::thread_main(std::shared_ptr<ReadyQueue> queue) -> void { </del> <ins> auto Engine::thread_main(std::shared_ptr<ReadyQueue> queue, int device) -> void { </ins> THInferNumThreads(); <ins> AutoGPU guard(device); </ins> while (1) { FunctionTask task = queue->pop_back(); if (!task.base->has_error.load()) {
num_devices = 0; } #endif <del> ready_queues = std::vector<std::shared_ptr<ReadyQueue>>(num_devices + 1); for (auto& queue : ready_queues) { </del> <ins> int num_threads = num_devices + 1; ready_queues = std::vector<std::shared_ptr<ReadyQueue>>(num_threads); for (int i = 0; i < num_threads; ++i) { auto& queue = ready_queues[i]; </ins> queue.reset(new ReadyQueue()); <del> std::thread t(&Engine::thread_main, this, queue); </del> <ins> std::thread t(&Engine::thread_main, this, queue, i - 1); </ins> t.detach(); } }
void evaluate_function(FunctionTask& task); ReadyQueue& ready_queue(int device); void start_threads(); <del> virtual void thread_main(std::shared_ptr<ReadyQueue> queue); </del> <ins> virtual void thread_main(std::shared_ptr<ReadyQueue> queue, int device); </ins> virtual void thread_on_exception(FunctionTask& task, std::exception& e); std::once_flag start_threads_flag;
}; struct PythonEngine : public Engine { <del> virtual void thread_main(std::shared_ptr<ReadyQueue> queue) override { </del> <ins> virtual void thread_main(std::shared_ptr<ReadyQueue> queue, int device) override { </ins> // Create a PyThreadState, but release the GIL. This lets AutoGIL calls // inside thread_main acquire the GIL without having to create a new // PyThreadState each time. AutoGIL gil; AutoNoGIL no_gil; <del> Engine::thread_main(queue); </del> <ins> Engine::thread_main(queue, device); </ins> } virtual void thread_on_exception(FunctionTask& task, std::exception& e) override {
#define IS_CUDA false #define CUDA_FLOAT false #else #define IS_CUDA true <del> #define CUDA_BYTE defined(THC_REAL_IS_BYTE) #define CUDA_CHAR defined(THC_REAL_IS_CHAR) #define CUDA_SHORT defined(THC_REAL_IS_SHORT) #define CUDA_INT defined(THC_REAL_IS_INT) #define CUDA_LONG defined(THC_REAL_IS_LONG) #define CUDA_FLOAT defined(THC_REAL_IS_FLOAT) #define CUDA_DOUBLE defined(THC_REAL_IS_DOUBLE) #define CUDA_HALF defined(THC_REAL_IS_HALF) </del> <ins> #if defined(THC_REAL_IS_BYTE) #define CUDA_BYTE 1 #else #define CUDA_BYTE 0 </ins> #endif <ins> #if defined(THC_REAL_IS_CHAR) #define CUDA_CHAR 1 #else #define CUDA_CHAR 0 #endif #if defined(THC_REAL_IS_SHORT) #define CUDA_SHORT 1 #else #define CUDA_SHORT 0 #endif #if defined(THC_REAL_IS_INT) #define CUDA_INT 1 #else #define CUDA_INT 0 #endif #if defined(THC_REAL_IS_LONG) #define CUDA_LONG 1 #else #define CUDA_LONG 0 #endif #if defined(THC_REAL_IS_FLOAT) #define CUDA_FLOAT 1 #else #define CUDA_FLOAT 0 #endif #if defined(THC_REAL_IS_DOUBLE) #define CUDA_DOUBLE 1 #else #define CUDA_DOUBLE 0 #endif #if defined(THC_REAL_IS_HALF) #define CUDA_HALF 1 #else #define CUDA_HALF 0 #endif #endif // ifndef THC_GENERIC_FILE </ins> #if IS_CUDA #define THIndexTensor THCudaLongTensor #define THIndexTensor_(NAME) TH_CONCAT_2(THCudaLongTensor_,NAME)
Best practices -------------- <ins> Device-agnostic code ^^^^^^^^^^^^^^^^^^^^ Due to the structure of PyTorch, you may need to explicitly write device-agnostic (CPU or GPU) code; an example may be creating a new tensor as the initial hidden state of a recurrent neural network. The first step is to determine whether the GPU should be used or not. A common pattern is to use Python's `argparse` module to read in user arguments, and have a flag that can be used to disable CUDA, in combination with `torch.cuda.is_available()`. In the following, `args.cuda` results in a flag that can be used to cast tensors and modules to CUDA if desired:: import argparse import torch parser = argparse.ArgumentParser(description='PyTorch Example') parser.add_argument('--disable-cuda', action='store_true', help='Disable CUDA') args = parser.parse_args() args.cuda = not args.disable_cuda and torch.cuda.is_available() If modules or tensors need to be sent to the GPU, `args.cuda` can be used as follows:: x = torch.Tensor(8, 42) net = Network() if args.cuda: x = x.cuda() net.cuda() When creating tensors, an alternative to the if statement is to have a default datatype defined, and cast all tensors using that. An example when using a dataloader would be as follows:: dtype = torch.cuda.FloatTensor for i, x in enumerate(train_loader): x = Variable(x.type(dtype)) When working with multiple GPUs on a system, you can use the `CUDA_VISIBLE_DEVICES` environment flag to manage which GPUs are available to PyTorch. To manually control which GPU a tensor is created on, the best practice is to use the `torch.cuda.device()` context manager:: print("Outside device is 0")  # On device 0 (default in most scenarios) with torch.cuda.device(1): print("Inside device is 1")  # On device 1 print("Outside device is still 0")  # On device 0 If you have a tensor and would like to create a new tensor of the same type on the same device, then you can use the `.new()` function, which acts the same as a normal tensor constructor. Whilst the previously mentioned methods depend on the current GPU context, `new()` preserves the device of the original tensor. This is the recommended practice when creating modules in which new tensors/variables need to be created internally during the forward pass:: x_cpu = torch.FloatTensor(1) x_gpu = torch.cuda.FloatTensor(1) x_cpu_long = torch.LongTensor(1) y_cpu = x_cpu.new(8, 10, 10).fill_(0.3) y_gpu = x_gpu.new(x_gpu.size()).fill_(-5) y_cpu_long = x_cpu_long.new([[1, 2, 3]]) If you want to create a tensor of the same type and size of another tensor, and fill it with either ones or zeros, `torch.ones_like()` or `torch.zeros_like()` are provided as more convenient functions (which also preserve device):: x_cpu = torch.FloatTensor(1) x_gpu = torch.cuda.FloatTensor(1) y_cpu = torch.ones_like(x_cpu) y_gpu = torch.zeros_like(x_gpu) </ins> Use pinned memory buffers ^^^^^^^^^^^^^^^^^^^^^^^^^
res2 += i * j self.assertEqual(res1, res2) <ins> # Test 0-strided for tname, _prec in types.items(): v1 = torch.randn(1).type(tname).expand(100) v2 = torch.randn(100).type(tname) res1 = torch.dot(v1, v2) res2 = 0 for i, j in zip(v1, v2): res2 += i * j self.assertEqual(res1, res2) def test_ger(self): types = { 'torch.DoubleTensor': 1e-8, 'torch.FloatTensor': 1e-4, } for tname, _prec in types.items(): v1 = torch.randn(100).type(tname) v2 = torch.randn(100).type(tname) res1 = torch.ger(v1, v2) res2 = torch.zeros(100, 100).type(tname) for i in range(100): for j in range(100): res2[i, j] = v1[i] * v2[j] self.assertEqual(res1, res2) # Test 0-strided for tname, _prec in types.items(): v1 = torch.randn(1).type(tname).expand(100) v2 = torch.randn(100).type(tname) res1 = torch.ger(v1, v2) res2 = torch.zeros(100, 100).type(tname) for i in range(100): for j in range(100): res2[i, j] = v1[i] * v2[j] self.assertEqual(res1, res2) def test_addmv(self): types = { 'torch.DoubleTensor': 1e-8, 'torch.FloatTensor': 1e-4, } for tname, _prec in types.items(): t = torch.randn(10).type(tname) m = torch.randn(10, 100).type(tname) v = torch.randn(100).type(tname) res1 = torch.addmv(t, m, v) res2 = torch.zeros(10).type(tname) res2 += t for i in range(10): for j in range(100): res2[i] += m[i, j] * v[j] self.assertEqual(res1, res2) # Test 0-strided for tname, _prec in types.items(): t = torch.randn(1).type(tname).expand(10) m = torch.randn(10, 1).type(tname).expand(10, 100) v = torch.randn(100).type(tname) res1 = torch.addmv(t, m, v) res2 = torch.zeros(10).type(tname) res2 += t for i in range(10): for j in range(100): res2[i] += m[i, j] * v[j] self.assertEqual(res1, res2) def test_addmm(self): types = { 'torch.DoubleTensor': 1e-8, 'torch.FloatTensor': 1e-4, } for tname, _prec in types.items(): M = torch.randn(10, 25).type(tname) m1 = torch.randn(10, 50).type(tname) m2 = torch.randn(50, 25).type(tname) res1 = torch.addmm(M, m1, m2) res2 = torch.zeros(10, 25).type(tname) res2 += M for i in range(10): for j in range(25): for k in range(50): res2[i, j] += m1[i, k] * m2[k, j] self.assertEqual(res1, res2) # Test 0-strided for tname, _prec in types.items(): M = torch.randn(10, 1).type(tname).expand(10, 25) m1 = torch.randn(10, 1).type(tname).expand(10, 50) m2 = torch.randn(50, 25).type(tname) res1 = torch.addmm(M, m1, m2) res2 = torch.zeros(10, 25).type(tname) res2 += M for i in range(10): for j in range(25): for k in range(50): res2[i, j] += m1[i, k] * m2[k, j] self.assertEqual(res1, res2) </ins> def _testMath(self, torchfn, mathfn): size = (10, 5) # contiguous
#define THC_GENERIC_FILE "generic/THCTensorIndex.cu" #else <ins> // Check tensor dimensions for index operations, and return the slice size. // src can be nullptr in case of indexFill: in that case it is ignored. static ptrdiff_t THCTensor_(getSliceSize)(THCState *state, THCTensor *dst, int dim, THCudaLongTensor *index, THCTensor *src) { int dstDims = THCTensor_(nDimension)(state, dst); int srcDims = (src == nullptr) ? dstDims : THCTensor_(nDimension)(state, src); THArgCheck(THCudaLongTensor_nDimension(state, index) == 1, 4, "expecting vector of indices"); THArgCheck(dim >= 0 && dim < dstDims, 2, "Indexing dim is out of bounds"); ptrdiff_t dstSliceSize = 1; for (int d = 0; d < dstDims; d++) { if (d != dim) { dstSliceSize *= dst->size[d]; } } if (src == nullptr) return dstSliceSize; THArgCheck(dim < srcDims, 3, "Indexing dim is out of bounds"); THArgCheck(THCudaLongTensor_nElement(state, index) == src->size[dim], 4, "length of src.size[dim] is not equal to length of indices"); ptrdiff_t srcSliceSize = 1; bool mismatch = false; if (dstDims != srcDims) mismatch = true; for (int d = 0; d < srcDims; d++) { if (d != dim) { srcSliceSize *= src->size[d]; if (!mismatch && dst->size[d] != src->size[d]) mismatch = true; } } THArgCheck(dstSliceSize == srcSliceSize, 2, "Source/destination tensor have different slice sizes (%ld vs %ld)", dstSliceSize, srcSliceSize); if (mismatch) { static bool warningShown = false; if (!warningShown) { warningShown = true; fprintf(stderr, "Warning: source/destination slices have same size but different " "shape for an index operation.  This behavior is deprecated.n"); } } return dstSliceSize; } </ins> void THCTensor_(indexCopy_long)(THCState *state, THCTensor *dst, int dim, THLongTensor *indices, THCTensor *src) { THCAssertSameGPU(THCTensor_(checkGPU)(state, 2, dst, src));
dims = THCudaLongTensor_nDimension(state, indices); THArgCheck(dims <= MAX_CUTORCH_DIMS, 4, CUTORCH_DIM_WARNING); <del> ptrdiff_t numIndices = THCudaLongTensor_nElement(state, indices); int srcDims = THCTensor_(nDimension)(state, src); cudaStream_t stream = THCState_getCurrentStream(state); THArgCheck(THCudaLongTensor_nDimension(state, indices) == 1, 3, "expecting vector of indices"); THArgCheck(dim < srcDims, 4, "Indexing dim is out of bounds"); THArgCheck(srcDims > 0, 2, "Source tensor is empty"); THArgCheck(numIndices == src->size[dim], 4, "length of src.size[dim] is not equal to length of indices"); int indContig = THCudaLongTensor_isContiguous(state, indices); </del> // The `src` is partitioned into two parts: // -the size of each slice we are indexing, which is the // total size of the tensor ignoring dimension `dim`; // -the number of indices we are choosing, which is the total size // of the tensor `indices`. <ins> ptrdiff_t sliceSize = THCTensor_(getSliceSize)(state, dst, dim, indices, src); </ins> ptrdiff_t srcTotalSize = THCTensor_(nElement)(state, src); int64_t dstCopyDimSize = THCTensor_(size)(state, dst, dim); <del> ptrdiff_t sliceSize = srcTotalSize / numIndices; </del> <ins> ptrdiff_t numIndices = THCudaLongTensor_nElement(state, indices); cudaStream_t stream = THCState_getCurrentStream(state); int indContig = THCudaLongTensor_isContiguous(state, indices); </ins> int mpc = THCState_getCurrentDeviceProperties(state)->multiProcessorCount;
dims = THCudaLongTensor_nDimension(state, indices); THArgCheck(dims <= MAX_CUTORCH_DIMS, 4, CUTORCH_DIM_WARNING); <del> ptrdiff_t numIndices = THCudaLongTensor_nElement(state, indices); int srcDims = THCTensor_(nDimension)(state, src); cudaStream_t stream = THCState_getCurrentStream(state); THArgCheck(THCudaLongTensor_nDimension(state, indices) == 1, 3, "expecting vector of indices"); THArgCheck(dim < srcDims, 4, "Indexing dim is out of bounds"); THArgCheck(srcDims > 0, 2, "Source tensor is empty"); THArgCheck(numIndices == src->size[dim], 4, "length of src.size[dim] is not equal to length of indices"); int indContig = THCudaLongTensor_isContiguous(state, indices); </del> // The `src` is partitioned into two parts: // -the size of each slice we are indexing, which is the // total size of the tensor ignoring dimension `dim`; // -the number of indices we are choosing, which is the total size // of the tensor `indices`. <ins> ptrdiff_t sliceSize = THCTensor_(getSliceSize)(state, dst, dim, indices, src); </ins> ptrdiff_t srcTotalSize = THCTensor_(nElement)(state, src); int64_t dstAddDimSize = THCTensor_(size)(state, dst, dim); <del> ptrdiff_t sliceSize = srcTotalSize / numIndices; </del> <ins> ptrdiff_t numIndices = THCudaLongTensor_nElement(state, indices); cudaStream_t stream = THCState_getCurrentStream(state); int indContig = THCudaLongTensor_isContiguous(state, indices); </ins> int mpc = THCState_getCurrentDeviceProperties(state)->multiProcessorCount;
dims = THCudaLongTensor_nDimension(state, indices); THArgCheck(dims <= MAX_CUTORCH_DIMS, 4, CUTORCH_DIM_WARNING); <del> ptrdiff_t numIndices = THCudaLongTensor_nElement(state, indices); int srcDims = THCTensor_(nDimension)(state, dst); cudaStream_t stream = THCState_getCurrentStream(state); THArgCheck(THCudaLongTensor_nDimension(state, indices) == 1, 3, "expecting vector of indices"); THArgCheck(dim < srcDims, 4, "Indexing dim is out of bounds"); THArgCheck(srcDims > 0, 2, "Source tensor is empty"); int indContig = THCudaLongTensor_isContiguous(state, indices); </del> // The `src` is partitioned into two parts: // -the size of each slice we are indexing, which is the // total size of the tensor ignoring dimension `dim`; // -the number of indices we are choosing, which is the total size // of the tensor `indices`. <ins> ptrdiff_t sliceSize = THCTensor_(getSliceSize)(state, dst, dim, indices, nullptr); </ins> ptrdiff_t dstTotalSize = THCTensor_(nElement)(state, dst); int64_t dstFillDimSize = THCTensor_(size)(state, dst, dim); <del> ptrdiff_t sliceSize = dstTotalSize / dstFillDimSize; </del> <ins> ptrdiff_t numIndices = THCudaLongTensor_nElement(state, indices); cudaStream_t stream = THCState_getCurrentStream(state); int indContig = THCudaLongTensor_isContiguous(state, indices); </ins> int mpc = THCState_getCurrentDeviceProperties(state)->multiProcessorCount;
THLongStorage_free(topKSize); #define RUN_K(INDEX_T, DIM, DIR)                                         <del> gatherTopK<real, INDEX_T, DIM, DIR>                                          </del> <ins> gatherTopK<real, INDEX_T, DIM, DIR>                                    </ins> <<<grid, block, 0, THCState_getCurrentStream(state)>>>(              inputInfo,                                                         sliceSize,
} #define RUN_T(INDEX_T)                                                   <del> TensorInfo<real, INDEX_T> inputInfo =                                 getTensorInfo<THCTensor, INDEX_T>(state, input);                  TensorInfo<real, INDEX_T> topKInfo =                                  getTensorInfo<THCTensor, INDEX_T>(state, topK);                   </del> <ins> TensorInfo<real, INDEX_T> inputInfo =                                  getTensorInfo<THCTensor, INDEX_T>(state, input);                     TensorInfo<real, INDEX_T> topKInfo =                                   getTensorInfo<THCTensor, INDEX_T>(state, topK);                      </ins> TensorInfo<int64_t, INDEX_T> indicesInfo =                             getTensorInfo<THCudaLongTensor, INDEX_T>(state, indices);