[LinalgExt] Share FFT rewriting between Torch and StableHLO

[giacs-epic]

Take FFT rewriting code out of StableHLO to LinalgExt conversion pass and put it into a shared utility.
Implement a conversion pass for the Torch to IREE pipeline in which, whenever the input criterion is met (size is power of 2), aten.fft_rfft is converted to linalgext.fft. Decomposition of aten.fft_rfft is disabled for the rewriting to work.

[kuhar]

I wonder if we should have a custom linalg_ext op to represent fft such that all input dialect can first lower to it. cc: @MaheshRavishankar @rsuderman

[kuhar]

This should probably return FailureOr<std::pair<Value, Value>> instead

[kuhar]

Use an llvm helper to check for powers of two: https://github.com/llvm/llvm-project/blob/7d1c661381d36018fd105f4ad4c2d6dc45e7288b/llvm/include/llvm/Support/MathExtras.h#L289-L298

[giacs-epic]

@kuhar I have to force-push the rebased commits due to missing sign-offs. I'm addressing your feedback in the latest one.

[ScottTodd]

I wonder if we should have a custom linalg_ext op to represent fft such that all input dialect can first lower to it. cc: @MaheshRavishankar @rsuderman

Assuming this PR is still active, can the assigned reviewers please comment on this?

BTW @giacs-epic , feel free to ping if a PR sits without much activity for several days.

[giacs-epic]

@ScottTodd @kuhar @MaheshRavishankar @rsuderman
I was away for a while, I just updated the PR hoping to fix the CI failure. Aside from that, this is ready for review.

[qedawkins]

A few nits. I'm not very familiar with FFT so I can't comment on the core logic, but overall makes sense to me.

If another reviewer has comments I'll defer to them, otherwise I'd be happy to approve.

[qedawkins]

Can you add more details about the API here? In particular describe how op, operand, and fftLength are used. An example using stableHLO or Torch would help too.

[giacs-epic]

Sure, adding more details.

[ScottTodd]

Drive-by since I'm trying to clean up some other StableHLO code in #19792

There is another lowering for stablehlo's FFT op to Linalg:

iree/compiler/plugins/input/StableHLO/Conversion/StableHLOToIREEInputDialects.cpp

Lines 163 to 165 in bbe7f5c

	/// Converts stablehlo.fft operation to Linalg ops.
	struct FftOpConversion final : OpConversionPattern<mlir::stablehlo::FftOp> {
	using OpConversionPattern::OpConversionPattern;

I think that Linalg lowering won't run since this LinalgExt lowering runs first. It could be the case that this LinalgExt lowering doesn't support all variants and some fall through to the Linalg version though.

I'm planning on deleting that lowering to Linalg as part of my cleanup, but please speak up (or contribute test cases) if I'm overlooking some load bearing usage.

[giacs-epic]

@ScottTodd Thank you for writing.
Indeed the conversion I'm refactoring here applies only to cases where the input has power of 2 size (so it's a special case), all the other cases fall under the lowering you are removing in your PR.
I'm not knowledgeable about the StableHLO side of things, as I'm mainly focused on the Torch path. But it seems to me like the lowering your are removing is indeed the "main" one, although I'm surprised to see there are no tests exercising it currently (I only looked through lit tests though). Will it be replaced by the upstream code on openxla/stablehlo#1817 ?

[ScottTodd]

Okay, I'll keep the StableHLO lowering then. Glad I checked.

Indeed, the LinalgExt lowering only powers power of 2 sizes:

iree/compiler/plugins/input/StableHLO/Conversion/StableHLOToLinalgExt.cpp

Lines 372 to 385 in 9a34131

	// Only handle 2^n fft length.
	auto operandType =
	llvm::dyn_cast<RankedTensorType>(adaptor.getOperand().getType());
	if (!operandType || !operandType.hasStaticShape()) {
	return failure();
	}
	if (!llvm::all_equal(op.getFftLength())) {
	return rewriter.notifyMatchFailure(op, "non-splat length");
	}
	int fftLength = op.getFftLength().front();
	if (fftLength & (fftLength - 1)) {
	return rewriter.notifyMatchFailure(
	op, "expected FFT length to be a power of two");
	}

The original change that added a lowering path added only lit and e2e tests for powers of two: #5450, then when a special lowering was added for LinalgExt , the e2e tests in https://github.com/iree-org/iree/blob/main/tests/e2e/stablehlo_ops/fft.mlir switched to using that path and the old path now appears to be untested.

[giacs-epic]

A few nits. I'm not very familiar with FFT so I can't comment on the core logic, but overall makes sense to me.

If another reviewer has comments I'll defer to them, otherwise I'd be happy to approve.

@qedawkins Thank you for the review. I've applied the changes.