Integrate internal changes.

This PR integrates the following internal changes into OSS.

commit c6f9d4b57ae793b08ddf415529a83e2aca6df5ce
Author: Sagar Shelke <[email protected]>
Date:   Tue Jun 10 14:27:27 2025 -0700

    [Dialect/Plan] Add constant foldable subgraph analysis pass.

    This MR adds a generic pass for constant foldable subgraph analysis pass. This pass
    implements the following three things,
    1. A forward data flow analysis to find constant foldable ops.
    2. Use clustering to find constant foldable subgraphs.
    3. Outline constant foldable subgraphs to `func.func` with
    `plan.constant_foldable` attributes.
    Pass option `skipClustering` allows user to extend default pass behavior
    and skip certain ops from outlining.

    MLIR test cases are added including region ops tests.

commit 5fad6480597ee15bdb4a8c04c033b7ef6beff2d0
Author: Christopher Bate <[email protected]>
Date:   Tue Jun 10 19:27:04 2025 +0000

    NFC: remove unused functions

commit 7eb853f52525756e399a7fc6fdc4ab25244d6f94
Author: Chris Bate <[email protected]>
Date:   Tue Jun 10 12:06:51 2025 -0700

    [compiler] Create InferTensorValueRangeInterface

    Previously, our TensorValueBoundsAnalysis was only able to analyze values
    that were either constant or were produced by the result of `plan.with_values`.
    This worked because we rely on `plan.with_values` to link integer range
    analysis and the TensorValueBoundsAnalysis. However, after clustering there
    can occur edge cases where the `plan.with_values` op is not present.
    Furthermore, eventually we would like to avoid creating `plan.with_values`
    ops in the first place since they temporarily bloat the IR.

    This commit adds a new interface, InferTensorValueRangeInterface, that allows
    operations to participate in the TensorValueBoundsAnalysis. The logic
    for `plan.with_values` is implemented using the interface, and this change
    adds an implementation and test for `stablehlo.convert`. Besides this
    new behavior for `stablehlo.convert`, this change is a pure refactor and
    no other new behavior is added.

commit d49fe69e08d79db07bf9f801427c97f6bfbdfcda
Author: Christopher Bate <[email protected]>
Date:   Fri Jun 6 17:03:45 2025 +0000

    NFC: Move `python` under `integrations/python`

    Moves the `python` directory under `integrations/python` to better
    reflect that, like PJRT, the Python packages are a downstream consumer of
    the MLIR-TensorRT compiler/runtime libraries.

commit ef5490368b2235e80bb21a6aca39db43d24f1c47
Author: Sagar Shelke <[email protected]>
Date:   Sat Jun 7 14:27:21 2025 -0700

    [tensorrt] Implement inliner inferface for TensorRT dialect

    This MR implements `DialectInlinerInterface` for the TensorRT dialect.
    All ops in TensorRT dialect are pure and thus can be inlined, except `module`,
    `call` and `call_alloc`. Upstream `--inline` pass doesn't inline operation with
     callop interface (`call` and `call_alloc`) because these operations in TensorRT dialect
    are referring to callee in different symbol table comapred to their own.

    MLIR valid and invalid test is added.

commit cd56aa6a511e2091fcd86106f20d27ff3673db75
Author: Christopher Bate <[email protected]>
Date:   Wed Jun 11 14:26:03 2025 +000

    Fix build with BUILD_SHARED_LIBS=ON

    The new InferTensorValueRangeInterface was used without correctly
    specifying the library dependency the PlanIR and StablehloExtIR
    libraries.

GitOrigin-RevId: cd56aa6a511e2091fcd86106f20d27ff3673db75

Author: christopherbate

mlir-tensorrt/CMakeLists.txt

@@ -249,4 +249,4 @@ include_directories(${CMAKE_CURRENT_LIST_DIR}/tensorrt/include)
 include_directories(${CMAKE_CURRENT_BINARY_DIR}/tensorrt/include)
 
 add_subdirectory(compiler)
-add_subdirectory(python)
+add_subdirectory(integrations)

mlir-tensorrt/build_tools/docker/Dockerfile

@@ -74,8 +74,8 @@ EOF
 ARG PYTHON_VERSION=3.10
 ENV PYENV_ROOT="/pyenv"
 ENV PATH="/pyenv/bin:/pyenv/shims:$PATH"
-COPY python/requirements-dev.txt /tmp/requirements-dev.txt
-COPY python/requirements.txt /tmp/requirements.txt
+COPY integrations/python/requirements-dev.txt /tmp/requirements-dev.txt
+COPY integrations/python/requirements.txt /tmp/requirements.txt
 RUN <<EOF
 set -e
 case "${LINUX_DISTRO}" in

mlir-tensorrt/build_tools/scripts/build_wheels.sh

@@ -9,7 +9,7 @@ py_version=${PY_VERSION:-3.10}
 PATH=$PATH:/pyenv/bin
 mkdir -p .private.wheels || true
 pyenv local ${py_version}
-python${py_version} -m pip install -r python/requirements-dev.txt
+python${py_version} -m pip install -r integrations/python/requirements-dev.txt
 
 export DOWNLOAD_TENSORRT_VERSION=${DOWNLOAD_TENSORRT_VERSION:-10.9}
 

