[Torch][FX] Rewrite as_strided before Torch IR import (#4585)

Move `aten.as_strided.default` handling out of backend/Torch IR
decomposition and into FX import. The new
`torch_mlir.extras.fx_as_strided` pass runs before Torch IR
construction, while FakeTensor metadata still carries storage offset,
physical strides, and layout information. Supported static `as_strided`
reads are rewritten to `aten.index.Tensor`, raw
`aten.as_strided.default` is rejected if it reaches generic ATen import.

Delete the existing Torch `DecomposeComplexOps` lowering and the direct
TorchToTOSA lowering for `aten.as_strided`. Those lowerings flattened
the immediate tensor value and could not recover storage/layout metadata
after import, which is unsafe for channels-last tensors and some
view/copy boundaries.

Add FX importer tests and e2e coverage for normal view chains, stepped
slices, `view -> contiguous -> as_strided`, channels-last inputs and
parameters, and `to(memory_format=channels_last) -> as_strided`. Update
xfail/crashing sets for ONNX cases where export has already lost storage
semantics or aborts before the runner can report an xfail.

Also retag StableHLO integer resource literals to signless result types.
The FX rewrite emits literal index tensors, and StableHLO constants must
use signless integer element types even when the resource blob came from
a signed Torch integer tensor.

Author: Lallapallooza

lib/Conversion/TorchToStablehlo/Basic.cpp

@@ -15,6 +15,7 @@
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Shape/IR/Shape.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/DialectResourceBlobManager.h"
 #include "stablehlo/dialect/ChloOps.h"
 #include "stablehlo/dialect/StablehloOps.h"
 #include "torch-mlir/Conversion/TorchToStablehlo/StablehloLegalizeUtils.h"
@@ -894,8 +895,16 @@ LogicalResult ConvertAtenOp<ValueTensorLiteralOp>::matchAndRewrite(
     return success();
   }
 
-  rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(op, resultType,
-                                                     adaptor.getValue());
+  ElementsAttr attr = cast<ElementsAttr>(adaptor.getValue());
+  if (auto res = dyn_cast<DenseResourceElementsAttr>(attr)) {
+    // Resource-backed integer literals keep the Torch signedness in the
+    // attribute type. StableHLO integer tensors are signless, so retag the blob
+    // to the converted result type before constructing the constant.
+    auto shapedAttrTy = cast<ShapedType>(res.getType());
+    if (isa<IntegerType>(shapedAttrTy.getElementType()))
+      attr = DenseResourceElementsAttr::get(resultType, res.getRawHandle());
+  }
+  rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(op, resultType, attr);
   return success();
 }
 

lib/Conversion/TorchToTosa/TorchToTosa.cpp

@@ -9082,108 +9082,6 @@ LogicalResult ConvertAtenOp<AtenThresholdBackwardOp>::matchAndRewriteImpl(
   return success();
 }
 
