[linalg] : Use (-)realmax instead of (-)inf to avoid usage of non-finites. (#4363)

This change replaces usage of non-finite value `inf` with finite value
`realmax` for init value of various max/min operations depending on
whether `allow-non-finites` option is disabled through the
`fx.export_and_import` API -- no change in semantics of the ops. The
default behavior of emitting `inf` is preserved.

Author: sahas3

include/torch-mlir/Conversion/Passes.td

@@ -103,6 +103,11 @@ def ConvertTorchToLinalg : Pass<"convert-torch-to-linalg", "func::FuncOp"> {
     compelling for modeling effects more broadly.
   }];
   let constructor = "mlir::torch::createConvertTorchToLinalgPass()";
+  let options = [
+    Option<"allowNonFinites", "allow-non-finites",
+            "bool", /*default=*/"true",
+            "When enabled (default), some ops may emit non-finites, for example, max pooling may compare values to an initial value of `-inf`. When disabled, non-finites will be replaced with the closest finite value for a given dtype.">,
+  ];
 }
 
 def ConvertTorchToTensor : Pass<"convert-torch-to-tensor", "func::FuncOp"> {
@@ -141,6 +146,11 @@ def ConvertTorchToTMTensor : Pass<"convert-torch-to-tmtensor", "func::FuncOp"> {
     pass also makes use of TMTensor Dialect, which the former one doesn't.
   }];
   let constructor = "mlir::torch::createConvertTorchToTMTensorPass()";
+  let options = [
+    Option<"allowNonFinites", "allow-non-finites",
+            "bool", /*default=*/"true",
+            "When enabled (default), some ops may emit non-finites, for example, max pooling may compare values to an initial value of `-inf`. When disabled, non-finites will be replaced with the closest finite value for a given dtype.">,
+  ];
 }
 
 def ConvertTorchConversionToMLProgram : Pass<"convert-torch-conversion-to-mlprogram", "ModuleOp"> {
@@ -168,6 +178,9 @@ def ConvertTorchToStablehlo : Pass<"convert-torch-to-stablehlo", "func::FuncOp">
     // are unlikely to exceed the range of i32(4GiB)
     Option<"enableI32Index", "enable-i32-index", "bool", /*default=*/"false",
            "Enable truncate index from i64 to i32(unsafely)">,
+    Option<"allowNonFinites", "allow-non-finites",
+            "bool", /*default=*/"true",
+            "When enabled (default), some ops may emit non-finites, for example, max pooling may compare values to an initial value of `-inf`. When disabled, non-finites will be replaced with the closest finite value for a given dtype.">,
   ];
 }
 #endif

include/torch-mlir/Conversion/TorchOnnxToTorch/Passes.h

@@ -17,7 +17,12 @@
 
 namespace mlir::torch::onnx_c {
 
+#define GEN_PASS_DECL_CONVERTTORCHONNXTOTORCH
+#include "torch-mlir/Conversion/TorchOnnxToTorch/Passes.h.inc"
+
 std::unique_ptr<OperationPass<func::FuncOp>> createTorchOnnxToTorchPass();
+std::unique_ptr<OperationPass<func::FuncOp>>
+createTorchOnnxToTorchPass(bool allowNonFinites);
 
 /// Registers all torch-mlir conversion passes.
 void registerTorchOnnxToTorchPasses();

include/torch-mlir/Conversion/TorchOnnxToTorch/Passes.td

@@ -21,6 +21,11 @@ def ConvertTorchOnnxToTorch : Pass<"convert-torch-onnx-to-torch", "func::FuncOp"
     }];
 
     let constructor = "mlir::torch::onnx_c::createTorchOnnxToTorchPass()";
+    let options = [
+    Option<"allowNonFinites", "allow-non-finites",
+            "bool", /*default=*/"true",
+            "When enabled (default), some ops may emit non-finites, for example, max pooling may compare values to an initial value of `-inf`. When disabled, non-finites will be replaced with the closest finite value for a given dtype.">,
+  ];
 }
 
 #endif  // TORCHMLIR_CONVERSION_TORCHONNX_TO_TORCH_PASSES

include/torch-mlir/Conversion/TorchOnnxToTorch/Patterns.h

