1// RUN: mlir-opt %s -split-input-file -canonicalize="test-convergence" | FileCheck %s 2 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_canonicalize 9// 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 : index 19 %c1 = arith.constant 1 : index 20 %c4 = arith.constant 4 : index 21 %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_canonicalize 33// 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 2 36 : 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 return 41} 42