From b5c028f8138fedcf7caa43c2c2245ad1deef7aad Mon Sep 17 00:00:00 2001 From: Kristijan Mitrovic Date: Wed, 15 Jan 2025 17:42:26 +0000 Subject: [PATCH] Code review --- .../StableHLOToTTIR/StableHLOToTTIRPass.cpp | 13 ++- .../ArithToStableHLO/constant_op.mlir | 61 +++++++++++- .../StableHLOToTTIR/constant_op.mlir | 92 +++++++++++++++++++ test/ttmlir/Dialect/TTNN/simple_constant.mlir | 14 +++ .../StableHLO/Constant/constant_ui16.mlir | 43 +++++++++ .../StableHLO/Constant/constant_ui32.mlir | 43 +++++++++ .../StableHLO/Constant/constant_ui64.mlir | 43 +++++++++ test/ttmlir/Silicon/TTNN/simple_constant.mlir | 13 ++- 8 files changed, 313 insertions(+), 9 deletions(-) create mode 100644 test/ttmlir/Silicon/StableHLO/Constant/constant_ui16.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Constant/constant_ui32.mlir create mode 100644 test/ttmlir/Silicon/StableHLO/Constant/constant_ui64.mlir diff --git a/lib/Conversion/StableHLOToTTIR/StableHLOToTTIRPass.cpp b/lib/Conversion/StableHLOToTTIR/StableHLOToTTIRPass.cpp index 8fc3538421..6a6e7e9d52 100644 --- a/lib/Conversion/StableHLOToTTIR/StableHLOToTTIRPass.cpp +++ b/lib/Conversion/StableHLOToTTIR/StableHLOToTTIRPass.cpp @@ -58,12 +58,19 @@ class StablehloTypeConverter : public TypeConverter { changed = true; } else if (bitWidth == 64) { // Convert 64 bit integer element type to 32 bit integer. - if (isa(type.getElementType())) { - elementType = IntegerType::get(context, 32); + // If element is unsigned, we explicitly assign it Unsigned semantics to + // it (like `ui32`). Otherwise, we don't explicitly use Signed semantics + // (like `si32`), but rather Signless (like `i32`) which is the default. + if (isa(elementType)) { + elementType = IntegerType::get( + context, 32, + elementType.isUnsignedInteger() + ? IntegerType::SignednessSemantics::Unsigned + : IntegerType::SignednessSemantics::Signless); changed = true; } // Convert 64 bit float element type to 32 bit float. - else if (isa(type.getElementType())) { + else if (isa(elementType)) { elementType = FloatType::getF32(context); changed = true; } diff --git a/test/ttmlir/Conversion/ArithToStableHLO/constant_op.mlir b/test/ttmlir/Conversion/ArithToStableHLO/constant_op.mlir index 0cbe0385d0..b388daeb15 100644 --- a/test/ttmlir/Conversion/ArithToStableHLO/constant_op.mlir +++ b/test/ttmlir/Conversion/ArithToStableHLO/constant_op.mlir @@ -1,15 +1,66 @@ // REQUIRES: stablehlo // RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | FileCheck %s module @jit_constant attributes {} { - func.func public @test_splat() -> tensor<64xf32> { - %0 = arith.constant dense<0.3> : tensor<64xf32> - // CHECK: %[[C:.*]] = "ttir.constant"[[C:.*]] + func.func public @test_scalar_float() -> tensor { + // CHECK: %{{[0-9]+}} = "ttir.constant"() <{value = dense<3.000000e+00> : tensor<1xf32>}> : () -> tensor<1xf32> + %0 = arith.constant dense<3.0> : tensor + // CHECK: return %{{[0-9]+}} : tensor<1xf32> + return %0 : tensor + } + + func.func public @test_splat_float() -> tensor<64xf32> { + // CHECK: %{{[0-9]+}} = "ttir.constant"() <{value = dense<3.000000e+00> : tensor<64xf32>}> : () -> tensor<64xf32> + %0 = arith.constant dense<3.0> : tensor<64xf32> + // CHECK: return %{{[0-9]+}} : tensor<64xf32> return %0 : tensor<64xf32> } - func.func public @test_multiple() -> tensor<2x2xf32> { + func.func public @test_multiple_float() -> tensor<2x2xf32> { + // CHECK: %{{[0-9]+}} = "ttir.constant"() <{value = dense<{{([[])}}[0.000000e+00, 1.000000e+00], [2.000000e+00, 3.000000e+00]]> : tensor<2x2xf32>}> : () -> tensor<2x2xf32> %0 = arith.constant dense<[[0.0, 1.0], [2.0, 3.0]]> : tensor<2x2xf32> - // CHECK: %[[C:.*]] = "ttir.constant"[[C:.*]] + // CHECK: return %{{[0-9]+}} : tensor<2x2xf32> return %0 : tensor<2x2xf32> } + + func.func public @test_scalar_int() -> tensor { + // CHECK: %{{[0-9]+}} = "ttir.constant"() <{value = dense<3> : tensor<1xi32>}> : () -> tensor<1xi32> + %0 = arith.constant dense<3> : tensor + // CHECK: return %{{[0-9]+}} : tensor<1xi32> + return %0 : tensor + } + + func.func public @test_splat_int() -> tensor<64xi32> { + // CHECK: %{{[0-9]+}} = "ttir.constant"() <{value = dense<3> : tensor<64xi32>}> : () -> tensor<64xi32> + %0 = arith.constant dense<3> : tensor<64xi32> + // CHECK: return %{{[0-9]+}} : tensor<64xi32> + return %0 : tensor<64xi32> + } + + func.func public @test_multiple_int() -> tensor<2x2xi32> { + // CHECK: %{{[0-9]+}} = "ttir.constant"() <{value = dense<{{([[])}}[0, 1], [2, 3]]> : tensor<2x2xi32>}> : () -> tensor<2x2xi32> + %0 = arith.constant dense<[[0, 1], [2, 3]]> : tensor<2x2xi32> + // CHECK: return %{{[0-9]+}} : tensor<2x2xi32> + return %0 : tensor<2x2xi32> + } + + func.func public @test_scalar_uint() -> tensor { + // CHECK: %{{[0-9]+}} = "ttir.constant"() <{value = dense<3> : tensor<1xui32>}> : () -> tensor<1xui32> + %0 = arith.constant dense<3> : tensor + // CHECK: return %{{[0-9]+}} : tensor<1xui32> + return %0 : tensor + } + + func.func public @test_splat_uint() -> tensor<64xui32> { + // CHECK: %{{[0-9]+}} = "ttir.constant"() <{value = dense<3> : tensor<64xui32>}> : () -> tensor<64xui32> + %0 = arith.constant dense<3> : tensor<64xui32> + // CHECK: return %{{[0-9]+}} : tensor<64xui32> + return %0 : tensor<64xui32> + } + + func.func public @test_multiple_uint() -> tensor<2x2xui32> { + // CHECK: %{{[0-9]+}} = "ttir.constant"() <{value = dense<{{([[])}}[0, 1], [2, 3]]> : tensor<2x2xui32>}> : () -> tensor<2x2xui32> + %0 = arith.constant dense<[[0, 1], [2, 3]]> : tensor<2x2xui32> + // CHECK: return %{{[0-9]+}} : tensor<2x2xui32> + return %0 : tensor<2x2xui32> + } } diff --git a/test/ttmlir/Conversion/StableHLOToTTIR/constant_op.mlir b/test/ttmlir/Conversion/StableHLOToTTIR/constant_op.mlir index eb0aa5b951..878d9f1652 100644 --- a/test/ttmlir/Conversion/StableHLOToTTIR/constant_op.mlir +++ b/test/ttmlir/Conversion/StableHLOToTTIR/constant_op.mlir @@ -214,4 +214,96 @@ module @jit_constant attributes {} { // CHECK: return %{{[0-9]+}} : tensor<2x2xi32> return %0 : tensor<2x2xi64> } + + func.func public @test_uint8_scalar() -> tensor { + // CHECK: %{{[0-9]+}} = "ttir.constant"() <{value = dense<3> : tensor<1xui8>}> : () -> tensor<1xui8> + %0 = stablehlo.constant dense<3> : tensor + // CHECK: return %{{[0-9]+}} : tensor<1xui8> + return %0 : tensor + } + + func.func public @test_uint8_splat() -> tensor<64xui8> { + // CHECK: %{{[0-9]+}} = "ttir.constant"() <{value = dense<3> : tensor<64xui8>}> : () -> tensor<64xui8> + %0 = stablehlo.constant dense<3> : tensor<64xui8> + // CHECK: return %{{[0-9]+}} : tensor<64xui8> + return %0 : tensor<64xui8> + } + + func.func public @test_uint8_multiple() -> tensor<2x2xui8> { + // The ugly regex after `dense` is necessary because double square opening + // brackets indicate substitution block in FileCheck syntax. + // CHECK: %{{[0-9]+}} = "ttir.constant"() <{value = dense<{{([[])}}[0, 1], [2, 3]]> : tensor<2x2xui8>}> : () -> tensor<2x2xui8> + %0 = stablehlo.constant dense<[[0, 1], [2, 3]]> : tensor<2x2xui8> + // CHECK: return %{{[0-9]+}} : tensor<2x2xui8> + return %0 : tensor<2x2xui8> + } + + func.func public @test_uint16_scalar() -> tensor { + // CHECK: %{{[0-9]+}} = "ttir.constant"() <{value = dense<3> : tensor<1xui16>}> : () -> tensor<1xui16> + %0 = stablehlo.constant dense<3> : tensor + // CHECK: return %{{[0-9]+}} : tensor<1xui16> + return %0 : tensor + } + + func.func public @test_uint16_splat() -> tensor<64xui16> { + // CHECK: %{{[0-9]+}} = "ttir.constant"() <{value = dense<3> : tensor<64xui16>}> : () -> tensor<64xui16> + %0 = stablehlo.constant dense<3> : tensor<64xui16> + // CHECK: return %{{[0-9]+}} : tensor<64xui16> + return %0 : tensor<64xui16> + } + + func.func public @test_uint16_multiple() -> tensor<2x2xui16> { + // The ugly regex after `dense` is necessary because double square opening + // brackets indicate substitution block in FileCheck syntax. + // CHECK: %{{[0-9]+}} = "ttir.constant"() <{value = dense<{{([[])}}[0, 1], [2, 3]]> : tensor<2x2xui16>}> : () -> tensor<2x2xui16> + %0 = stablehlo.constant dense<[[0, 1], [2, 3]]> : tensor<2x2xui16> + // CHECK: return %{{[0-9]+}} : tensor<2x2xui16> + return %0 : tensor<2x2xui16> + } + + func.func public @test_uint32_scalar() -> tensor { + // CHECK: %{{[0-9]+}} = "ttir.constant"() <{value = dense<3> : tensor<1xui32>}> : () -> tensor<1xui32> + %0 = stablehlo.constant dense<3> : tensor + // CHECK: return %{{[0-9]+}} : tensor<1xui32> + return %0 : tensor + } + + func.func public @test_uint32_splat() -> tensor<64xui32> { + // CHECK: %{{[0-9]+}} = "ttir.constant"() <{value = dense<3> : tensor<64xui32>}> : () -> tensor<64xui32> + %0 = stablehlo.constant dense<3> : tensor<64xui32> + // CHECK: return %{{[0-9]+}} : tensor<64xui32> + return %0 : tensor<64xui32> + } + + func.func public @test_uint32_multiple() -> tensor<2x2xui32> { + // The ugly regex after `dense` is necessary because double square opening + // brackets indicate substitution block in FileCheck syntax. + // CHECK: %{{[0-9]+}} = "ttir.constant"() <{value = dense<{{([[])}}[0, 1], [2, 3]]> : tensor<2x2xui32>}> : () -> tensor<2x2xui32> + %0 = stablehlo.constant dense<[[0, 1], [2, 3]]> : tensor<2x2xui32> + // CHECK: return %{{[0-9]+}} : tensor<2x2xui32> + return %0 : tensor<2x2xui32> + } + + func.func public @test_uint64_scalar() -> tensor { + // CHECK: %{{[0-9]+}} = "ttir.constant"() <{value = dense<3> : tensor<1xui32>}> : () -> tensor<1xui32> + %0 = stablehlo.constant dense<3> : tensor + // CHECK: return %{{[0-9]+}} : tensor<1xui32> + return %0 : tensor + } + + func.func public @test_uint64_splat() -> tensor<64xui64> { + // CHECK: %{{[0-9]+}} = "ttir.constant"() <{value = dense<3> : tensor<64xui32>}> : () -> tensor<64xui32> + %0 = stablehlo.constant dense<3> : tensor<64xui64> + // CHECK: return %{{[0-9]+}} : tensor<64xui32> + return %0 : tensor<64xui64> + } + + func.func public @test_uint64_multiple() -> tensor<2x2xui64> { + // The ugly regex after `dense` is necessary because double square opening + // brackets indicate substitution block in FileCheck syntax. + // CHECK: %{{[0-9]+}} = "ttir.constant"() <{value = dense<{{([[])}}[0, 1], [2, 3]]> : tensor<2x2xui32>}> : () -> tensor<2x2xui32> + %0 = stablehlo.constant dense<[[0, 1], [2, 3]]> : tensor<2x2xui64> + // CHECK: return %{{[0-9]+}} : tensor<2x2xui32> + return %0 : tensor<2x2xui64> + } } diff --git a/test/ttmlir/Dialect/TTNN/simple_constant.mlir b/test/ttmlir/Dialect/TTNN/simple_constant.mlir index 53de9a5ee1..b212164905 100644 --- a/test/ttmlir/Dialect/TTNN/simple_constant.mlir +++ b/test/ttmlir/Dialect/TTNN/simple_constant.mlir @@ -18,6 +18,12 @@ module attributes {} { return %0 : tensor<64x128xi32> } + func.func @test_empty_uint() -> tensor<64x128xui32> { + %0 = "ttir.constant"() <{value = dense<0> : tensor<64x128xui32>}> : () -> tensor<64x128xui32> + // CHECK: %{{[0-9]+}} = "ttnn.full" + return %0 : tensor<64x128xui32> + } + func.func @test_empty_bfloat16() -> tensor<64x128xbf16> { %0 = "ttir.constant"() <{value = dense<0.000000e+00> : tensor<64x128xbf16>}> : () -> tensor<64x128xbf16> // CHECK: %{{[0-9]+}} = "ttnn.full" @@ -54,6 +60,14 @@ module attributes {} { return %0 : tensor<64x128xi32> } + func.func @test_full_uint() -> tensor<64x128xui32> { + // CHECK: %{{[0-9]+}} = "ttnn.full" + // CHECK-SAME: fillValue = 1.000000e+00 : f32 + // CHECK-SAME: tensor<64x128xui32 + %0 = "ttir.constant"() <{value = dense<1> : tensor<64x128xui32>}> : () -> tensor<64x128xui32> + return %0 : tensor<64x128xui32> + } + func.func @test_full_bfloat16() -> tensor<64x128xbf16> { // CHECK: %{{[0-9]+}} = "ttnn.full" // CHECK-SAME: fillValue = 1.000000e+00 : f32 diff --git a/test/ttmlir/Silicon/StableHLO/Constant/constant_ui16.mlir b/test/ttmlir/Silicon/StableHLO/Constant/constant_ui16.mlir new file mode 100644 index 0000000000..0144da3084 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Constant/constant_ui16.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_uint16_scalar() -> tensor { + // CHECK-LABEL: func.func public @test_uint16_scalar + // CHECK: ttnn.full + // CHECK-SAME: fillValue = 3.000000e+00 : f32 + // CHECK-SAME: -> tensor<1xui16 + %0 = stablehlo.constant dense<3> : tensor + return %0 : tensor + } + + func.func public @test_uint16_scalar_empty() -> tensor { + // CHECK-LABEL: func.func public @test_uint16_scalar_empty + // CHECK: ttnn.full + // CHECK-SAME: -> tensor<1xui16 + %0 = stablehlo.constant dense<0> : tensor + return %0 : tensor + } + + func.func public @test_uint16_empty() -> tensor<64x128xui16> { + // CHECK-LABEL: func.func public @test_uint16_empty + // CHECK: ttnn.full + // CHECK-SAME: -> tensor<64x128xui16 + %0 = stablehlo.constant dense<0> : tensor<64x128xui16> + return %0 : tensor<64x128xui16> + } + + func.func public @test_uint16_splat() -> tensor<64x128xui16> { + // CHECK-LABEL: func.func public @test_uint16_splat + // CHECK: ttnn.full + // CHECK-SAME: fillValue = 3.000000e+00 : f32 + // CHECK-SAME: -> tensor<64x128xui16 + %0 = stablehlo.constant dense<3> : tensor<64x128xui16> + return %0 : tensor<64x128xui16> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Constant/constant_ui32.mlir b/test/ttmlir/Silicon/StableHLO/Constant/constant_ui32.mlir new file mode 100644 index 0000000000..029d45f9ce --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Constant/constant_ui32.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_uint32_scalar() -> tensor { + // CHECK-LABEL: func.func public @test_uint32_scalar + // CHECK: ttnn.full + // CHECK-SAME: fillValue = 3.000000e+00 : f32 + // CHECK-SAME: -> tensor<1xui32 + %0 = stablehlo.constant dense<3> : tensor + return %0 : tensor + } + + func.func public @test_uint32_scalar_empty() -> tensor { + // CHECK-LABEL: func.func public @test_uint32_scalar_empty + // CHECK: ttnn.full + // CHECK-SAME: -> tensor<1xui32 + %0 = stablehlo.constant dense<0> : tensor + return %0 : tensor + } + + func.func public @test_uint32_empty() -> tensor<64x128xui32> { + // CHECK-LABEL: func.func public @test_uint32_empty + // CHECK: ttnn.full + // CHECK-SAME: -> tensor<64x128xui32 + %0 = stablehlo.constant dense<0> : tensor<64x128xui32> + return %0 : tensor<64x128xui32> + } + + func.func public @test_uint32_splat() -> tensor<64x128xui32> { + // CHECK-LABEL: func.func public @test_uint32_splat + // CHECK: ttnn.full + // CHECK-SAME: fillValue = 3.000000e+00 : f32 + // CHECK-SAME: -> tensor<64x128xui32 + %0 = stablehlo.constant dense<3> : tensor<64x128xui32> + return %0 : tensor<64x128xui32> + } +} diff --git a/test/ttmlir/Silicon/StableHLO/Constant/constant_ui64.mlir b/test/ttmlir/Silicon/StableHLO/Constant/constant_ui64.mlir new file mode 100644 index 0000000000..8f24ebdfe8 --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/Constant/constant_ui64.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_uint64_scalar() -> tensor { + // CHECK-LABEL: func.func public @test_uint64_scalar + // CHECK: ttnn.full + // CHECK-SAME: fillValue = 3.000000e+00 : f32 + // CHECK-SAME: -> tensor<1xui32 + %0 = stablehlo.constant dense<3> : tensor + return %0 : tensor + } + + func.func public @test_uint64_scalar_empty() -> tensor { + // CHECK-LABEL: func.func public @test_uint64_scalar_empty + // CHECK: ttnn.full + // CHECK-SAME: -> tensor<1xui32 + %0 = stablehlo.constant dense<0> : tensor + return %0 : tensor + } + + func.func public @test_uint64_empty() -> tensor<64x128xui32> { + // CHECK-LABEL: func.func public @test_uint64_empty + // CHECK: ttnn.full + // CHECK-SAME: -> tensor<64x128xui32 + %0 = stablehlo.constant dense<0> : tensor<64x128xui32> + return %0 : tensor<64x128xui32> + } + + func.func public @test_uint64_splat() -> tensor<64x128xui32> { + // CHECK-LABEL: func.func public @test_uint64_splat + // CHECK: ttnn.full + // CHECK-SAME: fillValue = 3.000000e+00 : f32 + // CHECK-SAME: -> tensor<64x128xui32 + %0 = stablehlo.constant dense<3> : tensor<64x128xui32> + return %0 : tensor<64x128xui32> + } +} diff --git a/test/ttmlir/Silicon/TTNN/simple_constant.mlir b/test/ttmlir/Silicon/TTNN/simple_constant.mlir index 35728f0a93..5e135e4521 100644 --- a/test/ttmlir/Silicon/TTNN/simple_constant.mlir +++ b/test/ttmlir/Silicon/TTNN/simple_constant.mlir @@ -8,6 +8,12 @@ module @sysmem_creation attributes {} { return %0 : tensor<64x128xi32> } + func.func @test_empty_uint() -> tensor<64x128xui32> { + %0 = "ttir.constant"() <{value = dense<0> : tensor<64x128xui32>}> : () -> tensor<64x128xui32> + // CHECK: %[[C:.*]] = "ttnn.full"[[C:.*]] + return %0 : tensor<64x128xui32> + } + func.func @test_empty_float() -> tensor<64x128xf32> { %0 = "ttir.constant"() <{value = dense<0.000000e+00> : tensor<64x128xf32>}> : () -> tensor<64x128xf32> // CHECK: %[[C:.*]] = "ttnn.full"[[C:.*]] @@ -20,13 +26,18 @@ module @sysmem_creation attributes {} { return %0 : tensor<1x1xf32> } - func.func @test_full_int() -> tensor<64x128xi32> { %0 = "ttir.constant"() <{value = dense<1> : tensor<64x128xi32>}> : () -> tensor<64x128xi32> // CHECK: %[[C:.*]] = "ttnn.full"[[C:.*]] return %0 : tensor<64x128xi32> } + func.func @test_full_uint() -> tensor<64x128xui32> { + %0 = "ttir.constant"() <{value = dense<1> : tensor<64x128xui32>}> : () -> tensor<64x128xui32> + // CHECK: %[[C:.*]] = "ttnn.full"[[C:.*]] + return %0 : tensor<64x128xui32> + } + func.func @test_full_float() -> tensor<64x128xf32> { %0 = "ttir.constant"() <{value = dense<1.000000e+00> : tensor<64x128xf32>}> : () -> tensor<64x128xf32> // CHECK: %[[C:.*]] = "ttnn.full"[[C:.*]]