// RUN: mlir-tpu-opt -- %s \
// RUN:   --tpu-infer-memref-layout="hardware-generation=4" \
// RUN:   --tpu-infer-vector-layout \
// RUN:   --tpu-apply-vector-layout="hardware-generation=4" \
// RUN:   --canonicalize \
// RUN:   --lower-to-llo="mock-target=4" \
// RUN:   --cse | FileCheck %s

#map = affine_map<(d0, d1) -> (d0, d1)>

module {
  //
  // Loads a vector from singly tiled 512x512 i32 buffer (at
  // kernel level), adds one to the vector and then outputs
  // to another singly tiled 512x512 i32 buffer (at kernel
  // level). The lowering decomposes the vector operations,
  // within the kernel, along the x-axis and y-axis into
  // fully unrolled code on the basic 8x128 vector operations.
  // This implies that the fully unrolled version yields
  //   (512 / 8) * (512 / 128) = 64 * 4 = 256
  // vector load, add, and store operations.
  //
  // Even though unrolling is good in general compared to
  // a fully rolled loop, since it can hide latency and
  // exploit multiple functional units (e.g. tmu_x for
  // matrix multiplication), it would still be nice to
  // have some level of control on partially unrolling
  // the loop to save some code size.
  //
  // CHECK-LABEL: func.func @vecaddone
  // CHECK-COUNT-256:  llo.vector_load
  // CHECK-COUNT-256:  llo.vadd
  // CHECK-COUNT-256:  llo.vector_store
  // CHECK:       }
  // CHECK:     }
  //
  func.func @vecaddone(%t0: i32,
                       %t1: i32,
                       %INP: memref<512x512xi32>,
                       %OUT: memref<512x512xi32>)
     attributes {
       dimension_semantics =
         [#tpu.dimension_semantics<parallel>,
          #tpu.dimension_semantics<parallel>],
       iteration_bounds = array<i64: 1, 1>,
       window_params = [{transform_indices = #map},
                        {transform_indices = #map}]} {
    %c0 = arith.constant 0 : index
    %one = arith.constant 1 : i32
    %vone = vector.broadcast %one : i32 to vector<512x512xi32>
    %0 = vector.load %INP[%c0, %c0] : memref<512x512xi32>, vector<512x512xi32>
    %1 = arith.addi %0, %vone : vector<512x512xi32>
    vector.store %1, %OUT[%c0, %c0] : memref<512x512xi32>, vector<512x512xi32>
    return
  }
}