Skip to content

Commit f5dec55

Browse files
PeiyuLauPeiyuLau
authored
[MLU] support gather_nd and transpose int64 dtype (#1343)
1 parent 3931a65 commit f5dec55

File tree

5 files changed

+102
-22
lines changed

5 files changed

+102
-22
lines changed

‎backends/mlu/kernels/gather_nd_kernel.cc

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,7 @@ PD_REGISTER_PLUGIN_KERNEL(gather_nd,
141141
mlu,
142142
ALL_LAYOUT,
143143
custom_kernel::GatherNdKernel,
144+
int64_t,
144145
float,
145146
phi::dtype::float16) {}
146147

‎backends/mlu/kernels/set_value_kernel.cc

Lines changed: 68 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -147,10 +147,8 @@ void SetValueKernel(const Context& dev_ctx,
147147
index_out.Resize(index_dims);
148148
dev_ctx.template Alloc<int64_t>(&index_out);
149149
cnnlScatterRefMode_t mode = CNNL_SCATTERREF_UPDATE;
150-
MLUCnnlTensorDesc x_desc(in_temp);
151150
MLUCnnlTensorDesc indices_desc(index_temp);
152151
MLUCnnlTensorDesc indices_out_desc(index_out);
153-
MLUCnnlTensorDesc updates_desc(val_temp);
154152
MLUCnnlTensorDesc out_desc(*out);
155153
MLUCnnl::StridedSlice(dev_ctx,
156154
starts_indices,
@@ -165,22 +163,77 @@ void SetValueKernel(const Context& dev_ctx,
165163
phi::product(slice_dims_for_assign),
166164
phi::errors::InvalidArgument(
167165
"OP(set_value) error index indices and value update not match "));
166+
167+
// when input x and input value's dtype is int64,
168+
// cast datadtype to int32 for cnnlScatterRef usage
169+
Tensor Ref, Update;
170+
Ref.Resize(in_temp.dims());
171+
Update.Resize(val_temp.dims());
172+
173+
if (in_temp.dtype() != DataType::INT64 &&
174+
val_temp.dtype() != DataType::INT64 &&
175+
in_temp.dtype() != DataType::BOOL && val_temp.dtype() != DataType::BOOL) {
176+
Ref = in_temp;
177+
Update = val_temp;
178+
} else {
179+
dev_ctx.template Alloc<int32_t>(&Ref);
180+
dev_ctx.template Alloc<int32_t>(&Update);
181+
MLUCnnlTensorDesc in_temp_desc(in_temp);
182+
MLUCnnlTensorDesc Ref_desc(Ref);
183+
MLUCnnlTensorDesc val_temp_desc(val_temp);
184+
MLUCnnlTensorDesc Update_desc(Update);
185+
cnnlCastDataType_t cast_type = GetCastDataType(x.dtype(), DataType::INT32);
186+
MLUCnnl::Cast(dev_ctx,
187+
cast_type,
188+
in_temp_desc.get(),
189+
GetBasePtr(&in_temp),
190+
Ref_desc.get(),
191+
GetBasePtr(&Ref));
192+
MLUCnnl::Cast(dev_ctx,
193+
cast_type,
194+
val_temp_desc.get(),
195+
GetBasePtr(&val_temp),
196+
Update_desc.get(),
197+
GetBasePtr(&Update));
198+
}
199+
200+
MLUCnnlTensorDesc Ref_desc(Ref);
201+
MLUCnnlTensorDesc Update_desc(Update);
168202
Tensor index_final;
169203
index_final = index_out;
170204
int64_t indices_numel = phi::product(index_dims);
171205
auto new_index_dims = phi::make_ddim({indices_numel});
172206
index_final.Resize(new_index_dims);
173207
MLUCnnlTensorDesc indices_final_desc(index_final);
174208
MLUCnnl::ScatterRefFunctor(dev_ctx,
175-
x_desc.get(),
176-
GetBasePtr(&in_temp),
177-
updates_desc.get(),
178-
GetBasePtr(&val_temp),
209+
Ref_desc.get(),
210+
GetBasePtr(&Ref),
211+
Update_desc.get(),
212+
GetBasePtr(&Update),
179213
indices_final_desc.get(),
180214
GetBasePtr(&index_final),
181215
mode);
182-
in_temp.Resize(in_dims);
183-
TensorCopy(dev_ctx, in_temp, false, out);
216+
Ref.Resize(in_dims);
217+
218+
// When input x and input value's dtype is int64,
219+
// cast ScatterRef output datatype from int32 to int64
220+
Tensor in_temp_out;
221+
in_temp_out.Resize(Ref.dims());
222+
if (x.dtype() != DataType::INT64 && x.dtype() != DataType::BOOL) {
223+
in_temp_out = Ref;
224+
} else {
225+
dev_ctx.template Alloc<T>(&in_temp_out);
226+
MLUCnnlTensorDesc in_temp_desc(Ref);
227+
MLUCnnlTensorDesc in_temp_out_desc(in_temp_out);
228+
cnnlCastDataType_t cast_type = GetCastDataType(Ref.dtype(), x.dtype());
229+
MLUCnnl::Cast(dev_ctx,
230+
cast_type,
231+
in_temp_desc.get(),
232+
GetBasePtr(&Ref),
233+
in_temp_out_desc.get(),
234+
GetBasePtr(&in_temp_out));
235+
}
236+
TensorCopy(dev_ctx, in_temp_out, false, out);
184237
}
185238

186239
template <typename T, typename Context>
@@ -398,8 +451,13 @@ void SetTensorValueKernel(const Context& dev_ctx,
398451

399452
} // namespace custom_kernel
400453

401-
PD_REGISTER_PLUGIN_KERNEL(
402-
set_value, mlu, ALL_LAYOUT, custom_kernel::SetValueKernel, float, int) {}
454+
PD_REGISTER_PLUGIN_KERNEL(set_value,
455+
mlu,
456+
ALL_LAYOUT,
457+
custom_kernel::SetValueKernel,
458+
float,
459+
int,
460+
int64_t) {}
403461

404462
PD_REGISTER_PLUGIN_KERNEL(set_value_with_tensor,
405463
mlu,

‎backends/mlu/kernels/transpose_kernel.cc

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ PD_REGISTER_PLUGIN_KERNEL(transpose,
5050
ALL_LAYOUT,
5151
custom_kernel::TransposeKernel,
5252
int,
53+
int64_t,
5354
uint8_t,
5455
int8_t,
5556
int16_t,
@@ -62,6 +63,7 @@ PD_REGISTER_PLUGIN_KERNEL(transpose_grad,
6263
ALL_LAYOUT,
6364
custom_kernel::TransposeGradKernel,
6465
int,
66+
int64_t,
6567
uint8_t,
6668
int8_t,
6769
int16_t,

‎backends/mlu/tests/unittests/test_gather_nd_op_mlu.py

Lines changed: 10 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ def test_check_output(self):
5757
self.check_output_with_place(self.place)
5858

5959
def test_check_grad(self):
60-
if typename == "float16":
60+
if typename == "float16" or typename == "int64":
6161
self.__class__.no_need_check_grad = True
6262
else:
6363
self.check_grad_with_place(self.place, ["X"], "Out")
@@ -84,7 +84,7 @@ def test_check_output(self):
8484
self.check_output_with_place(self.place)
8585

8686
def test_check_grad(self):
87-
if typename == "float16":
87+
if typename == "float16" or typename == "int64":
8888
self.__class__.no_need_check_grad = True
8989
else:
9090
self.check_grad_with_place(self.place, ["X"], "Out")
@@ -116,7 +116,7 @@ def test_check_output(self):
116116
self.check_output_with_place(self.place)
117117

118118
def test_check_grad(self):
119-
if typename == "float16":
119+
if typename == "float16" or typename == "int64":
120120
self.__class__.no_need_check_grad = True
121121
else:
122122
self.check_grad_with_place(
@@ -149,7 +149,7 @@ def test_check_output(self):
149149
self.check_output_with_place(self.place)
150150

151151
def test_check_grad(self):
152-
if typename == "float16":
152+
if typename == "float16" or typename == "int64":
153153
self.__class__.no_need_check_grad = True
154154
else:
155155
self.check_grad_with_place(self.place, ["X"], "Out")
@@ -180,7 +180,7 @@ def test_check_output(self):
180180
self.check_output_with_place(self.place)
181181

182182
def test_check_grad(self):
183-
if typename == "float16":
183+
if typename == "float16" or typename == "int64":
184184
self.__class__.no_need_check_grad = True
185185
else:
186186
self.check_grad_with_place(self.place, ["X"], "Out")
@@ -212,7 +212,7 @@ def test_check_output(self):
212212
self.check_output_with_place(self.place)
213213

214214
def test_check_grad(self):
215-
if typename == "float16":
215+
if typename == "float16" or typename == "int64":
216216
self.__class__.no_need_check_grad = True
217217
else:
218218
self.check_grad_with_place(self.place, ["X"], "Out")
@@ -245,7 +245,7 @@ def test_check_output(self):
245245
self.check_output_with_place(self.place)
246246

247247
def test_check_grad(self):
248-
if typename == "float16":
248+
if typename == "float16" or typename == "int64":
249249
self.__class__.no_need_check_grad = True
250250
else:
251251
self.check_grad_with_place(self.place, ["X"], "Out")
@@ -255,13 +255,12 @@ def test_check_grad(self):
255255
globals()[cls_name] = TestGatherNdOpWithHighRankDiff
256256

257257

258-
# Test Python API
259258
class TestGatherNdAPI2(unittest.TestCase):
260259
def test_imperative(self):
261260
paddle.disable_static()
262261
paddle.set_device("mlu")
263-
input_1 = np.array([[1, 2], [3, 4], [5, 6]]).astype("float32")
264-
index_1 = np.array([[1]]).astype("int32")
262+
input_1 = np.array([[1, 2], [3, 4], [5, 6]])
263+
index_1 = np.array([[1]])
265264
input = base.dygraph.to_variable(input_1)
266265
index = base.dygraph.to_variable(index_1)
267266
output = paddle.gather(input, index)
@@ -271,7 +270,7 @@ def test_imperative(self):
271270
paddle.enable_static()
272271

273272

274-
for _typename in {"float16", "float32"}:
273+
for _typename in {"float16", "float32", "int64"}:
275274
test_class1("gather_nd", _typename)
276275
test_class2("gather_nd", _typename)
277276
test_class3("gather_nd", _typename)

‎backends/mlu/tests/unittests/test_transpose_op_mlu.py

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,8 @@ def setUp(self):
2929
self.init_op_type()
3030
self.initKernelType()
3131
self.initTestCase()
32-
self.inputs = {"X": np.random.random(self.shape).astype("float32")}
32+
self.init_dtype()
33+
self.inputs = {"X": np.random.random(self.shape).astype(self.dtype)}
3334
self.attrs = {
3435
"axis": list(self.axis),
3536
}
@@ -38,6 +39,9 @@ def setUp(self):
3839
def init_op_type(self):
3940
self.op_type = "transpose2"
4041

42+
def init_dtype(self):
43+
self.dtype = np.float32
44+
4145
def test_check_output(self):
4246
self.check_output_with_place(self.place)
4347

@@ -113,6 +117,22 @@ def initTestCase(self):
113117
self.axis = (6, 1, 3, 5, 0, 2, 4, 7)
114118

115119

120+
class TestTransposeOpFP16(TestTransposeOp):
121+
def init_dtype(self):
122+
self.dtype = np.float16
123+
124+
def test_check_grad(self):
125+
pass
126+
127+
128+
class TestTransposeOpInt64(TestTransposeOp):
129+
def init_dtype(self):
130+
self.dtype = np.int64
131+
132+
def test_check_grad(self):
133+
pass
134+
135+
116136
class TestTransposeOpBool(TestTransposeOp):
117137
def test_check_grad(self):
118138
pass

0 commit comments

Comments
 (0)