[mlir] Fix for MSVC bool splat issue encountered in Windows Bazel build

[GleasonK]

Summary

When building MLIR using bazel on windows with MSVC2019, bool splats
were being created incorrectly:

dense<[true,true,true,true]> : tensor<4xi1>
-(parse with mlir-opt)-> dense<[true, false, false, false]> : tensor<4xi1>

Appears that a windows bazel build produces a corrupt DenseIntOrFPElementsAttr.
Not able to repro using MSVC with cmake.

Issue first discovered here:
https://github.com/google/jax/issues/16394

Added test point for reproduction:

$ bazel test @llvm-project//mlir/unittests:ir_tests --test_arg=--gtest_filter=DenseSplatTest.BoolSplatSmall

[rriddle]

mlir/lib/IR/AttributeDetail.h

196–204

Could we just switch to using an enum for this?

[GleasonK]

mlir/lib/IR/AttributeDetail.h

196–204

Meaning switch from using these constexpr values to enum values? I believe the reason static member variables are used is since KeyTy gets passed around, but uses an ArrayRef<char> for data, so needs these values to persist between calls to getKey and uses of the key, like construct.

[rriddle]

mlir/lib/IR/AttributeDetail.h

196–204

Yeah, I'm paging back in why I even wrote it this way in the first place. Do we need to if/def _MSC_VER? Can we just align around that version?

[jpienaar]

mlir/lib/IR/AttributeDetail.h

196–204

I think just for discussion, I think we should just have one here. I almost wonder if we should just use pre-C++17 approach of static const char declaration here and definition in body. (apparently this failed with const added MSC_VER side)

[GleasonK]

mlir/lib/IR/AttributeDetail.h

196–204

(apparently this failed with const added MSC_VER side)

That's correct, even static inline const char failed as well.

I almost wonder if we should just use pre-C++17 approach of static const char declaration here and definition in body.

Just confirmed that this works, tested by putting the definition in BuiltinAttributes.cpp and changed this to just declaration. This is probably a better solution to align on than what I have coded up now, since these member variables really should be const.

[GleasonK]

Use pre-C++17 declaration for static member variables for splats

[GleasonK]

• Use pre-C++17 declaration of static member variables for splats
• clang-format

mlir/lib/IR/AttributeDetail.h

static KeyTy getKeyForBoolData(ShapedType ty, ArrayRef<char> data,
     return llvm::all_of(splatData, [mask](char c) { return c == mask; })
                ? getKeyForSplatBoolData(ty, splatValue)
                : KeyTy(ty, data, llvm::hash_value(data));
   }
 
   /// Return a key to use for a boolean splat of the given value.
   static KeyTy getKeyForSplatBoolData(ShapedType type, bool splatValue) {
     const char &splatData = splatValue ? kSplatTrue : kSplatFalse;
-    llvm::dbgs() << "getKeyFor " << (void*)&splatData << " -> " << (static_cast<int>(splatData)&0xFF) << '\n';
     return KeyTy(type, splatData, llvm::hash_value(splatData),
                  /*isSplat=*/true);
   }
 
   /// Hash the key for the storage.
   static llvm::hash_code hashKey(const KeyTy &key) {
     return llvm::hash_combine(key.type, key.hashCode);
   }
 
   /// Construct a new storage instance.
   static DenseIntOrFPElementsAttrStorage *
   construct(AttributeStorageAllocator &allocator, KeyTy key) {
     // If the data buffer is non-empty, we copy it into the allocator with a
     // 64-bit alignment.
-    llvm::dbgs() << "construct " << (void*)key.data.data() << " -> " << (static_cast<int>(key.data[0])&0xFF) << '\n';
     ArrayRef<char> copy, data = key.data;
     if (!data.empty()) {
       char *rawData = reinterpret_cast<char *>(
           allocator.allocate(data.size(), alignof(uint64_t)));
       std::memcpy(rawData, data.data(), data.size());
       copy = ArrayRef<char>(rawData, data.size());
     }
 
     return new (allocator.allocate<DenseIntOrFPElementsAttrStorage>())
         DenseIntOrFPElementsAttrStorage(key.type, copy, key.isSplat);
   }
 
   ArrayRef<char> data;
 
   /// The values used to denote a boolean splat value.
