Merge remote-tracking branch 'origin/move_internal_changes' into sam2

Author: parthchadha

mlir-tensorrt/compiler/include/mlir-tensorrt/Conversion/Passes.td

@@ -25,6 +25,7 @@ include "mlir/Pass/PassBase.td"
 //===----------------------------------------------------------------------===//
 // StablehloToTensorRT
 //===----------------------------------------------------------------------===//
+
 #ifdef MLIR_TENSORRT_ENABLE_HLO
 def ConvertStablehloToTensorRTPass : Pass<"convert-stablehlo-to-tensorrt"> {
   let summary = "Convert Stable HLO dialect to TensorRT dialect";
@@ -44,7 +45,30 @@ def ConvertStablehloToTensorRTPass : Pass<"convert-stablehlo-to-tensorrt"> {
     "target TensorRT version for conversion">
   ];
 }
+#endif // MLIR_TENSORRT_ENABLE_HLO
+
+//===----------------------------------------------------------------------===//
+// ChloToStableHloExt
+//===----------------------------------------------------------------------===//
+
+#ifdef MLIR_TENSORRT_ENABLE_HLO
+def ConvertChloToStableHloExtPass : Pass<"convert-chlo-to-stablehlo-ext"> {
+  let summary = "Convert specific CHLO operations to stablehlo";
+  let description = [{
+    This pass converts a CHLO operations to StableHlo while also allowing
+    for some CHLO operations to be preserved (see options).
+  }];
+  let dependentDialects = [
+    "::mlir::stablehlo::StablehloDialect"
+  ];
 
+  let options = [
+    Option<"preserveErf", "preserve-erf", "bool", "true",
+      "do not convert chlo.erf ops">,
+    Option<"preserveTopK", "preserve-topk", "bool", "true",
+      "do not convert chlo.topk ops">,
+  ];
+}
 #endif // MLIR_TENSORRT_ENABLE_HLO
 
 //===----------------------------------------------------------------------===//

mlir-tensorrt/compiler/include/mlir-tensorrt/Conversion/StablehloToTensorRT/StablehloToTensorRT.h

@@ -25,6 +25,7 @@
 #define MLIR_TENSORRT_CONVERSION_HLOTOTENSORRT_HLOTOTENSORRT_H
 
 #include "mlir-tensorrt/Conversion/TensorRTCommon/ConvertToTensorRTCommon.h"
