Add TMTensor dialect to torch-mlir

This is intended to explore support for non-structured ops that can't
be modeled by Linalg dialect. `tm_tensor.scan` and `tm_tensor.scatter`
are added as the first such ops. The dialect should aim to be
upstreamed in the future.

.github/workflows/buildAndTest.yml

@@ -53,8 +53,9 @@ jobs:
           -DPython3_EXECUTABLE=$(which python) \
           -DLLVM_ENABLE_ASSERTIONS=ON \
           -DLLVM_ENABLE_PROJECTS=mlir \
-          -DLLVM_EXTERNAL_PROJECTS=torch-mlir \
+          -DLLVM_EXTERNAL_PROJECTS="torch-mlir;torch-mlir-dialects" \
           -DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR="$GITHUB_WORKSPACE" \
+          -DLLVM_EXTERNAL_TORCH_MLIR_DIALECTS_SOURCE_DIR="${GITHUB_WORKSPACE}/external/llvm-external-projects/torch-mlir-dialects" \
           -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
           -DLLVM_TARGETS_TO_BUILD=host
         ninja check-torch-mlir-all

CMakeLists.txt

@@ -25,6 +25,22 @@ project(torch-mlir LANGUAGES CXX C)
 set(CMAKE_C_STANDARD 11)
 set(CMAKE_CXX_STANDARD 14)
 
+macro(torch_mlir_add_llvm_external_project name identifier location)
+  message(STATUS "Adding LLVM external project ${name} (${identifier}) -> ${location}")
+  if(NOT EXISTS "${location}/CMakeLists.txt")
+    message(FATAL_ERROR "External project location ${location} is not valid")
+  endif()
+  list(APPEND LLVM_EXTERNAL_PROJECTS ${name})
+  list(REMOVE_DUPLICATES LLVM_EXTERNAL_PROJECTS)
+  set(LLVM_EXTERNAL_${identifier}_SOURCE_DIR ${location} CACHE STRING "" FORCE)
+  set(LLVM_EXTERNAL_PROJECTS ${LLVM_EXTERNAL_PROJECTS} CACHE STRING "" FORCE)
+endmacro()
+
+torch_mlir_add_llvm_external_project(
+  torch-mlir-dialects
+  TORCH_MLIR_DIALECTS
+  ${CMAKE_CURRENT_SOURCE_DIR}/external/llvm-external-projects/torch-mlir-dialects)
+
 if(CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR)
   # Out-of-tree build
 
@@ -129,6 +145,7 @@ add_subdirectory(tools)
 add_custom_target(check-torch-mlir-all)
 add_dependencies(check-torch-mlir-all
   check-torch-mlir
+  check-torch-mlir-dialects
   )
 
 if(MLIR_ENABLE_BINDINGS_PYTHON)

README.md

@@ -62,8 +62,9 @@ cmake -GNinja -Bbuild \
   -DCMAKE_CXX_COMPILER=clang++ \
   -DPython3_FIND_VIRTUALENV=ONLY \
   -DLLVM_ENABLE_PROJECTS=mlir \
-  -DLLVM_EXTERNAL_PROJECTS=torch-mlir \
+  -DLLVM_EXTERNAL_PROJECTS=torch-mlir;torch-mlir-dialects \
   -DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR=`pwd` \
+  -DLLVM_EXTERNAL_TORCH_MLIR_DIALECTS_SOURCE_DIR=`pwd`/external/llvm-external-projects/torch-mlir-dialects \
   -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
   -DLLVM_TARGETS_TO_BUILD=host \
   external/llvm-project/llvm

build_tools/build_standalone.sh

@@ -19,8 +19,9 @@ cmake -GNinja -B"$build_dir" "$llvm_project_dir/llvm" \
   -DCMAKE_BUILD_TYPE=Release \
   -DCMAKE_C_COMPILER_LAUNCHER=ccache -DCMAKE_CXX_COMPILER_LAUNCHER=ccache \
   -DLLVM_ENABLE_PROJECTS=mlir \
-  -DLLVM_EXTERNAL_PROJECTS=torch-mlir \
+  -DLLVM_EXTERNAL_PROJECTS=torch-mlir;torch-mlir-dialects \
   -DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR="$project_dir" \
+  -DLLVM_EXTERNAL_TORCH_MLIR_DIALECTS_SOURCE_DIR=${project_dir}/external/llvm-external-projects/torch-mlir-dialects \
   -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
   -DLLVM_ENABLE_ASSERTIONS=ON \
   -DLLVM_TARGETS_TO_BUILD=host

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

