[mlir][Arithmetic] Add common constant folder function for type cast ops.

[jacquesguan]

Summary

This revision replaces current type cast constant folder with a new common type cast constant folder function template.
It will cover all former folder and support fold the constant splat and vector.

[ftynse]

Could you please add a longer commit description that motivates this change? From what I can see, it replaces a small amount of boilerplate with a longer and complex utility function; not sure if the trade off is the right one here.

mlir/include/mlir/Dialect/CommonFolders.h

117

Nit: either const & or &&.

118

140

Nit: no else after return.

[jacquesguan]

In D123489#3444780, @ftynse wrote:

Could you please add a longer commit description that motivates this change? From what I can see, it replaces a small amount of boilerplate with a longer and complex utility function; not sure if the trade off is the right one here.

OK, after the replacing, we could fold constant splat and contant vector, not only the constant integer or float.

[ftynse]

Then please (1) update the commit/diff description accordingly (the short version better highlight the added functionality, and the long version should explain why the change is made) and (2) add tests that exercise the vector versions.

[jacquesguan]

Address comment.

[jacquesguan]

In D123489#3444796, @ftynse wrote:

Then please (1) update the commit/diff description accordingly (the short version better highlight the added functionality, and the long version should explain why the change is made) and (2) add tests that exercise the vector versions.

Done, thanks.

mlir/include/mlir/Dialect/CommonFolders.h

if (operands[0].isa<SplatElementsAttr>()) {
     elementResults.reserve(op.getNumElements());
     for (size_t i = 0, e = op.getNumElements(); i < e; ++i, ++opIt)
       elementResults.push_back(calculate(*opIt));
     return DenseElementsAttr::get(op.getType(), elementResults);
   }
   return {};
 }
 
+template <
+    class AttrElementT, class TargetAttrElementT,
+    class ElementValueT = typename AttrElementT::ValueType,
+    class TargetElementValueT = typename TargetAttrElementT::ValueType,
+    class CalculationT = function_ref<TargetElementValueT(ElementValueT, bool)>>
+Attribute constFoldCastOp(ArrayRef<Attribute> operands, Type resType,
+                          const CalculationT &&calculate) {
+  assert(operands.size() == 1 && "Cast op takes one operands");
+  if (!operands[0])
+    return {};
+
+  if (operands[0].isa<AttrElementT>()) {
+    auto op = operands[0].cast<AttrElementT>();
+    bool castStatus = true;
+    auto res = calculate(op.getValue(), castStatus);
+    if (!castStatus)
+      return {};
+    return TargetAttrElementT::get(resType, res);
+  }
+  if (operands[0].isa<SplatElementsAttr>()) {
+    // Both operands are splats so we can avoid expanding the values out and
+    // just fold based on the splat value.
+    auto op = operands[0].cast<SplatElementsAttr>();
+    bool castStatus = true;
+    auto elementResult =
+        calculate(op.getSplatValue<ElementValueT>(), castStatus);
+    if (!castStatus)
+      return {};
+    return DenseElementsAttr::get(resType, elementResult);
+  } else if (operands[0].isa<ElementsAttr>()) {
+    // Operands are ElementsAttr-derived; perform an element-wise fold by
+    // expanding the values.
+    auto op = operands[0].cast<ElementsAttr>();
+    bool castStatus = true;
+    auto opIt = op.value_begin<ElementValueT>();
+    SmallVector<TargetElementValueT> elementResults;
+    elementResults.reserve(op.getNumElements());
+    for (size_t i = 0, e = op.getNumElements(); i < e; ++i, ++opIt) {
+      auto elt = calculate(*opIt, castStatus);
+      if (!castStatus)
+        return {};
+      elementResults.push_back(elt);
+    }
+
+    return DenseElementsAttr::get(resType, elementResults);
+  }
+  return {};
+}
+
 } // namespace mlir
 
 #endif // MLIR_DIALECT_COMMONFOLDERS_H

mlir/lib/Dialect/Arithmetic/IR/ArithmeticOps.cpp

return WidthComparator<unsigned>()(dstType.getIntOrFloatBitWidth(),
                                      srcType.getIntOrFloatBitWidth());
 }
 
 //===----------------------------------------------------------------------===//
 // ExtUIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::ExtUIOp::fold(ArrayRef<Attribute> operands) {
