Update LinalgExt to encompass iree-llvm-sandbox additions. (#8554)
Moves LinalgExt related developments in `iree-llvm-sandbox` into `iree-dialects`. With this revision, iree-llvm-sandbox is able to integrate more closely with IREE.
See PSA https://github.com/google/iree-llvm-sandbox/issues/373
Co-authored-by: Mahesh Ravishankar <ravishankarm@google.com>
diff --git a/build_tools/bazel_to_cmake/bazel_to_cmake_targets.py b/build_tools/bazel_to_cmake/bazel_to_cmake_targets.py
index aad06d4..e35479c 100644
--- a/build_tools/bazel_to_cmake/bazel_to_cmake_targets.py
+++ b/build_tools/bazel_to_cmake/bazel_to_cmake_targets.py
@@ -19,7 +19,7 @@
"//llvm-external-projects/iree-dialects:IREELinalgExtDialect": [
"IREELinalgExtDialect"
],
- "//llvm-external-projects/iree-dialects:IREELinalgExtTransforms": [
+ "//llvm-external-projects/iree-dialects:IREELinalgExtPasses": [
"IREELinalgExtPasses"
],
"//llvm-external-projects/iree-dialects:IREEPyDMDialect": [
diff --git a/iree/compiler/Codegen/Common/BUILD b/iree/compiler/Codegen/Common/BUILD
index 5859f09..80649f3 100644
--- a/iree/compiler/Codegen/Common/BUILD
+++ b/iree/compiler/Codegen/Common/BUILD
@@ -71,7 +71,7 @@
"//iree/compiler/Dialect/HAL/IR",
"//iree/compiler/Dialect/Util/IR",
"//llvm-external-projects/iree-dialects:IREELinalgExtDialect",
- "//llvm-external-projects/iree-dialects:IREELinalgExtTransforms",
+ "//llvm-external-projects/iree-dialects:IREELinalgExtPasses",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:Affine",
"@llvm-project//mlir:AffineBufferizableOpInterfaceImpl",
diff --git a/iree/compiler/Codegen/Common/TileAndDistributeToWorkgroupsPass.cpp b/iree/compiler/Codegen/Common/TileAndDistributeToWorkgroupsPass.cpp
index 40eaef0..176ed5a 100644
--- a/iree/compiler/Codegen/Common/TileAndDistributeToWorkgroupsPass.cpp
+++ b/iree/compiler/Codegen/Common/TileAndDistributeToWorkgroupsPass.cpp
@@ -14,7 +14,7 @@
//
//===---------------------------------------------------------------------===//
#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.h"
-#include "iree-dialects/Dialect/LinalgExt/Transforms/Transforms.h"
+#include "iree-dialects/Dialect/LinalgExt/Passes/Transforms.h"
#include "iree/compiler/Codegen/Common/DestructiveUpdateUtils.h"
#include "iree/compiler/Codegen/Dialect/LoweringConfig.h"
#include "iree/compiler/Codegen/PassDetail.h"
diff --git a/iree/compiler/Codegen/LLVMCPU/BUILD b/iree/compiler/Codegen/LLVMCPU/BUILD
index e101f53..2a7956d 100644
--- a/iree/compiler/Codegen/LLVMCPU/BUILD
+++ b/iree/compiler/Codegen/LLVMCPU/BUILD
@@ -41,7 +41,7 @@
"//iree/compiler/Dialect/Util/IR",
"//iree/compiler/Utils",
"//llvm-external-projects/iree-dialects:IREELinalgExtDialect",
- "//llvm-external-projects/iree-dialects:IREELinalgExtTransforms",
+ "//llvm-external-projects/iree-dialects:IREELinalgExtPasses",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:AffineToStandardTransforms",
"@llvm-project//mlir:Analysis",
diff --git a/iree/compiler/Codegen/LLVMCPU/Passes.cpp b/iree/compiler/Codegen/LLVMCPU/Passes.cpp
index 9d91798..1d376b2 100644
--- a/iree/compiler/Codegen/LLVMCPU/Passes.cpp
+++ b/iree/compiler/Codegen/LLVMCPU/Passes.cpp
@@ -7,7 +7,7 @@
#include "iree/compiler/Codegen/Passes.h"
#include "iree-dialects/Dialect/LinalgExt/IR/TiledOpInterface.h"
-#include "iree-dialects/Dialect/LinalgExt/Transforms/Passes.h"
+#include "iree-dialects/Dialect/LinalgExt/Passes/Passes.h"
#include "iree/compiler/Codegen/LLVMCPU/KernelDispatch.h"
#include "iree/compiler/Codegen/PassDetail.h"
#include "iree/compiler/Codegen/Sandbox/Passes.h"
diff --git a/iree/compiler/Codegen/LLVMGPU/BUILD b/iree/compiler/Codegen/LLVMGPU/BUILD
index 87af66b..9ffbdf1 100644
--- a/iree/compiler/Codegen/LLVMGPU/BUILD
+++ b/iree/compiler/Codegen/LLVMGPU/BUILD
@@ -45,7 +45,7 @@
"//iree/compiler/Dialect/HAL/IR",
"//iree/compiler/Dialect/Util/IR",
"//llvm-external-projects/iree-dialects:IREELinalgExtDialect",
- "//llvm-external-projects/iree-dialects:IREELinalgExtTransforms",
+ "//llvm-external-projects/iree-dialects:IREELinalgExtPasses",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:Affine",
"@llvm-project//mlir:AffineToStandard",
diff --git a/iree/compiler/Codegen/LLVMGPU/LLVMGPUTileAndDistribute.cpp b/iree/compiler/Codegen/LLVMGPU/LLVMGPUTileAndDistribute.cpp
index 9cd86f6..4325d90 100644
--- a/iree/compiler/Codegen/LLVMGPU/LLVMGPUTileAndDistribute.cpp
+++ b/iree/compiler/Codegen/LLVMGPU/LLVMGPUTileAndDistribute.cpp
@@ -5,7 +5,7 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.h"
-#include "iree-dialects/Dialect/LinalgExt/Transforms/Transforms.h"
+#include "iree-dialects/Dialect/LinalgExt/Passes/Transforms.h"
#include "iree/compiler/Codegen/Dialect/LoweringConfig.h"
#include "iree/compiler/Codegen/LLVMGPU/KernelConfig.h"
#include "iree/compiler/Codegen/LLVMGPU/LLVMGPUUtils.h"
diff --git a/iree/compiler/Codegen/LLVMGPU/Passes.cpp b/iree/compiler/Codegen/LLVMGPU/Passes.cpp
index 0cb1e6a..7c1ec65 100644
--- a/iree/compiler/Codegen/LLVMGPU/Passes.cpp
+++ b/iree/compiler/Codegen/LLVMGPU/Passes.cpp
@@ -6,7 +6,7 @@
#include "iree/compiler/Codegen/Passes.h"
-#include "iree-dialects/Dialect/LinalgExt/Transforms/Passes.h"
+#include "iree-dialects/Dialect/LinalgExt/Passes/Passes.h"
#include "iree/compiler/Codegen/PassDetail.h"
#include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
#include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"
diff --git a/iree/compiler/Codegen/SPIRV/BUILD b/iree/compiler/Codegen/SPIRV/BUILD
index c3a6370..dbb19b1 100644
--- a/iree/compiler/Codegen/SPIRV/BUILD
+++ b/iree/compiler/Codegen/SPIRV/BUILD
@@ -49,7 +49,7 @@
"//iree/compiler/Dialect/HAL/IR:HALDialect",
"//iree/compiler/Dialect/Util/IR",
"//llvm-external-projects/iree-dialects:IREELinalgExtDialect",
- "//llvm-external-projects/iree-dialects:IREELinalgExtTransforms",
+ "//llvm-external-projects/iree-dialects:IREELinalgExtPasses",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:Affine",
"@llvm-project//mlir:AffineAnalysis",
diff --git a/iree/compiler/Codegen/SPIRV/Passes.cpp b/iree/compiler/Codegen/SPIRV/Passes.cpp
index 954ecb5..fd600d8 100644
--- a/iree/compiler/Codegen/SPIRV/Passes.cpp
+++ b/iree/compiler/Codegen/SPIRV/Passes.cpp
@@ -13,7 +13,7 @@
#include "iree/compiler/Codegen/Passes.h"
-#include "iree-dialects/Dialect/LinalgExt/Transforms/Passes.h"
+#include "iree-dialects/Dialect/LinalgExt/Passes/Passes.h"
#include "iree/compiler/Codegen/PassDetail.h"
#include "iree/compiler/Codegen/Passes.h"
#include "iree/compiler/Codegen/SPIRV/MemorySpace.h"
diff --git a/iree/compiler/Codegen/SPIRV/SPIRVTileAndDistribute.cpp b/iree/compiler/Codegen/SPIRV/SPIRVTileAndDistribute.cpp
index 86931f0..138d50b 100644
--- a/iree/compiler/Codegen/SPIRV/SPIRVTileAndDistribute.cpp
+++ b/iree/compiler/Codegen/SPIRV/SPIRVTileAndDistribute.cpp
@@ -12,7 +12,7 @@
//===----------------------------------------------------------------------===//
#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.h"
-#include "iree-dialects/Dialect/LinalgExt/Transforms/Transforms.h"
+#include "iree-dialects/Dialect/LinalgExt/Passes/Transforms.h"
#include "iree/compiler/Codegen/Dialect/LoweringConfig.h"
#include "iree/compiler/Codegen/PassDetail.h"
#include "iree/compiler/Codegen/Passes.h"
diff --git a/iree/compiler/Dialect/Flow/Transforms/BUILD b/iree/compiler/Dialect/Flow/Transforms/BUILD
index 22205dd..6347d95 100644
--- a/iree/compiler/Dialect/Flow/Transforms/BUILD
+++ b/iree/compiler/Dialect/Flow/Transforms/BUILD
@@ -76,7 +76,7 @@
"//iree/compiler/Dialect/Util/Transforms",
"//iree/compiler/Utils",
"//llvm-external-projects/iree-dialects:IREELinalgExtDialect",
- "//llvm-external-projects/iree-dialects:IREELinalgExtTransforms",
+ "//llvm-external-projects/iree-dialects:IREELinalgExtPasses",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:Affine",
"@llvm-project//mlir:ArithmeticDialect",
diff --git a/iree/compiler/Dialect/Flow/Transforms/DispatchLinalgOnTensors.cpp b/iree/compiler/Dialect/Flow/Transforms/DispatchLinalgOnTensors.cpp
index 333b32c..9b92be2 100644
--- a/iree/compiler/Dialect/Flow/Transforms/DispatchLinalgOnTensors.cpp
+++ b/iree/compiler/Dialect/Flow/Transforms/DispatchLinalgOnTensors.cpp
@@ -5,7 +5,7 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.h"
-#include "iree-dialects/Dialect/LinalgExt/Transforms/Transforms.h"
+#include "iree-dialects/Dialect/LinalgExt/Passes/Transforms.h"
#include "iree/compiler/Dialect/Flow/IR/FlowDialect.h"
#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
#include "iree/compiler/Dialect/Flow/IR/FlowTypes.h"
diff --git a/iree/compiler/Dialect/Modules/VMVX/Transforms/BUILD b/iree/compiler/Dialect/Modules/VMVX/Transforms/BUILD
index 42ed489..72631ac 100644
--- a/iree/compiler/Dialect/Modules/VMVX/Transforms/BUILD
+++ b/iree/compiler/Dialect/Modules/VMVX/Transforms/BUILD
@@ -32,7 +32,7 @@
"//iree/compiler/Dialect/Util/IR",
"//iree/compiler/Dialect/Util/Transforms",
"//iree/compiler/Dialect/VM/IR",
- "//llvm-external-projects/iree-dialects:IREELinalgExtTransforms",
+ "//llvm-external-projects/iree-dialects:IREELinalgExtPasses",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:Affine",
"@llvm-project//mlir:AffineToStandardTransforms",
diff --git a/iree/compiler/Dialect/Modules/VMVX/Transforms/Passes.cpp b/iree/compiler/Dialect/Modules/VMVX/Transforms/Passes.cpp
index 2797737..5a96054 100644
--- a/iree/compiler/Dialect/Modules/VMVX/Transforms/Passes.cpp
+++ b/iree/compiler/Dialect/Modules/VMVX/Transforms/Passes.cpp
@@ -8,7 +8,7 @@
#include <memory>
-#include "iree-dialects/Dialect/LinalgExt/Transforms/Passes.h"
+#include "iree-dialects/Dialect/LinalgExt/Passes/Passes.h"
#include "iree/compiler/Codegen/Passes.h"
#include "iree/compiler/Dialect/HAL/Transforms/Passes.h"
#include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
diff --git a/iree/compiler/InputConversion/MHLO/BUILD b/iree/compiler/InputConversion/MHLO/BUILD
index e8093bc..ed5fdc7 100644
--- a/iree/compiler/InputConversion/MHLO/BUILD
+++ b/iree/compiler/InputConversion/MHLO/BUILD
@@ -68,7 +68,7 @@
"//iree/compiler/Dialect/Util/IR",
"//iree/compiler/InputConversion/Common",
"//llvm-external-projects/iree-dialects:IREELinalgExtDialect",
- "//llvm-external-projects/iree-dialects:IREELinalgExtTransforms",
+ "//llvm-external-projects/iree-dialects:IREELinalgExtPasses",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:Affine",
"@llvm-project//mlir:AffineUtils",
diff --git a/iree/tools/BUILD b/iree/tools/BUILD
index 6c4b3b2..753c641 100644
--- a/iree/tools/BUILD
+++ b/iree/tools/BUILD
@@ -119,7 +119,7 @@
"//iree/compiler/Translation:IREEVM",
"//llvm-external-projects/iree-dialects:IREEInputDialect",
"//llvm-external-projects/iree-dialects:IREELinalgExtDialect",
- "//llvm-external-projects/iree-dialects:IREELinalgExtTransforms",
+ "//llvm-external-projects/iree-dialects:IREELinalgExtPasses",
"//llvm-external-projects/iree-dialects:IREEPyDMDialect",
"@llvm-project//mlir:IR",
],
diff --git a/iree/tools/init_iree_passes.h b/iree/tools/init_iree_passes.h
index 66ce300..f5e0863 100644
--- a/iree/tools/init_iree_passes.h
+++ b/iree/tools/init_iree_passes.h
@@ -14,7 +14,7 @@
#include <cstdlib>
-#include "iree-dialects/Dialect/LinalgExt/Transforms/Passes.h"
+#include "iree-dialects/Dialect/LinalgExt/Passes/Passes.h"
#include "iree/compiler/Bindings/Native/Transforms/Passes.h"
#include "iree/compiler/Bindings/TFLite/Transforms/Passes.h"
#include "iree/compiler/ConstEval/Passes.h"
diff --git a/llvm-external-projects/iree-dialects/BUILD b/llvm-external-projects/iree-dialects/BUILD
index 7f99cb5..350a209 100644
--- a/llvm-external-projects/iree-dialects/BUILD
+++ b/llvm-external-projects/iree-dialects/BUILD
@@ -32,7 +32,7 @@
srcs = glob([
"include/iree-dialects/Dialect/Input/*.td",
"include/iree-dialects/Dialect/LinalgExt/IR/*.td",
- "include/iree-dialects/Dialect/LinalgExt/Transforms/*.td",
+ "include/iree-dialects/Dialect/LinalgExt/Passes/*.td",
"include/iree-dialects/Dialect/PyDM/IR/*.td",
"include/iree-dialects/Dialect/PyDM/Transforms/*.td",
]),
@@ -43,7 +43,7 @@
srcs = glob([
"include/iree-dialects/Dialect/Input/*.td",
"include/iree-dialects/Dialect/LinalgExt/IR/*.td",
- "include/iree-dialects/Dialect/LinalgExt/Transforms/*.td",
+ "include/iree-dialects/Dialect/LinalgExt/Passes/*.td",
"include/iree-dialects/Dialect/PyDM/IR/*.td",
"python/iree/compiler/dialects/*.td",
]) + [
@@ -175,6 +175,8 @@
":TdFiles",
"@llvm-project//mlir:CallInterfacesTdFiles",
"@llvm-project//mlir:ControlFlowInterfacesTdFiles",
+ "@llvm-project//mlir:TilingInterfaceTdFiles",
+ "@llvm-project//mlir:ViewLikeInterfaceTdFiles",
],
)
@@ -232,19 +234,19 @@
tbl_outs = [
(
["-gen-pass-decls"],
- "include/iree-dialects/Dialect/LinalgExt/Transforms/Passes.h.inc",
+ "include/iree-dialects/Dialect/LinalgExt/Passes/Passes.h.inc",
),
(
["-gen-pass-capi-header"],
- "include/iree-dialects/Dialect/LinalgExt/Transforms/Passes.capi.h.inc",
+ "include/iree-dialects/Dialect/LinalgExt/Passes/Passes.capi.h.inc",
),
(
["-gen-pass-capi-impl"],
- "include/iree-dialects/Dialect/LinalgExt/Transforms/Passes.capi.cpp.inc",
+ "include/iree-dialects/Dialect/LinalgExt/Passes/Passes.capi.cpp.inc",
),
],
tblgen = "@llvm-project//mlir:mlir-tblgen",
- td_file = "include/iree-dialects/Dialect/LinalgExt/Transforms/Passes.td",
+ td_file = "include/iree-dialects/Dialect/LinalgExt/Passes/Passes.td",
deps = [
":TdFiles",
"@llvm-project//mlir:PassBaseTdFiles",
@@ -286,12 +288,12 @@
)
cc_library(
- name = "IREELinalgExtTransforms",
+ name = "IREELinalgExtPasses",
srcs = glob([
- "lib/Dialect/LinalgExt/Transforms/*.cpp",
+ "lib/Dialect/LinalgExt/Passes/*.cpp",
]),
hdrs = glob([
- "include/iree-dialects/Dialect/LinalgExt/Transforms/*.h",
+ "include/iree-dialects/Dialect/LinalgExt/Passes/*.h",
]),
deps = [
":IREEInputDialect",
@@ -502,6 +504,7 @@
includes = ["include"],
deps = [
":IREEInputDialect",
+ ":IREELinalgExtDialect",
":IREEPyDMDialect",
":IREEPyDMTransforms",
"@llvm-project//mlir:CAPIIR",
@@ -523,7 +526,7 @@
deps = [
":IREEInputDialect",
":IREELinalgExtDialect",
- ":IREELinalgExtTransforms",
+ ":IREELinalgExtPasses",
":IREEPyDMDialect",
":IREEPyDMTransforms",
"@llvm-project//llvm:Support",
diff --git a/llvm-external-projects/iree-dialects/include/iree-dialects-c/Dialects.h b/llvm-external-projects/iree-dialects/include/iree-dialects-c/Dialects.h
index 5b5d93d..eb6276b 100644
--- a/llvm-external-projects/iree-dialects/include/iree-dialects-c/Dialects.h
+++ b/llvm-external-projects/iree-dialects/include/iree-dialects-c/Dialects.h
@@ -21,6 +21,12 @@
MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(IREEInput, iree_input);
+//===--------------------------------------------------------------------===//
+// IREELinalgExt
+//===--------------------------------------------------------------------===//
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(IREELinalgExt, iree_linalg_ext);
+
//===----------------------------------------------------------------------===//
// IREEPyDMDialect
//===----------------------------------------------------------------------===//
diff --git a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/CMakeLists.txt b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/CMakeLists.txt
index 9f57627..5a7289b 100644
--- a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/CMakeLists.txt
+++ b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/CMakeLists.txt
@@ -1,2 +1,2 @@
add_subdirectory(IR)
-add_subdirectory(Transforms)
+add_subdirectory(Passes)
diff --git a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.td b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.td
index b182052..75c0cad 100644
--- a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.td
+++ b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.td
@@ -10,8 +10,11 @@
include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtBase.td"
include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtInterfaces.td"
include "iree-dialects/Dialect/LinalgExt/IR/TiledOpInterface.td"
-include "mlir/Interfaces/SideEffectInterfaces.td"
include "mlir/Interfaces/ControlFlowInterfaces.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/Interfaces/TilingInterface.td"
+include "mlir/Interfaces/ViewLikeInterface.td"
+
//===----------------------------------------------------------------------===//
// Base class.
@@ -349,8 +352,8 @@
def IREELinalgExt_YieldOp : IREELinalgExt_PureOp<"yield", [NoSideEffect, ReturnLike, Terminator]> {
let summary = "LinalgExt yield op";
let description = [{
- `linalg_ext.yield` is a special terminator operation for blocks inside
- regions in `linalg_ext` ops.
+ `iree_linalg_ext.yield` is a special terminator operation for blocks inside
+ regions in `iree_linalg_ext` ops.
}];
let arguments = (ins Variadic<AnyType>:$operands);
@@ -362,4 +365,268 @@
let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
}
+//===----------------------------------------------------------------------===//
+// Ops supporting concurrency with tensors.
+//===----------------------------------------------------------------------===//
+
+def IREELinalgExt_TileOp : IREELinalgExt_PureOp<"tile",
+ [
+ // TODO: enable to allow hoisting, LICM and isDefinedOutside
+ // DeclareOpInterfaceMethods<LoopLikeOpInterface>,
+ SingleBlockImplicitTerminator<"::mlir::iree_compiler::IREE::LinalgExt::TileYieldOp">,
+ RecursiveSideEffects
+ ]> {
+ let summary = "tile operation";
+ let description = [{
+ `iree_linalg_ext.tile` is a 1-D loop construct that operates on tensors and
+ evaluates its body once for each tile. The number and size of tiles is
+ specified by the `tile_size` operand.
+
+ The `tile` op takes a list of destination-passing style tensors and returns
+ a matching list of tensors of the same size.
+
+ Every instance of the body is expected to return a tile with leading
+ dimension matching the corresponding tile size.
+
+ The default terminator behavior is such that tiles yielded by individual
+ iterations are concatenated along the `tiled_dim` dimension.
+ This is the canonical way to perform "subset insertions".
+ Note, if `tiled_dim` has the value `0`, it may be elided from pretty
+ pinting and parsing.
+
+ All return tiles are concatenated into forming the matching full result
+ tensor according to the terminator.
+
+ When the `tile_size` operand is a `tensor<..index>`, the `tile` op
+ evaluates its body `dim(tile_size, 0)` times. Each iteration `i` produces a
+ tile of leading size `tile_size[i]`.
+
+ The induced `offset` block argument captures the running sum of `tile_size`
+ for all the previous iterations.
+
+ When the `tile_size` operand is a single index, it is interpreted as a
+ sequence of tile sizes given by the following formula:
+ ```
+ N = tensor.dim(...)
+ S = sizes
+ T, R = divmod(N, S)
+ [T] * S + ([R] if R != 0 else [])
+ ```
+
+ All tiles except the last are of the same size.
+ }];
+ let arguments = (ins AnyTypeOf<[// TODO: allow TensorOf<[Index]>,
+ Index]>:$tile_size,
+ Variadic<AnyRankedTensor>:$outs,
+ I64Attr:$tiled_dim);
+ let results = (outs Variadic<AnyType>:$results);
+ let regions = (region SizedRegion<1>:$region);
+ let skipDefaultBuilders = 1;
+ let builders = [
+ // Builder that builds a tile on the implicit first dimension (i.e. `0`).
+ OpBuilder<(ins "Value":$tileSizes, "ValueRange":$outs,
+ CArg<"function_ref<void(OpBuilder &, Location, Value, Value, ValueRange)>",
+ "nullptr">)>,
+ // Builder that builds a tile with a specified integral dimension.
+ OpBuilder<(ins "Value":$tileSizes, "ValueRange":$outs, "int64_t":$tiledDims,
+ CArg<"function_ref<void(OpBuilder &, Location, Value, Value, ValueRange)>",
+ "nullptr">)>,
+ ];
+
+ let extraClassDeclaration = [{
+ static StringRef getTiledDimAttrName() { return "tiled_dim";}
+ using TileOpBodyBuilderFn =
+ function_ref<void(OpBuilder &, Location, Value /*offset*/, Value /*size*/,
+ ValueRange /*outs*/)>;
+ // TODO: helper for getting named region args without magic constants etc.
+ }];
+
+ let hasCustomAssemblyFormat = 1;
+ let hasVerifier = 1;
+}
+
+def IREELinalgExt_TileYieldOp : IREELinalgExt_PureOp<"tile_yield", [
+ NoSideEffect, ReturnLike, Terminator]> {
+ let summary = "LinalgExt tile_yield op";
+ let description = [{
+ `iree_linalg_ext.tile_yield` is a special terminator operation for blocks inside
+ regions in `iree_linalg_ext.tile`.
+ The tiles yielded by individual iterations are concatenated along the first
+ dimension. This is the canonical way to perform "subset insertions"
+ (TODO: allow dim permutations).
+ }];
+
+ let arguments = (ins Variadic<AnyType>:$operands);
+
+ let builders = [
+ OpBuilder<(ins), [{ /* nothing to do */ }]>,
+ ];
+
+ let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
+}
+
+def IREELinalgExt_InParallelOp : IREELinalgExt_PureOp<"in_parallel", [
+ SingleBlockImplicitTerminator<"::mlir::iree_compiler::IREE::LinalgExt::PerformConcurrentlyOp">,
+ RecursiveSideEffects,
+ AutomaticAllocationScope,
+ ]> {
+ let summary = "evaluate a block multiple times in parallel";
+ let description = [{
+ `iree_linalg_ext.in_parallel` is a target-independent parallel function application
+ operation. It has exactly one block that represents the parallel function body
+ and it takes a single index operand that indicates how many parallel instances
+ of that function should get instantiated.
+
+ When the parallel function body is pure (i.e. has no side effects) then the only
+ allowed terminator is `iree_linalg_ext.perform_concurrently`, which dictates
+ how the results of all parallel invocations should be reconciled into a full
+ value that will be returned from `in_parallel`. Multi-value returns are encoded
+ by including multiple operations inside the `perform_concurrently` block.
+
+ When the parallel function body has side effects, the order of reads and writes
+ to memory is unspecified across iterations.
+
+ This op resembles `scf.for` to a large degree, but crucially differs in that it
+ (1) doesn't have `iter_args` and (2) has a special terminator, both of which
+ enable reasoning about its parallel semantics. Another difference is that
+ `in_parallel` always iterates over a range between 0 and an upper bound, but
+ that's insignificant.
+ }];
+ let arguments = (ins Index:$num_threads);
+
+ let results = (outs Variadic<AnyType>:$results);
+ let regions = (region SizedRegion<1>:$region);
+
+ let hasCustomAssemblyFormat = 1;
+ let hasVerifier = 1;
+
+ // The default builder does not add the proper body BBargs, roll our own.
+ let skipDefaultBuilders = 1;
+ let builders = [
+ // Bodyless builder, result types must be specified.
+ OpBuilder<(ins "TypeRange":$resultTypes, "Value":$num_threads)>,
+ // Builder that takes a bodyBuilder lambda, result types are inferred from
+ // the terminator.
+ OpBuilder<(ins "Value":$num_threads,
+ "function_ref<void(OpBuilder &, Location, Value)>":$bodyBuilder)>
+ ];
+ let extraClassDeclaration = [{
+ Value getThreadIndex() { return getBody()->getArgument(0); }
+ static void ensureTerminator(Region ®ion, Builder &builder, Location loc);
+ PerformConcurrentlyOp getTerminator();
+ }];
+}
+
+def IREELinalgExt_PerformConcurrentlyOp : IREELinalgExt_PureOp<"perform_concurrently", [
+ NoSideEffect,
+ Terminator,
+ SingleBlockImplicitTerminator<"::mlir::iree_compiler::IREE::LinalgExt::EndPerformConcurrentlyOp">,
+ ]> {
+ let summary = "terminates a `in_parallel` block";
+ let description = [{
+ `iree_linalg_ext.perform_concurrently` is a designated terminator for the blocks
+ of `iree_linalg_ext.in_parallel` operations. The terminator contains a single block
+ itself, which describes how the results of each parallel invocation are to be
+ reconciled into a single value to be returned from the parallel invocation.
+ One operation in this terminator's block corresponds to a single return of
+ `in_parallel`.
+ }];
+
+ let regions = (region SizedRegion<1>:$region);
+
+ let hasCustomAssemblyFormat = 1;
+ let hasVerifier = 1;
+
+ // TODO(apaszke, ntv): Add an interface for ops that can appear inside
+ // perform_concurrently.
+ let extraClassDeclaration = [{
+ SmallVector<Type> yieldedTypes();
+ SmallVector<ParallelInsertSliceOp> yieldingOps();
+ }];
+}
+
+def IREELinalgExt_EndPerformConcurrentlyOp : IREELinalgExt_PureOp<"end_perform_concurrently", [
+ NoSideEffect, Terminator]> {
+ let summary = "terminates a `perform_concurrently` block";
+ let description = [{
+ A designated terminator for `perform_concurrently`. It's not expected to appear
+ in the textual form of the IR.
+ }];
+}
+
+def IREELinalgExt_ParallelInsertSliceOp : IREELinalgExt_PureOp<"parallel_insert_slice", [
+ AttrSizedOperandSegments, OffsetSizeAndStrideOpInterface]> {
+ let summary = "updates slices of a tensor concurrently";
+ let description = [{
+ Updates slices of a full tensor with multiple sub-slices concurrently.
+
+ Conflicting writes result in undefined semantics, in that the indices written
+ to by multiple parallel updates might contain data from any of the updates, or
+ even a malformed bit pattern (in reality the semantics might end up depending
+ on the memory model of the parallel hardware that `in_parallel` will be lowered to).
+
+ If an index is updated by exactly one updates, the value contained at that index
+ in the resulting tensor will be equal to the value at a corresponding index of a
+ slice that was used for the updated. If an index is not updated at all, its value
+ will be equal to the one in the original tensor.
+
+ Note that we cannot mark this operation as pure (NoSideEffects), even
+ though it has no side effects, because it will get DCEd during
+ canonicalization. Ideally we would use attributes instead of those funny
+ terminating ops, but attributes cannot refer to SSA values at the moment, so
+ it's the best we can do for now.
+ }];
+
+ let arguments = (ins
+ AnyRankedTensor:$source,
+ AnyRankedTensor:$dest,
+ Variadic<Index>:$offsets,
+ Variadic<Index>:$sizes,
+ Variadic<Index>:$strides,
+ I64ArrayAttr:$static_offsets,
+ I64ArrayAttr:$static_sizes,
+ I64ArrayAttr:$static_strides
+ );
+ let assemblyFormat = [{
+ $source `into` $dest ``
+ custom<OperandsOrIntegersOffsetsOrStridesList>($offsets, $static_offsets)
+ custom<OperandsOrIntegersSizesList>($sizes, $static_sizes)
+ custom<OperandsOrIntegersOffsetsOrStridesList>($strides, $static_strides)
+ attr-dict `:` type($source) `into` type($dest)
+ }];
+
+ let extraClassDeclaration = [{
+ Type yieldedType() { return dest().getType(); }
+
+ RankedTensorType getSourceType() {
+ return source().getType().cast<RankedTensorType>();
+ }
+
+ /// Return the expected rank of each of the `static_offsets`, `static_sizes`
+ /// and `static_strides` attributes.
+ std::array<unsigned, 3> getArrayAttrMaxRanks() {
+ unsigned rank = getSourceType().getRank();
+ return {rank, rank, rank};
+ }
+
+ /// Return the number of leading operands before `offsets`, `sizes` and
+ /// `strides` operands.
+ static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 1; }
+ }];
+
+ let builders = [
+ // Build a ParallelInsertSliceOp with mixed static and dynamic entries.
+ OpBuilder<(ins "Value":$source, "Value":$dest,
+ "ArrayRef<OpFoldResult>":$offsets, "ArrayRef<OpFoldResult>":$sizes,
+ "ArrayRef<OpFoldResult>":$strides,
+ CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,
+ // Build a ParallelInsertSliceOp with dynamic entries.
+ OpBuilder<(ins "Value":$source, "Value":$dest,
+ "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides,
+ CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>
+ ];
+
+ let hasCanonicalizer = 1;
+}
+
#endif // IREE_DIALECT_LINALGEXT_OPS
diff --git a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/LinalgExtBufferization.h b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/LinalgExtBufferization.h
new file mode 100644
index 0000000..c1b60b6
--- /dev/null
+++ b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/LinalgExtBufferization.h
@@ -0,0 +1,27 @@
+//===-- LinalgExtBufferization.h - Linalg Extension bufferization ---------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef IREE_DIALECTS_DIALECT_LINALGEXT_BUFFERIZATION_H_
+#define IREE_DIALECTS_DIALECT_LINALGEXT_BUFFERIZATION_H_
+
+namespace mlir {
+
+class DialectRegistry;
+
+namespace iree_compiler {
+namespace IREE {
+namespace LinalgExt {
+
+void registerBufferizableOpInterfaceExternalModels(DialectRegistry ®istry);
+
+} // namespace LinalgExt
+} // namespace IREE
+} // namespace iree_compiler
+} // namespace mlir
+
+#endif // IREE_DIALECTS_DIALECT_LINALGEXT_BUFFERIZATION_H_
diff --git a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Transforms/CMakeLists.txt b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Passes/CMakeLists.txt
similarity index 75%
rename from llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Transforms/CMakeLists.txt
rename to llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Passes/CMakeLists.txt
index 29737fc..07379ca 100644
--- a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Transforms/CMakeLists.txt
+++ b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Passes/CMakeLists.txt
@@ -2,4 +2,4 @@
mlir_tablegen(Passes.h.inc -gen-pass-decls)
mlir_tablegen(Passes.capi.h.inc -gen-pass-capi-header)
mlir_tablegen(Passes.capi.cpp.inc -gen-pass-capi-impl)
-add_public_tablegen_target(IREELinalgExtTransformsPassesIncGen)
+add_public_tablegen_target(IREELinalgExtPassesIncGen)
diff --git a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Transforms/PassDetail.h b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Passes/PassDetail.h
similarity index 83%
rename from llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Transforms/PassDetail.h
rename to llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Passes/PassDetail.h
index 4457f87..e5c044d 100644
--- a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Transforms/PassDetail.h
+++ b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Passes/PassDetail.h
@@ -9,7 +9,7 @@
namespace LinalgExt {
#define GEN_PASS_CLASSES
-#include "iree-dialects/Dialect/LinalgExt/Transforms/Passes.h.inc" // IWYU pragma: keep
+#include "iree-dialects/Dialect/LinalgExt/Passes/Passes.h.inc" // IWYU pragma: keep
} // namespace LinalgExt
} // namespace IREE
diff --git a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Transforms/Passes.h b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Passes/Passes.h
similarity index 92%
rename from llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Transforms/Passes.h
rename to llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Passes/Passes.h
index d3f6144..fb857f3 100644
--- a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Transforms/Passes.h
+++ b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Passes/Passes.h
@@ -20,6 +20,8 @@
std::unique_ptr<OperationPass<>> createPadContractionToBlockSizePass();
+void registerTilingInterfaceExternalModels(DialectRegistry ®istry);
+
void registerPasses();
} // namespace LinalgExt
diff --git a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Transforms/Passes.td b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Passes/Passes.td
similarity index 100%
rename from llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Transforms/Passes.td
rename to llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Passes/Passes.td
diff --git a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Passes/Transforms.h b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Passes/Transforms.h
new file mode 100644
index 0000000..6fa1f51
--- /dev/null
+++ b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Passes/Transforms.h
@@ -0,0 +1,93 @@
+// Copyright 2021 The IREE Authors
+//
+// 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
+
+#ifndef IREE_DIALECTS_DIALECT_LINALGEXT_TRANSFORMS_TRANSFORMS_H_
+#define IREE_DIALECTS_DIALECT_LINALGEXT_TRANSFORMS_TRANSFORMS_H_
+
+#include "iree-dialects/Dialect/LinalgExt/IR/TiledOpInterface.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
+
+namespace mlir {
+namespace iree_compiler {
+namespace IREE {
+namespace LinalgExt {
+
+/// Structure to represent the result of tiling operation.
+struct TiledOp {
+ /// Tiled op.
+ Operation *op;
+ /// Loops generated during tiling.
+ SmallVector<Operation *> loops;
+ /// Values that are replacements for the untiled operations.
+ SmallVector<Value> results;
+};
+
+/// Main entry point for tiling LinalgExtOps using TiledOpInterface.
+FailureOr<TiledOp> tileLinalgExtOp(OpBuilder &b, TiledOpInterface tilableOp,
+ const linalg::LinalgTilingOptions &options);
+
+/// Base rewrite pattern to tile and distribute operations that implement the
+/// `TiledOpInterface`.
+/// Base pattern for tiling TiledOpInterfaceOps.
+struct TiledOpInterfaceBaseTilingPattern
+ : public OpInterfaceRewritePattern<TiledOpInterface> {
+ TiledOpInterfaceBaseTilingPattern(MLIRContext *context,
+ linalg::LinalgTilingOptions options,
+ linalg::LinalgTransformationFilter filter =
+ linalg::LinalgTransformationFilter(),
+ PatternBenefit benefit = 1)
+ : OpInterfaceRewritePattern(context, benefit),
+ filter(filter),
+ options(options) {}
+
+ LogicalResult matchAndRewriteBase(TiledOpInterface tilableOp,
+ PatternRewriter &rewriter,
+ TiledOp &result) const;
+
+ private:
+ /// LinalgTransformMarker handles special attribute manipulations.
+ linalg::LinalgTransformationFilter filter;
+ /// Options to control tiling;
+ linalg::LinalgTilingOptions options;
+};
+
+struct TiledOpInterfaceTilingPattern
+ : public TiledOpInterfaceBaseTilingPattern {
+ TiledOpInterfaceTilingPattern(MLIRContext *context,
+ linalg::LinalgTilingOptions options,
+ linalg::LinalgTransformationFilter filter =
+ linalg::LinalgTransformationFilter(),
+ PatternBenefit benefit = 1)
+ : TiledOpInterfaceBaseTilingPattern(context, options, filter, benefit) {}
+
+ LogicalResult matchAndRewrite(TiledOpInterface tilableOp,
+ PatternRewriter &rewriter) const override {
+ TiledOp tiledOp;
+ // Check for failure.
+ if (failed(TiledOpInterfaceBaseTilingPattern::matchAndRewriteBase(
+ tilableOp, rewriter, tiledOp))) {
+ return failure();
+ }
+ // Check for do-nothing case.
+ if (!tiledOp.op) return failure();
+ if (tiledOp.op != tilableOp) {
+ if (tiledOp.results.empty()) {
+ rewriter.eraseOp(tilableOp);
+ } else {
+ rewriter.replaceOp(tilableOp, tiledOp.results);
+ }
+ }
+ return success();
+ }
+};
+
+} // namespace LinalgExt
+} // namespace IREE
+} // namespace iree_compiler
+} // namespace mlir
+
+#endif // IREE_DIALECTS_DIALECT_LINALGEXT_TRANSFORMS_TRANSFORMS_H_
diff --git a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Transforms/Transforms.h b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Transforms/Transforms.h
index 6fa1f51..3099515 100644
--- a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Transforms/Transforms.h
+++ b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Transforms/Transforms.h
@@ -1,87 +1,93 @@
-// Copyright 2021 The IREE Authors
+//===- Transforms.h - LinalgExt transformations as patterns -----*- C++ -*-===//
//
-// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// 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
+//
+//===----------------------------------------------------------------------===//
#ifndef IREE_DIALECTS_DIALECT_LINALGEXT_TRANSFORMS_TRANSFORMS_H_
#define IREE_DIALECTS_DIALECT_LINALGEXT_TRANSFORMS_TRANSFORMS_H_
-#include "iree-dialects/Dialect/LinalgExt/IR/TiledOpInterface.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
-#include "mlir/Dialect/Linalg/Utils/Utils.h"
+#include "mlir/IR/PatternMatch.h"
namespace mlir {
+namespace scf {
+class ForOp;
+}
+
namespace iree_compiler {
namespace IREE {
namespace LinalgExt {
-/// Structure to represent the result of tiling operation.
-struct TiledOp {
- /// Tiled op.
- Operation *op;
- /// Loops generated during tiling.
- SmallVector<Operation *> loops;
- /// Values that are replacements for the untiled operations.
- SmallVector<Value> results;
-};
+/// Pattern to tile a TilingInterface op using a TileOp.
+struct LinalgExtTilingPattern
+ : public OpInterfaceRewritePattern<TilingInterface> {
+ LinalgExtTilingPattern(MLIRContext *context, linalg::LinalgTilingOptions opt)
+ : OpInterfaceRewritePattern<TilingInterface>(context), options(opt) {}
-/// Main entry point for tiling LinalgExtOps using TiledOpInterface.
-FailureOr<TiledOp> tileLinalgExtOp(OpBuilder &b, TiledOpInterface tilableOp,
- const linalg::LinalgTilingOptions &options);
+ FailureOr<Operation *> returningMatchAndRewrite(
+ TilingInterface op, PatternRewriter &rewriter) const;
-/// Base rewrite pattern to tile and distribute operations that implement the
-/// `TiledOpInterface`.
-/// Base pattern for tiling TiledOpInterfaceOps.
-struct TiledOpInterfaceBaseTilingPattern
- : public OpInterfaceRewritePattern<TiledOpInterface> {
- TiledOpInterfaceBaseTilingPattern(MLIRContext *context,
- linalg::LinalgTilingOptions options,
- linalg::LinalgTransformationFilter filter =
- linalg::LinalgTransformationFilter(),
- PatternBenefit benefit = 1)
- : OpInterfaceRewritePattern(context, benefit),
- filter(filter),
- options(options) {}
-
- LogicalResult matchAndRewriteBase(TiledOpInterface tilableOp,
- PatternRewriter &rewriter,
- TiledOp &result) const;
+ LogicalResult matchAndRewrite(TilingInterface op,
+ PatternRewriter &rewriter) const override {
+ return returningMatchAndRewrite(op, rewriter);
+ }
private:
- /// LinalgTransformMarker handles special attribute manipulations.
- linalg::LinalgTransformationFilter filter;
- /// Options to control tiling;
linalg::LinalgTilingOptions options;
};
-struct TiledOpInterfaceTilingPattern
- : public TiledOpInterfaceBaseTilingPattern {
- TiledOpInterfaceTilingPattern(MLIRContext *context,
- linalg::LinalgTilingOptions options,
- linalg::LinalgTransformationFilter filter =
- linalg::LinalgTransformationFilter(),
- PatternBenefit benefit = 1)
- : TiledOpInterfaceBaseTilingPattern(context, options, filter, benefit) {}
+/// Pattern to rewrite a TileOp to an scf::ForOp.
+struct TileOpToSCFRewriter : public OpRewritePattern<TileOp> {
+ using OpRewritePattern::OpRewritePattern;
- LogicalResult matchAndRewrite(TiledOpInterface tilableOp,
+ FailureOr<scf::ForOp> returningMatchAndRewrite(
+ TileOp tileOp, PatternRewriter &rewriter) const;
+
+ LogicalResult matchAndRewrite(TileOp tileOp,
PatternRewriter &rewriter) const override {
- TiledOp tiledOp;
- // Check for failure.
- if (failed(TiledOpInterfaceBaseTilingPattern::matchAndRewriteBase(
- tilableOp, rewriter, tiledOp))) {
- return failure();
- }
- // Check for do-nothing case.
- if (!tiledOp.op) return failure();
- if (tiledOp.op != tilableOp) {
- if (tiledOp.results.empty()) {
- rewriter.eraseOp(tilableOp);
- } else {
- rewriter.replaceOp(tilableOp, tiledOp.results);
- }
- }
- return success();
+ return returningMatchAndRewrite(tileOp, rewriter);
+ }
+};
+
+/// Pattern to rewrite a TileOp to a InParallelOp.
+struct TileOpToInParallelRewriter : public OpRewritePattern<TileOp> {
+ using OpRewritePattern::OpRewritePattern;
+
+ FailureOr<InParallelOp> returningMatchAndRewrite(
+ TileOp tileOp, PatternRewriter &rewriter) const;
+
+ LogicalResult matchAndRewrite(TileOp tileOp,
+ PatternRewriter &rewriter) const override {
+ return returningMatchAndRewrite(tileOp, rewriter);
+ }
+};
+
+/// Pattern to rewrite a InParallelOp to the async dialect.
+struct InParallelOpToAsyncRewriter : public OpRewritePattern<InParallelOp> {
+ using OpRewritePattern::OpRewritePattern;
+
+ FailureOr<Operation *> returningMatchAndRewrite(
+ InParallelOp inParallelOp, PatternRewriter &rewriter) const;
+
+ LogicalResult matchAndRewrite(InParallelOp inParallelOp,
+ PatternRewriter &rewriter) const override {
+ return returningMatchAndRewrite(inParallelOp, rewriter);
+ }
+};
+
+/// Pattern to rewrite a InParallelOp to an scf::ForOp.
+struct InParallelOpToScfForRewriter : public OpRewritePattern<InParallelOp> {
+ using OpRewritePattern::OpRewritePattern;
+
+ FailureOr<scf::ForOp> returningMatchAndRewrite(
+ InParallelOp inParallelOp, PatternRewriter &rewriter) const;
+
+ LogicalResult matchAndRewrite(InParallelOp inParallelOp,
+ PatternRewriter &rewriter) const override {
+ return returningMatchAndRewrite(inParallelOp, rewriter);
}
};
diff --git a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Transforms/Utils.h b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Transforms/Utils.h
new file mode 100644
index 0000000..534e794
--- /dev/null
+++ b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Transforms/Utils.h
@@ -0,0 +1,124 @@
+//===- Utils.h - Utils for simplifying writing transformations -*- 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
+//
+//===---------------------------------------------------------------------===//
+
+#ifndef IREE_DIALECTS_DIALECT_LINALGEXT_TRANSFORMS_UTILS_H_
+#define IREE_DIALECTS_DIALECT_LINALGEXT_TRANSFORMS_UTILS_H_
+
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/Support/LLVM.h"
+
+namespace mlir {
+class Location;
+class OpBuilder;
+class Operation;
+class Value;
+
+namespace tensor {
+class ExtractSliceOp;
+}
+
+namespace iree_compiler {
+namespace IREE {
+namespace LinalgExt {
+
+/// Helper function which auto-completes the missing trailing dimensions to
+/// always be offset = 0, size = dim, stride = 1.
+void completeOffsetsSizesAndStrides(OpBuilder &b, Location loc, Value tensor,
+ ArrayRef<Value> leadingOffsets,
+ ArrayRef<Value> leadingSizes,
+ ArrayRef<Value> leadingStrides,
+ SmallVectorImpl<Value> &offsets,
+ SmallVectorImpl<Value> &sizes,
+ SmallVectorImpl<Value> &strides);
+
+/// Create a tensor::ExtractSliceOp by auto-completing the missing trailing
+/// dimensions to always be offset = 0, size = dim, stride = 1.
+Value createSubsetExtractOpFromLeadingOffsetsSizesAndStrides(
+ OpBuilder &b, Location loc, Value tensor,
+ llvm::ArrayRef<Value> leadingOffsets, ArrayRef<Value> leadingSizes,
+ ArrayRef<Value> leadingStrides);
+
+/// Create a tensor::InsertSliceOp by auto-completing the missing trailing
+/// dimensions to always be offset = 0, size = dim, stride = 1.
+Value createSubsetInsertOpFromLeadingOffsetsSizesAndStrides(
+ OpBuilder &b, Location loc, Value tensor, Value dest,
+ ArrayRef<Value> leadingOffsets, ArrayRef<Value> leadingSizes,
+ ArrayRef<Value> leadingStrides);
+
+/// Create a linalg_ext::ParallelInsertSliceOp by auto-completing the missing
+/// trailing dimensions to always be offset = 0, size = dim, stride = 1.
+Operation *createParallelInsertSliceOpFromLeadingOffsetsSizesAndStrides(
+ OpBuilder &b, Location loc, Value tensor, Value dest,
+ ArrayRef<Value> leadingOffsets, ArrayRef<Value> leadingSizes,
+ ArrayRef<Value> leadingStrides);
+
+/// Insert the `source` tensor into the `dest` tensor by creating the relevant
+/// `subset_insert` op. The details of the `subset_insert` op are retrieved
+/// from the `subset_extract` op so that they form a matching extract/insert
+/// pair.
+Value createMatchingSubsetInsertOp(OpBuilder &b, Location loc,
+ tensor::ExtractSliceOp subsetExtractOp,
+ Value source, Value dest);
+
+struct AffineValueExpr {
+ explicit AffineValueExpr(AffineExpr e) : e(e) {}
+ AffineValueExpr bind(Value v) {
+ this->v = v;
+ return *this;
+ }
+ operator AffineExpr() const { return e; }
+ operator Value() const { return v; }
+ AffineExpr e;
+ Value v;
+};
+
+/// Helper struct to build simple arithmetic quantiAffineValueExprs with minimal
+/// type inference support.
+// TODO: move into ArithBuilder once ops have been moved into arith.
+struct AffineBuilder {
+ AffineBuilder(OpBuilder &b, Location loc) : b(b), loc(loc) {}
+
+ Value add(AffineValueExpr lhs, AffineValueExpr rhs) {
+ return b.createOrFold<AffineApplyOp>(
+ loc, ArrayRef<AffineExpr>{lhs.e + rhs.e}, ValueRange{lhs, rhs});
+ }
+ Value sub(AffineValueExpr lhs, AffineValueExpr rhs) {
+ return b.createOrFold<AffineApplyOp>(
+ loc, ArrayRef<AffineExpr>{lhs.e - rhs.e}, ValueRange{lhs, rhs});
+ }
+ Value mul(AffineValueExpr lhs, AffineValueExpr rhs) {
+ return b.createOrFold<AffineApplyOp>(
+ loc, ArrayRef<AffineExpr>{lhs.e * rhs.e}, ValueRange{lhs, rhs});
+ }
+ Value ceil(AffineValueExpr lhs, AffineValueExpr rhs) {
+ return b.createOrFold<AffineApplyOp>(
+ loc, ArrayRef<AffineExpr>{lhs.e.ceilDiv(rhs.e)}, ValueRange{lhs, rhs});
+ }
+ Value min(ValueRange vals) {
+ return b.createOrFold<AffineMinOp>(
+ loc, AffineMap::getMultiDimIdentityMap(vals.size(), b.getContext()),
+ vals);
+ }
+ Value max(ValueRange vals) {
+ return b.createOrFold<AffineMinOp>(
+ loc, AffineMap::getMultiDimIdentityMap(vals.size(), b.getContext()),
+ vals);
+ }
+
+ private:
+ OpBuilder &b;
+ Location loc;
+};
+
+} // namespace LinalgExt
+} // namespace IREE
+} // namespace iree_compiler
+} // namespace mlir
+
+#endif // IREE_DIALECTS_DIALECT_LINALGEXT_TRANSFORMS_UTILS_H_
diff --git a/llvm-external-projects/iree-dialects/lib/CAPI/CMakeLists.txt b/llvm-external-projects/iree-dialects/lib/CAPI/CMakeLists.txt
index fde1221..5c0e24d 100644
--- a/llvm-external-projects/iree-dialects/lib/CAPI/CMakeLists.txt
+++ b/llvm-external-projects/iree-dialects/lib/CAPI/CMakeLists.txt
@@ -4,6 +4,7 @@
LINK_LIBS PUBLIC
MLIRIR
IREEInputDialect
+ IREELinalgExtDialect
IREEPyDMDialect
IREEPyDMPasses
)
diff --git a/llvm-external-projects/iree-dialects/lib/CAPI/Dialects.cpp b/llvm-external-projects/iree-dialects/lib/CAPI/Dialects.cpp
index ac169f1..569e530 100644
--- a/llvm-external-projects/iree-dialects/lib/CAPI/Dialects.cpp
+++ b/llvm-external-projects/iree-dialects/lib/CAPI/Dialects.cpp
@@ -7,6 +7,7 @@
#include "iree-dialects-c/Dialects.h"
#include "iree-dialects/Dialect/Input/InputDialect.h"
+#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtDialect.h"
#include "iree-dialects/Dialect/PyDM/IR/PyDMDialect.h"
#include "iree-dialects/Dialect/PyDM/Transforms/Passes.h"
#include "mlir/CAPI/IR.h"
@@ -27,6 +28,14 @@
MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(
IREEInput, iree_input, mlir::iree_compiler::IREE::Input::IREEInputDialect)
+//===--------------------------------------------------------------------===//
+// IREELinalgExt
+//===--------------------------------------------------------------------===//
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(
+ IREELinalgExt, iree_linalg_ext,
+ mlir::iree_compiler::IREE::LinalgExt::IREELinalgExtDialect)
+
//===----------------------------------------------------------------------===//
// IREEPyDMDialect
//===----------------------------------------------------------------------===//
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/CMakeLists.txt b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/CMakeLists.txt
index 9f57627..126b878 100644
--- a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/CMakeLists.txt
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/CMakeLists.txt
@@ -1,2 +1,3 @@
add_subdirectory(IR)
+add_subdirectory(Passes)
add_subdirectory(Transforms)
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/IR/LinalgExtOps.cpp b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/IR/LinalgExtOps.cpp
index b692d30..cabc5c4 100644
--- a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/IR/LinalgExtOps.cpp
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/IR/LinalgExtOps.cpp
@@ -24,6 +24,7 @@
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/FunctionImplementation.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/IR/OperationSupport.h"
@@ -1242,6 +1243,388 @@
} // namespace
//===----------------------------------------------------------------------===//
+// TileOp
+//===----------------------------------------------------------------------===//
+
+void TileOp::build(mlir::OpBuilder &builder, mlir::OperationState &result,
+ Value tileSize, ValueRange outs, int64_t tiledDim,
+ TileOp::TileOpBodyBuilderFn bodyBuilder) {
+ result.addOperands(tileSize);
+ result.addOperands(outs);
+ result.addAttribute(TileOp::getTiledDimAttrName(),
+ builder.getI64IntegerAttr(tiledDim));
+ result.addTypes(outs.getType());
+
+ Region *bodyRegion = result.addRegion();
+ bodyRegion->push_back(new Block);
+ Block &bodyBlock = bodyRegion->front();
+ // TODO: Pass a better location here.
+ Location loc = tileSize.getLoc();
+ bodyBlock.addArgument(builder.getIndexType(), loc);
+ bodyBlock.addArgument(builder.getIndexType(), loc);
+ // Handle the sliced out types in a conservative fashion: all dimensions
+ // become dynamic and a later canonicalization is expected to recover static
+ // types.
+ // TODO: should we relax this and use something less strict?
+ auto dynamicTypes =
+ llvm::to_vector(llvm::map_range(outs.getTypes(), [](Type t) -> Type {
+ auto rankedTensorType = t.cast<RankedTensorType>();
+ RankedTensorType::Builder rttb(rankedTensorType);
+ SmallVector<int64_t> dynamicShape(rankedTensorType.getRank(),
+ ShapedType::kDynamicSize);
+ return rttb.setShape(dynamicShape);
+ }));
+ SmallVector<Location> locs(dynamicTypes.size(), loc);
+ bodyBlock.addArguments(dynamicTypes, locs);
+
+ OpBuilder::InsertionGuard guard(builder);
+ builder.setInsertionPointToStart(&bodyBlock);
+ bodyBuilder(builder, result.location, bodyBlock.getArgument(0),
+ bodyBlock.getArgument(1), bodyBlock.getArguments().drop_front(2));
+}
+
+void TileOp::build(mlir::OpBuilder &builder, mlir::OperationState &result,
+ Value tileSize, ValueRange outs,
+ TileOp::TileOpBodyBuilderFn bodyBuilder) {
+ TileOp::build(builder, result, tileSize, outs, 0, bodyBuilder);
+}
+
+// TODO(#81): Impl me.
+LogicalResult TileOp::verify() { return success(); }
+
+void TileOp::print(OpAsmPrinter &p) {
+ p << ' ' << tile_size() << ' ';
+ if (tiled_dim() > 0) p << "tiled_dim = " << tiled_dim() << ' ';
+ if (!outs().empty()) {
+ p << "outs(";
+ llvm::interleaveComma(outs(), p,
+ [&p](Value v) { p << v << ": " << v.getType(); });
+ p << ')';
+ }
+ p << " -> (" << getResultTypes() << ") ";
+ p.printRegion(region(),
+ /*printEntryBlockArgs=*/true,
+ /*printBlockTerminators=*/true);
+ p.printOptionalAttrDict(getOperation()->getAttrs(),
+ /*elidedAttrs=*/{TileOp::getTiledDimAttrName()});
+}
+
+ParseResult TileOp::parse(OpAsmParser &parser, OperationState &result) {
+ auto &builder = parser.getBuilder();
+
+ OpAsmParser::OperandType tileSizes;
+ // TODO: also allow tensor<..xindex> and figure out a good syntax.
+ // Type tensorOfIndexType =
+ // RankedTensorType::get({ShapedType::kDynamicSize}, indexType);
+ Type tileSizesType = builder.getIndexType();
+ SmallVector<Type> outsTypes;
+ SmallVector<OpAsmParser::OperandType, 4> outsOperands;
+
+ llvm::SMLoc outputsOperandsLoc;
+ if (parser.parseOperand(tileSizes) ||
+ parser.resolveOperand(tileSizes, tileSizesType, result.operands))
+ return failure();
+
+ // Parse the `tiled_dim` attribute or set it to 0 implicitly when elided.
+ if (succeeded(parser.parseOptionalKeyword(TileOp::getTiledDimAttrName()))) {
+ outputsOperandsLoc = parser.getCurrentLocation();
+ Attribute valueAttr;
+ parser.parseAttribute(valueAttr, TileOp::getTiledDimAttrName(),
+ result.attributes);
+ } else {
+ result.attributes.append(TileOp::getTiledDimAttrName(),
+ parser.getBuilder().getI64IntegerAttr(0));
+ }
+
+ if (succeeded(parser.parseOptionalKeyword("outs"))) {
+ bool _1;
+ SmallVector<NamedAttrList> _2;
+ SmallVector<Location> _3;
+ outputsOperandsLoc = parser.getCurrentLocation();
+ if (mlir::function_interface_impl::parseFunctionArgumentList(
+ parser,
+ /*allowAttributes=*/false,
+ /*allowVariadic=*/false, outsOperands, outsTypes, /*argAttrs=*/_2,
+ /*argLocations=*/_3,
+ /*isVariadic=*/_1) ||
+ parser.resolveOperands(outsOperands, outsTypes, outputsOperandsLoc,
+ result.operands))
+ return failure();
+ }
+ if (parser.parseArrowTypeList(result.types)) return failure();
+
+ SmallVector<OpAsmParser::OperandType, 8> regionOperands;
+ std::unique_ptr<Region> region = std::make_unique<Region>();
+ SmallVector<Type, 8> operandTypes, regionTypes;
+ if (parser.parseRegion(*region, regionOperands, regionTypes))
+ return failure();
+
+ // Parse the optional attribute list.
+ if (parser.parseOptionalAttrDict(result.attributes)) return failure();
+
+ TileOp::ensureTerminator(*region, builder, result.location);
+ result.addRegion(std::move(region));
+
+ return success();
+}
+
+//===----------------------------------------------------------------------===//
+// InParallelOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult InParallelOp::verify() {
+ // Check that the body defines as single block argument for the thread index.
+ auto *body = getBody();
+ if (body->getNumArguments() != 1)
+ return emitOpError("body expects exactly one argument");
+ if (!body->getArgument(0).getType().isIndex())
+ return emitOpError(
+ "expected body first argument to be an index argument for "
+ "the thread index");
+
+ // Verify consistency between the result types and the terminator.
+ auto terminatorTypes = getTerminator().yieldedTypes();
+ auto opResults = getResults();
+ if (opResults.size() != terminatorTypes.size())
+ return emitOpError("produces ")
+ << opResults.size() << " results, but its terminator yields "
+ << terminatorTypes.size() << " values";
+ unsigned i = 0;
+ for (auto e : llvm::zip(terminatorTypes, opResults)) {
+ if (std::get<0>(e) != std::get<1>(e).getType())
+ return emitOpError() << "type mismatch between " << i
+ << "th result of in_parallel (" << std::get<0>(e)
+ << ") and " << i << "th result yielded by its "
+ << "terminator (" << std::get<1>(e).getType() << ")";
+ i++;
+ }
+
+ return success();
+}
+
+void InParallelOp::print(OpAsmPrinter &p) {
+ p << ' ' << num_threads() << ' ';
+ p << " -> (" << getResultTypes() << ") ";
+ p.printRegion(region(),
+ /*printEntryBlockArgs=*/true,
+ /*printBlockTerminators=*/true);
+ p.printOptionalAttrDict(getOperation()->getAttrs());
+}
+
+ParseResult InParallelOp::parse(OpAsmParser &parser, OperationState &result) {
+ auto &builder = parser.getBuilder();
+
+ OpAsmParser::OperandType numThreads;
+ Type indexType = builder.getIndexType();
+
+ if (parser.parseOperand(numThreads) ||
+ parser.resolveOperand(numThreads, indexType, result.operands))
+ return failure();
+ if (parser.parseArrowTypeList(result.types)) return failure();
+
+ SmallVector<OpAsmParser::OperandType, 8> regionOperands;
+ SmallVector<Type, 8> regionTypes;
+ std::unique_ptr<Region> region = std::make_unique<Region>();
+ if (parser.parseRegion(*region, regionOperands, regionTypes))
+ return failure();
+ InParallelOp::ensureTerminator(*region, builder, result.location);
+ result.addRegion(std::move(region));
+
+ // Parse the optional attribute list.
+ if (parser.parseOptionalAttrDict(result.attributes)) return failure();
+ return success();
+}
+
+// Bodyless builder, result types must be specified.
+void InParallelOp::build(mlir::OpBuilder &builder, mlir::OperationState &result,
+ TypeRange resultTypes, Value numThreads) {
+ // TODO: Pass better location.
+ Location loc = numThreads.getLoc();
+ result.addOperands(numThreads);
+
+ Region *bodyRegion = result.addRegion();
+ bodyRegion->push_back(new Block);
+ Block &bodyBlock = bodyRegion->front();
+ bodyBlock.addArgument(builder.getIndexType(), loc);
+
+ // Create the default terminator if the builder is not provided and if the
+ // iteration arguments are not provided. Otherwise, leave this to the caller
+ // because we don't know which values to return from the loop.
+ InParallelOp::ensureTerminator(*bodyRegion, builder, result.location);
+ result.addTypes(resultTypes);
+}
+
+// Builder that takes a bodyBuilder lambda, result types are inferred from
+// the terminator.
+void InParallelOp::build(
+ mlir::OpBuilder &builder, mlir::OperationState &result, Value numThreads,
+ function_ref<void(OpBuilder &, Location, Value)> bodyBuilder) {
+ // TODO: Pass better location.
+ Location loc = numThreads.getLoc();
+ result.addOperands(numThreads);
+
+ Region *bodyRegion = result.addRegion();
+ bodyRegion->push_back(new Block);
+ Block &bodyBlock = bodyRegion->front();
+ bodyBlock.addArgument(builder.getIndexType(), loc);
+
+ OpBuilder::InsertionGuard guard(builder);
+ builder.setInsertionPointToStart(&bodyBlock);
+ bodyBuilder(builder, result.location, bodyBlock.getArgument(0));
+ auto terminator =
+ llvm::cast<PerformConcurrentlyOp>(bodyBlock.getTerminator());
+ result.addTypes(terminator.yieldedTypes());
+}
+
+// The ensureTerminator method generated by SingleBlockImplicitTerminator is
+// unaware of the fact that our terminator also needs a region to be well
+// formed. We override it here to ensure that we do the right thing.
+void InParallelOp::ensureTerminator(Region ®ion, Builder &builder,
+ Location loc) {
+ OpTrait::SingleBlockImplicitTerminator<PerformConcurrentlyOp>::Impl<
+ InParallelOp>::ensureTerminator(region, builder, loc);
+ auto terminator =
+ llvm::dyn_cast<PerformConcurrentlyOp>(region.front().getTerminator());
+ PerformConcurrentlyOp::ensureTerminator(terminator.getRegion(), builder, loc);
+}
+
+PerformConcurrentlyOp InParallelOp::getTerminator() {
+ return cast<PerformConcurrentlyOp>(getBody()->getTerminator());
+}
+
+//===----------------------------------------------------------------------===//
+// ParallelInsertSliceOp
+//===----------------------------------------------------------------------===//
+
+// Build a ParallelInsertSliceOp with mixed static and dynamic entries.
+void ParallelInsertSliceOp::build(OpBuilder &b, OperationState &result,
+ Value source, Value dest,
+ ArrayRef<OpFoldResult> offsets,
+ ArrayRef<OpFoldResult> sizes,
+ ArrayRef<OpFoldResult> strides,
+ ArrayRef<NamedAttribute> attrs) {
+ SmallVector<int64_t> staticOffsets, staticSizes, staticStrides;
+ SmallVector<Value> dynamicOffsets, dynamicSizes, dynamicStrides;
+ dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets,
+ ShapedType::kDynamicStrideOrOffset);
+ dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes,
+ ShapedType::kDynamicSize);
+ dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides,
+ ShapedType::kDynamicStrideOrOffset);
+ build(b, result, {}, source, dest, dynamicOffsets, dynamicSizes,
+ dynamicStrides, b.getI64ArrayAttr(staticOffsets),
+ b.getI64ArrayAttr(staticSizes), b.getI64ArrayAttr(staticStrides));
+ result.addAttributes(attrs);
+}
+
+// Build a ParallelInsertSliceOp with dynamic entries.
+void ParallelInsertSliceOp::build(OpBuilder &b, OperationState &result,
+ Value source, Value dest, ValueRange offsets,
+ ValueRange sizes, ValueRange strides,
+ ArrayRef<NamedAttribute> attrs) {
+ SmallVector<OpFoldResult> offsetValues = llvm::to_vector<4>(
+ llvm::map_range(offsets, [](Value v) -> OpFoldResult { return v; }));
+ SmallVector<OpFoldResult> sizeValues = llvm::to_vector<4>(
+ llvm::map_range(sizes, [](Value v) -> OpFoldResult { return v; }));
+ SmallVector<OpFoldResult> strideValues = llvm::to_vector<4>(
+ llvm::map_range(strides, [](Value v) -> OpFoldResult { return v; }));
+ build(b, result, source, dest, offsetValues, sizeValues, strideValues);
+}
+
+namespace {
+/// Pattern to rewrite a parallel_insert_slice op with constant arguments.
+class ParallelInsertSliceOpConstantArgumentFolder final
+ : public OpRewritePattern<ParallelInsertSliceOp> {
+ public:
+ using OpRewritePattern<ParallelInsertSliceOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(ParallelInsertSliceOp insertSliceOp,
+ PatternRewriter &rewriter) const override {
+ // No constant operand, just return.
+ if (llvm::none_of(insertSliceOp.getOperands(), [](Value operand) {
+ return matchPattern(operand, matchConstantIndex());
+ }))
+ return failure();
+
+ // At least one of offsets/sizes/strides is a new constant.
+ // Form the new list of operands and constant attributes from the
+ // existing.
+ SmallVector<OpFoldResult> mixedOffsets(insertSliceOp.getMixedOffsets());
+ SmallVector<OpFoldResult> mixedSizes(insertSliceOp.getMixedSizes());
+ SmallVector<OpFoldResult> mixedStrides(insertSliceOp.getMixedStrides());
+ canonicalizeSubViewPart(mixedOffsets, ShapedType::isDynamicStrideOrOffset);
+ canonicalizeSubViewPart(mixedSizes, ShapedType::isDynamic);
+ canonicalizeSubViewPart(mixedStrides, ShapedType::isDynamicStrideOrOffset);
+
+ // Create the new op in canonical form.
+ rewriter.replaceOpWithNewOp<ParallelInsertSliceOp>(
+ insertSliceOp, insertSliceOp.source(), insertSliceOp.dest(),
+ mixedOffsets, mixedSizes, mixedStrides);
+ return success();
+ }
+};
+} // namespace
+
+void ParallelInsertSliceOp::getCanonicalizationPatterns(
+ RewritePatternSet &results, MLIRContext *context) {
+ results.add<ParallelInsertSliceOpConstantArgumentFolder>(context);
+}
+
+//===----------------------------------------------------------------------===//
+// PerformConcurrentlyOp
+//===----------------------------------------------------------------------===//
+
+// TODO(ntv,apaszke): Implement this
+LogicalResult PerformConcurrentlyOp::verify() { return success(); }
+
+void PerformConcurrentlyOp::print(OpAsmPrinter &p) {
+ p << " ";
+ p.printRegion(region(),
+ /*printEntryBlockArgs=*/false,
+ /*printBlockTerminators=*/false);
+ p.printOptionalAttrDict(getOperation()->getAttrs());
+}
+
+ParseResult PerformConcurrentlyOp::parse(OpAsmParser &parser,
+ OperationState &result) {
+ auto &builder = parser.getBuilder();
+
+ SmallVector<OpAsmParser::OperandType, 8> regionOperands;
+ SmallVector<Type, 8> regionTypes;
+ std::unique_ptr<Region> region = std::make_unique<Region>();
+ if (parser.parseRegion(*region, regionOperands, regionTypes))
+ return failure();
+ PerformConcurrentlyOp::ensureTerminator(*region, builder, result.location);
+ result.addRegion(std::move(region));
+
+ // Parse the optional attribute list.
+ if (parser.parseOptionalAttrDict(result.attributes)) return failure();
+ return success();
+}
+
+SmallVector<Type> PerformConcurrentlyOp::yieldedTypes() {
+ return llvm::to_vector(llvm::map_range(
+ this->yieldingOps(),
+ [](ParallelInsertSliceOp op) { return op.yieldedType(); }));
+}
+
+SmallVector<ParallelInsertSliceOp> PerformConcurrentlyOp::yieldingOps() {
+ SmallVector<ParallelInsertSliceOp> ret;
+ for (Operation &op : *getBody()) {
+ // TODO: interface when this grows up.
+ if (auto sliceOp = llvm::dyn_cast<ParallelInsertSliceOp>(op)) {
+ ret.push_back(sliceOp);
+ continue;
+ }
+ if (auto endPerformOp = llvm::dyn_cast<EndPerformConcurrentlyOp>(op)) {
+ continue;
+ }
+ llvm_unreachable("Unexpected operation in perform_concurrently");
+ }
+ return ret;
+}
+
+//===----------------------------------------------------------------------===//
// LinalgExtDialect
//===----------------------------------------------------------------------===//
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Passes/CMakeLists.txt b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Passes/CMakeLists.txt
new file mode 100644
index 0000000..e26003e
--- /dev/null
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Passes/CMakeLists.txt
@@ -0,0 +1,25 @@
+add_mlir_library(IREELinalgExtPasses
+ ConvertToLoops.cpp
+ PadContractionToBlockSize.cpp
+ Passes.cpp
+ Tiling.cpp
+
+ DEPENDS
+ IREELinalgExtPassesIncGen
+
+ LINK_LIBS PUBLIC
+ IREEInputDialect
+ IREELinalgExtDialect
+ MLIRAffine
+ MLIRIR
+ MLIRLinalg
+ MLIRLinalgTransforms
+ MLIRMath
+ MLIRMemRef
+ MLIRPass
+ MLIRSCF
+ MLIRFunc
+ MLIRSupport
+ MLIRTensor
+ MLIRTransforms
+)
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/ConvertToLoops.cpp b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Passes/ConvertToLoops.cpp
similarity index 96%
rename from llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/ConvertToLoops.cpp
rename to llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Passes/ConvertToLoops.cpp
index 52c9dcf..da62126 100644
--- a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/ConvertToLoops.cpp
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Passes/ConvertToLoops.cpp
@@ -6,8 +6,8 @@
#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtDialect.h"
#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.h"
-#include "iree-dialects/Dialect/LinalgExt/Transforms/PassDetail.h"
-#include "iree-dialects/Dialect/LinalgExt/Transforms/Passes.h"
+#include "iree-dialects/Dialect/LinalgExt/Passes/PassDetail.h"
+#include "iree-dialects/Dialect/LinalgExt/Passes/Passes.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/PadContractionToBlockSize.cpp b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Passes/PadContractionToBlockSize.cpp
similarity index 97%
rename from llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/PadContractionToBlockSize.cpp
rename to llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Passes/PadContractionToBlockSize.cpp
index b050cc7..a2fe9bd 100644
--- a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/PadContractionToBlockSize.cpp
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Passes/PadContractionToBlockSize.cpp
@@ -6,8 +6,8 @@
#include "iree-dialects/Dialect/Input/InputDialect.h"
#include "iree-dialects/Dialect/Input/InputOps.h"
-#include "iree-dialects/Dialect/LinalgExt/Transforms/PassDetail.h"
-#include "iree-dialects/Dialect/LinalgExt/Transforms/Passes.h"
+#include "iree-dialects/Dialect/LinalgExt/Passes/PassDetail.h"
+#include "iree-dialects/Dialect/LinalgExt/Passes/Passes.h"
#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/Passes.cpp b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Passes/Passes.cpp
similarity index 82%
rename from llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/Passes.cpp
rename to llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Passes/Passes.cpp
index c41b9ed..f038541 100644
--- a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/Passes.cpp
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Passes/Passes.cpp
@@ -4,7 +4,7 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-#include "iree-dialects/Dialect/LinalgExt/Transforms/Passes.h"
+#include "iree-dialects/Dialect/LinalgExt/Passes/Passes.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassRegistry.h"
@@ -20,7 +20,7 @@
namespace detail {
#define GEN_PASS_REGISTRATION
-#include "iree-dialects/Dialect/LinalgExt/Transforms/Passes.h.inc" // IWYU pragma: export
+#include "iree-dialects/Dialect/LinalgExt/Passes/Passes.h.inc" // IWYU pragma: export
} // namespace detail
} // namespace LinalgExt
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Passes/Tiling.cpp b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Passes/Tiling.cpp
new file mode 100644
index 0000000..fd66bff
--- /dev/null
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Passes/Tiling.cpp
@@ -0,0 +1,360 @@
+// Copyright 2021 The IREE Authors
+//
+// 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
+
+#include "iree-dialects/Dialect/Input/InputDialect.h"
+#include "iree-dialects/Dialect/Input/InputOps.h"
+#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtDialect.h"
+#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.h"
+#include "iree-dialects/Dialect/LinalgExt/Passes/PassDetail.h"
+#include "iree-dialects/Dialect/LinalgExt/Passes/Passes.h"
+#include "iree-dialects/Dialect/LinalgExt/Passes/Transforms.h"
+#include "llvm/ADT/TypeSwitch.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+namespace IREE = mlir::iree_compiler::IREE;
+using namespace IREE::LinalgExt;
+
+//===----------------------------------------------------------------------===//
+// Utility methods for tiling a linalg_ext operation that implements a
+// TiledOpInterface
+//===----------------------------------------------------------------------===//
+
+/// Returns failure if the options are unsupported.
+static LogicalResult verifySupportedTilingOptions(
+ PatternRewriter &rewriter, Operation *op,
+ const linalg::LinalgTilingOptions &options) {
+ if (!options.interchangeVector.empty()) {
+ return rewriter.notifyMatchFailure(op,
+ "unsupported interchange during tiling");
+ }
+ if (options.loopType != linalg::LinalgTilingLoopType::Loops) {
+ return rewriter.notifyMatchFailure(op,
+ "only tiling with scf.for is supported");
+ }
+ if (options.distribution) {
+ if (llvm::any_of(options.distribution->distributionMethod,
+ [](linalg::DistributionMethod method) {
+ return method != linalg::DistributionMethod::Cyclic;
+ })) {
+ return rewriter.notifyMatchFailure(op,
+ "only cyclic distibution is allowed");
+ }
+ }
+ return success();
+}
+
+/// Converts an `OpFoldResult` to a `Value` by building a constant op if
+/// if the `OpFoldResult` is an `IntegerAttr`.
+static Value getValue(OpBuilder &builder, Location loc,
+ OpFoldResult valueOrAttr) {
+ if (auto attr = valueOrAttr.dyn_cast<Attribute>()) {
+ return builder.create<arith::ConstantIndexOp>(
+ loc, attr.cast<IntegerAttr>().getInt());
+ }
+ return valueOrAttr.get<Value>();
+}
+
+/// Returns true if loop is untiled. Only checks if the value is statically
+/// zero. It is assumed that a `Value` defined by a constant op is already
+/// converted to an `IntegerAttr` of that value. So here just return true if
+/// this is an attribute with a zero value.
+static bool isUntiledLoop(OpFoldResult valueOrAttr) {
+ Optional<int64_t> intVal = getConstantIntValue(valueOrAttr);
+ return intVal && *intVal == 0;
+}
+
+/// Generates the tiled loops and the body by invoking the interface methods of
+/// TiledOpInterface.
+/// - `outputs` are the operands to use for outputs of the tiled operation.
+/// - `tileSizes` are tile sizes specified for all loops of the operation. If a
+/// loop is to be untiled it is set to 0.
+/// - `iteratorType` is the type of the loop iterator returned by the
+/// TiledOpInterface.
+/// - `loopBounds` are the bounds of all the loops of the op returned by the
+/// TiledOpInterface.
+/// - `loopDepth` is the current loop depth being processed.
+/// - `offsets` are the `Value`s that represent the position of the tile being
+/// operated on. The offsets are computed as the tiled loops are being
+/// generated.
+/// - `distributionInfo` is the proc_id and nprocs `Value`s to be used for
+/// distributed loops. It is a stack, and once an entry at the top of the
+/// stack is used for distribution it is popped before processing the inner
+/// loops.
+static FailureOr<TiledOp> tileInterfaceOpImpl(
+ OpBuilder &builder, TiledOpInterface tilableOp, ValueRange outputs,
+ MutableArrayRef<OpFoldResult> tileSizes, ArrayRef<StringRef> iteratorTypes,
+ ArrayRef<Range> loopBounds, unsigned loopDepth,
+ SmallVectorImpl<OpFoldResult> &offsets,
+ ArrayRef<linalg::ProcInfo> distributionInfo) {
+ Location loc = tilableOp.getLoc();
+ // If this is the innermost loop, then generated the tiled implementation of
+ // the op by invoking the TiledOpInterface methods.
+ if (loopDepth == tileSizes.size()) {
+ TiledOp ret;
+ ret.op = tilableOp.getTiledImplementation(builder, outputs, offsets,
+ tileSizes, ret.results);
+ if (!ret.op) {
+ return static_cast<LogicalResult>(
+ tilableOp.emitOpError("failed to get tiled implementation"));
+ }
+ return ret;
+ }
+
+ // If tile size at this depth is empty, do nothing.
+ if (isUntiledLoop(tileSizes[loopDepth])) {
+ auto zeroAttr = builder.getI64IntegerAttr(0);
+ offsets.push_back(zeroAttr);
+ assert(matchPattern(loopBounds[loopDepth].offset, m_Zero()) &&
+ "expected loop bounds to have lower bound of zero");
+ tileSizes[loopDepth] = getAsOpFoldResult(loopBounds[loopDepth].size);
+ return tileInterfaceOpImpl(builder, tilableOp, outputs, tileSizes,
+ iteratorTypes, loopBounds, loopDepth + 1,
+ offsets, distributionInfo);
+ }
+
+ // Generate an scf.for for the current loop depth.
+ Value lb = loopBounds[loopDepth].offset;
+ Value ub = loopBounds[loopDepth].size;
+ // TODO(#7073): Put the check back. This is required by tiling linalg_ext.fft
+ // op. We can put the check back after updating linalg_ext.fft semantics.
+ // if (!matchPattern(loopBounds[loopDepth].stride, m_One())) {
+ // return static_cast<LogicalResult>(
+ // tilableOp.emitOpError("expected stride to be 1"));
+ //}
+ Value step = getValue(builder, loc, tileSizes[loopDepth]);
+
+ // Update lb, ub and step for cyclic distribution.
+ if (!distributionInfo.empty() &&
+ iteratorTypes[loopDepth] == getParallelIteratorTypeName()) {
+ linalg::updateBoundsForCyclicDistribution(
+ builder, loc, distributionInfo.front().procId,
+ distributionInfo.front().nprocs, lb, ub, step);
+ distributionInfo = distributionInfo.drop_front();
+ }
+ FailureOr<TiledOp> innerReturnValue;
+ bool isBufferTiling = tilableOp->getNumResults() == 0;
+ ValueRange initValues(isBufferTiling ? ValueRange{} : outputs);
+ auto forOp = builder.create<scf::ForOp>(
+ loc, lb, ub, step, initValues,
+ [&](OpBuilder &b, Location loc, Value iv, ValueRange args) {
+ offsets.push_back(iv);
+ auto affineMaps = AffineMap::inferFromExprList({ArrayRef<AffineExpr>{
+ b.getAffineSymbolExpr(0),
+ b.getAffineSymbolExpr(1) - b.getAffineDimExpr(0)}})[0];
+ // Similar to linalg tiling, the tile size is the min(tileSizes, ub -
+ // iv) to account for cases where tile size does not divide (ub - lb)
+ // exactly.
+ Value inBoundsTileSize = b.create<AffineMinOp>(
+ loc, affineMaps,
+ ValueRange{iv, getValue(builder, loc, tileSizes[loopDepth]), ub});
+ tileSizes[loopDepth] = getAsOpFoldResult(inBoundsTileSize);
+ // Recursively proceed to generate the tiled loop for the next level.
+ innerReturnValue =
+ tileInterfaceOpImpl(b, tilableOp, (isBufferTiling ? outputs : args),
+ tileSizes, iteratorTypes, loopBounds,
+ loopDepth + 1, offsets, distributionInfo);
+ if (failed(innerReturnValue)) return;
+ b.create<scf::YieldOp>(loc, innerReturnValue->results);
+ });
+ if (failed(innerReturnValue)) {
+ return innerReturnValue;
+ }
+ innerReturnValue->loops.insert(innerReturnValue->loops.begin(),
+ forOp.getOperation());
+ innerReturnValue->results = forOp.getResults();
+ return innerReturnValue;
+}
+
+FailureOr<TiledOp> tileInterfaceOp(OpBuilder &b, TiledOpInterface tilableOp,
+ const linalg::LinalgTilingOptions &options) {
+ SmallVector<Value> dest = tilableOp.getDestinationOperands(b);
+ if (dest.empty()) {
+ return static_cast<LogicalResult>(tilableOp.emitOpError(
+ "cannot tile operation without destination operands"));
+ }
+
+ SmallVector<StringRef> iteratorTypes = tilableOp.getLoopIteratorTypes();
+ SmallVector<Value, 4> tileSizesVals =
+ options.tileSizeComputationFunction(b, tilableOp);
+ auto zeroAttr = b.getI64IntegerAttr(0);
+
+ // The actual tile sizes used converts `Value` defined as constant 0, to a
+ // zero integer attributes. Currently if the iterator type is not "parallel",
+ // the tile size is forced to zero as well.
+ auto tileSizes = getAsOpFoldResult(tileSizesVals);
+ tileSizes.resize(iteratorTypes.size(), zeroAttr);
+ for (auto en : llvm::enumerate(iteratorTypes)) {
+ if (en.value() == getParallelIteratorTypeName()) continue;
+ if (!isUntiledLoop(tileSizes[en.index()])) {
+ return static_cast<LogicalResult>(tilableOp.emitOpError(
+ "unimplemented tiling of non-parallel loop iterator type"));
+ }
+ }
+
+ // Trivial early exit case of tile sizes being zero for all parallel loops.
+ if (llvm::all_of(tileSizes, isUntiledLoop)) {
+ return TiledOp{tilableOp, {}, {}};
+ }
+
+ SmallVector<Range> loopBounds = tilableOp.getIterationDomain(b);
+ SmallVector<linalg::ProcInfo> distributionInfo;
+ // If the tiled loops are distributed, get the proc_id and nprocs for the
+ // distributed loops. First collect the parallel loops by iterating over the
+ // tileSizes and getting the loops that are distribute, i.e.,
+ // - parallel, i.e. iteratorTypes is "parallel"
+ // - tiled, i.e. tileSize != 0
+ if (options.distribution) {
+ SmallVector<Range> distributedLoopRange;
+ for (auto i : llvm::seq<unsigned>(0, tileSizes.size())) {
+ if (isUntiledLoop(tileSizes[i])) continue;
+ if (iteratorTypes[i] != getParallelIteratorTypeName()) continue;
+ distributedLoopRange.push_back(loopBounds[i]);
+ }
+ distributionInfo = options.distribution->procInfo(b, tilableOp.getLoc(),
+ distributedLoopRange);
+ }
+
+ SmallVector<OpFoldResult> offsets;
+ return tileInterfaceOpImpl(b, tilableOp, dest, tileSizes, iteratorTypes,
+ loopBounds, 0, offsets, distributionInfo);
+}
+
+LogicalResult TiledOpInterfaceBaseTilingPattern::matchAndRewriteBase(
+ TiledOpInterface tilableOp, PatternRewriter &rewriter,
+ TiledOp &result) const {
+ if (failed(filter.checkAndNotify(rewriter, tilableOp))) {
+ return failure();
+ }
+ if (failed(verifySupportedTilingOptions(rewriter, tilableOp, options))) {
+ return failure();
+ }
+
+ FailureOr<TiledOp> res = tileInterfaceOp(rewriter, tilableOp, options);
+ if (failed(res)) return res;
+ result = *res;
+ if (result.op) {
+ filter.replaceLinalgTransformationFilter(rewriter, result.op);
+ }
+ return success();
+}
+
+//===----------------------------------------------------------------------===//
+// Test pass for tiling Linalg Ext ops
+//===----------------------------------------------------------------------===//
+
+namespace {
+struct TiledOpInterfaceTilingPass
+ : public TiledOpInterfaceTilingBase<TiledOpInterfaceTilingPass> {
+ void getDependentDialects(DialectRegistry ®istry) const override {
+ registry.insert<
+ AffineDialect, IREE::Input::IREEInputDialect, linalg::LinalgDialect,
+ IREE::LinalgExt::IREELinalgExtDialect, memref::MemRefDialect,
+ func::FuncDialect, mlir::arith::ArithmeticDialect, math::MathDialect,
+ tensor::TensorDialect, scf::SCFDialect>();
+ }
+ void runOnOperation() override;
+};
+} // namespace
+
+template <typename OpTy>
+static Value buildFlowWorkgroupInfoOp(OpBuilder &b, unsigned dim) {
+ return b.template create<OpTy>(b.getInsertionPoint()->getLoc(), dim);
+}
+
+void TiledOpInterfaceTilingPass::runOnOperation() {
+ FuncOp funcOp = getOperation();
+ MLIRContext *context = funcOp.getContext();
+
+ RewritePatternSet patterns(context);
+ patterns.add<TiledOpInterfaceTilingPattern>(
+ context, linalg::LinalgTilingOptions().setTileSizes({10, 20}),
+ linalg::LinalgTransformationFilter(
+ StringAttr::get(context, "tiling_input"),
+ StringAttr::get(context, "tiling_output")));
+ patterns.add<TiledOpInterfaceTilingPattern>(
+ context, linalg::LinalgTilingOptions().setTileSizes(ArrayRef<int64_t>{0}),
+ linalg::LinalgTransformationFilter(
+ StringAttr::get(context, "no_tiling_input"),
+ StringAttr::get(context, "no_tiling_output")));
+
+ patterns.add<TiledOpInterfaceTilingPattern>(
+ context, linalg::LinalgTilingOptions().setTileSizes({0, 20}),
+ linalg::LinalgTransformationFilter(
+ StringAttr::get(context, "outer_reduce_input"),
+ StringAttr::get(context, "outer_reduce_output")));
+ patterns.add<TiledOpInterfaceTilingPattern>(
+ context, linalg::LinalgTilingOptions().setTileSizes({10, 0, 0}),
+ linalg::LinalgTransformationFilter(
+ StringAttr::get(context, "inner_reduce_input"),
+ StringAttr::get(context, "inner_reduce_output")));
+
+ static linalg::LinalgLoopDistributionOptions workgroupDistributionOptions = {
+ [](OpBuilder &builder, Location loc, ArrayRef<Range> parallelLoopRanges) {
+ auto numParallelDims = parallelLoopRanges.size();
+
+ SmallVector<linalg::ProcInfo, 3> procInfo(numParallelDims);
+ for (size_t dim = 0; dim < numParallelDims; ++dim) {
+ procInfo[numParallelDims - dim - 1] = {
+ buildFlowWorkgroupInfoOp<IREE::Input::DispatchWorkgroupIDOp>(
+ builder, dim),
+ buildFlowWorkgroupInfoOp<IREE::Input::DispatchWorkgroupCountOp>(
+ builder, dim)};
+ }
+ return procInfo;
+ },
+ {linalg::DistributionMethod::Cyclic, linalg::DistributionMethod::Cyclic,
+ linalg::DistributionMethod::Cyclic},
+ DenseMap<StringRef,
+ std::function<linalg::ProcInfo(OpBuilder &, Location)>>()};
+
+ patterns.add<TiledOpInterfaceTilingPattern>(
+ context,
+ linalg::LinalgTilingOptions()
+ .setTileSizes(ArrayRef<int64_t>{10, 0, 30})
+ .setDistributionOptions(workgroupDistributionOptions),
+ linalg::LinalgTransformationFilter(
+ StringAttr::get(context, "distribute_input"),
+ StringAttr::get(context, "distribute_output")));
+
+ patterns.add<TiledOpInterfaceTilingPattern>(
+ context,
+ linalg::LinalgTilingOptions().setTileSizes(ArrayRef<int64_t>{32}),
+ linalg::LinalgTransformationFilter(
+ StringAttr::get(context, "tiling_1d_stage5_fft_input"),
+ StringAttr::get(context, "tiling_1d_stage5_fft_output")));
+
+ patterns.add<TiledOpInterfaceTilingPattern>(
+ context,
+ linalg::LinalgTilingOptions().setTileSizes(ArrayRef<int64_t>{10, 32}),
+ linalg::LinalgTransformationFilter(
+ StringAttr::get(context, "tiling_2d_stage5_fft_input"),
+ StringAttr::get(context, "tiling_2d_stage5_fft_output")));
+
+ patterns.add<TiledOpInterfaceTilingPattern>(
+ context, linalg::LinalgTilingOptions().setTileSizes({0, 20}),
+ linalg::LinalgTransformationFilter(
+ StringAttr::get(context, "tiling_repeated_indices_scatter_input"),
+ StringAttr::get(context, "tiling_repeated_indices_scatter_output")));
+
+ if (failed(applyPatternsAndFoldGreedily(funcOp, std::move(patterns)))) {
+ return signalPassFailure();
+ }
+}
+
+std::unique_ptr<OperationPass<FuncOp>>
+IREE::LinalgExt::createTiledOpInterfaceTilingPass() {
+ return std::make_unique<TiledOpInterfaceTilingPass>();
+}
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/CMakeLists.txt b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/CMakeLists.txt
index 0cd7fd0..a174ba1 100644
--- a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/CMakeLists.txt
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/CMakeLists.txt
@@ -1,25 +1,44 @@
-add_mlir_library(IREELinalgExtPasses
- ConvertToLoops.cpp
- PadContractionToBlockSize.cpp
- Passes.cpp
+add_mlir_library(IREELinalgExtTransforms
+ InParallelToAsync.cpp
+ InParallelToSequentialFor.cpp
+ TilingExternalModels.cpp
+ TileToSequentialFor.cpp
+ TileToInParallel.cpp
Tiling.cpp
+ TilingToTileOp.cpp
+ Utils.cpp
+ PARTIAL_SOURCES_INTENDED
DEPENDS
- IREELinalgExtTransformsPassesIncGen
+ mlir-headers
+ IREELinalgExtDialect
LINK_LIBS PUBLIC
- IREEInputDialect
IREELinalgExtDialect
- MLIRAffine
+
+ MLIRAffineToStandard
+ MLIRAsync
+ MLIRSCFToControlFlow
+ MLIRLinalgToLLVM
+ MLIRVectorToLLVM
+ MLIRMathToLLVM
+ MLIRMemRefToLLVM
MLIRIR
+ MLIRMath
MLIRLinalg
MLIRLinalgTransforms
- MLIRMath
- MLIRMemRef
MLIRPass
MLIRSCF
- MLIRFunc
- MLIRSupport
- MLIRTensor
MLIRTransforms
)
+
+add_mlir_library(IREELinalgExtOpInterfaceImpl
+ LinalgExtBufferization.cpp
+
+ PARTIAL_SOURCES_INTENDED
+ LINK_LIBS PUBLIC
+ IREELinalgExtDialect
+
+ MLIRBufferization
+ MLIRTensorTransforms
+)
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/InParallelToAsync.cpp b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/InParallelToAsync.cpp
new file mode 100644
index 0000000..64514bb
--- /dev/null
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/InParallelToAsync.cpp
@@ -0,0 +1,91 @@
+//===- InParallelToAsync.cpp - Rewrite InParallel as Async ----------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include <cstdlib>
+
+#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.h"
+#include "iree-dialects/Dialect/LinalgExt/Transforms/Transforms.h"
+#include "iree-dialects/Dialect/LinalgExt/Transforms/Utils.h"
+#include "llvm/ADT/STLExtras.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Async/IR/Async.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/BlockAndValueMapping.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+using namespace mlir::iree_compiler::IREE::LinalgExt;
+
+FailureOr<Operation *> mlir::iree_compiler::IREE::LinalgExt::
+ InParallelOpToAsyncRewriter::returningMatchAndRewrite(
+ iree_compiler::IREE::LinalgExt::InParallelOp inParallelOp,
+ PatternRewriter &rewriter) const {
+ assert(inParallelOp.getNumResults() == 0 &&
+ "expected bufferized InParallelOp");
+
+ // Only consider the top level InParallelOp op and skip if it already
+ // contains an ExecuteOp.
+ if (inParallelOp
+ ->getParentOfType<iree_compiler::IREE::LinalgExt::InParallelOp>() ||
+ llvm::any_of(inParallelOp.getBody()->getOperations(),
+ [](Operation &op) { return isa<async::ExecuteOp>(&op); }))
+ return failure();
+
+ auto *ctx = inParallelOp.getContext();
+ Location loc = inParallelOp.getLoc();
+ Value zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
+ Value one = rewriter.create<arith::ConstantIndexOp>(loc, 1);
+ Value numThreads = inParallelOp.num_threads();
+
+ // Wrap the linalg_ext.in_parallel into an async::ExecuteOp.
+ // 1. Create the async::GroupType object on which we synchronize.
+ Value asyncGroup = rewriter.create<async::CreateGroupOp>(
+ loc, async::GroupType::get(ctx), numThreads);
+
+ // 2. Create a bodyless forOp.
+ scf::ForOp forOp = rewriter.create<scf::ForOp>(loc, zero, numThreads, one);
+ rewriter.setInsertionPointToStart(forOp.getBody());
+
+ // 3. Create an empty executeOp, nested within the forOp.
+ auto noopExec = [&](OpBuilder &executeBuilder, Location executeLoc,
+ ValueRange executeArgs) {};
+ auto executeOp =
+ rewriter.create<async::ExecuteOp>(loc, /*resultTypes=*/TypeRange(),
+ /*dependencies=*/ValueRange(),
+ /*operands=*/ValueRange(), noopExec);
+
+ // 3. Steal the iree_compiler::IREE::LinalgExt::InParallel ops, except the
+ // terminator, into the body of the async::ExecuteOp, just before the
+ // terminator.
+ SmallVector<Value> bbArgsTranslated{forOp.getInductionVar()};
+ rewriter.mergeBlocks(&inParallelOp.region().front(), executeOp.getBody(),
+ bbArgsTranslated);
+ // 3.b. Erase the terminator stolen from inParallelOp.
+ rewriter.eraseOp(&executeOp.getBody()->back());
+ // 3.c. Erase inParallelOp.
+ rewriter.eraseOp(inParallelOp);
+ // 3.d. Add ExecuteOp terminator.
+ rewriter.setInsertionPointToEnd(executeOp.getBody());
+ rewriter.create<async::YieldOp>(loc, ValueRange{});
+ // 3.e. Add to group within the loop.
+ rewriter.setInsertionPoint(forOp.getBody()->getTerminator());
+ rewriter.create<async::AddToGroupOp>(loc, rewriter.getIndexType(),
+ executeOp.token(), asyncGroup);
+
+ // 4. After the iree_compiler::IREE::LinalgExt::InParallel, await all async
+ // tasks in `asyncGroup`.
+ rewriter.setInsertionPointAfter(forOp);
+ return rewriter.create<async::AwaitAllOp>(loc, asyncGroup).getOperation();
+}
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/InParallelToSequentialFor.cpp b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/InParallelToSequentialFor.cpp
new file mode 100644
index 0000000..683629b
--- /dev/null
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/InParallelToSequentialFor.cpp
@@ -0,0 +1,111 @@
+//===- InParallelToSequentialFor.cpp.cpp - Rewrite InParallel as ForOp ---===//
+//
+// 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
+//
+//===---------------------------------------------------------------------===//
+
+#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.h"
+#include "iree-dialects/Dialect/LinalgExt/Transforms/Transforms.h"
+#include "iree-dialects/Dialect/LinalgExt/Transforms/Utils.h"
+#include "llvm/ADT/STLExtras.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/BlockAndValueMapping.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+using namespace mlir::iree_compiler::IREE::LinalgExt;
+
+namespace {
+
+SmallVector<Value> getValuesToYield(PerformConcurrentlyOp op) {
+ return llvm::to_vector(llvm::map_range(
+ op.yieldingOps(), [](ParallelInsertSliceOp op) { return op.dest(); }));
+}
+
+} // namespace
+
+FailureOr<scf::ForOp> InParallelOpToScfForRewriter::returningMatchAndRewrite(
+ InParallelOp inParallelOp, PatternRewriter &rewriter) const {
+ // Construct the loop bounds based on the canonical arithmetic progression.
+ Location loc = inParallelOp.getLoc();
+ Value zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
+ Value one = rewriter.create<arith::ConstantIndexOp>(loc, 1);
+ Value numThreads = inParallelOp.num_threads();
+
+ // Construct the op without a body builder: we need to clone the ops in the
+ // body explicitly after having access to the new bbArgs.
+ // As a consequence, `ensureTerminator` is not called and the `forOp` body
+ // has no terminator.
+ PerformConcurrentlyOp performConcurrentlyOp = inParallelOp.getTerminator();
+ SmallVector<Value> valuesToYield = getValuesToYield(performConcurrentlyOp);
+ scf::ForOp forOp =
+ rewriter.create<scf::ForOp>(loc, zero, numThreads, one, valuesToYield);
+
+ // Move the body while replacing the threadId by the forOp iv.
+ SmallVector<Value> bbArgsTranslated{forOp.getInductionVar()};
+ Block *body = forOp.getBody();
+ bool hasTerminator =
+ !body->empty() && body->back().hasTrait<OpTrait::IsTerminator>();
+ if (hasTerminator) {
+ rewriter.mergeBlockBefore(&inParallelOp.region().front(),
+ body->getTerminator(), bbArgsTranslated);
+ } else {
+ rewriter.mergeBlocks(&inParallelOp.region().front(), body,
+ bbArgsTranslated);
+ }
+
+ rewriter.setInsertionPointToStart(body);
+ BlockAndValueMapping bvm;
+ bvm.map(valuesToYield, forOp.getRegionIterArgs());
+
+ // Create sequential insertSlice ops.
+ SmallVector<Value> toYield;
+ rewriter.setInsertionPoint(performConcurrentlyOp);
+ for (ParallelInsertSliceOp op : performConcurrentlyOp.yieldingOps()) {
+ toYield.push_back(rewriter.createOrFold<tensor::InsertSliceOp>(
+ loc, op.source(), bvm.lookup(op.dest()), op.getMixedOffsets(),
+ op.getMixedSizes(), op.getMixedStrides()));
+ }
+
+ // performConcurrentlyOp.yieldedValues come from above, not from bbArgs.
+ // There is no rewriter method to make mergeBlocks update non-bbArgs.
+ // Need to manually clone + bvm all uses that are now nested under forOp.
+ // Warning: this replacement is currently optimistic and may change the
+ // semantics as explained in the pass description in Passes.td.
+ SmallVector<Operation *> opsToReplace;
+ for (Value toReplace : valuesToYield) {
+ for (OpOperand &u : toReplace.getUses()) {
+ Operation *op = u.getOwner();
+ if (!forOp->isProperAncestor(op)) continue;
+ opsToReplace.push_back(op);
+ }
+ }
+ for (Operation *op : opsToReplace) {
+ OpBuilder::InsertionGuard g(rewriter);
+ rewriter.setInsertionPoint(op);
+ Operation *cloned = rewriter.clone(*op, bvm);
+ rewriter.replaceOp(op, cloned->getResults());
+ }
+
+ // Insert terminator.
+ if (!hasTerminator) {
+ rewriter.setInsertionPointToEnd(body);
+ rewriter.create<scf::YieldOp>(loc, toYield);
+ }
+
+ // Cleanup and replace.
+ rewriter.eraseOp(performConcurrentlyOp);
+ rewriter.replaceOp(inParallelOp, forOp.getResults());
+
+ return forOp;
+}
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/LinalgExtBufferization.cpp b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/LinalgExtBufferization.cpp
new file mode 100644
index 0000000..6a03048
--- /dev/null
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/LinalgExtBufferization.cpp
@@ -0,0 +1,347 @@
+//===-- LinalgExtBufferization.cpp - Linalg Extension bufferization -------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "iree-dialects/Dialect/LinalgExt/LinalgExtBufferization.h"
+
+#include <mlir/IR/BuiltinOps.h>
+
+#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/IR/PatternMatch.h"
+
+using namespace mlir;
+using namespace mlir::iree_compiler::IREE::LinalgExt;
+
+/// Return the destinations that an InParallelOp is inserting into. One per
+/// ParallelInsertSliceOp.
+static SmallVector<OpOperand *> getInsertionDest(InParallelOp inParallelOp) {
+ Operation *terminator = inParallelOp.region().front().getTerminator();
+ auto performConcOp = dyn_cast<PerformConcurrentlyOp>(terminator);
+ assert(performConcOp && "expected PerformConcurrentlyOp as terminator");
+
+ SmallVector<OpOperand *> result;
+ performConcOp.walk([&](ParallelInsertSliceOp insertOp) {
+ result.push_back(&insertOp->getOpOperand(1) /*dest*/);
+ });
+
+ return result;
+}
+
+namespace mlir {
+
+using bufferization::BufferizableOpInterface;
+using bufferization::BufferizationState;
+using bufferization::BufferRelation;
+using bufferization::getMemRefType;
+using bufferization::replaceOpWithBufferizedValues;
+using bufferization::replaceOpWithNewBufferizedOp;
+using tensor::ExtractSliceOp;
+
+namespace iree_compiler {
+namespace IREE {
+namespace LinalgExt {
+
+/// Bufferization of InParallelOp. This also bufferizes the terminator of the
+/// region. There are op interfaces for the terminators (PerformConcurrentlyOp
+/// and ParallelInsertSliceOp), but these are only used during analysis. Not
+/// for bufferization.
+struct InParallelOpInterface
+ : public BufferizableOpInterface::ExternalModel<InParallelOpInterface,
+ InParallelOp> {
+ SmallVector<OpOperand *> getAliasingOpOperand(
+ Operation *op, OpResult opResult, const BufferizationState &state) const {
+ // Get OpOperand (dest) from corresponding ParallelInsertSliceOp.
+ auto inParallelOp = cast<InParallelOp>(op);
+ return {getInsertionDest(inParallelOp)[opResult.getResultNumber()]};
+ }
+
+ bool isMemoryWrite(Operation *op, OpResult opResult,
+ const BufferizationState &state) const {
+ // This op is a memory write. Stop lookup here to avoid finding false
+ // conflicts involving this op and one of the ops in the region. This is
+ // similar to how scf.if ops are analyzed.
+ return true;
+ }
+
+ bool isAllocationHoistingBarrier(Operation *op) const { return true; }
+
+ BufferRelation bufferRelation(Operation *op, OpResult opResult,
+ const BufferizationState &state) const {
+ return BufferRelation::Equivalent;
+ }
+
+ LogicalResult bufferize(Operation *op, RewriterBase &b,
+ const BufferizationState &state) const {
+ OpBuilder::InsertionGuard g(b);
+ auto inParallelOp = cast<InParallelOp>(op);
+ Block *body = &inParallelOp.region().front();
+ Operation *oldTerminator = body->getTerminator();
+ assert(isa<PerformConcurrentlyOp>(oldTerminator) &&
+ "unexpected terminator");
+
+ // Gather new results of the InParallelOp.
+ SmallVector<Value> newResults;
+ for (OpResult opResult : inParallelOp->getOpResults()) {
+ SmallVector<OpOperand *> insertDestOperands =
+ state.getAliasingOpOperand(opResult);
+ assert(insertDestOperands.size() == 1 &&
+ "expected exactly one aliasing OpOperand");
+ // Insert copies right before the PerformConcurrentlyOp terminator. They
+ // should not be inside terminator (which would be the default insertion
+ // point).
+ Value buffer = *state.getBuffer(
+ b, *insertDestOperands.front(), /*forceInPlace=*/false,
+ /*customCopyInsertionPoint=*/oldTerminator);
+ newResults.push_back(buffer);
+ Value destTensor = insertDestOperands.front()->get();
+
+ // Replace all uses of the insert dest tensor inside the InParallelOp
+ // with the result buffer.
+ OpBuilder::InsertionGuard g(b);
+ b.setInsertionPointToStart(body);
+ Value toTensorOp =
+ b.create<bufferization::ToTensorOp>(inParallelOp.getLoc(), buffer);
+ for (OpOperand &use : destTensor.getUses())
+ if (body->findAncestorOpInBlock(*use.getOwner()))
+ // This is a use inside the InParallelOp.
+ use.set(toTensorOp);
+ }
+
+ // Create new InParallelOp without any results.
+ TypeRange newResultTypes;
+ auto newInParallelOp = b.create<InParallelOp>(
+ inParallelOp.getLoc(), newResultTypes, inParallelOp.num_threads());
+
+ // Delete terminator.
+ newInParallelOp.getBody()->getTerminator()->erase();
+
+ // Move over block contents of the old op.
+ b.mergeBlocks(inParallelOp.getBody(), newInParallelOp.getBody(),
+ {newInParallelOp.getBody()->getArgument(0)});
+
+ // Bufferize terminator.
+ auto performConcurrentlyOp =
+ cast<PerformConcurrentlyOp>(newInParallelOp.getBody()->getTerminator());
+ b.setInsertionPoint(performConcurrentlyOp);
+ WalkResult walkResult =
+ performConcurrentlyOp.walk([&](ParallelInsertSliceOp insertOp) {
+ Location loc = insertOp.getLoc();
+ Type srcType = getMemRefType(
+ insertOp.source().getType().cast<RankedTensorType>(),
+ state.getOptions());
+ Type destType =
+ getMemRefType(insertOp.dest().getType().cast<RankedTensorType>(),
+ state.getOptions());
+ // ParallelInsertSliceOp bufferizes to a copy.
+ auto srcMemref = b.create<bufferization::ToMemrefOp>(
+ loc, srcType, insertOp.source());
+ auto destMemref = b.create<bufferization::ToMemrefOp>(
+ loc, destType, insertOp.dest());
+ Value subview = b.create<memref::SubViewOp>(
+ loc, destMemref, insertOp.getMixedOffsets(),
+ insertOp.getMixedSizes(), insertOp.getMixedStrides());
+ // This memcpy will fold away if everything bufferizes in-place.
+ if (failed(createMemCpy(b, insertOp.getLoc(), srcMemref, subview,
+ state.getOptions())))
+ return WalkResult::interrupt();
+ b.eraseOp(insertOp);
+ return WalkResult::advance();
+ });
+ if (walkResult.wasInterrupted()) return failure();
+
+ // Replace the op.
+ replaceOpWithBufferizedValues(b, op, newResults);
+
+ return success();
+ }
+};
+
+/// Nothing to do for PerformConcurrentlyOp.
+struct PerformConcurrentlyOpInterface
+ : public BufferizableOpInterface::ExternalModel<
+ PerformConcurrentlyOpInterface, PerformConcurrentlyOp> {
+ LogicalResult bufferize(Operation *op, RewriterBase &b,
+ const BufferizationState &state) const {
+ llvm_unreachable("op does not have any tensor OpOperands / OpResults");
+ return failure();
+ }
+};
+
+/// Return true if the (ExtractSliceOp, ParallelInsertSliceOp) pair match (i.e.
+/// equivalent operand / result and same offset/sizes/strides specification).
+static bool areEquivalentExtractSliceOps(const BufferizationState &state,
+ ExtractSliceOp st,
+ ParallelInsertSliceOp sti) {
+ if (!st || !sti) return false;
+ if (st != sti &&
+ !state.areEquivalentBufferizedValues(st.source(), sti.dest()))
+ return false;
+ if (!sameOffsetsSizesAndStrides(st, sti, isEqualConstantIntOrValue))
+ return false;
+ return true;
+}
+
+/// Return true if `value` is originating from an ExtractSliceOp that matches
+/// the given InsertSliceOp.
+static bool hasMatchingExtractSliceOp(const BufferizationState &state,
+ Value value,
+ ParallelInsertSliceOp insertOp) {
+ auto condition = [&](Value val) {
+ if (auto extractOp = val.getDefiningOp<ExtractSliceOp>())
+ if (areEquivalentExtractSliceOps(state, extractOp, insertOp)) return true;
+ return false;
+ };
+
+ return llvm::all_of(state.findValueInReverseUseDefChain(value, condition),
+ condition);
+}
+
+/// Analysis of ParallelInsertSliceOp.
+struct ParallelInsertSliceOpInterface
+ : public BufferizableOpInterface::ExternalModel<
+ ParallelInsertSliceOpInterface, ParallelInsertSliceOp> {
+ SmallVector<OpResult> getAliasingOpResult(
+ Operation *op, OpOperand &opOperand,
+ const BufferizationState &state) const {
+ if (&opOperand != &op->getOpOperand(1) /*dest*/) return {};
+
+ // ParallelInsertSliceOp itself has no results. Tensors are returned via
+ // the parent op.
+ auto inParallelOp = op->getParentOfType<InParallelOp>();
+ assert(inParallelOp &&
+ "could not find valid owner of parallel_insert_slice");
+
+ // The i-th ParallelInsertSliceOp result is returned via the i-th OpResult
+ // of the parent InParallelOp.
+ Block *block = op->getBlock();
+ unsigned int opIdx = 0;
+ for (ParallelInsertSliceOp insertOp :
+ block->getOps<ParallelInsertSliceOp>()) {
+ if (insertOp.getOperation() == op) break;
+ ++opIdx;
+ }
+ assert(opIdx < inParallelOp->getNumResults() &&
+ "could not find op inside terminator op");
+
+ return {inParallelOp->getResult(opIdx)};
+ }
+
+ bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+ const BufferizationState &state) const {
+ return true;
+ }
+
+ bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+ const BufferizationState &state) const {
+ return &opOperand == &op->getOpOperand(1) /*dest*/;
+ }
+
+ BufferRelation bufferRelation(Operation *op, OpResult opResult,
+ const BufferizationState &state) const {
+ return BufferRelation::Equivalent;
+ }
+
+ LogicalResult bufferize(Operation *op, RewriterBase &b,
+ const BufferizationState &state) const {
+ // Will be bufferized as part of InParallelOp.
+ return failure();
+ }
+
+ // TODO: This is copied from TensorInterfaceImpl.cpp. Find a way to share
+ // the code.
+ bool isNotConflicting(Operation *op, OpOperand *uRead,
+ OpOperand *uConflictingWrite,
+ const BufferizationState &state) const {
+ Operation *readingOp = uRead->getOwner();
+ Operation *conflictingWritingOp = uConflictingWrite->getOwner();
+
+ // Special rules for matching ExtractSliceOp/InsertSliceOp pairs. If
+ // uRead is an InsertSliceOp...
+ if (auto insertSliceOp = dyn_cast<ParallelInsertSliceOp>(readingOp)) {
+ // As an example, consider the following IR.
+ //
+ // %0 = tensor.extract_slice %t[%a, %b][%c, %d][1, 1] {inplace = [true] }
+ // %1 = linalg.fill %cst, %0 {inplace= [true] }
+ // %2 = tensor.insert_slice %1 into %t[%a, %b][%c, %d][1, 1]
+ // {inplace= [true] }
+
+ // TODO: Use insertSliceOp.getDestOpOperand etc. when available.
+ if (uRead == &insertSliceOp->getOpOperand(1) /*dest*/ &&
+ hasMatchingExtractSliceOp(state, uConflictingWrite->get(),
+ insertSliceOp))
+ // Case 1: The main insight is that InsertSliceOp reads only part of
+ // the destination tensor. The overwritten area is not read. If
+ // uConflictingWrite writes into exactly the memory location that is
+ // being read by uRead, this is not a conflict.
+ //
+ // In the above example:
+ // uRead = OpOperand 1 (%t) of tensor.insert_slice
+ // uConflictingWrite = OpOperand 1 (%0) of linalg.fill
+ //
+ // The read of %t does not conflict with the write of the FillOp
+ // (same aliases!) because the area that the FillOp operates on is
+ // exactly the one that is *not* read via %t.
+ return true;
+
+ if (uRead == &insertSliceOp->getOpOperand(0) /*source*/ &&
+ uConflictingWrite == &insertSliceOp->getOpOperand(1) /*dest*/ &&
+ hasMatchingExtractSliceOp(state, uRead->get(), insertSliceOp))
+ // Case 2: The read of the source tensor and the write to the dest
+ // tensor via an InsertSliceOp is not a conflict if the read is
+ // reading exactly that part of an equivalent tensor that the
+ // InsertSliceOp is writing.
+ //
+ // In the above example:
+ // uRead = OpOperand 0 (%1) of tensor.insert_slice
+ // uConflictingWrite = OpOperand 1 (%t) of tensor.insert_slice
+ return true;
+ }
+
+ // If uConflictingWrite is an InsertSliceOp...
+ if (auto insertSliceOp =
+ dyn_cast<ParallelInsertSliceOp>(conflictingWritingOp))
+ // As an example, consider the following IR.
+ //
+ // %0 = tensor.extract_slice %t[%a, %b][%c, %d][1, 1] {inplace = [true] }
+ // %1 = linalg.fill %cst, %0 {inplace= [true] }
+ // %2 = tensor.insert_slice %1 into %t[%a, %b][%c, %d][1, 1]
+ // {inplace= [true] }
+ // %3 = vector.transfer_read %1, %cst
+ //
+ // In the above example:
+ // uRead = OpOperand 0 (%1) of vector.transfer_read
+ // uConflictingWrite = OpOperand 1 (%t) of tensor.insert_slice
+ // lastWrite = %1
+ //
+ // This is not a conflict because the InsertSliceOp overwrites the
+ // memory segment of %1 with the exact same data. (Effectively, there
+ // is no memory write here.)
+ if (uConflictingWrite == &insertSliceOp->getOpOperand(1) /*dest*/ &&
+ state.areEquivalentBufferizedValues(uRead->get(),
+ insertSliceOp.source()) &&
+ hasMatchingExtractSliceOp(state, insertSliceOp.source(),
+ insertSliceOp))
+ return true;
+
+ return false;
+ }
+};
+} // namespace LinalgExt
+} // namespace IREE
+} // namespace iree_compiler
+} // namespace mlir
+
+void mlir::iree_compiler::IREE::LinalgExt::
+ registerBufferizableOpInterfaceExternalModels(DialectRegistry ®istry) {
+ registry.addOpInterface<InParallelOp, InParallelOpInterface>();
+ registry
+ .addOpInterface<PerformConcurrentlyOp, PerformConcurrentlyOpInterface>();
+ registry
+ .addOpInterface<ParallelInsertSliceOp, ParallelInsertSliceOpInterface>();
+}
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/TileToInParallel.cpp b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/TileToInParallel.cpp
new file mode 100644
index 0000000..83ece71
--- /dev/null
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/TileToInParallel.cpp
@@ -0,0 +1,132 @@
+//===- TileToInParallel.cpp.cpp - Rewrite TileOp as InParallel -----------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.h"
+#include "iree-dialects/Dialect/LinalgExt/Transforms/Transforms.h"
+#include "iree-dialects/Dialect/LinalgExt/Transforms/Utils.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+using namespace mlir::iree_compiler::IREE::LinalgExt;
+
+FailureOr<iree_compiler::IREE::LinalgExt::InParallelOp> mlir::iree_compiler::
+ IREE::LinalgExt::TileOpToInParallelRewriter::returningMatchAndRewrite(
+ iree_compiler::IREE::LinalgExt::TileOp tileOp,
+ PatternRewriter &rewriter) const {
+ // TODO: verifier.
+ assert(tileOp.getNumResults() > 0 &&
+ tileOp.outs().size() == tileOp.getNumResults());
+
+ // TODO: when supported, iterate over the tensor of sizes. This will be
+ // iterating through a level of indirection.
+
+ int64_t tiledDim = tileOp.tiled_dim();
+
+ // Construct the loop bounds based on the canonical arithmetic progression.
+ Location loc = tileOp.getLoc();
+ Value zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
+ Value tiledDimValue = rewriter.create<arith::ConstantIndexOp>(loc, tiledDim);
+ Value one = rewriter.create<arith::ConstantIndexOp>(loc, 1);
+ Value totalSize =
+ rewriter.create<tensor::DimOp>(loc, tileOp.outs().front(), tiledDimValue);
+ Value step = tileOp.tile_size();
+ assert(step.getType().isa<IndexType>() && "NYI: not an index type");
+
+ using AV = AffineValueExpr;
+ AffineBuilder ab(rewriter, loc);
+ AffineExpr i, j, M;
+ bindDims(rewriter.getContext(), i, j);
+ bindSymbols(rewriter.getContext(), M);
+ Value numThreads = ab.ceil(AV(i).bind(totalSize), AV(M).bind(step));
+
+ // Construct the op without a body builder: we need to clone the ops in the
+ // body explicitly after having access to the new bbArgs.
+ // As a consequence, `ensureTerminator` is not called and the body has no
+ // terminator.
+ iree_compiler::IREE::LinalgExt::InParallelOp inParallelOp =
+ rewriter.create<iree_compiler::IREE::LinalgExt::InParallelOp>(
+ loc, tileOp->getResultTypes(), numThreads);
+
+ // At the beginning of the InParallelOp, compute offset and sizes.
+ rewriter.setInsertionPointToStart(inParallelOp.getBody());
+
+ // Materialize the implicit subtensors as explicit subset_extract.
+ // TODO: generalize to multiple offset/chunk_size bbargs if needed.
+ // TODO: generalize the subset op.
+ SmallVector<Value> leadingOffsets, leadingSizes, leadingStrides;
+ for (int64_t i = 0; i < tiledDim; ++i) {
+ leadingOffsets.push_back(zero);
+ leadingSizes.push_back(
+ rewriter.createOrFold<tensor::DimOp>(loc, tileOp.outs().front(), i));
+ leadingStrides.push_back(one);
+ }
+ // clang-format off
+ Value offset = ab.mul(AV(i).bind(inParallelOp.getThreadIndex()),
+ AV(M).bind(step));
+ Value size = ab.min(
+ ValueRange{ab.sub(AV(i).bind(totalSize), AV(j).bind(offset)),
+ step});
+ // clang-format on
+ leadingOffsets.push_back(offset);
+ leadingSizes.push_back(size);
+ leadingStrides.push_back(one);
+
+ SmallVector<Value> implicitSubtensorExtracts;
+ for (Value tensor : tileOp.outs()) {
+ implicitSubtensorExtracts.push_back(
+ createSubsetExtractOpFromLeadingOffsetsSizesAndStrides(
+ rewriter, loc, tensor, leadingOffsets, leadingSizes,
+ leadingStrides));
+ }
+
+ // Get a reference to the TileOp terminator before the body is merged and it
+ // becomes too hard to get to the terminator.
+ auto tileYieldOp = cast<TileYieldOp>(tileOp.getBody()->getTerminator());
+
+ // Regroup the values that replace the tileOp's bbArg and move the body.
+ SmallVector<Value> bbArgsTranslated{offset, size};
+ llvm::append_range(bbArgsTranslated, implicitSubtensorExtracts);
+ rewriter.mergeBlockBefore(&tileOp.region().front(),
+ inParallelOp.getBody()->getTerminator(),
+ bbArgsTranslated);
+
+ // tileOp's terminator is not the terminator, insert explicit subset_insert
+ // ops and feed them to a new scf.yield terminator that we can now add.
+ PerformConcurrentlyOp performConcurrentlyOp = inParallelOp.getTerminator();
+
+ for (auto it : llvm::zip(tileYieldOp->getOperands(), tileOp.outs())) {
+ SmallVector<Value> offsets, sizes, strides;
+ completeOffsetsSizesAndStrides(rewriter, loc, std::get<0>(it),
+ leadingOffsets, leadingSizes, leadingStrides,
+ offsets, sizes, strides);
+ OpBuilder::InsertionGuard g(rewriter);
+ rewriter.setInsertionPoint(
+ performConcurrentlyOp.getBody()->getTerminator());
+ createParallelInsertSliceOpFromLeadingOffsetsSizesAndStrides(
+ rewriter, loc, std::get<0>(it), std::get<1>(it), offsets, sizes,
+ strides);
+ }
+
+ // Cleanup and replace.
+ rewriter.eraseOp(tileYieldOp);
+ rewriter.replaceOp(tileOp, inParallelOp.getResults());
+
+ return inParallelOp;
+}
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/TileToSequentialFor.cpp b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/TileToSequentialFor.cpp
new file mode 100644
index 0000000..657eedd
--- /dev/null
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/TileToSequentialFor.cpp
@@ -0,0 +1,106 @@
+//===- LowerToSCF.cpp.cpp - Lower to SCF ---------------------------------===//
+//
+// 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
+//
+//===---------------------------------------------------------------------===//
+
+#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.h"
+#include "iree-dialects/Dialect/LinalgExt/Transforms/Transforms.h"
+#include "iree-dialects/Dialect/LinalgExt/Transforms/Utils.h"
+#include "llvm/ADT/STLExtras.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+using namespace mlir::iree_compiler::IREE::LinalgExt;
+
+FailureOr<scf::ForOp> mlir::iree_compiler::IREE::LinalgExt::
+ TileOpToSCFRewriter::returningMatchAndRewrite(
+ iree_compiler::IREE::LinalgExt::TileOp tileOp,
+ PatternRewriter &rewriter) const {
+ // TODO: verifier.
+ assert(tileOp.getNumResults() > 0 &&
+ tileOp.outs().size() == tileOp.getNumResults());
+
+ // TODO: when supported, iterate over the tensor of sizes. This will be
+ // iterating through a level of indirection.
+
+ // Construct the loop bounds based on the canonical arithmetic progression.
+ Location loc = tileOp.getLoc();
+ Value zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
+ Value one = rewriter.create<arith::ConstantIndexOp>(loc, 1);
+ Value totalSize =
+ rewriter.create<tensor::DimOp>(loc, tileOp.outs().front(), zero);
+ Value step = tileOp.tile_size();
+ assert(step.getType().isa<IndexType>() && "NYI: not an index type");
+
+ // Construct the op without a body builder: we need to clone the ops in the
+ // body explicitly after having access to the new bbArgs.
+ // As a consequence, `ensureTerminator` is not called and the body has no
+ // terminator.
+ scf::ForOp forOp =
+ rewriter.create<scf::ForOp>(loc, zero, totalSize, step, tileOp.outs());
+
+ rewriter.setInsertionPointToStart(forOp.getBody());
+
+ // TODO: when supported, also compute from the tensor of sizes.
+ using AV = AffineValueExpr;
+ AffineBuilder ab(rewriter, loc);
+ AffineExpr i, j, M;
+ bindDims(rewriter.getContext(), i, j);
+ bindSymbols(rewriter.getContext(), M);
+
+ // Materialize the implicit subtensors as explicit subset_extract.
+ // TODO: generalize to multiple offset/chunk_size bbargs if needed.
+ // TODO: generalize the subset op.
+ Value offset = forOp.getInductionVar();
+ // clang-format off
+ Value size = ab.min(
+ ValueRange{ab.sub(AV(i).bind(totalSize), AV(j).bind(offset)),
+ step});
+ // clang-format on
+ SmallVector<Value> implicitSubtensorExtracts;
+ for (Value tensor : forOp.getRegionIterArgs()) {
+ implicitSubtensorExtracts.push_back(
+ createSubsetExtractOpFromLeadingOffsetsSizesAndStrides(
+ rewriter, loc, tensor, offset, size, one));
+ }
+
+ // Regroup the values that replace the tileOp's bbArg and move the body.
+ SmallVector<Value> bbArgsTranslated{offset, size};
+ llvm::append_range(bbArgsTranslated, implicitSubtensorExtracts);
+ rewriter.mergeBlocks(&tileOp.region().front(), forOp.getBody(),
+ bbArgsTranslated);
+ // tileOp's terminator is not the terminator, insert explicit subset_insert
+ // ops and feed them to a new scf.yield terminator that we can now add.
+ auto tileYieldOp = cast<TileYieldOp>(&forOp.getBody()->back());
+ SmallVector<Value> implicitSubtensorInserts;
+ for (auto it : llvm::zip(implicitSubtensorExtracts, tileYieldOp.getOperands(),
+ forOp.getRegionIterArgs())) {
+ implicitSubtensorInserts.push_back(createMatchingSubsetInsertOp(
+ rewriter, loc,
+ /*subsetExtractOp=*/
+ std::get<0>(it).getDefiningOp<tensor::ExtractSliceOp>(),
+ /*source=*/std::get<1>(it), /*dest=*/std::get<2>(it)));
+ }
+ // Insert terminator.
+ rewriter.setInsertionPointToEnd(forOp.getBody());
+ rewriter.create<scf::YieldOp>(loc, implicitSubtensorInserts);
+
+ // Cleanup and replace.
+ rewriter.eraseOp(tileYieldOp);
+ rewriter.replaceOp(tileOp, forOp.getResults());
+
+ return forOp;
+}
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/Tiling.cpp b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/Tiling.cpp
index 25df1f8..0e55970 100644
--- a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/Tiling.cpp
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/Tiling.cpp
@@ -1,360 +1,216 @@
-// Copyright 2021 The IREE Authors
+//===- Tiling.cpp - Tiling using TilingInterface --------------------------===//
//
-// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// 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
+//
+//===----------------------------------------------------------------------===//
-#include "iree-dialects/Dialect/Input/InputDialect.h"
-#include "iree-dialects/Dialect/Input/InputOps.h"
-#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtDialect.h"
#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.h"
-#include "iree-dialects/Dialect/LinalgExt/Transforms/PassDetail.h"
-#include "iree-dialects/Dialect/LinalgExt/Transforms/Passes.h"
-#include "iree-dialects/Dialect/LinalgExt/Transforms/Transforms.h"
-#include "llvm/ADT/TypeSwitch.h"
-#include "mlir/Dialect/Affine/IR/AffineOps.h"
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Dialect/Linalg/IR/Linalg.h"
-#include "mlir/Dialect/MemRef/IR/MemRef.h"
-#include "mlir/Dialect/SCF/SCF.h"
+#include "iree-dialects/Dialect/LinalgExt/Transforms/Utils.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/Dialect/Utils/StaticValueUtils.h"
-#include "mlir/IR/Matchers.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/OperationSupport.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
using namespace mlir;
-namespace IREE = mlir::iree_compiler::IREE;
-using namespace IREE::LinalgExt;
+using namespace mlir::iree_compiler::IREE::LinalgExt;
-//===----------------------------------------------------------------------===//
-// Utility methods for tiling a linalg_ext operation that implements a
-// TiledOpInterface
-//===----------------------------------------------------------------------===//
+// TODO: connect these patterns to PDL. Either via the transform dialect or via
+// PDLL.
-/// Returns failure if the options are unsupported.
-static LogicalResult verifySupportedTilingOptions(
- PatternRewriter &rewriter, Operation *op,
- const linalg::LinalgTilingOptions &options) {
- if (!options.interchangeVector.empty()) {
- return rewriter.notifyMatchFailure(op,
- "unsupported interchange during tiling");
+static bool isZero(Value v) {
+ if (auto cst = v.getDefiningOp<arith::ConstantIndexOp>())
+ return cst.value() == 0;
+ return false;
+}
+
+SmallVector<Value> tileToSCF(PatternRewriter &rewriter, TilingInterface op,
+ TilingInterface clonedOp, ValueRange tileSizes) {
+ // Compute lower and upper bounds of the loop nest.
+ SmallVector<Range> ranges = clonedOp.getIterationDomain(rewriter);
+ assert(tileSizes.size() <= ranges.size() &&
+ "expected tile sizes to match the number of loops");
+
+ // Fill the tile sizes with zeros for the untiled dimensions.
+ Location loc = op->getLoc();
+ SmallVector<Value> tileSizesVec(tileSizes.begin(), tileSizes.end());
+ if (ranges.size() != tileSizes.size()) {
+ Value zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
+ tileSizesVec.resize(ranges.size(), zero);
}
- if (options.loopType != linalg::LinalgTilingLoopType::Loops) {
- return rewriter.notifyMatchFailure(op,
- "only tiling with scf.for is supported");
- }
- if (options.distribution) {
- if (llvm::any_of(options.distribution->distributionMethod,
- [](linalg::DistributionMethod method) {
- return method != linalg::DistributionMethod::Cyclic;
- })) {
- return rewriter.notifyMatchFailure(op,
- "only cyclic distibution is allowed");
+
+ SmallVector<Value> lbs, dims, allDims, steps;
+ for (auto it : llvm::enumerate(ranges)) {
+ allDims.push_back(it.value().size);
+ if (!isZero(tileSizesVec[it.index()])) {
+ lbs.push_back(it.value().offset);
+ dims.push_back(it.value().size);
+ steps.push_back(tileSizesVec[it.index()]);
}
}
- return success();
-}
-/// Converts an `OpFoldResult` to a `Value` by building a constant op if
-/// if the `OpFoldResult` is an `IntegerAttr`.
-static Value getValue(OpBuilder &builder, Location loc,
- OpFoldResult valueOrAttr) {
- if (auto attr = valueOrAttr.dyn_cast<Attribute>()) {
- return builder.create<arith::ConstantIndexOp>(
- loc, attr.cast<IntegerAttr>().getInt());
- }
- return valueOrAttr.get<Value>();
-}
-
-/// Returns true if loop is untiled. Only checks if the value is statically
-/// zero. It is assumed that a `Value` defined by a constant op is already
-/// converted to an `IntegerAttr` of that value. So here just return true if
-/// this is an attribute with a zero value.
-static bool isUntiledLoop(OpFoldResult valueOrAttr) {
- Optional<int64_t> intVal = getConstantIntValue(valueOrAttr);
- return intVal && *intVal == 0;
-}
-
-/// Generates the tiled loops and the body by invoking the interface methods of
-/// TiledOpInterface.
-/// - `outputs` are the operands to use for outputs of the tiled operation.
-/// - `tileSizes` are tile sizes specified for all loops of the operation. If a
-/// loop is to be untiled it is set to 0.
-/// - `iteratorType` is the type of the loop iterator returned by the
-/// TiledOpInterface.
-/// - `loopBounds` are the bounds of all the loops of the op returned by the
-/// TiledOpInterface.
-/// - `loopDepth` is the current loop depth being processed.
-/// - `offsets` are the `Value`s that represent the position of the tile being
-/// operated on. The offsets are computed as the tiled loops are being
-/// generated.
-/// - `distributionInfo` is the proc_id and nprocs `Value`s to be used for
-/// distributed loops. It is a stack, and once an entry at the top of the
-/// stack is used for distribution it is popped before processing the inner
-/// loops.
-static FailureOr<TiledOp> tileInterfaceOpImpl(
- OpBuilder &builder, TiledOpInterface tilableOp, ValueRange outputs,
- MutableArrayRef<OpFoldResult> tileSizes, ArrayRef<StringRef> iteratorTypes,
- ArrayRef<Range> loopBounds, unsigned loopDepth,
- SmallVectorImpl<OpFoldResult> &offsets,
- ArrayRef<linalg::ProcInfo> distributionInfo) {
- Location loc = tilableOp.getLoc();
- // If this is the innermost loop, then generated the tiled implementation of
- // the op by invoking the TiledOpInterface methods.
- if (loopDepth == tileSizes.size()) {
- TiledOp ret;
- ret.op = tilableOp.getTiledImplementation(builder, outputs, offsets,
- tileSizes, ret.results);
- if (!ret.op) {
- return static_cast<LogicalResult>(
- tilableOp.emitOpError("failed to get tiled implementation"));
- }
- return ret;
- }
-
- // If tile size at this depth is empty, do nothing.
- if (isUntiledLoop(tileSizes[loopDepth])) {
- auto zeroAttr = builder.getI64IntegerAttr(0);
- offsets.push_back(zeroAttr);
- assert(matchPattern(loopBounds[loopDepth].offset, m_Zero()) &&
- "expected loop bounds to have lower bound of zero");
- tileSizes[loopDepth] = getAsOpFoldResult(loopBounds[loopDepth].size);
- return tileInterfaceOpImpl(builder, tilableOp, outputs, tileSizes,
- iteratorTypes, loopBounds, loopDepth + 1,
- offsets, distributionInfo);
- }
-
- // Generate an scf.for for the current loop depth.
- Value lb = loopBounds[loopDepth].offset;
- Value ub = loopBounds[loopDepth].size;
- // TODO(#7073): Put the check back. This is required by tiling linalg_ext.fft
- // op. We can put the check back after updating linalg_ext.fft semantics.
- // if (!matchPattern(loopBounds[loopDepth].stride, m_One())) {
- // return static_cast<LogicalResult>(
- // tilableOp.emitOpError("expected stride to be 1"));
- //}
- Value step = getValue(builder, loc, tileSizes[loopDepth]);
-
- // Update lb, ub and step for cyclic distribution.
- if (!distributionInfo.empty() &&
- iteratorTypes[loopDepth] == getParallelIteratorTypeName()) {
- linalg::updateBoundsForCyclicDistribution(
- builder, loc, distributionInfo.front().procId,
- distributionInfo.front().nprocs, lb, ub, step);
- distributionInfo = distributionInfo.drop_front();
- }
- FailureOr<TiledOp> innerReturnValue;
- bool isBufferTiling = tilableOp->getNumResults() == 0;
- ValueRange initValues(isBufferTiling ? ValueRange{} : outputs);
- auto forOp = builder.create<scf::ForOp>(
- loc, lb, ub, step, initValues,
- [&](OpBuilder &b, Location loc, Value iv, ValueRange args) {
- offsets.push_back(iv);
- auto affineMaps = AffineMap::inferFromExprList({ArrayRef<AffineExpr>{
- b.getAffineSymbolExpr(0),
- b.getAffineSymbolExpr(1) - b.getAffineDimExpr(0)}})[0];
- // Similar to linalg tiling, the tile size is the min(tileSizes, ub -
- // iv) to account for cases where tile size does not divide (ub - lb)
- // exactly.
- Value inBoundsTileSize = b.create<AffineMinOp>(
- loc, affineMaps,
- ValueRange{iv, getValue(builder, loc, tileSizes[loopDepth]), ub});
- tileSizes[loopDepth] = getAsOpFoldResult(inBoundsTileSize);
- // Recursively proceed to generate the tiled loop for the next level.
- innerReturnValue =
- tileInterfaceOpImpl(b, tilableOp, (isBufferTiling ? outputs : args),
- tileSizes, iteratorTypes, loopBounds,
- loopDepth + 1, offsets, distributionInfo);
- if (failed(innerReturnValue)) return;
- b.create<scf::YieldOp>(loc, innerReturnValue->results);
+ // Generate loop nest: One loop per dimension.
+ llvm::SmallPtrSet<Operation *, 1> preservedUses;
+ SmallVector<Value> destOperand = clonedOp.getDestinationOperands(rewriter);
+ auto loopNest = mlir::scf::buildLoopNest(
+ rewriter, loc, lbs, /*ubs=*/dims, steps, ValueRange(destOperand),
+ [&](OpBuilder &b, Location loc, ValueRange localIvs,
+ ValueRange iterArgs) -> scf::ValueVector {
+ // Compute offsets and sizes of ExtractSliceOp.
+ SmallVector<Value> offsets =
+ linalg::computeTileOffsets(b, loc, localIvs, tileSizesVec);
+ SmallVector<Value> sizes =
+ linalg::computeTileSizes(b, loc, localIvs, tileSizesVec, allDims);
+ // Create ExtractSliceOp: Extract a tile from the PadOp.
+ // Note: The PadOp is located outside of the loop nest. It is
+ // later moved inside by ExtractSliceOfPadTensorSwapPattern.
+ auto map =
+ AffineMap::getMultiDimIdentityMap(ranges.size(), b.getContext());
+ assert(clonedOp->getNumResults() == 1 && "expected single result op");
+ Value tiledOutput =
+ linalg::makeTiledShape(b, loc, clonedOp->getResult(0), tileSizesVec,
+ map, offsets, allDims, sizes);
+ auto sliceOp = tiledOutput.getDefiningOp<tensor::ExtractSliceOp>();
+ preservedUses.insert(sliceOp);
+ assert(sliceOp && "expected ExtractSliceOp");
+ // Insert the tile into the output tensor.
+ Value yieldValue =
+ createMatchingSubsetInsertOp(b, loc, sliceOp, sliceOp, iterArgs[0]);
+ return scf::ValueVector({yieldValue});
});
- if (failed(innerReturnValue)) {
- return innerReturnValue;
- }
- innerReturnValue->loops.insert(innerReturnValue->loops.begin(),
- forOp.getOperation());
- innerReturnValue->results = forOp.getResults();
- return innerReturnValue;
+ return loopNest.getResults();
}
-FailureOr<TiledOp> tileInterfaceOp(OpBuilder &b, TiledOpInterface tilableOp,
- const linalg::LinalgTilingOptions &options) {
- SmallVector<Value> dest = tilableOp.getDestinationOperands(b);
- if (dest.empty()) {
- return static_cast<LogicalResult>(tilableOp.emitOpError(
- "cannot tile operation without destination operands"));
- }
-
- SmallVector<StringRef> iteratorTypes = tilableOp.getLoopIteratorTypes();
- SmallVector<Value, 4> tileSizesVals =
- options.tileSizeComputationFunction(b, tilableOp);
- auto zeroAttr = b.getI64IntegerAttr(0);
-
- // The actual tile sizes used converts `Value` defined as constant 0, to a
- // zero integer attributes. Currently if the iterator type is not "parallel",
- // the tile size is forced to zero as well.
- auto tileSizes = getAsOpFoldResult(tileSizesVals);
- tileSizes.resize(iteratorTypes.size(), zeroAttr);
- for (auto en : llvm::enumerate(iteratorTypes)) {
- if (en.value() == getParallelIteratorTypeName()) continue;
- if (!isUntiledLoop(tileSizes[en.index()])) {
- return static_cast<LogicalResult>(tilableOp.emitOpError(
- "unimplemented tiling of non-parallel loop iterator type"));
- }
- }
-
- // Trivial early exit case of tile sizes being zero for all parallel loops.
- if (llvm::all_of(tileSizes, isUntiledLoop)) {
- return TiledOp{tilableOp, {}, {}};
- }
-
- SmallVector<Range> loopBounds = tilableOp.getIterationDomain(b);
- SmallVector<linalg::ProcInfo> distributionInfo;
- // If the tiled loops are distributed, get the proc_id and nprocs for the
- // distributed loops. First collect the parallel loops by iterating over the
- // tileSizes and getting the loops that are distribute, i.e.,
- // - parallel, i.e. iteratorTypes is "parallel"
- // - tiled, i.e. tileSize != 0
- if (options.distribution) {
- SmallVector<Range> distributedLoopRange;
- for (auto i : llvm::seq<unsigned>(0, tileSizes.size())) {
- if (isUntiledLoop(tileSizes[i])) continue;
- if (iteratorTypes[i] != getParallelIteratorTypeName()) continue;
- distributedLoopRange.push_back(loopBounds[i]);
- }
- distributionInfo = options.distribution->procInfo(b, tilableOp.getLoc(),
- distributedLoopRange);
- }
-
- SmallVector<OpFoldResult> offsets;
- return tileInterfaceOpImpl(b, tilableOp, dest, tileSizes, iteratorTypes,
- loopBounds, 0, offsets, distributionInfo);
-}
-
-LogicalResult TiledOpInterfaceBaseTilingPattern::matchAndRewriteBase(
- TiledOpInterface tilableOp, PatternRewriter &rewriter,
- TiledOp &result) const {
- if (failed(filter.checkAndNotify(rewriter, tilableOp))) {
- return failure();
- }
- if (failed(verifySupportedTilingOptions(rewriter, tilableOp, options))) {
- return failure();
- }
-
- FailureOr<TiledOp> res = tileInterfaceOp(rewriter, tilableOp, options);
- if (failed(res)) return res;
- result = *res;
- if (result.op) {
- filter.replaceLinalgTransformationFilter(rewriter, result.op);
- }
- return success();
-}
-
-//===----------------------------------------------------------------------===//
-// Test pass for tiling Linalg Ext ops
-//===----------------------------------------------------------------------===//
-
namespace {
-struct TiledOpInterfaceTilingPass
- : public TiledOpInterfaceTilingBase<TiledOpInterfaceTilingPass> {
- void getDependentDialects(DialectRegistry ®istry) const override {
- registry.insert<
- AffineDialect, IREE::Input::IREEInputDialect, linalg::LinalgDialect,
- IREE::LinalgExt::IREELinalgExtDialect, memref::MemRefDialect,
- func::FuncDialect, mlir::arith::ArithmeticDialect, math::MathDialect,
- tensor::TensorDialect, scf::SCFDialect>();
+
+/// The tiling here works by two steps. The first step is to create a loop based
+/// on the loop bounds of the operation obtained from `TilingInterface`.
+///
+/// ```mlir
+/// %1 = <tiling interface op> ins(...) outs(%0 : ...)
+/// ... <use_op> ... %1 ...
+/// ```
+///
+/// is rewritten using a "noop" subtensor extract/insert pair
+///
+/// ```mlir
+/// %1 = <tiling interface op> ins(...) outs(%0 : ...)
+/// %2 = scf.for %iv0 = ... iter_args(%arg0 = %0) {
+/// %3 = scf.for %iv1 = ... iter_args(%arg1 = %arg0) {
+/// ...
+/// %4 = tensor.extract_slice %1[%iv0, %iv1]....
+/// %5 = tensor.insert_slice %4 into %arg1[%iv0, %iv1]...
+/// scf.yield %5
+/// }
+/// scf.yield %3
+/// }
+/// ... <use_op> ... %2 ...
+/// ```
+///
+/// Following this the `TilingInterface` -> `tensor::ExtractSliceOp` pattern is
+/// replaced with
+///
+/// /// ```mlir
+/// %2 = scf.for %iv0 = ... iter_args(%arg0 = %0) {
+/// %3 = scf.for %iv1 = ... iter_args(%arg1 = %arg0) {
+/// ...
+/// %4 = tensor.extract_slice %0[%iv0, %iv1]
+/// %5 = <tiling interface op> ins(...) outs(%4 : ...)
+/// %6 = tensor.insert_slice %5 into %arg1[%iv0, %iv1]...
+/// scf.yield %6
+/// }
+/// scf.yield %3
+/// }
+/// ... <use_op> ... %2 ...
+/// ```
+///
+/// TODO(ravishankarm): The current approach seems to work for only tiling the
+/// parallel loops of the operation. Specifically,
+/// 1) the `%0` in the third snippet needs to be `%arg1`, for cases where the
+/// tiled loop is a reduction.
+/// 2) Current implementation is using the `getIterationDomain` method to get
+/// the
+/// initial loop structure as described in the second snippet. If any of
+/// those loops are reductions, then that IR snippet itself is wrong (replace
+/// this with the case of `linalg.matmul` and the error becomes apparent).
+
+/// First pattern to introduce the loop nests.
+struct OpTilingPattern : public OpInterfaceRewritePattern<TilingInterface> {
+ OpTilingPattern(MLIRContext *context, linalg::LinalgTilingOptions opt,
+ linalg::LinalgTransformationFilter filt)
+ : OpInterfaceRewritePattern<TilingInterface>(context),
+ options(opt),
+ filter(filt) {}
+
+ LogicalResult matchAndRewrite(TilingInterface op,
+ PatternRewriter &rewriter) const override {
+ if (failed(filter.checkAndNotify(rewriter, op))) return failure();
+
+ /// Currently only handle single result operations.
+ if (op->getNumResults() != 1) return failure();
+
+ Location loc = op->getLoc();
+ // Get rank and tile sizes.
+ SmallVector<Value> tileSizes =
+ options.tileSizeComputationFunction(rewriter, op);
+ auto iteratorTypes = op.getLoopIteratorTypes();
+ Value zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
+ tileSizes.resize(iteratorTypes.size(), zero);
+
+ /// Currently only handle operations with all parallel iterator types.
+ for (auto iteratorType : enumerate(iteratorTypes)) {
+ if (iteratorType.value() != getParallelIteratorTypeName() &&
+ !isZero(tileSizes[iteratorType.index()])) {
+ return rewriter.notifyMatchFailure(
+ op, "unhandled tiling of non-parallel iterator");
+ }
+ }
+
+ auto clonedOp = cast<TilingInterface>(rewriter.clone(*op.getOperation()));
+ SmallVector<Value> results = tileToSCF(rewriter, op, clonedOp, tileSizes);
+
+ filter.replaceLinalgTransformationFilter(rewriter, clonedOp);
+ rewriter.replaceOp(op, results);
+ return success();
}
- void runOnOperation() override;
+
+ private:
+ linalg::LinalgTilingOptions options;
+ linalg::LinalgTransformationFilter filter;
};
-} // namespace
-template <typename OpTy>
-static Value buildFlowWorkgroupInfoOp(OpBuilder &b, unsigned dim) {
- return b.template create<OpTy>(b.getInsertionPoint()->getLoc(), dim);
-}
+/// Second pattern to implement the switch of `TilingInterface ->
+/// tensor.extract_slice` to `tensor.extract_slice -> `TilingInterface`.
+struct SliceOpTiledOpSwapPattern
+ : public OpRewritePattern<tensor::ExtractSliceOp> {
+ SliceOpTiledOpSwapPattern(MLIRContext *context,
+ linalg::LinalgTilingOptions opt,
+ linalg::LinalgTransformationFilter filt)
+ : OpRewritePattern<tensor::ExtractSliceOp>(context),
+ options(opt),
+ filter(filt) {}
-void TiledOpInterfaceTilingPass::runOnOperation() {
- FuncOp funcOp = getOperation();
- MLIRContext *context = funcOp.getContext();
-
- RewritePatternSet patterns(context);
- patterns.add<TiledOpInterfaceTilingPattern>(
- context, linalg::LinalgTilingOptions().setTileSizes({10, 20}),
- linalg::LinalgTransformationFilter(
- StringAttr::get(context, "tiling_input"),
- StringAttr::get(context, "tiling_output")));
- patterns.add<TiledOpInterfaceTilingPattern>(
- context, linalg::LinalgTilingOptions().setTileSizes(ArrayRef<int64_t>{0}),
- linalg::LinalgTransformationFilter(
- StringAttr::get(context, "no_tiling_input"),
- StringAttr::get(context, "no_tiling_output")));
-
- patterns.add<TiledOpInterfaceTilingPattern>(
- context, linalg::LinalgTilingOptions().setTileSizes({0, 20}),
- linalg::LinalgTransformationFilter(
- StringAttr::get(context, "outer_reduce_input"),
- StringAttr::get(context, "outer_reduce_output")));
- patterns.add<TiledOpInterfaceTilingPattern>(
- context, linalg::LinalgTilingOptions().setTileSizes({10, 0, 0}),
- linalg::LinalgTransformationFilter(
- StringAttr::get(context, "inner_reduce_input"),
- StringAttr::get(context, "inner_reduce_output")));
-
- static linalg::LinalgLoopDistributionOptions workgroupDistributionOptions = {
- [](OpBuilder &builder, Location loc, ArrayRef<Range> parallelLoopRanges) {
- auto numParallelDims = parallelLoopRanges.size();
-
- SmallVector<linalg::ProcInfo, 3> procInfo(numParallelDims);
- for (size_t dim = 0; dim < numParallelDims; ++dim) {
- procInfo[numParallelDims - dim - 1] = {
- buildFlowWorkgroupInfoOp<IREE::Input::DispatchWorkgroupIDOp>(
- builder, dim),
- buildFlowWorkgroupInfoOp<IREE::Input::DispatchWorkgroupCountOp>(
- builder, dim)};
- }
- return procInfo;
- },
- {linalg::DistributionMethod::Cyclic, linalg::DistributionMethod::Cyclic,
- linalg::DistributionMethod::Cyclic},
- DenseMap<StringRef,
- std::function<linalg::ProcInfo(OpBuilder &, Location)>>()};
-
- patterns.add<TiledOpInterfaceTilingPattern>(
- context,
- linalg::LinalgTilingOptions()
- .setTileSizes(ArrayRef<int64_t>{10, 0, 30})
- .setDistributionOptions(workgroupDistributionOptions),
- linalg::LinalgTransformationFilter(
- StringAttr::get(context, "distribute_input"),
- StringAttr::get(context, "distribute_output")));
-
- patterns.add<TiledOpInterfaceTilingPattern>(
- context,
- linalg::LinalgTilingOptions().setTileSizes(ArrayRef<int64_t>{32}),
- linalg::LinalgTransformationFilter(
- StringAttr::get(context, "tiling_1d_stage5_fft_input"),
- StringAttr::get(context, "tiling_1d_stage5_fft_output")));
-
- patterns.add<TiledOpInterfaceTilingPattern>(
- context,
- linalg::LinalgTilingOptions().setTileSizes(ArrayRef<int64_t>{10, 32}),
- linalg::LinalgTransformationFilter(
- StringAttr::get(context, "tiling_2d_stage5_fft_input"),
- StringAttr::get(context, "tiling_2d_stage5_fft_output")));
-
- patterns.add<TiledOpInterfaceTilingPattern>(
- context, linalg::LinalgTilingOptions().setTileSizes({0, 20}),
- linalg::LinalgTransformationFilter(
- StringAttr::get(context, "tiling_repeated_indices_scatter_input"),
- StringAttr::get(context, "tiling_repeated_indices_scatter_output")));
-
- if (failed(applyPatternsAndFoldGreedily(funcOp, std::move(patterns)))) {
- return signalPassFailure();
+ LogicalResult matchAndRewrite(tensor::ExtractSliceOp sliceOp,
+ PatternRewriter &rewriter) const override {
+ auto sourceOp = sliceOp.source().getDefiningOp<TilingInterface>();
+ if (!sourceOp || !filter.hasReplacementFilter(sourceOp)) return failure();
+ SmallVector<Operation *> tiledOps = sourceOp.getTiledImplementation(
+ rewriter, sourceOp.getDestinationOperands(rewriter),
+ sliceOp.getMixedOffsets(), sliceOp.getMixedSizes(),
+ /*tileDestOperands=*/true);
+ assert(tiledOps.size() && "expected single tiled op");
+ Operation *tiledOp = tiledOps.front();
+ rewriter.replaceOp(sliceOp, tiledOp->getResults());
+ return success();
}
-}
-std::unique_ptr<OperationPass<FuncOp>>
-IREE::LinalgExt::createTiledOpInterfaceTilingPass() {
- return std::make_unique<TiledOpInterfaceTilingPass>();
-}
+ private:
+ linalg::LinalgTilingOptions options;
+ linalg::LinalgTransformationFilter filter;
+};
+
+} // namespace
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/TilingExternalModels.cpp b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/TilingExternalModels.cpp
new file mode 100644
index 0000000..7174daa
--- /dev/null
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/TilingExternalModels.cpp
@@ -0,0 +1,178 @@
+//===- TilingExternalModels.cpp - External models for TilingInterface -----===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "iree-dialects/Dialect/LinalgExt/Passes/Passes.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Debug.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
+#include "mlir/Interfaces/TilingInterface.h"
+
+#define DEBUG_TYPE "linalg-ext-tiling"
+
+using namespace mlir;
+using namespace mlir::linalg;
+using namespace mlir::iree_compiler::IREE::LinalgExt;
+
+static Value getAsValue(OpBuilder &b, Location loc, OpFoldResult ofr) {
+ if (auto v = ofr.dyn_cast<Value>()) return v;
+ return b.create<arith::ConstantIndexOp>(
+ loc, ofr.get<Attribute>().cast<IntegerAttr>().getInt());
+}
+static SmallVector<Value> getAsValues(OpBuilder &b, Location loc,
+ ArrayRef<OpFoldResult> ofrs) {
+ SmallVector<Value> vals;
+ vals.reserve(ofrs.size());
+ for (auto ofr : ofrs) vals.push_back(getAsValue(b, loc, ofr));
+ return vals;
+}
+
+static SmallVector<Value, 4> makeTiledInputShapes(OpBuilder &b, Location loc,
+ LinalgOp linalgOp,
+ ArrayRef<Value> valuesToTile,
+ ArrayRef<Value> ivsRef,
+ ArrayRef<Value> tileSizesRef,
+ ArrayRef<Value> sizeBounds) {
+ assert(static_cast<int64_t>(valuesToTile.size()) == linalgOp.getNumInputs() &&
+ "expected one value to tile for every operand");
+
+ Value zero = b.create<arith::ConstantIndexOp>(loc, 0);
+ SmallVector<Value> tileSizes{tileSizesRef.begin(), tileSizesRef.end()};
+ tileSizes.append(sizeBounds.size() - tileSizes.size(), zero);
+
+ // Construct (potentially temporary) mins and maxes on which to apply maps
+ // that define tile subshapes.
+ SmallVector<Value> lbs = computeTileOffsets(b, loc, ivsRef, tileSizes);
+ SmallVector<Value> subShapeSizes =
+ computeTileSizes(b, loc, ivsRef, tileSizes, sizeBounds);
+
+ SmallVector<Value, 4> tiledShapes;
+ tiledShapes.reserve(valuesToTile.size());
+ for (OpOperand *opOperand : linalgOp.getInputOperands()) {
+ Value shapedOp = valuesToTile[opOperand->getOperandNumber()];
+ LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for operand " << shapedOp);
+ AffineMap map = linalgOp.getTiedIndexingMap(opOperand);
+ LLVM_DEBUG(llvm::dbgs() << ": tiled: figure out subshape...\n");
+ tiledShapes.push_back(makeTiledShape(b, loc, shapedOp, tileSizes, map, lbs,
+ sizeBounds, subShapeSizes));
+ }
+
+ return tiledShapes;
+}
+
+namespace {
+
+/// External model implementation of TilingInterface for LinalgOps. This is
+/// templated on the actual Linalg named op for now since the registration of
+/// the external model requires the original operation.
+template <typename LinalgOpTy>
+struct LinalgOpTilingInterface
+ : public TilingInterface::ExternalModel<LinalgOpTilingInterface<LinalgOpTy>,
+ LinalgOpTy> {
+ SmallVector<Value> getDestinationOperands(Operation *op, OpBuilder &b) const {
+ LinalgOp linalgOp = cast<LinalgOp>(op);
+ return linalgOp.getOutputOperands();
+ }
+
+ SmallVector<StringRef> getLoopIteratorTypes(Operation *op) const {
+ LinalgOp linalgOp = cast<LinalgOp>(op);
+ SmallVector<StringRef> iteratorTypes;
+ iteratorTypes.reserve(linalgOp.iterator_types().size());
+ for (Attribute iteratorAttr : linalgOp.iterator_types()) {
+ iteratorTypes.push_back(iteratorAttr.cast<StringAttr>().getValue());
+ }
+ return iteratorTypes;
+ }
+
+ SmallVector<Range> getIterationDomain(Operation *op, OpBuilder &b) const {
+ LinalgOp linalgOp = cast<LinalgOp>(op);
+ return linalgOp.createLoopRanges(b, op->getLoc());
+ }
+
+ SmallVector<Operation *> getTiledImplementation(
+ Operation *op, OpBuilder &b, ValueRange tiledDest,
+ ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes,
+ bool tileDestOperands) const {
+ LinalgOp linalgOp = cast<LinalgOp>(op);
+ if (op->getNumResults() != 1) {
+ // TODO: Need a failure message here, but `notifyMatchFailure` is only a
+ // method on `PatternRewriter`.
+ return {};
+ }
+ Location loc = op->getLoc();
+ AffineMap shapeSizesToLoopsMap = linalgOp.getShapesToLoopsMap();
+ auto allShapeSizes = linalgOp.createFlatListOfOperandDims(b, loc);
+ if (!shapeSizesToLoopsMap) return {};
+
+ OpOperand *outOperand = linalgOp.getOutputOperand(0);
+ AffineMap indexingMap = linalgOp.getTiedIndexingMap(outOperand);
+ if (!indexingMap.isProjectedPermutation()) return {};
+
+ SmallVector<Value> offsetsVals = getAsValues(b, loc, offsets);
+ SmallVector<Value> sizeVals = getAsValues(b, loc, sizes);
+ SmallVector<Value> sizeBounds =
+ applyMapToValues(b, loc, shapeSizesToLoopsMap, allShapeSizes);
+
+ // The offsets and sizes form the slice operation only give you the tile
+ // size of the output. Use that compute the tile sizes and offsets of the
+ // loops. For loops not used to access the output, set the tile sizes to
+ // loop bounds and set the offset to 0.
+ Value zero = b.create<arith::ConstantIndexOp>(loc, 0);
+ SmallVector<Value> tileOffsets(sizeBounds.size(), zero);
+ SmallVector<Value> tileSizes = sizeBounds;
+ for (auto result : enumerate(indexingMap.getResults())) {
+ unsigned position = result.value().cast<AffineDimExpr>().getPosition();
+ tileOffsets[position] = offsetsVals[result.index()];
+ tileSizes[position] = sizeVals[result.index()];
+ }
+
+ SmallVector<Value> valuesToTile = linalgOp.getInputOperands();
+ SmallVector<Value> tiledOperands;
+ if (tileDestOperands) {
+ // Append the outputs then tile both the inputs and outputs.
+ valuesToTile.append(tiledDest.begin(), tiledDest.end());
+ tiledOperands = makeTiledShapes(b, loc, linalgOp, valuesToTile,
+ tileOffsets, tileSizes, sizeBounds);
+ } else {
+ // Only tile the inputs, then apped the outputs.
+ int64_t dim = offsets.size();
+ ArrayRef<Value> tileOffsetsRef{tileOffsets.begin(), tileOffsets.end()};
+ ArrayRef<Value> tileSizesRef{tileSizes.begin(), tileSizes.end()};
+ tiledOperands = makeTiledInputShapes(
+ b, loc, linalgOp, valuesToTile, tileOffsetsRef.take_front(dim + 1),
+ tileSizesRef.take_front(dim + 1), sizeBounds);
+ tiledOperands.append(tiledDest.begin(), tiledDest.end());
+ }
+ return {linalgOp.clone(b, loc, tiledDest.getTypes(), tiledOperands)};
+ }
+};
+} // namespace
+
+template <typename OpType>
+void registerOne(DialectRegistry ®istry) {
+ registry.addOpInterface<OpType, LinalgOpTilingInterface<OpType>>();
+}
+
+/// Variadic helper function.
+template <typename... OpTypes>
+void registerAll(DialectRegistry ®istry) {
+ // FIXME: In c++17 this can be simplified by using 'fold expressions'.
+ (void)std::initializer_list<int>{0, (registerOne<OpTypes>(registry), 0)...};
+}
+
+#define GET_OP_LIST
+
+void mlir::iree_compiler::IREE::LinalgExt::
+ registerTilingInterfaceExternalModels(DialectRegistry ®istry) {
+ registerOne<linalg::GenericOp>(registry);
+ registerAll<
+#include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
+ >(registry);
+}
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/TilingToTileOp.cpp b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/TilingToTileOp.cpp
new file mode 100644
index 0000000..ba8cc4d
--- /dev/null
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/TilingToTileOp.cpp
@@ -0,0 +1,106 @@
+//===- TilingToTileOp.cpp - Tiling using to TileOp TilingInterface --------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.h"
+#include "iree-dialects/Dialect/LinalgExt/Transforms/Transforms.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/OperationSupport.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+using namespace mlir::iree_compiler::IREE::LinalgExt;
+
+struct TilingResult {
+ TileOp tileOp;
+ Operation *tiledOp;
+};
+
+static TilingResult tileToTileOp(PatternRewriter &rewriter, TilingInterface op,
+ int64_t tiledDim, Value tileSize) {
+ Location loc = op->getLoc();
+ OpBuilder::InsertionGuard g(rewriter);
+ // TODO: Handle the case where the `loopRanges` are empty.
+ SmallVector<Range> loopRanges = op.getIterationDomain(rewriter);
+ assert(loopRanges.size() >= 1 &&
+ "expected at least a single loop in operation");
+ auto destOperands = op.getDestinationOperands(rewriter);
+ Operation *tiledOp = nullptr;
+ auto tileOp = rewriter.create<TileOp>(
+ loc, tileSize, destOperands, tiledDim,
+ [&](OpBuilder &b, Location loc, Value offset, Value size,
+ ValueRange outSlices) {
+ // TODO: support `getTiledImplementation` with >1 produced tiled ops.
+ int64_t nLoops = loopRanges.size();
+ SmallVector<OpFoldResult> tiledOffsets, tiledSizes;
+ tiledOffsets.reserve(nLoops);
+ tiledSizes.reserve(nLoops);
+ for (unsigned i = 0; i < nLoops; ++i) {
+ if (i == tiledDim) {
+ tiledOffsets.push_back(offset);
+ tiledSizes.push_back(size);
+ } else {
+ tiledOffsets.push_back(loopRanges[i].offset);
+ tiledSizes.push_back(loopRanges[i].size);
+ }
+ }
+ SmallVector<Operation *> tiledOps = op.getTiledImplementation(
+ b, outSlices, tiledOffsets, tiledSizes, /*tileDestOperands=*/false);
+ assert(tiledOps.size() == 1 && "expected single tiled op");
+ tiledOp = tiledOps.front();
+ b.create<TileYieldOp>(loc, tiledOp->getResults());
+ });
+ return TilingResult{tileOp, tiledOp};
+}
+
+FailureOr<Operation *> mlir::iree_compiler::IREE::LinalgExt::
+ LinalgExtTilingPattern::returningMatchAndRewrite(
+ TilingInterface op, PatternRewriter &rewriter) const {
+ /// Currently only handle single result operations.
+ if (op->getNumResults() != 1)
+ return rewriter.notifyMatchFailure(op, "Not a single result");
+
+ // Get rank and tile sizes.
+ // TODO: consider moving these checks to a common place that the TransformOp
+ // verifier can also use.
+ SmallVector<Value> tileSizes =
+ options.tileSizeComputationFunction(rewriter, op);
+ int64_t dim = -1;
+ for (auto en : llvm::enumerate(tileSizes)) {
+ Optional<int64_t> maybeTileSize = getConstantIntValue(en.value());
+ if (maybeTileSize && *maybeTileSize == 0) continue;
+ if (maybeTileSize && *maybeTileSize < 0)
+ return rewriter.notifyMatchFailure(op, "Negative tile size");
+ if (dim >= 0)
+ return rewriter.notifyMatchFailure(op,
+ "Could not find a single tiling dim");
+ dim = en.index();
+ }
+ if (dim < 0)
+ return rewriter.notifyMatchFailure(op,
+ "Could not find a single tiling dim");
+
+ /// Currently only handle tiling operations on a parallel iterator type.
+ auto loopIteratorTypes = op.getLoopIteratorTypes();
+ // Scalar operation, nothing to do, so just return.
+ if (loopIteratorTypes.empty())
+ return rewriter.notifyMatchFailure(op, "Scalar op, no tiling possible");
+ ArrayRef<StringRef> loopIteratorTypesRef(loopIteratorTypes);
+ if (loopIteratorTypesRef[dim] != getParallelIteratorTypeName())
+ return rewriter.notifyMatchFailure(op, "Trying to tile a non-parallel dim");
+
+ TilingResult tilingResult = tileToTileOp(rewriter, op, dim, tileSizes[dim]);
+ rewriter.replaceOp(op, tilingResult.tileOp->getResults());
+
+ return tilingResult.tiledOp;
+}
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/Utils.cpp b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/Utils.cpp
new file mode 100644
index 0000000..9b250b8
--- /dev/null
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/Utils.cpp
@@ -0,0 +1,104 @@
+//===- Utils.cpp - LinalgExt transform utils ------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "iree-dialects/Dialect/LinalgExt/Transforms/Utils.h"
+
+#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/OperationSupport.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+using namespace mlir::iree_compiler::IREE::LinalgExt;
+
+void mlir::iree_compiler::IREE::LinalgExt::completeOffsetsSizesAndStrides(
+ OpBuilder &b, Location loc, Value tensor, ArrayRef<Value> leadingOffsets,
+ ArrayRef<Value> leadingSizes, ArrayRef<Value> leadingStrides,
+ SmallVectorImpl<Value> &offsets, SmallVectorImpl<Value> &sizes,
+ SmallVectorImpl<Value> &strides) {
+ assert(leadingOffsets.size() == leadingSizes.size() &&
+ "expected matching lengths");
+ assert(leadingSizes.size() == leadingStrides.size() &&
+ "expected matching lengths");
+
+ auto rankedTensorType = tensor.getType().cast<RankedTensorType>();
+ int64_t tensorRank = rankedTensorType.getRank();
+ int64_t leadingRank = leadingOffsets.size();
+ offsets = SmallVector<Value>(leadingOffsets.begin(), leadingOffsets.end());
+ sizes = SmallVector<Value>(leadingSizes.begin(), leadingSizes.end());
+ strides = SmallVector<Value>(leadingStrides.begin(), leadingStrides.end());
+ if (leadingRank >= tensorRank) return;
+ Value zero = b.create<arith::ConstantIndexOp>(loc, 0);
+ Value one = b.create<arith::ConstantIndexOp>(loc, 1);
+ for (int64_t i = leadingRank, e = tensorRank; i < e; ++i) {
+ offsets.push_back(zero);
+ sizes.push_back(b.createOrFold<tensor::DimOp>(loc, tensor, i));
+ strides.push_back(one);
+ }
+}
+
+/// Create a tensor::ExtractSliceOp by auto-completing the missing trailing
+/// dimensions to always be offset = 0, size = dim, stride = 1.
+Value mlir::iree_compiler::IREE::LinalgExt::
+ createSubsetExtractOpFromLeadingOffsetsSizesAndStrides(
+ OpBuilder &b, Location loc, Value tensor,
+ ArrayRef<Value> leadingOffsets, ArrayRef<Value> leadingSizes,
+ ArrayRef<Value> leadingStrides) {
+ SmallVector<Value> offsets, sizes, strides;
+ completeOffsetsSizesAndStrides(b, loc, tensor, leadingOffsets, leadingSizes,
+ leadingStrides, offsets, sizes, strides);
+ return b.createOrFold<tensor::ExtractSliceOp>(loc, tensor, offsets, sizes,
+ strides);
+}
+
+/// Create a tensor::InsertSliceOp by auto-completing the missing trailing
+/// dimensions to always be offset = 0, size = dim, stride = 1.
+Value mlir::iree_compiler::IREE::LinalgExt::
+ createSubsetInsertOpFromLeadingOffsetsSizesAndStrides(
+ OpBuilder &b, Location loc, Value tensor, Value dest,
+ ArrayRef<Value> leadingOffsets, ArrayRef<Value> leadingSizes,
+ ArrayRef<Value> leadingStrides) {
+ SmallVector<Value> offsets, sizes, strides;
+ completeOffsetsSizesAndStrides(b, loc, tensor, leadingOffsets, leadingSizes,
+ leadingStrides, offsets, sizes, strides);
+ return b.createOrFold<tensor::InsertSliceOp>(loc, tensor, dest, offsets,
+ sizes, strides);
+}
+
+/// Create a iree_compiler::IREE::LinalgExt::ParallelInsertSliceOp by
+/// auto-completing the missing trailing dimensions to always be offset = 0,
+/// size = dim, stride = 1.
+Operation *mlir::iree_compiler::IREE::LinalgExt::
+ createParallelInsertSliceOpFromLeadingOffsetsSizesAndStrides(
+ OpBuilder &b, Location loc, Value tensor, Value dest,
+ ArrayRef<Value> leadingOffsets, ArrayRef<Value> leadingSizes,
+ ArrayRef<Value> leadingStrides) {
+ SmallVector<Value> offsets, sizes, strides;
+ completeOffsetsSizesAndStrides(b, loc, tensor, leadingOffsets, leadingSizes,
+ leadingStrides, offsets, sizes, strides);
+ return b.createOrFold<iree_compiler::IREE::LinalgExt::ParallelInsertSliceOp>(
+ loc, tensor, dest, offsets, sizes, strides);
+}
+
+/// Insert the `source` tensor into the `dest` tensor by creating the relevant
+/// `subset_insert` op. The details of the `subset_insert` op are retrieved
+/// from the `subset_extract` op so that they form a matching extract/insert
+/// pair.
+Value mlir::iree_compiler::IREE::LinalgExt::createMatchingSubsetInsertOp(
+ OpBuilder &b, Location loc, tensor::ExtractSliceOp subsetExtractOp,
+ Value source, Value dest) {
+ return b.create<tensor::InsertSliceOp>(
+ loc, subsetExtractOp.source().getType(), source, dest,
+ subsetExtractOp.offsets(), subsetExtractOp.sizes(),
+ subsetExtractOp.strides(), subsetExtractOp.static_offsets(),
+ subsetExtractOp.static_sizes(), subsetExtractOp.static_strides());
+}
diff --git a/llvm-external-projects/iree-dialects/python/CMakeLists.txt b/llvm-external-projects/iree-dialects/python/CMakeLists.txt
index ab6fd08..724982b 100644
--- a/llvm-external-projects/iree-dialects/python/CMakeLists.txt
+++ b/llvm-external-projects/iree-dialects/python/CMakeLists.txt
@@ -26,6 +26,14 @@
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT IREEDialectsPythonSources.Dialects
ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/iree/compiler"
+ TD_FILE dialects/IreeLinalgExtBinding.td
+ SOURCES dialects/iree_linalg_ext.py
+ DIALECT_NAME iree_linalg_ext
+)
+
+declare_mlir_dialect_python_bindings(
+ ADD_TO_PARENT IREEDialectsPythonSources.Dialects
+ ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/iree/compiler"
TD_FILE dialects/IreePyDmBinding.td
SOURCES
dialects/_iree_pydm_ops_ext.py
diff --git a/llvm-external-projects/iree-dialects/python/IREEDialectsModule.cpp b/llvm-external-projects/iree-dialects/python/IREEDialectsModule.cpp
index b3efba8..3647c47 100644
--- a/llvm-external-projects/iree-dialects/python/IREEDialectsModule.cpp
+++ b/llvm-external-projects/iree-dialects/python/IREEDialectsModule.cpp
@@ -90,6 +90,21 @@
py::arg("context") = py::none(), py::arg("load") = true);
//===--------------------------------------------------------------------===//
+ // IREELinalgExt
+ //===--------------------------------------------------------------------===//
+ auto iree_linalg_ext_m = m.def_submodule("iree_linalg_ext");
+ iree_linalg_ext_m.def(
+ "register_dialect",
+ [](MlirContext context, bool load) {
+ MlirDialectHandle handle = mlirGetDialectHandle__iree_linalg_ext__();
+ mlirDialectHandleRegisterDialect(handle, context);
+ if (load) {
+ mlirDialectHandleLoadDialect(handle, context);
+ }
+ },
+ py::arg("context") = py::none(), py::arg("load") = true);
+
+ //===--------------------------------------------------------------------===//
// IREEPyDMDialect
//===--------------------------------------------------------------------===//
auto iree_pydm_m = m.def_submodule("iree_pydm");
diff --git a/llvm-external-projects/iree-dialects/python/iree/compiler/dialects/IreeLinalgExtBinding.td b/llvm-external-projects/iree-dialects/python/iree/compiler/dialects/IreeLinalgExtBinding.td
new file mode 100644
index 0000000..da2ceae
--- /dev/null
+++ b/llvm-external-projects/iree-dialects/python/iree/compiler/dialects/IreeLinalgExtBinding.td
@@ -0,0 +1,13 @@
+// Copyright 2021 The IREE Authors
+//
+// 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
+
+#ifndef PYTHON_BINDINGS_IREE_LINALGEXT_OPS
+#define PYTHON_BINDINGS_IREE_LINALGEXT_OPS
+
+include "mlir/Bindings/Python/Attributes.td"
+include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.td"
+
+#endif // PYTHON_BINDINGS_IREE_LINALGEXT_OPS
diff --git a/llvm-external-projects/iree-dialects/python/iree/compiler/dialects/iree_linalg_ext.py b/llvm-external-projects/iree-dialects/python/iree/compiler/dialects/iree_linalg_ext.py
new file mode 100644
index 0000000..01fb430
--- /dev/null
+++ b/llvm-external-projects/iree-dialects/python/iree/compiler/dialects/iree_linalg_ext.py
@@ -0,0 +1,8 @@
+# Copyright 2021 The IREE Authors
+#
+# 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
+
+from ._iree_linalg_ext_ops_gen import *
+from .._mlir_libs._ireeDialects.iree_linalg_ext import *
diff --git a/llvm-external-projects/iree-dialects/test/iree_linalgext/canonicalize.mlir b/llvm-external-projects/iree-dialects/test/iree_linalgext/canonicalize.mlir
index acb8344..b8434d2 100644
--- a/llvm-external-projects/iree-dialects/test/iree_linalgext/canonicalize.mlir
+++ b/llvm-external-projects/iree-dialects/test/iree_linalgext/canonicalize.mlir
@@ -19,3 +19,24 @@
return %1: tensor<3x5xi32>
}
+
+// CHECK-LABEL: func @canonicalize_insert_slice_indices(
+// CHECK-SAME: %[[arg0:.*]]: tensor<?x?xf32>, %[[arg1:.*]]: tensor<?x?xf32>,
+// CHECK-SAME: %[[idx:.*]]: index
+func @canonicalize_insert_slice_indices(
+ %arg0 : tensor<?x?xf32>, %arg1: tensor<?x?xf32>,
+ %idx : index) -> tensor<?x?xf32>
+{
+ %cst = arith.constant 4.200000e+01 : f32
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+
+ %2 = iree_linalg_ext.in_parallel %idx -> (tensor<?x?xf32>) {
+ ^bb0(%arg3: index): // no predecessors
+ iree_linalg_ext.perform_concurrently {
+ // CHECK: iree_linalg_ext.parallel_insert_slice %[[arg0]] into %arg1[%[[idx]], 0] [1, 5] [1, 1]
+ iree_linalg_ext.parallel_insert_slice %arg0 into %arg1[%idx, %c0] [%c1, 5] [%c1, %c1] : tensor<?x?xf32> into tensor<?x?xf32>
+ }
+ }
+ return %2 : tensor<?x?xf32>
+}
diff --git a/llvm-external-projects/iree-dialects/test/iree_linalgext/invalid.mlir b/llvm-external-projects/iree-dialects/test/iree_linalgext/invalid.mlir
index f74461e..517e9c2 100644
--- a/llvm-external-projects/iree-dialects/test/iree_linalgext/invalid.mlir
+++ b/llvm-external-projects/iree-dialects/test/iree_linalgext/invalid.mlir
@@ -403,3 +403,46 @@
outs(%init : tensor<3x5xi32>) : tensor<3x5xi32>
return %0 : tensor<3x5xi32>
}
+
+// -----
+
+func @not_enough_results() -> () {
+ %num_threads = arith.constant 100 : index
+ // expected-error@+1 {{'iree_linalg_ext.in_parallel' op produces 1 results, but its terminator yields 0 values}}
+ %result = iree_linalg_ext.in_parallel %num_threads -> tensor<100xf32> {
+ ^bb0(%thread_idx : index):
+ iree_linalg_ext.perform_concurrently {}
+ }
+}
+
+// -----
+
+func @too_many_results(%1 : tensor<1xf32>, %out : tensor<100xf32>) -> () {
+ %num_threads = arith.constant 100 : index
+ // expected-error@+1 {{'iree_linalg_ext.in_parallel' op produces 1 results, but its terminator yields 2 values}}
+ %result = iree_linalg_ext.in_parallel %num_threads -> tensor<100xf32> {
+ ^bb0(%thread_idx : index):
+ %0 = arith.constant 1 : index
+ iree_linalg_ext.perform_concurrently {
+ iree_linalg_ext.parallel_insert_slice %1 into %out[%thread_idx][%0][%0] :
+ tensor<1xf32> into tensor<100xf32>
+ iree_linalg_ext.parallel_insert_slice %1 into %out[%thread_idx][%0][%0] :
+ tensor<1xf32> into tensor<100xf32>
+ }
+ }
+}
+
+// -----
+
+func @type_mismatch(%1 : tensor<1xf32>, %out : tensor<200xf32>) -> () {
+ %num_threads = arith.constant 100 : index
+ // expected-error@+1 {{'iree_linalg_ext.in_parallel' op type mismatch between 0th result of in_parallel ('tensor<200xf32>') and 0th result yielded by its terminator ('tensor<100xf32>')}}
+ %result = iree_linalg_ext.in_parallel %num_threads -> tensor<100xf32> {
+ ^bb0(%thread_idx : index):
+ %0 = arith.constant 1 : index
+ iree_linalg_ext.perform_concurrently {
+ iree_linalg_ext.parallel_insert_slice %1 into %out[%thread_idx][%0][%0] :
+ tensor<1xf32> into tensor<200xf32>
+ }
+ }
+}
diff --git a/llvm-external-projects/iree-dialects/test/iree_linalgext/roundtrip.mlir b/llvm-external-projects/iree-dialects/test/iree_linalgext/roundtrip.mlir
index a4c9fc6..98b2c71 100644
--- a/llvm-external-projects/iree-dialects/test/iree_linalgext/roundtrip.mlir
+++ b/llvm-external-projects/iree-dialects/test/iree_linalgext/roundtrip.mlir
@@ -539,3 +539,58 @@
// CHECK-SAME: dimensions(dense<[0, 1]> : tensor<2xi64>)
// CHECK-SAME: ins(%[[ARG0]]
// CHECK-SAME: outs(%[[INIT]]
+
+// -----
+
+// CHECK-LABEL: func @static_tile
+func @static_tile(%chunk_size: index, %in: tensor<?xf32>, %out: tensor<?xf32>, %out2: tensor<?xf32>) -> (tensor<?xf32>) {
+ %c0 = arith.constant 0: index
+ //%d0 = tensor.dim %out, %c0: tensor<?xf32>
+
+ // CHECK: iree_linalg_ext.tile %{{.*}} outs(%{{.*}}: tensor<?xf32>, %{{.*}}: tensor<?xf32>)
+ // CHECK: ^bb0(%{{.*}}: index, %{{.*}}: index, %{{.*}}: tensor<?xf32>, %{{.*}}: tensor<?xf32>):
+ %0:2 = iree_linalg_ext.tile %chunk_size outs(%out: tensor<?xf32>, %out2: tensor<?xf32>)
+ -> (tensor<?xf32>, tensor<?xf32>) {
+ // TODO: one offset and one size per tensor?
+ // If not necessary in the dense strided-array world, what about the rest?
+ ^bb0(%offset: index, %size: index, %st1: tensor<?xf32>, %st2: tensor<?xf32>):
+ // TODO: atm this is just 1-1: out-chunk-size -> in-size.
+ %1 = tensor.extract_slice %in[%offset][%size][1] : tensor<?xf32> to tensor<?xf32>
+ %3 = linalg.generic {
+ indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
+ iterator_types = ["parallel"]}
+ ins(%1: tensor<?xf32>) outs(%st1: tensor<?xf32>) {
+ ^bb0(%a: f32, %b:f32): // no predecessors
+ %f42 = arith.constant 42.0: f32
+ %tmp = arith.mulf %a, %f42: f32
+ linalg.yield %tmp: f32
+ } -> tensor<?xf32>
+ iree_linalg_ext.tile_yield %3, %st2: tensor<?xf32>, tensor<?xf32> // assumes dim is 0 and stacks
+ }
+ return %0#0: tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @simple_example
+func @simple_example(%in: tensor<100xf32>, %out: tensor<100xf32>) -> (tensor<100xf32>) {
+ %num_threads = arith.constant 100 : index
+ %result = iree_linalg_ext.in_parallel %num_threads -> tensor<100xf32> {
+ ^bb0(%thread_idx : index):
+ %0 = arith.constant 0 : index
+ %1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32>
+ iree_linalg_ext.perform_concurrently {
+ iree_linalg_ext.parallel_insert_slice %1 into %out[%thread_idx][%0][%0] :
+ tensor<1xf32> into tensor<100xf32>
+ }
+ }
+ return %result : tensor<100xf32>
+}
+
+func @no_terminator() -> () {
+ %num_threads = arith.constant 100 : index
+ iree_linalg_ext.in_parallel %num_threads -> () {
+ ^bb0(%thread_idx : index):
+ }
+ return
+}
diff --git a/llvm-external-projects/iree-dialects/test/python/smoketest.py b/llvm-external-projects/iree-dialects/test/python/smoketest.py
index cef2e6c..6804fec 100644
--- a/llvm-external-projects/iree-dialects/test/python/smoketest.py
+++ b/llvm-external-projects/iree-dialects/test/python/smoketest.py
@@ -2,10 +2,12 @@
import iree.compiler.ir
from iree.compiler.dialects import iree_input as iree_d
+from iree.compiler.dialects import iree_linalg_ext
from iree.compiler.dialects import iree_pydm as pydm_d
with iree.compiler.ir.Context() as ctx:
iree_d.register_dialect()
+ iree_linalg_ext.register_dialect()
pydm_d.register_dialect()
# iree_pydm types.
diff --git a/llvm-external-projects/iree-dialects/tools/iree-dialects-opt/CMakeLists.txt b/llvm-external-projects/iree-dialects/tools/iree-dialects-opt/CMakeLists.txt
index 735f948..6aecef3 100644
--- a/llvm-external-projects/iree-dialects/tools/iree-dialects-opt/CMakeLists.txt
+++ b/llvm-external-projects/iree-dialects/tools/iree-dialects-opt/CMakeLists.txt
@@ -13,6 +13,7 @@
IREEInputDialect
IREELinalgExtDialect
IREELinalgExtPasses
+ IREELinalgExtTransforms
IREEPyDMDialect
IREEPyDMPasses
)
diff --git a/llvm-external-projects/iree-dialects/tools/iree-dialects-opt/iree-dialects-opt.cpp b/llvm-external-projects/iree-dialects/tools/iree-dialects-opt/iree-dialects-opt.cpp
index fadfd96..d1e844b 100644
--- a/llvm-external-projects/iree-dialects/tools/iree-dialects-opt/iree-dialects-opt.cpp
+++ b/llvm-external-projects/iree-dialects/tools/iree-dialects-opt/iree-dialects-opt.cpp
@@ -7,7 +7,7 @@
#include "iree-dialects/Dialect/Input/InputDialect.h"
#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtDialect.h"
#include "iree-dialects/Dialect/LinalgExt/IR/TiledOpInterface.h"
-#include "iree-dialects/Dialect/LinalgExt/Transforms/Passes.h"
+#include "iree-dialects/Dialect/LinalgExt/Passes/Passes.h"
#include "iree-dialects/Dialect/PyDM/IR/PyDMDialect.h"
#include "iree-dialects/Dialect/PyDM/Transforms/Passes.h"
#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
