diff --git a/monai/networks/blocks/dynunet_block.py b/monai/networks/blocks/dynunet_block.py index 801b49de8b..9f2a729ee4 100644 --- a/monai/networks/blocks/dynunet_block.py +++ b/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): diff --git a/monai/networks/nets/dynunet.py b/monai/networks/nets/dynunet.py index d130b886a7..1c835e218f 100644 --- a/monai/networks/nets/dynunet.py +++ b/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) diff --git a/tests/networks/blocks/test_dynunet_block.py b/tests/networks/blocks/test_dynunet_block.py index 4cf68912be..1eb75077be 100644 --- a/tests/networks/blocks/test_dynunet_block.py +++ b/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() diff --git a/tests/networks/nets/test_dynunet.py b/tests/networks/nets/test_dynunet.py index c2c9369923..9da6ca6525 100644 --- a/tests/networks/nets/test_dynunet.py +++ b/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()