Skip to content

torch.aten.as_strided to StableHLO #3972

New issue
New issue
Open
Open
torch.aten.as_strided to StableHLO#3972
@yankee624

Description

@yankee624
yankee624
opened on Jan 20, 2025

torch.aten.as_strided operation in torch dialect fails to get converted to StableHLO.
This is preventing the conversion of many PyTorch models to StableHLO. I checked that torch.nn.MultiheadAttention operation fails to get converted because of this. Language models from transformers library (e.g., Bert, Qwen) also fail, though I haven't checked whether the underlying cause is same.
How should I enable the conversion?

Runnable example:

  1. PyTorch to StableHLO
    There are some PyTorch operations that uses torch.aten.as_strided when converted to torch dialect. One example is torch.chunk.
    When I run the code below, I get the error: error: failed to legalize operation 'torch.constant.int'
import torch
from torch_mlir import fx

output_type = "stablehlo"

class Chunk(torch.nn.Module):
    def forward(self, x):
        return x.chunk(2)
model = Chunk()
module = fx.export_and_import(
    model,
    torch.rand(4),
    output_type=output_type,
    func_name=model.__class__.__name__,
)

with open("model.mlir", "w", encoding="utf-8") as outf:
    outf.write(module.operation.get_asm())

When I change the output_type to torch, the resulting model.mlir file is as follows:

module {
  func.func @Chunk(%arg0: !torch.vtensor<[4],f32>) -> (!torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>) {
    %int0 = torch.constant.int 0
    %int1 = torch.constant.int 1
    %int2 = torch.constant.int 2
    %0 = torch.prim.ListConstruct %int2 : (!torch.int) -> !torch.list<int>
    %1 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
    %2 = torch.aten.as_strided %arg0, %0, %1, %int0 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
    %3 = torch.prim.ListConstruct %int2 : (!torch.int) -> !torch.list<int>
    %4 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
    %5 = torch.aten.as_strided %arg0, %3, %4, %int2 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
    return %2, %5 : !torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>
  }
}
  1. torch dialect to StableHLO
    With the above model.mlir file, I ran
./build/bin/torch-mlir-opt ./model.mlir \
    --convert-torch-to-stablehlo \
    --shape-legalize-to-stablehlo \
    --convert-torch-to-arith \
    --canonicalize \
    --cse \
    -mlir-print-ir-after-all -mlir-disable-threading --debug  --mlir-print-ir-before-all \
    > tmp 2>&1

The output tmp file shows that the torch dialect operations fail to get converted to stablehlo operations.
Below is the content of the tmp file:

