Do not store bool constants to file (#2714)

* Do not store bool constants to file Signed-off-by: Tung D. Le <tung@jp.ibm.com> --------- Signed-off-by: Tung D. Le <tung@jp.ibm.com>
onnx · Feb 21, 2024 · 7dc9f73 · 7dc9f73
1 parent b599d36
commit 7dc9f73
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 0 deletions.
diff --git a/src/Conversion/KrnlToLLVM/ConvertKrnlToLLVM.cpp b/src/Conversion/KrnlToLLVM/ConvertKrnlToLLVM.cpp
@@ -478,6 +478,15 @@ bool extractConstantsToFile(ModuleOp &module, std::string filepath,
     if (isReturnedValue)
       return WalkResult::advance();
 
+    // Ignore constants of bool.
+    // For an unknown reason, enabling constants of bool caused segfault in the
+    // IBM granite.20B model (The model with KV cache) at 1265 input tokens.
+    // See issue https://github.com/onnx/onnx-mlir/issues/2713.
+    if (llvm::cast<MemRefType>(op->getResult(0).getType())
+            .getElementType()
+            .isInteger(1))
+      return WalkResult::advance();
+
     // Get raw data from DenseElementsAttr or DenseResourceElementsAttr.
     ArrayRef<char> rawData = getRawData(op);
     if (rawData.empty())

diff --git a/src/Tools/binary-decoder/BinaryDecoder.cpp b/src/Tools/binary-decoder/BinaryDecoder.cpp
@@ -67,6 +67,20 @@ int printBuffer(std::vector<char> buffer) {
   auto data = std::vector<T>(ptr, ptr + buffer.size() / sizeof(T));
   for (const auto &elem : data)
     std::cout << elem << " ";
+  std::cout << std::endl;
+  return 0;
+}
+
+template <>
+int printBuffer<bool>(std::vector<char> buffer) {
+  const char *rawData = buffer.data();
+  for (unsigned i = 0; i < buffer.size() * 8; i++) {
+    bool b = (rawData[i / CHAR_BIT] & (1 << (i % CHAR_BIT))) != 0;
+    printf("%d", b);
+    if ((i + 1) % 8 == 0)
+      std::cout << " ";
+  }
+  std::cout << std::endl;
   return 0;
 }
 
@@ -87,6 +101,7 @@ int main(int argc, char **argv) {
   if (DataType == ONNX_TYPE)                                                   \
     return printBuffer<CPP_TYPE>(buffer);
 
+  PRINT_BUFFER_FOR_TYPE(onnx::TensorProto::BOOL, bool);
   PRINT_BUFFER_FOR_TYPE(onnx::TensorProto::UINT8, u_int8_t);
   PRINT_BUFFER_FOR_TYPE(onnx::TensorProto::UINT16, u_int16_t);
   PRINT_BUFFER_FOR_TYPE(onnx::TensorProto::INT16, int16_t);