Skip to content

Commit

Permalink
Region dialect: update bufferization test
Browse files Browse the repository at this point in the history
  • Loading branch information
tkarna committed Dec 3, 2024
1 parent fe09a3e commit 379b666
Showing 1 changed file with 13 additions and 23 deletions.
36 changes: 13 additions & 23 deletions test/Dialect/Region/Transforms/RegionBufferize.mlir
Original file line number Diff line number Diff line change
@@ -1,31 +1,21 @@
// RUN: imex-opt %s -region-bufferize --split-input-file | FileCheck %s
// RUN: imex-opt %s -one-shot-bufferize --split-input-file | FileCheck %s

#map = affine_map<(d0) -> (d0)>
module {
func.func @sharpy_jit() -> memref<16xi64, strided<[?], offset: ?>> attributes {llvm.emit_c_interface} {
%cst = arith.constant 0.000000e+00 : f64
func.func @test_bufferize() -> memref<16xi64, strided<[?], offset: ?>> {
%c1_i64 = arith.constant 1 : i64
%0 = region.env_region #region.gpu_env<device = "gpu"> -> tensor<16xi64> {
%alloc = memref.alloc() {alignment = 64 : i64} : memref<16xi64>
linalg.generic {indexing_maps = [#map], iterator_types = ["parallel"]} outs(%alloc : memref<16xi64>) {
^bb0(%out: i64):
%3 = linalg.index 0 : index
%4 = arith.index_cast %3 : index to i64
%5 = arith.sitofp %4 : i64 to f64
%6 = arith.addf %5, %cst : f64
%7 = arith.fptosi %6 : f64 to i64
linalg.yield %7 : i64
}
%2 = bufferization.to_tensor %alloc : memref<16xi64>
region.env_region_yield %2 : tensor<16xi64>
%2 = bufferization.alloc_tensor() : tensor<16xi64>
%3 = linalg.fill ins(%c1_i64 : i64) outs(%2 : tensor<16xi64>) -> tensor<16xi64>
region.env_region_yield %3 : tensor<16xi64>
}
%1 = bufferization.to_memref %0 : memref<16xi64, strided<[?], offset: ?>>
return %1 : memref<16xi64, strided<[?], offset: ?>>
}
}
// CHECK-LABEL: func.func @sharpy_jit() -> memref<16xi64, strided<[?], offset: ?>> attributes {llvm.emit_c_interface} {
// CHECK: region.env_region #region.gpu_env<device = "gpu"> -> memref<16xi64> {
// CHECK: memref.alloc() {alignment = 64 : i64} : memref<16xi64>
// CHECK: linalg.generic {indexing_maps = [#map], iterator_types = ["parallel"]} outs(%alloc : memref<16xi64>) {
// CHECK: region.env_region_yield %4 : memref<16xi64>
// CHECK: return
// CHECK-SAME: : memref<16xi64, strided<[?], offset: ?>>
// CHECK-LABEL: func.func @test_bufferize() -> memref<16xi64, strided<[?], offset: ?>> {
// CHECK: [[R1:%.*]] = region.env_region #region.gpu_env<device = "gpu"> -> memref<16xi64> {
// CHECK-NEXT: [[V1:%.*]] = memref.alloc() {alignment = 64 : i64} : memref<16xi64>
// CHECK-NEXT: linalg.fill
// CHECK-NEXT: region.env_region_yield [[V1]] : memref<16xi64>
// CHECK: [[V2:%.*]] = memref.cast [[R1]] : memref<16xi64> to memref<16xi64, strided<[?], offset: ?>>
// CHECK-NEXT: return [[V2]] : memref<16xi64, strided<[?], offset: ?>>

0 comments on commit 379b666

Please sign in to comment.