1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
|
#ifndef THC_GENERIC_FILE
#define THC_GENERIC_FILE "generic/THCTensorMasked.cu"
#else
THC_API void
THCTensor_(maskedFill)(THCState* state,
THCTensor *tensor, THCudaByteTensor *mask, real value)
{
THAssert(THCTensor_(checkGPU)(state, 2, tensor, mask));
THArgCheck(THCTensor_(nElement)(state, tensor) ==
THCudaByteTensor_nElement(state, mask),
2, "sizes do not match");
if (!THC_pointwiseApply2(state, tensor, mask,
TensorMaskedFillOp<real, unsigned char>(value))) {
THArgCheck(false, 2, CUTORCH_DIM_WARNING);
}
THCudaCheck(cudaGetLastError());
}
THC_API void
THCTensor_(maskedFillByte)(THCState* state,
THCTensor *tensor, THByteTensor *mask, real value)
{
THAssert(THCTensor_(checkGPU)(state, 1, tensor));
THLongStorage* maskSizes = THByteTensor_newSizeOf(mask);
THCudaByteTensor* maskCuda = THCudaByteTensor_newWithSize(state, maskSizes, NULL);
THLongStorage_free(maskSizes);
THCudaByteTensor_copyByte(state, maskCuda, mask);
THCTensor_(maskedFill)(state, tensor, maskCuda, value);
THCudaByteTensor_free(state, maskCuda);
}
THC_API void
THCTensor_(maskedCopy)(THCState* state,
THCTensor *tensor, THCudaByteTensor *mask, THCTensor *src)
{
THAssert(THCTensor_(checkGPU)(state, 3, tensor, src, mask));
long maskSize = THCudaByteTensor_nElement(state, mask);
long tensorSize = THCTensor_(nElement)(state, tensor);
long srcSize = THCTensor_(nElement)(state, src);
// `mask` and `tensor` must have the same number of elements
THArgCheck(maskSize == tensorSize, 2,
"mask and tensor must have the same number of elements");
// Determine our output size
long totalElements = THCudaByteTensor_sumall(state, mask);
// The number of `1` elements present in the mask must be <= the
// number of elements available in `src`
if (totalElements > srcSize) {
THArgCheck(false, 2, "source nElements must be == mask `1` elements");
}
// FIXME: there appears to be a bug in Thrust (CUDA 7.0) for mixed
// iterator prefix sums? Convert `mask` to the same datatype as what
// we're accumulating the prefix sum in (long) to get around it
THCudaLongTensor* maskLong = THCudaLongTensor_new(state);
THLongStorage* maskSizes = THCudaByteTensor_newSizeOf(state, mask);
THCudaLongTensor_resize(state, maskLong, maskSizes, NULL);
THCudaLongTensor_copyCudaByte(state, maskLong, mask);
// Use a prefix sum to determine the output locations of the masked elements
THCudaLongTensor* maskPrefixSum = THCudaLongTensor_new(state);
THCudaLongTensor_resize(state, maskPrefixSum, maskSizes, NULL);
THLongStorage_free(maskSizes);
thrust::device_ptr<long>
maskData(THCudaLongTensor_data(state, maskLong));
thrust::device_ptr<long>
maskPrefixSumData(THCudaLongTensor_data(state, maskPrefixSum));
thrust::exclusive_scan(
#if CUDA_VERSION >= 7000
thrust::cuda::par.on(THCState_getCurrentStream(state)),
#endif
maskData,
maskData + THCudaLongTensor_nElement(state, maskLong),
maskPrefixSumData);
// We are getting elements from `src` based on an offset from
// `maskPrefixSum`, so that should be made contiguous too
THCTensor* contigSrc = THCTensor_(newContiguous)(state, src);
// update `tensor` where `mask` == 1 but pull from `src` at
// maskPrefixSum
bool status = THC_pointwiseApply3(
state, tensor, mask, maskPrefixSum,
TensorMaskedCopyOp<real, unsigned char, long>(
THCTensor_(data)(state, contigSrc)));
THCTensor_(free)(state, contigSrc);
THCudaLongTensor_free(state, maskLong);
THCudaLongTensor_free(state, maskPrefixSum);
THArgCheck(status, 2, CUTORCH_DIM_WARNING);
THCudaCheck(cudaGetLastError());
}
THC_API void
THCTensor_(maskedCopyByte)(THCState* state,
THCTensor *tensor, THByteTensor *mask, THCTensor *src) {
THAssert(THCTensor_(checkGPU)(state, 2, tensor, src));
THLongStorage* maskSizes = THByteTensor_newSizeOf(mask);
THCudaByteTensor* maskCuda = THCudaByteTensor_newWithSize(state, maskSizes, NULL);
THLongStorage_free(maskSizes);
THCudaByteTensor_copyByte(state, maskCuda, mask);
THCTensor_(maskedCopy)(state, tensor, maskCuda, src);
THCudaByteTensor_free(state, maskCuda);
}
THC_API void
THCTensor_(maskedSelect)(THCState* state,
THCTensor* tensor, THCTensor* src, THCudaByteTensor* mask) {
THAssert(THCTensor_(checkGPU)(state, 3, tensor, src, mask));
THArgCheck(THCudaByteTensor_nElement(state, mask) ==
THCTensor_(nElement)(state, src),
2, "sizes do not match");
// Determine our output size
long totalElements = THCudaByteTensor_sumall(state, mask);
THCTensor* tensorContig = THCTensor_(newContiguous)(state, tensor);
THCTensor_(resize1d)(state, tensorContig, totalElements);
if (tensor != tensorContig) {
THCTensor_(resize1d)(state, tensor, totalElements);
}
// FIXME: there appears to be a bug in Thrust (CUDA 7.0) for mixed
// iterator prefix sums? Convert `mask` to the same datatype as what
// we're accumulating the prefix sum in (long) to get around it
THCudaLongTensor* maskLong = THCudaLongTensor_new(state);
THLongStorage* maskSizes = THCudaByteTensor_newSizeOf(state, mask);
THCudaLongTensor_resize(state, maskLong, maskSizes, NULL);
THCudaLongTensor_copyCudaByte(state, maskLong, mask);
// Use a prefix sum to determine the output locations of the masked elements
THCudaLongTensor* maskPrefixSum = THCudaLongTensor_new(state);
THCudaLongTensor_resize(state, maskPrefixSum, maskSizes, NULL);
THLongStorage_free(maskSizes);
thrust::device_ptr<long>
maskData(THCudaLongTensor_data(state, maskLong));
thrust::device_ptr<long>
maskPrefixSumData(THCudaLongTensor_data(state, maskPrefixSum));
thrust::exclusive_scan(
#if CUDA_VERSION >= 7000
thrust::cuda::par.on(THCState_getCurrentStream(state)),
#endif
maskData,
maskData + THCudaLongTensor_nElement(state, maskLong),
maskPrefixSumData);
// Then copy over the masked elements at their desired output index
bool status = THC_pointwiseApply3(
state, mask, maskPrefixSum,
src, TensorMaskedSelectOp<real, unsigned char, long>(
THCTensor_(data)(state, tensor)));
THCudaLongTensor_free(state, maskLong);
THCudaLongTensor_free(state, maskPrefixSum);
if (tensor != tensorContig) {
THCTensor_(freeCopyTo)(state, tensorContig, tensor);
} else {
THCTensor_(free)(state, tensorContig);
}
THArgCheck(status, 2, CUTORCH_DIM_WARNING);
THCudaCheck(cudaGetLastError());
}
// FIXME: remove now that we have THCudaByteTensor?
THC_API void
THCTensor_(maskedSelectByte)(THCState* state,
THCTensor *tensor, THCTensor *src, THByteTensor *mask)
{
THAssert(THCTensor_(checkGPU)(state, 2, tensor, src));
THLongStorage* maskSizes = THByteTensor_newSizeOf(mask);
THCudaByteTensor* maskCuda = THCudaByteTensor_newWithSize(state, maskSizes, NULL);
THLongStorage_free(maskSizes);
THCudaByteTensor_copyByte(state, maskCuda, mask);
THCTensor_(maskedSelect)(state, tensor, src, maskCuda);
THCudaByteTensor_free(state, maskCuda);
}
#endif
|