-// Legalization for aten.as_strided
-template <>
-LogicalResult ConvertAtenOp<AtenAsStridedOp>::matchAndRewriteImpl(
-    AtenAsStridedOp op, OpAdaptor adaptor,
-    ConversionPatternRewriter &rewriter) const {
-  // To lower aten.as_strided to TOSA, we will first reshape the input tensor to
-  // an 1-D tensor, then calculate the indices of result elements based on the
-  // output size, stride and storage offset. With the reshaped 1-D tensor and
-  // the indices, we can apply Gather to extract the required elements into a
-  // new tensor and then reshape it back to the desired output shape.
-  auto self = adaptor.getSelf();
-
-  // Not a tensor type
-  auto selfType = dyn_cast<TensorType>(self.getType());
-  if (!selfType)
-    return rewriter.notifyMatchFailure(op, "Only tensor types are supported");
-  auto selfElemTy = selfType.getElementType();
-  auto selfShape = selfType.getShape();
-
-  auto resultType =
-      dyn_cast<TensorType>(typeConverter->convertType(op.getType()));
-  auto resultElemTy = resultType.getElementType();
-
-  // Get output size
-  SmallVector<int64_t> outputSize;
-  if (!matchPattern(op.getSize(), m_TorchListOfConstantInts(outputSize)))
-    return rewriter.notifyMatchFailure(
-        op, "Only a constant list form of output size is supported");
-
-  // Get stride
-  SmallVector<int64_t> stride;
-  if (!matchPattern(op.getStride(), m_TorchListOfConstantInts(stride)))
-    return rewriter.notifyMatchFailure(
-        op, "Only a constant list form of stride is supported");
-
-  // Get storage offset
-  int64_t offset;
-  if (!matchPattern(op.getStorageOffset(), m_TorchConstantInt(&offset)))
-    offset = 0;
-
-  // Reshape input tensor into an 1-D tensor
-  int64_t selfNumElems = std::accumulate(selfShape.begin(), selfShape.end(), 1,
-                                         std::multiplies<int64_t>());
-
-  auto self1D = tosa::ReshapeOp::create(
-      rewriter, op->getLoc(), RankedTensorType::get({selfNumElems}, selfElemTy),
-      self, tosa::getTosaConstShape(rewriter, op->getLoc(), {selfNumElems}));
-
-  // Calculate the target elements indices
-  SmallVector<int32_t> targetIndicesVec;
-  int64_t outputRank = outputSize.size();
-  int64_t outputNumElems = std::accumulate(outputSize.begin(), outputSize.end(),
-                                           1, std::multiplies<int64_t>());
-
-  for (int64_t i = 0; i < outputNumElems; i++) {
-    // Index formula:
-    // index[i] = coord_i_0 * stride[0] + coord_i_1 * stride[1] + ... +
-    //              coord_i_n * stride[n]
-    int32_t index = offset;
-    int64_t coordFinder = i;
-    for (int64_t dim = 0; dim < outputRank; dim++) {
-      int64_t indexCoord = coordFinder % outputSize[outputRank - dim - 1];
-      index += indexCoord * stride[outputRank - dim - 1];
-      coordFinder /= outputSize[outputRank - dim - 1];
-    }
-    targetIndicesVec.push_back(index);
-  }
-
-  auto targetIndices =
-      tosa::getConstTensor<int32_t>(rewriter, op, targetIndicesVec,
-                                    makeShapeTorchCompatible({outputNumElems}))
-          .value();
-
-  // Convert PyTorch-style indices and dim into TensorFlow-style indices
-  auto targetIndicesTf = tosa::convertTorchIndexToTfIndices(
-      rewriter, op, self1D.getResult(), targetIndices, 0);
-  if (!targetIndicesTf)
-    return rewriter.notifyMatchFailure(op,
-                                       "Convert PyTorch-style indices and dim "
-                                       "to TensorFlow-style indices failed");
-
-  // Gather the target elements from 1-D input tensor
-  // Apply TensorFlow GatherNdOp with TensorFlow-style indices to retrieve the
-  // target elements
-  auto gatherOp = tosa::convertGatherNdOp(
-      rewriter, op,
-      RankedTensorType::get(makeShapeTorchCompatible({outputNumElems}),
-                            resultElemTy),
-      self1D.getResult(), targetIndicesTf.value());
-
-  if (!gatherOp)
-    return rewriter.notifyMatchFailure(op, "Convert GatherNdOp failed");
-
-  auto result = tosa::ReshapeOp::create(
-      rewriter, op->getLoc(), resultType, gatherOp.value(),
-      tosa::getTosaConstShape(rewriter, op->getLoc(), outputSize));
-
-  rewriter.replaceOp(op, {result.getResult()});
-
-  return success();
-}
-
 // Legalization for torch.prims.collapse
 template <>
 LogicalResult ConvertAtenOp<PrimsCollapseOp>::matchAndRewriteImpl(
@@ -11472,7 +11370,6 @@ std::set<StringRef> populateTorchToTosaConversionPatternsAndIllegalOps(
   INSERT_ATENOP_PATTERN(AtenDiagEmbedOp);
   INSERT_ATENOP_PATTERN(AtenUniformOp);
   INSERT_ATENOP_PATTERN(AtenThresholdBackwardOp);
-  INSERT_ATENOP_PATTERN(AtenAsStridedOp);
   INSERT_ATENOP_PATTERN(AtenClampTensorOp);
   INSERT_ATENOP_PATTERN(PrimsCollapseOp);
   INSERT_ATENOP_PATTERN(AtenReflectionPad1dOp);

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

@@ -13176,205 +13176,6 @@ class DecomposeAtenBroadcastTensorsOp
 };
 } // namespace
 