@@ -0,0 +1,54 @@
+if(CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR)
+message(FATAL_ERROR
+  "This project is intended to be built as part of LLVM via "
+  "-DLLVM_EXTERNAL_PROJECTS=torch-mlir-dialects "
+  "-DLLVM_EXTERNAL_TORCH_MLIR_DIALECTS_SOURCE_DIR=${CMAKE_CURRENT_SOURCE_DIR}")
+endif()
+
+option(MLIR_ENABLE_BINDINGS_PYTHON "Enables MLIR Python Bindings" OFF)
+
+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})")
+
+# TODO: Fix this upstream so that global include directories are not needed.
+set(MLIR_MAIN_SRC_DIR ${LLVM_MAIN_SRC_DIR}/../mlir)
+set(MLIR_INCLUDE_DIR ${LLVM_MAIN_SRC_DIR}/../mlir/include)
+set(MLIR_GENERATED_INCLUDE_DIR ${LLVM_BINARY_DIR}/tools/mlir/include)
+
+# TODO: Needed for tablegen. Remove.
+include_directories(SYSTEM ${MLIR_INCLUDE_DIR})
+include_directories(SYSTEM ${MLIR_GENERATED_INCLUDE_DIR})
+include_directories(SYSTEM ${TORCH_MLIR_DIALECTS_SOURCE_DIR}/include)
+
+function(torch_mlir_dialects_target_includes target)
+  set(_dirs
+    $<BUILD_INTERFACE:${MLIR_INCLUDE_DIR}>
+    $<BUILD_INTERFACE:${MLIR_GENERATED_INCLUDE_DIR}>
+    $<BUILD_INTERFACE:${TORCH_MLIR_DIALECTS_SOURCE_DIR}/include>
+    $<BUILD_INTERFACE:${TORCH_MLIR_DIALECTS_BINARY_DIR}/include>
+  )
+  # In LLVM parlance, the actual target may just be an interface and may not
+  # be responsible for actually compiling anything. The corresponding obj.
+  # target, when present, is just used for compilation and does not
+  # contribute to the interface properties.
+  # TODO: Normalize this upstream.
+  target_include_directories(${target} PUBLIC ${_dirs})
+  if(TARGET obj.${target})
+    target_include_directories(obj.${target} PRIVATE ${_dirs})
+  endif()
+endfunction()
+
+# Configure CMake and tablegen.
+list(APPEND CMAKE_MODULE_PATH ${MLIR_MAIN_SRC_DIR}/cmake/modules)
+list(APPEND CMAKE_MODULE_PATH ${LLVM_MAIN_SRC_DIR}/cmake)
+set(MLIR_TABLEGEN_EXE mlir-tblgen)
+
+include(TableGen)
+include(AddLLVM)
+include(AddMLIR)
+
+add_subdirectory(include)
+add_subdirectory(lib)
+add_subdirectory(tools)
+add_subdirectory(test)

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

@@ -0,0 +1 @@
+add_subdirectory(torch-mlir-dialects)

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

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

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

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

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

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

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

@@ -0,0 +1,33 @@
+function(_add_interfaces)
+  set(LLVM_TARGET_DEFINITIONS TMTensorInterfaces.td)
+  mlir_tablegen(TMTensorOpInterfaces.h.inc -gen-op-interface-decls)
+  mlir_tablegen(TMTensorOpInterfaces.cpp.inc -gen-op-interface-defs)
+  mlir_tablegen(TMTensorTypeInterfaces.h.inc -gen-type-interface-decls)
+  mlir_tablegen(TMTensorTypeInterfaces.cpp.inc -gen-type-interface-defs)
+  add_public_tablegen_target(TorchMLIRTMTensorInterfacesIncGen)
+  add_dependencies(TorchMLIRTMTensorOpsIncGen TorchMLIRTMTensorInterfacesIncGen)
+endfunction()
+
+function(_add_scalar_loop_op_interface)
+  set(LLVM_TARGET_DEFINITIONS ScalarLoopOpInterface.td)
+  mlir_tablegen(ScalarLoopOpInterface.h.inc -gen-op-interface-decls)
+  mlir_tablegen(ScalarLoopOpInterface.cpp.inc -gen-op-interface-defs)
+  add_public_tablegen_target(TorchMLIRTMTensorScalarLoopOpInterfaceIncGen)
+  add_dependencies(TorchMLIRTMTensorOpsIncGen TorchMLIRTMTensorScalarLoopOpInterfaceIncGen)
+endfunction()
+
+function(_add_dialect)
+  set(LLVM_TARGET_DEFINITIONS TMTensorOps.td)
+  mlir_tablegen(TMTensorOps.h.inc -gen-op-decls)
+  mlir_tablegen(TMTensorOps.cpp.inc -gen-op-defs)
+  mlir_tablegen(TMTensorTypes.h.inc -gen-typedef-decls)
+  mlir_tablegen(TMTensorTypes.cpp.inc -gen-typedef-defs)
+  mlir_tablegen(TMTensorDialect.h.inc -gen-dialect-decls -dialect=tm_tensor)
+  mlir_tablegen(TMTensorDialect.cpp.inc -gen-dialect-defs -dialect=tm_tensor)
+  add_public_tablegen_target(TorchMLIRTMTensorOpsIncGen)
+  add_dependencies(mlir-headers TorchMLIRTMTensorOpsIncGen)
+endfunction()
+
+_add_dialect()
+_add_interfaces()
+_add_scalar_loop_op_interface()