+#include "mlir-tensorrt/Dialect/StableHloExt/Utils/GatherScatterUtils.h"
 #include "mlir/IR/PatternMatch.h"
 
 namespace mlir {
@@ -33,7 +34,8 @@ class ConversionTarget;
 // Collection of rewrite patterns for lowering of Stable HLO to TensorRT
 // dialect.
 void populateStablehloToTensorRtConversionPattern(
-    TensorRTTypeConverter &typeConverter, RewritePatternSet &patterns);
+    TensorRTTypeConverter &typeConverter, RewritePatternSet &patterns,
+    ShapeInfoCallbacks shapeInfoCallbacks = {});
 
 /// Populate patterns for convert Chlo ops to TensorRT ops.
 void populateChloToTensorRtLegalityAndPatterns(

mlir-tensorrt/compiler/include/mlir-tensorrt/Dialect/Plan/Analysis/BoundsAnalysis.h

@@ -74,6 +74,8 @@ class BoundsArray {
   static BoundsArray fromIntegerValueBounds(unsigned bitwidth,
                                             ArrayRef<int64_t> min,
                                             ArrayRef<int64_t> max);
+  static BoundsArray fromIntegerValueBounds(ArrayRef<llvm::APInt> min,
+                                            ArrayRef<llvm::APInt> max);
 
   /// For the given tensor-typed value, return the most conservative bounds for
   /// the shape of `v`. For each unknown dimension of the shape of `v` the

mlir-tensorrt/compiler/include/mlir-tensorrt/Dialect/Plan/IR/PlanAttributes.td

@@ -24,7 +24,8 @@ def Plan_HostClusterKindAttr : Plan_Attr<"HostClusterKind", "host_cluster",
 }
 
 
-def Plan_BoundsAttr : Plan_Attr<"Bounds", "bounds">{
+def Plan_BoundsAttr : Plan_Attr<"Bounds", "bounds", [
+    DeclareAttrInterfaceMethods<TensorBoundsAttrInterface>]>{
   let parameters = (ins
       EnumParameter<Plan_BoundsKind>:$kind,
       OptionalParameter<"DenseI64ArrayAttr">:$min_shape,
@@ -46,17 +47,17 @@ def Plan_BoundsAttr : Plan_Attr<"Bounds", "bounds">{
 
   let extraClassDeclaration = [{
     /// Returns true if this bounds is for shape dimension extents.
-    bool isShapeBound() {
+    bool isShapeBound() const {
       return getKind() == BoundsKind::Shape;
     }
 
     /// Returns true if this bounds is a 'none' bounds kind.
-    bool isNone() {
+    bool isNone() const {
       return getKind() == BoundsKind::None;
     }
 
     /// Returns true if this bounds is for values of a tensor.
-    bool isValueBound() {
+    bool isValueBound() const {
       return getKind() == BoundsKind::Value;
     }
 

mlir-tensorrt/compiler/include/mlir-tensorrt/Dialect/Plan/IR/PlanDialect.td

@@ -15,6 +15,8 @@ def Plan_Dialect : Dialect {
 
   }];
   let cppNamespace = "::mlir::plan";
+  
+  let hasRegionArgAttrVerify = 1;
 
   let extraClassDeclaration = [{
 
@@ -72,6 +74,18 @@ def Plan_Dialect : Dialect {
       (addExtensionOperation<Ops>(), ...);
     }
 
+    /// Return the name of the function arg/result attributes that encode
+    /// host tensor value bounds. It should have a type `plan::BoundsAttr`.
+    static StringRef getValueBoundsAttrName() {
+      return "plan.value_bounds";
+    }
+
+    /// Return the name of the function arg/result attributes that encode
+    /// the shape bounds. It should have a type `plan::BoundsAttr`.
+    static StringRef getShapeBoundsAttrName() {
+      return "plan.shape_profile";
+    }
+
   private:
     ::llvm::StringMap<AttrParsingHook> attrParsingHooks;
     ::llvm::DenseMap<::mlir::TypeID, AttrPrintingHook> attrPrintingHooks;

mlir-tensorrt/compiler/include/mlir-tensorrt/Dialect/Plan/IR/PlanInterfaces.td

@@ -3,6 +3,34 @@
 
 include "mlir/IR/OpBase.td"
 
+//===----------------------------------------------------------------------===//
+// TensorBoundsAttrInterface
+//===----------------------------------------------------------------------===//
+
+def TensorBoundsAttrInterface : AttrInterface<"TensorBoundsAttrInterface"> {
+  let cppNamespace = "::mlir::plan";
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/"Return the shape bounds associated with the attribute",
+      /*retTy=*/"LogicalResult",
+      "getShapeBounds",
+      (ins "llvm::SmallVectorImpl<int64_t> &":$min,
+           "llvm::SmallVectorImpl<int64_t> &":$max),
+      /*body=*/"",
+      ""
+    >,
+    InterfaceMethod<
+      /*desc=*/"Return the integer value bounds associated with the attribute",
+      /*retTy=*/"LogicalResult",
+      "getIntegerValueBounds",
+      (ins "llvm::SmallVectorImpl<llvm::APInt> &":$min,
+           "llvm::SmallVectorImpl<llvm::APInt> &":$max),
+      /*body=*/"",
+      ""
+    >
+  ];
+}
+
 //===----------------------------------------------------------------------===//
 // ClusterKindInterface
 //===----------------------------------------------------------------------===//

mlir-tensorrt/compiler/include/mlir-tensorrt/Dialect/StableHloExt/Transforms/Passes.td

@@ -202,32 +202,11 @@ def LowerSpecialCustomCalls : Pass<"stablehlo-ext-lower-special-custom-calls"> {
 }
 
 //===----------------------------------------------------------------------===//
-// StablehloInputPreprocessingPass
+// CanonicalizeConvolutionPass
 //===----------------------------------------------------------------------===//
 
-def StablehloInputPreprocessingPass : Pass<"tensorrt-stablehlo-input-preprocessing"> {
-  let summary = "Prepares Stable HLO dialect operations for conversion to TensorRT";
-
-  let description = [{
-   This pass contains a set of patterns for simplifying or transforming Stable HLO
-   input IR so that conversion to the TensorRT dialect is more straightforward.
-
-   In particular:
-
-    - Simplify certain patterns commonly found in IR emitted for JAX programs
-      but not covered by existing Stable HLO canonicalizations/transforms.
-
-    - Prepare convolutions to be NCHW/FCRS format and have at least two
-      "spatial" dimensions.
-
-    - Canonicalize `stablehlo.scatter` operations so that they can be converted to
-      `tensorrt.scatter` in a straightforward manner.
-  }];
-
-  let dependentDialects = [
-    "::mlir::tensor::TensorDialect",
-    "::mlir::stablehlo::StablehloDialect"
-  ];
+def CanonicalizeConvolutionPass : Pass<"stablehlo-ext-canonicalize-convolution"> {
+  let summary = "Canonicalizes stablehlo convolution operations";
 }
 
 

mlir-tensorrt/compiler/include/mlir-tensorrt/Dialect/StableHloExt/Utils/GatherScatterUtils.h

@@ -26,6 +26,8 @@
 #ifndef MLIR_TENSORRT_DIALECT_STABLEHLOEXT_UTILS_GATHERSCATTERUTILS_H
 #define MLIR_TENSORRT_DIALECT_STABLEHLOEXT_UTILS_GATHERSCATTERUTILS_H
 
+#include "mlir-tensorrt/Utils/ShapeInfo.h"
+#include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 #include "mlir/IR/Value.h"
 #include <optional>
 
@@ -34,6 +36,7 @@ namespace mlir {
 class OpBuilder;
 
 namespace stablehlo {
+class DynamicGatherOp;
 class GatherOp;
 class ScatterOp;
 
@@ -124,6 +127,15 @@ namespace stablehlo_ext {
 std::optional<int64_t>
 isSingleDimSimpleGatherWithImplicitIndexDim(stablehlo::GatherOp op);
 
+/// Returns the "gather dimension" if `op` is a 'simple, single dimension'
+/// gather op with implicit index vector dimension (see above for definition).
+/// This version works for `stablehlo.dynamic_gather` using pattern matching
+/// against the expected canonical form when the operand shape along some
+/// "offset dimensions" is dynamic.
+std::optional<int64_t> isSingleDimSimpleGatherWithImplicitIndexDim(
+    stablehlo::DynamicGatherOp op,
+    const ShapeInfoCallbacks &shapeInfoCallbacks);
+
 /// Returns the "gather dimension" if `op` is a 'simple, single dimension'
 /// gather op with explicit size-1 index vector dimension (see above for
 /// definition).
@@ -138,6 +150,21 @@ bool isSimpleLeadingMultiDimGather(stablehlo::GatherOp op);
 /// gather' (see definition above).
 bool isSimpleLeadingMultiDimGatherWithDegenerateDims(stablehlo::GatherOp op);
 
+/// Attempts to construct a `stablehlo.reshape` if result type is statically
+/// shaped, otherwise creates `stablehlo.dynamic_reshape`.
+Value createCollapsingReshape(OpBuilder &b, Location loc, Value input,
+                              ArrayRef<ReassociationIndices> reassociation);
+
+/// Attempts to construct a `stablehlo.reshape` if `resultType` is statically
+/// shaped, otherwise creates a `stablehlo.dynamic_reshape`.
+Value createExpandingReshape(OpBuilder &b, Location loc,
+                             RankedTensorType resultType, Value input,
+                             ArrayRef<ReassociationIndices> reassociation);
+
+/// Returns true if the `scatterOp` has a configuration that corresponds to the
+/// ONNX ScatterNd operation semantic.
+bool isCanonicalScatterNd(stablehlo::ScatterOp scatterOp);
+
 //===----------------------------------------------------------------------===//
 // Code below this point was adapted from the MLIR-HLO project (part of OpenXLA
 // project) `xla/mlir_hlo/mhlo/utils/mhlo_scatter_gather_utils.h` and has the
@@ -155,10 +182,6 @@ bool isSimpleLeadingMultiDimGatherWithDegenerateDims(stablehlo::GatherOp op);
 // - scatter_dims_to_operand_dims is [0, 1, ...]
 bool isCanonicalScatter(stablehlo::ScatterOp scatterOp);
 
-/// Returns true if the `scatterOp` has a configuration that corresponds to the
-/// ONNX ScatterNd operation semantic.
-bool isCanonicalScatterNd(stablehlo::ScatterOp scatterOp);
-
 // Checks if the gather has the following characteristics:
 // - start_indices is a two-dimensional tensor
 // - index_vector_dim is 1

mlir-tensorrt/compiler/include/mlir-tensorrt/Pipelines/StableHloInputPipelines.h

@@ -32,12 +32,19 @@ class OpPassManager;
 struct StableHloInputOptions {
   /// Whether to lower Stablehlo control flow ops to SCF dialect ops.
   bool legalizeControlFlowToSCF = false;
-  /// Whether to lower chlo.erf into primitive stablehlo operations.
-  bool legalizeChloErfToStablehlo = false;
+
+  /// Whether to preserve 'chlo.erf' ops or lower them to 'stablehlo' ops.
+  /// By default, we preserve since it has a 1-1 correspondence with a TensorRT
+  /// op.
+  bool preserveChloErf = true;
+
+  /// Whether to preserve 'chlo.top_k' ops or lower them to 'stablehlo' ops.
+  /// By default, we preserve since it has a 1-1 correspondence with a TensorRT
+  /// op.
+  bool preserveChloTopK = true;
+
   /// Whether to disable running the inliner.
   bool disableInliner = false;
-  /// Whether to lower chlo to stablehlo.
-  bool convertChloToStablehlo = false;
 };
 
 /// Construct a pipeline for preprocessing StableHLO IR to convert it into the

mlir-tensorrt/compiler/include/mlir-tensorrt/Utils/ShapeInfo.h

@@ -0,0 +1,96 @@
+//===- ShapeInfo.h ---------------------------------------------*- C++ -*-===//
+//
+// SPDX-FileCopyrightText: Copyright 2024 NVIDIA CORPORATION & AFFILIATES.
+// All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+//===----------------------------------------------------------------------===//
+///
+/// Declarations for callback types are used to abstract away how to infer
+/// shape knowledge from a pass or transformation. For example, a pass operating
+/// on StableHlo IR may need to check whether the *values* of tensor A represent
+/// the actual *shape* of tensor B, whose shape may not be known statically at
+/// compile time.
+///
+/// The specific mechanism that one may use to determine the validity of a
+/// specific proposition like the example above (which must be reported as
+/// "unknown", "true", or "false") may depend on the context. In the case
+/// of the StableHlo example above, we could try to naively pattern match
+/// whether tensor A is the result of `stablehlo.concat` of appropriate
+/// `stablehlo.get_dimensions_size %A, dim = ...` results. In other cases,
+/// we may have access to an analysis that assists with more robustly
+/// checking the proposition.
+///
+/// This file just contains callback types that a Pass or rewrite/transform can
+/// accept as a parameter, allowing the creator or caller to hand in a
+/// particular implementation.
+///
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_TENSORRT_UTILS_SHAPEINFO
+#define MLIR_TENSORRT_UTILS_SHAPEINFO
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/Value.h"
+
+namespace mlir {
+
+/// TensorElementValue identifies a particular scalar element value of a
+/// statically-shaped tensor.
+struct TensorElementValue {
+  TensorElementValue(Value value, ArrayRef<int64_t> coord);
+
+  TypedValue<RankedTensorType> getTensor() const { return tensor; }
+  int64_t getLinearIndex() const { return linearIndex; }
+
+  /// A value of type (must be statically-shaped) RankedTensorType.
+  TypedValue<RankedTensorType> tensor;
+
+  /// The linear coordinate of the value.
+  int64_t linearIndex;
+};
+
+/// TensorShapeDimExtent identifies a (potentially dynamically shaped) size
+/// of a particular dimension of a tensor's shape.
+struct TensorShapeDimExtent {
+  TensorShapeDimExtent(Value value, int64_t dim);
+
+  std::optional<int64_t> getConstantSize() const;
+
+  /// A value of type  RankedTensorType.
+  TypedValue<RankedTensorType> tensor;
+
+  /// The dimension.
+  int64_t dim;
+};
+
+struct ShapeInfoCallbacks {
+  // Check whether 'tensorElementValue' is provably equivalent to
+  // `tensorShapeDimExtent`. Returning 'nullopt' means "unknown", true means
+  // "equal", false means "not equal".
+  std::function<std::optional<bool>(TensorElementValue tensorElementValue,
+                                    TensorShapeDimExtent tensorShapeDimExtent)>
+      isElementValueEqualToShapeDimExtent;
+
+  // Check whether 'tensorElementValue' is provably equivalent to the given
+  // static value. Returning 'nullopt' means "unknown", true means "equal",
+  // false means "not equal".
+  std::function<std::optional<bool>(TensorElementValue tensorElementValue,
+                                    Attribute constantValue)>
+      isElementValueEqualToConstant;
+};
+
+} // namespace mlir
+
+#endif // MLIR_TENSORRT_UTILS_SHAPEINFO