[inductor] Fix a bmm related bug

[desertfire]

Summary: The bug is exposed with pytorch_struct training run. We need to
be conservative on the optimization for bs=1 bmm, when any operand is not consecutive.

[desertfire]

@jansel I guess another option is to generate a ReinterpretView instead of View when the batch size is 1, but I am not sure if that is safe to do here.

[ngimel]

this doesn't look right, MatrixMultiply should be able to handle transposed inputs

[desertfire]

The problem is because View.create does not return a ReinterpretView because of the checking at

torchdynamo/torchinductor/ir.py

Line 695 in 044e8bd

if is_contiguous_storage_and_layout(x):

, and thus later failed an assertion at

torchdynamo/torchinductor/ir.py

Line 1570 in 044e8bd

assert isinstance(x, (Buffer, ReinterpretView)), x

.

I think making a special case in View.create might work. @jansel , any suggestion?

[jansel]

Hrm, in this case perhaps we should be treating it as contiguous, and ignore size=1 dims in contiguous checks. The bs=1 doesn't actually matter.

[ngimel]

size=1 is already not checked, but one (or both) of the 2d matrices resulting from removing batch dimension might not be contiguous, and thus View.create returns View and not ReinterpretView. I think we can return ReinterpretView in more cases by specifying proper strides in the new_layout

[ngimel]

I think we should extend SqueezeView.create to accept a specific dimension, and use Squeezeview.create instead of View.create below.

[jansel]

I agree with @ngimel here we should be able to map bmm to mm in the transposed case. It might be a bit annoying to do that because not all of the view types work with extern kernels without adding a copy.

Accepting because this fixes a bug, and converting bmm to mm is a minor optimization that it is likely ok to miss.

[ngimel]

Here's what I had in mind #548, it still allows bmm -> mm optimization for discontiguous inputs, and won't cause more copies than bmm itself would have caused.