external/llvm-external-projects/torch-mlir-dialects/include/torch-mlir-dialects/Dialect/TMTensor/IR/ScalarLoopOpInterface.h

@@ -0,0 +1,29 @@
+//===------------------------------------------------------------*- 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_TMTENSOR_IR_SCALARLOOPOPINTERFACE_H_
+#define TORCH_MLIR_DIALECTS_DIALECT_TMTENSOR_IR_SCALARLOOPOPINTERFACE_H_
+
+#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/Interfaces/ViewLikeInterface.h"
+#include "mlir/Support/LLVM.h"
+
+/// Include the ODS generated interface header files.
+#include "torch-mlir-dialects/Dialect/TMTensor/IR/ScalarLoopOpInterface.h.inc"
+
+namespace mlir {
+namespace torch {
+namespace TMTensor {} // namespace TMTensor
+} // namespace torch
+} // namespace mlir
+
+#endif // TORCH_MLIR_DIALECTS_DIALECT_TMTENSOR_IR_SCALARLOOPOPINTERFACE_H_

external/llvm-external-projects/torch-mlir-dialects/include/torch-mlir-dialects/Dialect/TMTensor/IR/ScalarLoopOpInterface.td

@@ -0,0 +1,79 @@
+//===-------------------------------------------------------*- 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_TMTENSOR_SCALARLOOPOPINTERFACE
+#define TORCH_MLIR_DIALECT_TMTENSOR_SCALARLOOPOPINTERFACE
+
+include "mlir/IR/OpBase.td"
+
+def ScalarLoopOpInterface : OpInterface<"ScalarLoopOpInterface"> {
+  let description = [{
+    Interface for allowing operations to expose information needed to
+    lower it to for loops
+  }];
+  let cppNamespace = "::mlir::torch::TMTensor";
+  let methods = [
+      InterfaceMethod<
+        /*desc=*/[{
+          Returns the destination operands. For op with `memref`
+          operands, this is the result buffers. For op with `tensor`
+          operands, this is the operands that contain the initial
+          value for the result. These are "tied" to the result
+          buffers. For example, for a `LinalgOp`/`TMTensor` ops, it
+          is the `outs` parameters. For `tensor.insert_slice`, it is
+          the `dest` parameter.
+        }],
+        /*retType=*/"SmallVector<Value>",
+        /*methodName=*/"getDestinationOperands",
+        /*args=*/(ins "OpBuilder &":$b),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/"return ValueRange{};"
+      >,
+      InterfaceMethod<
+        /*desc=*/[{
+          Returns a list of `StringRef`s that describe the number of
+          loops and the iterator types of the operation. The list is
+          expected to use
+          `getParallelIteratorTypeName()`/`getReductionIteratorTypeName()`
+          from MLIR Structured Op Utils.
+        }],
+        /*retType=*/"SmallVector<StringRef>",
+        /*methodName=*/"getLoopIteratorTypes"
+      >,
+      InterfaceMethod<
+        /*desc=*/[{
+          Returns a list of ranges that describe the loop bounds and
+          step for the loops of the operation.
+        }],
+        /*retTy=*/"SmallVector<Range>",
+        /*methodName=*/"getIterationDomain",
+        /*args=*/(ins "OpBuilder &":$b)
+      >,
+      InterfaceMethod<
+        /*desc=*/[{
+          Generates the loop body implementation. Assume that all the parallel
+          loops and reduction loops are created and the insertion point of the
+          build is set to the innermost of the loop. This method implements the
+          loop body IRs.
+        }],
+        /*retType=*/"LogicalResult",
+        /*methodName=*/"generateScalarImplementation",
+        /*args=*/(ins
+            "OpBuilder &":$b,
+            "Location ":$loc,
+            "ValueRange ":$ivs),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/[{
+          return failure();
+        }]
+      >
+  ];
+}
+
+#endif  // TORCH_MLIR_DIALECT_TMTENSOR_SCALARLOOPOPINTERFACES

