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

Add aten.stft.center and decomposition #3880

Draft
wants to merge 19 commits into
base: main
Choose a base branch
from
Draft

Add aten.stft.center and decomposition #3880

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
19 commits
Select commit Hold shift + click to select a range
e5f20fd
Emit aten.stft.center
giacs-epic Oct 18, 2024
7becc57
Add WIP code to decompose STFT into fft or rfft
giacs-epic Oct 18, 2024
34af833
Further decompose stack into unsqueeze+cat
giacs-epic Oct 21, 2024
ae145aa
Add simplifications to unroll+fold list generation
giacs-epic Oct 24, 2024
d2a5d52
Fix unsqueeze axis
giacs-epic Nov 5, 2024
476573b
Fix output sizes
giacs-epic Nov 12, 2024
794df5a
Fix onesided
giacs-epic Nov 12, 2024
abe7dcf
Add e2e test case
giacs-epic Nov 12, 2024
d00a51e
Fix result type check
giacs-epic Nov 14, 2024
412c5af
Fix frame padding constant
giacs-epic Nov 14, 2024
5621908
Add aten.stft.center shape and dtype functions
giacs-epic Nov 14, 2024
39ad96d
Add more tests
giacs-epic Nov 14, 2024
bc33b19
Merge branch 'main' into decompose_STFT
giacs-epic Nov 15, 2024
da353c4
Switch to implementation that doesn't depend on simplification
giacs-epic Nov 18, 2024
15cf656
Merge branch 'main' into decompose_STFT
giacs-epic Nov 18, 2024
80d3eb2
Regenerate AbstractInterpLibrary.cpp
giacs-epic Nov 18, 2024
fd1563a
Fix empty/dynamic-sized window and dynamic-sized signal handling
giacs-epic Nov 25, 2024
bac4d77
Add unit tests
giacs-epic Nov 26, 2024
4bfc92c
Revert "Add simplifications to unroll+fold list generation"
giacs-epic Nov 26, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 58 additions & 0 deletions include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13323,6 +13323,32 @@ def Torch_AtenFftFftOp : Torch_Op<"aten.fft_fft", [
}];
}

def Torch_AtenFftRfftOp : Torch_Op<"aten.fft_rfft", [
AllowsTypeRefinement,
HasValueSemantics,
ReadOnly
]> {
let summary = "Generated op for `aten::fft_rfft : (Tensor, int?, int, str?) -> (Tensor)`";
let arguments = (ins
AnyTorchTensorType:$self,
AnyTorchOptionalIntType:$n,
Torch_IntType:$dim,
AnyTorchOptionalStringType:$norm
);
let results = (outs
AnyTorchOptionalTensorType:$result
);
let hasCustomAssemblyFormat = 1;
let extraClassDefinition = [{
ParseResult AtenFftRfftOp::parse(OpAsmParser &parser, OperationState &result) {
return parseDefaultTorchOp(parser, result, 4, 1);
}
void AtenFftRfftOp::print(OpAsmPrinter &printer) {
printDefaultTorchOp(printer, *this, 4, 1);
}
}];
}

def Torch_AtenFftIfftOp : Torch_Op<"aten.fft_ifft", [
AllowsTypeRefinement,
HasValueSemantics,
Expand Down Expand Up @@ -13571,6 +13597,38 @@ def Torch_AtenStftOp : Torch_Op<"aten.stft", [
}];
}

