summaryrefslogtreecommitdiff
path: root/mlir/test/Dialect/Bufferization/ops.mlir
blob: ddb597a334d814bec5fc503be609e409afc610da (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
// RUN: mlir-opt %s | mlir-opt | FileCheck %s
// RUN: mlir-opt %s --mlir-print-op-generic | mlir-opt | FileCheck %s

#CSR = #sparse_tensor.encoding<{
  lvlTypes = ["dense", "compressed"]
}>

// CHECK-LABEL: func @test_clone
func.func @test_clone(%buf : memref<*xf32>) -> memref<*xf32> {
  %clone = bufferization.clone %buf : memref<*xf32> to memref<*xf32>
  return %clone : memref<*xf32>
}

// CHECK-LABEL: test_to_memref
func.func @test_to_memref(%arg0: tensor<?xi64>, %arg1: tensor<*xi64>)
    -> (memref<?xi64, affine_map<(d0) -> (d0 + 7)>>, memref<*xi64, 1>) {
  %0 = bufferization.to_memref %arg0
    : memref<?xi64, affine_map<(d0) -> (d0 + 7)>>
  %1 = bufferization.to_memref %arg1
    : memref<*xi64, 1>
  return %0, %1 : memref<?xi64, affine_map<(d0) -> (d0 + 7)>>, memref<*xi64, 1>
}

// CHECK-LABEL: func @test_to_tensor
func.func @test_to_tensor(%buf : memref<2xf32>) -> tensor<2xf32> {
  %tensor = bufferization.to_tensor %buf restrict writable : memref<2xf32>
  return %tensor : tensor<2xf32>
}

// CHECK-LABEL: func @test_alloc_tensor_op
func.func @test_alloc_tensor_op(%t: tensor<?x5xf32>, %sz: index)
  -> tensor<?x5xf32>
{
  // CHECK: bufferization.alloc_tensor(%{{.*}}) : tensor<?x5xf32>
  %0 = bufferization.alloc_tensor(%sz) : tensor<?x5xf32>
  // CHECK: bufferization.alloc_tensor() copy(%{{.*}}) : tensor<?x5xf32>
  %1 = bufferization.alloc_tensor() copy(%t) : tensor<?x5xf32>
  // CHECK: bufferization.alloc_tensor() : tensor<5x6xf32>
  %2 = bufferization.alloc_tensor() : tensor<5x6xf32>
  // CHECK: bufferization.alloc_tensor(%{{.*}}, %{{.*}}) : tensor<?x?xf32>
  %3 = bufferization.alloc_tensor(%sz, %sz) : tensor<?x?xf32>
  // CHECK: bufferization.alloc_tensor() copy(%{{.*}}) {escape = true} : tensor<?x5xf32>
  %4 = bufferization.alloc_tensor() copy(%t) {escape = true} : tensor<?x5xf32>
  // CHECK: bufferization.alloc_tensor() copy(%{{.*}}) {escape = false} : tensor<?x5xf32>
  %5 = bufferization.alloc_tensor() copy(%t) {escape = false} : tensor<?x5xf32>
  %c100 = arith.constant 100 : index
  // CHECK: bufferization.alloc_tensor() size_hint=
  %6 = bufferization.alloc_tensor() size_hint=%c100 : tensor<100x100xf64, #CSR>
  // CHECK: bufferization.alloc_tensor(%{{.+}}) {memory_space = "foo"} : tensor<?xf32>
  %7 = bufferization.alloc_tensor(%sz) {memory_space = "foo"} : tensor<?xf32>
  return %1 : tensor<?x5xf32>
}

// CHECK-LABEL: func @test_dealloc_tensor_op
func.func @test_dealloc_tensor_op(%arg0: tensor<4xi32>) {
  // CHECK: bufferization.dealloc_tensor {{.*}} : tensor<4xi32>
  bufferization.dealloc_tensor %arg0 : tensor<4xi32>
  return
}