mirror of https://github.com/open-mmlab/mmcv.git
[Fix] Fix MLU code format (#2887)
parent
9036241eac
commit
c5233598a5
|
@ -34,9 +34,9 @@ void ball_query_forward_mlu(int b, int n, int m, float min_radius,
|
|||
auto idx_ptr = idx_impl->cnnlMalloc();
|
||||
|
||||
auto handle = mluOpGetCurrentHandle();
|
||||
TORCH_MLUOP_CHECK(mluOpBallQuery(handle, new_xyz_desc.desc(), new_xyz_ptr, xyz_desc.desc(),
|
||||
xyz_ptr, min_radius, max_radius, nsample, idx_desc.desc(),
|
||||
idx_ptr));
|
||||
TORCH_MLUOP_CHECK(mluOpBallQuery(
|
||||
handle, new_xyz_desc.desc(), new_xyz_ptr, xyz_desc.desc(), xyz_ptr,
|
||||
min_radius, max_radius, nsample, idx_desc.desc(), idx_ptr));
|
||||
}
|
||||
|
||||
void ball_query_forward_impl(int b, int n, int m, float min_radius,
|
||||
|
|
|
@ -38,9 +38,9 @@ void BoxIouRotatedMLUKernelLauncher(const Tensor boxes1, const Tensor boxes2,
|
|||
auto ious_ptr = ious_impl->cnnlMalloc();
|
||||
|
||||
CNLOG(INFO) << "Call mluOpBoxIouRotated().";
|
||||
TORCH_MLUOP_CHECK(mluOpBoxIouRotated(handle, mode_flag, aligned, boxes1_desc.desc(), boxes1_ptr,
|
||||
boxes2_desc.desc(), boxes2_ptr, ious_desc.desc(),
|
||||
ious_ptr));
|
||||
TORCH_MLUOP_CHECK(mluOpBoxIouRotated(
|
||||
handle, mode_flag, aligned, boxes1_desc.desc(), boxes1_ptr,
|
||||
boxes2_desc.desc(), boxes2_ptr, ious_desc.desc(), ious_ptr));
|
||||
}
|
||||
|
||||
void box_iou_rotated_mlu(const Tensor boxes1, const Tensor boxes2, Tensor ious,
|
||||
|
|
|
@ -72,12 +72,12 @@ void CARAFEForwardMLUKernelLauncher(const Tensor input, const Tensor mask,
|
|||
auto handle = mluOpGetCurrentHandle();
|
||||
mluOpCarafeDescriptor_t carafe_desc;
|
||||
TORCH_MLUOP_CHECK(mluOpCreateCarafeDescriptor(&carafe_desc));
|
||||
TORCH_MLUOP_CHECK(mluOpSetCarafeDescriptor(carafe_desc, input.dim(), kernel_size, group_size,
|
||||
scale_factor));
|
||||
TORCH_MLUOP_CHECK(mluOpSetCarafeDescriptor(
|
||||
carafe_desc, input.dim(), kernel_size, group_size, scale_factor));
|
||||
// launch kernel
|
||||
TORCH_MLUOP_CHECK(mluOpCarafeForward(handle, carafe_desc, input_desc.desc(), input_ptr,
|
||||
mask_desc.desc(), mask_ptr, output_desc.desc(),
|
||||
output_ptr));
|
||||
TORCH_MLUOP_CHECK(mluOpCarafeForward(handle, carafe_desc, input_desc.desc(),
|
||||
input_ptr, mask_desc.desc(), mask_ptr,
|
||||
output_desc.desc(), output_ptr));
|
||||
// destroy op descriptor
|
||||
TORCH_MLUOP_CHECK(mluOpDestroyCarafeDescriptor(carafe_desc));
|
||||
|
||||
|
@ -160,13 +160,14 @@ void CARAFEBackwardMLUKernelLauncher(
|
|||
auto handle = mluOpGetCurrentHandle();
|
||||
mluOpCarafeDescriptor_t carafe_desc;
|
||||
TORCH_MLUOP_CHECK(mluOpCreateCarafeDescriptor(&carafe_desc));
|
||||
TORCH_MLUOP_CHECK(mluOpSetCarafeDescriptor(carafe_desc, grad_output.dim(), kernel_size,
|
||||
group_size, scale_factor));
|
||||
TORCH_MLUOP_CHECK(mluOpSetCarafeDescriptor(
|
||||
carafe_desc, grad_output.dim(), kernel_size, group_size, scale_factor));
|
||||
// launch kernel
|
||||
TORCH_MLUOP_CHECK(mluOpCarafeBackward(handle, carafe_desc, input_desc.desc(), input_ptr,
|
||||
mask_desc.desc(), mask_ptr, grad_output_desc.desc(),
|
||||
grad_output_ptr, grad_input_desc.desc(), grad_input_ptr,
|
||||
grad_mask_desc.desc(), grad_mask_ptr));
|
||||
TORCH_MLUOP_CHECK(mluOpCarafeBackward(
|
||||
handle, carafe_desc, input_desc.desc(), input_ptr, mask_desc.desc(),
|
||||
mask_ptr, grad_output_desc.desc(), grad_output_ptr,
|
||||
grad_input_desc.desc(), grad_input_ptr, grad_mask_desc.desc(),
|
||||
grad_mask_ptr));
|
||||
// destroy op descriptor
|
||||
TORCH_MLUOP_CHECK(mluOpDestroyCarafeDescriptor(carafe_desc));
|
||||
|
||||
|
|
|
@ -51,9 +51,9 @@ void DeformRoIPoolForwardMLUKernelLauncher(Tensor input, Tensor rois,
|
|||
// get compute handle
|
||||
auto handle = mluOpGetCurrentHandle();
|
||||
TORCH_MLUOP_CHECK(mluOpDeformRoiPoolForward(
|
||||
handle, input_desc.desc(), input_ptr, rois_desc.desc(), rois_ptr,
|
||||
offset_real_desc, offset_ptr, pooled_height, pooled_width, spatial_scale,
|
||||
sampling_ratio, gamma, output_desc.desc(), output_ptr));
|
||||
handle, input_desc.desc(), input_ptr, rois_desc.desc(), rois_ptr,
|
||||
offset_real_desc, offset_ptr, pooled_height, pooled_width, spatial_scale,
|
||||
sampling_ratio, gamma, output_desc.desc(), output_ptr));
|
||||
|
||||
output.copy_(output_contiguous);
|
||||
}
|
||||
|
@ -114,11 +114,11 @@ void DeformRoIPoolBackwardMLUKernelLauncher(
|
|||
// get compute handle
|
||||
auto handle = mluOpGetCurrentHandle();
|
||||
TORCH_MLUOP_CHECK(mluOpDeformRoiPoolBackward(
|
||||
handle, grad_output_desc.desc(), grad_output_ptr, input_desc.desc(),
|
||||
input_ptr, rois_desc.desc(), rois_ptr, offset_real_desc, offset_ptr,
|
||||
pooled_height, pooled_width, spatial_scale, sampling_ratio, gamma,
|
||||
grad_input_desc.desc(), grad_input_ptr, grad_offset_real_desc,
|
||||
grad_offset_ptr));
|
||||
handle, grad_output_desc.desc(), grad_output_ptr, input_desc.desc(),
|
||||
input_ptr, rois_desc.desc(), rois_ptr, offset_real_desc, offset_ptr,
|
||||
pooled_height, pooled_width, spatial_scale, sampling_ratio, gamma,
|
||||
grad_input_desc.desc(), grad_input_ptr, grad_offset_real_desc,
|
||||
grad_offset_ptr));
|
||||
grad_input.copy_(grad_input_);
|
||||
}
|
||||
|
||||
|
|
|
@ -43,8 +43,8 @@ Tensor diff_iou_rotated_sort_vertices_forward_mlu(Tensor vertices, Tensor mask,
|
|||
|
||||
// launch kernel
|
||||
TORCH_MLUOP_CHECK(mluOpDiffIouRotatedSortVerticesForward(
|
||||
handle, vertices_desc.desc(), vertices_ptr, mask_desc.desc(), mask_ptr,
|
||||
num_valid_desc.desc(), num_valid_ptr, idx_desc.desc(), idx_ptr));
|
||||
handle, vertices_desc.desc(), vertices_ptr, mask_desc.desc(), mask_ptr,
|
||||
num_valid_desc.desc(), num_valid_ptr, idx_desc.desc(), idx_ptr));
|
||||
return idx;
|
||||
}
|
||||
|
||||
|
|
|
@ -30,7 +30,8 @@ void IoU3DNMS3DMLUKernelLauncher(Tensor boxes, Tensor &keep, Tensor &keep_num,
|
|||
// workspace
|
||||
size_t workspace_size = 0;
|
||||
auto handle = mluOpGetCurrentHandle();
|
||||
TORCH_MLUOP_CHECK(mluOpGetNmsWorkspaceSize(handle, boxes_desc.desc(), NULL, &workspace_size));
|
||||
TORCH_MLUOP_CHECK(mluOpGetNmsWorkspaceSize(handle, boxes_desc.desc(), NULL,
|
||||
&workspace_size));
|
||||
auto workspace = at::empty(workspace_size, boxes.options().dtype(at::kByte));
|
||||
|
||||
// get compute queue
|
||||
|
@ -57,14 +58,14 @@ void IoU3DNMS3DMLUKernelLauncher(Tensor boxes, Tensor &keep, Tensor &keep_num,
|
|||
const float offset = 0.0;
|
||||
|
||||
TORCH_MLUOP_CHECK(mluOpCreateNmsDescriptor(&nms_desc));
|
||||
TORCH_MLUOP_CHECK(mluOpSetNmsDescriptor(nms_desc, box_mode, output_mode, algo, method_mode,
|
||||
iou_threshold, soft_nms_sigma, max_output_size,
|
||||
confidence_threshold, offset, input_layout,
|
||||
pad_to_max_output_size));
|
||||
TORCH_MLUOP_CHECK(mluOpSetNmsDescriptor(
|
||||
nms_desc, box_mode, output_mode, algo, method_mode, iou_threshold,
|
||||
soft_nms_sigma, max_output_size, confidence_threshold, offset,
|
||||
input_layout, pad_to_max_output_size));
|
||||
|
||||
TORCH_MLUOP_CHECK(mluOpNms(handle, nms_desc, boxes_desc.desc(), boxes_ptr, NULL, NULL,
|
||||
workspace_ptr, workspace_size, output_desc.desc(), output_ptr,
|
||||
output_size_ptr));
|
||||
TORCH_MLUOP_CHECK(mluOpNms(handle, nms_desc, boxes_desc.desc(), boxes_ptr,
|
||||
NULL, NULL, workspace_ptr, workspace_size,
|
||||
output_desc.desc(), output_ptr, output_size_ptr));
|
||||
TORCH_MLUOP_CHECK(mluOpDestroyNmsDescriptor(nms_desc));
|
||||
}
|
||||
|
||||
|
|
|
@ -123,7 +123,8 @@ void MluOpTensorDescriptor::set_desc(const at::Tensor& t,
|
|||
mluOpDataType_t dtype,
|
||||
std::vector<int>& dims) {
|
||||
int dimNb = dims.size();
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(desc_, layout, dtype, dimNb, dims.data()));
|
||||
TORCH_MLUOP_CHECK(
|
||||
mluOpSetTensorDescriptor(desc_, layout, dtype, dimNb, dims.data()));
|
||||
}
|
||||
|
||||
// Handles
|
||||
|
|
|
@ -35,13 +35,13 @@
|
|||
auto NAME##_ptr = NAME##_impl->cnnlMalloc();
|
||||
|
||||
#ifndef TORCH_MLUOP_CHECK
|
||||
#define TORCH_MLUOP_CHECK(EXPR) \
|
||||
do { \
|
||||
mluOpStatus_t status = EXPR; \
|
||||
if (status != MLUOP_STATUS_SUCCESS) { \
|
||||
CNLOG(ERROR) << ""; \
|
||||
TORCH_CHECK(false, "MLUOPS error: ", mluOpGetErrorString(status)); \
|
||||
} \
|
||||
#define TORCH_MLUOP_CHECK(EXPR) \
|
||||
do { \
|
||||
mluOpStatus_t status = EXPR; \
|
||||
if (status != MLUOP_STATUS_SUCCESS) { \
|
||||
CNLOG(ERROR) << ""; \
|
||||
TORCH_CHECK(false, "MLUOPS error: ", mluOpGetErrorString(status)); \
|
||||
} \
|
||||
} while (0);
|
||||
#endif
|
||||
|
||||
|
@ -65,8 +65,12 @@ mluOpReduceMode_t getMluOpReduceMode(const reduce_t reduce_type);
|
|||
|
||||
class MluOpTensorDescriptor {
|
||||
public:
|
||||
MluOpTensorDescriptor() { TORCH_MLUOP_CHECK(mluOpCreateTensorDescriptor(&desc_)); };
|
||||
~MluOpTensorDescriptor() { TORCH_MLUOP_CHECK(mluOpDestroyTensorDescriptor(desc_)); }
|
||||
MluOpTensorDescriptor() {
|
||||
TORCH_MLUOP_CHECK(mluOpCreateTensorDescriptor(&desc_));
|
||||
};
|
||||
~MluOpTensorDescriptor() {
|
||||
TORCH_MLUOP_CHECK(mluOpDestroyTensorDescriptor(desc_));
|
||||
}
|
||||
|
||||
void set(at::Tensor);
|
||||
void set_with_layout(at::Tensor, mluOpTensorLayout_t layout);
|
||||
|
@ -89,7 +93,9 @@ class MluOpHandle {
|
|||
handle = nullptr;
|
||||
}
|
||||
}
|
||||
void setQueue(cnrtQueue_t queue) { TORCH_MLUOP_CHECK(mluOpSetQueue(handle, queue)); }
|
||||
void setQueue(cnrtQueue_t queue) {
|
||||
TORCH_MLUOP_CHECK(mluOpSetQueue(handle, queue));
|
||||
}
|
||||
mluOpHandle_t handle;
|
||||
};
|
||||
|
||||
|
|
|
@ -36,11 +36,11 @@ Tensor MsDeformAttnForwardLauncher(const Tensor& value,
|
|||
INITIAL_MLU_PARAM_WITH_TENSOR(attn_weight);
|
||||
|
||||
TORCH_MLUOP_CHECK(mluOpMsDeformAttnForward(
|
||||
handle, value_desc.desc(), value_ptr, spatial_shapes_int_desc.desc(),
|
||||
spatial_shapes_int_ptr, level_start_index_int_desc.desc(),
|
||||
level_start_index_int_ptr, sampling_loc_desc.desc(), sampling_loc_ptr,
|
||||
attn_weight_desc.desc(), attn_weight_ptr, im2col_step, output_desc.desc(),
|
||||
output_ptr));
|
||||
handle, value_desc.desc(), value_ptr, spatial_shapes_int_desc.desc(),
|
||||
spatial_shapes_int_ptr, level_start_index_int_desc.desc(),
|
||||
level_start_index_int_ptr, sampling_loc_desc.desc(), sampling_loc_ptr,
|
||||
attn_weight_desc.desc(), attn_weight_ptr, im2col_step, output_desc.desc(),
|
||||
output_ptr));
|
||||
|
||||
output = output.view({batch_size, num_queries, num_heads * channels});
|
||||
return output;
|
||||
|
|
|
@ -34,8 +34,8 @@ Tensor NMSMLUKernelLauncher(Tensor boxes, Tensor scores, float iou_threshold,
|
|||
// workspace
|
||||
size_t workspace_size = 0;
|
||||
auto handle = mluOpGetCurrentHandle();
|
||||
TORCH_MLUOP_CHECK(mluOpGetNmsWorkspaceSize(handle, boxes_desc.desc(), scores_desc.desc(),
|
||||
&workspace_size));
|
||||
TORCH_MLUOP_CHECK(mluOpGetNmsWorkspaceSize(
|
||||
handle, boxes_desc.desc(), scores_desc.desc(), &workspace_size));
|
||||
auto workspace = at::empty(workspace_size, boxes.options().dtype(at::kByte));
|
||||
|
||||
// get compute queue
|
||||
|
@ -63,14 +63,15 @@ Tensor NMSMLUKernelLauncher(Tensor boxes, Tensor scores, float iou_threshold,
|
|||
const int max_output_size = max_output_boxes;
|
||||
|
||||
TORCH_MLUOP_CHECK(mluOpCreateNmsDescriptor(&nms_desc));
|
||||
TORCH_MLUOP_CHECK(mluOpSetNmsDescriptor(nms_desc, box_mode, output_mode, algo, method_mode,
|
||||
iou_threshold, soft_nms_sigma, max_output_size,
|
||||
confidence_threshold, (float)offset, input_layout,
|
||||
pad_to_max_output_size));
|
||||
TORCH_MLUOP_CHECK(mluOpSetNmsDescriptor(
|
||||
nms_desc, box_mode, output_mode, algo, method_mode, iou_threshold,
|
||||
soft_nms_sigma, max_output_size, confidence_threshold, (float)offset,
|
||||
input_layout, pad_to_max_output_size));
|
||||
|
||||
TORCH_MLUOP_CHECK(mluOpNms(handle, nms_desc, boxes_desc.desc(), boxes_ptr, scores_desc.desc(),
|
||||
scores_ptr, workspace_ptr, workspace_size, output_desc.desc(),
|
||||
output_ptr, output_size_ptr));
|
||||
TORCH_MLUOP_CHECK(mluOpNms(handle, nms_desc, boxes_desc.desc(), boxes_ptr,
|
||||
scores_desc.desc(), scores_ptr, workspace_ptr,
|
||||
workspace_size, output_desc.desc(), output_ptr,
|
||||
output_size_ptr));
|
||||
TORCH_MLUOP_CHECK(mluOpDestroyNmsDescriptor(nms_desc));
|
||||
int output_num = *static_cast<int *>(output_size.cpu().data_ptr());
|
||||
auto ret = output.to(boxes.options().dtype(at::kLong));
|
||||
|
|
|
@ -30,7 +30,8 @@ Tensor nms_rotated_mlu(Tensor boxes, Tensor scores, float iou_threshold) {
|
|||
// workspace
|
||||
size_t workspace_size = 0;
|
||||
auto handle = mluOpGetCurrentHandle();
|
||||
TORCH_MLUOP_CHECK(mluOpGetNmsRotatedWorkspaceSize(handle, boxes_desc.desc(), &workspace_size));
|
||||
TORCH_MLUOP_CHECK(mluOpGetNmsRotatedWorkspaceSize(handle, boxes_desc.desc(),
|
||||
&workspace_size));
|
||||
auto workspace = at::empty(workspace_size, boxes.options().dtype(at::kByte));
|
||||
|
||||
auto boxes_impl = torch_mlu::getMluTensorImpl(boxes_);
|
||||
|
@ -44,9 +45,10 @@ Tensor nms_rotated_mlu(Tensor boxes, Tensor scores, float iou_threshold) {
|
|||
auto output_size_impl = torch_mlu::getMluTensorImpl(output_size);
|
||||
auto output_size_ptr = output_size_impl->cnnlMalloc();
|
||||
|
||||
TORCH_MLUOP_CHECK(mluOpNmsRotated(handle, iou_threshold, boxes_desc.desc(), boxes_ptr,
|
||||
scores_desc.desc(), scores_ptr, workspace_ptr, workspace_size,
|
||||
output_desc.desc(), output_ptr, (int *)output_size_ptr));
|
||||
TORCH_MLUOP_CHECK(mluOpNmsRotated(
|
||||
handle, iou_threshold, boxes_desc.desc(), boxes_ptr, scores_desc.desc(),
|
||||
scores_ptr, workspace_ptr, workspace_size, output_desc.desc(), output_ptr,
|
||||
(int *)output_size_ptr));
|
||||
int output_num = *static_cast<int *>(output_size.cpu().data_ptr());
|
||||
auto ret = output.to(boxes.options().dtype(at::kLong));
|
||||
return ret.slice(0, 0, output_num);
|
||||
|
|
|
@ -35,8 +35,8 @@ void PSAMaskForwardMLUKernelLauncher(const int psa_type, const Tensor x,
|
|||
auto y_impl = torch_mlu::getMluTensorImpl(y_tmp);
|
||||
auto y_ptr = y_impl->cnnlMalloc();
|
||||
|
||||
TORCH_MLUOP_CHECK(mluOpPsamaskForward(handle, psa_type, x_desc.desc(), x_ptr, h_mask, w_mask,
|
||||
y_desc.desc(), y_ptr));
|
||||
TORCH_MLUOP_CHECK(mluOpPsamaskForward(handle, psa_type, x_desc.desc(), x_ptr,
|
||||
h_mask, w_mask, y_desc.desc(), y_ptr));
|
||||
|
||||
y.copy_(y_tmp);
|
||||
}
|
||||
|
@ -67,7 +67,8 @@ void PSAMaskBackwardMLUKernelLauncher(const int psa_type, const Tensor dy,
|
|||
auto dy_impl = torch_mlu::getMluTensorImpl(dy_tensor);
|
||||
auto dy_ptr = dy_impl->cnnlMalloc();
|
||||
|
||||
TORCH_MLUOP_CHECK(mluOpPsamaskBackward(handle, psa_type, dy_desc.desc(), dy_ptr, h_mask, w_mask,
|
||||
TORCH_MLUOP_CHECK(mluOpPsamaskBackward(handle, psa_type, dy_desc.desc(),
|
||||
dy_ptr, h_mask, w_mask,
|
||||
dx_tmp_desc.desc(), dx_ptr));
|
||||
|
||||
dx.copy_(dx_tmp);
|
||||
|
|
|
@ -49,9 +49,9 @@ void ROIAlignForwardMLUKernelLauncher(Tensor input, Tensor rois, Tensor output,
|
|||
|
||||
mluOpRoiAlignForwardDescriptor_t roialign_desc;
|
||||
TORCH_MLUOP_CHECK(mluOpCreateRoiAlignForwardDescriptor(&roialign_desc));
|
||||
TORCH_MLUOP_CHECK(mluOpSetRoiAlignForwardDescriptor_v2(roialign_desc, aligned_height,
|
||||
aligned_width, sampling_ratio,
|
||||
spatial_scale, pool_mode, aligned));
|
||||
TORCH_MLUOP_CHECK(mluOpSetRoiAlignForwardDescriptor_v2(
|
||||
roialign_desc, aligned_height, aligned_width, sampling_ratio,
|
||||
spatial_scale, pool_mode, aligned));
|
||||
|
||||
auto handle = mluOpGetCurrentHandle();
|
||||
if (pool_mode == 0) {
|
||||
|
@ -65,16 +65,16 @@ void ROIAlignForwardMLUKernelLauncher(Tensor input, Tensor rois, Tensor output,
|
|||
auto argmax_y_ptr = argmax_y_impl->cnnlMalloc();
|
||||
argmax_y_desc.set_with_layout(argmax_x_contiguous, MLUOP_LAYOUT_NHWC);
|
||||
argmax_x_desc.set_with_layout(argmax_x_contiguous, MLUOP_LAYOUT_NHWC);
|
||||
TORCH_MLUOP_CHECK(mluOpRoiAlignForward_v2(handle, roialign_desc, input_desc.desc(), self_ptr,
|
||||
rois_desc.desc(), rois_ptr, output_desc.desc(),
|
||||
output_ptr, argmax_x_desc.desc(), argmax_x_ptr,
|
||||
argmax_y_desc.desc(), argmax_y_ptr);
|
||||
TORCH_MLUOP_CHECK(mluOpRoiAlignForward_v2(
|
||||
handle, roialign_desc, input_desc.desc(), self_ptr, rois_desc.desc(),
|
||||
rois_ptr, output_desc.desc(), output_ptr, argmax_x_desc.desc(),
|
||||
argmax_x_ptr, argmax_y_desc.desc(), argmax_y_ptr));
|
||||
argmax_x.copy_(argmax_x_contiguous);
|
||||
argmax_y.copy_(argmax_y_contiguous);
|
||||
} else {
|
||||
TORCH_MLUOP_CHECK(mluOpRoiAlignForward_v2(handle, roialign_desc, input_desc.desc(), self_ptr,
|
||||
rois_desc.desc(), rois_ptr, output_desc.desc(),
|
||||
output_ptr, NULL, NULL, NULL, NULL);
|
||||
TORCH_MLUOP_CHECK(mluOpRoiAlignForward_v2(
|
||||
handle, roialign_desc, input_desc.desc(), self_ptr, rois_desc.desc(),
|
||||
rois_ptr, output_desc.desc(), output_ptr, NULL, NULL, NULL, NULL));
|
||||
}
|
||||
TORCH_MLUOP_CHECK(mluOpDestroyRoiAlignForwardDescriptor(roialign_desc));
|
||||
output.copy_(output_contiguous);
|
||||
|
@ -136,16 +136,16 @@ void ROIAlignBackwardMLUKernelLauncher(Tensor grad, Tensor rois,
|
|||
auto argmax_y_ptr = argmax_y_impl->cnnlMalloc();
|
||||
argmax_y_desc.set_with_layout(argmax_x_contiguous, MLUOP_LAYOUT_NHWC);
|
||||
argmax_x_desc.set_with_layout(argmax_x_contiguous, MLUOP_LAYOUT_NHWC);
|
||||
TORCH_MLUOP_CHECK(mluOpRoiAlignBackward_v2(handle, grads_desc.desc(), grad_ptr,
|
||||
rois_desc.desc(), rois_ptr, argmax_y_desc.desc(),
|
||||
argmax_x_ptr, argmax_y_desc.desc(), argmax_y_ptr,
|
||||
spatial_scale, sampling_ratio, aligned, pool_mode,
|
||||
grad_input_desc.desc(), grad_input_ptr));
|
||||
TORCH_MLUOP_CHECK(mluOpRoiAlignBackward_v2(
|
||||
handle, grads_desc.desc(), grad_ptr, rois_desc.desc(), rois_ptr,
|
||||
argmax_y_desc.desc(), argmax_x_ptr, argmax_y_desc.desc(), argmax_y_ptr,
|
||||
spatial_scale, sampling_ratio, aligned, pool_mode,
|
||||
grad_input_desc.desc(), grad_input_ptr));
|
||||
} else {
|
||||
TORCH_MLUOP_CHECK(mluOpRoiAlignBackward_v2(handle, grads_desc.desc(), grad_ptr,
|
||||
rois_desc.desc(), rois_ptr, NULL, NULL, NULL, NULL,
|
||||
spatial_scale, sampling_ratio, aligned, pool_mode,
|
||||
grad_input_desc.desc(), grad_input_ptr));
|
||||
TORCH_MLUOP_CHECK(mluOpRoiAlignBackward_v2(
|
||||
handle, grads_desc.desc(), grad_ptr, rois_desc.desc(), rois_ptr, NULL,
|
||||
NULL, NULL, NULL, spatial_scale, sampling_ratio, aligned, pool_mode,
|
||||
grad_input_desc.desc(), grad_input_ptr));
|
||||
}
|
||||
grad_input.copy_(grad_input_);
|
||||
}
|
||||
|
|
|
@ -41,9 +41,9 @@ void ROIAlignRotatedForwardMLUKernelLauncher(Tensor input, Tensor rois,
|
|||
// get compute handle
|
||||
auto handle = mluOpGetCurrentHandle();
|
||||
TORCH_MLUOP_CHECK(mluOpRoiAlignRotatedForward(
|
||||
handle, input_desc.desc(), input_ptr, rois_desc.desc(), rois_ptr,
|
||||
pooled_height, pooled_width, sampling_ratio, spatial_scale, aligned,
|
||||
clockwise, output_desc.desc(), output_ptr);
|
||||
handle, input_desc.desc(), input_ptr, rois_desc.desc(), rois_ptr,
|
||||
pooled_height, pooled_width, sampling_ratio, spatial_scale, aligned,
|
||||
clockwise, output_desc.desc(), output_ptr));
|
||||
|
||||
output.copy_(output_contiguous);
|
||||
}
|
||||
|
@ -77,9 +77,9 @@ void ROIAlignRotatedBackwardMLUKernelLauncher(
|
|||
// get compute handle
|
||||
auto handle = mluOpGetCurrentHandle();
|
||||
TORCH_MLUOP_CHECK(mluOpRoiAlignRotatedBackward(
|
||||
handle, top_grad_desc.desc(), top_grad_ptr, rois_desc.desc(), rois_ptr,
|
||||
pooled_height, pooled_width, sampling_ratio, spatial_scale, aligned,
|
||||
clockwise, bottom_grad_desc.desc(), bottom_grad_ptr);
|
||||
handle, top_grad_desc.desc(), top_grad_ptr, rois_desc.desc(), rois_ptr,
|
||||
pooled_height, pooled_width, sampling_ratio, spatial_scale, aligned,
|
||||
clockwise, bottom_grad_desc.desc(), bottom_grad_ptr));
|
||||
bottom_grad.copy_(bottom_grad_);
|
||||
}
|
||||
|
||||
|
|
|
@ -45,8 +45,8 @@ void RoiawarePool3dForwardMLUKernelLauncher(
|
|||
// allocate extra space for workspace
|
||||
size_t workspace_size = 0;
|
||||
TORCH_MLUOP_CHECK(mluOpGetRoiawarePool3dForwardWorkspaceSize(
|
||||
handle, rois_desc.desc(), pts_desc.desc(), pts_feature_desc.desc(),
|
||||
&workspace_size));
|
||||
handle, rois_desc.desc(), pts_desc.desc(), pts_feature_desc.desc(),
|
||||
&workspace_size));
|
||||
|
||||
auto workspace = at::empty(workspace_size, rois.options().dtype(at::kByte));
|
||||
auto workspace_impl = torch_mlu::getMluTensorImpl(workspace);
|
||||
|
|
|
@ -76,10 +76,10 @@ void RotatedFeatureAlignBackwardMLUKernelLauncher(const Tensor top_grad,
|
|||
|
||||
// get compute handle
|
||||
auto handle = mluOpGetCurrentHandle();
|
||||
TORCH_MLUOP_CHECK(mluOpRotatedFeatureAlignBackward(handle, top_grad_desc.desc(), top_grad_ptr,
|
||||
best_bboxes_desc.desc(), best_bboxes_ptr,
|
||||
spatial_scale, points,
|
||||
bottom_grad_desc.desc(), bottom_grad_ptr));
|
||||
TORCH_MLUOP_CHECK(mluOpRotatedFeatureAlignBackward(
|
||||
handle, top_grad_desc.desc(), top_grad_ptr, best_bboxes_desc.desc(),
|
||||
best_bboxes_ptr, spatial_scale, points, bottom_grad_desc.desc(),
|
||||
bottom_grad_ptr));
|
||||
bottom_grad.copy_(bottom_grad_);
|
||||
}
|
||||
|
||||
|
|
|
@ -86,17 +86,17 @@ std::vector<torch::Tensor> GetIndicePairsForwardMLUKernelLauncher(
|
|||
mluOpDataType_t dtype = MLUOP_DTYPE_INT32;
|
||||
std::vector<int> dims;
|
||||
dims = {numAct, coorDim + 1};
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(indices_desc.desc(), layout, dtype, dims.size(),
|
||||
dims.data()));
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(
|
||||
indices_desc.desc(), layout, dtype, dims.size(), dims.data()));
|
||||
dims = {kernelVolume, 2, numAct};
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(indicePairs_desc.desc(), layout, dtype,
|
||||
dims.size(), dims.data()));
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(
|
||||
indicePairs_desc.desc(), layout, dtype, dims.size(), dims.data()));
|
||||
dims = {kernelVolume};
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(indiceNum_desc.desc(), layout, dtype, dims.size(),
|
||||
dims.data()));
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(
|
||||
indiceNum_desc.desc(), layout, dtype, dims.size(), dims.data()));
|
||||
dims = {out_size, coorDim + 1};
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(out_indices_desc.desc(), layout, dtype,
|
||||
dims.size(), dims.data()));
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(
|
||||
out_indices_desc.desc(), layout, dtype, dims.size(), dims.data()));
|
||||
}
|
||||
|
||||
mluOpSparseConvolutionDescriptor_t sparse_conv_desc;
|
||||
|
@ -127,13 +127,14 @@ std::vector<torch::Tensor> GetIndicePairsForwardMLUKernelLauncher(
|
|||
auto indiceNum_ptr = indiceNum_impl->cnnlMalloc();
|
||||
auto indice_workspace_ptr = indice_workspace_impl->cnnlMalloc();
|
||||
|
||||
TORCH_MLUOP_CHECK(mluOpGetIndicePairs(handle, sparse_conv_desc, indices_desc.desc(),
|
||||
indices_ptr, indice_workspace_ptr, workspace_size,
|
||||
indicePairs_desc.desc(), indicePairs_ptr,
|
||||
out_indices_desc.desc(), out_indices_ptr,
|
||||
indiceNum_desc.desc(), indiceNum_ptr));
|
||||
TORCH_MLUOP_CHECK(mluOpGetIndicePairs(
|
||||
handle, sparse_conv_desc, indices_desc.desc(), indices_ptr,
|
||||
indice_workspace_ptr, workspace_size, indicePairs_desc.desc(),
|
||||
indicePairs_ptr, out_indices_desc.desc(), out_indices_ptr,
|
||||
indiceNum_desc.desc(), indiceNum_ptr));
|
||||
int num_act_out = 0;
|
||||
TORCH_MLUOP_CHECK(mluOpGetSparseConvolutionNumActOut(sparse_conv_desc, &num_act_out));
|
||||
TORCH_MLUOP_CHECK(
|
||||
mluOpGetSparseConvolutionNumActOut(sparse_conv_desc, &num_act_out));
|
||||
TORCH_MLUOP_CHECK(mluOpDestroySparseConvolutionDescriptor(sparse_conv_desc));
|
||||
if (!sub_m) {
|
||||
return {out_indices.slice(0, 0, num_act_out), indicePairs, indiceNum};
|
||||
|
@ -179,25 +180,28 @@ torch::Tensor IndiceConvForwardMLUKernelLauncher(
|
|||
int dims[8];
|
||||
|
||||
// features_desc
|
||||
TORCH_MLUOP_CHECK(mluOpGetTensorDescriptor(features_desc.desc(), &layout, &dtype, &dim, dims));
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(features_desc.desc(), MLUOP_LAYOUT_ARRAY, dtype,
|
||||
dim, dims));
|
||||
TORCH_MLUOP_CHECK(mluOpGetTensorDescriptor(features_desc.desc(), &layout,
|
||||
&dtype, &dim, dims));
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(
|
||||
features_desc.desc(), MLUOP_LAYOUT_ARRAY, dtype, dim, dims));
|
||||
|
||||
// filters_desc
|
||||
TORCH_MLUOP_CHECK(mluOpGetTensorDescriptor(filters_desc.desc(), &layout, &dtype, &dim, dims));
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(filters_desc.desc(), MLUOP_LAYOUT_ARRAY, dtype,
|
||||
dim, dims));
|
||||
TORCH_MLUOP_CHECK(mluOpGetTensorDescriptor(filters_desc.desc(), &layout,
|
||||
&dtype, &dim, dims));
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(
|
||||
filters_desc.desc(), MLUOP_LAYOUT_ARRAY, dtype, dim, dims));
|
||||
|
||||
// indice_pairs_desc
|
||||
TORCH_MLUOP_CHECK(mluOpGetTensorDescriptor(indice_pairs_desc.desc(), &layout, &dtype, &dim,
|
||||
dims));
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(indice_pairs_desc.desc(), MLUOP_LAYOUT_ARRAY,
|
||||
dtype, dim, dims));
|
||||
TORCH_MLUOP_CHECK(mluOpGetTensorDescriptor(indice_pairs_desc.desc(),
|
||||
&layout, &dtype, &dim, dims));
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(
|
||||
indice_pairs_desc.desc(), MLUOP_LAYOUT_ARRAY, dtype, dim, dims));
|
||||
|
||||
// output_desc
|
||||
TORCH_MLUOP_CHECK(mluOpGetTensorDescriptor(output_desc.desc(), &layout, &dtype, &dim, dims));
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(output_desc.desc(), MLUOP_LAYOUT_ARRAY, dtype, dim,
|
||||
dims));
|
||||
TORCH_MLUOP_CHECK(mluOpGetTensorDescriptor(output_desc.desc(), &layout,
|
||||
&dtype, &dim, dims));
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(
|
||||
output_desc.desc(), MLUOP_LAYOUT_ARRAY, dtype, dim, dims));
|
||||
}
|
||||
|
||||
auto handle = mluOpGetCurrentHandle();
|
||||
|
@ -290,37 +294,39 @@ std::vector<torch::Tensor> IndiceConvBackwardMLUKernelLauncher(
|
|||
int dims[8];
|
||||
|
||||
// features_desc
|
||||
TORCH_MLUOP_CHECK(mluOpGetTensorDescriptor(features_desc.desc(), &layout, &dtype, &dim, dims));
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(features_desc.desc(), MLUOP_LAYOUT_ARRAY, dtype,
|
||||
dim, dims));
|
||||
TORCH_MLUOP_CHECK(mluOpGetTensorDescriptor(features_desc.desc(), &layout,
|
||||
&dtype, &dim, dims));
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(
|
||||
features_desc.desc(), MLUOP_LAYOUT_ARRAY, dtype, dim, dims));
|
||||
|
||||
// filters_desc
|
||||
TORCH_MLUOP_CHECK(mluOpGetTensorDescriptor(filters_desc.desc(), &layout, &dtype, &dim, dims));
|
||||
TORCH_MLUOP_CHECK(mluOpGetTensorDescriptor(filters_desc.desc(), &layout,
|
||||
&dtype, &dim, dims));
|
||||
if (dim == 4) {
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(filters_desc.desc(), MLUOP_LAYOUT_HWCN, dtype,
|
||||
dim, dims));
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(
|
||||
filters_desc.desc(), MLUOP_LAYOUT_HWCN, dtype, dim, dims));
|
||||
} else {
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(filters_desc.desc(), MLUOP_LAYOUT_ARRAY, dtype,
|
||||
dim, dims));
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(
|
||||
filters_desc.desc(), MLUOP_LAYOUT_ARRAY, dtype, dim, dims));
|
||||
}
|
||||
|
||||
// output_grad_desc
|
||||
TORCH_MLUOP_CHECK(mluOpGetTensorDescriptor(output_grad_desc.desc(), &layout, &dtype, &dim,
|
||||
dims));
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(output_grad_desc.desc(), MLUOP_LAYOUT_ARRAY, dtype,
|
||||
dim, dims));
|
||||
TORCH_MLUOP_CHECK(mluOpGetTensorDescriptor(output_grad_desc.desc(), &layout,
|
||||
&dtype, &dim, dims));
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(
|
||||
output_grad_desc.desc(), MLUOP_LAYOUT_ARRAY, dtype, dim, dims));
|
||||
|
||||
// indice_pairs_desc
|
||||
TORCH_MLUOP_CHECK(mluOpGetTensorDescriptor(indice_pairs_desc.desc(), &layout, &dtype, &dim,
|
||||
dims));
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(indice_pairs_desc.desc(), MLUOP_LAYOUT_ARRAY,
|
||||
dtype, dim, dims));
|
||||
TORCH_MLUOP_CHECK(mluOpGetTensorDescriptor(indice_pairs_desc.desc(),
|
||||
&layout, &dtype, &dim, dims));
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(
|
||||
indice_pairs_desc.desc(), MLUOP_LAYOUT_ARRAY, dtype, dim, dims));
|
||||
|
||||
// input_grad_desc
|
||||
TORCH_MLUOP_CHECK(mluOpGetTensorDescriptor(input_grad_desc.desc(), &layout, &dtype, &dim,
|
||||
dims));
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(input_grad_desc.desc(), MLUOP_LAYOUT_ARRAY, dtype,
|
||||
dim, dims));
|
||||
TORCH_MLUOP_CHECK(mluOpGetTensorDescriptor(input_grad_desc.desc(), &layout,
|
||||
&dtype, &dim, dims));
|
||||
TORCH_MLUOP_CHECK(mluOpSetTensorDescriptor(
|
||||
input_grad_desc.desc(), MLUOP_LAYOUT_ARRAY, dtype, dim, dims));
|
||||
}
|
||||
|
||||
auto handle = mluOpGetCurrentHandle();
|
||||
|
|
|
@ -30,8 +30,8 @@ void ThreeNNMLUKernelLauncher(int b, int n, int m, const Tensor unknown,
|
|||
|
||||
auto handle = mluOpGetCurrentHandle();
|
||||
size_t workspace_size = 0;
|
||||
TORCH_MLUOP_CHECK(mluOpGetThreeNNForwardWorkspaceSize(handle, known_desc.desc(),
|
||||
&workspace_size));
|
||||
TORCH_MLUOP_CHECK(mluOpGetThreeNNForwardWorkspaceSize(
|
||||
handle, known_desc.desc(), &workspace_size));
|
||||
auto known_workspace =
|
||||
at::empty(workspace_size, known.options().dtype(at::kByte));
|
||||
|
||||
|
@ -46,10 +46,10 @@ void ThreeNNMLUKernelLauncher(int b, int n, int m, const Tensor unknown,
|
|||
auto idx_ptr = idx_impl->cnnlMalloc();
|
||||
auto workspace_ptr = workspace_impl->cnnlMalloc();
|
||||
|
||||
TORCH_MLUOP_CHECK(mluOpThreeNNForward(handle, unknown_desc.desc(), unknown_ptr,
|
||||
known_desc.desc(), known_ptr, workspace_ptr,
|
||||
workspace_size, dist2_desc.desc(), dist2_ptr,
|
||||
idx_desc.desc(), idx_ptr));
|
||||
TORCH_MLUOP_CHECK(mluOpThreeNNForward(
|
||||
handle, unknown_desc.desc(), unknown_ptr, known_desc.desc(), known_ptr,
|
||||
workspace_ptr, workspace_size, dist2_desc.desc(), dist2_ptr,
|
||||
idx_desc.desc(), idx_ptr));
|
||||
}
|
||||
|
||||
void three_nn_forward_mlu(int b, int n, int m, const Tensor unknown,
|
||||
|
|
|
@ -62,14 +62,14 @@ int HardVoxelizeForwardMLUKernelLauncher(
|
|||
at::empty(workspace_size, points.options().dtype(at::kByte));
|
||||
INITIAL_MLU_PARAM_WITH_TENSOR(workspace_tensor);
|
||||
|
||||
TORCH_MLUOP_CHECK(mluOpVoxelization(handle, points_desc.desc(), points_ptr,
|
||||
voxel_size_tensor_desc.desc(), voxel_size_tensor_ptr,
|
||||
coors_range_tensor_desc.desc(), coors_range_tensor_ptr,
|
||||
max_points, max_voxels, NDim, true, workspace_tensor_ptr,
|
||||
workspace_size, voxels_desc.desc(), voxels_ptr,
|
||||
coors_desc.desc(), coors_ptr,
|
||||
num_points_per_voxel_desc.desc(), num_points_per_voxel_ptr,
|
||||
voxel_num_tensor_desc.desc(), voxel_num_tensor_ptr));
|
||||
TORCH_MLUOP_CHECK(mluOpVoxelization(
|
||||
handle, points_desc.desc(), points_ptr, voxel_size_tensor_desc.desc(),
|
||||
voxel_size_tensor_ptr, coors_range_tensor_desc.desc(),
|
||||
coors_range_tensor_ptr, max_points, max_voxels, NDim, true,
|
||||
workspace_tensor_ptr, workspace_size, voxels_desc.desc(), voxels_ptr,
|
||||
coors_desc.desc(), coors_ptr, num_points_per_voxel_desc.desc(),
|
||||
num_points_per_voxel_ptr, voxel_num_tensor_desc.desc(),
|
||||
voxel_num_tensor_ptr));
|
||||
auto voxel_num_cpu = voxel_num_tensor.to(at::kCPU);
|
||||
int voxel_num_int = voxel_num_cpu.data_ptr<int>()[0];
|
||||
return voxel_num_int;
|
||||
|
|
Loading…
Reference in New Issue