-#if defined(_MSC_VER)
-  // Not safe in MSVC to use static constexpr variables by reference,
-  // it seems like they are not persisted, similar for `static inline const`.
-  static inline char kSplatTrue = ~0;
-  static inline char kSplatFalse = 0;
-#else
-  static constexpr char kSplatTrue = ~0;
-  static constexpr char kSplatFalse = 0;
-#endif
+  // Not using constexpr declaration since Windows may inline these values,
+  // which makes it unsafe to refer to them by reference in KeyTy.
+  static const char kSplatTrue;
+  static const char kSplatFalse;
 };
 
 /// An attribute representing a reference to a dense vector or tensor object
 /// containing strings.
 struct DenseStringElementsAttrStorage : public DenseElementsAttributeStorage {
   DenseStringElementsAttrStorage(ShapedType ty, ArrayRef<StringRef> data,
                                  bool isSplat = false)
       : DenseElementsAttributeStorage(ty, isSplat), data(data) {}
 
   struct KeyTy {
     KeyTy(ShapedType type, ArrayRef<StringRef> data, llvm::hash_code hashCode,
           bool isSplat = false)
         : type(type), data(data), hashCode(hashCode), isSplat(isSplat) {}

mlir/lib/IR/BuiltinAttributes.cpp

LogicalResult OpaqueAttr::verify(function_ref<InFlightDiagnostic()> emitError,
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // DenseElementsAttr Utilities
 //===----------------------------------------------------------------------===//
 
+const char DenseIntOrFPElementsAttrStorage::kSplatTrue = ~0;
+const char DenseIntOrFPElementsAttrStorage::kSplatFalse = 0;
+
 /// Get the bitwidth of a dense element type within the buffer.
 /// DenseElementsAttr requires bitwidths greater than 1 to be aligned by 8.
 static size_t getDenseElementStorageWidth(size_t origWidth) {
   return origWidth == 1 ? origWidth : llvm::alignTo<8>(origWidth);
 }
 static size_t getDenseElementStorageWidth(Type elementType) {
   return getDenseElementStorageWidth(getDenseElementBitWidth(elementType));
 }

mlir/unittests/IR/AttributeTest.cpp

TEST(DenseSplatTest, BoolSplatRawRoundtrip) {
   DenseElementsAttr trueSplatFromRaw =
       DenseElementsAttr::getFromRawBuffer(shape, trueSplat.getRawData());
   EXPECT_TRUE(trueSplatFromRaw.isSplat());
 
   EXPECT_EQ(trueSplat, trueSplatFromRaw);
 }
 
 TEST(DenseSplatTest, BoolSplatSmall) {
-  // There is a windows bug when using semi-packed values like `0b1111`
-  // to represent a 4xi1 splat and building using bazel
   MLIRContext context;
   Builder builder(&context);
-  std::vector<char> data{15};
+
+  // Check that splats that don't fill entire byte are handled properly.
   auto tensorType = RankedTensorType::get({4}, builder.getI1Type());
+  std::vector<char> data{0b00001111};
   auto trueSplatFromRaw = DenseIntOrFPElementsAttr::getFromRawBuffer(tensorType, data);
   EXPECT_TRUE(trueSplatFromRaw.isSplat());
   DenseElementsAttr trueSplat = DenseElementsAttr::get(tensorType, true);
   EXPECT_EQ(trueSplat, trueSplatFromRaw);
 }
 
 TEST(DenseSplatTest, LargeBoolSplat) {
   constexpr int64_t boolCount = 56;

mlir/unittests/IR_AttributeTest.md

-# Debug info:
-
-Reproduced with MSVC2019 / MSC2022, using compile flags:
-
-C:\Program Files (x86)\Microsoft Visual Studio\2019\BuildTools\VC\Tools\MSVC\14.29.30133\bin\HostX64\x64\cl.exe /nologo /DCOMPILER_MSVC /DNOMINMAX /D_WIN32_WINNT=0x0601 /D_CRT_SECURE_NO_DEPRECATE /D_CRT_SECURE_NO_WARNINGS /bigobj /Zm500 /EHsc /wd4351 /wd4291 /wd4250 /wd4996 /Iexternal/llvm-project /Ibazel-out/x64_windows-opt/bin/external/llvm-project /Ibazel-out/x64_windows-opt/bin/external/llvm-project/mlir/_virtual_includes/BuiltinAttributeInterfacesIncGen /Ibazel-out/x64_windows-opt/bin/external/llvm-project/mlir/_virtual_includes/BuiltinAttributesIncGen /Ibazel-out/x64_windows-opt/bin/external/llvm-project/mlir/_virtual_includes/BuiltinDialectBytecodeGen /Ibazel-out/x64_windows-opt/bin/external/llvm-project/mlir/_virtual_includes/BuiltinDialectIncGen /Ibazel-out/x64_windows-opt/bin/external/llvm-project/mlir/_virtual_includes/BuiltinLocationAttributesIncGen /Ibazel-out/x64_windows-opt/bin/external/llvm-project/mlir/_virtual_includes/BuiltinOpsIncGen /Ibazel-out/x64_windows-opt/bin/external/llvm-project/mlir/_virtual_includes/BuiltinTypeInterfacesIncGen /Ibazel-out/x64_windows-opt/bin/external/llvm-project/mlir/_virtual_includes/BuiltinTypesIncGen /Ibazel-out/x64_windows-opt/bin/external/llvm-project/mlir/_virtual_includes/BytecodeOpInterfaceIncGen /Ibazel-out/x64_windows-opt/bin/external/llvm-project/mlir/_virtual_includes/CallOpInterfacesIncGen /Ibazel-out/x64_windows-opt/bin/external/llvm-project/mlir/_virtual_includes/FunctionInterfacesIncGen /Ibazel-out/x64_windows-opt/bin/external/llvm-project/mlir/_virtual_includes/InferTypeOpInterfaceIncGen /Ibazel-out/x64_windows-opt/bin/external/llvm-project/mlir/_virtual_includes/OpAsmInterfaceIncGen /Ibazel-out/x64_windows-opt/bin/external/llvm-project/mlir/_virtual_includes/RegionKindInterfaceIncGen /Ibazel-out/x64_windows-opt/bin/external/llvm-project/mlir/_virtual_includes/SideEffectInterfacesIncGen /Ibazel-out/x64_windows-opt/bin/external/llvm-project/mlir/_virtual_includes/SymbolInterfacesIncGen /Ibazel-out/x64_windows-opt/bin/external/llvm-project/mlir/_virtual_includes/TensorEncodingIncGen /Iexternal/llvm-project/mlir/include /Ibazel-out/x64_windows-opt/bin/external/llvm-project/mlir/include /Iexternal/llvm-project/llvm/include /Ibazel-out/x64_windows-opt/bin/external/llvm-project/llvm/include /D_CRT_SECURE_NO_DEPRECATE /D_CRT_SECURE_NO_WARNINGS /D_CRT_NONSTDC_NO_DEPRECATE /D_CRT_NONSTDC_NO_WARNINGS /D_SCL_SECURE_NO_DEPRECATE /D_SCL_SECURE_NO_WARNINGS /DUNICODE /D_UNICODE /DLTDL_SHLIB_EXT=".dll" /DLLVM_PLUGIN_EXT=".dll" /DLLVM_NATIVE_ARCH="X86" /DLLVM_NATIVE_ASMPARSER=LLVMInitializeX86AsmParser /DLLVM_NATIVE_ASMPRINTER=LLVMInitializeX86AsmPrinter /DLLVM_NATIVE_DISASSEMBLER=LLVMInitializeX86Disassembler /DLLVM_NATIVE_TARGET=LLVMInitializeX86Target /DLLVM_NATIVE_TARGETINFO=LLVMInitializeX86TargetInfo /DLLVM_NATIVE_TARGETMC=LLVMInitializeX86TargetMC /DLLVM_NATIVE_TARGETMCA=LLVMInitializeX86TargetMCA /DLLVM_HOST_TRIPLE="x86_64-pc-win32" /DLLVM_DEFAULT_TARGET_TRIPLE="x86_64-pc-win32" /DLLVM_VERSION_MAJOR=17 /DLLVM_VERSION_MINOR=0 /DLLVM_VERSION_PATCH=0 /DLLVM_VERSION_STRING="17.0.0git" /D__STDC_LIMIT_MACROS /D__STDC_CONSTANT_MACROS /D__STDC_FORMAT_MACROS /DBLAKE3_USE_NEON=0 /DBLAKE3_NO_AVX2 /DBLAKE3_NO_AVX512 /DBLAKE3_NO_SSE2 /DBLAKE3_NO_SSE41 /DBAZEL_CURRENT_REPOSITORY="llvm-project" /showIncludes /MD /O2 /Oy- /DNDEBUG /wd4117 -D__DATE__="redacted" -D__TIMESTAMP__="redacted" -D__TIME__="redacted" /Gy /Gw /D_USE_MATH_DEFINES -DWIN32_LEAN_AND_MEAN -DNOGDI /Zc:preprocessor -DMLIR_PYTHON_PACKAGE_PREFIX=jaxlib.mlir. /std:c++17 /Fobazel-out/x64_windows-opt/bin/external/llvm-project/mlir/_objs/IR/Attributes.obj /c external/llvm-project/mlir/lib/IR/Attributes.cpp
-
-# Test Log:
-
-```
-$ bazel test @llvm-project//mlir/unittests:ir_tests --test_arg=--gtest_filter=DenseSplatTest.BoolSplatSmall
-
-## Without fix, using `static constexpr`:
-[ RUN      ] DenseSplatTest.BoolSplatSmall
-getKeyFor 0xa52db8f650 -> 255  // <-- trueSplatFromRaw 
-construct 0xa52db8f650 -> 208  // <-- trueSplatFromRaw 
-getKeyFor 0xa52db8f6b0 -> 255  // <-- trueSplat
-construct 0xa52db8f6b0 -> 1    // <-- trueSplat
-external/llvm-project/mlir/unittests/IR/AttributeTest.cpp(88): error: Expected equality of these values:
-  trueSplat
-    Which is: dense<true> : tensor<4xi1>
-  trueSplatFromRaw
-    Which is: dense<false> : tensor<4xi1>
-[  FAILED  ] DenseSplatTest.BoolSplatSmall (1 ms)[ RUN      ] DenseSplatTest.BoolSplatSmall
-
-## With fix, using `static inline`:
-[ RUN      ] DenseSplatTest.BoolSplatSmall
-getKeyFor 0x7ff763ac8010 -> 255  // <-- trueSplatFromRaw
-construct 0x7ff763ac8010 -> 255  // <-- trueSplatFromRaw
-getKeyFor 0x7ff763ac8010 -> 255  // <-- trueSplat
-[       OK ] DenseSplatTest.BoolSplatSmall (0 ms)
-```