-  if (auto lhs = operands[0].dyn_cast_or_null<IntegerAttr>())
-    return IntegerAttr::get(
-        getType(), lhs.getValue().zext(getType().getIntOrFloatBitWidth()));
-
   if (auto lhs = getIn().getDefiningOp<ExtUIOp>()) {
     getInMutable().assign(lhs.getIn());
     return getResult();
   }
-
-  return {};
+  Type resType = getType();
+  return constFoldCastOp<IntegerAttr, IntegerAttr>(
+      operands, getType(), [&resType](const APInt &a, bool &castStatus) {
+        return a.zext(resType.getIntOrFloatBitWidth());
+      });
 }
 
 bool arith::ExtUIOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
   return checkWidthChangeCast<std::greater, IntegerType>(inputs, outputs);
 }
 
 LogicalResult arith::ExtUIOp::verify() {
   return verifyExtOp<IntegerType>(*this);
 }
 
 //===----------------------------------------------------------------------===//
 // ExtSIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::ExtSIOp::fold(ArrayRef<Attribute> operands) {
-  if (auto lhs = operands[0].dyn_cast_or_null<IntegerAttr>())
-    return IntegerAttr::get(
-        getType(), lhs.getValue().sext(getType().getIntOrFloatBitWidth()));
-
   if (auto lhs = getIn().getDefiningOp<ExtSIOp>()) {
     getInMutable().assign(lhs.getIn());
     return getResult();
   }