Args: ./build/bin/torch-mlir-opt ./sample/model.mlir --convert-torch-to-stablehlo --shape-legalize-to-stablehlo --convert-torch-to-arith --canonicalize --cse -mlir-print-ir-after-all -mlir-disable-threading --debug --mlir-print-ir-before-all 
Load new dialect in Context builtin
ImplicitTypeIDRegistry::lookupOrInsert(mlir::ShapedType)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::MemRefLayoutAttrInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::TypedAttr)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::ElementsAttr)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::DistinctAttr)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::BytecodeOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::SymbolOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpAsmOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::RegionKindInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::ConditionallySpeculatable)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::MemoryEffectOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::ResourceBlobManagerDialectInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpAsmDialectInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::BytecodeDialectInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::detail::AffineBinaryOpExprStorage)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::detail::AffineConstantExprStorage)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::detail::AffineDimExprStorage)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::detail::AffineMapStorage)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::detail::IntegerSetStorage)
Load new dialect in Context builtin
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::ZeroOperands<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::OneRegion<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::ZeroResults<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::ZeroSuccessors<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::NoRegionArguments<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::NoTerminator<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::SingleBlock<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::OpInvariants<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::BytecodeOpInterface::Trait<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::AffineScope<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::IsIsolatedFromAbove<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::SymbolTable<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::SymbolOpInterface::Trait<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpAsmOpInterface::Trait<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::RegionKindInterface::Trait<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::HasOnlyGraphRegion<Empty>)
Load new dialect in Context func
ImplicitTypeIDRegistry::lookupOrInsert(mlir::CallOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::SymbolUserOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::CallableOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::FunctionOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::RegionBranchTerminatorOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::DialectInlinerInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::ConvertToLLVMPatternInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::bufferization::BufferizableOpInterface)
Load new dialect in Context cf
Load new dialect in Context arith
ImplicitTypeIDRegistry::lookupOrInsert(mlir::arith::ArithFastMathInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::VectorUnrollOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::InferTypeOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::InferIntRangeInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::arith::ArithIntegerOverflowFlagsInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::CastOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::arith::ArithRoundingModeInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::SelectLikeOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::bufferization::BufferDeallocationOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::ValueBoundsOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::BranchOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::AutomaticAllocationScope<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::CallableOpInterface::Trait<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::FunctionOpInterface::Trait<Empty>)
Load new dialect in Context torch
ImplicitTypeIDRegistry::lookupOrInsert(mlir::RegionBranchOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::ZeroRegions<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::OneResult<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::OneTypedResult<mlir::torch::Torch::IntType>::Impl<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::ConstantLike<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::ConditionallySpeculatable::Trait<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::AlwaysSpeculatableImplTrait<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::MemoryEffectOpInterface::Trait<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::torch::Torch::OpTrait::AllowedInModuleInitializer<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::InferTypeOpInterface::Trait<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::OneTypedResult<mlir::torch::Torch::ListType>::Impl<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::VariadicOperands<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::torch::Torch::OpTrait::AllowsTypeRefinement<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::OneTypedResult<mlir::Type>::Impl<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::NOperands<4>::Impl<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::torch::Torch::OpTrait::ReadOnly<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::HasParent<mlir::func::FuncOp>::Impl<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::MemRefsNormalizable<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::RegionBranchTerminatorOpInterface::Trait<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::ReturnLike<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::IsTerminator<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::detail::OpToOpPassAdaptor)
Load new dialect in Context chlo
ImplicitTypeIDRegistry::lookupOrInsert(mlir::InferShapedTypeOpInterface)
Load new dialect in Context math
Load new dialect in Context shape
Load new dialect in Context tensor
Load new dialect in Context affine
Load new dialect in Context ub
ImplicitTypeIDRegistry::lookupOrInsert(mlir::ub::PoisonAttrInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::affine::AffineDmaStartOp)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::affine::AffineMapAccessInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::affine::AffineDmaWaitOp)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::LoopLikeOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::affine::AffineReadOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::affine::AffineWriteOpInterface)
Load new dialect in Context complex
ImplicitTypeIDRegistry::lookupOrInsert(mlir::ReifyRankedShapedTypeOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::ShapedDimOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OffsetSizeAndStrideOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::DestinationStyleOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::transform::FindPayloadReplacementOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::SubsetOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::SubsetInsertionOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::SubsetExtractionOpInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::TilingInterface)
Load new dialect in Context stablehlo
ImplicitTypeIDRegistry::lookupOrInsert(mlir::hlo::HloDialectInterface)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::VerifiableTensorEncoding)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::hlo::BoundedAttrInterface)
Load new dialect in Context torch_c
// -----// IR Dump Before ConvertTorchToStablehlo (convert-torch-to-stablehlo) //----- //
func.func @Chunk(%arg0: !torch.vtensor<[4],f32>) -> (!torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>) {
  %int0 = torch.constant.int 0
  %int1 = torch.constant.int 1
  %int2 = torch.constant.int 2
  %0 = torch.prim.ListConstruct %int2 : (!torch.int) -> !torch.list<int>
  %1 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
  %2 = torch.aten.as_strided %arg0, %0, %1, %int0 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
  %3 = torch.prim.ListConstruct %int2 : (!torch.int) -> !torch.list<int>
  %4 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
  %5 = torch.aten.as_strided %arg0, %3, %4, %int2 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
  return %2, %5 : !torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>
}


