Cannot define an MmaOp where batch dimension is Broadcast

[jacobhinkle]

The following test fails when trying to create an MmaOp

// Single batch dimension which is broadcast
TEST_F(GPUTTensorCoreTest, FusionAmpereBroadcastBatchMatmul_CUDA) {
  auto layout = MmaLayout::TN;

  Fusion fusion;
  FusionGuard fg(&fusion);

  auto shapes = matmulAtInputShape3DTuring(-1, -1, -1, layout);

  auto tv0 = makeContigConcreteTensor(shapes.first, DataType::Half);
  auto tv1 = makeContigConcreteTensor(shapes.second, DataType::Half);

  fusion.addInput(tv0);
  fusion.addInput(tv1);

  tv0 = canonicalizeInputToBMNK(tv0, layout, MmaOperand::A);
  tv1 = canonicalizeInputToBMNK(tv1, layout, MmaOperand::B);
  auto tv2 = fusedMultiplySum(
      broadcast(tv0, {true, false, false, false}),
      broadcast(tv1, {true, false, false, false}),
      {-1});
/*
C++ exception with description "details.bcasts.empty() INTERNAL ASSERT FAILED at "/opt/pytorch/nvfuser/csrc/ir/utils.cpp":1
268, please report a bug with repro script to NVFuser at https://github.com/NVIDIA/Fuser/issues. MmaOp output: has broadcast domains.                                                                                                                 
Exception raised from operator() at /opt/pytorch/nvfuser/csrc/ir/utils.cpp:1268 (most recent call first):
*/

  fusion.addOutput(tv2);
}

This caused the failure of

Fuser/tests/cpp/test_combine_mul_sum.cpp

Line 121 in 5228f89

TEST_F(CombineMulSumAsMmaTest, AmpereMulSumToMatmul_Fail2) {

which is why that test currently checks that we cannot translate that case. However, I think that case should be covered and we should instead fix the MmaOp ctor to not balk at such cases.

[jacobhinkle]

Relevant comment from the PR introducing this check: https://github.com/NVIDIA/Fuser/pull/131/files#r1164511926. It seems we do plan to support this broadcast batch dims, but getMmaOpDetails and getInputLayout don't currently support it. This might change with #2272 since that uses IdModel and allocation domain to determine layout instead of pattern matching.

[kevinstephano]

Is this still an issue?

[jacobhinkle]

This does still fail. Now it fails in getInputLayout with Unsupported input layout. However, thanks for bringing this to my attention. I just commented out this and the test succeeds, and we can schedule the fusion.

Fuser/csrc/ir/utils.cpp

Lines 1410 to 1415 in 621e146

	details.input_layout = getInputLayout(
	in_a_details,
	in_b_details,
	details.m_axes,
	details.n_axes,
	details.k_axes);

details.input_layout is not used, so we can remove it. I will do that in a new PR that will close this issue.