123 lines · plain
1// RUN: mlir-opt %s --pre-sparsification-rewrite --sparse-reinterpret-map | FileCheck %s --check-prefix=CHECK-FOLD2// RUN: mlir-opt %s --pre-sparsification-rewrite --sparse-reinterpret-map --sparsification | FileCheck %s3 4#trait = {5 indexing_maps = [6 affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>,7 affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>,8 affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>,9 affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>10 ],11 iterator_types = ["parallel", "parallel", "parallel", "parallel"]12}13 14#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>15 16#COO = #sparse_tensor.encoding<{map = (d0, d1, d2) -> (d0 : compressed(nonunique), d1 : singleton(nonunique, soa), d2 : singleton(soa))}>17#CCCD = #sparse_tensor.encoding<{ map = (d0, d1, d2, d3) -> (d0 : compressed, d1 : compressed, d2 : compressed, d3 : dense) }>18 19// CHECK-LABEL: func.func @fold_convert(20// CHECK: scf.for21// CHECK: scf.for22// CHECK: scf.for23// CHECK: scf.if24// CHECK-NEXT: tensor.insert25// CHECK-NEXT: scf.yield26// CHECK-NEXT: else27// CHECK-NEXT: scf.yield28// CHECK: scf.yield29// CHECK: scf.yield30// CHECK: scf.yield31// CHECK: sparse_tensor.load32 33// CHECK-FOLD-LABEL: func.func @fold_convert(34// CHECK-FOLD-NOT: sparse_tensor.convert35func.func @fold_convert(%arg0: tensor<128x32x32x1xf32>, %arg1: tensor<128x32x32x1xf32>, %arg2: tensor<128x32x32x1xf32>) -> tensor<128x32x32x1xf32, #CCCD> {36 %cst = arith.constant 0.000000e+00 : f3237 %cst_0 = arith.constant 1.000000e+00 : f3238 %cst_1 = arith.constant 1.000000e+00 : f3239 %0 = tensor.empty() : tensor<128x32x32x1xf32>40 %1 = linalg.generic #trait41 ins(%arg0, %arg1, %arg2 : tensor<128x32x32x1xf32>, tensor<128x32x32x1xf32>, tensor<128x32x32x1xf32>)42 outs(%0 : tensor<128x32x32x1xf32>) {43 ^bb0(%in: f32, %in_2: f32, %in_3: f32, %out: f32):44 %3 = arith.subf %cst_0, %in_2 : f3245 %4 = arith.mulf %in, %3 : f3246 %5 = arith.mulf %4, %cst_1 : f3247 %6 = arith.addf %5, %in_3 : f3248 %7 = arith.subf %6, %cst_0 : f3249 %8 = arith.cmpf uge, %7, %cst : f3250 %9 = arith.uitofp %8 : i1 to f3251 linalg.yield %9 : f3252 } -> tensor<128x32x32x1xf32>53 %2 = sparse_tensor.convert %1 : tensor<128x32x32x1xf32> to tensor<128x32x32x1xf32, #CCCD>54 return %2 : tensor<128x32x32x1xf32, #CCCD>55}56 57#trait_bin = {58 indexing_maps = [59 affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>,60 affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>,61 affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>62 ],63 iterator_types = ["parallel", "parallel", "parallel", "parallel"]64}65 66// CHECK-FOLD-LABEL: func.func @fold_convert_multi_use(67// CHECK-FOLD: tensor.empty() : tensor<128x32x32x1xf32>68// CHECK-FOLD: linalg.generic69// CHECK-FOLD: tensor.empty() : tensor<128x32x32x1xf32, #sparse>70// CHECK-FOLD: linalg.generic71// CHECK-FOLD-NOT: sparse_tensor.convert72func.func @fold_convert_multi_use(%arg0: tensor<128x32x32x1xf32>, %arg1: tensor<128x32x32x1xf32>,73 %arg2: tensor<128x32x32x1xf32>, %arg3: tensor<128x32x32x1xf32>) -> (tensor<128x32x32x1xf32>, tensor<128x32x32x1xf32, #CCCD>) {74 %cst = arith.constant 0.000000e+00 : f3275 %cst_0 = arith.constant 1.000000e+00 : f3276 %cst_1 = arith.constant 1.000000e+00 : f3277 78 %0 = tensor.empty() : tensor<128x32x32x1xf32>79 %1 = linalg.generic #trait_bin80 ins(%arg0, %arg1 : tensor<128x32x32x1xf32>, tensor<128x32x32x1xf32>)81 outs(%0 : tensor<128x32x32x1xf32>) {82 ^bb0(%in: f32, %in_1: f32, %out: f32):83 %3 = arith.mulf %in, %in_1 : f3284 linalg.yield %3 : f3285 } -> tensor<128x32x32x1xf32>86 87 // A second kernel that uses %0 as the init operand.88 %3 = linalg.generic #trait_bin89 ins(%arg2, %arg3 : tensor<128x32x32x1xf32>, tensor<128x32x32x1xf32>)90 outs(%0 : tensor<128x32x32x1xf32>) {91 ^bb0(%in: f32, %in_1: f32, %out: f32):92 %3 = arith.mulf %in, %in_1 : f3293 linalg.yield %3 : f3294 } -> tensor<128x32x32x1xf32>95 %4 = sparse_tensor.convert %3 : tensor<128x32x32x1xf32> to tensor<128x32x32x1xf32, #CCCD>96 97 return %1, %4 : tensor<128x32x32x1xf32>, tensor<128x32x32x1xf32, #CCCD>98}99 100 101 102// FIXME: The following kernel is not sparsifiable because `arith.select`103// operations is not handled by the sparse compiler at the moment.104//105// CHECK-FOLD-LABEL: func.func @fold_cast(106// CHECK-FOLD-NOT: sparse_tensor.convert107func.func @fold_cast(%0: tensor<10x20x30xf64, #COO>) -> tensor<10x20x30xf64, #COO> {108 %cst = arith.constant 0.000000e+00 : f64109 %1 = tensor.empty() : tensor<10x20x30xf64>110 %2 = linalg.generic { indexing_maps = [#map, #map],111 iterator_types = ["parallel", "parallel", "parallel"]112 }113 ins (%0 : tensor<10x20x30xf64, #COO>)114 outs(%1 : tensor<10x20x30xf64>) {115 ^bb0(%in: f64, %out: f64):116 %4 = arith.cmpf ugt, %in, %cst : f64117 %5 = arith.select %4, %in, %cst : f64118 linalg.yield %5 : f64119 } -> tensor<10x20x30xf64>120 %cast = tensor.cast %2 : tensor<10x20x30xf64> to tensor<10x20x30xf64, #COO>121 return %cast : tensor<10x20x30xf64, #COO>122}123