Aten roi align

py/torch_migraphx/fx/converters/acc_ops_converters.py

@@ -200,7 +200,66 @@ def acc_ops_nll_loss(mgx_module, node, args, kwargs):
         return MGXInstruction(loss), MGXInstruction(weight_sum)
 
     return MGXInstruction(loss)
+
+
+# ROI_Align converter invokes MigraphX op of the same name.  See comment below on the
+# idiosyncratic use of the batch dimension for this op.
+@migraphx_converter(acc_ops.roi_align)
+def acc_ops_roi_align(mgx_module, node, args, kwargs):
+    inp = kwargs['input']
+    inp_instr_ref = inp.instr_ref
+    boxes_ref = kwargs['boxes'].instr_ref
+    output_size = kwargs['output_size']
+
+    transformation_mode = 'half_pixel' if kwargs['aligned'] else 'output_half_pixel'    
+    spatial_scale = kwargs['spatial_scale']
+    sampling_ratio = kwargs['sampling_ratio']
 
+    # "boxes" and "roi" both refer to the same region of interest boxes.  Each box has
+    # associated with it an index that tells which layer (in the batch dimension) it
+    # is to be applied to.  A box is _not_ applied equally to each batch layer, unless
+    # duplicate boxes are explicitly provided.  This makes the term "batch" rather
+    # misleading for this operation.
+    if boxes_ref.shape().lens()[1] == 5:
+        roi_indices = mgx_module.add_literal(
+                torch.tensor([1, 2, 3, 4], dtype=torch.int64).numpy())
+        # split off the 0'th column of boxes
+        boxes2 = mgx_module.add_instruction(
+            migraphx.op('gather', axis=1), [boxes_ref, roi_indices])
+        zero_indices = mgx_module.add_literal(
+                torch.tensor([0,], dtype=torch.int64).numpy())
+        batch_indices_1 = mgx_module.add_instruction(
+            migraphx.op('gather', axis=1), [boxes_ref, zero_indices])
+        batch_indices2 =  mgx_module.add_instruction(
+            migraphx.op('squeeze', axes=0), [batch_indices_1])
+        batch_indices = mgx_module.add_instruction( migraphx.op('convert', target_type=migraphx.shape.type_t.int32_type),
+            [batch_indices2])
+    elif boxes_ref.shape().lens()[1] == 4:
+        # TODO:  Add support for boxes input in form List[Tensor[L, 4]].  We will
+        # need to create a default index list (0, 1, 2...) something like this:
+        # batch_indices=range(boxes_ref.shape().lens()[0])
+        # boxes2 = boxes_ref
+        # This input format is not allowed by aten package's roi_align.
+        raise RuntimeError('List[Tensor[L, 4]] boxes input for roi_align() not currently supported')
+    else:
+        raise RuntimeError('boxes input must be Tensor[K, 5]')
+
+    temp_instr = mgx_module.add_instruction(
+        migraphx.op('roialign', coordinate_transformation_mode=transformation_mode,
+                    output_height=output_size[0],
+                    output_width=output_size[1],
+                    spatial_scale = spatial_scale,
+                    sampling_ratio = sampling_ratio),
+                    [inp_instr_ref, boxes2, batch_indices])
+
+    # It's not clear why this returns the right results in the wrong shape, but
+    # it now must be reshaped:
+    lens = temp_instr.shape().lens()
+    result =  mgx_module.add_instruction(
+        migraphx.op('reshape', dims=[lens[0], lens[1], lens[3], lens[2]]), [temp_instr])
+
+    return MGXInstruction(result)
+
 
 @migraphx_converter(acc_ops.hardtanh)
 @migraphx_converter(acc_ops.clamp)

py/torch_migraphx/fx/converters/aten_ops_converters.py

@@ -383,6 +383,19 @@ def aten_ops_clamp(mgx_module, node, args, _kwargs):
         acc_kwargs["max"] = args[2] if len(args) == 3 else None
     return acc_ops_converters.acc_ops_clamp(mgx_module, node, (), acc_kwargs)
 
