Cherry pick: Allow creation of TCP groups where an op has multiple uses

BUILD

@@ -23,7 +23,9 @@ td_library(
         "include/mlir-tcp/Dialect/IR/TcpEnums.td",
         "include/mlir-tcp/Dialect/IR/TcpOps.td",
         "include/mlir-tcp/Dialect/IR/TcpTypes.td",
-    ],
+        "include/mlir-tcp/Dialect/IR/TcpOpsCruiseInternal.td",
+        "include/mlir-tcp/Dialect/IR/TcpTypesCruiseInternal.td",
+        ],
     includes = ["include"],
     deps = [
         "@llvm-project//mlir:OpBaseTdFiles",
@@ -202,12 +204,15 @@ cc_library(
         "lib/Conversion/TorchToTcp/PopulatePatterns.h",
         "lib/Conversion/TorchToTcp/TcpCustomOp.cpp",
         "lib/Conversion/TorchToTcp/TorchToTcp.cpp",
+        "lib/Conversion/TorchToTcp/TorchToTcpCruiseInternal.cpp",
+        "lib/Conversion/TorchToTcp/CruiseInternalPatterns.cpp",
         "lib/Conversion/TorchToTcp/TorchToTcpCustomOp.cpp",
         "lib/Conversion/TorchToTcp/Utils.cpp",
         "lib/Conversion/TorchToTcp/Utils.h",
     ],
     hdrs = [
         "include/mlir-tcp/Conversion/TorchToTcp/TorchToTcp.h",
+        "include/mlir-tcp/Conversion/TorchToTcp/TorchToTcpCruiseInternal.h",
         "include/mlir-tcp/Conversion/TorchToTcp/TorchToTcpCustomOp.h",
     ],
     strip_include_prefix = "include",

include/mlir-tcp/Conversion/Passes.td

@@ -46,6 +46,18 @@ def ConvertTorchToTcpCustomOp : Pass<"convert-torch-to-tcp-custom-op", "func::Fu
   ];
 }
 
+//===----------------------------------------------------------------------===//
+// TorchToTcpCruiseInternal
+//===----------------------------------------------------------------------===//
+
+def ConvertTorchToTcpCruiseInternal : Pass<"convert-torch-to-tcp-cruise-internal", "func::FuncOp"> {
+  let summary = "Convert Torch ops to Tcp ops with cruise internal patterns only";
+  let description = [{
+    Convert Torch ops to Tcp ops with cruise internal patterns only.
+  }];
+  let constructor = "mlir::tcp::createConvertTorchToTcpCruiseInternalPass()";
+}
+
 //===----------------------------------------------------------------------===//
 // StablehloToTcp
 //===----------------------------------------------------------------------===//

include/mlir-tcp/Conversion/TorchToTcp/TorchToTcpCruiseInternal.h

@@ -0,0 +1,26 @@
+//===------------------------------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// Also available under a BSD-style license. See LICENSE.
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+
+#define GEN_PASS_DECL_CONVERTTORCHTOTCPCRUISEINTERNAL
+#include "mlir-tcp/Conversion/Passes.h.inc"
+
+namespace tcp {
+
+std::unique_ptr<OperationPass<func::FuncOp>>
+createConvertTorchToTcpCruiseInternalPass();
+
+} // namespace tcp
+} // namespace mlir

include/mlir-tcp/Dialect/IR/TcpOps.td

@@ -16,6 +16,9 @@ include "mlir-tcp/Dialect/IR/TcpBase.td"
 include "mlir-tcp/Dialect/IR/TcpEnums.td"
 include "mlir-tcp/Dialect/IR/TcpTypes.td"
 
