// RUN: mlir-opt %s --transform-interpreter -canonicalize --split-input-file --verify-diagnostics | FileCheck %s

func.func @simple_depth_2_unpeeled(%global: memref<?xf32>, %result: memref<?xf32> ) {
  %c0 = arith.constant 0 : index
  %c100 = arith.constant 100 : index
  %c4 = arith.constant 4 : index
  %shared = memref.alloc(%c100) : memref<?xf32, #gpu.address_space<workgroup>>
  %c0f = arith.constant 0.0 : f32
  // Predication is not currently implemented for transfer_read/write, so this is expected to fail.
  // expected-note @below {{couldn't predicate}}
  scf.for %i = %c0 to %c100 step %c4 iter_args(%accum = %c0f) -> f32 {
    %mem = vector.transfer_read %global[%i], %c0f : memref<?xf32>, vector<4xf32>
    vector.transfer_write %mem, %shared[%i] : vector<4xf32>, memref<?xf32, #gpu.address_space<workgroup>>
    %0 = arith.addf %accum, %accum : f32
    scf.yield %0 : f32
  }
  return
}

!t = !transform.any_op

module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(%arg0: !t {transform.readonly}) {
    %loop = transform.structured.match ops{["scf.for"]} in %arg0 : (!t) -> !t
    // expected-error @below {{irreversible pipelining failure}}
    // expected-note @below {{try setting "peel_epilogue"}}
    transform.nvgpu.pipeline_shared_memory_copies failures(propagate) %loop { depth = 2 } : (!t) -> !t
    transform.yield
  }
}

// -----

// Loop pipeliner is tested separately, just verify the overall shape of the IR here.

func.func private @body(index, memref<?xf32, #gpu.address_space<workgroup>>)

// CHECK-LABEL: @simple_depth_2_peeled
// CHECK-SAME: %[[ARG:.+]]: memref
func.func @simple_depth_2_peeled(%global: memref<?xf32>) {
  %c0 = arith.constant 0 : index
  %c100 = arith.constant 100 : index
  %c200 = arith.constant 200 : index
  %c4 = arith.constant 4 : index
  // CHECK: memref.alloc
  %shared = memref.alloc(%c200) : memref<?xf32, #gpu.address_space<workgroup>>
  %c0f = arith.constant 0.0 : f32
  // CHECK: %[[LOADED1:.+]] = vector.transfer_read %[[ARG]]
  // CHECK: %[[LOADED2:.+]] = vector.transfer_read %[[ARG]]
  // CHECK: %[[LOOP:.+]]:2 = scf.for {{.*}} iter_args(%[[IA1:.+]] = %[[LOADED1]], %[[IA2:.+]] = %[[LOADED2]])
  // CHECK:   vector.transfer_write %[[IA1]]
  // CHECK:   func.call @body
  // CHECK:   %[[LOCAL_LOADED:.+]] = vector.transfer_read %[[ARG]]
  // CHECK:   scf.yield %[[IA2]], %[[LOCAL_LOADED]]
  scf.for %i = %c0 to %c100 step %c4 {
    %mem = vector.transfer_read %global[%i], %c0f : memref<?xf32>, vector<4xf32>
    vector.transfer_write %mem, %shared[%i] : vector<4xf32>, memref<?xf32, #gpu.address_space<workgroup>>
    func.call @body(%i, %shared) : (index, memref<?xf32, #gpu.address_space<workgroup>>) -> ()
  }
  // CHECK: vector.transfer_write %[[LOOP]]#0
  // CHECK: call @body
  // CHECK: vector.transfer_write %[[LOOP]]#1
  // CHECK: call @body
  return
}

!t = !transform.any_op

module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(%arg0: !t {transform.readonly}) {
    %loop = transform.structured.match ops{["scf.for"]} in %arg0 : (!t) -> !t
    transform.nvgpu.pipeline_shared_memory_copies failures(propagate) %loop { depth = 2, peel_epilogue } : (!t) -> !t
    transform.yield
  }
}

// -----

// CHECK-LABEL: @async_depth_2_predicated
// CHECK-SAME: %[[GLOBAL:.+]]: memref
func.func @async_depth_2_predicated(%global: memref<?xf32>, %alloc_size: index) {
  %c0 = arith.constant 0 : index
  %c98 = arith.constant 98 : index
  %c100 = arith.constant 100 : index
  // CHECK-DAG: %[[C4:.+]] = arith.constant 4
  // CHECK-DAG:   %[[C90:.+]] = arith.constant 90
  // CHECK-DAG:   %[[C96:.+]] = arith.constant 96
  // CHECK-DAG:   %[[C8:.+]] = arith.constant 8
  // CHECK-DAG:   %[[C2:.+]] = arith.constant 2
  // CHECK-DAG:   %[[C0:.+]] = arith.constant 0
  %c4 = arith.constant 4 : index
  // CHECK: %[[SHARED:.+]] = memref.alloc{{.*}} #gpu.address_space<workgroup>
  %shared = memref.alloc(%alloc_size) : memref<?xf32, #gpu.address_space<workgroup>>
  %c0f = arith.constant 0.0 : f32
  // CHECK: %[[TOKEN0:.+]] = nvgpu.device_async_copy
  // CHECK: %[[TOKEN1:.+]] = nvgpu.device_async_copy
  // CHECK: scf.for %[[I:.+]] = {{.*}} iter_args
  // CHECK-SAME: %[[ITER_ARG0:.+]] = %[[TOKEN0]]
  // CHECK-SAME: %[[ITER_ARG1:.+]] = %[[TOKEN1]]
  scf.for %i = %c0 to %c98 step %c4 {
    // Condition for the predication "select" below.
    // CHECK:   %[[CMP0:.+]] = arith.cmpi slt, %[[I]], %[[C90]]
    // CHECK:   nvgpu.device_async_wait %[[ITER_ARG0]] {numGroups = 1
    // Original "select" with updated induction variable.
    // CHECK:   %[[I_PLUS_8:.+]] = arith.addi %[[I]], %[[C8]]
    // CHECK:   %[[CMP1:.+]] = arith.cmpi slt, %[[I_PLUS_8]], %[[C96]]
    // CHECK:   %[[SELECTED0:.+]] = arith.select %[[CMP1]], %[[C4]], %[[C2]]
    %c96 = arith.constant 96 : index
    %cond = arith.cmpi slt, %i, %c96 : index
    %c2 = arith.constant 2 : index
    %read_size = arith.select %cond, %c4, %c2 : index

    // Updated induction variables (two more) for the device_async_copy below.
    // These are generated repeatedly by the pipeliner.
    // CHECK:   %[[I_PLUS_8_2:.+]] = arith.addi %[[I]], %[[C8]]
    // CHECK:   %[[I_PLUS_8_3:.+]] = arith.addi %[[I]], %[[C8]]

    // The second "select" is generated by predication and selects 0 for
    // the two last iterations.
    // CHECK:   %[[SELECTED1:.+]] = arith.select %[[CMP0]], %[[SELECTED0]], %[[C0]]
    // CHECK:   %[[ASYNC_TOKEN:.+]] = nvgpu.device_async_copy %[[GLOBAL]][%[[I_PLUS_8_3]]], %[[SHARED]][%[[I_PLUS_8_2]]], 4, %[[SELECTED1]]
    %token = nvgpu.device_async_copy %global[%i], %shared[%i], 4, %read_size
      : memref<?xf32> to memref<?xf32, #gpu.address_space<workgroup>>

    nvgpu.device_async_wait %token

    // CHECK: scf.yield %[[ITER_ARG1]], %[[ASYNC_TOKEN]]
  }
  // There is no need to wait for the last copies as it it was fully predicated
  // out and doesn't load the original data.
  // CHECK-NOT: nvgpu.device_async_wait
  return
}


!t = !transform.any_op

module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(%arg0: !t {transform.readonly}) {
    %loop = transform.structured.match ops{["scf.for"]} in %arg0 : (!t) -> !t
    transform.nvgpu.pipeline_shared_memory_copies failures(propagate) %loop { depth = 2 } : (!t) -> !t
    transform.yield
  }
}

// -----

// CHECK-LABEL: @async_depth_2_peeled
func.func @async_depth_2_peeled(%global: memref<?xf32>) {
  %c0 = arith.constant 0 : index
  %c98 = arith.constant 98 : index
  %c100 = arith.constant 100 : index
  %c4 = arith.constant 4 : index
  %shared = memref.alloc(%c100) : memref<?xf32, #gpu.address_space<workgroup>>
  %c0f = arith.constant 0.0 : f32
  // CHECK: nvgpu.device_async_copy
  // CHECK: nvgpu.device_async_copy
  // CHECK: scf.for
  // CHECK:   nvgpu.device_async_wait %{{.*}} {numGroups = 1
  // CHECK:   arith.select
  // CHECK:   nvgpu.device_async_copy
  // CHECK:   scf.yield
  // CHECK: nvgpu.device_async_wait %{{.*}} {numGroups = 1
  // CHECK: nvgpu.device_async_wait %{{.*}} {numGroups = 0
  scf.for %i = %c0 to %c98 step %c4 {
    %c96 = arith.constant 96 : index
    %cond = arith.cmpi slt, %i, %c96 : index
    %c2 = arith.constant 2 : index
    %read_size = arith.select %cond, %c4, %c2 : index
    %token = nvgpu.device_async_copy %global[%i], %shared[%i], 4, %read_size
      : memref<?xf32> to memref<?xf32, #gpu.address_space<workgroup>>
    nvgpu.device_async_wait %token
  }
  return
}


!t = !transform.any_op

module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(%arg0: !t {transform.readonly}) {
    %loop = transform.structured.match ops{["scf.for"]} in %arg0 : (!t) -> !t
    transform.nvgpu.pipeline_shared_memory_copies failures(propagate) %loop { depth = 2, peel_epilogue } : (!t) -> !t
    transform.yield
  }
}