Optimize complex multiplication where imaginary part of the first operand is zero

[Antipath1]

closes #1988
This PR adds a rewrite pattern for stablehlo.multiply to optimize complex multiplication when the imaginary component of the first operand is zero.Mathematical simplification:Given
$(a + bi) \times (c + di)$ : If $b = 0$, the expression simplifies:
$(a \cdot c - b \cdot d) + (c \cdot b + a \cdot d)i \implies (a \cdot c - 0 \cdot d) + (c \cdot 0 + a \cdot d)i = a \cdot c + a \cdot d \cdot i$

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 25.10%. Comparing base (003d2cd) to head (0b68664).
⚠️ Report is 4 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2501      +/-   ##
==========================================
- Coverage   25.47%   25.10%   -0.38%     
==========================================
  Files         220      220              
  Lines       44609    44600       -9     
==========================================
- Hits        11366    11197     -169     
- Misses      33243    33403     +160     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[mofeing]

maybe you should make use of the PurelyRealResultAnalysis and PurelyImagResultAnalysis, so that this pattern can propagate further

Enzyme-JAX/src/enzyme_ad/jax/Utils.h

Lines 867 to 893 in 14b57c2

	class PurelyRealResultAnalysis
	: public GuaranteedResultAnalysisBase<PurelyRealResultAnalysis> {
	public:
	State localGuaranteed(Value val, SmallVectorImpl<Value> &localtodo,
	PatternRewriter &rewriter);
	bool constantComplexCheck(DenseElementsAttr attr);
	bool constantFloatCheck(DenseElementsAttr attr) { return true; }
	bool constantIntCheck(DenseElementsAttr attr) { return true; }
	StringRef getAttrName() const { return "enzymexla.complex_is_purely_real"; }
	};
	class PurelyImagResultAnalysis
	: public GuaranteedResultAnalysisBase<PurelyImagResultAnalysis> {
	public:
	State localGuaranteed(Value val, SmallVectorImpl<Value> &localtodo,
	PatternRewriter &rewriter);
	bool constantComplexCheck(DenseElementsAttr attr);
	bool constantFloatCheck(DenseElementsAttr attr) { return true; }
	bool constantIntCheck(DenseElementsAttr attr) { return true; }
	StringRef getAttrName() const {
	return "enzymexla.complex_is_purely_imaginary";
	}
	};

the way I would do is check if the the ops from the values you're checking (i.e. the ancestor ops to the stablehlo.multiply) actually contain such attributes, and run the pattern after the analysis.

[Antipath1]

Okay, I have addressed this for the case of multiplication for a purely real number and a complex number. I have not gotten around to the case of a purely Imaginary number with a complex number, but I think other than that, I have addressed your feedback.

I had to edit two unrelated tests since using the pattern fired in two of the tests:

• conj_convert.mlir: I removed the enzymexla.complex_is_purely_real = [#enzymexla] attribute from the second check since the new pattern firing does result in the exact same ops generated. The attribute is not carried with it, but since the result is a real number anyway, I do not think it really makes sense to have that attribute there, and at a cursory glance I do not think it will enable any other patterns that would not work anyway.
• split_complex_scatter_setindex.mlir: The pattern also fires there, which is correct, but that, of course, leads to me also having to update the test's expected IR. I think the tested optimization there works exactly as before, there is just a previous part of the IR that gets changed, which I decided to change.

(cc @avik-pal: since you wrote these tests, could you double-check that my updates to the expected IR look right to you and tell me if I messed anything up there?)