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