42 lines · plain
1// RUN: mlir-opt %s -split-input-file -canonicalize="test-convergence" | FileCheck %s2 3#BCOO = #sparse_tensor.encoding<{4 map = (d0, d1, d2) -> (d0 : dense, d1 : loose_compressed(nonunique), d2 : singleton)5}>6 7// CHECK-DAG: #[[$BCOO:.*]] = #sparse_tensor.encoding<{ map = (d0, d1, d2) -> (d0 : dense, d1 : loose_compressed(nonunique), d2 : singleton) }>8// CHECK-LABEL: func @sparse_slice_canonicalize9// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?xf32, #[[$BCOO]]>10// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[ARG0]][0, %{{[a-zA-Z0-9_]+}}, 1]11// CHECK-SAME: [4, 1, %{{[a-zA-Z0-9_]+}}] [1, 1, 1]12// CHECK-SAME: : tensor<?x?x?xf32, #[[$BCOO]]> to tensor<4x1x?xf32, #[[$BCOO]]>13// CHECK: %[[RESULT:.+]] = tensor.cast %[[SLICE]]14// CHECK: return %[[RESULT]]15func.func @sparse_slice_canonicalize(%arg0 : tensor<?x?x?xf32, #BCOO>, %arg1 : index,16 %arg2 : index) -> tensor<?x?x?xf32, #BCOO>17{18 %c0 = arith.constant 0 : index19 %c1 = arith.constant 1 : index20 %c4 = arith.constant 4 : index21 %0 = tensor.extract_slice %arg0[%c0, %arg1, %c1] [%c4, %c1, %arg2] [%c1, %c1, %c1] : tensor<?x?x?xf32, #BCOO> to tensor<?x?x?xf32, #BCOO>22 return %0 : tensor<?x?x?xf32, #BCOO>23}24 25// -----26 27#CSR = #sparse_tensor.encoding<{28 map = (i, j) -> (i : dense, j : compressed)29}>30 31// Make sure that the first unused coordinate is optimized.32// CHECK-LABEL: @sparse_iterate_canonicalize33// CHECK: sparse_tensor.iterate {{.*}} at(_, %{{.*}})34func.func @sparse_iterate_canonicalize(%sp : tensor<?x?xf64, #CSR>) {35 %l1 = sparse_tensor.extract_iteration_space %sp lvls = 0 to 236 : tensor<?x?xf64, #CSR> -> !sparse_tensor.iter_space<#CSR, lvls = 0 to 2>37 sparse_tensor.iterate %it1 in %l1 at (%coord0, %coord1) : !sparse_tensor.iter_space<#CSR, lvls = 0 to 2> {38 "test.op"(%coord1) : (index) -> ()39 }40 return41}42