Fix broadcast InferCorrectLayout

[lazycal]

Thanks for contributing to TVM! Please refer to guideline https://tvm.apache.org/docs/contribute/ for useful information and tips. After the pull request is submitted, please request code reviews from Reviewers by @ them in the pull request thread.

I'm reviewing this function after PR #10118, and found several buggy places. For instance,

• it does not respect the original input tensor's layout but always reset it to the other input's layout (see test_broadcast_respect_input_layouts);
• it assumes one can always only transform one layout to adapt to the other. However, this is not always possible, see test_broadcast_non_adaptable.
• AdjustSubordinateFactors should also be called in the else branch

  tvm/src/relay/transforms/infer_layout_utils.h

  Lines 222 to 234 in 6a274af

  	AdjustSubordinateFactors(layouts[large_idx], old_small_layout, old_small_shape);
  	layouts.Set(small_idx, new_small_layout);
  	} else {
  	// Support scenarios where original operands were of type [N, H, W, C] and [C]. In this case,
  	// while transforming the layout, we expand dims to make C go to NHWC, and then use the
  	// modified layout of the first operator to call the layout transform. E.g.
  	// a in NCHWC16c
  	// b in C
  	// b = expand_dims(b) from C -> NHWC
  	// b = layout_transform(b) from NHWC -> NCHW16c
  	// add(a, b)
  	layouts.Set(small_idx, ret);
  	}

  . (see test_alter_layout_nonscalar_broadcast)

This PR rewrites BinaryBroadcastLayoutHelper a bit, and I tried to make it as backward compatible as possible.

Misc

BTW I seem to spot another bug when I was working on this PR. So I sometimes I need to create a scalar layout, and I used Layout(""), but it returned an undefined layout. From the doc it should return a scalar layout right? I checked the constructor, and it seems in this if statement:

tvm/src/tir/ir/data_layout.cc

Line 99 in 6a274af

if (name.empty()) return; // scalar

, it directly returns without setting the data_ field as in normal path:

tvm/src/tir/ir/data_layout.cc

Line 144 in 6a274af

data_ = std::move(node);

. I didn't fix this issue but work it around somehow, but it does seem like a bug to me :-).

[lazycal]

@masahi

[lazycal]

@masahi The code is changed according to your suggestions.

[masahi]

I didn't actually try to understand the code, but why not, ship it