This revision removes the LinalgPromotion pattern and adds a `transform.structured.promotion` op. Since the LinalgPromotion transform allows the injection of arbitrary C++ via lambdas, the current transform op does not handle it. It is left for future work to decide what the right transform op control is for those cases. Note the underlying implementation remains unchanged and the mechanism is still controllable by lambdas from the API. During this refactoring it was also determined that the `dynamicBuffers` option does not actually connect to a change of behavior in the algorithm. This also exhibits that the related test is wrong (and dangerous). Both the option and the test are therefore removed. Lastly, a test that connects patterns using the filter-based mechanism is removed: all the independent pieces are already tested separately. Context: https://discourse.llvm.org/t/psa-retire-linalg-filter-based-patterns/63785 Differential Revision: https://reviews.llvm.org/D129649
200 lines
10 KiB
MLIR
200 lines
10 KiB
MLIR
// RUN: mlir-opt %s -test-linalg-transform-patterns=test-patterns -split-input-file | FileCheck %s
|
|
|
|
// CHECK-DAG: #[[$STRIDED_1D:.*]] = affine_map<(d0)[s0] -> (d0 + s0)>
|
|
// Map corresponding to a 2D memory access where the stride along the last dim is known to be 1.
|
|
// CHECK-DAG: #[[$STRIDED_2D_u_1:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
|
|
// CHECK-DAG: #[[$kn:.*]] = affine_map<(d0, d1, d2) -> (d2, d1)>
|
|
// CHECK-DAG: #[[$nm:.*]] = affine_map<(d0, d1, d2) -> (d1, d0)>
|
|
// CHECK-DAG: #[[$km:.*]] = affine_map<(d0, d1, d2) -> (d2, d0)>
|
|
|
|
func.func @dot(%x: memref<?xf32, offset: ?, strides: [1]>,
|
|
%y: memref<?xf32, offset: ?, strides: [1]>,
|
|
%v: memref<f32>) {
|
|
linalg.dot { __internal_linalg_transform__ = "MEM" }
|
|
ins(%x, %y: memref<?xf32, offset: ?, strides: [1]>,
|
|
memref<?xf32, offset: ?, strides: [1]>)
|
|
outs(%v: memref<f32>)
|
|
|
|
return
|
|
}
|
|
// CHECK-LABEL: func @dot
|
|
// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
|
|
// CHECK-DAG: %[[c1:.*]] = arith.constant 1 : index
|
|
// CHECK-DAG: %[[c8000:.*]] = arith.constant 8000 : index
|
|
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c8000]] {
|
|
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c1]] {
|
|
// CHECK: load
|
|
// CHECK: load
|
|
// CHECK: load
|
|
// CHECK: arith.mulf
|
|
// CHECK: arith.addf
|
|
// CHECK: store
|
|
|
|
func.func @matvec(%A: memref<?x?xf32, offset: ?, strides: [?, 1]>,
|
|
%x: memref<?xf32, offset: ?, strides: [1]>,
|
|
%y: memref<?xf32, offset: ?, strides: [1]>) {
|
|
linalg.matvec
|
|
ins(%A, %x: memref<?x?xf32, offset: ?, strides: [?, 1]>,
|
|
memref<?xf32, offset: ?, strides: [1]>)
|
|
outs(%y: memref<?xf32, offset: ?, strides: [1]>)
|
|
return
|
|
}
|
|
// CHECK-LABEL: func @matvec
|
|
// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
|
|
// CHECK-DAG: %[[c5:.*]] = arith.constant 5 : index
|
|
// CHECK-DAG: %[[c6:.*]] = arith.constant 6 : index
|
|
// CHECK: scf.parallel {{.*}} step (%[[c5]])
|
|
// CHECK: scf.for {{.*}} step %[[c6]]
|
|
// CHECK: linalg.matvec
|
|
// CHECK: ins({{.*}}: memref<?x?xf32, #[[$STRIDED_2D_u_1]]>, memref<?xf32, #[[$STRIDED_1D]]>)
|
|
// CHECK: outs({{.*}}: memref<?xf32, #[[$STRIDED_1D]]>)
|
|
|
|
func.func @matmul(%A: memref<?x?xf32, offset: ?, strides: [?, 1]>,
|
|
%B: memref<?x?xf32, offset: ?, strides: [?, 1]>,
|
|
%C: memref<?x?xf32, offset: ?, strides: [?, 1]>) {
|
|
linalg.matmul { __internal_linalg_transform__ = "MEM" }
|
|
ins(%A, %B: memref<?x?xf32, offset: ?, strides: [?, 1]>,
|
|
memref<?x?xf32, offset: ?, strides: [?, 1]>)
|
|
outs(%C: memref<?x?xf32, offset: ?, strides: [?, 1]>)
|
|
return
|
|
}
|
|
// CHECK-LABEL: func @matmul
|
|
// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
|
|
// CHECK-DAG: %[[c2:.*]] = arith.constant 2 : index
|
|
// CHECK-DAG: %[[c3:.*]] = arith.constant 3 : index
|
|
// CHECK-DAG: %[[c4:.*]] = arith.constant 4 : index
|
|
// CHECK-DAG: %[[c20:.*]] = arith.constant 20 : index
|
|
// CHECK-DAG: %[[c30:.*]] = arith.constant 30 : index
|
|
// CHECK-DAG: %[[c40:.*]] = arith.constant 40 : index
|
|
// CHECK-DAG: %[[c200:.*]] = arith.constant 200 : index
|
|
// CHECK-DAG: %[[c300:.*]] = arith.constant 300 : index
|
|
// CHECK-DAG: %[[c400:.*]] = arith.constant 400 : index
|
|
// CHECK-DAG: %[[c2000:.*]] = arith.constant 2000 : index
|
|
// CHECK-DAG: %[[c3000:.*]] = arith.constant 3000 : index
|
|
// CHECK-DAG: %[[c4000:.*]] = arith.constant 4000 : index
|
|
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c2000]] {
|
|
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c3000]] {
|
|
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c4000]] {
|
|
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c200]] {
|
|
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c300]] {
|
|
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c400]] {
|
|
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c20]] {
|
|
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c30]] {
|
|
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c40]] {
|
|
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c2]] {
|
|
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c3]] {
|
|
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c4]] {
|
|
// CHECK: linalg.matmul
|
|
// CHECK: ins({{.*}}: memref<?x?xf32, #[[$STRIDED_2D_u_1]]>, memref<?x?xf32, #[[$STRIDED_2D_u_1]]>)
|
|
// CHECK: outs({{.*}}: memref<?x?xf32, #[[$STRIDED_2D_u_1]]>)
|
|
|
|
#matmul_accesses = [
|
|
affine_map<(m, n, k) -> (m, k)>,
|
|
affine_map<(m, n, k) -> (k, n)>,
|
|
affine_map<(m, n, k) -> (m, n)>
|
|
]
|
|
#generic_matmul_trait = {
|
|
args_in = 2,
|
|
args_out = 1,
|
|
indexing_maps = #matmul_accesses,
|
|
library_call = "linalg_matmul",
|
|
iterator_types = ["parallel", "parallel", "reduction"]
|
|
}
|
|
func.func @permute_generic(%A: memref<?x?xf32, offset: ?, strides: [?, 1]>,
|
|
%B: memref<?x?xf32, offset: ?, strides: [?, 1]>,
|
|
%C: memref<?x?xf32, offset: ?, strides: [?, 1]>) {
|
|
linalg.generic #generic_matmul_trait
|
|
ins(%A, %B : memref<?x?xf32, offset: ?, strides: [?, 1]>,
|
|
memref<?x?xf32, offset: ?, strides: [?, 1]>)
|
|
outs(%C : memref<?x?xf32, offset: ?, strides: [?, 1]>) {
|
|
^bb(%a: f32, %b: f32, %c: f32):
|
|
%d = arith.mulf %a, %b: f32
|
|
%e = arith.addf %c, %d: f32
|
|
linalg.yield %e: f32
|
|
}
|
|
return
|
|
}
|
|
// CHECK-LABEL: func @permute_generic
|
|
// CHECK: linalg.generic {
|
|
// CHECK-SAME: indexing_maps = [#[[$kn]], #[[$nm]], #[[$km]]],
|
|
// CHECK-SAME: iterator_types = ["parallel", "reduction", "parallel"],
|
|
// CHECK-SAME: library_call = "linalg_matmul"}
|
|
// CHECK: memref<?x?xf32, #[[$STRIDED_2D_u_1]]>,
|
|
// CHECK-SAME: memref<?x?xf32, #[[$STRIDED_2D_u_1]]>
|
|
// CHECK-SAME: memref<?x?xf32, #[[$STRIDED_2D_u_1]]>
|
|
|
|
func.func @matvec_perm(%A: memref<?x?xf32, offset: ?, strides: [?, 1]>,
|
|
%x: memref<?xf32, offset: ?, strides: [1]>,
|
|
%y: memref<?xf32, offset: ?, strides: [1]>) {
|
|
linalg.matvec {__internal_linalg_transform__ = "__with_perm__"}
|
|
ins(%A, %x: memref<?x?xf32, offset: ?, strides: [?, 1]>,
|
|
memref<?xf32, offset: ?, strides: [1]>)
|
|
outs(%y: memref<?xf32, offset: ?, strides: [1]>)
|
|
return
|
|
}
|
|
// CHECK-LABEL: func @matvec_perm
|
|
// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
|
|
// CHECK-DAG: %[[c5:.*]] = arith.constant 5 : index
|
|
// CHECK-DAG: %[[c6:.*]] = arith.constant 6 : index
|
|
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c6]]
|
|
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c5]]
|
|
// CHECK: linalg.matvec
|
|
// CHECK: ins({{.*}}: memref<?x?xf32, #[[$STRIDED_2D_u_1]]>, memref<?xf32, #[[$STRIDED_1D]]>)
|
|
// CHECK: outs({{.*}}: memref<?xf32, #[[$STRIDED_1D]]>)
|
|
|
|
func.func @matmul_perm(%A: memref<?x?xf32, offset: ?, strides: [?, 1]>,
|
|
%B: memref<?x?xf32, offset: ?, strides: [?, 1]>,
|
|
%C: memref<?x?xf32, offset: ?, strides: [?, 1]>) {
|
|
linalg.matmul {__internal_linalg_transform__ = "__with_perm__"}
|
|
ins(%A, %B: memref<?x?xf32, offset: ?, strides: [?, 1]>,
|
|
memref<?x?xf32, offset: ?, strides: [?, 1]>)
|
|
outs(%C : memref<?x?xf32, offset: ?, strides: [?, 1]>)
|
|
return
|
|
}
|
|
// CHECK-LABEL: func @matmul_perm
|
|
// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
|
|
// CHECK-DAG: %[[c20:.*]] = arith.constant 20 : index
|
|
// CHECK-DAG: %[[c30:.*]] = arith.constant 30 : index
|
|
// CHECK-DAG: %[[c40:.*]] = arith.constant 40 : index
|
|
// CHECK-DAG: %[[c200:.*]] = arith.constant 200 : index
|
|
// CHECK-DAG: %[[c300:.*]] = arith.constant 300 : index
|
|
// CHECK-DAG: %[[c400:.*]] = arith.constant 400 : index
|
|
// CHECK-DAG: %[[c2000:.*]] = arith.constant 2000 : index
|
|
// CHECK-DAG: %[[c3000:.*]] = arith.constant 3000 : index
|
|
// CHECK-DAG: %[[c4000:.*]] = arith.constant 4000 : index
|
|
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c3000]] {
|
|
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c4000]] {
|
|
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c2000]] {
|
|
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c300]] {
|
|
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c200]] {
|
|
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c400]] {
|
|
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c20]] {
|
|
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c30]] {
|
|
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c40]] {
|
|
// CHECK: linalg.matmul
|
|
// CHECK: ins({{.*}}: memref<?x?xf32, #[[$STRIDED_2D_u_1]]>, memref<?x?xf32, #[[$STRIDED_2D_u_1]]>)
|
|
// CHECK: outs({{.*}}: memref<?x?xf32, #[[$STRIDED_2D_u_1]]>)
|
|
|
|
func.func @tile_permute_parallel_loop(%arg0: memref<?x?xf32>,
|
|
%arg1: memref<?x?xf32>,
|
|
%arg2: memref<?x?xf32>) {
|
|
linalg.matmul {__internal_linalg_transform__ = "par__with_perm__"}
|
|
ins(%arg0, %arg1: memref<?x?xf32>, memref<?x?xf32>)
|
|
outs(%arg2: memref<?x?xf32>)
|
|
return
|
|
}
|
|
// CHECK-LABEL: func @tile_permute_parallel_loop
|
|
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<?x?xf32>
|
|
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: memref<?x?xf32>
|
|
// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: memref<?x?xf32>
|
|
// CHECK-DAG: %[[C16:.*]] = arith.constant 16 : index
|
|
// CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
|
|
// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index
|
|
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
|
|
// CHECK-DAG: %[[D0:.*]] = memref.dim %[[ARG0]], %c0
|
|
// CHECK-DAG: %[[D1:.*]] = memref.dim %[[ARG0]], %c1
|
|
// CHECK-DAG: %[[D2:.*]] = memref.dim %[[ARG1]], %c1
|
|
// CHECK: scf.parallel (%{{.*}}) = (%[[C0]]) to (%[[D2]]) step (%[[C8]])
|
|
// CHECK: scf.for %{{.*}} = %[[C0]] to %[[D1]] step %[[C4]]
|
|
// CHECK: scf.parallel (%{{.*}}) = (%[[C0]]) to (%[[D0]]) step (%[[C16]])
|