147 lines
5.6 KiB
MLIR
147 lines
5.6 KiB
MLIR
// 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>
|
|
}
|