xref: /llvm-project/mlir/test/Dialect/Arith/bufferize.mlir (revision ced2fc7819d5ddea616ec330f18e08ff284c1868)
1// RUN: mlir-opt %s --one-shot-bufferize="dialect-filter=arith,bufferization copy-before-write unknown-type-conversion=identity-layout-map" -split-input-file -verify-diagnostics | FileCheck %s
2
3// CHECK-LABEL:   func @index_cast(
4// CHECK-SAME:  %[[TENSOR:.*]]: tensor<i32>, %[[SCALAR:.*]]: i32
5func.func @index_cast(%tensor: tensor<i32>, %scalar: i32) -> (tensor<index>, index) {
6  %index_tensor = arith.index_cast %tensor : tensor<i32> to tensor<index>
7  %index_scalar = arith.index_cast %scalar : i32 to index
8  return %index_tensor, %index_scalar : tensor<index>, index
9}
10// CHECK:  %[[MEMREF:.*]] = bufferization.to_memref %[[TENSOR]] : tensor<i32>
11// CHECK-NEXT: %[[INDEX_MEMREF:.*]] = arith.index_cast %[[MEMREF]]
12// CHECK-SAME:   memref<i32> to memref<index>
13// CHECK-NEXT: %[[INDEX_TENSOR:.*]] = bufferization.to_tensor %[[INDEX_MEMREF]]
14// CHECK: return %[[INDEX_TENSOR]]
15
16// -----
17
18// CHECK-LABEL: module {
19
20// We check the debug name too since we put some effort into making that readable.
21// The name isn't load-bearing though.
22
23// CHECK: memref.global "private" constant @__constant_3x4xf32 : memref<3x4xf32> = dense<7.000000e+00>
24// CHECK-SAME: {alignment = 64 : i64}
25
26// CHECK: @basic
27func.func @basic() -> tensor<3x4xf32> {
28  // CHECK: %[[MEMREF:.*]] = memref.get_global @__constant_3x4xf32 : memref<3x4xf32>
29  // CHECK: %[[TENSOR:.*]] = bufferization.to_tensor %[[MEMREF]]
30  %0 = arith.constant dense<7.0> : tensor<3x4xf32>
31  // CHECK: return %[[TENSOR]]
32  return %0 : tensor<3x4xf32>
33}
34
35// CHECK: }
36
37// -----
38
39// CHECK-LABEL: module {
40
41// Only one global is created.
42// CHECK: memref.global
43// CHECK-NOT: memref.global
44func.func @duplicate_constants() -> (tensor<3x4xf32>, tensor<3x4xf32>) {
45  %0 = arith.constant dense<7.0> : tensor<3x4xf32>
46  %1 = arith.constant dense<7.0> : tensor<3x4xf32>
47  return %0, %1 : tensor<3x4xf32>, tensor<3x4xf32>
48}
49
50// CHECK: }
51
52// -----
53
54// CHECK-LABEL: module {
55
56// Two globals are created.
57// CHECK: memref.global
58// CHECK: memref.global
59// CHECK-NOT: memref.global
60func.func @multiple_constants() -> (tensor<3x4xf32>, tensor<3x4xf32>) {
61  %0 = arith.constant dense<7.0> : tensor<3x4xf32>
62  %1 = arith.constant dense<8.0> : tensor<3x4xf32>
63  return %0, %1 : tensor<3x4xf32>, tensor<3x4xf32>
64}
65
66// CHECK: }
67
68// -----
69
70// CHECK-LABEL: module {
71// We don't convert non-tensor globals.
72// CHECK-NOT: memref.global
73func.func @non_tensor() {
74    %0 = arith.constant 7 : i32
75    return
76}
77
78// CHECK: }
79
80// -----
81
82// CHECK-LABEL:   func @select(
83// CHECK-SAME:                 %[[PRED:.*]]: i1,
84// CHECK-SAME:                 %[[TRUE_VAL:.*]]: tensor<f32>,
85// CHECK-SAME:                 %[[FALSE_VAL:.*]]: tensor<f32>) -> tensor<f32> {
86// CHECK-DAG:           %[[TRUE_VAL_MEMREF:.*]] = bufferization.to_memref %[[TRUE_VAL]] : tensor<f32>
87// CHECK-DAG:           %[[FALSE_VAL_MEMREF:.*]] = bufferization.to_memref %[[FALSE_VAL]] : tensor<f32>
88// CHECK:           %[[RET_MEMREF:.*]] = arith.select %[[PRED]], %[[TRUE_VAL_MEMREF]], %[[FALSE_VAL_MEMREF]] : memref<f32>
89// CHECK:           %[[RET:.*]] = bufferization.to_tensor %[[RET_MEMREF]] : memref<f32>
90// CHECK:           return %[[RET]] : tensor<f32>
91func.func @select(%arg0: i1, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
92  %0 = arith.select %arg0, %arg1, %arg2 : tensor<f32>
93  return %0 : tensor<f32>
94}
95
96// -----
97
98func.func @elementwise_select(%arg0: tensor<5xi1>, %arg1: tensor<5xi32>, %arg2: tensor<5xi32>) -> tensor<5xi32> {
99  // expected-error @below{{only i1 condition values are supported}}
100  // expected-error @below{{failed to bufferize op}}
101  %0 = arith.select %arg0, %arg1, %arg2 : tensor<5xi1>, tensor<5xi32>
102  return %0 : tensor<5xi32>
103}
104