187 lines · plain
1// RUN: mlir-opt %s --transform-interpreter --split-input-file \2// RUN: --verify-diagnostics | FileCheck %s3 4// CHECK-LABEL: @schedule_with_nondet_knobs5module attributes {transform.with_named_sequence} {6 transform.named_sequence @schedule_with_nondet_knobs(%arg0: !transform.any_op {transform.readonly}) {7 // CHECK: %[[HEADS_OR_TAILS:.*]] = transform.tune.knob<"coin"> options = [true, false] -> !transform.any_param8 %heads_or_tails = transform.tune.knob<"coin"> options = [true, false] -> !transform.any_param9 // CHECK: transform.tune.knob<"animal"> options = ["cat", "dog", unit] -> !transform.any_param10 %chosen_category = transform.tune.knob<"animal"> options = ["cat", "dog", unit] -> !transform.any_param11 // CHECK: transform.tune.knob<"tile_size"> options = [2, 4, 8, 16, 24, 32] -> !transform.any_param12 %chosen_tile_size = transform.tune.knob<"tile_size"> options = [2, 4, 8, 16, 24, 32] -> !transform.any_param13 // CHECK: transform.tune.knob<"magic_value"> options = [2.000000e+00 : f32, 2.250000e+00 : f32, 2.500000e+00 : f32, 2.750000e+00 : f32, 3.000000e+00 : f32] -> !transform.any_param14 %chosen_constant = transform.tune.knob<"magic_value"> options = [2.0 : f32, 2.25 : f32, 2.5 : f32, 2.75 : f32, 3.0 : f32] -> !transform.any_param15 // CHECK: transform.debug.emit_param_as_remark %[[HEADS_OR_TAILS]]16 transform.debug.emit_param_as_remark %heads_or_tails : !transform.any_param17 transform.yield18 }19 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {20 // Dummy sequence to appease -transform-interpreter invocation21 transform.yield22 }23}24 25// -----26 27// Schedule where non-determinism on knobs has been resolved by selecting a valid option.28 29// CHECK-LABEL: payload_for_schedule_with_selected_knobs30func.func private @payload_for_schedule_with_selected_knobs()31 32module attributes {transform.with_named_sequence} {33 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {34 // CHECK: %[[HEADS_OR_TAILS:.*]] = transform.tune.knob<"coin"> = true from options = [true, false] -> !transform.any_param35 %heads_or_tails = transform.tune.knob<"coin"> = true from options = [true, false] -> !transform.any_param36 // expected-remark@below {{true}}37 transform.debug.emit_param_as_remark %heads_or_tails : !transform.any_param38 39 // CHECK: transform.tune.knob<"animal"> = "dog" from options = ["cat", "dog", unit] -> !transform.any_param40 %chosen_category = transform.tune.knob<"animal"> = "dog" from options = ["cat", "dog", unit] -> !transform.any_param41 // CHECK: transform.tune.knob<"tile_size"> = 8 : i64 from options = [2, 4, 8, 16, 24, 32] -> !transform.any_param42 %chosen_tile_size = transform.tune.knob<"tile_size"> = 8 from options = [2, 4, 8, 16, 24, 32] -> !transform.any_param43 // CHECK: transform.tune.knob<"magic_value"> = 2.500000e+00 : f32 from options = [2.000000e+00 : f32, 2.250000e+00 : f32, 2.500000e+00 : f32, 2.750000e+00 : f32, 3.000000e+00 : f32] -> !transform.any_param44 %chosen_constant = transform.tune.knob<"magic_value"> = 2.5 : f32 from options = [2.0 : f32, 2.25 : f32, 2.5 : f32, 2.75 : f32, 3.0 : f32] -> !transform.any_param45 transform.yield46 }47}48 49// -----50 51// CHECK: #[[AFFINE_SET:.*]] = affine_set<(d0) : (d0 - 2 >= 0)>52// CHECK: payload_for_schedule_where_selected_knob_being_a_member_of_options_is_unverified53func.func private @payload_for_schedule_where_selected_knob_being_a_member_of_options_is_unverified()54 55module attributes {transform.with_named_sequence} {56 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {57 // CHECK: transform.tune.knob<"bounded"> = 4242 : i64 from options = #[[AFFINE_SET]] -> !transform.any_param58 %value_in_half_range = transform.tune.knob<"bounded"> = 4242 from options = affine_set<(d0) : (d0 - 2 >= 0)> -> !transform.any_param59 transform.yield60 }61}62 63 64// -----65 66// CHECK-LABEL: schedule_with_two_independent_choices_already_made67func.func @schedule_with_two_independent_choices_already_made(68 %arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>, %arg2: tensor<128x128xf32>)69 -> tensor<128x128xf32> {70// CHECK-NOT: scf.forall71// CHECK: scf.for72// CHECK-NOT: scf.for73// CHECK: scf.forall74// CHECK-NOT: scf.for75// CHECK: tensor.extract_slice76// CHECK: tensor.extract_slice77// CHECK: tensor.extract_slice78// CHECK: linalg.matmul79// CHECK: scf.forall.in_parallel80// CHECK: tensor.parallel_insert_slice81// CHECK: tensor.insert_slice82// CHECK: scf.yield83 %0 = linalg.matmul ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)84 outs(%arg2: tensor<128x128xf32>) -> tensor<128x128xf32>85 return %0 : tensor<128x128xf32>86}87 88module attributes {transform.with_named_sequence} {89 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {90 %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg0 : (!transform.any_op) -> !transform.any_op91 92 %tiled_matmul = transform.tune.alternatives<"outer_par_or_seq_tiling"> selected_region = 0 -> !transform.any_op93 { // First alternative/region, with index = 094 %contained_matmul, %loop = transform.structured.tile_using_for %matmul tile_sizes [8] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)95 transform.yield %contained_matmul : !transform.any_op96 }, { // Second alternative/region, with index = 197 %contained_matmul, %loop = transform.structured.tile_using_forall %matmul tile_sizes [8] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)98 transform.yield %contained_matmul : !transform.any_op99 }100 101 transform.tune.alternatives<"inner_par_or_seq_tiling"> selected_region = 1 -> !transform.any_op {102 %contained_matmul, %loop = transform.structured.tile_using_for %tiled_matmul tile_sizes [0, 16] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)103 transform.yield %contained_matmul : !transform.any_op104 }, {105 %contained_matmul, %loop = transform.structured.tile_using_forall %tiled_matmul tile_sizes [0, 16] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)106 transform.yield %contained_matmul : !transform.any_op107 }108 109 transform.yield110 }111}112 113// -----114 115// CHECK-LABEL: subschedule_with_choice_resolved_in_main_schedule116func.func @subschedule_with_choice_resolved_in_main_schedule(117 %arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>, %arg2: tensor<128x128xf32>)118 -> tensor<128x128xf32> {119// CHECK-NOT: scf.for120// CHECK: scf.forall121// CHECK-NOT: scf.forall122// CHECK: scf.for123// CHECK-NOT: scf.forall124// CHECK: tensor.extract_slice125// CHECK: tensor.extract_slice126// CHECK: tensor.extract_slice127// CHECK: linalg.matmul128// CHECK: tensor.insert_slice129// CHECK: scf.yield130// CHECK: scf.forall.in_parallel131// CHECK: tensor.parallel_insert_slice132 %0 = linalg.matmul ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)133 outs(%arg2: tensor<128x128xf32>) -> tensor<128x128xf32>134 return %0 : tensor<128x128xf32>135}136 137module attributes {transform.with_named_sequence} {138 transform.named_sequence @subschedule_with_embedded_choice(%matmul: !transform.any_op {transform.readonly},139 %par_or_seq: !transform.param<i64> {transform.readonly},140 %tile_size: !transform.param<i64> {transform.readonly}) -> !transform.any_op {141 %tiled_matmul = transform.tune.alternatives<"par_or_seq_tiling"> selected_region = %par_or_seq : !transform.param<i64> -> !transform.any_op {142 %contained_matmul, %loop = transform.structured.tile_using_for %matmul tile_sizes [%tile_size] : (!transform.any_op, !transform.param<i64>) -> (!transform.any_op, !transform.any_op)143 transform.yield %contained_matmul : !transform.any_op144 }, {145 %contained_matmul, %loop = transform.structured.tile_using_forall %matmul tile_sizes [%tile_size] : (!transform.any_op, !transform.param<i64>) -> (!transform.any_op, !transform.any_op)146 transform.yield %contained_matmul : !transform.any_op147 }148 transform.yield %tiled_matmul : !transform.any_op149 }150 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {151 %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg0 : (!transform.any_op) -> !transform.any_op152 %outer_par = transform.param.constant 1 -> !transform.param<i64>153 %outer_tile_size = transform.param.constant 32 -> !transform.param<i64>154 %inner_seq = transform.tune.knob<"inner_par_or_seq"> = 0 from options = [0, 1] -> !transform.param<i64>155 %inner_tile_size = transform.param.constant 8 -> !transform.param<i64>156 %tiled_matmul = transform.include @subschedule_with_embedded_choice failures(propagate) (%matmul, %outer_par, %outer_tile_size) : (!transform.any_op, !transform.param<i64>, !transform.param<i64>) -> !transform.any_op157 %tiled_tiled_matmul = transform.include @subschedule_with_embedded_choice failures(propagate) (%tiled_matmul, %inner_seq, %inner_tile_size) : (!transform.any_op, !transform.param<i64>, !transform.param<i64>) -> !transform.any_op158 transform.yield159 }160}161 162// -----163 164// CHECK-LABEL: eeny_meeny_miny_moe165func.func private @eeny_meeny_miny_moe()166 167module attributes {transform.with_named_sequence} {168 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {169 %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg0 : (!transform.any_op) -> !transform.any_op170 171 %tiled_matmul = transform.tune.alternatives<"4way"> selected_region = 3 -> !transform.any_param172 { // First alternative/region, with index = 0173 %out = transform.param.constant "eeny" -> !transform.any_param174 transform.yield %out : !transform.any_param175 }, { // Second alternative/region, with index = 1176 %out = transform.param.constant "meeny" -> !transform.any_param177 transform.yield %out : !transform.any_param178 }, { // Third alternative/region, with index = 2179 %out = transform.param.constant "miny" -> !transform.any_param180 transform.yield %out : !transform.any_param181 }, { // Fourth alternative/region, with index = 3182 %out = transform.param.constant "moe" -> !transform.any_param183 transform.yield %out : !transform.any_param184 }185 transform.yield186 }187}