fixing constant for TTIR->TTNN conversion

tenstorrent · Nov 6, 2024 · 90c5561 · 90c5561
1 parent 6bad515
commit 90c5561
Showing 1 changed file with 23 additions and 3 deletions.
diff --git a/lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp b/lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp
@@ -23,6 +23,7 @@
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/LogicalResult.h"
+#include <cstdint>
 
 using namespace mlir;
 using namespace mlir::tt;
@@ -503,9 +504,11 @@ class ConstantOpConversionPattern
 
     if (valueAttr.isSplat()) {
       Value device = getOrInsertDevice(rewriter, op);
-      float fillValue = valueAttr.getElementType().isInteger()
-                            ? static_cast<float>(valueAttr.getSplatValue<int>())
-                            : valueAttr.getSplatValue<float>();
+      float fillValue =
+          valueAttr.getElementType().isInteger()
+              ? getIntegerValue(
+                    valueAttr) // static_cast<float>(valueAttr.getSplatValue<int>())
+              : valueAttr.getSplatValue<mlir::APFloat>().convertToFloat();
       if (fillValue == 0) {
         rewriter.replaceOpWithNewOp<tensor::EmptyOp>(
             op, this->getTypeConverter()->convertType(op.getType()), device);
@@ -536,6 +539,23 @@ class ConstantOpConversionPattern
 
     return success();
   }
+
+  float getIntegerValue(mlir::ElementsAttr valueAttr) const {
+    size_t bitWidth = valueAttr.getElementType().getIntOrFloatBitWidth();
+    switch (bitWidth) {
+    case 1:
+      return static_cast<float>(valueAttr.getSplatValue<bool>());
+    case 8:
+      return static_cast<float>(valueAttr.getSplatValue<int8_t>());
+    case 16:
+      return static_cast<float>(valueAttr.getSplatValue<int16_t>());
+    case 32:
+      return static_cast<float>(valueAttr.getSplatValue<int>());
+    case 64:
+      return static_cast<float>(valueAttr.getSplatValue<int64_t>());
+    }
+    return 0.0;
+  }
 };
 
 } // namespace