Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

torch.aten.as_strided to StableHLO #3972

Open
yankee624 opened this issue Jan 20, 2025 · 0 comments
Open

torch.aten.as_strided to StableHLO #3972

yankee624 opened this issue Jan 20, 2025 · 0 comments

Comments

@yankee624
Copy link

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>
  }
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

No branches or pull requests
1 participant
@yankee624