[OnnxToTorch] Support non-rank-0 Loop index tensors (#4098)

The `onnx.Loop` index value isn't required to be rank-0 (e.g. it may
have a shape of `[1]`), so we can't use `NumToTensor.Scalar` since it
creates rank-0 tensors.
This switches to using `aten.full` instead, which eventually
canonicalizes to `NumToTensor.Scalar` in the rank-0 case.

Author: rkayaith

lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp

@@ -376,20 +376,25 @@ void mlir::torch::onnx_c::populateDefaultDomainGtoP(
         // primLoopOp loopBody expects torch.int as first arg
         // insert torch.int arg in loop body, convert to tensor,
         // replace all uses of old arg, delete old arg.
-        auto loopVarArg = loop.getRegion().front().getArgument(0);
+        auto loopVar = loop.getRegion().front().getArgument(0);
         // insert new Arg
         loop.getRegion().front().insertArgument(
             0U, rewriter.getType<Torch::IntType>(), binder.getLoc());
         auto newLoopVarArg = loop.getRegion().front().getArgument(0);
 
         // convert int arg to tensor of original Type
         rewriter.setInsertionPointToStart(&loop.getRegion().front());
-        Value loopVarVal = BlockArgument::Value(loopVarArg);
-        auto newTensor = rewriter.create<Torch::PrimNumToTensorScalarOp>(
-            loop.getRegion().op_begin()->getLoc(), loopVarVal.getType(),
-            newLoopVarArg);
-
-        loopVarArg.replaceAllUsesWith(newTensor);
+        auto loopVarType = dyn_cast<Torch::BaseTensorType>(loopVar.getType());
+        if (!loopVarType || !loopVarType.areAllSizesKnown())
+          return rewriter.notifyMatchFailure(
+              loopVar.getLoc(),
+              "loop iteration value must be a tensor with known sizes");
+        Value sizes = Torch::toIntListConstruct(rewriter, loopVar.getLoc(),
+                                                loopVarType.getSizes());
+        auto newTensor = torch::Torch::createInitTensor(
+            rewriter, loopVar.getLoc(), loopVarType, newLoopVarArg, sizes);
+
+        loopVar.replaceAllUsesWith(newTensor);
         loop.getRegion().eraseArgument(1);
 
         // primLoopOp loopBody has no condition arg

test/Conversion/TorchOnnxToTorch/simple_ops_g_to_p.mlir

@@ -1977,7 +1977,10 @@ func.func @test_loop_forlike(%arg0: !torch.vtensor<[],si64>, %arg1: !torch.vtens
   // CHECK:           %[[TRUE:.*]] = torch.constant.bool true
   // CHECK:           %[[LOOP:.*]] = torch.prim.Loop %[[MAX_TRIP_COUNT_INT]], %[[TRUE]], init(%[[LCD_1]]) {
   // CHECK:           ^bb0(%[[ITER_NUM:.*]]: !torch.int, %[[LCD_1_BODY:.*]]: !torch.vtensor<[1],f32>):
-  // CHECK:             %[[ITER_NUM_T:.*]] = torch.prim.NumToTensor.Scalar %[[ITER_NUM]] : !torch.int -> !torch.vtensor<[],si64>
+  // CHECK:             %[[SHAPE:.*]] = torch.prim.ListConstruct : () -> !torch.list
+  // CHECK:             %[[NONE_ITER_NUM_T:.*]] = torch.constant.none
+  // CHECK:             %[[DTYPE:.*]] = torch.constant.int 4
+  // CHECK:             %[[ITER_NUM_T:.*]] = torch.aten.full %[[SHAPE]], %[[ITER_NUM]], %[[DTYPE]], %[[NONE_ITER_NUM_T]], %[[NONE_ITER_NUM_T]], %[[NONE_ITER_NUM_T]] : !torch.list<int>, !torch.int, !torch.int, !torch.none, !torch.none, !torch.none -> !torch.vtensor<[],si64>
   // CHECK:             %[[NONE_1:.*]] = torch.constant.none
   // CHECK:             %[[CLONE_INP_COND:.*]] = torch.aten.clone %[[CONDITION_INP]], %[[NONE_1]] : !torch.vtensor<[],i1>, !torch.none -> !torch.vtensor<[],i1>
   // CHECK:             %[[CONST_ARR:.*]] = torch.vtensor.literal(dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00]> : tensor<5xf32>) : !torch.vtensor<[5],f32>