Add raw TorchScript->Linalg example

[ramiro050]

The purpose of this PR is to add a simple and quick way of performing an e2e test for PyTorch ops that don't have Python binding. An example of this is the op prim.NumToTensor.Scalar #393. (We could eventually add something like this to the e2e test suite).

While there is a way to test raw MLIR in torch-mlir, those tests skip many passes that the e2e tests have to go through, such as the RefineTypes pass, which can lead to uncaught bugs.

This example works by taking TorchScript as a string, importing it into a torch.jit.ScriptFunction, then feeding this to torch-mlir so that it gets lowered and compiled using the RefBackendLinalgOnTensorsBackend.

Below is the output of this example:

TorchScript Graph
-----------------
graph(%p0 : Tensor,
      %p1 : Tensor):
  %2 : int = prim::Constant[value=1]()
  %3 : Tensor = aten::add(%p0, %p1, %2)
  return (%3)


Torch-MLIR
----------
module  {
  func @my_method(%arg0: !torch.tensor {torch.type_bound = !torch.vtensor<[3,4],f32>}, %arg1: !torch.tensor {torch.type_bound = !torch.vtensor<[3,4],f32>}) -> !torch.tensor {
    %int1 = torch.constant.int 1
    %0 = torch.aten.add.Tensor %arg0, %arg1, %int1 : !torch.tensor, !torch.tensor, !torch.int -> !torch.tensor
    return %0 : !torch.tensor
  }
}

Linalg-MLIR
-----------
#map = affine_map<(d0, d1) -> (d0, d1)>
module  {
  func @my_method(%arg0: tensor<3x4xf32>, %arg1: tensor<3x4xf32>) -> tensor<?x?xf32> {
    %c1_i64 = arith.constant 1 : i64
    %0 = linalg.init_tensor [3, 4] : tensor<3x4xf32>
    %1 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg0, %arg1 : tensor<3x4xf32>, tensor<3x4xf32>) outs(%0 : tensor<3x4xf32>) {
    ^bb0(%arg2: f32, %arg3: f32, %arg4: f32):  // no predecessors
      %3 = arith.sitofp %c1_i64 : i64 to f32
      %4 = arith.mulf %arg3, %3 : f32
      %5 = arith.addf %arg2, %4 : f32
      linalg.yield %5 : f32
    } -> tensor<3x4xf32>
    %2 = tensor.cast %1 : tensor<3x4xf32> to tensor<?x?xf32>
    return %2 : tensor<?x?xf32>
  }
}


Running Compiled Graph
----------------------
Expected output:
tensor([[ 0.1539, -3.9302, -0.7747, -1.1965],
        [ 2.5253, -0.2176, -1.0435, -1.2479],
        [-0.2396,  0.5140,  0.1770,  0.6730]])
Output from compiled MLIR:
[[ 0.15393919 -3.9302053  -0.7747382  -1.1964501 ]
 [ 2.5252578  -0.21755877 -1.043476   -1.247885  ]
 [-0.23960173  0.51403964  0.17700636  0.6729531 ]]

[ramiro050]

@cathyzhyi @silvasean PTAL

[silvasean]

I'm trying to understand what your goal is here, as I'd like to keep the number of free-floating scripts to a minimum -- they are a pain to update, and aren't part of CI.

Usually if one has TorchScript already, it's not hard to write a Python-level test by just looking at the source code locations to find where a particular op comes from. In our intermediate state before LTC is brought up, that might mean finding locations from LTC manually for the moment, which isn't great, but the benefits are large to not create divergence here.

[ramiro050]

I think any time we have a model where the easiest way to get the ops from is by some form of tracing (which is what happened with Bert), we won't be able to look at the TorchScript to get source location information because torch.jit.trace does not capture it. This is also true for the current method of getting ops using LTC.

I understand the concern of having too many free-floating scripts. If this is too big of an issue, I can host the script on a repo owned by me.

[silvasean]

Yeah, that's sad about torch.jit.trace. Thankfully all the ops traced from LTC will be dispatcher ops and should be available in torch.ops.aten.* trivially. And I guess almost all the torch.jit.trace ops will be dispatcher ops too? Not sure why NumToTensor is so special here.

[ramiro050]

Yeah, it will probably be uncommon that this will be needed. Sometimes it is difficult to decipher what the right op should be for ops without documentation. This was also true for torch.int.Tensor #387, which I only figured out because I stumbled upon it by accident.

I will move this to a repo of my own, since it does not seem crucial to have here.