[Mlir-commits] [mlir] 998003e - [mlir] Fix for MSVC bool splat issue encountered.

[Jacques Pienaar]

Author: Kevin Gleason
Date: 2023-07-24T20:45:43-07:00
New Revision: 998003eaf7b91b5620eeb0b77c81654b9f36efb7

URL: https://github.com/llvm/llvm-project/commit/998003eaf7b91b5620eeb0b77c81654b9f36efb7
DIFF: https://github.com/llvm/llvm-project/commit/998003eaf7b91b5620eeb0b77c81654b9f36efb7.diff

LOG: [mlir] Fix for MSVC bool splat issue encountered.

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.
Unable to repro using MSVC and 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
```

Differential Revision: https://reviews.llvm.org/D155745

Added: 
    

Modified: 
    mlir/lib/IR/AttributeDetail.h
    mlir/lib/IR/BuiltinAttributes.cpp
    mlir/unittests/IR/AttributeTest.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/IR/AttributeDetail.h b/mlir/lib/IR/AttributeDetail.h
index b76da1b91bff9d..dcd24af0107ddf 100644
--- a/mlir/lib/IR/AttributeDetail.h
+++ b/mlir/lib/IR/AttributeDetail.h
@@ -193,8 +193,11 @@ struct DenseIntOrFPElementsAttrStorage : public DenseElementsAttributeStorage {
   ArrayRef<char> data;
 
   /// The values used to denote a boolean splat value.
-  static constexpr char kSplatTrue = ~0;
-  static constexpr char kSplatFalse = 0;
+  // This is not using constexpr declaration due to compilation failure
+  // encountered with MSVC where it would inline these values, which makes it
+  // unsafe to refer by reference in KeyTy.
+  static const char kSplatTrue;
+  static const char kSplatFalse;
 };
 
 /// An attribute representing a reference to a dense vector or tensor object

diff  --git a/mlir/lib/IR/BuiltinAttributes.cpp b/mlir/lib/IR/BuiltinAttributes.cpp
index c670359b02ec9c..7328fccfdcfe6f 100644
--- a/mlir/lib/IR/BuiltinAttributes.cpp
+++ b/mlir/lib/IR/BuiltinAttributes.cpp
@@ -438,6 +438,9 @@ LogicalResult OpaqueAttr::verify(function_ref<InFlightDiagnostic()> emitError,
 // 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) {

diff  --git a/mlir/unittests/IR/AttributeTest.cpp b/mlir/unittests/IR/AttributeTest.cpp
index f01cc026b72cc2..a707dfa58be197 100644
--- a/mlir/unittests/IR/AttributeTest.cpp
+++ b/mlir/unittests/IR/AttributeTest.cpp
@@ -74,6 +74,20 @@ TEST(DenseSplatTest, BoolSplatRawRoundtrip) {
   EXPECT_EQ(trueSplat, trueSplatFromRaw);
 }
 
+TEST(DenseSplatTest, BoolSplatSmall) {
+  MLIRContext context;
+  Builder builder(&context);
+
+  // 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;