[release/2.0] Fix rnn_packed format check (#1592)

* add UT for lstm weight reorder

* use is_opaque instead of is_rnn_packed

* update ideep commit to include is_opaque API

Author: chunyuan-w

csrc/cpu/aten/WeightPack.cpp

@@ -61,14 +61,14 @@ bool is_packed(const at::Tensor& weight) {
 std::tuple<ideep::tensor, ideep::tensor> CommonLstmWeightDesc::
     get_and_save_lstm_packed_weight() {
   ideep::tensor cached_weight_ih, cached_weight_hh;
-  // Don't pack when the weight is of rnn_packed format
+  // Don't pack when the weight is of opaque format (rnn_packed format).
   // When the weight is of rnn_packed format, if the seq_lens of
   // the input changes, the format of weight also changes.
   // oneDNN does not support reorder from rnn_packed back to public
   // format. LSTM based on BRGEMM kernel (on AVX512 and newest ISAs) will
   // use blocked format for weight of LSTM, which won't change when the
   // input seq_lens changes.
-  if (packed_desc_ih_.is_rnn_packed() || packed_desc_hh_.is_rnn_packed()) {
+  if (packed_desc_ih_.is_opaque() || packed_desc_hh_.is_opaque()) {
     return std::make_tuple(w1_src_, w2_src_);
   }
 

tests/cpu/test_weight_prepack.py

@@ -1295,6 +1295,35 @@ def forward(self, x, h=None):
             os.remove('origin_checkpoint.pth')
             os.remove('ipex_checkpoint.pth')
 
+    def test_lstm_weight_reorder(self):
+        class Lstm(torch.nn.Module):
+            def __init__(self, input_size, hidden_size, num_layers, bidirectional, bias, dropout, batch_first):
+                super(Lstm, self).__init__()
+                self.lstm = torch.nn.LSTM(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers, bidirectional=bidirectional, bias=bias, dropout=dropout, batch_first=batch_first)
+
+            def forward(self, x, h=None):
+                x, h = self.lstm(x, h)
+                return x, h
+
+        test_dtypes = []
+        if core.onednn_has_bf16_support():
+            test_dtypes.append(torch.bfloat16)
+        for dtype in test_dtypes:
+            m = Lstm(2, 3, 1, False, False, 0, False)
+            x = torch.randn(2, 1, 2)
+            x_var = torch.randn(5, 1, 2)
+
+            origin_model = copy.deepcopy(m).eval()
+            ipex_model = ipex.optimize(origin_model, dtype=dtype)
+
+            with torch.cpu.amp.autocast(enabled=True, dtype=dtype):            
+                # run with 2 different shapes to verify weight prepack works when weight format changes
+                y = ipex_model(x)
+                y_var = ipex_model(x_var)
+
+                y_ref = origin_model(x_var)
+                self.assertEqual(y_var, y_ref)     
+
 if __name__ == '__main__':
     torch.manual_seed(2020)
     test = unittest.main()