DynUNet: opt-in, ROCm-only GEMM decomposition for decoder ConvTranspose3d upsamples

monai/networks/blocks/dynunet_block.py

@@ -209,9 +209,13 @@ def __init__(
         act_name: tuple | str = ("leakyrelu", {"inplace": True, "negative_slope": 0.01}),
         dropout: tuple | str | float | None = None,
         trans_bias: bool = False,
+        use_gemm_transpose: bool = False,
     ):
         super().__init__()
         upsample_stride = upsample_kernel_size
+        # AMD MI300X: when kernel_size == stride, ConvTranspose3d admits an exact
+        # pixel-shuffle GEMM decomposition which is more efficient 
+        self._use_gemm_transpose = bool(use_gemm_transpose) and torch.version.hip is not None
         self.transp_conv = get_conv_layer(
             spatial_dims,
             in_channels,
@@ -238,11 +242,48 @@ def __init__(
 
     def forward(self, inp, skip):
         # number of channels for skip should equals to out_channels
-        out = self.transp_conv(inp)
+        if self._use_gemm_transpose:
+            out = self._transp_conv_gemm(inp)
+        else:
+            out = self.transp_conv(inp)
         out = torch.cat((out, skip), dim=1)
         out = self.conv_block(out)
         return out
 
+    def _transp_conv_gemm(self, x):
+        """Pixel-shuffle GEMM decomposition of ConvTranspose3d, valid only when
+        kernel_size == stride (zero output-window overlap). Falls through to the
+        stock transposed convolution for any shape that violates the decomposition
+        preconditions (k != s, dilation, output_padding, groups, non-5D input)."""
+        conv = self.transp_conv.conv
+        k = tuple(conv.kernel_size)
+        s = tuple(conv.stride)
+        if (
+            hasattr(self.transp_conv, "adn")  # extra act/norm/dropout layer — decomp covers conv only
+            or k != s
+            or any(d != 1 for d in conv.dilation)
+            or any(p != 0 for p in conv.output_padding)
+            or conv.groups != 1
+            or x.dim() != 5
+        ):
+            return self.transp_conv(x)
+
+        n, ic, d, h, w = x.shape
+        oc = conv.out_channels
+        kd, kh, kw = int(k[0]), int(k[1]), int(k[2])
+
+        x_flat = x.contiguous().permute(0, 2, 3, 4, 1).reshape(n * d * h * w, ic)
+        w_flat = conv.weight.reshape(ic, oc * kd * kh * kw)  # weight: (IC, OC, kD, kH, kW)
+        out_flat = x_flat @ w_flat
+
+        out = out_flat.reshape(n, d, h, w, oc, kd, kh, kw)
+        out = out.permute(0, 4, 1, 5, 2, 6, 3, 7).contiguous()
+        out = out.reshape(n, oc, d * kd, h * kh, w * kw)
+
+        if conv.bias is not None:
+            out = out + conv.bias.view(1, oc, 1, 1, 1)
+        return out
+
 
 class UnetOutBlock(nn.Module):
 

monai/networks/nets/dynunet.py

@@ -125,6 +125,9 @@ class DynUNet(nn.Module):
         res_block: whether to use residual connection based convolution blocks during the network.
             Defaults to ``False``.
         trans_bias: whether to set the bias parameter in transposed convolution layers. Defaults to ``False``.
+        use_gemm_transpose: AMD MI300X (ROCm) only. Replace the decoder ConvTranspose3d upsamples with
+            an exact pixel-shuffle GEMM decomposition when ``kernel_size == stride``,
+            Defaults to ``False``.
     """
 
     def __init__(
@@ -143,6 +146,7 @@ def __init__(
         deep_supr_num: int = 1,
         res_block: bool = False,
         trans_bias: bool = False,
+        use_gemm_transpose: bool = False,
     ):
         super().__init__()
         self.spatial_dims = spatial_dims
@@ -156,6 +160,7 @@ def __init__(
         self.dropout = dropout
         self.conv_block = UnetResBlock if res_block else UnetBasicBlock
         self.trans_bias = trans_bias
+        self.use_gemm_transpose = use_gemm_transpose
         if filters is not None:
             self.filters = filters
             self.check_filters()
@@ -319,6 +324,7 @@ def get_upsamples(self):
             UnetUpBlock,  # type: ignore
             upsample_kernel_size,
             trans_bias=self.trans_bias,
+            use_gemm_transpose=self.use_gemm_transpose,
         )
 
     def get_module_list(
@@ -330,6 +336,7 @@ def get_module_list(
         conv_block: nn.Module,
         upsample_kernel_size: Sequence[Sequence[int] | int] | None = None,
         trans_bias: bool = False,
+        use_gemm_transpose: bool = False,
     ):
         layers = []
         if upsample_kernel_size is not None:
@@ -347,6 +354,7 @@ def get_module_list(
                     "dropout": self.dropout,
                     "upsample_kernel_size": up_kernel,
                     "trans_bias": trans_bias,
+                    "use_gemm_transpose": use_gemm_transpose,
                 }
                 layer = conv_block(**params)
                 layers.append(layer)

tests/networks/blocks/test_dynunet_block.py

@@ -18,7 +18,7 @@
 
 from monai.networks import eval_mode
 from monai.networks.blocks.dynunet_block import UnetBasicBlock, UnetResBlock, UnetUpBlock, get_padding
-from tests.test_utils import dict_product, test_script_save
+from tests.test_utils import assert_allclose, dict_product, test_script_save
 
 TEST_CASE_RES_BASIC_BLOCK = []
 for params in dict_product(
@@ -109,5 +109,59 @@ def test_script(self):
         test_script_save(net, test_data, skip_data)
 
 
+class TestUpBlockGemmTranspose(unittest.TestCase):
+    """AMD MI300X: the opt-in pixel-shuffle GEMM decomposition of the decoder
+    ConvTranspose3d (kernel_size == stride) must be numerically identical to the
+    stock transposed convolution it replaces."""
+
+    def test_gemm_decomposition_equivalence(self):
+        # exercise the decomposition math directly so the check is meaningful on
+        # any platform (the runtime gate is ROCm-only, but the math is not).
+        net = UnetUpBlock(
+            spatial_dims=3, in_channels=4, out_channels=2, kernel_size=3, stride=2, norm_name="instance", upsample_kernel_size=2
+        )
+        x = torch.randn(1, 4, 5, 6, 7)
+        with eval_mode(net):
+            expected = net.transp_conv(x)
+            result = net._transp_conv_gemm(x)
+        self.assertEqual(result.shape, expected.shape)
+        assert_allclose(result, expected, atol=1e-4, rtol=1e-4)
+
+    def test_gemm_decomposition_falls_through_when_k_ne_s(self):
+        # kernel_size != stride violates the zero-overlap precondition; the
+        # decomposition must fall back to the stock transposed convolution.
+        net = UnetUpBlock(
+            spatial_dims=3, in_channels=4, out_channels=2, kernel_size=3, stride=2, norm_name="instance", upsample_kernel_size=3
+        )
+        x = torch.randn(1, 4, 5, 6, 7)
+        with eval_mode(net):
+            expected = net.transp_conv(x)
+            result = net._transp_conv_gemm(x)
+        assert_allclose(result, expected, atol=1e-4, rtol=1e-4)
+
+    def test_forward_equivalence(self):
+        # two blocks sharing weights, GEMM path on vs off, must agree end-to-end.
+        params = dict(
+            spatial_dims=3, in_channels=4, out_channels=2, kernel_size=3, stride=2, norm_name="instance", upsample_kernel_size=2
+        )
+        net_gemm = UnetUpBlock(use_gemm_transpose=True, **params)
+        net_ref = UnetUpBlock(use_gemm_transpose=False, **params)
+        net_gemm.load_state_dict(net_ref.state_dict())
+        inp = torch.randn(1, 4, 4, 4, 4)
+        skip = torch.randn(1, 2, 8, 8, 8)
+        with eval_mode(net_gemm), eval_mode(net_ref):
+            out_gemm = net_gemm(inp, skip)
+            out_ref = net_ref(inp, skip)
+        assert_allclose(out_gemm, out_ref, atol=1e-4, rtol=1e-4)
+
+    def test_gate_requires_rocm(self):
+        # the runtime gate must be off on non-ROCm builds even when opted in.
+        net = UnetUpBlock(
+            spatial_dims=3, in_channels=4, out_channels=2, kernel_size=3, stride=2,
+            norm_name="instance", upsample_kernel_size=2, use_gemm_transpose=True,
+        )
+        self.assertEqual(net._use_gemm_transpose, torch.version.hip is not None)
+
+
 if __name__ == "__main__":
     unittest.main()

tests/networks/nets/test_dynunet.py

@@ -185,5 +185,37 @@ def test_shape(self, input_param, input_shape, expected_shape):
             self.assertEqual(results.shape, expected_shape)
 
 
+class TestDynUNetGemmTranspose(unittest.TestCase):
+    """AMD MI300X: use_gemm_transpose must thread from DynUNet down to every
+    decoder upsample block (this flag is what the ROCm bundle overlay flips)."""
+
+    def test_flag_threaded_to_upsample_blocks(self):
+        net = DynUNet(
+            spatial_dims=3,
+            in_channels=1,
+            out_channels=2,
+            kernel_size=[3, 3, 3],
+            strides=[1, 2, 2],
+            upsample_kernel_size=[2, 2],
+            use_gemm_transpose=True,
+        )
+        self.assertTrue(len(net.upsamples) > 0)
+        expected = torch.version.hip is not None
+        for block in net.upsamples:
+            self.assertEqual(block._use_gemm_transpose, expected)
+
+    def test_flag_defaults_off(self):
+        net = DynUNet(
+            spatial_dims=3,
+            in_channels=1,
+            out_channels=2,
+            kernel_size=[3, 3, 3],
+            strides=[1, 2, 2],
+            upsample_kernel_size=[2, 2],
+        )
+        for block in net.upsamples:
+            self.assertFalse(block._use_gemm_transpose)
+
+
 if __name__ == "__main__":
     unittest.main()