[inductor][cpu] maml_omniglot AMP Dynamic shape default wrapper accuracy and performance crashed

[zxd1997066]

🐛 Describe the bug

maml_omniglot AMP Dynamic shape default wrapper accuracy and performance crashed

suite	name	thread	accuracy	perf	reason(reference only)
torchbench	maml_omniglot	multiple	X	N/A	maml_omniglot, AssertionError: expected size 64==64 stride 1==676 at dim=1
torchbench	maml_omniglot	single	X	N/A	maml_omniglot, AssertionError: expected size 64==64 stride 1==676 at dim=1

loading model: 0it [00:00, ?it/s]cpu  eval  maml_omniglot                      

skipping cudagraphs due to skipping cudagraphs due to cpu device. Found from : 
   File "benchmarks/dynamo/torchbench.py", line 357, in forward_pass
    return mod(*inputs)

ERROR:common:Backend dynamo failed in warmup()
Traceback (most recent call last):
  File "/workspace/pytorch/benchmarks/dynamo/common.py", line 2639, in warmup
    fn(model, example_inputs)
  File "/workspace/pytorch/torch/_dynamo/eval_frame.py", line 390, in _fn
    return fn(*args, **kwargs)
  File "benchmarks/dynamo/torchbench.py", line 355, in forward_pass
    def forward_pass(self, mod, inputs, collect_outputs=True):
  File "/workspace/pytorch/torch/_dynamo/eval_frame.py", line 390, in _fn
    return fn(*args, **kwargs)
  File "/workspace/pytorch/torch/_dynamo/external_utils.py", line 36, in inner
    return fn(*args, **kwargs)
  File "/workspace/pytorch/torch/_functorch/aot_autograd.py", line 917, in forward
    return compiled_fn(full_args)
  File "/workspace/pytorch/torch/_functorch/_aot_autograd/utils.py", line 89, in g
    return f(*args)
  File "/workspace/pytorch/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 106, in runtime_wrapper
    all_outs = call_func_at_runtime_with_args(
  File "/workspace/pytorch/torch/_functorch/_aot_autograd/utils.py", line 113, in call_func_at_runtime_with_args
    out = normalize_as_list(f(args))
  File "/workspace/pytorch/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py", line 152, in rng_functionalization_wrapper
    return compiled_fw(args)
  File "/workspace/pytorch/torch/_inductor/compile_fx.py", line 1093, in wrapper
    return optimized_function(args_new)
  File "/workspace/pytorch/torch/_inductor/codecache.py", line 937, in __call__
    return self.current_callable(inputs)
  File "/tmp/tmp2dbylmum/b6/cb6xejcvhq24bivlii5l4bnrokyshg25eaegfundva3emkz5hllr.py", line 251, in call
    assert_size_stride(buf1, (s0, 64, 26, 26), (43264, 676, 26, 1))
AssertionError: expected size 64==64, stride 1==676 at dim=1
Run failed with return code:  255
Output:  None
Error:  None

Versions

SW info

name	target_branch	target_commit	refer_branch	refer_commit
torchbench	main	d6015d42	main	1ef0a39e
torch	main	5bc7f7f	main	41286f1
torchvision	main	0.18.0a0+0325175	main	0.18.0a0+2c127da
torchtext	main	0.16.0a0+b0ebddc	main	0.16.0a0+b0ebddc
torchaudio	main	2.2.0a0+87aeb55	main	2.2.0a0+87aeb55
torchdata	main	0.7.1a0+0790338	main	0.7.1a0+0790338
dynamo_benchmarks	main	nightly	main	nightly

Repro:
inductor_single_run.sh
bash inductor_single_run.sh thread inference accuracy/performance torchbench maml_omniglot AMP first dynamic default 0

Suspected guilty commit: f2f8eee
torchbench-maml_omniglot-inference-amp-dynamic-default-single-accuracy-crash_guilty_commit.log
cc @ezyang @msaroufim @bdhirsh @anijain2305 @zou3519 @chauhang @WeizhuoZhang-intel @chuanqi129

[chunyuan-w]

In brief:

For this model, the first Conv kernel will choose channels first for FP32 mode but channels last for AMP mode when the input to Conv has iC = 1 and contiguous format. Need further investigation on the fix.

Details:

In FP32 case, x is a FixedLayout with size=[s0, 1, 28, 28], stride=[784, 784, 28, 1], when calling x = cls.require_stride_order(x, req_stride_order), it runs into this x.get_layout().is_stride_ordered(order) and is_stride_ordered will ignore the dimensions of size 1 so it considered that the layout of x is already as expected. After calling x = cls.require_stride_order(x, req_stride_order), the size and stride keeps unchanged.

In AMP mode, x becomes a FlexibleLayout with size=[s0, 1, 28, 28], stride=[784, 784, 28, 1], since x is the output of a generated kernel which converts it from FP32 to BF16. As a FlexibleLayout, when calling x = cls.require_stride_order(x, req_stride_order), it runs into as_storage_and_layout, which will always freeze_layout_with_stride_order and will always call as_stride_order, without checking if the stride order is already as expected (link). Thus after calling x = cls.require_stride_order(x, req_stride_order), x becomes: size=[s0, 1, 28, 28], stride=[784, 1, 28, 1]).

The mkldnn_conv_use_channels_last check in the Conv kernel returns false for size=[s0, 1, 28, 28], stride=[784, 784, 28, 1] but true for size=[s0, 1, 28, 28], stride=[784, 1, 28, 1]).

[zxd1997066]

fp32 multiple thread dynaimc shape cpp wrapper performance regression has the same guilty commit.

suite	name	thread	batch_size_new	speed_up_new	inductor_new	eager_new	compilation_latency_new	batch_size_old	speed_up_old	inductor_old	eager_old	compilation_latency_old	Ratio Speedup(New/old)	Eager Ratio(old/new)	Inductor Ratio(old/new)	Compilation_latency_Ratio(old/new)
torchbench	maml_omniglot	multiple	5.0	1.172697	0.000494197	0.0005795433393090001	15.577837	5.0	1.372539	0.000428888	0.000588665506632	15.412311	0.85	1.02	0.87	0.99

torchbench-maml_omniglot-inference-float32-dynamic-cpp-multiple-performance-drop_guilty_commit.log

bash inductor_single_run.sh multiple inference performance torchbench maml_omniglot float32 first dynamic cpp 0

[chunyuan-w]

fp32 multiple thread dynaimc shape cpp wrapper performance regression has the same guilty commit.

suite name thread batch_size_new speed_up_new inductor_new eager_new compilation_latency_new batch_size_old speed_up_old inductor_old eager_old compilation_latency_old Ratio Speedup(New/old) Eager Ratio(old/new) Inductor Ratio(old/new) Compilation_latency_Ratio(old/new)
torchbench maml_omniglot multiple 5.0 1.172697 0.000494197 0.0005795433393090001 15.577837 5.0 1.372539 0.000428888 0.000588665506632 15.412311 0.85 1.02 0.87 0.99
torchbench-maml_omniglot-inference-float32-dynamic-cpp-multiple-performance-drop_guilty_commit.log

bash inductor_single_run.sh multiple inference performance torchbench maml_omniglot float32 first dynamic cpp 0

This is a fix for the accuracy of fp32 multiple thread dynaimc shape cpp wrapper, which will bring performance drop and is as expected. The previous version has accuracy problem.