From 138ae232ea09997454c98b8a5dc2dfe7fd8211a5 Mon Sep 17 00:00:00 2001 From: Muhammad Asif Manzoor Date: Tue, 5 Nov 2024 01:54:05 +0000 Subject: [PATCH] Stable HLO runtime tests --- .../Silicon/StableHLO/Binary/add_op.mlir | 21 +++++ .../Silicon/StableHLO/Binary/compare_op.mlir | 81 ++++++++++++++++++ .../Silicon/StableHLO/Binary/concat_op.mlir | 85 +++++++++++++++++++ .../Silicon/StableHLO/Binary/divide_op.mlir | 21 +++++ .../Silicon/StableHLO/Binary/logical_op.mlir | 45 ++++++++++ .../Silicon/StableHLO/Binary/maximum_op.mlir | 21 +++++ .../Silicon/StableHLO/Binary/minimum_op.mlir | 21 +++++ .../Silicon/StableHLO/Binary/multiply_op.mlir | 21 +++++ .../StableHLO/Binary/remainder_op.mlir | 21 +++++ .../Silicon/StableHLO/Binary/subtract_op.mlir | 21 +++++ .../StableHLO/Constant/constant_bf16.mlir | 43 ++++++++++ .../StableHLO/Constant/constant_bool.mlir | 43 ++++++++++ .../StableHLO/Constant/constant_f32.mlir | 43 ++++++++++ .../StableHLO/Constant/constant_i16.mlir | 43 ++++++++++ .../StableHLO/Constant/constant_i32.mlir | 43 ++++++++++ .../StableHLO/Constant/constant_i64.mlir | 43 ++++++++++ .../Silicon/StableHLO/Unary/absolute_op.mlir | 20 +++++ .../Silicon/StableHLO/Unary/cbrt_op.mlir | 20 +++++ .../Silicon/StableHLO/Unary/ceil_op.mlir | 20 +++++ .../Silicon/StableHLO/Unary/clamp_op.mlir | 39 +++++++++ .../Silicon/StableHLO/Unary/cosine_op.mlir | 20 +++++ .../Unary/exponential_minus_one_op.mlir | 20 +++++ .../StableHLO/Unary/exponential_op.mlir | 20 +++++ .../Silicon/StableHLO/Unary/floor_op.mlir | 20 +++++ .../Silicon/StableHLO/Unary/isfinite_op.mlir | 20 +++++ .../StableHLO/Unary/log_plus_one_op.mlir | 20 +++++ .../Silicon/StableHLO/Unary/logical_op.mlir | 20 +++++ .../Silicon/StableHLO/Unary/negate_op.mlir | 20 +++++ .../Silicon/StableHLO/Unary/rsqrt_op.mlir | 20 +++++ .../Silicon/StableHLO/Unary/sign_op.mlir | 20 +++++ .../Silicon/StableHLO/Unary/sine_op.mlir | 20 +++++ .../Silicon/StableHLO/Unary/sqrt_op.mlir | 20 +++++ .../Silicon/StableHLO/Unary/tranpose_op.mlir | 19 +++++ .../Silicon/StableHLO/broadcast_op.mlir | 18 ++++ .../Silicon/StableHLO/composite_op.mlir | 25 ++++++ test/ttmlir/Silicon/StableHLO/conv2d_op.mlir | 24 ++++++ test/ttmlir/Silicon/StableHLO/convert_op.mlir | 32 +++++++ .../Silicon/StableHLO/dot_general_op.mlir | 21 +++++ .../StableHLO/get_dimension_size_op.mlir | 18 ++++ .../Silicon/StableHLO/maxpool2d_op.mlir | 19 +++++ .../Silicon/StableHLO/mnist_inference.mlir | 44 ++++++++++ .../Silicon/StableHLO/reduce_add_op.mlir | 22 +++++ .../Silicon/StableHLO/reduce_maximum_op.mlir | 22 +++++ test/ttmlir/Silicon/StableHLO/rehsape_op.mlir | 19 +++++ .../Silicon/StableHLO/scalar_add_op.mlir | 21 +++++ test/ttmlir/Silicon/StableHLO/select_op.mlir | 28 ++++++ test/ttmlir/Silicon/StableHLO/slice_op.mlir | 27 ++++++ 47 files changed, 1314 insertions(+) create mode 100644 test/ttmlir/Silicon/StableHLO/Binary/add_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Binary/compare_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Binary/concat_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Binary/divide_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Binary/logical_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Binary/maximum_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Binary/minimum_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Binary/multiply_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Binary/remainder_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Binary/subtract_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Constant/constant_bf16.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Constant/constant_bool.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Constant/constant_f32.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Constant/constant_i16.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Constant/constant_i32.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Constant/constant_i64.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Unary/absolute_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Unary/cbrt_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Unary/ceil_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Unary/clamp_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Unary/cosine_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Unary/exponential_minus_one_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Unary/exponential_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Unary/floor_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Unary/isfinite_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Unary/log_plus_one_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Unary/logical_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Unary/negate_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Unary/rsqrt_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Unary/sign_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Unary/sine_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Unary/sqrt_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Unary/tranpose_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/broadcast_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/composite_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/conv2d_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/convert_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/dot_general_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/get_dimension_size_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/maxpool2d_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/mnist_inference.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/reduce_add_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/reduce_maximum_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/rehsape_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/scalar_add_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/select_op.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/slice_op.mlir diff --git a/test/ttmlir/Silicon/StableHLO/Binary/add_op.mlir b/test/ttmlir/Silicon/StableHLO/Binary/add_op.mlir new file mode 100644 index 000000000..e6524689c --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Binary/add_op.mlir @@ -0,0 +1,21 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_add attributes {} { + func.func public @test_add(%arg0: tensor<64x128xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xf32> { + // CHECK-LABEL: func.func public @test_add + // CHECK: ttnn.empty + // CHECK: ttnn.add + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: -> tensor<64x128xf32, + %0 = stablehlo.add %arg0, %arg1 : tensor<64x128xf32> + return %0 : tensor<64x128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Binary/compare_op.mlir b/test/ttmlir/Silicon/StableHLO/Binary/compare_op.mlir new file mode 100644 index 000000000..4ba35b125 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Binary/compare_op.mlir @@ -0,0 +1,81 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_compare attributes {} { + func.func public @test_eq(%arg0: tensor<64x128xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xi1> { + // CHECK-LABEL: func.func public @test_eq + // CHECK: ttnn.empty + // CHECK: ttnn.eq + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xbf16, + // CHECK-SAME: -> tensor<64x128xbf16, + %0 = stablehlo.compare EQ, %arg0, %arg1 : (tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xi1> + return %0 : tensor<64x128xi1> + } + + func.func public @test_ne(%arg0: tensor<64x128xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xi1> { + // CHECK-LABEL: func.func public @test_ne + // CHECK: ttnn.empty + // CHECK: ttnn.ne + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xbf16, + // CHECK-SAME: -> tensor<64x128xbf16, + %0 = stablehlo.compare NE, %arg0, %arg1 : (tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xi1> + return %0 : tensor<64x128xi1> + } + + func.func public @test_ge(%arg0: tensor<64x128xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xi1> { + // CHECK-LABEL: func.func public @test_ge + // CHECK: ttnn.empty + // CHECK: ttnn.ge + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xbf16, + // CHECK-SAME: -> tensor<64x128xbf16, + %0 = stablehlo.compare GE, %arg0, %arg1 : (tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xi1> + return %0 : tensor<64x128xi1> + } + + func.func public @test_gt(%arg0: tensor<64x128xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xi1> { + // CHECK-LABEL: func.func public @test_gt + // CHECK: ttnn.empty + // CHECK: ttnn.gt + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xbf16, + // CHECK-SAME: -> tensor<64x128xbf16, + %0 = stablehlo.compare GT, %arg0, %arg1 : (tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xi1> + return %0 : tensor<64x128xi1> + } + + func.func public @test_le(%arg0: tensor<64x128xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xi1> { + // CHECK-LABEL: func.func public @test_le + // CHECK: ttnn.empty + // CHECK: ttnn.le + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xbf16, + // CHECK-SAME: -> tensor<64x128xbf16, + %0 = stablehlo.compare LE, %arg0, %arg1 : (tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xi1> + return %0 : tensor<64x128xi1> + } + + func.func public @test_lt(%arg0: tensor<64x128xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xi1> { + // CHECK-LABEL: func.func public @test_lt + // CHECK: ttnn.empty + // CHECK: ttnn.lt + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xbf16, + // CHECK-SAME: -> tensor<64x128xbf16, + %0 = stablehlo.compare LT, %arg0, %arg1 : (tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xi1> + return %0 : tensor<64x128xi1> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Binary/concat_op.mlir b/test/ttmlir/Silicon/StableHLO/Binary/concat_op.mlir new file mode 100644 index 000000000..484bdfc99 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Binary/concat_op.mlir @@ -0,0 +1,85 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_concat attributes {} { + func.func public @test_concat_0(%arg0: tensor<32x32xf32>, %arg1: tensor<64x32xf32>) -> tensor<96x32xf32> { + // CHECK-LABEL: func.func public @test_concat_0 + // CHECK: ttnn.empty + // CHECK: ttnn.concat + // CHECK-SAME: dim = 0 + // CHECK-SAME: tensor<32x32xf32, + // CHECK-SAME: tensor<64x32xf32, + // CHECK-SAME: tensor<96x32xf32, + // CHECK-SAME: -> tensor<96x32xf32, + %0 = "stablehlo.concatenate"(%arg0, %arg1) { + dimension = 0 : i64 + } : (tensor<32x32xf32>, tensor<64x32xf32>) -> tensor<96x32xf32> + return %0 : tensor<96x32xf32> + } + + func.func public @test_concat_1(%arg0: tensor<32x32xf32>, %arg1: tensor<32x64xf32>) -> tensor<32x96xf32> { + // CHECK-LABEL: func.func public @test_concat_1 + // CHECK: ttnn.empty + // CHECK: ttnn.concat + // CHECK-SAME: dim = 1 + // CHECK-SAME: tensor<32x32xf32, + // CHECK-SAME: tensor<32x64xf32, + // CHECK-SAME: tensor<32x96xf32, + // CHECK-SAME: -> tensor<32x96xf32, + %0 = "stablehlo.concatenate"(%arg0, %arg1) { + dimension = 1 : i64 + } : (tensor<32x32xf32>, tensor<32x64xf32>) -> tensor<32x96xf32> + return %0 : tensor<32x96xf32> + } + + + func.func public @test_concat_2(%arg0: tensor<128x64xf32>, %arg1: tensor<128x96xf32>) -> tensor<128x160xf32> { + // CHECK-LABEL: func.func public @test_concat_2 + // CHECK: ttnn.empty + // CHECK: ttnn.concat + // CHECK-SAME: dim = 1 + // CHECK-SAME: tensor<128x64xf32, + // CHECK-SAME: tensor<128x96xf32, + // CHECK-SAME: tensor<128x160xf32, + // CHECK-SAME: -> tensor<128x160xf32, + %0 = "stablehlo.concatenate"(%arg0, %arg1) { + dimension = 1 : i64 + } : (tensor<128x64xf32>, tensor<128x96xf32>) -> tensor<128x160xf32> + return %0 : tensor<128x160xf32> + } + + func.func public @test_concat_3(%arg0: tensor<64x32xf32>, %arg1: tensor<64x64xf32>) -> tensor<64x96xf32> { + // CHECK-LABEL: func.func public @test_concat_3 + // CHECK: ttnn.empty + // CHECK: ttnn.concat + // CHECK-SAME: dim = 1 + // CHECK-SAME: tensor<64x32xf32, + // CHECK-SAME: tensor<64x64xf32, + // CHECK-SAME: tensor<64x96xf32, + // CHECK-SAME: -> tensor<64x96xf32, + %0 = "stablehlo.concatenate"(%arg0, %arg1) { + dimension = 1 : i64 + } : (tensor<64x32xf32>, tensor<64x64xf32>) -> tensor<64x96xf32> + return %0 : tensor<64x96xf32> + } + + func.func public @test_concat_4(%arg0: tensor<32x32x32x32xf32>, %arg1: tensor<32x32x32x64xf32>) -> tensor<32x32x32x96xf32> { + // CHECK-LABEL: func.func public @test_concat_4 + // CHECK: ttnn.empty + // CHECK: ttnn.concat + // CHECK-SAME: dim = 3 + // CHECK-SAME: tensor<32x32x32x32xf32, + // CHECK-SAME: tensor<32x32x32x64xf32, + // CHECK-SAME: tensor<32x32x32x96xf32, + // CHECK-SAME: -> tensor<32x32x32x96xf32, + %0 = "stablehlo.concatenate"(%arg0, %arg1) { + dimension = 3 : i64 + } : (tensor<32x32x32x32xf32>, tensor<32x32x32x64xf32>) -> tensor<32x32x32x96xf32> + return %0 : tensor<32x32x32x96xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Binary/divide_op.mlir b/test/ttmlir/Silicon/StableHLO/Binary/divide_op.mlir new file mode 100644 index 000000000..6d27885e6 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Binary/divide_op.mlir @@ -0,0 +1,21 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_divice attributes {} { + func.func public @test_divide(%arg0: tensor<64x128xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xf32> { + // CHECK-LABEL: func.func public @test_divide + // CHECK: ttnn.empty + // CHECK: ttnn.div + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: -> tensor<64x128xf32, + %0 = stablehlo.divide %arg0, %arg1 : tensor<64x128xf32> + return %0 : tensor<64x128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Binary/logical_op.mlir b/test/ttmlir/Silicon/StableHLO/Binary/logical_op.mlir new file mode 100644 index 000000000..dcfc073f9 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Binary/logical_op.mlir @@ -0,0 +1,45 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_compare attributes {} { + func.func public @logical_and(%arg0: tensor<64x128xi1>, %arg1: tensor<64x128xi1>) -> tensor<64x128xi1> { + // CHECK-LABEL: func.func public @logical_and + // CHECK: ttnn.empty + // CHECK: ttnn.logical_and + // CHECK-SAME: tensor<64x128xbf16, + // CHECK-SAME: tensor<64x128xbf16, + // CHECK-SAME: tensor<64x128xbf16, + // CHECK-SAME: -> tensor<64x128xbf16, + %0 = stablehlo.and %arg0, %arg1 : tensor<64x128xi1> + return %0 : tensor<64x128xi1> + } + + func.func public @logical_or(%arg0: tensor<64x128xi1>, %arg1: tensor<64x128xi1>) -> tensor<64x128xi1> { + // CHECK-LABEL: func.func public @logical_or + // CHECK: ttnn.empty + // CHECK: ttnn.logical_or + // CHECK-SAME: tensor<64x128xbf16, + // CHECK-SAME: tensor<64x128xbf16, + // CHECK-SAME: tensor<64x128xbf16, + // CHECK-SAME: -> tensor<64x128xbf16, + %0 = stablehlo.or %arg0, %arg1 : tensor<64x128xi1> + return %0 : tensor<64x128xi1> + } + + func.func public @logical_xor(%arg0: tensor<64x128xi1>, %arg1: tensor<64x128xi1>) -> tensor<64x128xi1> { + // CHECK-LABEL: func.func public @logical_xor + // CHECK: ttnn.empty + // CHECK: ttnn.logical_xor + // CHECK-SAME: tensor<64x128xbf16, + // CHECK-SAME: tensor<64x128xbf16, + // CHECK-SAME: tensor<64x128xbf16, + // CHECK-SAME: -> tensor<64x128xbf16, + %0 = stablehlo.xor %arg0, %arg1 : tensor<64x128xi1> + return %0 : tensor<64x128xi1> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Binary/maximum_op.mlir b/test/ttmlir/Silicon/StableHLO/Binary/maximum_op.mlir new file mode 100644 index 000000000..a5c9bf6e6 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Binary/maximum_op.mlir @@ -0,0 +1,21 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_maximum attributes {} { + func.func public @test_maximum(%arg0: tensor<64x128xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xf32> { + // CHECK-LABEL: func.func public @test_maximum + // CHECK: ttnn.empty + // CHECK: ttnn.maximum + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: -> tensor<64x128xf32, + %0 = stablehlo.maximum %arg0, %arg1 : tensor<64x128xf32> + return %0 : tensor<64x128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Binary/minimum_op.mlir b/test/ttmlir/Silicon/StableHLO/Binary/minimum_op.mlir new file mode 100644 index 000000000..b72cc78d5 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Binary/minimum_op.mlir @@ -0,0 +1,21 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_minimum attributes {} { + func.func public @test_minimum(%arg0: tensor<64x128xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xf32> { + // CHECK-LABEL: func.func public @test_minimum + // CHECK: ttnn.empty + // CHECK: ttnn.minimum + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: -> tensor<64x128xf32, + %0 = stablehlo.minimum %arg0, %arg1 : tensor<64x128xf32> + return %0 : tensor<64x128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Binary/multiply_op.mlir b/test/ttmlir/Silicon/StableHLO/Binary/multiply_op.mlir new file mode 100644 index 000000000..b5937abe7 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Binary/multiply_op.mlir @@ -0,0 +1,21 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_multiply attributes {} { + func.func public @test_multiply(%arg0: tensor<64x128xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xf32> { + // CHECK-LABEL: func.func public @test_multiply + // CHECK: ttnn.empty + // CHECK: ttnn.multiply + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: -> tensor<64x128xf32, + %0 = stablehlo.multiply %arg0, %arg1 : tensor<64x128xf32> + return %0 : tensor<64x128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Binary/remainder_op.mlir b/test/ttmlir/Silicon/StableHLO/Binary/remainder_op.mlir new file mode 100644 index 000000000..9444d20f7 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Binary/remainder_op.mlir @@ -0,0 +1,21 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_remainder attributes {} { + func.func public @test_remainder(%arg0: tensor<64x128xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xf32> { + // CHECK-LABEL: func.func public @test_remainder + // CHECK: ttnn.empty + // CHECK: ttnn.remainder + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: -> tensor<64x128xf32, + %0 = stablehlo.remainder %arg0, %arg1 : tensor<64x128xf32> + return %0 : tensor<64x128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Binary/subtract_op.mlir b/test/ttmlir/Silicon/StableHLO/Binary/subtract_op.mlir new file mode 100644 index 000000000..e9169110c --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Binary/subtract_op.mlir @@ -0,0 +1,21 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_subtract attributes {} { + func.func public @test_subtract(%arg0: tensor<64x128xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xf32> { + // CHECK-LABEL: func.func public @test_subtract + // CHECK: ttnn.empty + // CHECK: ttnn.subtract + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: -> tensor<64x128xf32, + %0 = stablehlo.subtract %arg0, %arg1 : tensor<64x128xf32> + return %0 : tensor<64x128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Constant/constant_bf16.mlir b/test/ttmlir/Silicon/StableHLO/Constant/constant_bf16.mlir new file mode 100644 index 000000000..1a24e0759 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Constant/constant_bf16.mlir @@ -0,0 +1,43 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_constant attributes {} { + func.func public @test_bfloat16_scalar() -> tensor { + // CHECK-LABEL: func.func public @test_bfloat16_scalar + // CHECK: ttnn.full + // CHECK-SAME: fillValue = 3.000000e+00 : f32 + // CHECK-SAME: -> tensor<1xbf16 + %0 = stablehlo.constant dense<3.0> : tensor + return %0 : tensor + } + + func.func public @test_bfloat16_scalar_empty() -> tensor { + // CHECK-LABEL: func.func public @test_bfloat16_scalar_empty + // CHECK: ttnn.empty + // CHECK-SAME: -> tensor<1xbf16 + %0 = stablehlo.constant dense<0.0> : tensor + return %0 : tensor + } + + func.func public @test_bfloat16_empty() -> tensor<64x128xbf16> { + // CHECK-LABEL: func.func public @test_bfloat16_empty + // CHECK: ttnn.empty + // CHECK-SAME: -> tensor<64x128xbf16 + %0 = stablehlo.constant dense<0.0> : tensor<64x128xbf16> + return %0 : tensor<64x128xbf16> + } + + func.func public @test_bfloat16_splat() -> tensor<64x128xbf16> { + // CHECK-LABEL: func.func public @test_bfloat16_splat + // CHECK: ttnn.full + // CHECK-SAME: fillValue = 3.000000e+00 : f32 + // CHECK-SAME: -> tensor<64x128xbf16 + %0 = stablehlo.constant dense<3.0> : tensor<64x128xbf16> + return %0 : tensor<64x128xbf16> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Constant/constant_bool.mlir b/test/ttmlir/Silicon/StableHLO/Constant/constant_bool.mlir new file mode 100644 index 000000000..0c51294e3 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Constant/constant_bool.mlir @@ -0,0 +1,43 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_constant attributes {} { + func.func public @test_boolean_scalar() -> tensor { + // CHECK-LABEL: func.func public @test_boolean_scalar + // CHECK: ttnn.full + // CHECK-SAME: fillValue = 1.000000e+00 : f32 + // CHECK-SAME: -> tensor<1xbf16 + %0 = stablehlo.constant dense : tensor + return %0 : tensor + } + + func.func public @test_boolean_scalar_empty() -> tensor { + // CHECK-LABEL: func.func public @test_boolean_scalar_empty + // CHECK: ttnn.empty + // CHECK-SAME: -> tensor<1xbf16 + %0 = stablehlo.constant dense : tensor + return %0 : tensor + } + + func.func public @test_boolean_empty() -> tensor<64x128xi1> { + // CHECK-LABEL: func.func public @test_boolean_empty + // CHECK: ttnn.empty + // CHECK-SAME: -> tensor<64x128xbf16 + %0 = stablehlo.constant dense : tensor<64x128xi1> + return %0 : tensor<64x128xi1> + } + + func.func public @test_boolean_splat() -> tensor<64x128xi1> { + // CHECK-LABEL: func.func public @test_boolean_splat + // CHECK: ttnn.full + // CHECK-SAME: fillValue = 1.000000e+00 : f32 + // CHECK-SAME: -> tensor<64x128xbf16 + %0 = stablehlo.constant dense : tensor<64x128xi1> + return %0 : tensor<64x128xi1> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Constant/constant_f32.mlir b/test/ttmlir/Silicon/StableHLO/Constant/constant_f32.mlir new file mode 100644 index 000000000..5a29facc7 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Constant/constant_f32.mlir @@ -0,0 +1,43 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_constant attributes {} { + func.func public @test_float_scalar() -> tensor { + // CHECK-LABEL: func.func public @test_float_scalar + // CHECK: ttnn.full + // CHECK-SAME: fillValue = 3.000000e+00 : f32 + // CHECK-SAME: -> tensor<1xf32 + %0 = stablehlo.constant dense<3.0> : tensor + return %0 : tensor + } + + func.func public @test_float_scalar_empty() -> tensor { + // CHECK-LABEL: func.func public @test_float_scalar_empty + // CHECK: ttnn.empty + // CHECK-SAME: -> tensor<1xf32 + %0 = stablehlo.constant dense<0.0> : tensor + return %0 : tensor + } + + func.func public @test_float_empty() -> tensor<64x128xf32> { + // CHECK-LABEL: func.func public @test_float_empty + // CHECK: ttnn.empty + // CHECK-SAME: -> tensor<64x128xf32 + %0 = stablehlo.constant dense<0.0> : tensor<64x128xf32> + return %0 : tensor<64x128xf32> + } + + func.func public @test_float_splat() -> tensor<64x128xf32> { + // CHECK-LABEL: func.func public @test_float_splat + // CHECK: ttnn.full + // CHECK-SAME: fillValue = 3.000000e+00 : f32 + // CHECK-SAME: -> tensor<64x128xf32 + %0 = stablehlo.constant dense<3.0> : tensor<64x128xf32> + return %0 : tensor<64x128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Constant/constant_i16.mlir b/test/ttmlir/Silicon/StableHLO/Constant/constant_i16.mlir new file mode 100644 index 000000000..8f4dc247f --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Constant/constant_i16.mlir @@ -0,0 +1,43 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_constant attributes {} { + func.func public @test_int16_scalar() -> tensor { + // CHECK-LABEL: func.func public @test_int16_scalar + // CHECK: ttnn.full + // CHECK-SAME: fillValue = 3.000000e+00 : f32 + // CHECK-SAME: -> tensor<1xi16 + %0 = stablehlo.constant dense<3> : tensor + return %0 : tensor + } + + func.func public @test_int16_scalar_empty() -> tensor { + // CHECK-LABEL: func.func public @test_int16_scalar_empty + // CHECK: ttnn.empty + // CHECK-SAME: -> tensor<1xi16 + %0 = stablehlo.constant dense<0> : tensor + return %0 : tensor + } + + func.func public @test_int16_empty() -> tensor<64x128xi16> { + // CHECK-LABEL: func.func public @test_int16_empty + // CHECK: ttnn.empty + // CHECK-SAME: -> tensor<64x128xi16 + %0 = stablehlo.constant dense<0> : tensor<64x128xi16> + return %0 : tensor<64x128xi16> + } + + func.func public @test_int16_splat() -> tensor<64x128xi16> { + // CHECK-LABEL: func.func public @test_int16_splat + // CHECK: ttnn.full + // CHECK-SAME: fillValue = 3.000000e+00 : f32 + // CHECK-SAME: -> tensor<64x128xi16 + %0 = stablehlo.constant dense<3> : tensor<64x128xi16> + return %0 : tensor<64x128xi16> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Constant/constant_i32.mlir b/test/ttmlir/Silicon/StableHLO/Constant/constant_i32.mlir new file mode 100644 index 000000000..b5c73da0b --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Constant/constant_i32.mlir @@ -0,0 +1,43 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_constant attributes {} { + func.func public @test_int32_scalar() -> tensor { + // CHECK-LABEL: func.func public @test_int32_scalar + // CHECK: ttnn.full + // CHECK-SAME: fillValue = 3.000000e+00 : f32 + // CHECK-SAME: -> tensor<1xi32 + %0 = stablehlo.constant dense<3> : tensor + return %0 : tensor + } + + func.func public @test_int32_scalar_empty() -> tensor { + // CHECK-LABEL: func.func public @test_int32_scalar_empty + // CHECK: ttnn.empty + // CHECK-SAME: -> tensor<1xi32 + %0 = stablehlo.constant dense<0> : tensor + return %0 : tensor + } + + func.func public @test_int32_empty() -> tensor<64x128xi32> { + // CHECK-LABEL: func.func public @test_int32_empty + // CHECK: ttnn.empty + // CHECK-SAME: -> tensor<64x128xi32 + %0 = stablehlo.constant dense<0> : tensor<64x128xi32> + return %0 : tensor<64x128xi32> + } + + func.func public @test_int32_splat() -> tensor<64x128xi32> { + // CHECK-LABEL: func.func public @test_int32_splat + // CHECK: ttnn.full + // CHECK-SAME: fillValue = 3.000000e+00 : f32 + // CHECK-SAME: -> tensor<64x128xi32 + %0 = stablehlo.constant dense<3> : tensor<64x128xi32> + return %0 : tensor<64x128xi32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Constant/constant_i64.mlir b/test/ttmlir/Silicon/StableHLO/Constant/constant_i64.mlir new file mode 100644 index 000000000..bf4a3e8cb --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Constant/constant_i64.mlir @@ -0,0 +1,43 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_constant attributes {} { + func.func public @test_int64_scalar() -> tensor { + // CHECK-LABEL: func.func public @test_int64_scalar + // CHECK: ttnn.full + // CHECK-SAME: fillValue = 3.000000e+00 : f32 + // CHECK-SAME: -> tensor<1xi32 + %0 = stablehlo.constant dense<3> : tensor + return %0 : tensor + } + + func.func public @test_int64_scalar_empty() -> tensor { + // CHECK-LABEL: func.func public @test_int64_scalar_empty + // CHECK: ttnn.empty + // CHECK-SAME: -> tensor<1xi32 + %0 = stablehlo.constant dense<0> : tensor + return %0 : tensor + } + + func.func public @test_int64_empty() -> tensor<64x128xi64> { + // CHECK-LABEL: func.func public @test_int64_empty + // CHECK: ttnn.empty + // CHECK-SAME: -> tensor<64x128xi32 + %0 = stablehlo.constant dense<0> : tensor<64x128xi64> + return %0 : tensor<64x128xi64> + } + + func.func public @test_int64_splat() -> tensor<64x128xi64> { + // CHECK-LABEL: func.func public @test_int64_splat + // CHECK: ttnn.full + // CHECK-SAME: fillValue = 3.000000e+00 : f32 + // CHECK-SAME: -> tensor<64x128xi32 + %0 = stablehlo.constant dense<3> : tensor<64x128xi64> + return %0 : tensor<64x128xi64> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Unary/absolute_op.mlir b/test/ttmlir/Silicon/StableHLO/Unary/absolute_op.mlir new file mode 100644 index 000000000..915238179 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Unary/absolute_op.mlir @@ -0,0 +1,20 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_abs attributes {} { + func.func public @test_abs(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> { + // CHECK-LABEL: func.func public @test_abs + // CHECK: ttnn.empty + // CHECK: ttnn.abs + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: -> tensor<64x128xf32, + %0 = stablehlo.abs %arg0 : tensor<64x128xf32> + return %0 : tensor<64x128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Unary/cbrt_op.mlir b/test/ttmlir/Silicon/StableHLO/Unary/cbrt_op.mlir new file mode 100644 index 000000000..5bfbaf9f6 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Unary/cbrt_op.mlir @@ -0,0 +1,20 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_rsqrt attributes {} { + func.func public @test_cbrt(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> { + // CHECK-LABEL: func.func public @test_cbrt + // CHECK: ttnn.empty + // CHECK: ttnn.cbrt + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: -> tensor<64x128xf32, + %0 = stablehlo.cbrt %arg0 : tensor<64x128xf32> + return %0 : tensor<64x128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Unary/ceil_op.mlir b/test/ttmlir/Silicon/StableHLO/Unary/ceil_op.mlir new file mode 100644 index 000000000..911b775da --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Unary/ceil_op.mlir @@ -0,0 +1,20 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_ceil attributes {} { + func.func public @test_ceil(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> { + // CHECK-LABEL: func.func public @test_ceil + // CHECK: ttnn.empty + // CHECK: ttnn.ceil + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: -> tensor<64x128xf32, + %0 = stablehlo.ceil %arg0 : tensor<64x128xf32> + return %0 : tensor<64x128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Unary/clamp_op.mlir b/test/ttmlir/Silicon/StableHLO/Unary/clamp_op.mlir new file mode 100644 index 000000000..9e625b549 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Unary/clamp_op.mlir @@ -0,0 +1,39 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_transpose attributes {} { + func.func public @test_clamp_constant(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> { + // CHECK-LABEL: func.func public @test_clamp_constant + // CHECK: ttnn.clamp + // CHECK-SAME: {max = 3.000000e+00 : f32, min = 2.000000e+00 : f32} + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: -> tensor<64x128xf32, + %cst = stablehlo.constant dense<2.000000e+00> : tensor<64x128xf32> + %cst_0 = stablehlo.constant dense<3.000000e+00> : tensor<64x128xf32> + %0 = stablehlo.clamp %cst, %arg0, %cst_0 : tensor<64x128xf32> + return %0 : tensor<64x128xf32> + } + + func.func public @test_clamp_tensor(%arg0: tensor<64x128xf32>, %arg1: tensor<64x128xf32>, %arg2: tensor<64x128xf32>) -> tensor<64x128xf32> { + // CHECK-LABEL: func.func public @test_clamp_tensor + // CHECK: ttnn.empty + // CHECK: %[[MAX:.*]] = "ttnn.maximum" + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: -> tensor<64x128xf32, + // CHECK: ttnn.empty + // CHECK: "ttnn.minimum"(%[[MAX]] + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: -> tensor<64x128xf32, + %0 = stablehlo.clamp %arg1, %arg0, %arg2 : tensor<64x128xf32> + return %0 : tensor<64x128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Unary/cosine_op.mlir b/test/ttmlir/Silicon/StableHLO/Unary/cosine_op.mlir new file mode 100644 index 000000000..885a293e7 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Unary/cosine_op.mlir @@ -0,0 +1,20 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_cosine attributes {} { + func.func public @test_cosine(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> { + // CHECK-LABEL: func.func public @test_cosine + // CHECK: ttnn.empty + // CHECK: ttnn.cos + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: -> tensor<64x128xf32, + %0 = stablehlo.cosine %arg0 : tensor<64x128xf32> + return %0 : tensor<64x128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Unary/exponential_minus_one_op.mlir b/test/ttmlir/Silicon/StableHLO/Unary/exponential_minus_one_op.mlir new file mode 100644 index 000000000..30cd1e8c7 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Unary/exponential_minus_one_op.mlir @@ -0,0 +1,20 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_expm1 attributes {} { + func.func public @test_expm1(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> { + // CHECK-LABEL: func.func public @test_expm1 + // CHECK: ttnn.empty + // CHECK: ttnn.expm1 + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: -> tensor<64x128xf32, + %0 = stablehlo.exponential_minus_one %arg0 : tensor<64x128xf32> + return %0 : tensor<64x128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Unary/exponential_op.mlir b/test/ttmlir/Silicon/StableHLO/Unary/exponential_op.mlir new file mode 100644 index 000000000..47a2d6443 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Unary/exponential_op.mlir @@ -0,0 +1,20 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_exp attributes {} { + func.func public @test_exp(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> { + // CHECK-LABEL: func.func public @test_exp + // CHECK: ttnn.empty + // CHECK: ttnn.exp + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: -> tensor<64x128xf32, + %0 = stablehlo.exponential %arg0 : tensor<64x128xf32> + return %0 : tensor<64x128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Unary/floor_op.mlir b/test/ttmlir/Silicon/StableHLO/Unary/floor_op.mlir new file mode 100644 index 000000000..9dbde4c26 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Unary/floor_op.mlir @@ -0,0 +1,20 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_floor attributes {} { + func.func public @test_floor(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> { + // CHECK-LABEL: func.func public @test_floor + // CHECK: ttnn.empty + // CHECK: ttnn.floor + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: -> tensor<64x128xf32, + %0 = stablehlo.floor %arg0 : tensor<64x128xf32> + return %0 : tensor<64x128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Unary/isfinite_op.mlir b/test/ttmlir/Silicon/StableHLO/Unary/isfinite_op.mlir new file mode 100644 index 000000000..04b9f1fef --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Unary/isfinite_op.mlir @@ -0,0 +1,20 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_isfinite attributes {} { + func.func public @test_isfinite(%arg0: tensor<64x128xf32>) -> tensor<64x128xi1> { + // CHECK-LABEL: func.func public @test_isfinite + // CHECK: ttnn.empty + // CHECK: ttnn.isfinite + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xbf16, + // CHECK-SAME: -> tensor<64x128xbf16, + %0 = stablehlo.is_finite %arg0 : (tensor<64x128xf32>) -> tensor<64x128xi1> + return %0 : tensor<64x128xi1> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Unary/log_plus_one_op.mlir b/test/ttmlir/Silicon/StableHLO/Unary/log_plus_one_op.mlir new file mode 100644 index 000000000..cd1d68344 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Unary/log_plus_one_op.mlir @@ -0,0 +1,20 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_log_plus_one attributes {} { + func.func public @test_log_plus_one(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> { + // CHECK-LABEL: func.func public @test_log_plus_one + // CHECK: ttnn.empty + // CHECK: ttnn.log1p + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: -> tensor<64x128xf32, + %0 = stablehlo.log_plus_one %arg0 : tensor<64x128xf32> + return %0 : tensor<64x128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Unary/logical_op.mlir b/test/ttmlir/Silicon/StableHLO/Unary/logical_op.mlir new file mode 100644 index 000000000..489ad2e58 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Unary/logical_op.mlir @@ -0,0 +1,20 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_compare attributes {} { + func.func public @logical_not(%arg0: tensor<64x128xi1>) -> tensor<64x128xi1> { + // CHECK-LABEL: func.func public @logical_not + // CHECK: ttnn.empty + // CHECK: ttnn.logical_not + // CHECK-SAME: tensor<64x128xbf16, + // CHECK-SAME: tensor<64x128xbf16, + // CHECK-SAME: -> tensor<64x128xbf16, + %0 = stablehlo.not %arg0 : tensor<64x128xi1> + return %0 : tensor<64x128xi1> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Unary/negate_op.mlir b/test/ttmlir/Silicon/StableHLO/Unary/negate_op.mlir new file mode 100644 index 000000000..3d970b28e --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Unary/negate_op.mlir @@ -0,0 +1,20 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_neg attributes {} { + func.func public @test_neg(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> { + // CHECK-LABEL: func.func public @test_neg + // CHECK: ttnn.empty + // CHECK: ttnn.neg + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: -> tensor<64x128xf32, + %0 = stablehlo.negate %arg0 : tensor<64x128xf32> + return %0 : tensor<64x128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Unary/rsqrt_op.mlir b/test/ttmlir/Silicon/StableHLO/Unary/rsqrt_op.mlir new file mode 100644 index 000000000..6d0d2efc9 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Unary/rsqrt_op.mlir @@ -0,0 +1,20 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_rsqrt attributes {} { + func.func public @test_rsqrt(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> { + // CHECK-LABEL: func.func public @test_rsqrt + // CHECK: ttnn.empty + // CHECK: ttnn.rsqrt + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: -> tensor<64x128xf32, + %0 = stablehlo.rsqrt %arg0 : tensor<64x128xf32> + return %0 : tensor<64x128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Unary/sign_op.mlir b/test/ttmlir/Silicon/StableHLO/Unary/sign_op.mlir new file mode 100644 index 000000000..6c31b87cc --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Unary/sign_op.mlir @@ -0,0 +1,20 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_sign attributes {} { + func.func public @test_sign(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> { + // CHECK-LABEL: func.func public @test_sign + // CHECK: ttnn.empty + // CHECK: ttnn.sign + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: -> tensor<64x128xf32, + %0 = stablehlo.sign %arg0 : tensor<64x128xf32> + return %0 : tensor<64x128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Unary/sine_op.mlir b/test/ttmlir/Silicon/StableHLO/Unary/sine_op.mlir new file mode 100644 index 000000000..36531bb0c --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Unary/sine_op.mlir @@ -0,0 +1,20 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_sine attributes {} { + func.func public @test_sine(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> { + // CHECK-LABEL: func.func public @test_sine + // CHECK: ttnn.empty + // CHECK: ttnn.sin + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: -> tensor<64x128xf32, + %0 = stablehlo.sine %arg0 : tensor<64x128xf32> + return %0 : tensor<64x128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Unary/sqrt_op.mlir b/test/ttmlir/Silicon/StableHLO/Unary/sqrt_op.mlir new file mode 100644 index 000000000..feb6bd8ac --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Unary/sqrt_op.mlir @@ -0,0 +1,20 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_sqrt attributes {} { + func.func public @test_sqrt(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> { + // CHECK-LABEL: func.func public @test_sqrt + // CHECK: ttnn.empty + // CHECK: ttnn.sqrt + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: -> tensor<64x128xf32, + %0 = stablehlo.sqrt %arg0 : tensor<64x128xf32> + return %0 : tensor<64x128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Unary/tranpose_op.mlir b/test/ttmlir/Silicon/StableHLO/Unary/tranpose_op.mlir new file mode 100644 index 000000000..54795a063 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Unary/tranpose_op.mlir @@ -0,0 +1,19 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_transpose attributes {} { + func.func public @test_transpose(%arg0: tensor<64x128xf32>) -> tensor<128x64xf32> { + // CHECK-LABEL: func.func public @test_transpose + // CHECK: ttnn.transpose + // CHECK-SAME: {dim0 = 1 : si32, dim1 = 0 : si32} + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: -> tensor<128x64xf32, + %0 = stablehlo.transpose %arg0, dims = [1,0] : (tensor<64x128xf32>) -> tensor<128x64xf32> + return %0 : tensor<128x64xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/broadcast_op.mlir b/test/ttmlir/Silicon/StableHLO/broadcast_op.mlir new file mode 100644 index 000000000..c6f7f2e97 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/broadcast_op.mlir @@ -0,0 +1,18 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_broadcast attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} { + func.func public @main(%arg0: tensor<1xf32> {mhlo.layout_mode = "default"}, %arg1: tensor<64x128xf32> {mhlo.layout_mode = "default"}) -> (tensor<64x128xf32> {jax.result_info = "", mhlo.layout_mode = "default"}) { + // CHECK-LABEL: module @jit_broadcast + // CHECK-NOT: broadcast + // CHECK: ttnn.add + %0 = stablehlo.broadcast_in_dim %arg0, dims = [1] : (tensor<1xf32>) -> tensor<64x128xf32> + %1 = stablehlo.add %0, %arg1 : tensor<64x128xf32> + return %1 : tensor<64x128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/composite_op.mlir b/test/ttmlir/Silicon/StableHLO/composite_op.mlir new file mode 100644 index 000000000..9bb62c9e8 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/composite_op.mlir @@ -0,0 +1,25 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_add attributes {} { + // CHECK-NOT: func.func.private @add_impl + func.func private @add_impl(%arg0: tensor<64x128xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xf32> { + %0 = stablehlo.add %arg0, %arg1 : tensor<64x128xf32> + return %0 : tensor<64x128xf32> + } + + // CHECK-LABEL: func.func public @main + func.func public @main(%arg0: tensor<64x128xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xf32> { + // CHECK: ttnn.empty + // CEHCK: ttnn.add + %results = stablehlo.composite "jit_eltwise_add.my_add" %arg0, %arg1 { + decomposition = @add_impl + } : (tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xf32> + return %results : tensor<64x128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/conv2d_op.mlir b/test/ttmlir/Silicon/StableHLO/conv2d_op.mlir new file mode 100644 index 000000000..08b8a0799 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/conv2d_op.mlir @@ -0,0 +1,24 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s +// UNSUPPORTED: true + +module @jit_convolution attributes {} { + func.func public @test_convolution(%arg0: tensor<1x128x128x32xf32>, %arg1: tensor<64x32x3x3xf32>) -> tensor<1x128x128x64xf32> { + %0 = stablehlo.convolution(%arg0, %arg1) + dim_numbers = [b, 0, 1, f]x[o, i, 0, 1]->[b, 0, 1, f], + window = { + stride = [1, 1], + pad = [[1, 1], [1, 1]], + } { + feature_group_count = 1 : i64, + batch_group_count = 1 : i64, + precision_config = [#stablehlo, #stablehlo] + } : (tensor<1x128x128x32xf32>, tensor<64x32x3x3xf32>) -> tensor<1x128x128x64xf32> + return %0 : tensor<1x128x128x64xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/convert_op.mlir b/test/ttmlir/Silicon/StableHLO/convert_op.mlir new file mode 100644 index 000000000..2b87911cf --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/convert_op.mlir @@ -0,0 +1,32 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_convert attributes {} { + func.func public @test_convert(%arg0: tensor<64x128xbf16>) -> tensor<64x128xf32> { + // CHECK-LABEL: func.func public @test_convert + // CHECK: ttnn.typecast + // CHECK-SAME: tensor<64x128xbf16, + // CHECK-SAME: tensor<64x128xf32, + %0 = stablehlo.convert %arg0 : (tensor<64x128xbf16>) -> tensor<64x128xf32> + return %0 : tensor<64x128xf32> + } + + func.func public @test_add(%arg0: tensor<64x128xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xbf16> { + // CHECK-LABEL: func.func public @test_add + // CHECK: ttnn.typecast + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xbf16, + %0 = stablehlo.convert %arg0 : (tensor<64x128xf32>) -> tensor<64x128xbf16> + // CHECK: ttnn.typecast + // CHECK-SAME: tensor<64x128xf32, + // CHECK-SAME: tensor<64x128xbf16, + %1 = stablehlo.convert %arg1 : (tensor<64x128xf32>) -> tensor<64x128xbf16> + %2 = stablehlo.add %0, %1 : tensor<64x128xbf16> + return %2 : tensor<64x128xbf16> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/dot_general_op.mlir b/test/ttmlir/Silicon/StableHLO/dot_general_op.mlir new file mode 100644 index 000000000..57a0bdcd8 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/dot_general_op.mlir @@ -0,0 +1,21 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_dot_general attributes {} { + func.func public @test_dot_general(%arg0 : tensor<16x32xf32>, %arg1 : tensor<32x8xf32>) -> tensor<16x8xf32> { + // CHECK-LABEL: func.func public @test_dot_general + // CHECK: ttnn.empty + // CHECK: ttnn.matmul + // CHECK-SAME: tensor<16x32xf32, + // CHECK-SAME: tensor<32x8xf32, + // CHECK-SAME: tensor<16x8xf32, + // CHECK-SAME: -> tensor<16x8xf32 + %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<16x32xf32>, tensor<32x8xf32>) -> tensor<16x8xf32> + return %0 : tensor<16x8xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/get_dimension_size_op.mlir b/test/ttmlir/Silicon/StableHLO/get_dimension_size_op.mlir new file mode 100644 index 000000000..d7a7bd6cb --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/get_dimension_size_op.mlir @@ -0,0 +1,18 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_get_dimension_size attributes {} { + func.func public @test_get_dimension_size(%arg0: tensor<64x128xf32>) -> tensor { + // CHECK-LABEL: func.func public @test_get_dimension_size + // CHECK: ttnn.full + // CHECK-SAME: {fillValue = 1.280000e+02 : f32} + // CHECK-SAME: -> tensor<1xi32 + %0 = stablehlo.get_dimension_size %arg0, dim = 1 : (tensor<64x128xf32>) -> tensor + return %0 : tensor + } +} diff --git a/test/ttmlir/Silicon/StableHLO/maxpool2d_op.mlir b/test/ttmlir/Silicon/StableHLO/maxpool2d_op.mlir new file mode 100644 index 000000000..e20bb0375 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/maxpool2d_op.mlir @@ -0,0 +1,19 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s +// UNSUPPORTED: true + +func.func public @test_maxpool2d(%arg0: tensor<1x128x128x32xbf16>) -> tensor<1x64x64x32xbf16> { + %0 = stablehlo.constant dense<0xFF80> : tensor + %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor) -> tensor + %2 = "stablehlo.reduce_window"(%arg0, %1) <{padding = dense<[[0, 0], [1, 1], [1, 1], [0, 0]]> : tensor<4x2xi64>, window_dimensions = array, window_strides = array}> ({ + ^bb0(%arg2: tensor, %arg3: tensor): + %3 = stablehlo.maximum %arg2, %arg3 : tensor + stablehlo.return %3 : tensor + }) : (tensor<1x128x128x32xbf16>, tensor) -> tensor<1x64x64x32xbf16> + return %2 : tensor<1x64x64x32xbf16> +} diff --git a/test/ttmlir/Silicon/StableHLO/mnist_inference.mlir b/test/ttmlir/Silicon/StableHLO/mnist_inference.mlir new file mode 100644 index 000000000..849f08bb1 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/mnist_inference.mlir @@ -0,0 +1,44 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s +// TODO: Enable when all ops are supported. +// UNSUPPORTED: true + +module @jit_predict attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} { + func.func public @main(%arg0: tensor<512x784xf32> {mhlo.layout_mode = "default"}, %arg1: tensor<512xf32> {mhlo.layout_mode = "default"}, %arg2: tensor<10x512xf32> {mhlo.layout_mode = "default"}, %arg3: tensor<10xf32> {mhlo.layout_mode = "default"}, %arg4: tensor<128x784xui8> {mhlo.layout_mode = "default"}) -> (tensor<128x10xf32> {jax.result_info = "", mhlo.layout_mode = "default"}) { + %0 = stablehlo.convert %arg0 : tensor<512x784xf32> + %1 = stablehlo.convert %arg4 : (tensor<128x784xui8>) -> tensor<128x784xf32> + %2 = stablehlo.dot_general %0, %1, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<512x784xf32>, tensor<128x784xf32>) -> tensor<512x128xf32> + %3 = stablehlo.transpose %2, dims = [1, 0] : (tensor<512x128xf32>) -> tensor<128x512xf32> + %4 = stablehlo.broadcast_in_dim %arg1, dims = [1] : (tensor<512xf32>) -> tensor<1x512xf32> + %5 = stablehlo.broadcast_in_dim %4, dims = [0, 1] : (tensor<1x512xf32>) -> tensor<128x512xf32> + %6 = stablehlo.add %3, %5 : tensor<128x512xf32> + %cst = stablehlo.constant dense<0.000000e+00> : tensor + %7 = stablehlo.broadcast_in_dim %cst, dims = [] : (tensor) -> tensor<128x512xf32> + %8 = stablehlo.maximum %7, %6 : tensor<128x512xf32> + %9 = stablehlo.dot_general %arg2, %8, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<10x512xf32>, tensor<128x512xf32>) -> tensor<10x128xf32> + %10 = stablehlo.transpose %9, dims = [1, 0] : (tensor<10x128xf32>) -> tensor<128x10xf32> + %11 = stablehlo.broadcast_in_dim %arg3, dims = [1] : (tensor<10xf32>) -> tensor<1x10xf32> + %12 = stablehlo.broadcast_in_dim %11, dims = [0, 1] : (tensor<1x10xf32>) -> tensor<128x10xf32> + %13 = stablehlo.add %10, %12 : tensor<128x10xf32> + %cst_0 = stablehlo.constant dense<0xFF800000> : tensor + %14 = stablehlo.reduce(%13 init: %cst_0) applies stablehlo.maximum across dimensions = [1] : (tensor<128x10xf32>, tensor) -> tensor<128xf32> + %cst_1 = stablehlo.constant dense<0xFF800000> : tensor + %15 = stablehlo.broadcast_in_dim %cst_1, dims = [] : (tensor) -> tensor<128xf32> + %16 = stablehlo.maximum %15, %14 : tensor<128xf32> + %17 = stablehlo.broadcast_in_dim %16, dims = [0] : (tensor<128xf32>) -> tensor<128x1xf32> + %18 = stablehlo.broadcast_in_dim %17, dims = [0, 1] : (tensor<128x1xf32>) -> tensor<128x10xf32> + %19 = stablehlo.subtract %13, %18 : tensor<128x10xf32> + %20 = stablehlo.exponential %19 : tensor<128x10xf32> + %cst_2 = stablehlo.constant dense<0.000000e+00> : tensor + %21 = stablehlo.reduce(%20 init: %cst_2) applies stablehlo.add across dimensions = [1] : (tensor<128x10xf32>, tensor) -> tensor<128xf32> + %22 = stablehlo.broadcast_in_dim %21, dims = [0] : (tensor<128xf32>) -> tensor<128x1xf32> + %23 = stablehlo.broadcast_in_dim %22, dims = [0, 1] : (tensor<128x1xf32>) -> tensor<128x10xf32> + %24 = stablehlo.divide %20, %23 : tensor<128x10xf32> + return %24 : tensor<128x10xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/reduce_add_op.mlir b/test/ttmlir/Silicon/StableHLO/reduce_add_op.mlir new file mode 100644 index 000000000..9da138bbb --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/reduce_add_op.mlir @@ -0,0 +1,22 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s +// UNSUPPORTED: true +// error: keepdim=False is not supported + +module @jit_reduce_add attributes {} { + func.func public @test_reduce_add(%arg0: tensor<128x10xf32>, %cst_0: tensor) -> tensor<128xf32> { + // CHECK-LABEL: func.func public @test_reduce_add + // CHECK: ttnn.sum + // CHECK-SAME: dim_arg = [1 : i32], + // CHECK-SAME: keep_dim = false + // CHECK-SAME: tensor<128x10xf32 + // CHECK-SAME: -> tensor<128xf32, + %0 = stablehlo.reduce(%arg0 init: %cst_0) applies stablehlo.add across dimensions = [1] : (tensor<128x10xf32>, tensor) -> tensor<128xf32> + return %0 : tensor<128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/reduce_maximum_op.mlir b/test/ttmlir/Silicon/StableHLO/reduce_maximum_op.mlir new file mode 100644 index 000000000..57318948e --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/reduce_maximum_op.mlir @@ -0,0 +1,22 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s +// UNSUPPORTED: true +// error: keepdim=False is not supported + +module @jit_reduce_maximum attributes {} { + func.func public @test_reduce_maximum(%arg0: tensor<128x10xf32>, %cst_0: tensor) -> tensor<128xf32> { + // CHECK-LABEL: func.func public @test_reduce_maximum + // CHECK: ttnn.max + // CHECK-SAME: dim_arg = [1 : i32], + // CHECK-SAME: keep_dim = false} + // CHECK-SAME: tensor<128x10xf32, + // CHECK-SAME: -> tensor<128xf32 + %0 = stablehlo.reduce(%arg0 init: %cst_0) applies stablehlo.maximum across dimensions = [1] : (tensor<128x10xf32>, tensor) -> tensor<128xf32> + return %0 : tensor<128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/rehsape_op.mlir b/test/ttmlir/Silicon/StableHLO/rehsape_op.mlir new file mode 100644 index 000000000..8545b0263 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/rehsape_op.mlir @@ -0,0 +1,19 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_module_reshape attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} { + func.func public @test_reshape(%arg0: tensor<1x64x64x64xf32> {mhlo.layout_mode = "default", mhlo.sharding = "{replicated}"}) -> (tensor<1x1x4096x64xf32> {jax.result_info = "", mhlo.layout_mode = "default"}) { + // CHECK-LABEL: func.func public @test_reshape + // CHECK: ttnn.reshape + // CHECK-SAME: {shape = [1 : i32, 1 : i32, 4096 : i32, 64 : i32]} + // CHECK-SAME: tensor<1x64x64x64xf32 + // CHECK-SAME: -> tensor<1x1x4096x64xf32 + %0 = stablehlo.reshape %arg0 : (tensor<1x64x64x64xf32>) -> tensor<1x1x4096x64xf32> + return %0 : tensor<1x1x4096x64xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/scalar_add_op.mlir b/test/ttmlir/Silicon/StableHLO/scalar_add_op.mlir new file mode 100644 index 000000000..bae9b2b4a --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/scalar_add_op.mlir @@ -0,0 +1,21 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_scalar_add attributes {} { + func.func public @test_scalar_add(%arg0: tensor, %arg1: tensor) -> tensor { + // CHECK-LABEL: func.func public @test_scalar_add + // CHECK: ttnn.empty + // CHECK: ttnn.add + // CHECK-SAME: tensor<1xf32, + // CHECK-SAME: tensor<1xf32, + // CHECK-SAME: tensor<1xf32, + // CHECK-SAME: -> tensor<1xf32, + %0 = stablehlo.add %arg0, %arg1 : tensor + return %0 : tensor + } +} diff --git a/test/ttmlir/Silicon/StableHLO/select_op.mlir b/test/ttmlir/Silicon/StableHLO/select_op.mlir new file mode 100644 index 000000000..23b7182ce --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/select_op.mlir @@ -0,0 +1,28 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn + +module @jit_eltwise_select attributes {} { + func.func public @test_select(%arg0: tensor<64x128xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xf32> { + // CHECK-LABEL: func.func public @test_select + // CHECK: tensor.empty + // CHECK: [[EQ:{{0-9}}+]] = "ttnn.eq" + // CHECK-SAME: tensor<64x128xf32 + // CHECK-SAME: tensor<64x128xf32 + // CHECK-SAME: tensor<64x128xbf16 + // CHECK-SAME: -> tensor<64x128xbf16 + %0 = stablehlo.compare EQ, %arg0, %arg1 : (tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xi1> + // CHECK: ttnn.where + // CHECK-SAME: [[EQ]] + // CHECK-SAME: tensor<64x128xbf16 + // CHECK-SAME: tensor<64x128xf32 + // CHECK-SAME: tensor<64x128xf32 + // CHECK-SAME: tensor<64x128xf32 + // CHECK-SAME: -> tensor<64x128xf32 + %1 = stablehlo.select %0, %arg0, %arg1 : (tensor<64x128xi1>, tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xf32> + return %1 : tensor<64x128xf32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/slice_op.mlir b/test/ttmlir/Silicon/StableHLO/slice_op.mlir new file mode 100644 index 000000000..ee5cf91a9 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/slice_op.mlir @@ -0,0 +1,27 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | \ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module @jit_eltwise_subtract attributes {} { + func.func public @test_slice(%arg0: tensor<32x64xbf16>) -> tensor<8x8xbf16> { + // CHECK-LABEL: func.func public @test_slice + // CHECK: ttnn.empty + // CHECK: ttnn.slice + // CHECK-SAME: begins = [0 : i32, 16 : i32], + // CHECK-SAME: ends = [16 : i32, 32 : i32], + // CHECK-SAME: step = [2 : i32, 2 : i32] + // CHECK-SAME: tensor<32x64xbf16, + // CHECK-SAME: tensor<8x8xbf16, + // CHECK-SAME: -> tensor<8x8xbf16 + %result = "stablehlo.slice"(%arg0) { + start_indices = array, + limit_indices = array, + strides = array + } : (tensor<32x64xbf16>) -> tensor<8x8xbf16> + return %result : tensor<8x8xbf16> + } +}