+@migraphx_converter(torch.ops.torchvision.roi_align.default)
+def aten_ops_roi_align(mgx_module, node, args, _kwargs):
+    assert len(args) == 7
+    acc_kwargs = {
+	        "input": args[0],
+	        "boxes": args[1],
+	        "spatial_scale": args[2],
+	        "output_size": [args[3], args[4]],
+	        "sampling_ratio": args[5],
+	        "aligned": args[6]
+	    }          
+    return acc_ops_converters.acc_ops_roi_align(mgx_module, node, (), acc_kwargs)
+
 
 @migraphx_converter(torch.ops.aten.nll_loss_forward.default)
 def aten_ops_nll_loss_forward(mgx_module, node, args, _kwargs):

py/torch_migraphx/fx/tracer/acc_tracer/acc_ops.py

@@ -34,6 +34,12 @@
 import torch
 from typing import cast, Iterable, List, Sequence
 
+from sys import modules
+try:
+    import torchvision
+except ImportError:
+    # torchvision is not a mandatory dependency, so may not be present
+    pass
 import torch.nn as nn
 from torch.fx.passes.shape_prop import _extract_tensor_metadata
 from packaging import version
@@ -373,6 +379,28 @@ def clamp(*, input, min=None, max=None):
     return torch.clamp(input=input, min=min, max=max)
 
 
+if 'torchvision' in modules:
+    @register_acc_op_mapping(
+        op_and_target=("call_function", torchvision.ops.roi_align),
+        arg_replacement_tuples=[
+            ("input", "input", False),
+            ("boxes", "boxes", False),
+            ("output_size", "output_size", False),
+            ("spatial_scale", "spatial_scale", this_arg_is_optional),
+            ("sampling_ratio", "sampling_ratio", this_arg_is_optional),
+            ("aligned", "aligned", this_arg_is_optional),
+        ],
+    )
+    @register_acc_op
+    def roi_align(*, input, boxes, output_size,
+                spatial_scale = 1.0,
+                sampling_ratio = -1,
+                aligned = False):
+        return torchvision.ops.roi_align(input=input, boxes=boxes, output_size = output_size, 
+                                        spatial_scale = spatial_scale, 
+                                        sampling_ratio = sampling_ratio, aligned = aligned)
+
+
 @register_acc_op_properties(AccOpProperty.unary)
 @register_acc_op_mapping(op_and_target=("call_method", "tile"))
 @register_acc_op_mapping(op_and_target=("call_function", torch.tile))

tests/dynamo/converters/test_pooling_dynamo.py

@@ -1,12 +1,10 @@
 import pytest
 import torch
-from dynamo_test_utils import FuncModule, FuncModuleFirstOut, convert_to_mgx, verify_outputs
+from dynamo_test_utils import FuncModule, MultiInFuncModule, FuncModuleFirstOut, convert_to_mgx, verify_outputs
 import torch_migraphx
 
 if not hasattr(torch_migraphx, "dynamo"):
     pytest.skip(allow_module_level=True)