//===-------------------------------------------===//
Legalizing operation : 'func.func'(0x55bb62f94d30) {
  * Fold {
ImplicitTypeIDRegistry::lookupOrInsert(mlir::DialectFoldInterface)
  } -> FAILURE : unable to fold
} -> FAILURE : no matched legalization pattern
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.constant.int'(0x55bb62f3da90) {
  %0 = "torch.constant.int"() <{value = 0 : i64}> : () -> !torch.int

  * Fold {
  } -> FAILURE : unable to fold
} -> FAILURE : no matched legalization pattern
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.constant.int'(0x55bb62f95540) {
  %1 = "torch.constant.int"() <{value = 1 : i64}> : () -> !torch.int

  * Fold {
  } -> FAILURE : unable to fold
} -> FAILURE : no matched legalization pattern
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.constant.int'(0x55bb62f95670) {
  %2 = "torch.constant.int"() <{value = 2 : i64}> : () -> !torch.int

  * Fold {
  } -> FAILURE : unable to fold
} -> FAILURE : no matched legalization pattern
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.prim.ListConstruct'(0x55bb62fc8030) {
  %3 = "torch.prim.ListConstruct"(%2) : (!torch.int) -> !torch.list<int>

  * Fold {
  } -> FAILURE : unable to fold
} -> FAILURE : no matched legalization pattern
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.prim.ListConstruct'(0x55bb62fc85a0) {
  %4 = "torch.prim.ListConstruct"(%1) : (!torch.int) -> !torch.list<int>

  * Fold {
  } -> FAILURE : unable to fold
} -> FAILURE : no matched legalization pattern
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.aten.as_strided'(0x55bb62fc8690) {
  %5 = "torch.aten.as_strided"(%arg0, %3, %4, %0) : (!torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int) -> !torch.vtensor<[2],f32>

  * Fold {
  } -> FAILURE : unable to fold
} -> FAILURE : no matched legalization pattern
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.prim.ListConstruct'(0x55bb62fc87e0) {
  %6 = "torch.prim.ListConstruct"(%2) : (!torch.int) -> !torch.list<int>

  * Fold {
  } -> FAILURE : unable to fold
} -> FAILURE : no matched legalization pattern
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.prim.ListConstruct'(0x55bb62fc88d0) {
  %7 = "torch.prim.ListConstruct"(%1) : (!torch.int) -> !torch.list<int>

  * Fold {
  } -> FAILURE : unable to fold
} -> FAILURE : no matched legalization pattern
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.aten.as_strided'(0x55bb62fc89c0) {
  %8 = "torch.aten.as_strided"(%arg0, %6, %7, %2) : (!torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int) -> !torch.vtensor<[2],f32>

  * Fold {
  } -> FAILURE : unable to fold
} -> FAILURE : no matched legalization pattern
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'func.return'(0x55bb62fc8b20) {
  "func.return"(%5, %8) : (!torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>) -> ()

  * Fold {
  } -> FAILURE : unable to fold
} -> FAILURE : no matched legalization pattern
//===-------------------------------------------===//
ImplicitTypeIDRegistry::lookupOrInsert(mlir::detail::PreservedAnalyses::AllAnalysesType)
// -----// IR Dump After ConvertTorchToStablehlo (convert-torch-to-stablehlo) //----- //
func.func @Chunk(%arg0: !torch.vtensor<[4],f32>) -> (!torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>) {
  %int0 = torch.constant.int 0
  %int1 = torch.constant.int 1
  %int2 = torch.constant.int 2
  %0 = torch.prim.ListConstruct %int2 : (!torch.int) -> !torch.list<int>
  %1 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
  %2 = torch.aten.as_strided %arg0, %0, %1, %int0 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
  %3 = torch.prim.ListConstruct %int2 : (!torch.int) -> !torch.list<int>
  %4 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
  %5 = torch.aten.as_strided %arg0, %3, %4, %int2 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
  return %2, %5 : !torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>
}

// -----// IR Dump Before ShapeLegalizeToStablehloPass (shape-legalize-to-stablehlo) //----- //
func.func @Chunk(%arg0: !torch.vtensor<[4],f32>) -> (!torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>) {
  %int0 = torch.constant.int 0
  %int1 = torch.constant.int 1
  %int2 = torch.constant.int 2
  %0 = torch.prim.ListConstruct %int2 : (!torch.int) -> !torch.list<int>
  %1 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
  %2 = torch.aten.as_strided %arg0, %0, %1, %int0 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
  %3 = torch.prim.ListConstruct %int2 : (!torch.int) -> !torch.list<int>
  %4 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
  %5 = torch.aten.as_strided %arg0, %3, %4, %int2 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
  return %2, %5 : !torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>
}


//===-------------------------------------------===//
Legalizing operation : 'func.func'(0x55bb62f94d30) {
  * Fold {
  } -> FAILURE : unable to fold
} -> FAILURE : no matched legalization pattern
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.constant.int'(0x55bb62f3da90) {
  %0 = "torch.constant.int"() <{value = 0 : i64}> : () -> !torch.int

  * Fold {
  } -> FAILURE : unable to fold
} -> FAILURE : no matched legalization pattern
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.constant.int'(0x55bb62f95540) {
  %1 = "torch.constant.int"() <{value = 1 : i64}> : () -> !torch.int

  * Fold {
  } -> FAILURE : unable to fold
} -> FAILURE : no matched legalization pattern
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.constant.int'(0x55bb62f95670) {
  %2 = "torch.constant.int"() <{value = 2 : i64}> : () -> !torch.int

  * Fold {
  } -> FAILURE : unable to fold
} -> FAILURE : no matched legalization pattern
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.prim.ListConstruct'(0x55bb62fc8030) {
  %3 = "torch.prim.ListConstruct"(%2) : (!torch.int) -> !torch.list<int>

  * Fold {
  } -> FAILURE : unable to fold
} -> FAILURE : no matched legalization pattern
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.prim.ListConstruct'(0x55bb62fc85a0) {
  %4 = "torch.prim.ListConstruct"(%1) : (!torch.int) -> !torch.list<int>

  * Fold {
  } -> FAILURE : unable to fold
} -> FAILURE : no matched legalization pattern
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.aten.as_strided'(0x55bb62fc8690) {
  %5 = "torch.aten.as_strided"(%arg0, %3, %4, %0) : (!torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int) -> !torch.vtensor<[2],f32>

  * Fold {
  } -> FAILURE : unable to fold
} -> FAILURE : no matched legalization pattern
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.prim.ListConstruct'(0x55bb62fc87e0) {
  %6 = "torch.prim.ListConstruct"(%2) : (!torch.int) -> !torch.list<int>

  * Fold {
  } -> FAILURE : unable to fold
} -> FAILURE : no matched legalization pattern
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.prim.ListConstruct'(0x55bb62fc88d0) {
  %7 = "torch.prim.ListConstruct"(%1) : (!torch.int) -> !torch.list<int>

  * Fold {
  } -> FAILURE : unable to fold
} -> FAILURE : no matched legalization pattern
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.aten.as_strided'(0x55bb62fc89c0) {
  %8 = "torch.aten.as_strided"(%arg0, %6, %7, %2) : (!torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int) -> !torch.vtensor<[2],f32>

  * Fold {
  } -> FAILURE : unable to fold
} -> FAILURE : no matched legalization pattern
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'func.return'(0x55bb62fc8b20) {
  "func.return"(%5, %8) : (!torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>) -> ()

  * Fold {
  } -> FAILURE : unable to fold
} -> FAILURE : no matched legalization pattern
//===-------------------------------------------===//
// -----// IR Dump After ShapeLegalizeToStablehloPass (shape-legalize-to-stablehlo) //----- //
func.func @Chunk(%arg0: !torch.vtensor<[4],f32>) -> (!torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>) {
  %int0 = torch.constant.int 0
  %int1 = torch.constant.int 1
  %int2 = torch.constant.int 2
  %0 = torch.prim.ListConstruct %int2 : (!torch.int) -> !torch.list<int>
  %1 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
  %2 = torch.aten.as_strided %arg0, %0, %1, %int0 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
  %3 = torch.prim.ListConstruct %int2 : (!torch.int) -> !torch.list<int>
  %4 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
  %5 = torch.aten.as_strided %arg0, %3, %4, %int2 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
  return %2, %5 : !torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>
}

