// RUN: mlir-opt %s -allow-unregistered-dialect -arm-sve-legalize-vector-storage -split-input-file -verify-diagnostics | FileCheck %s

/// This tests the basic functionality of the -arm-sve-legalize-vector-storage pass.

// -----

// CHECK-LABEL: @store_and_reload_sve_predicate_nxv1i1(
// CHECK-SAME:                                         %[[MASK:.*]]: vector<[1]xi1>)
func.func @store_and_reload_sve_predicate_nxv1i1(%mask: vector<[1]xi1>) -> vector<[1]xi1> {
  // CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca() {alignment = 2 : i64} : memref<vector<[16]xi1>>
  %alloca = memref.alloca() : memref<vector<[1]xi1>>
  // CHECK-NEXT: %[[SVBOOL:.*]] = arm_sve.convert_to_svbool %[[MASK]] : vector<[1]xi1>
  // CHECK-NEXT: memref.store %[[SVBOOL]], %[[ALLOCA]][] : memref<vector<[16]xi1>>
  memref.store %mask, %alloca[] : memref<vector<[1]xi1>>
  // CHECK-NEXT: %[[RELOAD:.*]] = memref.load %[[ALLOCA]][] : memref<vector<[16]xi1>>
  // CHECK-NEXT: %[[MASK:.*]] = arm_sve.convert_from_svbool %[[RELOAD]] : vector<[1]xi1>
  %reload = memref.load %alloca[] : memref<vector<[1]xi1>>
  // CHECK-NEXT: return %[[MASK]] : vector<[1]xi1>
  return %reload : vector<[1]xi1>
}

// -----

// CHECK-LABEL: @store_and_reload_sve_predicate_nxv2i1(
// CHECK-SAME:                                         %[[MASK:.*]]: vector<[2]xi1>)
func.func @store_and_reload_sve_predicate_nxv2i1(%mask: vector<[2]xi1>) -> vector<[2]xi1> {
  // CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca() {alignment = 2 : i64} : memref<vector<[16]xi1>>
  %alloca = memref.alloca() : memref<vector<[2]xi1>>
  // CHECK-NEXT: %[[SVBOOL:.*]] = arm_sve.convert_to_svbool %[[MASK]] : vector<[2]xi1>
  // CHECK-NEXT: memref.store %[[SVBOOL]], %[[ALLOCA]][] : memref<vector<[16]xi1>>
  memref.store %mask, %alloca[] : memref<vector<[2]xi1>>
  // CHECK-NEXT: %[[RELOAD:.*]] = memref.load %[[ALLOCA]][] : memref<vector<[16]xi1>>
  // CHECK-NEXT: %[[MASK:.*]] = arm_sve.convert_from_svbool %[[RELOAD]] : vector<[2]xi1>
  %reload = memref.load %alloca[] : memref<vector<[2]xi1>>
  // CHECK-NEXT: return %[[MASK]] : vector<[2]xi1>
  return %reload : vector<[2]xi1>
}

// -----

// CHECK-LABEL: @store_and_reload_sve_predicate_nxv4i1(
// CHECK-SAME:                                         %[[MASK:.*]]: vector<[4]xi1>)
func.func @store_and_reload_sve_predicate_nxv4i1(%mask: vector<[4]xi1>) -> vector<[4]xi1> {
  // CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca() {alignment = 2 : i64} : memref<vector<[16]xi1>>
  %alloca = memref.alloca() : memref<vector<[4]xi1>>
  // CHECK-NEXT: %[[SVBOOL:.*]] = arm_sve.convert_to_svbool %[[MASK]] : vector<[4]xi1>
  // CHECK-NEXT: memref.store %[[SVBOOL]], %[[ALLOCA]][] : memref<vector<[16]xi1>>
  memref.store %mask, %alloca[] : memref<vector<[4]xi1>>
  // CHECK-NEXT: %[[RELOAD:.*]] = memref.load %[[ALLOCA]][] : memref<vector<[16]xi1>>
  // CHECK-NEXT: %[[MASK:.*]] = arm_sve.convert_from_svbool %[[RELOAD]] : vector<[4]xi1>
  %reload = memref.load %alloca[] : memref<vector<[4]xi1>>
  // CHECK-NEXT: return %[[MASK]] : vector<[4]xi1>
  return %reload : vector<[4]xi1>
}

// -----

// CHECK-LABEL: @store_and_reload_sve_predicate_nxv8i1(
// CHECK-SAME:                                         %[[MASK:.*]]: vector<[8]xi1>)
func.func @store_and_reload_sve_predicate_nxv8i1(%mask: vector<[8]xi1>) -> vector<[8]xi1> {
  // CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca() {alignment = 2 : i64} : memref<vector<[16]xi1>>
  %alloca = memref.alloca() : memref<vector<[8]xi1>>
  // CHECK-NEXT: %[[SVBOOL:.*]] = arm_sve.convert_to_svbool %[[MASK]] : vector<[8]xi1>
  // CHECK-NEXT: memref.store %[[SVBOOL]], %[[ALLOCA]][] : memref<vector<[16]xi1>>
  memref.store %mask, %alloca[] : memref<vector<[8]xi1>>
  // CHECK-NEXT: %[[RELOAD:.*]] = memref.load %[[ALLOCA]][] : memref<vector<[16]xi1>>
  // CHECK-NEXT: %[[MASK:.*]] = arm_sve.convert_from_svbool %[[RELOAD]] : vector<[8]xi1>
  %reload = memref.load %alloca[] : memref<vector<[8]xi1>>
  // CHECK-NEXT: return %[[MASK]] : vector<[8]xi1>
  return %reload : vector<[8]xi1>
}

// -----

// CHECK-LABEL: @store_and_reload_sve_predicate_nxv16i1(
// CHECK-SAME:                                         %[[MASK:.*]]: vector<[16]xi1>)
func.func @store_and_reload_sve_predicate_nxv16i1(%mask: vector<[16]xi1>) -> vector<[16]xi1> {
  // CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca() {alignment = 2 : i64} : memref<vector<[16]xi1>>
  %alloca = memref.alloca() : memref<vector<[16]xi1>>
  // CHECK-NEXT: memref.store %[[MASK]], %[[ALLOCA]][] : memref<vector<[16]xi1>>
  memref.store %mask, %alloca[] : memref<vector<[16]xi1>>
  // CHECK-NEXT: %[[RELOAD:.*]] = memref.load %[[ALLOCA]][] : memref<vector<[16]xi1>>
  %reload = memref.load %alloca[] : memref<vector<[16]xi1>>
  // CHECK-NEXT: return %[[RELOAD]] : vector<[16]xi1>
  return %reload : vector<[16]xi1>
}

// -----

/// This is not a valid SVE mask type, so is ignored by the
// `-arm-sve-legalize-vector-storage` pass.

// CHECK-LABEL: @store_and_reload_unsupported_type(
// CHECK-SAME:                                         %[[MASK:.*]]: vector<[7]xi1>)
func.func @store_and_reload_unsupported_type(%mask: vector<[7]xi1>) -> vector<[7]xi1> {
  // CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca() {alignment = 2 : i64} : memref<vector<[7]xi1>>
  %alloca = memref.alloca() : memref<vector<[7]xi1>>
  // CHECK-NEXT: memref.store %[[MASK]], %[[ALLOCA]][] : memref<vector<[7]xi1>>
  memref.store %mask, %alloca[] : memref<vector<[7]xi1>>
  // CHECK-NEXT: %[[RELOAD:.*]] = memref.load %[[ALLOCA]][] : memref<vector<[7]xi1>>
  %reload = memref.load %alloca[] : memref<vector<[7]xi1>>
  // CHECK-NEXT: return %[[RELOAD]] : vector<[7]xi1>
  return %reload : vector<[7]xi1>
}

// -----

// CHECK-LABEL: @store_2d_mask_and_reload_slice(
// CHECK-SAME:                                  %[[MASK:.*]]: vector<3x[8]xi1>)
func.func @store_2d_mask_and_reload_slice(%mask: vector<3x[8]xi1>) -> vector<[8]xi1> {
  // CHECK-NEXT: %[[C0:.*]] = arith.constant 0 : index
  %c0 = arith.constant 0 : index
  // CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca() {alignment = 2 : i64} : memref<vector<3x[16]xi1>>
  %alloca = memref.alloca() : memref<vector<3x[8]xi1>>
  // CHECK-NEXT: %[[SVBOOL:.*]] = arm_sve.convert_to_svbool %[[MASK]] : vector<3x[8]xi1>
  // CHECK-NEXT: memref.store %[[SVBOOL]], %[[ALLOCA]][] : memref<vector<3x[16]xi1>>
  memref.store %mask, %alloca[] : memref<vector<3x[8]xi1>>
  // CHECK-NEXT: %[[UNPACK:.*]] = vector.type_cast %[[ALLOCA]] : memref<vector<3x[16]xi1>> to memref<3xvector<[16]xi1>>
  %unpack = vector.type_cast %alloca : memref<vector<3x[8]xi1>> to memref<3xvector<[8]xi1>>
  // CHECK-NEXT: %[[RELOAD:.*]] = memref.load %[[UNPACK]][%[[C0]]] : memref<3xvector<[16]xi1>>
  // CHECK-NEXT: %[[SLICE:.*]] = arm_sve.convert_from_svbool %[[RELOAD]] : vector<[8]xi1>
  %slice = memref.load %unpack[%c0] : memref<3xvector<[8]xi1>>
  // CHECK-NEXT: return %[[SLICE]] : vector<[8]xi1>
  return %slice : vector<[8]xi1>
}

// -----