-
-
 @pytest.mark.parametrize('op_alias', [torch.ops.aten.avg_pool2d.default])
 @pytest.mark.parametrize(
     "kernel_size, stride, padding, ceil_mode, count_include_pad", [

tests/dynamo/converters/test_torchvision_ops_dynamo.py

@@ -0,0 +1,40 @@
+
+import pytest
+import torch
+from dynamo_test_utils import FuncModule, MultiInFuncModule, FuncModuleFirstOut, convert_to_mgx, verify_outputs
+import torch_migraphx
+import sys
+try:
+    import torchvision
+except ImportError:
+    pass
+
+@pytest.mark.skipif('torchvision' not in sys.modules,
+                    reason="requires the torchvision library")
+@pytest.mark.parametrize('op_alias', [torch.ops.torchvision.roi_align.default])
+@pytest.mark.parametrize("spatial_scale, sampling_ratio", [(1.5, 3), (0.5, 2)])
+@pytest.mark.parametrize("aligned", [(True), (False)])
+@pytest.mark.parametrize(
+    "input, boxes, output_size", [((1, 2, 3, 4),  ([[0, 1.1, 1.2,  0.6, 2.6]]), [2, 3]),
+                                  #   boxes as List[Tensor[L, 4]] is not supported
+                                  #############  ((2, 2, 3, 3),  ([(1.1, 1.2,  0.6, 2.6), (1.13, 1.23,  0.63, 2.63)]), [2, 2]),
+                                  ((2, 2, 3, 4),  ([[1, 1.1, 1.2,  0.6, 2.6], [0, 1.13, 1.23,  0.63, 2.63]]), [5, 2]),
+                                  ((4, 5, 256, 256),
+                                   ([[0, 10.6, 11.2,  21.1, 22.6],
+                                     [2, 10.63, 11.23,  21.63, 22.63],
+                                     [3, 10.64, 11.23,  21.67, 22.63],
+                                     [1, 10.65, 11.23,  21.68, 22.63],
+                                     ]),
+                                   [7, 6])
+                                  ]
+    )
+def test_roialign(op_alias, input, boxes, output_size, spatial_scale, sampling_ratio, aligned):
+    assert(input[0] == len(boxes))
+    inp = torch.randn(input).cuda()
+    roi = torch.tensor(boxes).cuda()
+
+    mod = MultiInFuncModule(op_alias, spatial_scale, output_size[0], output_size[1], sampling_ratio, aligned)
+    mgx_mod = convert_to_mgx(mod, [inp, roi])
+    verify_outputs(mod, mgx_mod, [inp, roi])
+
+

tests/fx/converters/test_pooling_fx.py

@@ -1,7 +1,9 @@
 import pytest
 import torch
-from fx_test_utils import convert_to_mgx, verify_outputs
+import torchvision
+import numpy as np
 
+from fx_test_utils import convert_to_mgx, verify_outputs
 
 @pytest.mark.parametrize(
     "kernel_size, stride, padding, dilation, ceil_mode",

tests/fx/converters/test_torchvision_ops_fx.py

@@ -18,3 +18,34 @@ def test_stochastic_depth():
                      training=False).eval()
     mgx_mod = convert_to_mgx(mod, [inp])
     verify_outputs(mod, mgx_mod, inp)
+
+
+@pytest.mark.skipif('torchvision' not in sys.modules,
+                    reason="requires the torchvision library")
+@pytest.mark.parametrize("spatial_scale, sampling_ratio", [(1.5, 3), (0.5, 2)])
+@pytest.mark.parametrize("aligned", [(True), (False)])
+@pytest.mark.parametrize(
+    "input, boxes, output_size", [((1, 2, 3, 4),  ([[0, 1.1, 1.2,  0.6, 2.6]]), [3, 2]),
+                                  # boxes as List[Tensor[L, 4]] is not supported
+                                  #   ((2, 2, 3, 3),  ([(1.1, 1.2,  0.6, 2.6), (1.13, 1.23,  0.63, 2.63)]), [2, 2]),
+                                  ((2, 2, 3, 2),  ([[1, 1.1, 1.2,  0.6, 2.6], [0, 1.13, 1.23,  0.63, 2.63]]), [3, 2]),
+                                  ((4, 2, 256, 256),  
+                                   ([[0, 10.6, 11.2,  21.1, 22.6], 
+                                     [2, 10.63, 11.23,  21.63, 22.63],
+                                     [3, 10.64, 11.23,  21.67, 22.63],
+                                     [1, 10.65, 11.23,  21.68, 22.63],
+                                     ]), 
+                                   [7, 6])]
+    )
+def test_roialign(input, boxes, output_size, spatial_scale, sampling_ratio, aligned):
+    assert(input[0] == len(boxes))
+    inp = torch.randn(input).cuda()
+    roi = torch.tensor(boxes).cuda()
+    outputs = torch.tensor(output_size)
+
+    roi_mod = torchvision.ops.RoIAlign(output_size=output_size, spatial_scale=spatial_scale, sampling_ratio=sampling_ratio, aligned=aligned)
+
+    mgx_mod = convert_to_mgx(roi_mod, [inp, roi, outputs])
+    verify_outputs(roi_mod, mgx_mod, (inp, roi))
+
+

tools/install_migraphx.sh

@@ -21,4 +21,4 @@ cd AMDMIGraphX
 rbuild build -d depend -DBUILD_TESTING=Off -DCMAKE_INSTALL_PREFIX=/opt/rocm/ --cxx=/opt/rocm/llvm/bin/clang++ -DGPU_TARGETS=$GPU_ARCH
 
 cd build
-make install
+make -j88 install