// -----// IR Dump Before ConvertTorchToArith (convert-torch-to-arith) //----- //
func.func @Chunk(%arg0: !torch.vtensor<[4],f32>) -> (!torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>) {
  %int0 = torch.constant.int 0
  %int1 = torch.constant.int 1
  %int2 = torch.constant.int 2
  %0 = torch.prim.ListConstruct %int2 : (!torch.int) -> !torch.list<int>
  %1 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
  %2 = torch.aten.as_strided %arg0, %0, %1, %int0 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
  %3 = torch.prim.ListConstruct %int2 : (!torch.int) -> !torch.list<int>
  %4 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
  %5 = torch.aten.as_strided %arg0, %3, %4, %int2 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
  return %2, %5 : !torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>
}


//===-------------------------------------------===//
Legalizing operation : 'func.func'(0x55bb62f94d30) {
} -> SUCCESS : operation marked legal by the target
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.constant.int'(0x55bb62f3da90) {
  %0 = "torch.constant.int"() <{value = 0 : i64}> : () -> !torch.int

  * Fold {
  } -> FAILURE : unable to fold

  * Pattern : 'torch.constant.int -> ()' {
Trying to match "(anonymous namespace)::ConvertTorchConstantIntOp"
ImplicitTypeIDRegistry::lookupOrInsert(mlir::arith::detail::ConstantOpGenericAdaptorBase::Properties)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::InferIntRangeInterface::Trait<Empty>)
    ** Insert  : 'arith.constant'(0x55bb630073e0)
    ** Replace : 'torch.constant.int'(0x55bb62f3da90)
"(anonymous namespace)::ConvertTorchConstantIntOp" result 1

    //===-------------------------------------------===//
    Legalizing operation : 'arith.constant'(0x55bb630073e0) {
      %0 = "arith.constant"() <{value = 0 : i64}> : () -> i64

    } -> SUCCESS : operation marked legal by the target
    //===-------------------------------------------===//
  } -> SUCCESS : pattern applied successfully
// *** IR Dump After Pattern Application ***
func.func @Chunk(%arg0: !torch.vtensor<[4],f32>) -> (!torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>) {
  %c0_i64 = arith.constant 0 : i64
  %int0 = torch.constant.int 0
  %int1 = torch.constant.int 1
  %int2 = torch.constant.int 2
  %0 = torch.prim.ListConstruct %int2 : (!torch.int) -> !torch.list<int>
  %1 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
  %2 = torch.aten.as_strided %arg0, %0, %1, %int0 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
  %3 = torch.prim.ListConstruct %int2 : (!torch.int) -> !torch.list<int>
  %4 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
  %5 = torch.aten.as_strided %arg0, %3, %4, %int2 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
  return %2, %5 : !torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>
}


} -> SUCCESS
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.constant.int'(0x55bb62f95540) {
  %2 = "torch.constant.int"() <{value = 1 : i64}> : () -> !torch.int

  * Fold {
  } -> FAILURE : unable to fold

  * Pattern : 'torch.constant.int -> ()' {
Trying to match "(anonymous namespace)::ConvertTorchConstantIntOp"
    ** Insert  : 'arith.constant'(0x55bb63007a90)
    ** Replace : 'torch.constant.int'(0x55bb62f95540)
"(anonymous namespace)::ConvertTorchConstantIntOp" result 1

    //===-------------------------------------------===//
    Legalizing operation : 'arith.constant'(0x55bb63007a90) {
      %2 = "arith.constant"() <{value = 1 : i64}> : () -> i64

    } -> SUCCESS : operation marked legal by the target
    //===-------------------------------------------===//
  } -> SUCCESS : pattern applied successfully
