Add scheduling support for SDPAOp

[Priya2698]

Stacked on #2294.

1. Adds the producer-consumer mapping to root domain map.
2. Adds SDPAOp to ExprEvalScheduler.
3. Modifies ExprEvalSched::canSchedule to skip computeAt checks and only use the compile time check since expression evaluator scheduler will only accept segments with a single expression of type MatmulOp / LinearOp / SdpaOp.,

Issue #2278

[Priya2698]

!build

[naoyam]

Is the root map change required for the scheduling support?

[Priya2698]

Is the root map change required for the scheduling support?

Yes. Otherwise, we noticed with the earlier Matmul and LinearOps, that the exactMappedExtentSubstitution presegmentation pass runs into error since the producer and consumers are not trivially aligned. I will verify that the error persists if this change is removed for SDPA as well.

[jacobhinkle]

Is the root map change required for the scheduling support?

I suspect there might be problems with ExpressionEvaluator::propagateBoundValuesThroughExactMaps if we did not map any IterDomains for this node (or MatmulOporLinearOp`), but maybe it would be OK since we always re-use in the input extents directly so no mapping is needed?

[naoyam]

Thanks. Looks good to me. I'll let @jacobhinkle give a stamp.

[Priya2698]

Is the root map change required for the scheduling support?

I suspect there might be problems with ExpressionEvaluator::propagateBoundValuesThroughExactMaps if we did not map any IterDomains for this node (or MatmulOporLinearOp`), but maybe it would be OK since we always re-use in the input extents directly so no mapping is needed?

It would still not work since, the third dimension (L/S) only maps from query to output.

C++ exception with description "known_size == this_size INTERNAL ASSERT FAILED at "/opt/pytorch/nvfuser/csrc/expr_evaluator.cpp":302, please report a bug with repro script to NVFuser at https://github.com/NVIDIA/Fuser/issues. Conflicting sizes: 128, 64
Exception raised from propagateBoundValuesThroughExactMaps at /opt/pytorch/nvfuser/csrc/expr_evaluator.cpp:302 (most recent call first):

[Priya2698]

!build

[jacobhinkle]

Test looks good. My comments are mostly minor

[jacobhinkle]

For idx == 0 and consumer_tv_->sameAs(op->output(2)) || consumer_tv_->sameAs(op->output(3)), these should not map since the extents differ by 1. I would just put a check before this that consumer is the output or logsumexp. Btw you might want to add accessors like op->logSumExp() to make this new condition more readable.

[Priya2698]

Even if I use resize?
I saw an error that there was no mapped iterdomain from producer.

[jacobhinkle]

Oh ok so you are mapping the consumer root then you have an rfactor domain for that consumer which uses a Resize?

[jacobhinkle]

That is right then.

[Priya2698]

!build