Skip to content

Commit 9d415ba

Browse files
Ritwik Dasfacebook-github-bot
authored andcommitted
Add roi_align_box_processor DSP op
Summary: Add the necessary infra code to get the roi_align_box_processor DSP op implemented. This op is responsible for converting the roi_align box coordinates from {x1, y1, x2, y2} format to the 80 bytes encoded Turing format (described here: https://docs.google.com/document/d/1IwPvQeRq34rXg-v2qDYO_91PTZPUg1pO4ZaNN1GcjDQ/edit?fbclid=IwY2xjawLCeilleHRuA2FlbQIxMQABHkwVHvvUSalyry20nNFMMU7KbvpExm6i92gPvV_jO1gNI71Zn8WNTp2Rn7kf_aem_nbrzf940O3zNRhet8DJvHg&tab=t.0#heading=h.iiq3n6kursay). This new op will run on the DSP and will execute as a predecessor of the turing::roi_align op which will now work with the Turing native box coordinates. Reviewed By: mcremon-meta, zonglinpeng, bingcy Differential Revision: D78197430
1 parent 8f062d3 commit 9d415ba

File tree

1 file changed

+19
-0
lines changed

1 file changed

+19
-0
lines changed

backends/cadence/aot/ops_registrations.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -276,6 +276,14 @@
276276
"requantize.per_tensor_out(Tensor input, float in_scale, int in_zero_point, float out_scale, "
277277
"int out_zero_point, ScalarType out_dtype, *, Tensor(a!) out) -> Tensor(a!)"
278278
)
279+
lib.define(
280+
"roi_align_box_processor.out(Tensor rois, int output_size_h, int output_size_w, "
281+
"int sampling_ratio, bool aligned, *, Tensor(a!) out) -> Tensor(a!)"
282+
)
283+
lib.define(
284+
"roi_align_box_processor(Tensor rois, int output_size_h, int output_size_w, "
285+
"int sampling_ratio, bool aligned) -> (Tensor out)"
286+
)
279287

280288
# Custom ops with aten namespace. Need to specify the lib var as FRAGMENT type as aten library is already defined
281289
aten_lib = Library("aten", "FRAGMENT")
@@ -1038,3 +1046,14 @@ def idma_store_impl(
10381046
channel: int = 0,
10391047
) -> torch.Tensor:
10401048
return copy_idma_copy_impl(src, task_num, channel)
1049+
1050+
1051+
@register_fake("cadence::roi_align_box_processor")
1052+
def roi_align_box_processor_meta(
1053+
rois: torch.Tensor,
1054+
output_size_h: int,
1055+
output_size_w: int,
1056+
sampling_ratio: int,
1057+
aligned: bool
1058+
) -> torch.Tensor:
1059+
return rois.new_empty((rois.shape[0], 80), dtype=torch.uint8)

0 commit comments

Comments
 (0)