// *** IR Dump After Pattern Application ***
func.func @Chunk(%arg0: !torch.vtensor<[4],f32>) -> (!torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>) {
  %c0_i64 = arith.constant 0 : i64
  %int0 = torch.constant.int 0
  %c1_i64 = arith.constant 1 : i64
  %int1 = torch.constant.int 1
  %int2 = torch.constant.int 2
  %0 = torch.prim.ListConstruct %int2 : (!torch.int) -> !torch.list<int>
  %1 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
  %2 = torch.aten.as_strided %arg0, %0, %1, %int0 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
  %3 = torch.prim.ListConstruct %int2 : (!torch.int) -> !torch.list<int>
  %4 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
  %5 = torch.aten.as_strided %arg0, %3, %4, %int2 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
  return %2, %5 : !torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>
}


} -> SUCCESS
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.constant.int'(0x55bb62f95670) {
  %4 = "torch.constant.int"() <{value = 2 : i64}> : () -> !torch.int

  * Fold {
  } -> FAILURE : unable to fold

  * Pattern : 'torch.constant.int -> ()' {
Trying to match "(anonymous namespace)::ConvertTorchConstantIntOp"
    ** Insert  : 'arith.constant'(0x55bb630177f0)
    ** Replace : 'torch.constant.int'(0x55bb62f95670)
"(anonymous namespace)::ConvertTorchConstantIntOp" result 1

    //===-------------------------------------------===//
    Legalizing operation : 'arith.constant'(0x55bb630177f0) {
      %4 = "arith.constant"() <{value = 2 : i64}> : () -> i64

    } -> SUCCESS : operation marked legal by the target
    //===-------------------------------------------===//
  } -> SUCCESS : pattern applied successfully
// *** IR Dump After Pattern Application ***
func.func @Chunk(%arg0: !torch.vtensor<[4],f32>) -> (!torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>) {
  %c0_i64 = arith.constant 0 : i64
  %int0 = torch.constant.int 0
  %c1_i64 = arith.constant 1 : i64
  %int1 = torch.constant.int 1
  %c2_i64 = arith.constant 2 : i64
  %int2 = torch.constant.int 2
  %0 = torch.prim.ListConstruct %int2 : (!torch.int) -> !torch.list<int>
  %1 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
  %2 = torch.aten.as_strided %arg0, %0, %1, %int0 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
  %3 = torch.prim.ListConstruct %int2 : (!torch.int) -> !torch.list<int>
  %4 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
  %5 = torch.aten.as_strided %arg0, %3, %4, %int2 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
  return %2, %5 : !torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>
}


} -> SUCCESS
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.prim.ListConstruct'(0x55bb62fc8030) {
  %6 = "torch.prim.ListConstruct"(%5) : (!torch.int) -> !torch.list<int>

} -> SUCCESS : operation marked legal by the target
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.prim.ListConstruct'(0x55bb62fc85a0) {
  %7 = "torch.prim.ListConstruct"(%3) : (!torch.int) -> !torch.list<int>

} -> SUCCESS : operation marked legal by the target
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.aten.as_strided'(0x55bb62fc8690) {
  %8 = "torch.aten.as_strided"(%arg0, %6, %7, %1) : (!torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int) -> !torch.vtensor<[2],f32>

} -> SUCCESS : operation marked legal by the target
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.prim.ListConstruct'(0x55bb62fc87e0) {
  %9 = "torch.prim.ListConstruct"(%5) : (!torch.int) -> !torch.list<int>

} -> SUCCESS : operation marked legal by the target
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.prim.ListConstruct'(0x55bb62fc88d0) {
  %10 = "torch.prim.ListConstruct"(%3) : (!torch.int) -> !torch.list<int>

} -> SUCCESS : operation marked legal by the target
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'torch.aten.as_strided'(0x55bb62fc89c0) {
  %11 = "torch.aten.as_strided"(%arg0, %9, %10, %5) : (!torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int) -> !torch.vtensor<[2],f32>

} -> SUCCESS : operation marked legal by the target
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'func.return'(0x55bb62fc8b20) {
  "func.return"(%8, %11) : (!torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>) -> ()

} -> SUCCESS : operation marked legal by the target
//===-------------------------------------------===//
// -----// IR Dump After ConvertTorchToArith (convert-torch-to-arith) //----- //
func.func @Chunk(%arg0: !torch.vtensor<[4],f32>) -> (!torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>) {
  %c0_i64 = arith.constant 0 : i64
  %0 = torch_c.from_i64 %c0_i64
  %c1_i64 = arith.constant 1 : i64
  %1 = torch_c.from_i64 %c1_i64
  %c2_i64 = arith.constant 2 : i64
  %2 = torch_c.from_i64 %c2_i64
  %3 = torch.prim.ListConstruct %2 : (!torch.int) -> !torch.list<int>
  %4 = torch.prim.ListConstruct %1 : (!torch.int) -> !torch.list<int>
  %5 = torch.aten.as_strided %arg0, %3, %4, %0 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
  %6 = torch.prim.ListConstruct %2 : (!torch.int) -> !torch.list<int>
  %7 = torch.prim.ListConstruct %1 : (!torch.int) -> !torch.list<int>
  %8 = torch.aten.as_strided %arg0, %6, %7, %2 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
  return %5, %8 : !torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>
}

ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::OneOperand<Empty>)
// -----// IR Dump Before Canonicalizer (canonicalize) //----- //
ImplicitTypeIDRegistry::lookupOrInsert(mlir::AttributeTrait::IsLocation<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::TypedAttr::Trait<Empty>)
module {
  func.func @Chunk(%arg0: !torch.vtensor<[4],f32>) -> (!torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>) {
    %c0_i64 = arith.constant 0 : i64
    %0 = torch_c.from_i64 %c0_i64
    %c1_i64 = arith.constant 1 : i64
    %1 = torch_c.from_i64 %c1_i64
    %c2_i64 = arith.constant 2 : i64
    %2 = torch_c.from_i64 %c2_i64
    %3 = torch.prim.ListConstruct %2 : (!torch.int) -> !torch.list<int>
    %4 = torch.prim.ListConstruct %1 : (!torch.int) -> !torch.list<int>
    %5 = torch.aten.as_strided %arg0, %3, %4, %0 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
    %6 = torch.prim.ListConstruct %2 : (!torch.int) -> !torch.list<int>
    %7 = torch.prim.ListConstruct %1 : (!torch.int) -> !torch.list<int>
    %8 = torch.aten.as_strided %arg0, %6, %7, %2 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
    return %5, %8 : !torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>
  }
}



