test/Dialect/OpenACC/acc-implicit-data.mlir

// RUN: mlir-opt %s -acc-implicit-data -split-input-file | FileCheck %s

// -----

// acc.kernel_environment is not a compute construct. Host-side preparation in
// its body and operands captured by a nested acc.compute_region must not cause
// implicit mappings.
func.func @test_kernel_environment_is_not_compute_construct(
    %arg: memref<f32>) {
  acc.kernel_environment {
    %host = memref.load %arg[] : memref<f32>
    acc.compute_region ins(%deviceArg = %arg) : (memref<f32>) {
      %device = memref.load %deviceArg[] : memref<f32>
      acc.yield
    } <{origin = "acc.parallel"}>
  }
  return
}

// CHECK-LABEL: func.func @test_kernel_environment_is_not_compute_construct
// CHECK-SAME: %[[ARG:.*]]: memref<f32>
// CHECK-NOT: acc.copyin
// CHECK-NOT: acc.firstprivate
// CHECK: acc.kernel_environment {
// CHECK: memref.load %[[ARG]][]
// CHECK: acc.compute_region ins(%{{.*}} = %[[ARG]]) : (memref<f32>)

// -----

// Test scalar in serial construct - should generate firstprivate
func.func @test_scalar_in_serial() {
  %alloc = memref.alloca() : memref<i64>
  acc.serial {
    %load = memref.load %alloc[] : memref<i64>
    acc.yield
  }
  return
}

// CHECK-LABEL: func.func @test_scalar_in_serial
// CHECK: acc.firstprivate varPtr({{.*}} : memref<i64>) recipe({{.*}}) implicit(true) name("") -> memref<i64>

// -----

// Test scalar in parallel construct - should generate firstprivate
func.func @test_scalar_in_parallel() {
  %alloc = memref.alloca() : memref<f32>
  acc.parallel {
    %load = memref.load %alloc[] : memref<f32>
    acc.yield
  }
  return
}

// CHECK-LABEL: func.func @test_scalar_in_parallel
// CHECK: acc.firstprivate varPtr({{.*}} : memref<f32>) recipe({{.*}}) implicit(true) name("") -> memref<f32>

// -----

// Test scalar in kernels construct - should generate copyin/copyout
func.func @test_scalar_in_kernels() {
  %alloc = memref.alloca() : memref<f64>
  acc.kernels {
    %load = memref.load %alloc[] : memref<f64>
    acc.terminator
  }
  return
}

// CHECK-LABEL: func.func @test_scalar_in_kernels
// CHECK: %[[COPYIN:.*]] = acc.copyin varPtr({{.*}} : memref<f64>) dataClause(acc_copy) implicit(true) name("") -> memref<f64>
// CHECK: acc.copyout accPtr(%[[COPYIN]] : memref<f64>) to varPtr({{.*}} : memref<f64>) dataClause(acc_copy) implicit(true) name("")

// -----

// Test scalar in parallel with default(none) - should NOT generate implicit data
func.func @test_scalar_parallel_defaultnone() {
  %alloc = memref.alloca() : memref<f32>
  acc.parallel {
    %load = memref.load %alloc[] : memref<f32>
    acc.yield
  } defaultAttr(none)
  return
}

// CHECK-LABEL: func.func @test_scalar_parallel_defaultnone
// CHECK-NOT: acc.firstprivate
// CHECK-NOT: acc.copyin

// -----

// Test array in parallel - should generate copyin/copyout
func.func @test_array_in_parallel() {
  %alloc = memref.alloca() : memref<10xf32>
  acc.parallel {
    %c0 = arith.constant 0 : index
    %load = memref.load %alloc[%c0] : memref<10xf32>
    acc.yield
  }
  return
}

// CHECK-LABEL: func.func @test_array_in_parallel
// CHECK: %[[COPYIN:.*]] = acc.copyin varPtr({{.*}} : memref<10xf32>) dataClause(acc_copy) implicit(true) name("") -> memref<10xf32>
// CHECK: acc.copyout accPtr(%[[COPYIN]] : memref<10xf32>) to varPtr({{.*}} : memref<10xf32>) dataClause(acc_copy) implicit(true) name("")

// -----

// Test array in kernels - should generate copyin/copyout
func.func @test_array_in_kernels() {
  %alloc = memref.alloca() : memref<20xi32>
  acc.kernels {
    %c0 = arith.constant 0 : index
    %load = memref.load %alloc[%c0] : memref<20xi32>
    acc.terminator
  }
  return
}

// CHECK-LABEL: func.func @test_array_in_kernels
// CHECK: %[[COPYIN:.*]] = acc.copyin varPtr({{.*}} : memref<20xi32>) dataClause(acc_copy) implicit(true) name("") -> memref<20xi32>
// CHECK: acc.copyout accPtr(%[[COPYIN]] : memref<20xi32>) to varPtr({{.*}} : memref<20xi32>) dataClause(acc_copy) implicit(true) name("")

// -----

// Test array with default(present) - should generate present
func.func @test_array_parallel_defaultpresent() {
  %alloc = memref.alloca() : memref<10xf32>
  acc.parallel {
    %c0 = arith.constant 0 : index
    %load = memref.load %alloc[%c0] : memref<10xf32>
    acc.yield
  } defaultAttr(present)
  return
}

// CHECK-LABEL: func.func @test_array_parallel_defaultpresent
// CHECK: %[[PRESENT:.*]] = acc.present varPtr({{.*}} : memref<10xf32>) implicit(true) name("") -> memref<10xf32> {acc.from_default}
// CHECK: acc.delete accPtr(%[[PRESENT]] : memref<10xf32>) dataClause(acc_present) implicit(true) name("")

// -----

// Test scalar with default(present) - should still generate firstprivate (scalars ignore default(present))
func.func @test_scalar_parallel_defaultpresent() {
  %alloc = memref.alloca() : memref<f32>
  acc.parallel {
    %load = memref.load %alloc[] : memref<f32>
    acc.yield
  } defaultAttr(present)
  return
}

// CHECK-LABEL: func.func @test_scalar_parallel_defaultpresent
// CHECK: acc.firstprivate varPtr({{.*}} : memref<f32>) recipe({{.*}}) implicit(true) name("") -> memref<f32>

// -----

// Test multidimensional array
func.func @test_multidim_array_in_parallel() {
  %alloc = memref.alloca() : memref<8x16xf32>
  acc.parallel {
    %c0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    %load = memref.load %alloc[%c0, %c1] : memref<8x16xf32>
    acc.yield
  }
  return
}

// CHECK-LABEL: func.func @test_multidim_array_in_parallel
// CHECK: %[[COPYIN:.*]] = acc.copyin varPtr({{.*}} : memref<8x16xf32>) dataClause(acc_copy) implicit(true) name("") -> memref<8x16xf32>
// CHECK: acc.copyout accPtr(%[[COPYIN]] : memref<8x16xf32>) to varPtr({{.*}} : memref<8x16xf32>) dataClause(acc_copy) implicit(true) name("")

// -----

// Test dynamic size array
func.func @test_dynamic_array(%size: index) {
  %alloc = memref.alloca(%size) : memref<?xf64>
  acc.parallel {
    %c0 = arith.constant 0 : index
    %load = memref.load %alloc[%c0] : memref<?xf64>
    acc.yield
  }
  return
}

// CHECK-LABEL: func.func @test_dynamic_array
// CHECK: %[[COPYIN:.*]] = acc.copyin varPtr({{.*}} : memref<?xf64>) dataClause(acc_copy) implicit(true) name("") -> memref<?xf64>
// CHECK: acc.copyout accPtr(%[[COPYIN]] : memref<?xf64>) to varPtr({{.*}} : memref<?xf64>) dataClause(acc_copy) implicit(true) name("")

// -----

// Test variable with explicit data clause - implicit should recognize it
func.func @test_with_explicit_copyin() {
  %alloc = memref.alloca() : memref<100xf32>
  %copyin = acc.copyin varPtr(%alloc : memref<100xf32>) name("explicit") -> memref<100xf32>
  acc.parallel dataOperands(%copyin : memref<100xf32>) {
    %c0 = arith.constant 0 : index
    %load = memref.load %alloc[%c0] : memref<100xf32>
    acc.yield
  }
  acc.copyout accPtr(%copyin : memref<100xf32>) to varPtr(%alloc : memref<100xf32>) name("explicit")
  return
}

// CHECK-LABEL: func.func @test_with_explicit_copyin
// CHECK: acc.present varPtr({{.*}} : memref<100xf32>) implicit(true) name("") -> memref<100xf32>

// -----

// Test multiple variables
func.func @test_multiple_variables() {
  %alloc1 = memref.alloca() : memref<f32>
  %alloc2 = memref.alloca() : memref<10xi32>
  acc.parallel {
    %load1 = memref.load %alloc1[] : memref<f32>
    %c0 = arith.constant 0 : index
    %load2 = memref.load %alloc2[%c0] : memref<10xi32>
    acc.yield
  }
  return
}

// CHECK-LABEL: func.func @test_multiple_variables
// CHECK: acc.firstprivate varPtr({{.*}} : memref<f32>) recipe({{.*}}) implicit(true) name("") -> memref<f32>
// CHECK: %[[COPYIN:.*]] = acc.copyin varPtr({{.*}} : memref<10xi32>) dataClause(acc_copy) implicit(true) name("") -> memref<10xi32>
// CHECK: acc.copyout accPtr(%[[COPYIN]] : memref<10xi32>) to varPtr({{.*}} : memref<10xi32>) dataClause(acc_copy) implicit(true) name("")

// -----

// Test memref.view aliasing - view of explicitly copied buffer should generate present
func.func @test_memref_view(%size: index) {
  %c0 = arith.constant 0 : index
  %buffer = memref.alloca(%size) : memref<?xi8>
  %copyin = acc.copyin varPtr(%buffer : memref<?xi8>) name("buffer") -> memref<?xi8>
  %view = memref.view %buffer[%c0][] : memref<?xi8> to memref<8x64xf32>
  acc.kernels dataOperands(%copyin : memref<?xi8>) {
    %c0_0 = arith.constant 0 : index
    %c0_1 = arith.constant 0 : index
    %load = memref.load %view[%c0_0, %c0_1] : memref<8x64xf32>
    acc.terminator
  }
  acc.copyout accPtr(%copyin : memref<?xi8>) to varPtr(%buffer : memref<?xi8>) name("buffer")
  return
}

// CHECK-LABEL: func.func @test_memref_view
// CHECK: acc.present varPtr({{.*}} : memref<8x64xf32>) implicit(true) name("") -> memref<8x64xf32>

// -----

// Test device data (memref with GPU address space) - should generate deviceptr
func.func @test_device_data_in_parallel() {
  %alloc = memref.alloca() : memref<10xf32, #gpu.address_space<global>>
  acc.parallel {
    %c0 = arith.constant 0 : index
    %load = memref.load %alloc[%c0] : memref<10xf32, #gpu.address_space<global>>
    acc.yield
  }
  return
}

// CHECK-LABEL: func.func @test_device_data_in_parallel
// CHECK: acc.deviceptr varPtr({{.*}} : memref<10xf32, #gpu.address_space<global>>) implicit(true) name("") -> memref<10xf32, #gpu.address_space<global>>
// CHECK-NOT: acc.copyin
// CHECK-NOT: acc.copyout

// -----

// Test device global (memref.global with GPU address space) - should generate deviceptr
memref.global @device_global : memref<10xf32, #gpu.address_space<global>>

func.func @test_device_global_in_parallel() {
  %global = memref.get_global @device_global : memref<10xf32, #gpu.address_space<global>>
  acc.parallel {
    %c0 = arith.constant 0 : index
    %load = memref.load %global[%c0] : memref<10xf32, #gpu.address_space<global>>
    acc.yield
  }
  return
}

// CHECK-LABEL: func.func @test_device_global_in_parallel
// CHECK: acc.deviceptr varPtr({{.*}} : memref<10xf32, #gpu.address_space<global>>) implicit(true) name("device_global") -> memref<10xf32, #gpu.address_space<global>>
// CHECK-NOT: acc.copyin
// CHECK-NOT: acc.copyout

// -----

// Test memref.view tagged with acc.declare deviceptr and used directly in region.
func.func @test_declare_deviceptr_arg_in_parallel(%arg0: memref<?xi8>) {
  %c0 = arith.constant 0 : index
  %view = memref.view %arg0[%c0][] {acc.declare = #acc.declare<dataClause = acc_deviceptr>} : memref<?xi8> to memref<10xf32>
  %devptr = acc.deviceptr varPtr(%view : memref<10xf32>) name("arg0") -> memref<10xf32>
  %token = acc.declare_enter dataOperands(%devptr : memref<10xf32>)
  acc.parallel {
    %c0_1 = arith.constant 0 : index
    %load = memref.load %arg0[%c0_1] : memref<?xi8>
    acc.yield
  }
  acc.declare_exit token(%token) dataOperands(%devptr : memref<10xf32>)
  return
}

// CHECK-LABEL: func.func @test_declare_deviceptr_arg_in_parallel
// CHECK: %[[VIEW:.*]] = memref.view %{{.*}}[{{.*}}][] {acc.declare = #acc.declare<dataClause = acc_deviceptr>} : memref<?xi8> to memref<10xf32>
// CHECK: %[[DEVPTR:.*]] = acc.deviceptr varPtr(%[[VIEW]] : memref<10xf32>) name("arg0") -> memref<10xf32>
// CHECK: %[[TOKEN:.*]] = acc.declare_enter dataOperands(%[[DEVPTR]] : memref<10xf32>)
// CHECK: %[[IMPLICIT_DEVPTR:.*]] = acc.deviceptr varPtr(%{{.*}} : memref<?xi8>) implicit(true) name("") -> memref<?xi8>
// CHECK: acc.parallel dataOperands(%[[IMPLICIT_DEVPTR]] : memref<?xi8>) {
// CHECK: memref.load %[[IMPLICIT_DEVPTR]][{{.*}}] : memref<?xi8>
// CHECK: acc.declare_exit token(%[[TOKEN]]) dataOperands(%[[DEVPTR]] : memref<10xf32>)
// CHECK-NOT: acc.copyin
// CHECK-NOT: acc.copyout