Move internal changes

GitOrigin-RevId: 8dd523bab228ccfb83aeba17793cb75895d31bc6

Author: Copybara Bot

mlir-tensorrt/compiler/CMakeLists.txt

@@ -31,6 +31,7 @@ add_mlir_tensorrt_compiler_dependency(MLIRNVVMTarget)
 add_mlir_tensorrt_compiler_dependency(MLIRPtrDialect)
 add_mlir_tensorrt_compiler_dependency(MLIRTargetLLVM)
 add_mlir_tensorrt_compiler_dependency(MLIRTensorTransformOps)
+add_mlir_tensorrt_compiler_dependency(MLIREmitCExtDataLayoutImpl)
 
 add_subdirectory(include)
 add_subdirectory(lib)

mlir-tensorrt/compiler/include/mlir-tensorrt/Backends/Host/HostBackend.td

@@ -5,7 +5,7 @@ include "mlir-tensorrt/Dialect/Plan/IR/PlanDialect.td"
 include "mlir-tensorrt/Dialect/Plan/IR/PlanInterfaces.td"
 
 def Plan_HostClusterKindAttr : Plan_Attr<"HostClusterKind", "host_cluster",
-      [DeclareAttrInterfaceMethods<ClusterKindAttrInterface>]> {
+      [DeclareAttrInterfaceMethods<ClusterKindAttrInterface, ["getDefaultMemorySpace"]>]> {
   let parameters = (ins "int64_t":$benefit);
   let assemblyFormat = "`<` struct(params) `>`";
 }

mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/Extension.h

@@ -73,11 +73,6 @@ class TaskExtensionBase {
   template <typename T, typename... Mods>
   using ListOption = mlir::detail::PassOptions::ListOption<T, Mods...>;
 
-protected:
-  /// Whether this extension is disabled. Should default to false and be
-  /// associated with a flag `--disable-[name]-extension`.
-  bool disabled{false};
-
 private:
   mlir::TypeID typeID;
 

mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/OptionsProviders.h

@@ -21,8 +21,8 @@
 /// Data structures and functions for manipulating compiler options.
 ///
 //===----------------------------------------------------------------------===//
-#ifndef MLIR_TENSORRT_COMPILER_OPTIONS
-#define MLIR_TENSORRT_COMPILER_OPTIONS
+#ifndef MLIR_TENSORRT_COMPILER_OPTIONSPROVIDERS
+#define MLIR_TENSORRT_COMPILER_OPTIONSPROVIDERS
 
 #include "mlir-executor/Support/DeviceInfo.h"
 #include "mlir/Pass/PassManager.h"
@@ -210,7 +210,7 @@ struct DeviceOptions : public OptionsProvider {
 
 private:
   /// Stores host device info. This is populated by the callback of
-  /// `shouldInfoFromHost`. If present, then it will also override the other
+  /// `shouldInferFromHost`. If present, then it will also override the other
   /// options in their callbacks.
   std::optional<DeviceInfo> hostDeviceInfo{};
 };
@@ -316,7 +316,6 @@ class CompilationTaskOptionsBase
                                  llvm::cl::desc("entrypoint function name")};
 
 protected:
-  std::vector<std::unique_ptr<CompilationTaskOptionsBase>> extensions;
   std::unique_ptr<DebugOptions> debugOptions{nullptr};
 };
 
@@ -363,4 +362,4 @@ class CompilationTaskOptions : public CompilationTaskOptionsBase {
 
 } // namespace mlirtrt::compiler
 
-#endif // MLIR_TENSORRT_COMPILER_OPTIONS
+#endif // MLIR_TENSORRT_COMPILER_OPTIONSPROVIDERS

mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/StablehloToExecutable/StablehloToExecutable.h

@@ -138,13 +138,6 @@ class StablehloToExecutableTask
 
   static void populatePassManager(mlir::OpPassManager &pm,
                                   const StablehloToExecutableOptions &options);
-
-  /// Compile a StableHLO module into a MLIR-TensorRT Runtime executable.
-  /// This is the "functional" entrypoint that will allocate a new PassManager
-  /// for a single run.
-  static mlirtrt::StatusOr<std::unique_ptr<runtime::Executable>>
-  compileStableHLOToExecutable(CompilerClient &client, mlir::ModuleOp module,
-                               const StablehloToExecutableOptions &options);
 };
 
 /// Register the task/options with the client's registry.

mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/StablehloToExecutable/TensorRTExtension.h

@@ -60,8 +60,8 @@ class StablehloToExecutableTensorRTExtension
     this->workspaceMemoryPoolLimit = options.workspaceMemoryPoolLimit;
   }
 
-  Option<bool> disable{this->ctx, "disable-tensorrt-extension",
-                       llvm::cl::init(false)};
+  Option<bool> disabled{this->ctx, "disable-tensorrt-extension",
+                        llvm::cl::init(false)};
 
   Option<TensorRTTargetFormat> format{
       this->ctx, "tensorrt-target",

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

@@ -42,7 +42,9 @@ def ConvertStablehloToTensorRTPass : Pass<"convert-stablehlo-to-tensorrt"> {
     Option<"convertConditionals", "convert-conditionals", "bool", "true",
       "convert conditionals to TensorRT's conditional layer">,
     Option<"trtMajorVersion", "trt-major-version", "int64_t", "10",
-    "target TensorRT version for conversion">
+    "target TensorRT version for conversion">,
+    Option<"preferEinsum", "prefer-einsum", "bool", "false",
+      "prefer converting to 'tensorrt.einsum' over 'tensorrt.matrix_multiply'">
   ];
 }
 #endif // MLIR_TENSORRT_ENABLE_HLO
