Handle different value types for constant op explicitly for TTNN backend

tenstorrent · Nov 6, 2024 · 2759638 · 2759638
1 parent 90c5561
commit 2759638
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 4 deletions.
diff --git a/lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp b/lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp
@@ -23,7 +23,6 @@
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/LogicalResult.h"
-#include <cstdint>
 
 using namespace mlir;
 using namespace mlir::tt;
@@ -506,8 +505,7 @@ class ConstantOpConversionPattern
       Value device = getOrInsertDevice(rewriter, op);
       float fillValue =
           valueAttr.getElementType().isInteger()
-              ? getIntegerValue(
-                    valueAttr) // static_cast<float>(valueAttr.getSplatValue<int>())
+              ? getIntegerValue(valueAttr)
               : valueAttr.getSplatValue<mlir::APFloat>().convertToFloat();
       if (fillValue == 0) {
         rewriter.replaceOpWithNewOp<tensor::EmptyOp>(
@@ -554,7 +552,7 @@ class ConstantOpConversionPattern
     case 64:
       return static_cast<float>(valueAttr.getSplatValue<int64_t>());
     }
-    return 0.0;
+    assert(false && "Unsupported integer type.");
   }
 };
 

diff --git a/test/ttmlir/Dialect/TTNN/simple_constant.mlir b/test/ttmlir/Dialect/TTNN/simple_constant.mlir
@@ -1,24 +1,60 @@
 // RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline %s | FileCheck %s
 #any_device = #tt.operand_constraint<dram|l1|scalar|tile|any_device|any_device_tile>
 module attributes {} {
+  func.func @test_empty_int8() -> tensor<64x128xi8> {
+    %0 = "ttir.constant"() <{value = dense<0> : tensor<64x128xi8>}> : () -> tensor<64x128xi8>
+    // CHECK: %[[C:.*]] = "ttnn.empty"[[C:.*]]
+    return %0 : tensor<64x128xi8>
+  }
+
+  func.func @test_empty_int16() -> tensor<64x128xi16> {
+    %0 = "ttir.constant"() <{value = dense<0> : tensor<64x128xi16>}> : () -> tensor<64x128xi16>
+    // CHECK: %[[C:.*]] = "ttnn.empty"[[C:.*]]
+    return %0 : tensor<64x128xi16>
+  }
+
   func.func @test_empty_int() -> tensor<64x128xi32> {
     %0 = "ttir.constant"() <{value = dense<0> : tensor<64x128xi32>}> : () -> tensor<64x128xi32>
     // CHECK: %[[C:.*]] = "ttnn.empty"[[C:.*]]
     return %0 : tensor<64x128xi32>
   }
 
+  func.func @test_empty_bfloat16() -> tensor<64x128xbf16> {
+    %0 = "ttir.constant"() <{value = dense<0.000000e+00> : tensor<64x128xbf16>}> : () -> tensor<64x128xbf16>
+    // CHECK: %[[C:.*]] = "ttnn.empty"[[C:.*]]
+    return %0 : tensor<64x128xbf16>
+  }
+
   func.func @test_empty_float() -> tensor<64x128xf32> {
     %0 = "ttir.constant"() <{value = dense<0.000000e+00> : tensor<64x128xf32>}> : () -> tensor<64x128xf32>
     // CHECK: %[[C:.*]] = "ttnn.empty"[[C:.*]]
     return %0 : tensor<64x128xf32>
   }
 
+  func.func @test_full_int8() -> tensor<64x128xi8> {
+    %0 = "ttir.constant"() <{value = dense<1> : tensor<64x128xi8>}> : () -> tensor<64x128xi8>
+    // CHECK: %[[C:.*]] = "ttnn.full"[[C:.*]]
+    return %0 : tensor<64x128xi8>
+  }
+
+  func.func @test_full_int16() -> tensor<64x128xi16> {
+    %0 = "ttir.constant"() <{value = dense<1> : tensor<64x128xi16>}> : () -> tensor<64x128xi16>
+    // CHECK: %[[C:.*]] = "ttnn.full"[[C:.*]]
+    return %0 : tensor<64x128xi16>
+  }
+
   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_bfloat16() -> tensor<64x128xbf16> {
+    %0 = "ttir.constant"() <{value = dense<1.000000e+00> : tensor<64x128xbf16>}> : () -> tensor<64x128xbf16>
+    // CHECK: %[[C:.*]] = "ttnn.full"[[C:.*]]
+    return %0 : tensor<64x128xbf16>
+  }
+
   func.func @test_full_float() -> tensor<64x128xf32> {
     %0 = "ttir.constant"() <{value = dense<1.000000e+00> : tensor<64x128xf32>}> : () -> tensor<64x128xf32>
     // CHECK: %[[C:.*]] = "ttnn.full"[[C:.*]]