117 lines · plain
1// DEFINE: %{compile} = mlir-opt %s \2// DEFINE: -transform-interpreter -test-transform-dialect-erase-schedule \3// DEFINE: --lower-vector-mask |\4// DEFINE: mlir-opt \5// DEFINE: -test-lower-to-llvm -o %t6// DEFINE: %{entry_point} = main7// DEFINE: %{run} = mlir-runner %t -e %{entry_point} -entry-point-result=void \8// DEFINE: -shared-libs=%mlir_runner_utils,%mlir_c_runner_utils9 10// RUN: rm -f %t && %{compile} && %{run} | FileCheck %s11 12/// End-to-end test for linalg.unpack where one of the inner tile sizes is13/// dynamic.14 15func.func @main() {16 // Allocate and initialise the inputs17 %A_alloc = tensor.empty() : tensor<7x3xi32>18 19 %A = arith.constant dense<[20 [[[1],21 [2],22 [3],23 [4],24 [5],25 [6],26 [7],27 [123]],28 [[8],29 [9],30 [10],31 [11],32 [12],33 [13],34 [14],35 [123]],36 [[15],37 [16],38 [17],39 [18],40 [19],41 [20],42 [21],43 [123]]]44 ]> : tensor<1x3x8x1xi32>45 46 %A_cast = tensor.cast %A : tensor<1x3x8x1xi32> to tensor<?x3x?x1xi32>47 func.call @unpack(%A_cast) : (tensor<?x3x?x1xi32>) -> ()48 49 return50}51 52func.func private @unpack(%A: tensor<?x3x?x1xi32>) {53 %c1 = arith.constant 1 : index54 %pad_val = arith.constant 123 : i3255 56 // Dynamic tile size57 %tile_size = arith.constant 8 : index58 %A_unpack_empty = tensor.empty() : tensor<7x3xi32>59 60 %A_unpack = linalg.unpack %A61 inner_dims_pos = [0, 1]62 inner_tiles = [%tile_size, 1]63 into %A_unpack_empty : tensor<?x3x?x1xi32> -> tensor<7x3xi32>64 %A_cast = tensor.cast %A_unpack : tensor<7x3xi32> to tensor<*xi32>65 66 // Print the results67 // CHECK: Unranked Memref base@ = 0x{{.*}} rank = 2 offset = 0 sizes = [7, 3] strides = [3, 1] data =68 // CHECK-NEXT: [1, 8, 15],69 // CHECK-NEXT: [2, 9, 16],70 // CHECK-NEXT: [3, 10, 17],71 // CHECK-NEXT: [4, 11, 18],72 // CHECK-NEXT: [5, 12, 19],73 // CHECK-NEXT: [6, 13, 20],74 // CHECK-NEXT: [7, 14, 21]75 call @printMemrefI32(%A_cast) : (tensor<*xi32>) -> ()76 77 return78}79 80module @transforms attributes { transform.with_named_sequence } {81 transform.named_sequence @__transform_main(%module: !transform.any_op {transform.consume}) {82 %pack = transform.structured.match ops{["linalg.unpack"]} in %module : (!transform.any_op) -> !transform.any_op83 84 // 1. Tile so that we can decompose linalg.pack85 // Ops (see step 2)86 %c8 = transform.param.constant 8 : i64 -> !transform.param<i64>87 %tiled_pack_op_p, %loops:2 = transform.structured.tile_using_for %pack tile_sizes [%c8, 1]88 : (!transform.any_op, !transform.param<i64>) -> (!transform.any_op, !transform.any_op, !transform.any_op)89 90 // 2. Decompose the tiled unpack Op into tensor.extract_slice + tensor.insert_slice:91 %func_op = transform.get_parent_op %tiled_pack_op_p {isolated_from_above} : (!transform.any_op) -> !transform.op<"func.func">92 transform.apply_patterns to %func_op {93 transform.apply_patterns.linalg.decompose_pack_unpack94 transform.apply_patterns.canonicalization95 } : !transform.op<"func.func">96 97 // 3. Vectorize tensor.insert_slice98 // Vector sizes match the inner tiles in the payload IR.99 %slice = transform.structured.match ops{["tensor.insert_slice"]} in %func_op : (!transform.op<"func.func">) -> !transform.any_op100 transform.structured.vectorize %slice vector_sizes [8, 1] : !transform.any_op101 102 // 4. Bufferize before lowering to LLVM103 %bufferize = transform.bufferization.one_shot_bufferize %module104 {bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op105 106 // 5. Canonicalize107 %func_op_bufferized = transform.structured.match ops{["func.func"]} in %bufferize : (!transform.any_op) -> !transform.op<"func.func">108 transform.apply_patterns to %func_op_bufferized {109 transform.apply_patterns.canonicalization110 } : !transform.op<"func.func">111 112 transform.yield113 }114}115 116func.func private @printMemrefI32(%ptr : tensor<*xi32>)117