[Tcp] Add boilerplate for TCP dialect (#1375)

* Initial boilerplate for TCP with a dummy op.

* Conditional flag to enable TCP

build_tools/python_deploy/build_linux_packages.sh

@@ -205,6 +205,7 @@ function build_in_tree() {
       -DLLVM_TARGETS_TO_BUILD=host \
       -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
       -DTORCH_MLIR_ENABLE_LTC=ON \
+      -DTORCH_MLIR_DIALECTS_ENABLE_TCP=OFF \
       -DTORCH_MLIR_USE_INSTALLED_PYTORCH="$torch_from_bin" \
       -DTORCH_MLIR_SRC_PYTORCH_REPO=${TORCH_MLIR_SRC_PYTORCH_REPO} \
       -DTORCH_MLIR_SRC_PYTORCH_BRANCH=${TORCH_MLIR_SRC_PYTORCH_BRANCH} \

externals/llvm-external-projects/torch-mlir-dialects/CMakeLists.txt

@@ -10,6 +10,11 @@ endif()
 
 option(MLIR_ENABLE_BINDINGS_PYTHON "Enables MLIR Python Bindings" OFF)
 
+option(TORCH_MLIR_DIALECTS_ENABLE_TCP "Add TCP dialect" OFF)
+if(TORCH_MLIR_DIALECTS_ENABLE_TCP)
+  add_definitions(-DTORCH_MLIR_DIALECTS_ENABLE_TCP)
+endif()
+
 set(TORCH_MLIR_DIALECTS_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}")
 set(TORCH_MLIR_DIALECTS_BINARY_DIR "${CMAKE_CURRENT_BINARY_DIR}")
 message(STATUS "Building torch-mlir-dialects project at ${TORCH_MLIR_DIALECTS_SOURCE_DIR} (into ${TORCH_MLIR_DIALECTS_BINARY_DIR})")

externals/llvm-external-projects/torch-mlir-dialects/README.md

@@ -9,3 +9,15 @@ This project is intended to be used via LLVM's external projects setup:
 * `-DLLVM_EXTERNAL_TORCH_MLIR_DIALECTS_SOURCE_DIR={this_directory}`
 
 It depends on the `mlir` project.
+
+## TCP Dialect
+Tensor Compute Primitives (TCP) is a mid-level transformation oriented IR for deep learning & similar applications.
+
+TCP is being bootstrapped under Torch-MLIR Dialects.
+
+All technical discussions regarding TCP will be done in the discourse [TCP-WG](https://discourse.llvm.org/c/mlir/mlir-tcp-wg/36) category.
+### References
+* Discourse [thread](https://discourse.llvm.org/t/rfc-proposal-for-a-high-level-ml-dialect-in-mlir/64249) regarding TCP proposal.
+* TCP [proposal document](https://docs.google.com/document/d/1f3KVsXA4xm6W7gd2cKx9ThGA52XMfZKsPV3ZOlsztC4/edit)
+* TCP spec [draft](https://docs.google.com/document/d/1Twyph8jU_f1QDoBInr8OkUXxcyRw-KqFdyclkK1UQbc/edit)
+* Torch-MLIR [issue](https://github.com/llvm/torch-mlir/issues/1366) to track bootstrapping TCP.

externals/llvm-external-projects/torch-mlir-dialects/include/torch-mlir-dialects/Dialect/CMakeLists.txt

@@ -1 +1,5 @@
 add_subdirectory(TMTensor)
+if(TORCH_MLIR_DIALECTS_ENABLE_TCP)
+  add_subdirectory(Tcp)
+endif()
+

externals/llvm-external-projects/torch-mlir-dialects/include/torch-mlir-dialects/Dialect/Tcp/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(IR)

externals/llvm-external-projects/torch-mlir-dialects/include/torch-mlir-dialects/Dialect/Tcp/IR/CMakeLists.txt

@@ -0,0 +1,11 @@
+function(_add_dialect)
+  set(LLVM_TARGET_DEFINITIONS TcpOps.td)
+  mlir_tablegen(TcpOps.h.inc -gen-op-decls)
+  mlir_tablegen(TcpOps.cpp.inc -gen-op-defs)
+  mlir_tablegen(TcpDialect.h.inc -gen-dialect-decls -dialect=tcp)
+  mlir_tablegen(TcpDialect.cpp.inc -gen-dialect-defs -dialect=tcp)
+  add_public_tablegen_target(TorchMLIRTcpOpsIncGen)
+  add_dependencies(mlir-headers TorchMLIRTcpOpsIncGen)
+endfunction()
+
+_add_dialect()

externals/llvm-external-projects/torch-mlir-dialects/include/torch-mlir-dialects/Dialect/Tcp/IR/TcpBase.td

@@ -0,0 +1,43 @@
+//===-------------------------------------------------------*- tablegen -*-===//
+//
+// 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.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef TORCH_MLIR_DIALECT_TCP_BASE
+#define TORCH_MLIR_DIALECT_TCP_BASE
+
+include "mlir/IR/OpBase.td"
+
+def Tcp_Dialect : Dialect {
+  let name = "tcp";
+  let cppNamespace = "::mlir::tcp";
+  let description = [{
+    Tensor Compute Primitives (TCP) dialect.
+
+    TCP is a mid-level transformation oriented IR for deep learning & similar
+    applications.
+  }];
+
+  let dependentDialects = [];
+}
+
+//===----------------------------------------------------------------------===//
+// Tcp Type Definitions.
+//===----------------------------------------------------------------------===//
+
+def Tcp_Scalar : AnyTypeOf<[AnyFloat, AnySignlessInteger, AnyComplex]>;
+def Tcp_Tensor : RankedTensorOf<[Tcp_Scalar]>;
+
+//===----------------------------------------------------------------------===//
+// Tcp Ops Base.
+//===----------------------------------------------------------------------===//
+
+class Tcp_Op<string mnemonic, list<Trait> traits = []> :
+    Op<Tcp_Dialect, mnemonic, traits> {
+}
+
+#endif // TORCH_MLIR_DIALECT_TCP_BASE

externals/llvm-external-projects/torch-mlir-dialects/include/torch-mlir-dialects/Dialect/Tcp/IR/TcpDialect.h

@@ -0,0 +1,18 @@
+//===------------------------------------------------------------*- 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.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef TORCH_MLIR_DIALECTS_DIALECT_TCP_IR_TCPDIALECT_H_
+#define TORCH_MLIR_DIALECTS_DIALECT_TCP_IR_TCPDIALECT_H_
+
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/OpDefinition.h"
+
+#include "torch-mlir-dialects/Dialect/Tcp/IR/TcpDialect.h.inc"
+
+#endif // TORCH_MLIR_DIALECTS_DIALECT_TCP_IR_TCPDIALECT_H_

externals/llvm-external-projects/torch-mlir-dialects/include/torch-mlir-dialects/Dialect/Tcp/IR/TcpOps.h

@@ -0,0 +1,21 @@
+//===------------------------------------------------------------*- 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.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef TORCH_MLIR_DIALECTS_DIALECT_TCP_IR_TCPOPS_H_
+#define TORCH_MLIR_DIALECTS_DIALECT_TCP_IR_TCPOPS_H_
+
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/OpDefinition.h"
+
+#define GET_OP_CLASSES
+#include "torch-mlir-dialects/Dialect/Tcp/IR/TcpOps.h.inc"
+
+#endif // TORCH_MLIR_DIALECTS_DIALECT_TCP_IR_TCPOPS_H_

externals/llvm-external-projects/torch-mlir-dialects/include/torch-mlir-dialects/Dialect/Tcp/IR/TcpOps.td

@@ -0,0 +1,35 @@
+//===-------------------------------------------------------*- tablegen -*-===//
+//
+// 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.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef TORCH_MLIR_DIALECT_TCP_OPS
+#define TORCH_MLIR_DIALECT_TCP_OPS
+
+include "torch-mlir-dialects/Dialect/Tcp/IR/TcpBase.td"
+
+include "mlir/IR/OpBase.td"
+
+def Tcp_DummyOp : Tcp_Op<"dummy", []> {
+  let summary = "dummy op";
+
+  let description = [{
+    Dummy identity op that returns the input as output
+  }];
+
+  let arguments = (ins
+    Tcp_Tensor:$in
+  );
+
+  let results = (outs
+    Tcp_Tensor:$out
+  );
+
+  let assemblyFormat = "$in attr-dict `:` type($in) `->` type($out)";
+}
+
+#endif // TORCH_MLIR_DIALECT_TCP_OPS

externals/llvm-external-projects/torch-mlir-dialects/lib/Dialect/CMakeLists.txt

@@ -1 +1,5 @@
 add_subdirectory(TMTensor)
+if(TORCH_MLIR_DIALECTS_ENABLE_TCP)
+  add_subdirectory(Tcp)
+endif()
+

externals/llvm-external-projects/torch-mlir-dialects/lib/Dialect/Tcp/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(IR)

externals/llvm-external-projects/torch-mlir-dialects/lib/Dialect/Tcp/IR/CMakeLists.txt

@@ -0,0 +1,20 @@
+add_mlir_library(TorchMLIRTcpDialect
+  TcpDialect.cpp
+  TcpOps.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${TORCH_MLIR_DIALECTS_SOURCE_DIR}/include
+
+  DEPENDS
+  TorchMLIRTcpOpsIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRDialectUtils
+  MLIRIR
+  MLIRPass
+  MLIRSideEffectInterfaces
+  MLIRSupport
+  MLIRViewLikeInterface
+)
+
+torch_mlir_dialects_target_includes(TorchMLIRTcpDialect)

externals/llvm-external-projects/torch-mlir-dialects/lib/Dialect/Tcp/IR/TcpDialect.cpp

@@ -0,0 +1,30 @@
+//===------------------------------------------------------------*- 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.
+//
+//===----------------------------------------------------------------------===//
+
+#include "torch-mlir-dialects/Dialect/Tcp/IR/TcpDialect.h"
+#include "torch-mlir-dialects/Dialect/Tcp/IR/TcpOps.h"
+
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/DialectImplementation.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/OpImplementation.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/SourceMgr.h"
+
+using namespace mlir;
+using namespace mlir::tcp;
+
+void TcpDialect::initialize() {
+#define GET_OP_LIST
+  addOperations<
+#include "torch-mlir-dialects/Dialect/Tcp/IR/TcpOps.cpp.inc"
+      >();
+}
+
+#include "torch-mlir-dialects/Dialect/Tcp/IR/TcpDialect.cpp.inc"

externals/llvm-external-projects/torch-mlir-dialects/lib/Dialect/Tcp/IR/TcpOps.cpp

@@ -0,0 +1,18 @@
+//===------------------------------------------------------------*- 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.
+//
+//===----------------------------------------------------------------------===//
+
+#include "torch-mlir-dialects/Dialect/Tcp/IR/TcpOps.h"
+
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/OperationSupport.h"
+#include "mlir/IR/Value.h"
+
+#define GET_OP_CLASSES
+#include "torch-mlir-dialects/Dialect/Tcp/IR/TcpOps.cpp.inc"

externals/llvm-external-projects/torch-mlir-dialects/test/CMakeLists.txt

@@ -1,3 +1,5 @@
+llvm_canonicalize_cmake_booleans(TORCH_MLIR_DIALECTS_ENABLE_TCP)
+
 # lit needs to find the tools, and the location differs for in-tree and out-of-tree builds.
 get_target_property(TORCH_MLIR_DIALECTS_TOOLS_DIR torch-mlir-dialects-opt RUNTIME_OUTPUT_DIRECTORY)
 configure_lit_site_cfg(

externals/llvm-external-projects/torch-mlir-dialects/test/Dialect/Tcp/dummy-op.mlir

@@ -0,0 +1,8 @@
+// RUN: torch-mlir-dialects-opt %s | torch-mlir-dialects-opt | FileCheck %s
+
+// CHECK-LABEL: func.func @test_dummy(%{{.*}}: tensor<10x20xf32>) -> tensor<10x20xf32>
+func.func @test_dummy(%arg0 : tensor<10x20xf32>) -> tensor<10x20xf32> {
+  // CHECK: %{{.*}} = tcp.dummy %{{.*}} : tensor<10x20xf32> -> tensor<10x20xf32>
+  %0 = tcp.dummy %arg0 : tensor<10x20xf32> -> tensor<10x20xf32>
+  return %0 : tensor<10x20xf32>
+}

externals/llvm-external-projects/torch-mlir-dialects/test/Dialect/Tcp/lit.local.cfg

@@ -0,0 +1,2 @@
+if not config.enable_tcp:
+  config.unsupported = True

externals/llvm-external-projects/torch-mlir-dialects/test/lit.site.cfg.py.in

@@ -19,6 +19,7 @@ config.llvm_shlib_ext = "@SHLIBEXT@"
 config.llvm_exe_ext = "@EXEEXT@"
 config.lit_tools_dir = "@LLVM_LIT_TOOLS_DIR@"
 config.python_executable = sys.executable
+config.enable_tcp = @TORCH_MLIR_DIALECTS_ENABLE_TCP@
 
 import lit.llvm
 lit.llvm.initialize(lit_config, config)

externals/llvm-external-projects/torch-mlir-dialects/tools/torch-mlir-dialects-opt/CMakeLists.txt

@@ -13,6 +13,10 @@ set(LIBS
   TorchMLIRTMTensorPasses
 )
 
+if(TORCH_MLIR_DIALECTS_ENABLE_TCP)
+  list(APPEND LIBS TorchMLIRTcpDialect)
+endif()
+
 add_llvm_tool(torch-mlir-dialects-opt
   torch-mlir-dialects-opt.cpp
 

externals/llvm-external-projects/torch-mlir-dialects/tools/torch-mlir-dialects-opt/torch-mlir-dialects-opt.cpp

@@ -18,6 +18,9 @@
 #include "mlir/IR/Dialect.h"
 #include "mlir/Tools/mlir-opt/MlirOptMain.h"
 #include "mlir/Transforms/Passes.h"
+#ifdef TORCH_MLIR_DIALECTS_ENABLE_TCP
+#include "torch-mlir-dialects/Dialect/Tcp/IR/TcpDialect.h"
+#endif // TORCH_MLIR_DIALECTS_ENABLE_TCP
 #include "torch-mlir-dialects/Dialect/TMTensor/IR/ScalarLoopOpInterface.h"
 #include "torch-mlir-dialects/Dialect/TMTensor/IR/TMTensorDialect.h"
 #include "torch-mlir-dialects/Dialect/TMTensor/Transforms/Passes.h"
@@ -38,6 +41,9 @@ int main(int argc, char **argv) {
   registry.insert<
       // Local dialects
       mlir::torch::TMTensor::TMTensorDialect,
+#ifdef TORCH_MLIR_DIALECTS_ENABLE_TCP
+      mlir::tcp::TcpDialect,
+#endif // TORCH_MLIR_DIALECTS_ENABLE_TCP
       // Upstream dialects
       mlir::arith::ArithDialect, mlir::linalg::LinalgDialect,
       mlir::func::FuncDialect, mlir::memref::MemRefDialect,