//===-------------------------------------------===//
Processing operation : 'func.func'(0x55bb62f94d30) {
} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'arith.constant'(0x55bb630073e0) {
  %2 = "arith.constant"() <{value = 0 : i64}> : () -> i64

} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'torch_c.from_i64'(0x55bb6300ec40) {
  %3 = "torch_c.from_i64"(%2) : (i64) -> !torch.int

} -> success : operation was folded
//===-------------------------------------------===//
ImplicitTypeIDRegistry::lookupOrInsert(mlir::torch::Torch::detail::ConstantIntOpGenericAdaptorBase::Properties)
** Insert  : 'torch.constant.int'(0x55bb62fdb6e0)
** Replace : 'torch_c.from_i64'(0x55bb6300ec40)
** Modified: 'torch.aten.as_strided'(0x55bb62fc8690)
** Erase   : 'torch_c.from_i64'(0x55bb6300ec40)
// *** IR Dump After Successful Folding ***
func.func @Chunk(%arg0: !torch.vtensor<[4],f32>) -> (!torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>) {
  %c2_i64 = arith.constant 2 : i64
  %c1_i64 = arith.constant 1 : i64
  %c0_i64 = arith.constant 0 : i64
  %int0 = torch.constant.int 0
  %0 = torch_c.from_i64 %c1_i64
  %1 = torch_c.from_i64 %c2_i64
  %2 = torch.prim.ListConstruct %1 : (!torch.int) -> !torch.list<int>
  %3 = torch.prim.ListConstruct %0 : (!torch.int) -> !torch.list<int>
  %4 = torch.aten.as_strided %arg0, %2, %3, %int0 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
  %5 = torch.prim.ListConstruct %1 : (!torch.int) -> !torch.list<int>
  %6 = torch.prim.ListConstruct %0 : (!torch.int) -> !torch.list<int>
  %7 = torch.aten.as_strided %arg0, %5, %6, %1 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
  return %4, %7 : !torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>
}



//===-------------------------------------------===//
Processing operation : 'arith.constant'(0x55bb630073e0) {
  %2 = "arith.constant"() <{value = 0 : i64}> : () -> i64

ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::HasRecursiveMemoryEffects<Empty>)
  ** Erase   : 'arith.constant'(0x55bb630073e0)
} -> success : operation is trivially dead
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'func.func'(0x55bb62f94d30) {
} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'torch.constant.int'(0x55bb62fdb6e0) {
  %2 = "torch.constant.int"() <{value = 0 : i64}> : () -> !torch.int

} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'arith.constant'(0x55bb63007a90) {
  %1 = "arith.constant"() <{value = 1 : i64}> : () -> i64

} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'torch_c.from_i64'(0x55bb6300ecd0) {
  %3 = "torch_c.from_i64"(%1) : (i64) -> !torch.int

} -> success : operation was folded
//===-------------------------------------------===//
** Insert  : 'torch.constant.int'(0x55bb62fd7900)
** Replace : 'torch_c.from_i64'(0x55bb6300ecd0)
** Modified: 'torch.prim.ListConstruct'(0x55bb62fc88d0)
** Modified: 'torch.prim.ListConstruct'(0x55bb62fc85a0)
** Erase   : 'torch_c.from_i64'(0x55bb6300ecd0)
// *** IR Dump After Successful Folding ***
func.func @Chunk(%arg0: !torch.vtensor<[4],f32>) -> (!torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>) {
  %c2_i64 = arith.constant 2 : i64
  %c1_i64 = arith.constant 1 : i64
  %int0 = torch.constant.int 0
  %int1 = torch.constant.int 1
  %0 = torch_c.from_i64 %c2_i64
  %1 = torch.prim.ListConstruct %0 : (!torch.int) -> !torch.list<int>
  %2 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
  %3 = torch.aten.as_strided %arg0, %1, %2, %int0 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
  %4 = torch.prim.ListConstruct %0 : (!torch.int) -> !torch.list<int>
  %5 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
  %6 = torch.aten.as_strided %arg0, %4, %5, %0 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
  return %3, %6 : !torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>
}



