[mlir][arith] Check for valid IR in BitcastOp::fold.

This PR prevents `BitcastOp::fold` to create invalid `IntegerAttr` and
`FloatAttr` values, which result in failed assertions. This can happen
if the input IR is invalid. The PR adds tests for whether the
to-be-created attribute verifies and returns early from `fold`if it
doesn't.

Signed-off-by: Ingo Müller <[email protected]>

Author: ingomueller-net

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

@@ -1740,9 +1740,24 @@ OpFoldResult arith::BitcastOp::fold(FoldAdaptor adaptor) {
                    ? llvm::cast<FloatAttr>(operand).getValue().bitcastToAPInt()
                    : llvm::cast<IntegerAttr>(operand).getValue();
 
-  if (auto resFloatType = llvm::dyn_cast<FloatType>(resType))
-    return FloatAttr::get(resType,
-                          APFloat(resFloatType.getFloatSemantics(), bits));
+  /// If bitwidth aren't the same, don't fold.
+  if (resType.getIntOrFloatBitWidth() != bits.getBitWidth())
+    return {};
+
+  MLIRContext *ctx = getContext();
+  auto emitErrorFn = [=] { return ::emitError(UnknownLoc::get(ctx)); };
+
+  if (auto resFloatType = llvm::dyn_cast<FloatType>(resType)) {
+    /// If bits don't represent a valid float, don't fold.
+    APFloat floatBits(resFloatType.getFloatSemantics(), bits);
+    if (failed(FloatAttr::verify(emitErrorFn, resType, floatBits)))
+      return {};
+    return FloatAttr::get(resType, floatBits);
+  }
+
+  /// If bits don't represent a valid integer, don't fold.
+  if (failed(IntegerAttr::verify(emitErrorFn, resType, bits)))
+    return {};
   return IntegerAttr::get(resType, bits);
 }
 

mlir/unittests/Dialect/Arith/ArithOpsFoldersTest.cpp

@@ -0,0 +1,46 @@
+//===- ArithOpsFoldersTest.cpp - unit tests for arith op folders ----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Verifier.h"
+#include "gtest/gtest.h"
+
+using namespace mlir;
+
+namespace {
+// Tests a regression that made `BitcastOp::fold` crash on invalid input IR, see
+// #100743,
+TEST(BitcastOpTest, FoldInteger) {
+  MLIRContext context;
+  context.loadDialect<arith::ArithDialect>();
+  auto loc = UnknownLoc::get(&context);
+  auto module = ModuleOp::create(loc);
+  OpBuilder builder(module.getBodyRegion());
+  Value i32Val = builder.create<arith::ConstantOp>(
+      loc, builder.getI32Type(), builder.getI32IntegerAttr(0));
+  // This would create an invalid op: `bitcast` can't cast different bitwidths.
+  builder.createOrFold<arith::BitcastOp>(loc, builder.getI64Type(), i32Val);
+  ASSERT_TRUE(failed(verify(module)));
+}
+
+// Tests a regression that made `BitcastOp::fold` crash on invalid input IR, see
+// #100743,
+TEST(BitcastOpTest, FoldFloat) {
+  MLIRContext context;
+  context.loadDialect<arith::ArithDialect>();
+  auto loc = UnknownLoc::get(&context);
+  auto module = ModuleOp::create(loc);
+  OpBuilder builder(module.getBodyRegion());
+  Value f32Val = builder.create<arith::ConstantOp>(loc, builder.getF32Type(),
+                                                   builder.getF32FloatAttr(0));
+  // This would create an invalid op: `bitcast` can't cast different bitwidths.
+  builder.createOrFold<arith::BitcastOp>(loc, builder.getF64Type(), f32Val);
+  ASSERT_TRUE(failed(verify(module)));
+}
+} // namespace

mlir/unittests/Dialect/Arith/CMakeLists.txt

@@ -0,0 +1,7 @@
+add_mlir_unittest(MLIRArithOpsTests
+  ArithOpsFoldersTest.cpp
+)
+target_link_libraries(MLIRArithOpsTests
+  PRIVATE
+  MLIRArithDialect
+)

mlir/unittests/Dialect/CMakeLists.txt

@@ -6,6 +6,7 @@ target_link_libraries(MLIRDialectTests
   MLIRIR
   MLIRDialect)
 
+add_subdirectory(Arith)
 add_subdirectory(ArmSME)
 add_subdirectory(Index)
 add_subdirectory(LLVMIR)