external/llvm-external-projects/torch-mlir-dialects/include/torch-mlir-dialects/Dialect/TMTensor/IR/TMTensorBase.td

@@ -0,0 +1,59 @@
+//===-------------------------------------------------------*- 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_TMTENSOR_BASE
+#define TORCH_MLIR_DIALECT_TMTENSOR_BASE
+
+include "mlir/IR/OpBase.td"
+
+//===----------------------------------------------------------------------===//
+// Dialect definition
+//===----------------------------------------------------------------------===//
+
+def TMTensor_Dialect : Dialect {
+  let name = "tm_tensor";
+  let cppNamespace = "::mlir::torch::TMTensor";
+  let description = [{
+    The tm_tensor (tm = torch-mlir) dialect is a temporary staging ground in
+    the torch-mlir project for a set of widely-accepted tensor compute
+    operations that are not well-served by existing representations in MLIR
+    upstream. These ops are currently heavily inspired by the linalg_ext
+    dialect (which itself is heavily inspired by the structured ops of the
+    linalg dialect). But while linalg_ext is meant to power specific codegen
+    transformations, the tm_tensor dialect is a much more pure "interface
+    dialect" agnostic to any particular set of transformations applied to
+    the operations. We simply require a way to name the specified operations
+    for interchange between projects, without taking strong opinions on the
+    mechanics of transformations.
+
+    The dialect does include interfaces to generate scalar reference code for
+    the operations, which simultaneously provides a precise definition of their
+    semantics, and aids in producing executable reference implementations of
+    the operations.
+
+    The goal of this dialect is to eventually either be upstreamed or to be
+    subsumed by functionality included by upstream MLIR. It should also be kept
+    consistent with the linalg_ext dialect unless there is a good reason not
+    to.
+  }];
+  let hasCanonicalizer = 1;
+}
+
+//===----------------------------------------------------------------------===//
+// Type definitions
+//===----------------------------------------------------------------------===//
+
+class RankedTensorOrMemRefOf<list<Type> allowedTypes> :
+  ShapedContainerType<allowedTypes,
+      Or<[IsMemRefTypePred, And<[IsTensorTypePred, HasRankPred]>]>,
+  "ranked tensor or memref", "::mlir::ShapedType">;
+
+def AnyRankedTensorOrMemRefType : RankedTensorOrMemRefOf<[AnyType]>;
+
+#endif // TORCH_MLIR_DIALECT_TMTENSOR_BASE

external/llvm-external-projects/torch-mlir-dialects/include/torch-mlir-dialects/Dialect/TMTensor/IR/TMTensorDialect.h

@@ -0,0 +1,20 @@
+//===------------------------------------------------------------*- 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_TMTENSOR_IR_TMTENSORDIALECT_H_
+#define TORCH_MLIR_DIALECTS_DIALECT_TMTENSOR_IR_TMTENSORDIALECT_H_
+
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/OpDefinition.h"
+
+// clang-format off: must be included after all LLVM/MLIR headers
+#include "torch-mlir-dialects/Dialect/TMTensor/IR/TMTensorDialect.h.inc" // IWYU pragma: keep
+// clang-format on
+
+#endif // TORCH_MLIR_DIALECTS_DIALECT_TMTENSOR_IR_TMTENSORDIALECT_H_

external/llvm-external-projects/torch-mlir-dialects/include/torch-mlir-dialects/Dialect/TMTensor/IR/TMTensorInterfaces.h

@@ -0,0 +1,42 @@
+//===------------------------------------------------------------*- 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_TMTENSOR_IR_TMTENSORINTERFACES_H_
+#define TORCH_MLIR_DIALECTS_DIALECT_TMTENSOR_IR_TMTENSORINTERFACES_H_
+
+#include "mlir/IR/BlockAndValueMapping.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/Support/LLVM.h"
+
+namespace mlir {
+namespace torch {
+namespace TMTensor {
+class TMTensorOp;
+
+/// OpOperand vector that implicitly converts to a Value vector.
+struct OpOperandVector : public SmallVector<OpOperand *> {
+  operator SmallVector<Value>();
+};
+
+namespace detail {
+LogicalResult verifyTMTensorOpInterface(Operation *op);
+}
+
+#include "torch-mlir-dialects/Dialect/TMTensor/IR/TMTensorOps.h.inc" // IWYU pragma: export
+
+/// Include the generated interface declarations.
+#include "torch-mlir-dialects/Dialect/TMTensor/IR/TMTensorOpInterfaces.h.inc" // IWYU pragma: export
+
+} // namespace TMTensor
+} // namespace torch
+} // namespace mlir
+
+#endif // TORCH_MLIR_DIALECTS_DIALECT_TMTENSOR_IR_TMTENSORINTERFACES_H_