2c1c67674c
This patch is a first pass at making consistent syntax across the `LinalgTransformOp`s that use dynamic index lists for size parameters. Previously, there were two different forms: inline types in the list, or place them in the functional style tuple. This patch goes for the latter. In order to do this, the `printPackedOrDynamicIndexList`, `printDynamicIndexList` and their `parse` counterparts were modified so that the types can be optionally provided to the corresponding custom directives. All affected ops now use tablegen `assemblyFormat`, so custom `parse`/`print` functions have been removed. There are a couple ops that will likely add dynamic size support, and once that happens it should be made sure that the assembly remains consistent with the changes in this patch. The affected ops are as follows: `pack`, `pack_greedily`, `tile_using_forall`. The `tile_using_for` and `vectorize` ops already used this syntax, but their custom assembly was removed. --------- Co-authored-by: Oleksandr "Alex" Zinenko <ftynse@gmail.com>
167 lines
6.5 KiB
MLIR
167 lines
6.5 KiB
MLIR
// RUN: mlir-opt %s --transform-interpreter --split-input-file | FileCheck %s
|
|
|
|
// CHECK-LABEL: func.func @matmul_tensors_1(
|
|
func.func @matmul_tensors_1(
|
|
%arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>,
|
|
%arg2: tensor<128x128xf32>)
|
|
-> tensor<128x128xf32> {
|
|
// This operation is marked for tiling only.
|
|
// CHECK-COUNT-3: scf.for
|
|
// CHECK-COUNT-3: tensor.extract_slice
|
|
// CHECK: linalg.matmul
|
|
// CHECK-SAME: -> tensor<4x4xf32>
|
|
%0 = linalg.matmul { test.attrA }
|
|
ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
|
|
outs(%arg2: tensor<128x128xf32>)
|
|
-> tensor<128x128xf32>
|
|
func.return %0 : tensor<128x128xf32>
|
|
}
|
|
|
|
func.func @matmul_tensors_2(
|
|
%arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>,
|
|
%arg2: tensor<128x128xf32>)
|
|
-> tensor<128x128xf32> {
|
|
// This operation is marked f
|
|
// This operation is marked for tiling and vectorization.
|
|
// CHECK-COUNT-3: scf.for
|
|
// CHECK-COUNT-3: vector.transfer_read
|
|
// CHECK: vector.contract
|
|
// CHECK-NOT: linalg.matmul
|
|
// CHECK: vector.transfer_write
|
|
%0 = linalg.matmul { test.attrA, test.attrC }
|
|
ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
|
|
outs(%arg2: tensor<128x128xf32>)
|
|
-> tensor<128x128xf32>
|
|
func.return %0 : tensor<128x128xf32>
|
|
}
|
|
|
|
func.func @matmul_tensors_3(
|
|
%arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>,
|
|
%arg2: tensor<128x128xf32>)
|
|
-> tensor<128x128xf32> {
|
|
// This operation is marked for vectorization only.
|
|
// CHECK-NOT: scf.for
|
|
// CHECK-COUNT-3: vector.transfer_read
|
|
// CHECK: vector.contract
|
|
// CHECK-SAME: into vector<128x128xf32>
|
|
// CHECK: vector.transfer_write
|
|
%0 = linalg.matmul { test.attrC }
|
|
ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
|
|
outs(%arg2: tensor<128x128xf32>)
|
|
-> tensor<128x128xf32>
|
|
func.return %0 : tensor<128x128xf32>
|
|
}
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%root : !transform.any_op) {
|
|
transform.with_pdl_patterns %root : !transform.any_op {
|
|
^bb0(%arg0: !transform.any_op):
|
|
// Match matmul operations inside @matmul_tensors with test.attrA set.
|
|
pdl.pattern @pdl_target_attrA : benefit(1) {
|
|
%args = operands
|
|
%results = types
|
|
%attr = attribute
|
|
%0 = operation "linalg.matmul"(%args : !pdl.range<value>) {"test.attrA" = %attr}-> (%results : !pdl.range<type>)
|
|
// TODO: we don't want this, but it is the required terminator for pdl.pattern
|
|
rewrite %0 with "transform.dialect"
|
|
}
|
|
|
|
// Match matmul operations inside @matmul_tensors with test.attrC set.
|
|
pdl.pattern @pdl_target_attrC : benefit(1) {
|
|
%args = operands
|
|
%results = types
|
|
%attr = attribute
|
|
%0 = operation "linalg.matmul"(%args : !pdl.range<value>) {"test.attrC" = %attr}-> (%results : !pdl.range<type>)
|
|
// TODO: we don't want this, but it is the required terminator for pdl.pattern
|
|
rewrite %0 with "transform.dialect"
|
|
}
|
|
|
|
transform.sequence %arg0 : !transform.any_op failures(propagate) {
|
|
^bb1(%arg1: !transform.any_op):
|
|
%0 = pdl_match @pdl_target_attrA in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.structured.tile_using_for %0 tile_sizes [4, 4, 4] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
|
|
%1 = pdl_match @pdl_target_attrC in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
%2 = get_parent_op %1 {isolated_from_above} : (!transform.any_op) -> !transform.any_op
|
|
transform.structured.vectorize_children_and_apply_patterns %2 : (!transform.any_op) -> !transform.any_op
|
|
}
|
|
}
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
// CHECK-LABEL: @vectorize_one
|
|
func.func @vectorize_one(
|
|
%arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>,
|
|
%arg2: tensor<128x128xf32>)
|
|
-> tensor<128x128xf32> {
|
|
// CHECK: vector.contract
|
|
%0 = linalg.matmul {test.attrA}
|
|
ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
|
|
outs(%arg2: tensor<128x128xf32>)
|
|
-> tensor<128x128xf32>
|
|
func.return %0 : tensor<128x128xf32>
|
|
}
|
|
|
|
func.func @vectorize_none(
|
|
%arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>,
|
|
%arg2: tensor<128x128xf32>)
|
|
-> tensor<128x128xf32> {
|
|
// CHECK: linalg.matmul
|
|
%0 = linalg.matmul ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
|
|
outs(%arg2: tensor<128x128xf32>)
|
|
-> tensor<128x128xf32>
|
|
func.return %0 : tensor<128x128xf32>
|
|
}
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%root : !transform.any_op) {
|
|
transform.with_pdl_patterns %root : !transform.any_op {
|
|
^bb0(%arg0: !transform.any_op):
|
|
pdl.pattern @pdl_target : benefit(1) {
|
|
%args = operands
|
|
%results = types
|
|
%attr = attribute
|
|
%0 = operation "linalg.matmul"(%args : !pdl.range<value>) {"test.attrA" = %attr}-> (%results : !pdl.range<type>)
|
|
// TODO: we don't want this, but it is the required terminator for pdl.pattern
|
|
rewrite %0 with "transform.dialect"
|
|
}
|
|
|
|
transform.sequence %arg0 : !transform.any_op failures(propagate) {
|
|
^bb1(%arg1: !transform.any_op):
|
|
%0 = pdl_match @pdl_target in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
%1 = get_parent_op %0 {isolated_from_above} : (!transform.any_op) -> !transform.any_op
|
|
transform.structured.vectorize_children_and_apply_patterns %1 : (!transform.any_op) -> !transform.any_op
|
|
}
|
|
}
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
// CHECK-LABEL: @vectorize_all
|
|
func.func @vectorize_all(
|
|
%arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>, %arg2: tensor<128x128xf32>,
|
|
%arg3: tensor<128x128xf32>)
|
|
-> tensor<128x128xf32> {
|
|
// CHECK: vector.contract
|
|
%0 = linalg.matmul {test.attrA}
|
|
ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
|
|
outs(%arg2: tensor<128x128xf32>)
|
|
-> tensor<128x128xf32>
|
|
// CHECK: vector.contract
|
|
%1 = linalg.matmul ins(%arg0, %0: tensor<128x128xf32>, tensor<128x128xf32>)
|
|
outs(%arg3: tensor<128x128xf32>)
|
|
-> tensor<128x128xf32>
|
|
return %1 : tensor<128x128xf32>
|
|
}
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg0: !transform.any_op) {
|
|
transform.structured.vectorize_children_and_apply_patterns %arg0 : (!transform.any_op) -> !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|