ATen scheduler for the new Matmul/LinearOp IR nodes

[Priya2698]

1. Adds a new scheduler -- ExprEvalScheduler that accepts the MatmulOp and LinearOp (next PR) for ATen evaluation.
2. Modify the matmul input generator to test for all cases supported by Thunder.
3. The eagerMatmul API is renamed and replaces the existing matmul API. fd.ops.matmul now creates a MatmulOp (except in a few special cases such as scalar dot product, for eg: [K] x [K].

Issue #2149, #2092.

[jacobhinkle]

Looking great. Sprinkle in a few tests once you merge #2175 and we'll be on our way.

[Priya2698]

!build

[jacobhinkle]

Looks good overall. I am just slightly concerned about removing stuff from PairwiseRootDomainMap::map.

[jacobhinkle]

Should NoOp come before ExprEval?

[Priya2698]

Some cases get accepted by NoOp scheduler, which is why I prioritized ExprEval scheduler.

We may need to change the heuristics of NoOp if we want to switch the order.

[jacobhinkle]

Oh! Thanks for mentioning that. Does NoOp scheduler accept the cases where you have a single scalar output? Because it seems to me that it would do so based on this code:

Fuser/csrc/fusion.cpp

Lines 341 to 359 in 8c18701

	bool Fusion::isNoOp() {
	if (exprs().empty()) {
	return true;
	}
	for (auto out_tv : ir_utils::filterByType<TensorView>(outputs())) {
	const std::vector<IterDomain*>& root_dom =
	TensorDomain::noReductions(out_tv->getMaybeRFactorDomain());
	const bool size_zero =
	std::any_of(root_dom.begin(), root_dom.end(), [](IterDomain* id) {
	return id->extent()->isConstScalar() && id->extent()->evaluate() == 0;
	});
	if (!size_zero) {
	return false;
	}
	}
	return true;
	}

We should add a special case for zero-dimensional outputs there On second look, it seems like size_zero would be false in the case that root_dom.empty(). However, the code below might not properly handle zero-dimensional outputs:

Fuser/csrc/scheduler/no_op.cpp

Lines 71 to 80 in 8c18701

	// Check that all outputs are either broadcast or ignored reduction.
	for (auto out_tv : ir_utils::filterByType<TensorView>(fusion->outputs())) {
	auto concrete_dimension = TensorDomain::noReductions(
	TensorDomain::noBroadcasts(out_tv->getLeafDomain()));
	if (!concrete_dimension.empty()) {
	scheduler_debug_utils::canScheduleRejectReason(
	heuristicType(), "output has a concrete dimension");
	return false;
	}
	}

[jacobhinkle]

#2241

[Priya2698]

[  FAILED  ] 6 tests, listed below:
[  FAILED  ] ATenNodesParametrizedTest.MatmulNodeConcrete/2, where GetParam() = ({ 32 }, { 32, 1 })
[  FAILED  ] ATenNodesParametrizedTest.MatmulNodeConcrete/8, where GetParam() = ({ 1, 32 }, { 32 })
[  FAILED  ] ATenNodesParametrizedTest.MatmulNodeConcrete/10, where GetParam() = ({ 1, 32 }, { 32, 1 })
[  FAILED  ] ATenNodesParametrizedTest.MatmulNodeSymbolic/2, where GetParam() = ({ 32 }, { 32, 1 })
[  FAILED  ] ATenNodesParametrizedTest.MatmulNodeSymbolic/8, where GetParam() = ({ 1, 32 }, { 32 })
[  FAILED  ] ATenNodesParametrizedTest.MatmulNodeSymbolic/10, where GetParam() = ({ 1, 32 }, { 32, 1 })

It is likely because there no reductions identified since we use ATen, and all the dimensions in the output are broadcast dimensions. So the cases where M/N = 1 get picked by NoOp

[jacobhinkle]

LGTM after tests pass and you add some broadcasts in test.

[Priya2698]

!build

[jjsjann123]

🚀