// RUN: mlir-opt %s -transform-interpreter | FileCheck %s

  // CHECK: @optimize_shmem([[arg0:%.+]]: memref<{{.*}}>, [[readRow:%.+]]: index, [[readCol:%.+]]: index, [[writeRow:%.+]]: index, [[writeCol:%.+]]: index, [[fragRow:%.+]]: index, [[fragCol:%.+]]: index, [[fragColPerm:%.+]]: index, [[stRow:%.+]]: index, [[stCol:%.+]]: index)
  func.func @optimize_shmem(%arg0: memref<4096x4096xf16>,
                    %readRow: index, %readCol: index,
                    %writeRow: index, %writeCol: index,
                    %fragRow: index, %fragCol: index,
                    %fragColPerm: index,
                    %stRow: index, %stCol: index) {
    %cst = arith.constant 0.000000e+00 : f16

    %shmA = memref.alloc() {alignment = 64 : i64} : memref<128x32xf16, 3>
    %shmB = memref.alloc() {alignment = 64 : i64} : memref<256x32xf16, 3>

    %0 = vector.transfer_read %arg0[%readRow, %readCol], %cst {in_bounds = [true, true]} : memref<4096x4096xf16>, vector<1x8xf16>
    // CHECK: [[c6:%.+]] = arith.constant 6 : index
    // CHECK: [[srcBits:%.+]] = arith.andi [[stRow:%.+]], [[c6]]
    // CHECK: [[c2:%.+]] = arith.constant 2 : index
    // CHECK: [[xorBits:%.+]] = arith.shli [[srcBits]], [[c2]]
    // CHECK: [[stColPerm:%.+]] = arith.xori [[stCol:%.+]], [[xorBits]]
    vector.transfer_write %0, %shmB[%writeRow, %writeCol] {in_bounds = [true, true]} : vector<1x8xf16>, memref<256x32xf16, 3>
    gpu.barrier
    gpu.barrier
    // CHECK: [[c6:%.+]] = arith.constant 6 : index
    // CHECK: [[srcBits:%.+]] = arith.andi [[fragRow]], [[c6]]
    // CHECK: [[c2:%.+]] = arith.constant 2 : index
    // CHECK: [[xorBits:%.+]] = arith.shli [[srcBits]], [[c2]]
    // CHECK: [[fragColPerm:%.+]] = arith.xori [[fragCol:%.+]], [[xorBits]]
    %1 = vector.load %shmB[%fragRow, %fragColPerm] : memref<256x32xf16, 3>, vector<8xf16>
    %2 = vector.transfer_read %arg0[%readRow, %readCol], %cst {in_bounds = [true, true]} : memref<4096x4096xf16>, vector<1x8xf16>
    // CHECK: [[c6:%.+]] = arith.constant 6 : index
    // CHECK: [[srcBits:%.+]] = arith.andi [[stRow:%.+]], [[c6]]
    // CHECK: [[c2:%.+]] = arith.constant 2 : index
    // CHECK: [[xorBits:%.+]] = arith.shli [[srcBits]], [[c2]]
    // CHECK: [[stColPerm:%.+]] = arith.xori [[stCol:%.+]], [[xorBits]]
    vector.transfer_write %2, %shmA[%writeRow, %writeCol] {in_bounds = [true, true]} : vector<1x8xf16>, memref<128x32xf16, 3>
    gpu.barrier
    gpu.barrier
    // CHECK: [[c6:%.+]] = arith.constant 6 : index
    // CHECK: [[srcBits:%.+]] = arith.andi [[fragRow]], [[c6]]
    // CHECK: [[c2:%.+]] = arith.constant 2 : index
    // CHECK: [[xorBits:%.+]] = arith.shli [[srcBits]], [[c2]]
    // CHECK: [[fragColPerm:%.+]] = arith.xori [[fragCol:%.+]], [[xorBits]]
    %3 = vector.load %shmA[%fragRow, %fragColPerm] : memref<128x32xf16, 3>, vector<8xf16>
    return
  }

module attributes { transform.with_named_sequence } {
  transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
    %0 = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.any_op
    transform.amdgpu.optimize_shared_memory_reads_and_writes %0 {sharedMemoryLineSizeBytes = 128, defaultVectorSizeBits = 128}: (!transform.any_op) -> ()
    transform.yield
  } // @__transform_main
} // module