@@ -422,6 +422,10 @@ struct OpBinder {
   Operation *op;
 };
 
+struct OnnxTorchToTorchOptions {
+  bool allowNonFinites = true;
+};
+
 /// We use a single pattern per ONNX domain to handle all named custom
 /// ops.
 /// This allows us to avoid the n^2 problem on pattern application by
@@ -431,19 +435,25 @@ struct OpBinder {
 class OnnxCustomOpConversionPattern
     : public OpConversionPattern<Torch::OperatorOp> {
 public:
-  using HandlerFn = LogicalResult (*)(OpBinder binder,
-                                      ConversionPatternRewriter &rewriter);
+  using HandlerFn =
+      std::function<LogicalResult(OpBinder, ConversionPatternRewriter &)>;
+
+  using HandlerFnWithOptions =
+      LogicalResult (*)(OpBinder binder, ConversionPatternRewriter &rewriter,
+                        const OnnxTorchToTorchOptions &options);
+
   struct HandlerReg {
     HandlerReg(HandlerFn callback, int64_t sinceVersion)
-        : callback(callback), sinceVersion(sinceVersion) {}
+        : callback(std::move(callback)), sinceVersion(sinceVersion) {}
     HandlerFn callback;
     int64_t sinceVersion;
   };
 
   OnnxCustomOpConversionPattern(MLIRContext *context, std::string domainPrefix,
-                                int64_t domainVersion)
+                                int64_t domainVersion,
+                                const OnnxTorchToTorchOptions &options)
       : OpConversionPattern(context), domainPrefix(std::move(domainPrefix)),
-        domainVersion(domainVersion) {
+        domainVersion(domainVersion), options(options) {
     // Onnx lowerings could produce other Onnx operations during the rewrite.
     setHasBoundedRewriteRecursion();
   }
@@ -463,19 +473,25 @@ class OnnxCustomOpConversionPattern
   /// Multiple conversions can be registered for the same op, most
   /// commonly differing by their `sinceVersion`.
   void onOp(StringRef name, int64_t sinceVersion, HandlerFn callback);
+  void onOp(StringRef name, int64_t sinceVersion,
+            HandlerFnWithOptions callback);
+
+  const OnnxTorchToTorchOptions &getOptions() const { return options; }
 
 private:
   std::string domainPrefix;
   int64_t domainVersion;
   DenseMap<StringAttr, SmallVector<HandlerReg, 1>> namedHandlers;
+  OnnxTorchToTorchOptions options;
 };
 
 // Patterns are split into chunks to speed compile time and reduce some
 // contention on the same source files.
 void populateComMicrosoftDomain(OnnxCustomOpConversionPattern &patterns);
 void populateDefaultDomainAtoF(OnnxCustomOpConversionPattern &patterns);
 void populateDefaultDomainGtoP(OnnxCustomOpConversionPattern &patterns);
-void populateDefaultDomainQtoZ(OnnxCustomOpConversionPattern &patterns);
+void populateDefaultDomainQtoZ(OnnxCustomOpConversionPattern &patterns,
+                               const OnnxTorchToTorchOptions &options);
 
 } // namespace mlir::torch::onnx_c
 

include/torch-mlir/Conversion/TorchToLinalg/TorchToLinalg.h

@@ -17,8 +17,16 @@
 
 namespace mlir {
 namespace torch {
+
+#define GEN_PASS_DECL_CONVERTTORCHTOLINALG
+#include "torch-mlir/Conversion/Passes.h.inc"
+
 std::unique_ptr<OperationPass<func::FuncOp>> createConvertTorchToLinalgPass();
-}
+
+std::unique_ptr<OperationPass<func::FuncOp>>
+createConvertTorchToLinalgPass(bool allowNonFinites);
+
+} // namespace torch
 } // namespace mlir
 
 #endif // TORCHMLIR_CONVERSION_ATENTOLINALG_ATENTOLINALG_H

include/torch-mlir/Conversion/TorchToStablehlo/TorchToStablehlo.h

@@ -26,7 +26,8 @@ createConvertTorchToStablehloPass();
 // Convenience wrapper for users who want to pass options as individual
 // parameters
 std::unique_ptr<OperationPass<func::FuncOp>>
