1f77f01c65
This operation is a NavigationOp that simplifies the writing of transform IR. Since there is no way of refering to an interface by name, the current implementation uses an EnumAttr and depends on the interfaces it supports. In the future, it would be worthwhile to remove this dependence and generalize. Differential Revision: https://reviews.llvm.org/D130267
328 lines
13 KiB
MLIR
328 lines
13 KiB
MLIR
// RUN: mlir-opt %s --test-transform-dialect-interpreter --split-input-file -verify-diagnostics | FileCheck %s
|
|
// RUN: mlir-opt %s --test-transform-dialect-interpreter --canonicalize --split-input-file -verify-diagnostics | FileCheck %s --check-prefix=CANON
|
|
|
|
transform.with_pdl_patterns {
|
|
^bb0(%arg0: !pdl.operation):
|
|
transform.sequence %arg0 {
|
|
^bb1(%arg1: !pdl.operation):
|
|
%0 = transform.structured.match ops{["linalg.generic"]} in %arg1
|
|
%1:2 = transform.structured.split %0 after 42 { dimension = 0 }
|
|
}
|
|
}
|
|
|
|
func.func private @elem(%arg0: f32, %arg1: index, %arg2: index) -> f32
|
|
|
|
// CHECK: #[[$ADD_42_MAP:.+]] = affine_map<(d0) -> (d0 + 42)>
|
|
// CHECK: #[[$ADD_10_MAP:.+]] = affine_map<(d0) -> (d0 + 10)>
|
|
|
|
// CHECK-LABEL: @one_d_static
|
|
// CHECK-SAME: %[[IN:.+]]: tensor<100xf32>, %[[OUT:.+]]: tensor<100xf32>
|
|
func.func @one_d_static(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> {
|
|
// CHECK: %[[IN_SLICE_LOW:.+]] = tensor.extract_slice %[[IN]][0] [42] [1] : tensor<100xf32> to tensor<42xf32>
|
|
// CHECK: %[[OUT_SLICE_LOW:.+]] = tensor.extract_slice %[[OUT]][0] [42] [1] : tensor<100xf32> to tensor<42xf32>
|
|
// CHECK: %[[RES_SLICE_LOW:.+]] = linalg.generic
|
|
// CHECK: ins(%[[IN_SLICE_LOW]]
|
|
// CHECK: outs(%[[OUT_SLICE_LOW]]
|
|
// CHECK: linalg.index 0
|
|
// CHECK: func.call @elem
|
|
// CHECK: %[[RES_PARTIAL:.+]] = tensor.insert_slice %[[RES_SLICE_LOW]] into %[[OUT]][0] [42] [1]
|
|
//
|
|
// CHECK: %[[IN_SLICE_HIGH:.+]] = tensor.extract_slice %[[IN]][42] [58] [1] : tensor<100xf32> to tensor<58xf32>
|
|
// CHECK: %[[OUT_SLICE_HIGH:.+]] = tensor.extract_slice %[[RES_PARTIAL]][42] [58] [1] : tensor<100xf32> to tensor<58xf32>
|
|
// CHECK: %[[RES_SLICE_HIGH:.+]] = linalg.generic
|
|
// CHECK: ins(%[[IN_SLICE_HIGH]]
|
|
// CHECK: outs(%[[OUT_SLICE_HIGH]]
|
|
// CHECK: %[[IDX:.+]] = linalg.index 0
|
|
// CHECK: affine.apply #[[$ADD_42_MAP]](%[[IDX]])
|
|
// CHECK: func.call @elem
|
|
// CHECK: %[[RES:.+]] = tensor.insert_slice %[[RES_SLICE_HIGH]] into %[[RES_PARTIAL]][42] [58] [1]
|
|
%0 = linalg.generic {
|
|
indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>],
|
|
iterator_types = ["parallel"]
|
|
}
|
|
ins(%arg0: tensor<100xf32>) outs(%arg1: tensor<100xf32>) {
|
|
^bb0(%0: f32, %1: f32):
|
|
%i = linalg.index 0 : index
|
|
%call_res = func.call @elem(%0, %i, %i) : (f32, index, index) -> f32
|
|
linalg.yield %call_res : f32
|
|
} -> tensor<100xf32>
|
|
|
|
// CHECK: return %[[RES]]
|
|
return %0 : tensor<100xf32>
|
|
}
|
|
|
|
// CHECK-LABEL: @one_d_static_overflow
|
|
// CHECK-SAME: %[[IN:.+]]: tensor<10xf32>, %[[OUT:.+]]: tensor<10xf32>
|
|
// CANON-LABEL: @one_d_static_overflow
|
|
// CANON-SAME: %[[IN:.+]]: tensor<10xf32>, %[[OUT:.+]]: tensor<10xf32>
|
|
func.func @one_d_static_overflow(%arg0: tensor<10xf32>, %arg1: tensor<10xf32>) -> tensor<10xf32> {
|
|
// CHECK: %[[IN_SLICE_LOW:.+]] = tensor.extract_slice %[[IN]][0] [10] [1] : tensor<10xf32> to tensor<10xf32>
|
|
// CHECK: %[[OUT_SLICE_LOW:.+]] = tensor.extract_slice %[[OUT]][0] [10] [1] : tensor<10xf32> to tensor<10xf32>
|
|
// CHECK: %[[RES_SLICE_LOW:.+]] = linalg.generic
|
|
// CHECK: ins(%[[IN_SLICE_LOW]]
|
|
// CHECK: outs(%[[OUT_SLICE_LOW]]
|
|
// CHECK: linalg.index 0
|
|
// CHECK: func.call @elem
|
|
// CHECK: %[[RES_PARTIAL:.+]] = tensor.insert_slice %[[RES_SLICE_LOW]] into %[[OUT]][0] [10] [1]
|
|
//
|
|
// Due to overflow, the first part of the split computes everything and the
|
|
// insert/extract slices are folded away by the canonicalizer.
|
|
// CANON: %[[RES_PARTIAL:.+]] = linalg.generic
|
|
// CANON: ins(%[[IN]]
|
|
// CANON: outs(%[[OUT]]
|
|
// CANON: linalg.index 0
|
|
// CANON: func.call @elem
|
|
// The second part operates on zero-sized slices that are not currently
|
|
// folded away.
|
|
//
|
|
// CHECK: %[[IN_SLICE_HIGH:.+]] = tensor.extract_slice %[[IN]][10] [0] [1] : tensor<10xf32> to tensor<0xf32>
|
|
// CHECK: %[[OUT_SLICE_HIGH:.+]] = tensor.extract_slice %[[RES_PARTIAL]][10] [0] [1] : tensor<10xf32> to tensor<0xf32>
|
|
// CHECK: %[[RES_SLICE_HIGH:.+]] = linalg.generic
|
|
// CHECK: ins(%[[IN_SLICE_HIGH]]
|
|
// CHECK: outs(%[[OUT_SLICE_HIGH]]
|
|
// CHECK: %[[IDX:.+]] = linalg.index 0
|
|
// CHECK: affine.apply #[[$ADD_10_MAP]](%[[IDX]])
|
|
// CHECK: func.call @elem
|
|
// CHECK: %[[RES:.+]] = tensor.insert_slice %[[RES_SLICE_HIGH]] into %[[RES_PARTIAL]][10] [0] [1]
|
|
%0 = linalg.generic {
|
|
indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>],
|
|
iterator_types = ["parallel"]
|
|
}
|
|
ins(%arg0: tensor<10xf32>) outs(%arg1: tensor<10xf32>) {
|
|
^bb0(%0: f32, %1: f32):
|
|
%i = linalg.index 0 : index
|
|
%call_res = func.call @elem(%0, %i, %i) : (f32, index, index) -> f32
|
|
linalg.yield %call_res : f32
|
|
} -> tensor<10xf32>
|
|
return %0 : tensor<10xf32>
|
|
}
|
|
|
|
// -----
|
|
|
|
transform.with_pdl_patterns {
|
|
^bb0(%arg0: !pdl.operation):
|
|
transform.sequence %arg0 {
|
|
^bb1(%arg1: !pdl.operation):
|
|
%0 = transform.structured.match ops{["linalg.generic"]} in %arg1
|
|
%1 = transform.structured.match ops{["func.call"]} in %arg1
|
|
transform.structured.split %0 after %1 { dimension = 0 }
|
|
}
|
|
}
|
|
|
|
func.func private @get_size() -> index
|
|
|
|
// CHECK: #[[$MAP_MIN_100:.+]] = affine_map<()[s0] -> (s0, 100)>
|
|
// CHECK: #[[$MAP_S_MINUS_100:.+]] = affine_map<()[s0] -> (-s0 + 100)>
|
|
|
|
// CHECK-LABEL: @dynamic
|
|
func.func @dynamic(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> {
|
|
// CHECK: %[[SPLIT:.+]] = call @get_size
|
|
// CHECK: %[[SPLIT_LOW:.+]] = affine.min #[[$MAP_MIN_100]]()[%[[SPLIT]]
|
|
// CHECK: %[[IN_SLICE_LOW:.+]] = tensor.extract_slice %[[IN:.+]][0] [%[[SPLIT_LOW]]] [1] : tensor<100xf32> to tensor<?xf32>
|
|
// CHECK: %[[OUT_SLICE_LOW:.+]] = tensor.extract_slice %[[OUT:.+]][0] [%[[SPLIT_LOW]]] [1] : tensor<100xf32> to tensor<?xf32>
|
|
// CHECK: %[[RES_SLICE_LOW:.+]] = linalg.generic
|
|
// CHECK: ins(%[[IN_SLICE_LOW]]
|
|
// CHECK: outs(%[[OUT_SLICE_LOW]]
|
|
// CHECK: %[[PARTIAL:.+]] = tensor.insert_slice %[[RES_SLICE_LOW]] into %[[OUT]][0] [%[[SPLIT_LOW]]] [1]
|
|
//
|
|
// CHECK: %[[SPLIT_HIGH_1:.+]] = affine.apply #[[$MAP_S_MINUS_100]]()[%[[SPLIT_LOW]]]
|
|
// CHECK: %[[SPLIT_HIGH_2:.+]] = affine.apply #[[$MAP_S_MINUS_100]]()[%[[SPLIT_LOW]]]
|
|
// CHECK: %[[IN_SLICE_HIGH:.+]] = tensor.extract_slice %[[IN:.+]][%[[SPLIT_LOW]]] [%[[SPLIT_HIGH_2]]] [1] : tensor<100xf32> to tensor<?xf32>
|
|
// CHECK: %[[SPLIT_HIGH_3:.+]] = affine.apply #[[$MAP_S_MINUS_100]]()[%[[SPLIT_LOW]]]
|
|
// CHECK: %[[OUT_SLICE_HIGH:.+]] = tensor.extract_slice %[[PARTIAL:.+]][%[[SPLIT_LOW]]] [%[[SPLIT_HIGH_3]]] [1] : tensor<100xf32> to tensor<?xf32>
|
|
// CHECK: %[[RES_SLICE_HIGH:.+]] = linalg.generic
|
|
// CHECK: ins(%[[IN_SLICE_HIGH]]
|
|
// CHECK: outs(%[[OUT_SLICE_HIGH]]
|
|
// CHECK: tensor.insert_slice %[[RES_SLICE_HIGH]] into %[[PARTIAL]][%[[SPLIT_LOW]]] [%[[SPLIT_HIGH_3]]] [1]
|
|
%0 = func.call @get_size() : () -> index
|
|
%1 = linalg.generic {
|
|
indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>],
|
|
iterator_types = ["parallel"]
|
|
}
|
|
ins(%arg0: tensor<100xf32>) outs(%arg1: tensor<100xf32>) {
|
|
^bb0(%3: f32, %4: f32):
|
|
%5 = arith.addf %3, %4 : f32
|
|
linalg.yield %5 : f32
|
|
} -> tensor<100xf32>
|
|
return %1 : tensor<100xf32>
|
|
}
|
|
|
|
// -----
|
|
|
|
transform.with_pdl_patterns {
|
|
^bb0(%arg0: !pdl.operation):
|
|
transform.sequence %arg0 {
|
|
^bb1(%arg1: !pdl.operation):
|
|
%0 = transform.structured.match ops{["linalg.generic"]} in %arg1
|
|
%1:2 = transform.structured.split %0 after 4 { dimension = 0}
|
|
%2:2 = transform.structured.split %1#1 after 16 { dimension = 1 }
|
|
}
|
|
}
|
|
|
|
func.func private @elem(%arg0: f32, %arg1: index, %arg2: index) -> f32
|
|
|
|
// CHECK-LABEL: @two_d
|
|
func.func @two_d(%arg0: tensor<10x34xf32>,
|
|
%arg1: tensor<10x34xf32>) -> tensor<10x34xf32> {
|
|
// Check the overall structure: split along the dimension 0, and then split
|
|
// the second half only along the dimension 1.
|
|
// CHECK: %[[IN_1:.+]] = tensor.extract_slice %[[IN:.+]][0, 0]
|
|
// CHECK: %[[OUT_1:.+]] = tensor.extract_slice %[[OUT:.+]][0, 0]
|
|
// CHECK: %[[RES_1:.+]] = linalg.generic
|
|
// CHECK-SAME: ins(%[[IN_1]] : tensor<4x34xf32>)
|
|
// CHECK-SAME: outs(%[[OUT_1]] : tensor<4x34xf32>)
|
|
// CHECK: %[[PARTIAL_1:.+]] = tensor.insert_slice %[[RES_1]] into %[[OUT]]
|
|
//
|
|
// CHECK: %[[IN_2:.+]] = tensor.extract_slice %[[IN]]
|
|
// CHECK: %[[OUT_2:.+]] = tensor.extract_slice %[[PARTIAL_1]]
|
|
// CHECK: %[[IN_21:.+]] = tensor.extract_slice %[[IN_2]]
|
|
// CHECK: %[[OUT_21:.+]] = tensor.extract_slice %[[OUT_2]]
|
|
// CHECK: %[[RES_21:.+]] = linalg.generic
|
|
// CHECK-SAME: ins(%[[IN_21]] : tensor<6x16xf32>)
|
|
// CHECK-SAME: outs(%[[OUT_21]] : tensor<6x16xf32>)
|
|
// CHECK: %[[PARTIAL_21:.+]] = tensor.insert_slice %[[RES_21]] into %[[OUT_2]]
|
|
//
|
|
// CHECK: %[[IN_22:.+]] = tensor.extract_slice %[[IN_2]]
|
|
// CHECK: %[[OUT_22:.+]] = tensor.extract_slice %[[PARTIAL_21]]
|
|
// CHECK: %[[RES_22:.+]] = linalg.generic
|
|
// CHECK-SAME: ins(%[[IN_22]] : tensor<6x18xf32>)
|
|
// CHECK-SAME: outs(%[[OUT_22]] : tensor<6x18xf32>)
|
|
// CHECK: %[[PARTIAL_22:.+]] = tensor.insert_slice %[[RES_22]] into %[[PARTIAL_21]]
|
|
// CHECK: %[[PARTIAL_2:.+]] = tensor.insert_slice %[[PARTIAL_22]] into %[[PARTIAL_1]]
|
|
%0 = linalg.generic {
|
|
indexing_maps = [affine_map<(i, j) -> (i, j)>,
|
|
affine_map<(i, j) -> (i, j)>],
|
|
iterator_types = ["parallel", "parallel"]
|
|
}
|
|
ins(%arg0: tensor<10x34xf32>)
|
|
outs(%arg1: tensor<10x34xf32>) {
|
|
^bb0(%0: f32, %1: f32):
|
|
%i = linalg.index 0 : index
|
|
%j = linalg.index 1 : index
|
|
%call_res = func.call @elem(%0, %i, %j) : (f32, index, index) -> f32
|
|
linalg.yield %call_res : f32
|
|
} -> tensor<10x34xf32>
|
|
return %0 : tensor<10x34xf32>
|
|
}
|
|
|
|
// -----
|
|
|
|
transform.sequence {
|
|
^bb1(%arg1: !pdl.operation):
|
|
// expected-error @below {{expects either a dynamic or a static split point to be provided}}
|
|
%0:2 = "transform.structured.split"(%arg1) { dimension = 1, static_split_point = -1 } : (!pdl.operation) -> (!pdl.operation, !pdl.operation)
|
|
}
|
|
|
|
// -----
|
|
|
|
transform.with_pdl_patterns {
|
|
^bb0(%arg0: !pdl.operation):
|
|
transform.sequence %arg0 {
|
|
^bb1(%arg1: !pdl.operation):
|
|
%0 = transform.structured.match ops{["linalg.generic"]} in %arg1
|
|
%1 = transform.structured.match ops{["func.call"]} in %arg1
|
|
// expected-error @below {{expected dynamic split point handle to point to a single-result index-typed op}}
|
|
transform.structured.split %0 after %1 { dimension = 0 }
|
|
}
|
|
}
|
|
|
|
func.func private @get_size() -> i64
|
|
|
|
func.func @dynamic(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> {
|
|
// expected-note @below {{dynamic split point}}
|
|
%0 = func.call @get_size() : () -> i64
|
|
%1 = linalg.generic {
|
|
indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>],
|
|
iterator_types = ["parallel"]
|
|
}
|
|
ins(%arg0: tensor<100xf32>) outs(%arg1: tensor<100xf32>) {
|
|
^bb0(%3: f32, %4: f32):
|
|
linalg.yield %3 : f32
|
|
} -> tensor<100xf32>
|
|
return %1 : tensor<100xf32>
|
|
}
|
|
|
|
// -----
|
|
|
|
transform.with_pdl_patterns {
|
|
^bb0(%arg0: !pdl.operation):
|
|
transform.sequence %arg0 {
|
|
^bb1(%arg1: !pdl.operation):
|
|
%0 = transform.structured.match ops{["linalg.generic"]} in %arg1
|
|
%1 = transform.structured.match ops{["func.call"]} in %arg1
|
|
// expected-error @below {{expected the dynamic split point handle to point to as many operations (0) as the target handle (1)}}
|
|
transform.structured.split %0 after %1 { dimension = 0 }
|
|
}
|
|
}
|
|
|
|
func.func private @get_size() -> i64
|
|
|
|
func.func @dynamic(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> {
|
|
%1 = linalg.generic {
|
|
indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>],
|
|
iterator_types = ["parallel"]
|
|
}
|
|
ins(%arg0: tensor<100xf32>) outs(%arg1: tensor<100xf32>) {
|
|
^bb0(%3: f32, %4: f32):
|
|
linalg.yield %3 : f32
|
|
} -> tensor<100xf32>
|
|
return %1 : tensor<100xf32>
|
|
}
|
|
|
|
// -----
|
|
|
|
transform.with_pdl_patterns {
|
|
^bb0(%arg0: !pdl.operation):
|
|
pdl.pattern @func_return : benefit(1) {
|
|
%0 = pdl.operands
|
|
%1 = pdl.types
|
|
%2 = pdl.operation "func.return"(%0 : !pdl.range<value>) -> (%1 : !pdl.range<type>)
|
|
pdl.rewrite %2 with "transform.dialect"
|
|
}
|
|
|
|
transform.sequence %arg0 {
|
|
^bb1(%arg1: !pdl.operation):
|
|
%0 = transform.structured.match ops{["func.return"]} in %arg1
|
|
// expected-error @below {{only applies to structured ops}}
|
|
transform.structured.split %0 after 16 { dimension = 1 }
|
|
}
|
|
}
|
|
|
|
func.func @noop(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> {
|
|
// expected-note @below {{target op}}
|
|
return %arg0 : tensor<100xf32>
|
|
}
|
|
|
|
// -----
|
|
|
|
transform.with_pdl_patterns {
|
|
^bb0(%arg0: !pdl.operation):
|
|
pdl.pattern @linalg_generic : benefit(1) {
|
|
%0 = pdl.operands
|
|
%1 = pdl.types
|
|
%2 = pdl.operation "linalg.generic"(%0 : !pdl.range<value>) -> (%1 : !pdl.range<type>)
|
|
pdl.rewrite %2 with "transform.dialect"
|
|
}
|
|
|
|
transform.sequence %arg0 {
|
|
^bb1(%arg1: !pdl.operation):
|
|
%0 = transform.structured.match ops{["linalg.generic"]} in %arg1
|
|
// expected-error @below {{dimension 1 does not exist in target op}}
|
|
transform.structured.split %0 after 16 { dimension = 1 }
|
|
}
|
|
}
|
|
|
|
func.func @one_d_static(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> {
|
|
// expected-note @below {{target op}}
|
|
%0 = linalg.generic {
|
|
indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>],
|
|
iterator_types = ["parallel"]
|
|
}
|
|
ins(%arg0: tensor<100xf32>) outs(%arg1: tensor<100xf32>) {
|
|
^bb0(%0: f32, %1: f32):
|
|
linalg.yield %0 : f32
|
|
} -> tensor<100xf32>
|
|
return %0 : tensor<100xf32>
|
|
}
|
|
|