mlir-tensorrt/build_tools/scripts/cicd_build.sh

@@ -10,7 +10,7 @@ export ENABLE_ASAN=${ENABLE_ASAN:-OFF}
 export CPM_SOURCE_CACHE=${CPM_SOURCE_CACHE:-/.cache.cpm}
 export CCACHE_DIR=${CCACHE_DIR:-/ccache}
 
-python3 -m pip install -r python/requirements-dev.txt
+python3 -m pip install -r integrations/python/requirements-dev.txt
 
 ccache --zero-stats || true
 rm -rf ${BUILD_DIR}  || true

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

@@ -25,96 +25,13 @@
 #ifndef MLIR_TENSORRT_DIALECT_PLAN_ANALYSIS_BOUNDSANALYSIS
 #define MLIR_TENSORRT_DIALECT_PLAN_ANALYSIS_BOUNDSANALYSIS
 
+#include "mlir-tensorrt/Interfaces/InferTensorValueRangeInterface.h"
 #include "mlir/Analysis/DataFlow/IntegerRangeAnalysis.h"
 #include "mlir/Analysis/DataFlow/SparseAnalysis.h"
-#include "mlir/Interfaces/InferIntRangeInterface.h"
-#include "llvm/Support/raw_ostream.h"
 
 namespace mlir::plan {
 
-//===----------------------------------------------------------------------===//
-// BoundsArray
-//===----------------------------------------------------------------------===//
-
-/// A BoundsArray is simply an array of ConstantIntRanges used to represent
-/// either the bounds on a shape of a tensor-typed SSA value or the bounds
-/// of the element values of a statically shaped integer tensor-typed SSA value.
-/// When it is used to represent the bounds for the value of a tensor, we use
-/// a canonical packed generalized row-major layout mapping from tensor
-/// coordinates to storage index.
-class BoundsArray {
-public:
-  BoundsArray(
-      std::optional<SmallVector<ConstantIntRanges>> value = std::nullopt)
-      : value(std::move(value)) {}
-
-  bool isUninitialized() const { return !value.has_value(); }
-
-  bool operator==(const BoundsArray &rhs) const { return value == rhs.value; }
-
-  ArrayRef<ConstantIntRanges> getValue() const {
-    assert(!isUninitialized());
-    return *value;
-  }
-
-  /// Return the most conservative integer scalar bounds for an dynamic/unknown
-  /// dimension extent.
-  static ConstantIntRanges getMaxDimRange();
-
-  /// Create a BoundsValue from the min/max bounds of shape. Using this method
-  /// ensures that the `value` are created with the correct storage bitwidth
-  /// (an implementation detail of the analysis).
-  static BoundsArray fromShapeBounds(ArrayRef<int64_t> min,
-                                     ArrayRef<int64_t> max);
-
-  /// Create a `BoundsValue` using the given scalar values encoded as int64_t
-  /// values. However, when storing the bounds, use the given bitwidth.
-  /// TODO: remove this when we migrate away from using
-  /// `#tensorrt.shape_profile` for value bounds.
-  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
-  /// `getMaxDimRange()` bound is used.
-  static BoundsArray getMaxRangeForShapeBounds(Value v);
-
-  /// For the given statically shaped integer tensor-typed value, return the
-  /// most conservative bounds for the value of `v`.
-  static BoundsArray getMaxRangeForValueBounds(Value v);
-
-  /// For the given DenseIntElementsAttr, return a corresponding BoudnsValue
-  /// representing constant bounds as indicated by the attribute.
-  static BoundsArray getFromConstantValue(DenseIntElementsAttr attr);
-
-  /// Join two BoundsValues by performing a pointwise union of the integer
-  /// scalar a ranges.
-  static BoundsArray join(const BoundsArray &lhs, const BoundsArray &rhs);
-
-  /// Meet two BoundsValues by performing a pointwise intersection of the
-  /// integer scalar a ranges.
-  static BoundsArray meet(const BoundsArray &lhs, const BoundsArray &rhs);
-
-  /// Print a human-readable representation of the bounds.
-  void print(raw_ostream &os) const;
-
-  /// Return the min/max bounds representation as two DenseElementsAttrs.
-  std::pair<DenseElementsAttr, DenseElementsAttr>
-  getAsElementsAttr(RankedTensorType type) const;
-
-  /// Returns DenseElementsAttr representation if the element ranges are all
-  /// constant (single-value) ranges, otherwise nullopt.
-  std::optional<DenseElementsAttr>
-  getConstantValues(RankedTensorType type) const;
-
-private:
-  std::optional<SmallVector<ConstantIntRanges>> value;
-};
-
-llvm::raw_ostream &operator<<(llvm::raw_ostream &os, const BoundsArray &v);
+using BoundsArray = mlirtrt::compiler::BoundsArray;
 
 //===----------------------------------------------------------------------===//
 // Shape Bounds Analyses

mlir-tensorrt/compiler/include/mlir-tensorrt/Dialect/Plan/IR/CMakeLists.txt

@@ -17,4 +17,6 @@ add_public_tablegen_target(MLIRTensorRTPlanDialectAttributesIncGen)
 set(LLVM_TARGET_DEFINITIONS PlanInterfaces.td)
 mlir_tablegen(PlanAttrInterfaces.h.inc -gen-attr-interface-decls)
 mlir_tablegen(PlanAttrInterfaces.cpp.inc -gen-attr-interface-defs)
+mlir_tablegen(PlanOpInterfaces.h.inc -gen-op-interface-decls)
+mlir_tablegen(PlanOpInterfaces.cpp.inc -gen-op-interface-defs)
 add_public_tablegen_target(MLIRTensorRTPlanDialectAttrInterfacesIncGen)

mlir-tensorrt/compiler/include/mlir-tensorrt/Dialect/Plan/IR/Plan.h

@@ -27,6 +27,7 @@
 #include "mlir-tensorrt-dialect/Interface/TensorKindOpInterface.h"
 #include "mlir-tensorrt/Compiler/Extension.h"
 #include "mlir-tensorrt/Dialect/Plan/IR/PlanInterfaces.h"
+#include "mlir-tensorrt/Interfaces/InferTensorValueRangeInterface.h"
 #include "mlir/Bytecode/BytecodeOpInterface.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinOps.h"

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

@@ -8,6 +8,7 @@ include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/DestinationStyleOpInterface.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
+include "mlir-tensorrt/Interfaces/InferTensorValueRangeInterface.td"
 include "mlir/IR/OpAsmInterface.td"
 
 class Plan_NativeOpTrait<string name,
@@ -433,6 +434,8 @@ def Plan_WithShapeOp : Plan_Op<"with_shape",
 def Plan_WithValuesOp : Plan_Op<"with_values",
     [Pure,
      DeclareOpInterfaceMethods<TensorKindOpInterface>,
+     DeclareOpInterfaceMethods<InferTensorValueRangeInterface,
+      ["inferResultRangesFromOptional"]>,
      AllTypesMatch<["operand", "result"]>]> {
   let summary =
     "Ties a tensor value with index SSA values representing its element values";

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

@@ -258,9 +258,9 @@ def ClusteringPass : Pass<"plan-clustering", "::mlir::ModuleOp"> {
     operations will be compiled.
 
     The kinds of clusters that can be formed and the specific rules for
-    clustering are defined by the clustering configuration specified 
+    clustering are defined by the clustering configuration specified
     by the module's `plan.cluster_kinds` attribute. This is an array of
-    attributes which all implement the 
+    attributes which all implement the
     [ClusterKindAttrInterface](../IR/PlanInterfaces.td).
   }];
 
@@ -585,5 +585,35 @@ def PlanOwnershipBasedBufferDeallocationPass : Pass<
   ];
 }
 
