// RUN: mlir-opt --arith-int-narrowing="int-bitwidths-supported=1,8,16,32" \
// RUN:          --verify-diagnostics %s | FileCheck %s

// Check that we can calculate `linalg.index` value bounds and use them to
// optimize index casts.

//===----------------------------------------------------------------------===//
// arith.index_cast
//===----------------------------------------------------------------------===//

// CHECK-LABEL: func @linalg_indexcast_dim_0_i8
// CHECK:         %[[IDX:.+]] = linalg.index 0 : index
// CHECK-NEXT:    %[[INT:.+]] = arith.index_cast %[[IDX]] : index to i8
// CHECK-NEXT:    %[[FP:.+]]  = arith.sitofp %[[INT]] : i8 to f16
// CHECK-NEXT:    linalg.yield %[[FP]] : f16
func.func @linalg_indexcast_dim_0_i8(%arg0: tensor<f16>) -> tensor<128xf16> {
  %init = tensor.empty() : tensor<128xf16>
  %res = linalg.generic {
      indexing_maps = [affine_map<(d0) -> ()>, affine_map<(d0) -> (d0)>],
      iterator_types = ["parallel"]
    }
    ins(%arg0 : tensor<f16>)
    outs(%init : tensor<128xf16>) {
  ^bb0(%in: f16, %out: f16):
    %idx = linalg.index 0 : index
    %int = arith.index_cast %idx : index to i64
    %fp = arith.sitofp %int : i64 to f16
    linalg.yield %fp : f16
  } -> tensor<128xf16>

  return %res : tensor<128xf16>
}

// CHECK-LABEL: func @linalg_indexcast_dim_1_i16
// CHECK:         %[[IDX:.+]] = linalg.index 1 : index
// CHECK-NEXT:    %[[INT:.+]] = arith.index_cast %[[IDX]] : index to i16
// CHECK-NEXT:    %[[FP:.+]]  = arith.sitofp %[[INT]] : i16 to f16
// CHECK-NEXT:    linalg.yield %[[FP]] : f16
func.func @linalg_indexcast_dim_1_i16(%arg0: tensor<f16>, %arg1: tensor<?x129xf16>) -> tensor<?x129xf16> {
  %res = linalg.generic {
      indexing_maps = [affine_map<(d0, d1) -> ()>, affine_map<(d0, d1) -> (d0, d1)>],
      iterator_types = ["parallel", "parallel"]
    }
    ins(%arg0 : tensor<f16>)
    outs(%arg1 : tensor<?x129xf16>) {
  ^bb0(%in: f16, %out: f16):
    %idx = linalg.index 1 : index
    %int = arith.index_cast %idx : index to i64
    %fp = arith.sitofp %int : i64 to f16
    linalg.yield %fp : f16
  } -> tensor<?x129xf16>

  return %res : tensor<?x129xf16>
}

// CHECK-LABEL: func @linalg_indexcast_dynamic_dim_i64
// CHECK:         %[[IDX:.+]] = linalg.index 0 : index
// CHECK-NEXT:    %[[INT:.+]] = arith.index_cast %[[IDX]] : index to i64
// CHECK-NEXT:    %[[FP:.+]]  = arith.sitofp %[[INT]] : i64 to f16
// CHECK-NEXT:    linalg.yield %[[FP]] : f16
func.func @linalg_indexcast_dynamic_dim_i64(%arg0: tensor<f16>, %arg1: tensor<?xf16>) -> tensor<?xf16> {
  %res = linalg.generic {
      indexing_maps = [affine_map<(d0) -> ()>, affine_map<(d0) -> (d0)>],
      iterator_types = ["parallel"]
    }
    ins(%arg0 : tensor<f16>)
    outs(%arg1 : tensor<?xf16>) {
  ^bb0(%in: f16, %out: f16):
    %idx = linalg.index 0 : index
    %int = arith.index_cast %idx : index to i64
    %fp = arith.sitofp %int : i64 to f16
    linalg.yield %fp : f16
  } -> tensor<?xf16>

  return %res : tensor<?xf16>
}

