Add aten.stft.center and decomposition

include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td

@@ -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,
@@ -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,

include/torch-mlir/Dialect/Torch/Transforms/Passes.h

@@ -118,6 +118,9 @@ std::unique_ptr<OperationPass<ModuleOp>> createRefinePublicReturnPass();
 std::unique_ptr<OperationPass<func::FuncOp>>
 createDecomposeComplexOpsPass(ArrayRef<std::string> legalOps);
 
+std::unique_ptr<OperationPass<func::FuncOp>>
+createSimplifyListGeneratingLoopsPass();
+
 std::unique_ptr<OperationPass<func::FuncOp>> createScalarizeShapesPass();
 
 std::unique_ptr<OperationPass<func::FuncOp>> createRecomposeComplexOpsPass();

include/torch-mlir/Dialect/Torch/Transforms/Passes.td

@@ -235,6 +235,18 @@ def DecomposeComplexOps : Pass<"torch-decompose-complex-ops", "func::FuncOp"> {
   }];
 }
 
+def SimplifyListGeneratingLoops : Pass<"torch-simplify-list-generating-loops", "func::FuncOp"> {
+  let summary = "Simplify loops that generate lists";
+  let constructor = "mlir::torch::Torch::createSimplifyListGeneratingLoopsPass()";
+  let description = [{
+  }];
+  let description = [{
+    Unroll loops that generate lists and greedily simplify list operations until
+    either they have been folded into a single prim.ListConstruct or maxIterations
+    are reached.
+  }];
+}
+
 def ScalarizeShapes : Pass<"torch-scalarize-shapes", "func::FuncOp"> {
   let summary = "Takes common shape computation operations and scalarizes them.";
   let constructor = "mlir::torch::Torch::createScalarizeShapesPass()";

include/torch-mlir/Dialect/Torch/Utils/Utils.h

@@ -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`.

lib/Conversion/TorchOnnxToTorch/DefaultDomainQtoZ.cpp

@@ -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());
@@ -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]}),