+//===----------------------------------------------------------------------===//
+// PlanOutlineConstantFoldableSubgraphs
+//===----------------------------------------------------------------------===//
+
+def PlanOutlineConstantFoldableSubgraphsPass : Pass<
+      "plan-outline-constant-foldable-subgraphs",
+      "::mlir::ModuleOp"> {
+  let summary = "Analyze and outline constant foldable subgraphs";
+
+  let description = [{
+    This pass implements forward dataflow analysis (named `SparseConstantFoldabilityAnalysis`)
+    to find out constant foldable ops. This analysis, unlike upstream
+    `ConstantPropagationAnalysis` is very simple and works only for pure ops.
+    If all operands of an operation are constant foldable, all results are marked
+    as constant foldable.
+    Constant foldability analysis is then used along with clustering to
+    find constant foldable subgraphs. These constant foldable subgraphs are
+    finally outlined to a private function with `plan.constant_foldable` attribute.
+  }];
+
+  let options = [
+    Option<"skipClustering", "skip-clustering",
+           "std::function<bool(Operation*)>", /*default=*/"nullptr",
+           "This option enables user to extend default pass behavior and skip "
+           "more ops from clustering. If this method returns true, `op` is not "
+           "clustered. When op is not clustered, it is not outlined for constant "
+           "folding. This is helpful in avoiding clustering of ops that can't be "
+           "run e2e at compile time, in the workflow of user's choice.">,
+  ];
+}
 
 #endif // MLIR_TENSORRT_DIALECT_PLAN_TRANSFORMS_PASSES_TD

mlir-tensorrt/compiler/include/mlir-tensorrt/Dialect/StablehloExt/IR/StableHloExt.h

@@ -34,6 +34,10 @@ void registerTensorKindOpInterfaceExternalModels(DialectRegistry &registry);
 /// Register StableHlo op implementations for ReifyRankedShapedTypeOpInterface.
 void registerTypeInferenceExternalModels(DialectRegistry &registry);
 
+/// Register StableHlo op implementations for InferTensorValueRangeInterface.
+void registerInferTensorValueRangeInterfaceExternalModels(
+    DialectRegistry &registry);
+
 } // namespace mlir::stablehlo
 
 #endif // MLIR_TENSORRT_DIALECT_STABLEHLOEXT_IR_STABLEHLOEXT_H