Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
OpenDAS
MMCV
Commits
c5233598
Unverified
Commit
c5233598
authored
Aug 07, 2023
by
Chris Jiang
Committed by
GitHub
Aug 07, 2023
Browse files
[Fix] Fix MLU code format (#2887)
parent
9036241e
Changes
19
Hide whitespace changes
Inline
Side-by-side
Showing
19 changed files
with
178 additions
and
159 deletions
+178
-159
mmcv/ops/csrc/pytorch/mlu/ball_query_mlu.cpp
mmcv/ops/csrc/pytorch/mlu/ball_query_mlu.cpp
+3
-3
mmcv/ops/csrc/pytorch/mlu/box_iou_rotated.cpp
mmcv/ops/csrc/pytorch/mlu/box_iou_rotated.cpp
+3
-3
mmcv/ops/csrc/pytorch/mlu/carafe_mlu.cpp
mmcv/ops/csrc/pytorch/mlu/carafe_mlu.cpp
+12
-11
mmcv/ops/csrc/pytorch/mlu/deform_roi_pool_mlu.cpp
mmcv/ops/csrc/pytorch/mlu/deform_roi_pool_mlu.cpp
+8
-8
mmcv/ops/csrc/pytorch/mlu/diff_iou_rotated_mlu.cpp
mmcv/ops/csrc/pytorch/mlu/diff_iou_rotated_mlu.cpp
+2
-2
mmcv/ops/csrc/pytorch/mlu/iou3d_mlu.cpp
mmcv/ops/csrc/pytorch/mlu/iou3d_mlu.cpp
+9
-8
mmcv/ops/csrc/pytorch/mlu/mlu_common_helper.cpp
mmcv/ops/csrc/pytorch/mlu/mlu_common_helper.cpp
+2
-1
mmcv/ops/csrc/pytorch/mlu/mlu_common_helper.h
mmcv/ops/csrc/pytorch/mlu/mlu_common_helper.h
+16
-10
mmcv/ops/csrc/pytorch/mlu/ms_deform_attn_mlu.cpp
mmcv/ops/csrc/pytorch/mlu/ms_deform_attn_mlu.cpp
+5
-5
mmcv/ops/csrc/pytorch/mlu/nms_mlu.cpp
mmcv/ops/csrc/pytorch/mlu/nms_mlu.cpp
+10
-9
mmcv/ops/csrc/pytorch/mlu/nms_rotated_mlu.cpp
mmcv/ops/csrc/pytorch/mlu/nms_rotated_mlu.cpp
+6
-4
mmcv/ops/csrc/pytorch/mlu/psamask_mlu.cpp
mmcv/ops/csrc/pytorch/mlu/psamask_mlu.cpp
+4
-3
mmcv/ops/csrc/pytorch/mlu/roi_align_mlu.cpp
mmcv/ops/csrc/pytorch/mlu/roi_align_mlu.cpp
+19
-19
mmcv/ops/csrc/pytorch/mlu/roi_align_rotated_mlu.cpp
mmcv/ops/csrc/pytorch/mlu/roi_align_rotated_mlu.cpp
+6
-6
mmcv/ops/csrc/pytorch/mlu/roiaware_pool3d_mlu.cpp
mmcv/ops/csrc/pytorch/mlu/roiaware_pool3d_mlu.cpp
+2
-2
mmcv/ops/csrc/pytorch/mlu/rotated_feature_align_mlu.cpp
mmcv/ops/csrc/pytorch/mlu/rotated_feature_align_mlu.cpp
+4
-4
mmcv/ops/csrc/pytorch/mlu/sparse_conv_mlu.cpp
mmcv/ops/csrc/pytorch/mlu/sparse_conv_mlu.cpp
+53
-47
mmcv/ops/csrc/pytorch/mlu/three_nn_mlu.cpp
mmcv/ops/csrc/pytorch/mlu/three_nn_mlu.cpp
+6
-6
mmcv/ops/csrc/pytorch/mlu/voxelization_mlu.cpp
mmcv/ops/csrc/pytorch/mlu/voxelization_mlu.cpp
+8
-8
No files found.
mmcv/ops/csrc/pytorch/mlu/ball_query_mlu.cpp
View file @
c5233598
...
...
@@ -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
,
...
...
mmcv/ops/csrc/pytorch/mlu/box_iou_rotated.cpp
View file @
c5233598
...
...
@@ -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
,
boxes
2
_desc
.
desc
(),
boxes
2
_ptr
,
ious_desc
.
desc
(),
ious_ptr
));
TORCH_MLUOP_CHECK
(
mluOpBoxIouRotated
(
handle
,
mode_flag
,
aligned
,
boxes
1
_desc
.
desc
(),
boxes
1
_ptr
,
boxes2_desc
.
desc
(),
boxes2_ptr
,
ious_desc
.
desc
(),
ious_ptr
));
}
void
box_iou_rotated_mlu
(
const
Tensor
boxes1
,
const
Tensor
boxes2
,
Tensor
ious
,
...
...
mmcv/ops/csrc/pytorch/mlu/carafe_mlu.cpp
View file @
c5233598
...
...
@@ -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
));
...
...
mmcv/ops/csrc/pytorch/mlu/deform_roi_pool_mlu.cpp
View file @
c5233598
...
...
@@ -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_
);
}
...
...
mmcv/ops/csrc/pytorch/mlu/diff_iou_rotated_mlu.cpp
View file @
c5233598
...
...
@@ -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
;
}
...
...
mmcv/ops/csrc/pytorch/mlu/iou3d_mlu.cpp
View file @
c5233598
...
...
@@ -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
));
}
...
...
mmcv/ops/csrc/pytorch/mlu/mlu_common_helper.cpp
View file @
c5233598
...
...
@@ -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
...
...
mmcv/ops/csrc/pytorch/mlu/mlu_common_helper.h
View file @
c5233598
...
...
@@ -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
;
};
...
...
mmcv/ops/csrc/pytorch/mlu/ms_deform_attn_mlu.cpp
View file @
c5233598
...
...
@@ -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
;
...
...
mmcv/ops/csrc/pytorch/mlu/nms_mlu.cpp
View file @
c5233598
...
...
@@ -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
));
...
...
mmcv/ops/csrc/pytorch/mlu/nms_rotated_mlu.cpp
View file @
c5233598
...
...
@@ -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
);
...
...
mmcv/ops/csrc/pytorch/mlu/psamask_mlu.cpp
View file @
c5233598
...
...
@@ -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
);
...
...
mmcv/ops/csrc/pytorch/mlu/roi_align_mlu.cpp
View file @
c5233598
...
...
@@ -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
,
roi
s_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
,
grad
s_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_
);
}
...
...
mmcv/ops/csrc/pytorch/mlu/roi_align_rotated_mlu.cpp
View file @
c5233598
...
...
@@ -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_
);
}
...
...
mmcv/ops/csrc/pytorch/mlu/roiaware_pool3d_mlu.cpp
View file @
c5233598
...
...
@@ -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
);
...
...
mmcv/ops/csrc/pytorch/mlu/rotated_feature_align_mlu.cpp
View file @
c5233598
...
...
@@ -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_
);
}
...
...
mmcv/ops/csrc/pytorch/mlu/sparse_conv_mlu.cpp
View file @
c5233598
...
...
@@ -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
();
...
...
mmcv/ops/csrc/pytorch/mlu/three_nn_mlu.cpp
View file @
c5233598
...
...
@@ -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
,
un
known_desc
.
desc
(),
un
known_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
,
...
...
mmcv/ops/csrc/pytorch/mlu/voxelization_mlu.cpp
View file @
c5233598
...
...
@@ -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
;
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment