19409bbb2Sbixia1// RUN: mlir-opt %s -split-input-file --sparse-buffer-rewrite --canonicalize --cse | FileCheck %s 2062e515bSbixia1 3654bbbdeSbixia1// CHECK-LABEL: func @sparse_push_back( 4988733c6SPeiming Liu// CHECK-SAME: %[[A:.*]]: index, 5654bbbdeSbixia1// CHECK-SAME: %[[B:.*]]: memref<?xf64>, 6988733c6SPeiming Liu// CHECK-SAME: %[[C:.*]]: f64) -> (memref<?xf64>, index) { 7654bbbdeSbixia1// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index 8654bbbdeSbixia1// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index 9654bbbdeSbixia1// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index 10d45be887Sbixia1// CHECK: %[[P1:.*]] = memref.dim %[[B]], %[[C0]] 11988733c6SPeiming Liu// CHECK: %[[S2:.*]] = arith.addi %[[A]], %[[C1]] : index 12d45be887Sbixia1// CHECK: %[[T:.*]] = arith.cmpi ugt, %[[S2]], %[[P1]] 13654bbbdeSbixia1// CHECK: %[[M:.*]] = scf.if %[[T]] -> (memref<?xf64>) { 14d45be887Sbixia1// CHECK: %[[P2:.*]] = arith.muli %[[P1]], %[[C2]] 15d45be887Sbixia1// CHECK: %[[M2:.*]] = memref.realloc %[[B]](%[[P2]]) 16654bbbdeSbixia1// CHECK: scf.yield %[[M2]] : memref<?xf64> 17654bbbdeSbixia1// CHECK: } else { 18654bbbdeSbixia1// CHECK: scf.yield %[[B]] : memref<?xf64> 19654bbbdeSbixia1// CHECK: } 20988733c6SPeiming Liu// CHECK: memref.store %[[C]], %[[M]]{{\[}}%[[A]]] 21988733c6SPeiming Liu// CHECK: return %[[M]], %[[S2]] 22988733c6SPeiming Liufunc.func @sparse_push_back(%arg0: index, %arg1: memref<?xf64>, %arg2: f64) -> (memref<?xf64>, index) { 23988733c6SPeiming Liu %0:2 = sparse_tensor.push_back %arg0, %arg1, %arg2 : index, memref<?xf64>, f64 24988733c6SPeiming Liu return %0#0, %0#1 : memref<?xf64>, index 25654bbbdeSbixia1} 26654bbbdeSbixia1 279409bbb2Sbixia1// ----- 289409bbb2Sbixia1 29d45be887Sbixia1// CHECK-LABEL: func @sparse_push_back_n( 30988733c6SPeiming Liu// CHECK-SAME: %[[S1:.*]]: index, 31d45be887Sbixia1// CHECK-SAME: %[[B:.*]]: memref<?xf64>, 32d45be887Sbixia1// CHECK-SAME: %[[C:.*]]: f64, 33988733c6SPeiming Liu// CHECK-SAME: %[[D:.*]]: index) -> (memref<?xf64>, index) { 34d45be887Sbixia1// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index 35d45be887Sbixia1// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index 36d45be887Sbixia1// CHECK: %[[P1:.*]] = memref.dim %[[B]], %[[C0]] 37d45be887Sbixia1// CHECK: %[[S2:.*]] = arith.addi %[[S1]], %[[D]] : index 38d45be887Sbixia1// CHECK: %[[T:.*]] = arith.cmpi ugt, %[[S2]], %[[P1]] 39d45be887Sbixia1// CHECK: %[[M:.*]] = scf.if %[[T]] -> (memref<?xf64>) { 40d45be887Sbixia1// CHECK: %[[P2:.*]] = scf.while (%[[I:.*]] = %[[P1]]) : (index) -> index { 41d45be887Sbixia1// CHECK: %[[P3:.*]] = arith.muli %[[I]], %[[C2]] : index 42d45be887Sbixia1// CHECK: %[[T2:.*]] = arith.cmpi ugt, %[[S2]], %[[P3]] : index 43d45be887Sbixia1// CHECK: scf.condition(%[[T2]]) %[[P3]] : index 44d45be887Sbixia1// CHECK: } do { 45d45be887Sbixia1// CHECK: ^bb0(%[[I2:.*]]: index): 46d45be887Sbixia1// CHECK: scf.yield %[[I2]] : index 47d45be887Sbixia1// CHECK: } 48d45be887Sbixia1// CHECK: %[[M2:.*]] = memref.realloc %[[B]](%[[P2]]) 49d45be887Sbixia1// CHECK: scf.yield %[[M2]] : memref<?xf64> 50d45be887Sbixia1// CHECK: } else { 51d45be887Sbixia1// CHECK: scf.yield %[[B]] : memref<?xf64> 52d45be887Sbixia1// CHECK: } 53d45be887Sbixia1// CHECK: %[[S:.*]] = memref.subview %[[M]]{{\[}}%[[S1]]] {{\[}}%[[D]]] [1] 54d45be887Sbixia1// CHECK: linalg.fill ins(%[[C]] : f64) outs(%[[S]] 55988733c6SPeiming Liu// CHECK: return %[[M]], %[[S2]] : memref<?xf64>, index 56988733c6SPeiming Liufunc.func @sparse_push_back_n(%arg0: index, %arg1: memref<?xf64>, %arg2: f64, %arg3: index) -> (memref<?xf64>, index) { 57988733c6SPeiming Liu %0:2 = sparse_tensor.push_back %arg0, %arg1, %arg2, %arg3 : index, memref<?xf64>, f64, index 58988733c6SPeiming Liu return %0#0, %0#1 : memref<?xf64>, index 59d45be887Sbixia1} 60d45be887Sbixia1 61d45be887Sbixia1// ----- 62d45be887Sbixia1 631c835b5aSbixia1// CHECK-LABEL: func @sparse_push_back_inbound( 64988733c6SPeiming Liu// CHECK-SAME: %[[S1:.*]]: index, 651c835b5aSbixia1// CHECK-SAME: %[[B:.*]]: memref<?xf64>, 66988733c6SPeiming Liu// CHECK-SAME: %[[C:.*]]: f64) -> (memref<?xf64>, index) { 671c835b5aSbixia1// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index 68d45be887Sbixia1// CHECK: %[[S2:.*]] = arith.addi %[[S1]], %[[C1]] 69d45be887Sbixia1// CHECK: memref.store %[[C]], %[[B]]{{\[}}%[[S1]]] 70988733c6SPeiming Liu// CHECK: return %[[B]], %[[S2]] : memref<?xf64>, index 71988733c6SPeiming Liufunc.func @sparse_push_back_inbound(%arg0: index, %arg1: memref<?xf64>, %arg2: f64) -> (memref<?xf64>, index) { 72988733c6SPeiming Liu %0:2 = sparse_tensor.push_back inbounds %arg0, %arg1, %arg2 : index, memref<?xf64>, f64 73988733c6SPeiming Liu return %0#0, %0#1 : memref<?xf64>, index 741c835b5aSbixia1} 751c835b5aSbixia1 769409bbb2Sbixia1// ----- 779409bbb2Sbixia1 78bfa3bc43SPeiming Liu#ID_MAP=affine_map<(d0, d1) -> (d0, d1)> 799409bbb2Sbixia1 809409bbb2Sbixia1// Only check the generated supporting functions. We have integration test to 819409bbb2Sbixia1// verify correctness of the generated code. 829409bbb2Sbixia1// 83bfa3bc43SPeiming Liu// CHECK-DAG: func.func private @_sparse_partition_0_1_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>) -> index { 84bfa3bc43SPeiming Liu// CHECK-DAG: func.func private @_sparse_qsort_0_1_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>) { 85a1507668Sbixia1// CHECK-LABEL: func.func @sparse_sort_coo_quick 86a1507668Sbixia1func.func @sparse_sort_coo_quick(%arg0: index, %arg1: memref<100xindex>, %arg2: memref<?xf32>, %arg3: memref<10xi32>) -> (memref<100xindex>, memref<?xf32>, memref<10xi32>) { 87*0083f833SPeiming Liu sparse_tensor.sort quick_sort %arg0, %arg1 jointly %arg2, %arg3 {perm_map = #ID_MAP, ny = 1: index} : memref<100xindex> jointly memref<?xf32>, memref<10xi32> 88a1507668Sbixia1 return %arg1, %arg2, %arg3 : memref<100xindex>, memref<?xf32>, memref<10xi32> 89a1507668Sbixia1} 90a1507668Sbixia1 91a1507668Sbixia1// ----- 92a1507668Sbixia1 93bfa3bc43SPeiming Liu#ID_MAP=affine_map<(d0, d1) -> (d0, d1)> 94bfa3bc43SPeiming Liu 95a1507668Sbixia1// Only check the generated supporting functions. We have integration test to 96a1507668Sbixia1// verify correctness of the generated code. 97a1507668Sbixia1// 98bfa3bc43SPeiming Liu// CHECK-DAG: func.func private @_sparse_binary_search_0_1_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>) -> index { 99bfa3bc43SPeiming Liu// CHECK-DAG: func.func private @_sparse_sort_stable_0_1_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>) { 100bfa3bc43SPeiming Liu// CHECK-DAG: func.func private @_sparse_shift_down_0_1_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>, %arg5: index) { 101bfa3bc43SPeiming Liu// CHECK-DAG: func.func private @_sparse_heap_sort_0_1_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>) { 102bfa3bc43SPeiming Liu// CHECK-DAG: func.func private @_sparse_partition_0_1_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>) -> index { 103bfa3bc43SPeiming Liu// CHECK-DAG: func.func private @_sparse_hybrid_qsort_0_1_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>, %arg5: i64) { 104a1507668Sbixia1// CHECK-LABEL: func.func @sparse_sort_coo_hybrid 105a1507668Sbixia1func.func @sparse_sort_coo_hybrid(%arg0: index, %arg1: memref<100xindex>, %arg2: memref<?xf32>, %arg3: memref<10xi32>) -> (memref<100xindex>, memref<?xf32>, memref<10xi32>) { 106*0083f833SPeiming Liu sparse_tensor.sort hybrid_quick_sort %arg0, %arg1 jointly %arg2, %arg3 {perm_map = #ID_MAP, ny = 1: index} : memref<100xindex> jointly memref<?xf32>, memref<10xi32> 1074f729d5aSbixia1 return %arg1, %arg2, %arg3 : memref<100xindex>, memref<?xf32>, memref<10xi32> 1084f729d5aSbixia1} 1094f729d5aSbixia1 1104f729d5aSbixia1// ----- 1114f729d5aSbixia1 112bfa3bc43SPeiming Liu#ID_MAP=affine_map<(d0, d1) -> (d0, d1)> 113bfa3bc43SPeiming Liu 1144f729d5aSbixia1// Only check the generated supporting functions. We have integration test to 1154f729d5aSbixia1// verify correctness of the generated code. 1164f729d5aSbixia1// 117bfa3bc43SPeiming Liu// CHECK-DAG: func.func private @_sparse_binary_search_0_1_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>) -> index { 118bfa3bc43SPeiming Liu// CHECK-DAG: func.func private @_sparse_sort_stable_0_1_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>) { 1194f729d5aSbixia1// CHECK-LABEL: func.func @sparse_sort_coo_stable 1204f729d5aSbixia1func.func @sparse_sort_coo_stable(%arg0: index, %arg1: memref<100xindex>, %arg2: memref<?xf32>, %arg3: memref<10xi32>) -> (memref<100xindex>, memref<?xf32>, memref<10xi32>) { 121*0083f833SPeiming Liu sparse_tensor.sort insertion_sort_stable %arg0, %arg1 jointly %arg2, %arg3 {perm_map = #ID_MAP, ny = 1: index} : memref<100xindex> jointly memref<?xf32>, memref<10xi32> 1224f729d5aSbixia1 return %arg1, %arg2, %arg3 : memref<100xindex>, memref<?xf32>, memref<10xi32> 1234f729d5aSbixia1} 1244f729d5aSbixia1 1253b1c86cdSbixia1// ----- 1263b1c86cdSbixia1 127bfa3bc43SPeiming Liu#ID_MAP=affine_map<(d0, d1) -> (d0, d1)> 128bfa3bc43SPeiming Liu 1293b1c86cdSbixia1// Only check the generated supporting functions. We have integration test to 1303b1c86cdSbixia1// verify correctness of the generated code. 1313b1c86cdSbixia1// 132bfa3bc43SPeiming Liu// CHECK-DAG: func.func private @_sparse_shift_down_0_1_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>, %arg5: index) { 133bfa3bc43SPeiming Liu// CHECK-DAG: func.func private @_sparse_heap_sort_0_1_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>) { 1343b1c86cdSbixia1// CHECK-LABEL: func.func @sparse_sort_coo_heap 1353b1c86cdSbixia1func.func @sparse_sort_coo_heap(%arg0: index, %arg1: memref<100xindex>, %arg2: memref<?xf32>, %arg3: memref<10xi32>) -> (memref<100xindex>, memref<?xf32>, memref<10xi32>) { 136*0083f833SPeiming Liu sparse_tensor.sort heap_sort %arg0, %arg1 jointly %arg2, %arg3 {perm_map = #ID_MAP, ny = 1: index} : memref<100xindex> jointly memref<?xf32>, memref<10xi32> 1373b1c86cdSbixia1 return %arg1, %arg2, %arg3 : memref<100xindex>, memref<?xf32>, memref<10xi32> 1383b1c86cdSbixia1} 139