Integrate internal changes

GitOrigin-RevId: 630a69d8e14506db43cfefe4be2c790f9352da4f

Author: Copybara Bot

mlir-tensorrt/common/CMakeLists.txt

@@ -0,0 +1,36 @@
+cmake_minimum_required(VERSION 3.25)
+project(mlir-tensorrt-common LANGUAGES CXX)
+
+# Depdendencies
+find_package(LLVM REQUIRED CONFIG)
+find_package(MLIR REQUIRED CONFIG)
+include(HandleLLVMOptions)
+include_directories(${LLVM_INCLUDE_DIRS})
+include_directories(${MLIR_INCLUDE_DIRS})
+
+if(MLIR_TRT_TARGET_TENSORRT)
+  find_package(TensorRT REQUIRED)
+endif()
+
+set(MLIR_TENSORRT_COMMON_SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR})
+set(MLIR_TENSORRT_COMMON_BINARY_DIR ${CMAKE_CURRENT_BINARY_DIR})
+
+include_directories(include ${CMAKE_CURRENT_BINARY_DIR}/include)
+
+add_library(MLIRTensorRTCommonIncludes INTERFACE)
+target_include_directories(MLIRTensorRTCommonIncludes INTERFACE
+  "$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>"
+  "$<BUILD_INTERFACE:${CMAKE_CURRENT_BINARY_DIR}/include>"
+)
+
+add_subdirectory(include/mlir-tensorrt-common)
+add_subdirectory(lib)
+
+install(TARGETS MLIRTensorRTCommonIncludes
+  EXPORT MLIRTensorRTCommonTargets
+  RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR}
+  LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
+  ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
+  INCLUDES DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}
+)
+

mlir-tensorrt/common/include/mlir-tensorrt-common/CMakeLists.txt

@@ -1,6 +1,6 @@
-//===- LuaRegistration.h ----------------------------------------*- C++ -*-===//
+//===- Passes.h -------------------------------------------------*- C++ -*-===//
 //
-// SPDX-FileCopyrightText: Copyright 2024 NVIDIA CORPORATION & AFFILIATES.
+// SPDX-FileCopyrightText: Copyright 2025 NVIDIA CORPORATION & AFFILIATES.
 // All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
 //
@@ -18,22 +18,22 @@
 //
 //===----------------------------------------------------------------------===//
 ///
-/// Registration for the Lua runtime methods.
+/// This file contains the declarations for the common conversion passes.
 ///
 //===----------------------------------------------------------------------===//
+#ifndef MLIR_TENSORRT_COMMON_CONVERSION_PASSES
+#define MLIR_TENSORRT_COMMON_CONVERSION_PASSES
 
-#include "mlir-executor/Runtime/API/API.h"
+#include "mlir/Pass/Pass.h"
+#include <memory>
 
-struct lua_State;
-
-namespace mlirtrt::runtime {
-/// Register various external functions with the given Lua state using a
-/// directly specified device number, total device count, and a pre-determined
-/// NCCL uuid.
-void registerLuaRuntimeMethods(lua_State *state,
-                               const RuntimeSessionOptions &options,
-                               PinnedMemoryAllocator *pinnedMemoryAllocator,
-                               AllocTracker *allocTracker,
-                               ResourceTracker *resourceTracker);
+//===----------------------------------------------------------------------===//
+// Add Tablegen'd pass declarations and registration methods.
+//===----------------------------------------------------------------------===//
+namespace mlir {
+#define GEN_PASS_DECL
+#define GEN_PASS_REGISTRATION
+#include "mlir-tensorrt-common/Conversion/Passes.h.inc"
+} // namespace mlir
 
-} // namespace mlirtrt::runtime
+#endif // MLIR_TENSORRT_COMMON_CONVERSION_PASSES

mlir-tensorrt/executor/include/mlir-executor/Runtime/Backend/Lua/LuaRegistration.h renamed to mlir-tensorrt/common/include/mlir-tensorrt-common/Conversion/Passes.h

@@ -0,0 +1,18 @@
+#ifndef MLIR_TENSORRT_COMMON_CONVERSION_PASSES
+#define MLIR_TENSORRT_COMMON_CONVERSION_PASSES
+
+include "mlir/Pass/PassBase.td"
+
+def ConvertToLoops : Pass<"convert-to-loops"> {
+  let summary = "Convert a LoopLikeOpInterface to loops";
+  let description = [{
+    This pass converts a LoopLikeOpInterface to loops.
+  }];
+
+  let dependentDialects = [
+    "::mlir::tensor::TensorDialect",
+    "::mlir::scf::SCFDialect",
+  ];
+}
+
+#endif // MLIR_TENSORRT_COMMON_CONVERSION_PASSES

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

@@ -0,0 +1,36 @@
+//===- DataLayoutImpl.h -----------------------------------------*- C++ -*-===//
+//
+// SPDX-FileCopyrightText: Copyright 2025 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.
+//
+//===----------------------------------------------------------------------===//
+///
+/// This file contains the declarations for the DataLayout extensions to
+/// the EmitC dialect.
+/// TODO: These interfaces should be upstreamed to the EmitC dialect so that
+/// external models are not required.
+///
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_TENSORRT_COMMON_DIALECT_EMITCEXT_IR_DATALAYOUTIMPL_H
+#define MLIR_TENSORRT_COMMON_DIALECT_EMITCEXT_IR_DATALAYOUTIMPL_H
+
+#include "mlir/IR/DialectRegistry.h"
+
+namespace mlir::emitc_ext {
+void registerDataLayoutInterfaceExternalModels(DialectRegistry &registry);
+}
+
+#endif // MLIR_TENSORRT_COMMON_DIALECT_EMITCEXT_IR_DATALAYOUTIMPL_H

