2798b72ae7
Introduce a new extension for simple print-debugging of the transform dialect scripts. The initial version of this extension consists of two ops that are printing the payload objects associated with transform dialect values. Similar ops were already available in the test extenion and several downstream projects, and were extensively used for testing.
82 lines
3.6 KiB
MLIR
82 lines
3.6 KiB
MLIR
// RUN: mlir-opt %s --transform-interpreter --split-input-file \
|
|
// RUN: --verify-diagnostics | FileCheck %s
|
|
|
|
// UNSUPPORTED: target=aarch64-pc-windows-msvc
|
|
|
|
// CHECK-LABEL: func @test_loop_invariant_subset_hoisting(
|
|
// CHECK-SAME: %[[arg:.*]]: tensor<?xf32>
|
|
func.func @test_loop_invariant_subset_hoisting(%arg: tensor<?xf32>) -> tensor<?xf32> {
|
|
%lb = "test.foo"() : () -> (index)
|
|
%ub = "test.foo"() : () -> (index)
|
|
%step = "test.foo"() : () -> (index)
|
|
// CHECK: %[[extract:.*]] = tensor.extract_slice %[[arg]]
|
|
// CHECK: %[[for:.*]]:2 = scf.for {{.*}} iter_args(%[[t:.*]] = %[[arg]], %[[hoisted:.*]] = %[[extract]])
|
|
// expected-remark @below{{new loop op}}
|
|
%0 = scf.for %iv = %lb to %ub step %step iter_args(%t = %arg) -> (tensor<?xf32>) {
|
|
%1 = tensor.extract_slice %t[0][5][1] : tensor<?xf32> to tensor<5xf32>
|
|
// CHECK: %[[foo:.*]] = "test.foo"(%[[hoisted]])
|
|
%2 = "test.foo"(%1) : (tensor<5xf32>) -> (tensor<5xf32>)
|
|
// Obfuscate the IR by inserting at offset %sub instead of 0; both of them
|
|
// have the same value.
|
|
%3 = tensor.insert_slice %2 into %t[0][5][1] : tensor<5xf32> into tensor<?xf32>
|
|
// CHECK: scf.yield %[[t]], %[[foo]]
|
|
scf.yield %3 : tensor<?xf32>
|
|
}
|
|
// CHECK: %[[insert:.*]] = tensor.insert_slice %[[for]]#1 into %[[for]]#0
|
|
// CHECK: return %[[insert]]
|
|
return %0 : tensor<?xf32>
|
|
}
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
|
|
%0 = transform.structured.match ops{["scf.for"]} in %arg0 : (!transform.any_op) -> !transform.any_op
|
|
%1 = transform.structured.match ops{["tensor.extract_slice"]} in %arg0 : (!transform.any_op) -> !transform.any_op
|
|
%2 = transform.structured.match ops{["tensor.insert_slice"]} in %arg0 : (!transform.any_op) -> !transform.any_op
|
|
|
|
transform.loop.hoist_loop_invariant_subsets %0 : !transform.any_op
|
|
// Make sure that the handles are still valid (and were updated in case of
|
|
// the loop).
|
|
|
|
%p = transform.num_associations %0 : (!transform.any_op) -> !transform.param<i64>
|
|
// expected-remark @below{{1}}
|
|
transform.debug.emit_param_as_remark %p : !transform.param<i64>
|
|
transform.debug.emit_remark_at %0, "new loop op" : !transform.any_op
|
|
%p2 = transform.num_associations %1 : (!transform.any_op) -> !transform.param<i64>
|
|
// expected-remark @below{{1}}
|
|
transform.debug.emit_param_as_remark %p2 : !transform.param<i64>
|
|
%p3 = transform.num_associations %2 : (!transform.any_op) -> !transform.param<i64>
|
|
// expected-remark @below{{1}}
|
|
transform.debug.emit_param_as_remark %p3 : !transform.param<i64>
|
|
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
// Checks that transform ops from LoopExtensionOps and SCFTransformOps can be
|
|
// used together.
|
|
|
|
// CHECK-LABEL: func @test_mixed_loop_extension_scf_transform(
|
|
func.func @test_mixed_loop_extension_scf_transform(%arg: tensor<?xf32>) -> tensor<?xf32> {
|
|
%lb = "test.foo"() : () -> (index)
|
|
%ub = "test.foo"() : () -> (index)
|
|
%step = "test.foo"() : () -> (index)
|
|
// CHECK: scf.for
|
|
// CHECK: scf.for
|
|
%0 = scf.for %iv = %lb to %ub step %step iter_args(%t = %arg) -> (tensor<?xf32>) {
|
|
%1 = "test.foo"(%t) : (tensor<?xf32>) -> (tensor<?xf32>)
|
|
scf.yield %1 : tensor<?xf32>
|
|
}
|
|
return %0 : tensor<?xf32>
|
|
}
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
|
|
%0 = transform.structured.match ops{["scf.for"]} in %arg0 : (!transform.any_op) -> !transform.any_op
|
|
transform.loop.hoist_loop_invariant_subsets %0 : !transform.any_op
|
|
transform.loop.unroll %0 { factor = 4 } : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|