-namespace {
-class DecomposeAtenAsStridedOp : public OpRewritePattern<AtenAsStridedOp> {
-public:
-  using OpRewritePattern<AtenAsStridedOp>::OpRewritePattern;
-  LogicalResult matchAndRewrite(AtenAsStridedOp op,
-                                PatternRewriter &rewriter) const override {
-
-    // The `aten.as_strided` operation is decomposed into a series of
-    // operations that compute the indices based on the provided sizes and
-    // strides, and then index into the flattened input tensor as follows:
-
-    // input_flat = input.view(-1)
-    //
-    // for dim, s in enumerate(self.size):
-    //     arange = torch.arange(s)
-    //     view_shape = []
-    //     for i in range(len(self.size)):
-    //         if i == dim:
-    //             view_shape.append(-1)
-    //         else:
-    //             view_shape.append(1)
-    //     arange = arange.view(view_shape)
-    //     if dim != 0:
-    //         idx = idx + arange * self.stride[dim]
-    //
-    // # Flatten indices and add offset
-    // final_indices = idx.reshape(-1) + self.storage_offset
-    //
-    // # Index the flattened input tensor
-    // output = input_flat[final_indices]
-    //
-    // # Reshape to desired output size
-    // return output.view(self.size)
-
-    Location loc = op.getLoc();
-    MLIRContext *context = op->getContext();
-    Value input = op.getSelf();
-    auto inputType = dyn_cast<BaseTensorType>(input.getType());
-
-    if (!inputType || !inputType.hasSizes())
-      return rewriter.notifyMatchFailure(op, "input must have sizes");
-
-    SmallVector<int64_t> sizesInts;
-    if (!matchPattern(op.getSize(), m_TorchListOfConstantInts(sizesInts)))
-      return rewriter.notifyMatchFailure(
-          op, "sizes must be a list of constant ints");
-
-    SmallVector<int64_t> stridesInts;
-    if (!matchPattern(op.getStride(), m_TorchListOfConstantInts(stridesInts)))
-      return rewriter.notifyMatchFailure(
-          op, "strides must be a list of constant ints");
-
-    int64_t storageOffset = 0;
-    if (!isa<Torch::NoneType>(op.getStorageOffset().getType())) {
-      if (!matchPattern(op.getStorageOffset(),
-                        m_TorchConstantInt(&storageOffset)))
-        return rewriter.notifyMatchFailure(
-            op, "storage_offset must be a constant integer");
-    }
-
-    ArrayRef<int64_t> inputSizes = inputType.getSizes();
-    int64_t inputRank = inputSizes.size();
-    int64_t resultRank = sizesInts.size();
-
-    Value cstZero =
-        ConstantIntOp::create(rewriter, loc, rewriter.getI64IntegerAttr(0));
-    if (inputRank > 1) {
-      // If the input is not a 1-d tensor, we need to flatten it
-      // to a 1D tensor before applying the strided indexing.
-      int64_t flattenedInputSize = 1;
-      for (int64_t size : inputSizes) {
-        if (size == kUnknownSize) {
-          flattenedInputSize = kUnknownSize;
-          break;
-        }
-        flattenedInputSize *= size;
-      }
-
-      auto flattenedInputTy =
-          cast<BaseTensorType>(inputType.getWithSizesAndDtype(
-              {flattenedInputSize}, inputType.getOptionalDtype()));
-
-      Value end = ConstantIntOp::create(
-          rewriter, loc, rewriter.getI64IntegerAttr(inputRank - 1));
-      input = AtenFlattenUsingIntsOp::create(rewriter, loc, flattenedInputTy,
-                                             input, cstZero, end);
-    }
-
-    Value cstOne =
-        ConstantIntOp::create(rewriter, loc, rewriter.getI64IntegerAttr(1));
-    Value cstMinusOne =
-        ConstantIntOp::create(rewriter, loc, rewriter.getI64IntegerAttr(-1));
-
-    SmallVector<int64_t> viewShapeInts(resultRank, 1);
-    SmallVector<Value> viewShapeListElems(resultRank, cstOne);
-
-    auto si64Type = IntegerType::get(context, 64, IntegerType::Signed);
-    Value finalIndices;
-    for (unsigned dim = 0; dim < sizesInts.size(); dim++) {
-      int64_t size = sizesInts[dim];
-      Value cstNone = ConstantNoneOp::create(rewriter, loc);
-      Value end = ConstantIntOp::create(rewriter, loc,
-                                        rewriter.getI64IntegerAttr(size));
-
-      auto arangeType =
-          ValueTensorType::get(context, llvm::ArrayRef(size), si64Type);
-      Value index = Torch::AtenArangeOp::create(
-          rewriter, loc, arangeType, end, cstNone, cstNone, cstNone, cstNone);
-
-      // Set the current dimension to -1 for broadcasting
-      viewShapeInts[dim] = -1;
-      viewShapeListElems[dim] = cstMinusOne;
-
-      Value viewShapeList = Torch::PrimListConstructOp::create(
-          rewriter, loc, Torch::ListType::get(Torch::IntType::get(context)),
-          viewShapeListElems);
-
-      auto viewType = ValueTensorType::get(
-          context, llvm::ArrayRef(viewShapeInts), si64Type);
-      index = AtenViewOp::create(rewriter, loc, viewType, index, viewShapeList);
-
-      // Multiply the index with the stride for the current dimension
-      Value cstStride = ConstantIntOp::create(
-          rewriter, loc, rewriter.getI64IntegerAttr(stridesInts[dim]));
-      index =
-          AtenMulScalarOp::create(rewriter, loc, viewType, index, cstStride);
-
-      // Reset the current dimension to 1 for the next iteration
-      viewShapeInts[dim] = 1;
-      viewShapeListElems[dim] = cstOne;
-
-      if (dim == 0) {
-        finalIndices = index;
-        continue;
-      }
-
-      // calculate common shape for broadcast
-      SmallVector<int64_t> broadcastShape;
-      SmallVector<Value> broadcastShapeValue;
-      computeBroadcastShape(rewriter, loc, {finalIndices, index},
-                            broadcastShape, broadcastShapeValue);
-      Type broadcastType = ValueTensorType::get(
-          context, llvm::ArrayRef(broadcastShape), si64Type);
-
-      finalIndices = AtenAddTensorOp::create(rewriter, loc, broadcastType,
-                                             finalIndices, index, cstOne);
-    }
-
-    int64_t flattenedResultSize = 1;
-    for (int64_t size : sizesInts)
-      flattenedResultSize *= size;
-
-    // Flattening the indices and adding the storage offset
-    finalIndices = AtenFlattenUsingIntsOp::create(
-        rewriter, loc,
-        ValueTensorType::get(context, llvm::ArrayRef(flattenedResultSize),
-                             si64Type),
-        finalIndices, cstZero, cstMinusOne); // -1 means flatten all
-
-    if (storageOffset != 0) {
-      Value cstStorageOffset = ConstantIntOp::create(
-          rewriter, loc, rewriter.getI64IntegerAttr(storageOffset));
-      finalIndices =
-          AtenAddScalarOp::create(rewriter, loc, finalIndices.getType(),
-                                  finalIndices, cstStorageOffset, cstOne);
-    }
-
-    // Index the flattened input tensor
-    Type listElemType =
-        inputType.getWithSizesAndDtype(/*optionalSizes=*/std::nullopt,
-                                       /*optionalDtype=*/nullptr);
-    Value indicesList = Torch::PrimListConstructOp::create(
-        rewriter, loc, Torch::ListType::get(listElemType),
-        SmallVector<Value>{finalIndices});
-
-    auto flattenedResultTy =
-        ValueTensorType::get(context, llvm::ArrayRef(flattenedResultSize),
-                             inputType.getOptionalDtype());
-    Value result = AtenIndexTensorOp::create(rewriter, loc, flattenedResultTy,
-                                             input, indicesList);
-
-    // Reshape the result to the desired output size
-    SmallVector<Value> sizesIntsValues;
-    for (int64_t size : sizesInts) {
-      sizesIntsValues.push_back(ConstantIntOp::create(
-          rewriter, loc, rewriter.getI64IntegerAttr(size)));
-    }
-    Value resultSizeList = Torch::PrimListConstructOp::create(
-        rewriter, loc, Torch::ListType::get(Torch::IntType::get(context)),
-        sizesIntsValues);
-    result =
-        AtenViewOp::create(rewriter, loc, op.getType(), result, resultSizeList);
-
-    rewriter.replaceOp(op, result);
-    return success();
-  }
-};
-} // namespace
-
 namespace {
 class DecomposeComplexOpsPass
     : public impl::DecomposeComplexOpsBase<DecomposeComplexOpsPass> {
@@ -13709,7 +13510,6 @@ class DecomposeComplexOpsPass
         patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAten_AssertScalarOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenRoundDecimalsOp>(patterns);
-    addPatternIfTargetOpIsIllegal<DecomposeAtenAsStridedOp>(patterns);
 
     GreedyRewriteConfig config;
     config.setUseTopDownTraversal(true);