From 5f91c4e6949ab43da6779afddcde679f426560f0 Mon Sep 17 00:00:00 2001 From: Liam Semeria Date: Fri, 5 Jun 2026 16:20:37 +0200 Subject: [PATCH] mlir: added scheduled relu tests --- .../tensor_dialect/test_conv2d_relu_tensor.py | 329 ++++++++++++++++++ .../tensor_dialect/test_relu_tensor.py | 149 ++++++++ .../backends/test_conv2d_relu_mlir.py | 198 +++++++++++ tests/filecheck/backends/test_relu_mlir.py | 100 ++++++ .../gen_assembly/skylake_generic_conv2d.mlir | 51 +++ 5 files changed, 827 insertions(+) create mode 100644 tests/filecheck/backends/tensor_dialect/test_conv2d_relu_tensor.py create mode 100644 tests/filecheck/backends/tensor_dialect/test_relu_tensor.py create mode 100644 tests/filecheck/backends/test_conv2d_relu_mlir.py create mode 100644 tests/filecheck/backends/test_relu_mlir.py create mode 100644 tests/filecheck/mlir_loop/gen_assembly/skylake_generic_conv2d.mlir diff --git a/tests/filecheck/backends/tensor_dialect/test_conv2d_relu_tensor.py b/tests/filecheck/backends/tensor_dialect/test_conv2d_relu_tensor.py new file mode 100644 index 000000000..c90d60124 --- /dev/null +++ b/tests/filecheck/backends/tensor_dialect/test_conv2d_relu_tensor.py @@ -0,0 +1,329 @@ +# RUN: python %s 2>&1 | filecheck %s +# UNSUPPORTED: mlir-target=nvgpu + +import xtc.graphs.xtc.op as O +from xtc.backends.mlir import Backend + +N, H, W, F, R, S, C, SH, SW, dtype = 1, 8, 8, 16, 3, 3, 3, 1, 1, "float32" +a = O.tensor((N, H + R - 1, W + S - 1, C), dtype, name="I") +b = O.tensor((R, S, C, F), dtype, name="W") + +with O.graph(name="conv2d_nhwc_mini") as gb: + c = O.conv2d(a, b, stride=(SH, SW), name="O") + O.relu(c, name="relu") + +graph = gb.graph +print(graph) + +impl = Backend(graph, use_tensor_dialect=True) + +sch = impl.get_scheduler(default_node="relu") +sch.tile("i", {"i1": 16}) +sch.interchange([ "i", "i1"]) +sch.vectorize(["i1"]) +sched = sch.schedule() + +comp = impl.get_compiler( + shared_lib=True, + dump_file="conv2d_relu_mlir_tensor", + print_source_ir=True, + print_transformed_ir=True, + print_bufferization_ir=True, +) +module = comp.compile(sched) +executor = module.get_executor(validate=True) +res = executor.execute() +print(f"CODE: {res}") + +# CHECK: // -----// IR Dump Before transform //----- // +# CHECK-NEXT: #map = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 + d4, d2 + d5, d6)> +# CHECK-NEXT: #map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d4, d5, d6, d3)> +# CHECK-NEXT: #map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)> +# CHECK-NEXT: #map3 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> +# CHECK-NEXT: #map4 = affine_map<(d0, d1, d2, d3) -> ()> +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @conv2d_nhwc_mini(%arg0: tensor<1x10x10x3xf32> {llvm.noalias}, %arg1: tensor<3x3x3x16xf32> {llvm.noalias}, %arg2: memref<1x8x8x16xf32> {llvm.noalias}) { +# CHECK-NEXT: %0 = tensor.empty() : tensor<1x8x8x16xf32> +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: %1 = linalg.fill {__xtc_id_O_0_} ins(%cst : f32) outs(%0 : tensor<1x8x8x16xf32>) -> tensor<1x8x8x16xf32> +# CHECK-NEXT: %2 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction"]} ins(%arg0, %arg1 : tensor<1x10x10x3xf32>, tensor<3x3x3x16xf32>) outs(%1 : tensor<1x8x8x16xf32>) attrs = {__xtc_id_O_} { +# CHECK-NEXT: ^bb0(%in: f32, %in_1: f32, %out: f32): +# CHECK-NEXT: %5 = arith.mulf %in, %in_1 fastmath : f32 +# CHECK-NEXT: %6 = arith.addf %out, %5 fastmath : f32 +# CHECK-NEXT: linalg.yield %6 : f32 +# CHECK-NEXT: } -> tensor<1x8x8x16xf32> +# CHECK-NEXT: %3 = tensor.empty() : tensor<1x8x8x16xf32> +# CHECK-NEXT: %cst_0 = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: %4 = linalg.generic {indexing_maps = [#map3, #map4, #map3], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%2, %cst_0 : tensor<1x8x8x16xf32>, f32) outs(%3 : tensor<1x8x8x16xf32>) attrs = {__xtc_id_relu_} { +# CHECK-NEXT: ^bb0(%in: f32, %in_1: f32, %out: f32): +# CHECK-NEXT: %5 = arith.maximumf %in, %in_1 : f32 +# CHECK-NEXT: linalg.yield %5 : f32 +# CHECK-NEXT: } -> tensor<1x8x8x16xf32> +# CHECK-NEXT: bufferization.materialize_in_destination %4 in restrict writable %arg2 : (tensor<1x8x8x16xf32>, memref<1x8x8x16xf32>) -> () +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @_vecto(%arg0: !transform.any_op {transform.consumed}) { +# CHECK-NEXT: transform.structured.vectorize %arg0 : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @_post_bufferize(%arg0: !transform.any_op {transform.readonly}) { +# CHECK-NEXT: %0 = transform.structured.match attributes {sym_name = "conv2d_nhwc_mini"} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: transform.apply_patterns to %0 { +# CHECK-NEXT: transform.apply_patterns.vector.lower_outerproduct +# CHECK-NEXT: transform.apply_patterns.vector.lower_contraction +# CHECK-NEXT: } : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { +# CHECK-NEXT: %0 = transform.structured.match attributes {__xtc_id_O_0_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_linalg_op, %loops = transform.structured.tile_using_for %0 tile_sizes [1, 0, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops "./b" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_0, %loops_1 = transform.structured.tile_using_for %tiled_linalg_op tile_sizes [0, 1, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_1 "./h" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_2, %loops_3 = transform.structured.tile_using_for %tiled_linalg_op_0 tile_sizes [0, 0, 1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_3 "./w" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_4, %loops_5 = transform.structured.tile_using_for %tiled_linalg_op_2 tile_sizes [0, 0, 0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_5 "./f" : !transform.any_op +# CHECK-NEXT: %1 = transform.structured.match attributes {__xtc_id_O_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_6, %loops_7 = transform.structured.tile_using_for %1 tile_sizes [1, 0, 0, 0, 0, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_7 "./b" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_8, %loops_9 = transform.structured.tile_using_for %tiled_linalg_op_6 tile_sizes [0, 1, 0, 0, 0, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_9 "./h" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_10, %loops_11 = transform.structured.tile_using_for %tiled_linalg_op_8 tile_sizes [0, 0, 1, 0, 0, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_11 "./w" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_12, %loops_13 = transform.structured.tile_using_for %tiled_linalg_op_10 tile_sizes [0, 0, 0, 1, 0, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_13 "./f" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_14, %loops_15 = transform.structured.tile_using_for %tiled_linalg_op_12 tile_sizes [0, 0, 0, 0, 1, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_15 "./r" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_16, %loops_17 = transform.structured.tile_using_for %tiled_linalg_op_14 tile_sizes [0, 0, 0, 0, 0, 1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_17 "./s" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_18, %loops_19 = transform.structured.tile_using_for %tiled_linalg_op_16 tile_sizes [0, 0, 0, 0, 0, 0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_19 "./c" : !transform.any_op +# CHECK-NEXT: %2 = transform.structured.match attributes {__xtc_id_relu_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_20, %loops_21 = transform.structured.tile_using_for %2 tile_sizes [16] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_21 "./i" : !transform.any_op +# CHECK-NEXT: %3 = transform.get_parent_op %tiled_linalg_op_20 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: transform.apply_patterns to %3 { +# CHECK-NEXT: transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices +# CHECK-NEXT: } : !transform.any_op +# CHECK-NEXT: %4 = transform.structured.match interface{LinalgOp} in %3 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: transform.include @_vecto failures(suppress) (%4) : (!transform.any_op) -> () +# CHECK-NEXT: %5 = transform.get_parent_op %loops_21 {isolated_from_above} : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: transform.apply_patterns to %5 { +# CHECK-NEXT: transform.apply_patterns.vector.reduction_to_contract +# CHECK-NEXT: transform.apply_patterns.vector.transfer_permutation_patterns +# CHECK-NEXT: } : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: // -----// IR Dump After transform //----- // +# CHECK-NEXT: #map = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 + d4, d2 + d5, d6)> +# CHECK-NEXT: #map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d4, d5, d6, d3)> +# CHECK-NEXT: #map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)> +# CHECK-NEXT: #map3 = affine_map<(d0) -> (-d0 + 1, 16)> +# CHECK-NEXT: #map4 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> +# CHECK-NEXT: #map5 = affine_map<(d0, d1, d2, d3) -> ()> +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @conv2d_nhwc_mini(%arg0: tensor<1x10x10x3xf32> {llvm.noalias}, %arg1: tensor<3x3x3x16xf32> {llvm.noalias}, %arg2: memref<1x8x8x16xf32> {llvm.noalias}) { +# CHECK-NEXT: %c3 = arith.constant 3 : index +# CHECK-NEXT: %c16 = arith.constant 16 : index +# CHECK-NEXT: %c8 = arith.constant 8 : index +# CHECK-NEXT: %c1 = arith.constant 1 : index +# CHECK-NEXT: %c0 = arith.constant 0 : index +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: %0 = tensor.empty() : tensor<1x8x8x16xf32> +# CHECK-NEXT: %1 = scf.for %arg3 = %c0 to %c1 step %c1 iter_args(%arg4 = %0) -> (tensor<1x8x8x16xf32>) { +# CHECK-NEXT: %extracted_slice = tensor.extract_slice %arg4[%arg3, 0, 0, 0] [1, 8, 8, 16] [1, 1, 1, 1] : tensor<1x8x8x16xf32> to tensor<1x8x8x16xf32> +# CHECK-NEXT: %5 = scf.for %arg5 = %c0 to %c8 step %c1 iter_args(%arg6 = %extracted_slice) -> (tensor<1x8x8x16xf32>) { +# CHECK-NEXT: %extracted_slice_0 = tensor.extract_slice %arg6[0, %arg5, 0, 0] [1, 1, 8, 16] [1, 1, 1, 1] : tensor<1x8x8x16xf32> to tensor<1x1x8x16xf32> +# CHECK-NEXT: %6 = scf.for %arg7 = %c0 to %c8 step %c1 iter_args(%arg8 = %extracted_slice_0) -> (tensor<1x1x8x16xf32>) { +# CHECK-NEXT: %extracted_slice_2 = tensor.extract_slice %arg8[0, 0, %arg7, 0] [1, 1, 1, 16] [1, 1, 1, 1] : tensor<1x1x8x16xf32> to tensor<1x1x1x16xf32> +# CHECK-NEXT: %7 = scf.for %arg9 = %c0 to %c16 step %c1 iter_args(%arg10 = %extracted_slice_2) -> (tensor<1x1x1x16xf32>) { +# CHECK-NEXT: %extracted_slice_4 = tensor.extract_slice %arg10[0, 0, 0, %arg9] [1, 1, 1, 1] [1, 1, 1, 1] : tensor<1x1x1x16xf32> to tensor<1x1x1x1xf32> +# CHECK-NEXT: %8 = linalg.fill {__xtc_id_O_0_} ins(%cst : f32) outs(%extracted_slice_4 : tensor<1x1x1x1xf32>) -> tensor<1x1x1x1xf32> +# CHECK-NEXT: %inserted_slice_5 = tensor.insert_slice %8 into %arg10[0, 0, 0, %arg9] [1, 1, 1, 1] [1, 1, 1, 1] : tensor<1x1x1x1xf32> into tensor<1x1x1x16xf32> +# CHECK-NEXT: scf.yield %inserted_slice_5 : tensor<1x1x1x16xf32> +# CHECK-NEXT: } {"./f"} +# CHECK-NEXT: %inserted_slice_3 = tensor.insert_slice %7 into %arg8[0, 0, %arg7, 0] [1, 1, 1, 16] [1, 1, 1, 1] : tensor<1x1x1x16xf32> into tensor<1x1x8x16xf32> +# CHECK-NEXT: scf.yield %inserted_slice_3 : tensor<1x1x8x16xf32> +# CHECK-NEXT: } {"./w"} +# CHECK-NEXT: %inserted_slice_1 = tensor.insert_slice %6 into %arg6[0, %arg5, 0, 0] [1, 1, 8, 16] [1, 1, 1, 1] : tensor<1x1x8x16xf32> into tensor<1x8x8x16xf32> +# CHECK-NEXT: scf.yield %inserted_slice_1 : tensor<1x8x8x16xf32> +# CHECK-NEXT: } {"./h"} +# CHECK-NEXT: %inserted_slice = tensor.insert_slice %5 into %arg4[%arg3, 0, 0, 0] [1, 8, 8, 16] [1, 1, 1, 1] : tensor<1x8x8x16xf32> into tensor<1x8x8x16xf32> +# CHECK-NEXT: scf.yield %inserted_slice : tensor<1x8x8x16xf32> +# CHECK-NEXT: } {"./b"} +# CHECK-NEXT: %2 = scf.for %arg3 = %c0 to %c1 step %c1 iter_args(%arg4 = %1) -> (tensor<1x8x8x16xf32>) { +# CHECK-NEXT: %extracted_slice = tensor.extract_slice %arg0[%arg3, 0, 0, 0] [1, 10, 10, 3] [1, 1, 1, 1] : tensor<1x10x10x3xf32> to tensor<1x10x10x3xf32> +# CHECK-NEXT: %extracted_slice_0 = tensor.extract_slice %arg4[%arg3, 0, 0, 0] [1, 8, 8, 16] [1, 1, 1, 1] : tensor<1x8x8x16xf32> to tensor<1x8x8x16xf32> +# CHECK-NEXT: %5 = scf.for %arg5 = %c0 to %c8 step %c1 iter_args(%arg6 = %extracted_slice_0) -> (tensor<1x8x8x16xf32>) { +# CHECK-NEXT: %extracted_slice_1 = tensor.extract_slice %extracted_slice[0, %arg5, 0, 0] [1, 3, 10, 3] [1, 1, 1, 1] : tensor<1x10x10x3xf32> to tensor<1x3x10x3xf32> +# CHECK-NEXT: %extracted_slice_2 = tensor.extract_slice %arg6[0, %arg5, 0, 0] [1, 1, 8, 16] [1, 1, 1, 1] : tensor<1x8x8x16xf32> to tensor<1x1x8x16xf32> +# CHECK-NEXT: %6 = scf.for %arg7 = %c0 to %c8 step %c1 iter_args(%arg8 = %extracted_slice_2) -> (tensor<1x1x8x16xf32>) { +# CHECK-NEXT: %extracted_slice_4 = tensor.extract_slice %extracted_slice_1[0, 0, %arg7, 0] [1, 3, 3, 3] [1, 1, 1, 1] : tensor<1x3x10x3xf32> to tensor<1x3x3x3xf32> +# CHECK-NEXT: %extracted_slice_5 = tensor.extract_slice %arg8[0, 0, %arg7, 0] [1, 1, 1, 16] [1, 1, 1, 1] : tensor<1x1x8x16xf32> to tensor<1x1x1x16xf32> +# CHECK-NEXT: %7 = scf.for %arg9 = %c0 to %c16 step %c1 iter_args(%arg10 = %extracted_slice_5) -> (tensor<1x1x1x16xf32>) { +# CHECK-NEXT: %extracted_slice_7 = tensor.extract_slice %arg1[0, 0, 0, %arg9] [3, 3, 3, 1] [1, 1, 1, 1] : tensor<3x3x3x16xf32> to tensor<3x3x3x1xf32> +# CHECK-NEXT: %extracted_slice_8 = tensor.extract_slice %arg10[0, 0, 0, %arg9] [1, 1, 1, 1] [1, 1, 1, 1] : tensor<1x1x1x16xf32> to tensor<1x1x1x1xf32> +# CHECK-NEXT: %8 = scf.for %arg11 = %c0 to %c3 step %c1 iter_args(%arg12 = %extracted_slice_8) -> (tensor<1x1x1x1xf32>) { +# CHECK-NEXT: %extracted_slice_10 = tensor.extract_slice %extracted_slice_4[0, %arg11, 0, 0] [1, 1, 3, 3] [1, 1, 1, 1] : tensor<1x3x3x3xf32> to tensor<1x1x3x3xf32> +# CHECK-NEXT: %extracted_slice_11 = tensor.extract_slice %extracted_slice_7[%arg11, 0, 0, 0] [1, 3, 3, 1] [1, 1, 1, 1] : tensor<3x3x3x1xf32> to tensor<1x3x3x1xf32> +# CHECK-NEXT: %9 = scf.for %arg13 = %c0 to %c3 step %c1 iter_args(%arg14 = %arg12) -> (tensor<1x1x1x1xf32>) { +# CHECK-NEXT: %extracted_slice_12 = tensor.extract_slice %extracted_slice_10[0, 0, %arg13, 0] [1, 1, 1, 3] [1, 1, 1, 1] : tensor<1x1x3x3xf32> to tensor<1x1x1x3xf32> +# CHECK-NEXT: %extracted_slice_13 = tensor.extract_slice %extracted_slice_11[0, %arg13, 0, 0] [1, 1, 3, 1] [1, 1, 1, 1] : tensor<1x3x3x1xf32> to tensor<1x1x3x1xf32> +# CHECK-NEXT: %10 = scf.for %arg15 = %c0 to %c3 step %c1 iter_args(%arg16 = %arg14) -> (tensor<1x1x1x1xf32>) { +# CHECK-NEXT: %extracted_slice_14 = tensor.extract_slice %extracted_slice_12[0, 0, 0, %arg15] [1, 1, 1, 1] [1, 1, 1, 1] : tensor<1x1x1x3xf32> to tensor<1x1x1x1xf32> +# CHECK-NEXT: %extracted_slice_15 = tensor.extract_slice %extracted_slice_13[0, 0, %arg15, 0] [1, 1, 1, 1] [1, 1, 1, 1] : tensor<1x1x3x1xf32> to tensor<1x1x1x1xf32> +# CHECK-NEXT: %11 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction"]} ins(%extracted_slice_14, %extracted_slice_15 : tensor<1x1x1x1xf32>, tensor<1x1x1x1xf32>) outs(%arg16 : tensor<1x1x1x1xf32>) attrs = {__xtc_id_O_} { +# CHECK-NEXT: ^bb0(%in: f32, %in_16: f32, %out: f32): +# CHECK-NEXT: %12 = arith.mulf %in, %in_16 fastmath : f32 +# CHECK-NEXT: %13 = arith.addf %out, %12 fastmath : f32 +# CHECK-NEXT: linalg.yield %13 : f32 +# CHECK-NEXT: } -> tensor<1x1x1x1xf32> +# CHECK-NEXT: scf.yield %11 : tensor<1x1x1x1xf32> +# CHECK-NEXT: } {"./c"} +# CHECK-NEXT: scf.yield %10 : tensor<1x1x1x1xf32> +# CHECK-NEXT: } {"./s"} +# CHECK-NEXT: scf.yield %9 : tensor<1x1x1x1xf32> +# CHECK-NEXT: } {"./r"} +# CHECK-NEXT: %inserted_slice_9 = tensor.insert_slice %8 into %arg10[0, 0, 0, %arg9] [1, 1, 1, 1] [1, 1, 1, 1] : tensor<1x1x1x1xf32> into tensor<1x1x1x16xf32> +# CHECK-NEXT: scf.yield %inserted_slice_9 : tensor<1x1x1x16xf32> +# CHECK-NEXT: } {"./f"} +# CHECK-NEXT: %inserted_slice_6 = tensor.insert_slice %7 into %arg8[0, 0, %arg7, 0] [1, 1, 1, 16] [1, 1, 1, 1] : tensor<1x1x1x16xf32> into tensor<1x1x8x16xf32> +# CHECK-NEXT: scf.yield %inserted_slice_6 : tensor<1x1x8x16xf32> +# CHECK-NEXT: } {"./w"} +# CHECK-NEXT: %inserted_slice_3 = tensor.insert_slice %6 into %arg6[0, %arg5, 0, 0] [1, 1, 8, 16] [1, 1, 1, 1] : tensor<1x1x8x16xf32> into tensor<1x8x8x16xf32> +# CHECK-NEXT: scf.yield %inserted_slice_3 : tensor<1x8x8x16xf32> +# CHECK-NEXT: } {"./h"} +# CHECK-NEXT: %inserted_slice = tensor.insert_slice %5 into %arg4[%arg3, 0, 0, 0] [1, 8, 8, 16] [1, 1, 1, 1] : tensor<1x8x8x16xf32> into tensor<1x8x8x16xf32> +# CHECK-NEXT: scf.yield %inserted_slice : tensor<1x8x8x16xf32> +# CHECK-NEXT: } {"./b"} +# CHECK-NEXT: %3 = tensor.empty() : tensor<1x8x8x16xf32> +# CHECK-NEXT: %4 = scf.for %arg3 = %c0 to %c1 step %c16 iter_args(%arg4 = %3) -> (tensor<1x8x8x16xf32>) { +# CHECK-NEXT: %5 = affine.min #map3(%arg3) +# CHECK-NEXT: %extracted_slice = tensor.extract_slice %2[%arg3, 0, 0, 0] [%5, 8, 8, 16] [1, 1, 1, 1] : tensor<1x8x8x16xf32> to tensor +# CHECK-NEXT: %extracted_slice_0 = tensor.extract_slice %arg4[%arg3, 0, 0, 0] [%5, 8, 8, 16] [1, 1, 1, 1] : tensor<1x8x8x16xf32> to tensor +# CHECK-NEXT: %6 = linalg.generic {indexing_maps = [#map4, #map5, #map4], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%extracted_slice, %cst : tensor, f32) outs(%extracted_slice_0 : tensor) attrs = {__xtc_id_relu_} { +# CHECK-NEXT: ^bb0(%in: f32, %in_1: f32, %out: f32): +# CHECK-NEXT: %7 = arith.maximumf %in, %in_1 : f32 +# CHECK-NEXT: linalg.yield %7 : f32 +# CHECK-NEXT: } -> tensor +# CHECK-NEXT: %inserted_slice = tensor.insert_slice %6 into %arg4[%arg3, 0, 0, 0] [%5, 8, 8, 16] [1, 1, 1, 1] : tensor into tensor<1x8x8x16xf32> +# CHECK-NEXT: scf.yield %inserted_slice : tensor<1x8x8x16xf32> +# CHECK-NEXT: } {"./i"} +# CHECK-NEXT: bufferization.materialize_in_destination %4 in restrict writable %arg2 : (tensor<1x8x8x16xf32>, memref<1x8x8x16xf32>) -> () +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @_vecto(%arg0: !transform.any_op {transform.consumed}) { +# CHECK-NEXT: transform.structured.vectorize %arg0 : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @_post_bufferize(%arg0: !transform.any_op {transform.readonly}) { +# CHECK-NEXT: %0 = transform.structured.match attributes {sym_name = "conv2d_nhwc_mini"} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: transform.apply_patterns to %0 { +# CHECK-NEXT: transform.apply_patterns.vector.lower_outerproduct +# CHECK-NEXT: transform.apply_patterns.vector.lower_contraction +# CHECK-NEXT: } : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: // -----// IR Dump After Tensor Lowering //----- // +# CHECK-NEXT: #map = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 + d4, d2 + d5, d6)> +# CHECK-NEXT: #map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d4, d5, d6, d3)> +# CHECK-NEXT: #map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)> +# CHECK-NEXT: #map3 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> +# CHECK-NEXT: #map4 = affine_map<(d0, d1, d2, d3) -> ()> +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @conv2d_nhwc_mini(%arg0: memref<1x10x10x3xf32> {llvm.noalias}, %arg1: memref<3x3x3x16xf32> {llvm.noalias}, %arg2: memref<1x8x8x16xf32> {llvm.noalias}) { +# CHECK-NEXT: %c3 = arith.constant 3 : index +# CHECK-NEXT: %c16 = arith.constant 16 : index +# CHECK-NEXT: %c8 = arith.constant 8 : index +# CHECK-NEXT: %c1 = arith.constant 1 : index +# CHECK-NEXT: %c0 = arith.constant 0 : index +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: %alloc = memref.alloc() {alignment = 256 : i64} : memref<1x8x8x16xf32> +# CHECK-NEXT: %0 = scf.for %arg3 = %c0 to %c8 step %c1 iter_args(%arg4 = %alloc) -> (memref<1x8x8x16xf32>) { +# CHECK-NEXT: %subview = memref.subview %arg4[0, %arg3, 0, 0] [1, 1, 8, 16] [1, 1, 1, 1] : memref<1x8x8x16xf32> to memref<1x1x8x16xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: %2 = scf.for %arg5 = %c0 to %c8 step %c1 iter_args(%arg6 = %subview) -> (memref<1x1x8x16xf32, strided<[1024, 128, 16, 1], offset: ?>>) { +# CHECK-NEXT: %subview_1 = memref.subview %arg6[0, 0, %arg5, 0] [1, 1, 1, 16] [1, 1, 1, 1] : memref<1x1x8x16xf32, strided<[1024, 128, 16, 1], offset: ?>> to memref<1x1x1x16xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: %3 = scf.for %arg7 = %c0 to %c16 step %c1 iter_args(%arg8 = %subview_1) -> (memref<1x1x1x16xf32, strided<[1024, 128, 16, 1], offset: ?>>) { +# CHECK-NEXT: %subview_3 = memref.subview %arg8[0, 0, 0, %arg7] [1, 1, 1, 1] [1, 1, 1, 1] : memref<1x1x1x16xf32, strided<[1024, 128, 16, 1], offset: ?>> to memref<1x1x1x1xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: linalg.fill {__xtc_id_O_0_} ins(%cst : f32) outs(%subview_3 : memref<1x1x1x1xf32, strided<[1024, 128, 16, 1], offset: ?>>) +# CHECK-NEXT: %subview_4 = memref.subview %arg8[0, 0, 0, %arg7] [1, 1, 1, 1] [1, 1, 1, 1] : memref<1x1x1x16xf32, strided<[1024, 128, 16, 1], offset: ?>> to memref<1x1x1x1xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: memref.copy %subview_3, %subview_4 : memref<1x1x1x1xf32, strided<[1024, 128, 16, 1], offset: ?>> to memref<1x1x1x1xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: scf.yield %arg8 : memref<1x1x1x16xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: } {"./f"} +# CHECK-NEXT: %subview_2 = memref.subview %arg6[0, 0, %arg5, 0] [1, 1, 1, 16] [1, 1, 1, 1] : memref<1x1x8x16xf32, strided<[1024, 128, 16, 1], offset: ?>> to memref<1x1x1x16xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: memref.copy %3, %subview_2 : memref<1x1x1x16xf32, strided<[1024, 128, 16, 1], offset: ?>> to memref<1x1x1x16xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: scf.yield %arg6 : memref<1x1x8x16xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: } {"./w"} +# CHECK-NEXT: %subview_0 = memref.subview %arg4[0, %arg3, 0, 0] [1, 1, 8, 16] [1, 1, 1, 1] : memref<1x8x8x16xf32> to memref<1x1x8x16xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: memref.copy %2, %subview_0 : memref<1x1x8x16xf32, strided<[1024, 128, 16, 1], offset: ?>> to memref<1x1x8x16xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: scf.yield %arg4 : memref<1x8x8x16xf32> +# CHECK-NEXT: } {"./h"} +# CHECK-NEXT: %1 = scf.for %arg3 = %c0 to %c8 step %c1 iter_args(%arg4 = %0) -> (memref<1x8x8x16xf32>) { +# CHECK-NEXT: %subview = memref.subview %arg0[0, %arg3, 0, 0] [1, 3, 10, 3] [1, 1, 1, 1] : memref<1x10x10x3xf32> to memref<1x3x10x3xf32, strided<[300, 30, 3, 1], offset: ?>> +# CHECK-NEXT: %subview_0 = memref.subview %arg4[0, %arg3, 0, 0] [1, 1, 8, 16] [1, 1, 1, 1] : memref<1x8x8x16xf32> to memref<1x1x8x16xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: %2 = scf.for %arg5 = %c0 to %c8 step %c1 iter_args(%arg6 = %subview_0) -> (memref<1x1x8x16xf32, strided<[1024, 128, 16, 1], offset: ?>>) { +# CHECK-NEXT: %subview_2 = memref.subview %subview[0, 0, %arg5, 0] [1, 3, 3, 3] [1, 1, 1, 1] : memref<1x3x10x3xf32, strided<[300, 30, 3, 1], offset: ?>> to memref<1x3x3x3xf32, strided<[300, 30, 3, 1], offset: ?>> +# CHECK-NEXT: %subview_3 = memref.subview %arg6[0, 0, %arg5, 0] [1, 1, 1, 16] [1, 1, 1, 1] : memref<1x1x8x16xf32, strided<[1024, 128, 16, 1], offset: ?>> to memref<1x1x1x16xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: %3 = scf.for %arg7 = %c0 to %c16 step %c1 iter_args(%arg8 = %subview_3) -> (memref<1x1x1x16xf32, strided<[1024, 128, 16, 1], offset: ?>>) { +# CHECK-NEXT: %subview_5 = memref.subview %arg1[0, 0, 0, %arg7] [3, 3, 3, 1] [1, 1, 1, 1] : memref<3x3x3x16xf32> to memref<3x3x3x1xf32, strided<[144, 48, 16, 1], offset: ?>> +# CHECK-NEXT: %subview_6 = memref.subview %arg8[0, 0, 0, %arg7] [1, 1, 1, 1] [1, 1, 1, 1] : memref<1x1x1x16xf32, strided<[1024, 128, 16, 1], offset: ?>> to memref<1x1x1x1xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: %4 = scf.for %arg9 = %c0 to %c3 step %c1 iter_args(%arg10 = %subview_6) -> (memref<1x1x1x1xf32, strided<[1024, 128, 16, 1], offset: ?>>) { +# CHECK-NEXT: %subview_8 = memref.subview %subview_2[0, %arg9, 0, 0] [1, 1, 3, 3] [1, 1, 1, 1] : memref<1x3x3x3xf32, strided<[300, 30, 3, 1], offset: ?>> to memref<1x1x3x3xf32, strided<[300, 30, 3, 1], offset: ?>> +# CHECK-NEXT: %subview_9 = memref.subview %subview_5[%arg9, 0, 0, 0] [1, 3, 3, 1] [1, 1, 1, 1] : memref<3x3x3x1xf32, strided<[144, 48, 16, 1], offset: ?>> to memref<1x3x3x1xf32, strided<[144, 48, 16, 1], offset: ?>> +# CHECK-NEXT: %5 = scf.for %arg11 = %c0 to %c3 step %c1 iter_args(%arg12 = %arg10) -> (memref<1x1x1x1xf32, strided<[1024, 128, 16, 1], offset: ?>>) { +# CHECK-NEXT: %subview_10 = memref.subview %subview_8[0, 0, %arg11, 0] [1, 1, 1, 3] [1, 1, 1, 1] : memref<1x1x3x3xf32, strided<[300, 30, 3, 1], offset: ?>> to memref<1x1x1x3xf32, strided<[300, 30, 3, 1], offset: ?>> +# CHECK-NEXT: %subview_11 = memref.subview %subview_9[0, %arg11, 0, 0] [1, 1, 3, 1] [1, 1, 1, 1] : memref<1x3x3x1xf32, strided<[144, 48, 16, 1], offset: ?>> to memref<1x1x3x1xf32, strided<[144, 48, 16, 1], offset: ?>> +# CHECK-NEXT: %6 = scf.for %arg13 = %c0 to %c3 step %c1 iter_args(%arg14 = %arg12) -> (memref<1x1x1x1xf32, strided<[1024, 128, 16, 1], offset: ?>>) { +# CHECK-NEXT: %subview_12 = memref.subview %subview_10[0, 0, 0, %arg13] [1, 1, 1, 1] [1, 1, 1, 1] : memref<1x1x1x3xf32, strided<[300, 30, 3, 1], offset: ?>> to memref<1x1x1x1xf32, strided<[300, 30, 3, 1], offset: ?>> +# CHECK-NEXT: %subview_13 = memref.subview %subview_11[0, 0, %arg13, 0] [1, 1, 1, 1] [1, 1, 1, 1] : memref<1x1x3x1xf32, strided<[144, 48, 16, 1], offset: ?>> to memref<1x1x1x1xf32, strided<[144, 48, 16, 1], offset: ?>> +# CHECK-NEXT: linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction"]} ins(%subview_12, %subview_13 : memref<1x1x1x1xf32, strided<[300, 30, 3, 1], offset: ?>>, memref<1x1x1x1xf32, strided<[144, 48, 16, 1], offset: ?>>) outs(%arg14 : memref<1x1x1x1xf32, strided<[1024, 128, 16, 1], offset: ?>>) attrs = {__xtc_id_O_} { +# CHECK-NEXT: ^bb0(%in: f32, %in_14: f32, %out: f32): +# CHECK-NEXT: %7 = arith.mulf %in, %in_14 fastmath : f32 +# CHECK-NEXT: %8 = arith.addf %out, %7 fastmath : f32 +# CHECK-NEXT: linalg.yield %8 : f32 +# CHECK-NEXT: } +# CHECK-NEXT: scf.yield %arg14 : memref<1x1x1x1xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: } {"./c"} +# CHECK-NEXT: scf.yield %6 : memref<1x1x1x1xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: } {"./s"} +# CHECK-NEXT: scf.yield %5 : memref<1x1x1x1xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: } {"./r"} +# CHECK-NEXT: %subview_7 = memref.subview %arg8[0, 0, 0, %arg7] [1, 1, 1, 1] [1, 1, 1, 1] : memref<1x1x1x16xf32, strided<[1024, 128, 16, 1], offset: ?>> to memref<1x1x1x1xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: memref.copy %4, %subview_7 : memref<1x1x1x1xf32, strided<[1024, 128, 16, 1], offset: ?>> to memref<1x1x1x1xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: scf.yield %arg8 : memref<1x1x1x16xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: } {"./f"} +# CHECK-NEXT: %subview_4 = memref.subview %arg6[0, 0, %arg5, 0] [1, 1, 1, 16] [1, 1, 1, 1] : memref<1x1x8x16xf32, strided<[1024, 128, 16, 1], offset: ?>> to memref<1x1x1x16xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: memref.copy %3, %subview_4 : memref<1x1x1x16xf32, strided<[1024, 128, 16, 1], offset: ?>> to memref<1x1x1x16xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: scf.yield %arg6 : memref<1x1x8x16xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: } {"./w"} +# CHECK-NEXT: %subview_1 = memref.subview %arg4[0, %arg3, 0, 0] [1, 1, 8, 16] [1, 1, 1, 1] : memref<1x8x8x16xf32> to memref<1x1x8x16xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: memref.copy %2, %subview_1 : memref<1x1x8x16xf32, strided<[1024, 128, 16, 1], offset: ?>> to memref<1x1x8x16xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: scf.yield %arg4 : memref<1x8x8x16xf32> +# CHECK-NEXT: } {"./h"} +# CHECK-NEXT: linalg.generic {indexing_maps = [#map3, #map4, #map3], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%1, %cst : memref<1x8x8x16xf32>, f32) outs(%arg2 : memref<1x8x8x16xf32>) attrs = {__xtc_id_relu_} { +# CHECK-NEXT: ^bb0(%in: f32, %in_0: f32, %out: f32): +# CHECK-NEXT: %2 = arith.maximumf %in, %in_0 : f32 +# CHECK-NEXT: linalg.yield %2 : f32 +# CHECK-NEXT: } +# CHECK-NEXT: memref.copy %arg2, %arg2 : memref<1x8x8x16xf32> to memref<1x8x8x16xf32> +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: graph: +# CHECK-NEXT: name: conv2d_nhwc_mini +# CHECK-NEXT: inputs: +# CHECK-NEXT: - %0 : 1x10x10x3xfloat32 +# CHECK-NEXT: - %1 : 3x3x3x16xfloat32 +# CHECK-NEXT: outputs: +# CHECK-NEXT: - %3 : 1x8x8x16xfloat32 +# CHECK-NEXT: nodes: +# CHECK-NEXT: - %2: conv2d(%0, %1, stride=(1, 1)) {name = 'O'} : [1x10x10x3xfloat32, 3x3x3x16xfloat32] -> [1x8x8x16xfloat32] +# CHECK-NEXT: - %3: relu(%2) {name = 'relu'} : [1x8x8x16xfloat32] -> [1x8x8x16xfloat32] +# CHECK-NEXT: +# CHECK-NEXT: CODE: 0 diff --git a/tests/filecheck/backends/tensor_dialect/test_relu_tensor.py b/tests/filecheck/backends/tensor_dialect/test_relu_tensor.py new file mode 100644 index 000000000..fd02a3498 --- /dev/null +++ b/tests/filecheck/backends/tensor_dialect/test_relu_tensor.py @@ -0,0 +1,149 @@ +# RUN: python %s 2>&1 | filecheck %s +# UNSUPPORTED: mlir-target=nvgpu + +import xtc.graphs.xtc.op as O +from xtc.backends.mlir import Backend + +I, dtype = 128, "float32" +a = O.tensor((I,), dtype, name="A") + +with O.graph(name="relu") as gb: + O.relu(a, name="relu") + +graph = gb.graph +print(graph) + +impl = Backend(graph, use_tensor_dialect=True) + +sch = impl.get_scheduler(default_node="relu") +sch.tile("i", {"i1": 16}) +sch.interchange([ "i", "i1"]) +sch.vectorize(["i1"]) +sched = sch.schedule() + +comp = impl.get_compiler( + shared_lib=True, + dump_file="relu_mlir_tensor", + print_source_ir=True, + print_transformed_ir=True, + print_bufferization_ir=True, +) +module = comp.compile(sched) +executor = module.get_executor(validate=True) +res = executor.execute() +print(f"CODE: {res}") + +# CHECK: // -----// IR Dump Before transform //----- // +# CHECK-NEXT: #map = affine_map<(d0) -> (d0)> +# CHECK-NEXT: #map1 = affine_map<(d0) -> ()> +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @relu(%arg0: tensor<128xf32> {llvm.noalias}, %arg1: memref<128xf32> {llvm.noalias}) { +# CHECK-NEXT: %0 = tensor.empty() : tensor<128xf32> +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: %1 = linalg.generic {indexing_maps = [#map, #map1, #map], iterator_types = ["parallel"]} ins(%arg0, %cst : tensor<128xf32>, f32) outs(%0 : tensor<128xf32>) attrs = {__xtc_id_relu_} { +# CHECK-NEXT: ^bb0(%in: f32, %in_0: f32, %out: f32): +# CHECK-NEXT: %2 = arith.maximumf %in, %in_0 : f32 +# CHECK-NEXT: linalg.yield %2 : f32 +# CHECK-NEXT: } -> tensor<128xf32> +# CHECK-NEXT: bufferization.materialize_in_destination %1 in restrict writable %arg1 : (tensor<128xf32>, memref<128xf32>) -> () +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @_vecto(%arg0: !transform.any_op {transform.consumed}) { +# CHECK-NEXT: transform.structured.vectorize %arg0 : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @_post_bufferize(%arg0: !transform.any_op {transform.readonly}) { +# CHECK-NEXT: %0 = transform.structured.match attributes {sym_name = "relu"} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: transform.apply_patterns to %0 { +# CHECK-NEXT: transform.apply_patterns.vector.lower_outerproduct +# CHECK-NEXT: transform.apply_patterns.vector.lower_contraction +# CHECK-NEXT: } : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { +# CHECK-NEXT: %0 = transform.structured.match attributes {__xtc_id_relu_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_linalg_op, %loops = transform.structured.tile_using_for %0 tile_sizes [16] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops "./i" : !transform.any_op +# CHECK-NEXT: %1 = transform.get_parent_op %tiled_linalg_op : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: transform.apply_patterns to %1 { +# CHECK-NEXT: transform.apply_patterns.linalg.fold_unit_extent_dims_via_slices +# CHECK-NEXT: } : !transform.any_op +# CHECK-NEXT: %2 = transform.structured.match interface{LinalgOp} in %1 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: transform.include @_vecto failures(suppress) (%2) : (!transform.any_op) -> () +# CHECK-NEXT: %3 = transform.get_parent_op %loops {isolated_from_above} : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: transform.apply_patterns to %3 { +# CHECK-NEXT: transform.apply_patterns.vector.reduction_to_contract +# CHECK-NEXT: transform.apply_patterns.vector.transfer_permutation_patterns +# CHECK-NEXT: } : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: // -----// IR Dump After transform //----- // +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @relu(%arg0: tensor<128xf32> {llvm.noalias}, %arg1: memref<128xf32> {llvm.noalias}) { +# CHECK-NEXT: %cst = arith.constant dense<0.000000e+00> : vector<16xf32> +# CHECK-NEXT: %0 = ub.poison : f32 +# CHECK-NEXT: %c16 = arith.constant 16 : index +# CHECK-NEXT: %c128 = arith.constant 128 : index +# CHECK-NEXT: %c0 = arith.constant 0 : index +# CHECK-NEXT: %1 = tensor.empty() : tensor<128xf32> +# CHECK-NEXT: %2 = scf.for %arg2 = %c0 to %c128 step %c16 iter_args(%arg3 = %1) -> (tensor<128xf32>) { +# CHECK-NEXT: %extracted_slice = tensor.extract_slice %arg0[%arg2] [16] [1] : tensor<128xf32> to tensor<16xf32> +# CHECK-NEXT: %extracted_slice_0 = tensor.extract_slice %arg3[%arg2] [16] [1] : tensor<128xf32> to tensor<16xf32> +# CHECK-NEXT: %3 = vector.transfer_read %extracted_slice[%c0], %0 {in_bounds = [true]} : tensor<16xf32>, vector<16xf32> +# CHECK-NEXT: %4 = arith.maximumf %3, %cst : vector<16xf32> +# CHECK-NEXT: %5 = vector.transfer_write %4, %extracted_slice_0[%c0] {in_bounds = [true]} : vector<16xf32>, tensor<16xf32> +# CHECK-NEXT: %inserted_slice = tensor.insert_slice %5 into %arg3[%arg2] [16] [1] : tensor<16xf32> into tensor<128xf32> +# CHECK-NEXT: scf.yield %inserted_slice : tensor<128xf32> +# CHECK-NEXT: } {"./i"} +# CHECK-NEXT: bufferization.materialize_in_destination %2 in restrict writable %arg1 : (tensor<128xf32>, memref<128xf32>) -> () +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @_vecto(%arg0: !transform.any_op {transform.consumed}) { +# CHECK-NEXT: transform.structured.vectorize %arg0 : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @_post_bufferize(%arg0: !transform.any_op {transform.readonly}) { +# CHECK-NEXT: %0 = transform.structured.match attributes {sym_name = "relu"} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: transform.apply_patterns to %0 { +# CHECK-NEXT: transform.apply_patterns.vector.lower_outerproduct +# CHECK-NEXT: transform.apply_patterns.vector.lower_contraction +# CHECK-NEXT: } : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: // -----// IR Dump After Tensor Lowering //----- // +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @relu(%arg0: memref<128xf32> {llvm.noalias}, %arg1: memref<128xf32> {llvm.noalias}) { +# CHECK-NEXT: %cst = arith.constant dense<0.000000e+00> : vector<16xf32> +# CHECK-NEXT: %0 = ub.poison : f32 +# CHECK-NEXT: %c16 = arith.constant 16 : index +# CHECK-NEXT: %c128 = arith.constant 128 : index +# CHECK-NEXT: %c0 = arith.constant 0 : index +# CHECK-NEXT: %1 = scf.for %arg2 = %c0 to %c128 step %c16 iter_args(%arg3 = %arg1) -> (memref<128xf32>) { +# CHECK-NEXT: %subview = memref.subview %arg0[%arg2] [16] [1] : memref<128xf32> to memref<16xf32, strided<[1], offset: ?>> +# CHECK-NEXT: %subview_0 = memref.subview %arg3[%arg2] [16] [1] : memref<128xf32> to memref<16xf32, strided<[1], offset: ?>> +# CHECK-NEXT: %2 = vector.transfer_read %subview[%c0], %0 {in_bounds = [true]} : memref<16xf32, strided<[1], offset: ?>>, vector<16xf32> +# CHECK-NEXT: %3 = arith.maximumf %2, %cst : vector<16xf32> +# CHECK-NEXT: vector.transfer_write %3, %subview_0[%c0] {in_bounds = [true]} : vector<16xf32>, memref<16xf32, strided<[1], offset: ?>> +# CHECK-NEXT: %subview_1 = memref.subview %arg3[%arg2] [16] [1] : memref<128xf32> to memref<16xf32, strided<[1], offset: ?>> +# CHECK-NEXT: memref.copy %subview_0, %subview_1 : memref<16xf32, strided<[1], offset: ?>> to memref<16xf32, strided<[1], offset: ?>> +# CHECK-NEXT: scf.yield %arg3 : memref<128xf32> +# CHECK-NEXT: } {"./i"} +# CHECK-NEXT: memref.copy %1, %arg1 : memref<128xf32> to memref<128xf32> +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: graph: +# CHECK-NEXT: name: relu +# CHECK-NEXT: inputs: +# CHECK-NEXT: - %0 : 128xfloat32 +# CHECK-NEXT: outputs: +# CHECK-NEXT: - %1 : 128xfloat32 +# CHECK-NEXT: nodes: +# CHECK-NEXT: - %1: relu(%0) {name = 'relu'} : [128xfloat32] -> [128xfloat32] +# CHECK-NEXT: +# CHECK-NEXT: CODE: 0 diff --git a/tests/filecheck/backends/test_conv2d_relu_mlir.py b/tests/filecheck/backends/test_conv2d_relu_mlir.py new file mode 100644 index 000000000..fa8dce602 --- /dev/null +++ b/tests/filecheck/backends/test_conv2d_relu_mlir.py @@ -0,0 +1,198 @@ +# RUN: python %s 2>&1 | filecheck %s +# UNSUPPORTED: mlir-target=nvgpu + +import xtc.graphs.xtc.op as O +from xtc.backends.mlir import Backend + +N, H, W, F, R, S, C, SH, SW, dtype = 1, 8, 8, 16, 3, 3, 3, 1, 1, "float32" +a = O.tensor((N, H + R - 1, W + S - 1, C), dtype, name="I") +b = O.tensor((R, S, C, F), dtype, name="W") + +with O.graph(name="conv2d_nhwc_mini") as gb: + c = O.conv2d(a, b, stride=(SH, SW), name="O") + O.relu(c, name="relu") + +graph = gb.graph +print(graph) + +impl = Backend(graph) + +sch = impl.get_scheduler(default_node="relu") +sch.tile("i", {"i1": 16}) +sch.interchange([ "i", "i1"]) +sch.vectorize(["i1"]) +sched = sch.schedule() + +comp = impl.get_compiler( + shared_lib=True, + dump_file="conv2d_relu_mlir_tensor", + print_source_ir=True, + print_transformed_ir=True, +) +module = comp.compile(sched) +executor = module.get_executor(validate=True) +res = executor.execute() +print(f"CODE: {res}") + +# CHECK: // -----// IR Dump Before transform //----- // +# CHECK-NEXT: #map = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 + d4, d2 + d5, d6)> +# CHECK-NEXT: #map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d4, d5, d6, d3)> +# CHECK-NEXT: #map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)> +# CHECK-NEXT: #map3 = affine_map<(d0) -> (d0)> +# CHECK-NEXT: #map4 = affine_map<(d0) -> ()> +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @conv2d_nhwc_mini(%arg0: memref<1x10x10x3xf32> {llvm.noalias}, %arg1: memref<3x3x3x16xf32> {llvm.noalias}, %arg2: memref<1x8x8x16xf32> {llvm.noalias}) { +# CHECK-NEXT: %alloca = memref.alloca() {alignment = 256 : i64} : memref<1x8x8x16xf32> +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: linalg.fill {__xtc_id_O_0_} ins(%cst : f32) outs(%alloca : memref<1x8x8x16xf32>) +# CHECK-NEXT: linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction"]} ins(%arg0, %arg1 : memref<1x10x10x3xf32>, memref<3x3x3x16xf32>) outs(%alloca : memref<1x8x8x16xf32>) attrs = {__xtc_id_O_} { +# CHECK-NEXT: ^bb0(%in: f32, %in_2: f32, %out: f32): +# CHECK-NEXT: %0 = arith.mulf %in, %in_2 fastmath : f32 +# CHECK-NEXT: %1 = arith.addf %out, %0 fastmath : f32 +# CHECK-NEXT: linalg.yield %1 : f32 +# CHECK-NEXT: } +# CHECK-NEXT: %collapse_shape = memref.collapse_shape %alloca [[0, 1, 2, 3]] : memref<1x8x8x16xf32> into memref<1024xf32> +# CHECK-NEXT: %collapse_shape_0 = memref.collapse_shape %arg2 [[0, 1, 2, 3]] : memref<1x8x8x16xf32> into memref<1024xf32> +# CHECK-NEXT: %cst_1 = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: linalg.generic {indexing_maps = [#map3, #map4, #map3], iterator_types = ["parallel"]} ins(%collapse_shape, %cst_1 : memref<1024xf32>, f32) outs(%collapse_shape_0 : memref<1024xf32>) attrs = {__xtc_id_relu_} { +# CHECK-NEXT: ^bb0(%in: f32, %in_2: f32, %out: f32): +# CHECK-NEXT: %0 = arith.maximumf %in, %in_2 : f32 +# CHECK-NEXT: linalg.yield %0 : f32 +# CHECK-NEXT: } +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @_vecto(%arg0: !transform.any_op {transform.consumed}) { +# CHECK-NEXT: transform.structured.vectorize %arg0 : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { +# CHECK-NEXT: %0 = transform.structured.match attributes {__xtc_id_O_0_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_linalg_op, %loops = transform.structured.tile_using_for %0 tile_sizes [1, 0, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops "./b" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_0, %loops_1 = transform.structured.tile_using_for %tiled_linalg_op tile_sizes [0, 1, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_1 "./h" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_2, %loops_3 = transform.structured.tile_using_for %tiled_linalg_op_0 tile_sizes [0, 0, 1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_3 "./w" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_4, %loops_5 = transform.structured.tile_using_for %tiled_linalg_op_2 tile_sizes [0, 0, 0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_5 "./f" : !transform.any_op +# CHECK-NEXT: %1 = transform.structured.match attributes {__xtc_id_O_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_6, %loops_7 = transform.structured.tile_using_for %1 tile_sizes [1, 0, 0, 0, 0, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_7 "./b" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_8, %loops_9 = transform.structured.tile_using_for %tiled_linalg_op_6 tile_sizes [0, 1, 0, 0, 0, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_9 "./h" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_10, %loops_11 = transform.structured.tile_using_for %tiled_linalg_op_8 tile_sizes [0, 0, 1, 0, 0, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_11 "./w" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_12, %loops_13 = transform.structured.tile_using_for %tiled_linalg_op_10 tile_sizes [0, 0, 0, 1, 0, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_13 "./f" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_14, %loops_15 = transform.structured.tile_using_for %tiled_linalg_op_12 tile_sizes [0, 0, 0, 0, 1, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_15 "./r" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_16, %loops_17 = transform.structured.tile_using_for %tiled_linalg_op_14 tile_sizes [0, 0, 0, 0, 0, 1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_17 "./s" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_18, %loops_19 = transform.structured.tile_using_for %tiled_linalg_op_16 tile_sizes [0, 0, 0, 0, 0, 0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_19 "./c" : !transform.any_op +# CHECK-NEXT: %2 = transform.structured.match attributes {__xtc_id_relu_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_20, %loops_21 = transform.structured.tile_using_for %2 tile_sizes [16] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_21 "./i" : !transform.any_op +# CHECK-NEXT: transform.include @_vecto failures(suppress) (%tiled_linalg_op_20) : (!transform.any_op) -> () +# CHECK-NEXT: %3 = transform.get_parent_op %loops_21 {isolated_from_above} : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: transform.apply_patterns to %3 { +# CHECK-NEXT: transform.apply_patterns.vector.reduction_to_contract +# CHECK-NEXT: transform.apply_patterns.vector.transfer_permutation_patterns +# CHECK-NEXT: } : !transform.any_op +# CHECK-NEXT: transform.apply_patterns to %3 { +# CHECK-NEXT: transform.apply_patterns.vector.lower_outerproduct +# CHECK-NEXT: transform.apply_patterns.vector.lower_contraction +# CHECK-NEXT: } : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: // -----// IR Dump After transform //----- // +# CHECK-NEXT: #map = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 + d4, d2 + d5, d6)> +# CHECK-NEXT: #map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d4, d5, d6, d3)> +# CHECK-NEXT: #map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)> +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @conv2d_nhwc_mini(%arg0: memref<1x10x10x3xf32> {llvm.noalias}, %arg1: memref<3x3x3x16xf32> {llvm.noalias}, %arg2: memref<1x8x8x16xf32> {llvm.noalias}) { +# CHECK-NEXT: %cst = arith.constant dense<0.000000e+00> : vector<16xf32> +# CHECK-NEXT: %0 = ub.poison : f32 +# CHECK-NEXT: %c1024 = arith.constant 1024 : index +# CHECK-NEXT: %c3 = arith.constant 3 : index +# CHECK-NEXT: %c16 = arith.constant 16 : index +# CHECK-NEXT: %c8 = arith.constant 8 : index +# CHECK-NEXT: %c1 = arith.constant 1 : index +# CHECK-NEXT: %c0 = arith.constant 0 : index +# CHECK-NEXT: %cst_0 = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: %alloca = memref.alloca() {alignment = 256 : i64} : memref<1x8x8x16xf32> +# CHECK-NEXT: scf.for %arg3 = %c0 to %c1 step %c1 { +# CHECK-NEXT: %subview = memref.subview %alloca[%arg3, 0, 0, 0] [1, 8, 8, 16] [1, 1, 1, 1] : memref<1x8x8x16xf32> to memref<1x8x8x16xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg4 = %c0 to %c8 step %c1 { +# CHECK-NEXT: %subview_2 = memref.subview %subview[0, %arg4, 0, 0] [1, 1, 8, 16] [1, 1, 1, 1] : memref<1x8x8x16xf32, strided<[1024, 128, 16, 1], offset: ?>> to memref<1x1x8x16xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg5 = %c0 to %c8 step %c1 { +# CHECK-NEXT: %subview_3 = memref.subview %subview_2[0, 0, %arg5, 0] [1, 1, 1, 16] [1, 1, 1, 1] : memref<1x1x8x16xf32, strided<[1024, 128, 16, 1], offset: ?>> to memref<1x1x1x16xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg6 = %c0 to %c16 step %c1 { +# CHECK-NEXT: %subview_4 = memref.subview %subview_3[0, 0, 0, %arg6] [1, 1, 1, 1] [1, 1, 1, 1] : memref<1x1x1x16xf32, strided<[1024, 128, 16, 1], offset: ?>> to memref<1x1x1x1xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: linalg.fill {__xtc_id_O_0_} ins(%cst_0 : f32) outs(%subview_4 : memref<1x1x1x1xf32, strided<[1024, 128, 16, 1], offset: ?>>) +# CHECK-NEXT: } {"./f"} +# CHECK-NEXT: } {"./w"} +# CHECK-NEXT: } {"./h"} +# CHECK-NEXT: } {"./b"} +# CHECK-NEXT: scf.for %arg3 = %c0 to %c1 step %c1 { +# CHECK-NEXT: %subview = memref.subview %arg0[%arg3, 0, 0, 0] [1, 10, 10, 3] [1, 1, 1, 1] : memref<1x10x10x3xf32> to memref<1x10x10x3xf32, strided<[300, 30, 3, 1], offset: ?>> +# CHECK-NEXT: %subview_2 = memref.subview %arg1[0, 0, 0, 0] [3, 3, 3, 16] [1, 1, 1, 1] : memref<3x3x3x16xf32> to memref<3x3x3x16xf32, strided<[144, 48, 16, 1]>> +# CHECK-NEXT: %subview_3 = memref.subview %alloca[%arg3, 0, 0, 0] [1, 8, 8, 16] [1, 1, 1, 1] : memref<1x8x8x16xf32> to memref<1x8x8x16xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg4 = %c0 to %c8 step %c1 { +# CHECK-NEXT: %subview_4 = memref.subview %subview[0, %arg4, 0, 0] [1, 3, 10, 3] [1, 1, 1, 1] : memref<1x10x10x3xf32, strided<[300, 30, 3, 1], offset: ?>> to memref<1x3x10x3xf32, strided<[300, 30, 3, 1], offset: ?>> +# CHECK-NEXT: %subview_5 = memref.subview %subview_3[0, %arg4, 0, 0] [1, 1, 8, 16] [1, 1, 1, 1] : memref<1x8x8x16xf32, strided<[1024, 128, 16, 1], offset: ?>> to memref<1x1x8x16xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg5 = %c0 to %c8 step %c1 { +# CHECK-NEXT: %subview_6 = memref.subview %subview_4[0, 0, %arg5, 0] [1, 3, 3, 3] [1, 1, 1, 1] : memref<1x3x10x3xf32, strided<[300, 30, 3, 1], offset: ?>> to memref<1x3x3x3xf32, strided<[300, 30, 3, 1], offset: ?>> +# CHECK-NEXT: %subview_7 = memref.subview %subview_5[0, 0, %arg5, 0] [1, 1, 1, 16] [1, 1, 1, 1] : memref<1x1x8x16xf32, strided<[1024, 128, 16, 1], offset: ?>> to memref<1x1x1x16xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg6 = %c0 to %c16 step %c1 { +# CHECK-NEXT: %subview_8 = memref.subview %subview_2[0, 0, 0, %arg6] [3, 3, 3, 1] [1, 1, 1, 1] : memref<3x3x3x16xf32, strided<[144, 48, 16, 1]>> to memref<3x3x3x1xf32, strided<[144, 48, 16, 1], offset: ?>> +# CHECK-NEXT: %subview_9 = memref.subview %subview_7[0, 0, 0, %arg6] [1, 1, 1, 1] [1, 1, 1, 1] : memref<1x1x1x16xf32, strided<[1024, 128, 16, 1], offset: ?>> to memref<1x1x1x1xf32, strided<[1024, 128, 16, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg7 = %c0 to %c3 step %c1 { +# CHECK-NEXT: %subview_10 = memref.subview %subview_6[0, %arg7, 0, 0] [1, 1, 3, 3] [1, 1, 1, 1] : memref<1x3x3x3xf32, strided<[300, 30, 3, 1], offset: ?>> to memref<1x1x3x3xf32, strided<[300, 30, 3, 1], offset: ?>> +# CHECK-NEXT: %subview_11 = memref.subview %subview_8[%arg7, 0, 0, 0] [1, 3, 3, 1] [1, 1, 1, 1] : memref<3x3x3x1xf32, strided<[144, 48, 16, 1], offset: ?>> to memref<1x3x3x1xf32, strided<[144, 48, 16, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg8 = %c0 to %c3 step %c1 { +# CHECK-NEXT: %subview_12 = memref.subview %subview_10[0, 0, %arg8, 0] [1, 1, 1, 3] [1, 1, 1, 1] : memref<1x1x3x3xf32, strided<[300, 30, 3, 1], offset: ?>> to memref<1x1x1x3xf32, strided<[300, 30, 3, 1], offset: ?>> +# CHECK-NEXT: %subview_13 = memref.subview %subview_11[0, %arg8, 0, 0] [1, 1, 3, 1] [1, 1, 1, 1] : memref<1x3x3x1xf32, strided<[144, 48, 16, 1], offset: ?>> to memref<1x1x3x1xf32, strided<[144, 48, 16, 1], offset: ?>> +# CHECK-NEXT: scf.for %arg9 = %c0 to %c3 step %c1 { +# CHECK-NEXT: %subview_14 = memref.subview %subview_12[0, 0, 0, %arg9] [1, 1, 1, 1] [1, 1, 1, 1] : memref<1x1x1x3xf32, strided<[300, 30, 3, 1], offset: ?>> to memref<1x1x1x1xf32, strided<[300, 30, 3, 1], offset: ?>> +# CHECK-NEXT: %subview_15 = memref.subview %subview_13[0, 0, %arg9, 0] [1, 1, 1, 1] [1, 1, 1, 1] : memref<1x1x3x1xf32, strided<[144, 48, 16, 1], offset: ?>> to memref<1x1x1x1xf32, strided<[144, 48, 16, 1], offset: ?>> +# CHECK-NEXT: linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction"]} ins(%subview_14, %subview_15 : memref<1x1x1x1xf32, strided<[300, 30, 3, 1], offset: ?>>, memref<1x1x1x1xf32, strided<[144, 48, 16, 1], offset: ?>>) outs(%subview_9 : memref<1x1x1x1xf32, strided<[1024, 128, 16, 1], offset: ?>>) attrs = {__xtc_id_O_} { +# CHECK-NEXT: ^bb0(%in: f32, %in_16: f32, %out: f32): +# CHECK-NEXT: %1 = arith.mulf %in, %in_16 fastmath : f32 +# CHECK-NEXT: %2 = arith.addf %out, %1 fastmath : f32 +# CHECK-NEXT: linalg.yield %2 : f32 +# CHECK-NEXT: } +# CHECK-NEXT: } {"./c"} +# CHECK-NEXT: } {"./s"} +# CHECK-NEXT: } {"./r"} +# CHECK-NEXT: } {"./f"} +# CHECK-NEXT: } {"./w"} +# CHECK-NEXT: } {"./h"} +# CHECK-NEXT: } {"./b"} +# CHECK-NEXT: %collapse_shape = memref.collapse_shape %alloca [[0, 1, 2, 3]] : memref<1x8x8x16xf32> into memref<1024xf32> +# CHECK-NEXT: %collapse_shape_1 = memref.collapse_shape %arg2 [[0, 1, 2, 3]] : memref<1x8x8x16xf32> into memref<1024xf32> +# CHECK-NEXT: scf.for %arg3 = %c0 to %c1024 step %c16 { +# CHECK-NEXT: %subview = memref.subview %collapse_shape[%arg3] [16] [1] : memref<1024xf32> to memref<16xf32, strided<[1], offset: ?>> +# CHECK-NEXT: %subview_2 = memref.subview %collapse_shape_1[%arg3] [16] [1] : memref<1024xf32> to memref<16xf32, strided<[1], offset: ?>> +# CHECK-NEXT: %1 = vector.transfer_read %subview[%c0], %0 {in_bounds = [true]} : memref<16xf32, strided<[1], offset: ?>>, vector<16xf32> +# CHECK-NEXT: %2 = arith.maximumf %1, %cst : vector<16xf32> +# CHECK-NEXT: vector.transfer_write %2, %subview_2[%c0] {in_bounds = [true]} : vector<16xf32>, memref<16xf32, strided<[1], offset: ?>> +# CHECK-NEXT: } {"./i"} +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: graph: +# CHECK-NEXT: name: conv2d_nhwc_mini +# CHECK-NEXT: inputs: +# CHECK-NEXT: - %0 : 1x10x10x3xfloat32 +# CHECK-NEXT: - %1 : 3x3x3x16xfloat32 +# CHECK-NEXT: outputs: +# CHECK-NEXT: - %3 : 1x8x8x16xfloat32 +# CHECK-NEXT: nodes: +# CHECK-NEXT: - %2: conv2d(%0, %1, stride=(1, 1)) {name = 'O'} : [1x10x10x3xfloat32, 3x3x3x16xfloat32] -> [1x8x8x16xfloat32] +# CHECK-NEXT: - %3: relu(%2) {name = 'relu'} : [1x8x8x16xfloat32] -> [1x8x8x16xfloat32] +# CHECK-NEXT: +# CHECK-NEXT: CODE: 0 diff --git a/tests/filecheck/backends/test_relu_mlir.py b/tests/filecheck/backends/test_relu_mlir.py new file mode 100644 index 000000000..807ff581a --- /dev/null +++ b/tests/filecheck/backends/test_relu_mlir.py @@ -0,0 +1,100 @@ +# RUN: python %s 2>&1 | filecheck %s +# UNSUPPORTED: mlir-target=nvgpu + +import xtc.graphs.xtc.op as O +from xtc.backends.mlir import Backend + +I, dtype = 128, "float32" +a = O.tensor((I,), dtype, name="A") + +with O.graph(name="relu") as gb: + O.relu(a, name="relu") + +graph = gb.graph +print(graph) + +impl = Backend(graph) + +sch = impl.get_scheduler(default_node="relu") +sch.tile("i", {"i1": 16}) +sch.interchange([ "i", "i1"]) +sch.vectorize(["i1"]) +sched = sch.schedule() + +comp = impl.get_compiler( + shared_lib=True, + dump_file="relu_mlir_tensor", + print_source_ir=True, + print_transformed_ir=True, +) +module = comp.compile(sched) +executor = module.get_executor(validate=True) +res = executor.execute() +print(f"CODE: {res}") + +# CHECK: // -----// IR Dump Before transform //----- // +# CHECK-NEXT: #map = affine_map<(d0) -> (d0)> +# CHECK-NEXT: #map1 = affine_map<(d0) -> ()> +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @relu(%arg0: memref<128xf32> {llvm.noalias}, %arg1: memref<128xf32> {llvm.noalias}) { +# CHECK-NEXT: %collapse_shape = memref.collapse_shape %arg0 [[0]] : memref<128xf32> into memref<128xf32> +# CHECK-NEXT: %collapse_shape_0 = memref.collapse_shape %arg1 [[0]] : memref<128xf32> into memref<128xf32> +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: linalg.generic {indexing_maps = [#map, #map1, #map], iterator_types = ["parallel"]} ins(%collapse_shape, %cst : memref<128xf32>, f32) outs(%collapse_shape_0 : memref<128xf32>) attrs = {__xtc_id_relu_} { +# CHECK-NEXT: ^bb0(%in: f32, %in_1: f32, %out: f32): +# CHECK-NEXT: %0 = arith.maximumf %in, %in_1 : f32 +# CHECK-NEXT: linalg.yield %0 : f32 +# CHECK-NEXT: } +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @_vecto(%arg0: !transform.any_op {transform.consumed}) { +# CHECK-NEXT: transform.structured.vectorize %arg0 : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { +# CHECK-NEXT: %0 = transform.structured.match attributes {__xtc_id_relu_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_linalg_op, %loops = transform.structured.tile_using_for %0 tile_sizes [16] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops "./i" : !transform.any_op +# CHECK-NEXT: transform.include @_vecto failures(suppress) (%tiled_linalg_op) : (!transform.any_op) -> () +# CHECK-NEXT: %1 = transform.get_parent_op %loops {isolated_from_above} : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: transform.apply_patterns to %1 { +# CHECK-NEXT: transform.apply_patterns.vector.reduction_to_contract +# CHECK-NEXT: transform.apply_patterns.vector.transfer_permutation_patterns +# CHECK-NEXT: } : !transform.any_op +# CHECK-NEXT: transform.apply_patterns to %1 { +# CHECK-NEXT: transform.apply_patterns.vector.lower_outerproduct +# CHECK-NEXT: transform.apply_patterns.vector.lower_contraction +# CHECK-NEXT: } : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: // -----// IR Dump After transform //----- // +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @relu(%arg0: memref<128xf32> {llvm.noalias}, %arg1: memref<128xf32> {llvm.noalias}) { +# CHECK-NEXT: %cst = arith.constant dense<0.000000e+00> : vector<16xf32> +# CHECK-NEXT: %0 = ub.poison : f32 +# CHECK-NEXT: %c16 = arith.constant 16 : index +# CHECK-NEXT: %c128 = arith.constant 128 : index +# CHECK-NEXT: %c0 = arith.constant 0 : index +# CHECK-NEXT: scf.for %arg2 = %c0 to %c128 step %c16 { +# CHECK-NEXT: %subview = memref.subview %arg0[%arg2] [16] [1] : memref<128xf32> to memref<16xf32, strided<[1], offset: ?>> +# CHECK-NEXT: %subview_0 = memref.subview %arg1[%arg2] [16] [1] : memref<128xf32> to memref<16xf32, strided<[1], offset: ?>> +# CHECK-NEXT: %1 = vector.transfer_read %subview[%c0], %0 {in_bounds = [true]} : memref<16xf32, strided<[1], offset: ?>>, vector<16xf32> +# CHECK-NEXT: %2 = arith.maximumf %1, %cst : vector<16xf32> +# CHECK-NEXT: vector.transfer_write %2, %subview_0[%c0] {in_bounds = [true]} : vector<16xf32>, memref<16xf32, strided<[1], offset: ?>> +# CHECK-NEXT: } {"./i"} +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: graph: +# CHECK-NEXT: name: relu +# CHECK-NEXT: inputs: +# CHECK-NEXT: - %0 : 128xfloat32 +# CHECK-NEXT: outputs: +# CHECK-NEXT: - %1 : 128xfloat32 +# CHECK-NEXT: nodes: +# CHECK-NEXT: - %1: relu(%0) {name = 'relu'} : [128xfloat32] -> [128xfloat32] +# CHECK-NEXT: +# CHECK-NEXT: CODE: 0 diff --git a/tests/filecheck/mlir_loop/gen_assembly/skylake_generic_conv2d.mlir b/tests/filecheck/mlir_loop/gen_assembly/skylake_generic_conv2d.mlir new file mode 100644 index 000000000..7839e031e --- /dev/null +++ b/tests/filecheck/mlir_loop/gen_assembly/skylake_generic_conv2d.mlir @@ -0,0 +1,51 @@ +// RUN: mlir-loop --no-alias --arch x86-64 --cpu skylake --print-assembly --hide-jumps %s 2>&1 | grep -v '\(nop\|ret\)' | filecheck %s +// REQUIRES: mlir-target=llvmir +// Assembly output will differ a bit when using C. + +func.func @myfun( + %I: memref<1x30x66x64xf32>, + %K: memref<3x3x64x128xf32>, + %O: memref<1x28x64x128xf32> +) { + linalg.generic { + indexing_maps = [ + affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 + d4, d2 + d5, d6)>, + affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d4, d5, d6, d3)>, + affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)> + ], + iterator_types = [ + "parallel", + "parallel", + "parallel", + "parallel", + "reduction", + "reduction", + "reduction" + ] + } + ins(%I, %K : memref<1x30x66x64xf32>, memref<3x3x64x128xf32>) + outs(%O : memref<1x28x64x128xf32>) + attrs = { + loop.dims = ["n","h","w","f","r","s","c"], + loop.schedule = { + "n", + "h", + "f", + "r", + "s", + "c", + "w", + "f#1" = {"unroll"}, + "c#8" = {"unroll"}, + "w#64" = {"vectorize"} + } + } + { + ^bb0(%in: f32, %in_0: f32, %out: f32): + %1 = arith.mulf %in, %in_0 : f32 + %2 = arith.addf %out, %1 : f32 + linalg.yield %2 : f32 + } + return +} +// CHECK: Disassembly of section .text: