Lowering to linalg for AtenCol2ImOp

[gdehame]

Added a lowering to linalg for the torch.aten.col2im operation.
Added a unit test to verify the lowering.

[sahas3]

Hi @gdehame, thanks for the submission.

I'm yet to review the code -- just leaving one high-level feedback. How does this op appear in the IR through the fx.export_and_import API usage? With the source pytorch program, it'll be good to also add an e2e test for numerical verification as well.

[sahas3]

I've taken a first pass and left some comments. Thanks!

[sahas3]

Suggested change

	if (!(col2imOp.getOutputSize().getDefiningOp() &&
	isa<Torch::PrimListConstructOp>(
	col2imOp.getOutputSize().getDefiningOp())))
	return failure();
	Torch::PrimListConstructOp outputSizes = cast<Torch::PrimListConstructOp>(
	col2imOp.getOutputSize().getDefiningOp());
	auto outputSizes = col2imOp.getOutputSize().getDefiningOp<Torch::PrimListConstructOp>();
	if (!outputSizes)
	return failure();

[gdehame]

Done, thanks!

[sahas3]

Here also similar suggestion about grabbing the def op at once and verifying it exists can be implemented as mentioned above.

[gdehame]

Done, thanks!

[sahas3]

Similar comment here and elsewhere too.

[gdehame]

Done, thanks!

[sahas3]

I think assumption here is that outputType dims 0 and 1 are static but that is not explicitly checked anywhere.

[gdehame]

I added a check, thanks!

[sahas3]

These nested tertiary ops are hard to read (here and in other places). Can you break this down?

[gdehame]

Done, thanks!

[sahas3]

This can be checked before starting to create any new ops in the IR and failure() can be returned for this unsupported case instead of an assert.

[gdehame]

Done, thanks!

[sahas3]

integer/complex types are not being tested. I'd suggest adding e2e tests for testing those combinations

[gdehame]

I started doing so and realized the integer case is bugged due to using arith.constant to generate a 0 constant of possibly (un)signed int which is not allowed.
I'll work on this further next week-end

[gdehame]

I updated the PR with additional tests and fixed the lowering

[gdehame]

Hi @gdehame, thanks for the submission.

I'm yet to review the code -- just leaving one high-level feedback. How does this op appear in the IR through the fx.export_and_import API usage? With the source pytorch program, it'll be good to also add an e2e test for numerical verification as well.

This came up during an internship I was doing last year, I completely forgot about the exact setup and example. I'll try to recreate one during the week-end

[gdehame]

Hi @gdehame, thanks for the submission.
I'm yet to review the code -- just leaving one high-level feedback. How does this op appear in the IR through the fx.export_and_import API usage? With the source pytorch program, it'll be good to also add an e2e test for numerical verification as well.

This came up during an internship I was doing last year, I completely forgot about the exact setup and example. I'll try to recreate one during the week-end

Here is an example:

from torch_mlir import fx
from torch_mlir.compiler_utils import OutputType
import torch
import numpy as np

fold = torch.nn.Fold((14, 30), (2, 2), (1, 1), (1, 1), (2, 2))

print(
    fx.export_and_import(
        torch.export.export(fold, (torch.tensor(np.zeros((1, 12, 128))),)),
        output_type=OutputType.LINALG_ON_TENSORS,
    )
)