-
-  return {};
+  Type resType = getType();
+  return constFoldCastOp<IntegerAttr, IntegerAttr>(
+      operands, getType(), [&resType](const APInt &a, bool &castStatus) {
+        return a.sext(resType.getIntOrFloatBitWidth());
+      });
 }
 
 bool arith::ExtSIOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
   return checkWidthChangeCast<std::greater, IntegerType>(inputs, outputs);
 }
 
 void arith::ExtSIOp::getCanonicalizationPatterns(
     RewritePatternSet &patterns, MLIRContext *context) {
if (matchPattern(getOperand(), m_Op<arith::ExtUIOp>()) ||
     return getOperand().getDefiningOp()->getOperand(0);
 
   // trunci(trunci(a)) -> trunci(a))
   if (matchPattern(getOperand(), m_Op<arith::TruncIOp>())) {
     setOperand(getOperand().getDefiningOp()->getOperand(0));
     return getResult();
   }
 
-  if (!operands[0])
-    return {};
-
-  if (auto lhs = operands[0].dyn_cast<IntegerAttr>()) {
-    return IntegerAttr::get(
-        getType(), lhs.getValue().trunc(getType().getIntOrFloatBitWidth()));
-  }
-
-  return {};
+  Type resType = getType();
+  return constFoldCastOp<IntegerAttr, IntegerAttr>(
+      operands, getType(), [&resType](const APInt &a, bool &castStatus) {
+        return a.trunc(resType.getIntOrFloatBitWidth());
+      });
 }
 
 bool arith::TruncIOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
   return checkWidthChangeCast<std::less, IntegerType>(inputs, outputs);
 }
 
 LogicalResult arith::TruncIOp::verify() {
   return verifyTruncateOp<IntegerType>(*this);

 // UIToFPOp
 //===----------------------------------------------------------------------===//
 
 bool arith::UIToFPOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
   return checkIntFloatCast<IntegerType, FloatType>(inputs, outputs);
 }
 
 OpFoldResult arith::UIToFPOp::fold(ArrayRef<Attribute> operands) {
-  if (auto lhs = operands[0].dyn_cast_or_null<IntegerAttr>()) {
-    const APInt &api = lhs.getValue();
-    FloatType floatTy = getType().cast<FloatType>();
+  Type resType = getType();
+  return constFoldCastOp<IntegerAttr, FloatAttr>(
+      operands, getType(), [&resType](const APInt &a, bool &castStatus) {
+        FloatType floatTy = resType.cast<FloatType>();
         APFloat apf(floatTy.getFloatSemantics(),
                     APInt::getZero(floatTy.getWidth()));
-    apf.convertFromAPInt(api, /*IsSigned=*/false, APFloat::rmNearestTiesToEven);
-    return FloatAttr::get(floatTy, apf);
-  }
-  return {};
+        apf.convertFromAPInt(a, /*IsSigned=*/false,
+                             APFloat::rmNearestTiesToEven);
+        return apf;
+      });
 }
 
 //===----------------------------------------------------------------------===//
 // SIToFPOp
 //===----------------------------------------------------------------------===//
 
 bool arith::SIToFPOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
   return checkIntFloatCast<IntegerType, FloatType>(inputs, outputs);
 }
 
 OpFoldResult arith::SIToFPOp::fold(ArrayRef<Attribute> operands) {
-  if (auto lhs = operands[0].dyn_cast_or_null<IntegerAttr>()) {
-    const APInt &api = lhs.getValue();
-    FloatType floatTy = getType().cast<FloatType>();
+  Type resType = getType();
+  return constFoldCastOp<IntegerAttr, FloatAttr>(
+      operands, getType(), [&resType](const APInt &a, bool &castStatus) {
+        FloatType floatTy = resType.cast<FloatType>();
         APFloat apf(floatTy.getFloatSemantics(),
                     APInt::getZero(floatTy.getWidth()));
-    apf.convertFromAPInt(api, /*IsSigned=*/true, APFloat::rmNearestTiesToEven);
-    return FloatAttr::get(floatTy, apf);
-  }
-  return {};
+        apf.convertFromAPInt(a, /*IsSigned=*/true,
+                             APFloat::rmNearestTiesToEven);
+        return apf;
+      });
 }
 //===----------------------------------------------------------------------===//
 // FPToUIOp
 //===----------------------------------------------------------------------===//
 
 bool arith::FPToUIOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
   return checkIntFloatCast<FloatType, IntegerType>(inputs, outputs);
 }
 
 OpFoldResult arith::FPToUIOp::fold(ArrayRef<Attribute> operands) {
-  if (auto lhs = operands[0].dyn_cast_or_null<FloatAttr>()) {
-    const APFloat &apf = lhs.getValue();
-    IntegerType intTy = getType().cast<IntegerType>();
+  Type resType = getType();
+  return constFoldCastOp<FloatAttr, IntegerAttr>(
+      operands, getType(), [&resType](const APFloat &a, bool &castStatus) {
+        IntegerType intTy = resType.cast<IntegerType>();
         bool ignored;
         APSInt api(intTy.getWidth(), /*isUnsigned=*/true);
-    if (APFloat::opInvalidOp ==
-        apf.convertToInteger(api, APFloat::rmTowardZero, &ignored)) {
-      // Undefined behavior invoked - the destination type can't represent
-      // the input constant.
-      return {};
-    }
-    return IntegerAttr::get(getType(), api);
-  }
-
-  return {};
+        castStatus = APFloat::opInvalidOp !=
+                     a.convertToInteger(api, APFloat::rmTowardZero, &ignored);
+        return api;
+      });
 }
 
 //===----------------------------------------------------------------------===//
 // FPToSIOp
 //===----------------------------------------------------------------------===//
 
 bool arith::FPToSIOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
   return checkIntFloatCast<FloatType, IntegerType>(inputs, outputs);
 }
 
 OpFoldResult arith::FPToSIOp::fold(ArrayRef<Attribute> operands) {
-  if (auto lhs = operands[0].dyn_cast_or_null<FloatAttr>()) {
-    const APFloat &apf = lhs.getValue();
-    IntegerType intTy = getType().cast<IntegerType>();
+  Type resType = getType();
+  return constFoldCastOp<FloatAttr, IntegerAttr>(
+      operands, getType(), [&resType](const APFloat &a, bool &castStatus) {
+        IntegerType intTy = resType.cast<IntegerType>();
         bool ignored;
         APSInt api(intTy.getWidth(), /*isUnsigned=*/false);
-    if (APFloat::opInvalidOp ==
-        apf.convertToInteger(api, APFloat::rmTowardZero, &ignored)) {
-      // Undefined behavior invoked - the destination type can't represent
-      // the input constant.
-      return {};
-    }
-    return IntegerAttr::get(getType(), api);
-  }
-
-  return {};
+        castStatus = APFloat::opInvalidOp !=
+                     a.convertToInteger(api, APFloat::rmTowardZero, &ignored);
+        return api;
+      });
 }
 
 //===----------------------------------------------------------------------===//
 // IndexCastOp
 //===----------------------------------------------------------------------===//
 
 bool arith::IndexCastOp::areCastCompatible(TypeRange inputs,
                                            TypeRange outputs) {