@@ -321,7 +323,10 @@ def ConvertStablehloToScfPass : Pass<"convert-stablehlo-to-scf"> {
   }];
   let dependentDialects = [
     "::mlir::tensor::TensorDialect",
-    "::mlir::scf::SCFDialect"
+    "::mlir::scf::SCFDialect",
+    "::mlir::tensor::TensorDialect",
+    "::mlir::arith::ArithDialect",
+    "::mlir::math::MathDialect"
   ];
 }
 

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

@@ -31,11 +31,20 @@
 namespace mlir {
 class ConversionTarget;
 
-// Collection of rewrite patterns for lowering of Stable HLO to TensorRT
-// dialect.
+/// Populate patterns for converting Stablehlo reduction and contraction ops to
+/// TensorRT.
+void populateStablehloReductionAndContractionToTensorRtConversionPattern(
+    TensorRTTypeConverter &typeConverter, RewritePatternSet &patterns,
+    PatternBenefit benefit = 1, PatternBenefit dotToEinsumBenefit = 0);
+
+/// Collection of rewrite patterns for lowering of Stable HLO to TensorRT
+/// dialect.
+/// The `preferEinsum` parameter controls whether `tensorrt.einsum` is used
+/// as the primary method for converting `stablehlo.dot_general` or only for
+/// fallback when conversion to `tensorrt.matrix_multiply` is not possible.
 void populateStablehloToTensorRtConversionPattern(
     TensorRTTypeConverter &typeConverter, RewritePatternSet &patterns,
-    ShapeInfoCallbacks shapeInfoCallbacks = {});
+    ShapeInfoCallbacks shapeInfoCallbacks = {}, bool preferEinsum = false);
 
 /// Populate patterns for convert Chlo ops to TensorRT ops.
 void populateChloToTensorRtLegalityAndPatterns(

mlir-tensorrt/compiler/include/mlir-tensorrt/Conversion/TensorRTCommon/ConvertToTensorRTCommon.h

@@ -208,26 +208,6 @@ class ConvertOpToTensorRTPattern : public ConvertToTensorRTPattern {
       : ConvertToTensorRTPattern(typeConverter, SourceOp::getOperationName(),
                                  benefit, context) {}
 
-  /// Wrappers around the ConversionPattern methods that pass the derived op
-  /// type.
-  LogicalResult match(Operation *op) const final {
-    return match(cast<SourceOp>(op));
-  }
-  void rewrite(Operation *op, ArrayRef<Value> operands,
-               ConversionPatternRewriter &rewriter) const final {
-    if constexpr (SourceOp::hasProperties())
-      return rewrite(cast<SourceOp>(op),
-                     OpAdaptor(operands, op->getAttrDictionary(),
-                               cast<SourceOp>(op).getProperties()),
-                     rewriter);
-    rewrite(cast<SourceOp>(op), OpAdaptor(operands, op->getAttrDictionary()),
-            rewriter);
-  }
-  void rewrite(Operation *op, ArrayRef<ValueRange> operands,
-               ConversionPatternRewriter &rewriter) const final {
-    auto sourceOp = cast<SourceOp>(op);
-    rewrite(sourceOp, OneToNOpAdaptor(operands, sourceOp), rewriter);
-  }
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
@@ -248,32 +228,10 @@ class ConvertOpToTensorRTPattern : public ConvertToTensorRTPattern {
                            rewriter);
   }
 
-  /// Rewrite and Match methods that operate on the SourceOp type. These must be
-  /// overridden by the derived pattern class.
-  virtual LogicalResult match(SourceOp op) const {
-    (void)op;
-    llvm_unreachable("must override match or matchAndRewrite");
-  }
-  virtual void rewrite(SourceOp op, OpAdaptor adaptor,
-                       ConversionPatternRewriter &rewriter) const {
-    (void)op;
-    (void)adaptor;
-    (void)rewriter;
-    llvm_unreachable("must override matchAndRewrite or a rewrite method");
-  }
-  virtual void rewrite(SourceOp op, OneToNOpAdaptor adaptor,
-                       ConversionPatternRewriter &rewriter) const {
-    SmallVector<Value> oneToOneOperands =
-        getOneToOneAdaptorOperands(adaptor.getOperands());
-    rewrite(op, OpAdaptor(oneToOneOperands, adaptor), rewriter);
-  }
   virtual LogicalResult
   matchAndRewrite(SourceOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const {
-    if (failed(match(op)))
-      return failure();
-    rewrite(op, adaptor, rewriter);
-    return success();
+    llvm_unreachable("must override matchAndRewrite");
   }
   virtual LogicalResult
   matchAndRewrite(SourceOp op, OneToNOpAdaptor adaptor,

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

@@ -26,6 +26,7 @@
 
 #include "mlir-tensorrt-dialect/Interface/TensorKindOpInterface.h"
 #include "mlir-tensorrt/Compiler/Extension.h"
+#include "mlir-tensorrt/Dialect/Plan/IR/PlanEnums.h"
 #include "mlir-tensorrt/Dialect/Plan/IR/PlanInterfaces.h"
 #include "mlir-tensorrt/Interfaces/InferTensorValueRangeInterface.h"
 #include "mlir/Bytecode/BytecodeOpInterface.h"
@@ -134,11 +135,6 @@ class PlanDialectExtension
 };
 } // namespace mlir::plan
 
-//===----------------------------------------------------------------------===//
-// Plan Enums
-//===----------------------------------------------------------------------===//
-#include "mlir-tensorrt/Dialect/Plan/IR/PlanEnums.h.inc"
-
 //===----------------------------------------------------------------------===//
 // Plan Attributes
 //===----------------------------------------------------------------------===//