def Torch_AtenStftCenterOp : Torch_Op<"aten.stft.center", [
AllowsTypeRefinement,
HasValueSemantics,
ReadOnly
]> {
let summary = "Generated op for `aten::stft.center : (Tensor, int, int?, int?, Tensor?, bool, str, bool, bool?, bool?) -> (Tensor)`";
let arguments = (ins
AnyTorchTensorType:$self,
Torch_IntType:$n_fft,
AnyTorchOptionalIntType:$hop_length,
AnyTorchOptionalIntType:$win_length,
AnyTorchOptionalTensorType:$window,
Torch_BoolType:$center,
Torch_StringType:$pad_mode,
Torch_BoolType:$normalized,
AnyTorchOptionalBoolType:$onesided,
AnyTorchOptionalBoolType:$return_complex
);
let results = (outs
AnyTorchOptionalTensorType:$result
);
let hasCustomAssemblyFormat = 1;
let extraClassDefinition = [{
ParseResult AtenStftCenterOp::parse(OpAsmParser &parser, OperationState &result) {
return parseDefaultTorchOp(parser, result, 10, 1);
}
void AtenStftCenterOp::print(OpAsmPrinter &printer) {
printDefaultTorchOp(printer, *this, 10, 1);
}
}];
}

def Torch_AtenAliasCopyOp : Torch_Op<"aten.alias_copy", [
AllowsTypeRefinement,
HasValueSemantics,
Expand Down
2 changes: 2 additions & 0 deletions include/torch-mlir/Dialect/Torch/Utils/Utils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@ namespace Torch {

int64_t toPositiveDim(int64_t dim, int64_t inputRank);
bool isValidDim(int64_t dim, int64_t inputRank);
Value toIntListConstruct(PatternRewriter &rewriter, Location loc,
ArrayRef<int64_t> cstInput, Torch::IntType intType);
bool getListConstructElements(Value v, SmallVectorImpl<Value> &elems);
/// Returns the index indicated by `v` for a list of given `length`.
/// If the index is negative, it is adjusted to `length` + `v`.
Expand Down
8 changes: 5 additions & 3 deletions lib/Conversion/TorchOnnxToTorch/DefaultDomainQtoZ.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3526,6 +3526,8 @@ void mlir::torch::onnx_c::populateDefaultDomainQtoZ(
rewriter.create<Torch::ConstantBoolOp>(binder.getLoc(), false);
Value trueVal =
rewriter.create<Torch::ConstantBoolOp>(binder.getLoc(), true);
Value padMode =
rewriter.create<Torch::ConstantStrOp>(binder.getLoc(), "reflect");
auto stftTy = complexSignalTy.getWithSizesAndDtype(
ArrayRef<int64_t>({resultShape[0], resultShape[2], resultShape[1]}),
complexSignalTy.getDtype());
Expand All @@ -3538,10 +3540,10 @@ void mlir::torch::onnx_c::populateDefaultDomainQtoZ(
// shape of stft to match the shape of resultType. Also, it is
// immaterial whether torch.view_as_real is called after or before the
// permutation; both outputs will be equivalent.
Value stft = rewriter.create<Torch::AtenStftOp>(
Value stft = rewriter.create<Torch::AtenStftCenterOp>(
binder.getLoc(), stftTy, signal, frameLengthItem, frameStepItem,
windowLen, window, falseVal, onesided ? trueVal : falseVal,
trueVal);
windowLen, window, falseVal, padMode, falseVal,
onesided ? trueVal : falseVal, trueVal);

auto permuteStftTy = complexSignalTy.getWithSizesAndDtype(
ArrayRef<int64_t>({resultShape[0], resultShape[1], resultShape[2]}),
Expand Down
265 changes: 265 additions & 0 deletions lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11055,6 +11055,134 @@ StringRef mlir::torch::Torch::getAbstractInterpLibrary() {
" }\n"
" return %14 : !torch.list<int>\n"
" }\n"
" func.func @\"__torch_mlir_shape_fn.aten.stft.center\"(%arg0: !torch.list<int>, %arg1: !torch.int, %arg2: !torch.optional<int>, %arg3: !torch.optional<int>, %arg4: !torch.optional<list<int>>, %arg5: !torch.bool, %arg6: !torch.str, %arg7: !torch.bool, %arg8: !torch.optional<bool>, %arg9: !torch.optional<bool>) -> !torch.list<int> {\n"
" %str = torch.constant.str \"AssertionError: Expected hop_length to be greater than 0\"\n"
" %str_0 = torch.constant.str \"AssertionError: Expected that 0 < n_fft <= len\"\n"
" %false = torch.constant.bool false\n"
" %none = torch.constant.none\n"
" %str_1 = torch.constant.str \"AssertionError: Expected input tensor to be of shape (B?,L), where B is an optional batch dimension\"\n"
" %true = torch.constant.bool true\n"
" %int1 = torch.constant.int 1\n"
" %int2 = torch.constant.int 2\n"
" %int0 = torch.constant.int 0\n"
" %int4 = torch.constant.int 4\n"
" %0 = torch.aten.len.t %arg0 : !torch.list<int> -> !torch.int\n"
" %1 = torch.aten.eq.int %0, %int1 : !torch.int, !torch.int -> !torch.bool\n"
" %2 = torch.prim.If %1 -> (!torch.bool) {\n"
" torch.prim.If.yield %true : !torch.bool\n"
" } else {\n"
" %21 = torch.aten.len.t %arg0 : !torch.list<int> -> !torch.int\n"
" %22 = torch.aten.eq.int %21, %int2 : !torch.int, !torch.int -> !torch.bool\n"
" torch.prim.If.yield %22 : !torch.bool\n"
" }\n"
" torch.prim.If %2 -> () {\n"
" torch.prim.If.yield\n"
" } else {\n"
" torch.prim.RaiseException %str_1, %none : !torch.str, !torch.none\n"
" torch.prim.If.yield\n"
" }\n"
" %3 = torch.aten.len.t %arg0 : !torch.list<int> -> !torch.int\n"
" %4 = torch.aten.eq.int %3, %int1 : !torch.int, !torch.int -> !torch.bool\n"
" %5 = torch.prim.If %4 -> (!torch.optional<int>) {\n"
" %21 = torch.derefine %none : !torch.none to !torch.optional<int>\n"
" torch.prim.If.yield %21 : !torch.optional<int>\n"
" } else {\n"
" %21 = torch.aten.__getitem__.t %arg0, %int0 : !torch.list<int>, !torch.int -> !torch.int\n"
" %22 = torch.derefine %21 : !torch.int to !torch.optional<int>\n"
" torch.prim.If.yield %22 : !torch.optional<int>\n"
" }\n"
" %6 = torch.aten.len.t %arg0 : !torch.list<int> -> !torch.int\n"
" %7 = torch.aten.eq.int %6, %int1 : !torch.int, !torch.int -> !torch.bool\n"
" %8 = torch.prim.If %7 -> (!torch.int) {\n"
" %21 = torch.aten.__getitem__.t %arg0, %int0 : !torch.list<int>, !torch.int -> !torch.int\n"
" torch.prim.If.yield %21 : !torch.int\n"
" } else {\n"
" %21 = torch.aten.__getitem__.t %arg0, %int1 : !torch.list<int>, !torch.int -> !torch.int\n"
" torch.prim.If.yield %21 : !torch.int\n"
" }\n"
" %9 = torch.aten.__is__ %arg2, %none : !torch.optional<int>, !torch.none -> !torch.bool\n"
" %10 = torch.prim.If %9 -> (!torch.int) {\n"
" %21 = torch.aten.floordiv.int %arg1, %int4 : !torch.int, !torch.int -> !torch.int\n"
" torch.prim.If.yield %21 : !torch.int\n"
" } else {\n"
" %21 = torch.prim.unchecked_cast %arg2 : !torch.optional<int> -> !torch.int\n"
" torch.prim.If.yield %21 : !torch.int\n"
" }\n"
" %11 = torch.aten.gt.int %arg1, %int0 : !torch.int, !torch.int -> !torch.bool\n"
" %12 = torch.prim.If %11 -> (!torch.bool) {\n"
" %21 = torch.aten.le.int %arg1, %8 : !torch.int, !torch.int -> !torch.bool\n"
" torch.prim.If.yield %21 : !torch.bool\n"
" } else {\n"
" torch.prim.If.yield %false : !torch.bool\n"
" }\n"
" torch.prim.If %12 -> () {\n"
" torch.prim.If.yield\n"
" } else {\n"
" torch.prim.RaiseException %str_0, %none : !torch.str, !torch.none\n"
" torch.prim.If.yield\n"
" }\n"
" %13 = torch.aten.gt.int %10, %int0 : !torch.int, !torch.int -> !torch.bool\n"
" torch.prim.If %13 -> () {\n"
" torch.prim.If.yield\n"
" } else {\n"
" torch.prim.RaiseException %str, %none : !torch.str, !torch.none\n"
" torch.prim.If.yield\n"
" }\n"
" %14 = torch.prim.ListConstruct : () -> !torch.list<int>\n"
" %15 = torch.aten.__isnot__ %5, %none : !torch.optional<int>, !torch.none -> !torch.bool\n"
" torch.prim.If %15 -> () {\n"
" %21 = torch.prim.unchecked_cast %5 : !torch.optional<int> -> !torch.int\n"
" %22 = torch.aten.append.t %14, %21 : !torch.list<int>, !torch.int -> !torch.list<int>\n"
" torch.prim.If.yield\n"
" } else {\n"
" torch.prim.If.yield\n"
" }\n"
" %16 = torch.aten.__is__ %arg8, %none : !torch.optional<bool>, !torch.none -> !torch.bool\n"
" %17 = torch.prim.If %16 -> (!torch.bool) {\n"
" torch.prim.If.yield %true : !torch.bool\n"
" } else {\n"
" %21 = torch.prim.unchecked_cast %arg8 : !torch.optional<bool> -> !torch.bool\n"
" %22 = torch.aten.eq.bool %21, %true : !torch.bool, !torch.bool -> !torch.bool\n"
" torch.prim.If.yield %22 : !torch.bool\n"
" }\n"
" torch.prim.If %17 -> () {\n"
" %21 = torch.aten.floordiv.int %arg1, %int2 : !torch.int, !torch.int -> !torch.int\n"
" %22 = torch.aten.add.int %21, %int1 : !torch.int, !torch.int -> !torch.int\n"
" %23 = torch.aten.append.t %14, %22 : !torch.list<int>, !torch.int -> !torch.list<int>\n"
" torch.prim.If.yield\n"
" } else {\n"
" %21 = torch.aten.append.t %14, %arg1 : !torch.list<int>, !torch.int -> !torch.list<int>\n"
" torch.prim.If.yield\n"
" }\n"
" %18 = torch.aten.eq.bool %arg5, %true : !torch.bool, !torch.bool -> !torch.bool\n"
" torch.prim.If %18 -> () {\n"
" %21 = torch.aten.floordiv.int %8, %10 : !torch.int, !torch.int -> !torch.int\n"
" %22 = torch.aten.add.int %int1, %21 : !torch.int, !torch.int -> !torch.int\n"
" %23 = torch.aten.append.t %14, %22 : !torch.list<int>, !torch.int -> !torch.list<int>\n"
" torch.prim.If.yield\n"
" } else {\n"
" %21 = torch.aten.sub.int %8, %arg1 : !torch.int, !torch.int -> !torch.int\n"
" %22 = torch.aten.floordiv.int %21, %10 : !torch.int, !torch.int -> !torch.int\n"
" %23 = torch.aten.add.int %int1, %22 : !torch.int, !torch.int -> !torch.int\n"
" %24 = torch.aten.append.t %14, %23 : !torch.list<int>, !torch.int -> !torch.list<int>\n"
" torch.prim.If.yield\n"
" }\n"
" %19 = torch.aten.__isnot__ %arg9, %none : !torch.optional<bool>, !torch.none -> !torch.bool\n"
" %20 = torch.prim.If %19 -> (!torch.bool) {\n"
" %21 = torch.prim.unchecked_cast %arg9 : !torch.optional<bool> -> !torch.bool\n"
" %22 = torch.aten.eq.bool %21, %false : !torch.bool, !torch.bool -> !torch.bool\n"
" torch.prim.If.yield %22 : !torch.bool\n"
" } else {\n"
" torch.prim.If.yield %false : !torch.bool\n"
" }\n"
" torch.prim.If %20 -> () {\n"
" %21 = torch.aten.append.t %14, %int2 : !torch.list<int>, !torch.int -> !torch.list<int>\n"
" torch.prim.If.yield\n"
" } else {\n"
" torch.prim.If.yield\n"
" }\n"
" return %14 : !torch.list<int>\n"
" }\n"
" func.func @\"__torch_mlir_shape_fn.aten.fft_ifft\"(%arg0: !torch.list<int>, %arg1: !torch.optional<int>, %arg2: !torch.int, %arg3: !torch.optional<str>) -> !torch.list<int> {\n"
" return %arg0 : !torch.list<int>\n"
" }\n"
Expand Down Expand Up @@ -13214,6 +13342,143 @@ StringRef mlir::torch::Torch::getAbstractInterpLibrary() {
" }\n"
" return %6 : !torch.int\n"
" }\n"
" func.func @\"__torch_mlir_dtype_fn.aten.stft.center\"(%arg0: !torch.tuple<int, int>, %arg1: !torch.int, %arg2: !torch.optional<int>, %arg3: !torch.optional<int>, %arg4: !torch.optional<tuple<int, int>>, %arg5: !torch.bool, %arg6: !torch.str, %arg7: !torch.bool, %arg8: !torch.optional<bool>, %arg9: !torch.optional<bool>) -> !torch.int {\n"
" %str = torch.constant.str \"AssertionError: Unsupported dtype\"\n"
" %int7 = torch.constant.int 7\n"
" %int10 = torch.constant.int 10\n"
" %int6 = torch.constant.int 6\n"
" %int9 = torch.constant.int 9\n"
" %int5 = torch.constant.int 5\n"
" %int8 = torch.constant.int 8\n"
" %none = torch.constant.none\n"
" %false = torch.constant.bool false\n"
" %true = torch.constant.bool true\n"
" %0 = torch.prim.Uninitialized : !torch.int\n"
" %1:2 = torch.prim.TupleUnpack %arg0 : !torch.tuple<int, int> -> !torch.int, !torch.int\n"
" %2 = call @__torch__.torch_mlir.jit_ir_importer.build_tools.library_generator.is_complex_dtype(%1#1) : (!torch.int) -> !torch.bool\n"
" %3 = torch.prim.If %2 -> (!torch.bool) {\n"
" %7 = torch.aten.__isnot__ %arg9, %none : !torch.optional<bool>, !torch.none -> !torch.bool\n"
" torch.prim.If.yield %7 : !torch.bool\n"
" } else {\n"
" torch.prim.If.yield %false : !torch.bool\n"
" }\n"
" %4 = torch.prim.If %3 -> (!torch.bool) {\n"
" %7 = torch.prim.unchecked_cast %arg9 : !torch.optional<bool> -> !torch.bool\n"
" torch.prim.If.yield %7 : !torch.bool\n"
" } else {\n"
" torch.prim.If.yield %false : !torch.bool\n"
" }\n"
" %5:2 = torch.prim.If %4 -> (!torch.bool, !torch.int) {\n"
" torch.prim.If.yield %true, %1#1 : !torch.bool, !torch.int\n"
" } else {\n"
" %7 = func.call @__torch__.torch_mlir.jit_ir_importer.build_tools.library_generator.is_complex_dtype(%1#1) : (!torch.int) -> !torch.bool\n"
" %8 = torch.prim.If %7 -> (!torch.bool) {\n"
" %11 = torch.aten.__isnot__ %arg9, %none : !torch.optional<bool>, !torch.none -> !torch.bool\n"
" torch.prim.If.yield %11 : !torch.bool\n"
" } else {\n"
" torch.prim.If.yield %false : !torch.bool\n"
" }\n"
" %9 = torch.prim.If %8 -> (!torch.bool) {\n"
" %11 = torch.prim.unchecked_cast %arg9 : !torch.optional<bool> -> !torch.bool\n"
" %12 = torch.aten.ne.bool %11, %true : !torch.bool, !torch.bool -> !torch.bool\n"
" torch.prim.If.yield %12 : !torch.bool\n"
" } else {\n"
" torch.prim.If.yield %false : !torch.bool\n"
" }\n"
" %10:2 = torch.prim.If %9 -> (!torch.bool, !torch.int) {\n"
" %11 = torch.aten.eq.int %1#1, %int8 : !torch.int, !torch.int -> !torch.bool\n"
" %12:2 = torch.prim.If %11 -> (!torch.bool, !torch.int) {\n"
" torch.prim.If.yield %true, %int5 : !torch.bool, !torch.int\n"
" } else {\n"
" %13 = torch.aten.eq.int %1#1, %int9 : !torch.int, !torch.int -> !torch.bool\n"
" %14:2 = torch.prim.If %13 -> (!torch.bool, !torch.int) {\n"
" torch.prim.If.yield %true, %int6 : !torch.bool, !torch.int\n"
" } else {\n"
" %15 = torch.aten.eq.int %1#1, %int10 : !torch.int, !torch.int -> !torch.bool\n"
" %16:2 = torch.prim.If %15 -> (!torch.bool, !torch.int) {\n"
" torch.prim.If.yield %true, %int7 : !torch.bool, !torch.int\n"
" } else {\n"
" torch.prim.If.yield %false, %0 : !torch.bool, !torch.int\n"
" }\n"
" torch.prim.If.yield %16#0, %16#1 : !torch.bool, !torch.int\n"
" }\n"
" torch.prim.If.yield %14#0, %14#1 : !torch.bool, !torch.int\n"
" }\n"
" torch.prim.If.yield %12#0, %12#1 : !torch.bool, !torch.int\n"
" } else {\n"
" %11 = func.call @__torch__.torch_mlir.jit_ir_importer.build_tools.library_generator.is_float_dtype(%1#1) : (!torch.int) -> !torch.bool\n"
" %12 = torch.prim.If %11 -> (!torch.bool) {\n"
" %15 = torch.aten.__isnot__ %arg9, %none : !torch.optional<bool>, !torch.none -> !torch.bool\n"
" torch.prim.If.yield %15 : !torch.bool\n"
" } else {\n"
" torch.prim.If.yield %false : !torch.bool\n"
" }\n"
" %13 = torch.prim.If %12 -> (!torch.bool) {\n"
" %15 = torch.prim.unchecked_cast %arg9 : !torch.optional<bool> -> !torch.bool\n"
" torch.prim.If.yield %15 : !torch.bool\n"
" } else {\n"
" torch.prim.If.yield %false : !torch.bool\n"
" }\n"
" %14:2 = torch.prim.If %13 -> (!torch.bool, !torch.int) {\n"
" %15 = torch.aten.eq.int %1#1, %int5 : !torch.int, !torch.int -> !torch.bool\n"
" %16:2 = torch.prim.If %15 -> (!torch.bool, !torch.int) {\n"
" torch.prim.If.yield %true, %int8 : !torch.bool, !torch.int\n"
" } else {\n"
" %17 = torch.aten.eq.int %1#1, %int6 : !torch.int, !torch.int -> !torch.bool\n"
" %18:2 = torch.prim.If %17 -> (!torch.bool, !torch.int) {\n"
" torch.prim.If.yield %true, %int9 : !torch.bool, !torch.int\n"
" } else {\n"
" %19 = torch.aten.eq.int %1#1, %int7 : !torch.int, !torch.int -> !torch.bool\n"
" %20:2 = torch.prim.If %19 -> (!torch.bool, !torch.int) {\n"
" torch.prim.If.yield %true, %int10 : !torch.bool, !torch.int\n"
" } else {\n"
" torch.prim.If.yield %false, %0 : !torch.bool, !torch.int\n"
" }\n"
" torch.prim.If.yield %20#0, %20#1 : !torch.bool, !torch.int\n"
" }\n"
" torch.prim.If.yield %18#0, %18#1 : !torch.bool, !torch.int\n"
" }\n"
" torch.prim.If.yield %16#0, %16#1 : !torch.bool, !torch.int\n"
" } else {\n"
" %15 = func.call @__torch__.torch_mlir.jit_ir_importer.build_tools.library_generator.is_float_dtype(%1#1) : (!torch.int) -> !torch.bool\n"
" %16 = torch.prim.If %15 -> (!torch.bool) {\n"
" %19 = torch.aten.__isnot__ %arg9, %none : !torch.optional<bool>, !torch.none -> !torch.bool\n"
" torch.prim.If.yield %19 : !torch.bool\n"
" } else {\n"
" torch.prim.If.yield %false : !torch.bool\n"
" }\n"
" %17 = torch.prim.If %16 -> (!torch.bool) {\n"
" %19 = torch.prim.unchecked_cast %arg9 : !torch.optional<bool> -> !torch.bool\n"
" %20 = torch.aten.ne.bool %19, %true : !torch.bool, !torch.bool -> !torch.bool\n"
" torch.prim.If.yield %20 : !torch.bool\n"
" } else {\n"
" torch.prim.If.yield %false : !torch.bool\n"
" }\n"
" %18:2 = torch.prim.If %17 -> (!torch.bool, !torch.int) {\n"
" torch.prim.If.yield %true, %1#1 : !torch.bool, !torch.int\n"
" } else {\n"
" %19 = func.call @__torch__.torch_mlir.jit_ir_importer.build_tools.library_generator.is_integer_dtype(%1#1) : (!torch.int) -> !torch.bool\n"
" %20:2 = torch.prim.If %19 -> (!torch.bool, !torch.int) {\n"
" torch.prim.If.yield %true, %int9 : !torch.bool, !torch.int\n"
" } else {\n"
" torch.prim.If.yield %false, %0 : !torch.bool, !torch.int\n"
" }\n"
" torch.prim.If.yield %20#0, %20#1 : !torch.bool, !torch.int\n"
" }\n"
" torch.prim.If.yield %18#0, %18#1 : !torch.bool, !torch.int\n"
" }\n"
" torch.prim.If.yield %14#0, %14#1 : !torch.bool, !torch.int\n"
" }\n"
" torch.prim.If.yield %10#0, %10#1 : !torch.bool, !torch.int\n"
" }\n"
" %6 = torch.prim.If %5#0 -> (!torch.int) {\n"
" torch.prim.If.yield %5#1 : !torch.int\n"
" } else {\n"
" torch.prim.RaiseException %str, %none : !torch.str, !torch.none\n"
" torch.prim.If.yield %0 : !torch.int\n"
" }\n"
" return %6 : !torch.int\n"
" }\n"
" func.func @\"__torch_mlir_dtype_fn.aten.fft_ifft\"(%arg0: !torch.tuple<int, int>, %arg1: !torch.optional<int>, %arg2: !torch.int, %arg3: !torch.optional<str>) -> !torch.int {\n"
" %none = torch.constant.none\n"
" %str = torch.constant.str \"AssertionError: Unsupported dtype\"\n"
Expand Down
Loading
Loading