//===-------------------------------------------===//
Processing operation : 'arith.constant'(0x55bb63007a90) {
  %1 = "arith.constant"() <{value = 1 : i64}> : () -> i64

  ** Erase   : 'arith.constant'(0x55bb63007a90)
} -> success : operation is trivially dead
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'func.func'(0x55bb62f94d30) {
} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'torch.constant.int'(0x55bb62fd7900) {
  %2 = "torch.constant.int"() <{value = 1 : i64}> : () -> !torch.int

} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'arith.constant'(0x55bb630177f0) {
  %0 = "arith.constant"() <{value = 2 : i64}> : () -> i64

} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'torch_c.from_i64'(0x55bb62ffe6b0) {
  %3 = "torch_c.from_i64"(%0) : (i64) -> !torch.int

} -> success : operation was folded
//===-------------------------------------------===//
** Insert  : 'torch.constant.int'(0x55bb62fd7820)
** Replace : 'torch_c.from_i64'(0x55bb62ffe6b0)
** Modified: 'torch.aten.as_strided'(0x55bb62fc89c0)
** Modified: 'torch.prim.ListConstruct'(0x55bb62fc87e0)
** Modified: 'torch.prim.ListConstruct'(0x55bb62fc8030)
** Erase   : 'torch_c.from_i64'(0x55bb62ffe6b0)
// *** IR Dump After Successful Folding ***
func.func @Chunk(%arg0: !torch.vtensor<[4],f32>) -> (!torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>) {
  %c2_i64 = arith.constant 2 : i64
  %int0 = torch.constant.int 0
  %int1 = torch.constant.int 1
  %int2 = torch.constant.int 2
  %0 = torch.prim.ListConstruct %int2 : (!torch.int) -> !torch.list<int>
  %1 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
  %2 = torch.aten.as_strided %arg0, %0, %1, %int0 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
  %3 = torch.prim.ListConstruct %int2 : (!torch.int) -> !torch.list<int>
  %4 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
  %5 = torch.aten.as_strided %arg0, %3, %4, %int2 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
  return %2, %5 : !torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>
}



//===-------------------------------------------===//
Processing operation : 'arith.constant'(0x55bb630177f0) {
  %0 = "arith.constant"() <{value = 2 : i64}> : () -> i64

  ** Erase   : 'arith.constant'(0x55bb630177f0)
} -> success : operation is trivially dead
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'func.func'(0x55bb62f94d30) {
} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'torch.constant.int'(0x55bb62fd7820) {
  %2 = "torch.constant.int"() <{value = 2 : i64}> : () -> !torch.int

} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'torch.prim.ListConstruct'(0x55bb62fc8030) {
  %3 = "torch.prim.ListConstruct"(%2) : (!torch.int) -> !torch.list<int>

} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'torch.prim.ListConstruct'(0x55bb62fc85a0) {
  %4 = "torch.prim.ListConstruct"(%1) : (!torch.int) -> !torch.list<int>

} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'torch.aten.as_strided'(0x55bb62fc8690) {
  %5 = "torch.aten.as_strided"(%arg0, %3, %4, %0) : (!torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int) -> !torch.vtensor<[2],f32>

} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'torch.prim.ListConstruct'(0x55bb62fc87e0) {
  %6 = "torch.prim.ListConstruct"(%2) : (!torch.int) -> !torch.list<int>

} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'torch.prim.ListConstruct'(0x55bb62fc88d0) {
  %7 = "torch.prim.ListConstruct"(%1) : (!torch.int) -> !torch.list<int>

} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'torch.aten.as_strided'(0x55bb62fc89c0) {
  %8 = "torch.aten.as_strided"(%arg0, %6, %7, %2) : (!torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int) -> !torch.vtensor<[2],f32>

} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'func.return'(0x55bb62fc8b20) {
  "func.return"(%5, %8) : (!torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>) -> ()

} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'func.func'(0x55bb62f94d30) {
} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'torch.constant.int'(0x55bb62fdb6e0) {
  %0 = "torch.constant.int"() <{value = 0 : i64}> : () -> !torch.int

} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'torch.constant.int'(0x55bb62fd7900) {
  %1 = "torch.constant.int"() <{value = 1 : i64}> : () -> !torch.int

} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'torch.constant.int'(0x55bb62fd7820) {
  %2 = "torch.constant.int"() <{value = 2 : i64}> : () -> !torch.int

} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'torch.prim.ListConstruct'(0x55bb62fc8030) {
  %3 = "torch.prim.ListConstruct"(%2) : (!torch.int) -> !torch.list<int>

} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'torch.prim.ListConstruct'(0x55bb62fc85a0) {
  %4 = "torch.prim.ListConstruct"(%1) : (!torch.int) -> !torch.list<int>

} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'torch.aten.as_strided'(0x55bb62fc8690) {
  %5 = "torch.aten.as_strided"(%arg0, %3, %4, %0) : (!torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int) -> !torch.vtensor<[2],f32>

} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'torch.prim.ListConstruct'(0x55bb62fc87e0) {
  %6 = "torch.prim.ListConstruct"(%2) : (!torch.int) -> !torch.list<int>

} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'torch.prim.ListConstruct'(0x55bb62fc88d0) {
  %7 = "torch.prim.ListConstruct"(%1) : (!torch.int) -> !torch.list<int>

} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'torch.aten.as_strided'(0x55bb62fc89c0) {
  %8 = "torch.aten.as_strided"(%arg0, %6, %7, %2) : (!torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int) -> !torch.vtensor<[2],f32>

} -> failure : pattern failed to match
//===-------------------------------------------===//