-createConvertTorchToStablehloPass(bool enableStaticShape, bool enableI32Index);
+createConvertTorchToStablehloPass(bool enableStaticShape, bool enableI32Index,
+                                  bool allowNonFinites);
 
 } // namespace torch
 } // namespace mlir

include/torch-mlir/Conversion/TorchToTMTensor/TorchToTMTensor.h

@@ -15,8 +15,16 @@
 
 namespace mlir {
 namespace torch {
+
+#define GEN_PASS_DECL_CONVERTTORCHTOTMTENSOR
+#include "torch-mlir/Conversion/Passes.h.inc"
+
 std::unique_ptr<OperationPass<func::FuncOp>> createConvertTorchToTMTensorPass();
-}
+
+std::unique_ptr<OperationPass<func::FuncOp>>
+createConvertTorchToTMTensorPass(bool allowNonFinites);
+
+} // namespace torch
 } // namespace mlir
 
 #endif // TORCHMLIR_CONVERSION_TORCHTOTMTENSOR_TORCHTOTMTENSOR_H

include/torch-mlir/Conversion/Utils/Utils.h

@@ -122,6 +122,9 @@ void getZeroPoint(Value value, Value &zeropoint);
 LogicalResult getQuantizationParams(Value value, Value &zeropoint, Value &scale,
                                     int64_t &axis);
 
+APFloat getFloatInf(mlir::FloatType fpType, bool negative,
+                    bool allowNonFinites);
+
 } // namespace Torch
 } // namespace torch
 } // namespace mlir

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

@@ -66,6 +66,15 @@ struct TorchLoweringPipelineOptions
       *this, "extra-library",
       llvm::cl::desc("Filename of MLIR module for splicing into the abstract "
                      "interpretation library.")};
+
+  Option<bool> allowNonFinites{
+      *this, "allow-non-finites",
+      llvm::cl::desc(
+          "When enabled (default), some ops may emit non-finites, for example, "
+          "max pooling may compare values to an initial value of `-inf`. When "
+          "disabled, non-finites will be replaced with the closest finite "
+          "value for a given dtype."),
+      llvm::cl::init(true)};
 };
 
 /// Creates a pipeline that lowers the object graph IR that is produced by

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

@@ -21,10 +21,22 @@ class ModuleOp;
 
 namespace torch {
 namespace TorchConversion {
+struct LinalgOnTensorsBackendPipelineOptions
+    : public PassPipelineOptions<LinalgOnTensorsBackendPipelineOptions> {
+  Option<bool> allowNonFinites{
+      *this, "allow-non-finites",
+      llvm::cl::desc(
+          "When enabled (default), some ops may emit non-finites, for example, "
+          "max pooling may compare values to an initial value of `-inf`. When "
+          "disabled, non-finites will be replaced with the closest finite "
+          "value for a given dtype."),
+      llvm::cl::init(true)};
+};
 
 /// Creates a pipeline that lowers from the torch backend contract to the
 /// linalg-on-tensors backend contract.
-void createTorchBackendToLinalgOnTensorsBackendPipeline(OpPassManager &pm);
+void createTorchBackendToLinalgOnTensorsBackendPipeline(
+    OpPassManager &pm, const LinalgOnTensorsBackendPipelineOptions &options);
 
 // Do not register the TOSA options if the TOSA target is disabled
 #ifdef TORCH_MLIR_ENABLE_TOSA
@@ -59,6 +71,14 @@ struct StablehloBackendPipelineOptions
       *this, "enable-i32-index",
       llvm::cl::desc("Enable truncate index from i64 to i32(unsafely)"),
       llvm::cl::init(false)};
+  Option<bool> allowNonFinites{
+      *this, "allow-non-finites",
+      llvm::cl::desc(
+          "When enabled (default), some ops may emit non-finites, for example, "
+          "max pooling may compare values to an initial value of `-inf`. When "
+          "disabled, non-finites will be replaced with the closest finite "
+          "value for a given dtype."),
+      llvm::cl::init(true)};
 };
 
 void createTorchBackendToStablehloBackendPipeline(