mlir-tensorrt/common/include/mlir-tensorrt-common/Dialect/EmitCExt/IR/DataLayoutImpl.h

@@ -0,0 +1,53 @@
+//===- ToLoopsOpInterfaceImpl.h ---------------------------------*- C++ -*-===//
+//
+// SPDX-FileCopyrightText: Copyright 2025 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.
+//
+//===----------------------------------------------------------------------===//
+///
+/// This file contains the declarations for the ToLoopsOpInterface extensions to
+/// the Linalg dialect.
+///
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_TENSORRT_COMMON_DIALECT_LINALGEXT_TRANSFORMS_TOLOOPSOPINTERFACEIMPL
+#define MLIR_TENSORRT_COMMON_DIALECT_LINALGEXT_TRANSFORMS_TOLOOPSOPINTERFACEIMPL
+
+#include "mlir-tensorrt-common/Interfaces/ToLoopsOpInterface.h"
+#include "mlir/IR/DialectRegistry.h"
+
+namespace mlir::linalg {
+class LinalgOp;
+}
+
+namespace mlir::scf {
+class ForOp;
+}
+
+namespace mlir::linalg_ext {
+
+/// Register the ToLoopsOpInterface external models for GenericOp. For other
+/// kinds of operations that are LinalgOps, we don't register an external model
+/// because there are so many; instead use the below function to perform
+/// conversion.
+void registerToLoopsOpInterfaceExternalModels(DialectRegistry &registry);
+
+/// Convert a LinalgOp (on tensors) to SCF loops.
+FailureOr<SmallVector<Operation *>>
+convertLinalgOpToLoops(RewriterBase &rewriter, linalg::LinalgOp op);
+
+} // namespace mlir::linalg_ext
+
+#endif // MLIR_TENSORRT_COMMON_DIALECT_LINALGEXT_TRANSFORMS_TOLOOPSOPINTERFACEIMPL

mlir-tensorrt/common/include/mlir-tensorrt-common/Dialect/LinalgExt/Transforms/ToLoopsOpInterfaceImpl.h

@@ -0,0 +1,38 @@
+//===- ToLoopsOpInterface.h -------------------------------*- C++ -*-===//
+//
+// SPDX-FileCopyrightText: Copyright 2025 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.
+//
+//===----------------------------------------------------------------------===//
+///
+/// This file contains the declarations for the `ToLoopsOpInterface` interface.
+///
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_TENSORRT_COMMON_INTERFACES_TOLOOPSOPINTERFACE
+#define MLIR_TENSORRT_COMMON_INTERFACES_TOLOOPSOPINTERFACE
+
+#include "mlir/IR/OpDefinition.h"
+
+namespace mlir {
+class RewriterBase;
+namespace scf {
+class ForOp;
+} // namespace scf
+} // namespace mlir
+
+#include "mlir-tensorrt-common/Interfaces/ToLoopsOpInterface.h.inc"
+
+#endif // MLIR_TENSORRT_COMMON_INTERFACES_TOLOOPSOPINTERFACE

mlir-tensorrt/common/include/mlir-tensorrt-common/Interfaces/ToLoopsOpInterface.h

@@ -0,0 +1,28 @@
+#ifndef MLIR_TENSORRT_COMMON_INTERFACES_TOLOOPSOPINTERFACE
+#define MLIR_TENSORRT_COMMON_INTERFACES_TOLOOPSOPINTERFACE
+
+include "mlir/IR/OpBase.td"
+
+def ToLoopsOpInterface : OpInterface<"ToLoopsOpInterface"> {
+  let description = "Interface for lowering to loops";
+
+  let cppNamespace = "::mlir";
+
+  let methods = [
+    InterfaceMethod<[{
+        Lower the operation to a loop nest. Returns
+        the outermost loop that should replace the original
+        op, but does not actually perform the replacement.
+      }],
+        "::mlir::FailureOr<::mlir::Operation*>",
+        "lowerToLoops",
+        (ins "::mlir::RewriterBase&":$rewriter),
+        "",
+        [{
+          llvm_unreachable("Not implemented");
+        }]
+    >
+  ];
+}
+
+#endif // MLIR_TENSORRT_COMMON_INTERFACES_TOLOOPSOPINTERFACE

mlir-tensorrt/common/include/mlir-tensorrt-common/Interfaces/ToLoopsOpInterface.td

@@ -0,0 +1,5 @@
+add_subdirectory(Conversion)
+add_subdirectory(Dialect)
+add_subdirectory(Interfaces)
+add_subdirectory(Utils)
+

mlir-tensorrt/common/lib/CMakeLists.txt

@@ -0,0 +1,11 @@
+set(LLVM_TARGET_DEFINITIONS
+  "${MLIR_TENSORRT_COMMON_SOURCE_DIR}/include/mlir-tensorrt-common/Conversion/Passes.td")
+set(OUTPUT_DIR
+  "${MLIR_TENSORRT_COMMON_BINARY_DIR}/include/mlir-tensorrt-common/Conversion")
+
+cmake_path(SET OUTPUT_DIR NORMALIZE "${OUTPUT_DIR}")
+cmake_path(RELATIVE_PATH OUTPUT_DIR BASE_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}")
+mlir_tablegen("${OUTPUT_DIR}/Passes.h.inc" -gen-pass-decls -name MLIRTensorRTCommonConversion)
+add_public_tablegen_target(MLIRTensorRTCommonConversionPassesIncGen)
+
+add_subdirectory(ToLoops)