Skip to content
This repository has been archived by the owner on Nov 17, 2023. It is now read-only.

Commit

Permalink
change template to argument for req types
Browse files Browse the repository at this point in the history
  • Loading branch information
@haojin2
haojin2 committed Mar 15, 2019
1 parent 8d82354 commit 4e181d9
Show file tree
Hide file tree
Showing 2 changed files with 42 additions and 44 deletions.
43 changes: 21 additions & 22 deletions src/operator/contrib/index_copy.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,6 @@ void IndexCopyForward<cpu>(const nnvm::NodeAttrs& attrs,
});
}

template<int orig_req, int new_req>
struct index_copy_bwd_cpu {
template<typename DType, typename IType>
static void Map(int i,
Expand All @@ -79,7 +78,9 @@ struct index_copy_bwd_cpu {
DType* new_tensor_grad,
const IType* idx,
int dim_size,
int idx_size) {
int idx_size,
OpReqType orig_req,
OpReqType new_req) {
const int index = idx[i];
DType* new_ptr = new_tensor_grad + i * dim_size;
DType* orig_ptr = orig_tensor_grad + index * dim_size;
Expand Down Expand Up @@ -116,29 +117,27 @@ void IndexCopyBackward<cpu>(const nnvm::NodeAttrs& attrs,
const TBlob& in_grad_2 = outputs[2];
int dim_size = inputs[3].Size() / inputs[2].Size();
int index_size = inputs[2].Size();
OpReqType orig_req = req[0];
OpReqType new_req = req[2];
// index_copy_backward
MSHADOW_TYPE_SWITCH(out_grad.type_flag_, DType, {
MSHADOW_TYPE_SWITCH(index.type_flag_, IType, {
MXNET_REQ_TYPE_SWITCH(req[0], orig_req, {
MXNET_REQ_TYPE_SWITCH(req[2], new_req, {
switch (orig_req) {
case kNullOp:
break;
case kWriteTo:
case kWriteInplace:
copy(s, in_grad_1, out_grad);
break;
case kAddTo:
Kernel<op_with_req<op::mshadow_op::plus, kWriteInplace>, cpu>::Launch(
s, out_grad.Size(), in_grad_1.dptr<DType>(),
out_grad.dptr<DType>(), in_grad_1.dptr<DType>());
}
Kernel<index_copy_bwd_cpu<orig_req, new_req>, cpu>::Launch(
s, index_size, out_grad.dptr<DType>(),
in_grad_1.dptr<DType>(), in_grad_2.dptr<DType>(),
index.dptr<IType>(), dim_size, index_size);
});
});
switch (orig_req) {
case kNullOp:
break;
case kWriteTo:
case kWriteInplace:
copy(s, in_grad_1, out_grad);
break;
case kAddTo:
Kernel<op_with_req<op::mshadow_op::plus, kWriteInplace>, cpu>::Launch(
s, out_grad.Size(), in_grad_1.dptr<DType>(),
out_grad.dptr<DType>(), in_grad_1.dptr<DType>());
}
Kernel<index_copy_bwd_cpu, cpu>::Launch(
s, index_size, out_grad.dptr<DType>(),
in_grad_1.dptr<DType>(), in_grad_2.dptr<DType>(),
index.dptr<IType>(), dim_size, index_size, orig_req, new_req);
});
});
}
Expand Down
43 changes: 21 additions & 22 deletions src/operator/contrib/index_copy.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,6 @@ void IndexCopyForward<gpu>(const nnvm::NodeAttrs& attrs,
});
}

template<int orig_req, int new_req>
struct index_copy_bwd_gpu {
template<typename DType, typename IType>
MSHADOW_XINLINE static void Map(int i,
Expand All @@ -78,7 +77,9 @@ struct index_copy_bwd_gpu {
DType* new_grad,
const IType* idx,
int dim_size,
int idx_size) {
int idx_size,
OpReqType orig_req,
OpReqType new_req) {
int index = idx[i / dim_size];
KERNEL_ASSIGN(new_grad[i], new_req, out_grad[index * dim_size + i % dim_size]);
if (orig_req == kAddTo) {
Expand Down Expand Up @@ -108,29 +109,27 @@ void IndexCopyBackward<gpu>(const nnvm::NodeAttrs& attrs,
const TBlob& in_grad_2 = outputs[2];
int dim_size = inputs[3].Size() / inputs[2].Size();
int index_size = inputs[2].Size();
OpReqType orig_req = req[0];
OpReqType new_req = req[2];
// index_copy_backward
MSHADOW_TYPE_SWITCH(out_grad.type_flag_, DType, {
MSHADOW_TYPE_SWITCH(index.type_flag_, IType, {
MXNET_REQ_TYPE_SWITCH(req[0], orig_req, {
MXNET_REQ_TYPE_SWITCH(req[2], new_req, {
switch (orig_req) {
case kNullOp:
break;
case kWriteTo:
case kWriteInplace:
copy(s, in_grad_1, out_grad);
break;
case kAddTo:
Kernel<op_with_req<op::mshadow_op::plus, kWriteInplace>, gpu>::Launch(
s, out_grad.Size(), in_grad_1.dptr<DType>(),
out_grad.dptr<DType>(), in_grad_1.dptr<DType>());
}
Kernel<index_copy_bwd_gpu<orig_req, new_req>, gpu>::Launch(
s, in_grad_2.Size(), out_grad.dptr<DType>(),
in_grad_1.dptr<DType>(), in_grad_2.dptr<DType>(),
index.dptr<IType>(), dim_size, index_size);
});
});
switch (orig_req) {
case kNullOp:
break;
case kWriteTo:
case kWriteInplace:
copy(s, in_grad_1, out_grad);
break;
case kAddTo:
Kernel<op_with_req<op::mshadow_op::plus, kWriteInplace>, gpu>::Launch(
s, out_grad.Size(), in_grad_1.dptr<DType>(),
out_grad.dptr<DType>(), in_grad_1.dptr<DType>());
}
Kernel<index_copy_bwd_gpu, gpu>::Launch(
s, in_grad_2.Size(), out_grad.dptr<DType>(),
in_grad_1.dptr<DType>(), in_grad_2.dptr<DType>(),
index.dptr<IType>(), dim_size, index_size, orig_req, new_req);
});
});
}
Expand Down

0 comments on commit 4e181d9

Please sign in to comment.