//===-------------------------------------------===//
Processing operation : 'func.return'(0x55bb62fc8b20) {
  "func.return"(%5, %8) : (!torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>) -> ()

} -> failure : pattern failed to match
//===-------------------------------------------===//
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
module {
  func.func @Chunk(%arg0: !torch.vtensor<[4],f32>) -> (!torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>) {
    %int0 = torch.constant.int 0
    %int1 = torch.constant.int 1
    %int2 = torch.constant.int 2
    %0 = torch.prim.ListConstruct %int2 : (!torch.int) -> !torch.list<int>
    %1 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
    %2 = torch.aten.as_strided %arg0, %0, %1, %int0 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
    %3 = torch.prim.ListConstruct %int2 : (!torch.int) -> !torch.list<int>
    %4 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
    %5 = torch.aten.as_strided %arg0, %3, %4, %int2 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
    return %2, %5 : !torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>
  }
}


// -----// IR Dump Before CSE (cse) //----- //
module {
  func.func @Chunk(%arg0: !torch.vtensor<[4],f32>) -> (!torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>) {
    %int0 = torch.constant.int 0
    %int1 = torch.constant.int 1
    %int2 = torch.constant.int 2
    %0 = torch.prim.ListConstruct %int2 : (!torch.int) -> !torch.list<int>
    %1 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
    %2 = torch.aten.as_strided %arg0, %0, %1, %int0 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
    %3 = torch.prim.ListConstruct %int2 : (!torch.int) -> !torch.list<int>
    %4 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
    %5 = torch.aten.as_strided %arg0, %3, %4, %int2 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
    return %2, %5 : !torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>
  }
}


ImplicitTypeIDRegistry::lookupOrInsert(mlir::DominanceInfo)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::OpTrait::IsCommutative<Empty>)
ImplicitTypeIDRegistry::lookupOrInsert(mlir::PostDominanceInfo)
// -----// IR Dump After CSE (cse) //----- //
module {
  func.func @Chunk(%arg0: !torch.vtensor<[4],f32>) -> (!torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>) {
    %int0 = torch.constant.int 0
    %int1 = torch.constant.int 1
    %int2 = torch.constant.int 2
    %0 = torch.prim.ListConstruct %int2 : (!torch.int) -> !torch.list<int>
    %1 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
    %2 = torch.aten.as_strided %arg0, %0, %1, %int0 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
    %3 = torch.aten.as_strided %arg0, %0, %1, %int2 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
    return %2, %3 : !torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>
  }
}


module {
  func.func @Chunk(%arg0: !torch.vtensor<[4],f32>) -> (!torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>) {
    %int0 = torch.constant.int 0
    %int1 = torch.constant.int 1
    %int2 = torch.constant.int 2
    %0 = torch.prim.ListConstruct %int2 : (!torch.int) -> !torch.list<int>
    %1 = torch.prim.ListConstruct %int1 : (!torch.int) -> !torch.list<int>
    %2 = torch.aten.as_strided %arg0, %0, %1, %int0 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
    %3 = torch.aten.as_strided %arg0, %0, %1, %int2 : !torch.vtensor<[4],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2],f32>
    return %2, %3 : !torch.vtensor<[2],f32>, !torch.vtensor<[2],f32>
  }
}

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions