// RUN: mlir-opt %s --bubble-down-memory-space-casts | FileCheck %s

#map = affine_map<(d0, d1)[s0] -> (d1 * s0 + d0)>

// CHECK: #[[$ATTR_0:.+]] = affine_map<(d0, d1)[s0] -> (d1 * s0 + d0)>
// CHECK-LABEL:   func.func @load_store(
// CHECK-SAME:      %[[ARG0:.*]]: memref<?xf32, 1>,
// CHECK-SAME:      %[[ARG1:.*]]: index) {
// CHECK:           %[[VAL_0:.*]] = memref.load %[[ARG0]]{{\[}}%[[ARG1]]] : memref<?xf32, 1>
// CHECK:           memref.store %[[VAL_0]], %[[ARG0]]{{\[}}%[[ARG1]]] : memref<?xf32, 1>
// CHECK:           return
// CHECK:         }
func.func @load_store(%arg0: memref<?xf32, 1>, %arg1: index) {
  %memspacecast = memref.memory_space_cast %arg0 : memref<?xf32, 1> to memref<?xf32>
  %0 = memref.load %memspacecast[%arg1] : memref<?xf32>
  memref.store %0, %memspacecast[%arg1] : memref<?xf32>
  return
}

// CHECK-LABEL:   func.func @load_store_unfoldable(
// CHECK-SAME:      %[[ARG0:.*]]: memref<?xf32, 1>,
// CHECK-SAME:      %[[ARG1:.*]]: index) {
// CHECK:           %[[VAL_0:.*]] = memref.memory_space_cast %[[ARG0]] : memref<?xf32, 1> to memref<?xf32, 2>
// CHECK:           %[[VAL_1:.*]] = memref.load %[[VAL_0]]{{\[}}%[[ARG1]]] : memref<?xf32, 2>
// CHECK:           memref.store %[[VAL_1]], %[[VAL_0]]{{\[}}%[[ARG1]]] : memref<?xf32, 2>
// CHECK:           return
// CHECK:         }
func.func @load_store_unfoldable(%arg0: memref<?xf32, 1>, %arg1: index) {
  %memspacecast = memref.memory_space_cast %arg0 : memref<?xf32, 1> to memref<?xf32, 2>
  %0 = memref.load %memspacecast[%arg1] : memref<?xf32, 2>
  memref.store %0, %memspacecast[%arg1] : memref<?xf32, 2>
  return
}

// CHECK-LABEL:   func.func @cast(
// CHECK-SAME:                    %[[ARG0:.*]]: memref<2xf32, 1>,
// CHECK-SAME:                    %[[ARG1:.*]]: memref<*xf32, 1>) -> (memref<*xf32>, memref<3x2xf32>) {
// CHECK:           %[[VAL_0:.*]] = memref.cast %[[ARG0]] : memref<2xf32, 1> to memref<*xf32, 1>
// CHECK:           %[[VAL_1:.*]] = memref.memory_space_cast %[[VAL_0]] : memref<*xf32, 1> to memref<*xf32>
// CHECK:           %[[VAL_2:.*]] = memref.cast %[[ARG1]] : memref<*xf32, 1> to memref<3x2xf32, 1>
// CHECK:           %[[VAL_3:.*]] = memref.memory_space_cast %[[VAL_2]] : memref<3x2xf32, 1> to memref<3x2xf32>
// CHECK:           return %[[VAL_1]], %[[VAL_3]] : memref<*xf32>, memref<3x2xf32>
// CHECK:         }
func.func @cast(%arg0: memref<2xf32, 1>, %arg1: memref<*xf32, 1>) -> (memref<*xf32>, memref<3x2xf32>) {
  %memspacecast = memref.memory_space_cast %arg0 : memref<2xf32, 1> to memref<2xf32>
  %1 = memref.cast %memspacecast : memref<2xf32> to memref<*xf32>
  %memspacecast_1 = memref.memory_space_cast %arg1 : memref<*xf32, 1> to memref<*xf32>
  %2 = memref.cast %memspacecast_1 : memref<*xf32> to memref<3x2xf32>
  return %1, %2 : memref<*xf32>, memref<3x2xf32>
}

// CHECK-LABEL:   func.func @view(
// CHECK-SAME:                    %[[ARG0:.*]]: memref<?xi8, 1>,
// CHECK-SAME:                    %[[ARG1:.*]]: index, %[[ARG2:.*]]: index) -> memref<?x?xi8> {
// CHECK:           %[[VAL_0:.*]] = arith.constant 100 : index
// CHECK:           %[[VAL_1:.*]] = memref.view %[[ARG0]]{{\[}}%[[ARG1]]]{{\[}}%[[ARG2]], %[[VAL_0]]] : memref<?xi8, 1> to memref<?x?xi8, 1>
// CHECK:           %[[VAL_2:.*]] = memref.memory_space_cast %[[VAL_1]] : memref<?x?xi8, 1> to memref<?x?xi8>
// CHECK:           return %[[VAL_2]] : memref<?x?xi8>
// CHECK:         }
func.func @view(%arg0: memref<?xi8, 1>, %arg1: index, %arg2: index) -> memref<?x?xi8> {
  %memspacecast = memref.memory_space_cast %arg0 : memref<?xi8, 1> to memref<?xi8>
  %c100 = arith.constant 100 : index
  %view = memref.view %memspacecast[%arg1][%arg2, %c100] : memref<?xi8> to memref<?x?xi8>
  return %view : memref<?x?xi8>
}

// CHECK-LABEL:   func.func @subview(
// CHECK-SAME:      %[[ARG0:.*]]: memref<?x?xf32, 1>,
// CHECK-SAME:      %[[ARG1:.*]]: index) -> memref<8x2xf32, strided<[?, 2], offset: ?>> {
// CHECK:           %[[VAL_0:.*]] = memref.subview %[[ARG0]][4, 2] [8, 2] [3, 2] : memref<?x?xf32, 1> to memref<8x2xf32, strided<[?, 2], offset: ?>, 1>
// CHECK:           %[[VAL_1:.*]] = memref.memory_space_cast %[[VAL_0]] : memref<8x2xf32, strided<[?, 2], offset: ?>, 1> to memref<8x2xf32, strided<[?, 2], offset: ?>>
// CHECK:           return %[[VAL_1]] : memref<8x2xf32, strided<[?, 2], offset: ?>>
// CHECK:         }
func.func @subview(%arg0: memref<?x?xf32, 1>, %arg1: index) -> memref<8x2xf32, strided<[?, 2], offset: ?>> {
  %memspacecast = memref.memory_space_cast %arg0 : memref<?x?xf32, 1> to memref<?x?xf32>
  %subview = memref.subview %memspacecast[4, 2] [8, 2] [3, 2] : memref<?x?xf32> to memref<8x2xf32, strided<[?, 2], offset: ?>>
  return %subview : memref<8x2xf32, strided<[?, 2], offset: ?>>
}

// CHECK-LABEL:   func.func @reinterpret_cast(
// CHECK-SAME:      %[[ARG0:.*]]: memref<?xf32, 1>,
// CHECK-SAME:      %[[ARG1:.*]]: index) -> memref<10x?xf32, strided<[?, 1], offset: ?>> {
// CHECK-DAG:       %[[VAL_0:.*]] = arith.constant 10 : index
// CHECK-DAG:       %[[VAL_1:.*]] = arith.constant 0 : index
// CHECK:           %[[VAL_2:.*]] = memref.reinterpret_cast %[[ARG0]] to offset: {{\[}}%[[VAL_1]]], sizes: [10, %[[VAL_0]]], strides: {{\[}}%[[VAL_0]], 1] : memref<?xf32, 1> to memref<10x?xf32, strided<[?, 1], offset: ?>, 1>
// CHECK:           %[[VAL_3:.*]] = memref.memory_space_cast %[[VAL_2]] : memref<10x?xf32, strided<[?, 1], offset: ?>, 1> to memref<10x?xf32, strided<[?, 1], offset: ?>>
// CHECK:           return %[[VAL_3]] : memref<10x?xf32, strided<[?, 1], offset: ?>>
// CHECK:         }
func.func @reinterpret_cast(%arg0: memref<?xf32, 1>, %arg1: index) -> memref<10x?xf32, strided<[?, 1], offset: ?>> {
  %memspacecast = memref.memory_space_cast %arg0 : memref<?xf32, 1> to memref<?xf32>
  %c0 = arith.constant 0 : index
  %c10 = arith.constant 10 : index
  %reinterpret_cast = memref.reinterpret_cast %memspacecast to offset: [%c0], sizes: [10, %c10], strides: [%c10, 1] : memref<?xf32> to memref<10x?xf32, strided<[?, 1], offset: ?>>
  return %reinterpret_cast : memref<10x?xf32, strided<[?, 1], offset: ?>>
}

// CHECK-LABEL:   func.func @reshape(
// CHECK-SAME:      %[[ARG0:.*]]: memref<?x?xf32, 1>,
// CHECK-SAME:      %[[ARG1:.*]]: memref<1xindex>) -> memref<?xf32> {
// CHECK:           %[[VAL_0:.*]] = memref.reshape %[[ARG0]](%[[ARG1]]) : (memref<?x?xf32, 1>, memref<1xindex>) -> memref<?xf32, 1>
// CHECK:           %[[VAL_1:.*]] = memref.memory_space_cast %[[VAL_0]] : memref<?xf32, 1> to memref<?xf32>
// CHECK:           return %[[VAL_1]] : memref<?xf32>
// CHECK:         }
func.func @reshape(%arg0: memref<?x?xf32, 1>, %arg1: memref<1xindex>) -> memref<?xf32> {
  %memspacecast = memref.memory_space_cast %arg0 : memref<?x?xf32, 1> to memref<?x?xf32>
  %reshape = memref.reshape %memspacecast(%arg1) : (memref<?x?xf32>, memref<1xindex>) -> memref<?xf32>
  return %reshape : memref<?xf32>
}

// CHECK-LABEL:   func.func @expand_shape(
// CHECK-SAME:      %[[ARG0:.*]]: memref<12xf32, 1>) -> memref<3x4xf32> {
// CHECK:           %[[VAL_0:.*]] = memref.expand_shape %[[ARG0]] {{\[\[}}0, 1]] output_shape [3, 4] : memref<12xf32, 1> into memref<3x4xf32, 1>
// CHECK:           %[[VAL_1:.*]] = memref.memory_space_cast %[[VAL_0]] : memref<3x4xf32, 1> to memref<3x4xf32>
// CHECK:           return %[[VAL_1]] : memref<3x4xf32>
// CHECK:         }
func.func @expand_shape(%arg0: memref<12xf32, 1>) -> memref<3x4xf32> {
  %memspacecast = memref.memory_space_cast %arg0 : memref<12xf32, 1> to memref<12xf32>
  %expand_shape = memref.expand_shape %memspacecast [[0, 1]] output_shape [3, 4] : memref<12xf32> into memref<3x4xf32>
  return %expand_shape : memref<3x4xf32>
}

// CHECK-LABEL:   func.func @collapse_shape(
// CHECK-SAME:      %[[ARG0:.*]]: memref<3x4xf32, 1>) -> memref<12xf32> {
// CHECK:           %[[VAL_0:.*]] = memref.collapse_shape %[[ARG0]] {{\[\[}}0, 1]] : memref<3x4xf32, 1> into memref<12xf32, 1>
// CHECK:           %[[VAL_1:.*]] = memref.memory_space_cast %[[VAL_0]] : memref<12xf32, 1> to memref<12xf32>
// CHECK:           return %[[VAL_1]] : memref<12xf32>
// CHECK:         }
func.func @collapse_shape(%arg0: memref<3x4xf32, 1>) -> memref<12xf32> {
  %memspacecast = memref.memory_space_cast %arg0 : memref<3x4xf32, 1> to memref<3x4xf32>
  %collapse_shape = memref.collapse_shape %memspacecast [[0, 1]] : memref<3x4xf32> into memref<12xf32>
  return %collapse_shape : memref<12xf32>
}

// CHECK-LABEL:   func.func @transpose(
// CHECK-SAME:      %[[ARG0:.*]]: memref<?x?xf32, 1>) -> memref<?x?xf32, #[[$ATTR_0]]> {
// CHECK:           %[[VAL_0:.*]] = memref.transpose %[[ARG0]] (d0, d1) -> (d1, d0) : memref<?x?xf32, 1> to memref<?x?xf32, #[[$ATTR_0]], 1>
// CHECK:           %[[VAL_1:.*]] = memref.memory_space_cast %[[VAL_0]] : memref<?x?xf32, #[[$ATTR_0]], 1> to memref<?x?xf32, #[[$ATTR_0]]>
// CHECK:           return %[[VAL_1]] : memref<?x?xf32, #[[$ATTR_0]]>
// CHECK:         }
func.func @transpose(%arg0: memref<?x?xf32, 1>) -> memref<?x?xf32, #map> {
  %memspacecast = memref.memory_space_cast %arg0 : memref<?x?xf32, 1> to memref<?x?xf32>
  %transpose = memref.transpose %memspacecast (d0, d1) -> (d1, d0) : memref<?x?xf32> to memref<?x?xf32, #map>
  return %transpose : memref<?x?xf32, #map>
}

// CHECK-LABEL:   func.func @atomic_rmw(
// CHECK-SAME:      %[[ARG0:.*]]: memref<?xf32, 1>,
// CHECK-SAME:      %[[ARG1:.*]]: index,
// CHECK-SAME:      %[[ARG2:.*]]: f32) -> f32 {
// CHECK:           %[[VAL_0:.*]] = memref.atomic_rmw addf %[[ARG2]], %[[ARG0]]{{\[}}%[[ARG1]]] : (f32, memref<?xf32, 1>) -> f32
// CHECK:           return %[[VAL_0]] : f32
// CHECK:         }
func.func @atomic_rmw(%arg0: memref<?xf32, 1>, %arg1: index, %arg2: f32) -> f32 {
  %memspacecast = memref.memory_space_cast %arg0 : memref<?xf32, 1> to memref<?xf32>
  %0 = memref.atomic_rmw addf %arg2, %memspacecast[%arg1] : (f32, memref<?xf32>) -> f32
  return %0 : f32
}

// CHECK-LABEL:   func.func @assume_alignment(
// CHECK-SAME:      %[[ARG0:.*]]: memref<?xf32, 1>) -> memref<?xf32> {
// CHECK:           %[[VAL_0:.*]] = memref.assume_alignment %[[ARG0]], 16 : memref<?xf32, 1>
// CHECK:           %[[VAL_1:.*]] = memref.memory_space_cast %[[VAL_0]] : memref<?xf32, 1> to memref<?xf32>
// CHECK:           return %[[VAL_1]] : memref<?xf32>
// CHECK:         }
func.func @assume_alignment(%arg0: memref<?xf32, 1>) -> memref<?xf32> {
  %memspacecast = memref.memory_space_cast %arg0 : memref<?xf32, 1> to memref<?xf32>
  %1 = memref.assume_alignment %memspacecast, 16 : memref<?xf32>
  return %1 : memref<?xf32>
}

// CHECK-LABEL:   func.func @op_with_cast_sequence(
// CHECK-SAME:      %[[ARG0:.*]]: memref<4x4xf32, 1>,
// CHECK-SAME:      %[[ARG1:.*]]: index,
// CHECK-SAME:      %[[ARG2:.*]]: f32) -> memref<16xf32> {
// CHECK-DAG:       %[[VAL_0:.*]] = arith.constant 4 : index
// CHECK-DAG:       %[[VAL_1:.*]] = arith.constant 0 : index
// CHECK:           %[[VAL_2:.*]] = memref.expand_shape %[[ARG0]] {{\[\[}}0], [1, 2]] output_shape [4, 2, 2] : memref<4x4xf32, 1> into memref<4x2x2xf32, 1>
// CHECK:           %[[VAL_3:.*]] = memref.collapse_shape %[[VAL_2]] {{\[\[}}0, 1, 2]] : memref<4x2x2xf32, 1> into memref<16xf32, 1>
// CHECK:           %[[VAL_4:.*]] = memref.memory_space_cast %[[VAL_3]] : memref<16xf32, 1> to memref<16xf32>
// CHECK:           %[[VAL_5:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_1]]] : memref<16xf32, 1>
// CHECK:           %[[VAL_6:.*]] = arith.addf %[[VAL_5]], %[[ARG2]] : f32
// CHECK:           memref.store %[[VAL_6]], %[[VAL_3]]{{\[}}%[[VAL_1]]] : memref<16xf32, 1>
// CHECK:           %[[VAL_7:.*]] = memref.atomic_rmw addf %[[ARG2]], %[[VAL_3]]{{\[}}%[[VAL_0]]] : (f32, memref<16xf32, 1>) -> f32
// CHECK:           return %[[VAL_4]] : memref<16xf32>
// CHECK:         }
func.func @op_with_cast_sequence(%arg0: memref<4x4xf32, 1>, %arg1: index, %arg2: f32) -> memref<16xf32> {
  %memspacecast = memref.memory_space_cast %arg0 : memref<4x4xf32, 1> to memref<4x4xf32>
  %c0 = arith.constant 0 : index
  %c4 = arith.constant 4 : index
  %expanded = memref.expand_shape %memspacecast [[0], [1, 2]] output_shape [4, 2, 2] : memref<4x4xf32> into memref<4x2x2xf32>
  %collapsed = memref.collapse_shape %expanded [[0, 1, 2]] : memref<4x2x2xf32> into memref<16xf32>
  %loaded = memref.load %collapsed[%c0] : memref<16xf32>
  %added = arith.addf %loaded, %arg2 : f32
  memref.store %added, %collapsed[%c0] : memref<16xf32>
  %atomic_result = memref.atomic_rmw addf %arg2, %collapsed[%c4] : (f32, memref<16xf32>) -> f32
  return %collapsed : memref<16xf32>
}

// CHECK-LABEL:   func.func @transfer_read_write(
// CHECK-SAME:      %[[ARG0:.*]]: memref<?xf32, 1>,
// CHECK-SAME:      %[[ARG1:.*]]: index) {
// CHECK:           %[[VAL_0:.*]] = arith.constant 0.000000e+00 : f32
// CHECK:           %[[VAL_1:.*]] = vector.transfer_read %[[ARG0]]{{\[}}%[[ARG1]]], %[[VAL_0]] : memref<?xf32, 1>, vector<4xf32>
// CHECK:           vector.transfer_write %[[VAL_1]], %[[ARG0]]{{\[}}%[[ARG1]]] : vector<4xf32>, memref<?xf32, 1>
// CHECK:           return
// CHECK:         }
func.func @transfer_read_write(%arg0: memref<?xf32, 1>, %arg1: index) {
  %memspacecast = memref.memory_space_cast %arg0 : memref<?xf32, 1> to memref<?xf32>
  %c0 = arith.constant 0.0 : f32
  %0 = vector.transfer_read %memspacecast[%arg1], %c0 : memref<?xf32>, vector<4xf32>
  vector.transfer_write %0, %memspacecast[%arg1] : vector<4xf32>, memref<?xf32>
  return
}

// NOTE: The operations disappear because they can get folded.
// CHECK-LABEL:   func.func @transfer_read_write_tensor(
// CHECK-SAME:      %[[ARG0:.*]]: tensor<?xf32>,
// CHECK-SAME:      %[[ARG1:.*]]: index) -> tensor<?xf32> {
// CHECK:           return %[[ARG0]] : tensor<?xf32>
// CHECK:         }
func.func @transfer_read_write_tensor(%arg0: tensor<?xf32>, %arg1: index) -> tensor<?xf32> {
  %c0 = arith.constant 0.0 : f32
  %0 = vector.transfer_read %arg0[%arg1], %c0 : tensor<?xf32>, vector<4xf32>
  %1 = vector.transfer_write %0, %arg0[%arg1] : vector<4xf32>, tensor<?xf32>
  return %1 : tensor<?xf32>
}

// CHECK-LABEL:   func.func @vector_load_store(
// CHECK-SAME:      %[[ARG0:.*]]: memref<?xf32, 1>,
// CHECK-SAME:      %[[ARG1:.*]]: index) {
// CHECK:           %[[VAL_0:.*]] = vector.load %[[ARG0]]{{\[}}%[[ARG1]]] : memref<?xf32, 1>, vector<4xf32>
// CHECK:           vector.store %[[VAL_0]], %[[ARG0]]{{\[}}%[[ARG1]]] : memref<?xf32, 1>, vector<4xf32>
// CHECK:           return
// CHECK:         }
func.func @vector_load_store(%arg0: memref<?xf32, 1>, %arg1: index) {
  %memspacecast = memref.memory_space_cast %arg0 : memref<?xf32, 1> to memref<?xf32>
  %0 = vector.load %memspacecast[%arg1] : memref<?xf32>, vector<4xf32>
  vector.store %0, %memspacecast[%arg1] : memref<?xf32>, vector<4xf32>
  return
}

// CHECK-LABEL:   func.func @masked_load_store(
// CHECK-SAME:      %[[ARG0:.*]]: memref<?xf32, 1>,
// CHECK-SAME:      %[[ARG1:.*]]: index) {
// CHECK-DAG:       %[[VAL_0:.*]] = arith.constant dense<0.000000e+00> : vector<4xf32>
// CHECK-DAG:       %[[VAL_1:.*]] = arith.constant dense<[true, true, false, false]> : vector<4xi1>
// CHECK:           %[[VAL_2:.*]] = vector.maskedload %[[ARG0]]{{\[}}%[[ARG1]]], %[[VAL_1]], %[[VAL_0]] : memref<?xf32, 1>, vector<4xi1>, vector<4xf32> into vector<4xf32>
// CHECK:           vector.maskedstore %[[ARG0]]{{\[}}%[[ARG1]]], %[[VAL_1]], %[[VAL_2]] : memref<?xf32, 1>, vector<4xi1>, vector<4xf32>
// CHECK:           return
// CHECK:         }
func.func @masked_load_store(%arg0: memref<?xf32, 1>, %arg1: index) {
  %memspacecast = memref.memory_space_cast %arg0 : memref<?xf32, 1> to memref<?xf32>
  %mask = arith.constant dense<[true, true, false, false]> : vector<4xi1>
  %passthrough = arith.constant dense<0.0> : vector<4xf32>
  %0 = vector.maskedload %memspacecast[%arg1], %mask, %passthrough : memref<?xf32>, vector<4xi1>, vector<4xf32> into vector<4xf32>
  vector.maskedstore %memspacecast[%arg1], %mask, %0 : memref<?xf32>, vector<4xi1>, vector<4xf32>
  return
}

// CHECK-LABEL:   func.func @gather_scatter(
// CHECK-SAME:      %[[ARG0:.*]]: memref<?xf32, 1>,
// CHECK-SAME:      %[[ARG1:.*]]: index) {
// CHECK-DAG:       %[[VAL_0:.*]] = arith.constant dense<0.000000e+00> : vector<4xf32>
// CHECK-DAG:       %[[VAL_1:.*]] = arith.constant dense<true> : vector<4xi1>
// CHECK-DAG:       %[[VAL_2:.*]] = arith.constant dense<[0, 1, 2, 3]> : vector<4xindex>
// CHECK-DAG:       %[[VAL_3:.*]] = arith.constant 0 : index
// CHECK:           %[[VAL_4:.*]] = vector.gather %[[ARG0]]{{\[}}%[[VAL_3]]] {{\[}}%[[VAL_2]]], %[[VAL_1]], %[[VAL_0]] : memref<?xf32, 1>, vector<4xindex>, vector<4xi1>, vector<4xf32> into vector<4xf32>
// CHECK:           vector.scatter %[[ARG0]]{{\[}}%[[VAL_3]]] {{\[}}%[[VAL_2]]], %[[VAL_1]], %[[VAL_4]] : memref<?xf32, 1>, vector<4xindex>, vector<4xi1>, vector<4xf32>
// CHECK:           return
// CHECK:         }
func.func @gather_scatter(%arg0: memref<?xf32, 1>, %arg1: index) {
  %memspacecast = memref.memory_space_cast %arg0 : memref<?xf32, 1> to memref<?xf32>
  %c0 = arith.constant 0 : index
  %indices = arith.constant dense<[0, 1, 2, 3]> : vector<4xindex>
  %mask = arith.constant dense<true> : vector<4xi1>
  %passthrough = arith.constant dense<0.0> : vector<4xf32>
  %0 = vector.gather %memspacecast[%c0] [%indices], %mask, %passthrough : memref<?xf32>, vector<4xindex>, vector<4xi1>, vector<4xf32> into vector<4xf32>
  vector.scatter %memspacecast[%c0] [%indices], %mask, %0 : memref<?xf32>, vector<4xindex>, vector<4xi1>, vector<4xf32>
  return
}

// CHECK-LABEL:   func.func @expandload_compressstore(
// CHECK-SAME:      %[[ARG0:.*]]: memref<?xf32, 1>,
// CHECK-SAME:      %[[ARG1:.*]]: index) {
// CHECK-DAG:       %[[VAL_0:.*]] = arith.constant dense<0.000000e+00> : vector<4xf32>
// CHECK-DAG:       %[[VAL_1:.*]] = arith.constant dense<[true, true, false, false]> : vector<4xi1>
// CHECK:           %[[VAL_2:.*]] = vector.expandload %[[ARG0]]{{\[}}%[[ARG1]]], %[[VAL_1]], %[[VAL_0]] : memref<?xf32, 1>, vector<4xi1>, vector<4xf32> into vector<4xf32>
// CHECK:           vector.compressstore %[[ARG0]]{{\[}}%[[ARG1]]], %[[VAL_1]], %[[VAL_2]] : memref<?xf32, 1>, vector<4xi1>, vector<4xf32>
// CHECK:           return
// CHECK:         }
func.func @expandload_compressstore(%arg0: memref<?xf32, 1>, %arg1: index) {
  %memspacecast = memref.memory_space_cast %arg0 : memref<?xf32, 1> to memref<?xf32>
  %mask = arith.constant dense<[true, true, false, false]> : vector<4xi1>
  %passthrough = arith.constant dense<0.0> : vector<4xf32>
  %0 = vector.expandload %memspacecast[%arg1], %mask, %passthrough : memref<?xf32>, vector<4xi1>, vector<4xf32> into vector<4xf32>
  vector.compressstore %memspacecast[%arg1], %mask, %0 : memref<?xf32>, vector<4xi1>, vector<4xf32>
  return
}