TorchToTMTensor CumSum Lowering Only Supports Constant Axis

[renxida]

Here,

torch-mlir/lib/Conversion/TorchToTMTensor/TorchToTMTensor.cpp

Line 1532 in 34f6948

op, "unimplemented: only constant dim value is supported");

the TorchToTMTensor lowering checks to make sure that dim is a TorchConstantInt.

However, ONNX forces the axis to be variable because it allows axis to be a tensor.

This causes (at least part of) the lowering failure described in nod-ai/iree-amd-aie#103

opt-125M.fp32.onnx.mlir:231:12: error: failed to legalize operation 'torch.aten.cumsum' that was explicitly marked illegal
    %228 = torch.operator "onnx.CumSum"(%226, %227) : (!torch.vtensor<[1,6],si64>, !torch.vtensor<[],si32>) -> !torch.vtensor<[1,6],si64>
           ^
opt-125M.fp32.onnx.mlir:231:12: note: see current operation: %287 = "torch.aten.cumsum"(%18, %286, %16) : (!torch.vtensor<[1,6],si64>, !torch.int, !torch.int) -> !torch.vtensor<[1,6],si64> loc("opt-125M.fp32.onnx.mlir":231:12)

[renxida]

It looks like the main reason why TorchToTMTensor requires a constant dim is because it uses the dim input to compute the shape of an intermediate value and the output.

SmallVector<Value> accSizes(sizes);
    accSizes.erase(accSizes.begin() + dim);
    SmallVector<int64_t> accStatic(
        makeShapeTorchCompatible(resultType.getShape()));
    accStatic.erase(accStatic.begin() + dim);
    Value acc = createZeroInitTensor(rewriter, loc, accSizes, elementType);
    Type accType =
        RankedTensorType::get(makeShapeLLVMCompatible(accStatic), elementType);
    acc = rewriter.create<tensor::CastOp>(loc, accType, acc);

But maybe we can just use the output shape of the AtenCumsumOp being lowered and skip that.