133 lines · plain
1// RUN: mlir-opt %s -sparse-tensor-codegen -cse | FileCheck %s2 3#SparseVector = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>4 5// CHECK-LABEL: func.func @for(6// CHECK-SAME: %[[VAL_1:.*0]]: memref<?xindex>,7// CHECK-SAME: %[[VAL_2:.*1]]: memref<?xindex>,8// CHECK-SAME: %[[VAL_3:.*2]]: memref<?xf32>,9// CHECK-SAME: %[[VAL_4:.*3]]: !sparse_tensor.storage_specifier10// CHECK-SAME: %[[VAL_5:.*4]]: index,11// CHECK-SAME: %[[VAL_6:.*5]]: index,12// CHECK-SAME: %[[VAL_7:.*6]]: index) -> (memref<?xindex>, memref<?xindex>, memref<?xf32>, !sparse_tensor.storage_specifier13// CHECK: %[[VAL_8:.*]]:4 = scf.for %[[VAL_9:.*]] = %[[VAL_5]] to %[[VAL_6]] step %[[VAL_7]] iter_args(14// CHECK-SAME: %[[VAL_11:.*]] = %[[VAL_1]],15// CHECK-SAME: %[[VAL_12:.*]] = %[[VAL_2]],16// CHECK-SAME: %[[VAL_13:.*]] = %[[VAL_3]],17// CHECK-SAME: %[[VAL_14:.*]] = %[[VAL_4]])18// CHECK: scf.yield %[[VAL_11]], %[[VAL_12]], %[[VAL_13]], %[[VAL_14]] :19// CHECK: }20// CHECK: return %[[VAL_8]]#0, %[[VAL_8]]#1, %[[VAL_8]]#2, %[[VAL_8]]#321func.func @for(%in: tensor<1024xf32, #SparseVector>,22 %lb: index, %ub: index, %step: index) -> tensor<1024xf32, #SparseVector> {23 %1 = scf.for %i = %lb to %ub step %step iter_args(%vin = %in)24 -> tensor<1024xf32, #SparseVector> {25 scf.yield %vin : tensor<1024xf32, #SparseVector>26 }27 return %1 : tensor<1024xf32, #SparseVector>28}29 30// CHECK-LABEL: func.func @if(31// CHECK-SAME: %[[VAL_1:.*0]]: memref<?xindex>,32// CHECK-SAME: %[[VAL_2:.*1]]: memref<?xindex>,33// CHECK-SAME: %[[VAL_3:.*2]]: memref<?xf32>,34// CHECK-SAME: %[[VAL_4:.*3]]: !sparse_tensor.storage_specifier35// CHECK-SAME: %[[VAL_6:.*4]]: memref<?xindex>,36// CHECK-SAME: %[[VAL_7:.*5]]: memref<?xindex>,37// CHECK-SAME: %[[VAL_8:.*6]]: memref<?xf32>,38// CHECK-SAME: %[[VAL_9:.*7]]: !sparse_tensor.storage_specifier39// CHECK-SAME: %[[VAL_10:.*]]: i1)40// CHECK: %[[VAL_11:.*]]:4 = scf.if %[[VAL_10]]41// CHECK: scf.yield %[[VAL_1]], %[[VAL_2]], %[[VAL_3]], %[[VAL_4]]42// CHECK: } else {43// CHECK: scf.yield %[[VAL_6]], %[[VAL_7]], %[[VAL_8]], %[[VAL_9]]44// CHECK: }45// CHECK: return %[[VAL_11]]#0, %[[VAL_11]]#1, %[[VAL_11]]#2, %[[VAL_11]]#3 :46// CHECK-SAME: memref<?xindex>, memref<?xindex>, memref<?xf32>, !sparse_tensor.storage_specifier47func.func @if(%t: tensor<1024xf32, #SparseVector>,48 %f: tensor<1024xf32, #SparseVector>,49 %c: i1) -> tensor<1024xf32, #SparseVector> {50 %1 = scf.if %c -> tensor<1024xf32, #SparseVector> {51 scf.yield %t : tensor<1024xf32, #SparseVector>52 } else {53 scf.yield %f : tensor<1024xf32, #SparseVector>54 }55 return %1 : tensor<1024xf32, #SparseVector>56}57 58 59// CHECK-LABEL: func.func @while(60// CHECK-SAME: %[[VAL_1:.*0]]: memref<?xindex>,61// CHECK-SAME: %[[VAL_2:.*1]]: memref<?xindex>,62// CHECK-SAME: %[[VAL_3:.*2]]: memref<?xf32>,63// CHECK-SAME: %[[VAL_4:.*3]]: !sparse_tensor.storage_specifier64// CHECK-SAME: %[[VAL_5:.*4]]: i1)65// CHECK: %[[VAL_6:.*]]:4 = scf.while (66// CHECK-SAME: %[[VAL_8:.*]] = %[[VAL_1]],67// CHECK-SAME: %[[VAL_9:.*]] = %[[VAL_2]],68// CHECK-SAME: %[[VAL_10:.*]] = %[[VAL_3]],69// CHECK-SAME: %[[VAL_11:.*]] = %[[VAL_4]])70// CHECK: scf.condition(%[[VAL_5]]) %[[VAL_8]], %[[VAL_9]], %[[VAL_10]], %[[VAL_11]]71// CHECK: } do {72// CHECK: ^bb0(%[[VAL_13:.*5]]: memref<?xindex>,73// CHECK-SAME: %[[VAL_14:.*6]]: memref<?xindex>,74// CHECK-SAME: %[[VAL_15:.*7]]: memref<?xf32>,75// CHECK-SAME: %[[VAL_16:.*8]]: !sparse_tensor.storage_specifier76// CHECK: scf.yield %[[VAL_13]], %[[VAL_14]], %[[VAL_15]], %[[VAL_16]]77// CHECK: }78// CHECK: return %[[VAL_6]]#0, %[[VAL_6]]#1, %[[VAL_6]]#2, %[[VAL_6]]#3 :79// CHECK-SAME: memref<?xindex>, memref<?xindex>, memref<?xf32>, !sparse_tensor.storage_specifier80func.func @while(%arg0: tensor<1024xf32, #SparseVector>, %c: i1) -> tensor<1024xf32, #SparseVector> {81 %0 = scf.while (%in = %arg0) : (tensor<1024xf32, #SparseVector>) -> tensor<1024xf32, #SparseVector> {82 scf.condition(%c) %in : tensor<1024xf32, #SparseVector>83 } do {84 ^bb0(%arg1: tensor<1024xf32, #SparseVector>):85 scf.yield %arg1 : tensor<1024xf32, #SparseVector>86 }87 return %0: tensor<1024xf32, #SparseVector>88}89 90// CHECK-LABEL: func.func @index_switch(91// CHECK-SAME: %[[PRED:.*0]]: index,92// CHECK-SAME: %[[VAL_A_1:.*1]]: memref<?xindex>,93// CHECK-SAME: %[[VAL_A_2:.*2]]: memref<?xindex>,94// CHECK-SAME: %[[VAL_A_3:.*3]]: memref<?xf32>,95// CHECK-SAME: %[[VAL_A_4:.*4]]: !sparse_tensor.storage_specifier96// CHECK-SAME: %[[VAL_B_1:.*5]]: memref<?xindex>,97// CHECK-SAME: %[[VAL_B_2:.*6]]: memref<?xindex>,98// CHECK-SAME: %[[VAL_B_3:.*7]]: memref<?xf32>,99// CHECK-SAME: %[[VAL_B_4:.*8]]: !sparse_tensor.storage_specifier100// CHECK-SAME: %[[VAL_C_1:.*9]]: memref<?xindex>,101// CHECK-SAME: %[[VAL_C_2:.*10]]: memref<?xindex>,102// CHECK-SAME: %[[VAL_C_3:.*11]]: memref<?xf32>,103// CHECK-SAME: %[[VAL_C_4:.*12]]: !sparse_tensor.storage_specifier104 105// CHECK: %[[RES:.*]]:4 = scf.index_switch %[[PRED]]106// CHECK-SAME: -> memref<?xindex>, memref<?xindex>, memref<?xf32>, !sparse_tensor.storage_specifier107// CHECK: case 1 {108// CHECK: scf.yield %[[VAL_A_1]], %[[VAL_A_2]], %[[VAL_A_3]], %[[VAL_A_4]]109// CHECK: case 2 {110// CHECK: scf.yield %[[VAL_B_1]], %[[VAL_B_2]], %[[VAL_B_3]], %[[VAL_B_4]]111// CHECK: default {112// CHECK: scf.yield %[[VAL_C_1]], %[[VAL_C_2]], %[[VAL_C_3]], %[[VAL_C_4]]113 114// CHECK: return %[[RES]]#0, %[[RES]]#1, %[[RES]]#2, %[[RES]]#3 :115// CHECK-SAME: memref<?xindex>, memref<?xindex>, memref<?xf32>, !sparse_tensor.storage_specifier116 117func.func @index_switch(%pred: index, %a: tensor<5xf32, #SparseVector>,118 %b: tensor<5xf32, #SparseVector>,119 %c: tensor<5xf32, #SparseVector>) -> tensor<5xf32, #SparseVector> {120 %0 = scf.index_switch %pred -> tensor<5xf32, #SparseVector>121 case 1 {122 scf.yield %a : tensor<5xf32, #SparseVector>123 }124 case 2 {125 scf.yield %b : tensor<5xf32, #SparseVector>126 }127 default {128 scf.yield %c : tensor<5xf32, #SparseVector>129 }130 131 return %0 : tensor<5xf32, #SparseVector>132}133