Update on "[inductor] weight prepack for _convolution_transpose_point…

…wise" cc VitalyFedyunin jgong5 mingfeima XiaobingSuper sanchitintel ashokei jingxu10 mlazos soumith voznesenskym yanboliang penguinwu anijain2305 EikanWang Guobing-Chen zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 desertfire [ghstack-poisoned]
pytorch · chunyuan-w · Dec 6, 2022 · Dec 6, 2022 · Dec 6, 2022 · Dec 6, 2022
commit 0b90c29fb379c7a23e884e48cf43341985a82090
diff --git a/test/test_mkldnn_fusion.py b/test/test_mkldnn_fusion.py
@@ -363,6 +363,21 @@ def forward(self, x):
                         attr = pointwise_info.attr
                         scalars = pointwise_info.scalars
                         algorithm = pointwise_info.algorithm
+
+                        # Prepack weight
+                        packed_weight = torch._C._nn.mkldnn_reorder_conv_transpose2d_weight(
+                            mod.conv_transpose.weight.to_mkldnn(),
+                            mod.conv_transpose.padding,
+                            mod.conv_transpose.stride,
+                            mod.conv_transpose.dilation,
+                            mod.conv_transpose.groups,
+                            mod.conv_transpose.output_padding,
+                            x.size())
+                        mod.conv_transpose.weight = torch.nn.Parameter(
+                            packed_weight,
+                            requires_grad=mod.conv_transpose.weight.requires_grad,
+                        )
+
                         fused = torch.ops.mkldnn._convolution_transpose_pointwise(
                             x,
                             mod.conv_transpose.weight,