Fix SCF Forloop fails to convert to linalg when a tensor argument is supplied to the loop block

[renxida]

No description provided.

[rsuderman]

We should included a simple execution test in the e2e example

[renxida]

something i'm concerned about the e2e test case is it doesnt work when i give an input to the function as the iteration count. i'm currently giving it a range(50)

like this:

# ==============================================================================
class TorchPrimLoop_ForLike_TensorArg_NoUnroll_Module(torch.nn.Module):
    def __init__(self):
        super().__init__()

    @export
    @annotate_args([
        None,
        ([7,9], torch.float32, True),
        ([1], torch.int64, False)
    ])
    def forward(self, x: torch.Tensor, n) -> torch.Tensor:
        # currently our test case uses 
        # for i in range(50):
        # instead
        for i in range(int(n[0])):
            x = x + i
        return x

@register_test_case(module_factory=lambda: TorchPrimLoop_ForLike_TensorArg_NoUnroll_Module())
def TorchPrimLoop_ForLike_TensorArg_NoUnroll_Module_basic(module, tu: TestUtils):
    x_test = torch.zeros([7, 9]).float()
    n_test = torch.tensor([50], dtype=torch.int64)

    module.forward(x_test, n_test)

I suspect there's something going on with unrolling

[renxida]

Error message:

        torch_mlir.compiler_utils.TorchMlirCompilerError: Lowering TorchScript IR -> Torch Backend IR failed with the following diagnostics:


        python exception: Failure while executing pass pipeline:
        error: unknown: unsupported by backend contract: tensor with unknown rank
        note: unknown: see current operation: %3 = "torch.tensor_static_info_cast"(%arg0) : (!torch.vtensor<[7,9],f32>) -> !torch.vtensor
        note: unknown: this is likely due to a missing transfer function in abstract_interp_lib_gen.py

        For Torch-MLIR developers, the error can be reproduced with:
        $ torch-mlir-opt -pass-pipeline='builtin.module(torchscript-module-to-torch-backend-pipeline{backend-legal-ops=aten.flatten.using_ints,aten.adaptive_avg_pool1d,aten.unflatten.int extra-library=})' /tmp/TorchPrimLoop_ForLike_TensorArg_NoUnroll_Module.mlir
        Add '-mlir-print-ir-after-all -mlir-disable-threading' to get the IR dump for debugging purpose.



Summary:
    Failed: 1

[renxida]

I banged my head against it a bunch without luck. Current guess is the forloop somehow removes shape information but unrolling somehow preserves it.