+// Import cruise specific ops
+include "mlir-tcp/Dialect/IR/TcpOpsCruiseInternal.td"
+
 def Tcp_TanhOp : Tcp_UnaryElementwiseOp<"tanh", [SameOperandsAndResultElementType]> {
   let summary = "Computes tanh of input, elementwise";
 

include/mlir-tcp/Dialect/IR/TcpOpsCruiseInternal.td

@@ -0,0 +1,83 @@
+//===-------------------------------------------------------*- tablegen -*-===//
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// Also available under a BSD-style license. See LICENSE.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef TCP_OPS_CRUISE_INTERNAL
+#define TCP_OPS_CRUISE_INTERNAL
+
+def Tcp_CreateIndexArrayOp : Tcp_Op<"create_index_array"> {
+    let summary = "Creates index_arry type given variadic number of input indices";
+
+    let arguments = (ins
+        Variadic<Index>:$inputs
+    );
+
+    let results = (outs
+        Tcp_IndexArrayType : $outputs
+    );
+
+  let assemblyFormat = [{
+    `(` operands `)` `:` attr-dict-with-keyword functional-type(operands, results)
+  }];
+}
+
+def Tcp_BindTensorShape : Tcp_Op<"bind_tensor_shape", []> {
+    let summary = "Binds a tensor to its index_array shape type.";
+
+    let arguments = (ins
+        AnyType:$tensor,
+        Tcp_IndexArrayType : $inputs
+    );
+}
+
+// caspr CasprCreateTensorFromIndexOp : used for encapsulating result of tensor_dim_op in tensor
+def Tcp_CasprCreateTensorFromIndexOp : Tcp_Op<"caspr_create_tensor_from_index"> {
+    let arguments = (ins
+      Index:$index
+    );
+
+    let results = (outs
+      AnyType:$result
+    );
+
+  let assemblyFormat = [{
+    `(` operands `)` `:` attr-dict-with-keyword functional-type(operands, results)
+  }];
+}
+
+// caspr CasprCreateTensorFromIndexArrayOp : used for encapsulating result of tensor_dim_op in tensor
+def Tcp_CasprCreateTensorFromIndexArrayOp : Tcp_Op<"caspr_create_tensor_from_index_array"> {
+    let arguments = (ins
+      Tcp_IndexArrayType:$shape_array
+    );
+
+    let results = (outs
+      AnyType:$result
+    );
+
+  let assemblyFormat = [{
+    `(` operands `)` `:` attr-dict-with-keyword functional-type(operands, results)
+  }];
+}
+
+def Tcp_CasprIndexFromTensorOp : Tcp_Op<"caspr_index_from_tensor"> {
+    let arguments = (ins
+      AnyType:$input
+    );
+
+    let results = (outs
+      Index:$index
+    );
+
+  let assemblyFormat = [{
+    `(` operands `)` `:` attr-dict-with-keyword functional-type(operands, results)
+  }];
+}
+
+
+#endif // TCP_OPS_CRUISE_INTERNAL

include/mlir-tcp/Dialect/IR/TcpTypes.td

@@ -14,7 +14,7 @@ include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/DialectBase.td"
 
 include "mlir-tcp/Dialect/IR/TcpBase.td"
-
+include "mlir-tcp/Dialect/IR/TcpTypesCruiseInternal.td"
 
 //===----------------------------------------------------------------------===//
 // Tcp Quantized Types
@@ -58,13 +58,4 @@ def Tcp_FloatTensor : RankedTensorOf<[AnyFloat]>;
 def Tcp_IntTensor : RankedTensorOf<[AnySignlessInteger]>;
 def Tcp_FloatOrIntTensor : RankedTensorOf<[AnyFloat, AnySignlessInteger]>;
 
-
-//===----------------------------------------------------------------------===//
-// Tcp Custom Types
-//===----------------------------------------------------------------------===//
-
-def Tcp_IndexArrayType : Tcp_Type<"IndexArray", "index_array"> {
-    let summary = "IndexArray TCP type, to holds a list of index builtin type to represent shape";
-}
-
 #endif // TCP_TYPES

include/mlir-tcp/Dialect/IR/TcpTypesCruiseInternal.td

@@ -0,0 +1,27 @@
+//===-------------------------------------------------------*- tablegen -*-===//
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// Also available under a BSD-style license. See LICENSE.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef TCP_TYPES_CRUISE_INTERNAL
+#define TCP_TYPES_CRUISE_INTERNAL
+
+//===----------------------------------------------------------------------===//
+// Tcp Custom Types
+//===----------------------------------------------------------------------===//
+
+def Tcp_IndexArrayType : Tcp_Type<"IndexArray", "index_array"> {
+    let summary = "IndexArray TCP type, to holds a list of index builtin type to represent shape";
+
+    let parameters = (ins
+        "int":$elements
+    );
+
+    let assemblyFormat = "`<` `[` $elements `]` `>`";
+}
+
+#endif // TCP_TYPES_CRUISE_INTERNAL

lib/Conversion/Passes.cpp

@@ -13,6 +13,7 @@
 #include "mlir-tcp/Conversion/TcpToArith/TcpToArith.h"
 #include "mlir-tcp/Conversion/TcpToLinalg/TcpToLinalg.h"
 #include "mlir-tcp/Conversion/TorchToTcp/TorchToTcp.h"
+#include "mlir-tcp/Conversion/TorchToTcp/TorchToTcpCruiseInternal.h"
 #include "mlir-tcp/Conversion/TorchToTcp/TorchToTcpCustomOp.h"
 
 //===----------------------------------------------------------------------===//