// CHECK-LABEL: @set_sve_alloca_alignment
func.func @set_sve_alloca_alignment() {
  /// This checks the alignment of alloca's of scalable vectors will be
  /// something the backend can handle. Currently, the backend sets the
  /// alignment of scalable vectors to their base size (i.e. their size at
  /// vscale = 1). This works for hardware-sized types, which always get a
  /// 16-byte alignment. The problem is larger types e.g. vector<[8]xf32> end up
  /// with alignments larger than 16-bytes (e.g. 32-bytes here), which are
  /// unsupported. The `-arm-sve-legalize-vector-storage` pass avoids this
  /// issue by explicitly setting the alignment to 16-bytes for all scalable
  /// vectors.

  // CHECK-COUNT-6: alignment = 16
  %a1 = memref.alloca() : memref<vector<[32]xi8>>
  %a2 = memref.alloca() : memref<vector<[16]xi8>>
  %a3 = memref.alloca() : memref<vector<[8]xi8>>
  %a4 = memref.alloca() : memref<vector<[4]xi8>>
  %a5 = memref.alloca() : memref<vector<[2]xi8>>
  %a6 = memref.alloca() : memref<vector<[1]xi8>>

  // CHECK-COUNT-6: alignment = 16
  %b1 = memref.alloca() : memref<vector<[32]xi16>>
  %b2 = memref.alloca() : memref<vector<[16]xi16>>
  %b3 = memref.alloca() : memref<vector<[8]xi16>>
  %b4 = memref.alloca() : memref<vector<[4]xi16>>
  %b5 = memref.alloca() : memref<vector<[2]xi16>>
  %b6 = memref.alloca() : memref<vector<[1]xi16>>

  // CHECK-COUNT-6: alignment = 16
  %c1 = memref.alloca() : memref<vector<[32]xi32>>
  %c2 = memref.alloca() : memref<vector<[16]xi32>>
  %c3 = memref.alloca() : memref<vector<[8]xi32>>
  %c4 = memref.alloca() : memref<vector<[4]xi32>>
  %c5 = memref.alloca() : memref<vector<[2]xi32>>
  %c6 = memref.alloca() : memref<vector<[1]xi32>>

  // CHECK-COUNT-6: alignment = 16
  %d1 = memref.alloca() : memref<vector<[32]xi64>>
  %d2 = memref.alloca() : memref<vector<[16]xi64>>
  %d3 = memref.alloca() : memref<vector<[8]xi64>>
  %d4 = memref.alloca() : memref<vector<[4]xi64>>
  %d5 = memref.alloca() : memref<vector<[2]xi64>>
  %d6 = memref.alloca() : memref<vector<[1]xi64>>

  // CHECK-COUNT-6: alignment = 16
  %e1 = memref.alloca() : memref<vector<[32]xf32>>
  %e2 = memref.alloca() : memref<vector<[16]xf32>>
  %e3 = memref.alloca() : memref<vector<[8]xf32>>
  %e4 = memref.alloca() : memref<vector<[4]xf32>>
  %e5 = memref.alloca() : memref<vector<[2]xf32>>
  %e6 = memref.alloca() : memref<vector<[1]xf32>>

  // CHECK-COUNT-6: alignment = 16
  %f1 = memref.alloca() : memref<vector<[32]xf64>>
  %f2 = memref.alloca() : memref<vector<[16]xf64>>
  %f3 = memref.alloca() : memref<vector<[8]xf64>>
  %f4 = memref.alloca() : memref<vector<[4]xf64>>
  %f5 = memref.alloca() : memref<vector<[2]xf64>>
  %f6 = memref.alloca() : memref<vector<[1]xf64>>

  "prevent.dce"(
    %a1, %a2, %a3, %a4, %a5, %a6,
    %b1, %b2, %b3, %b4, %b5, %b6,
    %c1, %c2, %c3, %c4, %c5, %c6,
    %d1, %d2, %d3, %d4, %d5, %d6,
    %e1, %e2, %e3, %e4, %e5, %e6,
    %f1, %f2, %f3, %f4, %f5, %f6)
    : (memref<vector<[32]xi8>>, memref<vector<[16]xi8>>, memref<vector<[8]xi8>>, memref<vector<[4]xi8>>, memref<vector<[2]xi8>>, memref<vector<[1]xi8>>,
       memref<vector<[32]xi16>>, memref<vector<[16]xi16>>, memref<vector<[8]xi16>>, memref<vector<[4]xi16>>, memref<vector<[2]xi16>>, memref<vector<[1]xi16>>,
       memref<vector<[32]xi32>>, memref<vector<[16]xi32>>, memref<vector<[8]xi32>>, memref<vector<[4]xi32>>, memref<vector<[2]xi32>>, memref<vector<[1]xi32>>,
       memref<vector<[32]xi64>>, memref<vector<[16]xi64>>, memref<vector<[8]xi64>>, memref<vector<[4]xi64>>, memref<vector<[2]xi64>>, memref<vector<[1]xi64>>,
       memref<vector<[32]xf32>>, memref<vector<[16]xf32>>, memref<vector<[8]xf32>>, memref<vector<[4]xf32>>, memref<vector<[2]xf32>>, memref<vector<[1]xf32>>,
       memref<vector<[32]xf64>>, memref<vector<[16]xf64>>, memref<vector<[8]xf64>>, memref<vector<[4]xf64>>, memref<vector<[2]xf64>>, memref<vector<[1]xf64>>) -> ()
  return
}