[LinalgExt] Add tests related to parallel abstractions (#8623)
This PR adds unit tests related to LinalgExt parallel operations.
In the process, improve the error message for the TileToLinalgExtTileOp transformation
which used to fail silently.
This exposed a missing external model registration which is also fixed.
diff --git a/llvm-external-projects/iree-dialects/BUILD b/llvm-external-projects/iree-dialects/BUILD
index 3280a69..0e20c94 100644
--- a/llvm-external-projects/iree-dialects/BUILD
+++ b/llvm-external-projects/iree-dialects/BUILD
@@ -767,6 +767,7 @@
],
tags = ["hostonly"],
deps = [
+ "IREELinalgExtTransforms",
":IREEDialectsTest",
":IREEInputDialect",
":IREELinalgExtDialect",
@@ -777,6 +778,7 @@
":IREEPyDMTransforms",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ArithmeticDialect",
+ "@llvm-project//mlir:Async",
"@llvm-project//mlir:ControlFlowOps",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
diff --git a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgTransform/LinalgTransformOps.td b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgTransform/LinalgTransformOps.td
index d4a8f26..9da9a0c 100644
--- a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgTransform/LinalgTransformOps.td
+++ b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgTransform/LinalgTransformOps.td
@@ -357,7 +357,7 @@
let extraClassDeclaration = [{
::mlir::FailureOr<Operation *> applyToOne(
- ::mlir::TilingInterface target);
+ ::mlir::Operation *target);
}];
}
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgTransform/IR/LinalgTransformOps.cpp b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgTransform/IR/LinalgTransformOps.cpp
index 163bef8..2d8079c 100644
--- a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgTransform/IR/LinalgTransformOps.cpp
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgTransform/IR/LinalgTransformOps.cpp
@@ -823,7 +823,7 @@
//===----------------------------------------------------------------------===//
FailureOr<Operation *>
-transform::TileToLinalgExtTileOp::applyToOne(TilingInterface target) {
+transform::TileToLinalgExtTileOp::applyToOne(Operation *target) {
LinalgTilingOptions tilingOptions;
SmallVector<int64_t> tileSizes = extractI64Array(sizes());
if (!tileSizes.empty())
@@ -831,12 +831,13 @@
LinalgExt::LinalgExtTilingPattern pattern(this->getContext(), tilingOptions);
auto functionalTile =
- [&](TilingInterface op,
- PatternRewriter &rewriter) -> FailureOr<Operation *> {
- auto result = pattern.returningMatchAndRewrite(op, rewriter);
- if (failed(result))
+ [&](Operation *op, PatternRewriter &rewriter) -> FailureOr<Operation *> {
+ auto tilingInterfaceOp = dyn_cast<TilingInterface>(op);
+ if (!tilingInterfaceOp) {
+ op->emitError("Cannot tile op: Not a TilingInterface");
return failure();
- return result;
+ }
+ return pattern.returningMatchAndRewrite(tilingInterfaceOp, rewriter);
};
auto tileSeq = functional::SequenceBuilder().begin(std::move(functionalTile));
diff --git a/llvm-external-projects/iree-dialects/test/Dialect/iree_linalg_ext/bufferize-in-parallel.mlir b/llvm-external-projects/iree-dialects/test/Dialect/iree_linalg_ext/bufferize-in-parallel.mlir
new file mode 100644
index 0000000..4e5043c
--- /dev/null
+++ b/llvm-external-projects/iree-dialects/test/Dialect/iree_linalg_ext/bufferize-in-parallel.mlir
@@ -0,0 +1,105 @@
+// RUN: iree-dialects-opt %s -linalg-interp-transforms -canonicalize | FileCheck %s
+
+// CHECK-LABEL: func @parallel_insert_slice_no_conflict(
+// CHECK-SAME: %[[idx:.*]]: index, %[[idx2:.*]]: index,
+// CHECK-SAME: %[[arg1:.*]]: memref<?xf32, #{{.*}}>,
+// CHECK-SAME: %[[arg2:.*]]: memref<?xf32, #{{.*}}>
+func @parallel_insert_slice_no_conflict(
+ %idx: index, %idx2: index,
+ %arg1: tensor<?xf32> {linalg.inplaceable=true},
+ %arg2: tensor<?xf32> {linalg.inplaceable=true}) -> (tensor<?xf32>, f32)
+{
+ %cst = arith.constant 4.200000e+01 : f32
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+
+ // CHECK: iree_linalg_ext.in_parallel %[[idx2]] -> ()
+ %2 = iree_linalg_ext.in_parallel %idx2 -> (tensor<?xf32>) {
+ ^bb0(%arg3: index): // no predecessors
+ // CHECK: %[[subview:.*]] = memref.subview %[[arg2]][5] [%[[idx]]] [1]
+ %6 = tensor.extract_slice %arg2[5] [%idx] [%c1] : tensor<?xf32> to tensor<?xf32>
+ // CHECK: linalg.fill ins(%{{.*}}) outs(%[[subview]] : memref<?xf32
+ %8 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?xf32>) -> tensor<?xf32>
+
+ // CHECK: iree_linalg_ext.perform_concurrently
+ // CHECK-NOT: parallel_insert_slice
+ iree_linalg_ext.perform_concurrently {
+ iree_linalg_ext.parallel_insert_slice %8 into %arg2[5] [%idx] [%c1] : tensor<?xf32> into tensor<?xf32>
+ }
+ }
+
+ // CHECK: %[[load:.*]] = memref.load %[[arg2]]
+ %f = tensor.extract %2[%c0] : tensor<?xf32>
+
+ // CHECK: return %[[load]] : f32
+ return %2, %f : tensor<?xf32>, f32
+}
+
+// -----
+
+module {
+
+// CHECK-LABEL: func @parallel_insert_slice_with_conflict(
+// CHECK-SAME: %[[idx:.*]]: index, %[[idx2:.*]]: index,
+// CHECK-SAME: %[[arg1:.*]]: memref<?xf32, #{{.*}}>,
+// CHECK-SAME: %[[arg2:.*]]: memref<?xf32, #{{.*}}>
+func @parallel_insert_slice_with_conflict(
+ %idx: index, %idx2: index,
+ %arg1: tensor<?xf32> {linalg.inplaceable=true},
+ %arg2: tensor<?xf32> {linalg.inplaceable=true}) -> (f32, f32)
+{
+ %cst = arith.constant 4.200000e+01 : f32
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+
+ // The parallel_insert_slice_op bufferizes out-of-place, so we need an allocation.
+ // CHECK: %[[alloc1:.*]] = memref.alloc
+ // CHECK: iree_linalg_ext.in_parallel %[[idx2]] -> ()
+ %2 = iree_linalg_ext.in_parallel %idx2 -> (tensor<?xf32>) {
+ ^bb0(%arg3: index): // no predecessors
+ // Another alloc for the extract_slice op.
+ // CHECK: %[[alloc2:.*]] = memref.alloc
+ %6 = tensor.extract_slice %arg2[5] [%idx] [%c1] : tensor<?xf32> to tensor<?xf32>
+
+ // CHECK: linalg.fill ins(%{{.*}}) outs(%[[alloc2]] : memref<?xf32
+ %8 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?xf32>) -> tensor<?xf32>
+
+ // parallel_insert_slice buffer was already allocated but not copied yet.
+ //
+ // CHECK: linalg.generic {{.*}} ins(%[[arg2]]{{.*}}outs(%[[alloc1]]
+
+ // Now the copy of the actual insert_slice.
+ // CHECK: %[[subview1:.*]] = memref.subview %[[arg2]][5] [%[[idx]]] [1]
+ //
+ // CHECK: linalg.generic {{.*}} ins(%[[alloc2]]{{.*}}outs(%[[subview1]]
+ // CHECK: memref.dealloc %[[alloc2]]
+
+ // The terminator is empty.
+ // CHECK: iree_linalg_ext.perform_concurrently
+ // CHECK-NOT: parallel_insert_slice
+ iree_linalg_ext.perform_concurrently {
+ iree_linalg_ext.parallel_insert_slice %8 into %arg2[5] [%idx] [%c1] : tensor<?xf32> into tensor<?xf32>
+ }
+ }
+
+ // CHECK: %[[load:.*]] = memref.load %[[arg2]]
+ // CHECK: %[[load2:.*]] = memref.load %[[alloc1]]
+ // CHECK: memref.dealloc %[[alloc1]]
+ %f = tensor.extract %arg2[%c0] : tensor<?xf32>
+ %f2 = tensor.extract %2[%c0] : tensor<?xf32>
+
+ // CHECK: return %[[load2]], %[[load]] : f32, f32
+ return %f2, %f : f32, f32
+}
+
+pdl.pattern @pdl_target_2 : benefit(1) {
+ %0 = operation "func"
+ // TODO: we don't want this, but it is the required terminator for pdl.pattern
+ rewrite %0 with "iree_linalg_transform.apply"
+}
+
+iree_linalg_transform.sequence {
+ bufferize
+}
+
+}
diff --git a/llvm-external-projects/iree-dialects/test/Dialect/iree_linalg_ext/in-parallel-to-async.mlir b/llvm-external-projects/iree-dialects/test/Dialect/iree_linalg_ext/in-parallel-to-async.mlir
new file mode 100644
index 0000000..0ddf5d1
--- /dev/null
+++ b/llvm-external-projects/iree-dialects/test/Dialect/iree_linalg_ext/in-parallel-to-async.mlir
@@ -0,0 +1,68 @@
+// RUN: iree-dialects-opt %s -linalg-interp-transforms --split-input-file | FileCheck %s
+
+// CHECK-DAG: #[[$MUL_MAP:.*]] = affine_map<(d0)[s0] -> (d0 * s0)>
+// CHECK-DAG: #[[$SUB_MAP:.*]] = affine_map<(d0)[s0, s1] -> (-(d0 * s0) + s1, s0)>
+// CHECK-DAG: #[[$ID1_MAP:.*]] = affine_map<(d0) -> (d0)>
+#map0 = affine_map<(d0)[s0] -> (d0 ceildiv s0)>
+#map1 = affine_map<(d0)[s0] -> (d0 * s0)>
+#map2 = affine_map<(d0, d1) -> (d0 - d1)>
+#map3 = affine_map<(d0, d1) -> (d0, d1)>
+#map4 = affine_map<(d0) -> (d0)>
+
+module {
+ // CHECK-LABEL: func @static_tile
+ // CHECK-SAME: %[[CHUNK_SIZE:[0-9a-z]+]]: index
+ // CHECK-SAME: %[[IN:[0-9a-z]+]]: memref<?xf32>
+ // CHECK-SAME: %[[OUT:[0-9a-z]+]]: memref<?xf32>
+ func @static_tile(%arg0: index, %arg1: memref<?xf32>, %arg2: memref<?xf32>) {
+ %cst = arith.constant 4.200000e+01 : f32
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %0 = memref.dim %arg2, %c0 : memref<?xf32>
+ %1 = affine.apply #map0(%0)[%arg0]
+
+ // CHECK: %[[M:.*]] = memref.dim %{{.*}}, %{{.*}} : memref<?xf32>
+ // CHECK: %[[group:.*]] = async.create_group {{.*}}: !async.group
+ // CHECK: scf.for %[[IV:.*]] = {{.*}}
+ // CHECK: %[[token:.*]] = async.execute {
+ // CHECK: subview
+ // CHECK: subview
+ // CHECK: linalg.generic
+ // CHECK: async.yield
+ // CHECK: }
+ // CHECK: async.add_to_group %[[token]], %[[group]] : !async.token
+ // CHECK: }
+ // CHECK: async.await_all %[[group]]
+ iree_linalg_ext.in_parallel %1 -> () {
+ ^bb0(%arg3: index): // no predecessors
+ %3 = affine.apply #map1(%arg3)[%arg0]
+ %4 = affine.apply #map2(%0, %3)
+ %5 = affine.min #map3(%4, %arg0)
+
+ %6 = memref.subview %arg2[%3] [%5] [%c1] : memref<?xf32> to memref<?xf32, offset:?, strides:[?]>
+ %7 = memref.subview %arg1[%3] [%5] [1] : memref<?xf32> to memref<?xf32, offset:?, strides:[?]>
+
+ linalg.generic {indexing_maps = [#map4, #map4], iterator_types = ["parallel"]}
+ ins(%7 : memref<?xf32, offset:?, strides:[?]>) outs(%6 : memref<?xf32, offset:?, strides:[?]>) {
+ ^bb0(%arg4: f32, %arg5: f32): // no predecessors
+ %9 = arith.mulf %arg4, %cst : f32
+ linalg.yield %9 : f32
+ }
+
+ iree_linalg_ext.perform_concurrently {
+ }
+ }
+ return
+ }
+
+ pdl.pattern @match_iree_linalg_ext_in_parallel : benefit(1) {
+ %0 = operands
+ %1 = types
+ %2 = operation "iree_linalg_ext.in_parallel"(%0 : !pdl.range<value>) -> (%1 : !pdl.range<type>)
+ rewrite %2 with "iree_linalg_transform.apply"
+ }
+ iree_linalg_transform.sequence {
+ %0 = match @match_iree_linalg_ext_in_parallel
+ %1 = rewrite_iree_linalg_ext_in_parallel_to_async %0
+ }
+}
diff --git a/llvm-external-projects/iree-dialects/test/Dialect/iree_linalg_ext/in-parallel-to-sequential-for.mlir b/llvm-external-projects/iree-dialects/test/Dialect/iree_linalg_ext/in-parallel-to-sequential-for.mlir
new file mode 100644
index 0000000..de3614a
--- /dev/null
+++ b/llvm-external-projects/iree-dialects/test/Dialect/iree_linalg_ext/in-parallel-to-sequential-for.mlir
@@ -0,0 +1,101 @@
+// RUN: iree-dialects-opt %s -linalg-interp-transforms --split-input-file | FileCheck %s
+
+// CHECK-DAG: #[[$MUL_MAP:.*]] = affine_map<(d0)[s0] -> (d0 * s0)>
+// CHECK-DAG: #[[$SUB_MAP:.*]] = affine_map<(d0)[s0, s1] -> (-(d0 * s0) + s1, s0)>
+// CHECK-DAG: #[[$ID1_MAP:.*]] = affine_map<(d0) -> (d0)>
+#map0 = affine_map<(d0)[s0] -> (d0 ceildiv s0)>
+#map1 = affine_map<(d0)[s0] -> (d0 * s0)>
+#map2 = affine_map<(d0, d1) -> (d0 - d1)>
+#map3 = affine_map<(d0, d1) -> (d0, d1)>
+#map4 = affine_map<(d0) -> (d0)>
+
+module {
+
+ // CHECK-LABEL: func @static_tile_tensors
+ // CHECK-SAME: %[[CHUNK_SIZE:[0-9a-z]+]]: index
+ // CHECK-SAME: %[[IN:[0-9a-z]+]]: tensor<?xf32>
+ // CHECK-SAME: %[[OUT:[0-9a-z]+]]: tensor<?xf32>
+ func @static_tile_tensors(%arg0: index, %arg1: tensor<?xf32>, %arg2: tensor<?xf32>) -> tensor<?xf32> {
+ %cst = arith.constant 4.200000e+01 : f32
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %0 = tensor.dim %arg2, %c0 : tensor<?xf32>
+ %1 = affine.apply #map0(%0)[%arg0]
+
+ // CHECK: %[[M:.*]] = tensor.dim %{{.*}}, %{{.*}} : tensor<?xf32>
+ // CHECK: scf.for %[[IV:.*]] = {{.*}} iter_args(%[[OUT:.*]] = %{{.*}}) -> (tensor<?xf32>) {
+ %2 = iree_linalg_ext.in_parallel %1 -> (tensor<?xf32>) {
+ ^bb0(%arg3: index): // no predecessors
+ %3 = affine.apply #map1(%arg3)[%arg0]
+ %4 = affine.apply #map2(%0, %3)
+ %5 = affine.min #map3(%4, %arg0)
+
+ // Check the iter_arg is properly propagated
+ // CHECK: %[[O:.*]] = tensor.extract_slice %[[OUT]][{{.*}}] : tensor<?xf32> to tensor<?xf32>
+ %6 = tensor.extract_slice %arg2[%3] [%5] [%c1] : tensor<?xf32> to tensor<?xf32>
+ %7 = tensor.extract_slice %arg1[%3] [%5] [1] : tensor<?xf32> to tensor<?xf32>
+
+ %8 = linalg.generic {indexing_maps = [#map4, #map4], iterator_types = ["parallel"]} ins(%7 : tensor<?xf32>) outs(%6 : tensor<?xf32>) {
+ ^bb0(%arg4: f32, %arg5: f32): // no predecessors
+ %9 = arith.mulf %arg4, %cst : f32
+ linalg.yield %9 : f32
+ } -> tensor<?xf32>
+
+ // CHECK: %[[RES:.*]] = tensor.insert_slice %{{.*}} into %[[OUT]][{{.*}}] : tensor<?xf32> into tensor<?xf32>
+ // CHECK: scf.yield %[[RES]] : tensor<?xf32>
+ iree_linalg_ext.perform_concurrently {
+ iree_linalg_ext.parallel_insert_slice %8 into %arg2[%3] [%5] [%c1] : tensor<?xf32> into tensor<?xf32>
+ }
+ }
+ return %2 : tensor<?xf32>
+ }
+
+ // CHECK-LABEL: func @static_tile_buffers
+ // CHECK-SAME: %[[CHUNK_SIZE:[0-9a-z]+]]: index
+ // CHECK-SAME: %[[IN:[0-9a-z]+]]: memref<?xf32>
+ // CHECK-SAME: %[[OUT:[0-9a-z]+]]: memref<?xf32>
+ func @static_tile_buffers(%arg0: index, %arg1: memref<?xf32>, %arg2: memref<?xf32>) {
+ %cst = arith.constant 4.200000e+01 : f32
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %0 = memref.dim %arg2, %c0 : memref<?xf32>
+ %1 = affine.apply #map0(%0)[%arg0]
+
+ // CHECK: %[[C1:.*]] = arith.constant 1 : index
+ // CHECK: %[[M:.*]] = memref.dim %{{.*}}, %{{.*}} : memref<?xf32>
+ // CHECK: scf.for %[[IV:.*]] = {{.*}} step %[[C1]] {
+ iree_linalg_ext.in_parallel %1 -> () {
+ ^bb0(%arg3: index): // no predecessors
+ %3 = affine.apply #map1(%arg3)[%arg0]
+ %4 = affine.apply #map2(%0, %3)
+ %5 = affine.min #map3(%4, %arg0)
+
+ %6 = memref.subview %arg2[%3] [%5] [%c1] : memref<?xf32> to memref<?xf32, offset:?, strides:[?]>
+ %7 = memref.subview %arg1[%3] [%5] [1] : memref<?xf32> to memref<?xf32, offset:?, strides:[?]>
+
+ linalg.generic {indexing_maps = [#map4, #map4], iterator_types = ["parallel"]}
+ ins(%7 : memref<?xf32, offset:?, strides:[?]>) outs(%6 : memref<?xf32, offset:?, strides:[?]>) {
+ ^bb0(%arg4: f32, %arg5: f32): // no predecessors
+ %9 = arith.mulf %arg4, %cst : f32
+ linalg.yield %9 : f32
+ }
+
+ // Nothing is yielded, skip the terminator.
+ // CHECK-NOT: scf.yield
+ iree_linalg_ext.perform_concurrently {
+ }
+ }
+ return
+ }
+
+ pdl.pattern @match_iree_linalg_ext_in_parallel : benefit(1) {
+ %0 = operands
+ %1 = types
+ %2 = operation "iree_linalg_ext.in_parallel"(%0 : !pdl.range<value>) -> (%1 : !pdl.range<type>)
+ rewrite %2 with "iree_linalg_transform.apply"
+ }
+ iree_linalg_transform.sequence {
+ %0 = match @match_iree_linalg_ext_in_parallel
+ %1 = rewrite_iree_linalg_ext_in_parallel_to_scf_for %0
+ }
+}
diff --git a/llvm-external-projects/iree-dialects/test/Dialect/iree_linalg_ext/tile-to-in-parallel.mlir b/llvm-external-projects/iree-dialects/test/Dialect/iree_linalg_ext/tile-to-in-parallel.mlir
new file mode 100644
index 0000000..4617dda
--- /dev/null
+++ b/llvm-external-projects/iree-dialects/test/Dialect/iree_linalg_ext/tile-to-in-parallel.mlir
@@ -0,0 +1,62 @@
+// RUN: iree-dialects-opt %s -linalg-interp-transforms --split-input-file | FileCheck %s
+
+// CHECK-DAG: #[[$CEIL_MAP:.*]] = affine_map<()[s0, s1] -> (s1 ceildiv s0)>
+// CHECK-DAG: #[[$MUL_MAP:.*]] = affine_map<(d0)[s0] -> (d0 * s0)>
+// CHECK-DAG: #[[$SUB_MAP:.*]] = affine_map<(d0)[s0, s1] -> (-(d0 * s0) + s1, s0)>
+// CHECK-DAG: #[[$ID1_MAP:.*]] = affine_map<(d0) -> (d0)>
+
+module {
+ // CHECK-LABEL: func @static_tile
+ // CHECK-SAME: %[[CHUNK_SIZE:[0-9a-z]+]]: index
+ // CHECK-SAME: %[[IN:[0-9a-z]+]]: tensor<?xf32>
+ // CHECK-SAME: %[[OUT:[0-9a-z]+]]: tensor<?xf32>
+ func @static_tile(%chunk_size: index, %in: tensor<?xf32>, %out: tensor<?xf32>) -> (tensor<?xf32>) {
+ %c0 = arith.constant 0: index
+
+ // CHECK: %[[M:.*]] = tensor.dim %{{.*}}, %{{.*}} : tensor<?xf32>
+ // CHECK: %[[CEIL:.*]] = affine.apply #[[$CEIL_MAP]]()[%[[CHUNK_SIZE]], %[[M]]]
+ // CHECK: iree_linalg_ext.in_parallel %[[CEIL]] -> (tensor<?xf32>) {
+ %0 = iree_linalg_ext.tile %chunk_size outs(%out: tensor<?xf32>) -> (tensor<?xf32>) {
+
+ // CHECK: ^bb0(%[[TIDX:.*]]: index):
+ // CHECK: %[[OFFSET:.*]] = affine.apply #[[$MUL_MAP]](%[[TIDX]])[%[[CHUNK_SIZE]]]
+ // CHECK: %[[SIZE:.*]] = affine.min #[[$SUB_MAP]](%[[TIDX]])[%[[CHUNK_SIZE]], %[[M]]]
+ // CHECK: %[[O:.*]] = tensor.extract_slice %[[OUT]][%[[OFFSET]]] [%[[SIZE]]] [{{.*}}] : tensor<?xf32> to 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>):
+
+ // TODO: atm this is just 1-1: out-chunk-size -> in-size.
+ // CHECK: %[[I:.*]] = tensor.extract_slice %[[IN]][%[[OFFSET]]] [%[[SIZE]]] [{{.*}}] : tensor<?xf32> to tensor<?xf32>
+ %1 = tensor.extract_slice %in[%offset][%size][1] : tensor<?xf32> to tensor<?xf32>
+
+ // CHECK: %[[R:.*]] = linalg.generic {{.*}} ins(%[[I]] : tensor<?xf32>) outs(%[[O]] : 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>
+
+ // CHECK: iree_linalg_ext.perform_concurrently {
+ // CHECK: iree_linalg_ext.parallel_insert_slice %[[R]] into %[[OUT]][%[[OFFSET]]] [%[[SIZE]]] [{{.*}}] : tensor<?xf32> into tensor<?xf32>
+ iree_linalg_ext.tile_yield %3: tensor<?xf32>
+ }
+ return %0: tensor<?xf32>
+ }
+
+ pdl.pattern @match_iree_linalg_ext_tile : benefit(1) {
+ %0 = operands
+ %1 = types
+ %2 = operation "iree_linalg_ext.tile"(%0 : !pdl.range<value>) -> (%1 : !pdl.range<type>)
+ rewrite %2 with "iree_linalg_transform.apply"
+ }
+ iree_linalg_transform.sequence {
+ %0 = match @match_iree_linalg_ext_tile
+ %1 = rewrite_iree_linalg_ext_tile_to_in_parallel %0
+ }
+}
diff --git a/llvm-external-projects/iree-dialects/test/Dialect/iree_linalg_ext/tile-to-sequential-for.mlir b/llvm-external-projects/iree-dialects/test/Dialect/iree_linalg_ext/tile-to-sequential-for.mlir
new file mode 100644
index 0000000..9d6d4c3
--- /dev/null
+++ b/llvm-external-projects/iree-dialects/test/Dialect/iree_linalg_ext/tile-to-sequential-for.mlir
@@ -0,0 +1,58 @@
+// RUN: iree-dialects-opt %s -linalg-interp-transforms --split-input-file | FileCheck %s
+
+// CHECK-DAG: #[[$SUB_MAP:.*]] = affine_map<(d0)[s0, s1] -> (-d0 + s1, s0)>
+// CHECK-DAG: #[[$ID1_MAP:.*]] = affine_map<(d0) -> (d0)>
+
+module {
+ // CHECK-LABEL: func @static_tile
+ // CHECK-SAME: %[[CHUNK_SIZE:[0-9a-z]+]]: index
+ // CHECK-SAME: %[[IN:[0-9a-z]+]]: tensor<?xf32>
+ func @static_tile(%chunk_size: index, %in: tensor<?xf32>, %out: tensor<?xf32>, %out2: tensor<?xf32>) -> (tensor<?xf32>) {
+ %c0 = arith.constant 0: index
+
+ // CHECK: %[[M:.*]] = tensor.dim %{{.*}}, %{{.*}} : tensor<?xf32>
+ // CHECK: scf.for %[[IV:.*]] = {{.*}} iter_args(%[[OUT:.*]] = %{{.*}}, %[[OUT2:.*]] = %{{.*}}) -> (tensor<?xf32>, tensor<?xf32>) {
+ %0:2 = iree_linalg_ext.tile %chunk_size outs(%out: tensor<?xf32>, %out2: tensor<?xf32>)
+ -> (tensor<?xf32>, tensor<?xf32>) {
+
+ // CHECK: %[[SIZE:.*]] = affine.min #[[$SUB_MAP]](%[[IV]])[%[[CHUNK_SIZE]], %[[M]]]
+ // CHECK: %[[O:.*]] = tensor.extract_slice %[[OUT]][%[[IV]]] [%[[SIZE]]] [{{.*}}] : tensor<?xf32> to tensor<?xf32>
+ // CHECK: %[[O2:.*]] = tensor.extract_slice %[[OUT2]][%[[IV]]] [%[[SIZE]]] [{{.*}}] : tensor<?xf32> to 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.
+ // CHECK: %[[I:.*]] = tensor.extract_slice %[[IN]][%[[IV]]] [%[[SIZE]]] [{{.*}}] : tensor<?xf32> to tensor<?xf32>
+ %1 = tensor.extract_slice %in[%offset][%size][1] : tensor<?xf32> to tensor<?xf32>
+
+ // CHECK: %[[R:.*]] = linalg.generic {{.*}} ins(%[[I]] : tensor<?xf32>) outs(%[[O]] : 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>
+
+ // CHECK: %[[RES:.*]] = tensor.insert_slice %[[R]] into %[[OUT]][%[[IV]]] [%[[SIZE]]] [{{.*}}] : tensor<?xf32> into tensor<?xf32>
+ // CHECK: %[[RES2:.*]] = tensor.insert_slice %[[O2]] into %[[OUT2]][%[[IV]]] [%[[SIZE]]] [{{.*}}] : tensor<?xf32> into tensor<?xf32>
+ // CHECK: scf.yield %[[RES]], %[[RES2]] : tensor<?xf32>, tensor<?xf32>
+ iree_linalg_ext.tile_yield %3, %st2: tensor<?xf32>, tensor<?xf32> // assumes dim is 0 and stacks
+ }
+ return %0#0: tensor<?xf32>
+ }
+ pdl.pattern @match_iree_linalg_ext_tile : benefit(1) {
+ %0 = operands
+ %1 = types
+ %2 = operation "iree_linalg_ext.tile"(%0 : !pdl.range<value>) -> (%1 : !pdl.range<type>)
+ rewrite %2 with "iree_linalg_transform.apply"
+ }
+ iree_linalg_transform.sequence {
+ %0 = match @match_iree_linalg_ext_tile
+ %1 = rewrite_iree_linalg_ext_tile_to_scf_for %0
+ }
+}
diff --git a/llvm-external-projects/iree-dialects/test/Dialect/iree_linalg_ext/tiling-to-tile-op.mlir b/llvm-external-projects/iree-dialects/test/Dialect/iree_linalg_ext/tiling-to-tile-op.mlir
new file mode 100644
index 0000000..0e0ea93
--- /dev/null
+++ b/llvm-external-projects/iree-dialects/test/Dialect/iree_linalg_ext/tiling-to-tile-op.mlir
@@ -0,0 +1,68 @@
+// RUN: iree-dialects-opt %s -linalg-interp-transforms --split-input-file | FileCheck %s
+
+// CHECK: #[[$MAP:.+]] = affine_map<(d0, d1)[s0] -> (d0, -d1 + s0)>
+module {
+// CHECK-LABEL: matmul(
+// CHECK-SAME: %[[A:[0-9a-z]+]]: tensor<?x?xf32>
+// CHECK-SAME: %[[B:[0-9a-z]+]]: tensor<?x?xf32>
+// CHECK-SAME: %[[C:[0-9a-z]+]]: tensor<?x?xf32>
+ func @matmul(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<?x?xf32>) -> tensor<?x?xf32> {
+ // CHECK: %[[C10:.*]] = arith.constant 10 : index
+ // CHECK: iree_linalg_ext.tile %[[C10]] outs(%[[C]]: tensor<?x?xf32>) -> (tensor<?x?xf32>) {
+ // CHECK: ^bb0(%[[OFF:.*]]: index, %[[SZ:.*]]: index, %[[C_ITER:.*]]: tensor<?x?xf32>):
+ // CHECK: %[[tA:.*]] = tensor.extract_slice %[[A]]{{.*}} : tensor<?x?xf32> to tensor<?x?xf32>
+ // CHECK: %[[tB:.*]] = tensor.extract_slice %[[B]]{{.*}} : tensor<?x?xf32> to tensor<?x?xf32>
+ // CHECK: %[[RES:.*]] = linalg.matmul
+ // CHECK-SAME: ins(%[[tA]], %[[tB]] : tensor<?x?xf32>, tensor<?x?xf32>)
+ // CHECK-SAME: outs(%[[C_ITER]] : tensor<?x?xf32>) -> tensor<?x?xf32>
+ // CHECK: iree_linalg_ext.tile_yield %[[RES]] : tensor<?x?xf32>
+ %0 = linalg.matmul ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
+ outs(%C : tensor<?x?xf32>) -> (tensor<?x?xf32>)
+ return %0 : tensor<?x?xf32>
+ }
+ pdl.pattern @match_linalg_matmul : benefit(1) {
+ %0 = operands
+ %1 = types
+ %2 = operation "linalg.matmul"(%0 : !pdl.range<value>) -> (%1 : !pdl.range<type>)
+ rewrite %2 with "iree_linalg_transform.apply"
+ }
+ iree_linalg_transform.sequence {
+ %0 = match @match_linalg_matmul
+ %1 = tile_to_iree_linalg_ext_tile_op %0 {sizes = [10]}
+ }
+}
+
+// -----
+
+// CHECK: #[[$MAP:.+]] = affine_map<(d0, d1) -> (d0, -d1 + 100)>
+module {
+// CHECK-LABEL: matmul_static(
+// CHECK-SAME: %[[A:[0-9a-z]+]]: tensor<100x200xf32>
+// CHECK-SAME: %[[B:[0-9a-z]+]]: tensor<200x300xf32>
+// CHECK-SAME: %[[C:[0-9a-z]+]]: tensor<100x300xf32>
+ func @matmul_static(%A: tensor<100x200xf32>, %B: tensor<200x300xf32>, %C: tensor<100x300xf32>) -> tensor<100x300xf32> {
+ // CHECK: %[[C10:.*]] = arith.constant 10 : index
+ // CHECK: iree_linalg_ext.tile %[[C10]] outs(%[[C]]: tensor<100x300xf32>) -> (tensor<100x300xf32>) {
+ // CHECK: ^bb0(%[[OFF:.*]]: index, %[[SZ:.*]]: index, %[[C_ITER:.*]]: tensor<?x?xf32>):
+ // CHECK: %[[M:.*]] = affine.min #[[$MAP]](%[[SZ]], %[[OFF]])
+ // CHECK: %[[tA:.*]] = tensor.extract_slice %[[A]]{{.*}} : tensor<100x200xf32> to tensor<?x200xf32>
+ // CHECK: %[[tC:.*]] = tensor.cast %[[C_ITER]] : tensor<?x?xf32> to tensor<?x300xf32>
+ // CHECK: %[[RES:.*]] = linalg.matmul
+ // CHECK-SAME: ins(%[[tA]], %[[B]] : tensor<?x200xf32>, tensor<200x300xf32>)
+ // CHECK-SAME: outs(%[[tC]] : tensor<?x300xf32>) -> tensor<?x300xf32>
+ // CHECK: %[[RES_DYN:.*]] = tensor.cast %[[RES]] : tensor<?x300xf32> to tensor<?x?xf32>
+ // CHECK: iree_linalg_ext.tile_yield %[[RES_DYN]] : tensor<?x?xf32>
+ %0 = linalg.matmul ins(%A, %B : tensor<100x200xf32>, tensor<200x300xf32>) outs(%C : tensor<100x300xf32>) -> (tensor<100x300xf32>)
+ return %0 : tensor<100x300xf32>
+ }
+ pdl.pattern @match_linalg_matmul : benefit(1) {
+ %0 = operands
+ %1 = types
+ %2 = operation "linalg.matmul"(%0 : !pdl.range<value>) -> (%1 : !pdl.range<type>)
+ rewrite %2 with "iree_linalg_transform.apply"
+ }
+ iree_linalg_transform.sequence {
+ %0 = match @match_linalg_matmul
+ %1 = tile_to_iree_linalg_ext_tile_op %0 {sizes = [10]}
+ }
+}
diff --git a/llvm-external-projects/iree-dialects/test/Dialect/linalg_transform/bufferize.mlir b/llvm-external-projects/iree-dialects/test/Dialect/linalg_transform/bufferize.mlir
index 5ca985e..b4b7bde 100644
--- a/llvm-external-projects/iree-dialects/test/Dialect/linalg_transform/bufferize.mlir
+++ b/llvm-external-projects/iree-dialects/test/Dialect/linalg_transform/bufferize.mlir
@@ -19,7 +19,6 @@
// CHECK: }
}
-
pdl.pattern @pdl_target : benefit(1) {
%args = operands
%results = 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 60e1afd..8c739a0 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
@@ -2,6 +2,7 @@
# Local dialects.
IREEInputDialect
IREELinalgExtDialect
+ IREELinalgExtOpInterfaceImpl
IREELinalgExtPasses
IREELinalgExtTransforms
IREELinalgTransformDialect
@@ -12,6 +13,7 @@
IREEPyDMPasses
# Core dialects.
MLIRArithmetic
+ MLIRAsync
MLIRControlFlow
MLIRDialect
MLIRFunc
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 e18941d..0534deb 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,12 +7,14 @@
#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/LinalgExtBufferization.h"
#include "iree-dialects/Dialect/LinalgExt/Passes/Passes.h"
#include "iree-dialects/Dialect/LinalgTransform/LinalgTransformOps.h"
#include "iree-dialects/Dialect/LinalgTransform/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"
+#include "mlir/Dialect/Async/IR/Async.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
@@ -51,7 +53,8 @@
mlir::iree_compiler::IREE::PYDM::IREEPyDMDialect,
mlir::linalg::transform::LinalgTransformDialect,
// Upstream dialects
- mlir::arith::ArithmeticDialect,
+ mlir::async::AsyncDialect,
+ mlir::arith::ArithmeticDialect,
mlir::AffineDialect,
mlir::cf::ControlFlowDialect,
mlir::func::FuncDialect,
@@ -79,6 +82,8 @@
// External models.
IREE::LinalgExt::registerTiledOpInterfaceExternalModels(registry);
+ IREE::LinalgExt::registerTilingInterfaceExternalModels(registry);
+ IREE::LinalgExt::registerBufferizableOpInterfaceExternalModels(registry);
return mlir::asMainReturnCode(
mlir::MlirOptMain(argc, argv, "MLIR modular optimizer driver\n", registry,
