From 995aad9377e814787733da5f0622efefbf24b24c Mon Sep 17 00:00:00 2001 From: Lewis Panos Date: Thu, 6 Mar 2025 18:22:05 -0500 Subject: [PATCH] DenseResourceElementsAttr as constant value E2E (#2381) ### Ticket Closes https://github.com/tenstorrent/tt-mlir/issues/2378 ### Problem description When the value of a `stalbehlo/ttir/ttnn` constant op is a `DenseResourceElementsAttr` instead of a `DenseElementsAttr`, we fail to lower stablehlo to ttir. We also do not yet handle the case where the value of a constant is in a `DenseResourceElementsAttr` when generating the flatbuffer. ### What's changed - Legalize lowering of `stablehlo.constant` when the value type is `DenseResourceElementsAttr` - Handle case where value type is `DenseResourceElementsAttr` when generating flatbuffer - Add verifier to `ttnn.constant` to ensure that the `ElementsAttr` holding the constant data is either `DenseResourceElementsAttr` or `DenseElementsAttr` --- include/ttmlir/Dialect/TTNN/IR/TTNNOps.td | 3 +++ .../StableHLOToTTIRPatterns.cpp | 2 +- lib/Dialect/TTNN/IR/TTNNOps.cpp | 15 +++++++++++ lib/Target/TTNN/TTNNToFlatbuffer.cpp | 15 ++++++++--- .../StableHLOToTTIR/constant_op.mlir | 18 +++++++++++++ .../simple_constant_sparse_resource.mlir | 9 +++++++ .../positive}/simple_constant.mlir | 18 +++++++++++++ .../n150/Constant/constant_bf16.mlir | 3 +-- .../n150/Constant/constant_bool.mlir | 3 +-- .../constant_dense_resource_bf16.mlir | 26 +++++++++++++++++++ .../StableHLO/n150/Constant/constant_f32.mlir | 3 +-- .../StableHLO/n150/Constant/constant_f64.mlir | 3 +-- .../StableHLO/n150/Constant/constant_i16.mlir | 3 +-- .../StableHLO/n150/Constant/constant_i32.mlir | 3 +-- .../StableHLO/n150/Constant/constant_i64.mlir | 3 +-- .../n150/Constant/constant_ui16.mlir | 3 +-- .../n150/Constant/constant_ui32.mlir | 3 +-- .../n150/Constant/constant_ui64.mlir | 3 +-- 18 files changed, 112 insertions(+), 24 deletions(-) create mode 100644 test/ttmlir/Dialect/TTNN/constant/negative/simple_constant_sparse_resource.mlir rename test/ttmlir/Dialect/TTNN/{ => constant/positive}/simple_constant.mlir (90%) create mode 100644 test/ttmlir/Silicon/StableHLO/n150/Constant/constant_dense_resource_bf16.mlir diff --git a/include/ttmlir/Dialect/TTNN/IR/TTNNOps.td b/include/ttmlir/Dialect/TTNN/IR/TTNNOps.td index 62bf275300..68674f5ee8 100644 --- a/include/ttmlir/Dialect/TTNN/IR/TTNNOps.td +++ b/include/ttmlir/Dialect/TTNN/IR/TTNNOps.td @@ -15,6 +15,7 @@ include "mlir/Interfaces/InferTypeOpInterface.td" include "mlir/Interfaces/DestinationStyleOpInterface.td" include "mlir/Interfaces/ControlFlowInterfaces.td" include "mlir/Interfaces/SideEffectInterfaces.td" +include "mlir/IR/BuiltinAttributes.td" include "mlir/IR/CommonTypeConstraints.td" include "mlir/IR/CommonAttrConstraints.td" @@ -1604,6 +1605,8 @@ def TTNN_ConstantOp : TTNN_Op<"constant", [AllShapesMatch<["value", "result"]>]> return wa::TTNNOperandsWorkaroundsFactory::createConstantOpOperandsWorkarounds(); } }]; + + let hasVerifier = 1; } #endif diff --git a/lib/Conversion/StableHLOToTTIR/StableHLOToTTIRPatterns.cpp b/lib/Conversion/StableHLOToTTIR/StableHLOToTTIRPatterns.cpp index de1c9b1ba9..3d790fbfa8 100644 --- a/lib/Conversion/StableHLOToTTIR/StableHLOToTTIRPatterns.cpp +++ b/lib/Conversion/StableHLOToTTIR/StableHLOToTTIRPatterns.cpp @@ -636,7 +636,7 @@ class StableHLOToTTIRConstantOpConversionPattern private: LogicalResult checkBasicLegality(mlir::stablehlo::ConstantOp &srcOp, ConversionPatternRewriter &rewriter) const { - if (srcOp.getValue().getShapedType().getShape().empty() && + if (isa(srcOp.getValue()) && !srcOp.getValue().getElementType().isIntOrFloat()) { return rewriter.notifyMatchFailure(srcOp, "Unsupported element type."); } diff --git a/lib/Dialect/TTNN/IR/TTNNOps.cpp b/lib/Dialect/TTNN/IR/TTNNOps.cpp index 476530265c..27987066be 100644 --- a/lib/Dialect/TTNN/IR/TTNNOps.cpp +++ b/lib/Dialect/TTNN/IR/TTNNOps.cpp @@ -10,6 +10,7 @@ #include "ttmlir/Utils.h" #include "mlir/Dialect/Traits.h" +#include "mlir/IR/BuiltinAttributes.h" #include #include @@ -19,6 +20,20 @@ namespace mlir::tt::ttnn { +//===----------------------------------------------------------------------===// +// ConstantOp +//===----------------------------------------------------------------------===// + +::mlir::LogicalResult mlir::tt::ttnn::ConstantOp::verify() { + + if (!isa(getValue())) { + return emitOpError("value attribute must be one of " + "DenseResourceElementsAttr or DenseElementsAttr."); + } + + return success(); +} + //===----------------------------------------------------------------------===// // ClampOp //===----------------------------------------------------------------------===// diff --git a/lib/Target/TTNN/TTNNToFlatbuffer.cpp b/lib/Target/TTNN/TTNNToFlatbuffer.cpp index 3f8189d8c0..7420916ab2 100644 --- a/lib/Target/TTNN/TTNNToFlatbuffer.cpp +++ b/lib/Target/TTNN/TTNNToFlatbuffer.cpp @@ -944,10 +944,19 @@ ::flatbuffers::Offset<::tt::target::ttnn::ConstantOp> createOp(FlatbufferObjectCache &cache, ttnn::ConstantOp op) { auto output = cache.getOrCreate(op.getResult(), tensorValueToFlatbuffer, kHostAllocatedSize); + std::vector rawVector; + if (auto data = + mlir::dyn_cast(op.getValue())) { + ArrayRef rawData = data.getData(); + rawVector = std::vector(rawData.begin(), rawData.end()); + } else if (auto data = + mlir::dyn_cast(op.getValue())) { + ArrayRef rawData = data.getRawData(); + rawVector = std::vector(rawData.begin(), rawData.end()); + } else { + llvm_unreachable("Unknown constant value attribute type"); + } - auto rawData = - mlir::dyn_cast(op.getValue()).getRawData(); - auto rawVector = std::vector(rawData.begin(), rawData.end()); return ::tt::target::ttnn::CreateConstantOpDirect(*cache.fbb, output, &rawVector); } diff --git a/test/ttmlir/Conversion/StableHLOToTTIR/constant_op.mlir b/test/ttmlir/Conversion/StableHLOToTTIR/constant_op.mlir index 1ec5c11371..7257ba8966 100644 --- a/test/ttmlir/Conversion/StableHLOToTTIR/constant_op.mlir +++ b/test/ttmlir/Conversion/StableHLOToTTIR/constant_op.mlir @@ -339,4 +339,22 @@ module @jit_constant attributes {} { // CHECK: return %[[CONSTANT]] : tensor<1xi32> return %0 : tensor } + + func.func @test_dense_attr() -> tensor<1x2xbf16> { + // CHECK: %{{[0-9]+}} = "ttir.constant"() <{value = dense_resource : tensor<1x2xbf16>}> : () -> tensor<1x2xbf16> + %0 = stablehlo.constant dense_resource : tensor<1x2xbf16> + // CHECK: return %{{[0-9]+}} : tensor<1x2xbf16> + return %0 : tensor<1x2xbf16> + } } +{-# + dialect_resources: { + builtin: { + // This should encode for two bfloat16 values which are both 2.0 + // 0x020000000 is a hex string blob + // 0x0040 is 2.0 in bfloat16 + // 0x00400040 is 2.0, 2.0 + dense_attr: "0x0200000000400040" + } + } +#-} diff --git a/test/ttmlir/Dialect/TTNN/constant/negative/simple_constant_sparse_resource.mlir b/test/ttmlir/Dialect/TTNN/constant/negative/simple_constant_sparse_resource.mlir new file mode 100644 index 0000000000..071b76489c --- /dev/null +++ b/test/ttmlir/Dialect/TTNN/constant/negative/simple_constant_sparse_resource.mlir @@ -0,0 +1,9 @@ +// RUN: not ttmlir-opt --ttir-to-ttnn-backend-pipeline %s 2>&1 | FileCheck %s + +module attributes {} { + func.func @test_dense_attr() -> tensor<1x2xbf16> { + // CHECK: error: 'ttnn.constant' op value attribute must be one of DenseResourceElementsAttr or DenseElementsAttr. + %0 = "ttir.constant"() <{value = sparse<[[0, 0], [0, 1]], [2.0, 2.0]> : tensor<1x2xbf16>}> : () -> tensor<1x2xbf16> + return %0 : tensor<1x2xbf16> + } +} diff --git a/test/ttmlir/Dialect/TTNN/simple_constant.mlir b/test/ttmlir/Dialect/TTNN/constant/positive/simple_constant.mlir similarity index 90% rename from test/ttmlir/Dialect/TTNN/simple_constant.mlir rename to test/ttmlir/Dialect/TTNN/constant/positive/simple_constant.mlir index cb085cd71b..2f5bfab502 100644 --- a/test/ttmlir/Dialect/TTNN/simple_constant.mlir +++ b/test/ttmlir/Dialect/TTNN/constant/positive/simple_constant.mlir @@ -137,4 +137,22 @@ module attributes {} { %0 = "ttir.constant"() <{value = dense<[[[-1, 2, 3]]]> : tensor<1x1x3xi32>}> : () -> tensor<1x1x3xi32> return %0 : tensor<1x1x3xi32> } + + func.func @test_dense_attr() -> tensor<1x2xbf16> { + %0 = "ttir.constant"() <{value = dense_resource : tensor<1x2xbf16>}> : () -> tensor<1x2xbf16> + // CHECK: "ttnn.constant" + // CHECK-SAME: value = dense_resource + return %0 : tensor<1x2xbf16> + } } +{-# + dialect_resources: { + builtin: { + // This should encode for two bfloat16 values which are both 2.0 + // 0x020000000 is a hex string blob + // 0x0040 is 2.0 in bfloat16 + // 0x00400040 is 2.0, 2.0 + dense_attr: "0x0200000000400040" + } + } +#-} diff --git a/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_bf16.mlir b/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_bf16.mlir index 636ea27167..6a9e01e101 100644 --- a/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_bf16.mlir +++ b/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_bf16.mlir @@ -1,8 +1,7 @@ // 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-opt --stablehlo-to-ttir-pipeline --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" %s > %t.mlir // RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn // RUN: FileCheck --input-file=%t.mlir %s diff --git a/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_bool.mlir b/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_bool.mlir index 6486ff99c6..98981102b0 100644 --- a/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_bool.mlir +++ b/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_bool.mlir @@ -1,8 +1,7 @@ // 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-opt --stablehlo-to-ttir-pipeline --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" %s > %t.mlir // RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn // RUN: FileCheck --input-file=%t.mlir %s diff --git a/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_dense_resource_bf16.mlir b/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_dense_resource_bf16.mlir new file mode 100644 index 0000000000..2e592db7eb --- /dev/null +++ b/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_dense_resource_bf16.mlir @@ -0,0 +1,26 @@ +// REQUIRES: stablehlo +// RUN: rm -rf %t.ttnn +// RUN: rm -rf %t.mlir +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" %s > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +// RUN: FileCheck --input-file=%t.mlir %s + +module attributes {} { + func.func @test_dense_attr() -> tensor<1x2xbf16> { + // CHECK: ttnn.constant + // CHECK-SAME: dense_resource + %0 = stablehlo.constant dense_resource : tensor<1x2xbf16> + return %0 : tensor<1x2xbf16> + } +} +{-# + dialect_resources: { + builtin: { + // This should encode for two bfloat16 values which are both 2.0 + // 0x020000000 is a hex string blob + // 0x0040 is 2.0 in bfloat16 + // 0x00400040 is 2.0, 2.0 + dense_attr: "0x0200000000400040" + } + } +#-} diff --git a/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_f32.mlir b/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_f32.mlir index 3fecd90fb0..778169cef5 100644 --- a/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_f32.mlir +++ b/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_f32.mlir @@ -1,8 +1,7 @@ // 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-opt --stablehlo-to-ttir-pipeline --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" %s > %t.mlir // RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn // RUN: FileCheck --input-file=%t.mlir %s diff --git a/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_f64.mlir b/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_f64.mlir index c286745a09..8f14aaa224 100644 --- a/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_f64.mlir +++ b/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_f64.mlir @@ -1,8 +1,7 @@ // 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-opt --stablehlo-to-ttir-pipeline --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" %s > %t.mlir // RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn // RUN: FileCheck --input-file=%t.mlir %s diff --git a/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_i16.mlir b/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_i16.mlir index 9da6371060..37b1b4d925 100644 --- a/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_i16.mlir +++ b/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_i16.mlir @@ -1,8 +1,7 @@ // 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-opt --stablehlo-to-ttir-pipeline --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" %s > %t.mlir // RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn // RUN: FileCheck --input-file=%t.mlir %s diff --git a/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_i32.mlir b/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_i32.mlir index 8750a832bd..44318a2c6a 100644 --- a/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_i32.mlir +++ b/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_i32.mlir @@ -1,8 +1,7 @@ // 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-opt --stablehlo-to-ttir-pipeline --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" %s > %t.mlir // RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn // RUN: FileCheck --input-file=%t.mlir %s diff --git a/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_i64.mlir b/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_i64.mlir index b19421085a..0f0bf0da7a 100644 --- a/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_i64.mlir +++ b/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_i64.mlir @@ -1,8 +1,7 @@ // 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-opt --stablehlo-to-ttir-pipeline --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" %s > %t.mlir // RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn // RUN: FileCheck --input-file=%t.mlir %s diff --git a/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_ui16.mlir b/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_ui16.mlir index 0144da3084..8117503740 100644 --- a/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_ui16.mlir +++ b/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_ui16.mlir @@ -1,8 +1,7 @@ // 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-opt --stablehlo-to-ttir-pipeline --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" %s > %t.mlir // RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn // RUN: FileCheck --input-file=%t.mlir %s diff --git a/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_ui32.mlir b/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_ui32.mlir index 029d45f9ce..80ff3c6176 100644 --- a/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_ui32.mlir +++ b/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_ui32.mlir @@ -1,8 +1,7 @@ // 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-opt --stablehlo-to-ttir-pipeline --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" %s > %t.mlir // RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn // RUN: FileCheck --input-file=%t.mlir %s diff --git a/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_ui64.mlir b/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_ui64.mlir index 8f24ebdfe8..c0832fd918 100644 --- a/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_ui64.mlir +++ b/test/ttmlir/Silicon/StableHLO/n150/Constant/constant_ui64.mlir @@ -1,8 +1,7 @@ // 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-opt --stablehlo-to-ttir-pipeline --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" %s > %t.mlir // RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn // RUN: FileCheck --input-file=%t.mlir %s