Add EvalTranspose pattern to StablehloAggressiveFolder

[mvpant]

This patch folds stablehlo.transpose operation with constant operand into stablehlo.constant.

I have considered doing this by iterating over source index space, i.e.

auto initialValue = *std::begin(data);
SmallVector<ElementType> result(resultType.getNumElements(), initialValue);

for (int64_t i = 0; i < operandType.getNumElements(); ++i) {
  auto srcDimIndex = delinearize(i, operandStrides);
  auto dstDimIndex = applyPermutation(srcDimIndex, permutation);
  auto dstLinearIndex = linearize(dstDimIndex, resultStrides);
  result[dstLinearIndex] = data[i];
}

but that requires preinitialization of result vector with some value, which is twice as slow on simple case:

func.func @eval_transpose() -> (tensor<5000x80x30xi32>) {
  %0 = stablehlo.iota dim = 0 : tensor<30x80x5000xi32>
  %1 = stablehlo.transpose %0, dims = [2, 1, 0] : (tensor<30x80x5000xi32>) -> tensor<5000x80x30xi32>
  func.return %1 : tensor<5000x80x30xi32>
}

[GleasonK]

@mvpant wanted to give a heads up that I'm going to be doing some re-organizing of StablehloAggressiveSimplification in the coming days (not planning to touch folders this week), and add some patterns from MHLO canonicalizers. If you don't have any active changes..maybe hold off on starting any for a few days to avoid a nasty merge conflict.

[mvpant]

@mvpant wanted to give a heads up that I'm going to be doing some re-organizing of StablehloAggressiveSimplification in the coming days (not planning to touch folders this week), and add some patterns from MHLO canonicalizers. If you don't have any active changes..maybe hold off on starting any for a few days to avoid a nasty merge conflict.

@GleasonK sure. By the way, out of curiosity, what is the reason why stablehlo operations do not have <builtin> canonicalization using let hasCanonicalizer = 1; and folder let hasFolder = 1;? My guess is that this approach is intended to keep it stable for external users, ensuring that the generated output is predictable?

[GleasonK]

My guess is that this approach is intended to keep it stable for external users, ensuring that the generated output is predictable?

Exactly this. I've had too many issues with folks not being able to opt-out of certain patterns being applied and then they invent "clever" solutions like barrier ops to prevent constant folding, or mirrors of convert ops to avoid folding ops that would indicate that model weights are quantized, etc. In the current design point I'm trying to make everything as opt-in as possible, and have a stable input format which serialized and deserializes without changing.

Also I had a followup question, I see a few uses of tensor::EmptyOp, and I'm thinking to replace these with stablehlo::ConstantOp : dense<>. Any objections to this? I figure if tensor ops are needed it can be an extension pattern which converts empty constants to tensor::empty that lives in a part of the codebase that supports tensors, as opposed to a builtin canonicalization pattern

[mvpant]

... invent "clever" solutions like barrier ops to prevent constant folding, or mirrors of convert ops to avoid folding ops that would indicate that model weights are quantized, etc.

Eh, doesn't sound good, makes much more sense now to keep emitted operations as is.

Also I had a followup question, I see a few uses of tensor::EmptyOp, and I'm thinking to replace these with stablehlo::ConstantOp : dense<>. Any objections to this?

Nope. I think tensor::EmptyOp is handy at further stages because of the built-in bufferization stuff, but at stablehlo level it shouldn't make a difference. Having stablehlo.constant everywhere seems fine and more consistent.