Optimize folding transposes that operate on splats (#2863)

Instead of expanding and iterating over the splat to create the new
constant, we now just replace the splat constant's dimensions as
specified by the transpose op.

When tested on the same input program that brought this issue to light,
this fix improved the optimizer's execution time from 8.68 s to 0.80 s,
a 985% speedup.

Author: mrguenther

stablehlo/transforms/optimization/StablehloAggressiveFolder.cpp

@@ -1549,9 +1549,11 @@ struct FoldTransposeOpPattern : public FoldOpRewritePattern<TransposeOp> {
       return rewriter.notifyMatchFailure(
           op, "expected constant integer or float operand");
 
-    // TODO: Does this expand splat values? Should we special case splats?
     DenseElementsAttr resAttr;
-    if (auto data = els.tryGetValues<APInt>())
+    if (auto splat = dyn_cast<SplatElementsAttr>(els))
+      resAttr =
+          DenseElementsAttr::get(resultType, splat.getSplatValue<Attribute>());
+    else if (auto data = els.tryGetValues<APInt>())
       resAttr = transposeType(op, *data);
     else if (auto data = els.tryGetValues<APFloat>())
       resAttr = transposeType(op, *data);