//===----------------------------------------------------------------------===//
// arith.index_castui
//===----------------------------------------------------------------------===//

// CHECK-LABEL: func @linalg_indexcastui_dim_0_i8
// CHECK:         %[[IDX:.+]] = linalg.index 0 : index
// CHECK-NEXT:    %[[INT:.+]] = arith.index_castui %[[IDX]] : index to i8
// CHECK-NEXT:    %[[FP:.+]]  = arith.uitofp %[[INT]] : i8 to f16
// CHECK-NEXT:    linalg.yield %[[FP]] : f16
func.func @linalg_indexcastui_dim_0_i8(%arg0: tensor<f16>) -> tensor<256xf16> {
  %init = tensor.empty() : tensor<256xf16>
  %res = linalg.generic {
      indexing_maps = [affine_map<(d0) -> ()>, affine_map<(d0) -> (d0)>],
      iterator_types = ["parallel"]
    }
    ins(%arg0 : tensor<f16>)
    outs(%init : tensor<256xf16>) {
  ^bb0(%in: f16, %out: f16):
    %idx = linalg.index 0 : index
    %int = arith.index_castui %idx : index to i64
    %fp = arith.uitofp %int : i64 to f16
    linalg.yield %fp : f16
  } -> tensor<256xf16>

  return %res : tensor<256xf16>
}

// CHECK-LABEL: func @linalg_indexcastui_dim_1_i16
// CHECK:         %[[IDX:.+]] = linalg.index 1 : index
// CHECK-NEXT:    %[[INT:.+]] = arith.index_castui %[[IDX]] : index to i16
// CHECK-NEXT:    %[[FP:.+]]  = arith.uitofp %[[INT]] : i16 to f16
// CHECK-NEXT:    linalg.yield %[[FP]] : f16
func.func @linalg_indexcastui_dim_1_i16(%arg0: tensor<f16>, %arg1: tensor<?x257xf16>) -> tensor<?x257xf16> {
  %res = linalg.generic {
      indexing_maps = [affine_map<(d0, d1) -> ()>, affine_map<(d0, d1) -> (d0, d1)>],
      iterator_types = ["parallel", "parallel"]
    }
    ins(%arg0 : tensor<f16>)
    outs(%arg1 : tensor<?x257xf16>) {
  ^bb0(%in: f16, %out: f16):
    %idx = linalg.index 1 : index
    %int = arith.index_castui %idx : index to i64
    %fp = arith.uitofp %int : i64 to f16
    linalg.yield %fp : f16
  } -> tensor<?x257xf16>

  return %res : tensor<?x257xf16>
}

// CHECK-LABEL: func @linalg_indexcastui_dynamic_dim_i64
// CHECK:         %[[IDX:.+]] = linalg.index 0 : index
// CHECK-NEXT:    %[[INT:.+]] = arith.index_castui %[[IDX]] : index to i64
// CHECK-NEXT:    %[[FP:.+]]  = arith.uitofp %[[INT]] : i64 to f16
// CHECK-NEXT:    linalg.yield %[[FP]] : f16
func.func @linalg_indexcastui_dynamic_dim_i64(%arg0: tensor<f16>, %arg1: tensor<?xf16>) -> tensor<?xf16> {
  %res = linalg.generic {
      indexing_maps = [affine_map<(d0) -> ()>, affine_map<(d0) -> (d0)>],
      iterator_types = ["parallel"]
    }
    ins(%arg0 : tensor<f16>)
    outs(%arg1 : tensor<?xf16>) {
  ^bb0(%in: f16, %out: f16):
    %idx = linalg.index 0 : index
    %int = arith.index_castui %idx : index to i64
    %fp = arith.uitofp %int : i64 to f16
    linalg.yield %fp : f16
  } -> tensor<?xf16>

  return %res : tensor<?xf16>
}